strain-relief 0.4.dev0__tar.gz → 1.0.dev0__tar.gz

{strain_relief-0.4.dev0 → strain_relief-1.0.dev0}/.gitignore

@@ -9,3 +9,5 @@ __pycache__/
 *.pyc
 
 tests/models/eSEN.pt
+
+outputs/

{strain_relief-0.4.dev0 → strain_relief-1.0.dev0}/.pre-commit-config.yaml

@@ -1,15 +1,11 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v6.0.0
     hooks:
       - id: check-yaml
       - id: trailing-whitespace
       - id: end-of-file-fixer
       - id: debug-statements
-  - repo: https://github.com/kynan/nbstripout
-    rev: 0.6.0
-    hooks:
-      - id: nbstripout
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.2.2
     hooks:

strain_relief-1.0.dev0/.vscode/launch.json

@@ -0,0 +1,68 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python Debugger: Current File",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal",
+            "justMyCode": false
+        },
+        {
+            "name": "train (mmff94)",
+            "type": "debugpy",
+            "program": "${workspaceFolder}/src/strain_relief/cmdline/_strain_relief.py",
+            "request": "launch",
+            "console": "integratedTerminal",
+            "justMyCode": true,
+            "args": [
+                "experiment=mmff94s",
+                "conformers.numConfs=1",
+                "io.input.parquet_path=${workspaceFolder}/data/example_ligboundconf_input.parquet",
+            ],
+            "env": {
+                "HYDRA_FULL_ERROR": "1",
+            }
+        },
+        {
+            "name": "train (mace - gpu)",
+            "type": "debugpy",
+            "program": "${workspaceFolder}/src/strain_relief/cmdline/_strain_relief.py",
+            "request": "launch",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "args": [
+                "experiment=mace",
+                "conformers.numConfs=1",
+                "io.input.parquet_path=${workspaceFolder}/data/example_ligboundconf_input.parquet",
+                "device=cuda",
+            ],
+            "env": {
+                "HYDRA_FULL_ERROR": "1",
+                "CUDA_LAUNCH_BLOCKING": "1"
+            }
+        },
+        {
+            "name": "train (mace- cpu)",
+            "type": "debugpy",
+            "program": "${workspaceFolder}/src/strain_relief/cmdline/_strain_relief.py",
+            "request": "launch",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "args": [
+                "experiment=mace",
+                "conformers.numConfs=1",
+                "io.input.parquet_path=${workspaceFolder}/data/example_ligboundconf_input.parquet",
+                "device=cpu",
+            ],
+            "env": {
+                "HYDRA_FULL_ERROR": "1",
+                "CUDA_LAUNCH_BLOCKING": "1"
+            }
+        },
+    ]
+}

strain_relief-1.0.dev0/PKG-INFO

@@ -0,0 +1,192 @@
+Metadata-Version: 2.4
+Name: strain-relief
+Version: 1.0.dev0
+Summary: Calculates ligand strain of small molecules from their docked poses.
+Author-email: Ewan Wallace <erw@wallace2.com>
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/prescient-design/StrainRelief
+Requires-Python: ==3.11.*
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: rdkit
+Requires-Dist: ase
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: omegaconf
+Requires-Dist: hydra-core
+Requires-Dist: pyarrow
+Requires-Dist: fastparquet
+Requires-Dist: universal-pathlib
+Requires-Dist: boto3
+Requires-Dist: loguru
+Requires-Dist: rich
+Requires-Dist: neural-optimiser
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-mock; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: ipykernel; extra == "dev"
+Dynamic: license-file
+
+# StrainRelief 💊
+StrainRelief calculates the ligand strain of docked poses and has a suite of different force fields with which to do this. This includes our own MACE neural network potential trained on SPICE2 and also Meta's FAIRChem models, such as e-SEN and UMA. Optimisations are run in parallel using the neural-optimiser package.
+
+- 📄 The publication can be found [here](https://pubs.acs.org/doi/10.1021/acs.jcim.5c00586).
+- 📊 All relevant datasets [here](https://huggingface.co/datasets/erwallace/LigBoundConf2.0).
+- 💬 RAG [chatbot](https://strain-relief.streamlit.app/) for questions about the paper and references.
+- 💻 Chatbot source [code](https://github.com/erwallace/paper_query).
+- 📉 neural-optimiser [package](https://github.com/erwallace/neural-optimiser).
+
+![Strain Relief Logo](assets/strain_relief_plus_logo.png)
+
+## Update 1.0.0
+
+This update focuses on accelerating the StrainRelief backend and increasig usability through simpler configurations and workflows:
+
+- Molecular optimisation are now paralellised using the `neural-optimiser` [package](https://github.com/erwallace/neural-optimiser) providing a 3x speed-up.
+- New simplified hyra configuration. An example of the full configuration is given [here](./examples/example_config.yaml).
+
+**Note:** this introduces changes that are not backwards-compatible with previous versions.
+
+## Installation
+Pre-requisities: Python 3.11, PyTorch and PyTorch Geometric compatible with your envirnment
+
+```bash
+# For example
+uv pip install torch==2.8.0 -f https://data.pyg.org/whl/torch-2.8.0+cu128.html
+uv pip install torch-geometric==2.7.0 torch-cluster -f https://data.pyg.org/whl/torch-2.8.0+cu128.html
+```
+
+### Installation from PyPi
+
+```
+pip install strain-relief
+```
+
+### Installation from source (uv)
+```bash
+uv sync --extra dev --editable
+```
+or create a virtual environment and install the package and its dependencies in editable mode:
+```bash
+uv venv
+source .venv/bin/activate
+# [install torch and torch-geometric as above]
+uv pip install -e ".[dev]"
+uv pip install --force-reinstall e3nn==0.5 fairchem-core
+
+uv run pre-commit install
+```
+**Note:** `mace-torch==0.3.x` requires `e3nn==0.4.4` (only for training, not inference). `fairchem-core` requires `e3nn>=0.5`. So until `mace-torch==0.4` is released we will have to do this finicky way of installing ([GitHub issue](https://github.com/ACEsuit/mace/issues/555)).
+
+## The Protocol
+
+The protocol used in StrainRelief is designed to be simple, fast and model agnostic - all that is needed to apply a new force field is to write a `neural-optimiser` `Calculator` class (such as an ASE calculator wrapper). Additionally, the package is already compatible with all MACE and Meta FAIRChem models.
+
+![Strain Relief Protocol](assets/strain_relief_protocol.png)
+
+The protocol consists of 5 steps:
+
+1. Minimise the docked pose with a loose convergence criteria to give a local minimum.
+2. Generate 20 conformers from the docked ligand pose.
+3. Minimise the generated conformers (and the original docked pose) with a stricter convergence criteria.
+4. Evaluate the energy of all conformers and choose the lowest energy as an approximation of the global minimum.
+5. Calculate `E(ligand strain) = E(local minimum) - E(global minimum)` and apply threshold.
+
+**N.B.** energies returned are in kcal/mol.
+
+## Usage
+
+StrainRelief runs are configured using hydra configs.
+
+### Python Package
+
+```
+from strain_relief import compute_strain
+
+computed = compute_strain(poses: list[RDKit.Mol], config: DictConfig)
+
+for i, r in computed.iterrows():
+    print(f"Pose {r['id']} has a strain of {r['ligand_strain']:.2f} kcal/mol")
+```
+For a complete set of examples see the tutorial [notebook](./examples/tutorial.ipynb).
+
+### Command Line
+
+```bash
+strain-relief \
+    experiment=mmff94 \
+    io.input.parquet_path=data/example_ligboundconf_input.parquet \
+    io.output.parquet_path=data/example_ligboundconf_output.parquet \
+    conformers.numConfs=1 \
+```
+
+More examples are given [here](./examples/example_scripts.sh), including the command used for the calculations in the StrainRelief paper.
+
+### Adding Your Own Calculator
+
+Add a new `Calculator` [class](https://github.com/erwallace/neural-optimiser/blob/main/src/neural_optimiser/calculators/base.py) from `neural-optimiser` to the `strain_relief/calculators/` directory. This can be as simple as implementing a wrapper around an existing ASE calculator:
+```python
+from neural_optimiser.calculators.base import Calculator
+
+class YourCalculator(Calculator):
+
+    def __init__(self, **kwargs):
+        self.calculator = YourASECalculator(**kwargs)
+
+    def _calculate(self, batch: Data | Batch) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return (energies, forces) from the calculator."""
+        ase_atoms = batch.to_ase()
+        ase_atoms.calc = self.calculator
+        return ase_atoms.get_potential_energy(), ase_atoms.get_gradients()
+```
+**Note:** `MACECalculator` and `FAIRChemCalculator` from `neural-optimiser` use a `from_atomic_data` helper method. This converts `ConformerBatch` objects to `AtomicData` model inputs in a batched process; a workflow bottleneck not handled by either the MACE or FAIRChem internal ASE calculators. I would recommend implementing something similar for high throughput workflows.
+
+Add a new config to `hydra_config/calculator/your_calculator.yaml`:
+
+```
+_target_: strain_relief.calculators.your_calculator.YourCalculator
+model_paths: null
+device: ${device}
+any_other_kwargs: null
+```
+
+### Configurations
+
+**Common kwargs**
+- `threshold` (set by default to 16.1 kcal/mol - calibrated using [LigBoundConf 2.0](https://huggingface.co/datasets/erwallace/LigBoundConf2.0))
+- `conformers.numConfs`
+- `global_optimiser.steps`/`local_optimiser.steps`
+- `global_optimiser.fmax`/`local_optimiser.fmax`
+- `io.input.include_charged`
+- `hydra.verbose`
+- `seed`
+
+### Logging
+
+Logging is set to the `INFO` level by default which logs only aggregate information. `hydra.verbose=true` can be used to activate `DEBUG` level logging which includes information for every molecule and conformer.
+
+## Unit Tests
+
+- `uv run pytest tests/` - runs all tests (unit and integration)
+- `uv run pytest tests/ -m "not integration"` - runs all unit tests
+
+**Note:** Unit tests will run on a GPU if available.
+
+## Citations
+If you use StrainRelief or adapt the StrainRelief code for any purpose, please cite:
+
+```bibtex
+@article{wallace2025strain,
+  title={Strain Problems Got You in a Twist? Try StrainRelief: A Quantum-Accurate Tool for Ligand Strain Calculations},
+  author={Wallace, Ewan RS and Frey, Nathan C and Rackers, Joshua A},
+  journal={Journal of Chemical Information and Modeling},
+  year={2025},
+  publisher={ACS Publications},
+  url={https://pubs.acs.org/doi/10.1021/acs.jcim.5c00586}
+}
+```
+
+## More information
+For any questions, please reach out to [Ewan Wallace](https://www.linkedin.com/in/ewan-wallace-82297318a/): ewan.wallace@roche.com

strain_relief-1.0.dev0/README.md

@@ -0,0 +1,161 @@
+# StrainRelief 💊
+StrainRelief calculates the ligand strain of docked poses and has a suite of different force fields with which to do this. This includes our own MACE neural network potential trained on SPICE2 and also Meta's FAIRChem models, such as e-SEN and UMA. Optimisations are run in parallel using the neural-optimiser package.
+
+- 📄 The publication can be found [here](https://pubs.acs.org/doi/10.1021/acs.jcim.5c00586).
+- 📊 All relevant datasets [here](https://huggingface.co/datasets/erwallace/LigBoundConf2.0).
+- 💬 RAG [chatbot](https://strain-relief.streamlit.app/) for questions about the paper and references.
+- 💻 Chatbot source [code](https://github.com/erwallace/paper_query).
+- 📉 neural-optimiser [package](https://github.com/erwallace/neural-optimiser).
+
+![Strain Relief Logo](assets/strain_relief_plus_logo.png)
+
+## Update 1.0.0
+
+This update focuses on accelerating the StrainRelief backend and increasig usability through simpler configurations and workflows:
+
+- Molecular optimisation are now paralellised using the `neural-optimiser` [package](https://github.com/erwallace/neural-optimiser) providing a 3x speed-up.
+- New simplified hyra configuration. An example of the full configuration is given [here](./examples/example_config.yaml).
+
+**Note:** this introduces changes that are not backwards-compatible with previous versions.
+
+## Installation
+Pre-requisities: Python 3.11, PyTorch and PyTorch Geometric compatible with your envirnment
+
+```bash
+# For example
+uv pip install torch==2.8.0 -f https://data.pyg.org/whl/torch-2.8.0+cu128.html
+uv pip install torch-geometric==2.7.0 torch-cluster -f https://data.pyg.org/whl/torch-2.8.0+cu128.html
+```
+
+### Installation from PyPi
+
+```
+pip install strain-relief
+```
+
+### Installation from source (uv)
+```bash
+uv sync --extra dev --editable
+```
+or create a virtual environment and install the package and its dependencies in editable mode:
+```bash
+uv venv
+source .venv/bin/activate
+# [install torch and torch-geometric as above]
+uv pip install -e ".[dev]"
+uv pip install --force-reinstall e3nn==0.5 fairchem-core
+
+uv run pre-commit install
+```
+**Note:** `mace-torch==0.3.x` requires `e3nn==0.4.4` (only for training, not inference). `fairchem-core` requires `e3nn>=0.5`. So until `mace-torch==0.4` is released we will have to do this finicky way of installing ([GitHub issue](https://github.com/ACEsuit/mace/issues/555)).
+
+## The Protocol
+
+The protocol used in StrainRelief is designed to be simple, fast and model agnostic - all that is needed to apply a new force field is to write a `neural-optimiser` `Calculator` class (such as an ASE calculator wrapper). Additionally, the package is already compatible with all MACE and Meta FAIRChem models.
+
+![Strain Relief Protocol](assets/strain_relief_protocol.png)
+
+The protocol consists of 5 steps:
+
+1. Minimise the docked pose with a loose convergence criteria to give a local minimum.
+2. Generate 20 conformers from the docked ligand pose.
+3. Minimise the generated conformers (and the original docked pose) with a stricter convergence criteria.
+4. Evaluate the energy of all conformers and choose the lowest energy as an approximation of the global minimum.
+5. Calculate `E(ligand strain) = E(local minimum) - E(global minimum)` and apply threshold.
+
+**N.B.** energies returned are in kcal/mol.
+
+## Usage
+
+StrainRelief runs are configured using hydra configs.
+
+### Python Package
+
+```
+from strain_relief import compute_strain
+
+computed = compute_strain(poses: list[RDKit.Mol], config: DictConfig)
+
+for i, r in computed.iterrows():
+    print(f"Pose {r['id']} has a strain of {r['ligand_strain']:.2f} kcal/mol")
+```
+For a complete set of examples see the tutorial [notebook](./examples/tutorial.ipynb).
+
+### Command Line
+
+```bash
+strain-relief \
+    experiment=mmff94 \
+    io.input.parquet_path=data/example_ligboundconf_input.parquet \
+    io.output.parquet_path=data/example_ligboundconf_output.parquet \
+    conformers.numConfs=1 \
+```
+
+More examples are given [here](./examples/example_scripts.sh), including the command used for the calculations in the StrainRelief paper.
+
+### Adding Your Own Calculator
+
+Add a new `Calculator` [class](https://github.com/erwallace/neural-optimiser/blob/main/src/neural_optimiser/calculators/base.py) from `neural-optimiser` to the `strain_relief/calculators/` directory. This can be as simple as implementing a wrapper around an existing ASE calculator:
+```python
+from neural_optimiser.calculators.base import Calculator
+
+class YourCalculator(Calculator):
+
+    def __init__(self, **kwargs):
+        self.calculator = YourASECalculator(**kwargs)
+
+    def _calculate(self, batch: Data | Batch) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return (energies, forces) from the calculator."""
+        ase_atoms = batch.to_ase()
+        ase_atoms.calc = self.calculator
+        return ase_atoms.get_potential_energy(), ase_atoms.get_gradients()
+```
+**Note:** `MACECalculator` and `FAIRChemCalculator` from `neural-optimiser` use a `from_atomic_data` helper method. This converts `ConformerBatch` objects to `AtomicData` model inputs in a batched process; a workflow bottleneck not handled by either the MACE or FAIRChem internal ASE calculators. I would recommend implementing something similar for high throughput workflows.
+
+Add a new config to `hydra_config/calculator/your_calculator.yaml`:
+
+```
+_target_: strain_relief.calculators.your_calculator.YourCalculator
+model_paths: null
+device: ${device}
+any_other_kwargs: null
+```
+
+### Configurations
+
+**Common kwargs**
+- `threshold` (set by default to 16.1 kcal/mol - calibrated using [LigBoundConf 2.0](https://huggingface.co/datasets/erwallace/LigBoundConf2.0))
+- `conformers.numConfs`
+- `global_optimiser.steps`/`local_optimiser.steps`
+- `global_optimiser.fmax`/`local_optimiser.fmax`
+- `io.input.include_charged`
+- `hydra.verbose`
+- `seed`
+
+### Logging
+
+Logging is set to the `INFO` level by default which logs only aggregate information. `hydra.verbose=true` can be used to activate `DEBUG` level logging which includes information for every molecule and conformer.
+
+## Unit Tests
+
+- `uv run pytest tests/` - runs all tests (unit and integration)
+- `uv run pytest tests/ -m "not integration"` - runs all unit tests
+
+**Note:** Unit tests will run on a GPU if available.
+
+## Citations
+If you use StrainRelief or adapt the StrainRelief code for any purpose, please cite:
+
+```bibtex
+@article{wallace2025strain,
+  title={Strain Problems Got You in a Twist? Try StrainRelief: A Quantum-Accurate Tool for Ligand Strain Calculations},
+  author={Wallace, Ewan RS and Frey, Nathan C and Rackers, Joshua A},
+  journal={Journal of Chemical Information and Modeling},
+  year={2025},
+  publisher={ACS Publications},
+  url={https://pubs.acs.org/doi/10.1021/acs.jcim.5c00586}
+}
+```
+
+## More information
+For any questions, please reach out to [Ewan Wallace](https://www.linkedin.com/in/ewan-wallace-82297318a/): ewan.wallace@roche.com

strain_relief-1.0.dev0/examples/example_config.yaml

@@ -0,0 +1,44 @@
+seed: -1
+threshold: 16.1
+num_workers: 0
+device: cpu
+batch_size: 1
+
+io:
+  input:
+    parquet_path: data/example_ligboundconf_input.parquet
+    mol_col_name: null
+    id_col_name: null
+    include_charged: true
+  output:
+    parquet_path: null
+    mol_col_name: ${..input.mol_col_name}
+    id_col_name: ${..input.id_col_name}
+    molecule_attr: ${..input.id_col_name}
+
+conformers:
+  randomSeed: ${seed}
+  numConfs: 1
+  maxAttempts: 10
+  pruneRmsThresh: 0.1
+  clearConfs: false
+  numThreads: ${num_workers}
+
+calculator:
+  _target_: neural_optimiser.calculators.MACECalculator
+  model_paths: models/MACE_SPICE2_NEUTRAL.model
+  device: ${device}
+
+local_optimiser:
+  _target_: neural_optimiser.optimisers.BFGS
+  max_step: 0.04
+  steps: 250
+  fmax: 0.5
+  fexit: 10
+
+global_optimiser:
+  _target_: neural_optimiser.optimisers.BFGS
+  max_step: 0.04
+  steps: 250
+  fmax: 0.05
+  fexit: 50

strain_relief-1.0.dev0/examples/example_scripts.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+
+#SBATCH -J strain
+#SBATCH -o example_script-%j.out
+#SBATCH -p gpu2
+#SBATCH --gpus-per-node 1
+#SBATCH --mem 10GB
+
+# scripts should be run from the StrainRelief root directory
+
+source ~/StrainRelief/.venv/bin/activate  # with uv
+# mamba activate strain  # with conda/mamba
+
+# You must choose a minimisation and energy evaluation method from "mmff94",
+# "mmff94s", "mace" or "fairchem". The calculator works best when the same
+# force field is used for both methods. If this is the case, "energy_eval"
+# does not need to be specified.
+
+# This is the simplest and fastest implementation of StrainRelief using MMFF94s
+# and a minimial example dataset.
+strain-relief \
+    io.input.parquet_path=./data/example_ligboundconf_input.parquet \
+    io.output.parquet_path=./data/example_ligboundconf_output.parquet \
+    optimiser@local_optimiser=bfgs \
+    optimiser@global_optimiser=bfgs \
+    calculator=mmff94 \
+    batch_size=1 \
+    device=cpu
+
+# This script demonstrates using different force fields for minimisation
+# (MMFF94s) and energy evaluations (MACE).
+strain-relief \
+    io.input.parquet_path=./data/example_ligboundconf_input.parquet \
+    io.output.parquet_path=./data/example_ligboundconf_output.parquet \
+    optimiser@local_optimiser=bfgs \
+    optimiser@global_optimiser=bfgs \
+    calculator=mmff94 \
+    +calculator@energy_evaluation.calculator=mace \
+    +energy_evaluation.calculator.model_paths=./tests/models/MACE.model \
+    batch_size=1 \
+    device=cpu
+
+# This is the script as used for most calculations in the StrainRelief paper.
+# MACE is used for minimisation (and energy evalutions implicitly). A looser
+# convergence criteria is used for local minimisation.
+strain-relief \
+    io.input.parquet_path=../data/example_ligboundconf_input.parquet \
+    io.output.parquet_path=../data/example_ligboundconf_output.parquet \
+    optimiser@local_optimiser=bfgs \
+    optimiser@global_optimiser=bfgs \
+    calculator=mace \
+    calculator.model_paths=./tests/models/MACE.model \
+    hydra.verbose=true \
+    device=cuda