Source code for couplingfunct

"""
Coupling Calculations Module
----------------------------

This module contains functions for coupling calculations in vibrational coupling Hamiltonians.

Functions
---------
- :func:`linear_coupling`: Computes the linear coupling.

Variables
---------
- ``coupling_funct``: Dictionary of coupling functions.
- ``n_var``: Dictionary of the number of variables for each function.
- ``initial_guesses``: Dictionary of initial guesses for each potential type.
"""

import numpy as np
import tensorflow as tf

# Constants
DEFAULT_INITIAL_GUESS = 1e-3  # Default initial guess for coupling constants

# Coupling Functions
[docs]@tf.function def linear_coupling(q: tf.Tensor, k1: tf.Tensor) -> tf.Tensor: """ Compute the linear coupling function. Parameters ---------- q : tf.Tensor Tensor of displacements with dtype tf.float32. k1 : tf.Tensor or float Tensor representing the coupling constant with dtype tf.float32. Returns ------- tf.Tensor Tensor of linear coupling values with dtype tf.float32. Example ------- >>> q = tf.constant([1.0, 2.0, 3.0], dtype=tf.float32) >>> k1 = tf.constant(0.5, dtype=tf.float32) >>> linear_coupling(q, k1) <tf.Tensor: shape=(3,), dtype=float32, numpy=array([0.5, 1.0, 1.5], dtype=float32)> """ # Ensure inputs are float32 q = tf.cast(q, tf.float32) k1 = tf.cast(k1, tf.float32) # Compute linear coupling using TensorFlow operations return tf.multiply(q, k1)
# List of types of coupling functions COUPLING_TYPES = ["linear"] # Coupling functions dictionary for using them in other modules coupling_funct = { "linear": linear_coupling } # Number of variables for each function n_var = { "linear": 1 } # Initial guesses for each potential type initial_guesses = { "linear": DEFAULT_INITIAL_GUESS }