# Path: fastpm/core.py
# class Solver(object):
#     def __init__(self, pm, cosmology, B=1, a_linear=1.0):
#         """
#             a_linear : float
#                 scaling factor at the time of the linear field.
#                 The growth function will be calibrated such that at a_linear D1 == 0.
# 
#         """
#         if not isinstance(cosmology, Cosmology):
#             raise TypeError("only nbodykit.cosmology object is supported")
# 
#         fpm = ParticleMesh(Nmesh=pm.Nmesh * B, BoxSize=pm.BoxSize, dtype=pm.dtype, comm=pm.comm, resampler=pm.resampler)
#         self.pm = pm
#         self.fpm = fpm
#         self.cosmology = cosmology
#         self.a_linear = a_linear
# 
#     # override nbodystep in subclasses
#     @property
#     def nbodystep(self):
#         return FastPMStep(self)
# 
#     def whitenoise(self, seed, unitary=False):
#         return self.pm.generate_whitenoise(seed, type='complex', unitary=unitary)
# 
#     def linear(self, whitenoise, Pk):
#         return whitenoise.apply(lambda k, v:
#                         Pk(sum(ki ** 2 for ki in k)**0.5) ** 0.5 * v / v.BoxSize.prod() ** 0.5)
# 
#     def lpt(self, linear, Q, a, order=2):
#         """ This computes the 'force' from LPT as well. """
#         assert order in (1, 2)
# 
#         from .force.lpt import lpt1, lpt2source
# 
#         state = StateVector(self, Q)
#         pt = MatterDominated(self.cosmology.Om0, a=[a], a_normalize=self.a_linear)
#         DX1 = pt.D1(a) * lpt1(linear, Q)
# 
#         V1 = a ** 2 * pt.f1(a) * pt.E(a) * DX1
#         if order == 2:
#             DX2 = pt.D2(a) * lpt1(lpt2source(linear), Q)
#             V2 = a ** 2 * pt.f2(a) * pt.E(a) * DX2
#             state.S[...] = DX1 + DX2
#             state.P[...] = V1 + V2
#             state.F[...] = a ** 2 * pt.E(a) * (pt.gf(a) / pt.D1(a) * DX1 + pt.gf2(a) / pt.D2(a) * DX2)
#         else:
#             state.S[...] = DX1
#             state.P[...] = V1
#             state.F[...] = a ** 2 * pt.E(a) * (pt.gf(a) / pt.D1(a) * DX1)
# 
#         state.a['S'] = a
#         state.a['P'] = a
# 
#         return state
# 
#     def nbody(self, state, stepping, monitor=None):
#         step = self.nbodystep
#         for action, ai, ac, af in stepping:
#             step.run(action, ai, ac, af, state, monitor)
# 
#         return state
# 
# 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


from numpy.testing import assert_allclose
from fastpm.core import Solver, leapfrog
from nbodykit.cosmology import Planck15, EHPower
from nbodykit.lab import FOF
from nbodykit.lab import ArrayCatalog
from pmesh.pm import ParticleMesh
import numpy
# This script tests if full reversability like
# JANUS is achieved by FastPM, without
# using int64 fixed point.
#
# JANUS: https://arxiv.org/pdf/1704.07715v1.pdf
#
# We do not need the high precision int64 representable
# because we only use very few steps (I think)
# Using float32 increases the error to about 1e-5.





pm = ParticleMesh(BoxSize=64, Nmesh=[64, 64, 64], dtype='f8')
Q = pm.generate_uniform_particle_grid()

stages = numpy.linspace(0.1, 1.0, 10, endpoint=True)


# Path: fastpm/core.py
# class Solver(object):
#     def __init__(self, pm, cosmology, B=1, a_linear=1.0):
#         """
#             a_linear : float
#                 scaling factor at the time of the linear field.
#                 The growth function will be calibrated such that at a_linear D1 == 0.
# 
#         """
#         if not isinstance(cosmology, Cosmology):
#             raise TypeError("only nbodykit.cosmology object is supported")
# 
#         fpm = ParticleMesh(Nmesh=pm.Nmesh * B, BoxSize=pm.BoxSize, dtype=pm.dtype, comm=pm.comm, resampler=pm.resampler)
#         self.pm = pm
#         self.fpm = fpm
#         self.cosmology = cosmology
#         self.a_linear = a_linear
# 
#     # override nbodystep in subclasses
#     @property
#     def nbodystep(self):
#         return FastPMStep(self)
# 
#     def whitenoise(self, seed, unitary=False):
#         return self.pm.generate_whitenoise(seed, type='complex', unitary=unitary)
# 
#     def linear(self, whitenoise, Pk):
#         return whitenoise.apply(lambda k, v:
#                         Pk(sum(ki ** 2 for ki in k)**0.5) ** 0.5 * v / v.BoxSize.prod() ** 0.5)
# 
#     def lpt(self, linear, Q, a, order=2):
#         """ This computes the 'force' from LPT as well. """
#         assert order in (1, 2)
# 
#         from .force.lpt import lpt1, lpt2source
# 
#         state = StateVector(self, Q)
#         pt = MatterDominated(self.cosmology.Om0, a=[a], a_normalize=self.a_linear)
#         DX1 = pt.D1(a) * lpt1(linear, Q)
# 
#         V1 = a ** 2 * pt.f1(a) * pt.E(a) * DX1
#         if order == 2:
#             DX2 = pt.D2(a) * lpt1(lpt2source(linear), Q)
#             V2 = a ** 2 * pt.f2(a) * pt.E(a) * DX2
#             state.S[...] = DX1 + DX2
#             state.P[...] = V1 + V2
#             state.F[...] = a ** 2 * pt.E(a) * (pt.gf(a) / pt.D1(a) * DX1 + pt.gf2(a) / pt.D2(a) * DX2)
#         else:
#             state.S[...] = DX1
#             state.P[...] = V1
#             state.F[...] = a ** 2 * pt.E(a) * (pt.gf(a) / pt.D1(a) * DX1)
# 
#         state.a['S'] = a
#         state.a['P'] = a
# 
#         return state
# 
#     def nbody(self, state, stepping, monitor=None):
#         step = self.nbodystep
#         for action, ai, ac, af in stepping:
#             step.run(action, ai, ac, af, state, monitor)
# 
#         return state
# 
# 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


# Path: examples/janus.py
from numpy.testing import assert_allclose
from fastpm.core import Solver, leapfrog
from nbodykit.cosmology import Planck15, EHPower
from nbodykit.lab import FOF
from nbodykit.lab import ArrayCatalog
from pmesh.pm import ParticleMesh
import numpy

# This script tests if full reversability like
# JANUS is achieved by FastPM, without
# using int64 fixed point.
#
# JANUS: https://arxiv.org/pdf/1704.07715v1.pdf
#
# We do not need the high precision int64 representable
# because we only use very few steps (I think)
# Using float32 increases the error to about 1e-5.





pm = ParticleMesh(BoxSize=64, Nmesh=[64, 64, 64], dtype='f8')
Q = pm.generate_uniform_particle_grid()

stages = numpy.linspace(0.1, 1.0, 10, endpoint=True)


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