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