add files

.gitattributes

@@ -0,0 +1,2 @@
+# Auto detect text files and perform LF normalization
+* text=auto

.gitignore

@@ -0,0 +1,13 @@
+*.pt
+*.pth
+#*.dset
+**/__pycache__/
+.idea/
+phase_space_reconstruction.egg-info/
+**/.ipynb_checkpoints/
+*.data
+*.h5
+*.dat
+*.sdds
+*.opal
+*.stat

environment.yml

@@ -0,0 +1,27 @@
+name: ps-reconstruction
+channels:
+  - conda-forge
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - python=3.10
+  - ipython
+  - ipykernel
+  - pytorch
+  - torchvision
+  - torchaudio
+  - pytorch-cuda=11.7
+  - numpy
+  - matplotlib
+  - scipy
+  - h5py
+  - scikit-image
+  - tqdm
+  - distgen
+  - pip
+  - openpmd-beamphysics
+  - pip:
+    - torchensemble==0.1.7
+    - pyro-ppl
+    - blitz-bayesian-pytorch

phase_space_reconstruction/analysis.py

@@ -0,0 +1,129 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from bmadx.bmad_torch.track_torch import Beam
+from bmadx.structures import Particle
+from pmd_beamphysics.particles import ParticleGroup
+
+# --------------------------------------------------------------------------
+
+
+def screen_stats(image, bins_x, bins_y):
+    """
+    Returns screen stats
+
+    Parameters
+    ----------
+    image: 2D array-like
+        screen image of size [n_x, n_y].
+
+    bins_x: 1D array-like
+        x axis bins physical locations of size [n_x]
+
+    bins_y: 2D array-like
+        x axis bins physical locations of size [n_y]
+
+    Returns
+    -------
+    dictionary with 'avg_x', 'avg_y', 'std_x' and 'std_y'.
+    """
+    proj_x = image.sum(axis=1)
+    proj_y = image.sum(axis=0)
+
+    # stats
+    avg_x = (bins_x * proj_x).sum() / proj_x.sum()
+    avg_y = (bins_y * proj_y).sum() / proj_y.sum()
+
+    std_x = (((bins_x * proj_x - avg_x) ** 2).sum() / proj_x.sum()) ** (1 / 2)
+    std_y = (((bins_y * proj_y - avg_y) ** 2).sum() / proj_y.sum()) ** (1 / 2)
+
+    return {"avg_x": avg_x, "avg_y": avg_y, "std_x": std_x, "std_y": std_y}
+
+
+# --------------------------------------------------------------------------
+
+
+def calculate_beam_matrix(beam_distribution: ParticleGroup, beam_fraction: float = 1.0):
+    fractional_beam = get_beam_fraction_openpmd_par(beam_distribution, beam_fraction)
+    return fractional_beam.cov("x", "py", "y", "py", "t", "pz")
+
+
+def get_beam_fraction_openpmd_par(beam_distribution: ParticleGroup, beam_fraction):
+    """get core of the beam according to 6D normalized beam coordinates"""
+    vnames = ["x", "px", "y", "py", "t", "pz"]
+    data = np.copy(np.stack([beam_distribution[name] for name in vnames]).T)
+    data[:, -2] = data[:, -2] - np.mean(data[:, -2])
+    data[:, -1] = data[:, -1] - np.mean(data[:, -1])
+    cov = np.cov(data.T)
+
+    # get inverse cholesky decomp
+    t_data = (np.linalg.inv(np.linalg.cholesky(cov)) @ data.T).T
+
+    J = np.linalg.norm(t_data, axis=1)
+    sort_idx = np.argsort(J)
+    frac_beam = beam_distribution[sort_idx][
+        : int(len(beam_distribution) * beam_fraction)
+    ]
+
+    return frac_beam
+
+
+def get_beam_fraction_bmadx_beam(beam_distribution: Beam, beam_fraction):
+    """get core of the beam according to 6D normalized beam coordinates"""
+    data = beam_distribution.data.detach().clone().numpy()
+    # data[:, -2] = data[:, -2] - np.mean(data[:, -2])
+    # data[:, -1] = data[:, -1] - np.mean(data[:, -1])
+    cov = np.cov(data.T)
+
+    # get inverse cholesky decomp
+    t_data = (np.linalg.inv(np.linalg.cholesky(cov)) @ data.T).T
+
+    J = np.linalg.norm(t_data, axis=1)
+    sort_idx = np.argsort(J)
+    frac_coords = data[sort_idx][: int(len(data) * beam_fraction)]
+    frac_beam = Beam(
+        torch.tensor(frac_coords),
+        p0c=beam_distribution.p0c,
+        s=beam_distribution.s,
+        mc2=beam_distribution.mc2,
+    )
+
+    return frac_beam
+
+
+def get_beam_fraction_bmadx_particle(beam_distribution: Particle, beam_fraction):
+    """get core of the beam according to 6D normalized beam coordinates"""
+    data = np.stack(beam_distribution[:6]).T
+    # data[:, -2] = data[:, -2] - np.mean(data[:, -2])
+    # data[:, -1] = data[:, -1] - np.mean(data[:, -1])
+    cov = np.cov(data.T)
+
+    # get inverse cholesky decomp
+    t_data = (np.linalg.inv(np.linalg.cholesky(cov)) @ data.T).T
+
+    J = np.linalg.norm(t_data, axis=1)
+    sort_idx = np.argsort(J)
+    frac_coords = data[sort_idx][: int(len(data) * beam_fraction)]
+    frac_particle = Particle(
+        *frac_coords.T,
+        p0c=beam_distribution.p0c,
+        s=beam_distribution.s,
+        mc2=beam_distribution.mc2
+    )
+
+    return frac_particle
+
+
+def get_beam_fraction_numpy_coords(beam_distribution, beam_fraction):
+    """get core of the beam according to 6D normalized beam coordinates"""
+    data = np.stack(beam_distribution[:6]).T
+    cov = np.cov(data.T)
+
+    # get inverse cholesky decomp
+    t_data = (np.linalg.inv(np.linalg.cholesky(cov)) @ data.T).T
+
+    J = np.linalg.norm(t_data, axis=1)
+    sort_idx = np.argsort(J)
+    frac_coords = data[sort_idx][: int(len(data) * beam_fraction)].T
+
+    return frac_coords

phase_space_reconstruction/beams/parameteric_models.py

@@ -0,0 +1,31 @@
+import torch
+from torch import nn
+
+
+class NNBeam(nn.Module):
+    def __init__(
+        self,
+        n_hidden,
+        width,
+        dropout=0.0,
+        activation=torch.nn.Tanh(),
+        output_scale=1e-2,
+    ):
+        super(NNBeam, self).__init__()
+
+        # create input layer
+        layers = [nn.Linear(6, width), activation]
+
+        # create hidden layers
+        for _ in range(n_hidden):
+            layers.append(nn.Linear(width, width))
+            layers.append(torch.nn.Dropout(dropout))
+            layers.append(activation)
+
+        layers.append(nn.Linear(width, 6))
+
+        self.stack = nn.Sequential(*layers)
+        self.register_buffer("output_scale", torch.tensor(output_scale))
+
+    def forward(self, X):
+        return self.stack(X) * self.output_scale

phase_space_reconstruction/diagnostics.py

@@ -0,0 +1,53 @@
+import torch
+from bmadx.bmad_torch.track_torch import Beam
+from torch.nn import Module
+
+from phase_space_reconstruction.histogram import histogram2d
+
+
+class ImageDiagnostic(Module):
+    def __init__(
+        self,
+        bins_x: torch.Tensor,
+        bins_y: torch.Tensor,
+        bandwidth: torch.Tensor,
+        x="x",
+        y="y",
+    ):
+        """
+        Parameters
+        ----------
+        bins_x : Tensor
+            A 'n' mesh of pixel centers that correspond to the physical diagnostic.
+
+        bins_y : Tensor
+            A 'm' mesh of pixel centers that correspond to the physical diagnostic.
+
+        bandwidth : Tensor
+            Bandwidth uses for kernel density estimation
+
+        x : str, optional
+            Beam attribute coorsponding to the horizontal image axis. Default: `x`
+
+        y : str, optional
+            Beam attribute coorsponding to the vertical image axis. Default: `y`
+        """
+
+        super(ImageDiagnostic, self).__init__()
+        self.x = x
+        self.y = y
+
+        self.register_buffer("bins_x", bins_x)
+        self.register_buffer("bins_y", bins_y)
+        self.register_buffer("bandwidth", bandwidth)
+
+    def forward(self, beam: Beam):
+        x_vals = getattr(beam, self.x)
+        y_vals = getattr(beam, self.y)
+        if not x_vals.shape == y_vals.shape:
+            raise ValueError("x,y coords must be the same shape")
+
+        if len(x_vals.shape) == 1:
+            raise ValueError("coords must be at least 2D")
+
+        return histogram2d(x_vals, y_vals, self.bins_x, self.bins_y, self.bandwidth)