[Bug]: Restricted trajectory training uses at most one trajectory #72

if traj in trajectories_sampled:
      global_indices = global_indices + list(
          range(0, self.n_windows_per_trajectory[file_index])
      )
  current_index += self.n_windows_per_trajectory[file_index]

  if traj in trajectories_sampled:
      n_windows = self.n_windows_per_trajectory[file_index]
      global_indices = global_indices + list(
          range(current_index, current_index + n_windows)
      )
  current_index += self.n_windows_per_trajectory[file_index]

import numpy as np
from the_well.data import WellDataset

dataset = WellDataset(
    well_base_path="data/the-well/datasets",
    well_dataset_name="active_matter",
    well_split_name="train",
    n_steps_input=1,
    n_steps_output=1,
    restrict_num_trajectories=2,
)

print(f"\nrestriction_set: {dataset.restriction_set}")
print(f"Unique indices in restriction_set: {np.unique(dataset.restriction_set)}")

time_per_traj = 80 # active_matter

sample1 = dataset[0]["output_fields"]
sample2 = dataset[time_per_traj]["output_fields"]

print(f"{np.allclose(sample1, sample2)}")