Small adjustments

Author: robin-janssen

codes/surrogates/LatentNeuralODE/latent_neural_ode.py

@@ -372,6 +372,7 @@ class ModelWrapper(torch.nn.Module):
     Wraps the encoder, decoder, and neural ODE into a single model.
     Chooses architecture based on the config.model_version flag.
     """
+
     def __init__(self, config, n_quantities: int):
         super().__init__()
         self.config = config
@@ -383,7 +384,9 @@ def __init__(self, config, n_quantities: int):
             self.encoder = OldEncoder(
                 in_features=n_quantities,
                 latent_features=config.latent_features,
-                layers_factor=getattr(config, "layers_factor", 8),  # for backward compatibility
+                layers_factor=getattr(
+                    config, "layers_factor", 8
+                ),  # for backward compatibility
                 activation=config.activation,
             )
             self.decoder = OldDecoder(
@@ -640,7 +643,8 @@ def __init__(
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass to encode the input into the latent space."""
-        return self.mlp(x)
+        out = self.mlp(x)
+        return out
 
 
 class Decoder(torch.nn.Module):
@@ -683,7 +687,8 @@ def __init__(
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass to decode the latent representation into output features."""
-        return self.mlp(x)
+        out = self.mlp(x)
+        return out
 
 
 class OldEncoder(torch.nn.Module):

codes/surrogates/LatentNeuralODE/latent_neural_ode_config.py

@@ -39,6 +39,6 @@ class LatentNeuralODEBaseConfig:
     ode_layers: int = 4
     ode_width: int = 64
     ode_tanh_reg: bool = True
-    rtol: float = 1e-5
-    atol: float = 1e-5
+    rtol: float = 1e-6
+    atol: float = 1e-6
     learning_rate: float = 1e-3

codes/train/train_fcts.py

@@ -82,6 +82,9 @@ 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)

config.yaml

@@ -1,10 +1,10 @@
 # Global settings for the benchmark
-training_id: "srtest3"
-surrogates: ["MultiONet", "FullyConnected", "LatentPoly", "LatentNeuralODE", ]
-batch_size: [4096, 4096, 256, 256]
-epochs: [10,10,10,10] # [12000, 10000, 10000, 7000]
+training_id: "primordialtest"
+surrogates: ["LatentNeuralODE"]
+batch_size: [128]
+epochs: [20,100] # [12000, 10000, 10000, 7000]
 dataset: 
-  name: "simple_reaction"
+  name: "primordial"
   log10_transform: True
   normalise: "minmax" # "minmax" # "standardise", "minmax", "disable"
   use_optimal_params: True

datasets/_data_analysis/analyse_all_datasets.py

@@ -27,6 +27,7 @@ def main():
     Main function to analyse the dataset. It checks the dataset and loads the data.
     """
     datasets = [
+        "primordial",
         "simple_reaction",
         "osutest2",
         "coupled_oscillators",

datasets/_data_analysis/dataset_dict.py

@@ -51,4 +51,9 @@
         "qpp": 5,
         "tol": 1e-30,
     },
+    "primordial": {
+        "log": True,
+        "qpp": 5,
+        "tol": 1e-30,
+    },
 }

datasets/primordial/surrogates_config.py

@@ -0,0 +1,52 @@
+from dataclasses import dataclass
+
+from torch import nn
+
+
+@dataclass
+class MultiONetConfig:
+    """Model config for MultiONet for the simple_ode dataset"""
+
+    branch_hidden_layers: int = 6
+    trunk_hidden_layers: int = 4
+    hidden_size: int = 347
+    output_factor: int = 90
+    learning_rate: float = 1.2e-5
+    activation: nn.Module = nn.ReLU()
+
+
+@dataclass
+class LatentNeuralODEConfig:
+    """Model config for LatentNeuralODE for the simple_ode dataset"""
+
+    latent_features: int = 9
+    coder_layers: int = 4
+    coder_width: int = 230
+    learning_rate: float = 0.0004
+    ode_layers: int = 5
+    ode_width: int = 150
+    ode_tanh_reg: bool = True
+    activation: nn.Module = nn.LeakyReLU()
+    model_version: str = "v2"
+
+
+@dataclass
+class FullyConnectedConfig:
+    """Model config for FullyConnected for the simple_ode dataset"""
+
+    hidden_size: int = 800
+    num_hidden_layers: int = 8
+    learning_rate: float = 3e-5
+    activation: nn.Module = nn.GELU()
+
+
+@dataclass
+class LatentPolyConfig:
+    """Model config for LatentPoly for the simple_ode dataset"""
+
+    latent_features: int = 6
+    degree: int = 2
+    learning_rate: float = 0.002
+    coder_layers: int = 4
+    coder_width: int = 230
+    activation: nn.Module = nn.LeakyReLU()