Helper functions to convert bit-wise determinants into lists of orbitals

src/pytrexio.i

@@ -36,6 +36,10 @@
 /* Return num variables as part of the output tuple */
 %apply int *OUTPUT { int32_t* const num};
 %apply int *OUTPUT { int64_t* const num};
+%apply int *OUTPUT { int32_t* const num_up};
+%apply int *OUTPUT { int32_t* const num_dn};
+%apply int *OUTPUT { int64_t* const num_up};
+%apply int *OUTPUT { int64_t* const num_dn};
 %apply float *OUTPUT { float* const num};
 %apply float *OUTPUT { double* const num};
 /* Return TREXIO exit code from trexio_open as part of the output tuple */
@@ -45,6 +49,10 @@
 /* Return number of sparse data points read from the file as part of the output tuple */
 %apply int *INOUT { int64_t* const buffer_size_read};
 
+/*
+%apply int { bitfield_t };
+*/
+
 /* Does not work for arrays (SIGSEGV) */
 
 /* This enables access to trexio_[...]_read_dset_str_low set of functions
@@ -93,6 +101,11 @@ import_array();
 
 %apply (double* ARGOUT_ARRAY1, int64_t DIM1) {(double* const value_sparse_read, const int64_t size_value_read)};
 %apply (int32_t* ARGOUT_ARRAY1, int64_t DIM1) {(int32_t* const index_sparse_read, const int64_t size_index_read)};
+/* Enable write|read_safe functions to convert numpy arrays from orbital list arrays */
+%apply (int32_t* ARGOUT_ARRAY1, int64_t DIM1) {(int32_t* const dset_up_out, const int64_t dim_up_out)};
+%apply (int32_t* ARGOUT_ARRAY1, int64_t DIM1) {(int32_t* const dset_dn_out, const int64_t dim_dn_out)};
+%apply (int64_t* ARGOUT_ARRAY1, int64_t DIM1) {(int64_t* const dset_up_out, const int64_t dim_up_out)};
+%apply (int64_t* ARGOUT_ARRAY1, int64_t DIM1) {(int64_t* const dset_dn_out, const int64_t dim_dn_out)};
 
 /* This tells SWIG to treat char ** dset_in pattern as a special case
    Enables access to trexio_[...]_write_dset_str set of functions directly, i.e.

src/templates_front/templator_front.org

@@ -43,6 +43,7 @@ typedef int32_t trexio_exit_code;
 #include <stdlib.h>
 #include <stdbool.h>
 #include <string.h>
+#include <err.h>
 
 #include "trexio.h"
 #include "trexio_private.h"
@@ -5154,7 +5155,7 @@ def has_determinant_coefficient(trexio_file) -> bool:
     return rc == TREXIO_SUCCESS
      #+end_src
 
-* TODO General helper functions
+* General helper functions
 
   This section contains general helper functions like ~trexio_info~.
 
@@ -5172,18 +5173,158 @@ def has_determinant_coefficient(trexio_file) -> bool:
   However, if the user validated that the file is correct (e.g. using ~trexio-tools~),
   then value of the ~metadata_unsafe~ attribute can be changed using the aforementioned function.
 
-  TODO:
+ ~trexio_to_orbital_list~ function converts the list of integer bit fields of a
- ~trexio_bitfield_to_list~ functions converts the bit field representation of a
+  given determinant into a list of indices of the occupied orbitals (for one spin component).
-  given determinants into a list of ~mo.num~ occupation numbers.
+ ~trexio_to_orbital_list_up_dn~ function does the same but for both up- and down-spin components
-  (adapt from slaterlib or QP)
+  of the determinant and returns two list of orbitals each corresponding to a different component.
 
 ** C
 
-     #+begin_src c :tangle prefix_front.h :exports none
+     #+begin_src c :tangle prefix_front.h
 trexio_exit_code trexio_info(void);
 trexio_exit_code trexio_mark_safety(trexio_t* const file, const int32_t safety_flag);
      #+end_src
 
+     #+begin_src c :tangle prefix_front.h
+typedef int64_t  bitfield_t;
+
+#define ORBITAL_SHIFT  1
+#define INT_SIZE       64
+#define NORB_PER_INT   ( 8*sizeof(bitfield_t) )
+
+/* Popcount and trailz */
+#if INT_SIZE == 64
+
+  extern int __builtin_popcountll (unsigned long long x_0);
+  #define popcnt(X) __builtin_popcountll((unsigned long long) X)
+
+  extern int __builtin_ctzll (unsigned long long x_0);
+  #define trailz(X) __builtin_ctzll((unsigned long long) X)
+
+#elif INT_SIZE == 32
+
+  extern int __builtin_popcountl (unsigned long x_0);
+  #define popcnt(X) __builtin_popcountl((unsigned long) X)
+
+  extern int __builtin_ctzl(unsigned long x_0);
+  #define trailz(X) __builtin_ctzl((unsigned long) X)
+
+#elif INT_SIZE == 16
+
+  extern int __builtin_popcount (unsigned int x_0);
+  #define popcnt(X) __builtin_popcount((unsigned int) X)
+
+  extern int __builtin_ctz (unsigned int x_0);
+  #define trailz(X) __builtin_ctz((unsigned int) X)
+
+#else
+
+ #error("Invalid INT_SIZE")
+
+#endif
+
+trexio_exit_code trexio_to_orbital_list (const int32_t N_int, const bitfield_t* d1, int32_t* const list, int32_t* const occupied_num);
+trexio_exit_code trexio_to_orbital_list_up_dn (const int32_t N_int, const bitfield_t* d1, int32_t* const list_up, int32_t* const list_dn, int32_t* const occ_num_up, int32_t* const occ_num_dn);
+trexio_exit_code trexio_safe_to_orbital_list (const int32_t N_int, const int64_t* dset_in, const int64_t dim_in, int32_t* const dset_out, const int64_t dim_out, int32_t* const num);
+trexio_exit_code trexio_safe_to_orbital_list_up_dn (const int32_t N_int, const int64_t* dset_in, const int64_t dim_in, int32_t* const dset_up_out, const int64_t dim_up_out, int32_t* const dset_dn_out, const int64_t dim_dn_out, int32_t* const num_up, int32_t* const num_dn);
+     #+end_src
+
+     #+begin_src c :tangle prefix_front.c
+trexio_exit_code trexio_to_orbital_list(const int32_t N_int,
+                                        const bitfield_t* d1,
+                                        int32_t* const list,
+                                        int32_t* const occupied_num)
+{
+  if (N_int <= 0) return TREXIO_INVALID_ARG_1;
+  if (d1 == NULL) return TREXIO_INVALID_ARG_2;
+  if (list == NULL) return TREXIO_INVALID_ARG_3;
+  if (occupied_num == NULL) return TREXIO_INVALID_ARG_4;
+
+  bitfield_t  tmp;
+  int32_t     shift;
+  int32_t     k;
+  int32_t     pos;
+
+  k = 0;
+  shift = ORBITAL_SHIFT;
+
+  for (int32_t i=0 ; i<N_int ; i++)
+  {
+      tmp = d1[i];
+
+      while (tmp != (bitfield_t) 0)
+      {
+          pos = trailz(tmp);
+          if (pos < 0) return TREXIO_FAILURE;
+
+          list[k] = ( (int32_t) pos) + shift;
+          tmp ^= ( ((bitfield_t) 1) << pos);
+          k++;
+      }
+      shift += NORB_PER_INT;
+  }
+
+  ,*occupied_num = (int32_t) k;
+  return TREXIO_SUCCESS;
+}
+    #+end_src
+
+    #+begin_src c :tangle prefix_front.c
+trexio_exit_code trexio_to_orbital_list_up_dn(const int32_t N_int,
+                                              const bitfield_t* d1,
+                                              int32_t* const list_up,
+                                              int32_t* const list_dn,
+                                              int32_t* const occ_num_up,
+                                              int32_t* const occ_num_dn)
+{
+  if (N_int <= 0) return TREXIO_INVALID_ARG_1;
+  if (d1 == NULL) return TREXIO_INVALID_ARG_2;
+  if (list_up == NULL) return TREXIO_INVALID_ARG_3;
+  if (list_dn == NULL) return TREXIO_INVALID_ARG_4;
+  if (occ_num_up == NULL) return TREXIO_INVALID_ARG_5;
+  if (occ_num_dn == NULL) return TREXIO_INVALID_ARG_6;
+
+  trexio_exit_code rc;
+
+ /* First process up-spin electrons */
+  rc = trexio_to_orbital_list(N_int, &d1[0], list_up, occ_num_up);
+  if (rc != TREXIO_SUCCESS) return rc;
+
+ /* Now process down-spin electrons */
+  rc = trexio_to_orbital_list(N_int, &d1[N_int], list_dn, occ_num_dn);
+  if (rc != TREXIO_SUCCESS) return rc;
+
+  return TREXIO_SUCCESS;
+}
+    #+end_src
+
+    #+begin_src c :tangle prefix_front.c
+trexio_exit_code
+trexio_safe_to_orbital_list (const int32_t N_int,
+			     const int64_t* dset_in,
+                             const int64_t dim_in,
+			     int32_t* const dset_out,
+                             const int64_t dim_out,
+			     int32_t* const num)
+{
+  return trexio_to_orbital_list(N_int, dset_in, dset_out, num);
+}
+
+trexio_exit_code
+trexio_safe_to_orbital_list_up_dn (const int32_t N_int,
+			           const int64_t* dset_in,
+                                   const int64_t dim_in,
+			           int32_t* const dset_up_out,
+                                   const int64_t dim_up_out,
+			           int32_t* const dset_dn_out,
+                                   const int64_t dim_dn_out,
+			           int32_t* const num_up,
+			           int32_t* const num_dn)
+{
+  return trexio_to_orbital_list_up_dn(N_int, dset_in, dset_up_out, dset_dn_out, num_up, num_dn);
+}
+    #+end_src
+
     #+begin_src c :tangle prefix_front.c
 trexio_exit_code
 trexio_info (void)
@@ -5233,18 +5374,104 @@ interface
 end interface
    #+end_src
 
+
+   #+begin_src f90 :tangle prefix_fortran.f90
+interface
+   integer(trexio_exit_code) function trexio_to_orbital_list(N_int, d1, list, occupied_num) bind(C)
+     use, intrinsic :: iso_c_binding
+     import
+     integer(c_int32_t), intent(in), value :: N_int
+     integer(c_int64_t), intent(in)        :: d1(*)
+     integer(c_int32_t), intent(out)       :: list(*)
+     integer(c_int32_t), intent(inout)     :: occupied_num
+   end function trexio_to_orbital_list
+end interface
+
+interface
+   integer(trexio_exit_code) function trexio_to_orbital_list_up_dn(N_int, d1, list_up, list_dn, occ_num_up, occ_num_dn) bind(C)
+     use, intrinsic :: iso_c_binding
+     import
+     integer(c_int32_t), intent(in), value :: N_int
+     integer(c_int64_t), intent(in)        :: d1(*)
+     integer(c_int32_t), intent(out)       :: list_up(*)
+     integer(c_int32_t), intent(out)       :: list_dn(*)
+     integer(c_int32_t), intent(inout)     :: occ_num_up
+     integer(c_int32_t), intent(inout)     :: occ_num_dn
+   end function trexio_to_orbital_list_up_dn
+end interface
+   #+end_src
+
 ** Python
 
    #+begin_src python :tangle basic_python.py
 def info():
-    """Print the info about the installed TREXIO library.
+    """Print the info about the installed TREXIO library."""
-    """
 
     rc = pytr.trexio_info()
     if rc != TREXIO_SUCCESS:
         raise Error(rc)
    #+end_src
 
+
+   #+begin_src python :tangle basic_python.py
+def to_orbital_list(n_int: int, determinant: list) -> list:
+    """Convert a given determinant into a list of occupied orbitals.
+
+    Input:
+ ~determinant~ - list of bit fields (integers)
+ ~n_int~ - number of bit fields per determinant of a given spin
+
+    Returns:
+ ~orbital_list~: list
+
+    Raises:
+        - Exception from AssertionError if TREXIO return code ~rc~ is different from TREXIO_SUCCESS and prints the error message using trexio_string_of_error.
+        - Exception from some other error (e.g. RuntimeError).
+    """
+
+    # max possible size of the orbital list per spin component (upper limit on the number of MOs)
+    size_max = n_int * 64
+
+    rc, orbital_list, occ_num = pytr.trexio_safe_to_orbital_list(n_int, determinant, size_max)
+    if rc != TREXIO_SUCCESS:
+        raise Error(rc)
+    if len(orbital_list) < occ_num:
+        raise Exception("Inconsistent size of the orbital_list.")
+
+    return orbital_list
+
+
+def to_orbital_list_up_dn(n_int: int, determinant: list) -> tuple:
+    """Convert a given determinant into two lists of occupied orbitals.
+
+    Input:
+ ~determinant~ - list of bit fields (integers)
+ ~n_int~ - number of bit fields per determinant of a given spin
+
+    Returns:
+        result: tuple with the following items:
+ ~orbital_list_up~: list of orbitals occupied by up-spin electrons
+ ~orbital_list_dn~: list of orbitals occupied by down-spin electrons
+
+    Raises:
+        - Exception from AssertionError if TREXIO return code ~rc~ is different from TREXIO_SUCCESS and prints the error message using trexio_string_of_error.
+        - Exception from some other error (e.g. RuntimeError).
+    """
+
+    # max possible size of the orbital list per spin component (upper limit on the number of MOs)
+    size_max = n_int * 64
+
+    rc, orbital_list_up, orbital_list_dn, occ_num_up, occ_num_dn = pytr.trexio_safe_to_orbital_list_up_dn(n_int, determinant, size_max, size_max)
+    if rc != TREXIO_SUCCESS:
+        raise Error(rc)
+    if len(orbital_list_up) < occ_num_up:
+        raise Exception("Inconsistent size of the orbital_list for up-spin electrons.")
+    if len(orbital_list_dn) < occ_num_dn:
+        raise Exception("Inconsistent size of the orbital_list for down-spin electrons.")
+
+    return (orbital_list_up[0:occ_num_up], orbital_list_dn[0:occ_num_dn])
+   #+end_src
+
 * Fortran helper/wrapper functions
 
   The function below adapts the original C-based ~trexio_open~ for Fortran.