from __future__ import annotations
from logging import warning
import pandas as pd
try:
import cPickle as pickle
except ImportError:
import pickle
[docs]
def write_cost(self, run_id, j_b, j_o, zcost):
"""Append one cost-function record to ``cost.txt`` and ``cost.csv``.
Args:
self (Plugin): simulator plugin instance (provides ``cost_file`` and
``cost_file_csv`` path attributes).
run_id (int): minimiser iteration identifier.
j_b (float): background cost term :math:`J_b`.
j_o (float): observation cost term :math:`J_o`.
zcost (float): total cost :math:`J = J_b + J_o`.
"""
now = pd.Timestamp.now().isoformat()
# Legacy txt file
with open(self.cost_file, "a") as f:
f.write(f"{run_id},{j_b},{j_o},{zcost},{now}\n")
# csv file
if self.cost_file_csv.exists():
df = pd.read_csv(self.cost_file_csv, index_col="run_id")
else:
df = pd.DataFrame(columns=["j_b", "j_o", "zcost", "time"])
df.index.name = "run_id"
df.loc[run_id, "j_b"] = j_b
df.loc[run_id, "j_o"] = j_o
df.loc[run_id, "zcost"] = zcost
df.loc[run_id, "time"] = now
df.to_csv(self.cost_file_csv)
[docs]
def write_gradcost(self, run_id, znorm_grad_b, znorm_grad_o, znorm_grad):
"""Append one gradient-norm record to ``gradcost.txt`` and ``gradcost.csv``.
Args:
self (Plugin): simulator plugin instance (provides ``gradcost_file``
and ``gradcost_file_csv`` path attributes).
run_id (int): minimiser iteration identifier.
znorm_grad_b (float): norm of the background gradient :math:`\|\nabla J_b\|`.
znorm_grad_o (float): norm of the observation gradient :math:`\|\nabla J_o\|`.
znorm_grad (float): norm of the total gradient :math:`\|\nabla J\|`.
"""
now = pd.Timestamp.now().isoformat()
# Legacy txt file
with open(self.gradcost_file, "a") as f:
f.write(f"{run_id},{znorm_grad_b},{znorm_grad_o},{znorm_grad},{now}\n")
# csv file
if self.gradcost_file_csv.exists():
df = pd.read_csv(self.gradcost_file_csv, index_col="run_id")
else:
df = pd.DataFrame(
columns=["znorm_grad_b", "znorm_grad_o", "znorm_grad", "time"]
)
df.index.name = "run_id"
df.loc[run_id, "znorm_grad_b"] = znorm_grad_b
df.loc[run_id, "znorm_grad_o"] = znorm_grad_o
df.loc[run_id, "znorm_grad"] = znorm_grad
df.loc[run_id, "time"] = now
df.to_csv(self.gradcost_file_csv)
[docs]
def dump_cost(
self,
run_id: int,
j_b: float,
j_o: float,
zgrad_b: float,
zgrad_o: float,
) -> None:
"""Pickle the cost-function components and gradient vectors for restart.
Saves ``(j_b, j_o, zgrad_b, zgrad_o)`` to
``$workdir/simulator/cost_grad.NNNN.pickle`` so that a restarted
inversion can skip recomputing the forward and adjoint runs.
Args:
self (Plugin): simulator plugin instance (provides
``simulator_dir``).
run_id (int): minimiser iteration identifier; used to name the file.
j_b (float): background cost :math:`J_b`.
j_o (float): observation cost :math:`J_o`.
zgrad_b (np.ndarray): background gradient component.
zgrad_o (np.ndarray): observation gradient component.
"""
cost_file = self.simulator_dir / f"cost_grad.{run_id:04d}.pickle"
with open(cost_file, "wb") as f:
pickle.dump((j_b, j_o, zgrad_b, zgrad_o), f)
[docs]
def load_cost(
self,
run_id: int,
) -> tuple[float, float, float, float] | tuple[None, None, None, None]:
"""Reload the cost-function components and gradient vectors from a pickle file.
Looks for ``cost_grad.NNNN.pickle`` (zero-padded) first, then falls
back to the legacy non-padded name. Returns four ``None`` values when
the file does not exist or cannot be read.
Args:
self (Plugin): simulator plugin instance (provides
``simulator_dir``).
run_id (int): minimiser iteration identifier.
Returns:
tuple: ``(j_b, j_o, zgrad_b, zgrad_o)`` if the file exists and is
readable; ``(None, None, None, None)`` otherwise.
"""
cost_file = self.simulator_dir / f"cost_grad.{run_id:04d}.pickle"
legacy_cost_file = self.simulator_dir / f"cost_grad.{run_id}.pickle"
path = cost_file if cost_file.exists() else legacy_cost_file
if path.exists():
with open(path, "rb") as f:
try:
j_b, j_o, zgrad_b, zgrad_o = pickle.load(f)
except EOFError:
warning(f"Could not load cost function from '{path}'")
return None, None, None, None
return j_b, j_o, zgrad_b, zgrad_o
else:
return None, None, None, None
