WIP Gain smoothing (do not merge)

.github/workflows/ci.yaml

@@ -18,7 +18,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-20.04, ubuntu-22.04]
-        python-version: ["3.9", "3.10"]
+        python-version: ["3.10"]
 
     steps:
       - name: Set up Python ${{ matrix.python-version }}
@@ -77,10 +77,10 @@ jobs:
     # Run on tagged push (generated by tbump).
     if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
     steps:
-      - name: Set up Python 3.8
+      - name: Set up Python 3.10
         uses: actions/setup-python@v4
         with:
-          python-version: 3.8
+          python-version: 3.10
 
       - name: Install poetry
         uses: abatilo/actions-poetry@v2
@@ -120,4 +120,4 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v3
       - name: Release
-        uses: softprops/[email protected]
+        uses: softprops/[email protected]

pyproject.toml

@@ -24,7 +24,7 @@ include = [
 ]
 
 [tool.poetry.dependencies]
-python = "^3.9"
+python = "^3.10"
 astro-tigger-lsm = ">=1.7.2, <=1.7.3"
 codex-africanus = {extras = ["dask", "scipy", "astropy", "python-casacore"], version = ">=0.3.4, <=0.3.4"}
 colorama = ">=0.4.6, <=0.4.6"
@@ -40,6 +40,12 @@ requests = ">=2.31.0, <=2.31.0"
 "ruamel.yaml" = ">=0.17.26, <=0.17.40"
 stimela = "2.0rc8"
 tbump = ">=6.10.0, <=6.11.0"
+ducc0 = "^0.31.0"
+sympy = "^1.12"
+nifty8 = {git = "https://gitlab.mpcdf.mpg.de/ift/nifty.git", branch = "NIFTy_8"}
+
+[tool.poetry.extras]
+degrid = ["ducc0", "sympy"]
 
 [tool.poetry.scripts]
 goquartical = 'quartical.executor:execute'

quartical/apps/backup.py

@@ -67,7 +67,7 @@ def backup():
     # This call exists purely to get the relevant shape and dtype info.
     data_xds_list = xds_from_storage_ms(
         args.ms_path,
-        columns=column_name,
+        columns=(column_name, "ROWID"),
         index_cols=("TIME",),
         group_cols=("FIELD_ID", "DATA_DESC_ID", "SCAN_NUMBER"),
     )
@@ -81,7 +81,7 @@ def backup():
     # Repeat of above call but now with correct chunking information.
     data_xds_list = xds_from_storage_ms(
         args.ms_path,
-        columns=column_name,
+        columns=(column_name, "ROWID"),
         index_cols=("TIME",),
         group_cols=("FIELD_ID", "DATA_DESC_ID", "SCAN_NUMBER"),
         chunks=chunks

quartical/apps/summary.py

@@ -1,6 +1,9 @@
 import argparse
 from pathlib import Path
-from daskms import xds_from_storage_ms, xds_from_storage_table
+from daskms.experimental.fragments import (
+    xds_from_ms_fragment,
+    xds_from_table_fragment
+)
 from daskms.fsspec_store import DaskMSStore
 import numpy as np
 import dask.array as da
@@ -44,7 +47,7 @@ def configure_loguru(output_dir):
 def antenna_info(path):
 
     # NOTE: Assume one dataset for now.
-    ant_xds = xds_from_storage_table(path + "::ANTENNA")[0]
+    ant_xds = xds_from_table_fragment(path + "::ANTENNA")[0]
 
     antenna_names = ant_xds.NAME.values
     antenna_mounts = ant_xds.MOUNT.values
@@ -64,7 +67,7 @@ def antenna_info(path):
 
 def data_desc_info(path):
 
-    dd_xds_list = xds_from_storage_table(  # noqa
+    dd_xds_list = xds_from_table_fragment(  # noqa
         path + "::DATA_DESCRIPTION",
         group_cols=["__row__"],
         chunks={"row": 1, "chan": -1}
@@ -76,7 +79,7 @@ def data_desc_info(path):
 
 def feed_info(path):
 
-    feed_xds_list = xds_from_storage_table(
+    feed_xds_list = xds_from_table_fragment(
         path + "::FEED",
         group_cols=["SPECTRAL_WINDOW_ID"],
         chunks={"row": -1}
@@ -106,15 +109,15 @@ def feed_info(path):
 
 def flag_cmd_info(path):
 
-    flag_cmd_xds = xds_from_storage_table(path + "::FLAG_CMD")  # noqa
+    flag_cmd_xds = xds_from_table_fragment(path + "::FLAG_CMD")  # noqa
 
     # Not printing any summary information for this subtable yet - not sure
     # what is relevant.
 
 
 def field_info(path):
 
-    field_xds = xds_from_storage_table(path + "::FIELD")[0]
+    field_xds = xds_from_table_fragment(path + "::FIELD")[0]
 
     field_ids = list(range(field_xds.sizes['row']))
     source_ids = [i for i in field_xds.SOURCE_ID.values]
@@ -148,23 +151,23 @@ def field_info(path):
 
 def history_info(path):
 
-    history_xds = xds_from_storage_table(path + "::HISTORY")[0]  # noqa
+    history_xds = xds_from_table_fragment(path + "::HISTORY")[0]  # noqa
 
     # Not printing any summary information for this subtable yet - not sure
     # what is relevant.
 
 
 def observation_info(path):
 
-    observation_xds = xds_from_storage_table(path + "::OBSERVATION")[0]  # noqa
+    observation_xds = xds_from_table_fragment(path + "::OBSERVATION")[0]  # noqa
 
     # Not printing any summary information for this subtable yet - not sure
     # what is relevant.
 
 
 def polarization_info(path):
 
-    polarization_xds = xds_from_storage_table(path + "::POLARIZATION")[0]
+    polarization_xds = xds_from_table_fragment(path + "::POLARIZATION")[0]
 
     corr_types = polarization_xds.CORR_TYPE.values
 
@@ -182,15 +185,15 @@ def polarization_info(path):
 
 def processor_info(path):
 
-    processor_xds = xds_from_storage_table(path + "::PROCESSOR")[0]  # noqa
+    processor_xds = xds_from_table_fragment(path + "::PROCESSOR")[0]  # noqa
 
     # Not printing any summary information for this subtable yet - not sure
     # what is relevant.
 
 
 def spw_info(path):
 
-    spw_xds_list = xds_from_storage_table(
+    spw_xds_list = xds_from_table_fragment(
         path + "::SPECTRAL_WINDOW",
         group_cols=["__row__"],
         chunks={"row": 1, "chan": -1}
@@ -214,7 +217,7 @@ def spw_info(path):
 
 def state_info(path):
 
-    state_xds = xds_from_storage_table(path + "::STATE")[0]  # noqa
+    state_xds = xds_from_table_fragment(path + "::STATE")[0]  # noqa
 
     # Not printing any summary information for this subtable yet - not sure
     # what is relevant.
@@ -233,7 +236,7 @@ def source_info(path):
 
 def pointing_info(path):
 
-    pointing_xds = xds_from_storage_table(path + "::POINTING")[0]  # noqa
+    pointing_xds = xds_from_table_fragment(path + "::POINTING")[0]  # noqa
 
     # Not printing any summary information for this subtable yet - not sure
     # what is relevant.
@@ -362,7 +365,7 @@ def summary():
     # Open the data, grouping by the usual columns. Use these datasets to
     # produce some useful summaries.
 
-    data_xds_list = xds_from_storage_ms(
+    data_xds_list = xds_from_ms_fragment(
         path,
         index_cols=("TIME",),
         columns=("TIME", "FLAG", "FLAG_ROW", "DATA"),

quartical/calibration/calibrate.py

@@ -1,15 +1,13 @@
 # -*- coding: utf-8 -*-
 import numpy as np
 import dask.array as da
-from quartical.calibration.mapping import make_mapping_datasets
 from quartical.gains.general.generics import (compute_residual,
                                               compute_corrected_residual,
                                               compute_corrected_weights)
 from quartical.calibration.constructor import construct_solver
 from quartical.gains.datasets import (make_gain_xds_lod,
                                       make_net_xds_lod,
                                       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
@@ -114,7 +112,10 @@ def add_calibration_graph(
     stats_xds_list,
     solver_opts,
     chain,
-    output_opts
+    mapping_xds_list,
+    gain_xds_lod,
+    output_opts,
+    dask_opts
 ):
     """Given data graph and options, adds the steps necessary for calibration.
 
@@ -144,24 +145,6 @@ def add_calibration_graph(
             "indicate user error."
         )
 
-    # Create a list of dicts of xarray.Dataset objects which will describe the
-    # gains per data xarray.Dataset.
-    gain_xds_lod = make_gain_xds_lod(data_xds_list, chain)
-
-    # Create a list of datasets containing mappings. TODO: Is this the best
-    # place to do this?
-    mapping_xds_list = make_mapping_datasets(data_xds_list, chain)
-
-    # If there are gains to be loaded from disk, this will load an interpolate
-    # them to be consistent with this calibration run. TODO: This needs to
-    # be substantially improved to handle term specific behaviour/utilize
-    # mappings.
-    gain_xds_lod = load_and_interpolate_gains(
-        gain_xds_lod,
-        chain,
-        output_opts.gain_directory
-    )
-
     # Poplulate the gain xarray.Datasets with solutions and convergence info.
     gain_xds_lod, data_xds_list, stats_xds_list = construct_solver(
         data_xds_list,

quartical/config/__init__.py

@@ -27,8 +27,8 @@
     class _GainSchema(object):
         gain: Dict[str, Parameter]
 
-    # The gain section is loaded explicitly, since we need to form up multiple
-    # instances.
+    # The gain and model sections are loaded explicitly, since we need to form
+    # up multiple instances.
     gain_schema = oc.merge(
         oc.structured(_GainSchema),
         oc.load(f"{dirname}/gain_schema.yaml")
@@ -42,3 +42,21 @@ class _GainSchema(object):
         bases=(BaseConfigSection,),
         post_init=POST_INIT_MAP['gain']
     )
+
+    @dataclass
+    class _ModelComponentSchema(object):
+        model_component: Dict[str, Parameter]
+
+    model_component_schema = oc.merge(
+        oc.structured(_ModelComponentSchema),
+        oc.load(f"{dirname}/model_component_schema.yaml")
+    )
+    model_component_schema = model_component_schema.model_component
+
+    # Create model dataclass.
+    ModelComponent = schema_utils.schema_to_dataclass(
+        model_component_schema,
+        "ModelComponent",
+        bases=(BaseConfigSection,),
+        post_init=POST_INIT_MAP['model_component']
+    )

quartical/config/argument_schema.yaml

@@ -119,6 +119,22 @@ input_model:
       given by the dE tagged clusters in LSM1. Leaving this value unset
       (the default) will use an identity model.
 
+  advanced_recipe:
+    dtype: bool
+    default: false
+    info:
+      Enable advanced recipe specification for use with degridder/more
+      complicated modes of model construction.
+
+  threads:
+    dtype: int
+    default: 1
+    info:
+      Controls the number of threads used internally by the various predict
+      methods. This should typically be set to the same value as
+      solver.threads if solver.threads is in use. Currently only supported
+      for degridding.
+
   beam:
     dtype: Optional[str]
     info:
@@ -177,6 +193,15 @@ output:
       Name of directory in which QuartiCal logging outputs will be stored.
       s3 is not currently supported for these outputs.
 
+  fragment_path:
+    dtype: Optional[str]
+    info:
+      If set, instead of mutating the input by e.g. writing flags, instead
+      writes a fragment to this location. A fragment is a zarr backed data
+      format that is read and dynamically combined with any parent datasets.
+      This allows QuartiCal to operate in an entirely read-only fashion.
+      This option is experimental.
+
   log_to_terminal:
     default: true
     dtype: bool
@@ -392,6 +417,13 @@ solver:
       Specify as the integer index of the antenna - antenna names are not
       currently supported.
 
+  bypass:
+    dtype: bool
+    default: false
+    info:
+      Disable executing the solver graph (eg. if only gain interpolation is
+      required).
+
 dask:
   threads:
     dtype: Optional[int]
@@ -413,7 +445,7 @@ dask:
 
   scheduler:
     default: threads
-    choices: ["threads", "single-threaded", "distributed"]
+    choices: ["processes", "threads", "single-threaded", "distributed"]
     info:
       Which dask scheduler to use. The default, threads, is the most
       appropriate for non-cluster use.

quartical/config/config_classes.py

@@ -112,6 +112,11 @@ def __input_model_post_init__(self):
     self.__validate_element_choices__()
 
 
+def __model_component_post_init__(self):
+    self.__validate_choices__()
+    self.__validate_element_choices__()
+
+
 def __output_post_init__(self):
 
     self.__validate_choices__()
@@ -182,6 +187,7 @@ def __gain_post_init__(self):
 POST_INIT_MAP = {
     "input_ms": __input_ms_post_init__,
     "input_model": __input_model_post_init__,
+    "model_component": __model_component_post_init__,
     "output": __output_post_init__,
     "mad_flags": __mad_flags_post_init__,
     "solver": __solver_post_init__,