Add a function for evaluating q.

1 files changed, 24 insertions, 0 deletions

diff --git a/brill_data.py b/brill_data.py
index 460b0d2..192f0aa 100644
--- a/brill_data.py
+++ b/brill_data.py
@@ -113,3 +113,27 @@ class BrillData(object):
             raise RuntimeError('Error evaluating the metric')
 
         return metric
+
+    def eval_q(self, rho, z, diff_order = None):
+        if rho.ndim != 1 or z.ndim != 1:
+            raise TypeError('rho and z must be 1-dimensional NumPy arrays')
+
+        if diff_order is None:
+            diff_order = [0, 0]
+
+        c_diff_order = (ctypes.c_uint * 2)()
+        c_diff_order[0] = diff_order[0]
+        c_diff_order[1] = diff_order[1]
+
+        q = np.empty((rho.shape[0], z.shape[0]))
+
+        c_q   = np.ctypeslib.as_ctypes(q)
+        c_rho = np.ctypeslib.as_ctypes(rho)
+        c_z   = np.ctypeslib.as_ctypes(z)
+
+        ret = self._libbd.bd_eval_q(self._bdctx, c_rho, len(c_rho),
+                                    c_z, len(c_z), c_diff_order, c_q)
+        if ret < 0:
+            raise RuntimeError('Error evaluating q')
+
+        return q