Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           0 : import logging
       5           0 : from typing import TYPE_CHECKING, Optional, Union
       6             : 
       7           0 : from attr import dataclass
       8             : 
       9           0 : from pairinteraction import (
      10             :     complex as pi_complex,
      11             :     real as pi_real,
      12             : )
      13           0 : from pairinteraction_gui.calculate.calculate_base import Parameters, Results
      14           0 : from pairinteraction_gui.worker import run_in_other_process
      15             : 
      16             : if TYPE_CHECKING:
      17             :     from typing_extensions import Self
      18             : 
      19             :     from pairinteraction_gui.page import TwoAtomsPage
      20             : 
      21           0 : logger = logging.getLogger(__name__)
      22             : 
      23             : 
      24           0 : @dataclass
      25           0 : class ParametersTwoAtoms(Parameters["TwoAtomsPage"]):
      26             :     """Parameters for the two atoms calculation."""
      27             : 
      28           0 :     pair_basis_energy_delta: float = 0
      29           0 :     order: int = 3
      30             : 
      31           0 :     @classmethod
      32           0 :     def from_page(cls, page: "TwoAtomsPage") -> "Self":
      33           0 :         obj = super().from_page(page)
      34           0 :         obj.pair_basis_energy_delta = page.basis_config.delta_pair_energy.value()
      35           0 :         obj.order = page.system_config.order.value()
      36           0 :         return obj
      37             : 
      38           0 :     def to_replacement_dict(self) -> dict[str, str]:
      39           0 :         replacements = super().to_replacement_dict()
      40           0 :         replacements["$MULTIPOLE_ORDER"] = str(self.order)
      41           0 :         replacements["$PAIR_ENERGY_DELTA"] = str(self.pair_basis_energy_delta)
      42           0 :         return replacements
      43             : 
      44             : 
      45           0 : @dataclass
      46           0 : class ResultsTwoAtoms(Results):
      47           0 :     basis_0_label: Optional[str] = None
      48             : 
      49             : 
      50           0 : @run_in_other_process
      51           0 : def calculate_two_atoms(parameters: ParametersTwoAtoms) -> ResultsTwoAtoms:
      52             :     """Calculate the energy plot for two atoms.
      53             : 
      54             :     This means, given a Paramaters object, do the pairinteraction calculations and return an ResultsTwoAtoms object.
      55             :     """
      56           0 :     pi = pi_real if parameters.is_real else pi_complex
      57           0 :     n_atoms = 2
      58             : 
      59           0 :     kets = tuple(pi.KetAtom(parameters.get_species(i), **parameters.get_quantum_numbers(i)) for i in range(n_atoms))
      60           0 :     bases = tuple(
      61             :         pi.BasisAtom(parameters.get_species(i), **parameters.get_quantum_number_restrictions(i)) for i in range(n_atoms)
      62             :     )
      63             : 
      64           0 :     fields = {k: v for k, v in parameters.ranges.items() if k in ["Ex", "Ey", "Ez", "Bx", "By", "Bz"]}
      65             : 
      66             :     basis_pair_list: Union[list[pi_real.BasisPair], list[pi_complex.BasisPair]]
      67           0 :     if all(v[0] == v[-1] for v in fields.values()):
      68             :         # If all fields are constant, we can only have to diagonalize one SystemAtom per atom
      69             :         # and can construct one BasisPair, which we can use for all steps
      70           0 :         systems = tuple(
      71             :             pi.SystemAtom(bases[i])
      72             :             .set_electric_field(parameters.get_efield(0), unit="V/cm")
      73             :             .set_magnetic_field(parameters.get_bfield(0), unit="G")
      74             :             for i in range(n_atoms)
      75             :         )
      76           0 :         logger.debug("Diagonalizing SystemAtoms...")
      77           0 :         pi.diagonalize(systems, **parameters.diagonalize_kwargs)
      78           0 :         logger.debug("Done diagonalizing SystemAtoms.")
      79           0 :         ket_pair_energy_0 = sum(systems[i].get_corresponding_energy(kets[i], "GHz") for i in range(n_atoms))
      80           0 :         delta_energy = parameters.pair_basis_energy_delta
      81           0 :         basis_pair = pi.BasisPair(
      82             :             systems,
      83             :             energy=(ket_pair_energy_0 - delta_energy, ket_pair_energy_0 + delta_energy),
      84             :             energy_unit="GHz",
      85             :         )
      86             :         # not very elegant, but works (note that importantly this does not copy the basis_pair objects)
      87           0 :         basis_pair_list = parameters.steps * [basis_pair]
      88             :     else:
      89             :         # Otherwise, we have to diagonalize one SystemAtom per atom and per step
      90             :         # and construct one BasisPair per step
      91           0 :         systems_list = []
      92           0 :         for step in range(parameters.steps):
      93           0 :             systems = tuple(
      94             :                 pi.SystemAtom(bases[i])
      95             :                 .set_electric_field(parameters.get_efield(step), unit="V/cm")
      96             :                 .set_magnetic_field(parameters.get_bfield(step), unit="G")
      97             :                 for i in range(n_atoms)
      98             :             )
      99           0 :             systems_list.append(systems)
     100           0 :         systems_flattened = [system for systems in systems_list for system in systems]
     101           0 :         logger.debug("Diagonalizing SystemAtoms...")
     102           0 :         pi.diagonalize(systems_flattened, **parameters.diagonalize_kwargs)
     103           0 :         logger.debug("Done diagonalizing SystemAtoms.")
     104           0 :         delta_energy = parameters.pair_basis_energy_delta
     105           0 :         basis_pair_list = []
     106           0 :         for step in range(parameters.steps):
     107           0 :             ket_pair_energy = sum(
     108             :                 systems_list[step][i].get_corresponding_energy(kets[i], "GHz") for i in range(n_atoms)
     109             :             )
     110           0 :             basis_pair = pi.BasisPair(
     111             :                 systems_list[step],
     112             :                 energy=(ket_pair_energy - delta_energy, ket_pair_energy + delta_energy),
     113             :                 energy_unit="GHz",
     114             :             )
     115           0 :             basis_pair_list.append(basis_pair)
     116           0 :         ket_pair_energy_0 = sum(systems_list[-1][i].get_corresponding_energy(kets[i], "GHz") for i in range(n_atoms))
     117             : 
     118           0 :     system_pair_list: Union[list[pi_real.SystemPair], list[pi_complex.SystemPair]] = []
     119           0 :     for step in range(parameters.steps):
     120           0 :         system = pi.SystemPair(basis_pair_list[step])
     121           0 :         system.set_interaction_order(parameters.order)
     122           0 :         if "Distance" in parameters.ranges:
     123           0 :             distance = parameters.ranges["Distance"][step]
     124           0 :             angle: float = 0
     125           0 :             if "Angle" in parameters.ranges:
     126           0 :                 angle = parameters.ranges["Angle"][step]
     127           0 :             system.set_distance(distance, angle, unit="micrometer")
     128           0 :         system_pair_list.append(system)
     129             : 
     130           0 :     logger.debug("Diagonalizing SystemPairs...")
     131           0 :     pi.diagonalize(
     132             :         system_pair_list,
     133             :         **parameters.diagonalize_kwargs,
     134             :         **parameters.get_diagonalize_energy_range(ket_pair_energy_0),
     135             :     )
     136           0 :     logger.debug("Done diagonalizing SystemPairs.")
     137             : 
     138           0 :     results = ResultsTwoAtoms.from_calculate(system_pair_list, kets, ket_pair_energy_0)
     139           0 :     results.basis_0_label = (
     140             :         str(basis_pair_list[-1]) + f"\n  ⇒ Basis consists of {basis_pair_list[-1].number_of_kets} kets"
     141             :     )
     142             : 
     143           0 :     return results