Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4           1 : from typing import Literal
       5             : 
       6           1 : from pairinteraction import _backend
       7             : 
       8           1 : FloatType = Literal["float32", "float64"]
       9           1 : _FloatTypeDict: dict[FloatType, _backend.FloatType] = {
      10             :     "float32": _backend.FloatType.FLOAT32,
      11             :     "float64": _backend.FloatType.FLOAT64,
      12             : }
      13             : 
      14           1 : OperatorType = Literal[
      15             :     "zero",
      16             :     "energy",
      17             :     "electric_dipole",
      18             :     "electric_quadrupole",
      19             :     "electric_quadrupole_zero",
      20             :     "electric_octupole",
      21             :     "magnetic_dipole",
      22             :     "identity",
      23             :     "arbitrary",
      24             : ]
      25           1 : _OperatorTypeDict: dict[OperatorType, _backend.OperatorType] = {
      26             :     "zero": _backend.OperatorType.ZERO,
      27             :     "energy": _backend.OperatorType.ENERGY,
      28             :     "electric_dipole": _backend.OperatorType.ELECTRIC_DIPOLE,
      29             :     "electric_quadrupole": _backend.OperatorType.ELECTRIC_QUADRUPOLE,
      30             :     "electric_quadrupole_zero": _backend.OperatorType.ELECTRIC_QUADRUPOLE_ZERO,
      31             :     "electric_octupole": _backend.OperatorType.ELECTRIC_OCTUPOLE,
      32             :     "magnetic_dipole": _backend.OperatorType.MAGNETIC_DIPOLE,
      33             :     "identity": _backend.OperatorType.IDENTITY,
      34             :     "arbitrary": _backend.OperatorType.ARBITRARY,
      35             : }
      36             : 
      37           1 : Parity = Literal["even", "odd"]
      38           1 : _ParityToCPP: dict[Parity, _backend.Parity] = {"even": _backend.Parity.EVEN, "odd": _backend.Parity.ODD}
      39           1 : _ParityToInt: dict[Parity, int] = {"even": 1, "odd": -1}
      40             : 
      41             : 
      42           1 : def get_cpp_float_type(float_type: FloatType) -> _backend.FloatType:
      43             :     """Convert a python FloatType string to a cpp FloatType enum."""
      44           1 :     if float_type not in _FloatTypeDict:
      45           0 :         raise ValueError(f"Unknown float_type '{float_type}', should be one of {list(_FloatTypeDict.keys())}")
      46           1 :     return _FloatTypeDict[float_type]
      47             : 
      48             : 
      49           1 : def get_cpp_operator_type(operator_type: OperatorType) -> _backend.OperatorType:
      50             :     """Convert a python OperatorType string to a cpp OperatorType enum."""
      51           1 :     if operator_type not in _OperatorTypeDict:
      52           0 :         raise ValueError(f"Unknown operator_type '{operator_type}', should be one of {list(_OperatorTypeDict.keys())}")
      53           1 :     return _OperatorTypeDict[operator_type]
      54             : 
      55             : 
      56           1 : def get_cpp_parity(parity: Parity) -> _backend.Parity:
      57             :     """Convert a python Parity string to a cpp Parity enum."""
      58           1 :     if parity not in _ParityToCPP:
      59           0 :         raise ValueError(f"Unknown parity '{parity}', should be one of {list(_ParityToCPP.keys())}")
      60           1 :     return _ParityToCPP[parity]
      61             : 
      62             : 
      63           1 : def get_python_parity(parity: _backend.Parity) -> Parity:
      64             :     """Convert a cpp Parity enum to a python Parity string."""
      65           0 :     if parity not in _ParityToCPP.values():
      66           0 :         raise ValueError(f"Unknown parity '{parity}', should be one of {list(_ParityToCPP.values())}")
      67           0 :     return next(k for k, v in _ParityToCPP.items() if v == parity)
      68             : 
      69             : 
      70           1 : def parity_to_int(parity: Parity) -> int:
      71             :     """Convert a python Parity string to its integer value (+1 for even, -1 for odd)."""
      72           0 :     if parity not in _ParityToInt:
      73           0 :         raise ValueError(f"Unknown parity '{parity}', should be one of {list(_ParityToInt.keys())}")
      74           0 :     return _ParityToInt[parity]
      75             : 
      76             : 
      77           1 : def int_to_parity(parity: int) -> Parity:
      78             :     """Convert an integer parity (+1, -1) to a python Parity string."""
      79           1 :     if parity not in _ParityToInt.values():
      80           0 :         raise ValueError(f"Unknown parity '{parity}', should be one of {list(_ParityToInt.values())}")
      81           1 :     return next(k for k, v in _ParityToInt.items() if v == parity)