<commit_before>#!/usr/bin/python

# =============================================================================================
# MODULE DOCSTRING
# =============================================================================================

"""
Test restraints module.

"""

# =============================================================================================
# GLOBAL IMPORTS
# =============================================================================================

import tempfile
import os
import shutil
import math
import numpy as np
import netCDF4 as netcdf

from nose.plugins.attrib import attr

import yank.restraints
from yank.repex import ThermodynamicState
from yank.yamlbuild import YamlBuilder
from yank.utils import get_data_filename
from yank import analyze

from simtk import unit
from openmmtools import testsystems

# =============================================================================================
# UNIT TESTS
# =============================================================================================

from openmmtools.testsystems import HostGuestVacuum


class HostGuestNoninteracting(HostGuestVacuum):
    """CB7:B2 host-guest system in vacuum with no nonbonded interactions.

    Parameters
    ----------
    Same as HostGuestVacuum

    Examples
    --------
    Create host:guest system with no nonbonded interactions.
    >>> testsystem = HostGuestVacuumNoninteracting()
    >>> (system, positions) = testsystem.system, testsystem.positions

    Properties
    ----------
    receptor_atoms : list of int
        Indices of receptor atoms
    ligand_atoms : list of int
        Indices of ligand atoms

    """
    def __init__(self, **kwargs):
        super(HostGuestNoninteracting, self).__init__(**kwargs)

        # Store receptor and ligand atom indices
        self.receptor_atoms = range(0,126)
        self.ligand_atoms = range(126,156)

        # Remove nonbonded interactions
        force_indices = { self.system.getForce(index).__class__.__name__ : index for index in range(self.system.getNumForces()) }
        self.system.removeForce(force_indices['NonbondedForce'])

expected_restraints = {
    'Harmonic' : yank.restraints.Harmonic,
    'FlatBottom' : yank.restraints.FlatBottom,
    'Boresch' : yank.restraints.Boresch,
}

restraint_test_yaml = """
---
options:
  minimize: no
  verbose: yes
  output_dir: %(output_directory)s
  number_of_iterations: %(number_of_iter)s
  nsteps_per_iteration: 10
  temperature: 300*kelvin
  pressure: null
  anisotropic_dispersion_correction: no
  platform: Reference

solvents:
  vacuum:
    nonbonded_method: PME
    nonbonded_cutoff: 0.59 * nanometer

systems:
  ship:
    phase1_path: [data/benzene-toluene-standard-state/standard_state_complex.inpcrd, data/benzene-toluene-standard-state/standard_state_complex.prmtop]
    phase2_path: [data/benzene-toluene-standard-state/standard_state_complex.inpcrd, data/benzene-toluene-standard-state/standard_state_complex.prmtop]
    ligand_dsl: resname ene
    solvent: vacuum

protocols:
  absolute-binding:
    complex:
      alchemical_path:
        lambda_restraints:     [0.0, 0.25, 0.5, 0.75, 1.0]
        lambda_electrostatics: [0.0, 0.00, 0.0, 0.00, 0.0]
        lambda_sterics:        [0.0, 0.00, 0.0, 0.00, 0.0]
    solvent:
      alchemical_path:
        lambda_electrostatics: [1.0, 1.0]
        lambda_sterics:        [1.0, 1.0]

experiments:
  system: ship
  protocol: absolute-binding
  restraint:
    type: %(restraint_type)s
"""


def general_restraint_run(options):
    """
    Generalized restraint simulation run to test free energy = standard state correction.

    options : Dict. A dictionary of substitutions for restraint_test_yaml
    """
    output_directory = tempfile.mkdtemp()
    options['output_directory'] = output_directory
    # run both setup and experiment
    yaml_builder = YamlBuilder(restraint_test_yaml % options)
    yaml_builder.build_experiments()
    # Estimate Free Energies
    ncfile_path = os.path.join(output_directory, 'experiments', 'complex.nc')
    ncfile = netcdf.Dataset(ncfile_path, 'r')
    (Deltaf_ij, dDeltaf_ij) = analyze.estimate_free_energies(ncfile)
    # Correct the sign for the fact that we are adding vs removing the restraints
    DeltaF_simulated = Deltaf_ij[-1, 0]
    dDeltaF_simulated = dDeltaf_ij[-1, 0]
    DeltaF_restraints = ncfile.groups['metadata'].variables['standard_state_correction'][0]
    ncfile.close()
    # Clean up
    shutil.rmtree(output_directory)
    # Check if they are close
    assert np.allclose(DeltaF_restraints, DeltaF_simulated, rtol=dDeltaF_simulated)


@attr('slow')  # Skip on Travis-CI
def test_harmonic_free_energy():
    """
    Test that the harmonic restraint simulated free energy equals the standard state correction
    """
    options = {'number_of_iter': '500',
               'restraint_type': 'Harmonic'}
    general_restraint_run(options)


@attr('slow')  # Skip on Travis-CI
def test_flat_bottom_free_energy():
    """
    Test that the harmonic restraint simulated free energy equals the standard state correction
    """
    options = {'number_of_iter': '500',
               'restraint_type': 'FlatBottom'}
    general_restraint_run(options)


@attr('slow')  # Skip on Travis-CI
def test_Boresch_free_energy():
    """
    Test that the harmonic restraint simulated free energy equals the standard state correction
    """
    options = {'number_of_iter': '500',
               'restraint_type': 'Boresch'}
    general_restraint_run(options)

<commit_msg>Documentation lines, changed default platform on restraint tests<commit_after>#!/usr/bin/python

# =============================================================================================
# MODULE DOCSTRING
# =============================================================================================

"""
Test restraints module.

"""

# =============================================================================================
# GLOBAL IMPORTS
# =============================================================================================

import tempfile
import os
import shutil
import math
import numpy as np
import netCDF4 as netcdf

from nose.plugins.attrib import attr

import yank.restraints
from yank.repex import ThermodynamicState
from yank.yamlbuild import YamlBuilder
from yank.utils import get_data_filename
from yank import analyze

from simtk import unit
from openmmtools import testsystems

# =============================================================================================
# UNIT TESTS
# =============================================================================================

from openmmtools.testsystems import HostGuestVacuum


class HostGuestNoninteracting(HostGuestVacuum):
    """CB7:B2 host-guest system in vacuum with no nonbonded interactions.

    Parameters
    ----------
    Same as HostGuestVacuum

    Examples
    --------
    Create host:guest system with no nonbonded interactions.
    >>> testsystem = HostGuestVacuumNoninteracting()
    >>> (system, positions) = testsystem.system, testsystem.positions

    Properties
    ----------
    receptor_atoms : list of int
        Indices of receptor atoms
    ligand_atoms : list of int
        Indices of ligand atoms

    """
    def __init__(self, **kwargs):
        super(HostGuestNoninteracting, self).__init__(**kwargs)

        # Store receptor and ligand atom indices
        self.receptor_atoms = range(0,126)
        self.ligand_atoms = range(126,156)

        # Remove nonbonded interactions
        force_indices = { self.system.getForce(index).__class__.__name__ : index for index in range(self.system.getNumForces()) }
        self.system.removeForce(force_indices['NonbondedForce'])

expected_restraints = {
    'Harmonic' : yank.restraints.Harmonic,
    'FlatBottom' : yank.restraints.FlatBottom,
    'Boresch' : yank.restraints.Boresch,
}

restraint_test_yaml = """
---
options:
  minimize: no
  verbose: yes
  output_dir: %(output_directory)s
  number_of_iterations: %(number_of_iter)s
  nsteps_per_iteration: 10
  temperature: 300*kelvin
  pressure: null
  anisotropic_dispersion_correction: no
  platform: OpenCL

solvents:
  vacuum:
    nonbonded_method: PME
    nonbonded_cutoff: 0.59 * nanometer

systems:
  ship:
    phase1_path: [data/benzene-toluene-standard-state/standard_state_complex.inpcrd, data/benzene-toluene-standard-state/standard_state_complex.prmtop]
    phase2_path: [data/benzene-toluene-standard-state/standard_state_complex.inpcrd, data/benzene-toluene-standard-state/standard_state_complex.prmtop]
    ligand_dsl: resname ene
    solvent: vacuum

protocols:
  absolute-binding:
    complex:
      alchemical_path:
        lambda_restraints:     [0.0, 0.25, 0.5, 0.75, 1.0]
        lambda_electrostatics: [0.0, 0.00, 0.0, 0.00, 0.0]
        lambda_sterics:        [0.0, 0.00, 0.0, 0.00, 0.0]
    solvent:
      alchemical_path:
        lambda_electrostatics: [1.0, 1.0]
        lambda_sterics:        [1.0, 1.0]

experiments:
  system: ship
  protocol: absolute-binding
  restraint:
    type: %(restraint_type)s
"""


def general_restraint_run(options):
    """
    Generalized restraint simulation run to test free energy = standard state correction.

    options : Dict. A dictionary of substitutions for restraint_test_yaml
    """
    output_directory = tempfile.mkdtemp()
    options['output_directory'] = output_directory
    # run both setup and experiment
    yaml_builder = YamlBuilder(restraint_test_yaml % options)
    yaml_builder.build_experiments()
    # Estimate Free Energies
    ncfile_path = os.path.join(output_directory, 'experiments', 'complex.nc')
    ncfile = netcdf.Dataset(ncfile_path, 'r')
    (Deltaf_ij, dDeltaf_ij) = analyze.estimate_free_energies(ncfile)
    # Correct the sign for the fact that we are adding vs removing the restraints
    DeltaF_simulated = Deltaf_ij[-1, 0]
    dDeltaF_simulated = dDeltaf_ij[-1, 0]
    DeltaF_restraints = ncfile.groups['metadata'].variables['standard_state_correction'][0]
    ncfile.close()
    # Clean up
    shutil.rmtree(output_directory)
    # Check if they are close
    assert np.allclose(DeltaF_restraints, DeltaF_simulated, rtol=dDeltaF_simulated)


@attr('slow')  # Skip on Travis-CI
def test_harmonic_free_energy():
    """
    Test that the harmonic restraint simulated free energy equals the standard state correction
    """
    options = {'number_of_iter': '500',
               'restraint_type': 'Harmonic'}
    general_restraint_run(options)


@attr('slow')  # Skip on Travis-CI
def test_flat_bottom_free_energy():
    """
    Test that the harmonic restraint simulated free energy equals the standard state correction
    """
    options = {'number_of_iter': '500',
               'restraint_type': 'FlatBottom'}
    general_restraint_run(options)


@attr('slow')  # Skip on Travis-CI
def test_Boresch_free_energy():
    """
    Test that the harmonic restraint simulated free energy equals the standard state correction
    """
    # These need more samples to converge
    options = {'number_of_iter': '1000',
               'restraint_type': 'Boresch'}
    general_restraint_run(options)


{'fraction_non_alphanumeric': 0.18702460850111857, 'fraction_numerical': 0.01458612975391499, 'mean_word_length': 4.08231014097317, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5498591', 'n_tokens_mistral': 3507, 'n_tokens_neox': 3216, 'n_words': 761}