Lukas/more optimizers

.github/workflows/tests.yml

@@ -37,6 +37,30 @@ jobs:
       - run: "pip install -e '.[dev]'"
       - run: ruff format --check .
   pytest:
+    # https://ericmjl.github.io/blog/2021/12/30/better-conda-environments-on-github-actions/
+    runs-on: ubuntu-22.04
+    defaults:
+      run:
+        shell: bash -l {0}
+    steps:
+      - uses: actions/checkout@v3
+      - uses: conda-incubator/setup-miniconda@v2
+        with:
+          auto-update-conda: true
+          miniforge-variant: Mambaforge
+          channels: conda-forge,defaults
+          python-version: 3.11
+          activate-environment: pytest
+          use-mamba: true
+      - name: Build environment
+        run: |
+          conda activate pytest
+          conda install -c conda-forge openmm openmmtools
+          python -m pip install --upgrade pip
+          python -m pip install --upgrade setuptools
+          python -m pip install -e '.[dev]'
+      - run: pytest
+  doctest:
     runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v3
@@ -45,4 +69,4 @@ jobs:
           python-version: "3.11"
           cache: "pip"
       - run: "pip install -e '.[dev]'"
-      - run: pytest
+      - run: make -C docs doctest

.readthedocs.yml

@@ -1,16 +1,12 @@
 version: 2
 
 build:
-    os: "ubuntu-22.04"
-    tools:
-      python: "3.11"
+  os: "ubuntu-22.04"
+  tools:
+    python: "mambaforge-latest"
 
 sphinx:
-    configuration: "docs/source/conf.py"
+  configuration: "docs/source/conf.py"
 
-python:
-    install:
-        - method: "pip"
-          path: .
-          extra_requirements:
-              - "dev"
+conda:
+  environment: "docs/environment.yml"

README.rst

@@ -23,21 +23,24 @@ for ``stk``, which can be joined through https://discord.gg/zbCUzuxe2B.
 Installation
 ============
 
-To get ``stko``, you can install it with pip::
+To get ``stko``, you can install it with pip:
 
-    $ pip install stko
+.. code-block:: bash
 
-You can also install xtb, which is an optional dependency::
+  pip install stko
 
-    $ conda install -c conda-forge xtb
+Some optional dependencies are only available through conda:
 
-You can also install openbabel, which is an optional dependency::
+.. code-block:: bash
 
-    $ conda install -c conda-forge openbabel
-
-You can also install MDAnalysis, which is an optional dependency::
-
-    $ conda install -c mdanalysis
+  # for OpenMM
+  mamba install openff-toolkit openmm openmmtools
+  # for xtb
+  mamba install xtb
+  # for openbabel
+  mamba install openbabel
+  # for mdanalysis
+  mamba install mdanalysis
 
 Examples
 ========

docs/environment.yml

@@ -0,0 +1,15 @@
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - openff-toolkit
+  - openmm
+  - openmmtools
+  - python=3.11
+  - pip
+  - pip:
+      - numpy
+      - rdkit
+      - stk
+      - networkx
+      - rmsd

docs/source/calculators.rst

@@ -7,6 +7,7 @@ Calculators
   MMFFEnergy <_autosummary/stko.MMFFEnergy>
   UFFEnergy <_autosummary/stko.UFFEnergy>
   OpenBabelEnergy <_autosummary/stko.OpenBabelEnergy>
+  OpenMMEnergy <_autosummary/stko.OpenMMEnergy>
   XTBEnergy <_autosummary/stko.XTBEnergy>
   OrcaEnergy <_autosummary/stko.OrcaEnergy>
   RmsdCalculator <_autosummary/stko.RmsdCalculator>

docs/source/conf.py

@@ -35,6 +35,11 @@
     "python": ("https://docs.python.org/3", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
     "pandas": ("https://pandas.pydata.org/docs/", None),
+    "openff": (
+        "https://docs.openforcefield.org/projects/toolkit/en/stable/",
+        None,
+    ),
+    "openmm": ("http://docs.openmm.org/latest/api-python/", None),
 }
 
 templates_path = ["_templates"]

docs/source/index.rst

@@ -27,9 +27,24 @@ GitHub: https://github.com/JelfsMaterialsGroup/stko
 Install
 =======
 
-:mod:`.stko` can be installed directly with pip::
+:mod:`.stko` can be installed directly with pip:
 
-   $ pip install stko
+.. code-block:: bash
+
+  pip install stko
+
+Some optional dependencies are only available through conda:
+
+.. code-block:: bash
+
+  # for OpenMM
+  mamba install openff-toolkit openmm openmmtools
+  # for xtb
+  mamba install xtb
+  # for openbabel
+  mamba install openbabel
+  # for mdanalysis
+  mamba install mdanalysis
 
 
 Dependencies
@@ -42,14 +57,16 @@ on the desired functions used. These are:
 * `GULP <http://gulp.curtin.edu.au/gulp/>`_
 * `XTB <https://xtb-docs.readthedocs.io/en/latest/>`_
 * `OpenBabel <https://github.com/openbabel/openbabel>`_
+* `OpenMM <https://openmm.org/>`_
+* `OpenFF <https://openforcefield.org/>`_
 
 
 Overview
 ========
 
-`stko <https://github.com/lukasturcani/stk>`_ is a Python library which
+`stko <https://github.com/JelfsMaterialsGroup/stko>`_ is a Python library which
 performs optimizations and calculations on complex molecules built using
-`stk <https://github.com/JelfsMaterialsGroup/stko>`_. In the case of
+`stk <https://github.com/lukasturcani/stk>`_. In the case of
 optimizations, a clone of :class:`stk.Molecule` is returned. For
 calculators, a ``Results`` class are used to calculate and extract
 properties of an :class:`stk.Molecule`.

docs/source/optimizers.rst

@@ -20,6 +20,8 @@ Optimizers
   MMFF <_autosummary/stko.MMFF>
   MetalOptimizer <_autosummary/stko.MetalOptimizer>
   OpenBabel <_autosummary/stko.OpenBabel>
+  OpenMMForceField <_autosummary/stko.OpenMMForceField>
+  OpenMMMD <_autosummary/stko.OpenMMMD>
   GulpUFFOptimizer <_autosummary/stko.GulpUFFOptimizer>
   MacroModelForceField <_autosummary/stko.MacroModelForceField>
   XTB <_autosummary/stko.XTB>

examples/aligner_example.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import stk
+
 import stko
 
 

examples/basic_example.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 import stk
+
 import stko
 
 

examples/cage_analysis_example.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import stk
+
 import stko
 
 

examples/calculators_example.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 
 import stk
+
 import stko
 
 

examples/gulp_test_example.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import stk
+
 import stko
 
 

examples/mdanalysis_example.py

@@ -1,6 +1,7 @@
 # ruff: noqa: T201
 import numpy as np
 import stk
+
 import stko
 
 

examples/molecule_splitter_example.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 
 import stk
+
 import stko
 
 

examples/obabel_example.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import stk
+
 import stko
 
 

examples/openmm_example.py

@@ -0,0 +1,102 @@
+# ruff: noqa: T201
+
+from pathlib import Path
+
+import openmm
+import stk
+from openff.toolkit import ForceField
+
+import stko
+
+
+def main() -> None:
+    """Run the example."""
+    bb1 = stk.BuildingBlock("C1CCC(C(C1)N)N", [stk.PrimaryAminoFactory()])
+    bb2 = stk.BuildingBlock(
+        "C1=C(C=C(C=C1C=O)C=O)C=O", [stk.AldehydeFactory()]
+    )
+    output = Path("openmm_example_directory")
+    output.mkdir(exist_ok=True, parents=True)
+    cage = stk.ConstructedMolecule(stk.cage.FourPlusSix([bb1, bb2]))
+    cage.write(output / "unopt_cage.mol")
+    ff_optimizer = stko.OpenMMForceField(
+        # Load the openff-2.1.0 force field appropriate for
+        # vacuum calculations (without constraints)
+        force_field=ForceField("openff_unconstrained-2.1.0.offxml"),
+        partial_charges_method="mmff94",
+        max_iterations=10,
+    )
+    ff_cage = ff_optimizer.optimize(cage)
+    ff_cage.write(output / "ff_opt_cage.mol")
+    print(
+        "ff_opt_cage",
+        stko.OpenMMEnergy(
+            force_field=ForceField("openff_unconstrained-2.1.0.offxml"),
+            partial_charges_method="mmff94",
+        ).get_energy(ff_cage),
+    )
+
+    ff_optimizer = stko.OpenMMForceField(
+        # Load the openff-2.1.0 force field appropriate for
+        # vacuum calculations (without constraints)
+        force_field=ForceField("openff_unconstrained-2.1.0.offxml"),
+        restricted=False,
+        partial_charges_method="mmff94",
+    )
+    ff_cage = ff_optimizer.optimize(cage)
+    ff_cage.write(output / "ffunrest_opt_cage.mol")
+    print(
+        "ffunrest_opt_cage",
+        stko.OpenMMEnergy(
+            force_field=ForceField("openff_unconstrained-2.1.0.offxml"),
+            partial_charges_method="mmff94",
+        ).get_energy(ff_cage),
+    )
+
+    ff_optimizer = stko.OpenMMForceField(
+        # Load the openff-2.1.0 force field appropriate for
+        # vacuum calculations (without constraints)
+        force_field=ForceField("openff_unconstrained-2.1.0.offxml"),
+        restricted=True,
+        partial_charges_method="mmff94",
+    )
+    ff_cage = ff_optimizer.optimize(cage)
+    ff_cage.write(output / "ffrest_opt_cage.mol")
+    print(
+        "ffrest_opt_cage",
+        stko.OpenMMEnergy(
+            force_field=ForceField("openff_unconstrained-2.1.0.offxml"),
+            partial_charges_method="mmff94",
+        ).get_energy(ff_cage),
+    )
+
+    temperature = 750 * openmm.unit.kelvin
+    friction = 10 / openmm.unit.picoseconds
+    time_step = 1 * openmm.unit.femtoseconds
+    integrator = openmm.LangevinIntegrator(temperature, friction, time_step)
+    integrator.setRandomNumberSeed(127)
+    md_optimizer = stko.OpenMMMD(
+        force_field=ForceField("openff_unconstrained-2.1.0.offxml"),
+        output_directory=output / "md_optimisation",
+        integrator=integrator,
+        random_seed=275,
+        partial_charges_method="mmff94",
+        reporting_freq=5,
+        trajectory_freq=5,
+        num_steps=10000,
+        num_conformers=50,
+        platform="CUDA",
+    )
+    md_cage = md_optimizer.optimize(cage)
+    md_cage.write(output / "md_opt_cage.mol")
+    print(
+        "md_opt_cage",
+        stko.OpenMMEnergy(
+            force_field=ForceField("openff_unconstrained-2.1.0.offxml"),
+            partial_charges_method="mmff94",
+        ).get_energy(md_cage),
+    )
+
+
+if __name__ == "__main__":
+    main()