blend.py

Go to the documentation of this file.

from __future__ import annotations
 
import logging
from dataclasses import dataclass
from functools import cached_property
from typing import Any
 
import numpy as np
from numpy.typing import DTypeLike
 
from ..bbox import Box
from ..blend import Blend
from ..observation import Observation
from .blend_base import ScarletBlendBaseData
from .migration import PRE_SCHEMA, MigrationRegistry, migration
from .source import ScarletSourceBaseData
from .utils import decode_metadata, encode_metadata, extract_from_metadata
 
__all__ = ["ScarletBlendData"]
 
CURRENT_SCHEMA = "1.0.0"
BLEND_TYPE = "blend"
MigrationRegistry.set_current(BLEND_TYPE, CURRENT_SCHEMA)
logger = logging.getLogger(__name__)
 
 
@dataclass(kw_only=True)
class ScarletBlendData(ScarletBlendBaseData):
    """Data for an entire blend.
 
    Attributes
    ----------
    blend_type :
        The type of blend being stored.
    metadata :
        Metadata associated with the blend,
        for example the order of bands, the PSF, etc.
    origin :
        The lower bound of the blend's bounding box.
    shape :
        The shape of the blend's bounding box.
    sources :
        Data for the sources contained in the blend,
        indexed by the source id.
    version :
        The schema version of the stored data.
    """
 
    blend_type: str = BLEND_TYPE
    origin: tuple[int, int]
    shape: tuple[int, int]
    sources: dict[Any, ScarletSourceBaseData]
    version: str = CURRENT_SCHEMA
 
    @cached_property
    def bbox(self) -> Box:
        """The bounding box of the blend"""
        return Box(self.shape, origin=self.origin)
 
    def as_dict(self) -> dict:
        """Return the object encoded into a dict for JSON serialization
 
        Returns
        -------
        result :
            The object encoded as a JSON compatible dict
        """
        result: dict[str, Any] = {
            "blend_type": self.blend_type,
            "origin": self.origin,
            "shape": self.shape,
            "sources": {bid: source.as_dict() for bid, source in self.sources.items()},
            "version": self.version,
        }
        if self.metadata is not None:
            result["metadata"] = encode_metadata(self.metadata)
        return result
 
    @classmethod
    def from_dict(cls, data: dict, dtype: DTypeLike = np.float32) -> ScarletBlendData:
        """Reconstruct `ScarletBlendData` from JSON compatible
        dict.
 
        Parameters
        ----------
        data :
            Dictionary representation of the object
        dtype :
            Datatype of the resulting model.
 
        Returns
        -------
        result :
            The reconstructed object
        """
        data = MigrationRegistry.migrate(BLEND_TYPE, data)
        metadata = data.get("metadata", None)
 
        return cls(
            origin=tuple(data["origin"]),  # type: ignore
            shape=tuple(data["shape"]),  # type: ignore
            sources={
                bid: ScarletSourceBaseData.from_dict(source, dtype=dtype)
                for bid, source in data["sources"].items()
            },
            metadata=decode_metadata(metadata),
        )
 
    def minimal_data_to_blend(
        self,
        model_psf: np.ndarray | None = None,
        psf: np.ndarray | None = None,
        bands: tuple[str] | None = None,
        dtype: DTypeLike = np.float32,
    ) -> Blend:
        """Convert the storage data model into a scarlet lite blend
 
        Parameters
        ----------
        model_psf :
            PSF in model space (usually a nyquist sampled circular Gaussian).
        psf :
            The PSF of the observation.
            If not provided, the PSF stored in the blend data is used.
        bands :
            The bands in the blend model.
            If not provided, the bands stored in the blend data are used.
        dtype :
            The data type of the model that is generated.
 
        Returns
        -------
        blend :
            A scarlet blend model extracted from persisted data.
        """
 
        _model_psf: np.ndarray = extract_from_metadata(model_psf, self.metadata, "model_psf")
        _psf: np.ndarray = extract_from_metadata(psf, self.metadata, "psf")
        _bands: tuple[str] = extract_from_metadata(bands, self.metadata, "bands")
        model_box = self.bbox
        observation = Observation.empty(
            bands=_bands,
            psfs=_psf,
            model_psf=_model_psf,
            bbox=model_box,
            dtype=dtype,
        )
        return self.to_blend(observation)
 
    def to_blend(self, observation: Observation) -> Blend:
        """Convert the storage data model into a scarlet lite blend
 
        Parameters
        ----------
        observation :
            The observation that contains the blend.
            If `observation` is ``None`` then an `Observation` containing
            no image data is initialized.
 
        Returns
        -------
        blend :
            A scarlet blend model extracted from persisted data.
        """
        sources = []
        for source_data in self.sources.values():
            source = source_data.to_source(observation)
            sources.append(source)
 
        return Blend(sources=sources, observation=observation, metadata=self.metadata)
 
    @staticmethod
    def from_blend(blend: Blend) -> ScarletBlendData:
        """Deprecated: Convert a scarlet lite blend into a storage data model.
 
        Parameters
        ----------
        blend :
            The blend to convert.
        Returns
        -------
        result :
            The storage data model representing the blend.
        """
        logger.warning("ScarletBlendData.from_blend is deprecated. Use blend.to_data() instead.")
        return blend.to_data()
 
 
ScarletBlendData.register()
 
 
@migration(BLEND_TYPE, PRE_SCHEMA)
def _to_1_0_0(data: dict) -> dict:
    """Migrate a pre-schema blend to schema version 1.0.0
 
    Parameters
    ----------
    data :
        The data to migrate.
 
    Returns
    -------
    result :
        The migrated data.
    """
    # Support legacy models before metadata was used
    if "metadata" not in data and "psf" in data:
        data["metadata"] = {
            "psf": data["psf"],
            "psf_shape": data["psf_shape"],
            "bands": tuple(data["bands"]),
            "array" "_keys": ["psf"],
        }
    data["version"] = "1.0.0"
    return data