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))