Cleanup trainer and SmartSimMeshMotionSolver

Allwmake

@@ -39,7 +39,7 @@ elif [ $WM_COMPILER == "Icx" ]; then
 fi
 make lib
 cd "$_REPO_ROOT" || exit 1
-cp $FOAM_SMARTREDIS/install/lib/libsmartredis.so $FOAM_USER_LIBBIN 
+cp $FOAM_SMARTREDIS/install/lib64/libsmartredis.so $FOAM_USER_LIBBIN 
 cp $FOAM_SMARTREDIS/build/Release/hiredis-prefix/src/hiredis-build/libhiredis.a $FOAM_USER_LIBBIN
 cp $FOAM_SMARTREDIS/build/Release/redis++-prefix/src/redis++-build/libredis++.a $FOAM_USER_LIBBIN
 export FOAM_CODE_TEMPLATES=$_REPO_ROOT/etc/dynamicCode/

run/meshMotion/.gitignore

@@ -2,3 +2,4 @@
 *.stl 
 *.csv
 *.pdf
+ellipsoid3D_MachineLearningMeshMotion/*

run/meshMotion/ml_model_training.py

@@ -1,129 +1,81 @@
 import argparse
-from smartredis import Client
 import torch
-import torch.nn as nn
 import numpy as np
 import io
-from sklearn.model_selection import train_test_split
-import torch.optim as optim 
-import time
-from typing import Tuple, Union
-from matplotlib import pyplot as plt
-
-from sklearn.metrics import mean_squared_error
-
-class MLP(nn.Module):
-    def __init__(self, num_layers, layer_width, input_size, output_size, activation_fn):
-        super(MLP, self).__init__()
-
-        layers = []
-        layers.append(nn.Linear(input_size, layer_width))
-        layers.append(activation_fn)
+import torch.optim as optim
 
-        for _ in range(num_layers - 2):
-            layers.append(nn.Linear(layer_width, layer_width))
-            layers.append(activation_fn)
+from matplotlib import pyplot as plt
+from smartredis import Client
 
-        layers.append(nn.Linear(layer_width, output_size))
-        self.layers = nn.Sequential(*layers)
+from MLP import MLP, MLPTrainer
 
-    def forward(self, x):
-        return self.layers(x)
 
-def train(num_mpi_ranks):
+def train(args):
     client = Client()
     torch.set_default_dtype(torch.float64)
 
     # Read the solution direction from a database
-    dimension = int(client.get_tensor("solution_dim"))
+    dimension = int(client.get_tensor("solution_dim")[0])
 
     print (f"Solution dimension = {dimension}.")
-
     # Initialize the model
-    model = MLP(
-        num_layers=3, 
-        layer_width=50, 
-        input_size=dimension, 
-        output_size=dimension, 
-        activation_fn=torch.nn.ReLU()
-    )
+    if args.model_name == "mlp":
+        model = MLP(
+            input_size=dimension,
+            output_size=dimension,
+            num_layers=3,
+            layer_width=10,
+            activation_fn=torch.nn.ELU()
+        )
+        trainer = MLPTrainer(model, args.radius_power)
 
-    # Initialize the optimizer
-    learning_rate = 1e-03
-    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
-
+    data_ready = client.poll_key("points", 1, 10000)
+    points = client.get_tensor("points")
+    interior_points = np.vstack([client.get_tensor(f"points_MPI_{i}" for i in range(4))])
+    X = torch.from_numpy(points).to(torch.float64)
     # Make sure all datasets are avaialble in the smartredis database.
+
+    epochs = 5000
     iteration = 1
     while True:
-    
+
         print (f"Iteration {iteration}")
 
         data_ready = client.poll_key("data_ready", 1, 10000)
         if (not data_ready):
             raise RuntimeError("Data not found in SmartRedis; aborting training.")
 
-        points = client.get_tensor("points")
         displacements = client.get_tensor("displacements")
+        interior_points = client.get_tensor
         client.delete_tensor("data_ready")
 
-        # Split training and validation data 
-        points_train, points_val, displ_train, displ_val = train_test_split(
-            points,
-            displacements,
-            test_size=0.2,
-            random_state=42
-        )
+        y = torch.from_numpy(displacements).to(torch.float64)
+
 
-        # Convert to torch.Tensor 
-        points_train = torch.from_numpy(points_train).to(torch.float64)
-        points_val   = torch.from_numpy(points_val).to(torch.float64)
-        displ_train  = torch.from_numpy(displ_train).to(torch.float64)
-        displ_val    = torch.from_numpy(displ_val).to(torch.float64)
-
-        loss_func = nn.MSELoss()
-
-        mean_mag_displ = torch.mean(torch.norm(displ_train, dim=1))
         validation_rmse = []
-        model.train()
-        epochs = 2000
         n_epochs = 0
-        rmse_loss_val = 1
-
-        for epoch in range(epochs):    
-            # Zero the gradients
-            optimizer.zero_grad()
-
-            # Forward pass on the training data
-            displ_pred = model(points_train)
-
-            # Compute loss on the training data
-            loss_train = loss_func(displ_pred, displ_train)
-
-            # Backward pass and optimization
-            loss_train.backward()
-            optimizer.step()
-
-            n_epochs = n_epochs + 1
-            # Forward pass on the validation data, with torch.no_grad() for efficiency
-            with torch.no_grad():
-                displ_pred_val = model(points_val)
-                mse_loss_val = loss_func(displ_pred_val, displ_val)
-                rmse_loss_val = torch.sqrt(mse_loss_val)
-                validation_rmse.append(rmse_loss_val)
-                if (mse_loss_val < 1e-04):
-                    break
-
-        print (f"RMSE {validation_rmse[-1]}, number of epochs {n_epochs}")
+
+        for epoch in range(epochs):
+            loss, model = trainer.training_step(X, y)
+            if trainer.converged():
+                break
+
+        print(f"MSE {loss.item()}, number of epochs {epoch}", flush=True)
+        np.savez(
+            f"data_{iteration:02d}.npz",
+            points=points,
+            displacements=displacements,
+        )
 
         # Uncomment to visualize validation RMSE
         plt.loglog()
         plt.title("Validation loss RMSE")
         plt.xlabel("Epochs")
         plt.plot(validation_rmse)
-        plt.savefig(f"validation_rmse_{epoch:04d}.png")
-    
+        plt.savefig(f"validation_rmse_{iteration:04d}.png")
+
         # Store the model into SmartRedis
-        # Put the model in evaluation mode. 
+        # Put the model in evaluation mode.
         model.eval() # TEST
         # Prepare a sample input
         example_forward_input = torch.rand(dimension)
@@ -136,18 +88,24 @@ def train(num_mpi_ranks):
         print("Saving model")
         client.set_model("model", model_buffer.getvalue(), "TORCH", "CPU")
         client.put_tensor("model_ready", np.array([0]))
-    
-        # Increase CFD+ML iteration 
+
+        # Increase CFD+ML iteration
         iteration = iteration + 1
 
-        # Check final iteration index and break 
+        # Check final iteration index and break
         if client.poll_key("final_iteration", 10, 10):
            print ("final iteration reached.")
            break
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Training script for mesh motion")
     parser.add_argument("mpi_ranks", help="number of mpi ranks", type=int)
+    parser.add_argument("radius_power", help="power law to weight losses", type=float)
+    parser.add_argument("model_name",
+                        help="which model to use to calculate interior displacements",
+                        choices=["mlp"],
+                        type=str
+    )
     args = parser.parse_args()
 
-    train(args.mpi_ranks)
+    train(args)

run/meshMotion/networks/MLP.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+class MLP(nn.Module):
+  def __init__(self, input_size, output_size, activation_fn, num_layers, layer_width):
+      super(MLP, self).__init__()
+
+      layers = []
+      layers.append(nn.Linear(input_size, layer_width))
+      layers.append(activation_fn)
+
+      for _ in range(num_layers - 2):
+          layers.append(nn.Linear(layer_width, layer_width))
+          layers.append(activation_fn)
+
+      layers.append(nn.Linear(layer_width, output_size))
+      self.layers = nn.Sequential(*layers)
+
+  def forward(self, x):
+      return self.layers(x)
+
+class MLPTrainer:
+  def __init__(self, model, radius_power, lr=1e-3, loss_stop=5e-5):
+    self.model = model
+    self.optimizer = optim.Adam(model.parameters(), lr=lr)
+    self.loss_stop = loss_stop
+    self.loss_value = None
+    self.radius_power = radius_power
+
+  def loss(self, X, y_true):
+    inner = y_true != 0.
+    center = torch.mean(X[inner], dim=0)
+    scaled_dist = torch.sqrt(torch.sum((X-center)**2, dim=1))**self.radius_power
+    wts = scaled_dist/torch.sum(scaled_dist)
+
+    y_pred = self.model(X)
+    return torch.sum(wts*torch.sum(torch.sqrt((y_true-y_pred)**2), dim=1))
+
+  def training_step(self, X, y_true):
+     self.optimizer.zero_grad()
+     loss_value = self.loss(X, y_true)
+     self.loss_value = loss_value
+     loss_value.backward()
+     self.optimizer.step()
+
+     return loss_value, self.model
+
+  def converged(self):
+     if self.loss_value.item() < self.loss_stop:
+        return True
+     return False
+
+
+
+

run/meshMotion/requirements.txt

@@ -0,0 +1,2 @@
+gmsh
+PyFoam

run/meshMotion/smartsim_driver.py

@@ -8,33 +8,30 @@
 
 from smartsim import Experiment
 
-def main():
-    parser = argparse.ArgumentParser(
-        description="Run a SmartSim Machine-Learning mesh deformation experiment"
-    )
-    parser.add_argument(
-        "--experiment", "-e",
-        required=True,
-        help="Name of the SmartSim experiment (e.g., mesh_deformation)"
-    )
-    parser.add_argument(
-        "--case", "-c",
-        required=True,
-        help="Name of the OpenFOAM case folder (e.g., ellipsoid3D)"
-    )
-    args = parser.parse_args()
+platform_config = {
+    "local": {
+        "launcher": "local",
+        "interface": "lo"
+    },
+    "hotlum": {
+        "launcher": "slurm",
+        "interface": "bond0"
+    }
+}
+
+def main(args):
 
     # ----------------------------------------------------------------
     # Create the SmartSim experiment
     # ----------------------------------------------------------------
 
-    exp = Experiment(args.experiment, launcher="local")
+    exp = Experiment(args.experiment, launcher=platform_config[args.platform]["launcher"])
 
     # ----------------------------------------------------------------
     # Launch the database
     # ----------------------------------------------------------------
 
-    db = exp.create_database(port=8000, interface="lo")
+    db = exp.create_database(port=8000, interface=platform_config[args.platform]["interface"])
     exp.generate(db, overwrite=True)
     exp.start(db)
     print(f"Database started at: {db.get_address()}")
@@ -55,13 +52,12 @@ def main():
     # ----------------------------------------------------------------
     # Configure and create the OpenFOAM mesh-motion model
     # ----------------------------------------------------------------
-    
-    # Create OpenFOAM moveDynamicMesh run settings 
+
+    # Create OpenFOAM moveDynamicMesh run settings
     openfoam_rs = exp.create_run_settings(
         exe="moveDynamicMesh",
         exe_args="-parallel",
-        run_command="mpirun",
-        run_args={"n": f"{num_mpi_ranks}"}
+        run_command="mpirun"
     )
     openfoam_rs.set_tasks(num_mpi_ranks)
     openfoam_rs.set_nodes(1)
@@ -79,7 +75,7 @@ def main():
 
     training_rs = exp.create_run_settings(
         exe="python",
-        exe_args=f"ml_model_training.py {num_mpi_ranks}"
+        exe_args=f"ml_model_training.py {num_mpi_ranks} {args.radius_power} mlp"
     )
     training_rs.set_tasks(1)
     training_rs.set_nodes(1)
@@ -88,7 +84,9 @@ def main():
         name="ml_model_training",
         run_settings=training_rs
     )
-    ml_model_training.attach_generator_files(to_copy="ml_model_training.py")
+    ml_model_training.attach_generator_files(
+        to_copy=["ml_model_training.py", "networks/MLP.py"]
+    )
 
     exp.generate(ml_model_training, overwrite=True)
 
@@ -111,4 +109,29 @@ def main():
         exp.stop(db)
 
 if __name__ == "__main__":
-    main()
+    parser = argparse.ArgumentParser(
+        description="Run a SmartSim Machine-Learning mesh deformation experiment"
+    )
+    parser.add_argument(
+        "--experiment", "-e",
+        required=True,
+        help="Name of the SmartSim experiment (e.g., mesh_deformation)"
+    )
+    parser.add_argument(
+        "--case", "-c",
+        required=True,
+        help="Name of the OpenFOAM case folder (e.g., ellipsoid3D)"
+    )
+    parser.add_argument(
+        "--radius_power",
+        default=0,
+        help="Power law associated with the loss function"
+    )
+    parser.add_argument(
+        "--platform",
+        choices=["slurm", "hotlum"],
+        default="local",
+        help="The platform on which this is being run"
+    )
+    args = parser.parse_args()
+    main(args)