Merge pull request #6 from chahak13/2d

Implement correct 2d JIT solve!

Author: kks32

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
@@ -142,7 +147,7 @@ def _step(pid, args):
         )
         _, self.nodes.momentum, _, _ = lax.fori_loop(0, len(particles), _step, args)
         self.nodes.momentum = jnp.where(
-            self.nodes.momentum < 1e-12,
+            jnp.abs(self.nodes.momentum) < 1e-12,
             jnp.zeros_like(self.nodes.momentum),
             self.nodes.momentum,
         )
@@ -154,7 +159,7 @@ def compute_velocity(self, particles):
             self.nodes.momentum / self.nodes.mass,
         )
         self.nodes.velocity = jnp.where(
-            self.nodes.velocity < 1e-12,
+            jnp.abs(self.nodes.velocity) < 1e-12,
             jnp.zeros_like(self.nodes.velocity),
             self.nodes.velocity,
         )
@@ -296,12 +301,12 @@ def update_nodal_acceleration_velocity(self, particles, dt: float, *args):
             self.nodes.mass * self.nodes.velocity
         )
         self.nodes.velocity = jnp.where(
-            self.nodes.velocity < 1e-12,
+            jnp.abs(self.nodes.velocity) < 1e-12,
             jnp.zeros_like(self.nodes.velocity),
             self.nodes.velocity,
         )
         self.nodes.acceleration = jnp.where(
-            self.nodes.acceleration < 1e-12,
+            jnp.abs(self.nodes.acceleration) < 1e-12,
             jnp.zeros_like(self.nodes.acceleration),
             self.nodes.acceleration,
         )
@@ -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

@@ -27,8 +27,10 @@ def parse(self):
         self._parse_output(self._fileconfig)
         self._parse_materials(self._fileconfig)
         self._parse_particles(self._fileconfig)
-        self._parse_math_functions(self._fileconfig)
-        self._parse_external_loading(self._fileconfig)
+        if "math_functions" in self._fileconfig:
+            self._parse_math_functions(self._fileconfig)
+        if "external_loading" in self._fileconfig:
+            self._parse_external_loading(self._fileconfig)
         mesh = self._parse_mesh(self._fileconfig)
         return mesh
 
@@ -78,33 +80,39 @@ def _parse_math_functions(self, config):
     def _parse_external_loading(self, config):
         external_loading = {}
         external_loading["gravity"] = jnp.array(config["external_loading"]["gravity"])
-        cnf_list = []
-        for cnfconfig in config["external_loading"]["concentrated_nodal_forces"]:
-            if "math_function_id" in cnfconfig:
-                fn = self.parsed_config["math_functions"][cnfconfig["math_function_id"]]
-            else:
-                fn = Unit(-1)
-            cnf = NodalForce(
-                node_ids=jnp.array(cnfconfig["node_ids"]),
-                function=fn,
-                dir=cnfconfig["dir"],
-                force=cnfconfig["force"],
-            )
-            cnf_list.append(cnf)
+        external_loading["concentrated_nodal_forces"] = []
+        external_loading["particle_surface_traction"] = []
+        if "concentrated_nodal_forces" in config["external_loading"]:
+            cnf_list = []
+            for cnfconfig in config["external_loading"]["concentrated_nodal_forces"]:
+                if "math_function_id" in cnfconfig:
+                    fn = self.parsed_config["math_functions"][
+                        cnfconfig["math_function_id"]
+                    ]
+                else:
+                    fn = Unit(-1)
+                cnf = NodalForce(
+                    node_ids=jnp.array(cnfconfig["node_ids"]),
+                    function=fn,
+                    dir=cnfconfig["dir"],
+                    force=cnfconfig["force"],
+                )
+                cnf_list.append(cnf)
+            external_loading["concentrated_nodal_forces"] = cnf_list
 
-        pst_list = []
-        for pstconfig in config["external_loading"]["particle_surface_traction"]:
-            pst = ParticleTraction(
-                pset=jnp.array(pstconfig["pset"]),
-                function=self.parsed_config["math_functions"][
-                    pstconfig["math_function_id"]
-                ],
-                dir=pstconfig["dir"],
-                traction=pstconfig["traction"],
-            )
-            pst_list.append(pst)
-        external_loading["concentrated_nodal_forces"] = cnf_list
-        external_loading["particle_surface_traction"] = pst_list
+        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"][
+                        pstconfig["math_function_id"]
+                    ],
+                    dir=pstconfig["dir"],
+                    traction=pstconfig["traction"],
+                )
+                pst_list.append(pst)
+            external_loading["particle_surface_traction"] = pst_list
         self.parsed_config["external_loading"] = external_loading
 
     def _parse_mesh(self, config):
@@ -117,6 +125,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,8 +1,7 @@
 from typing import Tuple
 
 import jax.numpy as jnp
-from jax import jit, vmap, lax
-import jax.debug as db
+from jax import vmap, lax
 from jax.tree_util import register_pytree_node_class
 
 from diffmpm.element import _Element
@@ -144,9 +143,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]
@@ -243,9 +243,7 @@ def compute_strain(self, elements: _Element, dt: float):
         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
@@ -296,7 +294,7 @@ def _step(pid, args):
         args = (dn_dx, temp, strain_rate)
         _, _, strain_rate = lax.fori_loop(0, self.loc.shape[0], _step, args)
         strain_rate = jnp.where(
-            strain_rate < 1e-12, jnp.zeros_like(strain_rate), strain_rate
+            jnp.abs(strain_rate) < 1e-12, jnp.zeros_like(strain_rate), strain_rate
         )
         return strain_rate
 

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]