Add field interpolation along LOS (#112)

* Add attemtp

* Add import

* Update LOS code

* Add better docstrings

* Update nb

csiborgtools/field/__init__.py

@@ -17,8 +17,9 @@ from .density import (DensityField, PotentialField, TidalTensorField,
                       overdensity_field)                                        # noqa
 from .enclosed_mass import (particles_enclosed_mass,                            # noqa
                             particles_enclosed_momentum, field_enclosed_mass)   # noqa
-from .interp import (evaluate_cartesian, evaluate_sky, field2rsp,               # noqa
-                     fill_outside, make_sky, observer_peculiar_velocity,        # noqa
-                     smoothen_field, field_at_distance)                         # noqa
+from .interp import (evaluate_cartesian, evaluate_sky, evaluate_los,            # noqa
+                     field2rsp, fill_outside, make_sky,                         # noqa
+                     observer_peculiar_velocity, smoothen_field,                # noqa
+                     field_at_distance)                                         # noqa
 from .corr import bayesian_bootstrap_correlation                                # noqa
 from .utils import nside2radec                                                  # noqa

csiborgtools/field/interp.py

@@ -47,7 +47,7 @@ def evaluate_cartesian(*fields, pos, smooth_scales=None, verbose=False):
 
     Returns
     -------
-    (list of) 1-dimensional array of shape `(n_samples, len(smooth_scales))`
+    (list of) 2-dimensional array of shape `(n_samples, len(smooth_scales))`
     """
     pos = force_single_precision(pos)
 
@@ -121,6 +121,102 @@ def observer_peculiar_velocity(velocity_field, smooth_scales=None,
 
     return numpy.vstack([vx, vy, vz]).T
 
+###############################################################################
+#                   Evaluating the fields along a LOS                         #
+###############################################################################
+
+
+def evaluate_los(*fields, sky_pos, boxsize, rmax, dr, smooth_scales=None,
+                 verbose=False):
+    """
+    Interpolate the fields for a set of lines of sights from the observer
+    in the centre of the box.
+
+    Parameters
+    ----------
+    *fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
+        Fields to be interpolated.
+    sky_pos : 2-dimensional array of shape `(n_samples, 2)`
+        Query positions in spherical coordinates (RA, dec) in degrees.
+    boxsize : float
+        Box size in `Mpc / h`.
+    rmax : float
+        Maximum radial distance in `Mpc / h`.
+    dr : float
+        Radial distance step in `Mpc / h`.
+    smooth_scales : (list of) float, optional
+        Smoothing scales in `Mpc / h`.
+    verbose : bool, optional
+        Smoothing verbosity flag.
+
+    Returns
+    -------
+    rdist : 1-dimensional array
+        Radial positions in `Mpc / h` where the fields were evaluated.
+    field_interp : (list of) 2- or 3-dimensional arrays of shape `(n_query, len(rdist), len(smooth_scales))`  # noqa
+        The interpolated fields. If `smooth_scales` is `None`, the last
+        is omitted.
+    """
+    mpc2box = 1. / boxsize
+
+    if not isinstance(sky_pos, numpy.ndarray) and sky_pos.ndim != 2:
+        raise ValueError("`sky_pos` must be a 2D array.")
+    nsamples = len(sky_pos)
+
+    if rmax > 3**0.5 / 2 * boxsize:
+        raise ValueError("The maximum radius must be within the box.")
+
+    # Radial positions to evaluate for each line of sight.
+    rdist = numpy.arange(0, rmax, dr, dtype=fields[0].dtype)
+
+    # Create an array of radial positions and sky coordinates of each line of
+    # sight.
+    pos = numpy.empty((nsamples * len(rdist), 3), dtype=fields[0].dtype)
+    for i in range(nsamples):
+        start, end = i * len(rdist), (i + 1) * len(rdist)
+        pos[start:end, 0] = rdist * mpc2box
+        pos[start:end, 1] = sky_pos[i, 0]
+        pos[start:end, 2] = sky_pos[i, 1]
+
+    pos = force_single_precision(pos)
+    # Convert the spherical coordinates to Cartesian coordinates.
+    pos = radec_to_cartesian(pos) + 0.5
+
+    if smooth_scales is not None:
+        if isinstance(smooth_scales, (int, float)):
+            smooth_scales = [smooth_scales]
+
+        if isinstance(smooth_scales, list):
+            smooth_scales = numpy.array(smooth_scales, dtype=numpy.float32)
+
+        smooth_scales *= mpc2box
+
+    field_interp = evaluate_cartesian(*fields, pos=pos,
+                                      smooth_scales=smooth_scales,
+                                      verbose=verbose)
+
+    # Now we reshape the interpolated field to have the same shape as the
+    # input `sky_pos`.
+    if smooth_scales is None:
+        shape_single = (nsamples, len(rdist))
+    else:
+        shape_single = (nsamples, len(rdist), len(smooth_scales))
+
+    field_interp_reshaped = [None] * len(fields)
+    for i in range(len(fields)):
+        samples = numpy.full(shape_single, numpy.nan,
+                             dtype=field_interp[i].dtype)
+        for j in range(nsamples):
+            start, end = j * len(rdist), (j + 1) * len(rdist)
+            samples[j] = field_interp[i][start:end, ...]
+
+        field_interp_reshaped[i] = samples
+
+    if len(fields) == 1:
+        return rdist, field_interp_reshaped[0]
+
+    return rdist, field_interp_reshaped
+
 
 ###############################################################################
 #                              Sky maps                                       #