save file from collapse and spectral extraction plugins (spacetelesco…

…pe#2586)

CHANGES.rst

@@ -8,6 +8,9 @@ Cubeviz
 ^^^^^^^
 - Calculated moments can now be output in velocity units. [#2584]
 
+- Added functionality to Collapse and Spectral Extraction plugins to save results to FITS file. [#2586]
+
+
 Imviz
 ^^^^^
 

jdaviz/components/tooltip.vue

@@ -97,6 +97,8 @@ const tooltips = {
   'plugin-line-analysis-sync-identify': 'Lock/unlock selection with identified line',
   'plugin-line-analysis-assign': 'Assign the centroid wavelength and update the redshift',
   'plugin-moment-save-fits': 'Save moment map as FITS file',
+  'plugin-extract-save-fits': 'Save spectral extraction as FITS file',
+  'plugin-collapse-save-fits': 'Save collapsed cube as FITS file',
   'plugin-link-apply': 'Apply linking to data',
   'plugin-footprints-color-picker': 'Change the color of the footprint overlay',
 }

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

@@ -88,7 +88,7 @@
       </span>
     </v-row>
 
-    <j-plugin-section-header v-if="export_enabled">Results</j-plugin-section-header>
+    <j-plugin-section-header v-if="moment_available && export_enabled">Results</j-plugin-section-header>
 
     <div style="display: grid; position: relative"> <!-- overlay container -->
       <div style="grid-area: 1/1">

jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py

@@ -1,15 +1,19 @@
+import os
+from pathlib import Path
+
 from packaging.version import Version
 import numpy as np
 import astropy
 import astropy.units as u
 from astropy.nddata import (
     NDDataArray, StdDevUncertainty, NDUncertainty
 )
-from traitlets import List, Unicode, observe
+from traitlets import Bool, List, Unicode, observe
 
 from jdaviz.core.events import SnackbarMessage
 from jdaviz.core.registries import tray_registry
 from jdaviz.core.template_mixin import (PluginTemplateMixin,
+                                        DatasetSelectMixin,
                                         SelectPluginComponent,
                                         SpatialSubsetSelectMixin,
                                         AddResultsMixin,
@@ -25,7 +29,8 @@
 @tray_registry(
     'cubeviz-spectral-extraction', label="Spectral Extraction", viewer_requirements='spectrum'
 )
-class SpectralExtraction(PluginTemplateMixin, SpatialSubsetSelectMixin, AddResultsMixin):
+class SpectralExtraction(PluginTemplateMixin, DatasetSelectMixin,
+                         SpatialSubsetSelectMixin, AddResultsMixin):
     """
     See the :ref:`Spectral Extraction Plugin Documentation <spex>` for more details.
 
@@ -43,6 +48,13 @@ class SpectralExtraction(PluginTemplateMixin, SpatialSubsetSelectMixin, AddResul
     template_file = __file__, "spectral_extraction.vue"
     function_items = List().tag(sync=True)
     function_selected = Unicode('Sum').tag(sync=True)
+    filename = Unicode().tag(sync=True)
+    extracted_spec_available = Bool(False).tag(sync=True)
+    overwrite_warn = Bool(False).tag(sync=True)
+    # export_enabled controls whether saving to a file is enabled via the UI.  This
+    # is a temporary measure to allow server-installations to disable saving server-side until
+    # saving client-side is supported
+    export_enabled = Bool(True).tag(sync=True)
 
     def __init__(self, *args, **kwargs):
 
@@ -52,6 +64,8 @@ def __init__(self, *args, **kwargs):
 
         super().__init__(*args, **kwargs)
 
+        self.extracted_spec = None
+
         self.function = SelectPluginComponent(
             self,
             items='function_items',
@@ -64,6 +78,11 @@ def __init__(self, *args, **kwargs):
         if ASTROPY_LT_5_3_2:
             self.disabled_msg = "Spectral Extraction in Cubeviz requires astropy>=5.3.2"
 
+        if self.app.state.settings.get('server_is_remote', False):
+            # when the server is remote, saving the file in python would save on the server, not
+            # on the user's machine, so export support in cubeviz should be disabled
+            self.export_enabled = False
+
     @property
     def user_api(self):
         return PluginUserApi(
@@ -157,6 +176,12 @@ def collapse_to_spectrum(self, add_data=True, **kwargs):
             mask=mask
         )
 
+        # stuff for exporting to file
+        self.extracted_spec = collapsed_spec
+        self.extracted_spec_available = True
+        fname_label = self.dataset_selected.replace("[", "_").replace("]", "")
+        self.filename = f"extracted_{self.function_selected.lower()}_{fname_label}.fits"
+
         if add_data:
             self.add_results.add_results_from_plugin(
                 collapsed_spec, label=self.results_label, replace=False
@@ -173,6 +198,50 @@ def collapse_to_spectrum(self, add_data=True, **kwargs):
     def vue_spectral_extraction(self, *args, **kwargs):
         self.collapse_to_spectrum(add_data=True)
 
+    def vue_save_as_fits(self, *args):
+        self._save_extracted_spec_to_fits()
+
+    def vue_overwrite_fits(self, *args):
+        """Attempt to force writing the spectral extraction if the user
+        confirms the desire to overwrite."""
+        self.overwrite_warn = False
+        self._save_extracted_spec_to_fits(overwrite=True)
+
+    def _save_extracted_spec_to_fits(self, overwrite=False, *args):
+
+        if not self.export_enabled:
+            # this should never be triggered since this is intended for UI-disabling and the
+            # UI section is hidden, but would prevent any JS-hacking
+            raise ValueError("Writing out extracted spectrum to file is currently disabled")
+
+        # Make sure file does not end up in weird places in standalone mode.
+        path = os.path.dirname(self.filename)
+        if path and not os.path.exists(path):
+            raise ValueError(f"Invalid path={path}")
+        elif (not path or path.startswith("..")) and os.environ.get("JDAVIZ_START_DIR", ""):  # noqa: E501 # pragma: no cover
+            filename = Path(os.environ["JDAVIZ_START_DIR"]) / self.filename
+        else:
+            filename = Path(self.filename).resolve()
+
+        if filename.exists():
+            if overwrite:
+                # Try to delete the file
+                filename.unlink()
+                if filename.exists():
+                    # Warn the user if the file still exists
+                    raise FileExistsError(f"Unable to delete {filename}. Check user permissions.")
+            else:
+                self.overwrite_warn = True
+                return
+
+        filename = str(filename)
+        self.extracted_spec.write(filename)
+
+        # Let the user know where we saved the file.
+        self.hub.broadcast(SnackbarMessage(
+            f"Extracted spectrum saved to {os.path.abspath(filename)}",
+                           sender=self, color="success"))
+
     @observe('spatial_subset_selected')
     def _set_default_results_label(self, event={}):
         label = "Spectral extraction"

jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.vue

@@ -40,5 +40,59 @@
       @click:action="spectral_extraction"
     ></plugin-add-results>
 
+    <j-plugin-section-header v-if="extracted_spec_available && export_enabled">Results</j-plugin-section-header>
+
+    <div style="display: grid; position: relative"> <!-- overlay container -->
+      <div style="grid-area: 1/1">
+        <div v-if="extracted_spec_available && export_enabled">
+
+          <v-row>
+            <v-text-field
+            v-model="filename"
+            label="Filename"
+            hint="Export the latest extracted spectrum."
+            :rules="[() => !!filename || 'This field is required']"
+            persistent-hint>
+            </v-text-field>
+          </v-row>
+
+          <v-row justify="end">
+            <j-tooltip tipid='plugin-extract-save-fits'>
+              <v-btn color="primary" text @click="save_as_fits">Save as FITS</v-btn>
+
+            </j-tooltip>
+          </v-row>
+
+        </div>
+      </div>
+
+      <v-overlay
+        absolute
+        opacity=1.0
+        :value="overwrite_warn && export_enabled"
+        :zIndex=3
+        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">
+            A file with this name is already on disk. Overwrite?
+          </div>
+        </v-card-text>
+
+        <v-card-actions>
+          <v-row justify="end">
+            <v-btn tile small color="primary" class="mr-2" @click="overwrite_warn=false">Cancel</v-btn>
+            <v-btn tile small color="accent" class="mr-4" @click="overwrite_fits" >Overwrite</v-btn>
+          </v-row>
+        </v-card-actions>
+      </v-card>
+
+      </v-overlay>
+
+
+    </div>
   </j-tray-plugin>
 </template>

jdaviz/configs/cubeviz/plugins/spectral_extraction/tests/test_spectral_extraction.py

@@ -1,3 +1,4 @@
+import os
 import pytest
 from packaging.version import Version
 import numpy as np
@@ -87,3 +88,53 @@ def test_subset(
     collapsed_spec_2 = plg.collapse_to_spectrum()
 
     assert np.all(np.equal(collapsed_spec_2.uncertainty.array, expected_uncert))
+
+
+def test_save_collapsed_to_fits(cubeviz_helper, spectrum1d_cube_with_uncerts, tmpdir):
+
+    cubeviz_helper.load_data(spectrum1d_cube_with_uncerts)
+
+    extract_plugin = cubeviz_helper.plugins['Spectral Extraction']
+
+    # make sure export enabled is true, and that before the collapse function
+    # is run `collapsed_spec_available` is correctly set to False
+    assert extract_plugin._obj.export_enabled
+    assert extract_plugin._obj.extracted_spec_available is False
+
+    # run extract function, and make sure `extracted_spec_available` was set to True
+    extract_plugin._obj.vue_spectral_extraction()
+    assert extract_plugin._obj.extracted_spec_available
+
+    # check that default filename is correct, then change path
+    fname = 'extracted_sum_Unknown spectrum object_FLUX.fits'
+    assert extract_plugin._obj.filename == fname
+    extract_plugin._obj.filename = os.path.join(tmpdir, fname)
+
+    # save output file with default name, make sure it exists
+    extract_plugin._obj.vue_save_as_fits()
+    assert os.path.isfile(os.path.join(tmpdir, fname))
+
+    # read file back in, make sure it matches
+    dat = Spectrum1D.read(os.path.join(tmpdir, fname))
+    assert np.all(dat.data == extract_plugin._obj.extracted_spec.data)
+    assert dat.unit == extract_plugin._obj.extracted_spec.unit
+
+    # make sure correct error message is raised when export_enabled is False
+    # this won't appear in UI, but just to be safe.
+    extract_plugin._obj.export_enabled = False
+    msg = "Writing out extracted spectrum to file is currently disabled"
+    with pytest.raises(ValueError, match=msg):
+        extract_plugin._obj.vue_save_as_fits()
+    extract_plugin._obj.export_enabled = True  # set back to True
+
+    # check that trying to overwrite without overwrite=True sets overwrite_warn to True, to
+    # display popup in UI
+    assert extract_plugin._obj.overwrite_warn is False
+    extract_plugin._obj.vue_save_as_fits()
+    assert extract_plugin._obj.overwrite_warn
+
+    # check that writing out to a non existent directory fails as expected
+    extract_plugin._obj.filename = '/this/path/doesnt/exist.fits'
+    with pytest.raises(ValueError, match="Invalid path=/this/path/doesnt"):
+        extract_plugin._obj.vue_save_as_fits()
+    extract_plugin._obj.filename == fname  # set back to original filename

jdaviz/configs/default/plugins/collapse/collapse.py

@@ -1,9 +1,12 @@
+import os
+from pathlib import Path
 import warnings
 
+from astropy.nddata import CCDData
 from glue.core import Data
 from specutils import Spectrum1D
 from specutils.manipulation import spectral_slab
-from traitlets import List, Unicode, observe
+from traitlets import Bool, List, Unicode, observe
 
 from jdaviz.core.events import SnackbarMessage
 from jdaviz.core.registries import tray_registry
@@ -39,13 +42,22 @@ class Collapse(PluginTemplateMixin, DatasetSelectMixin, SpectralSubsetSelectMixi
     template_file = __file__, "collapse.vue"
     function_items = List().tag(sync=True)
     function_selected = Unicode('Sum').tag(sync=True)
+    filename = Unicode().tag(sync=True)
+    collapsed_spec_available = Bool(False).tag(sync=True)
+    overwrite_warn = Bool(False).tag(sync=True)
+    # export_enabled controls whether saving to a file is enabled via the UI.  This
+    # is a temporary measure to allow server-installations to disable saving server-side until
+    # saving client-side is supported
+    export_enabled = Bool(True).tag(sync=True)
 
     def __init__(self, *args, **kwargs):
 
         super().__init__(*args, **kwargs)
 
         self._label_counter = 0
 
+        self.collapsed_spec = None
+
         self.function = SelectPluginComponent(self,
                                               items='function_items',
                                               selected='function_selected',
@@ -54,6 +66,11 @@ def __init__(self, *args, **kwargs):
         self.dataset.add_filter('is_cube')
         self.add_results.viewer.filters = ['is_image_viewer']
 
+        if self.app.state.settings.get('server_is_remote', False):
+            # when the server is remote, saving the file in python would save on the server, not
+            # on the user's machine, so export support in cubeviz should be disabled
+            self.export_enabled = False
+
     @property
     def _default_spectrum_viewer_reference_name(self):
         return self.jdaviz_helper._default_spectrum_viewer_reference_name
@@ -93,6 +110,12 @@ def collapse(self, add_data=True):
             # Spatial-spatial image only.
             collapsed_spec = spec.collapse(self.function_selected.lower(), axis=-1).T  # Quantity
 
+            # stuff for exporting to file
+            self.collapsed_spec = collapsed_spec
+            self.collapsed_spec_available = True
+            fname_label = self.dataset_selected.replace("[", "_").replace("]", "")
+            self.filename = f"collapsed_{self.function_selected.lower()}_{fname_label}.fits"
+
         if add_data:
             data = Data()
             data['flux'] = collapsed_spec.value
@@ -110,3 +133,46 @@ def collapse(self, add_data=True):
 
     def vue_collapse(self, *args, **kwargs):
         self.collapse(add_data=True)
+
+    def vue_save_as_fits(self, *args):
+        self._save_collapsed_spec_to_fits()
+
+    def vue_overwrite_fits(self, *args):
+        """Attempt to force writing the file if the user confirms the desire to overwrite."""
+        self.overwrite_warn = False
+        self._save_collapsed_spec_to_fits(overwrite=True)
+
+    def _save_collapsed_spec_to_fits(self, overwrite=False, *args):
+
+        if not self.export_enabled:
+            # this should never be triggered since this is intended for UI-disabling and the
+            # UI section is hidden, but would prevent any JS-hacking
+            raise ValueError("Writing out collapsed cube to file is currently disabled")
+
+        # Make sure file does not end up in weird places in standalone mode.
+        path = os.path.dirname(self.filename)
+        if path and not os.path.exists(path):
+            raise ValueError(f"Invalid path={path}")
+        elif (not path or path.startswith("..")) and os.environ.get("JDAVIZ_START_DIR", ""):  # noqa: E501 # pragma: no cover
+            filename = Path(os.environ["JDAVIZ_START_DIR"]) / self.filename
+        else:
+            filename = Path(self.filename).resolve()
+
+        if filename.exists():
+            if overwrite:
+                # Try to delete the file
+                filename.unlink()
+                if filename.exists():
+                    # Warn the user if the file still exists
+                    raise FileExistsError(f"Unable to delete {filename}. Check user permissions.")
+            else:
+                self.overwrite_warn = True
+                return
+
+        filename = str(filename)
+        CCDData(self.collapsed_spec).write(filename)
+
+        # Let the user know where we saved the file.
+        self.hub.broadcast(SnackbarMessage(
+            f"Collapsed cube saved to {os.path.abspath(filename)}",
+                           sender=self, color="success"))