normalize poscorr/centr and fixing pytest to work with new perturbation API

Author: jorgemarpa

src/psfmachine/machine.py

@@ -4,18 +4,15 @@
 import numpy as np
 import pandas as pd
 from scipy import sparse
-from scipy.ndimage import gaussian_filter1d
 import astropy.units as u
 from tqdm import tqdm
 import matplotlib.pyplot as plt
 from astropy.stats import sigma_clip
 
 from .utils import (
     _make_A_polar,
-    _make_A_cartesian,
     solve_linear_model,
     sparse_lessthan,
-    _combine_A,
     threshold_bin,
     get_breaks,
     smooth_vector,
@@ -638,14 +635,14 @@ def build_time_model(
         ----------
         plot: boolean
             Plot a diagnostic figure.
-        bin_method: boolean
-            If True the `time` and `pos_corr` arrays will be downsampled instead of
-            binned.
-        split_time_segments : boolean, list/array of ints
+        bin_method: string
+            Type of bin method, options are "bin" and "downsample".
+        split_time_segments : boolean
             If `True` will split the light curve into segments to fit different time
-            models with a commong pixel normalization. If `False` will fit the full
-            time series as one segment. If list or numpy array of ints, will use the
-            index values as segment breaks.
+            models with a common pixel normalization. If `False` will fit the full
+            time series as one segment.
+        focus_component: boolean
+            Add a component that models th focus change at the begining of a segment.
         """
         # create the time and space basis
         _whitened_time = (self.time - self.time.mean()) / (
@@ -677,11 +674,18 @@ def build_time_model(
                 mpc1 = self.ra_centroid.to("arcsec").value / 4
                 mpc2 = self.dec_centroid.to("arcsec").value / 4
 
-            # we smooth the vectors
+            # smooth the vectors
             mpc1_smooth, mpc2_smooth = smooth_vector(
                 [mpc1, mpc2], splits=splits, filter_size=self.poscorr_filter_size
             )
-            # we combine them as the first order
+            # normalize components
+            mpc1_smooth = (mpc1_smooth - mpc1_smooth.mean()) / (
+                mpc1_smooth.max() - mpc1_smooth.mean()
+            )
+            mpc2_smooth = (mpc2_smooth - mpc2_smooth.mean()) / (
+                mpc2_smooth.max() - mpc2_smooth.mean()
+            )
+            # combine them as the first order
             other_vectors = [mpc1_smooth, mpc2_smooth, mpc1_smooth * mpc2_smooth]
 
         # create a 3D perturbation matrix

tests/test_machine.py

@@ -120,19 +120,13 @@ def test_compute_aperture_photometry():
 @pytest.mark.remote_data
 def test_poscorr_smooth():
     machine = TPFMachine.from_TPFs(tpfs, apply_focus_mask=False)
-    machine._get_source_mask()
-    machine.poscorr_filter_size = 0.4
-    (
-        time_original,
-        time_binned,
-        flux_binned_raw,
-        flux_binned,
-        flux_err_binned,
-        poscorr1_smooth,
-        poscorr2_smooth,
-        poscorr1_binned,
-        poscorr2_binned,
-    ) = machine._time_bin(npoints=3)
+    machine.build_shape_model(plot=False)
+    # no segments
+    machine.time_corrector = "pos_corr"
+    machine.poscorr_filter_size = 1
+    machine.build_time_model(
+        split_time_segments=False, bin_method="bin", focus_component=False
+    )
 
     median_pc1 = np.nanmedian(machine.pos_corr1, axis=0)
     median_pc2 = np.nanmedian(machine.pos_corr2, axis=0)
@@ -143,24 +137,29 @@ def test_poscorr_smooth():
         median_pc2.max() - median_pc2.mean()
     )
 
-    assert np.isclose(poscorr1_smooth, median_pc1, atol=1e-3).all()
-    assert np.isclose(poscorr2_smooth, median_pc2, atol=1e-3).all()
+    assert ((machine.P3.other_vectors[:, 0] - median_pc1) ** 2).sum() ** 0.5 < 1
+    assert ((machine.P3.other_vectors[:, 1] - median_pc2) ** 2).sum() ** 0.5 < 1
 
 
 @pytest.mark.remote_data
 def test_segment_time_model():
     # testing segment with the current test dataset we have that only has 10 cadences
     # isn't the best, but we can still do some sanity checks.
-    machine = TPFMachine.from_TPFs(tpfs, apply_focus_mask=False, n_time_points=3)
+    machine = TPFMachine.from_TPFs(
+        tpfs, apply_focus_mask=False, n_time_points=3, time_corrector="polynomial"
+    )
     machine.build_shape_model(plot=False)
     # no segments
-    machine.build_time_model(split_time_model=False, downsample=True)
-    assert machine.time_model_w.shape[0] == machine.seg_splits.shape[0] - 1
-    assert machine.time_model_w.shape[0] == machine._time_masked.shape[0]
+    machine.build_time_model(
+        split_time_segments=False, bin_method="bin", focus_component=False
+    )
+    assert machine.P3.vectors.shape == (10, 4)
 
-    # need tighter knot spacing
-    machine.n_time_points = 2
+    # fake 2 time breaks
+    machine.time[4:] += 0.5
+    machine.time[7:] += 0.41
     # user defined segments
-    machine.build_time_model(split_time_model=[5], downsample=True)
-    assert machine.time_model_w.shape[0] == machine.seg_splits.shape[0] - 1
-    assert machine.time_model_w.shape[0] == machine._time_masked.shape[0]
+    machine.build_time_model(
+        split_time_segments=True, bin_method="bin", focus_component=False
+    )
+    assert machine.P3.vectors.shape == (10, 4 * 3)