def estimate_dt_nsteps(w0, potential, nperiods, nsteps_per_period, dE_threshold=1E-9, return_periods=False):
"""
Estimate the timestep and number of steps to integrate for given a potential
and set of initial conditions.
Parameters
----------
w0 : array_like
potential : :class:`~gary.potential.Potential`
nperiods : int
Number of (max) periods to integrate.
nsteps_per_period : int
Number of steps to take per (max) orbital period.
dE_threshold : numeric (optional)
Maximum fractional energy difference -- used to determine initial timestep.
Set to ``None`` to ignore this.
return_periods : bool (optional)
Also return the estimated periods for the orbit.
"""
# integrate orbit
dt = _autodetermine_initial_dt(w0, potential, dE_threshold=dE_threshold)
nsteps = int(round(10000 / dt))
t,w = potential.integrate_orbit(w0, dt=dt, nsteps=nsteps)
# if loop, align circulation with Z and take R period
loop = gd.classify_orbit(w[:,0])
if np.any(loop):
w = gd.align_circulation_with_z(w[:,0], loop)
# convert to cylindrical coordinates
R = np.sqrt(w[:,0]**2 + w[:,1]**2)
phi = np.arctan2(w[:,1], w[:,0])
z = w[:,2]
T = np.array([gd.peak_to_peak_period(t, f) for f in [R, phi, z]])
else:
T = np.array([gd.peak_to_peak_period(t, f) for f in w.T[:3,0]])
# timestep from number of steps per period
Tmax = T.max()
if np.isnan(Tmax):
T = T[np.isfinite(T)]
Tmax = T.max()
if np.isnan(Tmax):
raise RuntimeError("Failed to find period.")
dt = float(Tmax) / float(nsteps_per_period)
nsteps = int(round(nperiods * Tmax / dt))
if dt == 0.:
raise ValueError("Timestep is zero!")
if return_periods:
return dt, nsteps, T
else:
return dt, nsteps
def elliptical_orbit_to_events(t, w):
"""
Convert an orbit to MIDI events using Cartesian coordinates and rules.
Parameters
----------
t : array_like
w : array_like
midi_pool : array_like
"""
loop = gd.classify_orbit(w)
# apocenters
x,y,z = w.T[:3]
r = np.sqrt(x**2 + y**2 + z**2)
apo = np.array([argrelmax(rr)[0] for rr in r])
# get periods
periods = []
for i in range(w.shape[1]):
if np.any(loop[i] == 1):
w2 = gd.align_circulation_with_z(w[:,i], loop[i])
R = np.sqrt(w2[:,0]**2 + w2[:,1]**2)
phi = np.arctan2(w2[:,1], w2[:,0]) % (2*np.pi)
z = w2[:,2]
# loop
T1 = gd.peak_to_peak_period(t, R)
T2 = gd.peak_to_peak_period(t, phi)
T3 = gd.peak_to_peak_period(t, z)
else:
# box
T1 = gd.peak_to_peak_period(t, w[:,i,0])
T2 = gd.peak_to_peak_period(t, w[:,i,1])
T3 = gd.peak_to_peak_period(t, w[:,i,2])
periods.append([T1,T2,T3])
freqs = (2*np.pi / np.array(periods)) * 10000.
delays = []
notes = []
for j in range(w.shape[0]):
_no = []
for i in range(w.shape[1]):
if j in apo[i]:
_no.append(freqs[i].tolist())
if len(_no) > 0:
delays.append(t[j])
notes.append(np.unique(_no).tolist())
delays = np.array(delays)
notes = np.array(notes)
return delays, notes
def elliptical_orbit_to_events(t, w):
"""
Convert an orbit to MIDI events using Cartesian coordinates and rules.
Parameters
----------
t : array_like
w : array_like
midi_pool : array_like
"""
loop = gd.classify_orbit(w)
# apocenters
x, y, z = w.T[:3]
r = np.sqrt(x**2 + y**2 + z**2)
apo = np.array([argrelmax(rr)[0] for rr in r])
# get periods
periods = []
for i in range(w.shape[1]):
if np.any(loop[i] == 1):
w2 = gd.align_circulation_with_z(w[:, i], loop[i])
R = np.sqrt(w2[:, 0]**2 + w2[:, 1]**2)
phi = np.arctan2(w2[:, 1], w2[:, 0]) % (2 * np.pi)
z = w2[:, 2]
# loop
T1 = gd.peak_to_peak_period(t, R)
T2 = gd.peak_to_peak_period(t, phi)
T3 = gd.peak_to_peak_period(t, z)
else:
# box
T1 = gd.peak_to_peak_period(t, w[:, i, 0])
T2 = gd.peak_to_peak_period(t, w[:, i, 1])
T3 = gd.peak_to_peak_period(t, w[:, i, 2])
periods.append([T1, T2, T3])
freqs = (2 * np.pi / np.array(periods)) * 10000.
delays = []
notes = []
for j in range(w.shape[0]):
_no = []
for i in range(w.shape[1]):
if j in apo[i]:
_no.append(freqs[i].tolist())
if len(_no) > 0:
delays.append(t[j])
notes.append(np.unique(_no).tolist())
delays = np.array(delays)
notes = np.array(notes)
return delays, notes
def compute_all_freqs(t, ws, hamming_p=1, nintvec=10, force_cartesian=False):
"""
Compute the fundamental frequencies and amplitudes for all
specified orbits.
This assumes that all orbits have the same geometry as the first
orbit in the orbit array. That is, (if ``force_cartesian`` is
``False``) if the first orbit is a tube orbit, it assumes all orbits
are tubes.
Parameters
----------
t : array_like
ws : array_like
hamming_p : int (optional)
nintvec : int (optional)
force_cartesian : bool (optional)
Returns
-------
freqs : :class:`numpy.ndarray`
amps : :class:`numpy.ndarray`
"""
# classify parent orbit
circ = gd.classify_orbit(ws[:,0])
is_tube = np.any(circ)
allfreqs = []
allamps = []
for i in range(ws.shape[1]):
ww = ws[:,i]
if is_tube and not force_cartesian:
# need to flip coordinates until circulation is around z axis
new_ws = gd.align_circulation_with_z(ww, circ)
new_ws = gc.cartesian_to_poincare_polar(new_ws)
else:
new_ws = ww
fs = [(new_ws[:,j] + 1j*new_ws[:,j+ws.shape[-1]//2]) for j in range(ws.shape[-1]//2)]
sf = SuperFreq(t, p=hamming_p)
try:
freqs,d,ixs = sf.find_fundamental_frequencies(fs, nintvec=nintvec)
except:
allfreqs.append([np.nan,np.nan,np.nan])
allamps.append([np.nan,np.nan,np.nan])
continue
allfreqs.append(freqs.tolist())
allamps.append(d['|A|'][ixs].tolist())
allfreqs = np.array(allfreqs)
allamps = np.array(allamps)
return allfreqs, allamps
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
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
def run(cls, w0, potential, **kwargs):
c = dict()
for k in cls.config_defaults.keys():
if k not in kwargs:
c[k] = cls.config_defaults[k]
else:
c[k] = kwargs[k]
# return dict
result = dict()
# get timestep and nsteps for integration
try:
dt, nsteps = estimate_dt_nsteps(w0.copy(), potential,
c['nperiods'],
c['nsteps_per_period'])
except RuntimeError:
logger.warning("Failed to integrate orbit when estimating dt,nsteps")
result['freqs'] = np.ones((2,3))*np.nan
result['success'] = False
result['error_code'] = 1
return result
except:
logger.warning("Unexpected failure!")
result['freqs'] = np.ones((2,3))*np.nan
result['success'] = False
result['error_code'] = 4
return result
# integrate orbit
logger.debug("Integrating orbit with dt={0}, nsteps={1}".format(dt, nsteps))
try:
t,ws = potential.integrate_orbit(w0.copy(), dt=dt, nsteps=nsteps,
Integrator=gi.DOPRI853Integrator,
Integrator_kwargs=dict(atol=1E-11))
except RuntimeError: # ODE integration failed
logger.warning("Orbit integration failed.")
dEmax = 1E10
else:
logger.debug('Orbit integrated successfully, checking energy conservation...')
# check energy conservation for the orbit
E = potential.total_energy(ws[:,0,:3].copy(), ws[:,0,3:].copy())
dE = np.abs(E[1:] - E[0])
dEmax = dE.max() / np.abs(E[0])
logger.debug('max(∆E) = {0:.2e}'.format(dEmax))
if dEmax > c['energy_tolerance']:
logger.warning("Failed due to energy conservation check.")
result['freqs'] = np.ones((2,3))*np.nan
result['success'] = False
result['error_code'] = 2
result['dE_max'] = dEmax
return result
# start finding the frequencies -- do first half then second half
sf1 = SuperFreq(t[:nsteps//2+1], p=c['hamming_p'])
sf2 = SuperFreq(t[nsteps//2:], p=c['hamming_p'])
# classify orbit full orbit
circ = gd.classify_orbit(ws)
is_tube = np.any(circ)
# define slices for first and second parts
sl1 = slice(None,nsteps//2+1)
sl2 = slice(nsteps//2,None)
if is_tube and not c['force_cartesian']:
# first need to flip coordinates so that circulation is around z axis
new_ws = gd.align_circulation_with_z(ws, circ)
new_ws = gc.cartesian_to_poincare_polar(new_ws)
fs1 = [(new_ws[sl1,j] + 1j*new_ws[sl1,j+3]) for j in range(3)]
fs2 = [(new_ws[sl2,j] + 1j*new_ws[sl2,j+3]) for j in range(3)]
else: # box
fs1 = [(ws[sl1,0,j] + 1j*ws[sl1,0,j+3]) for j in range(3)]
fs2 = [(ws[sl2,0,j] + 1j*ws[sl2,0,j+3]) for j in range(3)]
logger.debug("Running SuperFreq on the orbits")
try:
freqs1,d1,ixs1 = sf1.find_fundamental_frequencies(fs1, nintvec=c['nintvec'])
freqs2,d2,ixs2 = sf2.find_fundamental_frequencies(fs2, nintvec=c['nintvec'])
except:
result['freqs'] = np.ones((2,3))*np.nan
result['success'] = False
result['error_code'] = 3
return result
result['freqs'] = np.vstack((freqs1, freqs2))
result['dE_max'] = dEmax
result['is_tube'] = float(is_tube)
result['dt'] = float(dt)
result['nsteps'] = nsteps
result['amps'] = np.vstack((d1['|A|'][ixs1], d2['|A|'][ixs2]))
result['success'] = True
result['error_code'] = 0
return result
def run(cls, w0, potential, **kwargs):
c = dict()
for k in cls.config_defaults.keys():
if k not in kwargs:
c[k] = cls.config_defaults[k]
else:
c[k] = kwargs[k]
# container for return
result = dict()
# automatically estimate dt, nsteps
try:
dt, nsteps = estimate_dt_nsteps(w0.copy(), potential,
c['total_nperiods'], c['nsteps_per_period'])
except RuntimeError:
logger.warning("Failed to integrate orbit when estimating dt,nsteps")
result['freqs'] = np.nan
result['success'] = False
result['error_code'] = 1
return result
logger.debug("Integrating orbit with dt={0}, nsteps={1}".format(dt, nsteps))
try:
t,ws = potential.integrate_orbit(w0.copy(), dt=dt, nsteps=nsteps,
Integrator=gi.DOPRI853Integrator,
Integrator_kwargs=dict(atol=1E-11))
except RuntimeError: # ODE integration failed
logger.warning("Orbit integration failed.")
dEmax = 1E10
else:
logger.debug('Orbit integrated successfully, checking energy conservation...')
# check energy conservation for the orbit
E = potential.total_energy(ws[:,0,:3].copy(), ws[:,0,3:].copy())
dE = np.abs(E[1:] - E[0])
dEmax = dE.max() / np.abs(E[0])
logger.debug('max(∆E) = {0:.2e}'.format(dEmax))
if dEmax > c['energy_tolerance']:
logger.warning("Failed due to energy conservation check.")
result['freqs'] = np.nan
result['success'] = False
result['error_code'] = 2
return result
# windowing properties - convert from period to steps
window_width = int(c['window_width'] * c['nsteps_per_period'])
window_stride = int(c['window_stride'] * c['nsteps_per_period'])
# classify orbit full orbit
circ = gd.classify_orbit(ws[:,0])
is_tube = np.any(circ)
logger.debug("Running SuperFreq on each window:")
allfreqs = []
allamps = []
for (i1,i2),ww in rolling_window(ws[:,0], window_size=window_width, stride=window_stride, return_idx=True):
if i2 >= nsteps:
break
logger.debug("Window: {0}:{1}".format(i1,i2))
if is_tube and not c['force_cartesian']:
# need to flip coordinates until circulation is around z axis
new_ws = gd.align_circulation_with_z(ww, circ)
new_ws = gc.cartesian_to_poincare_polar(new_ws)
else:
new_ws = ww
fs = [(new_ws[:,j] + 1j*new_ws[:,j+3]) for j in range(3)]
naff = SuperFreq(t[i1:i2], p=c['hamming_p'])
try:
freqs,d,ixs = naff.find_fundamental_frequencies(fs, nintvec=c['nintvec'])
except:
result['freqs'] = np.nan
result['success'] = False
result['error_code'] = 3
return result
allfreqs.append(freqs.tolist())
allamps.append(d['|A|'][ixs].tolist())
allfreqs = np.array(allfreqs)
allamps = np.array(allamps)
result['freqs'] = allfreqs
result['amps'] = allamps
result['dE_max'] = dEmax
result['dt'] = float(dt)
result['nsteps'] = nsteps
result['is_tube'] = is_tube
result['success'] = True
result['error_code'] = 0
return result