Update mathieu evp to use rebuilding matrices, fix cartesian ncc corn…

…er case where ncc is zero
DedalusProject · Aug 17, 2022 · 7a114a5 · 7a114a5
1 parent 27f3b46
commit 7a114a5
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 25 deletions.
diff --git a/dedalus/core/arithmetic.py b/dedalus/core/arithmetic.py
@@ -416,19 +416,23 @@ def build_cartesian_ncc_matrix(self, subproblem, ncc_cutoff, max_ncc_terms):
         subproblem_shape = subproblem.coeff_shape(out.domain)
         ncc_rank = len(ncc.tensorsig)
         select_all_comps = tuple(slice(None) for i in range(ncc_rank))
-        for ncc_mode in np.ndindex(self._ncc_data.shape[ncc_rank:]):
-            ncc_coeffs = self._ncc_data[select_all_comps + ncc_mode]
-            if np.max(np.abs(ncc_coeffs)) > ncc_cutoff:
-                mode_matrix = self.cartesian_mode_matrix(subproblem_shape, ncc.domain, arg.domain, out.domain, ncc_mode)
-                matrix = sparse.kron(np.dot(Gamma, ncc_coeffs.ravel()), mode_matrix, format='coo')
-                shape = matrix.shape
-                data.append(matrix.data)
-                rows.append(matrix.row)
-                cols.append(matrix.col)
-        data = np.concatenate(data)
-        rows = np.concatenate(rows)
-        cols = np.concatenate(cols)
-        matrix = sparse.coo_matrix((data, (rows, cols)), shape=shape).tocsr()
+        if np.any(self._ncc_data):
+            for ncc_mode in np.ndindex(self._ncc_data.shape[ncc_rank:]):
+                ncc_coeffs = self._ncc_data[select_all_comps + ncc_mode]
+                if np.max(np.abs(ncc_coeffs)) > ncc_cutoff:
+                    mode_matrix = self.cartesian_mode_matrix(subproblem_shape, ncc.domain, arg.domain, out.domain, ncc_mode)
+                    matrix = sparse.kron(np.dot(Gamma, ncc_coeffs.ravel()), mode_matrix, format='coo')
+                    shape = matrix.shape
+                    data.append(matrix.data)
+                    rows.append(matrix.row)
+                    cols.append(matrix.col)
+            data = np.concatenate(data)
+            rows = np.concatenate(rows)
+            cols = np.concatenate(cols)
+            matrix = sparse.coo_matrix((data, (rows, cols)), shape=shape).tocsr()
+        else:
+            shape = (subproblem.field_size(out), subproblem.field_size(arg))
+            matrix = sparse.csr_matrix(shape, dtype=self.dtype)
         return matrix
 
     @classmethod

diff --git a/examples/evp_1d_mathieu/mathieu_evp.py b/examples/evp_1d_mathieu/mathieu_evp.py
@@ -34,22 +34,24 @@
 
 # Substitutions
 x = dist.local_grid(basis)
+q = dist.Field()
 cos_2x = dist.Field(bases=basis)
 cos_2x['g'] = np.cos(2 * x)
 dx = lambda A: d3.Differentiate(A, coord)
-namespace = locals()
 
-# Eigenvalue solver
-def eigenvalues(q):
-    # Problem
-    problem = d3.EVP([y], eigenvalue=a, namespace=namespace)
-    problem.namespace['q'] = q
-    problem.add_equation("dx(dx(y)) + (a - 2*q*cos_2x)*y = 0")
-    # Solver
-    solver = problem.build_solver()
-    solver.solve_dense(solver.subproblems[0])
-    return np.sort(solver.eigenvalues.real)
-evals = np.array([eigenvalues(q)[:10] for q in q_list])
+# Problem
+problem = d3.EVP([y], eigenvalue=a, namespace=locals())
+problem.add_equation("dx(dx(y)) + (a - 2*q*cos_2x)*y = 0")
+
+# Solver
+solver = problem.build_solver()
+evals = []
+for qi in q_list:
+    q['g'] = qi
+    solver.solve_dense(solver.subproblems[0], rebuild_matrices=True)
+    sorted_evals = np.sort(solver.eigenvalues.real)
+    evals.append(sorted_evals[:10])
+evals = np.array(evals)
 
 # Plot
 fig = plt.figure(figsize=(6, 4))