def Jdark(x, xi, xp):
jacob = jacfwd(L, 1)(z, xi, xp)
return np.hstack((jacob[k] for k in xi.keys()))
J = jit(lambda z, xi, xp: Jdark(z, xi, xp), static_argnums=[
0,
])
# Create the residual and jacobians
xi1 = onp.zeros(H(z).shape[1])
xi2 = onp.zeros(H(z).shape[1])
y = onp.zeros(1)
yd = onp.zeros(1)
xi = TFCDict({'xi1': xi1, 'xi2': xi2, 'y': y, 'yd': yd})
# Create the NLLS
def cond(val):
return np.all(
np.array([
np.max(np.abs(L(z, val['xi'], val['xp']))) > tol, val['it'] < 30,
np.max(np.abs(val['dxi'])) > tol
]))
def body(val):
val['dxi'] = -np.dot(np.linalg.pinv(J(z, val['xi'], val['xp'])),
L(z, val['xi'], val['xp']))
val['xi'] += val['dxi']
L1 = lambda xi: ypps1(xs1, xi) + ys1(xs1, xi) - f(xs1)
L2 = lambda xi: ypps2(xs2, xi) + ys2(xs2, xi)*yps2(xs2, xi) \
- f(xs2)
L = jit(lambda xi: np.hstack((L1(xi), L2(xi))))
## SOLVE THE SYSTEM *************************************************
xi1 = onp.zeros(Hs1(xs1).shape[1])
xi2 = onp.zeros(Hs2(xs2).shape[1])
m = (yf - y0) / (xf - x0)
y1 = onp.ones(1) * (m * xs1[-1] + y0)
y1d = onp.ones(1) * m
xi0 = TFCDict({'xi1': xi1, 'xi2': xi2, 'y1': y1, 'y1d': y1d})
xi = TFCDict({'xi1': xi1, 'xi2': xi2, 'y1': y1, 'y1d': y1d})
xi, it, time = NLLS(xi, L, timer=True)
## COMPUTE ERROR AND RESIDUAL ***************************************
x = np.hstack((xs1, xs2))
yinit = np.hstack((ys1(xs1, xi0), ys2(xs2, xi0)))
y = np.hstack((ys1(xs1, xi), ys2(xs2, xi)))
yp = np.hstack((yps1(xs1, xi), yps2(xs2, xi)))
err = onp.abs(y - ytrue(x))
res = onp.abs(L(xi))
print()
import numpy as onp
import scipy as sp
from time import process_time
import jax.numpy as np
from jax import jacfwd, jacrev, jit
from jax.lax import fori_loop
from tfc.utils import TFCDict, Latex
# Initial solution
X0 = TFCDict({
'x': np.array([2. / 3.]),
'y': np.array([1. / 3.]),
'z': np.array([1. / 3.])
})
# Create function, Jacobian, and Hessian
f = lambda X: np.squeeze(X['x'] * X['y'] * (X['x'] * X['y'] + 6. * X[
'y'] - 8. * X['x'] - 48.) + X['z']**2 - 8. * X['z'] + 9. * X['y']**2 - 72.
* X['y'] + 16. * X['x']**2 + 96. * X['x'] + 160)
J = lambda X: np.hstack([val for val in jacfwd(f)(X).values()])
H = lambda X: np.hstack([val for val in jacrev(J)(X).values()])
# Equality constraint
A = np.array([[1., 2., -1.], [1., 0., 1.]])
b = np.array([[1.], [1.]])
N = sp.linalg.null_space(A)
# Iterate to find the solution (use a jax for loop)
X = X0
## 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)))
xi_x = onp.zeros(Hx(z).shape[1])
xi_u = onp.zeros(Hu(z).shape[1])
b = onp.ones(1) * np.sqrt(2.)
xi = TFCDict({'xi_x': xi_x, 'xi_u': xi_u, 'b': b})
## SOLVE THE SYSTEM *************************************************
xi, it, time = NLLS(xi, L, timer=True, tol=tol, maxIter=iterMax)
t = (z - z[0]) / xi['b']**2
X = x(z, xi['xi_x'])
U = u(z, xi['xi_u'])
Ham = onp.zeros(len(t))
int = onp.zeros(len(t))
for i in range(0, len(t)):
int[i] = 0.5 * (X[i]**2 + U[i]**2)
Ham[i] = int[i] + -U[i] / beta * (alfa * X[i] + beta * U[i])
opacity=opacity)
p.Surface(x=xEll(*matEll),
y=yEll(*matEll),
z=zEll(*matEll),
showscale=False,
colorscale=[[0.,"black"],[1.,"black"]],
opacity=opacity)
# Plot the constrained expressions
onp.random.seed(1)
ind = np.array([0,-1])
m = H(t).shape[1]
for k in range(nCEs):
xi = TFCDict({'xi':onp.random.randn(m)/3.,
'psi':onp.random.randn(1),
'phi':onp.random.randn(1),
'v':onp.random.rand(1)*6.-3.,
'th':onp.random.randn(1),
'n':onp.random.randn(1)})
X = x(xi,t)
Y = y(xi,t)
Z = z(xi,t)
p.Scatter3d(x=X,y=Y,z=Z,
mode="lines",line=dict(color=colors[k],width=3))
p.Scatter3d(x=X[ind],y=Y[ind],z=Z[ind],
mode="markers",
marker=dict(size=4),
line=dict(color=colors[k]))
# Set plot properties and display it
p.fig.update_layout(showlegend=False,scene_aspectmode='cube')
p.view(-45,20)
yd2 = lambda z, xi: ydz2(z, xi) * c2(xi)
ydd2 = lambda z, xi: yddz2(z, xi) * c2(xi)**2
L1 = lambda z, xi: ydd1(z, xi) - Pe * yd1(z, xi)
L2 = lambda z, xi: ydd2(z, xi) - Pe * yd2(z, xi)
L = lambda xi: np.hstack((L1(z, xi), L2(z, xi)))
# Create the residual and jacobians
xi1 = onp.zeros(H(z).shape[1])
xi2 = onp.zeros(H(z).shape[1])
y = onp.zeros(1)
yd = onp.zeros(1)
b = onp.ones(1) * np.sqrt(2. / 0.5)
xi = TFCDict({'xi1': xi1, 'xi2': xi2, 'y': y, 'yd': yd, 'b': b})
## SOLVE THE SYSTEM *************************************************
xi, it, time = NLLS(xi, L, timer=True)
X = np.hstack((x1(z, xi), x2(z, xi)))
Y = np.hstack((y1(z, xi), y2(z, xi)))
# p1 = MakePlot(onp.array([['x (m)']]),onp.array([['y (m)']]))
# p1.ax[0].plot(X,Y,label='TFC Solution')
# p1.ax[0].plot(X,soln(X),label='Analytical Solution')
# p1.ax[0].legend()
# p1.show()
print('{:.2e} & {:.2e} & {:.5f} & {:.2f}'.format(np.max(np.abs(Y - soln(X))),
np.max(np.abs(L(xi))),
L0 = lambda xi: ys1(xs1, xi)[0] - y0
L1 = lambda xi: (ypps1(xs1, xi) + ys1(xs1, xi) - f(xs1))[1:-1]
Li = lambda xi: ys1(xs1, xi)[-1] - ys2(xs2, xi)[0]
Lip = lambda xi: yps1(xs1, xi)[-1] - yps2(xs2, xi)[0]
L2 = lambda xi: (ypps2(xs2, xi) + ys2(xs2, xi) * yps2(xs2, xi) - f(xs2))[1:-1]
Lf = lambda xi: ys2(xs2, xi)[-1] - yf
L = jit(lambda xi: np.hstack(
(L0(xi), L1(xi), Li(xi), Lip(xi), L2(xi), Lf(xi))))
## SOLVE THE SYSTEM ******************************************************************
xi1 = onp.zeros(Hs1(xs1).shape[1])
xi2 = onp.zeros(Hs2(xs2).shape[1])
xi0 = TFCDict({'xi1': xi1, 'xi2': xi2})
xi, it, time = NLLS(xi0, L, timer=True)
## COMPUTE ERROR AND RESIDUAL ******************************************************************
x = np.hstack((xs1, xs2))
y = np.hstack((ys1(xs1, xi), ys2(xs2, xi)))
yp = np.hstack((yps1(xs1, xi), yps2(xs2, xi)))
print()
print('Max Error: ' + str(np.max(np.abs(y - ytrue(x)))))
print('y0 Error: ' + str(np.abs(y[0] - y0)))
print('y1 Error: ' + str(np.abs(Li(xi))))
print('y1_x Error: ' + str(np.abs(Lip(xi))))
print('yf Error: ' + str(np.abs(y[-1] - yf)))
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)
xi = TFCDict({'xiu':xiu,'xiv':xiv})
if xTfc:
xi,it,time = NLLS(xi,L,maxIter=maxIter,method='lstsq',timer=True)
else:
xi,it,time = NLLS(xi,L,maxIter=maxIter,method='pinv',timer=True)
xiu = xi['xiu']; xiv = xi['xiv']
# Calcualte u and plot for different times
n = 100
X = onp.matlib.repmat(onp.reshape(onp.linspace(0,xf[0],num=n),(n,1)),n,1).flatten()
Y = onp.reshape(onp.matlib.repmat(onp.reshape(onp.linspace(-Hb/2.,Hb/2.,num=n),(n,1)),1,n),(n**2,1)).flatten()
xTest = onp.zeros((3,n**2*3))
xTest[0,:] = onp.hstack([X,]*3)
xTest[1,:] = onp.hstack([Y,]*3)
xTest[2,:] = onp.hstack([onp.ones(n**2)*0.01,onp.ones(n**2)*0.1,onp.ones(n**2)*tf])
## LOSS FUNCTIONS AND JACOB *****************************************
Lx = lambda z, xi, c: -c * xp(z, xi['xi_x']) - alfa * x(z, xi[
'xi_x']) - beta * u(z, xi['xi_u'])
Lu = lambda z, xi, c: -c * 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, c: np.hstack((Lx(z, xi, c), Lu(z, xi, c)))
xi_x = onp.zeros(Hx(z).shape[1])
xi_u = onp.zeros(Hu(z).shape[1])
xi = TFCDict({'xi_x': xi_x, 'xi_u': xi_u})
# Using dummy value for c here, 1
nlls = NllsClass(xi, L, 1, tol=tol, maxIter=2, timer=True)
def InnerLoop(tf, z, xi):
xi['xi_x'] = onp.zeros(Hx(z).shape[1])
xi['xi_u'] = onp.zeros(Hu(z).shape[1])
c = 2. / tf
xi, _, time = nlls.run(xi, c)
loss = np.max(np.abs(H(z, xi)))
return loss
def CalculateSolutionSplit(Pe):
if Pe > 1e3:
xpBoundL = 0. + 1.e-3
xpBoundU = 1. - 1e-3
else:
xpBoundL = 0. + 1.e-1
xpBoundU = 1. - 1e-1
# Create the ToC Class:
N = 200
m = 190
nC = 3
tfc = utfc(N, nC, m, basis='LeP', x0=-1., xf=1.)
# Get the Chebyshev polynomials
H = tfc.H
dH = tfc.dH
H0 = H(np.array([-1.]))
Hf = H(np.array([1.]))
Hd0 = dH(np.array([-1.]))
Hdf = dH(np.array([1.]))
# Create the constraint expression and its derivatives
z = tfc.z
xp = lambda xi: xi['xpHat'] + (xpBoundU - xi['xpHat']) * step(xi[
'xpHat'] - xpBoundU) + (xpBoundL - xi['xpHat']) * step(xpBoundL - xi[
'xpHat'])
c1 = lambda xi: 2. / (xp(xi))
c2 = lambda xi: 2. / (1. - xp(xi))
x1 = lambda z, xi: (z + 1.) / c1(xi)
x2 = lambda z, xi: (z + 1.) / c2(xi) + xp(xi)
y1 = lambda z,xi: np.dot(H(z),xi['xi1'])+(1.-2.*z+z**2)/4.*(1.-np.dot(H0,xi['xi1']))\
+(3.+2.*z-z**2)/4.*(xi['y']-np.dot(Hf,xi['xi1']))\
+(-1.+z**2)/2.*(xi['yd']/c1(xi)-np.dot(Hdf,xi['xi1']))
ydz1 = egrad(y1, 0)
yddz1 = egrad(ydz1, 0)
yd1 = lambda z, xi: ydz1(z, xi) * c1(xi)
ydd1 = lambda z, xi: yddz1(z, xi) * c1(xi)**2
y2 = lambda z,xi: np.dot(H(z),xi['xi2'])+(3.-2.*z-z**2)/4.*(xi['y']-np.dot(H0,xi['xi2']))\
+(1.-z**2)/2.*(xi['yd']/c2(xi)-np.dot(Hd0,xi['xi2']))\
+(1.+2.*z+z**2)/4.*(0.-np.dot(Hf,xi['xi2']))
ydz2 = egrad(y2, 0)
yddz2 = egrad(ydz2, 0)
yd2 = lambda z, xi: ydz2(z, xi) * c2(xi)
ydd2 = lambda z, xi: yddz2(z, xi) * c2(xi)**2
# Solve the problem
xi = TFCDict({
'xi1': onp.zeros(H(z).shape[1]),
'xi2': onp.zeros(H(z).shape[1]),
'xpHat': onp.array([0.99]),
'y': onp.array([0.]),
'yd': onp.array([0.])
})
L1 = lambda xi: ydd1(z, xi) - Pe * yd1(z, xi)
L2 = lambda xi: ydd2(z, xi) - Pe * yd2(z, xi)
L = lambda xi: np.hstack([L1(xi), L2(xi)])
xi, it = NLLS(xi, L)
# Create the test set:
N = 1000
z = np.linspace(-1., 1., N)
# Calculate the error and return the results
X = np.hstack([x1(z, xi), x2(z, xi)])
Y = np.hstack([y1(z, xi), y2(z, xi)])
err = np.abs(Y - soln(X, Pe))
return np.max(err), np.mean(err)
const['g'] = atmData['Gravity (m/s^2)'][k]
const['L'] = 208*const['g']
const['atm_density'] = atmData['Density (kg/m^3)'][k]
const['Msg'] = atmData['Gas Mass (kg).1'][k]
# Calcualted constants:
const['b'] = const['g']*const['atm_density']*(1.-mw_gas/mw_atm)
# Initial guess
if k == atmData.shape[0]-1:
init = pickle.load(open('TfcInit.pkl','rb'))
m = H(x).shape[1]
m1 = H1(x).shape[1]
xi = TFCDict({'xiTh':onp.zeros((m)), 'xiQ':onp.zeros((m1)), 'xiR':onp.zeros((m)),
'xiZ':onp.zeros((m1)), 'beta':onp.zeros((1)), 'ld':onp.zeros((1))})
xi['beta'] = init['beta']
xi['ld'] = init['ld']
xp = init['x'].flatten()
ro = onp.expand_dims(onp.interp(x,xp,init['r'].flatten()),1)
xi['xiR'] = onp.dot(onp.linalg.pinv(jacfwd(r,1)(x,xi)['xiR']),ro.flatten()-r(x,xi))
tho = onp.expand_dims(onp.interp(x,xp,init['th'].flatten()),1)
xi['xiTh'] = onp.dot(onp.linalg.pinv(jacfwd(th,1)(x,xi)['xiTh']),tho.flatten()-th(x,xi))
zo = onp.expand_dims(onp.interp(x,xp,init['z'].flatten()),1)
xi['xiZ']= onp.dot(onp.linalg.pinv(jacfwd(z,1)(x,xi)['xiZ']),zo.flatten()-z(x,xi))