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

Timestamp:

Jul 7, 2025, 6:00:37 PM (5 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

File:

• trunk/LMDZ.MARS/util/display_netcdf.py (modified) (5 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: