Synthetic polarimetry image generator

gen_pol.py

# This file contains a simple code to generate synthetic image to simulate the data obtained by a polarimeter like IAGPOL
# It also reduce the images using astropop routines.

from astropy.io import fits
from astropy.table import vstack
from astropy.stats import gaussian_fwhm_to_sigma
from astropop.py_utils import mkdir_p
from astropop.pipelines.polarimetry_scripts import run_pccdpack, process_polarimetry
import numpy as np

def gaussian_2d(x, y, x0, y0, sigma_x, sigma_y, theta, flux, sky):
    cost2 = cos(theta)**2
    sint2 = sin(theta)**2
    sin2t = sin(2*theta)
    sigx2 = 2*sigma_x**2
    sigy2 = 2*sigma_y**2
    a = (cost2/sigx2) + (sint2/sigy2)
    b = -(sin2t/(2*sigx2)) + (sin2t/(2*sigy2))
    c = (sint2/sigx2) + (cost2/sigy2)
    xi = x - x0
    yi = y - y0
    return sky + (flux/(sqrt(2*pi*sigma_x*sigma_y))) * exp(-(a*xi**2 + 2*b*xi*yi + c*yi**2))

def gen_image(scale, psi, p, theta):
    rdnoise = 18.5
    bias = 300
    coords = [(80, 135), (150, 47), (60, 60), (161, 105), (225, 150)]
    scales = [1, 0.2, 0.05, 0.1, 0.01]
    sigma = 5.0*gaussian_fwhm_to_sigma

    q = p*np.cos(2*np.radians(theta))
    u = p*np.sin(2*np.radians(theta))
    
    z = q*np.cos(4*np.radians(psi)) + u*np.sin(4*np.radians(psi))
    c = (z+1)/2
    
    base_im = np.random.normal(bias, rdnoise, (256, 256)).astype('uint16')
    
    i_y, i_x = np.indices((256, 256))
    for coo, sca in zip(coords, scales):
        im_o = gaussian_2d(i_x, i_y, coo[0], coo[1], sigma, sigma, 0, sca*scale, 0)*c
        im_e = gaussian_2d(i_x, i_y, coo[0]-20, coo[1]+20, sigma, sigma, 0, sca*scale, 0)*(1-c)
        im_o = np.round(im_o).astype('uint16')
        im_e = np.round(im_e).astype('uint16')
        
        base_im += im_o
        base_im += im_e
    
    return np.random.poisson(base_im).astype('uint16')

p = 0.1
theta = 15
results = {'pccdpack': {}, 'astropop': {}}

for f in [1500, 2000,
          2500, 3000, 3500, 4000, 5000, 10000, 15000,
          20000, 30000, 50000, 100000,
          200000, 300000, 500000
          ]:
    mkdir_p(f'pol_models/{f}/')
    images = []
    for i in np.arange(16):
        ik = hex(i)[2]
        im = gen_image(f, i*22.5, p, theta)
        head = fits.Header()
        head['LAMINA'] = f'L{ik}'
        head['LAM-POS'] = f'{i}'
        head['FILTER'] = 'F2'
        head['FIL-POS'] = '2'
        head['GAIN'] = 1.0
        head['RDNOISE'] = 18.5
        hdu = fits.PrimaryHDU(im.astype('uint16'), header=head)
        images.append(hdu)
        #hdu.writeto(f'pol_models/{f}/model_{ik}.fits')
    try:
        results['pccdpack'][f] = run_pccdpack(images, retarder_type='half', retarder_direction=1,
                                              retarder_rotation=22.5,
                                              retarder_key='LAM-POS', rdnoise_key='RDNOISE',
                                              astrometry_calib=False,
                                              r=np.arange(1, 20, 1), r_in=50, r_out=60,
                                              detect_fwhm=None, detect_snr=5)
        results['astropop'][f]  = process_polarimetry(images, align_images=False, retarder_type='half',
                                                      retarder_direction=1,
                                                      retarder_rotation=22.5, retarder_key='LAM-POS',
                                                      rdnoise_key='RDNOISE', astrometry_calib=False,
                                                      r=np.arange(1, 20, 1), r_in=50, r_out=60,
                                                      detect_fwhm=None, detect_snr=3,
                                                      photometry_type='aperture',
                                                      delta_x=-20, delta_y=20)
    except Exception as e:
        # just pass the error
        print(e)

# save the tables
pop = None
pccd = None

for i in results['pccdpack'].keys():
    if i in results['astropop'].keys():
        p = results['astropop'][i][0]['aperture']
        if pop is None:
            pop = p
        else:
            pop = vstack([pop, p])
    p = results['pccdpack'][i][2]
    if pccd is None:
        pccd = p
    else:
        pccd = vstack([pccd, p])
        
pop.write('pol_models/astropop_synt.fits', format='fits')
pccd.write('pol_models/pccdpack_synt.fits', format='fits')

from matplotlib import pyplot as plt
# generate the comparison figure
f = plt.figure(figsize=(12,6))

ax = f.add_subplot(121)
ax.plot(pccd['P']/pccd['SIGMA'], pccd['P'], 'k.')
ax.set_title('pccdpack')
ax.set_xlim(0.1, 650)
ax.set_ylim(0, 0.5)
ax.set_xlabel('SNR')
ax.set_ylabel('P')
ax.set_xscale('log')
ax = f.add_subplot(122)

ax.plot(pop['p']/pop['p_error'], pop['p'], 'k.')
ax.set_xlim(0.1, 650)
ax.set_ylim(0, 0.5)
ax.set_xlabel('SNR')
ax.set_ylabel('P')
ax.set_xscale('log')
ax.set_title('astropop')

plt.show()