Use absolute values to check tolerance threshold

Author: chahak13

diffmpm/element.py

@@ -147,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,
         )
@@ -159,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,
         )
@@ -301,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,
         )

diffmpm/particle.py

@@ -1,9 +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 import vmap, lax
 from jax.tree_util import register_pytree_node_class
 
 from diffmpm.element import _Element
@@ -245,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
@@ -298,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