Update example and store more solver results

Author: chahak13

diffmpm/solver.py

@@ -30,9 +30,6 @@ def tree_unflatten(cls, aux_data, children):
         return cls(*children, aux_data[0], scheme=aux_data[1])
 
     def solve(self, nsteps: int, gravity: float):
-        # self.initialize_materials()
-        # self.initialize_mesh()
-        # self.initialize_particles()
         result = {"position": [], "velocity": []}
         for step in tqdm(range(nsteps)):
             self.mpm_scheme.compute_nodal_kinematics()
@@ -52,13 +49,19 @@ def solve_jit(self, nsteps: int, gravity: float):
         result = {
             "position": jnp.zeros((nsteps, nparticles)),
             "velocity": jnp.zeros((nsteps, nparticles)),
+            "strain_energy": jnp.zeros((nsteps, nparticles)),
+            "kinetic_energy": jnp.zeros((nsteps, nparticles)),
+            "total_energy": jnp.zeros((nsteps, nparticles)),
+            "stress": jnp.zeros((nsteps, nparticles)),
+            "strain": jnp.zeros((nsteps, nparticles)),
         }
 
         def _step(i, data):
             self, result = data
             self.mpm_scheme.compute_nodal_kinematics()
             self.mpm_scheme.precompute_stress_strain()
             self.mpm_scheme.compute_forces(gravity)
+            self.mpm_scheme.update_nodal_momentum()
             self.mpm_scheme.compute_particle_kinematics()
             self.mpm_scheme.postcompute_stress_strain()
 
@@ -76,6 +79,38 @@ def _step(i, data):
                     .at[i, idl:idu]
                     .set(self.mesh.particles[j].velocity.squeeze())
                 )
+                result["stress"] = (
+                    result["stress"]
+                    .at[i, idl:idu]
+                    .set(self.mesh.particles[j].stress[:, 0, :].squeeze())
+                )
+                result["strain"] = (
+                    result["strain"]
+                    .at[i, idl:idu]
+                    .set(self.mesh.particles[j].strain[:, 0, :].squeeze())
+                )
+                strain_energy = (
+                    0.5
+                    * self.mesh.particles[j].stress[:, 0, :].squeeze()
+                    * self.mesh.particles[j].strain[:, 0, :].squeeze()
+                    * self.mesh.particles[j].volume.squeeze()
+                )
+                kinetic_energy = (
+                    0.5
+                    * self.mesh.particles[j].velocity.squeeze() ** 2
+                    * self.mesh.particles[j].mass.squeeze()
+                )
+                result["strain_energy"] = (
+                    result["strain_energy"].at[i, idl:idu].set(strain_energy)
+                )
+                result["kinetic_energy"] = (
+                    result["kinetic_energy"].at[i, idl:idu].set(kinetic_energy)
+                )
+                result["total_energy"] = (
+                    result["total_energy"]
+                    .at[i, idl:idu]
+                    .set(strain_energy + kinetic_energy)
+                )
             return (self, result)
 
         _, result = lax.fori_loop(0, nsteps, _step, (self, result))

examples/optim_1d.py

@@ -19,8 +19,8 @@
 b1 = jnp.pi * 0.5
 particles.velocity += 0.1
 particles.set_mass_volume(1.0)
-dt = 0.001
-nsteps = 2000
+dt = 0.01
+nsteps = 1000
 mesh = _MeshBase({"particles": [particles], "elements": elements})
 
 mpm = MPMExplicit(mesh, dt, scheme="usl")
@@ -32,11 +32,9 @@
 
 @jit
 def compute_loss(E, mpm, target_vel):
-    # debug.breakpoint()
     mpm.mesh.particles[0].material.properties["E"] = E
     mpm.mesh.particles[0].velocity = mesh.particles[0].velocity.at[:].set(0.1)
-    # debug.breakpoint()
-    nsteps = 2000
+    nsteps = 1000
     result = mpm.solve_jit(nsteps, 0)
     vel = result["velocity"]
     loss = jnp.linalg.norm(vel - target_vel)
@@ -63,7 +61,7 @@ def optax_adam(params, niter, mpm, target_vel):
     return param_list, loss_list
 
 
-params = 95.0
+params = 107.5
 material = SimpleMaterial({"E": params, "density": 1})
 elements = Linear1D(1, 1, jnp.array([0]))
 particles = Particles(

examples/simple_1d.py

@@ -16,14 +16,28 @@
 # particles.velocity = velocity
 particles.velocity += 0.1
 particles.set_mass_volume(1.0)
-dt = 0.001
-nsteps = 2500
+dt = 0.01
+nsteps = 1000
 mesh = _MeshBase({"particles": [particles], "elements": elements})
 
 mpm = MPMExplicit(mesh, dt, scheme="usl")
 result = mpm.solve_jit(nsteps, 0)
 
 
+res_json = {
+    "x_p": result["position"].squeeze().tolist(),
+    "v_p": result["velocity"].squeeze().tolist(),
+    "strain_p": result["strain"].squeeze().tolist(),
+    "stress_p": result["stress"].squeeze().tolist(),
+    "ke_p": result["kinetic_energy"].squeeze().tolist(),
+    "se_p": result["strain_energy"].squeeze().tolist(),
+    "te_p": result["total_energy"].squeeze().tolist(),
+    "f_int_n": result["f_int"].squeeze().tolist(),
+    "f_ext_n": result["f_ext"].squeeze().tolist(),
+    "f_total_n": result["f_total"].squeeze().tolist(),
+}
+
+
 def analytical_vibration(E, rho, v0, x_loc, L, dt, nsteps):
     tt, vt, xt = [], [], []
     t = 0
@@ -54,4 +68,33 @@ def analytical_vibration(E, rho, v0, x_loc, L, dt, nsteps):
 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()