More code tidying

src/koopmans/ml_utils/_compute_decomposition.py

@@ -1,4 +1,3 @@
-import math
 from functools import partial
 from pathlib import Path
 from typing import Callable, Dict, List, Tuple
@@ -7,41 +6,19 @@
 from ase import Atoms, units
 from numpy.linalg import norm
 
-from koopmans import utils
 from koopmans.bands import Bands
 from koopmans.utils import read_xml_array, read_xml_nr
 
 from ._basis_functions import g as g_basis
 from ._basis_functions import \
     real_spherical_harmonics as real_spherical_harmonics_basis_functions
+from ._misc import cart2sph_array, compute_3d_integral_naive
 
-# Functions for the basis functions
+# Basis functions types
 RadialBasisFunctions = Callable[[np.ndarray, int, int, int], np.ndarray]
 SphericalBasisFunctions = Callable[[np.ndarray, np.ndarray, int, int], np.ndarray]
 
 
-# def radial_basis_function(r: np.ndarray, n: int, n_max: int, l: int, betas: np.ndarray, alphas: np.ndarray) -> np.ndarray:
-#     """
-#     Wrapper for the radial basis function.
-#
-#     Currently only Gaussian Type Basis Functions (see Hilmanen et al 2020) are implemented as radial basis functions.
-#     This wrapper allows to simply exchange them with other choices of radial basis functions.
-#     """
-#
-#     return g_basis(r, n, n_max, l, betas, alphas)
-#
-#
-# def spherical_basis_function(theta: np.ndarray, phi: np.ndarray, l: int, m: int) -> np.ndarray:
-#     """
-#     Wrapper for the spherical basis function.
-#
-#     Currently only real spherical harmonics (see Hilmanen et al 2020) are implemented as spherical basis functions.
-#     This wrapper allows to simply exchange them with other choices of spherical basis functions.
-#     """
-#
-#     return real_spherical_harmonics_basis(theta, phi, l, m)
-
-
 def precompute_basis_function(radial_basis_functions: RadialBasisFunctions,
                               spherical_basis_functions: SphericalBasisFunctions,
                               r_cartesian: np.ndarray, r_spherical: np.ndarray, n_max: int, l_max: int) -> np.ndarray:
@@ -76,35 +53,6 @@ def precompute_basis_function(radial_basis_functions: RadialBasisFunctions,
     return total_basis_function_array
 
 
-# Functions to convert cartesian to spherical coordinates
-
-def cart2sph(x: float, y: float, z: float) -> Tuple[float, float, float]:
-    """
-    Converts one cartesian coordinate into the corresponding spherical one.
-    """
-
-    XsqPlusYsq = x**2 + y**2
-    r = math.sqrt(XsqPlusYsq + z**2)
-    theta = math.atan2(z, math.sqrt(XsqPlusYsq))+np.pi/2.0
-    phi = math.atan2(y, x) + np.pi
-    return r, theta, phi
-
-
-def cart2sph_array(r_cartesian: np.ndarray) -> np.ndarray:
-    """
-    Converts an array of cartesian coordinates into the corresponding spherical ones.
-    """
-
-    (k_max, j_max, i_max, _) = np.shape(r_cartesian)
-    r_spherical = np.zeros_like(r_cartesian)
-    for k in range(k_max):
-        for j in range(j_max):
-            for i in range(i_max):
-                r_spherical[k, j, i, :] = cart2sph(
-                    r_cartesian[k, j, i, 2], r_cartesian[k, j, i, 1], r_cartesian[k, j, i, 0])
-    return r_spherical
-
-
 # functions to make sure that integration domain is chosen such that the center of the orbital density is the center of the integration domain
 
 def get_index(r: np.ndarray, vec: np.ndarray) -> Tuple[int, int, int]:
@@ -163,18 +111,6 @@ def translate_to_new_integration_domain(f: np.ndarray, wfc_center_index: Tuple[i
 
 # functions to compute the expansion coefficients
 
-def compute_3d_integral_naive(f: np.ndarray, r: np.ndarray) -> np.ndarray:
-    """
-    Computes the 3d-integral of an array of functions f on r with a simple trapezoidal rule.
-    """
-
-    z = r[:, 0, 0, 0]
-    y = r[0, :, 0, 1]
-    x = r[0, 0, :, 2]
-    result_n = np.sum(f, axis=(0, 1, 2))*(x[-1]-x[0])*(y[-1]-y[0])*(z[-1]-z[0])/((len(x)-1)*(len(y)-1)*(len(z)-1))
-    return result_n
-
-
 def get_coefficients(rho: np.ndarray, rho_total: np.ndarray, r_cartesian: np.ndarray, total_basis_function_array: np.ndarray) -> Tuple[List[float], List[float]]:
     """
     Computes the expansion coefficients of rho and rho_total wrt the basis defined in total_basis_function_array.
@@ -271,27 +207,3 @@ def compute_decomposition(n_max: int, l_max: int, r_min: float, r_max: float, r_
         # save the decomposition coefficients in files
         np.savetxt(dirs['coeff_orb'] / f'coff.orbital.{filled_str}.{band.index}.txt', coefficients_orbital)
         np.savetxt(dirs['coeff_tot'] / f'coff.total.{filled_str}.{band.index}.txt', coefficients_total)
-
-
-def get_reconstructed_orbital_densities(total_basis_array: np.ndarray, coefficients: List[float]) -> np.ndarray:
-    """
-    Reconstruct the density with the truncated expansion coefficients multiplied with the corresponding basis functions.
-    """
-
-    rho_r_reconstruced = np.einsum('ijkl,l->ijk', total_basis_array, coefficients)
-    return rho_r_reconstruced
-
-
-def map_again_to_original_grid(f_new: np.ndarray, wfc_center_index: Tuple[int, int, int], nr_xml: Tuple[int, int, int], nr_new_integration_domain: Tuple[int, int, int]):
-    """
-    Maps the function f_new defined on the new integration domain back to the unit cell.
-    """
-
-    f_on_reg_grid = np.zeros((nr_xml[2]-1, nr_xml[1]-1, nr_xml[0]-1), dtype=float)
-
-    for k_new, k in enumerate(range(wfc_center_index[0]-nr_new_integration_domain[2], wfc_center_index[0]+nr_new_integration_domain[2]+1)):
-        for j_new, j in enumerate(range(wfc_center_index[1]-nr_new_integration_domain[1], wfc_center_index[1]+nr_new_integration_domain[1]+1)):
-            for i_new, i in enumerate(range(wfc_center_index[2]-nr_new_integration_domain[0], wfc_center_index[2]+nr_new_integration_domain[0]+1)):
-                f_on_reg_grid[k % (nr_xml[2]-1), j % (nr_xml[1]-1), i % (nr_xml[0]-1)] = f_new[k_new, j_new, i_new]
-
-    return f_on_reg_grid

src/koopmans/ml_utils/_misc.py

@@ -3,6 +3,33 @@
 import numpy as np
 
 
+def cart2sph_array(r_cartesian: np.ndarray) -> np.ndarray:
+    """
+    Converts an array of cartesian coordinates into the corresponding spherical ones.
+
+    Note that cartesian is z, y, x; spherical is r, theta, phi
+    """
+
+    xy2 = r_cartesian[:, :, :, 2]**2 + r_cartesian[:, :, :, 1]**2
+    r_spherical = np.zeros_like(r_cartesian)
+    r_spherical[:, :, :, 0] = np.linalg.norm(r_cartesian, axis=-1)
+    r_spherical[:, :, :, 1] = np.arctan2(r_cartesian[:, :, :, 0], np.sqrt(xy2)) + np.pi / 2.0
+    r_spherical[:, :, :, 2] = np.arctan2(r_cartesian[:, :, :, 1], r_cartesian[:, :, :, 2]) + np.pi
+    return r_spherical
+
+
+def compute_3d_integral_naive(f: np.ndarray, r: np.ndarray) -> np.ndarray:
+    """
+    Computes the 3d-integral of an array of functions f on r with a simple trapezoidal rule.
+    """
+
+    z = r[:, 0, 0, 0]
+    y = r[0, :, 0, 1]
+    x = r[0, 0, :, 2]
+    result_n = np.sum(f, axis=(0, 1, 2))*(x[-1]-x[0])*(y[-1]-y[0])*(z[-1]-z[0])/((len(x)-1)*(len(y)-1)*(len(z)-1))
+    return result_n
+
+
 def reconstruct_orbital_densities(total_basis_array: np.ndarray, coefficients: List[float]) -> np.ndarray:
     """
     Reconstruct the density with the truncated expansion coefficients multiplied with the corresponding basis functions