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