Changeset 3839 for trunk/LMDZ.MARS/util

Timestamp:

Jul 7, 2025, 6:00:37 PM (4 days ago)

Author:

jbclement

Message:

Mars PCM:
In "display_netcdf.py", addition of the possibility to display altitude as function of time for 1D variables + showing value at cursor for 2D heatmaps.
JBC

Location:

trunk/LMDZ.MARS/util

Files:

• display_netcdf.py (modified) (5 diffs)
• display_netcdf.yml (modified) (9 diffs)

trunk/LMDZ.MARS/util/display_netcdf.py

@@ -85 +85 @@
 # Attempt to load MOLA topography
 try:
-    MOLA = np.load('MOLA_1px_per_deg.npy')  # shape (nlat, nlon) at 1° per pixel: lat from -90 to 90, lon from 0 to 360
+    MOLA = np.load('MOLA_1px_per_deg.npy') # shape (nlat, nlon) at 1° per pixel: lat from -90 to 90, lon from 0 to 360
     nlat, nlon = MOLA.shape
     topo_lats = np.linspace(90 - 0.5, -90 + 0.5, nlat)
@@ -171 +171 @@
         transform=transform
     )
+
+
+def attach_format_coord(ax, mat, x, y, is_pcolormesh=True):
+    """
+    Attach a format_coord function to the axes to display x, y, and value at cursor.
+    Works for both pcolormesh and imshow style grids.
+    """
+    # Determine dimensions
+    if mat.ndim == 2:
+        ny, nx = mat.shape
+    elif mat.ndim == 3 and mat.shape[2] in (3, 4):
+        ny, nx, nc = mat.shape
+    else:
+        raise ValueError(f"Unsupported mat shape {mat.shape}")
+    # Edges or extents
+    if is_pcolormesh:
+        xedges, yedges = x, y
+    else:
+        x0, x1 = x.min(), x.max()
+        y0, y1 = y.min(), y.max()
+
+    def format_coord(xp, yp):
+        # Map to indices
+        if is_pcolormesh:
+            col = np.searchsorted(xedges, xp) - 1
+            row = np.searchsorted(yedges, yp) - 1
+        else:
+            col = int((xp - x0) / (x1 - x0) * nx)
+            row = int((yp - y0) / (y1 - y0) * ny)
+        # Within bounds?
+        if 0 <= row < ny and 0 <= col < nx:
+            if mat.ndim == 2:
+                v = mat[row, col]
+                return f"x={xp:.3g}, y={yp:.3g}, val={v:.3g}"
+            else:
+                vals = mat[row, col]
+                txt = ", ".join(f"{vv:.3g}" for vv in vals[:3])
+                return f"x={xp:.3g}, y={yp:.3g}, val=({txt})"
+        return f"x={xp:.3g}, y={yp:.3g}"
+
+    ax.format_coord = format_coord
 
 
@@ -396 +437 @@
         data_full = np.where(data_full.mask, np.nan, data_full.data)
 
+    # If Time and altitude are both present and neither indexed,
+    # and every other dim has size 1:
+    t_idx = find_dim_index(dims, TIME_DIMS)
+    a_idx = find_dim_index(dims, ALT_DIMS)
+    shape = var.shape
+    if (t_idx is not None and a_idx is not None
+        and time_index is None and alt_index is None
+        and all(size == 1 for i, size in enumerate(shape) if i not in (t_idx, a_idx))):
+
+        # Build a slicer that keeps Time & altitude, drops other singletons
+        slicer = [0] * len(dims)
+        slicer[t_idx] = slice(None)
+        slicer[a_idx] = slice(None)
+        data2d = data_full[tuple(slicer)]  # shape (ntime, nalt)
+
+        # Coordinate arrays
+        tvar = find_coord_var(dataset, TIME_DIMS)
+        avar = find_coord_var(dataset, ALT_DIMS)
+        tvals = dataset.variables[tvar][:]
+        avals = dataset.variables[avar][:]
+
+        # Unmask if necessary
+        if hasattr(tvals, "mask"):
+            tvals = np.where(tvals.mask, np.nan, tvals.data)
+        if hasattr(avals, "mask"):
+            avals = np.where(avals.mask, np.nan, avals.data)
+
+        # Plot heatmap with x=time, y=altitude
+        fig, ax = plt.subplots(figsize=(10, 6))
+        T, A = np.meshgrid(tvals, avals)
+        im = ax.pcolormesh(
+            T, A, data2d.T,
+            shading="auto", cmap=colormap
+        )
+        dt = tvals[1] - tvals[0]
+        da = avals[1] - avals[0]
+        x_edges = np.concatenate([tvals - dt/2, [tvals[-1] + dt/2]])
+        y_edges = np.concatenate([avals - da/2, [avals[-1] + da/2]])
+        attach_format_coord(ax, data2d.T, x_edges, y_edges, is_pcolormesh=True)
+        ax.set_xlabel(tvar)
+        ax.set_ylabel(avar)
+        cbar = fig.colorbar(im, ax=ax)
+        cbar.set_label(varname + (f" ({getattr(var, 'units','')})"))
+        ax.set_title(f"{varname} — {avar} vs {tvar}", fontweight="bold")
+
+        if output_path:
+            fig.savefig(output_path, bbox_inches="tight")
+            print(f"Saved to {output_path}")
+        else:
+            plt.show()
+        return
+
     # Pure 1D time series
     if len(dims) == 1 and find_dim_index(dims, TIME_DIMS) is not None:
@@ -434 +527 @@
         if hasattr(lons, "mask"):
             lons = np.where(lons.mask, np.nan, lons.data)
-        plt.figure(figsize=(10, 6))
-        plt.pcolormesh(lons, tvals, data_avg, shading="auto", cmap=colormap)
-        plt.xlabel(f"Longitude ({getattr(dataset.variables[lon_var], 'units', 'deg')})")
-        plt.ylabel(time_var)
-        cbar = plt.colorbar()
+        fig, ax = ax.subplots(figsize=(10, 6))
+        im = plt.pcolormesh(lons, tvals, data_avg, shading="auto", cmap=colormap)
+        dx = lons[1] - lons[0]
+        dy = tvals[1] - tvals[0]
+        x_edges = np.concatenate([lons - dx/2, [lons[-1] + dx/2]])
+        y_edges = np.concatenate([tvals - dy/2, [tvals[-1] + dy/2]])
+        attach_format_coord(ax, data_avg.T, x_edges, y_edges, is_pcolormesh=True)
+        ax.set_xlabel(f"Longitude ({getattr(dataset.variables[lon_var], 'units', 'deg')})")
+        ax.set_ylabel(time_var)
+        cbar = fig.colorbar(im, ax=ax)
         cbar.set_label(varname + (f" ({var.units})" if hasattr(var, "units") else ""))
-        plt.title(f"{varname} averaged over latitude", fontweight='bold')
+        ax.set_title(f"{varname} averaged over latitude", fontweight='bold')
         if output_path:
-            plt.savefig(output_path, bbox_inches="tight")
+            fig.savefig(output_path, bbox_inches="tight")
             print(f"Saved to {output_path}")
         else:
@@ -700 +798 @@
 
         # Generic 2D
-        plt.figure(figsize=(8, 6))
-        plt.imshow(dslice, aspect="auto")
-        plt.colorbar(label=varname + (f" ({var.units})" if hasattr(var, "units") else ""))
-        plt.xlabel("Dim 2 index")
-        plt.ylabel("Dim 1 index")
-        plt.title(f"{varname} (2D)", fontweight='bold')
+        fig, ax = plt.subplots(figsize=(8, 6))
+        im = ax.imshow(
+            dslice,
+            aspect="auto",
+            interpolation='nearest'
+        )
+        x0, x1 = 0, dslice.shape[1] - 1
+        y0, y1 = 0, dslice.shape[0] - 1
+        x_centers = np.linspace(x0, x1, dslice.shape[1])
+        y_centers = np.linspace(y0, y1, dslice.shape[0])
+        attach_format_coord(ax, dslice, x_centers, y_centers, is_pcolormesh=False)
+        cbar = fig.colorbar(im, ax=ax, orientation='vertical')
+        cbar.set_label(varname + (f" ({var.units})" if hasattr(var, "units") else ""))
+        ax.set_xlabel("Dim 2 index")
+        ax.set_ylabel("Dim 1 index")
+        ax.set_title(f"{varname} (2D)")
+
         if output_path:
-            plt.savefig(output_path, bbox_inches="tight")
+            fig.savefig(output_path, bbox_inches="tight")
             print(f"Saved to {output_path}")
         else:

trunk/LMDZ.MARS/util/display_netcdf.yml

@@ -8 +8 @@
   - aiohappyeyeballs=2.6.1=pyhd8ed1ab_0
   - aiohttp=3.12.13=py312h178313f_0
-  - aiosignal=1.3.2=pyhd8ed1ab_0
+  - aiosignal=1.4.0=pyhd8ed1ab_0
   - alsa-lib=1.2.14=hb9d3cd8_0
   - aom=3.9.1=hac33072_0
@@ -38 +38 @@
   - fonts-conda-ecosystem=1=0
   - fonts-conda-forge=1=0
-  - fonttools=4.58.4=py312h178313f_0
+  - fonttools=4.58.5=py312h178313f_0
   - freetype=2.13.3=ha770c72_1
   - fribidi=1.0.10=h36c2ea0_0
@@ -44 +44 @@
   - gdk-pixbuf=2.42.12=hb9ae30d_0
   - geos=3.13.1=h97f6797_0
-  - gettext=0.24.1=h5888daf_0
-  - gettext-tools=0.24.1=h5888daf_0
+  - gettext=0.25.1=h5888daf_0
+  - gettext-tools=0.25.1=h5888daf_0
   - gl2ps=1.4.2=hae5d5c5_1
   - gmp=6.3.0=hac33072_2
@@ -56 +56 @@
   - jsoncpp=1.9.6=hf42df4d_1
   - keyutils=1.6.1=h166bdaf_0
-  - kiwisolver=1.4.8=py312h84d6215_0
+  - kiwisolver=1.4.8=py312h68727a3_1
   - krb5=1.21.3=h659f571_0
   - lame=3.100=h166bdaf_1003
   - lcms2=2.17=h717163a_0
-  - ld_impl_linux-64=2.43=h1423503_5
+  - ld_impl_linux-64=2.44=h1423503_0
   - lerc=4.0.0=h0aef613_1
   - level-zero=1.23.0=h84d6215_0
   - libabseil=20250127.1=cxx17_hbbce691_0
   - libaec=1.1.4=h3f801dc_0
-  - libasprintf=0.24.1=h8e693c7_0
-  - libasprintf-devel=0.24.1=h8e693c7_0
+  - libasprintf=0.25.1=h8e693c7_0
+  - libasprintf-devel=0.25.1=h8e693c7_0
   - libass=0.17.3=h52826cd_2
   - libblas=3.9.0=32_h59b9bed_openblas
@@ -91 +91 @@
   - libgcc-ng=15.1.0=h69a702a_3
   - libgcrypt-lib=1.11.1=hb9d3cd8_0
-  - libgettextpo=0.24.1=h5888daf_0
-  - libgettextpo-devel=0.24.1=h5888daf_0
+  - libgettextpo=0.25.1=h5888daf_0
+  - libgettextpo-devel=0.25.1=h5888daf_0
   - libgfortran=15.1.0=h69a702a_3
   - libgfortran5=15.1.0=hcea5267_3
@@ -129 +129 @@
   - libopus=1.5.2=hd0c01bc_0
   - libpciaccess=0.18=hb9d3cd8_0
-  - libpng=1.6.49=h943b412_0
+  - libpng=1.6.50=h943b412_0
   - libpq=17.5=h27ae623_0
   - libprotobuf=5.29.3=h501fc15_1
@@ -179 +179 @@
   - pango=1.56.4=hadf4263_0
   - pcre2=10.45=hc749103_0
-  - pillow=11.2.1=py312h80c1187_0
+  - pillow=11.3.0=py312h80c1187_0
   - pip=25.1.1=pyh8b19718_0
   - pixman=0.46.2=h29eaf8c_0
@@ -211 +211 @@
   - tk=8.6.13=noxft_hd72426e_102
   - tornado=6.5.1=py312h66e93f0_0
-  - typing_extensions=4.14.0=pyhe01879c_0
+  - typing_extensions=4.14.1=pyhe01879c_0
   - tzdata=2025b=h78e105d_0
   - unicodedata2=16.0.0=py312h66e93f0_0
@@ -219 +219 @@
   - vtk-base=9.4.2=py312h28cca35_1
   - vtk-io-ffmpeg=9.4.2=he57672f_1
-  - wayland=1.23.1=h3e06ad9_1
+  - wayland=1.24.0=h3e06ad9_0
   - wayland-protocols=1.45=hd8ed1ab_0
   - wheel=0.45.1=pyhd8ed1ab_1