# Path: fastpm/background.py
# class MatterDominated(Perturbation):
#     """
# 
#     Perturbation with matter dominated initial condition.
# 
#     This is usually referred to the single fluid approximation as well.
# 
#     The result here is accurate upto numerical precision. If Omega0_m 
# 
#     Parameters
#     ----------
#     Omega0_m :
#         matter density at redshift 0
# 
#     Omega0_lambda :
#         Lambda density at redshift 0, default None; set to ensure flat universe.
# 
#     Omega0_k:
#         Curvature density at redshift 0, default : 0
# 
#     a : array_like
#         a list of time steps where the factors are exact.
#         other a values are interpolated.
#     """
#     def __init__(self, Omega0_m, Omega0_lambda=None, Omega0_k=0, a=None, a_normalize=1.0):
#         if Omega0_lambda is None:
#             Omega0_lambda = 1 - Omega0_k - Omega0_m
# 
#         self.Omega0_lambda = Omega0_lambda
#         self.Omega0_m = Omega0_m
#         self.Omega0_k = Omega0_k
#         # Om0 is added for backward compatiblity
#         self.Om0 = Omega0_m
#         Perturbation.__init__(self, a, a_normalize)
# 
#     def get_initial_condition(self):
#         a0 = np.exp(self.lna[0])
# 
#         # matter dominated initial conditions
#         y0 = [a0, a0, -3./7 * a0**2, -6. / 7 *a0**2]
#         return y0
# 
#     def efunc(self, a):
#         return (self.Omega0_m / a**3+ self.Omega0_k / a**2 + self.Omega0_lambda) ** 0.5
# 
#     def efunc_prime(self, a):
#         return 0.5 / self.efunc(a) * (-3 * self.Omega0_m / (a**4) + -2 * self.Omega0_k / (a**3))
# 
#     def Om(self, a):
#         return (self.Omega0_m / a ** 3) / self.efunc(a) ** 2


import numpy
from pmesh.pm import ParticleMesh
from .background import MatterDominated
from nbodykit.cosmology import Cosmology
        from nbodykit.lab import ArrayCatalog
        from bigfile import FileMPI
        from .force.lpt import lpt1, lpt2source
        from .force.gravity import longrange
                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)


        state = StateVector(self, Q)

# Path: fastpm/background.py
# class MatterDominated(Perturbation):
#     """
# 
#     Perturbation with matter dominated initial condition.
# 
#     This is usually referred to the single fluid approximation as well.
# 
#     The result here is accurate upto numerical precision. If Omega0_m 
# 
#     Parameters
#     ----------
#     Omega0_m :
#         matter density at redshift 0
# 
#     Omega0_lambda :
#         Lambda density at redshift 0, default None; set to ensure flat universe.
# 
#     Omega0_k:
#         Curvature density at redshift 0, default : 0
# 
#     a : array_like
#         a list of time steps where the factors are exact.
#         other a values are interpolated.
#     """
#     def __init__(self, Omega0_m, Omega0_lambda=None, Omega0_k=0, a=None, a_normalize=1.0):
#         if Omega0_lambda is None:
#             Omega0_lambda = 1 - Omega0_k - Omega0_m
# 
#         self.Omega0_lambda = Omega0_lambda
#         self.Omega0_m = Omega0_m
#         self.Omega0_k = Omega0_k
#         # Om0 is added for backward compatiblity
#         self.Om0 = Omega0_m
#         Perturbation.__init__(self, a, a_normalize)
# 
#     def get_initial_condition(self):
#         a0 = np.exp(self.lna[0])
# 
#         # matter dominated initial conditions
#         y0 = [a0, a0, -3./7 * a0**2, -6. / 7 *a0**2]
#         return y0
# 
#     def efunc(self, a):
#         return (self.Omega0_m / a**3+ self.Omega0_k / a**2 + self.Omega0_lambda) ** 0.5
# 
#     def efunc_prime(self, a):
#         return 0.5 / self.efunc(a) * (-3 * self.Omega0_m / (a**4) + -2 * self.Omega0_k / (a**3))
# 
#     def Om(self, a):
#         return (self.Omega0_m / a ** 3) / self.efunc(a) ** 2


# Path: fastpm/core.py
import numpy
from pmesh.pm import ParticleMesh
from .background import MatterDominated
from nbodykit.cosmology import Cosmology
        from nbodykit.lab import ArrayCatalog
        from bigfile import FileMPI
        from .force.lpt import lpt1, lpt2source
        from .force.gravity import longrange

                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)


        state = StateVector(self, Q)

        pt = MatterDominated(self.cosmology.Om0, a=[a], a_normalize=self.a_linear)