硅片非线性反射系数的计算

Question

该脚本的目的是计算晶体硅片的非线性反射系数。

它接受一些输入文件(由空格分隔的数据列)，将该数据转换为NumPy矩阵，然后通过编码成不同函数的一些公式对该数据进行操作。最后，它将最终结果输出到一个文件中，作为由空格分隔的列。

我继承了一个用于同样目的的旧FORTRAN程序。即使它是完全难以辨认的，它就像飞出地狱的蝙蝠一样快！它在大约0.020秒内执行，而我的脚本在1.25秒左右执行。

请提供反馈意见，以提高总体质量。我将这段代码保存在GitHub上，所以任何反馈都会对我和其他人很有帮助。

"""
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()

使用导入语句运行脚本需要大约0.2秒，这样就占用了一小部分时间。

这是一些测试数据。所有数据都是完全相同的格式。

.1100 .38602E-04 -.11343E-02 .1200 .55456E-04 -.40092E-03 .1300 .51653 E-04 -.81893E-03 .1400 .66445 E-04-.19650 E-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-.177E-02-.2100.83163E04-15244-1524E-15244-15457E.02 .2200 .97154E-04 -.89845 E-03 .2300 .85164E-04 -.18913 E-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

避免重复计算同一件事情。例如：

• 定义全局常量，如SIN_THETA_RAD_SQUARED = sin(THETA_RAD) ** 2
• 将数据文件加载到全局常量或nonlinear_reflection中的局部变量中，并将它们作为参数传递给reflection_components。
• 将epsilon(energy)分配给reflection_components中的局部变量，并在冗长的表达式中使用。