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