Esempio n. 1
0
def test_2d_2d(comm):
    import pfft
    pm = ParticleMesh(BoxSize=8.0,
                      Nmesh=[8, 8],
                      np=pfft.split_size_2d(comm.size),
                      comm=comm,
                      dtype='f8')
    real = pm.generate_whitenoise(seed=123, mode='real')
    complex = pm.generate_whitenoise(seed=123, mode='complex')
    assert_array_equal(real, complex.c2r())
Esempio n. 2
0
def test_2d_2d(comm):
    import pfft
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], np=pfft.split_size_2d(comm.size), comm=comm, dtype='f8')
    pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], np=pfft.split_size_2d(comm.size), comm=comm, dtype='f8')
    assert pm._use_padded == False
    real = pm.generate_whitenoise(seed=123, type='real')
    complex = pm.generate_whitenoise(seed=123, type='complex')
    assert_array_equal(real, complex.c2r())

    real2 = pm2.generate_whitenoise(seed=123, type='real')

    assert real2.shape[:2] == real.shape
Esempio n. 3
0
def test_cdot_types(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
    comp1 = pm.generate_whitenoise(1234, type='complex')
    comp2 = pm.generate_whitenoise(1239, type='untransposedcomplex')

    with pytest.raises(TypeError):
        norm1 = comp1.cdot(comp2)
        norm2 = comp2.cdot(comp1)

    # this should work though fragile, because only on the Nmesh and 1 rank
    # the shape of values is the same.
    norm1 = comp1.cdot(comp2.value)
    norm2 = comp2.cdot(comp1.value)
Esempio n. 4
0
def test_cdot_types(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
    comp1 = pm.generate_whitenoise(1234, type='complex')
    comp2 = pm.generate_whitenoise(1239, type='untransposedcomplex')

    with pytest.raises(TypeError):
        norm1 = comp1.cdot(comp2)
        norm2 = comp2.cdot(comp1)

    # this should work though fragile, because only on the Nmesh and 1 rank
    # the shape of values is the same.
    norm1 = comp1.cdot(comp2.value)
    norm2 = comp2.cdot(comp1.value)
Esempio n. 5
0
def test_cnorm_log(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
    comp1 = pm.generate_whitenoise(1234, type='complex', mean=1.0)

    norm2 = comp1.cnorm(norm = lambda x: numpy.log(x.real ** 2 + x.imag ** 2))
    norm3 = (numpy.log(abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value))) ** 2)).sum()
    assert_allclose(norm2, norm3)
Esempio n. 6
0
def test_whitenoise_mean(comm):
    # the whitenoise shall preserve the large scale.
    pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8')

    complex1 = pm0.generate_whitenoise(seed=8, unitary=True, mean=1.0)

    assert_allclose(complex1.c2r().cmean(), 1.0)
Esempio n. 7
0
def test_whitenoise(comm):
    # the whitenoise shall preserve the large scale.
    pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8')
    pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[16, 16, 16], comm=comm, dtype='f8')
    pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[32, 32, 32], comm=comm, dtype='f8')
    complex1_down = ComplexField(pm0)
    complex2_down = ComplexField(pm0)

    complex1 = pm1.generate_whitenoise(seed=8, unitary=True)
    complex2 = pm2.generate_whitenoise(seed=8, unitary=True)

    complex1.resample(complex1_down)
    complex2.resample(complex2_down)

    mask1 = complex1_down.value != complex2_down.value
    assert_array_equal(complex1_down.value, complex2_down.value)
Esempio n. 8
0
def test_kdk(comm):
    # Or use an object from astropy.
    cosmo = lambda : None
    cosmo.Om0 = 0.3
    cosmo.Ode0 = 0.7
    cosmo.Ok0 = 0.0

    pm = ParticleMesh(BoxSize=128.0, Nmesh=(4,4,4), comm=comm, dtype='f8')
    vm = fastpm.Evolution(pm, B=1, shift=0.5)
    dlink = pm.generate_whitenoise(12345, mode='complex', unitary=True)

    dlink, ampl = _addampl(dlink)

    data = dlink.c2r()
    data[...] = 0
    sigma = data.copy()
    sigma[...] = 1.0

    code = vm.code()

    code.Decompress(C='dlin_k')
    code.KDKSimulation(dlin_k='dlin_k', mesh='mesh', cosmo=cosmo, astart=0.1, aend=1.0, Nsteps=5)
    code.Decompose()
    code.Paint()
    code.Residual(data_x=data, sigma_x=sigma)
    code.Chi2(variable='residual')

    _test_model(code, dlink, ampl)
Esempio n. 9
0
def test_preview(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')

    comp1 = pm.generate_whitenoise(1234, type='real')

    preview = comp1.preview(axes=(0, 1, 2))

    preview = comp1.preview(Nmesh=4, axes=(0, 1, 2))

    for ind1 in numpy.ndindex(*(list(comp1.cshape))):
        assert_allclose(preview[ind1], comp1.cgetitem(ind1))

    preview1 = comp1.preview(Nmesh=4, axes=(0, 1))
    previewsum1 = preview.sum(axis=2)
    assert_allclose(preview1, previewsum1)

    preview2 = comp1.preview(Nmesh=4, axes=(1, 2))
    previewsum2 = preview.sum(axis=0)
    assert_allclose(preview2, previewsum2)

    preview3 = comp1.preview(Nmesh=4, axes=(0, 2))
    previewsum3 = preview.sum(axis=1)
    assert_allclose(preview3, previewsum3)

    preview4 = comp1.preview(Nmesh=4, axes=(2, 0))
    previewsum4 = preview.sum(axis=1).T
    assert_allclose(preview4, previewsum4)

    preview5 = comp1.preview(Nmesh=4, axes=(0, ))
    previewsum5 = preview.sum(axis=(1, 2))
    assert_allclose(preview5, previewsum5)

    preview6 = comp1.preview(Nmesh=8, axes=(0, ))
Esempio n. 10
0
def test_whitenoise(comm):
    # the whitenoise shall preserve the large scale.
    pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8')
    pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[16, 16, 16], comm=comm, dtype='f8')
    pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[32, 32, 32], comm=comm, dtype='f8')
    complex1_down = ComplexField(pm0)
    complex2_down = ComplexField(pm0)

    complex1 = pm1.generate_whitenoise(seed=8, unitary=True)
    complex2 = pm2.generate_whitenoise(seed=8, unitary=True)

    complex1.resample(complex1_down)
    complex2.resample(complex2_down)

    mask1 = complex1_down.value != complex2_down.value
    assert_array_equal(complex1_down.value, complex2_down.value)
Esempio n. 11
0
def test_nody():
    """ Checking end to end nbody
  """
    a0 = 0.1

    pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
    grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)
    solver = Solver(pm, Planck15, B=1)
    stages = np.linspace(0.1, 1.0, 10, endpoint=True)

    # Generate initial state with fastpm
    whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
    lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
        sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)
    statelpt = solver.lpt(lineark, grid, a0, order=1)
    finalstate = solver.nbody(statelpt, leapfrog(stages))
    final_cube = pm.paint(finalstate.X)

    # Same thing with flowpm
    tlinear = tf.expand_dims(np.array(lineark.c2r()), 0)
    state = tfpm.lpt_init(tlinear, a0, order=1)
    state = tfpm.nbody(state, stages, nc)
    tfread = pmutils.cic_paint(tf.zeros_like(tlinear), state[0]).numpy()

    assert_allclose(final_cube, tfread[0], atol=1.2)
Esempio n. 12
0
def test_gravity(comm):
    from pmesh.pm import ParticleMesh
    import fastpm.operators as operators

    pm = ParticleMesh(BoxSize=4.0, Nmesh=(4, 4), comm=comm, method='cic', dtype='f8')

    vm = fastpm.Evolution(pm, shift=0.5)

    dlink = pm.generate_whitenoise(12345, mode='complex')

    code = vm.code()
    code.Decompress(C='dlin_k')
    code.LPTDisplace(dlin_k='dlin_k', D1=1.0, v1=0, D2=0.0, v2=0.0)
    code.ForcePaint()
    code.Force(factor=0.1)
    code.Chi2(variable='f')
    print(code)
    # FIXME: without the shift some particles have near zero dx1.
    # or near 1 dx1.
    # the gradient is not well approximated by the numerical if
    # any of the left or right value shifts beyond the support of
    # the window.
    #

    dlink, ampl = _addampl(dlink)
    _test_model(code, dlink, ampl)
Esempio n. 13
0
def test_cnorm_grad(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')

    comp1 = pm.generate_whitenoise(1234, type='complex')

    def objective(comp1):
        return comp1.cnorm()

    grad_comp1 = comp1 * 2
    grad_comp1.decompress_vjp(grad_comp1)

    ng = []
    ag = []
    ind = []
    dx = 1e-7
    for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])):
        dx1, c1 = perturb(comp1, ind1, dx)
        ng1 = (objective(c1) - objective(comp1)) / dx
        ag1 = grad_comp1.cgetitem(ind1) * dx1 / dx
        if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))):
            print (ind1, 'a', ag1, 'n', ng1)
        comm.barrier()
        ng.append(ng1)
        ag.append(ag1)
        ind.append(ind1)

    assert_allclose(ng, ag, rtol=1e-5, atol=1e-6)
Esempio n. 14
0
def test_paint(comm):

    from pmesh.pm import ParticleMesh

    # initialize a random white noise field in real space
    pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm)
    real = pm.generate_whitenoise(type='real', seed=3333) # a RealField

    # FFT to a ComplexField
    complex = real.r2c()

    # gather to all ranks --> shape is now (8,8,8) on all ranks
    data = numpy.concatenate(comm.allgather(numpy.array(real.ravel())))\
                 .reshape(real.cshape)

    # mesh from data array
    realmesh = ArrayMesh(data, BoxSize=128., comm=comm)

    # paint() must yield the RealField
    assert_array_equal(real, realmesh.to_field(mode='real'))

    # gather complex to all ranks --> shape is (8,8,5) on all ranks
    cdata = numpy.concatenate(comm.allgather(numpy.array(complex.ravel())))\
                 .reshape(complex.cshape)

    # mesh from complex array
    cmesh = ArrayMesh(cdata, BoxSize=128., comm=comm)

    # paint() yields the ComplexField
    assert_allclose(complex, cmesh.to_field(mode='complex'), atol=1e-8)
Esempio n. 15
0
def test_c2r_vjp(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')

    real = pm.generate_whitenoise(1234, mode='real')
    comp = real.r2c()

    def objective(comp):
        real = comp.c2r()
        obj = (real.value ** 2).sum()
        return comm.allreduce(obj)

    grad_real = RealField(pm)
    grad_real[...] = real[...] * 2
    grad_comp = ComplexField(pm)
    grad_comp = grad_real.c2r_vjp(grad_real)
    grad_comp.decompress_vjp(grad_comp)

    ng = []
    ag = []
    ind = []
    dx = 1e-7
    for ind1 in numpy.ndindex(*(list(grad_comp.cshape) + [2])):
        dx1, c1 = perturb(comp, ind1, dx)
        ng1 = (objective(c1) - objective(comp)) / dx
        ag1 = grad_comp.cgetitem(ind1) * dx1 / dx
        comm.barrier()
        ng.append(ng1)
        ag.append(ag1)
        ind.append(ind1)

    assert_allclose(ng, ag, rtol=1e-5)
Esempio n. 16
0
def test_c2r_vjp(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')

    real = pm.generate_whitenoise(1234, type='real', mean=1.0)
    comp = real.r2c()

    def objective(comp):
        real = comp.c2r()
        obj = (real.value ** 2).sum()
        return comm.allreduce(obj)

    grad_real = RealField(pm)
    grad_real[...] = real[...] * 2
    grad_comp = ComplexField(pm)
    grad_comp = grad_real.c2r_vjp(grad_real)
    grad_comp.decompress_vjp(grad_comp)

    ng = []
    ag = []
    ind = []
    dx = 1e-7
    for ind1 in numpy.ndindex(*(list(grad_comp.cshape) + [2])):
        dx1, c1 = perturb(comp, ind1, dx)
        ng1 = (objective(c1) - objective(comp)) / dx
        ag1 = grad_comp.cgetitem(ind1) * dx1 / dx
        comm.barrier()
        ng.append(ng1)
        ag.append(ag1)
        ind.append(ind1)

    assert_allclose(ng, ag, rtol=1e-5)
Esempio n. 17
0
def test_cnorm_grad(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')

    comp1 = pm.generate_whitenoise(1234, mode='complex')

    def objective(comp1):
        return comp1.cnorm()

    grad_comp1 = comp1 * 2
    grad_comp1.decompress_vjp(grad_comp1)

    ng = []
    ag = []
    ind = []
    dx = 1e-7
    for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])):
        dx1, c1 = perturb(comp1, ind1, dx)
        ng1 = (objective(c1) - objective(comp1)) / dx
        ag1 = grad_comp1.cgetitem(ind1) * dx1 / dx
        if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))):
            print (ind1, 'a', ag1, 'n', ng1)
        comm.barrier()
        ng.append(ng1)
        ag.append(ag1)
        ind.append(ind1)

    assert_allclose(ng, ag, rtol=1e-5, atol=1e-6)
Esempio n. 18
0
def test_whitenoise_mean(comm):
    # the whitenoise shall preserve the large scale.
    pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8')

    complex1 = pm0.generate_whitenoise(seed=8, unitary=True, mean=1.0)

    assert_allclose(complex1.c2r().cmean(), 1.0)
Esempio n. 19
0
def test_preview(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')

    comp1 = pm.generate_whitenoise(1234, type='real')

    preview = comp1.preview(axes=(0, 1, 2))

    preview = comp1.preview(Nmesh=4, axes=(0, 1, 2))

    for ind1 in numpy.ndindex(*(list(comp1.cshape))):
        assert_allclose(preview[ind1], comp1.cgetitem(ind1))

    preview1 = comp1.preview(Nmesh=4, axes=(0, 1))
    previewsum1 = preview.sum(axis=2)
    assert_allclose(preview1, previewsum1)

    preview2 = comp1.preview(Nmesh=4, axes=(1, 2))
    previewsum2 = preview.sum(axis=0)
    assert_allclose(preview2, previewsum2)

    preview3 = comp1.preview(Nmesh=4, axes=(0, 2))
    previewsum3 = preview.sum(axis=1)
    assert_allclose(preview3, previewsum3)

    preview4 = comp1.preview(Nmesh=4, axes=(2, 0))
    previewsum4 = preview.sum(axis=1).T
    assert_allclose(preview4, previewsum4)

    preview5 = comp1.preview(Nmesh=4, axes=(0,))
    previewsum5 = preview.sum(axis=(1, 2))
    assert_allclose(preview5, previewsum5)

    preview6 = comp1.preview(Nmesh=8, axes=(0,))
Esempio n. 20
0
def test_cnorm_log(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
    comp1 = pm.generate_whitenoise(1234, type='complex', mean=1.0)

    norm2 = comp1.cnorm(norm=lambda x: numpy.log(x.real**2 + x.imag**2))
    norm3 = (numpy.log(abs(numpy.fft.fftn(numpy.fft.irfftn(
        comp1.value)))**2)).sum()
    assert_allclose(norm2, norm3)
Esempio n. 21
0
def test_cdot_cnorm(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
    comp1 = pm.generate_whitenoise(1234, type='complex')

    norm1 = comp1.cdot(comp1)
    norm2 = comp1.cnorm()
    norm3 = (abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value)))**2).sum()
    assert_allclose(norm2, norm3)
    assert_allclose(norm2, norm1)
Esempio n. 22
0
def test_cdot_cnorm(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
    comp1 = pm.generate_whitenoise(1234, type='complex')

    norm1 = comp1.cdot(comp1)
    norm2 = comp1.cnorm()
    norm3 = (abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value))) ** 2).sum()
    assert_allclose(norm2, norm3)
    assert_allclose(norm2, norm1)
Esempio n. 23
0
def test_whitenoise_untransposed(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4')

    f1 = pm.generate_whitenoise(seed=3333, type='untransposedcomplex')
    f2 = pm.generate_whitenoise(seed=3333, type='transposedcomplex')

    f1r = numpy.concatenate(comm.allgather(numpy.array(f1.ravel())))
    f2r = numpy.concatenate(comm.allgather(numpy.array(f2.ravel())))

    assert_array_equal(f1r, f2r)

    # this should have asserted r2c transforms as well.
    r1 = f1.c2r()
    r2 = f2.c2r()

    r1r = numpy.concatenate(comm.allgather(numpy.array(r1.ravel())))
    r2r = numpy.concatenate(comm.allgather(numpy.array(r2.ravel())))

    assert_array_equal(r1r, r2r)
Esempio n. 24
0
def test_whitenoise_untransposed(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4')

    f1 = pm.generate_whitenoise(seed=3333, type='untransposedcomplex')
    f2 = pm.generate_whitenoise(seed=3333, type='transposedcomplex')

    f1r = numpy.concatenate(comm.allgather(numpy.array(f1.ravel())))
    f2r = numpy.concatenate(comm.allgather(numpy.array(f2.ravel())))

    assert_array_equal(f1r, f2r)

    # this should have asserted r2c transforms as well.
    r1 = f1.c2r()
    r2 = f2.c2r()

    r1r = numpy.concatenate(comm.allgather(numpy.array(r1.ravel())))
    r2r = numpy.concatenate(comm.allgather(numpy.array(r2.ravel())))

    assert_array_equal(r1r, r2r)
Esempio n. 25
0
def make_data(bs, nc, seed, nsteps, path='', z=0):
    #initiate
    pm = ParticleMesh(BoxSize=bs, Nmesh=(nc, nc, nc), dtype='f8')

    ofolder = path + 'z%02d/L%04d_N%04d_S%04d_%02dstep/'%(z*10, bs, nc, seed, nsteps)
#        
    
    if pm.comm.rank == 0:
        print('Make data for seed = %d'%seed)
        print('Data to be saved at path - %s'%ofolder)

    klin, plin = numpy.loadtxt('ics_matterpow_0.dat', unpack = True)
    pk = interpolate(klin, plin)
    cosmo = Planck15.clone(Omega_cdm = 0.2685, h = 0.6711, Omega_b = 0.049)
    #pk = EHPower(Planck15, redshift=0)
    #cosmo = Planck15

    s_truth = pm.generate_whitenoise(seed, type='complex')\
            .apply(lambda k, v: v * (pk(sum(ki **2 for ki in k) **0.5) / v.BoxSize.prod()) ** 0.5)\
            .c2r()

    #dynamics
    aa = 1.0/(1+z)
    if pm.comm.rank == 0:
        print('Evolve to redshift = %0.1f, scale factor = %0.2f'%(z, aa))
    stages = numpy.linspace(0.1, aa, nsteps, endpoint=True)

    start = time()
    dynamic_model = NBodyModel(cosmo, pm, B=2, steps=stages)
    #dynamic_model = LPTModel(cosmo, pm, B=2, steps=stages)

    #Save data

    X, V, final = dynamic_model.get_code().compute(['X', 'V', 'final'], init={'parameters':s_truth})
    end = time()
    print('Time taken = ', end-start)

    save_map(s_truth, ofolder + 'mesh', 's')
    save_map(final, ofolder + 'mesh', 'd')

    if pm.comm.rank == 0:
        print('X, V computed')
    cat = ArrayCatalog({'Position': X, 'Velocity' : V}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
    kdd = KDDensity(cat).density
    cat['KDDensity'] = kdd
    cat['InitPosition'] = dynamic_model.get_code().engine.q
    cat.save(ofolder + 'dynamic/1', ('InitPosition', 'Position', 'Velocity', 'KDDensity'))
    if pm.comm.rank == 0:
        print('dynamic model created')

    #FOF
    fof = FOF(cat, linking_length=0.2, nmin=12)
    fofcat = fof.find_features(peakcolumn='KDDensity')
    fofcat['Mass'] = fofcat['Length'] * cosmo.Om0 * 27.7455 * pm.BoxSize.prod() / pm.Nmesh.prod()
    fofcat.save(ofolder+'FOF', ('CMPosition', 'CMVelocity', 'PeakPosition', 'PeakVelocity', 'Length', 'Mass'))
Esempio n. 26
0
def test_2d_2d(comm):
    import pfft
    pm = ParticleMesh(BoxSize=8.0,
                      Nmesh=[8, 8],
                      np=pfft.split_size_2d(comm.size),
                      comm=comm,
                      dtype='f8')
    pm2 = ParticleMesh(BoxSize=8.0,
                       Nmesh=[8, 8, 8],
                       np=pfft.split_size_2d(comm.size),
                       comm=comm,
                       dtype='f8')
    assert pm._use_padded == False
    real = pm.generate_whitenoise(seed=123, type='real')
    complex = pm.generate_whitenoise(seed=123, type='complex')
    assert_array_equal(real, complex.c2r())

    real2 = pm2.generate_whitenoise(seed=123, type='real')

    assert real2.shape[:2] == real.shape
Esempio n. 27
0
def test_cdot_grad(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')

    comp1 = pm.generate_whitenoise(1234, type='complex', mean=1)
    comp2 = pm.generate_whitenoise(1235, type='complex', mean=1)

    def objective(comp1, comp2):
        return comp1.cdot(comp2).real

    y = objective(comp1, comp2)

    grad_comp2 = comp1.cdot_vjp(1.0)
    grad_comp1 = comp2.cdot_vjp(1.0)

    grad_comp1.decompress_vjp(grad_comp1)
    grad_comp2.decompress_vjp(grad_comp2)

    print(grad_comp1)
    print("comp1")
    ng = []
    ag = []
    ind = []
    dx = 1e-7
    for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])):
        dx1, c1 = perturb(comp1, ind1, dx)
        ng1 = (objective(c1, comp2) - objective(comp1, comp2)) / dx
        ag1 = grad_comp1.cgetitem(ind1) * dx1 / dx
        if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))):
            print (ind1, 'a', ag1, 'n', ng1)
        comm.barrier()
        ng.append(ng1)
        ag.append(ag1)
        ind.append(ind1)

    assert_allclose(ng, ag, rtol=1e-5)

    print("comp2")

    ng = []
    ag = []
    ind = []
    dx = 1e-7
    for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])):
        dx1, c2 = perturb(comp2, ind1, dx)
        ng1 = (objective(comp1, c2) - objective(comp1, comp2)) / dx
        ag1 = grad_comp2.cgetitem(ind1) * dx1 / dx
        if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))):
            print (ind1, 'a', ag1, 'n', ng1)
        comm.barrier()
        ng.append(ng1)
        ag.append(ag1)
        ind.append(ind1)

    assert_allclose(ng, ag, rtol=1e-5)
Esempio n. 28
0
def test_cdot(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
    comp1 = pm.generate_whitenoise(1234, type='complex')
    comp2 = pm.generate_whitenoise(1239, type='complex')

    norm1 = comp1.cdot(comp2)
    norm2 = comp2.cdot(comp1)

    norm_r = comp1.c2r().cdot(comp2.c2r()) / pm.Nmesh.prod()

    assert_allclose(norm2.real, norm_r)
    assert_allclose(norm1.real, norm_r)
    assert_allclose(norm1.real, norm2.real)
    assert_allclose(norm1.imag, -norm2.imag)
Esempio n. 29
0
def test_cdot(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
    comp1 = pm.generate_whitenoise(1234, type='complex')
    comp2 = pm.generate_whitenoise(1239, type='complex')

    norm1 = comp1.cdot(comp2)
    norm2 = comp2.cdot(comp1)

    norm_r = comp1.c2r().cdot(comp2.c2r()) / pm.Nmesh.prod()

    assert_allclose(norm2.real, norm_r)
    assert_allclose(norm1.real, norm_r)
    assert_allclose(norm1.real, norm2.real)
    assert_allclose(norm1.imag, -norm2.imag)
Esempio n. 30
0
def test_prior(comm):

    from pmesh.pm import ParticleMesh

    pm = ParticleMesh(BoxSize=128.0, Nmesh=(4, 4), comm=comm)
    vm = fastpm.Evolution(pm, shift=0.5)

    dlink = pm.generate_whitenoise(1234, mode='complex')
    code = vm.code()
    code.Prior(dlin_k='dlin_k', powerspectrum=lambda k : 1.0)
    code.CopyVariable(x='prior', y='chi2')
    dlink, ampl = _addampl(dlink)

    _test_model(code, dlink, ampl)
Esempio n. 31
0
def simulate(comm, ns):

    pm = ParticleMesh(BoxSize=ns.BoxSize,
                      Nmesh=[ns.Nmesh, ns.Nmesh, ns.Nmesh],
                      dtype='f8',
                      comm=comm)
    gaussian = pm.generate_whitenoise(ns.seed, unitary=True)
    time_steps = numpy.linspace(ns.ainit, ns.afinal, ns.steps, endpoint=True)

    Q = pm.generate_uniform_particle_grid(shift=0)
    print(Q.min(axis=0), Q.max(axis=0))

    def convolve(k, v):
        kmag = sum(ki**2 for ki in k)**0.5
        ampl = (PowerSpectrum(kmag) / v.BoxSize.prod())**0.5
        return v * ampl

    dlinear = gaussian.apply(convolve)

    DX1 = numpy.zeros_like(Q)
    layout = pm.decompose(Q)
    # Fill it in one dimension at a time.
    for d in range(pm.ndim):
        DX1[..., d] = dlinear \
                      .apply(dx1_transfer(d)) \
                      .c2r().readout(Q, layout=layout)

    a0 = time_steps[0]

    # 1-LPT Displacement and Veloicty; scaled back from z=0 to the first time step.
    S = DX1 * pt.D1(a=a0)
    V = S * a0**2 * pt.f1(a0) * pt.E(a0)
    state = State(Q, S, V)

    fpm = ParticleMesh(BoxSize=pm.BoxSize,
                       Nmesh=pm.Nmesh * ns.boost,
                       resampler='tsc',
                       dtype='f8')

    ns.scheme(fpm, state, time_steps, ns.factors)

    r = Result()
    r.Q = Q
    r.DX1 = DX1
    r.S = S
    r.V = V
    r.dlinear = dlinear

    return r
Esempio n. 32
0
def test_paint(comm):
    from pmesh.pm import ParticleMesh

    pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm)
    real = pm.generate_whitenoise(type='real', seed=3333)
    complex = real.r2c()

    realmesh = FieldMesh(real)
    complexmesh = FieldMesh(complex)

    assert_array_equal(real, realmesh.to_field(mode='real'))
    assert_array_equal(complex, complexmesh.to_field(mode='complex'))

    assert_allclose(complex, realmesh.to_field(mode='complex'))
    assert_allclose(real, complexmesh.to_field(mode='real'))
Esempio n. 33
0
def test_paint(comm):
    from pmesh.pm import ParticleMesh

    pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm)
    real = pm.generate_whitenoise(mode='real', seed=3333)
    complex = real.r2c()

    realmesh = FieldMesh(real)
    complexmesh = FieldMesh(complex)

    assert_array_equal(real, realmesh.to_field(mode='real'))
    assert_array_equal(complex, complexmesh.to_field(mode='complex'))

    assert_allclose(complex, realmesh.to_field(mode='complex'))
    assert_allclose(real, complexmesh.to_field(mode='real'))
Esempio n. 34
0
def test_lpt(comm):

    from pmesh.pm import ParticleMesh

    pm = ParticleMesh(BoxSize=128.0, Nmesh=(4, 4), comm=comm)
    vm = fastpm.Evolution(pm, shift=0.5)

    dlink = pm.generate_whitenoise(1234, mode='complex')
    code = vm.code()
    code.Decompress(C='dlin_k')
    code.LPTDisplace(dlin_k='dlin_k', D1=1.0, v1=0, D2=0.0, v2=0.0)
    code.Chi2(variable='s')

    dlink, ampl = _addampl(dlink)

    _test_model(code, dlink, ampl)
Esempio n. 35
0
def simulate(comm, ns):

    pm = ParticleMesh(BoxSize=ns.BoxSize, Nmesh=[ns.Nmesh, ns.Nmesh, ns.Nmesh], dtype='f8', comm=comm)
    gaussian = pm.generate_whitenoise(ns.seed, unitary=True)
    time_steps = numpy.linspace(ns.ainit, ns.afinal, ns.steps, endpoint=True)

    Q = pm.generate_uniform_particle_grid(shift=0)
    print(Q.min(axis=0), Q.max(axis=0))
    def convolve(k, v):
        kmag = sum(ki**2 for ki in k) ** 0.5
        ampl = (PowerSpectrum(kmag) / v.BoxSize.prod()) ** 0.5
        return v * ampl
        
    dlinear = gaussian.apply(convolve)


    DX1 = numpy.zeros_like(Q)
    layout = pm.decompose(Q)
    # Fill it in one dimension at a time.
    for d in range(pm.ndim):
        DX1[..., d] = dlinear \
                      .apply(dx1_transfer(d)) \
                      .c2r().readout(Q, layout=layout)


    a0 = time_steps[0]

    # 1-LPT Displacement and Veloicty; scaled back from z=0 to the first time step.
    S = DX1 * pt.D1(a=a0)
    V = S * a0 ** 2 * pt.f1(a0) * pt.E(a0)
    state = State(Q, S, V)

    fpm = ParticleMesh(BoxSize=pm.BoxSize, Nmesh=pm.Nmesh * ns.boost, resampler='tsc', dtype='f8')

    ns.scheme(fpm, state, time_steps, ns.factors)
        
    r = Result()
    r.Q = Q    
    r.DX1 = DX1
    r.S = S
    r.V = V
    r.dlinear = dlinear
    
    return r
Esempio n. 36
0
def test_lpt2():
    """ Checking lpt2_source, this also checks the laplace and gradient kernels
  """
    pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
    grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)

    whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
    lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
        sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)

    # Compute lpt1 from fastpm with matching kernel order
    source = fpmops.lpt2source(lineark).c2r()

    # Same thing from tensorflow
    tfsource = tfpm.lpt2_source(
        pmutils.r2c3d(tf.expand_dims(np.array(lineark.c2r()), axis=0)))
    tfread = pmutils.c2r3d(tfsource).numpy()

    assert_allclose(source, tfread[0], atol=1e-5)
Esempio n. 37
0
def test_resample(comm):
    if comm.size == 4:
        comm.barrier()
    from pmesh.pm import ParticleMesh

    # testing the gradient of down sampling, which we use the most.
    pm = ParticleMesh(BoxSize=128.0, Nmesh=(8, 8), comm=comm)
    vm = fastpm.Evolution(pm, shift=0.5)

    dlink = pm.generate_whitenoise(1234, mode='complex')

    dlink, ampl = _addampl(dlink)
    code = vm.code()
    code.Decompress(C='dlin_k')
    code.C2R(C='dlin_k', R='mesh')
    code.Resample(Neff=4)
    code.Chi2(variable='mesh')

    _test_model(code, dlink, ampl)
Esempio n. 38
0
def test_lpt1():
    """ Checking lpt1, this also checks the laplace and gradient kernels
  """
    pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
    grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)

    whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
    lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
        sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)

    # Compute lpt1 from fastpm with matching kernel order
    lpt = fpmops.lpt1(lineark, grid)

    # Same thing from tensorflow
    with tf.Session() as sess:
        state = tfpm.lpt1(
            pmutils.r2c3d(tf.expand_dims(tf.constant(lineark.c2r()), axis=0)),
            grid.reshape((1, -1, 3)) * nc / bs)
        tfread = sess.run(state)

    assert_allclose(lpt, tfread[0] * bs / nc, atol=1e-5)
Esempio n. 39
0
def test_lpt_init():
    """
  Checking lpt init
  """
    a0 = 0.1

    pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
    grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)
    solver = Solver(pm, Planck15, B=1)

    # Generate initial state with fastpm
    whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
    lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
        sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)
    statelpt = solver.lpt(lineark, grid, a0, order=1)

    # Same thing with flowpm
    tlinear = tf.expand_dims(np.array(lineark.c2r()), 0)
    tfread = tfpm.lpt_init(tlinear, a0, order=1).numpy()

    assert_allclose(statelpt.X, tfread[0, 0] * bs / nc, rtol=1e-2)
Esempio n. 40
0
def test_lpt1_64():
    """ Checking lpt1, this also checks the laplace and gradient kernels
  This variant of the test checks that it works for cubes of size 64
  """
    nc = 64
    pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
    grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)

    whitec = pm.generate_whitenoise(100, mode='complex', unitary=False)
    lineark = whitec.apply(lambda k, v: Planck15.get_pklin(
        sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5)

    # Compute lpt1 from fastpm with matching kernel order
    lpt = fpmops.lpt1(lineark, grid)

    # Same thing from tensorflow
    tfread = tfpm.lpt1(
        pmutils.r2c3d(tf.expand_dims(np.array(lineark.c2r()), axis=0)),
        grid.reshape((1, -1, 3)) * nc / bs).numpy()

    assert_allclose(lpt, tfread[0] * bs / nc, atol=5e-5)
Esempio n. 41
0
def test_lpt_prior(comm):

    from pmesh.pm import ParticleMesh

    pm = ParticleMesh(BoxSize=128.0, Nmesh=(4, 4), comm=comm)
    vm = fastpm.Evolution(pm, shift=0.5)

    dlink = pm.generate_whitenoise(1234, mode='complex')
    code = vm.code()
    code.Decompress(C='dlin_k')
    code.LPTDisplace(dlin_k='dlin_k', D1=1.0, v1=0, D2=0.0, v2=0.0)
    code.Decompose()
    code.Paint()
    code.Chi2(variable='mesh')
    code.Prior(powerspectrum=lambda k : 1.0)
    code.Multiply(a='prior', b='prior', f=0.1)
    code.Add(a='prior', b='chi2', c='chi2')

    dlink, ampl = _addampl(dlink)

    _test_model(code, dlink, ampl)
Esempio n. 42
0
def test_transpose(comm):
    pm = ParticleMesh(BoxSize=[8.0, 16.0, 32.0], Nmesh=[4, 6, 8], comm=comm, dtype='f8')

    comp1 = pm.generate_whitenoise(1234, type='real')

    comp1t = comp1.ctranspose([0, 1, 2])

    assert_array_equal(comp1t.Nmesh, comp1.Nmesh)
    assert_array_equal(comp1t.BoxSize, comp1.BoxSize)

    assert_array_equal(comp1t.cnorm(), comp1.cnorm())

    comp1t = comp1.ctranspose([1, 2, 0])

    assert_array_equal(comp1t.Nmesh, comp1.Nmesh[[1, 2, 0]])
    assert_array_equal(comp1t.BoxSize, comp1.BoxSize[[1, 2, 0]])

    comp1tt = comp1t.ctranspose([1, 2, 0])
    comp1ttt = comp1tt.ctranspose([1, 2, 0])

    assert_allclose(comp1ttt, comp1)
Esempio n. 43
0
def test_transpose(comm):
    pm = ParticleMesh(BoxSize=[8.0, 16.0, 32.0],
                      Nmesh=[4, 6, 8],
                      comm=comm,
                      dtype='f8')

    comp1 = pm.generate_whitenoise(1234, type='real')

    comp1t = comp1.ctranspose([0, 1, 2])

    assert_array_equal(comp1t.Nmesh, comp1.Nmesh)
    assert_array_equal(comp1t.BoxSize, comp1.BoxSize)

    assert_array_equal(comp1t.cnorm(), comp1.cnorm())

    comp1t = comp1.ctranspose([1, 2, 0])

    assert_array_equal(comp1t.Nmesh, comp1.Nmesh[[1, 2, 0]])
    assert_array_equal(comp1t.BoxSize, comp1.BoxSize[[1, 2, 0]])

    comp1tt = comp1t.ctranspose([1, 2, 0])
    comp1ttt = comp1tt.ctranspose([1, 2, 0])

    assert_allclose(comp1ttt, comp1)
Esempio n. 44
0
def test_cdot_cnorm(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
    comp1 = pm.generate_whitenoise(1234, mode='complex')

    norm1 = comp1.cdot(comp1)
    norm2 = comp1.cnorm()
    norm3 = (abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value)))**2).sum()
    assert_allclose(norm2, norm3)

    # p-2 norm of the full complex matrix (norm2) is
    # NORM = Self + Lower + Upper
    # CDOT = Self + Lower
    # Upper = Lower

    self_conj_sum = 0
    for ind1 in numpy.ndindex(*(list(comp1.cshape))):
        c = [(4 - i) % 4 for i, n in zip(ind1, comp1.cshape)]
        ci = numpy.ravel_multi_index(c, comp1.Nmesh)
        ii = numpy.ravel_multi_index(ind1, comp1.Nmesh)
        v = comp1.cgetitem(ind1)
        if ci == ii:
            self_conj_sum += abs(v)**2

    assert_allclose(norm1 * 2, norm2 + self_conj_sum)
Esempio n. 45
0
def test_1d(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[8], comm=comm, dtype='f8')
    real = pm.generate_whitenoise(seed=123, type='real')
    complex = pm.generate_whitenoise(seed=123, type='complex')
    assert_array_equal(real, complex.c2r())
Esempio n. 46
0
def test_readout_gradients(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8', resampler='cic')

    real = pm.generate_whitenoise(1234, type='real')

    def objective(real, pos, layout):
        value = real.readout(pos, layout=layout)
        obj = (value ** 2).sum()
        return comm.allreduce(obj)

    def forward_gradient(real, pos, layout, v_real=None, v_pos=None):
        value = real.readout(pos, layout=layout)
        v_value = real.readout_jvp(pos, v_self=v_real, v_pos=v_pos, layout=layout)
        return comm.allreduce((v_value * value * 2).sum())

    def backward_gradient(real, pos, layout):
        value = real.readout(pos, layout=layout)
        return real.readout_vjp(pos, v=value * 2, layout=layout)

    pos = numpy.array(numpy.indices(real.shape), dtype='f8').reshape(real.value.ndim, -1).T
    pos += real.start
    # avoid sitting at the pmesh points
    # cic gradient is zero on them, the numerical gradient fails.
    pos += 0.5
    pos *= pm.BoxSize / pm.Nmesh

    layout = pm.decompose(pos)

    obj = objective(real, pos, layout)

    grad_real, grad_pos = backward_gradient(real, pos, layout)

    ng = []
    f*g = []
    bag = []
    ind = []
    dx = 1e-6
    for ind1 in numpy.ndindex(*grad_real.cshape):
        dx1, r1 = perturb(real, ind1, dx)
        ng1 = (objective(r1, pos, layout) - obj)

        bag1 = grad_real.cgetitem(ind1) * dx1
        fag1 = forward_gradient(real, pos, layout, v_real=r1 - real)

        # print (dx1, dx, ind1, ag1, ng1)
        ng.append(ng1)
        f*g.append(fag1)
        bag.append(bag1)
        ind.append(ind1)

    assert_allclose(bag, f*g, rtol=1e-7)
    assert_allclose(ng, bag, rtol=1e-4)

    # FIXME
    ng = []
    bag = []
    f*g = []
    ind = []
    dx = 1e-6

    for ind1 in numpy.ndindex((real.csize, real.ndim)):
        dx1, pos1, old = perturb_pos(pos, ind1, dx, comm)
        layout1 = pm.decompose(pos1)

        ng1 = (objective(real, pos1, layout1) - obj)
        bag1 = get_pos(grad_pos, ind1, comm) * dx1

        fag1 = forward_gradient(real, pos, layout, v_pos=pos1 - pos)

        # print ('pos', old, ind1, 'a', ag1, 'n', ng1)
        ng.append(ng1)
        bag.append(bag1)
        f*g.append(fag1)
        ind.append(ind1)

    assert_allclose(bag, f*g, rtol=1e-7)
    assert_allclose(ng, bag, rtol=1e-4)
Esempio n. 47
0
def test_paint_gradients(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8', resampler='cic')

    real = pm.generate_whitenoise(1234, type='real')

    def objective(pos, mass, layout):
        real = pm.paint(pos, mass=mass, layout=layout)
        obj = (real[...] ** 2).sum()
        return comm.allreduce(obj)

    def forward_gradient(pos, mass, layout, v_pos=None, v_mass=None):
        real = pm.paint(pos, mass=mass, layout=layout)
        jvp = pm.paint_jvp(pos, mass=mass, v_mass=v_mass, v_pos=v_pos, layout=layout)
        return comm.allreduce((jvp * real * 2)[...].sum())

    def backward_gradient(pos, mass, layout):
        real = pm.paint(pos, mass=mass, layout=layout)
        return pm.paint_vjp(real * 2, pos, mass=mass, layout=layout)

    pos = numpy.array(numpy.indices(real.shape), dtype='f8').reshape(real.value.ndim, -1).T
    pos += real.start
    numpy.random.seed(9999)
    # avoid sitting at the pmesh points
    # cic gradient is zero on them, the numerical gradient fails.
    pos += (numpy.random.uniform(size=pos.shape)) * 0.8 + 0.1
    pos *= pm.BoxSize / pm.Nmesh
    mass = numpy.ones(len(pos)) * 2

    layout = pm.decompose(pos)

    obj = objective(pos, mass, layout)

    grad_pos, grad_mass = backward_gradient(pos, mass, layout)

    ng = []
    f*g = []
    bag = []
    ind = []
    dx = 1e-6
    for ind1 in numpy.ndindex(real.csize):
        dx1, mass1, old = perturb_mass(mass, ind1[0], dx, comm)
        ng1 = (objective(pos, mass1, layout) - obj)

        bag1 = get_mass(grad_mass, ind1[0], comm) * dx1
        fag1 = forward_gradient(pos, mass, layout, v_mass=mass1 - mass)

    #    print (dx1, dx, ind1, fag1, bag1, ng1)
        ng.append(ng1)
        f*g.append(fag1)
        bag.append(bag1)
        ind.append(ind1)

    assert_allclose(bag, f*g, rtol=1e-7)
    assert_allclose(ng, bag, rtol=1e-4)

    # FIXME
    ng = []
    bag = []
    f*g = []
    ind = []
    dx = 1e-6

    for ind1 in numpy.ndindex((real.csize, real.ndim)):
        dx1, pos1, old = perturb_pos(pos, ind1, dx, comm)
        layout1 = pm.decompose(pos1)

        ng1 = (objective(pos1, mass, layout1) - obj)
        bag1 = get_pos(grad_pos, ind1, comm) * dx1

        fag1 = forward_gradient(pos, mass, layout, v_pos=pos1 - pos)

    #    print ('pos', old, ind1, 'v_pos', (pos1 - pos)[ind1[0]], 'f*g', fag1, 'bag', bag1, 'n', ng1)
        ng.append(ng1)
        bag.append(bag1)
        f*g.append(fag1)
        ind.append(ind1)

    assert_allclose(bag, f*g, rtol=1e-7)
    assert_allclose(ng, bag, rtol=1e-4)
Esempio n. 48
0
def test_1d(comm):
    pm = ParticleMesh(BoxSize=8.0, Nmesh=[8], comm=comm, dtype='f8')
    real = pm.generate_whitenoise(seed=123, type='real')
    complex = pm.generate_whitenoise(seed=123, type='complex')
    assert_array_equal(real, complex.c2r())