Dgl nagl update (jthorton#13)

* update to nagl main and add new charge models * add charge checkpoint files * update workflows and lint * run isort * fix lint action, update env for missing nagl deps * fix pip install command * remove plugin tests for now, fix lint install * update lint workflow * fix python setup * fix flake8 errors
bismuthadams1 · Dec 1, 2023 · bfeb1f6 · bfeb1f6
1 parent c067a6f
commit bfeb1f6
Show file tree

Hide file tree

Showing 26 changed files with 1,570 additions and 1,113 deletions.
diff --git a/.github/workflows/CI.yaml b/.github/workflows/CI.yaml
@@ -19,11 +19,11 @@ jobs:
       matrix:
         cfg:
           - conda-env: env
-            python-version: 3.8
+            python-version: 3.11
             os: ubuntu-latest
 
           - conda-env: env
-            python-version: 3.8
+            python-version: 3.11
             os: macOS-latest
 
     env:
@@ -42,18 +42,10 @@ jobs:
           ulimit -a
 
       - name: Create Environment
-        uses: conda-incubator/setup-[email protected]
+        uses: mamba-org/setup-micromamba@v1
         with:
-          activate-environment: test
           environment-file: devtools/conda-envs/${{ matrix.cfg.conda-env }}.yaml
-          python-version: ${{ matrix.cfg.python-version }}
-          auto-update-conda: true
-          auto-activate-base: false
-          show-channel-urls: true
-          mamba-version: "*"
-          miniforge-version: latest
-          miniforge-variant: Mambaforge
-          use-mamba: true
+          create-args: python=${{ matrix.python-version }}
 
       - name: Environment Information
         shell: bash -l {0}
@@ -64,12 +56,12 @@ jobs:
       - name: Install NAGL-MBIS
         shell: bash -l {0}
         run: |
-          python setup.py develop --no-deps
+          pip install -e . --no-build-isolation 
 
       - name: PyTest
         shell: bash -l {0}
         run: |
-          pytest -v --cov=naglmbis --cov-config=setup.cfg naglmbis/tests/ --cov-report=xml
+          pytest -v --cov=naglmbis --cov-config=pyproject.toml naglmbis/tests/ --cov-report=xml --color=yes
 
       - name: Codecov
         uses: codecov/codecov-action@v1

diff --git a/.github/workflows/Lint.yaml b/.github/workflows/Lint.yaml
@@ -7,28 +7,32 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7]
+        python-version: [3.11]
 
     steps:
       - uses: actions/checkout@v2
 
       - name: Python Setup
-        uses: actions/setup-python@v1
+        uses: actions/setup-python@v3
         with:
           python-version: ${{ matrix.python-version }}
 
-      - name: Create Environment
+      - name: Install lint
         shell: bash
         run: |
-          python -m pip install --upgrade pip
-          python setup.py develop --no-deps
+          pip install black isort Flake8-pyproject
 
-      - name: Install black
+      - name: Run black
         shell: bash
         run: |
-          pip install black
+          black naglmbis --check
 
-      - name: Run black
+      - name: Run isort
+        shell: bash
+        run: |
+          isort --check-only naglmbis
+
+      - name: Run flake8
         shell: bash
         run: |
-          black naglmbis --check
+          flake8 naglmbis
diff --git a/README.md b/README.md
@@ -9,10 +9,10 @@ package by SimonBoothroyd.
 
 ## Installation
 
-The required dependencies to run these models can be installed using ``conda``:
+The required dependencies to run these models can be installed using ``mamba`` and the provided environment file:
 
 ```bash
-conda install -c conda-forge -c dglteam nagl "dgl >=0.7" openff-toolkit pytorch-lightning qubekit
+mamba env create -f devtools/conda-envs/env.yaml
 ```
 
 You will then need to install this package from source, first clone the repository from github:
@@ -25,33 +25,35 @@ cd nagl-mbis
 With the nagl environment activate install the models via:
 
 ```bash
-python setup.py install
+pip install -e . --no-build-isolation 
 ```
 
 ## Quick start
-NAGL-MBIS offers some pre-trained models to compute conformation independent MBIS charges and volumes, these can be loaded
+NAGL-MBIS offers some pre-trained models to compute conformation independent MBIS charges, these can be loaded
 using the following code in a script
 
 ```python
-from naglmbis.models import load_volume_model, load_charge_model
+from naglmbis.models import load_charge_model
 
-# load two pre-trained models
-charge_model = load_charge_model(charge_model=1)
-volume_model = load_volume_model(volume_model=1)
+# load two pre-trained charge models
+charge_model = load_charge_model(charge_model="nagl-v1-mbis")
+# load a model trained to scf dipole and mbis charges
+charge_model_2 = load_charge_model(charge_model="nagl-v1-mbis-dipole")
 ```
 
-we can then use these models to predict the corresponding properties for a given [openff-toolkit](https://github.com/openforcefield/openff-toolkit) [Molecule object](https://docs.openforcefield.org/projects/toolkit/en/stable/users/molecule_cookbook.html#cookbook-every-way-to-make-a-molecule).
+we can then use these models to predict the corresponding properties for a given [openff-toolkit](https://github.com/openforcefield/openff-toolkit) [Molecule object](https://docs.openforcefield.org/projects/toolkit/en/stable/users/molecule_cookbook.html#cookbook-every-way-to-make-a-molecule) or rdkit `Chem.Mol`.
 
 ```python
 from openff.toolkit.topology import Molecule
 
 # create ethanol
 ethanol = Molecule.from_smiles("CCO")
 # predict the charges (in e) and atomic volumes in (bohr ^3)
-charges = charge_model.compute_properties(ethanol)["mbis-charges"]
-volumes = volume_model.compute_properties(ethanol)["mbis-volumes"]
+charges = charge_model.compute_properties(ethanol.to_rdkit())["mbis-charges"]
+volumes = charge_model_2.compute_properties(ethanol.to_rdkit())["mbis-volumes"]
 ```
 
+# This is currently broken, due to plugins changing in the openff stack!
 Alternatively we provide an openff-toolkit parameter handler plugin which allows you to create an openmm system
 using the normal python pathway with a modified force field which requests that the ``NAGMBIS`` model be used to 
 predict charges and LJ parameters. We provide a function which can modify any offxml to add the custom handler
@@ -70,21 +72,21 @@ openmm_system = nagl_sage.create_openmm_system(topology=methanol.to_topology())
 
 # Models
 
-## MBISGraphModelV1
+## MBISGraphMode
 
 This model uses a minimal set of basic atomic features including
 
 - one hot encoded element
 - the number of bonds
-- ring membership of size 3-6
+- ring membership of size 3-8
 - n_gcn_layers 5
 - n_gcn_hidden_features 128
 - n_mbis_layers 2
 - n_mbis_hidden_features 64
 - learning_rate 0.001
-- n_epochs 100
+- n_epochs 1000
 
 These models were trained on the [OpenFF ESP Fragment Conformers v1.0](https://github.com/openforcefield/qca-dataset-submission/tree/master/submissions/2022-01-16-OpenFF-ESP-Fragment-Conformers-v1.0) dataset
-which is on QCArchive. The dataset was computed using HF/6-31G* with PSI4.  
+which is on QCArchive. The dataset was computed using HF/6-31G* with PSI4 and was split 80:10:10 using the deepchem maxmin spliter.  
 
 
diff --git a/devtools/conda-envs/env.yaml b/devtools/conda-envs/env.yaml
@@ -1,7 +1,6 @@
-name: test
+name: naglmbis
 channels:
   - conda-forge
-  - dglteam
 
 dependencies:
   - python
@@ -12,16 +11,23 @@ dependencies:
   - pytest-cov
 
   # core deps
-  - nagl
-  - openff-toolkit-base <0.11.0, >=0.10.6
+#  - nagl >=0.0.9
+  - pyarrow
+  - dgl >=1
+  - pytorch
+  - pytorch-lightning
+  - rich
+  - click
+  - click-option-group
+  - openff-toolkit-base
   - qcengine >=0.18.0
   - jinja2
   - chemper
   - torsiondrive
   # - qubekit
-  - pytorch-lightning
-  - dgl <1.0.0, >=0.7
-  - openff-utilities <=0.1.3
+  - openff-utilities
+  - pydantic <2
   - pip:
+      - git+https://github.com/SimonBoothroyd/nagl.git@main
       - espaloma_charge
-      - git+https://github.com/qubekit/QUBEKit.git
+      - git+https://github.com/qubekit/QUBEKit.git@main
diff --git a/naglmbis/__init__.py b/naglmbis/__init__.py
@@ -2,10 +2,10 @@
 naglmbis
 Models built with NAGL to predict MBIS properties.
 """
+from . import _version
 
-from ._version import get_versions
+__version__ = _version.get_versions()["version"]
+# make sure all custom features are loaded
+import naglmbis.features
 
-versions = get_versions()
-__version__ = versions["version"]
-__git_revision__ = versions["full-revisionid"]
-del get_versions, versions
+__all__ = [naglmbis.features]