diff --git a/CHANGELOG.md b/CHANGELOG.md index 8da32d65..25e15106 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ # Changelog +## Unreleased + +### Added +- Onboarded `hyper_solvate`, `hyper_run` and `hyper_minimize` + ## 7.0.0 No new changes since 7.0.0rc3. diff --git a/docs/modules.rst b/docs/modules.rst index 9e530464..e5de2e24 100644 --- a/docs/modules.rst +++ b/docs/modules.rst @@ -7,6 +7,9 @@ Rush Modules rush.exess_optimization rush.exess_qmmm rush.nnxtb + rush.hyper + rush.hyper_minimize_sumo + rush.hyper_run_sumo rush.prepare rush.auto3d rush.mmseqs2 diff --git a/docs/rush.hyper.rst b/docs/rush.hyper.rst new file mode 100644 index 00000000..91cb4e70 --- /dev/null +++ b/docs/rush.hyper.rst @@ -0,0 +1,31 @@ +:tocdepth: 3 + +Hyper Solvation +=============== + +.. automodule:: rush.hyper._hyper_solvate_sumo + +.. currentmodule:: rush.hyper + +Run Submission +-------------- + +.. autofunction:: hyper_solvate_sumo + +Input Types +----------- + +.. autoclass:: HyperConfig + :members: + :undoc-members: + +Result Types +------------ + +.. autoclass:: ItemError + :members: + :undoc-members: + +.. autoclass:: TRCBatchResultRef + :members: + :undoc-members: diff --git a/docs/rush.hyper_minimize_sumo.rst b/docs/rush.hyper_minimize_sumo.rst new file mode 100644 index 00000000..c0916d43 --- /dev/null +++ b/docs/rush.hyper_minimize_sumo.rst @@ -0,0 +1,30 @@ +:tocdepth: 3 + +Hyper Minimize +============== + +.. automodule:: rush.hyper._hyper_minimize_sumo + +.. currentmodule:: rush.hyper + +Run Submission +-------------- + +.. autofunction:: hyper_minimize_sumo + +Input Types +----------- + +.. autoclass:: HyperMinimizeConfig + :members: + :undoc-members: + +.. autoclass:: MinimizeInput + :members: + :undoc-members: + +Result Types +------------ + +Minimization returns :class:`rush.hyper.TRCBatchResultRef` with per-item +results that can be fetched as ``TRC`` objects or ``ItemError`` values. \ No newline at end of file diff --git a/docs/rush.hyper_run_sumo.rst b/docs/rush.hyper_run_sumo.rst new file mode 100644 index 00000000..361f034b --- /dev/null +++ b/docs/rush.hyper_run_sumo.rst @@ -0,0 +1,43 @@ +:tocdepth: 3 + +Hyper Run +========= + +.. automodule:: rush.hyper._hyper_run_sumo + +.. currentmodule:: rush.hyper + +Run Submission +-------------- + +.. autofunction:: hyper_run_sumo + +Input Types +----------- + +.. autoclass:: HyperRunConfig + :members: + :undoc-members: + +.. autoclass:: RunInput + :members: + :undoc-members: + +Result Types +------------ + +.. autoclass:: RunOutput + :members: + :undoc-members: + +.. autoclass:: RunOutputRef + :members: + :undoc-members: + +.. autoclass:: RunOutputPaths + :members: + :undoc-members: + +.. autoclass:: RunResultRef + :members: + :undoc-members: diff --git a/docs/tutorials.rst b/docs/tutorials.rst index d231e9cf..31481481 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -11,3 +11,6 @@ Tutorials tutorials/05-exess-interaction-energy tutorials/06-exess-qmmm tutorials/07-nnxtb-energy + tutorials/08-hyper-solvate_sumo + tutorials/09-hyper-minimize_sumo + tutorials/10-hyper-run_sumo diff --git a/docs/tutorials/08-hyper-solvate_sumo.md b/docs/tutorials/08-hyper-solvate_sumo.md new file mode 100644 index 00000000..59d03cc3 --- /dev/null +++ b/docs/tutorials/08-hyper-solvate_sumo.md @@ -0,0 +1,51 @@ +# Tutorial 8: Hyper Solvate Sumo + +**What you get:** A solvated `TRC` structure that you can pass directly into minimization or MD. + +| | | +|---|---| +| **Time** | ~2-5 minutes | +| **Skill level** | Beginner | +| **Prerequisites** | Python 3.12+, `rush-py` installed, `RUSH_TOKEN` and `RUSH_PROJECT` set | + +--- + +## Quick Start + +```python +from pathlib import Path + +from rush import RunOpts, TRC, hyper + +input_trc = Path("tests/data/hyper/valid_trc.json") + +run = hyper.hyper_solvate_sumo( + [input_trc], + config=hyper.HyperConfig(max_inputs=8, padding_nm=0.8, seed=12345, timeout_seconds=120), + run_opts=RunOpts(name="Tutorial: Hyper Solvate", tags=["rush-py", "tutorial", "hyper", "solvate"]), +) + +item = run.fetch()[0] +if not isinstance(item, TRC): + raise RuntimeError(f"Unexpected per-item error: {item}") + +print("Solvated atoms:", len(item.topology.symbols)) +``` + +--- + +## Output Contract + +`hyper_solvate_sumo()` returns `Run[hyper.TRCBatchResultRef]`. + +- `run.collect()` -> `TRCBatchResultRef` +- `result_ref.fetch()` -> `list[TRC | ItemError]` +- `result_ref.save()` -> `list[Path | ItemError]` + +--- + +## See Also + +- {doc}`Hyper Minimize tutorial <09-hyper-minimize_sumo>` +- {doc}`Hyper Run tutorial <10-hyper-run_sumo>` +- {doc}`Hyper API reference <../rush.hyper>` diff --git a/docs/tutorials/09-hyper-minimize_sumo.md b/docs/tutorials/09-hyper-minimize_sumo.md new file mode 100644 index 00000000..0c1eaffb --- /dev/null +++ b/docs/tutorials/09-hyper-minimize_sumo.md @@ -0,0 +1,56 @@ +# Tutorial 9: Hyper Minimize Sumo + +**What you get:** An energy-minimized structure from a structure/topology pair. + +| | | +|---|---| +| **Time** | ~2-8 minutes | +| **Skill level** | Beginner | +| **Prerequisites** | Python 3.12+, `rush-py` installed, `RUSH_TOKEN` and `RUSH_PROJECT` set | + +--- + +## Quick Start + +```python +from pathlib import Path + +from rush import RunOpts, TRC, hyper + +data_dir = Path("tests/data/hyper") + +run = hyper.hyper_minimize_sumo( + [ + hyper.MinimizeInput( + structure=data_dir / "methanol_trc.json", + topology=data_dir / "methanol_topology.json", + ) + ], + config=hyper.HyperMinimizeConfig(max_inputs=4, steps=100, gtol=100.0, timeout_seconds=900), + run_opts=RunOpts(name="Tutorial: Hyper Minimize", tags=["rush-py", "tutorial", "hyper", "minimize"]), +) + +item = run.fetch()[0] +if not isinstance(item, TRC): + raise RuntimeError(f"Unexpected per-item error: {item}") + +print("Minimized atoms:", len(item.topology.symbols)) +``` + +--- + +## Output Contract + +`hyper_minimize_sumo()` returns `Run[hyper.TRCBatchResultRef]`. + +- `run.collect()` -> `TRCBatchResultRef` +- `result_ref.fetch()` -> `list[TRC | ItemError]` +- `result_ref.save()` -> `list[Path | ItemError]` + +--- + +## See Also + +- {doc}`Hyper Solvate tutorial <08-hyper-solvate_sumo>` +- {doc}`Hyper Run tutorial <10-hyper-run_sumo>` +- {doc}`Hyper Minimize API reference <../rush.hyper_minimize_sumo>` diff --git a/docs/tutorials/10-hyper-run_sumo.md b/docs/tutorials/10-hyper-run_sumo.md new file mode 100644 index 00000000..570543d7 --- /dev/null +++ b/docs/tutorials/10-hyper-run_sumo.md @@ -0,0 +1,69 @@ +# Tutorial 10: Hyper Run Sumo + +**What you get:** Trajectory and optional checkpoint artifacts from short molecular dynamics runs. + +| | | +|---|---| +| **Time** | ~3-10 minutes | +| **Skill level** | Intermediate | +| **Prerequisites** | Python 3.12+, `rush-py` installed, `RUSH_TOKEN` and `RUSH_PROJECT` set | + +--- + +## Quick Start + +```python +from pathlib import Path + +from rush import RunOpts, hyper + +data_dir = Path("tests/data/hyper") + +run = hyper.hyper_run_sumo( + [ + hyper.RunInput( + sim_config=data_dir / "sim_config.json", + topology=data_dir / "methanol_topology.json", + coordinates=data_dir / "methanol_trc.json", + ) + ], + config=hyper.HyperRunConfig( + max_inputs=4, + nsteps=20, + dt_ps=0.001, + temperature_k=300.0, + ensemble="Nvt", + minimize_before_run=False, + solvate_before_run=False, + use_gpu=False, + nthreads=1, + timeout_seconds=900, + ), + run_opts=RunOpts(name="Tutorial: Hyper Run", tags=["rush-py", "tutorial", "hyper", "run"]), +) + +item = run.fetch()[0] +if isinstance(item, hyper.ItemError): + raise RuntimeError(f"Hyper run failed: {item}") + +print("Trajectory bytes:", len(item.trajectory)) +print("Checkpoint bytes:", 0 if item.checkpoint is None else len(item.checkpoint)) +``` + +--- + +## Output Contract + +`hyper_run_sumo()` returns `Run[hyper.RunResultRef]`. + +- `run.collect()` -> `RunResultRef` +- `result_ref.fetch()` -> `list[RunOutput | ItemError]` +- `result_ref.save()` -> `list[RunOutputPaths | ItemError]` + +--- + +## See Also + +- {doc}`Hyper Solvate tutorial <08-hyper-solvate_sumo>` +- {doc}`Hyper Minimize tutorial <09-hyper-minimize_sumo>` +- {doc}`Hyper Run API reference <../rush.hyper_run_sumo>` diff --git a/examples/hyper-minimize_sumo/09_hyper_minimize_sumo.py b/examples/hyper-minimize_sumo/09_hyper_minimize_sumo.py new file mode 100644 index 00000000..31a1d33a --- /dev/null +++ b/examples/hyper-minimize_sumo/09_hyper_minimize_sumo.py @@ -0,0 +1,39 @@ +"""Example: Hyper minimization workflow.""" + +from pathlib import Path + +from rush import RunOpts, TRC, hyper + +DATA_DIR = Path(__file__).parent / "data" + +run = hyper.hyper_minimize_sumo( + [ + hyper.MinimizeInput( + structure=DATA_DIR / "methanol_trc.json", + topology=DATA_DIR / "methanol_topology.json", + ) + ], + config=hyper.HyperMinimizeConfig(max_inputs=4, steps=100, gtol=100.0, timeout_seconds=900), + run_opts=RunOpts(name="Example: Hyper Minimize", tags=["rush-py", "example", "hyper", "minimize"]), +) + +result_ref = run.collect() +fetched = result_ref.fetch() +if len(fetched) != 1: + raise RuntimeError(f"Expected 1 output item, got {len(fetched)}") + +item = fetched[0] +if not isinstance(item, TRC): + raise RuntimeError(f"Expected TRC output, got {item}") + +print("Minimized atom count:", len(item.topology.symbols)) + +saved = result_ref.save() +if len(saved) != 1: + raise RuntimeError(f"Expected 1 saved item, got {len(saved)}") + +saved_item = saved[0] +if isinstance(saved_item, hyper.ItemError): + raise RuntimeError(f"Unexpected per-item save error: {saved_item}") + +print("Saved output:", saved_item) diff --git a/examples/hyper-minimize_sumo/README.md b/examples/hyper-minimize_sumo/README.md new file mode 100644 index 00000000..2e0e2600 --- /dev/null +++ b/examples/hyper-minimize_sumo/README.md @@ -0,0 +1,27 @@ +# Hyper Minimize Sumo Example + +Demonstrates `hyper.hyper_minimize_sumo()` with one structure/topology job. + +## Quick Start + +```bash +export RUSH_TOKEN="your-token" +export RUSH_PROJECT="your-project" + +python 09_hyper_minimize_sumo.py +``` + +## What This Example Covers + +1. Building a `MinimizeInput` job from local JSON files +2. Running Hyper minimization with explicit config bounds +3. Fetching and saving successful TRC output + +## Input Data + +- `data/methanol_trc.json` +- `data/methanol_topology.json` + +## Tutorial + +See: [Tutorial 9: Hyper Minimize Sumo](../../docs/tutorials/09-hyper-minimize_sumo.md) diff --git a/examples/hyper-minimize_sumo/data/methanol_topology.json b/examples/hyper-minimize_sumo/data/methanol_topology.json new file mode 100644 index 00000000..5d9c6f83 --- /dev/null +++ b/examples/hyper-minimize_sumo/data/methanol_topology.json @@ -0,0 +1,333 @@ +{ + "atom_type_count": 4, + "charges": [ + 0.28, + -0.68, + 0.0, + 0.0, + 0.0, + 0.4 + ], + "masses": [ + 12.011, + 15.999, + 1.008, + 1.008, + 1.008, + 1.008 + ], + "atom_types": [ + 0, + 1, + 2, + 2, + 2, + 3 + ], + "bonds": [ + { + "i": 0, + "j": 1, + "b0": 0.14179999999999998, + "kb": 303937.193826 + }, + { + "i": 0, + "j": 2, + "b0": 0.1093, + "kb": 287014.992228 + }, + { + "i": 0, + "j": 3, + "b0": 0.1093, + "kb": 287014.992228 + }, + { + "i": 0, + "j": 4, + "b0": 0.1093, + "kb": 287014.992228 + }, + { + "i": 1, + "j": 5, + "b0": 0.0972, + "kb": 469365.2642520001 + } + ], + "angles": [ + { + "i": 1, + "j": 0, + "k": 2, + "theta0": 1.895026141937883, + "ktheta": 470.32508605816537 + }, + { + "i": 1, + "j": 0, + "k": 3, + "theta0": 1.895026141937883, + "ktheta": 470.32508605816537 + }, + { + "i": 1, + "j": 0, + "k": 4, + "theta0": 1.895026141937883, + "ktheta": 470.32508605816537 + }, + { + "i": 2, + "j": 0, + "k": 3, + "theta0": 1.8995465447005484, + "ktheta": 310.7397495595562 + }, + { + "i": 2, + "j": 0, + "k": 4, + "theta0": 1.8995465447005484, + "ktheta": 310.7397495595562 + }, + { + "i": 3, + "j": 0, + "k": 4, + "theta0": 1.8995465447005484, + "ktheta": 310.7397495595562 + }, + { + "i": 0, + "j": 1, + "k": 5, + "theta0": 1.8588280132515207, + "ktheta": 477.55159186187603 + } + ], + "proper_dihedrals": [ + { + "i": 2, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 1.246832, + "mult": 1 + }, + { + "i": 2, + "j": 0, + "k": 1, + "l": 5, + "phi0": 3.141592653589793, + "kphi": -0.577392, + "mult": 2 + }, + { + "i": 2, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 0.7238319999999999, + "mult": 3 + }, + { + "i": 3, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 1.246832, + "mult": 1 + }, + { + "i": 3, + "j": 0, + "k": 1, + "l": 5, + "phi0": 3.141592653589793, + "kphi": -0.577392, + "mult": 2 + }, + { + "i": 3, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 0.7238319999999999, + "mult": 3 + }, + { + "i": 4, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 1.246832, + "mult": 1 + }, + { + "i": 4, + "j": 0, + "k": 1, + "l": 5, + "phi0": 3.141592653589793, + "kphi": -0.577392, + "mult": 2 + }, + { + "i": 4, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 0.7238319999999999, + "mult": 3 + } + ], + "improper_dihedrals": [], + "pairs14": [ + { + "i": 2, + "j": 5, + "c6": 4.274187696047065e-05, + "c12": 1.0119320214366275e-08 + }, + { + "i": 3, + "j": 5, + "c6": 4.274187696047065e-05, + "c12": 1.0119320214366275e-08 + }, + { + "i": 4, + "j": 5, + "c6": 4.274187696047065e-05, + "c12": 1.0119320214366275e-08 + } + ], + "cmap_entries": [], + "lj_params": [ + { + "c6": 0.002115016015813491, + "c12": 3.94244116886151e-06 + }, + { + "c6": 0.0016677018383901895, + "c12": 2.311249945030495e-06 + }, + { + "c6": 0.0005432157553294566, + "c12": 4.6102104135361286e-07 + }, + { + "c6": 0.0003954334320525218, + "c12": 2.4429962496901167e-07 + }, + { + "c6": 0.0016677018383901895, + "c12": 2.311249945030495e-06 + }, + { + "c6": 0.00129488850377234, + "c12": 1.3138538180583043e-06 + }, + { + "c6": 0.0004104134700407302, + "c12": 2.481369517052484e-07 + }, + { + "c6": 0.00029313999406590564, + "c12": 1.2658970268056165e-07 + }, + { + "c6": 0.0005432157553294566, + "c12": 4.6102104135361286e-07 + }, + { + "c6": 0.0004104134700407302, + "c12": 2.481369517052484e-07 + }, + { + "c6": 0.00012386870467525756, + "c12": 4.249495325323946e-08 + }, + { + "c6": 8.54837539209413e-05, + "c12": 2.023864042873255e-08 + }, + { + "c6": 0.0003954334320525218, + "c12": 2.4429962496901167e-07 + }, + { + "c6": 0.00029313999406590564, + "c12": 1.2658970268056165e-07 + }, + { + "c6": 8.54837539209413e-05, + "c12": 2.023864042873255e-08 + }, + { + "c6": 5.756897163937681e-05, + "c12": 9.17890990287509e-09 + } + ], + "cmap_grid_size": 24, + "cmap_grids": [], + "exclusions": [ + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ] + ] +} diff --git a/examples/hyper-minimize_sumo/data/methanol_trc.json b/examples/hyper-minimize_sumo/data/methanol_trc.json new file mode 100644 index 00000000..d753464a --- /dev/null +++ b/examples/hyper-minimize_sumo/data/methanol_trc.json @@ -0,0 +1,39 @@ +{ + "topology": { + "schema_version": "0.2.0", + "symbols": ["C", "O", "H", "H", "H", "H"], + "geometry": [ + -0.37, 0.01, -0.01, + 0.92, -0.5, -0.3, + -0.52, 0.03, 1.07, + -0.47, 1.01, -0.42, + -1.12, -0.64, -0.47, + 1.56, 0.09, 0.12 + ], + "labels": null, + "partial_charges": null, + "formal_charges": null, + "connectivity": null, + "stereochemistry": null, + "velocities": null, + "fragments": null, + "fragment_formal_charges": null, + "fragment_partial_charges": null, + "fragment_multiplicities": null + }, + "residues": { + "residues": [[0, 1, 2, 3, 4, 5]], + "seqs": ["LIG"], + "seq_ns": [1], + "insertion_codes": [""], + "labeled": null, + "labels": null + }, + "chains": { + "chains": [[0]], + "alpha_helices": null, + "beta_sheets": null, + "labeled": null, + "labels": null + } +} diff --git a/examples/hyper-run_sumo/10_hyper_run_sumo.py b/examples/hyper-run_sumo/10_hyper_run_sumo.py new file mode 100644 index 00000000..ae026543 --- /dev/null +++ b/examples/hyper-run_sumo/10_hyper_run_sumo.py @@ -0,0 +1,53 @@ +"""Example: Hyper molecular dynamics run workflow.""" + +from pathlib import Path + +from rush import RunOpts, hyper + +DATA_DIR = Path(__file__).parent / "data" + +run = hyper.hyper_run_sumo( + [ + hyper.RunInput( + sim_config=DATA_DIR / "sim_config.json", + topology=DATA_DIR / "methanol_topology.json", + coordinates=DATA_DIR / "methanol_trc.json", + ) + ], + config=hyper.HyperRunConfig( + max_inputs=4, + nsteps=20, + dt_ps=0.001, + temperature_k=300.0, + ensemble="Nvt", + minimize_before_run=False, + solvate_before_run=False, + use_gpu=False, + nthreads=1, + timeout_seconds=900, + ), + run_opts=RunOpts(name="Example: Hyper Run", tags=["rush-py", "example", "hyper", "run"]), +) + +result_ref = run.collect() +fetched = result_ref.fetch() +if len(fetched) != 1: + raise RuntimeError(f"Expected 1 output item, got {len(fetched)}") + +item = fetched[0] +if isinstance(item, hyper.ItemError): + raise RuntimeError(f"Unexpected per-item run error: {item}") + +print("Trajectory bytes:", len(item.trajectory)) +print("Checkpoint bytes:", 0 if item.checkpoint is None else len(item.checkpoint)) + +saved = result_ref.save() +if len(saved) != 1: + raise RuntimeError(f"Expected 1 saved item, got {len(saved)}") + +saved_item = saved[0] +if isinstance(saved_item, hyper.ItemError): + raise RuntimeError(f"Unexpected per-item save error: {saved_item}") + +print("Saved trajectory:", saved_item.trajectory) +print("Saved checkpoint:", saved_item.checkpoint) diff --git a/examples/hyper-run_sumo/README.md b/examples/hyper-run_sumo/README.md new file mode 100644 index 00000000..1a4585a4 --- /dev/null +++ b/examples/hyper-run_sumo/README.md @@ -0,0 +1,28 @@ +# Hyper Run Sumo Example + +Demonstrates `hyper.hyper_run_sumo()` for a short MD run using prebuilt simulation config and topology artifacts. + +## Quick Start + +```bash +export RUSH_TOKEN="your-token" +export RUSH_PROJECT="your-project" + +python 10_hyper_run_sumo.py +``` + +## What This Example Covers + +1. Submitting one `RunInput` job with explicit `HyperRunConfig` +2. Fetching trajectory/checkpoint payloads as bytes +3. Saving run artifacts to the Rush workspace + +## Input Data + +- `data/sim_config.json` +- `data/methanol_topology.json` +- `data/methanol_trc.json` + +## Tutorial + +See: [Tutorial 10: Hyper Run Sumo](../../docs/tutorials/10-hyper-run_sumo.md) diff --git a/examples/hyper-run_sumo/data/methanol_topology.json b/examples/hyper-run_sumo/data/methanol_topology.json new file mode 100644 index 00000000..5d9c6f83 --- /dev/null +++ b/examples/hyper-run_sumo/data/methanol_topology.json @@ -0,0 +1,333 @@ +{ + "atom_type_count": 4, + "charges": [ + 0.28, + -0.68, + 0.0, + 0.0, + 0.0, + 0.4 + ], + "masses": [ + 12.011, + 15.999, + 1.008, + 1.008, + 1.008, + 1.008 + ], + "atom_types": [ + 0, + 1, + 2, + 2, + 2, + 3 + ], + "bonds": [ + { + "i": 0, + "j": 1, + "b0": 0.14179999999999998, + "kb": 303937.193826 + }, + { + "i": 0, + "j": 2, + "b0": 0.1093, + "kb": 287014.992228 + }, + { + "i": 0, + "j": 3, + "b0": 0.1093, + "kb": 287014.992228 + }, + { + "i": 0, + "j": 4, + "b0": 0.1093, + "kb": 287014.992228 + }, + { + "i": 1, + "j": 5, + "b0": 0.0972, + "kb": 469365.2642520001 + } + ], + "angles": [ + { + "i": 1, + "j": 0, + "k": 2, + "theta0": 1.895026141937883, + "ktheta": 470.32508605816537 + }, + { + "i": 1, + "j": 0, + "k": 3, + "theta0": 1.895026141937883, + "ktheta": 470.32508605816537 + }, + { + "i": 1, + "j": 0, + "k": 4, + "theta0": 1.895026141937883, + "ktheta": 470.32508605816537 + }, + { + "i": 2, + "j": 0, + "k": 3, + "theta0": 1.8995465447005484, + "ktheta": 310.7397495595562 + }, + { + "i": 2, + "j": 0, + "k": 4, + "theta0": 1.8995465447005484, + "ktheta": 310.7397495595562 + }, + { + "i": 3, + "j": 0, + "k": 4, + "theta0": 1.8995465447005484, + "ktheta": 310.7397495595562 + }, + { + "i": 0, + "j": 1, + "k": 5, + "theta0": 1.8588280132515207, + "ktheta": 477.55159186187603 + } + ], + "proper_dihedrals": [ + { + "i": 2, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 1.246832, + "mult": 1 + }, + { + "i": 2, + "j": 0, + "k": 1, + "l": 5, + "phi0": 3.141592653589793, + "kphi": -0.577392, + "mult": 2 + }, + { + "i": 2, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 0.7238319999999999, + "mult": 3 + }, + { + "i": 3, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 1.246832, + "mult": 1 + }, + { + "i": 3, + "j": 0, + "k": 1, + "l": 5, + "phi0": 3.141592653589793, + "kphi": -0.577392, + "mult": 2 + }, + { + "i": 3, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 0.7238319999999999, + "mult": 3 + }, + { + "i": 4, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 1.246832, + "mult": 1 + }, + { + "i": 4, + "j": 0, + "k": 1, + "l": 5, + "phi0": 3.141592653589793, + "kphi": -0.577392, + "mult": 2 + }, + { + "i": 4, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 0.7238319999999999, + "mult": 3 + } + ], + "improper_dihedrals": [], + "pairs14": [ + { + "i": 2, + "j": 5, + "c6": 4.274187696047065e-05, + "c12": 1.0119320214366275e-08 + }, + { + "i": 3, + "j": 5, + "c6": 4.274187696047065e-05, + "c12": 1.0119320214366275e-08 + }, + { + "i": 4, + "j": 5, + "c6": 4.274187696047065e-05, + "c12": 1.0119320214366275e-08 + } + ], + "cmap_entries": [], + "lj_params": [ + { + "c6": 0.002115016015813491, + "c12": 3.94244116886151e-06 + }, + { + "c6": 0.0016677018383901895, + "c12": 2.311249945030495e-06 + }, + { + "c6": 0.0005432157553294566, + "c12": 4.6102104135361286e-07 + }, + { + "c6": 0.0003954334320525218, + "c12": 2.4429962496901167e-07 + }, + { + "c6": 0.0016677018383901895, + "c12": 2.311249945030495e-06 + }, + { + "c6": 0.00129488850377234, + "c12": 1.3138538180583043e-06 + }, + { + "c6": 0.0004104134700407302, + "c12": 2.481369517052484e-07 + }, + { + "c6": 0.00029313999406590564, + "c12": 1.2658970268056165e-07 + }, + { + "c6": 0.0005432157553294566, + "c12": 4.6102104135361286e-07 + }, + { + "c6": 0.0004104134700407302, + "c12": 2.481369517052484e-07 + }, + { + "c6": 0.00012386870467525756, + "c12": 4.249495325323946e-08 + }, + { + "c6": 8.54837539209413e-05, + "c12": 2.023864042873255e-08 + }, + { + "c6": 0.0003954334320525218, + "c12": 2.4429962496901167e-07 + }, + { + "c6": 0.00029313999406590564, + "c12": 1.2658970268056165e-07 + }, + { + "c6": 8.54837539209413e-05, + "c12": 2.023864042873255e-08 + }, + { + "c6": 5.756897163937681e-05, + "c12": 9.17890990287509e-09 + } + ], + "cmap_grid_size": 24, + "cmap_grids": [], + "exclusions": [ + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ] + ] +} diff --git a/examples/hyper-run_sumo/data/methanol_trc.json b/examples/hyper-run_sumo/data/methanol_trc.json new file mode 100644 index 00000000..d753464a --- /dev/null +++ b/examples/hyper-run_sumo/data/methanol_trc.json @@ -0,0 +1,39 @@ +{ + "topology": { + "schema_version": "0.2.0", + "symbols": ["C", "O", "H", "H", "H", "H"], + "geometry": [ + -0.37, 0.01, -0.01, + 0.92, -0.5, -0.3, + -0.52, 0.03, 1.07, + -0.47, 1.01, -0.42, + -1.12, -0.64, -0.47, + 1.56, 0.09, 0.12 + ], + "labels": null, + "partial_charges": null, + "formal_charges": null, + "connectivity": null, + "stereochemistry": null, + "velocities": null, + "fragments": null, + "fragment_formal_charges": null, + "fragment_partial_charges": null, + "fragment_multiplicities": null + }, + "residues": { + "residues": [[0, 1, 2, 3, 4, 5]], + "seqs": ["LIG"], + "seq_ns": [1], + "insertion_codes": [""], + "labeled": null, + "labels": null + }, + "chains": { + "chains": [[0]], + "alpha_helices": null, + "beta_sheets": null, + "labeled": null, + "labels": null + } +} diff --git a/examples/hyper-run_sumo/data/sim_config.json b/examples/hyper-run_sumo/data/sim_config.json new file mode 100644 index 00000000..d59922b6 --- /dev/null +++ b/examples/hyper-run_sumo/data/sim_config.json @@ -0,0 +1 @@ +{"force_field":"AMBERFF19SB","ensemble":"NVT","n_steps":40,"dt":0.001,"temperature":300.0,"initial_temp":300.0,"tau_t":0.1,"electrostatics":"CUTOFF","cutoff":0.0,"dispersion_correction":false,"constraints":"NONE","box_x":5.0,"box_y":5.0,"box_z":5.0,"output":{"trajectory_freq":1,"energy_freq":1,"checkpoint_freq":20,"log_freq":10}} diff --git a/examples/hyper-solvate/08_hyper_solvate.py b/examples/hyper-solvate/08_hyper_solvate.py new file mode 100644 index 00000000..eed6c1f0 --- /dev/null +++ b/examples/hyper-solvate/08_hyper_solvate.py @@ -0,0 +1,70 @@ +""" +Example: Hyper Solvation + +This script demonstrates how to: +1. Submit a Hyper solvation run for one TRC input +2. Fetch structured output with per-item error handling +3. Save the solvated structure to the Rush workspace + +Tutorial: https://exess.qdx.co/docs/tutorials/08-hyper-solvate.html + +Prerequisites: + - Set RUSH_TOKEN and RUSH_PROJECT environment variables + - Input file: valid_trc.json (provided in data/) +""" + +from pathlib import Path + +from rush import TRC, hyper +from rush import RunOpts + +DATA_DIR = Path(__file__).parent / "data" +INPUT_TRC = DATA_DIR / "valid_trc.json" + +# ===== Submit Hyper solvation run ===== +print("=" * 60) +print("Hyper Solvation") +print("=" * 60) + +run = hyper.hyper_solvate_sumo( + [INPUT_TRC], + config=hyper.HyperConfig(max_inputs=8, padding_nm=0.8, seed=12345, timeout_seconds=120), + run_opts=RunOpts( + name="Tutorial: Hyper Solvate", + tags=["rush-py", "tutorial", "hyper", "solvate"], + ), +) + +result_ref = run.collect() + +# ===== Fetch parsed output ===== +fetched = result_ref.fetch() +if len(fetched) != 1: + raise RuntimeError(f"Expected 1 output item, got {len(fetched)}") + +item = fetched[0] +if isinstance(item, hyper.ItemError): + raise RuntimeError(f"Hyper returned per-item error: {item}") +if not isinstance(item, TRC): + raise TypeError(f"Expected TRC output, got {type(item).__name__}") + +# ===== Print output summary ===== +print() +print("Results:") +print("-" * 60) +print(f"Input file: {INPUT_TRC.name}") +print(f"Solvated atom count: {len(item.topology.symbols)}") +print(f"Residue count: {len(item.residues.residues)}") + +# ===== Save output artifact ===== +saved = result_ref.save() +if len(saved) != 1: + raise RuntimeError(f"Expected 1 saved item, got {len(saved)}") + +saved_item = saved[0] +if isinstance(saved_item, hyper.ItemError): + raise RuntimeError(f"Hyper returned per-item error while saving: {saved_item}") + +print(f"Saved output path: {saved_item}") +print("-" * 60) +print("Done!") diff --git a/examples/hyper-solvate/README.md b/examples/hyper-solvate/README.md new file mode 100644 index 00000000..f8e6fb69 --- /dev/null +++ b/examples/hyper-solvate/README.md @@ -0,0 +1,27 @@ +# Hyper Solvation Example + +Demonstrates how to run `hyper.hyper_solvate_sumo()` with rush-py and handle batched per-item outputs. + +## Quick Start + +```bash +export RUSH_TOKEN="your-token" +export RUSH_PROJECT="your-project" + +python 08_hyper_solvate.py +``` + +## What This Example Covers + +1. Submitting a Hyper solvation run with explicit `HyperConfig` +2. Collecting and fetching a `TRCBatchResultRef` +3. Handling `ItemError` per batch item +4. Saving successful TRC output to the Rush workspace + +## Input Data + +- `data/valid_trc.json` — Minimal TRC structure used as the solvation input + +## Tutorial + +See the full tutorial: [Hyper Solvation](../../docs/tutorials/08-hyper-solvate.md) diff --git a/examples/hyper-solvate/data/valid_trc.json b/examples/hyper-solvate/data/valid_trc.json new file mode 100644 index 00000000..65ae683a --- /dev/null +++ b/examples/hyper-solvate/data/valid_trc.json @@ -0,0 +1,32 @@ +{ + "topology": { + "schema_version": "0.2.0", + "symbols": ["O", "H", "H"], + "geometry": [0.0, 0.0, 0.0, 0.96, 0.0, 0.0, -0.24, 0.93, 0.0], + "labels": null, + "partial_charges": null, + "formal_charges": null, + "connectivity": null, + "stereochemistry": null, + "velocities": null, + "fragments": null, + "fragment_formal_charges": null, + "fragment_partial_charges": null, + "fragment_multiplicities": null + }, + "residues": { + "residues": [[0, 1, 2]], + "seqs": ["HOH"], + "seq_ns": [1], + "insertion_codes": [""], + "labeled": null, + "labels": null + }, + "chains": { + "chains": [[0]], + "alpha_helices": null, + "beta_sheets": null, + "labeled": null, + "labels": null + } +} diff --git a/examples/hyper-solvate_sumo/08_hyper_solvate_sumo.py b/examples/hyper-solvate_sumo/08_hyper_solvate_sumo.py new file mode 100644 index 00000000..847410e8 --- /dev/null +++ b/examples/hyper-solvate_sumo/08_hyper_solvate_sumo.py @@ -0,0 +1,35 @@ +"""Example: Hyper solvation workflow.""" + +from pathlib import Path + +from rush import RunOpts, TRC, hyper + +DATA_DIR = Path(__file__).parent / "data" +INPUT_TRC = DATA_DIR / "valid_trc.json" + +run = hyper.hyper_solvate_sumo( + [INPUT_TRC], + config=hyper.HyperConfig(max_inputs=8, padding_nm=0.8, seed=12345, timeout_seconds=120), + run_opts=RunOpts(name="Example: Hyper Solvate", tags=["rush-py", "example", "hyper", "solvate"]), +) + +result_ref = run.collect() +fetched = result_ref.fetch() +if len(fetched) != 1: + raise RuntimeError(f"Expected 1 output item, got {len(fetched)}") + +item = fetched[0] +if not isinstance(item, TRC): + raise RuntimeError(f"Expected TRC output, got {item}") + +print("Solvated atom count:", len(item.topology.symbols)) + +saved = result_ref.save() +if len(saved) != 1: + raise RuntimeError(f"Expected 1 saved item, got {len(saved)}") + +saved_item = saved[0] +if isinstance(saved_item, hyper.ItemError): + raise RuntimeError(f"Unexpected per-item save error: {saved_item}") + +print("Saved output:", saved_item) diff --git a/examples/hyper-solvate_sumo/README.md b/examples/hyper-solvate_sumo/README.md new file mode 100644 index 00000000..b3a95d36 --- /dev/null +++ b/examples/hyper-solvate_sumo/README.md @@ -0,0 +1,26 @@ +# Hyper Solvate Sumo Example + +Demonstrates `hyper.hyper_solvate_sumo()` on a small TRC input and validates successful output handling. + +## Quick Start + +```bash +export RUSH_TOKEN="your-token" +export RUSH_PROJECT="your-project" + +python 08_hyper_solvate_sumo.py +``` + +## What This Example Covers + +1. Submitting a Hyper solvation run with explicit `HyperConfig` +2. Fetching parsed output as `TRC` +3. Saving the solvated artifact to the Rush workspace + +## Input Data + +- `data/valid_trc.json` + +## Tutorial + +See: [Tutorial 8: Hyper Solvate Sumo](../../docs/tutorials/08-hyper-solvate_sumo.md) diff --git a/examples/hyper-solvate_sumo/data/valid_trc.json b/examples/hyper-solvate_sumo/data/valid_trc.json new file mode 100644 index 00000000..65ae683a --- /dev/null +++ b/examples/hyper-solvate_sumo/data/valid_trc.json @@ -0,0 +1,32 @@ +{ + "topology": { + "schema_version": "0.2.0", + "symbols": ["O", "H", "H"], + "geometry": [0.0, 0.0, 0.0, 0.96, 0.0, 0.0, -0.24, 0.93, 0.0], + "labels": null, + "partial_charges": null, + "formal_charges": null, + "connectivity": null, + "stereochemistry": null, + "velocities": null, + "fragments": null, + "fragment_formal_charges": null, + "fragment_partial_charges": null, + "fragment_multiplicities": null + }, + "residues": { + "residues": [[0, 1, 2]], + "seqs": ["HOH"], + "seq_ns": [1], + "insertion_codes": [""], + "labeled": null, + "labels": null + }, + "chains": { + "chains": [[0]], + "alpha_helices": null, + "beta_sheets": null, + "labeled": null, + "labels": null + } +} diff --git a/src/rush/hyper/__init__.py b/src/rush/hyper/__init__.py new file mode 100644 index 00000000..5d0c5324 --- /dev/null +++ b/src/rush/hyper/__init__.py @@ -0,0 +1,46 @@ +"""Hyper module for the Rush Python client.""" + +# --- Shared result error type --- +from ._common import ItemError + +# --- Solvation --- +from ._hyper_solvate_sumo import HyperConfig, TRCBatchResultRef, hyper_solvate_sumo + +# --- Minimization --- +from ._hyper_minimize_sumo import ( + HyperMinimizeConfig, + MinimizeInput, + hyper_minimize_sumo, +) + +# --- Molecular dynamics run --- +from ._hyper_run_sumo import ( + HyperRunConfig, + RunInput, + RunOutput, + RunOutputPaths, + RunOutputRef, + RunResultRef, + hyper_run_sumo, +) + +__all__ = [ + # Shared + "ItemError", + # Solvation + "HyperConfig", + "TRCBatchResultRef", + "hyper_solvate_sumo", + # Minimization + "HyperMinimizeConfig", + "MinimizeInput", + "hyper_minimize_sumo", + # Molecular dynamics run + "HyperRunConfig", + "RunInput", + "RunOutput", + "RunOutputRef", + "RunOutputPaths", + "RunResultRef", + "hyper_run_sumo", +] diff --git a/src/rush/hyper/_common.py b/src/rush/hyper/_common.py new file mode 100644 index 00000000..7b996f88 --- /dev/null +++ b/src/rush/hyper/_common.py @@ -0,0 +1,145 @@ +"""Shared parsing and upload helpers for Hyper entrypoints.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Literal, TypeVar + +from ..convert import from_json, to_dict +from ..mol import TRC +from ..objects import RushObject, upload_object + +JsonObjectInput = Path | str | RushObject | dict[str, Any] +TRCInput = TRC | Path | str | RushObject | dict[str, Any] + +ErrorStage = Literal["InputDecode", "Execution", "OutputParse"] +ErrorCategory = Literal["InvalidInput", "ToolInput", "OutputFormat"] + +TSuccessRef = TypeVar("TSuccessRef") + + +@dataclass(frozen=True) +class ItemError: + """Per-item error returned by Hyper batch wrappers.""" + + stage: ErrorStage + category: ErrorCategory + message: str + input_index: int + + @classmethod + def from_raw_output(cls, raw: Any) -> "ItemError": + if not isinstance(raw, dict): + raise ValueError(f"Expected ItemError object, got {type(raw).__name__}") + return cls( + stage=raw["stage"], + category=raw["category"], + message=raw["message"], + input_index=int(raw["input_index"]), + ) + + +def _is_result_type(raw: Any) -> bool: + return isinstance(raw, dict) and len(raw) == 1 and ("Ok" in raw or "Err" in raw) + + +def _format_user_error(err: Any) -> str: + if isinstance(err, str): + return err + if isinstance(err, dict) and len(err) == 1: + key, value = next(iter(err.items())) + if value is None: + return str(key) + if isinstance(value, dict): + details = ", ".join(f"{k}={v}" for k, v in value.items()) + return f"{key}({details})" + return f"{key}({value})" + return json.dumps(err) + + +def _parse_batch_outputs( + raw: Any, + on_success: Callable[[Any], TSuccessRef], + label: str, +) -> list[TSuccessRef | ItemError]: + if not isinstance(raw, list) or len(raw) != 1: + raise ValueError( + f"{label} should return a single-element list, got {type(raw).__name__}" + ) + + payload = raw[0] + if _is_result_type(payload): + if "Err" in payload: + raise ValueError(f"{label} top-level error: {_format_user_error(payload['Err'])}") + payload = payload["Ok"] + + items = payload if isinstance(payload, list) else [payload] + + parsed: list[TSuccessRef | ItemError] = [] + for item in items: + if _is_result_type(item): + if "Err" in item: + parsed.append(ItemError.from_raw_output(item["Err"])) + continue + item = item["Ok"] + parsed.append(on_success(item)) + return parsed + + +def _upload_json_object(input_object: JsonObjectInput) -> RushObject: + match input_object: + case RushObject(): + return input_object + case Path() | str() | dict(): + return RushObject.from_dict(upload_object(input_object)) + case _: + raise TypeError( + "Expected Path | str | RushObject | dict input for Hyper JSON object" + ) + + +def _upload_trc_object(input_object: TRCInput) -> RushObject: + match input_object: + case RushObject(): + return input_object + case TRC(): + trc_dict = to_dict(input_object) + if not isinstance(trc_dict, dict): + raise TypeError("Expected single TRC object") + return RushObject.from_dict(upload_object(trc_dict)) + case Path() | str() | dict(): + return RushObject.from_dict(upload_object(input_object)) + case _: + raise TypeError( + "Expected TRC | Path | str | RushObject | dict input for Hyper TRC object" + ) + + +def _fetch_trc(obj: RushObject) -> TRC: + parsed = from_json(obj.fetch_json()) + if isinstance(parsed, list): + if len(parsed) != 1: + raise ValueError(f"Expected one TRC object, got {len(parsed)}") + item = parsed[0] + if not isinstance(item, TRC): + raise TypeError("Expected TRC item in parsed list") + return item + if not isinstance(parsed, TRC): + raise TypeError(f"Expected TRC output, got {type(parsed).__name__}") + return parsed + + +def _to_rex_json_obj(obj: RushObject) -> str: + return ( + 'VirtualObject { path = "' + + str(obj.path) + + '", format = ObjectFormat::json, size = 0 }' + ) + + +def _format_rex_list(items: list[str]) -> str: + if not items: + return "[]" + return "[\n " + ",\n ".join(items) + "\n ]" diff --git a/src/rush/hyper/_hyper_minimize_sumo.py b/src/rush/hyper/_hyper_minimize_sumo.py new file mode 100644 index 00000000..21bfb5af --- /dev/null +++ b/src/rush/hyper/_hyper_minimize_sumo.py @@ -0,0 +1,97 @@ +"""Hyper minimization wrapper for the Rush Python client.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass +from string import Template + +from gql.transport.exceptions import TransportQueryError + +from .._rex import optional_str +from ..runs import Run, RunOpts, RunSpec +from ..session import _submit_rex +from ._common import ( + JsonObjectInput, + TRCInput, + _format_rex_list, + _to_rex_json_obj, + _upload_json_object, + _upload_trc_object, +) +from ._hyper_solvate_sumo import TRCBatchResultRef + + +@dataclass(frozen=True) +class HyperMinimizeConfig: + """Config for :func:`hyper_minimize_sumo`.""" + + max_inputs: int | None = None + steps: int | None = None + gtol: float | None = None + timeout_seconds: int | None = None + + +@dataclass(frozen=True) +class MinimizeInput: + """Input item for :func:`hyper_minimize_sumo`.""" + + structure: TRCInput + topology: JsonObjectInput + + +def _to_rex_minimize_config(config: HyperMinimizeConfig | None) -> str: + if config is None: + return "None" + return Template( + """Some (hyper_minimize_sumo::HyperMinimizeConfig { + max_inputs = $max_inputs, + steps = $steps, + gtol = $gtol, + timeout_seconds = $timeout_seconds, + })""" + ).substitute( + max_inputs=optional_str(config.max_inputs), + steps=optional_str(config.steps), + gtol=optional_str(config.gtol), + timeout_seconds=optional_str(config.timeout_seconds), + ) + + +def hyper_minimize_sumo( + jobs: list[MinimizeInput], + config: HyperMinimizeConfig | None = None, + run_spec: RunSpec = RunSpec(target="Bullet"), + run_opts: RunOpts = RunOpts(), +) -> Run[TRCBatchResultRef]: + """Submit Hyper minimization for one or more structures.""" + job_exprs: list[str] = [] + for job in jobs: + structure = _to_rex_json_obj(_upload_trc_object(job.structure)) + topology = _to_rex_json_obj(_upload_json_object(job.topology)) + job_exprs.append( + "(" + + "hyper_minimize_sumo::MinimizeInput { " + + f"structure = {structure}, topology = {topology} " + + "}" + + ")" + ) + + rex = Template( + """hyper_minimize_sumo_s + ($run_spec) + ($config) + $jobs""" + ).substitute( + run_spec=run_spec._to_rex(), + config=_to_rex_minimize_config(config), + jobs=_format_rex_list(job_exprs), + ) + + try: + return Run(_submit_rex(rex, run_opts), TRCBatchResultRef) + except TransportQueryError as e: + if e.errors: + for error in e.errors: + print(f"Error: {error['message']}", file=sys.stderr) + raise diff --git a/src/rush/hyper/_hyper_run_sumo.py b/src/rush/hyper/_hyper_run_sumo.py new file mode 100644 index 00000000..191e5ee5 --- /dev/null +++ b/src/rush/hyper/_hyper_run_sumo.py @@ -0,0 +1,255 @@ +"""Hyper molecular dynamics wrapper for the Rush Python client.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass +from pathlib import Path +from string import Template +from typing import Any, Literal, Self + +from gql.transport.exceptions import TransportQueryError + +from .._rex import optional_str +from ..objects import RushObject +from ..runs import Run, RunOpts, RunSpec +from ..session import _submit_rex +from ._common import ( + ItemError, + JsonObjectInput, + TRCInput, + _format_rex_list, + _parse_batch_outputs, + _to_rex_json_obj, + _upload_json_object, + _upload_trc_object, +) + + +@dataclass(frozen=True) +class HyperRunConfig: + """Config for :func:`hyper_run_sumo`.""" + + max_inputs: int | None = None + nsteps: int | None = None + dt_ps: float | None = None + temperature_k: float | None = None + ensemble: Literal["Nve", "Nvt", "Npt"] | None = None + minimize_before_run: bool | None = None + solvate_before_run: bool | None = None + use_gpu: bool | None = None + nthreads: int | None = None + timeout_seconds: int | None = None + + +@dataclass(frozen=True) +class RunInput: + """Input item for :func:`hyper_run_sumo`.""" + + sim_config: JsonObjectInput + topology: JsonObjectInput + coordinates: TRCInput + + +@dataclass(frozen=True) +class RunOutputRef: + """Reference to one successful run output.""" + + trajectory: RushObject + checkpoint: RushObject | None + + +@dataclass(frozen=True) +class RunOutput: + """Fetched bytes for one successful run output.""" + + trajectory: bytes + checkpoint: bytes | None + + +@dataclass(frozen=True) +class RunOutputPaths: + """Workspace paths for one saved run output.""" + + trajectory: Path + checkpoint: Path | None + + +@dataclass(frozen=True) +class RunResultRef: + """Result reference for :func:`hyper_run_sumo`.""" + + items: list[RunOutputRef | ItemError] + + @classmethod + def from_raw_output(cls, raw: Any) -> Self: + parsed = _parse_batch_outputs(raw, _parse_run_item, "hyper run batch") + return cls(items=parsed) + + def __getitem__(self, index: int) -> RunOutputRef | ItemError: + return self.items[index] + + def __len__(self) -> int: + return len(self.items) + + def fetch(self) -> list[RunOutput | ItemError]: + out: list[RunOutput | ItemError] = [] + for item in self.items: + if isinstance(item, ItemError): + out.append(item) + continue + trajectory = _fetch_run_artifact_bytes(item.trajectory, "trajectory") + checkpoint = ( + _fetch_run_artifact_bytes(item.checkpoint, "checkpoint") + if item.checkpoint is not None + else None + ) + out.append(RunOutput(trajectory=trajectory, checkpoint=checkpoint)) + return out + + def save(self) -> list[RunOutputPaths | ItemError]: + out: list[RunOutputPaths | ItemError] = [] + for item in self.items: + if isinstance(item, ItemError): + out.append(item) + continue + out.append( + RunOutputPaths( + trajectory=_save_run_artifact( + item.trajectory, + ext="xtc", + label="trajectory", + ), + checkpoint=( + _save_run_artifact( + item.checkpoint, + ext="bin", + label="checkpoint", + ) + if item.checkpoint is not None + else None + ), + ) + ) + return out + + +def _fetch_run_artifact_bytes(obj: RushObject, label: str) -> bytes: + if obj.format.lower() == "bin": + return obj.fetch_bytes() + + if obj.format.lower() != "json": + raise TypeError( + f"hyper_run_sumo {label} object has unsupported format {obj.format!r}" + ) + + payload = obj.fetch_list() + try: + return bytes(payload) + except (TypeError, ValueError) as exc: + raise TypeError( + f"hyper_run_sumo {label} JSON output must be a list of byte values" + ) from exc + + +def _save_run_artifact(obj: RushObject, ext: str, label: str) -> Path: + out_path = obj.save(ext=ext) + if obj.format.lower() == "json": + payload = _fetch_run_artifact_bytes(obj, label) + with out_path.open("wb") as out_file: + out_file.write(payload) + return out_path + + +def _parse_run_item(raw: Any) -> RunOutputRef: + if not isinstance(raw, dict): + raise ValueError(f"Expected RunOutput object, got {type(raw).__name__}") + + trajectory = RushObject.from_dict(raw["trajectory"]) + checkpoint_raw = raw.get("checkpoint") + checkpoint = RushObject.from_dict(checkpoint_raw) if checkpoint_raw is not None else None + return RunOutputRef(trajectory=trajectory, checkpoint=checkpoint) + + +def _to_rex_run_ensemble(value: Literal["Nve", "Nvt", "Npt"] | None) -> str: + if value is None: + return "None" + variants = { + "Nve": "hyper_run_sumo::RunEnsemble::Nve", + "Nvt": "hyper_run_sumo::RunEnsemble::Nvt", + "Npt": "hyper_run_sumo::RunEnsemble::Npt", + } + return f"Some {variants[value]}" + + +def _to_rex_run_config(config: HyperRunConfig | None) -> str: + if config is None: + return "None" + return Template( + """Some (hyper_run_sumo::HyperRunConfig { + max_inputs = $max_inputs, + nsteps = $nsteps, + dt_ps = $dt_ps, + temperature_k = $temperature_k, + ensemble = $ensemble, + minimize_before_run = $minimize_before_run, + solvate_before_run = $solvate_before_run, + use_gpu = $use_gpu, + nthreads = $nthreads, + timeout_seconds = $timeout_seconds, + })""" + ).substitute( + max_inputs=optional_str(config.max_inputs), + nsteps=optional_str(config.nsteps), + dt_ps=optional_str(config.dt_ps), + temperature_k=optional_str(config.temperature_k), + ensemble=_to_rex_run_ensemble(config.ensemble), + minimize_before_run=optional_str(config.minimize_before_run), + solvate_before_run=optional_str(config.solvate_before_run), + use_gpu=optional_str(config.use_gpu), + nthreads=optional_str(config.nthreads), + timeout_seconds=optional_str(config.timeout_seconds), + ) + + +def hyper_run_sumo( + jobs: list[RunInput], + config: HyperRunConfig | None = None, + run_spec: RunSpec = RunSpec(target="Bullet"), + run_opts: RunOpts = RunOpts(), +) -> Run[RunResultRef]: + """Submit Hyper molecular dynamics runs for one or more jobs.""" + job_exprs: list[str] = [] + for job in jobs: + sim_config = _to_rex_json_obj(_upload_json_object(job.sim_config)) + topology = _to_rex_json_obj(_upload_json_object(job.topology)) + coordinates = _to_rex_json_obj(_upload_trc_object(job.coordinates)) + job_exprs.append( + "(" + + "hyper_run_sumo::RunInput { " + + ( + f"sim_config = {sim_config}, topology = {topology}, " + f"coordinates = {coordinates} " + ) + + "}" + + ")" + ) + + rex = Template( + """hyper_run_sumo_s + ($run_spec) + ($config) + $jobs""" + ).substitute( + run_spec=run_spec._to_rex(), + config=_to_rex_run_config(config), + jobs=_format_rex_list(job_exprs), + ) + + try: + return Run(_submit_rex(rex, run_opts), RunResultRef) + except TransportQueryError as e: + if e.errors: + for error in e.errors: + print(f"Error: {error['message']}", file=sys.stderr) + raise diff --git a/src/rush/hyper/_hyper_solvate_sumo.py b/src/rush/hyper/_hyper_solvate_sumo.py new file mode 100644 index 00000000..dbe29355 --- /dev/null +++ b/src/rush/hyper/_hyper_solvate_sumo.py @@ -0,0 +1,121 @@ +"""Hyper solvation wrapper for the Rush Python client.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass +from pathlib import Path +from string import Template +from typing import Any, Self + +from gql.transport.exceptions import TransportQueryError + +from .._rex import optional_str +from ..mol import TRC +from ..objects import RushObject +from ..runs import Run, RunOpts, RunSpec +from ..session import _submit_rex +from ._common import ( + ItemError, + TRCInput, + _fetch_trc, + _format_rex_list, + _parse_batch_outputs, + _to_rex_json_obj, + _upload_trc_object, +) + + +@dataclass(frozen=True) +class HyperConfig: + """Config for :func:`hyper_solvate_sumo`.""" + + max_inputs: int | None = None + padding_nm: float | None = None + seed: int | None = None + timeout_seconds: int | None = None + + +@dataclass(frozen=True) +class TRCBatchResultRef: + """Result reference for TRC-producing Hyper batch entrypoints.""" + + items: list[RushObject | ItemError] + + @classmethod + def from_raw_output(cls, raw: Any) -> Self: + parsed = _parse_batch_outputs(raw, _parse_trc_item, "hyper TRC batch") + return cls(items=parsed) + + def __getitem__(self, index: int) -> RushObject | ItemError: + return self.items[index] + + def __len__(self) -> int: + return len(self.items) + + def fetch(self) -> list[TRC | ItemError]: + out: list[TRC | ItemError] = [] + for item in self.items: + if isinstance(item, ItemError): + out.append(item) + else: + out.append(_fetch_trc(item)) + return out + + def save(self) -> list[Path | ItemError]: + return [ + item if isinstance(item, ItemError) else item.save(ext="json") + for item in self.items + ] + + +def _parse_trc_item(raw: Any) -> RushObject: + if not isinstance(raw, dict): + raise ValueError(f"Expected TRC output object, got {type(raw).__name__}") + return RushObject.from_dict(raw) + +def _to_rex_solvate_config(config: HyperConfig | None) -> str: + if config is None: + return "None" + return Template( + """Some (hyper_solvate_sumo::HyperConfig { + max_inputs = $max_inputs, + padding_nm = $padding_nm, + seed = $seed, + timeout_seconds = $timeout_seconds, + })""" + ).substitute( + max_inputs=optional_str(config.max_inputs), + padding_nm=optional_str(config.padding_nm), + seed=optional_str(config.seed), + timeout_seconds=optional_str(config.timeout_seconds), + ) + + +def hyper_solvate_sumo( + input_trcs: list[TRCInput], + config: HyperConfig | None = None, + run_spec: RunSpec = RunSpec(target="Bullet"), + run_opts: RunOpts = RunOpts(), +) -> Run[TRCBatchResultRef]: + """Submit Hyper solvation for one or more TRC inputs.""" + input_exprs = [_to_rex_json_obj(_upload_trc_object(item)) for item in input_trcs] + + rex = Template( + """hyper_solvate_sumo_s + ($run_spec) + ($config) + $inputs""" + ).substitute( + run_spec=run_spec._to_rex(), + config=_to_rex_solvate_config(config), + inputs=_format_rex_list(input_exprs), + ) + + try: + return Run(_submit_rex(rex, run_opts), TRCBatchResultRef) + except TransportQueryError as e: + if e.errors: + for error in e.errors: + print(f"Error: {error['message']}", file=sys.stderr) + raise diff --git a/src/rush/session.py b/src/rush/session.py index 24ea7e98..3d6ad016 100644 --- a/src/rush/session.py +++ b/src/rush/session.py @@ -98,6 +98,9 @@ def _get_project_id() -> str: "exess_rex": "github:talo/tengu-exess/133781d71c493900a82121729c18994b4a184197#exess_rex", "exess_geo_opt_rex": "github:talo/tengu-exess/133781d71c493900a82121729c18994b4a184197#exess_geo_opt_rex", "exess_qmmm_rex": "github:talo/tengu-exess/133781d71c493900a82121729c18994b4a184197#exess_qmmm_rex", + "hyper_minimize_sumo": "github:talo/tengu-hyper/f2906f085c7a3a07104a98a263dc9876f08ee4e7#hyper_minimize_sumo", + "hyper_run_sumo": "github:talo/tengu-hyper/f2906f085c7a3a07104a98a263dc9876f08ee4e7#hyper_run_sumo", + "hyper_solvate_sumo": "github:talo/tengu-hyper/f2906f085c7a3a07104a98a263dc9876f08ee4e7#hyper_solvate_sumo", "mmseqs2_rex": "github:talo/tengu-colabfold/749a096d082efdac3ac13de4aaa98aee3347d79d#mmseqs2_rex", "nnxtb_rex": "github:talo/tengu-nnxtb/4e733660264d38faab5d23eadc41ca86fd6ff97a#nnxtb_rex", "pbsa_rex": "github:talo/pbsa-cuda/f8b1c357fddfebf7e0c51a84f8d4e70958440c00#pbsa_rex", @@ -111,6 +114,9 @@ def _get_project_id() -> str: "exess_rex": "github:talo/tengu-exess/133781d71c493900a82121729c18994b4a184197#exess_rex", "exess_geo_opt_rex": "github:talo/tengu-exess/133781d71c493900a82121729c18994b4a184197#exess_geo_opt_rex", "exess_qmmm_rex": "github:talo/tengu-exess/133781d71c493900a82121729c18994b4a184197#exess_qmmm_rex", + "hyper_minimize_sumo": "github:talo/tengu-hyper/f2906f085c7a3a07104a98a263dc9876f08ee4e7#hyper_minimize_sumo", + "hyper_run_sumo": "github:talo/tengu-hyper/f2906f085c7a3a07104a98a263dc9876f08ee4e7#hyper_run_sumo", + "hyper_solvate_sumo": "github:talo/tengu-hyper/f2906f085c7a3a07104a98a263dc9876f08ee4e7#hyper_solvate_sumo", "mmseqs2_rex": "github:talo/tengu-colabfold/0b6ca8b9dc97fc6380d334169a6faae51d85fac7#mmseqs2_rex", "nnxtb_rex": "github:talo/tengu-nnxtb/4e733660264d38faab5d23eadc41ca86fd6ff97a#nnxtb_rex", "pbsa_rex": "github:talo/pbsa-cuda/f8b1c357fddfebf7e0c51a84f8d4e70958440c00#pbsa_rex", diff --git a/tests/conftest.py b/tests/conftest.py index 1d2ac797..60347381 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,6 +21,7 @@ "tests/module_output_helpers/test_mmseqs2_output_helpers.py", "tests/module_output_helpers/test_nnxtb_output_helpers.py", "tests/module_output_helpers/test_pbsa_output_helpers.py", + "tests/module_output_helpers/test_hyper_output_helpers.py", "tests/module_output_helpers/test_prepare_output_helpers.py", "tests/test_client_collect_run.py", "tests/test_convert_mmcif.py", diff --git a/tests/data/hyper/methanol_topology.json b/tests/data/hyper/methanol_topology.json new file mode 100644 index 00000000..5d9c6f83 --- /dev/null +++ b/tests/data/hyper/methanol_topology.json @@ -0,0 +1,333 @@ +{ + "atom_type_count": 4, + "charges": [ + 0.28, + -0.68, + 0.0, + 0.0, + 0.0, + 0.4 + ], + "masses": [ + 12.011, + 15.999, + 1.008, + 1.008, + 1.008, + 1.008 + ], + "atom_types": [ + 0, + 1, + 2, + 2, + 2, + 3 + ], + "bonds": [ + { + "i": 0, + "j": 1, + "b0": 0.14179999999999998, + "kb": 303937.193826 + }, + { + "i": 0, + "j": 2, + "b0": 0.1093, + "kb": 287014.992228 + }, + { + "i": 0, + "j": 3, + "b0": 0.1093, + "kb": 287014.992228 + }, + { + "i": 0, + "j": 4, + "b0": 0.1093, + "kb": 287014.992228 + }, + { + "i": 1, + "j": 5, + "b0": 0.0972, + "kb": 469365.2642520001 + } + ], + "angles": [ + { + "i": 1, + "j": 0, + "k": 2, + "theta0": 1.895026141937883, + "ktheta": 470.32508605816537 + }, + { + "i": 1, + "j": 0, + "k": 3, + "theta0": 1.895026141937883, + "ktheta": 470.32508605816537 + }, + { + "i": 1, + "j": 0, + "k": 4, + "theta0": 1.895026141937883, + "ktheta": 470.32508605816537 + }, + { + "i": 2, + "j": 0, + "k": 3, + "theta0": 1.8995465447005484, + "ktheta": 310.7397495595562 + }, + { + "i": 2, + "j": 0, + "k": 4, + "theta0": 1.8995465447005484, + "ktheta": 310.7397495595562 + }, + { + "i": 3, + "j": 0, + "k": 4, + "theta0": 1.8995465447005484, + "ktheta": 310.7397495595562 + }, + { + "i": 0, + "j": 1, + "k": 5, + "theta0": 1.8588280132515207, + "ktheta": 477.55159186187603 + } + ], + "proper_dihedrals": [ + { + "i": 2, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 1.246832, + "mult": 1 + }, + { + "i": 2, + "j": 0, + "k": 1, + "l": 5, + "phi0": 3.141592653589793, + "kphi": -0.577392, + "mult": 2 + }, + { + "i": 2, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 0.7238319999999999, + "mult": 3 + }, + { + "i": 3, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 1.246832, + "mult": 1 + }, + { + "i": 3, + "j": 0, + "k": 1, + "l": 5, + "phi0": 3.141592653589793, + "kphi": -0.577392, + "mult": 2 + }, + { + "i": 3, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 0.7238319999999999, + "mult": 3 + }, + { + "i": 4, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 1.246832, + "mult": 1 + }, + { + "i": 4, + "j": 0, + "k": 1, + "l": 5, + "phi0": 3.141592653589793, + "kphi": -0.577392, + "mult": 2 + }, + { + "i": 4, + "j": 0, + "k": 1, + "l": 5, + "phi0": 0.0, + "kphi": 0.7238319999999999, + "mult": 3 + } + ], + "improper_dihedrals": [], + "pairs14": [ + { + "i": 2, + "j": 5, + "c6": 4.274187696047065e-05, + "c12": 1.0119320214366275e-08 + }, + { + "i": 3, + "j": 5, + "c6": 4.274187696047065e-05, + "c12": 1.0119320214366275e-08 + }, + { + "i": 4, + "j": 5, + "c6": 4.274187696047065e-05, + "c12": 1.0119320214366275e-08 + } + ], + "cmap_entries": [], + "lj_params": [ + { + "c6": 0.002115016015813491, + "c12": 3.94244116886151e-06 + }, + { + "c6": 0.0016677018383901895, + "c12": 2.311249945030495e-06 + }, + { + "c6": 0.0005432157553294566, + "c12": 4.6102104135361286e-07 + }, + { + "c6": 0.0003954334320525218, + "c12": 2.4429962496901167e-07 + }, + { + "c6": 0.0016677018383901895, + "c12": 2.311249945030495e-06 + }, + { + "c6": 0.00129488850377234, + "c12": 1.3138538180583043e-06 + }, + { + "c6": 0.0004104134700407302, + "c12": 2.481369517052484e-07 + }, + { + "c6": 0.00029313999406590564, + "c12": 1.2658970268056165e-07 + }, + { + "c6": 0.0005432157553294566, + "c12": 4.6102104135361286e-07 + }, + { + "c6": 0.0004104134700407302, + "c12": 2.481369517052484e-07 + }, + { + "c6": 0.00012386870467525756, + "c12": 4.249495325323946e-08 + }, + { + "c6": 8.54837539209413e-05, + "c12": 2.023864042873255e-08 + }, + { + "c6": 0.0003954334320525218, + "c12": 2.4429962496901167e-07 + }, + { + "c6": 0.00029313999406590564, + "c12": 1.2658970268056165e-07 + }, + { + "c6": 8.54837539209413e-05, + "c12": 2.023864042873255e-08 + }, + { + "c6": 5.756897163937681e-05, + "c12": 9.17890990287509e-09 + } + ], + "cmap_grid_size": 24, + "cmap_grids": [], + "exclusions": [ + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ], + [ + 0, + 1, + 2, + 3, + 4, + 5 + ] + ] +} diff --git a/tests/data/hyper/methanol_trc.json b/tests/data/hyper/methanol_trc.json new file mode 100644 index 00000000..d753464a --- /dev/null +++ b/tests/data/hyper/methanol_trc.json @@ -0,0 +1,39 @@ +{ + "topology": { + "schema_version": "0.2.0", + "symbols": ["C", "O", "H", "H", "H", "H"], + "geometry": [ + -0.37, 0.01, -0.01, + 0.92, -0.5, -0.3, + -0.52, 0.03, 1.07, + -0.47, 1.01, -0.42, + -1.12, -0.64, -0.47, + 1.56, 0.09, 0.12 + ], + "labels": null, + "partial_charges": null, + "formal_charges": null, + "connectivity": null, + "stereochemistry": null, + "velocities": null, + "fragments": null, + "fragment_formal_charges": null, + "fragment_partial_charges": null, + "fragment_multiplicities": null + }, + "residues": { + "residues": [[0, 1, 2, 3, 4, 5]], + "seqs": ["LIG"], + "seq_ns": [1], + "insertion_codes": [""], + "labeled": null, + "labels": null + }, + "chains": { + "chains": [[0]], + "alpha_helices": null, + "beta_sheets": null, + "labeled": null, + "labels": null + } +} diff --git a/tests/data/hyper/sim_config.json b/tests/data/hyper/sim_config.json new file mode 100644 index 00000000..bab0d9be --- /dev/null +++ b/tests/data/hyper/sim_config.json @@ -0,0 +1 @@ +{"force_field":"AMBERFF19SB","ensemble":"NVT","n_steps":40,"dt":0.001,"temperature":300.0,"initial_temp":300.0,"tau_t":0.1,"electrostatics":"CUTOFF","cutoff":0.0,"dispersion_correction":false,"constraints":"NONE","box_x":5.0,"box_y":5.0,"box_z":5.0,"output":{"trajectory_freq":1,"energy_freq":1,"checkpoint_freq":20,"log_freq":10}} \ No newline at end of file diff --git a/tests/data/hyper/valid_trc.json b/tests/data/hyper/valid_trc.json new file mode 100644 index 00000000..65ae683a --- /dev/null +++ b/tests/data/hyper/valid_trc.json @@ -0,0 +1,32 @@ +{ + "topology": { + "schema_version": "0.2.0", + "symbols": ["O", "H", "H"], + "geometry": [0.0, 0.0, 0.0, 0.96, 0.0, 0.0, -0.24, 0.93, 0.0], + "labels": null, + "partial_charges": null, + "formal_charges": null, + "connectivity": null, + "stereochemistry": null, + "velocities": null, + "fragments": null, + "fragment_formal_charges": null, + "fragment_partial_charges": null, + "fragment_multiplicities": null + }, + "residues": { + "residues": [[0, 1, 2]], + "seqs": ["HOH"], + "seq_ns": [1], + "insertion_codes": [""], + "labeled": null, + "labels": null + }, + "chains": { + "chains": [[0]], + "alpha_helices": null, + "beta_sheets": null, + "labeled": null, + "labels": null + } +} diff --git a/tests/module_output_helpers/test_hyper_output_helpers.py b/tests/module_output_helpers/test_hyper_output_helpers.py new file mode 100644 index 00000000..3c11eaf3 --- /dev/null +++ b/tests/module_output_helpers/test_hyper_output_helpers.py @@ -0,0 +1,99 @@ +import json +from pathlib import Path + +from rush import TRC +from rush import hyper + + +def test_trc_batch_result_ref_fetch_and_save(monkeypatch): + fixture_path = Path(__file__).parent.parent / "data" / "hyper" / "valid_trc.json" + trc_payload = fixture_path.read_text() + + monkeypatch.setattr("rush.objects.RushObject.fetch_json", lambda _self: json.loads(trc_payload)) + monkeypatch.setattr( + "rush.objects.RushObject.save", + lambda self, ext="json", **_kw: Path("workspace") / f"{self.path}.{ext}", + ) + + ref = hyper.TRCBatchResultRef.from_raw_output( + [ + { + "Ok": [ + {"Ok": {"path": "solvated", "size": 0, "format": "Json"}}, + { + "Err": { + "stage": "Execution", + "category": "ToolInput", + "message": "bad input", + "input_index": 1, + } + }, + ] + } + ] + ) + + fetched = ref.fetch() + assert isinstance(fetched[0], TRC) + assert isinstance(fetched[1], hyper.ItemError) + assert fetched[1].input_index == 1 + + saved = ref.save() + assert saved[0] == Path("workspace/solvated.json") + assert isinstance(saved[1], hyper.ItemError) + + +def test_run_result_ref_fetch_and_save(monkeypatch): + monkeypatch.setattr("rush.objects.RushObject.fetch_bytes", lambda _self: b"binary-data") + monkeypatch.setattr( + "rush.objects.RushObject.save", + lambda self, ext="bin", **_kw: Path("workspace") / f"{self.path}.{ext}", + ) + + ref = hyper.RunResultRef.from_raw_output( + [ + { + "Ok": [ + { + "Ok": { + "trajectory": {"path": "traj", "size": 0, "format": "Bin"}, + "checkpoint": { + "path": "checkpoint", + "size": 0, + "format": "Bin", + }, + } + }, + { + "Err": { + "stage": "OutputParse", + "category": "OutputFormat", + "message": "empty output", + "input_index": 1, + } + }, + ] + } + ] + ) + + fetched = ref.fetch() + assert fetched[0] == hyper.RunOutput(trajectory=b"binary-data", checkpoint=b"binary-data") + assert isinstance(fetched[1], hyper.ItemError) + + saved = ref.save() + assert saved[0] == hyper.RunOutputPaths( + trajectory=Path("workspace/traj.xtc"), + checkpoint=Path("workspace/checkpoint.bin"), + ) + assert isinstance(saved[1], hyper.ItemError) + + +def test_batch_result_ref_rejects_top_level_error(): + with_error = [{"Err": {"TooManyInputs": {"count": 200, "max": 128}}}] + try: + hyper.TRCBatchResultRef.from_raw_output(with_error) + except ValueError as exc: + assert "TooManyInputs" in str(exc) + else: + raise AssertionError("Expected ValueError for top-level UserError output") diff --git a/tests/test_hyper.py b/tests/test_hyper.py new file mode 100644 index 00000000..dfd503bc --- /dev/null +++ b/tests/test_hyper.py @@ -0,0 +1,101 @@ +from pathlib import Path + +from rush import TRC, hyper +from rush import RunOpts +from tests._module_test_utils import assert_run_collects_and_caches + + +def test_hyper_solvate_sumo(test_data_dir: Path): + run = hyper.hyper_solvate_sumo( + [test_data_dir / "hyper" / "valid_trc.json"], + config=hyper.HyperConfig(max_inputs=8, padding_nm=0.8, seed=12345), + run_opts=RunOpts( + name="Rush-Py Test Hyper Solvate", + tags=["rush-py", "test", "hyper", "solvate"], + ), + ) + + assert_run_collects_and_caches(run, hyper.TRCBatchResultRef) + + fetched = run.fetch() + assert len(fetched) == 1 + assert all(not isinstance(item, hyper.ItemError) for item in fetched) + assert isinstance(fetched[0], TRC) + + saved = run.save() + assert len(saved) == 1 + assert all(not isinstance(item, hyper.ItemError) for item in saved) + assert isinstance(saved[0], Path) + assert saved[0].exists() + + +def test_hyper_minimize_sumo(test_data_dir: Path): + run = hyper.hyper_minimize_sumo( + [ + hyper.MinimizeInput( + structure=test_data_dir / "hyper" / "methanol_trc.json", + topology=test_data_dir / "hyper" / "methanol_topology.json", + ) + ], + config=hyper.HyperMinimizeConfig(max_inputs=4, steps=100, gtol=100.0), + run_opts=RunOpts( + name="Rush-Py Test Hyper Minimize", + tags=["rush-py", "test", "hyper", "minimize"], + ), + ) + + assert_run_collects_and_caches(run, hyper.TRCBatchResultRef) + + fetched = run.fetch() + assert len(fetched) == 1 + assert all(not isinstance(item, hyper.ItemError) for item in fetched) + assert isinstance(fetched[0], TRC) + + saved = run.save() + assert len(saved) == 1 + assert all(not isinstance(item, hyper.ItemError) for item in saved) + assert isinstance(saved[0], Path) + assert saved[0].exists() + + +def test_hyper_run_sumo(test_data_dir: Path): + run = hyper.hyper_run_sumo( + [ + hyper.RunInput( + sim_config=test_data_dir / "hyper" / "sim_config.json", + topology=test_data_dir / "hyper" / "methanol_topology.json", + coordinates=test_data_dir / "hyper" / "methanol_trc.json", + ) + ], + config=hyper.HyperRunConfig( + max_inputs=4, + nsteps=20, + dt_ps=0.001, + temperature_k=300.0, + ensemble="Nvt", + minimize_before_run=False, + solvate_before_run=False, + use_gpu=False, + nthreads=1, + ), + run_opts=RunOpts( + name="Rush-Py Test Hyper Run", + tags=["rush-py", "test", "hyper", "run"], + ), + ) + + assert_run_collects_and_caches(run, hyper.RunResultRef) + + fetched = run.fetch() + assert len(fetched) == 1 + assert all(not isinstance(item, hyper.ItemError) for item in fetched) + assert isinstance(fetched[0], hyper.RunOutput) + assert fetched[0].trajectory + + saved = run.save() + assert len(saved) == 1 + assert all(not isinstance(item, hyper.ItemError) for item in saved) + assert isinstance(saved[0], hyper.RunOutputPaths) + assert saved[0].trajectory.exists() + if saved[0].checkpoint is not None: + assert saved[0].checkpoint.exists()