Merge pull request #5 from stella-bourdin/read_limit_files

Add function to write in limit.nc files

dynamicopy/cartoplot.py

@@ -59,6 +59,12 @@ def lon_lat_plot_map(lon, lat, var, lon_axis=-1, lat_axis=-2, fig_ax=None, title
     # Obtain 2D variable to plot
     var2D = _var2d(var, lon_axis, lat_axis)
 
+    # Truncate if latitude coordinates array go too far
+    mask = ~((lat > 88) | (lat < -88))
+    var2D = apply_mask_axis(var2D, mask, axis=lat_axis)
+    lat = lat[mask]
+    if any(mask == False) : print("Warning, some values too close to the pole(s) were not displayed.")
+
     # Plotting
     if fig_ax == None:
         fig = plt.figure()

dynamicopy/ncload.py

@@ -73,6 +73,90 @@ def get_lon_lat(file_path, lon_name='lon', lat_name='lat'):
 
     return np.array(lon), np.array(lat)
 
+def var_load_from_limit(varname, limit_file='limit.nc', lon_name='longitude', lat_name='latitude'):
+    """ Loads a field from a limit.nc file (LMDZ standard structure)
+
+    Parameters
+    ----------
+    varname : str
+        Name of the variable to load in the file
+    limit_file : str
+        The path (relative or absolute) to the limit.nc file;
+    lon_name : str, optional
+        name of the longitude coordinate in the file (default 'longitude');
+    lat_name : str, optional
+        name of the latitude coordinate in the file (default 'latitute');
+
+    Returns
+    -------
+    list of np.ndarray (lon, lat, var)
+        The field reshaped in lon/lat coordinates and the corresponding coordinate arrays;
+    """
+
+    lon = var_load(lon_name, limit_file)
+    lat = var_load(lat_name, limit_file)
+    var = var_load(varname, limit_file)
+
+    lat_dim_no_pole = len(np.unique(lat)) - 2
+    lon_dim = len(np.unique(lon))
+
+    if (len(np.shape(var)) == 1) :
+        var_north_pole = var[0]
+        var_south_pole = var[-1]
+        var_no_pole = var[1:-1]
+
+        var_no_pole_reshape = np.reshape(var_no_pole, [lat_dim_no_pole, lon_dim])
+        var_north_pole_reshape = [[var_north_pole] * lon_dim]
+        var_south_pole_reshape = [[var_south_pole] * lon_dim]
+        var_reshape = np.concatenate([var_north_pole_reshape, var_no_pole_reshape, var_south_pole_reshape], 0)
+
+    elif (len(np.shape(var)) == 2) :
+        var_north_pole = var[:,0]
+        var_south_pole = var[:,-1]
+        var_no_pole = var[:,1:-1]
+
+        var_no_pole_reshape = [np.reshape(var_no_pole[i], [lat_dim_no_pole, lon_dim]) for i in range(len(var))]
+        var_north_pole_reshape = [[[var_north_pole[i]] * lon_dim] for i in range(len(var))]
+        var_south_pole_reshape = [[[var_south_pole[i]] * lon_dim] for i in range(len(var))]
+        var_reshape = np.concatenate([var_north_pole_reshape, var_no_pole_reshape, var_south_pole_reshape], 1)
+
+    else :
+        print("Problem with the variable dimensions")
+
+    lat_reshape = np.flip(np.unique(lat))
+    lon_reshape = np.unique(lon)
+
+    return lon_reshape, lat_reshape, var_reshape
+
+def change_limit(newfield, fieldname, limit_file='limit.nc'):
+    """ Function to change a field in a limit.nc file (LMDZ standard structure)
+
+    Parameters
+    ----------
+    newfield : np.ndarray
+        2D or 3D field corresponding to the new values for a limit variable.
+    fieldname : np.ndarray
+        name of the field to change.
+    limit_file : str
+        The path (relative or absolute) to the limit.nc file;
+
+    Returns
+    -------
+    """
+    north_pole = np.transpose([np.mean(newfield[:,0], -1)])
+    south_pole = np.transpose([np.mean(newfield[:,-1], -1)])
+    if len(np.shape(newfield)) == 3 : # If time dimension
+        newfield_flat = [newfield[t, 1:-1].flatten() for t in range(len(newfield))]
+    else :
+        newfield_flat = newfield[1:-1].flatten()
+
+    newfield_flat = np.concatenate([north_pole, newfield_flat, south_pole], 1)
+
+    f_in = Dataset(limit_file, 'a')
+    f_in.variables[fieldname][:]=newfield_flat
+    f_in.close()
+
+    return None
 
 if __name__ == "__main__":
     pass