# 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/state.py
# class StateVector(object):
#     def __init__(self, cosmology, species, comm):
#         """ A state vector is a dict of Species """
#         self.cosmology = cosmology
#         self.species = OrderedDict(sorted(species.items(), key=lambda t: t[0]))
#         self.comm = comm
#         self.a = dict(S=None, P=None, F=None)
# 
#     def __getitem__(self, spname):
#         return self.species[spname]
# 
#     def __iter__(self):
#         return iter(self.species)
# 
#     def __contains__(self, name):
#         return name in self.species
# 
#     def copy(self):
#         vc = {}
#         for k, v in self.species.items():
#             vc[k] = v.copy()
#         return StateVector(self.cosmology, vc, self.comm)
# 
#     def to_catalog(self, **kwargs):
#         from nbodykit.lab import MultipleSpeciesCatalog
#         names = []
#         sources = []
# 
#         for spname, sp in self.species.items():
#             sources.append(sp.to_catalog())
#             names.append(spname)
# 
#         cat = MultipleSpeciesCatalog(names, *sources, **kwargs)
#         return cat
# 
#     def save(self, filename, attrs={}):
#         from bigfile import FileMPI
#         a = self.a['S']
#         with FileMPI(self.comm, filename, create=True) as ff:
#             with ff.create('Header') as bb:
#                 keylist = ['Om0', 'Tcmb0', 'Neff', 'Ob0', 'Ode0']
#                 if getattr(self.cosmology, 'm_nu', None) is not None:
#                     keylist.insert(3,'m_nu')
#                 for key in keylist:
#                     bb.attrs[key] = getattr(self.cosmology, key)
#                 bb.attrs['Time'] = a
#                 bb.attrs['h'] = self.cosmology.H0 / H0 # relative h
#                 bb.attrs['RSDFactor'] = 1.0 / (H0 * a * self.cosmology.efunc(1.0 / a - 1))
#                 for key in attrs:
#                     try:
#                         #best effort
#                         bb.attrs[key] = attrs[key]
#                     except:
#                         pass
# 
#             for k, v in sorted(self.species.items()):
#                 v.save(filename, k)
# 
# class Matter(Species):
#     def Omega(self, a):
#         return self.cosmology.Om(z=1. / a - 1)
# 
# class Baryon(Species):
#     def Omega(self, a):
#         return self.cosmology.Ob(z=1. / a - 1)
# 
# class CDM(Species):
#     def Omega(self, a):
#         return self.cosmology.Odm(z=1. / a - 1)
# 
# class NCDM(Species):
#     def Omega(self, a):
#         # FIXME: using Omega_ncdm after switching to new nbodykit cosmology.
#         return self.cosmology.Onu(z=1. / a - 1)


import numpy
from pmesh.pm import ParticleMesh
from .background import PerturbationGrowth
from nbodykit.cosmology import Cosmology
from .state import StateVector, Matter, Baryon, CDM, NCDM
        from .force.lpt import lpt1, lpt2source
        from .force.gravity import longrange_batch
            source = primordial.apply(apply_density)
            source1[...] += sp.Omega(a) * source
            Omega_tot += sp.Omega(a)

            DX1 = lpt1(source, Q)
            sp.S[...] = DX1

            def apply_velocity(k, v):
                return a * self.cosmology.efunc(z) * dd(get_k(k)) * v

            source = primordial.apply(apply_velocity)
            sp.P[...] = a ** 2 * lpt1(source, Q)
            sp.F[...] = 0

            sp.a['S'] = a
            sp.a['P'] = a

            species[spname] = sp

        if order == 2:
            source1[...] /= Omega_tot
            source1[...] /= pt.D1(a) # goto z=0

            source2 = lpt2source(source1)

            for spname, sp in species.items():
                DX2 = lpt1(source2, Q)
                sp.S[...] += pt.D2(a) * DX2
                sp.V[...] += a ** 2 * pt.f2(a) * pt.E(a) * pt.D2(a) * DX2


# 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/state.py
# class StateVector(object):
#     def __init__(self, cosmology, species, comm):
#         """ A state vector is a dict of Species """
#         self.cosmology = cosmology
#         self.species = OrderedDict(sorted(species.items(), key=lambda t: t[0]))
#         self.comm = comm
#         self.a = dict(S=None, P=None, F=None)
# 
#     def __getitem__(self, spname):
#         return self.species[spname]
# 
#     def __iter__(self):
#         return iter(self.species)
# 
#     def __contains__(self, name):
#         return name in self.species
# 
#     def copy(self):
#         vc = {}
#         for k, v in self.species.items():
#             vc[k] = v.copy()
#         return StateVector(self.cosmology, vc, self.comm)
# 
#     def to_catalog(self, **kwargs):
#         from nbodykit.lab import MultipleSpeciesCatalog
#         names = []
#         sources = []
# 
#         for spname, sp in self.species.items():
#             sources.append(sp.to_catalog())
#             names.append(spname)
# 
#         cat = MultipleSpeciesCatalog(names, *sources, **kwargs)
#         return cat
# 
#     def save(self, filename, attrs={}):
#         from bigfile import FileMPI
#         a = self.a['S']
#         with FileMPI(self.comm, filename, create=True) as ff:
#             with ff.create('Header') as bb:
#                 keylist = ['Om0', 'Tcmb0', 'Neff', 'Ob0', 'Ode0']
#                 if getattr(self.cosmology, 'm_nu', None) is not None:
#                     keylist.insert(3,'m_nu')
#                 for key in keylist:
#                     bb.attrs[key] = getattr(self.cosmology, key)
#                 bb.attrs['Time'] = a
#                 bb.attrs['h'] = self.cosmology.H0 / H0 # relative h
#                 bb.attrs['RSDFactor'] = 1.0 / (H0 * a * self.cosmology.efunc(1.0 / a - 1))
#                 for key in attrs:
#                     try:
#                         #best effort
#                         bb.attrs[key] = attrs[key]
#                     except:
#                         pass
# 
#             for k, v in sorted(self.species.items()):
#                 v.save(filename, k)
# 
# class Matter(Species):
#     def Omega(self, a):
#         return self.cosmology.Om(z=1. / a - 1)
# 
# class Baryon(Species):
#     def Omega(self, a):
#         return self.cosmology.Ob(z=1. / a - 1)
# 
# class CDM(Species):
#     def Omega(self, a):
#         return self.cosmology.Odm(z=1. / a - 1)
# 
# class NCDM(Species):
#     def Omega(self, a):
#         # FIXME: using Omega_ncdm after switching to new nbodykit cosmology.
#         return self.cosmology.Onu(z=1. / a - 1)


# Path: fastpm/multi.py
import numpy
from pmesh.pm import ParticleMesh
from .background import PerturbationGrowth
from nbodykit.cosmology import Cosmology
from .state import StateVector, Matter, Baryon, CDM, NCDM
        from .force.lpt import lpt1, lpt2source
        from .force.gravity import longrange_batch

            source = primordial.apply(apply_density)
            source1[...] += sp.Omega(a) * source
            Omega_tot += sp.Omega(a)

            DX1 = lpt1(source, Q)
            sp.S[...] = DX1

            def apply_velocity(k, v):
                return a * self.cosmology.efunc(z) * dd(get_k(k)) * v

            source = primordial.apply(apply_velocity)
            sp.P[...] = a ** 2 * lpt1(source, Q)
            sp.F[...] = 0

            sp.a['S'] = a
            sp.a['P'] = a

            species[spname] = sp

        if order == 2:
            source1[...] /= Omega_tot
            source1[...] /= pt.D1(a) # goto z=0

            source2 = lpt2source(source1)

            for spname, sp in species.items():
                DX2 = lpt1(source2, Q)
                sp.S[...] += pt.D2(a) * DX2
                sp.V[...] += a ** 2 * pt.f2(a) * pt.E(a) * pt.D2(a) * DX2


        state = StateVector(self.cosmology, species, self.pm.comm)