class Test_r2c_c2r(BaseScalarTest):
to_scalar = staticmethod(fastpm.to_scalar)
pm = fastpm.ParticleMesh(Nmesh=[4, 4], BoxSize=8.0, comm=MPI.COMM_SELF)
x = pm.generate_whitenoise(seed=300, unitary=True, type='real')
y = x.cnorm()
x_ = create_bases(x)
def model(self, x):
c = fastpm.r2c(x)
r = fastpm.c2r(c)
return r
class Test_cdot(BaseScalarTest):
to_scalar = staticmethod(linalg.to_scalar)
pm = fastpm.ParticleMesh(Nmesh=[4, 4, 4], BoxSize=8.0, comm=MPI.COMM_SELF)
x = pm.generate_whitenoise(seed=300, unitary=True, type='real')
x1 = pm.generate_whitenoise(seed=301, unitary=True, type='real')
y = (x.cdot(x1))**2 * pm.Nmesh.prod()**-2.
x_ = create_bases(x)
def model(self, x):
x1 = fastpm.r2c(self.x1)
x = fastpm.r2c(x)
return fastpm.cdot(x, x1)
class Test_lpt2src(BaseScalarTest):
to_scalar = staticmethod(fastpm.to_scalar)
pm = fastpm.ParticleMesh(Nmesh=[4, 4, 4], BoxSize=8.0, comm=MPI.COMM_SELF)
x = pm.generate_whitenoise(seed=300, unitary=True, type='real')
y = NotImplemented
x_ = create_bases(x)
epsilon = 1e-4
rtol = 1e-6
def model(self, x):
return fastpm.lpt2src(fastpm.r2c(x), pm=self.pm)
class Test_lpt(BaseScalarTest):
to_scalar = staticmethod(linalg.to_scalar)
pm = fastpm.ParticleMesh(Nmesh=[4, 4, 4], BoxSize=8.0, comm=MPI.COMM_SELF)
x = pm.generate_whitenoise(seed=300, unitary=True, type='real')
pos = pm.generate_uniform_particle_grid(shift=0.5)
y = NotImplemented
x_ = create_bases(x)
epsilon = 1e-4
def model(self, x):
dx1, dx2 = fastpm.lpt(fastpm.r2c(x), q=self.pos, pm=self.pm)
return linalg.add(dx1, dx2)
class Test_readout_mesh(BaseScalarTest):
to_scalar = staticmethod(linalg.to_scalar)
pm = fastpm.ParticleMesh(Nmesh=[4, 4, 4], BoxSize=8.0, comm=MPI.COMM_SELF)
x = pm.generate_whitenoise(seed=300, unitary=True, type='real')
pos = pm.generate_uniform_particle_grid(shift=0.5)
y = (x.readout(pos)**2).sum()
x_ = create_bases(x)
epsilon = 1e-3
def model(self, x):
y = fastpm.readout(x, self.pos, layout=None)
return y
class Test_paint_mass(BaseScalarTest):
to_scalar = staticmethod(fastpm.to_scalar)
pm = fastpm.ParticleMesh(Nmesh=[4, 4, 4], BoxSize=8.0, comm=MPI.COMM_SELF)
mesh = pm.generate_whitenoise(seed=300, unitary=True, type='real')
pos = pm.generate_uniform_particle_grid(shift=0.1)
x = numpy.ones(len(pos))
y = x.sum()
x_ = create_bases(x)
epsilon = 1e-3
def model(self, x):
y = fastpm.paint(self.pos, layout=None, mass=x, pm=self.pm)
return y
class Test_paint_x(BaseScalarTest):
to_scalar = staticmethod(fastpm.to_scalar)
pm = fastpm.ParticleMesh(Nmesh=[4, 4, 4], BoxSize=8.0, comm=MPI.COMM_SELF)
mesh = pm.generate_whitenoise(seed=300, unitary=True, type='real')
x = pm.generate_uniform_particle_grid(shift=0.1)
y = (1 + mesh).cnorm()
x_ = create_bases(x)
epsilon = 1e-3
def model(self, x):
y = fastpm.paint(x, layout=None, mass=1.0, pm=self.pm)
y = linalg.add(y,
self.mesh) # biasing a bit to get non-zero derivatives.
return y
class Test_apply_digitized_amp_x(BaseScalarTest):
to_scalar = staticmethod(fastpm.to_scalar)
pm = fastpm.ParticleMesh(Nmesh=[4, 4], BoxSize=4.0, comm=MPI.COMM_SELF)
kedges = numpy.linspace(0, 2 * numpy.pi / 4.0 * 3, 8)
tf = numpy.arange(len(kedges) - 1) + 1
x = pm.generate_whitenoise(seed=300, unitary=True, type='real')
x_ = create_bases(x)
y = NotImplemented
def model(self, x):
c = fastpm.r2c(x)
digitizer = fastpm.apply_digitized.isotropic_wavenumber(self.kedges)
c = fastpm.apply_digitized(c,
tf=self.tf,
digitizer=digitizer,
kind='wavenumber')
r = fastpm.c2r(c)
return r
def test_add_scalar():
from vmad import autooperator
from vmad.core.stdlib import eval
from numpy.testing import assert_array_equal
pm = fastpm.ParticleMesh(Nmesh=[4, 4, 4], BoxSize=8.0, comm=MPI.COMM_SELF)
x = pm.generate_whitenoise(seed=300, unitary=True, type='real')
x[...] = 1.0
@autooperator('x, n->y')
def func(x, n):
a = linalg.take(n, 0, axis=0)
a = linalg.broadcast_to(a, eval(x, lambda x: x.shape))
return x + a
n = numpy.array([3, 4])
(y, ), (_x, _n) = (func.build().compute_with_vjp(init=dict(x=x, n=n),
v=dict(_y=x)))
assert_array_equal(y, 4)
assert_array_equal(_x, 1)
assert_array_equal(_n, (64, 0))
class Test_nbody(BaseScalarTest):
to_scalar = staticmethod(linalg.to_scalar)
pm = fastpm.ParticleMesh(Nmesh=[4, 4, 4], BoxSize=8.0, comm=MPI.COMM_SELF)
x = pm.generate_whitenoise(seed=300, unitary=True, type='real')
pos = pm.generate_uniform_particle_grid(shift=0.5)
y = NotImplemented
x_ = create_bases(x)
epsilon = 1e-4
def model(self, x):
from nbodykit.cosmology import Planck15
dx, p, f = fastpm.nbody(fastpm.r2c(x),
q=self.pos,
stages=[0.1, 0.5, 1.0],
pm=self.pm,
cosmology=Planck15)
return linalg.stack([dx, p, f], axis=-1)
class Test_apply_digitized_phase_tf(BaseScalarTest):
to_scalar = staticmethod(fastpm.to_scalar)
pm = fastpm.ParticleMesh(Nmesh=[4, 4, 4], BoxSize=4.0, comm=MPI.COMM_SELF)
kedges = numpy.linspace(0, 2 * numpy.pi / 4.0 * 3, 8)
y = NotImplemented
x = numpy.linspace(-1 * numpy.pi, 4 * numpy.pi, (len(kedges) - 1))
x_ = create_bases(x)
atol = 1e-9
def model(self, x):
c0 = self.pm.generate_whitenoise(seed=300,
unitary=True,
type='complex',
mean=1.0)
digitizer = fastpm.apply_digitized.isotropic_wavenumber(self.kedges)
c = fastpm.apply_digitized(c0,
tf=x,
digitizer=digitizer,
kind='wavenumber',
mode='phase')
r = fastpm.c2r(c)
return r
from vmad import Builder
from fastpm.force.lpt import lpt1, lpt2source
from vmad.lib import fastpm
import numpy
from nbodykit.cosmology import Planck15, LinearPower
pm = fastpm.ParticleMesh([32, 32, 32], BoxSize=128.)
powerspectrum = LinearPower(Planck15, 0)
q = pm.generate_uniform_particle_grid()
with Builder() as model:
x = model.input('x')
wnk = fastpm.as_complex_field(x, pm)
rhok = fastpm.induce_correlation(wnk, powerspectrum, pm)
dx1, dx2 = fastpm.lpt(rhok, q, pm)
dx, p, f = fastpm.nbody(rhok, q, [0.1, 0.6, 1.0], Planck15, pm)
dx0, p0, f0 = fastpm.nbody(rhok, q, [0.1], Planck15, pm)
model.output(dx1=dx1, dx2=dx2, dx=dx, p=p, dx0=dx0, p0=p0, f0=f0, f=f)
wn = pm.generate_whitenoise(555, unitary=True)
x = wn[...]
x = numpy.stack([x.real, x.imag], -1)
from fastpm.core import Solver, leapfrog
def run_wl_sim(params, num, cosmo, randseed=187):
'''
params: dictionary, of run specific settings
num: int, number of this run (which run out of #N_maps)
cosmo: nbodykit cosmology object, see https://nbodykit.readthedocs.io/en/latest/index.html
label: string, label of this run. used in filename if 3D matter distribution is saved
randseed:random seed for generating initial conditions
'''
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
if params['logging']:
a_logger = logging.getLogger()
lformat = logging.Formatter('%(asctime)s - %(message)s')
output_file_handler = logging.FileHandler("proc%d.log" % rank)
stdout_handler = logging.StreamHandler(sys.stdout)
output_file_handler.setFormatter(lformat)
a_logger.addHandler(output_file_handler)
a_logger.addHandler(stdout_handler)
else:
a_logger = None
# particle mesh for fastpm simulation
pm = fastpm.ParticleMesh(Nmesh=params['Nmesh'],
BoxSize=params['BoxSize'],
comm=MPI.COMM_WORLD,
resampler='cic')
# 2D FOV in radians
BoxSize2D = [deg / 180. * np.pi for deg in params['BoxSize2D']]
np.random.seed(randseed)
randseeds = np.random.randint(0, 1e6, 100)
# generate initial conditions
cosmo = cosmo.clone(P_k_max=30)
rho = pm.generate_whitenoise(seed=randseeds[num],
unitary=False,
type='complex')
#set zero mode to zero
rho.csetitem([0, 0, 0], 0)
rho = rho.c2r()
if params['logging']:
logging.info('simulations starts')
# weak lensing simulation object
wlsim = WLSimulation(stages=numpy.linspace(0.1,
1.0,
params['N_steps'],
endpoint=True),
cosmology=cosmo,
pm=pm,
boxsize2D=BoxSize2D,
params=params,
logger=a_logger)
#build
if params['interpolate']:
model = wlsim.run_interpolated.build()
else:
model = wlsim.run.build()
# results
kmap_vjp, kmap_jvp = [None, None]
# compute
if params['forward'] and (params['vjp'] or params['jvp']):
kmaps, tape = model.compute(vout='kmaps',
init=dict(rhok=rho),
return_tape=True)
if params['vjp']:
vjp = tape.get_vjp()
kmap_vjp = vjp.compute(init=dict(_kmaps=kmaps), vout='_rhok')
if params['jvp']:
jvp = tape.get_jvp()
kmap_jvp = jvp.compute(init=dict(rhok_=rho), vout='kmaps_')
if params['forward']:
kmaps = model.compute(vout='kmaps', init=dict(rhok=rho))
return kmaps, [kmap_vjp, kmap_jvp], pm
