diff --git a/observations.py b/observations.py index 7c94e82ccdb68b963c55e676cc9c3caecbbdcbfb..f2f9cbea1188f5d79239f48900ce75529e27e557 100644 --- a/observations.py +++ b/observations.py @@ -39,7 +39,7 @@ def decode_observations(observation_files): return observations,ocoords -def random_sort_observations(observations,ocoords): +def random_sort_observations(observations,ocoords,truths=None): # Check that input arrays have consistent dimensions assert (observations.shape == ocoords.shape),\ @@ -50,14 +50,32 @@ def random_sort_observations(observations,ocoords): shuffled_obs = np.zeros(observations.shape)+np.nan shuffled_coords = np.zeros(observations.shape)+np.nan - # Randomize observation order, differently at each time - indices = np.arange(nobs) - for i in range(nassim): - np.random.shuffle(indices) - shuffled_obs[i,:] = observations[i,indices] - shuffled_coords[i,:] = ocoords[i,indices] - - return shuffled_obs,shuffled_coords + # Randomize arrays, differently at each time + if truths is None: + indices = np.arange(nobs) + for i in range(nassim): + np.random.shuffle(indices) + shuffled_obs[i,:] = observations[i,indices] + shuffled_coords[i,:] = ocoords[i,indices] + + return shuffled_obs,shuffled_coords + + else: + # Check that truths array has consistent dimensions + assert (truths.shape == ocoords.shape),\ + f'observation and coordinate arrays must have the same size' + + # Preallocate truths array + shuffled_truths = np.zeros(ocoords.shape)+np.nan + + # Randomize arrays, differently at each time + indices = np.arange(nobs) + for i in range(nassim): + np.random.shuffle(indices) + shuffled_obs[i,:] = observations[i,indices] + shuffled_coords[i,:] = ocoords[i,indices] + shuffled_truths[i,:] = truths[i,indices] + return shuffled_obs,shuffled_coords,shuffled_truths if __name__ == '__main__':