Merge develop branch into main (#1)

* Alpha pontibus code - including first iteration of the ASFEProtocol

.github/workflows/black.yml

@@ -0,0 +1,12 @@
+name: Lint
+
+on: [push, pull_request]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: psf/black@stable
+        with:
+          src: "./src"

.github/workflows/ci.yaml

@@ -0,0 +1,86 @@
+name: "CI"
+on:
+  pull_request:
+    branches:
+      - main
+      - develop
+    paths-ignore:
+      - "docs/*"
+  push:
+    branches:
+      - main
+      - develop
+  schedule:
+    # At 07:00 UTC on Monday and Thursday
+    - cron: "0 7 * * *"
+  workflow_dispatch:
+
+
+concurrency:
+  group: "${{ github.workflow }}-${{ github.ref }}"
+  cancel-in-progress: true
+
+defaults:
+  run:
+    shell: bash -l {0}
+
+jobs:
+  tests:
+    runs-on: ${{ matrix.os }}
+    name: "${{ matrix.os }} - ${{ matrix.python-version }}"
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest"]
+        python-version:
+          - "3.10"
+          - "3.11"
+          - "3.12"
+        include:
+          - os: "macos-latest"
+            python-version: "3.12"
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: "Setup Micromamba"
+        uses: mamba-org/setup-micromamba@v1
+        with:
+          environment-file: environment.yml
+          environment-name: pontibus
+          create-args: >-
+            python=${{ matrix.python-version }}
+          init-shell: bash
+
+      - name: "Install"
+        run: python -m pip install --no-deps -e .
+
+      - name: "Test imports"
+        run: |
+          # if we add more to this, consider changing to for + env vars
+          python -Ic "import pontibus; print(pontibus.__version__)"
+
+      - name: "Environment Information"
+        run: |
+          micromamba info
+          micromamba list
+          pip list
+
+      - name: "Run tests"
+        env:
+          PONTIBUS_SLOW_TESTS: ${{ fromJSON('{"false":"false","true":"true"}')[github.event_name != 'pull_request'] }}
+          DUECREDIT_ENABLE: 'yes'
+        run: |
+          pytest -n logical -v --cov=pontibus --cov-report=xml --durations=10
+
+      - name: codecov-pr
+        if: ${{ github.repository == 'OpenFreeEnergy/pontibus'
+                && github.event_name == 'pull_request' }}
+        uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          file: coverage.xml
+          fail_ci_if_error: False
+          verbose: True
+          flags: fast-tests

AUTHORS.md

@@ -0,0 +1,18 @@
+# Authors
+
+All contributing authors are listed in the file below.
+The repository history and CHANGELOG show individual code contributions.
+
+<!--
+The rules for this file:
+  * Authors are sorted chronologically, earliest to latest
+  * Please format it each entry as "Preferred name <GitHub username>"
+  * Your preferred name is whatever you wish to go by --
+    it does *not* have to be your legal name!
+  * Please start a new section for each new year
+  * Don't ever delete anything
+-->
+
+**2024**
+- Irfan Alibay <@IAlibay>
+- Matt Thompson <@mattwthompson>

README.md

@@ -1,2 +1,4 @@
 # pontibus
 Experimental OpenFE Protocols for Force Field Evaluation
+
+!!WARNING!! This code is experimental and subject to change !!WARNING!!

benchmarks/subsampled/gen_systems.py

@@ -0,0 +1,79 @@
+import gzip
+import json
+import pathlib
+
+from gufe import ChemicalSystem, SmallMoleculeComponent
+from gufe.tokenization import JSON_HANDLER
+from openff.toolkit import Molecule
+from pontibus.components import ExtendedSolventComponent
+from rdkit import Chem
+
+
+def add_chemical_systems(
+    sdffile: str,
+    dataset_name: str,
+    solvents: dict[str, SmallMoleculeComponent],
+    systems: list[ChemicalSystem],
+) -> None:
+    """
+    Add Solute + Solvent ChemicalSystems to running list.
+
+    Parameters
+    ----------
+    sdffile : str
+      The SDF file to read entries from.
+    dataset_name : str
+      The name of the dataset.
+    solvents: dict[str, SmallMoleculeComponent]
+      Running dictionary of solvents to draw & store prepared solvent
+      molecules from/to.
+    systems: list[ChemicalSystem]
+      Runing list of ChemicalSystems we are appending to.
+    """
+    for i, rdmol in enumerate(Chem.SDMolSupplier(sdffile, removeHs=False)):
+        offmol = Molecule.from_rdkit(rdmol)
+        offmol.assign_partial_charges(partial_charge_method="am1bccelf10")
+        solvent_smi = rdmol.GetProp("solvent")
+        if solvent_smi not in solvents.keys():
+            solvent_offmol = Molecule.from_smiles(solvent_smi)
+            solvent_offmol.generate_conformers()
+            solvent_offmol.assign_partial_charges(partial_charge_method="am1bccelf10")
+            solvents[solvent_smi] = SmallMoleculeComponent.from_openff(solvent_offmol)
+
+        systems.append(
+            ChemicalSystem(
+                {
+                    "solute": SmallMoleculeComponent.from_openff(offmol),
+                    "solvent": ExtendedSolventComponent(
+                        solvent_molecule=solvents[solvent_smi]
+                    ),
+                },
+                name=f"molecule{i}_{dataset_name}",
+            )
+        )
+
+
+def store_chemical_systems(systems: list[ChemicalSystem], outdir: pathlib.Path):
+    """
+    Store ChemicalSystems to gzip file.
+
+    Parameters
+    ----------
+    systems: list[ChemicalSystem]
+      List of ChemicalSystems to store to file.
+    """
+    for system in systems:
+        with gzip.open(outdir / f"{system.name}_chemicalsystem.gz", "wt") as zipfile:
+            json.dump(system.to_dict(), zipfile, cls=JSON_HANDLER.encoder)
+
+
+if __name__ == "__main__":
+    solvents: dict[str, SmallMoleculeComponent] = {}
+    systems: list[ChemicalSystem] = []
+
+    add_chemical_systems("sub_sampled_fsolv.sdf", "fsolv", solvents, systems)
+    add_chemical_systems("sub_sampled_mnsol.sdf", "mnsol", solvents, systems)
+
+    outdir = pathlib.Path("chemicalsystems")
+    outdir.mkdir(exist_ok=True)
+    store_chemical_systems(systems, outdir)

benchmarks/subsampled/gen_transforms_2.0.0_single_repeat.py

@@ -0,0 +1,102 @@
+import gzip
+import json
+import pathlib
+
+from gufe import AlchemicalNetwork, ChemicalSystem, Transformation
+from gufe.tokenization import JSON_HANDLER
+from openff.toolkit import Molecule
+from openff.units import unit
+from pontibus.protocols.solvation import ASFEProtocol
+
+
+def deserialize_system(file: pathlib.Path):
+    with gzip.open(file, "r") as f:
+        d = json.loads(f.read().decode(), cls=JSON_HANDLER.decoder)
+    return ChemicalSystem.from_dict(d)
+
+
+def get_water_settings():
+    settings = ASFEProtocol.default_settings()
+    settings.protocol_repeats = 1
+    settings.solvent_forcefield_settings.forcefields = [
+        "openff-2.0.0.offxml",
+        "tip3p.offxml",
+    ]
+    settings.vacuum_forcefield_settings.forcefields = [
+        "openff-2.0.0.offxml",
+    ]
+    settings.solvent_simulation_settings.time_per_iteration = 5 * unit.picosecond
+    settings.vacuum_simulation_settings.time_per_iteration = 5 * unit.picosecond
+    settings.vacuum_engine_settings.compute_platform = "CPU"
+    settings.solvent_engine_settings.compute_platform = "CUDA"
+    settings.solvation_settings.box_shape = "dodecahedron"
+    return settings
+
+
+def get_nonwater_settings():
+    settings = ASFEProtocol.default_settings()
+    settings.protocol_repeats = 1
+    settings.solvent_forcefield_settings.forcefields = [
+        "openff-2.0.0.offxml",
+        "tip3p.offxml",
+    ]
+    settings.vacuum_forcefield_settings.forcefields = [
+        "openff-2.0.0.offxml",
+    ]
+    settings.solvent_simulation_settings.time_per_iteration = 5 * unit.picosecond
+    settings.vacuum_simulation_settings.time_per_iteration = 5 * unit.picosecond
+    settings.vacuum_engine_settings.compute_platform = "CUDA"
+    settings.solvent_engine_settings.compute_platform = "CUDA"
+    settings.solvation_settings.box_shape = "dodecahedron"
+    settings.solvation_settings.assign_solvent_charges = True
+    return settings
+
+
+def get_transformation(stateA):
+    solvent = stateA.components["solvent"]
+
+    water = Molecule.from_smiles("O")
+    if water.is_isomorphic_with(solvent.solvent_molecule.to_openff()):
+        settings = get_water_settings()
+    else:
+        settings = get_nonwater_settings()
+
+    stateB = ChemicalSystem({"solvent": solvent})
+    protocol = ASFEProtocol(settings=settings)
+    return Transformation(
+        stateA=stateA, stateB=stateB, mapping=None, protocol=protocol, name=stateA.name
+    )
+
+
+def run(outdir: pathlib.Path):
+    systems: list[ChemicalSystem] = []
+
+    systems_dir = pathlib.Path("chemicalsystems")
+    system_files = systems_dir.glob("*.gz")
+
+    for file in system_files:
+        systems.append(deserialize_system(file))
+
+    transformations = []
+    for system in systems:
+        transformations.append(get_transformation(system))
+
+    alchemical_network = AlchemicalNetwork(transformations)
+    an_outfile = outdir / "alchemical_network.json"
+    json.dump(
+        alchemical_network.to_dict(),
+        an_outfile.open(mode="w"),
+        cls=JSON_HANDLER.encoder,
+    )
+
+    transforms_dir = outdir / "transformations"
+    transforms_dir.mkdir(exist_ok=True, parents=True)
+
+    for transform in alchemical_network.edges:
+        transform.dump(transforms_dir / f"{transform.name}.json")
+
+
+if __name__ == "__main__":
+    outdir = pathlib.Path("2.0.0_single_repeat_inputs")
+    outdir.mkdir(exist_ok=False)
+    run(outdir)