fennol.utils.atomic_units

  1from math import pi
  2import math
  3from typing import Dict, Union, Optional
  4from .unit_parser import parse_unit_string
  5
  6
  7class _AtomicUnits:
  8    """Class to hold atomic unit constants and conversions.
  9    Definition:
 10        - e = 1. (elementary charge)
 11        - hbar = 1. (reduced Planck's constant)
 12        - me = 1. (electron mass in atomic units)
 13        - a0 = 1. (bohr radius)
 14    """
 15
 16    LIGHTSPEED = 137.035999177  # Speed of light in atomic units
 17    FSC = 1.0 / LIGHTSPEED  # Fine structure constant
 18    HBAR = 1.0  # Reduced Planck's constant
 19    PLANCK = 2.0 * pi * HBAR  # Planck's constant (not used directly)
 20    ME = 1.0  # Electron mass in atomic units (AU)
 21    MPROT = 1836.1526734252586  # proton/electron mass ratio
 22    NA = 6.02214129e23  # Avogadro's number
 23    AVOGADRO = NA  # Alias for Avogadro's number
 24    MOL = 1.0 / NA  # number of particules to Mole
 25    EPS0 = 1.0 / (4.0 * pi)  # Vacuum permittivity in atomic units
 26    K_E = 1.0  # Coulomb's constant in atomic units (k_e = 1/(4πε₀) = 1 in atomic units)
 27    COULOMB = K_E  # Alias for backward compatibility
 28    A0 = 1.0  # Bohr radius in atomic units (AU)
 29    K_B = 3.166811563e-6  # Boltzmann constant in atomic units Ha/K (AU)
 30
 31    CONSTANTS = {
 32        "LIGHTSPEED": (LIGHTSPEED, {"L": 1, "T": -1}),
 33        "FSC": (FSC, {}),
 34        "HBAR": (HBAR, {"E": 1, "T": 1}),
 35        "PLANCK": (PLANCK, {"E": 1, "T": 1}),
 36        "ME": (ME, {"M": 1}),
 37        "MPROT": (MPROT, {"M": 1}),
 38        "NA": (NA, {}),
 39        "AVOGADRO": (NA, {}),
 40        "EPS0": (EPS0, {"E": -1, "L": -1}),
 41        "K_E": (K_E, {"E": 1, "L": 1}),
 42        "COULOMB": (COULOMB, {"E": 1, "L": 1}),
 43        "A0": (A0, {"L": 1}),
 44        "K_B": (K_B, {"E": 1}),
 45    }
 46
 47    ### LENGTH UNITS
 48    ANGSTROM = 0.52917721  # Bohr to Angstrom
 49    ANG = ANGSTROM  # Alias for ANGSTROM
 50    NM = ANGSTROM * 1e-1  # 1 nm = 10 Angstroms
 51    CM = ANGSTROM * 1e-8
 52    LENGTH_UNITS = {
 53        "A0": 1.0,
 54        "BOHR": 1.0,
 55        "ANGSTROM": ANGSTROM,
 56        "ANG": ANGSTROM,
 57        "A": ANGSTROM,
 58        "NM": NM,
 59        "CM": CM,
 60    }
 61
 62    ### TIME UNITS
 63    FS = 2.4188843e-2  # AU time to femtoseconds
 64    PS = FS / 1000  # AU time to picoseconds
 65    NS = PS / 1000  # AU time to nanoseconds
 66    TIME_UNITS = {
 67        "AU_T": 1.0,
 68        "FS": FS,
 69        "FEMTOSECONDS": FS,
 70        "PS": PS,
 71        "PICOSECONDS": PS,
 72        "NS": NS,
 73        "NANOSECONDS": NS,
 74    }
 75
 76    ### ENERGY UNITS
 77    EV = 27.211386024367243  # Hartree to eV
 78    KCALPERMOL = 627.5096080305927  # Hartree to kcal/mol
 79    KJPERMOL = 2625.5002  # Hartree to kJ/mol
 80    RY = 2.0  # Hartree to Rydberg
 81    KELVIN = 1.0 / K_B  # Hartree to Kelvin (K/Ha)
 82    KCAL = KCALPERMOL * MOL
 83    KJ = KJPERMOL * MOL
 84    ENERGY_UNITS = {
 85        "HARTREE": 1.0,
 86        "HA": 1.0,
 87        "EV": EV,
 88        "MEV": 1.0e3 * EV,  # milli electronvolt
 89        "KCALPERMOL": KCALPERMOL,
 90        "KJPERMOL": KJPERMOL,
 91        "RY": RY,
 92        "KELVIN": KELVIN,
 93        "KCAL": KCAL,
 94        "KJ": KJ,
 95    }
 96
 97    # MASS UNITS
 98    DA = 1.0 / MPROT  # amu to Dalton (~ g/mol)
 99    GRAM = DA * MOL  # amu to gram
100    MASS_UNITS = {
101        "AU_M": 1.0,
102        "ME": 1.0,
103        "DA": DA,
104        "GRAM": GRAM,
105        "KG": GRAM * 1e-3,
106        "KILOGRAM": GRAM * 1e-3,
107    }
108
109    ### OTHER UNITS
110    DEBYE = 2.541746  # e.Bohr to Debye
111
112    # FREQUENCY UNITS
113    THZ = 1000.0 / FS  # AU frequency to THz
114    CM1 = 219471.52  # AU angular frequency to cm-1 (spectroscopist's wavenumber)
115
116    # PRESSURE UNITS
117    KBAR = 294210.2648438959  # Hartree/bohr**3 to kbar
118    ATM = KBAR * 1000.0 / 1.01325  # Hartree/bohr**3 to atm
119    GPA = 0.1 * KBAR  # Hartree/bohr**3 to GPa
120
121    # FORCE UNITS
122    NNEWTON = 82.387  # Ha/bohr to nNewton
123
124    OTHER_UNITS = {
125        "MOL": (MOL, {}),
126        "DEBYE": (DEBYE, {"L": 1}),
127        "THZ": (THZ, {"T": -1}),
128        "CM-1": (CM1, {"T": -1}),
129        "CM1": (CM1, {"T": -1}),
130        "KBAR": (KBAR, {"L": -3, "E": 1}),
131        "ATM": (ATM, {"L": -3, "E": 1}),
132        "GPA": (GPA, {"L": -3, "E": 1}),
133        "NNEWTON": (NNEWTON, {"L": -1, "E": 1}),
134    }
135
136    def __setattr__(self, name, value):
137        raise AttributeError(
138            f"_AtomicUnits is immutable. Cannot set attribute '{name}'."
139        )
140
141
142class UnitSystem:
143    """Class to handle unit systems and conversions.
144    A unit system is defined by its base units which are a set of THREE units among:
145        - Length (L)
146        - Time (T)
147        - Energy (E)
148        - Mass (M)
149    that are defined at initialization.
150    In all unit systems, we use the elementary charge and Kelvin as base units.
151
152    Other units are derived from these base units.
153
154    Unit multipliers are made accessible as attributes of the class.
155    They correspond to the conversion factors from the unit system to the unit.
156    Example:
157        us = UnitSystem(L='BOHR', T='AU_T', E='HARTREE')
158        print(us.BOHR) # 1.0
159        print(us.ANGSTROM) # 0.52917721
160        print(au.FS) # 2.4188843e-2
161
162        us = UnitSystem(L='ANGSTROM', E='EV',M='DA')
163        print(us.BOHR)  # 1.889725989
164        print(us.fs) # 10.217477
165
166    WARNING: This is the opposite convention as ASE units.
167
168    shorthand notations:
169        - au: Atomic Units
170        - us: Unit/User System
171
172    """
173
174    _initialized = False
175
176    def __init__(
177        self,
178        L: Optional[str] = None,
179        T: Optional[str] = None,
180        E: Optional[str] = None,
181        M: Optional[str] = None,
182    ):
183
184        self.unit_system = {}
185        self.us_to_au = {}
186        if L is not None:
187            L = L.strip().upper()
188            self.unit_system["L"] = L
189            self.us_to_au["L"] = 1.0 / _AtomicUnits.LENGTH_UNITS[L]
190        if T is not None:
191            T = T.strip().upper()
192            self.unit_system["T"] = T
193            self.us_to_au["T"] = 1.0 / _AtomicUnits.TIME_UNITS[T]
194        if E is not None:
195            E = E.strip().upper()
196            self.unit_system["E"] = E
197            self.us_to_au["E"] = 1.0 / _AtomicUnits.ENERGY_UNITS[E]
198        if M is not None:
199            M = M.strip().upper()
200            self.unit_system["M"] = M
201            self.us_to_au["M"] = 1.0 / _AtomicUnits.MASS_UNITS[M]
202        assert (
203            len(self.unit_system) == 3
204        ), f"Unit system must have exactly 3 base units to be consistent. Got: {self.unit_system}"
205
206        # Derive the last base unit from the other ones.
207        if L is None:
208            self.us_to_au["L"] = self.base_us2au_converter(E=0.5, T=1, M=-0.5)
209        elif T is None:
210            self.us_to_au["T"] = self.base_us2au_converter(E=-0.5, L=1, M=0.5)
211        elif E is None:
212            self.us_to_au["E"] = self.base_us2au_converter(L=2, T=-2, M=1)
213        elif M is None:
214            self.us_to_au["M"] = self.base_us2au_converter(E=1, T=2, L=-2)
215
216        self.us_to_unit = {}
217        ## LENGTH UNITS
218        for unit, au_to_unit in _AtomicUnits.LENGTH_UNITS.items():
219            self.us_to_unit[unit.upper()] = self.us_to_au["L"] * au_to_unit
220
221        ## TIME UNITS
222        for unit, au_to_unit in _AtomicUnits.TIME_UNITS.items():
223            self.us_to_unit[unit.upper()] = self.us_to_au["T"] * au_to_unit
224
225        ## ENERGY UNITS
226        for unit, au_to_unit in _AtomicUnits.ENERGY_UNITS.items():
227            self.us_to_unit[unit.upper()] = self.us_to_au["E"] * au_to_unit
228
229        ## MASS UNITS
230        for unit, au_to_unit in _AtomicUnits.MASS_UNITS.items():
231            self.us_to_unit[unit.upper()] = self.us_to_au["M"] * au_to_unit
232
233        ### OTHER UNITS
234        for unit, (au_to_unit, powers) in _AtomicUnits.OTHER_UNITS.items():
235            self.us_to_unit[unit.upper()] = au_to_unit * self.base_us2au_converter(
236                **powers
237            )
238
239        ##CONSTANTS
240        self.CONSTANTS = {}
241        for const_name, (value_au, powers) in _AtomicUnits.CONSTANTS.items():
242            self.CONSTANTS[const_name.upper()] = value_au / self.base_us2au_converter(
243                **powers
244            )
245
246        # ## SET ATTRIBUTES FOR EASY ACCESS
247        for const_name, value in self.CONSTANTS.items():
248            self.__setattr__(const_name, value)
249        for unit, value in self.us_to_unit.items():
250            self.__setattr__(unit, value)
251
252        self.UNITS_LIST = list(self.us_to_unit.keys())
253
254        # Add some common constants for get_multiplier
255        self.us_to_unit["1"] = 1.0
256        self.us_to_unit["2PI"] = 2.0 * pi
257
258        self._initialized = True
259
260    def __setattr__(self, name, value):
261        if not self._initialized:
262            super().__setattr__(name, value)
263        else:
264            raise AttributeError(
265                f"UnitSystem is immutable after initialization. Cannot set attribute '{name}'."
266            )
267
268    def list_units(self):
269        """List all units in the unit system."""
270        return self.UNITS_LIST
271
272    def list_constants(self):
273        """List all constants in the unit system."""
274        return list(self.CONSTANTS.keys())
275
276    def base_us2au_converter(self, **powers: Dict[str, Union[int, float]]):
277        """
278        Get the conversion multiplier from the unit system to atomic units based on powers of base units.
279        """
280        return math.prod(
281            [self.us_to_au[unit.upper()] ** power for unit, power in powers.items()]
282        )
283
284    def base_au2us_converter(self, **powers: Dict[str, Union[int, float]]):
285        """
286        Get the conversion multiplier from atomic units to the unit system based on powers of base units.
287        """
288        return 1.0 / self.base_us2au_converter(**powers)
289
290    def get_multiplier(self, unit_string: str) -> float:
291        """
292        Parse a unit string and return the conversion multiplier from the unit system to that unit.
293
294        Example:
295          us = UnitSystem('BOHR', 'AU_T', 'HA')
296          multiplier = us.get_multiplier('ANGSTROM')
297          print(multiplier)  # Should print 0.5291.. (the bohr radius in Angstroms)
298          multiplier = us.get_multiplier('EV')
299          print(multiplier)  # Should print 27.211386024367243
300
301        Supports syntax like:
302        - Simple units: "EV", "ANGSTROM"
303        - Powers: "ANGSTROM^{2}", "FS^{-1}", "ANGSTROM^2"
304        - Products: "EV*ANGSTROM", "KBAR*FS"
305        - Quotients: "EV/ANGSTROM", "KBAR/FS"
306        - Complex: "EV*ANGSTROM^{2}/FS^{3}"
307        - Floating point powers: "ANGSTROM^{2.5}"
308        - Parentheses: "(EV*ANGSTROM)^2", "EV/(ANGSTROM*FS)"
309        """
310        return parse_unit_string(unit_string, self.us_to_unit)
311
312
313AtomicUnits = au = UnitSystem(L="BOHR", T="AU_T", E="HARTREE")