Implement correct 2d JIT solve!

To correct the 2D JIT solve with correct volume computation, 2 major
changes were made:
1. The elements class constructor now takes the total number of elements
as an argument along with number of elements in both directions. This is
so that `jnp.bincount` can work correctly without having to pass a
tracer value to its `length` parameter.
2. Update the `MPMExplicit` flattening/unflattening to include the
`velocity_update` parameter to make it work with JIT.

Author: chahak13

diffmpm/element.py

@@ -52,12 +52,14 @@ def id_to_node_vel(self, id: int):
         return self.nodes.velocity[node_ids]
 
     def tree_flatten(self):
-        children = (self.nodes,)
+        children = (self.nodes, self.volume)
         aux_data = (
             self.nelements,
+            self.total_elements,
             self.el_len,
             self.constraints,
             self.concentrated_nodal_forces,
+            self.initialized,
         )
         return children, aux_data
 
@@ -67,8 +69,11 @@ def tree_unflatten(cls, aux_data, children):
             aux_data[0],
             aux_data[1],
             aux_data[2],
+            aux_data[3],
             nodes=children[0],
-            concentrated_nodal_forces=aux_data[3],
+            concentrated_nodal_forces=aux_data[4],
+            initialized=aux_data[5],
+            volume=children[1],
         )
 
     @abc.abstractmethod
@@ -350,7 +355,7 @@ def __init__(
             IDs of nodes that are supposed to be fixed (boundary).
         """
         self.nelements = nelements
-        self.ids = jnp.arange(nelements)
+        self.total_elements = nelements
         self.el_len = el_len
         if nodes is None:
             self.nodes = Nodes(
@@ -495,7 +500,7 @@ def f(x):
 
     def compute_volume(self):
         vol = jnp.ediff1d(self.nodes.loc)
-        self.volume = jnp.ones((len(self.ids), 1, 1)) * vol
+        self.volume = jnp.ones((self.total_elements, 1, 1)) * vol
 
 
 @register_pytree_node_class
@@ -523,10 +528,13 @@ class Quadrilateral4Node(_Element):
     def __init__(
         self,
         nelements: Tuple[int, int],
+        total_elements: int,
         el_len: Tuple[float, float],
         constraints: List[Tuple[jnp.ndarray, Constraint]],
         nodes: Nodes = None,
         concentrated_nodal_forces=[],
+        initialized: bool = None,
+        volume: jnp.ndarray = None,
     ):
         """Initialize Quadrilateral4Node.
 
@@ -539,8 +547,7 @@ def __init__(
         """
         self.nelements = jnp.asarray(nelements)
         self.el_len = jnp.asarray(el_len)
-        total_elements = jnp.product(self.nelements)
-        self.ids = jnp.arange(total_elements)
+        self.total_elements = total_elements
 
         if nodes is None:
             total_nodes = jnp.product(self.nelements + 1)
@@ -561,6 +568,11 @@ def __init__(
 
         self.constraints = constraints
         self.concentrated_nodal_forces = concentrated_nodal_forces
+        if initialized is None:
+            self.volume = jnp.ones((self.total_elements, 1, 1))
+        else:
+            self.volume = volume
+        self.initialized = True
 
     def id_to_node_ids(self, id: int):
         """
@@ -773,12 +785,12 @@ def _step(pid, args):
         )
         self.nodes.f_int, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args)
 
-    def compute_volume(self):
+    def compute_volume(self, *args):
         a = c = self.el_len[1]
         b = d = self.el_len[0]
         p = q = jnp.sqrt(a**2 + b**2)
         vol = 0.25 * jnp.sqrt(4 * p * p * q * q - (a * a + c * c - b * b - d * d) ** 2)
-        self.volume = jnp.ones((len(self.ids), 1, 1)) * vol
+        self.volume = self.volume.at[:].set(vol)
 
 
 if __name__ == "__main__":

diffmpm/io.py

@@ -117,6 +117,7 @@ def _parse_mesh(self, config):
         if config["mesh"]["type"] == "generator":
             elements = element_cls(
                 config["mesh"]["nelements"],
+                jnp.product(jnp.array(config["mesh"]["nelements"])),
                 config["mesh"]["element_length"],
                 constraints,
                 concentrated_nodal_forces=self.parsed_config["external_loading"][

diffmpm/particle.py

@@ -1,6 +1,7 @@
 from typing import Tuple
 
 import jax.numpy as jnp
+from functools import partial
 from jax import jit, vmap, lax
 import jax.debug as db
 from jax.tree_util import register_pytree_node_class
@@ -144,9 +145,10 @@ def set_mass_volume(self, m: float | jnp.ndarray):
             )
         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))
+    def compute_volume(self, elements, total_elements):
+        particles_per_element = jnp.bincount(
+            self.element_ids, length=elements.total_elements
+        )
         vol = (
             elements.volume.squeeze((1, 2))[self.element_ids]
             / particles_per_element[self.element_ids]

diffmpm/solver.py

@@ -23,7 +23,7 @@ def __init__(self, filepath):
             raise ValueError("Wrong type of solver specified.")
 
     def solve(self):
-        res = self.solver.solve(
+        res = self.solver.solve_jit(
             self._config.parsed_config["meta"]["nsteps"],
             self._config.parsed_config["external_loading"]["gravity"],
         )
@@ -42,22 +42,27 @@ def __init__(self, mesh, dt, scheme="usf", velocity_update=False):
         self.mesh = mesh
         self.dt = dt
         self.scheme = scheme
+        self.velocity_update = velocity_update
         self.mesh.apply_on_elements("set_particle_element_ids")
-        self.mesh.apply_on_particles("compute_volume")
+        self.mesh.apply_on_elements("compute_volume")
+        self.mesh.apply_on_particles(
+            "compute_volume", args=(self.mesh.elements.total_elements,)
+        )
 
     def tree_flatten(self):
         children = (self.mesh,)
-        aux_data = (self.dt, self.scheme)
+        aux_data = (self.dt, self.scheme, self.velocity_update)
         return children, aux_data
 
     @classmethod
     def tree_unflatten(cls, aux_data, children):
-        return cls(*children, aux_data[0], scheme=aux_data[1])
+        return cls(
+            *children, aux_data[0], scheme=aux_data[1], velocity_update=aux_data[2]
+        )
 
     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)
@@ -75,21 +80,17 @@ def solve(self, nsteps: int, gravity: float | jnp.ndarray):
     def solve_jit(self, nsteps: int, gravity: float | jnp.ndarray):
         nparticles = sum(pset.loc.shape[0] for pset in self.mesh.particles)
         result = {
-            "position": jnp.zeros((nsteps, nparticles)),
-            "velocity": jnp.zeros((nsteps, nparticles)),
-            "strain_energy": jnp.zeros((nsteps, nparticles)),
-            "kinetic_energy": jnp.zeros((nsteps, nparticles)),
-            "total_energy": jnp.zeros((nsteps, nparticles)),
-            "stress": jnp.zeros((nsteps, nparticles)),
-            "strain": jnp.zeros((nsteps, nparticles)),
+            "position": jnp.zeros((nsteps, nparticles, 2)),
+            "velocity": jnp.zeros((nsteps, nparticles, 2)),
+            "stress": jnp.zeros((nsteps, nparticles, 6)),
+            "strain": jnp.zeros((nsteps, nparticles, 6)),
         }
 
         def _step(i, data):
             self, result = data
             self.mpm_scheme.compute_nodal_kinematics()
             self.mpm_scheme.precompute_stress_strain()
-            self.mpm_scheme.compute_forces(gravity)
-            # self.mpm_scheme.update_nodal_momentum()
+            self.mpm_scheme.compute_forces(gravity, i)
             self.mpm_scheme.compute_particle_kinematics()
             self.mpm_scheme.postcompute_stress_strain()
 
@@ -99,45 +100,23 @@ def _step(i, data):
                 idu += len(self.mesh.particles[j])
                 result["position"] = (
                     result["position"]
-                    .at[i, idl:idu]
+                    .at[i, idl:idu, :]
                     .set(self.mesh.particles[j].loc.squeeze())
                 )
                 result["velocity"] = (
                     result["velocity"]
-                    .at[i, idl:idu]
+                    .at[i, idl:idu, :]
                     .set(self.mesh.particles[j].velocity.squeeze())
                 )
                 result["stress"] = (
                     result["stress"]
-                    .at[i, idl:idu]
-                    .set(self.mesh.particles[j].stress[:, 0, :].squeeze())
+                    .at[i, idl:idu, :]
+                    .set(self.mesh.particles[j].stress[:, :, 0].squeeze())
                 )
                 result["strain"] = (
                     result["strain"]
-                    .at[i, idl:idu]
-                    .set(self.mesh.particles[j].strain[:, 0, :].squeeze())
-                )
-                strain_energy = (
-                    0.5
-                    * self.mesh.particles[j].stress[:, 0, :].squeeze()
-                    * self.mesh.particles[j].strain[:, 0, :].squeeze()
-                    * self.mesh.particles[j].volume.squeeze()
-                )
-                kinetic_energy = (
-                    0.5
-                    * self.mesh.particles[j].velocity.squeeze() ** 2
-                    * self.mesh.particles[j].mass.squeeze()
-                )
-                result["strain_energy"] = (
-                    result["strain_energy"].at[i, idl:idu].set(strain_energy)
-                )
-                result["kinetic_energy"] = (
-                    result["kinetic_energy"].at[i, idl:idu].set(kinetic_energy)
-                )
-                result["total_energy"] = (
-                    result["total_energy"]
-                    .at[i, idl:idu]
-                    .set(strain_energy + kinetic_energy)
+                    .at[i, idl:idu, :]
+                    .set(self.mesh.particles[j].strain[:, :, 0].squeeze())
                 )
             return (self, result)
 

examples/mpm-nodal-forces.toml

@@ -16,7 +16,7 @@ type = "MPMExplicit"
 dimension = 2
 scheme = "usf"
 dt = 0.001
-nsteps = 20
+nsteps = 301
 velocity_update = true
 
 [output]
@@ -59,7 +59,7 @@ gravity = [0, 0]
 node_ids = [3, 7]
 math_function_id = 0
 dir = 0
-force = 0.01
+force = 0.05
 
 [[external_loading.particle_surface_traction]]
 pset = [1]

examples/simple_2d_file.py

@@ -6,7 +6,7 @@
 mpm = MPM(sys.argv[1])
 result = mpm.solve()
 # breakpoint()
-print(result["stress"][-1])
+print(result["stress"][-1][:, :2])
 exit()