Merge pull request #15 from chahak13/cli

Add a CLI to the repo

Author: kks32

benchmarks/2d/__init__.py

@@ -1,30 +1,32 @@
 import os
 from pathlib import Path
 import jax.numpy as jnp
-from diffmpm.solver import MPM
+from diffmpm import MPM
 
-curr_filepath = Path(__file__).absolute()
-curr_dir = curr_filepath.parent
-os.chdir(curr_dir)
-mpm = MPM("mpm-nodal-forces.toml")
-mpm.solve()
 
-result = jnp.load("results/uniaxial-nodal-forces/particles_0300.npz")
-## Step 300
-assert jnp.round(result["stress"][0, :, 0].min() - 0.5925210678182377, 6) == 0.0
-assert jnp.round(result["stress"][0, :, 0].max() - 0.5974539476363379, 6) == 0.0
+def test_benchmarks():
+    curr_filepath = Path(__file__).absolute()
+    curr_dir = curr_filepath.parent
+    os.chdir(curr_dir)
+    mpm = MPM("mpm-nodal-forces.toml")
+    mpm.solve()
 
-## Step 510
-result = jnp.load("results/uniaxial-nodal-forces/particles_0510.npz")
-assert jnp.round(result["stress"][0, :, 0].min() - 1.0026665338366039, 5) == 0.0
-assert jnp.round(result["stress"][0, :, 0].max() - 1.0112231542459431, 5) == 0.0
+    result = jnp.load("results/uniaxial-nodal-forces/particles_0300.npz")
+    ## Step 300
+    assert jnp.round(result["stress"][0, :, 0].min() - 0.5925210678182377, 6) == 0.0
+    assert jnp.round(result["stress"][0, :, 0].max() - 0.5974539476363379, 6) == 0.0
 
-## Step 750
-result = jnp.load("results/uniaxial-nodal-forces/particles_0750.npz")
-assert jnp.round(result["stress"][0, :, 0].min() - 1.0000053000532143, 5) == 0.0
-assert jnp.round(result["stress"][0, :, 0].max() - 1.0000225119807862, 5) == 0.0
+    ## Step 510
+    result = jnp.load("results/uniaxial-nodal-forces/particles_0510.npz")
+    assert jnp.round(result["stress"][0, :, 0].min() - 1.0026665338366039, 5) == 0.0
+    assert jnp.round(result["stress"][0, :, 0].max() - 1.0112231542459431, 5) == 0.0
 
-## Step 990
-result = jnp.load("results/uniaxial-nodal-forces/particles_0990.npz")
-assert jnp.round(result["stress"][0, :, 0].min() - 0.9999990078443788, 5) == 0.0
-assert jnp.round(result["stress"][0, :, 0].max() - 0.9999990292713694, 5) == 0.0
+    ## Step 750
+    result = jnp.load("results/uniaxial-nodal-forces/particles_0750.npz")
+    assert jnp.round(result["stress"][0, :, 0].min() - 1.0000053000532143, 5) == 0.0
+    assert jnp.round(result["stress"][0, :, 0].max() - 1.0000225119807862, 5) == 0.0
+
+    ## Step 990
+    result = jnp.load("results/uniaxial-nodal-forces/particles_0990.npz")
+    assert jnp.round(result["stress"][0, :, 0].min() - 0.9999990078443788, 5) == 0.0
+    assert jnp.round(result["stress"][0, :, 0].max() - 0.9999990292713694, 5) == 0.0

benchmarks/2d/uniaxial_nodal_forces/__init__.py

@@ -1,29 +1,31 @@
 import os
 from pathlib import Path
 import jax.numpy as jnp
-from diffmpm.solver import MPM
+from diffmpm import MPM
 
-curr_filepath = Path(__file__).absolute()
-curr_dir = curr_filepath.parent
-os.chdir(curr_dir)
-mpm = MPM("mpm-particle-traction.toml")
-mpm.solve()
-result = jnp.load("results/uniaxial-particle-traction/particles_0300.npz")
-## Step 300
-assert jnp.round(result["stress"][0, :, 0].min() - 0.4450086768966724, 5) == 0.0
-assert jnp.round(result["stress"][0, :, 0].max() - 0.5966527842046769, 5) == 0.0
 
-## Step 510
-result = jnp.load("results/uniaxial-particle-traction/particles_0510.npz")
-assert jnp.round(result["stress"][0, :, 0].min() - 0.7528092313640623, 5) == 0.0
-assert jnp.round(result["stress"][0, :, 0].max() - 1.0109599915279937, 5) == 0.0
+def test_benchmarks():
+    curr_filepath = Path(__file__).absolute()
+    curr_dir = curr_filepath.parent
+    os.chdir(curr_dir)
+    mpm = MPM("mpm-particle-traction.toml")
+    mpm.solve()
+    result = jnp.load("results/uniaxial-particle-traction/particles_0300.npz")
+    ## Step 300
+    assert jnp.round(result["stress"][0, :, 0].min() - 0.4450086768966724, 5) == 0.0
+    assert jnp.round(result["stress"][0, :, 0].max() - 0.5966527842046769, 5) == 0.0
 
-## Step 750
-result = jnp.load("results/uniaxial-particle-traction/particles_0750.npz")
-assert jnp.round(result["stress"][0, :, 0].min() - 0.7500090055681591, 5) == 0.0
-assert jnp.round(result["stress"][0, :, 0].max() - 1.0000224746314728, 5) == 0.0
+    ## Step 510
+    result = jnp.load("results/uniaxial-particle-traction/particles_0510.npz")
+    assert jnp.round(result["stress"][0, :, 0].min() - 0.7528092313640623, 5) == 0.0
+    assert jnp.round(result["stress"][0, :, 0].max() - 1.0109599915279937, 5) == 0.0
 
-## Step 990
-result = jnp.load("results/uniaxial-particle-traction/particles_0990.npz")
-assert jnp.round(result["stress"][0, :, 0].min() - 0.750002924022295, 5) == 0.0
-assert jnp.round(result["stress"][0, :, 0].max() - 0.9999997782938734, 5) == 0.0
+    ## Step 750
+    result = jnp.load("results/uniaxial-particle-traction/particles_0750.npz")
+    assert jnp.round(result["stress"][0, :, 0].min() - 0.7500090055681591, 5) == 0.0
+    assert jnp.round(result["stress"][0, :, 0].max() - 1.0000224746314728, 5) == 0.0
+
+    ## Step 990
+    result = jnp.load("results/uniaxial-particle-traction/particles_0990.npz")
+    assert jnp.round(result["stress"][0, :, 0].min() - 0.750002924022295, 5) == 0.0
+    assert jnp.round(result["stress"][0, :, 0].max() - 0.9999997782938734, 5) == 0.0

benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py

@@ -1,17 +1,19 @@
 import os
 from pathlib import Path
 import jax.numpy as jnp
-from diffmpm.solver import MPM
+from diffmpm import MPM
 
-curr_filepath = Path(__file__).absolute()
-curr_dir = curr_filepath.parent
-os.chdir(curr_dir)
-mpm = MPM("mpm-uniaxial-stress.toml")
-mpm.solve()
 
-result = jnp.load("results/uniaxial-stress-usf/particles_09.npz")
-true_stress_yy = -1.0
-true_stress_xx = 0.0
+def test_benchmarks():
+    curr_filepath = Path(__file__).absolute()
+    curr_dir = curr_filepath.parent
+    os.chdir(curr_dir)
+    mpm = MPM("mpm-uniaxial-stress.toml")
+    mpm.solve()
 
-assert jnp.round(result["stress"][0, :, 1].max() - true_stress_yy, 8) == 0.0
-assert jnp.round(result["stress"][0, :, 0].max() - true_stress_xx, 8) == 0.0
+    result = jnp.load("results/uniaxial-stress-usf/particles_09.npz")
+    true_stress_yy = -1.0
+    true_stress_xx = 0.0
+
+    assert jnp.round(result["stress"][0, :, 1].max() - true_stress_yy, 8) == 0.0
+    assert jnp.round(result["stress"][0, :, 0].max() - true_stress_xx, 8) == 0.0

benchmarks/2d/uniaxial_particle_traction/__init__.py

@@ -0,0 +1 @@
+

benchmarks/2d/uniaxial_particle_traction/test_benchmark.py

@@ -1 +1,47 @@
-#!/usr/bin/env python3
+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__"]
+
+__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."""
+        arrays = self.solver.solve_jit(
+            self._config.parsed_config["external_loading"]["gravity"],
+        )
+        return arrays

benchmarks/2d/uniaxial_stress/__init__.py

@@ -0,0 +1,13 @@
+import click
+
+from diffmpm import MPM
+
+
+@click.command()
+@click.option(
+    "-f", "--file", "filepath", required=True, type=str, help="Input TOML file"
+)
+@click.version_option(package_name="diffmpm")
+def mpm(filepath):
+    solver = MPM(filepath)
+    solver.solve()