Update documentation and clean up repo

.github/workflows/CI.yaml

@@ -26,7 +26,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python-version: [3.8, 3.9]
+        python-version: [3.9,"3.10", "3.11"]
 
     steps:
     - uses: actions/checkout@v1

README.md

@@ -1,4 +1,4 @@
-Endstate correction from MM to QML potential
+Endstate correction from MM to NNP
 ==============================
 [//]: # (Badges)
 [![GitHub Actions Build Status](https://github.com/wiederm/endstate_correction/workflows/CI/badge.svg)](https://github.com/wiederm/endstate_correction/actions?query=workflow%3ACI)

docs/conf.py

@@ -186,7 +186,7 @@
         "endstate_correction Documentation",
         author,
         "endstate_correction",
-        "Endstate reweighting from MM to QML potential",
+        "Endstate reweighting from MM to NNP",
         "Miscellaneous",
     ),
 ]

docs/getting_started.rst

@@ -1,6 +1,6 @@
 Getting Started
 ===============
-This page details how to perform free energy calculations between a molecular mechanics force field (in this case the open-forcefield is sued) and a neural network potential (ANI-2x).
+This page details how to perform free energy calculations between a molecular mechanics force field (in this case the open-forcefield is used) and a neural network potential (ANI-2x).
 
 Installation
 -----------------
@@ -12,9 +12,9 @@ This package can be installed using:
 How to use this package
 -----------------
 We have prepared two scripts that should help to use this package, both are located in :code:`endstate_correction/scripts`.
-We will start by describing the use of the :code:`sampling.py` script and then discuss the :code:`perform_correction.py` script.
+We will start by describing the use of the :code:`generate_endstate_samples.py` script and then discuss the :code:`perform_correction.py` script.
 
-A typical NEQ workflow
+A typical Non-equilibrium (NEQ) switching workflow
 -----------------
 Generate the equilibrium distribution :math:`\pi(x, \lambda=0)`
 ~~~~~~~~~~~~~~~~~~~~~~
@@ -57,49 +57,53 @@ Note that we explicitly define the atoms that should be perturbed from the refer
 
 If you want to perform bidirectional FEP or NEQ you need to sample at :math:`\pi(x, \lambda=0)` *and* :math:`\pi(x, \lambda=1)`. 
 This can be controlled by setting the number using the variable :code:`nr_lambda_states`.
-The default value set in the :code:`sampling.py` script is :code:`nr_lambda_states=2`, generating samples from both endstate equilibrium distributions.
+The default value set in the :code:`generate_endstate_samples.py` script is :code:`nr_lambda_states=2`, generating samples from both endstate equilibrium distributions.
 
 .. code:: python
+
     nr_lambda_states = 2  # samples equilibrium distribution at both endstates
     lambs = np.linspace(0, 1, nr_lambda_states)
 
 Perform unidirectional NEQ from :math:`\pi(x, \lambda=0)`
 ~~~~~~~~~~~~~~~~~~~~~~
 The endstate correction can be performed using the script :code:`perform_correction.py`.
 The script will calculate the free energy estimate using the samples generated with the :code:`generate_endstate_samples.py` script.
-Subsequently, the relevant section of the :code:`perform_correction.py` script are explained --- but they should work for for the testsystem without any modifications. 
+Subsequently, the relevant sections of the :code:`perform_correction.py` script are explained --- but they should work for for the testsystem without any modifications. 
 
 To perform a specific endstate correction we need to define a protocol 
-(some standard protocols are shown :ref:`here<Available protocols>`) 
+(some standard protocols are shown :ref:`here <available protocols>`) 
 with:
 
 .. code:: python
 
-    neq_protocol = Protocol(
-        method="NEQ",
+    neq_protocol = NEQProtocol(
         sim=sim,
         reference_samples=mm_samples,
         nr_of_switches=100,
         neq_switching_length=1_000,
+        save_endstates=False,
+        save_trajs=False
     )
 
 This protocol is then passed to the actual function performing the protocol: :code:`perform_endstate_correction(neq_protocol)`.
 The particular code above defines unidirectional NEQ switching using 100 switches and a switching length of 1 ps.
 The direciton of the switching simulation is controlled by the sampels that are provided: 
-if `reference_samples` are provided, switching is performed from the reference to the target level of theory, if `target_samples` are provided, switching is performed from the target level to the reference level.
+if :code:`reference_samples` are provided, switching is performed from the reference to the target level of theory, if :code:`target_samples` are provided, switching is performed from the target level to the reference level.
 If both samples are provided, bidirectional NEQ switching is performed (for an example see below).
+There is the possibility to save the endstate of each switch if :code:`save_endstates` is set to :code:`True`. 
+Additionally, the switching trajectory of each switch can be saved if :code:`save_trajs=True`.
+Both options only make sense for a NEQ protocol.
 
 Perform bidirectional NEQ from :math:`\pi(x, \lambda=0)` and :math:`\pi(x, \lambda=1)`
 ~~~~~~~~~~~~~~~~~~~~~~
 The endstate correction can be performed using the script :code:`perform_correction.py` and the following protocol.
 
 .. code:: python
 
-    neq_protocol = Protocol(
-        method="NEQ",
+    neq_protocol = NEQProtocol(
         sim=sim,
         reference_samples=mm_samples,
-        target_samples=qml_samples,
+        target_samples=nnp_samples,
         nr_of_switches=100,
         neq_switching_length=1_000,
     )
@@ -112,45 +116,110 @@ Perform unidirectional FEP from :math:`\pi(x, \lambda=0)`
 The protocol has to be adopted slightly:
 
 .. code:: python
-    fep_protocol = Protocol(
-        method="FEP",
+
+    fep_protocol = FEPProtocol(
         sim=sim,
         reference_samples=mm_samples,
         nr_of_switches=1_000,
     )
 
 This protocol is then passed to the actual function performing the protocol: :code:`perform_endstate_correction(fep_protocol)`.
 
+Perform unidirectional SMC from :math:`\pi(x, \lambda=0)`
+~~~~~~~~~~~~~~~~~~~~~~
+To perform a unidirectional Sequential Monte-Carlo (SMC) protocol, the following adaptations have to be made:
+
+.. code:: python
+
+    smc_protocol = SMCProtocol(
+        sim=sim,
+        reference_samples=mm_samples,
+        nr_of_walkers=100,
+        protocol_length=1_000
+        nr_of_resampling_steps=100,
+        save_endstates=True
+    )
+Here, :code:`nr_of_walkers` indicates the number of walkers selected from the reference distribution, which are propagated for 1ps (:code:`protocol_length`). 
+:code:`nr_of nr_of_resampling_steps` states that the walkers are resampled 100 times.
+With the option :code:`save_endstates` samples at the end of the protocol can be saved.
+
+Perform multiple protocols
+~~~~~~~~~~~~~~~~~~~~~~
+It is possible to combine several protocols into one with :code:`AllProtocol`. 
+
+.. code:: python
+
+    protocol = AllProtocol(
+        fep_protocol=fep_protocol, 
+        neq_protocol=neq_protocol, 
+        smc_protocol=smc_protocol
+    )
+
+This :code:`protocol` is then passed to :code:`perform_endstate_correction()` to perform all endstate corrections.
 
 Analyse results of an unidirection NEQ protocol
 ~~~~~~~~~~~~~~~~~~~~~~
 To analyse the results generated by :code:`r = perform_endstate_correction()` pass the return value to :code:`plot_endstate_correction_results(system_name, r, "results_neq_unidirectional.png")` and results will be plotted and printed.
+It should be noted, that :code:`perform_endstate_correction()` returns an instance of :code:`AllResults()`, which, depending on the passed protocol, contains results of an FEP, NEQ switching and SMC protocol (instances of the :code:`FEPResults()`, :code:`NEQResults()` and :code:`SMCResults()` class, respectively).
+It is also possible to add equilibrium sampling results to the instance of :code:`AllResults()` retrospectively (:code:`EquResults()`).
 
-
+.. _available protocols:
 Available protocols
 -----------------
 
+unidirectional NEQ protocol (reference to target)
+
 .. code:: python
 
-    neq_protocol = Protocol(
-        method="NEQ",
+    neq_protocol = NEQProtocol(
         sim=sim,
         reference_samples=mm_samples,
-        target_samples=qml_samples,
         nr_of_switches=100,
         neq_switching_length=1_000,
     )
 
+bidirectional NEQ protocol (save switching trajectory and endstate of each switch)
+
 .. code:: python
 
-    neq_protocol = Protocol(
-        method="FEP",
+    neq_protocol = NEQProtocol(
         sim=sim,
         reference_samples=mm_samples,
-        target_samples=qml_samples,
+        target_samples=nnp_samples,
         nr_of_switches=100,
         neq_switching_length=1_000,
         save_endstates=True,
         save_trajs=True,
     )
 
+unidirectional FEP protocol (reference to target)
+
+.. code:: python
+
+    fep_protocol = FEPProtocol(
+        sim=sim,
+        reference_samples=mm_samples,
+        nr_of_switches=1_000,
+    )
+
+bidirectional FEP protocol
+
+.. code:: python
+
+    fep_protocol = FEPProtocol(
+        sim=sim,
+        reference_samples=mm_samples,
+        target_samples=nnp_samples,
+        nr_of_switches=1_000,
+    )
+
+unidirectional SMC protocol (reference to target)
+
+.. code:: python
+
+    smc_protocol = SMCProtocol(
+        sim=sim,
+        reference_samples=mm_samples,
+        nr_of_walkers=100,
+        nr_of_resampling_steps=1_000,
+    )

docs/requirements.yaml

@@ -20,7 +20,7 @@ dependencies:
   - pytorch=1.11.0
 
   # Documentation
-  - sphinx
+  - sphinx<7.0
   - sphinx_rtd_theme
   - myst-parser
   # Testing

docs/theory.rst

@@ -2,21 +2,19 @@ Theory
 ===============
 
 
-Equilibrium free energy endstate corrections
+Multistage equilibrium free energy calculations
 -----------------
 
-In equilibrium free energy calculations samples are drawn from the Boltzmann distrubtion 
+In multistage equilibrium free energy calculations samples are drawn from the Boltzmann distrubtion 
 at specific interpolation states between thermodynamic states (in our specific case: different energetic
-descriptions of the molecular system, i.e. the source level of theory and the target level of theroy) and, 
+descriptions of the molecular system, i.e. the reference level of theory and the target level of theroy) and, 
 given sufficient overlap of its pdfs, a free energy can be estimated. This protocol is expensive 
 (it needs iid samples at each lambda state connecting the Boltzmann distribution at the endstates) 
 but also reliable and accureate (with low variance).
 
-.. figure:: images/equi.png
 
 
-
-Non-equilibrium work protocol 
+Non-equilibrium (NEQ) work protocol 
 -----------------
 
 Non-equilibrium work protocols, and the fluctuation theorems connecting non-equilibrium driven 
@@ -26,11 +24,11 @@ A specific NEQ protocol typically consists of a series of perturbation kernel  :
 propagation kernel  :math:`\kappa_t(x,y)`, which are used in an alternating pattern to drive the system
 out of equilibrium.
 Each perturbation kernel :math:`\alpha` drives an alchemical coupling parameter :math:`\lambda`, and each 
-propagation kernel :math:`\kappa`$` propagates the coordinates of the system at fixed :math:`\lambda`$` according 
+propagation kernel :math:`\kappa` propagates the coordinates of the system at fixed :math:`\lambda` according 
 to a defined MD process.
 The free energy difference can then be recovered using either the Jarzynski equation (if initial conformations 
-to seed the NEQ protocol are only drawn from :math:`\pi(x, \lambda=0)` and the NEQ protocol perturbations only 
-from :math:`\lambda=0` to :math:`\lambda=1`) or the Crooks' fluctuation theorem (if samples to seed the NEQ protocol 
+to seed the NEQ protocol are only drawn from :math:`\pi(x, \lambda=0)` (OR :math:`\pi(x, \lambda=1)`)) and the NEQ protocol perturbations only 
+from :math:`\lambda=0` to :math:`\lambda=1` (OR :math:`\lambda=1` to :math:`\lambda=0`) or the Crooks' fluctuation theorem (if samples to seed the NEQ protocol 
 are drawn from :math:`\pi(x, \lambda=0)` and :math:`\pi(x, \lambda=1)` and the perturbation kernels are set for a bidirectinoal 
 protocol).
 
@@ -43,4 +41,4 @@ consists only of a single perturbation kernel :math:`\alpha_t(x,y)`, without a p
 to another 'endstate'. 
 In the limiting cases of infinitely fast switching the Jarzynski equality reduces to the well-known FEP equation:
 :math:`e^{-\beta \Delta F} = \langle e^{−β[E(x,\lambda=1)− E(x,\lambda=0)]} \rangle_{\lambda=0}`.
-:math:`\langle \rangle_{\lambda=0}` indicate that samples are drawn from the equilibrium distribution :math:`\pi(x, \lambda=0)`.
+:math:`\langle \rangle_{\lambda=0}` indicates that samples are drawn from the equilibrium distribution :math:`\pi(x, \lambda=0)`.

endstate_correction/__init__.py

@@ -1,4 +1,4 @@
-"""Endstate reweighting from MM to QML potential"""
+"""Endstate reweighting from MM to NNP"""
 
 # Add imports here
 from .endstate_correction import *