From 40271b5cd680586bbca38c65801e1467d7b7d3a2 Mon Sep 17 00:00:00 2001
From: Jonathan Kenyon <jonathan.simon.kenyon@gmail.com>
Date: Wed, 11 Dec 2024 08:54:39 +0200
Subject: [PATCH] Expose parameter to control chain collapse. Fix formatting.

---
 quartical/calibration/solver.py       |  8 ++++----
 quartical/config/argument_schema.yaml | 10 ++++++++++
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/quartical/calibration/solver.py b/quartical/calibration/solver.py
index 13ae3b3b..0c3b0aec 100644
--- a/quartical/calibration/solver.py
+++ b/quartical/calibration/solver.py
@@ -208,8 +208,6 @@ def solver_wrapper(
             **{k: chain_kwargs.get(k, None) for k in chain_fields}
         )
 
-        solver_opts.collapse_chain = True  # TODO: Expose + default to this.
-
         if solver_opts.collapse_chain:
             mapping_inputs, chain_inputs, collapsed_term = get_collapsed_inputs(
                 ms_kwargs,
@@ -344,7 +342,8 @@ def get_collapsed_inputs(
 
     if l_terms:
         n_l_d = max([s.shape[3] for s in l_terms])
-        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.
+        dir_map_func = np.zeros if n_l_d > 1 else np.arange
+        l_dir_map = dir_map_func(n_d, dtype=np.int32)
 
         # TODO: Cache array to avoid allocation?
         l_gain = combine_gains(
@@ -369,7 +368,8 @@ def get_collapsed_inputs(
 
     if r_terms:
         n_r_d = max([s.shape[3] for s in r_terms])
-        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.
+        dir_map_func = np.zeros if n_r_d > 1 else np.arange
+        r_dir_map = dir_map_func(n_d, dtype=np.int32)
 
         r_gain = combine_gains(
             chain_kwargs["gains"][active_term + 1:],
diff --git a/quartical/config/argument_schema.yaml b/quartical/config/argument_schema.yaml
index af3b369a..46164086 100644
--- a/quartical/config/argument_schema.yaml
+++ b/quartical/config/argument_schema.yaml
@@ -357,6 +357,16 @@ solver:
       almost always be enabled so that data associated with diverging gains
       is properly flagged.
 
+  collapse_chain:
+    dtype: bool
+    default: True
+    info:
+      Determines whether the chain is collapsed into the minimum number of
+      terms inside the solver. This will typically increase memory footprint,
+      but may speed up calibration when utilising many gain terms. Set to false
+      to apply every term on-the-fly inside the solver (behaviour prior to
+      v0.2.6).
+
   robust:
     dtype: bool
     default: False