Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix memory usage issue in evaluate_density (#121) #191

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Fix memory usage issue in evaluate_density (#121) #191

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9c0be30
Fix memory usage issue in `evaluate_density` (#121)
msricher Jun 26, 2024
f73b556
Add chunked algorithm
msricher Jul 8, 2024
a061e5f
Fix chunking
msricher Jul 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 32 additions & 11 deletions gbasis/evals/density.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,11 @@
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


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
Expand All @@ -16,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
-------
Expand Down Expand Up @@ -57,9 +60,9 @@ 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)
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):
Expand Down Expand Up @@ -94,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(
Expand Down
Loading