# 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/hold.py
# class Solver(core.Solver):
#     def __init__(self, pm, cosmology, B=1, r_split=None, NTimeBin=4):
#         core.Solver.__init__(self, pm, cosmology, B)
# 
#         if r_split is None:
#             r_split = self.fpm.BoxSize[0] / self.fpm.Nmesh[0]
# 
#         self.r_split = r_split
#         self.r_cut = r_split * 4.5
#         self.r_smth = r_split / 32
#         self.NTimeBin = NTimeBin
# 
#     @property
#     def nbodystep(self):
#         return PPPMStep(self)
# 
#     def compute_shortrange(self, tree1, tree2, factor, out=None):
#         from .force.gravity import shortrange
#         return shortrange(tree1, tree2,
#             self.r_split, self.r_cut, self.r_smth,
#             factor=factor, out=out)
# 
#     def compute_longrange(self, X1, delta_k, factor):
#         from .force.gravity import longrange
#         return longrange(X1, delta_k, split=self.r_split, factor=factor)
# 
#     def compute_stepsize(self, tree, P, ac, E, Eprime, factor, out=None):
#         from .force.gravity import compute_stepsize
#         return compute_stepsize(
#                 tree,
#                 P, ac, E, Eprime,
#                 self.r_cut, self.r_smth,
#                 factor=factor, out=out)


from fastpm.core import leapfrog, autostages
from fastpm.hold import Solver
from pmesh.pm import ParticleMesh
from nbodykit.cosmology import Planck15, LinearPower
import numpy

pm = ParticleMesh(BoxSize=32., Nmesh=[16, 16, 16])

def test_solver():
    Plin = LinearPower(Planck15, redshift=0, transfer='EisensteinHu')

# 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/hold.py
# class Solver(core.Solver):
#     def __init__(self, pm, cosmology, B=1, r_split=None, NTimeBin=4):
#         core.Solver.__init__(self, pm, cosmology, B)
# 
#         if r_split is None:
#             r_split = self.fpm.BoxSize[0] / self.fpm.Nmesh[0]
# 
#         self.r_split = r_split
#         self.r_cut = r_split * 4.5
#         self.r_smth = r_split / 32
#         self.NTimeBin = NTimeBin
# 
#     @property
#     def nbodystep(self):
#         return PPPMStep(self)
# 
#     def compute_shortrange(self, tree1, tree2, factor, out=None):
#         from .force.gravity import shortrange
#         return shortrange(tree1, tree2,
#             self.r_split, self.r_cut, self.r_smth,
#             factor=factor, out=out)
# 
#     def compute_longrange(self, X1, delta_k, factor):
#         from .force.gravity import longrange
#         return longrange(X1, delta_k, split=self.r_split, factor=factor)
# 
#     def compute_stepsize(self, tree, P, ac, E, Eprime, factor, out=None):
#         from .force.gravity import compute_stepsize
#         return compute_stepsize(
#                 tree,
#                 P, ac, E, Eprime,
#                 self.r_cut, self.r_smth,
#                 factor=factor, out=out)


# Path: fastpm/tests/test_hold.py
from fastpm.core import leapfrog, autostages
from fastpm.hold import Solver
from pmesh.pm import ParticleMesh
from nbodykit.cosmology import Planck15, LinearPower
import numpy


pm = ParticleMesh(BoxSize=32., Nmesh=[16, 16, 16])

def test_solver():
    Plin = LinearPower(Planck15, redshift=0, transfer='EisensteinHu')

    solver = Solver(pm, Planck15, B=2)