def execute(self, field):
if any(field.shape != self.tile.shape):
raise AttributeError("Field passed to PSF incorrect shape")
outfield = np.zeros_like(field, dtype='float')
zc, yc, xc = self.tile.coords(form='flat')
# here's the plan. we are going to rotate the field so that the current
# plane of interest is in the center. we then crop the image to the
# size of the support so that we are only convolving a small region.
# finally, take the mid plane back out as the solution.
for i, z in enumerate(zc):
# pad the psf slice for the convolution
fs = np.array(self.tile.shape)
fs[0] = self.support[0]
if z < self.zrange[0] or z > self.zrange[1]:
continue
zslice = int(np.clip(z, *self.zrange) - self.zrange[0])
middle = field.shape[0] // 2
subpsf = self._kpad(self.slices[zslice], fs, norm=True)
subfield = np.roll(field, middle - i, axis=0)
subfield = subfield[middle - fs[0] // 2:middle + fs[0] // 2 + 1]
kshape = subfield.shape
kfield = fft.rfftn(subfield, **fftkwargs)
outfield[i] = np.real(
fft.irfftn(kfield * subpsf, s=kshape,
**fftkwargs))[self.support[0] // 2]
return outfield
def _kpad(self, field, finalshape, zpad=False, norm=True):
"""
fftshift and pad the field with zeros until it has size finalshape.
if zpad is off, then no padding is put on the z direction. returns
the fourier transform of the field
"""
currshape = np.array(field.shape)
if any(finalshape < currshape):
raise IndexError("PSF tile size is less than minimum support size")
d = finalshape - currshape
# fix off-by-one issues when going odd to even tile sizes
o = d % 2
d = np.floor_divide(d, 2)
if not zpad:
o[0] = 0
axes = None
pad = tuple((d[i] + o[i], d[i]) for i in [0, 1, 2])
rpsf = np.pad(field, pad, mode='constant', constant_values=0)
rpsf = np.fft.ifftshift(rpsf, axes=axes)
kpsf = fft.rfftn(rpsf, **fftkwargs)
if norm:
kpsf /= kpsf[0, 0, 0]
return kpsf
def execute(self, field):
if any(field.shape != self.tile.shape):
raise AttributeError("Field passed to PSF incorrect shape")
outfield = np.zeros_like(field, dtype='float')
zc, yc, xc = self.tile.coords(form='flat')
kshape = field.shape
kfield = fft.rfftn(field, **fftkwargs)
for k, c in enumerate(self.cheb.coefficients):
pad = self._kpad(c,
finalshape=self.tile.shape,
zpad=True,
norm=False)
cov = np.real(fft.irfftn(kfield * pad, s=kshape, **fftkwargs))
outfield += self.cheb.tk(k, zc)[:, None, None] * cov
return outfield