def main():
setup_timestamp_logging()
# Retrieve the current version.
version = evaluator.__version__.replace(".", "-").replace("v", "")
if "+" in version:
version = "latest"
# Create a new directory to run the current versions results in.
os.makedirs(os.path.join(version, "results"))
with temporarily_change_directory(version):
with DaskLSFBackend(
minimum_number_of_workers=1,
maximum_number_of_workers=12,
resources_per_worker=QueueWorkerResources(
number_of_gpus=1,
preferred_gpu_toolkit=QueueWorkerResources.GPUToolkit.CUDA,
per_thread_memory_limit=5 * unit.gigabyte,
wallclock_time_limit="05:59",
),
setup_script_commands=[
f"conda activate openff-evaluator-{version}",
"module load cuda/10.0",
],
queue_name="gpuqueue",
) as calculation_backend:
with EvaluatorServer(
calculation_backend,
working_directory="outputs",
storage_backend=LocalFileStorage("cached-data"),
):
client = EvaluatorClient()
for allowed_layer in ["SimulationLayer", "ReweightingLayer"]:
data_set = define_data_set(
allowed_layer == "ReweightingLayer")
options = RequestOptions()
options.calculation_layers = [allowed_layer]
options.calculation_schemas = {
property_type: {}
for property_type in data_set.property_types
}
if allowed_layer == "SimulationLayer":
options.add_schema(
"SimulationLayer",
"SolvationFreeEnergy",
solvation_free_energy_schema(),
)
request, _ = client.request_estimate(
data_set,
ForceField("openff-1.2.0.offxml"),
options,
parameter_gradient_keys=[
ParameterGradientKey("vdW", smirks, attribute)
for smirks in [
"[#1:1]-[#6X4]",
"[#1:1]-[#6X4]-[#7,#8,#9,#16,#17,#35]",
"[#1:1]-[#8]",
"[#6X4:1]",
"[#8X2H1+0:1]",
"[#1]-[#8X2H2+0:1]-[#1]",
] for attribute in ["epsilon", "rmin_half"]
],
)
results, _ = request.results(synchronous=True,
polling_interval=60)
results.json(
os.path.join("results", f"{allowed_layer}.json"))
def _generate_request_options(
cls, target: EvaluatorTarget,
training_set: "PhysicalPropertyDataSet") -> "RequestOptions":
"""Generates the request options to use when estimating an evaluator
optimization targets.
Parameters
----------
target
The evaluator target which will spawn the request.
training_set
The training set which will be estimated.
Returns
-------
The request options.
"""
import inspect
from openff.evaluator.client import RequestOptions
from openff.evaluator.layers import registered_calculation_schemas
request_options = RequestOptions()
# Specify the calculation layers to use.
request_options.calculation_layers = []
if target.allow_reweighting:
request_options.calculation_layers.append("ReweightingLayer")
if target.allow_direct_simulation:
request_options.calculation_layers.append("SimulationLayer")
# Check if a non-default option has been specified.
if target.n_molecules is None and target.n_effective_samples is None:
return request_options
# Generate estimation schemas for each of the properties if a non-default
# option has been specified in the optimization options.
property_types = training_set.property_types
request_options.calculation_schemas = defaultdict(dict)
for property_type in property_types:
default_reweighting_schemas = registered_calculation_schemas.get(
"ReweightingLayer", {})
if (target.allow_reweighting
and target.n_effective_samples is not None
and property_type in default_reweighting_schemas
and callable(default_reweighting_schemas[property_type])):
default_schema = default_reweighting_schemas[property_type]
if "n_effective_samples" in inspect.getfullargspec(
default_schema).args:
default_schema = default_schema(
n_effective_samples=target.n_effective_samples)
request_options.calculation_schemas[property_type][
"ReweightingLayer"] = default_schema
default_simulation_schemas = registered_calculation_schemas.get(
"SimulationLayer", {})
if (target.allow_direct_simulation
and target.n_molecules is not None
and property_type in default_simulation_schemas
and callable(default_simulation_schemas[property_type])):
default_schema = default_simulation_schemas[property_type]
if "n_molecules" in inspect.getfullargspec(
default_schema).args:
default_schema = default_schema(
n_molecules=target.n_molecules)
request_options.calculation_schemas[property_type][
"SimulationLayer"] = default_schema
return request_options
