From f86f3099e2041a4dc44255c48e3d5ed6e16f5423 Mon Sep 17 00:00:00 2001
From: "Keaton J. Burns" <keaton.burns@gmail.com>
Date: Sat, 30 Mar 2024 10:09:31 -0400
Subject: [PATCH] Add print_subproblem_ranks for NLBVP

---
 dedalus/core/solvers.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/dedalus/core/solvers.py b/dedalus/core/solvers.py
index 5e34bddc..364eea56 100644
--- a/dedalus/core/solvers.py
+++ b/dedalus/core/solvers.py
@@ -456,6 +456,17 @@ def __init__(self, problem, **kw):
         self.F = F_handler.fields
         logger.debug('Finished NLBVP instantiation')
 
+    def print_subproblem_ranks(self, subproblems=None):
+        """Print rank of each subproblem LHS."""
+        if subproblems is None:
+            subproblems = self.subproblems
+        # Check matrix rank
+        for i, sp in enumerate(subproblems):
+            if not hasattr(sp, 'dF_min'):
+                continue
+            dF = sp.dF_min.A
+            print(f"MPI rank: {self.dist.comm.rank}, subproblem: {i}, group: {sp.group}, matrix rank: {np.linalg.matrix_rank(dF)}/{dF.shape[0]}, cond: {np.linalg.cond(dF):.1e}")
+
     def newton_iteration(self, damping=1):
         """Update solution with a Newton iteration."""
         # Compute RHS