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)