Add tests

Author: chahak13

tests/test_element.py

@@ -202,3 +202,11 @@ def test_compute_internal_force(self, elements, particles):
     def test_compute_volume(self, elements):
         elements.compute_volume()
         assert jnp.allclose(elements.volume, jnp.array([1]).reshape(1, 1, 1))
+
+    def test_apply_particle_traction_forces(self, elements, particles):
+        particles.traction += jnp.array([1, 0])
+        elements.apply_particle_traction_forces(particles)
+        assert jnp.allclose(
+            elements.nodes.f_ext,
+            jnp.array([[0.5, 0], [0.5, 0], [0.5, 0], [0.5, 0]]).reshape(4, 1, 2),
+        )

tests/test_particle.py

@@ -47,3 +47,15 @@ def test_compute_strain(self, elements, particles):
     def test_compute_volume(self, elements, particles):
         particles.compute_volume(elements, elements.total_elements)
         assert jnp.allclose(particles.volume, jnp.array([0.5, 0.5]).reshape(2, 1, 1))
+
+    def test_assign_traction(self, elements, particles):
+        particles.compute_volume(elements, elements.total_elements)
+        particles.assign_traction(jnp.array([0]), 1, 10)
+        assert jnp.allclose(
+            particles.traction, jnp.array([[0, 7.071068], [0, 0]]).reshape(2, 1, 2)
+        )
+
+    def test_zero_traction(self, particles):
+        particles.traction += 1
+        particles.zero_traction()
+        assert jnp.all(particles.traction == 0)