Add beginnings of baseline-based corrections. (#235)

* Add beginnings of baseline-based corrections. More for analysis than application.

* Fix bug in baseline corrections.

quartical/calibration/calibrate.py

@@ -11,6 +11,8 @@
                                       make_net_xds_list,
                                       populate_net_xds_list)
 from quartical.interpolation.interpolate import load_and_interpolate_gains
+from quartical.gains.baseline import (compute_baseline_corrections,
+                                      apply_baseline_corrections)
 from loguru import logger  # noqa
 from collections import namedtuple
 
@@ -163,6 +165,17 @@ def add_calibration_graph(
     else:
         net_xds_lod = []
 
+    if output_opts.compute_baseline_corrections:
+        bl_corr_xds_list = compute_baseline_corrections(
+            data_xds_list,
+            gain_xds_lod,
+            t_map_list,
+            f_map_list,
+            d_map_list
+        )
+    else:
+        bl_corr_xds_list = None
+
     # Update the data xarray.Datasets with visibility outputs.
     data_xds_list = make_visibility_output(
         data_xds_list,
@@ -173,8 +186,20 @@ def add_calibration_graph(
         output_opts
     )
 
+    if output_opts.apply_baseline_corrections:
+        data_xds_list = apply_baseline_corrections(
+            data_xds_list,
+            bl_corr_xds_list
+        )
+
     # Return the resulting graphs for the gains and updated xds.
-    return gain_xds_lod, net_xds_lod, data_xds_list, stats_xds_list
+    return (
+        gain_xds_lod,
+        net_xds_lod,
+        data_xds_list,
+        stats_xds_list,
+        bl_corr_xds_list
+    )
 
 
 def make_visibility_output(data_xds_list, solved_gain_xds_lod, t_map_list,

quartical/config/external.py

@@ -111,6 +111,8 @@ class Outputs(Input):
     apply_p_jones_inv: bool = False
     subtract_directions: Optional[List[int]] = None
     net_gains: Optional[List[Any]] = None
+    compute_baseline_corrections: bool = False
+    apply_baseline_corrections: bool = False
 
     def __post_init__(self):
         self.validate_choice_fields()
@@ -131,6 +133,10 @@ def __post_init__(self):
                      "Must be strictly a list or list of lists.")
                 # In the non-nested case, introduce outer list (consistent).
                 self.net_gains = [self.net_gains]
+        if self.apply_baseline_corrections:
+            assert self.compute_baseline_corrections, \
+                ("output.compute_baseline_corrections must be enabled if "
+                 "output.apply_baseline corrections is enabled.")
 
 
 @dataclass

quartical/config/helpstrings.yaml

@@ -129,6 +129,13 @@ output:
     Which model directions to subtract when generating residuals. Must be
     specified as a list of integers e.g. [0, 5, 7]. The default will subtract
     all directions.
+  compute_baseline_corrections:
+    Enable or disable computation of baseline-based corrections. Functionality
+    is currently limited to a solution per-channel, per-chunk. These solutions
+    are useful for analysis and are stored in output.gain_directory.
+  apply_baseline_corrections:
+    Enable or disable application of baseline-based corrections. Extreme
+    caution advised - this can and will lead to overfitting.
 
 
 mad_flags:

quartical/executor.py

@@ -19,6 +19,7 @@
 from quartical.flagging.flagging import finalise_flags, add_mad_graph
 from quartical.scheduling import install_plugin
 from quartical.gains.datasets import write_gain_datasets
+from quartical.gains.baseline import write_baseline_datasets
 from quartical.utils.dask import compute_context
 
 
@@ -128,7 +129,11 @@ def _execute(exitstack):
         output_opts
     )
 
-    gain_xds_lod, net_xds_lod, data_xds_list, stats_xds_list = cal_outputs
+    (gain_xds_lod,
+     net_xds_lod,
+     data_xds_list,
+     stats_xds_list,
+     bl_corr_xds_list) = cal_outputs
 
     if mad_flag_opts.enable:
         data_xds_list = add_mad_graph(data_xds_list, mad_flag_opts)
@@ -149,6 +154,9 @@ def _execute(exitstack):
                                       net_xds_lod,
                                       output_opts)
 
+    bl_corr_writes = write_baseline_datasets(bl_corr_xds_list,
+                                             output_opts)
+
     logger.success("{:.2f} seconds taken to build graph.", time.time() - t0)
 
     logger.info("Computation starting. Please be patient - log messages will "
@@ -158,10 +166,11 @@ def _execute(exitstack):
 
     with compute_context(dask_opts, output_opts, time_str):
 
-        _, _, stats_xds_list = dask.compute(
+        _, _, stats_xds_list, _ = dask.compute(
             ms_writes,
             gain_writes,
             stats_xds_list,
+            bl_corr_writes,
             num_workers=dask_opts.threads,
             optimize_graph=True,
             scheduler=dask_opts.scheduler