From 04323c596c73547be5207687e82e904b4d0b56f1 Mon Sep 17 00:00:00 2001
From: Ishaan Desai <ishaandesai@gmail.com>
Date: Mon, 3 Feb 2025 16:24:08 +0100
Subject: [PATCH] Simplify the algorithm for domain decomposition to make sure
 it works for all cases

---
 micro_manager/domain_decomposition.py   | 39 +++++++------------------
 tests/unit/test_domain_decomposition.py | 19 +++++++++++-
 2 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/micro_manager/domain_decomposition.py b/micro_manager/domain_decomposition.py
index b80b87d..0dcd7b8 100644
--- a/micro_manager/domain_decomposition.py
+++ b/micro_manager/domain_decomposition.py
@@ -48,40 +48,23 @@ def decompose_macro_domain(self, macro_bounds: list, ranks_per_axis: list) -> li
             np.prod(ranks_per_axis) == self._size
         ), "Total number of processors provided in the Micro Manager configuration and in the MPI execution command do not match."
 
+        for z in range(ranks_per_axis[2]):
+            for y in range(ranks_per_axis[1]):
+                for x in range(ranks_per_axis[0]):
+                    n = (
+                        x
+                        + y * ranks_per_axis[0]
+                        + z * ranks_per_axis[0] * ranks_per_axis[1]
+                    )
+                    if n == self._rank:
+                        rank_in_axis = [x, y, z]
+
         dx = []
         for d in range(self._dims):
             dx.append(
                 abs(macro_bounds[d * 2 + 1] - macro_bounds[d * 2]) / ranks_per_axis[d]
             )
 
-        rank_in_axis: list[int] = [0] * self._dims
-        if ranks_per_axis[0] == 1:  # if serial in x axis
-            rank_in_axis[0] = 0
-        else:
-            rank_in_axis[0] = self._rank % ranks_per_axis[0]  # x axis
-
-        if self._dims == 2:
-            if ranks_per_axis[1] == 1:  # if serial in y axis
-                rank_in_axis[1] = 0
-            else:
-                rank_in_axis[1] = int(self._rank / ranks_per_axis[0])  # y axis
-        elif self._dims == 3:
-            if ranks_per_axis[2] == 1:  # if serial in z axis
-                rank_in_axis[2] = 0
-            else:
-                rank_in_axis[2] = int(
-                    self._rank / (ranks_per_axis[0] * ranks_per_axis[1])
-                )  # z axis
-
-            if ranks_per_axis[1] == 1:  # if serial in y axis
-                rank_in_axis[1] = 0
-            else:
-                rank_in_axis[1] = (
-                    self._rank - ranks_per_axis[0] * ranks_per_axis[1] * rank_in_axis[2]
-                ) % ranks_per_axis[
-                    2
-                ]  # y axis
-
         mesh_bounds = []
         for d in range(self._dims):
             if rank_in_axis[d] > 0:
diff --git a/tests/unit/test_domain_decomposition.py b/tests/unit/test_domain_decomposition.py
index ef92a93..49330ae 100644
--- a/tests/unit/test_domain_decomposition.py
+++ b/tests/unit/test_domain_decomposition.py
@@ -1,5 +1,5 @@
+import unittest
 from unittest import TestCase
-from unittest.mock import MagicMock
 
 import numpy as np
 
@@ -17,6 +17,23 @@ def setUp(self) -> None:
             8,
         ]  # Cuboid which is not symmetric around origin
 
+    def test_rank1_out_of_4_3d(self):
+        """
+        Check bounds for rank 1 in a setting of axis-wise ranks: [2, 2, 1]
+        """
+        rank = 1
+        size = 4
+        ranks_per_axis = [2, 2, 1]
+        domain_decomposer = DomainDecomposer(3, rank, size)
+        domain_decomposer._dims = 3
+        mesh_bounds = domain_decomposer.decompose_macro_domain(
+            self._macro_bounds_3d, ranks_per_axis
+        )
+
+        print("mesh_bounds", mesh_bounds)
+
+        self.assertTrue(np.allclose(mesh_bounds, [0.0, 1, -2, 0.0, -2, 8]))
+
     def test_rank5_outof_10_3d(self):
         """
         Test domain decomposition for rank 5 in a setting of axis-wise ranks: [1, 2, 5]