def test_temporary_cd():
"""Tests that temporary cd works as expected"""
original_directory = os.getcwd()
# Move to the parent directory
with temporary_cd(os.pardir):
current_directory = os.getcwd()
expected_directory = os.path.abspath(
os.path.join(original_directory, os.pardir))
assert compare_paths(current_directory, expected_directory)
assert compare_paths(os.getcwd(), original_directory)
# Move to a temporary directory
with temporary_cd():
assert not compare_paths(os.getcwd(), original_directory)
assert compare_paths(os.getcwd(), original_directory)
# Move to the same directory
with temporary_cd(""):
assert compare_paths(os.getcwd(), original_directory)
assert compare_paths(os.getcwd(), original_directory)
with temporary_cd(os.curdir):
assert compare_paths(os.getcwd(), original_directory)
assert compare_paths(os.getcwd(), original_directory)
def test_prepare_restart(
original_existing_results: Optional[RequestResult],
restart: bool,
expected_message: Optional[str],
caplog,
):
with temporary_cd():
# Create a mock data set
original_data_set = DataSetCollection(
data_sets=[create_data_set("data-set-1", 1)])
original_data_set.to_file("test-set-collection.json")
if original_existing_results is not None:
original_existing_results.json("results.json")
with caplog.at_level(logging.INFO):
data_set, existing_result = _prepare_restart(restart)
if original_existing_results is not None and restart:
assert existing_result.json() == original_existing_results.json()
else:
assert existing_result is None
if expected_message is not None:
assert expected_message in caplog.text
else:
assert caplog.text == ""
expected_n_data_points = (1 if existing_result is None else len(
existing_result.unsuccessful_properties))
assert len(data_set) == expected_n_data_points
def mock_target(tmpdir) -> Tuple[Optimization, EvaluatorTarget, str]:
"""Create a mock evaluator target directory which is populated with a dummy
set of results.
Returns
-------
A tuple of the parent optimization, the mock target and the path to the
directory in which the files were created.
"""
with temporary_cd(str(tmpdir)):
# Mock the target to analyze.
target = create_evaluator_target("evaluator-target-1", ["data-set-1"])
optimization = create_optimization("project-1", "study-1",
"optimization-1", [target])
optimization.analysis_environments = []
# Create a dummy data set and estimated result.
reference_data_set = create_data_set("data-set-1", 1)
DataSetCollection(data_sets=[reference_data_set]).to_evaluator().json(
"training-set.json")
results = RequestResult()
results.estimated_properties = reference_data_set.to_evaluator()
results.json("results.json")
lp_dump({"X": 1.0}, "objective.p")
return optimization, target, str(tmpdir)
def mock_target(tmpdir) -> Tuple[Optimization, RechargeTarget, str]:
"""Create a mock recharge target directory which is populated with a dummy
set of results.
Returns
-------
A tuple of the parent optimization, the mock target and the path to the
directory in which the files were created.
"""
with temporary_cd(str(tmpdir)):
# Mock the target to analyze.
target = create_recharge_target("recharge-target-1", ["qc-data-set-1"])
optimization = create_optimization("project-1", "study-1",
"optimization-1", [target])
optimization.analysis_environments = [
ChemicalEnvironment.Alkane,
ChemicalEnvironment.Alcohol,
]
# Create a dummy set of residuals.
with open("residuals.json", "w") as file:
json.dump({"C": 9.0, "CO": 4.0}, file)
lp_dump({"X": 1.0}, "objective.p")
return optimization, target, str(tmpdir)
def test_analysis_n_iteration(monkeypatch, force_field):
"""Test that the correction exception is raised in the case where a refit
force field is found but no target outputs are."""
optimization = create_optimization(
"project-1",
"study-1",
"optimization-1",
[
create_evaluator_target("evaluator-target-1", ["data-set-1"]),
create_recharge_target("recharge-target-1", ["qc-data-set-1"]),
],
)
optimization.force_field = force_field
with temporary_cd():
# Save the mock optimization file.
with open("optimization.json", "w") as file:
file.write(optimization.json())
# Mock successfully reading a refit force field.
monkeypatch.setattr(OptimizationAnalysisFactory,
"_load_refit_force_field", lambda: force_field)
with pytest.raises(RuntimeError) as error_info:
OptimizationAnalysisFactory.analyze(True)
assert "No iteration results could be found" in str(error_info.value)
def test_retrieve_force_field(self, requests_mock):
# Mock a refit force field to retrieve.
result = create_optimization_result(
"project-1",
"study-1",
"optimization-1",
["evaluator-target-1"],
[],
)
result.refit_force_field = ForceField(
inner_content=('<SMIRNOFF version="0.3" '
'aromaticity_model="OEAroModel_MDL"></SMIRNOFF>'))
mock_get_optimization_result(requests_mock, result)
# Mock a benchmark which targets the refit force field.
benchmark = create_benchmark(
"project-1",
"study-1",
"benchmark-1",
data_set_ids=["data-set-1"],
optimization_id="optimization-1",
force_field=None,
)
with temporary_cd():
BenchmarkInputFactory._retrieve_force_field(benchmark, None)
assert os.path.isfile("force-field.offxml")
def test_run_command(runner, monkeypatch):
from nonbonded.cli.projects.benchmark import run
monkeypatch.setattr(run, "_prepare_restart", lambda *args:
(None, successful_result()))
monkeypatch.setattr(run, "_load_force_field", lambda *args: None)
monkeypatch.setattr(run, "_run_calculations",
lambda *args: RequestResult())
# Save a copy of the result model.
with temporary_cd():
# Create mock inputs
with open("server-config.json", "w") as file:
file.write(
EvaluatorServerConfig(backend_config=DaskLocalClusterConfig(
resources_per_worker=ComputeResources())).json())
RequestOptions().json("estimation-options.json")
result = runner.invoke(run_command())
with open("results.json") as file:
assert successful_result().json() == file.read()
if result.exit_code != 0:
raise result.exception
def test_study_no_children(requests_mock):
study = create_study("project-1", "study-1")
mock_get_study(requests_mock, study)
with temporary_cd():
InputFactory.generate(study, "test-env", "12:34", "lilac-dask", 8000, 1, False)
def test_project_no_children(requests_mock):
project = create_project(project_id="project-1")
mock_get_project(requests_mock, project)
with temporary_cd():
InputFactory.generate(
project, "test-env", "12:34", "lilac-dask", 8000, 1, False
)
def test_analysis_missing_result(monkeypatch, force_field):
"""Test that the correction exception is raised in the case where a the
expected results of a target are missing."""
optimization = create_optimization(
"project-1",
"study-1",
"optimization-1",
[
create_evaluator_target("evaluator-target-1", ["data-set-1"]),
create_recharge_target("recharge-target-1", ["qc-data-set-1"]),
],
)
optimization.force_field = force_field
with temporary_cd():
# Save the expected results files.
os.makedirs(os.path.join("result", "optimize"))
for target in optimization.targets:
os.makedirs(os.path.join("targets", target.id))
os.makedirs(os.path.join("optimize.tmp", target.id, "iter_0000"))
lp_dump(
{"X": 1.0},
os.path.join("optimize.tmp", target.id, "iter_0000",
"objective.p"),
)
with open("optimization.json", "w") as file:
file.write(optimization.json())
monkeypatch.setattr(OptimizationAnalysisFactory,
"_load_refit_force_field", lambda: force_field)
# Mock a missing target result.
monkeypatch.setattr(
EvaluatorAnalysisFactory,
"analyze",
lambda *args, **kwargs: EvaluatorTargetResult(
objective_function=1.0, statistic_entries=[]),
)
monkeypatch.setattr(
RechargeAnalysisFactory,
"analyze",
lambda *args, **kwargs: None,
)
with pytest.raises(RuntimeError) as error_info:
OptimizationAnalysisFactory.analyze(True)
assert "The results of the recharge-target-1 target could not be found" in str(
error_info.value)
def test_benchmark_analysis(caplog, monkeypatch, dummy_conda_env):
from openff.evaluator.client import RequestResult
from openff.evaluator.datasets import PhysicalPropertyDataSet
benchmark = create_benchmark(
"project-1", "study-1", "benchmark-1", ["data-set-1"], "optimization-1", None
)
# Create a reference data set.
reference_data_set = create_data_set("data-set-1")
reference_data_set.entries.append(reference_data_set.entries[0].copy())
reference_data_set.entries[0].id = 1
reference_data_set.entries[1].id = 2
# Create a set of evaluator results
estimated_data_set = PhysicalPropertyDataSet()
estimated_data_set.add_properties(reference_data_set.entries[0].to_evaluator())
unsuccessful_properties = PhysicalPropertyDataSet()
unsuccessful_properties.add_properties(reference_data_set.entries[1].to_evaluator())
results = RequestResult()
results.estimated_properties = estimated_data_set
results.unsuccessful_properties = unsuccessful_properties
with temporary_cd(os.path.dirname(dummy_conda_env)):
# Save the expected input files.
with open("benchmark.json", "w") as file:
file.write(benchmark.json())
with open("test-set-collection.json", "w") as file:
file.write(DataSetCollection(data_sets=[reference_data_set]).json())
results.json("results.json")
with caplog.at_level(logging.WARNING):
BenchmarkAnalysisFactory.analyze(True)
assert (
"1 properties could not be estimated and so were not analyzed"
in caplog.text
)
assert os.path.isdir("analysis")
assert os.path.isfile(os.path.join("analysis", "benchmark-results.json"))
results_object = BenchmarkResult.parse_file(
os.path.join("analysis", "benchmark-results.json")
)
assert len(results_object.calculation_environment) > 0
assert len(results_object.analysis_environment) > 0
def test_server_config_to_evaluator():
server_config = EvaluatorServerConfig(
backend_config=DaskLocalClusterConfig(
resources_per_worker=ComputeResources(n_gpus=0)))
backend = server_config.to_backend()
backend._started = True
with temporary_cd():
assert isinstance(backend, dask.DaskLocalCluster)
assert isinstance(server_config.to_server(backend),
server.EvaluatorServer)
def test_load_refit_force_field(tmpdir, smirnoff_force_field):
with temporary_cd(str(tmpdir)):
os.makedirs(os.path.join("result", "optimize"))
smirnoff_force_field.to_file(
os.path.join(
tmpdir, os.path.join("result", "optimize",
"force-field.offxml")))
refit_force_field = OptimizationAnalysisFactory._load_refit_force_field(
)
assert refit_force_field.inner_content == smirnoff_force_field.to_string(
)
def test_prepare_restart_finished(caplog):
optimization = create_optimization(
"project-1",
"study-1",
"optimization-1",
[
create_recharge_target("recharge-target-1", ["qc-data-set-1"]),
create_recharge_target("recharge-target-2", ["qc-data-set-1"]),
],
)
with temporary_cd():
directories = [
os.path.join("optimize.tmp", "recharge-target-1", "iter_0000"),
os.path.join("optimize.tmp", "recharge-target-1", "iter_0001"),
os.path.join("optimize.tmp", "recharge-target-2", "iter_0000"),
os.path.join("optimize.tmp", "recharge-target-2", "iter_0001"),
]
for directory in directories:
os.makedirs(directory)
for file_name in [
"mvals.txt", "force-field.offxml", "objective.p"
]:
with open(os.path.join(directory, file_name), "w") as file:
file.write("")
assert len(glob(os.path.join("optimize.tmp", "recharge-target-1",
"*"))) == 2
assert len(glob(os.path.join("optimize.tmp", "recharge-target-2",
"*"))) == 2
with caplog.at_level(logging.INFO):
_prepare_restart(optimization)
assert len(glob(os.path.join("optimize.tmp", "recharge-target-1",
"*"))) == 2
assert len(glob(os.path.join("optimize.tmp", "recharge-target-2",
"*"))) == 2
assert (
"2 iterations had previously been completed. The optimization will be "
"restarted from iteration 0002") in caplog.text
def test_generate(self, benchmark, monkeypatch):
# Mock the already tested functions
monkeypatch.setattr(BenchmarkInputFactory, "_retrieve_force_field",
lambda *args: None)
monkeypatch.setattr(BenchmarkInputFactory, "_retrieve_data_sets",
lambda *args: None)
monkeypatch.setattr(BenchmarkInputFactory,
"_generate_submission_script", lambda *args: None)
monkeypatch.setattr(BenchmarkInputFactory, "_retrieve_results",
lambda *args: None)
with temporary_cd():
BenchmarkInputFactory.generate(benchmark, "env", "01:00",
"lilac-local", 8000, 1, True)
def test_retrieve_data_sets(self, benchmark, requests_mock):
# Mock the data set to retrieve.
data_set = create_data_set("data-set-1", 1)
mock_get_data_set(requests_mock, data_set)
with temporary_cd():
BenchmarkInputFactory._retrieve_data_sets(benchmark, None)
assert os.path.isfile("test-set-collection.json")
from nonbonded.library.models.datasets import DataSetCollection
data_set_collection = DataSetCollection.parse_file(
"test-set-collection.json")
assert data_set_collection.data_sets[0].json() == data_set.json()
def test_project_with_children(requests_mock):
project = create_project(project_id="project-1")
project.studies = [create_study(project.id, "study-1")]
mock_get_project(requests_mock, project)
mock_get_study(requests_mock, project.studies[0])
with temporary_cd():
InputFactory.generate(
project, "test-env", "12:34", "lilac-dask", 8000, 1, False
)
assert os.path.isdir(project.id)
assert os.path.isdir(os.path.join(project.id, "studies"))
def test_retrieve_results(self, benchmark, requests_mock):
result = create_benchmark_result(
benchmark.project_id,
benchmark.study_id,
benchmark.id,
create_data_set("data-set-1", 1),
)
mock_get_benchmark_result(requests_mock, result)
with temporary_cd():
BenchmarkInputFactory._retrieve_results(benchmark)
stored_result = BenchmarkResult.parse_file(
os.path.join("analysis", "benchmark-results.json"))
assert stored_result.json() == result.json()
def test_generate_validate_result(model, optimization_result, expected_raises):
with temporary_cd():
with expected_raises:
InputFactory._generate(
model,
"mock-env",
"01:00",
"lilac-dask",
8000,
1,
False,
None,
optimization_result,
)
def test_run_command(restart: bool, create_save: bool, runner, monkeypatch):
from nonbonded.cli.projects.optimization import run
monkeypatch.setattr(run, "_remove_previous_files", lambda: print("REMOVE"))
monkeypatch.setattr(run, "_prepare_restart",
lambda *args: print("PREPARE"))
monkeypatch.setattr(subprocess, "check_call", lambda *args, **kwargs: None)
optimization = create_optimization(
"project-1",
"study-1",
"optimization-1",
[create_recharge_target("recharge-target-1", ["qc-data-set-1"])],
)
# Save a copy of the result model.
with temporary_cd():
with open("optimization.json", "w") as file:
file.write(optimization.json())
if create_save:
with open("optimize.sav", "w") as file:
file.write("")
arguments = [] if not restart else ["--restart", True]
result = runner.invoke(run_command(), arguments)
if restart and create_save:
assert "REMOVE" not in result.output
assert "PREPARE" in result.output
elif restart and not create_save:
assert "REMOVE" in result.output
assert "PREPARE" not in result.output
if not restart:
assert "REMOVE" in result.output
assert "PREPARE" not in result.output
if result.exit_code != 0:
raise result.exception
def test_generate_evaluator_target(self, requests_mock):
data_set = create_data_set("data-set-1")
mock_get_data_set(requests_mock, data_set)
target = create_evaluator_target("evaluator-target-1", [data_set.id])
with temporary_cd():
OptimizationInputFactory._generate_evaluator_target(
target, 8000, None)
assert os.path.isfile("training-set.json")
off_data_set = PhysicalPropertyDataSet.from_json(
"training-set.json")
assert off_data_set.json() == data_set.to_evaluator().json()
assert os.path.isfile("options.json")
def test_retrieve_results(self, optimization, requests_mock):
result = create_optimization_result(
optimization.project_id,
optimization.study_id,
optimization.id,
["evaluator-target-1"],
[],
)
mock_get_optimization_result(requests_mock, result)
with temporary_cd():
OptimizationInputFactory._retrieve_results(optimization)
stored_result = OptimizationResult.parse_file(
os.path.join("analysis", "optimization-results.json"))
assert stored_result.json() == result.json()
def test_generate_recharge_target(self, requests_mock):
qc_data_set = create_qc_data_set("qc-data-set-1")
mock_get_qc_data_set(requests_mock, qc_data_set)
target = create_recharge_target("recharge-target-1", [qc_data_set.id])
with temporary_cd():
OptimizationInputFactory._generate_recharge_target(target, None)
with open("training-set.json") as file:
training_entries = json.load(file)
assert training_entries == qc_data_set.entries
with open("grid-settings.json") as file:
assert file.read() == target.grid_settings.json()
def test_load_force_field(force_fields, expected_raises,
expected_error_message):
with temporary_cd():
# Create any required force field files.
for force_field_path, force_field in force_fields:
if isinstance(force_field, ForceField):
force_field.to_file(force_field_path)
elif isinstance(force_field, ForceFieldSource):
force_field.json(force_field_path)
# Make sure the right error (if any is raised) on loading
with expected_raises as error_info:
_load_force_field()
if expected_error_message is not None:
assert expected_error_message in str(error_info.value)
def test_plot(force_field, monkeypatch):
from nonbonded.library.plotting.seaborn import optimization as optimization_module
# Mock the required file inputs
data_set = create_data_set("data-set-1", 1)
data_set_collection = DataSetCollection(data_sets=[data_set])
optimization = create_optimization(
"project-1",
"study-1",
"optimization-1",
[create_evaluator_target("target-1", ["data-set-1"])],
)
optimization_result = create_optimization_result("project-1", "study-1",
"optimization-1",
["target-1"], [])
# Mock the already tested plotting methods.
monkeypatch.setattr(optimization_module, "plot_parameter_changes",
lambda *args: None)
monkeypatch.setattr(optimization_module, "plot_objective_per_iteration",
lambda *args: None)
monkeypatch.setattr(optimization_module, "plot_rmse_change",
lambda *args: None)
if "nonbonded.library.factories.plots.optimization" in sys.modules:
sys.modules.pop("nonbonded.library.factories.plots.optimization")
from nonbonded.library.factories.plots.optimization import OptimizationPlotFactory
with temporary_cd():
# Save the inputs in their expected locations.
data_set_collection.to_file("test-set-collection.json")
optimization.to_file("optimization.json")
os.makedirs("analysis")
optimization_result.to_file(
os.path.join("analysis", "optimization-results.json"))
OptimizationPlotFactory.plot([""], "png")
assert os.path.isdir("plots")
def test_study_with_children(requests_mock, monkeypatch):
# Overwrite the child factories so we don't need to provide
# sensible children and wait for them to be buit.
def mock_generate(model, **_):
os.makedirs(model.id, exist_ok=True)
monkeypatch.setattr(OptimizationInputFactory, "generate", mock_generate)
monkeypatch.setattr(BenchmarkInputFactory, "generate", mock_generate)
mock_get_data_set(requests_mock, create_data_set("data-set-1"))
optimization = create_optimization(
"project-1",
"study-1",
"optimization-1",
targets=[create_evaluator_target("evaluator-target", ["data-set-1"])],
)
benchmark = create_benchmark(
"project-1",
"study-1",
"benchmark-1",
["data-set-1"],
None,
create_force_field(),
)
study = create_study("project-1", "study-1")
study.optimizations = [optimization]
study.benchmarks = [benchmark]
mock_get_study(requests_mock, study)
with temporary_cd():
InputFactory.generate(study, "test-env", "12:34", "lilac-dask", 8000, 1, False)
assert os.path.isdir(study.id)
assert os.path.isdir(os.path.join(study.id, "optimizations"))
assert os.path.isdir(os.path.join(study.id, "benchmarks"))
assert os.path.isdir(os.path.join(study.id, "optimizations", optimization.id))
assert os.path.isdir(os.path.join(study.id, "benchmarks", benchmark.id))
def test_remove_previous_files():
with temporary_cd():
with open("optimize.sav", "w") as file:
file.write("")
for directory_name in [
"optimize.tmp",
"optimize.bak",
"result",
"working-data",
]:
os.makedirs(directory_name)
assert len(glob("*")) == 5
_remove_previous_files()
assert len(glob("*")) == 0
def test_prepare_force_field(self, optimization):
"""Test that the correct cosmetic attributes are attached to the FF, especially
in the special case of BCC handlers."""
optimization.parameters_to_train.append(
Parameter(
handler_type="ChargeIncrementModel",
smirks="[#6:1]-[#6:2]",
attribute_name="charge_increment1",
))
optimization.parameters_to_train.append(
Parameter(
handler_type="vdW",
smirks=None,
attribute_name="scale14",
))
with temporary_cd():
OptimizationInputFactory._prepare_force_field(optimization, None)
assert os.path.isfile(
os.path.join("forcefield", "force-field.offxml"))
off_force_field = OFFForceField(
os.path.join("forcefield", "force-field.offxml"),
allow_cosmetic_attributes=True,
)
vdw_handler = off_force_field.get_parameter_handler("vdW")
assert vdw_handler._parameterize == "scale14"
assert len(vdw_handler.parameters) == 1
parameter = vdw_handler.parameters["[#6:1]"]
assert parameter._parameterize == "epsilon, sigma"
bcc_handler = off_force_field.get_parameter_handler(
"ChargeIncrementModel")
assert len(bcc_handler.parameters) == 1
parameter = bcc_handler.parameters["[#6:1]-[#6:2]"]
assert len(parameter.charge_increment) == 1
assert parameter._parameterize == "charge_increment1"
def test_retrieve_refit_force_field(self, force_field):
benchmark = create_benchmark(
"project-1",
"study-1",
"benchmark-1",
data_set_ids=["data-set-1"],
optimization_id=None,
force_field=force_field,
)
is_smirnoff = "SMIRNOFF" in force_field.inner_content
with temporary_cd():
BenchmarkInputFactory._retrieve_force_field(benchmark, None)
if is_smirnoff:
assert os.path.isfile("force-field.offxml")
else:
assert os.path.isfile("force-field.json")
def test_generate_target(self, target, caplog, monkeypatch):
monkeypatch.setattr(
OptimizationInputFactory,
"_generate_evaluator_target",
lambda *args: logging.info("EvaluatorTarget"),
)
monkeypatch.setattr(
OptimizationInputFactory,
"_generate_recharge_target",
lambda *args: logging.info("RechargeTarget"),
)
with caplog.at_level(logging.INFO):
with temporary_cd():
OptimizationInputFactory._generate_target(target, 8000, None)
assert os.path.isdir(os.path.join("targets", target.id))
assert target.__class__.__name__ in caplog.text
