Key error, multiple particle test cases added

Author: SachinJalan

diffmpm/material.py

@@ -251,15 +251,15 @@ def __compute_stress_per_particle(
 
             shear_rate = jnp.sqrt(
                 2.0 * (strain_r.T @ (strain_r) + strain_r[-3:].T @ strain_r[-3:])
-            )
+            ).squeeze()
 
             # Apparent_viscosity maps shear rate to shear stress
             # Check if shear rate is 0
 
             apparent_viscosity_true = 2.0 * (
-                (self.properties["tau0"] / shear_rate[0, 0]) + self.properties["mu"]
+                (self.properties["tau0"] / shear_rate) + self.properties["mu"]
             )
-            condition = (shear_rate[0, 0] * shear_rate[0, 0]) > (
+            condition = (shear_rate * shear_rate) > (
                 self.properties["critical_shear_rate"]
                 * self.properties["critical_shear_rate"]
             )

tests/test_bingham.py

@@ -146,3 +146,72 @@ def test_compute_stress(particles, state_vars, element, target, dt):
     particles.compute_strain(element, dt)
     stress = particles.material.compute_stress(None, particles, state_vars)
     assert jnp.allclose(stress, target)
+
+
+def test_key_not_present_in_material_properties():
+    with pytest.raises(KeyError):
+        material = Bingham(
+            {
+                "density": 1000,
+                "youngs_modulus": 1.0e7,
+                "poisson_ratio": 0.3,
+                "tau0": 771.8,
+                "critical_shear_rate": 0.2,
+                "ndim": 2,
+            }
+        )
+
+
+@pytest.mark.parametrize(
+    "particles, element, target, state_vars",
+    [
+        (Particles(
+            jnp.array([[0.5, 0.5, 0.5, 0.5]]).reshape(2, 1, 2),
+            (
+                Bingham(
+                    {
+                        "density": 1000,
+                        "youngs_modulus": 1.0e7,
+                        "poisson_ratio": 0.3,
+                        "tau0": 200.0,
+                        "mu": 200.0,
+                        "critical_shear_rate": 0.2,
+                        "ndim": 2,
+                    }
+                )
+            ),
+            jnp.array([0,0]),
+        ),
+        Quadrilateral4Node(
+            (1, 1),
+            1,
+            (4.0, 4.0),
+            [(0, Constraint(0, 2.0)), (0, Constraint(1, 3.0))],
+            Nodes(4, jnp.array([-2, -2, 2, -2, -2, 2, 2, 2]).reshape((4, 1, 2))),
+        ),
+        jnp.array(
+            [
+                -5208520.35574006,
+                -5208613.86694342,
+                0.0,
+                -233.778008402801,
+                0.0,
+                0.0,
+                -5208520.35574006,
+                -5208613.86694342,
+                0.0,
+                -233.778008402801,
+                0.0,
+                0.0,
+            ]
+        ).reshape((2, 6, 1)),
+        {"pressure": jnp.zeros((2, 1))}),
+    ],
+)
+def test_compute_stress_two_particles(particles, state_vars, element, target, dt):
+    particles.update_natural_coords(element)
+    if element.constraints:
+        element.apply_boundary_constraints()
+    particles.compute_strain(element, dt)
+    stress = particles.material.compute_stress(None, particles, state_vars)
+    assert jnp.allclose(stress, target)