1
0
mirror of https://github.com/TREX-CoE/qmckl.git synced 2024-11-03 20:54:09 +01:00
qmckl/org/qmckl_examples.org
2023-09-26 17:34:57 +02:00

26 KiB

Code examples

In this section, we provide hands-on examples to demonstrate the capabilities of the QMCkl library. We furnish code samples in C, Fortran, and Python, serving as exhaustive tutorials for effectively leveraging QMCkl. For simplicity, we assume that the wave function parameters are stored in a TREXIO file.

Overlap matrix in the MO basis

The focal point of this example is the numerical evaluation of the overlap matrix in the MO basis. Utilizing QMCkl, we approximate the MOs at discrete grid points to compute the overlap matrix \( S_{ij} \) as follows: \[ S_{ij} = \int \phi_i(\mathbf{r})\, \phi_j(\mathbf{r}) \text{d}\mathbf{r} \approx \sum_k \phi_i(\mathbf{r}_k)\, \phi_j(\mathbf{r}_k) \delta\mathbf{r} \]

The code starts by reading a wave function from a TREXIO file. This is accomplished using the qmckl_trexio_read function, which populates a qmckl_context with the necessary wave function parameters. The context serves as the primary interface for interacting with the QMCkl library, encapsulating the state and configurations of the system. Subsequently, the code retrieves various attributes such as the number of nuclei nucl_num and coordinates nucl_coord. These attributes are essential for setting up the integration grid.

The core of the example lies in the numerical computation of the overlap matrix. To achieve this, the code employs a regular grid in three-dimensional space, and the grid points are then populated into the qmckl_context using the qmckl_set_point function.

The MO values at these grid points are computed using the qmckl_get_mo_basis_mo_value function. These values are then used to calculate the overlap matrix through a matrix multiplication operation facilitated by the qmckl_dgemm function.

The code is also instrumented to measure the execution time for the MO value computation, providing an empirical assessment of the computational efficiency. Error handling is robustly implemented at each stage to ensure the reliability of the simulation.

In summary, this example serves as a holistic guide for leveraging the QMCkl library, demonstrating its ease of use. It provides a concrete starting point for researchers and developers interested in integrating QMCkl into their QMC code.

Python

In this example, we will compute numerically the overlap between the molecular orbitals:

\[ S_{ij} = \int \phi_i(\mathbf{r}) \phi_j(\mathbf{r}) \text{d}\mathbf{r} \sim \sum_{k=1}^{N} \phi_i(\mathbf{r}_k) \phi_j(\mathbf{r}_k) \delta \mathbf{r} \] \[ S_{ij} = \langle \phi_i | \phi_j \rangle \sim \sum_{k=1}^{N} \langle \phi_i | \mathbf{r}_k \rangle \langle \mathbf{r}_k | \phi_j \rangle \]

import numpy as np
import qmckl
trexio_filename = "..//share/qmckl/test_data/h2o_5z.h5"

context = qmckl.context_create()
qmckl.trexio_read(context, trexio_filename)

We now define the grid points $\mathbf{r}_k$ as a regular grid around the molecule.

We fetch the nuclear coordinates from the context,

nucl_num = qmckl.get_nucleus_num(context)

nucl_charge = qmckl.get_nucleus_charge(context, nucl_num)

nucl_coord = qmckl.get_nucleus_coord(context, 'N', nucl_num*3)
nucl_coord = np.reshape(nucl_coord, (3, nucl_num))

for i in range(nucl_num):
 print("%d  %+f %+f %+f"%(int(nucl_charge[i]),
                          nucl_coord[i,0],
                          nucl_coord[i,1],
                          nucl_coord[i,2]) )
8  +0.000000 +0.000000 +0.000000
1  -1.430429 +0.000000 -1.107157
1  +1.430429 +0.000000 -1.107157

and compute the coordinates of the grid points:

nx = ( 120, 120, 120 )
shift = np.array([5.,5.,5.])
point_num = nx[0] * nx[1] * nx[2]

rmin  = np.array( list([ np.min(nucl_coord[:,a]) for a in range(3) ]) )
rmax  = np.array( list([ np.max(nucl_coord[:,a]) for a in range(3) ]) )


linspace = [ None for i in range(3) ]
step     = [ None for i in range(3) ]
for a in range(3):
 linspace[a], step[a] = np.linspace(rmin[a]-shift[a],
                                    rmax[a]+shift[a],
                                    num=nx[a],
                                    retstep=True)

dr = step[0] * step[1] * step[2]
point = []
for x in linspace[0]:
 for y in linspace[1]:
     for z in linspace[2]:
         point += [ [x, y, z] ]

point = np.array(point)
point_num = len(point)
qmckl.set_point(context, 'N', point_num, np.reshape(point, (point_num*3)))

Then, we evaluate all the MOs at the grid points (and time the execution), and thus obtain the matrix $M_{ki} = \langle \mathbf{r}_k | \phi_i \rangle = \phi_i(\mathbf{r}_k)$.

import time

mo_num = qmckl.get_mo_basis_mo_num(context)

before   = time.time()
mo_value = qmckl.get_mo_basis_mo_value(context, point_num*mo_num)
after    = time.time()

mo_value = np.reshape( mo_value, (point_num, mo_num) ).T   # Transpose to get mo_num x point_num

print("Number of MOs: ", mo_num)
print("Number of grid points: ", point_num)
print("Execution time : ", (after - before), "seconds")
Number of MOs:  201
Number of grid points:  1728000
Execution time :  5.577778577804565 seconds

and finally we compute the overlap between all the MOs as $M.M^\dagger$.

overlap = mo_value @ mo_value.T * dr
print (overlap)
[[ 9.88693941e-01  2.34719693e-03 -1.50518232e-08 ...  3.12084178e-09
  -5.81064929e-10  3.70130091e-02]
 [ 2.34719693e-03  9.99509628e-01  3.18930040e-09 ... -2.46888958e-10
  -1.06064273e-09 -7.65567973e-03]
 [-1.50518232e-08  3.18930040e-09  9.99995073e-01 ... -5.84882580e-06
  -1.21598117e-06  4.59036468e-08]
 ...
 [ 3.12084178e-09 -2.46888958e-10 -5.84882580e-06 ...  1.00019107e+00
  -2.03342837e-04 -1.36954855e-08]
 [-5.81064929e-10 -1.06064273e-09 -1.21598117e-06 ... -2.03342837e-04
   9.99262427e-01  1.18264754e-09]
 [ 3.70130091e-02 -7.65567973e-03  4.59036468e-08 ... -1.36954855e-08
   1.18264754e-09  8.97215950e-01]]

C

In this example, electron-nucleus cusp fitting is added.

In this example, we will compute numerically the overlap between the molecular orbitals:

\[ S_{ij} = \int \phi_i(\mathbf{r}) \phi_j(\mathbf{r}) \text{d}\mathbf{r} \sim \sum_{k=1}^{N} \phi_i(\mathbf{r}_k) \phi_j(\mathbf{r}_k) \delta \mathbf{r} \] \[ S_{ij} = \langle \phi_i | \phi_j \rangle \sim \sum_{k=1}^{N} \langle \phi_i | \mathbf{r}_k \rangle \langle \mathbf{r}_k | \phi_j \rangle \]

We apply the cusp fitting procedure, so the MOs might deviate slightly from orthonormality.

#include <qmckl.h>
#include <stdio.h>
#include <string.h>
#include <sys/time.h>

int main(int argc, char** argv)
{
const char* trexio_filename = "..//share/qmckl/test_data/h2o_5z.h5";
qmckl_exit_code rc = QMCKL_SUCCESS;

First, we create a context for the QMCkl calculation, and load the wave function stored in h2o_5z.h5 inside it. It is a Hartree-Fock determinant for the water molecule in the cc-pV5Z basis set.

qmckl_context context = qmckl_context_create();

rc = qmckl_trexio_read(context, trexio_filename, strlen(trexio_filename));

if (rc != QMCKL_SUCCESS) {
 fprintf(stderr, "Error reading TREXIO file:\n%s\n", qmckl_string_of_error(rc));
 exit(1);
}

We impose the electron-nucleus cusp fitting to occur when the electrons are close to the nuclei. The critical distance is 0.5 atomic units for hydrogens and 0.1 for the oxygen. To identify which atom is an oxygen and which are hydrogens, we fetch the nuclear charges from the context.

int64_t nucl_num;

rc = qmckl_get_nucleus_num(context, &nucl_num);

if (rc != QMCKL_SUCCESS) {
 fprintf(stderr, "Error getting nucl_num:\n%s\n", qmckl_string_of_error(rc));
 exit(1);
}


double nucl_charge[nucl_num];

rc = qmckl_get_nucleus_charge(context, &(nucl_charge[0]), nucl_num);

if (rc != QMCKL_SUCCESS) {
 fprintf(stderr, "Error getting nucl_charge:\n%s\n", qmckl_string_of_error(rc));
 exit(1);
}


double r_cusp[nucl_num];

for (size_t i=0 ; i<nucl_num ; ++i) {

 switch ((int) nucl_charge[i]) {

 case 1:
   r_cusp[i] = 0.5;
   break;

 case 8:
   r_cusp[i] = 0.1;
   break;
 }

}


rc = qmckl_set_mo_basis_r_cusp(context, &(r_cusp[0]), nucl_num);

if (rc != QMCKL_SUCCESS) {
 fprintf(stderr, "Error setting r_cusp:\n%s\n", qmckl_string_of_error(rc));
 exit(1);
}

We now define the grid points $\mathbf{r}_k$ as a regular grid around the molecule. We fetch the nuclear coordinates from the context,

double nucl_coord[nucl_num][3];

rc = qmckl_get_nucleus_coord(context, 'N', &(nucl_coord[0][0]), nucl_num*3);

if (rc != QMCKL_SUCCESS) {
 fprintf(stderr, "Error getting nucl_coord:\n%s\n", qmckl_string_of_error(rc));
 exit(1);
}

for (size_t i=0 ; i<nucl_num ; ++i) {
 printf("%d  %+f %+f %+f\n",
        (int32_t) nucl_charge[i],
        nucl_coord[i][0],
        nucl_coord[i][1],
        nucl_coord[i][2]);
}
8  +0.000000 +0.000000 +0.000000
1  -1.430429 +0.000000 -1.107157
1  +1.430429 +0.000000 -1.107157

and compute the coordinates of the grid points:

size_t nx[3] = { 120, 120, 120 };
double shift[3] = {5.,5.,5.};
int64_t point_num = nx[0] * nx[1] * nx[2];

double rmin[3] = { nucl_coord[0][0], nucl_coord[0][1], nucl_coord[0][2] } ;
double rmax[3] = { nucl_coord[0][0], nucl_coord[0][1], nucl_coord[0][2] } ;

for (size_t i=0 ; i<nucl_num ; ++i) {
 for (int j=0 ; j<3 ; ++j) {
   rmin[j] = nucl_coord[i][j] < rmin[j] ? nucl_coord[i][j] : rmin[j];
   rmax[j] = nucl_coord[i][j] > rmax[j] ? nucl_coord[i][j] : rmax[j];
 }
}

rmin[0] -= shift[0]; rmin[1] -= shift[1]; rmin[2] -= shift[2];
rmax[0] += shift[0]; rmax[1] += shift[1]; rmax[2] += shift[2];

double step[3];

double* linspace[3];
for (int i=0 ; i<3 ; ++i) {

 linspace[i] = (double*) calloc( sizeof(double), nx[i] );

 if (linspace[i] == NULL) {
   fprintf(stderr, "Allocation failed (linspace)\n");
   exit(1);
 }

 step[i] = (rmax[i] - rmin[i]) / ((double) (nx[i]-1));

 for (size_t j=0 ; j<nx[i] ; ++j) {
   linspace[i][j] = rmin[i] + j*step[i];
 }

}

double dr = step[0] * step[1] * step[2];

Now the grid is ready, we can create the list of grid points $\mathbf{r}_k$ on which the MOs $\phi_i$ will be evaluated, and transfer them to the QMCkl context:

double* point = (double*) calloc(sizeof(double), 3*point_num);

if (point == NULL) {
 fprintf(stderr, "Allocation failed (point)\n");
 exit(1);
}

size_t m = 0;
for (size_t i=0 ; i<nx[0] ; ++i) {
 for (size_t j=0 ; j<nx[1] ; ++j) {
   for (size_t k=0 ; k<nx[2] ; ++k) {

     point[m] = linspace[0][i];
     m++;

     point[m] = linspace[1][j];
     m++;

     point[m] = linspace[2][k];
     m++;

   }
 }
}

rc = qmckl_set_point(context, 'N', point_num, point, (point_num*3));

if (rc != QMCKL_SUCCESS) {
 fprintf(stderr, "Error setting points:\n%s\n", qmckl_string_of_error(rc));
 exit(1);
}

Then, we evaluate all the MOs at the grid points (and time the execution), and thus obtain the matrix $M_{ki} = \langle \mathbf{r}_k | \phi_i \rangle = \phi_i(\mathbf{r}_k)$.

int64_t mo_num;
rc = qmckl_get_mo_basis_mo_num(context, &mo_num);

long before, after;
struct timeval timecheck;

double* mo_value = (double*) calloc(sizeof(double), point_num*mo_num);
if (mo_value == NULL) {
 fprintf(stderr, "Allocation failed (mo_value)\n");
 exit(1);
}

gettimeofday(&timecheck, NULL);
before = (long)timecheck.tv_sec * 1000 + (long)timecheck.tv_usec / 1000;

rc = qmckl_get_mo_basis_mo_value(context, mo_value, point_num*mo_num);
if (rc != QMCKL_SUCCESS) {
 fprintf(stderr, "Error getting mo_value)\n");
 exit(1);
}

gettimeofday(&timecheck, NULL);
after = (long)timecheck.tv_sec * 1000 + (long)timecheck.tv_usec / 1000;

printf("Number of MOs: %ld\n", (long) mo_num);
printf("Number of grid points: %ld\n", (long) point_num);
printf("Execution time : %f seconds\n", (after - before)*1.e-3);
Number of MOs:  201
Number of grid points:  1728000
Execution time :  5.608000 seconds

and finally we compute the overlap between all the MOs as $M.M^\dagger$.

double* overlap = (double*) malloc (mo_num*mo_num*sizeof(double));

rc = qmckl_dgemm(context, 'N', 'T', mo_num, mo_num, point_num, dr,
                mo_value, mo_num, mo_value, mo_num, 0.0,
                overlap, mo_num);

for (size_t i=0 ; i<mo_num ; ++i) {
 printf("%4ld", i);
 for (size_t j=0 ; j<mo_num ; ++j) {
   printf(" %f", overlap[i*mo_num+j]);
 }
 printf("\n");
}

// Clean-up and exit
free(mo_value);
free(overlap);

rc = qmckl_context_destroy(context);
if (rc != QMCKL_SUCCESS) {
 fprintf(stderr, "Error destroying context)\n");
 exit(1);
}

return 0;
}
0 0.988765 0.002336 0.000000 -0.000734 0.000000 0.000530 0.000000 0.000446 0.000000 -0.000000 -0.000511 -0.000000 -0.000267 0.000000 0.000000 0.001007 0.000000 0.000168 -0.000000 -0.000000 -0.000670 -0.000000 0.000000 -0.000251 -0.000261 -0.000000 -0.000000 -0.000000 -0.000397 -0.000000 -0.000810 0.000000 0.000231 -0.000000 -0.000000 0.000000 -0.000000
...
200 0.039017 -0.013066 -0.000000 -0.001935 -0.000000 -0.003829 -0.000000 0.000996 -0.000000 0.000000 -0.003733 0.000000 0.000065 -0.000000 -0.000000 -0.002220 -0.000000 -0.001961 0.000000 0.000000 -0.004182 0.000000 -0.000000 -0.000165 -0.002445 0.000000 -0.000000 0.000000 0.001985 0.000000 0.001685 -0.000000 -0.002899 0.000000 0.000000 0.000000 -0.000000 0.002591 0.000000 -0.000000 0.000000 0.002385 0.000000 -0.011086 0.000000 -0.003885 0.000000 -0.000000 0.003602 -0.000000 0.000000 -0.003241 0.000000 0.000000 0.002613 -0.007344 -0.000000 -0.000000 0.000000 0.000017 0.000000 0.000000 0.000000 -0.008719 0.000000 -0.001358 -0.003233 0.000000 -0.000000 -0.000000 -0.000000 0.000000 0.003778 0.000000 0.000000 -0.000000 0.000000 0.000000 -0.001190 0.000000 0.000000 -0.000000 0.005578 -0.000000 -0.001502 0.003899 -0.000000 -0.000000 0.000000 -0.003291 -0.001775 -0.000000 -0.002374 0.000000 -0.000000 -0.000000 -0.000000 -0.001290 -0.000000 0.002178 0.000000 0.000000 0.000000 -0.001252 0.000000 -0.000000 -0.000926 0.000000 -0.000000 -0.013130 -0.000000 0.012124 0.000000 -0.000000 -0.000000 -0.000000 0.000000 0.025194 0.000343 -0.000000 0.000000 -0.000000 -0.004421 0.000000 0.000000 -0.000599 -0.000000 0.005289 0.000000 -0.000000 0.012826 -0.000000 0.000000 0.006190 0.000000 0.000000 -0.000000 0.000000 -0.000321 0.000000 -0.000000 -0.000000 0.000000 -0.000000 0.001499 -0.006629 0.000000 0.000000 0.000000 -0.000000 0.008737 -0.000000 0.006880 0.000000 -0.001851 -0.000000 -0.000000 0.000000 -0.007464 0.000000 0.010298 -0.000000 -0.000000 -0.000000 -0.000000 -0.000000 0.000000 0.000540 0.000000 -0.006616 -0.000000 0.000000 -0.002927 -0.000000 0.000000 0.010352 0.000000 -0.003103 -0.000000 -0.007640 -0.000000 -0.000000 0.005302 0.000000 0.000000 -0.000000 -0.000000 -0.010181 0.000000 -0.001108 0.000000 0.000000 -0.000000 0.000000 0.000000 -0.000998 -0.009754 0.013562 0.000000 -0.000000 0.887510

Fortran

Here is the same piece of code translated in Fortran

#include <qmckl_f.F90>

program main
use iso_c_binding
use qmckl
implicit none

! Declare variables
integer :: argc
character(128) :: trexio_filename, err_msg
integer(qmckl_exit_code) :: rc
integer(qmckl_context) :: context
integer*8 :: nucl_num, mo_num, point_num
double precision, allocatable :: nucl_coord(:,:)
integer*8  :: nx(3)
double precision, dimension(3) :: shift, step, rmin, rmax
double precision, allocatable :: mo_value(:,:), overlap(:,:), point(:), linspace(:,:)
double precision :: before, after, dr
integer*8 :: i,j,k,m

! Initialize variables
err_msg = ' ' 
argc = command_argument_count()
if (argc /= 1) then
 print *, "Usage: ./program <TREXIO filename>"
 stop -1
end if
call get_command_argument(1, trexio_filename)
rc = QMCKL_SUCCESS

! Create a QMCkl context
context = qmckl_context_create()

! Read the TREXIO file into the context
rc = qmckl_trexio_read(context, trim(trexio_filename), len(trexio_filename)*1_8)
if (rc /= QMCKL_SUCCESS) then
 call qmckl_string_of_error(rc, err_msg)
 write(*,*) "Error reading TREXIO file:", err_msg
 stop -1
end if

! Retrieve the number of nuclei
rc = qmckl_get_nucleus_num(context, nucl_num)
if (rc /= QMCKL_SUCCESS) then
 call qmckl_string_of_error(rc, err_msg)
 write(*,*) "Error getting nucl_num:", err_msg
 stop -1
end if

! Retrieve the nuclear coordinates
allocate(nucl_coord(3, nucl_num))
rc = qmckl_get_nucleus_coord(context, 'N', nucl_coord, nucl_num * 3_8)
if (rc /= QMCKL_SUCCESS) then
 call qmckl_string_of_error(rc, err_msg)
 write(*,*) "Error getting nucl_coord:", err_msg
 stop -1
end if

! Retrieve the number of MOs
rc = qmckl_get_mo_basis_mo_num(context, mo_num)
if (rc /= QMCKL_SUCCESS) then
 call qmckl_string_of_error(rc, err_msg)
 write(*,*) "Error getting mo_num:", err_msg
 stop -1
end if

! Initialize grid points for the calculation
nx = (/ 120, 120, 120 /)
shift = (/ 5.d0, 5.d0, 5.d0 /)
point_num = nx(1) * nx(2) * nx(3)

! Initialize rmin and rmax
rmin = nucl_coord(:,1)
rmax = nucl_coord(:,1)

! Update rmin and rmax based on nucl_coord
do i = 1, 3
 do j = 1, nucl_num
   rmin(i) = min(nucl_coord(i,j), rmin(i))
   rmax(i) = max(nucl_coord(i,j), rmax(i))
 end do
end do

! Apply shift
rmin = rmin - shift
rmax = rmax + shift

! Initialize linspace and step
allocate(linspace(3, maxval(nx)))

do i = 1, 3
 step(i) = (rmax(i) - rmin(i)) / real(nx(i) - 1, 8)
 do j = 1, nx(i)
   linspace(i, j) = rmin(i) + (j - 1) * step(i)
 end do
end do

! Initialize point array
allocate(point(3 * point_num))
m = 1
do i = 1, nx(1)
 do j = 1, nx(2)
   do k = 1, nx(3)
     point(m) = linspace(1, i); m = m + 1
     point(m) = linspace(2, j); m = m + 1
     point(m) = linspace(3, k); m = m + 1
   end do
 end do
end do

deallocate(linspace)


! Set points in QMCKL context
rc = qmckl_set_point(context, 'N', point_num, point, point_num * 3)
if (rc /= QMCKL_SUCCESS) then
 call qmckl_string_of_error(rc, err_msg)
 write(*,*) "Error setting point:", err_msg
 stop -1
end if




! Perform the actual calculation and measure the time taken
call cpu_time(before)
allocate(mo_value(point_num, mo_num))
rc = qmckl_get_mo_basis_mo_value(context, mo_value, point_num * mo_num)
if (rc /= QMCKL_SUCCESS) then
 call qmckl_string_of_error(rc, err_msg)
 write(*,*) "Error getting mo_value:", err_msg
 stop
end if
call cpu_time(after)

write(*,*) "Number of MOs:", mo_num
write(*,*) "Number of grid points:", point_num
write(*,*) "Execution time:", (after - before), "seconds"

! Compute the overlap matrix
dr = step(1) * step(2) * step(3)

allocate(overlap(mo_num, mo_num))
rc = qmckl_dgemm(context, 'N', 'T', mo_num, mo_num, point_num, dr, &
                mo_value, mo_num, mo_value, mo_num, 0.d0, overlap, mo_num)

! Print the overlap matrix
do i = 1, mo_num
 write(*,'(i4)', advance='no') i
 do j = 1, mo_num
   write(*,'(f8.4)', advance='no') overlap(i, j)
 end do
 write(*,*)
end do

! Clean-up and exit
deallocate(mo_value, overlap)
rc = qmckl_context_destroy(context)
if (rc /= QMCKL_SUCCESS) then
 call qmckl_string_of_error(rc, err_msg)
 write(*,*) "Error destroying context:", err_msg
 stop -1
end if

end program main

Fortran

Checking errors

All QMCkl functions return an error code. A convenient way to handle errors is to write an error-checking function that displays the error in text format and exits the program.

subroutine qmckl_check_error(rc, message)
use qmckl
implicit none
integer(qmckl_exit_code), intent(in) :: rc
character(len=*)        , intent(in) :: message
character(len=128)                   :: str_buffer
if (rc /= QMCKL_SUCCESS) then
  print *, message
  call qmckl_string_of_error(rc, str_buffer)
  print *, str_buffer
  call exit(rc)
end if
end subroutine qmckl_check_error

Computing an atomic orbital on a grid

The following program, in Fortran, computes the values of an atomic orbital on a regular 3-dimensional grid. The 100^3 grid points are automatically defined, such that the molecule fits in a box with 5 atomic units in the borders.

This program uses the qmckl_check_error function defined above.

To use this program, run

$ ao_grid <trexio_file> <AO_id> <point_num>
<<qmckl_check_error>>

program ao_grid
use qmckl
implicit none

integer(qmckl_context)    :: qmckl_ctx  ! QMCkl context
integer(qmckl_exit_code)  :: rc         ! Exit code of QMCkl functions

character(len=128)            :: trexio_filename
character(len=128)            :: str_buffer
integer                       :: ao_id
integer                       :: point_num_x

integer(c_int64_t)            :: nucl_num
double precision, allocatable :: nucl_coord(:,:)

integer(c_int64_t)            :: point_num
integer(c_int64_t)            :: ao_num
integer(c_int64_t)            :: ipoint, i, j, k
double precision              :: x, y, z, dr(3)
double precision              :: rmin(3), rmax(3)
double precision, allocatable :: points(:,:)
double precision, allocatable :: ao_vgl(:,:,:)

Start by fetching the command-line arguments:

if (iargc() /= 3) then
  print *, 'Syntax: ao_grid <trexio_file> <AO_id> <point_num>'
  call exit(-1)
end if
call getarg(1, trexio_filename)
call getarg(2, str_buffer)
read(str_buffer, *) ao_id
call getarg(3, str_buffer)
read(str_buffer, *) point_num_x

if (point_num_x < 0 .or. point_num_x > 300) then
  print *, 'Error: 0 < point_num < 300'
  call exit(-1)
end if

Create the QMCkl context and initialize it with the wave function present in the TREXIO file:

qmckl_ctx = qmckl_context_create()
rc  = qmckl_trexio_read(qmckl_ctx, trexio_filename, 1_8*len(trim(trexio_filename)))
call qmckl_check_error(rc, 'Read TREXIO')

We need to check that ao_id is in the range, so we get the total number of AOs from QMCkl:

rc = qmckl_get_ao_basis_ao_num(qmckl_ctx, ao_num)
call qmckl_check_error(rc, 'Getting ao_num')

if (ao_id < 0 .or. ao_id > ao_num) then
  print *, 'Error: 0 < ao_id < ', ao_num
  call exit(-1)
end if

Now we will compute the limits of the box in which the molecule fits. For that, we first need to ask QMCkl the coordinates of nuclei.

rc = qmckl_get_nucleus_num(qmckl_ctx, nucl_num)
call qmckl_check_error(rc, 'Get nucleus num')

allocate( nucl_coord(3, nucl_num) )
rc = qmckl_get_nucleus_coord(qmckl_ctx, 'N', nucl_coord, 3_8*nucl_num)
call qmckl_check_error(rc, 'Get nucleus coord')

We now compute the coordinates of opposite points of the box, and the distance between points along the 3 directions:

rmin(1) = minval( nucl_coord(1,:) ) - 5.d0
rmin(2) = minval( nucl_coord(2,:) ) - 5.d0
rmin(3) = minval( nucl_coord(3,:) ) - 5.d0

rmax(1) = maxval( nucl_coord(1,:) ) + 5.d0
rmax(2) = maxval( nucl_coord(2,:) ) + 5.d0
rmax(3) = maxval( nucl_coord(3,:) ) + 5.d0

dr(1:3) = (rmax(1:3) - rmin(1:3)) / dble(point_num_x-1)

We now produce the list of point coordinates where the AO will be evaluated:

point_num = point_num_x**3
allocate( points(point_num, 3) )
ipoint=0
z = rmin(3)
do k=1,point_num_x
  y = rmin(2)
  do j=1,point_num_x
     x = rmin(1)
     do i=1,point_num_x
        ipoint = ipoint+1
        points(ipoint,1) = x
        points(ipoint,2) = y
        points(ipoint,3) = z
        x = x + dr(1)
     end do
     y = y + dr(2)
  end do
  z = z + dr(3)
end do

We give the points to QMCkl:

rc = qmckl_set_point(qmckl_ctx, 'T', point_num, points, size(points)*1_8 )
call qmckl_check_error(rc, 'Setting points')

We allocate the space required to retrieve the values, gradients and Laplacian of all AOs, and ask to retrieve the values of the AOs computed at the point positions.

allocate( ao_vgl(ao_num, 5, point_num) )
rc = qmckl_get_ao_basis_ao_vgl(qmckl_ctx, ao_vgl, ao_num*5_8*point_num)
call qmckl_check_error(rc, 'Setting points')

We finally print the value and Laplacian of the AO:

do ipoint=1, point_num
  print '(3(F10.6,X),2(E20.10,X))', points(ipoint, 1:3), ao_vgl(ao_id,1,ipoint), ao_vgl(ao_id,5,ipoint)
end do
deallocate( nucl_coord, points, ao_vgl )
end program ao_grid