From 1fa1e04f6b5ef60547fe8ec5190a49891137907e Mon Sep 17 00:00:00 2001 From: Andrew Shao Date: Mon, 29 Sep 2025 12:43:18 -0500 Subject: [PATCH 1/3] Add weighted loss function Adjust the loss function to weight by the distance from the "center" of the ellipsoid by an exponent d (e.g. r^{d}). Additionally, reduce the default size of the MLP to avoid having an overdetermined system (# of free parameters > samples). --- Allwmake | 2 +- run/meshMotion/ml_model_training.py | 106 ++++++++++++++-------------- run/meshMotion/smartsim_driver.py | 67 ++++++++++++------ 3 files changed, 99 insertions(+), 76 deletions(-) diff --git a/Allwmake b/Allwmake index 4dc38d0..e0430ee 100755 --- a/Allwmake +++ b/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/ diff --git a/run/meshMotion/ml_model_training.py b/run/meshMotion/ml_model_training.py index 7983477..12b3643 100644 --- a/run/meshMotion/ml_model_training.py +++ b/run/meshMotion/ml_model_training.py @@ -5,7 +5,7 @@ import numpy as np import io from sklearn.model_selection import train_test_split -import torch.optim as optim +import torch.optim as optim import time from typing import Tuple, Union from matplotlib import pyplot as plt @@ -30,7 +30,14 @@ def __init__(self, num_layers, layer_width, input_size, output_size, activation_ def forward(self, x): return self.layers(x) -def train(num_mpi_ranks): +def loss_weighted_center(y_true, y_pred, weights, weights_power): + weights_normed = torch.pow(weights, weights_power) + weights_normed = weights_normed/torch.sum(weights_normed) + + return torch.sum(torch.sum((y_true-y_pred)**2, dim=1)*weights_normed) + + +def train(args): client = Client() torch.set_default_dtype(torch.float64) @@ -38,24 +45,24 @@ def train(num_mpi_ranks): dimension = int(client.get_tensor("solution_dim")) 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() + num_layers=3, + layer_width=10, + input_size=dimension, + output_size=dimension, + activation_fn=torch.nn.ELU() ) # Initialize the optimizer - learning_rate = 1e-03 + learning_rate = 1e-3 optimizer = optim.Adam(model.parameters(), lr=learning_rate) - + # Make sure all datasets are avaialble in the smartredis database. iteration = 1 while True: - + print (f"Iteration {iteration}") data_ready = client.poll_key("data_ready", 1, 10000) @@ -64,66 +71,60 @@ def train(num_mpi_ranks): points = client.get_tensor("points") displacements = client.get_tensor("displacements") + 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 - ) + X = torch.from_numpy(points).to(torch.float64) + y = torch.from_numpy(displacements).to(torch.float64) + + # Find the center of the shape as the average of all the points on the inner boundary + r = torch.sqrt(torch.sum(X**2, dim=1)) + inner = r < 5 + center = torch.mean(X[inner], dim=0) + + dist = torch.sqrt(torch.sum((X-center)**2, dim=1)) + wts = dist/torch.sum(dist) - # 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 + epochs = 5000 n_epochs = 0 - rmse_loss_val = 1 - for epoch in range(epochs): + for epoch in range(epochs): # Zero the gradients optimizer.zero_grad() - + # Forward pass on the training data - displ_pred = model(points_train) - + displ_pred = model(X) + # Compute loss on the training data - loss_train = loss_func(displ_pred, displ_train) - + loss_train = loss_weighted_center(displ_pred, y, wts, args.radius_power) + + if (loss_train < 5e-05): + break + # 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}") + + print (f"MSE {loss_train.item()}, number of epochs {n_epochs}", 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,11 +137,11 @@ 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 @@ -148,6 +149,7 @@ def train(num_mpi_ranks): 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) args = parser.parse_args() - train(args.mpi_ranks) + train(args) diff --git a/run/meshMotion/smartsim_driver.py b/run/meshMotion/smartsim_driver.py index 863a627..e612e08 100644 --- a/run/meshMotion/smartsim_driver.py +++ b/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,16 +52,15 @@ 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}"} ) openfoam_rs.set_tasks(num_mpi_ranks) openfoam_rs.set_nodes(1) + openfoam_rs.set("exclusive") # Create the model from the OpenFOAM case argument openfoam_model = exp.create_model( @@ -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}" ) training_rs.set_tasks(1) training_rs.set_nodes(1) @@ -111,4 +107,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) From a07cea6eb442bafdda48d999880803c757545073 Mon Sep 17 00:00:00 2001 From: Andrew Shao Date: Tue, 7 Oct 2025 13:16:52 -0500 Subject: [PATCH 2/3] Only post boundary points on initialization Modify the solve routine and add a new method called during the, meshMotion constructor so that the boundary points are only posted once. --- .../displacementSmartSimMotionSolver.C | 228 ++++++++++-------- .../displacementSmartSimMotionSolver.H | 31 +-- 2 files changed, 150 insertions(+), 109 deletions(-) diff --git a/src/fvMotionSolvers/displacementSmartSimMotionSolver/displacementSmartSimMotionSolver.C b/src/fvMotionSolvers/displacementSmartSimMotionSolver/displacementSmartSimMotionSolver.C index 7714640..5fc0758 100644 --- a/src/fvMotionSolvers/displacementSmartSimMotionSolver/displacementSmartSimMotionSolver.C +++ b/src/fvMotionSolvers/displacementSmartSimMotionSolver/displacementSmartSimMotionSolver.C @@ -5,7 +5,7 @@ \\ / A nd | www.openfoam.com \\/ M anipulation | ------------------------------------------------------------------------------- - Copyright (C) 2023 Tomislav Maric, TU Darmstadt + Copyright (C) 2023 Tomislav Maric, TU Darmstadt ------------------------------------------------------------------------------- License This file is part of OpenFOAM. @@ -56,12 +56,12 @@ namespace Foam ); } -Foam::labelList Foam::displacementSmartSimMotionSolver::filterValidCmpts(const Vector