update check_unit

Author: PabloCorcho

src/pst/observables.py

@@ -291,54 +291,6 @@ def interpolate(self, wavelength=None):
         self.wavelength= wavelength
         return self.response
 
-    def _check_spectra(self, spectra, default_unit=u.Lsun / u.angstrom / u.cm**2):
-        """
-        Ensure that ``spectra`` is an astropy Quantity and return it as flux density
-        per wavelength (F_lambda), using `self.wavelength` for conversions if
-        necessary.
-
-        Parameters
-        ----------
-        spectra : array-like or ~astropy.units.Quantity
-            Spectrum values. If array-like (dimensionless), it will be multiplied by
-            `default_unit`. If a Quantity, it may be either F_nu (per frequency) or
-            F_lambda (per wavelength).
-        default_unit : ~astropy.units.Unit, optional
-            Target per-wavelength flux-density unit to return, e.g. Lsun / Å / cm²
-            (must be equivalent to power / area / length).
-
-        Returns
-        -------
-        astropy.units.Quantity
-            Spectrum expressed in `default_unit` (per-wavelength flux density).
-
-        Notes
-        -----
-        - Requires `self.wavelength` to be defined as an astropy Quantity with
-        length units.
-        - If `spectra` is already per wavelength, only a unit conversion is applied.
-        - If `spectra` is per frequency, the conversion uses
-        `equivalencies=u.spectral_density(self.wavelength)`.
-        """
-        if spectra is None:
-            return None
-
-        if not isinstance(spectra, u.Quantity):
-            spectra = u.Quantity(spectra, default_unit)
-        # Check flux density units
-        if spectra.unit.is_equivalent(u.W / (u.m**2 * u.m)):
-            return spectra
-        elif spectra.unit.is_equivalent(u.W / (u.m**2 * u.Hz)):
-            if self.wavelength is None:
-                raise AttributeError("Attribute `self.wavelength` must be defined"
-                                     + " to convert spectra to Flam units.")
-            return spectra.to(default_unit, equivalencies=u.spectral_density(self.wavelength))
-        else:
-            raise ValueError(
-           f"Cannot interpret spectral flux units {spectra.unit}. "
-            "Expected per-frequency or per-wavelength flux density."
-            )   
-
     def get_photons(self, spectra, spectra_err=None, mask_nan=True):
         r"""Compute the photon flux from an input spectra.
         
@@ -367,8 +319,9 @@ def get_photons(self, spectra, spectra_err=None, mask_nan=True):
         photon_flux_err : :class:``astropy.units.Quantity``
             Filter photon flux associated error.
         """
-        spectra = self._check_spectra(spectra)
-        spectra_err = self._check_spectra(spectra_err)
+        spectra = utils.check_unit(spectra, default_unit=u.Lsun / u.angstrom / u.cm**2,
+                                   equivalence=u.spectral_density, wav=self.wavelength)
+
         if mask_nan:
             mask = np.isfinite(spectra)
             photon_flux = np.trapz(
@@ -382,6 +335,10 @@ def get_photons(self, spectra, spectra_err=None, mask_nan=True):
                 x=self.wavelength)
 
         if spectra_err is not None:
+
+            spectra_err = utils.check_unit(spectra_err,
+                                           default_unit=u.Lsun / u.angstrom / u.cm**2,
+                                           equivalence=u.spectral_density, wav=self.wavelength)
             if mask_nan:
                 mask = mask & np.isfinite(spectra_err)
             else:
@@ -497,8 +454,9 @@ def get_flambda_vegamag(self, spectra, spectra_err=None, mask_nan=True):
         --------
         :func:`get_photons`
         """
-        spectra = self._check_spectra(spectra)
-        spectra_err = self._check_spectra(spectra_err)
+        spectra = spectra = utils.check_unit(spectra, default_unit=u.Lsun / u.angstrom / u.cm**2,
+                                   equivalence=u.spectral_density, wav=self.wavelength)
+
         if mask_nan:
             mask = np.isfinite(spectra)
         else:
@@ -508,6 +466,9 @@ def get_flambda_vegamag(self, spectra, spectra_err=None, mask_nan=True):
                                ) / np.trapz(self.response[mask] * self.wavelength[mask], x=self.wavelength[mask])
 
         if spectra_err is not None:
+
+            spectra_err = utils.check_unit(spectra_err, default_unit=u.Lsun / u.angstrom / u.cm**2,
+                                           equivalence=u.spectral_density, wav=self.wavelength)
             if mask_nan:
                 mask = mask & np.isfinite(spectra_err)
             else:

src/pst/utils.py

@@ -67,7 +67,7 @@ def gaussian1d_conv(f, sigma, deltax):
             f_convolved[pixel] = np.sum(f * g)
     return f_convolved
 
-def check_unit(quantity, default_unit=None):
+def check_unit(quantity, default_unit=None, equivalence=None, **equiv_kwargs):
     """Check the units of an input quantity.
     
     Parameters
@@ -81,12 +81,17 @@ def check_unit(quantity, default_unit=None):
     isq = isinstance(quantity, u.Quantity)
     if isq and default_unit is not None:
         if not quantity.unit.is_equivalent(default_unit):
-            raise u.UnitTypeError(
-                "Input quantity does not have the appropriate units")
+            if equivalence is not None:
+                return quantity.to(default_unit,
+                equivalencies=equivalence(**equiv_kwargs))
+            else:
+                raise u.UnitTypeError(
+                    "Input quantity does not have the appropriate units")
         else:
-            return quantity
+            return quantity.to(default_unit)
+
     elif not isq and default_unit is not None:
-        return quantity * default_unit
+        return quantity << default_unit
     elif not isq and default_unit is None:
         raise ValueError("Input value must be a astropy.units.Quantity")
     else: