Add dummy timestep mechanism

Author: robin-janssen

codes/surrogates/AbstractSurrogate/surrogates.py

@@ -158,6 +158,7 @@ def prepare_data(
         timesteps: np.ndarray,
         batch_size: int,
         shuffle: bool,
+        dummy_timesteps: bool = True,
     ) -> tuple[DataLoader, DataLoader | None, DataLoader | None]:
         """
         Prepare the data for training, testing, and validation. This method should
@@ -170,6 +171,7 @@ def prepare_data(
             timesteps (np.ndarray): The timesteps.
             batch_size (int): The batch size.
             shuffle (bool): Whether to shuffle the data.
+            dummy_timesteps (bool): Whether to use dummy timesteps. Defaults to True.
 
         Returns:
             tuple[DataLoader, DataLoader, DataLoader]: The DataLoader objects for the

codes/surrogates/DeepONet/deeponet.py

@@ -237,6 +237,7 @@ def prepare_data(
         timesteps: np.ndarray,
         batch_size: int,
         shuffle: bool = True,
+        dummy_timesteps: bool = True,
     ) -> tuple[DataLoader, DataLoader, DataLoader | None]:
         """
         Prepare the data for the predict or fit methods.
@@ -249,11 +250,17 @@ def prepare_data(
             timesteps (np.ndarray): The timesteps.
             batch_size (int, optional): The batch size.
             shuffle (bool, optional): Whether to shuffle the data.
+            dummy_timesteps (bool, optional): Whether to create a dummy timestep array.
 
         Returns:
             tuple: The training, test, and validation DataLoaders.
         """
         dataloaders = []
+
+        # Create dummy timesteps
+        if dummy_timesteps:
+            timesteps = np.linspace(0, 1, dataset_train.shape[1])
+
         # Create the train dataloader
         dataloader_train = self.create_dataloader(
             dataset_train,
@@ -297,7 +304,7 @@ def fit(
             epochs (int, optional): The number of epochs to train the model.
             position (int): The position of the progress bar.
             description (str): The description for the progress bar.
-            multi_objective (bool): Whether multi-objective optimization is used. 
+            multi_objective (bool): Whether multi-objective optimization is used.
                                     If True, trial.report is not used (not supported by Optuna).
 
         Returns:

codes/surrogates/FCNN/fcnn.py

@@ -126,19 +126,32 @@ def forward(self, inputs: tuple) -> torch.Tensor:
     def prepare_data(
         self,
         dataset_train: np.ndarray,
-        dataset_test: np.ndarray,
+        dataset_test: np.ndarray | None,
         dataset_val: np.ndarray | None,
         timesteps: np.ndarray,
         batch_size: int,
         shuffle: bool = True,
+        dummy_timesteps: bool = True,
     ) -> tuple[DataLoader, DataLoader, DataLoader | None]:
         """
         Prepare the data for the predict or fit methods.
-        All datasets: shape (n_samples, n_timesteps, n_quantities)
 
-        Returns: train_loader, test_loader, val_loader
+        Args:
+            dataset_train (np.ndarray): Training data.
+            dataset_test (np.ndarray | None): Test data (optional).
+            dataset_val (np.ndarray | None): Validation data (optional).
+            timesteps (np.ndarray): Timesteps.
+            batch_size (int): Batch size.
+            shuffle (bool, optional): Whether to shuffle the data. Defaults to True.
+            dummy_timesteps (bool, optional): Whether to use dummy timesteps. Defaults to True.
+
+        Returns:
+            tuple[DataLoader, DataLoader | None, DataLoader | None]:
+                DataLoader for training, test, and validation data.
         """
         dataloaders = []
+        if dummy_timesteps:
+            timesteps = np.linspace(0, 1, dataset_train.shape[1])
         loader = self.create_dataloader(dataset_train, timesteps, batch_size, shuffle)
         dataloaders.append(loader)
         for dataset in [dataset_test, dataset_val]:
@@ -207,7 +220,7 @@ def fit(
                 # Update progress bar postfix
                 postfix = {
                     "train_loss": f"{train_losses[index]:.2e}",
-                    "test_loss": f"{test_losses[index]:.2e}"
+                    "test_loss": f"{test_losses[index]:.2e}",
                 }
                 progress_bar.set_postfix(postfix)
 

codes/surrogates/LatentNeuralODE/latent_neural_ode.py

@@ -78,6 +78,7 @@ def prepare_data(
         timesteps: np.ndarray,
         batch_size: int = 128,
         shuffle: bool = True,
+        dummy_timesteps: bool = True,
     ) -> tuple[DataLoader, DataLoader | None, DataLoader | None]:
         """
         Prepares the data for training by creating DataLoader objects.
@@ -101,6 +102,9 @@ def prepare_data(
             shuffled_indices = np.random.permutation(len(dataset_train))
             dataset_train = dataset_train[shuffled_indices]
 
+        if dummy_timesteps:
+            timesteps = np.linspace(0, 1, dataset_train.shape[1])
+
         # Create training DataLoader
         dset_train = ChemDataset(dataset_train, timesteps, device=self.device)
         dataloader_train = DataLoader(

codes/surrogates/LatentPolynomial/latent_poly.py

@@ -84,6 +84,7 @@ def prepare_data(
         timesteps: np.ndarray,
         batch_size: int = 128,
         shuffle: bool = True,
+        dummy_timesteps: bool = True,
     ) -> tuple[DataLoader, DataLoader | None, DataLoader | None]:
         """
         Prepare DataLoaders for training, testing, and validation.
@@ -95,10 +96,13 @@ def prepare_data(
             timesteps (np.ndarray): Array of timesteps.
             batch_size (int): Batch size.
             shuffle (bool): Whether to shuffle training data.
+            dummy_timesteps (bool): Whether to use dummy timesteps.
 
         Returns:
             tuple: DataLoaders for training, test, and validation datasets.
         """
+        if dummy_timesteps:
+            timesteps = np.linspace(0, 1, dataset_train.shape[1])
         if shuffle:
             shuffled_indices = np.random.permutation(len(dataset_train))
             dataset_train = dataset_train[shuffled_indices]

codes/train/train_fcts.py

@@ -82,9 +82,6 @@ def train_and_save_model(
 
     _, n_timesteps, n_quantities = train_data.shape
 
-    # # Replace timesteps with dummy timesteps between 0 and 1
-    # timesteps = np.linspace(0, 1, n_timesteps)
-
     # Get the surrogate class
     surrogate_class = get_surrogate(surr_name)
     model_config = get_model_config(surr_name, config)
@@ -107,6 +104,7 @@ def train_and_save_model(
             timesteps=timesteps,
             batch_size=batch_size,
             shuffle=True,
+            dummy_timesteps=True,
         )
 
     description = make_description(mode, device, str(metric), surr_name)

codes/tune/optuna_fcts.py

@@ -14,7 +14,7 @@
     measure_inference_time,
 )
 from codes.utils import check_and_load_data, make_description, set_random_seeds
-from codes.utils.data_utils import get_data_subset
+from codes.utils.data_utils import download_data, get_data_subset
 
 
 def load_yaml_config(config_path: str) -> dict:

config.yaml

@@ -1,5 +1,5 @@
 # Global settings for the benchmark
-training_id: "primordialtest"
+training_id: "primordialtest3"
 surrogates: ["LatentNeuralODE"]
 batch_size: [128]
 epochs: [20,100] # [12000, 10000, 10000, 7000]
@@ -10,7 +10,8 @@ dataset:
   use_optimal_params: True
   tolerance: 1e-30
   subset_factor: 1
-devices: ["cuda:1", "cuda:6"] # ["cuda:2", "cuda:3", "cuda:4", "cuda:5", "cuda:6", "cuda:7", "cuda:8", "cuda:9"]
+  log_timesteps: True
+devices: ["cuda:0", "cuda:1"] # ["cuda:2", "cuda:3", "cuda:4", "cuda:5", "cuda:6", "cuda:7", "cuda:8", "cuda:9"]
 seed: 42
 verbose: False