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