"""Adjoint reconstruction for Caputo-Wismer initial-pressure imaging."""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from time import perf_counter
from typing import Any, Literal
import numpy as np
from scipy.optimize import minimize
from yonderdrake.applications.caputo_wismer.forward import (
AttenuationMode,
CaputoWismerForwardOperator,
_nonnegative_real,
_positive_integer,
)
from yonderdrake.applications.caputo_wismer.model import (
CaputoWismerMaterial,
SensorArray,
)
ReconstructionMethod = Literal["kaltenbacher", "adjoint"]
[docs]
@dataclass(frozen=True, slots=True)
class CaputoWismerReconstruction:
"""Result and diagnostics from an adjoint reconstruction."""
pressure: Any
converged: bool
iterations: int
objective: float
objective_history: tuple[float, ...]
message: str
function_evaluations: int
forward_seconds: float
adjoint_seconds: float
elapsed_seconds: float
[docs]
class CaputoWismerInverseProblem:
"""Kaltenbacher-style regularized initial-pressure reconstruction."""
def __init__(
self,
space: Any,
sensor_data: Any,
sensors: SensorArray,
*,
materials: Sequence[CaputoWismerMaterial],
dt: float,
num_modes: int = 32,
representation: Any = None,
absorbing_speed: Any = None,
attenuation: AttenuationMode = "dissipative",
stiffness_theta: float = 0.0,
regularization: float = 1.0e-6,
solver_parameters: Any = None,
) -> None:
values = np.asarray(sensor_data, dtype=np.float64)
if (
values.ndim != 2
or values.shape[0] < 2
or values.shape[1] != sensors.num_sensors
):
raise ValueError(
"sensor_data must have shape (num_times, num_sensors)"
)
if not np.all(np.isfinite(values)):
raise ValueError("sensor_data must be finite")
self.space = space
self.sensor_data = values.copy()
self.regularization = _nonnegative_real(
regularization,
"regularization",
)
self.operator = CaputoWismerForwardOperator(
space,
sensors,
materials=materials,
dt=dt,
num_steps=values.shape[0] - 1,
num_modes=num_modes,
representation=representation,
absorbing_speed=absorbing_speed,
attenuation=attenuation,
stiffness_theta=stiffness_theta,
solver_parameters=solver_parameters,
)
self._candidate = self.operator._function(
"initial_pressure_candidate"
)
self._time_weights = np.full(values.shape[0], float(dt))
self._time_weights[[0, -1]] *= 0.5
self._function_evaluations = 0
self._forward_seconds = 0.0
self._adjoint_seconds = 0.0
[docs]
def objective_gradient(self, candidate: Any) -> tuple[float, Any]:
"""Return the Tikhonov objective and its coefficient-space gradient."""
if candidate.function_space() != self.space:
raise ValueError("candidate must belong to the inversion space")
started = perf_counter()
residual = self.operator.forward(candidate) - self.sensor_data
self._forward_seconds += perf_counter() - started
weighted_residual = self._time_weights[:, None] * residual
objective = 0.5 * float(np.sum(residual * weighted_residual))
started = perf_counter()
gradient = self.operator.adjoint_covector(weighted_residual)
self._adjoint_seconds += perf_counter() - started
self._function_evaluations += 1
if self.regularization:
objective += 0.5 * self.regularization * float(
self.operator._fd.assemble(candidate * candidate
* self.operator._fd.dx)
)
self.operator._matrix_axpy(
self.operator._mass,
candidate,
gradient,
self.regularization,
)
return objective, gradient
def _scaled_adjoint_initial_guess(
self,
*,
positivity: bool,
) -> Any:
weighted_data = self._time_weights[:, None] * self.sensor_data
started = perf_counter()
direction = self.operator.adjoint(weighted_data)
self._adjoint_seconds += perf_counter() - started
if positivity:
direction.dat.data[:] = np.maximum(
direction.dat.data_ro,
0.0,
)
started = perf_counter()
traces = self.operator.forward(direction)
self._forward_seconds += perf_counter() - started
numerator = float(
np.sum(self.sensor_data * self._time_weights[:, None] * traces)
)
denominator = float(
np.sum(traces * self._time_weights[:, None] * traces)
)
if self.regularization:
denominator += self.regularization * float(
self.operator._fd.assemble(
direction * direction * self.operator._fd.dx
)
)
scale = (
max(0.0, numerator / denominator)
if denominator > 0.0
else 0.0
)
direction *= scale
return direction
def _solve_distributed(
self,
*,
max_iterations: int | None,
tolerance: float,
positivity: bool,
started: float,
) -> CaputoWismerReconstruction:
from petsc4py import PETSc
history: list[float] = []
solution = self._candidate.copy(deepcopy=True)
gradient_buffer = self.operator._covector(
"tao_objective_gradient"
)
tao = PETSc.TAO().create(comm=self.space.mesh().comm)
tao.setType(
PETSc.TAO.Type.BLMVM
if positivity
else PETSc.TAO.Type.LMVM
)
tao.setTolerances(
gatol=tolerance,
grtol=tolerance,
gttol=tolerance,
)
if max_iterations is not None:
tao.setMaximumIterations(max_iterations)
def objective_gradient(
_tao: Any,
coefficients: Any,
output_gradient: Any,
) -> float:
with self._candidate.dat.vec as candidate_vector:
coefficients.copy(candidate_vector)
value, gradient = self.objective_gradient(self._candidate)
with gradient.dat.vec_ro as gradient_vector:
gradient_vector.copy(output_gradient)
history.append(value)
return value
lower = self.operator._function("tao_lower_bound")
upper = self.operator._function("tao_upper_bound")
lower.assign(0.0)
upper.assign(PETSc.INFINITY)
with (
solution.dat.vec as solution_vector,
gradient_buffer.dat.vec as gradient_vector,
lower.dat.vec_ro as lower_vector,
upper.dat.vec_ro as upper_vector,
):
tao.setObjectiveGradient(
objective_gradient,
gradient_vector,
)
if positivity:
tao.setVariableBounds((lower_vector, upper_vector))
tao.solve(solution_vector)
iterations, objective, _, _, _, reason = tao.getSolutionStatus()
converged = int(reason) > 0
pressure = solution.copy(deepcopy=True)
pressure.rename("kaltenbacher_initial_pressure")
result = CaputoWismerReconstruction(
pressure=pressure,
converged=converged,
iterations=int(iterations),
objective=float(objective),
objective_history=tuple(history),
message=str(reason),
function_evaluations=self._function_evaluations,
forward_seconds=self._forward_seconds,
adjoint_seconds=self._adjoint_seconds,
elapsed_seconds=perf_counter() - started,
)
tao.destroy()
return result
[docs]
def solve(
self,
*,
initial_guess: Any = None,
max_iterations: int | None = None,
tolerance: float = 1.0e-7,
positivity: bool = False,
warm_start: bool = False,
) -> CaputoWismerReconstruction:
"""Minimize the regularized data misfit with L-BFGS-B."""
if max_iterations is not None:
_positive_integer(max_iterations, "max_iterations")
tolerance_value = _nonnegative_real(tolerance, "tolerance")
if tolerance_value == 0.0:
raise ValueError("tolerance must be positive")
self._function_evaluations = 0
self._forward_seconds = 0.0
self._adjoint_seconds = 0.0
started = perf_counter()
if initial_guess is not None:
if initial_guess.function_space() != self.space:
raise ValueError(
"initial_guess must belong to the inversion space"
)
self._candidate.assign(initial_guess)
elif warm_start:
self._candidate.assign(
self._scaled_adjoint_initial_guess(
positivity=positivity,
)
)
else:
self._candidate.assign(0.0)
if self.space.mesh().comm.size > 1:
return self._solve_distributed(
max_iterations=max_iterations,
tolerance=tolerance_value,
positivity=positivity,
started=started,
)
initial_values = np.asarray(
self._candidate.dat.data_ro,
dtype=np.float64,
).copy()
history: list[float] = []
def objective(values: np.ndarray) -> tuple[float, np.ndarray]:
self._candidate.dat.data[:] = values
value, gradient = self.objective_gradient(self._candidate)
history.append(value)
return value, np.asarray(
gradient.dat.data_ro,
dtype=np.float64,
).copy()
options: dict[str, Any] = {
"ftol": tolerance_value,
"gtol": tolerance_value,
}
if max_iterations is not None:
options["maxiter"] = max_iterations
result = minimize(
objective,
initial_values,
method="L-BFGS-B",
jac=True,
bounds=(
[(0.0, None)] * initial_values.size
if positivity
else None
),
options=options,
)
pressure = self.operator._function(
"kaltenbacher_initial_pressure"
)
pressure.dat.data[:] = result.x
return CaputoWismerReconstruction(
pressure=pressure,
converged=bool(result.success),
iterations=int(result.nit),
objective=float(result.fun),
objective_history=tuple(history),
message=str(result.message),
function_evaluations=self._function_evaluations,
forward_seconds=self._forward_seconds,
adjoint_seconds=self._adjoint_seconds,
elapsed_seconds=perf_counter() - started,
)
[docs]
def reconstruct_initial_pressure(
space: Any,
sensor_data: Any,
sensors: SensorArray,
*,
materials: Sequence[CaputoWismerMaterial],
dt: float,
method: ReconstructionMethod = "kaltenbacher",
attenuation: AttenuationMode = "dissipative",
num_modes: int = 32,
representation: Any = None,
absorbing_speed: Any = None,
stiffness_theta: float = 0.0,
regularization: float = 1.0e-6,
initial_guess: Any = None,
max_iterations: int | None = None,
tolerance: float = 1.0e-7,
positivity: bool = False,
warm_start: bool = False,
solver_parameters: Any = None,
) -> Any:
"""Reconstruct initial pressure with the selected Python imaging method."""
if method == "adjoint":
values = np.asarray(sensor_data, dtype=np.float64)
if (
values.ndim != 2
or values.shape[0] < 2
or values.shape[1] != sensors.num_sensors
):
raise ValueError(
"sensor_data must have shape (num_times, num_sensors)"
)
if not np.all(np.isfinite(values)):
raise ValueError("sensor_data must be finite")
operator = CaputoWismerForwardOperator(
space,
sensors,
materials=materials,
dt=dt,
num_steps=values.shape[0] - 1,
num_modes=num_modes,
representation=representation,
absorbing_speed=absorbing_speed,
attenuation=attenuation,
stiffness_theta=stiffness_theta,
solver_parameters=solver_parameters,
)
time_weights = np.full(values.shape[0], float(dt))
time_weights[[0, -1]] *= 0.5
return operator.adjoint(time_weights[:, None] * values)
if method != "kaltenbacher":
raise ValueError("method must be 'kaltenbacher' or 'adjoint'")
problem = CaputoWismerInverseProblem(
space,
sensor_data,
sensors,
materials=materials,
dt=dt,
num_modes=num_modes,
representation=representation,
absorbing_speed=absorbing_speed,
attenuation=attenuation,
stiffness_theta=stiffness_theta,
regularization=regularization,
solver_parameters=solver_parameters,
)
result = problem.solve(
initial_guess=initial_guess,
max_iterations=max_iterations,
tolerance=tolerance,
positivity=positivity,
warm_start=warm_start,
)
if max_iterations is None and not result.converged:
raise RuntimeError(
"initial-pressure reconstruction did not converge: "
f"{result.message}"
)
return result.pressure