Merge pull request #23 from austin-hoover/dev

Add pre-commit hooks

Author: austin-hoover

.pre-commit-config.yaml

@@ -0,0 +1,11 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+    -   id: check-yaml
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+-   repo: https://github.com/psf/black
+    rev: 22.10.0
+    hooks:
+    -   id: black

examples/cov/fit_cov_2d.py

@@ -78,13 +78,13 @@
 
 # Fit covariance matrix
 # --------------------------------------------------------------------------------------
-    
+
 fitter = ment.CholeskyCovFitter(
-    ndim=ndim, 
+    ndim=ndim,
     transforms=transforms,
     projections=projections,
     nsamp=args.nsamp,
-    bound=1.00e+06,
+    bound=1.00e06,
     verbose=True,
 )
 cov_matrix, fit_results = fitter.fit(iters=args.iters, method=args.method)
@@ -100,11 +100,11 @@
 projections_meas = unravel(fitter.projections)
 
 fig, axs = plt.subplots(
-    ncols=args.nmeas, 
-    figsize=(11.0, 1.0), 
+    ncols=args.nmeas,
+    figsize=(11.0, 1.0),
     sharey=True,
     sharex=True,
-    constrained_layout=True
+    constrained_layout=True,
 )
 for i, ax in enumerate(axs):
     values_pred = projections_pred[i].values
@@ -113,8 +113,3 @@
     ax.plot(values_meas / values_meas.max(), color="black", lw=0.0, marker=".", ms=2.0)
 plt.savefig(os.path.join(output_dir, "fig_results.png"), dpi=300)
 plt.close()
-
-
-
-
-

examples/cov/fit_cov_4d.py

@@ -85,11 +85,11 @@
 
 # Run optimizer
 fitter = ment.CholeskyCovFitter(
-    ndim=ndim, 
+    ndim=ndim,
     transforms=transforms,
     projections=projections,
     nsamp=args.nsamp,
-    bound=1.00e+02,
+    bound=1.00e02,
     verbose=True,
 )
 cov_matrix, fit_result = fitter.fit(method=args.method, iters=args.iters)
@@ -105,14 +105,14 @@ 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)
@@ -125,9 +125,9 @@ def rms_ellipse_params(cov_matrix: np.ndarray) -> tuple[float, float, float]:
 ncols = min(args.nmeas, 7)
 nrows = int(np.ceil(args.nmeas / ncols))
 fig, axs = plt.subplots(
-    ncols=ncols, 
-    nrows=nrows, 
-    figsize=(1.1 * ncols, 1.1 * nrows), 
+    ncols=ncols,
+    nrows=nrows,
+    figsize=(1.1 * ncols, 1.1 * nrows),
     constrained_layout=True,
     sharex=True,
     sharey=True,
@@ -140,14 +140,14 @@ def rms_ellipse_params(cov_matrix: np.ndarray) -> tuple[float, float, float]:
     for i, proj in enumerate([proj_true, proj_pred]):
         color = ["white", "red"][i]
         ls = ["-", "-"][i]
-        
+
         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))
+        ax.add_patch(
+            Ellipse(center, cx, cy, angle=angle, color=color, fill=False, ls=ls)
+        )
 plt.savefig(os.path.join(output_dir, "fig_results.png"), dpi=300)
 plt.close()
-
-

examples/cov/fit_cov_nd.py

@@ -59,6 +59,7 @@
 # Forward model
 # --------------------------------------------------------------------------------------
 
+
 class ProjectionTransform:
     def __init__(self, direction: np.ndarray) -> None:
         self.direction = direction
@@ -94,11 +95,11 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
 
 # Run optimizer
 fitter = ment.CholeskyCovFitter(
-    ndim=ndim, 
+    ndim=ndim,
     transforms=transforms,
     projections=projections,
     nsamp=args.nsamp,
-    bound=1.00e+02,
+    bound=1.00e02,
     verbose=args.verbose,
 )
 cov_matrix, fit_result = fitter.fit(method=args.method, iters=args.iters)
@@ -120,10 +121,10 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
 fig, axs = plt.subplots(
     ncols=ncols,
     nrows=nrows,
-    figsize=(ncols * 1.1, nrows * 1.1), 
+    figsize=(ncols * 1.1, nrows * 1.1),
     sharey=True,
     sharex=True,
-    constrained_layout=True
+    constrained_layout=True,
 )
 for i, ax in enumerate(axs.flat):
     values_pred = projections_pred[i].values
@@ -132,8 +133,3 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
     ax.plot(values_meas / values_meas.max(), color="black", lw=0.0, marker=".", ms=2.0)
 plt.savefig(os.path.join(output_dir, "fig_results.png"), dpi=300)
 plt.close()
-
-
-
-
-

examples/ct/train.py

@@ -30,10 +30,7 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument(
-    "--im", 
-    type=str, 
-    default="tree", 
-    choices=["shepp", "leaf", "tree", "brain"]
+    "--im", type=str, default="tree", choices=["shepp", "leaf", "tree", "brain"]
 )
 parser.add_argument("--im-blur", type=float, default=0.0)
 parser.add_argument("--im-pad", type=int, default=0)
@@ -126,7 +123,7 @@
 # import scipy.interpolate
 
 # interp = scipy.interpolate.RegularGridInterpolator(
-#     grid_coords, 
+#     grid_coords,
 #     grid_values_true,
 #     method="linear",
 #     fill_value=0.0,
@@ -150,7 +147,6 @@
 #     projections.append([projection])
 
 
-
 # Plot sinogram
 fig, ax = plt.subplots()
 ax.pcolormesh(sinogram)
@@ -228,7 +224,7 @@ def evaluate_model(model: ment.MENT) -> dict:
 
     discrepancy = np.mean(np.abs(sinogram_pred - sinogram_true))
 
-    # Absolute entropy    
+    # Absolute entropy
     p = values_pred
     q = np.ones(p.shape)
     q = q / np.sum(q) / cell_volume
@@ -397,7 +393,7 @@ def evaluate_model(model: ment.MENT) -> dict:
 for j, name in enumerate(results):
     for i, key in enumerate(["image", "sinogram"]):
         ax = axs[i, j]
-        image = results[name][key]        
+        image = results[name][key]
         ax.pcolormesh(image.T, vmin=vmin, vmax=vmax)
 for j, name in enumerate(results):
     axs[0, j].set_title(name.upper())

examples/ct/utils.py

@@ -39,20 +39,20 @@ def gen_image(key: str, res: int = None, blur: float = 0.0, pad: int = 0) -> Non
         if key == "tree":
             pad = max(pad, 25)
 
-    if pad:    
+    if pad:
         shape = image.shape
         new_shape = tuple(np.add(shape, pad * 2))
         new_image = np.zeros(new_shape)
         new_image[pad:-pad, pad:-pad] = image.copy()
         image = new_image.copy()
-        
+
     if res:
         shape = (res, res)
         image = skimage.transform.resize(image, shape, anti_aliasing=True)
 
     if blur:
         image = skimage.filters.gaussian(image, blur)
-        
+
     return image
 
 
@@ -64,7 +64,9 @@ def radon_transform(image: np.ndarray, angles: np.ndarray) -> np.ndarray:
     return sinogram
 
 
-def rec_sart(sinogram: np.ndarray, angles: np.ndarray, iterations: int = 1) -> np.ndarray:
+def rec_sart(
+    sinogram: np.ndarray, angles: np.ndarray, iterations: int = 1
+) -> np.ndarray:
     theta = -np.copy(np.degrees(angles))
     image = skimage.transform.iradon_sart(sinogram, theta=theta)
     for _ in range(iterations - 1):
@@ -73,8 +75,10 @@ def rec_sart(sinogram: np.ndarray, angles: np.ndarray, iterations: int = 1) -> n
     return image
 
 
-def rec_fbp(sinogram: np.ndarray, angles: np.ndarray, iterations: int = 1) -> np.ndarray:
+def rec_fbp(
+    sinogram: np.ndarray, angles: np.ndarray, iterations: int = 1
+) -> np.ndarray:
     theta = -np.copy(np.degrees(angles))
     image = skimage.transform.iradon(sinogram, theta=theta)
     image = image.T
-    return image
+    return image

examples/longitudinal/train.py

@@ -1,13 +1,13 @@
 """Reconstruct longitudinal phase space distribution from turn-by-turn projections.
 
-This script uses a PyORBIT [https://github.com/PyORBIT-Collaboration/PyORBIT3] lattice 
+This script uses a PyORBIT [https://github.com/PyORBIT-Collaboration/PyORBIT3] lattice
 model consisting of a harmonic RF cavity surrounded by two drifts. Things are a bit slow
 because we have to repeatedly convert between NumPy arrays and Bunch objects, but it works.
 
 Note that one MENT iteration requires simulating all projectionos. If projectiono k
 is measured after k turns, then we must first track the bunch 1 turn, then resample
 and track 2 turns, then resample and track 3 turns, etc. In total, we must track
-n * (n + 1) / 2 turns. For a significant number of turns, ART may be the better 
+n * (n + 1) / 2 turns. For a significant number of turns, ART may be the better
 option.
 """
 import os
@@ -51,6 +51,7 @@
 # Forward model
 # --------------------------------------------------------------------------------------
 
+
 def get_part_coords(bunch: Bunch, index: int) -> list[float]:
     x = bunch.x(index)
     y = bunch.y(index)
@@ -81,7 +82,9 @@ def get_bunch_coords(bunch: Bunch, axis: tuple[int, ...] = None) -> np.ndarray:
     return x
 
 
-def set_bunch_coords(bunch: Bunch, x: np.ndarray, axis: tuple[int, ...] = None) -> Bunch:
+def set_bunch_coords(
+    bunch: Bunch, x: np.ndarray, axis: tuple[int, ...] = None
+) -> Bunch:
     if axis is None:
         axis = tuple(range(6))
 
@@ -128,7 +131,7 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
         bunch = self.track_bunch()
         x_out = get_bunch_coords(bunch, axis=self.axis)
         return x_out
-    
+
 
 # Create accelerator lattice (drift, rf, drift)
 drift_node_1 = DriftTEAPOT()
@@ -142,7 +145,9 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
 rf_synchronous_de = 0.0
 rf_voltage = 300.0e-06
 rf_phase = 0.0
-rf_node = Harmonic_RFNode(z_to_phi, rf_synchronous_de, rf_hnum, rf_voltage, rf_phase, rf_length)
+rf_node = Harmonic_RFNode(
+    z_to_phi, rf_synchronous_de, rf_hnum, rf_voltage, rf_phase, rf_length
+)
 
 lattice = TEAPOT_Ring()
 lattice.addNode(drift_node_1)
@@ -178,11 +183,8 @@ def __call__(self, x: np.ndarray) -> np.ndarray:
     transform = ORBITTransform(lattice, bunch, nturns=nturns, axis=(4, 5))
     transforms.append(transform)
 
-limits = [
-    (-0.5 * lattice.getLength(), +0.5 * lattice.getLength()), 
-    (-0.030, 0.030)
-]
-    
+limits = [(-0.5 * lattice.getLength(), +0.5 * lattice.getLength()), (-0.030, 0.030)]
+
 # Create a list of histogram diagnostics for each transform.
 bin_edges = np.linspace(limits[0][0], limits[0][1], 100)
 diagnostics = []
@@ -235,19 +237,25 @@ def plot_model(model):
     projections_pred = ment.utils.unravel(projections_pred)
 
     fig, axs = plt.subplots(
-        ncols=nmeas, figsize=(11.0, 1.0), sharey=True, sharex=True, constrained_layout=True
+        ncols=nmeas,
+        figsize=(11.0, 1.0),
+        sharey=True,
+        sharex=True,
+        constrained_layout=True,
     )
     for i, ax in enumerate(axs):
         values_pred = projections_pred[i].values
         values_true = projections_true[i].values
         ax.plot(values_pred / values_true.max(), color="lightgray")
-        ax.plot(values_true / values_true.max(), color="black", lw=0.0, marker=".", ms=2.0)
+        ax.plot(
+            values_true / values_true.max(), color="black", lw=0.0, marker=".", ms=2.0
+        )
     return fig
 
 
 for epoch in range(4):
     print("epoch =", epoch)
-    
+
     if epoch > 0:
         model.gauss_seidel_step(learning_rate=0.90)
 
@@ -264,4 +272,3 @@ def plot_model(model):
     ax.hist2d(x[:, 0], x[:, 1], bins=100, range=limits)
 fig.savefig(os.path.join(output_dir, f"fig_dist_{epoch:02.0f}.png"))
 plt.close()
-

examples/nonlinear_ring_4d/README.md

@@ -5,5 +5,3 @@
 In this example we measure the $x-p_x$ and $y-p_y$ distribution after every 20 turns in a linear periodic lattice + nonlinear rotationally symmetric kick [1]. MENT finds an initial 4D phase space distribution consistent with these projections.
 
 [1] https://arxiv.org/abs/2405.05657
-
-