Add lightweight ParaView plotting interface

Co-authored-by: Cursor <cursoragent@cursor.com>

src/pv_plot/README.md

@@ -0,0 +1,53 @@
+# pv_plot
+
+`pv_plot` 是 DynamisLab 内部的轻量 ParaView 画图接口，只负责读取流场、应用少量预设并输出 PNG/MP4。它不是独立 pip 项目，没有服务、socket、worker pool 或命令行框架。
+
+## 环境
+
+总控使用 `pycuda_3_10`，ParaView 由仓库根目录的 `ParaView/bin/pvpython` 独立运行。在仓库根目录运行模块时无需安装：
+
+```bash
+CUDA_VISIBLE_DEVICES=2 conda run -n pycuda_3_10 python -m your_module
+```
+
+若显式将源码目录加入 `PYTHONPATH`，也可使用顶层导入 `from pv_plot import plot`：
+
+```bash
+CUDA_VISIBLE_DEVICES=2 PYTHONPATH="$PWD/src" conda run -n pycuda_3_10 python your_script.py
+```
+
+ParaView 5.11.2 自带 Python 3.9。不要把 `ParaView/lib/python3.9/site-packages` 加到 conda Python 3.10 的 `PYTHONPATH`。如果 ParaView 不在仓库根，可设置 `PV_PLOT_PVPYTHON`；FFmpeg 可用 `PV_PLOT_FFMPEG` 覆盖。
+
+## 使用
+
+仓库内模块使用：
+
+```python
+from src.pv_plot import plot
+
+plot("run/fields.npz", "run/vorticity.png")
+plot("run/fields.npz", "run/speed.mp4", preset="speed", fps=10)
+plot(
+    "run/fields.npz",
+    "run/error.png",
+    preset="comparison",
+    reference="run/q_in.npz",
+    scalar_range=(0.0, 0.5),
+)
+```
+
+新 CelerisLab 和 DRL eval 的 `ux/uy` 是 `(T, NY, NX)`，使用默认 `axis_order="tyx"`。Legacy CCD/OID 的 `(T, NX, NY)` 必须传 `axis_order="txy"`。单帧支持 `yx` 和 `xy`。工具不会猜测或静默转换未知方向。
+
+Legacy Tecplot 文件序列可直接传列表：
+
+```python
+plot(sorted(field_files), "legacy.mp4", axis_order="tyx")
+```
+
+支持的预设只有：
+
+- `vorticity`：有符号二维涡量；
+- `speed`：速度模长；
+- `comparison`：主场与参考场的速度误差。
+
+输出为 `.mp4` 时自动渲染全部帧并调用 FFmpeg；其他后缀按单帧截图处理。

src/pv_plot/__init__.py

@@ -0,0 +1,6 @@
+"""Small ParaView plotting interface for DynamisLab flow fields."""
+
+from .api import plot
+from .presets import PRESETS
+
+__all__ = ["PRESETS", "plot"]

src/pv_plot/api.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import Literal, Sequence
+
+from .presets import PRESETS
+
+Preset = Literal["vorticity", "speed", "comparison"]
+AxisOrder = Literal["tyx", "txy", "yx", "xy"]
+
+
+def _executable(value: str | Path | None, environment: str, fallback: Path | str) -> str:
+    configured = str(value or os.environ.get(environment) or fallback)
+    path = Path(configured).expanduser()
+    if path.parent != Path("."):
+        path = path.resolve()
+        if not path.is_file() or not os.access(path, os.X_OK):
+            raise FileNotFoundError(f"executable not found: {path}")
+        return str(path)
+    resolved = shutil.which(configured)
+    if not resolved:
+        raise FileNotFoundError(f"executable not found on PATH: {configured}")
+    return resolved
+
+
+def plot(
+    source: str | Path | Sequence[str | Path],
+    output: str | Path,
+    *,
+    preset: Preset = "vorticity",
+    reference: str | Path | Sequence[str | Path] | None = None,
+    axis_order: AxisOrder = "tyx",
+    scalar_range: tuple[float, float] | None = None,
+    frame: int = 0,
+    fps: int = 10,
+    field_names: tuple[str, str] = ("ux", "uy"),
+    show_colorbar: bool = True,
+    pvpython: str | Path | None = None,
+    ffmpeg: str | Path | None = None,
+    timeout: float = 1800,
+) -> Path:
+    """Render an NPZ field or Legacy Tecplot series with one ParaView subprocess.
+
+    NPZ files must contain velocity arrays named by ``field_names``. Their axis
+    order is explicit: new CelerisLab eval files use ``tyx`` and Legacy
+    CCD/OID files use ``txy``. A source sequence is treated as Tecplot files.
+    ``output`` ending in ``.mp4`` renders every frame; otherwise one PNG frame
+    is rendered.
+    """
+    if preset not in PRESETS:
+        raise ValueError(f"unknown preset {preset!r}; choose from {sorted(PRESETS)}")
+    if reference is not None and preset != "comparison":
+        raise ValueError("reference is only used by the comparison preset")
+    if preset == "comparison" and reference is None:
+        raise ValueError("comparison requires reference")
+    if axis_order not in {"tyx", "txy", "yx", "xy"}:
+        raise ValueError("axis_order must be tyx, txy, yx, or xy")
+    if scalar_range is not None and scalar_range[0] >= scalar_range[1]:
+        raise ValueError("scalar_range must satisfy low < high")
+
+    root = Path(__file__).resolve().parents[2]
+    pvpython_exe = _executable(pvpython, "PV_PLOT_PVPYTHON", root / "ParaView" / "bin" / "pvpython")
+    ffmpeg_exe = _executable(ffmpeg, "PV_PLOT_FFMPEG", "ffmpeg") if Path(output).suffix.lower() == ".mp4" else None
+
+    def paths(value):
+        values = value if isinstance(value, (list, tuple)) else [value]
+        resolved = [str(Path(item).expanduser().resolve()) for item in values]
+        missing = [item for item in resolved if not Path(item).is_file()]
+        if missing:
+            raise FileNotFoundError(f"input file not found: {missing[0]}")
+        return resolved
+
+    sources = paths(source)
+    references = paths(reference) if reference is not None else []
+    source_format = "npz" if len(sources) == 1 and Path(sources[0]).suffix.lower() == ".npz" else "tecplot"
+    if source_format == "npz" and len(sources) != 1:
+        raise ValueError("NPZ input accepts one archive containing all frames")
+
+    target = Path(output).expanduser().resolve()
+    target.parent.mkdir(parents=True, exist_ok=True)
+    config = {
+        "source": sources,
+        "reference": references,
+        "format": source_format,
+        "axis_order": axis_order,
+        "fields": list(field_names),
+        "preset": preset,
+        "range": list(scalar_range or PRESETS[preset]["range"]),
+        "colorbar": PRESETS[preset]["colorbar"] if show_colorbar else None,
+        "frame": int(frame),
+        "fps": int(fps),
+        "output": str(target),
+        "ffmpeg": ffmpeg_exe,
+    }
+    worker = Path(__file__).with_name("render.py")
+    with tempfile.NamedTemporaryFile("w", suffix=".json", encoding="utf-8", delete=False) as stream:
+        json.dump(config, stream)
+        config_path = Path(stream.name)
+    try:
+        completed = subprocess.run(
+            [pvpython_exe, str(worker), str(config_path)],
+            text=True,
+            capture_output=True,
+            timeout=timeout,
+        )
+    finally:
+        config_path.unlink(missing_ok=True)
+    if completed.returncode:
+        raise RuntimeError((completed.stderr or completed.stdout)[-4000:])
+    if not target.is_file():
+        raise RuntimeError(f"ParaView reported success but did not create {target}")
+    return target

src/pv_plot/presets.py

@@ -0,0 +1,7 @@
+"""Conservative plot defaults; callers may override every value."""
+
+PRESETS = {
+    "vorticity": {"range": [-0.1, 0.1], "colorbar": "Vorticity"},
+    "speed": {"range": [0.0, 2.0], "colorbar": "Speed"},
+    "comparison": {"range": [0.0, 1.0], "colorbar": "Velocity error"},
+}

src/pv_plot/render.py

@@ -0,0 +1,166 @@
+"""ParaView-side renderer. Run only with pvpython."""
+
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import paraview.simple as pv
+
+os.environ.setdefault("PV_DEBUG_SKIP_OPENGL_VERSION_CHECK", "1")
+
+
+def npz_source(config, files, name):
+    source = pv.ProgrammableSource(registrationName=name)
+    source.OutputDataSetType = "vtkImageData"
+    import numpy as np
+
+    with np.load(files[0], mmap_mode="r", allow_pickle=False) as archive:
+        u_name, v_name = config["fields"]
+        if u_name not in archive or v_name not in archive:
+            raise ValueError("NPZ does not contain configured velocity fields")
+        if archive[u_name].shape != archive[v_name].shape:
+            raise ValueError("velocity field shapes differ")
+        shape = archive[u_name].shape
+    order = config["axis_order"]
+    frames = shape[0] if order in ("tyx", "txy") and len(shape) == 3 else 1
+    if (order in ("tyx", "txy") and len(shape) != 3) or (order in ("yx", "xy") and len(shape) != 2):
+        raise ValueError("NPZ shape is incompatible with axis_order")
+    source.ScriptRequestInformation = """
+from vtkmodules.vtkCommonExecutionModel import vtkStreamingDemandDrivenPipeline
+info = self.GetExecutive().GetOutputInformation(0)
+times = list(range(%d))
+info.Set(vtkStreamingDemandDrivenPipeline.TIME_STEPS(), times, len(times))
+info.Set(vtkStreamingDemandDrivenPipeline.TIME_RANGE(), [times[0], times[-1]], 2)
+""" % frames
+    payload = json.dumps({"path": files[0], "fields": config["fields"], "order": order, "frames": frames})
+    source.Script = """
+import json
+import numpy as np
+from vtkmodules.numpy_interface import dataset_adapter as dsa
+from vtkmodules.vtkCommonExecutionModel import vtkStreamingDemandDrivenPipeline
+cfg = json.loads(%r)
+info = self.GetExecutive().GetOutputInformation(0)
+time = info.Get(vtkStreamingDemandDrivenPipeline.UPDATE_TIME_STEP())
+index = max(0, min(cfg['frames'] - 1, int(round(time or 0))))
+with np.load(cfg['path'], allow_pickle=False) as archive:
+    def read(name):
+        array = np.asarray(archive[name])
+        if cfg['order'] == 'tyx': array = array[index]
+        elif cfg['order'] == 'txy': array = array[index].T
+        elif cfg['order'] == 'xy': array = array.T
+        return np.ascontiguousarray(array)
+    u, v = read(cfg['fields'][0]), read(cfg['fields'][1])
+    output.SetDimensions(u.shape[1], u.shape[0], 1)
+    wrapped = dsa.WrapDataObject(output)
+    wrapped.PointData.append(u.ravel(order='C'), 'U')
+    wrapped.PointData.append(v.ravel(order='C'), 'V')
+""" % payload
+    return source
+
+
+def source(config, files, name):
+    if config["format"] == "npz":
+        return npz_source(config, files, name)
+    reader = pv.TecplotReader(registrationName=name, FileNames=files)
+    reader.DataArrayStatus = ["flag", "U", "V"]
+    return reader
+
+
+def vector_field(input_proxy, name="Velocity", u="U", v="V"):
+    zero = pv.Calculator(registrationName=name + "ZeroZ", Input=input_proxy)
+    zero.ResultArrayName = name + "Z"
+    zero.Function = "0"
+    vector = pv.MergeVectorComponents(registrationName=name, Input=zero)
+    vector.XArray, vector.YArray, vector.ZArray = u, v, name + "Z"
+    return vector
+
+
+def scalar_field(input_proxy, preset, reference=None):
+    if preset == "comparison":
+        difference = pv.ProgrammableFilter(registrationName="Difference", Input=[input_proxy, reference])
+        difference.Script = """import numpy as np
+u = inputs[0].PointData['U'] - inputs[1].PointData['U']
+v = inputs[0].PointData['V'] - inputs[1].PointData['V']
+output.PointData.append(np.sqrt(u*u + v*v), 'VelocityError')
+"""
+        return difference, "VelocityError", False
+    vector = vector_field(input_proxy)
+    if preset == "speed":
+        speed = pv.Calculator(registrationName="Speed", Input=vector)
+        speed.ResultArrayName, speed.Function = "Speed", "mag(Velocity)"
+        return speed, "Speed", False
+    gradient = pv.Gradient(registrationName="Vorticity", Input=vector)
+    gradient.ScalarArray = ["POINTS", "Velocity"]
+    gradient.ComputeGradient = 0
+    gradient.ComputeVorticity = 1
+    vorticity = pv.Calculator(registrationName="ZVorticity", Input=gradient)
+    vorticity.ResultArrayName, vorticity.Function = "Z_Vort", "Vorticity_Z"
+    return vorticity, "Z_Vort", True
+
+
+def show(proxy, view, array_name, value_range, diverging, colorbar):
+    low, high = value_range
+    middle = 0.0 if diverging and low < 0 < high else (low + high) / 2
+    if diverging:
+        points = [low, 0.231, 0.298, 0.753, middle, 1, 1, 1, high, 0.706, 0.016, 0.149]
+    else:
+        points = [low, 0.267, 0.005, 0.329, middle, 0.128, 0.567, 0.551, high, 0.993, 0.906, 0.144]
+    lut = pv.GetColorTransferFunction(array_name)
+    lut.AutomaticRescaleRangeMode, lut.RGBPoints = "Never", points
+    display = pv.Show(proxy, view)
+    display.ColorArrayName, display.LookupTable = ["POINTS", array_name], lut
+    display.SetScalarBarVisibility(view, bool(colorbar))
+    if colorbar:
+        pv.GetScalarBar(lut, view).Title = colorbar
+
+
+def main(config_path):
+    config = json.loads(Path(config_path).read_text())
+    primary = source(config, config["source"], "Flow")
+    reference = source(config, config["reference"], "Reference") if config["reference"] else None
+    if reference is not None:
+        primary.UpdatePipelineInformation()
+        reference.UpdatePipelineInformation()
+        if len(times(primary)) != len(times(reference)):
+            raise ValueError("source and reference frame counts differ")
+    scalar, array_name, diverging = scalar_field(primary, config["preset"], reference)
+    view = pv.CreateView("RenderView")
+    view.ViewSize = [1200, 480]
+    view.InteractionMode, view.OrientationAxesVisibility = "2D", 0
+    view.CameraParallelProjection, view.Background, view.UseLight = 1, [1, 1, 1], 0
+    show(scalar, view, array_name, config["range"], diverging, config["colorbar"])
+    values = times(primary)
+    output = Path(config["output"])
+    if output.suffix.lower() != ".mp4":
+        index = config["frame"]
+        if index < 0 or index >= len(values):
+            raise IndexError("frame is outside the available range")
+        primary.UpdatePipeline(time=values[index])
+        view.ResetCamera()
+        pv.Render(view)
+        pv.SaveScreenshot(str(output), view=view)
+        return
+    with tempfile.TemporaryDirectory(prefix=output.stem + "-", dir=output.parent) as directory:
+        directory = Path(directory)
+        for index, value in enumerate(values):
+            primary.UpdatePipeline(time=value)
+            if index == 0:
+                view.ResetCamera()
+            pv.Render(view)
+            pv.SaveScreenshot(str(directory / ("frame.%06d.png" % index)), view=view)
+        temporary = output.with_name(output.stem + ".tmp" + output.suffix)
+        subprocess.run([config["ffmpeg"], "-y", "-framerate", str(config["fps"]), "-i", str(directory / "frame.%06d.png"), "-c:v", "libx264", "-pix_fmt", "yuv420p", str(temporary)], check=True)
+        os.replace(temporary, output)
+
+
+def times(proxy):
+    proxy.UpdatePipelineInformation()
+    return list(getattr(proxy, "TimestepValues", []) or [0.0])
+
+
+if __name__ == "__main__":
+    main(sys.argv[1])