How do I repeat the same task on different data without duplication

Question

I am new to Python and have been writing scripts that do some tasks to certain FITS files. The script that I currently use is this:

# General routines
import numpy as np
import math 
import pyfits
import matplotlib.pyplot as plt
import pylab as py
from scipy.optimize import curve_fit

# Load the FITS file into the program
hdulist1 = pyfits.open('/home/ssridhar/mock_test_files/most_massive_halo_density.fits')
hdulist2 = pyfits.open('/home/ssridhar/mock_test_files/less_massive_halo_density.fits')

tbdata1 = hdulist1[1].data
tbdata2 = hdulist2[1].data

# variables for table1
ra_1 = tbdata1.field('ra')
dec_1 = tbdata1.field('dec')
zcosmo_1 = tbdata1.field('zcosmo')
r200_1 = tbdata1.field('halo_r200')
r_1 = tbdata1.field('dist_to_center')
m200_1 = tbdata1.field('halo_mass')
rho_0_1 = tbdata1.field('rho_0')
rho_1 = tbdata1.field('density')

# variables for table2
ra_2 = tbdata2.field('ra')
dec_2 = tbdata2.field('dec')
zcosmo_2 = tbdata2.field('zcosmo')
r200_2 = tbdata2.field('halo_r200')
r_2 = tbdata2.field('dist_to_center')
m200_2 = tbdata2.field('halo_mass')
rho_0_2 = tbdata2.field('rho_0')
rho_2= tbdata2.field('density')

# global variables
pi = math.pi
rad = pi/180 # converting degrees to radians

c_m = 25
delta_c_m = (200*c_m**3)/(3*(math.log(1+c_m)-(c_m/(1+c_m))))
c_l = 5
delta_c_l = (200*c_l**3)/(3*(math.log(1+c_l)-(c_l/(1+c_l))))

r_s_1 = r200_1/c_m
r_s_2 = r200_2/c_l

# finding x = r/r_s
x_1 = np.linspace(0.0,3.5,num=1242)/r_s_1
x_2 = np.linspace(0.0,2.7,num=135)/r_s_2

# splitting values to find sigma(x)

a_11 = (2*delta_c_m*rho_0_1*r_s_1)/(x_1**2-1)
b1_1 = (2/(np.sqrt(1-x_1**2)))
c1_1 = np.arctanh(np.sqrt((1-x_1)/(1+x_1)))
b2_1 = 2/(np.sqrt(x_1**2-1))
c2_1 = np.arctan(np.sqrt((x_1-1)/(1+x_1)))

a_12 = (2*delta_c_l*rho_0_2*r_s_2)/(x_2**2-1)
b1_2 = (2/(np.sqrt(1-x_2**2)))
c1_2 = np.arctanh(np.sqrt((1-x_2)/(1+x_2)))
b2_2 = 2/(np.sqrt(x_2**2-1))
c2_2 = np.arctan(np.sqrt((x_2-1)/(1+x_2)))


# implementing the conditions for x
a_1_1 = a_11[x_1<1]
b_1_1 = b1_1[x_1<1]
c_1_1 = c1_1[x_1<1]

a_2_1 = a_11[x_1>1]
b_2_1 = b2_1[x_1>1]
c_2_1 = c2_1[x_1>1]

a_1_2 = a_12[x_2<1]
b_1_2 = b1_2[x_2<1]
c_1_2 = c1_2[x_2<1]

a_2_2 = a_12[x_2>1]
b_2_2 = b2_2[x_2>1]
c_2_2 = c2_2[x_2>1]


# finding sigma(x), the projected NFW profile

sigma_x_m1 = a_1_1*(1-(b_1_1*c_1_1))
sigma_x_m2 = a_2_1*(1-(b_2_1*c_2_1))
sigma_x_m = np.concatenate((sigma_x_m1,sigma_x_m2))

sigma_x_l1 = a_1_2*(1-(b_1_2*c_1_2))
sigma_x_l2 = a_2_2*(1-(b_2_2*c_2_2))
sigma_x_l = np.concatenate((sigma_x_l1,sigma_x_l2))

# fits

A_m = 400
B_m = (sigma_x_m/m200_1)
sigma_fit_m = A_m*B_m


A_l = 40
B_l = (sigma_x_l/m200_2)
sigma_fit_l = A_l*B_l

# finding the projected number density profile

r_hist_1 = np.histogram(r_1/r200_1, bins=20, range=None, normed=False, weights=None)

r_hist_2 = np.histogram(r_2/r200_2, bins=20, range=None, normed=False, weights=None)


N_1 = r_hist_1[0]  # the array of number of galaxies within specific (r) bins 
dist_1 = r_hist_1[1]
area_1 = math.pi*((dist_1[1:2]-dist_1[0:1])**2)
sigma_num_1 = N_1/area_1

N_2 = r_hist_2[0]  # the array of number of galaxies within specific (r) bins 
dist_2 = r_hist_2[1]
area_2 = math.pi*((dist_2[1:2]-dist_2[0:1])**2)
sigma_num_2 = N_2/area_2

The programme takes in two FITS files which have the same column names and perform tasks on the data in the columns. The code works fine and I get my results.

The problem here is that since the columns of the FITS have the same name, I had to load them all separately by assigning variables to each of the columns differently and perform the tasks.

Is it possible to write a single program that will take in as the input the FITS files and perform the tasks without having to assign variables to each column separately?

Answer 1

Do you want something like:

fields = ['ra', 'dec', 'zcosmo', 'halo_r200', 'dist_to_center', 
          'halo_mass', 'rho_0', 'density']

variables1 = dict((f, tbdata1.field(f)) for f in fields)

You can then access the variables like:

variables1['ra']