def get_code(self):
import numpy
pm = self.mock_model.pm
code = self.mock_model.get_code()
data = self.data.mapp
code.add(x1='model', x2=Literal(data * -1), y='residual')
if self.noise_model is not None:
code.multiply(x1='residual',
x2=Literal(self.noise_model.ivar2d**0.5),
y='residual')
code.r2c(real='residual', complex='C')
smooth_window = lambda k: numpy.exp(-self.sml**2 * sum(ki**2
for ki in k))
code.transfer(complex='C', tf=smooth_window)
code.c2r(real='residual', complex='C')
code.to_scalar(x='residual', y='chi2')
code.create_whitenoise(dlinear_k='dlinear_k',
powerspectrum=self.prior_ps,
whitenoise='pvar')
code.to_scalar(x='pvar', y='prior')
# the whitenoise is not properly normalized as d_k / P**0.5
code.multiply(x1='prior', x2=Literal(pm.Nmesh.prod()**-1.), y='prior')
code.add(x1='prior', x2='chi2', y='objective')
return code
def get_code(self):
pm = self.mock_model.pm
code = base.Objective.get_code(self)
data = self.data.mapp
M0 = self.M0
#likelihood is in log(M + M0)
logdataM0 = self.pm.create(mode='real')
logdataM0.value[...] = numpy.log(data + M0)
code.add(x1='model', x2=Literal(M0), y='modelM0')
code.log(x='modelM0', y='logmodelM0')
code.add(x1='logmodelM0', x2=Literal(logdataM0 * -1.), y='residual')
code.multiply(x1='residual',
x2=Literal(self.noise_model.ivar2d**0.5),
y='residual')
code.to_scalar(x='residual', y='chi2')
code.create_whitenoise(dlinear_k='dlinear_k',
powerspectrum=self.prior_ps,
whitenoise='pvar')
code.to_scalar(x='pvar', y='prior')
# the whitenoise is not properly normalized as d_k / P**0.5
code.multiply(x1='prior', x2=Literal(pm.Nmesh.prod()**-1.), y='prior')
code.add(x1='prior', x2='chi2', y='objective')
return code
def get_code(self):
import numpy
pm = self.mock_model.pm
code = self.mock_model.get_code()
data = self.data.mapp
code.add(x1='model', x2=Literal(data * -1), y='residual')
if self.noise_model is not None:
code.multiply(x1='residual',
x2=Literal(self.noise_model.ivar2d**0.5),
y='residual')
code.r2c(real='residual', complex='C')
smooth_window = lambda k: numpy.exp(-self.sml**2 * sum(ki**2
for ki in k))
code.transfer(complex='C', tf=smooth_window)
def tf(k):
k2 = sum(ki**2 for ki in k)
r = (powerspectrum(k2**0.5) / engine.pm.BoxSize.prod())**-0.5
r[k2 == 0] = 1.0
return r
code.assign(x='dlinear_k', y='tmp')
code.create_whitenoise(dlinear_k='C',
powerspectrum=self.error_ps,
whitenoise='perror')
def tf(k):
kmesh = sum(ki**2 for ki in k)**0.5
mask = [numpy.ones_like(ki) for ki in k]
mask[2] *= abs(k[2]) >= self.kmin
mask = numpy.prod(mask)
mask2 = numpy.ones_like(mask)
if self.angle > 0:
kperp = (k[0]**2 + k[1]**2)**0.5
kangle = abs(k[2]) / (kperp + 1e-10)
angles = (numpy.arctan(kangle) * 180 / numpy.pi)
mask2[angles < self.angle] = 0
fgmask = mask * mask2
return fgmask
code.r2c(real='perror', complex='perrorc')
code.transfer(complex='perrorc', tf=tf)
code.c2r(complex='perrorc', real='perror')
code.to_scalar(x='perror', y='chi2')
#code.c2r(real='residual', complex='C')
#code.to_scalar(x='residual', y='chi2')
code.create_whitenoise(dlinear_k='dlinear_k',
powerspectrum=self.prior_ps,
whitenoise='pvar')
code.to_scalar(x='pvar', y='prior')
# the whitenoise is not properly normalized as d_k / P**0.5
code.multiply(x1='prior', x2=Literal(pm.Nmesh.prod()**-1.), y='prior')
code.add(x1='prior', x2='chi2', y='objective')
return code
def get_code(self):
import numpy
pm = self.mock_model.pm
code = self.mock_model.get_code()
data = self.data.mapp
code.add(x1='model', x2=Literal(data * -1), y='residual')
if self.noise_model is not None:
code.multiply(x1='residual',
x2=Literal(self.noise_model.ivar2d**0.5),
y='residual')
code.r2c(real='residual', complex='C')
smooth_window = lambda k: numpy.exp(-self.sml**2 * sum(ki**2
for ki in k))
code.transfer(complex='C', tf=smooth_window)
if self.ivarmesh is None:
code.create_whitenoise(dlinear_k='C',
powerspectrum=self.error_ps,
whitenoise='perror')
else:
code.multiply(x1='C', x2=self.ivarmesh**0.5, y='perrorc')
code.c2r(real='perrorc', complex='perror')
def tf(k):
kmesh = sum(ki**2 for ki in k)**0.5
mask = [numpy.ones_like(ki) for ki in k]
mask[2] *= abs(k[2]) >= self.kmin
mask = numpy.prod(mask)
mask2 = numpy.ones_like(mask)
if self.angle > 0:
kperp = (k[0]**2 + k[1]**2)**0.5
kangle = abs(k[2]) / (kperp + 1e-10)
angles = (numpy.arctan(kangle) * 180 / numpy.pi)
mask2[angles < self.angle] = 0
fgmask = mask * mask2
# rankweight = fgmask.sum()
# totweight = pm.comm.allreduce(rankweight)
# rankweight = fgmask.size
# totweight2 = pm.comm.allreduce(rankweight)
# if pm.comm.rank == 0: print('Fraction of zeros on the mask : ', totweight/totweight2)
return fgmask
code.r2c(real='perror', complex='perrorc')
code.transfer(complex='perrorc', tf=tf)
code.c2r(complex='perrorc', real='perror')
code.to_scalar(x='perror', y='chi2')
#code.c2r(real='residual', complex='C')
#code.to_scalar(x='residual', y='chi2')
code.create_whitenoise(dlinear_k='dlinear_k',
powerspectrum=self.prior_ps,
whitenoise='pvar')
code.to_scalar(x='pvar', y='prior')
# the whitenoise is not properly normalized as d_k / P**0.5
code.multiply(x1='prior', x2=Literal(pm.Nmesh.prod()**-1.), y='prior')
code.add(x1='prior', x2='chi2', y='objective')
return code
def LogObjective(self, data):
#likelihood is log(M+M0)
M0 = self.M0
if self.pm.comm.rank == 0:
print('M0 is - %0.3e\n'%self.M0)
logdataM0 = self.pm.create(mode = 'real')
logdataM0.value[...] = numpy.log(data + M0)
code.add(x1='modelsm', x2=Literal(M0), y='modelM0')
code.log(x='modelM0', y='logmodelM0')
code.add(x1='logmodelM0', x2=Literal(logdataM0*-1.), y='residual')
def OvdObjective(self, data):
#likelihood is M/ (Mmean + M0)
M0 = self.M0
datamean = data.cmean()
data[...] /= (M0 + datamean)
inpoints = 1/self.pm.Nmesh.prod()
code.total(x='modelsm', y='modeltot')
code.multiply(x1='modeltot', x2=Literal(inpoints), y='modelmean')
code.add(x1='modelmean', x2=Literal(M0), y='modelmean')
code.pow(x='modelmean', y='invmodelmean', power=-1)
code.multiply(x1='modelsm', x2='invmodelmean', y='modelsm')
code.add(x1='modelsm', x2=Literal(data*-1.), y='residual')
def OvdLogObjective(data):
#likelihood is log( (M+M0)/Mmean )
M0 = self.M0
datamean = data.cmean()
logdataM0 = np.log((data + M0)/datamean())
inpoints = 1/self.pm.Nmesh.prod()
code.total(x='modelsm', y='modeltot')
code.multiply(x1='modeltot', x2=Literal(inpoints), y='modelmean')
code.pow(x='modelmean', y='invmodelmean', power=-1)
code.add(x1='modelsm', x2=Literal(M0), y='modelM0')
code.multiply(x1='modelM0', x2='invmodelmean', y='modelM0invmean')
code.log(x='modelM0invmean', y='logmodel')
code.add(x1='logmodel', x2=Literal(logdataM0*-1.), y='residual')
def get_code(self):
code = self.dynamic_model.get_code()
##Generate differet smoothed fields
code.r2c(real='final', complex='d_k')
code.de_cic(deconvolved='decic', d_k='d_k')
#subtract mean
code.add(x1='decic', x2=Literal(-1.), y='decic')
#
code.r2c(real='decic', complex='d_k')
code.fingauss_smoothing(smoothed='R1', R=self.R1p, d_k='d_k')
code.fingauss_smoothing(smoothed='R2', R=self.R2p, d_k='d_k')
code.multiply(x1='R2', x2=Literal(-1), y='negR2')
code.add(x1='R1', x2='negR2', y='R12')
##Create feature array of 27neighbor field for all
#names = self.mftname
N = len(self.engine.q)
Npf = len(self.pftname)
if self.pwts[0].shape[0] % 27:
Nnb = 1
else:
Nnb = 27
Np = Npf*Nnb
code.assign(x=Literal(numpy.zeros((N, Np))), y='pfeature')
grid = self.pm.generate_uniform_particle_grid(shift=0)
layout = self.engine.pm.decompose(grid)
#pos
for i in range(Npf):
#p
if Nnb == 27:
code.find_neighbours(field=self.pftname[i], features='ptmp')
else:
code.readout(x=Literal(grid), mesh=self.pftname[i], value='ptmp', layout=Literal(layout), resampler='nearest')
#normalize feature
code.add(x1='ptmp', x2=Literal(-1*self.pmx[i*Nnb:(i+1)*Nnb]), y='ptmp1')
code.multiply(x1='ptmp1', x2=Literal(self.psx[i*Nnb:(i+1)*Nnb]**-1), y='ptmp2')
if Nnb == 27:
code.assign_chunk(attribute='pfeature', value='ptmp2', start=i*Nnb, end=Nnb*(i+1))
else:
code.assign_component(attribute='pfeature', value='ptmp2', dim=i)
code.apply_nets(predict='ppredict', features='pfeature', arch=self.parch, Nd=N, t=0, w=self.pwidth)
#renormalize mass
code.reshape_scalar(x='ppredict', y='ppredict')
####Extra factor
#code.relu(x='mpredict', y='mpredict')
#paint
code.paint(x=Literal(grid), mesh='posmesh', layout=Literal(layout), mass='ppredict')
#Smooth
code.assign(x='posmesh', y='model')
return code
def get_code(self):
import numpy
pm = self.mock_model.pm
code = self.mock_model.get_code()
data = self.data.mapp
#sm
if self.smooth is not None:
code.r2c(real='model', complex='d_km')
code.fingauss_smoothing(smoothed='modelsm', R=self.smooth, d_k='d_km')
#
tf = fingauss(pm, self.smooth)
data = data.r2c().apply(lambda k, v: tf(k )*v).c2r()
else: code.assign(x='model', y='modelsm')
#likelihood is (M + M0)/Mmean
M0 = self.M0
data[...] += M0
datamean = data.cmean()
data[...] /= (datamean)
inpoints = 1/pm.Nmesh.prod()
code.add(x1='modelsm', x2=Literal(M0), y='modelsm')
code.total(x='modelsm', y='modeltot')
code.multiply(x1='modeltot', x2=Literal(inpoints), y='modelmean')
#code.add(x1='modelmean', x2=Literal(M0), y='modelmean')
code.pow(x='modelmean', y='invmodelmean', power=-1)
code.multiply(x1='modelsm', x2='invmodelmean', y='modelsm')
code.add(x1='modelsm', x2=Literal(data*-1.), y='residual')
if self.offset:
#offset array has data-model, so add to residual i.e. model-data
if pm.comm.rank == 0: print('Offset in the noise \n')
code.add(x1='residual', x2=Literal(self.noise_model.offset), y='residual')
if self.noised == 2:
if pm.comm.rank == 0: print('2D noise model\n\n')
code.multiply(x1='residual', x2=Literal(self.noise_model.ivar2d ** 0.5), y='residual')
elif self.noised == 3:
if pm.comm.rank == 0: print('3D noise model\n\n')
code.multiply(x1='residual', x2=Literal(self.noise_model.ivar3d ** 0.5), y='residual')
#Smooth
smooth_window = lambda k: numpy.exp(- self.sml ** 2 * sum(ki ** 2 for ki in k))
code.r2c(real='residual', complex='C')
code.transfer(complex='C', tf=smooth_window)
code.c2r(real='residual', complex='C')
#LNorm
if self.L1:
if pm.comm.rank == 0: print('L1 norm objective\n\n')
code.L1norm(x='residual', y='chi2')
else:
if pm.comm.rank == 0: print('L2 norm objective\n\n')
code.to_scalar(x='residual', y='chi2')
#Prior
code.create_whitenoise(dlinear_k='dlinear_k', powerspectrum=self.prior_ps, whitenoise='pvar')
## Smooth prior as well ##
if self.smoothp:
if pm.comm.rank == 0:
print('Smooth Prior')
code.r2c(real='pvar', complex='CC')
code.transfer(complex='CC', tf=smooth_window)
code.c2r(real='pvar', complex='CC')
code.to_scalar(x='pvar', y='prior')
# the whitenoise is not properly normalized as d_k / P**0.5
code.multiply(x1='prior', x2=Literal(pm.Nmesh.prod()**-1.), y='prior')
code.add(x1='prior', x2='chi2', y='objective')
return code
def get_code(self):
import numpy
pm = self.mock_model.pm
code = self.mock_model.get_code()
data = self.data.mapp
#sm
if self.smooth is not None:
code.r2c(real='model', complex='d_km')
code.fingauss_smoothing(smoothed='modelsm', R=self.smooth, d_k='d_km')
#
tf = fingauss(pm, self.smooth)
data = data.r2c().apply(lambda k, v: tf(k )*v).c2r()
else:
code.assign(x='model', y='modelsm')
#likelihood is log(M+M0)
M0 = self.M0
if self.pm.comm.rank == 0:
print('M0 is - %0.3e\n'%self.M0)
logdataM0 = self.pm.create(mode = 'real')
logdataM0.value[...] = numpy.log(data + M0)
code.add(x1='modelsm', x2=Literal(M0), y='modelM0')
code.log(x='modelM0', y='logmodelM0')
code.add(x1='logmodelM0', x2=Literal(logdataM0*-1.), y='residual')
if self.offset:
#offset array has data-model, so add to residual i.e. model-data
if pm.comm.rank == 0:
print('Offset in the noise \n')
code.add(x1='residual', x2=Literal(self.noise_model.offset), y='residual')
if self.noised == 2:
if pm.comm.rank == 0:
print('2D noise model\n\n')
code.multiply(x1='residual', x2=Literal(self.noise_model.ivar2d ** 0.5), y='residual')
elif self.noised == 3:
if pm.comm.rank == 0:
print('3D noise model\n\n')
code.multiply(x1='residual', x2=Literal(self.noise_model.ivar3d ** 0.5), y='residual')
#F**k up right here....this was uncommented until Dec 4!
#code.multiply(x1='residual', x2=Literal(self.noise_model.ivar2d ** 0.5), y='residual')
#Smooth
smooth_window = lambda k: numpy.exp(- self.sml ** 2 * sum(ki ** 2 for ki in k))
code.r2c(real='residual', complex='C')
code.transfer(complex='C', tf=smooth_window)
code.c2r(real='residual', complex='C')
#LNorm
if self.L1:
if pm.comm.rank == 0:
print('L1 norm objective\n\n')
code.L1norm(x='residual', y='chi2')
else:
if pm.comm.rank == 0:
print('L2 norm objective\n\n')
code.to_scalar(x='residual', y='chi2')
#Prior
code.create_whitenoise(dlinear_k='dlinear_k', powerspectrum=self.prior_ps, whitenoise='pvar')
## Smooth prior as well ##
if self.smoothp:
if pm.comm.rank == 0:
print('Smooth Prior')
code.r2c(real='pvar', complex='CC')
code.transfer(complex='CC', tf=smooth_window)
code.c2r(real='pvar', complex='CC')
code.to_scalar(x='pvar', y='prior')
# the whitenoise is not properly normalized as d_k / P**0.5
code.multiply(x1='prior', x2=Literal(pm.Nmesh.prod()**-1.), y='prior')
code.add(x1='prior', x2='chi2', y='objective')
return code
def get_code(self):
import numpy
pm = self.mock_model.pm
code = self.mock_model.get_code()
data = self.data.mapp
code.add(x1='model', x2=Literal(data * -1), y='residual')
code.r2c(real='residual', complex='C')
smooth_window = lambda k: numpy.exp(- self.sml ** 2 * sum(ki ** 2 for ki in k))
code.transfer(complex='C', tf=smooth_window)
#
#code.create_whitenoise(dlinear_k='C', powerspectrum=self.error_ps, whitenoise='perror')
if self.ivarmesh is None:
code.create_whitenoise(dlinear_k='C', powerspectrum=self.error_ps, whitenoise='perror')
else:
if pm.comm.rank == 0: print('Using ivarmesh')
code.multiply(x1='C', x2=Literal(self.ivarmeshc**0.5), y='perrorc')
code.c2r(complex='perrorc', real='perror')
def tf(k):
kmesh = sum(ki ** 2 for ki in k)**0.5
mask = [numpy.ones_like(ki) for ki in k]
mask[2] *= abs(k[2]) >= self.kmin
mask = numpy.prod(mask)
mask2 = numpy.ones_like(mask)
if self.angle > 0:
kperp = (k[0]**2 + k[1]**2)**0.5
kangle = abs(k[2])/(kperp+1e-10)
angles = (numpy.arctan(kangle)*180/numpy.pi)
mask2[angles < self.angle] = 0
fgmask = mask*mask2
return fgmask
code.r2c(real='perror', complex='perrorc')
code.transfer(complex='perrorc', tf=tf)
code.c2r(complex='perrorc', real='perror')
code.multiply(x1='perror', x2=Literal(self.h1wt), y='perror')
code.to_scalar(x='perror', y='chi2HI')
#code.multiply(x1='chi2HI', x2=Literal(self.h1wt), y='chi2HI')
#add LSST data
if self.datalsst is None: code.add(x1='model2', x2=Literal(self.data.mapp2 * -1), y='residuallsst')
else: code.add(x1='model2', x2=Literal(self.datalsst.mapp * -1), y='residuallsst')
##
def phototf(k): #Photoz smoothing
kmesh = sum(ki ** 2 for ki in k)**0.5
kmesh[kmesh == 0] = 1
mumesh = k[2]/(kmesh + 1e-10)
weights = numpy.exp(-0.5 * kmesh**2 * mumesh**2 * self.photosigma**2)
return weights
if pm.comm.rank == 0: print('\nsmoothing for photoz with smoothing = ', self.photosigma)
code.r2c(real='residuallsst', complex='residuallsstc')
code.transfer(complex='residuallsstc', tf=phototf)
smooth_window = lambda k: numpy.exp(- self.sml ** 2 * sum(ki ** 2 for ki in k))
code.transfer(complex='residuallsstc', tf=smooth_window)
code.create_whitenoise(dlinear_k='residuallsstc', powerspectrum=self.error_ps_lsst, whitenoise='perrorlsst')
#code.multiply(x1='residuallsstc', x2=Literal(self.ivarmeshc**0.5), y='perrorclsst')
#code.c2r(complex='perrorclsst', real='perrorlsst')
code.multiply(x1='perrorlsst', x2=Literal(self.lsstwt), y='perrorlsst')
code.to_scalar(x='perrorlsst', y='chi2lsst')
#code.multiply(x1='chi2lsst', x2=Literal(self.lsstwt), y='chi2lsst')
code.create_whitenoise(dlinear_k='dlinear_k', powerspectrum=self.prior_ps, whitenoise='pvar')
code.multiply(x1='pvar', x2=Literal(pm.Nmesh.prod()**-1.), y='pvar')
code.to_scalar(x='pvar', y='prior')
# the whitenoise is not properly normalized as d_k / P**0.5
#In v2 I think I am disabling this normalization since it is waaaay downweighted
#code.multiply(x1='prior', x2=Literal(pm.Nmesh.prod()**-1.), y='prior')
code.add(x1='chi2HI', x2='chi2lsst', y='chi2')
#code.add(x1='prior', x2='chi2', y='objective')
#code.add(x1='prior', x2='chi2HI', y='objective')
code.add(x1='perror', x2='perrorlsst', y='perrorall')
code.add(x1='pvar', x2='perrorall', y='loss')
#
low_pass_window = lambda k: 1 - numpy.exp(- (self.sml + self.smlmax) ** 2 * sum(ki ** 2 for ki in k))
lowpass = False
if self.smlmax is not None:
if self.sml < self.smlmax:
lowpass = True
if pm.comm.rank == 0: print('\nLow pass filtering\n')
if lowpass :
code.r2c(real='loss', complex='lossC')
code.transfer(complex='lossC', tf=low_pass_window)
code.c2r(complex='lossC', real='loss')
code.to_scalar(x='loss', y='objective')
return code
def get_code(self):
code = self.dynamic_model.get_code()
##Generate differet smoothed fields
code.r2c(real='final', complex='d_k')
code.de_cic(deconvolved='decic', d_k='d_k')
#subtract mean
code.add(x1='decic', x2=Literal(-1.), y='decic')
#
code.r2c(real='decic', complex='d_k')
code.fingauss_smoothing(smoothed='R1', R=self.R1p, d_k='d_k')
code.fingauss_smoothing(smoothed='R2', R=self.R2p, d_k='d_k')
code.multiply(x1='R2', x2=Literal(-1), y='negR2')
code.add(x1='R1', x2='negR2', y='R12')
##Create feature array of 27neighbor field for all
#names = self.mftname
N = len(self.engine.q)
Npf = len(self.pftname)
Nm = len(self.mftname)
if self.pwts[0].shape[0] % 27:
Nnb = 1
else:
Nnb = 27
Np = Npf * Nnb
code.assign(x=Literal(numpy.zeros((N, Np))), y='pfeature')
code.assign(x=Literal(numpy.zeros((N, Nm))), y='mfeature')
grid = self.pm.generate_uniform_particle_grid(shift=0)
layout = self.engine.pm.decompose(grid)
code.paint(x='X', mesh='tmp', layout=Literal(layout))
code.paint(x='V', mesh='tmp', layout=Literal(layout))
#pos
for i in range(Npf):
#p
if Nnb == 27:
code.find_neighbours(field=self.pftname[i], features='ptmp')
else:
code.readout(x=Literal(grid),
mesh=self.pftname[i],
value='ptmp',
layout=Literal(layout),
resampler='nearest')
#normalize feature
code.add(x1='ptmp',
x2=Literal(-1 * self.pmx[i * Nnb:(i + 1) * Nnb]),
y='ptmp1')
code.multiply(x1='ptmp1',
x2=Literal(self.psx[i * Nnb:(i + 1) * Nnb]**-1),
y='ptmp2')
if Nnb == 27:
code.assign_chunk(attribute='pfeature',
value='ptmp2',
start=i * Nnb,
end=Nnb * (i + 1))
else:
code.assign_component(attribute='pfeature',
value='ptmp2',
dim=i)
#mass
for i in range(Nm):
#m
code.readout(x=Literal(grid),
mesh=self.mftname[i],
value='mtmp',
layout=Literal(layout),
resampler='nearest')
#normalize feature
code.add(x1='mtmp', x2=Literal(-1 * self.mmx[i]), y='mtmp1')
code.multiply(x1='mtmp1', x2=Literal(self.msx[i]**-1), y='mtmp2')
code.assign_component(attribute='mfeature', value='mtmp2', dim=i)
code.apply_nets(predict='ppredict',
features='pfeature',
arch=self.parch,
Nd=N,
t=0,
w=self.pwidth)
code.apply_nets(predict='mpredict',
features='mfeature',
arch=self.march,
Nd=N,
t=0,
w=self.mwidth)
#renormalize mass
code.multiply(x1='mpredict', x2=Literal(self.msy), y='mpredict')
code.add(x1='mpredict', x2=Literal(self.mmy), y='mpredict')
code.reshape_scalar(x='ppredict', y='ppredict')
code.reshape_scalar(x='mpredict', y='mpredict')
####Extra factor
#code.relu(x='mpredict', y='mpredict')
#paint
code.paint(x=Literal(grid),
mesh='posmesh',
layout=Literal(layout),
mass='ppredict')
code.paint(x=Literal(grid),
mesh='massmesh',
layout=Literal(layout),
mass='mpredict')
code.multiply(x1='posmesh', x2='massmesh', y='premodel')
##RSD below
code.assign(x=Literal(numpy.zeros((N, 3))), y='displace')
#code.assign(x=Literal(numpy.array([0, 0, 1])), y='zdirect')
## Z-component
#code.assign_component(attribute='velz', value='V', dim=2)
#code.matrix_cmul(y='velz', x='v', W='zdirect')
code.matrix_cmul(y='velz',
x='v',
W=numpy.array([0, 0, 1]).reshape(-1, 1))
code.reshape_scalar(x='velz', y='velz')
code.paint(x='X', mesh='momentum', layout=Literal(layout), mass='velz')
code.paint(x='X', mesh='mass', layout=Literal(layout))
code.r2c(real='mass', complex='d_kd')
code.fingauss_smoothing(smoothed='mass', R=self.Rrsd, d_k='d_kd')
code.r2c(real='momentum', complex='d_km')
code.fingauss_smoothing(smoothed='momentum', R=self.Rrsd, d_k='d_km')
code.add(x1='mass', x2=Literal(1e-5), y='mass')
## Divide the shit here
code.pow(y='invmass', x='mass', power=-1)
code.multiply(x1='momentum', x2='invmass', y='velocity')
# #Smooth velocity!
# code.r2c(real='velocity', complex='d_kv')
# code.fingauss_smoothing(smoothed='velocity', R=3, d_k='d_kv')
#
#Readouts
code.readout(x=Literal(grid),
mesh='velocity',
value='velgrid',
layout=Literal(layout),
resampler='nearest')
#Take care of rsdfactor for redshift
code.multiply(x1='velgrid', x2=Literal(self.rsdfactor), y='velgrid')
code.readout(x=Literal(grid),
mesh='premodel',
value='predictmass',
layout=Literal(layout),
resampler='nearest')
code.assign_component(attribute='displace', value='velgrid', dim=2)
code.add(x1=Literal(grid), x2='displace', y='rsdpos')
code.paint(x='rsdpos',
mesh='model',
layout=Literal(layout),
mass='predictmass')
return code
