Add optimized C-code for column renomalization

.gitignore

@@ -134,3 +134,4 @@ docs/*.json
 docs/stubs/*
 docs/tutorials/
 bad_cals.json
+cals.json

mthree/matvec.pyx

@@ -21,7 +21,17 @@ from libcpp.map cimport map
 from libcpp.string cimport string
 from cython.operator cimport dereference, postincrement
 
-from mthree.compute cimport within_distance, compute_element, compute_col_norms
+from mthree.compute cimport within_distance, compute_element
+
+
+cdef extern from "src/col_renorm.h" nogil:
+    void compute_col_norms(float * col_norms,
+                           const unsigned char * bitstrings,
+                           const float * cals,
+                           unsigned int num_bits,
+                           unsigned int num_elems,
+                           unsigned int distance)
+
 
 
 cdef class M3MatVec():

mthree/src/col_renorm.h

@@ -0,0 +1,61 @@
+/* Copyright (c) 2024 mthree */
+#include <stddef.h>
+#include <stdbool.h>
+#include "distance.h"
+#include "elements.h"
+
+
+void compute_col_norms(float * col_norms,
+                       const unsigned char * __restrict bitstrings,
+                       const float * __restrict cals,
+                       unsigned int num_bits,
+                       unsigned int num_elems,
+                       unsigned int distance)
+    /**
+   * @brief Computes the renormalization factor for each column of A-matrix
+   *
+   * @param col_norms Pointer to where to store col norm data
+   * @param bitstrings Pointer to array of bitstrings
+   * @param cals Pointer to array containing calibration data
+   * @param num_bits Number of bits in a single bit-string
+   * @param num_elems Number of elements (dimension) of reduced A-matrix
+   * @param distance Max Hamming distance
+   */
+    {
+
+      size_t col;
+      bool MAX_DIST = false;
+      int num_terms = -1;
+
+      if (distance == num_bits)
+      {
+        MAX_DIST = true;
+      }
+      else{
+        num_terms = hamming_terms(num_bits, distance); 
+      }
+
+      #pragma omp parallel for
+      for (col = 0; col < num_elems; ++col)
+      {
+        bool flag = false;
+        float col_norm = 0.0F;
+        size_t row;
+        int terms = 0;
+        for (row = 0; row < num_elems; ++row)
+        {
+          if (flag) continue;
+          if (MAX_DIST || within_distance(row, col, bitstrings, num_bits, distance))
+          {
+            col_norm += compute_element(row, col, bitstrings, cals, num_bits);
+            terms += 1;
+            if (terms == num_terms)
+            {
+              flag = true;
+            }
+          }
+        }
+        col_norms[col] = col_norm;
+      }
+
+    }

mthree/src/distance.h

@@ -0,0 +1,70 @@
+/* Copyright (c) 2024 mthree */
+#include <stddef.h>
+#include <stdbool.h>
+
+
+static inline bool within_distance(unsigned int row,
+                                   unsigned int col,
+                                   const unsigned char * __restrict bitstrings,
+                                   unsigned int num_bits,
+                                   unsigned int distance)
+  /**
+   * @brief Computes if two bit-strings are within the specified Hamming distance.
+   *
+   * @param row Index for row bit-string
+   * @param col Index for col bit-string
+   * @param bitstrings Pointer to array of bit-strings
+   * @param num_bits Number of bits in a single bit-string
+   * @param distance Max Hamming distance
+   *
+   * @return Are the bit-strings within the Hamming distance 
+   */
+  {
+    size_t kk;
+    unsigned int sum=0;
+    const unsigned char * row_pos = &bitstrings[num_bits*row];
+    const unsigned char * col_pos = &bitstrings[num_bits*col];
+    #pragma omp simd reduction(+:sum)
+    for (kk = 0; kk < num_bits; ++kk)
+        {
+          sum += row_pos[kk] ^ col_pos[kk];
+        }
+    return sum <= distance;
+
+  }
+
+
+static inline unsigned int binomial_coeff(unsigned int n, unsigned int k)
+  {
+    if (k > n)
+      {
+        return 0U;
+      }
+    else if ((k == 0) || (k == n))
+      {
+        return 1U;
+      }
+    else if ((k == 1) || (k == (n-1)))
+      {
+        return n;
+      }
+    else if (k+k < n)
+      {
+        return (binomial_coeff(n-1, k-1) * n) / k;
+      }
+    else
+      {
+        return (binomial_coeff(n-1, k) * n) / (n-k);
+      }
+  }
+
+
+unsigned int hamming_terms(unsigned int num_bits, unsigned int distance)
+{
+  unsigned int kk, out = 0;
+  for (kk=0; kk < (distance+1); ++kk)
+    {
+      out += binomial_coeff(num_bits, kk);
+    }
+  return out;
+}

mthree/src/elements.h

@@ -0,0 +1,22 @@
+/* Copyright (c) 2024 mthree */
+#include <stddef.h>
+
+static inline float compute_element(unsigned int row,
+                                   unsigned int col,
+                                   const unsigned char * __restrict bitstrings,
+                                   const float * __restrict cals,
+                                   unsigned int num_bits)
+  {
+    float res = 1.0F;
+    size_t kk;
+    unsigned int offset;
+    const unsigned char * row_pos = &bitstrings[num_bits*row];
+    const unsigned char * col_pos = &bitstrings[num_bits*col];
+    #pragma omp simd reduction(*:res)
+    for (kk=0; kk < num_bits; ++kk)
+    {
+      offset = 2*row_pos[kk] + col_pos[kk];
+      res *= cals[4*kk+offset];
+    }
+    return res;
+  }