"""Optimized Caputo-Wismer forward propagation."""
from __future__ import annotations
from collections.abc import Sequence
from math import isfinite
from typing import Any, Literal
import numpy as np
from yonderdrake.applications.caputo_wismer.model import (
CaputoWismerMaterial,
SensorArray,
)
from yonderdrake.time.coefficients import exp_neg, phi1
from yonderdrake.time.representations import BirkSong
AttenuationMode = Literal["dissipative", "none"]
def _nonnegative_real(value: Any, name: str) -> float:
try:
result = float(value)
except (TypeError, ValueError) as error:
raise TypeError(f"{name} must be a real scalar") from error
if not isfinite(result) or result < 0.0:
raise ValueError(f"{name} must be finite and nonnegative")
return result
def _positive_integer(value: Any, name: str) -> int:
if (
not isinstance(value, int)
or isinstance(value, bool)
or value < 1
):
raise ValueError(f"{name} must be a positive integer")
return value
def _material_model(
materials: Sequence[CaputoWismerMaterial],
attenuation: AttenuationMode,
fd: Any,
) -> tuple[CaputoWismerMaterial, ...]:
values = tuple(materials)
if not values:
raise ValueError("materials must contain at least one material")
if any(not isinstance(value, CaputoWismerMaterial) for value in values):
raise TypeError("materials must contain CaputoWismerMaterial objects")
if attenuation == "dissipative":
return values
if attenuation == "none":
return tuple(
CaputoWismerMaterial(
value.indicator,
value.wave_speed,
fd.Constant(0.0),
value.alpha,
)
for value in values
)
raise ValueError("attenuation must be 'dissipative' or 'none'")
def _solver_parameters(parameters: Any) -> dict[str, Any]:
if parameters is not None:
return dict(parameters)
return {
"snes_type": "ksponly",
"ksp_type": "preonly",
"pc_type": "lu",
}
[docs]
class CaputoWismerForwardOperator:
"""Linear initial-pressure to sensor-trace map and its discrete adjoint."""
def __init__(
self,
space: Any,
sensors: SensorArray,
*,
materials: Sequence[CaputoWismerMaterial],
dt: float,
num_steps: int,
num_modes: int = 32,
representation: Any = None,
absorbing_speed: Any = None,
attenuation: AttenuationMode = "dissipative",
stiffness_theta: float = 0.0,
solver_parameters: Any = None,
) -> None:
try:
import firedrake as fd
except ImportError as error:
raise RuntimeError(
"CaputoWismerForwardOperator requires Firedrake"
) from error
if sensors.space != space:
raise ValueError("sensors must belong to the forward space")
if space.ufl_element().family() != "Lagrange":
raise NotImplementedError(
"CaputoWismerForwardOperator supports continuous "
"Lagrange spaces only"
)
if space.value_shape != ():
raise NotImplementedError(
"CaputoWismerForwardOperator supports scalar fields only"
)
step_size = _nonnegative_real(dt, "dt")
if step_size == 0.0:
raise ValueError("dt must be positive")
_positive_integer(num_steps, "num_steps")
_positive_integer(num_modes, "num_modes")
try:
theta = float(stiffness_theta)
except (TypeError, ValueError) as error:
raise TypeError(
"stiffness_theta must be a real scalar"
) from error
if not isfinite(theta) or not 0.0 <= theta <= 1.0:
raise ValueError("stiffness_theta must be between 0 and 1")
self._fd = fd
self.space = space
self.sensors = sensors
self.dt = step_size
self.num_steps = num_steps
self.num_modes = num_modes
self.representation = (
BirkSong(num_modes) if representation is None else representation
)
self.materials = _material_model(materials, attenuation, fd)
self.attenuation = attenuation
self.absorbing_speed = absorbing_speed
self.stiffness_theta = theta
trial = fd.TrialFunction(space)
test = fd.TestFunction(space)
mass_form = fd.inner(trial, test) * fd.dx
stiffness_form = sum(
(
fd.inner(
material.indicator
* material.wave_speed**2
* fd.grad(trial),
fd.grad(test),
)
* fd.dx
for material in self.materials
),
0,
)
damping_forms = tuple(
fd.inner(
material.indicator
* material.damping
* fd.grad(trial),
fd.grad(test),
)
* fd.dx
for material in self.materials
)
boundary_form = (
0
if absorbing_speed is None
else absorbing_speed * fd.inner(trial, test) * fd.ds
)
spectra = tuple(
self.representation.spectrum(float(material.alpha))
for material in self.materials
)
recurrence = []
implicit_weights = []
for spectrum in spectra:
arguments = spectrum.rates * step_size
decay = np.asarray(exp_neg(arguments), dtype=np.float64)
interpolation = np.asarray(phi1(arguments), dtype=np.float64)
recurrence.append((decay, interpolation, spectrum.weights))
implicit_weights.append(
float(np.dot(spectrum.weights, interpolation))
)
self._recurrence = tuple(recurrence)
self._implicit_weights = tuple(implicit_weights)
left_form = (1.0 / step_size**2) * mass_form
left_form += theta * stiffness_form
left_form += sum(
(
weight * form
for weight, form in zip(
self._implicit_weights,
damping_forms,
strict=True,
)
),
0,
)
if absorbing_speed is not None:
left_form += (1.0 / step_size) * boundary_form
self._mass = fd.assemble(mass_form, mat_type="aij")
self._stiffness = fd.assemble(stiffness_form, mat_type="aij")
self._damping = tuple(
fd.assemble(form, mat_type="aij") for form in damping_forms
)
self._boundary = (
None
if absorbing_speed is None
else fd.assemble(boundary_form, mat_type="aij")
)
self._left = fd.assemble(left_form, mat_type="aij")
parameters = _solver_parameters(solver_parameters)
self._left_solver = fd.LinearSolver(
self._left,
solver_parameters=parameters,
)
self._mass_solver = fd.LinearSolver(
self._mass,
solver_parameters=parameters,
)
def _function(self, name: str) -> Any:
return self._fd.Function(self.space, name=name)
def _covector(self, name: str) -> Any:
return self._fd.Cofunction(self.space.dual(), name=name)
@staticmethod
def _axpy(target: Any, scale: float, source: Any) -> None:
with (
target.dat.vec as target_vector,
source.dat.vec_ro as source_vector,
):
target_vector.axpy(float(scale), source_vector)
@staticmethod
def _matrix_axpy(
matrix: Any,
source: Any,
target: Any,
scale: float,
*,
transpose: bool = False,
work: Any = None,
) -> None:
temporary = target.copy(deepcopy=True) if work is None else work
temporary.assign(0.0)
with (
source.dat.vec_ro as source_vector,
temporary.dat.vec as temporary_vector,
):
if transpose:
matrix.petscmat.multTranspose(
source_vector,
temporary_vector,
)
else:
matrix.petscmat.mult(source_vector, temporary_vector)
CaputoWismerForwardOperator._axpy(target, scale, temporary)
def _solve_left(
self,
right_hand_side: Any,
*,
transpose: bool = False,
out: Any = None,
) -> Any:
result = (
self._function("caputo_wismer_linear_solution")
if out is None
else out
)
if not transpose:
self._left_solver.solve(result, right_hand_side)
return result
with (
right_hand_side.dat.vec_ro as source_vector,
result.dat.vec as result_vector,
):
result_vector.set(0.0)
self._left_solver.ksp.solveTranspose(
source_vector,
result_vector,
)
return result
[docs]
def forward(self, initial_pressure: Any) -> np.ndarray:
"""Propagate one initial pressure and return all sensor samples."""
if initial_pressure.function_space() != self.space:
raise ValueError("initial_pressure must belong to the forward space")
current = initial_pressure.copy(deepcopy=True)
previous = initial_pressure.copy(deepcopy=True)
modes = tuple(
np.zeros(
(decay.size, current.dat.data_ro.size),
dtype=np.float64,
)
for decay, _, _ in self._recurrence
)
data = np.empty(
(self.num_steps + 1, self.sensors.num_sensors),
dtype=np.float64,
)
right_hand_side = self._covector("forward_right_hand_side")
matrix_work = self._covector("forward_matrix_work")
material_state = self._function("forward_material_state")
next_field = self._function("forward_next")
increment = self._function("forward_increment")
data[0] = self.sensors.sample(current)
inverse_dt_squared = 1.0 / self.dt**2
for step in range(self.num_steps):
right_hand_side.assign(0.0)
self._matrix_axpy(
self._mass,
current,
right_hand_side,
2.0 * inverse_dt_squared,
work=matrix_work,
)
self._matrix_axpy(
self._mass,
previous,
right_hand_side,
-inverse_dt_squared,
work=matrix_work,
)
self._matrix_axpy(
self._stiffness,
current,
right_hand_side,
-(1.0 - self.stiffness_theta),
work=matrix_work,
)
if self._boundary is not None:
self._matrix_axpy(
self._boundary,
current,
right_hand_side,
1.0 / self.dt,
work=matrix_work,
)
for (
damping,
material_mode_values,
coefficients,
implicit_weight,
) in zip(
self._damping,
modes,
self._recurrence,
self._implicit_weights,
strict=True,
):
decay, _, weights = coefficients
np.matmul(
-(decay * weights),
material_mode_values,
out=material_state.dat.data,
)
material_state.dat.data[:] += (
implicit_weight * current.dat.data_ro
)
self._matrix_axpy(
damping,
material_state,
right_hand_side,
1.0,
work=matrix_work,
)
self._solve_left(right_hand_side, out=next_field)
increment.assign(next_field)
increment -= current
for material_mode_values, coefficients in zip(
modes,
self._recurrence,
strict=True,
):
decay, interpolation, _ = coefficients
material_mode_values *= decay[:, None]
material_mode_values += (
interpolation[:, None]
* increment.dat.data_ro[None, :]
)
previous.assign(current)
current.assign(next_field)
data[step + 1] = self.sensors.sample(current)
return data
[docs]
def adjoint_covector(self, sensor_values: Any) -> Any:
"""Apply the exact transpose of ``forward`` to sensor-time values."""
values = np.asarray(sensor_values, dtype=np.float64)
expected_shape = (
self.num_steps + 1,
self.sensors.num_sensors,
)
if values.shape != expected_shape:
raise ValueError(f"sensor_values must have shape {expected_shape}")
if not np.all(np.isfinite(values)):
raise ValueError("sensor_values must be finite")
adjoint_current = self.sensors.adjoint_covector(values[-1])
adjoint_previous = self._covector("adjoint_previous")
adjoint_modes = tuple(
np.zeros(
(decay.size, adjoint_current.dat.data_ro.size),
dtype=np.float64,
)
for decay, _, _ in self._recurrence
)
current_contribution = self._covector(
"adjoint_current_contribution"
)
previous_contribution = self._covector(
"adjoint_previous_contribution"
)
multiplier = self._function("adjoint_multiplier")
matrix_work = self._covector("adjoint_matrix_work")
damping_action = self._covector("adjoint_damping_action")
observation = self._covector("adjoint_observation")
mode_contribution = np.empty(
adjoint_current.dat.data_ro.size,
dtype=np.float64,
)
inverse_dt_squared = 1.0 / self.dt**2
for step in range(self.num_steps - 1, -1, -1):
current_contribution.assign(adjoint_previous)
previous_contribution.assign(0.0)
for material_mode_values, coefficients in zip(
adjoint_modes,
self._recurrence,
strict=True,
):
decay, interpolation, _ = coefficients
np.matmul(
interpolation,
material_mode_values,
out=mode_contribution,
)
adjoint_current.dat.data[:] += mode_contribution
current_contribution.dat.data[:] -= mode_contribution
self._solve_left(
adjoint_current,
transpose=True,
out=multiplier,
)
self._matrix_axpy(
self._mass,
multiplier,
current_contribution,
2.0 * inverse_dt_squared,
transpose=True,
work=matrix_work,
)
self._matrix_axpy(
self._mass,
multiplier,
previous_contribution,
-inverse_dt_squared,
transpose=True,
work=matrix_work,
)
self._matrix_axpy(
self._stiffness,
multiplier,
current_contribution,
-(1.0 - self.stiffness_theta),
transpose=True,
work=matrix_work,
)
if self._boundary is not None:
self._matrix_axpy(
self._boundary,
multiplier,
current_contribution,
1.0 / self.dt,
transpose=True,
work=matrix_work,
)
for (
damping,
material_mode_values,
coefficients,
implicit_weight,
) in zip(
self._damping,
adjoint_modes,
self._recurrence,
self._implicit_weights,
strict=True,
):
decay, _, weights = coefficients
damping_action.assign(0.0)
with (
multiplier.dat.vec_ro as source_vector,
damping_action.dat.vec as target_vector,
):
damping.petscmat.multTranspose(
source_vector,
target_vector,
)
self._axpy(
current_contribution,
implicit_weight,
damping_action,
)
material_mode_values *= decay[:, None]
material_mode_values -= (
(decay * weights)[:, None]
* damping_action.dat.data_ro[None, :]
)
self.sensors.adjoint_covector(
values[step],
out=observation,
)
self._axpy(current_contribution, 1.0, observation)
adjoint_current, current_contribution = (
current_contribution,
adjoint_current,
)
adjoint_previous, previous_contribution = (
previous_contribution,
adjoint_previous,
)
result = self._covector("initial_pressure_adjoint")
result.assign(adjoint_current)
self._axpy(result, 1.0, adjoint_previous)
return result
[docs]
def adjoint(self, sensor_values: Any) -> Any:
"""Return the spatial L2 representative of the discrete adjoint."""
covector = self.adjoint_covector(sensor_values)
result = self._function("initial_pressure_adjoint")
self._mass_solver.solve(result, covector)
return result
[docs]
def record_sensor_data(
initial_pressure: Any,
sensors: SensorArray,
*,
materials: Sequence[CaputoWismerMaterial],
dt: float,
num_steps: int,
num_modes: int = 32,
representation: Any = None,
absorbing_speed: Any = None,
stiffness_theta: float = 0.0,
solver_parameters: Any = None,
) -> np.ndarray:
"""Propagate initial pressure and record a chosen sensor array."""
if sensors.space != initial_pressure.function_space():
raise ValueError(
"sensors must belong to the initial-pressure function space"
)
if (
not isinstance(num_steps, int)
or isinstance(num_steps, bool)
or num_steps < 1
):
raise ValueError("num_steps must be a positive integer")
operator = CaputoWismerForwardOperator(
initial_pressure.function_space(),
sensors,
materials=materials,
dt=dt,
num_steps=num_steps,
num_modes=num_modes,
representation=representation,
absorbing_speed=absorbing_speed,
stiffness_theta=stiffness_theta,
solver_parameters=solver_parameters,
)
return operator.forward(initial_pressure)
__all__ = ["CaputoWismerForwardOperator", "record_sensor_data"]