From eba0ded6439182a4adaeceb1fba23718e1fcad59 Mon Sep 17 00:00:00 2001
From: Bas Rustenburg <bas.rustenburg@choderalab.org>
Date: Wed, 3 Apr 2019 15:56:33 -0400
Subject: [PATCH] Add sampling method option to tautomer drives

This adds the samplingmethod constructor argument to the tautomer
driver classes.
---
 protons/app/driver.py | 36 ++++++++++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/protons/app/driver.py b/protons/app/driver.py
index 35b219a..f68981f 100644
--- a/protons/app/driver.py
+++ b/protons/app/driver.py
@@ -4159,13 +4159,34 @@ def strip_in_unit_system(quant, unit_system=unit.md_unit_system, compatible_with
 class TautomerNCMCProtonDrive(NCMCProtonDrive):
     def __init__(
         self,
-        temperature,
-        topology,
-        system,
-        pressure=None,
-        perturbations_per_trial=0,
-        propagations_per_step=1,
+        temperature: unit.Quantity,
+        topology: app.Topology,
+        system: mm.System,
+        pressure: Optional[unit.Quantity] = None,
+        perturbations_per_trial: int = 0,
+        propagations_per_step: int = 1,
+        sampling_method: SamplingMethod = SamplingMethod.MCMC,
     ):
+        """
+        Initialize a Monte Carlo titration driver for simulation of protonation states and tautomers.
+
+        Parameters
+        ----------
+        temperature : simtk.unit.Quantity compatible with kelvin
+            Temperature at which the system is to be simulated.
+        topology : protons.app.Topology
+            OpenMM object containing the topology of system
+        system : simtk.openmm.System
+            System to be titrated, containing all possible protonation sites.
+        pressure : simtk.unit.Quantity compatible with atmospheres, optional
+            For explicit solvent simulations, the pressure.
+        perturbations_per_trial : int, optional, default=0
+            Number of perturbation steps per NCMC switching trial, or 0 if instantaneous Monte Carlo is to be used.
+        propagations_per_step : int, optional, default=1
+            Number of propagation steps in between perturbation steps.
+        sampling_method : The method of sampling that is used.
+            See the SamplingMethod enum for the set of supported options (including MCMC and importance sampling).
+        """
 
         super().__init__(
             temperature,
@@ -4174,6 +4195,7 @@ def __init__(
             pressure=pressure,
             perturbations_per_trial=perturbations_per_trial,
             propagations_per_step=propagations_per_step,
+            sampling_method=sampling_method,
         )
         # Store force object pointers.
         force_classes_to_update = [
@@ -4688,6 +4710,7 @@ def __init__(
         propagations_per_step=1,
         residues_by_name=None,
         residues_by_index=None,
+        sampling_method: SamplingMethod = SamplingMethod.MCMC,
     ):
 
         """
@@ -4741,6 +4764,7 @@ def __init__(
             pressure,
             perturbations_per_trial=perturbations_per_trial,
             propagations_per_step=propagations_per_step,
+            sampling_method=sampling_method,
         )
 
         ffxml_residues = self._parse_ffxml_files(ffxml_files)