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 main(ngrid, CJ, q, ecc=0., nu=0., max_xy=3., output_path=None,
run_name=None, overwrite=False, plot=False):
"""
"""
# create path
if output_path is None:
output_path = os.path.join(project_path, "output")
if run_name is None:
run_name = "n{}_cj{:.3f}_q{:.0e}_e{:.0e}_nu{:.0e}".format(ngrid, CJ, q, ecc, nu)
path = os.path.join(output_path, run_name)
else:
path = output_path
logger.info("Caching to: {}".format(path))
if not os.path.exists(path):
os.makedirs(path)
# path to initial conditions cache
w0path = os.path.join(path, 'w0.npy')
if os.path.exists(w0path) and overwrite:
os.remove(w0path)
if not os.path.exists(w0path):
# generate initial conditions
# box
# grid = np.linspace(-max_xy,max_xy,ngrid)
# z = np.zeros(ngrid*ngrid)
# xyz = np.vstack(map(np.ravel, np.meshgrid(grid, grid))+[z]).T.copy()
# disc
r = np.sqrt(np.random.uniform(0.,max_xy**2,ngrid*ngrid))
phi = np.random.uniform(0.,2*np.pi,ngrid*ngrid)
xyz = np.zeros((len(r),3))
xyz[:,0] = r*np.cos(phi)
xyz[:,1] = r*np.sin(phi)
U = r3bp_potential(xyz, q, ecc, nu)
# ignore all points above ZVC
ix = 2*U >= CJ
xyz = xyz[ix]
U = U[ix]
r = np.sqrt(np.sum(xyz**2, axis=-1))
vmag = np.sqrt(2*U - CJ)
vx = -vmag * xyz[:,1]/r
vy = vmag * xyz[:,0]/r
vxyz = np.zeros_like(xyz)
vxyz[:,0] = vx
vxyz[:,1] = vy
w0 = np.hstack((xyz,vxyz))
CJs = 2*r3bp_potential(xyz, q, ecc, nu) - vx**2 - vy**2
assert np.allclose(CJs, CJ)
else:
w0 = np.load(w0path)
logger.info("Initial conditions file already exists!\n\t{}".format(w0path))
if plot:
fig,ax = pl.subplots(1,1,figsize=(8,8))
ax.plot(w0[:,0], w0[:,1], ls='none', marker=',', color='k')
fig.savefig(os.path.join(path, 'w0.png'), dpi=300)
logger.info("Number of initial conditions: {}".format(len(w0)))
np.save(w0path, w0)
