Merge pull request #68 from roussel-ryan/67-small-changes

small changes in `beams.py` and `train.py`

Author: jp-ga

gpsr/beams.py

@@ -8,7 +8,6 @@
 from torch.distributions import MultivariateNormal, Distribution
 
 from cheetah.particles import ParticleBeam
-from cheetah.utils.bmadx import bmad_to_cheetah_coords
 
 
 class BeamGenerator(torch.nn.Module, ABC):
@@ -55,14 +54,15 @@ def __init__(
         energy: float,
         base_dist: Distribution = None,
         transformer: NNTransform = None,
+        output_scale: float = 1e-2,
         n_dim: int = 6,
     ):
         super(NNParticleBeamGenerator, self).__init__()
         self.base_dist = base_dist or MultivariateNormal(
             torch.zeros(n_dim), torch.eye(n_dim)
         )
         self.transformer = transformer or NNTransform(
-            2, 20, output_scale=1e-2, phase_space_dim=n_dim
+            2, 20, output_scale=output_scale, phase_space_dim=n_dim
         )
         self.register_buffer("beam_energy", torch.tensor(energy))
         self.register_buffer("particle_charges", torch.tensor(1.0))
@@ -80,14 +80,16 @@ def forward(self) -> ParticleBeam:
 
         # create near zero coordinates into which we deposit the transformed beam
         # Note: these need to be near zero to maintain finite emittances
-        bmad_coords = torch.randn(len(transformed_beam), 6).to(transformed_beam) * 1e-7
-        bmad_coords[:, : transformed_beam.shape[1]] = transformed_beam
+        coords = torch.randn(len(transformed_beam), 6).to(transformed_beam) * 1e-7
+        coords[:, : transformed_beam.shape[1]] = transformed_beam
 
-        transformed_beam = bmad_to_cheetah_coords(
-            bmad_coords, self.beam_energy, torch.tensor(0.511e6)
+        coords = torch.cat(
+            (coords, torch.ones_like(coords[:, 0].unsqueeze(dim=-1))), dim=-1
         )
+
         return ParticleBeam(
-            *transformed_beam,
+            particles=coords,
+            energy=self.beam_energy,
             particle_charges=self.particle_charges,
             survival_probabilities=self.survival_probabilities,
         )
@@ -249,10 +251,14 @@ def forward(self) -> tuple[ParticleBeam, torch.Tensor]:
 
         entropy = -torch.mean(log_p - log_q)
 
-        particles, ref_energy = bmad_to_cheetah_coords(x, self.energy, self.mass)
-        particles[:, 4] *= -1.0  # [TO DO] why is sign wrong?
+        coords = torch.randn(len(x), 6).to(x) * 1e-7
+        coords[:, : x.shape[1]] = x
+        coords = torch.cat(
+            (coords, torch.ones_like(coords[:, 0].unsqueeze(dim=-1))), dim=-1
+        )
+
         beam = ParticleBeam(
-            particles, energy=ref_energy, particle_charges=self.particle_charges
+            particles=coords, energy=self.energy, particle_charges=self.particle_charges
         )
         return (beam, entropy)
 

gpsr/train.py

@@ -37,14 +37,17 @@ def training_step(self, batch, batch_idx):
                     f"prediction {i} shape {pred[i].shape} does not match target shape {y[i].shape}"
                 )
 
-        # normalize images
-        y_normalized = [normalize_images(y_ele) for y_ele in y]
-        pred_normalized = [normalize_images(pred_ele) for pred_ele in pred]
-
         # add up the loss functions from each prediction
         loss = 0.0
-        for y_ele, pred_ele in zip(y_normalized, pred_normalized):
-            loss += torch.sum(self.loss_func(y_ele, pred_ele))
+        for y_ele, pred_ele in zip(y, pred):
+            # normalize images
+            y_ele_norm = normalize_images(y_ele)
+            pred_ele_norm = normalize_images(pred_ele)
+            # add loss
+            loss += self.loss_func(y_ele_norm, pred_ele_norm)
+
+        # normalize loss by number of outputs
+        loss /= len(y)
 
         # log the loss function at the end of each epoch
         self.log("loss", loss, on_epoch=True)
@@ -57,7 +60,7 @@ def configure_optimizers(self):
 
 
 def train_gpsr(
-    model,
+    gpsr_model: GPSR,
     train_dataloader,
     n_epochs: int = 100,
     lr: float = 1e-3,
@@ -72,7 +75,7 @@ def train_gpsr(
 
     Arguments
     ---------
-    model: GPSRModel
+    gpsr_model: GPSR
         GPSR model to be trained.
     train_dataloader: DataLoader
         DataLoader for the training data.
@@ -88,8 +91,8 @@ def train_gpsr(
 
     Returns
     -------
-    model: GPSRModel
-        Trained GPSR model.
+    lit_gpsr_model: LitGPSR
+        Trained LitGPSR model.
 
     """
 
@@ -112,13 +115,13 @@ def train_gpsr(
         **kwargs,
     )
 
-    gpsr_model = LitGPSR(model, lr, loss_func=loss_func)
+    lit_gpsr_model = LitGPSR(gpsr_model, lr, loss_func=loss_func)
     trainer.fit(
-        gpsr_model,
+        lit_gpsr_model,
         train_dataloader,
     )
 
-    return model
+    return lit_gpsr_model
 
 
 class EntropyLitGPSR(L.LightningModule, ABC):