"""Explicit Euler step implementation.
This module provides :class:`ExplicitEulerStep`, a single-stage,
non-adaptive forward Euler integrator compiled as a CUDA device
function through the :class:`CUDAFactory` pattern.
Published Classes
-----------------
:class:`ExplicitEulerStep`
Forward Euler integration step.
>>> from numpy import float32
>>> step = ExplicitEulerStep(precision=float32, n_states=4)
>>> step.order
1
>>> step.has_error_estimate
False
See Also
--------
:class:`~cubie.integrators.algorithms.ode_explicitstep.ODEExplicitStep`
Abstract base class for explicit integration steps.
:class:`~cubie.integrators.algorithms.base_algorithm_step.BaseAlgorithmStep`
Factory base managing compilation and caching.
"""
from typing import Callable, Optional
from cubie.cuda_simsafe import cuda, int32
from cubie.cuda_simsafe import unroll_if
from cubie.result_codes import CUBIE_RESULT_CODES
from cubie._utils import PrecisionDType, build_config
from cubie.integrators.algorithms.base_algorithm_step import StepCache, \
AlgorithmDefaults
from cubie.integrators.algorithms.ode_explicitstep import (
ExplicitStepConfig,
ODEExplicitStep,
)
EE_DEFAULTS = AlgorithmDefaults(
settings={
"step_controller": "fixed",
}
)
[docs]
class ExplicitEulerStep(ODEExplicitStep):
"""Forward Euler integration step for explicit ODE updates."""
[docs]
@classmethod
def family_defaults(cls, tableau=None) -> AlgorithmDefaults:
"""Return the explicit Euler defaults."""
return EE_DEFAULTS.copy()
[docs]
def __init__(
self,
precision: PrecisionDType,
n_states: int,
dxdt_fn: Optional[Callable] = None,
observables_fn: Optional[Callable] = None,
drivers_fn: Optional[Callable] = None,
get_solver_helper_fn: Optional[Callable] = None,
**kwargs,
) -> None:
"""Initialise the explicit Euler step configuration.
Parameters
----------
precision
Precision applied to device buffers.
n_states
Number of state entries advanced per step.
dxdt_fn
Device function for evaluating f(t, y) right-hand side.
observables_fn
Device function computing system observables.
drivers_fn
Optional device function evaluating drivers at arbitrary times.
get_solver_helper_fn
Present for interface parity with implicit steps and ignored here.
**kwargs
Optional parameters passed to config classes. See
ExplicitStepConfig for available parameters. None values are
ignored.
"""
config = build_config(
ExplicitStepConfig,
required={
'precision': precision,
'n_states': n_states,
'dxdt_fn': dxdt_fn,
'observables_fn': observables_fn,
'drivers_fn': drivers_fn,
'get_solver_helper_fn': get_solver_helper_fn,
},
**kwargs
)
super().__init__(config, EE_DEFAULTS.copy())
[docs]
def build_step(
self,
dxdt_fn: Callable,
observables_fn: Callable,
drivers_fn: Optional[Callable],
numba_precision: type,
n: int,
n_drivers: int,
) -> StepCache:
"""Build the device function for an explicit Euler step.
Parameters
----------
dxdt_fn
Device function for evaluating f(t, y).
observables_fn
Device function for computing observables.
drivers_fn
Optional device function for evaluating drivers at time t.
numba_precision
Numba type for device buffers.
n
State vector dimension.
n_drivers
Number of driver signals.
Returns
-------
StepCache
Compiled step function.
"""
has_evaluate_driver_at_t = drivers_fn is not None
n = int32(n)
unroll_step_element = self.compile_settings.unroll.unroll_step_element
success = int32(CUBIE_RESULT_CODES.SUCCESS)
# no cover: start
@cuda.jit(
# (
# numba_precision[::1],
# numba_precision[::1],
# numba_precision[::1],
# numba_precision[:, :, ::1],
# numba_precision[::1],
# numba_precision[::1],
# numba_precision[::1],
# numba_precision[::1],
# numba_precision[::1],
# numba_precision,
# numba_precision,
# int32,
# int32,
# numba_precision[::1],
# numba_precision[::1],
# int32[::1],
# ),
device=True,
inline=True,
**self.jit_kwargs,
)
def step(
state,
proposed_state,
parameters,
driver_coefficients,
drivers_buffer,
proposed_drivers,
observables,
proposed_observables,
error, # Non-adaptive algorithms receive a zero-length slice.
dt_scalar,
time_scalar,
first_step_flag,
accepted_flag,
shared,
persistent_local,
counters,
):
"""Advance the state with a single explicit Euler update.
Parameters
----------
state
Device array storing the current state.
proposed_state
Device array receiving the updated state.
parameters
Device array of static model parameters.
driver_coefficients
Device array containing spline driver coefficients.
drivers_buffer
Device array of time-dependent drivers.
proposed_drivers
Device array receiving proposed driver samples.
observables
Device array storing accepted observable outputs.
proposed_observables
Device array receiving proposed observable outputs.
error
Device array reserved for error estimates. Non-adaptive
algorithms receive a zero-length slice that can be reused
as scratch.
dt_scalar
Scalar containing the proposed step size.
time_scalar
Scalar containing the current simulation time.
first_step_flag : int32
Non-zero on the first step of the integration.
accepted_flag : int32
Non-zero when the previous step was accepted.
shared
Device array providing shared scratch buffers.
persistent_local
Device array for persistent local storage (unused here).
counters : int32 array
Diagnostic counter array (unused here).
Returns
-------
int
Status code indicating successful completion.
"""
# error buffer unused; stage dx/dt in proposed_state instead.
dxdt_buffer = proposed_state
dxdt_fn(
state,
parameters,
drivers_buffer,
observables,
dxdt_buffer,
time_scalar,
)
for i in unroll_if(range(n), unroll_step_element):
proposed_state[i] = state[i] + dt_scalar * dxdt_buffer[i]
next_time = time_scalar + dt_scalar
if has_evaluate_driver_at_t:
drivers_fn(
next_time,
driver_coefficients,
proposed_drivers,
)
observables_fn(
proposed_state,
parameters,
proposed_drivers,
proposed_observables,
next_time,
)
return success
# no cover: end
return StepCache(step_fn=step, nonlinear_solver_fn=None)
@property
def threads_per_step(self) -> int:
"""Return the number of threads used per step."""
return 1
@property
def is_multistage(self) -> bool:
"""Return ``False`` because explicit Euler is a single-stage method."""
return False
# Class attribute so alias queries need no instance.
has_error_estimate = False
@property
def order(self) -> int:
"""Return the classical order of the explicit Euler method."""
return 1