# Path: fastpm/core.py
# def leapfrog(stages):
#     """ Generate a leap frog stepping scheme.
# 
#         Parameters
#         ----------
#         stages : array_like
#             Time (a) where force computing stage is requested.
#     """
#     if len(stages) == 0:
#         return
# 
#     ai = stages[0]
#     # first force calculation for jump starting
#     yield 'F', ai, ai, ai
#     x, p, f = ai, ai, ai
# 
#     for i in range(len(stages) - 1):
#         a0 = stages[i]
#         a1 = stages[i + 1]
#         ah = (a0 * a1) ** 0.5
#         yield 'K', p, f, ah
#         p = ah
#         yield 'D', x, p, a1
#         x = a1
#         yield 'F', f, x, a1
#         f = a1
#         yield 'K', p, f, a1
#         p = a1
# 
# def autostages(knots, N, astart=None, N0=None):
#     """ Generate an optimized list of N stages that includes time steps at the knots.
# 
#         Parameters
#         ----------
#         astart : float, or None
#             starting time, default is knots[0]
#         N : int
#             total number of stages
#         N0 : int or None
#             at least add this many stages before the earlist knot, default None;
#             useful only if astart != min(knots), and len(knots) > 1
#         knots : list
#             stages that must exist
# 
# 
#         >>> autostages(0.1, N=11, knots=[0.1, 0.2, 0.5, 1.0])
# 
#     """
# 
#     knots = numpy.array(knots)
#     knots.sort()
# 
#     stages = numpy.array([], dtype='f8')
#     if astart is not None and astart != knots.min():
#         assert astart < knots.min()
#         if N0 is None: N0 = 1
#         knots = numpy.append([astart], knots)
#     else:
#         N0 = 1
# 
#     for i in range(0, len(knots) - 1):
#         da = (knots[-1] - knots[i]) / (N - len(stages) - 1)
# 
#         N_this_span = int((knots[i + 1] - knots[i]) / da + 0.5)
#         if i == 0 and N_this_span < N0:
#             N_this_span = N0
# 
#         add = numpy.linspace(knots[i], knots[i + 1], N_this_span, endpoint=False)
# 
#         #print('i = =====', i)
#         #print('knots[i]', knots[i], da, N_this_span, stages, add)
# 
#         stages = numpy.append(stages, add)
# 
#     stages = numpy.append(stages, [knots[-1]])
# 
#     return stages
# 
# Path: fastpm/background.py
# class Perturbation:
# class MatterDominated(Perturbation):
# class RadiationDominated(Perturbation):
#     def __init__(self, a, a_normalize=1.0):
#     def D1(self, a, order=0):
#     def D2(self, a, order=0):
#     def f1(self, a):
#     def f2(self, a):
#     def Gp(self, a):
#     def Gp2(self, a):
#     def gp(self, a):
#     def gp2(self, a):
#     def Gf(self, a):
#     def Gf2(self, a):
#     def gf(self, a):
#     def gf2(self, a):
#     def E(self, a, order=0):
#     def Hfac(self, a):
#     def ode(self, y, lna):
#     def _solve(self, y0, a_normalize):
#     def __init__(self, Omega0_m, Omega0_lambda=None, Omega0_k=0, a=None, a_normalize=1.0):
#     def get_initial_condition(self):
#     def efunc(self, a):
#     def efunc_prime(self, a):
#     def Om(self, a):
#     def __init__(self, cosmo, a=None, a_normalize=1.0):
#     def get_initial_condition(self):
#     def efunc(self, a):
#     def efunc_prime(self, a):
#     def Om(self, a):
#         D1, F1, D2, F2 = y
#             D1, F1, D2, F2 = yi
# 
# Path: fastpm/ncdm.py
# class Solver(core.Solver):
#     @property
#     def nbodystep(self):
#         return FastPMStep(self)


from runtests.mpi import MPITest
from fastpm.core import leapfrog, autostages
from fastpm.background import PerturbationGrowth
from pmesh.pm import ParticleMesh
from nbodykit.cosmology import Planck15, LinearPower
from numpy.testing import assert_allclose
from fastpm.ncdm import Solver
import numpy




@MPITest([1, 4])
def test_ncdm(comm):
    pm = ParticleMesh(BoxSize=512., Nmesh=[8, 8, 8], comm=comm)
    Plin = LinearPower(Planck15, redshift=0, transfer='EisensteinHu')
    solver = Solver(pm, Planck15, B=1)
    Q = pm.generate_uniform_particle_grid(shift=0)

    wn = solver.whitenoise(1234)
    dlin = solver.linear(wn, lambda k: Plin(k))
    state = solver.lpt(dlin, Q, a=1.0, order=2)


# Path: fastpm/core.py
# def leapfrog(stages):
#     """ Generate a leap frog stepping scheme.
# 
#         Parameters
#         ----------
#         stages : array_like
#             Time (a) where force computing stage is requested.
#     """
#     if len(stages) == 0:
#         return
# 
#     ai = stages[0]
#     # first force calculation for jump starting
#     yield 'F', ai, ai, ai
#     x, p, f = ai, ai, ai
# 
#     for i in range(len(stages) - 1):
#         a0 = stages[i]
#         a1 = stages[i + 1]
#         ah = (a0 * a1) ** 0.5
#         yield 'K', p, f, ah
#         p = ah
#         yield 'D', x, p, a1
#         x = a1
#         yield 'F', f, x, a1
#         f = a1
#         yield 'K', p, f, a1
#         p = a1
# 
# def autostages(knots, N, astart=None, N0=None):
#     """ Generate an optimized list of N stages that includes time steps at the knots.
# 
#         Parameters
#         ----------
#         astart : float, or None
#             starting time, default is knots[0]
#         N : int
#             total number of stages
#         N0 : int or None
#             at least add this many stages before the earlist knot, default None;
#             useful only if astart != min(knots), and len(knots) > 1
#         knots : list
#             stages that must exist
# 
# 
#         >>> autostages(0.1, N=11, knots=[0.1, 0.2, 0.5, 1.0])
# 
#     """
# 
#     knots = numpy.array(knots)
#     knots.sort()
# 
#     stages = numpy.array([], dtype='f8')
#     if astart is not None and astart != knots.min():
#         assert astart < knots.min()
#         if N0 is None: N0 = 1
#         knots = numpy.append([astart], knots)
#     else:
#         N0 = 1
# 
#     for i in range(0, len(knots) - 1):
#         da = (knots[-1] - knots[i]) / (N - len(stages) - 1)
# 
#         N_this_span = int((knots[i + 1] - knots[i]) / da + 0.5)
#         if i == 0 and N_this_span < N0:
#             N_this_span = N0
# 
#         add = numpy.linspace(knots[i], knots[i + 1], N_this_span, endpoint=False)
# 
#         #print('i = =====', i)
#         #print('knots[i]', knots[i], da, N_this_span, stages, add)
# 
#         stages = numpy.append(stages, add)
# 
#     stages = numpy.append(stages, [knots[-1]])
# 
#     return stages
# 
# Path: fastpm/background.py
# class Perturbation:
# class MatterDominated(Perturbation):
# class RadiationDominated(Perturbation):
#     def __init__(self, a, a_normalize=1.0):
#     def D1(self, a, order=0):
#     def D2(self, a, order=0):
#     def f1(self, a):
#     def f2(self, a):
#     def Gp(self, a):
#     def Gp2(self, a):
#     def gp(self, a):
#     def gp2(self, a):
#     def Gf(self, a):
#     def Gf2(self, a):
#     def gf(self, a):
#     def gf2(self, a):
#     def E(self, a, order=0):
#     def Hfac(self, a):
#     def ode(self, y, lna):
#     def _solve(self, y0, a_normalize):
#     def __init__(self, Omega0_m, Omega0_lambda=None, Omega0_k=0, a=None, a_normalize=1.0):
#     def get_initial_condition(self):
#     def efunc(self, a):
#     def efunc_prime(self, a):
#     def Om(self, a):
#     def __init__(self, cosmo, a=None, a_normalize=1.0):
#     def get_initial_condition(self):
#     def efunc(self, a):
#     def efunc_prime(self, a):
#     def Om(self, a):
#         D1, F1, D2, F2 = y
#             D1, F1, D2, F2 = yi
# 
# Path: fastpm/ncdm.py
# class Solver(core.Solver):
#     @property
#     def nbodystep(self):
#         return FastPMStep(self)


# Path: fastpm/tests/test_ncdm.py
from runtests.mpi import MPITest
from fastpm.core import leapfrog, autostages
from fastpm.background import PerturbationGrowth
from pmesh.pm import ParticleMesh
from nbodykit.cosmology import Planck15, LinearPower
from numpy.testing import assert_allclose
from fastpm.ncdm import Solver
import numpy





@MPITest([1, 4])
def test_ncdm(comm):
    pm = ParticleMesh(BoxSize=512., Nmesh=[8, 8, 8], comm=comm)
    Plin = LinearPower(Planck15, redshift=0, transfer='EisensteinHu')
    solver = Solver(pm, Planck15, B=1)
    Q = pm.generate_uniform_particle_grid(shift=0)

    wn = solver.whitenoise(1234)
    dlin = solver.linear(wn, lambda k: Plin(k))
    state = solver.lpt(dlin, Q, a=1.0, order=2)


    dnonlin = solver.nbody(state, leapfrog([0.1, 1.0]))