def sigp_prop(r1,r2,r3,nu,Sigfunc,M,Mcentral,beta,betaf,nupars,Mpars,\ betpars,\ Mstar_rad,Mstar_prof,Mstar,Arot,Rhalf,G,rmin,rmax): #Calculate projected velocity dispersion profiles #given input *functions* nu(r); M(r); beta(r); betaf(r). #Also input is an array Mstar_prof(Mstar_rad) describing the 3D #cumulative stellar mass profile. This should be normalised #so that it peaks at 1.0. The total stellar mass is passed in Mstar. #Set up theta integration array: intpnts = np.int(150) thmin = 0. bit = 1.e-5 thmax = np.pi/2.-bit th = np.linspace(thmin,thmax,intpnts) sth = np.sin(th) cth = np.cos(th) cth2 = cth**2. rint = np.logspace(np.log10(rmin),np.log10(rmax),intpnts) sigr2 = np.zeros(len(rint)) nur = nu(rint,nupars) betafunc = betaf(rint,betpars,Rhalf,Arot) for i in range(len(rint)): rq = rint[i]/cth Mq = M(rq,Mpars)+Mcentral(rq,Mpars) if (Mstar > 0): Mq = Mq+Mstar*np.interp(rq,Mstar_rad,Mstar_prof) nuq = nu(rq,nupars) betafuncq = betaf(rq,betpars,Rhalf,Arot) sigr2[i] = 1./nur[i]/rint[i]/betafunc[i] * \ integrator(G*Mq*nuq*betafuncq*sth,th) Sig = Sigfunc(rint,nupars) sigLOS2 = np.zeros(len(rint)) sigpmr2 = np.zeros(len(rint)) sigpmt2 = np.zeros(len(rint)) for i in range(len(rint)): rq = rint[i]/cth nuq = nu(rq,nupars) sigr2q = np.interp(rq,rint,sigr2,left=0,right=0) betaq = beta(rq,betpars) sigLOS2[i] = 2.0*rint[i]/Sig[i]*\ integrator((1.0-betaq*cth2)*nuq*sigr2q/cth2,th) sigpmr2[i] = 2.0*rint[i]/Sig[i]*\ integrator((1.0-betaq+betaq*cth2)*nuq*sigr2q/cth2,th) sigpmt2[i] = 2.0*rint[i]/Sig[i]*\ integrator((1.0-betaq)*nuq*sigr2q/cth2,th) sigr2out = np.interp(r2,rint,sigr2,left=0,right=0) sigLOS2out = np.interp(r2,rint,sigLOS2,left=0,right=0) sigpmr2out = np.interp(r3,rint,sigpmr2,left=0,right=0) sigpmt2out = np.interp(r3,rint,sigpmt2,left=0,right=0) Sigout = np.interp(r1,rint,Sig,left=0,right=0) return sigr2out,Sigout,sigLOS2out,sigpmr2out,sigpmt2out
def mee_ocp_central_difference(x): """ Central different computation of the Jacobian of the cost funtion. Parameters: =========== tspan -- vector containting initial and final time p0 -- initial costates guess states0 -- initial state vector (modified equinoctial elements) rho -- switch smoothing parameter eclipse -- boolean (true or false) Returns: ======== Jac -- Jacobian External: ========= numpy, spipy.integrate """ # TEMP tspan = np.array([0, 642.54]) rho = 1 eclipse = False meef = np.array([6.610864583184714, 0.0, 0.0, 0.0, 0.0, 0.0]) # Initialization IC1 = np.zeros(14) IC2 = np.zeros(14) grad = np.zeros(7) DEL = 1e-5 * (x[7:14] / np.linalg.norm(x[7:14])) for i in range(0, 7): IC1 = x IC2 = x IC1[i + 7] = IC1[i + 7] + DEL[i] IC2[i + 7] = IC2[i + 7] - DEL[i] # Integrate Dynamics sol1 = integrator( lambda t, y: eom_mee_twobodyJ2_minfuel(t, y, rho, eclipse), tspan, IC1, method='LSODA', rtol=1e-13) sol2 = integrator( lambda t, y: eom_mee_twobodyJ2_minfuel(t, y, rho, eclipse), tspan, IC2, method='LSODA', rtol=1e-13) res1 = np.linalg.norm(meef[0:5] - sol1.y[0:5, -1]) res2 = np.linalg.norm(meef[0:5] - sol2.y[0:5, -1]) grad[i] = (res1 - res2) / (2 * DEL[i]) return grad
def surf_renorm(rbin,surfden): #Calculate the integral of the surface density #so that it can then be renormalised. #Calcualte also Rhalf. ranal = np.linspace(0,10,np.int(5000)) surfden_ranal = np.interp(ranal,rbin,surfden,left=0,right=0) Menc_tot = 2.0*np.pi*integrator(surfden_ranal*ranal,ranal) Menc_half = 0.0 i = 3 while (Menc_half < Menc_tot/2.0): Menc_half = 2.0*np.pi*\ integrator(surfden_ranal[:i]*ranal[:i],ranal[:i]) i = i + 1 Rhalf = ranal[i-1] return Rhalf, Menc_tot
def sidm_novel(rc,M200,c,oden,rhocrit): #Calculate SIDM model parameters from the coreNFWtides #model fit. For this to be valid, the coreNFWtides fit #should assume a pure-core model, with n=1. See #Read et al. 2018 for further details. #Returns cross section/particle mass in cm^2 / g. GammaX = 0.005/(1e9*year) Guse = G*Msun/kpc rho_unit = Msun/kpc**3.0 rc = np.abs(rc)*10.0 gcon=1./(np.log(1.+c)-c/(1.+c)) deltachar=oden*c**3.*gcon/3. rv=(3./4.*M200/(np.pi*oden*rhocrit))**(1./3.) rs=rv/c rhos=rhocrit*deltachar rhorc = rhoNFW(rc,rhos,rs) r = np.logspace(np.log10(rc),np.log10(rs*5000.0),50000) rho = rhoNFW(r,rhos,rs) mass = M200*gcon*(np.log(1.0 + r/rs)-r/rs/(1.0+r/rs)) sigvtworc = Guse/rhorc*integrator(mass*rho/r**2.0,r) sigvrc = np.sqrt(sigvtworc) sigm = np.sqrt(np.pi)*GammaX/(4.0*rhorc*rho_unit*sigvrc) return sigm*100.0**2.0/1000.0
def residuals(v0, tspan, states0, statesf): """ Computes error between current final state and desired final state for the Keplerian Lambert type problem. This is a numerical method not an analytic Lambert solver. Parameters: =========== v0 -- current initial velocity tspan -- vector containing initial and final time states0 -- initial states statesf -- final states Returns: ======== res -- residual External: ========= numpy, scipy.integrate """ states_in = np.zeros(6) states_in[0:3] = states0[0:3] states_in[3:6] = v0[0:3] sol = integrator(eom_twobody, (tspan[0], tspan[-1]), states_in, method='LSODA', rtol=1e-12) res = np.linalg.norm(statesf[0:3] - sol.y[0:3, -1]) / cn.DU return res
def residuals(p0): # Residual function # TEMP: For now I have copied these parameters into the function to make it more inline with the Rosenbrock # function above, but eventually we'll need to pass them as input arguments tspan = np.array([0, 642.54]) rho = 1 eclipse = False mee0 = np.array( [1.8225634499541166, 0.725, 0.0, 0.06116262015048431, 0.0, 0, 100]) meef = np.array([6.610864583184714, 0.0, 0.0, 0.0, 0.0, 0.0]) ## Computes the error between the current and desired final state. states_in = np.zeros(14) states_in[0:7] = mee0[0:7] states_in[7:14] = p0[0:7] # sol = integrator(eom_mee_twobody_minfuel,(tspan[0],tspan[-1]),states_in,method='LSODA',rtol=1e-12) sol = integrator( lambda t, y: eom_mee_twobodyJ2_minfuel(t, y, rho, eclipse), tspan, states_in, method='LSODA', rtol=1e-12) # Free final mass and free final true longitude (angle) res = np.linalg.norm(meef[0:5] - sol.y[0:5, -1]) print(res) return res
def residuals_mee_ocp(p0, tspan, mee0, meef, rho, eclipse): """ Computes error between current final state and desired final state for the low thrust, min-fuel, optimal trajectory using modified equinoctial elements. Parameters: =========== p0 -- current initial costate tspan -- vector containing initial and final time mee0 -- initial states (modified equinoctial elements) meef -- final states (modified equinoctial elements) rho -- switch smoothing parameter Returns: ======== res -- residual External: ========= numpy, scipy.integrate """ ## Computes the error between the current and desired final state. states_in = np.zeros(14) states_in[0:7] = mee0[0:7] states_in[7:14] = p0[0:7] # sol = integrator(eom_mee_twobody_minfuel,(tspan[0],tspan[-1]),states_in,method='LSODA',rtol=1e-12) sol = integrator( lambda t, y: eom_mee_twobodyJ2_minfuel(t, y, rho, eclipse), tspan, states_in, method='LSODA', rtol=1e-12) # Free final mass and free final true longitude (angle) res = np.linalg.norm(meef[0:5] - sol.y[0:5, -1]) print(res) return res
def residuals(v0,tspan,states0,statesf): ## Computes the error between the current and desired final state. states_in = np.zeros(6) states_in[0:3] = states0[0:3] states_in[3:6] = v0[0:3] sol = integrator(eom_twobody,(tspan[0],tspan[-1]),states_in,method='LSODA',rtol=1e-12) res = np.linalg.norm(statesf[0:3] - sol.y[0:3,-1])/cn.DU print(res) return res
def alpbetgamsigr(r,rho0s,r0s,alps,bets,gams,rho0,r0,alp,bet,gam,ra): nu = alpbetgamden(r,rho0s,r0s,alps,bets,gams) mass = alpbetgammass(r,rho0,r0,alp,bet,gam) gf = gfunc_osipkov(r,ra) sigr = np.zeros(len(r)) for i in range(len(r)-3): sigr[i] = 1.0/nu[i]/gf[i] * \ integrator(Guse*mass[i:]*nu[i:]/r[i:]**2.0*\ gf[i:],r[i:]) return sigr
def residuals(p0,tspan,states0,statesf,rho): ## Computes the error between the current and desired final state. states_in = np.zeros(14) states_in[0:7] = states0[0:7] states_in[7:14] = p0[0:7] # sol = integrator(eom_mee_twobody_minfuel,(tspan[0],tspan[-1]),states_in,method='LSODA',rtol=1e-12) sol = integrator(lambda t,y: eom_mee_twobody_minfuel(t,y,rho),tspan,states_in,method='LSODA',rtol=1e-12) # Free final mass and free final true longitude (angle) res = np.linalg.norm(statesf[0:5] - sol.y[0:5,-1]) # print(res) return res
def alpbetgamvsp(rho0s,r0s,alps,bets,gams,rho0,r0,alp,bet,gam,ra): intpnts = np.int(1e4) r = np.logspace(np.log10(r0s/50.0),np.log10(500.0*r0s),\ np.int(intpnts)) nu = alpbetgamden(r,rho0s,r0s,alps,bets,gams) massnu = alpbetgamden(r,rho0s,r0s,alps,bets,gams) mass = alpbetgammass(r,rho0,r0,alp,bet,gam) sigr = alpbetgamsigr(r,rho0s,r0s,alps,bets,gams,rho0,\ r0,alp,bet,gam,ra) bet = osipkov(r,ra) sigstar = np.zeros(len(r)) for i in range(1,len(r)-3): sigstar[i] = 2.0*integrator(nu[i:]*r[i:]/\ np.sqrt(r[i:]**2.0-r[i-1]**2.0),\ r[i:]) #Normalise similarly to the data: norm = integrator(sigstar*2.0*np.pi*r,r) nu = nu / norm sigstar = sigstar / norm #VSPs: vsp1 = \ integrator(2.0/5.0*Guse*mass*nu*(5.0-2.0*bet)*\ sigr*r,r)/1.0e12 vsp2 = \ integrator(4.0/35.0*Guse*mass*nu*(7.0-6.0*bet)*\ sigr*r**3.0,r)/1.0e12 #Richardson & Fairbairn zeta parameters: Ntotuse = integrator(sigstar*r,r) sigint = integrator(sigstar*r**3.0,r) zeta_A = 9.0/10.0*Ntotuse*integrator(Guse*mass*nu*(\ 5.0-2.0*bet)*sigr*r,r)/\ (integrator(Guse*mass*nu*r,r))**2.0 zeta_B = 9.0/35.0*Ntotuse**2.0*\ integrator(Guse*mass*nu*(7.0-6.0*bet)*sigr*r**3.0,r)/\ ((integrator(Guse*mass*nu*r,r))**2.0*sigint) return vsp1, vsp2, zeta_A, zeta_B
def Rhalf_func(M1,M2,M3,a1,a2,a3): #Calculate projected half light radius for #the threeplum model: ranal = np.logspace(-3,1,np.int(500)) Mstar_surf = threeplumsurf(ranal,M1,M2,M3,a1,a2,a3) Menc_half = 0.0 i = 3 while (Menc_half < (M1+M2+M3)/2.0): Menc_half = 2.0*np.pi*\ integrator(Mstar_surf[:i]*ranal[:i],ranal[:i]) i = i + 1 Rhalf = ranal[i-1] return Rhalf
def mps_twobody(tspan, states0, statesf): # Method of particular solutions shooting method to solve Lambert's problem mps_err = 1 mps_tol = 1e-13 cnt = 0 IC = np.zeros(6) Del = np.zeros((3, 3)) DEL = 1e-5 v0 = states0[3:6] MatInv = np.zeros((3, 3)) res = np.zeros(3) del_v0 = np.zeros(3) while mps_err > mps_tol and cnt < 15: for tcnt in range(0, 4): if tcnt == 0: IC[0:3] = states0[0:3] IC[3:6] = v0[0:3] elif tcnt > 0: IC[0:3] = states0[0:3] IC[3:6] = v0[0:3] IC[tcnt + 2] = IC[tcnt + 2] + DEL # Integrate Dynamics sol = integrator(eom_twobody, (tspan[0], tspan[-1]), IC, method='LSODA', rtol=1e-12) soln = sol.y[0:3, -1] if tcnt == 0: ref = soln elif tcnt > 0: Del[0:3, tcnt - 1] = (soln[0:3] - ref[0:3]) / DEL # Compute Updated Variables MatInv = np.linalg.inv(Del) res = statesf[0:3] - ref[0:3] del_v0 = np.matmul(MatInv, res) v0 = v0 + del_v0 # Compute Error mps_err = np.linalg.norm(res) / cn.DU # print(mps_err) # Counter cnt = cnt + 1 return v0, mps_err
def mps_mee_ocp(tspan, p0, states0, statesf, rho, mps_tol): # Method of particular solutions to solve Lambert's problem mps_err = 1 cnt = 0 IC = np.zeros(14) Del = np.zeros((5, 7)) Del_plus = np.zeros((5, 7)) Del_minus = np.zeros((5, 7)) DEL = np.array([1e-5, 1e-5, 1e-5, 1e-5, 1e-5, 1e-5, 1e-5]) # p0 = states0[7:14] MatInv = np.zeros((7, 7)) res = np.zeros(7) del_p0 = np.zeros(7) while mps_err > mps_tol and cnt < 30: for tcnt in range(0, 15): if tcnt == 0: IC[0:7] = states0[0:7] IC[7:14] = p0[0:7] elif tcnt > 0 and tcnt < 8: IC[0:7] = states0[0:7] IC[7:14] = p0[0:7] IC[tcnt + 6] = IC[tcnt + 6] + DEL[tcnt - 1] elif tcnt > 7 and tcnt < 15: IC[0:7] = states0[0:7] IC[7:14] = p0[0:7] IC[tcnt + 6 - 7] = IC[tcnt + 6 - 7] - DEL[tcnt - 1 - 7] # Integrate Dynamics sol = integrator(lambda t, y: eom_mee_twobody_minfuel(t, y, rho), tspan, IC, method='LSODA', rtol=1e-13) # sol = integrator(eom_mee_twobody_minfuel,(tspan[0],tspan[-1]),IC,method='LSODA',rtol=1e-12) soln = sol.y[0:7, -1] if tcnt == 0: ref = soln traj = sol elif tcnt > 0 and tcnt < 8: # Del_plus[0:5,tcnt-1] = (soln[0:5]-ref[0:5])/DEL # Free final mass & free final true longitude (angle) Del_plus[0:5, tcnt - 1] = soln[ 0:5] # Free final mass & free final true longitude (angle) elif tcnt > 7 and tcnt < 15: Del_minus[0:5, tcnt - 1 - 7] = soln[ 0:5] # Free final mass & free final true longitude (angle) # Compute Updated Variables for i in range(5): for j in range(7): Del[i, j] = (Del_plus[i, j] - Del_minus[i, j]) / (2 * DEL[j]) MatInv = np.linalg.pinv(Del) res = statesf[0:5] - ref[ 0:5] # Free final mass & free final true longitude (angle) del_p0 = np.matmul(MatInv, res) p0 = p0 + del_p0 DEL = 1e-5 * (p0 / np.linalg.norm(p0)) # Compute Error mps_err = np.linalg.norm(res) print(mps_err) # Counter cnt = cnt + 1 return traj, p0, mps_err
# Scipy's Optimize v0_guess = states0[3:6] optres = minimize(residuals, v0_guess[0:3], args=(tspan, states0, statesf), method='Nelder-Mead', tol=1e-13) states0_new = np.zeros(6) states0_new[0:3] = states0[0:3] states0_new[3:6] = optres.x[0:3] print("opt_err:", optres.fun) # Check Solution sol = integrator(eom_twobody, (tspan[0], tspan[-1]), states0_new, method='LSODA', rtol=1e-12) # Plot fig = plt.figure() ax = plt.axes(projection='3d') ax.plot3D(sol.y[0, :], sol.y[1, :], sol.y[2, :]) ax.scatter3D(states0[0], states0[1], states0[2], 'r') ax.scatter3D(statesf[0], statesf[1], statesf[2], 'g') plt.show() # import pandas as pd # statesdf=pd.DataFrame([states0,statesf] ,columns=['x','y','z','vx','vy','vz']) # statesdf['cat']=['s','f'] # soldf=pd.DataFrame(sol.y.T ,columns=['x','y','z','vx','vy','vz']) # soldf['cat']='o'
def velfit_full(R,vz,vzerr,ms,Nbin,\ vfitmin,vfitmax,\ p0vin_min,p0vin_max,p0best,\ nsamples,outfile,nprocs): #Code to fit the velocity data with #the velpdf model. #Arrays to store binned: #radius (rbin), #<vlos> (vzmeanbin), #<vlos^2>^1/2 (vztwobin), #<vlos^4> (vzfourbin), #and their confidence intervals; and #the mean and dispersion of the background #model. rbin = np.zeros(len(R)) right_bin_edge = np.zeros(len(R)) vzmeanbin = np.zeros(len(R)) vzmeanbinlo = np.zeros(len(R)) vzmeanbinhi = np.zeros(len(R)) vztwobin = np.zeros(len(R)) vztwobinlo = np.zeros(len(R)) vztwobinhi = np.zeros(len(R)) vzfourbin = np.zeros(len(R)) vzfourbinlo = np.zeros(len(R)) vzfourbinhi = np.zeros(len(R)) backampbin = np.zeros(len(R)) backampbinlo = np.zeros(len(R)) backampbinhi = np.zeros(len(R)) backmeanbin = np.zeros(len(R)) backmeanbinlo = np.zeros(len(R)) backmeanbinhi = np.zeros(len(R)) backsigbin = np.zeros(len(R)) backsigbinlo = np.zeros(len(R)) backsigbinhi = np.zeros(len(R)) #This for storing the vzfour pdf #for calculating the VSPs: vzfour_pdf = np.zeros((nsamples, len(R))) #Loop through the bins, assuming Nbin stars #(weighted by ms) per bin: index = np.argsort(R) vzstore = np.zeros(len(R)) vzerrstore = np.zeros(len(R)) msstore = np.zeros(len(R)) cnt = 0 jsum = 0 js = 0 for i in range(len(R)): #Find stars in bin: if (jsum < Nbin): vzstore[js] = vz[index[i]] vzerrstore[js] = vzerr[index[i]] msstore[js] = ms[index[i]] right_bin_edge[cnt] = R[index[i]] jsum = jsum + ms[index[i]] js = js + 1 if (jsum >= Nbin): #Fit the velpdf model to these stars: vzuse = vzstore[:js] vzerruse = vzerrstore[:js] msuse = msstore[:js] #Cut back to fit range. If doing this, #perform an initial centre for this #bin first to ensure a symmetric #"haircut" on the data. if (vfitmin != 0 or vfitmax != 0): vzuse = vzuse - \ np.sum(vzuse*msuse)/np.sum(msuse) if (vfitmin != 0): vzuse_t = vzuse[vzuse > vfitmin] vzerruse_t = vzerruse[vzuse > vfitmin] msuse_t = msuse[vzuse > vfitmin] else: vzuse_t = vzuse vzerruse_t = vzerruse msuse_t = msuse if (vfitmax != 0): vzuse = vzuse_t[vzuse_t < vfitmax] vzerruse = vzerruse_t[vzuse_t < vfitmax] msuse = msuse_t[vzuse_t < vfitmax] else: vzuse = vzuse_t vzerruse = vzerruse_t msuse = msuse_t vzmeanbin[cnt],vzmeanbinlo[cnt],vzmeanbinhi[cnt],\ vztwobin[cnt],vztwobinlo[cnt],vztwobinhi[cnt],\ vzfourbin[cnt],vzfourbinlo[cnt],vzfourbinhi[cnt],\ backampbin[cnt],backampbinlo[cnt],backampbinhi[cnt],\ backmeanbin[cnt],backmeanbinlo[cnt],backmeanbinhi[cnt],\ backsigbin[cnt],backsigbinlo[cnt],backsigbinhi[cnt],\ vzfour_store,p0vbest = \ velfitbin(vzuse,vzerruse,msuse,\ p0vin_min,p0vin_max,nsamples,nprocs) vzfour_pdf[:, cnt] = vzfour_store #Calculate non-para versions for comparison: vztwo_nonpara = \ (np.sum(vzuse**2.0*msuse)-np.sum(vzerruse**2.0*msuse))/np.sum(msuse) vztwo_nonpara = np.sqrt(vztwo_nonpara) vzfour_nonpara = \ (np.sum(vzuse**4.0*msuse)-np.sum(3.0*vzerruse**4.0*msuse))/np.sum(msuse) #Make a plot of the fit: fig = plt.figure() ax = fig.add_subplot(111) for axis in ['top', 'bottom', 'left', 'right']: ax.spines[axis].set_linewidth(mylinewidth) ax.minorticks_on() ax.tick_params('both', length=10, width=2, which='major') ax.tick_params('both', length=5, width=1, which='minor') plt.xticks(fontsize=myfontsize) plt.yticks(fontsize=myfontsize) plt.xlabel(r'v\,[km/s]', fontsize=myfontsize) plt.ylabel(r'frequency',\ fontsize=myfontsize) n, bins, patches = plt.hist(vzuse,10,weights=msuse,\ facecolor='g',\ alpha=0.75) vplot = np.linspace(-50, 50, np.int(500)) vperr = np.zeros(len(vplot))+\ np.sum(vzerruse*msuse)/np.sum(msuse) pdf = velpdfuse(vplot, vperr, p0vbest) plt.plot(vplot,pdf/np.max(pdf)*np.max(n),\ linewidth=mylinewidth) plt.xlim([-50, 50]) plt.savefig(outfile+'hist_%d.pdf' % (cnt),\ bbox_inches='tight') #Calculate bin radius: if (cnt == 0): rbin[cnt] = right_bin_edge[cnt] / 2.0 else: rbin[cnt] = \ (right_bin_edge[cnt] + right_bin_edge[cnt-1])/2.0 #Output the fit values (with non-para comparison): print('Bin: %d | radius: %f | vztwo %.2f(%.2f)+%.2f-%.2f | vzfour %.2f(%.2f)+%.2f-%.2f' \ % (cnt,rbin[cnt],\ vztwobin[cnt],vztwo_nonpara,\ vztwobinhi[cnt]-vztwobin[cnt],\ vztwobin[cnt]-vztwobinlo[cnt],\ vzfourbin[cnt],vzfour_nonpara,\ vzfourbinhi[cnt]-vzfourbin[cnt],\ vzfourbin[cnt]-vzfourbinlo[cnt]),\ ) #Move on to the next bin: jsum = 0.0 js = 0 cnt = cnt + 1 #Cut back the output arrays: rbin = rbin[:cnt] vzmeanbin = vzmeanbin[:cnt] vzmeanbinlo = vzmeanbinlo[:cnt] vzmeanbinhi = vzmeanbinhi[:cnt] vztwobin = vztwobin[:cnt] vztwobinlo = vztwobinlo[:cnt] vztwobinhi = vztwobinhi[:cnt] vzfourbin = vzfourbin[:cnt] vzfourbinlo = vzfourbinlo[:cnt] vzfourbinhi = vzfourbinhi[:cnt] backampbin = backampbin[:cnt] backampbinlo = backampbinlo[:cnt] backampbinhi = backampbinhi[:cnt] backmeanbin = backmeanbin[:cnt] backmeanbinlo = backmeanbinlo[:cnt] backmeanbinhi = backmeanbinhi[:cnt] backsigbin = backsigbin[:cnt] backsigbinlo = backsigbinlo[:cnt] backsigbinhi = backsigbinhi[:cnt] #Calculate the VSPs with uncertainties. This #assumes negligible error in the surface density #profile as compared to the velocity uncertainties. #This is usually fine, but something to bear in mind. ranal = np.logspace(-5, 3, np.int(5e4)) surfden = threeplumsurf(ranal,p0best[0],p0best[1],p0best[2],\ p0best[3],p0best[4],p0best[5]) #This assumes a flat or linearly falling relation #beyond the last data point: vsp1 = np.zeros(nsamples) vsp2 = np.zeros(nsamples) vsp1_int = np.zeros(7) vsp2_int = np.zeros(7) vzfourstore = np.zeros((nsamples, len(ranal))) for i in range(nsamples): vzfour_thissample = vzfour_pdf[i, :cnt] vzfour = vzfourfunc(ranal, rbin, vzfour_thissample) vzfourstore[i, :] = vzfour vsp1[i] = integrator(surfden * vzfour * ranal, ranal) vsp2[i] = integrator(surfden * vzfour * ranal**3.0, ranal) vsp1_int[0], vsp1_int[1], vsp1_int[2], vsp1_int[3], \ vsp1_int[4], vsp1_int[5], vsp1_int[6] = \ calcmedquartnine(vsp1) vsp2_int[0], vsp2_int[1], vsp2_int[2], vsp2_int[3], \ vsp2_int[4], vsp2_int[5], vsp2_int[6] = \ calcmedquartnine(vsp2) return rbin,vzmeanbin,vzmeanbinlo,vzmeanbinhi,\ vztwobin,vztwobinlo,vztwobinhi,\ vzfourbin,vzfourbinlo,vzfourbinhi,\ backampbin,backampbinlo,backampbinhi,\ backmeanbin,backmeanbinlo,backmeanbinhi,\ backsigbin,backsigbinlo,backsigbinhi,\ vsp1_int[0],vsp1_int[1],vsp1_int[2],\ vsp2_int[0],vsp2_int[1],vsp2_int[2],\ ranal,vzfourstore,vsp1,vsp2
def velfit_easy(R,vz,vzerr,ms,Nbin,\ vfitmin,vfitmax,\ p0vin_min,p0vin_max,p0best,\ nsamples,outfile,nprocs): #Uses easy non-para estimates rather than VDF fit #**only use for rapid testing; not recommended #for real analysis**. Assumes only Poisson error #on each bin. Not really correct. rbin = np.zeros(len(R)) right_bin_edge = np.zeros(len(R)) vzmeanbin = np.zeros(len(R)) vzmeanbinlo = np.zeros(len(R)) vzmeanbinhi = np.zeros(len(R)) vztwobin = np.zeros(len(R)) vztwobinlo = np.zeros(len(R)) vztwobinhi = np.zeros(len(R)) vzfourbin = np.zeros(len(R)) vzfourbinlo = np.zeros(len(R)) vzfourbinhi = np.zeros(len(R)) backampbin = np.zeros(len(R)) backampbinlo = np.zeros(len(R)) backampbinhi = np.zeros(len(R)) backmeanbin = np.zeros(len(R)) backmeanbinlo = np.zeros(len(R)) backmeanbinhi = np.zeros(len(R)) backsigbin = np.zeros(len(R)) backsigbinlo = np.zeros(len(R)) backsigbinhi = np.zeros(len(R)) #This for storing the vzfour pdf #for calculating the VSPs: vzfour_pdf = np.zeros((nsamples, len(R))) #Loop through the bins, assuming Nbin stars #(weighted by ms) per bin: index = np.argsort(R) vzstore = np.zeros(len(R)) vzerrstore = np.zeros(len(R)) msstore = np.zeros(len(R)) cnt = 0 jsum = 0 js = 0 for i in range(len(R)): #Find stars in bin: if (jsum < Nbin): vzstore[js] = vz[index[i]] vzerrstore[js] = vzerr[index[i]] msstore[js] = ms[index[i]] right_bin_edge[cnt] = R[index[i]] jsum = jsum + ms[index[i]] js = js + 1 if (jsum >= Nbin): #Non-para estimate of vel moments for these stars: vzuse = vzstore[:js] vzerruse = vzerrstore[:js] msuse = msstore[:js] #Cut back to fit range. If doing this, #perform an initial centre for this #bin first to ensure a symmetric #"haircut" on the data. if (vfitmin != 0 or vfitmax != 0): vzuse = vzuse - \ np.sum(vzuse*msuse)/np.sum(msuse) if (vfitmin != 0): vzuse_t = vzuse[vzuse > vfitmin] vzerruse_t = vzerruse[vzuse > vfitmin] msuse_t = msuse[vzuse > vfitmin] else: vzuse_t = vzuse vzerruse_t = vzerruse msuse_t = msuse if (vfitmax != 0): vzuse = vzuse_t[vzuse_t < vfitmax] vzerruse = vzerruse_t[vzuse_t < vfitmax] msuse = msuse_t[vzuse_t < vfitmax] else: vzuse = vzuse_t vzerruse = vzerruse_t msuse = msuse_t #Non-param estimates assuming Poisson errors: #**only use for testing; not recommended for real analysis** vztwobin[cnt] = \ (np.sum(vzuse**2.0*msuse)-np.sum(vzerruse**2.0*msuse))/np.sum(msuse) vztwobin[cnt] = np.sqrt(vztwobin[cnt]) vztwobinlo[cnt] = vztwobin[cnt] - vztwobin[cnt] / np.sqrt(Nbin) vztwobinhi[cnt] = vztwobin[cnt] + vztwobin[cnt] / np.sqrt(Nbin) vzfourbin[cnt] = \ (np.sum(vzuse**4.0*msuse)-np.sum(3.0*vzerruse**4.0*msuse))/np.sum(msuse) vzfourbinlo[cnt] = vzfourbin[cnt] - vzfourbin[cnt] / np.sqrt(Nbin) vzfourbinhi[cnt] = vzfourbin[cnt] + vzfourbin[cnt] / np.sqrt(Nbin) #Assume vzfour distributed as Gaussian (again, for rapid testing only): vzfour_pdf[:,cnt] = np.random.normal(loc = vzfourbin[cnt],\ scale = vzfourbin[cnt]/np.sqrt(Nbin),\ size = nsamples) #Make a plot of the bin: fig = plt.figure() ax = fig.add_subplot(111) for axis in ['top', 'bottom', 'left', 'right']: ax.spines[axis].set_linewidth(mylinewidth) ax.minorticks_on() ax.tick_params('both', length=10, width=2, which='major') ax.tick_params('both', length=5, width=1, which='minor') plt.xticks(fontsize=myfontsize) plt.yticks(fontsize=myfontsize) plt.xlabel(r'v\,[km/s]', fontsize=myfontsize) plt.ylabel(r'frequency', fontsize=myfontsize) n, bins, patches = plt.hist(vzuse,10,weights=msuse,\ facecolor='g',\ alpha=0.75) plt.axvline(x=vztwobin[cnt],\ color='blue',linewidth=mylinewidth) plt.axvline(x=vztwobinlo[cnt],linestyle='dashed',\ color='blue',linewidth=mylinewidth) plt.axvline(x=vztwobinhi[cnt],linestyle='dashed',\ color='blue',linewidth=mylinewidth) plt.xlim([-50, 50]) plt.savefig(outfile+'hist_%d.pdf' % (cnt),\ bbox_inches='tight') #Calculate bin radius: if (cnt == 0): rbin[cnt] = right_bin_edge[cnt] / 2.0 else: rbin[cnt] = \ (right_bin_edge[cnt] + right_bin_edge[cnt-1])/2.0 #Output the fit values (with non-para comparison): print('Bin: %d | radius: %f | vztwo %.2f+%.2f-%.2f | vzfour %.2f+%.2f-%.2f' \ % (cnt,rbin[cnt],\ vztwobin[cnt],\ vztwobinhi[cnt]-vztwobin[cnt],\ vztwobin[cnt]-vztwobinlo[cnt],\ vzfourbin[cnt],\ vzfourbinhi[cnt]-vzfourbin[cnt],\ vzfourbin[cnt]-vzfourbinlo[cnt])) #Move on to the next bin: jsum = 0.0 js = 0 cnt = cnt + 1 #Cut back the output arrays: rbin = rbin[:cnt] vzmeanbin = vzmeanbin[:cnt] vzmeanbinlo = vzmeanbinlo[:cnt] vzmeanbinhi = vzmeanbinhi[:cnt] vztwobin = vztwobin[:cnt] vztwobinlo = vztwobinlo[:cnt] vztwobinhi = vztwobinhi[:cnt] vzfourbin = vzfourbin[:cnt] vzfourbinlo = vzfourbinlo[:cnt] vzfourbinhi = vzfourbinhi[:cnt] backampbin = backampbin[:cnt] backampbinlo = backampbinlo[:cnt] backampbinhi = backampbinhi[:cnt] backmeanbin = backmeanbin[:cnt] backmeanbinlo = backmeanbinlo[:cnt] backmeanbinhi = backmeanbinhi[:cnt] backsigbin = backsigbin[:cnt] backsigbinlo = backsigbinlo[:cnt] backsigbinhi = backsigbinhi[:cnt] #Calculate the VSPs with uncertainties. This #assumes negligible error in the surface density #profile as compared to the velocity uncertainties. #This is usually fine, but something to bear in mind. ranal = np.logspace(-5, 3, np.int(5e4)) surfden = threeplumsurf(ranal,p0best[0],p0best[1],p0best[2],\ p0best[3],p0best[4],p0best[5]) #This assumes a flat or linearly falling relation #beyond the last data point: vsp1 = np.zeros(nsamples) vsp2 = np.zeros(nsamples) vsp1_int = np.zeros(7) vsp2_int = np.zeros(7) vzfourstore = np.zeros((nsamples, len(ranal))) for i in range(nsamples): vzfour_thissample = vzfour_pdf[i, :cnt] vzfour = vzfourfunc(ranal, rbin, vzfour_thissample) vzfourstore[i, :] = vzfour vsp1[i] = integrator(surfden * vzfour * ranal, ranal) vsp2[i] = integrator(surfden * vzfour * ranal**3.0, ranal) vsp1_int[0], vsp1_int[1], vsp1_int[2], vsp1_int[3], \ vsp1_int[4], vsp1_int[5], vsp1_int[6] = \ calcmedquartnine(vsp1) vsp2_int[0], vsp2_int[1], vsp2_int[2], vsp2_int[3], \ vsp2_int[4], vsp2_int[5], vsp2_int[6] = \ calcmedquartnine(vsp2) return rbin,vzmeanbin,vzmeanbinlo,vzmeanbinhi,\ vztwobin,vztwobinlo,vztwobinhi,\ vzfourbin,vzfourbinlo,vzfourbinhi,\ backampbin,backampbinlo,backampbinhi,\ backmeanbin,backmeanbinlo,backmeanbinhi,\ backsigbin,backsigbinlo,backsigbinhi,\ vsp1_int[0],vsp1_int[1],vsp1_int[2],\ vsp2_int[0],vsp2_int[1],vsp2_int[2],\ ranal,vzfourstore,vsp1,vsp2
def sigp_vs(r1,r2,nu,Sigfunc,M,Mcentral,beta,betaf,nupars,Mpars,\ betpars,\ Mstar_rad,Mstar_prof,Mstar,Arot,Rhalf,G,rmin,rmax): #Calculate projected velocity dispersion profiles #given input *functions* nu(r); M(r); beta(r); betaf(r). #Also input is an array Mstar_prof(Mstar_rad) describing the 3D #cumulative stellar mass profile. This should be normalised #so that it peaks at 1.0. The total stellar mass is passed in Mstar. #Finally, the routine calculates a dimensional version of the #fourth order "virial shape" parmaeters in Richardson & Fairbairn 2014 #described in their equations 8 and 9. #Set up theta integration array: intpnts = np.int(150) thmin = 0. bit = 1.e-5 thmax = np.pi/2.-bit th = np.linspace(thmin,thmax,intpnts) sth = np.sin(th) cth = np.cos(th) cth2 = cth**2. #Set up rint interpolation array: rintpnts = np.int(150) rint = np.logspace(np.log10(rmin),\ np.log10(rmax),rintpnts) #First calc sigr2(rint): sigr2 = np.zeros(len(rint)) nur = nu(rint,nupars) betafunc = betaf(rint,betpars,Rhalf,Arot) for i in range(len(rint)): rq = rint[i]/cth Mq = M(rq,Mpars)+Mcentral(rq,Mpars) if (Mstar > 0): Mq = Mq+Mstar*np.interp(rq,Mstar_rad,Mstar_prof) nuq = nu(rq,nupars) betafuncq = betaf(rq,betpars,Rhalf,Arot) sigr2[i] = 1./nur[i]/rint[i]/betafunc[i] * \ integrator(G*Mq*nuq*betafuncq*sth,th) #And now the sig_LOS projection: Sig = Sigfunc(rint,nupars) sigLOS2 = np.zeros(len(rint)) for i in range(len(rint)): rq = rint[i]/cth nuq = nu(rq,nupars) sigr2q = np.interp(rq,rint,sigr2,left=0,right=0) betaq = beta(rq,betpars) sigLOS2[i] = 2.0*rint[i]/Sig[i]*\ integrator((1.0-betaq*cth2)*nuq*sigr2q/cth2,th) #And now the dimensional fourth order "virial shape" #parameters: betar = beta(rint,betpars) Mr = M(rint,Mpars)+Mstar*np.interp(rint,Mstar_rad,Mstar_prof) vs1 = 2.0/5.0*integrator(nur*(5.0-2.0*betar)*sigr2*\ G*Mr*rint,rint) vs2 = 4.0/35.0*integrator(nur*(7.0-6.0*betar)*sigr2*\ G*Mr*rint**3.0,rint) sigr2out = np.interp(r2,rint,sigr2,left=0,right=0) sigLOS2out = np.interp(r2,rint,sigLOS2,left=0,right=0) Sigout = np.interp(r1,rint,Sig,left=0,right=0) return sigr2out,Sigout,sigLOS2out,vs1,vs2
def mee_ocp_central_difference(p0): """ Central different computation of the Jacobian of the cost funtion. Parameters: =========== tspan -- vector containting initial and final time p0 -- initial costates guess states0 -- initial state vector (modified equinoctial elements) rho -- switch smoothing parameter eclipse -- boolean (true or false) Returns: ======== Jac -- Jacobian External: ========= numpy, spipy.integrate """ # TEMP tspan = np.array([0, 642.54]) rho = 1 eclipse = False states0 = np.array( [1.8225634499541166, 0.725, 0.0, 0.06116262015048431, 0.0, 0, 100]) statesf = np.array([6.610864583184714, 0.0, 0.0, 0.0, 0.0, 0.0]) # Initialization IC = np.zeros(14) Jac = np.zeros((14, 14)) Del_plus = np.zeros((7, 7)) Del_minus = np.zeros((7, 7)) DEL = 1e-5 * (p0 / np.linalg.norm(p0)) # DEL = np.array([1e-5,1e-5,1e-5,1e-5,1e-5,1e-5,1e-5]) for tcnt in range(0, 14): print(tcnt) if (tcnt < 7): IC[0:7] = states0[0:7] IC[7:14] = p0[0:7] IC[tcnt + 7] = IC[tcnt + 7] + DEL[tcnt] elif (tcnt >= 7): IC[0:7] = states0[0:7] IC[7:14] = p0[0:7] IC[tcnt] = IC[tcnt] - DEL[tcnt - 7] # Integrate Dynamics sol = integrator( lambda t, y: eom_mee_twobodyJ2_minfuel(t, y, rho, eclipse), tspan, IC, method='LSODA', rtol=1e-12) soln = sol.y[0:7, -1] if (tcnt < 7): Del_plus[0:7, tcnt] = soln[ 0:7] # Free final mass & free final true longitude (angle) elif (tcnt >= 7): Del_minus[0:7, tcnt - 7] = soln[ 0:7] # Free final mass & free final true longitude (angle) # Compute Jacobian (central differences) for i in range(7): for j in range(7): Jac[i + 7, j] = (Del_plus[i, j] - Del_minus[i, j]) / (2 * DEL[j]) # print(Jac) # sys.exit() # MatInv = np.linalg.pinv(Del) # res = statesf[0:5] - ref[0:5] # Free final mass & free final true longitude (angle) # del_p0 = np.matmul(MatInv,res) return Jac
def mps_twobody(tspan, states0, statesf): """ Shooting method to solve Lambert's problem Parameters: =========== tspan -- vector containing initial and final time states0 -- initial state vector (position & velocity) statesf -- final state vector (position & velocity) Returns: ======== v0 -- initial velocity for transfer trajectory mps_err -- shooting method convergence error External: ========= numpy, const.py, spipy.integrate """ mps_err = 1 mps_tol = 1e-13 cnt = 0 IC = np.zeros(6) Del = np.zeros((3, 3)) DEL = 1e-5 v0 = states0[3:6] MatInv = np.zeros((3, 3)) res = np.zeros(3) del_v0 = np.zeros(3) while mps_err > mps_tol and cnt < 15: for tcnt in range(0, 4): if tcnt == 0: IC[0:3] = states0[0:3] IC[3:6] = v0[0:3] elif tcnt > 0: IC[0:3] = states0[0:3] IC[3:6] = v0[0:3] IC[tcnt + 2] = IC[tcnt + 2] + DEL # Integrate Dynamics sol = integrator(eom_twobody, (tspan[0], tspan[-1]), IC, method='LSODA', rtol=1e-12) soln = sol.y[0:3, -1] if tcnt == 0: ref = soln elif tcnt > 0: Del[0:3, tcnt - 1] = (soln[0:3] - ref[0:3]) / DEL # Compute Updated Variables MatInv = np.linalg.inv(Del) res = statesf[0:3] - ref[0:3] del_v0 = np.matmul(MatInv, res) v0 = v0 + del_v0 # Compute Error mps_err = np.linalg.norm(res) / cn.DU # print(mps_err) # Counter cnt = cnt + 1 return v0, mps_err
print("rho:", rho, "mps err:", mps_err) rho = rho * rho_rate # Update iter = iter + 1 # Final Solution rho = rho / rho_rate # rho for converged solution mee0_new = np.zeros(14) mee0_new[0:7] = mee0[0:7] mee0_new[7:14] = p0[0:7] # Propagation Solution sol = integrator(lambda t, y: eom_mee_twobodyJ2_minfuel(t, y, rho, eclipse), tspan, mee0_new, method='LSODA', rtol=1e-13) data = np.array(soln.y, dtype=float).T time = np.array(soln.t, dtype=float).T # Plot Solution plot_mee_minfuel(time, data, rho, eclipse) # Save Solution df = pd.DataFrame(np.array([time, data]).T, columns=['time', 'trajectory']) df.to_csv("output_soln.csv", index=False) # Save Solution Parameters df = pd.DataFrame(np.array([rho, eclipse, sc.Isp, sc.A, sc.eff]).T, columns=[ 'rho', 'eclipse', 'Isp', 'Solar Array Area',
def mps_mee_ocp(tspan, p0, states0, statesf, rho, eclipse, mps_tol): """ Shooting method to solve min-fuel optimal control prolem Parameters: =========== tspan -- vector containing initial and final time p0 -- initial costates guess states0 -- initial state vector (modified equinoctial elements) statesf -- final state vector (modified equinoctial elements) rho -- switch smoothing parameter eclipse -- boolean (true or false) mps_tol -- user specified convergence tolerance Returns: ======== traj -- optimal trajectory p0 -- initial costate for optimal trajectory mps_err -- shooting method convergence error External: ========= numpy, const.py, spipy.integrate """ mps_err = 1 cnt = 0 IC = np.zeros(14) Del = np.zeros((5, 7)) Del_plus = np.zeros((5, 7)) Del_minus = np.zeros((5, 7)) DEL = np.array([1e-5, 1e-5, 1e-5, 1e-5, 1e-5, 1e-5, 1e-5]) # p0 = states0[7:14] MatInv = np.zeros((7, 7)) res = np.zeros(7) del_p0 = np.zeros(7) while mps_err > mps_tol and cnt < 30: for tcnt in range(0, 15): if tcnt == 0: IC[0:7] = states0[0:7] IC[7:14] = p0[0:7] elif tcnt > 0 and tcnt < 8: IC[0:7] = states0[0:7] IC[7:14] = p0[0:7] IC[tcnt + 6] = IC[tcnt + 6] + DEL[tcnt - 1] elif tcnt > 7 and tcnt < 15: IC[0:7] = states0[0:7] IC[7:14] = p0[0:7] IC[tcnt + 6 - 7] = IC[tcnt + 6 - 7] - DEL[tcnt - 1 - 7] # Integrate Dynamics sol = integrator( lambda t, y: eom_mee_twobodyJ2_minfuel(t, y, rho, eclipse), tspan, IC, method='LSODA', rtol=1e-13) # sol = integrator(eom_mee_twobody_minfuel,(tspan[0],tspan[-1]),IC,method='LSODA',rtol=1e-12) soln = sol.y[0:7, -1] if tcnt == 0: ref = soln traj = sol elif tcnt > 0 and tcnt < 8: # Del_plus[0:5,tcnt-1] = (soln[0:5]-ref[0:5])/DEL # Free final mass & free final true longitude (angle) Del_plus[0:5, tcnt - 1] = soln[ 0:5] # Free final mass & free final true longitude (angle) elif tcnt > 7 and tcnt < 15: Del_minus[0:5, tcnt - 1 - 7] = soln[ 0:5] # Free final mass & free final true longitude (angle) # Compute Updated Variables for i in range(5): for j in range(7): Del[i, j] = (Del_plus[i, j] - Del_minus[i, j]) / (2 * DEL[j]) MatInv = np.linalg.pinv(Del) res = statesf[0:5] - ref[ 0:5] # Free final mass & free final true longitude (angle) del_p0 = np.matmul(MatInv, res) p0 = p0 + del_p0 DEL = 1e-5 * (p0 / np.linalg.norm(p0)) # Compute Error mps_err = np.linalg.norm(res) print(mps_err) # Counter cnt = cnt + 1 return traj, p0, mps_err
def alpbetgammass(r,rho0,r0,alp,bet,gam): den = rho0*(r/r0)**(-gam)*(1.0+(r/r0)**alp)**((gam-bet)/alp) mass = np.zeros(len(r)) for i in range(3,len(r)): mass[i] = integrator(4.0*np.pi*r[:i]**2.*den[:i],r[:i]) return mass
index = np.argsort(surfbest) if (surfbest[index[0]] < 0): print('Warning! negative surface density at:',\ Rplot[index[0]],surfbest[index[0]]) #Numerical projection (test): intpnts = 250 theta = np.linspace(0, np.pi / 2.0 - 1e-6, num=intpnts) cth = np.cos(theta) cth2 = cth**2.0 surf = np.zeros(len(Rplot)) for i in range(len(Rplot)): q = Rplot[i] / cth y = threeplumden(q,p0best[0],p0best[1],p0best[2],\ p0best[3],p0best[4],p0best[5]) surf[i] = 2.0 * Rplot[i] * integrator(y / cth2, theta) plt.plot(Rplot,surf,color='orange',linewidth=2,linestyle='dashed',\ label='Best fit (numerical)') plt.axvline(x=Rhalf_int[0],color='blue',alpha=0.5,\ linewidth=mylinewidth) plt.xlim([xpltmin, xpltmax]) plt.ylim([surfpltmin, surfpltmax]) plt.legend(loc='upper right', fontsize=mylegendfontsize) plt.savefig(outfile + '_surfden.pdf', bbox_inches='tight') ##### Tracer 3D density profile (best fit) ##### fig = plt.figure(figsize=(figx, figy)) ax = fig.add_subplot(111)