Defining material models using Pytorch
COMMET allows for the use of material models implemented using pytorch. This is done by using torchscript (see Torchscript docs).
The learning outcomes are as follows:
How to define a model using Pytorch that can later be used in a COMMET simulation.
How to export a Pytorch model using torchscript.
How to read a torchscript file into a commet simulation.
Example files
The files for running the example can be downloaded as a single zip file here and it contains the following files:
Problem definition
The problem definition is identical to the Hello World example (see Hello world: running a simulation without an NCM). However, we will implement the Neohookean material model in Python using Pytorch, export the model as torchscript, and read in the model in the COMMET solver.
Defining a material model using Pytorch
We assume that a reader has some prior experience with Pytorch, in particular, one needs to be able to construct neural networks using Pytorch (see, e.g. Creating an NN in using Pytorch).
To implement a material model with Pytorch for later use in COMMET one must write a class that inherits from torch.nn.Module.
That class will then be serialized into a torchscript file which will later be read into COMMET.
When the model is read in, COMMET will search for one of several possible methods of the class which ultimately define the material behaviour – it is important that one of these methods is defined otherwise COMMET will not know where to look for the definition of the material behaviour.
These methods, with their signatures, are as displayed in the following code listing.
1 import torch
2 from torch import nn
3
4 class YourMaterialModel(nn.Module):
5 def __init__(self, ...):
6 ...
7
8 @torch.jit.export
9 def W_NN_from_C(self,
10 C: torch.Tensor,
11 structural_vectors: Optional[torch.Tensor] = None) -> torch.Tensor:
12 """Returns the strain energy density for a given right Cauchy-Green tensor and structural vectors.
13 ***Important***: COMMET takes advantage of symmetries in the right Cauchy-Green tensor when applying AD to calculate the stress and stiffness.
14 In order for this to work correctly, it is important that the first line of this method is:
15 C = 0.5*(C + C.transpose(1, 2))
16
17
18 Parameters
19 ----------
20 C : torch.Tensor
21 **shape (batch size, 3, 3)** Right Cauchy-Green tensors.
22 structural_vectors : torch.Tensor
23 **shape (batch size, number of structural vectors, 3)**
24
25 Returns
26 -------
27 torch.Tensor
28 **shape (batch size)** strain energy density
29
30 """
31 C = 0.5*(C + C.transpose(1, 2)) # !!! This is required for the autograd to pick-up the symmetries correctly !!!
32 ...
33
34 @torch.jit.export
35 def W_NN_from_F(self,
36 F: torch.Tensor,
37 structural_vectors: Optional[torch.Tensor] = None) -> torch.Tensor:
38 """Returns the strain energy density for a given deformation gradient and structural vectors.
39
40 Parameters
41 ----------
42 F : torch.Tensor
43 **shape (batch size, 3, 3)** Deformation gradients.
44 structural_vectors : torch.Tensor
45 **shape (batch size, number of structural vectors, 3)**
46
47 Returns
48 -------
49 torch.Tensor
50 **shape (batch size)** strain energy density
51
52 """
53 ...
54
55 @torch.jit.export
56 def psi_tau_cc_from_F(self,
57 F: torch.Tensor,
58 structural_tensors: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
59 """Returns the strain energy density, Kirchhoff stress, and spatial stiffness tensor in Voigt notation for a given deformation gradient and structural vectors.
60 The Voigt notation ordering is as follows (zero indexed):
61 vgt = [
62 (0, 0),
63 (1, 1),
64 (2, 2),
65 (1, 2),
66 (0, 2),
67 (0, 1)
68 ]
69
70 Parameters
71 ----------
72 F : torch.Tensor
73 **shape (batch size, 3, 3)** Deformation gradients.
74 structural_vectors : torch.Tensor
75 **shape (batch size, number of structural vectors, 3)**
76
77 Returns
78 -------
79 torch.Tensor
80 **shape (batch size)** strain energy density
81 torch.Tensor
82 **shape (batch size, 6)** Kirchhoff stress (2 F d\Psi/dC F^T) in Voigt notation
83 torch.Tensor
84 **shape (batch size, 6, 6)** Spatial stiffness tensor (4 d\Psi/(dC_{IJ} dC_{KL} ) F_{iI}F_{jJ}F_{kK}F_{lL}) in Voigt notation
85
86 """
87 ...
You may note that the methods W_NN_from_F and W_NN_from_C only require calculation of the strain energy density, whereas psi_tau_cc_from_F also requires calculation of the stress and stiffness.
This is because, if W_NN_from_F or W_NN_from_C is used, COMMET will use automatic differentiation (AD) to calculate the stress and stiffness.
In contrast, the method psi_tau_cc_from_F allows a user to implement the necessary derivatives themself, which may be more efficient than using AD.
The decorators @torch.jit.export are required to indicate that Pytorch must include these methods when generating the torchscript file (we’ll show how this is done below).
Finally, note that a user only needs to implement one of these methods, not all of them.
Below we show a full example of how to implement a simple Neohookean model using Pytorch and export an associated torchscript file – this is the contents of nh_torchscript_model.py.
1import torch
2from torch import nn
3from typing import Optional
4
5
6class NH(nn.Module):
7 def __init__(self,
8 mu: float,
9 lam: float):
10 super().__init__()
11 self.mu = torch.nn.Parameter(torch.Tensor([mu]))
12 self.lam = torch.nn.Parameter(torch.Tensor([lam]))
13
14 @torch.jit.export
15 def W_NN_from_C(self,
16 C: torch. Tensor,
17 structural_vectors: Optional[torch.Tensor] = None) -> torch.Tensor:
18 """Returns the strain energy density for a given right Cauchy-Green tensor and structural vectors.
19
20 Parameters
21 ----------
22 C : torch.Tensor
23 **shape (batch size, 3, 3)** Right Cauchy-Green tensors.
24 structural_vectors : torch.Tensor
25 **shape (batch size, number of structural vectors, 3)**
26
27 Returns
28 -------
29 torch.Tensor
30 **shape (batch size)** strain energy density
31
32 """
33
34 C = 0.5*(C + C.transpose(1, 2)) # !!! This is required for the autograd to pick-up the symmetries correctly !!!
35
36 I1 = torch.sum(C[:, [0, 1, 2], [0, 1, 2]], dim=1)
37 I3 = torch.det(C)
38 return 0.5 * self.mu * (I1 - 3 - torch.log(I3)) + 0.25 * self.lam * (I3 - 1 - torch.log(I3))
39
40 @torch.jit.export
41 def W_NN_from_F(self,
42 F: torch.Tensor,
43 structural_vectors: Optional[torch.Tensor] = None) -> torch.Tensor:
44 """Returns the strain energy density for a given deformation gradient and structural vectors.
45
46 Parameters
47 ----------
48 F : torch.Tensor
49 **shape (batch size, 3, 3)** Deformation gradients.
50 structural_vectors : torch.Tensor
51 **shape (batch size, number of structural vectors, 3)**
52
53 Returns
54 -------
55 torch.Tensor
56 **shape (batch size)** strain energy density
57
58 """
59
60 C = torch.bmm(F.transpose(1, 2), F)
61 return self.W_NN_from_C(C, structural_vectors)
62
63 def forward(self,
64 F: torch.Tensor,
65 structural_vectors: Optional[torch.Tensor] = None):
66 return self.W_NN_from_F(F, structural_vectors)
67
68
69def main():
70 mu: float = 77
71 lam: float = 115
72 nh: NH = NH(mu, lam)
73
74 dim = 3
75 F_mat = torch.zeros(100, dim, dim)
76
77 traced = torch.jit.trace(nh, (F_mat, ))
78 traced.save("nh.torchscript")
79
80
81if __name__ == "__main__":
82 main()
The above example also demonstrates how the model is converted to torchscript, i.e. the model and some example input is provided to torch.jit.trace, and the traced output is written to file using the save method.
It is important that the forward method is defined for the model since this is the default entry point for the torch.jit.trace function.
Next we’ll show how to use this in a COMMET simulation.
The input file
Since the problem definition is identical to the Hello World example (see Hello world: running a simulation without an NCM), the input file is largely the same. The only difference is the material definition.
1 {
2 // Defining the mesh
3 "mesh": {
4 // Indicating that we will use a built-in mesh (not reading one in from file)
5 "from": "built-in",
6 // The type of built-in mesh that we will use
7 "type": "quarter_plate_with_hole",
8 // The order of the finite elements to use
9 "order": 1,
10 // Radius of the hole in the plate
11 "radius": 0.2,
12 // Length of the plate
13 "length": 1,
14 // Thickness of the plate
15 "thickness": 0.1,
16 // The number of times to refine the mesh of the geometry in-plane
17 "planar_refinements": 1,
18 // The number of times to refine the mesh of the geometry globally (through the thickness and in-plane)
19 "global_refinements": 1
20 },
21 // Defining the materials
22 "materials": [
23 {
24 "id": 0,
25 // Technically, NCM stands for "neural constitutive network".
26 // This just indicates to commet_solve that the model must be loaded from torchscript.
27 "type": "ncm",
28 "vectorization": "batched", // batched | global
29 "batch_size": 512, // must be defined if "vectorization": "batched"
30 "evaluation_method": "using_C", // "optimized" | "using_F" | "using_C"
31 "path_to_torchscript": "nh.torchscript"
32 }
33 ],
34 // Defining the boundary conditions to be applied over time
35 "stages": [
36 {
37 "end_time": 1,
38 "time_increment_size": 0.1,
39 "dirichlet_boundary_conditions": [
40 {
41 "type": "standard", // Either "standard" or "pull_twist"
42 "boundary_id": 1, // The id of the boundary to which this condition applies
43 "components": [0], // The components (0->x, 1->y, 2->z) of the displacement to be prescribed.
44 "values": [0] // The values of the prescribed displacement. Must have the same number of entries as "components"
45 },
46 {
47 // Similar to above
48 "type": "standard",
49 "boundary_id": 2,
50 "components": [1],
51 "values": [0]
52 },
53 {
54 // Similar to above
55 "type": "standard",
56 "boundary_id": 3,
57 "components": [2],
58 "values": [0]
59 },
60 {
61 // Similar to above
62 "type": "standard",
63 "boundary_id": 4,
64 "components": [0, 1, 2],
65 "values": [1, 0, 0]
66 }
67 ],
68 "neumann_boundary_conditions": [],
69 "robin_boundary_conditions": []
70 }
71 ],
72 // Defining the fields to be output
73 "outputs": {
74 "scalar_outputs": ["I1"], // Scalar fields to output (here I1=\tr(C))
75 "vector_outputs": [ ], // Vector fields to output
76 "tensor_outputs": ["kirchhoff_stress", "C" ] // Tensor fields to output
77 }
78 }
The details of the of defining the material behaviour are as follows:
"type": "ncm"Settingtypeto code:”ncm” indicates to COMMET that the material model must be read in from torchscript"vectorization": "batched"If the model is read in from torchscript, COMMET uses vectorization to execute constitutive updates for multiple material points simultaneously, i.e. it uses vectorization. A user can choose whether that vectorization is applied globally, i.e. all material points are processed simultaneously, or in batches of a fixed size. In brief, using batched vectorization tends to be faster and more memory efficient than global vectorization. For more details see the COMMET paper."batch_size": 512If batch vectorization is chosen over global vectorization then the batch size (i.e. the number of material points to process simultaneously) must also be chosen. The most efficient batch size will likely vary from one problem to another. However, a batch size of 512 or 1024 is typically a good choice. Again, see the COMMET paper for more details."evaluation_method": "using_C"This is where we tell COMMET which method to use when computing the constitutive update:"using_F"->W_NN_from_F"using_C"->W_NN_from_C"optimized"->psi_tau_cc_from_F
"path_to_torchscript": "nh.torchscript"As you might expect, this is just the path to the torchscript file that must be read in.
Running the simulation
If you are using docker, you can run the example by going to the directory containing the input file and running
$ docker run --rm -v ./:/data -w /data commetcode/commet_solve mpirun -n <n_procs_to_use> commet_solve using_torchscript.jsonc
If, instead, you have build COMMET locally and it is in your path, you can run
$ mpirun -np <n_procs_to_use> commet_solve using_torchscript.jsonc