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 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