Allow access to Tk3D spline data (#924)

* first commit

* one more test

* gsl spline checks at every scale factor

* fallback for Tk3D.has_tsp

pyccl/ccl_utils.i

@@ -14,6 +14,8 @@
 %apply (int DIM1, double* ARGOUT_ARRAY1) {(int x_size, double* xarr)};
 %apply (int DIM1, double* ARGOUT_ARRAY1) {(int y_size, double* yarr)};
 %apply (int DIM1, double* ARGOUT_ARRAY1) {(int z_size, double* zarr)};
+%apply (int DIM1, double* ARGOUT_ARRAY1) {(int t_size, double* tarr)};
+%apply (int DIM1, double* ARGOUT_ARRAY1) {(int a_size, double* out_arr)};
 
 %apply (int* OUTPUT) {(int *size)};
 %apply int *OUTPUT { int *x_size, int *y_size };
@@ -78,6 +80,18 @@ void get_spline2d_array_sizes(gsl_spline2d *spline2d, int* x_size, int* y_size,
   *y_size = spline2d->interp_object.ysize;
 }
 
+// Not really 3d, because 3d is an array of 2d interps, but works with double pointers. //
+void get_spline3d_array_sizes(gsl_spline2d **spline2d, int* x_size, int* y_size,
+                              int* status) {
+  if((*spline2d) == NULL) {
+    *status = CCL_ERROR_MEMORY;
+    return;
+  }
+  // check only the first item of the array //
+  *x_size = (*spline2d)->interp_object.xsize;
+  *y_size = (*spline2d)->interp_object.ysize;
+}
+
 void get_spline1d_arrays(gsl_spline *spline,
                          int x_size, double* xarr,
                          int y_size, double* yarr,
@@ -122,4 +136,48 @@ void get_spline2d_arrays(gsl_spline2d *spline2d,
   memcpy(zarr, spline2d->zarr, sizeof(double)*x_size*y_size);
 }
 
+void get_spline3d_arrays(gsl_spline2d **spline2d,
+                         int x_size, double* xarr,
+                         int y_size, double* yarr,
+                         int t_size, double* tarr,
+                         int na,
+                         int *status)
+{
+  // check for inconsistencies //
+  if (spline2d == NULL) {
+    *status = CCL_ERROR_MEMORY;
+    return;
+  }
+
+  for (int ia = 0; ia < na; ia++) {
+    if (x_size != spline2d[ia]->interp_object.xsize
+        || y_size != spline2d[ia]->interp_object.ysize) {
+      *status = CCL_ERROR_INCONSISTENT;
+      return;
+    }
+  }
+
+  for (int ia = 0; ia < na; ia++) {
+    for (int ik = 0; ik < x_size*y_size; ik++) {
+      tarr[ia*x_size*y_size + ik] = spline2d[ia]->zarr[ik];
+    }
+  }
+
+  // no need to do this for every scale factor //
+  memcpy(xarr, spline2d[0]->xarr, sizeof(double)*x_size);
+  memcpy(yarr, spline2d[0]->yarr, sizeof(double)*y_size);
+}
+
+
+void get_array(double *arr,
+               int a_size, double* out_arr,
+               int *status)
+{
+  if(arr == NULL) {
+    *status = CCL_ERROR_MEMORY;
+    return;
+  }
+  memcpy(out_arr, arr, sizeof(double)*a_size);
+}
+
 %}

pyccl/pyutils.py

@@ -548,7 +548,7 @@ def _get_spline2d_arrays(gsl_spline):
     """Get array data from a 2D GSL spline.
 
     Args:
-        gsl_spline: `SWIGObject` of gsl_spline2d
+        gsl_spline: `SWIGObject` of gsl_spline2d *
             The SWIG object of the 2D GSL spline.
 
     Returns:
@@ -570,3 +570,33 @@ def _get_spline2d_arrays(gsl_spline):
     check(status)
 
     return yarr, xarr, zarr.reshape(y_size, x_size)
+
+
+def _get_spline3d_arrays(gsl_spline, length):
+    """Get array data from an array of 2D GSL splines.
+
+    Args:
+        gsl_spline (`SWIGObject` of gsl_spline2d **):
+            The SWIG object of the 2D GSL spline.
+        length (int):
+            The length of the 3rd dimension.
+
+    Returns:
+        xarr: array_like
+            The x array of the spline.
+        yarr: array_like
+            The y array of the spline.
+        zarr: array_like
+            The z array of the spline. The shape is (yarr.size, xarr.size).
+    """
+    status = 0
+    x_size, y_size, status = lib.get_spline3d_array_sizes(gsl_spline, status)
+    check(status)
+
+    z_size = x_size*y_size*length
+    xarr, yarr, zarr, status = lib.get_spline3d_arrays(gsl_spline,
+                                                       x_size, y_size, z_size,
+                                                       length, status)
+    check(status)
+
+    return xarr, yarr, zarr.reshape((length, x_size, y_size))

pyccl/tests/test_tk3d.py

@@ -137,3 +137,40 @@ def test_tk3d_eval(is_product):
     ptrue = tkkaf(ktest[None, :], ktest[:, None], atest)
     phere = tsp.eval(ktest, atest)
     assert_allclose(phere.flatten(), ptrue.flatten(), rtol=1E-6)
+
+
+@pytest.mark.parametrize('is_product', [True, False])
+def test_tk3d_spline_arrays(is_product):
+    (a_arr, lk_arr, fka1_arr, fka2_arr, tkka_arr) = get_arrays()
+    if is_product:
+        tsp = ccl.Tk3D(a_arr, lk_arr, pk1_arr=fka1_arr, pk2_arr=fka2_arr)
+    else:
+        tsp = ccl.Tk3D(a_arr, lk_arr, tkk_arr=tkka_arr)
+
+    a_get, lk_get1, lk_get2, out = tsp.get_spline_arrays()
+    assert np.allclose(a_get, a_arr, rtol=1e-15)
+    assert np.allclose(lk_get1, lk_arr, rtol=1e-15)
+    assert np.allclose(lk_get2, lk_arr, rtol=1e-15)
+
+    if is_product:
+        assert np.allclose(np.log(out[0]), fka1_arr, rtol=1e-15)
+        assert np.allclose(np.log(out[1]), fka2_arr, rtol=1e-15)
+    else:
+        assert np.allclose(np.log(out[0]), tkka_arr, rtol=1e-15)
+
+
+def test_tk3d_spline_arrays_raises():
+    (a_arr, lk_arr, fka1_arr, fka2_arr, tkka_arr) = get_arrays()
+    tsp = ccl.Tk3D(a_arr, lk_arr, tkk_arr=tkka_arr)
+
+    # PR923 aims to change this bit of code; the assertion is there to remind
+    # us to uncomment what is commented out.
+    assert not hasattr(tsp.__class__, "has_tsp")
+    # ccl.lib.f3d_t_free(tsp.tsp)
+    # delattr(tsp, "tsp")
+
+    # fool `Tk3D` into believing it doesn't have a `tsp`
+    tsp.has_tsp = False
+
+    with pytest.raises(ValueError):
+        tsp.get_spline_arrays()

pyccl/tk3d.py

@@ -1,6 +1,6 @@
 from . import ccllib as lib
 
-from .pyutils import check
+from .pyutils import check, _get_spline2d_arrays, _get_spline3d_arrays
 import numpy as np
 
 
@@ -167,3 +167,37 @@ def __del__(self):
         if hasattr(self, 'has_tsp'):
             if self.has_tsp and hasattr(self, 'tsp'):
                 lib.f3d_t_free(self.tsp)
+
+    def get_spline_arrays(self):
+        """Get the spline data arrays.
+
+        Returns:
+            a_arr (1D ``numpy.ndarray``):
+                Array of scale factors.
+            lk_arr1, lk_arr2 (1D ``numpy.ndarray``):
+                Arrays of :math:``ln(k)``.
+            out (list of ``numpy.ndarray``):
+                The trispectrum T(k1, k2, z) or its factors f(k1, z), f(k2, z).
+        """
+        if not self.has_tsp:
+            raise ValueError("Tk3D object does not have data.")
+
+        out = []
+        if self.tsp.is_product:
+            a_arr, lk_arr1, pk_arr1 = _get_spline2d_arrays(self.tsp.fka_1.fka)
+            _, lk_arr2, pk_arr2 = _get_spline2d_arrays(self.tsp.fka_2.fka)
+            out.append(pk_arr1)
+            out.append(pk_arr2)
+        else:
+            status = 0
+            a_arr, status = lib.get_array(self.tsp.a_arr, self.tsp.na, status)
+            check(status)
+            lk_arr1, lk_arr2, tkka_arr = _get_spline3d_arrays(self.tsp.tkka,
+                                                              self.tsp.na)
+            out.append(tkka_arr)
+
+        if self.tsp.is_log:
+            # exponentiate in-place
+            [np.exp(tk, out=tk) for tk in out]
+
+        return a_arr, lk_arr1, lk_arr2, out