def test_plot(self):
# Have to enable debug stuff in Cython file.
potential = sp.LeeSutoTriaxialNFWPotential(v_h=0.5, r_h=20., a=1., b=1., c=1., units=galactic)
ll = np.zeros((6000,self.nstars), dtype=float)
x,v = rewinder_likelihood(ll, -1., 6000, potential.c_instance,
self.prog, self.stars, 2.5E6, 0.,
1.25, self.betas, -0.3, True)
w = np.vstack((x.T,v.T)).T
fig = sd.plot_orbits(w, ix=0, marker=None, alpha=0.5)
plt.show()
def plot(self, **kwargs):
"""
TODO:
"""
if 'ls' not in kwargs and 'linestyle' not in kwargs:
kwargs['linestyle'] = 'none'
if 'marker' not in kwargs:
kwargs['marker'] = '.'
fig = gd.plot_orbits(self.X, **kwargs)
return fig
def make_orbit_files(potential, w0, N_max=6, suffix="", overwrite=False,
force_harmonic_oscillator=False):
orbit_filename = os.path.join(plot_path, "orbits{}.npy".format(suffix))
units = potential.units
nsteps = 100000
dt = 5.
#NT = 9*N_max**3/2
#every = nsteps // NT
if overwrite and os.path.exists(orbit_filename):
os.remove(orbit_filename)
toy_potential = None
if not os.path.exists(orbit_filename):
# integrate an orbit in a axisymmetric potential
logger.debug("Integrating orbit for {} steps...".format(nsteps))
acc = lambda t,w: np.hstack((w[...,3:],potential.acceleration(w[...,:3])))
integrator = si.DOPRI853Integrator(acc)
# acc = lambda t,x: potential.acceleration(x)
# integrator = si.LeapfrogIntegrator(acc)
t,w = integrator.run(w0, dt=dt, nsteps=nsteps)
logger.debug("...done!")
loop = sd.classify_orbit(w)
if np.any(loop == 1) and not force_harmonic_oscillator: # loop orbit
logger.debug("Orbit classified as LOOP")
m,b = sd.fit_isochrone(w, units=units)
toy_potential = sp.IsochronePotential(m=m, b=b, units=units)
else:
logger.debug("Orbit classified as BOX")
omegas = sd.fit_harmonic_oscillator(w, units=units)
toy_potential = sp.HarmonicOscillatorPotential(omega=omegas, units=units)
# also integrate the orbit in the best-fitting toy potential
toy_steps = w.shape[0]//10
logger.debug("Integrating toy orbit for {} steps...".format(toy_steps))
acc = lambda ts,ws: np.hstack((ws[...,3:],toy_potential.acceleration(ws[...,:3])))
integrator = si.DOPRI853Integrator(acc)
# acc = lambda t,x: toy_potential.acceleration(x)
# integrator = si.LeapfrogIntegrator(acc)
toy_t,toy_w = integrator.run(w0, dt=dt, nsteps=toy_steps)
logger.debug("...done!")
# cache the orbits
np.save(orbit_filename, (t,w,toy_t,toy_w))
logger.debug("Orbit computed and saved to file: {}".format(orbit_filename))
else:
t,w,toy_t,toy_w = np.load(orbit_filename)
logger.debug("Orbit read from file: {}".format(orbit_filename))
# for i in range(100):
# omega0 = np.random.uniform(0.,100.,size=3)
# omegas = sd.fit_harmonic_oscillator(w, units=units, omega=omega0)
# print(omegas)
# toy_potential = sp.HarmonicOscillatorPotential(omega=omegas, units=units)
# sys.exit(0)
if toy_potential is None:
loop = sd.classify_orbit(w)
if np.any(loop == 1) and not force_harmonic_oscillator: # loop orbit
m,b = sd.fit_isochrone(w, units=units)
toy_potential = sp.IsochronePotential(m=m, b=b, units=units)
else:
omegas = sd.fit_harmonic_oscillator(w, units=units)
toy_potential = sp.HarmonicOscillatorPotential(omega=omegas, units=units)
logger.debug("Toy potential: {}".format(toy_potential))
# plot a smaller section of the orbit in projections of XYZ
plot_w = w[:w.shape[0]//10]
fig = sd.plot_orbits(toy_w, marker=None, linestyle='-',
alpha=0.5, triangle=True, c='r')
fig = sd.plot_orbits(plot_w, axes=fig.axes, marker=None, linestyle='-',
alpha=0.8, triangle=True, c='k')
fig.savefig(os.path.join(plot_path, "orbit_xyz{}.png".format(suffix)))
# plot a smaller section of the orbit in the meridional plane
fig,ax = plt.subplots(1,1,figsize=(6,6))
R = np.sqrt(plot_w[:,0,0]**2 + plot_w[:,0,1]**2)
toy_R = np.sqrt(toy_w[:,0,0]**2 + toy_w[:,0,1]**2)
ax.plot(toy_R, toy_w[:,0,2], marker=None, linestyle='-', alpha=0.5, c='r')
ax.plot(R, plot_w[:,0,2], marker=None, linestyle='-', alpha=0.8, c='k')
ax.set_xlabel("$R$")
ax.set_ylabel("$Z$", rotation='horizontal')
fig.savefig(os.path.join(plot_path, "orbit_Rz{}.png".format(suffix)))
return t,w,toy_potential
