diff --git a/bin.src/plotDiaSourceLightcurve b/bin.src/plotDiaSourceLightcurve new file mode 100644 index 0000000..f84ac44 --- /dev/null +++ b/bin.src/plotDiaSourceLightcurve @@ -0,0 +1,6 @@ +#! /usr/bin/env python + +from lsst.analysis.ap.plotDiaSourceLightcurve import main + +if __name__ == "__main__": + main() diff --git a/python/lsst/analysis/ap/__init__.py b/python/lsst/analysis/ap/__init__.py index 9d7ab47..7e4ef9f 100644 --- a/python/lsst/analysis/ap/__init__.py +++ b/python/lsst/analysis/ap/__init__.py @@ -22,7 +22,13 @@ from .apdb import * from .apdbCassandra import * from .ppdb import * +from .nb_utils import * from .version import * # Generated by sconsUtils from .plotImageSubtractionCutouts import * -# NOTE: do not import from nb_utils in this file, as it depends on packages -# that are not available in the base environment. +from .plotDiaSourceLightcurve import * +from .apdbReconstruct import * +from .compare import * +from .imageQA import * +from .spatiallySampledMetricsQA import * +from .plotUtils import * +from .taskRuntimes import * diff --git a/python/lsst/analysis/ap/apdb.py b/python/lsst/analysis/ap/apdb.py index 606ebf5..7813d1f 100644 --- a/python/lsst/analysis/ap/apdb.py +++ b/python/lsst/analysis/ap/apdb.py @@ -26,11 +26,88 @@ import abc import contextlib +import os import warnings +from importlib.resources import files +import felis.datamodel import pandas as pd import sqlalchemy +from lsst.pipe.tasks.schemaUtils import column_dtype, readSdmSchemaFile + + +# Integer felis types whose default pandas read path silently coerces NULL +# rows through float64. Drive the dtype from the SDM schema for these types +# only and leave float / bool / string / timestamp columns to pandas' default +# inference so we don't perturb unrelated behavior. +_INT_FELIS_TYPES = frozenset({ + felis.datamodel.DataType.long, + felis.datamodel.DataType.int, + felis.datamodel.DataType.short, + felis.datamodel.DataType.byte, +}) + + +_apdb_schema_cache = None + + +def _apdb_schema(): + """Lazily load the APDB SDM schema (``sdm_schemas/apdb.yaml``) once.""" + global _apdb_schema_cache + if _apdb_schema_cache is None: + path = os.fspath(files("lsst.sdm.schemas").joinpath("apdb.yaml")) + _apdb_schema_cache = readSdmSchemaFile(path) + return _apdb_schema_cache + + +def _schema_int_dtypes(table_name): + """Return a {column_name: pandas-dtype} mapping for integer columns of + one APDB table. + + Returns an empty dict if the table is not in the SDM schema (e.g. an + older fixture with tables that have since been removed). + """ + table_def = _apdb_schema().get(table_name) + if table_def is None: + return {} + return { + cdef.name: column_dtype(cdef.datatype, nullable=cdef.nullable) + for cdef in table_def.columns + if cdef.datatype in _INT_FELIS_TYPES + } + + +def _read_query(connection, query, table_name=None): + """Run ``query`` on ``connection`` and return a DataFrame with correct + integer-column dtypes. + + ``pd.read_sql_query`` represents SQL NULLs as ``NaN`` and so forces any + nullable BIGINT column through ``float64``. + Here we fetch rows directly from the cursor and build each integer + column with the dtype declared by the SDM schema (looked up via + ``schemaUtils.column_dtype``); other columns fall through to pandas' + default inference. + """ + cursor = connection.execute(query) + columns = list(cursor.keys()) + rows = cursor.fetchall() + dtype_map = _schema_int_dtypes(table_name) if table_name else {} + data = {} + for i, col in enumerate(columns): + values = [row[i] for row in rows] + dtype = dtype_map.get(col) + if dtype is not None: + try: + data[col] = pd.array(values, dtype=dtype) + continue + except (TypeError, ValueError): + # Cast impossible (driver returned an unexpected type); + # fall through to pandas inference. + pass + data[col] = values + return pd.DataFrame(data) + class DbQuery(abc.ABC): """Abstract interface for APDB queries. @@ -255,17 +332,7 @@ def connection(self): pass def set_excluded_diaSource_flags(self, flag_list): - """Set flags of diaSources to exclude when loading diaSources. - - Any diaSources with configured flags are not returned - when calling `load_sources_for_object` or `load_sources` - with `exclude_flagged = True`. - - Parameters - ---------- - flag_list : `list` [`str`] - Flag names to exclude. - """ + # Docstring is inherited. for flag in flag_list: if flag not in self._tables["DiaSource"].columns: raise ValueError(f"flag {flag} not included in DiaSource flags") @@ -273,249 +340,233 @@ def set_excluded_diaSource_flags(self, flag_list): self.diaSource_flags_exclude = flag_list def _make_flag_exclusion_query(self, query, table, flag_list): - """Return an SQL where query that excludes sources with chosen flags. + """Attach a where clause excluding sources with any chosen flag set. Parameters ---------- - flag_list : `list` [`str`] - Flag names to exclude. - query : `sqlalchemy.sql.Query` - Query to include the where statement in. + query : `sqlalchemy.sql.Select` + Query to attach the where clause to. table : `sqlalchemy.schema.Table` - Table containing the column to be queried. + Reflected table containing the flag columns. + flag_list : `list` [`str`] + Flag column names to exclude. Returns ------- - query : `sqlalchemy.sql.Query` - Query that selects rows to exclude based on flags. + query : `sqlalchemy.sql.Select` + Query with the flag exclusion clause attached. """ - # Build a query that selects any source with one or more chosen flags, - # and return the opposite (`not_`) of that query. - query = query.where(sqlalchemy.and_(table.columns[flag_col] == False # noqa: E712 - for flag_col in flag_list)) - return query + return query.where(sqlalchemy.and_(table.columns[col] == False # noqa: E712 + for col in flag_list)) - def load_sources_for_object(self, dia_object_id, exclude_flagged=False, limit=100000): - """Load diaSources for a single diaObject. + def _load_table(self, table, *, where=None, exclude_flagged=False, + order_by=(), limit=None, fill_instrument=True): + """Run a parameterized SELECT and return the result as a DataFrame. Parameters ---------- - dia_object_id : `int` - Id of object to load sources for. + table : `sqlalchemy.schema.Table` + Reflected table to query. + where : `sqlalchemy.sql.ClauseElement`, optional + Extra where clause to attach. exclude_flagged : `bool`, optional - Exclude sources that have selected flags set. - Use `set_excluded_diaSource_flags` to configure which flags - are excluded. - limit : `int` - Maximum number of rows to return. + If True, attach the configured flag-exclusion clause. + order_by : `tuple` [`str`], optional + Column names to order by. + limit : `int`, optional + Maximum number of rows to return; None means no limit. + fill_instrument : `bool`, optional + If True, append an ``instrument`` column to the result. Returns ------- - data : `pandas.DataFrame` - A data frame of diaSources for the specified diaObject. + result : `pandas.DataFrame` """ - table = self._tables["DiaSource"] - query = table.select().where(table.columns["diaObjectId"] == dia_object_id) + query = table.select() + if where is not None: + query = query.where(where) if exclude_flagged: query = self._make_flag_exclusion_query(query, table, self.diaSource_flags_exclude) - query = query.order_by(table.columns["visit"], - table.columns["detector"], - table.columns["diaSourceId"]) + if order_by: + query = query.order_by(*[table.columns[c] for c in order_by]) + if limit is not None: + query = query.limit(limit) with self.connection as connection: - result = pd.read_sql_query(query, connection) - - self._fill_from_instrument(result) + result = _read_query(connection, query, table_name=table.name) + if fill_instrument: + self._fill_from_instrument(result) return result - def load_forced_sources_for_object(self, dia_object_id, exclude_flagged=False, limit=100000): - """Load diaForcedSources for a single diaObject. + def _load_one(self, table_name, id_column, id_value, fill_instrument=True): + """Load a single row from a table by id, raising if missing. Parameters ---------- - dia_object_id : `int` - Id of object to load sources for. - exclude_flagged : `bool`, optional - Exclude sources that have selected flags set. - Use `set_excluded_diaSource_flags` to configure which flags - are excluded. - limit : `int` - Maximum number of rows to return. + table_name : `str` + Key into ``self._tables`` for the table to query. + id_column : `str` + Name of the id column to filter on. + id_value : `int` + Id value to match. + fill_instrument : `bool`, optional + If True, append an ``instrument`` column to the result. Returns ------- - data : `pandas.DataFrame` - A data frame of diaSources for the specified diaObject. - """ - table = self._tables["DiaForcedSource"] - query = table.select().where(table.columns["diaObjectId"] == dia_object_id) - if exclude_flagged: - query = self._make_flag_exclusion_query(query, table, self.diaSource_flags_exclude) - query = query.order_by(table.columns["visit"], - table.columns["detector"], - table.columns["diaForcedSourceId"]) - with self.connection as connection: - result = pd.read_sql_query(query, connection) + row : `pandas.Series` - self._fill_from_instrument(result) - return result - - def load_source(self, id): - """Load one diaSource. - - Parameters - ---------- - id : `int` - The diaSourceId to load data for. - - Returns - ------- - data : `pandas.Series` - The requested diaSource. + Raises + ------ + RuntimeError + If no row matches. """ - table = self._tables["DiaSource"] - query = table.select().where(table.columns["diaSourceId"] == id) - with self.connection as connection: - result = pd.read_sql_query(query, connection) + table = self._tables[table_name] + result = self._load_table(table, + where=table.columns[id_column] == id_value, + fill_instrument=fill_instrument) if len(result) == 0: - raise RuntimeError(f"diaSourceId={id} not found in DiaSource table") - - self._fill_from_instrument(result) + raise RuntimeError(f"{id_column}={id_value} not found in {table_name} table") return result.iloc[0] - def load_sources(self, exclude_flagged=False, limit=100000): - """Load diaSources. - - Parameters - ---------- - exclude_flagged : `bool`, optional - Exclude sources that have selected flags set. - Use `set_excluded_diaSource_flags` to configure which flags - are excluded. - limit : `int` - Maximum number of rows to return. - - Returns - ------- - data : `pandas.DataFrame` - All available diaSources. - """ + def load_sources_for_object(self, dia_object_id, exclude_flagged=False, limit=100000): + # Docstring is inherited. table = self._tables["DiaSource"] - query = table.select() - if exclude_flagged: - query = self._make_flag_exclusion_query(query, table, self.diaSource_flags_exclude) - query = query.order_by(table.columns["visit"], - table.columns["detector"], - table.columns["diaSourceId"]) - if limit is not None: - query = query.limit(limit) + return self._load_table( + table, + where=table.columns["diaObjectId"] == dia_object_id, + exclude_flagged=exclude_flagged, + order_by=("visit", "detector", "diaSourceId"), + limit=limit, + ) - with self.connection as connection: - result = pd.read_sql_query(query, connection) - - self._fill_from_instrument(result) - return result - - def load_object(self, id): - """Load the most-recently updated version of one diaObject. + def load_forced_sources_for_object(self, dia_object_id, exclude_flagged=False, limit=100000): + # Docstring is inherited. + table = self._tables["DiaForcedSource"] + return self._load_table( + table, + where=table.columns["diaObjectId"] == dia_object_id, + exclude_flagged=exclude_flagged, + order_by=("visit", "detector", "diaForcedSourceId"), + limit=limit, + ) - Parameters - ---------- - id : `int` - The diaObjectId to load data for. + def load_source(self, id): + # Docstring is inherited. + return self._load_one("DiaSource", "diaSourceId", id) - Returns - ------- - data : `pandas.Series` - The requested object. + def load_sources(self, exclude_flagged=False, limit=100000): + # Docstring is inherited. + return self._load_table( + self._tables["DiaSource"], + exclude_flagged=exclude_flagged, + order_by=("visit", "detector", "diaSourceId"), + limit=limit, + ) + + @staticmethod + def _validity_end_column(table): + """Return the DiaObject "validity end" column. + + sdm_schemas renamed this to ``validityEndMjdTai`` (it's also nullable + double-precision MJD now, not a TIMESTAMP). Older fixtures still have + the original ``validityEnd`` name; this helper prefers the current + name and falls back to the legacy one. """ + for name in ("validityEndMjdTai", "validityEnd"): + if name in table.columns: + return table.columns[name] + raise KeyError("DiaObject has neither validityEndMjdTai nor validityEnd") + + def load_object(self, id): + # Docstring is inherited. table = self._tables["DiaObject"] - query = table.select().where(table.columns["validityEnd"] == None) # noqa: E711 - query = query.where(table.columns["diaObjectId"] == id) - with self.connection as connection: - result = pd.read_sql_query(query, connection) + result = self._load_table( + table, + where=sqlalchemy.and_( + self._validity_end_column(table) == None, # noqa: E711 + table.columns["diaObjectId"] == id, + ), + fill_instrument=False, + ) if len(result) == 0: raise RuntimeError(f"diaObjectId={id} not found in DiaObject table") - return result.iloc[0] def load_objects(self, limit=100000, latest=True): - """Load all diaObjects. - - Parameters - ---------- - limit : `int` - Maximum number of rows to return. - latest : `bool` - Only load diaObjects where validityEnd is None. - These are the most-recently updated diaObjects. - - Returns - ------- - data : `pandas.DataFrame` - All available diaObjects. - """ + # Docstring is inherited. table = self._tables["DiaObject"] - if latest: - query = table.select().where(table.columns["validityEnd"] == None) # noqa: E711 - else: - query = table.select() - query = query.order_by(table.columns["diaObjectId"]) - if limit is not None: - query = query.limit(limit) + where = self._validity_end_column(table) == None if latest else None # noqa: E711 + return self._load_table( + table, + where=where, + order_by=("diaObjectId",), + limit=limit, + fill_instrument=False, + ) - with self.connection as connection: - result = pd.read_sql_query(query, connection) + def load_forced_source(self, id): + # Docstring is inherited. + return self._load_one("DiaForcedSource", "diaForcedSourceId", id) - return result + def load_forced_sources(self, limit=100000): + # Docstring is inherited. + return self._load_table( + self._tables["DiaForcedSource"], + order_by=("visit", "detector", "diaForcedSourceId"), + limit=limit, + ) - def load_forced_source(self, id): - """Load one diaForcedSource. + def iter_sources(self, page_size=100000, reliability_min=None, reliability_max=None): + """Yield DiaSources in pages of ``page_size`` rows. Parameters ---------- - id : `int` - The diaForcedSourceId to load data for. + page_size : `int` + Number of rows per page. + reliability_min, reliability_max : `float`, optional + Inclusive bounds on the reliability column. - Returns - ------- - data : `pandas.Series` - The requested forced source. + Yields + ------ + page : `pandas.DataFrame` + One page of DiaSources, with the ``instrument`` column attached. """ - table = self._tables["DiaForcedSource"] - query = table.select().where(table.columns["diaForcedSourceId"] == id) - with self.connection as connection: - result = pd.read_sql_query(query, connection) - if len(result) == 0: - raise RuntimeError(f"diaForcedSourceId={id} not found in DiaForcedSource table") - - self._fill_from_instrument(result) - return result.iloc[0] - - def load_forced_sources(self, limit=100000): - """Load all diaForcedSources. - - Parameters - ---------- - limit : `int` - Maximum number of rows to return. + table = self._tables["DiaSource"] + clauses = [] + if reliability_min is not None: + clauses.append(table.columns["reliability"] >= reliability_min) + if reliability_max is not None: + clauses.append(table.columns["reliability"] <= reliability_max) + where = sqlalchemy.and_(*clauses) if clauses else None + + offset = 0 + while True: + query = table.select() + if where is not None: + query = query.where(where) + query = query.order_by(table.columns["visit"], + table.columns["detector"], + table.columns["diaSourceId"]) + query = query.limit(page_size).offset(offset) + with self.connection as connection: + page = _read_query(connection, query, table_name=table.name) + if len(page) == 0: + break + self._fill_from_instrument(page) + yield page + offset += page_size + + def count_sources(self): + """Return the total number of DiaSources in the database. Returns ------- - data : `pandas.DataFrame` - All available diaForcedSources. + count : `int` """ - table = self._tables["DiaForcedSource"] - query = table.select() - query = query.order_by(table.columns["visit"], - table.columns["detector"], - table.columns["diaForcedSourceId"]) - if limit is not None: - query = query.limit(limit) - + table = self._tables["DiaSource"] + query = sqlalchemy.select(sqlalchemy.func.count()).select_from(table) with self.connection as connection: - result = pd.read_sql_query(query, connection) - self._fill_from_instrument(result) - return result + return connection.execute(query).scalar() def _fill_from_instrument(self, diaSources): """Add an instrument column to a list of sources. diff --git a/python/lsst/analysis/ap/apdbReconstruct.py b/python/lsst/analysis/ap/apdbReconstruct.py new file mode 100644 index 0000000..5980c1c --- /dev/null +++ b/python/lsst/analysis/ap/apdbReconstruct.py @@ -0,0 +1,358 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Reconstruct APDB-shaped catalogs from DiaPipelineTask butler outputs. + +`lsst.ap.association.diaPipe.DiaPipelineTask` writes APDB-bound catalogs as +side-effects of its quantum execution: per-(visit, detector) DiaSource, +DiaObject, and DiaForcedSource datasets that mirror what gets inserted +into the APDB. `ApdbReconstructor` walks those datasets in a butler +collection and concatenates them into single DataFrames matching the APDB +SDM schema. + +By default the reconstructor uses the dataset names configured by the +production AP pipeline (``ap_pipe/pipelines/_ingredients/ApPipe.yaml``'s +``associateApdb`` task: ``dia_source_apdb``, ``dia_object_apdb``, +``dia_forced_source_apdb``). For runs that use the raw +``DiaPipelineConnections`` defaults instead, pass ``dataset_names`` +explicitly. +""" + +from __future__ import annotations + +__all__ = ["ApdbReconstruction", "ApdbReconstructor"] + +import dataclasses +import logging + +import pandas as pd + +from lsst.pipe.tasks.schemaUtils import convertDataFrameToSdmSchema + +from .apdb import DbQuery, _apdb_schema + +_log = logging.getLogger(__name__) + + +@dataclasses.dataclass +class ApdbReconstruction: + """APDB-shaped DataFrames reconstructed from DiaPipelineTask outputs. + + Attributes + ---------- + diaSources : `pandas.DataFrame` + DiaSource rows after association and standardization, deduped on + ``diaSourceId``. + diaObjects : `pandas.DataFrame` + DiaObject rows. By default deduped to the latest snapshot per + ``diaObjectId`` (``history=False`` in `ApdbReconstructor.reconstruct`). + diaForcedSources : `pandas.DataFrame` + DiaForcedSource rows, deduped on the schema primary key + ``(diaObjectId, visit, detector)``. + """ + diaSources: pd.DataFrame + diaObjects: pd.DataFrame + diaForcedSources: pd.DataFrame + + +class ApdbReconstructor: + """Reconstruct APDB-shaped catalogs from `DiaPipelineTask` butler outputs. + + Parameters + ---------- + butler : `lsst.daf.butler.Butler` + Butler initialized with the collections that hold the pipeline run. + dataset_names : `dict` [`str`, `list` [`str`]], optional + Override the default dataset types used for each table. Keys are + ``"diaSource"``, ``"diaObject"``, ``"diaForcedSource"``; each + value is a list of dataset-type names that get concatenated and + deduplicated. Defaults to ``DEFAULT_DATASET_NAMES``, which includes the + ``preloaded_*`` datasets so the reconstruction includes both "history" + rows from preload and the new ones. + collections : `list` [`str`], optional + Specify the butler collections to query if not using the default + configured for the butler. + where : `str`, optional + Butler ``where`` clause passed through to ``queryDatasets`` + (e.g. ``"instrument='LSSTComCam' AND visit > 1000"``). + """ + + DEFAULT_DATASET_NAMES = { + "diaSource": ["dia_source_apdb", "preloaded_dia_source"], + "diaObject": ["dia_object_apdb", "preloaded_dia_object"], + "diaForcedSource": ["dia_forced_source_apdb", "preloaded_dia_forced_source"], + } + + def __init__(self, butler, dataset_names=None, *, + collections=None, where=None): + self.butler = butler + self.dataset_names = (dataset_names if dataset_names is not None + else self.DEFAULT_DATASET_NAMES) + self.collections = collections + self.where = where + self.log = _log + + def _query_kwargs(self): + kwargs = {"findFirst": True} + if self.collections is not None: + kwargs["collections"] = self.collections + if self.where is not None: + kwargs["where"] = self.where + return kwargs + + def _load_tables(self, dataset_names): + """Load and concatenate one or more dataset types into a single + DataFrame. Returns an empty DataFrame if no dataset is present. + + Each dataset is loaded independently via `_load_table` and the + results are concatenated. Dedup happens later in `finalize`, so + overlap between sources (e.g. the same diaSource appearing in + both ``dia_source_apdb`` and ``preloaded_dia_source`` after a + prior pipeline run) is harmless. + """ + frames = [self._load_table(name) for name in dataset_names] + frames = [f for f in frames if len(f)] + if not frames: + return pd.DataFrame() + if len(frames) == 1: + return frames[0] + return pd.concat(frames, ignore_index=True) + + def _load_table(self, dataset_name): + """Load every instance of ``dataset_name`` from the butler and + concatenate into a single DataFrame. Returns an empty DataFrame if + the dataset type doesn't exist or no refs match. + """ + try: + refs = list(self.butler.registry.queryDatasets( + dataset_name, **self._query_kwargs())) + except Exception as e: + self.log.info("Skipping %s: query failed (%s)", + dataset_name, e) + return pd.DataFrame() + if not refs: + return pd.DataFrame() + frames = [self.butler.get(ref, storageClass="DataFrame") for ref in refs] + return pd.concat(frames, ignore_index=True) + + def reconstruct(self, *, coerce_to_schema=True, history=False): + """Load all per-quantum catalogs and return APDB-shaped DataFrames. + + Parameters + ---------- + coerce_to_schema : `bool`, optional + If True (default), coerce each output to the SDM ``apdb.yaml`` + schema. + history : `bool`, optional + If True, keep every diaObject row written across all quanta. + If False (default), dedupe to the atest row per ``diaObjectId``, + mirroring the "current APDB state" view that ``DiaObjectLast`` + would give. + + Returns + ------- + result : `ApdbReconstruction` + """ + diaSources = self._load_tables(self.dataset_names["diaSource"]) + diaObjects = self._load_tables(self.dataset_names["diaObject"]) + diaForcedSources = self._load_tables(self.dataset_names["diaForcedSource"]) + return self.finalize(diaSources, diaObjects, diaForcedSources, + coerce_to_schema=coerce_to_schema, + history=history) + + @staticmethod + def finalize(diaSources, diaObjects, diaForcedSources, *, + coerce_to_schema=True, history=False): + """Dedup and (optionally) schema-coerce already-loaded catalogs. + + Parameters + ---------- + diaSources, diaObjects, diaForcedSources : `pandas.DataFrame` + Concatenated per-quantum catalogs. + coerce_to_schema, history : see `reconstruct`. + + Returns + ------- + result : `ApdbReconstruction` + """ + # DiaSource: Primary Key (PK) is diaSourceId. + if len(diaSources) and "diaSourceId" in diaSources.columns: + diaSources = diaSources.drop_duplicates(subset="diaSourceId", + keep="last") + # DiaForcedSource: PK is (diaObjectId, visit, detector). + fkey = ["diaObjectId", "visit", "detector"] + if (len(diaForcedSources) + and set(fkey).issubset(diaForcedSources.columns)): + diaForcedSources = diaForcedSources.drop_duplicates( + subset=fkey, keep="last") + # DiaObject: each quantum that touches a diaObject emits a row for + # it. Many of those rows are "passthrough" snapshots: the diaObject + # was in the quantum's preloaded working set but wasn't actually + # updated, and the writer leaves ``nDiaSources`` (and other + # update-only fields) as NaN/NULL. The validity timestamps on those + # passthrough rows still advance to the quantum's processing time, + # so a naive "sort by validityStart, keep last" dedup picks them + # over the older snapshot that carries the real count. + # + # Fix: include ``nDiaSources`` as the primary sort key with NaN at + # the front, so dedup ``keep="last"`` prefers any informative + # snapshot (highest ``nDiaSources``, ties broken by latest validity) + # over a passthrough one. Falls back to validity-only sort when + # ``nDiaSources`` is absent. + if (len(diaObjects) and "diaObjectId" in diaObjects.columns + and not history): + validity_col = next( + (c for c in ("validityStartMjdTai", "validityStart") + if c in diaObjects.columns), None) + sort_keys = [] + if "nDiaSources" in diaObjects.columns: + sort_keys.append("nDiaSources") + if validity_col is not None: + sort_keys.append(validity_col) + if sort_keys: + diaObjects = diaObjects.sort_values(sort_keys, + na_position="first") + diaObjects = diaObjects.drop_duplicates(subset="diaObjectId", + keep="last") + + if coerce_to_schema: + schema = _apdb_schema() + if len(diaSources): + diaSources = convertDataFrameToSdmSchema( + schema, diaSources, "DiaSource", skipIndex=True) + if len(diaObjects): + diaObjects = convertDataFrameToSdmSchema( + schema, diaObjects, "DiaObject", skipIndex=True) + if len(diaForcedSources): + diaForcedSources = convertDataFrameToSdmSchema( + schema, diaForcedSources, "DiaForcedSource", + skipIndex=True) + + return ApdbReconstruction( + diaSources=diaSources.reset_index(drop=True), + diaObjects=diaObjects.reset_index(drop=True), + diaForcedSources=diaForcedSources.reset_index(drop=True), + ) + + def to_query(self, *, coerce_to_schema=True, history=False): + """Reconstruct and wrap the result as a `DbQuery`-compatible adapter + so the in-memory frames can be passed to `lightcurve`, + `PlotDiaSourceLightcurveTask`, and other tools that expect the + ``DbQuery`` interface. + + Returns + ------- + query : `InMemoryDbQuery` + """ + recon = self.reconstruct(coerce_to_schema=coerce_to_schema, + history=history) + return InMemoryDbQuery(recon.diaSources, + recon.diaObjects, + recon.diaForcedSources) + + +class InMemoryDbQuery(DbQuery): + """`DbQuery` backed by in-memory DataFrames (as from `ApdbReconstructor`). + + Implements the same load_* methods as `ApdbSqliteQuery`/`ApdbPostgresQuery` + so the reconstructed data can be passed directly to ``lightcurve`` and + ``PlotDiaSourceLightcurveTask``. + """ + + def __init__(self, diaSources, diaObjects, diaForcedSources): + self._diaSources = diaSources + self._diaObjects = diaObjects + self._diaForcedSources = diaForcedSources + self.diaSource_flags_exclude = [] + + def set_excluded_diaSource_flags(self, flag_list): + # Docstring inherited. + missing = [f for f in flag_list if f not in self._diaSources.columns] + if missing: + raise ValueError( + f"flag(s) {missing} not present in reconstructed DiaSource columns") + self.diaSource_flags_exclude = list(flag_list) + + def _apply_flag_exclusion(self, df): + if not self.diaSource_flags_exclude: + return df + mask = pd.Series(False, index=df.index) + for flag in self.diaSource_flags_exclude: + if flag in df.columns: + mask |= df[flag].fillna(False).astype(bool) + return df[~mask] + + def load_sources_for_object(self, dia_object_id, exclude_flagged=False, + limit=100000): + # Docstring inherited. + df = self._diaSources + result = df[df["diaObjectId"] == dia_object_id] + if exclude_flagged: + result = self._apply_flag_exclusion(result) + return result.head(limit).reset_index(drop=True) + + def load_forced_sources_for_object(self, dia_object_id, + exclude_flagged=False, limit=100000): + # Docstring inherited. + df = self._diaForcedSources + result = df[df["diaObjectId"] == dia_object_id] + return result.head(limit).reset_index(drop=True) + + def load_source(self, id): + # Docstring inherited. + match = self._diaSources[self._diaSources["diaSourceId"] == id] + if len(match) == 0: + raise RuntimeError(f"diaSourceId={id} not found in DiaSource table") + return match.iloc[0] + + def load_sources(self, exclude_flagged=False, limit=100000): + # Docstring inherited. + df = self._diaSources + if exclude_flagged: + df = self._apply_flag_exclusion(df) + return df.head(limit).reset_index(drop=True) + + def load_object(self, id): + # Docstring inherited. + match = self._diaObjects[self._diaObjects["diaObjectId"] == id] + if len(match) == 0: + raise RuntimeError(f"diaObjectId={id} not found in DiaObject table") + return match.iloc[0] + + def load_objects(self, limit=100000, latest=True): + # Docstring inherited. + return self._diaObjects.head(limit).reset_index(drop=True) + + def load_forced_source(self, id): + # Docstring inherited.. + if "diaForcedSourceId" not in self._diaForcedSources.columns: + raise RuntimeError("Reconstructed DiaForcedSource has no " + "diaForcedSourceId column") + match = self._diaForcedSources[ + self._diaForcedSources["diaForcedSourceId"] == id] + if len(match) == 0: + raise RuntimeError( + f"diaForcedSourceId={id} not found in DiaForcedSource table") + return match.iloc[0] + + def load_forced_sources(self, limit=100000): + # Docstring inherited. + return self._diaForcedSources.head(limit).reset_index(drop=True) diff --git a/python/lsst/analysis/ap/compare.py b/python/lsst/analysis/ap/compare.py new file mode 100644 index 0000000..470a66d --- /dev/null +++ b/python/lsst/analysis/ap/compare.py @@ -0,0 +1,320 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Catalog cross-matching and pair-wise comparison utilities for +DiaSource-like tables. +""" + +__all__ = ["match_catalogs", "flux_residuals", "match_to_truth"] + +import numpy as np +import pandas as pd + +import astropy.units as u +from astropy.coordinates import SkyCoord, match_coordinates_sky + + +def _as_arcsec(radius): + """Coerce a number-or-Quantity radius to a float in arcseconds.""" + if isinstance(radius, u.Quantity): + return float(radius.to_value(u.arcsec)) + return float(radius) + + +def _group_keys(srcs1, srcs2, on): + """Yield the union of distinct ``on``-tuples present in either frame.""" + if not on: + yield None + return + s1 = set(map(tuple, srcs1[list(on)].itertuples(index=False, name=None))) + s2 = set(map(tuple, srcs2[list(on)].itertuples(index=False, name=None))) + for key in sorted(s1 | s2): + yield key + + +def _select_group(df, on, key): + """Return the slice of ``df`` whose ``on`` columns equal ``key``.""" + if key is None or not on: + return df + mask = pd.Series(True, index=df.index) + for col, val in zip(on, key): + mask &= df[col] == val + return df[mask] + + +def match_catalogs(srcs1, srcs2, radius=0.5*u.arcsec, on=("visit", "detector"), + ra_col="ra", dec_col="dec", id_col="diaSourceId"): + """Spatially cross-match two DiaSource-like DataFrames. + + Sources are first partitioned by the columns in + ``on`` (e.g. matched only within the same (visit, detector)), then matched + via nearest-neighbor on the sphere. + + Parameters + ---------- + srcs1, srcs2 : `pandas.DataFrame` + Source tables. Each must contain ``ra_col``, ``dec_col``, ``id_col``, + and every column listed in ``on``. + radius : `astropy.units.Quantity` or `float` + Maximum separation for a pair to count as matched. A bare float is + interpreted as arcseconds. + on : `tuple` [`str`] + Columns to group on before matching. Pass an empty tuple to match the + full catalog with no grouping. + ra_col, dec_col : `str` + Column names for sky coordinates, in degrees. + id_col : `str` + Column name for the source id in both catalogs. + + Returns + ------- + matched : `pandas.DataFrame` + Rows from ``srcs1`` that found a partner in ``srcs2`` within + ``radius``. Two columns are added: ``_2`` with the partner's + id, and ``xmatch_dist_arcsec`` with the on-sky separation in arcsec. + unique1 : `pandas.DataFrame` + Rows from ``srcs1`` with no partner. + unique2 : `pandas.DataFrame` + Rows from ``srcs2`` not pointed at by any matched pair. + """ + rad_arcsec = _as_arcsec(radius) + on = tuple(on) + id2_col = f"{id_col}_2" + + matched_chunks = [] + unique1_chunks = [] + unique2_chunks = [] + + for key in _group_keys(srcs1, srcs2, on): + gs1 = _select_group(srcs1, on, key).copy() + gs2 = _select_group(srcs2, on, key).copy() + + if len(gs1) == 0: + unique2_chunks.append(gs2) + continue + if len(gs2) == 0: + unique1_chunks.append(gs1) + continue + + coords1 = SkyCoord(ra=gs1[ra_col].values*u.deg, + dec=gs1[dec_col].values*u.deg) + coords2 = SkyCoord(ra=gs2[ra_col].values*u.deg, + dec=gs2[dec_col].values*u.deg) + idx, sep, _ = match_coordinates_sky(coords1, coords2) + + gs1["xmatch_dist_arcsec"] = sep.to_value(u.arcsec) + gs1[id2_col] = gs2[id_col].values[idx] + + has_match = gs1["xmatch_dist_arcsec"] <= rad_arcsec + m = gs1[has_match] + u1 = gs1[~has_match].drop(columns=["xmatch_dist_arcsec", id2_col]) + u2 = gs2[~gs2[id_col].isin(set(m[id2_col]))] + + matched_chunks.append(m) + unique1_chunks.append(u1) + unique2_chunks.append(u2) + + matched = (pd.concat(matched_chunks) if matched_chunks + else srcs1.iloc[0:0].assign(**{"xmatch_dist_arcsec": np.nan, + id2_col: pd.NA})) + unique1 = pd.concat(unique1_chunks) if unique1_chunks else srcs1.iloc[0:0].copy() + unique2 = pd.concat(unique2_chunks) if unique2_chunks else srcs2.iloc[0:0].copy() + return matched, unique1, unique2 + + +def flux_residuals(matched, srcs2, flux_col="psfFlux", err_col="psfFluxErr", + id_col="diaSourceId", plot=False): + """Compute per-pair flux residuals from a `match_catalogs` result. + + Parameters + ---------- + matched : `pandas.DataFrame` + Output of `match_catalogs` (rows from catalog 1 with partner ids). + srcs2 : `pandas.DataFrame` + Catalog 2, indexed implicitly by ``id_col`` for partner lookup. + flux_col, err_col : `str` + Column names for the flux and its error in both catalogs. + id_col : `str` + Source-id column in both catalogs. ``matched`` is assumed to have + ``f"{id_col}_2"`` populated by `match_catalogs`. + plot : `bool` + If True, return a histogram + Q-Q plot in addition to the residuals. + + Returns + ------- + residuals : `pandas.DataFrame` + One row per pair with columns ``flux1``, ``flux2``, ``err1``, ``err2``, + ``delta_flux``, ``delta_flux_sigma``, and ``xmatch_dist_arcsec``. + fig : `matplotlib.figure.Figure`, optional + Only returned when ``plot=True``. + """ + id2_col = f"{id_col}_2" + s2 = srcs2.set_index(id_col) + partner_ids = matched[id2_col].values + + f1 = matched[flux_col].to_numpy(dtype=float, copy=True) + e1 = matched[err_col].to_numpy(dtype=float, copy=True) + f2 = s2.loc[partner_ids, flux_col].to_numpy(dtype=float, copy=True) + e2 = s2.loc[partner_ids, err_col].to_numpy(dtype=float, copy=True) + + delta = f1 - f2 + sigma = np.sqrt(e1**2 + e2**2) + with np.errstate(divide="ignore", invalid="ignore"): + delta_sigma = np.where(sigma > 0, delta / sigma, np.nan) + + residuals = pd.DataFrame({ + id_col: matched[id_col].values, + id2_col: partner_ids, + "flux1": f1, + "flux2": f2, + "err1": e1, + "err2": e2, + "delta_flux": delta, + "delta_flux_sigma": delta_sigma, + "xmatch_dist_arcsec": matched["xmatch_dist_arcsec"].values, + }) + + if not plot: + return residuals + return residuals, _plot_flux_residuals(residuals, flux_col) + + +def _plot_flux_residuals(residuals, flux_col): + """Histogram + Q-Q plot of normalized flux residuals.""" + import matplotlib.pyplot as plt + sigma = residuals["delta_flux_sigma"].dropna().values + if len(sigma) == 0: + fig, ax = plt.subplots() + ax.text(0.5, 0.5, "no finite residuals", ha="center", va="center") + return fig + + fig, axes = plt.subplots(1, 2, figsize=(10, 4)) + axes[0].hist(sigma, bins=50, color="C0") + axes[0].axvline(0, color="grey", lw=0.5) + axes[0].set_xlabel(rf"$\Delta {flux_col} / \sigma$") + axes[0].set_ylabel("count") + axes[0].set_title(f"N={len(sigma)}, " + f"med={np.median(sigma):.3f}, " + f"MAD={np.median(np.abs(sigma - np.median(sigma))):.3f}") + + # Quick Q-Q vs standard normal without depending on scipy. + sp = np.sort(sigma) + # Inverse-CDF approximation for the standard normal via erfinv. + quantiles = (np.arange(len(sp)) + 0.5) / len(sp) + expected = np.sqrt(2) * _erfinv(2*quantiles - 1) + axes[1].plot(expected, sp, ".", ms=2) + lim = max(abs(expected[0]), abs(expected[-1]), 3.0) + axes[1].plot([-lim, lim], [-lim, lim], "k--", lw=0.5) + axes[1].set_xlabel("expected (N(0,1))") + axes[1].set_ylabel("observed") + axes[1].set_title("Q-Q vs standard normal") + fig.tight_layout() + return fig + + +def _erfinv(y): + """Approximate inverse error function (vectorized). + + Uses the formula from Winitzki (2008); accurate to ~4e-3 across the + domain, which is plenty for plotting Q-Q lines. + """ + a = 0.147 + sign = np.sign(y) + ln1 = np.log(np.clip(1 - y*y, 1e-300, 1.0)) + term = 2/(np.pi*a) + ln1/2 + return sign * np.sqrt(np.sqrt(term*term - ln1/a) - term) + + +def match_to_truth(srcs, truth, radius=0.5*u.arcsec, + src_ra="ra", src_dec="dec", src_id="diaSourceId", + truth_ra="ra", truth_dec="dec", truth_id="injection_id"): + """Match a detected catalog against a truth/injection catalog. + + Performs the match in both directions to compute purity (fraction of + detections that correspond to a real injected source) and completeness + (fraction of injected sources recovered). + + Parameters + ---------- + srcs : `pandas.DataFrame` + Detected sources. + truth : `pandas.DataFrame` + Truth catalog, e.g. an injection catalog. + radius : `astropy.units.Quantity` or `float` + Maximum separation for a match. + src_ra, src_dec, src_id : `str` + Column names in ``srcs``. + truth_ra, truth_dec, truth_id : `str` + Column names in ``truth``. + + Returns + ------- + out : `dict` + With keys: + + - ``"srcs"``: a copy of ``srcs`` with three columns appended: + ``"is_real"`` (bool), ``"truth_dist_arcsec"`` (float), and + ``f"{truth_id}_match"`` (matched truth id, NA if no match). + - ``"truth"``: a copy of ``truth`` with three columns appended: + ``"detected"``, ``"detection_dist_arcsec"``, and + ``f"{src_id}_match"``. + - ``"purity"``: float, fraction of ``srcs`` rows with ``is_real``. + - ``"completeness"``: float, fraction of ``truth`` rows with + ``detected``. + """ + rad_arcsec = _as_arcsec(radius) + srcs_out = srcs.copy() + truth_out = truth.copy() + + truth_id_match = f"{truth_id}_match" + src_id_match = f"{src_id}_match" + + if len(srcs_out) == 0 or len(truth_out) == 0: + srcs_out["is_real"] = False + srcs_out["truth_dist_arcsec"] = np.nan + srcs_out[truth_id_match] = pd.NA + truth_out["detected"] = False + truth_out["detection_dist_arcsec"] = np.nan + truth_out[src_id_match] = pd.NA + return {"srcs": srcs_out, "truth": truth_out, + "purity": 0.0, "completeness": 0.0} + + sc_src = SkyCoord(srcs_out[src_ra].values*u.deg, + srcs_out[src_dec].values*u.deg) + sc_tru = SkyCoord(truth_out[truth_ra].values*u.deg, + truth_out[truth_dec].values*u.deg) + + idx_to_truth, sep_to_truth, _ = match_coordinates_sky(sc_src, sc_tru) + srcs_out["truth_dist_arcsec"] = sep_to_truth.to_value(u.arcsec) + srcs_out[truth_id_match] = truth_out[truth_id].values[idx_to_truth] + srcs_out["is_real"] = srcs_out["truth_dist_arcsec"] <= rad_arcsec + srcs_out.loc[~srcs_out["is_real"], truth_id_match] = pd.NA + + idx_to_src, sep_to_src, _ = match_coordinates_sky(sc_tru, sc_src) + truth_out["detection_dist_arcsec"] = sep_to_src.to_value(u.arcsec) + truth_out[src_id_match] = srcs_out[src_id].values[idx_to_src] + truth_out["detected"] = truth_out["detection_dist_arcsec"] <= rad_arcsec + truth_out.loc[~truth_out["detected"], src_id_match] = pd.NA + + purity = float(srcs_out["is_real"].sum() / len(srcs_out)) + completeness = float(truth_out["detected"].sum() / len(truth_out)) + return {"srcs": srcs_out, "truth": truth_out, + "purity": purity, "completeness": completeness} diff --git a/python/lsst/analysis/ap/imageQA.py b/python/lsst/analysis/ap/imageQA.py new file mode 100644 index 0000000..15d0001 --- /dev/null +++ b/python/lsst/analysis/ap/imageQA.py @@ -0,0 +1,242 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Pixel-level QA for AP difference images. + +Two utilities live here: + +- `image_diff_stats` reports basic statistics (median, MAD, stdev, kurtosis, + mask-plane fractions) of a single difference image. +- `pixel_compare` aligns two difference images for the same data id and + returns their pixel-wise difference, ratio, mask XOR, and a small summary. + +Both are intended for human-driven analysis from a notebook. The first is +also designed to be mapped over a list of dataIds to produce a per-detector +quality DataFrame. +""" + +from __future__ import annotations + +__all__ = ["image_diff_stats", "image_diff_stats_table", "pixel_compare"] + +from dataclasses import dataclass +from typing import Iterable, Mapping + +import numpy as np +import pandas as pd + + +def _finite_view(array): + """Return only the finite entries of ``array`` as a flat 1-D view.""" + flat = np.asarray(array).ravel() + return flat[np.isfinite(flat)] + + +def _basic_stats(array): + """Median, MAD, stdev, and kurtosis of the finite entries of ``array``. + + Kurtosis is the Pearson definition (``E[((x-mu)/sigma)**4] - 3``); zero + for a Gaussian. Computed without scipy so this module has no extra deps. + """ + finite = _finite_view(array) + if finite.size == 0: + return {"median": np.nan, "mad": np.nan, "stddev": np.nan, "kurtosis": np.nan} + median = float(np.median(finite)) + mad = float(np.median(np.abs(finite - median))) + stddev = float(np.std(finite)) + if stddev > 0: + z = (finite - finite.mean()) / stddev + kurtosis = float(np.mean(z**4) - 3.0) + else: + kurtosis = np.nan + return {"median": median, "mad": mad, "stddev": stddev, "kurtosis": kurtosis} + + +def _mask_plane_fractions(mask): + """Return ``{plane_name: fraction_of_pixels_set}`` for an afw mask. + + Parameters + ---------- + mask : `lsst.afw.image.Mask` + The mask whose planes are to be summed. + + Returns + ------- + fractions : `dict` [`str`, `float`] + Per-plane fraction of pixels where that plane bit is set. + """ + arr = mask.array + npix = arr.size + if npix == 0: + return {} + out = {} + for plane_name, plane_bit in mask.getMaskPlaneDict().items(): + bitmask = np.uint64(1) << np.uint64(plane_bit) + out[f"frac_{plane_name}"] = float(np.sum((arr & bitmask) != 0) / npix) + return out + + +def image_diff_stats(butler, visit, detector, dataset_name="difference_image"): + """Summary statistics on a single difference image. + + Parameters + ---------- + butler : `lsst.daf.butler.Butler` + Butler initialized with the relevant collections. + visit, detector : `int` + Data id selecting the exposure to load. + dataset_name : `str`, optional + Butler dataset type name to load. Defaults to ``"difference_image"``. + + Returns + ------- + stats : `dict` + ``{"visit", "detector", "median", "mad", "stddev", "kurtosis"}`` plus + a ``frac_`` entry per mask plane on the exposure. + """ + diff = butler.get(dataset_name, {"visit": visit, "detector": detector}) + stats = {"visit": visit, "detector": detector} + stats.update(_basic_stats(diff.image.array)) + stats.update(_mask_plane_fractions(diff.mask)) + return stats + + +def image_diff_stats_table(butler, + data_ids: Iterable[Mapping[str, int]], + dataset_name="difference_image"): + """Run `image_diff_stats` over many data ids and return a DataFrame. + + Parameters + ---------- + butler : `lsst.daf.butler.Butler` + data_ids : iterable of mapping + Each mapping must contain at least ``visit`` and ``detector`` keys. + dataset_name : `str`, optional + See `image_diff_stats`. + + Returns + ------- + table : `pandas.DataFrame` + One row per dataId, indexed by ``(visit, detector)``. Mask-plane + columns may be missing on some rows if the corresponding plane is not + defined on every exposure; pandas fills those with NaN. + """ + rows = [] + for data_id in data_ids: + rows.append(image_diff_stats(butler, + data_id["visit"], + data_id["detector"], + dataset_name=dataset_name)) + if not rows: + return pd.DataFrame() + return pd.DataFrame(rows).set_index(["visit", "detector"]).sort_index() + + +@dataclass +class PixelCompareResult: + """Container for `pixel_compare` outputs. + + Attributes + ---------- + img1, img2 : `lsst.afw.image.Exposure` + The two loaded exposures, in case the caller wants their masks/WCS. + diff : `numpy.ndarray` + Pixel array of ``img1 - img2``. + ratio : `numpy.ndarray` + Pixel array of ``img1 / img2``; NaN where ``img2 == 0``. + mask_diff : `numpy.ndarray` + XOR of the two mask arrays: nonzero pixels are where the mask planes + differ between the two images. + summary : `dict` + ``{"visit", "detector", "diff_median", "diff_mad", "diff_stddev", + "n_mask_pixels_changed", "frac_mask_pixels_changed"}``. + """ + img1: object + img2: object + diff: np.ndarray + ratio: np.ndarray + mask_diff: np.ndarray + summary: dict + + +def pixel_compare(butler1, butler2, visit, detector, + dataset_name="difference_image"): + """Compare two difference images for the same data id, pixel by pixel. + + Pull the same dataId from two butlers (or two collections of one butler) + and inspect the residuals. + + Parameters + ---------- + butler1, butler2 : `lsst.daf.butler.Butler` + Two butlers, possibly the same instance with different collections. + visit, detector : `int` + Data id to load from both butlers. + dataset_name : `str`, optional + Butler dataset type name to load from each butler. Defaults to + ``"difference_image"``. + + Returns + ------- + result : `PixelCompareResult` + + Raises + ------ + ValueError + If the two loaded images differ in shape (cannot be aligned by simple + subtraction). + """ + img1 = butler1.get(dataset_name, {"visit": visit, "detector": detector}) + img2 = butler2.get(dataset_name, {"visit": visit, "detector": detector}) + + a1 = img1.image.array + a2 = img2.image.array + if a1.shape != a2.shape: + raise ValueError(f"Image shapes differ for visit={visit} detector={detector}: " + f"{a1.shape} vs {a2.shape}") + + diff = a1 - a2 + with np.errstate(divide="ignore", invalid="ignore"): + ratio = np.where(a2 != 0, a1 / a2, np.nan) + + mask_diff = img1.mask.array ^ img2.mask.array + + finite = _finite_view(diff) + if finite.size: + diff_median = float(np.median(finite)) + diff_mad = float(np.median(np.abs(finite - diff_median))) + diff_stddev = float(np.std(finite)) + else: + diff_median = diff_mad = diff_stddev = np.nan + + n_changed = int(np.count_nonzero(mask_diff)) + summary = { + "visit": visit, + "detector": detector, + "diff_median": diff_median, + "diff_mad": diff_mad, + "diff_stddev": diff_stddev, + "n_mask_pixels_changed": n_changed, + "frac_mask_pixels_changed": (float(n_changed / mask_diff.size) + if mask_diff.size else 0.0), + } + return PixelCompareResult(img1=img1, img2=img2, diff=diff, ratio=ratio, + mask_diff=mask_diff, summary=summary) diff --git a/python/lsst/analysis/ap/legacyPlotUtils.py b/python/lsst/analysis/ap/legacyPlotUtils.py index e902a86..e15b450 100644 --- a/python/lsst/analysis/ap/legacyPlotUtils.py +++ b/python/lsst/analysis/ap/legacyPlotUtils.py @@ -19,6 +19,8 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . +__all__ = [] + import numpy as np import matplotlib as mpl import matplotlib.pyplot as plt diff --git a/python/lsst/analysis/ap/nb_utils.py b/python/lsst/analysis/ap/nb_utils.py index 62f4cee..5fed6c6 100644 --- a/python/lsst/analysis/ap/nb_utils.py +++ b/python/lsst/analysis/ap/nb_utils.py @@ -19,22 +19,77 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . -__all__ = ["make_simbad_link", "compare_sources", "display_images", "get_xy_from_source_table"] +from __future__ import annotations + +__all__ = ["make_simbad_link", "compare_sources", "compare_objects", + "find_objects_sharing_sources", + "classify_association_clusters", + "plot_cutouts_with_object_markers", + "plot_objects_sharing_sources", + "display_images", "display_images_ab", "display_footprints", + "get_xy_from_source_table", "extract_timestamped_messages"] import astropy.coordinates as coord -from astropy.coordinates import SkyCoord, match_coordinates_sky from astroquery.simbad import Simbad import astropy.units as u import astropy.table +from datetime import datetime, timezone +import functools +import json +import numpy as np import os +import random import pandas as pd -import numpy as np +from typing import Any import lsst.afw.display +import lsst.afw.table +from lsst.daf.butler import DatasetNotFoundError from lsst.analysis.ap import plotImageSubtractionCutouts +from lsst.analysis.ap.compare import match_catalogs +from lsst.analysis.ap.skymapOverlay import draw_skymap_outlines_afw from IPython.display import display, Image, Markdown +# Maps the image_type kwarg used by `display_images` to butler dataset +# names. +_IMAGE_DATASETS = { + "science": "preliminary_visit_image", + "template": "template_detector", + "difference": "difference_image", +} + + +def _apply_fakes_prefix(image_datasets, use_fakes): + """Prepend the fake-source pipeline's prefixes to the science and + template image dataset names when ``use_fakes`` is True: ``fakes_`` + for the science image, ``injectedTemplate_`` for the template. The + difference image and all catalogs keep their non-prefixed names + (the fake-source pipeline injects into science + template but + re-uses the same downstream artifacts). + """ + if not use_fakes: + return image_datasets + return { + "science": f"fakes_{image_datasets['science']}", + "template": f"injectedTemplate_{image_datasets['template']}", + "difference": image_datasets["difference"], + } + + +def _cutout_exists(cpath, dia_source_id): + """Return True if a cutout PNG for this diaSourceId already exists. + + Parameters + ---------- + cpath : `~plotImageSubtractionCutouts.CutoutPath` + Path manager for the cutout directory. + dia_source_id : `int` + DiaSourceId whose ``{id}.png`` is checked. + """ + return cpath.exists(dia_source_id, f"{dia_source_id}.png") + + def make_simbad_link(ra, dec, radius_arcsec=3.0): """Search Simbad for associated sources within a 3 arcsecond region. @@ -131,9 +186,10 @@ def compare_sources(butler1, butler2, query1, query2, unique2 : `pandas.DataFrame` Data frame of sources only found in the second dataset. matched : `pandas.DataFrame` - Data frame of matched sources; actually only the sources from - the first dataset but with a new column pointing to the - diaSourceId of the match in the decond dataset. + Data frame of matched sources; the rows are sources from the first + dataset, with two columns added: ``src2_diaSourceId`` pointing to + the matched diaSourceId in the second dataset, and + ``xmatch_dist_arcsec`` giving the on-sky separation in arcseconds. """ if make_cutouts and (cutout_path1 is None or cutout_path2 is None): @@ -141,70 +197,35 @@ def compare_sources(butler1, butler2, query1, query2, raise ValueError(errstr) if bad_flag_list is not None: + # Snapshot and restore so we don't leave the caller's queries with a + # different exclusion list than they started with. + saved_flags1 = list(query1.diaSource_flags_exclude) + saved_flags2 = list(query2.diaSource_flags_exclude) query1.set_excluded_diaSource_flags(bad_flag_list) query2.set_excluded_diaSource_flags(bad_flag_list) - goodSrc1 = query1.load_sources(exclude_flagged=True) - goodSrc2 = query2.load_sources(exclude_flagged=True) + try: + goodSrc1 = query1.load_sources(exclude_flagged=True) + goodSrc2 = query2.load_sources(exclude_flagged=True) + finally: + query1.set_excluded_diaSource_flags(saved_flags1) + query2.set_excluded_diaSource_flags(saved_flags2) + else: + goodSrc1 = query1.load_sources(exclude_flagged=True) + goodSrc2 = query2.load_sources(exclude_flagged=True) if 'reliability' not in goodSrc1.columns: goodSrc1['reliability'] = None if 'reliability' not in goodSrc2.columns: goodSrc2['reliability'] = None - """ - We assume that the runs to be compared here will always be - over the same dataset with the same visit and detector numbers. - Then we can enforce the matching on that level, too. - We could relax that assumption later, though it would require - removing the visit, detector loop. - """ - - visit_det = set(zip(goodSrc1['visit'], goodSrc1['detector'])) - src1 = [] - src2 = [] - matches = [] - indices = {} - seps = {} - - for visit, detector in visit_det: - mask1 = (goodSrc1['visit'] == visit) &\ - (goodSrc1['detector'] == detector) - mask2 = (goodSrc2['visit'] == visit) &\ - (goodSrc2['detector'] == detector) - gs1 = goodSrc1[mask1].copy() - gs2 = goodSrc2[mask2].copy() - - coords1 = SkyCoord(ra=gs1['ra'].values*u.degree, - dec=gs1['dec'].values*u.degree) - coords2 = SkyCoord(ra=gs2['ra'].values*u.degree, - dec=gs2['dec'].values*u.degree) - - index, sep, _ = match_coordinates_sky(coords1, coords2) - indices[(visit, detector)] = index - seps[(visit, detector)] = sep - gs1['xmatch_dist_arcsec'] = sep.to_value(u.arcsecond) - gs1['src2_diaSourceId'] = gs2['diaSourceId'].values.astype(np.int64)[index] - - # Set the match ID to 0 if the distance is above threshold - gs1.loc[(gs1['xmatch_dist_arcsec'] > match_radius), - ['src2_diaSourceId']] = 0 - - # get the DiaSources in dataset 2 not matched to something in dataset 1 - uniqueid2 = set(gs2['diaSourceId']) - set(gs1['src2_diaSourceId']) - unique2 = gs2[gs2['diaSourceId'].isin(uniqueid2)] - - unique1 = gs1[(gs1['src2_diaSourceId'] == 0)] - - withmatch = gs1[(gs1['src2_diaSourceId'] > 0)] - - src1.append(unique1) - src2.append(unique2) - matches.append(withmatch) - - # Out of the loop, concatenate everything together. - unique1 = pd.concat(src1) - unique2 = pd.concat(src2) - matched = pd.concat(matches) + # Cross-match within each (visit, detector) group. + matched, unique1, unique2 = match_catalogs( + goodSrc1, goodSrc2, + radius=match_radius * u.arcsec, + on=("visit", "detector"), + ) + # Preserve the legacy column name on the returned `matched` DataFrame. + matched = matched.rename(columns={"diaSourceId_2": "src2_diaSourceId"}) print("{} matched sources; {} unique to set 1; {} unique to set 2.".format( len(matched), len(unique1), len(unique2))) @@ -239,19 +260,15 @@ def compare_sources(butler1, butler2, query1, query2, plotter2 = plotImageSubtractionCutouts.PlotImageSubtractionCutoutsTask( output_path=cutout_path2, config=config2) - # First figure out which cutouts already exist at the output path - unique1['pathexists'] = False - for i in range(len(unique1)): - dId = unique1.iloc[i]['diaSourceId'] - idx = unique1.index[i] - unique1.at[idx, 'pathexists'] = os.path.exists(cpath1(dId, f"{dId}.png")) + # First figure out which cutouts already exist at the output path. + # Series.apply passes one positional argument (the diaSourceId), but + # _cutout_exists also needs the per-dataset cpath; partial binds it. + unique1['pathexists'] = unique1['diaSourceId'].apply( + functools.partial(_cutout_exists, cpath1)) pathchk1 = unique1.loc[~unique1['pathexists']] - unique2['pathexists'] = False - for i in range(len(unique2)): - dId = unique2.iloc[i]['diaSourceId'] - idx = unique2.index[i] - unique2.at[idx, 'pathexists'] = os.path.exists(cpath2(dId, f"{dId}.png")) + unique2['pathexists'] = unique2['diaSourceId'].apply( + functools.partial(_cutout_exists, cpath2)) pathchk2 = unique2.loc[~unique2['pathexists']] # Only write those that don't exist yet @@ -277,19 +294,1183 @@ def compare_sources(butler1, butler2, query1, query2, return unique1, unique2, matched -def get_xy_from_source_table(table, wcs, degrees=False): +def compare_objects(query1, query2, match_radius=0.5): + """Compare two APDB datasets by spatially crossmatching diaObjects. + + Parameters + ---------- + query1 : `lsst.analysis.ap.DbQuery` + DbQuery to first APDB (postgresql or slite file; + NOT created in this function). + query2 : `lsst.analysis.ap.DbQuery` + DbQuery to second APDB (postgresql or slite file; + NOT created in this function). + match_radius : `double` + Maximum allowable distance in arcsec between an object in + data1 and data2. + + Returns + ------- + unique1 : `pandas.DataFrame` + Data frame of diaObjects only found in the first dataset. + unique2 : `pandas.DataFrame` + Data frame of diaObjects only found in the second dataset. + matched : `pandas.DataFrame` + Data frame of matched diaObjects; the rows are objects from the + first dataset, with two columns added: ``obj2_diaObjectId`` + pointing to the matched diaObjectId in the second dataset, and + ``xmatch_dist_arcsec`` giving the on-sky separation in arcseconds. + """ + obj1 = query1.load_objects() + obj2 = query2.load_objects() + + # diaObjects aren't tied to a single (visit, detector); match across + # the full catalog with `on=()`. + matched, unique1, unique2 = match_catalogs( + obj1, obj2, + radius=match_radius * u.arcsec, + on=(), + id_col="diaObjectId", + ) + matched = matched.rename(columns={"diaObjectId_2": "obj2_diaObjectId"}) + + print("{} matched objects; {} unique to set 1; {} unique to set 2.".format( + len(matched), len(unique1), len(unique2))) + + return unique1, unique2, matched + + +def find_objects_sharing_sources(diaObjectId, sources1, sources2, + objects1, objects2, + max_distance_arcsec=2): + """For a diaObjectId in run 1, return the full association cluster + of diaSources and diaObjects from both runs. + + Treats the (diaSource, run-1-diaObject, run-2-diaObject) links as a + graph -- each diaSource is connected to its owning diaObject in + each run -- and grows the connected component reachable from the + input diaObjectId until no new sources or objects are discovered. + Catches arbitrarily deep merge/split chains across the two runs + (e.g. run 2 merges A+B into Z, then a third source in B is split + into a fourth object in run 2, etc.). + + Assumes diaSourceIds are stable across the two runs. + + Parameters + ---------- + diaObjectId : `int` + A diaObjectId, typically from `unique1` returned by + `compare_objects`. + sources1, sources2 : `pandas.DataFrame` + Full diaSources catalogs from runs 1 and 2 (e.g. from + ``query.load_sources()``). Each must contain `diaSourceId`, + `diaObjectId`, `ra`, and `dec` columns. + objects1, objects2 : `pandas.DataFrame` + Full diaObjects catalogs from runs 1 and 2 (e.g. from + ``query.load_objects()``). Each must contain `diaObjectId`, + `ra`, and `dec` columns. + max_distance_arcsec : `float`, optional + If given, only include diaSources within this distance of the input + diaObject's (ra, dec) in the search. All diaSources of the final + diaObjects will still be returned, even if outside this distance. + + Returns + ------- + sources : `pandas.DataFrame` + Rows of `sources1` for every diaSource belonging to any of the + found run-2 diaObjects (in run 2's view). + related_objects1 : `pandas.DataFrame` + Rows of `objects1` for every run-1 diaObject containing any of + those diaSources in run 1. + related_objects2 : `pandas.DataFrame` + Rows of `objects2` for every run-2 diaObject containing any of + those diaSources in run 2. + """ + if max_distance_arcsec is not None: + ref_match = objects1[objects1["diaObjectId"] == diaObjectId] + if len(ref_match) == 0: + raise ValueError( + f"diaObjectId={diaObjectId} not found in objects1") + ref_row = ref_match.iloc[0] + ref = coord.SkyCoord(ra=ref_row["ra"] * u.deg, + dec=ref_row["dec"] * u.deg) + # diaSourceIds (and their sky positions) match across runs, so + # filtering once against sources1 suffices. + sep = ref.separation( + coord.SkyCoord(ra=sources1["ra"].values * u.deg, + dec=sources1["dec"].values * u.deg) + ).to_value(u.arcsec) + allowed_src_ids = set( + sources1.loc[sep <= max_distance_arcsec, "diaSourceId"]) + else: + allowed_src_ids = None + + # Breadth-first search for the connected component: + # alternately expand sources from the currently-known objects, + # then expand objects from the sources. + # Terminates because every iteration adds at least one source + # before the fixed-point check fires, and the source pool is finite. + src_ids = set() + obj1_ids = {diaObjectId} + obj2_ids = set() + + while True: + new_src_ids = set( + sources1.loc[sources1["diaObjectId"].isin(obj1_ids), + "diaSourceId"]) + new_src_ids.update( + sources2.loc[sources2["diaObjectId"].isin(obj2_ids), + "diaSourceId"]) + if allowed_src_ids is not None: + new_src_ids &= allowed_src_ids + if new_src_ids <= src_ids: + break + src_ids |= new_src_ids + obj1_ids |= set( + sources1.loc[sources1["diaSourceId"].isin(src_ids), + "diaObjectId"]) + obj2_ids |= set( + sources2.loc[sources2["diaSourceId"].isin(src_ids), + "diaObjectId"]) + + # Expand the final diaSource list to every source owned by any + # surviving diaObject. + final_src_ids = set( + sources1.loc[sources1["diaObjectId"].isin(obj1_ids), "diaSourceId"]) + final_src_ids |= set( + sources2.loc[sources2["diaObjectId"].isin(obj2_ids), "diaSourceId"]) + + sources = sources1[sources1["diaSourceId"].isin(final_src_ids)] + related_objects1 = objects1[objects1["diaObjectId"].isin(obj1_ids)] + related_objects2 = objects2[objects2["diaObjectId"].isin(obj2_ids)] + + return sources, related_objects1, related_objects2 + + +class _UnionFind: + """Disjoint-set with path compression and union-by-rank. + + Used by `classify_association_clusters` to quickly find connected + components of the (run-1 diaObject, run-2 diaObject) graph. + """ + + def __init__(self): + self._parent = {} + self._rank = {} + + def add(self, x): + if x not in self._parent: + self._parent[x] = x + self._rank[x] = 0 + + def find(self, x): + # Two-pass iterative find with path compression. + root = x + while self._parent[root] != root: + root = self._parent[root] + while self._parent[x] != root: + self._parent[x], x = root, self._parent[x] + return root + + def union(self, x, y): + rx, ry = self.find(x), self.find(y) + if rx == ry: + return + if self._rank[rx] < self._rank[ry]: + rx, ry = ry, rx + self._parent[ry] = rx + if self._rank[rx] == self._rank[ry]: + self._rank[rx] += 1 + + +def classify_association_clusters(sources1, sources2): + """Enumerate and classify every association-disagreement cluster + between two APDBs that share input diaSources. + + Builds the bipartite graph whose edges are + ``(diaSource -> its run-1 diaObject, diaSource -> its run-2 + diaObject)`` over all common diaSources, runs union-find over the + diaObjectIds to extract every connected component, and labels each + cluster: + + * ``matched`` -- one run-1 obj <-> one run-2 obj. + * ``split`` -- one run-1 obj split into multiple run-2 objs. + * ``merged`` -- multiple run-1 objs merged into one run-2 obj. + * ``tangled`` -- M run-1 objs <-> N run-2 objs, both > 1. + + Assumes diaSourceIds are stable across the two runs; diaSources + present in only one catalog are silently skipped via inner join. + + Parameters + ---------- + sources1, sources2 : `pandas.DataFrame` + Full diaSources catalogs from runs 1 and 2 (e.g. from + ``query.load_sources()``). Each must contain `diaSourceId`, + `diaObjectId`, `ra`, and `dec` columns. + + Returns + ------- + clusters : `pandas.DataFrame` + One row per cluster, with columns: + - ``kind``: matched / split / merged / tangled. + - ``n_obj1``, ``n_obj2``: distinct diaObject counts per run. + - ``n_sources``: distinct diaSources in the cluster. + - ``obj1_ids``, ``obj2_ids``: tuples of diaObjectIds. + - ``ra``, ``dec``: mean sky position of the cluster's + diaSources (degrees). + """ + # Pre-define the types so that value_counts() and groupby() + # include unused kinds with a count of 0. + kind_dtype = pd.CategoricalDtype( + categories=["matched", "split", "merged", "tangled"], ordered=True) + + paired = sources1[["diaSourceId", "diaObjectId", "ra", "dec"]].merge( + sources2[["diaSourceId", "diaObjectId"]].rename( + columns={"diaObjectId": "diaObjectId_2"}), + on="diaSourceId", how="inner") + + if len(paired) == 0: + empty = pd.DataFrame(columns=[ + "kind", "n_obj1", "n_obj2", "n_sources", + "obj1_ids", "obj2_ids", "ra", "dec"]) + empty["kind"] = empty["kind"].astype(kind_dtype) + return empty + + # Define a namespace for the two runs so identical numeric ids in run 1 and + # run 2 don't collide as keys. + keys1 = [("r1", int(i)) for i in paired["diaObjectId"].to_numpy()] + keys2 = [("r2", int(i)) for i in paired["diaObjectId_2"].to_numpy()] + + uf = _UnionFind() + for k in set(keys1): + uf.add(k) + for k in set(keys2): + uf.add(k) + for k1, k2 in zip(keys1, keys2): + uf.union(k1, k2) + + paired = paired.assign(_cluster=[uf.find(k) for k in keys1]) + + rows = [] + for _, grp in paired.groupby("_cluster", sort=False): + ids_a = tuple(sorted(int(i) for i in grp["diaObjectId"].unique())) + ids_b = tuple(sorted(int(i) for i in grp["diaObjectId_2"].unique())) + n1, n2 = len(ids_a), len(ids_b) + if n1 == 1 and n2 == 1: + kind = "matched" + elif n1 == 1: + kind = "split" + elif n2 == 1: + kind = "merged" + else: + kind = "tangled" + rows.append({ + "kind": kind, + "n_obj1": n1, "n_obj2": n2, + "n_sources": grp["diaSourceId"].nunique(), + "obj1_ids": ids_a, "obj2_ids": ids_b, + "ra": float(grp["ra"].mean()), + "dec": float(grp["dec"].mean()), + }) + + result = pd.DataFrame(rows) + result["kind"] = result["kind"].astype(kind_dtype) + return result + + +# Colors used by the cutout plotters to give each distinct +# diaObjectId its own marker color. +_OBJECT_PALETTE = ("lime", "red", "cyan", "magenta", "yellow", "orange", + "deepskyblue", "pink", "white", "violet", "gold", + "lightgreen") + + +def _prepare_object_overlays(objects, palette): + """Deduplicate `objects` by diaObjectId and assign one palette color + per distinct id, returning the parallel arrays the cutout renderer + needs: ``(ids, ras, decs, colors)``. Done once per call so the same + color identifies the same diaObject across every cutout. + """ + obj_unique = objects.drop_duplicates(subset="diaObjectId") + # Prefer the run-2 id when present (matched rows carry both); pandas + # concat promotes the column to float64 if any rows lack it, so cast + # back to int64 after filling. + if "obj2_diaObjectId" in obj_unique.columns: + obj_ids = obj_unique["obj2_diaObjectId"].combine_first( + obj_unique["diaObjectId"]).astype(np.int64).to_numpy() + else: + obj_ids = obj_unique["diaObjectId"].astype(np.int64).to_numpy() + obj_ras = np.asarray(obj_unique["ra"]) + obj_decs = np.asarray(obj_unique["dec"]) + obj_colors = [palette[i % len(palette)] for i in range(len(obj_ids))] + return obj_ids, obj_ras, obj_decs, obj_colors + + +def _load_cutout(butler, row, *, size, image_type, image_datasets): + """Fetch the requested image dataset for this row's (visit, + detector) and return a small dict with everything the renderer + needs: pixel data, dimensions, cutout origin, WCS, an + ImageNormalize tuned to the central source, and the `image_type` + label used in the cutout title. + + Loading is separated from rendering so callers that need to draw + the same cutout into multiple Axes can pay the butler.get + getCutout + cost once. + """ + import astropy.visualization as aviz + import lsst.geom + + dataset = image_datasets[image_type] + data_id = {"visit": int(row.visit), "detector": int(row.detector)} + exposure = butler.get(dataset, data_id) + + center = lsst.geom.SpherePoint(row.ra, row.dec, lsst.geom.degrees) + extent = lsst.geom.Extent2I(size, size) + cutout = exposure.getCutout(center, extent) + data = cutout.image.array + ny, nx = data.shape + + if image_type == "difference": + # Normalize on a small central window so the source dominates + # the dynamic range. + cy, cx = ny // 2, nx // 2 + half = min(7, cy, cx) + norm_data = data[cy - half:cy + half + 1, cx - half:cx + half + 1] + else: + norm_data = data + norm = aviz.ImageNormalize( + norm_data, interval=aviz.MinMaxInterval(), + stretch=aviz.AsinhStretch(a=0.1)) + + return { + "data": data, "ny": ny, "nx": nx, + "wcs": cutout.wcs, + "x0": cutout.getX0(), "y0": cutout.getY0(), + "norm": norm, + "image_type": image_type, + } + + +def _render_cutout_axes(ax, row, cutout_data, sources, + obj_ids, obj_ras, obj_decs, obj_colors, *, + marker_size, marker_symbol, + source_marker_size, current_source_marker_size, + current_source_color, + source_match_ids=None, + title=None, subtitle=""): + """Render one diaSource cutout onto an existing matplotlib Axes + using preloaded data from `_load_cutout`. + + Internal helper shared by `plot_cutouts_with_object_markers` and + `plot_objects_sharing_sources`. The caller owns figure creation, + layout, saving, and displaying. + + By default the axes title is built from `row` as + ``"diaSourceId=... (image_type, visit=..., det=...)"``. Pass + `title=` explicitly (including ``""``) to override or suppress + that line -- useful when a parent figure or subfigure already + carries the shared header. If `subtitle` is non-empty it is drawn + on a second title line. + """ + from matplotlib import cm + + data = cutout_data["data"] + ny = cutout_data["ny"] + nx = cutout_data["nx"] + x0 = cutout_data["x0"] + y0 = cutout_data["y0"] + wcs = cutout_data["wcs"] + norm = cutout_data["norm"] + image_type = cutout_data["image_type"] + + ax.imshow(data, cmap=cm.bone, interpolation="none", norm=norm, + origin="lower", aspect="equal", + extent=(0, nx, 0, ny)) + + this_id = int(row.diaSourceId) + + # Project every supplied diaSource into the cutout frame once, + # then split into "this cutout's diaSource" vs every other + # diaSource whose sky position falls inside this cutout + # (regardless of which image it was detected on). + if len(sources) > 0: + src_xs, src_ys = wcs.skyToPixelArray( + np.asarray(sources["ra"]), + np.asarray(sources["dec"]), + degrees=True) + src_xs = src_xs - x0 + src_ys = src_ys - y0 + src_in_frame = ( + (src_xs >= 0) & (src_xs < nx) + & (src_ys >= 0) & (src_ys < ny)) + id_arr = sources["diaSourceId"].to_numpy() + other_src_mask = src_in_frame & (id_arr != this_id) + current_src_mask = src_in_frame & (id_arr == this_id) + else: + src_xs = src_ys = None + other_src_mask = current_src_mask = None + + # Per-diaSource color resolution: map each diaSource to its + # owning diaObject (in this panel's view), then to that + # diaObject's palette color. Falls back to `current_source_color` + # for sources whose owner is not present in `obj_ids`. + if source_match_ids is None: + match_ids_list = sources["diaObjectId"].tolist() + else: + match_ids_list = list(source_match_ids) + src_to_match = dict(zip(sources["diaSourceId"].tolist(), + match_ids_list)) + id_to_color = dict(zip(obj_ids.tolist(), obj_colors)) + + def _color_for(diaSourceId): + return id_to_color.get( + src_to_match.get(int(diaSourceId)), current_source_color) + + if other_src_mask is not None and other_src_mask.any(): + other_indices = np.flatnonzero(other_src_mask) + other_colors = [_color_for(int(id_arr[i])) for i in other_indices] + ax.scatter(src_xs[other_src_mask], src_ys[other_src_mask], + s=source_marker_size, marker="+", + c=other_colors, linewidths=1.0, + label="other diaSource") + + if len(obj_ids) > 0: + xs, ys = wcs.skyToPixelArray(obj_ras, obj_decs, degrees=True) + xs = xs - x0 + ys = ys - y0 + in_bounds = (xs >= 0) & (xs < nx) & (ys >= 0) & (ys < ny) + else: + in_bounds = np.zeros(0, dtype=bool) + + for i in np.flatnonzero(in_bounds): + ax.scatter(xs[i], ys[i], + s=marker_size, marker=marker_symbol, + facecolors="none", edgecolors=obj_colors[i], + linewidths=1.5, + label=f"diaObjectId={int(obj_ids[i])}") + + # Current diaSource last so it stays on top of any overlapping + # diaObject marker at the cutout center. + if current_src_mask is not None and current_src_mask.any(): + ax.scatter(src_xs[current_src_mask], src_ys[current_src_mask], + s=current_source_marker_size, marker="x", + c=_color_for(this_id), linewidths=2.0, + label=f"current diaSourceId={this_id}") + + if title is None: + title = (f"diaSourceId={this_id} " + f"({image_type}, visit={int(row.visit)}, " + f"det={int(row.detector)})") + if subtitle: + title = f"{title}\n{subtitle}" if title else subtitle + if title: + ax.set_title(title, fontsize="small") + ax.set_xticks([]) + ax.set_yticks([]) + if ax.get_legend_handles_labels()[0]: + ax.legend(loc="upper right", fontsize="x-small", framealpha=0.7) + + +def plot_cutouts_with_object_markers(sources, butler, objects, *, + output_path=None, + display_cutouts=False, + size=51, + image_type="difference", + image_datasets=_IMAGE_DATASETS, + marker_size=80, + marker_symbol="o", + palette=_OBJECT_PALETTE, + source_marker_size=80, + current_source_marker_size=180, + current_source_color="yellow"): + """Plot per-diaSource cutouts with overlaid markers at given diaObject + sky positions. + + For each diaSource in `sources`, fetch a square cutout from `butler` + centered on the source's (ra, dec). On each cutout draw: + + * A small ``+`` marker at every other diaSource in `sources` + whose sky position lands inside the cutout, regardless of which + (visit, detector) it was detected on. + * A distinct ``x`` marker for the diaSource the cutout is + centered on (the "current" diaSource). + * One color-coded marker per distinct diaObjectId in `objects`, + cycling through `palette`; the same color identifies the same + diaObject across every cutout in the run. + + Markers that fall outside the cutout bounds are skipped. + + Typical use: visualize how a group of diaSources (all originally + associated with one diaObjectId in run 1) got redistributed across + diaObjects in run 2. `sources` and `objects` are usually built from + the output of `find_objects_sharing_sources`:: + + sources, ro1, ro2 = find_objects_sharing_sources( + diaObjectId, sources1, sources2, objects1, objects2) + objects = pd.concat([ro1, ro2]) + plot_cutouts_with_object_markers( + sources, butler1, objects, display_cutouts=True, + ) + + Parameters + ---------- + sources : `pandas.DataFrame` + DiaSources to cut out. Must contain `diaSourceId`, `ra`, `dec`, + `visit`, and `detector` columns. + butler : `lsst.daf.butler.Butler` + Butler containing the image datasets for these (visit, detector) + pairs. + objects : `pandas.DataFrame` + DiaObjects to mark. Must contain `diaObjectId`, `ra`, and `dec` + columns. Duplicate diaObjectIds are dropped (first row wins). + If an `obj2_diaObjectId` column is present (e.g. for rows from + a `matched` DataFrame returned by `compare_objects`), the run-2 + id is shown in the legend in preference to the run-1 + `diaObjectId`. + output_path : `str`, optional + Directory to write ``{diaSourceId}.png`` files to. Created if + missing. Pass None to skip writing. + display_cutouts : `bool`, optional + If True, display each cutout inline (notebook). + size : `int`, optional + Cutout side length in pixels. + image_type : {"science", "template", "difference"}, optional + Which image to render. + image_datasets : `dict` [`str`, `str`], optional + Mapping from image-type key to butler dataset name. + marker_size : `int`, optional + matplotlib scatter ``s`` parameter for diaObject markers. + marker_symbol : `str`, optional + matplotlib scatter ``marker`` parameter for diaObject markers. + palette : sequence of `str`, optional + Color cycle used to assign one color per diaObjectId. + source_marker_size : `int`, optional + Scatter ``s`` parameter for the small ``+`` markers drawn at + the positions of the *other* diaSources in `sources`. + current_source_marker_size : `int`, optional + Scatter ``s`` parameter for the distinct marker drawn at the + diaSource the cutout is centered on. + current_source_color : `str`, optional + Color of the current-diaSource marker. + """ + import matplotlib.pyplot as plt + + if image_type not in image_datasets: + raise ValueError( + f"image_type must be one of {sorted(image_datasets)}, " + f"got {image_type!r}") + + if output_path is not None: + os.makedirs(output_path, exist_ok=True) + + overlays = _prepare_object_overlays(objects, palette) + + for row in sources.itertuples(index=False): + cutout_data = _load_cutout(butler, row, size=size, + image_type=image_type, + image_datasets=image_datasets) + fig, ax = plt.subplots() + _render_cutout_axes( + ax, row, cutout_data, sources, *overlays, + marker_size=marker_size, marker_symbol=marker_symbol, + source_marker_size=source_marker_size, + current_source_marker_size=current_source_marker_size, + current_source_color=current_source_color) + + if output_path is not None: + fpath = os.path.join(output_path, f"{int(row.diaSourceId)}.png") + fig.savefig(fpath, bbox_inches="tight") + if display_cutouts: + display(fig) + plt.close(fig) + + +def plot_objects_sharing_sources(diaObjectId, sources1, sources2, + objects1, objects2, butler, *, + max_distance_arcsec=None, + output_path=None, + display_figure=True, + column_labels=("run 1", "run 2"), + figsize_per_row=4.0, + size=51, + image_type="difference", + image_datasets=_IMAGE_DATASETS, + marker_size=80, + marker_symbol="o", + palette=_OBJECT_PALETTE, + source_marker_size=80, + current_source_marker_size=180, + current_source_color="yellow"): + """Two-column cutout figure comparing the run-1 and run-2 views of + an association cluster. + + Calls `find_objects_sharing_sources` internally to identify the + cluster of diaSources and diaObjects reachable from the input + `diaObjectId`, then renders one row per diaSource in the cluster. + The same cutout image is loaded once per row and drawn into both + columns: the left panel is overlaid with run-1 diaObject markers + (from `objects1`) and the right panel with run-2 diaObject markers + (from `objects2`). Each column's diaObjects get their own palette + mapping, so the same color in the left and right columns does + *not* imply the same diaObject. + + Parameters + ---------- + diaObjectId : `int` + Starting diaObjectId for the cluster walk. + sources1, sources2, objects1, objects2 : `pandas.DataFrame` + Forwarded to `find_objects_sharing_sources`. + butler : `lsst.daf.butler.Butler` + Butler used to fetch the cutout images. The same image backs + both panels of a given row. + max_distance_arcsec : `float`, optional + Forwarded to `find_objects_sharing_sources`. + output_path : `str`, optional + Filename to write the combined figure to as a PNG. Parent + directories are created if missing. + display_figure : `bool`, optional + If True, display the figure inline (notebook). + column_labels : pair of `str`, optional + Labels appended to each cutout's title to identify the column. + figsize_per_row : `float`, optional + Height in inches allocated to each cutout row. + All other kwargs: + Forwarded to the cutout renderer; same meaning as in + `plot_cutouts_with_object_markers`. + + Returns + ------- + sources, related_objects1, related_objects2 : `pandas.DataFrame` + The catalogs returned by `find_objects_sharing_sources`. + """ + import matplotlib.pyplot as plt + + if image_type not in image_datasets: + raise ValueError( + f"image_type must be one of {sorted(image_datasets)}, " + f"got {image_type!r}") + + sources, ro1, ro2 = find_objects_sharing_sources( + diaObjectId, sources1, sources2, objects1, objects2, + max_distance_arcsec=max_distance_arcsec) + + if len(sources) == 0: + print(f"No diaSources in the cluster for " + f"diaObjectId={diaObjectId}") + return sources, ro1, ro2 + + left_overlays = _prepare_object_overlays(ro1, palette) + right_overlays = _prepare_object_overlays(ro2, palette) + + # diaSourceId -> run-2 diaObjectId, so the right panel can color + # each diaSource by its run-2 owner. The left panel uses the + # default (sources["diaObjectId"], the run-1 owner) since + # `sources` is a slice of `sources1`. + src_to_obj2 = dict(zip( + sources2["diaSourceId"].to_numpy(), + sources2["diaObjectId"].to_numpy())) + right_match_ids = [src_to_obj2.get(int(sid)) + for sid in sources["diaSourceId"]] + + n_rows = len(sources) + # One subfigure per row so each row can carry a single shared + # suptitle above both panels; the per-axes title is then just the + # column label. + fig = plt.figure(figsize=(8, figsize_per_row * n_rows), + constrained_layout=True) + subfigs = np.atleast_1d(fig.subfigures(n_rows, 1, squeeze=False).ravel()) + + common_kw = dict( + marker_size=marker_size, marker_symbol=marker_symbol, + source_marker_size=source_marker_size, + current_source_marker_size=current_source_marker_size, + current_source_color=current_source_color) + + for i, row in enumerate(sources.itertuples(index=False)): + sf = subfigs[i] + sf.suptitle( + f"diaSourceId={int(row.diaSourceId)} " + f"({image_type}, visit={int(row.visit)}, " + f"det={int(row.detector)})", + fontsize="small") + ax_left, ax_right = sf.subplots(1, 2) + # Load once; both panels in this row use the same image. + cutout_data = _load_cutout(butler, row, size=size, + image_type=image_type, + image_datasets=image_datasets) + _render_cutout_axes(ax_left, row, cutout_data, sources, + *left_overlays, + title="", subtitle=column_labels[0], + **common_kw) + _render_cutout_axes(ax_right, row, cutout_data, sources, + *right_overlays, + title="", subtitle=column_labels[1], + source_match_ids=right_match_ids, + **common_kw) + + if output_path is not None: + out_dir = os.path.dirname(output_path) + if out_dir: + os.makedirs(out_dir, exist_ok=True) + fig.savefig(output_path, bbox_inches="tight") + if display_figure: + display(fig) + plt.close(fig) + + return sources, ro1, ro2 + + +def get_xy_from_source_table(table, wcs, degrees=None): """Convert ra/dec coordinates in an astropy table/pandas data frame to pixel x/y positions. """ - ra = table['coord_ra'] - dec = table['coord_dec'] + try: + ra = table['ra'] + dec = table['dec'] + inferred_degrees = True + except KeyError: + ra = table['coord_ra'] + dec = table['coord_dec'] + inferred_degrees = False + if degrees is None: + degrees = inferred_degrees + x, y = wcs.skyToPixelArray(ra, dec, degrees=degrees) return astropy.table.Table.from_pandas(pd.DataFrame({'x': x, 'y': y})) -def display_images(butler, visit, detector, backend="firefly"): - """Display the science/template/difference images for a given - visit+detector, with sources and mask planes overlaid. +# Palette used by the `color_by` flag-bucketing mode. Sources with none of +# the requested flags set get the residual color "white", which is kept out +# of this palette so it never collides with a flagged bucket. +_FLAG_PALETTE = ("red", "orange", "yellow", "magenta", "cyan", "green") + + +def _group_sources_by_flag(table, flag_names, palette=_FLAG_PALETTE): + """Split a source table into per-flag buckets for color-coded overlay. + + Each row is assigned to the first flag in ``flag_names`` whose column + is True; remaining rows go into a residual "no flag" bucket. Names that + aren't present as columns in ``table`` are silently skipped. + + Parameters + ---------- + table : table-like + Anything that supports ``len(table)``, ``table[name]`` returning a + boolean-coercible column, and ``table[bool_array]`` row selection. + flag_names : sequence of str + Column names to group on. Order determines color *and* priority + when a row has multiple flags set. + palette : sequence of str, optional + Cycle of display ``ctype`` values to assign in order. + + Returns + ------- + buckets : list of ``(subset_table, ctype, legend)`` tuples. + """ + n = len(table) + if n == 0: + return [] + remaining = np.ones(n, dtype=bool) + buckets = [] + for i, flag in enumerate(flag_names): + try: + col = table[flag] + except KeyError: + continue + mask = np.asarray(col, dtype=bool) & remaining + if mask.any(): + buckets.append((table[mask], palette[i % len(palette)], flag)) + remaining = remaining & ~mask + if remaining.any(): + buckets.append((table[remaining], "white", "no flag")) + return buckets + + +def _line_segments_from(source_table, wcs, *, flag_col, angle_col, ctype, + length_col=None, fixed_length=None, + min_length=None, length_scale=1.0, thickness=1.5): + """Build centered line-segment endpoints from a source catalog. + + Sources are selected by ``flag_col`` (rows where the boolean column + is True) when given, then optionally by ``length_col > min_length`` + (only meaningful when ``length_col`` is set), then always by + ``sky_source == False`` when that column is present. Endpoints are + the source centroid ± half the (scaled) length along ``angle_col``. + + Exactly one of ``length_col`` (per-source measurement, e.g. + ``ext_trailedSources_Naive_length`` in pixels) or ``fixed_length`` + (constant, in pixels — used for markers whose real separation is + below display resolution) must be given. Length units are pixels + and angle units are radians in the detector frame, matching the + raw ``ip_diffim`` and ``ext_trailedSources`` measurement outputs + on ``dia_source_unfiltered``. ``length_scale`` multiplies measured + lengths after filtering; it is a no-op when ``fixed_length`` is + used, since a "fixed" length by definition is not magnified. + + ``thickness`` is the stroke width forwarded to ``afw_display.line``; + stored as ``size`` in the returned dict. + + Returns a dict ``{"x1", "y1", "x2", "y2", "ctype", "size"}`` of + numpy arrays and scalars, or ``None`` if the table is missing/empty + or any referenced column is absent. + """ + if (length_col is None) == (fixed_length is None): + raise ValueError("_line_segments_from: pass exactly one of " + "length_col or fixed_length") + if source_table is None or len(source_table) == 0: + return None + try: + angle = np.asarray(source_table[angle_col], dtype=float) + mask = (np.asarray(source_table[flag_col], dtype=bool) if flag_col is not None + else np.ones(len(source_table), dtype=bool)) + if length_col is not None: + length = np.asarray(source_table[length_col], dtype=float) + else: + length = np.full(len(source_table), fixed_length, dtype=float) + except KeyError: + return None + if length_col is not None and min_length is not None: + mask = mask & (length > min_length) + try: + mask = mask & ~np.asarray(source_table["sky_source"], dtype=bool) + except KeyError: + pass + if not mask.any(): + return None + xy = get_xy_from_source_table(source_table[mask], wcs) + x0 = xy["x"].data + y0 = xy["y"].data + effective_scale = length_scale if fixed_length is None else 1.0 + half = length[mask] * effective_scale / 2.0 + dx = half * np.cos(angle[mask]) + dy = half * np.sin(angle[mask]) + return {"x1": x0 - dx, "y1": y0 - dy, + "x2": x0 + dx, "y2": y0 + dy, + "ctype": ctype, "size": thickness} + + +def _collect_overlays(butler, data_id, wcs, *, + reliability_threshold, + show_unfiltered, show_trailed, + show_rejected, show_standardized, show_marginal, + show_kernel_sources, show_solar_system, show_apdb, + show_reliability_labels, show_dipoles, + show_trail_geometry, line_length_scale, + color_by): + """Load catalogs from one butler and build the overlay record list. + + Shared between `display_images` and `display_images_ab`. Catalogs that + aren't present for this dataId are silently skipped. + + Returns + ------- + overlays : list of ``(x_arr, y_arr, symbol, size, ctype, legend)`` tuples. + reliability_labels : dict or None + ``{"x", "y", "reliability"}`` arrays for the diaSources whose + reliability score should be drawn as text. + solar_system_labels : dict or None + ``{"x", "y", "designation"}`` arrays for matched solar-system + sources, suitable for drawing the designation as text next to each + marker. + dipole_segments, trail_segments : dict or None + Each ``{"x1", "y1", "x2", "y2", "ctype"}`` arrays describing line + segments through diaSource centroids to visualize dipole + (``isDipole`` / ``dipoleLength`` / ``dipoleAngle``) and trail + (``trailLength`` / ``trailAngle``) geometry from the standardized + or APDB catalog. + """ + def _try_get(dataset): + try: + return butler.get(dataset, data_id) + except DatasetNotFoundError: + return None + + overlays = [] + + def _add(table, *, symbol, size, ctype, legend, use_radec=True): + if table is None or len(table) == 0: + return + if use_radec: + xy = get_xy_from_source_table(table, wcs) + x_arr = xy["x"].data + y_arr = xy["y"].data + else: + x_arr = table["x"].data + y_arr = table["y"].data + overlays.append((x_arr, y_arr, symbol, size, ctype, legend)) + + # Catalogs are added in AP-pipeline-creation order + if show_kernel_sources: + _add(_try_get("difference_kernel_sources"), + symbol="o", size=12, ctype="green", + legend="psf-matching kernel source") + + # Load dia_source_unfiltered once — it backs the unfiltered marker + # overlay AND the dipole/trail line-segment overlays below (which + # always read from the unfiltered catalog because dipoles and long + # trails are filtered out of the downstream catalogs). + unfiltered = None + if show_unfiltered or show_dipoles or show_trail_geometry: + unfiltered = _try_get("dia_source_unfiltered") + + if show_unfiltered and unfiltered is not None and len(unfiltered) > 0: + non_sky = unfiltered[~unfiltered["sky_source"]] + if color_by: + for sub, ctype, flag in _group_sources_by_flag(non_sky, color_by): + _add(sub, symbol="+", size=10, ctype=ctype, + legend=f"unfiltered: {flag}") + else: + _add(non_sky, symbol="+", size=10, ctype="red", + legend="unfiltered candidate") + if show_rejected: + _add(_try_get("rejected_dia_source"), + symbol="+", size=10, ctype="orange", legend="rejected diaSource") + if show_trailed: + _add(_try_get("long_trailed_source_detector"), + symbol="x", size=30, ctype="magenta", legend="long-trailed source") + + # Stash the standardized catalog + projected xy for reuse by the + # geometry overlays below. `standardizeDiaSource` runs between + # filterDiaSource and associateApdb; when the pipeline stops before + # the APDB ingest, dia_source_detector is the last diaSource catalog + # available. + standardized_data = None + if show_standardized: + standardized = _try_get("dia_source_detector") + if standardized is not None and len(standardized) > 0: + xy = get_xy_from_source_table(standardized, wcs) + x_arr = xy["x"].data + y_arr = xy["y"].data + overlays.append((x_arr, y_arr, "+", 10, "blue", "standardized diaSource")) + standardized_data = {"catalog": standardized, "x": x_arr, "y": y_arr} + + # Load dia_source_apdb once: it backs the APDB reliability overlay and + # also supplies pixel x/y for the solar-system overlay (ss_source_detector + # carries only the matched diaSourceId, not coordinates). + dia_apdb = None + if show_solar_system or show_apdb: + dia_apdb = _try_get("dia_source_apdb") + + if show_apdb and dia_apdb is not None and len(dia_apdb) > 0: + good_mask = dia_apdb["reliability"] > reliability_threshold + _add(dia_apdb[good_mask], symbol="o", size=14, ctype="blue", use_radec=False, + legend=f"APDB, reliability > {reliability_threshold:g}") + _add(dia_apdb[~good_mask], symbol="o", size=14, ctype="red", use_radec=False, + legend=f"APDB, reliability <= {reliability_threshold:g}") + + solar_system_labels = None + if show_solar_system: + ss = _try_get("ss_source_detector") + if (ss is not None and len(ss) > 0 + and dia_apdb is not None and len(dia_apdb) > 0): + # ss_source_detector lacks coords; match each diaSourceId to + # the APDB row to recover its pixel x/y. + ss_ids = np.asarray(ss["diaSourceId"]) + apdb_ids = np.asarray(dia_apdb["diaSourceId"]) + idx_in_apdb = pd.Series(np.arange(len(apdb_ids)), index=apdb_ids).reindex(ss_ids) + keep = idx_in_apdb.notna().to_numpy() + if keep.any(): + apdb_idx = idx_in_apdb.dropna().astype(int).to_numpy() + x_arr = np.asarray(dia_apdb["x"])[apdb_idx] + y_arr = np.asarray(dia_apdb["y"])[apdb_idx] + designation = np.asarray(ss["designation"])[keep] + overlays.append((x_arr, y_arr, "o", 16, "cyan", "solar-system match")) + solar_system_labels = {"x": x_arr, "y": y_arr, "designation": designation, } + + if show_marginal: + _add(_try_get("marginal_new_dia_source"), + symbol="+", size=10, ctype="yellow", legend="marginal new diaSource") + + # Reliability text is drawn at most once per diaSource: prefer APDB + # good sources when APDB is displayed, and otherwise annotate every + # standardized row (they've been pipeline-filtered to high + # reliability already). The unfiltered catalog doesn't carry a final + # reliability score, so it never provides labels. + reliability_labels = None + apdb_shown = show_apdb and dia_apdb is not None and len(dia_apdb) > 0 + if show_reliability_labels: + if apdb_shown: + rel = np.asarray(dia_apdb["reliability"]) + mask = rel > reliability_threshold + if mask.any(): + reliability_labels = {"x": np.asarray(dia_apdb["x"])[mask], + "y": np.asarray(dia_apdb["y"])[mask], + "reliability": rel[mask]} + elif standardized_data is not None: + reliability_labels = { + "x": standardized_data["x"], + "y": standardized_data["y"], + "reliability": np.asarray(standardized_data["catalog"]["reliability"]), + } + + # Dipole and trail line segments both come from dia_source_unfiltered: + # it's the earliest AP-pipeline catalog and thus a superset of the + # downstream ones that get dipoles/long trails filtered out, and it + # carries the raw ip_diffim / ext_trailedSources measurement columns + # in native pixel + detector-radian units — exactly the coordinate + # system the endpoint math lives in. + dipole_segments = None + if show_dipoles: + # Fixed 10 px length: the measured ``ip_diffim_DipoleFit_separation`` + # is sub-pixel for the vast majority of classified dipoles (median + # ~0.09 px in typical data), so drawing at the measured length + # would hide them. The line here is a fixed-size *marker* of the + # dipole's orientation, not a physical extent. + dipole_segments = _line_segments_from( + unfiltered, wcs, + flag_col="ip_diffim_DipoleFit_classification", + angle_col="ip_diffim_DipoleFit_orientation", + ctype="white", fixed_length=10.0, thickness=3.0) + + # 3 px threshold: below this the trail measurement is dominated by + # noise on point-like sources. Long-trailed sources removed by + # filterDiaSource still show up via ``show_trailed`` as ``x`` markers. + trail_segments = None + if show_trail_geometry: + trail_segments = _line_segments_from( + unfiltered, wcs, flag_col=None, + length_col="ext_trailedSources_Naive_length", + angle_col="ext_trailedSources_Naive_angle", + ctype="magenta", min_length=3.0, + length_scale=line_length_scale) + + return overlays, reliability_labels, solar_system_labels, dipole_segments, trail_segments + + +def _print_overlay_legend(overlays, header, indent=""): + """Print a one-line-per-overlay legend for a single panel.""" + print(f"{indent}{header}") + for x_arr, _, symbol, _, ctype, legend in overlays: + print(f"{indent} {len(x_arr):5d} {ctype:>8s} {symbol} {legend}") + + +def _draw_line_segments(afw_display, segments): + """Draw one batch of centered line segments on the active frame.""" + if segments is None: + return + ctype = segments["ctype"] + size = segments["size"] + for x1, y1, x2, y2 in zip(segments["x1"], segments["y1"], + segments["x2"], segments["y2"]): + afw_display.line([(float(x1), float(y1)), (float(x2), float(y2))], + ctype=ctype, size=size) + + +def _draw_overlays_on_current_frame(afw_display, overlays, + reliability_labels, solar_system_labels, + dipole_segments=None, + trail_segments=None, + label_size=3): + """Stamp one set of overlays + optional reliability and solar-system + designation labels onto the active frame. + + ``label_size`` is the text size (in pixels) used for both label sets. + ``dipole_segments`` and ``trail_segments`` are optional line-segment + dicts (see `_line_segments_from`). + """ + # Scale the text offset with the size so larger labels still clear the + # circle markers they annotate. + label_offset = max(14, 2 * label_size) + with afw_display.Buffering(): + for x_arr, y_arr, symbol, size, ctype, _ in overlays: + for x, y in zip(x_arr, y_arr): + afw_display.dot(symbol, x, y, size=size, ctype=ctype) + # Trails first, dipoles on top: a source flagged as a dipole is + # the more actionable pipeline-quality issue, so it wins any + # pixel overlap with the trail line. + _draw_line_segments(afw_display, trail_segments) + _draw_line_segments(afw_display, dipole_segments) + if reliability_labels is not None: + # Offset the score text so it doesn't sit on top of the marker. + for r, x, y in zip(reliability_labels["reliability"], + reliability_labels["x"], + reliability_labels["y"]): + afw_display.dot(f"{r:.2f}", x + label_offset, y, + size=label_size, ctype="cyan") + if solar_system_labels is not None: + # Offset SS designations *below* the marker so they don't + # overplot any reliability score drawn to the right. + for desig, x, y in zip(solar_system_labels["designation"], + solar_system_labels["x"], + solar_system_labels["y"]): + afw_display.dot(str(desig), x, y + label_offset, + size=label_size, ctype="cyan") + + +def _strip_ds9_metadata(*exposures): + """Drop LTV1/LTV2 keys from each exposure's metadata in place.""" + for exp in exposures: + md = exp.metadata + for k in ("LTV1", "LTV2"): + if md.exists(k): + md.remove(k) + + +def _erase_current_frame_regions(afw_display): + """Clear all region markers on the currently-selected frame. + + The firefly backend caches ``_regionLayerId`` on its impl and only + refreshes it inside ``_flush()``. Calling ``erase()`` right after + switching frames therefore issues ``delete_region_layer`` with the + *previous* frame's layer id under the current frame's plot id, so + nothing gets removed. Refreshing the cached id from the current + frame before erasing fixes it. Backends that don't cache a layer + id (e.g. ds9) fall through to a plain ``erase()``. + """ + impl = getattr(afw_display, "_impl", None) + if impl is not None and hasattr(impl, "_getRegionLayerId"): + impl._regionLayerId = impl._getRegionLayerId() + afw_display.erase() + + +def display_images(butler, visit, detector, backend="firefly", *, + reliability_threshold=0.1, + show_unfiltered=True, + show_trailed=True, + show_rejected=True, + show_standardized=True, + show_marginal=True, + show_kernel_sources=True, + show_solar_system=True, + show_apdb=True, + show_reliability_labels=True, + show_dipoles=True, + show_trail_geometry=True, + line_length_scale=1.0, + label_size=3, + color_by=None, + mask_transparency=80, + strip_metadata=True, + skymap=None, + skymap_ctype="green", + skymap_label_size=1.5, + image_datasets=_IMAGE_DATASETS, + use_fakes=False): + """Display the science, template, and difference images for a given + visit+detector with diagnostic catalog markers overlaid. + + Three frames are produced (science, template, difference) and the same + overlays are drawn on each. Catalogs that are missing from the butler + are silently skipped, so the same call works against partial outputs. + + Default overlay key. Rows are in AP-pipeline creation order, so the + last marker drawn at any pixel reflects the latest classification the + pipeline assigned. Circle sizes step by 2 so successive ``o`` markers + nest rather than stack. + + ============================ ======= ==== =========================== + catalog symbol size color + ============================ ======= ==== =========================== + psf-matching kernel sources ``o`` 12 green + unfiltered candidates ``+`` 10 red + rejected diaSources ``+`` 10 orange + long-trailed sources ``x`` 30 magenta + standardized diaSources ``+`` 10 blue + APDB, reliability > threshold ``o`` 14 blue (+ score text) + APDB, reliability ≤ threshold ``o`` 14 red + solar-system matches ``o`` 16 cyan + marginal new diaSources ``+`` 10 yellow + ============================ ======= ==== =========================== Parameters ---------- @@ -297,65 +1478,597 @@ def display_images(butler, visit, detector, backend="firefly"): Butler to load data from. visit, detector : `int` Visit and detector ids to load data for. - backend : str, optional - afw display backend to display to (typically "firefly" or "ds9"). + backend : `str`, optional + afw display backend (typically "firefly" or "ds9"). + reliability_threshold : `float`, optional + APDB diaSources with reliability strictly greater than this are + drawn as "good" (blue); the rest as "bad" (red). + show_unfiltered, show_trailed, show_rejected, show_standardized, + show_marginal, show_kernel_sources, show_solar_system, + show_apdb : `bool`, optional + Toggle individual catalog overlays. ``show_kernel_sources`` + loads ``difference_kernel_sources``, the PSF-matching constraint + sources from image subtraction — useful for seeing where the + kernel was actually anchored vs extrapolated. + show_reliability_labels : `bool`, optional + If True, annotate each good APDB diaSource with its reliability score. + show_dipoles : `bool`, optional + If True, draw a 10-px white line segment through each source in + ``dia_source_unfiltered`` with ``ip_diffim_DipoleFit_classification`` + set, oriented along ``ip_diffim_DipoleFit_orientation`` (radians, + detector-frame). The length is fixed rather than measured because + ``ip_diffim_DipoleFit_separation`` is sub-pixel for the vast + majority of classified dipoles; the line is a fixed-size marker + of orientation rather than a physical extent. Sourced from the + unfiltered catalog rather than the standardized or APDB catalogs + because ``filterDiaSource`` removes dipoles before those stages. + show_trail_geometry : `bool`, optional + If True, draw a magenta line segment along + ``ext_trailedSources_Naive_angle`` with length + ``ext_trailedSources_Naive_length`` for every source in + ``dia_source_unfiltered`` whose trail length exceeds 3 px. + Long-trailed sources removed by ``filterDiaSource`` still show + up separately under ``show_trailed`` as ``x`` markers. + line_length_scale : `float`, optional + Multiplicative factor applied to the drawn length of the trail + line segments *after* the 3 px trail filter, so a below-threshold + trail stays hidden regardless of the scale. Does not affect the + dipole marker, which is drawn at a fixed 10 px length by design. + Default 1.0 (draw trails at their measured length); use larger + values to make short trails easier to see against the image. + label_size : `int`, optional + Text size (in pixels) for the reliability score and solar-system + designation annotations. + color_by : sequence of `str`, optional + Flag column names from ``dia_source_unfiltered``. When supplied, + the unfiltered-candidate overlay is split into buckets colored by + which named flag fires first (list order = color *and* priority), + with a residual white bucket for rows that match none of them. + Unknown column names are silently skipped. Example:: - Notes - ----- - There are some unused variables in here that could be made useable with - boolean kwargs to define what is being displayed. + color_by=["pixelFlags_bad", "pixelFlags_edge", + "ip_diffim_DipoleFit_classification", + "pixelFlags_saturated"] + mask_transparency : `int` or `None`, optional + Mask-plane transparency forwarded to the display (0 = opaque, + 100 = fully transparent). Pass ``None`` to leave the backend's + current setting untouched. + strip_metadata : `bool`, optional + Drop ``LTV1``/``LTV2`` keywords from each exposure's metadata + before sending to the backend. Needed for ds9 to align frames. + skymap : `lsst.skymap.BaseSkyMap`, optional + If supplied, overlay the boundaries of every tract/patch that + touches each frame, labeled ``tract,patch``. + skymap_ctype : `str`, optional + Display color for the tract/patch outlines and labels. + skymap_label_size : `float`, optional + Text size for the ``tract,patch`` labels. + image_datasets : `dict` [`str`, `str`], optional + Mapping from image-type key (``"science"``, ``"template"``, + ``"difference"``) to butler dataset name. Override to point at + alternate dataset types. + use_fakes : `bool`, optional + If True, load the fake-source-injected versions of the science + and template images: ``fakes_`` prefix on the science dataset + and ``injectedTemplate_`` prefix on the template dataset. The + difference image and every catalog keep their non-prefixed + names, per the fake-source pipeline's output convention. + Default False. """ - diffim = butler.get("difference_image", visit=visit, detector=detector) - science = butler.get("preliminary_visit_image", visit=visit, detector=detector) - template = butler.get("template_detector", visit=visit, detector=detector) - images = {"difference": diffim, "science": science, "template": template} - # red - unfiltered = butler.get("dia_source_unfiltered", visit=visit, detector=detector) - rejected = butler.get("rejected_dia_source", visit=visit, detector=detector) - trailed = butler.get("long_trailed_source_detector", visit=visit, detector=detector) - # yellow - candidate = butler.get("dia_source_unstandardized", visit=visit, detector=detector) - standardized = butler.get("dia_source_detector", visit=visit, detector=detector) # noqa - - dia_source = butler.get("dia_source_apdb", visit=visit, detector=detector) - good = dia_source['reliability'] > 0.1 - # blue - good_source = dia_source[good] - # red - bad_source = dia_source[~good] - print(f"{len(unfiltered)} unfiltered") - print(f"{len(trailed)} trailed") - print(f"{len(candidate)} candidate") - print(f"{len(bad_source)} low reliability diaSources") - print(f"{len(good_source)} good diaSources") - - marginal = butler.get("marginal_new_dia_source", visit=visit, detector=detector) # noqa - ss_source_detector = butler.get("ss_source_detector", visit=visit, detector=detector) # noqa - sky_source = unfiltered["sky_source"] - - rejected = get_xy_from_source_table(rejected, diffim.wcs) - candidate = get_xy_from_source_table(candidate, diffim.wcs) - unfiltered = get_xy_from_source_table(unfiltered[~sky_source], diffim.wcs) - trailed = get_xy_from_source_table(trailed, diffim.wcs) - display = lsst.afw.display.Display(backend=backend) - for frame, label in enumerate(("science", "template", "difference")): - display.frame = frame - display.image(images[label], title=label) - with display.Buffering(): - for x, y in zip(unfiltered["x"].data, unfiltered["y"].data): - display.dot("+", x, y, size=10, ctype="red") - for x, y in zip(trailed["x"].data, trailed["y"].data): - display.dot("x", x, y, size=30, ctype="red") - for x, y in zip(candidate["x"].data, candidate["y"].data): - display.dot("+", x, y, size=10, ctype="yellow") - for x, y in zip(dia_source["x"].data, dia_source["y"].data): - display.dot("+", x, y, size=10, ctype="blue") - for x, y in zip(good_source["x"].data, good_source["y"].data): - display.dot("o", x, y, size=10, ctype="blue") - for x, y in zip(bad_source["x"].data, bad_source["y"].data): - display.dot("o", x, y, size=10, ctype="red") + data_id = {"visit": visit, "detector": detector} + image_datasets = _apply_fakes_prefix(image_datasets, use_fakes) + + diffim = butler.get(image_datasets["difference"], data_id) + science = butler.get(image_datasets["science"], data_id) + template = butler.get(image_datasets["template"], data_id) + template = template[science.getBBox()] + if strip_metadata: + _strip_ds9_metadata(science, diffim, template) + images = {"science": science, "template": template, "difference": diffim} + + (overlays, reliability_labels, solar_system_labels, + dipole_segments, trail_segments) = _collect_overlays( + butler, data_id, diffim.wcs, + reliability_threshold=reliability_threshold, + show_unfiltered=show_unfiltered, + show_trailed=show_trailed, show_rejected=show_rejected, + show_standardized=show_standardized, + show_marginal=show_marginal, + show_kernel_sources=show_kernel_sources, + show_solar_system=show_solar_system, + show_apdb=show_apdb, + show_reliability_labels=show_reliability_labels, + show_dipoles=show_dipoles, + show_trail_geometry=show_trail_geometry, + line_length_scale=line_length_scale, + color_by=color_by, + ) + _print_overlay_legend( + overlays, f"visit={visit}, detector={detector} -- overlay legend:") + + afw_display = lsst.afw.display.Display(backend=backend) + if mask_transparency is not None: + afw_display.setMaskTransparency(mask_transparency) + for frame, image_name in enumerate(("science", "template", "difference")): + afw_display.frame = frame + # Wipe any markers left over from a previous call — `image()` + # only replaces the pixel data, region overlays persist otherwise. + _erase_current_frame_regions(afw_display) + image = images[image_name] + afw_display.image(image, title=image_name) + _draw_overlays_on_current_frame( + afw_display, overlays, reliability_labels, solar_system_labels, + dipole_segments=dipole_segments, + trail_segments=trail_segments, + label_size=label_size) + if skymap is not None: + draw_skymap_outlines_afw(afw_display, skymap, image.wcs, image.getBBox(), + ctype=skymap_ctype, label_size=skymap_label_size) try: - display.alignImages(match_type="Pixel") + afw_display.alignImages(match_type="Pixel") except NotImplementedError: - print(f"WARNING: cannot automatically align and lock images with backend={backend}!") + print(f"WARNING: cannot automatically align and lock images with backend={backend!r}.") + + +def display_images_ab(butler_a, butler_b, visit, detector, *, + image_type="difference", + labels=("A", "B"), + backend="firefly", + reliability_threshold=0.1, + show_unfiltered=True, + show_trailed=True, + show_rejected=True, + show_standardized=True, + show_marginal=True, + show_kernel_sources=True, + show_solar_system=True, + show_apdb=True, + show_reliability_labels=True, + show_dipoles=True, + show_trail_geometry=True, + line_length_scale=1.0, + label_size=3, + color_by=None, + mask_transparency=80, + strip_metadata=True, + skymap=None, + skymap_ctype="green", + skymap_label_size=1.5, + image_datasets=_IMAGE_DATASETS, + use_fakes=False): + """Display one image type side-by-side from two butlers, with overlays. + + Loads the same (visit, detector) from ``butler_a`` and ``butler_b``, + places them in two frames, and draws each butler's catalog overlays on + its own frame. Intended for A/B-testing pipeline-config changes that affect + detection or subtraction quality. + + Parameters + ---------- + butler_a, butler_b : `lsst.daf.butler.Butler` + Two butlers, typically from different pipeline runs of the same data. + visit, detector : `int` + Visit and detector ids to load data for. + image_type : {"science", "template", "difference"}, optional + Which image dataset to compare. Default ``"difference"``. + labels : pair of `str`, optional + Short tags for the two frames; appear in the image title and the + legend header. Default ``("A", "B")``. + backend : `str`, optional + afw display backend (typically "firefly" or "ds9"). + reliability_threshold, show_unfiltered, show_trailed, show_rejected, + show_standardized, show_marginal, show_kernel_sources, + show_solar_system, show_apdb, show_reliability_labels, show_dipoles, + show_trail_geometry, line_length_scale, label_size, color_by, + mask_transparency, strip_metadata, skymap, skymap_ctype, + skymap_label_size, image_datasets, use_fakes + Same meaning as in `display_images`. Applied to both frames; the + tract/patch overlay uses each frame's own exposure WCS. + """ + if image_type not in image_datasets: + raise ValueError( + f"image_type must be one of {sorted(image_datasets)}, got {image_type!r}") + image_datasets = _apply_fakes_prefix(image_datasets, use_fakes) + dataset = image_datasets[image_type] + data_id = {"visit": visit, "detector": detector} + + image_a = butler_a.get(dataset, data_id) + image_b = butler_b.get(dataset, data_id) + if image_type == "template": + # Templates are usually larger than the science footprint; clip them + # to the science bbox so the two frames have matching extents. + sci_a = butler_a.get(image_datasets["science"], data_id) + sci_b = butler_b.get(image_datasets["science"], data_id) + image_a = image_a[sci_a.getBBox()] + image_b = image_b[sci_b.getBBox()] + if strip_metadata: + _strip_ds9_metadata(image_a, image_b) + + common = dict( + reliability_threshold=reliability_threshold, + show_unfiltered=show_unfiltered, + show_trailed=show_trailed, show_rejected=show_rejected, + show_standardized=show_standardized, + show_marginal=show_marginal, + show_kernel_sources=show_kernel_sources, + show_solar_system=show_solar_system, + show_apdb=show_apdb, show_reliability_labels=show_reliability_labels, + show_dipoles=show_dipoles, + show_trail_geometry=show_trail_geometry, + line_length_scale=line_length_scale, + color_by=color_by, + ) + overlays_a, rel_a, ss_a, dip_a, tr_a = _collect_overlays( + butler_a, data_id, image_a.wcs, **common) + overlays_b, rel_b, ss_b, dip_b, tr_b = _collect_overlays( + butler_b, data_id, image_b.wcs, **common) + + label_a, label_b = labels + print(f"visit={visit}, detector={detector}: A/B comparison of {image_type!r}") + _print_overlay_legend(overlays_a, f"-- {label_a} overlay legend:", indent=" ") + _print_overlay_legend(overlays_b, f"-- {label_b} overlay legend:", indent=" ") + + afw_display = lsst.afw.display.Display(backend=backend) + if mask_transparency is not None: + afw_display.setMaskTransparency(mask_transparency) + for frame, (tag, image, overlays, rel, ss, dip, tr) in enumerate(( + (label_a, image_a, overlays_a, rel_a, ss_a, dip_a, tr_a), + (label_b, image_b, overlays_b, rel_b, ss_b, dip_b, tr_b))): + afw_display.frame = frame + # Wipe any markers left over from a previous call — `image()` + # only replaces the pixel data, region overlays persist otherwise. + _erase_current_frame_regions(afw_display) + afw_display.image(image, title=f"{image_type} ({tag})") + _draw_overlays_on_current_frame(afw_display, overlays, rel, ss, + dipole_segments=dip, + trail_segments=tr, + label_size=label_size) + if skymap is not None: + draw_skymap_outlines_afw(afw_display, skymap, image.wcs, image.getBBox(), + ctype=skymap_ctype, label_size=skymap_label_size) + + try: + afw_display.alignImages(match_type="Pixel") + except NotImplementedError: + print(f"WARNING: cannot automatically align and lock images with backend={backend!r}.") + + +def _ensure_deblend_nchild(catalog): + """Return `catalog` with a ``deblend_nChild`` column guaranteed present. + + Firefly's footprint overlay (``createFootprintsTable``) unconditionally + reads ``deblend_nChild`` to classify each footprint's deblend family, + but the diaSource detection catalogs are not deblended and omit that + column. When it is missing, copy the catalog into an augmented schema + with a ``deblend_nChild`` field defaulting to 0 (so every footprint + classifies as "isolated"), preserving each record's `Footprint`. + Catalogs that already carry the column are returned unchanged. + """ + if "deblend_nChild" in catalog.schema.getNames(): + return catalog + mapper = lsst.afw.table.SchemaMapper(catalog.schema) + mapper.addMinimalSchema(catalog.schema, True) + mapper.editOutputSchema().addField( + "deblend_nChild", type="I", + doc="Placeholder 0; source catalog was not deblended.") + augmented = lsst.afw.table.SourceCatalog(mapper.getOutputSchema()) + augmented.reserve(len(catalog)) + for record in catalog: + new_record = augmented.addNew() + new_record.assign(record, mapper) + new_record.setFootprint(record.getFootprint()) + return augmented + + +def _subset_catalog(catalog, indices): + """Build a `~lsst.afw.table.SourceCatalog` of the rows at `indices`. + + The subset shares the input's table and appends the existing records + (shallow), so each row keeps its attached `Footprint` rather than a + copy. Used to split a catalog into one sub-catalog per overlay color. + """ + subset = lsst.afw.table.SourceCatalog(catalog.table) + for i in indices: + subset.append(catalog[i]) + return subset + + +def _footprint_adjacency(bboxes): + """Adjacency lists for footprints whose bounding boxes overlap. + + Two footprints are adjacent when their bounding boxes overlap (share + at least one pixel). Bounding-box overlap can slightly over-count + versus true pixel touching, which only makes the coloring more + conservative (never assigns the same color to two touching footprints). + + Parameters + ---------- + bboxes : `list` [`lsst.geom.Box2I`] + Footprint bounding boxes, in the order the catalog iterates. + + Returns + ------- + adjacency : `list` [`set` [`int`]] + ``adjacency[i]`` is the set of indices whose footprints touch + footprint ``i``. + """ + n = len(bboxes) + adjacency = [set() for _ in range(n)] + for i in range(n): + for j in range(i + 1, n): + if bboxes[i].overlaps(bboxes[j]): + adjacency[i].add(j) + adjacency[j].add(i) + return adjacency + + +def _greedy_color(adjacency, n_colors, rng=None): + """Greedily color a graph so touching nodes differ, using `n_colors`. + + Nodes are colored highest-degree first (Welsh--Powell), and each node + takes a color chosen *at random* from those not already used by an + adjacent node. Randomizing the choice -- rather than always taking the + lowest free index -- spreads the coloring across the whole palette and + varies it between runs, while still guaranteeing adjacent footprints + differ. Footprint adjacency graphs are essentially planar, so with 12 + colors available a conflict-free coloring is found in practice. In the + pathological case where a node's neighbors already occupy all + `n_colors`, it falls back to a color least used among those neighbors + (ties broken randomly) rather than failing. + + Parameters + ---------- + adjacency : `list` [`set` [`int`]] + Adjacency lists from `_footprint_adjacency`. + n_colors : `int` + Number of available colors (palette length). + rng : `random.Random`, optional + Random source, for reproducible colorings in tests. Defaults to a + fresh unseeded `random.Random`. + + Returns + ------- + colors : `list` [`int`] + ``colors[i]`` is the palette index assigned to node ``i``. + """ + if rng is None: + rng = random.Random() + n = len(adjacency) + colors = [-1] * n + order = sorted(range(n), key=lambda i: len(adjacency[i]), reverse=True) + for node in order: + used = {colors[nbr] for nbr in adjacency[node] if colors[nbr] >= 0} + available = [c for c in range(n_colors) if c not in used] + if available: + chosen = rng.choice(available) + else: + # Every color is taken by a neighbor (needs >n_colors mutually + # touching footprints -- essentially never). Reuse a color that + # appears least among the neighbors, breaking ties randomly. + counts = [0] * n_colors + for nbr in adjacency[node]: + if colors[nbr] >= 0: + counts[colors[nbr]] += 1 + fewest = min(counts) + chosen = rng.choice( + [c for c in range(n_colors) if counts[c] == fewest]) + colors[node] = chosen + return colors + + +def display_footprints(butler=None, visit=None, detector=None, + backend="firefly", *, + exposure=None, catalog=None, + image_type="difference", + catalog_dataset="dia_source_unfiltered", + style="outline", + palette=_OBJECT_PALETTE, + frame=0, + mask_transparency=80, + strip_metadata=True, + image_datasets=_IMAGE_DATASETS): + """Overlay diaSource footprints on an exposure in Firefly, color-cycled + so that touching footprints get distinct colors. + + The footprints come from an afw `~lsst.afw.table.SourceCatalog` (the + diffim detection output, which still carries per-source `Footprint`\\ s + -- the transformed and APDB diaSource tables are DataFrames with the + footprints stripped). Supply the data one of two ways: + + * pass ``butler`` plus ``visit`` and ``detector`` to load the + exposure (``image_datasets[image_type]``) and catalog + (``catalog_dataset``) from the butler; or + * pass ``exposure`` and ``catalog`` directly, skipping the butler. + + The footprints are then drawn on a single Firefly frame using the + backend's native footprint overlay. + + Each footprint is assigned one of the `palette` colors by greedy + graph coloring over a bounding-box-touch adjacency graph, so no two + touching footprints share a color (see `_footprint_adjacency` and + `_greedy_color`). The color chosen for each footprint is randomized + among those its neighbors are not using, so the palette is spread + across the frame and re-running produces a different coloring. Because + Firefly's ``overlayFootprints`` takes a single color per call, the + catalog is split into one sub-catalog per color and each is overlaid as + its own Firefly layer. + + Re-running on the same frame overwrites each color layer in place. + Color layers left over from a previous run that used *more* colors are + not cleared automatically. + + Parameters + ---------- + butler : `lsst.daf.butler.Butler`, optional + Butler to load the exposure and catalog from. Required (with + ``visit`` and ``detector``) unless ``exposure`` and ``catalog`` are + given directly. + visit, detector : `int`, optional + Visit and detector ids to load data for. Required with ``butler``. + backend : `str`, optional + afw display backend. Only ``"firefly"`` is supported, since the + overlay uses Firefly's native footprint rendering. + exposure : `lsst.afw.image.Exposure`, optional + Exposure to draw on, supplied directly instead of via the butler. + Must be given together with ``catalog``; when set, ``butler``, + ``visit``, ``detector``, ``catalog_dataset``, ``image_type``, and + ``image_datasets`` are all ignored. + catalog : `lsst.afw.table.SourceCatalog`, optional + Footprint-bearing source catalog, supplied directly instead of via + the butler. Must be given together with ``exposure``. + image_type : {"science", "template", "difference"}, optional + Which image to display the footprints on (butler mode only). + Default ``"difference"``. + catalog_dataset : `str`, optional + Butler dataset of the footprint-bearing afw source catalog (butler + mode only). Default ``"dia_source_unfiltered"`` (the pre-filter + detection catalog, which still carries footprints; the + transformed/standardized diaSource tables have them stripped). + style : {"outline", "fill"}, optional + Footprint rendering style. ``"outline"`` (default) keeps the + color coding legible where footprints overlap; ``"fill"`` shades + the interior. + palette : sequence of `str`, optional + Colors cycled across footprints. Defaults to the 12-color + ``_OBJECT_PALETTE`` also used by the cutout plotters. + frame : `int`, optional + Display frame to draw the image and footprints in. Default ``0``. + mask_transparency : `int` or `None`, optional + Mask-plane transparency forwarded to the display (0 = opaque, + 100 = fully transparent). Pass ``None`` to leave it untouched. + strip_metadata : `bool`, optional + Drop ``LTV1``/``LTV2`` keywords from the exposure metadata before + sending to the backend. + image_datasets : `dict` [`str`, `str`], optional + Mapping from image-type key to butler dataset name. + """ + if backend != "firefly": + raise ValueError( + f"display_footprints only supports the 'firefly' backend " + f"(needs Firefly's native footprint overlay); got {backend!r}") + if style not in ("outline", "fill"): + raise ValueError(f"style must be 'outline' or 'fill', got {style!r}") + + # Two input modes: direct (exposure + catalog) or butler-loaded. + direct = exposure is not None or catalog is not None + if direct: + if exposure is None or catalog is None: + raise ValueError( + "supply BOTH exposure and catalog to draw directly") + title = "footprints" + location = "" + else: + if butler is None or visit is None or detector is None: + raise ValueError( + "supply either (butler, visit, detector) or " + "(exposure, catalog)") + if image_type not in image_datasets: + raise ValueError( + f"image_type must be one of {sorted(image_datasets)}, " + f"got {image_type!r}") + data_id = {"visit": visit, "detector": detector} + exposure = butler.get(image_datasets[image_type], data_id) + catalog = butler.get(catalog_dataset, data_id) + title = f"{image_type} footprints" + location = f"visit={visit}, detector={detector}: " + + if strip_metadata: + _strip_ds9_metadata(exposure) + if not isinstance(catalog, lsst.afw.table.SourceCatalog): + raise TypeError( + f"catalog is a {type(catalog).__name__}, not an afw " + "SourceCatalog. Footprints are only carried by the afw detection " + "catalog (storageClass 'SourceCatalog'); the " + "transformed/standardized diaSource tables (DataFrame or " + "ArrowAstropy, e.g. 'dia_source_detector') have them stripped.") + if len(catalog) > 0 and catalog[0].getFootprint() is None: + raise ValueError( + "catalog is an afw SourceCatalog but its records have no " + "Footprint attached, so there is nothing to draw.") + catalog = _ensure_deblend_nchild(catalog) + + # Color the footprints so touching ones differ, then group indices by + # assigned color for the per-color Firefly overlay calls below. + bboxes = [record.getFootprint().getBBox() for record in catalog] + color_indices = _greedy_color(_footprint_adjacency(bboxes), + len(palette)) + groups = {} + for i, c in enumerate(color_indices): + groups.setdefault(c, []).append(i) + + afw_display = lsst.afw.display.Display(backend=backend) + if mask_transparency is not None: + afw_display.setMaskTransparency(mask_transparency) + afw_display.frame = frame + # image() only replaces pixel data; wipe stale region markers first. + _erase_current_frame_regions(afw_display) + afw_display.image(exposure, title=title) + + print(f"{location}{len(catalog)} footprints in {len(groups)} colors") + # overlayFootprints is a Firefly-impl method reached via Display's + # attribute delegation, the same way display_images calls alignImages. + for c in sorted(groups): + indices = groups[c] + color = palette[c] + subset = _subset_catalog(catalog, indices) + afw_display.overlayFootprints( + subset, color=color, style=style, + layerString=f"footprints c{c} ", + titleString=f"footprints c{c} ") + + +def extract_timestamped_messages(log: str | dict[str, Any]) -> str: + """ + Extract records[*].(asctime, message) from an LSST-style JSON log and + format as: + 2026-02-25T04:15:35.092108Z Preparing execution... + one per line. + + Parameters + ---------- + log: + Either the JSON text (str) or a parsed dict. + sort: + If True, sort by asctime (robust if log fragments are concatenated). + + Returns + ------- + str + Joined lines. + """ + if isinstance(log, str): + s = log.strip() + + # Handle the case where the *JSON itself* is wrapped in quotes, like: + # '"{...}"' or "'{...}'" + if (len(s) >= 2) and (s[0] == s[-1]) and s[0] in ("'", '"'): + s = s[1:-1] + + try: + obj = json.loads(s) + except json.JSONDecodeError: + # One more attempt: sometimes a quoted-JSON string is itself + # JSON-encoded e.g. "\"{...}\"" + obj = json.loads(json.loads(s)) + else: + obj = log + + records = obj.get("records", []) + if not isinstance(records, list): + raise TypeError("Expected obj['records'] to be a list.") + + rows: list[tuple[datetime, str, str]] = [] + for rec in records: + if not isinstance(rec, dict): + continue + ts = rec.get("asctime") + msg = rec.get("message") + if not ts or msg is None: + continue + + # Parse ISO-8601 with trailing "Z" + dt = datetime.fromisoformat(ts.replace("Z", "+00:00")).astimezone(timezone.utc) + rows.append((dt, ts, str(msg))) + + return "\n".join(f"{ts} {msg}" for _, ts, msg in rows) diff --git a/python/lsst/analysis/ap/plotDiaSourceLightcurve.py b/python/lsst/analysis/ap/plotDiaSourceLightcurve.py new file mode 100644 index 0000000..a40fffd --- /dev/null +++ b/python/lsst/analysis/ap/plotDiaSourceLightcurve.py @@ -0,0 +1,401 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Render diaSource cutouts with a per-diaObject lightcurve panel. + +``PlotDiaSourceLightcurveTask`` extends ``PlotImageSubtractionCutoutsTask`` +by adding a lightcurve panel below the science/template/difference cutouts. +The panel shows the diaSource history for the associated diaObject and +overlays diaForcedSource measurements for any visit that has a forced +measurement but no diaSource (drawn with a distinct marker). +""" + +__all__ = ["PlotDiaSourceLightcurveConfig", "PlotDiaSourceLightcurveTask"] + +import argparse +import io +import logging +import os + +import lsst.daf.butler +import lsst.pex.config as pexConfig + +import pandas as pd + +from . import apdb as _apdb_mod +from . import plotUtils +from .plotImageSubtractionCutouts import ( + PlotImageSubtractionCutoutsConfig, + PlotImageSubtractionCutoutsTask, + _annotate_image, +) + +_log = logging.getLogger(__name__) + + +class PlotDiaSourceLightcurveConfig(PlotImageSubtractionCutoutsConfig): + lightcurve_height = pexConfig.Field( + doc="Height in inches reserved for the lightcurve panel below the cutouts.", + dtype=float, + default=2.5, + ) + lightcurve_exclude_flagged = pexConfig.Field( + doc="Pass exclude_flagged=True to the APDB query when loading " + "diaSources for the lightcurve. Defaults to False so the " + "lightcurve matches the row count of a direct APDB query; " + "set True to drop diaSources matching the configured bad-flag " + "list. DiaForcedSources are always loaded unfiltered.", + dtype=bool, + default=False, + ) + lightcurve_marker_source = pexConfig.Field( + doc="Matplotlib marker style for visits that have a diaSource " + "detection.", + dtype=str, + default="o", + ) + lightcurve_marker_forced_only = pexConfig.Field( + doc="Matplotlib marker style for visits that have a forced " + "measurement but no diaSource.", + dtype=str, + default="v", + ) + highlight_current_source = pexConfig.Field( + doc="Draw a vertical line and open ring on the lightcurve at the " + "MJD/flux of the diaSource being cut out.", + dtype=bool, + default=True, + ) + + +class PlotDiaSourceLightcurveTask(PlotImageSubtractionCutoutsTask): + """Generate cutouts plus a diaObject lightcurve panel for each diaSource. + + Parameters + ---------- + output_path : `str` + Path to write outputs to. Same convention as the parent task. + apdb_query : `lsst.analysis.ap.apdb.DbQuery`, optional + Query handle used to load the diaSource and diaForcedSource history + for each diaObject. If None, the lightcurve panel is rendered with + a placeholder message. + + Notes + ----- + The input ``data`` DataFrame must include the fields required by the + parent (``ra, dec, diaSourceId, detector, visit, instrument``) plus + ``diaObjectId`` (to load the lightcurve) and ``midpointMjdTai`` (to + highlight the current source on the lightcurve). + """ + ConfigClass = PlotDiaSourceLightcurveConfig + _DefaultName = "plotDiaSourceLightcurve" + + def __init__(self, *, output_path, apdb_query=None, **kwargs): + super().__init__(output_path=output_path, **kwargs) + self._apdb_query = apdb_query + # Per-diaObject lightcurve cache so adjacent diaSources on the same + # object don't re-query the APDB. Keyed by diaObjectId. + self._lightcurve_cache = {} + + def _reduce_kwargs(self): + kwargs = super()._reduce_kwargs() + kwargs["apdb_query"] = self._apdb_query + return kwargs + + def run(self, data, butler, njobs=0): + if njobs > 0: + self.log.warning("njobs=%d ignored; PlotDiaSourceLightcurveTask " + "runs single-process only.", njobs) + return super().run(data, butler, njobs=0) + + def write_images(self, data, butler, njobs=0): + if njobs > 0: + self.log.warning("njobs=%d ignored; PlotDiaSourceLightcurveTask " + "runs single-process only.", njobs) + return super().write_images(data, butler, njobs=0) + + def _load_lightcurve_data(self, dia_object_id): + """Return cached (sources, forced) for one diaObject. + + Returns ``(None, None)`` if no APDB query handle is configured. + """ + if self._apdb_query is None: + return None, None + if dia_object_id not in self._lightcurve_cache: + sources = self._apdb_query.load_sources_for_object( + dia_object_id, + exclude_flagged=self.config.lightcurve_exclude_flagged, + ) + # Always show all of the forced sources regardless of flags. + forced = self._apdb_query.load_forced_sources_for_object( + dia_object_id, + ) + self._lightcurve_cache[dia_object_id] = (sources, forced) + return self._lightcurve_cache[dia_object_id] + + def _plot_cutout(self, science, template, difference, scale, sizes, source=None): + import astropy.visualization as aviz + import matplotlib + matplotlib.use("AGG") + matplotlib.rcParams.update(matplotlib.rcParamsDefault) + import matplotlib.pyplot as plt + from matplotlib import cm + from matplotlib.gridspec import GridSpec + + len_sizes = len(sizes) + + sources_lc, forced_lc = (None, None) + if source is not None and "diaObjectId" in source.dtype.names: + dia_object_id = source["diaObjectId"] + # diaObjectId is NaN for diaSources not associated with any + # diaObject (e.g. single unassociated detections). Skip the + # lightcurve query — the panel will render its empty placeholder. + if not pd.isna(dia_object_id): + try: + sources_lc, forced_lc = self._load_lightcurve_data(int(dia_object_id)) + except Exception as e: + self.log.warning("Failed to load lightcurve for diaObjectId=%s: %s. " + "The DiaSource is likely unassociated.", + dia_object_id, e) + + cutout_height_in = max(1.7, 1.7 * len_sizes) + lc_height_in = float(self.config.lightcurve_height) + fig_height = cutout_height_in + lc_height_in + fig = plt.figure(figsize=(7, fig_height), constrained_layout=True) + + gs = GridSpec(2, 1, height_ratios=[cutout_height_in, lc_height_in], figure=fig) + cutout_gs = gs[0].subgridspec(len_sizes, 3) + cutout_axes = [[fig.add_subplot(cutout_gs[r, c]) for c in range(3)] + for r in range(len_sizes)] + lc_ax = fig.add_subplot(gs[1]) + + def plot_one_image(ax, data, size, name=None): + if name == "Difference": + norm = aviz.ImageNormalize( + data[data.shape[0] // 2 - 7:data.shape[0] // 2 + 8, + data.shape[1] // 2 - 7:data.shape[1] // 2 + 8], + interval=aviz.MinMaxInterval(), + stretch=aviz.AsinhStretch(a=0.1), + ) + else: + norm = aviz.ImageNormalize( + data, + interval=aviz.MinMaxInterval(), + stretch=aviz.AsinhStretch(a=0.1), + ) + ax.imshow(data, cmap=cm.bone, interpolation="none", norm=norm, + extent=(0, size, 0, size), origin="lower", aspect="equal") + x_line = 1 + y_line = 1 + ax.plot((x_line, x_line + 1.0/scale), (y_line, y_line), color="blue", lw=6) + ax.plot((x_line, x_line + 1.0/scale), (y_line, y_line), color="yellow", lw=2) + ax.axis("off") + if name is not None: + ax.set_title(name) + + try: + plot_one_image(cutout_axes[0][0], template[0].image.array, sizes[0], "Template") + plot_one_image(cutout_axes[0][1], science[0].image.array, sizes[0], "Science") + plot_one_image(cutout_axes[0][2], difference[0].image.array, sizes[0], "Difference") + for i in range(1, len_sizes): + plot_one_image(cutout_axes[i][0], template[i].image.array, sizes[i], None) + plot_one_image(cutout_axes[i][1], science[i].image.array, sizes[i], None) + plot_one_image(cutout_axes[i][2], difference[i].image.array, sizes[i], None) + + self._draw_lightcurve(lc_ax, sources_lc, forced_lc, current_source=source) + + if source is not None and self.config.add_metadata: + # Place metadata text above the figure top, matching the + # multi-size layout in the parent class. + # ``bbox_inches="tight"`` in savefig expands the saved area to + # include them. + _annotate_image(fig, source, len_sizes, + heights=[1.2, 1.15, 1.1, 1.05, 1.0]) + + output = io.BytesIO() + plt.savefig(output, bbox_inches="tight", format="png") + output.seek(0) + finally: + plt.close(fig) + return output + + def _draw_lightcurve(self, ax, sources, forced, current_source=None): + """Draw the lightcurve panel for a diaObject. + + Parameters + ---------- + ax : `matplotlib.axes.Axes` + Axes to draw into. + sources : `pandas.DataFrame` or None + DiaSources for this diaObject, or None if no APDB is configured. + forced : `pandas.DataFrame` or None + DiaForcedSources for this diaObject. Rows whose ``visit`` matches + one in ``sources`` are suppressed; the rest are drawn with the + ``lightcurve_marker_forced_only`` marker. + current_source : `numpy.record`, optional + The diaSource being cut out. If non-None, a vertical line and + open ring mark its MJD/psfFlux on the panel. + """ + if sources is None and forced is None: + ax.text(0.5, 0.5, "no APDB query configured", + ha="center", va="center", transform=ax.transAxes) + ax.set_xticks([]) + ax.set_yticks([]) + return + + n_src = 0 if sources is None else len(sources) + n_forced = 0 if forced is None else len(forced) + if n_src == 0 and n_forced == 0: + ax.text(0.5, 0.5, "no lightcurve data", + ha="center", va="center", transform=ax.transAxes) + ax.set_xticks([]) + ax.set_yticks([]) + return + + # DiaSources are point-like (with the exception of moving objects, + # which appear only once), so a single visit gives at most one + # diaSource per diaObject — dedup on ``visit`` is safe. + if n_forced and n_src: + forced_only = forced[~forced["visit"].isin(sources["visit"])] + elif n_forced: + forced_only = forced + else: + forced_only = None + + def _plot_group(group, color, marker, label): + time_col = plotUtils._time_column(group) + if "psfFluxErr" in group.columns and group["psfFluxErr"].notna().any(): + ax.errorbar(group[time_col], group["psfFlux"], + yerr=group["psfFluxErr"], fmt=marker, + color=color, label=label) + else: + ax.plot(group[time_col], group["psfFlux"], marker, + color=color, label=label) + + if n_src: + for band, group in sources.groupby("band"): + color = plotUtils.BAND_COLORS.get(band, "k") + _plot_group(group, color, + self.config.lightcurve_marker_source, + f"{band} (n={len(group)})") + + if forced_only is not None and len(forced_only): + # Plot per-band forced points with their band colors, but + # suppress them from the legend so we only emit one combined + # entry (in black) for the forced marker regardless of how many + # bands are present. + for band, group in forced_only.groupby("band"): + color = plotUtils.BAND_COLORS.get(band, "k") + _plot_group(group, color, + self.config.lightcurve_marker_forced_only, + "_nolegend_") + ax.plot([], [], self.config.lightcurve_marker_forced_only, + color="black", + label=f"forced (n={len(forced_only)})") + + if self.config.highlight_current_source and current_source is not None: + try: + x = float(current_source["midpointMjdTai"]) + y = float(current_source["psfFlux"]) + ax.axvline(x, color="grey", lw=0.5, ls="--") + ax.plot([x], [y], "o", markerfacecolor="none", + markeredgecolor="red", markersize=12, + markeredgewidth=1.5) + except (KeyError, ValueError): + pass + + ax.axhline(0, color="grey", lw=0.5) + ax.set_xlabel("MJD (TAI)") + ax.set_ylabel("psfFlux (nJy)") + ax.legend(frameon=True, fontsize=7, loc="best") + + +def _make_apdbQuery(sqlitefile=None, postgres_url=None, namespace=None): + """Return a query connection to the specified APDB.""" + if sqlitefile is not None: + return _apdb_mod.ApdbSqliteQuery(sqlitefile) + if postgres_url is not None and namespace is not None: + return _apdb_mod.ApdbPostgresQuery(namespace, postgres_url) + raise RuntimeError("Cannot handle database connection args: " + f"sqlitefile={sqlitefile}, postgres_url={postgres_url}, " + f"namespace={namespace}") + + +def build_argparser(): + """Argument parser for the ``plotDiaSourceLightcurve`` command.""" + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog="More information is available at https://pipelines.lsst.io.", + ) + apdbArgs = parser.add_mutually_exclusive_group(required=True) + apdbArgs.add_argument("--sqlitefile", default=None, + help="Path to sqlite APDB file.") + apdbArgs.add_argument("--namespace", default=None, + help="Postgres namespace (schema) to connect to.") + parser.add_argument( + "--postgres_url", + default="rubin@usdf-prompt-processing-dev.slac.stanford.edu/lsst-devl", + help="Postgres connection path.") + parser.add_argument("--limit", default=5, type=int, + help="Number of sources to load (default=5).") + parser.add_argument("-C", "--configFile", + help="PlotDiaSourceLightcurveConfig file to load.") + parser.add_argument("--collections", nargs="*", + help="Butler collection(s) to load image data from.") + parser.add_argument("repo", help="Path to butler repository.") + parser.add_argument("outputPath", + help="Path to write images to (under outputPath/images/).") + parser.add_argument("--reliabilityMin", type=float, default=None, + help="Minimum reliability for diaSource selection.") + parser.add_argument("--reliabilityMax", type=float, default=None, + help="Maximum reliability for diaSource selection.") + return parser + + +def run_lightcurves(args): + """Run PlotDiaSourceLightcurveTask from parsed command-line arguments.""" + logging.basicConfig(level=logging.INFO, + format="{name} {levelname}: {message}", style="{") + + butler = lsst.daf.butler.Butler(args.repo, collections=args.collections) + apdb_query = _make_apdbQuery(sqlitefile=args.sqlitefile, + postgres_url=args.postgres_url, + namespace=args.namespace) + + config = PlotDiaSourceLightcurveConfig() + if args.configFile is not None: + config.load(os.path.expanduser(args.configFile)) + config.freeze() + task = PlotDiaSourceLightcurveTask(config=config, + output_path=args.outputPath, + apdb_query=apdb_query) + + data = next(apdb_query.iter_sources(args.limit, + args.reliabilityMin, + args.reliabilityMax)) + sources = task.run(data, butler) + print(f"Generated {len(sources)} diaSource lightcurve plots to {args.outputPath}.") + + +def main(): + args = build_argparser().parse_args() + run_lightcurves(args) diff --git a/python/lsst/analysis/ap/plotImageSubtractionCutouts.py b/python/lsst/analysis/ap/plotImageSubtractionCutouts.py index 0e72da8..9631563 100644 --- a/python/lsst/analysis/ap/plotImageSubtractionCutouts.py +++ b/python/lsst/analysis/ap/plotImageSubtractionCutouts.py @@ -43,10 +43,11 @@ import lsst.utils import numpy as np import pandas as pd -import sqlalchemy from . import apdb +_log = logging.getLogger(__name__) + class _ButlerCache: """Global class to handle butler queries, to allow lru_cache and @@ -96,9 +97,9 @@ def get_exposures(self, instrument, detector, visit): try: science = self._butler.get(self._config.science_image_type, data_id) except DatasetNotFoundError as e: - self.log.error(f"Cannot load {self._config.science_image_type} with data_id {data_id}: {e}") - self.log.error("If you are working with data processed earlier than May 2025, try setting " - "config.science_image_type = 'initial_pvi' or 'calexp'.") + _log.error(f"Cannot load {self._config.science_image_type} with data_id {data_id}: {e}") + _log.error("If you are working with data processed earlier than May 2025, try setting " + "config.science_image_type = 'initial_pvi' or 'calexp'.") raise if self._config.diff_image_type is not None: @@ -311,11 +312,12 @@ def write_images(self, data, butler, njobs=0): with multiprocessing.Pool(njobs) as pool: sources = pool.map(self._do_one_source, data.to_records(index=indexNotInColumns)) else: - for i, source in enumerate(data.to_records(index=indexNotInColumns)): - if not self.cutout_path.exists(source["diaSourceId"], - f'{source["diaSourceId"]}.png'): - id = self._do_one_source(source) - sources.append(id) + for source in data.to_records(index=indexNotInColumns): + src_id = source["diaSourceId"] + if self.cutout_path.exists(src_id, f"{src_id}.png"): + sources.append(src_id) + else: + sources.append(self._do_one_source(source)) # restore numpy error message state np.seterr(**seterr_dict) @@ -533,7 +535,43 @@ def plot_one_image(ax, data, size, name=None): return output -def _annotate_image(fig, source, len_sizes): +# Flag groupings for the metadata legend on cutout images. A row's label is +# colored red if any column in its group is set on the source. +_FLAG_GROUPS = { + "psf": ["psfFlux_flag", "psfFlux_flag_noGoodPixels", "psfFlux_flag_edge"], + "aperture": ["apFlux_flag", "apFlux_flag_apertureTruncated"], + "forced": ["forced_PsfFlux_flag", "forced_PsfFlux_flag_noGoodPixels", + "forced_PsfFlux_flag_edge"], + "edge": ["pixelFlags_edge"], + "interp": ["pixelFlags_interpolated", "pixelFlags_interpolatedCenter"], + "saturated": ["pixelFlags_saturated", "pixelFlags_saturatedCenter"], + "cr": ["pixelFlags_cr", "pixelFlags_crCenter"], + "bad": ["pixelFlags_bad"], + "suspect": ["pixelFlags_suspect", "pixelFlags_suspectCenter"], + "centroid": ["centroid_flag"], + "shape": ["shape_flag", "shape_flag_no_pixels", "shape_flag_not_contained", + "shape_flag_parent_source"], +} + +# Flag-tag overlays drawn on top of the flux rows (rows 2 and 3). Each entry +# is (predicate, x, label, color, row_index). The predicate is either a key +# into ``_FLAG_GROUPS`` (the tag is drawn if any column in the group is set) +# or a callable returning a bool. +_FLAG_TAGS = [ + ("edge", 0.55, "EDGE", "goldenrod", 2), + ("interp", 0.62, "INTERP", "green", 2), + ("saturated", 0.72, "SAT", "green", 2), + ("cr", 0.77, "CR", "magenta", 2), + ("bad", 0.81, "BAD", "red", 2), + (lambda src: bool(src["isDipole"]), 0.87, "DIPOLE", "indigo", 2), + ("suspect", 0.55, "SUS", "goldenrod", 3), + ("centroid", 0.60, "CENTROID", "red", 3), + ("shape", 0.73, "SHAPE", "red", 3), +] +# Future option: add two more flag flavors at x = 0.80 and 0.87 on row 3. + + +def _annotate_image(fig, source, len_sizes, heights=None): """Annotate the cutouts image with metadata and flags. Parameters @@ -544,101 +582,99 @@ def _annotate_image(fig, source, len_sizes): DiaSource record of the object being plotted. len_sizes : `int` Length of the ``size`` array set in configuration. + heights : `list` [`float`], optional + Five figure-fraction y-coordinates for the metadata rows. If None, + the default heights are chosen based on ``len_sizes``. Subclasses + that add extra panels to the figure can pass their own positions. """ - # Names of flags fields to add a flag label to the image, using any(). - flags_psf = ["psfFlux_flag", "psfFlux_flag_noGoodPixels", "psfFlux_flag_edge"] - flags_aperture = ["apFlux_flag", "apFlux_flag_apertureTruncated"] - flags_forced = ["forced_PsfFlux_flag", "forced_PsfFlux_flag_noGoodPixels", - "forced_PsfFlux_flag_edge"] - flags_edge = ["pixelFlags_edge"] - flags_interp = ["pixelFlags_interpolated", "pixelFlags_interpolatedCenter"] - flags_saturated = ["pixelFlags_saturated", "pixelFlags_saturatedCenter"] - flags_cr = ["pixelFlags_cr", "pixelFlags_crCenter"] - flags_bad = ["pixelFlags_bad"] - flags_suspect = ["pixelFlags_suspect", "pixelFlags_suspectCenter"] - flags_centroid = ["centroid_flag"] - flags_shape = ["shape_flag", "shape_flag_no_pixels", "shape_flag_not_contained", - "shape_flag_parent_source"] - flag_color = "red" text_color = "grey" - if len_sizes == 1: - heights = [0.95, 0.91, 0.87, 0.83, 0.79] - else: - heights = [1.2, 1.15, 1.1, 1.05, 1.0] - - # NOTE: fig.text coordinates are in fractions of the figure. - fig.text(0, heights[0], "diaSourceId:", color=text_color) - fig.text(0.145, heights[0], f"{source['diaSourceId']}") - fig.text(0.43, heights[0], f"{source['instrument']}", fontweight="bold") - fig.text(0.64, heights[0], "detector:", color=text_color) - fig.text(0.74, heights[0], f"{source['detector']}") - fig.text(0.795, heights[0], "visit:", color=text_color) - fig.text(0.85, heights[0], f"{source['visit']}") - fig.text(0.95, heights[0], f"{source['band']}") - - fig.text(0.0, heights[1], "ra:", color=text_color) - fig.text(0.037, heights[1], f"{source['ra']:.8f}") - fig.text(0.21, heights[1], "dec:", color=text_color) - fig.text(0.265, heights[1], f"{source['dec']:+.8f}") - fig.text(0.50, heights[1], "detection S/N:", color=text_color) - fig.text(0.66, heights[1], f"{source['snr']:6.1f}") - fig.text(0.75, heights[1], "PSF chi2:", color=text_color) - fig.text(0.85, heights[1], f"{source['psfChi2']/source['psfNdata']:6.2f}") - - fig.text(0.0, heights[2], "PSF (nJy):", color=flag_color if any(source[flags_psf]) else text_color) - fig.text(0.25, heights[2], f"{source['psfFlux']:8.1f}", horizontalalignment='right') - fig.text(0.252, heights[2], "+/-", color=text_color) - fig.text(0.29, heights[2], f"{source['psfFluxErr']:8.1f}") - fig.text(0.40, heights[2], "S/N:", color=text_color) - fig.text(0.45, heights[2], f"{abs(source['psfFlux']/source['psfFluxErr']):6.2f}") - - # NOTE: yellow is hard to read on white; use goldenrod instead. - if any(source[flags_edge]): - fig.text(0.55, heights[2], "EDGE", color="goldenrod", fontweight="bold") - if any(source[flags_interp]): - fig.text(0.62, heights[2], "INTERP", color="green", fontweight="bold") - if any(source[flags_saturated]): - fig.text(0.72, heights[2], "SAT", color="green", fontweight="bold") - if any(source[flags_cr]): - fig.text(0.77, heights[2], "CR", color="magenta", fontweight="bold") - if any(source[flags_bad]): - fig.text(0.81, heights[2], "BAD", color="red", fontweight="bold") - if source['isDipole']: - fig.text(0.87, heights[2], "DIPOLE", color="indigo", fontweight="bold") - - fig.text(0.0, heights[3], "ap (nJy):", color=flag_color if any(source[flags_aperture]) else text_color) - fig.text(0.25, heights[3], f"{source['apFlux']:8.1f}", horizontalalignment='right') - fig.text(0.252, heights[3], "+/-", color=text_color) - fig.text(0.29, heights[3], f"{source['apFluxErr']:8.1f}") - fig.text(0.40, heights[3], "S/N:", color=text_color) - fig.text(0.45, heights[3], f"{abs(source['apFlux']/source['apFluxErr']):#6.2f}") - - if any(source[flags_suspect]): - fig.text(0.55, heights[3], "SUS", color="goldenrod", fontweight="bold") - if any(source[flags_centroid]): - fig.text(0.60, heights[3], "CENTROID", color="red", fontweight="bold") - if any(source[flags_shape]): - fig.text(0.73, heights[3], "SHAPE", color="red", fontweight="bold") - # Future option: to add two more flag flavors to the legend, - # use locations 0.80 and 0.87 - - # rb score + if heights is None: + if len_sizes == 1: + heights = [0.95, 0.91, 0.87, 0.83, 0.79] + else: + heights = [1.2, 1.15, 1.1, 1.05, 1.0] + + def label_color(group_key): + """Red label if any flag in the group is set, otherwise grey.""" + return flag_color if any(source[_FLAG_GROUPS[group_key]]) else text_color + + # Each row is a list of (x, text, kwargs) atoms drawn at heights[row_idx]. + # fig.text coordinates are in fractions of the figure. + rows = [ + # Row 0: identity (diaSourceId, instrument, detector, visit, band). + [ + (0.000, "diaSourceId:", {"color": text_color}), + (0.145, f"{source['diaSourceId']}", {}), + (0.430, f"{source['instrument']}", {"fontweight": "bold"}), + (0.640, "detector:", {"color": text_color}), + (0.740, f"{source['detector']}", {}), + (0.795, "visit:", {"color": text_color}), + (0.850, f"{source['visit']}", {}), + (0.950, f"{source['band']}", {}), + ], + # Row 1: coordinates and detection-quality numbers. + [ + (0.000, "ra:", {"color": text_color}), + (0.037, f"{source['ra']:.8f}", {}), + (0.210, "dec:", {"color": text_color}), + (0.265, f"{source['dec']:+.8f}", {}), + (0.500, "detection S/N:", {"color": text_color}), + (0.660, f"{source['snr']:6.1f}", {}), + (0.750, "PSF chi2:", {"color": text_color}), + (0.850, f"{source['psfChi2']/source['psfNdata']:6.2f}", {}), + ], + # Row 2: PSF flux. + [ + (0.000, "PSF (nJy):", {"color": label_color("psf")}), + (0.250, f"{source['psfFlux']:8.1f}", {"horizontalalignment": "right"}), + (0.252, "+/-", {"color": text_color}), + (0.290, f"{source['psfFluxErr']:8.1f}", {}), + (0.400, "S/N:", {"color": text_color}), + (0.450, f"{abs(source['psfFlux']/source['psfFluxErr']):6.2f}", {}), + ], + # Row 3: aperture flux. + [ + (0.000, "ap (nJy):", {"color": label_color("aperture")}), + (0.250, f"{source['apFlux']:8.1f}", {"horizontalalignment": "right"}), + (0.252, "+/-", {"color": text_color}), + (0.290, f"{source['apFluxErr']:8.1f}", {}), + (0.400, "S/N:", {"color": text_color}), + (0.450, f"{abs(source['apFlux']/source['apFluxErr']):#6.2f}", {}), + ], + # Row 4: forced-photometry flux + ABmag. + [ + (0.000, "sci (nJy):", {"color": label_color("forced")}), + (0.250, f"{source['scienceFlux']:8.1f}", {"horizontalalignment": "right"}), + (0.252, "+/-", {"color": text_color}), + (0.290, f"{source['scienceFluxErr']:8.1f}", {}), + (0.400, "S/N:", {"color": text_color}), + (0.450, f"{abs(source['scienceFlux']/source['scienceFluxErr']):6.2f}", {}), + (0.550, "ABmag:", {"color": text_color}), + (0.635, f"{(source['scienceFlux']*u.nanojansky).to_value(u.ABmag):.3f}", {}), + ], + ] + for row, y in zip(rows, heights): + for x, text, kwargs in row: + fig.text(x, y, text, **kwargs) + + # Draw flag-tag overlays after the rows so they sit on top. + for predicate, x, text, color, row_idx in _FLAG_TAGS: + if callable(predicate): + triggered = predicate(source) + else: + triggered = any(source[_FLAG_GROUPS[predicate]]) + if triggered: + fig.text(x, heights[row_idx], text, color=color, fontweight="bold") + + # Reliability score: color depends on the value, not on flags. if source['reliability'] is not None and np.isfinite(source['reliability']): - fig.text(0.73, heights[4], f"RB:{source['reliability']:.03f}", - color='#e41a1c' if source['reliability'] < 0.5 else '#4daf4a', + rb = source['reliability'] + fig.text(0.73, heights[4], f"RB:{rb:.03f}", + color='#e41a1c' if rb < 0.5 else '#4daf4a', fontweight="bold") - fig.text(0.0, heights[4], "sci (nJy):", color=flag_color if any(source[flags_forced]) else text_color) - fig.text(0.25, heights[4], f"{source['scienceFlux']:8.1f}", horizontalalignment='right') - fig.text(0.252, heights[4], "+/-", color=text_color) - fig.text(0.29, heights[4], f"{source['scienceFluxErr']:8.1f}") - fig.text(0.40, heights[4], "S/N:", color=text_color) - fig.text(0.45, heights[4], f"{abs(source['scienceFlux']/source['scienceFluxErr']):6.2f}") - fig.text(0.55, heights[4], "ABmag:", color=text_color) - fig.text(0.635, heights[4], f"{(source['scienceFlux']*u.nanojansky).to_value(u.ABmag):.3f}") - class CutoutPath: """Manage paths to image cutouts with filenames based on diaSourceId. @@ -865,72 +901,6 @@ def _make_apdbQuery(sqlitefile=None, postgres_url=None, namespace=None): return apdb_query -def select_sources(apdb_query, limit, reliabilityMin=None, reliabilityMax=None): - """Load an APDB and return n sources from it. - - Parameters - ---------- - apdb_query : `lsst.analysis.ap.ApdbQuery` - APDB query interface to load from. - limit : `int` - Number of sources to select from the APDB. - reliabilityMin : `float` - Minimum reliability value on which to filter the DiaSources. - reliabilityMax : `float` - Maximum reliability value on which to filter the DiaSources. - - Returns - ------- - sources : `pandas.DataFrame` - The loaded DiaSource data. - """ - offset = 0 - try: - while True: - with apdb_query.connection as connection: - table = apdb_query._tables["DiaSource"] - query = table.select() - if reliabilityMin is not None: - query = query.where(table.columns['reliability'] >= reliabilityMin) - if reliabilityMax is not None: - query = query.where(table.columns['reliability'] <= reliabilityMax) - query = query.order_by(table.columns["visit"], - table.columns["detector"], - table.columns["diaSourceId"]) - query = query.limit(limit).offset(offset) - sources = pd.read_sql_query(query, connection) - if len(sources) == 0: - break - apdb_query._fill_from_instrument(sources) - - yield sources - offset += limit - finally: - connection.close() - - -def len_sources(apdb_query, namespace=None): - """Return the number of DiaSources in the supplied APDB. - - Parameters - ---------- - apdb_query : `lsst.analysis.ap.ApdbQuery` - APDB query interface to load from. - namespace : `str`, optional - Postgres schema to load data from. - - Returns - ------- - count : `int` - Number of diaSources in this APDB. - """ - with apdb_query.connection as connection: - if namespace: - connection.execute(sqlalchemy.text(f"SET search_path TO {namespace}")) - count = connection.execute(sqlalchemy.text('select count(*) FROM "DiaSource";')).scalar() - return count - - def run_cutouts(args): """Run PlotImageSubtractionCutoutsTask on the parsed commandline arguments. @@ -957,7 +927,7 @@ def run_cutouts(args): if config.save_as_numpy: # save the RB output up front so we can use partial runs - data = select_sources(apdb_query, args.limit, args.reliabilityMin, args.reliabilityMax) + data = apdb_query.iter_sources(args.limit, args.reliabilityMin, args.reliabilityMax) cols_to_export = ["diaSourceId", "visit", "detector", "diaObjectId", "ssObjectId", "midpointMjdTai", "ra", "dec", "x", "y", "apFlux", "apFluxErr", "snr", "psfFlux", "psfFluxErr", @@ -972,14 +942,14 @@ def run_cutouts(args): all_data = pd.concat([d[cols_to_export] for d in data]) all_data.to_csv(os.path.join(args.outputPath, "all_diasources.csv.gz"), index=False) - getter = select_sources(apdb_query, args.limit, args.reliabilityMin, args.reliabilityMax) + getter = apdb_query.iter_sources(args.limit, args.reliabilityMin, args.reliabilityMax) # Process just one block of length "limit", or all sources in the database? if not args.all: data = next(getter) sources = cutouts.run(data, butler, njobs=args.jobs) else: sources = [] - count = len_sources(apdb_query, args.namespace) + count = apdb_query.count_sources() for i, data in enumerate(getter): sources.extend(cutouts.write_images(data, butler, njobs=args.jobs)) print(f"Completed {i+1} batches of {args.limit} size, out of {count} diaSources.") diff --git a/python/lsst/analysis/ap/plotUtils.py b/python/lsst/analysis/ap/plotUtils.py new file mode 100644 index 0000000..459876c --- /dev/null +++ b/python/lsst/analysis/ap/plotUtils.py @@ -0,0 +1,343 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Visualization helpers for AP analysis. + +Three tools live here: + +- `lightcurve` plots a per-band psfFlux light curve for a single diaObject, + optionally overlaying forced photometry. It works with either the APDB + `DbQuery` interface (pandas) or the PPDB `PpdbTap` interface (astropy + Tables). +- `cutout_grid` lays out science/template/difference cutouts for many + DiaSources in a single mosaic figure. +- `summarize_run` returns a per-visit summary DataFrame of an APDB run + (counts, dipole rate, reliability statistics, etc.) suitable for a quick + health check of a processing run. +""" + +from __future__ import annotations + +__all__ = ["lightcurve", "cutout_grid", "summarize_run", "BAND_COLORS"] + +import inspect +import io + +import numpy as np +import pandas as pd + +# Colors keyed by lower-case band character. Mirrors the choices in +# legacyPlotUtils.source_magnitude_histogram (g=C2, r=C1, i=C3, z=C5, y=k). +BAND_COLORS = { + "u": "C4", + "g": "C2", + "r": "C1", + "i": "C3", + "z": "C5", + "y": "k", +} + + +def _time_column(frame): + """Return the name of the MJD-like time column on a DiaSource frame. + + Older APDB schemas used ``midPointTai``; current ones use + ``midpointMjdTai``. This helper accepts either. + """ + for candidate in ("midpointMjdTai", "midPointTai"): + if candidate in frame.columns: + return candidate + raise KeyError("Expected one of 'midpointMjdTai' or 'midPointTai' " + f"in DataFrame; got columns: {list(frame.columns)}") + + +def _to_dataframe(table): + """Return ``table`` as a `pandas.DataFrame`. + + The APDB `DbQuery` loaders already return DataFrames; the PPDB `PpdbTap` + loaders return `astropy.table.Table`. Normalizing here lets the plotting + code use a single pandas (groupby-based) path regardless of which + interface produced the data. Masked astropy values become NaN. + """ + if isinstance(table, pd.DataFrame): + return table + return table.to_pandas() + + +def _load_object_sources(query, dia_object_id, exclude_flagged): + """Load one diaObject's DiaSources as a DataFrame across query interfaces. + """ + method = query.load_sources_for_object + if "exclude_flagged" in inspect.signature(method).parameters: + sources = method(dia_object_id, exclude_flagged=exclude_flagged) + elif exclude_flagged: + raise TypeError( + f"{type(query).__name__}.load_sources_for_object does not support " + "exclude_flagged; the PPDB public interface does not expose " + "diaSource flag filtering. Pass exclude_flagged=False.") + else: + sources = method(dia_object_id) + return _to_dataframe(sources) + + +def lightcurve(query, dia_object_id, ax=None, exclude_flagged=False, + include_forced=True): + """Plot a per-band psfFlux light curve for one diaObject. + + Parameters + ---------- + query : `lsst.analysis.ap.apdb.DbQuery` or \ + `lsst.analysis.ap.ppdb.PpdbTap` + dia_object_id : `int` + Object id to load. + ax : `matplotlib.axes.Axes`, optional + Axes to draw into; if None, a new figure is created. + exclude_flagged : `bool`, optional + Forwarded to `load_sources_for_object` when the query supports it. + Defaults to False so the lightcurve matches the row count of a direct + APDB query; pass True to drop diaSources matching the configured + bad-flag list. The PPDB `PpdbTap` interface does not expose flag + filtering, so True is rejected there. DiaForcedSources are always + loaded unfiltered. + include_forced : `bool`, optional + If True, also overlay diaForcedSources as small markers. + + Returns + ------- + fig : `matplotlib.figure.Figure` + ax : `matplotlib.axes.Axes` + sources : `pandas.DataFrame` + DiaSources used for the plot. + forced : `pandas.DataFrame` or None + DiaForcedSources used for the plot (None if ``include_forced`` is + False). + """ + import matplotlib.pyplot as plt + + if ax is None: + fig, ax = plt.subplots(figsize=(8, 5)) + else: + fig = ax.figure + + # diaObjectId is NaN for diaSources not associated with any diaObject; + # short-circuit before hitting the database. + if pd.isna(dia_object_id): + ax.text(0.5, 0.5, "no diaObjectId (NaN)", + ha="center", va="center", transform=ax.transAxes) + return fig, ax, pd.DataFrame(), None + + sources = _load_object_sources(query, dia_object_id, exclude_flagged) + # DiaForcedSource has a different (and smaller) flag schema than + # DiaSource: applying the diaSource exclusion list would key into + # columns that don't exist on the forced table. Forced photometry is + # also a measurement at a known location rather than a fresh detection, + # so showing it unfiltered is the right behavior. + forced = (_to_dataframe(query.load_forced_sources_for_object(dia_object_id)) + if include_forced else None) + + if len(sources) == 0: + ax.text(0.5, 0.5, f"no sources for diaObjectId={dia_object_id}", + ha="center", va="center", transform=ax.transAxes) + return fig, ax, sources, forced + + time_col = _time_column(sources) + for band, group in sources.groupby("band"): + color = BAND_COLORS.get(band, "k") + ax.errorbar(group[time_col], group["psfFlux"], yerr=group["psfFluxErr"], + fmt="o", color=color, label=f"{band} (n={len(group)})") + + if forced is not None and len(forced): + forced_time_col = _time_column(forced) + # Plot per-band forced points in their band colors, but suppress + # individual legend entries so the forced marker is represented + # once (in black) regardless of how many bands are present. + for band, group in forced.groupby("band"): + color = BAND_COLORS.get(band, "k") + ax.errorbar(group[forced_time_col], group["psfFlux"], + yerr=group["psfFluxErr"], fmt=".", ms=4, color=color, + alpha=0.4, label="_nolegend_") + ax.plot([], [], ".", color="black", ms=4, alpha=0.4, + label=f"forced (n={len(forced)})") + + ax.axhline(0, color="grey", lw=0.5) + ax.set_xlabel(time_col) + ax.set_ylabel("psfFlux (nJy)") + ax.set_title(f"diaObjectId = {dia_object_id}") + ax.legend(frameon=True) + return fig, ax, sources, forced + + +def cutout_grid(sources, butler, instrument, n_per_row=4, config=None, output=None, + figsize=None, ra_column='ra', dec_column='dec', detector_column='detector', + visit_column='visit', id_column='diaSourceId'): + """Render science/template/difference cutouts for many sources in a grid. + + This is a thin wrapper around `PlotImageSubtractionCutoutsTask`: it calls + `generate_image` for each source (which returns a PNG in memory) and + arranges the resulting rasters in a single matplotlib figure. + + Parameters + ---------- + sources : `pandas.DataFrame` + DiaSources to cut out. Must contain at least + ``ra, dec, diaSourceId, detector, visit, instrument`` plus whatever + annotation fields the task config requires (see + ``PlotImageSubtractionCutoutsConfig.add_metadata``). + butler : `lsst.daf.butler.Butler` + Butler initialized with the relevant collections. + instrument : `str` + Name of the instrument for the data being plotted. + n_per_row : `int` + Number of cutouts per row in the resulting figure. + config : `PlotImageSubtractionCutoutsConfig`, optional + Cutout config to use (see + ``plotImageSubtractionCutouts.PlotImageSubtractionCutoutsConfig``). + Defaults to a fresh instance with ``add_metadata=False`` + (annotations get cluttered in a grid). + output : `str`, optional + If given, save the figure to this path with ``bbox_inches="tight"``. + figsize : `tuple` [`float`, `float`], optional + Figure size in inches. Defaults to ``(n_per_row*3.5, n_rows*1.7)``. + + Returns + ------- + fig : `matplotlib.figure.Figure` + """ + import matplotlib.pyplot as plt + import PIL.Image + + import lsst.geom + + # Local import to avoid a circular dependency at module load time. + from . import plotImageSubtractionCutouts as cutouts_mod + + if config is None: + config = cutouts_mod.PlotImageSubtractionCutoutsConfig() + # Annotations get cramped in a grid; default them off here. + config.add_metadata = False + + task = cutouts_mod.PlotImageSubtractionCutoutsTask(config=config, output_path="") + cutouts_mod.butler_cache.set(butler, config) + + n_sources = len(sources) + n_rows = max(1, (n_sources + n_per_row - 1) // n_per_row) + if figsize is None: + figsize = (n_per_row * 3.5, n_rows * 1.7) + fig, axes = plt.subplots(n_rows, n_per_row, figsize=figsize, squeeze=False) + + records = sources.to_records(index=False) + for i, source in enumerate(records): + row, col = divmod(i, n_per_row) + ax = axes[row][col] + ax.set_axis_off() + try: + sci, tmpl, diff = cutouts_mod.butler_cache.get_exposures( + instrument, source[detector_column], source[visit_column]) + center = lsst.geom.SpherePoint(source[ra_column], source[dec_column], + lsst.geom.degrees) + scale = sci.wcs.getPixelScale(sci.getBBox().getCenter()).asArcseconds() + png = task.generate_image( + sci, tmpl, diff, center, scale, + source=source if config.add_metadata else None, + ) + with PIL.Image.open(io.BytesIO(png.getvalue())) as img: + ax.imshow(np.asarray(img)) + ax.set_title(f"{source[id_column]}", fontsize=7) + except Exception as exc: + ax.text(0.5, 0.5, f"{type(exc).__name__}", ha="center", va="center", + fontsize=8, transform=ax.transAxes) + + # Blank out any trailing axes in the last row. + for i in range(n_sources, n_rows * n_per_row): + row, col = divmod(i, n_per_row) + axes[row][col].set_axis_off() + + fig.tight_layout() + if output is not None: + fig.savefig(output, bbox_inches="tight") + # Remove the figure from pyplot's figure manager so the Jupyter inline + # backend does not auto-display it at end-of-cell in addition to the + # caller's own rendering of the returned Figure (which would duplicate + # the output the first time the function is run in a notebook). The + # returned Figure object remains valid and renders via its repr hooks. + plt.close(fig) + return fig + + +def summarize_run(query, bad_flag_list=None): + """Return a per-visit summary DataFrame for an APDB run. + + Useful as a quick health check after a processing run: one row per visit + with source counts, dipole rate, reliability statistics, and the fraction + of sources that fall in the bad-flag list. + + Parameters + ---------- + query : `lsst.analysis.ap.apdb.DbQuery` + APDB query interface (sqlite, postgres, or cassandra). + bad_flag_list : `list` [`str`], optional + Flag column names to count as "bad". If omitted, the query's + currently-configured exclusion list is used. The caller's exclusion + list is restored before returning. + + Returns + ------- + summary : `pandas.DataFrame` + Indexed by ``visit``. Columns: + ``n_sources``, ``n_unflagged``, ``bad_flag_fraction``, + ``median_reliability`` (if column present), + ``dipole_fraction`` (if column present), + ``median_psf_chi2_per_dof`` (if columns present). + """ + saved_flags = list(query.diaSource_flags_exclude) + try: + if bad_flag_list is not None: + query.set_excluded_diaSource_flags(bad_flag_list) + sources_all = query.load_sources(limit=None) + sources_clean = query.load_sources(exclude_flagged=True, limit=None) + finally: + query.set_excluded_diaSource_flags(saved_flags) + + if len(sources_all) == 0: + return pd.DataFrame() + + clean_per_visit = sources_clean.groupby("visit").size() if len(sources_clean) else pd.Series(dtype=int) + + rows = [] + for visit, group in sources_all.groupby("visit"): + n_all = len(group) + n_clean = int(clean_per_visit.get(visit, 0)) + row = { + "visit": visit, + "n_sources": n_all, + "n_unflagged": n_clean, + "bad_flag_fraction": 1.0 - n_clean / n_all if n_all else 0.0, + } + if "reliability" in group.columns: + row["median_reliability"] = group["reliability"].median() + if "isDipole" in group.columns: + row["dipole_fraction"] = float(group["isDipole"].mean()) + if "psfChi2" in group.columns and "psfNdata" in group.columns: + with np.errstate(divide="ignore", invalid="ignore"): + ratio = group["psfChi2"] / group["psfNdata"] + row["median_psf_chi2_per_dof"] = float(np.nanmedian(ratio)) + rows.append(row) + return pd.DataFrame(rows).set_index("visit").sort_index() diff --git a/python/lsst/analysis/ap/skymapOverlay.py b/python/lsst/analysis/ap/skymapOverlay.py new file mode 100644 index 0000000..69a718f --- /dev/null +++ b/python/lsst/analysis/ap/skymapOverlay.py @@ -0,0 +1,427 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Backend-agnostic skymap tract/patch overlays. + +The geometry of projecting overlapping tract/patch boundaries into a +display's pixel coordinate system is shared by two very different +renderers: + +- `draw_skymap_outlines_mpl`, which draws onto a Matplotlib axis whose + pixel<->sky relationship is only known approximately (e.g. the + spatially-sampled-metrics panels, which carry sampled ``(x, y)`` and + ``(coord_ra, coord_dec)`` but no WCS). Use `make_affine_sky_to_xy` to + build the required ``sky_to_xy`` callable from those samples. +- `draw_skymap_outlines_afw`, which draws onto an `lsst.afw.display` + frame showing an exposure with a real WCS, so the sky<->pixel mapping + is exact. + +Both renderers share `compute_tract_patch_outlines`, which does the +findTractPatchList lookup, projects every patch corner through a caller +supplied ``sky_to_xy`` map, and ranks the overlapping tracts by how much +of the display they cover. +""" + +from __future__ import annotations + +__all__ = ["make_affine_sky_to_xy", "compute_tract_patch_outlines", + "draw_skymap_outlines_mpl", "draw_skymap_outlines_afw"] + +import numpy as np + + +def make_affine_sky_to_xy(ra, dec, x, y): + """Build a least-squares affine map from sky to detector pixels. + + Useful when no WCS is available but matched ``(ra, dec)`` and + ``(x, y)`` samples are. For a single detector the affine fit is + typically accurate to well under a pixel -- enough for visualization + but not for science. + + Parameters + ---------- + ra, dec : array-like + Sky coordinates of the samples, in **radians**. + x, y : array-like + Detector pixel coordinates of the same samples. + + Returns + ------- + sky_to_xy : callable + Maps an `lsst.geom.SpherePoint` to an ``(x, y)`` tuple of + `float` in the same pixel system as the input ``x, y``. + """ + ra = np.asarray(ra) + dec = np.asarray(dec) + A = np.column_stack([np.ones(ra.size), ra, dec]) + coef_x, *_ = np.linalg.lstsq(A, np.asarray(x), rcond=None) + coef_y, *_ = np.linalg.lstsq(A, np.asarray(y), rcond=None) + + def sky_to_xy(sphere_point): + r = sphere_point.getRa().asRadians() + d = sphere_point.getDec().asRadians() + return (float(coef_x[0] + coef_x[1]*r + coef_x[2]*d), + float(coef_y[0] + coef_y[1]*r + coef_y[2]*d)) + + return sky_to_xy + + +def compute_tract_patch_outlines(skymap, sky_to_xy, sky_corners, clip_rect): + """Project overlapping tract/patch boundaries into display coordinates. + + Parameters + ---------- + skymap : `lsst.skymap.BaseSkyMap` + Skymap to query for overlapping tracts and patches. + sky_to_xy : callable + Maps an `lsst.geom.SpherePoint` to an ``(x, y)`` tuple in the + display's pixel coordinate system. + sky_corners : `list` [`lsst.geom.SpherePoint`] + Sky positions spanning the region of interest; passed to + ``skymap.findTractPatchList`` to enumerate the tracts/patches that + touch it. + clip_rect : `tuple` [`float`] + ``(xmin, xmax, ymin, ymax)`` display-pixel rectangle. Used only to + rank tracts by their visible (clipped) patch area, so that the + most-covering tract sorts first; it does not clip the returned + geometry. + + Returns + ------- + outlines : `list` [`dict`] + One entry per overlapping tract, sorted by descending visible + area (ties broken by tract id for determinism). Each is:: + + {"tract_id": int, + "rank": int, # 0-based position in the sorted list + "patches": [{"patch_index": int, # sequential index + "corners_xy": [(x, y), ...], # 4 corners + "center_xy": (x, y)}, # inner-bbox center + ...]} + """ + import lsst.geom as geom + + xmin, xmax, ymin, ymax = clip_rect + tract_patch_list = skymap.findTractPatchList(sky_corners) + + tract_entries = [] # (visible_area, tract_id, patches) + for tract_info, patches in tract_patch_list: + tract_wcs = tract_info.wcs + per_tract_patches = [] + per_tract_area = 0.0 + for patch in patches: + bbox = patch.getInnerBBox() + corners_px = [geom.Point2D(bbox.minX, bbox.minY), + geom.Point2D(bbox.maxX, bbox.minY), + geom.Point2D(bbox.maxX, bbox.maxY), + geom.Point2D(bbox.minX, bbox.maxY)] + xy_corners = [sky_to_xy(tract_wcs.pixelToSky(p)) for p in corners_px] + center_px = geom.Point2D(0.5*(bbox.minX + bbox.maxX), + 0.5*(bbox.minY + bbox.maxY)) + center_xy = sky_to_xy(tract_wcs.pixelToSky(center_px)) + + clipped = _clip_polygon_to_rect(xy_corners, xmin, xmax, ymin, ymax) + per_tract_area += _polygon_area(clipped) + per_tract_patches.append({"patch_index": patch.getSequentialIndex(), + "corners_xy": xy_corners, + "center_xy": center_xy}) + tract_entries.append((per_tract_area, tract_info.getId(), per_tract_patches)) + + # Sort by descending visible area; break ties by tract id so the order + # (and hence the mpl linestyle assignment) is stable across calls. + tract_entries.sort(key=lambda e: (-e[0], e[1])) + return [{"tract_id": tract_id, "rank": rank, "patches": patches} + for rank, (_area, tract_id, patches) in enumerate(tract_entries)] + + +def draw_skymap_outlines_mpl(ax, skymap, sky_to_xy, sky_corners, *, + label_fontsize=7): + """Overlay patch boundaries (with ``tract,patch`` labels) on a panel. + + Each overlapping patch gets a thin white outline (black-stroked so it + reads on any colormap) plus a ``tract,patch`` label placed along the + midpoint of its longest visible edge inside the current view. + + When more than one tract overlaps the panel (e.g. detectors near a + tract boundary), patches are distinguished by linestyle: tracts are + ranked by their total visible patch area and assigned solid, dashed, + dotted, dash-dotted in turn. + + Parameters + ---------- + ax : `matplotlib.axes.Axes` + Axis to draw on. Its current ``xlim``/``ylim`` define the visible + region and are restored on return. + skymap : `lsst.skymap.BaseSkyMap` + Skymap to query for overlapping tracts and patches. + sky_to_xy : callable + Maps an `lsst.geom.SpherePoint` to an ``(x, y)`` tuple in the + axis' data coordinate system (see `make_affine_sky_to_xy`). + sky_corners : `list` [`lsst.geom.SpherePoint`] + Sky positions spanning the region of interest, used to enumerate + the overlapping tracts/patches. + label_fontsize : `int` or `float`, optional + Font size for the per-patch ``tract,patch`` labels. + """ + import matplotlib.patheffects as pe + + xlim = ax.get_xlim() + ylim = ax.get_ylim() + xmin, xmax = xlim + ymin, ymax = ylim + clip_rect = (xmin, xmax, ymin, ymax) + + outlines = compute_tract_patch_outlines(skymap, sky_to_xy, sky_corners, clip_rect) + + # White lines with a thin black stroke read on any colormap background. + line_outline = [pe.withStroke(linewidth=2.0, foreground="black")] + text_outline = [pe.withStroke(linewidth=1.4, foreground="black")] + linestyles = ("-", "--", ":", "-.") + + for tract in outlines: + rank = tract["rank"] + tract_id = tract["tract_id"] + linestyle = linestyles[rank] if rank < len(linestyles) else "-" + patch_kwargs = dict(color="white", linewidth=0.5, alpha=0.85, + linestyle=linestyle, + path_effects=line_outline, zorder=6) + for patch in tract["patches"]: + xy_corners = patch["corners_xy"] + xs_c = [p[0] for p in xy_corners] + [xy_corners[0][0]] + ys_c = [p[1] for p in xy_corners] + [xy_corners[0][1]] + ax.plot(xs_c, ys_c, **patch_kwargs) + + # Place the label at the midpoint of the patch edge with the + # longest visible portion within the panel, offset slightly + # toward the patch center so the text sits inside. + cx, cy = patch["center_xy"] + best = None + for i in range(4): + (x0, y0), (x1, y1) = xy_corners[i], xy_corners[(i + 1) % 4] + clipped = _clip_segment_to_rect(x0, y0, x1, y1, + xmin, xmax, ymin, ymax) + if clipped is None: + continue + cx0, cy0, cx1, cy1 = clipped + length = float(np.hypot(cx1 - cx0, cy1 - cy0)) + if best is None or length > best[0]: + best = (length, cx0, cy0, cx1, cy1) + if best is None: + continue # entire patch is outside the panel + length, cx0, cy0, cx1, cy1 = best + mx = 0.5*(cx0 + cx1) + my = 0.5*(cy0 + cy1) + # Inward offset toward the projected patch center. Use the + # smaller of "fixed fraction of edge length" and "fraction of + # the midpoint-to-center distance" so the offset never lands + # outside the patch on slivers. + dx, dy = cx - mx, cy - my + d_center = float(np.hypot(dx, dy)) + if d_center > 0: + step = min(0.06*length, 0.4*d_center) + mx += dx/d_center * step + my += dy/d_center * step + + # Rotate the text to lie parallel to the visible edge, flipping + # to keep it reading right-side-up (angle clamped to [-90, 90]). + angle_deg = float(np.degrees(np.arctan2(cy1 - cy0, cx1 - cx0))) + if angle_deg > 90.0: + angle_deg -= 180.0 + elif angle_deg < -90.0: + angle_deg += 180.0 + + ax.text(mx, my, f"{tract_id},{patch['patch_index']}", + ha="center", va="center", + rotation=angle_deg, rotation_mode="anchor", + color="white", fontsize=label_fontsize, + path_effects=text_outline, zorder=7, + clip_on=True) + + ax.set_xlim(xlim) + ax.set_ylim(ylim) + + +def draw_skymap_outlines_afw(afw_display, skymap, wcs, bbox, *, + ctype="green", label_size=1.5, draw_labels=True): + """Overlay tract/patch boundaries on the current `afw.display` frame. + + Each overlapping patch is drawn as a closed polyline in the image's + parent pixel coordinates and, optionally, labeled ``tract,patch`` + near the center of its visible portion. + + Parameters + ---------- + afw_display : `lsst.afw.display.Display` + Display whose current frame already shows the exposure. The caller + is responsible for selecting the frame. + skymap : `lsst.skymap.BaseSkyMap` + Skymap to query for overlapping tracts and patches. + wcs : `lsst.afw.geom.SkyWcs` + WCS of the displayed exposure (``exposure.wcs``). + bbox : `lsst.geom.Box2I` or `lsst.geom.Box2D` + Parent bounding box of the displayed exposure. + Defines the footprint searched for overlapping patches + and the region used to place labels. + ctype : `str`, optional + Display color for both the outlines and labels. + label_size : `float`, optional + Text size for the ``tract,patch`` labels. + draw_labels : `bool`, optional + If False, draw only the outlines. + """ + import lsst.geom as geom + + def sky_to_xy(sphere_point): + p = wcs.skyToPixel(sphere_point) + return (p.getX(), p.getY()) + + xmin = bbox.getMinX() + xmax = bbox.getMaxX() + ymin = bbox.getMinY() + ymax = bbox.getMaxY() + clip_rect = (xmin, xmax, ymin, ymax) + + # Sky positions at the four image corners span the footprint for the + # tract/patch lookup. + corner_px = [geom.Point2D(xmin, ymin), geom.Point2D(xmax, ymin), + geom.Point2D(xmax, ymax), geom.Point2D(xmin, ymax)] + sky_corners = [wcs.pixelToSky(p) for p in corner_px] + + outlines = compute_tract_patch_outlines(skymap, sky_to_xy, sky_corners, clip_rect) + + with afw_display.Buffering(): + for tract in outlines: + tract_id = tract["tract_id"] + for patch in tract["patches"]: + xy_corners = patch["corners_xy"] + # Closed polyline: repeat the first corner. + afw_display.line(list(xy_corners) + [xy_corners[0]], ctype=ctype) + if not draw_labels: + continue + # Anchor the label at the centroid of the patch's visible + # portion so it doesn't land off-image for patches that + # mostly fall outside the detector. + clipped = _clip_polygon_to_rect(xy_corners, xmin, xmax, ymin, ymax) + if clipped: + lx = sum(p[0] for p in clipped)/len(clipped) + ly = sum(p[1] for p in clipped)/len(clipped) + else: + lx, ly = patch["center_xy"] + afw_display.dot(f"{tract_id},{patch['patch_index']}", + lx, ly, size=label_size, ctype=ctype) + + +def _clip_polygon_to_rect(polygon, xmin, xmax, ymin, ymax): + """Clip a convex polygon against an axis-aligned rectangle. + + Parameters + ---------- + polygon : sequence of ``(x, y)`` tuples + Vertices of the (convex) input polygon, in order. + xmin, xmax, ymin, ymax : `float` + The clipping rectangle. + + Returns + ------- + clipped : `list` of ``(x, y)`` tuples + The clipped polygon, or an empty list if the polygon lies + entirely outside the rectangle. + """ + # Each clip edge is parameterized by ("axis", value, keep_side) + # where keep_side is +1 if "inside" means coordinate >= value, + # -1 if "inside" means coordinate <= value. + edges = (("x", xmin, +1), ("x", xmax, -1), + ("y", ymin, +1), ("y", ymax, -1)) + + def _inside(point, axis, val, sign): + coord = point[0] if axis == "x" else point[1] + return (coord - val)*sign >= 0.0 + + def _intersect(p1, p2, axis, val): + x1, y1 = p1 + x2, y2 = p2 + if axis == "x": + t = (val - x1)/(x2 - x1) + return (val, y1 + t*(y2 - y1)) + t = (val - y1)/(y2 - y1) + return (x1 + t*(x2 - x1), val) + + output = list(polygon) + for axis, val, sign in edges: + if not output: + return [] + input_list = output + output = [] + for i in range(len(input_list)): + curr = input_list[i] + prev = input_list[i - 1] + curr_in = _inside(curr, axis, val, sign) + prev_in = _inside(prev, axis, val, sign) + if curr_in: + if not prev_in: + output.append(_intersect(prev, curr, axis, val)) + output.append(curr) + elif prev_in: + output.append(_intersect(prev, curr, axis, val)) + return output + + +def _polygon_area(polygon): + """Area of a polygon via the shoelace formula.""" + n = len(polygon) + if n < 3: + return 0.0 + s = 0.0 + for i in range(n): + x1, y1 = polygon[i] + x2, y2 = polygon[(i + 1) % n] + s += x1*y2 - x2*y1 + return abs(s)*0.5 + + +def _clip_segment_to_rect(x0, y0, x1, y1, xmin, xmax, ymin, ymax): + """Liang-Barsky line-segment clipping against an axis-aligned rect. + + Returns the clipped endpoints ``(x0', y0', x1', y1')`` or ``None`` if + the segment lies entirely outside the rectangle. + """ + dx = x1 - x0 + dy = y1 - y0 + p = (-dx, dx, -dy, dy) + q = (x0 - xmin, xmax - x0, y0 - ymin, ymax - y0) + u1, u2 = 0.0, 1.0 + for pi, qi in zip(p, q): + if pi == 0.0: + if qi < 0.0: + return None + else: + t = qi/pi + if pi < 0.0: + if t > u2: + return None + if t > u1: + u1 = t + else: + if t < u1: + return None + if t < u2: + u2 = t + return (x0 + u1*dx, y0 + u1*dy, x0 + u2*dx, y0 + u2*dy) diff --git a/python/lsst/analysis/ap/spatiallySampledMetricsQA.py b/python/lsst/analysis/ap/spatiallySampledMetricsQA.py new file mode 100644 index 0000000..11f4bba --- /dev/null +++ b/python/lsst/analysis/ap/spatiallySampledMetricsQA.py @@ -0,0 +1,396 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Notebook-friendly QA for ``SpatiallySampledMetricsTask`` output. + +The single public entry point `subtraction_quality_report` consumes a +``*_spatiallySampledMetrics`` table for one detector and produces: + +- A printed top-line summary of the three diagnostic scalars + (`diffim_chi2PerPix`, `psfMatchingKernel_residualNorm`, + `dipole_density`) at percentiles that make localized failures visible. +- A three-panel diagnostic figure with the three scalars rendered as + linear-interpolated heatmaps, with the kernel centroid offset quiver + (colored cyclically by angle, magnitude shown by a per-panel + quiverkey reference arrow) overlaid on each panel. + +See the module-level constants for the metric reference values and the +default mask-fraction columns used to filter samples sitting on bad +detector regions before computing statistics. +""" + +from __future__ import annotations + +__all__ = ["subtraction_quality_report"] + +import numpy as np +import pandas as pd + +from lsst.analysis.ap.skymapOverlay import make_affine_sky_to_xy, draw_skymap_outlines_mpl + +# Mask-fraction columns whose sum indicates a sample sits on an unusable +# region of the detector. The headline scalars are computed only on samples +# below ``bad_mask_threshold`` so the distribution tails reflect subtraction +# quality rather than edge / saturated pixels. +DEFAULT_BAD_MASK_COLUMNS = ( + "bad_mask_fraction", + "sat_mask_fraction", + "edge_mask_fraction", + "no_data_mask_fraction", +) + +# Default extra padding (in detector pixels) added to each panel beyond +# the autoscaled data limits, so labels anchored at the edge of the data have +# room to render without being clipped by the panel boundary. NOTE: the +# displayed figure is at significantly lower resolution than the original +# image. +_DEFAULT_PANEL_PADDING_PIX = 150 + +# (column, reference value, display label). Reference is the value the +# metric takes on a perfectly subtracted, well-decorrelated diffim. +_HEADLINE_METRICS = ( + ("diffim_chi2PerPix", 1.0, "Diffim chi^2/pix"), # noqa: E241 + ("psfMatchingKernel_residualNorm", 0.0, "PSF match residual"), # noqa: E241 + ("dipole_density", 0.0, "Dipoles / deg^2"), # noqa: E241 +) + + +def _coerce_to_frame(metrics): + """Accept an astropy Table, DataFrame, or any DataFrame-castable input.""" + if isinstance(metrics, pd.DataFrame): + return metrics + if hasattr(metrics, "to_pandas"): + return metrics.to_pandas() + return pd.DataFrame(metrics) + + +def _filter_clean(df, bad_mask_threshold, bad_mask_columns): + """Drop samples whose summed bad-mask fractions exceed the threshold.""" + cols = [c for c in bad_mask_columns if c in df.columns] + if not cols: + return df.copy() + return df[df[cols].sum(axis=1) < bad_mask_threshold].copy() + + +def _percentile_row(values): + """Return formatted (median, p84, p95, p99) strings for a 1-D array.""" + finite = values[np.isfinite(values)] + if finite.size == 0: + return ["nan"]*4 + q = np.percentile(finite, [50, 84, 95, 99]) + return [f"{x:.3f}" for x in q] + + +def _print_summary(clean, n_total, threshold): + """Print a fixed-width summary of the three headline metrics.""" + headers = ["metric", "ref", "median", "p84", "p95", "p99"] + rows = [] + for col, ref, _label in _HEADLINE_METRICS: + if col not in clean.columns: + rows.append([col, f"{ref:.2f}", "n/a", "n/a", "n/a", "n/a"]) + continue + rows.append([col, f"{ref:.2f}", *_percentile_row(clean[col].to_numpy())]) + widths = [max(len(str(r[i])) for r in [headers, *rows]) for i in range(len(headers))] + + def _fmt(row): + return " ".join(str(c).ljust(w) for c, w in zip(row, widths)) + + print(f"SpatiallySampledMetrics: {len(clean)}/{n_total} samples retained " + f"(bad-mask fraction sum < {threshold:g})") + print() + print(_fmt(headers)) + print(" ".join("-"*w for w in widths)) + for row in rows: + print(_fmt(row)) + + +def _metric_panel(ax, fig, clean, col, vmin, vmax, label, cmap, vcenter=None, + panel_padding_pix=_DEFAULT_PANEL_PADDING_PIX): + """Render one metric panel: linear-interpolated heatmap with sample + markers. + + Points whose (x, y, value) tuple contains any NaN are dropped before + interpolation. Grid cells outside the convex hull of the surviving + samples are left as NaN, which the colormap renders as transparent. + + If ``vcenter`` is provided and strictly between ``vmin`` and ``vmax``, + the panel uses a ``TwoSlopeNorm`` so the colormap midpoint always + maps to ``vcenter`` regardless of whether the (vmin, vmax) range is + symmetric about it. + """ + if col not in clean.columns or not clean[col].notna().any(): + ax.text(0.5, 0.5, f"{col}\n(not present)", + ha="center", va="center", transform=ax.transAxes) + ax.set_title(col) + return + + x = clean["x"].to_numpy() + y = clean["y"].to_numpy() + z = clean[col].to_numpy() + valid = np.isfinite(x) & np.isfinite(y) & np.isfinite(z) + if not valid.any(): + ax.text(0.5, 0.5, f"{col}\n(no valid samples)", + ha="center", va="center", transform=ax.transAxes) + ax.set_title(col) + return + xv, yv, zv = x[valid], y[valid], z[valid] + + # vmin and vmax autoscale independently: an explicit ``vmin=0.0`` from + # the caller is never overridden by the vmax-None branch. + if vmin is None: + vmin = 0.0 + if vmax is None: + vmax = float(np.nanpercentile(zv, 99)) + # Set a reasonable colorbar scale even if the data is completely constant. + if not (vmax > vmin): + vmax = vmin + 1.0 + + from scipy.interpolate import griddata + from matplotlib.colors import Normalize, TwoSlopeNorm + grid_size = 200 + xi = np.linspace(xv.min(), xv.max(), grid_size) + yi = np.linspace(yv.min(), yv.max(), grid_size) + XI, YI = np.meshgrid(xi, yi) + ZI = griddata((xv, yv), zv, (XI, YI), method="linear") + + # Center the colormap on ``vcenter`` when supplied and well-posed. + # TwoSlopeNorm requires vmin < vcenter < vmax strictly; if the + # autoscaled range collapses around vcenter we fall back to a plain + # clipping Normalize rather than raise. + if vcenter is not None and vmin < vcenter < vmax: + norm = TwoSlopeNorm(vcenter=vcenter, vmin=vmin, vmax=vmax) + else: + norm = Normalize(vmin=vmin, vmax=vmax, clip=True) + im = ax.imshow(ZI, origin="lower", + extent=(xv.min(), xv.max(), yv.min(), yv.max()), + norm=norm, cmap=cmap, aspect="equal") + fig.colorbar(im, ax=ax, label=label) + # Overlay the kernel centroid quiver so direction information is + # available next to every metric heatmap, with a quiverkey reference + # arrow. Sample markers go on top with higher "zorder" so the positions + # stay visible through arrows. + quiver_info = _overlay_kernel_quiver(ax, clean) + if quiver_info is not None: + q, ref_length = quiver_info + ax.quiverkey(q, 0.85, 1.05, ref_length, f"{ref_length:.2f}″", + labelpos="E", coordinates="axes") + ax.scatter(xv, yv, s=8, facecolors="white", edgecolors="black", + linewidths=0.3, zorder=5) + ax.set_title(col) + ax.set_xlabel("x [pix]") + ax.set_ylabel("y [pix]") + ax.set_aspect("equal") + # Pad the data limits relative to whatever the artists left them at, + # so labels (and quiver heads) sitting at the very edge of the data + # are not clipped by the panel boundary. + xmin, xmax = ax.get_xlim() + ymin, ymax = ax.get_ylim() + ax.set_xlim(xmin - panel_padding_pix, xmax + panel_padding_pix) + ax.set_ylim(ymin - panel_padding_pix, ymax + panel_padding_pix) + + +def _overlay_kernel_quiver(ax, clean): + """Draw the HSV-colored kernel centroid quiver onto ``ax``. + + Returns + ------- + info : tuple or None + ``(quiver_artist, ref_length_arcsec)`` on success, or None if the + required columns are missing or no samples are valid. The caller + decides whether to render a colorbar / quiverkey for it. + """ + needed = ("psfMatchingKernel_length", "psfMatchingKernel_direction", "x", "y") + if not all(c in clean.columns for c in needed): + return None + length = clean["psfMatchingKernel_length"].to_numpy() + direction = clean["psfMatchingKernel_direction"].to_numpy() + u = length*np.cos(direction) + v = length*np.sin(direction) + ok = np.isfinite(u) & np.isfinite(v) & (length > 0) + if not ok.any(): + return None + + ref_length = float(np.nanpercentile(length[ok], 95)) + direction_deg = np.degrees(direction[ok]) % 360.0 + x = clean["x"].to_numpy()[ok] + y = clean["y"].to_numpy()[ok] + q = ax.quiver(x, y, u[ok], v[ok], direction_deg, + cmap="hsv", clim=(0, 360), + angles="xy", scale_units="xy", + scale=ref_length/200, width=0.004, pivot="mid", + zorder=4) + return q, ref_length + + +def _draw_skymap_outlines(ax, skymap, clean, label_fontsize=7): + """Overlay patch boundaries (with tract,patch labels) on a panel. + + The metrics table carries no WCS, so the sky↔detector pixel mapping is + derived from the sample positions' ``(x, y)`` and + ``(coord_ra, coord_dec)`` columns via a least-squares affine fit (see + `~lsst.analysis.ap.skymapOverlay.make_affine_sky_to_xy`). For a single + detector this is typically accurate to well under a pixel — enough for + visualization but not for science. + + Silently no-ops when sky coordinates aren't available or fewer than + three samples are valid. + """ + if "coord_ra" not in clean.columns or "coord_dec" not in clean.columns: + return + + ra = clean["coord_ra"].to_numpy() + dec = clean["coord_dec"].to_numpy() + xs = clean["x"].to_numpy() + ys = clean["y"].to_numpy() + valid = np.isfinite(ra) & np.isfinite(dec) & np.isfinite(xs) & np.isfinite(ys) + if valid.sum() < 3: + return + + import lsst.geom as geom + + sky_to_xy = make_affine_sky_to_xy(ra[valid], dec[valid], xs[valid], ys[valid]) + + # Build a coord list spanning the sample footprint so findTractPatchList + # returns every tract / patch that touches it. + corner_pairs = [ + (ra[valid].min(), dec[valid].min()), + (ra[valid].max(), dec[valid].min()), + (ra[valid].max(), dec[valid].max()), + (ra[valid].min(), dec[valid].max()), + ] + sky_corners = [geom.SpherePoint(r, d, geom.radians) for r, d in corner_pairs] + + draw_skymap_outlines_mpl(ax, skymap, sky_to_xy, sky_corners, + label_fontsize=label_fontsize) + + +def _make_figure(clean, chi2_scale=1.0, skymap=None, label_fontsize=7, + panel_padding_pix=_DEFAULT_PANEL_PADDING_PIX): + """Build a three-panel diagnostic figure. + + Parameters + ---------- + clean : `pandas.DataFrame` + The filtered metrics table. + chi2_scale : `float`, optional + Half-width of the ``diffim_chi2PerPix`` colormap range, measured + multiplicatively about the nominal value of 1.0. ``vmax`` is set + to ``1 + chi2_scale`` and ``vmin`` to its reciprocal so the + colorbar covers the same factor above and below nominal in log space. + skymap : `lsst.skymap.BaseSkyMap`, optional + If supplied, overlay patch outlines (with tract,patch labels) on + every panel. + label_fontsize : `int` or `float`, optional + Font size for the per-patch ``tract,patch`` labels. + """ + import matplotlib.pyplot as plt + from matplotlib.colors import LinearSegmentedColormap + + # White-to-black sequential colormap for dipole density: pure white + # at the floor, pure black at vmax. NaN cells (outside the + # interpolation convex hull) are mapped to ``lightblue`` instead of + # any grey because every grey lives somewhere inside the white→black + # gradient and would otherwise read as a real mid-range value. + white_to_black = LinearSegmentedColormap.from_list( + "white_to_black", [(1.0, 1.0, 1.0), (0.0, 0.0, 0.0)]).copy() + white_to_black.set_bad("lightblue") + + chi2_vmax = 1.0 + chi2_scale + chi2_vmin = 1.0/chi2_vmax + + fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True) + # (col, vmin, vmax, label, cmap, vcenter). ``vcenter`` is non-None only + # for panels whose colormap should be anchored at a known reference + # value -- e.g. chi^2/pix is centered on its nominal value of 1.0. + panels = ( + ("diffim_chi2PerPix", chi2_vmin, chi2_vmax, "Diffim chi^2/pix", "RdBu_r", 1.0), # noqa: E241,E501 + ("psfMatchingKernel_residualNorm", 0.0, 0.2, "PSF match residual", "viridis", None), # noqa: E241,E501 + ("dipole_density", 0.0, None, "Dipoles / deg^2", white_to_black, None), # noqa: E241,E501 + ) + for ax, (col, vmin, vmax, label, cmap, vcenter) in zip(axes, panels): + _metric_panel(ax, fig, clean, col, vmin, vmax, label, cmap, vcenter=vcenter, + panel_padding_pix=panel_padding_pix) + if skymap is not None: + _draw_skymap_outlines(ax, skymap, clean, label_fontsize=label_fontsize) + return fig + + +def subtraction_quality_report(metrics, + bad_mask_threshold=0.2, + bad_mask_columns=DEFAULT_BAD_MASK_COLUMNS, + chi2_scale=1.0, + skymap=None, + label_fontsize=7, + panel_padding_pix=_DEFAULT_PANEL_PADDING_PIX): + """Print a headline metric summary and build the diagnostic plot. + + Parameters + ---------- + metrics : `astropy.table.Table` or `pandas.DataFrame` + The metrics table produced by ``SpatiallySampledMetricsTask`` for a + single detector. + bad_mask_threshold : `float`, optional + Samples whose summed mask fractions across ``bad_mask_columns`` + meet or exceed this value are dropped before computing statistics + and rendering the figure. + bad_mask_columns : iterable of `str`, optional + Names of mask-fraction columns to sum when deciding whether a + sample sits on a usable patch of the detector. Columns missing + from the input table are silently ignored. + chi2_scale : `float`, optional + Multiplicative half-width of the ``diffim_chi2PerPix`` colormap + around its nominal value of 1.0. The colorbar covers + ``[1/(1+chi2_scale), 1+chi2_scale]`` so a deviation by a factor + of ``(1+chi2_scale)`` above or below 1 sits at the colormap extremes. + skymap : `lsst.skymap.BaseSkyMap`, optional + If supplied, overlay the boundaries of every patch that touches + the detector footprint on each panel, with a ``tract,patch`` + label anchored just inside the lower-left corner of each patch. + The local sky↔pixel mapping is inferred from the sample + positions' ``(x, y)`` and ``(coord_ra, coord_dec)`` columns via + a least-squares affine fit, so the alignment is approximate + When omitted, no outlines are drawn. + label_fontsize : `int` or `float`, optional + Font size for the per-patch ``tract,patch`` labels. Only used + when ``skymap`` is supplied. Default 7. + panel_padding_pix : `float`, optional + Detector-pixel buffer added on each side of every panel beyond + the autoscaled data limits. + + Returns + ------- + fig : `matplotlib.figure.Figure` + The 1x3 diagnostic figure (`diffim_chi2PerPix`, PSF match + residual, dipole density). The kernel centroid offset quiver is + overlaid on every panel, with a quiverkey reference arrow. + + Notes + ----- + The headline summary is printed via ``print``, so the function is + intended for direct use in a notebook cell. Pass the returned figure + to ``fig.savefig(...)`` if you want to persist the diagnostic. + """ + df = _coerce_to_frame(metrics) + clean = _filter_clean(df, bad_mask_threshold, bad_mask_columns) + _print_summary(clean, n_total=len(df), threshold=bad_mask_threshold) + return _make_figure(clean, chi2_scale=chi2_scale, skymap=skymap, + label_fontsize=label_fontsize, + panel_padding_pix=panel_padding_pix) diff --git a/python/lsst/analysis/ap/taskRuntimes.py b/python/lsst/analysis/ap/taskRuntimes.py new file mode 100644 index 0000000..94181bf --- /dev/null +++ b/python/lsst/analysis/ap/taskRuntimes.py @@ -0,0 +1,160 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Collect per-quantum runtimes from the ``*_metadata`` datasets of a +butler run collection. + +The single public entry point, `collect_task_runtimes`, walks every +``_metadata`` dataset, pulls the timing fields via +`lsst.pipe.base.resource_usage.QuantumResourceUsage.from_task_metadata`, +applies a per-task threshold (any-quantum-over-threshold keeps the +whole task), and returns a tidy DataFrame, optionally with a box plot. +""" + +from __future__ import annotations + +__all__ = ["collect_task_runtimes"] + +import pandas as pd + +from lsst.pipe.base.resource_usage import QuantumResourceUsage + + +def collect_task_runtimes(butler, collections, threshold=1.0, *, + plot=False, ax=None): + """Per-task runtime and memory summary for a butler run collection. + + Each ``_metadata`` dataset under ``collections`` is loaded and + its timing fields extracted with + `~lsst.pipe.base.resource_usage.QuantumResourceUsage.from_task_metadata`. + Tasks whose every quantum runs faster than ``threshold`` seconds are + dropped; tasks with at least one quantum at or above ``threshold`` + contribute all their quanta to the per-task summary statistics, so the + summary reflects cross-quantum variability rather than a single outlier. + + Parameters + ---------- + butler : `lsst.daf.butler.Butler` + Butler used to query the registry and load metadata datasets. + collections : `str` or iterable of `str` + Collections to query, typically a single run collection name. + threshold : `float`, optional + Minimum task duration (seconds) for inclusion. A task is kept iff + at least one of its quanta has ``total_time`` >= ``threshold``. + Default ``1.0``. + plot : `bool`, optional + If True, render a horizontal box plot of per-quantum + ``total_time`` per surviving task (tasks ordered by max + ``total_time`` descending) and return the ``(df, fig)`` pair + instead of just ``df``. + ax : `matplotlib.axes.Axes` or None + Axes to plot onto. Only used when ``plot=True``. If None, a new + figure and axes are created. + + Returns + ------- + df : `pandas.DataFrame` + One row per task, with columns: + + - ``task``: pipeline task label + - ``n_quanta``: number of surviving quanta contributing to the row + - ``total_time_mean``, ``total_time_min``, ``total_time_max``, + ``total_time_std``: seconds + - ``memory_mean_``, ``memory_min_``, + ``memory_max_``, ``memory_std_``: where ``UNIT`` is + ``GB`` if any task's peak memory crosses 1 GB and ``MB`` + otherwise. The unit is chosen once across the whole table so + the columns remain numerically comparable. + fig : `matplotlib.figure.Figure` + Only returned when ``plot=True``. + """ + rows = [] + for dataset_type in butler.registry.queryDatasetTypes("*_metadata"): + task = dataset_type.name[:-len("_metadata")] + for ref in butler.registry.queryDatasets(dataset_type, collections=collections): + metadata = butler.get(ref) + try: + usage = QuantumResourceUsage.from_task_metadata(metadata) + except KeyError: + # Quantum block exists but is missing one of the expected + # fields (e.g. an aborted run); skip rather than fail. + continue + if usage is None: + continue + rows.append({"task": task, + "total_time": usage.total_time, + "memory": usage.memory}) + + if not rows: + return (pd.DataFrame(), None) if plot else pd.DataFrame() + + per_quantum = pd.DataFrame(rows) + per_task_max = per_quantum.groupby("task")["total_time"].max() + keep_tasks = per_task_max[per_task_max >= threshold].index + per_quantum = per_quantum[per_quantum["task"].isin(keep_tasks)] + + summary = (per_quantum.groupby("task") + .agg(n_quanta=("total_time", "count"), + total_time_mean=("total_time", "mean"), + total_time_min=("total_time", "min"), + total_time_max=("total_time", "max"), + total_time_std=("total_time", "std"), + memory_mean=("memory", "mean"), + memory_min=("memory", "min"), + memory_max=("memory", "max"), + memory_std=("memory", "std")) + .reset_index() + .sort_values("total_time_max", ascending=False) + .reset_index(drop=True)) + + mem_cols = ["memory_mean", "memory_min", "memory_max", "memory_std"] + # Use GB if any task's peak memory crosses 1 GB, otherwise MB. + if summary["memory_max"].max() >= 1024**3: + mem_unit, mem_divisor = "GB", 1024**3 + else: + mem_unit, mem_divisor = "MB", 1024**2 + summary[mem_cols] = summary[mem_cols] / mem_divisor + summary = summary.rename(columns={c: f"{c}_{mem_unit}" for c in mem_cols}) + + if not plot: + return summary + + import matplotlib.pyplot as plt + + # Bottom-up order so the slowest task sits at the top of the horizontal + # plot. + order = per_task_max.loc[keep_tasks].sort_values(ascending=True).index.tolist() + data = [per_quantum.loc[per_quantum["task"] == t, "total_time"].to_numpy() for t in order] + if ax is None: + fig, ax = plt.subplots(figsize=(8, max(3, 0.3 * len(order)))) + else: + fig = ax.figure + ax.boxplot(data, vert=False, tick_labels=order, showfliers=True) + ax.set_xlabel("total_time (s)") + if per_quantum["total_time"].max() > 100 * threshold: + ax.set_xscale("log") + ax.axvline(threshold, color="grey", linestyle=":", linewidth=1, + label=f"threshold = {threshold:g} s") + ax.set_title("Per-quantum total_time") + ax.grid(axis="x", linestyle=":", alpha=0.5) + ax.legend(loc="lower right", fontsize="small") + fig.tight_layout() + return summary, fig diff --git a/tests/test_apdb.py b/tests/test_apdb.py index c2d040d..d5eb2f4 100644 --- a/tests/test_apdb.py +++ b/tests/test_apdb.py @@ -156,6 +156,41 @@ def test_set_excluded_diaSource_flags(self): self.assertEqual(str(query.whereclause.compile(compile_kwargs={"literal_binds": True})), queryString) + def test_diaObjectId_int_when_null_present(self): + """diaObjectId must round-trip as an exact integer even when some + rows have SQL NULL diaObjectId (e.g. ssObject-only diaSources). + The 18-digit IDs exceed float64 precision (2**53), so a naive read + that downgrades the column to float64 would mangle them. The + fixture has no NULLs, so simulate one by patching a temp copy. + """ + import tempfile + import shutil + import sqlite3 + + src = os.path.join(os.path.dirname(os.path.abspath(__file__)), + "data", "apdb.sqlite3") + with tempfile.NamedTemporaryFile(suffix=".sqlite3", delete=False) as tmp: + shutil.copy(src, tmp.name) + path = tmp.name + try: + with sqlite3.connect(path) as conn: + conn.execute( + "UPDATE DiaSource SET diaObjectId = NULL " + "WHERE diaSourceId = 506428274000265217") + apdb = ApdbSqliteQuery(path) + result = apdb.load_sources(limit=None) + self.assertEqual(str(result["diaObjectId"].dtype), "Int64") + # The exact integer must be preserved (would round to + # 506428274000265216 if the column were float64). + row = result.loc[result["diaSourceId"] == 506428274000265218] + self.assertEqual(int(row["diaObjectId"].iloc[0]), + 506428274000265218) + # The nulled row must be pd.NA, not a mangled numeric value. + nulled = result.loc[result["diaSourceId"] == 506428274000265217] + self.assertTrue(pd.isna(nulled["diaObjectId"].iloc[0])) + finally: + os.unlink(path) + def test_fill_from_instrument(self): # an empty series should be unchanged empty = pd.Series() diff --git a/tests/test_apdbReconstruct.py b/tests/test_apdbReconstruct.py new file mode 100644 index 0000000..4abaf4f --- /dev/null +++ b/tests/test_apdbReconstruct.py @@ -0,0 +1,313 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +import unittest + +import lsst.utils.tests +import numpy as np +import pandas as pd + +from lsst.analysis.ap.apdbReconstruct import ( + ApdbReconstructor, + InMemoryDbQuery, +) + + +def _diaSources(): + """Two distinct diaSourceIds, plus one duplicate (later visit wins).""" + return pd.DataFrame({ + "diaSourceId": [100, 200, 100], + "diaObjectId": [10, 20, 10], + "visit": [1, 2, 3], + "detector": [50, 50, 50], + "ra": [1.0, 2.0, 1.5], + "dec": [-1.0, -2.0, -1.5], + "midpointMjdTai": [60100.0, 60101.0, 60110.0], + "psfFlux": [100.0, 200.0, 150.0], + "psfFluxErr": [10.0, 20.0, 15.0], + "band": ["g", "r", "g"], + "x": [1.0, 2.0, 1.0], + "y": [1.0, 2.0, 1.0], + "psfNdata": [10, 10, 10], + }) + + +def _diaObjects(): + """diaObjectId=10 appears twice with different validityStart; the later + snapshot should win after dedup. diaObjectId=20 appears once. + """ + return pd.DataFrame({ + "diaObjectId": [10, 20, 10], + "validityStartMjdTai": [60100.0, 60101.0, 60110.0], + "ra": [1.0, 2.0, 1.5], + "dec": [-1.0, -2.0, -1.5], + "nDiaSources": [1, 1, 2], + }) + + +def _diaForcedSources(): + """Three rows with the second duplicating the first on PK (latest wins).""" + return pd.DataFrame({ + "diaForcedSourceId": [1000, 1001, 1002], + "diaObjectId": [10, 10, 20], + "visit": [1, 1, 2], # row 0 and row 1 share PK + "detector": [50, 50, 50], + "midpointMjdTai": [60100.0, 60100.0, 60101.0], + "psfFlux": [100.0, 150.0, 200.0], # later (150) should win + "psfFluxErr": [10.0, 15.0, 20.0], + "band": ["g", "g", "r"], + "ra": [1.0, 1.0, 2.0], + "dec": [-1.0, -1.0, -2.0], + "scienceFlux": [50.0, 75.0, 100.0], + "scienceFluxErr": [5.0, 7.0, 10.0], + "timeProcessedMjdTai": [60101.0, 60102.0, 60102.0], + }) + + +class TestFinalize(lsst.utils.tests.TestCase): + """Tests for the staticmethod that does dedup + schema coercion.""" + + def test_diaSource_dedup_by_id(self): + result = ApdbReconstructor.finalize( + _diaSources(), _diaObjects(), _diaForcedSources(), + coerce_to_schema=False) + # 3 input rows -> 2 unique diaSourceIds. + self.assertEqual(len(result.diaSources), 2) + self.assertEqual(sorted(result.diaSources["diaSourceId"]), [100, 200]) + # The duplicate-keep="last" entry should win: visit 3 wins over visit 1 + row100 = result.diaSources.set_index("diaSourceId").loc[100] + self.assertEqual(int(row100["visit"]), 3) + + def test_diaForcedSource_dedup_by_pk(self): + result = ApdbReconstructor.finalize( + _diaSources(), _diaObjects(), _diaForcedSources(), + coerce_to_schema=False) + # 3 input rows -> 2 unique (diaObjectId, visit, detector). + self.assertEqual(len(result.diaForcedSources), 2) + # The later row for the dup PK should win (psfFlux=150, not 100). + dup = result.diaForcedSources[ + (result.diaForcedSources["diaObjectId"] == 10) + & (result.diaForcedSources["visit"] == 1) + & (result.diaForcedSources["detector"] == 50)] + self.assertEqual(len(dup), 1) + self.assertAlmostEqual(float(dup["psfFlux"].iloc[0]), 150.0) + + def test_diaObject_keeps_latest_by_validity(self): + result = ApdbReconstructor.finalize( + _diaSources(), _diaObjects(), _diaForcedSources(), + coerce_to_schema=False) + self.assertEqual(len(result.diaObjects), 2) + # diaObject 10 had two snapshots; the later (validityStart=60110) + # should win — nDiaSources=2, not 1. + row10 = result.diaObjects.set_index("diaObjectId").loc[10] + self.assertEqual(int(row10["nDiaSources"]), 2) + + def test_diaObject_dedup_skips_nan_nDiaSources(self): + """Passthrough snapshots from quanta that touch a diaObject but + don't actually update it leave ``nDiaSources`` as NaN. The dedup + must skip those in favor of an older snapshot that carries the + real count — otherwise the survivor's NaN gets fillna(0)'d during + schema coercion and the user sees ``nDiaSources=0`` even though + the diaObject has real diaSources. + """ + diaObjects = pd.DataFrame({ + "diaObjectId": [10, 10, 10, 20], # noqa: E241 + "validityStartMjdTai": [0.0, 61167.0, 61168.0, 61167.0], # noqa: E241 + "nDiaSources": [1.0, np.nan, np.nan, 2.0], # noqa: E241 + "ra": [1.0, 1.0, 1.0, 2.0], # noqa: E241 + "dec": [-1.0, -1.0, -1.0, -2.0], # noqa: E241 + }) + empty = pd.DataFrame() + result = ApdbReconstructor.finalize( + empty, diaObjects, empty, coerce_to_schema=False) + # diaObject 10: must survive with nDiaSources=1 (NOT NaN, NOT 0). + row10 = result.diaObjects.set_index("diaObjectId").loc[10] + self.assertEqual(int(row10["nDiaSources"]), 1) + # diaObject 20: untouched (only one row), should round-trip. + row20 = result.diaObjects.set_index("diaObjectId").loc[20] + self.assertEqual(int(row20["nDiaSources"]), 2) + + def test_diaObject_dedup_prefers_higher_nDiaSources(self): + """When a diaObject has multiple informative snapshots, dedup + picks the one with the most diaSources (which is also typically + the latest update; this ordering is robust against snapshot + re-ordering during concatenation). + """ + diaObjects = pd.DataFrame({ + "diaObjectId": [10, 10, 10], # noqa: E241 + "validityStartMjdTai": [60100.0, 60110.0, 60120.0], # noqa: E241 + "nDiaSources": [1.0, 3.0, 2.0], # noqa: E241 + "ra": [1.0, 1.0, 1.0], # noqa: E241 + "dec": [-1.0, -1.0, -1.0], # noqa: E241 + }) + empty = pd.DataFrame() + result = ApdbReconstructor.finalize( + empty, diaObjects, empty, coerce_to_schema=False) + row10 = result.diaObjects.set_index("diaObjectId").loc[10] + # The snapshot with nDiaSources=3 wins (highest count), even + # though a later validity is available. + self.assertEqual(int(row10["nDiaSources"]), 3) + + def test_diaObject_history_keeps_all(self): + result = ApdbReconstructor.finalize( + _diaSources(), _diaObjects(), _diaForcedSources(), + coerce_to_schema=False, history=True) + # Full update trail preserved. + self.assertEqual(len(result.diaObjects), 3) + + def test_schema_coercion_dtypes(self): + result = ApdbReconstructor.finalize( + _diaSources(), _diaObjects(), _diaForcedSources(), + coerce_to_schema=True) + # Integer IDs come out as Int64 (nullable long) or int64. + # diaSourceId is non-nullable -> int64; + # diaObjectId is nullable -> Int64. + self.assertEqual(str(result.diaSources["diaSourceId"].dtype), "int64") + self.assertEqual(str(result.diaSources["diaObjectId"].dtype), "Int64") + # Coercion also fills in schema columns that were missing — these + # ought to appear with sensible defaults. + self.assertIn("snr", result.diaSources.columns) + # Extra columns NOT in the schema are dropped by + # convertDataFrameToSdmSchema. + # (We didn't introduce any in the fixtures, but verify the call + # didn't add a bogus index column.) + self.assertNotIn("index", result.diaSources.columns) + + def test_empty_inputs(self): + empty = pd.DataFrame() + result = ApdbReconstructor.finalize(empty, empty, empty, + coerce_to_schema=False) + self.assertEqual(len(result.diaSources), 0) + self.assertEqual(len(result.diaObjects), 0) + self.assertEqual(len(result.diaForcedSources), 0) + + +class TestInMemoryDbQuery(lsst.utils.tests.TestCase): + """Tests that the DbQuery adapter routes queries against the underlying + DataFrames the way the SQL backends do. + """ + + def setUp(self): + recon = ApdbReconstructor.finalize( + _diaSources(), _diaObjects(), _diaForcedSources(), + coerce_to_schema=False) + self.query = InMemoryDbQuery(recon.diaSources, + recon.diaObjects, + recon.diaForcedSources) + + def test_load_sources_for_object(self): + result = self.query.load_sources_for_object(10) + self.assertEqual(len(result), 1) + self.assertEqual(int(result["diaSourceId"].iloc[0]), 100) + + def test_load_forced_sources_for_object_ignores_exclude_flagged(self): + # DiaForcedSource has no flag columns; exclude_flagged is a no-op + # for parity with the abstract interface. + result = self.query.load_forced_sources_for_object( + 10, exclude_flagged=True) + self.assertEqual(len(result), 1) + self.assertEqual(int(result["visit"].iloc[0]), 1) + + def test_load_source_raises_when_missing(self): + with self.assertRaisesRegex(RuntimeError, "diaSourceId=999999"): + self.query.load_source(999999) + + def test_load_object_round_trip(self): + obj = self.query.load_object(10) + self.assertEqual(int(obj["diaObjectId"]), 10) + self.assertEqual(int(obj["nDiaSources"]), 2) + + def test_load_forced_source_round_trip(self): + result = self.query.load_forced_source(1002) + self.assertEqual(int(result["diaForcedSourceId"]), 1002) + self.assertEqual(int(result["visit"]), 2) + + def test_excluded_flag_validation(self): + with self.assertRaisesRegex(ValueError, "not present"): + self.query.set_excluded_diaSource_flags(["pixelFlags_bad"]) + + def test_load_sources_with_exclude_flagged(self): + # Add a flag column to one row and verify it's excluded. + diaSrc = _diaSources() + diaSrc["pixelFlags_bad"] = [False, True, False] + recon = ApdbReconstructor.finalize( + diaSrc, _diaObjects(), _diaForcedSources(), + coerce_to_schema=False) + q = InMemoryDbQuery(recon.diaSources, recon.diaObjects, + recon.diaForcedSources) + q.set_excluded_diaSource_flags(["pixelFlags_bad"]) + # No exclusion requested: all 2 deduped rows returned. + self.assertEqual(len(q.load_sources()), 2) + # Exclusion requested: the flagged row is dropped. + flagged = q.load_sources(exclude_flagged=True) + self.assertEqual(len(flagged), 1) + self.assertNotIn(200, flagged["diaSourceId"].tolist()) + + +class TestDatasetNameDefaults(lsst.utils.tests.TestCase): + """Pin the dataset-name defaults to the ApPipe.yaml `associateApdb` + config, since downstream production tooling relies on these names. + """ + + def test_default_names_match_ap_pipe(self): + # The "apdb" entries come from ApPipe.yaml's `associateApdb` task + # config; the "preloaded_*" entries come from + # `lsst.ap.association.LoadDiaCatalogsTask` output names. Both + # are loaded so the reconstruction includes both prior history + # and the current run's new rows. + recon = ApdbReconstructor(butler=None) + self.assertEqual(recon.dataset_names["diaSource"], + ["dia_source_apdb", "preloaded_dia_source"]) + self.assertEqual(recon.dataset_names["diaObject"], + ["dia_object_apdb", "preloaded_dia_object"]) + self.assertEqual(recon.dataset_names["diaForcedSource"], + ["dia_forced_source_apdb", + "preloaded_dia_forced_source"]) + + def test_dataset_names_override(self): + # A full dict replaces DEFAULT_DATASET_NAMES wholesale. + override = { + "diaSource": ["goodSeeingDiff_assocDiaSrc"], + "diaObject": ["goodSeeingDiff_diaObject"], + "diaForcedSource": ["goodSeeingDiff_diaForcedSrc"], + } + recon = ApdbReconstructor(butler=None, dataset_names=override) + self.assertEqual(recon.dataset_names, override) + + def test_dataset_names_override_list(self): + # A list override replaces the default list entirely. + recon = ApdbReconstructor( + butler=None, + dataset_names={"diaSource": ["a", "b", "c"]}) + self.assertEqual(recon.dataset_names["diaSource"], ["a", "b", "c"]) + + +class MemoryTester(lsst.utils.tests.MemoryTestCase): + pass + + +def setup_module(module): + lsst.utils.tests.init() + + +if __name__ == "__main__": + lsst.utils.tests.init() + unittest.main() diff --git a/tests/test_compare.py b/tests/test_compare.py new file mode 100644 index 0000000..c0bf11e --- /dev/null +++ b/tests/test_compare.py @@ -0,0 +1,195 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +import unittest + +import astropy.units as u +import numpy as np +import pandas as pd + +import lsst.utils.tests +from lsst.analysis.ap.compare import ( + match_catalogs, flux_residuals, match_to_truth, +) + + +def _make_frame(rows, **defaults): + """Build a small DiaSource-like DataFrame from a list of partial dicts.""" + full = [] + for row in rows: + merged = dict(defaults) + merged.update(row) + full.append(merged) + return pd.DataFrame(full) + + +class TestMatchCatalogs(lsst.utils.tests.TestCase): + """Spatial cross-matching is the building block for compare_sources; + these tests pin down its behavior on small hand-crafted catalogs.""" + + def setUp(self): + # Three nearby points, two on (visit=1, detector=0), one on + # (visit=1, detector=1). All separations are tiny (~0.1 arcsec). + self.srcs1 = _make_frame([ + {"diaSourceId": 1, "ra": 10.0, "dec": -5.0, "visit": 1, "detector": 0}, + {"diaSourceId": 2, "ra": 10.001, "dec": -5.0, "visit": 1, "detector": 0}, + {"diaSourceId": 3, "ra": 11.0, "dec": -5.0, "visit": 1, "detector": 1}, + ]) + # srcs2: id 11 matches id 1 (same coord). id 12 is far from id 2. + # No row in srcs2 for (visit=1, detector=1), so id 3 must be unique1. + self.srcs2 = _make_frame([ + {"diaSourceId": 11, "ra": 10.0, "dec": -5.0, "visit": 1, "detector": 0}, + {"diaSourceId": 12, "ra": 10.5, "dec": -5.0, "visit": 1, "detector": 0}, + ]) + + def test_basic_match(self): + matched, unique1, unique2 = match_catalogs( + self.srcs1, self.srcs2, radius=1*u.arcsec) + # id=1 in srcs1 matches id=11 in srcs2. + self.assertEqual(list(matched["diaSourceId"]), [1]) + self.assertEqual(list(matched["diaSourceId_2"]), [11]) + # The match is essentially zero arcsec apart. + self.assertLess(float(matched["xmatch_dist_arcsec"].iloc[0]), 1e-6) + # ids 2, 3 from srcs1 are unique; ids 12 from srcs2 is unique. + self.assertEqual(set(unique1["diaSourceId"]), {2, 3}) + self.assertEqual(set(unique2["diaSourceId"]), {12}) + + def test_radius_units(self): + # Bare-float radius is interpreted as arcseconds. + matched_q, _, _ = match_catalogs(self.srcs1, self.srcs2, + radius=1*u.arcsec) + matched_f, _, _ = match_catalogs(self.srcs1, self.srcs2, radius=1.0) + self.assertEqual(len(matched_q), len(matched_f)) + + def test_no_grouping(self): + # With on=(), sources match across visit/detector boundaries. + srcs1 = _make_frame([ + {"diaSourceId": 1, "ra": 10.0, "dec": 0.0, "visit": 1, "detector": 0}, + ]) + srcs2 = _make_frame([ + {"diaSourceId": 2, "ra": 10.0, "dec": 0.0, "visit": 99, "detector": 99}, + ]) + # With grouping, no match (different visits/detectors). + matched, _, _ = match_catalogs(srcs1, srcs2, radius=1*u.arcsec) + self.assertEqual(len(matched), 0) + # Without grouping, the spatial match succeeds. + matched, _, _ = match_catalogs(srcs1, srcs2, radius=1*u.arcsec, on=()) + self.assertEqual(len(matched), 1) + + def test_empty_inputs(self): + empty = self.srcs1.iloc[0:0] + matched, u1, u2 = match_catalogs(empty, self.srcs2, radius=1*u.arcsec) + self.assertEqual(len(matched), 0) + self.assertEqual(len(u1), 0) + self.assertEqual(set(u2["diaSourceId"]), {11, 12}) + + matched, u1, u2 = match_catalogs(self.srcs1, empty, radius=1*u.arcsec) + self.assertEqual(len(matched), 0) + self.assertEqual(set(u1["diaSourceId"]), {1, 2, 3}) + self.assertEqual(len(u2), 0) + + +class TestFluxResiduals(lsst.utils.tests.TestCase): + """`flux_residuals` joins matched pairs with catalog 2 to compute + flux differences in units of sigma.""" + + def test_zero_residuals_when_identical(self): + srcs1 = _make_frame([ + {"diaSourceId": 1, "ra": 10.0, "dec": 0.0, "visit": 1, "detector": 0, + "psfFlux": 100.0, "psfFluxErr": 5.0}, + {"diaSourceId": 2, "ra": 10.001, "dec": 0.0, "visit": 1, "detector": 0, + "psfFlux": 200.0, "psfFluxErr": 10.0}, + ]) + srcs2 = _make_frame([ + {"diaSourceId": 11, "ra": 10.0, "dec": 0.0, "visit": 1, "detector": 0, + "psfFlux": 100.0, "psfFluxErr": 5.0}, + {"diaSourceId": 12, "ra": 10.001, "dec": 0.0, "visit": 1, "detector": 0, + "psfFlux": 200.0, "psfFluxErr": 10.0}, + ]) + matched, _, _ = match_catalogs(srcs1, srcs2, radius=1*u.arcsec) + residuals = flux_residuals(matched, srcs2) + np.testing.assert_array_equal(residuals["delta_flux"], [0.0, 0.0]) + np.testing.assert_array_equal(residuals["delta_flux_sigma"], [0.0, 0.0]) + + def test_known_offset(self): + # Set up a 3-sigma flux offset on a single matched pair. + srcs1 = _make_frame([ + {"diaSourceId": 1, "ra": 10.0, "dec": 0.0, "visit": 1, "detector": 0, + "psfFlux": 130.0, "psfFluxErr": 5.0}, + ]) + srcs2 = _make_frame([ + {"diaSourceId": 11, "ra": 10.0, "dec": 0.0, "visit": 1, "detector": 0, + "psfFlux": 100.0, "psfFluxErr": 5.0}, + ]) + matched, _, _ = match_catalogs(srcs1, srcs2, radius=1*u.arcsec) + residuals = flux_residuals(matched, srcs2) + self.assertEqual(float(residuals["delta_flux"].iloc[0]), 30.0) + # Combined sigma is sqrt(5^2 + 5^2); 30/sqrt(50) ~ 4.243. + self.assertAlmostEqual( + float(residuals["delta_flux_sigma"].iloc[0]), + 30.0 / np.sqrt(50.0), + places=10, + ) + + +class TestMatchToTruth(lsst.utils.tests.TestCase): + """`match_to_truth` reports purity and completeness against a truth + catalog by matching in both directions.""" + + def setUp(self): + # Three detected sources: two near truth, one in empty sky. + self.srcs = _make_frame([ + {"diaSourceId": 1, "ra": 10.0, "dec": 0.0}, + {"diaSourceId": 2, "ra": 11.0, "dec": 0.0}, + {"diaSourceId": 3, "ra": 99.0, "dec": -50.0}, # bogus, not real + ]) + # Three truth sources: two recovered, one missed. + self.truth = _make_frame([ + {"injection_id": 100, "ra": 10.0, "dec": 0.0}, + {"injection_id": 101, "ra": 11.0, "dec": 0.0}, + {"injection_id": 102, "ra": 50.0, "dec": -10.0}, # not recovered + ]) + + def test_purity_and_completeness(self): + result = match_to_truth(self.srcs, self.truth, radius=1*u.arcsec) + # 2 of 3 detections are real; 2 of 3 truths are detected. + self.assertAlmostEqual(result["purity"], 2/3, places=10) + self.assertAlmostEqual(result["completeness"], 2/3, places=10) + # The bogus detection has is_real=False and a NA partner. + bogus = result["srcs"].set_index("diaSourceId").loc[3] + self.assertFalse(bool(bogus["is_real"])) + self.assertTrue(pd.isna(bogus["injection_id_match"])) + # The unrecovered truth has detected=False. + miss = result["truth"].set_index("injection_id").loc[102] + self.assertFalse(bool(miss["detected"])) + + +class TestMemory(lsst.utils.tests.MemoryTestCase): + pass + + +def setup_module(module): + lsst.utils.tests.init() + + +if __name__ == "__main__": + lsst.utils.tests.init() + unittest.main() diff --git a/tests/test_plotDiaSourceLightcurve.py b/tests/test_plotDiaSourceLightcurve.py new file mode 100644 index 0000000..2119843 --- /dev/null +++ b/tests/test_plotDiaSourceLightcurve.py @@ -0,0 +1,332 @@ +# This file is part of analysis_ap. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +import unittest + +import lsst.afw.image +import lsst.geom +import lsst.meas.base.tests +import lsst.utils.tests +import numpy as np +import pandas as pd +import PIL + +from lsst.analysis.ap import plotDiaSourceLightcurve + +# Pull in the DATA fixture from the cutouts test so we get the full set of +# flag columns expected by _annotate_image. +from test_plotImageSubtractionCutouts import DATA, skyCenter + + +DIA_OBJECT_ID = 999999999999000001 + + +# Add the extra columns required by the lightcurve task. +def _augment_with_object_columns(data): + data = data.copy() + data["diaObjectId"] = [DIA_OBJECT_ID, DIA_OBJECT_ID] + data["midpointMjdTai"] = [60100.5, 60110.5] + return data + + +class _StubApdbQuery: + """Minimal DbQuery stub returning canned sources/forced DataFrames. + + Tracks the number of calls so tests can verify the per-diaObject cache. + """ + + def __init__(self, sources, forced): + self._sources = sources + self._forced = forced + self.source_calls = 0 + self.forced_calls = 0 + + def load_sources_for_object(self, dia_object_id, exclude_flagged=False, limit=100000): + self.source_calls += 1 + self.last_source_kwargs = {"exclude_flagged": exclude_flagged, "limit": limit} + return self._sources.copy() + + def load_forced_sources_for_object(self, dia_object_id, exclude_flagged=False, limit=100000): + self.forced_calls += 1 + self.last_forced_kwargs = {"exclude_flagged": exclude_flagged, "limit": limit} + return self._forced.copy() + + +def _make_sources_frame(visits, bands, mjds, fluxes, errs, with_err=True): + frame = pd.DataFrame({ + "visit": visits, + "band": bands, + "midpointMjdTai": mjds, + "psfFlux": fluxes, + }) + if with_err: + frame["psfFluxErr"] = errs + return frame + + +class TestPlotDiaSourceLightcurve(lsst.utils.tests.TestCase): + """Tests for PlotDiaSourceLightcurveTask.""" + + def setUp(self): + bbox = lsst.geom.Box2I(lsst.geom.Point2I(0, 0), lsst.geom.Point2I(100, 100)) + self.centroid = lsst.geom.Point2D(50, 50) + dataset = lsst.meas.base.tests.TestDataset(bbox, crval=skyCenter) + self.scale = 0.3 + dataset.addSource(instFlux=1e5, centroid=self.centroid) + self.science, _ = dataset.realize(noise=1000.0, schema=dataset.makeMinimalSchema()) + self.template, _ = dataset.realize(noise=5.0, schema=dataset.makeMinimalSchema()) + self.difference = lsst.afw.image.ExposureF(self.science, deep=True) + self.difference.image -= self.template.image + self.data = _augment_with_object_columns(DATA) + + def _make_task(self, sources=None, forced=None, **config_kwargs): + config = plotDiaSourceLightcurve.PlotDiaSourceLightcurveConfig() + for key, value in config_kwargs.items(): + setattr(config, key, value) + query = None + if sources is not None or forced is not None: + query = _StubApdbQuery( + sources if sources is not None else pd.DataFrame(), + forced if forced is not None else pd.DataFrame(), + ) + task = plotDiaSourceLightcurve.PlotDiaSourceLightcurveTask( + config=config, output_path="", apdb_query=query) + return task, query + + def _record(self, row): + """Return a single-row numpy.record matching DataFrame ``iloc``.""" + return self.data.iloc[[row]].to_records(index=False)[0] + + def test_generate_image_no_apdb(self): + """Without an APDB handle, the lightcurve panel renders a placeholder + but the figure still produces a valid PNG. + """ + task, _ = self._make_task() + cutout = task.generate_image(self.science, self.template, self.difference, + skyCenter, self.scale, + source=self._record(0)) + with PIL.Image.open(cutout) as im: + self.assertGreater(im.width, 0) + self.assertGreater(im.height, 0) + + def test_generate_image_with_lightcurve(self): + """With matching sources and forced rows, both groups render.""" + sources = _make_sources_frame( + visits=[1234, 5678, 9999], + bands=["r", "g", "r"], + mjds=[60100.5, 60110.5, 60120.5], + fluxes=[1234.5, 2345.6, 3456.7], + errs=[123.5, 234.5, 345.6], + ) + # Forced has a visit (8888) that does NOT appear in sources — that + # one should be drawn with the forced-only marker. Visit 1234 IS in + # sources, so the forced entry for it should be suppressed. + forced = _make_sources_frame( + visits=[1234, 8888], + bands=["r", "g"], + mjds=[60100.5, 60115.5], + fluxes=[1200.0, 800.0], + errs=[120.0, 80.0], + ) + task, query = self._make_task(sources=sources, forced=forced) + cutout = task.generate_image(self.science, self.template, self.difference, + skyCenter, self.scale, + source=self._record(0)) + self.assertEqual(query.source_calls, 1) + self.assertEqual(query.forced_calls, 1) + with PIL.Image.open(cutout) as im: + self.assertGreater(im.width, 0) + self.assertGreater(im.height, 0) + + def test_lightcurve_cache_reuses_query(self): + """Adjacent diaSources on the same diaObject should hit the cache.""" + sources = _make_sources_frame( + visits=[1234, 5678], + bands=["r", "g"], + mjds=[60100.5, 60110.5], + fluxes=[1234.5, 2345.6], + errs=[123.5, 234.5], + ) + forced = pd.DataFrame(columns=["visit", "band", "midpointMjdTai", + "psfFlux", "psfFluxErr"]) + task, query = self._make_task(sources=sources, forced=forced) + for i in range(2): + task.generate_image(self.science, self.template, self.difference, + skyCenter, self.scale, + source=self._record(i)) + # Both sources share the same diaObjectId; the second call must come + # from the cache. + self.assertEqual(query.source_calls, 1) + self.assertEqual(query.forced_calls, 1) + + def test_forced_only_dedup_by_visit(self): + """Forced rows whose ``visit`` matches a diaSource are suppressed.""" + sources = _make_sources_frame( + visits=[1234, 5678], + bands=["r", "g"], + mjds=[60100.5, 60110.5], + fluxes=[1234.5, 2345.6], + errs=[123.5, 234.5], + ) + forced = _make_sources_frame( + visits=[1234, 5678, 7777, 8888], + bands=["r", "g", "r", "g"], + mjds=[60100.5, 60110.5, 60112.5, 60115.5], + fluxes=[1200.0, 2400.0, 500.0, 800.0], + errs=[120.0, 240.0, 50.0, 80.0], + ) + task, _ = self._make_task(sources=sources, forced=forced) + forced_only = forced[~forced["visit"].isin(sources["visit"])] + self.assertEqual(set(forced_only["visit"]), {7777, 8888}) + # Also exercise the rendering path to confirm it does not raise. + task.generate_image(self.science, self.template, self.difference, + skyCenter, self.scale, source=self._record(0)) + + def test_no_psfFluxErr(self): + """Missing or NaN psfFluxErr should fall back to no error bars.""" + sources = _make_sources_frame( + visits=[1234, 5678], + bands=["r", "g"], + mjds=[60100.5, 60110.5], + fluxes=[1234.5, 2345.6], + errs=None, with_err=False, + ) + forced = _make_sources_frame( + visits=[7777], + bands=["r"], + mjds=[60112.5], + fluxes=[500.0], + errs=[np.nan], with_err=True, + ) + task, _ = self._make_task(sources=sources, forced=forced) + cutout = task.generate_image(self.science, self.template, self.difference, + skyCenter, self.scale, + source=self._record(0)) + with PIL.Image.open(cutout) as im: + self.assertGreater(im.width, 0) + + def test_forced_query_skips_exclude_flagged(self): + """DiaForcedSource lacks the diaSource flag columns; the forced + query must always be called with exclude_flagged=False, even when + the config asks to exclude flagged diaSources. + """ + sources = _make_sources_frame( + visits=[1234], bands=["r"], mjds=[60100.5], + fluxes=[1234.5], errs=[123.5]) + forced = pd.DataFrame(columns=["visit", "band", "midpointMjdTai", + "psfFlux", "psfFluxErr"]) + task, query = self._make_task(sources=sources, forced=forced, + lightcurve_exclude_flagged=True) + task.generate_image(self.science, self.template, self.difference, + skyCenter, self.scale, + source=self._record(0)) + self.assertTrue(query.last_source_kwargs["exclude_flagged"]) + self.assertFalse(query.last_forced_kwargs["exclude_flagged"]) + + def test_nan_dia_object_id(self): + """A NaN diaObjectId (unassociated diaSource) must not crash; the + lightcurve panel renders its empty placeholder. + """ + data = self.data.copy() + data["diaObjectId"] = data["diaObjectId"].astype(float) + data.loc[0, "diaObjectId"] = np.nan + sources = _make_sources_frame( + visits=[1234], bands=["r"], mjds=[60100.5], + fluxes=[1234.5], errs=[123.5]) + forced = pd.DataFrame(columns=["visit", "band", "midpointMjdTai", + "psfFlux", "psfFluxErr"]) + task, query = self._make_task(sources=sources, forced=forced) + record = data.iloc[[0]].to_records(index=False)[0] + cutout = task.generate_image(self.science, self.template, self.difference, + skyCenter, self.scale, source=record) + # APDB was never queried — diaObjectId is NaN. + self.assertEqual(query.source_calls, 0) + self.assertEqual(query.forced_calls, 0) + with PIL.Image.open(cutout) as im: + self.assertGreater(im.width, 0) + + def test_forced_legend_single_black_entry(self): + """Forced points are colored per-band, but contribute a single + black legend entry regardless of band count. + """ + sources = _make_sources_frame( + visits=[1234, 5678], + bands=["r", "g"], + mjds=[60100.5, 60110.5], + fluxes=[1234.5, 2345.6], + errs=[123.5, 234.5], + ) + # Forced-only visits span two bands: should still produce one entry. + forced = _make_sources_frame( + visits=[7777, 8888], + bands=["r", "g"], + mjds=[60112.5, 60115.5], + fluxes=[500.0, 800.0], + errs=[50.0, 80.0], + ) + task, _ = self._make_task(sources=sources, forced=forced) + import matplotlib.pyplot as plt + fig, ax = plt.subplots() + try: + task._draw_lightcurve(ax, sources, forced, + current_source=self._record(0)) + handles, labels = ax.get_legend_handles_labels() + forced_labels = [lbl for lbl in labels if "forced" in lbl] + self.assertEqual(len(forced_labels), 1) + self.assertIn("(n=2)", forced_labels[0]) + # The single forced legend handle should be drawn in black. + forced_handle = handles[labels.index(forced_labels[0])] + self.assertEqual(forced_handle.get_color(), "black") + finally: + plt.close(fig) + + def test_njobs_downgraded(self): + """Requesting multiprocessing should be downgraded with a warning.""" + sources = _make_sources_frame( + visits=[1234], bands=["r"], mjds=[60100.5], + fluxes=[1234.5], errs=[123.5]) + forced = pd.DataFrame(columns=["visit", "band", "midpointMjdTai", + "psfFlux", "psfFluxErr"]) + task, _ = self._make_task(sources=sources, forced=forced) + # write_images would normally attempt multiprocessing; the override + # should silently drop njobs to 0 and not raise. + with self.assertLogs(task.log.name, level="WARNING") as ctx: + # Use just one row so we don't need real files on disk; the + # cutouts task will try to look them up via butler_cache and + # fail, but the warning should fire first. + try: + task.write_images(self.data.head(0), butler=None, njobs=4) + except Exception: + pass + self.assertTrue(any("njobs" in msg for msg in ctx.output)) + + +class MemoryTester(lsst.utils.tests.MemoryTestCase): + pass + + +def setup_module(module): + lsst.utils.tests.init() + + +if __name__ == "__main__": + lsst.utils.tests.init() + unittest.main() diff --git a/ups/analysis_ap.table b/ups/analysis_ap.table index 00197d0..4bd9529 100644 --- a/ups/analysis_ap.table +++ b/ups/analysis_ap.table @@ -13,6 +13,8 @@ setupRequired(pipe_base) setupRequired(utils) setupRequired(analysis_tools) setupRequired(ap_association) +setupRequired(pipe_tasks) +setupRequired(sdm_schemas) # TODO: DM-39501: to mock a dimension packer, until detector/visit are in APDB. setupOptional(obs_lsst)