Merge pull request #482 from mwcraig/better-transit-fit

Better transit fitting notebooks

stellarphot/gui_tools/photometry_widget_functions.py

@@ -2,17 +2,137 @@
 
 import ipywidgets as ipw
 from ccdproc import ImageFileCollection
+from ipyautoui.custom import FileChooser
 
+from stellarphot import PhotometryData
 from stellarphot.settings import (
     PhotometryApertures,
     PhotometryFileSettings,
     ui_generator,
 )
+from stellarphot.settings.custom_widgets import Spinner
 
-__all__ = ["PhotometrySettings"]
+__all__ = ["TessAnalysisInputControls", "PhotometrySettingsOLDBAD"]
 
 
-class PhotometrySettings:
+class TessAnalysisInputControls(ipw.VBox):
+    """
+    A class to hold the widgets for choosing TESS input
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        hidden = ipw.Layout(display="none")
+
+        self.phot_chooser = FileChooser(filter_pattern=["*.csv", "*.fits", "*.ecsv"])
+        self._fits_openr = ipw.VBox(
+            children=[
+                ipw.HTML(value="<h3>Select your photometry/flux file</h3>"),
+                self.phot_chooser,
+            ]
+        )
+        self.tic_file_chooser = FileChooser(filter_pattern=["*.json"])
+        fits_openr2 = ipw.VBox(
+            children=[
+                ipw.HTML(value="<h3>Select your TESS info file</h3>"),
+                self.tic_file_chooser,
+            ],
+            layout=hidden,
+        )
+        self._passband = ipw.Dropdown(
+            description="Ccoose Filter",
+            options=["gp", "ip"],
+            disabled=True,
+            layout=hidden,
+        )
+
+        spinner = Spinner(message="<h4>Loading photometry...</h4>")
+
+        self.phot_data = None
+
+        def update_filter_list(_):
+            spinner.start()
+            self.phot_data = PhotometryData.read(self.phot_chooser.value)
+            passband_data = self.phot_data["passband"]
+            fits_openr2.layout.display = "flex"
+            self._passband.layout.display = "flex"
+            self._passband.options = sorted(set(passband_data))
+            self._passband.disabled = False
+            self._passband.value = self._passband.options[0]
+            spinner.stop()
+
+        self.phot_chooser.observe(update_filter_list, names="_value")
+        self.children = [self._fits_openr, spinner, fits_openr2, self._passband]
+
+    @property
+    def tic_info_file(self):
+        p = Path(self.tic_file_chooser.value)
+        selected_file = p.name
+        if not selected_file:
+            raise ValueError("No TIC info json file selected")
+        return p
+
+    @property
+    def photometry_data_file(self):
+        p = Path(self.phot_chooser.value)
+        selected_file = p.name
+        if not selected_file:
+            raise ValueError("No photometry data file selected")
+        return p
+
+    @property
+    def passband(self):
+        return self._passband.value
+
+
+def filter_by_dates(
+    phot_times=None,
+    use_no_data_before=None,
+    use_no_data_between=None,
+    use_no_data_after=None,
+):
+    n_dropped = 0
+
+    bad_data = phot_times < use_no_data_before
+
+    n_dropped = bad_data.sum()
+
+    if n_dropped > 0:
+        print(
+            f"👉👉👉👉 Dropping {n_dropped} data points before "
+            f"BJD {use_no_data_before}"
+        )
+
+    bad_data = bad_data | (
+        (use_no_data_between[0][0] < phot_times)
+        & (phot_times < use_no_data_between[0][1])
+    )
+
+    new_dropped = bad_data.sum() - n_dropped
+
+    if new_dropped:
+        print(
+            f"👉👉👉👉 Dropping {new_dropped} data points between "
+            f"BJD {use_no_data_between[0][0]} and {use_no_data_between[0][1]}"
+        )
+
+    n_dropped += new_dropped
+
+    bad_data = bad_data | (phot_times > use_no_data_after)
+
+    new_dropped = bad_data.sum() - n_dropped
+
+    if new_dropped:
+        print(
+            f"👉👉👉👉 Dropping {new_dropped} data points after "
+            f"BJD {use_no_data_after}"
+        )
+
+    n_dropped += new_dropped
+    return bad_data
+
+
+class PhotometrySettingsOLDBAD:
     """
     A class to hold the widgets for photometry settings.
 

stellarphot/io/tess.py

@@ -1,4 +1,5 @@
 import re
+import warnings
 from dataclasses import dataclass
 from pathlib import Path
 from tempfile import NamedTemporaryFile
@@ -356,6 +357,41 @@ def from_tic_id(cls, tic_id):
             tess_mag_error=toi_table["TESS Mag err"],
         )
 
+    def transit_time_for_observation(self, obs_times):
+        """
+        Calculate the transit time for a set of observation times.
+
+        Parameters
+        ----------
+
+        obs_times : `astropy.time.Time`
+            The times of the observations.
+
+        Returns
+        -------
+
+        `astropy.time.Time`
+            The transit times for the observations.
+        """
+        first_obs = obs_times[0]
+        # Three possible cases here. Either the first time is close to, but before, a
+        # transit, or it is close to, but just after a transit, or it is nowhere close
+        # to a transit.
+        # Assume that the first time is just before a transit
+        cycle_number = int((first_obs - self.epoch) / self.period + 1)
+        that_transit = cycle_number * self.period + self.epoch
+
+        # Check -- is the first time closer to this transit or the one before it?
+        previous_transit = that_transit - self.period
+        if abs(first_obs - previous_transit) < abs(first_obs - that_transit):
+            that_transit = previous_transit
+
+        # Check -- are we way, way, way off from a transit?
+        if abs(first_obs - that_transit) > 3 * self.duration:
+            warnings.warn("Observation times are far from a transit.", stacklevel=2)
+
+        return that_transit
+
 
 def tess_photometry_setup(tic_id=None, TOI_object=None, overwrite=False):
     """

stellarphot/io/tests/test_tess.py

@@ -3,8 +3,10 @@
 import warnings
 from pathlib import Path
 
+import numpy as np
 import pytest
 from astropy.coordinates import SkyCoord
+from astropy.time import Time
 from requests import ConnectionError, ReadTimeout
 
 from stellarphot.io.tess import (
@@ -187,6 +189,25 @@ def test_from_tic_id(self, tess_tic_expected_values):
 
         assert toi_info.coord.separation(new_toi.coord).arcsecond < 0.01
 
+    @pytest.mark.parametrize("start_before_midpoint", [True, False])
+    def test_transit_time_for_observation(self, sample_toi, start_before_midpoint):
+        # For this test we are checking that the correct transit time is identified
+        # for a given observation time.
+        # For the sake of the test, we are going to use the 124th transit of the TOI
+        # as the reference transit.
+        reference_midpoint = sample_toi.epoch + 124 * sample_toi.period
+        if start_before_midpoint:
+            # Start the observation before the midpoint (and before the transit)
+            test_time_start = Time(reference_midpoint - 0.8 * sample_toi.duration)
+        else:
+            # Start the observation after the midpoint
+            test_time_start = Time(reference_midpoint + 0.1 * sample_toi.duration)
+        obs_times = test_time_start + np.linspace(0, 2) * sample_toi.duration
+
+        assert sample_toi.transit_time_for_observation(obs_times).jd == pytest.approx(
+            reference_midpoint.jd
+        )
+
 
 class TestTessPhotometrySetup:
     # This auto-used fixture changes the working directory to the temporary directory