charged-shells/patch_size.py

import numpy as np
from scipy.optimize import bisect
import expansion
from matplotlib import pyplot as plt
import parameters
import potentials

Expansion = expansion.Expansion
Array = np.ndarray
ModelParams = parameters.ModelParams


@expansion.map_over_expansion
def charge_patch_size(ex: Expansion, phi: float = 0, theta0: Array | float = 0, theta1: Array | float = np.pi / 2):
    return bisect(lambda theta: ex.charge_value(theta, phi), theta0, theta1)


def potential_patch_size(ex: Expansion, params: ModelParams,
                         phi: float = 0, theta0: Array | float = 0, theta1: Array | float = np.pi / 2,
                         match_expansion_axis_to_params: int = None):

    meatp = match_expansion_axis_to_params

    @expansion.map_over_expansion
    def potential_zero(exp: Expansion, prms: ModelParams):
        return bisect(lambda theta: potentials.charged_shell_potential(theta, phi, 1, exp, prms), theta0, theta1)

    print(ex.shape)

    params_list = params.unravel()
    if meatp is not None:
        expansion_list = [Expansion(ex.l_array, np.take(ex.coefs, i, axis=meatp)) for i in range(ex.shape[meatp])]
    else:
        expansion_list = [ex]

    results = []
    for exp, prms in zip(expansion_list, params_list):
        results.append(potential_zero(exp, prms))


if __name__ == '__main__':

    a_bar = np.linspace(0.2, 0.8, 100)
    kappaR = np.array([0.26, 1, 3, 10, 26])
    params = ModelParams(R=150, kappaR=kappaR)
    ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_m=0.001, l_max=20, kappaR=kappaR[None, :])
    print(ex.shape)

    patch_size = charge_patch_size(ex)
    patch_size_pot = potential_patch_size(ex, params, match_expansion_axis_to_params=1)

    plt.plot(a_bar, patch_size * 180 / np.pi, label=kappaR)
    plt.legend()
    plt.show()