|
|
@@ -2,22 +2,15 @@ from __future__ import annotations
|
|
|
import numpy as np
|
|
|
from dataclasses import dataclass
|
|
|
from functools import cached_property
|
|
|
-import functions as fn
|
|
|
+import charged_shells.functions as fn
|
|
|
import quaternionic
|
|
|
import spherical
|
|
|
-import time
|
|
|
import copy
|
|
|
-import matplotlib.pyplot as plt
|
|
|
-from typing import Callable, TypeVar
|
|
|
from scipy.special import eval_legendre
|
|
|
-import plotly.graph_objects as go
|
|
|
-from pathlib import Path
|
|
|
|
|
|
|
|
|
Array = np.ndarray
|
|
|
Quaternion = quaternionic.array
|
|
|
-T = TypeVar('T')
|
|
|
-V = TypeVar('V')
|
|
|
|
|
|
|
|
|
class InvalidExpansion(Exception):
|
|
|
@@ -65,7 +58,7 @@ class Expansion:
|
|
|
@cached_property
|
|
|
def lm_arrays(self) -> (Array, Array):
|
|
|
"""Return l and m arrays containing all (l, m) pairs."""
|
|
|
- return full_fm_arrays(self.l_array)
|
|
|
+ return full_lm_arrays(self.l_array)
|
|
|
|
|
|
def repeat_over_m(self, arr: Array, axis=0) -> Array:
|
|
|
if not arr.shape[axis] == len(self.l_array):
|
|
|
@@ -108,24 +101,49 @@ class Expansion24(Expansion):
|
|
|
|
|
|
|
|
|
class MappedExpansionQuad(Expansion):
|
|
|
-
|
|
|
- def __init__(self, a_bar: Array | float, kappaR: Array | float, sigma_m: float, l_max: int = 20, sigma0: float = 0):
|
|
|
+ """Expansion that matches the outside potential of a quadrupolar impermeable particle with point charges inside."""
|
|
|
+
|
|
|
+ def __init__(self,
|
|
|
+ a_bar: Array | float,
|
|
|
+ kappaR: Array | float,
|
|
|
+ sigma_tilde: float,
|
|
|
+ l_max: int = 20,
|
|
|
+ sigma0: float | Array = 0):
|
|
|
+ """
|
|
|
+ :param a_bar: distance between the center and off center charges
|
|
|
+ :param kappaR: screening parameter
|
|
|
+ :param sigma_tilde: magnitude of off-center charges / 4pi R^2
|
|
|
+ :param l_max: maximal ell value for the expansion
|
|
|
+ :param sigma0: total (mean) charge density
|
|
|
+ """
|
|
|
a_bar, kappaR = np.broadcast_arrays(a_bar, kappaR)
|
|
|
|
|
|
l_array = np.array([l for l in range(l_max + 1) if l % 2 == 0])
|
|
|
- a_bar, kappaR, l_array_expanded = np.broadcast_arrays(a_bar[..., None], kappaR[..., None], l_array[None, :])
|
|
|
+ a_bar, kappaR, l_array_expanded = np.broadcast_arrays(a_bar[..., None],
|
|
|
+ kappaR[..., None],
|
|
|
+ l_array[None, :])
|
|
|
|
|
|
- coefs = (2 * sigma_m * fn.coef_C_diff(l_array_expanded, kappaR)
|
|
|
- * np.sqrt(4 * np.pi * (2 * l_array_expanded + 1)) * np.power(a_bar, l_array_expanded))
|
|
|
- coefs[..., 0] = sigma0 / np.sqrt(4 * np.pi)
|
|
|
+ coefs = (2 * sigma_tilde * fn.coef_C_diff(l_array_expanded, kappaR)
|
|
|
+ * np.sqrt(4 * np.pi * (2 * l_array_expanded + 1)) * np.power(a_bar, l_array_expanded))
|
|
|
coefs = np.squeeze(rot_sym_expansion(l_array, coefs))
|
|
|
+ coefs = expansion_total_charge(coefs, sigma0)
|
|
|
super().__init__(l_array, coefs)
|
|
|
|
|
|
|
|
|
class GaussianCharges(Expansion):
|
|
|
+ """Expansion for a collection of smeared charges on the sphere."""
|
|
|
|
|
|
def __init__(self, omega_k: Array, lambda_k: Array | float, sigma1: float, l_max: int,
|
|
|
- sigma0: float = 0, equal_charges: bool = True):
|
|
|
+ sigma0: float | Array = 0, equal_charges: bool = True):
|
|
|
+ """
|
|
|
+ :param omega_k: array of positions (theta, phi) of all charges
|
|
|
+ :param lambda_k: smear parameter for each charge or smear for different cases (if equal_charges = True)
|
|
|
+ :param sigma1: scaling
|
|
|
+ :param l_max: maximal ell value for the expansion
|
|
|
+ :param sigma0: total (mean) charge density
|
|
|
+ :param equal_charges: if this is False, length of lambda_k should be N. If True, theta0_k array will be treated
|
|
|
+ as different expansion cases
|
|
|
+ """
|
|
|
omega_k = omega_k.reshape(-1, 2)
|
|
|
if not isinstance(lambda_k, Array):
|
|
|
lambda_k = np.array([lambda_k])
|
|
|
@@ -137,7 +155,7 @@ class GaussianCharges(Expansion):
|
|
|
raise ValueError("Number of charges (length of omega_k) should match the last dimension of lambda_k array.")
|
|
|
lambda_k = lambda_k.reshape(-1, omega_k.shape[0])
|
|
|
l_array = np.arange(l_max + 1)
|
|
|
- full_l_array, full_m_array = full_fm_arrays(l_array)
|
|
|
+ full_l_array, full_m_array = full_lm_arrays(l_array)
|
|
|
theta_k = omega_k[:, 0]
|
|
|
phi_k = omega_k[:, 1]
|
|
|
summands = (lambda_k[:, None, :] / np.sinh(lambda_k[:, None, :])
|
|
|
@@ -145,15 +163,25 @@ class GaussianCharges(Expansion):
|
|
|
* np.conj(fn.sph_harm(full_l_array[None, :, None], full_m_array[None, :, None],
|
|
|
theta_k[None, None, :], phi_k[None, None, :])))
|
|
|
coefs = np.squeeze(4 * np.pi * sigma1 * np.sum(summands, axis=-1))
|
|
|
- coefs[..., 0] = sigma0 / np.sqrt(4 * np.pi)
|
|
|
+ coefs = expansion_total_charge(coefs, sigma0)
|
|
|
l_array, coefs = purge_unneeded_l(l_array, coefs)
|
|
|
super().__init__(l_array, coefs)
|
|
|
|
|
|
|
|
|
class SphericalCap(Expansion):
|
|
|
+ """Expansion for a collection of spherical caps."""
|
|
|
|
|
|
- def __init__(self, omega_k: Array, theta0_k: Array | float, sigma1: float, l_max: int, sigma0: float = 0,
|
|
|
+ def __init__(self, omega_k: Array, theta0_k: Array | float, sigma1: float, l_max: int, sigma0: float | Array = 0,
|
|
|
equal_sizes: bool = True):
|
|
|
+ """
|
|
|
+ :param omega_k: array of positions (theta, phi) of all spherical caps
|
|
|
+ :param theta0_k: sizes of each spherical caps or cap sizes for different cases (if equal_sizes = True)
|
|
|
+ :param sigma1: charge magnitude for the single cap, currently assumes that this is equal for all caps
|
|
|
+ :param l_max: maximal ell value for the expansion
|
|
|
+ :param sigma0: total (mean) charge density
|
|
|
+ :param equal_sizes: if this is False, length of theta0_k should be N. If True, theta0_k array will be treated as
|
|
|
+ different expansion cases
|
|
|
+ """
|
|
|
omega_k = omega_k.reshape(-1, 2)
|
|
|
if not isinstance(theta0_k, Array):
|
|
|
theta0_k = np.array(theta0_k)
|
|
|
@@ -179,29 +207,12 @@ class SphericalCap(Expansion):
|
|
|
# Rotating is implemented in such a way that it rotates every patch to every position,
|
|
|
# hence the redundancy of out of diagonal elements.
|
|
|
coefs_all = np.sum(np.diagonal(coefs_all_single_caps), axis=-1)
|
|
|
- if sigma0 is not None:
|
|
|
- coefs_all[..., 0] = sigma0 / np.sqrt(4 * np.pi)
|
|
|
+ coefs_all = expansion_total_charge(coefs_all, sigma0)
|
|
|
super().__init__(l_array, np.squeeze(coefs_all))
|
|
|
|
|
|
|
|
|
-def map_over_expansion(f: Callable[[Expansion, T], V]) -> Callable[[Expansion, T], V]:
|
|
|
- """Map a function f over the leading axes of an expansion. Uses for loops, so it is kinda slow."""
|
|
|
-
|
|
|
- def mapped_f(ex: Expansion, *args, **kwargs):
|
|
|
- og_shape = ex.shape
|
|
|
- flat_ex = ex.flatten()
|
|
|
- results = []
|
|
|
- for i in range(int(np.prod(og_shape))):
|
|
|
- results.append(f(flat_ex[i], *args, **kwargs))
|
|
|
- try:
|
|
|
- return np.array(results).reshape(og_shape + results[0].shape)
|
|
|
- except AttributeError:
|
|
|
- return np.array(results).reshape(og_shape)
|
|
|
-
|
|
|
- return mapped_f
|
|
|
-
|
|
|
-
|
|
|
-def full_fm_arrays(l_array: Array) -> (Array, Array):
|
|
|
+def full_lm_arrays(l_array: Array) -> (Array, Array):
|
|
|
+ """From an array of l_values get arrays of ell and m that give you all pairs (ell, m)."""
|
|
|
all_m_list = []
|
|
|
for l in l_array:
|
|
|
for i in range(2 * l + 1):
|
|
|
@@ -216,22 +227,30 @@ def rot_sym_expansion(l_array: Array, coefs: Array) -> Array:
|
|
|
return full_coefs
|
|
|
|
|
|
|
|
|
-def coefs_fill_missing_l(expansion: Expansion, target_l_array: Array) -> Expansion:
|
|
|
- missing_l = np.setdiff1d(target_l_array, expansion.l_array, assume_unique=True)
|
|
|
- fill = np.zeros(np.sum(2 * missing_l + 1))
|
|
|
- full_l_array1, _ = expansion.lm_arrays
|
|
|
- # we search for where to place missing coefs with the help of a boolean array and argmax function
|
|
|
- comparison_bool = (full_l_array1[:, None] - missing_l[None, :]) > 0
|
|
|
- indices = np.where(np.any(comparison_bool, axis=0), np.argmax(comparison_bool, axis=0), full_l_array1.shape[0])
|
|
|
- new_coefs = np.insert(expansion.coefs, np.repeat(indices, 2 * missing_l + 1), fill, axis=-1)
|
|
|
- return Expansion(target_l_array, new_coefs)
|
|
|
+def expansion_total_charge(coefs: Array, sigma0: float | Array):
|
|
|
+ """Adds a new axis to the expansion coefficients that modifies expansion based on given net charge density."""
|
|
|
+ if sigma0 is None:
|
|
|
+ return coefs
|
|
|
+ if not isinstance(sigma0, Array):
|
|
|
+ x = copy.deepcopy(coefs)
|
|
|
+ x[..., 0] = sigma0 / np.sqrt(4 * np.pi)
|
|
|
+ return x
|
|
|
+ sigma0 = sigma0.flatten()
|
|
|
+ x = np.repeat(np.expand_dims(coefs, -2), len(sigma0), axis=-2)
|
|
|
+ x[..., 0] = sigma0 / np.sqrt(4 * np.pi)
|
|
|
+ return x
|
|
|
|
|
|
|
|
|
def m_indices_at_l(l_arr: Array, l_idx: int):
|
|
|
+ """
|
|
|
+ For a given l_array and index l_idx for some ell in this array, get indices of all (ell, m) coefficients
|
|
|
+ in coefficients array.
|
|
|
+ """
|
|
|
return np.arange(np.sum(2 * l_arr[:l_idx] + 1), np.sum(2 * l_arr[:l_idx+1] + 1))
|
|
|
|
|
|
|
|
|
def purge_unneeded_l(l_array: Array, coefs: Array) -> (Array, Array):
|
|
|
+ """Remove ell values from expansion for which all (ell, m) coefficients are zero."""
|
|
|
def delete_zero_entries(l, l_arr, cfs):
|
|
|
l_idx = np.where(l_arr == l)[0][0]
|
|
|
m_indices = m_indices_at_l(l_arr, l_idx)
|
|
|
@@ -243,48 +262,20 @@ def purge_unneeded_l(l_array: Array, coefs: Array) -> (Array, Array):
|
|
|
return l_array, coefs
|
|
|
|
|
|
|
|
|
-def plot_theta_profile(ex: Expansion, phi: float = 0, num: int = 100, theta_start: float = 0, theta_end: float = np.pi):
|
|
|
- theta_vals = np.linspace(theta_start, theta_end, num)
|
|
|
- charge = ex.charge_value(theta_vals, phi)
|
|
|
- plt.plot(theta_vals, charge.T)
|
|
|
- plt.show()
|
|
|
-
|
|
|
-
|
|
|
-def plot_theta_profile_multiple(ex_list: list[Expansion], label_list, phi: float = 0, num: int = 100, theta_start: float = 0, theta_end: float = np.pi):
|
|
|
- theta_vals = np.linspace(theta_start, theta_end, num)
|
|
|
-
|
|
|
- fig, ax = plt.subplots()
|
|
|
- for ex, label in zip(ex_list, label_list):
|
|
|
- ax.plot(theta_vals, ex.charge_value(theta_vals, phi).T, label=label)
|
|
|
- ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
|
|
|
- ax.set_xlabel(r'$\theta$', fontsize=13)
|
|
|
- ax.set_ylabel(r'$\sigma$', fontsize=13)
|
|
|
- plt.legend(fontsize=12)
|
|
|
- plt.tight_layout()
|
|
|
- plt.savefig(Path("/home/andraz/ChargedShells/Figures/charge_shape_comparison.png"), dpi=600)
|
|
|
- plt.show()
|
|
|
-
|
|
|
-
|
|
|
-def plot_charge_3d(ex: Expansion, num_theta=100, num_phi=100):
|
|
|
- theta = np.linspace(0, np.pi, num_theta)
|
|
|
- phi = np.linspace(0, 2 * np.pi, num_phi)
|
|
|
- theta, phi = np.meshgrid(theta, phi)
|
|
|
- r = ex.charge_value(theta.flatten(), phi.flatten()).reshape(theta.shape)
|
|
|
-
|
|
|
- # Convert spherical coordinates to Cartesian coordinates
|
|
|
- x = np.sin(theta) * np.cos(phi)
|
|
|
- y = np.sin(theta) * np.sin(phi)
|
|
|
- z = np.cos(theta)
|
|
|
-
|
|
|
- # Create a heatmap on the sphere
|
|
|
- fig = go.Figure(data=go.Surface(x=x, y=y, z=z, surfacecolor=r, colorscale='Jet'))
|
|
|
- fig.update_layout(scene=dict(aspectmode='data'))
|
|
|
- fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
|
|
|
-
|
|
|
- fig.show()
|
|
|
+def coefs_fill_missing_l(expansion: Expansion, target_l_array: Array) -> Expansion:
|
|
|
+ """Explicitly add missing expansion coefficients so that expansion includes all ell values from the target array."""
|
|
|
+ missing_l = np.setdiff1d(target_l_array, expansion.l_array, assume_unique=True)
|
|
|
+ fill = np.zeros(np.sum(2 * missing_l + 1))
|
|
|
+ full_l_array1, _ = expansion.lm_arrays
|
|
|
+ # we search for where to place missing coefs with the help of a boolean array and argmax function
|
|
|
+ comparison_bool = (full_l_array1[:, None] - missing_l[None, :]) > 0
|
|
|
+ indices = np.where(np.any(comparison_bool, axis=0), np.argmax(comparison_bool, axis=0), full_l_array1.shape[0])
|
|
|
+ new_coefs = np.insert(expansion.coefs, np.repeat(indices, 2 * missing_l + 1), fill, axis=-1)
|
|
|
+ return Expansion(target_l_array, new_coefs)
|
|
|
|
|
|
|
|
|
def expansions_to_common_l(ex1: Expansion, ex2: Expansion) -> (Expansion, Expansion):
|
|
|
+ """Explicitly add zero expansion coefficients so that both expansions include coefficients for the same ell."""
|
|
|
common_l_array = np.union1d(ex1.l_array, ex2.l_array)
|
|
|
return coefs_fill_missing_l(ex1, common_l_array), coefs_fill_missing_l(ex2, common_l_array)
|
|
|
|
|
|
@@ -309,42 +300,3 @@ def expansion_rotation(rotations: Quaternion, coefs: Array, l_array: Array):
|
|
|
new_coefs[..., all_m_indices] = np.einsum('rnm, qm -> rqn',
|
|
|
wm, coefs_reshaped[:, all_m_indices])
|
|
|
return new_coefs.reshape(rotations.ndarray.shape[:-1] + coefs.shape)
|
|
|
-
|
|
|
-
|
|
|
-if __name__ == '__main__':
|
|
|
-
|
|
|
- # ex = MappedExpansionQuad(0.44, 3, 1, 10)
|
|
|
- # ex = Expansion(np.arange(3), np.array([1, -1, 0, 1, 2, 0, 3, 0, 2]))
|
|
|
- # ex = GaussianCharges(omega_k=np.array([[0, 0], [np.pi, 0]]), lambda_k=10, sigma1=0.001, l_max=10)
|
|
|
- ex = SphericalCap(np.array([[0, 0], [np.pi, 0]]), 0.5, 0.1, 70)
|
|
|
- # print(np.real(ex.coefs))
|
|
|
- # plot_theta_profile(ex, num=1000, theta_end=2 * np.pi, phi=0)
|
|
|
- plot_charge_3d(ex)
|
|
|
-
|
|
|
- # new_coeffs = expansion_rotation(Quaternion(np.arange(20).reshape(5, 4)).normalized, ex.coeffs, ex.l_array)
|
|
|
- # print(new_coeffs.shape)
|
|
|
- #
|
|
|
- # newnew_coeffs = expansion_rotation(Quaternion(np.arange(16).reshape(4, 4)).normalized, new_coeffs, ex.l_array)
|
|
|
- # print(newnew_coeffs.shape)
|
|
|
-
|
|
|
- # rot_angles = np.linspace(0, np.pi, 1000)
|
|
|
- # t0 = time.time()
|
|
|
- # ex.rotate_euler(0, np.pi / 2, -1)
|
|
|
- # t1 = time.time()
|
|
|
- # print(ex.coefs)
|
|
|
- # print(t1 - t0)
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
|
|
-
|
