Cannot append to exisiting trajectory via `ts.integrate`/`ts.optimize`

[danielzuegner]

When trying to append to existing trajectories, I construct my reporter like so:

reporter = TrajectoryReporter(
    filenames=relax_filenames,
    state_frequency=1,
    metadata=self.metadata,
    trajectory_kwargs=dict(mode="a"),
)

Here, the trajectories inside relax_filenames already contain partial trajectories. Now, if I continue from these trajectories into ts.integrate, I get an error in _validate_arrays():

  File "/opt/miniconda/envs/env/lib/python3.12/site-packages/torch_sim/runners.py", line 499, in optimize
    trajectory_reporter.report(state, step, model=model)
  File "/opt/miniconda/envs/env/lib/python3.12/site-packages/torch_sim/trajectory.py", line 264, in report
    self.trajectories[idx].write_state(substate, step, **self.state_kwargs)
  File "/opt/miniconda/envs/env/lib/python3.12/site-packages/torch_sim/trajectory.py", line 782, in write_state
    self.write_arrays(data, steps)
  File "/opt/miniconda/envs/env/lib/python3.12/site-packages/torch_sim/trajectory.py", line 510, in write_arrays
    self._validate_array(name, array, steps)
  File "/opt/miniconda/envs/env/lib/python3.12/site-packages/torch_sim/trajectory.py", line 587, in _validate_array
    raise ValueError(
ValueError: steps[0]=1 must be greater than the last recorded step 10 for array positions

The issue appears to be that in ts.optimize/ts.integrate we always start from step=1, even if we're in append mode. In that case, the validation will fail and lead to a crash. Thus, I don't see a way currently to actually append to a trajectory via ts.optimize/ts.integrate.

From my perspective, a solution could be to auto-detect the current step from the trajectory if we're in a mode and start the integration from there, not 1. If you agree I can draft a small PR to address that.

Thanks!

[orionarcher]

Yep, that makes sense to me! A PR would be awesome. Would you make sure to write a test that demonstrates the before/after difference? Happy to take a look when done.

[danielzuegner]

@orionarcher see here: #361. Wasn't able to add you as reviewer I think.

[thomasloux]

Can you describe the exact use case that you need?
Because I feel like you trying to provide a wrong solution to a problem that is relevant.
If you want to continue the optimization of some systems but others are converged, you shoudn't have to continue running the one that are converged.
Let's think quickly what is the best solution before jumping to implementation

[danielzuegner]

I want to cover the case that a optimization or integration gets interrupted mid-way. In that case, I want to be able to load the trajectory and complete the relaxation/ MD. I am not concerned right now getting the state exactly back to how it was before it got interrupted (e.g., barostat/ thermostat) -- just being able to point to the existing trajectory and continuing from there is good enough for now.

[thomasloux]

Ok let's separate the different fixes then. I would suggest:

• Let's first fix the append trajectory problem for optimize: you solved it. But I would just modify the traj mode appending in TorchSimTrajectory.load_new_trajectories. Either manually change mode to append here, or only close the necessary files.
• Restart MD. All steps will always be the same, at least with the current implementation of ts.integrate. So either assume it's the case, or assuming a weird failure of the MD simulation, start from the lowest step. Example: failure during saving, so that you have for half of the systems until step 50 and the other half until step 40. Restart from step 40 state and call it a day. It may not be straightforward to rewrite some steps already present in the traj file, so you may have to first copy the traj file to get a clean one with all systems at the same step and then start again normally. Try not to use step_so_far, keep n_step and make it start at n_step_initial that is set from the traj file.
• Restart optimize: Now that you have done the most difficult part: check individually that step is <= max_step and otherwise remove these systems. You should reload and reoptimize all states from the current states in the Traj. For converged state, they will just be evaluated one time and then the convergence criteria will pass. So actually you don't have much more to modify as the autobatcher will do the job.
  Otherwise, nice finding, thank you!

[orionarcher]

Let's first fix the append trajectory problem for optimize: you solved it. But I would just modify the traj mode appending in TorchSimTrajectory.load_new_trajectories. Either manually change mode to append here, or only close the necessary files.

I am not sure this is a bug, it may just be a poor/underdocumented default. If you instantiate the provided TrajectoryReporter in append mode, I don't think it will overwrite the files. It's just that by default a TrajectoryReporter will be created in write mode. I don't know that append should be the default, because the user might very well want to overwrite, especially if they are testing things.

EDIT: Sorry, I understand the issue now.

[danielzuegner]

@thomasloux thanks for the useful feedback. In this commit I'm tackling your first point:

• change mode inside load_new_trajectories as suggested
• to distinguish the "reloading" case from the first initialization of the reporter, I'm now constructing the trajectories right in the constructor of TrajectoryReporter, where indeed we do want to respect the "w" mode if specified.
• I found the name load_new_trajectories not ideal for what it does now and how it's used, so I suggest to change it to reopen_trajectories. This reflects how it's used in the code: reopen the input trajectories and close all others.
• Slightly orthogonal, I suggest to change self.filenames to a computed @property to ensure that it is always in sync with the filenames in self.trajectories.

Let me know if you have suggestions. Will move on to the other two points.