fix restricted trajectories

t-muser · t-muser · commit 8e6fd0885a7d · 2025-12-23T13:54:45.000+01:00
diff --git a/the_well/data/datasets.py b/the_well/data/datasets.py
@@ -389,8 +389,9 @@ def _build_restriction_set(
                     - 1
                 )
                 if traj in trajectories_sampled:
+                    n_windows = self.n_windows_per_trajectory[file_index]
                     global_indices = global_indices + list(
-                        range(0, self.n_windows_per_trajectory[file_index])
+                        range(current_index, current_index + n_windows)
                     )
                 current_index += self.n_windows_per_trajectory[file_index]
             global_indices = np.array(global_indices)

Original file line number	Diff line number	Diff line change
`@@ -389,8 +389,9 @@ def _build_restriction_set(`
`389`	`389`	`- 1`
`390`	`390`	`)`
`391`	`391`	`if traj in trajectories_sampled:`
	`392`	`+ n_windows = self.n_windows_per_trajectory[file_index]`
`392`	`393`	`global_indices = global_indices + list(`
`393`		`- range(0, self.n_windows_per_trajectory[file_index])`
	`394`	`+ range(current_index, current_index + n_windows)`
`394`	`395`	`)`
`395`	`396`	`current_index += self.n_windows_per_trajectory[file_index]`
`396`	`397`	`global_indices = np.array(global_indices)`