Turn half_plume.py into both a script and a module

metrics/half_plume.py

 import numpy as np
 import netCDF4 as nc
 import matplotlib.pyplot as plt
+import argparse
 
 CHIMERE_days = ['04', '05', '06', '07', '08', '09', '10', '11']
+radius, g, p0 = 6.371e6, 9.81, 1e5 # Earth radius, gravity, reference pressure for hybrid coordinate
+
+dynamico_nc      = '/data/PLT8/pub/smailler/multiscale/dynamico.nc'
+dynamico_hybrids = '/data/PLT8/pub/smailler/multiscale/apbp.nc'
+native_nc        = '/data/WORKSPACE/dubos/Transport/NetCDF/dynamico_native.nc'
+native_area      = '/data/WORKSPACE/dubos/Transport/NetCDF/Ai.nc'
+wrf_nc           = '/home/smailler/PUY60/geog_USGS_PUY60.nc'
+chimere_nc       = '/data/PLT8/pub/smailler/CHIMOUT-PUY60-tier2/out.201106%s00_24_PUY60-tier2.nc'
 
 def compute_pressure(ap, bp, p0, ps):
     lev, lon, lat = len(ap), ps.shape[1], ps.shape[0]
@@ -62,7 +71,6 @@ def to_radian(angle):
     return angle*(np.pi/180)
 
 def read_DYNAMICO_common(hybrids, hour, area, ps, pmid, phi, q):
-    radius, g, p0 = 6.371e6, 9.81, 1e5 # Earth radius, gravity, reference pressure for hybrid coordinate
     volume = compute_volume(area/g, phi) # m^3
 
 #    print('  geopotential last hour = ', phi[:,12,27])
@@ -83,10 +91,9 @@ def read_DYNAMICO_common(hybrids, hour, area, ps, pmid, phi, q):
 
 #    print("Total SO2 mass in DYNAMICO = ", (q*air_mass).sum())
 #    print("Max SO2 mixing ratio (kg/kg) in DYNAMICO = ", q.max())
-    return radius, volume, SO2_density
+    return volume, SO2_density
 
 def read_DYNAMICO(dynamico_nc, hybrids, day=7, hour=22):
-    radius, g, p0 = 6.371e6, 9.81, 1e5 # Earth radius, gravity, reference pressure for hybrid coordinate
     hour = day*24+hour # assume output every hour, first day is day 0  
 
 #    print('Reading DYNAMICO data from %s ...'%dynamico_nc)
@@ -105,8 +112,8 @@ def read_DYNAMICO(dynamico_nc, hybrids, day=7, hour=22):
     lon, lat = to_radian(lon_deg), to_radian(lat_deg) # radians
     area = (radius**2)*np.cos(lat)*(2*np.pi/nlon)*(np.pi/nlat)
 
-    radius, volume, SO2_density = read_DYNAMICO_common(hybrids, hour, area, ps, pmid, phi, q)
-    return radius, lon_deg, lat_deg, area, volume, SO2_density
+    volume, SO2_density = read_DYNAMICO_common(hybrids, hour, area, ps, pmid, phi, q)
+    return lon_deg, lat_deg, area, volume, SO2_density
 
 def read_DYNAMICO_native(dynamico_nc, area_nc, hybrids, day=7, hour=22):
     def extend(x):
@@ -125,13 +132,13 @@ def read_DYNAMICO_native(dynamico_nc, area_nc, hybrids, day=7, hour=22):
     phi  = ds['PHI'][hour, 1:, :] # geopotential'
     q    = ds['q']  [hour, 0, :, :]
     
-    print("Reading DYNAMICO native output with %d cells"%ncell)
-    print(area.shape)
+#    print("Reading DYNAMICO native output with %d cells"%ncell)
+#    print(area.shape)
 
     lon, lat, area, ps, pmid, phi, q = map(extend, (lon, lat, area, ps, pmid, phi, q))
 
-    radius, volume, SO2_density = read_DYNAMICO_common(hybrids, hour, area, ps, pmid, phi, q)
-    return radius, lon, lat, area, volume, SO2_density
+    volume, SO2_density = read_DYNAMICO_common(hybrids, hour, area, ps, pmid, phi, q)
+    return lon, lat, area, volume, SO2_density
 
 def read_CHIMERE(wrf_nc, chimere_nc, day=6, hour=24):
 
@@ -213,14 +220,6 @@ def check_area(area, volume, radius, label):
 
 # --------------------------------------------------------------
 
-import argparse
-import matplotlib.pyplot as plt
-
-getargs = argparse.ArgumentParser(description='Analyze 3D transport experiments.')
-getargs.add_argument("--amount-emitted", action='store_true', help='Total emitted SO2 vs time')
-getargs.add_argument("--half-mass-volume", action='store_true', help='Minimal volume containing half of emitted SO2 mass')
-args = getargs.parse_args()
-
 def amount(lon, lat, area, vol, dens):
     mass = vol*dens # kg
     total = mass.sum()  # kg
@@ -234,7 +233,7 @@ def plot_amount_emitted(read_DYN, read_CHI, ndays=4, hour_step=3):
 
     DYN, CHI = [],[]
     for i in range(n):
-        radius, lon_DYN, lat_DYN, area_DYN, volume_DYN, density_DYN = read_DYN(3, hours[i]) # we shift the DYNAMICO data by 3 days
+        lon_DYN, lat_DYN, area_DYN, volume_DYN, density_DYN = read_DYN(3, hours[i]) # we shift the DYNAMICO data by 3 days
         lon_CHI, lat_CHI, area_CHI, volume_CHI, density_CHI = read_CHI(0, hours[i])
         DYN.append(amount(lon_DYN, lat_DYN, area_DYN, volume_DYN, density_DYN)) # tuple (total, lon, lat)
         CHI.append(amount(lon_CHI, lat_CHI, area_CHI, volume_CHI, density_CHI))
@@ -278,12 +277,12 @@ def plot_amount_emitted(read_DYN, read_CHI, ndays=4, hour_step=3):
     plt.close()
         
 def main():
-    dynamico_nc      = '/data/PLT8/pub/smailler/multiscale/dynamico.nc'
-    dynamico_hybrids = '/data/PLT8/pub/smailler/multiscale/apbp.nc'
-    native_nc        = '/data/WORKSPACE/dubos/Transport/NetCDF/dynamico_native.nc'
-    native_area      = '/data/WORKSPACE/dubos/Transport/NetCDF/Ai.nc'
-    wrf_nc           = '/home/smailler/PUY60/geog_USGS_PUY60.nc'
-    chimere_nc       = '/data/PLT8/pub/smailler/CHIMOUT-PUY60-tier2/out.201106%s00_24_PUY60-tier2.nc'
+
+    getargs = argparse.ArgumentParser(description='Analyze 3D transport experiments.')
+    getargs.add_argument("--native", action='store_true', help='Read DYNAMICO native data instead of lon-lat')
+    getargs.add_argument("--amount-emitted", action='store_true', help='Total emitted SO2 vs time')
+    getargs.add_argument("--half-mass-volume", action='store_true', help='Minimal volume containing half of emitted SO2 mass')
+    args = getargs.parse_args()
 
     if args.amount_emitted:
         def read_DYN(day,hour):
@@ -292,10 +291,16 @@ def main():
             return read_DYNAMICO_native(native_nc, native_area, dynamico_hybrids, day, hour)
         def read_CHI(day,hour):
             return read_CHIMERE(wrf_nc, chimere_nc, day, hour)
-        plot_amount_emitted(read_DYN, read_CHI)
+        if args.native:
+            plot_amount_emitted(read_NAT, read_CHI)
+        else:
+            plot_amount_emitted(read_DYN, read_CHI)
     else:
-#        radius, _, _, area_DYNAMICO, volume_DYNAMICO, density_DYNAMICO = read_DYNAMICO(dynamico_nc, dynamico_hybrids)
-        radius, _, _, area_DYNAMICO, volume_DYNAMICO, density_DYNAMICO = read_DYNAMICO_native(native_nc, native_area, dynamico_hybrids)
+        if args.native:
+            _, _, area_DYNAMICO, volume_DYNAMICO, density_DYNAMICO = read_DYNAMICO_native(native_nc, native_area, dynamico_hybrids)
+        else:
+            _, _, area_DYNAMICO, volume_DYNAMICO, density_DYNAMICO = read_DYNAMICO(dynamico_nc, dynamico_hybrids)
+
         check_area(area_DYNAMICO, volume_DYNAMICO, radius, 'DYNAMICO')
         _, _, area_CHIMERE, volume_CHIMERE, density_CHIMERE = read_CHIMERE(wrf_nc, chimere_nc)    
         check_area(area_CHIMERE, volume_CHIMERE, radius, 'CHIMERE')
@@ -307,4 +312,5 @@ def main():
 
     #plots(metrics_CHIMERE, metrics_DYNAMICO)
 
-main()
+if __name__ == "__main__":
+    main()