Add check for screens resolutions and observations shapes

docs/examples/example_reconstruction_4d.ipynb

@@ -56,7 +56,7 @@
     "print(\n",
     "    dset.parameters.shape,\n",
     "    dset.observations[0].shape,\n",
-    "    dset.screen,\n",
+    "    dset.screens[0],\n",
     ")\n",
     "dset.plot_data();"
    ]
@@ -109,10 +109,12 @@
    "outputs": [],
    "source": [
     "# print screen information\n",
-    "print(train_dset.screen)\n",
+    "print(train_dset.screens[0])\n",
     "# create diagnostic lattice\n",
     "p0c = 43.36e6  # reference momentum in eV/c\n",
-    "gpsr_lattice = GPSRQuadScanLattice(l_quad=0.1, l_drift=1.0, screen=train_dset.screen)"
+    "gpsr_lattice = GPSRQuadScanLattice(\n",
+    "    l_quad=0.1, l_drift=1.0, screen=train_dset.screens[0]\n",
+    ")"
    ]
   },
   {
@@ -210,7 +212,7 @@
    "outputs": [],
    "source": [
     "test_pred = gpsr_model(test_dset.parameters)[0].detach()\n",
-    "test_pred_dset = QuadScanDataset(test_dset.parameters, (test_pred,), train_dset.screen)"
+    "test_pred_dset = QuadScanDataset(test_dset.parameters, (test_pred,), train_dset.screens)"
    ]
   },
   {

docs/examples/example_reconstruction_4d_ensemble.ipynb

@@ -37,7 +37,7 @@
     "print(\n",
     "    dset.parameters.shape,\n",
     "    dset.observations[0].shape,\n",
-    "    dset.screen,\n",
+    "    dset.screens[0],\n",
     ")\n",
     "dset.plot_data();"
    ]
@@ -75,7 +75,9 @@
    "source": [
     "# create diagnostic lattice\n",
     "p0c = 43.36e6  # reference momentum in eV/c\n",
-    "gpsr_lattice = GPSRQuadScanLattice(l_quad=0.1, l_drift=1.0, screen=train_dset.screen)"
+    "gpsr_lattice = GPSRQuadScanLattice(\n",
+    "    l_quad=0.1, l_drift=1.0, screen=train_dset.screens[0]\n",
+    ")"
    ]
   },
   {
@@ -142,8 +144,8 @@
     "\n",
     "mean, nci = compute_mean_and_confidence_interval(test_pred.cpu())\n",
     "\n",
-    "mean_dset = QuadScanDataset(test_dset.parameters, (mean,), train_dset.screen)\n",
-    "nci_dset = QuadScanDataset(test_dset.parameters, (nci,), train_dset.screen)"
+    "mean_dset = QuadScanDataset(test_dset.parameters, (mean,), train_dset.screens)\n",
+    "nci_dset = QuadScanDataset(test_dset.parameters, (nci,), train_dset.screens)"
    ]
   },
   {
@@ -189,7 +191,7 @@
     "    1, len(test_dset.parameters), figsize=(20, 3), sharex=True, sharey=True\n",
     ")\n",
     "\n",
-    "xbins, ybins = nci_dset.screen.pixel_bin_centers\n",
+    "xbins, ybins = nci_dset.screens[0].pixel_bin_centers\n",
     "xx = torch.meshgrid(xbins * 1e3, ybins * 1e3, indexing=\"ij\")\n",
     "\n",
     "for i in range(len(test_dset.parameters)):\n",

gpsr/datasets.py

@@ -11,28 +11,38 @@
 
 class ObservableDataset(Dataset):
     """
-    Define a PyTorch dataset to be passed to a TorchDataloader object for training.
+    Define a PyTorch dataset to be passed to a TorchDataloader object for
+    training.
 
     Parameters
     ----------
     parameters : Tensor
-        Tensor of beamline parameters that correspond to data observations. Should
-        have a shape of (B x N), B is a batch dimension and N is equal to the number of
-        different parameters changed in the beamline. If parameters control which observation is being used,
-        the shape should be (B x M x N) where M is the number of different observations. See Notes for an example.
-    observations : Tuple[Tensor]
-        Tuple of tensors contaning observed data. Tuple length should be M,
-        or the number of different observations. Tensor shapes should be (B x D)
-        where B is a batch dimension shape that corresponds to the `parameter`
-        tensor, and D is an arbitrary shape corresponding to the dimensionality of
-        the observable. The images must follow the matrix convention,
-        where axis -2 is Y and axis -1 is X.
+        Tensor of beamline parameters that correspond to data observations.
+        Should have a shape of (B x N), B is a batch dimension and N is equal
+        to the number of different parameters changed in the beamline. If
+        parameters control which observation is being used, the shape should
+        be (B x M x N) where M is the number of different observations. See
+        Notes for an example.
+    observations : Tuple[Tensor, ...]
+        Tuple of tensors containing observed data. Tuple length should be M,
+        or the number of different observations. Tensor shapes should be
+        (B x D) where B is a batch dimension shape that corresponds to the
+        `parameter` tensor, and D is an arbitrary shape corresponding to the
+        dimensionality of the observable. The images must follow the matrix
+        convention, where axis -2 is Y and axis -1 is X.
+    screens: Tuple[Screen, ...]
+        Tuple of cheetah screen objects that corresponds to the observed
+        images. len(screens) must be equal to len(observations), and for each
+        i, screens[i].resolution[-2] == observations[i].shape[-1] and
+        screens[i].resolution[-1] == observations[i].shape[-2]. Note that the
+        screen resolution is given in (x,y) order, while the observation shape
+        is given in (y,x) order.
 
     Notes
     -----
-
-    Example: 4D phase space reconstruction with a quadrupole scan on one screen
-    using 5 quadrupole strengths. See below for the correct shapes of each input.
+    Example: 4D phase space reconstruction with a quadrupole scan on one
+    screen using 5 quadrupole strengths. See below for the correct shapes of
+    each input.
 
     >>> import torch
     >>> parameters = torch.rand((5, 1)) # B = (5), N = 1
@@ -41,12 +51,12 @@ class ObservableDataset(Dataset):
     >>> )
     >>> dataset = ObservableDataset(parameters, observations)
 
-
-    Example: 6D phase space reconstruction with 2 screens of various pixel sizes.
-    Beamline parameters include the 2 dipole strengths (corresponding to each screen),
-    2 transverse deflecting cavity strengths, and 5 quadrupole strenghs. See below for
-    the correct shapes of each input. Note: that because the TCAV and quadrupole parameters
-    do not change the measurement diagnostic, they are flattened into a single dimension.
+    Example: 6D phase space reconstruction with 2 screens of various pixel
+    sizes. Beamline parameters include the 2 dipole strengths (corresponding
+    to each screen), 2 transverse deflecting cavity strengths, and 5
+    quadrupole strengths. See below for the correct shapes of each input.
+    Note: that because the TCAV and quadrupole parameters do not change the
+    measurement diagnostic, they are flattened into a single dimension.
 
     >>> import torch
     >>> parameters = torch.rand((10, 2, 3)) # B = 10, M = 2, N = 3
@@ -55,20 +65,24 @@ class ObservableDataset(Dataset):
     >>>     torch.rand((10, 150, 150))
     >>> )
     >>> dataset = ObservableDataset(parameters, observations, n_observation_dims=1)
-
     """
 
     def __init__(
         self,
         parameters: Tensor,
         observations: Tuple[Tensor, ...],
+        screens: Tuple[Screen, ...],
     ):
         self.parameters = parameters
         self.observations = observations
+        self.screens = screens
 
         if not isinstance(observations, tuple):
             raise ValueError("observations must be passed as a tuple of tensors")
 
+        if not isinstance(screens, tuple):
+            raise ValueError("screens must be passed as a tuple of tensors")
+
         # validate the input shapes
         if len(parameters.shape) == 2 or len(parameters.shape) == 3:
             # the leading dimension must match the number of observations
@@ -83,12 +97,29 @@ def __init__(
             for ele in self.observations:
                 if not parameters.shape[0] == ele.shape[0]:
                     raise ValueError(
-                        "Batch dimension of parameters must match batch dimension of observations"
+                        "Batch dimension of parameters must match batch "
+                        "dimension of observations"
                     )
 
         else:
             raise ValueError("parameters must be a 2D or 3D tensor")
 
+        # validate screens correspond to observations
+        if len(screens) != len(observations):
+            raise ValueError("Number of screens must match number of observations")
+        else:
+            for i in range(len(screens)):
+                if (
+                    screens[i].resolution[-2] != observations[i].shape[-1]
+                    or screens[i].resolution[-1] != observations[i].shape[-2]
+                ):
+                    raise ValueError(
+                        f"Screen {i} resolution {screens[i].resolution} must match "
+                        f"observation {i} shape {observations[i].shape[-2:]}.\n"
+                        "Note that the screen resolution is given in (x,y) order, "
+                        "while the observation shape is given in (y,x) order."
+                    )
+
     def __len__(self):
         return self.parameters.shape[0]
 
@@ -111,31 +142,30 @@ def plot_data(self):
 
 
 class QuadScanDataset(ObservableDataset):
-    def __init__(self, parameters: Tensor, observations: Tuple[Tensor], screen: Screen):
-        """
-        Light wrapper dataset class for 4D phase space reconstructions with
-        quadrupole. Checks for correct sizes of parameters and observations
-        and provides a plotting utility.
-
-        Parameters
-        ----------
-        parameters : Tensor
-            Tensor of beamline parameters that correspond to data observations.
-            Should have a shape of (K x 1) where K is the number of quadrupole strengths.
-        observations : Tuple[Tensor]
-            Tuple contaning tensor of observed images, where the tensor shape
-            should be (K x bins x bins).The images must follow the matrix convention,
-            where axis -2 is Y and axis -1 is X.
-        screen: Screen
-            Cheetah screen object that corresponds to the observed images.
-            Note: screen resolution is in (x,y) order and image shape is in (y,x) order,
-            so screen.resolution[0] == observations[0].shape[-1] and
-            screen.resolution[1] == observations[0].shape[-2] must be true.
+    """
+    Light wrapper dataset class for 4D phase space reconstructions with
+    quadrupole. Checks for correct sizes of parameters and observations
+    and provides a plotting utility.
 
-        """
+    Parameters
+    ----------
+    parameters : Tensor
+        Tensor of beamline parameters that correspond to data observations.
+        Should have a shape of (K x 1) where K is the number of quadrupole
+        strengths.
+    observations : Tuple[Tensor]
+        Tuple containing tensor of observed images, where the tensor shape
+        should be (K x bins x bins). The images must follow the matrix
+        convention, where axis -2 is Y and axis -1 is X.
+    screens: Tuple[Screen]
+        Tuple containing cheetah screen object that corresponds to the
+        observed images.
+    """
 
-        super().__init__(parameters, observations)
-        self.screen = screen
+    def __init__(
+        self, parameters: Tensor, observations: Tuple[Tensor], screens: Tuple[Screen]
+    ):
+        super().__init__(parameters, observations, screens)
 
     def plot_data(
         self,
@@ -164,10 +194,11 @@ def plot_data(
                 fig = ax[tuple(np.zeros(len(ax.shape), dtype=int))].get_figure()
             else:
                 fig = ax.get_figure()
+        screen = self.screens[0]
+        images = self.observations[0]
 
-        px_bin_centers = self.screen.pixel_bin_centers
+        px_bin_centers = screen.pixel_bin_centers
         px_bin_centers = px_bin_centers[0] * 1e3, px_bin_centers[1] * 1e3
-        images = self.observations[0]
 
         for i in range(n_k):
             ax[i].pcolormesh(
@@ -209,43 +240,40 @@ def plot_data(
 
 
 class SixDReconstructionDataset(ObservableDataset):
+    """
+    Light wrapper dataset class for 6D phase space reconstructions with
+    quadrupole, dipole, and TDC. Checks for correct sizes of parameters
+    and observations.
+
+    Parameters
+    ----------
+    six_d_parameters : Tensor
+        Tensor of beamline parameters that correspond to data observations.
+        Shape should be (K x N x 2 x 3) where K is the number of quadrupole
+        strengths, N the number of TDC voltages, and 2 is the number of
+        dipole strengths. The last dimension should be ordered as
+        (quadrupole focusing strengths, TDC voltages, dipole strengths).
+
+    six_d_observations : Tuple[Tensor, Tensor]
+        Tuple of tensors containing observed images, where the tensor shapes
+        should be (K x N x n_bins x n_bins). Here, K is the number of
+        quadrupole strengths, and N is the number of TDC voltages. The first
+        entry should be dipole-off images, and the second entry should be
+        dipole-on images. The images must follow the matrix convention,
+        where axis -2 is Y and axis -1 is X.
+
+    screens: Tuple[Screen, Screen]
+        Tuple of cheetah screen objects that correspond to the observed
+        images. First element corresponds to dipole-off screen, second
+        element corresponds to dipole-on screen.
+    """
+
     def __init__(
         self,
         six_d_parameters: Tensor,
         six_d_observations: Tuple[Tensor, Tensor],
         screens: Tuple[Screen, Screen],
     ):
-        """
-        Light wrapper dataset class for 6D phase space reconstructions with
-        quadrupole, dipole, and TDC. Checks for correct sizes of parameters
-        and observations.
-
-        Parameters
-        ----------
-        six_d_parameters : Tensor
-            Tensor of beamline parameters that correspond to data observations.
-            Shape should be (K x N x 2 x 3) where K is the number of quadrupole
-            strengths, N the number of TDC voltages, and 2 is the number of
-            dipole strengths. The last dimension should be ordered as
-            (quadrupole focusing strengths, TDC voltages, dipole strengths).
-
-        six_d_observations : Tuple[Tensor, Tensor]
-            Tuple of tensors containing observed images, where the tensor shapes
-            should be (K x N x n_bins x n_bins). Here, K is the number of
-            quadrupole strengths, and N is the number of TDC voltages. The first
-            entry should be dipole-off images, and the second entry should be
-            dipole-on images. The images must follow the matrix convention,
-            where axis -2 is Y and axis -1 is X.
-
-        screens: Tuple[Screen, Screen]
-            Tuple of cheetah screen objects that corresponds to the observed images.
-
-        Notes
-        -----
-
-        Only squared images are supported for now.
-        """
-
         # keep unflattened parameters and observations for visualization
         self.six_d_parameters = six_d_parameters
         self.six_d_observations = six_d_observations
@@ -256,8 +284,7 @@ def __init__(
             [ele.clone().flatten(end_dim=-3) for ele in self.six_d_observations]
         )
 
-        super().__init__(parameters, observations)
-        self.screens = screens
+        super().__init__(parameters, observations, screens)
 
     def plot_data(
         self,
@@ -474,13 +501,13 @@ def generate_test_indices(all_k_ids, train_k_ids, test_k_ids=None):
         train_dataset = QuadScanDataset(
             parameters=dataset.parameters[train_k_ids],
             observations=tuple([dataset.observations[0][train_k_ids]]),
-            screen=dataset.screen,
+            screens=dataset.screens,
         )
 
         test_dataset = QuadScanDataset(
             parameters=dataset.parameters[test_k_ids_copy],
             observations=tuple([dataset.observations[0][test_k_ids_copy]]),
-            screen=dataset.screen,
+            screens=dataset.screens,
         )
 
     else: