Changeset 3818 for trunk/LMDZ.MARS/util/display_netcdf.py

Timestamp:

Jul 1, 2025, 11:31:08 AM (2 weeks ago)

Author:

jbclement

Message:

Mars PCM:
In the script "display_netcdf.py" ("util" folder), addition of the possibility to display a 3D globe view of 2D maps. An environment file is also added to set up quickly the libraries and modules needed for the Python script.
JBC

File:

• trunk/LMDZ.MARS/util/display_netcdf.py (modified) (10 diffs)

trunk/LMDZ.MARS/util/display_netcdf.py

@@ -7 +7 @@
 This script can display any numeric variable from a NetCDF file.
 It supports the following cases:
-  - 1D time series (Time)
-  - 1D vertical profiles (e.g., subsurface_layers)
+  - Scalar output
+  - 1D time series
+  - 1D vertical profiles
   - 2D latitude/longitude map
-  - 2D (Time × another dimension)
-  - Variables with dimension “physical_points” displayed on a 2D map if lat/lon are present
+  - 2D cross-sections
   - Optionally average over latitude and plot longitude vs. time heatmap
-  - Scalar output (ndim == 0 after slicing)
-  - 2D cross-sections (altitude × latitude or altitude × longitude)
+  - Optionally display polar stereographic view of 2D maps
+  - Optionally display 3D globe view of 2D maps
+
+Automatic setup from the environment file found in the "LMDZ.MARS/util" folder:
+  1. Make sure Conda is installed.
+  2. In terminal, navigate to the folder containing this script.
+  4. Create the environment:
+       conda env create -f display_netcdf.yml
+  5. Activate the environment:
+       conda activate my_env
+  6. Run the script:
+     python display_netcdf.py
 
 Usage:
   1) Command-line mode:
-       python display_netcdf.py /path/to/your_file.nc --variable VAR_NAME \
-           [--time-index 0] [--alt-index 0] [--cmap viridis] [--avg-lat] \
-           [--slice-lon-index 10] [--slice-lat-index 20] [--show-topo] \
-           [--output out.png] [--extra-indices '{"nslope": 1}']
-
+       python display_netcdf.py /path/to/your_file.nc --variable VAR_NAME [options]
+     Options:
     --variable         : Name of the variable to visualize.
     --time-index       : Index along the Time dimension (0-based, ignored for purely 1D time series).
@@ -31 +38 @@
     --slice-lat-index  : Fixed latitude index for altitude×latitude cross-section.
     --show-topo        : Overlay MOLA topography on lat/lon maps.
+    --show-polar       :
+    --show-3d          :
     --output           : If provided, save the figure to this filename instead of displaying.
     --extra-indices    : JSON string to fix indices for any other dimensions.
@@ -38 +47 @@
   2) Interactive mode:
        python display_netcdf.py
-       (The script will prompt for everything, including averaging or slicing options.)
+     The script will prompt for everything.
 """
 
@@ -49 +58 @@
 import numpy as np
 import matplotlib.pyplot as plt
-import matplotlib.tri as mtri
 import matplotlib.path as mpath
+import matplotlib.colors as mcolors
 import cartopy.crs as ccrs
+import pandas as pd
 from netCDF4 import Dataset
+
+# Attempt vedo import early for global use
+try:
+    import vedo
+    from vedo import *
+    from scipy.interpolate import RegularGridInterpolator
+    vedo_available = True
+except ImportError:
+    vedo_available = False
 
 # Constants for recognized dimension names
@@ -59 +78 @@
 LAT_DIMS  = ("latitude", "lat")
 LON_DIMS  = ("longitude", "lon")
+
+# Paths for MOLA data
+MOLA_NPY = 'MOLA_1px_per_deg.npy'
+MOLA_CSV = 'molaTeam_contour_31rgb_steps.csv'
 
 # Attempt to load MOLA topography
@@ -68 +91 @@
     topo_lon2d, topo_lat2d = np.meshgrid(topo_lons, topo_lats)
     topo_loaded = True
-    print("MOLA topography loaded successfully from 'MOLA_1px_per_deg.npy'.")
 except Exception as e:
+    print(f"Warning: '{MOLA_NPY}' not found: {e}")
     topo_loaded = False
-    print(f"Warning: failed to load MOLA topography ('MOLA_1px_per_deg.npy'): {e}")
+
+
+# Attempt to load contour color table
+if os.path.isfile(MOLA_CSV):
+    color_table = pd.read_csv(MOLA_CSV)
+    csv_loaded = True
+else:
+    print(f"Warning: '{MOLA_CSV}' not found. 3D view colors disabled.")
+    csv_loaded = False
 
 
@@ -204 +235 @@
     # Show both figures
     plt.show()
+
+
+def plot_3D_globe(lon2d, lat2d, data2d, colormap, varname, units=None):
+    """
+    Plot a 3D globe view of the data using vedo, with surface coloring based on data2d
+    and overlaid contour lines from MOLA topography.
+    """
+    if not vedo_available:
+        print("3D view skipped: vedo missing.")
+        return
+    if not csv_loaded:
+        print("3D view skipped: color table missing.")
+        return
+
+    # Prepare MOLA grid
+    nlat, nlon = MOLA.shape
+    lats = np.linspace(90, -90, nlat)
+    lons = np.linspace(-180, 180, nlon)
+    lon_grid, lat_grid = np.meshgrid(lons, lats)
+
+    # Interpolate data2d onto MOLA grid
+    lat_data = np.linspace(-90, 90, data2d.shape[0])
+    lon_data = np.linspace(-180, 180, data2d.shape[1])
+    interp2d = RegularGridInterpolator((lat_data, lon_data), data2d,
+                                       bounds_error=False, fill_value=None)
+    newdata2d = interp2d((lat_grid, lon_grid))
+
+    # Generate contour lines from MOLA
+    cs = plt.contour(lon_grid, lat_grid, MOLA, levels=10, linewidths=0)
+    plt.clf()
+    contour_lines = []
+    radius = 3389500 # Mars average radius [m]
+    for segs, level in zip(cs.allsegs, cs.levels):
+        for verts in segs:
+            lon_c = verts[:, 0]
+            lat_c = verts[:, 1]
+            phi_c = np.radians(90 - lat_c)
+            theta_c = np.radians(lon_c)
+            elev = RegularGridInterpolator((lats, lons), MOLA,
+                                           bounds_error=False,
+                                           fill_value=0.0)((lat_c, lon_c))
+            r_cont = radius + elev * 10
+            x_c = r_cont * np.sin(phi_c) * np.cos(theta_c) * 1.002
+            y_c = r_cont * np.sin(phi_c) * np.sin(theta_c) * 1.002
+            z_c = r_cont * np.cos(phi_c) * 1.002
+            pts = np.column_stack([x_c, y_c, z_c])
+            if pts.shape[0] > 1:
+                contour_lines.append(Line(pts, c='k', lw=0.5))
+
+    # Create sphere surface mesh
+    phi = np.deg2rad(90 - lat_grid)
+    theta = np.deg2rad(lon_grid)
+    r = radius + MOLA * 10
+    x = r * np.sin(phi) * np.cos(theta)
+    y = r * np.sin(phi) * np.sin(theta)
+    z = r * np.cos(phi)
+    pts = np.stack([x.ravel(), y.ravel(), z.ravel()], axis=1)
+
+    # Build mesh faces
+    faces = []
+    for i in range(nlat - 1):
+        for j in range(nlon - 1):
+            p0 = i * nlon + j
+            p1 = p0 + 1
+            p2 = p0 + nlon
+            p3 = p2 + 1
+            faces.extend([(p0, p2, p1), (p1, p2, p3)])
+
+    mesh = Mesh([pts, faces])
+    mesh.cmap(colormap, newdata2d.ravel())
+    mesh.add_scalarbar(title=varname + (f' [{units}]' if units else ''), c='white')
+    mesh.lighting('default')
+
+    # Geographic grid lines
+    meridians, parallels, labels = [], [], []
+    zero_lon_offset = radius * 0.03
+    for lon in range(-150, 181, 30):
+        lat_line = np.linspace(-90, 90, nlat)
+        lon_line = np.full_like(lat_line, lon)
+        phi = np.deg2rad(90 - lat_line)
+        theta = np.deg2rad(lon_line)
+        elev = RegularGridInterpolator((lats, lons), MOLA)((lat_line, lon_line))
+        rr = radius + elev * 10
+        pts_line = np.column_stack([
+            rr * np.sin(phi) * np.cos(theta),
+            rr * np.sin(phi) * np.sin(theta),
+            rr * np.cos(phi)
+        ]) * 1.005
+        label_pos = pts_line[len(pts_line)//2]
+        norm = np.linalg.norm(label_pos)
+        label_pos_out = label_pos / norm * (norm + radius * 0.02)
+        if lon == 0:
+            label_pos_out[1] += zero_lon_offset
+        meridians.append(Line(pts_line, c='k', lw=1)#.flagpole(
+            #f"{lon}°",
+            #point=label_pos_out,
+            #offset=[0, 0, radius * 0.05],
+            #s=radius*0.01,
+            #c='yellow'
+        #).follow_camera()
+        )
+
+    for lat in range(-60, 91, 30):
+        lon_line = np.linspace(-180, 180, nlon)
+        lat_line = np.full_like(lon_line, lat)
+        phi = np.deg2rad(90 - lat_line)
+        theta = np.deg2rad(lon_line)
+        elev = RegularGridInterpolator((lats, lons), MOLA)((lat_line, lon_line))
+        rr = radius + elev * 10
+        pts_line = np.column_stack([
+            rr * np.sin(phi) * np.cos(theta),
+            rr * np.sin(phi) * np.sin(theta),
+            rr * np.cos(phi)
+        ]) * 1.005
+        label_pos = pts_line[len(pts_line)//2]
+        norm = np.linalg.norm(label_pos)
+        label_pos_out = label_pos / norm * (norm + radius * 0.02)
+        parallels.append(Line(pts_line, c='k', lw=1)#.flagpole(
+            #f"{lat}°",
+            #point=label_pos_out,
+            #offset=[0, 0, radius * 0.05],
+            #s=radius*0.01,
+            #c='yellow'
+        #).follow_camera()
+        )
+
+    # Create plotter
+    plotter = Plotter(title="3D topography view", bg="bb", axes=0)
+
+    # Configure camera
+    cam_dist = radius * 3
+    plotter.camera.SetPosition([cam_dist, 0, 0])
+    plotter.camera.SetFocalPoint([0, 0, 0])
+    plotter.camera.SetViewUp([0, 0, 1])
+
+    # Show the globe
+    plotter.show(mesh, *contour_lines, *meridians, *parallels)
+
 
 def plot_variable(dataset, varname, time_index=None, alt_index=None,
@@ -400 +569 @@
 
                     # Prompt for polar-stereo views if interactive
-                    if sys.stdin.isatty() and input("Display polar-stereo views? [y/n]: ").strip().lower() == "y":
+                    if input("Display polar-stereo views? [y/n]: ").strip().lower() == "y":
                         units = getattr(dataset.variables[varname], "units", None)
                         plot_polar_views(lon2d, lat2d, data2d, colormap, varname, units)
+
+                    # Prompt for 3D globe view if interactive
+                    if input("Display 3D globe view? [y/n]: ").strip().lower() == "y":
+                        units = getattr(dataset.variables[varname], "units", None)
+                        plot_3D_globe(lon2d, lat2d, data2d, colormap, varname, units)
 
                     if output_path:
@@ -510 +684 @@
                 plot_polar_views(lon2d, lat2d, dslice, colormap, varname, units)
 
+            # Prompt for 3D globe view if interactive
+            if sys.stdin.isatty() and input("Display 3D globe view? [y/n]: ").strip().lower() == "y":
+                units = getattr(dataset.variables[varname], "units", None)
+                plot_3D_globe(lon2d, lat2d, dslice, colormap, varname, units)
+
             if output_path:
                 plt.savefig(output_path, bbox_inches="tight")
@@ -753 +932 @@
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("nc_file", nargs="?", help="NetCDF file (omit for interactive)")
-    parser.add_argument("-v", "--variable", help="Variable name")
-    parser.add_argument("-t", "--time-index", type=int, help="Time index (0-based)")
-    parser.add_argument("-a", "--alt-index", type=int, help="Altitude index (0-based)")
-    parser.add_argument("-c", "--cmap", default="jet", help="Colormap")
-    parser.add_argument("--avg-lat", action="store_true",
-                        help="Average over latitude (time × lon heatmap)")
-    parser.add_argument("-o", "--output", help="Save figure path")
-    parser.add_argument("-e", "--extra-indices", help="JSON for other dims")
+    parser.add_argument('nc_file', nargs='?', help='NetCDF file (omit for interactive)')
+    parser.add_argument('-v','--variable', help='Variable name')
+    parser.add_argument('-t','--time-index', type=int, help='Time index (0-based)')
+    parser.add_argument('-a','--alt-index', type=int, help='Altitude index (0-based)')
+    parser.add_argument('-c','--cmap', default='jet', help='Colormap')
+    parser.add_argument('--avg-lat', action='store_true', help='Average over latitude')
+    parser.add_argument('--show-polar', action='store_true', help='Enable polar-stereo views')
+    parser.add_argument('--show-3d', action='store_true', help='Enable 3D globe view')
+    parser.add_argument('-o','--output', help='Save figure path')
+    parser.add_argument('-e','--extra-indices', help='JSON string for other dims')
     args = parser.parse_args()