ENH: Add the function tri_area_indexed(p, i)

Author: WarrenWeckesser

README.md

@@ -110,6 +110,7 @@ processed with the script in `ufunclab/tools/conv_template.py`.
 | [`cross2`](#cross2)                             | 2-d vector cross product (returns scalar)             |
 | [`cross3`](#cross3)                             | 3-d vector cross product                              |
 | [`tri_area`](#tri_area)                         | Area of triangles in n-dimensional space              |
+| [`tri_area_indexed`](#tri_area_indexed)         | Area of triangles in n-dimensional space              |
 | [`fillnan1d`](#fillnan1d)                       | Replace `nan` using linear interpolation              |
 | [`linear_interp1d`](#linear_interp1d)           | Linear interpolation, like `numpy.interp`             |
 | [`backlash`](#backlash)                         | Backlash operator                                     |
@@ -1530,7 +1531,7 @@ array([[[ -1.,   0.,   0.],
 
 #### `tri_area`
 
-`tri_area(p)` is a gufunc with signature `(3, n) - > ()`.  It computes the
+`tri_area(p)` is a gufunc with signature `(3, n) -> ()`.  It computes the
 area of a triangle defined by three points in n-dimensional space.
 
 ```
@@ -1550,6 +1551,37 @@ of two triangles in 4-dimensional space.
 array([1.73205081, 0.70710678])
 ```
 
+#### `tri_area_indexed`
+
+`tri_area_indexed(p, i)` is a gufunc with signature `(m, n),(3) -> ()`.
+It computes the area of a triangle defined by three points in n-dimensional
+space.  The first argument, `p`, is an array with shape `(m, n)` holding `m`
+points in n-dimensional space.  The second argument, `i`, is a 1-d array with
+length three that holds indices into `p`.  The core calculation is equivalent
+to `tri_area(p[i])`.
+
+```
+>>> import numpy as np
+>>> from ufunclab import tri_area_indexed, tri_area
+
+>>> p = np.array([[0.0, 0.0, 0.0, 6.0],
+                  [1.0, 2.0, 3.0, 6.0],
+                  [0.0, 2.0, 2.0, 6.0],
+                  [1.5, 1.0, 2.5, 2.0],
+                  [4.0, 1.0, 0.0, 2.5],
+                  [2.0, 1.0, 2.0, 2.5]])
+>>> tri_area_indexed(p, [0, 2, 3])
+6.224949798994367
+>>> tri_area(p[[0, 2, 3]])
+6.224949798994367
+
+Compute the areas of several triangles formed from points in `p`.
+Note that the last two are the same triangles.
+
+>>> tri_area_indexed(p, [[0, 2, 3], [1, 3, 4], [3, 4, 5], [-1, -2, -3]])
+array([6.2249498 , 7.46449931, 0.70710678, 0.70710678])
+```
+
 #### `fillnan1d`
 
 `fillnan1d(x)` is a gufunc with signature `(i)->(i)`.  It uses linear

pyproject.toml

@@ -8,7 +8,7 @@ requires = [
 
 [project]
 name = 'ufunclab'
-version = '0.0.8.dev9'
+version = '0.0.8.dev10'
 description = 'NumPy ufuncs and utilities.'
 readme = 'README.md'
 requires-python = '>=3.9'

src/tri_area/define_cxx_gufunc_extmod.py

@@ -36,22 +36,80 @@
 array([1.73205081, 0.70710678])
 """
 
-ufunc_src = UFuncSource(
+TRI_AREA_INDEXED_DOCSTRING = """\
+tri_area_indexed(p, i, /, ...)
+
+Compute the area of the triangle formed by three points
+in n-dimensional space.  `p` is an array of `m` points
+in n-dimensional space, and `i` is a 1-d array of
+three integers.  These integers are indices into `p`.
+
+Parameters
+----------
+p : array_like, shape (m, n)
+    Input array of points.
+i : array_like, shape (3,)
+    Indices into `i`.
+
+Returns
+-------
+out : ndarray
+    Area of the triangle formed by `p[i[0]]`, `p[i[1]]` and `p[i[2]]`.
+
+Examples
+--------
+
+>>> import numpy as np
+>>> from ufunclab import tri_area_indexed, tri_area
+
+>>> p = np.array([[0.0, 0.0, 0.0, 6.0],
+                  [1.0, 2.0, 3.0, 6.0],
+                  [0.0, 2.0, 2.0, 6.0],
+                  [1.5, 1.0, 2.5, 2.0],
+                  [4.0, 1.0, 0.0, 2.5],
+                  [2.0, 1.0, 2.0, 2.5]])
+>>> tri_area_indexed(p, [0, 2, 3])
+6.224949798994367
+>>> tri_area(p[[0, 2, 3]])
+6.224949798994367
+
+Compute the areas of several triangles formed from points in `p`.
+Note that the last two are the same triangles.
+
+>>> tri_area_indexed(p, [[0, 2, 3], [1, 3, 4], [3, 4, 5], [-1, -2, -3]])
+array([6.2249498 , 7.46449931, 0.70710678, 0.70710678])
+
+"""
+
+ta_ufunc_src = UFuncSource(
     funcname='tri_area_core_calc',
     typesignatures=['f->f', 'd->d', 'g->g'],
 )
 
-
-ufunc = UFunc(
+ta_ufunc = UFunc(
     name='tri_area',
     header='tri_area_gufunc.h',
     docstring=TRI_AREA_DOCSTRING,
     signature='(3,n)->()',
-    sources=[ufunc_src],
+    sources=[ta_ufunc_src],
+)
+
+tai_ufunc_src = UFuncSource(
+    funcname='tri_area_indexed_core_calc',
+    typesignatures=['fp->f', 'dp->d', 'gp->g'],
+)
+
+tai_ufunc = UFunc(
+    name='tri_area_indexed',
+    header='tri_area_gufunc.h',
+    docstring=TRI_AREA_INDEXED_DOCSTRING,
+    signature='(m,n),(3)->()',
+    sources=[tai_ufunc_src],
 )
 
 extmod = UFuncExtMod(
     module='_tri_area',
-    docstring="This extension module defines the gufunc 'tri_area'.",
-    ufuncs=[ufunc],
+    docstring=("This extension module defines the gufuncs 'tri_area' "
+               "and 'tri_area_indexed'."),
+    ufuncs=[ta_ufunc, tai_ufunc],
 )

src/tri_area/tri_area_gufunc.h

@@ -27,26 +27,28 @@ static T norm_diff(npy_intp n, T *p_p, const npy_intp p_strides[2], npy_intp i1,
 }
 
 //
-// `tri_area_core_calc`, the C++ core function
-// for the gufunc `tri_area` with signature '(3,n)->()'
-// for types ['f->f', 'd->d', 'g->g'].
+// This function is used by tri_area_core_calc and tri_area_indexed_core_calc.
 //
 template<typename T>
-static void tri_area_core_calc(
+static void
+tri_area_common_calc(
     npy_intp n,                   // core dimension n
-    T *p_p,                       // pointer to first element of p, a strided 2-d array with shape (3, n)
+    T *p_p,                       // pointer to first element of p, a strided 2-d array with shape (m, n)
     const npy_intp p_strides[2],  // array of length 2 of strides (in bytes) of p
+    npy_intp i0,                  // index of first point
+    npy_intp i1,                  // index of second point
+    npy_intp i2,                  // index of third point
     T *p_out                      // pointer to out
 )
 {
-    T a = norm_diff(n, p_p, p_strides, 0, 1);
-    T b = norm_diff(n, p_p, p_strides, 0, 2);
+    T a = norm_diff(n, p_p, p_strides, i0, i1);
+    T b = norm_diff(n, p_p, p_strides, i0, i2);
+    T c = norm_diff(n, p_p, p_strides, i1, i2);
     if (b > a) {
         T t = a;
         a = b;
         b = t;
     }
-    T c = norm_diff(n, p_p, p_strides, 1, 2);
     if (c > a) {
         T t = c;
         c = b;
@@ -69,4 +71,63 @@ static void tri_area_core_calc(
     *p_out = sqrt(f1*f2*f3*f4)/4;
 }
 
+//
+// `tri_area_core_calc`, the C++ core function
+// for the gufunc `tri_area` with signature '(3,n)->()'
+// for types ['f->f', 'd->d', 'g->g'].
+//
+template<typename T>
+static void tri_area_core_calc(
+    npy_intp n,                   // core dimension n
+    T *p_p,                       // pointer to first element of p, a strided 2-d array with shape (3, n)
+    const npy_intp p_strides[2],  // array of length 2 of strides (in bytes) of p
+    T *p_out                      // pointer to out
+)
+{
+    tri_area_common_calc(n, p_p, p_strides, 0, 1, 2, p_out);
+}
+
+
+static inline
+npy_intp get_index(npy_intp *p_i, npy_intp i_stride, npy_intp index, npy_intp m)
+{
+    npy_intp i = get(p_i, i_stride, index);
+    if (i < 0) {
+        i += m;
+    }
+    if (i < 0 || i >= m) {
+        return -1;
+    }
+    return i;
+}
+
+//
+// `tri_area_indexed_core_calc`, the C++ core function
+// for the gufunc `tri_area_indexed` with signature '(m, n),(3)->()'
+// for types ['fp->f', 'dp->d', 'gp->g'].
+//
+template<typename T>
+static void tri_area_indexed_core_calc(
+    npy_intp m,                   // core dimension m
+    npy_intp n,                   // core dimension n
+    T *p_p,                       // pointer to first element of p, a strided
+                                  // 2-d array with shape (m, n)
+    const npy_intp p_strides[2],  // array of length 2 of strides (in bytes) of p
+    npy_intp *p_i,                // pointer to first element of i, a strided 1-d array
+    npy_intp i_stride,            // stride of elements in i
+    T *p_out                      // pointer to out
+)
+{
+    npy_intp i0, i1, i2;
+    i0 = get_index(p_i, i_stride, 0, m);
+    i1 = get_index(p_i, i_stride, 1, m);
+    i2 = get_index(p_i, i_stride, 2, m);
+    if (i0 == -1 || i1 == -1 || i2 == -1) {
+        // XXX Should probably raise IndexError instead of returning NAN.
+        *p_out = NPY_NAN;
+        return;
+    }
+    tri_area_common_calc(n, p_p, p_strides, i0, i1, i2, p_out);
+}
+
 #endif

ufunclab/__init__.py

@@ -26,7 +26,7 @@
 from ._wjaccard import wjaccard
 from ._mad import mad, rmad, gini
 from ._vnorm import vnorm, rms, vdot
-from ._tri_area import tri_area
+from ._tri_area import tri_area, tri_area_indexed
 from ._backlash import backlash
 from ._fillnan1d import fillnan1d
 from ._linear_interp1d import linear_interp1d

ufunclab/tests/test_tri_area.py

@@ -1,7 +1,7 @@
 import pytest
 import numpy as np
 from numpy.testing import assert_array_equal, assert_allclose
-from ufunclab import tri_area
+from ufunclab import tri_area, tri_area_indexed
 
 
 @pytest.mark.parametrize('dt', [np.float32, np.float64, np.longdouble])
@@ -58,3 +58,20 @@ def test_4d():
                   [1.0, 2.0, 5.0, 0.0]])
     a = tri_area(p)
     assert_allclose(a, 1.0, rtol=1e-14)
+
+
+def test_tri_area_indexed_basic():
+    p = np.array([[2, 3, 0, 1],
+                  [0, 0, 0, 0],
+                  [1, 1, 3, 8],
+                  [4, 0, -1, 1],
+                  [3, 3, 3, 2]])
+    a1 = tri_area(p[:3])
+    a2 = tri_area(p[2:])
+    a3 = tri_area(p[::2])
+    a4 = tri_area(p[[-1, -2, 1]])
+    a = tri_area_indexed(p, [[0, 1, 2],
+                             [2, 3, 4],
+                             [0, 2, 4],
+                             [-1, -2, 1]])
+    assert_allclose(a, [a1, a2, a3, a4])