Scramble by cadence rather than time.

PiperOrigin-RevId: 208863882

research/astronet/light_curve_util/kepler_io.py

@@ -213,11 +213,6 @@ def read_kepler_light_curve(filenames,
       # Index into primary HDU header and get quarter.
       quarter = hdu_list[0].header["QUARTER"]
 
-      # Remove NaN flux values.
-      valid_indices = np.where(np.isfinite(flux))
-      time = time[valid_indices]
-      flux = flux[valid_indices]
-
       if time.size:
         all_time.append(time)
         all_flux.append(flux)
@@ -227,4 +222,11 @@ def read_kepler_light_curve(filenames,
     all_time, all_flux = scramble_light_curve(all_time, all_flux, all_quarters,
                                               scramble_type)
 
+  # Remove timestamps with NaN time or flux values.
+  for i, (time, flux) in enumerate(zip(all_time, all_flux)):
+    flux_and_time_finite = np.logical_and(np.isfinite(flux), np.isfinite(time))
+    valid_indices = np.where(flux_and_time_finite)
+    all_time[i] = time[valid_indices]
+    all_flux[i] = flux[valid_indices]
+
   return all_time, all_flux

research/astronet/light_curve_util/kepler_io_test.py

@@ -155,10 +155,12 @@ class KeplerIoTest(absltest.TestCase):
         filenames, scramble_type="SCR1")
     self.assertLen(all_time, 3)
     self.assertLen(all_flux, 3)
-    self.assertLen(all_time[0], 4486)
-    self.assertLen(all_flux[0], 4486)
-    self.assertLen(all_time[1], 4134)
-    self.assertLen(all_flux[1], 4134)
+
+    # Arrays are shorter than above due to separation of time and flux NaNs.
+    self.assertLen(all_time[0], 4344)
+    self.assertLen(all_flux[0], 4344)
+    self.assertLen(all_time[1], 4041)
+    self.assertLen(all_flux[1], 4041)
     self.assertLen(all_time[2], 1008)
     self.assertLen(all_flux[2], 1008)