import firedrake as fd
from firedrake.adjoint import *
import ufl
from .configuration import FrozenHasTraits
from .solver2d import FlowSolver2d
from .utility import create_directory, print_function_value_range, get_functionspace, unfrozen, domain_constant
from .log import print_output, info, warning
from .diagnostics import GradientRecoverer2D, HessianRecoverer2D
from .exporter import HDF5Exporter
from .callback import DiagnosticCallback
import abc
import numpy
import h5py
from scipy.interpolate import interp1d
import time as time_mod
from pyadjoint.optimization.optimization import SciPyConvergenceError
import os
from mpi4py import MPI
[docs]
class CostFunctionCallback(DiagnosticCallback):
def __init__(self, solver_obj, cost_function, **kwargs):
# Disable logging and HDF5 export
kwargs.setdefault('append_to_log', False)
kwargs.setdefault('export_to_hdf5', False)
super().__init__(solver_obj, **kwargs)
self.cost_function = cost_function
@property
def name(self):
return 'cost_function_callback'
@property
def variable_names(self):
return ['cost_function']
def __call__(self):
# Evaluate the cost function
cost_value = self.cost_function()
return [cost_value]
[docs]
def message_str(self, cost_value):
# Return a string representation of the cost function value
return f"Cost function value: {cost_value}"
[docs]
class ControlManager:
"""
Handles an individual control (spatially varying Function, masked combination of Functions,
or uniform value Function) and its export logic, used internally in InversionManager.
"""
def __init__(self, control, output_dir, index, no_exports=False, mappings=None):
self.output_dir = output_dir
self.index = index
self.no_exports = no_exports
# Wrap single controls in list
self.controls = control if isinstance(control, list) else [control]
self.mappings = mappings
self.is_masked_combination = mappings is not None
self.is_field = False
# Exporters
self.vtk_file = None
self.hdf5_exporter = None
self.gradient_vtk_file = None
# Identify type and setup projection space
if self.is_masked_combination:
for c in self.controls:
if not isinstance(c, fd.Function):
raise ValueError("All masked combination controls must be Functions")
self.mesh = self.controls[0].function_space().mesh()
self.projection_space = fd.FunctionSpace(self.mesh, "CG", 1, variant="equispaced")
self.is_field = False
elif self.controls[0].function_space().ufl_element().family() == "Real":
# domain_constant (uniform value Function in the real space)
self.mesh = self.controls[0].function_space().mesh()
self.projection_space = fd.FunctionSpace(self.mesh, "CG", 1, variant="equispaced")
self.is_field = False
else:
# Spatially varying field
if not isinstance(self.controls[0], fd.Function):
raise TypeError("Control must be a Firedrake Function")
self.mesh = self.controls[0].function_space().mesh()
self.projection_space = None
self.is_field = True
# Initialize exporters if needed
if not no_exports:
fs = self.projection_space if not self.is_field else self.controls[0].function_space()
self.vtk_file = fd.VTKFile(f"{output_dir}/control_progress_{index:02d}.pvd")
prefix = f"control_{index:02d}"
self.hdf5_exporter = HDF5Exporter(fs, output_dir + "/hdf5", prefix)
self.gradient_vtk_file = fd.VTKFile(f"{output_dir}/gradient_progress_{index:02d}.pvd")
[docs]
def project_control(self, updated_controls):
"""Project masked combination or domain_constant to CG1 field for export."""
field = fd.Function(self.projection_space, name="control")
field.assign(0)
if self.is_masked_combination:
for m_, mask_ in zip(updated_controls, self.mappings):
field += m_ * mask_
else:
field.project(updated_controls[0]) # domain_constant
return field
[docs]
def export(self, updated_control):
"""Export control to VTK/HDF5 files."""
if self.no_exports:
return
if self.is_field:
# spatially varying field written directly
if isinstance(updated_control, list):
assert len(updated_control) == 1, "Field export got multiple controls!"
updated_control = updated_control[0]
updated_control.rename(self.controls[0].name())
field = updated_control
else:
# masked combination or domain_constant -> project
if not isinstance(updated_control, list):
updated_control = [updated_control]
field = self.project_control(updated_control)
self.vtk_file.write(field)
self.hdf5_exporter.export(field)
[docs]
def export_gradient(self, updated_gradient):
"""Export gradient to VTK/HDF5 files."""
if self.no_exports:
return
if self.is_field:
if isinstance(updated_gradient, list):
assert len(updated_gradient) == 1, "Field export got multiple controls!"
updated_gradient = updated_gradient[0]
updated_gradient.rename("Gradient")
gradient = updated_gradient
else:
if not isinstance(updated_gradient, list):
updated_gradient = [updated_gradient]
gradient = fd.Function(self.projection_space, name="Gradient")
gradient.assign(0)
if self.is_masked_combination:
for g_, mask_ in zip(updated_gradient, self.mappings):
gradient += g_ * mask_
else: # domain_constant
gradient.project(updated_gradient[0])
self.gradient_vtk_file.write(gradient)
[docs]
class InversionManager(FrozenHasTraits):
"""
Class for handling inversion problems and stashing
the progress of the associated optimization routines.
"""
@unfrozen
def __init__(self, sta_manager, output_dir='outputs', no_exports=False, real=False,
penalty_parameters=[], test_consistency=True, test_gradient=True):
"""
:arg sta_manager: the :class:`StationObservationManagerBase` instance
:kwarg output_dir: model output directory
:kwarg no_exports: if True, nothing will be written to disk
:kwarg real: is the inversion in the Real space?
:kwarg penalty_parameters: a list of penalty parameters to pass
to the :class:`ControlRegularizationManager`
:kwarg test_consistency: toggle testing the correctness with
which the :class:`ReducedFunctional` can recompute values
:kwarg test_gradient: toggle testing the correctness with
which the :class:`ReducedFunctional` can recompute gradients
"""
assert isinstance(sta_manager, StationObservationManagerBase)
self.sta_manager = sta_manager
self.reg_manager = None
self.output_dir = output_dir
self.no_exports = no_exports or real
self.real = real
self.penalty_parameters = penalty_parameters
self.test_consistency = test_consistency
self.test_gradient = test_gradient
self._controls_wrapped = []
self.control_managers = []
self.initialized = False
self.J = 0 # cost function value (float)
self.J_reg = 0 # regularization term value (float)
self.J_misfit = 0 # misfit term value (float)
self.dJdm_list = None # cost function gradient (Function)
self.m_list = None # control (Function)
self.Jhat = None
self.m_progress = []
self.J_progress = []
self.J_reg_progress = []
self.J_misfit_progress = []
self.dJdm_progress = []
self.i = 0
self.tic = None
self.nb_grad_evals = 0
[docs]
def initialize(self):
if not self.no_exports:
if self.real:
raise ValueError("Exports are not supported in Real mode.")
create_directory(self.output_dir)
create_directory(self.output_dir + '/hdf5')
self.initialized = True
[docs]
def add_control(self, controls, mappings=None):
"""
Add control(s).
Can be called multiple times in case of multiparameter optimization.
Parameters
----------
controls : fd.Function or list/tuple of fd.Function
The Function(s) representing the control parameters.
mappings : list of fd.Function, optional
Masks for multi-control cases. If provided, `controls` must be a list of Functions
matching the masks.
Notes
-----
- Pyadjoint `Control` objects are created immediately upon adding the control.
This marks the current state on the tape.
- The value of a control Function should not be altered during a forward solve
outside of the inversion machinery. Changing it beforehand is allowed,
but changing it mid-solve may lead to incorrect gradients or mis-evaluation
of the ReducedFunctional.
- Masked combinations are supported via lists of Functions combined with masks.
"""
if not isinstance(controls, (list, tuple)):
controls = [controls]
for c in controls:
if not isinstance(c, fd.Function):
raise TypeError(f"Control {c} is not a Firedrake Function")
if mappings is not None:
index = len(self.control_managers)
cm = ControlManager(controls, self.output_dir, index, no_exports=self.no_exports, mappings=mappings)
self.control_managers.append(cm)
# store Pyadjoint Control objects
for c in controls:
self._controls_wrapped.append(Control(c))
else:
for f in controls:
index = len(self.control_managers)
cm = ControlManager(f, self.output_dir, index, no_exports=self.no_exports)
self.control_managers.append(cm)
self._controls_wrapped.append(Control(f))
@property
def control_coeff_list(self):
"""
Return a flat list of the underlying Firedrake Functions
used as controls, in the order they were added.
"""
return [c for cm in self.control_managers for c in cm.controls]
@property
def control_list(self):
return self._controls_wrapped
[docs]
def reset_counters(self):
self.nb_grad_evals = 0
[docs]
def set_control_state(self, j, djdm_list, m_list):
"""
Stores optimization state.
To call whenever variables are updated.
:arg j: error functional value
:arg djdm_list: list of gradient functions
:arg m_list: list of control coefficents
"""
self.J = j
self.dJdm_list = djdm_list
self.m_list = m_list
[docs]
def start_clock(self):
self.tic = time_mod.perf_counter()
[docs]
def stop_clock(self):
toc = time_mod.perf_counter()
return toc
[docs]
def set_initial_state(self, *state):
self.set_control_state(*state)
self.update_progress()
[docs]
def update_progress(self):
"""
Updates optimization progress and stores variables to disk.
To call after successful line searches.
"""
toc = self.stop_clock()
if self.i == 0:
for f in self.control_coeff_list:
print_function_value_range(f, name=f.name, prefix='Initial')
elapsed = '-' if self.tic is None else f'{toc - self.tic:.1f} s'
self.tic = toc
if not self.initialized:
self.initialize()
# cost function and gradient norm output
djdm = [fd.norm(f) for f in self.dJdm_list]
if self.real:
controls = [m.dat.data[0] for m in self.m_list]
self.m_progress.append(controls)
self.J_progress.append(self.J)
self.J_reg_progress.append(self.J_reg)
self.J_misfit_progress.append(self.J_misfit)
self.dJdm_progress.append(djdm)
comm = self.dJdm_list[0].comm
if comm.rank == 0 and not self.no_exports:
if self.real:
numpy.save(f'{self.output_dir}/m_progress', self.m_progress)
numpy.save(f'{self.output_dir}/J_progress', self.J_progress)
numpy.save(f'{self.output_dir}/J_reg_progress', self.J_reg_progress)
numpy.save(f'{self.output_dir}/J_misfit_progress', self.J_misfit_progress)
numpy.save(f'{self.output_dir}/dJdm_progress', self.dJdm_progress)
if len(djdm) > 10:
djdm = f"[{numpy.min(djdm):.4e} .. {numpy.max(djdm):.4e}]"
else:
djdm = "[" + ", ".join([f"{dj:.4e}" for dj in djdm]) + "]"
print_output(f'line search {self.i:2d}: '
f'J={self.J:.3e}, dJdm={djdm}, '
f'grad_ev={self.nb_grad_evals}, duration {elapsed}')
if not self.no_exports:
ref_index = 0
for cm in self.control_managers:
num_controls = len(cm.mappings) if cm.mappings is not None else 1
controls_slice = self.m_list[ref_index: ref_index + num_controls]
cm.export(controls_slice)
gradient_slice = self.dJdm_list[ref_index: ref_index + num_controls]
cm.export_gradient(gradient_slice)
ref_index += num_controls
self.i += 1
self.reset_counters()
@property
def rf_kwargs(self):
"""
Default keyword arguments to pass to the
:class:`ReducedFunctional` class.
"""
def functional_eval_cb(j, m):
"""Called after functional evaluation"""
self.J = float(j)
# collect additional data we want before it is discarded if we have garbage collection
tape = get_working_tape()
reg_blocks = tape.get_blocks(tag="reg_eval")
self.J_reg = sum([b.get_outputs()[0].saved_output for b in reg_blocks])
misfit_blocks = tape.get_blocks(tag="misfit_eval")
self.J_misfit = sum([b.get_outputs()[0].saved_output for b in misfit_blocks])
self.sta_manager.collect_time_series(self.i)
return j
def gradient_eval_cb(j, djdm, m):
"""Called after gradient evaluation"""
self.set_control_state(self.J, djdm, m) # must be self.J as j is reset by garbage collection
self.nb_grad_evals += 1
return djdm
params = {
'eval_cb_post': functional_eval_cb,
'derivative_cb_post': gradient_eval_cb,
}
return params
[docs]
def get_cost_function(self, solver_obj, weight_by_variance=False, regularisation_manager="Hessian",
use_local_element_size=True, regularisation_scaling_factor=1.0):
r"""
Get a sum of square errors cost function for the inverse problem.
.. math::
J(u) = J_\mathrm{reg} +
\sum_{i=1}^{n_{ts}}
\sum_{k=1}^{n_\mathrm{types}}
\sum_{j=1}^{n_{\mathrm{sta},k}}
w_k (u_{j,k}^{(i)} - u_{j,k,o}^{(i)})^2,
where :math:`u_{j,k,o}^{(i)}` and :math:`u_{j,k}^{(i)}` denote the observed
and computed values for observation type :math:`k` at timestep :math:`i`,
and :math:`n_{ts}` and :math:`n_{\mathrm{sta},k}` are the number of timesteps
and stations, respectively.
The total cost functional :math:`J(u)` may include:
* **Misfit terms** for one or more observation types (e.g. elevation, velocity),
combined via a :class:`MultiStationObservationManager`.
* **Station-wise weighting** by the inverse of the variance of the observed
data, ensuring stations with larger signal magnitudes do not dominate.
* **Regularization terms** on the control fields, computed once at setup
using either *Hessian*, *Gradient*, or *R-space* smoothness penalties.
The function returns a callable ``cost_fn()`` that should be evaluated at
each model export during the forward simulation to accumulate the
time-dependent misfit contribution.
:arg solver_obj: the :class:`FlowSolver2d` instance
:kwarg weight_by_variance: whether to weight each station by the inverse
of its data variance to balance different magnitude signals
:kwarg regularisation_manager: which regularisation to apply
(``"Hessian"`` or ``"Gradient"``; ``"R-space"`` is used automatically
for real-valued controls)
:kwarg use_local_element_size: whether to use local (spatially varying)
element size to normalise regularisation units
:kwarg regularisation_scaling_factor: global factor applied *only* to
regularization terms
:returns: a callable ``cost_fn`` that accumulates the misfit term at each
timestep during the forward model run
"""
assert isinstance(solver_obj, FlowSolver2d)
# --- Regularisation setup ---
if self.real:
calculator = RSpaceRegularizationCalculator
calculator_kwargs = None
print_output("R-space regularisation being applied (real case).")
else:
calculator_kwargs = {"use_local_element_size": use_local_element_size}
if regularisation_manager == "Hessian":
calculator = HessianRegularizationCalculator
print_output("Hessian regularisation being applied.")
elif regularisation_manager == "Gradient":
calculator = GradientRegularizationCalculator
print_output("Gradient regularisation being applied.")
else:
raise ValueError(
f"Unsupported regularisation_manager: '{regularisation_manager}'. "
"Must be one of: 'Hessian', 'Gradient'."
)
# --- Build regularisation manager ---
if len(self.penalty_parameters) > 0:
self.reg_manager = ControlRegularizationManager(
self.control_coeff_list,
self.penalty_parameters,
regularisation_scaling_factor,
calculator,
calculator_kwargs=calculator_kwargs
)
# --- Compute regularisation term once ---
self.J_reg = 0
if self.reg_manager is not None:
self.J_reg = self.reg_manager.eval_cost_function()
# --- Initialize J ---
self.J_misfit = 0
self.J = self.J_reg
# --- Station variance weighting ---
if isinstance(self.sta_manager, MultiStationObservationManager) and not weight_by_variance:
names = list(self.sta_manager.observation_managers.keys())
warning(
f"You are using a MultiStationObservationManager with weight_by_variance=False. "
f"Consider enabling weight_by_variance to avoid magnitude bias between observation types: {names}"
)
if weight_by_variance:
def apply_weighting(sta_mgr):
var = fd.Function(sta_mgr.fs_points_0d_scalar)
# in parallel access to .dat.data should be collective
if len(var.dat.data[:]) > 0:
for i, j in enumerate(sta_mgr.local_station_index):
if sta_mgr.observed_quantity_is_scalar:
obs = sta_mgr.observation_data[j]
var.dat.data[i] = numpy.var(obs)
else:
u_list, v_list = sta_mgr.observation_data
u = u_list[j]
v = v_list[j]
magnitude = numpy.sqrt(u ** 2 + v ** 2)
var.dat.data[i] = numpy.var(magnitude)
assert numpy.all(numpy.isfinite(var.dat.data[:])), \
(f"[{fd.COMM_WORLD.rank}] ERROR: Check for NaNs. Found non-finite variances of "
f"observation data: {var.dat.data[:]}")
sta_mgr.station_weight_0d.interpolate(1. / var)
if isinstance(self.sta_manager, MultiStationObservationManager):
for sub_mgr in self.sta_manager.observation_managers.values():
apply_weighting(sub_mgr)
else:
apply_weighting(self.sta_manager)
# --- Time-dependent cost function (only adds misfit) ---
def cost_fn():
t = solver_obj.simulation_time
misfit = self.sta_manager.eval_cost_function(t)
self.J_misfit += misfit
self.J += misfit
return cost_fn
@property
def reduced_functional(self):
"""
Create a Pyadjoint :class:`ReducedFunctional` for the optimization.
"""
if self.Jhat is None:
self.Jhat = ReducedFunctional(self.J, self.control_list, **self.rf_kwargs)
return self.Jhat
[docs]
def stop_annotating(self):
"""
Stop recording operations for the adjoint solver.
This method should be called after the :meth:`iterate`
method of :class:`FlowSolver2d`.
"""
assert self.reduced_functional is not None
if self.test_consistency:
self.consistency_test()
if self.test_gradient:
self.taylor_test()
pause_annotation()
[docs]
def get_optimization_callback(self):
"""
Get a callback for stashing optimization progress
after successful line search.
"""
def optimization_callback(m):
self.update_progress()
if not self.no_exports:
self.sta_manager.dump_time_series()
return optimization_callback
[docs]
def minimize(self, opt_method="BFGS", bounds=None, **opt_options):
"""
Minimize the reduced functional using a given optimization routine.
:kwarg opt_method: the optimization routine
:kwarg bounds: a list of bounds to pass to the optimization routine
:kwarg opt_options: other optimization parameters to pass
"""
print_output(f'Running {opt_method} optimization')
self.reset_counters()
self.start_clock()
J = float(self.reduced_functional(self.control_coeff_list))
self.set_initial_state(J, self.reduced_functional.derivative(apply_riesz=True), self.control_coeff_list)
if not self.no_exports:
self.sta_manager.collect_time_series(self.i)
self.sta_manager.dump_time_series()
try:
return minimize(
self.reduced_functional, method=opt_method, bounds=bounds,
callback=self.get_optimization_callback(), options=opt_options)
except SciPyConvergenceError as e:
if "TOTAL NO. OF ITERATIONS REACHED LIMIT" in str(e):
print_output("Optimization stopped: reached iteration limit.")
return self.m_list
else:
raise
[docs]
def consistency_test(self):
"""
Test that :attr:`reduced_functional` can correctly recompute the
objective value, assuming that none of the controls have changed
since it was created.
"""
print_output("Running consistency test")
J = self.reduced_functional(self.control_coeff_list)
if not numpy.isclose(J, self.J):
raise ValueError(f"Consistency test failed (expected {self.J}, got {J})")
print_output("Consistency test passed!")
[docs]
def taylor_test(self):
"""
Run a Taylor test to check that the :attr:`reduced_functional` can
correctly compute consistent gradients.
Note that the Taylor test is applied on the current control values.
"""
func_list = []
for f in self.control_coeff_list:
dc = f.copy(deepcopy=True)
func_list.append(dc)
minconv = taylor_test(self.reduced_functional, self.control_coeff_list, func_list)
if minconv < 1.9:
raise ValueError("Taylor test failed") # NOTE: Pyadjoint already prints the testing
print_output("Taylor test passed!")
[docs]
class StationObservationManagerBase(abc.ABC):
"""
Abstract base class defining the interface for station observation managers.
A station observation manager is responsible for evaluating the misfit
between model output and observational data at station locations during
a forward model run, and for storing the resulting model time series to
disk for diagnostic purposes.
Subclasses may handle a single observation variable
(:class:`StationObservationManager`) or combine several
(:class:`MultiStationObservationManager`). Both expose the same
interface so that :class:`InversionManager` can use either without
special-casing.
"""
[docs]
@abc.abstractmethod
def eval_cost_function(self, t) -> float:
"""
Evaluate the misfit cost function at time `t`.
Should be called at every export of the forward model. Implementations
are expected to accumulate contributions over time so that the total
misfit across the simulation is available at the end of the forward solve.
:arg t: model simulation time
:returns: scalar misfit value at this timestep
"""
[docs]
@abc.abstractmethod
def collect_time_series(self, current_iteration: int):
"""
Collect model time series from the current forward solve.
Should be called after the forward solve completes, before the
PyAdjoint tape is cleared. If called multiple times within the
same iteration (e.g. during a line search), the stored series
should be updated in-place rather than appended.
:arg current_iteration: current optimisation iteration index
"""
[docs]
@abc.abstractmethod
def dump_time_series(self):
"""
Write collected model time series to disk.
:meth:`collect_time_series` must have been called at least once
before this method is invoked.
"""
[docs]
class StationObservationManager(StationObservationManagerBase):
"""
Implements error functional based on observation time series.
The functional is the squared sum of error between the model and
observations.
This object evaluates the model fields at the station locations,
interpolates the observations time series to the model time, computes the
error functional, and also stores the model's time series data to disk.
"""
def __init__(self, mesh, output_directory='outputs'):
"""
:arg mesh: the 2D mesh object.
:kwarg output_directory: directory where model time series are stored.
"""
self.mesh = mesh
on_sphere = self.mesh.geometric_dimension == 3
if on_sphere:
raise NotImplementedError('Sphere meshes are not supported yet.')
self.cost_function_scaling = fd.Constant(1.0)
self.output_directory = output_directory
# keep observation time series in memory
self.obs_func_list = []
# keep model time series in memory during optimization progress
self.station_value_progress = []
# model time when cost function was evaluated
self.simulation_time = []
self.model_observation_field = None
self.initialized = False
self.last_recorded_iteration = -1
[docs]
def register_observation_data(self, station_names, variable, time,
x, y, data=None, start_times=None, end_times=None):
"""
Add station time series data to the object.
The `x`, and `y` coordinates must be such that
they allow extraction of model data at the same coordinates.
Supports arbitrary scalar or 2D vector data for observation data.
:arg list station_names: list of station names
:arg str variable: canonical variable name, e.g. 'elev', 'uv'
:arg list time: array of time stamps, one for each station
:arg list data:
- scalar: list of arrays, one per station
- vector: tuple (list of arrays, list of arrays)
:arg list x: list of station x coordinates
:arg list y: list of station y coordinates
:kwarg list start_times: optional start times for the observation periods
:kwarg list end_times: optional end times for the observation periods
"""
if data is None:
raise ValueError("Must provide observation data.")
if isinstance(data, tuple):
u, v = data
if not isinstance(u, list) or not isinstance(v, list):
raise ValueError("For vector data, 'data' must be a tuple/list of two lists.")
self.observed_quantity_is_scalar = False
elif isinstance(data, list):
# always treat as scalar
self.observed_quantity_is_scalar = True
else:
raise ValueError("Invalid data format. Must be list (scalar) or tuple (vector).")
print_output(f"Data registered as {'scalar' if self.observed_quantity_is_scalar else 'vector'} for variable "
f"'{variable}' with {len(station_names)} station(s).")
self.station_names = station_names
self.variable = variable
self.observation_time = time
self.observation_data = data
self.observation_x = x
self.observation_y = y
num_stations = len(station_names)
self._start_times = start_times or -numpy.ones(num_stations)*numpy.inf
self._end_times = end_times or numpy.ones(num_stations)*numpy.inf
[docs]
def set_model_field(self, function):
"""
Set the model field that will be evaluated.
"""
self.model_observation_field = function
[docs]
def load_scalar_observation_data(self, observation_data_dir, station_names, variable,
start_times=None, end_times=None):
"""
Load scalar observation data (e.g. elevation) from disk.
This assumes the data was stored with `TimeSeriesCallback2D` during the
forward run. For generic/custom sources (or vectors e.g. velocity), use
`register_observation_data` directly.
:arg str observation_data_dir: directory where observation data is stored
:arg list station_names: list of station names
:arg str variable: canonical variable name, e.g. 'elev'
:kwarg list start_times: optional start times for the observation periods
:kwarg list end_times: optional end times for the observation periods
"""
file_list = [
f'{observation_data_dir}/'
f'diagnostic_timeseries_{s}_{variable}.hdf5' for s in
station_names
]
# arrays of time stamps and values for each station
observation_time = []
observation_values = []
observation_coords = []
for f in file_list:
with h5py.File(f, 'r') as h5file:
t = h5file['time'][:].flatten()
v = h5file[variable][:].flatten()
x = h5file.attrs['x']
y = h5file.attrs['y']
observation_coords.append((x, y))
observation_time.append(t)
observation_values.append(v)
# list of station coordinates
observation_x, observation_y = numpy.array(observation_coords).T
self.register_observation_data(
station_names, variable, observation_time,
observation_x, observation_y, data=observation_values,
start_times=start_times, end_times=end_times,
)
self.construct_evaluator()
[docs]
def update_stations_in_use(self, t):
"""
Indicate which stations are in use at the current time.
An entry of unity indicates use, whereas zero indicates disuse.
"""
if not hasattr(self, 'obs_start_times'):
self.construct_evaluator()
in_use = fd.Function(self.fs_points_0d_scalar)
in_use.dat.data[:] = numpy.array(
numpy.bitwise_and(
self.obs_start_times <= t, t <= self.obs_end_times
), dtype=float)
self.indicator_0d.assign(in_use)
[docs]
def construct_evaluator(self):
"""
Builds evaluators needed to compute the error functional.
"""
# Create 0D mesh for station evaluation
xy = numpy.array((self.observation_x, self.observation_y)).T
mesh0d = fd.VertexOnlyMesh(self.mesh, xy)
self.fs_points_0d_scalar = fd.FunctionSpace(mesh0d, 'DG', 0)
if self.observed_quantity_is_scalar:
self.fs_points_0d = self.fs_points_0d_scalar
else:
self.fs_points_0d = fd.VectorFunctionSpace(mesh0d, 'DG', 0, dim=2)
self.obs_values_0d = fd.Function(self.fs_points_0d, name='observations')
self.mod_values_0d = fd.Function(self.fs_points_0d, name='model values')
self.indicator_0d = fd.Function(self.fs_points_0d_scalar, name='station use indicator')
self.indicator_0d.assign(1.0)
self.cost_function_scaling_0d = domain_constant(0.0, mesh0d)
self.cost_function_scaling_0d.assign(self.cost_function_scaling)
self.station_weight_0d = fd.Function(self.fs_points_0d_scalar, name='station-wise weighting')
self.station_weight_0d.assign(1.0)
interp_kw = {}
if numpy.isfinite(self._start_times).any() or numpy.isfinite(self._end_times).any():
interp_kw.update({'bounds_error': False, 'fill_value': 0.0})
# Construct timeseries interpolator
self.station_interpolators = []
self.local_station_index = []
if len(mesh0d.coordinates.dat.data[:]) > 0:
for i in range(self.fs_points_0d_scalar.dof_dset.size):
# loop over local DOFs and match coordinates to observations
# NOTE this must be done manually as VertexOnlyMesh reorders points
x_mesh, y_mesh = mesh0d.coordinates.dat.data[i, :]
xy_diff = xy - numpy.array([x_mesh, y_mesh])
xy_dist = numpy.sqrt(xy_diff[:, 0]**2 + xy_diff[:, 1]**2)
j = numpy.argmin(xy_dist)
self.local_station_index.append(j)
x, y = xy[j, :]
t = self.observation_time[j]
x_mesh, y_mesh = mesh0d.coordinates.dat.data[i, :]
msg = 'bad station location ' \
f'{j} {i} {x} {x_mesh} {y} {y_mesh} {x - x_mesh} {y - y_mesh}'
assert numpy.allclose([x, y], [x_mesh, y_mesh]), msg
if self.observed_quantity_is_scalar:
ip = interp1d(t, self.observation_data[j], **interp_kw)
self.station_interpolators.append(ip)
else:
u_data, v_data = self.observation_data
ip_u = interp1d(t, u_data[j], **interp_kw)
ip_v = interp1d(t, v_data[j], **interp_kw)
self.station_interpolators.append((ip_u, ip_v))
# Process start and end times for observations
self.obs_start_times = numpy.array([
self._start_times[i] for i in self.local_station_index
])
self.obs_end_times = numpy.array([
self._end_times[i] for i in self.local_station_index
])
# expressions for cost function
if self.observed_quantity_is_scalar:
self.misfit_expr = self.obs_values_0d - self.mod_values_0d
else:
diff = self.obs_values_0d - self.mod_values_0d
self.misfit_expr = fd.sqrt(fd.inner(diff, diff))
self.initialized = True
[docs]
def eval_observation_at_time(self, t):
"""
Evaluate observation time series at the given time.
:arg t: model simulation time
:returns: list of observation time series values at time `t`
"""
self.update_stations_in_use(t)
if self.observed_quantity_is_scalar:
ip = [float(ip(t)) for ip in self.station_interpolators]
return ip
else:
ip_u = [float(ip_u(t)) for ip_u, _ in self.station_interpolators]
ip_v = [float(ip_v(t)) for _, ip_v in self.station_interpolators]
return ip_u, ip_v
[docs]
def eval_cost_function(self, t):
"""
Evaluate the cost function.
Should be called at every export of the forward model.
:arg t: model simulation time
:returns: scalar misfit cost function value
"""
assert self.initialized, 'Not initialized, call construct_evaluator first.'
assert self.model_observation_field is not None, 'Model field not set.'
self.simulation_time.append(t)
# observed data
obs_func = fd.Function(self.fs_points_0d)
if self.observed_quantity_is_scalar:
obs_func.dat.data[:] = self.eval_observation_at_time(t)
self.obs_func_list.append(obs_func)
else:
obs_u_data, obs_v_data = self.eval_observation_at_time(t)
obs_func.dat.data[:, 0] = obs_u_data
obs_func.dat.data[:, 1] = obs_v_data
self.obs_func_list.append(obs_func)
self.obs_values_0d.assign(obs_func)
# modelled data
if self.observed_quantity_is_scalar:
P1_2d = get_functionspace(self.mesh, 'CG', 1)
mod_func = fd.Function(P1_2d, name=self.variable)
else:
P1_2vd = get_functionspace(self.mesh, 'CG', 1, vector=True)
mod_func = fd.Function(P1_2vd, name=self.variable)
mod_func.project(self.model_observation_field)
self.mod_values_0d.interpolate(mod_func, ad_block_tag=f'{self.variable} observation')
# compute square error
s = self.cost_function_scaling_0d * self.indicator_0d * self.station_weight_0d
self.J_misfit = fd.assemble(s * self.misfit_expr ** 2 * fd.dx, ad_block_tag='misfit_eval')
return self.J_misfit
[docs]
def collect_time_series(self, current_iteration):
"""
Collect the model time-series from the current forward solve.
This extracts the saved station-observation values from the PyAdjoint
tape *before it is cleared*, and stores them in memory.
If multiple forward solves occur within the same optimisation iteration
(e.g. due to a line search), the time-series for that iteration is
*updated* rather than appended — ensuring that only the final accepted
forward model state is stored.
"""
tape = get_working_tape()
blocks = tape.get_blocks(tag=f'{self.variable} observation')
ts_data = [b.get_outputs()[0].saved_output.dat.data for b in blocks] # .copy()?
ts_data = numpy.array(ts_data) # shape (ntimesteps, nstations, ndims) ndims = 2 for vector
if current_iteration == self.last_recorded_iteration:
self.station_value_progress[-1] = ts_data
else:
self.station_value_progress.append(ts_data)
self.last_recorded_iteration = current_iteration
[docs]
def dump_time_series(self):
"""
Write the collected model time-series to disk.
Station time series should be obtained from the last optimization
iteration using collect_time_series(iteration_number) before calling.
The output files have the format
`{odir}/diagnostic_timeseries_progress_{station_name}_{variable}.hdf5`
The file contains the simulation time in the `time` array, and the
station name and coordinates as attributes. The time series data is
stored as a rank 2 (n_iterations, n_time_steps) array for scalar,
or rank 3 (n_iterations, n_time_steps, 2) for vector quantities.
"""
assert self.station_names is not None
assert self.last_recorded_iteration >= 0, (
"dump_time_series() called before any time-series were collected. "
"Ensure collect_time_series() is called at least once (e.g. via "
"the functional post-evaluation callback)."
)
create_directory(self.output_directory)
var = self.variable
for ilocal, iglobal in enumerate(self.local_station_index):
name = self.station_names[iglobal]
ts = numpy.array([s[:, ilocal] for s in self.station_value_progress]) # shape (niters, ntimesteps, ?)
fn = f'diagnostic_timeseries_progress_{name}_{var}.hdf5'
fn = os.path.join(self.output_directory, fn)
with h5py.File(fn, 'w') as hdf5file:
if self.observed_quantity_is_scalar:
hdf5file.create_dataset(var, data=ts) # scalar data
else:
hdf5file.create_dataset(f'{var}_u_component', data=ts[:, :, 0])
hdf5file.create_dataset(f'{var}_v_component', data=ts[:, :, 1])
hdf5file.create_dataset('time', data=numpy.array(self.simulation_time))
attrs = {
'x': self.observation_x[iglobal],
'y': self.observation_y[iglobal],
'location_name': name,
}
hdf5file.attrs.update(attrs)
[docs]
class MultiStationObservationManager(StationObservationManagerBase):
"""
Manage multiple StationObservationManager instances (e.g. elevation + velocity).
Combines their cost functions into a single scalar J.
NOTE: It is recommended that `weight_by_variance=True` for each manager to
avoid magnitude bias between observation types.
"""
def __init__(self, managers):
"""
:param managers: dict {name: StationObservationManager}
"""
self.observation_managers = managers
# --- Diagnostic checks on scaling factors ---
scalings = {}
print_output("Initialising MultiStationObservationManager with station managers:")
for name, som in self.observation_managers.items():
if hasattr(som, "cost_function_scaling"):
try:
scal = float(som.cost_function_scaling)
except Exception:
scal = None
scalings[name] = scal
print_output(f" - {name}: cost_function_scaling = {scal}")
else:
warning(f" - {name}: missing cost_function_scaling attribute")
# --- Check for uninitialised or inconsistent scaling ---
for name, val in scalings.items():
if val == 1.0:
warning(
f"Station manager '{name}' has cost_function_scaling=1.0 "
"(default value). Did you set it before construct_evaluator()?"
)
unique_vals = {v for v in scalings.values() if v is not None}
if len(unique_vals) > 1:
diff_list = ", ".join([f"{n}: {v}" for n, v in scalings.items()])
warning(
f"Detected differing cost_function_scaling values between managers: {diff_list}. "
"Ensure these reflect your intended relative weighting."
)
else:
info("All station managers share consistent cost_function_scaling values.")
[docs]
def eval_cost_function(self, t):
"""
Evaluate combined cost function for all registered observation sets.
:arg t: model simulation time
:returns: scalar total cost function value (sum over all managers)
"""
J_total = 0.0
for som in self.observation_managers.values():
J_total += som.eval_cost_function(t)
return J_total
[docs]
def collect_time_series(self, current_iteration):
"""
Forward the collect_time_series call to each station manager.
"""
for som in self.observation_managers.values():
som.collect_time_series(current_iteration)
[docs]
def dump_time_series(self):
"""
Forward the dump_time_series call to each station manager.
"""
for som in self.observation_managers.values():
som.dump_time_series()
[docs]
class RegularizationCalculator(abc.ABC):
"""
Base class for computing regularization terms.
A derived class should set :attr:`regularization_expr` in
:meth:`__init__`. Whenever the cost function is evaluated,
the ratio of this expression and the total mesh area will
be added.
"""
@abc.abstractmethod
def __init__(self, function, scaling=1.0):
"""
:arg function: Control :class:`Function`
"""
self.scaling = scaling
self.regularization_expr = 0
self.mesh = function.function_space().mesh()
# calculate mesh area (to scale the cost function)
self.mesh_area = fd.assemble(fd.Constant(1.0) * fd.dx(domain=self.mesh))
self.name = function.name()
[docs]
def eval_cost_function(self):
expr = self.scaling * self.regularization_expr / self.mesh_area * fd.dx
return fd.assemble(expr, ad_block_tag="reg_eval")
[docs]
class GradientRegularizationCalculator(RegularizationCalculator):
r"""
Computes the following regularization term for a control Function `f`:
.. math::
J = \gamma \| (\Delta x) \nabla f \|^2,
where:
- :math:`\nabla f` is the gradient of the field `f`.
- :math:`\Delta x` is the characteristic element size (cell size)
used to scale the gradient regularization.
The element size defaults to varying spatially with the mesh,
but can be set constant with the smallest element size by specifying
the `use_local_element_size` flag to be False.
"""
def __init__(self, function, gradient_penalty, scaling=1.0, use_local_element_size=True):
"""
:arg function: Control :class:`Function`
:arg gradient_penalty: Gradient penalty coefficient
:kwarg scaling: Optional global scaling for the regularization term
"""
super().__init__(function, scaling=scaling)
# Setup continuous vector function space
P1v_2d = get_functionspace(self.mesh, "CG", 1, vector=True)
self.gradient_2d = fd.Function(P1v_2d, name=f"{self.name} gradient")
# Recover gradient using projection
self.gradient_recoverer = GradientRecoverer2D(
function, self.gradient_2d)
h = fd.CellSize(self.mesh)
if not use_local_element_size:
V = fd.FunctionSpace(self.mesh, "DG", 0)
h_ = fd.Function(V)
h_.project(h)
local_min = h_.dat.data.min()
global_min = self.mesh.comm.allreduce(local_min, op=MPI.MIN)
h = fd.Function(V).assign(global_min)
# Regularization term: |grad(f)|^2 * h^2
self.regularization_expr = gradient_penalty * fd.inner(self.gradient_2d, self.gradient_2d) * h**2
[docs]
def eval_cost_function(self):
self.gradient_recoverer.solve()
return super().eval_cost_function()
[docs]
class HessianRegularizationCalculator(RegularizationCalculator):
r"""
Computes the following regularization term for a cost function
involving a control :class:`Function` :math:`f`:
.. math::
J = \gamma \| (\Delta x)^2 H(f) \|^2,
where:
- :math:`H` is the Hessian of field :math:`f`.
- :math:`\Delta x` is the characteristic element size (cell size)
used to scale the gradient regularization.
The element size defaults to varying spatially with the mesh,
but can be set constant with the smallest element size by specifying
the `use_local_element_size` flag to be False.
"""
def __init__(self, function, hessian_penalty, scaling=1.0, use_local_element_size=True):
"""
:arg function: Control :class:`Function`
:arg hessian_penalty: Hessian penalty coefficient
"""
super().__init__(function, scaling=scaling)
# solvers to evaluate the gradient of the control
P1v_2d = get_functionspace(self.mesh, "CG", 1, vector=True)
P1t_2d = get_functionspace(self.mesh, "CG", 1, tensor=True)
gradient_2d = fd.Function(P1v_2d, name=f"{self.name} gradient")
hessian_2d = fd.Function(P1t_2d, name=f"{self.name} hessian")
self.hessian_calculator = HessianRecoverer2D(
function, hessian_2d, gradient_2d)
h = fd.CellSize(self.mesh)
if not use_local_element_size:
V = fd.FunctionSpace(self.mesh, "DG", 0)
h_ = fd.Function(V)
h_.project(h)
local_min = h_.dat.data.min()
global_min = self.mesh.comm.allreduce(local_min, op=MPI.MIN)
h = fd.Function(V).assign(global_min)
# regularization expression |hessian|^2
# NOTE this is normalized by the mesh element size
# d^2 u/dx^2 * dx^2 ~ du^2
self.regularization_expr = hessian_penalty * fd.inner(hessian_2d, hessian_2d) * h**4
[docs]
def eval_cost_function(self):
self.hessian_calculator.solve()
return super().eval_cost_function()
[docs]
class RSpaceRegularizationCalculator(RegularizationCalculator):
r"""
Computes the following regularization term for a cost function
involving a control :class:`Function` :math:`f` from an R-space:
.. math::
J = \gamma (f - f_0)^2,
where :math:`f_0` is a prior, taken to be the initial value of
:math:`f`.
"""
def __init__(self, function, gamma, eps=1.0e-03, scaling=1.0):
"""
:arg function: Control :class:`Function`
:arg gamma: penalty coefficient
:kwarg eps: tolerance for normalising by near-zero priors
"""
super().__init__(function, scaling=scaling)
R = function.function_space()
if R.ufl_element().family() != "Real":
raise ValueError("function must live in R-space")
prior = fd.Function(R, name=f"{self.name} prior")
prior.assign(function, annotate=False) # Set the prior to the initial value
self.regularization_expr = gamma * (function - prior) ** 2 / ufl.max_value(abs(prior), eps)
# NOTE: If the prior is small then dividing by prior**2 puts too much emphasis
# on the regularization. Therefore, we divide by abs(prior) instead.
[docs]
class ControlRegularizationManager:
"""
Handles regularization of multiple control fields
"""
def __init__(self, function_list, gamma_list, penalty_term_scaling=None,
calculator=HessianRegularizationCalculator, calculator_kwargs=None):
"""
:arg function_list: list of control functions
:arg gamma_list: list of penalty parameters
:kwarg penalty_term_scaling: Penalty term scaling factor
:kwarg calculator: class used for obtaining regularization
"""
if calculator_kwargs is None:
calculator_kwargs = {}
self.reg_calculators = []
assert len(function_list) == len(gamma_list), \
'Number of control functions and parameters must match'
self.reg_calculators = [
calculator(f, g, scaling=penalty_term_scaling, **calculator_kwargs)
for f, g in zip(function_list, gamma_list)
]
[docs]
def eval_cost_function(self):
return sum([r.eval_cost_function() for r in self.reg_calculators])
