I have a pretty long Python 3 code that oversamples many images, then fits them to 2D Gaussians to find the relevant parameters of the sources in the images (centroid coordinates, amplitude, etc), shifts it to the center, normalizes, and then stacks all of them at the end by a median. I am trying to do this with data from about 50,000 images. I already extracted the data and have it saved from a previous code, which gave me a binary file of all the image data (numpy arrays) in a file called initial_data.data, and I also used an astropy wcs package to find good guesses for the centroid coordinates, so those are also saved (in the same order) in centroid_coords.data. When I first tried running my code on all the images, it said Killed during the oversampling part of the code. I then decided I'd try to break up all my images into subsections and run the code on one subsection at a time. This is okay if it's only a few subsections, but would become too tedious and impractical if I'd have to operate on say only 5,000 images at a time! Below is the rest of my code, where I attempted subsections of thirds, and tried just operating on the first one (so I entered in 1 after the first user input line).
import pickle
import numpy as np
import re
from scipy import optimize, interpolate
from astropy.io import fits
from astropy.nddata import Cutout2D
frac = int(input('Subset 1, 2, or 3? ')) # select which subset of images to operate on
oversampled_file = 'oversampled_images_frac%d.data' % (frac)
params_file = 'gparams_2D_frac%d.data' % (frac)
final_images_file = 'images_b4stacking_2D_frac%d.data' % (frac)
FWHM_file = 'FWHM_2D_frac%d.data' % (frac)
fitsfilename = 'stacked_images_frac%d.fits' % (frac)
# Define subsets of total image data based on chosen fraction
if frac==1:
start, end = 0, 50000//3
print('Will operate on 1st subsection: %d elements' % (end-start))
elif frac==2:
start, end = 50000//3, 2*50000//3
print('Will operate on 2nd subsection: %d elements' % (end-start))
elif frac==3:
start, end = 2*50000//3, 49999
print('Will operate on 3rd subsection: %d elements' % (end-start))
else:
print('Subsection of 1, 2, or 3 not detected')
# Read in a subset of initial data and centroid coordinates
with open('initial_data.data', 'rb') as f_r:
initial_data = pickle.load(f_r)[start:end]
with open('centroid_coords.data', 'rb') as f_r:
centroid_coords = pickle.load(f_r)[start:end]
# Oversample images
def oversample(data): # pixels -> NxN pixels
oversampled = []
N = 5
c=1
for data_set in data:
Y, X = np.shape(data_set)
x = np.linspace(0, 0.5, X)
y = np.linspace(0, 0.5, Y)
f = interpolate.interp2d(x, y, data_set, kind='cubic')
Xnew = np.linspace(0, 0.5, X*N)
Ynew = np.linspace(0, 0.5, Y*N)
new_data = f(Xnew, Ynew)
oversampled.append(new_data)
if c%50==0:
print('Oversampling %f%% complete' % (c*100/len(data)))
c+=1
return np.array(oversampled) # array of 2D arrays
resampled_data = oversample(initial_data)
# Save oversampled image data -- array by array to save RAM
with open(oversampled_file, 'wb') as f:
for image in resampled_data:
pickle.dump(image, f)
# Fit to 2D Gaussian
def gaussian_func(xy, x0, y0, sigma_x, sigma_y, amp, theta, offset): # (x0, y0) is center
x, y = xy
a = (np.cos(theta))**2/(2*sigma_x**2) + (np.sin(theta))**2/(2*sigma_y**2)
b = -np.sin(2*theta)/(4*sigma_x**2) + np.sin(2*theta)/(4*sigma_y**2)
c = (np.sin(theta))**2/(2*sigma_x**2) + (np.cos(theta))**2/(2*sigma_y**2)
inner = a * (x-x0)**2
inner += 2*b*(x-x0)*(y-y0)
inner += c * (y-y0)**2
return (offset + amp * np.exp(-inner)).ravel()
def Sigma2width(sigma):
return 2 * np.sqrt(2*np.log(2)) * sigma
def generate(data_set):
xvec = np.arange(0, np.shape(data_set)[1], 1)
yvec = np.arange(0, np.shape(data_set)[0], 1)
X, Y = np.meshgrid(xvec, yvec)
return X, Y
def fit_to_model(data):
gaussian_params = []
N = 5
theta_guess = 6
c=1
for i in range(len(data)):
data_set = data[i]
# Guess parameters
# Centroid and amplitude
offset = np.median(data_set)
x0, y0 = centroid_coords[i][0]*N, centroid_coords[i][1]*N
x0, y0 = int(round(x0)), int(round(y0))
if x0<500 and x0>300 and y0<500 and y0>300: #skips images that have likely inaccurate centroids
subimage = data_set[y0-80:y0+80, x0-80:x0+80]
amp = np.max(subimage)
bg_sub_amp = amp-offset
# Sigmas and offset
ylim, xlim = np.shape(subimage)
x, y = np.arange(0, xlim, 1), np.arange(0, ylim, 1)
ypix, xpix = np.where(subimage==amp)
y_range = np.take(subimage-offset, ypix[0], axis=0)
x_range = np.take(subimage-offset, xpix[0], axis=1)
half_max = bg_sub_amp/2
d_x = x_range - half_max
d_y = y_range - half_max
indices_x = np.where(d_x > 0)[0]
indices_y = np.where(d_y > 0)[0]
width_x = len(indices_x) # estimate of integer pixels only
width_y = len(indices_y)
sigma_x = width_x/(2*np.sqrt(2*np.log(2)))
sigma_y = width_y/(2*np.sqrt(2*np.log(2)))
# Guesses
guesses = [x0, y0, sigma_x, sigma_y, amp, theta_guess, offset]
# Fit to Gaussian
xx, yy = generate(data_set)
try:
pred_params, uncert_cov = optimize.curve_fit(gaussian_func, (xx.ravel(), yy.ravel()), data_set.ravel(), p0=guesses)
except (OptimizeWarning, RuntimeError): #not working?
print('Failed to find fit; skipping...')
continue
gaussian_params.append(pred_params)
if c%10==0:
print('2D Gaussian fit complete for %f%% of images' % (c*100/len(data)))
c+=1
else: #not a good Gaussian fit (doesn't have centroid coords I want)
gaussian_params.append([])
return gaussian_params
params = fit_to_model(resampled_data)
# Save Gaussian params to data file
with open(params_file, 'wb') as f:
for param_set in params:
pickle.dump(param_set, f)
# Make centered cutouts of size 600x600 pixels
def cutout(data_list):
centered, FWHM_images = [], []
size = 600
c=1
for i in range(len(data_list)):
data_set = data_list[i]
if params[i] != []:
if np.shape(data_set)[0] == np.shape(data_set)[1]: #eliminates images that had objects too close to the edge
x, y = params[i][0], params[i][1]
FWHM_x, FWHM_y = Sigma2width(np.abs(params[i][2])), Sigma2width(np.abs(params[i][3]))
# Pass on images that have bad fit (too big FWHMs)
if FWHM_x > 50 or FWHM_y > 50:
continue
FWHM_images.append((FWHM_x, FWHM_y))
cutout = Cutout2D(data_set, (x, y), size).data
centered.append(cutout)
if c%50==0:
print('Centered cutouts complete for %f%% of images' % (c*100/len(data_list)))
c+=1
return centered, FWHM_images
centered_cutouts, FWHM_images = cutout(resampled_data)
# Normalize
def normalize(data):
img_count = 1
norm = []
for data_set in data:
data_set *= 1/data_set.max()
norm.append(data_set)
if img_count%50==0:
print('Normalization complete for %f%% of images' % (img_count*100/len(data)))
img_count+=1
print('NUMBER OF FINAL IMAGES:', img_count)
return np.array(norm)
final_images = normalize(centered_cutouts)
# Save final (normalized) image data and individual FWHMs before stacking
with open(final_images_file, 'wb') as f:
for image in final_images:
pickle.dump(image, f)
with open(FWHM_file, 'wb') as f:
for width in FWHM_images:
pickle.dump(width, f)
# Stack images
stacked_image = np.median(final_images, axis=0)
# Write to new FITS file
hdu = fits.PrimaryHDU(stacked_image)
hdu.writeto(fitsfilename, overwrite=True)
As you can see, I also have multiple checkpoints in my code where new data is saved to a binary data file. This is because if some error occurs, I want to be able to simply read in the last saved data instead of running everything again, since it takes a while. I believe it would take about 11 hours for a third of my 50,000 images, so I ran it overnight on screen. However, now I see another Killed message. Here are the last two output lines I see from my attempt of running this code:
2D Gaussian fit complete for 5.340214% of images
Killed
I generally have a lot of RAM available (close to 100 GB), as I'm using a large server at work. The numpy array data for each initial image generally has a shape of about 157x158, and then of course after oversampling they become larger than 700x700.
I don't want to have to divide into many small subsections, because then I'll just be sitting here running them one by one all day. Is there anything else I can do to ensure that the code runs through without being killed? Also, in reality I will probably just want to save all the image data at the end of my code to a file for each subsection, because what I really need to do is take the median of all the images together (from all subsections). Perhaps generators would prove to be useful? I've read things about them but not sure if that would actually work here. I also have access to multiple CPUs, but I'm not sure how to use multiprocessing in Python (if that would even help). Thank you!