Feature: TPF support (#82)

* Adding TPF translator and parser (#75) * TPF viewer (#81) * make use of upstream refactor to override indices in lcviz (#83) * fix creating phase-viewer when TPF is loaded (#86) * Time Selector (adapted version of cubeviz's slice) plugin (#85) * enable clone viewer for image/TPF viewer (#101) --------- Co-authored-by: Brett M. Morris <[email protected]>
spacetelescope · Apr 5, 2024 · 6982333 · 6982333
1 parent 4979a8b
commit 6982333
Show file tree

Hide file tree

Showing 17 changed files with 700 additions and 65 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -1,6 +1,8 @@
 0.4.0 - unreleased
 ------------------
 
+* Support loading, viewing, and slicing through TPF data cubes. [#82]
+
 0.3.0 - (04-05-2024)
 --------------------
 

diff --git a/docs/plugins.rst b/docs/plugins.rst
@@ -179,6 +179,40 @@ visible when the plugin is opened.
         Jdaviz documentation on the Markers plugin.
 
 
+.. _time-selector:
+
+Time Selector
+==============
+
+The time selector plugin allows defining the time indicated in all light curve viewers
+(time and phase viewers) as well as the time at which all image cubes are displayed.
+
+
+.. admonition:: User API Example
+    :class: dropdown
+
+    See the :class:`~lcviz.plugins.time_selector.time_selector.TimeSelector` user API documentation for more details.
+
+    .. code-block:: python
+
+      from lcviz import LCviz
+      lc = search_lightcurve("HAT-P-11", mission="Kepler",
+                             cadence="long", quarter=10).download().flatten()
+      lcviz = LCviz()
+      lcviz.load_data(lc)
+      lcviz.show()
+
+      ts = lcviz.plugins['Time Selector']
+      ts.open_in_tray()
+
+
+.. seealso::
+
+    :ref:`Jdaviz Slice Plugin <jdaviz:slice>`
+        Jdaviz documentation on the Slice plugin.
+
+
+
 .. _flatten:
 
 Flatten

diff --git a/docs/reference/api_plugins.rst b/docs/reference/api_plugins.rst
@@ -29,3 +29,6 @@ Plugins API
 
 .. automodapi:: lcviz.plugins.subset_plugin.subset_plugin
   :no-inheritance-diagram:
+
+.. automodapi:: lcviz.plugins.time_selector.time_selector
+  :no-inheritance-diagram:
diff --git a/lcviz/helper.py b/lcviz/helper.py
@@ -69,7 +69,8 @@ class LCviz(ConfigHelper):
         'toolbar': ['g-data-tools', 'g-subset-tools', 'g-viewer-creator', 'lcviz-coords-info'],
         'tray': ['lcviz-metadata-viewer', 'flux-column',
                  'lcviz-plot-options', 'lcviz-subset-plugin',
-                 'lcviz-markers', 'flatten', 'frequency-analysis', 'ephemeris',
+                 'lcviz-markers', 'time-selector',
+                 'flatten', 'frequency-analysis', 'ephemeris',
                  'binning', 'lcviz-export'],
         'viewer_area': [{'container': 'col',
                          'children': [{'container': 'row',
@@ -150,6 +151,13 @@ def default_time_viewer(self):
             raise ValueError("no time viewers exist")
         return tvs[0].user_api
 
+    @property
+    def _has_cube_data(self):
+        for data in self.app.data_collection:
+            if data.ndim == 3:
+                return True
+        return False
+
     @property
     def _tray_tools(self):
         """

diff --git a/lcviz/marks.py b/lcviz/marks.py
@@ -1,11 +1,21 @@
+from astropy import units as u
 import numpy as np
 
-from jdaviz.core.marks import PluginLine, PluginScatter
+from jdaviz.core.marks import PluginLine, PluginScatter, SliceIndicatorMarks
 from lcviz.viewers import PhaseScatterView
 
 __all__ = ['LivePreviewTrend', 'LivePreviewFlattened', 'LivePreviewBinning']
 
 
+def _slice_indicator_get_slice_axis(self, data):
+    if hasattr(data, 'time'):
+        return data.time.value * u.d
+    return [] * u.dimensionless_unscaled
+
+
+SliceIndicatorMarks._get_slice_axis = _slice_indicator_get_slice_axis
+
+
 class WithoutPhaseSupport:
     def update_ty(self, times, y):
         self.times = np.asarray(times)

diff --git a/lcviz/parsers.py b/lcviz/parsers.py
@@ -10,6 +10,13 @@
 
 @data_parser_registry("light_curve_parser")
 def light_curve_parser(app, file_obj, data_label=None, show_in_viewer=True, **kwargs):
+    # load a LightCurve or TargetPixelFile object:
+    cls_with_translator = (
+        lightkurve.LightCurve,
+        lightkurve.targetpixelfile.KeplerTargetPixelFile,
+        lightkurve.targetpixelfile.TessTargetPixelFile
+    )
+
     # load local FITS file from disk by its path:
     if isinstance(file_obj, str) and os.path.exists(file_obj):
         if data_label is None:
@@ -18,8 +25,7 @@ def light_curve_parser(app, file_obj, data_label=None, show_in_viewer=True, **kw
         # read the light curve:
         light_curve = lightkurve.read(file_obj)
 
-    # load a LightCurve object:
-    elif isinstance(file_obj, lightkurve.LightCurve):
+    elif isinstance(file_obj, cls_with_translator):
         light_curve = file_obj
 
     # make a data label:
@@ -30,7 +36,12 @@ def light_curve_parser(app, file_obj, data_label=None, show_in_viewer=True, **kw
 
     # handle flux_origin default
     flux_origin = light_curve.meta.get('FLUX_ORIGIN', None)  # i.e. PDCSAP or SAP
-    if flux_origin == 'flux' or (flux_origin is None and 'flux' in light_curve.columns):
+    if isinstance(light_curve, lightkurve.targetpixelfile.TargetPixelFile):
+        new_data_label += '[TPF]'
+    elif flux_origin is not None:
+        new_data_label += f'[{flux_origin}]'
+
+    if flux_origin == 'flux' or (flux_origin is None and 'flux' in getattr(light_curve, 'columns', [])):  # noqa
         # then make a copy of this column so it won't be lost when changing with the flux_column
         # plugin
         light_curve['flux:orig'] = light_curve['flux']
@@ -41,13 +52,33 @@ def light_curve_parser(app, file_obj, data_label=None, show_in_viewer=True, **kw
     data = _data_with_reftime(app, light_curve)
     app.add_data(data, new_data_label)
 
-    if show_in_viewer:
-        # add to any known time/phase viewers
-        for viewer_id, viewer in app._viewer_store.items():
-            if isinstance(viewer, TimeScatterView):
-                app.add_data_to_viewer(viewer_id, new_data_label)
-            elif isinstance(viewer, PhaseScatterView):
-                app.add_data_to_viewer(viewer_id, new_data_label)
+    if isinstance(light_curve, lightkurve.targetpixelfile.TargetPixelFile):
+        # ensure an image/cube/TPF viewer exists
+        # TODO: move this to an event listener on add_data so that we can also remove when empty?
+        from jdaviz.core.events import NewViewerMessage
+        from lcviz.viewers import CubeView
+        if show_in_viewer:
+            found_viewer = False
+            for viewer_id, viewer in app._viewer_store.items():
+                if isinstance(viewer, CubeView):
+                    app.add_data_to_viewer(viewer_id, new_data_label)
+                    found_viewer = True
+            if not found_viewer:
+                app._on_new_viewer(NewViewerMessage(CubeView, data=None, sender=app),
+                                   vid='image', name='image')
+                app.add_data_to_viewer('image', new_data_label)
+
+    else:
+        if show_in_viewer:
+            for viewer_id, viewer in app._viewer_store.items():
+                if isinstance(viewer, (TimeScatterView, PhaseScatterView)):
+                    app.add_data_to_viewer(viewer_id, new_data_label)
+
+            # add to any known phase viewers
+            ephem_plugin = app._jdaviz_helper.plugins.get('Ephemeris', None)
+            if ephem_plugin is not None:
+                for viewer in ephem_plugin._obj._get_phase_viewers():
+                    app.add_data_to_viewer(viewer.reference, new_data_label)
 
 
 def _data_with_reftime(app, light_curve):

diff --git a/lcviz/plugins/__init__.py b/lcviz/plugins/__init__.py
@@ -8,6 +8,7 @@
 from .flux_column.flux_column import *  # noqa
 from .frequency_analysis.frequency_analysis import *  # noqa
 from .markers.markers import *  # noqa
+from .time_selector.time_selector import *  # noqa
 from .metadata_viewer.metadata_viewer import *  # noqa
 from .plot_options.plot_options import *  # noqa
 from .subset_plugin.subset_plugin import *  # noqa
diff --git a/lcviz/plugins/coords_info/coords_info.py b/lcviz/plugins/coords_info/coords_info.py
@@ -5,14 +5,14 @@
 from jdaviz.core.events import ViewerRenamedMessage
 from jdaviz.core.registries import tool_registry
 
-from lcviz.viewers import TimeScatterView, PhaseScatterView
+from lcviz.viewers import TimeScatterView, PhaseScatterView, CubeView
 
 __all__ = ['CoordsInfo']
 
 
 @tool_registry('lcviz-coords-info')
 class CoordsInfo(CoordsInfo):
-    _supported_viewer_classes = (TimeScatterView, PhaseScatterView)
+    _supported_viewer_classes = (TimeScatterView, PhaseScatterView, CubeView)
     _viewer_classes_with_marker = (TimeScatterView, PhaseScatterView)
 
     def __init__(self, *args, **kwargs):
@@ -25,12 +25,19 @@ def __init__(self, *args, **kwargs):
     def _viewer_renamed(self, msg):
         self._marks[msg.new_viewer_ref] = self._marks.pop(msg.old_viewer_ref)
 
-    def update_display(self, viewer, x, y):
-        self._dict = {}
+    def _image_shape_inds(self, image):
+        if image.ndim == 3:
+            # exception to the upstream cubeviz case of (0, 1)
+            return (2, 1)
+        return super()._image_shape_inds(image)
 
-        if not len(viewer.state.layers):
-            return
+    def _get_cube_value(self, image, arr, x, y, viewer):
+        if image.ndim == 3:
+            # exception to the upstream cubeviz case of x, y, slice
+            return arr[viewer.state.slices[0], int(round(y)), int(round(x))]
+        return super()._get_cube_value(image, arr, x, y, viewer)
 
+    def _lc_viewer_update(self, viewer, x, y):
         is_phase = isinstance(viewer, PhaseScatterView)
         # TODO: update with display_unit when supported in lcviz
         x_unit = '' if is_phase else str(viewer.time_unit)
@@ -138,3 +145,14 @@ def _cursor_fallback():
 
         self.marks[viewer._reference_id].update_xy([closest_x], [closest_y])  # noqa
         self.marks[viewer._reference_id].visible = True
+
+    def update_display(self, viewer, x, y):
+        self._dict = {}
+
+        if not len(viewer.state.layers):
+            return
+
+        if isinstance(viewer, (TimeScatterView, PhaseScatterView)):
+            self._lc_viewer_update(viewer, x, y)
+        elif isinstance(viewer, CubeView):
+            self._image_viewer_update(viewer, x, y)
diff --git a/lcviz/plugins/ephemeris/ephemeris.py b/lcviz/plugins/ephemeris/ephemeris.py
@@ -307,6 +307,9 @@ def create_phase_viewer(self, ephem_component=None):
         # set default data visibility
         time_viewer_item = self.app._get_viewer_item(self.app._jdaviz_helper.default_time_viewer._obj.reference)  # noqa
         for data in dc:
+            if data.ndim > 1:
+                # skip image/cube entries
+                continue
             data_id = self.app._data_id_from_label(data.label)
             visible = time_viewer_item['selected_data_items'].get(data_id, 'hidden')
             self.app.set_data_visibility(phase_viewer_id, data.label, visible == 'visible')

diff --git a/lcviz/plugins/time_selector/__init__.py b/lcviz/plugins/time_selector/__init__.py
@@ -0,0 +1 @@
+from .time_selector import *  # noqa
diff --git a/lcviz/plugins/time_selector/time_selector.py b/lcviz/plugins/time_selector/time_selector.py
@@ -0,0 +1,82 @@
+from jdaviz.configs.cubeviz.plugins import Slice
+from jdaviz.core.registries import tray_registry
+
+from lcviz.events import EphemerisChangedMessage
+from lcviz.viewers import CubeView, PhaseScatterView
+
+__all__ = ['TimeSelector']
+
+
+@tray_registry('time-selector', label="Time Selector")
+class TimeSelector(Slice):
+    """
+    See the :ref:`Time Selector Plugin Documentation <time-selector>` for more details.
+
+    Only the following attributes and methods are available through the
+    :ref:`public plugin API <plugin-apis>`:
+
+    * :meth:`~jdaviz.core.template_mixin.PluginTemplateMixin.show`
+    * :meth:`~jdaviz.core.template_mixin.PluginTemplateMixin.open_in_tray`
+    * :meth:`~jdaviz.core.template_mixin.PluginTemplateMixin.close_in_tray`
+    * ``value``  Time of the indicator.  When setting this directly, it will
+      update automatically to the value corresponding to the nearest slice, if ``snap_to_slice`` is
+      enabled and a cube is loaded.
+    * ``show_indicator``
+      Whether to show indicator in spectral viewer when slice tool is inactive.
+    * ``show_value``
+      Whether to show slice value in label to right of indicator.
+    * ``snap_to_slice``
+      Whether the indicator (and ``value``) should snap to the value of the nearest slice in the
+      cube (if one exists).
+    """
+    _cube_viewer_cls = CubeView
+    _cube_viewer_default_label = 'image'
+
+    def __init__(self, *args, **kwargs):
+        """
+
+        """
+        super().__init__(*args, **kwargs)
+        self.docs_link = f"https://lcviz.readthedocs.io/en/{self.vdocs}/plugins.html#time-selector"
+        self.docs_description = "Select time to sync across all viewers (as an indicator in all time/phase viewers or to select the active slice in any image/cube viewers).  The slice can also be changed interactively in any time viewer by activating the slice tool."  # noqa
+        self.value_label = 'Time'
+        self.value_unit = 'd'
+        self.allow_disable_snapping = True
+
+        self.session.hub.subscribe(self, EphemerisChangedMessage,
+                                   handler=self._on_ephemeris_changed)
+
+    @property
+    def slice_axis(self):
+        # global display unit "axis" corresponding to the slice axis
+        return 'time'
+
+    @property
+    def valid_slice_att_names(self):
+        return ["time", "dt"]
+
+    @property
+    def user_api(self):
+        api = super().user_api
+        # can be removed after deprecated upstream attributes for wavelength/wavelength_value
+        # are removed in the lowest supported version of jdaviz
+        api._expose = [e for e in api._expose if e not in ('slice', 'wavelength',
+                                                           'wavelength_value', 'show_wavelength')]
+        return api
+
+    def _on_select_slice_message(self, msg):
+        viewer = msg.sender.viewer
+        if isinstance(viewer, PhaseScatterView):
+            prev_phase = viewer.times_to_phases(self.value)
+            new_phase = msg.value
+            self.value = self.value + (new_phase - prev_phase) * viewer.ephemeris.get('period', 1.0)
+        else:
+            super()._on_select_slice_message(msg)
+
+    def _on_ephemeris_changed(self, msg):
+        for viewer in self.slice_indicator_viewers:
+            if not isinstance(viewer, PhaseScatterView):
+                continue
+            if viewer._ephemeris_component != msg.ephemeris_label:
+                continue
+            viewer._set_slice_indicator_value(self.value)
diff --git a/lcviz/plugins/viewer_creator/viewer_creator.py b/lcviz/plugins/viewer_creator/viewer_creator.py
@@ -1,8 +1,10 @@
+from glue.core.message import (DataCollectionAddMessage,
+                               DataCollectionDeleteMessage)
 from jdaviz.configs.default.plugins import ViewerCreator
 from jdaviz.core.events import NewViewerMessage
 from jdaviz.core.registries import tool_registry
 from lcviz.events import EphemerisComponentChangedMessage
-from lcviz.viewers import TimeScatterView
+from lcviz.viewers import TimeScatterView, CubeView
 
 __all__ = ['ViewerCreator']
 
@@ -13,8 +15,11 @@ class ViewerCreator(ViewerCreator):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        self.hub.subscribe(self, EphemerisComponentChangedMessage,
-                           handler=self._rebuild_available_viewers)
+        for msg in (EphemerisComponentChangedMessage,
+                    DataCollectionAddMessage,
+                    DataCollectionDeleteMessage):
+            self.hub.subscribe(self, msg,
+                               handler=lambda x: self._rebuild_available_viewers())
         self._rebuild_available_viewers()
 
     def _rebuild_available_viewers(self, *args):
@@ -25,12 +30,18 @@ def _rebuild_available_viewers(self, *args):
         if self.app._jdaviz_helper is not None:
             phase_viewers = [{'name': f'lcviz-phase-viewer:{e}', 'label': f'flux-vs-phase:{e}'}
                               for e in self.app._jdaviz_helper.plugins['Ephemeris'].component.choices]  # noqa
+            if self.app._jdaviz_helper._has_cube_data:
+                cube_viewers = [{'name': 'lcviz-cube-viewer', 'label': 'image'}]
+            else:
+                cube_viewers = []
         else:
             phase_viewers = [{'name': 'lcviz-phase-viewer:default',
                               'label': 'flux-vs-phase:default'}]
+            cube_viewers = []
 
         self.viewer_types = [v for v in self.viewer_types if v['name'].startswith('lcviz')
-                             and not v['label'].startswith('flux-vs-phase')] + phase_viewers
+                             and not v['label'].startswith('flux-vs-phase')
+                             and not v['label'] in ('cube', 'image')] + phase_viewers + cube_viewers
         self.send_state('viewer_types')
 
     def vue_create_viewer(self, name):
@@ -45,5 +56,10 @@ def vue_create_viewer(self, name):
             self.app._on_new_viewer(NewViewerMessage(TimeScatterView, data=None, sender=self.app),
                                     vid=viewer_id, name=viewer_id)
             return
+        if name in ('image', 'lcviz-cube-viewer'):
+            viewer_id = self.app._jdaviz_helper._get_clone_viewer_reference('image')
+            self.app._on_new_viewer(NewViewerMessage(CubeView, data=None, sender=self.app),
+                                    vid=viewer_id, name=viewer_id)
+            return
 
         super().vue_create_viewer(name)