"""
A collection of dielectric physical property definitions.
"""
import copy
import numpy as np
import pint
from simtk import openmm
from simtk.openmm import XmlSerializer
from openff.evaluator import unit
from openff.evaluator.attributes import UNDEFINED
from openff.evaluator.datasets import PhysicalProperty, PropertyPhase
from openff.evaluator.datasets.thermoml import thermoml_property
from openff.evaluator.layers import register_calculation_schema
from openff.evaluator.layers.reweighting import ReweightingLayer, ReweightingSchema
from openff.evaluator.layers.simulation import SimulationLayer, SimulationSchema
from openff.evaluator.protocols import analysis, reweighting
from openff.evaluator.protocols.utils import (
generate_base_reweighting_protocols,
generate_base_simulation_protocols,
generate_gradient_protocol_group,
)
from openff.evaluator.thermodynamics import ThermodynamicState
from openff.evaluator.utils import timeseries
from openff.evaluator.utils.statistics import bootstrap
from openff.evaluator.workflow import WorkflowSchema, workflow_protocol
from openff.evaluator.workflow.attributes import InputAttribute, OutputAttribute
from openff.evaluator.workflow.utils import ProtocolPath
[docs]@workflow_protocol()
class ReweightDielectricConstant(reweighting.BaseMBARProtocol):
"""Reweights a set of dipole moments (`reference_observables`) and volumes
(`reference_volumes`) using MBAR, and then combines these to yeild the reweighted
dielectric constant. Uncertainties in the dielectric constant are determined
by bootstrapping.
"""
reference_dipole_moments = InputAttribute(
docstring="A Quantity wrapped np.ndarray of the dipole moments of each "
"of the reference states.",
type_hint=list,
default_value=UNDEFINED,
)
reference_volumes = InputAttribute(
docstring="A Quantity wrapped np.ndarray of the volumes of each of the "
"reference states.",
type_hint=list,
default_value=UNDEFINED,
)
thermodynamic_state = InputAttribute(
docstring="The thermodynamic state at which the trajectory was generated.",
type_hint=ThermodynamicState,
default_value=UNDEFINED,
)
[docs] def __init__(self, protocol_id):
super().__init__(protocol_id)
self.bootstrap_uncertainties = True
def _bootstrap_function(
self,
reference_reduced_potentials,
target_reduced_potentials,
**reference_observables,
):
assert len(reference_observables) == 3
transposed_observables = {}
for key in reference_observables:
transposed_observables[key] = np.transpose(reference_observables[key])
values, _, _ = self._reweight_observables(
np.transpose(reference_reduced_potentials),
np.transpose(target_reduced_potentials),
**transposed_observables,
)
average_squared_dipole = values["dipoles_sqr"]
average_dipole_squared = np.linalg.norm(values["dipoles"])
dipole_variance = (average_squared_dipole - average_dipole_squared) * (
unit.elementary_charge * unit.nanometers
) ** 2
volume = values["volumes"] * unit.nanometer ** 3
e0 = 8.854187817e-12 * unit.farad / unit.meter # Taken from QCElemental
dielectric_constant = 1.0 + dipole_variance / (
3
* unit.boltzmann_constant
* self.thermodynamic_state.temperature
* volume
* e0
)
return dielectric_constant
def _execute(self, directory, available_resources):
if len(self.reference_dipole_moments) == 0:
raise ValueError("There were no dipole moments to reweight.")
if len(self.reference_volumes) == 0:
raise ValueError("There were no volumes to reweight.")
if not isinstance(
self.reference_dipole_moments[0], pint.Quantity
) or not isinstance(self.reference_volumes[0], pint.Quantity):
raise ValueError(
"The reference observables should be a list of "
"pint.Quantity wrapped ndarray's.",
)
if len(self.reference_dipole_moments) != len(self.reference_volumes):
raise ValueError(
"The number of reference dipoles does not match the "
"number of reference volumes.",
)
for reference_dipoles, reference_volumes in zip(
self.reference_dipole_moments, self.reference_volumes
):
if len(reference_dipoles) == len(reference_volumes):
continue
raise ValueError(
"The number of reference dipoles does not match the "
"number of reference volumes.",
)
self._reference_observables = self.reference_dipole_moments
dipole_moments = self._prepare_observables_array(self.reference_dipole_moments)
dipole_moments_sqr = np.array(
[[np.dot(dipole, dipole) for dipole in np.transpose(dipole_moments)]]
)
volumes = self._prepare_observables_array(self.reference_volumes)
if self.bootstrap_uncertainties:
self._execute_with_bootstrapping(
unit.dimensionless,
dipoles=dipole_moments,
dipoles_sqr=dipole_moments_sqr,
volumes=volumes,
)
else:
raise ValueError(
"Dielectric constant can only be reweighted in conjunction "
"with bootstrapped uncertainties.",
)
[docs]@thermoml_property(
"Relative permittivity at zero frequency",
supported_phases=PropertyPhase.Liquid,
)
class DielectricConstant(PhysicalProperty):
"""A class representation of a dielectric property"""
[docs] @classmethod
def default_unit(cls):
return unit.dimensionless
[docs] @staticmethod
def default_simulation_schema(
absolute_tolerance=UNDEFINED, relative_tolerance=UNDEFINED, n_molecules=1000
):
"""Returns the default calculation schema to use when estimating
this class of property from direct simulations.
Parameters
----------
absolute_tolerance: pint.Quantity, optional
The absolute tolerance to estimate the property to within.
relative_tolerance: float
The tolerance (as a fraction of the properties reported
uncertainty) to estimate the property to within.
n_molecules: int
The number of molecules to use in the simulation.
Returns
-------
SimulationSchema
The schema to follow when estimating this property.
"""
assert absolute_tolerance == UNDEFINED or relative_tolerance == UNDEFINED
calculation_schema = SimulationSchema()
calculation_schema.absolute_tolerance = absolute_tolerance
calculation_schema.relative_tolerance = relative_tolerance
# Define the protocol which will extract the average dielectric constant
# from the results of a simulation.
extract_dielectric = ExtractAverageDielectric("extract_dielectric")
extract_dielectric.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global"
)
# Define the protocols which will run the simulation itself.
use_target_uncertainty = (
absolute_tolerance != UNDEFINED or relative_tolerance != UNDEFINED
)
protocols, value_source, output_to_store = generate_base_simulation_protocols(
extract_dielectric,
use_target_uncertainty,
n_molecules=n_molecules,
)
# Make sure the input of the analysis protcol is properly hooked up.
extract_dielectric.system_path = ProtocolPath(
"system_path", protocols.assign_parameters.id
)
# Dielectric constants typically take longer to converge, so we need to
# reflect this in the maximum number of convergence iterations.
protocols.converge_uncertainty.max_iterations = 400
# Set up the gradient calculations. For dielectric constants, we need to use
# a slightly specialised reweighting protocol which we set up here.
coordinate_source = ProtocolPath(
"output_coordinate_file", protocols.equilibration_simulation.id
)
trajectory_source = ProtocolPath(
"trajectory_file_path",
protocols.converge_uncertainty.id,
protocols.production_simulation.id,
)
statistics_source = ProtocolPath(
"statistics_file_path",
protocols.converge_uncertainty.id,
protocols.production_simulation.id,
)
gradient_mbar_protocol = ReweightDielectricConstant("gradient_mbar")
gradient_mbar_protocol.reference_dipole_moments = [
ProtocolPath(
"dipole_moments",
protocols.converge_uncertainty.id,
extract_dielectric.id,
)
]
gradient_mbar_protocol.reference_volumes = [
ProtocolPath(
"volumes", protocols.converge_uncertainty.id, extract_dielectric.id
)
]
gradient_mbar_protocol.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global"
)
gradient_mbar_protocol.reference_reduced_potentials = statistics_source
(
gradient_group,
gradient_replicator,
gradient_source,
) = generate_gradient_protocol_group(
gradient_mbar_protocol,
ProtocolPath("force_field_path", "global"),
coordinate_source,
trajectory_source,
statistics_source,
)
# Build the workflow schema.
schema = WorkflowSchema()
schema.protocol_schemas = [
protocols.build_coordinates.schema,
protocols.assign_parameters.schema,
protocols.energy_minimisation.schema,
protocols.equilibration_simulation.schema,
protocols.converge_uncertainty.schema,
protocols.extract_uncorrelated_trajectory.schema,
protocols.extract_uncorrelated_statistics.schema,
gradient_group.schema,
]
schema.protocol_replicators = [gradient_replicator]
schema.outputs_to_store = {"full_system": output_to_store}
schema.gradients_sources = [gradient_source]
schema.final_value_source = value_source
calculation_schema.workflow_schema = schema
return calculation_schema
[docs] @staticmethod
def default_reweighting_schema(
absolute_tolerance=UNDEFINED,
relative_tolerance=UNDEFINED,
n_effective_samples=50,
):
"""Returns the default calculation schema to use when estimating
this property by reweighting existing data.
Parameters
----------
absolute_tolerance: pint.Quantity, optional
The absolute tolerance to estimate the property to within.
relative_tolerance: float
The tolerance (as a fraction of the properties reported
uncertainty) to estimate the property to within.
n_effective_samples: int
The minimum number of effective samples to require when
reweighting the cached simulation data.
Returns
-------
ReweightingSchema
The schema to follow when estimating this property.
"""
assert absolute_tolerance == UNDEFINED or relative_tolerance == UNDEFINED
calculation_schema = ReweightingSchema()
calculation_schema.absolute_tolerance = absolute_tolerance
calculation_schema.relative_tolerance = relative_tolerance
data_replicator_id = "data_replicator"
# Set up a protocol to extract the dielectric constant from the stored data.
extract_dielectric = ExtractAverageDielectric(
f"calc_dielectric_$({data_replicator_id})"
)
# For the dielectric constant, we employ a slightly more advanced reweighting
# protocol set up for calculating fluctuation properties.
reweight_dielectric = ReweightDielectricConstant("reweight_dielectric")
reweight_dielectric.reference_dipole_moments = ProtocolPath(
"uncorrelated_values", extract_dielectric.id
)
reweight_dielectric.reference_volumes = ProtocolPath(
"uncorrelated_volumes", extract_dielectric.id
)
reweight_dielectric.thermodynamic_state = ProtocolPath(
"thermodynamic_state", "global"
)
reweight_dielectric.bootstrap_uncertainties = True
reweight_dielectric.bootstrap_iterations = 200
reweight_dielectric.required_effective_samples = n_effective_samples
protocols, data_replicator = generate_base_reweighting_protocols(
extract_dielectric, reweight_dielectric, data_replicator_id
)
# Make sure input is taken from the correct protocol outputs.
extract_dielectric.system_path = ProtocolPath(
"system_path", protocols.build_reference_system.id
)
extract_dielectric.thermodynamic_state = ProtocolPath(
"thermodynamic_state", protocols.unpack_stored_data.id
)
# Set up the gradient calculations
coordinate_path = ProtocolPath(
"output_coordinate_path", protocols.concatenate_trajectories.id
)
trajectory_path = ProtocolPath(
"output_trajectory_path", protocols.concatenate_trajectories.id
)
statistics_path = ProtocolPath(
"statistics_file_path", protocols.reduced_target_potential.id
)
reweight_dielectric_template = copy.deepcopy(reweight_dielectric)
(
gradient_group,
gradient_replicator,
gradient_source,
) = generate_gradient_protocol_group(
reweight_dielectric_template,
ProtocolPath("force_field_path", "global"),
coordinate_path,
trajectory_path,
statistics_path,
replicator_id="grad",
effective_sample_indices=ProtocolPath(
"effective_sample_indices", reweight_dielectric.id
),
)
schema = WorkflowSchema()
schema.protocol_schemas = [
*(x.schema for x in protocols),
gradient_group.schema,
]
schema.protocol_replicators = [data_replicator, gradient_replicator]
schema.gradients_sources = [gradient_source]
schema.final_value_source = ProtocolPath("value", protocols.mbar_protocol.id)
calculation_schema.workflow_schema = schema
return calculation_schema
# Register the properties via the plugin system.
register_calculation_schema(
DielectricConstant, SimulationLayer, DielectricConstant.default_simulation_schema
)
register_calculation_schema(
DielectricConstant, ReweightingLayer, DielectricConstant.default_reweighting_schema
)