2023-10-12

expansion.py

@@ -50,14 +50,14 @@ class Expansion:
         return self.coefs.shape[:-1]
 
     def flatten(self) -> Expansion:
-        new_expansion = self.clone()
+        new_expansion = self.clone()  # np.ndarray.flatten() also copies the array
         new_expansion.coefs = new_expansion.coefs.reshape(-1, new_expansion.coefs.shape[-1])
         new_expansion._rotations = new_expansion._rotations.reshape(-1, 4)
         return new_expansion
 
-    def reshape(self, shape):
+    def reshape(self, shape: tuple):
         self.coefs = self.coefs.reshape(shape + (self.coefs.shape[-1],))
-        self._rotations = self._rotations.reshape(shape + (4),)
+        self._rotations = self._rotations.reshape(shape + (4,))
 
     @cached_property
     def lm_arrays(self) -> (Array, Array):
@@ -148,7 +148,7 @@ class GaussianCharges(Expansion):
 
 
 def map_over_expansion(f: Callable[[Expansion, T], V]) -> Callable[[Expansion, T], V]:
-    """Map a function f over the leading axes of an expansion."""
+    """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

patch_size.py

@@ -19,24 +19,35 @@ 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
+    # this is more complicate to map over leading axes of the expansion as potential also depends on model parameters,
+    # with some, such as kappaR, also being the parameters of the expansion in the first place. When mapping,
+    # we must therefore provide the expansion axis that should match the collection of parameters in params.
+
+    meap = match_expansion_axis_to_params  # just a shorter variable name
 
     @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])]
+    if meap is not None:
+        expansion_list = [Expansion(ex.l_array, np.take(ex.coefs, i, axis=meap)) for i in range(ex.shape[meap])]
     else:
-        expansion_list = [ex]
+        expansion_list = [ex for _ in params_list]
+
+    if not len(expansion_list) == len(params_list):
+        raise ValueError(f'Axis of expansion that is supposed to match params does not have the same length, got '
+                         f'len(params.unravel()) = {len(params_list)} and '
+                         f'expansion.shape[{meap}] = {len(expansion_list)}')
 
     results = []
     for exp, prms in zip(expansion_list, params_list):
         results.append(potential_zero(exp, prms))
 
+    if meap is not None:
+        return np.array(results).swapaxes(0, meap)  # return the params-matched axis to where it belongs
+    return np.array(results)
+
 
 if __name__ == '__main__':
 
@@ -44,12 +55,11 @@ if __name__ == '__main__':
     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 = 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.plot(a_bar, patch_size_pot * 180 / np.pi, label=kappaR)
     plt.legend()
     plt.show()
 

potentials.py

@@ -26,6 +26,9 @@ def charged_shell_potential(theta: Array | float,
     :param params: ModelParams object specifying parameter values for the model
     """
     theta, phi = np.broadcast_arrays(theta, phi)
+    angles_shape = theta.shape
+    theta = theta.reshape(-1)  # ensures that arrays are 1D
+    phi = phi.reshape(-1)
 
     if not theta.shape == phi.shape:
         raise ValueError('theta and phi arrays should have the same shape.')
@@ -36,11 +39,11 @@ def charged_shell_potential(theta: Array | float,
                  / fn.sph_bessel_k(ex.l_array, params.kappaR))
     l_factors = ex.repeat_over_m(l_factors)
 
+    pot = (1 / (params.kappa * params.epsilon * uc.CONSTANTS.epsilon0)
+            * np.real(np.sum(l_factors[:, None] * ex.coefs[..., None]
+                             * fn.sph_harm(l_array[:, None], m_array[:, None], theta[None, :], phi[None, :]), axis=-2)))
 
-
-    return (1 / (params.kappa * params.epsilon * uc.CONSTANTS.epsilon0)
-            * np.real(np.sum(ex.repeat_over_m(l_factors)[None, :] * ex.coefs
-                             * fn.sph_harm(l_array[None, :], m_array[None, :], theta[:, None], phi[:, None]), axis=1)))
+    return pot.reshape(ex.shape + angles_shape)
 
 
 if __name__ == '__main__':
@@ -53,7 +56,8 @@ if __name__ == '__main__':
     dist = 1
 
     potential = charged_shell_potential(theta, phi, dist, ex, params)
+    print(potential.shape)
     # print(potential)
 
-    plt.plot(potential)
+    plt.plot(theta, potential.T)
     plt.show()