|
@@ -1,71 +1,49 @@
|
|
|
import expansion
|
|
|
import parameters
|
|
|
import functions as fn
|
|
|
-import numpy as np
|
|
|
from py3nj import wigner3j
|
|
|
+import numpy as np
|
|
|
import time
|
|
|
-import matplotlib.pyplot as plt
|
|
|
+from typing import Literal
|
|
|
+import units_and_constants as uc
|
|
|
|
|
|
|
|
|
Array = np.ndarray
|
|
|
Expansion = expansion.Expansion
|
|
|
ModelParams = parameters.ModelParams
|
|
|
+EnergyUnit = Literal['kT', 'eV', 'J']
|
|
|
|
|
|
|
|
|
-# def prefactor(R: float, kappaR: float, c0: float):
|
|
|
-# return R * kappaR * 1e4 / (12.04 * c0)
|
|
|
-#
|
|
|
-#
|
|
|
-# def c0(R: float, kappaR: float):
|
|
|
-# return 10 * kappaR ** 2 / (0.329 ** 2 * R ** 2)
|
|
|
-
|
|
|
-
|
|
|
-def expansions_to_common_l(ex1: Expansion, ex2: expansion) -> (Expansion, Expansion):
|
|
|
- common_l_array = np.union1d(ex1.l_array, ex2.l_array)
|
|
|
- missing_l1 = np.setdiff1d(common_l_array, ex1.l_array, assume_unique=True)
|
|
|
- missing_l2 = np.setdiff1d(common_l_array, ex2.l_array, assume_unique=True)
|
|
|
-
|
|
|
- fill_1 = np.zeros(np.sum(2 * missing_l1 + 1))
|
|
|
- fill_2 = np.zeros(np.sum(2 * missing_l2 + 1))
|
|
|
-
|
|
|
- full_l_array1, _ = ex1.lm_arrays
|
|
|
- full_l_array2, _ = ex2.lm_arrays
|
|
|
-
|
|
|
- # we search for where to place missing coeffs with the help of a boolean array and argmax function
|
|
|
- bool1 = (full_l_array1[:, None] - missing_l1[None, :]) > 0
|
|
|
- bool2 = (full_l_array2[:, None] - missing_l2[None, :]) > 0
|
|
|
-
|
|
|
- # if all elements of bool are false, i.e. missing_l > all existing_l, additional l values come to the end
|
|
|
- indices1 = np.where(np.any(bool1, axis=0), np.argmax(bool1, axis=0), full_l_array1.shape[0])
|
|
|
- indices2 = np.where(np.any(bool2, axis=0), np.argmax(bool2, axis=0), full_l_array2.shape[0])
|
|
|
-
|
|
|
- new_coeffs1 = np.insert(ex1.coeffs, np.repeat(indices1, 2 * missing_l1 + 1), fill_1)
|
|
|
- new_coeffs2 = np.insert(ex2.coeffs, np.repeat(indices2, 2 * missing_l2 + 1), fill_2)
|
|
|
+def energy_units(units: EnergyUnit, params: ModelParams) -> float:
|
|
|
+ match units:
|
|
|
+ case 'eV':
|
|
|
+ return 1 / (uc.CONSTANTS.e0 * uc.UNITS.voltage)
|
|
|
+ case 'kT':
|
|
|
+ return 1 / (params.temperature * uc.CONSTANTS.Boltzmann)
|
|
|
+ case 'J':
|
|
|
+ return uc.UNITS.energy
|
|
|
|
|
|
- return Expansion(common_l_array, new_coeffs1), Expansion(common_l_array, new_coeffs2)
|
|
|
|
|
|
+def charged_shell_energy(ex1: Expansion, ex2: Expansion, dist: float, params: ModelParams, units: EnergyUnit = 'kT'):
|
|
|
|
|
|
-def charged_shell_energy(ex1: Expansion, ex2: Expansion, dist: float, params: ModelParams):
|
|
|
-
|
|
|
- ex1, ex2 = expansions_to_common_l(ex1, ex2)
|
|
|
+ ex1, ex2 = expansion.expansions_to_common_l(ex1, ex2)
|
|
|
+ dist = dist * params.R
|
|
|
|
|
|
full_l_array, full_m_array = ex1.lm_arrays
|
|
|
|
|
|
- coefficient_C = fn.interaction_coeff_C(ex1.l_array[:, None], ex2.l_array[None, :], params.kappaR)
|
|
|
- full_coefficient_C = ex1.repeat_over_m(ex2.repeat_over_m(coefficient_C, axis=1), axis=0)
|
|
|
-
|
|
|
- charge_factor = np.real(ex1.coeffs[:, None] * np.conj(ex2.coeffs[None, :]) +
|
|
|
- (-1) ** (full_l_array[:, None] + full_l_array[None, :])
|
|
|
- * ex1.coeffs[None, :] * ex1.coeffs[:, None])
|
|
|
-
|
|
|
+ # 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]
|
|
|
|
|
|
+ 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, indices_s = np.nonzero(s_bool1 * s_bool2)
|
|
@@ -73,27 +51,39 @@ def charged_shell_energy(ex1: Expansion, ex2: Expansion, dist: float, params: Mo
|
|
|
l_vals = flat_l[indices_lpm]
|
|
|
p_vals = flat_p[indices_lpm]
|
|
|
m_vals = flat_m[indices_lpm]
|
|
|
- C_vals = full_coefficient_C[indices_l, indices_p][indices_lpm]
|
|
|
- charge_vals = charge_factor[indices_l, indices_p][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))
|
|
|
|
|
|
- wigner1 = wigner3j(l_vals, s_vals, p_vals,
|
|
|
- 0, 0, 0)
|
|
|
- wigner2 = wigner3j(l_vals, s_vals, p_vals,
|
|
|
- -m_vals, 0, m_vals)
|
|
|
+ # 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_coeff_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)
|
|
|
|
|
|
- return (params.R ** 2 / (params.kappa * params.epsilon * params.epsilon0)
|
|
|
- * np.sum(C_vals * (-1) ** (l_vals + m_vals) * charge_vals * lps_terms * bessel_vals * wigner1 * wigner2))
|
|
|
+ return 0.5 * constants * np.sum(broadcasted_lspm_vals * charge_vals, axis=-1)
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
|
- params = ModelParams(1, 3, 1, 1)
|
|
|
- ex1 = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=5)
|
|
|
- ex2 = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=5)
|
|
|
+ params = ModelParams(R=150, kappaR=3)
|
|
|
+ ex1 = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=20)
|
|
|
+ ex2 = ex1.clone()
|
|
|
|
|
|
dist = 2.
|
|
|
|