Electromagnetic pulse on a distorted mesh

"""An electromagnetic pulse on a distorted mesh, with Struphy's Maxwell model.

The Colella mapping x = Lx (eta1 + alpha sin(2 pi eta1) sin(2 pi eta2)), y = Ly (eta2 + alpha sin(2 pi eta2) sin(2 pi eta1)) leaves the
boundary of the box [0, Lx] x [0, Ly] where it is but bends the mesh lines inside it. Struphy's finite element spaces are built on the
logical cube and pushed forward with this map, so Maxwell's equations are solved on a curved mesh. Because the physical domain is
still the periodic rectangle, the exact solution is known: a Gaussian pulse of E_z, with B = 0 at first, is a sum of cosine modes,
each oscillating at omega = c |k|, and it can be compared with the numerical field at every point of the distorted mesh. The
energy, which the structure-preserving scheme conserves for any mesh, is followed as well.

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 maxwell-curved-mesh.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 struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, grids
from struphy.initial.base import GenericPerturbation
from struphy.models import Maxwell
from plasma_plots import save_figure

lx, ly = 2.0, 3.0
distortion = 0.1
width = 0.25  # of the Gaussian pulse
max_mode = 6  # cosine modes per direction in the Fourier series of the Gaussian
x0, y0 = 0.5 * lx, 0.5 * ly

# The Gaussian pulse as a cosine series about its centre, E_z = sum c_lm cos(kx (x - x0)) cos(ky (y - y0)) without the mean (l, m) = (0, 0).
modes = [(ell, m) for ell in range(max_mode + 1) for m in range(max_mode + 1) if (ell, m) != (0, 0)]
kx = np.array([2 * np.pi * ell / lx for ell, _ in modes])
ky = np.array([2 * np.pi * m / ly for _, m in modes])
weights = np.array([(1 if ell == 0 else 2) * (1 if m == 0 else 2) for ell, m in modes])
coefficients = weights * 2 * np.pi * width**2 / (lx * ly) * np.exp(-0.5 * width**2 * (kx**2 + ky**2))
frequencies = np.hypot(kx, ky)


def exact_field(x, y, t):
    """E_z(x, y, t): every cosine mode oscillates at c |k|, since B is zero and dE/dt is zero at the start."""
    x, y = np.asarray(x)[..., None], np.asarray(y)[..., None]
    return np.sum(coefficients * np.cos(kx * (x - x0)) * np.cos(ky * (y - y0)) * np.cos(frequencies * t), axis=-1)


def create_simulation() -> Simulation:
    time_opts = Time(dt=0.01, Tend=3.0)

    model = Maxwell()
    model.propagators.maxwell.options = model.propagators.maxwell.Options(algo="implicit")
    model.em_fields.e_field.add_perturbation(
        GenericPerturbation(lambda x, y, z: exact_field(x, y, 0.0), given_in_basis="physical", comp=2)
    )

    domain = domains.Colella(Lx=lx, Ly=ly, alpha=distortion, Lz=1.0)
    grid = grids.TensorProductGrid(num_elements=(24, 36, 1))
    derham_opts = DerhamOptions(degree=(3, 3, 1))

    sim = Simulation(
        model=model,
        name="Electromagnetic pulse on a distorted mesh",
        description=(
            "A Gaussian pulse of the electric field spreads as a ring across a periodic box, on a mesh that the Colella "
            "mapping has bent. The structure-preserving Maxwell solver keeps the energy constant and follows the exact "
            "solution, which is a sum of cosine modes."
            r" The initial pulse is a truncated, zero-mean Gaussian Fourier expansion: $$E_z(x,y,0)=\sum_{\substack{0\le l,m\le6\\(l,m)\ne(0,0)}}a_{lm}\cos[k_l(x-1)]\cos[q_m(y-1.5)],$$"
            r" where :math:`k_l=\pi l`, :math:`q_m=2\pi m/3` and :math:`a_{lm}=(2-\delta_{l0})(2-\delta_{m0})(2\pi\sigma^2/6)e^{-\sigma^2(k_l^2+q_m^2)/2}`."
            r" The width is :math:`\sigma=0.25`, :math:`\delta_{ij}` is the Kronecker delta, and the magnetic field starts at zero."
        ),
        env=EnvironmentOptions(out_folders="struphy_gallery_runs", sim_folder="maxwell_curved_mesh"),
        time_opts=time_opts,
        domain=domain,
        grid=grid,
        derham_opts=derham_opts,
    )
    return sim


def pproc(sim: Simulation, show: bool = False):
    from plotly.subplots import make_subplots
    from scipy.interpolate import griddata

    output = sim.output
    output.pproc(physical=True)

    # (t, e1, e2), on the mesh points
    e_z = output.evaluate("em_fields/e_field_xyz").isel(component=2, eta3=0).transpose("t", "eta1", "eta2")
    times = e_z.t.values
    mesh_x, mesh_y = e_z.X.values, e_z.Y.values
    numeric = e_z.values
    exact = np.array([exact_field(mesh_x, mesh_y, t) for t in times])
    scale = float(np.abs(exact[0]).max())
    # The rms error over the mesh points at each time, relative to the rms of the initial pulse.
    error = (e_z.plasma.analysis.error(exact) / np.sqrt(np.mean(exact[0] ** 2))).values
    if not np.isfinite(error).all():
        raise RuntimeError("Non-finite field")
    total = output.scalars["total_energy"].values
    energy_drift = float(output.scalars["total_energy"].plasma.analysis.relative_error().max())
    print(f"Largest rms error of E_z, relative to the initial rms: {error.max():.3e}")
    print(f"Maximum relative drift of the total energy: {energy_drift:.2e}")

    # Resample the values on the distorted mesh onto a regular grid, for the heatmap.
    x_plot, y_plot = np.linspace(0, lx, 64), np.linspace(0, ly, 96)
    xx, yy = np.meshgrid(x_plot, y_plot)
    points = np.column_stack([mesh_x.ravel(), mesh_y.ravel()])

    def resample(values):
        # Rounded, to keep the frames of the page small.
        return np.round(griddata(points, values.ravel(), (xx, yy), method="linear") / scale, 3)

    picks = np.unique(np.linspace(0, len(times) - 1, 45, dtype=int))
    mesh_lines = [
        go.Scatter(x=mesh_x[i, :], y=mesh_y[i, :], mode="lines", line={"color": "rgba(0,0,0,0.35)", "width": 0.7},
                   showlegend=False, hoverinfo="skip")
        for i in range(0, mesh_x.shape[0], 3)
    ] + [
        go.Scatter(x=mesh_x[:, j], y=mesh_y[:, j], mode="lines", line={"color": "rgba(0,0,0,0.35)", "width": 0.7},
                   showlegend=False, hoverinfo="skip")
        for j in range(0, mesh_x.shape[1], 3)
    ]
    frames_numeric = {i: resample(numeric[i]) for i in picks}
    figure = go.Figure(
        data=[go.Heatmap(z=frames_numeric[picks[0]], x=x_plot, y=y_plot, colorscale="RdBu", zmid=0.0, zmin=-0.5, zmax=0.5,
                         colorbar={"title": "E_z / initial peak"})] + mesh_lines,
        frames=[go.Frame(name=f"{times[i]:.2f}", data=[go.Heatmap(z=frames_numeric[i], x=x_plot, y=y_plot, colorscale="RdBu",
                                                                 zmid=0.0, zmin=-0.5, zmax=0.5)], traces=[0])
                for i in picks],
    )
    figure.update_layout(
        title="An electromagnetic pulse on a distorted mesh: E_z", template="plotly_white", autosize=True,
        xaxis_title="x", yaxis_title="y", margin={"l": 70, "r": 30, "t": 80, "b": 130},
        updatemenus=[{"type": "buttons", "showactive": False, "x": 0, "y": -0.1,
                      "buttons": [{"label": "Play", "method": "animate", "args": [None, {"frame": {"duration": 60, "redraw": True}, "transition": {"duration": 0}, "fromcurrent": True}]}]}],
        sliders=[{"active": 0, "x": 0.12, "len": 0.88, "y": -0.02, "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(range=[0, lx], constrain="domain")
    figure.update_yaxes(range=[0, ly], scaleanchor="x")
    # The image shows the pulse at a fifth of the run, when the ring has spread; the page keeps the animation.
    still = go.Figure(data=figure.data, layout=figure.layout)
    still.data[0].z = resample(numeric[len(times) // 5])
    save_figure(figure, "maxwell-curved-mesh", width=750, height=1000, still=still, show=show)

    diagnostics = make_subplots(rows=2, cols=1, vertical_spacing=0.18,
                                subplot_titles=("Error of E_z against the exact solution", "Total energy"))
    diagnostics.add_scatter(x=times, y=np.maximum(error, 1e-12), mode="lines", showlegend=False,
                            line={"color": "#d62828", "width": 2}, row=1, col=1)
    diagnostics.add_scatter(x=times, y=total / total[0] - 1.0, mode="lines", showlegend=False,
                            line={"color": "#168aad", "width": 2}, row=2, col=1)
    diagnostics.update_layout(template="plotly_white", autosize=True, margin={"l": 80, "r": 30, "t": 80, "b": 60})
    diagnostics.update_yaxes(title_text="rms error / initial rms", type="log", row=1, col=1)
    diagnostics.update_yaxes(title_text="relative change", exponentformat="e", row=2, col=1)
    diagnostics.update_xaxes(title_text="t", row=2, col=1)
    save_figure(diagnostics, "maxwell-curved-mesh-diagnostics", show=show)


if __name__ == "__main__":
    argparser = argparse.ArgumentParser(description="Run the maxwell curved mesh 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)