Source code for pylcp.fields

import numpy as np
from inspect import signature
from pylcp.common import cart2spherical, spherical2cart
from .integration_tools import parallelIntegrator
from scipy.spatial.transform import Rotation

import numba

@numba.njit
def dot2D(a, b):
    c = np.zeros((a.shape[1],), dtype=a.dtype)
    for ii in range(a.shape[1]):
        c[ii] = np.sum(a[:, ii]*b[:, ii])
    return c

@numba.njit
def electric_field(r, t, amp, pol, k, phase):
    return pol*amp*np.exp(-1j*(k[0]*r[0]+k[1]*r[1]+k[2]*r[2]) + 1j*phase)


def return_constant_val(R, t, val):
    if R.shape==(3,):
        return val
    elif R.shape[0] == 3:
        return val*np.ones(R[0].shape)
    else:
        raise ValueError('The first dimension of R should have length 3, ' +
                         'not %d.'%R.shape[0])

def return_constant_vector(R, t, vector):
    if R.shape==(3,):
        return vector
    elif R.shape[0] == 3:
        return np.outer(vector, np.ones(R[0].shape))
    else:
        raise ValueError('The first dimension of R should have length 3, ' +
                         'not %d.'% R.shape[0])

def return_constant_val_t(t, val):
    if isinstance(t, np.ndarray):
        return val*np.ones(t.shape)
    else:
        return val

def promote_to_lambda(val, var_name='', type='Rt'):
    """
    Promotes a constant or callable to a lambda function with proper arguments.

    Parameters
    ----------
        val : array_like or callable
            The value to promote.  Can either be a function, array_like
            (vector), or a scalar (constant).
        var_name : str, optional
            Name of the variable attempting to be promoted.  Useful for error
            messages. Default: empty string.
        type : str, optional
            The arguments of the lambda function we are creating.  If `Rt`,
            the lambda function returned has ``(R, t)`` as its arguments.  If
            `t`, it has only `t` as its arguments.  Default: `Rt`.


    Returns
    -------
        func : callable
            lambda function created by this function
        sig : string
            Either `(R,t)` or `(t)`
    """
    if type == 'Rt':
        if not callable(val):
            if isinstance(val, list) or isinstance(val, np.ndarray):
                func = lambda R=np.array([0., 0., 0.]), t=0.: return_constant_vector(R, t, val)
            else:
                func = lambda R=np.array([0., 0., 0.]), t=0.: return_constant_val(R, t, val)
            sig = '()'
        else:
            sig = str(signature(val))
            if ('(R)' in sig or '(r)' in sig or '(x)' in sig):
                func = lambda R=np.array([0., 0., 0.]), t=0.: val(R)
                sig = '(R)'
            elif ('(R, t)' in sig or '(r, t)' in sig or '(x, t)' in sig):
                func = lambda R=np.array([0., 0., 0.]), t=0.: val(R, t)
                sig = '(R, t)'
            elif '(t)' in sig:
                func = lambda R=np.array([0., 0., 0.]), t=0.: val(t)
                sig = '(R, t)'
            else:
                raise TypeError('Signature [%s] of function %s not'+
                                'understood.'% (sig, var_name))

        return func, sig
    elif type == 't':
        if not callable(val):
            func = lambda t=0.: return_constant_val_t(t, val)
            sig = '()'
        else:
            sig = str(signature(val))
            if '(t)' in sig:
                func = lambda t=0.: val(t)
            else:
                raise TypeError('Signature [%s] of function %s not '+
                                'understood.'% (sig, var_name))

        return func, sig

def return_dx_dy_dz(R, eps):
    if R.shape == (3,):
        dx = np.array([eps, 0., 0.])
        dy = np.array([0., eps, 0.])
        dz = np.array([0., 0., eps])
    else:
        dx = np.zeros(R.shape)
        dy = np.zeros(R.shape)
        dz = np.zeros(R.shape)

        dx[0] = eps
        dy[1] = eps
        dz[2] = eps

    return dx, dy, dz


[docs]class magField(object): """ Base magnetic field class Stores a magnetic defined magnetic field and calculates useful derivatives for `pylcp`. Parameters ---------- field : array_like with shape (3,) or callable If constant, the magnetic field vector, specified as either as an array_like with shape (3,). If a callable, it must have a signature like (R, t), (R), or (t) where R is an array_like with shape (3,) and t is a float and it must return an array_like with three elements. eps : float, optional Small distance to use in calculation of numerical derivatives. By default `eps=1e-5`. Attributes ---------- eps : float small epsilon used for computing derivatives """ def __init__(self, field, eps=1e-5): self.eps = eps R = np.random.rand(3) # Pick a random point for testing # Promote it to a lambda func: self.Field, self.FieldSig = promote_to_lambda(field, var_name='for field') # Try it out: response = self.Field(R, 0.) if (isinstance(response, float) or isinstance(response, int) or len(response) != 3): raise ValueError('Magnetic field function must return a vector.')
[docs] def FieldMag(self, R=np.array([0., 0., 0.]), t=0): """ Magnetic field magnitude at R and t: Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- B : float the magnetic field mangitude at position R and time t. """ return np.linalg.norm(self.Field(R, t))
[docs] def gradFieldMag(self, R=np.array([0., 0., 0.]), t=0): """ Gradient of the magnetic field magnitude at R and t: Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- dB : array_like, shape (3,) :math:`\\nabla|B|`, the gradient of the magnetic field magnitude at position :math:`R` and time :math:`t`. """ dx, dy, dz = return_dx_dy_dz(R, self.eps) return np.array([ (self.FieldMag(R+dx, t)-self.FieldMag(R-dx, t))/2/self.eps, (self.FieldMag(R+dy, t)-self.FieldMag(R-dy, t))/2/self.eps, (self.FieldMag(R+dz, t)-self.FieldMag(R-dz, t))/2/self.eps ])
[docs] def gradField(self, R=np.array([0., 0., 0.]), t=0): """ Full spaitial derivative of the magnetic field at R and t: Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- dB : array_like, shape (3, 3) the full gradient of the magnetic field, with elements .. math:: \\begin{pmatrix} \\frac{dB_x}{dx} & \\frac{dB_y}{dx} & \\frac{dB_z}{dx} \\\\ \\frac{dB_x}{dy} & \\frac{dB_y}{dy} & \\frac{dB_z}{dy} \\\\ \\frac{dB_x}{dz} & \\frac{dB_y}{dz} & \\frac{dB_z}{dz} \\\\ \\end{pmatrix} Notes ----- This method calculates the derivative stupidly, just using first order numerical differentiation using the `eps` parameter. """ dx, dy, dz = return_dx_dy_dz(R, self.eps) return np.array([ (self.Field(R+dx, t) - self.Field(R-dx, t))/2/self.eps, (self.Field(R+dy, t) - self.Field(R-dy, t))/2/self.eps, (self.Field(R+dz, t) - self.Field(R-dz, t))/2/self.eps ])
[docs]class iPMagneticField(magField): """ Ioffe-Pritchard trap magnetic field Generates a magnetic field of the form .. math:: \mathbf{B} = B_1 x \\hat{x} - B_1 y \\hat{y} + \\left(B_0 + \\frac{B_2}{2}z^2\\right)\\hat{z} Parameters ---------- B0 : float Constant offset field B1 : float Magnetic field gradient in x-y plane B2 : float Magnetic quadratic component along z direction. Notes ----- It is currently missing extra terms that are required for it to fulfill Maxwell's equations at second order. """ def __init__(self, B0, B1, B2, eps = 1e-5): super().__init__(lambda R, t: np.array([B1*R[0]-B2*R[0]*R[2]/2, -R[1]*B1-B2*R[1]*R[2]/2, B0+B2/2*(R[2]**2 - (R[0]**2+R[1]**2)/2)])) self.B0 = B0 self.B1 = B1 self.B2 = B2 #Analytical form, not numerical for this and gradField
[docs] def gradFieldMag(self, R=np.array([0., 0., 0.]), t=0): a = self.B0 b = self.B1 c = self.B2 x = R[0] y = R[1] z = R[2] mag = self.FieldMag(R, t) xcom = 0.5*(2*b**2*x-a*c*x+(c**2)*(x**3)/4+(c**2)*(y**2)*x/4-2*b*c*z*x)/mag ycom = 0.5*(2*b**2*(y)-a*c*y+(c**2)*(x**2)*y/4 + (c**2)*(y**3)/4+2*b*c*z*y)/mag zcom = 0.5*(0-b*c*(x**2)+b*c*(y**2)+2*a*c*z+(c**2)*(z**3))/mag return np.array([xcom, ycom, zcom])
[docs] def gradField(self, R=np.array([0., 0., 0.]), t=0): B0 = self.B0 B1 = self.B1 B2 = self.B2 x = R[0] y = R[1] z = R[2] xcom = np.array([B1-B2*z/2, 0, -B2*x/2]) ycom = np.array([0, -B1-B2*z/2, B2*y/2]) zcom = np.array([-B2*x/2, -B2*y/2, B2*z]) return np.array([ np.array([B1-B2*z/2, 0, -B2*x/2]), np.array([0, -B1-B2*z/2, B2*y/2]), np.array([-B2*x/2, -B2*y/2, B2*z]) ])
[docs]class constantMagneticField(magField): """ Spatially constant magnetic field Represents a magnetic field of the form .. math:: \\mathbf{B} = \mathbf{B}_0 Parameters ---------- val : array_like with shape (3,) The three-vector defintion of the constant magnetic field. """ def __init__(self, B0): super().__init__(lambda R, t: B0) self.constant_grad_field_mag = np.zeros((3,)) self.constant_grad_field = np.zeros((3,3))
[docs] def gradFieldMag(self, R=np.array([0., 0., 0.]), t=0): """ Gradient of the magnetic field magnitude at R and t: Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- dB : np.zeros((3,)) The gradient of a constant magnetic field magnitude is always zero. """ return self.constant_grad_field_mag
[docs] def gradField(self, R=np.array([0., 0., 0.]), t=0): """ Gradient of the magnetic field magnitude at R and t: Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- dB : np.zeros((3,)) :math:`\\nabla|B|=0`, the gradient of the magnitude of a constant magnetic field is always zero. """ return self.constant_grad_field
[docs]class quadrupoleMagneticField(magField): """ Spherical quadrupole magnetic field Represents a magnetic field of the form .. math:: \\mathbf{B} = \\alpha\\left(- \\frac{x\\hat{x}}{2} - \\frac{y\\hat{y}}{2} + z\\hat{z}\\right) Parameters ---------- alpha : float strength of the magnetic field gradient. """ def __init__(self, alpha, eps=1e-5): super().__init__(lambda R, t: alpha*np.array([-0.5*R[0], -0.5*R[1], R[2]])) self.alpha = alpha self.constant_grad_field = alpha*\ np.array([[-0.5, 0., 0.], [0., -0.5, 0.], [0., 0., 1.]])
[docs] def gradField(self, R=np.array([0., 0., 0.]), t=0): """ Full spaitial derivative of the magnetic field at R and t: Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- dB : array_like, shape (3, 3) the full gradient of the magnetic field, with elements .. math:: \\begin{pmatrix} -\\alpha/2 & 0 & 0 \\\\ 0 & -\\alpha/2 & 0 \\\\ 0 & 0 & \\alpha \\\\ \\end{pmatrix} """ return self.constant_grad_field
# First, define the laser beam class:
[docs]class laserBeam(object): """ The base class for a single laser beam Attempts to represent a laser beam as .. math:: \\frac{1}{2}\\hat{\\boldsymbol{\\epsilon}}(r, t) E_0(r, t) e^{i\\mathbf{k}(r,t)\\cdot\\mathbf{r}-i \\int dt\\Delta(t) + i\\phi(r, t)} where :math:`\\hat{\\boldsymbol{\\epsilon}}` is the polarization, :math:`E_0` is the electric field magnitude, :math:`\\mathbf{k}(r,t)` is the k-vector, :math:`\\mathbf{r}` is the position, :math:`\\Delta(t)` is the deutning, :math:`t` is the time, and :math:`\\phi` is the phase. Parameters ---------- kvec : array_like with shape (3,) or callable The k-vector of the laser beam, specified as either a three-element list or numpy array or as callable function. If a callable, it must have a signature like (R, t), (R), or (t) where R is an array_like with shape (3,) and t is a float and it must return an array_like with three elements. pol : int, float, array_like with shape (3,), or callable The polarization of the laser beam, specified as either an integer, float array_like with shape(3,), or as callable function. If an integer or float, if `pol<0` the polarization will be left circular polarized relative to the k-vector of the light. If `pol>0`, the polarization will be right circular polarized. If array_like, polarization will be specified by the vector, whose basis is specified by `pol_coord`. If a callable, it must have a signature like (R, t), (R), or (t) where R is an array_like with shape (3,) and t is a float and it must return an array_like with three elements. s : float or callable The intensity of the laser beam, normalized to the saturation intensity, specified as either a float or as callable function. If a callable, it must have a signature like (R, t), (R), or (t) where R is an array_like with shape (3,) and t is a float and it must return a float. delta: float or callable Detuning of the laser beam. If a callable, it must have a signature like (t) where t is a float and it must return a float. phase : float, optional Phase of laser beam. By default, zero. pol_coord : string, optional Polarization basis of the input polarization vector: 'cartesian' or 'spherical' (default). eps : float, optional Small distance to use in calculation of numerical derivatives. By default `eps=1e-5`. Attributes ---------- eps : float Small epsilon used for computing derivatives phase : float Overall phase of the laser beam. """ def __init__(self, kvec=None, s=None, pol=None, delta=None, phase=0., pol_coord='spherical', eps=1e-5): # Promote it to a lambda func: if not kvec is None: self.kvec, self.kvec_sig = promote_to_lambda(kvec, var_name='kvector') # Promote it to a lambda func: if not s is None: self.intensity, self.intensity_sig = promote_to_lambda(s, var_name='s') if not pol is None: if not callable(pol): pol = self.__parse_constant_polarization(pol, pol_coord) # Now, promote! self.pol, self.pol_sig = promote_to_lambda(pol, var_name='polarization') # Promote it to a lambda func: if not delta is None: self.delta, self.delta_sig = promote_to_lambda(delta, var_name='delta', type='t') if self.delta_sig == '(t)': self.delta_phase = parallelIntegrator(self.delta) elif self.delta_sig == '()': self.delta_phase = lambda t: delta*t # Promote it to a lambda func: if not phase is None: self.phase, self.phase_sig = promote_to_lambda(phase, var_name='phase', type='t') self.eps = eps def __parse_constant_polarization(self, pol, pol_coord): if isinstance(pol, float) or isinstance(pol, int): # If the polarization is defined by just a single number (+/-1), # we assume that the polarization is defined as sigma^+ or sigma^- # using the k-vector of the light as the axis defining z. In this # case, we want to project onto the actual z axis, which is # relatively simple as there is only one angle. # Set the polarization in this direction: if np.sign(pol)<0: self.pol = np.array([1., 0., 0.], dtype='complex') else: self.pol = np.array([0., 0., 1.], dtype='complex') # Promote to lambda: self.pol, self.pol_sig = promote_to_lambda(self.pol, var_name='polarization') # Project onto the actual k-vector: self.pol = self.project_pol(self.kvec()/np.linalg.norm(self.kvec()), invert=True).astype('complex128') elif isinstance(pol, np.ndarray): if pol.shape != (3,): raise ValueError("pol, when a vector, must be a (3,) array") # The user has specified a single polarization vector in # cartesian coordinates: if pol_coord=='cartesian': # Check for transverseness: if np.abs(np.dot(self.kvec(), pol)) > 1e-9: raise ValueError("I'm sorry; light is a transverse wave") # Always store in spherical basis. Matches Brown and Carrington. self.pol = cart2spherical(pol).astype('complex128') # The user has specified a single polarization vector in # spherical basis: elif pol_coord=='spherical': pol_cart = spherical2cart(pol) # Check for transverseness: if np.abs(np.dot(self.kvec(), pol_cart)) > 1e-9: raise ValueError("I'm sorry; light is a transverse wave") # Save the variable: self.pol = pol.astype('complex128') # Finally, normalize self.pol = self.pol/np.linalg.norm(self.pol) else: raise ValueError("pol must be +1, -1, or a numpy array") return self.pol
[docs] def kvec(self, R=np.array([0., 0., 0.]), t=0.): """ Returns the k-vector of the laser beam Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- kvec : array_like, size(3,) the k vector at position R and time t. """ pass
[docs] def intensity(self, R=np.array([0., 0., 0.]), t=0.): """ Returns the intensity of the laser beam at position R and t Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- s : float or array_like Saturation parameter of the laser beam at R and t. """ pass
[docs] def pol(self, R=np.array([0., 0., 0.]), t=0.): """ Returns the polarization of the laser beam at position R and t The polarization is returned in the spherical basis. Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- pol : array_like, size (3,) polarization of the laser beam at R and t in spherical basis. """ pass
[docs] def delta(self, t=0.): """ Returns the detuning of the laser beam at time t Parameters ---------- t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- delta : float or array like detuning of the laser beam at time t """ pass
# TODO: add testing of kvec/pol orthogonality.
[docs] def project_pol(self, quant_axis, R=np.array([0., 0., 0.]), t=0, treat_nans=False, calculate_norm=False, invert=False): """ Project the polarization onto a quantization axis. Parameters ---------- quant_axis : array_like, shape (3,) A normalized 3-vector of the quantization axis direction. R : array_like, shape (3,), optional If polarization is a function of R is the 3-vectors at which the polarization shall be calculated. calculate_norm : bool, optional If true, renormalizes the quant_axis. By default, False. treat_nans : bool, optional If true, every place that nan is encoutnered, replace with the $hat{z}$ axis as the quantization axis. By default, False. invert : bool, optional If true, invert the process to project the quantization axis onto the specified polarization. Returns ------- projected_pol : array_like, shape (3,) The polarization projected onto the quantization axis. """ # First, return the polarization at the desired R and t. pol = self.pol(R, t) # Second, check the quanitization axis if specified by the user. The fun # thing here is that we may only need to do this once, since it should # overwrite the original variable with the improper quant_axis with the # proper one. if calculate_norm: quant_axis2 = np.zeros(quant_axis.shape) quant_axis[2] = 1.0 # Make the third entry all ones. quant_axis_norm = np.linalg.norm(quant_axis, axis=0) for ii in range(3): quant_axis2[ii][quant_axis_norm!=0] = \ quant_axis2[ii][quant_axis_norm!=0]/\ quant_axis_norm[quant_axis_norm!=0] quant_axis=quant_axis2 elif treat_nans: for ii in range(quant_axis.shape[0]): if ii<quant_axis.shape[0]-1: quant_axis[ii][np.isnan(quant_axis[-1])] = 0.0 else: quant_axis[ii][np.isnan(quant_axis[-1])] = 1.0 # To project the full three-vector, we want to determine the Euler # angles alpha, beta, and gamma that rotate the z-axis into the # quantization axis. The final Euler angle, gamma, only sets the # phase of the -1 to +1 component, so it does not have any physical # meaning for the rate equations (where we expect this method to # mostly be used.) Thus, we just set that angle equal to zero here. cosbeta = quant_axis[2] sinbeta = np.sqrt(1-cosbeta**2) if isinstance(cosbeta, (float, int)): if np.abs(cosbeta)<1: gamma = np.arctan2(quant_axis[1], quant_axis[0]) else: gamma = 0 alpha = 0 else: gamma = np.zeros(cosbeta.shape) inds = np.abs(quant_axis[2])<1 gamma[inds] = np.arctan2(quant_axis[1][inds], quant_axis[0][inds]) alpha = np.zeros(cosbeta.shape) quant_axis = quant_axis.astype('float64') pol = pol.astype('complex128') D = np.array([ [(1+cosbeta)/2*np.exp(-1j*alpha + 1j*gamma), -sinbeta/np.sqrt(2)*np.exp(-1j*alpha), (1-cosbeta)/2*np.exp(-1j*alpha - 1j*gamma)], [sinbeta/np.sqrt(2)*np.exp(1j*gamma), cosbeta, -sinbeta/np.sqrt(2)*np.exp(-1j*gamma)], [(1-cosbeta)/2*np.exp(1j*alpha+1j*gamma), sinbeta/np.sqrt(2), (1+cosbeta)/2*np.exp(1j*alpha-1j*gamma)] ]) if invert: D = np.linalg.inv(D) # Tensordot is a cool function, it allow you to do # multiplication-sums across various axes. Because D is a # 3\times3\times\cdots array and likewise pol is a 3\times\cdots # array, we matrix multiply against the 1st dimension of D and the # 0th dimension of pol. # TODO: This probably won't work in the case pol.shape=3\times\cdots # case if pol.shape == (3,) and quant_axis.shape == (3,): return D @ pol else: return np.tensordot(D, pol, ([1],[0]))
[docs] def cartesian_pol(self, R=np.array([0., 0., 0.]), t=0): """ Returns the polarization in Cartesian coordinates. Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- pol : array_like, size (3,) polarization of the laser beam at R and t in Cartesian basis. """ pol = self.pol(R, t) return spherical2cart(pol)
[docs] def jones_vector(self, xp, yp, R=np.array([0., 0., 0.]), t=0): """ Returns the Jones vector at position Parameters ---------- xp : array_like, shape (3,) The x vector of the basis in which to calculate the Jones vector. Must be orthogonal to k. yp : array_like, shape (3,) The y vector of the basis in which to calculate the Jones vector. Must be orthogonal to k and `xp`. R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- pol : array_like, size (2,) Jones vector of the laser beam at R and t in Cartesian basis. """ # First, run some basic checks. if np.abs(np.dot(xp, yp)) > 1e-10: raise ValueError('xp and yp must be orthogonal.') if np.abs(np.dot(xp, self.kvec(R, t))) > 1e-10: raise ValueError('xp and k must be orthogonal.') if np.abs(np.dot(yp, self.kvec(R, t))) > 1e-10: raise ValueError('yp and k must be orthogonal.') if np.sum(np.abs(np.cross(xp, yp) - self.kvec(R, t))) > 1e-10: raise ValueError('xp, yp, and k must form a right-handed' + 'coordinate system.') pol_cart = self.cartesian_pol(R, t) if np.abs(np.dot(pol_cart, self.kvec(R, t))) > 1e-9: raise ValueError('Something is terribly, terribly wrong.') return np.array([np.dot(pol_cart, xp), np.dot(pol_cart, yp)])
[docs] def stokes_parameters(self, xp, yp, R=np.array([0., 0., 0.]), t=0): """ The Stokes Parameters of the laser beam at R and t Parameters ---------- xp : array_like, shape (3,) The x vector of the basis in which to calculate the Stokes parameters. Must be orthogonal to k. yp : array_like, shape (3,) The y vector of the basis in which to calculate the Stokes parameters. Must be orthogonal to k and `xp`. R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- pol : array_like, shape (3,) Stokes parameters for the laser beam, [Q, U, V] """ jones_vector = self.jones_vector(xp, yp, R, t) Q = np.abs(jones_vector[0])**2 - np.abs(jones_vector[1])**2 U = 2*np.real(jones_vector[0]*np.conj(jones_vector[1])) V = -2*np.imag(jones_vector[0]*np.conj(jones_vector[1])) return (Q, U, V)
[docs] def polarization_ellipse(self, xp, yp, R=np.array([0., 0., 0.]), t=0): """ The polarization ellipse parameters of the laser beam at R and t Parameters ---------- xp : array_like, shape (3,) The x vector of the basis in which to calculate the polarization ellipse. Must be orthogonal to k. yp : array_like, shape (3,) The y vector of the basis in which to calculate the polarization ellipse. Must be orthogonal to k and `xp`. R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- psi : float :math:`\\psi` parameter of the polarization ellipse chi : float :math:`\\chi` parameter of the polarization ellipse """ Q, U, V = self.stokes_parameters(xp, yp, R, t) psi = np.arctan2(U, Q) while psi<0: psi+=2*np.pi psi = psi%(2*np.pi)/2 if np.sqrt(Q**2+U**2)>1e-10: chi = 0.5*np.arctan(V/np.sqrt(Q**2+U**2)) else: chi = np.pi/4*np.sign(V) return (psi, chi)
[docs] def electric_field(self, R, t): """ The electric field at position R and t Parameters ---------- R : array_like, size (3,) vector of the position at which to return the kvector. By default, the origin. t : float time at which to return the k-vector. By default, t=0. Returns ------- Eq : array_like, shape (3,) electric field in the spherical basis. """ kvec = self.kvec(R, t) s = self.intensity(R, t) pol = self.pol(R, t) delta_phase = self.delta_phase(t) phase = self.phase(t) amp = np.sqrt(2*s) if isinstance(t, float): Eq = electric_field(R, t, amp, pol, kvec, delta_phase - phase) else: Eq = pol.reshape(3, t.size)*\ (amp*np.exp(-1j*dot2D(kvec, R) + 1j*delta_phase - 1j*phase)).reshape(1, t.size) return Eq
[docs] def electric_field_gradient(self, R, t): """ The full derivative of electric field at position R and t Parameters ---------- R : array_like, size (3,) vector of the position at which to return the kvector. By default, the origin. t : float time at which to return the k-vector. By default, t=0. Returns ------- dEq : array_like, shape (3, 3) The full gradient of the electric field, in spherical coordinates. .. math:: \\begin{pmatrix} \\frac{dE_{-1}}{dx} & \\frac{dE_0}{dx} & \\frac{dE_{+1}}{dx} \\\\ \\frac{dE_{-1}}{dy} & \\frac{dE_0}{dy} & \\frac{dE_{+1}}{dy} \\\\ \\frac{dE_{-1}}{dz} & \\frac{dE_0}{dz} & \\frac{dE_{+1}}{dz} \\\\ \\end{pmatrix} """ (dx, dy, dz) = return_dx_dy_dz(R, self.eps) delEq = np.array([ (self.electric_field(R+dx, t) - self.electric_field(R-dx, t))/2/self.eps, (self.electric_field(R+dy, t) - self.electric_field(R-dy, t))/2/self.eps, (self.electric_field(R+dz, t) - self.electric_field(R-dz, t))/2/self.eps ]) return delEq
[docs]class infinitePlaneWaveBeam(laserBeam): """ Infinte plane wave beam A beam which has spatially constant intensity, k-vector, and polarization. .. math:: \\frac{1}{2}\\hat{\\boldsymbol{\\epsilon}} E_0e^{i\\mathbf{k}\\cdot\\mathbf{r}-i \\int dt\\Delta(t) + i\\phi(r, t)} where :math:`\\hat{\\boldsymbol{\\epsilon}}` is the polarization, :math:`E_0` is the electric field magnitude, :math:`\\mathbf{k}(r,t)` is the k-vector, :math:`\\mathbf{r}` is the position, :math:`\\Delta(t)` is the deutning, :math:`t` is the time, and :math:`\\phi` is the phase. Parameters ---------- kvec : array_like with shape (3,) or callable The k-vector of the laser beam, specified as either a three-element list or numpy array. pol : int, float, array_like with shape (3,), or callable The polarization of the laser beam, specified as either an integer, float array_like with shape(3,). If an integer or float, if `pol<0` the polarization will be left circular polarized relative to the k-vector of the light. If `pol>0`, the polarization will be right circular polarized. If array_like, polarization will be specified by the vector, whose basis is specified by `pol_coord`. s : float or callable The intensity of the laser beam, specified as either a float or as callable function. delta: float or callable Detuning of the laser beam. If a callable, it must have a signature like (t) where t is a float and it must return a float. **kwargs : Additional keyword arguments to pass to laserBeam superclass. Notes ----- This implementation is much faster, when it can be used, compared to the base laserBeam class. """ def __init__(self, kvec, pol, s, delta, **kwargs): if callable(kvec): raise TypeError('kvec cannot be a function for an infinite plane wave.') if callable(s): raise TypeError('s cannot be a function for an infinite plane wave.') if callable(pol): raise TypeError('Polarization cannot be a function for an infinite plane wave.') # Use the super class to define the functions kvec, s, pol, and delta. super().__init__(kvec=kvec, s=s, pol=pol, delta=delta, **kwargs) # Save the constant values (might be useful): self.con_kvec = kvec self.con_s = s self.con_pol = self.pol(np.array([0., 0., 0.]), 0.) # Define attributes to speed up gradient calculation: self.amp = np.sqrt(2*self.con_s) self.dEq_prefactor = (-1j*self.amp*self.con_kvec.reshape(3, 1)* self.con_pol.reshape(1, 3))
[docs] def electric_field_gradient(self, R, t): # With a plane wave, this is simple: delta_phase = self.delta_phase(t) phase = self.phase(t) if isinstance(t, float) or (isinstance(t, np.ndarray) and t.size==1): delEq = self.dEq_prefactor*\ np.exp(-1j*np.dot(self.con_kvec, R) + 1j*delta_phase - 1j*phase) else: delEq = self.dEq_prefactor.reshape(3, 3, 1)*\ np.exp(-1j*np.dot(self.con_kvec, R) + 1j*delta_phase -1j*phase).reshape(1, 1, t.size) return delEq
[docs]class gaussianBeam(laserBeam): """ Collimated Gaussian beam A beam which has spatially constant k-vector and polarization, with a Gaussian intensity modulation. Specifically, .. math:: \\frac{1}{2}\\hat{\\boldsymbol{\\epsilon}} E_0 e^{-\\mathbf{r}^2/w_b^2} e^{i\\mathbf{k}\\cdot\\mathbf{r}-i \\int dt\\Delta(t) + i\\phi(r, t)} where :math:`\\hat{\\boldsymbol{\\epsilon}}` is the polarization, :math:`E_0` is the electric field magnitude, :math:`\\mathbf{k}(r,t)` is the k-vector, :math:`\\mathbf{r}` is the position, :math:`\\Delta(t)` is the deutning, :math:`t` is the time, and :math:`\\phi` is the phase. Note that because :math:`I\\propto E^2`, :math:`w_b` is the :math:`1/e^2` radius. Parameters ---------- kvec : array_like with shape (3,) or callable The k-vector of the laser beam, specified as either a three-element list or numpy array. pol : int, float, array_like with shape (3,), or callable The polarization of the laser beam, specified as either an integer, float array_like with shape(3,). If an integer or float, if `pol<0` the polarization will be left circular polarized relative to the k-vector of the light. If `pol>0`, the polarization will be right circular polarized. If array_like, polarization will be specified by the vector, whose basis is specified by `pol_coord`. s : float or callable The maximum intensity of the laser beam at the center, specified as either a float or as callable function. delta : float or callable Detuning of the laser beam. If a callable, it must have a signature like (t) where t is a float and it must return a float. wb : float The :math:`1/e^2` radius of the beam. **kwargs: Additional keyword arguments to pass to the laserBeam superclass. """ def __init__(self, kvec, pol, s, delta, wb, **kwargs): if callable(kvec): raise TypeError('kvec cannot be a function for a Gaussian beam.') if callable(pol): raise TypeError('Polarization cannot be a function for a Gaussian beam.') # Use super class to define kvec(R, t), pol(R, t), and delta(t) super().__init__(kvec=kvec, pol=pol, delta=delta, **kwargs) # Save the constant values (might be useful): self.con_kvec = kvec self.con_khat = kvec/np.linalg.norm(kvec) self.con_pol = self.pol(np.array([0., 0., 0.]), 0.) # Save the parameters specific to the Gaussian beam: self.s_max = s # central saturation parameter self.wb = wb # 1/e^2 radius self.define_rotation_matrix() def define_rotation_matrix(self): # Angles of rotation: th = np.arccos(self.con_khat[2]) phi = np.arctan2(self.con_khat[1], self.con_khat[0]) # Use scipy to define the rotation matrix self.rmat = Rotation.from_euler('ZY', [phi, th]).inv().as_matrix()
[docs] def intensity(self, R=np.array([0., 0., 0.]), t=0.): # Rotate up to the z-axis where we can apply formulas: Rp = np.einsum('ij,j...->i...', self.rmat, R) rho_sq=np.sum(Rp[:2]**2, axis=0) # Return the intensity: return self.s_max*np.exp(-2*rho_sq/self.wb**2)
[docs]class clippedGaussianBeam(gaussianBeam): """ Clipped, collimated Gaussian beam A beam which has spatially constant k-vector and polarization, with a Gaussian intensity modulation. Specifically, .. math:: \\frac{1}{2}\\hat{\\boldsymbol{\\epsilon}} E_0 e^{-\\mathbf{r}^2/w_b^2} (|\\mathbf{r}|<r_s) e^{i\\mathbf{k}\\cdot\\mathbf{r}-i \\int dt\\Delta(t) + i\\phi(r, t)} where :math:`\\hat{\\boldsymbol{\\epsilon}}` is the polarization, :math:`E_0` is the electric field magnitude, :math:`r_s` is the radius of the stop, :math:`\\mathbf{k}(r,t)` is the k-vector, :math:`\\mathbf{r}` is the position, :math:`\\Delta(t)` is the deutning, :math:`t` is the time, and :math:`\\phi` is the phase. Note that because :math:`I\\propto E^2`, :math:`w_b` is the :math:`1/e^2` radius. Parameters ---------- kvec : array_like with shape (3,) or callable The k-vector of the laser beam, specified as either a three-element list or numpy array. pol : int, float, array_like with shape (3,), or callable The polarization of the laser beam, specified as either an integer, float array_like with shape(3,). If an integer or float, if `pol<0` the polarization will be left circular polarized relative to the k-vector of the light. If `pol>0`, the polarization will be right circular polarized. If array_like, polarization will be specified by the vector, whose basis is specified by `pol_coord`. s : float or callable The maximum intensity of the laser beam at the center, specified as either a float or as callable function. delta : float or callable Detuning of the laser beam. If a callable, it must have a signature like (t) where t is a float and it must return a float. wb : float The :math:`1/e^2` radius of the beam. rs : float The radius of the stop. **kwargs: Additional keyword arguments to pass to the laserBeam superclass. """ def __init__(self, kvec, pol, s, delta, wb, rs, **kwargs): super().__init__(kvec=kvec, pol=pol, s=s, delta=delta, wb=wb, **kwargs) self.rs = rs # Save the radius of the stop.
[docs] def intensity(self, R=np.array([0., 0., 0.]), t=0.): Rp = np.einsum('ij,j...->i...', self.rmat, R) rho_sq = np.sum(Rp[:2]**2, axis=0) return self.s_max*np.exp(-2*rho_sq/self.wb**2)*(np.sqrt(rho_sq)<self.rs)
[docs]class laserBeams(object): """ The base class for a collection of laser beams Parameters ---------- laserbeamparams : array_like of laserBeam or array_like of dictionaries If array_like contains laserBeams, the laserBeams in the array will be joined together to form a collection. If array_like is a list of dictionaries, the dictionaries will be passed as keyword arguments to beam_type beam_type : laserBeam or laserBeam subclass, optional Type of beam to use in the collection of laserBeams. By default `beam_type=laserBeam`. """ def __init__(self, laserbeamparams=None, beam_type=laserBeam): if laserbeamparams is not None: if not isinstance(laserbeamparams, list): raise ValueError('laserbeamparams must be a list.') self.beam_vector = [] for laserbeamparam in laserbeamparams: if isinstance(laserbeamparam, dict): self.beam_vector.append(beam_type(**laserbeamparam)) elif isinstance(laserbeamparam, laserBeam): self.beam_vector.append(laserbeamparam) else: raise TypeError('Each element of laserbeamparams must either ' + 'be a list of dictionaries or list of ' + 'laserBeams') self.num_of_beams = len(self.beam_vector) else: self.beam_vector = [] self.num_of_beams = 0 def __iadd__(self, other): self.beam_vector += other.beam_vector self.num_of_beams = len(self.beam_vector) return self def __add__(self, other): return laserBeams(self.beam_vector + other.beam_vector)
[docs] def add_laser(self, new_laser): """ Add a laser to the collection Parameters ---------- new_laser : laserBeam or laserBeam subclass """ if isinstance(new_laser, laserBeam): self.beam_vector.append(new_laser) self.num_of_beams = len(self.beam_vector) elif isinstance(new_laser, dict): self.beam_vector.append(laserBeam(**new_laser)) else: raise TypeError('new_laser should by type laserBeam or a dictionary' + 'of arguments to initialize the laserBeam class.')
[docs] def pol(self, R=np.array([0., 0., 0.]), t=0.): """ Returns the polarization of the laser beam at position R and t The polarization is returned in the spherical basis. Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- pol : list of array_like, size (3,) polarization of each laser beam at R and t in spherical basis. """ return np.array([beam.pol(R, t) for beam in self.beam_vector])
[docs] def intensity(self, R=np.array([0., 0., 0.]), t=0.): """ Returns the intensity of the laser beam at position R and t Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- s : list of float or array_like Saturation parameters of all laser beams at R and t. """ return np.array([beam.intensity(R, t) for beam in self.beam_vector])
[docs] def kvec(self, R=np.array([0., 0., 0.]), t=0.): """ Returns the k-vector of the laser beam Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- kvec : list of array_like, size(3,) the k vector at position R and time t for each laser beam. """ return np.array([beam.kvec(R, t) for beam in self.beam_vector])
[docs] def delta(self, t=0): """ Returns the detuning of the laser beam at time t Parameters ---------- t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- delta : float or array like detuning of the laser beam at time t for all laser beams """ return np.array([beam.delta(t) for beam in self.beam_vector])
[docs] def electric_field(self, R=np.array([0., 0., 0.]), t=0.): """ Returns the electric field of the laser beams Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- E : list of array_like, size(3,) the electric field vectors at position R and time t for each laser beam. """ return np.array([beam.electric_field(R, t) for beam in self.beam_vector])
[docs] def electric_field_gradient(self, R=np.array([0., 0., 0.]), t=0.): """ Returns the gradient of the electric field of the laser beams Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- dE : list of array_like, size(3,) the electric field gradient matrices at position R and time t for each laser beam. """ return np.array([beam.electric_field_gradient(R, t) for beam in self.beam_vector])
[docs] def total_electric_field(self, R=np.array([0., 0., 0.]), t=0.): """ Returns the total electric field of the laser beams Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- E : array_like, size(3,) the total electric field vector at position R and time t of all the laser beams """ return np.sum(self.electric_field(R, t), axis=0)
[docs] def total_electric_field_gradient(self, R=np.array([0., 0., 0.]), t=0.): """ Returns the total gradient of the electric field of the laser beams Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- dE : array_like, size(3,) the total electric field gradient matrices at position R and time t of all laser beams. """ return np.sum(self.electric_field_gradient(R, t), axis=0)
[docs] def project_pol(self, quant_axis, R=np.array([0., 0., 0.]), t=0, **kwargs): """ Project the polarization onto a quantization axis. Parameters ---------- quant_axis : array_like, shape (3,) A normalized 3-vector of the quantization axis direction. R : array_like, shape (3,), optional If polarization is a function of R is the 3-vectors at which the polarization shall be calculated. calculate_norm : bool, optional If true, renormalizes the quant_axis. By default, False. treat_nans : bool, optional If true, every place that nan is encoutnered, replace with the $hat{z}$ axis as the quantization axis. By default, False. invert : bool, optional If true, invert the process to project the quantization axis onto the specified polarization. Returns ------- projected_pol : list of array_like, shape (3,) The polarization projected onto the quantization axis for all laser beams """ cosbeta = quant_axis[2] sinbeta = np.sqrt(1-cosbeta**2) if isinstance(cosbeta, float): if np.abs(cosbeta)<1: gamma = np.arctan2(quant_axis[1], quant_axis[0]) else: gamma = 0 alpha = 0 else: gamma = np.zeros(cosbeta.shape) inds = np.abs(quant_axis[2])<1 gamma[inds] = np.arctan2(quant_axis[1][inds], quant_axis[0][inds]) alpha = np.zeros(cosbeta.shape) quant_axis = quant_axis.astype('float64') D = np.array([ [(1+cosbeta)/2*np.exp(-1j*alpha + 1j*gamma), -sinbeta/np.sqrt(2)*np.exp(-1j*alpha), (1-cosbeta)/2*np.exp(-1j*alpha - 1j*gamma)], [sinbeta/np.sqrt(2)*np.exp(1j*gamma), cosbeta, -sinbeta/np.sqrt(2)*np.exp(-1j*gamma)], [(1-cosbeta)/2*np.exp(1j*alpha+1j*gamma), sinbeta/np.sqrt(2), (1+cosbeta)/2*np.exp(1j*alpha-1j*gamma)] ]) if quant_axis.shape == (3,) and R.shape == (3,): return [D @ beam.pol(R, t) for beam in self.beam_vector] else: return [np.tensordot(D, beam.pol(R, t), ([1],[0])) for beam in self.beam_vector]
[docs] def cartesian_pol(self, R=np.array([0., 0., 0.]), t=0): """ Returns the polarization of all laser beams in Cartesian coordinates. Parameters ---------- R : array_like, size (3,), optional vector of the position at which to return the polarization. By default, the origin. t : float, optional time at which to return the polarization. By default, t=0. Returns ------- pol : array_like, shape (num_of_beams, 3) polarization of the laser beam at R and t in Cartesian basis. """ return [beam.cartesian_pol(R, t) for beam in self.beam_vector]
[docs] def jones_vector(self, xp, yp, R=np.array([0., 0., 0.]), t=0): """ Jones vector at position R and time t Parameters ---------- xp : array_like, shape (3,) The x vector of the basis in which to calculate the Jones vector. Must be orthogonal to k. yp : array_like, shape (3,) The y vector of the basis in which to calculate the Jones vector. Must be orthogonal to k and `xp`. R : array_like, size (3,), optional vector of the position at which to evaluate the Jones vector. By default, the origin. t : float, optional time at which to evaluate the Jones vector. By default, t=0. Returns ------- pol : array_like, size (num_of_beams, 2) Jones vector of the laser beams at R and t in Cartesian basis. """ return [beam.jones_vector(xp, yp, R, t) for beam in self.beam_vector]
[docs] def stokes_parameters(self, xp, yp, R=np.array([0., 0., 0.]), t=0): """ The Stokes Parameters of the laser beam at R and t Parameters ---------- xp : array_like, shape (3,) The x vector of the basis in which to calculate the Stokes parameters. Must be orthogonal to k. yp : array_like, shape (3,) The y vector of the basis in which to calculate the Stokes parameters. Must be orthogonal to k and `xp`. R : array_like, size (3,), optional vector of the position at which to calculate the Stokes parameters. By default, the origin. t : float, optional time at which to calculate the Stokes parameters. By default, t=0. Returns ------- pol : array_like, shape (num_of_beams, 3) Stokes parameters for the laser beams, [Q, U, V] """ return [beam.stokes_parameters(xp, yp, R, t) for beam in self.beam_vector]
[docs] def polarization_ellipse(self, xp, yp, R=np.array([0., 0., 0.]), t=0): """ The polarization ellipse parameters of the laser beam at R and t Parameters ---------- xp : array_like, shape (3,) The x vector of the basis in which to calculate the polarization ellipse. Must be orthogonal to k. yp : array_like, shape (3,) The y vector of the basis in which to calculate the polarization ellipse. Must be orthogonal to k and `xp`. R : array_like, size (3,), optional vector of the position at which to return the kvector. By default, the origin. t : float, optional time at which to return the k-vector. By default, t=0. Returns ------- list of (psi, chi) : list of tuples list of (:math:`\\psi`, :math:`\\chi`) parameters of the polarization ellipses for each laser beam """ return [beam.polarization_ellipse(xp, yp, R, t) for beam in self.beam_vector]
[docs]class conventional3DMOTBeams(laserBeams): """ A collection of laser beams for 6-beam MOT The standard geometry is to generate counter-progagating beams along all orthogonal axes :math:`(\\hat{x}, \\hat{y}, \\hat{z})`. Parameters ---------- k : float, optional Magnitude of the k-vector for the six laser beams. Default: 1 pol : int or float, optional Sign of the circular polarization for the beams moving along :math:`\\hat{z}`. Default: +1. Orthogonal beams have opposite polarization by default. rotation_angles : array_like List of angles to define a rotated MOT. Default: [0., 0., 0.] rotation_spec : str String to define the convention of the Euler rotations. Default: 'ZYZ' beam_type : pylcp.laserBeam or subclass Type of beam to generate. **kwargs : other keyword arguments to pass to beam_type """ def __init__(self, k=1, pol=+1, rotation_angles=[0., 0., 0.], rotation_spec='ZYZ', beam_type=laserBeam, **kwargs): super().__init__() rot_mat = Rotation.from_euler(rotation_spec, rotation_angles).as_matrix() kvecs = [np.array([ 1., 0., 0.]), np.array([-1., 0., 0.]), np.array([ 0., 1., 0.]), np.array([ 0., -1., 0.]), np.array([ 0., 0., 1.]), np.array([ 0., 0., -1.])] pols = [-pol, -pol, -pol, -pol, +pol, +pol] for kvec, pol in zip(kvecs, pols): self.add_laser(beam_type(kvec=rot_mat @ (k*kvec), pol=pol, **kwargs))
if __name__ == '__main__': import matplotlib.pyplot as plt test_field = magField(lambda R: np.array([-0.5*R[0], -0.5*R[1], 1*R[2]])) print(test_field.Field()) print(test_field.gradField(np.array([5., 2., 1.]))) example_beams = laserBeams([ {'kvec':np.array([0., 0., 1.]), 'pol':np.array([0., 0., 1.]), 'pol_coord':'spherical', 'delta':-2, 's': 1.}, {'kvec':np.array([0., 0., -1.]), 'pol':np.array([0., 0., 1.]), 'pol_coord':'spherical', 'delta':-2, 's': 1.}, ]) print(example_beams.beam_vector[0].jones_vector(np.array([1., 0., 0.]), np.array([0., 1., 0.]))) print(example_beams.kvec()) print(example_beams.pol()) print(example_beams.intensity()) print(example_beams.electric_field_gradient(np.array([0., 0., 0.]), 0.5)) example_beams_2 = laserBeams([ {'kvec':np.array([0., 0., 1.]), 'pol':np.array([0., 0., 1.]), 'pol_coord':'spherical', 'delta':-2, 's': lambda R: 1.}, {'kvec':np.array([0., 0., -1.]), 'pol':np.array([0., 0., 1.]), 'pol_coord':'spherical', 'delta':-2, 's': lambda R: 1.}, ]) print(example_beams_2.electric_field_gradient(np.array([0., 0., 0.]), 0.5)) example_beam = gaussianBeam(np.array([1., 0., 0.]), +1, 5, -2, 1000) print(example_beam.s(np.array([0., 1000/np.sqrt(2), 1000/np.sqrt(2)]))) example_beam = infinitePlaneWaveBeam(np.array([1., 0., 0.]), +1, 5, -2) print(example_beam.electric_field_gradient(np.array([0., 0., 0.]), 0.)) R = np.random.rand(3, 101) t = np.linspace(0, 10, 101) print(example_beam.electric_field_gradient(R, t).shape) MOT_beams = conventional3DMOTBeams(-2, 1, beam_type=gaussianBeam, wb=1000) MOT_beams.beam_vector[1].kvec()