From a9199d6f2d37789cf649c4cbd9b5bc33fafd6b13 Mon Sep 17 00:00:00 2001 From: Lewis Anthony Boulton Date: Fri, 30 May 2025 16:51:00 +0200 Subject: [PATCH 1/8] Add class to manage running a GPSR training and loading the results --- gpsr/run.py | 234 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 234 insertions(+) create mode 100644 gpsr/run.py diff --git a/gpsr/run.py b/gpsr/run.py new file mode 100644 index 0000000..6de8c15 --- /dev/null +++ b/gpsr/run.py @@ -0,0 +1,234 @@ +import torch +from pprint import pprint +import yaml + +import lightning as L +from lightning.pytorch.loggers import CSVLogger +from lightning.pytorch.callbacks import ModelCheckpoint + +from gpsr.modeling import GPSR +from gpsr.beams import NNTransform, NNParticleBeamGenerator +from gpsr.train import LitGPSR + +import os +import re + + +class GPSRRun: + """ + A class to manage the setup and execution of a GPSR training run. + This includes preparing datasets, models, logging, checkpointing, and trainer setup. + """ + + def __init__(self, hparams, log_name="scans"): + """ + Initialize the GPSRRun with hyperparameters and logging setup. + + Args: + hparams (dict): Hyperparameters for the model and training. + log_name (str): Name for the log directory. + """ + self.hparams = hparams + self.log_name = log_name + self.train_dataset = None + + print("Hyperparameters:") + pprint(self.hparams) + + def setup_training(self, train_dataset): + """ + Setup the full training pipeline, including dataset, model, logger, checkpointing, and trainer. + + Args: + train_dataset (torch.utils.data.Dataset): The dataset to use for training. + """ + self.train_dataset = train_dataset + + # Initialize the GPSR lattice (likely defining the physics or geometry structure) + self.gpsr_lattice = self.setup_gpsr_lattice() + + # Initialize the GPSR model with the lattice and particle generator + self.gpsr_model = self.setup_gpsr_model() + + # Wrap the GPSR model in the LitGPSR Lightning module + self.litgpsr = self.setup_litgpsr() + + # Prepare the DataLoader for training + self.train_loader = self.setup_trainloader() + + # Setup logger for tracking metrics and checkpoints + self.logger = self.setup_logger() + self.logger.log_hyperparams(self.hparams) + + # Setup checkpointing to save model progress + self.checkpoint_callback = self.setup_checkpointing() + + # Setup PyTorch Lightning Trainer + self.trainer = self.setup_trainer() + + def train(self): + """ + Start the training process using the configured trainer and DataLoader. + """ + print(f"Running training - results will be saved in {self.logger.log_dir}") + self.trainer.fit(self.litgpsr, self.train_loader) + + def setup_gpsr_lattice(self): + """ + Setup the GPSR lattice component. + Should be implemented to define the lattice structure. + """ + pass # To be implemented by the user + + def setup_gpsr_model(self): + """ + Initialize the GPSR model using provided hyperparameters and lattice. + + Returns: + GPSR: The initialized GPSR model object. + """ + return GPSR( + NNParticleBeamGenerator( + self.hparams["N_particles"], + self.hparams["p0c"], + transformer=NNTransform( + self.hparams["n_hidden"], + self.hparams["width"], + output_scale=self.hparams["output_scale"], + ), + ), + self.gpsr_lattice, + ) + + def setup_litgpsr(self): + """ + Wrap the GPSR model in the LitGPSR Lightning module. + + Returns: + LitGPSR: The wrapped model ready for training. + """ + return LitGPSR(self.gpsr_model, self.hparams["learning_rate"]) + + def setup_trainloader(self): + """ + Create the DataLoader for the training dataset. + + Returns: + DataLoader: The PyTorch DataLoader object. + """ + return torch.utils.data.DataLoader( + self.train_dataset, batch_size=self.hparams["batch_size"] + ) + + def setup_logger(self): + """ + Setup the CSV logger for experiment tracking. + + Returns: + CSVLogger: The logger object. + """ + return CSVLogger("logs", name=self.log_name) + + def setup_checkpointing(self): + """ + Configure model checkpointing to save progress at specified intervals. + + Returns: + ModelCheckpoint: The checkpoint callback. + """ + dirpath = os.path.join(self.logger.log_dir, "checkpoints") + checkpoint_callback = ModelCheckpoint( + dirpath=dirpath, # Directory to save checkpoints + filename="{step}", # Filename pattern + save_weights_only=False, # Save full model (not just weights) + every_n_epochs=self.hparams["checkpoint_period_epochs"], + save_top_k=-1, # Save all checkpoints + monitor="loss", # Monitor loss for saving + ) + return checkpoint_callback + + def setup_trainer(self): + """ + Setup the PyTorch Lightning Trainer with max epochs, logger, and callbacks. + + Returns: + Trainer: The PyTorch Lightning Trainer. + """ + return L.Trainer( + max_epochs=self.hparams["max_epochs"], + logger=self.logger, + callbacks=[self.checkpoint_callback], + ) + + @classmethod + def from_checkpoint( + cls, log_name, version_no, checkpoint_number=-1, extra_hparams={} + ): + """ + Load a GPSRRun instance from a saved checkpoint. + + Args: + log_name (str): Name of the log directory. + version_no (int): Version number of the experiment. + checkpoint_number (int): Index of the checkpoint to load (-1 for the latest). + extra_hparams (dict): Extra hyperparameters to override. + + Returns: + GPSRRun: The loaded GPSRRun instance. + """ + # Load hyperparameters from saved YAML + with open(f"{log_name}/version_{version_no}/hparams.yaml") as stream: + hparams = yaml.safe_load(stream) + + # Initialize the run + run = cls(hparams, log_name) + run.hparams.update(extra_hparams) + + # Re-setup model components + run.gpsr_lattice = run.setup_gpsr_lattice() + run.gpsr_model = run.setup_gpsr_model() + + # Get checkpoint filename + checkpoint_file_name = run.list_checkpoint_filenames(version_no)[ + checkpoint_number + ] + + print(f"Loading checkpoint {checkpoint_file_name}...") + + # Load the Lightning module from checkpoint + run.litgpsr = LitGPSR.load_from_checkpoint( + f"{checkpoint_file_name}", + gpsr_model=run.gpsr_model, + strict=False, + map_location=torch.device("cpu"), + ) + + return run + + def list_checkpoint_filenames(self, version_no): + """ + List all checkpoint filenames for a given version, sorted by step number. + + Args: + version_no (int): Version number of the experiment. + + Returns: + list: Sorted list of checkpoint file paths. + """ + checkpoint_filenames = [ + f"{self.log_name}/version_{version_no}/checkpoints/" + name + for name in sorted( + os.listdir(f"{self.log_name}/version_{version_no}/checkpoints") + ) + ] + + # Helper function to extract step number from filename + def extract_epoch(filename): + match = re.search(r"step=(\d+).", filename) + step = int(match.group(1)) if match else float("inf") + return step + + # Sort filenames by extracted step number + sorted_filenames = sorted(checkpoint_filenames, key=extract_epoch) + + return sorted_filenames From 804a12afb0104c031fea227ed1cc5e3b34754d4a Mon Sep 17 00:00:00 2001 From: Lewis Anthony Boulton Date: Tue, 3 Jun 2025 13:49:37 +0200 Subject: [PATCH 2/8] Provide gpsr_lattice as a input at initialisation --- gpsr/run.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/gpsr/run.py b/gpsr/run.py index 6de8c15..79d9ba0 100644 --- a/gpsr/run.py +++ b/gpsr/run.py @@ -20,7 +20,9 @@ class GPSRRun: This includes preparing datasets, models, logging, checkpointing, and trainer setup. """ - def __init__(self, hparams, log_name="scans"): + def __init__(self, + hparams, + log_name="scans"): """ Initialize the GPSRRun with hyperparameters and logging setup. From 92bf49a4d9603e2ff886ff320356b233c641fd31 Mon Sep 17 00:00:00 2001 From: Lewis Anthony Boulton Date: Tue, 3 Jun 2025 13:56:12 +0200 Subject: [PATCH 3/8] ruff --- gpsr/run.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/gpsr/run.py b/gpsr/run.py index 79d9ba0..6de8c15 100644 --- a/gpsr/run.py +++ b/gpsr/run.py @@ -20,9 +20,7 @@ class GPSRRun: This includes preparing datasets, models, logging, checkpointing, and trainer setup. """ - def __init__(self, - hparams, - log_name="scans"): + def __init__(self, hparams, log_name="scans"): """ Initialize the GPSRRun with hyperparameters and logging setup. From e9738dadf6cb599b0a2663c49d7cc05ddb968538 Mon Sep 17 00:00:00 2001 From: Lewis Anthony Boulton Date: Tue, 3 Jun 2025 13:59:18 +0200 Subject: [PATCH 4/8] Actually implement the changes with gpsr_lattice now, wasn't done in previous commits --- gpsr/run.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/gpsr/run.py b/gpsr/run.py index 6de8c15..87d3174 100644 --- a/gpsr/run.py +++ b/gpsr/run.py @@ -20,7 +20,7 @@ class GPSRRun: This includes preparing datasets, models, logging, checkpointing, and trainer setup. """ - def __init__(self, hparams, log_name="scans"): + def __init__(self, gpsr_lattice, hparams, log_name="scans"): """ Initialize the GPSRRun with hyperparameters and logging setup. @@ -28,6 +28,7 @@ def __init__(self, hparams, log_name="scans"): hparams (dict): Hyperparameters for the model and training. log_name (str): Name for the log directory. """ + self.gpsr_lattice = gpsr_lattice self.hparams = hparams self.log_name = log_name self.train_dataset = None @@ -44,8 +45,8 @@ def setup_training(self, train_dataset): """ self.train_dataset = train_dataset - # Initialize the GPSR lattice (likely defining the physics or geometry structure) - self.gpsr_lattice = self.setup_gpsr_lattice() + # # Initialize the GPSR lattice (likely defining the physics or geometry structure) + # self.gpsr_lattice = self.setup_gpsr_lattice() # Initialize the GPSR model with the lattice and particle generator self.gpsr_model = self.setup_gpsr_model() @@ -73,12 +74,12 @@ def train(self): print(f"Running training - results will be saved in {self.logger.log_dir}") self.trainer.fit(self.litgpsr, self.train_loader) - def setup_gpsr_lattice(self): - """ - Setup the GPSR lattice component. - Should be implemented to define the lattice structure. - """ - pass # To be implemented by the user + # def setup_gpsr_lattice(self): + # """ + # Setup the GPSR lattice component. + # Should be implemented to define the lattice structure. + # """ + # pass # To be implemented by the user def setup_gpsr_model(self): """ @@ -162,7 +163,7 @@ def setup_trainer(self): @classmethod def from_checkpoint( - cls, log_name, version_no, checkpoint_number=-1, extra_hparams={} + cls, gpsr_lattice, log_name, version_no, checkpoint_number=-1, extra_hparams={} ): """ Load a GPSRRun instance from a saved checkpoint. @@ -185,7 +186,7 @@ def from_checkpoint( run.hparams.update(extra_hparams) # Re-setup model components - run.gpsr_lattice = run.setup_gpsr_lattice() + run.gpsr_lattice = gpsr_lattice run.gpsr_model = run.setup_gpsr_model() # Get checkpoint filename From e0dc0eac1ab7c0ac7e1be0a0fe126b9422a63dcd Mon Sep 17 00:00:00 2001 From: Lewis Anthony Boulton Date: Tue, 3 Jun 2025 16:57:24 +0200 Subject: [PATCH 5/8] Have the mandatory hyperparameters as named arguments; train dataset now provided at init --- gpsr/run.py | 70 +++++++++++++++++++++++++++++++++++------------------ 1 file changed, 47 insertions(+), 23 deletions(-) diff --git a/gpsr/run.py b/gpsr/run.py index 87d3174..cab9bdc 100644 --- a/gpsr/run.py +++ b/gpsr/run.py @@ -13,40 +13,69 @@ import os import re - class GPSRRun: """ A class to manage the setup and execution of a GPSR training run. This includes preparing datasets, models, logging, checkpointing, and trainer setup. """ - def __init__(self, gpsr_lattice, hparams, log_name="scans"): + def __init__(self, + gpsr_lattice, + log_name="scans", + train_dataset = None, + N_particles = int(5e4), + n_hidden = 2, + hidden_width = 20, + output_scale = 1e-4, + dropout = 0.0, + batch_size = 100, + max_epochs = 5000, + p0c = 1000*1e6, + learning_rate = 10e-3, + checkpoint_period_epochs = 100, + **extra_hparams): """ - Initialize the GPSRRun with hyperparameters and logging setup. + Initializes the GPSRRun object with model hyperparameters and training settings. Args: - hparams (dict): Hyperparameters for the model and training. - log_name (str): Name for the log directory. + gpsr_lattice: The beamline or lattice structure for the GPSR model. + log_name (str): Name of the directory where logs will be saved. + train_dataset (Dataset, optional): Dataset to be used for training. + N_particles (int): Number of particles in the simulation. + n_hidden (int): Number of hidden layers in the NNTransform. + hidden_width (int): Width of each hidden layer. + output_scale (float): Scaling factor for NNTransform outputs. + dropout (float): Dropout rate for the model. + batch_size (int): Batch size for training. + max_epochs (int): Number of epochs to train. + p0c (float): Reference momentum of the beam (in eV/c). + learning_rate (float): Learning rate for training. + checkpoint_period_epochs (int): Interval (in epochs) to save checkpoints. + extra_hparams (dict): Any additional hyperparameters. """ self.gpsr_lattice = gpsr_lattice - self.hparams = hparams + self.hparams = {'N_particles' : N_particles, + 'n_hidden': n_hidden, + 'hidden_width': hidden_width, + 'output_scale' : output_scale, + 'dropout' : dropout, + 'batch_size' : batch_size, + 'max_epochs' : max_epochs, + 'p0c' : p0c, + 'learning_rate' : learning_rate, + 'checkpoint_period_epochs' : checkpoint_period_epochs} | extra_hparams + self.log_name = log_name - self.train_dataset = None + self.train_dataset = train_dataset print("Hyperparameters:") pprint(self.hparams) - def setup_training(self, train_dataset): + def setup_training(self): """ Setup the full training pipeline, including dataset, model, logger, checkpointing, and trainer. - Args: - train_dataset (torch.utils.data.Dataset): The dataset to use for training. """ - self.train_dataset = train_dataset - - # # Initialize the GPSR lattice (likely defining the physics or geometry structure) - # self.gpsr_lattice = self.setup_gpsr_lattice() # Initialize the GPSR model with the lattice and particle generator self.gpsr_model = self.setup_gpsr_model() @@ -74,12 +103,6 @@ def train(self): print(f"Running training - results will be saved in {self.logger.log_dir}") self.trainer.fit(self.litgpsr, self.train_loader) - # def setup_gpsr_lattice(self): - # """ - # Setup the GPSR lattice component. - # Should be implemented to define the lattice structure. - # """ - # pass # To be implemented by the user def setup_gpsr_model(self): """ @@ -94,7 +117,7 @@ def setup_gpsr_model(self): self.hparams["p0c"], transformer=NNTransform( self.hparams["n_hidden"], - self.hparams["width"], + self.hparams["hidden_width"], output_scale=self.hparams["output_scale"], ), ), @@ -182,11 +205,12 @@ def from_checkpoint( hparams = yaml.safe_load(stream) # Initialize the run - run = cls(hparams, log_name) + run = cls(gpsr_lattice, + **hparams, + log_name = log_name) run.hparams.update(extra_hparams) # Re-setup model components - run.gpsr_lattice = gpsr_lattice run.gpsr_model = run.setup_gpsr_model() # Get checkpoint filename From dfd5dd6667e9e92cbe2378248e1116dc6fe2a7fc Mon Sep 17 00:00:00 2001 From: Lewis Anthony Boulton Date: Tue, 3 Jun 2025 16:57:52 +0200 Subject: [PATCH 6/8] ruff --- gpsr/run.py | 62 +++++++++++++++++++++++++++-------------------------- 1 file changed, 32 insertions(+), 30 deletions(-) diff --git a/gpsr/run.py b/gpsr/run.py index cab9bdc..f2947ca 100644 --- a/gpsr/run.py +++ b/gpsr/run.py @@ -13,27 +13,30 @@ import os import re + class GPSRRun: """ A class to manage the setup and execution of a GPSR training run. This includes preparing datasets, models, logging, checkpointing, and trainer setup. """ - def __init__(self, - gpsr_lattice, - log_name="scans", - train_dataset = None, - N_particles = int(5e4), - n_hidden = 2, - hidden_width = 20, - output_scale = 1e-4, - dropout = 0.0, - batch_size = 100, - max_epochs = 5000, - p0c = 1000*1e6, - learning_rate = 10e-3, - checkpoint_period_epochs = 100, - **extra_hparams): + def __init__( + self, + gpsr_lattice, + log_name="scans", + train_dataset=None, + N_particles=int(5e4), + n_hidden=2, + hidden_width=20, + output_scale=1e-4, + dropout=0.0, + batch_size=100, + max_epochs=5000, + p0c=1000 * 1e6, + learning_rate=10e-3, + checkpoint_period_epochs=100, + **extra_hparams, + ): """ Initializes the GPSRRun object with model hyperparameters and training settings. @@ -54,17 +57,19 @@ def __init__(self, extra_hparams (dict): Any additional hyperparameters. """ self.gpsr_lattice = gpsr_lattice - self.hparams = {'N_particles' : N_particles, - 'n_hidden': n_hidden, - 'hidden_width': hidden_width, - 'output_scale' : output_scale, - 'dropout' : dropout, - 'batch_size' : batch_size, - 'max_epochs' : max_epochs, - 'p0c' : p0c, - 'learning_rate' : learning_rate, - 'checkpoint_period_epochs' : checkpoint_period_epochs} | extra_hparams - + self.hparams = { + "N_particles": N_particles, + "n_hidden": n_hidden, + "hidden_width": hidden_width, + "output_scale": output_scale, + "dropout": dropout, + "batch_size": batch_size, + "max_epochs": max_epochs, + "p0c": p0c, + "learning_rate": learning_rate, + "checkpoint_period_epochs": checkpoint_period_epochs, + } | extra_hparams + self.log_name = log_name self.train_dataset = train_dataset @@ -103,7 +108,6 @@ def train(self): print(f"Running training - results will be saved in {self.logger.log_dir}") self.trainer.fit(self.litgpsr, self.train_loader) - def setup_gpsr_model(self): """ Initialize the GPSR model using provided hyperparameters and lattice. @@ -205,9 +209,7 @@ def from_checkpoint( hparams = yaml.safe_load(stream) # Initialize the run - run = cls(gpsr_lattice, - **hparams, - log_name = log_name) + run = cls(gpsr_lattice, **hparams, log_name=log_name) run.hparams.update(extra_hparams) # Re-setup model components From 8d6a057cd11e2afee041f8824a94dae036ef6ebb Mon Sep 17 00:00:00 2001 From: Lewis Boulton Date: Mon, 22 Dec 2025 12:08:17 +0100 Subject: [PATCH 7/8] Add load_metrics function --- gpsr/run.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/gpsr/run.py b/gpsr/run.py index f2947ca..1b78916 100644 --- a/gpsr/run.py +++ b/gpsr/run.py @@ -1,6 +1,7 @@ import torch from pprint import pprint import yaml +import pandas as pd import lightning as L from lightning.pytorch.loggers import CSVLogger @@ -259,3 +260,8 @@ def extract_epoch(filename): sorted_filenames = sorted(checkpoint_filenames, key=extract_epoch) return sorted_filenames + + def load_metrics(self , version_no): + + return pd.read_csv(f'{self.log_name}/version_{version_no}/metrics.csv') + From 0ad81450381943904163c47dd1345caf5294aaaf Mon Sep 17 00:00:00 2001 From: Lewis Boulton Date: Mon, 22 Dec 2025 12:08:55 +0100 Subject: [PATCH 8/8] Account for repeated elements --- gpsr/modeling.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/gpsr/modeling.py b/gpsr/modeling.py index 6ada595..9c11f0e 100644 --- a/gpsr/modeling.py +++ b/gpsr/modeling.py @@ -18,6 +18,8 @@ from cheetah.accelerator import Element from gpsr.beams import BeamGenerator +import time + class GPSRLattice(torch.nn.Module, ABC): @abstractmethod @@ -223,9 +225,14 @@ def __init__( super().__init__() for element in variable_elements: - if not hasattr(element[0], element[1]): + element, attr = element + + if isinstance(element, list): + element = element[0] + + if not hasattr(element, attr): raise AttributeError( - f"Variable element {element[0].name} does not have parameter '{element[1]}'." + f"Variable element {element.name} does not have parameter '{attr}'." ) for element in observable_elements: @@ -253,11 +260,15 @@ def track_and_observe(self, beam) -> Tuple[Tensor, ...]: merged_segment = self.segment.transfer_maps_merged(beam) # Apply the merged segment transformations to the beam + #start = time.time() merged_segment(beam) + #print(time.time() - start) # Collect observations from the observable elements observations = tuple([element.reading for element in self.observable_elements]) + # FOR constrained phase spaces: compute the phase spaces and add them as observations + return observations def set_lattice_parameters(self, settings: torch.Tensor): @@ -268,4 +279,9 @@ def set_lattice_parameters(self, settings: torch.Tensor): settings: A tensor containing the new parameter values for the variable elements. """ for i, element in enumerate(self.variable_elements, 0): - setattr(element[0], element[1], settings[..., i]) + element, attr = element + if isinstance(element, list): + for ele in element: + setattr(ele, attr, settings[..., i]) + else: + setattr(element, attr, settings[..., i])