From 97e5392f5f19d3e84b36e73cb53d74a43d1f8558 Mon Sep 17 00:00:00 2001
From: Anthony Scemama <scemama@irsamc.ups-tlse.fr>
Date: Mon, 11 Oct 2021 18:15:31 +0200
Subject: [PATCH] Added 1 function in trexio

---
 org/qmckl_trexio.org | 95 +++++++++++++++++++++++++++++++++-----------
 1 file changed, 71 insertions(+), 24 deletions(-)

diff --git a/org/qmckl_trexio.org b/org/qmckl_trexio.org
index d83f937..e4c2eec 100644
--- a/org/qmckl_trexio.org
+++ b/org/qmckl_trexio.org
@@ -71,7 +71,6 @@ int main() {
   - ~rc~: the return code for QMCkl functions
   - ~rcio~: the return code for TREXIO functions.
   
-  
 ** Open file
    
    We first define a helper function to open a file by first trying to
@@ -164,9 +163,9 @@ qmckl_trexio_read_nucleus_X(qmckl_context context, trexio_t* const file)
 *** Number of nuclei
 
    #+begin_src c :tangle (eval c)
-  int64_t num = 0;
+  int64_t nucleus_num = 0;
 
-  rcio = trexio_read_nucleus_num_64(file, &num);
+  rcio = trexio_read_nucleus_num_64(file, &nucleus_num);
   if (rcio != TREXIO_SUCCESS) {
     return qmckl_failwith( context,
                            QMCKL_FAILURE,
@@ -174,8 +173,8 @@ qmckl_trexio_read_nucleus_X(qmckl_context context, trexio_t* const file)
                            trexio_string_of_error(rcio));
   }
   
-  assert (num > 0);
-  rc = qmckl_set_nucleus_num(context, num);
+  assert (nucleus_num > 0);
+  rc = qmckl_set_nucleus_num(context, nucleus_num);
 
   if (rc != QMCKL_SUCCESS)
     return rc;
@@ -186,7 +185,7 @@ qmckl_trexio_read_nucleus_X(qmckl_context context, trexio_t* const file)
    #+begin_src c :tangle (eval c)
   {
     qmckl_memory_info_struct mem_info = qmckl_memory_info_struct_zero;
-    mem_info.size = num * sizeof(double);
+    mem_info.size = nucleus_num * sizeof(double);
 
     double* nucl_charge = (double*) qmckl_malloc(context, mem_info);
 
@@ -226,7 +225,7 @@ qmckl_trexio_read_nucleus_X(qmckl_context context, trexio_t* const file)
    #+begin_src c :tangle (eval c)
   {
     qmckl_memory_info_struct mem_info = qmckl_memory_info_struct_zero;
-    mem_info.size = num * 3 * sizeof(double);
+    mem_info.size = nucleus_num * 3 * sizeof(double);
 
     double* nucl_coord = (double*) qmckl_malloc(context, mem_info);
 
@@ -379,7 +378,7 @@ qmckl_trexio_read_ao_X(qmckl_context context, trexio_t* const file)
 
    #+end_src
 
-*** Nucleus index array
+*** Nucleus_index array
 
    #+begin_src c :tangle (eval c)
   {
@@ -415,6 +414,42 @@ qmckl_trexio_read_ao_X(qmckl_context context, trexio_t* const file)
   }
    #+end_src
 
+*** Nucleus_shell_num array
+
+   #+begin_src c :tangle (eval c)
+  {
+    qmckl_memory_info_struct mem_info = qmckl_memory_info_struct_zero;
+    mem_info.size = nucleus_num * sizeof(int64_t);
+
+    int64_t* nucleus_shell_num = (int64_t*) qmckl_malloc(context, mem_info);
+
+    if (nucleus_shell_num == NULL) {
+      return qmckl_failwith( context,
+                             QMCKL_ALLOCATION_FAILED,
+                             "qmckl_trexio_read_basis_nucleus_shell_num_X",
+                             NULL);
+    }
+
+    assert (nucleus_shell_num != NULL);
+
+    rcio = trexio_read_basis_nucleus_shell_num_64(file, nucleus_shell_num);
+    if (rcio != TREXIO_SUCCESS) {
+      return qmckl_failwith( context,
+                             QMCKL_FAILURE,
+                             "trexio_read_basis_nucleus_shell_num",
+                             trexio_string_of_error(rcio));
+    }
+
+    rc = qmckl_set_ao_basis_nucleus_shell_num(context, nucleus_shell_num);
+
+    qmckl_free(context, nucleus_shell_num);
+    nucleus_shell_num = NULL;
+
+    if (rc != QMCKL_SUCCESS) 
+      return rc;
+  }
+   #+end_src
+
 
    #+begin_src c :tangle (eval c)
 
@@ -516,41 +551,44 @@ assert ( rc == QMCKL_SUCCESS );
 
   #+begin_src c :tangle (eval c_test)
 printf("Electrons\n");
-int64_t num;
-rc = qmckl_get_electron_up_num(context, &num);
+int64_t up_num, dn_num;
+rc = qmckl_get_electron_up_num(context, &up_num);
 assert (rc == QMCKL_SUCCESS);
-assert (num == chbrclf_elec_up_num);
+assert (up_num == chbrclf_elec_up_num);
   
-rc = qmckl_get_electron_down_num(context, &num);
+rc = qmckl_get_electron_down_num(context, &dn_num);
 assert (rc == QMCKL_SUCCESS);
-assert (num == chbrclf_elec_dn_num);
-
-rc = qmckl_get_nucleus_num(context, &num);
-assert (rc == QMCKL_SUCCESS);
-assert (num == chbrclf_nucl_num);
+assert (dn_num == chbrclf_elec_dn_num);
 
   #+end_src
   
 *** Nuclei
 
   #+begin_src c :tangle (eval c_test)
+printf("Nuclei\n");
+
+int64_t nucl_num;
+rc = qmckl_get_nucleus_num(context, &nucl_num);
+assert (rc == QMCKL_SUCCESS);
+assert (nucl_num == chbrclf_nucl_num);
+
 printf("Nuclear charges\n");
-double * charge = (double*) malloc (num * sizeof(double));
+double * charge = (double*) malloc (nucl_num * sizeof(double));
 rc = qmckl_get_nucleus_charge(context, charge);
 assert (rc == QMCKL_SUCCESS);
-for (int i=0 ; i<num ; i++) {
+for (int i=0 ; i<nucl_num ; i++) {
   assert (charge[i] == chbrclf_charge[i]);
  }
 free(charge);
 charge = NULL;
 
 printf("Nuclear coordinates\n");
-double * coord  = (double*) malloc (num * 3 * sizeof(double));
+double * coord  = (double*) malloc (nucl_num * 3 * sizeof(double));
 rc = qmckl_get_nucleus_coord(context, 'T', coord);
 assert (rc == QMCKL_SUCCESS);
 int k=0;
 for (int j=0 ; j<3 ; ++j) {
-  for (int i=0 ; i<num ; ++i) {
+  for (int i=0 ; i<nucl_num ; ++i) {
 //    printf("%f  %f\n", coord[k], chbrclf_nucl_coord[j][i]);
     assert (coord[k] == chbrclf_nucl_coord[j][i]);
     ++k;
@@ -582,12 +620,21 @@ int64_t ao_num;
 rc = qmckl_get_ao_basis_ao_num(context, &ao_num);
 assert (ao_num == chbrclf_ao_num);
 
-int64_t* nucleus_index = (int64_t*) malloc (shell_num * sizeof(int64_t));
-rc = qmckl_get_ao_basis_nucleus_index(context, nucleus_index, num);
+int64_t* nucleus_index = (int64_t*) malloc (nucl_num * sizeof(int64_t));
+rc = qmckl_get_ao_basis_nucleus_index(context, nucleus_index, nucl_num);
 assert (rc == QMCKL_SUCCESS);
-for (int i=0 ; i<num ; i++) {
+for (int i=0 ; i<nucl_num ; i++) {
   assert (nucleus_index[i] == chbrclf_basis_nucleus_index[i]);
  }
+
+int64_t* nucleus_shell_num  = (int64_t*) malloc (nucl_num * sizeof(int64_t));
+rc = qmckl_get_ao_basis_nucleus_shell_num(context, nucleus_shell_num, nucl_num);
+assert (rc == QMCKL_SUCCESS);
+for (int i=0 ; i<nucl_num ; i++) {
+  assert (nucleus_shell_num[i] == chbrclf_basis_nucleus_shell_num[i]);
+ }
+
+
 free(charge);
 charge = NULL;
 #endif