Flux Column plugin

CHANGES.rst

@@ -3,6 +3,11 @@
 
 * Clone viewer tool. [#74]
 
+* Flux column plugin to choose which column is treated as the flux column for each dataset. [#77]
+
+* Flatten plugin no longer creates new data entries, but instead appends a new column to the input
+  light curve and selects as the flux column (origin). [#77]
+
 0.1.0 (12-14-2023)
 ------------------
 

docs/plugins.rst

@@ -35,6 +35,41 @@ This plugin allows viewing of any metadata associated with the selected data.
     :ref:`Jdaviz Metadata Viewer <jdaviz:imviz_metadata-viewer>`
         Jdaviz documentation on the Metadata Viewer plugin.
 
+.. _flux-column:
+
+Flux Column
+===========
+
+This plugin allows choosing which column in the underlying data should be used as the flux column
+(origin) throughout the app (when plotting and in any data analysis plugins).
+
+
+.. admonition:: User API Example
+    :class: dropdown
+
+    See the :class:`~lcviz.plugins.plot_options.plot_options.PlotOptions` 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()
+
+      flux_col = lcviz.plugins['Flux Column']
+      print(flux_col.flux_column.choices)
+      flux_col.flux_column = 'sap_flux'
+
+
+.. seealso::
+
+    This plugin reproduces the behavior also available in ``lightkurve`` as:
+
+    * :meth:`lightkurve.LightCurve.select_flux`
+
+
 .. _plot-options:
 
 Plot Options

lcviz/components/__init__.py

@@ -0,0 +1 @@
+from .components import *  # noqa

lcviz/components/components.py

@@ -0,0 +1,117 @@
+from astropy import units as u
+from ipyvuetify import VuetifyTemplate
+from glue.core import HubListener
+from traitlets import List, Unicode
+
+from jdaviz.core.template_mixin import SelectPluginComponent
+
+from lcviz.events import FluxColumnChangedMessage
+
+__all__ = ['FluxColumnSelect', 'FluxColumnSelectMixin']
+
+
+class FluxColumnSelect(SelectPluginComponent):
+    def __init__(self, plugin, items, selected, dataset):
+        super().__init__(plugin,
+                         items=items,
+                         selected=selected,
+                         dataset=dataset)
+
+        self.add_observe(selected, self._on_change_selected)
+        self.add_observe(self.dataset._plugin_traitlets['selected'],
+                         self._on_change_dataset)
+
+        # sync between instances in different plugins
+        self.hub.subscribe(self, FluxColumnChangedMessage,
+                           handler=self._on_flux_column_changed_msg)
+
+    def _on_change_dataset(self, *args):
+        def _include_col(lk_obj, col):
+            if col == 'flux' and lk_obj.meta.get('FLUX_ORIGIN') != 'flux':
+                # this is the currently active column (and should be copied elsewhere unless)
+                return False
+            if col in ('time', 'cadn', 'cadenceno', 'quality'):
+                return False
+            if col.startswith('phase:'):
+                # internal jdaviz ephemeris phase columns
+                return False
+            if col.startswith('time'):
+                return False
+            if col.startswith('centroid'):
+                return False
+            if col.startswith('cbv'):
+                # cotrending basis vector
+                return False
+            if col.endswith('_err'):
+                return False
+            if col.endswith('quality'):
+                return False
+            # TODO: need to think about flatten losing units in the flux column
+            return lk_obj[col].unit != u.pix
+
+        lk_obj = self.dataset.selected_obj
+        if lk_obj is None:
+            return
+        self.choices = [col for col in lk_obj.columns if _include_col(lk_obj, col)]
+        flux_column = lk_obj.meta.get('FLUX_ORIGIN')
+        if flux_column in self.choices:
+            self.selected = flux_column
+        else:
+            self.selected = ''
+
+    def _on_flux_column_changed_msg(self, msg):
+        if msg.dataset != self.dataset.selected:
+            return
+
+        # need to clear the cache due to the change in metadata made to the data-collection entry
+        self.dataset._clear_cache('selected_obj', 'selected_dc_item')
+        self._on_change_dataset()
+        self.selected = msg.flux_column
+
+    def _on_change_selected(self, *args):
+        if self.selected == '':
+            return
+
+        dc_item = self.dataset.selected_dc_item
+        old_flux_column = dc_item.meta.get('FLUX_ORIGIN')
+        if self.selected == old_flux_column:
+            # nothing to do here!
+            return
+
+        # instead of using lightkurve's select_flux and having to reparse the data entry, we'll
+        # manipulate the arrays in the data-collection directly, and modify FLUX_ORIGIN so that
+        # exporting back to a lightkurve object works as expected
+        self.app._jdaviz_helper._set_data_component(dc_item, 'flux', dc_item[self.selected])
+        self.app._jdaviz_helper._set_data_component(dc_item, 'flux_err', dc_item[self.selected+"_err"])  # noqa
+        dc_item.meta['FLUX_ORIGIN'] = self.selected
+
+        self.hub.broadcast(FluxColumnChangedMessage(dataset=self.dataset.selected,
+                                                    flux_column=self.selected,
+                                                    sender=self))
+
+    def add_new_flux_column(self, flux, flux_err, label, selected=False):
+        dc_item = self.dataset.selected_dc_item
+        self.app._jdaviz_helper._set_data_component(dc_item,
+                                                    label,
+                                                    flux)
+        self.app._jdaviz_helper._set_data_component(dc_item,
+                                                    f"{label}_err",
+                                                    flux_err)
+
+        # broadcast so all instances update to get the new column and selection (if applicable)
+        self.hub.broadcast(FluxColumnChangedMessage(dataset=self.dataset.selected,
+                                                    flux_column=label if selected else self.selected,  # noqa
+                                                    sender=self))
+
+
+class FluxColumnSelectMixin(VuetifyTemplate, HubListener):
+    flux_column_items = List().tag(sync=True)
+    flux_column_selected = Unicode().tag(sync=True)
+    # assumes DatasetSelectMixin is also used (DatasetSelectMixin must appear after in inheritance)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.flux_column = FluxColumnSelect(self,
+                                            'flux_column_items',
+                                            'flux_column_selected',
+                                            dataset='dataset')

lcviz/conftest.py

@@ -32,9 +32,12 @@ def light_curve_like_kepler_quarter(seed=42):
 
     quality = np.zeros(len(time), dtype=np.int32)
 
-    return LightCurve(
+    lc = LightCurve(
         time=time, flux=flux, flux_err=flux_err, quality=quality
     )
+    lc['flux_alt'] = flux + 1
+    lc['flux_alt_err'] = flux_err
+    return lc
 
 
 try:

lcviz/events.py

@@ -1,7 +1,8 @@
 from glue.core.message import Message
 
 __all__ = ['EphemerisComponentChangedMessage',
-           'EphemerisChangedMessage']
+           'EphemerisChangedMessage',
+           'FluxColumnChangedMessage']
 
 
 class EphemerisComponentChangedMessage(Message):
@@ -27,3 +28,12 @@ class EphemerisChangedMessage(Message):
     in the ephemeris plugin"""
     def __init__(self, ephemeris_label, *args, **kwargs):
         self.ephemeris_label = ephemeris_label
+
+
+class FluxColumnChangedMessage(Message):
+    """Message emitted by the FluxColumnSelect component when the selection has been changed.
+    To subscribe to a change for a particular dataset, consider using FluxColumnSelect directly
+    and observing the traitlet, rather than subscribing to this message"""
+    def __init__(self, dataset, flux_column, *args, **kwargs):
+        self.dataset = dataset
+        self.flux_column = flux_column

lcviz/helper.py

@@ -72,7 +72,8 @@ class LCviz(ConfigHelper):
                      'dense_toolbar': False,
                      'context': {'notebook': {'max_height': '600px'}}},
         'toolbar': ['g-data-tools', 'g-subset-tools', 'lcviz-coords-info'],
-        'tray': ['lcviz-metadata-viewer', 'lcviz-plot-options', 'lcviz-subset-plugin',
+        'tray': ['lcviz-metadata-viewer', 'flux-column',
+                 'lcviz-plot-options', 'lcviz-subset-plugin',
                  'lcviz-markers', 'flatten', 'frequency-analysis', 'ephemeris',
                  'binning', 'lcviz-export-plot'],
         'viewer_area': [{'container': 'col',
@@ -155,12 +156,17 @@ def _phase_comp_lbl(self, component):
         return f'phase:{component}'
 
     def _set_data_component(self, data, component_label, values):
-        if component_label not in self._component_ids:
-            self._component_ids[component_label] = ComponentID(component_label)
-
-        if self._component_ids[component_label] in data.components:
-            data.update_components({self._component_ids[component_label]: values})
+        if component_label in self._component_ids:
+            component_id = self._component_ids[component_label]
         else:
-            data.add_component(values, self._component_ids[component_label])
-
-        data.add_component(values, self._component_ids[component_label])
+            existing_components = [component.label for component in data.components]
+            if component_label in existing_components:
+                component_id = data.components[existing_components.index(component_label)]
+            else:
+                component_id = ComponentID(component_label)
+                self._component_ids[component_label] = component_id
+
+        if component_id in data.components:
+            data.update_components({component_id: values})
+        else:
+            data.add_component(values, component_id)

lcviz/parsers.py

@@ -27,9 +27,16 @@ def light_curve_parser(app, file_obj, data_label=None, show_in_viewer=True, **kw
         new_data_label = f'{data_label}'
     else:
         new_data_label = light_curve.meta.get('OBJECT', 'Light curve')
+
+    # handle flux_origin default
     flux_origin = light_curve.meta.get('FLUX_ORIGIN', None)  # i.e. PDCSAP or SAP
-    if flux_origin is not None:
-        new_data_label += f'[{flux_origin}]'
+    if flux_origin == 'flux' or (flux_origin is None and 'flux' in light_curve.columns):
+        # 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']
+        if 'flux_err' in light_curve.columns:
+            light_curve['flux:orig_err'] = light_curve['flux_err']
+        light_curve.meta['FLUX_ORIGIN'] = 'flux:orig'
 
     data = _data_with_reftime(app, light_curve)
     app.add_data(data, new_data_label)

lcviz/plugins/__init__.py

@@ -3,6 +3,7 @@
 from .ephemeris.ephemeris import *  # noqa
 from .export_plot.export_plot import *  # noqa
 from .flatten.flatten import *  # noqa
+from .flux_column.flux_column import *  # noqa
 from .frequency_analysis.frequency_analysis import *  # noqa
 from .markers.markers import *  # noqa
 from .metadata_viewer.metadata_viewer import *  # noqa

lcviz/plugins/binning/binning.py

@@ -13,6 +13,7 @@
                                         with_spinner)
 from jdaviz.core.user_api import PluginUserApi
 
+from lcviz.components import FluxColumnSelectMixin
 from lcviz.events import EphemerisChangedMessage
 from lcviz.helper import _default_time_viewer_reference_name
 from lcviz.marks import LivePreviewBinning
@@ -24,7 +25,8 @@
 
 
 @tray_registry('binning', label="Binning")
-class Binning(PluginTemplateMixin, DatasetSelectMixin, EphemerisSelectMixin, AddResultsMixin):
+class Binning(PluginTemplateMixin, FluxColumnSelectMixin, DatasetSelectMixin,
+              EphemerisSelectMixin, AddResultsMixin):
     """
     See the :ref:`Binning Plugin Documentation <binning>` for more details.
 
@@ -150,7 +152,8 @@ def _toggle_marks(self, event={}):
             # then the marks themselves need to be updated
             self._live_update(event)
 
-    @observe('dataset_selected', 'ephemeris_selected',
+    @observe('flux_column_selected', 'dataset_selected',
+             'ephemeris_selected',
              'n_bins', 'previews_temp_disable')
     @skip_if_no_updates_since_last_active()
     def _live_update(self, event={}):