source: branches/numpy_misc/tools/write_large_files/write_big_file.py @ 7244

Last change on this file since 7244 was 6817, checked in by rwilson, 16 years ago

Changed the 'large file' netCDF test programs to use numpy.

  • Property svn:executable set to *
File size: 6.1 KB
Line 
1#!/usr/bin/env python
2################################################################################
3# A test program to write large NetCDF files.
4################################################################################
5
6import sys
7import getopt
8import time
9import numpy as num
10import Scientific.IO.NetCDF as nc
11
12######
13# Various constants
14######
15
16# default modifier and number of files
17DEF_MODIFIER = 'M'
18
19# array slice definitions
20XDIM = 1024*1024
21SLICE_SIZE = XDIM*4     # 4 bytes / int
22MBYTE = 1024*1024
23
24# Dictionary for suffix modifiers
25Suffixes = {'M': 1024*1024,
26            'G': 1024*1024*1024
27           }
28
29######
30# Globals
31######
32
33TimeFirstWrite = False
34TimeVarWrite = False
35TimeFileWrite = False
36
37CloseAfterSlice = False
38CloseAfterVar = False
39
40Verbose = False
41
42######
43# mainline code
44######
45
46##
47# @brief Write a NetCDF file with set number of variables of a defined size.
48# @param variable_size Size of variable, in bytes.
49# @param num_variables Number of required variables.
50def write_file(variable_size, num_variables=1):
51    # set file and variable name masks
52    filename = 'test.nc'
53    varname_mask = 'var_%04d'
54
55    # create a data array slice
56    slice_array = num.ones((XDIM,), 'i')
57
58    # if timing file write, remember start time
59    if TimeFileWrite:
60        start_file_write = time.time()
61
62    fid = nc.NetCDFFile(filename, 'w')
63    fid.createDimension('y', None)
64    fid.createDimension('x', XDIM)
65
66    # num file bytes written counter
67    file_bytes_written = 0
68
69    for i in xrange(num_variables):
70        varname = varname_mask % i
71        fid.createVariable(varname, 'i', ('y', 'x'))
72
73        # num variable bytes written counter
74        var_bytes_written = 0
75
76        if TimeVarWrite:
77            start_var_write = time.time()
78
79        index = 0
80        while var_bytes_written < variable_size:
81            var_bytes_written += SLICE_SIZE
82            file_bytes_written += SLICE_SIZE
83
84            if TimeFirstWrite and index == 0:
85                start_first_var_write = time.time()
86
87            if Verbose:
88                print ('File %s, variable %s, writing slice %d: '
89                       'var=%.1fMiB, file=%.1fMiB' %
90                       (filename, varname, index, float(var_bytes_written)/MBYTE,
91                        float(file_bytes_written)/MBYTE))
92
93            fid.variables[varname][index,:] = slice_array
94
95            if TimeFirstWrite and index == 0:
96                stop_first_var_write = time.time()
97                print ("Time to do first write for variable %s: %.2f sec" %
98                       (varname,
99                        (stop_first_var_write - start_first_var_write)))
100
101            index += 1
102
103            if CloseAfterSlice:
104                fid.close()
105                fid = nc.NetCDFFile(filename, 'a')
106        if CloseAfterVar:
107            fid.close()
108            fid = nc.NetCDFFile(filename, 'a')
109
110        if TimeVarWrite:
111            stop_var_write = time.time()
112            print ('Time to write variable %s: %.2f sec' %
113                   (varname, (stop_first_var_write - start_first_var_write)))
114
115    fid.close()
116
117    if TimeFileWrite:
118        stop_file_write = time.time()
119        print ('Time to write file: %.2f sec' %
120               (stop_file_write - start_file_write))
121
122##
123# @brief Provide help for the befuddled user.
124# @return Doesn't, calls sys.exit().
125def usage(msg=None):
126    print "Usage: write_large_files <opts> <varsize> [<numvars>]"
127    print ""
128    print "where <varsize> is a number followed by an optional modifier:"
129    print "                    1024M or 4G"
130    print "                the assumed modifier is '%s' if none is given." \
131          % DEF_MODIFIER
132    print "  and <numvars> is the number of variables of the above size"
133    print "                    to write.  If not supplied, 1 is assumed."
134    print "  and <opts>    is zero or more of:"
135    print "                    -c s[lice]  close & open the output file after"
136    print "                                each variable slice is written,"
137    print "                    -c v[ar]    close & open the output file after"
138    print "                                each variable is written,"
139    print "                    -t f[irst]  time the first var slice write,"
140    print "                    -t v[ar]    time the complete var write,"
141    print "                    -t w[hole]  time the complete file write,"
142
143    if msg:
144        print "\n%s" % msg
145
146    sys.exit(10)
147
148def main(argv=None):
149    global TimeFirstWrite, TimeVarWrite, TimeFileWrite
150    global CloseAfterSlice, CloseAfterVar
151    global Verbose
152
153    if argv is None:
154        argv = sys.argv
155
156    # parse command line args
157    try:
158        opts, args = getopt.getopt(argv[1:], "c:ht:v", ["help"])
159    except getopt.error, msg:
160        usage(msg)
161
162    for (opt, optarg) in opts:
163        if opt == '-c':
164            optargchar = optarg[0].lower()
165            if optargchar == 's':
166                CloseAfterSlice = True
167            elif optargchar == 'v':
168                CloseAfterVar = True
169            else:
170                usage("Unrecognized -c suboption: %s" % optarg)
171        elif opt == '-t':
172            optargchar = optarg[0].lower()
173            if optargchar == 'f':
174                TimeFirstWrite = True
175            elif optargchar == 'v':
176                TimeVarWrite = True
177            elif optargchar == 'w':
178                TimeFileWrite = True
179            else:
180                usage("Unrecognized -c suboption: %s" % optarg)
181        elif opt == '-v':
182            Verbose = True
183        else:
184            usage("Unrecognized option: %s" % opt)
185
186    if len(args) != 1 and len(args) != 2:
187        usage()
188
189    var_size = args[0][:-1]
190    modifier =args[0][-1]
191
192    if modifier in '0123456789':
193        var_size = args[0]
194        modifier = DEF_MODIFIER
195    modifier = Suffixes.get(modifier, None)
196    if modifier is None:
197        usage()
198       
199    try:
200        var_size = int(var_size) * modifier
201    except:
202        usage()
203
204    num_vars = 1
205    if len(args) == 2:
206        try:
207            num_vars = int(args[1])
208        except:
209            usage()
210
211    # call production function
212    write_file(var_size, num_vars)
213       
214if __name__ == "__main__":
215    sys.exit(main())
Note: See TracBrowser for help on using the repository browser.