SPH in a Beltrami force field

"""Pressureless SPH markers orbiting in a stationary Beltrami force field.

The two-dimensional velocity

    u = (-cos(pi x) sin(pi y), sin(pi x) cos(pi y))

is a stationary solution of the pressureless momentum equation when the external
potential is p = (sin(pi x)^2 + sin(pi y)^2) / 2, because (u . grad)u = -grad p.
Markers initialized with this velocity should therefore keep following the same
Eulerian velocity field. Their particle Hamiltonian |v|^2 / 2 + p is conserved as
well. The animation shows the marker motion on stream-function contours while the
second panel reports both numerical errors.

Adapted from Struphy's ``tutorial_beltrami_sph.ipynb`` tutorial.

Install dependencies and compile the kernels (Python 3.10 or newer):

    pip install "struphy[pproc]>=3.4.0"
    pip install "plasma-plots[plotly]>=0.1.1"
    struphy compile

PNG exports require Chrome or Chromium. If Chrome is not installed, run:

    kaleido_get_chrome

Save this file and run it from the directory where you want the output:

    python beltrami-sph.py

Figures are saved in the current directory; add --show to display them before saving.
"""

import argparse

import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots

from struphy import (
    BoundaryParameters,
    DerhamOptions,
    EnvironmentOptions,
    KernelDensityPlot,
    LoadingParameters,
    SavingParameters,
    Simulation,
    SortingParameters,
    Time,
    WeightsParameters,
    domains,
    equils,
    grids,
)
from struphy.models import PressureLessSPH
from struphy.ode.utils import ButcherTableau
from plasma_plots import save_figure

box_min = -0.5
box_max = 0.5
boxes = 16
markers_per_box = 4
density_points = 65


def velocity(x, y, z):
    """Stationary cellular-flow velocity in physical coordinates."""
    del z
    return (
        -np.cos(np.pi * x) * np.sin(np.pi * y),
        np.sin(np.pi * x) * np.cos(np.pi * y),
        np.zeros_like(x),
    )


def potential(x, y, z):
    """Potential whose negative gradient accelerates the stationary flow."""
    del z
    return 0.5 * (np.sin(np.pi * x) ** 2 + np.sin(np.pi * y) ** 2)


def density(x, y, z):
    """Uniform initial density."""
    del y, z
    return np.ones_like(x)


def create_simulation() -> Simulation:
    beltrami_flow = equils.GenericCartesianFluidEquilibrium(
        u_xyz=velocity,
        p_xyz=potential,
        n_xyz=density,
    )

    model = PressureLessSPH(epsilon=1.0)
    model.propagators.push_eta.options = model.propagators.push_eta.Options(
        butcher=ButcherTableau(algo="forward_euler"),
    )
    model.propagators.push_v.potential = beltrami_flow.p0
    model.cold_fluid.set_markers(
        loading_params=LoadingParameters(ppb=markers_per_box, loading="tesselation"),
        weights_params=WeightsParameters(),
        boundary_params=BoundaryParameters(bc=("reflect", "reflect", "periodic")),
        sorting_params=SortingParameters(
            boxes_per_dim=(boxes, boxes, 1),
            dims_mask=(True, True, False),
        ),
        saving_params=SavingParameters(
            n_markers=1.0,
            kernel_density_plots=(
                KernelDensityPlot(
                    pts_e1=density_points,
                    pts_e2=density_points,
                    pts_e3=1,
                ),
            ),
        ),
        bufsize=0.5,
    )
    model.cold_fluid.var.add_background(beltrami_flow)

    domain = domains.Cuboid(
        l1=box_min,
        r1=box_max,
        l2=box_min,
        r2=box_max,
        l3=0.0,
        r3=1.0,
    )
    env = EnvironmentOptions(
        out_folders="struphy_gallery_runs",
        sim_folder="beltrami_sph",
    )
    sim = Simulation(
        model=model,
        name="SPH in a Beltrami force field",
        description=(
            "Pressureless SPH markers circulate through a stationary cellular flow driven by a prescribed "
            "Beltrami potential. The computed marker velocity is compared with the exact Eulerian velocity "
            "field, while conservation of each marker's kinetic-plus-potential energy checks the orbit integration."
            r" Initially :math:`n=1` and $$\mathbf{u}(x,y,0)=(-\cos(\pi x)\sin(\pi y),\sin(\pi x)\cos(\pi y),0).$$"
            r" The prescribed potential is :math:`V(x,y)=\tfrac12[\sin^2(\pi x)+\sin^2(\pi y)]` on :math:`[-0.5,0.5]^2`."
        ),
        env=env,
        time_opts=Time(dt=0.02, Tend=4.0, split_algo="Strang"),
        domain=domain,
        grid=grids.TensorProductGrid(num_elements=(32, 32, 1)),
        derham_opts=DerhamOptions(
            degree=(3, 3, 1),
            bcs=(("free", "free"), ("free", "free"), None),
        ),
    )
    return sim


def pproc(sim: Simulation, show: bool = False):
    output = sim.output
    output.pproc()

    orbits = output.orbits.cold_fluid.to_dataarray("quantity").transpose("t", "marker", "quantity")
    rho = output.evaluate("cold_fluid/view_0/n", eta3=0)
    times = np.asarray(orbits.t.values)
    x, y, z = (np.asarray(orbits.sel(quantity=name).values) for name in ("x", "y", "z"))
    v1, v2, v3 = (np.asarray(orbits.sel(quantity=name).values) for name in ("v1", "v2", "v3"))
    values = np.stack((x, y, z, v1, v2, v3), axis=-1)
    rho_values = np.asarray(rho.transpose("t", "eta2", "eta1").values)
    if not (np.isfinite(values).all() and np.isfinite(rho_values).all()):
        raise RuntimeError("Non-finite Beltrami result: refusing to publish the run")
    if not np.isclose(times[-1], sim.time_opts.Tend):
        raise RuntimeError(f"Incomplete Beltrami run: stopped at t={times[-1]:g}")
    if not np.allclose(times, rho.t.values):
        raise RuntimeError("Marker and density output times do not match")

    exact_v1, exact_v2, _ = velocity(x, y, z)
    initial_speed_rms = float(np.sqrt(np.mean(v1[0] ** 2 + v2[0] ** 2)))
    velocity_error = np.sqrt(np.mean((v1 - exact_v1) ** 2 + (v2 - exact_v2) ** 2, axis=1))
    velocity_error /= initial_speed_rms

    hamiltonian = 0.5 * (v1**2 + v2**2 + v3**2) + potential(x, y, z)
    hamiltonian_scale = np.maximum(np.abs(hamiltonian[0]), 1e-14)
    energy_error = np.max(np.abs(hamiltonian - hamiltonian[0]) / hamiltonian_scale, axis=1)
    max_velocity_error = float(np.max(velocity_error))
    max_energy_error = float(np.max(energy_error))
    if max_velocity_error > 0.1 or max_energy_error > 0.1:
        raise RuntimeError(
            "Beltrami verification failed: "
            f"{max_velocity_error=:.3g}, {max_energy_error=:.3g}"
        )
    print(
        f"maximum relative velocity RMS error {max_velocity_error:.3e}; "
        f"maximum per-marker Hamiltonian drift {max_energy_error:.3e}"
    )

    # The stream function psi has u = (d_y psi, -d_x psi); its contours are exact trajectories.
    mesh = np.linspace(box_min, box_max, 161)
    mesh_x, mesh_y = np.meshgrid(mesh, mesh)
    stream_function = np.cos(np.pi * mesh_x) * np.cos(np.pi * mesh_y) / np.pi
    marker_color = np.cos(np.pi * x[0]) * np.cos(np.pi * y[0]) / np.pi
    color_limit = float(np.max(np.abs(marker_color)))

    def marker_trace(index):
        return go.Scatter(
            x=x[index],
            y=y[index],
            mode="markers",
            name="SPH markers",
            marker={
                "size": 5,
                "color": marker_color,
                "colorscale": "Viridis",
                "cmin": 0.0,
                "cmax": color_limit,
                "line": {"width": 0},
                "colorbar": {
                    "title": "initial ψ",
                    "x": 1.03,
                    "xanchor": "left",
                    "y": 0.3,
                    "yanchor": "middle",
                    "len": 0.5,
                    "thickness": 14,
                },
            },
            hovertemplate="x=%{x:.3f}<br>y=%{y:.3f}<extra></extra>",
        )

    contour = go.Contour(
        x=mesh,
        y=mesh,
        z=stream_function,
        contours={"coloring": "none", "showlabels": False},
        line={"color": "rgba(40,55,65,.32)", "width": 1},
        showscale=False,
        hoverinfo="skip",
        name="exact streamlines",
    )

    def error_trace(index, values, name, color):
        """Error history up to one animation frame."""
        return go.Scatter(
            x=times[: index + 1],
            y=np.maximum(values[: index + 1], 1e-16),
            xaxis="x2",
            yaxis="y2",
            mode="lines",
            name=name,
            line={"color": color, "width": 2.5},
        )

    figure = make_subplots(
        rows=1,
        cols=2,
        column_widths=(0.58, 0.42),
        horizontal_spacing=0.14,
        subplot_titles=("Markers on exact streamlines", "Verification errors"),
    )
    figure.add_trace(contour, row=1, col=1)
    figure.add_trace(marker_trace(0), row=1, col=1)
    figure.add_trace(error_trace(0, velocity_error, "velocity RMS", "#00a884"), row=1, col=2)
    figure.add_trace(error_trace(0, energy_error, "Hamiltonian drift", "#9b51e0"), row=1, col=2)

    picks = np.unique(np.linspace(0, len(times) - 1, min(100, len(times)), dtype=int))
    figure.frames = [
        go.Frame(
            name=f"{times[index]:.2f}",
            data=[
                marker_trace(index),
                error_trace(index, velocity_error, "velocity RMS", "#00a884"),
                error_trace(index, energy_error, "Hamiltonian drift", "#9b51e0"),
            ],
            traces=[1, 2, 3],
        )
        for index in picks
    ]
    figure.update_layout(
        title="Pressureless SPH in a stationary Beltrami flow",
        template="plotly_white",
        margin={"l": 65, "r": 175, "t": 90, "b": 135},
        legend={
            "orientation": "v",
            "x": 1.02,
            "xanchor": "left",
            "y": 1,
            "yanchor": "top",
        },
        updatemenus=[
            {
                "type": "buttons",
                "showactive": False,
                "x": 0,
                "y": -0.2,
                "buttons": [
                    {
                        "label": "Play",
                        "method": "animate",
                        "args": [
                            None,
                            {
                                "frame": {"duration": 45, "redraw": True},
                                "transition": {"duration": 0},
                                "fromcurrent": True,
                            },
                        ],
                    }
                ],
            }
        ],
        sliders=[
            {
                "active": 0,
                "x": 0.12,
                "len": 0.88,
                "y": -0.12,
                "currentvalue": {"prefix": "t = "},
                "steps": [
                    {
                        "args": [
                            [frame.name],
                            {
                                "frame": {"duration": 0, "redraw": True},
                                "transition": {"duration": 0},
                                "mode": "immediate",
                            },
                        ],
                        "label": frame.name,
                        "method": "animate",
                    }
                    for frame in figure.frames
                ],
            }
        ],
    )
    figure.update_xaxes(title_text="x", range=[box_min, box_max], constrain="domain", row=1, col=1)
    figure.update_yaxes(
        title_text="y",
        range=[box_min, box_max],
        scaleanchor="x",
        scaleratio=1,
        row=1,
        col=1,
    )
    figure.update_xaxes(title_text="t", range=[0, float(times[-1])], row=1, col=2)
    figure.update_yaxes(title_text="relative error", type="log", range=[-7, -0.7], row=1, col=2)

    # The image shows the last frame, the end of the run (the contour stays as it is).
    save_figure(figure, "beltrami-sph", frame=-1, width=1100, height=680, show=show)

    # The kernel reconstruction shows the simulated mass density independently of the marker view.
    # The exact divergence-free Beltrami transport preserves the initially uniform rho = 1.
    density_x = box_min + (box_max - box_min) * np.asarray(rho.eta1.values)
    density_y = box_min + (box_max - box_min) * np.asarray(rho.eta2.values)
    density_min = float(np.min(rho_values))
    density_max = float(np.max(rho_values))
    max_density_deviation = float(np.max(np.abs(rho_values - 1.0)))

    def density_trace(index, *, colorbar=False):
        return go.Heatmap(
            x=density_x,
            y=density_y,
            z=rho_values[index],
            zmin=density_min,
            zmax=density_max,
            colorscale="Viridis",
            colorbar={"title": "ρ"} if colorbar else None,
            hovertemplate="x=%{x:.3f}<br>y=%{y:.3f}<br>ρ=%{z:.4f}<extra></extra>",
        )

    density_figure = go.Figure(data=[density_trace(0, colorbar=True)])
    density_figure.frames = [
        go.Frame(name=f"{times[index]:.2f}", data=[density_trace(index)]) for index in picks
    ]
    density_figure.update_layout(
        title="SPH density in the Beltrami flow",
        template="plotly_white",
        margin={"l": 70, "r": 60, "t": 80, "b": 130},
        updatemenus=[
            {
                "type": "buttons",
                "showactive": False,
                "x": 0,
                "y": -0.2,
                "buttons": [
                    {
                        "label": "Play",
                        "method": "animate",
                        "args": [
                            None,
                            {
                                "frame": {"duration": 45, "redraw": True},
                                "transition": {"duration": 0},
                                "fromcurrent": True,
                            },
                        ],
                    }
                ],
            }
        ],
        sliders=[
            {
                "active": 0,
                "x": 0.12,
                "len": 0.88,
                "y": -0.12,
                "currentvalue": {"prefix": "t = "},
                "steps": [
                    {
                        "args": [
                            [frame.name],
                            {
                                "frame": {"duration": 0, "redraw": True},
                                "transition": {"duration": 0},
                                "mode": "immediate",
                            },
                        ],
                        "label": frame.name,
                        "method": "animate",
                    }
                    for frame in density_figure.frames
                ],
            }
        ],
    )
    density_figure.update_xaxes(title_text="x", range=[box_min, box_max], constrain="domain")
    density_figure.update_yaxes(
        title_text="y",
        range=[box_min, box_max],
        scaleanchor="x",
        scaleratio=1,
    )

    # A signed density perturbation makes compression and rarefaction easier to distinguish than
    # the absolute-density movie. Clip only the colour scale at the 99th percentile; hover values
    # retain the full data and isolated extrema no longer wash out the rest of the field.
    compression = rho_values - 1.0
    compression_limit = float(np.percentile(np.abs(compression), 99.0))
    compression_figure = go.Figure(density_figure)
    compression_figure.data[0].update(
        z=compression[0],
        zmin=-compression_limit,
        zmax=compression_limit,
        colorscale="RdBu_r",
        colorbar={"title": "ρ − 1"},
        hovertemplate="x=%{x:.3f}<br>y=%{y:.3f}<br>ρ − 1=%{z:.4f}<extra></extra>",
    )
    for frame, index in zip(compression_figure.frames, picks, strict=True):
        frame.data[0].update(
            z=compression[index],
            zmin=-compression_limit,
            zmax=compression_limit,
            colorscale="RdBu_r",
        )
    compression_figure.update_layout(title="SPH compression and rarefaction: ρ − 1")

    # Recover the original rectangular marker lattice, then measure the signed area of every
    # quadrilateral throughout the run. Its ratio to the initial area is a discrete Lagrangian
    # Jacobian: an exactly incompressible flow keeps it equal to one.
    initial_x = np.round(x[0], decimals=12)
    initial_y = np.round(y[0], decimals=12)
    lattice_x = np.unique(initial_x)
    lattice_y = np.unique(initial_y)
    marker_count = x.shape[1]
    if lattice_x.size * lattice_y.size != marker_count:
        raise RuntimeError("Beltrami marker loading is not a rectangular tessellation")
    ix = np.searchsorted(lattice_x, initial_x)
    iy = np.searchsorted(lattice_y, initial_y)
    if np.unique(np.column_stack((ix, iy)), axis=0).shape[0] != marker_count:
        raise RuntimeError("Beltrami marker tessellation contains duplicate lattice sites")
    lattice_positions_x = np.empty((len(times), lattice_y.size, lattice_x.size))
    lattice_positions_y = np.empty_like(lattice_positions_x)
    lattice_positions_x[:, iy, ix] = x
    lattice_positions_y[:, iy, ix] = y

    def signed_cell_area(grid_x, grid_y):
        corners_x = (grid_x[:, :-1, :-1], grid_x[:, :-1, 1:], grid_x[:, 1:, 1:], grid_x[:, 1:, :-1])
        corners_y = (grid_y[:, :-1, :-1], grid_y[:, :-1, 1:], grid_y[:, 1:, 1:], grid_y[:, 1:, :-1])
        return 0.5 * sum(
            corner_x * corners_y[(index + 1) % 4] - corner_y * corners_x[(index + 1) % 4]
            for index, (corner_x, corner_y) in enumerate(zip(corners_x, corners_y, strict=True))
        )

    cell_area = signed_cell_area(lattice_positions_x, lattice_positions_y)
    if np.any(cell_area[0] <= 0):
        raise RuntimeError("Beltrami marker tessellation has non-positive initial cell areas")
    area_ratio = cell_area / cell_area[0]
    max_area_deviation = float(np.max(np.abs(area_ratio - 1.0)))
    area_color_limit = float(np.percentile(np.abs(area_ratio - 1.0), 99.0))
    cell_x = 0.5 * (lattice_x[:-1] + lattice_x[1:])
    cell_y = 0.5 * (lattice_y[:-1] + lattice_y[1:])

    area_figure = go.Figure(density_figure)
    area_figure.data[0].update(
        x=cell_x,
        y=cell_y,
        z=area_ratio[0],
        zmin=1.0 - area_color_limit,
        zmax=1.0 + area_color_limit,
        colorscale="RdBu_r",
        colorbar={"title": "A(t) / A(0)"},
        hovertemplate="x₀=%{x:.3f}<br>y₀=%{y:.3f}<br>A/A₀=%{z:.4f}<extra></extra>",
    )
    for frame, index in zip(area_figure.frames, picks, strict=True):
        frame.data[0].update(
            x=cell_x,
            y=cell_y,
            z=area_ratio[index],
            zmin=1.0 - area_color_limit,
            zmax=1.0 + area_color_limit,
            colorscale="RdBu_r",
        )
    area_figure.update_layout(title="Lagrangian marker-cell area deformation")
    area_figure.update_xaxes(title_text="initial x")
    area_figure.update_yaxes(title_text="initial y")

    # Six nested trajectories sample different streamline families without obscuring the field.
    trajectory_targets = np.column_stack((np.linspace(0.07, 0.43, 6), np.zeros(6)))
    trajectory_indices = np.array(
        [np.argmin((x[0] - target_x) ** 2 + (y[0] - target_y) ** 2) for target_x, target_y in trajectory_targets]
    )
    trajectory_colors = ("#440154", "#414487", "#2a788e", "#22a884", "#7ad151", "#fde725")
    trajectory_figure = go.Figure(
        go.Contour(
            x=mesh,
            y=mesh,
            z=stream_function,
            contours={"coloring": "none", "showlabels": False},
            line={"color": "rgba(40,55,65,.32)", "width": 1},
            showscale=False,
            hoverinfo="skip",
            showlegend=False,
        )
    )
    for marker, color in zip(trajectory_indices, trajectory_colors, strict=True):
        trajectory_figure.add_trace(
            go.Scatter(
                x=x[:, marker],
                y=y[:, marker],
                customdata=times,
                mode="lines",
                line={"color": color, "width": 2.5},
                showlegend=False,
                hovertemplate="t=%{customdata:.2f}<br>x=%{x:.3f}<br>y=%{y:.3f}<extra></extra>",
            )
        )
    trajectory_figure.add_trace(
        go.Scatter(
            x=x[0, trajectory_indices],
            y=y[0, trajectory_indices],
            mode="markers",
            name="start",
            marker={"color": "#00a884", "size": 9, "symbol": "circle", "line": {"color": "white", "width": 1}},
        )
    )
    trajectory_figure.add_trace(
        go.Scatter(
            x=x[-1, trajectory_indices],
            y=y[-1, trajectory_indices],
            mode="markers",
            name="tmax",
            marker={"color": "#d1495b", "size": 10, "symbol": "x"},
        )
    )
    trajectory_figure.update_layout(
        title="Selected SPH marker trajectories",
        template="plotly_white",
        margin={"l": 70, "r": 145, "t": 80, "b": 60},
        legend={"x": 1.02, "xanchor": "left", "y": 1, "yanchor": "top"},
    )
    trajectory_figure.update_xaxes(title_text="x", range=[box_min, box_max], constrain="domain")
    trajectory_figure.update_yaxes(
        title_text="y",
        range=[box_min, box_max],
        scaleanchor="x",
        scaleratio=1,
    )

    print(
        f"maximum density deviation from 1 {max_density_deviation:.3e}; "
        f"maximum marker-cell area deviation from 1 {max_area_deviation:.3e}"
    )

    save_figure(density_figure, "beltrami-sph-density", frame=-1, show=show)
    save_figure(compression_figure, "beltrami-sph-compression", frame=-1, show=show)
    save_figure(area_figure, "beltrami-sph-area-deformation", frame=-1, show=show)
    save_figure(trajectory_figure, "beltrami-sph-trajectories", show=show)


if __name__ == "__main__":
    argparser = argparse.ArgumentParser(description="Run the beltrami sph example.")
    argparser.add_argument(
        "--pproc-only",
        action="store_true",
        help="Run post-processing on an existing simulation instead of running a new one.",
    )
    argparser.add_argument("--show", action="store_true", help="Show the figures before saving them.")
    args = argparser.parse_args()

    simulation = create_simulation()
    if not args.pproc_only:
        simulation.run()
    pproc(simulation, show=args.show)