Calculate psi values on a sparse grid, to accelerate multiple calculations.

[pfebrer]

This PR addreses #493.

Concept and implementation details

The basic idea is to compute orbital values and store them in a sparse data structure. For now, the orbital values are computed in a new method -Geometry.orbital_values- and stored in a new class, SparseGrid. I didn't think too much on where to put the functionalities, so feel free to reorganize it however you think it fits best.

Once you have the SparseGrid data structure, you can call its project_values and project_products methods, which, through cython implementations, reduce the orbitals dimension and therefore project the values on a normal sisl.Grid. These routines accept coeficient arrays, which determine the weight of each value/product to be added according to the orbital indices from which it came. This generalization is then used to compute wavefunctions (using project_values with the coefficients being the eigenstate coefficients), and electronic density (using project_products with the coefficients being the density matrix).

Some considerations about these two cython routines:

1. project_products: It loops over grid points to facilitate finding the overlaping orbitals. This makes it trivially easy to parallelize I would say, because there are no race conditions or any other problem. In principle, it should be as simple as changing the range call by cython's prange and compiling with openmp support, but I haven't tried.
2. project_products uses a dense coefficients matrix, because accessing a sparse one using indices is painfully slow if you put it inside a loop. Perhaps a smarter implementation could be added in the future to manage sparse matrices in case memory is a problem. But for now I don't see it as a problem, even for huge DFT systems.
3. project_values: In this case, there was no need to loop over some specific quantity, so it just loops over values and adds them as they come. This avoids the need for sorting the SparseGrid, but probably makes parallelization more difficult, since multiple iterations modify the same item of the grid. This is easy to solve by looping over grid points as in project_products, as long as the benefits are bigger than the cost of sorting.
4. Both: They use raveled grids and raveled indices, so they are not specific for 3D grids. This means that as long as you can generate a SparseGrid for something else (a path or a simple collection of points, then it would be called SparseRealSpace I guess), you can use these same routines out of the box. This opens the door for adding other RealSpace classes in sisl, instead of having to always go through grids, which sometimes are unnecessarily expensive for the goals of the user.
5. In the same way, higher order products could be implemented. I don't know if they are useful for something.

Performance for computing the density

First, I needed to check that the implementation gives correct results. I took some SIESTA calculations that I had and compared the RHO file returned by SIESTA with the density calculated by this implementation using the DM file, in the exact same grid. Then took the absolute difference and computed histograms to check that the differences are close to 0. For both gamma point calculations and k-point calculations I get histograms like this:

Which indicates that it is basically the same electronic density (differences < 7e-7).

Here's the code used to plot differences

import sisl
import numpy as np
import plotly.express as px

def compare_siesta_sisl(fdf):
  
  DM = fdf.read_density_matrix()
  
  # RHO computed by SIESTA
  rho_SIESTA = fdf.read_grid("RHO")
  
  # Compute the density with sisl
  sisl_grid = DM.density(rho_SIESTA)
  
  # Compute a histogram of the difference
  counts, edges = np.histogram(abs((sisl_grid - rho_SIESTA)).grid, bins=100)

  # Plot it and show it
  px.bar(x=edges[:-1], y=counts / sisl_grid.size).update_layout(
      title="Histogram of difference between sisl and SIESTA RHO",
      xaxis_title="Difference", yaxis_title="Fraction of points"
  ).show()
  
fdf = sisl.get_sile("my.fdf")
compare_siesta_sisl(fdf) 

Now that I know that it works, I went on to check the performance. For this, I splitted the contributions of computing the orbital values and the density computation, given that one of the goals is to use the same orbital values to call DM.density multiple times. This is the scalability I get with grid points:

So, both computations seem to scale close to linearly and are quite cheap, specially if we compare it with the current implementation, which I include here for reference using a logarithmic time scale:

For 1000 grid points in this system the current implementation takes 46s while the new one takes around 500ms, so an improvement of almost two orders of magnitude.

What is the source of this improvement? I found two main points:

• Of course, the implementation in cython of the loops that compute products makes it run much faster.
• But also, the way indices inside a sphere are retrieved in the current implementation (grid.index(ia_atom.toSphere(ia_xyz))) seems to be very inneficient. By changing it to the approach in the new implementation (getting a bounding box and then filtering by radius) I was able to improve the performance of the current implementation from 46s to 12s. The new method is probably a bit worse for skewed grids, where the bounding box is overestimated, but still I expect that it should be better.
• Finally, I think getting values of the sparse DM by querying with indices is probably also slowing down the current implementation.

Here's the code used to plot the scalability

import sisl
import numpy as np
import time
import plotly.graph_objs as go

def time_density(DM, grid_shape):
    geometry = DM.geometry
    
    # Initialize grid
    grid = sisl.Grid(grid_shape, geometry=geometry)
    
    # Compute orbital values and time it
    start_psi = time.time()
    psi_values = geometry.orbital_values(grid.shape)
    end_psi = time.time()
    
    # Use orbital values to compute density and time it
    start_dens = time.time()
    DM.density(grid, psi_values=psi_values)
    end_dens = time.time()
    
    # Return both times.
    return (end_psi - start_psi, end_dens - start_dens)

def scalability_density(fdf, ref_old=False):
    
    DM = fdf.read_density_matrix()
    
    sizes = []
    psi_times = []
    dens_times = []

    for shape in np.arange(5, 50, 5):
        psi_time, dens_time = time_density(DM, grid_shape=(shape, shape, shape))

        sizes.append(shape**3)
        psi_times.append(psi_time)
        dens_times.append(dens_time)
    
    if ref_old:
        # Get the old performance for a coarse grid, for reference
        old_grid_shape = (10,10,10)
        old_grid = sisl.Grid(old_grid_shape, geometry=DM.geometry)

        start = time.time()
        DM.old_density(old_grid)
        end = time.time()

        old_time = end - start
        old_size = np.prod(old_grid_shape)
    
    fig = go.Figure()
    
    fig.add_scatter(
        x=sizes, y=np.array(psi_times) + dens_times, name="Total"
    ).add_scatter(
        x=sizes, y=dens_times, name="Density"
    ).add_scatter(
        x=sizes, y=psi_times, name="Orbitals on grid"
    ).update_layout(
        xaxis_title="Number of grid points",
        yaxis_title="Elapsed time (s)",
        title="Scaling with grid size",
        yaxis_showgrid=True, xaxis_showgrid=True
    )
    
    if ref_old:
        fig.add_scatter(
            x=[old_size], y=[old_time], name="Previous implementation"
        )
        fig.update_layout(yaxis_type="log", yaxis_title="log(Elapsed time) (s)")
    
    fig.show()
    
    return fig

fdf = sisl.get_sile("./Final_coor/mdsamp-2.0V/1000/Au_Water_MD.fdf")
fig = scalability_density(fdf)

Performance for computing wavefunctions

For wavefunctions, I also checked with the differences histogram that I get the same results as with the current implementation for gamma point and k-point calculations.

The performance of the current implementation and the implementation through stored psi values is very similar (the new implementation is slightly more expensive, about 10%). This is to be expected, since the expensive part is computing the orbital values, which is done in a similar way in both cases. However the new implementation uses much more memory, since it stores the values instead of adding them on the fly. This is why I kept the current implementation as the default. Computing the wavefunction through stored psi values should only be done when you are computing multiple wavefunctions, or you have already computed the psi values for something else. If that is the case, calculating wavefunctions is of course very cheap in terms of time.

Here's the code used to compare wavefunctions with both methods

import sisl
import numpy as np
import plotly.express as px

def compare_wf_strategies(H, k, grid_shape):
    
    geometry = H.geometry
    geometry.xyz = geometry.translate2uc().xyz
    
    eig = H.eigenstate(k=k)[1]

    wf_grid = sisl.Grid(grid_shape, geometry=geometry, dtype=eig.dtype)
    eig.wavefunction(wf_grid)
    
    psi_values = geometry.orbital_values(wf_grid.shape, pbc=True)
    new_wf_grid = eig.wavefunction(wf_grid, psi_values=psi_values)
    
    # Compute a histogram of the difference
    counts, edges = np.histogram(abs(new_wf_grid - wf_grid).grid, bins=100)

    # Plot it and show it
    px.bar(x=edges[:-1], y=counts / wf_grid.size).update_layout(
        title="Histogram of difference between different wavefunction strategies",
        xaxis_title="Difference", yaxis_title="Fraction of points"
    ).show()

fdf = sisl.get_sile("my.fdf")
compare_wf_strategies(fdf.read_hamiltonian(), k=(0.5,0,0), grid_shape=(10, 10, 10))  

Looking forward to your thoughts on this!

[pfebrer]

By the way you need cython 3.0 to compile it, which is still an alpha version, so you can install it with pip install --pre Cython

[zerothi]

By the way you need cython 3.0 to compile it, which is still an alpha version, so you can install it with pip install --pre Cython

I'll have a look! Thanks for this.

I am not too happy about the alpha dependency... What feature is it that you need? Can we use older cython and update when 3 comes out?

[pfebrer]

Cython 3.0 has been in alpha for like 2 years I think and it is now in its 11th iteration. I think it is pretty safe to use. Not using it is like using siesta 4.0 instead of 4.1b4 Anyway this is not a dependency that you will impose on users, since you will provide the sources already compiled. So I don't think it's that dramatic.

The problem is that fused types in the pure python syntax don't work in 0.29, so I would have to rewrite everything with cython syntax.

[zerothi]

I must admit, if scipy isn't doing this, then I would also like to wait... I would like to follow scipy here... Also given that there are so many bugs listed in the 3.0 road map of cython. While I partially agree that it should work, I think it is easier for many of the partial devs to not rely on alpha releases of cython. Sorry.

[pfebrer]

Ok, it now works with 0.29.32. I didn't know you could mix both syntaxes. Since you can, this is an easy fix that doesn't annoy me too much and is easy to revert.

I have also created an issue in cython to see if they can solve it in 0.29: cython/cython#5085

[lgtm-com]

This pull request introduces 2 alerts when merging d93d027 into c5f7d5d - view on LGTM.com

new alerts:

• 1 for Unused import
• 1 for Unreachable code

[pfebrer]

I've done some changes. I realized that if I stored the sparse grid as a CSR matrix where rows are the elements of the grid (if it is a 3D grid, it has been raveled to 1D) and the columns are the extra dimension it would be benefitial because:

• I don't have to create and maintain a new sparse structure.
• I can store orbitals in columns by their sc index, following the conventions in the rest of sisl.
• It requires less memory than my previous implementation. Because of this it also got a little faster.
• Projecting the wavefunctions now becomes a simple matrix multiplication (psi X state), where state can be 2D resulting in 1D grid elements (as opposed to 0D or scalar).

The structure is cleaner now, so generic routines can be implemented for useful things and they don't look ultra-complex. For example I implemented:

• Orbital products can take a sisl sparse matrix with dim > 1 and perform all orbital products in the same go, resulting in 1D grid elements. This is currently implemented by creating a dense matrix containing the pointers to nonzero data, so that the access can be as fast as possible. I also have a not so slow algorithm for avoiding creating a dense array (which is not in the PR, yet), but is still much slower.
• Reduction of the grid dimensionality while projecting the values to save memory. This plays nicely with the 1D grid elements. For example, I wanted to reduce the X and Y axes to keep only Z. If I can't reduce them on the fly, the memory requirements blow up if I have 1D grid elements, but they are below those of a regular 3D grid if I use on the fly reduction.

I also added lots of tests, so I now feel much more confident about its correctness and it will be safer to do modifications on it.

There is one main problem though: I don't know how to make the functionality here play nicely with the current API for density and wavefunction. The current API accepts a grid that modifies in place. However, if you have stored the psi values already in a SparseGrid4D, to me it makes more sense that this grid has a method, e.g. reduce_orbitals, that returns a new grid. I would like to hear your thoughts on how would you manage this. Also, because ND grids still do not exist in sisl, when multiple coefficients are computed I need to return a numpy array.

Finally here are some plots of using multiple density matrices in the same product computation vs using them one by one. The plots are for different grid shapes, for a system with 3000 orbitals:

Overall I think it's very nice because it has allowed me to calculate LDOS(E) maps for huge systems and 240 energy points in less than 3 minutes, and before this it took around 4-5h :) This could enable an LDOS method in tbtncSile.

[lgtm-com]

This pull request introduces 1 alert when merging d46f65e into c5f7d5d - view on LGTM.com

new alerts:

• 1 for Unreachable code

[zerothi]

I'll have another look when you say so, thanks!!!!

[pfebrer]

Great! I'll try to amend things as we discussed.

For now I was just making the branch functional again because I needed to use it.

[pfebrer]

If this went in, I'm not sure it is more expensive to compute the LDOS directly by squaring wavefunctions. It would be something like this:

$$\textbf{LDOS}(x, E) = ( \textbf{V} \textbf{C})^{2_{el-wise}}\textbf{N}$$

Where:

• C is the eigenstate coefficients matrix (orbitals X wavefunctions).
• V is the sparse matrix implemented in this PR storing the orbital values (point grids X orbitals).
• 2_el-wise represents the element-wise squaring of the VC matrix (to get the squared wavefunctions).
• N is the occupations matrix (wavefunctions X energies).

V and N are sparse, but C and VC are not.

You end up with a matrix LDOS (point grids X energies). I'm not sure this is more expensive than the DM way, specially in python since these are simple matrix multiplications that are already implemented in scipy/numpy and which are probably already implemented in GPUs, versus the custom code for finding products of orbitals that you need to compute the LDOS from the DM. In terms of memory, I don't know what's more expensive.

Also, C can contain a third dimension for the k points, and doing the same operation you end up with LDOS(x, E, k). I'm sure this is an operation GPUs would love doing :)

[pfebrer]

I noticed that $V^TV$ ($V$ being the matrix of psi values) was the overlap matrix.

I know this is not the most efficient way, it is faster use fourier transforms. But it was easy to add and since there is no way currently of computing the overlap in sisl I just added it here 😅

[pfebrer]

I also leave this here. A working example of how to compute the LDOS using the proposal that I explained two comments above. Even though the occupations could have been a sparse matrix and they aren't, it was surprisingly fast. It allowed me to compute a high resolution full LDOS spectrum of a DZP water molecule in 1s.

import sisl
import numpy as np
from pathlib import Path
import plotly.express as px

# Get fdf
fdf = sisl.get_sile("waterDZP.fdf")

# Read hamiltonian and get geometry
H = fdf.read_hamiltonian()
geom = H.geometry

# Compute eigenstates
eig = H.eigenstate()

# Compute psi values on a very fine grid
shape = (100, 300, 100)
psi_values = geom.orbital_values(shape, pbc=True)

# Compute wavefunctions on the grid
VC = psi_values.reduce_orbitals(eig.state.T)
# Square them
VC_2 = VC ** 2

# Get the occupations matrix (wavefunction X energy point)
gaussian = sisl.get_distribution("gaussian", smearing=0.2)
E = np.linspace(eig.eig[0] - 2, eig.eig[-1] + 2, 800)
N = np.array([gaussian(eig.eig - point_E) for point_E in E]).T

# Sum wavefunctions over X and Z, then multiply by occupations to get the LDOS
# on the Y direction
LDOS_y = VC_2.sum(axis=0).sum(axis=1) @ N

# Plot the LDOS
fig = px.imshow(
    LDOS_y.T,
    x=np.linspace(0, geom.cell[1,1], LDOS_y.shape[0]), 
    y=E,
    aspect="data",
).update_layout(
    title="LDOS of DZP water molecule",
    height=900, 
    yaxis_title="Energy [eV]",
    xaxis_title="Y axis [Ang]"
)

# Plot some lines to see where the atoms are.
colors = ["red", "gray", "gray"]
for i, y in enumerate(geom.xyz[:, 1]):
    fig.add_vline(x=y, line_color=colors[i])

fig

:)

[pfebrer]

This is now ready. As we discussed a long time ago, the only user-visible change is the addition of a method argument to DensityMatrix.density. The rest of things are private for now until we decide how to expose them exactly.

pytest seems to agree that pre-computing the orbital values leads to much faster density computation :)

[codecov]

Codecov Report

Attention: Patch coverage is 57.89474% with 304 lines in your changes missing coverage. Please review.

Project coverage is 87.07%. Comparing base (050126a) to head (c0986ad).

❗ Current head c0986ad differs from pull request most recent head 718b7ab

Please upload reports for the commit 718b7ab to get more accurate results.

Files	Patch %	Lines
src/sisl/_sparse_grid_ops.py	0.00%	263 Missing ⚠️
src/sisl/_sparse_grid.py	78.06%	34 Missing ⚠️
src/sisl/_core/geometry.py	92.13%	7 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #496      +/-   ##
==========================================
- Coverage   87.49%   87.07%   -0.42%     
==========================================
  Files         397      400       +3     
  Lines       50972    51651     +679     
==========================================
+ Hits        44597    44975     +378     
- Misses       6375     6676     +301     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[zerothi]

I'll have a look at this in the coming weeks, shouldn't we move this to pure-python mode? Or?

[pfebrer]

I'll have a look at this in the coming weeks, shouldn't we move this to pure-python mode? Or?

Yes, I can convert it

[zerothi]

I will merge this, but note that the methods can and probably will change at will in the future.

There are some design choices that needs some coalescing.

Thanks!

[pfebrer]

Thank youuu!!!

Yes, that's why everything is private except for the method argument on DensityMatrix.density. That is probably the only thing we should worry about making stable. Do you agree with the choice or is there some other option?

[zerothi]

You mean the method argument? In general it shouldn't be there, but perhaps it is ok for now...

[pfebrer]

What do you mean? How would you do it if there are two methods that do the same thing?

One of them demands more computation and the other demands more memory.

[zerothi]

Ah, that aspect, yeah, that is the problem of this method, ideally it should default to remove the memory after use. Exactly how this should work is not clear at this moment. So, for now I guess this is ok.

[pfebrer]

What do you mean "after use"? Memory is cleared once you have finished computing all the products.

This is unavoidable if you don't want to recompute the orbital values each time.

[zerothi]

What do you mean "after use"? Memory is cleared once you have finished computing all the products.

This is unavoidable if you don't want to recompute the orbital values each time.

sorry, I had pre-assumed the orbital values were stored on the geometry, in any case I think the way this is handled could be done a bit differently. Perhaps by having an internal slice of the geometry that gets queried by a batch size so that we only calculate for a subset of the grid. In this way one will have dynamic memory, but also a tad more complicated. Well, i think we should get 0.15 out soonish. this has a somewhat low priority. What is important is that we fix supercell issues, and get those things back on track...

[pfebrer]

sorry, I had pre-assumed the orbital values were stored on the geometry

Ah, no, _orbital_values returns an object containing all the values. And that is used in density to compute the density. Then it is garbage collected.

Perhaps by having an internal slice of the geometry that gets queried by a batch size so that we only calculate for a subset of the grid.

Yes, but that would be much more complicated. I have tested some big systems and there are no problems with memory. I guess if you run out of memory you can just use the direct method. That's why I left the option to use one or the other.

[pfebrer]

Should I add the speed up to the changelog?

[zerothi]

Should I add the speed up to the changelog?

I am still cautious on promoting these methods (density, wavefunction because of the previous problems I have had, I still have an issue #592 that isn't fully solved). So there are some fundamental things that we need to resolve. To add to this, that might be better solved once #553 and friends are also solved... Argh... time... ;-D

[pfebrer]

Do you have an example where the density is wrong? In this PR the density matrix is translated into the unit cell to compute the density, so with that most issues should be solved.

[zerothi]

Do you have an example where the density is wrong? In this PR the density matrix is translated into the unit cell to compute the density, so with that most issues should be solved.

Well, that might also cause errors because only the coordinates are translated, not the DM elements, and so there could be wrong-doings. This is my concern at least.

[pfebrer]

Nono, the matrix elements are translated as well.

At the beginning of the density method there is a line like

dm = self.translate2uc()

[zerothi]

see, I have no idea what I am doing ;-D

[zerothi]

I just recall I had some initial issues where i wasn't fully sure whether everything was correct. So my mind could play tricks.

[pfebrer]

I added this line on this PR, and I saw that it fixed some corner cases.

That's why I say that maybe things are solved in the density now.