Update interpolation funcs

Author: lukelbd

isobars2isentropes/metpy.py

@@ -1,78 +1,71 @@
 #!/usr/bin/env python3
-#------------------------------------------------------------------------------#
 # Interpolate to isentropic levels
-# NOTE: File cannot have any print statements! Need to capture output of
-# time command
-#------------------------------------------------------------------------------#
-# Imports
-# import cartopy.crs as ccrs
-# import cartopy.feature as cfeature
-# import matplotlib.pyplot as plt
+# NOTE: File cannot have any print statements! Need to capture output of time command.
 import os
 import sys
+
 import numpy as np
-import xarray as xr
+
 import metpy.calc as mcalc
+import xarray as xr
 from metpy.units import units
 
 # File, with optional decoding
-if len(sys.argv)!=3:
+if len(sys.argv) != 3:
     raise ValueError('Require two input args.')
 filename = sys.argv[1]
 dir = filename.split('/')[0]
-nt = int(sys.argv[2]) # number of time chunks; 0 to disable chunking
-if nt: # 1 chunk per level
-    N = 100000 # max chunk size for lon/lat dimensions; want to chunk each 2D slice separately
-    # N = 1000 # this made no difference
-    chunks = {'time':nt, 'plev':N, 'lat':N, 'lon':N} # one chunk per horizontal slice, works pretty well
+nt = int(sys.argv[2])  # number of time chunks; 0 to disable chunking
+if nt:  # 1 chunk per level
+    N = 100000  # max chunk size for space (want to chunk each timestep separately)
+    chunks = {'time': nt, 'plev': N, 'lat': N, 'lon': N}
     data = xr.open_dataset(filename, chunks=chunks, decode_times=False)
 else:
     data = xr.open_dataset(filename, decode_times=False)
 
 # Interpoland
-eps = 1
-# thlev = np.array([*range(260, 360, 10), *range(360, 400, 20), *range(400, 800, 50), *range(800, 1600, 200)])
 # thlev = np.array([*range(260, 360, 10), *range(360, 400, 20), *range(400, 800, 50)])
 # thlev = np.array([*range(300, 360, 10), *range(360, 400, 20)])
-# thlev = np.arange(240,400,20) # must be ascending!
-thlev = np.array([265, 275, 285, 300, 315, 330, 350, 370, 395, 430, 475, 530, 600, 700, 850])
+# thlev = np.arange(240,400,20)  # must be ascending!
+thlev = [265, 275, 285, 300, 315, 330, 350, 370, 395, 430, 475, 530, 600, 700, 850]
+thlev = np.array(thlev)
+
 # Coordinates
 time = data['time']
 plev = data['plev']
 lat = data['lat']
 lon = data['lon']
+
 # Variables
-# u = data['u'][0]
 u = data['u']
 v = data['v']
 t = data['t']
-t = t*units.kelvin
-thlev = thlev.astype('float32')*units.kelvin # assign units to input arg
-plev = plev.values.astype('float32')*units.hectopascals
+t = t * units.kelvin
+thlev = thlev.astype('float32') * units.kelvin
+plev = plev.values.astype('float32') * units.hectopascals
 
 # Interpolation
 # The returned 'p' are pressure along each isentrope
-# NOTE: If temps out of bounds get error:
-# raise ValueError('Input theta level out of data bounds')
-# This should just be a warning, or optionally disabled!
-p, t_th, u_th, v_th = mcalc.isentropic_interpolation(thlev, plev, t.values*units.kelvin, u.values, v.values, axis=1, tmpk_out=True) # output isentropes in Kelvin
+p, t_th, u_th, v_th = mcalc.isentropic_interpolation(
+    thlev, plev, t.values * units.kelvin, u.values, v.values, axis=1, tmpk_out=True
+)
 dims = ('time', 'thlev', 'lat', 'lon')
-p = xr.Variable(dims, p.astype('float32'), {'long_name':'pressure', 'units':'hPa'})
+p = xr.Variable(dims, p.astype('float32'), {'long_name': 'pressure', 'units': 'hPa'})
 t_th = xr.Variable(dims, t_th.astype('float32'), t.attrs)
 u_th = xr.Variable(dims, u_th.astype('float32'), u.attrs)
 v_th = xr.Variable(dims, v_th.astype('float32'), v.attrs)
+
 # Make dataset, add variables and coordinates
 # NOTE: Put coordinates in first, so they come first in ncdump
-thlev = xr.Variable('thlev', thlev, {'long_name':'potential_temperature', 'units':'K'})
-out = xr.Dataset({}, coords={'time':time, 'thlev':thlev, 'lat':lat, 'lon':lon})
-for name,data in zip(('p','t','u','v'),(p,t_th,u_th,v_th)):
+thlev = xr.Variable(
+    'thlev', thlev, {'long_name': 'potential_temperature', 'units': 'K'}
+)
+out = xr.Dataset({}, coords={'time': time, 'thlev': thlev, 'lat': lat, 'lon': lon})
+for name, data in zip(('p', 't', 'u', 'v'), (p, t_th, u_th, v_th)):
     out[name] = data
-# out = xr.Dataset({'p':p, 't':t_th, 'u':u_th, 'v':v_th},
-#                  {'thlev':thlev, 'lon':lon, 'lat':lat})
 
 # Save file
 outname = f'{dir}/isentropes_py{nt}.nc'
 if os.path.exists(outname):
     os.remove(outname)
-out.to_netcdf(outname, mode='w') # specify whether we did chunking
-
+out.to_netcdf(outname, mode='w')  # specify whether we did chunking

isobars2isentropes/ncl.ncl

@@ -1,19 +1,18 @@
-;------------------------------------------------------------------------------;
 ; Interpolate to isentropic levels
-;------------------------------------------------------------------------------;
-qq = integertochar(34) ; a double quote; only way to put inside string! yuck.
+qq = integertochar(34)  ; a double quote
 demo = "ncl 'filename=" + qq + "foobar" + qq + "' or " + qq + "filename=\" + qq + "foobar\" + qq + qq + "."
 if (.not. isvar("filename")) then
     print("fatal:File name must be passed as variable 'filename' as follows: " + demo)
-    exit ; almost impossible to put double-quote in string
+    exit
 end if
 if (.not. isvar("outname")) then
     print("fatal:Output name must be passed as variable 'outname' as follows: " + demo)
-    exit ; almost impossible to put double-quote in string
+    exit
 end if
 
 ; Read file
 f = addfile(filename, "r") ; just read data from here
+
 ; Coordinates
 time = f->time
 plev = f->plev
@@ -24,18 +23,15 @@ lat@axis = "Y"
 lon@axis = "X"
 
 ; Get potential temp
-; p0 = 1013.25
 t = f->t
-p0 = 1000.0 ; more common to make this the reference pressure
+p0 = 1000.0
 kappa = 0.286
-theta = t*conform(t, (p0/plev)^kappa, (/1/)) ; calculate potential temperature
+theta = t * conform(t, (p0 / plev)^kappa, (/1/))
 delete(t)
 
-; The theta coordinates
-; Should go in *same order* as pressure coordinates (pressure increasing,
-; so theta decreasing)
-; thlev = ispan(400,240,20)*1.0 ; specify desired isentropic levels
-thlev = (/265, 275, 285, 300, 315, 330, 350, 370, 395, 430, 475, 530, 600, 700, 850/)*1.0
+; The theta coordinates. Should go in *same order* as pressure coordinates.
+; thlev = ispan(400, 240, 20) *1.0  ; specify desired isentropic levels
+thlev = 1.0 * (/265, 275, 285, 300, 315, 330, 350, 370, 395, 430, 475, 530, 600, 700, 850/)
 thlev = thlev(::-1)
 thlev!0 = "thlev"
 thlev&thlev = thlev ; this is apparently necessary
@@ -46,46 +42,43 @@ thlev@axis = "Z"
 ; Output file
 ; Add coordinates so they appear first in order
 ; TODO: Add coordinates and whatnot
-; dir = str_split(filename, "/")
-; dir := dir(0)
-; system("rm " + dir + "/isentropes_ncl.nc 2>/dev/null") ; remove file
-system("rm " + outname + " 2>/dev/null") ; remove file
-o = addfile(outname, "c") ; create new file; don't want to read old values or anything
-filedimdef(o, "time", -1, True) ; unlimited dimension
-; print(time)
-; print(lat)
-; print(thlev)
+system("rm " + outname + " 2>/dev/null")  ; remove file
+o = addfile(outname, "c")  ; create new file, don't want to read old values or anything
+filedimdef(o, "time", -1, True)  ; unlimited dimension
 o->time = time
 o->thlev = thlev
 o->lat = lat
 o->lon = lon
 
 ; Interpolate
-; Don't need to save temperature of course, but would like pressure
 names = (/"p", "u", "v"/)
 do i=0,dimsizes(names)-1
   ; Special consideration for p levs
   name = names(i)
   if (name .eq. "p") then
-    var_orig = conform(f->t, plev, (/1/)) ; only dimension 1 matches
+    var_orig = conform(f->t, plev, (/1/))
     var_orig@long_name = "pressure"
     var_orig@units = "hPa"
   else
     var_orig = f->$name$
   end if
+
   ; Interpolate to isentropes
   var = int2p_n(theta, var_orig, thlev, 0, 1)
+
   ; Dimensions
-  var!0 = "time" ; name dimensions
+  var!0 = "time"
   var!1 = "thlev"
   var!2 = "lat"
   var!3 = "lon"
+
   ; Coordinates
   var&time = time
   var&thlev = thlev
   var&lat = lat
   var&lon = lon
   copy_VarAtts(var_orig, var)
+
   ; Add to file
   o->$name$ = var
   delete(var)