Merge pull request #10 from bmorris3/lightkurve-parser

First light curve parser

lcviz/conftest.py

@@ -4,7 +4,12 @@
 from astropy.utils.masked import Masked
 from lightkurve import LightCurve
 
-from lcviz import __version__
+from lcviz import __version__, LCviz
+
+
+@pytest.fixture
+def helper():
+    return LCviz()
 
 
 @pytest.fixture

lcviz/helper.py

@@ -1,5 +1,4 @@
 from jdaviz.core.helpers import ConfigHelper
-from glue.core import Data
 
 
 class LCviz(ConfigHelper):
@@ -23,21 +22,26 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._default_time_viewer_reference_name = 'time-viewer'
 
-    def load_data(self, flux, time, data_label):
-        '''
-        Loads two quantity arrays by constructing a glue data object
+    def load_data(self, data, data_label=None):
+        """
+        Load data into LCviz.
 
         Parameters
         ----------
-        flux : astropy.units.Quantity
-            An astropy quantity array designating the flux or profile axis
-        time : astropy.units.Quantity
-            An astropy quantity array designating the time axis
-        data_label : str
-            The Glue data label found in the ``DataCollection``.
-        '''
-        data_to_load = Data(x=time.value, flux=flux.value)
-        data_to_load.get_component('x').units = str(time.unit)
-        data_to_load.get_component('flux').units = str(flux.unit)
-        super().load_data(data=data_to_load, parser_reference='lcviz_manual_data_parser',
-                          data_label=data_label)
+        data : obj or str
+            File name or object to be loaded. Supported formats include:
+
+            * ``'filename.fits'`` (or any extension that ``astropy.io.fits``
+              supports)
+            * `~lightkurve.LightCurve` (extracts the default flux column)
+        data_label : str or `None`
+            Data label to go with the given data. If not given, this is
+            automatically determined from filename or randomly generated.
+            The final label shown in LCviz may have additional information
+            appended for clarity.
+        """
+        super().load_data(
+            data=data,
+            parser_reference='light_curve_parser',
+            data_label=data_label
+        )

lcviz/parsers.py

@@ -1,13 +1,48 @@
+import os
+from astropy.io import fits
 from jdaviz.core.registries import data_parser_registry
+from lightkurve import LightCurve, KeplerLightCurve, TessLightCurve
+from lightkurve.io.detect import detect_filetype
 
-__all__ = ["lcviz_manual_data_parser"]
+__all__ = ["light_curve_parser"]
 
 
-@data_parser_registry("lcviz_manual_data_parser")
-def lcviz_manual_data_parser(app, data, data_label=None, show_in_viewer=True):
+@data_parser_registry("light_curve_parser")
+def light_curve_parser(app, file_obj, data_label=None, show_in_viewer=True, **kwargs):
     time_viewer_reference_name = app._jdaviz_helper._default_time_viewer_reference_name
 
-    data._preferred_translation = True  # Triggers custom viewer.set_plot_axes()
+    # 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:
+            data_label = os.path.splitext(os.path.basename(file_obj))[0]
 
-    app.add_data(data, data_label)
-    app.add_data_to_viewer(time_viewer_reference_name, data_label)
+        # detect the type light curve in a FITS file:
+        with fits.open(file_obj) as hdulist:
+            filetype = detect_filetype(hdulist)
+
+        # get the constructor for this type of light curve:
+        filetype_to_cls = {
+            'KeplerLightCurve': KeplerLightCurve,
+            'TessLightCurve': TessLightCurve
+        }
+        cls = filetype_to_cls[filetype]
+        # read the light curve:
+        light_curve = cls.read(file_obj)
+
+    # load a LightCurve object:
+    elif isinstance(file_obj, LightCurve):
+        light_curve = file_obj
+
+    # make a data label:
+    if data_label is not None:
+        new_data_label = f'{data_label}'
+    else:
+        new_data_label = light_curve.meta.get('OBJECT', 'Light curve')
+    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}]'
+
+    app.add_data(light_curve, new_data_label)
+
+    if show_in_viewer:
+        app.add_data_to_viewer(time_viewer_reference_name, new_data_label)

lcviz/tests/test_parser.py

@@ -0,0 +1,65 @@
+# import pytest
+import numpy as np
+from astropy.time import Time
+# from astropy.utils.data import download_file
+from lightkurve import LightCurve
+# from lightkurve.io import kepler
+from gwcs.wcs import WCS
+import astropy.units as u
+
+
+# @pytest.mark.remote_data
+# def test_kepler_via_mast_local_file(helper):
+#     url = (
+#         'https://archive.stsci.edu/pub/kepler/'
+#         'lightcurves/0014/001429092/kplr001429092-2009166043257_llc.fits'
+#     )  # 188 KB
+#
+#     path = download_file(url, cache=True, timeout=100)
+#     helper.load_data(path)
+#
+#     data = helper.app.data_collection[0]
+#     flux_arr = data['flux']
+#     flux_unit = u.Unit(data.get_component('flux').units)
+#     flux = flux_arr * flux_unit
+#
+#     assert isinstance(data.coords, WCS)
+#     assert isinstance(flux, u.Quantity)
+#     assert flux.unit.is_equivalent(u.electron / u.s)
+#
+#
+# @pytest.mark.remote_data
+# def test_kepler_via_mast_preparsed(helper):
+#     url = (
+#         'https://archive.stsci.edu/pub/kepler/'
+#         'lightcurves/0014/001429092/kplr001429092-2009166043257_llc.fits'
+#     )  # 188 KB
+#
+#     light_curve = kepler.read_kepler_lightcurve(url)
+#     helper.load_data(light_curve)
+#
+#     data = helper.app.data_collection[0]
+#     flux_arr = data['flux']
+#     flux_unit = u.Unit(data.get_component('flux').units)
+#     flux = flux_arr * flux_unit
+#
+#     assert isinstance(data.coords, WCS)
+#     assert isinstance(flux, u.Quantity)
+#     assert flux.unit.is_equivalent(u.electron / u.s)
+
+
+def test_synthetic_lc(helper):
+    time = Time(np.linspace(2460050, 2460060), format='jd')
+    flux = np.ones(len(time)) * u.electron / u.s
+    flux_err = 0.1 * np.ones_like(flux)
+    lc = LightCurve(time=time, flux=flux, flux_err=flux_err)
+    helper.load_data(lc)
+
+    data = helper.app.data_collection[0]
+    flux_arr = data['flux']
+    flux_unit = u.Unit(data.get_component('flux').units)
+    flux = flux_arr * flux_unit
+
+    assert isinstance(data.coords, WCS)
+    assert isinstance(flux, u.Quantity)
+    assert flux.unit.is_equivalent(u.electron / u.s)

lcviz/viewers.py

@@ -33,18 +33,39 @@ def set_plot_axes(self):
         # Get data to be used for axes labels
         data = self.data()[0]
 
-        # TBF: Temp comps until actual glue-astronomy translators are developed
-        x_component = 'x'
-        y_component = 'flux'
+        # Get the lookup table from the time axis in the gwcs obj:
+        lookup_table = data.coords._pipeline[0].transform.lookup_table
+        x_unit = lookup_table.unit
+        reference_time = data.meta.get('reference_time', None)
 
-        x_unit = u.Unit(data.get_component(x_component).units)
-        self.figure.axes[0].label = f'{str(x_unit.physical_type).title()} ({x_unit})'
+        if reference_time is not None:
+            xlabel = f'{str(x_unit.physical_type).title()} from {reference_time.iso} ({x_unit})'
+        else:
+            xlabel = f'{str(x_unit.physical_type).title()} ({x_unit})'
 
-        y_unit = u.Unit(data.get_component(y_component).units)
-        self.figure.axes[1].label = f'{str(y_unit.physical_type).title()} ({y_unit})'
+        self.figure.axes[0].label = xlabel
 
-        # Make it so y axis label is not covering tick numbers.
+        y_unit = u.Unit(data.get_component('flux').units)
+        y_unit_physical_type = str(y_unit.physical_type).title()
+
+        common_count_rate_units = (u.electron / u.s, u.dn / u.s, u.ct / u.s)
+
+        if y_unit_physical_type == 'Unknown':
+            if y_unit.is_equivalent(common_count_rate_units):
+                y_unit_physical_type = 'Flux'
+        if y_unit_physical_type == 'Dimensionless':
+            y_unit_physical_type = 'Relative Flux'
+
+        ylabel = f'{y_unit_physical_type}'
+
+        if not y_unit.is_equivalent(u.dimensionless_unscaled):
+            ylabel += f' ({y_unit})'
+
+        self.figure.axes[1].label = ylabel
+
+        # Make it so y axis label is not covering tick numbers (sometimes)
         self.figure.axes[1].label_offset = "-50"
 
-        # Set Y-axis to scientific notation
-        self.figure.axes[1].tick_format = '0.1e'
+        # Set (X,Y)-axis to scientific notation if necessary:
+        self.figure.axes[0].tick_format = 'g'
+        self.figure.axes[1].tick_format = 'g'

setup.cfg

@@ -26,9 +26,9 @@ zip_safe = False
 setup_requires =
     setuptools_scm
 install_requires =
-    astropy>=5
+    astropy>=5.2
     jdaviz
-    lightkurve
+    lightkurve>=2
 python_requires = >=3.8
 
 [options.extras_require]
@@ -49,8 +49,11 @@ lcviz =
 [tool:pytest]
 minversion = 5.0
 norecursedirs = build docs/_build
+testpaths = "lcviz" "docs"
 astropy_header = True
 doctest_plus = enabled
+text_file_format = rst
+addopts = --doctest-rst --import-mode=append
 filterwarnings =
     error
     ignore:numpy\.ndarray size changed:RuntimeWarning