Viz¶

def render(self, plotter: pv.Plotter) -> None:  # pragma: no cover - needs pyvista (viz extra) + a VTK render context
    """Draw the precomputed field lines into ``plotter``."""
    assert self._actors is not None  # set in __post_init__

    options = self.options if self.options is not None else FieldLineOptions()
    draw_streamlines(self._actors, options, plotter)

def solid(self) -> dict[str, str]:
    """The color of every body not on the field colormap, by body: its own, or its share of the body colormap."""
    spread: list[str] = [style.body for style in self.body_styles if style.paint is BodyPaint.COLORMAP]
    shares: dict[str, str] = dict(zip(spread, spread_colors(self.body_colormap, len(spread)), strict=True)) if spread else {}
    colors: dict[str, str] = {}
    for style in self.body_styles:
        if style.paint is BodyPaint.COLORMAP:
            colors[style.body] = shares[style.body]
        elif style.paint is BodyPaint.SOLID and style.color is not None:
            colors[style.body] = style.color

    return colors

def body_opacity(self, body: str) -> float:
    """A body's drawn opacity: its own times the overall one."""
    own: float = next((style.opacity for style in self.body_styles if style.body == body), 1.0)
    return own * self.opacity

def unmatched(self, bodies: Iterable[str]) -> list[str]:
    """The ``hidden_bodies`` and ``body_styles`` names missing from ``bodies``, in spec order."""
    present: set[str] = set(bodies)
    named: list[str] = [*self.hidden_bodies, *(style.body for style in self.body_styles)]
    return [name for name in dict.fromkeys(named) if name not in present]

def fitted(self, fields: Sequence[FieldInfo]) -> DisplaySpec:
    """
    This spec snapped to what a run actually offers: an unset field takes the first available one and a mode the field disallows falls back to its first allowed mode.

    An unknown field name raises [`KeyError`][] naming the available arrays.
    """
    if not fields:
        raise KeyError("no displayable fields")

    by_name: dict[str, FieldInfo] = {info.name: info for info in fields}
    chosen: FieldInfo = fields[0] if not self.field else by_name[select_array(by_name, self.field)]

    allowed: tuple[DisplayMode, ...] = chosen.modes
    mode: DisplayMode = self.mode if self.mode in allowed else allowed[0]
    if chosen.name == self.field and mode is self.mode:
        return self

    return self.model_copy(update={"field": chosen.name, "mode": mode})

@classmethod
def from_run(cls, run: Run) -> FieldScene:  # pragma: no cover - needs pyvista (viz extra); proven by the gated integration test
    """Open a run bundle's VTU as a field scene, its bodies split and named from ``entities.json``."""
    import pyvista as pv

    if run.vtu_path is None:
        raise FileNotFoundError(f"{run.path}: no VTU output to display")

    bodies: Bodies = run.bodies
    if not len(bodies):
        raise FileNotFoundError(f"{run.path}: no named bodies ({BundleFile.ENTITIES} missing); re-solve the run so the bundle carries its entities")

    grid: pv.UnstructuredGrid = cast("pv.UnstructuredGrid", pv.read(str(run.vtu_path)))
    blocks: pv.MultiBlock = split_by_body(grid, bodies)
    fields: list[FieldInfo] = list_fields(grid)

    log.bind(run_id=run.path.name, bodies=len(blocks), fields=len(fields)).info("Field scene loaded.")
    return cls(grid=grid, blocks=blocks, fields=fields, model_dim=model_dim_of(grid), run_id=run.path.name, bodies=bodies)

def field_info(self, name: str) -> FieldInfo:  # pragma: no cover - needs pyvista (viz extra); proven by the gated integration test
    """The [`FieldInfo`][] for an array name, resolved the same way a field query resolves one."""
    by_name: dict[str, FieldInfo] = {info.name: info for info in self.fields}
    return by_name[select_array(by_name, name)]

def apply(self, spec: DisplaySpec) -> tuple[float, float]:  # pragma: no cover - needs pyvista (viz extra); proven by the gated integration test
    """Write the spec's reduction into every block's ``display_scalars`` array and return the color limits it produced."""
    import pyvista as pv

    fitted: DisplaySpec = spec.fitted(self.fields)
    info: FieldInfo = self.field_info(fitted.field)

    def raw_values(block: pv.DataSet) -> FloatArray:
        """One block's copy of the displayed array, zero-filled when the block never carried it."""
        arrays = block.point_data if info.association is Association.POINT else block.cell_data
        width: int = block.n_points if info.association is Association.POINT else block.n_cells
        if info.name not in arrays:
            return np.zeros((width, info.components), dtype=np.float64)

        return np.asarray(arrays[info.name], dtype=np.float64)

    names: list[str] = [str(key) for key in self.blocks.keys()]
    sources: list[FloatArray] = [raw_values(cast("pv.DataSet", self.blocks[name])) for name in names]
    reductions: list[AnyArray] = [reduce_field(values, fitted.mode) for values in sources]

    # Every body counts, each through its own range, whatever it is shown as: the scale belongs to the field, not the view.
    self.data_range = color_range(reductions, robust=False)
    clim: tuple[float, float] = fitted.clim if fitted.clim is not None else color_range(reductions, log_scale=fitted.log_scale, robust=fitted.robust)

    if fitted.mode is DisplayMode.POLAR_XY:
        cmap = pv.get_cmap_safe(fitted.colormap)  # pyright: ignore[reportUnknownMemberType]  # pyvista returns an untyped matplotlib Colormap

        def mapper(fractions: FloatArray) -> FloatArray:
            """RGBA rows for cyclic angle fractions, through the spec's colormap."""
            return np.asarray(cmap(fractions), dtype=np.float64)

        reductions = [polar_rgba(values, clim[0], clim[1], mapper) for values in sources]

    for name, scalars in zip(names, reductions, strict=True):
        block: pv.DataSet = cast("pv.DataSet", self.blocks[name])
        arrays = block.point_data if info.association is Association.POINT else block.cell_data
        arrays.set_scalars(scalars, name=SCALARS_ARRAY)  # pyright: ignore[reportUnknownMemberType]  # pyvista's set_scalars carries **kwargs: Unknown

    self.spec = fitted
    self.clim = clim
    return clim

def render(self, plotter: pv.Plotter) -> None:  # pragma: no cover - needs pyvista (viz extra) + a VTK render context; proven by the gated integration test
    """Draw the whole display into ``plotter``: the composite body actor, then everything [`update`][] owns; solid bodies an earlier scene left go first."""
    spec: DisplaySpec = self.spec
    _ = self.apply(spec)
    self._styled = {}
    self.drawn_lines = None

    actors = cast("Mapping[str, object]", plotter.renderer.actors)
    leftovers = [name for name in actors if name.startswith(SceneActor.STYLED.of(""))]
    for name in leftovers:
        _ = plotter.remove_actor(name)  # pyright: ignore[reportUnknownMemberType, reportArgumentType]  # pyvista removes by actor name, and its _Wrapped stub loses `self`

    self._actor = plotter.add_mesh(  # pyright: ignore[reportUnknownMemberType]
        self.blocks,
        scalars=SCALARS_ARRAY,
        show_edges=False,
        show_scalar_bar=False,
        name=str(SceneActor.BODIES),
    )
    # Translucent bodies sort per fragment only under depth peeling; without it a far body can paint over a near one.
    _ = cast("object", plotter.enable_depth_peeling())  # pyright: ignore[reportCallIssue, reportUnknownMemberType]  # pyvista's _Wrapped stub loses `self`
    self._redraw(plotter, spec)

def update(self, plotter: pv.Plotter, spec: DisplaySpec) -> None:  # pragma: no cover - needs pyvista (viz extra) + a VTK render context; proven by the gated integration test
    """Recolor an already-rendered scene in place — the cheap path a panel calls on every knob change; field lines are re-integrated only when their options changed."""
    if self._actor is None:
        self.render(plotter)
        return

    self._redraw(plotter, spec)

def shown_blocks(self, hidden: Iterable[str]) -> pv.MultiBlock:  # pragma: no cover - needs pyvista (viz extra)
    """The blocks of every body not in ``hidden``, under their names."""
    import pyvista as pv

    skipped: set[str] = set(hidden)
    shown: pv.MultiBlock = pv.MultiBlock()
    for name in (str(key) for key in self.blocks.keys()):
        if name not in skipped:
            shown.append(self.blocks[name], name=name)

    return shown

def integrated(self, options: FieldLineOptions) -> StreamlineActors | None:
    """The cached field lines when ``options`` only restyle them, None when they need integrating."""
    if self._lines is None or not integrates_like(self._lines[1], options):
        return None

    return self._lines[0]

def integrate_streamlines(self, options: FieldLineOptions) -> StreamlineActors:  # pragma: no cover - needs pyvista (viz extra); proven by the gated integration test
    """Seed and integrate the field lines through this scene's grid — the compute half, safe off the UI thread."""
    source, b_name = prepare_streamline_source(self.grid, options)
    actors: StreamlineActors = build_streamlines(source, b_name, options)
    return replace(actors, magnitudes=self.region_magnitudes(b_name, options.region_id))

def region_magnitudes(self, name: str, region_id: int | None) -> list[FloatArray]:  # pragma: no cover - needs pyvista (viz extra); proven by the gated integration test
    """The magnitudes of array ``name`` per body of a region (every body for None), the chunks a field-line range spans."""

    chunks: list[FloatArray] = []
    for block_name in (str(key) for key in self.blocks.keys()):
        body = self.bodies.get(block_name)
        if region_id is not None and (body is None or body.mesh_id != region_id):
            continue

        block: pv.DataSet = cast("pv.DataSet", self.blocks[block_name])
        arrays = block.cell_data if name in block.cell_data else block.point_data
        if name in arrays:
            chunks.append(reduce_field(np.asarray(arrays[name]), DisplayMode.MAGNITUDE))

    return chunks

def show_streamlines(
    self, plotter: pv.Plotter, actors: StreamlineActors, options: FieldLineOptions
) -> None:  # pragma: no cover - needs pyvista (viz extra) + a VTK render context; proven by the gated integration test
    """Draw integrated field lines into ``plotter`` in the display's mapping, record their options as the drawn ones and keep the lines for a restyle.

    Args:
        plotter: The plotter to draw into.
        actors: The integrated field lines.
        options: The options the lines were integrated with.
    """
    import pyvista as pv

    styled: FieldLineOptions = self.line_style(options, actors)
    draw_streamlines(actors, styled, plotter)
    self._lines = (actors, options)
    self.drawn_lines = styled
    self.spec = self.spec.model_copy(update={"streamlines": options})

    # The bar holds the mapper it was made with, and a redraw replaces the lines' actor, so it is made again every time.
    self._remove_lines_bar(plotter)
    drawn: object = plotter.renderer.actors.get(str(StreamlineActor.LINES))
    if not isinstance(drawn, pv.Actor):
        return

    self._lines_bar = self.lines_title(actors.b_name)
    _ = cast(
        "object",
        plotter.add_scalar_bar(  # pyright: ignore[reportCallIssue, reportUnknownMemberType]  # pyvista's _Wrapped stub loses `self`
            title=self._lines_bar,
            mapper=drawn.mapper,
            vertical=True,
            position_x=0.1,
            position_y=0.2,
            width=0.035,
            height=0.6,
        ),
    )

def lines_pinned(self, vectors: str) -> bool:
    """Whether field lines following ``vectors`` take the display's pinned range: it is pinned and shows that field's magnitude."""
    spec: DisplaySpec = self.spec
    return spec.clim is not None and spec.mode is DisplayMode.MAGNITUDE and spec.field == vectors

def lines_title(self, b_name: str) -> str:
    """The field lines' scalar bar title: the integrated field's name."""
    return f"{title(b_name)}{LINES_SUFFIX}"

def line_style(self, options: FieldLineOptions, actors: StreamlineActors) -> FieldLineOptions:
    """
    The field lines in the display's colormap, inversion, log scale and opacity map; a polar display leaves them as asked.

    Their range is the display's pinned one while it shows their field's magnitude, else the automatic range of that magnitude over their own region, whatever the bodies show.
    """
    spec: DisplaySpec = self.spec
    if spec.mode is DisplayMode.POLAR_XY:
        return options

    return options.model_copy(
        update={
            "colormap": spec.colormap,
            "invert": spec.invert,
            "log_scale": spec.log_scale,
            "opacity_map": spec.opacity_map,
            "clim": spec.clim if self.lines_pinned(actors.b_name) else color_range(actors.magnitudes, log_scale=spec.log_scale, robust=spec.robust),
        }
    )

def figure(self, spec: DisplaySpec) -> dict[str, object]:  # pragma: no cover - needs pyvista (viz extra); proven by the gated integration test
    """The plotly figure dict for the 2D path: one ``Mesh3d`` per body (hidden ones ``legendonly``) plus the outline, under an orthographic top-down camera.

    ``POLAR_XY`` falls back to the magnitude, a log scale plots ``log10`` with decade ticks, and ``opacity_map`` draws as ``opacity``.
    """
    fitted: DisplaySpec = spec.fitted(self.fields)
    plotted: DisplaySpec = fitted if fitted.mode is not DisplayMode.POLAR_XY else fitted.model_copy(update={"mode": DisplayMode.MAGNITUDE})
    info: FieldInfo = self.field_info(plotted.field)
    lo, hi = self.apply(plotted)

    colorscale: Colorscale = plotly_colorscale(sample_colormap(plotted.colormap), invert=plotted.invert)
    colorbar: TraceDict = {"title": {"text": info.display_name}, **(log_colorbar(lo, hi) if plotted.log_scale else {})}
    cmin, cmax = (math.log10(lo), math.log10(hi)) if plotted.log_scale else (lo, hi)
    hovertemplate: str = f"{info.display_name}: %{{customdata:.4g}}<extra>%{{fullData.name}}</extra>"

    solid: dict[str, str] = plotted.solid()
    traces: list[TraceDict] = []
    scale_shown: bool = False
    for name in (str(key) for key in self.blocks.keys()):
        block: pv.DataSet = cast("pv.DataSet", self.blocks[name])
        surface: pv.PolyData = cast("pv.PolyData", cast("pv.PolyData", block.extract_surface(algorithm=None)).triangulate())
        points, triangles = surface_triangles(surface)
        visible: bool = name not in plotted.hidden_bodies
        color: str | None = solid.get(name)
        if color is not None:
            traces.extend(mesh3d_traces(points, triangles, None, name=name, color=color, opacity=plotted.body_opacity(name), visible=visible))
            continue

        arrays = surface.point_data if info.association is Association.POINT else surface.cell_data
        values: FloatArray = np.asarray(arrays[SCALARS_ARRAY], dtype=np.float64)
        intensity: FloatArray = log_intensity(values, lo) if plotted.log_scale else values

        traces.extend(
            mesh3d_traces(
                points,
                triangles,
                intensity,
                name=name,
                opacity=plotted.body_opacity(name),
                intensitymode=info.association,
                colorscale=colorscale,
                cmin=cmin,
                cmax=cmax,
                showscale=visible and not scale_shown,
                hovertemplate=hovertemplate,
                customdata=values,
                colorbar=colorbar,
                visible=visible,
            )
        )
        scale_shown = scale_shown or visible

    edges: pv.PolyData | None = feature_outline(self.shown_blocks(plotted.hidden_bodies)) if plotted.outline else None
    if edges is not None:
        traces.extend(outline_traces(outline_segments(edges)))

    return {"data": traces, "layout": self._layout()}