562 lines
19 KiB
Python
562 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import Callable, Literal, Sequence
|
|
|
|
|
|
CancellationCheck = Callable[[], bool]
|
|
AcceptedStepCallback = Callable[[float], None]
|
|
IntegrationStatus = Literal["completed", "cancelled", "failed"]
|
|
|
|
|
|
class _IntegrationCancelled(Exception):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SolveIVPConfig:
|
|
t_start: float = 0.0
|
|
t_stop: float = 20.0
|
|
method: str = "BDF"
|
|
rtol: float = 1e-6
|
|
atol: float = 1e-8
|
|
max_step: float = 1e-3
|
|
first_step: float | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ODESolution:
|
|
t: list[float]
|
|
y: list[list[float]]
|
|
success: bool
|
|
message: str
|
|
status: IntegrationStatus = "completed"
|
|
error: Exception | None = None
|
|
|
|
|
|
def _vector_add(a: list[float], b: list[float], scale: float = 1.0) -> list[float]:
|
|
return [x + scale * y for x, y in zip(a, b)]
|
|
|
|
|
|
def _append_solution_sample(
|
|
times: list[float],
|
|
states: list[list[float]],
|
|
time: float,
|
|
state: list[float],
|
|
) -> None:
|
|
if times and time <= times[-1] + 1e-12:
|
|
return
|
|
times.append(float(time))
|
|
for index, value in enumerate(state):
|
|
states[index].append(float(value))
|
|
|
|
|
|
def _normalize_breakpoints(
|
|
config: SolveIVPConfig,
|
|
breakpoints: Sequence[float] | None,
|
|
) -> list[float]:
|
|
"""Return sorted, unique breakpoints strictly inside the integration span."""
|
|
|
|
if breakpoints is None or len(breakpoints) == 0:
|
|
return []
|
|
if config.t_stop < config.t_start:
|
|
raise ValueError("Segmented integration requires t_stop to follow t_start.")
|
|
|
|
normalized: list[float] = []
|
|
for raw_breakpoint in breakpoints:
|
|
breakpoint = float(raw_breakpoint)
|
|
if not math.isfinite(breakpoint):
|
|
raise ValueError("Integration breakpoints must be finite numbers.")
|
|
if config.t_start < breakpoint < config.t_stop:
|
|
normalized.append(breakpoint)
|
|
|
|
normalized.sort()
|
|
return [
|
|
breakpoint
|
|
for index, breakpoint in enumerate(normalized)
|
|
if index == 0 or breakpoint != normalized[index - 1]
|
|
]
|
|
|
|
|
|
def _runge_kutta_4(
|
|
rhs: Callable[[float, list[float]], list[float]],
|
|
initial_state: list[float],
|
|
config: SolveIVPConfig,
|
|
t_eval: list[float] | None,
|
|
cancel_check: CancellationCheck | None = None,
|
|
accepted_step_callback: AcceptedStepCallback | None = None,
|
|
) -> ODESolution:
|
|
if t_eval is None:
|
|
point_count = max(
|
|
2,
|
|
int((config.t_stop - config.t_start) / max(config.max_step, 1e-6)) + 1,
|
|
)
|
|
step = (config.t_stop - config.t_start) / (point_count - 1)
|
|
t_eval = [config.t_start + index * step for index in range(point_count)]
|
|
|
|
state = list(initial_state)
|
|
states = [[value] for value in state]
|
|
times = [float(t_eval[0])]
|
|
current_time = float(t_eval[0])
|
|
status: IntegrationStatus = "completed"
|
|
message = "Integrated with built-in RK4 fallback because SciPy is unavailable."
|
|
error: Exception | None = None
|
|
|
|
try:
|
|
for target_time in t_eval[1:]:
|
|
while current_time < target_time - 1e-15:
|
|
if cancel_check is not None and cancel_check():
|
|
raise _IntegrationCancelled
|
|
dt = min(config.max_step, target_time - current_time)
|
|
k1 = rhs(current_time, state)
|
|
k2 = rhs(current_time + 0.5 * dt, _vector_add(state, k1, 0.5 * dt))
|
|
k3 = rhs(current_time + 0.5 * dt, _vector_add(state, k2, 0.5 * dt))
|
|
k4 = rhs(current_time + dt, _vector_add(state, k3, dt))
|
|
state = [
|
|
value + (dt / 6.0) * (a + 2.0 * b + 2.0 * c + d)
|
|
for value, a, b, c, d in zip(state, k1, k2, k3, k4)
|
|
]
|
|
current_time += dt
|
|
if accepted_step_callback is not None:
|
|
accepted_step_callback(current_time)
|
|
|
|
_append_solution_sample(times, states, target_time, state)
|
|
except _IntegrationCancelled:
|
|
status = "cancelled"
|
|
message = "Simulation was stopped before reaching the requested end time."
|
|
_append_solution_sample(times, states, current_time, state)
|
|
except Exception as exc:
|
|
status = "failed"
|
|
message = str(exc)
|
|
error = exc
|
|
_append_solution_sample(times, states, current_time, state)
|
|
|
|
return ODESolution(
|
|
t=times,
|
|
y=states,
|
|
success=status == "completed",
|
|
message=message,
|
|
status=status,
|
|
error=error,
|
|
)
|
|
|
|
|
|
def _runge_kutta_4_segmented(
|
|
rhs: Callable[[float, list[float]], list[float]],
|
|
initial_state: list[float],
|
|
config: SolveIVPConfig,
|
|
t_eval: list[float] | None,
|
|
breakpoints: Sequence[float],
|
|
cancel_check: CancellationCheck | None = None,
|
|
accepted_step_callback: AcceptedStepCallback | None = None,
|
|
) -> ODESolution:
|
|
"""RK4 fallback that never evaluates a pre-breakpoint step at the breakpoint."""
|
|
|
|
if t_eval is None:
|
|
point_count = max(
|
|
2,
|
|
int((config.t_stop - config.t_start) / max(config.max_step, 1e-6)) + 1,
|
|
)
|
|
sample_step = (config.t_stop - config.t_start) / (point_count - 1)
|
|
sample_times = [
|
|
config.t_start + index * sample_step for index in range(point_count)
|
|
]
|
|
else:
|
|
sample_times = [float(time) for time in t_eval]
|
|
|
|
state = [float(value) for value in initial_state]
|
|
states = [[value] for value in state]
|
|
times = [float(config.t_start)]
|
|
current_time = float(config.t_start)
|
|
sample_index = 0
|
|
while (
|
|
sample_index < len(sample_times)
|
|
and sample_times[sample_index] <= config.t_start + 1e-12
|
|
):
|
|
sample_index += 1
|
|
|
|
status: IntegrationStatus = "completed"
|
|
message = "Integrated with built-in RK4 fallback because SciPy is unavailable."
|
|
error: Exception | None = None
|
|
last_reported_step: float | None = None
|
|
|
|
def report_step(time: float) -> None:
|
|
nonlocal last_reported_step
|
|
if accepted_step_callback is None:
|
|
return
|
|
if last_reported_step is not None and time <= last_reported_step:
|
|
return
|
|
accepted_step_callback(float(time))
|
|
last_reported_step = float(time)
|
|
|
|
def advance_to(
|
|
target_time: float, reported_terminal_time: float | None = None
|
|
) -> None:
|
|
nonlocal current_time, state
|
|
while current_time < target_time - 1e-15:
|
|
if cancel_check is not None and cancel_check():
|
|
raise _IntegrationCancelled
|
|
dt = min(config.max_step, target_time - current_time)
|
|
k1 = rhs(current_time, state)
|
|
k2 = rhs(
|
|
current_time + 0.5 * dt,
|
|
_vector_add(state, k1, 0.5 * dt),
|
|
)
|
|
k3 = rhs(
|
|
current_time + 0.5 * dt,
|
|
_vector_add(state, k2, 0.5 * dt),
|
|
)
|
|
k4 = rhs(current_time + dt, _vector_add(state, k3, dt))
|
|
state = [
|
|
value + (dt / 6.0) * (a + 2.0 * b + 2.0 * c + d)
|
|
for value, a, b, c, d in zip(state, k1, k2, k3, k4)
|
|
]
|
|
current_time += dt
|
|
report_time = current_time
|
|
if (
|
|
reported_terminal_time is not None
|
|
and current_time >= target_time - 1e-15
|
|
):
|
|
report_time = reported_terminal_time
|
|
report_step(report_time)
|
|
|
|
try:
|
|
segment_ends = [*breakpoints, float(config.t_stop)]
|
|
for segment_index, segment_end in enumerate(segment_ends):
|
|
is_breakpoint = segment_index < len(breakpoints)
|
|
integration_end = (
|
|
math.nextafter(segment_end, -math.inf) if is_breakpoint else segment_end
|
|
)
|
|
|
|
while (
|
|
sample_index < len(sample_times)
|
|
and sample_times[sample_index] <= integration_end
|
|
):
|
|
sample_time = float(sample_times[sample_index])
|
|
advance_to(sample_time)
|
|
_append_solution_sample(times, states, sample_time, state)
|
|
sample_index += 1
|
|
|
|
advance_to(
|
|
integration_end,
|
|
segment_end if is_breakpoint else None,
|
|
)
|
|
|
|
if is_breakpoint:
|
|
current_time = float(segment_end)
|
|
report_step(current_time)
|
|
while (
|
|
sample_index < len(sample_times)
|
|
and sample_times[sample_index] <= segment_end
|
|
):
|
|
sample_time = float(sample_times[sample_index])
|
|
_append_solution_sample(times, states, sample_time, state)
|
|
sample_index += 1
|
|
except _IntegrationCancelled:
|
|
status = "cancelled"
|
|
message = "Simulation was stopped before reaching the requested end time."
|
|
_append_solution_sample(times, states, current_time, state)
|
|
except Exception as exc:
|
|
status = "failed"
|
|
message = str(exc)
|
|
error = exc
|
|
_append_solution_sample(times, states, current_time, state)
|
|
|
|
return ODESolution(
|
|
t=times,
|
|
y=states,
|
|
success=status == "completed",
|
|
message=message,
|
|
status=status,
|
|
error=error,
|
|
)
|
|
|
|
|
|
def _integrate_scipy_stepwise(
|
|
rhs: Callable[[float, list[float]], list[float]],
|
|
initial_state: list[float],
|
|
config: SolveIVPConfig,
|
|
t_eval: list[float] | None,
|
|
cancel_check: CancellationCheck,
|
|
accepted_step_callback: AcceptedStepCallback | None,
|
|
breakpoints: Sequence[float] = (),
|
|
) -> ODESolution:
|
|
import numpy as np
|
|
from scipy.integrate import BDF, DOP853, LSODA, RK23, RK45, Radau
|
|
|
|
solver_types = {
|
|
"BDF": BDF,
|
|
"DOP853": DOP853,
|
|
"LSODA": LSODA,
|
|
"RK23": RK23,
|
|
"RK45": RK45,
|
|
"Radau": Radau,
|
|
}
|
|
solver_type = solver_types.get(config.method)
|
|
if solver_type is None:
|
|
raise ValueError(f"Unsupported integration method: {config.method}")
|
|
|
|
times = [float(config.t_start)]
|
|
states = [[float(value)] for value in initial_state]
|
|
last_accepted_time = float(config.t_start)
|
|
last_accepted_state = [float(value) for value in initial_state]
|
|
sample_times = [float(time) for time in (t_eval or [])]
|
|
sample_index = 0
|
|
while (
|
|
sample_index < len(sample_times)
|
|
and sample_times[sample_index] <= config.t_start + 1e-12
|
|
):
|
|
sample_index += 1
|
|
|
|
def cancellable_rhs(time, state):
|
|
if cancel_check():
|
|
raise _IntegrationCancelled
|
|
return rhs(float(time), [float(value) for value in state])
|
|
|
|
status: IntegrationStatus = "completed"
|
|
message = "The solver successfully reached the end of the integration interval."
|
|
error: Exception | None = None
|
|
last_reported_step: float | None = None
|
|
|
|
def report_step(time: float) -> None:
|
|
nonlocal last_reported_step
|
|
if accepted_step_callback is None:
|
|
return
|
|
if last_reported_step is not None and time <= last_reported_step:
|
|
return
|
|
accepted_step_callback(float(time))
|
|
last_reported_step = float(time)
|
|
|
|
segment_ends = [*breakpoints, float(config.t_stop)]
|
|
for segment_index, segment_end in enumerate(segment_ends):
|
|
if cancel_check():
|
|
status = "cancelled"
|
|
message = (
|
|
"Simulation was stopped before integration started."
|
|
if segment_index == 0
|
|
else "Simulation was stopped before reaching the requested end time."
|
|
)
|
|
break
|
|
|
|
is_breakpoint = segment_index < len(breakpoints)
|
|
integration_end = (
|
|
math.nextafter(segment_end, -math.inf) if is_breakpoint else segment_end
|
|
)
|
|
has_integration_interval = integration_end > last_accepted_time
|
|
|
|
if has_integration_interval:
|
|
solver_options = {
|
|
"rtol": config.rtol,
|
|
"atol": config.atol,
|
|
"max_step": config.max_step,
|
|
}
|
|
if config.first_step is not None:
|
|
solver_options["first_step"] = min(
|
|
config.first_step,
|
|
integration_end - last_accepted_time,
|
|
)
|
|
|
|
try:
|
|
solver = solver_type(
|
|
cancellable_rhs,
|
|
last_accepted_time,
|
|
np.asarray(last_accepted_state, dtype=float),
|
|
integration_end,
|
|
**solver_options,
|
|
)
|
|
except _IntegrationCancelled:
|
|
status = "cancelled"
|
|
message = (
|
|
"Simulation was stopped before integration started."
|
|
if segment_index == 0
|
|
else "Simulation was stopped before reaching the requested end time."
|
|
)
|
|
break
|
|
except Exception as exc:
|
|
status = "failed"
|
|
message = str(exc)
|
|
error = exc
|
|
break
|
|
|
|
while solver.status == "running":
|
|
if cancel_check():
|
|
status = "cancelled"
|
|
message = (
|
|
"Simulation was stopped before reaching the requested end time."
|
|
)
|
|
break
|
|
try:
|
|
step_message = solver.step()
|
|
except _IntegrationCancelled:
|
|
status = "cancelled"
|
|
message = (
|
|
"Simulation was stopped before reaching the requested end time."
|
|
)
|
|
break
|
|
except Exception as exc:
|
|
status = "failed"
|
|
message = str(exc)
|
|
error = exc
|
|
break
|
|
|
|
if solver.status == "failed":
|
|
status = "failed"
|
|
message = str(step_message or "Integration step failed.")
|
|
break
|
|
|
|
last_accepted_time = float(solver.t)
|
|
last_accepted_state = [float(value) for value in solver.y]
|
|
reported_time = (
|
|
float(segment_end)
|
|
if is_breakpoint and solver.status == "finished"
|
|
else last_accepted_time
|
|
)
|
|
if sample_times:
|
|
dense_output = solver.dense_output()
|
|
while (
|
|
sample_index < len(sample_times)
|
|
and sample_times[sample_index] <= last_accepted_time
|
|
):
|
|
sample_time = float(sample_times[sample_index])
|
|
sample_state = [
|
|
float(value) for value in dense_output(sample_time)
|
|
]
|
|
_append_solution_sample(
|
|
times,
|
|
states,
|
|
sample_time,
|
|
sample_state,
|
|
)
|
|
sample_index += 1
|
|
else:
|
|
_append_solution_sample(
|
|
times,
|
|
states,
|
|
reported_time,
|
|
last_accepted_state,
|
|
)
|
|
report_step(reported_time)
|
|
|
|
if status != "completed":
|
|
break
|
|
|
|
if is_breakpoint:
|
|
# The old equation is integrated only to the representable point just
|
|
# left of the event. The continuous state is then lifted to the exact
|
|
# event time, where the freshly constructed next solver sees the new
|
|
# equation immediately.
|
|
last_accepted_time = float(segment_end)
|
|
if sample_times:
|
|
while (
|
|
sample_index < len(sample_times)
|
|
and sample_times[sample_index] <= segment_end
|
|
):
|
|
sample_time = float(sample_times[sample_index])
|
|
_append_solution_sample(
|
|
times,
|
|
states,
|
|
sample_time,
|
|
last_accepted_state,
|
|
)
|
|
sample_index += 1
|
|
elif not has_integration_interval:
|
|
_append_solution_sample(
|
|
times,
|
|
states,
|
|
last_accepted_time,
|
|
last_accepted_state,
|
|
)
|
|
report_step(last_accepted_time)
|
|
|
|
if status != "completed":
|
|
_append_solution_sample(
|
|
times,
|
|
states,
|
|
last_accepted_time,
|
|
last_accepted_state,
|
|
)
|
|
|
|
return ODESolution(
|
|
t=times,
|
|
y=states,
|
|
success=status == "completed",
|
|
message=message,
|
|
status=status,
|
|
error=error,
|
|
)
|
|
|
|
|
|
def integrate_ode(
|
|
rhs: Callable[[float, list[float]], list[float]],
|
|
initial_state: list[float],
|
|
config: SolveIVPConfig,
|
|
t_eval: list[float] | None = None,
|
|
cancel_check: CancellationCheck | None = None,
|
|
accepted_step_callback: AcceptedStepCallback | None = None,
|
|
breakpoints: Sequence[float] | None = None,
|
|
):
|
|
"""Integrate an ODE, optionally restarting at equation discontinuities.
|
|
|
|
Breakpoints are interpreted as right-continuous equation changes: the old
|
|
equation is integrated to the floating-point left limit, then a fresh solver
|
|
starts at the exact breakpoint with the unchanged continuous state.
|
|
"""
|
|
|
|
if abs(config.t_stop - config.t_start) <= 1e-15:
|
|
return ODESolution(
|
|
t=[float(config.t_start)],
|
|
y=[[value] for value in initial_state],
|
|
success=True,
|
|
message="Skipped integration because t_start equals t_stop.",
|
|
)
|
|
|
|
normalized_breakpoints = _normalize_breakpoints(config, breakpoints)
|
|
|
|
try:
|
|
from scipy.integrate import solve_ivp
|
|
except ImportError:
|
|
if normalized_breakpoints:
|
|
return _runge_kutta_4_segmented(
|
|
rhs,
|
|
initial_state,
|
|
config,
|
|
t_eval,
|
|
normalized_breakpoints,
|
|
cancel_check,
|
|
accepted_step_callback,
|
|
)
|
|
return _runge_kutta_4(
|
|
rhs,
|
|
initial_state,
|
|
config,
|
|
t_eval,
|
|
cancel_check,
|
|
accepted_step_callback,
|
|
)
|
|
|
|
if cancel_check is not None or normalized_breakpoints:
|
|
return _integrate_scipy_stepwise(
|
|
rhs,
|
|
initial_state,
|
|
config,
|
|
t_eval,
|
|
cancel_check or (lambda: False),
|
|
accepted_step_callback,
|
|
normalized_breakpoints,
|
|
)
|
|
|
|
solve_options = {
|
|
"fun": rhs,
|
|
"t_span": (config.t_start, config.t_stop),
|
|
"y0": initial_state,
|
|
"method": config.method,
|
|
"rtol": config.rtol,
|
|
"atol": config.atol,
|
|
"max_step": config.max_step,
|
|
"t_eval": t_eval,
|
|
}
|
|
if config.first_step is not None:
|
|
solve_options["first_step"] = config.first_step
|
|
return solve_ivp(**solve_options)
|