From 9c0be30810a23c82f9acaaa229f013e26dd7d648 Mon Sep 17 00:00:00 2001 From: Michelle Richer Date: Wed, 26 Jun 2024 14:23:05 -0400 Subject: [PATCH 1/3] Fix memory usage issue in `evaluate_density` (#121) --- gbasis/evals/density.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/gbasis/evals/density.py b/gbasis/evals/density.py index c0712924..88baf60f 100644 --- a/gbasis/evals/density.py +++ b/gbasis/evals/density.py @@ -57,9 +57,7 @@ def evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval): " of the orbital evaluations." ) - density = one_density_matrix.dot(orb_eval) - density *= orb_eval - return np.sum(density, axis=0) + return np.einsum("pq,qn,pn->n", one_density_matrix, orb_eval, orb_eval, optimize=True) def evaluate_density(one_density_matrix, basis, points, transform=None, threshold=1.0e-8): From f73b556cd79f8da96f3013c026818621e6606da4 Mon Sep 17 00:00:00 2001 From: Michelle Richer Date: Mon, 8 Jul 2024 13:32:43 -0400 Subject: [PATCH 2/3] Add chunked algorithm --- gbasis/evals/density.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/gbasis/evals/density.py b/gbasis/evals/density.py index 88baf60f..ff86e033 100644 --- a/gbasis/evals/density.py +++ b/gbasis/evals/density.py @@ -2,6 +2,7 @@ from gbasis.evals.eval import evaluate_basis from gbasis.evals.eval_deriv import evaluate_deriv_basis import numpy as np +import psutil from scipy.special import comb @@ -57,7 +58,29 @@ def evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval): " of the orbital evaluations." ) - return np.einsum("pq,qn,pn->n", one_density_matrix, orb_eval, orb_eval, optimize=True) + # Output vector of densities + density = np.zeros(orb_eval.shape[1], dtype=one_density_matrix.dtype) + # Number of chunks (use 8.1 as the number of bytes in a double, just to have a buffer so that we + # don't run out of memory). If it fits in 1 chunk, 1 chunk will be used. + chunks = max( + 1, + int( + 8.1 + * one_density_matrix.shape[0] + * orb_eval.shape[1] + / psutil.virtual_memory().available + ), + ) + # Evaluate densities chunk-wise + for density_chunk, orb_eval_chunk in zip( + np.split(density, chunks, axis=0), np.split(orb_eval, chunks, axis=1) + ): + # Evaluate orbital densities + d = one_density_matrix.dot(orb_eval_chunk) + d *= orb_eval_chunk + # Sum orbital densities into output vector + np.sum(d, axis=0, out=density_chunk) + return density def evaluate_density(one_density_matrix, basis, points, transform=None, threshold=1.0e-8): From a061e5f9bb34e58e64159d0dfebbbfeaad2d8cc3 Mon Sep 17 00:00:00 2001 From: Michelle Richer Date: Wed, 10 Jul 2024 09:17:59 -0400 Subject: [PATCH 3/3] Fix chunking --- gbasis/evals/density.py | 62 ++++++++++++++++++++--------------------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/gbasis/evals/density.py b/gbasis/evals/density.py index ff86e033..4057b99c 100644 --- a/gbasis/evals/density.py +++ b/gbasis/evals/density.py @@ -6,7 +6,7 @@ from scipy.special import comb -def evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval): +def evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval, out=None): """Return the evaluation of the density given the evaluated orbitals. Parameters @@ -17,6 +17,8 @@ def evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval): Orbitals evaluated at :math:`N` grid points. The set of orbitals must be the same basis set used to build the one-electron density matrix. + out : np.ndarray(N,), optional + Output array. Returns ------- @@ -58,29 +60,9 @@ def evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval): " of the orbital evaluations." ) - # Output vector of densities - density = np.zeros(orb_eval.shape[1], dtype=one_density_matrix.dtype) - # Number of chunks (use 8.1 as the number of bytes in a double, just to have a buffer so that we - # don't run out of memory). If it fits in 1 chunk, 1 chunk will be used. - chunks = max( - 1, - int( - 8.1 - * one_density_matrix.shape[0] - * orb_eval.shape[1] - / psutil.virtual_memory().available - ), - ) - # Evaluate densities chunk-wise - for density_chunk, orb_eval_chunk in zip( - np.split(density, chunks, axis=0), np.split(orb_eval, chunks, axis=1) - ): - # Evaluate orbital densities - d = one_density_matrix.dot(orb_eval_chunk) - d *= orb_eval_chunk - # Sum orbital densities into output vector - np.sum(d, axis=0, out=density_chunk) - return density + d = one_density_matrix.dot(orb_eval) + d *= orb_eval + return np.sum(d, axis=0, out=out) def evaluate_density(one_density_matrix, basis, points, transform=None, threshold=1.0e-8): @@ -115,13 +97,31 @@ def evaluate_density(one_density_matrix, basis, points, transform=None, threshol Density evaluated at `N` grid points. """ - orb_eval = evaluate_basis(basis, points, transform=transform) - output = evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval) - # Fix #117: check magnitude of small negative density values, then use clip to remove them - min_output = np.min(output) - if min_output < 0.0 and abs(min_output) > threshold: - raise ValueError(f"Found negative density <= {-threshold}, got {min_output}.") - return output.clip(min=0.0) + # Output vector of densities + density = np.zeros(points.shape[0]) + # Number of chunks (use 8.1 as the number of bytes in a double, just to have a buffer so that we + # don't run out of memory). If it fits in 1 chunk, 1 chunk will be used. + chunks = max( + 1, + int( + 8.1 + * one_density_matrix.shape[0] + * points.shape[0] + / psutil.virtual_memory().available + ), + ) + # Evaluate densities chunk-wise + for density_chunk, points_chunk in zip( + np.array_split(density, chunks, axis=0), np.array_split(points, chunks, axis=0) + ): + orb_eval = evaluate_basis(basis, points_chunk, transform=transform) + evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval, out=density_chunk) + # Fix #117: check magnitude of small negative density values, then remove them + min_density = np.min(density) + if min_density <= -threshold: + raise ValueError(f"Found negative density <= {-threshold}, got {min_density}.") + density[density < 0] = 0 + return density def evaluate_deriv_reduced_density_matrix(