From 961ed0bc1a6bcf31e61f1216159741d613ec54fd Mon Sep 17 00:00:00 2001
From: lecoanet <lecoanet@users.noreply.github.com>
Date: Thu, 6 Jun 2024 00:19:12 +0200
Subject: [PATCH] Modified CartesianTrace to work for tensors with rank >2.
 Added rank 3 tests to test_cartesian_operators.py which pass with the changes
 to operators.py, but did not pass before. (#292)

---
 dedalus/core/operators.py                 |  7 ++--
 dedalus/tests/test_cartesian_operators.py | 40 +++++++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py
index d6a099fa..c78a6a14 100644
--- a/dedalus/core/operators.py
+++ b/dedalus/core/operators.py
@@ -1829,8 +1829,11 @@ class CartesianTrace(Trace):
     def subproblem_matrix(self, subproblem):
         dim = self.coordsys.dim
         trace = np.ravel(np.eye(dim))
-        # Assume all components have the same n_size
-        eye = sparse.identity(subproblem.coeff_size(self.domain), self.dtype, format='csr')
+        # Kronecker up identity for remaining tensor components
+        n_eye = prod(cs.dim for cs in self.tensorsig)
+        # Kronecker up identity for coeff size
+        n_eye *= subproblem.coeff_size(self.domain)
+        eye = sparse.identity(n_eye, self.dtype, format='csr')
         matrix = sparse.kron(trace, eye)
         return matrix
 
diff --git a/dedalus/tests/test_cartesian_operators.py b/dedalus/tests/test_cartesian_operators.py
index 555ec547..5bfbb1aa 100644
--- a/dedalus/tests/test_cartesian_operators.py
+++ b/dedalus/tests/test_cartesian_operators.py
@@ -147,6 +147,23 @@ def test_trace_explicit(basis, N, dealias, dtype, layout):
     assert np.allclose(g[layout], np.trace(f[layout]))
 
 
+@pytest.mark.parametrize('basis', [build_FF, build_FC, build_CC, build_FFF, build_FFC])
+@pytest.mark.parametrize('N', N_range)
+@pytest.mark.parametrize('dealias', dealias_range)
+@pytest.mark.parametrize('dtype', dtype_range)
+@pytest.mark.parametrize('layout', ['c', 'g'])
+def test_trace_rank3_explicit(basis, N, dealias, dtype, layout):
+    """Test explicit evaluation of trace operator for correctness."""
+    c, d, b, r = basis(N, dealias, dtype)
+    # Random tensor field
+    f = d.TensorField((c,c,c), bases=b)
+    f.fill_random(layout='g')
+    # Evaluate trace
+    f.change_layout(layout)
+    g = d3.trace(f).evaluate()
+    assert np.allclose(g[layout], np.trace(f[layout]))
+
+
 @pytest.mark.parametrize('basis', [build_FF, build_FC, build_CC, build_FFF, build_FFC])
 @pytest.mark.parametrize('N', N_range)
 @pytest.mark.parametrize('dealias', dealias_range)
@@ -170,6 +187,29 @@ def test_trace_implicit(basis, N, dealias, dtype):
     assert np.allclose(u['c'], f['c'])
 
 
+@pytest.mark.parametrize('basis', [build_FF, build_FC, build_CC, build_FFF, build_FFC])
+@pytest.mark.parametrize('N', N_range)
+@pytest.mark.parametrize('dealias', dealias_range)
+@pytest.mark.parametrize('dtype', dtype_range)
+def test_trace_rank3_implicit(basis, N, dealias, dtype):
+    """Test implicit evaluation of trace operator for correctness."""
+    c, d, b, r = basis(N, dealias, dtype)
+    # Random scalar field
+    f = d.VectorField(c, bases=b)
+    f.fill_random(layout='g')
+    # Trace LBVP
+    u = d.VectorField(c, bases=b)
+    I = d.TensorField((c,c))
+    dim = len(r)
+    for i in range(dim):
+        I['g'][i,i] = 1
+    problem = d3.LBVP([u], namespace=locals())
+    problem.add_equation("trace(I*u) = dim*f")
+    solver = problem.build_solver()
+    solver.solve()
+    assert np.allclose(u['c'], f['c'])
+
+
 @pytest.mark.parametrize('basis', [build_FF, build_FC, build_CC, build_FFF, build_FFC])
 @pytest.mark.parametrize('N', N_range)
 @pytest.mark.parametrize('dealias', dealias_range)