Merge pull request #244 from nonhermitian/better-array-creation

Cleanup array handling for iterative solver

mthree/converters.pxd

@@ -13,11 +13,20 @@
 from libcpp.map cimport map
 from libcpp.string cimport string
 
+
 cdef void counts_to_internal(map[string, float] * counts,
                              unsigned char * vec,
                              float * probs,
                              unsigned int num_bits,
                              float shots)
 
+
 cdef void internal_to_probs(map[string, float] * counts,
                             float * probs)
+
+
+cdef void _core_counts_to_bp(map[string, float] * counts_map,
+                             unsigned int num_bits,
+                             float shots,
+                             unsigned char * bitstrings,
+                             float * probs)

mthree/converters.pyx

@@ -11,6 +11,8 @@
 # that they have been altered from the originals.
 # cython: c_string_type=unicode, c_string_encoding=UTF-8
 cimport cython
+import numpy as np
+cimport numpy as np
 from libcpp.map cimport map
 from libcpp.string cimport string
 from cython.operator cimport dereference, postincrement
@@ -66,3 +68,49 @@ cdef void internal_to_probs(map[string, float] * counts_map,
         dereference(it).second = probs[idx]
         idx += 1
         postincrement(it)
+
+
+def counts_to_bitstrings_and_probs(object counts):
+    """Convert counts to NumPy arrays of bitstrings and probabilities
+
+    Parameters:
+        counts (object): Dict or Counts object of counts data
+
+    Returns:
+        ndarray: Array of unsigned char bitstrings
+        ndarray: Array of float probabilities
+    """
+    cdef float shots = sum(counts.values())
+    cdef map[string, float] counts_map = counts
+    cdef unsigned int num_elems = counts_map.size()
+    cdef unsigned int num_bits = len(next(iter(counts)))
+
+    cdef unsigned char[::1] bitstrings = np.empty(num_bits*num_elems, dtype=np.uint8)
+    cdef float[::1] probs = np.empty(num_elems, dtype=np.float32)
+
+    _core_counts_to_bp(&counts_map, num_bits, shots,
+                       &bitstrings[0], &probs[0])
+
+    return np.asarray(bitstrings), np.asarray(probs)
+
+
+@cython.cdivision(True)
+cdef void _core_counts_to_bp(map[string, float] * counts_map,
+                             unsigned int num_bits,
+                             float shots,
+                             unsigned char * bitstrings,
+                             float * probs):
+
+    cdef unsigned int idx, letter, start
+    cdef map[string, float].iterator end = counts_map.end()
+    cdef map[string, float].iterator it = counts_map.begin()
+    cdef string temp
+    idx = 0
+    while it != end:
+        start = num_bits*idx
+        probs[idx] = dereference(it).second / shots
+        temp = dereference(it).first
+        for letter in range(num_bits):
+            bitstrings[start+letter] = <unsigned char>temp[letter]-48
+        idx += 1
+        postincrement(it)

mthree/iterative.py

@@ -77,7 +77,7 @@ def precond_matvec(x):
         (M.num_elems, M.num_elems), precond_matvec, dtype=np.float32
     )
     st = time.perf_counter()
-    vec = counts_to_vector(M.sorted_counts)
+    vec = np.asarray(M.probs, np.float32)
     fin = time.perf_counter()
     logger.info(f"Counts to vector time: {fin-st}")
 

mthree/matvec.pyx

@@ -16,12 +16,14 @@ from cython.parallel cimport prange
 import numpy as np
 cimport numpy as np
 np.import_array()
-from libc.stdlib cimport malloc, free
+
 from libcpp cimport bool
 from libcpp.map cimport map
 from libcpp.string cimport string
 from cython.operator cimport dereference, postincrement
 
+from mthree.converters cimport _core_counts_to_bp
+
 
 cdef extern from "src/distance.h" nogil:
     unsigned int hamming_terms(unsigned int num_bits,
@@ -71,13 +73,14 @@ cdef extern from "src/matvec.h" nogil:
 logger = logging.getLogger(__name__)
 
 cdef class M3MatVec():
-    cdef unsigned char * bitstrings
-    cdef float * col_norms
+    cdef public unsigned char[::1] bitstrings
+    cdef public float[::1] probs
+    cdef float[::1] col_norms
     cdef bool MAX_DIST
     cdef unsigned int distance
     cdef public unsigned int num_elems
     cdef public unsigned int num_bits
-    cdef float * cals
+    cdef float[::1] cals
     cdef public dict sorted_counts
     cdef int num_terms
 
@@ -87,7 +90,7 @@ cdef class M3MatVec():
         cdef map[string, float] counts_map = counts
         self.num_elems = counts_map.size()
         self.num_bits = len(next(iter(counts)))
-        self.cals = &cals[0]
+        self.cals = cals
         self.sorted_counts = counts_map
         self.num_terms = -1
 
@@ -101,19 +104,16 @@ cdef class M3MatVec():
 
         logger.info(f"Number of Hamming terms: {self.num_terms}")
 
-        self.bitstrings = <unsigned char *>malloc(self.num_bits*self.num_elems*sizeof(unsigned char))
-        self.col_norms = <float *>malloc(self.num_elems*sizeof(float))
-
-        counts_to_bitstrings(&counts_map, self.bitstrings, self.num_bits)
+        self.bitstrings = np.empty(self.num_bits*self.num_elems, np.uint8)
+        self.probs = np.empty(self.num_elems, np.float32)
+        self.col_norms = np.empty(self.num_elems, np.float32)
+
+        _core_counts_to_bp(&counts_map, self.num_bits, shots,
+                           &self.bitstrings[0], &self.probs[0])
 
-        compute_col_norms(self.col_norms, self.bitstrings, self.cals,
+        compute_col_norms(&self.col_norms[0], &self.bitstrings[0], &self.cals[0],
                           self.num_bits, self.num_elems, distance)
-
-    def __dealloc__(self):
-        if self.bitstrings is not NULL:
-            free(self.bitstrings)
-        if self.col_norms is not NULL:
-            free(self.col_norms)
+
 
     @cython.boundscheck(False)
     def get_col_norms(self):
@@ -136,8 +136,8 @@ cdef class M3MatVec():
         cdef float temp_elem
         cdef float[::1] out = np.empty(self.num_elems, dtype=np.float32)
         for kk in range(self.num_elems):
-            temp_elem = compute_element(kk, kk,self.bitstrings,
-                                        self.cals, self.num_bits)
+            temp_elem = compute_element(kk, kk,&self.bitstrings[0],
+                                        &self.cals[0], self.num_bits)
             temp_elem /= self.col_norms[kk]
             out[kk] = temp_elem
         return np.asarray(out, dtype=np.float32)
@@ -150,9 +150,9 @@ cdef class M3MatVec():
         cdef float[::1] out = np.empty(self.num_elems, dtype=np.float32)
         matvec(&x[0],
                &out[0],
-               self.col_norms,
-               self.bitstrings,
-               self.cals,
+               &self.col_norms[0],
+               &self.bitstrings[0],
+               &self.cals[0],
                self.num_bits,
                self.num_elems,
                self.distance,
@@ -168,9 +168,9 @@ cdef class M3MatVec():
         cdef float[::1] out = np.empty(self.num_elems, dtype=np.float32)
         rmatvec(&x[0],
                 &out[0],
-                self.col_norms,
-                self.bitstrings,
-                self.cals,
+                &self.col_norms[0],
+                &self.bitstrings[0],
+                &self.cals[0],
                 self.num_bits,
                 self.num_elems,
                 self.distance,

mthree/mitigation.py

@@ -583,8 +583,8 @@ def _apply_correction(
         num_bits = len(qubits)
         num_elems = len(counts)
         if distance is None:
-            distance = min(num_bits,3)
-        elif distance == -1: # shortcut for setting max distance
+            distance = min(num_bits, 3)
+        elif distance == -1:  # shortcut for setting max distance
             distance = num_bits
 
         # check if len of bitstrings does not equal number of qubits passed.

mthree/test/test_full.py

@@ -26,7 +26,9 @@ def test_full_problem():
     qubits = [1, 4, 7, 10, 12, 2, 3, 5]
     mit = M3Mitigation(None)
     mit.cals_from_matrices(CALS)
-    mit_counts = mit.apply_correction(COUNTS, qubits, distance=-1, tol=1e-6, method="iterative")
+    mit_counts = mit.apply_correction(
+        COUNTS, qubits, distance=-1, tol=1e-6, method="iterative"
+    )
 
     # Compute using LU solver
     sorted_counts = dict(