def optimizations_and_results( smirnoff_force_field, ) -> Tuple[List[Optimization], List[OptimizationResult]]: optimization_1 = 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_1.name = "Optimization 1" optimization_1.force_field = ForceField.from_openff(smirnoff_force_field) optimization_2 = create_optimization( "project-1", "study-1", "optimization-2", [ create_evaluator_target("evaluator-target-1", ["data-set-1"]), create_recharge_target("recharge-target-1", ["qc-data-set-1"]), ], ) optimization_2.force_field = ForceField.from_openff(smirnoff_force_field) optimization_2.name = "Optimization 2" smirnoff_force_field.get_parameter_handler("vdW").parameters["[#6:1]"].epsilon *= 2 smirnoff_force_field.get_parameter_handler("vdW").parameters["[#6:1]"].sigma *= 3 optimization_result_1 = create_optimization_result( "project-1", "study-1", "optimization-1", ["evaluator-target-1"], ["recharge-target-1"], ) optimization_result_1.refit_force_field = ForceField.from_openff( smirnoff_force_field ) smirnoff_force_field.get_parameter_handler("vdW").parameters["[#6:1]"].epsilon /= 4 smirnoff_force_field.get_parameter_handler("vdW").parameters["[#6:1]"].sigma /= 6 optimization_result_2 = create_optimization_result( "project-1", "study-1", "optimization-2", ["evaluator-target-1"], ["recharge-target-1"], ) optimization_result_2.refit_force_field = ForceField.from_openff( smirnoff_force_field ) return ( [optimization_1, optimization_2], [optimization_result_1, optimization_result_2], )
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 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_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 create_model(cls, include_children=False, index=1): optimization = create_optimization( "project-1", "study-1", f"optimization-{index}", targets=[ create_evaluator_target("evaluator-target-1", ["data-set-1"]), create_recharge_target("recharge-target-1", ["qc-data-set-1"]), ], ) return optimization
def optimization(force_field) -> Optimization: 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 return optimization
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_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 commit_optimization( db: Session, ) -> Tuple[Project, Study, Optimization, DataSetCollection, QCDataSetCollection]: """Commits a new project and study to the current session and appends an empty optimization onto it. Additionally, this function commits two data sets to the session to use as the training set. Parameters ---------- db The current data base session. """ training_set = commit_data_set_collection(db) training_set_ids = [x.id for x in training_set.data_sets] qc_data_set = commit_qc_data_set_collection(db) qc_data_set_ids = [x.id for x in qc_data_set.data_sets] study = create_study("project-1", "study-1") study.optimizations = [ create_optimization( "project-1", "study-1", "optimization-1", [ create_evaluator_target("evaluator-target-1", training_set_ids), create_recharge_target("recharge-target-1", qc_data_set_ids), ], ) ] project = create_project(study.project_id) project.studies = [study] db_project = ProjectCRUD.create(db, project) db.add(db_project) db.commit() project = ProjectCRUD.db_to_model(db_project) return project, study, study.optimizations[0], training_set, qc_data_set
def test_plot_target_rmse(tmpdir): initial_result = RechargeTargetResult( objective_function=0.5, statistic_entries=[ Statistic( statistic_type=StatisticType.RMSE, value=1.0, lower_95_ci=0.95, upper_95_ci=1.05, category="Alcohol", ) ], ) final_result = RechargeTargetResult( objective_function=0.5, statistic_entries=[ Statistic( statistic_type=StatisticType.RMSE, value=0.5, lower_95_ci=0.4, upper_95_ci=0.6, category="Alcohol", ) ], ) figures = plot_target_rmse( [create_recharge_target("target-1", ["data-set-1"])] * 2, [initial_result, final_result], ["Initial", "Final"], ) assert "esp" in figures figure = figures["esp"] assert len(figure.subplots) == 1 assert len(figure.subplots[0].traces) == 2 assert figure is not None assert figure.to_plotly() is not None
def valid_optimization_kwargs(valid_sub_study_kwargs): force_field = ForceField(inner_content=" ") parameters = [Parameter(handler_type=" ", smirks=" ", attribute_name=" ")] return { **valid_sub_study_kwargs, "max_iterations": 1, "engine": ForceBalance(priors={" ": 1.0}), "targets": [ create_evaluator_target("evaluator-target", ["data-set-1"]), create_recharge_target("recharge-target", ["qc-data-set-1"]), ], "force_field": force_field, "parameters_to_train": parameters, "analysis_environments": [], }
from nonbonded.tests.utilities.factory import ( create_evaluator_target, create_optimization, create_recharge_target, ) @pytest.mark.parametrize( "optimization, server_config, expected_raises", [ ( create_optimization( "project-1", "study-1", "optimization-1", [create_recharge_target("recharge-target", ["qc-data-set-1"])], ), None, does_not_raise(), ), ( create_optimization( "project-1", "study-1", "optimization-1", [create_evaluator_target("evaluator-target", ["data-set-1"])], ), None, pytest.raises(RuntimeError), ), (
def test_analysis(monkeypatch, force_field, dummy_conda_env): 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(os.path.dirname(dummy_conda_env)): # 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")) os.makedirs(os.path.join("optimize.tmp", target.id, "iter_0001")) # Add enough output files to make it look like only one full iteration has # finished. lp_dump( {"X": 1.0}, os.path.join("optimize.tmp", target.id, "iter_0000", "objective.p"), ) lp_dump( {"X": 1.0}, os.path.join("optimize.tmp", optimization.targets[0].id, "iter_0001", "objective.p"), ) with open("optimization.json", "w") as file: file.write(optimization.json()) optimization.force_field.to_openff().to_file( os.path.join("result", "optimize", "force-field.offxml")) # Mock the already tested functions. monkeypatch.setattr(OptimizationAnalysisFactory, "_load_refit_force_field", lambda: force_field) monkeypatch.setattr( EvaluatorAnalysisFactory, "analyze", lambda *args, **kwargs: EvaluatorTargetResult( objective_function=1.0, statistic_entries=[]), ) monkeypatch.setattr( RechargeAnalysisFactory, "analyze", lambda *args, **kwargs: RechargeTargetResult( objective_function=1.0, statistic_entries=[]), ) OptimizationAnalysisFactory.analyze(True) for target in optimization.targets: assert os.path.isfile( os.path.join("analysis", target.id, "iteration-0.json")) assert not os.path.isfile( os.path.join("analysis", target.id, "iteration-1.json")) result = OptimizationResult.parse_file( os.path.join("analysis", "optimization-results.json")) assert len(result.target_results) == 1 assert all(target.id in result.target_results[0] for target in optimization.targets) assert result.refit_force_field.inner_content == force_field.inner_content
class TestOptimizationInputFactory: 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_generate_force_balance_input(self, optimization): with temporary_cd(): OptimizationInputFactory._generate_force_balance_input( optimization) assert os.path.isfile("optimize.in") @pytest.mark.parametrize("allow_reweighting", [False, True]) def test_generate_request_options_default(self, allow_reweighting): training_set = create_data_set("data-set-1", 1) target = create_evaluator_target("evaluator-target-1", ["data-set-1"]) target.allow_direct_simulation = True target.allow_reweighting = allow_reweighting request_options = OptimizationInputFactory._generate_request_options( target, training_set.to_evaluator()) if allow_reweighting: assert request_options.calculation_layers == [ "ReweightingLayer", "SimulationLayer", ] else: assert request_options.calculation_layers == ["SimulationLayer"] assert request_options.calculation_schemas == UNDEFINED def test_generate_request_options(self): training_set = create_data_set("data-set-1", 1) target = create_evaluator_target("evaluator-target-1", [training_set.id]) target.allow_direct_simulation = True target.allow_reweighting = True target.n_molecules = 512 target.n_effective_samples = 10 request_options = OptimizationInputFactory._generate_request_options( target, training_set.to_evaluator()) assert request_options.calculation_layers == [ "ReweightingLayer", "SimulationLayer", ] assert request_options.calculation_schemas != UNDEFINED expected_simulation_schema = Density.default_simulation_schema( n_molecules=512) expected_reweighting_schema = Density.default_reweighting_schema( n_effective_samples=10) assert ( request_options.calculation_schemas["Density"] ["SimulationLayer"].json() == expected_simulation_schema.json()) assert ( request_options.calculation_schemas["Density"] ["ReweightingLayer"].json() == expected_reweighting_schema.json()) 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_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() @pytest.mark.parametrize( "target", [ create_evaluator_target("evaluator-target-1", ["data-set-1"]), create_recharge_target("recharge-target-1", ["qc-data-set-1"]), ], ) 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 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(self, optimization, monkeypatch): logging.basicConfig(level=logging.INFO) # Mock the already tested functions monkeypatch.setattr(OptimizationInputFactory, "_prepare_force_field", lambda *args: None) monkeypatch.setattr( OptimizationInputFactory, "_generate_force_balance_input", lambda *args: None, ) monkeypatch.setattr(OptimizationInputFactory, "_generate_target", lambda *args: None) monkeypatch.setattr(OptimizationInputFactory, "_generate_submission_script", lambda *args: None) monkeypatch.setattr(OptimizationInputFactory, "_retrieve_results", lambda *args: None) with temporary_cd(): OptimizationInputFactory.generate(optimization, "env", "01:00", "lilac-local", 8000, 1, True)
def create_dependencies(db: Session, dependencies: List[str]): """Create any dependencies such as parent studies, projects, or data sets and commit them to the database. Parameters ---------- db The current database session. dependencies The required dependencies. """ project = None data_set_ids = [] qc_data_set_ids = [] if "data-set" in dependencies: data_set_ids.append("data-set-1") if "qc-data-set" in dependencies: qc_data_set_ids.append("qc-data-set-1") for data_set_id in data_set_ids: data_set = create_data_set(data_set_id) db_data_set = DataSetCRUD.create(db, data_set) db.add(db_data_set) for qc_data_set_id in qc_data_set_ids: qc_data_set = create_qc_data_set(qc_data_set_id) db_qc_data_set = QCDataSetCRUD.create(db, qc_data_set) db.add(db_qc_data_set) db.commit() if ("project" in dependencies or "study" in dependencies or "evaluator-target" in dependencies or "recharge-target" in dependencies or "benchmark" in dependencies): project = create_project("project-1") if ("study" in dependencies or "evaluator-target" in dependencies or "recharge-target" in dependencies or "benchmark" in dependencies): project.studies = [create_study(project.id, "study-1")] if "evaluator-target" in dependencies or "recharge-target" in dependencies: targets = [] if "evaluator-target" in dependencies: targets.append( create_evaluator_target("evaluator-target-1", ["data-set-1"])) if "recharge-target" in dependencies: targets.append( create_recharge_target("recharge-target-1", ["qc-data-set-1"])) optimization = create_optimization(project.id, project.studies[0].id, "optimization-1", targets) project.studies[0].optimizations = [optimization] if "benchmark" in dependencies: benchmark = create_benchmark( project.id, project.studies[0].id, "benchmark-1", ["data-set-1"], None, create_force_field(), ) project.studies[0].benchmarks = [benchmark] if project is not None: db_project = ProjectCRUD.create(db, project) db.add(db_project) db.commit()
def optimization_model_perturbations(): updated_engine_delete_prior = ForceBalance( priors={"vdW/Atom/epsilon": 0.1}) updated_engine_update_prior = ForceBalance(priors={ "vdW/Atom/epsilon": 0.2, "vdW/Atom/sigma": 2.0 }, ) updated_engine_add_prior = ForceBalance(priors={ "vdW/Atom/epsilon": 0.1, "vdW/Atom/sigma": 2.0, "vdW/Atom/r_min": 2.0 }, ) invalid_evaluator_target = create_evaluator_target("evaluator-target-1", ["data-set-999"]) invalid_recharge_target = create_recharge_target("recharge-target-1", ["qc-data-set-999"]) return [ ({ "name": "updated" }, lambda db: [], does_not_raise()), ({ "description": "updated" }, lambda db: [], does_not_raise()), ({ "max_iterations": 999 }, lambda db: [], does_not_raise()), ( { "analysis_environments": [ChemicalEnvironment.Hydroxy] }, lambda db: [], does_not_raise(), ), # Test updating the force field. ( { "force_field": create_force_field("updated") }, lambda db: [ "updated" in db.query(models.ForceField.inner_content).first()[ 0], db.query(models.ForceField.id).count() == 1, ], does_not_raise(), ), # Test updating the parameters to train. ( { "parameters_to_train": [ Parameter(handler_type="vdW", smirks="[#6:1]", attribute_name="epsilon"), ] }, lambda db: [db.query(models.Parameter.id).count() == 1], does_not_raise(), ), ( { "parameters_to_train": [ Parameter(handler_type="vdW", smirks="[#6:1]", attribute_name="epsilon"), Parameter(handler_type="vdW", smirks="[#6:1]", attribute_name="sigma"), Parameter(handler_type="vdW", smirks="[#1:1]", attribute_name="sigma"), ] }, lambda db: [db.query(models.Parameter.id).count() == 3], does_not_raise(), ), # Test updating an engine's priors ( { "engine": updated_engine_delete_prior }, lambda db: [db.query(models.ForceBalancePrior.id).count() == 1], does_not_raise(), ), ( { "engine": updated_engine_update_prior }, lambda db: [db.query(models.ForceBalancePrior.id).count() == 2], does_not_raise(), ), ( { "engine": updated_engine_add_prior }, lambda db: [db.query(models.ForceBalancePrior.id).count() == 3], does_not_raise(), ), # Test deleting a target ( { "targets": [ create_evaluator_target("evaluator-target-1", ["data-set-1"]) ] }, lambda db: [ db.query(models.EvaluatorTarget.id).count() == 1, db.query(models.RechargeTarget.id).count() == 0, ], does_not_raise(), ), ( { "targets": [ create_recharge_target("recharge-target-1", ["qc-data-set-1"]) ] }, lambda db: [ db.query(models.EvaluatorTarget.id).count() == 0, db.query(models.RechargeTarget.id).count() == 1, ], does_not_raise(), ), # Test adding a target ( { "targets": [ create_evaluator_target("evaluator-target-1", ["data-set-1"]), create_evaluator_target("evaluator-target-2", ["data-set-1"]), create_recharge_target("recharge-target-1", ["qc-data-set-1"]), ] }, lambda db: [ db.query(models.EvaluatorTarget.id).count() == 2, db.query(models.RechargeTarget.id).count() == 1, ], does_not_raise(), ), ( { "targets": [ create_evaluator_target("evaluator-target-1", ["data-set-1"]), create_recharge_target("recharge-target-1", ["qc-data-set-1"]), create_recharge_target("recharge-target-2", ["qc-data-set-1"]), ] }, lambda db: [ db.query(models.EvaluatorTarget.id).count() == 1, db.query(models.RechargeTarget.id).count() == 2, ], does_not_raise(), ), # Test invalidly updating a target's training set ( { "targets": [invalid_evaluator_target] }, lambda db: [], pytest.raises(DataSetNotFoundError), ), ( { "targets": [invalid_recharge_target] }, lambda db: [], pytest.raises(QCDataSetNotFoundError), ), ]
def test_prepare_restart_unfinished(partial_restart, 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-2", "iter_0000"), os.path.join("optimize.tmp", "recharge-target-1", "iter_0001"), os.path.join("optimize.tmp", "recharge-target-2", "iter_0001"), ] for index, directory in enumerate(directories): os.makedirs(directory) expected_files = ["mvals.txt"] if index < 3: expected_files.append("objective.p") if index < (3 if not partial_restart else 4): expected_files.append("force-field.offxml") for file_name in expected_files: 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) expected_directories = 2 if partial_restart else 1 assert (len( glob(os.path.join("optimize.tmp", "recharge-target-1", "*"))) == expected_directories) assert (len( glob(os.path.join("optimize.tmp", "recharge-target-2", "*"))) == expected_directories) if not partial_restart: assert ( f"Removing the {directories[2]} directory which was not expected to be " f"present") in caplog.text assert ( f"Removing the {directories[3]} directory which was not expected to be " f"present") in caplog.text else: assert "Removing the" not in caplog.text assert ( "1 iterations had previously been completed. The optimization will be " f"restarted from iteration {'0000' if not partial_restart else '0001'}" ) in caplog.text