Code cleanup

Author: chahak13

diffmpm/element.py

@@ -224,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"""
@@ -302,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))
         )
@@ -317,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."""
@@ -784,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

@@ -286,17 +286,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():

examples/simple_2d_file.py

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