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) xi = TFCDictRobust({'xis':xis,\ 'xic':xic,\ 'b':b}) IC = {'R0': np.zeros((3, )), 'V0': np.zeros((3, )), 'ag': np.zeros((3, ))} ## NONLINEAR LEAST-SQUARES CLASS ***************************************************************************** nlls = NllsClass(xi, L, IC, tol=tol, maxIter=maxIter, timer=True) R0 = np.array([-150000., 30000., 15000.]) V0 = np.array([915., 0., 0.]) ## scale initial conditons pscale = np.max(np.abs(R0)) tscale = pscale / np.max(np.abs(V0)) IC['R0'] = R0 / pscale IC['V0'] = V0 * tscale / pscale IC['ag'] = np.array([0., 0., -1.62]) * tscale**2 / pscale xi, it, time = nlls.run(xi, IC) ## CONSTRUCT SOLUTION **********************************************
# !!! 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ₓₓ + δ yₓ + α y + β y^3 - γ cos(ω x) = 0 L = jit(lambda xi,IC: ypp(x,xi,IC) + delta*yp(x,xi,IC) + alpha*y(x,xi,IC) + beta*y(x,xi,IC)**3 \ - gamma*np.cos(omega*x)) ## construct the least-squares class: ************************************************************** xi0 = np.zeros(H(x).shape[1]) IC = {'y0': np.array([y0]), 'y0p': np.array([y0p])} nlls = NllsClass(xi0, L, IC, timer=True, tol=tfcTol) ## initialize dictionary to record solution: ******************************************************* xSol = onp.zeros((Nstep, N)) ySol = onp.zeros_like(xSol) res = onp.zeros_like(xSol) time = onp.zeros(Nstep) xSol[0, :] = x[:-1] xFinal = x[-1] ## 'propagation' loop: ***************************************************************************** for i in tqdm.trange(Nstep): xi, it, time[i] = nlls.run(xi0, IC) # print solution to dictionary if i > 0:
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)) qo = onp.expand_dims(onp.interp(x,xp,init['q'].flatten()),1) xi['xiQ'] = onp.dot(onp.linalg.pinv(jacfwd(q,1)(x,xi,const)['xiQ']),qo.flatten()-q(x,xi,const)) # Create NLLS class nlls = NllsClass(xi,L,x,const,tol=tol,maxIter=maxIter,timer=True,constant_arg_nums=[1]) # Run the NLLS xi,_,time[k] = nlls.run(xi,const) # Plots and results print("Altitude: "+str(atmData['Alt (km)'][k])+"km \t Norm of the residual: "+str(np.linalg.norm(L(xi,x,const)))) th1 = np.linspace(0.,xi['beta'],num=100).flatten() x1 = -Rs*np.sin(th1) y1 = Rs*(1.-np.cos(th1)) x2 = -r(x,xi) y2 = z(x,xi) X = np.hstack([x1,x2]) Y = np.hstack([y1,y2])
'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 start = timer() sol = optim.differential_evolution(InnerLoop, [(0, 15)],
c = 2. * onp.ones(1) else: c = 10. * onp.ones(1) xi = TFCDictRobust({'xis':xis,\ 'xic':xic}) IC = {'R0': np.zeros((3,)), \ 'V0': np.zeros((3,)), \ 'ag': np.zeros((3,)), \ 'Gam': np.zeros((1,)), \ 'c': 2.*onp.ones(1)} ## NONLINEAR LEAST-SQUARES CLASS ***************************************************************************** nlls = NllsClass(xi, L, IC, maxIter=2, timer=True) R0 = np.array([500000., 100000., 50000.]) V0 = np.array([-3000., 0., 0.]) ## scale initial conditons pscale = np.max(np.abs(R0)) tscale = pscale / np.max(np.abs(V0)) IC['R0'] = R0 / pscale IC['V0'] = V0 * tscale / pscale IC['ag'] = np.array([0., 0., -5.314961]) * tscale**2 / pscale IC['Gam'] = Gam * tscale**4 / pscale**2 global it it = 0