def forward_func(z, x, v, δ):
δω = snp.sqrt(δ) * v
a = drift_func(x, z)
B = diff_coeff(x, z)
if noise_type == "diagonal":
B_dB_dx = snp.array(
[B[i, i] * B[i, i].diff(x[i]) for i in range(v.shape[0])]
)
else:
B_dB_dx = snp.array(
[(B * B[i].diff(x)).sum() for i in range(x.shape[0])]
)
x_ = x + δ * a + B @ δω + B_dB_dx * (δω ** 2 - δ) / 2
return snp.array([sympy.simplify(x_[i]) for i in range(x.shape[0])])
def transformed_diff_coeff(y, z):
x = symnum.named_array("x", y.shape)
B = diff_coeff(x, z)
x_y = backward_func(y)
B_y = (diffops.jacobian(forward_func)(x) @ B).subs([
(x_i, x_y_i) for x_i, x_y_i in zip(x.flatten(), x_y.flatten())
])
return snp.array(
[[sympy.simplify(B_y[i, j]) for j in range(B_y.shape[1])]
for i in range(B_y.shape[0])])
def diff_coeff(x, z):
α = snp.exp(x[2])
β, γ, ζ, ϵ = z
return snp.array(
[
[snp.sqrt(α * x[0] * x[1] / N), 0, 0],
[-snp.sqrt(α * x[0] * x[1] / N), snp.sqrt(β * x[1]), 0],
[0, 0, ϵ],
]
)
def transformed_drift_func(y, z):
x = symnum.named_array("x", y.shape)
a = drift_func(x, z)
B = diff_coeff(x, z)
x_y = backward_func(y)
a_y = (diffops.jacobian_vector_product(forward_func)(x)(a) +
diffops.matrix_hessian_product(forward_func)(x)(B @ B.T) /
2).subs([(x_i, x_y_i)
for x_i, x_y_i in zip(x.flatten(), x_y.flatten())])
return snp.array(
[sympy.simplify(a_y[i]) for i in range(a_y.shape[0])])
def forward_func(z, x, v, δ):
dim_noise = v.shape[0] // 2
δω = snp.sqrt(δ) * v[:dim_noise]
δζ = δ * snp.sqrt(δ) * (v[:dim_noise] + v[dim_noise:] / snp.sqrt(3)) / 2
x_ = (
x
+ δ * drift_func(x, z)
+ diff_coeff(x, z) @ δω
+ (δ ** 2 / 2)
* diffusion_operator(drift_func, diff_coeff)(drift_func)(x, z)
+ sum(
Lj_operator(diff_coeff, j)(drift_func)(x, z) * δζ[j]
for j in range(dim_noise)
)
)
return snp.array([sympy.simplify(x_[i]) for i in range(x.shape[0])])
def forward_func(z, x, v, δ):
δω = snp.sqrt(δ) * v[:1]
δζ = δ * snp.sqrt(δ) * (v[:1] + v[1:] / snp.sqrt(3)) / 2
x_ = (
x
+ δ * drift_func(x, z)
+ diff_coeff(x, z) @ δω
+ Lj_operator(diff_coeff, 0)(diff_coeff)(x, z) @ (δω ** 2 - δ) / 2
+ Lj_operator(diff_coeff, 0)(drift_func)(x, z) * δζ
+ diffusion_operator(drift_func, diff_coeff)(
lambda x, z: diff_coeff(x, z)[:, 0]
)(x, z)
* (δω * δ - δζ)
+ (δ ** 2 / 2)
* diffusion_operator(drift_func, diff_coeff)(drift_func)(x, z)
+ Lj_operator(diff_coeff, 0)(Lj_operator(diff_coeff, 0)(diff_coeff))(
x, z
)
@ (δω ** 3 / 3 - δ * δω)
)
return snp.array([sympy.simplify(x_[i]) for i in range(x.shape[0])])
def diff_coeff(x, z):
σ, ε, γ, β = z
return snp.array([[0], [σ]])
def drift_func(x, z):
σ, ε, γ, β = z
return snp.array([(x[0] - x[0] ** 3 - x[1]) / ε, γ * x[0] - x[1] + β])
def drift_func(x, z):
α = snp.exp(x[2])
β, γ, ζ, ϵ = z
return snp.array(
[-α * x[0] * x[1] / N, α * x[0] * x[1] / N - β * x[1], γ * (ζ - x[2])]
)
def diff_coeff(x, z):
α = snp.exp(x[2])
β, γ, ζ, ϵ = z
return snp.array(
[
[snp.sqrt(α * x[0] * x[1] / N), 0, 0],
[-snp.sqrt(α * x[0] * x[1] / N), snp.sqrt(β * x[1]), 0],
[0, 0, ϵ],
]
)
_forward_func = symnum.numpify_func(
sde.integrators.euler_maruyama_step(
*sde.transforms.transform_sde(
lambda x: snp.array([snp.log(x[0]), snp.log(x[1]), x[2]]),
lambda x: snp.array([snp.exp(x[0]), snp.exp(x[1]), x[2]]),
)(drift_func, diff_coeff)
),
(dim_z,),
(dim_x,),
(dim_v,),
None,
numpy_module=jnp,
)
def forward_func(z, x, v, δ):
# Clip first two state components below at -500, in original domain corresponding to
# exp(-500) ≈ 7 × 10^(-218) when updating state to prevent numerical NaN issues when
# these state components tends to negative infinity. 500 was chosen as the cutoff to