Line data    Source code

       1              : # SPDX-FileCopyrightText: 2025 PairInteraction Developers
       2              : # SPDX-License-Identifier: LGPL-3.0-or-later
       3              : 
       4            1 : import logging
       5            1 : from typing import TYPE_CHECKING
       6              : 
       7            1 : from attr import dataclass
       8              : 
       9            1 : import pairinteraction as pi_complex
      10            1 : import pairinteraction.real as pi_real
      11            1 : from pairinteraction_gui.calculate.calculate_base import Parameters, Results
      12              : 
      13              : if TYPE_CHECKING:
      14              :     from pairinteraction_gui.page import OneAtomPage  # noqa: F401  # related to ruff extend-generics
      15              : 
      16            1 : logger = logging.getLogger(__name__)
      17              : 
      18              : 
      19            1 : @dataclass
      20            1 : class ParametersOneAtom(Parameters["OneAtomPage"]):
      21              :     """Parameters for the single-atom calculation."""
      22              : 
      23              : 
      24            1 : @dataclass
      25            1 : class ResultsOneAtom(Results):
      26              :     """Results for the single-atom calculation."""
      27              : 
      28              : 
      29            1 : def calculate_one_atom(parameters: ParametersOneAtom) -> ResultsOneAtom:
      30              :     """Calculate the energy plot for one atom.
      31              : 
      32              :     This means, given a Parameters object, do the PairInteraction calculations and return an ResultsOneAtom object.
      33              :     """
      34            1 :     return _calculate_one_atom(parameters)
      35              : 
      36              : 
      37            1 : def _calculate_one_atom(parameters: ParametersOneAtom) -> ResultsOneAtom:
      38              :     """Make the unwrapped function available for testing."""
      39            1 :     pi = pi_real if parameters.is_real else pi_complex
      40              : 
      41            1 :     ket = pi.KetAtom(parameters.get_species(), **parameters.get_quantum_numbers())
      42            1 :     ket_energy = ket.get_energy("GHz")
      43            1 :     basis = pi.BasisAtom(parameters.get_species(), **parameters.get_quantum_number_restrictions())
      44              : 
      45            1 :     system_list = [
      46              :         pi.SystemAtom(basis)
      47              :         .set_electric_field(parameters.get_efield(step), unit="V/cm")
      48              :         .set_magnetic_field(parameters.get_bfield(step), unit="G")
      49              :         .set_diamagnetism_enabled(parameters.diamagnetism_enabled)
      50              :         for step in range(parameters.steps)
      51              :     ]
      52              : 
      53            1 :     logger.debug("Diagonalizing SystemAtoms...")
      54            1 :     pi.diagonalize(
      55              :         system_list,
      56              :         **parameters.diagonalize_kwargs,
      57              :         **parameters.get_diagonalize_energy_range_kwargs(ket_energy),
      58              :     )
      59            1 :     logger.debug("Done diagonalizing SystemAtoms.")
      60            1 :     return ResultsOneAtom.from_calculate(parameters, system_list, ket, ket_energy)