[Python] Stack multiple fits file with different WCS into a 3D array

I have generated a fits image of size 2048 x 2048 for each frequency channel, and it is saved as a separate fits file. I want to combine all of them into a single 3d array such that the WCS header of the 3D array is equivalent to the WCS header of the last frequency channel. However, if I do it the way prescribed in the code below, the individual fits do have different was, and the image is smaller than the coordinate grid (lower frequencies have smaller images while it increases with frequency). But once it is combined, it's not true. PS: I tried the reproject_interp function, but it doesn't preserve the flux, i.e. it introduces regridding artefacts.

Can somebody please help me?

import numpy as np
import os
from tqdm import tqdm
import astropy.io.fits as pyfits

#%%
f1_Hz = 138.88e+6; f2_Hz = 169.48e+6

delta_f = 0.120 #fine chan frequency in MHz
nchan = 256 #number of freuency channels
imsize= 2048 #image size
centre_freq_MHz = (f1_Hz/1e6+ (nchan/2)*delta_f)
freq_MHz = f1_Hz/1e6

dec_array=-26.70331940
phase_centre = 57.1863376776493 #phase value in degrees

xpos, ypos = 1600, 1440
flux = 1
image = np.zeros((imsize, imsize),dtype=np.float32)
SKY = np.zeros((imsize, imsize),dtype=np.float32)
SKY[xpos, ypos] = flux

freqs = np.arange(nchan)*delta_f+freq_MHz

#creating multiple fits file with different wcs (cdelt)
for i,freq in tqdm(enumerate(freqs)):
image = SKY*((freq/centre_freq_MHz)**-2.5)
hdu = pyfits.PrimaryHDU(image) # Extract 2D slice for this frequency

pixscale = 180.0 / (imsize / 2) # Base pixel scale
cdelt_value = (pixscale / np.pi) * (freq / target_freq)

# Update the FITS header for this frequency
hdu.header['NAXIS'] = 2 # Since it's now a 2D image
hdu.header['NAXIS1'] = imsize
hdu.header['NAXIS2'] = imsize
hdu.header['CRPIX1'] = imsize // 2 + 1
hdu.header['CDELT1'] = -cdelt_value # RA pixel scale
hdu.header['CRVAL1'] = phase_centre
hdu.header['CTYPE1'] = "RA---SIN"

hdu.header['CRPIX2'] = imsize // 2 + 1
hdu.header['CDELT2'] = cdelt_value # Frequency-dependent DEC pixel scale
hdu.header['CRVAL2'] = dec_array
hdu.header['CTYPE2'] = "DEC--SIN"


# Define the output filename
output_filename = os.path.join(f'sky_{freq:.3f}MHz.fits')

# Write the FITS file for this frequency channel
hdu.writeto(output_filename, overwrite=True)


# Trying to combine all the fits into a 3d array with a different cdelt value
image_cube=np.zeros((nchan, imsize, imsize),dtype=np.float32)
hdu = pyfits.PrimaryHDU(image_cube)
pixscale=180.0/(imsize/2) # for all-sky

hdu.header.update(NAXIS=3)
hdu.header.update(NAXIS1=imsize)
hdu.header.update(NAXIS2=imsize)
hdu.header.update(NAXIS3=nchan) # 256 frequency channels

hdu.header.update(CRPIX1= imsize//2+1)
hdu.header.update(CDELT1=-pixscale/np.pi)
hdu.header.update(CRVAL1=phase_centre)
hdu.header.update(CTYPE1="RA---SIN")

hdu.header.update(CRPIX2=imsize//2+1)
hdu.header.update(CDELT2=pixscale/np.pi)
hdu.header.update(CRVAL2=dec_array)
hdu.header.update(CTYPE2="DEC--SIN")

hdu.header.update(CRPIX3=1)
hdu.header.update(CDELT3=delta_f*1e6)
hdu.header.update(CRVAL3=f1_Hz)
hdu.header.update(CTYPE3='FREQ')
for i,freq in tqdm(enumerate(freqs)):
fname='sky_%.3fMHz.fits'%(freq)
# print('opening '+fname)
with pyfits.open(fname,mode='readonly') as hdul:
image_cube[i ] = hdul[0].data

sky_filename = "SKY_CUBE.fits"
hdu.writeto(sky_filename, overwrite=True)

Continue reading...