from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from ....structures.Lightcurve import Lightcurve
from dataclasses import dataclass, field
from pathlib import Path
from typing import Self
import astropy.units as u
from astropy.coordinates import SkyCoord
from astropy.units import Quantity
from bokeh.plotting import figure as Figure
from ..utilities.docstrings import get_docstring
from .Base import COMBINE_PLOTS, SPLIT_BY_SURVEY, SPLIT_BY_TARGET, Container, manage_inplace
from .Target import Target
[docs]
@dataclass
class DataSet:
#: Kind of data container that is being stored.
kind: str | None = None
#: :class:`~ATK.Models.Target`\ s for which data is stored.
targets: list[Target] | None = field(default_factory=list)
#: If ``True``, something unexpected occurred during data retrieval.
exception: bool | None = False
#: Stored data as a list of containers.
data: list = field(default_factory=list)
#: Bokeh figure, generated by :meth:`~ATK.Models.DataSet.plot` or :meth:`~ATK.Models.DataSet.open`.
figure: Figure | None = None
_stored_plot_params: dict = field(default_factory=dict)
_plotted_keys: list = field(default_factory=list)
# maps per-Target key to Target
_key_map: dict[str, Target] = field(init=False, default_factory=dict)
# maps per-Target alias to per-Target key
_alias_map: dict[str, str] = field(init=False, default_factory=dict)
_cache_key: str | None = None
@property
def empty(self) -> bool:
"""
If ``True``, no data containers are being stored.
"""
if not self.data:
return True
else:
return False
def __post_init__(self):
self._build_target_maps()
def __repr__(self):
if not self.data:
return "<Empty DataSet>"
else:
return f"<{self.kind} DataSet>"
def __str__(self):
return self.__repr__()
def _build_target_maps(self):
self._key_map = {t._key: t for t in self.targets}
self._alias_map = {}
for t in self.targets:
for alias in t._aliases:
self._alias_map[alias] = t._key
[docs]
def show(self, show_types: bool = False, show_all: bool = False, **kwargs) -> None:
from ..io.struct_stdout import pprint_structure
pprint_structure(self, show_types, show_all, **kwargs)
show.__doc__ = get_docstring("show")
[docs]
def store(self, path: Path | str) -> Self:
"""
Saves :class:`~ATK.Models.DataSet` to a local FITS file.
Parameters
----------
path: Path | str
Path to which local FITS file should be saved.
The resulting local FITS file can be used to recreate the original structure with :func:`~ATK.Tools.read`.
Returns
-------
``self``
"""
from ..io.files.writing import write_local
write_local(self, path)
return self
[docs]
def plot(self, **kwargs) -> Self:
"""
Plots all stored data containers in a grid.
The resulting figure is stored in the :attr:`~ATK.Models.DataSet.figure` attribute.
Parameters
----------
**kwargs
Accepted keyword arguments depend on the ``kind`` of data being plotted.
See the documentation of the data container that is being plotted for more information.
Returns
-------
``self``
"""
if not self._ctnr_kind:
return self
if self._ctnr_kind not in COMBINE_PLOTS:
raise ValueError(f"{self._ctnr_kind} containers do not support plotting.")
from ..io.plot_io import plot_data
keys = []
for target in self.targets:
keys.append(target._key)
self.figure = plot_data(self, **kwargs)
self._stored_plot_params = kwargs
self._plotted_keys = keys
return self
[docs]
def open(self, path: Path | str | None = None, **kwargs) -> Self:
"""
Opens the figure from the :attr:`~ATK.Models.DataSet.figure` attribute in the default browser.
If a figure has not yet been generated when :meth:`ATK.Models.DataSet.open` is called, one will be generated with :meth:`~ATK.Models.DataSet.plot`.
Optionally, the figure can also be saved to local files.
Parameters
----------
path: Path | str, optional
Path to which the figure should be saved.
If not provided, figures are saved to the ``~/.AstroToolkit/cached_figures`` directory.
**kwargs
If a plot needs to be generated (see above), additional keyword arguments are passed to :meth:`~ATK.Models.DataSet.plot`.
Accepted keyword arguments depend on the ``kind`` of data being plotted. See the documentation of the data container that is being plotted for more information.
Returns
-------
``self``
"""
from ..io.plot_io import open as open_html
keys = []
for target in self.targets:
keys.append(target._key)
open_html(self, fname=path, keys=keys, **kwargs)
return self
[docs]
def save(self, path: Path | str, **kwargs) -> Self:
"""
Saves the figure from the :attr:`~ATK.Models.DataSet.figure` attribute to local files.
If a figure has not yet been generated when :meth:`ATK.Models.DataSet.open` is called, one will be generated with :meth:`~ATK.Models.DataSet.plot`.
Parameters
----------
path: Path | str, optional
Path to which the figure should be saved.
**kwargs
If a plot needs to be generated (see above), additional keyword arguments are passed to :meth:`~ATK.Models.DataSet.plot`.
Accepted keyword arguments depend on the ``kind`` of data being plotted. See the documentation of the data container that is being plotted for more information.
Returns
-------
``self``
"""
from ..io.plot_io import save
keys = []
for target in self.targets:
keys.append(target._key)
save(self, fname=path, keys=keys, **kwargs)
return self
[docs]
def apply(self, method: str, *args, inplace=True, **kwargs) -> Self:
"""
Applies a **data method** to all stored containers.
The set of available **data methods** depends on the ``kind`` of data that is stored. See the documentation of the stored data container for more information.
Parameters
----------
method: str
Name of **data method** that should be applied.
inplace : bool, optional
If ``True``, modify the current :class:`~ATK.Models.DataSet` inplace. If ``False``, operate on and return a copy - leaving the original unchanged.
**kwargs
Accepted keyword arguments depend on the chosen **data method**. See the documentation of the **data method** for more information.
Returns
-------
``Self``
:class:`~ATK.Models.DataSet` with modified data containers. Returns ``self`` if ``inplace=True``, otherwise returns a new instance.
"""
from .methods.apply import apply_methods
struct = manage_inplace(self, inplace)
if len(set(type(ctnr) for ctnr in struct.data)) > 1:
raise ValueError("DataSet contains more than one kind of Container.")
struct = apply_methods(struct, method, *args, **kwargs)
return struct
# ====================
# Multi-Target Methods
# ====================
def _fetch_by_key(self, key: str):
return [ctnr for ctnr in self.data if ctnr._target_key == key].copy()
def _fetch_by_id(self, id: int):
key = self._alias_map.get(f"id:{id}")
if key is None:
return []
return self._fetch_by_key(key)
def _fetch_by_coord(self, coord: SkyCoord, radius: Quantity | None = 3 * u.arcsec):
for t in self.targets:
if coord.separation(t.initial_coords) < radius:
return self._fetch_by_key(t._key)
return []
def _fetch_by_target(self, target: Target):
return self._fetch_by_key(target._key)
[docs]
def split(
self,
targets: int | SkyCoord | Target | list[int | SkyCoord | Target],
radius: Quantity | float = 3 * u.arcsec,
inplace: bool = True,
) -> DataSet:
"""
Generates a :class:`~ATK.Models.DataSet` which only contains data from specified targets.
Parameters
----------
targets: int, :class:`~astropy.coordinates.SkyCoord`, :class:`~ATK.Models.Target`, or iterable of these
Targets that should be included in the returned :class:`~ATK.Models.DataSet`.
radius: :class:`~astropy.units.Quantity`, optional
Sets the radius that is used when matching input :class:`~astropy.coordinates.SkyCoord`\ s to :class:`~ATK.Models.DataSet` :class:`~ATK.Models.Target`\ s (this form of matching is not exact).
Only relevant if one ore more input ``targets`` is a :class:`~astropy.coordinates.SkyCoord`.
inplace : bool, optional
If ``True``, modify the current :class:`~ATK.Models.DataSet` inplace. If ``False``, operate on and return a copy - leaving the original unchanged.
Returns
-------
``Self``
The split :class:`~ATK.Models.DataSet`. Returns ``self`` if ``inplace=True``, otherwise returns a new instance.
"""
from ..queries.query_core import _normalise_targeting_input, setup_targeting
input_targets = _normalise_targeting_input(targets)
targets = setup_targeting(targets)
struct = manage_inplace(self, inplace)
ctnrs = []
for target in input_targets:
# already a Target
if isinstance(target, Target):
ctnrs.extend(struct._fetch_by_target(target))
# SkyCoord -> Target
elif isinstance(target, SkyCoord):
ctnrs.extend(struct._fetch_by_coord(target, radius=radius))
# id -> Target
elif isinstance(target, int):
ctnrs.extend(struct._fetch_by_id(target))
else:
raise TypeError(f"Unsupported target type: {type(target)}")
out_keys = []
for ctnr in ctnrs:
if ctnr._target_key not in out_keys:
out_keys.append(ctnr._target_key)
out_targets = [struct._key_map[key] for key in out_keys]
struct.data = ctnrs
struct.targets = out_targets
struct._build_target_maps()
return struct
[docs]
def merge(self, dataset: DataSet, inplace: bool = True) -> DataSet:
"""
Merges this :class:`~ATK.Models.DataSet` with another.
Parameters
----------
dataset: :class:`~ATK.Models.DataSet`
The :class:`~ATK.Models.DataSet` to be merged into this one.
Both :class:`~ATK.Models.DataSet`\ s must store the same ``kind`` of data.
inplace : bool, optional
If ``True``, modify the current :class:`~ATK.Models.DataSet` inplace. If ``False``, operate on and return a copy - leaving the original unchanged.
Returns
-------
``Self``
The merged :class:`~ATK.Models.DataSet`. Returns ``self`` if ``inplace=True``, otherwise returns a new instance.
"""
struct = manage_inplace(self, inplace)
struct.data = struct.data + dataset.data
for target in dataset.targets:
struct_keys = [t._key for t in struct.targets]
if target._key not in struct_keys:
struct.targets.append(target)
struct._build_target_maps()
return struct
# Other Stuff
# ===========
@property
def _title(self):
return f"ATK {self._ctnr_kind.upper()}"
@property
def _ctnr_kind(self):
ctnr_kinds = [type(ctnr).__name__.lower() for ctnr in self.data]
if len(set(ctnr_kinds)) > 1:
raise ValueError("DataSet contains multiple container types.")
if not ctnr_kinds:
return
return ctnr_kinds[0]
@property
def _plot_method(self):
if self._ctnr_kind not in COMBINE_PLOTS:
return
return COMBINE_PLOTS[self._ctnr_kind]
@property
def _split_by_target(self):
if self._ctnr_kind not in SPLIT_BY_TARGET:
return
return SPLIT_BY_TARGET[self._ctnr_kind]
@property
def _split_by_survey(self):
if self._ctnr_kind not in SPLIT_BY_SURVEY:
return
return SPLIT_BY_SURVEY[self._ctnr_kind]
[docs]
@classmethod
def from_target(cls, target: Target | int | SkyCoord) -> Self:
"""
Initialises an empty :class:`~ATK.Models.DataSet` for a given target.
Parameters
----------
target: Target, int or :class:`~astropy.coordinates.SkyCoord`
Target for which a :class:`~ATK.Models.DataSet` should be initialised.
Returns
-------
``Self``
:class:`~ATK.Models.DataSet` with ``data = []``.
"""
from ..queries.query_core import setup_targeting
# kind, targets, survey, radius, exception, data, figure
targets = setup_targeting(target)
return cls(kind=None, targets=targets, exception=False)
[docs]
def add(self, data: Container) -> Self:
"""
Adds a data container to a :class:`~ATK.Models.DataSet`.
The container must be of the same type as the currently stored data (unless the :class:`~ATK.Models.DataSet` is empty).
Parameters
----------
data: :class:`~ATK.base.Container`
Container to append to the :class:`~ATK.Models.DataSet`
Returns
-------
``Self``
:class:`~ATK.Models.DataSet` with an additional data container.
"""
if self.data and self._ctnr_kind != type(data).__name__.lower():
raise ValueError(f"Cannot add container of type '{type(data).__name__.lower()}' to DataSet containing {self._ctnr_kind} data.")
if not self.kind:
self.kind = type(data).__name__
self.data.append(data)