From 40a747edfa670cedaac91069bb0f75222624f8ff Mon Sep 17 00:00:00 2001
From: Stefano Serafin <serafin@jet01.img.univie.ac.at>
Date: Fri, 24 Jan 2025 09:14:42 +0100
Subject: [PATCH] added ability to randomize true system states (besides
 observations and coordinates)

---
 observations.py | 36 +++++++++++++++++++++++++++---------
 1 file changed, 27 insertions(+), 9 deletions(-)

diff --git a/observations.py b/observations.py
index 7c94e82..f2f9cbe 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__':
 
-- 
GitLab