Changeset 1947 in lmdz_wrf for trunk/tools/generic_tools.py

Timestamp:

Jul 19, 2018, 4:58:53 AM (6 years ago)

Author:

lfita

Message:

Adding:

• `draw_2Dshad_contdisc': 2D shadding field with a discrete one
• `draw_2Dshad_contdisc_time': 2D shadding field with a discrete one wtih a time-axis
• Fixing 'interpolation'
• Adding more variables

File:

• trunk/tools/generic_tools.py (modified) (4 diffs)

trunk/tools/generic_tools.py

@@ -5763 +5763 @@
     return angle
 
-def lonlat2D(lon,lat):
+def lonlat2D(lon, lat, error=True):
     """ Function to return lon, lat 2D matrices from any lon,lat matrix
       lon= matrix with longitude values
       lat= matrix with latitude values
+      error= whether should quit if there is an error
     """
     fname = 'lonlat2D'
 
-    if len(lon.shape) != len(lat.shape):
+    if len(lon.shape) == len(lat.shape):
+        if len(lon.shape) == 3:
+            lonvv = lon[0,:,:]
+            latvv = lat[0,:,:]
+        elif len(lon.shape) == 2:
+            lonvv = lon[:]
+            latvv = lat[:]
+        elif len(lon.shape) == 1:
+            lonlatv = np.meshgrid(lon[:],lat[:])
+            lonvv = lonlatv[0]
+            latvv = lonlatv[1]
+    else:
         print errormsg
         print '  ' + fname + ': longitude values with shape:', lon.shape,            \
           'is different than latitude values with shape:', lat.shape, '(dif. size) !!'
-        quit(-1)
-
-    if len(lon.shape) == 3:
-        lonvv = lon[0,:,:]
-        latvv = lat[0,:,:]
-    elif len(lon.shape) == 2:
-        lonvv = lon[:]
-        latvv = lat[:]
-    elif len(lon.shape) == 1:
-        lonlatv = np.meshgrid(lon[:],lat[:])
-        lonvv = lonlatv[0]
-        latvv = lonlatv[1]
+        if error: quit(-1)
+
+        # Not so weird case !!
+        print '    Trying to fix it!'
+        if len(lon.shape) == 3: 
+            lonv0 = lon[0,:,:]
+        else:
+            lonv0 = lon.copy()
+        if len(lat.shape) == 3: 
+            latv0 = lat[0,:,:]
+        else:
+            latv0 = lat.copy()
+        if len(lonv0.shape) == 1 and len(latv0.shape) == 2:
+            lonvv = np.zeros(latv0.shape, dtype=np.float)
+            if lonv0.shape[0] == latv0.shape[0]: 
+                for k in range(latv0.shape[1]):
+                    lonvv[:,k] = lonv0[:]
+            else:
+                for k in range(latv0.shape[0]):
+                    lonvv[k,:] = lonv0[:]
+            latvv = latv0.copy()
+
+        if len(lonv0.shape) == 2 and len(latv0.shape) == 1:
+            lonvv = lonv0.copy()
+            latvv = np.zeros(lonv0.shape, dtype=np.float)
+            if lonv0.shape[0] == latv0.shape[0]: 
+                for k in range(lonv0.shape[1]):
+                    latvv[:,k] = latv0[:]
+            else:
+                for k in range(lonv0.shape[0]):
+                    latvv[k,:] = latv0[:]
 
     return lonvv, latvv
 
-def interpolate_locs(locs,coords,kinterp):
+def interpolate_locs_old(locs,coords,kinterp):
     """ Function to provide interpolate locations on a given axis
     interpolate_locs(locs,axis,kinterp)
@@ -5823 +5854 @@
 
     for iloc in range(Nlocs):
+        found = False
+        a = None
+        b = None
+        c = None
         for icor in range(Ncoords-1):
             if locs[iloc] < minc and dcoords > 0.:
                 a = 0.
-                b = 1. / (coords[1] - coords[0])
+                b = coords[1] - coords[0]
                 c = coords[0]
             elif locs[iloc] > maxc and dcoords > 0.:
                 a = (Ncoords-1)*1.
-                b = 1. / (coords[Ncoords-1] - coords[Ncoords-2])
+                b = coords[Ncoords-1] - coords[Ncoords-2]
                 c = coords[Ncoords-2]
             elif locs[iloc] < minc and dcoords < 0.:
                 a = (Ncoords-1)*1.
-                b = 1. / (coords[Ncoords-1] - coords[Ncoords-2])
+                b = coords[Ncoords-1] - coords[Ncoords-2]
                 c = coords[Ncoords-2]
             elif locs[iloc] > maxc and dcoords < 0.:
                 a = 0.
-                b = 1. / (coords[1] - coords[0])
+                b = (coords[1] - coords[0])
                 c = coords[0]
             elif locs[iloc] >= coords[icor] and locs[iloc] < coords[icor+1] and dcoords > 0.:
                 a = icor*1.
-                b = 1. / (coords[icor+1] - coords[icor])
+                b = coords[icor+1] - coords[icor]
                 c = coords[icor]
-                print coords[icor], locs[iloc], coords[icor+1], ':', icor, '->', a, b
             elif locs[iloc] <= coords[icor] and locs[iloc] > coords[icor+1] and dcoords < 0.:
                 a = icor*1.
-                b = 1. / (coords[icor+1] - coords[icor])
+                b = coords[icor+1] - coords[icor]
                 c = coords[icor]
+        if a is not None: found = True
+        if not found:
+            print errormsg
+            print '  ' + fname + ": case not prepared !!"
+            print '    locs[iloc]:', locs[iloc], 'minc:', minc, 'maxc:', maxc, 'dcoords:', dcoords
+            quit(-1)
 
         if kinterp == 'lin':
-            intlocs[iloc] = a + (locs[iloc] - c)*b
+            intlocs[iloc] = a + (locs[iloc] - c)/b
         else:
             print errormsg
             print '  ' + fname + ": interpolation kind '" + kinterp + "' not ready !!!!!"
             quit(-1)
+
+    return intlocs
+
+def interpolate_locs(locs,coords,kinterp):
+    """ Function to provide interpolate locations on a given axis
+    interpolate_locs(locs,axis,kinterp)
+      locs= locations to interpolate
+      coords= axis values with the reference of coordinates
+      kinterp: kind of interpolation
+        'lin': linear
+    >>> coordinates = np.arange((10), dtype=np.float)
+    >>> values = np.array([-1.2, 2.4, 5.6, 7.8, 12.0])
+    >>> interpolate_locs(values,coordinates,'lin')
+    [ -1.2   2.4   5.6   7.8  12. ]
+    >>> coordinates[0] = 0.5
+    >>> coordinates[2] = 2.5
+    >>> interpolate_locs(values,coordinates,'lin')
+    [ -3.4          1.93333333   5.6          7.8         12.        ]
+    >>> coordinates = -10+np.arange((10), dtype=np.float)
+    >>> values = -np.array([-1.2, 2.4, 5.6, 7.8, 12.0])
+    >>> interpolate_locs(values,coordinates,'lin')
+    [ 11.2   7.6   4.4   2.2  -2. ]
+    """
+
+    fname = 'interpolate_locs'
+
+    if type(locs) == type('S') and locs == 'h':
+        print fname + '_____________________________________________________________'
+        print interpolate_locs.__doc__
+        quit()
+
+    Nlocs = locs.shape[0]
+    Ncoords = coords.shape[0]
+
+    origsing = np.abs(coords[Ncoords-1]-coords[0])/(coords[Ncoords-1]-coords[0])
+
+    intlocs = np.zeros((Nlocs), dtype=np.float)
+    minc = np.min(coords)
+    maxc = np.max(coords)
+
+    sortc = list(coords)
+    sortc.sort()
+
+    for iloc in range(Nlocs):
+        a = None
+        refi = None
+        if locs[iloc] < minc:
+            a = 0
+            refi = 0
+            sing = -1.
+        elif locs[iloc] > maxc:
+            a = Ncoords-1
+            refi = Ncoords-2
+            sing = 1.
+        else:
+            for icor in range(Ncoords-1):
+                if locs[iloc] >= sortc[icor] and locs[iloc] < sortc[icor+1]:
+                    a = icor
+                    refi = icor
+                    sing = 1.
+        if a is None:
+            print errormsg
+            print '  ' + fname + ': locs[iloc]:', locs[iloc], 'minc:', minc, 'maxc:',\
+              maxc, 'not ready!!'
+            quit(-1)
+        b = sortc[refi+1]-sortc[refi]
+        #print 'locs[iloc]:', locs[iloc], 'a:', a, 'refi:', refi, 'b:', b, 'np.abs(sortc[refi]-locs[iloc]):', np.abs(sortc[refi]-locs[iloc])
+
+        if kinterp == 'lin':
+            intlocs[iloc] = a + sing*np.abs(sortc[a]-locs[iloc])/b
+        else:
+            print errormsg
+            print '  ' + fname + ": interpolation kind '" + kinterp + "' not ready !!!!!"
+            quit(-1)
+
+    if origsing < 0.: intlocs = Ncoords - intlocs
 
     return intlocs
@@ -12301 +12417 @@
 
     # axis limits. They will be used to flip the axis in plot if necessary
-    if len(dxv.shape) == 1:
-        newxaxislim = np.array([dxv[0], dxv[len(dxv)-1]])
-    elif len(dxv.shape) == 2:
-        newxaxislim = np.array([dxv[0,0], dxv[0,dxv.shape[1]-1]])
-    else:
-        print errormsg
-        print '  ' + fname + ": wrong rank of data for axis 'x':", dxv.shape, '!!'
-        print '    must be of rank 2!!'
-        quit(-1)
-
-    if len(dyv.shape) == 1:
-        newyaxislim = np.array([dyv[0], dyv[len(dyv)-1]])
-    elif len(dyv.shape) == 2:
-        newyaxislim = np.array([dyv[0,0], dyv[dyv.shape[0]-1],0])
-    else:
-        print errormsg
-        print '  ' + fname + ": wrong rank of data for axis 'y':", dyv.shape, '!!'
-        print '    must be of rank 2!!'
-        quit(-1)
+    newxaxislim = [np.min(dxv), np.max(dxv)]
+    newyaxislim = [np.min(dyv), np.max(dyv)]
 
     for transform in transforms:
@@ -12376 +12475 @@
             flip = transform.split('@')[1]
             if flip == 'x':
-                if len(newxaxislim.shape) == 1:
-                    newxaxislim = newxaxislim[::-1]
-                else:
-                    for iy in len(newxaxislim.shape[0]):
-                        newxaxislim[iy,:] = newxaxislim[iy,::-1]
+                oldv = list(newxaxislim)
+                newxaxislim[0] = oldv[1]
+                newxaxislim[1] = oldv[0]
             elif flip == 'y':
-                if len(newyaxislim.shape) == 1:
-                    newyaxislim = newyaxislim[::-1]
-                else:
-                    for ix in len(newyaxislim.shape[1]):
-                        newyaxislim[:,ix] = newyaxislim[::-1,ix]
+                oldv = list(newyaxislim)
+                newyaxislim[0] = oldv[1]
+                newyaxislim[1] = oldv[0]
             elif flip == 'z':
                 newvals = newvals[...,::-1,:,:]