import os
import pickle
import sys
import time
from logging import info, warning
import numpy as np
import pandas as pd
from ....utils import path
from ....utils.yml import ordered_dump
[docs]
def init_check_files(self):
"""Set the paths of the files monitoring the convergence and back up the
ones from a previous run."""
simulator_dir = self.simulator.simulator_dir
self.check_cost_file = simulator_dir / "congrad_checks_cost.csv"
self.check_gradcost_file = simulator_dir / "congrad_checks_gradcost.csv"
for f in (self.check_cost_file, self.check_gradcost_file):
if f.exists():
f.replace(f.with_suffix(f"{f.suffix}.bak"))
self.pending_checks = []
[docs]
def write_predicted(self, niter, chi, cost, grad):
r"""Write the cost function and gradient predicted by the quadratic model
at the current estimate of the minimum.
Args:
niter (int): number of Lanczos iterations performed so far.
chi (np.ndarray): current estimate of the minimum.
cost (float): predicted cost function.
grad (np.ndarray): predicted gradient.
"""
run_id = _run_id(niter)
# Jb = 1/2 chi'.chi and grad(Jb) = chi are exact
j_b = 0.5 * np.dot(chi, chi)
znorm_grad = np.sqrt(np.dot(grad, grad))
info(
f"CONGRAD predicted values after {niter} iterations:\n"
f" Jb = {j_b}\n"
f" Jo = {cost - j_b}\n"
f" Total = {cost}\n"
f" grad(J) = {znorm_grad}"
)
_write_csv(
self.check_cost_file, run_id,
{"j_b_pred": j_b, "j_o_pred": cost - j_b, "zcost_pred": cost},
)
_write_csv(
self.check_gradcost_file, run_id,
{
"znorm_grad_b_pred": np.sqrt(np.dot(chi, chi)),
"znorm_grad_o_pred": np.sqrt(np.dot(grad - chi, grad - chi)),
"znorm_grad_pred": znorm_grad,
},
)
[docs]
def dump_state(self, chi, niter, **kwargs):
"""Dump the control vector at the current estimate of the minimum.
Args:
chi (np.ndarray): current estimate of the minimum.
niter (int): number of Lanczos iterations performed so far.
**kwargs: forwarded to ``sqrtbprod``.
"""
controlvect = self.simulator.controlvect
check_dir = _check_dir(self, niter)
path.init_dir(check_dir)
# Keep the control vector of the Lanczos iterations untouched
x_lanczos = controlvect.x
controlvect.x = controlvect.sqrtbprod(chi, **kwargs)
controlvect.dump(
f"{check_dir}/controlvect.pickle",
to_netcdf=controlvect.save_out_netcdf,
dir_netcdf=f"{check_dir}/controlvect/",
)
controlvect.x = x_lanczos
[docs]
def submit_check(self, chi, niter):
r"""Submit a job evaluating the cost function and its gradient at the
current estimate of the minimum.
Args:
chi (np.ndarray): current estimate of the minimum.
niter (int): number of Lanczos iterations performed so far.
"""
platform = self.platform
workdir = self.simulator.workdir
run_id = _run_id(niter)
# Separate working directory to keep the dumped state apart
eval_dir = f"{_check_dir(self, niter)}/evaluation/"
path.init_dir(eval_dir)
check = {
"run_id": run_id,
"niter": niter,
"job_id": None,
"result_file":
f"{eval_dir}/simulator/cost_grad.{run_id}.pickle",
}
# Already computed in a previous run
if os.path.isfile(check["result_file"]):
self.pending_checks.append(check)
return
chi_file = f"{eval_dir}/chi.npy"
np.save(chi_file, chi)
# Same configuration as the main run, without the mode and with the
# simulator at the top level to only evaluate the cost function
yml_dict = self.from_yaml(
self.reference_instances["reference_setup"].def_file)
mode_dict = yml_dict.pop("mode")
simulator_dict = mode_dict.get("minimizer", {}).get(
"simulator",
mode_dict.get("simulator", yml_dict.get("simulator", {}))
)
yml_dict["simulator"] = {
**simulator_dict,
"plugin": {"name": self.simulator.plugin.name,
"version": self.simulator.plugin.version},
}
yml_dict["workdir"] = eval_dir
yml_dict["return_config"] = True
yml_dict["obsvect"] = {
**yml_dict.get("obsvect", {}),
"dir_obsvect": f"{workdir}/obsvect/",
}
del yml_dict["def_file"]
yml_file = f"{eval_dir}/config_{run_id}.yml"
with open(yml_file, "w") as outfile:
ordered_dump(outfile, yml_dict)
python = getattr(platform, "python", sys.executable)
exe = (f"{python} -m pycif.plugins.minimizers.congrad.check_job "
f"{yml_file} {chi_file} {run_id}")
if self.force_linearize:
exe += " --linear"
info(f"Submitting CONGRAD check after {niter} iterations in {eval_dir}")
check["job_id"] = platform.submit_job(
exe, os.path.join(eval_dir, f"job_pycif_{run_id}"))
self.pending_checks.append(check)
[docs]
def collect_checks(self, wait=False):
"""Process the evaluated checks whose job is over.
Args:
wait (bool): wait for all pending checks to be over.
"""
platform = self.platform
while self.pending_checks:
remaining = []
for check in self.pending_checks:
if check["job_id"] is not None \
and not platform.check_jobs([check["job_id"]]):
remaining.append(check)
continue
_process_check(self, check)
self.pending_checks = remaining
if not wait or not remaining:
break
time.sleep(platform.sleep_time)
def _run_id(niter):
return f"check_{niter:04d}"
def _check_dir(self, niter):
return f"{self.simulator.workdir}/congrad_checks/{_run_id(niter)}/"
def _process_check(self, check):
run_id = check["run_id"]
# Checks are only diagnostics: a failure does not stop the minimization
if not os.path.isfile(check["result_file"]):
warning(
f"CONGRAD check {run_id} failed. Please check the simulation in "
f"{os.path.dirname(os.path.dirname(check['result_file']))}"
)
return
with open(check["result_file"], "rb") as f:
j_b, j_o, zgrad_b, zgrad_o = pickle.load(f)
zcost = j_b + j_o
zgrad = zgrad_b + zgrad_o
znorm_grad = np.sqrt(np.dot(zgrad, zgrad))
info(
f"CONGRAD evaluated values after {check['niter']} iterations:\n"
f" Jb = {j_b}\n"
f" Jo = {j_o}\n"
f" Total = {zcost}\n"
f" grad(J) = {znorm_grad}"
)
_write_csv(
self.check_cost_file, run_id,
{"j_b": j_b, "j_o": j_o, "zcost": zcost},
)
_write_csv(
self.check_gradcost_file, run_id,
{
"znorm_grad_b": np.sqrt(np.dot(zgrad_b, zgrad_b)),
"znorm_grad_o": np.sqrt(np.dot(zgrad_o, zgrad_o)),
"znorm_grad": znorm_grad,
},
)
def _write_csv(csv_file, run_id, values):
if csv_file.exists():
df = pd.read_csv(csv_file, index_col="run_id")
else:
df = pd.DataFrame(columns=[*values, "time"])
df.index.name = "run_id"
for key, value in values.items():
df.loc[run_id, key] = value
df.loc[run_id, "time"] = pd.Timestamp.now().isoformat()
df.to_csv(csv_file)
