Make a movie of the convergence of a VMEC++ run for W7-X.

examples/convergence_movie_make_runs.py

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH <[email protected]>
+#
+# SPDX-License-Identifier: MIT
+"""Run VMEC++ via the Python API and take snapshots along the run."""
+
+from pathlib import Path
+
+import numpy as np
+
+import vmecpp
+
+cache_folder = Path("/home/jons/results/vmec_w7x/movie_cache")
+Path.mkdir(cache_folder, parents=True, exist_ok=True)
+
+# NOTE: This resolves to src/vmecpp/cpp/vmecpp/test_data in the repo.
+TEST_DATA_DIR = (
+    Path(__file__).parent.parent / "src" / "vmecpp" / "cpp" / "vmecpp" / "test_data"
+)
+
+# input_file = TEST_DATA_DIR / "w7x.json"
+input_file = TEST_DATA_DIR / "w7x_generic_initial_guess.json"
+input = vmecpp.VmecInput.from_file(input_file)
+
+# adjust as needed - we don't vendor the mgrid file, since it is too large
+input.mgrid_file = "/home/jons/results/vmec_w7x/mgrid_w7x.nc"
+# input.mgrid_file = "/home/jons/results/vmec_w7x/mgrid_w7x_nv72.nc"
+
+input.return_outputs_even_if_not_converged = True
+
+# # higher-res for nicer plots
+# input.ntheta = 100
+# input.nzeta = 72
+
+maximum_iterations = 20000
+
+# step = 100
+step = 10
+
+verbose = False
+max_threads = 6
+
+saved_steps = []
+
+currently_allowed_num_iterations = 1
+while currently_allowed_num_iterations < maximum_iterations:
+    # only run up to given limit of number of iterations
+    input.niter_array[0] = currently_allowed_num_iterations
+
+    cpp_indata = input._to_cpp_vmecindatapywrapper()
+    # start all over again, because flow control flags are not saved (yet) for restarting
+    output = vmecpp._vmecpp.run(
+        cpp_indata,
+        max_threads=max_threads,
+        verbose=verbose,
+    )
+
+    # print convergence progress
+    print(
+        "% 5d | % .3e | % .3e | % .3e"
+        % (
+            currently_allowed_num_iterations,
+            output.wout.fsqr,
+            output.wout.fsqz,
+            output.wout.fsql,
+        )
+    )
+
+    # save outputs for later plotting
+    output.save(cache_folder / f"vmecpp_w7x_{currently_allowed_num_iterations:04d}.h5")
+    saved_steps.append(currently_allowed_num_iterations)
+
+    # early exis this loop when VMEC is converged
+    if (
+        output is not None
+        and output.wout.fsqr < input.ftol_array[0]
+        and output.wout.fsqz < input.ftol_array[0]
+        and output.wout.fsql < input.ftol_array[0]
+    ):
+        print("converged after", output.wout.maximum_iterations, "iterations")
+        break
+
+    currently_allowed_num_iterations += step
+
+np.savetxt(cache_folder / "saved_steps.dat", saved_steps, fmt="%d")

examples/convergence_movie_plots.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH <[email protected]>
+#
+# SPDX-License-Identifier: MIT
+"""Plot snapshots of VMEC++ taken along the run."""
+
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+import vmecpp
+
+cache_folder = Path("/home/jons/results/vmec_w7x/movie_cache")
+if not Path.exists(cache_folder):
+    raise RuntimeError(
+        "Output cache folder "
+        + cache_folder
+        + " does not exist. Run convergence_movie_make_runs.py to generate it."
+    )
+
+plots_folder = cache_folder / "plots"
+Path.mkdir(plots_folder, parents=True, exist_ok=True)
+
+saved_steps = np.loadtxt(cache_folder / "saved_steps.dat", dtype=int)
+
+
+def flux_surfaces_rz(
+    oq: vmecpp._vmecpp.OutputQuantities, phi: float = 0.0, num_theta: int = 101
+) -> np.ndarray:
+    """Evaluate the flux-surface geometry at a fixed toroidal angle.
+
+    Returned shape: [ns][2: R, Z][num_theta]
+    """
+    rz = np.zeros([oq.wout.ns, 2, num_theta])
+
+    theta_grid = np.linspace(0.0, 2.0 * np.pi, num_theta, endpoint=True)
+    phi_grid = np.zeros([num_theta]) + phi
+    kernel = np.outer(oq.wout.xm, theta_grid) - np.outer(oq.wout.xn, phi_grid)
+    for js in range(oq.wout.ns):
+        rz[js, 0, :] = oq.wout.rmnc[js, :] @ np.cos(kernel)
+        rz[js, 1, :] = oq.wout.zmns[js, :] @ np.sin(kernel)
+
+    return rz
+
+
+# plt.figure(figsize=(4, 6))
+plt.figure()
+for saved_step in saved_steps:
+    # print(saved_step, "/", saved_steps[-1])
+
+    vmecpp_out_filename = cache_folder / f"vmecpp_w7x_{saved_step:04d}.h5"
+    if not Path.exists(vmecpp_out_filename):
+        raise RuntimeError(
+            "VMEC++ output file "
+            + str(vmecpp_out_filename)
+            + " does not exist. Run convergence_movie_make_runs.py to generate it."
+        )
+
+    oq = vmecpp._vmecpp.OutputQuantities.load(vmecpp_out_filename)
+
+    # phi = 0.0
+    # num_theta = 101
+    # rz = flux_surfaces_rz(oq=oq, phi=phi, num_theta=num_theta)
+    # for js in range(oq.wout.ns):
+    #     plt.plot(rz[js, 0, :], rz[js, 1, :], color="k", lw=0.5)
+    # plt.axis("equal")
+    # plt.grid(True)
+    # plt.savefig(plots_folder / f"flux_surfaces_{saved_step:04d}.png")
+
+    ns = oq.wout.ns
+    ntheta = oq.indata.ntheta // 2 + 1
+    nzeta = oq.indata.nzeta
+    jxb_gradp = np.reshape(oq.jxbout.jxb_gradp, [ns, ntheta * nzeta])
+
+    plt.clf()
+    plt.semilogy(np.abs(np.average(jxb_gradp, axis=-1)[3:-1]))
+    plt.title(f"{saved_step} / {saved_steps[-1]}")
+    plt.savefig(plots_folder / f"jxb_gradp_fsa_{saved_step:04d}.png")

examples/plot_torus.py

@@ -0,0 +1,141 @@
+# This script demonstrates how to plot a scalar field mapped onto a toroidal surface in 3D
+# using matplotlib in Python. The scalar field is defined to be independent of the Z-coordinate.
+
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib import cm
+from matplotlib.colors import Normalize
+
+import vmecpp
+
+cache_folder = Path("/home/jons/results/vmec_w7x/movie_cache")
+if not Path.exists(cache_folder):
+    raise RuntimeError(
+        "Output cache folder "
+        + cache_folder
+        + " does not exist. Run convergence_movie_make_runs.py to generate it."
+    )
+
+plots_folder = cache_folder / "plots"
+Path.mkdir(plots_folder, parents=True, exist_ok=True)
+
+saved_steps = np.loadtxt(cache_folder / "saved_steps.dat", dtype=int)
+
+saved_step = saved_steps[-1]
+
+vmecpp_out_filename = cache_folder / f"vmecpp_w7x_{saved_step:04d}.h5"
+if not Path.exists(vmecpp_out_filename):
+    raise RuntimeError(
+        "VMEC++ output file "
+        + str(vmecpp_out_filename)
+        + " does not exist. Run convergence_movie_make_runs.py to generate it."
+    )
+
+oq = vmecpp._vmecpp.OutputQuantities.load(vmecpp_out_filename)
+
+ns = oq.wout.ns
+nfp = oq.wout.nfp
+
+ntheta1 = 2 * (oq.indata.ntheta // 2)
+ntheta3 = ntheta1 // 2 + 1
+nzeta = oq.indata.nzeta
+
+theta_grid = np.linspace(0.0, 2.0 * np.pi, ntheta1 + 1, endpoint=True)
+phi_grid = np.linspace(0.0, 2.0 * np.pi, nfp * nzeta + 1, endpoint=True)
+
+theta, phi = np.meshgrid(theta_grid, phi_grid)
+
+kernel = np.outer(oq.wout.xm, theta) - np.outer(oq.wout.xn, phi)
+
+
+# Create a new figure for 3D plotting
+fig = plt.figure(figsize=(16, 9))
+ax = fig.add_subplot(projection="3d")
+
+cmap = cm.jet
+# cmap = cm.inferno
+
+dk = 18
+
+for i, js in enumerate([1, 10, 35, 48]):
+    k_start = i * dk
+
+    # compute corresponding flux surface geometry
+    r = np.zeros([ntheta3, nzeta])
+    z = np.zeros([ntheta3, nzeta])
+
+    r = (oq.wout.rmnc[js, :] @ np.cos(kernel)).reshape([nfp * nzeta + 1, ntheta1 + 1])
+    x = r * np.cos(phi)
+    y = r * np.sin(phi)
+    z = (oq.wout.zmns[js, :] @ np.sin(kernel)).reshape([nfp * nzeta + 1, ntheta1 + 1])
+
+    # MHD force residual
+    jxb_gradp = np.reshape(oq.jxbout.jxb_gradp, [ns, nzeta, ntheta3])[js, :, :]
+
+    # extend to full poloidal range
+    jxb_gradp_full = np.zeros([nfp * nzeta + 1, ntheta1 + 1])
+    jxb_gradp_full[:nzeta, :ntheta3] = jxb_gradp
+    jxb_gradp_full[:nzeta, ntheta3:] = jxb_gradp[:, 1:][::-1, ::-1]
+
+    # extend to full toroidal range
+    for kp in range(1, nfp):
+        jxb_gradp_full[kp * nzeta : (kp + 1) * nzeta, :] = jxb_gradp_full[:nzeta, :]
+
+    # ensure periodicity
+    jxb_gradp_full[-1, :] = jxb_gradp_full[0, :]
+
+    X = x[k_start:, :]
+    Y = y[k_start:, :]
+    Z = z[k_start:, :]
+    scalar_field = jxb_gradp_full[k_start:, :]
+
+    # Normalize the scalar field values to a [0, 1] range
+    # FIXME: normalize over all occuring values in plot
+    norm = Normalize(scalar_field.min(), scalar_field.max())
+
+    # Map the normalized scalar field through a colormap
+    mapped_colors = cmap(norm(scalar_field))
+
+    # Plot the torus surface and use the mapped colors
+    surf = ax.plot_surface(
+        X,
+        Y,
+        Z,
+        rstride=1,
+        cstride=1,
+        facecolors=mapped_colors,
+        linewidth=0,
+        antialiased=False,
+        alpha=1.0,  # ensure fully opaque
+        zsort="average",  # or 'min' or 'max'
+    )
+
+    # Create a ScalarMappable for the colorbar using the same normalization and colormap
+    m = cm.ScalarMappable(cmap=cmap, norm=norm)
+    m.set_array([])  # Dummy array for the colorbar
+
+# Attach the colorbar to the same Axes to avoid "Unable to determine Axes" errors
+fig.colorbar(m, ax=ax, shrink=0.6, aspect=10, label="Scalar Field Value")
+
+# Set the viewpoint by specifying elevation and azimuth angles
+# elev is the angle (in degrees) above the xy-plane
+# azim is the angle (in degrees) around the z-axis
+ax.view_init(elev=30, azim=45)
+
+# Force a 1:1:1 aspect ratio by setting the axis limits to a cube.
+max_range = np.array([X.max() - X.min(), Y.max() - Y.min(), Z.max() - Z.min()]).max()
+
+mid_x = (X.max() + X.min()) * 0.5
+mid_y = (Y.max() + Y.min()) * 0.5
+mid_z = (Z.max() + Z.min()) * 0.5
+
+ax.set_xlim(mid_x - max_range / 2, mid_x + max_range / 2)
+ax.set_ylim(mid_y - max_range / 2, mid_y + max_range / 2)
+ax.set_zlim(mid_z - max_range / 2, mid_z + max_range / 2)
+
+plt.tight_layout()
+
+# Display the final plot
+plt.show()