import numpy as np

from astropy.units import Unit, si

from astropy.units import equivalencies as eq

from astropy.units.decorators import quantity_input

from astropy.units.quantity import Quantity, SpecificTypeQuantity

__all__ = ["SpectralQuantity"]

# We don't want to run doctests in the docstrings we inherit from Quantity

__doctest_skip__ = ["SpectralQuantity.*"]

KMS = si.km / si.s

SPECTRAL_UNITS = (si.Hz, si.m, si.J, si.m**-1, KMS)

DOPPLER_CONVENTIONS = {

"radio": eq.doppler_radio,

"optical": eq.doppler_optical,

"relativistic": eq.doppler_relativistic,

}

class SpectralQuantity(SpecificTypeQuantity):

"""

One or more value(s) with spectral units.

The spectral units should be those for frequencies, wavelengths, energies,

wavenumbers, or velocities (interpreted as Doppler velocities relative to a

rest spectral value). The advantage of using this class over the regular

`~astropy.units.Quantity` class is that in `SpectralQuantity`, the

``u.spectral`` equivalency is enabled by default (allowing automatic

conversion between spectral units), and a preferred Doppler rest value and

convention can be stored for easy conversion to/from velocities.

Parameters

----------

value : ndarray or `~astropy.units.Quantity` or `SpectralQuantity`

Spectral axis data values.

unit : unit-like

Unit for the given data.

doppler_rest : `~astropy.units.Quantity` ['speed'], optional

The rest value to use for conversions from/to velocities

doppler_convention : str, optional

The convention to use when converting the spectral data to/from

velocities.

"""

_equivalent_unit = SPECTRAL_UNITS

_include_easy_conversion_members = True

def __new__(

cls, value, unit=None, doppler_rest=None, doppler_convention=None, **kwargs

):

obj = super().__new__(cls, value, unit=unit, **kwargs)

# If we're initializing from an existing SpectralQuantity, keep any

# parameters that aren't being overridden

if doppler_rest is None:

doppler_rest = getattr(value, "doppler_rest", None)

if doppler_convention is None:

doppler_convention = getattr(value, "doppler_convention", None)

obj._doppler_rest = doppler_rest

obj._doppler_convention = doppler_convention

return obj

def __array_finalize__(self, obj):

super().__array_finalize__(obj)

self._doppler_rest = getattr(obj, "_doppler_rest", None)

self._doppler_convention = getattr(obj, "_doppler_convention", None)

def __quantity_subclass__(self, unit):

# Always default to just returning a Quantity, unless we explicitly

# choose to return a SpectralQuantity - even if the units match, we

# want to avoid doing things like adding two SpectralQuantity instances

# together and getting a SpectralQuantity back

if unit is self.unit:

return SpectralQuantity, True

else:

return Quantity, False

def __array_ufunc__(self, function, method, *inputs, **kwargs):

# We always return Quantity except in a few specific cases

result = super().__array_ufunc__(function, method, *inputs, **kwargs)

if (

(

function is np.multiply

or function is np.true_divide

and inputs[0] is self

)

and result.unit == self.unit

or (

function in (np.minimum, np.maximum, np.fmax, np.fmin)

and method in ("reduce", "reduceat")

)

):

result = result.view(self.__class__)

result.__array_finalize__(self)

else:

if result is self:

raise TypeError(

"Cannot store the result of this operation in"

f" {self.__class__.__name__}"

)

if result.dtype.kind == "b":

result = result.view(np.ndarray)

else:

result = result.view(Quantity)

return result

@property

def doppler_rest(self):

"""

The rest value of the spectrum used for transformations to/from

velocity space.

Returns

-------

`~astropy.units.Quantity` ['speed']

Rest value as an astropy `~astropy.units.Quantity` object.

"""

return self._doppler_rest

@doppler_rest.setter

@quantity_input(value=SPECTRAL_UNITS)

def doppler_rest(self, value):

"""

New rest value needed for velocity-space conversions.

Parameters

----------

value : `~astropy.units.Quantity` ['speed']

Rest value.

"""

if self._doppler_rest is not None:

raise AttributeError(

"doppler_rest has already been set, and cannot be changed. Use the"

" ``to`` method to convert the spectral values(s) to use a different"

" rest value"

)

self._doppler_rest = value

@property

def doppler_convention(self):

"""

The defined convention for conversions to/from velocity space.

Returns

-------

str

One of 'optical', 'radio', or 'relativistic' representing the

equivalency used in the unit conversions.

"""

return self._doppler_convention

@doppler_convention.setter

def doppler_convention(self, value):

"""

New velocity convention used for velocity space conversions.

Parameters

----------

value

Notes

-----

More information on the equations dictating the transformations can be

found in the astropy documentation [1]_.

References

----------

.. [1] Astropy documentation: https://docs.astropy.org/en/stable/units/equivalencies.html#spectral-doppler-equivalencies

"""

if self._doppler_convention is not None:

raise AttributeError(

"doppler_convention has already been set, and cannot be changed. Use"

" the ``to`` method to convert the spectral values(s) to use a"

" different convention"

)

if value is not None and value not in DOPPLER_CONVENTIONS:

raise ValueError(

"doppler_convention should be one of"

f" {'/'.join(sorted(DOPPLER_CONVENTIONS))}"

)

self._doppler_convention = value

@quantity_input(doppler_rest=SPECTRAL_UNITS)

def to(self, unit, equivalencies=[], doppler_rest=None, doppler_convention=None):

"""

Return a new `~astropy.coordinates.SpectralQuantity` object with the specified unit.

By default, the ``spectral`` equivalency will be enabled, as well as

one of the Doppler equivalencies if converting to/from velocities.

Parameters

----------

unit : unit-like

An object that represents the unit to convert to. Must be

an `~astropy.units.UnitBase` object or a string parseable

by the `~astropy.units` package, and should be a spectral unit.

equivalencies : list of `~astropy.units.equivalencies.Equivalency`, optional

A list of equivalence pairs to try if the units are not

directly convertible (along with spectral).

See :ref:`astropy:unit_equivalencies`.

If not provided or ``[]``, spectral equivalencies will be used.

If `None`, no equivalencies will be applied at all, not even any

set globally or within a context.

doppler_rest : `~astropy.units.Quantity` ['speed'], optional

The rest value used when converting to/from velocities. This will

also be set at an attribute on the output

`~astropy.coordinates.SpectralQuantity`.

doppler_convention : {'relativistic', 'optical', 'radio'}, optional

The Doppler convention used when converting to/from velocities.

This will also be set at an attribute on the output

`~astropy.coordinates.SpectralQuantity`.

Returns

-------

`SpectralQuantity`

New spectral coordinate object with data converted to the new unit.

"""

# Make sure units can be passed as strings

unit = Unit(unit)

# If equivalencies is explicitly set to None, we should just use the

# default Quantity.to with equivalencies also set to None

if equivalencies is None:

result = super().to(unit, equivalencies=None)

result = result.view(self.__class__)

result.__array_finalize__(self)

return result

# FIXME: need to consider case where doppler equivalency is passed in

# equivalencies list, or is u.spectral equivalency is already passed

if doppler_rest is None:

doppler_rest = self._doppler_rest

if doppler_convention is None:

doppler_convention = self._doppler_convention

elif doppler_convention not in DOPPLER_CONVENTIONS:

raise ValueError(

"doppler_convention should be one of"

f" {'/'.join(sorted(DOPPLER_CONVENTIONS))}"

)

if self.unit.is_equivalent(KMS) and unit.is_equivalent(KMS):

# Special case: if the current and final units are both velocity,

# and either the rest value or the convention are different, we

# need to convert back to frequency temporarily.

if doppler_convention is not None and self._doppler_convention is None:

raise ValueError("Original doppler_convention not set")

if doppler_rest is not None and self._doppler_rest is None:

raise ValueError("Original doppler_rest not set")

if doppler_rest is None and doppler_convention is None:

result = super().to(unit, equivalencies=equivalencies)

result = result.view(self.__class__)

result.__array_finalize__(self)

return result

elif (doppler_rest is None) is not (doppler_convention is None):

raise ValueError(

"Either both or neither doppler_rest and doppler_convention should"

" be defined for velocity conversions"

)

vel_equiv1 = DOPPLER_CONVENTIONS[self._doppler_convention](

self._doppler_rest

)

freq = super().to(si.Hz, equivalencies=equivalencies + vel_equiv1)

vel_equiv2 = DOPPLER_CONVENTIONS[doppler_convention](doppler_rest)

result = freq.to(unit, equivalencies=equivalencies + vel_equiv2)

else:

additional_equivalencies = eq.spectral()

if self.unit.is_equivalent(KMS) or unit.is_equivalent(KMS):

if doppler_convention is None:

raise ValueError(

"doppler_convention not set, cannot convert to/from velocities"

)

if doppler_rest is None:

raise ValueError(

"doppler_rest not set, cannot convert to/from velocities"

)

additional_equivalencies = (

additional_equivalencies

+ DOPPLER_CONVENTIONS[doppler_convention](doppler_rest)

)

result = super().to(

unit, equivalencies=equivalencies + additional_equivalencies

)

# Since we have to explicitly specify when we want to keep this as a

# SpectralQuantity, we need to convert it back from a Quantity to

# a SpectralQuantity here. Note that we don't use __array_finalize__

# here since we might need to set the output doppler convention and

# rest based on the parameters passed to 'to'

result = result.view(self.__class__)

result.__array_finalize__(self)

result._doppler_convention = doppler_convention

result._doppler_rest = doppler_rest

return result

def to_value(self, unit=None, *args, **kwargs):

if unit is None:

return self.view(np.ndarray)

return self.to(unit, *args, **kwargs).value