Add interpolation scheme for n-dimensional binning

src/c_lib/base/base_model.pxd

@@ -157,3 +157,15 @@ cdef extern from "simulation.h":
         bool_t *freq_contrib,
         double *affine_matrix,
         )
+
+cdef extern from "interpolation.h":
+    void multidimensional_linear_interpolation(
+        double *points,
+        int *nearest_points,
+        double *offset_vectors,
+        double *amp,
+        double *temp_amp,
+        int *dim_lengths,
+        int n_dims,
+        int n_points
+    )

src/c_lib/base/base_model.pyx

@@ -2,6 +2,7 @@ cimport base_model as clib
 from libcpp cimport bool as bool_t
 from numpy cimport ndarray
 import numpy as np
+cimport numpy as cnp
 import cython
 
 __author__ = "Deepansh J. Srivastava"
@@ -592,6 +593,59 @@ def calculate_transition_connect_weight(
     return complex(factor[0], factor[1])
 
 
+@cython.profile(False)
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def bin_with_linear_interp(
+    ndarray[cnp.float64_t, ndim=2] _points,
+    _grid_dims,  # list since 2d ndarrays can't have differnt shapes
+    # amp_out,
+):
+    """TODO: Docstring
+
+    Arguments:
+        (np.ndarray) points: random points
+        (list) grid_dims: Numpy array of each dimension on the grid
+        # (np.ndarray) amp_out: Numpy array to add binned pdf to
+
+    Return: A distribution of points binned with linear interpolation.
+    """
+    for i in range(len(_grid_dims)):
+        dim = _grid_dims[i]
+        start = dim[0]
+        step = dim[1] - start
+        _points[i] = (_points[i] - start) / step
+
+    _points = np.vstack(_points).T
+    nearest = np.floor(_points + 0.5).astype(np.int32)
+    distance = _points - nearest
+    cdef ndarray[cnp.float64_t, ndim=1] points = _points.flatten()
+    cdef ndarray[cnp.float64_t, ndim=1] nearest_distance = distance.flatten()
+    cdef ndarray[int, ndim=1] nearest_points = nearest.astype(np.int32).flatten()
+
+    cdef ndarray[int, ndim=1] dim_sizes = np.asarray(
+        [dim.size for dim in _grid_dims], dtype=np.int32
+    )
+    cdef ndarray[cnp.float64_t, ndim=1] amp = np.zeros(np.prod(dim_sizes), dtype=np.float64).flatten()
+    cdef ndarray[cnp.float64_t, ndim=1] temp_amp = np.zeros(np.prod(dim_sizes), dtype=np.float64)
+
+    cdef int n_dims = len(dim_sizes)
+    cdef int n_points = points.size / n_dims
+
+    clib.multidimensional_linear_interpolation(
+        &points[0],
+        &nearest_points[0],
+        &nearest_distance[0],
+        &amp[0],
+        &temp_amp[0],
+        &dim_sizes[0],
+        n_dims,
+        n_points,
+    )
+
+    return amp.reshape([dim.size for dim in _grid_dims])
+
+
 # @cython.profile(False)
 # @cython.boundscheck(False)
 # @cython.wraparound(False)

src/c_lib/include/interpolation.h

@@ -9,6 +9,22 @@
 
 #include "config.h"
 
+/**
+ * @brief Bin random points onto a n-dimensional grid while performing linear
+ * interpolation.
+ *
+ * @param points Pointer to flattened array of points
+ * @param nearest_points Pointer to array of nearest grid coord for each point
+ * @param offsets Pointer to array of distances from nearest grid coord
+ * @param dim_sizes Pointer to array holding number of point in each grid dim
+ * @param n_dims Integer specifying total number of dimensions
+ * @param n_points Integer specifying total number of points
+ */
+void multidimensional_linear_interpolation(double *points, int *nearest_points,
+                                           double *offsets, double *amp,
+                                           double *temp_amp, int *dim_sizes, int n_dims,
+                                           int n_points);
+
 /**
  * @brief Create a triangle with coordinates (f1, f2, f2) onto a 1D grid.
  *

src/c_lib/lib/interpolation.c

@@ -99,6 +99,149 @@ static void inline delta_fn_gauss_interpolation(const double *freq, const int *p
   if (p + 2 >= 0 && p + 2 < *points) spec[p2 + 4] += temp * a4;
 }
 
+static inline void int_to_bin_digit(int num, int count, int *out) {
+  unsigned int mask = 1U << (count - 1);
+  for (int i = count - 1; i >= 0; i--) {
+    out[i] = (num & mask) ? 1 : 0;
+    num <<= 1;
+  }
+}
+
+static inline double arr_min_double(double *arr, int count) {
+  int min_val = arr[0];
+  for (int i = 0; i < count; i++) {
+    if (arr[i] < min_val) min_val = arr[i];
+  }
+  return min_val;
+}
+
+static inline void arr_add_double(double *u, double *v, int count, double *out) {
+  for (int i = 0; i < count; i++) {
+    out[i] = u[i] + v[i];
+  }
+}
+
+static inline void arr_add_int(int *u, int *v, int count, int *out) {
+  for (int i = 0; i < count; i++) {
+    out[i] = u[i] + v[i];
+  }
+}
+
+static inline void arr_mult_int(int *u, int *v, int count, int *out) {
+  for (int i = 0; i < count; i++) {
+    out[i] = u[i] * v[i];
+  }
+}
+
+static inline int dot_prod_int(int *u, int *v, int count) {
+  int i, sum = 0;
+  for (i = 0; i < count; i++) {
+    sum += u[i] * v[i];
+  }
+  return sum;
+}
+
+static inline void fill_int(int *arr, int val, int count) {
+  for (int i = 0; i < count; i++) {
+    arr[i] = val;
+  }
+}
+
+static inline void fill_double(double *arr, double val, int count) {
+  for (int i = 0; i < count; i++) {
+    arr[i] = val;
+  }
+}
+
+/**
+ * @brief Bin random points onto a n-dimensional grid while performing linear
+ * interpolation.
+ *
+ * @param points Pointer to flattened array of points
+ * @param nearest_points Pointer to array of nearest grid coord for each point
+ * @param offsets Pointer to array of distances from nearest grid coord
+ * @param dim_sizes Pointer to array holding number of point in each grid dim
+ * @param n_dims Integer specifying total number of dimensions
+ * @param n_points Integer specifying total number of points
+ */
+void multidimensional_linear_interpolation(double *points, int *nearest_points,
+                                           double *offsets, double *amp,
+                                           double *temp_amp, int *dim_sizes, int n_dims,
+                                           int n_points) {
+  int i, j, k, dir, n_dirs, lin_idx, dir_idx, zero_offset_dims,
+      total_bins = 1, dim_strides[n_dims], prev_directions[n_dims],
+      temp_int_arr1[n_dims], temp_int_arr2[n_dims];
+  double delta, maximum_vector[n_dims], temp_double_arr[n_dims];
+
+  // maximum allowed vector, but negative
+  for (i = 0; i < n_dims; i++) {
+    maximum_vector[i] = (double)(1 - dim_sizes[i]);
+    total_bins *= dim_sizes[i];
+  }
+
+  // Roughly equivelent to numpy strides. Used to convert n-D to 1-D index
+  dim_strides[n_dims - 1] = 1;
+  for (i = n_dims - 2; i >= 0; i--) {
+    dim_strides[i] = dim_strides[i + 1] * dim_sizes[i + 1];
+  }
+
+  // Loop iterating over all points
+  for (i = 0; i < n_points * n_dims; i += n_dims) {
+    arr_add_double(points + i, maximum_vector, n_dims, temp_double_arr);
+    if (arr_min_double(points + i, n_dims) < -0.5 ||   // half bin width to "left"
+        arr_min_double(temp_double_arr, n_dims) > 0.5  // half bin width "right"
+        // if (arr_min_double(points + i, n_dims) < 0 ||
+        //     arr_min_double(temp_double_arr, n_dims) > 0
+    )
+      continue;
+
+    zero_offset_dims = 0;
+    fill_int(prev_directions, 0, n_dims);    // could use memset, but this
+    fill_double(temp_amp, 0.0, total_bins);  // is more comprehensible
+    temp_amp[dot_prod_int(dim_strides, nearest_points + i, n_dims)] = 1;
+
+    // iterate over all directions for a point
+    for (j = 0; j < n_dims; j++) {
+      n_dirs = (int)pow((double)2, (double)j);  // 2^j different combos
+      delta = offsets[i + j];
+
+      // Check if delta (offset) is zero, left or right
+      if (delta == 0.0) {
+        zero_offset_dims += 1 << j;
+        continue;
+      }
+      dir = 1;
+      if (delta < 0.0) {
+        dir = -1;
+        delta = -delta;
+      }
+
+      // Generate all combinations of directions (neighboring bins)
+      for (k = 0; k < n_dirs; k++) {
+        if (zero_offset_dims & k) continue;  // skip dirs with zero offsets
+
+        // put one combo of directions into temp_int_arr2
+        int_to_bin_digit(k, n_dims, temp_int_arr1);
+        arr_mult_int(temp_int_arr1, prev_directions, n_dims, temp_int_arr2);
+
+        // put cartesian index of bin with offset into temp_int_arr1
+        arr_add_int(temp_int_arr2, nearest_points + i, n_dims, temp_int_arr1);
+
+        // continue if new bin is outside of grid
+        if (temp_int_arr1[j] + dir < 0) continue;
+        if (temp_int_arr1[j] + dir >= dim_sizes[j]) continue;
+
+        lin_idx = dot_prod_int(dim_strides, temp_int_arr1, n_dims);
+        dir_idx = lin_idx + (dim_strides[j] * dir);
+        temp_amp[dir_idx] = temp_amp[lin_idx] * delta;
+        temp_amp[lin_idx] *= 1 - delta;
+      }
+      prev_directions[j] = dir;
+    }
+    arr_add_double(amp, temp_amp, total_bins, amp);
+  }
+}
+
 /**
  * @brief Get the clipping conditions object
  *

src/mrsimulator/models/czjzek.py

@@ -1,9 +1,11 @@
 import numpy as np
 from mrsimulator.spin_system.tensors import SymmetricTensor
+from mrsimulator.base_model import bin_with_linear_interp
 
 from .utils import get_Haeberlen_components
 from .utils import get_principal_components
 from .utils import x_y_from_zeta_eta
+from .utils import x_y_to_zeta_eta
 
 __author__ = "Deepansh J. Srivastava"
 __email__ = "[email protected]"
@@ -68,7 +70,7 @@ def _czjzek_random_distribution_tensors(sigma, n):
 
 
 class AbstractDistribution:
-    def pdf(self, pos, size: int = 400000):
+    def pdf(self, pos, size: int = 400000, interpolation: str = "none"):
         """Generates a probability distribution function by binning the random
         variates of length size onto the given grid system.
 
@@ -86,15 +88,29 @@ def pdf(self, pos, size: int = 400000):
             >>> eta = np.arange(21)/20
             >>> Cq_dist, eta_dist, amp = cz_model.pdf(pos=[cq, eta])
         """
-        delta_z = (pos[0][1] - pos[0][0]) / 2
-        delta_e = (pos[1][1] - pos[1][0]) / 2
-        x = [pos[0][0] - delta_z, pos[0][-1] + delta_z]
-        y = [pos[1][0] - delta_e, pos[1][-1] + delta_e]
-
-        x_size = pos[0].size
-        y_size = pos[1].size
+        if interpolation not in ["none", "linear"]:
+            raise ValueError(
+                f"Unrecognized interpolation type of {interpolation}. "
+                "Allowed values are \"none\" and \"linear\"."
+            )
+
         zeta, eta = self.rvs(size)
-        hist, _, _ = np.histogram2d(zeta, eta, bins=[x_size, y_size], range=[x, y])
+
+        if interpolation == "linear":
+            x_dim = pos[0].astype(np.float64, copy=True)
+            y_dim = pos[1].astype(np.float64, copy=True)
+            dims = [x_dim, y_dim]
+            hist = bin_with_linear_interp(np.array([zeta, eta]), dims)
+
+        if interpolation == "none":
+            delta_z = (pos[0][1] - pos[0][0]) / 2
+            delta_e = (pos[1][1] - pos[1][0]) / 2
+            x = [pos[0][0] - delta_z, pos[0][-1] + delta_z]
+            y = [pos[1][0] - delta_e, pos[1][-1] + delta_e]
+
+            x_size = pos[0].size
+            y_size = pos[1].size
+            hist, _, _ = np.histogram2d(zeta, eta, bins=[x_size, y_size], range=[x, y])
 
         hist /= hist.sum()
 
@@ -141,6 +157,65 @@ def __init__(self, sigma: float, polar=False):
         self.sigma = sigma
         self.polar = polar
 
+    def _analytical_czjzek_pdf(self, zeta, eta, sigma):
+        """Analytical probability distribution function of (zeta, eta) for Czjzek
+        distribution.
+
+        TODO: complete docstring
+        """
+        denom = (2 * np.pi) ** 0.5 * sigma**5
+        res = (zeta**4 * eta) * (1 - eta**2 / 9) / denom
+        res *= np.exp(-(zeta**2 * (1 + (eta**2 / 3))) / (2 * sigma**2))
+        res /= res.sum()
+
+        return res
+
+
+    def cartesian_pdf(self, pos):
+        """TODO: Docstring"""
+        sigma_ = 2 * self.sigma
+        z_range = pos[0]
+        e_range = pos[1]
+        VV, ee = np.meshgrid(z_range, e_range)
+        amp = self._analytical_czjzek_pdf(zeta=VV, eta=ee, sigma=sigma_)
+
+        return VV, ee, amp
+
+    def polar_pdf(self, pos):
+        """TODO: Docstring"""
+        sigma_ = 2 * self.sigma
+        x_range = pos[0]
+        y_range = pos[1]
+        xx, yy = np.meshgrid(x_range, y_range)
+        z_points, e_points = x_y_to_zeta_eta(xx, yy)
+
+        # Call pdf method accepting (zeta, eta) grid as argument
+        amp = self._analytical_czjzek_pdf(
+            zeta=z_points, 
+            eta=e_points,
+            sigma=sigma_
+        )
+        amp = amp.reshape(xx.shape)
+
+        return xx, yy, amp
+
+
+
+    # def pdf(self, pos):
+    #     """Analytical solution to distribution of Cq and eta"""
+    #     sigma_ = 2 * self.sigma
+    #     z_range = pos[0]
+    #     e_range = pos[1]
+    #     x_, y_ = np.meshgrid(z_range, e_range)
+
+    #     V, e = np.meshgrid(z_range, e_range)
+    #     denom = (2 * np.pi) ** 0.5 * sigma_**5
+    #     res = (V**4 * e) * (1 - e**2 / 9) / denom
+    #     res *= np.exp(-(V**2 * (1 + (e**2 / 3))) / (2 * sigma_**2))
+    #     res /= res.sum()
+
+    #     return x_, y_, res
+
     def rvs(self, size: int):
         """Draw random variates of length `size` from the distribution.