Code cleanup; doesn't change any functionality

Author: chahak13

diffmpm/element.py

@@ -31,6 +31,14 @@ def tree_flatten(self):
     def tree_unflatten(self):
         ...
 
+    @abc.abstractmethod
+    def shapefn(self):
+        ...
+
+    @abc.abstractmethod
+    def shapefn_grad(self):
+        ...
+
 
 @register_pytree_node_class
 class Linear1D(_Element):
@@ -64,10 +72,10 @@ def __init__(
             IDs of nodes that are supposed to be fixed (boundary).
         """
         self.nelements = nelements
-        self.ids: jnp.ndarray = jnp.arange(nelements)
-        self.el_len: float = el_len
+        self.ids = jnp.arange(nelements)
+        self.el_len = el_len
         if nodes is None:
-            self.nodes: Nodes = Nodes(
+            self.nodes = Nodes(
                 nelements + 1,
                 jnp.arange(nelements + 1).reshape(-1, 1, 1) * el_len,
             )
@@ -136,7 +144,7 @@ def id_to_node_vel(self, id: int):
         """
         return self.nodes.velocity[jnp.array([id, id + 1])].reshape(2, 1)
 
-    def shapefn(self, xi):
+    def shapefn(self, xi: float | jnp.ndarray):
         """
         Evaluate linear shape function.
 
@@ -164,7 +172,7 @@ def shapefn(self, xi):
         )
         return result
 
-    def _shapefn_natural_grad(self, xi):
+    def _shapefn_natural_grad(self, xi: float | jnp.ndarray):
         """
         Calculate the gradient of shape function.
 
@@ -197,7 +205,7 @@ def _shapefn_natural_grad(self, xi):
         # )
         return result.reshape(-1, 2)
 
-    def shapefn_grad(self, xi, coords):
+    def shapefn_grad(self, xi: float | jnp.ndarray, coords: jnp.ndarray):
         """
         Gradient of shape function in physical coordinates.
 
@@ -219,8 +227,6 @@ def shapefn_grad(self, xi, coords):
             raise ValueError(
                 f"`x` should be of size (npoints, 1, ndim); found {xi.shape}"
             )
-        # natural_grad = self._shapefn_natural_grad(x)
-        # result = natural_grad * 2 / (coords[1] - coords[0])
         grad_sf = self._shapefn_natural_grad(xi)
         _jacobian = grad_sf @ coords
 
@@ -286,8 +292,6 @@ def _step(pid, args):
             mapped_positions,
             mapped_nodes,
         )
-        # breakpoint()
-        # _step(0, args)
         _, self.nodes.mass, _, _ = lax.fori_loop(0, len(particles), _step, args)
 
     def compute_nodal_momentum(self, particles):
@@ -489,7 +493,8 @@ def _step(pid, args):
             0, len(particles), _step, args
         )
 
-    def update_nodal_momentum(self, particles, dt, *args):
+    def update_nodal_momentum(self, particles, dt: float, *args):
+        """Update the nodal momentum based on total force on nodes."""
         total_force = self.nodes.get_total_force()
         self.nodes.acceleration = self.nodes.acceleration.at[:].set(
             jnp.divide(total_force, self.nodes.mass)
@@ -500,13 +505,15 @@ def update_nodal_momentum(self, particles, dt, *args):
         self.nodes.momentum = self.nodes.momentum.at[:].add(total_force * dt)
 
     def apply_boundary_constraints(self, *args):
+        """Apply boundary conditions for nodal velocity."""
         self.nodes.velocity = self.nodes.velocity.at[self.boundary_nodes].set(0)
         self.nodes.momentum = self.nodes.momentum.at[self.boundary_nodes].set(0)
         self.nodes.acceleration = self.nodes.acceleration.at[
             self.boundary_nodes
         ].set(0)
 
     def apply_force_boundary_constraints(self, *args):
+        """Apply boundary conditions for nodal forces."""
         self.nodes.f_int = self.nodes.f_int.at[self.boundary_nodes].set(0)
         self.nodes.f_ext = self.nodes.f_ext.at[self.boundary_nodes].set(0)
         self.nodes.f_damp = self.nodes.f_damp.at[self.boundary_nodes].set(0)

diffmpm/mesh.py

@@ -1,22 +1,25 @@
-import itertools
-from typing import Sequence
+import abc
+from typing import Iterable
 
 import jax.numpy as jnp
-from jax import jit, lax, vmap
 from jax.tree_util import register_pytree_node_class
-from jax_tqdm import loop_tqdm
-from tqdm import tqdm
 
 from diffmpm.element import _Element
-from diffmpm.node import Nodes
-from diffmpm.shapefn import Linear1DShapeFn, Linear4NodeQuad
+from diffmpm.particle import Particles
 
 
-@register_pytree_node_class
-class _MeshBase:
+class _MeshBase(abc.ABC):
+    """
+    Base class for Meshes.
+
+    Note: If attributes other than elements and particles are added
+    then the child class should also implement `tree_flatten` and
+    `tree_unflatten` correctly or that information will get lost.
+    """
+
     def __init__(self, config: dict):
         """Initialize mesh using configuration."""
-        self.particles: Sequence = config["particles"]
+        self.particles: Iterable[Particles, ...] = config["particles"]
         self.elements: _Element = config["elements"]
 
     # TODO: Convert to using jax directives for loop
@@ -42,211 +45,32 @@ def tree_unflatten(cls, aux_data, children):
         return cls({"particles": children[0], "elements": children[1]})
 
 
-class Mesh1D:
-    """
-    1D Mesh class with nodes, elements, and particles.
-    """
-
-    def __init__(
-        self,
-        nelements,
-        material,
-        domain_size,
-        boundary_nodes,
-        *,
-        ppe=1,
-        particle_distribution="uniform",
-        elements=None,
-        nodes=None,
-        particles=None,
-        shapefn=None,
-        dim=1,
-    ):
-        """
-        Construct a 1D Mesh.
-
-        Arguments
-        ---------
-        nelements : int
-            Number of elements in the mesh.
-        material : diffmpm.material.Material
-            Material to meshed.
-        domain_size : float
-            The size of the domain in consideration.
-        boundary_nodes : array_like
-            Node ids of boundary nodes of the mesh. Needs to be a JAX
-        array.
-        ppe : int
-            Number of particles per element in Mesh.
-        """
-        self.dim = dim
-        self.material = material
-        self.shapefn = (
-            Linear1DShapeFn(self.dim)
-            if (
-                shapefn is None
-                or type(shapefn) is object
-                or isinstance(shapefn, Mesh1D)
-            )
-            else shapefn
-        )
-        self.domain_size = domain_size
-        self.nelements = nelements
-        self.element_length = domain_size / nelements
-        self.elements = jnp.arange(nelements) if elements is None else elements
-        nnodes = nelements + 1
-        self.nodes = (
-            Nodes(
-                nnodes,
-                jnp.arange(nelements + 1) * self.element_length,
-                jnp.zeros(nnodes),
-                jnp.zeros(nnodes),
-                jnp.zeros(nnodes),
-                jnp.zeros(nnodes),
-                jnp.zeros(nnodes),
-                jnp.zeros(nnodes),
-            )
-            if (
-                nodes is None
-                or type(nodes) is object
-                or isinstance(nodes, Mesh1D)
-            )
-            else nodes
-        )
-        self.boundary_nodes = boundary_nodes
-        self.ppe = ppe
-        self.particles = (
-            self._init_particles(particle_distribution)
-            if (
-                particles is None
-                or type(particles) is object
-                or isinstance(particles, Mesh1D)
-            )
-            else particles
-        )
-        return
-
-
-class Mesh2D:
-    """
-    2D Mesh class with nodes, elements, and particles.
-    """
+@register_pytree_node_class
+class Mesh1D(_MeshBase):
+    """1D Mesh class with nodes, elements, and particles."""
 
-    def __init__(
-        self,
-        nelements,
-        material,
-        domain_size,
-        boundary_nodes,
-        *,
-        ppe=1,
-        particle_distribution="uniform",
-        elements=None,
-        nodes=None,
-        particles=None,
-        shapefn=None,
-        dim=1,
-    ):
+    def __init__(self, config: dict):
         """
-        Construct a 2D Mesh using 4-Node Quadrilateral Elements.
-
-        Nodes and elements are numbered as
-
-                    0---0---0---0---0
-                    | 8 | 9 | 10| 11|
-                 10 0---0---0---0---0
-                    | 4 | 5 | 6 | 7 |
-                  5 0---0---0---0---0 9
-                    | 0 | 1 | 2 | 3 |
-                    0---0---0---0---0
-                    0   1   2   3   4
-
+        Initialize a 1D Mesh.
 
         Arguments
         ---------
-        nelements : array_like
-            Number of elements in the mesh in the x and y direction.
-        material : diffmpm.material.Material
-            Material to meshed.
-        domain_size : 4-tuple, array_like
-            The boundaries of the domain. Should be of the form
-        (x_min, x_max, y_min, y_max)
-        boundary_nodes : array_like
-            Node ids of boundary nodes of the mesh. Needs to be a JAX
-        array.
-        ppe : int
-            Number of particles per element in Mesh.
+        config: dict
+            Configuration to be used for initialization. It _should_
+        contain `elements` and `particles` keys.
         """
-        self.dim = 2
-        self.material = material
-        self.shapefn = (
-            Linear4NodeQuad()
-            if (
-                shapefn is None
-                or type(shapefn) is object
-                or isinstance(shapefn, Mesh1D)
-            )
-            else shapefn
-        )
-        self.domain_size = domain_size
-        self.nelements = jnp.asarray(nelements)
-        self.element_length = jnp.array(
-            [
-                (domain_size[1] - domain_size[0]) / nelements[0],
-                (domain_size[3] - domain_size[2]) / nelements[1],
-            ]
-        )
-        self.elements = (
-            jnp.arange(self.nelements[0] * self.nelements[1])
-            if elements is None
-            else elements
-        )
-        nnodes = jnp.product(self.nelements + 1)
-        coords = jnp.asarray(
-            list(
-                itertools.product(
-                    jnp.arange(nelements[1] + 1), jnp.arange(nelements[0] + 1)
-                )
-            )
-        )
-        node_positions = (
-            jnp.asarray([coords[:, 1], coords[:, 0]]).T * self.element_length
-        )
-
-        self.nodes = (
-            Nodes(
-                nnodes,
-                node_positions,
-                jnp.zeros((nnodes, 2)),
-                jnp.zeros(nnodes),
-                jnp.zeros((nnodes, 2)),
-                jnp.zeros((nnodes, 2)),
-                jnp.zeros((nnodes, 2)),
-                jnp.zeros((nnodes, 2)),
-            )
-            if (
-                nodes is None
-                or type(nodes) is object
-                or isinstance(nodes, Mesh1D)
-            )
-            else nodes
-        )
-        self.boundary_nodes = boundary_nodes
-        self.ppe = ppe
-        self.particles = particles
-        return
+        super().__init__(config)
 
 
 if __name__ == "__main__":
-    from diffmpm.utils import _show_example
-    from diffmpm.particle import Particles
     from diffmpm.element import Linear1D
     from diffmpm.material import SimpleMaterial
+    from diffmpm.utils import _show_example
 
     particles = Particles(
         jnp.array([[[1]]]),
         SimpleMaterial({"E": 2, "density": 1}),
         jnp.array([0]),
     )
     elements = Linear1D(2, 1, jnp.array([0]))
-    _show_example(_MeshBase({"particles": [particles], "elements": elements}))
+    _show_example(Mesh1D({"particles": [particles], "elements": elements}))