# SPDX-FileCopyrightText: 2024 PairInteraction Developers
# SPDX-License-Identifier: LGPL-3.0-or-later
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, overload
import numpy as np
from pairinteraction import _backend
from pairinteraction.basis import BasisAtom, BasisAtomReal
from pairinteraction.system.system_base import SystemBase
from pairinteraction.units import QuantityScalar
if TYPE_CHECKING:
from typing_extensions import Self
from pairinteraction.green_tensor import GreenTensorBase
from pairinteraction.ket import KetAtom
from pairinteraction.units import ArrayLike, PintArrayLike, PintFloat
logger = logging.getLogger(__name__)
[docs]
class SystemAtom(SystemBase[BasisAtom]):
"""System of a single atom.
Use the given BasisAtom object to create the system object.
You can set the electric and magnetic fields and enable diamagnetism afterwards via the corresponding methods.
Examples:
>>> import pairinteraction as pi
>>> ket = pi.KetAtom("Rb", n=60, l=0, m=0.5)
>>> basis = pi.BasisAtom("Rb", n=(58, 63), l=(0, 3))
>>> system = pi.SystemAtom(basis)
>>> system.set_magnetic_field([0, 0, 1], unit="gauss").set_electric_field([0.1, 0, 0.1], unit="V/cm")
SystemAtom(BasisAtom('Rb', n=(58, 63), l=(0, 3)), is_diagonal=False)
>>> system.diagonalize()
SystemAtom(BasisAtom(|Rb:58,S_1/2,-1/2⟩ ... |Rb:63,F_5/2,5/2⟩), is_diagonal=True)
>>> eigenenergies = system.get_eigenenergies(unit="GHz")
>>> print(f"{eigenenergies[0] - ket.get_energy(unit='GHz'):.5f}")
-75.51823
"""
_cpp: _backend.SystemAtomComplex
_cpp_type = _backend.SystemAtomComplex
_basis_class = BasisAtom
[docs]
def __init__(self, basis: BasisAtom) -> None:
"""Create a system object for a single atom.
Args:
basis: The :class:`pairinteraction.BasisAtom` object that describes the basis of the system.
"""
self._cpp = self._cpp_type(basis._cpp)
self._basis = basis
[docs]
def set_electric_field(
self: Self,
electric_field: PintArrayLike | ArrayLike,
unit: str | None = None,
) -> Self:
"""Set the electric field for the system.
Args:
electric_field: The electric field to set for the system.
unit: The unit of the electric field, e.g. "V/cm".
Default None expects a `pint.Quantity`.
"""
electric_field_au = [QuantityScalar.convert_user_to_au(v, unit, "electric_field") for v in electric_field]
self._cpp.set_electric_field(electric_field_au)
return self
[docs]
def set_magnetic_field(
self: Self,
magnetic_field: PintArrayLike | ArrayLike,
unit: str | None = None,
) -> Self:
"""Set the magnetic field for the system.
Args:
magnetic_field: The magnetic field to set for the system.
unit: The unit of the magnetic field, e.g. "gauss".
Default None expects a `pint.Quantity`.
"""
magnetic_field_au = [QuantityScalar.convert_user_to_au(v, unit, "magnetic_field") for v in magnetic_field]
self._cpp.set_magnetic_field(magnetic_field_au)
return self
[docs]
def set_diamagnetism_enabled(self: Self, enable: bool = True) -> Self:
"""Enable or disable diamagnetism for the system.
Args:
enable: Whether to enable or disable diamagnetism.
"""
self._cpp.set_diamagnetism_enabled(enable)
return self
[docs]
def set_distance_to_ion(
self: Self, distance: float | PintFloat, angle_degree: float = 0, unit: str | None = None
) -> Self:
distance_vector = [np.sin(np.deg2rad(angle_degree)) * distance, 0, np.cos(np.deg2rad(angle_degree)) * distance]
return self.set_ion_distance_vector(distance_vector, unit)
[docs]
def set_ion_distance_vector(
self: Self,
distance: PintArrayLike | ArrayLike,
unit: str | None = None,
) -> Self:
distance_au = [QuantityScalar.convert_user_to_au(v, unit, "distance") for v in distance]
self._cpp.set_ion_distance_vector(distance_au)
return self
[docs]
def set_ion_charge(self: Self, charge: float | PintFloat, unit: str | None = None) -> Self:
charge_au = QuantityScalar.convert_user_to_au(charge, unit, "charge")
self._cpp.set_ion_charge(charge_au)
return self
[docs]
def set_ion_interaction_order(self: Self, order: int) -> Self:
self._cpp.set_ion_interaction_order(order)
return self
[docs]
def set_green_tensor(self, green_tensor: GreenTensorBase) -> Self:
"""Set the Green tensor for the single atom system.
This will add self interaction terms to the Hamiltonian of the single atom system
based on the provided Green tensor.
This can be used to e.g. include the effect of a nearby surface on the energy levels of the atom.
Args:
green_tensor: The Green tensor to set for the system.
"""
if not green_tensor.static_limit:
raise ValueError("Setting a Green tensor that is not in the static limit is currently not supported.")
if not np.allclose(
green_tensor.pos1_au,
green_tensor.pos2_au,
rtol=0,
atol=1e-12,
):
raise ValueError(
"You can only set a Green tensor for the single atom system, where the positions are identical."
)
if not green_tensor.without_vacuum_contribution:
raise ValueError(
"The Green tensor for the single atom system needs to be provided without the vacuum contribution, "
"i.e., without_vacuum_contribution has to be set to True when creating the Green tensor."
)
use_real = isinstance(self, SystemAtomReal)
gti = green_tensor.get_interpolator(use_real=use_real)
self._cpp.set_green_tensor_interpolator(gti._cpp)
return self
@overload
def get_corresponding_energy(self: Self, ket: KetAtom, unit: None = None) -> PintFloat: ...
@overload
def get_corresponding_energy(self: Self, ket: KetAtom, unit: str) -> float: ...
[docs]
def get_corresponding_energy(self: Self, ket: KetAtom, unit: str | None = None) -> float | PintFloat:
state_idx = self.get_eigenbasis().get_corresponding_state_index(ket)
return self.get_eigenenergies(unit=unit)[state_idx] # type: ignore [index,no-any-return] # PintArray does not know it can be indexed
class SystemAtomReal(SystemAtom):
_cpp: _backend.SystemAtomReal # type: ignore [assignment]
_cpp_type = _backend.SystemAtomReal # type: ignore [assignment]
_basis_class = BasisAtomReal