Skip to content

Instantly share code, notes, and snippets.

@xkstein

xkstein/integrate_gaussian_sympy.md

Last active November 20, 2023 20:09
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save xkstein/112d43f0aad873099a8da5a8c0f54dce to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save xkstein/112d43f0aad873099a8da5a8c0f54dce to your computer and use it in GitHub Desktop.
Download ZIP
Gaussian integrals over all space in Sympy
Raw
integrate_gaussian_sympy.md

Symbolically integrate gaussians over all space!

Integrals that look like:

$$A \int_{-\infty}^{+\infty} x^{2n} e^{-x^2/a^2} dx$$

are really common in quantum mechanics (at least at the upper-division undergraduate level). Unfortunately, trying to perform one of these integrals using the built-in sympy integration function gives confusing results that are difficult to use.

This code produces more practical results by using the simple analytical solution for these integrals (shown in the back of Griffiths QM):

$$\int_{-\infty}^{+\infty} x^{2n} e^{-x^2/a^2} dx = 2 \sqrt{\pi} \frac{ (2n)! }{ n! } \Big( \frac{a}{2} \Big)^{2n + 1}$$

It identifies $n$ and $a$ by pattern matching and plugs into the analytical solution. This makes nice and clean symbolic solutions which are easily checkable by hand.

Note: Any minor deviation from this form, and this script won't work at all. It does, however, support coefficients that have a sum with different orders of x.

Raw
sympy_integrate_gaussian.py
from sympy import *
def check_gaussian(expr, x):
'''Internal (private) function that checks if a function is a gaussian
returns: True or False
'''
p = Wild('outer_coefficients')
a = Wild('alpha_gaussian_spacial_variable')
matches = expr.match( p * exp( -( x / a ) ** 2 ) )
try:
return not matches[a].has(x)
except TypeError:
return False
def split_by_sum(expr):
'''Internal (private) function that splits an expression into its sum terms
returns: list of terms in expr, split by sum
'''
mul_args = []
add_args = []
for arg in expr.args:
if arg.func == Add:
add_args.append(arg)
continue
mul_args.append(arg)
multiplied_args = Mul(*mul_args)
if not add_args:
return expr
if len(add_args) > 1:
return NotImplementedError('Sorry, can\'t handle polynomial args, easy fix, but I haven\'t')
return [ multiplied_args * addition_arg for addition_arg in add_args[0].args ]
def integrate_gaussian(expr, x):
'''Evaluates a basic gaussian from -inf to inf
The gausians that this expects are in the form: c * exp( -x**2 / a**2 )
This first finds what c and a are, then checks if c contains odd
order x--in this case the integral evaluates to 0. It checks to see
if there is a sum in the c term, and integrates those separately.
Then it uses the basic gaussian solution forms (found in the back of Griffiths EM)
Note: Don't use `outer_coefficients` or `alpha_gaussian_spacial_variable` as
variable names in your equations, doesn't seem to always work
'''
if not check_gaussian(expr, x):
raise TypeError('Not a gaussian')
p = Wild('outer_coefficients')
a = Wild('alpha_gaussian_spacial_variable')
matches = expr.match( p * exp( -( x / a ) ** 2 ) )
coeff = matches[p]
alpha = matches[a]
sum_split_coeffs = [ coeff ]
# Check if there are addition terms which must be delt with
if any([ type(arg) == Add for arg in coeff.args ]):
sum_split_coeffs = split_by_sum(coeff)
evaluated_terms = []
for coefficient in sum_split_coeffs:
non_x_coeff, pow_x = coefficient.as_coeff_exponent(x)
if pow_x % 2 != 0:
evaluated_terms.append(0)
continue
n = pow_x / 2
evaluated_terms.append( 2 * non_x_coeff * sqrt(pi) * ( factorial( 2 * n ) / factorial( n ) ) * ( alpha / 2 ) ** ( 2 * n + 1 ) )
return Add(*evaluated_terms)
if __name__ == '__main__':
var('a x')
assert ( integrate_gaussian(x ** 4 * exp(-(x ** 2) / a ** 2), x) - (3 / 4) * sqrt(pi) * a ** 5 ) == 0
print(integrate_gaussian(x ** 4 * exp(-(x ** 2) / a ** 2), x))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment