charged-shells/charged_shells/interactions.py

from charged_shells import expansion, parameters
import charged_shells.functions as fn
from py3nj import wigner3j
import numpy as np
from typing import Literal
import charged_shells.units_and_constants as uc

Array = np.ndarray
Expansion = expansion.Expansion
ModelParams = parameters.ModelParams
EnergyUnit = Literal['kT', 'eV', 'J']


def energy_units(units: EnergyUnit, params: ModelParams) -> float:
    match units:
        case 'eV':
            # return 1 / (uc.CONSTANTS.e0 * uc.UNITS.voltage)
            return 1.
        case 'kT':
            return 1 / (params.temperature * uc.CONSTANTS.Boltzmann)
        case 'J':
            return uc.UNITS.energy


def charged_shell_energy(ex1: Expansion, ex2: Expansion, params: ModelParams, dist: float = 2, units: EnergyUnit = 'kT',
                         chunk_size: int = 1000):

    ex1, ex2 = expansion.expansions_to_common_l(ex1, ex2)
    dist = dist * params.R

    full_l_array, full_m_array = ex1.lm_arrays

    # determine indices of relevant elements in the sum
    indices_l, indices_p = np.nonzero(full_m_array[:, None] == full_m_array[None, :])
    flat_l = full_l_array[indices_l]
    flat_p = full_l_array[indices_p]
    flat_m = full_m_array[indices_l]  # the same as full_m_array[indices_p]

    relevant_pairs, = np.nonzero(flat_l >= flat_p)
    flat_l = flat_l[relevant_pairs]
    flat_p = flat_p[relevant_pairs]
    flat_m = flat_m[relevant_pairs]
    indices_l = indices_l[relevant_pairs]
    indices_p = indices_p[relevant_pairs]

    charge_factor = np.real(ex1.coefs[..., indices_l] * np.conj(ex2.coefs[..., indices_p]) +
                            (-1) ** (flat_l + flat_p) * ex1.coefs[..., indices_p] * np.conj(ex2.coefs[..., indices_l]))

    all_s_array = np.arange(2 * ex1.max_l + 1)
    bessels = fn.sph_bessel_k(all_s_array, params.kappa * dist)

    # additional selection that excludes terms where Wigner 3j symbols are 0 by definition
    s_bool1 = np.abs(flat_l[:, None] - all_s_array[None, :]) <= flat_p[:, None]
    s_bool2 = flat_p[:, None] <= (flat_l[:, None] + all_s_array[None, :])
    indices_lpm_all, indices_s_all = np.nonzero(s_bool1 * s_bool2)

    # indices array can get really large (a lot of combinations) so we split the calculation into chunks to preserve RAM
    # interestingly, this also leads to performance improvements if chunks are still large enough
    if chunk_size is None:
        chunk_size = len(indices_lpm_all)
    num_sections = np.ceil(len(indices_lpm_all) / chunk_size)
    energy = 0
    for indices_lpm, indices_s in zip(np.array_split(indices_lpm_all, num_sections),
                                      np.array_split(indices_s_all, num_sections)):

        l_vals = flat_l[indices_lpm]
        p_vals = flat_p[indices_lpm]
        m_vals = flat_m[indices_lpm]
        s_vals = all_s_array[indices_s]
        bessel_vals = bessels[indices_s]

        # While all other arrays are 1D, this one can have extra leading axes corresponding to leading dimensions
        # of expansion coefficients. The last dimension is the same as other arrays.
        charge_vals = charge_factor[..., indices_lpm]

        lps_terms = (2 * s_vals + 1) * np.sqrt((2 * l_vals + 1) * (2 * p_vals + 1))

        # the same combination of l, s, and p is repeated many times
        _, unique_indices1, inverse1 = np.unique(np.stack((l_vals, s_vals, p_vals)),
                                                 axis=1, return_inverse=True, return_index=True)
        wigner1 = wigner3j(2 * l_vals[unique_indices1], 2 * s_vals[unique_indices1], 2 * p_vals[unique_indices1],
                           0, 0, 0)[inverse1]

        # all the combinations (l, s, p, m) are unique
        wigner2 = wigner3j(2 * l_vals, 2 * s_vals, 2 * p_vals,
                           -2 * m_vals, 0, 2 * m_vals)

        constants = params.R ** 2 / (params.kappa * params.epsilon * uc.CONSTANTS.epsilon0) * energy_units(units, params)

        C_vals = fn.interaction_coef_C(l_vals, p_vals, params.kappaR)
        lspm_vals = C_vals * (-1) ** (l_vals + m_vals) * lps_terms * bessel_vals * wigner1 * wigner2
        broadcasted_lspm_vals = np.broadcast_to(lspm_vals, charge_vals.shape)

        rescale_at_equal_lp = np.where(l_vals == p_vals, 0.5, 1)

        energy +=  constants * np.sum(rescale_at_equal_lp * broadcasted_lspm_vals * charge_vals, axis=-1)

    return energy