python_tools/restart_particles.py at master · tiagoams/python_tools · 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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
restart_particles - create a spawning file for particle tracking based on the
                    end position on a particle tracking run

    restart_particles -i Input_file -o Output_file  Date Time

    Input_file  - name of the netcdf file
    Output_file - name of the output csv file
    Date_string - String containing the Date and time of spawning

    example:
    restart_particles -i tn.nc -o Output_file spawnfile 1996-01-01 01:00:00

Context: Particle tracking over multiple years or to repeat years
         Written for Matthew Bone's PhD on Winter Nitrate NERC-SSB

Created on Tue May 10 16:44:56 2016

@author: TAMS00
"""

from netCDF4 import Dataset
import numpy as np
import argparse
import sys

#******************************************************************
def read_nc(VarName,FName):


    #print 'VarName, FName',VarName, FName
    try:
        grp = Dataset(FName)
    except:
        print("ERROR reading file "+FName)
        sys.exit()

    Var = grp.variables[VarName][:]

    #print "read var "+VarName+" from "+FName

    grp.close()

    return Var


#******************************************************************
def write_csv(FName,rows):
    import csv

    with open(FName,'wb') as csvfile:
        csvwriter = csv.writer(csvfile, delimiter=',')
        for row in rows:
            csvwriter.writerow(row)


#******************************************************************
# main() to take an optional 'argv' argument, which allows us to call it
# from the interactive Python prompt:
def main():
    # parse command line options

    parser = argparse.ArgumentParser()
    parser.add_argument('date', help='date for spawing')
    parser.add_argument('time', help='time for spawing')
    parser.add_argument('-i', '--infile', required=True, help='path of data input file')
    parser.add_argument('-o', '--outfile', required=True, help='path of data output file')
    args = parser.parse_args()

    strDateTime = args.date+" "+args.time
    InFName = args.infile
    OutFName = args.outfile

    link=InFName

    # Extract ipos
    try:
        ipos = read_nc('ipos',link)
    except:
        print("ERROR Can\'t find variable ipos")
        sys.exit()

    # Extract jpos
    try:
        jpos = read_nc('jpos',link)
    except:
        print("ERROR Cannot find variable jpos")
        sys.exit()

    # Extract vertical position. Negative is down.
    try:
        kpos = read_nc('kpos',link)
    except:
        print("ERROR Can\'t find variable kpos")
        sys.exit()

    # Extract health of superparticle
    try:
        wfact = read_nc('wfact',link)
    except:
        print("ERROR Can\'t find variable wfact")
        sys.exit()


    # Use final position
    [Ntime,Npart]=ipos.shape
    ipos = ipos[Ntime-1,:]
    jpos = jpos[Ntime-1,:]
    kpos = kpos[Ntime-1,:]
    wfact = wfact[Ntime-1,:]


    rows=[]
    # Header has number of particles
    rows.append([Npart])
    # Create a row for every particle
    for ipart in range(Npart):
        lon=ipos[ipart]
        lat=jpos[ipart]
        depth=kpos[ipart]
        health=wfact[ipart]
        row=[ipart,lon,lat,depth,strDateTime,health]
        rows.append(row)

    write_csv(OutFName,rows)


if __name__ == "__main__":
    main()