Add concentrated forces, math funcs & code cleanup

Author: chahak13

diffmpm/element.py

@@ -0,0 +1,6 @@
+from collections import namedtuple
+
+NodalForce = namedtuple("NodalForce", ("node_ids", "function", "dir", "force"))
+ParticleTraction = namedtuple(
+    "ParticleTraction", ("pset", "function", "dir", "traction")
+)

diffmpm/forces.py

@@ -0,0 +1,29 @@
+import abc
+import jax.numpy as jnp
+
+
+class Function(abc.ABC):
+    def __init__(self, id):
+        self.id = id
+
+    @abc.abstractmethod
+    def value(self):
+        ...
+
+
+class Unit(Function):
+    def __init__(self, id):
+        super().__init__(id)
+
+    def value(self, x):
+        return 1.0
+
+
+class Linear(Function):
+    def __init__(self, id, xvalues, fxvalues):
+        self.xvalues = xvalues
+        self.fxvalues = fxvalues
+        super().__init__(id)
+
+    def value(self, x):
+        return jnp.interp(x, self.xvalues, self.fxvalues)

diffmpm/functions.py

@@ -1,21 +1,22 @@
 import json
 import tomllib as tl
+from collections import namedtuple
 
 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.constraint import Constraint
-from diffmpm.node import Nodes
+from diffmpm.forces import NodalForce, ParticleTraction
+from diffmpm.functions import Unit, Linear
 from diffmpm.particle import Particles
 
 
 class Config:
     def __init__(self, filepath):
         self._filepath = filepath
-        self.config = {}
+        self.parsed_config = {}
         self.parse()
 
     def parse(self):
@@ -26,14 +27,16 @@ def parse(self):
         self._parse_output(self._fileconfig)
         self._parse_materials(self._fileconfig)
         self._parse_particles(self._fileconfig)
+        self._parse_math_functions(self._fileconfig)
+        self._parse_external_loading(self._fileconfig)
         mesh = self._parse_mesh(self._fileconfig)
         return mesh
 
     def _parse_meta(self, config):
-        self.config["meta"] = config["meta"]
+        self.parsed_config["meta"] = config["meta"]
 
     def _parse_output(self, config):
-        self.config["output"] = config["output"]
+        self.parsed_config["output"] = config["output"]
 
     def _parse_materials(self, config):
         materials = []
@@ -42,21 +45,67 @@ def _parse_materials(self, config):
             mat_cls = getattr(mpmat, mat_type)
             mat = mat_cls(mat_config)
             materials.append(mat)
-        self.config["materials"] = materials
+        self.parsed_config["materials"] = materials
 
     def _parse_particles(self, config):
         particle_sets = []
         for pset_config in config["particles"]:
-            pmat = self.config["materials"][pset_config["material_id"]]
+            pmat = self.parsed_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"]
-            )
+            pset.velocity = pset.velocity.at[:].set(pset_config["init_velocity"])
             particle_sets.append(pset)
-        self.config["particles"] = particle_sets
+        self.parsed_config["particles"] = particle_sets
+
+    def _parse_math_functions(self, config):
+        flist = []
+        for i, fnconfig in enumerate(config["math_functions"]):
+            if fnconfig["type"] == "Linear":
+                fn = Linear(
+                    i,
+                    jnp.array(fnconfig["xvalues"]),
+                    jnp.array(fnconfig["fxvalues"]),
+                )
+                flist.append(fn)
+            else:
+                raise NotImplementedError(
+                    "Function type other than `Linear` not yet supported"
+                )
+        self.parsed_config["math_functions"] = flist
+
+    def _parse_external_loading(self, config):
+        external_loading = {}
+        external_loading["gravity"] = jnp.array(config["external_loading"]["gravity"])
+        cnf_list = []
+        for cnfconfig in config["external_loading"]["concentrated_nodal_forces"]:
+            if "math_function_id" in cnfconfig:
+                fn = self.parsed_config["math_functions"][cnfconfig["math_function_id"]]
+            else:
+                fn = Unit(-1)
+            cnf = NodalForce(
+                node_ids=jnp.array(cnfconfig["node_ids"]),
+                math_function=fn,
+                dir=cnfconfig["dir"],
+                force=cnfconfig["force"],
+            )
+            cnf_list.append(cnf)
+
+        pst_list = []
+        for pstconfig in config["external_loading"]["particle_surface_traction"]:
+            pst = ParticleTraction(
+                pset=jnp.array(cnfconfig["pset"]),
+                math_function=self.parsed_config["math_functions"][
+                    pstconfig["math_function_id"]
+                ],
+                dir=cnfconfig["dir"],
+                traction=cnfconfig["traction"],
+            )
+            pst_list.append(pst)
+        external_loading["concentrated_nodal_forces"] = cnf_list
+        external_loading["particle_surface_traction"] = pst_list
+        self.parsed_config["external_loading"] = external_loading
 
     def _parse_mesh(self, config):
         element_cls = getattr(mpel, config["mesh"]["element"])
@@ -65,16 +114,19 @@ def _parse_mesh(self, config):
             (jnp.asarray(c["node_ids"]), Constraint(c["dir"], c["velocity"]))
             for c in config["mesh"]["constraints"]
         ]
-        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":
+        if config["mesh"]["type"] == "generator":
             elements = element_cls(
                 config["mesh"]["nelements"],
                 config["mesh"]["element_length"],
                 constraints,
+                concentrated_nodal_forces=self.parsed_config["external_loading"][
+                    "concentrated_nodal_forces"
+                ],
+            )
+        else:
+            raise NotImplementedError(
+                "Mesh type other than `generator` not yet supported."
             )
-        self.config["elements"] = elements
-        mesh = mesh_cls(self.config)
+        self.parsed_config["elements"] = elements
+        mesh = mesh_cls(self.parsed_config)
         return mesh

diffmpm/io.py

@@ -184,7 +184,7 @@ def update_natural_coords(self, elements: _Element):
         self.reference_loc = xi_coords
 
     def update_position_velocity(
-        self, elements: _Element, dt: float, velocity_update: bool = False
+        self, elements: _Element, dt: float, velocity_update: bool
     ):
         """
         Transfer nodal velocity to particles and update particle position.
@@ -281,14 +281,17 @@ def _compute_strain_rate(self, dn_dx: jnp.ndarray, elements: _Element):
 
         # TODO: This will need to change to be more general for ndim.
         # breakpoint()
-        # L = jnp.einsum("ijk, ikj -> ijk", dn_dx, mapped_vel).sum(axis=2)
+        # L = jnp.einsum("ijk, ikj -> ijk", dn_dx, mapped_vel.squeeze(-1)).sum(
+        #     axis=2
+        # )
         # strain_rate = strain_rate.at[:, 0, :].add(L)
 
         # For 2d
-        temp = mapped_vel.squeeze()
+        temp = mapped_vel.squeeze(2)
 
         def _step(pid, args):
             dndx, nvel, strain_rate = args
+            # breakpoint()
             matmul = dndx[pid].T @ nvel[pid]
             strain_rate = strain_rate.at[pid, 0].add(matmul[0, 0])
             strain_rate = strain_rate.at[pid, 1].add(matmul[1, 1])
@@ -298,7 +301,9 @@ def _step(pid, args):
             return dndx, nvel, strain_rate
 
         args = (dn_dx, temp, strain_rate)
+        # _step(0, args)
         _, _, strain_rate = lax.fori_loop(0, self.loc.shape[0], _step, args)
+        # breakpoint()
         return strain_rate
 
     def compute_stress(self, *args):

diffmpm/particle.py

@@ -4,8 +4,9 @@
 
 
 class _MPMScheme(abc.ABC):
-    def __init__(self, mesh, dt):
+    def __init__(self, mesh, dt, velocity_update):
         self.mesh = mesh
+        self.velocity_update = velocity_update
         self.dt = dt
 
     def compute_nodal_kinematics(self):
@@ -20,9 +21,12 @@ def compute_stress_strain(self):
         self.mesh.apply_on_particles("update_volume")
         self.mesh.apply_on_particles("compute_stress")
 
-    def compute_forces(self, gravity):
+    def compute_forces(self, gravity, step):
         self.mesh.apply_on_elements("compute_external_force")
         self.mesh.apply_on_elements("compute_body_force", args=(gravity,))
+        self.mesh.apply_on_elements(
+            "apply_concentrated_nodal_forces", args=(step * self.dt,)
+        )
         self.mesh.apply_on_elements("compute_internal_force")
         # self.mesh.apply_on_elements("apply_force_boundary_constraints")
 
@@ -31,7 +35,8 @@ def compute_particle_kinematics(self):
             "update_nodal_acceleration_velocity", args=(self.dt,)
         )
         self.mesh.apply_on_particles(
-            "update_position_velocity", args=(self.dt,)
+            "update_position_velocity",
+            args=(self.dt, self.velocity_update),
         )
         # TODO: Apply particle velocity constraints.
         self.mesh.apply_on_elements("compute_nodal_momentum")
@@ -49,8 +54,8 @@ def postcompute_stress_strain():
 class USF(_MPMScheme):
     """USF Scheme solver."""
 
-    def __init__(self, mesh, dt):
-        super().__init__(mesh, dt)
+    def __init__(self, mesh, dt, velocity_update):
+        super().__init__(mesh, dt, velocity_update)
 
     def precompute_stress_strain(self):
         self.compute_stress_strain()
@@ -62,8 +67,8 @@ def postcompute_stress_strain(self):
 class USL(_MPMScheme):
     """USL Scheme solver."""
 
-    def __init__(self, mesh, dt):
-        super().__init__(mesh, dt)
+    def __init__(self, mesh, dt, velocity_update):
+        super().__init__(mesh, dt, velocity_update)
 
     def precompute_stress_strain(self):
         pass

diffmpm/scheme.py

@@ -2,6 +2,7 @@
 from jax import lax
 from jax.tree_util import register_pytree_node_class
 from tqdm import tqdm
+from collections import defaultdict
 
 from diffmpm.io import Config
 from diffmpm.scheme import USF, USL, _schemes
@@ -12,30 +13,32 @@ 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"])
+        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"],
+            )
         else:
             raise ValueError("Wrong type of solver specified.")
 
     def solve(self):
         res = self.solver.solve(
-            self._config.config["meta"]["nsteps"],
-            self._config.config["meta"]["gravity"],
+            self._config.parsed_config["meta"]["nsteps"],
+            self._config.parsed_config["external_loading"]["gravity"],
         )
         return res
 
 
 @register_pytree_node_class
 class MPMExplicit:
-    def __init__(self, mesh, dt, scheme="usf"):
+    def __init__(self, mesh, dt, scheme="usf", velocity_update=False):
         if scheme == "usf":
-            self.mpm_scheme = USF(mesh, dt)
+            self.mpm_scheme = USF(mesh, dt, velocity_update)
         elif scheme == "usl":
-            self.mpm_scheme = USL(mesh, dt)
+            self.mpm_scheme = USL(mesh, dt, velocity_update)
         else:
-            raise ValueError(
-                f"Please select scheme from {_schemes}. Found {scheme}"
-            )
+            raise ValueError(f"Please select scheme from {_schemes}. Found {scheme}")
         self.mesh = mesh
         self.dt = dt
         self.scheme = scheme
@@ -51,16 +54,17 @@ def tree_unflatten(cls, aux_data, children):
         return cls(*children, aux_data[0], scheme=aux_data[1])
 
     def solve(self, nsteps: int, gravity: float | jnp.ndarray):
-        result = {"position": [], "velocity": []}
+        result = defaultdict(list)
         for step in tqdm(range(nsteps)):
             self.mpm_scheme.compute_nodal_kinematics()
             self.mpm_scheme.precompute_stress_strain()
-            self.mpm_scheme.compute_forces(gravity)
+            self.mpm_scheme.compute_forces(gravity, step)
             self.mpm_scheme.compute_particle_kinematics()
             self.mpm_scheme.postcompute_stress_strain()
             for pset in self.mesh.particles:
-                result["position"].append(pset.loc)
-                result["velocity"].append(pset.velocity)
+                # result["position"].append(pset.loc)
+                # result["velocity"].append(pset.velocity)
+                result["stress"].append(pset.stress)
 
         result = {k: jnp.asarray(v) for k, v in result.items()}
         return result

diffmpm/solver.py

@@ -18,6 +18,7 @@ scheme = "usl"
 dt = 0.01
 nsteps = 1000
 gravity = 0
+velocity_update = True
 
 [output]
 type = "hdf5"