Module tau2.rates_calculator

Functions

def calculate_detailed_balance(gamma, energies, T)

Expand source code

def calculate_detailed_balance(gamma, energies, T):
    """
    Calculates the detailed balance check for a given rate matrix.
    
    Args:
        gamma (np.ndarray): The (N, N) transition rate matrix.
                            gamma[f, i] is the rate from state i to state f.
        energies (np.ndarray): A (N,) array of state energies in cm^-1.
        T (float): The temperature in Kelvin.

    Returns:
        tuple:
            - det_balance_matrix (np.ndarray): The (N, N) matrix of
                                               absolute differences.
            - detailed_value_sum (float): The scalar sum of all unique pair
                                          differences, used for the check.
    """    
    
    # Suppress division by zero warnings for calculating the rate ratio
    with np.errstate(divide='ignore', invalid='ignore'):
        
        # Calculate the Boltzmann factor for all pairs (f, i)
        # energy_diffs[f, i] = E_i - E_f
        energy_diffs = energies - energies[:, np.newaxis]
        # Clip the argument to np.exp to avoid overflow, which happens for large
        # positive energy differences (large E_i, small E_f) at low T.
        # A value of 700 is safe, as exp(709) is near the float64 limit.
        clipped_arg = np.clip(energy_diffs / (KB_INV_CM * T), a_min=None, a_max=700)
        boltzmann_matrix = np.exp(clipped_arg)

        # Calculate the ratio of forward (i->f) to reverse (f->i) rates
        # gamma_ratio[f, i] = gamma[f, i] / gamma[i, f]
        gamma_ratio = np.divide(gamma, gamma.T, out=np.full_like(gamma, np.inf), 
                            where=(gamma.T != 0))

        # Calculate the absolute difference from the expected Boltzmann ratio
        det_balance_matrix = np.abs(gamma_ratio - boltzmann_matrix)

        # Correct for the case where both forward and reverse rates are zero
        # In this case, the ratio is effectively 1
        both_zero_mask = (gamma == 0.0) & (gamma.T == 0.0)
        det_balance_matrix[both_zero_mask] = np.abs(1.0 - boltzmann_matrix[both_zero_mask])

    # The detailed balance value is the sum of differences for all unique
    # pairs (i, f). We sum the lower triangle (where f > i) to get this.
    detailed_value_sum = np.sum(np.tril(det_balance_matrix, k=-1))

    # Check if the total sum exceeds the tolerance
    if detailed_value_sum > 1:
        print(f"!!! WARNING !!!: Detailed balance not obeyed at {T} K. Sum = {detailed_value_sum}")

    return det_balance_matrix, detailed_value_sum

Calculates the detailed balance check for a given rate matrix.

Args

gamma : np.ndarray

The (N, N) transition rate matrix. gamma[f, i] is the rate from state i to state f.

energies : np.ndarray

A (N,) array of state energies in cm^-1.

T : float

The temperature in Kelvin.

Returns

tuple: - det_balance_matrix (np.ndarray): The (N, N) matrix of absolute differences. - detailed_value_sum (float): The scalar sum of all unique pair differences, used for the check.

def calculate_rates_from_data(data, temperatures, max_state, rate_integrator='trapz')

Expand source code

def calculate_rates_from_data(data, temperatures, max_state, rate_integrator='trapz'):
    """
    Calculates relaxation rates from pre-computed Raman spectral density data.

    This function takes the spectral density integrands (J(omega)) for various
    Raman processes and integrates them over phonon energy (omega) to obtain
    the transition rates (Gamma_if). It then constructs the full rate matrix
    (Gamma) and extracts the final relaxation rate from its eigenvalues.

    Args:
        data (dict): A dictionary containing:
            - 'omega_grid' (np.ndarray): The phonon energy grid (cm^-1).
            - 'energies' (np.ndarray): Zeeman-split electronic energies (cm^-1).
            - 'integrands' (dict): A nested dictionary where keys are (i, f)
              transition tuples and values are dictionaries of process
              integrands ('pp', 'mm', 'pm', 'mp').
        temperatures (list or np.ndarray): List of temperatures (K) at which
                                           to calculate the rates.
        max_state (int): The maximum electronic state index to consider for
                         the Gamma matrix (states 0 to max_state-1).
        rate_integrator (str, optional): The numerical integration method.
                                         Supports 'trapz' (default) and 'simpson'.

    Returns:
        tuple: A tuple containing:
            - np.ndarray: An array of final relaxation rates (in s^-1) for each
                          temperature (typically the second smallest eigenvalue).
            - dict: A dictionary of the full Gamma matrices for each temperature,
                    with temperatures as keys.
    """
    omega_grid = data['omega_grid']
    energies = data['energies']
    
    # 1. Integrate spectral densities to get rates for each T
    all_transition_rates = {}
    
    # Filter transitions based on keys present in data['integrands']
    # We assume 'data' contains only the transitions requested.
    transitions = list(data['integrands'].keys())
    
    # Identify unique states involved in the transitions to map them to matrix indices
    involved_states = sorted(list(set([idx for pair in transitions for idx in pair])))
    
    # Create a mapping from state index -> matrix index (0, 1, 2...)
    # If the user requested states [2, 3], involved_states is [2, 3].
    # map: 2->0, 3->1.
    state_to_matrix_idx = {state_idx: i for i, state_idx in enumerate(involved_states)}
    matrix_dim = max(max_state, len(involved_states))

    for (i, f), integrands_dict in data['integrands'].items():
        delta_E = energies[f] - energies[i]
        rates_T = []

        # --- Direct Integration Methods ---
        for T in temperatures:
            if T <= 0:
                rates_T.append(0.0)
                continue
            
            n1 = get_thermal_pop(omega_grid, T)
            temp_dep_integrand = np.zeros_like(omega_grid)
            
            # Term 'pp' (emission-emission)
            if 'pp' in integrands_dict:
                omega_2 = -delta_E - omega_grid
                n2 = get_thermal_pop(omega_2, T)
                thermal_factor = (n1 + 1.0) * (n2 + 1.0)
                temp_dep_integrand += integrands_dict['pp'] * thermal_factor

            # Term 'mm' (absorption-absorption)
            if 'mm' in integrands_dict:
                omega_2 = delta_E - omega_grid
                n2 = get_thermal_pop(omega_2, T)
                thermal_factor = n1 * n2 
                temp_dep_integrand += integrands_dict['mm'] * thermal_factor

            # Term 'pm' (absorption-emission)
            if 'pm' in integrands_dict:
                omega_2 = delta_E + omega_grid
                n2 = get_thermal_pop(omega_2, T)
                thermal_factor = (n1 + 1.0) * n2
                temp_dep_integrand += integrands_dict['pm'] * thermal_factor
            
            # Term 'mp' (emission-absorption)
            if 'mp' in integrands_dict:
                omega_2 = -delta_E + omega_grid
                n2 = get_thermal_pop(omega_2, T)
                thermal_factor = n1 * (n2 + 1.0)
                temp_dep_integrand += integrands_dict['mp'] * thermal_factor

            if rate_integrator == 'trapz':
                total_integral = trapz(temp_dep_integrand, omega_grid)
            elif rate_integrator == 'simpson':
                total_integral = simpson(temp_dep_integrand, x=omega_grid)
            else:
                sys.exit(f"Error: rate_integrator '{rate_integrator}' is not supported. "
                         "Only 'trapz' and 'simpson' are currently implemented.")
                
            rates_T.append(total_integral * TWO_PI_OVER_HBAR)

        all_transition_rates[(i, f)] = np.array(rates_T)

    rate_values = []
    gamma_matrices = {}
    det_balance_matrices = {}
    det_values = {}

    for t_idx, T in enumerate(temperatures):
        gamma = np.zeros((matrix_dim, matrix_dim))

        for (i, f), rates_at_temps in all_transition_rates.items():
            # Use mapped indices if available, otherwise original indices
            # If state indices (e.g. 2, 3) are outside matrix_dim (e.g. 2), map them.
            idx_i = state_to_matrix_idx.get(i, i)
            idx_f = state_to_matrix_idx.get(f, f)
            
            # Safety check to keep within bounds
            if idx_i < matrix_dim and idx_f < matrix_dim:
                gamma[idx_f, idx_i] = rates_at_temps[t_idx]

        np.fill_diagonal(gamma, -np.sum(gamma, axis=0))
        gamma_matrices[T] = gamma

        eigs = np.sort(np.abs(np.real(np.linalg.eigvals(gamma))))
        # The first eigenvalue is always 0 (equilibrium). The second is the relaxation rate.
        rate_val = eigs[1] if len(eigs) > 1 else 0.0
        rate_values.append(rate_val)
        
        # Detailed balance check requires energies for the specific states in the matrix
        # Map indices back to original energy indices
        matrix_energies = np.zeros(matrix_dim)
        for state_idx, matrix_idx in state_to_matrix_idx.items():
            if matrix_idx < matrix_dim:
                matrix_energies[matrix_idx] = energies[state_idx]
                
        det_balance_matrices[T], det_values[T] = calculate_detailed_balance(gamma, matrix_energies, T)

    return np.array(rate_values), gamma_matrices, det_balance_matrices

Calculates relaxation rates from pre-computed Raman spectral density data.

This function takes the spectral density integrands (J(omega)) for various Raman processes and integrates them over phonon energy (omega) to obtain the transition rates (Gamma_if). It then constructs the full rate matrix (Gamma) and extracts the final relaxation rate from its eigenvalues.

Args

data : dict

A dictionary containing: - 'omega_grid' (np.ndarray): The phonon energy grid (cm^-1). - 'energies' (np.ndarray): Zeeman-split electronic energies (cm^-1). - 'integrands' (dict): A nested dictionary where keys are (i, f) transition tuples and values are dictionaries of process integrands ('pp', 'mm', 'pm', 'mp').

temperatures : list or np.ndarray

List of temperatures (K) at which to calculate the rates.

max_state : int

The maximum electronic state index to consider for the Gamma matrix (states 0 to max_state-1).

rate_integrator : str, optional

The numerical integration method. Supports 'trapz' (default) and 'simpson'.

Returns

tuple

A tuple containing: - np.ndarray: An array of final relaxation rates (in s^-1) for each temperature (typically the second smallest eigenvalue). - dict: A dictionary of the full Gamma matrices for each temperature, with temperatures as keys.

def get_thermal_pop(omega, T, omega_cutoff=0.001)

Expand source code

def get_thermal_pop(omega, T, omega_cutoff=0.001):
    """
    Computes the Bose-Einstein thermal population factor for phonons.

    This function calculates n(omega) = 1 / (exp(omega / (kB*T)) - 1),
    handling potential overflow for very large omega/T values and ensuring
    correct behavior for omega <= 0.

    Args:
        omega (float or np.ndarray): Phonon energy (in cm^-1).
        T (float): Temperature (in Kelvin).
        omega_cutoff (float, optional): A small positive value to avoid
                                        singularity when omega is very close to zero.
                                        Defaults to 0.001 cm^-1.

    Returns:
        float or np.ndarray: The Bose-Einstein thermal population factor(s).
    """
    is_scalar = np.isscalar(omega)
    omega_arr = np.atleast_1d(omega)

    # Default result is 0. This is the correct limit for omega <= 0
    # and for the high-energy/low-temp case that causes overflow.
    pop = np.zeros_like(omega_arr, dtype=float)

    # Only perform calculations for positive energies.
    mask = omega_arr > 0
    if np.any(mask):
        # Apply cutoff to avoid singularity near zero.
        omega_pos = np.maximum(omega_arr[mask], omega_cutoff)
        arg = omega_pos / (KB_INV_CM * T)

        # Temporarily ignore benign overflow warnings.
        # The result of 1.0 / np.expm1(inf) is correctly 0.0.
        with np.errstate(over='ignore'):
            pop[mask] = 1.0 / np.expm1(arg)

    return pop.item() if is_scalar else pop

Computes the Bose-Einstein thermal population factor for phonons.

This function calculates n(omega) = 1 / (exp(omega / (kB*T)) - 1), handling potential overflow for very large omega/T values and ensuring correct behavior for omega <= 0.

Args

omega : float or np.ndarray

Phonon energy (in cm^-1).

T : float

Temperature (in Kelvin).

omega_cutoff : float, optional

A small positive value to avoid singularity when omega is very close to zero. Defaults to 0.001 cm^-1.

Returns

float or np.ndarray

The Bose-Einstein thermal population factor(s).