BifurcationInference.jl

This library implements the method described in Szep, G. Dalchau, N. and Csikasz-Nagy, A. 2021. Parameter Inference with Bifurcation Diagrams using parameter continuation library BifurcationKit.jl and auto-differentiation ForwardDiff.jl. This implementation enables continuation methods can be used as layers in machine learning proceedures, and inference can be run end-to-end directly on the geometry of state space.

The model definition requires a distpatched method on F(z::BorderedArray,θ::AbstractVector) where BorderedArray is a type that contains the state vector u and control condition p used by the library BifurcationKit.jl. θ is a vector of parameters to be optimised.

using BifurcationInference, StaticArrays

F(z::BorderedArray,θ::AbstractVector) = F(z.u,(θ=θ,p=z.p))
function F(u::AbstractVector,parameters::NamedTuple)

	@unpack θ,p = parameters
	μ₁,μ₂, a₁,a₂, k = θ

	f = first(u)*first(p)*first(θ)
	F = similar(u,typeof(f))

	F[1] = ( 10^a₁ + (p*u[2])^2 ) / ( 1 + (p*u[2])^2 ) - u[1]*10^μ₁
	F[2] = ( 10^a₂ + (k*u[1])^2 ) / ( 1 + (k*u[1])^2 ) - u[2]*10^μ₂

	return F
end

The targets are specified with StateSpace( dimension::Integer, condition::AbstractRange, targets::AbstractVector ). It contains the dimension of the state space, which must match that of the defined model, the control condition range that we would like to perform the continuation method in, and a vector of target conditions we would like to match.

X = StateSpace( 2, 0.01:0.01:10, [4,5] )

The optimisation needs to be initialised using a NamedTuple containing the initial guess for θ and the initial value p from which to begin the continuation.

using Flux: Optimise
parameters = ( θ=SizedVector{5}(0.5,0.5,0.5470,2.0,7.5), p=minimum(X.parameter) )
train!( F, parameters, X;  iter=200, optimiser=Optimise.ADAM(0.01) )