def MakeMultiScaleCube(self):
if self.CubePSFScales is not None: return
print("Making MultiScale PSFs...", file=log)
LScales = self.GD["HMP"]["Scales"]
if 0 in LScales: LScales.remove(0)
LRatios = self.GD["HMP"]["Ratios"]
NTheta = self.GD["HMP"]["NTheta"]
_, _, nx, ny = self.SubPSF.shape
NScales = len(LScales)
NRatios = len(LRatios)
CubePSFScales = np.zeros(
(NScales + 1 + NRatios * NTheta * (NScales), nx, ny))
Scales = np.array(LScales)
Ratios = np.array(LRatios)
self.ListScales = []
CubePSFScales[0, :, :] = self.SubPSF[0, 0, :, :]
self.ListScales.append({"ModelType": "Delta"})
iSlice = 1
Support = 61
for i in range(NScales):
Minor = Scales[i] / (2. * np.sqrt(2. * np.log(2.)))
Major = Minor
PSFGaussPars = (Major, Minor, 0.)
CubePSFScales[iSlice, :, :] = ModFFTW.ConvolveGaussian(
self.SubPSF, CellSizeRad=1., GaussPars=[PSFGaussPars])[0, 0]
Gauss = ModFFTW.GiveGauss(Support,
CellSizeRad=1.,
GaussPars=PSFGaussPars)
self.ListScales.append({
"ModelType": "Gaussian",
"Model": Gauss,
"ModelParams": PSFGaussPars,
"Scale": i
})
iSlice += 1
Theta = np.arange(0., np.pi - 1e-3, np.pi / NTheta)
for iScale in range(NScales):
for ratio in Ratios:
for th in Theta:
Minor = Scales[iScale] / (2. * np.sqrt(2. * np.log(2.)))
Major = Minor * ratio
PSFGaussPars = (Major, Minor, th)
CubePSFScales[iSlice, :, :] = ModFFTW.ConvolveGaussian(
self.SubPSF, CellSizeRad=1.,
GaussPars=[PSFGaussPars])[0, 0]
Max = np.max(CubePSFScales[iSlice, :, :])
CubePSFScales[iSlice, :, :] /= Max
# pylab.clf()
# pylab.subplot(1,2,1)
# pylab.imshow(CubePSFScales[0,:,:],interpolation="nearest")
# pylab.subplot(1,2,2)
# pylab.imshow(CubePSFScales[iSlice,:,:],interpolation="nearest")
# pylab.title("Scale = %s"%str(PSFGaussPars))
# pylab.draw()
# pylab.show(False)
# pylab.pause(0.1)
iSlice += 1
Gauss = ModFFTW.GiveGauss(
Support, CellSizeRad=1., GaussPars=PSFGaussPars) / Max
self.ListScales.append({
"ModelType": "Gaussian",
"Model": Gauss,
"ModelParams": PSFGaussPars,
"Scale": iScale
})
# Max=np.max(np.max(CubePSFScales,axis=1),axis=1)
# Max=Max.reshape((Max.size,1,1))
# CubePSFScales=CubePSFScales/Max
self.CubePSFScales = np.float32(CubePSFScales)
self.WeightWidth = 6
CellSizeRad = 1.
PSFGaussPars = (self.WeightWidth, self.WeightWidth, 0.)
self.WeightFunction = ModFFTW.GiveGauss(self.SubPSF.shape[-1],
CellSizeRad=1.,
GaussPars=PSFGaussPars)
#self.WeightFunction.fill(1)
self.SupWeightWidth = 3. * self.WeightWidth
print(" ... Done", file=log)
def GiveSpectralIndexMap(self, GaussPars=[(1, 1, 0)], ResidCube=None,
GiveComponents=False, ChannelWeights=None):
# convert to radians
ex, ey, pa = GaussPars
ex *= np.pi/180/np.sqrt(2)/2
ey *= np.pi/180/np.sqrt(2)/2
epar = (ex + ey)/2.0
pa = 0.0
# get in terms of number of cells
CellSizeRad = self.GD['Image']['Cell'] * np.pi / 648000
# get Gaussian kernel
GaussKern = ModFFTW.GiveGauss(self.Npix, CellSizeRad=CellSizeRad, GaussPars=(epar, epar, pa), parallel=False)
# take FT
Fs = np.fft.fftshift
iFs = np.fft.ifftshift
# evaluate model
ModelImage = self.GiveModelImage(self.GridFreqs)
# pad GausKern and take FT
GaussKern = np.pad(GaussKern, self.Npad, mode='constant')
FTshape, _ = GaussKern.shape
from scipy import fftpack as FT
GaussKernhat = FT.fft2(iFs(GaussKern))
# pad and FT of ModelImage
ModelImagehat = np.zeros((self.Nchan, FTshape, FTshape), dtype=np.complex128)
ConvModelImage = np.zeros((self.Nchan, self.Npix, self.Npix), dtype=np.float64)
I = slice(self.Npad, -self.Npad)
for i in range(self.Nchan):
tmp_array = np.pad(ModelImage[i, 0], self.Npad, mode='constant')
ModelImagehat[i] = FT.fft2(iFs(tmp_array)) * GaussKernhat
ConvModelImage[i] = Fs(FT.ifft2(ModelImagehat[i]))[I, I].real
if ResidCube is not None:
ConvModelImage += ResidCube.squeeze()
RMS = np.std(ResidCube.flatten())
Threshold = self.GD["SPIMaps"]["AlphaThreshold"] * RMS
else:
AbsModel = np.abs(ModelImage).squeeze()
MinAbsImage = np.amin(AbsModel, axis=0)
RMS = np.min(np.abs(MinAbsImage.flatten())) # base cutoff on smallest value in model
Threshold = self.GD["SPIMaps"]["AlphaThreshold"] * RMS
# get minimum along any freq axis
MinImage = np.amin(ConvModelImage, axis=0)
MaskIndices = np.argwhere(MinImage > Threshold)
FitCube = ConvModelImage[:, MaskIndices[:, 0], MaskIndices[:, 1]]
if ChannelWeights is None:
weights = np.ones(self.Nchan, dtype=np.float32)
else:
weights = ChannelWeights.astype(np.float32)
if ChannelWeights.size != self.Nchan:
import warnings
warnings.warn("The provided channel weights are of incorrect length. Ignoring weights.", RuntimeWarning)
weights = np.ones(self.Nchan, dtype=np.float32)
try:
import traceback
from africanus.model.spi.dask import fit_spi_components
NCPU = self.GD["Parallel"]["NCPU"]
if NCPU:
from multiprocessing.pool import ThreadPool
import dask
dask.config.set(pool=ThreadPool(NCPU))
else:
import multiprocessing
NCPU = multiprocessing.cpu_count()
import dask.array as da
_, ncomps = FitCube.shape
FitCubeDask = da.from_array(FitCube.T.astype(np.float64), chunks=(ncomps//NCPU, self.Nchan))
weightsDask = da.from_array(weights.astype(np.float64), chunks=(self.Nchan))
freqsDask = da.from_array(np.array(self.GridFreqs).astype(np.float64), chunks=(self.Nchan))
alpha, varalpha, Iref, varIref = fit_spi_components(FitCubeDask, weightsDask,
freqsDask, self.RefFreq,
dtype=np.float64).compute()
except Exception as e:
traceback_str = traceback.format_exc(e)
print("Warning - Failed at importing africanus spi fitter. This could be an issue with the dask " \
"version. Falling back to (slow) scipy version", file=log)
print("Original traceback - ", traceback_str, file=log)
alpha, varalpha, Iref, varIref = self.FreqMachine.FitSPIComponents(FitCube, self.GridFreqs, self.RefFreq)
_, _, nx, ny = ModelImage.shape
alphamap = np.zeros([nx, ny])
Irefmap = np.zeros([nx, ny])
alphastdmap = np.zeros([nx, ny])
Irefstdmap = np.zeros([nx, ny])
alphamap[MaskIndices[:, 0], MaskIndices[:, 1]] = alpha
Irefmap[MaskIndices[:, 0], MaskIndices[:, 1]] = Iref
alphastdmap[MaskIndices[:, 0], MaskIndices[:, 1]] = np.sqrt(varalpha)
Irefstdmap[MaskIndices[:, 0], MaskIndices[:, 1]] = np.sqrt(varIref)
if GiveComponents:
return alphamap[None, None], alphastdmap[None, None], alpha
else:
return alphamap[None, None], alphastdmap[None, None]
def GiveSpectralIndexMap(self,
GaussPars=[(1, 1, 0)],
ResidCube=None,
GiveComponents=False,
ChannelWeights=None):
# convert to radians
ex, ey, pa = GaussPars
ex *= np.pi / 180
ey *= np.pi / 180
pa *= np.pi / 180
# get in terms of number of cells
CellSizeRad = self.GD['Image']['Cell'] * np.pi / 648000
# ex /= self.GD['Image']['Cell'] * np.pi / 648000
# ey /= self.GD['Image']['Cell'] * np.pi / 648000
# get Gaussian kernel
GaussKern = ModFFTW.GiveGauss(self.Npix,
CellSizeRad=CellSizeRad,
GaussPars=(ex, ey, pa),
parallel=False)
# take FT
Fs = np.fft.fftshift
iFs = np.fft.ifftshift
npad = self.Npad
FTarray = self.ScaleMachine.FTMachine.xhatim.view()
FTarray[...] = iFs(np.pad(GaussKern[None, None],
((0, 0), (0, 0), (npad, npad), (npad, npad)),
mode='constant'),
axes=(2, 3))
# this puts the FT in FTarray
self.ScaleMachine.FTMachine.FFTim()
# need to copy since FTarray and FTcube are views to the same array
FTkernel = FTarray.copy()
# evaluate model
ModelImage = self.GiveModelImage(self.GridFreqs)
# pad and take FT
FTcube = self.ScaleMachine.FTMachine.Chatim.view()
FTcube[...] = iFs(np.pad(ModelImage,
((0, 0), (0, 0), (npad, npad), (npad, npad)),
mode='constant'),
axes=(2, 3))
self.ScaleMachine.FTMachine.CFFTim()
# multiply by kernel
FTcube *= FTkernel
# take iFT
self.ScaleMachine.FTMachine.iCFFTim()
I = slice(npad, -npad)
ConvModelImage = Fs(FTcube, axes=(2, 3))[:, :, I, I].real
ConvModelMean = np.mean(ConvModelImage.squeeze(), axis=0)
ConvModelLow = ConvModelImage[-1, 0]
ConvModelHigh = ConvModelImage[0, 0]
if ResidCube is not None:
ConvModelImage += ResidCube
ConvModelImage = ConvModelImage.squeeze()
RMS = np.std(ResidCube.flatten())
Threshold = self.GD["SPIMaps"]["AlphaThreshold"] * RMS
# get minimum along any freq axis
MinImage = np.amin(ConvModelImage, axis=0)
MaskIndices = np.argwhere(MinImage > Threshold)
FitCube = ConvModelImage[:, MaskIndices[:, 0], MaskIndices[:, 1]]
# Initial guess for I0
I0i = ConvModelMean[MaskIndices[:, 0], MaskIndices[:, 1]]
# initial guess for alphas
Ilow = ConvModelLow[MaskIndices[:, 0], MaskIndices[:, 1]] / I0i
Ihigh = ConvModelHigh[MaskIndices[:, 0], MaskIndices[:, 1]] / I0i
alphai = (np.log(Ihigh) -
np.log(Ilow)) / (np.log(self.GridFreqs[0] / self.RefFreq) -
np.log(self.GridFreqs[-1] / self.RefFreq))
# import matplotlib.pyplot as plt
#
# for i in xrange(self.Nchan):
# plt.imshow(np.where(ConvModelImage[i] > Threshold, ConvModelImage[i], 0.0))
# plt.show()
if ChannelWeights is None:
weights = np.ones(self.Nchan, dtype=np.float32)
else:
weights = ChannelWeights.astype(np.float32)
if ChannelWeights.size != self.Nchan:
import warnings
warnings.warn(
"The provided channel weights are of incorrect length. Ignoring weights.",
RuntimeWarning)
weights = np.ones(self.Nchan, dtype=np.float32)
try:
import traceback
from africanus.model.spi.dask import fit_spi_components
NCPU = self.GD["Parallel"]["NCPU"]
if NCPU:
from multiprocessing.pool import ThreadPool
import dask
dask.config.set(pool=ThreadPool(NCPU))
else:
import multiprocessing
NCPU = multiprocessing.cpu_count()
import dask.array as da
_, ncomps = FitCube.shape
FitCubeDask = da.from_array(FitCube.T.astype(np.float64),
chunks=(ncomps // NCPU, self.Nchan))
weightsDask = da.from_array(weights.astype(np.float64),
chunks=(self.Nchan))
freqsDask = da.from_array(self.GridFreqs.astype(np.float64),
chunks=(self.Nchan))
alpha, varalpha, Iref, varIref = fit_spi_components(
FitCubeDask,
weightsDask,
freqsDask,
self.RefFreq,
dtype=np.float64,
I0i=I0i,
alphai=alphai).compute()
# from africanus.model.spi import fit_spi_components
#
# alpha, varalpha, Iref, varIref = fit_spi_components(FitCube.T.astype(np.float64), weights.astype(np.float64),
# self.GridFreqs.astype(np.float64), self.RefFreq.astype(np.float64),
# dtype=np.float64, I0i=I0i, alphai=alphai)
except Exception as e:
traceback_str = traceback.format_exc(e)
print>>log, "Warning - Failed at importing africanus spi fitter. This could be an issue with the dask " \
"version. Falling back to (slow) scipy version"
print >> log, "Original traceback - ", traceback_str
alpha, varalpha, Iref, varIref = self.FreqMachine.FitSPIComponents(
FitCube, self.GridFreqs, self.RefFreq)
_, _, nx, ny = ModelImage.shape
alphamap = np.zeros([nx, ny])
Irefmap = np.zeros([nx, ny])
alphastdmap = np.zeros([nx, ny])
Irefstdmap = np.zeros([nx, ny])
alphamap[MaskIndices[:, 0], MaskIndices[:, 1]] = alpha
Irefmap[MaskIndices[:, 0], MaskIndices[:, 1]] = Iref
alphastdmap[MaskIndices[:, 0], MaskIndices[:, 1]] = np.sqrt(varalpha)
Irefstdmap[MaskIndices[:, 0], MaskIndices[:, 1]] = np.sqrt(varIref)
if GiveComponents:
return alphamap[None, None], alphastdmap[None, None], alpha
else:
return alphamap[None, None], alphastdmap[None, None]