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.

Author: chahak13

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.pydiffmpm/material.py renamed to 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