Add visualize sos CLI command and heatmap --vmin/--vmax

samples/pipelines/sos_fv3chem_frames.yaml

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Render Science On a Sphere (SOS) frames from a sequence of gridded NetCDF
+# files (e.g., FV3-Chem output) and compose them into an MP4.
+#
+# Key points for production-quality SOS animations:
+#   * SOS frames are full-globe, PlateCarree, 2:1, edge-to-edge images.
+#   * A FIXED color range (vmin/vmax) keeps the color scaling identical across
+#     every frame, which avoids the flicker caused by per-frame self-scaling.
+#
+# Replace <VAR>, the input paths, and vmin/vmax with values appropriate for
+# your dataset.
+name: "sos: FV3-Chem frames"
+stages:
+  # 1) Render each NetCDF timestep to a 2:1 SOS frame with a fixed color scale.
+  - stage: visualize
+    command: sos
+    args:
+      inputs:
+        - /data/fv3chem/t000.nc
+        - /data/fv3chem/t001.nc
+        - /data/fv3chem/t002.nc
+      output_dir: /data/fv3chem/frames
+      var: <VAR>
+      cmap: YlOrBr
+      width: 4096
+      height: 2048
+      # Fixed range -> consistent colors across frames (no flicker)
+      vmin: 0
+      vmax: 50
+
+  # 2) Compose the rendered frames into an MP4.
+  - stage: visualize
+    command: compose-video
+    args:
+      frames: /data/fv3chem/frames
+      output: /data/output/fv3chem_sos.mp4
+      fps: 24

src/zyra/pipeline_runner.py

@@ -265,6 +265,16 @@ def to_flag(k: str) -> str:
                 argv.append(to_flag(k))
         elif v is None:
             continue
+        elif isinstance(v, (list, tuple)):
+            # Expand to a multi-valued flag (argparse nargs), e.g.
+            # ``inputs: [a, b]`` -> ``--inputs a b`` and
+            # ``extent: [-180, 180, -90, 90]`` -> ``--extent -180 180 -90 90``.
+            # Skip empty sequences: emitting a bare ``--inputs`` with no values
+            # is a hard parse error for ``nargs='+'`` options.
+            if not v:
+                continue
+            argv.append(to_flag(k))
+            argv.extend(str(item) for item in v)
         else:
             argv.extend([to_flag(k), str(v)])
 

src/zyra/visualization/__init__.py

@@ -47,6 +47,7 @@
 from zyra.visualization.cli_contour import handle_contour
 from zyra.visualization.cli_heatmap import handle_heatmap
 from zyra.visualization.cli_interactive import handle_interactive
+from zyra.visualization.cli_sos import register_sos_cli
 from zyra.visualization.cli_timeseries import handle_timeseries
 from zyra.visualization.cli_vector import handle_vector
 
@@ -95,6 +96,16 @@ def register_cli(subparsers: Any) -> None:
     p_hm.add_argument("--height", type=int, default=512)
     p_hm.add_argument("--dpi", type=int, default=96)
     p_hm.add_argument("--cmap", default="YlOrBr")
+    p_hm.add_argument(
+        "--vmin",
+        type=float,
+        help="Fixed minimum data value for color scaling (use across frame sequences to avoid flicker)",
+    )
+    p_hm.add_argument(
+        "--vmax",
+        type=float,
+        help="Fixed maximum data value for color scaling (use across frame sequences to avoid flicker)",
+    )
     p_hm.add_argument("--colorbar", action="store_true")
     p_hm.add_argument("--label")
     p_hm.add_argument("--units")
@@ -603,3 +614,6 @@ def _levels_arg(val):
         help="Vector: render streamlines image overlay",
     )
     p_int.set_defaults(func=handle_interactive)
+
+    # sos (Science On a Sphere)
+    register_sos_cli(subparsers)

src/zyra/visualization/cli_heatmap.py

@@ -42,6 +42,8 @@ def handle_heatmap(ns) -> int:
                 height=ns.height,
                 dpi=ns.dpi,
                 cmap=ns.cmap,
+                vmin=getattr(ns, "vmin", None),
+                vmax=getattr(ns, "vmax", None),
                 colorbar=getattr(ns, "colorbar", False),
                 label=getattr(ns, "label", None),
                 units=getattr(ns, "units", None),
@@ -89,6 +91,8 @@ def handle_heatmap(ns) -> int:
         height=ns.height,
         dpi=ns.dpi,
         cmap=ns.cmap,
+        vmin=getattr(ns, "vmin", None),
+        vmax=getattr(ns, "vmax", None),
         colorbar=getattr(ns, "colorbar", False),
         label=getattr(ns, "label", None),
         units=getattr(ns, "units", None),

src/zyra/visualization/cli_register.py

@@ -10,6 +10,7 @@
 from .cli_contour import handle_contour
 from .cli_heatmap import handle_heatmap
 from .cli_interactive import handle_interactive
+from .cli_sos import register_sos_cli
 from .cli_timeseries import handle_timeseries
 from .cli_vector import handle_vector
 
@@ -52,6 +53,16 @@ def register_cli(subparsers: Any) -> None:
     p_hm.add_argument("--height", type=int, default=512)
     p_hm.add_argument("--dpi", type=int, default=96)
     p_hm.add_argument("--cmap", default="YlOrBr")
+    p_hm.add_argument(
+        "--vmin",
+        type=float,
+        help="Fixed minimum data value for color scaling (use across frame sequences to avoid flicker)",
+    )
+    p_hm.add_argument(
+        "--vmax",
+        type=float,
+        help="Fixed maximum data value for color scaling (use across frame sequences to avoid flicker)",
+    )
     p_hm.add_argument("--colorbar", action="store_true")
     p_hm.add_argument("--label")
     p_hm.add_argument("--units")
@@ -395,3 +406,6 @@ def register_cli(subparsers: Any) -> None:
         help="Shell-style trace of key steps and external commands",
     )
     p_int.set_defaults(func=handle_interactive)
+
+    # sos (Science On a Sphere)
+    register_sos_cli(subparsers)

src/zyra/visualization/cli_sos.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CLI handler for the ``visualize sos`` (Science On a Sphere) subcommand.
+
+This wires :class:`zyra.visualization.plot_manager.PlotManager` to the CLI so
+that gridded data can be rendered as Science On a Sphere (SOS) frames. SOS
+frames are full-globe, PlateCarree, 2:1, edge-to-edge images. A fixed color
+range (``--vmin``/``--vmax``) keeps the color scaling identical across every
+frame in a sequence, which avoids the flicker that results from per-frame
+self-scaling.
+
+Both single-input (``--input``/``--output``) and batch (``--inputs``/
+``--output-dir``) rendering are supported so that frame sequences for an
+animation can be produced in one invocation.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Any
+
+from zyra.utils.cli_helpers import configure_logging_from_env
+from zyra.visualization.cli_utils import load_data_array, resolve_basemap_ref
+
+
+def _render_one(
+    ns,
+    src: str,
+    dest: str,
+) -> str | None:
+    """Render a single SOS frame from ``src`` to ``dest``.
+
+    Returns the output path on success, or ``None`` on failure.
+    """
+    from zyra.visualization.plot_manager import PlotManager
+
+    arr = load_data_array(
+        src,
+        var=getattr(ns, "var", None),
+        xarray_engine=getattr(ns, "xarray_engine", None),
+    )
+    # Treat NaN/inf as transparent so missing data does not skew the rendering.
+    import numpy as np
+
+    arr = np.ma.masked_invalid(arr)
+
+    bmap, guard = resolve_basemap_ref(getattr(ns, "basemap", None))
+    try:
+        pm = PlotManager(basemap=bmap, image_extent=ns.extent, base_cmap=ns.cmap)
+        return pm.sos_plot_data(
+            arr,
+            custom_cmap=ns.cmap,
+            output_path=dest,
+            width=ns.width,
+            height=ns.height,
+            dpi=ns.dpi,
+            flip_data=bool(getattr(ns, "flip", False)),
+            vmin=getattr(ns, "vmin", None),
+            vmax=getattr(ns, "vmax", None),
+        )
+    finally:
+        if guard is not None:
+            try:
+                guard.close()
+            except Exception:
+                pass
+
+
+def handle_sos(ns) -> int:
+    """Handle the ``visualize sos`` CLI subcommand.
+
+    Renders one or more gridded inputs as Science On a Sphere frames using
+    :class:`PlotManager`, with a fixed ``--vmin``/``--vmax`` color range for
+    flicker-free frame sequences.
+    """
+    if getattr(ns, "verbose", False):
+        os.environ["ZYRA_VERBOSITY"] = "debug"
+    elif getattr(ns, "quiet", False):
+        os.environ["ZYRA_VERBOSITY"] = "quiet"
+    if getattr(ns, "trace", False):
+        os.environ["ZYRA_SHELL_TRACE"] = "1"
+    configure_logging_from_env()
+
+    # Batch mode: --inputs with --output-dir
+    if getattr(ns, "inputs", None):
+        outdir = getattr(ns, "output_dir", None)
+        if not outdir:
+            raise SystemExit("--output-dir is required when using --inputs")
+        outdir_p = Path(outdir)
+        outdir_p.mkdir(parents=True, exist_ok=True)
+        outputs: list[str] = []
+        failures: list[str] = []
+        for src in ns.inputs:
+            dest = outdir_p / f"{Path(str(src)).stem}.png"
+            out = _render_one(ns, src, str(dest))
+            if out:
+                logging.info(out)
+                outputs.append(out)
+            else:
+                failures.append(str(src))
+        try:
+            print(json.dumps({"outputs": outputs}))
+        except Exception:
+            pass
+        # Surface render failures as a non-zero exit so that `zyra run`
+        # pipelines do not treat a failed/empty render as success.
+        if failures:
+            raise SystemExit(
+                f"Failed to render {len(failures)} SOS frame(s): {', '.join(failures)}"
+            )
+        return 0
+
+    # Single input mode
+    if not getattr(ns, "input", None):
+        raise SystemExit("--input or --inputs is required")
+    if not getattr(ns, "output", None):
+        raise SystemExit("--output is required when using --input")
+    out = _render_one(ns, ns.input, ns.output)
+    if not out:
+        raise SystemExit(f"Failed to render SOS frame from {ns.input}")
+    logging.info(out)
+    return 0
+
+
+def register_sos_cli(subparsers: Any) -> None:
+    """Register the ``sos`` subcommand on a provided subparsers object.
+
+    Shared by both the CLI registration path
+    (:func:`zyra.visualization.cli_register.register_cli`) and the API manifest
+    path (:func:`zyra.visualization.register_cli`) to keep the parser definition
+    in a single place.
+    """
+    p_sos = subparsers.add_parser(
+        "sos",
+        help="Visualization: render Science On a Sphere frames",
+        description=(
+            "Render gridded data as Science On a Sphere (SOS) frames: full-globe, "
+            "PlateCarree, 2:1, edge-to-edge PNGs. Use a fixed --vmin/--vmax range to "
+            "keep color scaling identical across a frame sequence (flicker-free). "
+            "Supports single (--input/--output) and batch (--inputs/--output-dir) modes."
+        ),
+    )
+    p_sos.add_argument("--input", help="Path to .nc or .npy input")
+    p_sos.add_argument(
+        "--inputs", nargs="+", help="Multiple input paths for batch rendering"
+    )
+    p_sos.add_argument("--var", help="Variable name for NetCDF inputs")
+    p_sos.add_argument(
+        "--output",
+        help="Output PNG path (required when using --input)",
+    )
+    p_sos.add_argument(
+        "--output-dir",
+        dest="output_dir",
+        help="Directory to write outputs (required when using --inputs)",
+    )
+    p_sos.add_argument(
+        "--basemap",
+        help="Optional basemap (path, bare image name, or pkg:ref) drawn under the data",
+    )
+    p_sos.add_argument(
+        "--extent",
+        nargs=4,
+        type=float,
+        default=[-180, 180, -90, 90],
+        help="west east south north (default: global -180 180 -90 90)",
+    )
+    p_sos.add_argument("--width", type=int, default=4096, help="Output width (px)")
+    p_sos.add_argument("--height", type=int, default=2048, help="Output height (px)")
+    p_sos.add_argument("--dpi", type=int, default=96)
+    p_sos.add_argument("--cmap", default="YlOrBr", help="Colormap name")
+    p_sos.add_argument(
+        "--vmin",
+        type=float,
+        help="Fixed minimum data value for color scaling (recommended for sequences)",
+    )
+    p_sos.add_argument(
+        "--vmax",
+        type=float,
+        help="Fixed maximum data value for color scaling (recommended for sequences)",
+    )
+    p_sos.add_argument(
+        "--flip",
+        action="store_true",
+        help="Flip data vertically before rendering (for north-up grids)",
+    )
+    p_sos.add_argument(
+        "--xarray-engine",
+        dest="xarray_engine",
+        help="xarray engine for NetCDF inputs (e.g., netcdf4, h5netcdf, scipy)",
+    )
+    p_sos.add_argument(
+        "--verbose", action="store_true", help="Verbose logging for this command"
+    )
+    p_sos.add_argument(
+        "--quiet", action="store_true", help="Quiet logging for this command"
+    )
+    p_sos.add_argument(
+        "--trace",
+        action="store_true",
+        help="Shell-style trace of key steps and external commands",
+    )
+    p_sos.set_defaults(func=handle_sos)