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(self, w0, H):
c = self.config
# return dict
result = self._empty_result
# get timestep and nsteps for integration
try:
dt, nsteps = estimate_dt_n_steps(
w0,
H,
n_periods=c.n_periods,
n_steps_per_period=c.n_steps_per_period,
func=np.nanmin,
Integrator=gi.DOPRI853Integrator)
except RuntimeError:
result['error_code'] = 2
return result
except:
result['error_code'] = 9
return result
# integrate orbit
logger.debug("Integrating orbit with dt={0}, nsteps={1}".format(
dt, nsteps))
try:
orbit = H.integrate_orbit(w0,
dt=dt,
n_steps=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 = orbit.energy()
dEmax = np.max(np.abs((E[1:] - E[0]) / E[0]))
logger.debug('max(∆E) = {0:.2e}'.format(dEmax))
if dEmax > c.energy_tolerance:
result['error_code'] = 4
result['dE_max'] = dEmax
return result
# start finding the frequencies -- do first half then second half
sf1 = SuperFreq(orbit.t[:nsteps // 2 + 1].value, p=c.hamming_p)
sf2 = SuperFreq(orbit.t[nsteps // 2:].value, p=c.hamming_p)
# classify orbit full orbit
circ = orbit.circulation()
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_orbit = orbit.align_circulation_with_z(circ)
fs1 = orbit_to_poincare_polar(new_orbit[sl1])
fs2 = orbit_to_poincare_polar(new_orbit[sl2])
else: # box
ws = orbit.w()
fs1 = [(ws[j, sl1] + 1j * ws[j + 3, sl1]) for j in range(3)]
fs2 = [(ws[j, sl2] + 1j * ws[j + 3, sl2]) for j in range(3)]
logger.debug("Running SuperFreq on the orbits")
try:
freqs1, d1, ixs1 = sf1.find_fundamental_frequencies(
fs1, nintvec=c.n_intvec)
freqs2, d2, ixs2 = sf2.find_fundamental_frequencies(
fs2, nintvec=c.n_intvec)
except:
result['error_code'] = 5
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'] = 1
return result
def run(cls, w0, **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
binary_period = 2*np.pi # Omega = 1
dt = binary_period / c['nsteps_per_period']
nsteps = c['nperiods'] * c['nsteps_per_period']
# integrate orbit
logger.debug("Integrating orbit with dt={0}, nsteps={1}".format(dt, nsteps))
try:
t = np.linspace(0., c['nperiods']*binary_period, nsteps)
t,ws = dop853_integrate_r3bp(np.atleast_2d(w0).copy(), t,
c['q'], c['ecc'], c['nu'],
atol=1E-11, rtol=1E-10, nmax=0)
except RuntimeError: # ODE integration failed
logger.warning("Orbit integration failed.")
dEmax = 1E10
else:
logger.debug('Orbit integrated successfully, checking energy conservation...')
# check Jacobi energy conservation for the orbit
E = 2*r3bp_potential(ws[:,0,:3].copy(), c['q'], c['ecc'], c['nu']) \
- (ws[:,0,3]**2 + ws[:,0,4]**2 + ws[:,0,5]**2)
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,2))*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'])
# define slices for first and second parts
sl1 = slice(None,nsteps//2+1)
sl2 = slice(nsteps//2,None)
# TODO: the 2's below should change if we do the full 3d problem
if c['force_cartesian']:
fs1 = [(ws[sl1,0,j] + 1j*ws[sl1,0,j+3]) for j in range(2)]
fs2 = [(ws[sl2,0,j] + 1j*ws[sl2,0,j+3]) for j in range(2)]
else: # use Poincare polars
# first need to flip coordinates so that circulation is around z axis
new_ws = gc.cartesian_to_poincare_polar(ws)
fs1 = [(new_ws[sl1,j] + 1j*new_ws[sl1,j+3]) for j in range(2)]
fs2 = [(new_ws[sl2,j] + 1j*new_ws[sl2,j+3]) for j in range(2)]
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,2))*np.nan
result['success'] = False
result['error_code'] = 3
return result
result['freqs'] = np.vstack((freqs1, freqs2))
result['dE_max'] = dEmax
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]
# 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
pst_arr = np.zeros((nstars_clusterone, len(t), ndim))
for j, wfile in enumerate(ws):
time_slice = np.loadtxt(wfile)
pst_arr[:, j, :] = time_slice[nstars_clusterzero:end_index]
all_freqs = []
for k in range(nstars_clusterone):
w = pst_arr[k, :, :]
fs = [(w[:, l] * 1j * (0.001022) * w[:, l + ndim // 2])
for l in range(ndim // 2)]
result = sf.find_fundamental_frequencies(fs)
all_freqs += list(result.fund_freqs)
all_freqs = np.sort(all_freqs)
plt.plot(all_freqs, [i / len(all_freqs) for i in range(len(all_freqs))],
linewidth=0.5,
label=r'$\log_{2}N_{\mathrm{clusters}} = %i$' % (np.log2(tag)))
plt.title('Cluster 1, Fundamental Frequencies', fontsize=16)
plt.xlabel(r'Fundamental Frequency (Myr$^{-1}$)', fontsize=12)
plt.ylabel('Cumulative Distribution Function', fontsize=12)
plt.legend(loc='center left', bbox_to_anchor=(1, 0.5), fontsize=10)
plt.tight_layout()
plt.savefig('fund_freqs_clusterone.pdf')