Merge pull request #189 from timcallow/fp_options

Give users an option for float 32 as default

atoMEC/config.py

@@ -1,5 +1,7 @@
 """Configuration file to store global parameters."""
 
+import numpy as np
+
 # physical constants
 mp_g = 1.6726219e-24  # mass of proton in grams
 
@@ -30,6 +32,7 @@
 force_bound = []
 
 grid_type = "log"
+fp = np.float64
 
 # parallelization
 numcores = 0  # defaults to serial

atoMEC/convergence.py

@@ -32,9 +32,13 @@ class SCF:
 
     def __init__(self, xgrid, grid_type):
         self._xgrid = xgrid
-        self._energy = np.zeros((2))
-        self._potential = np.zeros((2, config.spindims, config.grid_params["ngrid"]))
-        self._density = np.zeros((2, config.spindims, config.grid_params["ngrid"]))
+        self._energy = np.zeros((2), dtype=config.fp)
+        self._potential = np.zeros(
+            (2, config.spindims, config.grid_params["ngrid"]), dtype=config.fp
+        )
+        self._density = np.zeros(
+            (2, config.spindims, config.grid_params["ngrid"]), dtype=config.fp
+        )
         self.grid_type = grid_type
 
     def check_conv(self, E_free, pot, dens, iscf):

atoMEC/numerov.py

@@ -33,6 +33,8 @@
 
 # from . import writeoutput
 
+dtype_map = {32: np.float32, 64: np.float64}
+
 
 class Solver:
     """
@@ -46,6 +48,7 @@ class Solver:
 
     def __init__(self, grid_type):
         self.grid_type = grid_type
+        self.fp = config.fp
 
     def calc_eigs_min(self, v, xgrid, bc):
         """
@@ -130,8 +133,12 @@ def matrix_solve(self, v, xgrid, bc, solve_type="full", eigs_min_guess=None):
            solving Schrödinger’s equation, American Journal of Physics 80, 1017-1019
            (2012) `DOI:10.1119/1.4748813 <https://doi.org/10.1119/1.4748813>`__.
         """
+        if solve_type == "guess":
+            dtype = np.float64
+        else:
+            dtype = self.fp
         if eigs_min_guess is None:
-            eigs_min_guess = np.zeros((config.spindims, config.lmax))
+            eigs_min_guess = np.zeros((config.spindims, config.lmax), dtype=dtype)
 
         # define the spacing of the xgrid
         dx = xgrid[1] - xgrid[0]
@@ -145,11 +152,13 @@ def matrix_solve(self, v, xgrid, bc, solve_type="full", eigs_min_guess=None):
         # J(x)=exp(x/2) for log grid; J(x) = x**-1.5 for sqrt grid
 
         # off-diagonal matrices
-        I_minus = np.eye(N, k=-1)
-        I_zero = np.eye(N)
-        I_plus = np.eye(N, k=1)
+        I_minus = np.eye(N, k=-1, dtype=dtype)
+        I_zero = np.eye(N, dtype=dtype)
+        I_plus = np.eye(N, k=1, dtype=dtype)
 
-        p = np.zeros((N, N))  # transformation for kinetic term on log/sqrt grid
+        p = np.zeros(
+            (N, N), dtype=dtype
+        )  # transformation for kinetic term on log/sqrt grid
         if self.grid_type == "log":
             np.fill_diagonal(p, np.exp(-2 * xgrid))
         else:
@@ -240,15 +249,19 @@ def KS_matsolve_parallel(self, T, B, v, xgrid, bc, solve_type, eigs_min_guess):
         N.B. if `config.numcores=-N` then `joblib` detects the number of available cores
         `n_avail` and parallelizes into `n_avail + 1 - N` separate jobs.
         """
+        if solve_type == "guess":
+            dtype = np.float64
+        else:
+            dtype = self.fp
         # compute the number of grid points
         N = np.size(xgrid)
 
         # Compute the number pmax of distinct diagonizations to be solved
         pmax = config.spindims * config.lmax
 
         # now flatten the potential matrix over spins
-        v_flat = np.zeros((pmax, N))
-        eigs_guess_flat = np.zeros((pmax))
+        v_flat = np.zeros((pmax, N), dtype=dtype)
+        eigs_guess_flat = np.zeros((pmax), dtype=dtype)
         for i in range(np.shape(v)[0]):
             for l in range(config.lmax):
                 if self.grid_type == "log":
@@ -299,8 +312,8 @@ def KS_matsolve_parallel(self, T, B, v, xgrid, bc, solve_type, eigs_min_guess):
 
         if solve_type == "full":
             # retrieve the eigfuncs and eigvals from the joblib output
-            eigfuncs_flat = np.zeros((pmax, config.nmax, N))
-            eigvals_flat = np.zeros((pmax, config.nmax))
+            eigfuncs_flat = np.zeros((pmax, config.nmax, N), dtype=dtype)
+            eigvals_flat = np.zeros((pmax, config.nmax), dtype=dtype)
             for q in range(pmax):
                 eigfuncs_flat[q] = X[q][0]
                 eigvals_flat[q] = X[q][1]
@@ -350,15 +363,19 @@ def KS_matsolve_serial(self, T, B, v, xgrid, bc, solve_type, eigs_min_guess):
         eigvals : ndarray
             KS eigenvalues
         """
+        if solve_type == "guess":
+            dtype = np.float64
+        else:
+            dtype = self.fp
         # compute the number of grid points
         N = np.size(xgrid)
         # initialize empty potential matrix
-        V_mat = np.zeros((N, N))
+        V_mat = np.zeros((N, N), dtype=dtype)
 
         # initialize the eigenfunctions and their eigenvalues
-        eigfuncs = np.zeros((config.spindims, config.lmax, config.nmax, N))
-        eigvals = np.zeros((config.spindims, config.lmax, config.nmax))
-        eigs_guess = np.zeros((config.spindims, config.lmax))
+        eigfuncs = np.zeros((config.spindims, config.lmax, config.nmax, N), dtype=dtype)
+        eigvals = np.zeros((config.spindims, config.lmax, config.nmax), dtype=dtype)
+        eigs_guess = np.zeros((config.spindims, config.lmax), dtype=dtype)
 
         # A new Hamiltonian has to be re-constructed for every value of l and each spin
         # channel if spin-polarized
@@ -397,7 +414,7 @@ def KS_matsolve_serial(self, T, B, v, xgrid, bc, solve_type, eigs_min_guess):
                         tol=config.conv_params["eigtol"],
                     )
 
-                    K = np.zeros((N, config.nmax))
+                    K = np.zeros((N, config.nmax), dtype=dtype)
                     if self.grid_type == "log":
                         prefac = -2 * np.exp(2 * xgrid)
                     else:
@@ -414,7 +431,7 @@ def KS_matsolve_serial(self, T, B, v, xgrid, bc, solve_type, eigs_min_guess):
 
                     # sort the eigenvalues to find the lowest
                     idr = np.argsort(eigs_up)
-                    eigs_guess[i, l] = np.array(eigs_up[idr].real)[0]
+                    eigs_guess[i, l] = np.array(eigs_up[idr].real, dtype=dtype)[0]
 
                     # dummy variable for the null eigenfucntions
                     eigfuncs_null = eigfuncs
@@ -456,10 +473,14 @@ def diag_H(self, p, T, B, v, xgrid, nmax, bc, eigs_guess, solve_type):
         evals : ndarray
             the KS eigenvalues
         """
+        if solve_type == "guess":
+            dtype = np.float64
+        else:
+            dtype = self.fp
         # compute the number of grid points
         N = np.size(xgrid)
         # initialize empty potential matrix
-        V_mat = np.zeros((N, N))
+        V_mat = np.zeros((N, N), dtype=dtype)
 
         # fill potential matrices
         # np.fill_diagonal(V_mat, v + 0.5 * (l + 0.5) ** 2 * np.exp(-2 * xgrid))
@@ -491,7 +512,7 @@ def diag_H(self, p, T, B, v, xgrid, nmax, bc, eigs_guess, solve_type):
             )
 
             # sort and normalize
-            K = np.zeros((N, nmax))
+            K = np.zeros((N, nmax), dtype=dtype)
             for n in range(nmax):
                 if self.grid_type == "log":
                     K[:, n] = (
@@ -509,10 +530,10 @@ def diag_H(self, p, T, B, v, xgrid, nmax, bc, eigs_guess, solve_type):
 
             # sort the eigenvalues to find the lowest
             idr = np.argsort(evals)
-            evals = np.array(evals[idr].real)[0]
+            evals = np.array(evals[idr].real, dtype=dtype)[0]
 
             # dummy eigenvector for return statement
-            evecs_null = np.zeros((N))
+            evecs_null = np.zeros((N), dtype=dtype)
 
             return evecs_null, evals
 
@@ -541,13 +562,13 @@ def update_orbs(l_eigfuncs, l_eigvals, xgrid, bc, K, grid_type):
         """
         # Sort eigenvalues in ascending order
         idr = np.argsort(l_eigvals)
-        eigvals = np.array(l_eigvals[idr].real)
+        eigvals = np.array(l_eigvals[idr].real, dtype=np.float64)
 
         # resize l_eigfuncs from N-1 to N for dirichlet condition
         if bc == "dirichlet":
             N = np.size(xgrid)
             nmax = np.shape(l_eigfuncs)[1]
-            l_eigfuncs_dir = np.zeros((N, nmax))
+            l_eigfuncs_dir = np.zeros((N, nmax), dtype=np.float64)
             l_eigfuncs_dir[:-1] = l_eigfuncs.real
             l_eigfuncs = l_eigfuncs_dir
 
@@ -559,7 +580,7 @@ def update_orbs(l_eigfuncs, l_eigvals, xgrid, bc, K, grid_type):
         ) / (1 + h * K[-1])
 
         # convert to correct dimensions
-        eigfuncs = np.array(np.transpose(l_eigfuncs.real)[idr])
+        eigfuncs = np.array(np.transpose(l_eigfuncs.real)[idr], dtype=np.float64)
         if grid_type == "sqrt":
             eigfuncs *= xgrid**-1.5
         eigfuncs = mathtools.normalize_orbs(eigfuncs, xgrid, grid_type)  # normalize
@@ -601,9 +622,9 @@ def calc_wfns_e_grid(self, xgrid, v, e_arr, eigfuncs_l, eigfuncs_u):
         e_arr_flat = e_arr.flatten()
 
         # initialize the W (potential) and eigenfunction arrays
-        W_arr = np.zeros((N, nkpts, spindims, lmax, nmax))
+        W_arr = np.zeros((N, nkpts, spindims, lmax, nmax), dtype=np.float64)
         eigfuncs_init = np.zeros_like(W_arr)
-        eigfuncs_backup = np.zeros((nkpts, spindims, lmax, nmax, N))
+        eigfuncs_backup = np.zeros((nkpts, spindims, lmax, nmax, N), dtype=np.float64)
         if self.grid_type == "sqrt":
             for k in range(nkpts):
                 eigfuncs_backup[k] = (k * eigfuncs_u + (nkpts - k - 1) * eigfuncs_l) / (

atoMEC/postprocess/conductivity.py

@@ -25,7 +25,7 @@
 from scipy.integrate import quad
 
 # internal modules
-from atoMEC import mathtools
+from atoMEC import mathtools, config
 
 
 class KuboGreenwood:
@@ -320,7 +320,7 @@ def check_sum_rule(self, l, n, m):
 
         # initialize sum_mom and various indices
         nbands, nspin, lmax, nmax = np.shape(self._eigvals)
-        sum_mom = np.zeros((nbands))
+        sum_mom = np.zeros((nbands), dtype=config.fp)
 
         # compute the sum rule
         for k in range(nbands):
@@ -517,8 +517,10 @@ def calc_sig_func(
         omega_arr = np.linspace(1e-5, np.sqrt(omega_max), n_freq) ** 2
 
         # set up lorentzian: requires dummy array to get right shape
-        sig_omega = np.zeros((np.size(omega_arr), 2))
-        omega_dummy_mat = np.ones((nbands, lmax, nmax, lmax, nmax, n_freq))
+        sig_omega = np.zeros((np.size(omega_arr), 2), dtype=config.fp)
+        omega_dummy_mat = np.ones(
+            (nbands, lmax, nmax, lmax, nmax, n_freq), dtype=config.fp
+        )
         eig_diff_omega_mat = np.einsum(
             "nijkl,nijklm->nijklm", eig_diff_mat, omega_dummy_mat
         )
@@ -698,7 +700,7 @@ def P_mat_int(cls, func_int, lmax):
         P2 and P4 functions, ands the :func:`P2_func`, :func:`P4_func` and
         :func:`P_int` functions.
         """
-        P_mat = np.zeros((lmax, lmax, 2 * lmax + 1))
+        P_mat = np.zeros((lmax, lmax, 2 * lmax + 1), dtype=config.fp)
 
         for l1 in range(lmax):
             for l2 in range(lmax):

atoMEC/postprocess/localization.py

@@ -19,7 +19,7 @@
 import numpy as np
 
 # internal modules
-from atoMEC import staticKS, mathtools, check_inputs
+from atoMEC import staticKS, mathtools, check_inputs, config
 
 
 class ELFTools:
@@ -65,8 +65,8 @@ def __init__(self, Atom, model, orbitals, density, method="orbitals"):
         spindims = np.shape(self._eigfuncs)[1]
         ngrid = np.shape(self._eigfuncs)[4]
 
-        self._ELF = np.zeros((spindims, ngrid))
-        self._epdc = np.zeros((spindims, ngrid))
+        self._ELF = np.zeros((spindims, ngrid), dtype=config.fp)
+        self._epdc = np.zeros((spindims, ngrid), dtype=config.fp)
         self._N_shell = None
 
     @property
@@ -410,7 +410,7 @@ def calc_IPR_mat(eigfuncs, xgrid, grid_type=None):
     lmax = np.shape(eigfuncs)[2]
     nmax = np.shape(eigfuncs)[3]
 
-    IPR_mat = np.zeros((nkpts, spindims, lmax, nmax))
+    IPR_mat = np.zeros((nkpts, spindims, lmax, nmax), dtype=config.fp)
 
     # compute the IPR matrix
     # FIXME: add spherical harmonic term