diff --git a/AUTHORS.rst b/AUTHORS.rst index af934369..9c8005c9 100644 --- a/AUTHORS.rst +++ b/AUTHORS.rst @@ -5,3 +5,4 @@ Contributors * Nick Heilenkötter, nheilenkoetter * Tom Freudenberg, TomF98 * Daniel Kreuter, dkreuter +* Maren Geisel, mgei01 diff --git a/NOTICE b/NOTICE index 5b8ccbae..c06d9871 100644 --- a/NOTICE +++ b/NOTICE @@ -21,6 +21,7 @@ # Please keep the list sorted. Robert Bosch GmbH + Maren Geisel Felix Hildebrand Uwe Iben Daniel Christopher Kreuter diff --git a/README.rst b/README.rst index d5640d1f..7c4e0728 100644 --- a/README.rst +++ b/README.rst @@ -26,7 +26,7 @@ TorchPhysics is build upon the machine learning library PyTorch_. Features ======== -The Goal of this library is to create a basic framework that can be used in many +The goal of this library is to create a basic framework that can be used in many different applications and with different deep learning methods. To this end, TorchPhysics aims at a: @@ -61,6 +61,13 @@ Some built-in features are: .. _Shapely: https://github.com/shapely/shapely .. _`PyTorch Lightning`: https://www.pytorchlightning.ai/ +Additional modules: + +- TorchPhysics comes with a wrapper module for `NVIDIA Modulus`_. This module serves as a bridge between the two frameworks and allows users to train TorchPhysics models with the Modulus training framework with minimal changes to their existing code. The additional installation of Modulus is required and documented in the `Wrapper Readme`_. + +.. _`NVIDIA Modulus`: https://developer.nvidia.com/modulus +.. _`Wrapper Readme`: ./src/torchphysics/wrapper/TPModulusWrapper.rst + Getting Started =============== @@ -78,7 +85,7 @@ to have a look at the following sections: Installation ============ -TorchPhysics reqiueres the follwing dependencies to be installed: +TorchPhysics requires the following dependencies to be installed: - Python >= 3.8 - PyTorch_ >= 2.0.0 @@ -99,7 +106,7 @@ Or by git clone https://github.com/boschresearch/torchphysics cd path_to_torchphysics_folder - pip install .[all] + pip install -e . if you want to modify the code. diff --git a/examples/wrapper/data/heat-eq-inverse-data.npy b/examples/wrapper/data/heat-eq-inverse-data.npy new file mode 100644 index 00000000..f75fe905 Binary files /dev/null and b/examples/wrapper/data/heat-eq-inverse-data.npy differ diff --git a/examples/wrapper/fdm_heat_equation.py b/examples/wrapper/fdm_heat_equation.py new file mode 100644 index 00000000..4197731c --- /dev/null +++ b/examples/wrapper/fdm_heat_equation.py @@ -0,0 +1,150 @@ +'''Implements a basic FDM to construct validation data for the heat equation. +''' +import numpy as np +import torch +from torchphysics.utils import UserFunction +from torchphysics.problem.spaces import Points + + +def FDM(domain_bounds, step_width_list, time_interval, parameter_list, + initial_condition): + """A simple FDM with a explicit euler to create some validation/inverse + data for the heatequation examples. Only works for rectangular domains and + Dirichlet boundary conditions. + + Parameters + ---------- + domain_bounds : list + The bounds of the rectangular domain, in the form: + [min_x1, max_x1, min_x2, max_x2] + step_width_list : list + The step size for the domain disrectization, in the form: + [step_x1, step_x2] + time_interval: list + The time interval for which the solution should be computed. + parameter_list : list + The different parameters for the heat equation, will solve the + equation for each params in the list and append the solutions. + initial_condition: callable + The initail condition. + + Returns + ------- + domain : list + The list of domain points. The first entry contains the x1-axis, the + sconde entry the x2-axis. For plotting the meshgrid of the points + have to be created. + time_domains : list + A list containing the different time domains. + Since the time step size scales with the used coefficent for the heat + equation, the number of time points may vary. + The i-th entry in the list corresspond to the i-th coefficent. + solutions : list + The different solutions, in the same order as time_domains. + + Notes + ----- + Only can solve the eq: du/dt = D*laplcian(u). + """ + solution = [] + time_domains = [] + dx = step_width_list[0] + dy = step_width_list[1] + initial_condition = UserFunction(initial_condition) + domain = _create_domain(domain_bounds, step_width_list) + # solve for each coefficent the heat equation + for k in parameter_list: + # scale time step to make scheme stable + dt = dx**2 * dy**2 / (2 * k * (dx**2 + dy**2)) + time = _create_domain(time_interval, [dt]) + u = _create_solution(domain, time) + _set_initial_condition(u[0, :, :], domain, time[0], initial_condition) + for i in range(1, len(time[0])): + u[i, :, :] = do_timestep(u[i-1, :, :], dx, dy, dt, k) + solution.append(u) + time_domains.append(time[0]) + return domain, time_domains, solution + + +def do_timestep(u0, dx, dy, dt, variable): + '''Implements the time step and fdm scheme of the methode + ''' + D = variable + dx2 = dx**2 + dy2 = dy**2 + u = u0.detach().clone() + u[1:-1, 1:-1] = u0[1:-1, 1:-1] + D * dt * ( + (u0[2:, 1:-1] - 2*u0[1:-1, 1:-1] + u0[:-2, 1:-1])/dx2 + + (u0[1:-1, 2:] - 2*u0[1:-1, 1:-1] + u0[1:-1, :-2])/dy2) + + return u + + +def _create_domain(domain_list, step_width_list): + domain = [[] for _ in range(int(len(domain_list)/2))] + for dim in range(len(domain)): + step_number = int((domain_list[2*dim+1] - domain_list[2*dim]) + / step_width_list[dim]) + domain[dim][:] = torch.linspace(domain_list[2*dim], domain_list[2*dim+1], + step_number+1) + return domain + + +def _create_solution(domain, time): + # creates the 'empyt' tensor to later store the solution + array_dim = [] + array_dim.append(len(time[0])) + for i in range(len(domain)): + array_dim.append(len(domain[i])) + solution_array = torch.zeros(tuple(array_dim)) + return solution_array + + +def _set_initial_condition(u0, domain, time, initial_condition): + input_dict = {'t': time[0]} + for i in range(len(domain[0])): + for j in range(len(domain[1])): + input_dict['x'] = torch.tensor([domain[0][i], domain[1][j]]).reshape(-1, 2) + u0[i, j] = initial_condition(input_dict) + + +def transform_to_points(domain, time_list, solution_list, parameter_list, + param_is_input): + """Transforms the FDM-data to data that can be used by the conditions. + Only meant to be used for the heat equation example. + + Parameters + ---------- + domain : list + The domain list from the FDM. + time_list : list + The list of time points from the FDM. + soluion_list : list + The list of solutions from the FDM. + parameter_list : list + The list of used parameters. + param_is_input : bool + If the parameter is a input variable of the model. + + + Returns + ------- + inp_data : Points + The input points for the model. + out_data : Poitns + The expected output. + """ + points = torch.empty((0, 4)) + out_data = torch.empty((0, 1)) + for k in range(len(parameter_list)): + new_points = np.array(np.meshgrid( + domain[0], domain[1], time_list[k], parameter_list[k])).T.reshape(-1, 4) + new_points = torch.as_tensor(new_points) + new_data_u = solution_list[k].reshape(-1, 1) + points = torch.cat((points, new_points)).float() + out_data = torch.cat((out_data, new_data_u)).float() + inp_data = {'x': points[:, [0, 1]], 't': points[:, [2]]} + out_data = {'u': out_data} + if param_is_input: + inp_data['D'] = points[:, [3]] + return Points.from_coordinates(inp_data), Points.from_coordinates(out_data) diff --git a/examples/wrapper/heat-equation-wrapper.ipynb b/examples/wrapper/heat-equation-wrapper.ipynb new file mode 100644 index 00000000..8cce4ee0 --- /dev/null +++ b/examples/wrapper/heat-equation-wrapper.ipynb @@ -0,0 +1,514 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# PINN: Heat equation with variable diffusion\n", + "Solving the heat equation in 2D for variable diffusion D using the PINN-concept.\n", + "\n", + "This is the same example as in the folder examples/pinn, but in the second part of the notebook, the problem is solved in Modulus with the wrapper module of TorchPhysics." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "import torchphysics as tp\n", + "import pytorch_lightning as pl\n", + "import math" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "First, we create the spaces for our problem. These define the variable names which will be used in the remaining part of this code.\n", + "\n", + "In this example, x is the space variable, t corresponds to the time, D is an interval of diffusions and u is the variable for the (1D-)solution." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "X = tp.spaces.R2('x')\n", + "T = tp.spaces.R1('t')\n", + "D = tp.spaces.R1('D')\n", + "\n", + "U = tp.spaces.R1('u')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "As a next step, we build the domain of the problem. There are multiple options to build multi-dimensional domains - in this case, we simply create a rectangle in space and intervals in time and diffusion which will later be multiplied to obtain the cartesian product." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "h, w = 20, 20" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "A_x = tp.domains.Parallelogram(X, [0, 0], [w, 0], [0, h])\n", + "A_t = tp.domains.Interval(T, 0, 40)\n", + "A_D = tp.domains.Interval(D, 0.1, 1.0)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Before we visualize the created domain, we create Sampler objects which are iterators that sample points from the domain during the optimization task. There are again various options to sample from the domains, an easy way would be to sample uniformly distributed random points. In this example, we choose an adaptive sampler to sample points in the inner of the domain. It will sample more points in points where the loss is large.\n", + "\n", + "The amount of sampled points is defined by their density in the 3/2-dim subset, it could be increased to achieve better training results." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "inner_sampler = tp.samplers.AdaptiveRandomRejectionSampler(A_x*A_t*A_D, density=1)\n", + "# initial values should be sampled on the left boundary of the time interval and for every x and D\n", + "initial_v_sampler = tp.samplers.RandomUniformSampler(A_x*A_t.boundary_left*A_D, density=1)\n", + "\n", + "boundary_v_sampler = tp.samplers.RandomUniformSampler(A_x.boundary*A_t*A_D, density=1)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We visualize the domain through the points created by the samplers using matplotlib:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig = tp.utils.scatter(X*T, inner_sampler, initial_v_sampler, boundary_v_sampler)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the next step we define the NN-model we want to fit to the PDE. A normalization can improve convergence for large or small domains." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model = tp.models.Sequential(\n", + " tp.models.NormalizationLayer(A_x*A_t*A_D),\n", + " tp.models.FCN(input_space=X*T*D, output_space=U, hidden=(50,50,50))\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, we define a condition which aims to minimze the mean squared error of the residual of the poisson equation. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def heat_residual(u, x, t, D):\n", + " return D*tp.utils.laplacian(u, x) - tp.utils.grad(u, t)\n", + "\n", + "pde_condition = tp.conditions.PINNCondition(module=model,\n", + " sampler=inner_sampler,\n", + " residual_fn=heat_residual,\n", + " name='pde_condition')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Additionally, we add a boundary condition at the boundary of the domain:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def boundary_v_residual(u):\n", + " return u\n", + "\n", + "boundary_v_condition = tp.conditions.PINNCondition(module=model,\n", + " sampler=boundary_v_sampler,\n", + " residual_fn=boundary_v_residual,\n", + " name='boundary_condition')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The initial condition can be defined via a data function. Again, we minimize the mean squared error over the sampled points." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def f(x):\n", + " return torch.sin(math.pi/w*x[:, :1])*torch.sin(math.pi/h*x[:,1:])\n", + "\n", + "def initial_v_residual(u, f):\n", + " return u-f\n", + "\n", + "initial_v_condition = tp.conditions.PINNCondition(module=model,\n", + " sampler=initial_v_sampler,\n", + " residual_fn=initial_v_residual,\n", + " data_functions={'f': f},\n", + " name='initial_condition')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For comparison, we compute the solution via a finite difference scheme." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "sys.path.append('..')\n", + "\n", + "from fdm_heat_equation import FDM, transform_to_points\n", + "\n", + "fdm_domain, fdm_time_domains, fdm_solution = FDM([0, w, 0, h], 2*[2e-1], [0,5], [0.1,0.2,0.4,0.6,0.8, 1.0], f)\n", + "fdm_inp, fdm_out = transform_to_points(fdm_domain, fdm_time_domains, fdm_solution, [0.1,0.2,0.4,0.6,0.8, 1.0], True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Comparsion to measured or computed data can be performed via a DataCondition using data supplied via a PointsDataLoader." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "val_condition = tp.conditions.DataCondition(module=model,\n", + " dataloader=tp.utils.PointsDataLoader((fdm_inp, fdm_out), batch_size=80000),\n", + " norm='inf')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 1. Training in TorchPhysics" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, we optimize the conditions using a pytorch-lightning.LightningModule Solver and running the training. In the Solver, the training and validation conditions, as well as all optimizer options can be specified.\n", + "\n", + "The process of the training can be monitored with Tensorboard (e.g. tensorboard --logdir=lightning_logs --port=0). The losses are plotted under the tab \"SCALARS\" and the plots of the PlotterCallback can be found in the tab \"IMAGES\"." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "decay_rate = 0.95\n", + "decay_steps = 15000\n", + "optim_setting = tp.solver.OptimizerSetting(torch.optim.Adam, lr=0.001,\n", + " scheduler_class=torch.optim.lr_scheduler.ExponentialLR,\n", + " scheduler_args={\"gamma\": decay_rate ** (1.0 / decay_steps)})\n", + "\n", + "solver = tp.solver.Solver([pde_condition,\n", + " boundary_v_condition,\n", + " initial_v_condition], [val_condition],\n", + " optimizer_setting=optim_setting)\n", + "\n", + "run_name ='Heat_equation_paramD'\n", + "\n", + "plot_sampler = tp.samplers.PlotSampler(plot_domain=A_x, n_points=1000,device='cpu', data_for_other_variables={'D':1.0, 't': 5})\n", + "callbacks = [ tp.utils.PlotterCallback(model=model,plot_function=lambda u: u[:,0],point_sampler=plot_sampler, \n", + " plot_type='contour_surface',log_name=run_name, check_interval=50), \n", + " \n", + " pl.callbacks.LearningRateMonitor(logging_interval='step')\n", + " ]\n", + "\n", + "\n", + "\n", + "trainer = pl.Trainer(devices=1, accelerator=\"gpu\",\n", + " max_steps = 200, \n", + " benchmark=True, \n", + " val_check_interval=40,\n", + " enable_checkpointing=False,\n", + " callbacks = callbacks\n", + " )\n", + "\n", + "trainer.fit(solver)\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, we plot the obtained solution:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "anim_sampler = tp.samplers.AnimationSampler(A_x, A_t, 100, n_points=400, data_for_other_variables={'D': 1.0})\n", + "anim = tp.utils.animate(model, lambda u: u[:, 0], anim_sampler, ani_speed=10)\n", + "anim[1].save('heat-eq.gif')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 2. Training with the wrapper in Modulus" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, we want to do the training in Modulus with the usage of the TPModulusWrapper.\n", + "\n", + "Remember that an additional Modulus installation is required!\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We use a Modulus Fourier net architecture with the ModulusArchitectureWrapper." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model = tp.models.Sequential(\n", + " tp.models.NormalizationLayer(A_x*A_t*A_D), \n", + " tp.wrapper.ModulusArchitectureWrapper(input_space=X*T*D, output_space=U, arch_name='fourier', frequencies = ['axis',[0,1,2]])\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "New initialization of the same conditions as before.\n", + "\n", + "Since Modulus by default saves the resulting values of the condition objectives and validation in VTK-files, the training procedure will slow down, when the amount of validation data increase or if the frequency of validation (val_check_interval) or the plotter callback frequency (check_interval) is too high." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pde_condition = tp.conditions.PINNCondition(module=model,\n", + " sampler=inner_sampler,\n", + " residual_fn=heat_residual,\n", + " name='pde_condition')\n", + "boundary_v_condition = tp.conditions.PINNCondition(module=model,\n", + " sampler=boundary_v_sampler,\n", + " residual_fn=boundary_v_residual,\n", + " name='boundary_condition')\n", + "initial_v_condition = tp.conditions.PINNCondition(module=model,\n", + " sampler=initial_v_sampler,\n", + " residual_fn=initial_v_residual,\n", + " data_functions={'f': f},\n", + " name='initial_condition')\n", + "val_condition = tp.conditions.DataCondition(module=model,\n", + " dataloader=tp.utils.PointsDataLoader((fdm_inp, fdm_out), batch_size=80000),\n", + " norm='inf')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The learning rate monitor callback could be omitted here, as Modulus uses it by default.\n", + "\n", + "The process of the training can be monitored again with tensorboard, but in this example with --logdir=outputs_heat, as this is the name of the folder containing the Modulus output.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "decay_rate = 0.95\n", + "decay_steps = 15000\n", + "optim_setting = tp.solver.OptimizerSetting(torch.optim.Adam, lr=0.001,\n", + " scheduler_class=torch.optim.lr_scheduler.ExponentialLR,\n", + " scheduler_args={\"gamma\": decay_rate ** (1.0 / decay_steps)})\n", + "\n", + "solver = tp.solver.Solver([pde_condition,\n", + " boundary_v_condition,\n", + " initial_v_condition], \n", + " [val_condition],\n", + " optimizer_setting=optim_setting)\n", + "\n", + "run_name ='Heat_equation_paramD_Modulus'\n", + "\n", + "plot_sampler = tp.samplers.PlotSampler(plot_domain=A_x, n_points=1000,device='cpu', data_for_other_variables={'D':1.0, 't': 5})\n", + "callbacks = [ tp.utils.PlotterCallback(model=model,plot_function=lambda u: u[:,0],point_sampler=plot_sampler, \n", + " plot_type='contour_surface',log_name=run_name, check_interval=50), \n", + " \n", + " pl.callbacks.LearningRateMonitor(logging_interval='step')\n", + " ]\n", + "\n", + "\n", + "\n", + "trainer = pl.Trainer(devices=1, accelerator=\"gpu\",\n", + " max_steps = 200, \n", + " benchmark=True, \n", + " val_check_interval=40,\n", + " enable_checkpointing=False,\n", + " callbacks = callbacks\n", + " )\n", + "\n", + "\n", + "tp.wrapper.TPModulusWrapper(trainer,solver,confdir_name='conf_heat',outputdir_name='outputs_heat').train()\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "anim_sampler = tp.samplers.AnimationSampler(A_x, A_t, 100, n_points=400, data_for_other_variables={'D': 1.0})\n", + "anim = tp.utils.animate(model, lambda u: u[:, 0], anim_sampler, ani_speed=10)\n", + "anim[1].save('heat-eq-wrapper.gif')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "There are addtional options for the training with the wrapper, e.g. use of aggregating function like LR Annealing or using spatial loss weighting, e.g. with the signed distance function. \n", + "For further information use the help function" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "help(tp.wrapper.TPModulusWrapper)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Training with the use of LR Annealing and the sdf function for spatial weighting." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "tp.wrapper.TPModulusWrapper(trainer,solver,lambda_weighting=[\"sdf\"],aggregator='LRAnnealing',confdir_name='conf_heat',outputdir_name='outputs_heat').train()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.14" + }, + "widgets": { + "application/vnd.jupyter.widget-state+json": { + "state": {}, + "version_major": 2, + "version_minor": 0 + } + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/wrapper/inverse-heat-equation-D-function-wrapper.ipynb b/examples/wrapper/inverse-heat-equation-D-function-wrapper.ipynb new file mode 100644 index 00000000..79bbb137 --- /dev/null +++ b/examples/wrapper/inverse-heat-equation-D-function-wrapper.ipynb @@ -0,0 +1,483 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Inverse heat equation with space dependent diffusion $D$\n", + "========================================================\n", + "In this example we want to solve an inverse problem where the parameter we are looking for is also a function.\n", + "We consider, $\\Omega = [0, 10] \\times [0, 10]$ and:\n", + "\\begin{align*}\n", + " \\text{div}(D(x)\\nabla u(x, t)) &= \\partial_t u(x, t) \\text{ in } \\Omega \\times [0, 5]\\\\\n", + " u(x, 0) &= 100\\sin(\\tfrac{\\pi}{10}x_1)\\sin(\\tfrac{\\pi}{10}x_2) \\text{ in } \\Omega \\\\\n", + " u(x, t) &= 0 \\text{ on } \\partial \\Omega \\times [0, 5]\n", + "\\end{align*} \n", + "The diffusion will be given by: $D(x)=5e^{-\\tfrac{1}{20}((x_1-3)^2 + (x_2-3)^2)}$. \n", + "\n", + "With an external program, we computed a FEM solution of this equation. The data can be found under the data folder.\n", + "This forward solution will be used as input data for the inverse problem.\n", + "\n", + "Since $D$ is a function too, we will use two neural networks. One for the temperature $u$ and for $D$. They will be connected over the PDE.\n", + "\n", + "This is the same example as in the folder examples/pinn, but in the second part of the notebook, the inverse problem is solved in Modulus with the wrapper module of TorchPhysics." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torchphysics as tp\n", + "import pytorch_lightning as pl\n", + "import torch" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The spaces and domains are defined like always:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "X = tp.spaces.R2('x')\n", + "T = tp.spaces.R1('t')\n", + "\n", + "D = tp.spaces.R1('D')\n", + "U = tp.spaces.R1('u')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "w, h = 10, 10\n", + "t_0, t_end = 0, 5\n", + "\n", + "domain_x = tp.domains.Parallelogram(X, [0, 0], [w, 0], [0, h])\n", + "domain_t = tp.domains.Interval(T, t_0, t_end)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here we now have two networks, like previously mentioned.\n", + "In the PDE both networks have to be evaluated, therefore we use the parallel evaluation of ``tp.models.Parallel``.\n", + "This objects will evaluate all models in parallel and pass the output to the corresponding condition." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model_u = tp.models.Sequential(\n", + " tp.models.NormalizationLayer(domain_x*domain_t),\n", + " tp.models.FCN(input_space=X*T, output_space=U, hidden=(80,80,50,50)))\n", + "model_D = tp.models.Sequential(\n", + " tp.models.NormalizationLayer(domain_x),\n", + " tp.models.FCN(input_space=X, output_space=D, hidden=(80,80,50,50))) # D only has X as an Input\n", + "parallel_model = tp.models.Parallel(model_u, model_D)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here we load the data of the forward problem. There are around, 90000 data points, we take only some random ones of them. This is chosen with ``num_of_data ``. \n", + "The data is in the form: [time, x-coordinates, temperature].\n", + "\n", + "After we have picked some points, we transform them to a ``Points``-object, so we can use the data in training." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "solution_data = np.load('data/heat-eq-inverse-data.npy')\n", + "num_of_data = 30000\n", + "solution_data = solution_data[np.random.choice(len(solution_data), num_of_data, replace=False), :].astype(np.float32)\n", + "solution_data = torch.tensor(solution_data)\n", + "inp_data = tp.spaces.Points.from_coordinates({'x': solution_data[:, 1:3], 't': solution_data[:, :1]})\n", + "out_data = tp.spaces.Points.from_coordinates({'u': solution_data[:, 3:]})" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The data will be put in a ``PointsDataLoader``, that will pass the data to the condition. \n", + "As a condition, we will here use a ``DataCondition`` that just compares the output of the model with the expected values.\n", + "\n", + "In this condition, we only train the model of temperature, since we only have data of the temperature. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "data_loader = tp.utils.PointsDataLoader((inp_data, out_data), batch_size=num_of_data)\n", + "data_condition = tp.conditions.DataCondition(module=model_u,\n", + " dataloader=data_loader,\n", + " norm=2, \n", + " use_full_dataset=True) " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For the differential equation we have to create some points inside the domain and check there the PDE.\n", + "Here, both models will be used, since $D$ and $u$ have to fulfill this condition together. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "inner_sampler = tp.samplers.RandomUniformSampler(domain_x*domain_t, n_points=15000)\n", + "\n", + "def heat_residual(u, x, t, D):\n", + " return tp.utils.div(D*tp.utils.grad(u, x), x) - tp.utils.grad(u, t)\n", + "\n", + "pde_condition = tp.conditions.PINNCondition(module=parallel_model, # use parallel evaluation\n", + " sampler=inner_sampler,\n", + " residual_fn=heat_residual,\n", + " name='pde_condition', \n", + " weight=100)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This is all we have to do. In total, the definition is not much different from the inverse problem with a constant $D$. Now we can start the training:" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 1. Training with TorchPhysics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "optim = tp.OptimizerSetting(optimizer_class=torch.optim.Adam, lr=0.001)\n", + "\n", + "solver = tp.solver.Solver([pde_condition, data_condition])\n", + "\n", + "trainer = pl.Trainer(devices=1, accelerator=\"gpu\",\n", + " max_steps=200,#8000,\n", + " logger=True,\n", + " benchmark=True,\n", + " enable_checkpointing=False\n", + " ) \n", + " \n", + "trainer.fit(solver)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "optim = tp.OptimizerSetting(optimizer_class=torch.optim.LBFGS, lr=0.5, \n", + " optimizer_args={'max_iter': 2, 'history_size': 100})\n", + "\n", + "# make sampler static.\n", + "pde_condition.sampler = pde_condition.sampler.make_static()\n", + "\n", + "solver = tp.solver.Solver([pde_condition, data_condition], optimizer_setting=optim)\n", + "\n", + "trainer = pl.Trainer(devices=1, accelerator=\"gpu\",\n", + " max_steps=200,#10000, # number of training steps\n", + " logger=True,\n", + " benchmark=True,\n", + " enable_checkpointing=False)\n", + " \n", + "trainer.fit(solver)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we can have look at the learned diffusion:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_sampler = tp.samplers.PlotSampler(plot_domain=domain_x, n_points=760, device='cuda')\n", + "fig = tp.utils.plot(model_D, lambda D : D, plot_sampler, plot_type='contour_surface')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here the exact function:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def exact(x):\n", + " return 5*torch.exp(-1/20.0 * ((x[:, :1] - 3)**2 + (x[:, 1:] - 3)**2))\n", + "\n", + "fig = tp.utils.plot(model_D, exact, plot_sampler, plot_type='contour_surface')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "And the absolute error:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def error(D, x):\n", + " return torch.abs(5*torch.exp(-1/20.0 * ((x[:, :1] - 3)**2 + (x[:, 1:] - 3)**2)) - D)\n", + "\n", + "fig = tp.utils.plot(model_D, error, plot_sampler, plot_type='contour_surface')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "anim_sampler = tp.samplers.AnimationSampler(domain_x, domain_t, 200, n_points=760)\n", + "fig, anim = tp.utils.animate(model_u, lambda u: u, anim_sampler, ani_speed=10, ani_type='contour_surface')\n", + "anim.save('inverse-heat-eq.gif')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 2. Training with the TorchPhysics wrapper in Modulus" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, we want to do the training in Modulus with the usage of the TPModulusWrapper.\n", + "\n", + "Remember that an additional Modulus installation is required!\n", + "\n", + "First, new initialization of the model and the conditions:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model_u = tp.models.Sequential(\n", + " tp.models.NormalizationLayer(domain_x*domain_t),\n", + " tp.models.FCN(input_space=X*T, output_space=U, hidden=(80,80,50,50)))\n", + "model_D = tp.models.Sequential(\n", + " tp.models.NormalizationLayer(domain_x),\n", + " tp.models.FCN(input_space=X, output_space=D, hidden=(80,80,50,50))) # D only has X as an Input\n", + "parallel_model = tp.models.Parallel(model_u, model_D)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "data_loader = tp.utils.PointsDataLoader((inp_data, out_data), batch_size=num_of_data)\n", + "data_condition = tp.conditions.DataCondition(module=model_u,\n", + " dataloader=data_loader,\n", + " norm=2, \n", + " use_full_dataset=True) " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "As the Modulus losses are different from the TorchPhysics losses, we don't add an additional weight for the pde_condition." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "inner_sampler = tp.samplers.RandomUniformSampler(domain_x*domain_t, n_points=15000)\n", + "\n", + "def heat_residual(u, x, t, D):\n", + " return tp.utils.div(D*tp.utils.grad(u, x), x) - tp.utils.grad(u, t)\n", + "\n", + "pde_condition = tp.conditions.PINNCondition(module=parallel_model, # use parallel evaluation\n", + " sampler=inner_sampler,\n", + " residual_fn=heat_residual,\n", + " name='pde_condition')\n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "optim = tp.OptimizerSetting(optimizer_class=torch.optim.Adam, lr=0.001)\n", + "\n", + "solver = tp.solver.Solver([pde_condition, data_condition])\n", + "\n", + "trainer = pl.Trainer(devices=1, accelerator=\"gpu\",\n", + " max_steps=200,#8000,\n", + " logger=True,\n", + " benchmark=True,\n", + " enable_checkpointing=False\n", + " ) \n", + "\n", + "tp.wrapper.TPModulusWrapper(trainer,solver).train()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To use LBFGS with the wrapper (Modulus), the max_steps has to be set to 1, the number of iterations are set by max_iter!\n", + "As it one main iteration step, there is no output in between the max_iter iterations - iteration takes several minutes. Reduce max_iter for testing purposes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "optim = tp.OptimizerSetting(optimizer_class=torch.optim.LBFGS, lr=0.5, \n", + " optimizer_args={'max_iter': 200,#10000, # number of training steps\n", + " 'history_size': 100})\n", + "\n", + "# make sampler static.\n", + "pde_condition.sampler = pde_condition.sampler.make_static()\n", + "\n", + "solver = tp.solver.Solver([pde_condition, data_condition], optimizer_setting=optim)\n", + "\n", + "trainer = pl.Trainer(devices=1, accelerator=\"gpu\",\n", + " max_steps=1, \n", + " logger=True,\n", + " benchmark=True,\n", + " enable_checkpointing=False)\n", + "\n", + "tp.wrapper.TPModulusWrapper(trainer,solver).train() " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_sampler = tp.samplers.PlotSampler(plot_domain=domain_x, n_points=760, device='cuda')\n", + "fig = tp.utils.plot(model_D, lambda D : D, plot_sampler, plot_type='contour_surface')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def error(D, x):\n", + " return torch.abs(5*torch.exp(-1/20.0 * ((x[:, :1] - 3)**2 + (x[:, 1:] - 3)**2)) - D)\n", + "\n", + "fig = tp.utils.plot(model_D, error, plot_sampler, plot_type='contour_surface')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "anim_sampler = tp.samplers.AnimationSampler(domain_x, domain_t, 200, n_points=760)\n", + "fig, anim = tp.utils.animate(model_u, lambda u: u, anim_sampler, ani_speed=10, ani_type='contour_surface')\n", + "anim.save('inverse-heat-eq-wrapper.gif')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.14" + }, + "widgets": { + "application/vnd.jupyter.widget-state+json": { + "state": {}, + "version_major": 2, + "version_minor": 0 + } + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/setup.cfg b/setup.cfg index 7f36c77f..88c59ae9 100644 --- a/setup.cfg +++ b/setup.cfg @@ -46,8 +46,9 @@ package_dir = # new major versions. This works if the required packages follow Semantic Versioning. # For more information, check out https://semver.org/. install_requires = - torch>=2.0.0, <2.4 + torch>=2.0.0, <=2.4 pytorch-lightning>=2.0.0 + # nvidia-modulus numpy>=1.20.2, <2.0 matplotlib>=3.0.0 scipy>=1.6.3 diff --git a/src/torchphysics/__init__.py b/src/torchphysics/__init__.py index 542903f6..1cd07388 100644 --- a/src/torchphysics/__init__.py +++ b/src/torchphysics/__init__.py @@ -6,6 +6,7 @@ from .models import * from .utils import * from .solver import Solver, OptimizerSetting +from .wrapper import TPModulusWrapper, ModulusArchitectureWrapper if sys.version_info[:2] >= (3, 8): # TODO: Import directly (no need for conditional) when `python_requires = >= 3.8` diff --git a/src/torchphysics/utils/callbacks.py b/src/torchphysics/utils/callbacks.py index 753540cb..0a0f19e0 100644 --- a/src/torchphysics/utils/callbacks.py +++ b/src/torchphysics/utils/callbacks.py @@ -54,9 +54,7 @@ def on_train_start(self, trainer, pl_module): self.model.state_dict(), self.path + "/" + self.name + "_init.pt" ) - def on_train_batch_start( - self, trainer, pl_module, batch, batch_idx, dataloader_idx - ): + def on_train_batch_start(self, trainer, pl_module, batch, batch_idx): if (self.check_interval > 0 and batch_idx > 0) and ( (batch_idx - 1) % self.check_interval == 0 ): @@ -124,7 +122,7 @@ def on_train_start(self, trainer, pl_module): self.point_sampler.sample_points(device=pl_module.device) def on_train_batch_end( - self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx + self, trainer, pl_module, outputs, batch, batch_idx ): if batch_idx % self.check_interval == 0: fig = plot( @@ -146,7 +144,7 @@ def on_train_end(self, trainer, pl_module): class TrainerStateCheckpoint(Callback): """ - A callback to saves the current state of the trainer (a PyTorch Lightning checkpoint), + A callback to save the current state of the trainer (a PyTorch Lightning checkpoint), if the training has to be resumed at a later point in time. Parameters @@ -162,8 +160,9 @@ class TrainerStateCheckpoint(Callback): Note ---- - To continue from the checkpoint, use `resume_from_checkpoint ="path_to_ckpt_file"` as an - argument in the initialization of the trainer. + To continue from the checkpoint, use ckpt_path="some/path/to/my_checkpoint.ckpt" as + argument in the fit command of the trainer. + The PyTorch Lightning checkpoint would save the current epoch and restart from it. In TorchPhysics we dont use multiple epochs, instead we train with multiple iterations @@ -180,7 +179,7 @@ def __init__(self, path, name, check_interval=200, weights_only=False): self.weights_only = weights_only def on_train_batch_end( - self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx + self, trainer, pl_module, outputs, batch, batch_idx ): if batch_idx % self.check_interval == 0: trainer.save_checkpoint( diff --git a/src/torchphysics/utils/differentialoperators.py b/src/torchphysics/utils/differentialoperators.py index 588e43a8..0bf50853 100644 --- a/src/torchphysics/utils/differentialoperators.py +++ b/src/torchphysics/utils/differentialoperators.py @@ -1,6 +1,6 @@ """File contains differentialoperators -NOTE: We aim to make the computation of differential operaotrs more efficient +NOTE: We aim to make the computation of differential operators more efficient by building an intelligent framework that is able to keep already computed derivatives and therefore make the computations more efficient. """ diff --git a/src/torchphysics/wrapper/TPModulusWrapper.rst b/src/torchphysics/wrapper/TPModulusWrapper.rst new file mode 100644 index 00000000..b829aecd --- /dev/null +++ b/src/torchphysics/wrapper/TPModulusWrapper.rst @@ -0,0 +1,110 @@ +=============================== +TorchPhysics to Modulus Wrapper +=============================== + +This folder contains a wrapper module for TorchPhysics to `NVIDIA Modulus`_. +This module serves as a bridge between the two frameworks and allows users to train TorchPhysics models with the Modulus training framework with minimal changes to their existing code. + +Both libraries are based on Pytorch_, but instead of `PyTorch Lightning`_, Modulus uses its own *Distributedmanager*, `TorchScript`_ with Pytorch JIT backend and CUDA graphs as a further framework to optimize and accelerate the training process, especially on NVIDIA GPUs. + +Both libraries use ``torch.nn.Module`` as the base class for the model definition, so that any model architecture of Modulus can easily be used as a TorchPhysics model, which is one of the main purposes of this wrapper module. +Modulus offers a wide range of model architectures, including various types of Fourier neural networks. + +As a second purpose, the wrapper module can be used to automatically convert a TorchPhysics problem to Modulus and perform the training in Modulus to benefit from the optimizable Modulus training framework. + +In both libraries, geometries are defined for the spatial domain of the problem (PDE or ODE), but the Modulus geometry provides a signed distance function (SDF) that calculates the distance of any arbitrary point to the boundary of the geometry. +This makes it possible to weight the loss function of the problem with the distance to the boundary, which is recommended in Modulus, e.g. for sharp gradients near the boundary, as it can increase convergence speed and improve accuracy. + +The wrapper module generally provides spatial loss weighting and the loss balancing algorithms of Modulus, e.g. GradNorm or Learning Rate Annealing. + +Each of the two main classes of the wrapper module can be used independently, allowing the user to choose whether to use only the conversion of the problem or only the model architecture, or both. + +Variable conventions +==================== +The spatial and temporal variables in Modulus are always defined as 'x', 'y', 'z' and 't', so TorchPhysics models can only be converted if the spatial variable is defined as 'x', which can be 1D, 2D or 3D, or 'x', 'y' or 'x', 'y','z' and the temporal variable as 't'. + +Contents +======== +The wrapper module consists of two main classes: + +* ``TPModulusWrapper``: can be used to convert a TorchPhysics problem into Modulus and to perform training in Modulus. ``help(TPModulusWrapper)`` provides a detailed description of the class and its parameters. + +* ``ModulusArchitectureWrapper``: can be used to choose any implemented Modulus architecture as TorchPhysics model. ``help(ModulusArchitectureWrapper)`` provides a detailed description of the class and its parameters. + +Usage +===== +To use the wrapper module to run a TorchPhysics model in Modulus, you can add the following line to your existing TorchPhysics code after the definition +of the ``trainer`` and the ``solver`` object, replacing the ``trainer.fit(solver)`` call: + +.. code-block:: python + + torchphysics.wrapper.TPModulusWrapper(trainer,solver).train() + + +To use one of the Modulus model architectures, you can add the following line to your existing TorchPhysics code and replace the model definition, +e.g. if you want to use the Modulus Fourier architecture as a TorchPhysics model: + +.. code-block:: python + + model=torchphysics.wrapper.ModulusArchitectureWrapper(input_space=X*T, output_space=U,arch_name='fourier',frequencies = ['axis',[0,1,2]]) + + +Installation +============ +The wrapper module requires a working installation of TorchPhysics and Modulus Symbolic (Sym), which is a framework providing algorithms +and utilities to be used with Modulus core for physics informed model training. + +The installation of Modulus Sym is documented here: `NVIDIA Modulus Github Repository`_ + +We recommend to create a new conda environment and to first install NVIDIA Modulus with the following command: + +.. code-block:: python + + pip install nvidia-modulus.sym + +Then you can install TorchPhysics as described in the `TorchPhysics documentation`_. + +As Modulus Sym uses TorchScript_ by default to compile the model, it is important to have the correct version of PyTorch_ installed. The current Modulus Sym version requires 2.1.0a0+4136153. +If a different version is installed, a warning will be raised when starting the training and the use of TorchScript is disabled. +The version can be determined by the command + +.. code-block:: python + + torch.__version__ + +To circumvent disabling of TorchScript you can edit the file /modulus/sym/constants.py in your python site packages installation path, and change the constant JIT_PYTORCH_VERSION = "2.1.0a0+4136153" to the version you have installed, e.g. "2.4.0+cu121". + +.. _`PyTorch Lightning`: https://www.pytorchlightning.ai/ +.. _`NVIDIA Modulus`: https://developer.nvidia.com/modulus +.. _`NVIDIA Modulus Github Repository`: https://github.com/NVIDIA/modulus-sym/tree/main +.. _PyTorch: https://pytorch.org/ +.. _TorchScript: https://pytorch.org/docs/stable/jit.html +.. _`TorchPhysics documentation`: https://github.com/boschresearch/torchphysics/blob/main/README.rst + + + +Testing +======= +As the wrapper module needs additional installation steps and can not be used without Modulus, it is excluded from the automatic testing with pytest. To test the functionality of the wrapper, there are example notebooks in the folder examples/wrapper and tests in src/torchphysics/wrapper/tests that can be manually invoked by the command (requires the installation of pytest and pytest-cov): + +.. code-block:: python + + pytest src/torchphysics/wrapper/tests + + + +Some notes +========== +* The loss definition in Modulus is based on Monte Carlo integration and therefore the loss is scaled proportional to the corresponding area, i.e. it is usually different from the loss in TorchPhysics, where the loss is the mean value. +* Currently, ``stl``-file support in Modulus is only available for Docker installation, so ``shapely`` and ``Trimesh`` geometries in TorchPhysics can not be converted. +* Cross product domains are generally not supported in Modulus, so must be automatically converted by the wrapper to existing primary geometries, so not all combinations of domain operations are allowed, e.g. product domains only from the union of 1D or 0D domains and no further rotation and translation is allowed (must be done with the entire product). +* Physics-Informed Deep Operator Networks (PIDOns) are currently not supported in the wrapper. +* Fourier Neural Operators (FNOs) are currently not supported in the wrapper, but an FNO framework is currently being developed in TorchPhysics. +* Samplers other than random uniformn and Halton sequence are not supported in Modulus. +* The imposition of exact boundary conditions using hard constraints with Approximate Distance Functions (ADFs) is not yet supported in TorchPhysics. +* The Limited-memory Broyden-Fletcher-Goldfarb-Shanno (L-BFGS) optimizer can be used in Modulus by setting the maximum step size (``max_steps``) to 1 (one single optimization step), but using the maximum number of iterations per optimization step (``max_iter``) as the number of iterations instead. This is very slow, so it is recommended to use Adam instead. In TorchPhysics, ``max_iter`` is decreased and many optimization steps are performed. +* If the combination of the Adam and L-BFGS optimizers is used, then loading the L-BFGS optimizer checkpoint file (optim_checkpoint.0.pth) will result in an error regarding ``max_iter`` as Adam does not use ``max_iter``. This is a known issue for Modulus support and it is recommended to delete or rename the optim_checkpoint.0.pth file. Then it works, but Tensorboard cannot display the loss history correctly! +* If several losses with the same name of the objective variable are used, the losses are summarized in Tensorboard, e.g. initial condition for T and Dirichlet condition for T, then there is only one loss (sum) for T. +* In general, all TorchPhysics callbacks are supported, but for the ``WeightSaveCallback`` the check for minimial loss (parameter ``check_interval``) is not supported by the wrapper, only initial and final model states are saved. +* Modulus automatically provides Tensorboard logging of the losses. The corresponding logging folder is ``outputs`` by default, but can be set by the user with the parameter ``outputdir_name``. +* Modulus automatically provides ``.vtp``-files containing data computed on the collocation points of the conditions that can be found in subfolders of the output directory. These files can be viewed using visualization tools like Paraview. \ No newline at end of file diff --git a/src/torchphysics/wrapper/__init__.py b/src/torchphysics/wrapper/__init__.py new file mode 100644 index 00000000..6a95314f --- /dev/null +++ b/src/torchphysics/wrapper/__init__.py @@ -0,0 +1,7 @@ +""" Wrapper module that serves as a bridge between TorchPhysics and Nvidia Modulus and allows users to train TorchPhysics models with the Modulus training framework with minimal changes to their existing code. +""" +from .wrapper import TPModulusWrapper +from .solver import ModulusSolverWrapper +from .model import ModulusArchitectureWrapper +from .helper import convertDataModulus2TP, convertDataTP2Modulus + \ No newline at end of file diff --git a/src/torchphysics/wrapper/geometry.py b/src/torchphysics/wrapper/geometry.py new file mode 100644 index 00000000..a509970b --- /dev/null +++ b/src/torchphysics/wrapper/geometry.py @@ -0,0 +1,1280 @@ +# This file includes work covered by the following copyright and permission notices: +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# All rights reserved. +# Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import torch.nn as nn +from torch import atan2, sqrt +import numpy as np +from functools import reduce +from operator import mul + +from modulus.sym.geometry.geometry import Geometry, csg_curve_naming +from modulus.sym.geometry.curve import Curve, SympyCurve +from modulus.sym.geometry.parameterization import Parameterization, Bounds, Parameter +from modulus.sym.geometry.helper import _sympy_sdf_to_sdf + +from sympy import Symbol, Min, Max, Heaviside, sqrt, Abs, sign, Eq, Or, And, Rational, lambdify +from sympy.vector import CoordSys3D + + + +def GeometryAxischange( + self, + axis1: str="x", + axis2: str="y", + axis3: str="z", + parameterization=Parameterization(), + ): + """ + Exchange two axes of the geometry. Many primitive geometries in Modulus are defined in a specific coordinate system and this function allows to change the coordinate system of the geometry. + It transforms the boundary curves, bounds and computes a new sdf function as wrapper function. + Parameters + ---------- + axis1: str + name of the first axis + axis2: str + name of the second axis + axis3: str + name of the third axis + parameterization: Parameterization, optional + parameterization of the geometry + + Returns + ------- + Geometry + new geometry object with exchanged axes + + """ + # create wrapper function for sdf function + def _axischange_sdf(sdf, dims,axes): + def axischange_sdf(invar, params, compute_sdf_derivatives=False): + changed_invar = {**invar} + _changed_invar = {**changed_invar} + for i, key in enumerate(dims): + _changed_invar[key] = changed_invar[axes[i]] + + # compute sdf + computed_sdf = sdf(_changed_invar, params, compute_sdf_derivatives) + return computed_sdf + + return axischange_sdf + + axes = (axis1,axis2,axis3) + new_sdf = _axischange_sdf(self.sdf, self.dims, axes) + # add parameterization + new_parameterization = self.parameterization.union(parameterization) + + # change bounds + bound_ranges = {**self.bounds.bound_ranges} + _bound_ranges = {**bound_ranges} + + for i, key in enumerate(self.dims): + _bound_ranges[Parameter(key)] = bound_ranges[Parameter(axes[i])] + + new_bounds = self.bounds.copy() + new_bounds.bound_ranges = _bound_ranges + + # change curves + new_curves = [] + for c in self.curves: + new_c = c.axischange(axes, parameterization) + new_curves.append(new_c) + + # return rotated geometry + return Geometry( + new_curves, + new_sdf, + len(self.dims), + new_bounds, + new_parameterization, + interior_epsilon=self.interior_epsilon, + ) + +def CurveAxischange(self, axes, parameterization=Parameterization()): + """ + Exchange two axes of the curves. Many primitive geometry curves + in Modulus are defined in a specific coordinate system and this + function allows to change the coordinate system of the curve. + It computes a new sample function for the curves as wrapper + function. + + Parameters + ---------- + axes: tuple + names of the axes + parameterization: Parameterization, optional + parameterization of the geometry + Returns + ------- + Curve + new curve object with exchanged axes + """ + def _sample(internal_sample, dims, axes): + def sample( + nr_points, parameterization=Parameterization(), quasirandom=False + ): + # sample points + invar, params = internal_sample( + nr_points, parameterization, quasirandom + ) + changed_invar = {**invar} + _changed_invar = {**changed_invar} + + for i, key in enumerate(dims): + _changed_invar[key] = changed_invar[axes[i]] + _changed_invar['normal_'+key] = changed_invar['normal_'+axes[i]] + + return _changed_invar, params + + return sample + return Curve( + _sample(self._sample, self.dims,axes), + len(self.dims), + self.parameterization.union(parameterization), + ) + +# Modulus Curve class gets the additional method CurveAxischange +Curve.axischange=CurveAxischange +# Modulus Geometry class gets the additional method GeometryAxischange +Geometry.axischange=GeometryAxischange + +class TPGeometryWrapper(): + """ + Wrapper to convert dictionary with spatial domain decomposition + into Modulus geometry. The domain is recursively analyzed mainly + concerning domain operations and product domains to identify the + underlying spatial geometries. + A general product of domains as in TorchPhysics can not be + implemented in Modulus, because the domain defines a spatial + geometry that must have a sdf implementation which is not available + in general for an arbitrary product domain. + Therefore each product domain has to be analyzed and mapped on a + spatial geometry resulting on union/intersection/cut of a + translated/rotated primitive geometry of Modulus. + + Not supported types/combinations in TorchPhysics? + + Parameters + ---------- + domain_dict: dictionary + spatial domain decomposition + + + + """ + def __init__(self,domain_dict)-> None: + self.domain_dict = domain_dict + + def getModulusGeometry(self): + return self.RecursiveDomainDictGeoWrapper(self.domain_dict,translate_vec=[0,0,0],rotation_list=[[],[],[],0]) + + + def RecursiveDomainDictGeoWrapper(self,domain_dict,prod_factor=None,translate_vec=[0,0,0],rotation_list=[[],[],[],0]): + """ + Recursive function that analyzes a dictionary of TorchPhysics + (spatial) domains as a result of domain decomposition to map + the domain to a Modulus geometry. + + Parameters + ---------- + domain_dict: dictionary + dictionary with TorchPhysics (spatial) domains. + Keys are domain operations: + 'u' (union), 'ub' (boundary of union),'d' (decomposed), + 'i' (intersection), 'ib' (boundary of intersection), + 'c' (cut), 'cb' (boundary of cut), 't' (tranlation), + 'r' (rotation),'p' (product) + prod_factor: float, optional + domain factor of previous decomposed product (3D domain + can consist of two cross products) + + Returns + ------- + geometry object + Modulus geometry object + is_boundary: bool + True if domain is boundary + cond: sympy expression + condition to restrict on parts of geometry, resulting from + cross products + translate_vec: list + translation vector, containing the translation in x,y,z + rotation_list: list + list containing the rotation angles, rotation axes, rotation points and the priority of the rotation vs. translation + """ + + for key, val in domain_dict.items(): + if (key=='u') or (key=='ub'): + is_boundary1, geometry1, cond1, translate_vec, rotation_list = self.RecursiveDomainDictGeoWrapper(val[0],prod_factor,translate_vec,rotation_list) + is_boundary2, geometry2, cond2, translate_vec, rotation_list = self.RecursiveDomainDictGeoWrapper(val[1],prod_factor,translate_vec,rotation_list) + + assert is_boundary1==is_boundary2 + + if cond1==cond2: + if cond1!=None: + cond = cond1 + else: + cond =None + elif cond1==None: + cond = cond2 + elif cond2==None: + cond=cond1 + else: + cond = self._or_condition(self._geo_condition(geometry1.sdf,cond1,is_boundary1),self._geo_condition(geometry2.sdf,cond2,is_boundary2)) + return is_boundary1,geometry1+geometry2, cond, translate_vec, rotation_list + elif key=='i' or key=='ib': + is_boundary1, geometry1, cond1, translate_vec, rotation_list = self.RecursiveDomainDictGeoWrapper(val[0],prod_factor,translate_vec, rotation_list) + is_boundary2, geometry2, cond2, translate_vec, rotation_list = self.RecursiveDomainDictGeoWrapper(val[1],prod_factor,translate_vec, rotation_list) + assert is_boundary1==is_boundary2 + if cond1==cond2: + if cond1!=None: + cond = cond1 + else: + cond =None + elif cond1==None: + cond = cond2 + elif cond2==None: + cond=cond1 + else: + cond = "And(" + cond1 + "," + cond2 + ")" + return is_boundary1, geometry1&geometry2, cond, translate_vec, rotation_list + + elif key=='c' or key=='cb': + is_boundary1, geometry1, cond1, translate_vec, rotation_list = self.RecursiveDomainDictGeoWrapper(val[0],prod_factor,translate_vec, rotation_list) + is_boundary2, geometry2, cond2, translate_vec, rotation_list = self.RecursiveDomainDictGeoWrapper(val[1],prod_factor,translate_vec, rotation_list) + assert is_boundary1==is_boundary2 + if cond1==cond2: + if cond1!=None: + cond = cond1 + else: + cond =None + elif cond1==None: + cond = cond2 + elif cond2==None: + cond=cond1 + else: + cond = "Or(" + cond1 + "," + cond2 + ")" + return is_boundary1, geometry1-geometry2, cond, translate_vec, rotation_list + + elif key=='p': + if prod_factor is None: + return self.ProductDomainWrapper(val[0],val[1],translate_vec=translate_vec, rotation_list=rotation_list) + else: + return self.ProductDomainWrapper(val[0],val[1],prod_factor,translate_vec=translate_vec, rotation_list=rotation_list) + + + elif key=='t': + assert ((next(iter(val[0].keys()))=='d') | (next(iter(val[0].keys()))=='r') ), "Only translation of primitive or rotated domains allowed" + rot_angles, rot_axes, rot_points, rot_prio = rotation_list + if next(iter(val[0].keys()))=='d': + # translation has to be done first + rotation_list[3] = 1 + dom_space = val[0]['d'].space + else: + # rotation has to be done first + rotation_list[3] = 0 + assert (next(iter(val[0]['r'][0].keys()))=='d'), "Only translation of primitive or rotated domains allowed" + dom_space = val[0]['r'][0]['d'].space + + vec = [float(val_in) for val_in in val[1]] + translate_vec_new = self.compute_3D_translate_vec(dom_space,vec) + + if translate_vec !=[0,0,0]: + translate_vec =[translate_vec[ii]+translate_vec_new[ii] for ii in range(3)] + else: + translate_vec = translate_vec_new + + is_boundary, geometry, cond, translate_vec, rotation_list = self.RecursiveDomainDictGeoWrapper(val[0],prod_factor=prod_factor,translate_vec=translate_vec,rotation_list=rotation_list) + return is_boundary, geometry, cond, translate_vec, rotation_list + elif key=='r': + assert ((next(iter(val[0].keys()))=='d') | (next(iter(val[0].keys()))=='t')), "Only rotation of primitive or translated domains allowed" + rot_angles, rot_axes, rot_points, rot_prio = rotation_list + + + if next(iter(val[0].keys()))=='d': + # rotation has to be done first + rotation_list[3] = 0 + dom_space = val[0]['d'].space + else: + # translation has to be done first + rotation_list[3] = 1 + assert (next(iter(val[0]['t'][0].keys()))=='d'), "Only rotation of primitive or translated domains allowed" + dom_space = val[0]['t'][0]['d'].space + + rot_matrix = val[1] + rot_point = [float(val_in) for val_in in val[2]] + if len(rot_point)==2: + rot_axis, rot_point = self.compute_3D_rotation(dom_space,rot_point) + rot_angles.append(float(atan2(rot_matrix[0, 1, 0], rot_matrix[0, 0, 0]))) + rot_axes.append(rot_axis) + rot_points.append(rot_point) + else: + theta_x = float(atan2(rot_matrix[2, 1], rot_matrix[2, 2])) + theta_y = float(atan2(-rot_matrix[2, 0], sqrt((rot_matrix[2, 1] ** 2) + (rot_matrix[2, 2] ** 2)))) + theta_z = float(atan2(rot_matrix[1, 0], rot_matrix[0, 0])) + rot_angles.extend([theta_x, theta_y, theta_z]) + rot_axes.extend(['x', 'y', 'z']) + rot_points.extend([rot_point, rot_point, rot_point]) + + is_boundary, geometry, cond, translate_vec, rotation_list = self.RecursiveDomainDictGeoWrapper(val[0],prod_factor=prod_factor,translate_vec=translate_vec,rotation_list=rotation_list) + return is_boundary, geometry, cond, translate_vec, rotation_list + + elif key=='d': + if prod_factor is None: + return self.GeometryNameMapper(domain_dict['d'],translate_vec=translate_vec,rotation_list=rotation_list) + else: + return self.ProductDomainWrapper(domain_dict,prod_factor,translate_vec=translate_vec,rotation_list=rotation_list) + + + def _geo_condition(self, sdf_func,sympy_expression,is_boundary): + """ + Wrapper function to convert a general sympy_expression to a + condition that is only valid for points on the geometry + (interior or on the boundary). + It is based on evaluating the sdf function of the geometry and + the sympy expression. + It returns a condition function that can be evaluated for + arbitrary points, but is only true for points on the geometry + that satisfy the sympy expression. + + Parameters + ---------- + sdf_func: function + signed distance function + sympy_expression: sympy expression + sympy expression to restrict the geometry + is_boundary: bool + True if domain is boundary + Returns + ------- + function + condition function + """ + def geo_condition(invar,params): + x=Symbol('x') + y=Symbol('y') + z=Symbol('z') + + f2 = sdf_func(invar,params) + # Get the number of arguments that f1 takes + num_args = len(invar) + + if num_args == 3: + # If f1 takes 3 arguments, call it with 'x', 'y', 'z' + f1 = lambdify([(x,y,z)], sympy_expression) + result = f1((invar['x'], invar['y'], invar['z'])) + elif num_args == 2: + f1 = lambdify([(x,y)], sympy_expression) + result = f1((invar['x'], invar['y'])) + else: + f1 = lambdify([x], sympy_expression) + result = f1(invar['x']) + + if is_boundary: + return np.equal(f2["sdf"],0)&result + else: + return np.greater(f2["sdf"],0)& result + return geo_condition + + def _or_condition(self,f1,f2): + """ + Wrapper function to combine two condition functions with an OR + operation. + """ + def or_condition(invar,params): + return np.logical_or(f1(invar,params),f2(invar,params)) + return or_condition + + def changeVarNames(self,domain,dim): + """ + Change variable names of TorchPhysics variables to Modulus + variable names, if variable 'x' is of higher dimension than 1. + """ + vars = list(domain.space.keys()) + if vars==['x']: + if dim ==2: + return ['x','y'] + elif dim ==3: + return ['x','y','z'] + else: + return ['x'] + else: + return vars + + + def ProductDomainWrapper(self,dom1,dom2,dom3=None,translate_vec = [0,0,0],rotation_list=[[],[],[],0]): + """ + Analyzes the elements of a decomposition of a TorchPhysics + product domain (cross product of domains) and maps the whole + product domain to a Modulus geometry. + + Parameters + ---------- + dom1: domain object + domain object from TorchPhysics + dom2: domain object + domain object from TorchPhysics + dom3: domain object, optional + domain object from TorchPhysics + + Returns + ------- + geometry object + Modulus geometry object + is_boundary: bool + True if domain is boundary + cond: sympy expression + condition to restrict on parts of geometry + """ + imported_module_3d = importlib.import_module("modulus.sym.geometry.primitives_3d") + imported_module_2d = importlib.import_module("modulus.sym.geometry.primitives_2d") + + x, y, z = Symbol("x"), Symbol("y"), Symbol("z") + cond=None + if dom3 == None: + key1 = next(iter(dom1.keys())) + val1 = dom1[key1] + key2 = next(iter(dom2.keys())) + val2 = dom2[key2] + + # sorting by domain dimension: dom1.dim>=dom2.dim ? + # 2D geometries in 3D define boundaries! + if key1 == 'd': + if key2 == 'd': + if val1.dim < val2.dim: + dom1,dom2 = dom2,dom1 + key1 = next(iter(dom1.keys())) + val1 = dom1[key1] + key2 = next(iter(dom2.keys())) + val2 = dom2[key2] + var1 = self.changeVarNames(val1,val1.dim) + var2 = self.changeVarNames(val2,val2.dim) + # Cross product of 2D domain and point / interval boundary + if (val1.dim==2) & (val2.dim==0): + if var2[0] == 'z': + if var1 == ['x','y']: + varchange = False + else: + var1=['x','y'] + varchange = True + elif var2[0] == 'x': + if var1 == ['z','y']: + varchange = False + else: + var1=['z','y'] + varchange = True + else: + if var1 == ['x','z']: + varchange = False + else: + var1=['x','z'] + varchange = True + + is_boundary = True + #only single boundary point + if hasattr(val2,'side'): + point1=float(val2.side()) + point2=point1+1 + cond=Eq(Symbol(var2[0]),Rational(point1)) + elif (type(val2).__name__=='Point'): + point1=float(val2.point()) + point2=point1+1 + cond=Eq(Symbol(var2[0]),Rational(point1)) + else: + point1=float(val2.bounding_box()[0]) + point2=float(val2.bounding_box()[1]) + cond=Or(Eq(Symbol(var2[0]),Rational(point1)),Eq(Symbol(var2[0]),Rational(point2))) + if type(val1).__name__ == 'Parallelogram': + orig=self.varChange(tuple(element.item() for element in val1.origin()),varchange) + c1=self.varChange(tuple(element.item() for element in val1.corner_1()),varchange) + c2=self.varChange(tuple(element.item() for element in val1.corner_2()),varchange) + geometry=ParallelogramCylinder(orig+(point1,),c1+(point1,),c2+(point1,),height=point2-point1).axischange(var1[0],var1[1],var2[0]) + elif type(val1).__name__ == 'Circle': + center = self.varChange(tuple(element.item() for element in val1.center()),varchange) + radius = float(val1.radius()) + height = point2-point1 + geometry=getattr(imported_module_3d,"Cylinder")(center+(point1+height/2,),radius,height).axischange(var1[0],var1[1],var2[0]) + elif type(val1).__name__ == 'Triangle': + #only Isosceles triangle with axis of symmetry parallel to y-axis + assert (val1.origin()[1]==val1.corner_1()[1]), "Symmetry axis of triangle has to be y-axis parallel!" + assert (np.linalg.norm(val1.corner_2()-val1.origin()) == np.linalg.norm(val1.corner_2()-val1.corner_1())), "Triangle not Isosceles!" + base = float(val1.corner_1()[0]-val1.origin()[0]) + height = float(np.sqrt(np.linalg.norm(val1.corner_2()-val1.corner_1())**2-(base/2)**2)) + height_prism = float(point2-point1) + center = self.varChange((float(val1.origin()[0])+base/2,float(val1.origin()[1])),varchange) + geometry=getattr(imported_module_3d,"IsoTriangularPrism")(center+(point1+height_prism/2,),base,height,height_prism).axischange(var1[0],var1[1],var2[0]) + else: + assert False, "Type of product domain not supported: "+type(val1).__name__+" * "+type(val2).__name__ + # Cross product of 2D domain and interval + elif (val1.dim==2) & (val2.dim==1): + if var2[0] == 'z': + if var1 == ['x','y']: + varchange = False + else: + var1=['x','y'] + varchange = True + elif var2[0] == 'x': + if var1 == ['z','y']: + varchange = False + else: + var1=['z','y'] + varchange = True + else: + if var1 == ['x','z']: + varchange = False + else: + var1=['x','z'] + varchange = True + + assert (type(val2).__name__ == 'Interval'), "Type of product domain not supported: Should be 2D domain * Interval" + + is_boundary = False + point1=float(val2.bounding_box()[0]) + point2=float(val2.bounding_box()[1]) + + if type(val1).__name__ == 'Parallelogram': + orig= self.varChange(tuple(element.item() for element in val1.origin()),varchange) + c1= self.varChange(tuple(element.item() for element in val1.corner_1()),varchange) + c2= self.varChange(tuple(element.item() for element in val1.corner_2()) ,varchange) + geometry=ParallelogramCylinder(orig+(point1,),c1+(point1,),c2+(point1,),height=point2-point1).axischange(var1[0],var1[1],var2[0]) + elif type(val1).__name__ == 'Circle': + center = self.varChange(tuple(element.item() for element in val1.center()),varchange) + radius = float(val1.radius()) + height = point2-point1 + geometry=getattr(imported_module_3d,"Cylinder")(center+(point1+height/2,),radius,height).axischange(var1[0],var1[1],var2[0]) + elif type(val1).__name__ == 'Triangle': + #only Isosceles triangle with axis of symmetry parallel to y-axis + assert (val1.origin()[1]==val1.corner_1()[1]),"Symmetry axis of triangle has to be y-axis parallel!" + assert (np.linalg.norm(val1.corner_2()-val1.origin()) == np.linalg.norm(val1.corner_2()-val1.corner_1())), "Triangle not Isosceles!" + base = float(val1.corner_1()[0]-val1.origin()[0]) + height = float(np.sqrt(np.linalg.norm(val1.corner_2()-val1.corner_1())**2-(base/2)**2)) + height_prism = point2-point1 + center = self.varChange((float(val1.origin()[0])+base/2,float(val1.origin()[1])),varchange) + geometry=getattr(imported_module_3d,"IsoTriangularPrism")(center+(point1+height_prism/2,),base,height,height_prism).axischange(var1[0],var1[1],var2[0]) + else: + assert False, "Type of product domain not supported: "+type(val1).__name__+" * "+type(val2).__name__ + # Cross product of 1D domain (boundary of 2D domain) and interval + elif (val1.dim==1) & (val2.dim==1): + assert ((type(val2).__name__ == 'Interval')|(type(val1).__name__ == 'Interval')), "Product domain for these types of domain not allowed: Has to be 1D domain * Interval" + is_boundary = True + if type(val2).__name__ != 'Interval': + val1,val2=val2,val1 + var1 = self.changeVarNames(val1,val1.dim) + var2 = self.changeVarNames(val2,val2.dim) + if var2[0] == 'z': + if var1 == ['x','y']: + varchange = False + else: + var1=['x','y'] + varchange = True + elif var2[0] == 'x': + if var1 == ['z','y']: + varchange = False + else: + var1=['z','y'] + varchange = True + else: + if var1 == ['x','z']: + varchange = False + else: + var1=['x','z'] + varchange = True + + point1=float(val2.bounding_box()[0]) + point2=float(val2.bounding_box()[1]) + cond=And(Symbol(var2[0])Rational(point1)) + if type(val1).__name__ == 'ParallelogramBoundary': + orig= self.varChange(tuple(element.item() for element in val1.domain.origin()),varchange) + c1= self.varChange(tuple(element.item() for element in val1.domain.corner_1()),varchange) + c2= self.varChange(tuple(element.item() for element in val1.domain.corner_2()) ,varchange) + geometry=ParallelogramCylinder(orig+(point1,),c1+(point1,),c2+(point1,),height=point2-point1).axischange(var1[0],var1[1],var2[0]) + elif type(val1).__name__ == 'CircleBoundary': + center = self.varChange(tuple(element.item() for element in val1.domain.center()),varchange) + radius = float(val1.domain.radius()) + height = point2-point1 + geometry=getattr(imported_module_3d,"Cylinder")(center+(point1+height/2,),radius,height).axischange(var1[0],var1[1],var2[0]) + elif type(val1).__name__ == 'TriangleBoundary': + #only Isosceles triangle with axis of symmetry parallel to y-axis + assert (val1.domain.origin()[1]==val1.domain.corner_1()[1]), "Symmetry axis of triangle has to be y-axis parallel!" + assert (np.linalg.norm(val1.domain.corner_2()-val1.domain.origin()) == np.linalg.norm(val1.domain.corner_2()-val1.domain.corner_1())), "Triangle not Isosceles!" + base = float(val1.domain.corner_1()[0]-val1.domain.origin()[0]) + height = float(np.sqrt(np.linalg.norm(val1.domain.corner_2()-val1.domain.corner_1())**2-(base/2)**2)) + height_prism = point2-point1 + center = self.varChange((float(val1.domain.origin()[0])+base/2,float(val1.domain.origin()[1])),varchange) + geometry=getattr(imported_module_3d,"IsoTriangularPrism")((center+(point1+height_prism/2,),base,height,height_prism).axischange(var1[0],var1[1],var2[0])) + elif type(val1).__name__ == 'Interval': + #var1 = list(val1.space.keys()) + #var2 = list(val2.space.keys()) + var1 = self.changeVarNames(val1,val1.dim) + var2 = self.changeVarNames(val2,val2.dim) + assert (((var1[0]=='x')&(var2[0]=='y'))|((var1[0]=='y')&(var2[0]=='x'))), "Only x,y as coordinates for 2D problems allowed" + + point1=float(val1.bounding_box()[0]) + point2=float(val1.bounding_box()[1]) + point3=float(val2.bounding_box()[0]) + point4=float(val2.bounding_box()[1]) + geometry=getattr(imported_module_2d,"Rectangle")((point1,point3),(point2,point4)).axischange(var1[0],var2[0],[]) + cond = None + is_boundary=False + else: + assert False, "Type of product domain not supported: "+type(val1).__name__+" * "+type(val2).__name__ + + # Cross product of 1D domain (boundary of 2D domain or Interval) and point/boundary of interval + elif (val1.dim==1) & (val2.dim==0): + is_boundary = True + if var2[0] == 'z': + if var1 == ['x','y']: + varchange = False + else: + var1=['x','y'] + varchange = True + elif var2[0] == 'x': + if var1 == ['z','y']: + varchange = False + else: + var1=['z','y'] + varchange = True + else: + if var1 == ['x','z']: + varchange = False + else: + var1=['x','z'] + varchange = True + + + if hasattr(val2,'side'): + point1=float(val2.side()) + point2=point1+1 + cond=Eq(Symbol(var2[0]),Rational(point1)) + elif (type(val2).__name__=='Point'): + point1=float(val2.point()) + point2=point1+1 + cond=Eq(Symbol(var2[0]),Rational(point1)) + else: + point1=float(val2.bounding_box()[0]) + point2=float(val2.bounding_box()[1]) + cond=Or(Eq(Symbol(var2[0]),Rational(point1)),Eq(Symbol(var2[0]),Rational(point2))) + + if type(val1).__name__ == 'ParallelogramBoundary': + orig=self.varChange(tuple(element.item() for element in val1.domain.origin()),varchange) + c1=self.varChange(tuple(element.item() for element in val1.domain.corner_1()),varchange) + c2=self.varChange(tuple(element.item() for element in val1.domain.corner_2()) ,varchange) + geometry=ParallelogramCylinder(orig+(point1,),c1+(point1,),c2+(point1,),height=point2-point1).axischange(var1[0],var1[1],var2[0]) + elif type(val1).__name__ == 'CircleBoundary': + center = self.varChange(tuple(element.item() for element in val1.domain.center()),varchange) + radius = float(val1.domain.radius()) + height = point2-point1 + geometry=getattr(imported_module_3d,"Cylinder")(center+(point1+height/2,),radius,height).axischange(var1[0],var1[1],var2[0]) + elif type(val1).__name__ == 'TriangleBoundary': + #only Isosceles triangle with axis of symmetry parallel to y-axis + assert (val1.domain.origin()[1]==val1.domain.corner_1()[1]),"Symmetry axis of triangle has to be y-axis parallel!" + assert (np.linalg.norm(val1.domain.corner_2()-val1.domain.origin()) == np.linalg.norm(val1.domain.corner_2()-val1.domain.corner_1())), "Triangle not Isosceles!" + base = float(val1.domain.corner_1()[0]-val1.domain.origin()[0]) + height = float(np.sqrt(np.linalg.norm(val1.domain.corner_2()-val1.domain.corner_1())**2-(base/2)**2)) + height_prism = point2-point1 + center = self.varChange((float(val1.domain.origin()[0])+base/2,float(val1.domain.origin()[1])),varchange) + geometry=getattr(imported_module_3d,"IsoTriangularPrism")(center+(point1+height_prism/2,),base,height,height_prism).axischange(var1[0],var1[1],var2[0]) + elif type(val1).__name__ == 'Interval': + var1 = list(val1.space.keys()) + var2 = list(val2.space.keys()) + assert ((var1[0]!='z')&(var2[0]!='z')), "Only x,y as coordinates for 2D problems allowed" + point3=float(val1.bounding_box()[0]) + point4=float(val1.bounding_box()[1]) + geometry=getattr(imported_module_2d,"Rectangle")((point3,point1),(point4,point2)).axischange(var1[0],var2[0],[]) + else: + assert False, "Type of product domain not supported: "+type(val1).__name__+" * "+type(val2).__name__ + # Cross product of two 0D domains not allowed + elif (val1.dim==0) & (val2.dim==0): + assert(False), "2D or 3D points are not allowed for sampling due to zero surface" + else: + is_boundary, geometry, cond, translate_vec, rotation_list = self.RecursiveDomainDictGeoWrapper(dom2,dom1,translate_vec=translate_vec,rotation_list=rotation_list) + else: + is_boundary, geometry, cond, translate_vec, rotation_list = self.RecursiveDomainDictGeoWrapper(dom1,dom2,translate_vec=translate_vec,rotation_list=rotation_list) + else: # dom3 != None + key1 = next(iter(dom1.keys())) + val1 = dom1[key1] + key2 = next(iter(dom2.keys())) + val2 = dom2[key2] + key3 = next(iter(dom3.keys())) + val3 = dom3[key3] + + # fully decomposed product domain + if (key1 == 'd')& (key2 == 'd')&(key3=='d'): + # I1xI2xI3 I1xI2xP/IBoundary + # Edges and 3D points are not allowed + # sort by (x,y,z)-order + vals= sorted((val1,val2,val3), key=lambda x: list(x.space.variables)[0]) + assert (np.sum([vals[0].dim,vals[1].dim,vals[2].dim])>1), "3D points or edges are not allowed due to zero surface in 3D" + vars = ['x','y','z'] + points=[] + conds=[] + is_boundary = False + for index, val in enumerate(vals): + if hasattr(val,'side'): + points.append([float(val.side()),float(val.side())+1]) + conds.append(Eq(Symbol(vars[index]),Rational(float(val.side())))) + is_boundary = True + elif (type(val).__name__=='Point'): + points.append([float(val.point()),float(val.point())+1]) + conds.append(Eq(Symbol(vars[index]),Rational(float(val.point())))) + is_boundary = True + elif (type(val).__name__=='Interval'): + points.append([float(val.bounding_box()[0]),float(val.bounding_box()[1])]) + conds.append(True) + else: # IntervalBoundary with both sides + points.append([float(val.bounding_box()[0]),float(val.bounding_box()[1])]) + conds.append(Or(Eq(Symbol(vars[index]),Rational(float(val.bounding_box()[0]))),Eq(Symbol(vars[index]),Rational(float(val.bounding_box()[1]))))) + is_boundary=True + geometry=getattr(imported_module_3d,"Box")((points[0][0],points[1][0],points[2][0]),(points[0][1],points[1][1],points[2][1])) + cond_all = And(conds[0],conds[1],conds[2]) + cond = None if cond_all == True else cond_all + # further decomposition of product domain, but only union or union boundary of 1D or 0D domains allowed, that means no further rotation and translation will be done by call of ProductDomainWrapper + else: + #sort domains: domains with key other than 'd' first + doms,keys = zip(*sorted(list(zip([dom1,dom2,dom3],[key1,key2,key3])), key=lambda x: x[1] == 'd')) + #only union or union boundary of 1D or 0D domains allowed (key='u' or 'ub') + assert((keys[0]=='u')|(keys[0]=='ub')), "Other domain operations than Union or UnionBoundary not allowed" + + is_boundary1, geometry1, cond1, translate_vec, rotation_list = self.ProductDomainWrapper(doms[0][keys[0]][0],doms[1],doms[2],translate_vec, rotation_list) + is_boundary2, geometry2, cond2, translate_vec, rotation_list = self.ProductDomainWrapper(doms[0][keys[0]][1],doms[1],doms[2],translate_vec, rotation_list) + #sampler (interior or boundary) has to be the same for both domains, union of interior and boundary is not allowed + assert is_boundary1==is_boundary2 + + if cond1==cond2: + if cond1!=None: + cond = cond1 + else: + cond =None + + elif cond1==None: + cond = cond2 + elif cond2==None: + cond=cond1 + else: + cond = self._or_condition(self._geo_condition(geometry1.sdf,cond1,is_boundary1),self._geo_condition(geometry2.sdf,cond2,is_boundary2)) + is_boundary = is_boundary1 + geometry = geometry1+geometry2 + + + # check if rotation and translation have to be done + rot_angles, rot_axes, rot_points, rot_prio = rotation_list + if rot_angles != []: + if translate_vec != [0,0,0]: + if rot_prio == 0: + for angle, axis, point in zip(rot_angles,rot_axes,rot_points): + geometry = geometry.rotate(angle,axis,point) + geometry = geometry.translate(translate_vec) + cond = self.translate_condition(cond,translate_vec) + else: + geometry = geometry.translate(translate_vec) + cond = self.translate_condition(cond,translate_vec) + for angle, axis, point in zip(rot_angles,rot_axes,rot_points): + geometry = geometry.rotate(angle,axis,point) + rotation_list = [[],[],[],0] + translate_vec = [0,0,0] + else: + for angle, axis, point in zip(rot_angles,rot_axes,rot_points): + geometry = geometry.rotate(angle,axis,point) + rotation_list = [[],[],[],0] + else: + if translate_vec != [0,0,0]: + geometry = geometry.translate(translate_vec) + cond = self.translate_condition(cond,translate_vec) + translate_vec = [0,0,0] + + return is_boundary, geometry, cond, translate_vec, rotation_list + + + + def varChange(self,val,ischange): + """ + Convert 2D numpy array to tuple of floats and change the order + of the values if ischange is True. + + Parameters + ---------- + val: numpy array + array with 2 values + ischange: bool + change the order of the values + + Returns + ------- + tuple + tuple of floats + """ + if ischange: + return (float(val[1]), float(val[0])) + else: + return (float(val[0]), float(val[1])) + + def GeometryNameMapper(self,domain,translate_vec=[0,0,0],rotation_list=[[],[],[],0]): + """ + Mapping of single TorchPhysics (decomposed) domain to Modulus + geometry + + Parameters + ---------- + domain: domain object + domain object from TorchPhysics + + Returns + ------- + geometry object + Modulus geometry object + is_boundary: bool + True if domain is boundary + cond: sympy expression + condition to restrict on parts of geometry + """ + + if 'Boundary' in type(domain).__name__: + dim = domain.dim+1 + is_boundary = True + else: + dim = domain.dim + is_boundary = False + cond = None + # if variable x with higher dimension than 1, change to x,y (,z) + vars = self.changeVarNames(domain,dim) + if dim ==1: + imported_module = importlib.import_module("modulus.sym.geometry.primitives_1d") + if type(domain).__name__ == 'Interval': # IntervalBoundary + geometry = getattr(imported_module,"Line1D")(float(domain.lower_bound()),float(domain.upper_bound())) + elif type(domain).__name__ == 'IntervalBoundary': + geometry = getattr(imported_module,"Line1D")(float(domain.domain.lower_bound()),float(domain.domain.upper_bound())) + elif type(domain).__name__ == 'IntervalSingleBoundaryPoint': + geometry = getattr(imported_module,"Line1D")(float(domain.domain.lower_bound()),float(domain.domain.upper_bound())) + cond = Eq(Symbol('x'),Rational(float(domain.side()))) + else: + assert False, "Domain type not supported" + elif dim==2: + if vars == ['x','y']: + varchange = False + y_index =1 + else: + varchange = True + y_index = 0 + imported_module = importlib.import_module("modulus.sym.geometry.primitives_2d") + if (type(domain).__name__ == 'Parallelogram'): + origin = self.varChange(domain.origin(),varchange) + corner1 = self.varChange(domain.corner_1(),varchange) + corner2 = self.varChange(domain.corner_2(),varchange) + geometry = getattr(imported_module,"Polygon")([origin,corner1, np.subtract(np.add(corner1,corner2),origin),corner2]) + elif (type(domain).__name__ == 'ParallelogramBoundary'): + origin = self.varChange(domain.domain.origin(),varchange) + corner1 = self.varChange(domain.domain.corner_1(),varchange) + corner2 = self.varChange(domain.domain.corner_2(),varchange) + geometry = getattr(imported_module,"Polygon")([origin,corner1, np.subtract(np.add(corner1,corner2),origin),corner2]) + elif (type(domain).__name__ == 'Circle'): + center = self.varChange(domain.center(),varchange) + geometry = getattr(imported_module,"Circle")(center,float(domain.radius())) + elif (type(domain).__name__ == 'CircleBoundary'): + center = self.varChange(domain.domain.center(),varchange) + geometry = getattr(imported_module,"Circle")(center,float(domain.domain.radius())) + elif (type(domain).__name__ == 'ShapelyPolygon'): + geometry = getattr(imported_module,"Polygon")(list(zip(*domain.polygon.exterior.coords.xy))[:-1]) + elif (type(domain).__name__ == 'ShapelyBoundary'): + geometry = getattr(imported_module,"Polygon")(list(zip(*domain.domain.polygon.exterior.coords.xy))[:-1]) + elif (type(domain).__name__ == 'Triangle'): + #only Isosceles triangle with axis of symmetry parallel to y-axis + assert (domain.origin()[y_index]==domain.corner_1()[y_index]),"Symmetry axis of triangle has to be y-axis parallel!" + assert (np.linalg.norm(domain.corner_2()-domain.origin()) == np.linalg.norm(domain.corner_2()-domain.corner_1())), "Triangle not Isosceles!" + origin = self.varChange(domain.origin(),varchange) + corner1 = self.varChange(domain.corner_1(),varchange) + corner2 = self.varChange(domain.corner_2(),varchange) + base = float(corner1[0]-origin[0]) + height = float(np.sqrt(np.linalg.norm(np.subtract(corner2,corner1))**2-(base/2)**2)) + center = (float(origin[0])+base/2,float(origin[1])) + geometry=getattr(imported_module,"Triangle")(center,base,height) + elif (type(domain).__name__ == 'TriangleBoundary'): + #only Isosceles triangle with axis of symmetry parallel to y-axis + assert (domain.domain.origin()[y_index]==domain.domain.corner_1()[y_index]), "Symmetry axis of triangle has to be y-axis parallel!" + assert (np.linalg.norm(domain.domain.corner_2()-domain.domain.origin()) == np.linalg.norm(domain.domain.corner_2()-domain.domain.corner_1())), "Triangle not Isosceles!" + origin = self.varChange(domain.domain.origin(),varchange) + corner1 = self.varChange(domain.domain.corner_1(),varchange) + corner2 = self.varChange(domain.domain.corner_2(),varchange) + base = float(corner1[0]-origin[0]) + height = float(np.sqrt(np.linalg.norm(np.subtract(corner2,corner1))**2-(base/2)**2)) + center = (float(origin[0])+base/2,float(origin[1])) + geometry=getattr(imported_module,"Triangle")(center,base,height) + else: + assert (True), "Domain type not supported" + else: #dim==3 + imported_module = importlib.import_module("modulus.sym.geometry.primitives_3d") + if type(domain).__name__ == 'Sphere': + _, sorted_center = zip(*(sorted(zip(vars,domain.center())))) + sorted_center = tuple([float(elem) for elem in sorted_center]) + geometry = getattr(imported_module,"Sphere")(sorted_center,float(domain.radius())) + elif type(domain).__name__ == 'SphereBoundary': + _, sorted_center = zip(*(sorted(zip(vars,domain.domain.center())))) + sorted_center = tuple([float(elem) for elem in sorted_center]) + geometry = getattr(imported_module,"Sphere")(sorted_center,float(domain.domain.radius())) + elif type(domain).__name__ == 'TrimeshPolyhedron': + try: + geometry = getattr(importlib.import_module("modulus.sym.geometry.tessellation"),"Tessellation")(domain.mesh) + except: + raise Exception("Tessellation module only supported for Modulus docker installation due to missing pysdf installation!") + elif type(domain).__name__ == 'TrimeshBoundary': + try: + geometry = getattr(importlib.import_module("modulus.sym.geometry.tessellation"),"Tessellation")(domain.domain.mesh) + except: + raise Exception("Tessellation module only supported for Modulus docker installation due to missing pysdf installation!") + else: + assert (True), "Domain type not supported" + + # check if rotation and translation have to be done + rot_angles, rot_axes, rot_points, rot_prio = rotation_list + if rot_angles != []: + if translate_vec != [0,0,0]: + if rot_prio == 0: + for angle, axis, point in zip(rot_angles,rot_axes,rot_points): + geometry = geometry.rotate(angle,axis,point) + geometry = geometry.translate(translate_vec) + cond = self.translate_condition(cond,translate_vec) + else: + geometry = geometry.translate(translate_vec) + cond = self.translate_condition(cond,translate_vec) + for angle, axis, point in zip(rot_angles,rot_axes,rot_points): + geometry = geometry.rotate(angle,axis,point) + rotation_list = [[],[],[],0] + translate_vec = [0,0,0] + else: + for angle, axis, point in zip(rot_angles,rot_axes,rot_points): + geometry = geometry.rotate(angle,axis,point) + rotation_list = [[],[],[],0] + else: + if translate_vec != [0,0,0]: + geometry = geometry.translate(translate_vec) + cond = self.translate_condition(cond,translate_vec) + translate_vec = [0,0,0] + + return is_boundary, geometry, cond, translate_vec, rotation_list + + + def translate_condition(self,cond,translate_vec): + """ + Translate condition by vector + + Parameters + ---------- + cond: sympy expression + condition to translate + translate_vec: tuple of floats + translation vector + + Returns + ------- + sympy expression + translated condition + """ + if cond is not None: + return cond.subs(Symbol('x'),Symbol('x')-Rational(translate_vec[0])).subs(Symbol('y'),Symbol('y')-Rational(translate_vec[1])).subs(Symbol('z'),Symbol('z')-Rational(translate_vec[2])) + else: + return None + + def compute_3D_translate_vec(self,space,translate_vec): + """ + Compute 3D translation vector in the order "x","y","z" out of + 1D, 2D or 3D translation vector + + Parameters + ---------- + space: dict + translate_vec: tuple of floats + translation vector + + Returns + ------- + tuple of floats + translated geometry + """ + vec= [0,0,0] + if 'x' in space: + if space['x'] ==2: + vec= [translate_vec[0],translate_vec[1],0] + elif space['x'] ==3: + vec= translate_vec + else: + for ind,key in enumerate(space.keys()): + if key =='x': + vec[0]=translate_vec[ind] + elif key =='y': + vec[1]=translate_vec[ind] + elif key =='z': + vec[2]=translate_vec[ind] + else: + for ind,key in enumerate(space.keys()): + if key =='y': + vec[1]=translate_vec[ind] + elif key =='z': + vec[2]=translate_vec[ind] + + return vec + + def compute_3D_rotation(self,space,point): + """ + Compute 3D rotation point in the order "x","y","z" out of 2D + rotation point and also rotation axis + + Parameters + ---------- + space: dict + space of domain (variables "x","y" or "x","z" or "y","z" or + "x" (2D)) + point: tuple of floats (2D) + + Returns + ------- + rot_ax: str + rotation axis + rot_point: tuple of floats + rotation point + """ + rot_point = [0.0,0.0,0.0] + if 'x' in space: + if space['x'] ==2: + rot_point= [point[0],point[1],0.0] + rot_ax = 'z' + + else: + for ind,key in enumerate(space.keys()): + if key =='x': + rot_point[0]=point[ind] + elif key =='y': + rot_point[1]=point[ind] + rot_point[2]=0.0 + rot_ax = 'z' + elif key =='z': + rot_point[2]=point[ind] + rot_point[1]=0.0 + rot_ax = 'y' + else: + for ind,key in enumerate(space.keys()): + if key =='y': + rot_point[1]=point[ind] + elif key =='z': + rot_point[2]=point[ind] + rot_point[0]=0.0 + rot_ax = 'x' + + return rot_ax, rot_point + +class ParallelogramCylinder(Geometry): + """ + 3D Cylinder with Parallelogram base area perpendicular to z-axis + + Parameters + ---------- + origin, corner1,corner2 : tuples with 3 ints or floats + Three corners of the parallelogram, in the following order + + | corner_2 -------- x + | / / + | / / + | origin ----- corner_1 + height: z-coordinate of upper plane of parallelogram + parameterization : Parameterization + Parameterization of geometry. + + + """ + + def __init__(self, origin,corner1,corner2,height, parameterization=Parameterization()): + assert ((origin[2] == corner1[2]) & (origin[2] == corner2[2])), "Points must have same coordinate on normal dim:z" + + # make sympy symbols to use + x, y, z = Symbol("x"), Symbol("y"), Symbol("z") + + s1, s2 = Symbol(csg_curve_naming(0)), Symbol(csg_curve_naming(1)) + + + # surface + curve_parameterization = Parameterization({s1: (0, 1),s2: (0,1)}) + curve_parameterization = Parameterization.combine(curve_parameterization, parameterization) + + # area + N = CoordSys3D("N") + vec1 = tuple(np.subtract(corner1,origin)) + vec2 = tuple(np.subtract(corner2,origin)) + vvec1 = vec1[0]*N.i + vec1[1]*N.j + vec1[2]*N.k + vvec2 = vec2[0]*N.i + vec2[1]*N.j + vec2[2]*N.k + + corner3 = tuple(np.subtract(np.add(corner1,corner2),origin)) + + #compute cross product of the "base"-vectors for computing area and testing of right-handed system + cross_vec = vvec1.cross(vvec2) + sgn_normal = sign(cross_vec.dot(N.k)) + + area_p= sqrt(cross_vec.dot(cross_vec)) + + + bottom = SympyCurve( + functions={ + "x": origin[0]+ s1*vec1[0]+s2*vec2[0], + "y": origin[1]+ s1*vec1[1]+s2*vec2[1], + "z": origin[2], + "normal_x": 0, + "normal_y": 0, + "normal_z": -1, + }, + parameterization=curve_parameterization, + area=area_p, + ) + top = SympyCurve( + functions={ + "x": origin[0]+ s1*vec1[0]+s2*vec2[0], + "y": origin[1]+ s1*vec1[1]+s2*vec2[1], + "z": origin[2]+height, + "normal_x": 0, + "normal_y": 0, + "normal_z": 1, + }, + parameterization=curve_parameterization, + area=area_p, + ) + norm_l=np.linalg.norm([-vec2[1],vec2[0]]) + side1 = SympyCurve( + functions={ + "x": origin[0]+ s1*vec2[0], + "y": origin[1]+ s1*vec2[1], + "z": origin[2]+ s2*height, + "normal_x": -vec2[1]/norm_l*sgn_normal, + "normal_y": vec2[0]/norm_l*sgn_normal, + "normal_z": 0, + }, + parameterization=curve_parameterization, + area=np.linalg.norm(vec2)*height, + ) + norm_l=np.linalg.norm([vec1[1],-vec1[0]]) + side2 = SympyCurve( + functions={ + "x": origin[0]+ s1*vec1[0], + "y": origin[1]+ s1*vec1[1], + "z": origin[2]+ s2*height, + "normal_x": vec1[1]/norm_l*sgn_normal, + "normal_y": -vec1[0]/norm_l*sgn_normal, + "normal_z": 0, + }, + parameterization=curve_parameterization, + area=np.linalg.norm(vec1)*height, + ) + + norm_l=np.linalg.norm([-vec2[1],vec2[0]]) + side3 = SympyCurve( + functions={ + "x": origin[0]+ vec1[0]+s1*vec2[0], + "y": origin[1]+ vec1[1]+s1*vec2[1], + "z": origin[2]+ s2*height, + "normal_x": vec2[1]/norm_l*sgn_normal, + "normal_y": -vec2[0]/norm_l*sgn_normal, + "normal_z": 0, + }, + parameterization=curve_parameterization, + area=np.linalg.norm(vec2)*height, + ) + norm_l=np.linalg.norm([vec1[1],-vec1[0]]) + side4 = SympyCurve( + functions={ + "x": origin[0]+ s1*vec1[0]+vec2[0], + "y": origin[1]+ s1*vec1[1]+vec2[1], + "z": origin[2]+ s2*height, + "normal_x": -vec1[1]/norm_l*sgn_normal, + "normal_y": vec1[0]/norm_l*sgn_normal, + "normal_z": 0, + }, + parameterization=curve_parameterization, + area=np.linalg.norm(vec1)*height, + ) + + sides=[top,bottom,side1,side2,side3,side4] + + # compute SDF in 2D Parallelogram + points = [origin[0:2],corner1[0:2],corner3[0:2] ,corner2[0:2]] + # wrap points + wrapted_points = points + [points[0]] + + sdfs = [(x - wrapted_points[0][0]) ** 2 + (y - wrapted_points[0][1]) ** 2] + conds = [] + for v1, v2 in zip(wrapted_points[:-1], wrapted_points[1:]): + # sdf calculation + dx = v1[0] - v2[0] + dy = v1[1] - v2[1] + px = x - v2[0] + py = y - v2[1] + d_dot_d = dx**2 + dy**2 + p_dot_d = px * dx + py * dy + max_min = Max(Min(p_dot_d / d_dot_d, 1.0), 0.0) + vx = px - dx * max_min + vy = py - dy * max_min + sdf = vx**2 + vy**2 + sdfs.append(sdf) + + # winding calculation + cond_1 = Heaviside(y - v2[1]) + cond_2 = Heaviside(v1[1] - y) + cond_3 = Heaviside((dx * py) - (dy * px)) + all_cond = cond_1 * cond_2 * cond_3 + none_cond = (1.0 - cond_1) * (1.0 - cond_2) * (1.0 - cond_3) + cond = 1.0 - 2.0 * Min(all_cond + none_cond, 1.0) + conds.append(cond) + + # set inside outside + sdf_xy = Min(*sdfs) + cond = reduce(mul, conds) + sdf_xy = sqrt(sdf_xy) * -cond + + + #compute distance in z-direction in 3D + center_z = origin[2] + height / 2 + sdf_z = height / 2 - Abs(z - center_z) + + if (sdf_xy >=0) & (sdf_z >= 0): + sdf= Min(sdf_xy,sdf_z) + elif (sdf_xy <0)&(sdf_z>0): + sdf = sdf_xy + elif (sdf_xy>0)&sdf_z<0: + sdf = sdf_z + else: #min distance to all 12 boundary curves + sdf = -sqrt(sdf_xy**2+sdf_z**2) + + + # calculate bounds + max_x = max(origin[0],corner1[0],corner2[0],corner3[0]) + min_x = min(origin[0],corner1[0],corner2[0],corner3[0]) + max_y = max(origin[1],corner1[1],corner2[1],corner3[1]) + min_y = min(origin[1],corner1[1],corner2[1],corner3[1]) + + bounds = Bounds( + { + Parameter("x"): (min_x, max_x), + Parameter("y"): (min_y, max_y), + Parameter("z"): (float(origin[2]), float(origin[2]+height)), + }, + parameterization=parameterization, + ) + + # initialize ParallelogramCylinder + super().__init__( + sides, + _sympy_sdf_to_sdf(sdf), + dims=3, + bounds=bounds, + parameterization=parameterization, + ) diff --git a/src/torchphysics/wrapper/helper.py b/src/torchphysics/wrapper/helper.py new file mode 100644 index 00000000..52271167 --- /dev/null +++ b/src/torchphysics/wrapper/helper.py @@ -0,0 +1,653 @@ +# This file includes work covered by the following copyright and permission notices: +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# All rights reserved. +# Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import numpy as np + +from torchphysics.problem.spaces.points import Points +from torchphysics.utils.plotting.plot_functions import plot +from modulus.sym.loss import Loss +from modulus.sym.utils.io import InferencerPlotter + +from modulus.sym.domain.validator import Validator +from modulus.sym.node import Node + +from modulus.sym.dataset import DictInferencePointwiseDataset, DictPointwiseDataset + +from modulus.sym.domain.constraint import Constraint +from modulus.sym.graph import Graph +from modulus.sym.key import Key +from modulus.sym.node import Node +from modulus.sym.constants import TF_SUMMARY +from modulus.sym.distributed import DistributedManager +from modulus.sym.utils.io.vtk import var_to_polyvtk + +from typing import List, Dict +from typing import Dict +from torch import Tensor + +def convertDataTP2Modulus(points): + """ + Convert data from TorchPhysics to Modulus format + + Parameters + ---------- + + data: Points (torchphysics.problem.space.Points) + Dictionary containing the data points of the TorchPhysics + dataset with the TorchPhysics variable names as keys and the + corresponding data points as values. + + """ + data=points.coordinates + + outdata={} + + for var in points.space.variables: + if var =='x': + outdata['x'] = data[var][:,0].unsqueeze(1) + if points.space[var] > 1: + outdata['y'] = data[var][:,1].unsqueeze(1) + if points.space[var]== 3: + outdata['z'] = data[var][:,2].unsqueeze(1) + + else: + if points.space[var]>1: + for ind in range(points.space[var]): + outdata[var+str(ind+1)]=data[var][:,ind].unsqueeze(1) + else: + outdata[var]=data[var] + + return outdata + + +def convertDataModulus2TP(data,TP_space): + """ + Convert data from Modulus to TorchPhysics format. If a key of + TP_space is not present in data, it will be ignored. + + Parameters + ---------- + + data: dict[str,torch.Tensor] + Dictionary containing the data points of the Modulus dataset + with the Modulus variable names as keys and the corresponding + data points as values + + """ + outdata={} + + + for key in TP_space.variables: + if TP_space[key] > 1: + if key !='x': + if all([key+str(l+1) in data.keys() for l in range(TP_space[key])]): + cat_var = list(data[key+str(l+1)] for l in range(TP_space[key])) + outdata[key] = torch.cat(cat_var,dim=1) + else: + if TP_space[key] ==2: + if ('x' in data.keys()) and ('y' in data.keys()): + outdata['x']=torch.cat((data['x'],data['y']),dim=1) + else: + if ('x' in data.keys()) and ('y' in data.keys()) and ('z' in data.keys()): + outdata['x']=torch.cat((data['x'],data['y'],data['z']),dim=1) + + else: + if key in data.keys(): + outdata[key]=data[key] + + + return outdata + + +class PointwiseLossInfNorm(Loss): + """ + L-inf loss function for pointwise data + Computes the maximum norm loss of each output tensor + """ + + def __init__(self): + super().__init__() + + + @staticmethod + def _loss( + invar: Dict[str, Tensor], + pred_outvar: Dict[str, Tensor], + true_outvar: Dict[str, Tensor], + lambda_weighting: Dict[str, Tensor], + step: int, + ) -> Dict[str, Tensor]: + losses = {} + for key, value in pred_outvar.items(): + l = lambda_weighting[key] * torch.abs( + pred_outvar[key] - true_outvar[key] + ) + if "area" in invar.keys(): + l *= invar["area"] + losses[key] = torch.max(l) + return losses + + def forward( + self, + invar: Dict[str, Tensor], + pred_outvar: Dict[str, Tensor], + true_outvar: Dict[str, Tensor], + lambda_weighting: Dict[str, Tensor], + step: int, + ) -> Dict[str, Tensor]: + return PointwiseLossInfNorm._loss( + invar, pred_outvar, true_outvar, lambda_weighting, step + ) + +class PointwiseLossMean(Loss): + """ + Computes the mean of loss function values + + """ + + def __init__(self): + super().__init__() + + + @staticmethod + def _loss( + invar: Dict[str, Tensor], + pred_outvar: Dict[str, Tensor], + true_outvar: Dict[str, Tensor], + lambda_weighting: Dict[str, Tensor], + step: int, + ) -> Dict[str, Tensor]: + losses = {} + for key, value in pred_outvar.items(): + losses[key]= torch.mean(lambda_weighting[key] * torch.abs(pred_outvar[key] - true_outvar[key])) + return losses + + def forward( + self, + invar: Dict[str, Tensor], + pred_outvar: Dict[str, Tensor], + true_outvar: Dict[str, Tensor], + lambda_weighting: Dict[str, Tensor], + step: int, + ) -> Dict[str, Tensor]: + return PointwiseLossMean._loss( + invar, pred_outvar, true_outvar, lambda_weighting, step + ) + + + + +def OptimizerNameMapper(tp_name): + """ + Maps the optimizer name to intern Modulus names. If the optimizer + is not defined, it returns 'not defined'. + """ + if tp_name == 'Adam': + return 'adam' + elif tp_name == 'LBFGS': + return 'bfgs' + elif tp_name == 'SGD': + return 'sgd' + elif tp_name == 'Adahessian': + return 'adahessian' + elif tp_name == 'Adadelta': + return 'adadelta' + elif tp_name == 'Adagrad': + return 'adagrad' + elif tp_name == 'AdamW': + return 'adamw' + elif tp_name == 'SparseAdam': + return 'sparse_adam' + elif tp_name == 'Adamax': + return 'adamax' + elif tp_name == 'ASGD': + return 'asgd' + elif tp_name == 'NAdam': + return 'Nadam' + elif tp_name == 'RAdam': + return 'Radam' + elif tp_name == 'RMSprop': + return 'rmsprop' + elif tp_name == 'Rprop': + return 'rprop' + elif tp_name == 'A2GradExp': + return 'a2grad_exp' + elif tp_name == 'A2GradInc': + return 'a2grad_inc' + elif tp_name == 'A2GradUni': + return 'a2grad_uni' + elif tp_name == 'AccSGD': + return 'accsgd' + elif tp_name == 'AdaBelief': + return 'adabelief' + elif tp_name == 'AdaBound': + return 'adabound' + elif tp_name == 'AdaMod': + return 'adamod' + elif tp_name == 'AdaFactor': + return 'adafactor' + elif tp_name == 'AdamP': + return 'adamp' + elif tp_name == 'AggMo': + return 'aggmo' + elif tp_name == 'Apollo': + return 'apollo' + elif tp_name == 'DiffGrad': + return 'diffgrad' + elif tp_name == 'Lamb': + return 'lamb' + elif tp_name == 'NovoGrad': + return 'novograd' + elif tp_name == 'PID': + return 'pid' + elif tp_name == 'QHAdam': + return 'qhadam' + elif tp_name == 'MADGRAD': + return 'madgrad' + elif tp_name == 'QHM': + return 'qhm' + elif tp_name == 'Ranger': + return 'ranger' + elif tp_name == 'RangerQH': + return 'ranger_qh' + elif tp_name == 'RangerVA': + return 'ranger_va' + elif tp_name == 'SGDW': + return 'sgdw' + elif tp_name == 'SGDP': + return 'sgdp' + elif tp_name == 'SWATS': + return 'swats' + elif tp_name == 'Shampoo': + return 'shampoo' + elif tp_name == 'Yogi': + return 'yogi' + else: + return 'not defined' + +def SchedulerNameMapper(tp_name): + """ + Maps the scheduler name to intern Modulus names. If the scheduler + is not defined, it returns 'not defined'. + """ + + if tp_name == 'ExponentialLR': + return 'exponential_lr' + elif tp_name == 'TFExponentialLR': + return 'tf_exponential_lr' + elif tp_name == 'CosineAnnealingLR': + return 'cosine_annealing' + elif tp_name == 'CosineAnnealingWarmRestarts': + return 'cosine_annealing_warm_restarts' + else: + return 'not defined' + +def AggregatorNameMapper(tp_name): + """ + Maps the aggregator name to intern Modulus names. If the aggregator + is not defined, it returns 'not defined'. + """ + if tp_name == 'Sum' : + return 'sum' + elif tp_name == 'GradNorm': + return 'grad_norm' + elif tp_name == 'ResNorm': + return 'res_norm' + elif tp_name == 'Homoscedastic': + return 'homoscedastic' + elif tp_name == 'LRAnnealing': + return 'lr_annealing' + elif tp_name == 'SoftAdapt': + return 'soft_adapt' + elif tp_name == 'Relobralo': + return 'relobralo' + else: + return 'not defined' + +class CustomInferencerPlotter(InferencerPlotter): + """ + Custom inferencer plotter class for using TorchPhysics plot + callbacks + """ + def __init__(self,callback): + self.plot_function = callback.plot_function + self.plot_type = callback.plot_type + self.angle = callback.angle + self.point_sampler = callback.point_sampler + self.kwargs = callback.kwargs + self.model = callback.model + + def __call__(self, invar, outvar): + fig = plot(model=self.model, plot_function=self.plot_function, + point_sampler=self.point_sampler, + angle=self.angle, plot_type=self.plot_type, + device=next(self.model.parameters()).device, **self.kwargs) + + return [(fig,'')] + + + + +class PINNConditionValidator(Validator): + """ + Pointwise Validator that allows validating output variables of the + Graph on pointwise data. + + Parameters + ---------- + nodes : List[Node] + List of Modulus Nodes to unroll graph with. + invar : Dict[str, np.ndarray (N, 1)] + Dictionary of numpy arrays as input. + output_names : List[str] + List of desired outputs. + batch_size : int, optional + Batch size used when running validation, by default 1024 + requires_grad : bool = False + If automatic differentiation is needed for computing results. + """ + + def __init__( + self, + nodes: List[Node], + invar: Dict[str, np.array], + output_names: List[str], + batch_size: int = 1024, + requires_grad: bool = False, + ): + + # get dataset and dataloader + self.dataset = DictInferencePointwiseDataset( + invar=invar, output_names=output_names + ) + self.dataloader = Constraint.get_dataloader( + dataset=self.dataset, + batch_size=batch_size, + shuffle=False, + drop_last=False, + num_workers=0, + distributed=False, + infinite=False, + ) + + + # construct model from nodes + self.model = Graph( + nodes, + Key.convert_list(self.dataset.invar_keys), + Key.convert_list(self.dataset.outvar_keys), + ) + + self.manager = DistributedManager() + self.device = self.manager.device + self.model.to(self.device) + + # set forward method + self.requires_grad = requires_grad + self.forward = self.forward_grad if requires_grad else self.forward_nograd + + # set plotter + self.plotter = None + + + + def save_results(self, name, results_dir, writer, save_filetypes, step): + + invar_cpu = {key: [] for key in self.dataset.invar_keys} + #true_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} + pred_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} + # Loop through mini-batches + for i, (invar0,) in enumerate(self.dataloader): + # Move data to device (may need gradients in future, if so requires_grad=True) + invar = Constraint._set_device( + invar0, device=self.device, requires_grad=self.requires_grad + ) + + pred_outvar = self.forward(invar) + + # Collect minibatch info into cpu dictionaries + invar_cpu = { + key: value + [invar[key].cpu().detach()] + for key, value in invar_cpu.items() + } + + pred_outvar_cpu = { + key: value + [pred_outvar[key].cpu().detach()] + for key, value in pred_outvar_cpu.items() + } + + # Concat mini-batch tensors + invar_cpu = {key: torch.cat(value) for key, value in invar_cpu.items()} + + pred_outvar_cpu = { + key: torch.cat(value) for key, value in pred_outvar_cpu.items() + } + # compute losses on cpu + # TODO add metrics specific for validation + # TODO: add potential support for lambda_weighting + #losses = PointwiseValidator._l2_relative_error(true_outvar_cpu, pred_outvar_cpu) + losses = {key: torch.mean(value) for key, value in pred_outvar_cpu.items()} + + # convert to numpy arrays + invar = {k: v.numpy() for k, v in invar_cpu.items()} + pred_outvar = {k: v.numpy() for k, v in pred_outvar_cpu.items()} + + # save batch to vtk file TODO clean this up after graph unroll stuff + + named_pred_outvar = {"pred_" + k: v for k, v in pred_outvar.items()} + + # save batch to vtk/npz file TODO clean this up after graph unroll stuff + if "np" in save_filetypes: + np.savez( + results_dir + name, {**invar, **named_pred_outvar} + ) + if "vtk" in save_filetypes: + var_to_polyvtk( + {**invar, **named_pred_outvar}, results_dir + name + ) + + # add tensorboard plots + if self.plotter is not None: + self.plotter._add_figures( + "Validators", + name, + results_dir, + writer, + step, + invar, + pred_outvar, + ) + + # add tensorboard scalars + for k, loss in losses.items(): + if TF_SUMMARY: + writer.add_scalar("val/" + name + "/" + k, loss, step, new_style=True) + else: + writer.add_scalar( + "Validators/" + name , loss, step, new_style=True + ) + return losses + + +class DataConditionValidator(Validator): + """ + Validator that allows validating on pointwise data. + The validation error is the cumulative norm over all output + variables. The norm can be defined by the user. + + Parameters + ---------- + nodes : List[Node] + List of Modulus Nodes to unroll graph with. + invar : Dict[str, np.ndarray (N, 1)] + Dictionary of numpy arrays as input. + true_outvar : Dict[str, np.ndarray (N, 1)] + Dictionary of numpy arrays used to validate against validation. + batch_size : int, optional + Batch size used when running validation, by default 1024 + requires_grad : bool = False + If automatic differentiation is needed for computing results. + norm: int or 'inf', optional + The 'norm' which should be computed for evaluation. If 'inf', + maximum norm will be used. Else, the result will be taken to + the n-th potency (without computing the root!) + root: int, optional + The root of the norm. If norm is 'inf', this parameter will be + ignored. + """ + + def __init__( + self, + nodes: List[Node], + invar: Dict[str, np.array], + true_outvar: Dict[str, np.array], + batch_size: int = 1024, + requires_grad: bool = False, + norm: int = 2, + root: int = 1, + ): + + # get dataset and dataloader + self.dataset = DictPointwiseDataset(invar=invar, outvar=true_outvar) + self.dataloader = Constraint.get_dataloader( + dataset=self.dataset, + batch_size=batch_size, + shuffle=False, + drop_last=False, + num_workers=0, + distributed=False, + infinite=False, + ) + + # construct model from nodes + self.model = Graph( + nodes, + Key.convert_list(self.dataset.invar_keys), + Key.convert_list(self.dataset.outvar_keys), + ) + self.manager = DistributedManager() + self.device = self.manager.device + self.model.to(self.device) + + # set forward method + self.requires_grad = requires_grad + self.forward = self.forward_grad if requires_grad else self.forward_nograd + + # set plotter + self.plotter = None + + self.norm = norm + self.root = root + + def save_results(self, name, results_dir, writer, save_filetypes, step): + + invar_cpu = {key: [] for key in self.dataset.invar_keys} + true_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} + pred_outvar_cpu = {key: [] for key in self.dataset.outvar_keys} + # Loop through mini-batches + for i, (invar0, true_outvar0, lambda_weighting) in enumerate(self.dataloader): + # Move data to device (may need gradients in future, if so requires_grad=True) + invar = Constraint._set_device( + invar0, device=self.device, requires_grad=self.requires_grad + ) + true_outvar = Constraint._set_device( + true_outvar0, device=self.device, requires_grad=self.requires_grad + ) + pred_outvar = self.forward(invar) + + # Collect minibatch info into cpu dictionaries + invar_cpu = { + key: value + [invar[key].cpu().detach()] + for key, value in invar_cpu.items() + } + true_outvar_cpu = { + key: value + [true_outvar[key].cpu().detach()] + for key, value in true_outvar_cpu.items() + } + pred_outvar_cpu = { + key: value + [pred_outvar[key].cpu().detach()] + for key, value in pred_outvar_cpu.items() + } + + # Concat mini-batch tensors + invar_cpu = {key: torch.cat(value) for key, value in invar_cpu.items()} + true_outvar_cpu = { + key: torch.cat(value) for key, value in true_outvar_cpu.items() + } + pred_outvar_cpu = { + key: torch.cat(value) for key, value in pred_outvar_cpu.items() + } + # compute losses on cpu + # TODO add metrics specific for validation + # TODO: add potential support for lambda_weighting + #losses = PointwiseValidator._l2_relative_error(true_outvar_cpu, pred_outvar_cpu) + loss=torch.zeros(1) + for key in true_outvar_cpu.keys(): + if self.norm == 'inf': + loss = torch.max(loss,torch.max(torch.abs(true_outvar_cpu[key] - pred_outvar_cpu[key]))) + else: + loss += torch.mean(torch.abs(true_outvar_cpu[key] - pred_outvar_cpu[key])**self.norm) + + losses ={'': loss**1/self.root} + + + + # convert to numpy arrays + invar = {k: v.numpy() for k, v in invar_cpu.items()} + true_outvar = {k: v.numpy() for k, v in true_outvar_cpu.items()} + pred_outvar = {k: v.numpy() for k, v in pred_outvar_cpu.items()} + + # save batch to vtk file TODO clean this up after graph unroll stuff + named_true_outvar = {"true_" + k: v for k, v in true_outvar.items()} + named_pred_outvar = {"pred_" + k: v for k, v in pred_outvar.items()} + + # save batch to vtk/npz file TODO clean this up after graph unroll stuff + if "np" in save_filetypes: + np.savez( + results_dir + name, {**invar, **named_true_outvar, **named_pred_outvar} + ) + if "vtk" in save_filetypes: + var_to_polyvtk( + {**invar, **named_true_outvar, **named_pred_outvar}, results_dir + name + ) + + # add tensorboard plots + if self.plotter is not None: + self.plotter._add_figures( + "Validators", + name, + results_dir, + writer, + step, + invar, + true_outvar, + pred_outvar, + ) + + # add tensorboard scalars + for k, loss in losses.items(): + if TF_SUMMARY: + writer.add_scalar("val/" + name + "/" + k, loss, step, new_style=True) + else: + writer.add_scalar( + "Validators/" + name, loss, step, new_style=True + ) + return losses + \ No newline at end of file diff --git a/src/torchphysics/wrapper/model.py b/src/torchphysics/wrapper/model.py new file mode 100644 index 00000000..965f3653 --- /dev/null +++ b/src/torchphysics/wrapper/model.py @@ -0,0 +1,383 @@ +from torchphysics.models.model import Model +from torchphysics.problem.spaces import Points + +from modulus.sym.hydra.utils import compose +from modulus.sym.hydra import instantiate_arch +from modulus.sym.models.arch import Arch +from modulus.sym.key import Key + +from typing import Dict +from torch import Tensor +import torch + +import logging +import os + +class TPModelArch(Arch): + """ + Wrapper class to convert instances of TorchPhysics model base class + getting all necessary attributes of Modulus Arch models + + Parameters + ---------- + model: TorchPhysics model object + The model object to be converted into Modulus model object + + """ + + def __init__( self,model) -> None: + input_keys, output_keys = self.getVarKeys(model) + super().__init__( + input_keys=input_keys, + output_keys=output_keys, + detach_keys=[], + periodicity=None, + ) + + self._impl = model + + def forward(self, in_vars: Dict[str, Tensor]) -> Dict[str, Tensor]: + conv_in_vars_2 = self.convertSpatialVariables(in_vars) + y=self._tensor_forward(conv_in_vars_2) + + res= self.split_output(y, self.output_key_dict, dim=-1) + return res + + def _tensor_forward(self, x: Tensor) -> Tensor: + # there were an error message coming out of an x-dictionary with tensor.float64 type entries + # so is converted to normal float here + x_t = {k: v.float() for k, v in Points.from_coordinates(x).coordinates.items()} + points = Points.from_coordinates(x_t) + y = self._impl(points) + return y + + + def convertSpatialVariables(self,vars): + conv_vars = {} + for key in self._impl.input_space.keys(): + if key =='x': + if self._impl.input_space['x'] > 1: + if self._impl.input_space['x'] ==2: + conv_vars['x']=torch.cat((vars['x'],vars['y']),dim=1) + else: + conv_vars['x']=torch.cat((vars['x'],vars['y'],vars['z']),dim=1) + else: + conv_vars['x']=vars['x'] + else: + if self._impl.input_space[key] >1: + cat_var = list(vars[key+str(l+1)] for l in range(self._impl.input_space[key])) + conv_vars[key] = torch.cat(cat_var,dim=1) + else: + conv_vars[key]=vars[key] + return conv_vars + + + def getVarKeys(self,model): + inputkeys = [] + outputkeys = [] + for key in model.input_space.keys(): + if key =='x': + inputkeys.append(Key('x')) + if model.input_space['x'] > 1: + inputkeys.append(Key('y')) + if model.input_space['x']== 3: + inputkeys.append(Key('z')) + else: + if model.input_space[key] >1: + for l in range(model.input_space[key]): + inputkeys.append(Key(key+str(l+1))) + else: + inputkeys.append(Key(key)) + for key in model.output_space.keys(): + if model.output_space[key] >1: + for l in range(model.output_space[key]): + outputkeys.append(Key(key+str(l+1))) + else: + outputkeys.append(Key(key)) + + return inputkeys, outputkeys + + def getInputSpace(self): + return self._impl.input_space + + + + + +class ModulusArchitectureWrapper(Model): + """ + Wrapper class to use all model architectures implemented in Modulus + library as TorchPhysics model. The chosen Modulus architecture is + defined by the arch_name parameter and optional architecture + specific parameters can be passed as keyword arguments. + The model then gets all necessary attributes of TorchPhysics model + base class. The input points are converted to the Modulus input + format and the Modulus output points back to the TorchPhysics + output format. + + Parameters + ---------- + input_space : Space + The space of the points the can be put into this model. + output_space : Space + The space of the points returned by this model. + arch_name : {"afno","distributed_afno","deeponet","fno","fourier", + "fully_connected","conv_fully_connected", + "fused_fully_connected","fused_fourier","fused_hash_encoding", + "hash_encoding","highway_fourier","modified_fourier", + "multiplicative_fourier","multiscale_fourier","pix2pix", + "siren","super_res"} + Name of the Modulus architecture. + **kwargs : optional + Additional keyword arguments, depending on the chosen Modulus + architecture - listed with default values: + "afno": + img_shape: Tuple[int] = MISSING + patch_size: int = 16 + embed_dim: int = 256 + depth: int = 4 + num_blocks: int = 8 + "distributed_afno": + img_shape: Tuple[int] = MISSING + patch_size: int = 16 + embed_dim: int = 256 + depth: int = 4 + num_blocks: int = 8 + channel_parallel_inputs: bool = False + channel_parallel_outputs: bool = False + "deeponet": + trunk_dim: Any = None # Union[None, int] + branch_dim: Any = None # Union[None, int] + "fno": + dimension: int = MISSING + decoder_net: Arch + nr_fno_layers: int = 4 + fno_modes: Any = 16 # Union[int, List[int]] + padding: int = 8 + padding_type: str = "constant" + activation_fn: str = "gelu" + coord_features: bool = True + "fourier": + frequencies: Any = "('axis', [0, 1, 2, 3, 4, 5, 6, 7, + 8, 9])" + frequencies_params: Any = "('axis', [0, 1, 2, 3, 4, 5, + 6, 7, 8, 9])" + activation_fn: str = "silu" + layer_size: int = 512 + nr_layers: int = 6 + skip_connections: bool = False + weight_norm: bool = True + adaptive_activations: bool = False + "fully_connected": + layer_size: int = 512 + nr_layers: int = 6 + skip_connections: bool = False + activation_fn: str = "silu" + adaptive_activations: bool = False + weight_norm: bool = True + "conv_fully_connected": + layer_size: int = 512 + nr_layers: int = 6 + skip_connections: bool = False + activation_fn: str = "silu" + adaptive_activations: bool = False + weight_norm: bool = True + "fused_fully_connected": + layer_size: int = 128 + nr_layers: int = 6 + activation_fn: str = "sigmoid" + "fused_fourier": + layer_size: int = 128 + nr_layers: int = 6 + activation_fn: str = "sigmoid" + n_frequencies: int = 12 + "fused_hash_encoding": + layer_size: int = 128 + nr_layers: int = 6 + activation_fn: str = "sigmoid" + indexing: str = "Hash" + n_levels: int = 16 + n_features_per_level: int = 2 + log2_hashmap_size: int = 19 + base_resolution: int = 16 + per_level_scale: float = 2.0 + interpolation: str = "Smoothstep" + "hash_encoding": + layer_size: int = 64 + nr_layers: int = 3 + skip_connections: bool = False + weight_norm: bool = True + adaptive_activations: bool = False + bounds: Any = "[(1.0, 1.0), (1.0, 1.0)]" + nr_levels: int = 16 + nr_features_per_level: int = 2 + log2_hashmap_size: int = 19 + base_resolution: int = 2 + finest_resolution: int = 32 + "highway_fourier": + frequencies: Any = "('axis', [0, 1, 2, 3, 4, 5, 6, 7, + 8, 9])" + frequencies_params: Any = "('axis', [0, 1, 2, 3, 4, + 5, 6, 7, 8, 9])" + activation_fn: str = "silu" + layer_size: int = 512 + nr_layers: int = 6 + skip_connections: bool = False + weight_norm: bool = True + adaptive_activations: bool = False + transform_fourier_features: bool = True + project_fourier_features: bool = False + "modified_fourier": + frequencies: Any = "('axis', [0, 1, 2, 3, 4, 5, 6, 7, + 8, 9])" + frequencies_params: Any = "('axis', [0, 1, 2, 3, 4, + 5, 6, 7, 8, 9])" + activation_fn: str = "silu" + layer_size: int = 512 + nr_layers: int = 6 + skip_connections: bool = False + weight_norm: bool = True + adaptive_activations: bool = False + "multiplicative_fourier": + layer_size: int = 512 + nr_layers: int = 6 + skip_connections: bool = False + activation_fn: str = "identity" + filter_type: str = "fourier" + weight_norm: bool = True + input_scale: float = 10.0 + gabor_alpha: float = 6.0 + gabor_beta: float = 1.0 + normalization: Any = (None # Change to Union[None, + Dict[str, Tuple[float, float]]] when supported) + "multiscale_fourier": + frequencies: Any = field(default_factory=lambda: [32]) + frequencies_params: Any = None + activation_fn: str = "silu" + layer_size: int = 512 + nr_layers: int = 6 + skip_connections: bool = False + weight_norm: bool = True + adaptive_activations: bool = False + "pix2pix": + dimension: int = MISSING + conv_layer_size: int = 64 + n_downsampling: int = 3 + n_blocks: int = 3 + scaling_factor: int = 1 + batch_norm: bool = True + padding_type: str = "reflect" + activation_fn: str = "relu" + "siren": + layer_size: int = 512 + nr_layers: int = 6 + first_omega: float = 30.0 + omega: float = 30.0 + normalization: Any = (None # Change to Union[None, + Dict[str, Tuple[float, float]]] when supported) + "super_res": + large_kernel_size: int = 7 + small_kernel_size: int = 3 + conv_layer_size: int = 32 + n_resid_blocks: int = 8 + scaling_factor: int = 8 + activation_fn: str = "prelu" + + + Examples + -------- + >>> testmodel=ModulusArchitectureWrapper(input_space=X*T, + output_space=U,arch_name='fully_connected',layer_size=30, nr_layers=3) + >>> testmodel=ModulusArchitectureWrapper(input_space=X*T, + output_space=U,arch_name='fourier',frequencies = ['axis',[0,1,2]]) + + """ + def __init__(self,input_space, output_space,arch_name,**kwargs): + # get the absolute path of the conf-directory + caller_path=os.path.abspath(os.getcwd()+'/conf') + os.makedirs(caller_path, exist_ok=True) + # Get the relative path of the current file to the conf-directory + current_path = os.path.relpath(caller_path,os.path.dirname(os.path.abspath(__file__))) + with open(caller_path+'/config_model.yaml', 'w') as f: + f.write('defaults :\n - modulus_default\n - arch:\n - '+arch_name) + cfg3 = compose(config_path=current_path, config_name="config_model") + + #from modulus.sym.models.arch import arch_name + inputkeys = [] + input_keys = [] + for key in input_space.keys(): + if key =='x': + inputkeys.append(Key('x')) + input_keys.append('x') + + if input_space['x'] > 1: + inputkeys.append(Key('y')) + input_keys.append('y') + if input_space['x']== 3: + inputkeys.append(Key('z')) + input_keys.append('z') + else: + if input_space[key] >1: + for l in range(input_space[key]): + inputkeys.append(Key(key+str(l+1))) + input_keys.append(key+str(l+1)) + else: + inputkeys.append(Key(key)) + input_keys.append(key) + + outputkeys = [] + for key in output_space.keys(): + if output_space[key] >1: + for l in range(output_space[key]): + outputkeys.append(Key(key+str(l+1))) + else: + outputkeys.append(Key(key)) + + super().__init__(input_space,output_space) + + self.output_space = output_space + self.input_keys = input_keys + hydra_conf = eval(f"cfg3.arch.{arch_name}") + + for key, value in kwargs.items(): + hydra_conf[key] = value + + self.modulus_net = instantiate_arch(input_keys=inputkeys,output_keys=outputkeys,cfg = hydra_conf) + + def forward(self, points): + num_points = len(points) + dim = len(self.input_keys) + # values of TP input points should get the same order as the keys of the input_space + points = self._fix_points_order(points) + # ordered values get the corresponding Modulus input keys (if some of TP input keys are of higher dimension than 1, they are split into several Modulus input keys) + points_modulus = dict(zip(self.input_keys,[points.as_tensor[:,index].reshape(num_points,1) for index in range(0,dim)])) + output_modulus = self.modulus_net(points_modulus) + for key in self.output_space.keys(): + if self.output_space[key] >1: + cat_out = [output_modulus[key+str(l+1)] for l in range(self.output_space[key])] + output_modulus[key] = torch.cat(cat_out,dim=1) + for l in range(self.output_space[key]): + output_modulus.pop(key+str(l+1)) + + return Points.from_coordinates(output_modulus) + + +class ModulusNetWrapper(Model): + """ + Wrapper to convert objects of Modulus base class Arch into + TorchPhysics models + + Parameters + ---------- + arch: Modulus Arch object + + """ + def __init__(self,arch): + input_space = torchphysics.problem.spaces.Space(arch.input_key_dict) + output_space = torchphysics.problem.spaces.Space(arch.output_key_dict) + + super().__init__(input_space,output_space) + + + + \ No newline at end of file diff --git a/src/torchphysics/wrapper/nodes.py b/src/torchphysics/wrapper/nodes.py new file mode 100644 index 00000000..fd0f3e79 --- /dev/null +++ b/src/torchphysics/wrapper/nodes.py @@ -0,0 +1,228 @@ +import inspect +from typing import Dict +from torch import Tensor +import torch.nn as nn +import torch +import numpy as np +import logging + + +class TPNodeFunction(nn.Module): + """ + Module to evaluates a given TorchPhysics function (objective) with + the given input variables. + In the forward call it converts the input variables to the format + of the TP objective function and evaluates it. + It returns the result of the objective function as a dictionary + with the name of the objective function as key and the norm of the + result as value. + + Parameters + ---------- + tpfunction : callable + The TorchPhysics function to evaluate. + input_space : dict + The input space of the objective function. + output_space : dict + The output space of the objective function. + data_functions : dict + The data functions of the objective function. + params : list + The parameters of the objective function. + + """ + + def __init__(self,tpfunction,input_space,output_space,data_functions,params): + super().__init__() + self.objective = tpfunction + self.input_space = input_space + self.output_space = output_space + self.data_functions = data_functions + self.params = params + variables = set(inspect.signature(self.objective.fun).parameters) + input_keys, output_keys = self.getInOutputVariables() + parameter_variables=variables-self.input_space.keys()-self.output_space.keys() + + parameter_variables_new = set([]) + for param in parameter_variables: + if param in params[0].space.variables: + if params[0].space.dim>1: + parameter_variables_new = parameter_variables_new | set([param+str(l) for l in range(params[0].space.dim)]) + else: + parameter_variables_new = parameter_variables_new | set(param) + self.parameter_variables=parameter_variables_new + self.variables = ((parameter_variables_new)|input_keys|output_keys)-set(self.data_functions.keys()) + self.cond_names = [tpfunction.fun.__name__] + + def forward(self, in_vars: Dict[str, Tensor]) -> Dict[str, Tensor]: + invars = self.convertVariables(in_vars) + objectives = self.objective(invars) + res = {self.cond_names[0]: torch.norm(objectives,dim=1)} + return res + + def getInOutputVariables(self): + ''' + Returns the input keys, output keys and spatial keys of the + objective function. + ''' + + spatial_keys = set([]) + inputkeys = set([]) + outputkeys = set([]) + + if {'x','y','z'} <= set(self.input_space.keys()): + spatial_keys = {'x','y','z'} + elif {'x','y'} <= set(self.input_space.keys()): + spatial_keys = {'x','y'} + elif {'x'} <= set(self.input_space.keys()): + if self.input_space['x'] > 1: + if self.input_space['x'] ==2: + spatial_keys = {'x','y'} + else: + spatial_keys = {'x','y','z'} + else: + spatial_keys = {'x'} + + for key in set(self.input_space.keys())- {'x','y','z'}: + if self.input_space[key] >1: + inputkeys = {key+str(l+1) for l in range(self.input_space[key])} + else: + inputkeys = inputkeys|{key} + for key in set(self.output_space.keys()): + if self.output_space[key] >1: + outputkeys = outputkeys|{key+str(l+1) for l in range(self.output_space[key])} + else: + outputkeys = outputkeys|{key} + + + return spatial_keys|inputkeys, outputkeys + + def convertVariables(self,vars): + ''' + Converts the input variables to the format of the TP + objective function + ''' + conv_vars = vars.copy() + + if {'x','y','z'} <= set(self.input_space.keys()): + pass + elif {'x','y'} <= set(self.input_space.keys()): + pass + elif {'x'} <= set(self.input_space.keys()): + if self.input_space['x'] > 1: + if self.input_space['x'] ==2: + conv_vars['x']=torch.cat((vars['x'],vars['y']),dim=1) + conv_vars.pop('y') + else: + conv_vars['x']=torch.cat((vars['x'],vars['y'],vars['z']),dim=1) + conv_vars.pop('y') + conv_vars.pop('z') + + for key in set(self.input_space.keys())- {'x','y','z'}: + if self.input_space[key] >1: + cat_var = list(vars[key+str(l+1)] for l in range(self.input_space[key])) + conv_vars[key] = torch.cat(cat_var,dim=1) + for l in range(self.input_space[key]): + conv_vars.pop(key+str(l+1)) + + for key in self.output_space.keys(): + if self.output_space[key] >1: + if self.input_space['x'] > 1: + if self.input_space['x'] ==2: + cat_var = list(OutvarFunction.apply(vars[key+str(l+1)], conv_vars['x'],vars['x'],vars['y']) for l in range(self.output_space[key])) + else: + cat_var = list(OutvarFunction.apply(vars[key+str(l+1)], conv_vars['x'],vars['x'],vars['y'],vars['z']) for l in range(self.output_space[key])) + conv_vars[key] = torch.cat(cat_var,dim=1) + else: + cat_var = list(vars[key+str(l+1)] for l in range(self.output_space[key])) + conv_vars[key] = torch.cat(cat_var,dim=1) + for l in range(self.output_space[key]): + conv_vars.pop(key+str(l+1)) + else: + if self.input_space['x'] > 1: + if self.input_space['x'] ==2: + conv_vars[key] = OutvarFunction.apply(vars[key], conv_vars['x'],vars['x'],vars['y']) + else: + conv_vars[key] = OutvarFunction.apply(vars[key], conv_vars['x'],vars['x'],vars['y'],vars['z']) + + for param in self.params[0].space.variables: + if self.params[0].space.dim>1: + conv_vars[param] = torch.cat(list(vars[param+str(l)] for l in range(self.params[0].space.dim)),dim=1) + for l in range(self.params[0].space.dim): + conv_vars.pop(param+str(l)) + else: + conv_vars[param] = vars[param] + + for fun in self.data_functions.keys(): + conv_vars[fun] = self.data_functions[fun](conv_vars) + + return conv_vars + + +class OutvarFunction(torch.autograd.Function): + """ + Function that calculates the gradient of the output variable with + respect to a multi-dimensional spatial input variable, if the + gradient function is given for the corresponding one-dimensional + spatial input variable. + Variables x,y,z are the one-dimensional spatial input variables and + x_vec is the multi-dimensional spatial input variable. + """ + @staticmethod + def forward(ctx, u, x_vec, x, y, z=None): + """ + Forward pass of the node. + + Args: + ctx (torch.autograd.function._ContextMethodMixin): Context + object for autograd. + u (torch.Tensor): Input tensor. + x_vec (torch.Tensor): Input tensor. + x (torch.Tensor): Input tensor. + y (torch.Tensor): Input tensor. + z (torch.Tensor, optional): Input tensor. Defaults to None. + + Returns: + torch.Tensor: Output tensor. + """ + # Compute u(u, xy) (new u) during forward pass + result = u + + # Save u and xy for the backward pass + if z is not None: + ctx.save_for_backward(u, x_vec, x, y, z) + else: + ctx.save_for_backward(u, x_vec, x, y) + return result + + @staticmethod + def backward(ctx, grad_output): + """ + Computes the backward pass for the custom autograd function. + + Args: + ctx (torch.autograd.function._ContextMethodMixin): The + context object that holds the saved tensors. + grad_output (torch.Tensor): The gradient of the output + with respect to the function's output. + + Returns: + tuple: A tuple containing the gradients of the input + tensors with respect to the function's inputs. + """ + grad_u = grad_output + + # Load u and xy from the context + if len(ctx.saved_tensors)==5: + u, x_vec,x,y,z = ctx.saved_tensors + else: + u, x_vec, x,y = ctx.saved_tensors + # Compute the gradients of u (new) with respect to u and xy + if len(ctx.saved_tensors)==5: + grad_x, grad_y, grad_z = torch.autograd.grad(u, (x,y,z), grad_output, create_graph=True) + grad_xyz = torch.cat((grad_x,grad_y,grad_z),dim=1) + return grad_u, OutvarFunction.apply(grad_xyz,x_vec,x,y,z), OutvarFunction.apply(grad_x,x_vec,x,y,z), OutvarFunction.apply(grad_y,x_vec,x,y,z), OutvarFunction.apply(grad_z,x_vec,x,y,z) + else: + grad_x, grad_y = torch.autograd.grad(u, (x,y), grad_output, create_graph=True) + grad_xy = torch.cat((grad_x,grad_y),dim=1) + return grad_u, OutvarFunction.apply(grad_xy,x_vec,x,y), OutvarFunction.apply(grad_x,x_vec,x,y), OutvarFunction.apply(grad_y,x_vec,x,y) diff --git a/src/torchphysics/wrapper/solver.py b/src/torchphysics/wrapper/solver.py new file mode 100644 index 00000000..413386e4 --- /dev/null +++ b/src/torchphysics/wrapper/solver.py @@ -0,0 +1,539 @@ +from torchphysics.problem.conditions.condition import PINNCondition, DataCondition, ParameterCondition +from torchphysics.problem.domains.domainoperations.product import ProductDomain +from torchphysics.problem.domains.domainoperations.union import UnionDomain, UnionBoundaryDomain +from torchphysics.problem.domains.domainoperations.intersection import IntersectionDomain, IntersectionBoundaryDomain +from torchphysics.problem.domains.domainoperations.cut import CutDomain, CutBoundaryDomain +from torchphysics.problem.domains.domainoperations.rotate import Rotate +from torchphysics.problem.domains.domainoperations.translate import Translate +from torchphysics.models.parameter import Parameter +from torchphysics.problem.spaces import Points + +from modulus.sym.domain import Domain +from modulus.sym.domain.constraint import PointwiseBoundaryConstraint, PointwiseInteriorConstraint, PointwiseConstraint +from modulus.sym.domain.inferencer import PointwiseInferencer +from modulus.sym.domain.validator import PointwiseValidator +from modulus.sym.domain.monitor import PointwiseMonitor +from modulus.sym.loss import PointwiseLossNorm + + +from modulus.sym.node import Node +from modulus.sym.key import Key +from modulus.sym.models.fully_connected import FullyConnectedArch + +from sympy import Symbol +from functools import partial + +import warnings +import torch +import sympy + +from .model import TPModelArch +from .nodes import TPNodeFunction +from .geometry import TPGeometryWrapper +from .helper import convertDataModulus2TP, convertDataTP2Modulus, PointwiseLossInfNorm, PointwiseLossMean, CustomInferencerPlotter, PINNConditionValidator, DataConditionValidator + +class ModulusSolverWrapper(): + """ + Wrapper to use the solver class of Modulus with the rest + implemented in TorchPhysics + + Parameters + ---------- + + tp_solver: torchphysics.solver.Solver (pl.LightningModule) + TorchPhysics solver class + callbacks: list of torchphysics.callbacks.Callback + List of TorchPhysics callbacks + **kwargs: + lambda_weighting: list of floats, integers, sympy expressions + or the string "sdf" + List of lambda weightings for the conditions. If the + string "sdf" is used, the lambda weighting is multiplied + with the signed distance function of the boundary condition. + If only one value is given, it is used for all conditions. + If the length of the list is not equal to the number of + conditions, the default value of 1.0 is used for the + remaining conditions. + + Returns + ------- + ModulusSolverWrapper + Wrapper class for Modulus solver class containing the Modulus + domain with the necessary nodes, geometries, conditions and + parameters + + """ + + def __init__(self,tpsolver,callbacks,**kwargs): + self.nodes = [] + self.domain = Domain() + + self.lambda_weighting_vals = [1.0]*len(tpsolver.train_conditions) + for key, value in kwargs.items(): + if key == 'lambda_weighting': + assert type(value) == list, "lambda_weighting must be a list" + assert all((type(val) in (int,float)) or (isinstance(val, sympy.Expr)) or (val=="sdf") for val in value), "lambda_weighting must be a list of floats, integers, sympy expressions or the string ""sdf""" + if len(value) == 1: + value = value*len(tpsolver.train_conditions) + elif len(value) != len(tpsolver.train_conditions): + assert False, "lambda_weighting must have the same length as the number of conditions or 1" + self.lambda_weighting_vals = value + + self.device = "cuda:0" if torch.cuda.is_available() else "cpu" + + self.models = [] + num_models = 0 + self.orig_models = [] + self.Modmodels = [] + self.Geometries = [] + self.conds = [] + self.isBoundary = [] + self.parameters = [] + self.parameter_samples = [] + self.parameter_nets = [] + objectives = [] + exist_DataCondition = False + is_inverse_problem = False + not_seen_before = True + + # loop through all conditions to collect nodes out of objective functions + for condition in tpsolver.train_conditions+tpsolver.val_conditions: + if type(condition)!=ParameterCondition: + model = condition.module + self.orig_models.append(model) + if type(model).__name__=='Parallel': + models = model.models + else: + models = [model] + for mod in models: + if mod not in self.models: + self.models.append(mod) + Modmodel = TPModelArch(mod) + self.Modmodels.append(Modmodel) + num_models +=1 + self.nodes.append(Modmodel.make_node("model"+str(num_models-1))) + + if (type(condition) == PINNCondition): + if condition.sampler.is_static: + for fn in condition.data_functions: + condition.data_functions[fn].fun = condition.data_functions[fn].fun.to(self.device) + # construct node out of objective function + NodeFCN=TPNodeFunction(condition.residual_fn,model.input_space,model.output_space,condition.data_functions,condition.parameter) + objectives.append(NodeFCN) + if condition in tpsolver.train_conditions and condition in tpsolver.val_conditions: + if not_seen_before: + self.nodes.append(Node(NodeFCN.variables, NodeFCN.cond_names, NodeFCN)) + not_seen_before = False + else: + self.nodes.append(Node(NodeFCN.variables, NodeFCN.cond_names, NodeFCN)) + # if there is a learnable parameter in the condition, add a parameter net to learn this parameter as Modulus does not support additional learnable parameters in the conditions. + # The parameter is then learned as a function of the input variables of the TP model + if (condition.parameter!=Parameter.empty())& (condition.parameter not in self.parameters): + is_inverse_problem = True + inputkeys, _= TPModelArch(model).getVarKeys(model) + outputkeys = [Key(var+str(ind+1)) for var in condition.parameter.space.variables for ind in range(condition.parameter.space.dim)] if condition.parameter.space.dim>1 else [Key(var) for var in condition.parameter.space.variables] + # the parameter net is a fully connected network with 2 layers and 30 neurons per layer + # the input keys are the keys of the corresponding TP model input space + parameter_net = FullyConnectedArch(input_keys=inputkeys, output_keys=outputkeys, layer_size = 30, nr_layers = 2) + self.parameter_nets.append(parameter_net) + self.nodes.append(parameter_net.make_node("parameter_net")) + self.parameters.append(condition.parameter) + points = condition.sampler.sample_points() + # sort out the input keys of the TP sampler and the corresponding values to the order of TP model input keys + self.parameter_samples.append(points[..., list(model.input_space.keys())]) + elif type(condition) == DataCondition: + exist_DataCondition = True + objectives.append(None) + elif type(condition) == ParameterCondition: + objectives.append(None) + else: + assert False, "Only PINNCondition, DataCondition, ParameterCondition are allowed as conditions" + + assert(exist_DataCondition if is_inverse_problem else True), "DataCondition must be present for inverse problems" + + # loop over train conditions to build Modulus constraints out of TorchPhysics conditions + for condition, obj, weight in zip(tpsolver.train_conditions,objectives[0:len(tpsolver.train_conditions)],self.lambda_weighting_vals): + if type(condition) == PINNCondition: + + # identify sampler + # check if static sampler + is_static = condition.sampler.is_static + sampler = condition.sampler.sampler if is_static else condition.sampler + quasi_random = False + if type(sampler).__name__ != 'RandomUniformSampler': + if type(sampler).__name__ == 'LHSSampler': + warnings.warn("Modulus only supports RandomUniformSampler or Halton sequence. Using Halton sequence instead.") + quasi_random = True + else: + warnings.warn("Modulus only supports RandomUniformSampler or Halton sequence. Using RandomUniformSampler instead.") + + + # identify different types of domains and split them into spatial, time and parameter domains + # spatial_domain is a nested dictionary containing different domain operations and the underlying TorchPhysics domains + spatial_domain, time_domain, parameter_domain = self.TPDomainWrapper(sampler.domain) + + # identify parameter ranges + param_ranges={} + for dom in parameter_domain: + if dom.dim>1: + for l in range(0,dom.dim): + param_ranges[Symbol(list(dom.space.variables)[0]+str(l+1))]= tuple((dom.bounding_box()[2*l].item(),dom.bounding_box()[2*l+1].item())) + else: + param_ranges[Symbol(list(dom.space.variables)[0])]= tuple(dom.bounding_box().numpy()) + + # identify time variable ranges + if time_domain != []: + assert(all(dom == time_domain[0] for dom in time_domain)), "Only single time domain allowed" + time_domain = time_domain[0] + if time_domain.dim == 1: + param_ranges[Symbol("t")]=tuple(time_domain.bounding_box().numpy()) + elif time_domain.dim == 0: + param_ranges[Symbol("t")]= time_domain.bounding_box()[0].item() + + # Build geometry out of spatial domain + is_boundary, geometry, cond, _, _= TPGeometryWrapper(spatial_domain).getModulusGeometry() + self.Geometries.append(geometry) + self.conds.append(cond) + self.isBoundary.append(is_boundary) + + # determine lambda_weightings + if weight == "sdf": + if is_boundary: + lambda_weightings = {name : condition.weight for name in obj.cond_names} + else: + lambda_weightings = {name : condition.weight*Symbol("sdf") for name in obj.cond_names} + else: + lambda_weightings = {name : condition.weight*weight for name in obj.cond_names} + + assert (condition.track_gradients == True), "track_gradients must be True for PINNCondition" + + # add constraints to domain + if is_boundary: + constraint = PointwiseBoundaryConstraint( + nodes=self.nodes, + geometry=geometry, + outvar={name : 0.0 for name in obj.cond_names}, + batch_size=len(sampler.sample_points()), + parameterization=param_ranges, + lambda_weighting = lambda_weightings, + fixed_dataset = is_static, + criteria = cond, + batch_per_epoch = 1, + quasirandom = quasi_random, + ) + + else: + constraint = PointwiseInteriorConstraint( + nodes=self.nodes, + geometry=geometry, + outvar={name : 0.0 for name in obj.cond_names}, + batch_size=len(sampler.sample_points()), + parameterization=param_ranges, + lambda_weighting = lambda_weightings, + fixed_dataset = is_static, + batch_per_epoch = 1, + criteria = cond, + quasirandom = quasi_random, + ) + self.domain.add_constraint(constraint, condition.name) + + + elif type(condition) == DataCondition: + if condition.use_full_dataset == True: + batch_size = len(condition.dataloader.dataset.data_points[0]) + else: + batch_size = condition.dataloader.dataset.batch_size + + if condition.norm == "inf": + norm = PointwiseLossInfNorm() + else: + norm = PointwiseLossNorm(condition.norm) + assert (condition.root == 1), "Only root=1 is allowed for DataCondition" + + outvar=convertDataTP2Modulus(condition.dataloader.dataset.data_points[1]) + invar=convertDataTP2Modulus(condition.dataloader.dataset.data_points[0]) + + # determine lambda_weightings + # lambda_weightings has to be dict of numpy arrays in the same length as invar + lambda_weightings = {name : condition.weight*weight*torch.ones(len(outvar[list(outvar.keys())[0]]),1) for name in outvar.keys()} + + data_constraint = PointwiseConstraint.from_numpy( + nodes=self.nodes, + invar=invar, + outvar=outvar, + loss = norm, + batch_size=batch_size, + lambda_weighting = lambda_weightings, + ) + + # define new forward function for data constraint to include the constrain_fn + if condition.constrain_fn: + def create_new_forward(condition): + def new_forward(self): + self._output_vars = self.model(self._input_vars) + inputvars = convertDataModulus2TP(self._input_vars,condition.module.input_space) + outputvars = convertDataModulus2TP(self._output_vars,condition.module.output_space) + constraint_output = condition.constrain_fn({**outputvars, **inputvars}) + output_dict = {key: constraint_output[:,index:index+condition.module.output_space[key]] for index, key in enumerate(outputvars.keys())} + #output_dict = {key: constraint_output[:,index:index+condition.module.output_space[key]] for index, key in enumerate(condition.module.output_space.keys())} + output_vars_points = Points.from_coordinates(output_dict) + self._output_vars = convertDataTP2Modulus(output_vars_points) + + return new_forward + + data_constraint.forward = create_new_forward(condition).__get__(data_constraint, PointwiseConstraint) + + self.domain.add_constraint(data_constraint, condition.name) + + elif type(condition) == ParameterCondition: + assert (condition.parameter.space.dim ==1), "Only single parameter allowed for ParameterCondition" + param_index = self.parameters.index(condition.parameter) + + points = self.parameter_samples[param_index] + modpoints = dict(zip([str(key) for key in self.parameter_nets[param_index].input_keys],[points.as_tensor[:,index].reshape(len(points),1) for index in range(len(self.parameter_nets[param_index].input_keys))])) + + # determine lambda_weightings + # lambda_weightings has to be dict of numpy arrays with the same length as invar + lambda_weightings = {condition.parameter.variables.pop() : condition.weight*weight*torch.ones(len(points),1)} + + parameter_constraint = PointwiseConstraint.from_numpy( + nodes=self.nodes, + invar=modpoints, + outvar= {condition.parameter.variables.pop():torch.zeros(len(points),1)}, + batch_size=batch_size, + loss = PointwiseLossMean(), + lambda_weighting = lambda_weightings, + ) + + # create new forward function for parameter constraint to include the penalty function + def create_new_forward_paramCond(condition): + def new_forward(self): + self._output_vars = self.model(self._input_vars) + penalty_output = condition.penalty({**self._output_vars}) + self._output_vars = {condition.parameter.variables.pop(): penalty_output} + return new_forward + + parameter_constraint.forward = create_new_forward_paramCond(condition).__get__(parameter_constraint, PointwiseConstraint) + self.domain.add_constraint(parameter_constraint, condition.name) + + else: + assert False, "Condition type not yet supported! Only PINNCondition, DataCondition or ParameterCondition!" + + # loop over validation conditions to build Modulus constraints out of TorchPhysics conditions + for condition, obj in zip(tpsolver.val_conditions,objectives[len(tpsolver.train_conditions):]): + if type(condition) == PINNCondition: + # convert sample points to Modulus format + samples=convertDataTP2Modulus(sampler.sample_points()) + + # build validator + validator = PINNConditionValidator( + nodes=self.nodes, + invar = samples, + output_names = obj.cond_names, + batch_size=len(samples), + requires_grad = condition.track_gradients, + ) + + self.domain.add_validator(validator, condition.name) + + elif type(condition) == DataCondition: + batch_size = condition.dataloader.dataset.batch_size + + outvar=convertDataTP2Modulus(condition.dataloader.dataset.data_points[1]) + invar=convertDataTP2Modulus(condition.dataloader.dataset.data_points[0]) + + validator = DataConditionValidator( + nodes=self.nodes, + invar=invar, + true_outvar=outvar, + batch_size=batch_size, + requires_grad = condition.track_gradients, + norm = condition.norm, + root = condition.root + ) + + # define new forward function for data validator to include the constrain_fn + if condition.constrain_fn: + def create_new_forward(condition): + def new_forward(self,invar): + with torch.set_grad_enabled(condition.track_gradients): + pred_outvar = self.model(invar) + inputvars = convertDataModulus2TP(invar,condition.module.input_space) + outputvars = convertDataModulus2TP(pred_outvar,condition.module.output_space) + constraint_output = condition.constrain_fn({**outputvars, **inputvars}) + output_dict = {key: constraint_output[:,index:index+condition.module.output_space[key]] for index, key in enumerate(condition.module.output_space.keys())} + output_vars_points = Points.from_coordinates(output_dict) + return convertDataTP2Modulus(output_vars_points) + return new_forward + + validator.forward = create_new_forward(condition).__get__(validator, PointwiseValidator) + self.domain.add_validator(validator, condition.name) + + # if inverse problem with single parameters to identify, add parameter monitor for each parameter + # A parameter is learned by a parameter net, and the mean of the net output will be the parameter value and is monitored + if is_inverse_problem: + for index, param_obj in enumerate(zip(self.parameter_samples,self.parameter_nets)): + points = param_obj[0] + # ordered values get the corresponding Modulus input keys (if some of TP input keys are of higher dimension than 1, they are split into several Modulus input keys) + modpoints = dict(zip([str(key) for key in param_obj[1].input_keys],[points.as_tensor[:,index].reshape(len(points),1) for index in range(len(param_obj[1].input_keys))])) + + # define mean function to compute mean of parameter net over a fixed set of input points + def mean_func(var, key): + return torch.mean(var[str(key)], dim=0) + + metrics = {"mean_"+str(key): partial(mean_func, key=key) for key in param_obj[1].output_keys} + + + parameter_monitor = PointwiseMonitor( + nodes=self.nodes, + invar=modpoints, + output_names=[str(key) for key in param_obj[1].output_keys], + metrics = metrics, + ) + self.domain.add_monitor(parameter_monitor, "parameter_monitor"+str(index)) + + + # if plotter callback is present, an inferencer with plotter is added to the domain + if callbacks: + for callback in callbacks: + invars = { key: value.cpu().detach().numpy() for key, value in convertDataTP2Modulus(callback.point_sampler.sample_points()).items()} + callback.point_sampler.created_points = None + + plotter_inferencer = PointwiseInferencer( + nodes=self.nodes, + invar=invars, + output_names= [var+str(ind+1) if callback.model.output_space[var] >1 else var for var in callback.model.output_space.variables for ind in range(callback.model.output_space[var]) ], + batch_size=len(invars), + plotter=CustomInferencerPlotter(callback), + ) + self.domain.add_inferencer(plotter_inferencer, callback.log_name) + + + def TPDomainWrapper(self,domain): + """ + Function that parses domain and splits it into spatial, time + and parameter domains. + The spatial domain is recursively split into nested dictionary + containing the following keys: + -'u': union + -'i': intersection + -'p': product + -'c': cut + -'r': rotate + -'t': translate + -'d': final domain + The values are the partial/splitted domains + + It returns the spatial_domain, time_domain and parameter_domain. + + Parameters + ---------- + domain: torchphysics.problem.domains.domain.Domain + TorchPhysics domain + + Returns + ------- + spatial_domain: dict + Nested dictionary containing the splitted domain + time_domain: list + List of time domains + parameter_domain: list + List of parameter domains + + """ + + # we have to define these global variables due to the recursive nature of the function + global parameter_domain, time_domain + parameter_domain = [] + time_domain = [] + spatial_domain = self.splitDomains(domain) + return spatial_domain, time_domain, parameter_domain + + + def splitDomains(self, domain): + """ + Recursive function that splits the spatial domain into nested + dictionary containing the following keys: + -'u': union + -'i': intersection + -'p': product + -'c': cut + -'r': rotate + -'t': translate + -'d': final domain + The values are then the partial/splitted domains + """ + global parameter_domain, time_domain + if self.is_splittable(domain): + if not hasattr(domain,'domain_a'): + domain.domain_a = domain.domain.domain_a + domain.domain_b = domain.domain.domain_b + + if self.is_xspace(domain.domain_a): + if self.is_xspace(domain.domain_b): + if type(domain) == ProductDomain: + return {'p': (self.splitDomains(domain.domain_a),self.splitDomains(domain.domain_b))} + elif type(domain) == UnionDomain: + return {'u': (self.splitDomains(domain.domain_a),self.splitDomains(domain.domain_b))} + elif type(domain) == UnionBoundaryDomain: + return {'ub': (self.splitDomains(domain.domain.domain_a),self.splitDomains(domain.domain.domain_b))} + elif type(domain) == IntersectionDomain: + return {'i': (self.splitDomains(domain.domain_a),self.splitDomains(domain.domain_b))} + elif type(domain) == IntersectionBoundaryDomain: + return {'ib': (self.splitDomains(domain.domain.domain_a),self.splitDomains(domain.domain.domain_b))} + elif type(domain) == CutDomain: + return {'c': (self.splitDomains(domain.domain_a),self.splitDomains(domain.domain_b))} + elif type(domain) == CutBoundaryDomain: + return {'cb': (self.splitDomains(domain.domain.domain_a),self.splitDomains(domain.domain.domain_b))} + else: + if self.is_tspace(domain.domain_b): + time_domain.append(domain.domain_b) + else: + parameter_domain.append(domain.domain_b) + return self.splitDomains(domain.domain_a) + + else: + if self.is_xspace(domain.domain_b): + if self.is_tspace(domain.domain_a): + time_domain.append(domain.domain_a) + else: + parameter_domain.append(domain.domain_a) + return self.splitDomains(domain.domain_b) + else: + if self.is_tspace(domain.domain_a): + time_domain.append(domain.domain_a) + else: + parameter_domain.append(domain.domain_a) + if self.is_tspace(domain.domain_b): + time_domain.append(domain.domain_b) + else: + parameter_domain.append(domain.domain_b) + return None + + elif type(domain) == Rotate: + # in the case of rotation, we have to collect the rotation function and the rotation center for later use + return {'r': [self.splitDomains(domain.domain),domain.rotation_fn(),domain.rotate_around()]} + elif type(domain) == Translate: + # in the case of translation, we have to collect the translation function for later use + return {'t': [self.splitDomains(domain.domain),domain.translate_fn()]} + else: + if self.is_xspace(domain): + return {'d': domain} + else: + if self.is_tspace(domain): + time_domain.append(domain) + else: + parameter_domain.append(domain) + return None + + def is_xspace(self,domain): + return ('x' in domain.space.variables)|('y' in domain.space.variables)|('z' in domain.space.variables) + + def is_tspace(self,domain): + return 't' in domain.space.variables + + def is_splittable(self,domain): + return type(domain) in set((ProductDomain,UnionDomain,UnionBoundaryDomain,IntersectionDomain,IntersectionBoundaryDomain,CutBoundaryDomain,CutDomain)) + + + diff --git a/src/torchphysics/wrapper/tests/test_ParallelogramCylinder.py b/src/torchphysics/wrapper/tests/test_ParallelogramCylinder.py new file mode 100644 index 00000000..a7a68aba --- /dev/null +++ b/src/torchphysics/wrapper/tests/test_ParallelogramCylinder.py @@ -0,0 +1,77 @@ +import pytest +from math import sqrt +import numpy as np +from modulus.sym.geometry.geometry import Geometry +from modulus.sym.geometry.curve import Curve +from modulus.sym.geometry.parameterization import Parameterization +from torchphysics.wrapper.geometry import ParallelogramCylinder + + +def test_ParallelogramCylinder_creation(): + origin = (0, 0, 0) + corner1 = (1, 0, 0) + corner2 = (0, 1, 0) + height = 2 + parameterization = Parameterization() + + cylinder = ParallelogramCylinder(origin, corner1, corner2, height, parameterization) + + assert isinstance(cylinder, Geometry) + assert cylinder.parameterization == parameterization + assert callable(cylinder.sdf) + assert cylinder.dims == ['x', 'y', 'z'] + assert cylinder.sdf({'x': np.array([[0]]),'y':np.array([[0]]),'z':np.array([[0]])},params={})['sdf'][0][0] == 0 + assert cylinder.sdf({'x': np.array([[0.5]]),'y':np.array([[0.5]]),'z':np.array([[1]])},params={})['sdf'][0][0]==pytest.approx(0.5,rel=1e-12) + assert cylinder.sdf({'x': np.array([[1.5]]),'y':np.array([[0.5]]),'z':np.array([[2]])},params={})['sdf'][0][0]== -0.5 + + + +def test_ParallelogramCylinder_curves(): + origin = (0, 0, 0) + corner1 = (1, 0, 0) + corner2 = (0, 1, 0) + height = 2 + parameterization = Parameterization() + + cylinder = ParallelogramCylinder(origin, corner1, corner2, height, parameterization) + + assert len(cylinder.curves) == 6 + assert all(isinstance(curve, Curve) for curve in cylinder.curves) + + + +def test_ParallelogramCylinder_invalid_parameters(): + origin = (0, 0, 0) + corner1 = (1, 0, 1) + corner2 = (0, 1, 0) + height = 2 + parameterization = Parameterization() + + with pytest.raises(AssertionError): + ParallelogramCylinder(origin, corner1, corner2, height, parameterization) + +def test_ParallelogramCylinder_negative_height(): + origin = (0, 0, 0) + corner1 = (1, 0, 0) + corner2 = (0, 1, 0) + height = -2 # Negative height should raise an error or be considered invalid + parameterization = Parameterization() + + with pytest.raises(TypeError): + ParallelogramCylinder(origin, corner1, corner2, height, parameterization) + + + + +def test_ParallelogramCylinder_repr(): + origin = (0, 0, 0) + corner1 = (1, 0, 0) + corner2 = (0, 1, 0) + height = 2 + parameterization = Parameterization() + + cylinder = ParallelogramCylinder(origin, corner1, corner2, height, parameterization) + + assert isinstance(repr(cylinder), str) + assert "ParallelogramCylinder" in repr(cylinder) + diff --git a/src/torchphysics/wrapper/tests/test_TPGeometryWrapper.py b/src/torchphysics/wrapper/tests/test_TPGeometryWrapper.py new file mode 100644 index 00000000..c09b5db6 --- /dev/null +++ b/src/torchphysics/wrapper/tests/test_TPGeometryWrapper.py @@ -0,0 +1,269 @@ +import pytest +import torchphysics as tp +import numpy as np +import torch + +from torchphysics.wrapper.geometry import TPGeometryWrapper +from torchphysics.problem.domains import (Interval, Circle, Sphere, Triangle, Parallelogram) +from torchphysics.problem.domains.domain2D.shapely_polygon import ShapelyPolygon +from torchphysics.problem.domains.domain3D.trimesh_polyhedron import TrimeshPolyhedron +from modulus.sym.geometry.primitives_1d import (Line1D) +from modulus.sym.geometry.primitives_2d import Circle as Circle_modulus +from modulus.sym.geometry.primitives_2d import Triangle as Triangle_modulus +from modulus.sym.geometry.primitives_2d import Polygon +from modulus.sym.geometry.primitives_3d import Sphere as Sphere_modulus + +from sympy import Eq, Symbol + + + +def single_domain_dict(domain): + domain_dict = {'d': domain} + return domain_dict + +def union_domain_dict(domain1,domain2): + domain_dict = {'u': [{'d': domain1},{'d': domain2}]} + return domain_dict + +def intersection_domain_dict(domain1,domain2): + return {'i': [{'d': domain1},{'d': domain2}]} + + +def cut_domain_dict(domain1,domain2): + return {'c' : [{'d': domain1},{'d': domain2}]} + +def translate_domain_dict(domain,translation_params): + return {'t' : [{'d': domain},translation_params]} + +def rotate_domain_dict(domain,rotation_matrix,rotation_center): + return {'r' : [{'d': domain},rotation_matrix,rotation_center]} + +def combined_operations_dict(domain1,domain2,domain3,translation_params,rotation_matrix,rotation_center): + return {'i' : [{'t': [{'r' : [{'d': domain1},rotation_matrix,rotation_center]},translation_params]},{'p':[{'d': domain2},{'d': domain3}]}]} + + +def product_domain_dict(domain1,domain2): + return {'p': [{'d': domain1},{'d': domain2}]} + +# Fixtures for different domain types +@pytest.fixture +def interval_domain(): + X = tp.spaces.R1('x') + return Interval(X,0, 1) + +@pytest.fixture +def interval_domain2(): + X = tp.spaces.R1('x') + return Interval(X,2,4) + +@pytest.fixture +def interval_domain_y(): + Y = tp.spaces.R1('y') + return Interval(Y,0,4) + + +@pytest.fixture +def interval_domain_z(): + Z = tp.spaces.R1('z') + return Interval(Z,0,4) + + + +@pytest.fixture +def circle_domain(): + X = tp.spaces.R2('x') + return Circle(X,center=(0, 0), radius=1) + +@pytest.fixture +def circle_domain_xy(): + X = tp.spaces.R1('x') + Y = tp.spaces.R1('y') + return Circle(X*Y,center=(0, 0), radius=1) + +@pytest.fixture +def sphere_domain(): + X = tp.spaces.R3('x') + return Sphere(X,center=(0, 0, 0), radius=1) + +@pytest.fixture +def triangle_domain(): + X = tp.spaces.R2('x') + return Triangle(X,origin=(0, 0), corner_1=(1, 0), corner_2=(0.5, np.sqrt(3)/2)) + +@pytest.fixture +def parallelogram_domain(): + X = tp.spaces.R2('x') + return Parallelogram(X,origin=(0, 0), corner_1=(1, 0), corner_2=(0.5, 0.5)) + +@pytest.fixture +def parallelogram_domain_xy(): + XY = tp.spaces.R1('x')*tp.spaces.R1('y') + return Parallelogram(XY,origin=(0, 0), corner_1=(1, 0), corner_2=(0, 2)) + +@pytest.fixture +def trimesh_polyhedron_domain(): + vertices = [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [0.5, 0.5, -1]] + faces = [[0, 1, 3], [0, 2, 3], [1, 2, 3], [0, 2, 4], [0, 1, 4], [1, 2, 4]] + poly3D = TrimeshPolyhedron(tp.spaces.R3('x'), vertices=vertices, faces=faces) + return poly3D + +@pytest.fixture +def shapely_polygon_domain(): + X = tp.spaces.R2('x') + return ShapelyPolygon(X, vertices=[[0, 0], [1, 0], [1, 2], [0, 1]]) + +# Test functions for each domain type +def test_interval_domain(interval_domain): + is_boundary, geometry, cond, _, _= TPGeometryWrapper(single_domain_dict(interval_domain)).getModulusGeometry() + assert isinstance(geometry, Line1D) + assert not is_boundary + assert cond == None + +def test_circle_domain(circle_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(circle_domain)).getModulusGeometry() + assert isinstance(geometry, Circle_modulus) + assert not is_boundary + assert cond == None + +def test_sphere_domain(sphere_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(sphere_domain)).getModulusGeometry() + assert isinstance(geometry, Sphere_modulus) + assert not is_boundary + assert cond == None + +def test_triangle_domain(triangle_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(triangle_domain)).getModulusGeometry() + assert isinstance(geometry, Triangle_modulus) + assert not is_boundary + assert cond == None + +def test_parallelogram_domain(parallelogram_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(parallelogram_domain)).getModulusGeometry() + assert isinstance(geometry, Polygon) + assert not is_boundary + assert cond == None + +def test_trimesh_polyhedron_domain(trimesh_polyhedron_domain): + with pytest.raises(Exception) as excinfo: + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(trimesh_polyhedron_domain)).getModulusGeometry() + try: + assert isinstance(geometry, Tessellation) + assert not is_boundary + assert cond == None + except NameError: + assert "Tessellation module only supported for Modulus docker installation due to missing pysdf installation!" in str(excinfo.value) + + +def test_boundary_of_interval_domain(interval_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(interval_domain.boundary_left)).getModulusGeometry() + assert isinstance(geometry, Line1D) + assert is_boundary + assert (cond==Eq(Symbol('x'),0)) + +def test_boundary_of_circle_domain(circle_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(circle_domain.boundary)).getModulusGeometry() + assert isinstance(geometry, Circle_modulus) + assert is_boundary + +def test_boundary_of_sphere_domain(sphere_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(sphere_domain.boundary)).getModulusGeometry() + assert isinstance(geometry, Sphere_modulus) + assert is_boundary + +def test_boundary_of_triangle_domain(triangle_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(triangle_domain.boundary)).getModulusGeometry() + assert isinstance(geometry,Triangle_modulus) + assert is_boundary + +def test_boundary_of_parallelogram_domain(parallelogram_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(parallelogram_domain.boundary)).getModulusGeometry() + assert isinstance(geometry, Polygon) + assert is_boundary + +def test_boundary_of_trimesh_polyhedron_domain(trimesh_polyhedron_domain): + with pytest.raises(Exception) as excinfo: + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(single_domain_dict(trimesh_polyhedron_domain.boundary)).getModulusGeometry() + try: + assert isinstance(geometry, Tesselation) + assert is_boundary + assert cond == None + except NameError: + assert "Tessellation module only supported for Modulus docker installation due to missing pysdf installation!" in str(excinfo.value) + + + + +def test_union_operation(interval_domain,interval_domain2): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(union_domain_dict(interval_domain,interval_domain2)).getModulusGeometry() + assert not is_boundary + assert geometry.sdf({'x': [[1.5]]},params={})['sdf'][0][0] < 0 + assert geometry.sdf({'x': [[3]]},params={})['sdf'][0][0] > 0 + assert geometry.sdf({'x': [[0.5]]},params={})['sdf'][0][0]>0 + assert abs(geometry.sdf({'x': [[2]]},params={})['sdf'][0][0]) < 1e-12 + assert cond == None + + + +def test_intersection_operation(circle_domain,parallelogram_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(intersection_domain_dict(circle_domain,parallelogram_domain)).getModulusGeometry() + assert not is_boundary + assert abs(geometry.sdf({'x': np.array([[0.5]]),'y':np.array([[0.5]])},params={})['sdf'][0][0]) < 1e-16 + assert geometry.sdf({'x': np.array([[0.5]]),'y':np.array([[0.1]])},params={})['sdf'][0][0] > 0 + assert geometry.sdf({'x': np.array([[2]]),'y':np.array([[0]])},params={})['sdf'][0][0] < 0 + assert geometry.sdf({'x': np.array([[-0.5]]),'y':np.array([[0.5]])},params={})['sdf'][0][0] < 0 + assert cond == None + +def test_cut_operation(circle_domain,parallelogram_domain): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(cut_domain_dict(circle_domain,parallelogram_domain)).getModulusGeometry() + assert not is_boundary + assert abs(geometry.sdf({'x': np.array([[0.5]]),'y':np.array([[0.5]])},params={})['sdf'][0][0]) < 1e-16 + assert geometry.sdf({'x': np.array([[0.5]]),'y':np.array([[0.1]])},params={})['sdf'][0][0] < 0 + assert geometry.sdf({'x': np.array([[2]]),'y':np.array([[0]])},params={})['sdf'][0][0] < 0 + assert geometry.sdf({'x': np.array([[-0.5]]),'y':np.array([[0.5]])},params={})['sdf'][0][0] > 0 + assert cond == None + +def test_translate_operation(sphere_domain): + is_boundary, geometry, cond, translate_vec, rotation_list = TPGeometryWrapper(translate_domain_dict(sphere_domain,torch.tensor([5,0,0]))).getModulusGeometry() + assert not is_boundary + assert cond == None + assert rotation_list == [[],[],[],1] + assert translate_vec == [0,0,0] + assert geometry.sdf({'x': np.array([[0]]),'y':np.array([[0]]),'z':np.array([[0]])},params={})['sdf'][0][0] < 0 + assert abs(geometry.sdf({'x': np.array([[4]]),'y':np.array([[0]]),'z':np.array([[0]])},params={})['sdf'][0][0]) < 1e-12 + assert geometry.sdf({'x': np.array([[5.5]]),'y':np.array([[0.5]]),'z':np.array([[-0.5]])},params={})['sdf'][0][0] > 0 + + +def test_rotate_operation(circle_domain): + is_boundary, geometry, cond, translate_vec, rotation_list = TPGeometryWrapper(rotate_domain_dict(circle_domain,torch.tensor([[[-1.0000e+00, 0], + [0, -1.0000e+00]]]),torch.tensor([5,0]))).getModulusGeometry() + assert not is_boundary + assert cond == None + assert rotation_list == [[],[],[],0] + assert translate_vec == [0,0,0] + assert geometry.sdf({'x': np.array([[12]]),'y':np.array([[0]])},params={})['sdf'][0][0] < 0 + assert abs(geometry.sdf({'x': np.array([[11]]),'y':np.array([[0]])},params={})['sdf'][0][0]) <1e-12 + assert geometry.sdf({'x': np.array([[10]]),'y':np.array([[0.5]])},params={})['sdf'][0][0] > 0 + + +def test_product_operation(parallelogram_domain_xy,interval_domain_z): + is_boundary, geometry, cond, _, _ = TPGeometryWrapper(product_domain_dict(parallelogram_domain_xy,interval_domain_z)).getModulusGeometry() + assert not is_boundary + assert geometry.sdf({'x': np.array([[2]]),'y':np.array([[1]]),'z':np.array([[1]])},params={})['sdf'][0][0] < 0 + assert geometry.sdf({'x': np.array([[0.5]]),'y':np.array([[1]]),'z':np.array([[2]])},params={})['sdf'][0][0] > 0 + assert abs(geometry.sdf({'x': np.array([[1]]),'y':np.array([[1]]),'z':np.array([[1]])},params={})['sdf'][0][0]) <1e-12 + + assert cond == None + + +def test_multiple_operations(circle_domain_xy,interval_domain,interval_domain_y): + is_boundary, geometry, cond, translate_vec, rotation_list = TPGeometryWrapper(combined_operations_dict(circle_domain_xy,interval_domain,interval_domain_y,torch.tensor([1,-1]),torch.tensor([[[-1.0000e+00, 0],[0, -1.0000e+00]]]),torch.tensor([0,2]))).getModulusGeometry() + assert not is_boundary + assert cond == None + assert rotation_list == [[],[],[],0] + assert translate_vec == [0,0,0] + # left half circle with radius 1 and center at (1,3) + assert geometry.sdf({'x': np.array([[1.5]]),'y':np.array([[3]])},params={})['sdf'][0][0] < 0 + assert geometry.sdf({'x': np.array([[0.5]]),'y':np.array([[3]])},params={})['sdf'][0][0] > 0 + assert abs(geometry.sdf({'x': np.array([[1]]),'y':np.array([[3]])},params={})['sdf'][0][0]) <1e-12 + + diff --git a/src/torchphysics/wrapper/tests/test_model.py b/src/torchphysics/wrapper/tests/test_model.py new file mode 100644 index 00000000..1dc2f820 --- /dev/null +++ b/src/torchphysics/wrapper/tests/test_model.py @@ -0,0 +1,42 @@ +import pytest +import shutil +import os +import torch +from torchphysics.wrapper.model import ModulusArchitectureWrapper + +from torchphysics.problem.spaces import Space +from torchphysics.problem.spaces.points import Points + +@pytest.fixture +def input_space(): + return Space({'x': 2}) + +@pytest.fixture +def output_space(): + return Space({'output': 1}) + +@pytest.fixture +def modulus_architecture_wrapper(input_space, output_space): + return ModulusArchitectureWrapper(input_space, output_space, arch_name="fully_connected") + +def test_modulus_architecture_wrapper_init(modulus_architecture_wrapper): + assert isinstance(modulus_architecture_wrapper, ModulusArchitectureWrapper) + assert list(modulus_architecture_wrapper.input_space.keys()) == ['x'] + assert list(modulus_architecture_wrapper.output_space.keys()) == ['output'] + assert modulus_architecture_wrapper.input_space.dim == 2 + assert modulus_architecture_wrapper.output_space.dim == 1 + assert type(modulus_architecture_wrapper.modulus_net).__name__=='FullyConnectedArch' + +def test_modulus_architecture_wrapper_forward(modulus_architecture_wrapper): + in_vars = Points.from_coordinates({'x': torch.tensor([[1.0, 2.0], [3.0, 4.0]])}) + result = modulus_architecture_wrapper.forward(in_vars) + assert isinstance(result, Points) + assert list(result.coordinates.keys())==['output'] + assert result.coordinates['output'].shape == torch.Size([2,1]) + + +def teardown_module(module): + """This method is called after test completion.""" + conf_dir = os.path.abspath(os.getcwd() + '/conf') + if os.path.isdir(conf_dir): + shutil.rmtree(conf_dir) \ No newline at end of file diff --git a/src/torchphysics/wrapper/tests/test_wrapper.py b/src/torchphysics/wrapper/tests/test_wrapper.py new file mode 100644 index 00000000..62794608 --- /dev/null +++ b/src/torchphysics/wrapper/tests/test_wrapper.py @@ -0,0 +1,147 @@ +import pytest +import torchphysics as tp +import torch +import pytorch_lightning as pl +import os +import shutil +from tensorboard.backend.event_processing import event_accumulator + +from torchphysics.wrapper import TPModulusWrapper, ModulusSolverWrapper +from torchphysics.utils import PointsDataLoader +from torchphysics.problem.spaces import Points, R1 +from omegaconf import DictConfig +from modulus.sym.solver import Solver + + +def _create_dummies(): + fcn = tp.FCN(tp.spaces.R1('x'), tp.spaces.R1('u')) + ps = tp.samplers.RandomUniformSampler(tp.domains.Interval(tp.spaces.R1('x'), 0, 1), + n_points=10) + cond = tp.conditions.PINNCondition(fcn, ps, lambda u: u) + opti = tp.OptimizerSetting(optimizer_class=torch.optim.Adam, lr=0.1,) + solver = tp.Solver(train_conditions=[cond],optimizer_setting=opti) + trainer = pl.Trainer(max_steps=15) + + return fcn, trainer, solver + +def _create_dummies_with_params(): + model = tp.FCN(tp.spaces.R1('x'), tp.spaces.R1('u')) + ps = tp.samplers.RandomUniformSampler(tp.domains.Interval(tp.spaces.R1('x'), 0, 1), + n_points=10) + p = tp.models.parameter.Parameter(init=1.0, space=tp.spaces.R1('D')) + cond1 = tp.conditions.PINNCondition(model, ps, lambda u,D: u*D,parameter=p) + + loader = PointsDataLoader((Points(torch.tensor([[0.0], [2.0]]), R1('x')), + Points(torch.tensor([[0.0], [4.0]]), R1('u'))),batch_size=2) + + + cond2 = tp.conditions.DataCondition(module=model, dataloader=loader, norm=2) + + opti = tp.OptimizerSetting(optimizer_class=torch.optim.Adam, lr=0.1,) + solver = tp.Solver(train_conditions=[cond1,cond2],optimizer_setting=opti) + trainer = pl.Trainer(max_steps=15) + return trainer, solver, p + +def test_TPModulusWrapper_creation(): + _, trainer,solver = _create_dummies() + aggregator_args = {"alpha": 0.5} + scheduler_args = {"T_max": 999} + wrapper = TPModulusWrapper(trainer=trainer, solver=solver, lambda_weighting=["sdf"],aggregator="GradNorm", aggregator_args=aggregator_args,scheduler="CosineAnnealingLR", scheduler_args=scheduler_args,outputdir_name="my_output_dir",keep_output=True,confdir_name="my_conf_dir") + assert isinstance(wrapper, TPModulusWrapper) + + assert wrapper.cfg.loss._target_ == "modulus.sym.loss.aggregator.GradNorm" + assert wrapper.cfg.loss["alpha"] == 0.5 + assert wrapper.cfg.scheduler._target_ == "torch.optim.lr_scheduler.CosineAnnealingLR" + assert isinstance(wrapper.cfg, DictConfig) + assert wrapper.cfg.scheduler.T_max == 999 + assert wrapper.cfg.loss.alpha == 0.5 + assert wrapper.cfg.optimizer.lr == 0.1 + assert wrapper.cfg.network_dir == "my_output_dir" + assert wrapper.cfg.training.max_steps == 15 + assert wrapper.keep_output == True + assert wrapper.cfg.optimizer._target_ == "torch.optim.Adam" + + assert wrapper.Msolver.lambda_weighting_vals == ["sdf"] + assert isinstance(wrapper.Msolver,ModulusSolverWrapper) + assert isinstance(wrapper.slv,Solver) + + assert os.path.isdir(os.path.abspath(os.getcwd()+'/my_conf_dir')) + + + solver.optimizer_setting.scheduler_class=torch.optim.lr_scheduler.ExponentialLR + solver.optimizer_setting.scheduler_args={'gamma':0.8} + wrapper = TPModulusWrapper(trainer=trainer, solver=solver, lambda_weighting=["sdf"],aggregator="GradNorm", aggregator_args=aggregator_args,outputdir_name="my_output_dir",keep_output=False,confdir_name="my_conf_dir") + assert wrapper.cfg.scheduler._target_ == "torch.optim.lr_scheduler.ExponentialLR" + assert wrapper.cfg.scheduler.gamma == 0.8 + + + +def test_callbacks(): + model, trainer,solver = _create_dummies() + + trainer.callbacks.append(tp.utils.WeightSaveCallback(model= model,path=os.getcwd(),name='test',check_interval = 5, save_initial_model=True, + save_final_model = True)) + + with pytest.warns(UserWarning, match="The option check_interval of the WeightSaveCallback with check for minimial loss is not supported by Modulus. Only initial and final model state saves."): + wrapper = TPModulusWrapper(trainer=trainer, solver=solver,outputdir_name="my_output_dir",confdir_name="my_conf_dir") + wrapper.train() + + assert os.path.isfile(os.path.abspath(os.getcwd()+'/test_init.pt')) + assert os.path.isfile(os.path.abspath(os.getcwd()+'/test_final.pt')) + + model, trainer,solver = _create_dummies() + plot_sampler = tp.samplers.PlotSampler(tp.domains.Interval(tp.spaces.R1('x'), 0, 1), + n_points=50) + trainer.callbacks.append(tp.utils.PlotterCallback(model=model,plot_function=lambda u: u,point_sampler=plot_sampler, + log_name='plot_u', plot_type='plot',check_interval=5)) + wrapper = TPModulusWrapper(trainer=trainer, solver=solver,keep_output=True, outputdir_name="my_output_dir",confdir_name="my_conf_dir") + wrapper.train() + # path to tensorboard-events-file + log_dir = os.path.join(os.getcwd(), 'my_output_dir') + event_file = None + + # Check if tensorboard-events-file exists + for file in os.listdir(log_dir): + if file.startswith("events.out.tfevents"): + event_file = os.path.join(log_dir, file) + break + + assert event_file is not None, "tensorboard-events-file not found." + + # Usage of EventAccumulator to load events file + ea = event_accumulator.EventAccumulator(event_file) + ea.Reload() + + # Check if plot is contained + tags = ea.Tags()['images'] + assert 'Inferencers/plot_u/' in tags, "Plot 'plot_u' not found in tensorboard-events-file." + + + + +def test_train_function(): + _, trainer,solver = _create_dummies() + wrapper = TPModulusWrapper(trainer=trainer, solver=solver,outputdir_name="my_output_dir",confdir_name="my_conf_dir") + assert wrapper.train()==[] + assert os.path.isdir(os.path.abspath(os.getcwd()+'/my_output_dir')) + + trainer,solver, p = _create_dummies_with_params() + wrapper = TPModulusWrapper(trainer=trainer, solver=solver,keep_output=True, outputdir_name="my_output_dir",confdir_name="my_conf_dir") + wrapper.train() + assert abs(p.as_tensor- 1.0) >1e-15 + + + + +def teardown_module(module): + """This method is called after test completion.""" + conf_dir = os.path.abspath(os.getcwd() + '/my_conf_dir') + if os.path.isdir(conf_dir): + shutil.rmtree(conf_dir) + output_dir = os.path.abspath(os.getcwd() + '/my_output_dir') + if os.path.isdir(output_dir): + shutil.rmtree(output_dir) + if os.path.isfile(os.path.abspath(os.getcwd()+'/test_init.pt')): + os.remove(os.path.abspath(os.getcwd()+'/test_init.pt')) + if os.path.isfile(os.path.abspath(os.getcwd()+'/test_final.pt')): + os.remove(os.path.abspath(os.getcwd()+'/test_final.pt')) \ No newline at end of file diff --git a/src/torchphysics/wrapper/wrapper.py b/src/torchphysics/wrapper/wrapper.py new file mode 100644 index 00000000..450154d6 --- /dev/null +++ b/src/torchphysics/wrapper/wrapper.py @@ -0,0 +1,361 @@ +from modulus.sym.hydra.utils import compose +from modulus.sym.solver import Solver + +from torchphysics.models import Parameter +import torch + +import os +import csv +import numpy as np + +from .helper import OptimizerNameMapper, SchedulerNameMapper, AggregatorNameMapper +from .solver import ModulusSolverWrapper + +import shutil +import warnings +import logging +# Set the logging level for matplotlib to WARNING +matplotlib_logger = logging.getLogger('matplotlib') +matplotlib_logger.setLevel(logging.WARNING) + + +class TPModulusWrapper(): + ''' + Training of a TorchPhysics trainer/solver with the Modulus wrapper. + The wrapper is a bridge between TorchPhysics and Modulus. It uses + the Modulus configuration and the Modulus solver to train the + TorchPhysics solver/trainer/models. + Loss weighting algorithms can be selected by choosing an + aggregation function. The aggregation function can be selected by + the parameter "aggregator" and additional arguments can be set by + the parameter "aggregator_args". + A learning rate scheduler can be selected by the parameter + "scheduler" and additional arguments can be set by the parameter + "scheduler_args". + Pointwise weighting of the loss can be set by the parameter + "lambda_weighting". The pointwise weighting can be a list of + numbers or sympy expressions or the string 'sdf'. + Notes + ----- + The following conventions are important for the usage of the + wrapper: + Possible spatial variables: x, y, z or x (multidimensional) + Time variable: t + Geometries: TP geometries (domains) should be defined in the + space of x, y, z or x (multidimensional). + A general product of domains as in TorchPhysics + can not be implemented in Modulus, because the + domain defines a spatial geometry that must have a + sdf implementation which is not available in + general for an arbitrary product domain. + Cross products of domains and domain operations are + generally allowed, but too complicated + constructions should be avoided, e.g. a cross + product of 3 translated intervals is not allowed or + only isosceles triangle with axis of symmetry + parallel to y-axis in cross product with an interval + are supported. + Shapely polygons in 2D are supported, but currently + 3D geometries (TrimeshPolyhedron) defined in stl- + files are only supported in Modulus in the container + installation. + Translation of primitive domains is supported, but + not translation of domains resulting from domain + operations like union, intersection, difference. + + Parameters + ---------- + trainer : pytorch_lightning.Trainer + The Pytorch Lightning Trainer instance. + Supported parameters of trainer instance: + Modulus always uses GPU device if available. Trainer + settings concerning GPU devices or cuda handling, e.g. + 'accelerator' or 'devices', are not supported by this + wrapper. + Modulus automatically logs the training process with + tensorboard. The tensorboard logs are saved in the output + directory. + All TorchPhysics callbacks are supported by the wrapper. + The following Trainer parameters are supported by this + wrapper: + 'max_steps' : int, optional + The maximum number of training steps. If not + specified, the default value of Pytorch Lightning + Trainer is used. + 'val_check_interval' : int optional + How often to check the validation set. Default is + 1.0, meaning once per training epoch. + 'log_every_n_steps' : int, optional + How often to log within steps. Modulus/wrapper + default is 50. + Checkpoints, progress bar and model summary are + automatically used by Modulus. + + solver: torchphysics.solvers.Solver + The TorchPhysics solver instance. + All parameters of the TorchPhysics solver are supported by the + wrapper. + outputdir_name : str, optional + The name of the Modulus output directory, where the trained + models, the optimization configuration, tensorboard files, etc. + are saved. Default is 'outputs'. + If the directory contains the results of a previous run and the + configuration of a second call is mainly changed, there will be + a conflict loading existing models or configuration leading to + an error. + If the directory contains the results of a previous run and the + configuration of a second call is mainly unchanged, the new run + will continue the previous run with the already trained Modulus + models. + If not desired or in error case, it is recommended to remove + the content of the output directory before starting a new run. + confdir_name : str, optional + The name of a Modulus configuration directory, where initially + a hydra configuration file is saved. It is overwritten on each + call. Default is 'conf'. + keep_output : bool, optional. Default is True. + If True, the output directory is not deleted after the training + process. Otherwise, it is deleted after the training process. + **kwargs : optional + Additional keyword arguments: + "lambda_weighting": list[Union[int, float, sp.Basic]]=None + The spatial pointwise weighting of the constraint. It + is a list of numbers or sympy expressions or the string + 'sdf'. + If the list has more than one element, the length of + the list and the order has to match the number of + TorchPhysics conditions in the call + tp.solver.Solver([condition_1,condition_2,...]). + If the list has only one element, the same weighting is + applied to all conditions. + If it is a sympy expression, it has to be a function of + the spatial coordinates x, y, z. + If the TorchPhysics conditions contain weight + definitions with the keyword "weight", these are + additionally applied. + For example, + 'lambda_weighting=["sdf"]' would apply a pointwise + weighting of the loss by the signed distance function, + but only for interior sampling, not boundary sampling. + 'lambda_weighting=[100.0, 2.0] would apply a pointwise + weighting of the loss by 100 to the first TorchPhysics + condition and 2 to the second TorchPhysics condition. + 'lambda_weighting=[2.0*sympy.Symbol('x')]' would apply + a pointwise weighting to the loss of `2.0 * x`. + "aggregator" : str = None + The aggregation function for the loss. It is a string + with the name of the aggregation function. Default is + 'Sum'. + Possible values are 'Sum', 'GradNorm', 'ResNorm, + 'Homoscedastic','LRAnnealing','SoftAdapt','Relobralo'. + "aggregator_args" : dict = None + Additional arguments for the aggregation function. It + is a dictionary with the argument names as keys and the + argument values as values. Default is None. + Possible arguments with its default values are, + depending on the aggregator: + GradNorm: + alpha = 1.0 + ResNorm: + alpha = 1.0 + LRAnnealing: + update_freq = 1 + alpha = 0.01 + ref_key = None # Change to Union[None, str] when + supported by hydra + eps = 1e-8 + SoftAdapt: + eps = 1e-8 + Relobralo: + alpha = 0.95 + beta = 0.99 + tau = 1.0 + eps = 1e-8 + "scheduler" : str = None + The learning rate scheduler. It is a string with the + name of the scheduler. Default is constant learning + rate. + Possible values are 'ExponentialLR', + 'CosineAnnealingLR' or 'CosineAnnealingWarmRestarts'. + "scheduler_args" : dict = None + Additional arguments for the scheduler. It is a + dictionary with the argument names as keys and the + argument values as values. Default is None. + Possible arguments with its default values are, + depending on the scheduler: + ExponentialLR: + gamma = 0.99998718 + TFExponentialLR: + decay_rate = 0.95 + decay_steps = 1000 + CosineAnnealingLR: + T_max = 1000 + eta_min = 0.0 + last_epoch= -1 + CosineAnnealingWarmRestarts: + T_0 = 1000 + T_mult = 1 + eta_min = 0.0 + last_epoch = -1 + ''' + def __init__(self,trainer,solver,outputdir_name = 'outputs',confdir_name = 'conf',keep_output = True, **kwargs): + self.outputdir_name = outputdir_name + self.keep_output = keep_output + + self.logger=logging.getLogger() + self.ch = logging.StreamHandler() + self.logger.addHandler(self.ch) + + # get the absolute path of the conf-directory + caller_path=os.path.abspath(os.getcwd()+'/'+confdir_name) + + # Get the relative path of the current file to the conf-directory + current_path = os.path.relpath(caller_path,os.path.dirname(os.path.abspath(__file__))) + + if 'aggregator' in kwargs.keys(): + aggregator_name = AggregatorNameMapper(kwargs['aggregator']) + assert (aggregator_name != 'not defined'), "This aggregator class is currently not supported by Modulus!" + else: + aggregator_name = AggregatorNameMapper('Sum') + + optimizer_name = OptimizerNameMapper(solver.optimizer_setting.optimizer_class .__name__) + assert (optimizer_name != 'not defined'), "This optimizer class is currently not supported by Modulus!" + assert ((solver.optimizer_setting.scheduler_class is None) or (kwargs.get('scheduler') is None)), "The scheduler should either be defined in the optimizer settings or as additional parameter of the TPModulusWrapper!" + if solver.optimizer_setting.scheduler_class is not None: + scheduler_name = SchedulerNameMapper(solver.optimizer_setting.scheduler_class.__name__) + elif kwargs.get('scheduler') is not None: + scheduler_name = SchedulerNameMapper(kwargs.get('scheduler')) + else: + scheduler_name = 'exponential_lr' + + assert (scheduler_name != 'not defined'), "This scheduler class is currently not supported by Modulus!" + + os.makedirs(caller_path, exist_ok=True) + with open(caller_path+'/config_Modulus.yaml', 'w') as f: + f.write('defaults :\n - modulus_default\n - loss: '+aggregator_name+'\n - optimizer: '+optimizer_name+'\n - scheduler: '+scheduler_name+'\n - _self_\n') + self.cfg = compose(config_path=current_path, config_name="config_Modulus") + + training_rec_results_freq = self.cfg.training.rec_results_freq + + # as the initialization without scheduler leads to an error and additionally the constant LR scheduler is not implemented as class, + # we use the exponential LR scheduler with gamma=1 for constant lr + if (solver.optimizer_setting.scheduler_class is None) and (kwargs.get('scheduler') is None): + self.cfg.scheduler.gamma = 1.0 + else: + if solver.optimizer_setting.scheduler_args is not None: + for key, value in solver.optimizer_setting.scheduler_args.items(): + self.cfg.scheduler[key]=value + if kwargs.get('scheduler_args') is not None: + for key, value in kwargs.get('scheduler_args').items(): + self.cfg.scheduler[key]=value + + assert (solver.optimizer_setting.scheduler_frequency == 1), "The scheduler frequency is not supported in Modulus!" + + for key, value in solver.optimizer_setting.optimizer_args.items(): + self.cfg.optimizer[key]=value + + self.cfg.network_dir = self.outputdir_name + + self.cfg.training.max_steps = trainer.max_steps + self.cfg.optimizer.lr = solver.optimizer_setting.lr + + + self.cfg.training.rec_results_freq = min(training_rec_results_freq,trainer.max_steps) + + if (type(trainer.val_check_interval)) is float: + self.cfg.training.rec_validation_freq =int(trainer.val_check_interval*self.cfg.training.max_steps) + else: + self.cfg.training.rec_validation_freq = trainer.val_check_interval + self.cfg.training.summary_freq = trainer.log_every_n_steps + + self.weight_save_callback = None + self.checkpoint_callback = None + callbacks2Modulus=[] + + for callback in trainer.callbacks: + if type(callback).__name__=='WeightSaveCallback': + self.weight_save_callback = callback + if self.weight_save_callback.check_interval >0: + warnings.warn('The option check_interval of the WeightSaveCallback with check for minimial loss is not supported by Modulus. Only initial and final model state saves.') + elif type(callback).__name__=='PlotterCallback': + self.cfg.training.rec_inference_freq = callback.check_interval + callbacks2Modulus.append(callback) + elif type(callback).__name__=='TrainerStateCheckpoint': + self.checkpoint_callback = callback + self.cfg.training.save_network_freq = self.checkpoint_callback.check_interval + warnings.warn('TorchPhysics TrainerStateCheckpoint callback is requested. The checkpointing will be automatically done by Modulus and training can be restarted. The option weights_only is not supported by Modulus. Please use the WeightSaveCallback instead.') + + + + self.Msolver=ModulusSolverWrapper(solver,callbacks2Modulus,**kwargs) + + # adapt cfg-parameters + for key, value in kwargs.items(): + if key == 'aggregator_args': + for key2, value2 in value.items(): + self.cfg.loss[key2] = value2 + + + # Modulus solver instance is created + self.slv = Solver(self.cfg, self.Msolver.domain) + + + def train(self,resume_from_ckpt=False): + ''' + Call the training process of the Modulus solver. The training + process is started with the Modulus configuration and the + Modulus solver. + The TorchPhysics models are trained and the function + additionally returns trained parameters. + + Parameters + ---------- + resume_from_ckpt: bool, optional. Default is False. + If True, the training is resumed from the Modulus checkpoint files. + + ''' + # if a TorchPhysics checkpoint callback exists and the training is started without resume option, delete Modulus checkpoint and model files in the Modulus output directory. + if not resume_from_ckpt: + if os.path.isdir(os.path.abspath(os.getcwd()+'/'+self.outputdir_name)): + shutil.rmtree(os.path.abspath(os.getcwd()+'/'+self.outputdir_name)) + + # save initial model before training + if self.weight_save_callback: + if self.weight_save_callback.save_initial_model: + torch.save(self.weight_save_callback.model.state_dict(), self.weight_save_callback.path+'/' + self.weight_save_callback.name + '_init.pt') + # start training + self.slv.solve() + + + for model in self.Msolver.models: + model.to('cpu') + for model in self.Msolver.orig_models: + model.to('cpu') + + result_vec=[] + for index, param_obj in enumerate(zip(self.Msolver.parameters,self.Msolver.parameter_nets)): + for outvar in param_obj[1].output_keys: + with open(os.path.abspath(os.getcwd()+'/'+self.outputdir_name+'/monitors/mean_'+str(outvar)+'.csv'), 'r') as f: + reader = csv.reader(f) + key_line = next(reader) + last_line = list(reader)[-1] + result_vec.append(float(last_line[1])) + + self.Msolver.parameters[index]._t=torch.tensor([result_vec]) + if self.weight_save_callback: + if self.weight_save_callback.save_final_model: + torch.save(self.weight_save_callback.model.state_dict(), self.weight_save_callback.path+'/' + self.weight_save_callback.name + '_final.pt') + + # if no TorchPhysics checkpoint callback exists, delete Modulus checkpoint and model files in the Modulus output directory + if not self.checkpoint_callback and not self.keep_output: + if os.path.isdir(os.path.abspath(os.getcwd()+'/'+self.outputdir_name)): + shutil.rmtree(os.path.abspath(os.getcwd()+'/'+self.outputdir_name)) + result = self.Msolver.parameters + + + self.logger.removeHandler(self.ch) + + return result + + +