Add broken nodal forces

Author: chahak13

diffmpm/element.py

@@ -196,7 +196,9 @@ def compute_body_force(self, particles, gravity: float):
 
         def _step(pid, args):
             f_ext, pmass, mapped_pos, el_nodes, gravity = args
-            f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ (pmass * gravity))
+            f_ext = f_ext.at[el_nodes[pid]].add(
+                mapped_pos[pid] @ (pmass[pid] * gravity)
+            )
             return f_ext, pmass, mapped_pos, el_nodes, gravity
 
         mapped_positions = self.shapefn(particles.reference_loc)
@@ -210,17 +212,30 @@ def _step(pid, args):
         )
         self.nodes.f_ext, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args)
 
-    def apply_concentrated_nodal_force(self, particles, curr_time):
-        def _step(fid, args):
-            f_ext, cnf, curr_time = args
-            factor = cnf[fid].value(curr_time)
-            f_ext = f_ext.at[cnf[fid].node_ids].add(factor * cnf[fid].force)
-            return f_ext, cnf, curr_time
+    def apply_concentrated_nodal_forces(self, particles, curr_time):
+        # def _step(fid, args):
+        #     f_ext, cnf, curr_time = args
+        #     breakpoint()
+        #     factor = cnf[fid].function.value(curr_time)
+        #     f_ext = f_ext.at[cnf[fid].node_ids].add(factor * cnf[fid].force)
+        #     return f_ext, cnf, curr_time
+
+        # args = (self.nodes.f_ext, self.concentrated_nodal_forces, curr_time)
+        # self.nodes.f_ext, _, _ = lax.fori_loop(
+        #     0, len(self.concentrated_nodal_forces), _step, args
+        # )
+        # breakpoint()
+        import jax.debug as db
 
-        args = (self.nodes.f_ext, self.concentrated_nodal_forces, curr_time)
-        self.nodes.f_ext, _, _ = lax.fori_loop(
-            0, len(self.concentrated_nodal_forces), _step, args
-        )
+        for cnf in self.concentrated_nodal_forces:
+            factor = cnf.function.value(curr_time)
+            self.nodes.f_ext = self.nodes.f_ext.at[cnf.node_ids, 0, cnf.dir].add(
+                factor * cnf.force
+            )
+            db.print(
+                f"Factor: {factor}, curr_time: {curr_time}, "
+                f"f_ext[3]: {self.nodes.f_ext[3].squeeze()}"
+            )
 
     def compute_internal_force(self, particles):
         r"""
@@ -276,7 +291,11 @@ def _step(pid, args):
 
     def update_nodal_acceleration_velocity(self, particles, dt: float, *args):
         """Update the nodal momentum based on total force on nodes."""
+        import jax.debug as db
+
         total_force = self.nodes.get_total_force()
+        db.print(f"Before: nodes.velocity[3]: {self.nodes.velocity[3].squeeze()}")
+        breakpoint()
         self.nodes.acceleration = self.nodes.acceleration.at[:].set(
             jnp.nan_to_num(jnp.divide(total_force, self.nodes.mass))
         )
@@ -287,6 +306,7 @@ def update_nodal_acceleration_velocity(self, particles, dt: float, *args):
         self.nodes.momentum = self.nodes.momentum.at[:].set(
             self.nodes.mass * self.nodes.velocity
         )
+        db.print(f"After: nodes.velocity[3]: {self.nodes.velocity[3].squeeze()}")
 
     def apply_boundary_constraints(self, *args):
         """Apply boundary conditions for nodal velocity."""
@@ -720,7 +740,7 @@ def _step(pid, args):
             ) = args
             # TODO: correct matrix multiplication for n-d
             # update = -(pvol[pid]) * pstress[pid] @ mapped_grads[pid]
-            force = jnp.zeros((mapped_grads.shape[0], 1, 2))
+            force = jnp.zeros((mapped_grads.shape[1], 1, 2))
             force = force.at[:, 0, 0].set(
                 mapped_grads[pid][:, 0] * pstress[pid][0]
                 + mapped_grads[pid][:, 1] * pstress[pid][3]

diffmpm/forces.py

@@ -1,6 +1,17 @@
 from collections import namedtuple
+from jax.tree_util import register_pytree_node
 
 NodalForce = namedtuple("NodalForce", ("node_ids", "function", "dir", "force"))
 ParticleTraction = namedtuple(
     "ParticleTraction", ("pset", "function", "dir", "traction")
 )
+register_pytree_node(
+    NodalForce,
+    lambda xs: (tuple(xs), None),  # tell JAX how to unpack to an iterable
+    lambda _, xs: NodalForce(*xs),  # tell JAX how to pack back into a NodalForce
+)
+register_pytree_node(
+    ParticleTraction,
+    lambda xs: (tuple(xs), None),  # tell JAX how to unpack to an iterable
+    lambda _, xs: ParticleTraction(*xs),  # tell JAX how to pack back
+)

diffmpm/functions.py

@@ -1,5 +1,6 @@
 import abc
 import jax.numpy as jnp
+from jax.tree_util import register_pytree_node_class
 
 
 class Function(abc.ABC):
@@ -11,14 +12,24 @@ def value(self):
         ...
 
 
+@register_pytree_node_class
 class Unit(Function):
     def __init__(self, id):
         super().__init__(id)
 
     def value(self, x):
         return 1.0
 
+    def tree_flatten(self):
+        return ((), (self.id))
 
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        del children
+        return cls(*aux_data)
+
+
+@register_pytree_node_class
 class Linear(Function):
     def __init__(self, id, xvalues, fxvalues):
         self.xvalues = xvalues
@@ -27,3 +38,11 @@ def __init__(self, id, xvalues, fxvalues):
 
     def value(self, x):
         return jnp.interp(x, self.xvalues, self.fxvalues)
+
+    def tree_flatten(self):
+        return ((), (self.id, self.xvalues, self.fxvalues))
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        del children
+        return cls(*aux_data)

diffmpm/io.py

@@ -86,7 +86,7 @@ def _parse_external_loading(self, config):
                 fn = Unit(-1)
             cnf = NodalForce(
                 node_ids=jnp.array(cnfconfig["node_ids"]),
-                math_function=fn,
+                function=fn,
                 dir=cnfconfig["dir"],
                 force=cnfconfig["force"],
             )
@@ -95,12 +95,12 @@ def _parse_external_loading(self, config):
         pst_list = []
         for pstconfig in config["external_loading"]["particle_surface_traction"]:
             pst = ParticleTraction(
-                pset=jnp.array(cnfconfig["pset"]),
-                math_function=self.parsed_config["math_functions"][
+                pset=jnp.array(pstconfig["pset"]),
+                function=self.parsed_config["math_functions"][
                     pstconfig["math_function_id"]
                 ],
-                dir=cnfconfig["dir"],
-                traction=cnfconfig["traction"],
+                dir=pstconfig["dir"],
+                traction=pstconfig["traction"],
             )
             pst_list.append(pst)
         external_loading["concentrated_nodal_forces"] = cnf_list

diffmpm/particle.py

@@ -2,6 +2,7 @@
 
 import jax.numpy as jnp
 from jax import jit, vmap, lax
+import jax.debug as db
 from jax.tree_util import register_pytree_node_class
 
 from diffmpm.element import _Element
@@ -45,8 +46,7 @@ def __init__(
         self.element_ids = element_ids
         if len(loc.shape) != 3:
             raise ValueError(
-                f"`loc` should be of size (nparticles, 1, ndim); "
-                f"found {loc.shape}"
+                f"`loc` should be of size (nparticles, 1, ndim); " f"found {loc.shape}"
             )
         self.loc = loc
 
@@ -140,16 +140,13 @@ def set_mass_volume(self, m: float | jnp.ndarray):
             self.mass = m
         else:
             raise ValueError(
-                f"Incompatible shapes. Expected {self.mass.shape}, "
-                f"found {m.shape}."
+                f"Incompatible shapes. Expected {self.mass.shape}, " f"found {m.shape}."
             )
         self.volume = jnp.divide(self.mass, self.material.properties["density"])
 
     def compute_volume(self, elements: _Element):
         elements.compute_volume()
-        particles_per_element = jnp.bincount(
-            self.element_ids, length=len(elements.ids)
-        )
+        particles_per_element = jnp.bincount(self.element_ids, length=len(elements.ids))
         vol = (
             elements.volume.squeeze((1, 2))[self.element_ids]
             / particles_per_element[self.element_ids]
@@ -238,29 +235,26 @@ def compute_strain(self, elements: _Element, dt: float):
         dt : float
             Timestep.
         """
-        mapped_coords = vmap(elements.id_to_node_loc)(self.element_ids).squeeze(
-            2
-        )
+        mapped_coords = vmap(elements.id_to_node_loc)(self.element_ids).squeeze(2)
         dn_dx_ = vmap(elements.shapefn_grad)(
             self.reference_loc[:, jnp.newaxis, ...], mapped_coords
         )
         self.strain_rate = self._compute_strain_rate(dn_dx_, elements)
         self.dstrain = self.dstrain.at[:].set(self.strain_rate * dt)
+
+        # db.print(f"compute_strain() - dstrain: {self.dstrain.squeeze()[3, :2]}")
         self.strain = self.strain.at[:].add(self.dstrain)
+        # db.print(f"compute_strain() - strain: {self.strain.squeeze()[3, :2]}")
         centroids = jnp.zeros_like(self.loc)
         dn_dx_centroid_ = vmap(elements.shapefn_grad)(
             centroids[:, jnp.newaxis, ...], mapped_coords
         )
-        strain_rate_centroid = self._compute_strain_rate(
-            dn_dx_centroid_, elements
-        )
+        strain_rate_centroid = self._compute_strain_rate(dn_dx_centroid_, elements)
         ndim = self.loc.shape[-1]
-        self.dvolumetric_strain = dt * strain_rate_centroid[:, :ndim].sum(
-            axis=1
+        self.dvolumetric_strain = dt * strain_rate_centroid[:, :ndim].sum(axis=1)
+        self.volumetric_strain_centroid = self.volumetric_strain_centroid.at[:].add(
+            self.dvolumetric_strain
         )
-        self.volumetric_strain_centroid = self.volumetric_strain_centroid.at[
-            :
-        ].add(self.dvolumetric_strain)
 
     def _compute_strain_rate(self, dn_dx: jnp.ndarray, elements: _Element):
         """
@@ -279,6 +273,7 @@ def _compute_strain_rate(self, dn_dx: jnp.ndarray, elements: _Element):
             self.element_ids
         )  # (nparticles, 2, 1)
 
+        # db.print(f"mapped_val[3]: {mapped_vel[3, :, 0]}")
         # TODO: This will need to change to be more general for ndim.
         # breakpoint()
         # L = jnp.einsum("ijk, ikj -> ijk", dn_dx, mapped_vel.squeeze(-1)).sum(
@@ -295,9 +290,7 @@ def _step(pid, args):
             matmul = dndx[pid].T @ nvel[pid]
             strain_rate = strain_rate.at[pid, 0].add(matmul[0, 0])
             strain_rate = strain_rate.at[pid, 1].add(matmul[1, 1])
-            strain_rate = strain_rate.at[pid, 3].add(
-                matmul[0, 1] + matmul[1, 0]
-            )
+            strain_rate = strain_rate.at[pid, 3].add(matmul[0, 1] + matmul[1, 0])
             return dndx, nvel, strain_rate
 
         args = (dn_dx, temp, strain_rate)
@@ -314,18 +307,12 @@ def compute_stress(self, *args):
         particles. The stress calculated by the material is then
         added to the particles current stress values.
         """
-        self.stress = self.stress.at[:].add(
-            self.material.compute_stress(self.dstrain)
-        )
+        self.stress = self.stress.at[:].add(self.material.compute_stress(self.dstrain))
 
     def update_volume(self, *args):
         """Update volume based on central strain rate."""
-        self.volume = self.volume.at[:, 0, :].multiply(
-            1 + self.dvolumetric_strain
-        )
-        self.density = self.density.at[:, 0, :].divide(
-            1 + self.dvolumetric_strain
-        )
+        self.volume = self.volume.at[:, 0, :].multiply(1 + self.dvolumetric_strain)
+        self.density = self.density.at[:, 0, :].divide(1 + self.dvolumetric_strain)
 
 
 if __name__ == "__main__":

diffmpm/solver.py

@@ -56,15 +56,17 @@ def tree_unflatten(cls, aux_data, children):
     def solve(self, nsteps: int, gravity: float | jnp.ndarray):
         result = defaultdict(list)
         for step in tqdm(range(nsteps)):
+            # breakpoint()
             self.mpm_scheme.compute_nodal_kinematics()
             self.mpm_scheme.precompute_stress_strain()
             self.mpm_scheme.compute_forces(gravity, step)
             self.mpm_scheme.compute_particle_kinematics()
             self.mpm_scheme.postcompute_stress_strain()
             for pset in self.mesh.particles:
-                # result["position"].append(pset.loc)
-                # result["velocity"].append(pset.velocity)
-                result["stress"].append(pset.stress)
+                result["position"].append(pset.loc)
+                result["velocity"].append(pset.velocity)
+                result["stress"].append(pset.stress[:, :2, 0])
+                result["strain"].append(pset.strain[:, :2, 0])
 
         result = {k: jnp.asarray(v) for k, v in result.items()}
         return result

examples/mpm-nodal-forces.toml

@@ -0,0 +1,73 @@
+# 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-nodal-traction"
+type = "MPMExplicit"
+dimension = 2
+scheme = "usf"
+dt = 0.001
+nsteps = 20
+velocity_update = true
+
+[output]
+type = "hdf5"
+file = "results/example_2d_out.hdf5"
+step_frequency = 5
+
+[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 = "examples/particles-2d-nodal-force.json"
+material_id = 0
+init_velocity = 0.0
+
+[external_loading]
+gravity = [0, 0]
+
+[[external_loading.concentrated_nodal_forces]]
+node_ids = [3, 7]
+math_function_id = 0
+dir = 0
+force = 0.01
+
+[[external_loading.particle_surface_traction]]
+pset = [1]
+dir = 1
+math_function_id = 0
+traction = 10.5
+
+[[math_functions]]
+type = "Linear"
+xvalues = [0.0, 0.5, 1.0]
+fxvalues = [0.0, 1.0, 1.0]