77 lines
2.8 KiB
Python
77 lines
2.8 KiB
Python
from numbers import Number
|
|
|
|
import sympy as sp
|
|
|
|
x, y, z = sp.symbols("x, y, z", real=True, finite=True)
|
|
#: Beam waist in x-direction (collimated).
|
|
waist_x0 = sp.Symbol("W_{x0}", real=True, positive=True, finite=True, nonzero=True)
|
|
#: Beam waist in z-direction at the focus.
|
|
waist_z0 = sp.Symbol("W_{z0}", real=True, positive=True, finite=True, nonzero=True)
|
|
#: Offset from waist at origin
|
|
|
|
#: Power per beam. Total power is 2x.
|
|
power = sp.Symbol("P", real=True, positive=True, finite=True, nonzero=True)
|
|
#: Magnetic field gradients.
|
|
grad_r, grad_z = sp.symbols(
|
|
"\\nabla{B_x}, \\nabla{B_z}", real=True, finite=True, nonzero=True
|
|
)
|
|
#: Real part of polarizability :math:`\alpha`
|
|
a = sp.Symbol("a", real=True, finite=True, nonzero=True)
|
|
d = sp.Symbol("d", real=True, positive=True, finite=True, nonzero=True)
|
|
mu_b = sp.Symbol("mu_b", real=True, positive=True, finite=True, nonzero=True)
|
|
#: Wavelength
|
|
wvl = sp.Symbol("lambda", real=True, positive=True, finite=True, nonzero=True)
|
|
#: Angle of the incident beams
|
|
theta = sp.Symbol("theta", real=True, positive=True, finite=True, nonzero=True)
|
|
|
|
|
|
def BeamWaistZ(length: sp.Symbol | Number) -> sp.Expr: # noqa: N802
|
|
"""The beam is focused in z-direction, leading to a waist that
|
|
depends on the distance :math:`l` from the lens.
|
|
|
|
Depending on which beam (upper or lower), the distance from the
|
|
lens :math:`l` is given by
|
|
:math:`y \\cos(\\theta) \\mp z \\sin(\\theta)`
|
|
|
|
:param length: Distance :math:`l` from the lens.
|
|
"""
|
|
return waist_z0 * sp.sqrt(1 + (wvl * length / (sp.pi * waist_z0**2)) ** 2)
|
|
|
|
|
|
def UpperBeamIntensity() -> sp.Expr: # noqa: N802
|
|
"""Intensity distribution of the upper beam"""
|
|
length = y * sp.cos(theta) - z * sp.sin(theta)
|
|
waist_z = BeamWaistZ(length)
|
|
# Intensity of the beam at the origin
|
|
i_0 = 2 * power / (sp.pi * waist_x0 * waist_z0)
|
|
# Turn off automatic formatting to preserve readability
|
|
# fmt: off
|
|
return i_0 * waist_z0 / waist_z * sp.exp(
|
|
- 2 * x**2 / waist_x0**2
|
|
- 2 * (z * sp.cos(theta) + y * sp.sin(theta)) ** 2 / waist_z**2
|
|
) # fmt: on
|
|
|
|
|
|
def LowerBeamIntensity() -> sp.Expr: # noqa: N802
|
|
"""Intensity distribution of the lower beam"""
|
|
length = y * sp.cos(theta) + z * sp.sin(theta)
|
|
waist_z = BeamWaistZ(length)
|
|
# Intensity of the beam at the origin
|
|
i_0 = 2 * power / (sp.pi * waist_x0 * waist_z0)
|
|
# Turn off automatic formatting to preserve readability
|
|
# fmt: off
|
|
return i_0 * waist_z0 / waist_z * sp.exp(
|
|
- 2 * x**2 / waist_x0**2
|
|
- 2 * (z * sp.cos(theta) - y * sp.sin(theta)) ** 2 / waist_z**2
|
|
) # fmt: on
|
|
|
|
|
|
def TotalIntensity() -> sp.Expr: # noqa: N802
|
|
i_1 = UpperBeamIntensity()
|
|
i_2 = LowerBeamIntensity()
|
|
return i_1 + i_2 + 2 * sp.sqrt(i_1 * i_2) * sp.cos(2 * sp.pi * z / d)
|
|
|
|
|
|
def TotalPotential() -> sp.Expr: # noqa: N802
|
|
return -a * TotalIntensity()
|