def CalculateSolution(Pe):
# Create the TFC Class:
N = 200
m = 190
nC = 2
tfc = utfc(N, nC, m, basis='LeP', x0=0., xf=1.)
x = tfc.x
# Get the Chebyshev polynomials
H = tfc.H
H0 = H(np.array([0.]))
Hf = H(np.array([1.]))
# Create the constraint expression and its derivatives
y = lambda x, xi: np.dot(H(x), xi) + (1. - x) * (1. - np.dot(H0, xi)
) - x * np.dot(Hf, xi)
yd = egrad(y)
ydd = egrad(yd)
L = lambda xi: ydd(x, xi) - Pe * yd(x, xi)
# Calculate the solution
zXi = np.zeros(H(x).shape[1])
xi, it = NLLS(zXi, L)
# Create the test set:
N = 1000
xTest = np.linspace(0., 1., N)
err = np.abs(y(xTest, xi) - soln(xTest, Pe))
return np.max(err), np.mean(err)
def test_dstep():
A = np.array([5., 4., -1., 1., 0.])
B = lambda A: step(A)
C = lambda A: step(A)
D = egrad(B, 0)(A)
E = egrad(C, 0)(A)
assert (np.all(D == np.zeros_like(D)))
assert (np.all(E == D))
A = np.array([[5., 4., -1., 1., 0.], [-1.5, 0., 0., 2.2, 5.6]])
D = egrad(B, 0)(A)
E = egrad(C, 0)(A)
assert (np.all(D == np.zeros_like(D)))
assert (np.all(E == D))
B = lambda A: A * step(A) + A
C = lambda A: A * step(A) + A
D = egrad(B, 0)(A)
E = egrad(C, 0)(A)
true = np.array([[1., 1., 0., 1., 0.], [0., 0., 0., 1., 1.]
]) + np.ones_like(A)
assert (np.all(D == true))
assert (np.all(E == D))
B = lambda A: A**3 * step(A) + np.sin(A) * step(A)
C = lambda A: A**3 * step(A) + np.sin(A) * step(A)
D = egrad(egrad(B, 0), 0)(A)
E = egrad(egrad(C, 0), 0)(A)
true = 6. * A * np.array([[1., 1., 0., 1., 0.], [0., 0., 0., 1., 1.]
]) - np.sin(A) * np.array([[1., 1., 0., 1., 0.],
[0., 0., 0., 1., 1.]])
assert (np.all(D == true))
assert (np.all(E == D))
def test_ELMSin():
x = np.linspace(0, 1, num=10)
elm = ELMSin(0., 1., np.array([], dtype=np.int32), 10)
Fc1 = elm.H(x, 0, False)
Fc2 = elm.H(x, 1, False)
Fc3 = elm.H(x, 2, False)
Fc4 = elm.H(x, 3, False)
Fc5 = elm.H(x, 4, False)
x = x.reshape(10, 1)
x = np.ones((10, 10)) * x
w = elm.w.reshape(1, 10)
b = elm.b.reshape(1, 10)
sin = lambda x: np.sin(w * x + b)
dsin = egrad(sin)
d2sin = egrad(dsin)
d3sin = egrad(d2sin)
d4sin = egrad(d3sin)
d5sin = egrad(d4sin)
Fp1 = sin(x)
Fp2 = dsin(x)
Fp3 = d2sin(x)
Fp4 = d3sin(x)
Fp5 = d4sin(x)
assert (np.linalg.norm(Fc1 - Fp1, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc2 - Fp2, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc3 - Fp3, ord='fro') < 5e-14)
assert (np.linalg.norm(Fc4 - Fp4, ord='fro') < 5e-12)
assert (np.linalg.norm(Fc5 - Fp5, ord='fro') < 5e-12)
def BVP_tfc(N, m, basis, iterMax, tol):
## Unpack Paramters: *********************************************************
x0 = 0.
xf = np.pi
## Initial Conditions: *******************************************************
y0 = 0.
yf = 0.
nC = 2 # number of constraints
## Determine call tfc class needs to be 1 for ELMs
if basis == 'CP' or basis == 'LeP':
c = 2. / (xf - x0)
elif basis == 'FS':
c = 2. * np.pi / (xf - x0)
else:
c = 1. / (xf - x0)
## Compute true solution
ytrue = lambda x: np.exp(-x) * np.sin(x)
err = onp.ones_like(m) * np.nan
res = onp.ones_like(m) * np.nan
## GET CHEBYSHEV VALUES: *********************************************
tfc = utfc(N, nC, int(m), basis=basis, x0=x0, xf=xf)
x = tfc.x
H = tfc.H
H0 = H(tfc.x[0])
Hf = H(tfc.x[-1])
## DEFINE THE ASSUMED SOLUTION: *************************************
phi1 = lambda x: (np.pi - x) / np.pi
phi2 = lambda x: x / np.pi
f = lambda x: np.exp(-2. * x) * np.sin(x) * (np.cos(x) - np.sin(
x)) - 2. * np.exp(-x) * np.cos(x)
y = lambda x, xi: np.dot(H(x), xi) + phi1(x) * (y0 - np.dot(
H0, xi)) + phi2(x) * (yf - np.dot(Hf, xi))
yp = egrad(y)
ypp = egrad(yp)
## DEFINE LOSS AND JACOB ********************************************
L = lambda xi: ypp(x, xi) + y(x, xi) * yp(x, xi) - f(x)
## SOLVE THE SYSTEM *************************************************
xi = onp.zeros(H(x).shape[1])
xi, _, time = NLLS(xi, L, timer=True, maxIter=iterMax)
## COMPUTE ERROR AND RESIDUAL ***************************************
err = onp.linalg.norm(y(x, xi) - ytrue(x))
res = onp.linalg.norm(L(xi))
return err, res, time
def IVP2BVP(N, m, gamma, basis, iterMax, tol):
## Unpack Paramters: *********************************************************
x0 = -1.
xf = 1.
## Initial Conditions: *******************************************************
y0 = -2.
y0p = -2.
yf = 2.
nC = 2 # number of constraints
## Determine call tfc class needs to be 1 for ELMs
if basis == 'CP' or basis == 'LeP':
c = 2./ (xf - x0)
elif basis == 'FS':
c = 2. * np.pi / (xf - x0)
else:
c = 1./ (xf - x0)
## GET CHEBYSHEV VALUES: *********************************************
tfc = utfc(N,nC,m,basis = basis, x0=-1., xf=1.)
x = tfc.x
H = tfc.H
dH = tfc.dH
H0 = H(tfc.z[0])
Hf = H(tfc.z[-1])
H0p = dH(tfc.z[0])
## DEFINE THE ASSUMED SOLUTION: *************************************
phi1 = lambda a: 1./(1. + 4.*gamma - gamma**2) * ( (1. + gamma) - 2.*gamma*a )
phi2 = lambda a: 1./(1. + 4.*gamma - gamma**2) * ( (1. - gamma)**2 + (1. - gamma**2)*a)
phi3 = lambda a: 1./(1. + 4.*gamma - gamma**2) * ( -gamma*(gamma-3.) + 2.*gamma*a )
y = lambda x, xi: np.dot(H(x),xi) + phi1(x)*(y0 - np.dot(H0,xi)) \
+ phi2(x)*(y0p - np.dot(H0p,xi)) \
+ phi3(x)*(yf - np.dot(Hf,xi))
yp = egrad(y,0)
ypp = egrad(yp,0)
## DEFINE LOSS AND JACOB ********************************************
L = lambda xi: ypp(x,xi) + (np.cos(3.*x**2) -3.*x + 1.)*yp(x,xi) \
+ (6.*np.sin(4.*x**2) - np.exp(np.cos(3.*x)))*y(x,xi) \
- 2. * (1.-np.sin(3.*x))*(3.*x-np.pi)/(4.-x)
## SOLVE THE SYSTEM *************************************************
xi = onp.zeros(H(x).shape[1])
xi,_,_ = NLLS(xi,L,timer=True)
return y(x,xi), L(xi), x
def test_ODE():
# This script will solve the non-linear differential equation
# of the form: y''+f(t)*y*y' = f2(t)
# Constants used in the differential equation:
f = lambda t: np.ones(t.shape)
f2 = lambda t: np.exp(-2. * t) * np.sin(t) * (np.cos(t) - np.sin(
t)) - 2. * np.exp(-t) * np.cos(t)
ti = 0.
tf = np.pi
yi = 0.
yf = 0.
# Real analytical solution:
real = lambda t: np.exp(-t) * np.sin(t)
# Create the ToC Class:
N = 100
m = 30
nC = 2
tfc = TFC(N, nC, m, x0=ti, xf=tf, basis='LeP')
t = tfc.x
# Get the Chebyshev polynomials
H = tfc.H
dH = tfc.dH
Zero = np.zeros_like(t)
End = tf * np.ones_like(t)
H0 = H(Zero)
Hd0 = dH(Zero)
Hf = H(End)
# Create the constraint expression and its derivatives
beta0 = lambda t: (t - tf) / (ti - tf)
beta1 = lambda t: (ti - t) / (ti - tf)
y = lambda t, xi: np.dot(H(t), xi) + beta0(t) * (yi - np.dot(
H0, xi)) + beta1(t) * (yf - np.dot(Hf, xi))
yd = egrad(y)
ydd = egrad(yd)
# Create the residual and jacobians
r = lambda xi, t: ydd(t, xi) + f(t) * y(t, xi) * yd(t, xi) - f2(t)
xi = np.zeros(H(t).shape[1])
xi, it = NLLS(xi, r, t, constant_arg_nums=[1])
assert (np.max(np.abs(r(xi, t))) < 1e-10)
def test_ELMSwish():
x = np.linspace(0, 1, num=10)
elm = ELMSwish(0., 1., np.array([], dtype=np.int32), 10)
Fc1 = elm.H(x, 0, False)
Fc2 = elm.H(x, 1, False)
Fc3 = elm.H(x, 2, False)
Fc4 = elm.H(x, 3, False)
Fc5 = elm.H(x, 4, False)
Fc6 = elm.H(x, 5, False)
Fc7 = elm.H(x, 6, False)
Fc8 = elm.H(x, 7, False)
Fc9 = elm.H(x, 8, False)
x = x.reshape(10, 1)
x = np.ones((10, 10)) * x
w = elm.w.reshape(1, 10)
b = elm.b.reshape(1, 10)
swish = lambda x: (w * x + b) * (1. / (1. + np.exp(-w * x - b)))
dswish = egrad(swish)
d2swish = egrad(dswish)
d3swish = egrad(d2swish)
d4swish = egrad(d3swish)
d5swish = egrad(d4swish)
#d6swish = egrad(d5swish)
#d7swish = egrad(d6swish)
#d8swish = egrad(d7swish)
#d9swish = egrad(d8swish)
Fp1 = swish(x)
Fp2 = dswish(x)
Fp3 = d2swish(x)
Fp4 = d3swish(x)
Fp5 = d4swish(x)
#Fp6 = d5swish(x)
#Fp7 = d6swish(x)
#Fp8 = d7swish(x)
#Fp9 = d8swish(x)
assert (np.linalg.norm(Fc1 - Fp1, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc2 - Fp2, ord='fro') < 1e-13)
assert (np.linalg.norm(Fc3 - Fp3, ord='fro') < 5e-10)
assert (np.linalg.norm(Fc4 - Fp4, ord='fro') < 5e-10)
assert (np.linalg.norm(Fc5 - Fp5, ord='fro') < 5e-9)
def test_nELMSwish():
dim = 2
nC = -1 * np.ones(1, dtype=np.int32)
d = np.zeros(dim, dtype=np.int32)
c = np.ones(dim)
d2 = np.array([2, 3], dtype=np.int32)
nC2 = np.array([4], dtype=np.int32)
n = np.array([10] * dim)
N = np.prod(n)
z = np.linspace(0, 1, num=n[0])
X = onp.zeros((N, dim))
for k in range(dim):
nProd = np.prod(n[k + 1:])
nStack = np.prod(n[0:k])
dark = np.hstack([z] * nProd)
X[:, k] = onp.array([dark] * nStack).flatten()
c = 1. / (X[-1, :] - X[0, :])
z = (X - X[0, :]) * c
elm1 = nELMSwish(X[0, :], X[-1, :], nC, 10)
w = elm1.w
b = elm1.b
elm2 = nELMSwish(X[0, :], X[-1, :], nC2, 10)
elm2.w = w
elm2.b = b
Fc1 = elm1.H(X.T, d, False)
Fc2 = elm2.H(X.T, d2, False)
x = np.ones((100, 10)) * z[:, 0:1]
y = np.ones((100, 10)) * z[:, 1:2]
w1 = w[0, :].reshape((1, 10))
w2 = w[1, :].reshape((1, 10))
b = b.reshape((1, 10))
swish = lambda x, y: (w1 * x + w2 * y + b) * (1. /
(1. + np.exp(-(x * w1) -
(y * w2) - b)))
mydswish = egrad(egrad(egrad(egrad(egrad(swish, 0), 0), 1), 1), 1)
Fp1 = swish(x, y)
Fp2 = onp.delete(mydswish(x, y), nC2[0], axis=1)
assert (np.linalg.norm(Fc1 - Fp1, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc2 - Fp2, ord='fro') < 1e-12)
def test_ELMSigmoid():
x = np.linspace(0, 1, num=10)
elm = ELMSigmoid(0., 1., np.array([], dtype=np.int32), 10)
Fc1 = elm.H(x, 0, False)
Fc2 = elm.H(x, 1, False)
Fc3 = elm.H(x, 2, False)
Fc4 = elm.H(x, 3, False)
Fc5 = elm.H(x, 4, False)
Fc6 = elm.H(x, 5, False)
Fc7 = elm.H(x, 6, False)
Fc8 = elm.H(x, 7, False)
Fc9 = elm.H(x, 8, False)
x = x.reshape(10, 1)
x = np.ones((10, 10)) * x
w = elm.w.reshape(1, 10)
b = elm.b.reshape(1, 10)
sig = lambda x: 1. / (1. + np.exp(-w * x - b))
dsig = egrad(sig)
d2sig = egrad(dsig)
d3sig = egrad(d2sig)
d4sig = egrad(d2sig)
#d5sig = egrad(d2sig)
#d6sig = egrad(d2sig)
#d7sig = egrad(d2sig)
#d8sig = egrad(d2sig)
#d9sig = egrad(d2sig)
Fp1 = sig(x)
Fp2 = dsig(x)
Fp3 = d2sig(x)
Fp4 = d3sig(x)
#Fp5 = d4sig(x)
#Fp6 = d5sig(x)
#Fp7 = d6sig(x)
#Fp8 = d7sig(x)
#Fp9 = d8sig(x)
assert (np.linalg.norm(Fc1 - Fp1, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc2 - Fp2, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc3 - Fp3, ord='fro') < 5e-10)
assert (np.linalg.norm(Fc4 - Fp4, ord='fro') < 5e-10)
def test_ELMTanh():
x = np.linspace(0, 1, num=10)
elm = ELMTanh(0., 1., np.array([], dtype=np.int32), 10)
Fc1 = elm.H(x, 0, False)
Fc2 = elm.H(x, 1, False)
Fc3 = elm.H(x, 2, False)
Fc4 = elm.H(x, 3, False)
Fc5 = elm.H(x, 4, False)
Fc6 = elm.H(x, 5, False)
Fc7 = elm.H(x, 6, False)
Fc8 = elm.H(x, 7, False)
Fc9 = elm.H(x, 8, False)
x = x.reshape(10, 1)
x = np.ones((10, 10)) * x
w = elm.w.reshape(1, 10)
b = elm.b.reshape(1, 10)
Tanh = lambda x: np.tanh(w * x + b)
dTanh = egrad(Tanh)
d2Tanh = egrad(dTanh)
d3Tanh = egrad(d2Tanh)
d4Tanh = egrad(d3Tanh)
#d5Tanh = egrad(d4Tanh)
#d6Tanh = egrad(d5Tanh)
#d7Tanh = egrad(d6Tanh)
#d8Tanh = egrad(d7Tanh)
#d9Tanh = egrad(d8Tanh)
Fp1 = Tanh(x)
Fp2 = dTanh(x)
Fp3 = d2Tanh(x)
Fp4 = d3Tanh(x)
#Fp5 = d4Tanh(x)
#Fp6 = d5Tanh(x)
#Fp7 = d6Tanh(x)
#Fp8 = d7Tanh(x)
#Fp9 = d8Tanh(x)
assert (np.linalg.norm(Fc1 - Fp1, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc2 - Fp2, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc3 - Fp3, ord='fro') < 1e-13)
assert (np.linalg.norm(Fc4 - Fp4, ord='fro') < 5e-10)
def Solve(m, xtfc):
# TFC Constants:
if xtfc:
nC = -1
basis = 'ELMTanh'
else:
nC = [2, 2, 2]
basis = 'CP'
# Create the TFC Class:
N = [n, n, n]
myTfc = mtfc(N, nC, m, dim=3, basis=basis, x0=x0, xf=xf)
# Create the TFC constrained expression
x = myTfc.x
H = myTfc.H
u1 = lambda xi, *x: np.dot(H(*x), xi) - (1. - x[0]) * np.dot(
H(np.zeros_like(x[0]), x[1], x[2]), xi) - x[0] * np.dot(
H(np.ones_like(x[0]), x[1], x[2]), xi)
u2 = lambda xi, *x: u1(xi, *x) - (1. - x[1]) * u1(
xi, x[0], np.zeros_like(x[1]), x[2]) - x[1] * u1(
xi, x[0], np.ones_like(x[1]), x[2])
du2dt = egrad(u2, 3)
u = lambda xi, *x: u2(xi, *x) + np.sin(np.pi * x[0]) * np.sin(np.pi * x[
1]) - u2(xi, x[0], x[1], np.zeros_like(x[2])) - x[2] * du2dt(
xi, x[0], x[1], np.zeros_like(x[2]))
# Create the residual
uxx = egrad(egrad(u, 1), 1)
uyy = egrad(egrad(u, 2), 2)
utt = egrad(egrad(u, 3), 3)
r = lambda xi: uxx(xi, *x) + uyy(xi, *x) - 64. * utt(xi, *x)
xi = np.zeros(H(*x).shape[1])
if xtfc:
LS = lambda xi: np.linalg.lstsq(jacfwd(r, 0)
(xi), -r(xi), rcond=None)[0]
else:
LS = lambda xi: np.dot(np.linalg.pinv(jacfwd(r, 0)(xi)), -r(xi))
xi = LS(xi)
# Calculate the test set error
dark = np.meshgrid(np.linspace(x0[0], xf[0], nTest),
np.linspace(x0[1], xf[1], nTest),
np.linspace(x0[2], xf[2], nTest))
xTest = tuple([k.flatten() for k in dark])
err = np.abs(real(*xTest) - u(xi, *xTest))
return np.max(err), np.mean(err), myTfc.basisClass.numBasisFunc
def test_ELMReLU():
from jax.nn import relu as ReLU
x = np.linspace(0, 1, num=10)
elm = ELMReLU(0., 1., np.array([], dtype=np.int32), 10)
Fc1 = elm.H(x, 0, False)
Fc2 = elm.H(x, 1, False)
Fc3 = elm.H(x, 2, False)
Fc4 = elm.H(x, 3, False)
x = x.reshape(10, 1)
x = np.ones((10, 10)) * x
w = elm.w.reshape(1, 10)
b = elm.b.reshape(1, 10)
relu = lambda x: ReLU(w * x + b)
drelu = egrad(relu)
d2relu = egrad(drelu)
Fp1 = relu(x)
Fp2 = drelu(x)
Fp3 = d2relu(x)
assert (np.linalg.norm(Fc1 - Fp1, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc2 - Fp2, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc3 - Fp3, ord='fro') < 1e-14)
def test_nELMReLU():
dim = 2
nC = -1 * np.ones(1, dtype=np.int32)
d = np.zeros(dim, dtype=np.int32)
c = np.ones(dim)
d2 = np.array([0, 1], dtype=np.int32)
d3 = np.array([1, 1], dtype=np.int32)
d4 = np.array([0, 2], dtype=np.int32)
nC2 = np.array([4], dtype=np.int32)
n = np.array([10] * dim)
N = np.prod(n)
z = np.linspace(0, 1, num=n[0])
X = onp.zeros((N, dim))
for k in range(dim):
nProd = np.prod(n[k + 1:])
nStack = np.prod(n[0:k])
dark = np.hstack([z] * nProd)
X[:, k] = onp.array([dark] * nStack).flatten()
c = 1. / (X[-1, :] - X[0, :])
z = (X - X[0, :]) * c
elm1 = nELMReLU(X[0, :], X[-1, :], nC, 10)
w = elm1.w
b = elm1.b
elm2 = nELMReLU(X[0, :], X[-1, :], nC2, 10)
elm2.w = w
elm2.b = b
Fc1 = elm1.H(X.T, d, False)
Fc2 = elm2.H(X.T, d2, False)
Fc3 = elm2.H(X.T, d3, False)
Fc4 = elm2.H(X.T, d4, False)
x = np.ones((100, 10)) * z[:, 0:1]
y = np.ones((100, 10)) * z[:, 1:2]
w1 = w[0, :].reshape((1, 10))
w2 = w[1, :].reshape((1, 10))
b = b.reshape((1, 10))
relu = lambda x, y: np.maximum(np.zeros_like(x), w1 * x + w2 * y + b)
Fp1 = relu(x, y)
Fp2 = onp.delete(egrad(relu, 1)(x, y), nC2[0], axis=1)
assert (np.linalg.norm(Fc1 - Fp1, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc2 - Fp2, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc3, ord='fro') < 1e-14)
assert (np.linalg.norm(Fc4, ord='fro') < 1e-14)
def test_PDE():
## TFC Parameters
maxIter = 10
tol = 1e-13
# Constants and switches:
n = 20
m = 20
x0 = np.array([0., 0.])
xf = np.array([1., 1.])
# Real analytical solution:
real = lambda x, y: y**2 * np.sin(np.pi * x)
# Create the TFC Class:
N = np.array([n, n])
nC = np.array([2, 2])
tfc = TFC(N, nC, m, x0=x0, xf=xf, dim=2, basis='CP')
x = tfc.x
Zero = np.zeros_like(x[0])
One = np.ones_like(x[0])
# Get the basis functions
H = tfc.H
Hy = tfc.Hy
z1 = lambda xi, *x: np.dot(H(*x), xi) - (1. - x[0]) * np.dot(
H(*(Zero, x[1])), xi) - x[0] * np.dot(H(*(One, x[1])), xi)
z = lambda xi, *x: z1(xi, *x) - z1(xi, x[0], Zero) + x[1] * (2. * np.sin(
np.pi * x[0]) - egrad(z1, 2)(xi, x[0], One))
# Create the residual
zxx = egrad(egrad(z, 1), 1)
zyy = egrad(egrad(z, 2), 2)
zy = egrad(z, 2)
r = lambda xi, *x: zxx(xi, *x) + zyy(xi, *x) + z(xi, *x) * zy(
xi, *x) - np.sin(np.pi * x[0]) * (2. - np.pi**2 * x[1]**2 + 2. * x[1]**
3 * np.sin(np.pi * x[0]))
xi = np.zeros(H(*x).shape[1])
xi, it = NLLS(xi, r, *x, constant_arg_nums=[1, 2])
zr = real(x[0], x[1])
ze = z(xi, *x)
err = zr - ze
maxErr = np.max(np.abs(err))
assert (maxErr < 1e-10)
dH = tfc.dH
H0 = H(x[0:1])
H0p = dH(x[0:1])
## define tfc constrained expression and derivatives: **********************************************
# switching function
phi1 = lambda x: np.ones_like(x)
phi2 = lambda x: x
# tfc constrained expression
y = lambda x,xi,IC: np.dot(H(x),xi) + phi1(x)*(IC['y0'] - np.dot(H0,xi)) \
+ phi2(x)*(IC['y0p'] - np.dot(H0p,xi))
# !!! notice here that the initial conditions are passed as a dictionary (i.e. IC['y0'])
# this will be important so that the least-squares does not need to be re-JITed
yp = egrad(y)
ypp = egrad(yp)
## define the loss function: ***********************************************************************
# yₓₓ + w^2 y = 0
L = jit(lambda xi, x, IC: ypp(x, xi, IC) + w**2 * y(x, xi, IC))
## construct the least-squares class: **************************************************************
xi0 = np.zeros(H(x).shape[1])
IC = {'y0': np.array([y0]), 'y0p': np.array([y0p])}
ls = LsClass(xi0, L, x, IC, timer=True, constant_arg_nums=[1])
## initialize dictionary to record solution: *******************************************************
xSol = onp.zeros((Nstep, N))
ySol = onp.zeros_like(xSol)
def runLaneEmden(N, m, basis, k, xf):
## user defined parameters: ************************************************************************
# N - number of discretization points
# m - number of basis function terms
# basis - basis function type
# k - specific problem type, k >=0 (analytical solution known for k = 0, 1, and 5)
## problem initial conditions: *****************************************************************
xspan = [0., xf] # problem domain range [x0, xf], where x₀ > 0
y0 = 1. # y(x0) = 1
y0p = 0. # y'(x0) = 0
nC = 2 # number of constraints
## construct univariate tfc class: *************************************************************
tfc = utfc(N, nC, int(m), basis=basis, x0=xspan[0], xf=xspan[1])
x = tfc.x
H = tfc.H
dH = tfc.dH
H0 = H(x[0:1])
H0p = dH(x[0:1])
## define tfc constrained expression and derivatives: ******************************************
# switching function
phi1 = lambda x: np.ones_like(x)
phi2 = lambda x: x
# tfc constrained expression
y = lambda x, xi: np.dot(H(x), xi) + phi1(x) * (y0 - np.dot(
H0, xi)) + phi2(x) * (y0p - np.dot(H0p, xi))
yp = egrad(y)
ypp = egrad(yp)
## define the loss function: *******************************************************************
L = lambda xi: x * ypp(x, xi) + 2. * yp(x, xi) + x * y(x, xi)**k
## solve the problem via nonlinear least-squares ***********************************************
xi = np.zeros(H(x).shape[1])
# if k==0 or k==1, the problem is linear
if k == 0 or k == 1:
xi, time = LS(xi, L, timer=True)
iter = 1
else:
xi, iter, time = NLLS(xi, L, timer=True)
## compute the error (if k = 0, 1, or 5): ******************************************************
if k == 0:
ytrue = 1. - 1. / 6. * x**2
elif k == 1:
ytrue = onp.ones_like(x)
ytrue[1:] = np.sin(x[1:]) / x[1:]
elif k == 5:
ytrue = (1. + x**2 / 3)**(-1 / 2)
else:
ytrue = np.empty_like(x)
err = np.abs(y(x, xi) - ytrue)
## compute the residual of the loss vector: ****************************************************
res = np.abs(L(xi))
return x, y(x, xi), err, res
def laneEmden_tfc(N, m, type, xspan, basis, iterMax, tol):
## Unpack Paramters: *********************************************************
x0 = xspan[0]
xf = xspan[1]
## Initial Conditions: *******************************************************
y0 = 1.
y0p = 0.
nC = 2 # number of constraints
## Determine call tfc class needs to be 1 for ELMs
if basis == 'CP' or basis == 'LeP':
c = 2. / (xf - x0)
elif basis == 'FS':
c = 2. * np.pi / (xf - x0)
else:
c = 1. / (xf - x0)
## Compute true solution
if type == 0:
maxIter = 1
def ytrue(x):
val = onp.zeros_like(x)
val[0] = 1.
val[1:] = 1. - 1. / 6. * x[1:]**2
return val
elif type == 1:
maxIter = 1
def ytrue(x):
val = onp.zeros_like(x)
val[0] = 1.
val[1:] = np.sin(x[1:]) / x[1:]
return val
elif type == 5:
maxIter = iterMax
def ytrue(x):
val = onp.zeros_like(x)
val[0] = 1.
val[1:] = (1. + x[1:]**2 / 3)**(-1 / 2)
return val
else:
def ytrue(x):
return np.nan * np.ones_like(x)
err = np.ones_like(m) * np.nan
res = np.ones_like(m) * np.nan
## GET CHEBYSHEV VALUES: *********************************************
tfc = utfc(N, nC, int(m), basis=basis, x0=x0, xf=xf)
x = tfc.x
H = tfc.H
dH = tfc.dH
H0 = H(x[0])
H0p = dH(x[0])
## DEFINE THE ASSUMED SOLUTION: *************************************
phi1 = lambda x: np.ones_like(x)
phi2 = lambda x: x
y = lambda x,xi: np.dot(H(x),xi) \
+ phi1(x)*(y0 - np.dot(H0,xi)) \
+ phi2(x)*(y0p - np.dot(H0p,xi))
yp = egrad(y)
ypp = egrad(yp)
## DEFINE LOSS AND JACOB ********************************************
L = jit(lambda xi: x * ypp(x, xi) + 2. * yp(x, xi) + x * y(x, xi)**type)
## SOLVE THE SYSTEM *************************************************
# Solve the problem
xi = np.zeros(H(x).shape[1])
xi, _, time = NLLS(xi, L, timer=True, maxIter=maxIter)
## COMPUTE ERROR AND RESIDUAL ***************************************
err = np.linalg.norm(y(x, xi) - ytrue(x))
res = np.linalg.norm(L(xi))
return err, res, time
H0ps2 = dHs2(tfc2.x[0])
## DEFINE THE ASSUMED SOLUTION: *************************************
# First segment
phi11 = lambda a: 1. / (x1 - x0)**2 * (x1**2 - 2. * x1 * a + a**2)
phi12 = lambda a: 1./(x1 - x0)**2 * (x0*(x0-2.*x1) \
+ 2.*x1*a - a**2)
phi13 = lambda a: 1. / (x1 - x0) * (x0 * x1 - (x0 + x1) * a + a**2)
ys1 = lambda x, xi: np.dot(Hs1(x),xi['xi1']) \
+ phi11(x)*(y0 - np.dot(H0s1, xi['xi1'])) \
+ phi12(x)*(xi['y1'] - np.dot(Hfs1, xi['xi1'])) \
+ phi13(x)*(xi['y1d'] - np.dot(Hfps1,xi['xi1']))
yps1 = egrad(ys1, 0)
ypps1 = egrad(yps1, 0)
# Second segment
phi21 = lambda a: 1./(xf - x1)**2 * (xf*(xf-2.*x1) \
+ 2.*x1*a - a**2)
phi22 = lambda a: 1. / (xf - x1) * (-xf * x1 + (xf + x1) * a - a**2)
phi23 = lambda a: 1. / (xf - x1)**2 * (x1**2 - 2. * x1 * a + a**2)
ys2 = lambda x, xi: np.dot(Hs2(x),xi['xi2']) \
+ phi21(x)*(xi['y1'] - np.dot(H0s2,xi['xi2'])) \
+ phi22(x)*(xi['y1d'] - np.dot(H0ps2,xi['xi2'])) \
+ phi23(x)*(yf - np.dot(Hfs2,xi['xi2']))
yps2 = egrad(ys2, 0)
ypps2 = egrad(yps2, 0)
1./(zf-z0)**3 * (-zf**2*(3.*z0-zf) + 6.*z0*zf*a - 3.*(z0+zf)*a**2 + 2.*a**3),1)
phi2 = lambda a: np.expand_dims(\
1./(zf-z0)**3 * (-z0**2*(z0-3.*zf) - 6.*z0*zf*a + 3.*(z0+zf)*a**2 - 2.*a**3),1)
phi3 = lambda a: np.expand_dims(\
1./(zf-z0)**2 * (-z0*zf**2 + zf*(2.*z0+zf)*a - (z0+2.*zf)*a**2 + a**3),1)
phi4 = lambda a: np.expand_dims(\
1./(zf-z0)**2 * (-z0**2*zf + z0*(z0+2.*zf)*a - (2.*z0+zf)*a**2 + a**3),1)
## DEFINE CONSTRAINED EXPRESSION *************************************
r = lambda z, xi, IC: np.dot(Hs(z),xi['xis']) \
+ phi1(z)*(IC['R0'] - np.dot(Hs0, xi['xis'])) \
+ phi2(z)*( - np.dot(Hsf, xi['xis'])) \
+ phi3(z)*(IC['V0']/xi['b']**2 - np.dot(pHs0,xi['xis'])) \
+ phi4(z)*( - np.dot(pHsf,xi['xis']))
v = egrad(r)
a = egrad(v)
lam = lambda z, xi: np.dot(Hc(z), xi['xic'])
## FORM LOSS AND JACOBIAN ***********************************************************************************
Ls = lambda xi, IC: xi['b']**4 * a(z, xi, IC) - IC['ag'] + lam(z, xi)
Htf = lambda xi, IC: np.dot(
lam(z, xi)[-1, :], (-1. / 2. * lam(z, xi)[-1, :] + IC['ag']))
L = jit(lambda xi, IC: np.hstack([Ls(xi, IC).flatten(), Htf(xi, IC)]))
## INITIALIZE VARIABLES *************************************************************************************
xis = onp.zeros((Hs(z).shape[1], 3))
xic = onp.zeros((Hc(z).shape[1], 3))
b = np.sqrt(2) * onp.ones(1)
# y(pi) = exp(-pi/2) ################################################################################ from tfc import utfc from tfc.utils import MakePlot, TFCDict, egrad, NLLS, step import jax.numpy as np from jax import jit, jacfwd, lax import numpy as onp import tqdm from time import process_time as timer import scipy.optimize as optim global time, it #: Analytical solution soln = lambda x: (1. - np.exp(Pe * (x - 1.))) / (1. - np.exp(-Pe)) dsoln = lambda x: egrad(soln, 0)(x) ddsoln = lambda x: egrad(dsoln, 0)(x) # Constants used in the differential equation: Pe = 10**6 tol = 1e-13 xI = 0. xf = 1. yi = 1. yf = 0. xpBound = 1. - 1. * 10**-6 # Create the ToC Class: N = 200 c = 1.
def Ainv(*x):
dark = np.dot(ainv, np.vstack([*x]) - p0)
return (dark[0, :], dark[1, :])
# Create the constrained expression:
n1 = np.array([2., 1.])
n1mag = np.linalg.norm(n1)
n1 = n1 / n1mag
n2 = np.array([1., 1.])
n2mag = np.linalg.norm(n2)
g = lambda xi, *x: np.sin(x[0]) * np.cos(x[1])
dgn1 = lambda xi, *x: egrad(g, 1)(xi, *x) * n1[0] + egrad(g, 2)(xi, *x) * n1[1]
intgn2 = lambda t: -1. / 2. * np.cos(t)**2 * np.sqrt(2.)
u1 = lambda xi,*x: g(xi,*x)+\
n1mag*(Ainv(*x)[0]-1.)*(0.5-dgn1(xi,*A(np.zeros_like(x[0]),Ainv(*x)[1])))+\
1./n2mag*(-2.-intgn2(5.)+intgn2(4.))
uslow = lambda xi,*x: u1(xi,*x)+\
(1./2.-Ainv(*x)[1])*(u1(xi,*A(Ainv(*x)[0],np.ones_like(x[1])))-u1(xi,*A(Ainv(*x)[0],np.zeros_like(x[1]))))
u = jit(uslow)
# Plot:
x = A(*X)
xi = np.zeros(5)
xmat = [k.reshape(n) for k in x]
p = MakePlot('x', 'y', zlabs='u(x,y,g(x,y))')
from tfc.utils import egrad
from tfc.utils.PlotlyMakePlot import MakePlot
# Create X:
n = [50, 50]
xMat = np.meshgrid(np.linspace(-4., 4., 50), np.linspace(-4., 4., 50))
x = xMat[0].flatten() + 1.j * xMat[1].flatten()
# Create the constrained expression:
g = lambda x: 0.25 * x + np.cos(x / 4.) * 0.3j
uslow = lambda x: g(x)\
+((-44.+8.j)*x**2+(52.+36.j)*x+(132.-24.j))/125.*(1.+np.pi*1.j-g(0.5j*np.ones_like(x)))\
+((69.+67.j)/125.*x**2-(129.+397.j)/250.*x+(-82.+49.j)/125.)*(g(np.ones_like(x))-g(1.j*np.ones_like(x)))\
+((44.-8.j)*x**2-(52.+36.j)*x+(-7.+24.j))/125.*(2.j-g(np.ones_like(x)*(2.+1.j))-egrad(g)(np.ones_like(x)))
u = jit(uslow)
# Plot:
p = MakePlot('Re[x]', 'Im[x]', zlabs='Re[u(x,g(x))]')
p2 = MakePlot('Re[x]', 'Im[x]', zlabs='Im[u(x,g(x))]')
# Add constraints
p.Scatter3d(
x=[0.],
y=[1. / 2.],
z=[1.],
mode="markers",
marker=dict(color="red", size=5),
)
p2.Scatter3d(
Hx0 = Hx(xtfc.z[0])
Hxf = Hx(xtfc.z[-1])
Hu = utfc.H
## DEFINE THE ASSUMED SOLUTION: *************************************
z = xtfc.z
z0 = z[0]
zf = z[-1]
phi1 = lambda a: (zf - a) / (zf - z0)
phi2 = lambda a: (a - z0) / (zf - z0)
x = lambda z, xi: np.dot(Hx(z), xi) + phi1(z) * (x0 - np.dot(Hx0, xi)) + phi2(
z) * (xf - np.dot(Hxf, xi))
xp = egrad(x, 0)
u = lambda z, xi: np.dot(Hu(z), xi)
up = egrad(u, 0)
## LOSS FUNCTIONS AND JACOB *****************************************
Lx = lambda z, xi: -xi['b']**2 * xp(z, xi['xi_x']) - alfa * x(z, xi[
'xi_x']) - beta * u(z, xi['xi_u'])
Lu = lambda z, xi: -xi['b']**2 * up(z, xi['xi_u']) - beta * x(z, xi[
'xi_x']) + alfa * u(z, xi['xi_u'])
H = lambda z, xi: 0.5 * x(z, xi['xi_x'])**2 - 0.5 * u(z, xi[
'xi_u'])**2 - alfa / beta * x(z, xi['xi_x']) * u(z, xi['xi_u'])
Lf = lambda z, xi: H(z, xi)[-1]
L = lambda xi: np.hstack((Lx(z, xi), Lu(z, xi), H(z, xi)))
# Create the TFC Class: N = [n,n,n] myTfc = mtfc(N,nC,m,dim=3,basis=basis,x0=x0,xf=xf) x = myTfc.x # Get the basis functions H = myTfc.H Hx = myTfc.Hx # Create the TFC constrained expression (here f stands as a placeholder for u and v) f1 = lambda xi,*x: np.dot(H(*x),xi)-np.dot(H(np.zeros_like(x[0]),x[1],x[2]),xi)-x[0]*np.dot(Hx(xend*np.ones_like(x[0]),x[1],x[2]),xi) f2 = lambda xi,*x: f1(xi,*x)-(Hb-2.*x[1])/(2.*Hb)*f1(xi,x[0],-Hb/2.*np.ones_like(x[1]),x[2])-(Hb+2.*x[1])/(2.*Hb)*f1(xi,x[0],Hb/2.*np.ones_like(x[1]),x[2]) f = lambda xi,*x: f2(xi,*x)-f2(xi,x[0],x[1],np.zeros_like(x[2])) fx = egrad(f,1) f2x = egrad(fx,1) fy = egrad(f,2) f2y = egrad(fy,2) ft = egrad(f,3) # Create the residual and jacobian L1 = lambda xiu,xiv,*x: fx(xiu,*x)+fy(xiv,*x) L2 = lambda xiu,xiv,*x: rho*(ft(xiu,*x)+f(xiu,*x)*fx(xiu,*x)+f(xiv,*x)*fy(xiu,*x))+P-mu*(f2x(xiu,*x)+f2y(xiu,*x)) L3 = lambda xiu,xiv,*x: rho*(ft(xiv,*x)+f(xiu,*x)*fx(xiv,*x)+f(xiv,*x)*fy(xiv,*x))-mu*(f2x(xiv,*x)+f2y(xiv,*x)) L = lambda xi: np.hstack([L1(xi['xiu'],xi['xiv'],*x),L2(xi['xiu'],xi['xiv'],*x),L3(xi['xiu'],xi['xiv'],*x)]) # Calculate the xi values M = H(*x).shape[1] xiu = np.zeros(M) xiv = np.zeros(M)
testErr[j,k] = np.nan
continue
# Get the basis functions
H = tfc.H
# Create the constrained expression
u1 = lambda xi,*x: np.dot(H(*x),xi)\
+(1.-x[0])*(x[1]**3-np.dot(H(np.zeros_like(x[0]),x[1]),xi))\
+x[0]*((1.+x[1]**3)*np.exp(-1.)-np.dot(H(np.ones_like(x[0]),x[1]),xi))
u = lambda xi,*x: u1(xi,*x)\
+(1.-x[1])*(x[0]*np.exp(-x[0])-u1(xi,x[0],np.zeros_like(x[1])))\
+x[1]*(np.exp(-x[0])*(x[0]+1.)-u1(xi,x[0],np.ones_like(x[1])))
# Create the residual
laplace = lambda xi,*x: egrad(egrad(u,1),1)(xi,*x)+egrad(egrad(u,2),2)(xi,*x)
L = lambda xi,*x: laplace(xi,*x)-np.exp(-x[0])*(x[0]-2.+x[1]**3+6.*x[1])
# Calculate the xi values
zXi = np.zeros(H(*x).shape[1])
A = jacfwd(L,0)(zXi,*x)
B = -L(zXi,*x)
xi = np.dot(np.linalg.pinv(A),B)
# Calculate the error
dark = np.meshgrid(np.linspace(x0[0],xf[0],n),np.linspace(x0[1],xf[1],n))
x = (dark[0].flatten(),dark[1].flatten())
ur = real(*x)
ue = u(xi,*x)
err = ur-ue
## GET CHEBYSHEV VALUES ********************************************** stfc = utfc(N, nCx, ms, basis='CP', x0=-1, xf=1.) ctfc = utfc(N, nCy, mc, basis='CP', x0=-1, xf=1.) Hs = stfc.H Hc = ctfc.H ## DEFINE THE ASSUMED SOLUTION ************************************** z = stfc.z z0 = z[0] zf = z[-1] ## DEFINE CONSTRAINED EXPRESSION ************************************* r = lambda z, xi, IC: np.dot(Hs(z), xi['xis']) v = egrad(r, 0) a = egrad(v, 0) lam = lambda z, xi: np.dot(Hc(z), xi['xic']) lamr = egrad(lam, 0) ## FORM LOSS AND JACOBIAN *********************************************************************************** L0 = lambda xi, IC: r(z, xi, IC)[0, :] - IC['R0'] Ld0 = lambda xi, IC: xi['b']**2 * v(z, xi, IC)[0, :] - IC['V0'] Lf = lambda xi, IC: r(z, xi, IC)[-1, :] Ldf = lambda xi, IC: xi['b']**2 * v(z, xi, IC)[-1, :] Ls = lambda xi, IC: xi['b']**4 * a(z, xi, IC) - IC['ag'] + lam(z, xi) # Htf = lambda xi,IC: np.dot(lam(z,xi)[-1,:],(-1./2.*lam(z,xi)[-1,:] + IC['ag'])) # Updated because need to at lam_r * v term for spectral method
from tfc.utils import LS, egrad, MakePlot
# Constants and switches:
x0 = np.array([1., 0.])
xf = np.array([4., 2. * np.pi])
n = 30
nTest = 100
xTFC = False # Set to True to use X-TFC rather than TFC
usePlotly = True # Set to true to use plotly rather than matplotlib
# Real analytical solution:
real = lambda r, th: r**2 / 8. + np.pi * np.cos(th) / r + r**2 / 4. * np.sin(
2. * th) + r**3 / 16. * np.sin(3. * th)
d2realdr2 = egrad(egrad(real, 0), 0)
# MC constants:
if xTFC:
m = 335
else:
m = 30
# Create the TFC Class:
N = [n, n]
if xTFC:
nC = -1
myTfc = mtfc(N, nC, m, dim=2, basis='ELMTanh', x0=x0, xf=xf)
else:
nC = [4, 4]
myTfc = mtfc(N, nC, m, dim=2, basis='CP', x0=x0, xf=xf)
# form constrained expressions
r = lambda z, xi: np.dot(H(z),xi['xis']) + phi1(z)*(R0(xi) - np.dot(H0, xi['xis'])) \
+ phi2(z)*(R0(xi) - np.dot(Hf, xi['xis'])) \
+ phi3(z)*(V0(xi) - np.dot(Hp0,xi['xis'])) \
+ phi4(z)*(V0(xi) - np.dot(Hpf,xi['xis']))
# terms defining distance from third-body to the primary (R1) and the secondary (R2)
R1 = lambda z, xi: np.sqrt((r(z, xi)[:, 0] + mu)**2 + r(z, xi)[:, 1]**2)
R2 = lambda z, xi: np.sqrt((r(z, xi)[:, 0] + mu - 1.)**2 + r(z, xi)[:, 1]**2)
# scaling term dr/dt = dr/dz * dz/dζ * dζ/dt = b²/R₂ dr/dz
d2 = lambda z, xi: 1. / R2(z, xi)
# !!! where t is the problem domain, ζ is the scaled domain, and z is the basis domain
# derivatives of the constrained expressions
v = lambda z, xi: np.expand_dims(d2(z, xi), 1) * xi['b']**2 * egrad(r)(z, xi)
a = lambda z, xi: np.expand_dims(d2(z, xi), 1) * xi['b']**2 * egrad(v)(z, xi)
## form the loss vector: ***************************************************************************
# jacobi constant equation
Jc = lambda Z, xi: (r(z,xi)[:,0]**2 + r(z,xi)[:,1]**2) \
+ 2.*(1.-mu)/R1(z,xi) \
+ 2.*mu/R2(z,xi) \
+ (1.-mu)*mu \
- (v(z,xi)[:,0]**2 + v(z,xi)[:,1]**2)
# terms of the differential equation
Psi1 = lambda z, xi: np.hstack(( np.expand_dims( 2.*v(z,xi)[:,1] + r(z,xi)[:,0],1), \
np.expand_dims(-2.*v(z,xi)[:,0] + r(z,xi)[:,1],1) ))
r0 = 0.
rf = 5.
# Create TFC class:
myTfc = utfc(n, nC, m, x0=th0, xf=thf)
th = myTfc.x
H = myTfc.H
# Create constrained expression:
g = lambda th, xi: np.dot(H(th), xi)
r = lambda th,xi: g(th,xi)+\
(th-thf)/(th0-thf)*(r0-g(th0*np.ones_like(th),xi))+\
(th-th0)/(thf-th0)*(rf-g(thf*np.ones_like(th),xi))
# Create loss function:
dr = egrad(r)
d2r = egrad(dr)
L = lambda xi: -r(th,xi)**2*(dr(th,xi)*np.tan(th)+2.*d2r(th,xi))+\
-np.tan(th)*dr(th,xi)**3+3.*r(th,xi)*dr(th,xi)**2+r(th,xi)**3
# Solve the problem:
xi = np.zeros(H(th).shape[1])
xi, _, time = NLLS(xi, L, timer=True)
# Print out statistics:
print("Solution time: {0} seconds".format(time))
# Plot the solution and residual
p = MakePlot([r"$y$"], [r"$x$"])
p.ax[0].plot(r(th, xi) * np.sin(th), r(th, xi) * np.cos(th), "k")
p.ax[0].axis("equal")
1./(zf-z0)**3 * (-zf**2*(3.*z0-zf) + 6.*z0*zf*a - 3.*(z0+zf)*a**2 + 2.*a**3),1)
phi2 = lambda a: np.expand_dims(\
1./(zf-z0)**3 * (-z0**2*(z0-3.*zf) - 6.*z0*zf*a + 3.*(z0+zf)*a**2 - 2.*a**3),1)
phi3 = lambda a: np.expand_dims(\
1./(zf-z0)**2 * (-z0*zf**2 + zf*(2.*z0+zf)*a - (z0+2.*zf)*a**2 + a**3),1)
phi4 = lambda a: np.expand_dims(\
1./(zf-z0)**2 * (-z0**2*zf + z0*(z0+2.*zf)*a - (2.*z0+zf)*a**2 + a**3),1)
## DEFINE CONSTRAINED EXPRESSION *************************************
r = lambda z, xi, IC: np.dot(Hs(z),xi['xis']) \
+ phi1(z)*(IC['R0'] - np.dot(Hs0, xi['xis'])) \
+ phi2(z)*( - np.dot(Hsf, xi['xis'])) \
+ phi3(z)*(IC['V0']/xi['b']**2 - np.dot(pHs0,xi['xis'])) \
+ phi4(z)*( - np.dot(pHsf,xi['xis']))
v = egrad(r)
a = egrad(v)
lam = lambda z, xi: np.dot(Hc(z), xi['xic'])
## FORM LOSS AND JACOBIAN ***********************************************************************************
Ls = lambda xi, IC: xi['b']**4 * a(z, xi, IC) - IC['ag'] + lam(z, xi)
Htf = lambda xi, IC: np.dot(
lam(z, xi)[-1, :], (-1. / 2. * lam(z, xi)[-1, :] + IC['ag']))
L = jit(lambda xi, IC: np.hstack([Ls(xi, IC).flatten(), Htf(xi, IC)]))
## INITIALIZE VARIABLES *************************************************************************************
xis = onp.zeros((Hs(z).shape[1], 3))
xic = onp.zeros((Hc(z).shape[1], 3))
b = np.sqrt(2) * onp.ones(1)
