diabayes.solver.ODESolver#

class diabayes.solver.ODESolver(forward_model: Forward, rtol: float = 1e-07, atol: float = 1e-10, checkpoints: int = 100)[source]#

The main solver class that contains forward and inverse modelling methods.

__init__(forward_model: Forward, rtol: float = 1e-07, atol: float = 1e-10, checkpoints: int = 100) → None[source]#

Methods

__init__(forward_model[, rtol, atol, ...])

bayesian_inversion(t, mu, noise_std, y0, ...)

A Bayesian inversion routine using the Stein Variational Inference method.

evaluate_at_t(sol, t_eval)

generate_sequence(t_steps, v_steps, dt, y0, ...)

Solve the forward problem for a sequence of velocity steps.

max_likelihood_inversion(t, mu, y0, params, ...)

Minimises the least-squares residuals between the observed friction curve and the parameterised one, using the Levenberg-Marquardt algorithm.

solve_forward(t, y0, params, ...[, method, ...])

Solve a forward problem using SciPy's solve_ivp routine.

Attributes

learning_rate

The initial learning rate provided to the Adam algorithm for the Stein Variational Inference.

forward_model

An instantiated diabayes.Forward class, including friction law, state evolution equations, and stress transfer equation

rtol

The absolute tolerance used by the ODE solver

atol

The absolute tolerance used by the ODE solver

checkpoints

The number of checkpoints to use to compute the (adjoint) gradients through the ODE routine.
bayesian_inversion(t: Float[Array, 'Nt'], mu: Float[Array, 'Nt'], noise_std: Float, y0: Variables, params: RSFParams | CNSParams, friction_constants: RSFConstants | CNSConstants, block_constants: SpringBlockConstants | InertialSpringBlockConstants, Nparticles: int = 1000, Nsteps: int = 150, rng: None | int | Array = None) → BayesianSolution[source]#

A Bayesian inversion routine using the Stein Variational Inference method.

Parameters:

  • t (Array) – A vector or time values (in units of seconds). The time steps do not need to be uniform

  • mu (Array) – The observed friction curve sampled at t

  • noise_std (Float) – An estimate of the standard deviation of the noise in the measured friction curve. A conservative value is recommended, i.e. if the noise has a standard deviation of $10^{-3}$, a good starting point would be to set noise_std = 0.5e-3

  • y0 (Variables) – The initial values for the modelled friction and any state variables

  • params (_Params) – The initial guess for the invertible parameters that characterise the forward problem. It is recommended to use the result from ODESolver.max_likelihood_inversion. The prior distribution will be centered around this initial guess

  • friction_constants (_Constants) – The non-invertible constants that characterise the forward problem

  • block_constants (_BlockConstants) – The stress transfer constants (e.g. stiffness and loading rate)

  • Nparticles (int) – The number of particles (= posterior samples) to include. A higher value gives a more accurate estimation of the posterior distribution, at a higher computational cost.

  • Nsteps (int) – The number of iterations before convergence is expected to be achieved. It is recommended to start with a value of 100 and then see if an equilibrium was actually achieved. Increasing this value beyond the point of equilibrium does not do anything.

  • rng (None, int, jax.random.PRNGKey) – The random seed used to initialise the particle swarm distribution. If None, the current time will be used as a seed, which leads to different result for each realisation. When an integer is provided, a new jax.random.PRNGKey is generated.

Returns:

result – The result of the inversion incapsulated in a BayesianSolution container, which provides access to the full convergence chains, as well as diagnostic and visualisation routines.

Return type:

diabayes.BayesianSolution

Notes

This method is currently only compatible with standard rate-and-state friction models.

evaluate_at_t(sol: Solution, t_eval: Float[Array, 'Nt']) → Variables[source]#
generate_sequence(t_steps: Float[Array, 'Nsteps'], v_steps: Float[Array, 'Nsteps'], dt: Float, y0: Variables, params: RSFParams | CNSParams, friction_constants: RSFConstants | CNSConstants, block_constants: SpringBlockConstants | InertialSpringBlockConstants, method: str = 'RK45') → Tuple[Variables, Float[Array, 'Nt']][source]#

Solve the forward problem for a sequence of velocity steps. For each load-point velocity in v_steps, a forward simulation is run using the previous step’s final state as the initial state.

This routine can be used to generate a sequence of (up/down) velocity steps, or a slide-hold-slide sequence (by setting a given v_step to zero).

Parameters:

  • t_steps (Float[Array, "Nsteps"]) – A vector of change point times of each step in the sequence

  • v_steps (Float[Array, "Nsteps"]) – A vector of the load-point velocity values for each step in the sequence

  • dt (Float) – The desired time sample spacing of the solution

  • y0 (Variables) – The initial values (fricton and state) wrapped in a Variables container.

  • params (_Params) – The (invertible) parameters that govern the dynamics, wrapped in a Params container.

  • friction_constants (_Constants) – A container object containing the friction constants

  • block_constants (_BlockConstants) – A container object containing the block constants. Note that the load-point velocity will be updated for each step

Returns:

  • result (Variables) – Solution time series of friction and state

  • t (Float[Array, “Nt”]) – Solution time samples

max_likelihood_inversion(t: Float[Array, 'Nt'], mu: Float[Array, 'Nt'], y0: Variables, params: RSFParams | CNSParams, friction_constants: RSFConstants | CNSConstants, block_constants: SpringBlockConstants | InertialSpringBlockConstants, verbose: bool = False, retries: int = 3, seed: int = 42) → Solution | None[source]#

Minimises the least-squares residuals between the observed friction curve and the parameterised one, using the Levenberg-Marquardt algorithm.

Parameters:

  • t (Array) – A vector or time values (in units of seconds). The time steps do not need to be uniform

  • mu (Array) – The observed friction curve sampled at t

  • y0 (Variables) – The initial values for the modelled friction and any state variables

  • params (_Params) – The initial guess for the invertible parameters that characterise the forward problem. These need to be sufficiently close to the “true” values for the algorithm to converge

  • friction_constants (_Constants) – The non-invertible constants that characterise the forward problem

  • block_constants (_BlockConstants) – The stress transfer constants (e.g. stiffness and loading rate)

  • verbose (bool) – Whether or not to output detailed progress of the inversion. Defaults to False

  • retries (int) – The maximum number of inversion attempts. When the inversion fails to converge, it will retry up to retries times with randomly perturbed initial parameters.

  • seed (int) – Seed for the random number generator. This is only used when the initial inversion attempt fails, and the initial parameters are randomly perturbed before the next attempt.

Returns:

sol – The inversion result, including various diagnostics. The inverted parameter values can be accessed as sol.values

Return type:

optimistix.Solution

solve_forward(t: Float[Array, 'Nt'], y0: Variables, params: RSFParams | CNSParams, friction_constants: RSFConstants | CNSConstants, block_constants: SpringBlockConstants | InertialSpringBlockConstants, method: str = 'RK45', interpolate_time: bool = True) → Variables[source]#

Solve a forward problem using SciPy’s solve_ivp routine. While this routine doesn’t propagate any gradients, it is much faster to initialise and to perform a single forward run. Hence for playing around with different parameters, it is preferred over a JITed JAX implementation.

Parameters:

  • t (Float[Array, "Nt"]) – A vector of time samples where a solution is requested.

  • y0 (Variables) – The initial values (fricton and state) wrapped in a Variables container.

  • params (_Params) – The (invertible) parameters that govern the dynamics, wrapped in a Params container.

  • friction_constants (_Constants) – A container object containing the friction constants

  • block_constants (_BlockConstants) – A container object containing the block constants

  • method (str) – The solver used by SciPy’s solve_ivp. Default: RK45

  • interpolate_time (bool) – Whether to interpolate the result to the user-provided time samples (True), or to use the samples from the adaptive ODE solver (False). Default: True

Returns:

result – Solution time series of friction and state

Return type:

Variables

atol: float#

The absolute tolerance used by the ODE solver

checkpoints: int#

The number of checkpoints to use to compute the (adjoint) gradients through the ODE routine. A higher number increases stability and speed, at the expense of more GPU memory

forward_model: Forward#

An instantiated diabayes.Forward class, including friction law, state evolution equations, and stress transfer equation

learning_rate: float = 0.01#

The initial learning rate provided to the Adam algorithm for the Stein Variational Inference. The default value of 1e-2 seems like a sensible choice for most models.

rtol: float#

The absolute tolerance used by the ODE solver