Add convenience methods to extract ligand trajectories by replica

proteneer / timemachine

Differentiate all the things!

Other

138 stars 17 forks source link

Add convenience methods to extract ligand trajectories by replica #1320

Closed mcwitt closed 2 months ago

mcwitt commented 2 months ago

Adds a new subclass HREXSimulationResult(SimulationResult), moves HREX-specific fields from generic SimulationResult
Adds convenience methods:
- extract_trajectories_by_replica(atom_idxs)
- extract_ligand_trajectories_by_replica()
to HREXSimulationResult

mcwitt commented 2 months ago

This should go in after #1321

mcwitt commented 2 months ago

(also adding @jkausrelay as reviewer, since we might want to save ligand trajectories downstream)

jkausrelay commented 2 months ago

Does the return type of https://github.com/proteneer/timemachine/pull/1320/files#diff-49ef8469d9fb5c9d39e05a0cef32710703096c53f657e55eff49d5c3ddd47b13R459 need to be updated (estimate_relative_free_energy_bisection_or_hrex)?

mcwitt commented 2 months ago

Does the return type of https://github.com/proteneer/timemachine/pull/1320/files#diff-49ef8469d9fb5c9d39e05a0cef32710703096c53f657e55eff49d5c3ddd47b13R459 need to be updated (estimate_relative_free_energy_bisection_or_hrex)?

I don't think so? HREXSimulationResult is a subclass of SimulationResult, so it is valid for the return type to be the former. I don't think it's straightforward to make the annotation more precise, since estimate_relative_free_energy_bisection_or_hrex could return a plain SimulationResult if md_params.hrex_params is None.

mcwitt commented 2 months ago

Does the return type of https://github.com/proteneer/timemachine/pull/1320/files#diff-49ef8469d9fb5c9d39e05a0cef32710703096c53f657e55eff49d5c3ddd47b13R459 need to be updated (estimate_relative_free_energy_bisection_or_hrex)?

I don't think so? HREXSimulationResult is a subclass of SimulationResult, so it is valid for the return type to be the former. I don't think it's straightforward to make the annotation more precise, since estimate_relative_free_energy_bisection_or_hrex could return a plain SimulationResult if md_params.hrex_params is None.

@jkausrelay we could also annotate it as returning SimulationResult | HREXSimulationResult - maybe that's what you had in mind? I could see the argument for that being more readable (though I think it's equivalent from a type-checking perspective)