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/.gitignore b/run/meshMotion/.gitignore index a8c86b9..bc09812 100644 --- a/run/meshMotion/.gitignore +++ b/run/meshMotion/.gitignore @@ -2,3 +2,4 @@ *.stl *.csv *.pdf +ellipsoid3D_MachineLearningMeshMotion/* diff --git a/run/meshMotion/ml_model_training.py b/run/meshMotion/ml_model_training.py index 7983477..9dfc346 100644 --- a/run/meshMotion/ml_model_training.py +++ b/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,11 +88,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 +100,12 @@ 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) + 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) diff --git a/run/meshMotion/networks/MLP.py b/run/meshMotion/networks/MLP.py new file mode 100644 index 0000000..4b72270 --- /dev/null +++ b/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 + + + + diff --git a/run/meshMotion/requirements.txt b/run/meshMotion/requirements.txt new file mode 100644 index 0000000..8ea9fb0 --- /dev/null +++ b/run/meshMotion/requirements.txt @@ -0,0 +1,2 @@ +gmsh +PyFoam diff --git a/run/meshMotion/smartsim_driver.py b/run/meshMotion/smartsim_driver.py index 863a627..bc61602 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,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) 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