Source code for diabfunct

"""
This module contains functions for diabatic calculations.

Functions
---------
- :func:`harmonic_oscillator`: Computes the harmonic oscillator potential.
- :func:`general_quartic_potential`: Computes the general quartic potential.
- :func:`general_morse_np`: Computes the general Morse potential using NumPy.
- :func:`general_morse_tf`: Computes the general Morse potential using TensorFlow.

Variables
---------
- ``potential_functions``: Dictionary of potential functions.
- ``n_var``: Dictionary of the number of variables for each function.
- ``initial_guesses``: Dictionary of initial guesses for each potential type.
- ``kappa_compatible``: Dictionary indicating compatibility with the kappa function.
"""

import numpy as np
import tensorflow as tf

# Constants
HALF = tf.constant(0.5, dtype=tf.float32)
ONE = tf.constant(1.0, dtype=tf.float32)
ONE_OVER_24 = tf.constant(1.0 / 24.0, dtype=tf.float32)

# Potential Functions
# @tf.function
[docs]def harmonic_oscillator(q: tf.Tensor, omega: tf.Tensor, params: list) -> tf.Tensor: """ Compute the harmonic oscillator potential. Parameters ---------- q : tf.Tensor Tensor of displacements. omega : tf.Tensor Tensor representing the frequency of the mode. Returns ------- tf.Tensor Tensor of potential values. """ k1 = params[0] q = tf.cast(q, tf.float32) omega = tf.cast(omega, tf.float32) return HALF * omega * tf.math.square(q) + HALF * k1 * tf.math.square(q)
# @tf.function
[docs]def general_quartic_potential(q: tf.Tensor, omega: tf.Tensor, params: list) -> tf.Tensor: """ Compute the general quartic potential. Parameters ---------- q : tf.Tensor Tensor of displacements. omega : tf.Tensor Tensor representing the frequency of the mode. params : list List containing quadratic (k2) and quartic (k3) coefficients. Returns ------- tf.Tensor Tensor of potential values. """ # k2, k3 = params k2, k3 = params[0], params[1] q = tf.cast(q, tf.float32) omega = tf.cast(omega, tf.float32) k2 = tf.cast(k2, tf.float32) k3 = tf.cast(k3, tf.float32) return ( HALF * omega * tf.math.square(q) + HALF * k2 * tf.math.square(q) + ONE_OVER_24 * k3 * tf.math.pow(q, 4) )
[docs]def general_morse_np(q: np.ndarray, De: float, alpha: float, q0: float) -> np.ndarray: """ Compute the general Morse potential using NumPy. Parameters ---------- q : np.ndarray Array of displacements. De : float Dissociation energy. alpha : float Range parameter. q0 : float Equilibrium bond distance. Returns ------- np.ndarray Array of potential values. """ morse = np.exp(-alpha * (q - q0)) - 1.0 return De * morse ** 2
# @tf.function
[docs]def general_morse_tf(q: tf.Tensor, params: list) -> tf.Tensor: """ Compute the general Morse potential using TensorFlow. Parameters ---------- q : tf.Tensor Tensor of displacements. params : list List containing dissociation energy (De), range parameter (alpha), and equilibrium bond distance (q0). Returns ------- tf.Tensor Tensor of potential values. """ De, alpha, q0 = params[0], params[1], params[2] q = tf.cast(q, tf.float32) De = tf.cast(De, tf.float32) alpha = tf.cast(alpha, tf.float32) q0 = tf.cast(q0, tf.float32) morse = tf.exp(-alpha * (q - q0)) - ONE return De * tf.math.square(morse)
# Potential Functions Dictionary potential_functions = { "ho": harmonic_oscillator, "morse": general_morse_tf, "antimorse": general_morse_tf, "quartic": general_quartic_potential, } # Number of Variables for Each Function n_var = { "ho": 1, "morse": 3, "antimorse": 3, "quartic": 2, } # Initial Guesses for Each Potential Type initial_guesses = { "ho": np.array([0.5]), "morse": np.array([5.0, 0.5, 0.0]), "antimorse": np.array([5.0, -0.5, 0.0]), "quartic": np.array([0.0, 0.0005]), } # Kappa Compatibility kappa_compatible = { "ho": True, "morse": False, # Morse potential has in its definition q - q0 "antimorse": False, # Morse potential has in its definition q - q0 "quartic": True, }