'''

Implements the calculation of the power spectral density of an irregular

image using Mexican-hat filtering.

For astrophysical rotation measure maps, using some knowledge about

the observed galaxy cluster, also allows deprojection into 3D to estimate

the 3D PDS of fluctuating magnetic fields in the cluster.

'''

def __init__(self, input_path, output_dir='out'):

'''

Loads and preprocesses the data, generates the mask.

Args:

input_path (str): path to the input file, overrides the config file

output_dir (str): path to the output directory

'''

# Input (path to the input image)

self.input_path = input_path

input_fname_parts = os.path.split(self.input_path)[-1].split('.')

if len(input_fname_parts)==1:

print('Error: the input file has no extension, please use .npy or .fits')

sys.exit()

input_fname=('.').join(input_fname_parts[:-1]) # file name w/o extension

# output folder

if not os.path.exists(output_dir): os.mkdir(output_dir)

# output file names

self.output_paths = {}

out_files = ['deproj','smooth','mask', 'smooth3d']

for f in out_files:

self.output_paths[f] = os.path.join(output_dir,

f+ '_'+input_fname+'.npy')

self.output_paths['pds'] = os.path.join(output_dir,

'pds_'+input_fname+'.txt')

#input

# mask_path = config.get("input", "mask_path")

#

# #output

# rm_out_path = config.get("output", "rm_out_path")

#

# if regime=="2d":

# #gen2d_params

# lx = config.getint("gen2d_params", "Lx")

# ly = config.getint("gen2d_params", "Ly")

# p1 = config.getfloat("gen2d_params", "p1")

# p2 = config.getfloat("gen2d_params", "p2")

# kb = config.getfloat("gen2d_params", "kb")

# C = config.getfloat("gen2d_params","C")

# apply_mask = config.getboolean("gen2d_params", "apply_mask")

# bw = config.getfloat("gen2d_params", "beam_width")

#

# elif regime=="3d":

# #gen3d_params

# lx = config.getint("gen3d_params", "Lx")

# ly = config.getint("gen3d_params", "Ly")

# lz = config.getint("gen3d_params", "Lz")

# p1 = config.getfloat("gen3d_params", "p1")

# p2 = config.getfloat("gen3d_params", "p2")

# kb = config.getfloat("gen3d_params", "kb")

# C = config.getfloat("gen3d_params", "C")

# apply_mask = config.getboolean("gen3d_params", "apply_mask")

# inclin = config.getfloat("gen3d_params", "inclin")

# alpha = config.getfloat("gen3d_params", "alpha")

ftype = self.input_path.split('.')[-1].lower()

# load data

d = load_data(self.input_path, ftype=ftype, print_info=True)

# crop it

ir, jr = np.nonzero(np.invert(np.isnan(d)))

imin, jmin, imax, jmax = ir.min(),jr.min(), ir.max(),jr.max()

self.data = d[imin:imax+1, jmin:jmax+1]

self.data_dim = self.data.shape

# adjust the cluster center location after the cropping

self.cluster_params = {}

self.cluster_params['center'] = [-imin,-jmin]

# make the mask given the cropped input data

self.mask = np.array(np.invert(np.isnan(self.data)), dtype=float)

# save the mask

write_data(self.output_paths['mask'], self.mask, ftype='npy')

def deproject(self, config_path='params.cfg'):

'''

Deproject a rotation measure map by dividing it by a smooth model

of the cluster with the parameters taken from the configuration file.

Args:

config_path (str): path to the configuration file

'''

# Observed cluster parameters parameters (use only when need to deproject 3D PDS)

config = configparser.ConfigParser()

try:

config.read(config_path)

except IOError:

print('Error: No config file found at '+config_path)

return

# beta model parameters

self.cluster_params['ne0'] = config.getfloat("cluster_params", "ne0")

self.cluster_params['rc'] = config.getfloat("cluster_params", "rc")

self.cluster_params['beta'] = config.getfloat("cluster_params", "beta")

# size of pixel in kiloparsecs

self.cluster_params['kpc_px'] = config.getfloat("cluster_params", "kpc_px")

# Galactic background to subtract from the observed RM map

self.cluster_params['gal_bg'] = config.getfloat("cluster_params", "gal_bg")

# location of the cluster center on the sky in pixels

self.cluster_params['center'][0] += config.getint("cluster_params", "ix0")

self.cluster_params['center'][1] += config.getint("cluster_params", "iy0")

#recovery

self.recovery_params = {}

self.recovery_params['alpha'] = config.getfloat("recovery_params", "alpha")

self.recovery_params['inclin'] = config.getfloat("recovery_params", "inclin")

# depth of the 3D box used to generate the smooth model

self.recovery_params['lz'] = config.getint("recovery_params", "Lz")

# subtract the background

self.data -= self.cluster_params['gal_bg']

kpc_px = self.cluster_params['kpc_px']

lx,ly = self.data.shape

print('generate a 3D smooth model of the cluster...')

smod = gen3d_smooth_model(shape=(lx,ly),

cluster_params=self.cluster_params,

recovery_params=self.recovery_params)

write_data(self.output_paths['smooth3d'], smod[:,ly//2,:], ftype='npy')

print('deproject the image...')

# integrate the smooth model along the line of sight

I0 = np.sqrt( (smod**2).sum(axis=2)) * kpc_px * 812.

write_data(self.output_paths['smooth'], I0, ftype='npy')

# divide the image by the smooth model

ind = (I0!=0.) * np.invert(np.isnan(self.data))

self.data[ind] /= I0[ind]

write_data(self.output_paths['deproj'], self.data, ftype='npy')

print('deprojection done\n')

def get_spectrum(self, ns=30, smin=3, smax=100, p_corr=0., nt=1):

'''

Calculate the isotropic power spectral density (PDS) for an image of

irregular shape.

Args:

ns (int): number of length-scales at which the PDS is calculated

smin (int): minimum length-scale in pixels (>=2)

smax (int): maximum length-scale in pixels

p_corr: spectral slope if known beforehand (to correct normalization)

nt: number of threads for multiprocessing

Returns:

kr: radial wavenumbers in px^(-1)

pds: PDS measured at kr

'''

print('calculate the PDS...')

lx,ly = self.data.shape

# subtract the mean component

self.data[self.mask==1.] -= self.data[self.mask==1.].mean()

# replace NaNs by zeros, add padding

dp = add_padding(self.data, width=0.5)

mp = add_padding(self.mask, width=0.5)

# Fourier transform of the padded image and mask

dpf = fftpack.fftn(dp)/np.sqrt(lx*ly)

mpf = fftpack.fftn(mp)/np.sqrt(lx*ly)

#pds = mp.Array('d', np.zeros(N))

pds = np.zeros(ns)

# set the range of radial wave numbers

lx,ly = dp.shape

krmin = 1./smax if smax<=0.5*max(lx,ly) else 2./max(lx,ly) # mind zero padding

krmax = 1./smin # 1/2 is the Nyquist frequency

fk = (krmax/krmin)**(1./(ns-1))

kr = krmin * np.power(fk, np.arange(ns))

eps = 1e-3

# Split calculation for different wave numbers between threads

params = [{'i':i, 'kri':kr[i], 'datf':dpf, 'mskf':mpf, 'msk':mp,

'eps':eps, 'p_corr':p_corr} for i in range(ns)]

if nt>1:

with multiprocessing.Pool(nt) as pool:

pds = pool.map(calc_pds_single, params)

pool.join()

else:

pds = map(calc_pds_single, params)

if np.any(np.isnan(pds)): print('Error: NaN in the PDS')

self.kr = kr

self.pds = pds

np.savetxt(self.output_paths['pds'], np.asarray([self.kr,self.pds]))

print('PDS calculated and saved as '+self.output_paths['pds'])

'''

Calculate the PDS value for a single radial wavenumber kri=kr[i].

Args:

params (dict) with items:

kri (float): radial wavenumber

datf, mskf (2d arrays): Fourier transforms of the padded image and mask

msk (2d array): padded mask

eps (float): sets the width difference between the two Gaussian filters

p_corr (float): spectral slope if known beforehand (to correct normalization)

Returns:

float: PDS value at kri

kri sets the width of the Mexican hat filter applied to the image.

The Mexican hat is the difference of two Gaussian filters of similar

width: G1 and G2.

The PDS value is obtained from the total variance of the convolved image.

'''

i,kri,eps,p_corr = params['i'],params['kri'],params['eps'],params['p_corr']

datf,mskf,msk = params['datf'],params['mskf'],params['msk']

(lx,ly) = datf.shape

Ikr = np.zeros((lx,ly))

# widths of the two Gaussian filters whose difference is the Mexican hat filter

sigma = 0.225079/kri

sigma1 = sigma/np.sqrt(1.+eps)

sigma2 = sigma*np.sqrt(1.+eps)

# Set Gaussian filters on the grid in k-space

fG1 = lambda kx,ky: np.exp(-2 * np.pi**2 * (kx**2+ky**2) * sigma1**2)

fG2 = lambda kx,ky: np.exp(-2 * np.pi**2 * (kx**2+ky**2) * sigma2**2)

kx,ky = np.meshgrid(np.linspace(0,1.-1./lx,lx),

np.linspace(0,1.-1./ly,ly), indexing='ij')

G1 = fG1(kx,ky)+fG1(1.-kx,ky)+fG1(kx,1.-ky)+fG1(1.-kx,1.-ky)

G2 = fG2(kx,ky)+fG2(1.-kx,ky)+fG2(kx,1.-ky)+fG2(1.-kx,1.-ky)

# Mexican hat filter power

fpwr = (np.abs(G1-G2)**2).sum()

if not np.all(msk==1.):

# general case

# multiply the data and mask by the filter in k-space

# and take the inverse Fourier transform

conv_dG1 = fftpack.ifftn(G1*datf)

conv_dG2 = fftpack.ifftn(G2*datf)

conv_mG1 = fftpack.ifftn(G1*mskf)

conv_mG2 = fftpack.ifftn(G2*mskf)

#print abs(datf).max()

'''

Divide the G1/G2 image convolutions by the respective mask convolutions

and get the difference beteen G1 and G2. Division by the mask helps

remove the spurious harmonics resulting from the sharp edges of

the irregular image.

'''

nz = msk==1.

Ikr[nz] =(conv_dG1[nz].real / conv_mG1[nz].real -

conv_dG2[nz].real / conv_mG2[nz].real)

else:

'''

Case of a periodic rectangular image.

No need to divide by the mask -- simlpy convolve with the Mexican hat.

'''

conv_dG = fftpack.ifftn((G1-G2) * datf)

Ikr = conv_dG * np.sqrt(lx*ly)

vr = np.sqrt((np.abs(Ikr)**2).mean())

vf = np.sqrt((np.abs((G1-G2) * datf)**2).mean())

print(f'checking the Parseval theorem (backward FFT):\nvr={vr}\nvf={vf}')

# total variance

var = ( lx*ly / msk.sum() * (np.abs(Ikr)**2).sum() )

# divide by the filter power

pds_tilda = var / fpwr

# correction coefficient in case the spectral slope is known beforehand

# (used for testing on mock images to fix normalization)

corr_f = (2**(p_corr/2) *

special.gamma(3-p_corr/2)/

special.gamma(3))

return pds_tilda / corr_f

#print 'kr='+str(kr)+' P(kr)='+str(pds[i])

#make an example of the filtered image

#if i == 20:

# figure()

# imshow(Ikr)

# savefig('../img/filt.ps')

# close()

# Ikr1 = trim(Ikr).copy()

# Ikr1[Ikr1==0.]=NaN