Merge pull request #5 from chahak13/2d

2d `solve` works now. Bug fixed!

Author: kks32

diffmpm/element.py

@@ -127,25 +127,36 @@ def compute_nodal_momentum(self, particles):
         """
 
         def _step(pid, args):
-            pmom, pvel, mom, vel, mapped_pos, el_nodes = args
+            pmom, mom, mapped_pos, el_nodes = args
             mom = mom.at[el_nodes[pid]].add(mapped_pos[pid] @ pmom[pid])
-            vel = vel.at[el_nodes[pid]].add(mapped_pos[pid] @ pvel[pid])
-            return pmom, pvel, mom, vel, mapped_pos, el_nodes
+            return pmom, mom, mapped_pos, el_nodes
 
         self.nodes.momentum = self.nodes.momentum.at[:].set(0)
-        self.nodes.velocity = self.nodes.velocity.at[:].set(0)
         mapped_positions = self.shapefn(particles.reference_loc)
         mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1)
         args = (
             particles.mass * particles.velocity,
-            particles.velocity,
             self.nodes.momentum,
-            self.nodes.velocity,
             mapped_positions,
             mapped_nodes,
         )
-        _, _, self.nodes.momentum, self.nodes.velocity, _, _ = lax.fori_loop(
-            0, len(particles), _step, args
+        _, self.nodes.momentum, _, _ = lax.fori_loop(0, len(particles), _step, args)
+        self.nodes.momentum = jnp.where(
+            self.nodes.momentum < 1e-12,
+            jnp.zeros_like(self.nodes.momentum),
+            self.nodes.momentum,
+        )
+
+    def compute_velocity(self, particles):
+        self.nodes.velocity = jnp.where(
+            self.nodes.mass == 0,
+            self.nodes.velocity,
+            self.nodes.momentum / self.nodes.mass,
+        )
+        self.nodes.velocity = jnp.where(
+            self.nodes.velocity < 1e-12,
+            jnp.zeros_like(self.nodes.velocity),
+            self.nodes.velocity,
         )
 
     def compute_external_force(self, particles):
@@ -213,29 +224,11 @@ def _step(pid, args):
         self.nodes.f_ext, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args)
 
     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
-
         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"""
@@ -291,11 +284,7 @@ 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))
         )
@@ -306,7 +295,16 @@ 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()}")
+        self.nodes.velocity = jnp.where(
+            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.zeros_like(self.nodes.acceleration),
+            self.nodes.acceleration,
+        )
 
     def apply_boundary_constraints(self, *args):
         """Apply boundary conditions for nodal velocity."""
@@ -773,7 +771,6 @@ def _step(pid, args):
             mapped_nodes,
             particles.stress,
         )
-        _step(0, args)
         self.nodes.f_int, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args)
 
     def compute_volume(self):

diffmpm/particle.py

@@ -208,11 +208,12 @@ def update_position_velocity(
             mapped_positions * elements.nodes.acceleration[mapped_ids],
             axis=1,
         )
-        self.velocity = self.velocity.at[:].add(
+        self.velocity = self.velocity.at[:].set(
             lax.cond(
                 velocity_update,
-                lambda nv, na, t: nv,
-                lambda nv, na, t: na * t,
+                lambda sv, nv, na, t: nv,
+                lambda sv, nv, na, t: sv + na * t,
+                self.velocity,
                 nodal_velocity,
                 nodal_acceleration,
                 dt,
@@ -286,17 +287,17 @@ def _compute_strain_rate(self, dn_dx: jnp.ndarray, elements: _Element):
 
         def _step(pid, args):
             dndx, nvel, strain_rate = args
-            # breakpoint()
             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])
             return dndx, nvel, strain_rate
 
         args = (dn_dx, temp, strain_rate)
-        # _step(0, args)
         _, _, strain_rate = lax.fori_loop(0, self.loc.shape[0], _step, args)
-        # breakpoint()
+        strain_rate = jnp.where(
+            strain_rate < 1e-12, jnp.zeros_like(strain_rate), strain_rate
+        )
         return strain_rate
 
     def compute_stress(self, *args):

diffmpm/scheme.py

@@ -14,6 +14,7 @@ def compute_nodal_kinematics(self):
         self.mesh.apply_on_particles("update_natural_coords")
         self.mesh.apply_on_elements("compute_nodal_mass")
         self.mesh.apply_on_elements("compute_nodal_momentum")
+        self.mesh.apply_on_elements("compute_velocity")
         self.mesh.apply_on_elements("apply_boundary_constraints")
 
     def compute_stress_strain(self):
@@ -39,8 +40,6 @@ def compute_particle_kinematics(self):
             args=(self.dt, self.velocity_update),
         )
         # TODO: Apply particle velocity constraints.
-        self.mesh.apply_on_elements("compute_nodal_momentum")
-        self.mesh.apply_on_elements("apply_boundary_constraints")
 
     @abc.abstractmethod
     def precompute_stress_strain():

diffmpm/solver.py

@@ -42,6 +42,7 @@ def __init__(self, mesh, dt, scheme="usf", velocity_update=False):
         self.mesh = mesh
         self.dt = dt
         self.scheme = scheme
+        self.mesh.apply_on_elements("set_particle_element_ids")
         self.mesh.apply_on_particles("compute_volume")
 
     def tree_flatten(self):

examples/simple_2d_file.py

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