Merge pull request #14 from chahak13/traction

Add particle surface traction

Author: kks32

benchmarks/2d/uniaxial_particle_traction/mpm-particle-traction.toml

@@ -0,0 +1,68 @@
+# The `meta` group contains top level attributes that govern the
+# behaviour of the MPM Solver.
+#
+# Attributes:
+# title:    The title of the experiment. This is just for the user's
+#           reference.
+# type:     The type of simulation to be used. Allowed values are
+#           {"MPMExplicit"}
+# scheme:   The MPM Scheme used for simulation. Allowed values are
+#           {"usl", "usf"}
+# dt:       Timestep used in the simulation.
+# nsteps:   Number of steps to run the simulation for.
+[meta]
+title = "uniaxial-particle-traction"
+type = "MPMExplicit"
+dimension = 2
+scheme = "usf"
+dt = 0.001
+nsteps = 1000
+velocity_update = true
+
+[output]
+format = "npz"
+folder = "results/"
+step_frequency = 10
+
+[mesh]
+# type = "file"
+# file = "mesh-1d.txt"
+# boundary_nodes = "boundary-1d.txt"
+# particle_element_ids = "particles-elements.txt"
+type = "generator"
+nelements = [3, 1]
+element_length = [0.1, 0.1]
+particle_element_ids = [0]
+element = "Quadrilateral4Node"
+
+[[mesh.constraints]]
+node_ids = [0, 4]
+dir = 0
+velocity = 0.0
+
+[[materials]]
+id = 0
+density = 1000
+poisson_ratio = 0
+youngs_modulus = 1000000
+type = "LinearElastic"
+
+[[particles]]
+file = "particles-2d-particle-traction.json"
+material_id = 0
+init_velocity = 0.0
+
+[external_loading]
+gravity = [0, 0]
+
+[[external_loading.particle_surface_traction]]
+pset = [0]
+pids = [[5, 11]]
+math_function_id = 0
+dir = 0
+traction = 1
+
+[[math_functions]]
+type = "Linear"
+xvalues = [0.0, 0.5, 1.0]
+fxvalues = [0.0, 1.0, 1.0]

benchmarks/2d/uniaxial_particle_traction/particles-2d-particle-traction.json

@@ -0,0 +1,74 @@
+[
+  [
+    [
+      0.025,
+      0.025
+    ]
+  ],
+  [
+    [
+      0.075,
+      0.025
+    ]
+  ],
+  [
+    [
+      0.125,
+      0.025
+    ]
+  ],
+  [
+    [
+      0.175,
+      0.025
+    ]
+  ],
+  [
+    [
+      0.225,
+      0.025
+    ]
+  ],
+  [
+    [
+      0.275,
+      0.025
+    ]
+  ],
+  [
+    [
+      0.025,
+      0.075
+    ]
+  ],
+  [
+    [
+      0.075,
+      0.075
+    ]
+  ],
+  [
+    [
+      0.125,
+      0.075
+    ]
+  ],
+  [
+    [
+      0.175,
+      0.075
+    ]
+  ],
+  [
+    [
+      0.225,
+      0.075
+    ]
+  ],
+  [
+    [
+      0.275,
+      0.075
+    ]
+  ]
+]

benchmarks/2d/uniaxial_particle_traction/test_benchmark.py

@@ -0,0 +1,29 @@
+import os
+from pathlib import Path
+import jax.numpy as jnp
+from diffmpm.solver import MPM
+
+curr_filepath = Path(__file__).absolute()
+curr_dir = curr_filepath.parent
+os.chdir(curr_dir)
+mpm = MPM("mpm-particle-traction.toml")
+mpm.solve()
+result = jnp.load("results/uniaxial-particle-traction/particles_0300.npz")
+## Step 300
+assert jnp.round(result["stress"][0, :, 0].min() - 0.4450086768966724, 5) == 0.0
+assert jnp.round(result["stress"][0, :, 0].max() - 0.5966527842046769, 5) == 0.0
+
+## Step 510
+result = jnp.load("results/uniaxial-particle-traction/particles_0510.npz")
+assert jnp.round(result["stress"][0, :, 0].min() - 0.7528092313640623, 5) == 0.0
+assert jnp.round(result["stress"][0, :, 0].max() - 1.0109599915279937, 5) == 0.0
+
+## Step 750
+result = jnp.load("results/uniaxial-particle-traction/particles_0750.npz")
+assert jnp.round(result["stress"][0, :, 0].min() - 0.7500090055681591, 5) == 0.0
+assert jnp.round(result["stress"][0, :, 0].max() - 1.0000224746314728, 5) == 0.0
+
+## Step 990
+result = jnp.load("results/uniaxial-particle-traction/particles_0990.npz")
+assert jnp.round(result["stress"][0, :, 0].min() - 0.750002924022295, 5) == 0.0
+assert jnp.round(result["stress"][0, :, 0].max() - 0.9999997782938734, 5) == 0.0

diffmpm/element.py

@@ -239,6 +239,17 @@ def apply_concentrated_nodal_forces(self, particles, curr_time):
                 factor * cnf.force
             )
 
+    def apply_particle_traction_forces(self, particles):
+        def _step(pid, args):
+            f_ext, ptraction, mapped_pos, el_nodes = args
+            f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ ptraction[pid])
+            return f_ext, ptraction, mapped_pos, el_nodes
+
+        mapped_positions = self.shapefn(particles.reference_loc)
+        mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1)
+        args = (self.nodes.f_ext, particles.traction, mapped_positions, mapped_nodes)
+        self.nodes.f_ext, _, _, _ = lax.fori_loop(0, len(particles), _step, args)
+
     def update_nodal_acceleration_velocity(self, particles, dt: float, *args):
         """Update the nodal momentum based on total force on nodes."""
         total_force = self.nodes.get_total_force()

diffmpm/forces.py

@@ -3,7 +3,7 @@
 
 NodalForce = namedtuple("NodalForce", ("node_ids", "function", "dir", "force"))
 ParticleTraction = namedtuple(
-    "ParticleTraction", ("pset", "function", "dir", "traction")
+    "ParticleTraction", ("pset", "pids", "function", "dir", "traction")
 )
 register_pytree_node(
     NodalForce,

diffmpm/io.py

@@ -80,7 +80,7 @@ def _parse_external_loading(self, config):
         external_loading = {}
         external_loading["gravity"] = jnp.array(config["external_loading"]["gravity"])
         external_loading["concentrated_nodal_forces"] = []
-        external_loading["particle_surface_traction"] = []
+        particle_surface_traction = []
         if "concentrated_nodal_forces" in config["external_loading"]:
             cnf_list = []
             for cnfconfig in config["external_loading"]["concentrated_nodal_forces"]:
@@ -102,17 +102,23 @@ def _parse_external_loading(self, config):
         if "particle_surface_traction" in config["external_loading"]:
             pst_list = []
             for pstconfig in config["external_loading"]["particle_surface_traction"]:
-                pst = ParticleTraction(
-                    pset=jnp.array(pstconfig["pset"]),
-                    function=self.parsed_config["math_functions"][
+                if "math_function_id" in pstconfig:
+                    fn = self.parsed_config["math_functions"][
                         pstconfig["math_function_id"]
-                    ],
+                    ]
+                else:
+                    fn = Unit(-1)
+                pst = ParticleTraction(
+                    pset=pstconfig["pset"],
+                    pids=jnp.array(pstconfig["pids"]),
+                    function=fn,
                     dir=pstconfig["dir"],
                     traction=pstconfig["traction"],
                 )
                 pst_list.append(pst)
-            external_loading["particle_surface_traction"] = pst_list
+            particle_surface_traction.extend(pst_list)
         self.parsed_config["external_loading"] = external_loading
+        self.parsed_config["particle_surface_traction"] = particle_surface_traction
 
     def _parse_mesh(self, config):
         element_cls = getattr(mpel, config["mesh"]["element"])

diffmpm/mesh.py

@@ -21,6 +21,12 @@ def __init__(self, config: dict):
         """Initialize mesh using configuration."""
         self.particles: Iterable[Particles, ...] = config["particles"]
         self.elements: _Element = config["elements"]
+        self.particle_tractions = config["particle_surface_traction"]
+
+    @property
+    @abc.abstractmethod
+    def ndim(self):
+        ...
 
     # TODO: Convert to using jax directives for loop
     def apply_on_elements(self, function, args=()):
@@ -34,15 +40,33 @@ def apply_on_particles(self, function, args=()):
             f = getattr(particle_set, function)
             f(self.elements, *args)
 
+    def apply_traction_on_particles(self, curr_time):
+        self.apply_on_particles("zero_traction")
+        for ptraction in self.particle_tractions:
+            factor = ptraction.function.value(curr_time)
+            traction_val = factor * ptraction.traction
+            for i, pset_id in enumerate(ptraction.pset):
+                self.particles[pset_id].assign_traction(
+                    ptraction.pids[i], ptraction.dir, traction_val
+                )
+
+        # breakpoint()
+        self.apply_on_elements("apply_particle_traction_forces")
+
     def tree_flatten(self):
         children = (self.particles, self.elements)
-        aux_data = tuple()
+        aux_data = self.particle_tractions
         return (children, aux_data)
 
     @classmethod
     def tree_unflatten(cls, aux_data, children):
-        del aux_data
-        return cls({"particles": children[0], "elements": children[1]})
+        return cls(
+            {
+                "particles": children[0],
+                "elements": children[1],
+                "particle_surface_traction": aux_data,
+            }
+        )
 
 
 @register_pytree_node_class
@@ -61,6 +85,10 @@ def __init__(self, config: dict):
         """
         super().__init__(config)
 
+    @property
+    def ndim(self):
+        return 1
+
 
 @register_pytree_node_class
 class Mesh2D(_MeshBase):
@@ -78,6 +106,10 @@ def __init__(self, config: dict):
         """
         super().__init__(config)
 
+    @property
+    def ndim(self):
+        return 2
+
 
 if __name__ == "__main__":
     from diffmpm.element import Linear1D

diffmpm/particle.py

@@ -55,6 +55,7 @@ def __init__(
                 jnp.ones_like(self.mass) * self.material.properties["density"]
             )
             self.volume = jnp.ones_like(self.mass)
+            self.size = jnp.zeros_like(self.loc)
             self.velocity = jnp.zeros_like(self.loc)
             self.acceleration = jnp.zeros_like(self.loc)
             self.momentum = jnp.zeros_like(self.loc)
@@ -63,13 +64,15 @@ def __init__(
             self.strain_rate = jnp.zeros((self.loc.shape[0], 6, 1))
             self.dstrain = jnp.zeros((self.loc.shape[0], 6, 1))
             self.f_ext = jnp.zeros_like(self.loc)
+            self.traction = jnp.zeros_like(self.loc)
             self.reference_loc = jnp.zeros_like(self.loc)
             self.volumetric_strain_centroid = jnp.zeros((self.loc.shape[0], 1))
         else:
             (
                 self.mass,
                 self.density,
                 self.volume,
+                self.size,
                 self.velocity,
                 self.acceleration,
                 self.momentum,
@@ -78,6 +81,7 @@ def __init__(
                 self.strain_rate,
                 self.dstrain,
                 self.f_ext,
+                self.traction,
                 self.reference_loc,
                 self.volumetric_strain_centroid,
             ) = data
@@ -91,6 +95,7 @@ def tree_flatten(self):
             self.mass,
             self.density,
             self.volume,
+            self.size,
             self.velocity,
             self.acceleration,
             self.momentum,
@@ -99,6 +104,7 @@ def tree_flatten(self):
             self.strain_rate,
             self.dstrain,
             self.f_ext,
+            self.traction,
             self.reference_loc,
             self.volumetric_strain_centroid,
         )
@@ -152,6 +158,7 @@ def compute_volume(self, elements, total_elements):
             / particles_per_element[self.element_ids]
         )
         self.volume = self.volume.at[:, 0, 0].set(vol)
+        self.size = self.size.at[:].set(self.volume ** (1 / self.size.shape[-1]))
         self.mass = self.mass.at[:, 0, 0].set(vol * self.density.squeeze())
 
     def update_natural_coords(self, elements: _Element):
@@ -304,6 +311,14 @@ def update_volume(self, *args):
         self.volume = self.volume.at[:, 0, :].multiply(1 + self.dvolumetric_strain)
         self.density = self.density.at[:, 0, :].divide(1 + self.dvolumetric_strain)
 
+    def assign_traction(self, pids, dir, traction_):
+        self.traction = self.traction.at[pids, 0, dir].add(
+            traction_ * self.volume[pids, 0, 0] / self.size[pids, 0, dir]
+        )
+
+    def zero_traction(self, *args):
+        self.traction = self.traction.at[:].set(0)
+
 
 if __name__ == "__main__":
     from diffmpm.material import SimpleMaterial

diffmpm/scheme.py

@@ -25,6 +25,7 @@ def compute_stress_strain(self):
     def compute_forces(self, gravity, step):
         self.mesh.apply_on_elements("compute_external_force")
         self.mesh.apply_on_elements("compute_body_force", args=(gravity,))
+        self.mesh.apply_traction_on_particles(step * self.dt)
         self.mesh.apply_on_elements(
             "apply_concentrated_nodal_forces", args=(step * self.dt,)
         )