<jupyter_start><jupyter_text>Check we reproduce the same answer as galpy for some random potential<jupyter_code>q=0.5

#Getting units correct is painful. with ro=1 vo=1 and turn_physical_off then everything should be just G=1
galpy_pot = galpy.potential.TwoPowerTriaxialPotential(c=q,ro=1,vo=1)
galpy_pot.turn_physical_off()
pot = astro_dynamo.analytic_potentials.SpheroidalPotential(lambda m: galpy_pot._amp*galpy_pot._mdens(m),q=q)
x=np.linspace(0,10,100)
plt.semilogy(x,list(map(lambda x: -galpy_pot.Rforce(x,1),x)),'r',label='galpy FR')
plt.semilogy(x,-pot.f_r_cyl(x,np.array([1.])),'--k')
plt.semilogy(x,list(map(lambda x: -galpy_pot.zforce(x,1),x)),'y',label='galpy Fz')
plt.semilogy(x,-pot.f_z(x,np.array([1.])),'--k',label='astro-dynamo')
plt.legend()
plt.ylabel('Force')
plt.xlabel('R')
x=np.linspace(0,10,100)
plt.plot(x,list(map(lambda x: galpy_pot.vcirc(x,0),x)),'r',label='galpy FR')
plt.plot(x,torch.sqrt(pot.vc2(x,np.array([0.]))),'--k',label='astro-dynamo')
plt.ylabel('$V_c$')
plt.xlabel('$R$')<jupyter_output><empty_output><jupyter_text>Try replacing the dark matter particles in a snapshot by an analytic profile<jupyter_code>snap=astro_dynamo.snap.SnapShot('../inputmodels/M85_0.gz',
              particle_type_mapping={0:ParticleType.DarkMatter,1:ParticleType.Star})
q,qerr = astro_dynamo.analytic_potentials.fit_q_to_snapshot(snap,plot=True,r_bins=50)
print(f'q={q:.3f}+-{qerr:.3f}')<jupyter_output><empty_output><jupyter_text>Define and test a spheriodal potential based on this fit<jupyter_code>def ein(m,rhor0,m0,alpha):
    rho0 = rhor0 / (np.exp(-(2 / alpha) * ((8.2 / m0) ** alpha - 1)))
    return rho0 * np.exp(-(2 / alpha) * ((m / m0) ** alpha - 1))

pot = astro_dynamo.analytic_potentials.fit_potential_to_snap(snap.dm,ein,init_parms=[1e-3,8.0,0.7],plot=True)
r,dm_vc2 = mwtools.nemo.rotationcurve(snap.dm.as_numpy_array(),rrange=(0, 40))
r,stellar_vc2 = mwtools.nemo.rotationcurve(snap.stars.as_numpy_array(),rrange=(0, 40))

i = (np.abs(snap.dm.positions[:,0]) < 10) & \
    (np.abs(snap.dm.positions[:,1]) < 10) & \
    (np.abs(snap.dm.positions[:,2]) < 10)
r,dm_vc2_trunc = mwtools.nemo.rotationcurve(snap.dm[i].as_numpy_array(),rrange=(0, 40))

i = (np.abs(snap.stars.positions[:,0]) < 10) & \
    (np.abs(snap.stars.positions[:,1]) < 10) & \
    (np.abs(snap.stars.positions[:,2]) < 10)
r,stellar_vc2_trunc = mwtools.nemo.rotationcurve(snap.stars[i].as_numpy_array(),rrange=(0, 40))

f,ax = plt.subplots(1,1)
ax.plot(r,np.sqrt(dm_vc2),label = 'DM Particles')
ax.plot(r,np.sqrt(stellar_vc2),label = 'Stellar Particles')
ax.plot(r,np.sqrt(dm_vc2_trunc),label = 'DM Particles in 10kpc box')
x=np.linspace(0.,40,100)
ax.plot(x,np.sqrt(pot.vc2(x,torch.tensor(0.0,dtype=torch.float64))),label = 'Einasto Fit')

r=r.copy()
ax.plot(r,np.sqrt(stellar_vc2+pot.vc2(r,torch.tensor(0.0,dtype=torch.float64)).numpy()),label = 'Total Vc: Einasto Fit')
ax.plot(r,np.sqrt(stellar_vc2+dm_vc2),label = 'Total Vc: Particles')
ax.set_xlim((0,20))
ax.set_ylabel('$V_c$')
ax.set_xlabel('$R$')

ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))<jupyter_output><empty_output><jupyter_text>Test the gridding of the potential<jupyter_code>pot.grid_acc()

maxi=1000
positions = snap.stars.positions
r_cyl = snap.stars.rcyl
z = snap.stars.positions[..., 2]

f_r_cyl,f_z = pot.get_accelerations_cyl(positions[:maxi,:]).t()

f,ax = plt.subplots(1,2, figsize = (8,4), sharey = 'row')

ax[0].plot(r_cyl[:maxi],np.abs((pot.f_r_cyl(r_cyl[:maxi],z[:maxi])-f_r_cyl)/f_r_cyl),'.',label='$F_r$')
ax[0].plot(r_cyl[:maxi],np.abs((pot.f_z(r_cyl[:maxi],z[:maxi])-f_z)/f_z),'.',label='$F_z$')
ax[0].semilogy()
ax[0].legend()
ax[0].set_ylabel('Fractional Difference')
ax[0].set_xlabel('R')

ax[1].plot(z[:maxi],np.abs((pot.f_r_cyl(r_cyl[:maxi],z[:maxi])-f_r_cyl)/f_r_cyl),'.',label='$F_r$')
ax[1].plot(z[:maxi],np.abs((pot.f_z(r_cyl[:maxi],z[:maxi])-f_z)/f_z),'.',label='$F_z$')
ax[1].semilogy()
ax[1].legend()
ax[1].set_xlabel('z')

maxi=1000
positions = snap.stars.positions
r_cyl = snap.stars.rcyl
z = snap.stars.positions[..., 2]

acc = pot.get_accelerations(positions)
f_r_cyl = -torch.sqrt( acc[..., 0]**2 + acc[..., 1]**2 )
f_z = acc[..., 2]
f_r_cyl=f_r_cyl[:maxi]
f_z=f_z[:maxi]

f,ax = plt.subplots(1,2, figsize = (8,4), sharey = 'row')

ax[0].plot(r_cyl[:maxi],np.abs((pot.f_r_cyl(r_cyl[:maxi],z[:maxi])-f_r_cyl)/f_r_cyl),'.',label='$F_r$')
ax[0].plot(r_cyl[:maxi],np.abs((pot.f_z(r_cyl[:maxi],z[:maxi])-f_z)/f_z),'.',label='$F_z$')
ax[0].semilogy()
ax[0].legend()
ax[0].set_ylabel('Fractional Difference')
ax[0].set_xlabel('R')

ax[1].plot(z[:maxi],np.abs((pot.f_r_cyl(r_cyl[:maxi],z[:maxi])-f_r_cyl)/f_r_cyl),'.',label='$F_r$')
ax[1].plot(z[:maxi],np.abs((pot.f_z(r_cyl[:maxi],z[:maxi])-f_z)/f_z),'.',label='$F_z$')
ax[1].semilogy()
ax[1].legend()
ax[1].set_xlabel('z')
gpu_pot = pot.to('cuda')
acc = gpu_pot.get_accelerations(positions)
f_r_cyl = -torch.sqrt( acc[..., 0]**2 + acc[..., 1]**2 )
f_z = acc[..., 2]
f_r_cyl=f_r_cyl[:maxi]
f_z=f_z[:maxi]

f,ax = plt.subplots(1,2, figsize = (8,4), sharey = 'row')

ax[0].plot(r_cyl[:maxi],np.abs((pot.f_r_cyl(r_cyl[:maxi],z[:maxi])-f_r_cyl)/f_r_cyl),'.',label='$F_r$')
ax[0].plot(r_cyl[:maxi],np.abs((pot.f_z(r_cyl[:maxi],z[:maxi])-f_z)/f_z),'.',label='$F_z$')
ax[0].semilogy()
ax[0].legend()
ax[0].set_ylabel('Fractional Difference')
ax[0].set_xlabel('R')

ax[1].plot(z[:maxi],np.abs((pot.f_r_cyl(r_cyl[:maxi],z[:maxi])-f_r_cyl)/f_r_cyl),'.',label='$F_r$')
ax[1].plot(z[:maxi],np.abs((pot.f_z(r_cyl[:maxi],z[:maxi])-f_z)/f_z),'.',label='$F_z$')
ax[1].semilogy()
ax[1].legend()
ax[1].set_xlabel('z')<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup