Calculating the nonlinear reflection coefficient of a crystalline silicon slab

Question

The purpose of this script is to calculate the nonlinear reflection coefficient of a crystalline silicon slab.

It takes some input files (columns of data separated by whitespace), converts that data to NumPy matrices, then operates on that data via some formulas that are coded into different functions. Finally, it spits out the final results to a file, as columns separated by whitespace.

I inherited an old FORTRAN program that serves the same purpose. Even though it's completely illegible it is fast as a bat out of hell! It executes in ~0.020 seconds, while my script clocks in at around ~1.25 seconds.

Please offer feedback for improving the general quality overall. I keep this code on GitHub so any and all feedback will be very helpful for me and others.

"""
nrc.py is a python program designed to calculate the Nonlinear reflection
coefficient for silicon surfaces. It works in conjunction with the matrix
elements calculated using ABINIT, and open source ab initio software,
and TINIBA, our in-house optical calculation software.

The work codified in this software can be found in Phys.Rev.B66, 195329(2002).
"""

from math import sin, cos, radians
from scipy import constants, interpolate
from numpy import loadtxt, savetxt, column_stack, absolute, \
                  sqrt, linspace, ones, complex128

########### user input ###########
OUT = "data/nrc/"
CHI1 = "data/res/chi1"
ZZZ = "data/res/zzz"
ZXX = "data/res/zxx"
XXZ = "data/res/xxz"
XXX = "data/res/xxx"
# Angles
THETA_RAD = radians(65)
PHI_RAD = radians(30)
# Misc
ELEC_DENS = 1e-28 # electronic density and scaling factor (1e-7 * 1e-21)
ENERGIES = linspace(0.01, 12.00, 1200)

########### functions ###########
def nonlinear_reflection():
    """ calls the different math functions and returns matrix,
    which is written to file """
    onee = linspace(0.01, 12.00, 1200)
    twoe = 2 * onee
    polarization = [["p", "p"], ["p", "s"], ["s", "p"], ["s", "s"]]
    for state in polarization:
        nrc = rif_constants(onee) * absolute(fresnel_vs(state[1], twoe) *
              fresnel_sb(state[1], twoe) * ((fresnel_vs(state[0], onee) *
              fresnel_sb(state[0], onee)) ** 2) *
              reflection_components(state[0], state[1], onee, twoe)) ** 2
        nrc = column_stack((onee, nrc))
        out = OUT + "R" + state[0] + state[1]
        save_matrix(out, nrc)

def chi_one(part, energy):
    """ creates spline from real part of chi1 matrix"""
    chi1 = load_matrix(CHI1)
    interpolated = \
    interpolate.InterpolatedUnivariateSpline(ENERGIES, getattr(chi1, part))
    return interpolated(energy)

def epsilon(energy):
    """ combines splines for real and imaginary parts of chi1 """
    chi1 = chi_one("real", energy) + 1j * chi_one("imag", energy)
    linear = 1 + (4 * constants.pi * chi1)
    return linear

def wave_vector(energy):
    """ math for wave vectors """
    k = sqrt(epsilon(energy) - (sin(THETA_RAD) ** 2))
    return k

def rif_constants(energy):
    """ math for constant term """
    const = (32 * (constants.pi ** 3) * (energy ** 2)) / (ELEC_DENS *
              ((constants.c * 100) ** 3) * (cos(THETA_RAD) ** 2) *
              (constants.value("Planck constant over 2 pi in eV s") ** 2))
    return const

def electrostatic_units(energy):
    """ coefficient to convert to appropriate electrostatic units """
    complex_esu = 1j * \
           ((2 * constants.value("Rydberg constant times hc in eV")) ** 5) * \
           ((0.53e-8 / (constants.value("lattice parameter of silicon") * 100))
            ** 5) / ((2 * sqrt(3)) / ((2 * sqrt(2)) ** 2))
    factor = (complex_esu * 2.08e-15 *
        (((constants.value("lattice parameter of silicon") * 100) /
            1e-8) ** 3)) / (energy ** 3)
    return factor

def fresnel_vs(polarization, energy):
    """ math for fresnel factors from vacuum to surface """
    if polarization == "s":
        fresnel = (2 * cos(THETA_RAD)) / (cos(THETA_RAD) +
                   wave_vector(energy))
    elif polarization == "p":
        fresnel = (2 * cos(THETA_RAD)) / (epsilon(energy) *
                   cos(THETA_RAD) + wave_vector(energy))
    return fresnel

def fresnel_sb(polarization, energy):
    """ math for fresnel factors from surface to bulk. Fresnel model """
    if polarization == "s":
        fresnel = ones(1200, dtype=complex128)
        #fresnel = (2 * wave_vector(energy)) / (wave_vector(energy)
        #           + wave_vector(energy))
    elif polarization == "p":
        fresnel = 1 / epsilon(energy)
        #fresnel = (2 * wave_vector(energy)) / (epsilon(energy) *
        #wave_vector(energy) + epsilon(energy) * wave_vector(energy))
    return fresnel

def reflection_components(polar_in, polar_out, energy, twoenergy):
    """ math for different r factors. loads in different component matrices """
    zzz = electrostatic_units(energy) * load_matrix(ZZZ)
    zxx = electrostatic_units(energy) * load_matrix(ZXX)
    xxz = electrostatic_units(energy) * load_matrix(XXZ)
    xxx = electrostatic_units(energy) * load_matrix(XXX)
    if polar_in == "p" and polar_out == "p":
        r_factor = sin(THETA_RAD) * epsilon(twoenergy) * \
                (((sin(THETA_RAD) ** 2) * (epsilon(energy) ** 2) * zzz) +
                (wave_vector(energy) ** 2) * (epsilon(energy) ** 2) * zxx) \
                 + epsilon(energy) * epsilon(twoenergy) * \
                 wave_vector(energy) * wave_vector(twoenergy) * \
                 (-2 * sin(THETA_RAD) * epsilon(energy) * xxz +
                wave_vector(energy) * epsilon(energy) * xxx *
                cos(3 * PHI_RAD))
    elif polar_in == "s" and polar_out == "p":
        r_factor = sin(THETA_RAD) * epsilon(twoenergy) * zxx - \
               wave_vector(twoenergy) * epsilon(twoenergy) * \
               xxx * cos(3 * PHI_RAD)
    elif polar_in == "p" and polar_out == "s":
        r_factor = -(wave_vector(energy) ** 2) * (epsilon(energy) ** 2) * \
                     xxx * sin(3 * PHI_RAD)
    elif polar_in == "s" and polar_out == "s":
        r_factor = xxx * sin(3 * PHI_RAD)
    return r_factor

def load_matrix(in_file):
    """ loads files into matrices and extracts columns """
    real, imaginary = loadtxt(in_file, unpack=True, usecols=[1, 2])
    data = real + 1j * imaginary
    return data

def save_matrix(out_file, data):
    """ saves matrix to file """
    savetxt(out_file, data, fmt=('%5.14e'), delimiter='\t')

nonlinear_reflection()

Running a script with just the import statements takes ~0.2 seconds, so that accounts for a little time.

Here's some test data. All data is of the exact same format.

.1100 .38602E-04 -.11343E-02
.1200 .55456E-04 -.40092E-03
.1300 .51653E-04 -.81893E-03
.1400 .66445E-04 -.19650E-03
.1500 .52905E-04 -.15417E-02
.1600 .62693E-04 -.11310E-02
.1700 .69121E-04 -.10427E-02
.1800 .55286E-04 -.25198E-02
.1900 .70385E-04 -.16457E-02
.2000 .74872E-04 -.17719E-02
.2100 .83163E-04 -.15324E-02
.2200 .97154E-04 -.89845E-03
.2300 .85164E-04 -.18913E-02
.2400 .94601E-04 -.18361E-02
.2500 .97245E-04 -.19024E-02
.2600 .10928E-03 -.13463E-02
.2700 .11207E-03 -.16597E-02
.2800 .10805E-03 -.22026E-02
.2900 .11929E-03 -.17817E-02
.3000 .12704E-03 -.16619E-02
.3100 .11721E-03 -.26010E-02

Answer 1

Avoid computing the same thing more than once. For example:

• Define global constants such as SIN_THETA_RAD_SQUARED = sin(THETA_RAD) ** 2
• Load the data files just once into global constants, or into local variables in nonlinear_reflection and pass them to reflection_components as parameters.
• Assign epsilon(energy) to a local variable in reflection_components and use that in the lengthy expression.