Create import_region method for loading subsets

jdaviz/configs/default/plugins/subset_plugin/subset_plugin.py

@@ -1,23 +1,27 @@
 import os
+import json
 
 import numpy as np
 
 from astropy.time import Time
 import astropy.units as u
 from glue.core.message import EditSubsetMessage, SubsetUpdateMessage
 from glue.core.edit_subset_mode import (AndMode, AndNotMode, OrMode,
-                                        ReplaceMode, XorMode)
-from glue.core.roi import CircularROI, CircularAnnulusROI, EllipticalROI, RectangularROI
+                                        ReplaceMode, XorMode, NewMode)
+
+from glue.core.roi import CircularROI, CircularAnnulusROI, EllipticalROI, RectangularROI, XRangeROI
 from glue.core.subset import RoiSubsetState, RangeSubsetState, CompositeSubsetState
 from glue.icons import icon_path
 from glue_jupyter.widgets.subset_mode_vuetify import SelectionModeMenu
 from glue_jupyter.common.toolbar_vuetify import read_icon
 from traitlets import Any, List, Unicode, Bool, observe
+from specutils import SpectralRegion
 
 from jdaviz.core.events import SnackbarMessage, GlobalDisplayUnitChanged, LinkUpdatedMessage
 from jdaviz.core.registries import tray_registry
 from jdaviz.core.template_mixin import PluginTemplateMixin, DatasetSelectMixin, SubsetSelect
 from jdaviz.core.tools import ICON_DIR
+from jdaviz.core import region_translators
 from jdaviz.utils import MultiMaskSubsetState
 
 from jdaviz.configs.default.plugins.subset_plugin import utils
@@ -655,3 +659,53 @@
             if self.subset_definitions[index][i]['name'] == name:
                 self.subset_definitions[index][i]['value'] = new_value
                 break
+
+    def _import_spectral_regions(self, spec_region, mode=NewMode):
+        """
+        Method for importing a SpectralRegion object or list of SpectralRegion objects.
+
+        Parameters
+        ----------
+        spec_region : ~`specutils.SpectralRegion` object or list
+            The object that contains bounds for the spectral region or regions
+        mode : AndMode, AndNotMode, OrMode, ReplaceMode, XorMode, or NewMode
+            By default this is set to NewMode which creates a new subset per subregion.
+            OrMode will create a composite subset with each subregion corresponding
+            to a subregion of a single subset.
+        """
+        viewer_name = self.app._jdaviz_helper._default_spectrum_viewer_reference_name
+        spectrum_viewer = self.app.get_viewer(viewer_name)
+        for sub_region in spec_region:
+            self.app.session.edit_subset_mode.mode = mode
+            spectrum_viewer.apply_roi(XRangeROI(sub_region.lower.value, sub_region.upper.value))
+
+    def import_region(self, spec_region):
+        if isinstance(spec_region, list):
+            if len(spec_region) < 1 or ('region' in spec_region[0] and
+                                        not spec_region[0]['region']):
+                return
+            elif isinstance(spec_region[0], SpectralRegion):
+                self._import_spectral_regions(spec_region)
+            else:
+                self._import_regions_list(spec_region)
+        elif isinstance(spec_region, dict):
+            for key, value in spec_region.items():
+                self.import_region(value)
+        elif isinstance(spec_region, str):
+            try:
+                dict_region = json.loads(spec_region)
+            except ValueError:
+                raise ValueError
+            self.import_region(dict_region)
+        elif isinstance(spec_region, SpectralRegion):
+            self._import_spectral_regions(spec_region)
+
+    def _import_regions_list(self, spec_region):
+        viewer_name = self.app._jdaviz_helper._default_flux_viewer_reference_name
+        flux_viewer = self.app.get_viewer(viewer_name)
+        for subregion in spec_region:
+            if subregion['glue_state'] == 'RoiSubsetState':
+                self.app.session.edit_subset_mode.mode = NewMode
+            else:
+                self.app.session.edit_subset_mode.mode = SUBSET_MODES[subregion['glue_state']]
+            flux_viewer.apply_roi(region_translators.regions2roi(subregion['region']))

jdaviz/configs/default/plugins/subset_plugin/tests/test_subset_plugin.py

@@ -3,12 +3,19 @@
 
 import numpy as np
 from astropy.nddata import NDData
+import astropy.units as u
+from specutils import SpectralRegion
+
 from glue.core.roi import EllipticalROI, CircularROI, CircularAnnulusROI, RectangularROI, XRangeROI
+from glue.core.edit_subset_mode import OrMode, NewMode
 from numpy.testing import assert_allclose
 
 from jdaviz.configs.default.plugins.subset_plugin import utils
 from jdaviz.core.region_translators import regions2roi
 
+from regions import (CirclePixelRegion, EllipsePixelRegion, RectanglePixelRegion,
+                     CircleAnnulusPixelRegion, PixCoord)
+
 
 @pytest.mark.filterwarnings('ignore')
 def test_plugin(specviz_helper, spectrum1d):
@@ -168,3 +175,91 @@ def test_circle_recenter_linking(roi_class, subset_info, imviz_helper, image_2d_
     for i, attr in enumerate(subset_info):
         assert subset_defs[0][i+1]['name'] == subset_info[attr]['wcs_name']
         assert_allclose(subset_defs[0][i+1]['value'], true_values_final[i]['value'])
+
+
+@pytest.mark.parametrize(
+    ('spec_regions', 'mode', 'len_subsets', 'len_subregions'),
+    [([SpectralRegion(5.772486091213352 * u.um, 6.052963676101135 * u.um),
+       SpectralRegion(6.494371022809778 * u.um, 6.724270682553864 * u.um),
+       SpectralRegion(7.004748267441649 * u.um, 7.3404016303483965 * u.um)], NewMode, 3, 1),
+     ((SpectralRegion(5.772486091213352 * u.um, 6.052963676101135 * u.um) +
+       SpectralRegion(6.494371022809778 * u.um, 6.724270682553864 * u.um) +
+       SpectralRegion(7.004748267441649 * u.um, 7.3404016303483965 * u.um)), OrMode, 1, 3)]
+)
+def test_import_spectral_region(cubeviz_helper, spectrum1d_cube, spec_regions, mode, len_subsets,
+                                len_subregions):
+    cubeviz_helper.load_data(spectrum1d_cube)
+    cubeviz_helper.app.get_tray_item_from_name('g-subset-plugin')
+    plg = cubeviz_helper.plugins['Subset Tools']._obj
+    plg._import_spectral_regions(spec_region=spec_regions, mode=mode)
+    subsets = cubeviz_helper.app.get_subsets()
+    assert len(subsets) == len_subsets
+    assert len(subsets['Subset 1']) == len_subregions
+
+
+test_subset_dict = {'Subset 1': [
+    {'name': 'TrueCircularROI',
+     'glue_state': 'AndState',
+     'region': CirclePixelRegion(center=PixCoord(x=20.6072998046875, y=23.749065399169922),
+                                 radius=4.079286151278358),
+     'sky_region': None,
+     'subset_state': None},
+    {'name': 'RectangularROI',
+     'glue_state': 'OrState',
+     'region': RectanglePixelRegion(center=PixCoord(x=20.68303102318066, y=30.615620480255785),
+                                    width=6.815691577924664, height=4.272523078699049,
+                                    angle=0.0 * u.rad),
+     'sky_region': None,
+     'subset_state': None},
+    {'name': 'EllipticalROI',
+     'glue_state': 'AndNotState',
+     'region': EllipsePixelRegion(center=PixCoord(x=20.30438232421875, y=27.716407775878906),
+                                  width=3.427145004272461, height=8.200210571289062,
+                                  angle=0.0 * u.rad),
+     'sky_region': None,
+     'subset_state': None}],
+                       'Subset 2': [
+                           {'name': 'CircularAnnulusROI',
+                            'glue_state': 'RoiSubsetState',
+                            'region': CircleAnnulusPixelRegion(
+                                center=PixCoord(x=30.452190399169922, y=22.070573806762695),
+                                inner_radius=2.4145185947418213, outer_radius=4.829037189483643),
+                            'sky_region': None,
+                            'subset_state': None}],
+                       'Subset 3': SpectralRegion(6.939254409861322 * u.um,
+                                                  7.224590119773996 * u.um)
+}
+
+test_subset_string = str(test_subset_dict)
+
+region = [{'name': 'TrueCircularROI',
+           'glue_state': 'RoiSubsetState',
+           'region': CirclePixelRegion(center=PixCoord(x=25.888399124145508, y=22.078184127807617),
+                                       radius=3.7648651881914112),
+           'sky_region': None,
+           'subset_state': ''},
+          {'name': 'TrueCircularROI',
+           'glue_state': 'OrState',
+           'region': CirclePixelRegion(center=PixCoord(x=19.990379333496094, y=30.782743453979492),
+                                       radius=2.756227940196304),
+           'sky_region': None,
+           'subset_state': ''}]
+
+
+@pytest.mark.parametrize(
+    ('spec_region', 'len_subsets', 'len_subregions'),
+    [(test_subset_dict, 3, 3),
+     # (test_subset_string, 3, 3),
+     (region, 1, 2)]
+)
+def test_add_subset_with_import_region(cubeviz_helper, spectrum1d_cube, spec_region,
+                                       len_subsets, len_subregions):
+    with warnings.catch_warnings():
+        warnings.simplefilter('ignore')
+        cubeviz_helper.load_data(spectrum1d_cube)
+    plg = cubeviz_helper.plugins['Subset Tools']._obj
+
+    plg.import_region(spec_region)
+    subsets = cubeviz_helper.app.get_subsets()
+    assert len(subsets) == len_subsets
+    assert len(subsets['Subset 1']) == len_subregions