"""
    Test the integrators.
"""

# Third-party
import pytest
import matplotlib.pyplot as pl
import numpy as np

# Project
from .. import LeapfrogIntegrator, RK5Integrator, DOPRI853Integrator

# Integrators to test
integrator_list = [RK5Integrator, DOPRI853Integrator, LeapfrogIntegrator]

# ----------------------------------------------------------------------------

@pytest.mark.parametrize("Integrator", integrator_list)
def test_sho_forward_backward(Integrator):
    def sho(t,w,T):
        q,p = w
        return np.array([p, -(2*np.pi/T)**2*q])

    integrator = Integrator(sho, func_args=(1.,))

    dt = 0.01
    n_steps = 100
    if Integrator == LeapfrogIntegrator:
        dt = 1E-4
        n_steps = int(1E4)

    forw = integrator.run([0., 1.], dt=dt, n_steps=n_steps)
    back = integrator.run([0., 1.], dt=-dt, n_steps=n_steps)

    assert np.allclose(forw.w()[:,-1], back.w()[:,-1], atol=1E-6)

@pytest.mark.parametrize("Integrator", integrator_list)
def test_point_mass(Integrator):
    def F(t,w):
        x,y,px,py = w
        a = -1./(x*x+y*y)**1.5
        return np.array([px, py, x*a, y*a])

    q0 = np.array([1., 0.])
    p0 = np.array([0., 1.])
    T = 1.

    integrator = Integrator(F)
    orbit = integrator.run(np.append(q0,p0), t1=0., t2=2*np.pi, n_steps=1E4)

    assert np.allclose(orbit.w()[:,0], orbit.w()[:,-1], atol=1E-6)

@pytest.mark.parametrize("Integrator", integrator_list)
def test_point_mass_multiple(Integrator):
    def F(t,w):
        x,y,px,py = w
        a = -1/(x*x+y*y)**1.5
        return np.array([px, py, x*a, y*a])

    w0 = np.array([[1.0, 0.0, 0.0, 1.],
                   [0.8, 0.0, 0.0, 1.1],
                   [2., 1.0, -1.0, 1.1]]).T

    integrator = Integrator(F)
    orbit = integrator.run(w0, dt=1E-3, n_steps=1E4)

@pytest.mark.parametrize("Integrator", integrator_list)
def test_driven_pendulum(Integrator):
    def F(t,w,A,omega_d):
        q,p = w
        return np.array([p,-np.sin(q) + A*np.cos(omega_d*t)])

    integrator = Integrator(F, func_args=(0.07, 0.75))
    orbit = integrator.run([3., 0.], dt=1E-2, n_steps=1E4)

@pytest.mark.parametrize("Integrator", integrator_list)
def test_lorenz(Integrator):
    def F(t,w,sigma,rho,beta):
        x,y,z,px,py,pz = w
        return np.array([sigma*(y-x), x*(rho-z)-y, x*y-beta*z, 0., 0., 0.]).reshape(w.shape)

    sigma, rho, beta = 10., 28., 8/3.
    integrator = Integrator(F, func_args=(sigma, rho, beta))

    orbit = integrator.run([0.5,0.5,0.5,0,0,0], dt=1E-2, n_steps=1E4)

    # pl.plot(ws[0], ws[1])
    # pl.show()

@pytest.mark.parametrize("Integrator", integrator_list)
def test_memmap(tmpdir, Integrator):
    dt = 0.1
    n_steps = 1000
    nw0 = 10000
    mmap = np.memmap("/tmp/test_memmap.npy", mode='w+', shape=(2, n_steps+1, nw0))

    def sho(t,w,T):
        q,p = w
        return np.array([p, -(2*np.pi/T)**2*q])

    w0 = np.random.uniform(-1,1,size=(2,nw0))

    integrator = Integrator(sho, func_args=(1.,))

    orbit = integrator.run(w0, dt=dt, n_steps=n_steps, mmap=mmap)

cea8457907071d8ff3936d4c974e588a3c52bd08
<mask token>


@pytest.mark.parametrize('Integrator', integrator_list)
def test_sho_forward_backward(Integrator):

    def sho(t, w, T):
        q, p = w
        return np.array([p, -(2 * np.pi / T) ** 2 * q])
    integrator = Integrator(sho, func_args=(1.0,))
    dt = 0.01
    n_steps = 100
    if Integrator == LeapfrogIntegrator:
        dt = 0.0001
        n_steps = int(10000.0)
    forw = integrator.run([0.0, 1.0], dt=dt, n_steps=n_steps)
    back = integrator.run([0.0, 1.0], dt=-dt, n_steps=n_steps)
    assert np.allclose(forw.w()[:, -1], back.w()[:, -1], atol=1e-06)


<mask token>


@pytest.mark.parametrize('Integrator', integrator_list)
def test_point_mass_multiple(Integrator):

    def F(t, w):
        x, y, px, py = w
        a = -1 / (x * x + y * y) ** 1.5
        return np.array([px, py, x * a, y * a])
    w0 = np.array([[1.0, 0.0, 0.0, 1.0], [0.8, 0.0, 0.0, 1.1], [2.0, 1.0, -
        1.0, 1.1]]).T
    integrator = Integrator(F)
    orbit = integrator.run(w0, dt=0.001, n_steps=10000.0)


<mask token>


@pytest.mark.parametrize('Integrator', integrator_list)
def test_lorenz(Integrator):

    def F(t, w, sigma, rho, beta):
        x, y, z, px, py, pz = w
        return np.array([sigma * (y - x), x * (rho - z) - y, x * y - beta *
            z, 0.0, 0.0, 0.0]).reshape(w.shape)
    sigma, rho, beta = 10.0, 28.0, 8 / 3.0
    integrator = Integrator(F, func_args=(sigma, rho, beta))
    orbit = integrator.run([0.5, 0.5, 0.5, 0, 0, 0], dt=0.01, n_steps=10000.0)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_memmap(tmpdir, Integrator):
    dt = 0.1
    n_steps = 1000
    nw0 = 10000
    mmap = np.memmap('/tmp/test_memmap.npy', mode='w+', shape=(2, n_steps +
        1, nw0))

    def sho(t, w, T):
        q, p = w
        return np.array([p, -(2 * np.pi / T) ** 2 * q])
    w0 = np.random.uniform(-1, 1, size=(2, nw0))
    integrator = Integrator(sho, func_args=(1.0,))
    orbit = integrator.run(w0, dt=dt, n_steps=n_steps, mmap=mmap)

<mask token>


@pytest.mark.parametrize('Integrator', integrator_list)
def test_sho_forward_backward(Integrator):

    def sho(t, w, T):
        q, p = w
        return np.array([p, -(2 * np.pi / T) ** 2 * q])
    integrator = Integrator(sho, func_args=(1.0,))
    dt = 0.01
    n_steps = 100
    if Integrator == LeapfrogIntegrator:
        dt = 0.0001
        n_steps = int(10000.0)
    forw = integrator.run([0.0, 1.0], dt=dt, n_steps=n_steps)
    back = integrator.run([0.0, 1.0], dt=-dt, n_steps=n_steps)
    assert np.allclose(forw.w()[:, -1], back.w()[:, -1], atol=1e-06)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_point_mass(Integrator):

    def F(t, w):
        x, y, px, py = w
        a = -1.0 / (x * x + y * y) ** 1.5
        return np.array([px, py, x * a, y * a])
    q0 = np.array([1.0, 0.0])
    p0 = np.array([0.0, 1.0])
    T = 1.0
    integrator = Integrator(F)
    orbit = integrator.run(np.append(q0, p0), t1=0.0, t2=2 * np.pi, n_steps
        =10000.0)
    assert np.allclose(orbit.w()[:, 0], orbit.w()[:, -1], atol=1e-06)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_point_mass_multiple(Integrator):

    def F(t, w):
        x, y, px, py = w
        a = -1 / (x * x + y * y) ** 1.5
        return np.array([px, py, x * a, y * a])
    w0 = np.array([[1.0, 0.0, 0.0, 1.0], [0.8, 0.0, 0.0, 1.1], [2.0, 1.0, -
        1.0, 1.1]]).T
    integrator = Integrator(F)
    orbit = integrator.run(w0, dt=0.001, n_steps=10000.0)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_driven_pendulum(Integrator):

    def F(t, w, A, omega_d):
        q, p = w
        return np.array([p, -np.sin(q) + A * np.cos(omega_d * t)])
    integrator = Integrator(F, func_args=(0.07, 0.75))
    orbit = integrator.run([3.0, 0.0], dt=0.01, n_steps=10000.0)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_lorenz(Integrator):

    def F(t, w, sigma, rho, beta):
        x, y, z, px, py, pz = w
        return np.array([sigma * (y - x), x * (rho - z) - y, x * y - beta *
            z, 0.0, 0.0, 0.0]).reshape(w.shape)
    sigma, rho, beta = 10.0, 28.0, 8 / 3.0
    integrator = Integrator(F, func_args=(sigma, rho, beta))
    orbit = integrator.run([0.5, 0.5, 0.5, 0, 0, 0], dt=0.01, n_steps=10000.0)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_memmap(tmpdir, Integrator):
    dt = 0.1
    n_steps = 1000
    nw0 = 10000
    mmap = np.memmap('/tmp/test_memmap.npy', mode='w+', shape=(2, n_steps +
        1, nw0))

    def sho(t, w, T):
        q, p = w
        return np.array([p, -(2 * np.pi / T) ** 2 * q])
    w0 = np.random.uniform(-1, 1, size=(2, nw0))
    integrator = Integrator(sho, func_args=(1.0,))
    orbit = integrator.run(w0, dt=dt, n_steps=n_steps, mmap=mmap)

<mask token>
integrator_list = [RK5Integrator, DOPRI853Integrator, LeapfrogIntegrator]


@pytest.mark.parametrize('Integrator', integrator_list)
def test_sho_forward_backward(Integrator):

    def sho(t, w, T):
        q, p = w
        return np.array([p, -(2 * np.pi / T) ** 2 * q])
    integrator = Integrator(sho, func_args=(1.0,))
    dt = 0.01
    n_steps = 100
    if Integrator == LeapfrogIntegrator:
        dt = 0.0001
        n_steps = int(10000.0)
    forw = integrator.run([0.0, 1.0], dt=dt, n_steps=n_steps)
    back = integrator.run([0.0, 1.0], dt=-dt, n_steps=n_steps)
    assert np.allclose(forw.w()[:, -1], back.w()[:, -1], atol=1e-06)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_point_mass(Integrator):

    def F(t, w):
        x, y, px, py = w
        a = -1.0 / (x * x + y * y) ** 1.5
        return np.array([px, py, x * a, y * a])
    q0 = np.array([1.0, 0.0])
    p0 = np.array([0.0, 1.0])
    T = 1.0
    integrator = Integrator(F)
    orbit = integrator.run(np.append(q0, p0), t1=0.0, t2=2 * np.pi, n_steps
        =10000.0)
    assert np.allclose(orbit.w()[:, 0], orbit.w()[:, -1], atol=1e-06)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_point_mass_multiple(Integrator):

    def F(t, w):
        x, y, px, py = w
        a = -1 / (x * x + y * y) ** 1.5
        return np.array([px, py, x * a, y * a])
    w0 = np.array([[1.0, 0.0, 0.0, 1.0], [0.8, 0.0, 0.0, 1.1], [2.0, 1.0, -
        1.0, 1.1]]).T
    integrator = Integrator(F)
    orbit = integrator.run(w0, dt=0.001, n_steps=10000.0)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_driven_pendulum(Integrator):

    def F(t, w, A, omega_d):
        q, p = w
        return np.array([p, -np.sin(q) + A * np.cos(omega_d * t)])
    integrator = Integrator(F, func_args=(0.07, 0.75))
    orbit = integrator.run([3.0, 0.0], dt=0.01, n_steps=10000.0)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_lorenz(Integrator):

    def F(t, w, sigma, rho, beta):
        x, y, z, px, py, pz = w
        return np.array([sigma * (y - x), x * (rho - z) - y, x * y - beta *
            z, 0.0, 0.0, 0.0]).reshape(w.shape)
    sigma, rho, beta = 10.0, 28.0, 8 / 3.0
    integrator = Integrator(F, func_args=(sigma, rho, beta))
    orbit = integrator.run([0.5, 0.5, 0.5, 0, 0, 0], dt=0.01, n_steps=10000.0)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_memmap(tmpdir, Integrator):
    dt = 0.1
    n_steps = 1000
    nw0 = 10000
    mmap = np.memmap('/tmp/test_memmap.npy', mode='w+', shape=(2, n_steps +
        1, nw0))

    def sho(t, w, T):
        q, p = w
        return np.array([p, -(2 * np.pi / T) ** 2 * q])
    w0 = np.random.uniform(-1, 1, size=(2, nw0))
    integrator = Integrator(sho, func_args=(1.0,))
    orbit = integrator.run(w0, dt=dt, n_steps=n_steps, mmap=mmap)

<mask token>
import pytest
import matplotlib.pyplot as pl
import numpy as np
from .. import LeapfrogIntegrator, RK5Integrator, DOPRI853Integrator
integrator_list = [RK5Integrator, DOPRI853Integrator, LeapfrogIntegrator]


@pytest.mark.parametrize('Integrator', integrator_list)
def test_sho_forward_backward(Integrator):

    def sho(t, w, T):
        q, p = w
        return np.array([p, -(2 * np.pi / T) ** 2 * q])
    integrator = Integrator(sho, func_args=(1.0,))
    dt = 0.01
    n_steps = 100
    if Integrator == LeapfrogIntegrator:
        dt = 0.0001
        n_steps = int(10000.0)
    forw = integrator.run([0.0, 1.0], dt=dt, n_steps=n_steps)
    back = integrator.run([0.0, 1.0], dt=-dt, n_steps=n_steps)
    assert np.allclose(forw.w()[:, -1], back.w()[:, -1], atol=1e-06)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_point_mass(Integrator):

    def F(t, w):
        x, y, px, py = w
        a = -1.0 / (x * x + y * y) ** 1.5
        return np.array([px, py, x * a, y * a])
    q0 = np.array([1.0, 0.0])
    p0 = np.array([0.0, 1.0])
    T = 1.0
    integrator = Integrator(F)
    orbit = integrator.run(np.append(q0, p0), t1=0.0, t2=2 * np.pi, n_steps
        =10000.0)
    assert np.allclose(orbit.w()[:, 0], orbit.w()[:, -1], atol=1e-06)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_point_mass_multiple(Integrator):

    def F(t, w):
        x, y, px, py = w
        a = -1 / (x * x + y * y) ** 1.5
        return np.array([px, py, x * a, y * a])
    w0 = np.array([[1.0, 0.0, 0.0, 1.0], [0.8, 0.0, 0.0, 1.1], [2.0, 1.0, -
        1.0, 1.1]]).T
    integrator = Integrator(F)
    orbit = integrator.run(w0, dt=0.001, n_steps=10000.0)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_driven_pendulum(Integrator):

    def F(t, w, A, omega_d):
        q, p = w
        return np.array([p, -np.sin(q) + A * np.cos(omega_d * t)])
    integrator = Integrator(F, func_args=(0.07, 0.75))
    orbit = integrator.run([3.0, 0.0], dt=0.01, n_steps=10000.0)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_lorenz(Integrator):

    def F(t, w, sigma, rho, beta):
        x, y, z, px, py, pz = w
        return np.array([sigma * (y - x), x * (rho - z) - y, x * y - beta *
            z, 0.0, 0.0, 0.0]).reshape(w.shape)
    sigma, rho, beta = 10.0, 28.0, 8 / 3.0
    integrator = Integrator(F, func_args=(sigma, rho, beta))
    orbit = integrator.run([0.5, 0.5, 0.5, 0, 0, 0], dt=0.01, n_steps=10000.0)


@pytest.mark.parametrize('Integrator', integrator_list)
def test_memmap(tmpdir, Integrator):
    dt = 0.1
    n_steps = 1000
    nw0 = 10000
    mmap = np.memmap('/tmp/test_memmap.npy', mode='w+', shape=(2, n_steps +
        1, nw0))

    def sho(t, w, T):
        q, p = w
        return np.array([p, -(2 * np.pi / T) ** 2 * q])
    w0 = np.random.uniform(-1, 1, size=(2, nw0))
    integrator = Integrator(sho, func_args=(1.0,))
    orbit = integrator.run(w0, dt=dt, n_steps=n_steps, mmap=mmap)

"""
    Test the integrators.
"""

# Third-party
import pytest
import matplotlib.pyplot as pl
import numpy as np

# Project
from .. import LeapfrogIntegrator, RK5Integrator, DOPRI853Integrator

# Integrators to test
integrator_list = [RK5Integrator, DOPRI853Integrator, LeapfrogIntegrator]

# ----------------------------------------------------------------------------

@pytest.mark.parametrize("Integrator", integrator_list)
def test_sho_forward_backward(Integrator):
    def sho(t,w,T):
        q,p = w
        return np.array([p, -(2*np.pi/T)**2*q])

    integrator = Integrator(sho, func_args=(1.,))

    dt = 0.01
    n_steps = 100
    if Integrator == LeapfrogIntegrator:
        dt = 1E-4
        n_steps = int(1E4)

    forw = integrator.run([0., 1.], dt=dt, n_steps=n_steps)
    back = integrator.run([0., 1.], dt=-dt, n_steps=n_steps)

    assert np.allclose(forw.w()[:,-1], back.w()[:,-1], atol=1E-6)

@pytest.mark.parametrize("Integrator", integrator_list)
def test_point_mass(Integrator):
    def F(t,w):
        x,y,px,py = w
        a = -1./(x*x+y*y)**1.5
        return np.array([px, py, x*a, y*a])

    q0 = np.array([1., 0.])
    p0 = np.array([0., 1.])
    T = 1.

    integrator = Integrator(F)
    orbit = integrator.run(np.append(q0,p0), t1=0., t2=2*np.pi, n_steps=1E4)

    assert np.allclose(orbit.w()[:,0], orbit.w()[:,-1], atol=1E-6)

@pytest.mark.parametrize("Integrator", integrator_list)
def test_point_mass_multiple(Integrator):
    def F(t,w):
        x,y,px,py = w
        a = -1/(x*x+y*y)**1.5
        return np.array([px, py, x*a, y*a])

    w0 = np.array([[1.0, 0.0, 0.0, 1.],
                   [0.8, 0.0, 0.0, 1.1],
                   [2., 1.0, -1.0, 1.1]]).T

    integrator = Integrator(F)
    orbit = integrator.run(w0, dt=1E-3, n_steps=1E4)

@pytest.mark.parametrize("Integrator", integrator_list)
def test_driven_pendulum(Integrator):
    def F(t,w,A,omega_d):
        q,p = w
        return np.array([p,-np.sin(q) + A*np.cos(omega_d*t)])

    integrator = Integrator(F, func_args=(0.07, 0.75))
    orbit = integrator.run([3., 0.], dt=1E-2, n_steps=1E4)

@pytest.mark.parametrize("Integrator", integrator_list)
def test_lorenz(Integrator):
    def F(t,w,sigma,rho,beta):
        x,y,z,px,py,pz = w
        return np.array([sigma*(y-x), x*(rho-z)-y, x*y-beta*z, 0., 0., 0.]).reshape(w.shape)

    sigma, rho, beta = 10., 28., 8/3.
    integrator = Integrator(F, func_args=(sigma, rho, beta))

    orbit = integrator.run([0.5,0.5,0.5,0,0,0], dt=1E-2, n_steps=1E4)

    # pl.plot(ws[0], ws[1])
    # pl.show()

@pytest.mark.parametrize("Integrator", integrator_list)
def test_memmap(tmpdir, Integrator):
    dt = 0.1
    n_steps = 1000
    nw0 = 10000
    mmap = np.memmap("/tmp/test_memmap.npy", mode='w+', shape=(2, n_steps+1, nw0))

    def sho(t,w,T):
        q,p = w
        return np.array([p, -(2*np.pi/T)**2*q])

    w0 = np.random.uniform(-1,1,size=(2,nw0))

    integrator = Integrator(sho, func_args=(1.,))

    orbit = integrator.run(w0, dt=dt, n_steps=n_steps, mmap=mmap)
