EXPERIMENTAL: Add support for fragments and degridding

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/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

@@ -120,6 +120,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:
@@ -183,6 +199,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

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__,

quartical/config/external.py

@@ -2,7 +2,7 @@
 from omegaconf import OmegaConf as oc
 from typing import Dict, Any
 from scabha.cargo import Parameter
-from quartical.config import Gain, BaseConfig, gain_schema
+from quartical.config import Gain, ModelComponent, BaseConfig, gain_schema
 
 
 def finalize_structure(additional_config):
@@ -18,6 +18,21 @@ def finalize_structure(additional_config):
     # Use the default terms if no alternative is specified.
     terms = terms or BaseConfig.solver.terms
 
+    recipe = None
+    models = []  # No components by default.
+
+    # Get last specified version of input_model.recipe.
+    for cfg in additional_config[::-1]:
+        advanced_recipe = oc.select(cfg, "input_model.advanced_recipe")
+        recipe = oc.select(cfg, "input_model.recipe")
+        if recipe is not None and advanced_recipe:
+            ingredients = re.split(r'([\+~:])', recipe)
+            ingredients = [
+                i for i in ingredients if not bool(re.search(r'([\+~:])', i))
+            ]
+            models = list(dict.fromkeys(i.split("@")[0] for i in ingredients))
+            break
+
     FinalConfig = make_dataclass(
         "FinalConfig",
         [(t, Gain, field(default_factory=Gain)) for t in terms],

quartical/config/helper.py

@@ -79,8 +79,15 @@ def help():
     if len(sys.argv) != 1 and not help_arg:
         return
 
-    # Include a generic gain term in the help config.
-    additional_config = [oc.from_dotlist(["solver.terms=['gain']"])]
+    # Include a generic gain term and model component in the help config.
+    additional_config = [
+        oc.from_dotlist(
+            [
+                "input_model.advanced_recipe=model_component",
+                "solver.terms=['gain']"
+            ]
+        )
+    ]
     help_class = finalize_structure(additional_config)
     HelpConfig = help_class()
 

quartical/config/internal.py

@@ -1,5 +1,15 @@
 from daskms.fsspec_store import DaskMSStore
 from quartical.gains import TERM_TYPES
+from quartical.config import ModelComponent
+
+
+def get_component_dict(opts):
+
+    return {
+        k: getattr(opts, k)
+        for k in opts.__dataclass_fields__.keys()
+        if isinstance(getattr(opts, k), ModelComponent)
+    }
 
 
 def gains_to_chain(opts):

quartical/config/model_component_schema.yaml

@@ -0,0 +1,61 @@
+model_component:
+  path_or_name:
+    dtype: str
+    required: true
+    info:
+      Path to model/name of column.
+  type:
+    dtype: str
+    choices:
+      - mds  # Replace with awesome acronym eventually.
+      - tigger-lsm
+      - column
+    default: column
+    info:
+      Type of model component  
+  tags:
+    dtype: Optional[List[str]]
+    info:
+      Tag for use in the predict.
+  # region:  # Add when implemented.
+  #   dtype: Optional[str]
+  #   info:
+  #     Name of region file to use in degridder.
+  npix_x:
+    dtype: Optional[int]
+    info:
+      Image x size in pixels for use in degridding.
+  npix_y:
+    dtype: Optional[int]
+    info:
+      Image y size in pixels for use in degridding.
+  cellsize_x:
+    dtype: Optional[float]
+    info:
+      Pixel x cellsize in radians.
+  cellsize_y:
+    dtype: Optional[float]
+    info:
+      Pixel y cellsize in radians.
+  centre_x:
+    dtype: int
+    default: 0
+    info:
+      x coordinate of central pixel.
+  centre_y:
+    dtype: int
+    default: 0
+    info:
+      y coordinate of central pixel. 
+  integrations_per_image:
+    dtype: int
+    default: 0
+    info:
+      Number of integrations per image to use in degridding. The default
+      (zero) is all times in a chunk.
+  channels_per_image:
+    dtype: int
+    default: 0
+    info:
+      Number of channels per image to use in degridding. The default
+      (zero) is all channels in a chunk.