diff --git a/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py b/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py index ae72923..4dcb077 100644 --- a/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py +++ b/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(): diff --git a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py index 356d0a3..995ca16 100644 --- a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py +++ b/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(): diff --git a/benchmarks/2d/uniaxial_stress/test_benchmark.py b/benchmarks/2d/uniaxial_stress/test_benchmark.py index f04e820..0dd6af8 100644 --- a/benchmarks/2d/uniaxial_stress/test_benchmark.py +++ b/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(): diff --git a/diffmpm/__init__.py b/diffmpm/__init__.py index faa8316..a138300 100644 --- a/diffmpm/__init__.py +++ b/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 diff --git a/diffmpm/cli/mpm.py b/diffmpm/cli/mpm.py index aebc4ba..0b4b9d7 100644 --- a/diffmpm/cli/mpm.py +++ b/diffmpm/cli/mpm.py @@ -1,6 +1,6 @@ import click -from diffmpm import MPM +from diffmpm.mpm import MPM @click.command() # type: ignore diff --git a/diffmpm/io.py b/diffmpm/io.py index d6e4573..66447f5 100644 --- a/diffmpm/io.py +++ b/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 diff --git a/diffmpm/materials/__init__.py b/diffmpm/materials/__init__.py new file mode 100644 index 0000000..ce35083 --- /dev/null +++ b/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 diff --git a/diffmpm/materials/_base.py b/diffmpm/materials/_base.py new file mode 100644 index 0000000..d30b15b --- /dev/null +++ b/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." + ) diff --git a/diffmpm/material.py b/diffmpm/materials/linear_elastic.py similarity index 55% rename from diffmpm/material.py rename to diffmpm/materials/linear_elastic.py index 09230d4..098c10e 100644 --- a/diffmpm/material.py +++ b/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"] diff --git a/diffmpm/materials/newtonian.py b/diffmpm/materials/newtonian.py new file mode 100644 index 0000000..e5a0d9b --- /dev/null +++ b/diffmpm/materials/newtonian.py @@ -0,0 +1 @@ +#!/usr/bin/env python3 diff --git a/diffmpm/materials/simple.py b/diffmpm/materials/simple.py new file mode 100644 index 0000000..d9cf15d --- /dev/null +++ b/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"] diff --git a/diffmpm/mpm.py b/diffmpm/mpm.py new file mode 100644 index 0000000..b06eaee --- /dev/null +++ b/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 diff --git a/diffmpm/particle.py b/diffmpm/particle.py index 1bb3d70..586fec0 100644 --- a/diffmpm/particle.py +++ b/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 diff --git a/tests/test_element.py b/tests/test_element.py index 50881d9..ff8d92e 100644 --- a/tests/test_element.py +++ b/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 diff --git a/tests/test_material.py b/tests/test_material.py index 2e041d7..66cb4dc 100644 --- a/tests/test_material.py +++ b/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 = [ ( diff --git a/tests/test_particle.py b/tests/test_particle.py index d7dedaa..d67bc2f 100644 --- a/tests/test_particle.py +++ b/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