[BUG] Big memory usage for density evaluation

[leila-pujal]

Describe the bug

Big memory usage for density evaluation

To Reproduce

I attach the system I used to test:
iron_hexacarbonyl.tar.gz

mol=load_one("iron_hexacarbonyl.fchk")
dm = mol.one_rdms['scf']
basis, coord_types = from_iodata(mol)
becke = BeckeWeights(order=3)

oned = GaussChebyshevType2(100)
rgrid = BeckeTF(1e-4, 1.5).transform_1d_grid(oned)
grid = MolGrid.horton_molgrid(mol.atcoords, mol.atnums, rgrid, 1030, becke)

evaluate_density(dm, basis, grid.points, coord_type=coord_types)

Expected behaviour

I don't know if this could be improved, I am just reporting the behaviour as an issue.

Screenshots

System Information:

• OS: 18.04.5 LTS
• Python version: 3.6.10
• NumPy version: 1.18.1
• SciPy version: 1.4.1

Additional context

For computing the density it starts with evaluate_density(), then evaluates each orbital by computing the density for the basis functions and then it obtains the total density by using the density matrix. For evaluating the basis functions density, it is done in _deriv.py module inside _eval_deriv_contractions(). To monitor the memory I used several prints with nbytes for arrays and @profile from memory_profiler import profile. To monitor the memory of each contraction I printed the associated memory of matrix_contraction in base one

gbasis/gbasis/base_one.py

Line 239 in c57de8f

matrix_contraction = np.concatenate(matrix_contraction, axis=0)

. Then I profiled the memory in _eval_deriv_contractions(), evaluate_density() and evaluate_density_using_evaluated_orbs. With the information I gathered I think the memory dependcy in the code is as follows. Through the evaluation of all the generalized contractions, the memory will scale as a function of AO orbitals generated for each angular momentum. S-generalized contractions are the lowest memory cost and to obtain the memory cost for higher angular momentum generalized contractions you have to multiply the memory cost of the S-generalized contractions by the number of generated atomic orbitals (e.g l=1, 3p orbitals. Increase in memory = S-generalized contractions-memory x 3) The base memory corresponding to S-generalized contractions scales linearly with the number of points of the molecular grid. I fitted some values for not so big molecular sizes and I got the following regression S-generalized contractions memory (MB)=(0.000008150229943 *n_points) - 0.24118548. I checked to predict for some bigger ones and the predicted value was around 2MB wrong.
To show the agreement of the memory predicted with the one monitored I used iron_hexacarbonyl:

For a radial grid =100points
angular = 1030
Total = 100 * 1202 * 13atoms = 1562600

Iron hexacarbonyl
Total atoms	13
Total electrons	108
Total number generalized contractions	113

Atom types	Fe	oxygen	carbon
total number	1	6	6
s shell	5	4	4
p shell	4	2	3
d shell	2	2	2
f shell	1

Base memory for S-generalized contractions = 12.49MB

	Memory = base x AO_generated x number_shells
	nitrogen	oxygen	carbon
s shell	60 MB	48 MB	48 MB
p shell	144 MB	72 MB	108 MB
d shell	144 MB	144 MB	144 MB
f shell	120 MB	0	0
total (1atom)	468 MB	264MB	300MB	total memory
Total all atoms	468 MB	1584 MB	1800 MB	3852 MB

The values I obtained profiling the code are shown in the following pictures:

this is the profiling for evaluate density()

You can see evaluates basis increases almost the same memory that it has been calculated in line 78.

But also, there's a spike of memory that it is not shown here that occurs in evaluate_density_using_evaluated_orbs. The profiling for this is shown in the following image.

Here you can see in line 63 where the memory is doubled.

[PaulWAyers]

The memory-doubling should be avoidable I think.

A bigger issue is that we should be using some screening to not evaluate all the orbital overlaps, because we don't need them all. That is going to take a bit of refactoring, though perhaps @BradenDKelly has already looked at this?

[BradenDKelly]

@leila-pujal @PaulWAyers yes, there are two different PRs for screening. Either one would save some time/calculations. I will investigate once home.

[PaulWAyers]

@leila-pujal @marco-2023 this is the issue Leila was talking about. If there is a pull request open, it may be worth reviewing and merging it.

[leila-pujal]

I checked and the PR related to this seems to be #110 and #111. I will check them and see if they can help us to fix this problem.

[leila-pujal]

I opened a PR #166 that fixes this bug. Below there is my testing code that uses the new grid (iron_hexacarbonyl.fchk is provided at the top of the issue):

import numpy as np
from iodata import load_one

from grid.becke import BeckeWeights
from grid.molgrid import MolGrid
from grid.onedgrid import GaussChebyshev
from grid.rtransform import BeckeRTransform

from gbasis.evals.density import evaluate_density
from gbasis.wrappers import from_iodata

mol=load_one("iron_hexacarbonyl.fchk")
dm = mol.one_rdms['scf']
basis = from_iodata(mol)
becke = BeckeWeights(order=3)

aim_weights = BeckeWeights(radii=None, order=3)
oned = GaussChebyshev(100)
rgrid = BeckeRTransform(1e-4, 1.5).transform_1d_grid(oned)
grid = MolGrid.from_preset(mol.atnums, mol.atcoords, 'insane', rgrid, aim_weights=aim_weights)

d=evaluate_density(dm, basis, grid.points)

This gives a smaller grid than the one used when this issue was opened. With the old code, we get the doubling in memory in line 45 in evaluate_density_using_evaluated_orbs:

With the new code we avoid creating the intermediate matrix density and the memory stays the same:

[leila-pujal]

I commented on PR #166, I talked with @Ali-Tehrani earlier and he pointed out that condensing the lines of code does not solve the problem. I confirmed this earlier this week because I actually had a similar problem, the profiler only tracks memory increased/decreased before and after lines of code but doing the operations inside the np.sum it creates the array and deletes it at the same time. At this point, I am not sure if this doubling in memory can be avoided. I don't remember if we talked about some other alternative for this part of the code to avoid this happening.

[PaulWAyers]

There are external libraries to fix this, like numexpr. We used to use the linear algebra factories for this.

I think there is a way to do this with numpy "in place" operations but I am no expert.

[Ali-Tehrani]

Don't think "in-place" will work for this specific case. The other suggestions alongside using "in-place" would be:

1. To use einsum, setting optimize="optimal" or "greedy". This will find intermediate optimizations, see here for how. Don't think it will help in this specific case.

2. Knowing how much available memory there is, then calculating optimal number of points to fit in memory (minus a safe margin) then computing slices at a time. In this example, assuming double precision it would be solving the number of points N in K= (M^2 + 2NM + N)*8, where K=available memory in bytes, 8 is from 8 bytes in a double and M=number of MO. M^2 comes from one-rdm, NM comes from orb_eval, NM comes from density-intermediate array and N is the size of the output.