def test_operators(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4')
numpy.random.seed(1234)
real = pm.create(type='real', value=0)
complex = pm.create(type='complex', value=0)
real = real + 1
real = 1 + real
real = real + real.value
real = real * real.value
real = real * real
assert isinstance(real, RealField)
complex = 1 + complex
assert isinstance(complex, ComplexField)
complex = complex + 1
assert isinstance(complex, ComplexField)
complex = complex + complex.value
assert isinstance(complex, ComplexField)
complex = numpy.conj(complex) * complex
assert isinstance(complex, ComplexField)
assert (real == real).dtype == numpy.dtype('?')
assert not isinstance(real == real, RealField)
assert not isinstance(complex == complex, ComplexField)
assert not isinstance(numpy.sum(real), RealField)
complex = numpy.conj(complex)
def test_indices_c2c(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='c16')
comp = pm.create(type='complex')
real = pm.create(type='real')
assert_almost_equal(comp.x[0], [[0], [0.785], [-1.571], [-0.785]], decimal=3)
assert_almost_equal(comp.x[1], [[0, 0.785, -1.571, -0.785]], decimal=3)
assert_almost_equal(real.x[0], [[0], [2], [-4], [-2]], decimal=3)
assert_almost_equal(real.x[1], [[0, 2, -4, -2]], decimal=3)
def test_c2c_r2c_edges(comm):
pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[5, 7, 9], comm=comm, dtype='c16')
pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[5, 7, 9], comm=comm, dtype='f8')
real1 = pm1.create(type='real')
real2 = pm2.create(type='real')
assert_allclose(real1.x[0], real2.x[0])
assert_allclose(real1.x[1], real2.x[1])
assert_allclose(real1.x[2], real2.x[2])
def test_create_typenames(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4')
numpy.random.seed(1234)
real = pm.create(type=RealField, value=0)
from pmesh.pm import _typestr_to_type
real = pm.create(type=RealField, value=0)
real = pm.create(type=_typestr_to_type('real'), value=0)
real.cast(type=_typestr_to_type('real'))
real.cast(type=RealField)
real.cast(type=RealField)
def main():
from argparse import ArgumentParser
ap = ArgumentParser()
ap.add_argument("--ndim", type=int, choices=[2, 3], default=2, help="Number of dimensions, 2 or 3")
ap.add_argument("--nmesh", type=int, default=256, help="Size of FFT mesh")
ns = ap.parse_args()
pm = ParticleMesh(BoxSize=32.0, Nmesh=[ns.nmesh] * ns.ndim, comm=MPI.COMM_WORLD)
u = pm.create(mode='real')
def transfer(i, v):
r = [(ii - 0.5 * ni) * (Li / ni) for ii, ni, Li in zip(i, v.Nmesh, v.BoxSize)]
r2 = sum(ri ** 2 for ri in r)
return 4.0 * numpy.arctan(numpy.exp(3 - r2))
u = u.apply(transfer, kind='index')
du = pm.create(mode='real')
du[...] = 0
steps = numpy.linspace(0, 16, 321, endpoint=True)
tmonitor = [0, 4, 8, 11.5, 15]
def monitor(t, dt, u_k, dv_k):
norm = u_k.cnorm()
if pm.comm.rank == 0:
print("---- timestep %5.3f, step size %5.4f" % (t, dt))
print("norm of u_k is %g." % norm)
for tm in tmonitor.copy():
if abs(t - tm) > dt * 0.5: continue
preview = u_k.c2r().preview(Nmesh=min([512, ns.nmesh]), axes=(0, 1))
if pm.comm.rank == 0:
print("writing a snapshot")
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg
fig = Figure(figsize=(8, 8))
ax = fig.add_subplot(111)
ax.imshow(preview.T, origin='lower', extent=(0, pm.BoxSize[0], 0, pm.BoxSize[1]))
canvas = FigureCanvasAgg(fig)
fig.savefig('klein-gordon-result-%05.3f.png' % t, dpi=128)
tmonitor.remove(tm)
kgsolver(steps, u, du, lambda u : numpy.sin(u), monitor=monitor)
def test_negnyquist(comm):
# the nyquist mode wave number in the hermitian complex field must be negative.
# nbodykit depends on this behavior.
# see https://github.com/bccp/nbodykit/pull/459
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
c = pm.create(type='complex')
assert (c.x[-1][0][-1] < 0).all()
assert (c.x[-1][0][:-1] >= 0).all()
def test_lic(comm):
from pmesh.pm import ParticleMesh
pm = ParticleMesh(Nmesh=[8, 8], comm=comm, BoxSize=8.0)
vx = pm.create(type='real')
vy = pm.create(type='real')
vx = vx.apply(lambda r, v: r[0])
vy = vy.apply(lambda r, v: 1.0)
# FIXME: This only tests if the code 'runs', but no correctness
# is tested.
r = lic([vx, vy], kernel=lambda s: 1.0, length=2.0, ds=0.1)
def test_real_resample(comm):
from functools import reduce
pmh = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
pml = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')
reall = pml.create(type='real')
reall.apply(lambda i, v: (i[0] % 2) * (i[1] %2 ), kind='index', out=Ellipsis)
for resampler in ['nearest', 'cic', 'tsc', 'cubic']:
realh = pmh.upsample(reall, resampler=resampler, keep_mean=False)
reall2 = pml.downsample(realh, resampler=resampler)
# print(resampler, comm.rank, comm.size, reall, realh)
assert_almost_equal(reall.csum(), realh.csum())
assert_almost_equal(reall.csum(), reall2.csum())
def test_fft(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4')
numpy.random.seed(1234)
real = pm.create(type='real', value=0)
raw = real._base.view_raw()
real[...] = 2
real[::2, ::2] = -2
real3 = real.copy()
complex = real.r2c()
assert_almost_equal(numpy.asarray(real), numpy.asarray(real3), decimal=7)
real2 = complex.c2r()
assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_real_resample(comm):
from functools import reduce
pmh = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
pml = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8')
reall = pml.create(type='real')
reall.apply(lambda i, v: (i[0] % 2) * (i[1] % 2),
kind='index',
out=Ellipsis)
for resampler in ['nearest', 'cic', 'tsc', 'cubic']:
realh = pmh.upsample(reall, resampler=resampler, keep_mean=False)
reall2 = pml.downsample(realh, resampler=resampler)
# print(resampler, comm.rank, comm.size, reall, realh)
assert_almost_equal(reall.csum(), realh.csum())
assert_almost_equal(reall.csum(), reall2.csum())
def test_cic_readout():
bs = 50
nc = 16
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
nparticle = 100
pos = bs * np.random.random(3 * nparticle).reshape(-1, 3).astype(
np.float32)
base = 100 * np.random.random(nc**3).reshape(nc, nc, nc).astype(np.float32)
pmmesh = pm.create(mode='real', value=base)
pmread = pmmesh.readout(pos)
mesh = cic_readout(
tf.constant(base.reshape((1, nc, nc, nc)), dtype=tf.float32),
(pos * nc / bs).reshape((1, nparticle, 3)))
tfread = mesh.numpy()
assert_allclose(pmread, tfread[0], rtol=1e-06)
def createmesh(self, bs, nc, positions, weights):
'''use this to create mesh of HI
'''
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc])
mesh = pm.create(mode='real', value=0)
comm = pm.comm
rankweight = sum([wt.sum() for wt in weights])
totweight = comm.allreduce(rankweight)
for wt in weights:
wt /= totweight / float(nc)**3
for i in range(len(positions)):
lay = pm.decompose(positions[i])
mesh.paint(positions[i], mass=weights[i], layout=lay, hold=True)
return mesh
def test_cic_readout():
bs = 50
nc = 16
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
nparticle = 100
pos = bs * np.random.random(3 * nparticle).reshape(-1, 3).astype(
np.float32)
base = 100 * np.random.random(nc**3).reshape(nc, nc, nc).astype(np.float32)
pmmesh = pm.create(mode='real', value=base)
pmread = pmmesh.readout(pos)
with tf.Session() as sess:
mesh = cic_readout(
tf.constant(base.reshape((1, nc, nc, nc)), dtype=tf.float32),
(pos * nc / bs).reshape((1, nparticle, 3)))
sess.run(tf.global_variables_initializer())
tfread = sess.run(mesh)
assert_allclose(pmread, tfread[0], rtol=1e-06)
color = int(world_rank % 2)
time_comm = comm.Split(color=color)
time_size = time_comm.Get_size()
time_rank = time_comm.Get_rank()
print(world_rank, time_rank, space_rank)
t0 = time.time()
if time_rank == 0:
pm = ParticleMesh(BoxSize=1.0,
Nmesh=[nvars] * 2,
dtype='f8',
plan_method='measure',
comm=space_comm)
u = pm.create(type='real')
u = u.apply(doublesine, kind='index', out=Ellipsis)
time_comm.send(u.value, dest=1, tag=11)
else:
pm = ParticleMesh(BoxSize=1.0,
Nmesh=[nvars] * 2,
dtype='f8',
plan_method='measure',
comm=space_comm)
tmp = time_comm.recv(source=0, tag=11)
u = pm.create(type='real', value=tmp)
t1 = time.time()
print(f'PMESH setup time: {t1 - t0:6.4f} sec.')
exit()
def circle(i, v):
r = [ii * (Li / ni) - 0.5 * Li for ii, ni, Li in zip(i, v.Nmesh, v.BoxSize)]
r2 = sum(ri ** 2 for ri in r)
return np.tanh((0.25 - np.sqrt(r2)) / (np.sqrt(2) * 0.04))
nvars = 128
nruns = 1000
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
t0 = time.time()
pm = ParticleMesh(BoxSize=1.0, Nmesh=[nvars] * 2, dtype='f8', plan_method='measure', comm=comm)
tmp = pm.create(type='real')
t1 = time.time()
# a = pmesh_datatype((pm, (2, (4, 4))))
print(f'PMESH setup time: {t1 - t0:6.4f} sec.')
dt = 0.121233
res = 0.0
t0 = time.time()
# for n in range(nruns):
#
# # set initial condition
# # u = pmesh_datatype(init=pm)
# # u.values.apply(circle, kind='index', out=Ellipsis)
class Data_Generation():
log=classlogger
def __init__(self, parameters_box, parameters_igm, use_baryonic=False, transfer="EisensteinHu", path_transfer=None, column=None, cosmo=None):
self.redshift= parameters_box['redshift']
self.width=parameters_box['width']
chosen_cosmo=parameters_box['cosmology']
#Use the baryonic power spectrum instead
#Then the resulting overdensity field does not correspond to cdm-density, but to baryonic density
#Jeans length needed [h**-1 Mpc]
#Note that Jeans length scales with redshift
#For proper values look: Choudhury, Sraianda ... 2001
self.use_baryonic=use_baryonic
self.Jeans_length=parameters_box['jeans_length']
#specifies specific cosmological model
if chosen_cosmo=='Planck15':
self.cosmo = cosmology.Planck15
elif chosen_cosmo=='UserDefined':
self.cosmo = cosmo
else:
print(chosen_cosmo, 'not implemented right now')
print('Continue with Planck15 cosmology instead')
self.cosmo = cosmology.Planck15
self.comoving_distance=self.cosmo.comoving_distance(self.redshift)
#linear power spectrum
Plin = LinearPower(self.cosmo, self.redshift, transfer, path_transfer, column)
if self.use_baryonic:
self.Plin=lambda k: 1/(1+self.Jeans_length**2*k**2)**2*Plin(k)
self.Plin.redshift=self.redshift
self.Plin.sigma8=self.cosmo.sigma8
self.Plin.cosmo=self.cosmo
if self.log.isEnabledFor(logging.INFO):
self.log.info('Start computation of baryonic density field')
else:
self.Plin=Plin
if self.log.isEnabledFor(logging.INFO):
self.log.info('Start computation of cdm density field')
#Parameters describing the box
self.background_field_x=parameters_box['background_box'][0]
self.background_field_y=parameters_box['background_box'][1]
self.background_sampling_x=parameters_box['background_sampling'][0]
self.background_sampling_y=parameters_box['background_sampling'][1]
if self.log.isEnabledFor(logging.INFO):
self.log.info('Background has been specified:')
self.log.info('-->Comoving length: {}'.format([self.background_field_x, self.background_field_y]))
self.log.info('-->sampled:{}'.format([self.background_sampling_x, self.background_sampling_y]))
#initializes pseudo-gaussian random generator
self.seed=parameters_box['seed']
#if True, the seed gaussian has unitary_amplitude
self.unitary_amplitude=True
self.inverted_phase=False
#Also compute displacement
self.compute_displacement=True
self.BoxSize=[self.background_field_x, self.background_field_y, self.width]
self.N_pix=parameters_box['nr_pix']
self.bias=2
#The following values describe the state of the IGM and have to match whatever specified in the inversion
self.beta=parameters_igm['beta']
self.alpha=2-1.4*self.beta
self.tildeC=parameters_igm['tilde_c']
if self.log.isEnabledFor(logging.INFO):
self.log.info('Generator initialized with parameters')
self.log.info('-->Cosmology: {}'.format(chosen_cosmo))
self.log.info('-->Box size: {}'.format(self.BoxSize))
self.log.info('-->Number pixels: {}'.format(self.N_pix))
self.log.info('-->Redshift: {}'.format(self.redshift))
self.log.info('-->seed: {}'.format(self.seed))
self.log.info('-->Jeans: {}'.format(self.Jeans_length))
# the particle mesh for gridding purposes
_Nmesh = np.empty(3, dtype='i8')
_Nmesh[0] = self.background_sampling_x
_Nmesh[1] = self.background_sampling_y
_Nmesh[2] = self.N_pix
self.pm = ParticleMesh(BoxSize=self.BoxSize, Nmesh=_Nmesh, dtype='f4')
return
#Compute Density Field. The function strongly matches the implementation in nbodykit.
#Return:
#delta: Density Perturbation (rho/rho_med)
#comoving: Comoving distance to each point at line of sight [h^-1 Mpc]
#vel: Peculiar velocity field [km/s]
def Compute_Density_Field(self):
# growth rate to do RSD in the Zel'dovich approximation
f = self.cosmo.scale_independent_growth_rate(self.redshift)
if self.log.isEnabledFor(logging.INFO):
self.log.info('Growth rate is {}'.format(f))
#Lineat density and displacement field
delta, disp = gaussian_real_fields(self.pm, self.Plin, self.seed,
unitary_amplitude=self.unitary_amplitude,
inverted_phase=self.inverted_phase,
compute_displacement=self.compute_displacement)
if self.log.isEnabledFor(logging.INFO):
self.log.info('Gaussian field generated')
# apply the lognormal transformation to the initial conditions density
# this creates a positive-definite delta (necessary for Poisson sampling)
lagrangian_bias = self.bias - 1
delta = lognormal_transform(delta, bias=lagrangian_bias)
if self.log.isEnabledFor(logging.INFO):
self.log.info('Density field projected to lognormal')
#Compute peculiar velocity field from linear displacement field
if self.compute_displacement:
self.displacement=disp
velocity_norm = f * 100 * self.cosmo.efunc(self.redshift) / (1+self.redshift)
vel_x = velocity_norm * disp[0]
vel_y = velocity_norm * disp[1]
vel_z = velocity_norm * disp[2]
vel=[vel_x.value, vel_y.value, vel_z.value]
if self.log.isEnabledFor(logging.INFO):
self.log.info('Velocity field computed')
#Computes comoving distances along line of sight
comoving=self.cosmo.comoving_distance(self.redshift)-self.width+self.width/self.N_pix*np.linspace(1, self.N_pix, self.N_pix)
if self.log.isEnabledFor(logging.INFO):
self.log.info('Density field succesfully computed')
return [delta, vel, comoving]
#Computes the Covariance matrix, depending on the linear autocorrelation
#Correlation may look different in nonlinear and mildly nonlinear regime (due to lognormal transform)
#In case of non biased gaussian fields (expectation value equal to zero) the output matches the covariance matrix
#diff: The difference value between two points in real space
#Autocorrelation is given by the inverse Fourier transform of power spectrum (Wiener-Khinchin-Theorem)
#comoving: vector of comoving coordinates on which the density field is evaluated
#WARNING: Due to rebinning the comoving field could have changed and is not equal binned anymore
#As Covariance matrix is symmetric, only the upper half is computed, then both merged together by transposition
def Compute_Linear_Covariance(self, comoving):
if self.log.isEnabledFor(logging.INFO):
self.log.info('Computation of prior covariance started')
Covariance=np.zeros((np.size(comoving), np.size(comoving)))
k = np.linspace(10**(-5), 10, 10**6)
size = len(k)//2
Pk = self.Plin(k)
fourier_coeff = np.abs(np.fft.fftn(Pk)[0:size+1])
frqs = np.linspace(0, 0.1*size, size+1)
cf_lin = Spline(frqs, fourier_coeff)
diff=np.zeros((np.size(comoving), np.size(comoving)))
for i in range(np.size(comoving)):
for j in range(np.size(comoving)):
diff[i, j]=np.abs(comoving[i]-comoving[j])
Covariance=cf_lin(diff)
Covariance /= Covariance[0, 0]
if self.log.isEnabledFor(logging.INFO):
self.log.info('Prior Covariance computed')
return Covariance
#This Function computes the selection of a single line of sight from the delta sample
#i, j: indices of background sources at maximal redshift
#Return: Density-field along one line-of-sight, evaluated at the pixels of comoving
def Select_LOS(self, delta, vel, index):
i=index[0].astype(int)
j=index[1].astype(int)
if self.log.isEnabledFor(logging.INFO):
self.log.info('Line of sight selected: {}'.format(index))
return [delta[i, j, :], vel[i, j, :]]
#Poisson sample to overdensity field,
#matches the nbodykit routine
def PoissonSample(self, delta, parameters_sampling):
nbar=parameters_sampling['nbar']
seed1=parameters_sampling['seed1']
seed2=parameters_sampling['seed2']
comm = self.pm.comm
# mean number of objects per cell
H = self.BoxSize / self.pm.Nmesh
overallmean = H.prod() * nbar
# number of objects in each cell (per rank, as a RealField)
cellmean = delta * overallmean
# create a random state with the input seed
rng = MPIRandomState(seed=seed1, comm=comm, size=delta.size)
# generate poissons. Note that we use ravel/unravel to
# maintain MPI invariane.
Nravel = rng.poisson(lam=cellmean.ravel())
N = self.pm.create(type='real')
N.unravel(Nravel)
Ntot = N.csum()
if self.log.isEnabledFor(logging.INFO):
self.log.info('Poisson sampling done, total number of objects is {}'.format(Ntot))
pos_mesh = self.pm.generate_uniform_particle_grid(shift=0.0)
disp_mesh = np.empty_like(pos_mesh)
# no need to do decompose because pos_mesh is strictly within the
# local volume of the RealField.
N_per_cell = N.readout(pos_mesh, resampler='nnb')
for i in range(N.ndim):
disp_mesh[:, i] = self.displacement[i].readout(pos_mesh, resampler='nnb')
# fight round off errors, if any
N_per_cell = np.int64(N_per_cell + 0.5)
pos = pos_mesh.repeat(N_per_cell, axis=0)
disp = disp_mesh.repeat(N_per_cell, axis=0)
del pos_mesh
del disp_mesh
if self.log.isEnabledFor(logging.INFO):
self.log.info("Catalog produced. Assigning in cell shift.")
# FIXME: after pmesh update, remove this
orderby = np.int64(pos[:, 0] / H[0] + 0.5)
for i in range(1, delta.ndim):
orderby[...] *= self.pm.Nmesh[i]
orderby[...] += np.int64(pos[:, i] / H[i] + 0.5)
# sort by ID to maintain MPI invariance.
pos = mpsort.sort(pos, orderby=orderby, comm=comm)
disp = mpsort.sort(disp, orderby=orderby, comm=comm)
if self.log.isEnabledFor(logging.INFO):
self.log.info("Sorting done")
rng_shift = MPIRandomState(seed=seed2, comm=comm, size=len(pos))
in_cell_shift = rng_shift.uniform(0, H[i], itemshape=(delta.ndim,))
pos[...] += in_cell_shift
pos[...] %= self.pm.BoxSize
if self.log.isEnabledFor(logging.INFO):
self.log.info("Catalog shifted.")
#Catalog needs to be shifted in z-coordinate, such that pos and comoving match
pos[...,0]+=(self.comoving_distance-self.width)*np.ones(pos.shape[0])
return pos, disp
#Projects baryonic density field to neutral hydrogen overdensity field
def Find_Neutral_Hydrogen_Fraction(self, density, redshift):
density_h1=self.tildeC*(1+redshift)**6*density**(self.alpha)
if self.log.isEnabledFor(logging.INFO):
self.log.info("Density field projected to neutral hydrogen density field")
return density_h1
def Project_Velocity(self, vel, direction):
return vel[direction]
print('Reading Files')
meshdict, dummy = ntools.readfiles(
pm,
proj + '/data/z%02d/L%04d_N%04d_S%04d_05step/' % (zz * 10, bs, nc, seed),
R1=R1,
R2=R2,
abund=abund,
doexp=doexp,
mexp=mexp,
cc=cc,
stellar=stellar)
ftt = ntools.createdata(pm, meshdict, pdict['pftname'], plocal)
mftt = ntools.createdata(pm, meshdict, mftname, mlocal)
nnpred = ntools.applynet(ftt, ptup).reshape(nc, nc, nc)
nnmass = ntools.applynet(mftt, mtup).reshape(nc, nc, nc)
predict = pm.create(mode='real')
predict[...] = nnpred * nnmass
predictr = ntools.relu(predict[...])
#predictR = ft.smooth(predict, Rsm, 'fingauss')
#data
print('Generating data')
dpath = proj + '/data/z%02d/L%04d_N%04d_S%04d_40step/' % (zz * 10, bs,
sfine * nc, seed)
galcat = BigFileCatalog(dpath + 'galaxies_n%02d/galcat' % (numd * 1e4),
header='Header',
comm=pm.comm)
datap = pm.paint(galcat['Position'], mass=galcat['Mass'])
datapR = ft.smooth(datap, Rsm, 'fingauss')
#model
print('Reading Files')
meshdict, dummy = ntools.readfiles(pm,
scratch +
'/data/z%02d/L%04d_N%04d_S%04d_05step/' %
(zz * 10, bs, nc, seed),
R1=R1,
R2=R2,
abund=abund)
ftt = ntools.createdata(pm, meshdict, pdict['pftname'], plocal)
mftt = ntools.createdata(pm, meshdict, mftname, mlocal)
nnpred = ntools.applynet(ftt, ptup).reshape(nc, nc, nc)
nnmass = ntools.applynet(mftt, mtup).reshape(nc, nc, nc)
predict = pm.create(mode='real', value=nnpred * nnmass)
predictR = ft.smooth(predict, Rsm, 'fingauss')
#data
print('Generating data')
hdictf = ntools.gridhalos(pm,
dpath=scratch +
'/data/z%02d/L%04d_N%04d_S%04d_40step/' %
(zz * 10, bs, sfine * nc, seed),
R1=R1,
R2=R2,
pmesh=False,
abund=abund)[1]
datap = pm.paint(hdictf['position'][:num], mass=hdictf['mass'][:num])
R2 = R1 * sfac
#model
print('Reading Files')
meshdict, dummy = ntools.readfiles(pm,
scratch +
'/data/z%02d/L%04d_N%04d_S%04d_05step/' %
(zz * 10, bs, nc, seed),
R1=R1,
R2=R2,
abund=abund)
ftt = ntools.createdata(pm, meshdict, pdict['pftname'], plocal)
mftt = ntools.createdata(pm, meshdict, mftname, mlocal)
nnpred = ntools.applynet(ftt, ptup).reshape(nc, nc, nc)
predict = pm.create(mode='real', value=nnpred)
predictR = ft.smooth(predict, Rsm, 'fingauss')
if ovd: predictR[...] = (predictR[...] - predictR.cmean()) / (predictR.cmean())
#data
print('Generating data')
hdictf = ntools.gridhalos(pm,
dpath=scratch +
'/data/z%02d/L%04d_N%04d_S%04d_40step/' %
(zz * 10, bs, sfine * nc, seed),
R1=R1,
R2=R2,
pmesh=False,
abund=abund)[1]
class allencahn_imex(ptype):
"""
Example implementing Allen-Cahn equation in 2-3D using PMESH for solving linear parts, IMEX time-stepping
PMESH: https://github.com/rainwoodman/pmesh
Attributes:
xvalues: grid points in space
dx: mesh width
"""
def __init__(self, problem_params, dtype_u=pmesh_datatype, dtype_f=rhs_imex_pmesh):
"""
Initialization routine
Args:
problem_params (dict): custom parameters for the example
dtype_u: pmesh data type (will be passed to parent class)
dtype_f: pmesh data type wuth implicit and explicit parts (will be passed to parent class)
"""
if 'L' not in problem_params:
problem_params['L'] = 1.0
if 'init_type' not in problem_params:
problem_params['init_type'] = 'circle'
if 'comm' not in problem_params:
problem_params['comm'] = None
if 'dw' not in problem_params:
problem_params['dw'] = 0.0
# these parameters will be used later, so assert their existence
essential_keys = ['nvars', 'eps', 'L', 'radius', 'dw']
for key in essential_keys:
if key not in problem_params:
msg = 'need %s to instantiate problem, only got %s' % (key, str(problem_params.keys()))
raise ParameterError(msg)
if not (isinstance(problem_params['nvars'], tuple) and len(problem_params['nvars']) > 1):
raise ProblemError('Need at least two dimensions')
# Creating ParticleMesh structure
self.pm = ParticleMesh(BoxSize=problem_params['L'], Nmesh=list(problem_params['nvars']), dtype='f8',
plan_method='measure', comm=problem_params['comm'])
# create test RealField to get the local dimensions (there's probably a better way to do that)
tmp = self.pm.create(type='real')
# invoke super init, passing the communicator and the local dimensions as init
super(allencahn_imex, self).__init__(init=(self.pm.comm, tmp.value.shape), dtype_u=dtype_u, dtype_f=dtype_f,
params=problem_params)
# Need this for diagnostics
self.dx = self.params.L / problem_params['nvars'][0]
self.dy = self.params.L / problem_params['nvars'][1]
self.xvalues = [i * self.dx - problem_params['L'] / 2 for i in range(problem_params['nvars'][0])]
self.yvalues = [i * self.dy - problem_params['L'] / 2 for i in range(problem_params['nvars'][1])]
def eval_f(self, u, t):
"""
Routine to evaluate the RHS
Args:
u (dtype_u): current values
t (float): current time
Returns:
dtype_f: the RHS
"""
def Laplacian(k, v):
k2 = sum(ki ** 2 for ki in k)
return -k2 * v
f = self.dtype_f(self.init)
tmp_u = self.pm.create(type='real', value=u.values)
f.impl.values = tmp_u.r2c().apply(Laplacian, out=Ellipsis).c2r(out=Ellipsis).value
if self.params.eps > 0:
f.expl.values = - 2.0 / self.params.eps ** 2 * u.values * (1.0 - u.values) * (1.0 - 2.0 * u.values) - \
6.0 * self.params.dw * u.values * (1.0 - u.values)
return f
def solve_system(self, rhs, factor, u0, t):
"""
Simple FFT solver for the diffusion part
Args:
rhs (dtype_f): right-hand side for the linear system
factor (float) : abbrev. for the node-to-node stepsize (or any other factor required)
u0 (dtype_u): initial guess for the iterative solver (not used here so far)
t (float): current time (e.g. for time-dependent BCs)
Returns:
dtype_u: solution as mesh
"""
def linear_solve(k, v):
k2 = sum(ki ** 2 for ki in k)
return 1.0 / (1.0 + factor * k2) * v
me = self.dtype_u(self.init)
tmp_rhs = self.pm.create(type='real', value=rhs.values)
me.values = tmp_rhs.r2c().apply(linear_solve, out=Ellipsis).c2r(out=Ellipsis).value
return me
def u_exact(self, t):
"""
Routine to compute the exact solution at time t
Args:
t (float): current time
Returns:
dtype_u: exact solution
"""
def circle(i, v):
r = [ii * (Li / ni) - 0.5 * Li for ii, ni, Li in zip(i, v.Nmesh, v.BoxSize)]
r2 = sum(ri ** 2 for ri in r)
return 0.5 * (1.0 + np.tanh((self.params.radius - np.sqrt(r2)) / (np.sqrt(2) * self.params.eps)))
def circle_rand(i, v):
L = [int(l) for l in v.BoxSize]
r = [ii * (Li / ni) - 0.5 * Li for ii, ni, Li in zip(i, v.Nmesh, L)]
rshift = r.copy()
ndim = len(r)
data = 0
# get random radii for circles/spheres
np.random.seed(1)
lbound = 3.0 * self.params.eps
ubound = 0.5 - self.params.eps
rand_radii = (ubound - lbound) * np.random.random_sample(size=tuple(L)) + lbound
# distribnute circles/spheres
if ndim == 2:
for indexi, i in enumerate(range(-L[0] + 1, L[0], 2)):
for indexj, j in enumerate(range(-L[1] + 1, L[1], 2)):
# shift x and y coordinate depending on which box we are in
rshift[0] = r[0] + i/2
rshift[1] = r[1] + j/2
# build radius
r2 = sum(ri ** 2 for ri in rshift)
# add this blob, shifted by 1 to avoid issues with adding up negative contributions
data += np.tanh((rand_radii[indexi, indexj] - np.sqrt(r2)) / (np.sqrt(2) * self.params.eps)) + 1
# get rid of the 1
data *= 0.5
assert np.all(data <= 1.0)
return data
assert t == 0, 'ERROR: u_exact only valid for t=0'
me = self.dtype_u(self.init)
if self.params.init_type == 'circle':
tmp_u = self.pm.create(type='real', value=0.0)
me.values = tmp_u.apply(circle, kind='index').value
elif self.params.init_type == 'circle_rand':
tmp_u = self.pm.create(type='real', value=0.0)
me.values = tmp_u.apply(circle_rand, kind='index').value
else:
raise NotImplementedError('type of initial value not implemented, got %s' % self.params.init_type)
return me
R2=R2,
abund=abund,
doexp=doexp,
mexp=mexp,
cc=cc,
stellar=stellar)
ftt = ntools.createdata(pm, meshdict, pdict['pftname'], plocal)
mftt = ntools.createdata(pm, meshdict, mftname, mlocal)
nnpred = ntools.applynet(ftt, ptup).reshape(nc, nc, nc)
if regression:
nnmass = ntools.applymassreg(mftt, mtup).reshape(nc, nc, nc)
else:
nnmass = ntools.applynet(mftt, mtup).reshape(nc, nc, nc)
if doexp:
nnmass = mf.fmexp(ntools.relu(nnmass), mexp, cc)
predict = pm.create(mode='real')
predict[...] = nnpred * nnmass
predictR = ft.smooth(predict, Rsm, 'fingauss')
#data
hdictf = ntools.gridhalos(pm,
scratch + '/data/L%04d_N%04d_S%04d_40step/' %
(bs, fine * nc, seed),
R1=R1,
R2=R2,
pmesh=True,
abund=abund,
doexp=doexp,
mexp=mexp,
cc=cc,
LDLmap *= 1e-5 #The learned LDL map is actually 1e5*map_ne.
save = args.save + '/ne_snap' + str(args.snapNum).zfill(3) + '_Nmesh%d_Nstep%d_n%.2f_map' % (args.Nmesh, args.Nstep, args.n)
FieldMesh(LDLmap).save(save)
if args.FastPMpath:
X = File(args.FastPMpath)['Position']
X = np.array(X[start:end]).astype(np.float32)
Vz = File(args.FastPMpath)['Velocity']
Vz = np.array(Vz[start:end])[:,2].astype(np.float32)
else:
from readTNG import scalefactor
a = scalefactor(args.TNGDarkpath, args.snapNum)
Vz = load_TNG_data(TNG_basepath=args.TNGDarkpath, snapNum=args.snapNum, partType='dm', field='Velocities', mdi=2) * a**0.5
layout = pm.decompose(X)
X = layout.exchange(X)
Vz = layout.exchange(Vz)
map_vz = pm.create(type="real")
map_vz.paint(X, mass=Vz, layout=None, hold=False)
map_delta = pm.create(type="real")
map_delta.paint(X, mass=1., layout=None, hold=False)
select = map_delta > 0
map_vz[select] = map_vz[select] / map_delta[select]
map_vz[~select] = 0
LDLmap = LDLmap * map_vz
save = args.save + '/' + args.target + '_snap' + str(args.snapNum).zfill(3) + '_Nmesh%d_Nstep%d_n%.2f_map' % (args.Nmesh, args.Nstep, args.n)
FieldMesh(LDLmap).save(save)
elif args.target in ['tSZ_T', 'Xray_T']:
param = np.loadtxt(args.restore_ne)
assert len(param) % 5 == 3
Nstep = len(param) // 5