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