Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 Pairinteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           1 : import re
       5           1 : from typing import Literal
       6             : 
       7           1 : from pairinteraction.cli import download_databases
       8           1 : from pairinteraction_gui.worker import run_in_other_process
       9             : 
      10           1 : AVAILABLE_SPECIES = [
      11             :     "Rb",
      12             :     "Li",
      13             :     "Na",
      14             :     "K",
      15             :     "Cs",
      16             :     "Sr88_singlet",
      17             :     "Sr88_triplet",
      18             :     "Sr87_mqdt",
      19             :     "Sr88_mqdt",
      20             :     "Yb171_mqdt",
      21             :     "Yb173_mqdt",
      22             :     "Yb174_mqdt",
      23             : ]
      24           1 : SpeciesTypes = Literal["sqdt_duplet", "sqdt_singlet", "sqdt_triplet", "mqdt_halfint", "mqdt_int"]
      25             : 
      26             : 
      27           1 : class DatabaseMissingError(Exception):
      28           1 :     def __init__(self, err: RuntimeError) -> None:
      29           0 :         super().__init__(str(err))
      30           0 :         table = next(w for w in str(err).split(" ") if "states" in w)
      31           0 :         self.species = table.replace("_states", "")
      32             : 
      33             : 
      34           1 : class NoStateFoundError(Exception):
      35           1 :     def __init__(self, err: ValueError) -> None:
      36           1 :         super().__init__(str(err))
      37             : 
      38             : 
      39           1 : def get_custom_error(err: Exception) -> Exception:
      40             :     """Get a custom error message based on the type of error."""
      41           1 :     if isinstance(err, RuntimeError) and "Table" in str(err) and "not found" in str(err):
      42           0 :         return DatabaseMissingError(err)
      43           1 :     if isinstance(err, ValueError) and ("No state found" in str(err) or "quantum number m must be" in str(err)):
      44           1 :         return NoStateFoundError(err)
      45           0 :     return err
      46             : 
      47             : 
      48           1 : @run_in_other_process
      49           1 : def download_databases_mp(species: list[str]) -> None:
      50             :     """Download the databases in a separate process."""
      51           0 :     download_databases(species)
      52             : 
      53             : 
      54           1 : def get_species_type(species: str) -> SpeciesTypes:
      55             :     """Return the species type based on the species name of the ... atom."""
      56           0 :     if "mqdt" in species:
      57           0 :         match = re.search(r"\d+", species)
      58           0 :         if match:
      59           0 :             if int(match.group()) % 2 == 0:
      60           0 :                 return "mqdt_int"
      61           0 :             return "mqdt_halfint"
      62           0 :         raise ValueError(f"Invalid species name: {species}")
      63           0 :     if "singlet" in species:
      64           0 :         return "sqdt_singlet"
      65           0 :     if "triplet" in species:
      66           0 :         return "sqdt_triplet"
      67           0 :     return "sqdt_duplet"