I have been trying to develop and run a custom model for my physics informed NN, but I don’t understand the issue faced
Script:
# libraries
from sympy import Symbol, Function, exp, Number
from modulus.sym.eq.pde import PDE
from modulus.sym.node import Node
from sympy import Symbol, Eq
import numpy as np
import modulus.sym
from modulus.sym.hydra import instantiate_arch, ModulusConfig
from modulus.sym.models.layers import Activation
from modulus.sym.solver import Solver
from modulus.sym.domain import Domain
from modulus.sym.geometry.primitives_2d import Rectangle, Circle
from modulus.sym.domain.constraint import (
PointwiseBoundaryConstraint,
PointwiseInteriorConstraint,
)
from modulus.sym.domain.validator import PointwiseValidator
from modulus.sym.domain.inferencer import PointwiseInferencer
from modulus.sym.key import Key
from modulus.sym.node import Node
from modulus.sym.models.fully_connected import FullyConnectedArch
from modulus.sym.models.deeponet import DeepONetArch
E = 10
nu = 0.33
class damage_model(PDE):
name = "damage_model"
def __init__(self, E, nu, rho = 1e4, beta=1/8100, r1 = 1e4, r2 = 1e4, dim =2, time =True):
self.sim = dim
self.time = time
# space coordinates
x = Symbol("x")
y = Symbol("y")
# plane normals
normal_x, normal_y= (
Symbol("normal_x"),
Symbol("normal_y"),
)
# time coordinate
t = Symbol("t")
# functionality
input_variables = {"x": x, "y": y, "t": t}
# material properties
if E is None:
E = Function("E", positive = True)(*input_variables)
if E is None:
nu = Function("nu", positive = True)(*input_variables)
# LE terms
lambda_ = nu * E / ((1 + nu) * (1 - 2 * nu))
mu = E / (2 * (1 + nu))
# displacement fields
u = Function("u")(*input_variables)
v = Function("v")(*input_variables)
# damage field
d = Function("d", nonnegative = True)(*input_variables)
# damage function:
f = 1 - exp(-d)
# strains
e_xx = u.diff(x)
e_yy = v.diff(y)
e_xy = 0.5*(v.diff(x) + u.diff(y))
# stress fields
w_z = -lambda_ / (lambda_ + 2 * mu) * (u.diff(x) + v.diff(y))
sigma_xx = lambda_ * (u.diff(x) + v.diff(y) + w_z) + 2 * mu * u.diff(x)
sigma_yy = lambda_ * (u.diff(x) + v.diff(y) + w_z) + 2 * mu * v.diff(y)
sigma_xy = mu * (u.diff(y) + v.diff(x))
# set equations
self.equations = {}
# stress equations as per Linear Elasticity
self.equations["LE_sigma_xx"] = sigma_xx
self.equations["LE_sigma_yy"] = sigma_yy
self.equations["LE_sigma_xy"] = sigma_xy
# Equations of equilibrium
self.equations["eq_x"] = rho * ((u.diff(t)).diff(t)) - (sigma_xx.diff(x) + sigma_xy.diff(y))
self.equations["eq_y"] = rho * ((v.diff(t)).diff(t)) - (sigma_xy.diff(x) + sigma_yy.diff(y))
# Traction equations
self.equations["tr_x"] = normal_x * sigma_xx + normal_y * sigma_xy
self.equations["tr_y"] = normal_x * sigma_xy + normal_y * sigma_yy
# LE energy: phi_{0}
self.equations["phi_0"] = 0.5*(sigma_xx*e_xx + sigma_yy*e_yy + 2*sigma_xy*e_xy)
# damge function
self.equations["f"] = f
# distorsion energy
self.equations["phi"] = f*self.equations["phi_0"] + 0.5*beta*((f.diff(x))**2 + (f.diff(y))**2)
# /delta^{diss}
self.equations["diss"] = r1*(((f.diff(t))**2)**0.5) + 0.5*r2*((f.diff(t))**2)
# p
self.equations["p"] = -1*self.equations["phi_0"] + beta*(f.diff(x, 2) + f.diff(y, 2))
# p_hat
self.equations["p_hat"] = self.equations["p"] - r1
# p_plus
self.equations["p_plus"] = 0.5*(self.equations["p_hat"] + (self.equations["p_hat"]**2)**0.5)
# DM
self.equations["DM"] = r2*(f.diff(t)) + self.equations["p_plus"]
# running the model
@modulus.sym.main(config_path="config", config_name="conf")
def run(cfg: ModulusConfig) -> None:
equations = damage_model(E = E, nu = nu)
damage_net = instantiate_arch(
input_keys = [
Key("x"),
Key("y"),
Key("t")
],
output_keys = [
Key("u"),
Key("v"),
Key("d")
],
cfg = cfg.arch.fully_connected,
activation_fn=Activation.TANH,
)
nodes = equations.make_nodes() + [damage_net.make_node(name="damage_net")]
panel_origin = (0, 0)
panel_left_bottom = (-1, -2)
panel_right_top = (1, 2)
panel = Rectangle(panel_left_bottom, panel_right_top)
cut = Circle(panel_origin, radius = 0.2)
geo = panel - cut
x = Symbol("x")
y = Symbol("y")
t_symbol = Symbol("t")
bound_x = (panel_left_bottom[0], panel_right_top[0])
bound_y = (panel_left_bottom[1], panel_right_top[1])
bound_t = (0.0, 1.0)
time_range = {t_symbol: bound_t}
domain = Domain()
left = PointwiseBoundaryConstraint(
nodes = nodes,
geometry = geo,
outvar = {
"tr_x" : 0.0,
"tr_y": 0.0,
},
batch_size = cfg.batch_size.panel_left,
criteria = Eq(x, panel_left_bottom[0]),
parameterization=time_range,
batch_per_epoch=500,
)
domain.add_constraint(left, "BC_left")
right = PointwiseBoundaryConstraint(
nodes = nodes,
geometry = geo,
outvar = {
"tr_x" : 0.0,
"tr_y": 0.0,
},
batch_size = cfg.batch_size.panel_right,
criteria = Eq(x, panel_right_top[0]),
parameterization=time_range,
batch_per_epoch=500,
)
domain.add_constraint(right, "BC_right")
cutting = PointwiseBoundaryConstraint(
nodes = nodes,
geometry = cut,
outvar = {
"tr_x" : 0.0,
"tr_y": 0.0,
},
batch_size = cfg.batch_size.panel_right,
parameterization=time_range,
batch_per_epoch=500,
)
domain.add_constraint(cutting, "BC_cutting")
panel_bottom = PointwiseBoundaryConstraint(
nodes=nodes,
geometry=geo,
outvar = {"v": 0.0, "d": 0.0},
batch_size = cfg.batch_size.panel_bottom,
criteria = Eq(y, panel_left_bottom[1]),
parameterization=time_range,
batch_per_epoch=500,
)
domain.add_constraint(panel_bottom, "panel_bottom")
panel_top = PointwiseBoundaryConstraint(
nodes=nodes,
geometry=geo,
outvar={"v": 0.01*t_symbol},
batch_size = cfg.batch_size.panel_top,
criteria = Eq(y, panel_right_top[1]),
parameterization=time_range,
batch_per_epoch=500,
)
domain.add_constraint(panel_top, "panel_top")
IC = PointwiseInteriorConstraint(
nodes = nodes,
geometry = geo,
outvar={"u": 0.0, "v":0.0, "d": 0.0},
batch_size=cfg.batch_size.interior,
lambda_weighting={
"u": 100.0,
"v": 100.0,
"d": 100.0,
},
parameterization=time_range,
batch_per_epoch=500,
)
domain.add_constraint(IC, "IC")
# low-resolution interior
lr_interior = PointwiseInteriorConstraint(
nodes=nodes,
geometry=geo,
outvar={
"eq_x": 0.0,
"eq_y": 0.0,
# "stress_disp_xx": 0.0,
# "stress_disp_yy": 0.0,
# "stress_disp_xy": 0.0,
"DM": 0.0
},
batch_size=cfg.batch_size.interior,
bounds={x: bound_x, y: bound_y,},
lambda_weighting={
"eq_x": Symbol("sdf"),
"eq_y": Symbol("sdf"),
# "stress_disp_xx": Symbol("sdf"),
# "stress_disp_yy": Symbol("sdf"),
# "stress_disp_xy": Symbol("sdf"),
"DM": Symbol("sdf"),
},
parameterization=time_range,
)
domain.add_constraint(lr_interior, "lr_interior")
for i, specific_time in enumerate(np.linspace(0, bound_t[1], 10)):
invar_numpy = geo.sample_interior(
1000000,
bounds={x: bound_x, y: bound_y},
parameterization={t_symbol: specific_time},
)
grid_inference = PointwiseInferencer(
nodes=nodes,
invar=invar_numpy,
output_names=["u", "v", "d", "f", "phi", "LE_sigma_xx", "LE_sigma_yy", "LE_sigma_xy"],
batch_size=4096,
)
domain.add_inferencer(grid_inference, name="time_slice_" + str(i+1).zfill(3))
# make solver
slv = Solver(cfg, domain)
# start solver
slv.solve()
if __name__ == "__main__":
run()
The error I am facing:
grad_outputs: List[Optional[torch.Tensor]] = [torch.ones_like(y, device=y.device)]
grad = torch.autograd.grad(
~~~~~~~~~~~~~~~~~~~ <--- HERE
[
y,
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
which arises in the side script:
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# 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 itertools
import torch
import numpy as np
import logging
from torch.autograd import Function
from modulus.sym.constants import diff
from modulus.sym.key import Key
from modulus.sym.node import Node
from modulus.sym.eq.mfd import FirstDeriv, SecondDeriv, ThirdDeriv, ForthDeriv
from typing import Dict, List, Set, Optional, Union, Callable
Tensor = torch.Tensor
logger = logging.getLogger(__name__)
# ==== Autodiff ====
@torch.jit.script
def gradient(y: torch.Tensor(), x: List[torch.Tensor]) -> List[torch.Tensor]:
"""
TorchScript function to compute the gradient of a tensor wrt multiple inputs
"""
grad_outputs: List[Optional[torch.Tensor]] = [torch.ones_like(y, device=y.device)]
grad = torch.autograd.grad(
[
y,
],
x,
grad_outputs=grad_outputs,
create_graph=True,
allow_unused=True,
)
if grad is None:
grad = [torch.zeros_like(xx) for xx in x]
assert grad is not None
grad = [g if g is not None else torch.zeros_like(x[i]) for i, g in enumerate(grad)]
return grad
class Derivative(torch.nn.Module):
"""
Module to compute derivatives using backward automatic differentiation
"""
def __init__(self, bwd_derivative_dict: Dict[Key, List[Key]]):
"""
Constructor of the Derivative class.
Parameters
----------
inputs : List[Key]
A list of keys of the available variables to compute the required variables
This list should contain both the variables that need to be differentiated
and the variables to differentiate with respect to.
derivatives : List[Key]
A list of keys of the required derivatives
"""
super().__init__()
self.gradient_dict: Dict[str, Dict[str, int]] = {
str(k): {str(w): w.size for w in v} for k, v in bwd_derivative_dict.items()
}
self.gradient_names: Dict[str, List[str]] = {
k: [diff(k, der) for der in v.keys()] for k, v in self.gradient_dict.items()
}
self.nvtx_str: str = f"Auto-Diff Node: {list(self.gradient_dict.keys())}"
@staticmethod
def prepare_input(
input_variables: Dict[str, torch.Tensor], mask: List[str]
) -> List[torch.Tensor]:
return [input_variables[x] for x in mask]
@staticmethod
def dict_output(
output_tensors: List[torch.Tensor], sizes: List[str], var_name: str
) -> Dict[str, torch.Tensor]:
return {diff(var_name, name): output_tensors[i] for i, name in enumerate(sizes)}
def forward(self, input_var: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
output_var = {}
for var_name, grad_sizes in self.gradient_dict.items():
var = input_var[var_name]
grad_var = self.prepare_input(input_var, grad_sizes.keys())
grad = gradient(var, grad_var)
grad_dict = {
name: grad[i] for i, name in enumerate(self.gradient_names[var_name])
}
output_var.update(grad_dict)
return output_var
@classmethod
def make_node(cls, inputs: List[Key], derivatives: List[Key], name=None, jit=True):
derivatives = [d for d in derivatives if d not in inputs]
bwd_derivative_dict = _derivative_dict(inputs, derivatives, forward=False)
output_derivatives = []
for key, value in bwd_derivative_dict.items():
output_derivatives += [
Key(key.name, key.size, key.derivatives + [x]) for x in value
]
evaluate = cls(bwd_derivative_dict)
nvtx_str = evaluate.nvtx_str
if jit:
evaluate = torch.jit.script(evaluate)
derivative_node = Node(
inputs,
output_derivatives,
evaluate,
name=(nvtx_str if name is None else str(name)),
)
return derivative_node
def _derivative_dict(var, derivatives, forward=False):
needed = derivatives
while True: # break apart diff to see if first order needed
break_loop = True
for n in needed:
l_n = Key(n.name, n.size, n.derivatives[:-1])
if (len(n.derivatives) > 1) and l_n not in needed and l_n not in var:
needed.append(l_n)
break_loop = False
if break_loop:
break
current = var
diff_dict = {}
for c, n in itertools.product(current, needed):
c_under_n = Key(n.name, n.size, n.derivatives[0 : len(c.derivatives)])
if (c == c_under_n) and (len(n.derivatives) == len(c.derivatives) + 1):
if forward:
if n.derivatives[len(c.derivatives)] not in diff_dict:
diff_dict[n.derivatives[len(c.derivatives)]] = set()
diff_dict[n.derivatives[len(c.derivatives)]].add(c)
else:
if c not in diff_dict:
diff_dict[c] = set()
diff_dict[c].add(n.derivatives[len(c.derivatives)])
diff_dict = {key: list(value) for key, value in diff_dict.items()}
return diff_dict
# ==== Meshless finite derivs ====
class MeshlessFiniteDerivative(torch.nn.Module):
"""
Module to compute derivatives using meshless finite difference
Parameters
----------
model : torch.nn.Module
Forward torch module for calculating stencil values
derivatives : List[Key]
List of derivative keys to calculate
dx : Union[float, Callable]
Spatial discretization of all axis, can be function with parameter `count` which is
the number of forward passes for dynamically adjusting dx
order : int, optional
Order of derivative, by default 2
max_batch_size : Union[int, None], optional
Max batch size of stencil calucations, by default uses batch size of inputs
double_cast : bool, optional
Cast fields to double precision to calculate derivatives, by default True
jit : bool, optional
Use torch script for finite deriv calcs, by default True
"""
def __init__(
self,
model: torch.nn.Module,
derivatives: List[Key],
dx: Union[float, Callable],
order: int = 2,
max_batch_size: Union[int, None] = None,
double_cast: bool = True,
input_keys: Union[List[Key], None] = None,
):
super().__init__()
self.model = model
self._dx = dx
self.double_cast = double_cast
self.max_batch_size = max_batch_size
self.input_keys = input_keys
self.count = 0
self.derivatives = {1: [], 2: [], 3: [], 4: []}
for key in derivatives:
try:
self.derivatives[len(key.derivatives)].append(key)
except:
raise NotImplementedError(
f"{len(key.derivatives)}th derivatives not supported"
)
self.first_deriv = FirstDeriv(self.derivatives[1], self.dx, order=order)
self.second_deriv = SecondDeriv(self.derivatives[2], self.dx, order=order)
self.third_deriv = ThirdDeriv(self.derivatives[3], self.dx, order=order)
self.forth_deriv = ForthDeriv(self.derivatives[4], self.dx, order=order)
@torch.jit.ignore()
def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
self.count += 1
dx = self.dx
self.first_deriv.dx = dx
self.second_deriv.dx = dx
self.third_deriv.dx = dx
self.forth_deriv.dx = dx
torch.cuda.nvtx.range_push(f"Calculating meshless finite derivatives")
# Assemble global stencil
global_stencil = []
for deriv in [
self.first_deriv,
self.second_deriv,
self.third_deriv,
self.forth_deriv,
]:
stencil_list = deriv.stencil
# Remove centered stencil points if already in input dictionary
for i, point in enumerate(stencil_list):
if point.split("::")[1] == str(0) and point.split("::")[0] in inputs:
stencil_list.pop(i)
global_stencil.extend(stencil_list)
global_stencil = list(set(global_stencil))
# Number of stencil points to fit into a forward pass
input_batch_size = next(iter(inputs.values())).size(0)
if self.max_batch_size is None:
num_batch = 1
else:
num_batch = max([self.max_batch_size, input_batch_size]) // input_batch_size
# Stencil forward passes
index = 0
finite_diff_inputs = inputs.copy()
while index < len(global_stencil):
torch.cuda.nvtx.range_push(f"Running stencil forward pass")
# Batch up stencil inputs
stencil_batch = [global_stencil[index]]
index += 1
for j in range(1, min([len(global_stencil) - (index - 1), num_batch])):
stencil_batch.append(global_stencil[index])
index += 1
model_inputs = self._get_stencil_input(inputs, stencil_batch)
# Model forward
outputs = self.model(model_inputs)
# Dissassemble batched inputs
for key, value in outputs.items():
outputs[key] = torch.split(value.view(-1, len(stencil_batch)), 1, dim=1)
for i, stencil_str in enumerate(stencil_batch):
for key, value in outputs.items():
finite_diff_inputs[f"{key}>>{stencil_str}"] = value[i]
torch.cuda.nvtx.range_pop()
# Calc finite diff grads
torch.cuda.nvtx.range_push(f"Calc finite difference")
if self.double_cast: # Cast tensors to doubles for finite diff calc
for key, value in finite_diff_inputs.items():
finite_diff_inputs[key] = value.double()
outputs_first = self.first_deriv(finite_diff_inputs)
outputs_second = self.second_deriv(finite_diff_inputs)
outputs_third = self.third_deriv(finite_diff_inputs)
outputs_forth = self.forth_deriv(finite_diff_inputs)
outputs = inputs
if self.double_cast:
dtype = torch.get_default_dtype()
for key, value in outputs_first.items():
outputs_first[key] = value.type(dtype)
for key, value in outputs_second.items():
outputs_second[key] = value.type(dtype)
for key, value in outputs_third.items():
outputs_third[key] = value.type(dtype)
for key, value in outputs_forth.items():
outputs_forth[key] = value.type(dtype)
outputs = {
**inputs,
**outputs_first,
**outputs_second,
**outputs_third,
**outputs_forth,
}
torch.cuda.nvtx.range_pop()
torch.cuda.nvtx.range_pop()
return outputs
@property
def dx(self):
if hasattr(self._dx, "__call__"):
return self._dx(self.count)
else:
return self._dx
def _get_stencil_input(
self, inputs: Dict[str, Tensor], stencil_strs: List[str]
) -> Dict[str, Tensor]:
"""Creates a copy of the inputs tensor and adjusts its values based on
the stencil str.
Parameters
----------
inputs : Dict[str, Tensor]
Input tensor dictionary
stencil_strs : List[str]
batch list of stencil string from derivative class
Returns
-------
Dict[str, Tensor]
Modified input tensor dictionary
Example
-------
A stencil string `x::1` will modify inputs['x'] = inputs['x'] + dx
A stencil string `y::-1,z::1` will modify inputs['y'] = inputs['y'] - dx, inputs['z'] = inputs['z'] + dx
"""
if self.input_keys is None:
outputs = inputs.copy()
else:
outputs = {str(key): inputs[str(key)].clone() for key in self.input_keys}
for key, value in outputs.items():
outputs[key] = value.repeat(1, len(stencil_strs))
for i, stencil_str in enumerate(stencil_strs):
# Loop through points
for point in stencil_str.split("&&"):
var_name = point.split("::")[0]
spacing = int(point.split("::")[1])
outputs[var_name][:, i] = outputs[var_name][:, i] + spacing * self.dx
for key, value in outputs.items():
outputs[key] = value.view(-1, 1)
return outputs
@classmethod
def make_node(
cls,
node_model: Union[Node, torch.nn.Module],
derivatives: List[Key],
dx: Union[float, Callable],
order: int = 2,
max_batch_size: Union[int, None] = None,
name: str = None,
double_cast: bool = True,
input_keys: Union[List[Key], List[str], None] = None,
):
"""Makes a meshless finite derivative node.
Parameters
----------
node_model : Union[Node, torch.nn.Module]
Node or torch.nn.Module for computing FD stencil values.
Part of the inputs to this model should consist of the independent
variables and output the functional value
derivatives : List[Key]
List of derivatives to be computed
dx : Union[float, Callable]
Spatial discretization for finite diff calcs, can be function
order : int, optional
Order of accuracy of finite diff calcs, by default 2
max_batch_size : Union[int, None], optional
Maximum batch size to used with the stenicl foward passes, by default None
name : str, optional
Name of node, by default None
double_cast : bool, optional
Cast tensors to double precision for derivatives, by default True
input_keys : Union[List[Key], List[str], None], optional
List of input keys to be used for input of forward model.
Should be used if node_model is not a :obj:`Node`, by default None
"""
# We have two sets of input keys:
# input_keys: which are the list of inputs to the model for stencil points
# mfd_input_keys: input keys for the MFD node
if input_keys is None:
input_keys = []
mfd_input_keys = []
else:
input_keys = [str(key) for key in input_keys]
mfd_input_keys = [str(key) for key in input_keys]
for derivative in derivatives:
mfd_input_keys.append(derivative.name)
for dstr in derivative.derivatives:
mfd_input_keys.append(dstr.name)
input_keys.append(dstr.name)
if isinstance(node_model, Node):
model = node_model.evaluate
input_keys = input_keys + [str(key) for key in node_model.inputs]
else:
model = node_model
# Remove duplicate keys
mfd_input_keys = Key.convert_list(list(set(mfd_input_keys)))
input_keys = Key.convert_list(list(set(input_keys)))
evaluate = cls(
model,
derivatives,
dx=dx,
order=order,
max_batch_size=max_batch_size,
double_cast=double_cast,
input_keys=input_keys,
)
derivative_node = Node(
mfd_input_keys,
derivatives,
evaluate,
name=(
"Meshless-Finite-Derivative Node" + "" if name is None else f": {name}"
),
)
return derivative_node
You see, the variable “y” is defined in the side script, and I am unable to enable require_grad as recommended online.
Please help me out with this issue. I would be really obliged if you could please help me out @ngeneva
According to me this file should run without any issues
Thank you