msptools 0.1.0__py3-none-any.whl

msptools/GreenTensor_Electric.py

@@ -0,0 +1,236 @@
+from .backend import get_backend
+from numpy.typing import ArrayLike
+import numpy as np
+from scipy.constants import pi
+from cmath import exp
+
+def G_0_function(r: float | ArrayLike, wave_number: float) -> complex | ArrayLike:
+    """
+    Computes the G_0 function for a given distance r and wave number.
+
+    Parameters
+    ----------
+    r :
+        The distance between two points.
+    wave_number :
+        The wave number.
+
+    Returns
+    -------
+    complex | ArrayLike
+        The value of the G_0 function.
+    """
+    if isinstance(r, (float, int)):
+        if r > 0:
+            kr = wave_number * r
+            return exp(1j * wave_number * r) / (4 * pi * r) * (1 + 1j/kr - 1/kr**2)
+        else:
+            return 0.0j
+    else:
+        xp = get_backend(r)
+        mask = r > 0
+        result = xp.zeros_like(r, dtype=complex)
+        kr = wave_number * r[mask]
+        result[mask] = xp.exp(1j * kr) / (4 * pi * r[mask]) * (1 + 1j/kr - 1/kr**2)
+        return result
+
+def G_1_function(r: float | ArrayLike, wave_number: float) -> complex | ArrayLike:
+    """
+    Computes the G_1 function for a given distance r and wave number.
+
+    Parameters
+    ----------
+    r :
+        The distance between two points.
+    wave_number :
+        The wave number.
+
+    Returns
+    -------
+    complex | ArrayLike
+        The value of the G_1 function.
+    """
+    if isinstance(r, (float, int)):
+        if r > 0:
+            kr = wave_number * r
+            return -exp(1j * wave_number * r) / (4 * pi * r**3) * (1 + 3j/kr - 3/kr**2)
+        else:
+            return 0.0j
+    else:
+        xp = get_backend(r)
+        mask = r > 0
+        result = xp.zeros_like(r, dtype=complex)
+        kr = wave_number * r[mask]
+        result[mask] = -xp.exp(1j * kr) / (4 * xp.pi * r[mask]**3) * (1 + 3j/kr - 3/kr**2)
+        return result
+
+def G_0_derivative_function(r: float | ArrayLike, wave_number: float) -> complex:
+    """
+    Computes the derivative of the G_0 function with respect to r.
+
+    Parameters
+    ----------
+    r : float
+        The distance between two points.
+    wave_number : float
+        The wave number.
+
+    Returns
+    -------
+    complex
+        The value of the derivative of the G_0 function.
+    """
+    if isinstance(r, (float, int)):
+        xp = np
+    else:
+        xp = get_backend(r)
+    return wave_number * xp.exp(1j * wave_number * r) / (4 * xp.pi * r) * \
+           (1j - 2/(wave_number * r) - 3j/(wave_number * r)**2 + 3/(wave_number * r)**3)
+
+def G_1_derivative_function(r: float | ArrayLike, wave_number: float) -> complex:
+    """
+    Computes the derivative of the G_1 function with respect to r.
+
+    Parameters
+    ----------
+    r : float
+        The distance between two points.
+    wave_number : float
+        The wave number.
+
+    Returns
+    -------
+    complex
+        The value of the derivative of the G_1 function.
+    """
+    if isinstance(r, (float, int)):
+        xp = np
+    else:
+        xp = get_backend(r)
+    return -wave_number * xp.exp(1j * wave_number * r) / (4 * xp.pi * r**3) * \
+           (1j - 6/(wave_number * r) - 15j/(wave_number * r)**2 + 15/(wave_number * r)**3)
+
+def v_cross_derivative(r_vec: ArrayLike, coordinate: int) -> np.ndarray:
+    """
+    Computes the derivative of a vector cross dyadic product with respect to a specific coordinate.
+
+    Parameters
+    ----------
+    r_vec : 
+        The vector for which the derivative is computed.
+    coordinate : 
+        The coordinate with respect to which the derivative is taken (0, 1, or 2).
+
+    Returns
+    -------
+    np.ndarray
+        The derivative of the cross product with respect to the specified coordinate.
+    """
+    xp = get_backend(r_vec)
+
+    dimensions = r_vec.shape[0]
+    if coordinate < 0 or coordinate >= dimensions:
+        raise ValueError("Coordinate must be in the range [0, {}]".format(dimensions - 1))
+
+    der_R_cross = xp.zeros((dimensions, dimensions))
+
+    for i in range(dimensions):
+        if i == coordinate:
+            der_R_cross[i, i] = 2 * r_vec[i]
+        else:
+            der_R_cross[i, coordinate] = r_vec[i]
+            der_R_cross[coordinate, i] = r_vec[i]
+
+    return der_R_cross
+
+def construct_green_tensor(positions : np.ndarray, wave_number: float) -> np.ndarray:
+    """
+    Constructs the Green's tensor for a given set of positions and wave number.
+
+    Parameters
+    ----------
+    positions : np.ndarray
+        Array of shape (num_particles, dimension) containing the positions of the particles.
+    wave_number : float
+        The wave number.
+
+    Returns
+    -------
+    np.ndarray
+        Green's tensor of shape (num_particles, num_particles, dimension, dimension).
+    """
+    xp = get_backend(positions)
+    dimensions = positions.shape[1]
+    rel_vec_matrix = positions[:, None, :] - positions[None, :, :]
+    distances = xp.linalg.norm(rel_vec_matrix, axis=-1)
+    G_0_matrix = G_0_function(distances, wave_number)
+    G_1_matrix = G_1_function(distances, wave_number)
+    R_cross_matrix = rel_vec_matrix[:, :, :, None] * rel_vec_matrix[:, :, None, :]
+    green_tensor = G_0_matrix[:, :, None, None] * xp.eye(dimensions) + G_1_matrix[:, :, None, None] * R_cross_matrix
+    
+    return green_tensor
+
+
+def pair_green_tensor_derivative(pos_i: np.ndarray, pos_j: np.ndarray, coordinate : int,  wave_number: float):
+    """
+    Constructs the derivative of the pair Green's tensor with respect to a specific coordinate.
+
+    Parameters
+    ----------
+    pos_i : np.ndarray
+        Position of the first particle.
+    pos_j : np.ndarray
+        Position of the second particle.
+    coordinate : int
+        The coordinate with respect to which the derivative is taken (0, 1, or 2).
+    wave_number : float
+        The wave number.
+
+    Returns
+    -------
+    np.ndarray
+        Derivative of the pair Green's tensor with respect to the specified coordinate.
+    """
+    xp = get_backend(pos_i)
+    dimensions = pos_i.shape[0]
+    R_vec = pos_i - pos_j
+    r = xp.linalg.norm(R_vec)
+
+    g_1 = G_1_function(r, wave_number)
+    der_g_0 = G_0_derivative_function(r, wave_number) * R_vec[coordinate] / r
+    der_g_1 = G_1_derivative_function(r, wave_number) * R_vec[coordinate] / r
+    R_cross = R_vec[:, None] @ R_vec[None, :]
+    der_R_cross = v_cross_derivative(R_vec, coordinate)
+
+    derivative_tensor = der_g_0 * xp.eye(dimensions) + der_g_1 * R_cross + g_1 * der_R_cross
+    
+    return derivative_tensor 
+
+def construct_green_tensor_gradient(positions : np.ndarray, wave_number: float) -> np.ndarray:
+    """
+    Constructs the derivative of the Green's tensor for a given set of positions and wave number.
+
+    Parameters
+    ----------
+    positions : np.ndarray
+        Array of shape (num_particles, dimension) containing the positions of the particles.
+    wave_number : float
+        The wave number.
+
+    Returns
+    -------
+    np.ndarray
+        Derivative of Green's tensor of shape (num_particles, num_particles, dimension, dimension, dimension).
+    """
+    xp = get_backend(positions)
+    
+    num_particles, dimensions = positions.shape
+    green_tensor_derivative = xp.zeros((num_particles, num_particles, dimensions, dimensions, dimensions), dtype=xp.complex128)
+
+    for i in range(num_particles):
+        for j in range(i + 1, num_particles):
+            for coord in range(dimensions):
+                green_tensor_derivative[i, j, coord, :, :] = pair_green_tensor_derivative(positions[i], positions[j], coord, wave_number)
+                green_tensor_derivative[j, i, coord, :, :] = -green_tensor_derivative[i, j, coord, :, :]
+    return green_tensor_derivative
+    

msptools/MSP.py

@@ -0,0 +1,175 @@
+from .backend import get_backend
+from typing import Iterable
+from numpy.typing import ArrayLike
+
+from msptools.dipole_moments import calculate_dipole_moments_linear
+
+def solve_MSP_from_arrays(polarizability: ArrayLike,
+                          external_field : ArrayLike,
+                          wave_number : float,
+                          green_tensor : ArrayLike,
+                          method : str = 'Inverse',
+                          **kwargs) -> ArrayLike:
+
+    """
+    Solve the Multiple Scattering Problem (MSP) using the provided arrays.
+
+    Parameters
+    ----------
+    polarizability :
+        Polarizability of the particles, can be a complex number, float, int, list, or numpy array.
+    external_field :
+        External field on particles positions.
+    wave_number :
+        Wave number of the incident wave.
+    green_tensor :
+        Green's tensor for the system.
+    method :
+        Method to solve the MSP, either 'Iterative' or 'Inverse'. The default is 'Iterative'.
+
+    Returns
+    -------
+    ArrayLike
+        The solution to the MSP.
+
+    """
+    
+    
+    if green_tensor.ndim != 4 or green_tensor.shape[0] != green_tensor.shape[1] or green_tensor.shape[2] != green_tensor.shape[3]:
+        raise ValueError("Invalid green_tensor shape. Expected shape (N, N, d, d), got {}".format(green_tensor.shape))
+    if green_tensor.shape[0] != external_field.shape[0]:
+        raise ValueError("The first dimension of green_tensor must match the number of particles in external_field. Expected {}, got {}".format(external_field.shape[0], green_tensor.shape[0]))
+    if green_tensor.shape[2] != external_field.shape[1]:
+        raise ValueError("The third dimension of green_tensor must match the system dimensionality. Expected {}, got {}".format(external_field.shape[1], green_tensor.shape[2]))
+
+    if method == 'Iterative':
+        if 'tolerance' in kwargs:
+            tolerance = kwargs['tolerance']
+            return array_MSP_iterative(polarizability, external_field, wave_number, green_tensor, tolerance=tolerance)
+        else:
+            return array_MSP_iterative(polarizability, external_field, wave_number, green_tensor)
+    elif method == 'Inverse':
+        return array_MSP_inverse(polarizability, external_field, wave_number, green_tensor)
+    else:
+        raise ValueError("Unknown method: {}".format(method))
+
+def array_MSP_iterative(polarizability : ArrayLike,
+                          external_field : ArrayLike,
+                          wave_number : float,
+                          green_tensor : ArrayLike,
+                          num_iterations : int = 500,
+                          tolerance : float = 1e-6) -> ArrayLike:
+    
+    """
+    Solve the MSP using an iterative method.
+
+    Parameters
+    ----------
+    polarizability :
+        Polarizability of the particles.
+    external_field :
+        External field on particles positions.
+    wave_number :
+        Wave number of the incident wave.
+    green_tensor :
+        Green's tensor for the system.
+    num_iterations : optional
+        Maximum number of iterations for the iterative method. Default is 500.
+    tolerance : optional
+        Convergence tolerance for the iterative method. Default is 1e-6.
+
+    Returns
+    -------
+    xp.ndarray
+        The solution to the MSP.
+    """
+
+    xp = get_backend(external_field)
+    old_field = external_field.copy()
+
+    for iteration in range(num_iterations):
+        
+        dipole_moments = calculate_dipole_moments_linear(polarizability, old_field)
+        scattered_field = wave_number**2 * xp.einsum('ijmn,jn->im', green_tensor, dipole_moments)
+        new_field = external_field + scattered_field
+
+        if xp.allclose(new_field, old_field, rtol=tolerance):
+            break
+        old_field = new_field.copy()
+
+    if iteration == num_iterations - 1:
+        print(f"Warning: MSP iterative solution did not converge within {num_iterations} iterations.")
+
+    return new_field
+
+def array_MSP_inverse(polarizability : ArrayLike,
+                        external_field : ArrayLike,
+                        wave_number : float,
+                        green_tensor : ArrayLike) -> ArrayLike:
+        """
+        Solve the MSP using the inverse method.
+    
+        Parameters
+        ----------
+        polarizability :
+            Polarizability of the particles.
+        external_field :
+            External field on particles positions.
+        wave_number :
+            Wave number of the incident wave.
+        green_tensor :
+            Green's tensor for the system.
+    
+        Returns
+        -------
+        xp.ndarray
+            The solution to the MSP.
+        """
+        xp = get_backend(external_field)
+        num_particles = external_field.shape[0]
+        dimensions = external_field.shape[1]
+        
+        scattering_matrix = wave_number**2 * xp.einsum('ijmk,jkl->ijml', green_tensor, polarizability)
+
+        MSP_matrix = xp.eye(num_particles * dimensions) - scattering_matrix.transpose(0,2,1,3).reshape(num_particles * dimensions, num_particles * dimensions)
+        
+        total_field = xp.linalg.solve(MSP_matrix, external_field.flatten())
+        total_field = total_field.reshape(num_particles, dimensions)
+        return total_field
+
+def MSP_gradient_from_arrays(dipole_moments: ArrayLike,
+                             external_gradient : ArrayLike,
+                             wave_number : float,
+                             green_tensor_derivative : ArrayLike) -> ArrayLike:
+    """
+    Compute the gradient of the MSP solution with respect to particle positions.
+
+    Parameters
+    ----------
+    polarizability :
+        Polarizability of the particles.
+    MS_field :
+        Multiple scattering field on particles positions.
+    external_gradient :
+        Gradient of the external field on particles positions.
+    wave_number :
+        Wave number of the incident wave.
+    green_tensor_derivative :
+        Derivative of the Green's tensor with respect to particle positions.
+
+    Returns
+    -------
+    xp.ndarray
+        The gradient of the MSP solution with respect to particle positions.
+    
+    Notes
+    -----
+    The gradient is returned as an array of shape (N, d, d) where N is the number of particles and d is the dimensionality.
+    """
+
+    xp = get_backend(external_gradient)
+    scattered_gradient = wave_number**2 * xp.einsum('ijcmn,jn->icm', green_tensor_derivative, dipole_moments)
+    
+    MSP_gradient = external_gradient + scattered_gradient
+
+    return MSP_gradient

msptools/OFO_calculations.py

@@ -0,0 +1,36 @@
+from typing import Iterable
+from .backend import get_backend
+from numpy.typing import ArrayLike
+
+def calculate_forces_eppgrad(medium_permittivity: float, dipole_moments: ArrayLike, field_gradient: ArrayLike) -> ArrayLike:
+    """
+    Calculate the force on a set of dipoles in an electric field gradient.
+
+    Parameters
+    ----------
+    medium_permittivity :
+        The permittivity of the medium in which the dipoles are located.
+    dipole_moments :
+        An array representing the dipole moments of the particles. Shape should be (N, d), 
+        where N is the number of dipoles and d is the dimensionality.
+    field_gradient :
+        An array representing the electric field gradient at the location of the dipoles. 
+        Shape should be (N, d, d), where N is the number of dipoles and d is the dimensionality.
+
+    Returns
+    -------
+    Forces :
+        An array representing the force on each dipole.
+    
+    Notes
+    -----
+    The force is calculated using the formula:
+        F = (ε/2) * Re{ p · ∇E* }
+    where ε is the medium permittivity, p is the dipole moment, and ∇E* is the complex conjugate of the electric field gradient.
+    """
+    xp = get_backend(dipole_moments)
+
+    forces = (medium_permittivity / 2) * xp.real(xp.einsum('im,inm->in', dipole_moments, xp.conj(field_gradient)))
+
+    return forces
+

msptools/__init__.py

@@ -0,0 +1,192 @@
+from dataclasses import Field
+import types
+from .OFO_calculations import *
+from .dipole_moments import *
+from .polarizability_mod import *
+from .particle_types import *
+from .particles_mod import *
+from .permittivity import *
+from .field_mod import *
+from .tools.unit_calcs import *
+from .GreenTensor_Electric import *
+from .MSP import *
+from .backend import get_backend
+from numpy.typing import ArrayLike, NDArray
+from typing import List
+
+__all__ = [
+    "OFO_calculations",
+    "dipole_moments",
+    "polarizability_mod",
+    "particle_types",
+    "particles_mod",
+    "permittivity",
+    "field_mod",
+    "unit_calcs",
+    "GreenTensor_Electric",
+    "MSP"
+]
+
+class System:
+    """Class representing a Optical_Forces physical system containing particles."""
+
+    def __init__(self, device: str = "GPU") -> None:
+        """Initialize a System object by specifying the device to use for calculations."""
+        
+        print(f"Initializing System with device: {device}")
+        if device not in ["GPU", "CPU"]:
+            raise ValueError("Invalid device specified. Use 'GPU' or 'CPU'.")
+        elif device == "GPU":
+            try:
+                import cupy as xp
+            except ImportError:
+                print("CuPy is not available. Falling back to CPU.")
+                import numpy as xp
+        else:
+            import numpy as xp
+        self.xp = xp
+
+    def set_system(self, particle_types : ParticleType | List[ParticleType], field: Field, positions_unit: str, medium_permittivity: float = 1.0) -> None:
+        """
+        Set a System object by specifying the particle types, the field and the medium permittivity.
+        
+        Parameters
+        ----------
+        particle_types :
+            The type(s) of the particles in the system. This can be a single ParticleType or a list of ParticleType objects.
+        field :
+            The field in which the particles are placed. This should be an instance of the Field class.
+        positions_unit :
+            The unit of the positions of the particles. This should be a string representing a unit of length (e.g., 'nm', 'um', 'm').
+        medium_permittivity :
+            The permittivity of the medium in which the particles are placed. This should be a float.
+        """
+        if not isinstance(particle_types, list):
+            particle_types = [particle_types]
+        self.particle_types = particle_types
+        self.field = field
+        self.field.set_medium_permittivity(medium_permittivity)
+        self.medium_permittivity = medium_permittivity
+        self.positions_unit = positions_unit
+        self.particles = Particles(self.xp)
+        self.medium_wave_number_nm = frequency_to_wavenumber_nm(self.field.frequency_eV) * self.xp.sqrt(self.medium_permittivity)
+        for ptype in self.particle_types:
+            ptype.compute_polarizability(frequency = self.field.frequency_eV, medium_permittivity=self.medium_permittivity)
+    
+    def add_particles(self,
+                     positions: ArrayLike,
+                     particle_type: ParticleType | None = None) -> None:
+        """
+        Add particles to the system at specified positions.
+
+        Parameters
+        ----------
+        positions :
+            The position of the particles to add. This can be a 1D-three-element or 2D array-like.
+        type :
+            The type of the particles to add. If not specified, and there is only one type in the system, that type will be used.
+        """
+
+        if particle_type is None and len(self.particle_types) > 1:
+            raise ValueError("When adding particles to a multi-type system, the 'particle_type' parameter must be specified.")
+        else:
+            particle_type = self.particle_types[0]
+
+        if particle_type is not None and particle_type not in self.particle_types:
+            raise ValueError("The specified particle type is not part of the system's types.")
+
+        positions = positions* get_multiplier_nanometers(self.positions_unit)
+        if positions.ndim == 1:
+            positions = self.xp.array(positions)
+ 
+        polarizability = polarizability_to_matrix(particle_type.polarizability, positions.shape[0], 3, self.xp)
+        print(f"positions type: {type(positions)}, shape: {positions.shape}")
+        print(f"polarizability type: {type(polarizability)}, shape: {polarizability.shape}")
+        self.particles.add_particles(positions=positions, polarizabilities=polarizability)
+    
+    def get_field_in_particles(self, method : str = 'Inverse') -> ArrayLike:
+        """
+        Get the electric field at specified positions by solving the Multiple Scattering Problem (MSP).
+
+        Returns
+        -------
+        xp.ndarray
+            The electric field at the specified positions.
+        """
+        
+        positions = self.particles.positions
+        external_field = self.field.get_external_field_in_positions(positions)
+        green_tensor = construct_green_tensor(positions, self.medium_wave_number_nm)
+        field_solution = solve_MSP_from_arrays(polarizability=self.particles.polarizabilities,
+                                   external_field=external_field,
+                                   wave_number=self.medium_wave_number_nm,
+                                   green_tensor=green_tensor,
+                                   method=method)
+        return field_solution
+    
+    def get_field_gradient_in_particles(self, current_field: ArrayLike) -> ArrayLike:
+        """
+        Get the electric field gradient at specified positions by solving the Multiple Scattering Problem (MSP) for the gradient.
+
+        Returns
+        -------
+        ArrayLike
+            The electric field gradient at the specified positions.
+        """
+        
+        external_gradient = self.field.get_external_gradient_in_positions(self.particles.positions)
+        green_tensor_derivative = construct_green_tensor_gradient(self.particles.positions, self.medium_wave_number_nm)
+        dipole_moments = calculate_dipole_moments_linear(self.particles.polarizabilities,
+                                                         current_field) 
+        gradient_solution = MSP_gradient_from_arrays(dipole_moments=dipole_moments,
+                                                     external_gradient=external_gradient,
+                                                     wave_number=self.medium_wave_number_nm,
+                                                     green_tensor_derivative=green_tensor_derivative)
+        return gradient_solution
+    
+    def set_position(self, index: int, position: ArrayLike) -> None:
+        """
+        Set the position of a particle at a specified index.
+
+        Parameters
+        ----------
+        index :
+            The index of the particle to set the position for.
+
+        position :
+            The new position of the particle. This can be a 1D-three-element array-like.
+        """
+        position = xp.array(position)* get_multiplier_nanometers(self.positions_unit)
+        if position.ndim != 1 or position.shape[0] != 3:
+            raise ValueError("Position must be a 1D-three-element array-like.")
+        self.particles.set_position(index, position.tolist())
+
+
+class ForceCalculator:
+    """Class to compute optical forces on particles in a System."""
+    
+    def __init__(self, system: System) -> None:
+        """
+        Initialize a ForceCalculator object by specifying the System.
+        """
+        self.system = system
+
+
+    def compute_forces(self) -> xp.ndarray:
+        """
+        Compute the optical forces on particles at specified positions.
+
+        Returns
+        -------
+        xp.ndarray
+            The computed optical forces on the particles.
+        """
+
+        E_field = self.system.get_field_in_particles()
+        E_grad = self.system.get_field_gradient_in_particles(E_field)
+        dipole_moments = calculate_dipole_moments_linear(self.system.particles.polarizabilities, E_field)
+        forces = calculate_forces_eppgrad(self.system.medium_permittivity, dipole_moments, E_grad)
+
+        return forces
+
+

msptools/backend.py

@@ -0,0 +1,22 @@
+from numpy.typing import ArrayLike
+import numpy as np
+try:
+    import cupy as cp
+except ImportError:
+    cp = None
+
+def get_backend(array: ArrayLike) -> None:
+    _backend = None
+    if isinstance(array, np.ndarray):
+        _backend = np
+    elif isinstance(array, cp.ndarray) and cp is not None:
+        _backend = cp
+    else:
+        raise ValueError("Unsupported array type. Only NumPy and CuPy arrays are supported.")
+    return _backend
+
+
+
+    
+
+

msptools/dipole_moments.py

@@ -0,0 +1,27 @@
+from .backend import get_backend
+from numpy.typing import ArrayLike
+
+
+def calculate_dipole_moments_linear(polarizability: ArrayLike,
+                                    electric_field : ArrayLike) -> ArrayLike:
+    """
+    Calculate the dipole moments of particles in an electric field using a linear polarizability model.
+    
+    Parameters
+    ----------
+    polarizability :
+        The polarizability of the particles. This is in general an (N, d, d) array, where N is the number of particles and d is the dimensionality of the system. It can also be a scalar (complex, float, or int) which will be applied to all particles.
+    electric_field :
+        The electric field at the location of the particles. This should be an array of shape (N, d), where N is the number of particles and d is the dimensionality of the system.
+    
+    Returns
+    -------
+    ArrayLike
+        An array of shape (N, d) representing the dipole moments of the particles.
+    """
+    
+    xp = get_backend(electric_field)
+    dipole_moments = xp.einsum('ikl,il->ik', polarizability, electric_field)
+
+
+    return dipole_moments