Better handling of DD case.

ratt-ru · Dec 10, 2024 · 7efa134 · 7efa134
1 parent 0dedc11
commit 7efa134
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/quartical/calibration/solver.py b/quartical/calibration/solver.py
@@ -344,7 +344,7 @@ def get_collapsed_inputs(
 
     if l_terms:
         n_l_d = max([s.shape[3] for s in l_terms])
-        l_dir_map = np.arange(n_l_d, dtype=np.int32)  # TODO: Wrong.
+        l_dir_map = np.zeros(n_d, dtype=np.int32) if n_l_d > 1 else np.arange(n_d, dtype=np.int32)  # TODO: Wrong.
 
         # TODO: Cache array to avoid allocation?
         l_gain = combine_gains(
@@ -369,7 +369,7 @@ def get_collapsed_inputs(
 
     if r_terms:
         n_r_d = max([s.shape[3] for s in r_terms])
-        r_dir_map = np.arange(n_r_d, dtype=np.int32)  # TODO: Wrong.
+        r_dir_map = np.zeros(n_d, dtype=np.int32) if n_r_d > 1 else np.arange(n_d, dtype=np.int32)  # TODO: Wrong.
 
         r_gain = combine_gains(
             chain_kwargs["gains"][active_term + 1:],