Merge pull request #11 from austin-hoover/dev

Organize examples

Author: austin-hoover

examples/fit_cov_2d.py renamed to examples/cov/fit_cov_2d.py

@@ -1,7 +1,10 @@
 """Fit 2D covariance matrix to 1D measurements."""
 import argparse
+import os
+import pathlib
 from typing import Callable
 from typing import Optional
+
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -13,7 +16,7 @@
 from ment.utils import rotation_matrix
 
 
-# Setup
+# Arguments
 # --------------------------------------------------------------------------------------
 
 parser = argparse.ArgumentParser()
@@ -25,6 +28,14 @@
 parser.add_argument("--method", type=str, default="differential_evolution")
 args = parser.parse_args()
 
+
+# Setup
+# --------------------------------------------------------------------------------------
+
+path = pathlib.Path(__file__)
+output_dir = os.path.join("outputs", path.stem)
+os.makedirs(output_dir, exist_ok=True)
+
 ndim = 2
 
 
@@ -84,7 +95,7 @@
 print(fit_results)
 
 # Plot results
-x = fitter.sample(10_000)
+x = fitter.sample(100_000)
 projections_pred = unravel(simulate(x, fitter.transforms, fitter.diagnostics))
 projections_meas = unravel(fitter.projections)
 
@@ -100,7 +111,9 @@
     values_meas = projections_meas[i].values
     ax.plot(values_pred / values_meas.max(), color="lightgray")
     ax.plot(values_meas / values_meas.max(), color="black", lw=0.0, marker=".", ms=2.0)
-plt.show()
+plt.savefig(os.path.join(output_dir, "fig_results.png"), dpi=300)
+plt.close()
+
 
 
 

examples/fit_cov_4d.py renamed to examples/cov/fit_cov_4d.py

@@ -1,11 +1,12 @@
 """Fit 4D covariance matrix to 2D measurements."""
 import argparse
+import os
+import pathlib
 from typing import Callable
 from typing import Optional
-import matplotlib.pyplot as plt
+
 import numpy as np
-import psdist as ps
-import scipy.optimize
+from matplotlib import pyplot as plt
 from matplotlib.patches import Ellipse
 
 import ment
@@ -16,7 +17,7 @@
 from ment.utils import rotation_matrix
 
 
-# Setup
+# Arguments
 # --------------------------------------------------------------------------------------
 
 parser = argparse.ArgumentParser()
@@ -31,6 +32,14 @@
 args = parser.parse_args()
 
 
+# Setup
+# --------------------------------------------------------------------------------------
+
+path = pathlib.Path(__file__)
+output_dir = os.path.join("outputs", path.stem)
+os.makedirs(output_dir, exist_ok=True)
+
+
 # Source distribution
 # --------------------------------------------------------------------------------------
 
@@ -83,7 +92,7 @@
     bound=1.00e+02,
     verbose=True,
 )
-cov_matrix, fit_result = fitter.fit(method=args.method)
+cov_matrix, fit_result = fitter.fit(method=args.method, iters=args.iters)
 
 
 # Print results
@@ -92,6 +101,23 @@
 
 
 # Plot results
+def rms_ellipse_params(cov_matrix: np.ndarray) -> tuple[float, float, float]:
+    sii = cov_matrix[0, 0]
+    sjj = cov_matrix[1, 1]
+    sij = cov_matrix[0, 1]
+    
+    angle = -0.5 * np.arctan2(2 * sij, sii - sjj)
+    
+    _sin = np.sin(angle)
+    _cos = np.cos(angle)
+    _sin2 = _sin**2
+    _cos2 = _cos**2
+    
+    c1 = np.sqrt(abs(sii * _cos2 + sjj * _sin2 - 2 * sij * _sin * _cos))
+    c2 = np.sqrt(abs(sii * _sin2 + sjj * _cos2 + 2 * sij * _sin * _cos))
+    return (c1, c2, angle)
+
+
 x = fitter.sample(100_000)
 projections_pred = unravel(simulate(x, fitter.transforms, fitter.diagnostics))
 projections_true = unravel(fitter.projections)
@@ -101,7 +127,7 @@
 fig, axs = plt.subplots(
     ncols=ncols, 
     nrows=nrows, 
-    figsize=(1.5 * ncols, 1.1 * nrows), 
+    figsize=(1.1 * ncols, 1.1 * nrows), 
     constrained_layout=True,
     sharex=True,
     sharey=True,
@@ -115,14 +141,13 @@
         color = ["white", "red"][i]
         ls = ["-", "-"][i]
 
-        cx, cy, angle = ps.cov.rms_ellipse_params(proj.cov(), axis=(0, 1))
+        cx, cy, angle = rms_ellipse_params(proj.cov())
         angle = -np.degrees(angle)
         center = (0.0, 0.0)
         cx *= 4.0
         cy *= 4.0
         ax.add_patch(Ellipse(center, cx, cy, angle=angle, color=color, fill=False, ls=ls))
-plt.show()
-
-
+plt.savefig(os.path.join(output_dir, "fig_results.png"), dpi=300)
+plt.close()
 
 

examples/fit_cov_nd.py renamed to examples/cov/fit_cov_nd.py

@@ -1,11 +1,12 @@
 """Fit ND covariance matrix to random 1D projections."""
 import argparse
+import os
+import pathlib
 from typing import Callable
 from typing import Optional
-import matplotlib.pyplot as plt
+
 import numpy as np
-import psdist as ps
-import scipy.optimize
+from matplotlib import pyplot as plt
 from matplotlib.patches import Ellipse
 
 import ment
@@ -16,7 +17,7 @@
 from ment.utils import rotation_matrix
 
 
-# Setup
+# Arguments
 # --------------------------------------------------------------------------------------
 
 parser = argparse.ArgumentParser()
@@ -34,6 +35,14 @@
 args = parser.parse_args()
 
 
+# Setup
+# --------------------------------------------------------------------------------------
+
+path = pathlib.Path(__file__)
+output_dir = os.path.join("outputs", path.stem)
+os.makedirs(output_dir, exist_ok=True)
+
+
 # Source distribution
 # --------------------------------------------------------------------------------------
 
@@ -111,7 +120,7 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
 fig, axs = plt.subplots(
     ncols=ncols,
     nrows=nrows,
-    figsize=(ncols * 11.0, nrows * 1.0), 
+    figsize=(ncols * 1.1, nrows * 1.1), 
     sharey=True,
     sharex=True,
     constrained_layout=True
@@ -121,7 +130,8 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
     values_meas = projections_meas[i].values
     ax.plot(values_pred / values_meas.max(), color="lightgray")
     ax.plot(values_meas / values_meas.max(), color="black", lw=0.0, marker=".", ms=2.0)
-plt.show()
+plt.savefig(os.path.join(output_dir, "fig_results.png"), dpi=300)
+plt.close()
 
 
 

examples/ct/plot_image.py

@@ -1,16 +1,17 @@
-import sys
+import os
 import matplotlib.pyplot as plt
-import numpy as np
-import skimage as ski
 
 from utils import gen_image
 
-name = None
+
+os.makedirs("outputs/plot_image", exist_ok=True)
+
+name = "tree"
 if len(sys.argv) > 1:
     name = sys.argv[1]
 
 im = gen_image(name)
 
 fig, ax = plt.subplots()
 ax.pcolormesh(im.T)
-plt.show()
+plt.savefig(f"outputs/plot_image/fig_{name}.png", dpi=300)

examples/rec_long.py renamed to examples/longitudinal/train.py

@@ -85,7 +85,6 @@ def set_bunch_coords(bunch: Bunch, x: np.ndarray, axis: tuple[int, ...] = None)
     if axis is None:
         axis = tuple(range(6))
 
-   # Resize
     size = x.shape[0]
     size_error = size - bunch.getSize()
     if size_error > 0:
@@ -137,7 +136,6 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
 drift_node_1.setLength(124.0)
 drift_node_2.setLength(124.0)
 
-
 z_to_phi = 2.0 * np.pi / 248.0
 rf_hnum = 1.0
 rf_length = 0.0
@@ -162,6 +160,13 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
     bunch.addParticle(0.0, 0.0, 0.0, 0.0, x_true[i, 0], x_true[i, 1])
 
 
+# Evolve forward a few turns; this will be our ground-truth distribution.
+for _ in range(250):
+    lattice.trackBunch(bunch)
+
+x_true = get_bunch_coords(bunch, axis=(4, 5))
+
+
 # Create transform functions
 turn_min = 0
 turn_max = 500
@@ -189,23 +194,19 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
 # Training data
 # --------------------------------------------------------------------------------------
 
-# Here we simulate the projections; in real life the projections would be measured.
 projections = ment.sim.simulate(x_true, transforms, diagnostics)
 
 
 # Reconstruction model
 # --------------------------------------------------------------------------------------
 
-# Define prior distribution for relative entropy calculation
 prior = ment.prior.GaussianPrior(ndim=2, scale=[200.0, 0.020])
 
-# Define particle sampler (if mode="sample")
 sampler = ment.samp.GridSampler(
     grid_limits=limits,
     grid_shape=(128, 128),
 )
 
-# Set up MENT model
 model = ment.MENT(
     ndim=ndim,
     transforms=transforms,
@@ -254,10 +255,11 @@ def plot_model(model):
     fig.savefig(os.path.join(output_dir, f"fig_proj_{epoch:02.0f}.png"))
     plt.close()
 
+
 # Plot final distribution
 x_pred = model.sample(x_true.shape[0])
 
-fig, axs = plt.subplots(ncols=2, constrained_layout=True)
+fig, axs = plt.subplots(ncols=2, figsize=(6, 3), constrained_layout=True)
 for ax, x in zip(axs, [x_pred, x_true]):
     ax.hist2d(x[:, 0], x[:, 1], bins=100, range=limits)
 fig.savefig(os.path.join(output_dir, f"fig_dist_{epoch:02.0f}.png"))