Handling larger range for BUNIT conversions

spectral_cube/base_class.py

@@ -479,6 +479,15 @@ def beams(self, obj):
 
         self._beams = obj
 
+    @property
+    @cached
+    def pixels_per_beam(self):
+        pixels_per_beam = [(beam.sr /
+                (astropy.wcs.utils.proj_plane_pixel_area(self.wcs) *
+                 u.deg**2)).to(u.one).value for beam in self.beams]
+
+        return pixels_per_beam
+
     @property
     def unmasked_beams(self):
         return self._beams

spectral_cube/cube_utils.py

@@ -2,6 +2,8 @@
 
 import contextlib
 import warnings
+from copy import deepcopy
+
 try:
     import builtins
 except ImportError:
@@ -10,6 +12,7 @@
 
 import dask.array as da
 import numpy as np
+from astropy.wcs.utils import proj_plane_pixel_area
 from astropy.wcs import (WCSSUB_SPECTRAL, WCSSUB_LONGITUDE, WCSSUB_LATITUDE)
 from . import wcs_utils
 from .utils import FITSWarning, AstropyUserWarning, WCSCelestialError
@@ -535,3 +538,185 @@ def world_take_along_axis(cube, position_plane, axis):
     out = out.squeeze()
 
     return out
+
+
+def bunit_converters(obj, unit, equivalencies=(), freq=None):
+        '''
+        Handler for all brightness unit conversions, including: K, Jy/beam, Jy/pix, Jy/sr.
+        This also includes varying resolution spectral cubes, where the beam size varies along
+        the frequency axis.
+
+        Parameters
+        ----------
+        obj : {SpectralCube, LowerDimensionalObject}
+            A spectral cube or any other lower dimensional object.
+        unit : `~astropy.units.Unit`
+            Unit to convert `obj` to.
+        equivalencies : tuple, optional
+            Initial list of equivalencies.
+        freq : `~astropy.unit.Quantity`, optional
+            Frequency to use for spectral conversions. If the spectral axis is available, the
+            frequencies will already be defined.
+
+        Outputs
+        -------
+        factor : `~numpy.ndarray`
+            Array of factors for the unit conversion.
+
+        '''
+
+        # Add a simple check it the new unit is already equivalent, and so we don't need
+        # any additional unit equivalencies
+        if obj.unit.is_equivalent(unit):
+            # return equivalencies
+            factor = obj.unit.to(unit, equivalencies=equivalencies)
+            return np.array([factor])
+
+        # Determine the bunit "type". This will determine what information we need for the unit conversion.
+        has_btemp = obj.unit.is_equivalent(u.K) or unit.is_equivalent(u.K)
+        has_perbeam = obj.unit.is_equivalent(u.Jy/u.beam) or unit.is_equivalent(u.Jy/u.beam)
+        has_perangarea = obj.unit.is_equivalent(u.Jy/u.sr) or unit.is_equivalent(u.Jy/u.sr)
+        has_perpix = obj.unit.is_equivalent(u.Jy/u.pix) or unit.is_equivalent(u.Jy/u.pix)
+
+        # Is there any beam object defined?
+        has_beam = hasattr(obj, 'beam') or hasattr(obj, 'beams')
+
+        # Set if this is a varying resolution object
+        has_beams = hasattr(obj, 'beams')
+
+        # Define freq, if needed:
+        if any([has_perangarea, has_perbeam, has_btemp]):
+            # Create a beam equivalency for brightness temperature
+            # This requires knowing the frequency along the spectral axis.
+            if freq is None:
+                try:
+                    freq = obj.with_spectral_unit(u.Hz).spectral_axis
+                except AttributeError:
+                    raise TypeError("Object of type {0} has no spectral "
+                                    "information. `freq` must be provided for"
+                                    " unit conversion from Jy/beam"
+                                    .format(type(obj)))
+            else:
+                if not freq.unit.is_equivalent(u.Hz):
+                    raise u.UnitsError("freq must be given in equivalent "
+                                        "frequency units.")
+
+                freq = freq.reshape((-1,))
+
+        else:
+            freq = [None]
+
+        # To handle varying resolution objects, loop through "channels"
+        # Default to a single iteration for a 2D spatial object or when a beam is not defined
+        # This allows handling all 1D, 2D, and 3D data products.
+        if has_beams:
+            iter = range(len(obj.beams))
+            beams = obj.beams
+        elif has_beam:
+            iter = range(0, 1)
+            beams = [obj.beam]
+        else:
+            iter = range(0, 1)
+            beams = [None]
+
+        # Append the unit conversion factors
+        factors = []
+
+        # Iterate through spectral channels.
+        for ii in iter:
+
+            beam = beams[ii]
+
+            # Use the range of frequencies when the beam does not change. Otherwise, select the
+            # frequency corresponding to this beam.
+            if has_beams:
+                thisfreq = freq[ii]
+            else:
+                thisfreq = freq
+
+            # Changes in beam require a new equivalency for each.
+            this_equivalencies = deepcopy(equivalencies)
+
+            # Equivalencies for Jy per ang area.
+            if has_perangarea:
+                bmequiv_angarea = u.brightness_temperature(thisfreq)
+
+                this_equivalencies = list(this_equivalencies) + bmequiv_angarea
+
+            # Beam area equivalencies for Jy per beam and/or Jy per ang area
+            if has_perbeam or has_perangarea:
+                if not has_beam:
+                    raise ValueError("To convert cubes with Jy/beam units, "
+                                    "the cube needs to have a beam defined.")
+
+                # create a beam equivalency for brightness temperature
+                bmequiv = beam.jtok_equiv(thisfreq)
+
+                # TODO: Remove check once `beamarea_equiv` is in a radio-beam release.
+                if hasattr(beam, 'beamarea_equiv'):
+                    bmarea_equiv = beam.beamarea_equiv
+                else:
+                    bmarea_equiv = u.beam_angular_area(beam.sr)
+
+                this_equivalencies = list(this_equivalencies) + bmequiv + bmarea_equiv
+
+            # Equivalencies for Jy per pixel area.
+            if has_perpix:
+
+                if not obj.wcs.has_celestial:
+                    raise ValueError("Spatial WCS information is required for unit conversions"
+                                    " involving spatial areas (e.g., Jy/pix, Jy/sr)")
+
+                pix_area = (proj_plane_pixel_area(obj.wcs.celestial) * u.deg**2).to(u.sr)
+
+                pix_area_equiv = [(u.Jy / u.pix, u.Jy / u.sr,
+                                lambda x: x / pix_area.value,
+                                lambda x: x * pix_area.value)]
+
+                this_equivalencies = list(this_equivalencies) + pix_area_equiv
+
+                # Define full from brightness temp to Jy / pix.
+                # Otherwise isn't working in 1 step
+                if has_btemp:
+                    if not has_beam:
+                        raise ValueError("Conversions between K and Jy/beam or Jy/pix"
+                                        "requires the cube to have a beam defined.")
+
+                    jtok_factor = beam.jtok(thisfreq) / (u.Jy / u.beam)
+
+                    # We're going to do this piecemeal because it's easier to conceptualize
+                    # We specifically anchor these conversions based on the beam area. So from
+                    # beam to pix, this is beam -> angular area -> area per pixel
+                    # Altogether:
+                    # K ->  Jy/beam -> Jy /sr - > Jy / pix
+                    forward_factor = 1 / (jtok_factor * (beam.sr / u.beam) / (pix_area / u.pix))
+                    reverse_factor = jtok_factor * (beam.sr / u.beam) / (pix_area / u.pix)
+
+                    pix_area_btemp_equiv = [(u.K, u.Jy / u.pix,
+                                            lambda x: x * forward_factor.value,
+                                            lambda x: x * reverse_factor.value)]
+
+                    this_equivalencies = list(this_equivalencies) + pix_area_btemp_equiv
+
+                # Equivalencies between pixel and angular areas.
+                if has_perbeam:
+                    if not has_beam:
+                        raise ValueError("Conversions between Jy/beam or Jy/pix"
+                                        "requires the cube to have a beam defined.")
+
+                    beam_area = beam.sr
+
+                    pix_area_btemp_equiv = [(u.Jy / u.pix, u.Jy / u.beam,
+                                            lambda x: x * (beam_area / pix_area).value,
+                                            lambda x: x * (pix_area / beam_area).value)]
+
+                    this_equivalencies = list(this_equivalencies) + pix_area_btemp_equiv
+
+            factor = obj.unit.to(unit, equivalencies=this_equivalencies)
+            factors.append(factor)
+
+        if has_beams:
+            return factors
+        else:
+            # Slice along first axis to return a 1D array.
+            return factors[0]

spectral_cube/lower_dimensional_structures.py

@@ -17,7 +17,7 @@
 from . import spectral_axis
 from .io.core import LowerDimensionalObjectWrite
 from .utils import SliceWarning, BeamWarning, SmoothingWarning, FITSWarning
-from .cube_utils import convert_bunit
+from . import cube_utils
 from . import wcs_utils
 from .masks import BooleanArrayMask, MaskBase
 
@@ -26,7 +26,6 @@
                          MultiBeamMixinClass, BeamMixinClass,
                          HeaderMixinClass
                         )
-from . import cube_utils
 
 __all__ = ['LowerDimensionalObject', 'Projection', 'Slice', 'OneDSpectrum']
 class LowerDimensionalObject(u.Quantity, BaseNDClass, HeaderMixinClass):
@@ -162,48 +161,15 @@ def to(self, unit, equivalencies=[], freq=None):
             # No copying
             return self
 
-        if ((self.unit.is_equivalent(u.Jy / u.beam) and
-             not any({u.Jy/u.beam, u.K}.issubset(set(eq)) for eq in equivalencies))):
-            # the 'not any' above checks that there is not already a defined
-            # Jy<->K equivalency.  If there is, the code below is redundant
-            # and will cause problems.
+        if hasattr(self, 'with_spectral_unit'):
+            freq = self.with_spectral_unit(u.Hz).spectral_axis
 
-            if hasattr(self, 'beams'):
-                factor = (self.jtok_factors(equivalencies=equivalencies) *
-                          (self.unit*u.beam).to(u.Jy))
-            else:
-                # replace "beam" with the actual beam
-                if not hasattr(self, 'beam'):
-                    raise ValueError("To convert objects with Jy/beam units, "
-                                     "the object needs to have a beam defined.")
-                brightness_unit = self.unit * u.beam
-
-                # create a beam equivalency for brightness temperature
-                if freq is None:
-                    try:
-                        freq = self.with_spectral_unit(u.Hz).spectral_axis
-                    except AttributeError:
-                        raise TypeError("Object of type {0} has no spectral "
-                                        "information. `freq` must be provided for"
-                                        " unit conversion from Jy/beam"
-                                        .format(type(self)))
-                else:
-                    if not freq.unit.is_equivalent(u.Hz):
-                        raise u.UnitsError("freq must be given in equivalent "
-                                           "frequency units.")
-
-                bmequiv = self.beam.jtok_equiv(freq)
-                # backport to handle astropy < 3: the beam equivalency was only
-                # modified to handle jy/beam in astropy 3
-                if bmequiv[0] == u.Jy:
-                    bmequiv.append([u.Jy/u.beam, u.K, bmequiv[2], bmequiv[3]])
-
-                factor = brightness_unit.to(unit,
-                                            equivalencies=bmequiv + list(equivalencies))
+        if freq is None and 'RESTFRQ' in self.header:
+            freq = self.header['RESTFRQ'] * u.Hz
 
-        else:
-            # scaling factor
-            factor = self.unit.to(unit, equivalencies=equivalencies)
+        # Create the tuple of unit conversions needed.
+        factor = cube_utils.bunit_converters(self, unit, equivalencies=equivalencies,
+                                             freq=freq)
 
         converted_array = (self.quantity * factor).value
 
@@ -412,7 +378,7 @@ def from_hdu(hdu):
         mywcs = wcs.WCS(hdu.header)
 
         if "BUNIT" in hdu.header:
-            unit = convert_bunit(hdu.header["BUNIT"])
+            unit = cube_utils.convert_bunit(hdu.header["BUNIT"])
             meta["BUNIT"] = hdu.header["BUNIT"]
         else:
             unit = None
@@ -673,7 +639,7 @@ def from_hdu(hdu):
         mywcs = wcs.WCS(hdu.header)
 
         if "BUNIT" in hdu.header:
-            unit = convert_bunit(hdu.header["BUNIT"])
+            unit = cube_utils.convert_bunit(hdu.header["BUNIT"])
             meta["BUNIT"] = hdu.header["BUNIT"]
         else:
             unit = None

spectral_cube/spectral_cube.py

@@ -2480,20 +2480,8 @@ def to(self, unit, equivalencies=()):
             # No copying
             return self
 
-        if self.unit.is_equivalent(u.Jy/u.beam):
-            # replace "beam" with the actual beam
-            if not hasattr(self, 'beam') or self._beam is None:
-                raise ValueError("To convert cubes with Jy/beam units, "
-                                 "the cube needs to have a beam defined.")
-            brightness_unit = self.unit * u.beam
-
-            # create a beam equivalency for brightness temperature
-            bmequiv = self.beam.jtok_equiv(self.with_spectral_unit(u.Hz).spectral_axis)
-            factor = brightness_unit.to(unit,
-                                        equivalencies=bmequiv+list(equivalencies))
-        else:
-            # scaling factor
-            factor = self.unit.to(unit, equivalencies=equivalencies)
+        # Create the tuple of unit conversions needed.
+        factor = cube_utils.bunit_converters(self, unit, equivalencies=equivalencies)
 
         # special case: array in equivalencies
         # (I don't think this should have to be special cased, but I don't know
@@ -4062,15 +4050,9 @@ def to(self, unit, equivalencies=()):
             # No copying
             return self
 
-        if self.unit.is_equivalent(u.Jy/u.beam) and unit.is_equivalent(u.K):
-            # replace "beam" with the actual beam
-            if not hasattr(self, 'beams'):
-                raise ValueError("To convert cubes with Jy/beam units, "
-                                 "the cube needs to have beams defined.")
-            factor = self.jtok_factors(equivalencies=equivalencies) * (self.unit*u.beam).to(u.Jy)
-        else:
-            # scaling factor
-            factor = self.unit.to(unit, equivalencies=equivalencies)
+        # Create the tuple of unit conversions needed.
+        factor = cube_utils.bunit_converters(self, unit, equivalencies=equivalencies)
+        factor = np.array(factor)
 
         # special case: array in equivalencies
         # (I don't think this should have to be special cased, but I don't know