Fix nested product Jacobian tests (#614)

SleipnirGroup · Jul 27, 2024 · 8817f89 · 8817f89
1 parent 42eb9b1
commit 8817f89
Show file tree

Hide file tree

Showing 2 changed files with 117 additions and 59 deletions.
diff --git a/jormungandr/test/autodiff/jacobian_test.py b/jormungandr/test/autodiff/jacobian_test.py
@@ -1,6 +1,5 @@
 from jormungandr.autodiff import Jacobian, VariableMatrix
 import numpy as np
-import pytest
 
 
 def test_y_eq_x():
@@ -71,38 +70,69 @@ def test_products():
     assert (J.value() == expected_J).all()
 
 
-@pytest.mark.skip(reason="Fails")
 def test_nested_products():
-    z = VariableMatrix(1)
-    z[0].set_value(1)
-    x = VariableMatrix(3)
-    x[0] = 1 * z[0]
-    x[1] = 2 * z[0]
-    x[2] = 3 * z[0]
-
-    J = Jacobian(x, z).value()
-    assert J[0, 0] == 1.0
-    assert J[1, 0] == 2.0
-    assert J[2, 0] == 3.0
+    x = VariableMatrix(1)
+    x[0].set_value(3)
+    assert x.value(0) == 3.0
 
-    #     [x₁x₂]
-    # y = [x₂x₃]
-    #     [x₁x₃]
-    #
-    #         [x₂  x₁  0 ]
-    # dy/dx = [0   x₃  x₂]
-    #         [x₃  0   x₁]
-    #
-    #         [2  1  0]
-    # dy/dx = [0  3  2]
-    #         [3  0  1]
+    #     [ 5x]   [15]
+    # y = [ 7x] = [21]
+    #     [11x]   [33]
     y = VariableMatrix(3)
-    y[0] = x[0] * x[1]
-    y[1] = x[1] * x[2]
-    y[2] = x[0] * x[2]
+    y[0] = 5 * x[0]
+    y[1] = 7 * x[0]
+    y[2] = 11 * x[0]
+    assert y.value(0) == 15.0
+    assert y.value(1) == 21.0
+    assert y.value(2) == 33.0
+
+    #     [y₁y₂]   [15⋅21]   [315]
+    # z = [y₂y₃] = [21⋅33] = [693]
+    #     [y₁y₃]   [15⋅33]   [495]
+    z = VariableMatrix(3)
+    z[0] = y[0] * y[1]
+    z[1] = y[1] * y[2]
+    z[2] = y[0] * y[2]
+    assert z.value(0) == 315.0
+    assert z.value(1) == 693.0
+    assert z.value(2) == 495.0
+
+    #     [ 5x]
+    # y = [ 7x]
+    #     [11x]
+    #
+    #         [ 5]
+    # dy/dx = [ 7]
+    #         [11]
     J = Jacobian(y, x)
+    assert J.get().value()[0, 0] == 5.0
+    assert J.get().value()[1, 0] == 7.0
+    assert J.get().value()[2, 0] == 11.0
+    assert J.value()[0, 0] == 5.0
+    assert J.value()[1, 0] == 7.0
+    assert J.value()[2, 0] == 11.0
+
+    #     [y₁y₂]
+    # z = [y₂y₃]
+    #     [y₁y₃]
+    #
+    #         [y₂  y₁  0 ]   [21  15   0]
+    # dz/dy = [0   y₃  y₂] = [ 0  33  21]
+    #         [y₃  0   y₁]   [33   0  15]
+    J = Jacobian(z, y)
+    expected_J = np.array([[21.0, 15.0, 0.0], [0.0, 33.0, 21.0], [33.0, 0.0, 15.0]])
+    assert (J.get().value() == expected_J).all()
+    assert (J.value() == expected_J).all()
 
-    expected_J = np.array([[2.0, 1.0, 0.0], [0.0, 3.0, 2.0], [3.0, 0.0, 1.0]])
+    #     [y₁y₂]   [5x⋅ 7x]   [35x²]
+    # z = [y₂y₃] = [7x⋅11x] = [77x²]
+    #     [y₁y₃]   [5x⋅11x]   [55x²]
+    #
+    #         [ 70x]   [210]
+    # dz/dx = [154x] = [462]
+    #         [110x] = [330]
+    J = Jacobian(z, x)
+    expected_J = np.array([[210.0], [462.0], [330.0]])
     assert (J.get().value() == expected_J).all()
     assert (J.value() == expected_J).all()
 

diff --git a/test/src/autodiff/JacobianTest.cpp b/test/src/autodiff/JacobianTest.cpp
@@ -83,44 +83,72 @@ TEST_CASE("Jacobian - Products", "[Jacobian]") {
 }
 
 TEST_CASE("Jacobian - Nested products", "[Jacobian]") {
-  SKIP("Fails");
-
   sleipnir::scope_exit exit{
       [] { CHECK(sleipnir::GlobalPoolResource().blocks_in_use() == 0u); }};
 
-  sleipnir::VariableMatrix z{1};
-  z(0).SetValue(1);
-  sleipnir::VariableMatrix x{3};
-  x(0) = 1 * z(0);
-  x(1) = 2 * z(0);
-  x(2) = 3 * z(0);
-
-  auto J = sleipnir::Jacobian(x, z);
-  CHECK(J.Get().Value(0, 0) == 1.0);
-  CHECK(J.Get().Value(1, 0) == 2.0);
-  CHECK(J.Get().Value(2, 0) == 3.0);
-  CHECK(J.Value().coeff(0, 0) == 1.0);
-  CHECK(J.Value().coeff(1, 0) == 2.0);
-  CHECK(J.Value().coeff(2, 0) == 3.0);
+  sleipnir::VariableMatrix x{1};
+  x(0).SetValue(3);
+  CHECK(x.Value(0) == 3.0);
 
-  //     [x₁x₂]
-  // y = [x₂x₃]
-  //     [x₁x₃]
+  //     [ 5x]   [15]
+  // y = [ 7x] = [21]
+  //     [11x]   [33]
+  sleipnir::VariableMatrix y{3};
+  y(0) = 5 * x(0);
+  y(1) = 7 * x(0);
+  y(2) = 11 * x(0);
+  CHECK(y.Value(0) == 15.0);
+  CHECK(y.Value(1) == 21.0);
+  CHECK(y.Value(2) == 33.0);
+
+  //     [y₁y₂]   [15⋅21]   [315]
+  // z = [y₂y₃] = [21⋅33] = [693]
+  //     [y₁y₃]   [15⋅33]   [495]
+  sleipnir::VariableMatrix z{3};
+  z(0) = y(0) * y(1);
+  z(1) = y(1) * y(2);
+  z(2) = y(0) * y(2);
+  CHECK(z.Value(0) == 315.0);
+  CHECK(z.Value(1) == 693.0);
+  CHECK(z.Value(2) == 495.0);
+
+  //     [ 5x]
+  // y = [ 7x]
+  //     [11x]
   //
-  //         [x₂  x₁  0 ]
-  // dy/dx = [0   x₃  x₂]
-  //         [x₃  0   x₁]
+  //         [ 5]
+  // dy/dx = [ 7]
+  //         [11]
+  auto J = sleipnir::Jacobian(y, x);
+  CHECK(J.Get().Value(0, 0) == 5.0);
+  CHECK(J.Get().Value(1, 0) == 7.0);
+  CHECK(J.Get().Value(2, 0) == 11.0);
+  CHECK(J.Value().coeff(0, 0) == 5.0);
+  CHECK(J.Value().coeff(1, 0) == 7.0);
+  CHECK(J.Value().coeff(2, 0) == 11.0);
+
+  //     [y₁y₂]
+  // z = [y₂y₃]
+  //     [y₁y₃]
   //
-  //         [2  1  0]
-  // dy/dx = [0  3  2]
-  //         [3  0  1]
-  sleipnir::VariableMatrix y{3};
-  y(0) = x(0) * x(1);
-  y(1) = x(1) * x(2);
-  y(2) = x(0) * x(2);
-  J = sleipnir::Jacobian(y, x);
+  //         [y₂  y₁  0 ]   [21  15   0]
+  // dz/dy = [0   y₃  y₂] = [ 0  33  21]
+  //         [y₃  0   y₁]   [33   0  15]
+  J = sleipnir::Jacobian(z, y);
+  Eigen::MatrixXd expectedJ{
+      {21.0, 15.0, 0.0}, {0.0, 33.0, 21.0}, {33.0, 0.0, 15.0}};
+  CHECK(J.Get().Value() == expectedJ);
+  CHECK(J.Value().toDense() == expectedJ);
 
-  Eigen::MatrixXd expectedJ{{2.0, 1.0, 0.0}, {0.0, 3.0, 2.0}, {3.0, 0.0, 1.0}};
+  //     [y₁y₂]   [5x⋅ 7x]   [35x²]
+  // z = [y₂y₃] = [7x⋅11x] = [77x²]
+  //     [y₁y₃]   [5x⋅11x]   [55x²]
+  //
+  //         [ 70x]   [210]
+  // dz/dx = [154x] = [462]
+  //         [110x] = [330]
+  J = sleipnir::Jacobian(z, x);
+  expectedJ = Eigen::MatrixXd{{210.0}, {462.0}, {330.0}};
   CHECK(J.Get().Value() == expectedJ);
   CHECK(J.Value().toDense() == expectedJ);
 }