Add simple version of file input and parser

Author: chahak13

diffmpm/io.py

@@ -0,0 +1,75 @@
+import json
+import tomllib as tl
+
+import jax.numpy as jnp
+import numpy as np
+
+from diffmpm import element as mpel
+from diffmpm import material as mpmat
+from diffmpm import mesh as mpmesh
+from diffmpm.node import Nodes
+from diffmpm.particle import Particles
+
+
+class Config:
+    def __init__(self, filepath):
+        self._filepath = filepath
+        self.config = {}
+        self.parse()
+
+    def parse(self):
+        with open(self._filepath, "rb") as f:
+            self._fileconfig = tl.load(f)
+
+        self._parse_meta(self._fileconfig)
+        self._parse_output(self._fileconfig)
+        self._parse_materials(self._fileconfig)
+        self._parse_particles(self._fileconfig)
+        mesh = self._parse_mesh(self._fileconfig)
+        return mesh
+
+    def _parse_meta(self, config):
+        self.config["meta"] = config["meta"]
+
+    def _parse_output(self, config):
+        self.config["output"] = config["output"]
+
+    def _parse_materials(self, config):
+        materials = []
+        for mat_config in config["materials"]:
+            mat_type = mat_config.pop("type")
+            mat_cls = getattr(mpmat, mat_type)
+            mat = mat_cls(mat_config)
+            materials.append(mat)
+        self.config["materials"] = materials
+
+    def _parse_particles(self, config):
+        particle_sets = []
+        for pset_config in config["particles"]:
+            pmat = self.config["materials"][pset_config["material_id"]]
+            with open(pset_config["file"], "r") as f:
+                ploc = jnp.asarray(json.load(f))
+            peids = jnp.zeros(ploc.shape[0], dtype=jnp.int32)
+            pset = Particles(ploc, pmat, peids)
+            pset.velocity = pset.velocity.at[:].set(
+                pset_config["init_velocity"]
+            )
+            particle_sets.append(pset)
+        self.config["particles"] = particle_sets
+
+    def _parse_mesh(self, config):
+        element_cls = getattr(mpel, config["mesh"]["element"])
+        mesh_cls = getattr(mpmesh, f"Mesh{config['meta']['dimension']}D")
+        if config["mesh"]["type"] == "file":
+            nodes_loc = jnp.asarray(np.loadtxt(config["mesh"]["file"]))
+            nodes = Nodes(len(nodes_loc), nodes_loc)
+            elements = element_cls(nelements, el_len, boundary_nodes)
+        elif config["mesh"]["type"] == "generator":
+            elements = element_cls(
+                config["mesh"]["nelements"],
+                config["mesh"]["element_length"],
+                jnp.asarray(config["mesh"]["boundary_nodes"]),
+            )
+        self.config["elements"] = elements
+        mesh = mesh_cls(self.config)
+        return mesh

diffmpm/particle.py

@@ -51,9 +51,11 @@ def __init__(
         self.loc = loc
 
         if initialized is None:
-            self.mass = jnp.zeros((self.loc.shape[0], 1, 1))
-            self.density = jnp.zeros_like(self.mass)
-            self.volume = jnp.zeros_like(self.mass)
+            self.mass = jnp.ones((self.loc.shape[0], 1, 1))
+            self.density = (
+                jnp.ones_like(self.mass) * self.material.properties["density"]
+            )
+            self.volume = jnp.divide(self.mass, self.density)
             self.velocity = jnp.zeros_like(self.loc)
             self.acceleration = jnp.zeros_like(self.loc)
             self.momentum = jnp.zeros_like(self.loc)

diffmpm/solver.py

@@ -1,8 +1,28 @@
 import jax.numpy as jnp
 from jax import lax
-from diffmpm.scheme import _schemes, USF, USL
-from tqdm import tqdm
 from jax.tree_util import register_pytree_node_class
+from tqdm import tqdm
+
+from diffmpm.io import Config
+from diffmpm.scheme import USF, USL, _schemes
+
+
+# TODO: Move to __init__.py mostly.
+class MPM:
+    def __init__(self, filepath):
+        self._config = Config(filepath)
+        mesh = self._config.parse()
+        if self._config.config["meta"]["type"] == "MPMExplicit":
+            self.solver = MPMExplicit(mesh, self._config.config["meta"]["dt"])
+        else:
+            raise ValueError("Wrong type of solver specified.")
+
+    def solve(self):
+        res = self.solver.solve_jit(
+            self._config.config["meta"]["nsteps"],
+            self._config.config["meta"]["gravity"],
+        )
+        return res
 
 
 @register_pytree_node_class

examples/input_1d.toml

@@ -0,0 +1,49 @@
+# The `meta` group contains top level attributes that govern the
+# behaviour of the MPM Solver.
+#
+# Attributes:
+# title:    The title of the experiment. This is just for the user's
+#           reference.
+# type:     The type of simulation to be used. Allowed values are
+#           {"MPMExplicit"}
+# scheme:   The MPM Scheme used for simulation. Allowed values are
+#           {"usl", "usf"}
+# dt:       Timestep used in the simulation.
+# nsteps:   Number of steps to run the simulation for.
+[meta]
+title = "Example TOML Input for 1D MPM"
+type = "MPMExplicit"
+dimension = 1
+scheme = "usl"
+dt = 0.01
+nsteps = 1000
+gravity = 0
+
+[output]
+type = "hdf5"
+file = "results/example_1d_out.hdf5"
+step_frequency = 5
+
+[mesh]
+# type = "file"
+# file = "mesh-1d.txt"
+# boundary_nodes = "boundary-1d.txt"
+# particle_element_ids = "particles-elements.txt"
+type = "generator"
+nelements = 1
+element_length = 1.0
+boundary_nodes = [0]
+particle_element_ids = [0]
+element = "Linear1D"
+
+[[materials]]
+id = 0
+density = 1
+poisson_ratio = 1
+E = 39.5
+type = "SimpleMaterial"
+
+[[particles]]
+file = "examples/particles-1d.json"
+material_id = 0
+init_velocity = 0.1

examples/particles-1d.json

@@ -0,0 +1 @@
+[[[0.5]]]

examples/simple_1d_file.py

@@ -0,0 +1,71 @@
+import sys
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+from diffmpm.solver import MPM
+
+mpm = MPM(sys.argv[1])
+result = mpm.solve()
+
+
+def analytical_vibration(E, rho, v0, x_loc, L, dt, nsteps):
+    tt, vt, xt = [], [], []
+    t = 0
+    for _ in range(nsteps):
+        omega = 1.0 / L * jnp.sqrt(E / rho)
+        v = v0 * jnp.cos(omega * t)
+        x = x_loc * jnp.exp(v0 / (L * omega) * jnp.sin(omega * t))
+        vt.append(v)
+        xt.append(x)
+        tt.append(t)
+        t += dt
+    return tt, vt, xt
+
+
+E = mpm._config._fileconfig["materials"][0]["E"]
+nsteps = mpm._config._fileconfig["meta"]["nsteps"]
+dt = mpm._config._fileconfig["meta"]["dt"]
+
+# analytical solution at the end of the bar
+ta, va, xa = analytical_vibration(
+    E=E, rho=1, v0=0.1, x_loc=0.5, nsteps=nsteps, dt=dt, L=1.0
+)
+fig, ax = plt.subplots(1, 2, figsize=(16, 6))
+ax[0].plot(ta, va, "r", label="analytical")
+ax[0].plot(ta, result["velocity"].squeeze(), "ob", markersize=2, label="mpm")
+ax[0].legend()
+ax[0].set_title("Velocity")
+
+ax[1].plot(ta, xa, "r", label="analytical")
+ax[1].plot(ta, result["position"].squeeze(), "ob", markersize=2, label="mpm")
+ax[1].legend()
+ax[1].set_title("Position")
+
+fig, ax = plt.subplots()
+ax.plot(
+    ta,
+    result["strain_energy"].squeeze(),
+    "r",
+    linewidth=1,
+    alpha=0.7,
+    label="Strain energy",
+)
+ax.plot(
+    ta,
+    result["kinetic_energy"].squeeze(),
+    "b",
+    linewidth=1,
+    alpha=0.7,
+    label="Kinetic energy",
+)
+ax.plot(
+    ta,
+    result["total_energy"].squeeze(),
+    "k",
+    linewidth=1,
+    alpha=0.7,
+    label="Total energy",
+)
+ax.legend()
+ax.set_title("Energies")
+
+plt.show()