reconstruct

def backproject(
    data: xr.DataArray,
    tx_position: xr.DataArray,
    rx_position: xr.DataArray,
    image_x: ArrayLike,
    image_y: ArrayLike,
    image_z: ArrayLike = 0,
    c: float = 343.0,
    rvg: bool = True,
    complex_image: bool = True,
    channel_dim: str = "channel",
    time_dim: str = "time",
    position_dim: str = "xyz",
) -> xr.DataArray:
    """Reconstruct an image using the backprojection algorithm.

    The backprojection algorithm takes each trace of data and projects it back onto the
    positions on the reconstruction surface that it could have originated from. For each
    trace, this results in an arc of data. When summed over all pings and channels, the
    data constructively adds in true target locations and destructively cancels
    elsewhere.

    Parameters
    ----------
    data
        The data to backproject.
    tx_position
        The position of the transmitter for each trace in `data`.
    rx_position
        The position of each receiver for each trace in `data`.
    image_x
        The x coordinates of the focal points to reconstruct at. This must be
        one-dimensional.
    image_y
        The y coordinates of the focal points to reconstruct at. This must be
        one-dimensional.
    image_z
        The z coordinates of the focal points to reconstruct at. This may be either
        scalar if z is constant for all focal points, or a two-dimensional array giving
        a value for each focal point.
    c
        The speed of sound to use during reconstruction.
    rvg
        Whether to apply a range-varying gain during reconstruction to compensate for
        spherical spreading losses. Note that backprojection has an inherent one-way
        gain. Setting this parameter to True applies the remaining one-way gain to
        complete this compensation. Setting it to False does not apply the remaining
        gain, leaving the image with its inherent compensation.
    complex_image
        If True and `data` is real, use [scipy.signal.hilbert][] to convert it to an
        analytic signal before backprojecting. If False, use the `data` array as given.
    channel_dim
        The dimension of `data` and `rx_position` corresponding to the receiver
        channels.
    time_dim
        The dimension of `data` corresponding to the sampling times.
    position_dim
        The dimension of `tx_position` and `rx_position` corresponding to the position
        coordinates.

    Returns
    -------
    image : xr.DataArray
        The complex image formed at the given focal points. This will be two-dimensional
        with the requested x and y coordinates.

    """
    # Check and convert coordinates to DataArrays.
    image_x = _ensure_1d(image_x, "x")
    image_y = _ensure_1d(image_y, "y")

    # Coerce z positions to either a float or a 2D DataArray.
    if not isinstance(image_z, xr.DataArray):
        image_z_arr = np.asarray(image_z)
        if image_z_arr.ndim == 2:
            image_z = xr.DataArray(image_z_arr, coords=[image_x, image_y])
        elif image_z_arr.ndim != 0:
            raise ValueError("image_z must be a float or a 2D array")

    # Decide on the output datatype based on the input and the complex_image flag.
    if data.dtype.kind == "c":
        output_dtype = data.dtype
    elif complex_image:
        output_dtype = complex  # type:ignore[assignment]
    else:
        output_dtype = data.dtype

    # Call the low-level ufunc on each block.
    image = xr.apply_ufunc(
        _backproject,
        data,
        tx_position,
        rx_position,
        input_core_dims=[
            [channel_dim, time_dim],
            [position_dim],
            [channel_dim, position_dim],
        ],
        kwargs={
            "image_x": image_x.values,
            "image_y": image_y.values,
            "image_z": image_z.values if isinstance(image_z, xr.DataArray) else image_z,
            "c": c,
            "t": data[time_dim].values,
            "complex_image": complex_image,
            "rvg": rvg,
        },
        output_core_dims=[["__image_x__", "__image_y__"]],
        dask="parallelized",
        dask_gufunc_kwargs={
            "meta": [np.empty(0, dtype=output_dtype)],
            "output_sizes": {"__image_x__": len(image_x), "__image_y__": len(image_y)},
        },
        keep_attrs=True,
    )

    # Sum over non-image dimensions.
    if image.ndim > 2:
        image = image.sum(image.dims[:-2])

    # And assign the dimensions and coordinates of the focal points.
    image = image.assign_coords(
        __image_x__=image_x.values,
        __image_y__=image_y.values,
    ).rename(
        __image_x__=image_x.dims[0],
        __image_y__=image_y.dims[0],
    )

    return image