Exemple #1
0
    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
Exemple #2
0
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)