Add a function for evaluating q.

3 files changed, 83 insertions, 0 deletions

diff --git a/brill_data.c b/brill_data.c
index 1f6378f..96b2d3a 100644
--- a/brill_data.c
+++ b/brill_data.c
@@ -803,3 +803,39 @@ fail:
 
     return ret;
 }
+
+int bd_eval_q(BDContext *bd, const double *rho, int nb_coords_rho,
+              const double *z, int nb_coords_z, const unsigned int diff_order[2],
+              double *out)
+{
+    BDPriv *s = bd->priv;
+    double (*eval)(struct BDContext *bd, double rho, double z);
+
+    if (diff_order[0] > 2 || diff_order[1] > 2) {
+        fprintf(stderr, "Derivatives of higher order than 2 are not supported.\n");
+        return -ENOSYS;
+    }
+
+    switch (diff_order[0]) {
+    case 0:
+        switch (diff_order[1]) {
+        case 0: eval = s->q_func->q;     break;
+        case 1: eval = s->q_func->dq_z;  break;
+        case 2: eval = s->q_func->d2q_z; break;
+        }
+        break;
+    case 1:
+        switch (diff_order[1]) {
+        case 0: eval = s->q_func->dq_rho;    break;
+        case 1: eval = s->q_func->d2q_rho_z; break;
+        }
+        break;
+    case 2:     eval = s->q_func->d2q_rho;   break;
+    }
+
+    for (int i = 0; i < nb_coords_z; i++)
+        for (int j = 0; j < nb_coords_rho; j++)
+            out[i * nb_coords_rho + j] = eval(bd, rho[j], z[i]);
+
+    return 0;
+}
diff --git a/brill_data.h b/brill_data.h
index cd26270..37b338a 100644
--- a/brill_data.h
+++ b/brill_data.h
@@ -190,4 +190,27 @@ int bd_eval_metric(BDContext *bd, const double *rho, int nb_coords_rho,
                    const unsigned int comp[2], const unsigned int diff_order[2],
                    double *out);
 
+/**
+ * Evaluate the q function at the specified rectangular grid (in cylindrical
+ * coordinates { ρ, z, φ }).
+ *
+ * @param bd the solver context
+ * @param rho the array of ρ coordinates.
+ * @param nb_coords_rho the number of elements in rho.
+ * @param z   the array of z coordinates.
+ * @param nb_coords_z the number of elements in z.
+ * @param diff_order the order of the derivatives of the q function to evaluate. The first element
+ *                   specifies the derivative wrt ρ, the second wrt z. I.e. diff_order = { 0, 0 }
+ *                   evaluates the q function itself, diff_order = { 0, 1 } evaluates ∂q/∂z etc.
+ * @param out the array into which the values of the q function will be written. The q function
+ *            is evaluated on the grid formed by the outer product of the rho
+ *            and z vectors. I.e. out[j * nb_coords_rho + i] = γ_comp[0]comp[1](rho[i], z[j]).
+ *            The length of out must be nb_coords_rho * nb_coords_z.
+ *
+ * @return >= 0 on success, a negative error code on failure.
+ */
+int bd_eval_q(BDContext *bd, const double *rho, int nb_coords_rho,
+              const double *z, int nb_coords_z, const unsigned int diff_order[2],
+              double *out);
+
 #endif /* BRILL_DATA_H */
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