Merge pull request #73 from till-m/fix-restricted-trajectories

mikemccabe210 · web-flow · commit 718df875d1b1 · 2026-02-25T13:29:31.000-05:00
fix restricted trajectories
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)`