pythonlib/asc_split_into_monthly.py at master · bartnijssen/pythonlib · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
#!/usr/bin/env python
import numpy as np
import sys
import os
import glob

if len(sys.argv) < 4:
    print 'Must provide inpath, outpath, and filemask'
    sys.exit()

inpath=sys.argv[1]
outpath=sys.argv[2]
filemask=sys.argv[3]
filelist = glob.glob('{}/{}*'.format(inpath, filemask))

# all the files are the same, so we need to determine the indices only once
infile = filelist[0]
data = np.loadtxt(infile,
                  dtype={'names': ('year', 'month', 'day', 'perc'),
                         'formats': ('i4', 'i4', 'i4', 'f4')})
startyear = data['year'][0]
startmonth = data['month'][0]
endyear = data['year'][-1]
endmonth = data['month'][-1]
year = startyear
month = startmonth

# determine the number of months
nmonths = 0
while year*100+month <= endyear*100+endmonth:
    month += 1
    nmonths += 1
    if month > 12:
        year += 1
        month = 1

# initialize empty arrays
idx = np.empty((data['year'].size, nmonths), dtype='bool')*False
months = np.ones(nmonths, dtype='i4')
years = np.ones(nmonths, dtype='i4')

# get the indices, year and months
year = startyear
month = startmonth
for i in range(nmonths):
    idx[:,i] = (data['year'] == year) & (data['month'] == month)
    months[i] = month
    years[i] = year
    month += 1
    if month > 12:
        year += 1
        month = 1

# loop over all the files
for infile in filelist:
    print infile
    data = np.loadtxt(infile, dtype='S')
    filename = os.path.basename(infile)
    for i in range(nmonths):
        x = data[idx[:,i]]
        np.savetxt('{}/{:04d}_{:02d}/{}'.format(outpath, years[i], months[i],
                                                filename),
                    x, '%s')