Esempio n. 1
0
    def __init_subclass__(cls, meta_copy: bool = True):
        """Control subclass creation.

        Parameters
        ----------
        meta_copy : bool, optional
            Whether the `~astropy.utils.metadata.MetaData` instance uses
            ``copy=True``.
            The value is stored in ``self._meta_copy``

        """
        cls.meta = MetaData(copy=meta_copy)
        cls._meta_copy = meta_copy
        super().__init_subclass__()
Esempio n. 2
0
class Cosmology(metaclass=abc.ABCMeta):
    """Base-class for all Cosmologies.

    Parameters
    ----------
    *args
        Arguments into the cosmology; used by subclasses, not this base class.
    name : str or None (optional, keyword-only)
        The name of the cosmology.
    meta : dict or None (optional, keyword-only)
        Metadata for the cosmology, e.g., a reference.
    **kwargs
        Arguments into the cosmology; used by subclasses, not this base class.

    Notes
    -----
    Class instances are static -- you cannot (and should not) change the values
    of the parameters.  That is, all of the above attributes (except meta) are
    read only.

    For details on how to create performant custom subclasses, see the
    documentation on :ref:`astropy-cosmology-fast-integrals`.
    """

    meta = MetaData()

    # Unified I/O object interchange methods
    from_format = UnifiedReadWriteMethod(CosmologyFromFormat)
    to_format = UnifiedReadWriteMethod(CosmologyToFormat)

    # Unified I/O read and write methods
    read = UnifiedReadWriteMethod(CosmologyRead)
    write = UnifiedReadWriteMethod(CosmologyWrite)

    # Parameters
    __parameters__ = ()
    __all_parameters__ = ()

    # ---------------------------------------------------------------

    def __init_subclass__(cls):
        super().__init_subclass__()

        # -------------------
        # Parameters

        # Get parameters that are still Parameters, either in this class or above.
        parameters = []
        derived_parameters = []
        for n in cls.__parameters__:
            p = getattr(cls, n)
            if isinstance(p, Parameter):
                derived_parameters.append(
                    n) if p.derived else parameters.append(n)

        # Add new parameter definitions
        for n, v in cls.__dict__.items():
            if n in parameters or n.startswith("_") or not isinstance(
                    v, Parameter):
                continue
            derived_parameters.append(n) if v.derived else parameters.append(n)

        # reorder to match signature
        ordered = [
            parameters.pop(parameters.index(n))
            for n in cls._init_signature.parameters.keys() if n in parameters
        ]
        parameters = ordered + parameters  # place "unordered" at the end
        cls.__parameters__ = tuple(parameters)
        cls.__all_parameters__ = cls.__parameters__ + tuple(derived_parameters)

        # -------------------
        # register as a Cosmology subclass
        _COSMOLOGY_CLASSES[cls.__qualname__] = cls

    @classproperty(lazy=True)
    def _init_signature(cls):
        """Initialization signature (without 'self')."""
        # get signature, dropping "self" by taking arguments [1:]
        sig = inspect.signature(cls.__init__)
        sig = sig.replace(parameters=list(sig.parameters.values())[1:])
        return sig

    # ---------------------------------------------------------------

    def __init__(self, name=None, meta=None):
        self._name = str(name) if name is not None else name
        self.meta.update(meta or {})

    @property
    def name(self):
        """The name of the Cosmology instance."""
        return self._name

    @property
    @abc.abstractmethod
    def is_flat(self):
        """
        Return bool; `True` if the cosmology is flat.
        This is abstract and must be defined in subclasses.
        """
        raise NotImplementedError("is_flat is not implemented")

    def clone(self, *, meta=None, **kwargs):
        """Returns a copy of this object with updated parameters, as specified.

        This cannot be used to change the type of the cosmology, so ``clone()``
        cannot be used to change between flat and non-flat cosmologies.

        Parameters
        ----------
        meta : mapping or None (optional, keyword-only)
            Metadata that will update the current metadata.
        **kwargs
            Cosmology parameter (and name) modifications. If any parameter is
            changed and a new name is not given, the name will be set to "[old
            name] (modified)".

        Returns
        -------
        newcosmo : `~astropy.cosmology.Cosmology` subclass instance
            A new instance of this class with updated parameters as specified.
            If no arguments are given, then a reference to this object is
            returned instead of copy.

        Examples
        --------
        To make a copy of the ``Planck13`` cosmology with a different matter
        density (``Om0``), and a new name:

            >>> from astropy.cosmology import Planck13
            >>> Planck13.clone(name="Modified Planck 2013", Om0=0.35)
            FlatLambdaCDM(name="Modified Planck 2013", H0=67.77 km / (Mpc s),
                          Om0=0.35, ...

        If no name is specified, the new name will note the modification.

            >>> Planck13.clone(Om0=0.35).name
            'Planck13 (modified)'
        """
        # Quick return check, taking advantage of the Cosmology immutability.
        if meta is None and not kwargs:
            return self

        # There are changed parameter or metadata values.
        # The name needs to be changed accordingly, if it wasn't already.
        _modname = self.name + " (modified)"
        kwargs.setdefault("name",
                          (_modname if self.name is not None else None))

        # mix new meta into existing, preferring the former.
        meta = meta if meta is not None else {}
        new_meta = {**self.meta, **meta}
        # Mix kwargs into initial arguments, preferring the former.
        new_init = {**self._init_arguments, "meta": new_meta, **kwargs}
        # Create BoundArgument to handle args versus kwargs.
        # This also handles all errors from mismatched arguments
        ba = self._init_signature.bind_partial(**new_init)
        # Instantiate, respecting args vs kwargs
        cloned = type(self)(*ba.args, **ba.kwargs)

        # Check if nothing has changed.
        # TODO! or should return self?
        if (cloned.name
                == _modname) and not meta and cloned.is_equivalent(self):
            cloned._name = self.name

        return cloned

    @property
    def _init_arguments(self):
        # parameters
        kw = {n: getattr(self, n) for n in self.__parameters__}

        # other info
        kw["name"] = self.name
        kw["meta"] = self.meta

        return kw

    # ---------------------------------------------------------------
    # comparison methods

    def is_equivalent(self, other, *, format=False):
        r"""Check equivalence between Cosmologies.

        Two cosmologies may be equivalent even if not the same class.
        For example, an instance of ``LambdaCDM`` might have :math:`\Omega_0=1`
        and :math:`\Omega_k=0` and therefore be flat, like ``FlatLambdaCDM``.

        Parameters
        ----------
        other : `~astropy.cosmology.Cosmology` subclass instance
            The object in which to compare.
        format : bool or None or str, optional keyword-only
            Whether to allow, before equivalence is checked, the object to be
            converted to a |Cosmology|. This allows, e.g. a |Table| to be
            equivalent to a Cosmology.
            `False` (default) will not allow conversion. `True` or `None` will,
            and will use the auto-identification to try to infer the correct
            format. A `str` is assumed to be the correct format to use when
            converting.

        Returns
        -------
        bool
            True if cosmologies are equivalent, False otherwise.

        Examples
        --------
        Two cosmologies may be equivalent even if not of the same class.
        In this examples the ``LambdaCDM`` has ``Ode0`` set to the same value
        calculated in ``FlatLambdaCDM``.

            >>> import astropy.units as u
            >>> from astropy.cosmology import LambdaCDM, FlatLambdaCDM
            >>> cosmo1 = LambdaCDM(70 * (u.km/u.s/u.Mpc), 0.3, 0.7)
            >>> cosmo2 = FlatLambdaCDM(70 * (u.km/u.s/u.Mpc), 0.3)
            >>> cosmo1.is_equivalent(cosmo2)
            True

        While in this example, the cosmologies are not equivalent.

            >>> cosmo3 = FlatLambdaCDM(70 * (u.km/u.s/u.Mpc), 0.3, Tcmb0=3 * u.K)
            >>> cosmo3.is_equivalent(cosmo2)
            False

        Also, using the keyword argument, the notion of equivalence is extended
        to any Python object that can be converted to a |Cosmology|.

            >>> from astropy.cosmology import Planck18
            >>> tbl = Planck18.to_format("astropy.table")
            >>> Planck18.is_equivalent(tbl, format=True)
            True

        The list of valid formats, e.g. the |Table| in this example, may be
        checked with ``Cosmology.from_format.list_formats()``.

        As can be seen in the list of formats, not all formats can be
        auto-identified by ``Cosmology.from_format.registry``. Objects of
        these kinds can still be checked for equivalence, but the correct
        format string must be used.

            >>> tbl = Planck18.to_format("yaml")
            >>> Planck18.is_equivalent(tbl, format="yaml")
            True
        """
        # Allow for different formats to be considered equivalent.
        if format is not False:
            format = None if format is True else format  # str->str, None/True->None
            try:
                other = Cosmology.from_format(other, format=format)
            except Exception:  # TODO! should enforce only TypeError
                return False

        # The options are: 1) same class & parameters; 2) same class, different
        # parameters; 3) different classes, equivalent parameters; 4) different
        # classes, different parameters. (1) & (3) => True, (2) & (4) => False.
        equiv = self.__equiv__(other)
        if equiv is NotImplemented and hasattr(other, "__equiv__"):
            equiv = other.__equiv__(self)  # that failed, try from 'other'

        return equiv if equiv is not NotImplemented else False

    def __equiv__(self, other):
        """Cosmology equivalence. Use ``.is_equivalent()`` for actual check!

        Parameters
        ----------
        other : `~astropy.cosmology.Cosmology` subclass instance
            The object in which to compare.

        Returns
        -------
        bool or `NotImplemented`
            `NotImplemented` if 'other' is from a different class.
            `True` if 'other' is of the same class and has matching parameters
            and parameter values. `False` otherwise.
        """
        if other.__class__ is not self.__class__:
            return NotImplemented  # allows other.__equiv__

        # check all parameters in 'other' match those in 'self' and 'other' has
        # no extra parameters (latter part should never happen b/c same class)
        params_eq = (set(self.__all_parameters__) == set(
            other.__all_parameters__) and all(
                np.all(getattr(self, k) == getattr(other, k))
                for k in self.__all_parameters__))
        return params_eq

    def __eq__(self, other):
        """Check equality between Cosmologies.

        Checks the Parameters and immutable fields (i.e. not "meta").

        Parameters
        ----------
        other : `~astropy.cosmology.Cosmology` subclass instance
            The object in which to compare.

        Returns
        -------
        bool
            `True` if Parameters and names are the same, `False` otherwise.
        """
        if other.__class__ is not self.__class__:
            return NotImplemented  # allows other.__eq__

        # check all parameters in 'other' match those in 'self'
        equivalent = self.__equiv__(other)
        # non-Parameter checks: name
        name_eq = (self.name == other.name)

        return equivalent and name_eq

    # ---------------------------------------------------------------

    def __repr__(self):
        namelead = f"{self.__class__.__qualname__}("
        if self.name is not None:
            namelead += f"name=\"{self.name}\", "
        # nicely formatted parameters
        fmtps = (f'{k}={getattr(self, k)}' for k in self.__parameters__)

        return namelead + ", ".join(fmtps) + ")"

    def __astropy_table__(self, cls, copy, **kwargs):
        """Return a `~astropy.table.Table` of type ``cls``.

        Parameters
        ----------
        cls : type
            Astropy ``Table`` class or subclass.
        copy : bool
            Ignored.
        **kwargs : dict, optional
            Additional keyword arguments. Passed to ``self.to_format()``.
            See ``Cosmology.to_format.help("astropy.table")`` for allowed kwargs.

        Returns
        -------
        `astropy.table.Table` or subclass instance
            Instance of type ``cls``.
        """
        return self.to_format("astropy.table", cls=cls, **kwargs)
class ExampleData:
    meta = MetaData()

    def __init__(self, meta=None):
        self.meta = meta
Esempio n. 4
0
class NDData(NDDataBase):
    """
    A container for `numpy.ndarray`-based datasets, using the
    `~astropy.nddata.NDDataBase` interface.

    The key distinction from raw `numpy.ndarray` is the presence of
    additional metadata such as uncertainty, mask, unit, a coordinate system
    and/or a dictionary containing further meta information. This class *only*
    provides a container for *storing* such datasets. For further functionality
    take a look at the ``See also`` section.

    See also: http://docs.astropy.org/en/stable/nddata/

    Parameters
    -----------
    data : `numpy.ndarray`-like or `NDData`-like
        The dataset.

    uncertainty : any type, optional
        Uncertainty in the dataset.
        Should have an attribute ``uncertainty_type`` that defines what kind of
        uncertainty is stored, for example ``"std"`` for standard deviation or
        ``"var"`` for variance. A metaclass defining such an interface is
        `NDUncertainty` - but isn't mandatory. If the uncertainty has no such
        attribute the uncertainty is stored as `UnknownUncertainty`.
        Defaults to ``None``.

    mask : any type, optional
        Mask for the dataset. Masks should follow the ``numpy`` convention that
        **valid** data points are marked by ``False`` and **invalid** ones with
        ``True``.
        Defaults to ``None``.

    wcs : any type, optional
        World coordinate system (WCS) for the dataset.
        Default is ``None``.

    meta : `dict`-like object, optional
        Additional meta information about the dataset. If no meta is provided
        an empty `collections.OrderedDict` is created.
        Default is ``None``.

    unit : `~astropy.units.Unit`-like or str, optional
        Unit for the dataset. Strings that can be converted to a
        `~astropy.units.Unit` are allowed.
        Default is ``None``.

    copy : `bool`, optional
        Indicates whether to save the arguments as copy. ``True`` copies
        every attribute before saving it while ``False`` tries to save every
        parameter as reference.
        Note however that it is not always possible to save the input as
        reference.
        Default is ``False``.

        .. versionadded:: 1.2

    Raises
    ------
    TypeError
        In case ``data`` or ``meta`` don't meet the restrictions.

    Notes
    -----
    Each attribute can be accessed through the homonymous instance attribute:
    ``data`` in a `NDData` object can be accessed through the `data`
    attribute::

        >>> from astropy.nddata import NDData
        >>> nd = NDData([1,2,3])
        >>> nd.data
        array([1, 2, 3])

    Given a conflicting implicit and an explicit parameter during
    initialization, for example the ``data`` is a `~astropy.units.Quantity` and
    the unit parameter is not ``None``, then the implicit parameter is replaced
    (without conversion) by the explicit one and a warning is issued::

        >>> import numpy as np
        >>> import astropy.units as u
        >>> q = np.array([1,2,3,4]) * u.m
        >>> nd2 = NDData(q, unit=u.cm)
        INFO: overwriting Quantity's current unit with specified unit. [astropy.nddata.nddata]
        >>> nd2.data  # doctest: +FLOAT_CMP
        array([1., 2., 3., 4.])
        >>> nd2.unit
        Unit("cm")

    See also
    --------
    NDDataRef
    NDDataArray
    """

    # Instead of a custom property use the MetaData descriptor also used for
    # Tables. It will check if the meta is dict-like or raise an exception.
    meta = MetaData(doc=_meta_doc, copy=False)

    def __init__(self,
                 data,
                 uncertainty=None,
                 mask=None,
                 wcs=None,
                 meta=None,
                 unit=None,
                 copy=False):

        # Rather pointless since the NDDataBase does not implement any setting
        # but before the NDDataBase did call the uncertainty
        # setter. But if anyone wants to alter this behavior again the call
        # to the superclass NDDataBase should be in here.
        super().__init__()

        # Check if data is any type from which to collect some implicitly
        # passed parameters.
        if isinstance(data, NDData):  # don't use self.__class__ (issue #4137)
            # Of course we need to check the data because subclasses with other
            # init-logic might be passed in here. We could skip these
            # tests if we compared for self.__class__ but that has other
            # drawbacks.

            # Comparing if there is an explicit and an implicit unit parameter.
            # If that is the case use the explicit one and issue a warning
            # that there might be a conflict. In case there is no explicit
            # unit just overwrite the unit parameter with the NDData.unit
            # and proceed as if that one was given as parameter. Same for the
            # other parameters.
            if (unit is not None and data.unit is not None
                    and unit != data.unit):
                log.info("overwriting NDData's current "
                         "unit with specified unit.")
            elif data.unit is not None:
                unit = data.unit

            if uncertainty is not None and data.uncertainty is not None:
                log.info("overwriting NDData's current "
                         "uncertainty with specified uncertainty.")
            elif data.uncertainty is not None:
                uncertainty = data.uncertainty

            if mask is not None and data.mask is not None:
                log.info("overwriting NDData's current "
                         "mask with specified mask.")
            elif data.mask is not None:
                mask = data.mask

            if wcs is not None and data.wcs is not None:
                log.info("overwriting NDData's current "
                         "wcs with specified wcs.")
            elif data.wcs is not None:
                wcs = data.wcs

            if meta is not None and data.meta is not None:
                log.info("overwriting NDData's current "
                         "meta with specified meta.")
            elif data.meta is not None:
                meta = data.meta

            data = data.data

        else:
            if hasattr(data, 'mask') and hasattr(data, 'data'):
                # Separating data and mask
                if mask is not None:
                    log.info("overwriting Masked Objects's current "
                             "mask with specified mask.")
                else:
                    mask = data.mask

                # Just save the data for further processing, we could be given
                # a masked Quantity or something else entirely. Better to check
                # it first.
                data = data.data

            if isinstance(data, Quantity):
                if unit is not None and unit != data.unit:
                    log.info("overwriting Quantity's current "
                             "unit with specified unit.")
                else:
                    unit = data.unit
                data = data.value

        # Quick check on the parameters if they match the requirements.
        if (not hasattr(data, 'shape') or not hasattr(data, '__getitem__')
                or not hasattr(data, '__array__')):
            # Data doesn't look like a numpy array, try converting it to
            # one.
            data = np.array(data, subok=True, copy=False)

        # Another quick check to see if what we got looks like an array
        # rather than an object (since numpy will convert a
        # non-numerical/non-string inputs to an array of objects).
        if data.dtype == 'O':
            raise TypeError("could not convert data to numpy array.")

        if unit is not None:
            unit = Unit(unit)

        if copy:
            # Data might have been copied before but no way of validating
            # without another variable.
            data = deepcopy(data)
            mask = deepcopy(mask)
            wcs = deepcopy(wcs)
            meta = deepcopy(meta)
            uncertainty = deepcopy(uncertainty)
            # Actually - copying the unit is unnecessary but better safe
            # than sorry :-)
            unit = deepcopy(unit)

        # Validate the wcs

        # Store the attributes
        self._data = data
        self.mask = mask
        self._wcs = wcs
        self.meta = meta  # TODO: Make this call the setter sometime
        self._unit = unit
        # Call the setter for uncertainty to further check the uncertainty
        self.uncertainty = uncertainty

    def __str__(self):
        return str(self.data)

    def __repr__(self):
        prefix = self.__class__.__name__ + '('
        body = np.array2string(self.data, separator=', ', prefix=prefix)
        return ''.join([prefix, body, ')'])

    @property
    def data(self):
        """
        `~numpy.ndarray`-like : The stored dataset.
        """
        return self._data

    @property
    def mask(self):
        """
        any type : Mask for the dataset, if any.

        Masks should follow the ``numpy`` convention that valid data points are
        marked by ``False`` and invalid ones with ``True``.
        """
        return self._mask

    @mask.setter
    def mask(self, value):
        self._mask = value

    @property
    def unit(self):
        """
        `~astropy.units.Unit` : Unit for the dataset, if any.
        """
        return self._unit

    @property
    def wcs(self):
        """
        any type : A world coordinate system (WCS) for the dataset, if any.
        """
        return self._wcs

    @wcs.setter
    def wcs(self, wcs):
        if self._wcs is not None and wcs is not None:
            raise ValueError(
                "You can only set the wcs attribute with a WCS if no WCS is present."
            )

        if wcs is None or isinstance(wcs, BaseHighLevelWCS):
            self._wcs = wcs
        elif isinstance(wcs, BaseLowLevelWCS):
            self._wcs = HighLevelWCSWrapper(wcs)
        else:
            raise TypeError(
                "The wcs argument must implement either the high or"
                " low level WCS API.")

    @property
    def uncertainty(self):
        """
        any type : Uncertainty in the dataset, if any.

        Should have an attribute ``uncertainty_type`` that defines what kind of
        uncertainty is stored, such as ``'std'`` for standard deviation or
        ``'var'`` for variance. A metaclass defining such an interface is
        `~astropy.nddata.NDUncertainty` but isn't mandatory.
        """
        return self._uncertainty

    @uncertainty.setter
    def uncertainty(self, value):
        if value is not None:
            # There is one requirements on the uncertainty: That
            # it has an attribute 'uncertainty_type'.
            # If it does not match this requirement convert it to an unknown
            # uncertainty.
            if not hasattr(value, 'uncertainty_type'):
                log.info('uncertainty should have attribute uncertainty_type.')
                value = UnknownUncertainty(value, copy=False)

            # If it is a subclass of NDUncertainty we must set the
            # parent_nddata attribute. (#4152)
            if isinstance(value, NDUncertainty):
                # In case the uncertainty already has a parent create a new
                # instance because we need to assume that we don't want to
                # steal the uncertainty from another NDData object
                if value._parent_nddata is not None:
                    value = value.__class__(value, copy=False)
                # Then link it to this NDData instance (internally this needs
                # to be saved as weakref but that's done by NDUncertainty
                # setter).
                value.parent_nddata = self
        self._uncertainty = value
Esempio n. 5
0
class Cosmology(metaclass=abc.ABCMeta):
    """Base-class for all Cosmologies.

    Parameters
    ----------
    *args
        Arguments into the cosmology; used by subclasses, not this base class.
    name : str or None (optional, keyword-only)
        The name of the cosmology.
    meta : dict or None (optional, keyword-only)
        Metadata for the cosmology, e.g., a reference.
    **kwargs
        Arguments into the cosmology; used by subclasses, not this base class.

    Notes
    -----
    Class instances are static -- you cannot (and should not) change the values
    of the parameters.  That is, all of the above attributes (except meta) are
    read only.

    For details on how to create performant custom subclasses, see the
    documentation on :ref:`astropy-cosmology-fast-integrals`.
    """

    meta = MetaData()

    # Unified I/O object interchange methods
    from_format = UnifiedReadWriteMethod(CosmologyFromFormat)
    to_format = UnifiedReadWriteMethod(CosmologyToFormat)

    # Unified I/O read and write methods
    read = UnifiedReadWriteMethod(CosmologyRead)
    write = UnifiedReadWriteMethod(CosmologyWrite)

    # Parameters
    __parameters__: tuple[str, ...] = ()
    __all_parameters__: tuple[str, ...] = ()

    # ---------------------------------------------------------------

    def __init_subclass__(cls):
        super().__init_subclass__()

        # -------------------
        # Parameters

        # Get parameters that are still Parameters, either in this class or above.
        parameters = []
        derived_parameters = []
        for n in cls.__parameters__:
            p = getattr(cls, n)
            if isinstance(p, Parameter):
                derived_parameters.append(
                    n) if p.derived else parameters.append(n)

        # Add new parameter definitions
        for n, v in cls.__dict__.items():
            if n in parameters or n.startswith("_") or not isinstance(
                    v, Parameter):
                continue
            derived_parameters.append(n) if v.derived else parameters.append(n)

        # reorder to match signature
        ordered = [
            parameters.pop(parameters.index(n))
            for n in cls._init_signature.parameters.keys() if n in parameters
        ]
        parameters = ordered + parameters  # place "unordered" at the end
        cls.__parameters__ = tuple(parameters)
        cls.__all_parameters__ = cls.__parameters__ + tuple(derived_parameters)

        # -------------------
        # register as a Cosmology subclass
        _COSMOLOGY_CLASSES[cls.__qualname__] = cls

    @classproperty(lazy=True)
    def _init_signature(cls):
        """Initialization signature (without 'self')."""
        # get signature, dropping "self" by taking arguments [1:]
        sig = inspect.signature(cls.__init__)
        sig = sig.replace(parameters=list(sig.parameters.values())[1:])
        return sig

    # ---------------------------------------------------------------

    def __init__(self, name=None, meta=None):
        self._name = str(name) if name is not None else name
        self.meta.update(meta or {})

    @property
    def name(self):
        """The name of the Cosmology instance."""
        return self._name

    @property
    @abc.abstractmethod
    def is_flat(self):
        """
        Return bool; `True` if the cosmology is flat.
        This is abstract and must be defined in subclasses.
        """
        raise NotImplementedError("is_flat is not implemented")

    def clone(self, *, meta=None, **kwargs):
        """Returns a copy of this object with updated parameters, as specified.

        This cannot be used to change the type of the cosmology, so ``clone()``
        cannot be used to change between flat and non-flat cosmologies.

        Parameters
        ----------
        meta : mapping or None (optional, keyword-only)
            Metadata that will update the current metadata.
        **kwargs
            Cosmology parameter (and name) modifications. If any parameter is
            changed and a new name is not given, the name will be set to "[old
            name] (modified)".

        Returns
        -------
        newcosmo : `~astropy.cosmology.Cosmology` subclass instance
            A new instance of this class with updated parameters as specified.
            If no arguments are given, then a reference to this object is
            returned instead of copy.

        Examples
        --------
        To make a copy of the ``Planck13`` cosmology with a different matter
        density (``Om0``), and a new name:

            >>> from astropy.cosmology import Planck13
            >>> Planck13.clone(name="Modified Planck 2013", Om0=0.35)
            FlatLambdaCDM(name="Modified Planck 2013", H0=67.77 km / (Mpc s),
                          Om0=0.35, ...

        If no name is specified, the new name will note the modification.

            >>> Planck13.clone(Om0=0.35).name
            'Planck13 (modified)'
        """
        # Quick return check, taking advantage of the Cosmology immutability.
        if meta is None and not kwargs:
            return self

        # There are changed parameter or metadata values.
        # The name needs to be changed accordingly, if it wasn't already.
        _modname = self.name + " (modified)"
        kwargs.setdefault("name",
                          (_modname if self.name is not None else None))

        # mix new meta into existing, preferring the former.
        meta = meta if meta is not None else {}
        new_meta = {**self.meta, **meta}
        # Mix kwargs into initial arguments, preferring the former.
        new_init = {**self._init_arguments, "meta": new_meta, **kwargs}
        # Create BoundArgument to handle args versus kwargs.
        # This also handles all errors from mismatched arguments
        ba = self._init_signature.bind_partial(**new_init)
        # Instantiate, respecting args vs kwargs
        cloned = type(self)(*ba.args, **ba.kwargs)

        # Check if nothing has changed.
        # TODO! or should return self?
        if (cloned.name
                == _modname) and not meta and cloned.is_equivalent(self):
            cloned._name = self.name

        return cloned

    @property
    def _init_arguments(self):
        # parameters
        kw = {n: getattr(self, n) for n in self.__parameters__}

        # other info
        kw["name"] = self.name
        kw["meta"] = self.meta

        return kw

    # ---------------------------------------------------------------
    # comparison methods

    def is_equivalent(self,
                      other: Any,
                      /,
                      *,
                      format: _FormatType = False) -> bool:
        r"""Check equivalence between Cosmologies.

        Two cosmologies may be equivalent even if not the same class.
        For example, an instance of ``LambdaCDM`` might have :math:`\Omega_0=1`
        and :math:`\Omega_k=0` and therefore be flat, like ``FlatLambdaCDM``.

        Parameters
        ----------
        other : `~astropy.cosmology.Cosmology` subclass instance, positional-only
            The object to which to compare.
        format : bool or None or str, optional keyword-only
            Whether to allow, before equivalence is checked, the object to be
            converted to a |Cosmology|. This allows, e.g. a |Table| to be
            equivalent to a Cosmology.
            `False` (default) will not allow conversion. `True` or `None` will,
            and will use the auto-identification to try to infer the correct
            format. A `str` is assumed to be the correct format to use when
            converting.
            ``format`` is broadcast to match the shape of ``other``.
            Note that the cosmology arguments are not broadcast against
            ``format``, so it cannot determine the output shape.

        Returns
        -------
        bool
            True if cosmologies are equivalent, False otherwise.

        Examples
        --------
        Two cosmologies may be equivalent even if not of the same class.
        In this examples the ``LambdaCDM`` has ``Ode0`` set to the same value
        calculated in ``FlatLambdaCDM``.

            >>> import astropy.units as u
            >>> from astropy.cosmology import LambdaCDM, FlatLambdaCDM
            >>> cosmo1 = LambdaCDM(70 * (u.km/u.s/u.Mpc), 0.3, 0.7)
            >>> cosmo2 = FlatLambdaCDM(70 * (u.km/u.s/u.Mpc), 0.3)
            >>> cosmo1.is_equivalent(cosmo2)
            True

        While in this example, the cosmologies are not equivalent.

            >>> cosmo3 = FlatLambdaCDM(70 * (u.km/u.s/u.Mpc), 0.3, Tcmb0=3 * u.K)
            >>> cosmo3.is_equivalent(cosmo2)
            False

        Also, using the keyword argument, the notion of equivalence is extended
        to any Python object that can be converted to a |Cosmology|.

            >>> from astropy.cosmology import Planck18
            >>> tbl = Planck18.to_format("astropy.table")
            >>> Planck18.is_equivalent(tbl, format=True)
            True

        The list of valid formats, e.g. the |Table| in this example, may be
        checked with ``Cosmology.from_format.list_formats()``.

        As can be seen in the list of formats, not all formats can be
        auto-identified by ``Cosmology.from_format.registry``. Objects of
        these kinds can still be checked for equivalence, but the correct
        format string must be used.

            >>> tbl = Planck18.to_format("yaml")
            >>> Planck18.is_equivalent(tbl, format="yaml")
            True
        """
        from .funcs import cosmology_equal

        try:
            return cosmology_equal(self,
                                   other,
                                   format=(None, format),
                                   allow_equivalent=True)
        except Exception:
            # `is_equivalent` allows `other` to be any object and returns False
            # if `other` cannot be converted to a Cosmology, rather than
            # raising an Exception.
            return False
Esempio n. 6
0
class Cosmology(metaclass=abc.ABCMeta):
    """Base-class for all Cosmologies.

    Parameters
    ----------
    *args
        Arguments into the cosmology; used by subclasses, not this base class.
    name : str or None (optional, keyword-only)
        The name of the cosmology.
    meta : dict or None (optional, keyword-only)
        Metadata for the cosmology, e.g., a reference.
    **kwargs
        Arguments into the cosmology; used by subclasses, not this base class.

    Notes
    -----
    Class instances are static -- you cannot (and should not) change the values
    of the parameters.  That is, all of the above attributes (except meta) are
    read only.

    For details on how to create performant custom subclasses, see the
    documentation on :ref:`astropy-cosmology-fast-integrals`.
    """

    meta = MetaData()

    # Unified I/O object interchange methods
    from_format = UnifiedReadWriteMethod(CosmologyFromFormat)
    to_format = UnifiedReadWriteMethod(CosmologyToFormat)

    # Unified I/O read and write methods
    read = UnifiedReadWriteMethod(CosmologyRead)
    write = UnifiedReadWriteMethod(CosmologyWrite)

    # Parameters
    __parameters__ = ()
    __all_parameters__ = ()

    # ---------------------------------------------------------------

    def __init_subclass__(cls):
        super().__init_subclass__()

        # -------------------
        # Parameters

        # Get parameters that are still Parameters, either in this class or above.
        parameters = []
        derived_parameters = []
        for n in cls.__parameters__:
            p = getattr(cls, n)
            if isinstance(p, Parameter):
                derived_parameters.append(
                    n) if p.derived else parameters.append(n)

        # Add new parameter definitions
        for n, v in cls.__dict__.items():
            if n in parameters or n.startswith("_") or not isinstance(
                    v, Parameter):
                continue
            derived_parameters.append(n) if v.derived else parameters.append(n)

        # reorder to match signature
        ordered = [
            parameters.pop(parameters.index(n))
            for n in cls._init_signature.parameters.keys() if n in parameters
        ]
        parameters = ordered + parameters  # place "unordered" at the end
        cls.__parameters__ = tuple(parameters)
        cls.__all_parameters__ = cls.__parameters__ + tuple(derived_parameters)

        # -------------------
        # register as a Cosmology subclass
        _COSMOLOGY_CLASSES[cls.__qualname__] = cls

    @classproperty(lazy=True)
    def _init_signature(cls):
        """Initialization signature (without 'self')."""
        # get signature, dropping "self" by taking arguments [1:]
        sig = inspect.signature(cls.__init__)
        sig = sig.replace(parameters=list(sig.parameters.values())[1:])
        return sig

    # ---------------------------------------------------------------

    def __init__(self, name=None, meta=None):
        self._name = name
        self.meta.update(meta or {})

    @property
    def name(self):
        """The name of the Cosmology instance."""
        return self._name

    @property
    @abc.abstractmethod
    def is_flat(self):
        """
        Return bool; `True` if the cosmology is flat.
        This is abstract and must be defined in subclasses.
        """
        raise NotImplementedError("is_flat is not implemented")

    def clone(self, *, meta=None, **kwargs):
        """Returns a copy of this object with updated parameters, as specified.

        This cannot be used to change the type of the cosmology, so ``clone()``
        cannot be used to change between flat and non-flat cosmologies.

        Parameters
        ----------
        meta : mapping or None (optional, keyword-only)
            Metadata that will update the current metadata.
        **kwargs
            Cosmology parameter (and name) modifications.
            If any parameter is changed and a new name is not given, the name
            will be set to "[old name] (modified)".

        Returns
        -------
        newcosmo : `~astropy.cosmology.Cosmology` subclass instance
            A new instance of this class with updated parameters as specified.
            If no modifications are requested, then a reference to this object
            is returned instead of copy.

        Examples
        --------
        To make a copy of the ``Planck13`` cosmology with a different matter
        density (``Om0``), and a new name:

            >>> from astropy.cosmology import Planck13
            >>> newcosmo = Planck13.clone(name="Modified Planck 2013", Om0=0.35)

        If no name is specified, the new name will note the modification.

            >>> Planck13.clone(Om0=0.35).name
            'Planck13 (modified)'
        """
        # Quick return check, taking advantage of the Cosmology immutability.
        if meta is None and not kwargs:
            return self

        # There are changed parameter or metadata values.
        # The name needs to be changed accordingly, if it wasn't already.
        kwargs.setdefault(
            "name",
            (self.name + " (modified)" if self.name is not None else None))

        # mix new meta into existing, preferring the former.
        new_meta = {**self.meta, **(meta or {})}
        # Mix kwargs into initial arguments, preferring the former.
        new_init = {**self._init_arguments, "meta": new_meta, **kwargs}
        # Create BoundArgument to handle args versus kwargs.
        # This also handles all errors from mismatched arguments
        ba = self._init_signature.bind_partial(**new_init)
        # Return new instance, respecting args vs kwargs
        return self.__class__(*ba.args, **ba.kwargs)

    @property
    def _init_arguments(self):
        # parameters
        kw = {n: getattr(self, n) for n in self.__parameters__}

        # other info
        kw["name"] = self.name
        kw["meta"] = self.meta

        return kw

    # ---------------------------------------------------------------
    # comparison methods

    def is_equivalent(self, other):
        r"""Check equivalence between Cosmologies.

        Two cosmologies may be equivalent even if not the same class.
        For example, an instance of ``LambdaCDM`` might have :math:`\Omega_0=1`
        and :math:`\Omega_k=0` and therefore be flat, like ``FlatLambdaCDM``.

        Parameters
        ----------
        other : `~astropy.cosmology.Cosmology` subclass instance
            The object in which to compare.

        Returns
        -------
        bool
            True if cosmologies are equivalent, False otherwise.
        """
        # The options are: 1) same class & parameters; 2) same class, different
        # parameters; 3) different classes, equivalent parameters; 4) different
        # classes, different parameters. (1) & (3) => True, (2) & (4) => False.
        equiv = self.__equiv__(other)
        if equiv is NotImplemented and hasattr(other, "__equiv__"):
            equiv = other.__equiv__(self)  # that failed, try from 'other'

        return equiv if equiv is not NotImplemented else False

    def __equiv__(self, other):
        """Cosmology equivalence. Use ``.is_equivalent()`` for actual check!

        Parameters
        ----------
        other : `~astropy.cosmology.Cosmology` subclass instance
            The object in which to compare.

        Returns
        -------
        bool or `NotImplemented`
            `NotImplemented` if 'other' is from a different class.
            `True` if 'other' is of the same class and has matching parameters
            and parameter values. `False` otherwise.
        """
        if other.__class__ is not self.__class__:
            return NotImplemented  # allows other.__equiv__

        # check all parameters in 'other' match those in 'self' and 'other' has
        # no extra parameters (latter part should never happen b/c same class)
        params_eq = (set(self.__all_parameters__) == set(
            other.__all_parameters__) and all(
                np.all(getattr(self, k) == getattr(other, k))
                for k in self.__all_parameters__))
        return params_eq

    def __eq__(self, other):
        """Check equality between Cosmologies.

        Checks the Parameters and immutable fields (i.e. not "meta").

        Parameters
        ----------
        other : `~astropy.cosmology.Cosmology` subclass instance
            The object in which to compare.

        Returns
        -------
        bool
            True if Parameters and names are the same, False otherwise.
        """
        if other.__class__ is not self.__class__:
            return NotImplemented  # allows other.__eq__

        # check all parameters in 'other' match those in 'self'
        equivalent = self.__equiv__(other)
        # non-Parameter checks: name
        name_eq = (self.name == other.name)

        return equivalent and name_eq

    # ---------------------------------------------------------------

    def __repr__(self):
        ps = {k: getattr(self, k) for k in self.__parameters__}  # values
        cps = {k: getattr(self.__class__, k)
               for k in self.__parameters__}  # Parameter objects

        namelead = f"{self.__class__.__qualname__}("
        if self.name is not None:
            namelead += f"name=\"{self.name}\", "
        # nicely formatted parameters
        fmtps = (k + '=' +
                 format(v, cps[k].format_spec if v is not None else '')
                 for k, v in ps.items())

        return namelead + ", ".join(fmtps) + ")"

    def __astropy_table__(self, cls, copy, **kwargs):
        """Return a `~astropy.table.Table` of type ``cls``.

        Parameters
        ----------
        cls : type
            Astropy ``Table`` class or subclass.
        copy : bool
            Ignored.
        **kwargs : dict, optional
            Additional keyword arguments. Passed to ``self.to_format()``.
            See ``Cosmology.to_format.help("astropy.table")`` for allowed kwargs.

        Returns
        -------
        `astropy.table.Table` or subclass instance
            Instance of type ``cls``.
        """
        return self.to_format("astropy.table", cls=cls, **kwargs)
Esempio n. 7
0
class TablesList(HomogeneousList):
    """Grouped Tables.

    A subclass of list that contains only elements of a given type or
    types.  If an item that is not of the specified type is added to
    the list, a `TypeError` is raised. Also includes some pretty printing
    methods for an OrderedDict of :class:`~astropy.table.Table` objects.

    """

    _types = None

    meta = MetaData(copy=False)

    def __init__(
        self,
        inp: OrderedDictType = [],
        *,
        name: T.Optional[str] = None,
        reference: T.Optional[T.Any] = None,
        **metadata,
    ):
        """Astroquery-style table list.

        Parameters
        ----------
        inp : sequence, optional
            An initial set of tables.
        name : str, optional
            name of the list of tables.
        reference : citation, optional
            citation.
        **metadata : Any
            arguments into meta

        """
        # meta
        self.meta["name"] = name
        self.meta["reference"] = reference
        for k, v in metadata.items():
            self.meta[k] = v

        inp = self._validate(inp)  # ODict, & ensure can assign values

        # Convert input to correct to type
        # If None, can be anything
        if self._types is not None:
            for k, val in inp.items():
                inp[k] = self._types(val)  # TODO handle multiple "_types"

        # finally add the input
        # _dict store the indices for the keys
        self._dict = {k: i for i, k in enumerate(inp.keys())}

        # need to bypass HomogeneousList init, which uses ``extend``
        list.__init__(self, inp.values())

    # /def

    # -----------------

    def _assert(self, x):
        """Check `x` is correct type (set by _type)."""
        if self._types is None:  # allow any type
            return
        super()._assert(x)

    # /def

    def _validate(self, value):
        """Validate `value` compatible with table."""
        if isinstance(value, (TablesList, OrderedDict)):
            pass
        else:
            try:
                value = OrderedDict(value)
            except (TypeError, ValueError):
                raise ValueError(
                    "Input to TableList must be an OrderedDict "
                    "or list of (k,v) pairs"
                )

        return value

    # /def

    # -----------------
    # Properties

    @property
    def name(self) -> str:
        """Name."""
        return self.meta["name"]

    # /def

    @property
    def __reference__(self):
        """Get reference from metadata, if exists."""
        return self.meta.get("reference", None)

    # /def

    # -----------------
    # Dictionary methods

    def keys(self):
        """Set-like object giving table names."""
        return self._dict.keys()

    # /def

    def sortedkeys(self):
        """Set-like object giving table names.

        Ordered by value. Does not update with table.

        """
        sorted_dict = dict(
            sorted(self._dict.items(), key=lambda item: item[0])
        )
        return cabc.KeysView(sorted_dict.keys())

    # /def

    def values(self):
        """Tuple object providing a view on tables.

        Note that the tables can be edited

        """
        return tuple(self)

    # /def

    def items(self):
        """Generator providing iterator over name, table."""
        for key, value in zip(self.keys(), self.values()):
            yield key, value

    # /def

    # -----------------
    # Get / Set

    def index(self, key: str) -> int:
        """Index of `key`.

        Parameters
        ----------
        key : str

        Returns
        -------
        int

        """
        return self._dict[key]

    # /def

    def __getitem__(self, key: T.Union[int, slice, str]):
        """Get item or slice.

        Parameters
        ----------
        key : str or int or slice
            if str, the dictionary key.
            if int, the dictionary index
            if slice, slices dictionary.values
            supports string as slice start or stop

        Returns
        -------
        Table

        Raises
        ------
        TypeError
            if key is not int or key

        """
        if isinstance(key, int):
            return super().__getitem__(key)
        elif isinstance(key, slice):
            start, stop = key.start, key.stop
            # string replacement for start, stop values
            # replace by int
            if isinstance(start, str):
                start = self.index(start)
            if isinstance(stop, str):
                stop = self.index(stop)
            key = slice(start, stop, key.step)
            return super().__getitem__(key)
        else:
            return super().__getitem__(self.index(key))

    # /def

    def __setitem__(self, key: str, value):
        """Set item, but only if right type (managed by super)."""
        if not isinstance(key, str):
            raise TypeError

        # first try if exists
        if key in self._dict:
            ind = self._dict[key]
            super().__setitem__(ind, value)  # (super _assert)

        # else append to end
        else:
            ind = len(self)
            self._dict[key] = ind
            super().append(value)  # (super _assert)

    # /def

    def __delitem__(self, key: T.Union[str, int]):  # TODO test!
        """Delete Item. Forbidden."""
        if isinstance(key, str):
            i = self.index(key)
        elif isinstance(key, int):
            i = key
            key = [k for k, v in self.items() if v == i][0]  # get key
        else:
            raise TypeError

        super().__delitem__(i)
        self._dict.pop(key)

    # /def

    def update(self, other):
        """Update TableList using OrderedDict update method."""
        values = self._validate(other)  # first make sure adding a key-val pair
        for k, v in values.items():  # TODO better
            self[k] = v  # setitem manages _dict

    # /def

    def extend(self, other):
        """Extend TableList. Unlike update, cannot have duplicate keys."""
        values = self._validate(other)  # first make sure adding a key-val pair

        if any((k in self.keys() for k in values.keys())):
            raise ValueError("cannot have duplicate keys")

        self.update(values)

    # /def

    def __iadd__(self, other):
        """Add in-place."""
        return super().__iadd__(other)

    # /def

    def append(self, key: str, value):
        """Append, if unique key and right type (managed by super)."""
        if key in self._dict:
            raise ValueError("cannot append duplicate key.")

        self._dict[key] = len(self)
        return super().append(value)

    # /def

    def pop(self):
        """Pop. Forbidden."""
        raise NotImplementedError("Forbidden.")

    # /def

    def insert(self, value):
        """Insert. Forbidden."""
        raise NotImplementedError("Forbidden.")

    # /def

    # -----------------
    # string representation

    def __repr__(self):
        """String representation.

        Overrides the `OrderedDict.__repr__` method to return a simple summary
        of the `TableList` object.

        Returns
        -------
        str

        """
        return self.format_table_list()

    # /def

    def format_table_list(self) -> str:
        """String Representation of list of Tables.

        Prints the names of all :class:`~astropy.table.Table` objects, with
        their respective number of row and columns, contained in the
        `TableList` instance.

        Returns
        -------
        str

        """
        ntables = len(list(self.keys()))
        if ntables == 0:
            return "Empty {cls}".format(cls=self.__class__.__name__)

        header_str = "{cls} with {keylen} tables:".format(
            cls=self.__class__.__name__, keylen=ntables
        )
        body_str = "\n".join(
            [
                "\t'{t_number}:{t_name}' with {ncol} column(s) "
                "and {nrow} row(s) ".format(
                    t_number=t_number,
                    t_name=t_name,
                    nrow=len(self[t_number]),
                    ncol=len(self[t_number].colnames),
                )
                for t_number, t_name in enumerate(self.keys())
            ]
        )

        return "\n".join([header_str, body_str])

    # /def

    def print_table_list(self):
        """Print Table List.

        calls ``format_table_list``

        """
        print(self.format_table_list())

    # /def

    def pprint(self, **kwargs):
        """Helper function to make API more similar to astropy.Tables.

        .. todo::

            uses "kwargs"

        """
        if kwargs != {}:
            warnings.warn(
                "TableList is a container of astropy.Tables.", InputWarning
            )

        self.print_table_list()

    # /def

    # -----------------
    # I/O

    def _save_table_iter(self, format, **table_kw):
        for i, name in enumerate(self.keys()):  # order-preserving

            # get kwargs for table writer
            # first get all general keys (by filtering)
            # then update with table-specific dictionary (if present)
            kw = {
                k: v for k, v in table_kw.items() if not k.startswith("table_")
            }
            kw.update(table_kw.get("table_" + name, {}))

            if isinstance(format, str):
                fmt = format
            else:
                fmt = format[i]

            yield name, fmt, kw

    # /def

    def write(
        self,
        drct: str,
        format="asdf",
        split=True,
        serialize_method=None,
        **table_kw,
    ):
        """Write to ASDF.

        Parameters
        ----------
        drct : str
            The main directory path.
        format : str or list, optional
            save format. default "asdf"
            can be list of same length as TableList
        split : bool, optional
            *Applies to asdf `format` only*

            Whether to save the tables as individual file
            with `file` coordinating by reference.

        serialize_method : str, dict, optional
            Serialization method specifier for columns.

        **table_kw
            kwargs into each table.
            1. dictionary with table name as key
            2. General keys

        """

        # -----------
        # Path checks

        path = pathlib.Path(drct)

        if path.suffix == "":  # no suffix
            path = path.with_suffix(".asdf")

        if path.suffix != ".asdf":  # ensure only asdf
            raise ValueError("file type must be `.asdf`.")

        drct = path.parent  # directory in which to save

        # -----------
        # TableType

        if self._types is None:
            table_type = [
                tp.__class__.__module__ + "." + tp.__class__.__name__
                for tp in self.values()
            ]
        else:
            table_type = self._types.__module__ + "." + self._types.__name__

        # -----------
        # Saving

        TL = asdf.AsdfFile()
        TL.tree["meta"] = tuple(self.meta.items())
        TL.tree["table_names"] = tuple(self.keys())  # in order
        TL.tree["save_format"] = format
        TL.tree["table_type"] = table_type

        if format == "asdf" and not split:  # save as single file
            for name in self.keys():  # add to tree
                TL.tree[name] = self[name]

        else:  # save as individual files
            for name, fmt, kw in self._save_table_iter(format, **table_kw):

                # name of table
                table_path = drct.joinpath(name)
                if table_path.suffix == "":  # TODO currently always. CLEANUP
                    table_path = table_path.with_suffix(
                        "." + fmt.split(".")[-1]
                    )

                # internal save
                if format == "asdf":
                    kw["data_key"] = kw.get("data_key", name)
                self[name].write(
                    table_path,
                    format=fmt,
                    serialize_method=serialize_method,
                    **kw,
                )

                # save by relative reference
                if format == "asdf":
                    with asdf.open(table_path) as f:
                        TL.tree[name] = f.make_reference(path=[name])
                else:
                    TL.tree[name] = str(table_path.relative_to(drct))

        # Need to add a "data" key to not break asdf
        if not format == "asdf":
            TL.tree["data"] = TL.tree["table_names"]

        # /if
        TL.write_to(str(path))  # save directory

    # /def

    @classmethod
    def _read_table_iter(cls, f, format, **table_kw):
        names = f.tree["table_names"]
        # table type, for casting
        # so that QTableList can open a saved TableList correctly
        # TablesList specifies no type, so must rely on saved info
        if cls._types is None:
            table_type = f.tree["table_type"]
            if not isinstance(table_type, cabc.Sequence):
                table_type = [table_type] * len(names)
            ttypes = [resolve_name(t) for t in table_type]
        else:
            ttypes = [cls._types] * len(names)

        for i, name in enumerate(names):  # order-preserving

            if isinstance(format, str):
                fmt = format
            else:
                fmt = format[i]

            # get kwargs for table writer
            # first get all general keys (by filtering)
            # then update with table-specific dictionary (if present)
            kw = {
                k: v for k, v in table_kw.items() if not k.startswith("table_")
            }
            kw.update(table_kw.get("table_" + name, {}))

            yield name, ttypes[i], fmt, kw

    # /def

    @classmethod
    def read(
        cls,
        drct: str,
        format: T.Union[str, T.Sequence] = None,
        suffix: T.Optional[str] = None,
        **table_kw,
    ):
        """Write to ASDF.

        Parameters
        ----------
        drct : str
            The main directory path.
        format : str or list, optional
            read format. default "asdf"
            can be list of same length as TableList
        suffix : str, optional
            suffix to apply to table names.
            will be superceded by an "fnames" argument, when added

        **table_kw
            kwargs into each table.
            1. dictionary with table name as key
            2. General keys

        """
        # -----------
        # Path checks

        path = pathlib.Path(drct)

        if path.suffix == "":  # no suffix
            path = path.with_suffix(".asdf")

        if path.suffix != ".asdf":  # ensure only asdf
            raise ValueError("file type must be `.asdf`.")

        drct = path.parent  # directory

        # -----------
        # Read

        TL = cls()
        with asdf.open(path) as f:
            f.resolve_references()

            # load in the metadata
            TL.meta = OrderedDict(f.tree["meta"])

            if format is None:
                format = f.tree["save_format"]

            # iterate through tables
            for name, ttype, fmt, kw in cls._read_table_iter(
                f, format, **table_kw
            ):
                tl = f.tree[name]

                # TODO what if tuple of str as path to name?
                if not isinstance(tl, str):  # only for asdf
                    TL[name] = ttype(f.tree[name], **kw)  # TODO need kw?
                else:
                    table_path = drct.joinpath(tl)  # .with_suffix(suffix)
                    TL[name] = ttype.read(str(table_path), format=fmt, **kw)

        return TL

    # /def

    def copy(self):
        """Shallow copy."""
        out = self.__class__(self)
        out.meta = self.meta

        return out
Esempio n. 8
0
class Cosmology(metaclass=ABCMeta):
    """Base-class for all Cosmologies.

    Parameters
    ----------
    *args
        Arguments into the cosmology; used by subclasses, not this base class.
    name : str or None (optional, keyword-only)
        The name of the cosmology.
    meta : dict or None (optional, keyword-only)
        Metadata for the cosmology, e.g., a reference.
    **kwargs
        Arguments into the cosmology; used by subclasses, not this base class.

    Notes
    -----
    Class instances are static -- you cannot (and should not) change the values
    of the parameters.  That is, all of the above attributes (except meta) are
    read only.

    For details on how to create performant custom subclasses, see the
    documentation on :ref:`astropy-cosmology-fast-integrals`.
    """

    meta = MetaData()

    # Unified I/O object interchange methods
    from_format = UnifiedReadWriteMethod(CosmologyFromFormat)
    to_format = UnifiedReadWriteMethod(CosmologyToFormat)

    # Unified I/O read and write methods
    read = UnifiedReadWriteMethod(CosmologyRead)
    write = UnifiedReadWriteMethod(CosmologyWrite)

    def __init_subclass__(cls):
        super().__init_subclass__()

        _COSMOLOGY_CLASSES[cls.__qualname__] = cls

    def __new__(cls, *args, **kwargs):
        self = super().__new__(cls)

        # bundle and store initialization arguments on the instance
        ba = cls._init_signature.bind_partial(*args, **kwargs)
        ba.apply_defaults()  # and fill in the defaults
        self._init_arguments = ba.arguments

        return self

    def __init__(self, *args, name=None, meta=None, **kwargs):
        self._name = name
        self.meta.update(meta or {})

    @classproperty(lazy=True)
    def _init_signature(cls):
        """Initialization signature (without 'self')."""
        # get signature, dropping "self" by taking arguments [1:]
        sig = signature(cls.__init__)
        sig = sig.replace(parameters=list(sig.parameters.values())[1:])
        return sig

    @property
    def name(self):
        """The name of the Cosmology instance."""
        return self._name

    def clone(self, *, meta=None, **kwargs):
        """Returns a copy of this object with updated parameters, as specified.

        This cannot be used to change the type of the cosmology, so ``clone()``
        cannot be used to change between flat and non-flat cosmologies.

        Parameters
        ----------
        meta : mapping or None (optional, keyword-only)
            Metadata that will update the current metadata.
        **kwargs
            Cosmology parameter (and name) modifications.
            If any parameter is changed and a new name is not given, the name
            will be set to "[old name] (modified)".

        Returns
        -------
        newcosmo : `~astropy.cosmology.Cosmology` subclass instance
            A new instance of this class with updated parameters as specified.
            If no modifications are requested, then a reference to this object
            is returned instead of copy.

        Examples
        --------
        To make a copy of the ``Planck13`` cosmology with a different matter
        density (``Om0``), and a new name:

            >>> from astropy.cosmology import Planck13
            >>> newcosmo = Planck13.clone(name="Modified Planck 2013", Om0=0.35)

        If no name is specified, the new name will note the modification.

            >>> Planck13.clone(Om0=0.35).name
            'Planck13 (modified)'

        """
        # Quick return check, taking advantage of the Cosmology immutability.
        if meta is None and not kwargs:
            return self

        # There are changed parameter or metadata values.
        # The name needs to be changed accordingly, if it wasn't already.
        kwargs.setdefault(
            "name",
            (self.name + " (modified)" if self.name is not None else None))

        # mix new meta into existing, preferring the former.
        new_meta = {**self.meta, **(meta or {})}
        # Mix kwargs into initial arguments, preferring the former.
        new_init = {**self._init_arguments, "meta": new_meta, **kwargs}
        # Create BoundArgument to handle args versus kwargs.
        # This also handles all errors from mismatched arguments
        ba = self._init_signature.bind_partial(**new_init)
        # Return new instance, respecting args vs kwargs
        return self.__class__(*ba.args, **ba.kwargs)

    # -----------------------------------------------------

    def __eq__(self, other):
        """Check equality on all immutable fields (i.e. not "meta").

        Parameters
        ----------
        other : `~astropy.cosmology.Cosmology` subclass instance
            The object in which to compare.

        Returns
        -------
        bool
            True if all immutable fields are equal, False otherwise.

        """
        if not isinstance(other, Cosmology):
            return False

        sias, oias = self._init_arguments, other._init_arguments

        # check if the cosmologies have identical signatures.
        # this protects against one cosmology having a superset of input
        # parameters to another cosmology.
        if (sias.keys() ^ oias.keys()) - {'meta'}:
            return False

        # are all the non-excluded immutable arguments equal?
        return all(
            (np.all(oias[k] == v) for k, v in sias.items() if k != "meta"))
Esempio n. 9
0
class Base:
    """Base class of all tasks and generators.

    Following the design of `baseband` stream readers, features properties
    describing the size, shape, data type, sample rate and start/stop times of
    the task's output.  Also defines methods to move a sample pointer across
    the output data in units of either complete samples or time.

    Subclasses should define

      ``_read_frame``: method to read (or generate) a single block of data.

    Parameters
    ----------
    shape : tuple, optional
        Overall shape of the stream, with first entry the total number
        of complete samples, and the remainder the sample shape.
    start_time : `~astropy.time.Time`
        Start time of the stream.
    sample_rate : `~astropy.units.Quantity`
        Rate at which complete samples are produced.
    samples_per_frame : int, optional
        Number of samples dealt with in one go.  The number of complete
        samples (``shape[0]``) should be an integer multiple of this.
    dtype : `~numpy.dtype`, optional
        Dtype of the samples.

    --- **kwargs : meta data for the stream, which usually include

    frequency : `~astropy.units.Quantity`, optional
        Frequencies for each channel.  Should be broadcastable to the
        sample shape.
    sideband : array, optional
        Whether frequencies are upper (+1) or lower (-1) sideband.
        Should be broadcastable to the sample shape.
    polarization : array or (nested) list of char, optional
        Polarization labels.  Should broadcast to the sample shape,
        i.e., the labels are in the correct axis.  For instance,
        ``['X', 'Y']``, or ``[['L'], ['R']]``.
    """

    # Initial values for sample and frame pointers, etc.
    offset = 0
    _frame_index = None
    _frame = None
    closed = False

    meta = MetaData()

    def __init__(self,
                 shape,
                 start_time,
                 sample_rate,
                 *,
                 samples_per_frame=1,
                 dtype=np.complex64,
                 **kwargs):
        self._shape = shape
        self._start_time = start_time
        self._samples_per_frame = operator.index(samples_per_frame)
        self._sample_rate = sample_rate
        self._dtype = np.dtype(dtype, copy=False)

        if len({'frequency', 'sideband'}.difference(kwargs)) == 1:
            raise ValueError('frequency and sideband should both '
                             'be passed in.')

        attributes = {}
        for attr, value in kwargs.items():
            if attr in META_ATTRIBUTES:
                if value is not None:
                    if attr == 'sideband':
                        value = np.where(value > 0, np.int8(1), np.int8(-1))
                    attributes[attr] = self._check_shape(value)
            else:
                raise TypeError('__init__() got unexpected keyword argument '
                                f'{attr!r}')
        if attributes:
            self.meta.setdefault('__attributes__', {}).update(attributes)

    def __getattr__(self, attr):
        if attr in META_ATTRIBUTES:
            value = self.meta.get('__attributes__', {}).get(attr, None)
            if value is None:
                raise AttributeError(f"{attr} not set.")
            else:
                return value
        else:
            return super().__getattr__(attr)

    def __dir__(self):
        return sorted(META_ATTRIBUTES.union(super().__dir__()))

    def _repr_item(self, key, default, value=None):
        """Representation of one argument.

        Subclasses can override this, either to return something else than
        the base key=value or to set a different default for specific keys.

        """

        if value is None:
            value = getattr(self, key, None)
            if value is None:
                value = getattr(self, '_' + key, None)
                if value is None:
                    return None

        if default is not inspect._empty:
            try:
                if np.all(value == default):
                    return None
            except Exception:
                pass

        return f"{key}={value}".replace('\n', ',')

    def __str__(self):
        name = self.__class__.__name__
        pars = inspect.signature(self.__class__).parameters
        overrides = [
            self._repr_item(key, par.default) for key, par in pars.items()
        ]

        overrides = ', '.join([override for override in overrides if override])
        return f"{name}({overrides})"

    def __repr__(self):
        """Representation which lists non-default arguments.

        Finds possible arguments by inspection of the whole class hierarchy
        (as long as kwargs are passed along) and creates a list of all whose
        values on the instance are different from the default. Subclasses
        can override the assumed default and what to return in _repr_item.

        """
        name = self.__class__.__name__
        pars = {}
        for cls in self.__class__.__mro__:
            for key, par in inspect.signature(cls).parameters.items():
                pars.setdefault(key, par)
            if 'kwargs' not in pars or cls is Base:
                break

        overrides = [
            self._repr_item(key, par.default) for key, par in pars.items()
        ]
        if cls is Base and '__attributes__' in self.meta:
            overrides.extend([
                self._repr_item(key, None)
                for key in self.meta['__attributes__'].keys()
                if key not in pars
            ])

        overrides = (',\n ' + ' ' * len(name)).join(
            [override for override in overrides if override])
        return f"{name}({overrides})"

    def _check_shape(self, value):
        """Check that value can be broadcast to the sample shape."""
        broadcast = check_broadcast_to(value, self.sample_shape)
        return simplify_shape(broadcast)

    @property
    def shape(self):
        """Shape of the output."""
        return self._shape

    @property
    def sample_shape(self):
        """Shape of a complete sample."""
        return self.shape[1:]

    @property
    def samples_per_frame(self):
        """Number of samples per frame of data.

        For compatibility with file readers, to help indicate what
        a nominal chunk of data is.
        """
        return self._samples_per_frame

    @property
    def size(self):
        """Number of component samples in the output."""
        prod = 1
        for dim in self.shape:
            prod *= dim
        return prod

    @property
    def ndim(self):
        """Number of dimensions of the output."""
        return len(self.shape)

    @property
    def dtype(self):
        """Data type of the output."""
        return self._dtype

    @property
    def complex_data(self):
        return self._dtype.kind == 'c'

    @property
    def sample_rate(self):
        """Number of complete samples per second."""
        return self._sample_rate

    @property
    def start_time(self):
        """Start time of the output.

        See also `time` and `stop_time`.
        """
        # We don't just return self._start_time so classes like Integrate
        # can get correct results by just overriding _tell_time.
        return self._tell_time(0)

    @property
    def time(self):
        """Time of the sample pointer's current offset in the output.

        See also `start_time` and `stop_time`.
        """
        return self._tell_time(self.offset)

    @property
    def stop_time(self):
        """Time at the end of the output, just after the last sample.

        See also `start_time` and `time`.
        """
        return self._tell_time(self.shape[0])

    def seek(self, offset, whence=0):
        """Change the sample pointer position.

        This works like a normal filehandle seek, but the offset is in samples
        (or a relative or absolute time).

        Parameters
        ----------
        offset : int, `~astropy.units.Quantity`, or `~astropy.time.Time`
            Offset to move to.  Can be an (integer) number of samples,
            an offset in time units, or an absolute time.  For the latter
            two, the pointer will be moved to the nearest integer sample.
        whence : {0, 1, 2, 'start', 'current', or 'end'}, optional
            Like regular seek, the offset is taken to be from the start if
            ``whence=0`` (default), from the current position if 1,
            and from the end if 2.  One can alternativey use 'start',
            'current', or 'end' for 0, 1, or 2, respectively.  Ignored if
            ``offset`` is a time.
        """
        try:
            offset = operator.index(offset)
        except Exception:
            try:
                offset = offset - self.start_time
            except Exception:
                pass
            else:
                whence = 0

            offset = int((offset * self.sample_rate).to(u.one).round())

        if whence == 0 or whence == 'start':
            self.offset = offset
        elif whence == 1 or whence == 'current':
            self.offset += offset
        elif whence == 2 or whence == 'end':
            self.offset = self.shape[0] + offset
        else:
            raise ValueError("invalid 'whence'; should be 0 or 'start', 1 or "
                             "'current', or 2 or 'end'.")

        return self.offset

    def tell(self, unit=None):
        """Current offset in the file.

        Parameters
        ----------
        unit : `~astropy.units.Unit` or str, optional
            Time unit the offset should be returned in.  By default, no unit
            is used, i.e., an integer enumerating samples is returned. For the
            special string 'time', the absolute time is calculated.

        Returns
        -------
        offset : int, `~astropy.units.Quantity`, or `~astropy.time.Time`
             Offset in current file (or time at current position).
        """
        if unit is None:
            return self.offset

        # "isinstance" avoids costly comparisons of an actual unit with 'time'.
        if not isinstance(unit, u.UnitBase) and unit == 'time':
            return self._tell_time(self.offset)

        return (self.offset / self.sample_rate).to(unit)

    def _tell_time(self, offset):
        """Calculate time for given offset.

        Used for ``start_time``, ``time``, ``stop_time`` and
        ``tell(unit='time')``.  Simple implementation is present mostly so
        subclasses like Integration and Stack can override as appropriate.

        """
        return self._start_time + offset / self.sample_rate

    def read(self, count=None, out=None):
        """Read a number of complete samples.

        Parameters
        ----------
        count : int or None, optional
            Number of complete samples to read. If `None` (default) or
            negative, the number of samples left. Ignored if ``out`` is given.
        out : None or array, optional
            Array to store the samples in. If given, ``count`` will be inferred
            from the first dimension; the remaining dimensions should equal
            `sample_shape`.

        Returns
        -------
        out : `~numpy.ndarray` of float or complex
            The first dimension is sample-time, and the remaining ones are
            as given by `sample_shape`.
        """
        if self.closed:
            raise ValueError("I/O operation on closed stream.")

        samples_left = self.shape[0] - self.offset
        if out is None:
            if count is None or count < 0:
                count = max(0, samples_left)

            out = np.empty((count, ) + self.sample_shape, dtype=self.dtype)
        else:
            assert out.shape[1:] == self.sample_shape, (
                "'out' must have trailing shape {}".format(self.sample_shape))
            count = out.shape[0]

        if count > samples_left:
            raise EOFError("cannot read from beyond end of input.")

        offset0 = self.offset
        sample = 0
        while sample < count:
            # For current position, get frame plus offset in that frame.
            frame, sample_offset = self._get_frame(self.offset)
            nsample = min(count - sample, len(frame) - sample_offset)
            data = frame[sample_offset:sample_offset + nsample]
            # Copy to relevant part of output.
            out[sample:sample + nsample] = data
            sample += nsample
            # Explicitly set offset (leaving get_frame free to adjust it).
            self.offset = offset0 + sample

        return out

    def _get_frame(self, offset):
        """Get a frame that includes given offset.

        Finds the index corresponding to the needed frame, assuming frames
        are all the same length.  If not already cached, it retrieves a
        frame by calling ``self._read_frame(index)``.

        Parameters
        ----------
        offset : int
            Offset in the stream for which a frame should be found.

        Returns
        -------
        frame : `~baseband.base.frame.FrameBase`
            Frame holding the sample at ``offset``.
        sample_offset : int
            Offset within the frame corresponding to ``offset``.
        """
        frame_index, sample_offset = divmod(offset, self.samples_per_frame)
        if frame_index != self._frame_index:
            # Read the frame required. Set offset to start so that _read_frame
            # can count on tell() being correct.
            self.offset = frame_index * self.samples_per_frame
            self._frame = self._read_frame(frame_index)
            self._frame_index = frame_index

        return self._frame, sample_offset

    def __getitem__(self, item):
        from .shaping import GetSlice

        return GetSlice(self, item)

    def __array__(self, dtype=None):
        old_offset = self.tell()
        try:
            self.seek(0)
            return np.array(self.read(), dtype=dtype, copy=False)
        finally:
            self.seek(old_offset)

    def __array_ufunc__(self, *args, **kwargs):
        return NotImplemented

    def __array_function__(self, *args, **kwargs):
        return NotImplemented

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def close(self):
        self.closed = True
        self._frame = None  # clear possibly cached frame