Add some 1d changes

Author: chahak13

diffmpm/element.py

@@ -339,9 +339,13 @@ class Linear1D(_Element):
     def __init__(
         self,
         nelements: int,
+        total_elements,
         el_len: float,
         constraints: List[Tuple[jnp.ndarray, Constraint]],
         nodes: Nodes = None,
+        concentrated_nodal_forces=[],
+        initialized=None,
+        volume=None,
     ):
         """Initialize Linear1D.
 
@@ -367,6 +371,12 @@ def __init__(
 
         # self.boundary_nodes = boundary_nodes
         self.constraints = constraints
+        self.concentrated_nodal_forces = concentrated_nodal_forces
+        if initialized is None:
+            self.volume = jnp.ones((self.total_elements, 1, 1))
+        else:
+            self.volume = volume
+        self.initialized = True
 
     def id_to_node_ids(self, id: int):
         """
@@ -498,7 +508,7 @@ def f(x):
             ids[0] == ids[1], ids[0], jnp.ones_like(ids[0]) * -1
         )
 
-    def compute_volume(self):
+    def compute_volume(self, *args):
         vol = jnp.ediff1d(self.nodes.loc)
         self.volume = jnp.ones((self.total_elements, 1, 1)) * vol
 

diffmpm/particle.py

@@ -272,15 +272,6 @@ def _compute_strain_rate(self, dn_dx: jnp.ndarray, elements: _Element):
             self.element_ids
         )  # (nparticles, 2, 1)
 
-        # db.print(f"mapped_val[3]: {mapped_vel[3, :, 0]}")
-        # TODO: This will need to change to be more general for ndim.
-        # breakpoint()
-        # 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(2)
 
         def _step(pid, args):

examples/input_1d.toml

@@ -17,13 +17,12 @@ dimension = 1
 scheme = "usl"
 dt = 0.01
 nsteps = 1000
-gravity = 0
-velocity_update = True
+velocity_update = true
 
 [output]
-type = "hdf5"
-file = "results/example_1d_out.hdf5"
-step_frequency = 5
+format = "npz"
+folder = "examples/1d_ex/results/"
+step_frequency = 1
 
 [mesh]
 # type = "file"
@@ -52,3 +51,6 @@ type = "SimpleMaterial"
 file = "examples/particles-1d.json"
 material_id = 0
 init_velocity = 0.1
+
+[external_loading]
+gravity = 0

examples/optim_1d.py

@@ -5,19 +5,19 @@
 from diffmpm.material import SimpleMaterial
 from diffmpm.mesh import Mesh1D
 from diffmpm.particle import Particles
+from diffmpm.constraint import Constraint
 from diffmpm.solver import MPMExplicit
 from jax import value_and_grad, grad, jit
 from tqdm import tqdm
 
 E_true = 100
 material = SimpleMaterial({"E": E_true, "density": 1})
-elements = Linear1D(1, 1, jnp.array([0]))
-particles = Particles(
-    jnp.array([0.5]).reshape(1, 1, 1), material, jnp.array([0])
-)
+cons = [(jnp.array([0]), Constraint(0, 0.0))]
+elements = Linear1D(1, 1, 1, cons)
+particles = Particles(jnp.array([0.5]).reshape(1, 1, 1), material, jnp.array([0]))
 b1 = jnp.pi * 0.5
 particles.velocity += 0.1
-particles.set_mass_volume(1.0)
+# particles.set_mass_volume(1.0)
 dt = 0.01
 nsteps = 1000
 mesh = Mesh1D({"particles": [particles], "elements": elements})
@@ -60,20 +60,17 @@ def optax_adam(params, niter, mpm, target_vel):
     return param_list, loss_list
 
 
-params = 107.5
+params = 105.0
 material = SimpleMaterial({"E": params, "density": 1})
-elements = Linear1D(1, 1, jnp.array([0]))
-particles = Particles(
-    jnp.array([0.5]).reshape(1, 1, 1), material, jnp.array([0])
-)
+cons = [(jnp.array([0]), Constraint(0, 0.0))]
+elements = Linear1D(1, 1, 1, cons)
+particles = Particles(jnp.array([0.5]).reshape(1, 1, 1), material, jnp.array([0]))
 particles.velocity += 0.1
 particles.set_mass_volume(1.0)
 mesh = Mesh1D({"particles": [particles], "elements": elements})
 
 mpm = MPMExplicit(mesh, dt, scheme="usl")
-param_list, loss_list = optax_adam(
-    params, 400, mpm, target_vel
-)  # ADAM optimizer
+param_list, loss_list = optax_adam(params, 400, mpm, target_vel)  # ADAM optimizer
 # print("E: {}".format(result))
 
 fig, ax = plt.subplots(1, 2, figsize=(16, 6))

examples/simple_1d.py

@@ -10,10 +10,8 @@
 E = 100
 material = SimpleMaterial({"E": E, "density": 1})
 cons = [(jnp.array([0]), Constraint(0, 0.0))]
-elements = Linear1D(1, 1, cons)
-particles = Particles(
-    jnp.array([0.5]).reshape(1, 1, 1), material, jnp.array([0])
-)
+elements = Linear1D(1, 1, 1, cons)
+particles = Particles(jnp.array([0.5]).reshape(1, 1, 1), material, jnp.array([0]))
 # b1 = jnp.pi * 0.5
 # velocity = 0.1 * jnp.sin(b1 * particles.loc)
 # particles.velocity = velocity
@@ -24,7 +22,7 @@
 mesh = Mesh1D({"particles": [particles], "elements": elements})
 
 mpm = MPMExplicit(mesh, dt, scheme="usf", velocity_update=True)
-result = mpm.solve(nsteps, 0)
+result = mpm.solve_jit(nsteps, 0)
 
 
 def analytical_vibration(E, rho, v0, x_loc, L, dt, nsteps):
@@ -45,6 +43,7 @@ def analytical_vibration(E, rho, v0, x_loc, L, dt, nsteps):
 ta, va, xa = analytical_vibration(
     E=E, rho=1, v0=0.1, x_loc=0.5, nsteps=nsteps, dt=dt, L=1.0
 )
+breakpoint()
 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")