Restructure code and add materials subdir

To make the imports easier from the materials subdir, also restructured
other files. This moves `MPM` to a separate file so as to remove
circular imports for materials module.

benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py

@@ -3,7 +3,7 @@
 
 import jax.numpy as jnp
 
-from diffmpm import MPM
+from diffmpm.mpm import MPM
 
 
 def test_benchmarks():

benchmarks/2d/uniaxial_particle_traction/test_benchmark.py

@@ -3,7 +3,7 @@
 
 import jax.numpy as jnp
 
-from diffmpm import MPM
+from diffmpm.mpm import MPM
 
 
 def test_benchmarks():

benchmarks/2d/uniaxial_stress/test_benchmark.py

@@ -3,7 +3,7 @@
 
 import jax.numpy as jnp
 
-from diffmpm import MPM
+from diffmpm.mpm import MPM
 
 
 def test_benchmarks():

diffmpm/__init__.py

@@ -1,47 +1,5 @@
 from importlib.metadata import version
-from pathlib import Path
 
-import diffmpm.writers as writers
-from diffmpm.io import Config
-from diffmpm.solver import MPMExplicit
-
-__all__ = ["MPM", "__version__"]
+__all__ = ["__version__"]
 
 __version__ = version("diffmpm")
-
-
-class MPM:
-    def __init__(self, filepath):
-        self._config = Config(filepath)
-        mesh = self._config.parse()
-        out_dir = Path(self._config.parsed_config["output"]["folder"]).joinpath(
-            self._config.parsed_config["meta"]["title"],
-        )
-
-        write_format = self._config.parsed_config["output"].get("format", None)
-        if write_format is None or write_format.lower() == "none":
-            writer_func = None
-        elif write_format == "npz":
-            writer_func = writers.NPZWriter().write
-        else:
-            raise ValueError(f"Specified output format not supported: {write_format}")
-
-        if self._config.parsed_config["meta"]["type"] == "MPMExplicit":
-            self.solver = MPMExplicit(
-                mesh,
-                self._config.parsed_config["meta"]["dt"],
-                velocity_update=self._config.parsed_config["meta"]["velocity_update"],
-                sim_steps=self._config.parsed_config["meta"]["nsteps"],
-                out_steps=self._config.parsed_config["output"]["step_frequency"],
-                out_dir=out_dir,
-                writer_func=writer_func,
-            )
-        else:
-            raise ValueError("Wrong type of solver specified.")
-
-    def solve(self):
-        """Solve the MPM simulation using JIT solver."""
-        arrays = self.solver.solve_jit(
-            self._config.parsed_config["external_loading"]["gravity"],
-        )
-        return arrays

diffmpm/cli/mpm.py

@@ -1,6 +1,6 @@
 import click
 
-from diffmpm import MPM
+from diffmpm.mpm import MPM
 
 
 @click.command()  # type: ignore

diffmpm/io.py

@@ -1,11 +1,10 @@
 import json
 import tomllib as tl
-from collections import namedtuple
 
 import jax.numpy as jnp
 
 from diffmpm import element as mpel
-from diffmpm import material as mpmat
+from diffmpm import materials as mpmat
 from diffmpm import mesh as mpmesh
 from diffmpm.constraint import Constraint
 from diffmpm.forces import NodalForce, ParticleTraction

diffmpm/materials/__init__.py

@@ -0,0 +1,3 @@
+from diffmpm.materials._base import _Material
+from diffmpm.materials.simple import SimpleMaterial
+from diffmpm.materials.linear_elastic import LinearElastic

diffmpm/materials/_base.py

@@ -0,0 +1,48 @@
+import abc
+from typing import Tuple
+
+
+class _Material(abc.ABC):
+    """Base material class."""
+
+    _props: Tuple[str, ...]
+
+    def __init__(self, material_properties):
+        """Initialize material properties.
+
+        Parameters
+        ----------
+        material_properties: dict
+            A key-value map for various material properties.
+        """
+        self.properties = material_properties
+
+    # @abc.abstractmethod
+    def tree_flatten(self):
+        """Flatten this class as PyTree Node."""
+        return (tuple(), self.properties)
+
+    # @abc.abstractmethod
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        """Unflatten this class as PyTree Node."""
+        del children
+        return cls(aux_data)
+
+    @abc.abstractmethod
+    def __repr__(self):
+        """Repr for Material class."""
+        ...
+
+    @abc.abstractmethod
+    def compute_stress(self):
+        """Compute stress for the material."""
+        ...
+
+    def validate_props(self, material_properties):
+        for key in self._props:
+            if key not in material_properties:
+                raise KeyError(
+                    f"'{key}' should be present in `material_properties` "
+                    f"for {self.__class__.__name__} materials."
+                )

diffmpm/material.py → diffmpm/materials/linear_elastic.py

@@ -1,58 +1,11 @@
-import abc
-from typing import Tuple
-
 import jax.numpy as jnp
 from jax.tree_util import register_pytree_node_class
 
-
-class Material(abc.ABC):
-    """Base material class."""
-
-    _props: Tuple[str, ...]
-
-    def __init__(self, material_properties):
-        """Initialize material properties.
-
-        Parameters
-        ----------
-        material_properties: dict
-            A key-value map for various material properties.
-        """
-        self.properties = material_properties
-
-    # @abc.abstractmethod
-    def tree_flatten(self):
-        """Flatten this class as PyTree Node."""
-        return (tuple(), self.properties)
-
-    # @abc.abstractmethod
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        """Unflatten this class as PyTree Node."""
-        del children
-        return cls(aux_data)
-
-    @abc.abstractmethod
-    def __repr__(self):
-        """Repr for Material class."""
-        ...
-
-    @abc.abstractmethod
-    def compute_stress(self):
-        """Compute stress for the material."""
-        ...
-
-    def validate_props(self, material_properties):
-        for key in self._props:
-            if key not in material_properties:
-                raise KeyError(
-                    f"'{key}' should be present in `material_properties` "
-                    f"for {self.__class__.__name__} materials."
-                )
+from ._base import _Material
 
 
 @register_pytree_node_class
-class LinearElastic(Material):
+class LinearElastic(_Material):
     """Linear Elastic Material."""
 
     _props = ("density", "youngs_modulus", "poisson_ratio")
@@ -114,18 +67,3 @@ def compute_stress(self, dstrain):
         """Compute material stress."""
         dstress = self.de @ dstrain
         return dstress
-
-
-@register_pytree_node_class
-class SimpleMaterial(Material):
-    _props = ("E", "density")
-
-    def __init__(self, material_properties):
-        self.validate_props(material_properties)
-        self.properties = material_properties
-
-    def __repr__(self):
-        return f"SimpleMaterial(props={self.properties})"
-
-    def compute_stress(self, dstrain):
-        return dstrain * self.properties["E"]

diffmpm/materials/newtonian.py

@@ -0,0 +1 @@
+#!/usr/bin/env python3

diffmpm/materials/simple.py

@@ -0,0 +1,18 @@
+from jax.tree_util import register_pytree_node_class
+
+from ._base import _Material
+
+
+@register_pytree_node_class
+class SimpleMaterial(_Material):
+    _props = ("E", "density")
+
+    def __init__(self, material_properties):
+        self.validate_props(material_properties)
+        self.properties = material_properties
+
+    def __repr__(self):
+        return f"SimpleMaterial(props={self.properties})"
+
+    def compute_stress(self, dstrain):
+        return dstrain * self.properties["E"]

diffmpm/mpm.py

@@ -0,0 +1,42 @@
+from pathlib import Path
+
+import diffmpm.writers as writers
+from diffmpm.io import Config
+from diffmpm.solver import MPMExplicit
+
+
+class MPM:
+    def __init__(self, filepath):
+        self._config = Config(filepath)
+        mesh = self._config.parse()
+        out_dir = Path(self._config.parsed_config["output"]["folder"]).joinpath(
+            self._config.parsed_config["meta"]["title"],
+        )
+
+        write_format = self._config.parsed_config["output"].get("format", None)
+        if write_format is None or write_format.lower() == "none":
+            writer_func = None
+        elif write_format == "npz":
+            writer_func = writers.NPZWriter().write
+        else:
+            raise ValueError(f"Specified output format not supported: {write_format}")
+
+        if self._config.parsed_config["meta"]["type"] == "MPMExplicit":
+            self.solver = MPMExplicit(
+                mesh,
+                self._config.parsed_config["meta"]["dt"],
+                velocity_update=self._config.parsed_config["meta"]["velocity_update"],
+                sim_steps=self._config.parsed_config["meta"]["nsteps"],
+                out_steps=self._config.parsed_config["output"]["step_frequency"],
+                out_dir=out_dir,
+                writer_func=writer_func,
+            )
+        else:
+            raise ValueError("Wrong type of solver specified.")
+
+    def solve(self):
+        """Solve the MPM simulation using JIT solver."""
+        arrays = self.solver.solve_jit(
+            self._config.parsed_config["external_loading"]["gravity"],
+        )
+        return arrays

diffmpm/particle.py

@@ -6,7 +6,7 @@
 from jax.typing import ArrayLike
 
 from diffmpm.element import _Element
-from diffmpm.material import Material
+from diffmpm.materials import _Material
 
 
 @register_pytree_node_class
@@ -16,7 +16,7 @@ class Particles(Sized):
     def __init__(
         self,
         loc: ArrayLike,
-        material: Material,
+        material: _Material,
         element_ids: ArrayLike,
         initialized: Optional[bool] = None,
         data: Optional[Tuple[ArrayLike, ...]] = None,
@@ -27,7 +27,7 @@ def __init__(
         ----------
         loc: ArrayLike
             Location of the particles. Expected shape (nparticles, 1, ndim)
-        material: diffmpm.material.Material
+        material: diffmpm.materials._Material
             Type of material for the set of particles.
         element_ids: ArrayLike
             The element ids that the particles belong to. This contains

tests/test_element.py

@@ -5,7 +5,7 @@
 from diffmpm.element import Quadrilateral4Node
 from diffmpm.forces import NodalForce
 from diffmpm.functions import Unit
-from diffmpm.material import SimpleMaterial
+from diffmpm.materials import SimpleMaterial
 from diffmpm.particle import Particles
 
 

tests/test_material.py

@@ -1,7 +1,6 @@
 import jax.numpy as jnp
 import pytest
-
-from diffmpm.material import LinearElastic, SimpleMaterial
+from diffmpm.materials import LinearElastic, SimpleMaterial
 
 material_dstrain_stress_targets = [
     (

tests/test_particle.py

@@ -2,7 +2,7 @@
 import pytest
 
 from diffmpm.element import Quadrilateral4Node
-from diffmpm.material import SimpleMaterial
+from diffmpm.materials import SimpleMaterial
 from diffmpm.particle import Particles