Line data    Source code

       1             : # SPDX-FileCopyrightText: 2026 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           1 : from __future__ import annotations
       5             : 
       6           1 : from typing import TYPE_CHECKING
       7             : 
       8           1 : import numpy as np
       9             : 
      10           1 : from pairinteraction.units import QuantityScalar, ureg
      11             : 
      12             : if TYPE_CHECKING:
      13             :     from collections.abc import Callable
      14             :     from typing import TypeAlias
      15             : 
      16             :     from pairinteraction.units import NDArray, PintFloat
      17             : 
      18             :     PermittivityLike: TypeAlias = "complex | Callable[[PintFloat], complex]"
      19             : 
      20             : 
      21           1 : def evaluate_relative_permittivity(
      22             :     epsilon: PermittivityLike, transition_energy: float, transition_energy_unit: str | None = None
      23             : ) -> complex:
      24             :     """Get the electric permittivity for the given frequency.
      25             : 
      26             :     Args:
      27             :         epsilon: The electric permittivity (dimensionless) or a callable function that returns it.
      28             :         transition_energy: The angular frequency at which to evaluate the permittivity.
      29             :             Only needed if the permittivity is frequency dependent.
      30             :         transition_energy_unit: The unit of the angular frequency.
      31             :             Default None, which means that the angular frequency must be given as pint object.
      32             : 
      33             :     Returns:
      34             :         The electric permittivity at the given angular frequency.
      35             : 
      36             :     """
      37           1 :     omega_au = QuantityScalar.convert_user_to_au(transition_energy, transition_energy_unit, "energy")
      38           1 :     if np.isscalar(epsilon):
      39           1 :         return epsilon  # type: ignore [return-value]
      40           0 :     if callable(epsilon):
      41           0 :         return epsilon(ureg.Quantity(omega_au, "hartree"))
      42           0 :     raise TypeError("epsilon must be either a complex number or a callable function.")
      43             : 
      44             : 
      45           1 : def normalize(vector: NDArray) -> NDArray:
      46             :     """Return the normalized version of the input vector."""
      47           1 :     norm = np.linalg.norm(vector)
      48           1 :     if np.isclose(norm, 0):
      49           0 :         raise ValueError("Cannot normalize a zero vector.")
      50           1 :     return vector / norm  # type: ignore [no-any-return]
      51             : 
      52             : 
      53           1 : def get_lab_to_local_rotation_matrix(normal: NDArray) -> NDArray:
      54             :     """Return a rotation matrix that maps lab coordinates to a frame with local z parallel to normal."""
      55           1 :     if np.isclose(np.linalg.norm(normal), 0):
      56           0 :         raise ValueError("Normal vector cannot be zero.")
      57             : 
      58           1 :     z_axis = normalize(normal)
      59           1 :     reference_axis = np.eye(3)[np.argmin(np.abs(z_axis))]
      60           1 :     x_axis = normalize(np.cross(reference_axis, z_axis))
      61           1 :     y_axis = normalize(np.cross(z_axis, x_axis))
      62             : 
      63           1 :     return np.vstack((x_axis, y_axis, z_axis))
      64             : 
      65             : 
      66           1 : def rotate_vector_to_local(vector_lab: NDArray, lab_to_local_rotation: NDArray) -> NDArray:
      67           1 :     return lab_to_local_rotation @ vector_lab
      68             : 
      69             : 
      70           1 : def rotate_tensor_to_lab(tensor_local: NDArray, lab_to_local_rotation: NDArray) -> NDArray:
      71           1 :     return lab_to_local_rotation.T @ tensor_local @ lab_to_local_rotation