From 08f42b008c1620c1ad0726b06b8603d3948f4c42 Mon Sep 17 00:00:00 2001
From: Yannick Augenstein <yaugenst@google.com>
Date: Mon, 20 Nov 2023 09:47:49 -0800
Subject: [PATCH] Fix tests and add tests for temperature, compliance, and
 displacement grads

---
 tests/test_fea2d.py | 65 ++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 61 insertions(+), 4 deletions(-)

diff --git a/tests/test_fea2d.py b/tests/test_fea2d.py
index 6df1ff2..22eeac9 100644
--- a/tests/test_fea2d.py
+++ b/tests/test_fea2d.py
@@ -1,15 +1,29 @@
 import numpy as np
 import pytest
+from autograd.test_util import check_grads
 
 from tofea.fea2d import FEA2D_K, FEA2D_T
 
 
+@pytest.fixture()
+def rng():
+    seed = 36523523
+    return np.random.default_rng(seed)
+
+
 class TestFEA2DK:
     @pytest.fixture()
     def fea2d_k_instance(self):
         fixed = np.zeros((5, 5, 2), dtype=bool)
+        fixed[0] = 1
         return FEA2D_K(fixed)
 
+    @pytest.fixture()
+    def x_and_b(self, fea2d_k_instance, rng):
+        x = rng.random(fea2d_k_instance.shape)
+        b = rng.random(fea2d_k_instance.fixed.shape)
+        return x, b
+
     def test_shape(self, fea2d_k_instance):
         assert fea2d_k_instance.shape == (4, 4)
 
@@ -22,20 +36,50 @@ def test_dofs(self, fea2d_k_instance):
     def test_fixdofs(self, fea2d_k_instance):
         fixdofs = fea2d_k_instance.fixdofs
         assert isinstance(fixdofs, np.ndarray)
-        assert fixdofs.size == 0
+        assert fixdofs.size == fea2d_k_instance.fixed[0].size
 
     def test_freedofs(self, fea2d_k_instance):
         freedofs = fea2d_k_instance.freedofs
         assert isinstance(freedofs, np.ndarray)
-        assert np.all(freedofs == np.arange(50))
+        assert freedofs.size == 50 - fea2d_k_instance.fixdofs.size
+
+    def test_displacement_grads(self, fea2d_k_instance, x_and_b):
+        x, b = x_and_b
+        check_grads(
+            lambda x_: fea2d_k_instance.displacement(x_, b),
+            modes=["fwd", "rev"],
+            order=1,
+        )(x)
+        check_grads(
+            lambda b_: fea2d_k_instance.displacement(x, b_),
+            modes=["fwd", "rev"],
+            order=1,
+        )(b)
+
+    def test_compliance_grads(self, fea2d_k_instance, x_and_b, rng):
+        x, _ = x_and_b
+        d = rng.random(fea2d_k_instance.dofs.shape)
+        check_grads(
+            lambda x_: fea2d_k_instance.compliance(x_, d), modes=["fwd", "rev"], order=1
+        )(x)
+        check_grads(
+            lambda d_: fea2d_k_instance.compliance(x, d_), modes=["fwd", "rev"], order=1
+        )(d)
 
 
 class TestFEA2DT:
     @pytest.fixture()
     def fea2d_t_instance(self):
         fixed = np.zeros((5, 5), dtype=bool)
+        fixed[0, 0] = 1
         return FEA2D_T(fixed)
 
+    @pytest.fixture()
+    def x_and_b(self, fea2d_t_instance, rng):
+        x = rng.random(fea2d_t_instance.shape)
+        b = rng.random(fea2d_t_instance.fixed.shape)
+        return x, b
+
     def test_shape(self, fea2d_t_instance):
         assert fea2d_t_instance.shape == (4, 4)
 
@@ -48,9 +92,22 @@ def test_dofs(self, fea2d_t_instance):
     def test_fixdofs(self, fea2d_t_instance):
         fixdofs = fea2d_t_instance.fixdofs
         assert isinstance(fixdofs, np.ndarray)
-        assert fixdofs.size == 0
+        assert fixdofs.size == 1
 
     def test_freedofs(self, fea2d_t_instance):
         freedofs = fea2d_t_instance.freedofs
         assert isinstance(freedofs, np.ndarray)
-        assert np.all(freedofs == np.arange(25))
+        assert freedofs.size == 25 - fea2d_t_instance.fixdofs.size
+
+    def test_temperature_grads(self, fea2d_t_instance, x_and_b):
+        x, b = x_and_b
+        check_grads(
+            lambda x_: fea2d_t_instance.temperature(x_, b),
+            modes=["fwd", "rev"],
+            order=1,
+        )(x)
+        check_grads(
+            lambda b_: fea2d_t_instance.temperature(x, b_),
+            modes=["fwd", "rev"],
+            order=1,
+        )(b)