Changeset 3833 for trunk/LMDZ.PLUTO/util/script_figures/FV3_utils.py

Timestamp:

Jul 7, 2025, 1:39:13 PM (23 hours ago)

Author:

afalco

Message:

Pluto: updated plots scripts.
Fixed some issues with reading XIOS, etc.
Included display_netcdf.py tool from Mars PCM.
AF

File:

• trunk/LMDZ.PLUTO/util/script_figures/FV3_utils.py (modified) (5 diffs)

trunk/LMDZ.PLUTO/util/script_figures/FV3_utils.py

@@ -3 +3 @@
 pi=math.pi
 import matplotlib 
-matplotlib.use('TKAgg')
-
-name="../Xhistins2015_short"
+matplotlib.use('Agg')
+from input import *
+from display_netcdf import *
 
 def fmtsci(x, pos):
@@ -1201 +1201 @@
 
 def printvar(infile):
-    # Get all variable names for a netCDF file ---
+    '''Get all variable names for a netCDF file '''
     print("Variables:")
     variableNames = list(infile.variables.keys());
@@ -1208 +1208 @@
 
 def getind(myloc,field):
-    # get a specific index in the lat, lon, time or pfull 1D field
+    '''get a specific index in the lat, lon, time or pfull 1D field'''
     res=[]
     for loc in np.atleast_1d(myloc):
@@ -1220 +1220 @@
 # getinds=np.vectorize(getind, excluded="field")
 
-def getvar(nc1,var,times=None,tim=None,longitudes=None,lon=None,t_mean=False,l_mean=False):
+def getvar(nc1,var,times=None,longitudes=None,latitudes=None,altitudes=None,t_mean=False,l_mean=False):
+    """Get variables from netcdf file and slice it up!
+
+    Args:
+        nc1 : netcdf identifier
+        var (str): variable name
+        times (list, optional): times to select. Defaults to None.
+        longitudes (list, optional): longitudes to select. Defaults to None.
+        latitudes (list, optional): latitudes to select. Defaults to None.
+        t_mean (bool, optional): Apply mean on time dimension. Defaults to False.
+        l_mean (bool, optional): Apply mean on longitude dimension. Defaults to False.
+
+    Returns:
+        array: variable data array
+    """
+    
     if var == "latitude" and var not in nc1.variables: var="lat"
     if var == "longitude" and var not in nc1.variables: var="lon"
@@ -1227 +1242 @@
     if var == "phisinit" and var not in nc1.variables: var="phisfi"
 
-    myvar = nc1.variables[var][:]
-    try:
-        t=getind(times,tim)
-        if len(t)==2:
-            myvar = myvar[t[0]:t[1]+1]
-        else:
-            myvar = myvar[t]
-    except:
-        # print("no time selected")
-        pass
-    try:
-        l=getind(longitudes,lon)
-        if len(l)==2:
-            myvar = myvar[...,l[0]:l[1]+1]
-        else:
-            myvar = myvar[...,l]
-    except:
-        # print("no longitude selected")
-        pass
-    if t_mean: myvar=np.mean(myvar,axis=0) # temporal mean
-    if l_mean: myvar=np.mean(myvar,axis=-1) # longitudinal mean
+    var_data = nc1.variables[var]
+    dims = var_data.dimensions
+    myvar = var_data[:]
+
+    if times or longitudes or latitudes:
+        slicer = [slice(None)] * len(dims)
+
+        time_it = (times,TIME_DIMS)
+        alt_it = (altitudes,ALT_DIMS)
+        lon_it = (longitudes,LON_DIMS)
+        lat_it = (latitudes,LAT_DIMS)
+
+        # iterate over dimensions to pick indices to select
+        for values, dim in (time_it, lon_it, lat_it, alt_it):
+            try:
+                dim_idx = find_dim_index(dims, dim)
+                tmp_var = find_coord_var(nc1, dim)
+                file_values = nc1.variables[tmp_var][:]
+                i = getind(values, file_values)
+                print(dim, i)
+                if len(i)==2: # range of indices
+                    slicer[dim_idx] = slice(min(i),max(i)+1)
+                else: # specific index
+                    slicer[dim_idx] = i
+            except:
+                continue
+        
+        # slice the variable according to the indices
+        myvar = myvar[tuple(slicer)]
+        
+    if l_mean: # longitudinal mean
+        myvar=np.mean(myvar,axis=find_dim_index(dims, LON_DIMS))
+    if t_mean: # temporal mean
+        myvar=np.mean(myvar,axis=find_dim_index(dims, TIME_DIMS)) 
     print(np.shape(myvar))
     try: