Merge pull request #9 from austin-hoover/examples

Add 4D SNS example

Author: austin-hoover

examples/sns_ring_4d_1d/README.md

@@ -0,0 +1,7 @@
+# 4D phase space tomography from 1D measurements in the SNS accumulator ring
+
+<img src="saved/fig_sns_diagram.png">
+
+<img src="saved/fig_diagram.png">
+
+https://doi.org/10.1103/PhysRevAccelBeams.27.122802

examples/sns_ring_4d_1d/plot_actions.py

@@ -0,0 +1,160 @@
+"""Plot 4D reconstruction in action space."""
+
+import argparse
+import os
+import pathlib
+
+import numpy as np
+import ment
+import matplotlib.pyplot as plt
+import scipy.ndimage
+
+from tools.cov import normalization_matrix
+from tools.utils import list_paths
+
+
+plt.style.use("tools/style.mplstyle")
+
+
+# Arguments
+# --------------------------------------------------------------------------------------
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--nsamp", type=int, default=1_000_000)
+parser.add_argument("--nbins", type=int, default=64)
+parser.add_argument("--jmax", type=float, default=85.0)
+parser.add_argument("--contours", type=int, default=8)
+parser.add_argument("--show", action="store_true")
+args = parser.parse_args()
+
+
+# Setup
+# --------------------------------------------------------------------------------------
+
+input_dir = "./outputs/train/"
+
+path = pathlib.Path(__file__)
+output_dir = os.path.join("outputs", path.stem)
+if not os.path.exists(output_dir):
+    os.makedirs(output_dir)
+
+
+# Load model
+# --------------------------------------------------------------------------------------
+
+checkpoints_folder = os.path.join(input_dir, "checkpoints")
+checkpoint_filenames = list_paths(os.path.join(input_dir, "checkpoints"), sort=True)
+checkpoint_filename = checkpoint_filenames[-1]
+
+model = ment.MENT(
+    ndim=4,
+    transforms=None,
+    projections=None,
+    prior=None,
+    sampler=None,
+)
+model.load(checkpoint_filenames[-1])
+model.sampler.noise = 1.0
+
+
+# Sample particles from distribution
+# --------------------------------------------------------------------------------------
+
+x = model.unnormalize(model.sample(args.nsamp))
+
+cov_matrix = np.cov(x.T)
+norm_matrix = normalization_matrix(cov_matrix, scale=False)
+
+z = np.matmul(x, norm_matrix.T)
+cov_matrix_n = np.cov(z.T)
+
+print("covariance matrix:")
+print(np.round(cov_matrix, 5))
+
+print("normalized covariance matrix:")
+print(np.round(cov_matrix_n, 5))
+
+eps1 = np.sqrt(cov_matrix_n[0, 0] * cov_matrix_n[1, 1])
+eps2 = np.sqrt(cov_matrix_n[2, 2] * cov_matrix_n[3, 3])
+eps_avg = np.sqrt(eps1 * eps2)  # maintains 4D emittance but equal mode amplitudes
+
+z_pred = np.copy(z)
+
+
+# Plot actions
+# --------------------------------------------------------------------------------------
+
+def make_joint_grid(figwidth=5.0, panel_width=0.33) -> tuple:
+    fig, axs = plt.subplots(
+        ncols=2,
+        nrows=2,
+        sharex="col",
+        sharey="row",
+        figsize=(figwidth, figwidth),
+        gridspec_kw=dict(
+            width_ratios=[1.0, panel_width],
+            height_ratios=[panel_width, 1.0],
+        ),
+    )
+    axs[0, 1].axis("off")
+    return fig, axs
+
+
+def plot_hist(values: np.ndarray, edges: list[np.ndarray], contours: int = 0, **plot_kws) -> tuple:
+    fig, axs = make_joint_grid()
+
+    ax = axs[1, 0]
+    ax.pcolormesh(
+        edges[0],
+        edges[1],
+        values.T,
+        linewidth=0.0,
+        rasterized=True,
+        shading="auto",
+    )
+
+    if contours:
+        coords = [0.5 * (e[:-1] + e[1:]) for e in edges]
+        axs[1, 0].contour(
+            coords[0],
+            coords[1],
+            scipy.ndimage.gaussian_filter(values, 1.0).T,
+            levels=np.linspace(0.01, 1.0, contours),
+            colors="white",
+            linewidths=0.80,
+            alpha=0.15,
+        )
+
+    ax = axs[0, 0]
+    proj_edges = edges[0]
+    proj_values = np.sum(values, axis=1)
+    proj_values = proj_values / np.max(proj_values)
+    ax.stairs(proj_values, proj_edges, lw=1.5, color="black")
+    ax.set_ylim(0.0, 1.25)
+
+    ax = axs[1, 1]
+    proj_edges = edges[1]
+    proj_values = np.sum(values, axis=0)
+    proj_values = proj_values / np.max(proj_values)
+    ax.stairs(proj_values, proj_edges, orientation="horizontal", lw=1.5, color="black")
+    ax.set_xlim(0.0, 1.25)
+
+    return fig, axs
+
+
+z = z_pred
+j1 = np.sum(np.square(z[:, (0, 1)]), axis=1)
+j2 = np.sum(np.square(z[:, (2, 3)]), axis=1)
+
+edges = np.linspace(0.0, args.jmax, args.nbins + 1)
+edges = [edges, edges]
+values, _, _ = np.histogram2d(j1, j2, bins=edges)
+values = values / np.max(values)
+
+fig, axs = plot_hist(values, edges, contours=args.contours)
+axs[1, 0].set_xlabel(r"$J_1$")
+axs[1, 0].set_ylabel(r"$J_2$")
+
+filename = f"fig_action.pdf"
+filename = os.path.join(output_dir, filename)
+plt.savefig(filename)

examples/sns_ring_4d_1d/plot_dist.py

@@ -0,0 +1,173 @@
+"""Plot 4D reconstruction results."""
+import argparse
+import os
+import pathlib
+
+import numpy as np
+import ment
+import matplotlib.pyplot as plt
+
+# local
+from tools.utils import list_paths
+
+plt.style.use("tools/style.mplstyle")
+
+
+# Arguments
+# --------------------------------------------------------------------------------------
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--nsamp", type=int, default=1_000_000)
+args = parser.parse_args()
+
+
+# Setup
+# --------------------------------------------------------------------------------------
+
+input_dir = "./outputs/train/"
+
+path = pathlib.Path(__file__)
+output_dir = os.path.join("outputs", path.stem)
+if not os.path.exists(output_dir):
+    os.makedirs(output_dir)
+
+
+# Load model
+# --------------------------------------------------------------------------------------
+
+checkpoints_folder = os.path.join(input_dir, "checkpoints")
+checkpoint_filenames = list_paths(os.path.join(input_dir, "checkpoints"), sort=True)
+checkpoint_filename = checkpoint_filenames[-1]
+
+model = ment.MENT(
+    ndim=4,
+    transforms=None,
+    projections=None,
+    prior=None,
+    sampler=None,
+)
+model.load(checkpoint_filenames[-1])
+model.sampler.noise = 1.0
+
+
+# Sample particles from distribution
+# --------------------------------------------------------------------------------------
+
+x_pred = model.unnormalize(model.sample(args.nsamp))
+
+
+# Simulate data
+# --------------------------------------------------------------------------------------
+
+projections_meas = model.projections
+projections_pred = ment.simulate(x_pred, model.transforms, model.diagnostics)
+
+
+# Plot measured vs. simulated projections.
+# --------------------------------------------------------------------------------------
+
+fig, axs = plt.subplots(
+    nrows=6,
+    ncols=6,
+    figsize=(6.0, 4.5),
+    sharex=True,
+    sharey=True,
+    constrained_layout=True,
+)
+
+index = 0
+for i in range(len(projections_meas)):
+    for j in range(len(projections_meas[i])):
+        proj_meas = projections_meas[i][j].copy()
+        proj_pred = projections_pred[i][j].copy()
+        scale = np.max(proj_meas.values)
+
+        ax = axs.flat[index]
+        ax.plot(
+            proj_pred.coords, 
+            proj_pred.values / scale, 
+            label="pred",
+            color="red", 
+            alpha=0.3,
+        )
+        ax.plot(
+            proj_meas.coords, 
+            proj_meas.values / scale, 
+            label="meas",
+            color="black",
+            lw=0,
+            marker=".",
+            ms=1.0,
+        )
+        ax.annotate(
+            "{:02.0f}".format(index // 3),
+            xy=(0.03, 0.96),
+            xycoords="axes fraction",
+            horizontalalignment="left",
+            verticalalignment="top",
+        )
+        index += 1
+
+for ax in axs.flat:
+    ax.set_ylim(-0.05, 1.15)
+    ax.set_yticks([])
+    ax.set_xticks([-35.0, 0.0, 35.0])
+
+axs[-1, 0].set_xlabel(r"$x$ (mm)")
+axs[-1, 1].set_xlabel(r"$y$ (mm)")
+axs[-1, 2].set_xlabel(r"$u$ (mm)")
+axs[-1, 3].set_xlabel(r"$x$ (mm)")
+axs[-1, 4].set_xlabel(r"$y$ (mm)")
+axs[-1, 5].set_xlabel(r"$u$ (mm)")
+
+filename = "fig_profiles.pdf"
+filename = os.path.join(output_dir, filename)
+plt.savefig(filename, dpi=250)
+plt.close()
+
+
+# Plot 2D projections of 4D distribution
+# --------------------------------------------------------------------------------------
+
+axes_proj = [(0, 1), (2, 3), (0, 2), (0, 3), (2, 1), (1, 3)]
+dims = [r"$x$", "$x'$", "$y$", "$y'$"]
+units = ["mm", "mrad", "mm", "mrad"]
+
+xmax = 3.5 * np.std(x_pred, axis=0)
+limits = list(zip(-xmax, xmax))
+
+fig, axs = plt.subplots(
+    ncols=3,
+    nrows=2,
+    figsize=(7.0, 4.0),
+    sharex=False,
+    sharey=False,
+    constrained_layout=True,
+)
+for j, axis in enumerate(axes_proj):
+    values, edges = np.histogramdd(
+        x_pred[:, axis], 
+        bins=64, 
+        range=[limits[k] for k in axis], 
+        density=True,
+    )
+    
+    ax = axs.flat[j]
+    ax.pcolormesh(
+        edges[0],
+        edges[1],
+        values.T,
+        rasterized=True,
+        edgecolor="None",
+        linewidth=0.0,
+    )
+
+for j, axis in enumerate(axes_proj):
+    ax = axs.flat[j]
+    ax.set_xlabel(f"{dims[axis[0]]} ({units[axis[0]]})")
+    ax.set_ylabel(f"{dims[axis[1]]} ({units[axis[1]]})")
+
+filename = f"fig_proj2d.pdf"
+filename = os.path.join(output_dir, filename)
+plt.savefig(filename, dpi=250)
+plt.close("all")