Add option to calculate higher-order moments in velocity

CHANGES.rst

@@ -6,6 +6,7 @@ New Features
 
 Cubeviz
 ^^^^^^^
+- Calculated moments can now be output in velocity units. [#2584]
 
 Imviz
 ^^^^^

docs/cubeviz/plugins.rst

@@ -76,10 +76,10 @@ corresponding wavelength highlighted in the spectrum viewer.
 
 The slider can be grabbed to scrub through the cube.  To choose
 a specific slice, enter a slice number (integer) or an approximate
-wavelength (in which case the nearest slice will be selected and 
+wavelength (in which case the nearest slice will be selected and
 the wavelength entry will update to the exact value of that slice).
 
-The spectrum viewer also contains a tool to allow clicking and 
+The spectrum viewer also contains a tool to allow clicking and
 dragging in the spectrum plot to choose the currently selected slice.
 When the slice tool is active, clicking anywhere on the spectrum viewer
 will select the nearest slice across all viewers, even if the indicator
@@ -142,7 +142,7 @@ Model Fitting
         Specviz documentation on fitting spectral models.
 
 For Cubeviz, there is an additional option to fit the model over each individual spaxel by
-enabling the :guilabel:`Cube Fit` toggle before pressing :guilabel:`Fit Model`.  
+enabling the :guilabel:`Cube Fit` toggle before pressing :guilabel:`Fit Model`.
 The best-fit parameters for each spaxel are stored in planes and saved in a data structure.
 The resulting model itself is saved with the label specified in the :guilabel:`Output Data Label` field.
 
@@ -216,15 +216,16 @@ map is created. You can load this into any image
 viewer pane to inspect the result. You can also save the result to
 a FITS format file by pressing :guilabel:`SAVE AS FITS`.
 
-For example, the middle image viewer in the screenshot above shows the Moment 1 map
+For example, the right image viewer in the screenshot above shows the Moment 2 map
 for a continuum-subtracted cube.  Note that the cube should first be
 continuum-subtracted in order to create continuum-free moment maps of an
 emission line. Moment maps of continuum emission can also be created, but
 moments other than moment 0 may not be physically meaningful.  Also note
-that the units in the moment 1 and moment 2 maps reflect the units of the spectral
-axis (Angstroms in this case). The units of the input cube should first be
-converted to velocity units before running the plugin if those units are
-desired for the output moment maps.
+that by default, the units in the moment 1 and moment 2 maps reflect the units of the
+spectral axis (microns in this case). For moments higher than 0, the output units can
+instead be converted to velocity (e.g., m/s for moment 1, m2/s2 for moment 2, etc.) by
+selecting the :guilabel:`Velocity` radio button under :guilabel:`Output Units`
+and providing a reference wavelength, commonly that of the spectral line of interest.
 
 Line or Continuum Maps
 ----------------------

jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.py

@@ -2,18 +2,21 @@
 from pathlib import Path
 
 from astropy import units as u
+from astropy.constants import c
 from astropy.nddata import CCDData
+import numpy as np
 
-from traitlets import Unicode, Bool, observe
+from traitlets import Bool, List, Unicode, observe
 from specutils import manipulation, analysis
 
-from jdaviz.core.custom_traitlets import IntHandleEmpty
+from jdaviz.core.custom_traitlets import IntHandleEmpty, FloatHandleEmpty
 from jdaviz.core.events import SnackbarMessage
 from jdaviz.core.registries import tray_registry
 from jdaviz.core.template_mixin import (PluginTemplateMixin,
                                         DatasetSelectMixin,
                                         SpectralSubsetSelectMixin,
                                         AddResultsMixin,
+                                        SelectPluginComponent,
                                         with_spinner)
 from jdaviz.core.user_api import PluginUserApi
 
@@ -42,6 +45,10 @@ class MomentMap(PluginTemplateMixin, DatasetSelectMixin, SpectralSubsetSelectMix
     * ``spectral_subset`` (:class:`~jdaviz.core.template_mixin.SubsetSelect`):
       Subset to use for the line, or ``Entire Spectrum``.
     * ``n_moment``
+    * ``output_unit``
+      Choice of "Wavelength" or "Velocity", applicable for n_moment >= 1.
+    * ``reference_wavelength``
+      Reference wavelength for conversion of output to velocity units.
     * ``add_results`` (:class:`~jdaviz.core.template_mixin.AddResults`)
     * :meth:`calculate_moment`
     """
@@ -51,6 +58,10 @@ class MomentMap(PluginTemplateMixin, DatasetSelectMixin, SpectralSubsetSelectMix
     filename = Unicode().tag(sync=True)
     moment_available = Bool(False).tag(sync=True)
     overwrite_warn = Bool(False).tag(sync=True)
+    output_unit_items = List().tag(sync=True)
+    output_unit_selected = Unicode().tag(sync=True)
+    reference_wavelength = FloatHandleEmpty().tag(sync=True)
+    dataset_spectral_unit = Unicode().tag(sync=True)
 
     # export_enabled controls whether saving the moment map to a file is enabled via the UI.  This
     # is a temporary measure to allow server-installations to disable saving server-side until
@@ -62,6 +73,11 @@ def __init__(self, *args, **kwargs):
 
         self.moment = None
 
+        self.output_unit = SelectPluginComponent(self,
+                                                 items='output_unit_items',
+                                                 selected='output_unit_selected',
+                                                 manual_options=['Wavelength', 'Velocity'])
+
         self.dataset.add_filter('is_cube')
         self.add_results.viewer.filters = ['is_image_viewer']
 
@@ -83,6 +99,7 @@ def user_api(self):
         # NOTE: leaving save_as_fits out for now - we may want a more general API to do that
         # accross all plugins at some point
         return PluginUserApi(self, expose=('dataset', 'spectral_subset', 'n_moment',
+                                           'output_unit', 'reference_wavelength',
                                            'add_results', 'calculate_moment'))
 
     @observe("dataset_selected", "dataset_items", "n_moment")
@@ -93,6 +110,16 @@ def _set_default_results_label(self, event={}):
         label_comps += [f"moment {self.n_moment}"]
         self.results_label_default = " ".join(label_comps)
 
+    @observe("dataset_selected")
+    def _set_data_units(self, event={}):
+        if self.dataset_selected != "":
+            # Spectral axis is first in this list
+            if self.app.data_collection[self.dataset_selected].coords is not None:
+                unit = self.app.data_collection[self.dataset_selected].coords.world_axis_units[0]
+                self.dataset_spectral_unit = unit
+            else:
+                self.dataset_spectral_unit = ""
+
     @with_spinner()
     def calculate_moment(self, add_data=True):
         """
@@ -126,6 +153,17 @@ def calculate_moment(self, add_data=True):
         if data_wcs:
             data_wcs = data_wcs.swapaxes(0, 1)  # We also transpose WCS to match.
         self.moment = CCDData(analysis.moment(slab, order=n_moment).T, wcs=data_wcs)
+        if n_moment > 0 and self.output_unit_selected.lower() == "velocity":
+            # Catch this if called from API
+            if not self.reference_wavelength > 0.0:
+                raise ValueError("reference_wavelength must be set for output in velocity units.")
+            power_unit = f"{self.dataset_spectral_unit}{self.n_moment}"
+            self.moment = np.power(self.moment.convert_unit_to(power_unit), 1/self.n_moment)
+            self.moment = self.moment << u.Unit(self.dataset_spectral_unit)
+            ref_wavelength = self.reference_wavelength * u.Unit(self.dataset_spectral_unit)
+            relative_wavelength = (self.moment-ref_wavelength)/ref_wavelength
+            in_velocity = np.power(c*relative_wavelength, self.n_moment)
+            self.moment = CCDData(in_velocity, wcs=data_wcs)
 
         fname_label = self.dataset_selected.replace("[", "_").replace("]", "")
         self.filename = f"moment{n_moment}_{fname_label}.fits"

jdaviz/configs/cubeviz/plugins/moment_maps/moment_maps.vue

@@ -12,7 +12,7 @@
       hint="Select the data set."
     />
 
-    <plugin-subset-select 
+    <plugin-subset-select
       :items="spectral_subset_items"
       :selected.sync="spectral_subset_selected"
       :has_subregions="spectral_subset_selected_has_subregions"
@@ -35,6 +35,37 @@
       ></v-text-field>
     </v-row>
 
+    <div v-if="n_moment > 0 && dataset_spectral_unit !== ''">
+      <v-row>
+        <v-radio-group
+          label="Output Units"
+          hint="Choose whether calculated moment is in units of wavelength or velocity."
+          v-model="output_unit_selected"
+          column
+          class="no-hint">
+          <v-radio
+            v-for="item in output_unit_items"
+            :key="item.label"
+            :label="item.label"
+            :value="item.label"
+          ></v-radio>
+        </v-radio-group>
+      </v-row>
+      <v-row v-if="output_unit_selected === 'Velocity'">
+        <v-text-field
+        ref="reference_wavelength"
+        type="number"
+        label="Reference Wavelength"
+        v-model.number="reference_wavelength"
+        :suffix="dataset_spectral_unit.replace('Angstrom', 'A')"
+        hint="Rest wavelength of the line of interest"
+        persistent-hint
+        :rules="[() => reference_wavelength !== '' || 'This field is required',
+                 () => reference_wavelength > 0 || 'Wavelength must be positive']"
+        ></v-text-field>
+      </v-row>
+    </div>
+
     <plugin-add-results
       :label.sync="results_label"
       :label_default="results_label_default"
@@ -49,7 +80,7 @@
       :action_spinner="spinner"
       @click:action="calculate_moment"
     ></plugin-add-results>
-    
+
     <j-plugin-section-header v-if="export_enabled">Results</j-plugin-section-header>
 
     <div style="display: grid; position: relative"> <!-- overlay container -->
@@ -80,10 +111,10 @@
         opacity=1.0
         :value="overwrite_warn && export_enabled"
         :zIndex=3
-        style="grid-area: 1/1; 
+        style="grid-area: 1/1;
                margin-left: -24px;
                margin-right: -24px">
-      
+
       <v-card color="transparent" elevation=0 >
         <v-card-text width="100%">
           <div class="white--text">

jdaviz/configs/cubeviz/plugins/moment_maps/tests/test_moment_maps.py

@@ -28,6 +28,7 @@ def test_moment_calculation(cubeviz_helper, spectrum1d_cube, tmpdir):
 
     mm.n_moment = 0  # Collapsed sum, will get back 2D spatial image
     assert mm.results_label == 'moment 0'
+    assert mm.output_unit == "Wavelength"
 
     mm.add_results.viewer.selected = cubeviz_helper._default_uncert_viewer_reference_name
     mm.vue_calculate_moment()
@@ -99,6 +100,52 @@ def test_moment_calculation(cubeviz_helper, spectrum1d_cube, tmpdir):
                                          "204.9998877673 27.0001000000 (deg)")
 
 
+def test_moment_velocity_calculation(cubeviz_helper, spectrum1d_cube, tmpdir):
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message="No observer defined on WCS.*")
+        cubeviz_helper.load_data(spectrum1d_cube, data_label='test')
+
+    uncert_viewer = cubeviz_helper.app.get_viewer("uncert-viewer")
+
+    # Since we are not really displaying, need this to trigger GUI stuff.
+    uncert_viewer.shape = (100, 100)
+    uncert_viewer.state._set_axes_aspect_ratio(1)
+
+    mm = cubeviz_helper.plugins["Moment Maps"]
+    print(mm._obj.dataset_selected)
+    mm._obj.dataset_selected = 'test[FLUX]'
+
+    # Test moment 1 in velocity
+    mm.n_moment = 1
+    mm.add_results.viewer = "uncert-viewer"
+    assert mm._obj.results_label == 'moment 1'
+    mm.output_unit = "Velocity"
+
+    with pytest.raises(ValueError, match="reference_wavelength must be set"):
+        mm.calculate_moment()
+
+    mm.reference_wavelength = 4.63e-7
+    mm.calculate_moment()
+
+    # Make sure coordinate display works
+    label_mouseover = cubeviz_helper.app.session.application._tools['g-coords-info']
+    label_mouseover._viewer_mouse_event(uncert_viewer, {'event': 'mousemove',
+                                                        'domain': {'x': 0, 'y': 0}})
+    assert label_mouseover.as_text() == ("Pixel x=00.0 y=00.0 Value -4.14382e+05 m / s",
+                                         "World 13h39m59.9731s +27d00m00.3600s (ICRS)",
+                                         "204.9998877673 27.0001000000 (deg)")
+
+    # Test moment 2 in velocity
+    mm.n_moment = 2
+    mm.calculate_moment()
+
+    label_mouseover._viewer_mouse_event(uncert_viewer, {'event': 'mousemove',
+                                                        'domain': {'x': 0, 'y': 0}})
+    assert label_mouseover.as_text() == ("Pixel x=00.0 y=00.0 Value +8.98755e+16 m2 / s2",
+                                         "World 13h39m59.9731s +27d00m00.3600s (ICRS)",
+                                         "204.9998877673 27.0001000000 (deg)")
+
+
 def test_write_momentmap(cubeviz_helper, spectrum1d_cube, tmp_path):
     ''' Test writing a moment map out to a FITS file on disk '''