Swap locations of two functions

`set_particle_element_ids` is moved to element.py since the element id
mapping is based on what type of element the particles lie in. On the
other hand, `update_natural_coords` was moved to particle.py since it is
an intrinsic property of the particles and is just updated based on the
elements passed as argument.

Author: chahak13

diffmpm/element.py

@@ -233,33 +233,35 @@ def shapefn_grad(self, xi: float | jnp.ndarray, coords: jnp.ndarray):
         result = grad_sf.T @ jnp.linalg.inv(_jacobian)
         return result.T
 
-    # TODO: See if this is generic enough to be moved to
-    # `Particles` instead.
-    def update_particle_natural_coords(self, particles):
+    def set_particle_element_ids(self, particles):
         """
-        Update natural coordinates for the particles.
+        Set the element IDs for the particles.
 
-        Whenever the particles' physical coordinates change, their
-        natural coordinates need to be updated. This function updates
-        the natural coordinates of the particles based on the element
-        a particle is a part of. The update formula is
-
-        :math:`xi = (2x - (x_1^e + x_2^e))  / (x_2^e - x_1^e)`
+        If the particle doesn't lie between the boundaries of any
+        element, it sets the element index to -1.
+        """
 
-        If a particle is not in any element (element_id = -1), its
-        natural coordinate is set to 0.
+        @jit
+        def f(x):
+            idl = (
+                len(self.nodes.loc)
+                - 1
+                - jnp.asarray(self.nodes.loc[::-1] <= x).nonzero(
+                    size=1, fill_value=-1
+                )[0][-1]
+            )
+            idg = (
+                jnp.asarray(self.nodes.loc > x).nonzero(size=1, fill_value=-1)[
+                    0
+                ][0]
+                - 1
+            )
+            return (idl, idg)
 
-        Arguments
-        ---------
-        particles: diffmpm.particle.Particles
-            Particles whose natural coordinates need to be updated based
-        on these elements.
-        """
-        t = self.id_to_node_loc(particles.element_ids)
-        xi_coords = (particles.loc - (t[:, 0, ...] + t[:, 1, ...]) / 2) * (
-            2 / (t[:, 1, ...] - t[:, 0, ...])
+        ids = vmap(f)(particles.loc)
+        particles.element_ids = jnp.where(
+            ids[0] == ids[1], ids[0], jnp.ones_like(ids[0]) * -1
         )
-        particles.reference_loc = xi_coords
 
     # Mapping from particles to nodes (P2G)
     def compute_nodal_mass(self, particles):

diffmpm/particle.py

@@ -143,38 +143,31 @@ def set_mass_volume(self, m: float | jnp.ndarray):
             )
         self.volume = jnp.divide(self.mass, self.material.properties["density"])
 
-    # TODO: This needs to be in element so that each type of element
-    # can govern how ot check if particle is in an element.
-    # This implementation is true only for 1D.
-    def set_particle_element_ids(self, elements: _Element):
+    def update_natural_coords(self, elements: _Element):
         """
-        Set the element IDs for the particles.
+        Update natural coordinates for the particles.
 
-        If the particle doesn't lie between the boundaries of any
-        element, it sets the element index to -1.
-        """
+        Whenever the particles' physical coordinates change, their
+        natural coordinates need to be updated. This function updates
+        the natural coordinates of the particles based on the element
+        a particle is a part of. The update formula is
 
-        @jit
-        def f(x):
-            idl = (
-                len(elements.nodes.loc)
-                - 1
-                - jnp.asarray(elements.nodes.loc[::-1] <= x).nonzero(
-                    size=1, fill_value=-1
-                )[0][-1]
-            )
-            idg = (
-                jnp.asarray(elements.nodes.loc > x).nonzero(
-                    size=1, fill_value=-1
-                )[0][0]
-                - 1
-            )
-            return (idl, idg)
+        :math:`xi = (2x - (x_1^e + x_2^e))  / (x_2^e - x_1^e)`
+
+        If a particle is not in any element (element_id = -1), its
+        natural coordinate is set to 0.
 
-        ids = vmap(f)(self.loc)
-        self.element_ids = jnp.where(
-            ids[0] == ids[1], ids[0], jnp.ones_like(ids[0]) * -1
+        Arguments
+        ---------
+        elements: diffmpm.element._Element
+            Elements based on which to update the natural coordinates
+        of the particles.
+        """
+        t = elements.id_to_node_loc(self.element_ids)
+        xi_coords = (self.loc - (t[:, 0, ...] + t[:, 1, ...]) / 2) * (
+            2 / (t[:, 1, ...] - t[:, 0, ...])
         )
+        self.reference_loc = xi_coords
 
     def update_position_velocity(self, elements: _Element, dt: float):
         """

diffmpm/scheme.py

@@ -9,8 +9,8 @@ def __init__(self, mesh, dt):
         self.dt = dt
 
     def compute_nodal_kinematics(self):
-        self.mesh.apply_on_elements("update_particle_natural_coords")
-        self.mesh.apply_on_particles("set_particle_element_ids")
+        self.mesh.apply_on_particles("update_natural_coords")
+        self.mesh.apply_on_elements("set_particle_element_ids")
         self.mesh.apply_on_elements("compute_nodal_mass")
         self.mesh.apply_on_elements("compute_nodal_momentum")
         self.mesh.apply_on_elements("apply_boundary_constraints")

examples/optim_1d.py

@@ -3,7 +3,7 @@
 import optax
 from diffmpm.element import Linear1D
 from diffmpm.material import SimpleMaterial
-from diffmpm.mesh import _MeshBase
+from diffmpm.mesh import Mesh1D
 from diffmpm.particle import Particles
 from diffmpm.solver import MPMExplicit
 from jax import value_and_grad, grad, jit
@@ -20,7 +20,7 @@
 particles.set_mass_volume(1.0)
 dt = 0.01
 nsteps = 1000
-mesh = _MeshBase({"particles": [particles], "elements": elements})
+mesh = Mesh1D({"particles": [particles], "elements": elements})
 
 mpm = MPMExplicit(mesh, dt, scheme="usl")
 true_result = mpm.solve_jit(nsteps, 0)
@@ -68,7 +68,7 @@ def optax_adam(params, niter, mpm, target_vel):
 )
 particles.velocity += 0.1
 particles.set_mass_volume(1.0)
-mesh = _MeshBase({"particles": [particles], "elements": elements})
+mesh = Mesh1D({"particles": [particles], "elements": elements})
 
 mpm = MPMExplicit(mesh, dt, scheme="usl")
 param_list, loss_list = optax_adam(