source: branches/numpy_misc/tools/write_large_files/rw_big_file.py @ 6817

Last change on this file since 6817 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: 10.1 KB
Line 
1#!/usr/bin/env python
2################################################################################
3# A test program to write a large NetCDF file and see if we can read it.
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*8     # 8 bytes / float
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
37TimeFirstRead = False
38TimeVarRead = False
39TimeFileRead = False
40
41CloseAfterSlice = False
42CloseAfterVar = False
43
44Verbose = False
45
46# mask for variable names
47VarnameMask = 'var_%04d'
48
49
50######
51# mainline code
52######
53
54##
55# @brief Read a NetCDF file - see if it as we wrote it.
56# @param filename The name of the file to read.
57# @param variable_size Size of variable, in bytes.
58# @param num_variables Number of variables.
59def read_file(filename, variable_size, num_variables):
60    # create a data array slice
61    slice_array_mask = numpy.ones((XDIM,), 'd')
62
63    # if timing file read, get start time
64    if TimeFileRead:
65        start_file_read = time.time()
66
67    fid = nc.NetCDFFile(filename, 'r')
68
69    # num file bytes read counter
70    file_bytes_read = 0
71
72    result = True
73
74    for i in xrange(num_variables):
75        varname = VarnameMask % i
76        var = fid.variables[varname]
77
78        # num variable bytes written counter
79        var_bytes_read = 0
80
81        # create expected slice array
82        slice_array = slice_array_mask * (i + 1)
83        for x in xrange(XDIM):
84            slice_array[x] += x*1.0e-7
85
86        if TimeVarRead:
87            start_var_read = time.time()
88
89        index = 0
90        while var_bytes_read < variable_size:
91            var_bytes_read += SLICE_SIZE
92            file_bytes_read += SLICE_SIZE
93
94            if TimeFirstRead and index == 0:
95                start_first_var_read = time.time()
96
97            if Verbose:
98                print ('File %s, variable %s, reading slice %d: '
99                       'var=%.1fMiB, file=%.1fMiB' %
100                       (filename, varname, index, float(var_bytes_read)/MBYTE,
101                        float(file_bytes_read)/MBYTE))
102
103            var_array = fid.variables[varname][index,:]
104            if num.any(var_array != slice_array):
105                result = False
106                if Verbose:
107                    print 'Read variable %s, slice %d: got unexpected value' % \
108                          (varname, index)
109
110            if TimeFirstRead and index == 0:
111                stop_first_var_read = time.time()
112                print ("Time to do first read for variable %s: %.2f sec" %
113                       (varname, (stop_first_var_read - start_first_var_read)))
114
115            index += 1
116
117            if CloseAfterSlice:
118                fid.close()
119                fid = nc.NetCDFFile(filename, 'r')
120        if CloseAfterVar:
121            fid.close()
122            fid = nc.NetCDFFile(filename, 'r')
123
124        if TimeVarRead:
125            stop_var_read = time.time()
126            print ('Time to read variable %s: %.2f sec' %
127                   (varname, (stop_var_read - start_var_read)))
128
129    fid.close()
130
131    if TimeFileRead:
132        stop_file_read = time.time()
133        print ('Time to read file: %.2f sec' %
134               (stop_file_read - start_file_read))
135
136    return result
137
138
139##
140# @brief Write a NetCDF file with set number of variables of a defined size.
141# @param variable_size Size of variable, in bytes.
142# @param num_variables Number of required variables.
143# @return The filename of the created file.
144def write_file(variable_size, num_variables=1):
145    # set file and variable name masks
146    filename = 'test.nc'
147    VarnameMask = 'var_%04d'
148
149    # create a data array slice
150    slice_array_mask = numpy.ones((XDIM,), 'd')
151
152    # if timing file write, remember start time
153    if TimeFileWrite:
154        start_file_write = time.time()
155
156    fid = nc.NetCDFFile(filename, 'w')
157    fid.createDimension('y', None)
158    fid.createDimension('x', XDIM)
159
160    # num file bytes written counter
161    file_bytes_written = 0
162
163    for i in xrange(num_variables):
164        varname = VarnameMask % i
165        fid.createVariable(varname, 'd', ('y', 'x'))
166
167        # num variable bytes written counter
168        var_bytes_written = 0
169
170        # create unique slice array
171        slice_array = slice_array_mask * (i + 1)
172        for x in xrange(XDIM):
173            slice_array[x] = slice_array[x] + x*1.0e-7
174
175        if TimeVarWrite:
176            start_var_write = time.time()
177
178        index = 0
179        while var_bytes_written < variable_size:
180            var_bytes_written += SLICE_SIZE
181            file_bytes_written += SLICE_SIZE
182
183            if TimeFirstWrite and index == 0:
184                start_first_var_write = time.time()
185
186            if Verbose:
187                print ('File %s, variable %s, writing slice %d: '
188                       'var=%.1fMiB, file=%.1fMiB' %
189                       (filename, varname, index, float(var_bytes_written)/MBYTE,
190                        float(file_bytes_written)/MBYTE))
191
192            fid.variables[varname][index,:] = slice_array
193
194            if TimeFirstWrite and index == 0:
195                stop_first_var_write = time.time()
196                print ("Time to do first write for variable %s: %.2f sec" %
197                       (varname, (stop_first_var_write-start_first_var_write)))
198
199            index += 1
200
201            if CloseAfterSlice:
202                fid.close()
203                fid = nc.NetCDFFile(filename, 'a')
204        if CloseAfterVar:
205            fid.close()
206            fid = nc.NetCDFFile(filename, 'a')
207
208        if TimeVarWrite:
209            stop_var_write = time.time()
210            print ('Time to write variable %s: %.2f sec' %
211                   (varname, (stop_var_write - start_var_write)))
212
213    fid.close()
214
215    if TimeFileWrite:
216        stop_file_write = time.time()
217        print ('Time to write file: %.2f sec' %
218               (stop_file_write - start_file_write))
219
220    return filename
221
222
223##
224# @brief Provide help for the befuddled user.
225# @return Doesn't, calls sys.exit().
226def usage(msg=None):
227    print "Usage: write_large_files <opts> <varsize> [<numvars>]"
228    print ""
229    print "where <varsize> is a number followed by an optional modifier:"
230    print "                    1024M or 4G"
231    print "                the assumed modifier if none is given is '%s'." \
232          % DEF_MODIFIER
233    print "  and <numvars> is the number of variables of the above size"
234    print "                    to write.  If not supplied, 1 is assumed."
235    print "  and <opts>    is zero or more of:"
236    print "                    -c s[lice]  close & open the output file after"
237    print "                                each variable slice is read/written,"
238    print "                    -c v[ar]    close & open the output file after"
239    print "                                each variable is read/written,"
240    print "                    -t rf[irst] time the first var slice read,"
241    print "                    -t rv[ar]   time the complete var read,"
242    print "                    -t rw[hole] time the complete file read,"
243    print "                    -t wf[irst] time the first var slice write,"
244    print "                    -t wv[ar]   time the complete var write,"
245    print "                    -t ww[hole] time the complete file write,"
246
247    if msg:
248        derr, "\n%s" % msg
249
250    sys.exit(10)
251
252def main(argv=None):
253    global TimeFirstWrite, TimeVarWrite, TimeFileWrite
254    global TimeFirstRead, TimeVarRead, TimeFileRead
255    global CloseAfterSlice, CloseAfterVar
256    global Verbose
257
258    if argv is None:
259        argv = sys.argv
260
261    # parse command line args
262    try:
263        opts, args = getopt.getopt(argv[1:], "c:t:v", ["help"])
264    except getopt.error, msg:
265        usage(msg)
266
267    for (opt, optarg) in opts:
268        if opt == '-c':
269            optargchar = optarg[0].lower()
270            if optargchar == 's':
271                CloseAfterSlice = True
272            elif optargchar == 'v':
273                CloseAfterVar = True
274            else:
275                usage("Unrecognized -c suboption: %s" % optarg)
276        elif opt == '-t':
277            optargchar = optarg[0].lower()
278            if optargchar == 'r':
279                optargchar = optarg[1].lower()
280                if optargchar == 'f':
281                    TimeFirstRead = True
282                elif optargchar == 'v':
283                    TimeVarRead = True
284                elif optargchar == 'w':
285                    TimeFileRead = True
286                else:
287                    usage("Unrecognized -t suboption: %s" % optarg)
288            elif optargchar == 'w':
289                optargchar = optarg[1].lower()
290                if optargchar == 'f':
291                    TimeFirstWrite = True
292                elif optargchar == 'v':
293                    TimeVarWrite = True
294                elif optargchar == 'w':
295                    TimeFileWrite = True
296                else:
297                    usage("Unrecognized -t suboption: %s" % optarg)
298            else:
299                usage("Unrecognized -t suboption: %s" % optarg)
300        elif opt == '-v':
301            Verbose = True
302        else:
303            usage("Unrecognized option: %s" % opt)
304
305    if len(args) != 1 and len(args) != 2:
306        usage()
307
308    var_size = args[0][:-1]
309    modifier =args[0][-1]
310
311    if modifier in '0123456789':
312        var_size = args[0]
313        modifier = DEF_MODIFIER
314    modifier = Suffixes.get(modifier, None)
315    if modifier is None:
316        usage()
317       
318    try:
319        var_size = int(var_size) * modifier
320    except:
321        usage()
322
323    num_vars = 1
324    if len(args) == 2:
325        try:
326            num_vars = int(args[1])
327        except:
328            usage()
329
330    # write the required file
331    filename = write_file(var_size, num_vars)
332       
333    # read the file to see if is as expected
334    if not read_file(filename, var_size, num_vars):
335        print "Didn't read data that we read!?"
336        return 10
337
338    print 'Read/write of NectCDF file was correct'
339    return 0
340
341
342if __name__ == "__main__":
343    sys.exit(main())
Note: See TracBrowser for help on using the repository browser.