add functions related to coesa76

pyatmos/utils/utils.py

@@ -1,21 +1,122 @@
+import numpy as np
+import warnings
+
+from . import Const
+
+def vectorize(x):
+    '''
+    Vectorize a number(int, float) or a list to a numpy array.
+    '''
+    try:
+        n = len(x)
+        x = np.array(x)
+    except:
+        x = np.array([x])
+    return x    
+            
 def wraplon(lon):
     '''
-    Wrap a longitude in range of [0,360] to [-180,180]
+    Wrap a longitude in range of [0,360] to [-180,180].
+
+    Usage:
+    lon_wrap = wraplon(lon)
+
+    Inputs:
+    lon -> [float] longitude in range of [0,360]
+
+    Outputs:
+    lon_wrap -> [float] wrapped longitude in range of [-180,180]
     '''
     if lon > 180:
-        lonwrap = lon - 360
+        lon_wrap = lon - 360
     else:
-        lonwrap = lon
+        lon_wrap = lon
-    return lonwrap 
+    return lon_wrap 
 
-def hms2s(h,m,s):
+def wraplons(lons):
     '''
-    Convert hour/minute/second to seconds
+    Wrap a set of longitudes in range of [0,360] to [-180,180].
-    '''
-    return h*3.6E3 + m*60 + s
 
-def hms2h(h,m,s):
+    Usage:
+    lons_wrap = wraplons(lons)
+
+    Inputs:
+    lons -> [float list/array] longitudes in range of [0,360]
+
+    Outputs:
+    lons_wrap -> [float array] wrapped longitudes in range of [-180,180]
     '''
-    Convert hour/minute/second to hours
+    lons = vectorize(lons)
+    lons_wrap = lons.copy()
+    flags = lons > 180
+    lons_wrap[flags] = lons[flags] - 360
+
+    return lons_wrap     
+
+def hms_conver(h,m,s):
     '''
-    return h + m/60 + s/3.6E3
+    Convert the form of hour/minute/second to hours and seconds.
+
+    Uasge:
+    hours,seconds = hms_conversion(h,m,s)
+    '''
+    hours = h + m/60 + s/3.6e3
+    seconds = h*3.6e3 + m*60 + s
+    return hours,seconds
+
+def alt_conver(alts,alt_type='geometric'):
+    '''
+    Fulfill conversions between geometric altitudes and geopotential altitudes.  
+
+    Usage:
+    zs,hs = alt_conver(alts,'geometric') 
+    or
+    zs,hs = alt_conver(alts,'geopotential') 
+
+    Inputs:
+    alts -> [float list/array] geometric altitudes or geopotential altitudes, [km]
+    
+    Parameters:
+    alt_type -> [string] 'geometric' or 'geopotential'
+
+    Outputs:
+    zs -> [float array] geometric altitudes, [km]
+    hs -> [float array] geopotential altitudes, [km]
+    '''
+
+    alts = vectorize(alts)
+
+    r0 = Const.r0
+
+    if alt_type == 'geometric': 
+        zs = alts
+        # from geometric altitude to geopotential altitude
+        hs = zs*r0/(r0+zs)  
+
+    elif alt_type == 'geopotential':
+        hs = alts
+        # from geopotential altitude to geometric altitude
+        zs = hs*r0/(r0-hs)  
+    return zs,hs
+
+def check_altitude(zs,z_range,mode):
+    '''
+    Checks if altitudes are inside a valid range.
+
+    Inputs:
+    zs -> [float list/array] geometric altitude to be checked
+    lower_z -> [float] lower limit of geometric altitudes
+    upper_z -> [float] upper limit of geometric altitudes
+    '''
+    zs = np.array(zs)
+
+    lower_z,upper_z = z_range
+    # Assert in range
+
+    if (zs < lower_z).any() or (zs > upper_z).any():
+        msg_warning = "Geometric altitudes are outside the range of [{}, {}] km. Output values will be extrapolated for those heights.".format(lower_z,upper_z)
+        msg_error = "Geometric altitudes are outside the range of [{}, {}] km.".format(lower_z,upper_z)
+        if mode == 'warning':
+            warnings.warn(msg_warning)      
+        elif mode == 'error':
+            raise Exception(msg_error)