Remove spinup from sample indices (#44)

timothyas · Oct 20, 2023 · d06eff1 · d06eff1
1 parent 5c6a006
commit d06eff1
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/xesn/driver.py b/xesn/driver.py
@@ -243,7 +243,7 @@ def get_samples(self, xda, n_samples, n_steps, n_spinup, random_seed=None, sampl
         else:
             assert len(sample_indices) == n_samples, f"Driver.get_samples: found different values for len(sample_indices) and n_samples"
 
-        samples = [xda.isel(time=slice(ridx, ridx+n_steps+n_spinup+1))
+        samples = [xda.isel(time=slice(ridx-n_spinup, ridx+n_steps+1))
                    for ridx in sample_indices]
 
         return samples, sample_indices
@@ -265,7 +265,9 @@ def get_sample_indices(self, data_length, n_samples, n_steps, n_spinup, random_s
         rstate = np.random.RandomState(seed=random_seed)
         n_valid = data_length - (n_steps + n_spinup)
         sample_indices = rstate.choice(n_valid, n_samples, replace=False)
-        sample_indices = list(int(x) for x in sample_indices)
+
+        # add spinup here to get initial condition of the prediction, not including spinup
+        sample_indices = list(int(x+n_spinup) for x in sample_indices)
         return sample_indices