RCAL-1175: Implement multiband associations. (#2077)

Author: mairanteodoro

changes/2077.associations.rst

@@ -0,0 +1 @@
+Implemented MultibandAssociation class to create multiband catalog association files.

docs/roman/associations/commands.rst

@@ -10,3 +10,4 @@ Association Commands
    skycell_asn.rst
    mk_skycell_list.rst
    mk_skycell_asn_from_skycell_list.rst
+   multiband_asn.rst

docs/roman/associations/multiband_asn.rst

@@ -0,0 +1,70 @@
+
+.. _multiband_asn:
+
+multiband_asn
+=============
+
+Create multiband associations using either the command line tool
+``multiband_asn`` or through the Python API using
+:class:`romancal.associations.multiband_asn.MultibandAssociation`.
+
+Multiband Associations
+^^^^^^^^^^^^^^^^^^^^^^
+
+This module groups input files by their skycell identifier, creating an
+association file for each unique skycell. Each association file contains all
+filters (observations) that were used for that specific skycell, allowing for
+multiband data products to be generated per skycell. This enables efficient
+organization and processing of data across multiple filters for the same region
+of the sky.
+
+To create a multiband association, use the following command:
+
+.. code-block:: bash
+
+    multiband_asn r00001_*full*_coadd.asdf
+
+where the input files are in the current directory and the wildcard expands to
+all relevant files. The tool will group files by their skycell identifier and
+generate association files for each group. To get a complete list of options
+you can run the command with the ``-h`` option:
+
+.. code-block:: bash
+
+    multiband_asn -h
+
+The input filenames should follow the convention:
+
+.. code-block:: text
+
+    rPPPPP_<data_release_id>_<product_type>_<skycell_id>_coadd.asdf
+
+Where:
+    PPPPP = Program number
+    data_release_id = Data release identifier (e.g., 'p' for prompt)
+    product_type = Product type (e.g., 'full')
+    skycell_id = Skycell identifier (e.g., '270p65x48y69')
+
+The association files will be JSON files named:
+
+.. code-block:: text
+
+    rPPPPP_<data_release_id>_<product_type>_<skycell_id>_asn.json
+
+For example, to generate associations for all skycells in a program:
+
+.. code-block:: bash
+
+    multiband_asn r00001_*full*_coadd.asdf
+
+You can also use the Python API:
+
+.. code-block:: python
+
+    from romancal.associations.multiband_asn import MultibandAssociation
+    files = ["r00001_p_full_270p65x48y69_coadd.asdf", ...]
+    multiband = MultibandAssociation(files)
+    multiband.create_multiband_asn()
+
+The data release ID and product type are extracted from the filenames.
+Association files are generated for each skycell group.

pyproject.toml

@@ -91,6 +91,7 @@ roman_set_velocity_aberration = "romancal.scripts.set_velocity_aberration:main"
 roman_static_preview = "romancal.scripts.static_preview:command"
 roman_v1_calculate = "romancal.scripts.v1_calculate:main"
 skycell_asn = "romancal.associations.skycell_asn:_cli"
+multiband_asn = "romancal.associations.multiband_asn:_cli"
 
 [build-system]
 requires = [

romancal/associations/multiband_asn.py

@@ -0,0 +1,115 @@
+import argparse
+import glob
+import logging
+import re
+
+from . import asn_from_list
+
+# Configure logging
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
+logger.setLevel("INFO")
+
+__all__ = ["MultibandAssociation"]
+
+
+class MultibandAssociation:
+    """A class to create multiband associations."""
+
+    def __init__(self, files):
+        self.files = self._parse_file_list(files)
+        self.skycell_groups = self._get_skycell_groups(self.files)
+
+    def _parse_file_list(self, files):
+        """
+        Parse a file list, expanding wildcards if present.
+
+        If the input list contains a single string with wildcard characters
+        ('*' or '?'), expand it to the matching files using glob. Otherwise,
+        return the list as is.
+
+        Parameters
+        ----------
+        files : list of str
+            List of file paths or a single wildcard pattern.
+
+        Returns
+        -------
+        list of str
+            List of file paths, expanded if a wildcard was provided.
+        """
+        if len(files) == 1 and any(char in files[0] for char in ["*", "?"]):
+            return glob.glob(files[0])
+        return files
+
+    def _get_skycell_groups(self, filelist):
+        """
+        Create skycell groups based on the unique skycell identifiers from a list of filenames.
+        Parameters
+        ----------
+        filelist : list of str
+            List of filenames.
+        Returns
+        -------
+        dict
+            Dictionary mapping skycell identifiers to lists of filenames.
+        """
+        pattern = re.compile(
+            r".*_(?P<skycells>[0-9p]*x[0-9]*y[0-9]*)_f[0-9]*_coadd\.asdf$"
+        )
+        groups = {}
+        for filename in filelist:
+            match = pattern.match(filename)
+            if match:
+                key = match.group("skycells")
+                groups.setdefault(key, []).append(filename)
+        return groups
+
+    def create_multiband_asn(self):
+        """
+        Create a multiband association from a list of files.
+
+        Parameters:
+        files (list): List of file paths or pattern to include in the association.
+
+        Returns:
+        dict: The created association.
+        """
+        for skycell_id, filenames in self.skycell_groups.items():
+            # Get prefixes for all combinations of data_release_id and product_type from filenames
+            # (r00001_{data_release_id}_{product_type}_{skycell_id}_asn.json)
+            prefixes = {x.split(skycell_id)[0] for x in filenames}
+            for prefix in prefixes:
+                # Get all files that match this prefix (data_release_id + product_type) and skycell
+                files = [x for x in filenames if x.startswith(f"{prefix}{skycell_id}")]
+                args = [
+                    *files,
+                    "-o",
+                    f"{prefix}{skycell_id}_asn.json",
+                    "--product-name",
+                    f"{skycell_id}",
+                    "--data-release-id",
+                    prefix.split("_")[1],
+                ]
+                asn_from_list._cli(args)
+
+
+def _cli():
+    parser = argparse.ArgumentParser(
+        description="Create a multiband association from a list of files",
+        usage="multiband_asn file1.asdf file2.asdf ... fileN.asdf",
+    )
+    parser.add_argument(
+        "files",
+        type=str,
+        nargs="+",
+        help="List of files to include in the multiband association",
+    )
+
+    args = parser.parse_args()
+
+    multiband_asn = MultibandAssociation(args.files)
+
+    multiband_asn.create_multiband_asn()
+
+    logger.info("Multiband association creation complete.")

romancal/associations/tests/test_multiband_asn.py

@@ -0,0 +1,53 @@
+import os
+
+from romancal.associations import multiband_asn
+
+
+def test_parse_file_list_wildcard(tmp_path):
+    """Test that _parse_file_list expands wildcards and returns correct file list."""
+    # Create dummy files
+    filenames = ["test1.asdf", "test2.asdf"]
+    for fname in filenames:
+        (tmp_path / fname).touch()
+    files = [str(tmp_path / "*.asdf")]
+    assoc = multiband_asn.MultibandAssociation(files)
+    parsed_files = assoc._parse_file_list(files)
+    assert set(os.path.basename(f) for f in parsed_files) == set(filenames)
+
+
+def test_get_skycell_groups():
+    """Test that _get_skycell_groups correctly groups files by skycell id."""
+    files = [
+        "r00001_p_full_270p65x48y69_f123_coadd.asdf",
+        "r00001_p_full_270p65x48y69_f456_coadd.asdf",
+        "r00001_p_full_271p66x49y70_f123_coadd.asdf",
+    ]
+    assoc = multiband_asn.MultibandAssociation(files)
+    groups = assoc._get_skycell_groups(files)
+    assert set(groups.keys()) == {"270p65x48y69", "271p66x49y70"}
+    assert set(groups["270p65x48y69"]) == {
+        "r00001_p_full_270p65x48y69_f123_coadd.asdf",
+        "r00001_p_full_270p65x48y69_f456_coadd.asdf",
+    }
+
+
+def test_parse_file_list_no_wildcard():
+    """Test that _parse_file_list returns the input list if no wildcard is present."""
+    files = ["file1.asdf", "file2.asdf"]
+    assoc = multiband_asn.MultibandAssociation(files)
+    assert assoc._parse_file_list(files) == files
+
+
+def test_get_skycell_groups_empty():
+    """Test that _get_skycell_groups returns empty dict when no skycell patterns match."""
+    files = ["not_a_skycell_file.txt", "another_file.fits"]
+    assoc = multiband_asn.MultibandAssociation(files)
+    groups = assoc._get_skycell_groups(files)
+    assert groups == {}
+
+
+def test_multiband_association_empty_list():
+    """Test MultibandAssociation with an empty file list."""
+    assoc = multiband_asn.MultibandAssociation([])
+    assert assoc.files == []
+    assert assoc.skycell_groups == {}