Skip to content

array

triwarp.array

NumPy-style structural and elementwise ops on Warp arrays (ranges, gather, sort, masks).

DType module-attribute 

DType = TypeVar('DType')

SORT_ROWS_INSERTION_MAX_COLS module-attribute 

SORT_ROWS_INSERTION_MAX_COLS = 8

arange 

arange(
    start: int,
    stop: int | None = None,
    step: int = 1,
    *,
    device: DeviceLike,
) -> wp.array[wp.int32]

Evenly spaced integers over the half-open interval [start, stop) (numpy.arange).

Called with a varying number of positional arguments, exactly as numpy.arange: arange(stop, device=...) runs from 0, arange(start, stop, device=...) from start, and arange(start, stop, step, device=...) spaces the values by step. A step may be negative, in which case the interval runs downwards.

Parameters:

Name Type Description Default
start int

Start of the interval, included. Read as stop when stop is omitted, in which case the interval starts at 0.

 required
stop int | None

End of the interval, excluded.

 None
step int

Spacing between consecutive values, so out[i + 1] - out[i] == step. Must be non-zero.

 1
device DeviceLike

Warp device for the result. Keyword-only, unlike the first three, so the positional arguments stay numpy.arange's; required, because nothing in this package allocates onto Warp's ambient current device.

 required

Returns:

Type Description
array[int32]

1-D array of max(0, ceil((stop - start) / step)) values on device.

Raises:

Type Description
ValueError

If step is zero, or the first or last value does not fit in int32.
Notes

There is no dtype argument, which is the one place this departs from numpy.arange's signature. Every index buffer in this package is int32, so a second width has no caller, and registering the kernel overloads for widths nothing reaches would cost compile time on every rebuild. The range check above is what stands in for choosing a wider dtype: a range that does not fit raises rather than silently wrapping.

See Also

arange_repeat The same range with each value repeated a fixed number of times. numpy.arange

Source code in triwarp/array.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def arange(
    start: int, stop: int | None = None, step: int = 1, *, device: wp.DeviceLike
) -> wp.array[wp.int32]:
    """
    Evenly spaced integers over the half-open interval ``[start, stop)`` (``numpy.arange``).

    Called with a varying number of positional arguments, exactly as [`numpy.arange`][]:
    ``arange(stop, device=...)`` runs from ``0``, ``arange(start, stop, device=...)`` from
    ``start``, and ``arange(start, stop, step, device=...)`` spaces the values by ``step``. A
    ``step`` may be negative, in which case the interval runs downwards.

    Parameters
    ----------
    start
        Start of the interval, included. Read as ``stop`` when ``stop`` is omitted, in which case
        the interval starts at ``0``.
    stop
        End of the interval, excluded.
    step
        Spacing between consecutive values, so ``out[i + 1] - out[i] == step``. Must be non-zero.
    device
        Warp device for the result. Keyword-only, unlike the first three, so the positional
        arguments stay [`numpy.arange`][]'s; required, because nothing in this package allocates
        onto Warp's ambient current device.

    Returns
    -------
    wp.array[wp.int32]
        1-D array of ``max(0, ceil((stop - start) / step))`` values on ``device``.

    Raises
    ------
    ValueError
        If ``step`` is zero, or the first or last value does not fit in ``int32``.

    Notes
    -----
    There is no ``dtype`` argument, which is the one place this departs from [`numpy.arange`][]'s
    signature. Every index buffer in this package is ``int32``, so a second width has no caller,
    and registering the kernel overloads for widths nothing reaches would cost compile time on
    every rebuild. The range check above is what stands in for choosing a wider dtype: a range
    that does not fit raises rather than silently wrapping.

    See Also
    --------
    [`arange_repeat`][triwarp.array.arange_repeat]
        The same range with each value repeated a fixed number of times.
    [`numpy.arange`][]
    """
    if stop is None:
        start, stop = 0, start
    if step == 0:
        raise ValueError("step must be non-zero")
    # ``ceil((stop - start) / step)`` in integers, correct for either sign of ``step``: Python's
    # ``//`` floors, so negating both sides of the division ceils.
    n = max(0, -((start - stop) // step))
    if n > 0:
        _check_int32_fits(start, "start")
        _check_int32_fits(start + (n - 1) * step, "stop")
    out = wp.empty(n, dtype=wp.int32, device=device)
    if n == 0:
        return out
    if start == 0 and step == 1:
        wp.launch(kernel_array.ARANGE[wp.int32], dim=n, inputs=[out], device=device)
    else:
        wp.launch(
            kernel_array.ARANGE_AFFINE[wp.int32],
            dim=n,
            inputs=[wp.int32(start), wp.int32(step), out],
            device=device,
        )
    return out

arange_repeat 

arange_repeat(
    count: int, repeats: int, device: DeviceLike
) -> wp.array[wp.int32]

Fill out[i] = i // repeats (numpy.repeat of an index range).

Parameters:

Name Type Description Default
count int

Number of elements; must be non-negative.

 required
repeats int

How many consecutive entries share an index; must be positive.

 required
device DeviceLike

Warp device for the result.

 required

Returns:

Type Description
array[int32]

Length-count array on device.

Raises:

Type Description
ValueError

If count is negative, repeats is not positive, or the largest value does not fit in int32.
See Also

arange The range this repeats, and whose Notes explain why neither takes a dtype. numpy.repeat

Source code in triwarp/array.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def arange_repeat(count: int, repeats: int, device: wp.DeviceLike) -> wp.array[wp.int32]:
    """
    Fill ``out[i] = i // repeats`` (``numpy.repeat`` of an index range).

    Parameters
    ----------
    count
        Number of elements; must be non-negative.
    repeats
        How many consecutive entries share an index; must be positive.
    device
        Warp device for the result.

    Returns
    -------
    wp.array[wp.int32]
        Length-``count`` array on ``device``.

    Raises
    ------
    ValueError
        If ``count`` is negative, ``repeats`` is not positive, or the largest value does not fit
        in ``int32``.

    See Also
    --------
    [`arange`][triwarp.array.arange]
        The range this repeats, and whose ``Notes`` explain why neither takes a ``dtype``.
    [`numpy.repeat`][]
    """
    if count < 0:
        raise ValueError(f"count must be non-negative, got {count}")
    if repeats <= 0:
        raise ValueError(f"repeats must be positive, got {repeats}")
    if count > 0:
        _check_int32_fits((count - 1) // repeats, "count // repeats")
    out = wp.empty(count, dtype=wp.int32, device=device)
    if count > 0:
        wp.launch(
            kernel_array.ARANGE_REPEAT[wp.int32],
            dim=count,
            inputs=[wp.int32(repeats), out],
            device=device,
        )
    return out

sort_pair_indices 

sort_pair_indices(
    n: int, fill_value: int, device: DeviceLike
) -> wp.array[wp.int32]

Fill [0, 1, ..., n-1, fill_value, ..., fill_value] (length 2 * n).

The payload buffer warp.utils.radix_sort_pairs wants: the first half seeded with the identity permutation, the second half (its scratch) filled with a padding value.

Parameters:

Name Type Description Default
n int

Number of real entries; the result has length 2 * n.

 required
fill_value int

Padding written into the upper half.

 required
device DeviceLike

Warp device for the result.

 required

Returns:

Type Description
array[int32]

Length-2 * n array on device.

Raises:

Type Description
ValueError

If n is negative, or n - 1 / fill_value does not fit in int32.
See Also

sort_and_argsort arange Whose Notes explain why neither takes a dtype.

Source code in triwarp/array.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def sort_pair_indices(n: int, fill_value: int, device: wp.DeviceLike) -> wp.array[wp.int32]:
    """
    Fill ``[0, 1, ..., n-1, fill_value, ..., fill_value]`` (length ``2 * n``).

    The payload buffer ``warp.utils.radix_sort_pairs`` wants: the first half seeded with the
    identity permutation, the second half (its scratch) filled with a padding value.

    Parameters
    ----------
    n
        Number of real entries; the result has length ``2 * n``.
    fill_value
        Padding written into the upper half.
    device
        Warp device for the result.

    Returns
    -------
    wp.array[wp.int32]
        Length-``2 * n`` array on ``device``.

    Raises
    ------
    ValueError
        If ``n`` is negative, or ``n - 1`` / ``fill_value`` does not fit in ``int32``.

    See Also
    --------
    [`sort_and_argsort`][triwarp.array.sort_and_argsort]
    [`arange`][triwarp.array.arange]
        Whose ``Notes`` explain why neither takes a ``dtype``.
    """
    if n < 0:
        raise ValueError(f"n must be non-negative, got {n}")
    if n > 0:
        _check_int32_fits(n - 1, "n")
    _check_int32_fits(fill_value, "fill_value")
    out = wp.empty(2 * n, dtype=wp.int32, device=device)
    if n > 0:
        wp.launch(
            kernel_array.SORT_PAIR_INDICES[wp.int32],
            dim=2 * n,
            inputs=[wp.int32(n), wp.int32(fill_value), out],
            device=device,
        )
    return out

pack_1d_arrays 

pack_1d_arrays(
    arrays: Sequence[array[Scalar]], *, copy: bool = True
) -> tuple[wp.array[wp.Scalar], wp.array[wp.int32]]

Concatenate several 1-D warp.array instances into one buffer plus per-segment offsets.

Segment i starts at offsets[i] in flat; its length is arrays[i].size, so it occupies flat[offsets[i] : offsets[i] + arrays[i].size]. This is the usual packed representation for variable-length per-item lists on the device (no nested arrays).

The packed pair is spelled values first, offsets second

This function is where the package's convention is stated, because it is the primitive the others are built on: a packed buffer and its offsets are returned, and accepted, values first. Fifteen public functions hand back such a pair — boundary_loops_batched, successor_cycles, bfs_multi_source, query_ball_with_offsets, geodesic_ball, vertex_face_adjacency, vertex_one_rings and the triwarp.geodesic_walk tracers among them — and every function that takes one (split, trace_polylines, submeshes_from_face_groups) takes it in the same order. Both halves are wp.int32 in the common case, so a transposed unpack type-checks, runs, and indexes garbage; there is nothing but the convention to lean on. Where a third array rides along it is a per-item one and goes last, as in (ring_halfedges, offsets, is_boundary).

Parameters:

Name Type Description Default
arrays Sequence[array[Scalar]]

Non-empty sequence of 1-D arrays sharing the same dtype and device.

 required
copy bool

Keep False for a read-only result. arrays that are already consecutive non-empty views of one buffer -- what split returns with copy=False -- are then handed back as that buffer's span instead of being copied into a new one, so the split round trip costs nothing at all. The default copies, so writing into flat is safe; with copy=False such a write reaches the segments.

 True

Returns:

Type Description
flat

1-D array of length sum(a.size for a in arrays), same dtype and device as the inputs.
offsets

Length len(arrays) (the start offset of each segment, an exclusive scan of the segment sizes), dtype wp.int32, same device as the inputs. offsets[0] == 0. This is not a total-terminated CSR array — there is no offsets[-1] == flat.size terminator, so the last segment's length must be taken from arrays[-1].size (or flat.size - offsets[-1]).

Raises:

Type Description
ValueError

If arrays is empty, any input is not rank-1, or their dtype differs.
RuntimeError

If the segments are not all on one device.
See Also

split The inverse: recovers the per-segment arrays from (flat, offsets). concatenate

Source code in triwarp/array.py 
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
def pack_1d_arrays(
    arrays: Sequence[wp.array[wp.Scalar]], *, copy: bool = True
) -> tuple[wp.array[wp.Scalar], wp.array[wp.int32]]:
    """
    Concatenate several 1-D ``warp.array`` instances into one buffer plus per-segment offsets.

    Segment ``i`` starts at ``offsets[i]`` in ``flat``; its length is ``arrays[i].size``, so it
    occupies ``flat[offsets[i] : offsets[i] + arrays[i].size]``. This is the usual packed
    representation for variable-length per-item lists on the device (no nested arrays).

    !!! note "The packed pair is spelled values first, offsets second"
        This function is where the package's convention is stated, because it is the primitive the
        others are built on: **a packed buffer and its offsets are returned, and accepted, values
        first.** Fifteen public functions hand back such a pair —
        [`boundary_loops_batched`][triwarp.boundary.boundary_loops_batched],
        [`successor_cycles`][triwarp.graph.successor_cycles],
        [`bfs_multi_source`][triwarp.graph.bfs_multi_source],
        [`query_ball_with_offsets`][triwarp.neighbors.query_ball_with_offsets],
        [`geodesic_ball`][triwarp.neighbors.geodesic_ball],
        [`vertex_face_adjacency`][triwarp.adjacency.vertex_face_adjacency],
        [`vertex_one_rings`][triwarp.halfedge.vertex_one_rings] and the
        [`triwarp.geodesic_walk`][triwarp.geodesic_walk] tracers among them — and every function
        that *takes* one ([`split`][triwarp.array.split],
        [`trace_polylines`][triwarp.geodesic_walk.trace_polylines],
        [`submeshes_from_face_groups`][triwarp.selection.submeshes_from_face_groups]) takes it in
        the same order. Both halves are ``wp.int32`` in the common case, so a transposed unpack
        type-checks, runs, and indexes garbage; there is nothing but the convention to lean on.
        Where a third array rides along it is a *per-item* one and goes last, as in
        ``(ring_halfedges, offsets, is_boundary)``.

    Parameters
    ----------
    arrays
        Non-empty sequence of 1-D arrays sharing the same ``dtype`` and ``device``.
    copy
        Keep ``False`` for a read-only result. ``arrays`` that are already consecutive non-empty
        views of one buffer -- what [`split`][triwarp.array.split] returns with ``copy=False`` --
        are then handed back as that buffer's span instead of being copied into a new one, so the
        ``split`` round trip costs nothing at all. The default copies, so writing into ``flat`` is
        safe; with ``copy=False`` such a write reaches the segments.

    Returns
    -------
    flat
        1-D array of length ``sum(a.size for a in arrays)``, same ``dtype`` and ``device`` as
        the inputs.
    offsets
        Length ``len(arrays)`` (the start offset of each segment, an exclusive scan of the
        segment sizes), ``dtype`` ``wp.int32``, same ``device`` as the inputs. ``offsets[0] == 0``.
        This is *not* a total-terminated CSR array — there is no ``offsets[-1] == flat.size``
        terminator, so the last segment's length must be taken from ``arrays[-1].size`` (or
        ``flat.size - offsets[-1]``).

    Raises
    ------
    ValueError
        If ``arrays`` is empty, any input is not rank-1, or their ``dtype`` differs.
    RuntimeError
        If the segments are not all on one device.

    See Also
    --------
    [`split`][triwarp.array.split]
        The inverse: recovers the per-segment arrays from ``(flat, offsets)``.
    [`concatenate`][triwarp.array.concatenate]
    """
    flat, offsets = _pack_segments(arrays, caller="pack_1d_arrays", copy=copy)
    return flat, wp.array(offsets, dtype=wp.int32, device=flat.device)

concatenate 

concatenate(
    arrays: Sequence[array[DType]], *, copy: bool = True
) -> wp.array[DType]

Concatenate 1-D warp.array instances in order (numpy.concatenate).

Parameters:

Name Type Description Default
arrays Sequence[array[DType]]

Non-empty sequence of rank-1 arrays sharing the same dtype and device. Empty segments are allowed.

 required
copy bool

Keep False for a read-only result, which then costs nothing when the segments already tile one buffer -- see pack_1d_arrays.

 True

Returns:

Type Description
array

Contiguous 1-D array of length sum(a.size for a in arrays) on the input device. When arrays has a single element, that array is returned without copying whatever copy says, since there is nothing to concatenate it with.

Raises:

Type Description
ValueError

If arrays is empty, any input is not rank-1, or their dtype differs.
RuntimeError

If the segments are not all on one device.
See Also

pack_1d_arrays concatenate The mesh-level operation of the same name, which joins (vertices, faces) pairs and reindexes as it goes. Both names are required: this one mirrors numpy.concatenate, that one trimesh.util.concatenate. numpy.concatenate

Source code in triwarp/array.py 
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
def concatenate(arrays: Sequence[wp.array[DType]], *, copy: bool = True) -> wp.array[DType]:
    """
    Concatenate 1-D ``warp.array`` instances in order (``numpy.concatenate``).

    Parameters
    ----------
    arrays
        Non-empty sequence of rank-1 arrays sharing the same ``dtype`` and ``device``.
        Empty segments are allowed.
    copy
        Keep ``False`` for a read-only result, which then costs nothing when the segments already
        tile one buffer -- see [`pack_1d_arrays`][triwarp.array.pack_1d_arrays].

    Returns
    -------
    wp.array
        Contiguous 1-D array of length ``sum(a.size for a in arrays)`` on the input
        device. When ``arrays`` has a single element, that array is returned without
        copying whatever ``copy`` says, since there is nothing to concatenate it with.

    Raises
    ------
    ValueError
        If ``arrays`` is empty, any input is not rank-1, or their ``dtype`` differs.
    RuntimeError
        If the segments are not all on one device.

    See Also
    --------
    [`pack_1d_arrays`][triwarp.array.pack_1d_arrays]
    [`concatenate`][triwarp.combine.concatenate]
        The mesh-level operation of the same name, which joins ``(vertices, faces)`` pairs and
        reindexes as it goes. Both names are required: this one mirrors
        [`numpy.concatenate`][], that one [`trimesh.util.concatenate`][].
    [`numpy.concatenate`][]
    """
    if len(arrays) == 0:
        raise ValueError("arrays must be non-empty")
    if len(arrays) == 1:
        arr = arrays[0]
        if int(arr.ndim) != 1:
            raise ValueError(f"concatenate requires rank-1 arrays, got ndim={arr.ndim}")
        return arr
    return _pack_segments(arrays, caller="concatenate", copy=copy)[0]

split 

split(
    array: array[DType],
    offsets: array[int32],
    *,
    copy: bool = False,
) -> list[wp.array[DType]]

Break a packed 1-D array into its per-segment arrays (numpy.split).

The inverse of pack_1d_arrays: segment i is array[offsets[i] : offsets[i + 1]], with the last segment running to the end of array. One host readback of offsets, then zero copies by default — the segments are views into array, which they keep alive.

Parameters:

Name Type Description Default
array array[DType]

Rank-1 array to split, any dtype.

 required
offsets array[int32]

Length-n_segments wp.int32 exclusive prefix sum of the segment sizes, starting at 0 — exactly what pack_1d_arrays and counts_to_offsets return. The total-terminated n + 1 form (include_total=True) is also accepted; its trailing entry simply yields one final empty segment, so pass the length-n form when that matters.

 required
copy bool

When True, return independent wp.clone copies instead of views.

 False

Returns:

Type Description
list[array]

One array per segment, on array.device, in segment order. Empty list when offsets is empty.

Raises:

Type Description
ValueError

If array or offsets is not rank-1, or offsets is not a non-decreasing sequence starting at 0 and bounded by array's length.
RuntimeError

If array and offsets are not all on one device.
See Also

pack_1d_arrays The inverse: packs per-segment arrays into one buffer plus these offsets. split The mesh-level operation of the same name, which separates a mesh into connected components. Both names are required: this one mirrors numpy.split, that one trimesh.Trimesh.split. numpy.split

Source code in triwarp/array.py 
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
def split(
    array: wp.array[DType], offsets: wp.array[wp.int32], *, copy: bool = False
) -> list[wp.array[DType]]:
    """
    Break a packed 1-D array into its per-segment arrays (``numpy.split``).

    The inverse of [`pack_1d_arrays`][triwarp.array.pack_1d_arrays]: segment ``i`` is
    ``array[offsets[i] : offsets[i + 1]]``, with the last segment running to the end of
    ``array``. One host readback of ``offsets``, then zero copies by default — the segments are
    views into ``array``, which they keep alive.

    Parameters
    ----------
    array
        Rank-1 array to split, any ``dtype``.
    offsets
        Length-``n_segments`` ``wp.int32`` exclusive prefix sum of the segment sizes, starting
        at ``0`` — exactly what [`pack_1d_arrays`][triwarp.array.pack_1d_arrays] and
        [`counts_to_offsets`][triwarp.array.counts_to_offsets] return. The total-terminated
        ``n + 1`` form (``include_total=True``) is also accepted; its trailing entry simply
        yields one final empty segment, so pass the length-``n`` form when that matters.
    copy
        When ``True``, return independent ``wp.clone`` copies instead of views.

    Returns
    -------
    list[wp.array]
        One array per segment, on ``array.device``, in segment order. Empty list when
        ``offsets`` is empty.

    Raises
    ------
    ValueError
        If ``array`` or ``offsets`` is not rank-1, or ``offsets`` is not a non-decreasing
        sequence starting at ``0`` and bounded by ``array``'s length.
    RuntimeError
        If ``array`` and ``offsets`` are not all on one device.

    See Also
    --------
    [`pack_1d_arrays`][triwarp.array.pack_1d_arrays]
        The inverse: packs per-segment arrays into one buffer plus these offsets.
    [`split`][triwarp.combine.split]
        The mesh-level operation of the same name, which separates a mesh into connected
        components. Both names are required: this one mirrors [`numpy.split`][], that one
        ``trimesh.Trimesh.split``.
    [`numpy.split`][]
    """
    require_same_device(array=array, offsets=offsets)
    if int(array.ndim) != 1:
        raise ValueError(f"split requires a rank-1 array, got ndim={array.ndim}")
    if int(offsets.ndim) != 1:
        raise ValueError(f"split requires rank-1 offsets, got ndim={offsets.ndim}")

    n = int(array.shape[0])
    starts = [int(start) for start in offsets.list()]
    if not starts:
        return []
    bounds = [*starts, n]
    if starts[0] != 0 or any(a > b for a, b in itertools.pairwise(bounds)):
        raise ValueError(
            f"offsets must start at 0 and be non-decreasing within [0, {n}], got {starts}"
        )
    # Warp rejects a zero-length slice at the very end of a buffer (``arr[n:n]``) while accepting
    # an interior one, so an empty trailing segment needs its own allocation.
    #
    # Views are built with ``array[begin:end]`` rather than a raw ``wp.array(ptr=..., shape=...,
    # strides=...)`` construction: the latter re-implements ``wp.array.__getitem__``'s contract (20
    # attributes) and already gets one of them wrong, dropping the ``grad`` view a slice of a
    # ``requires_grad`` array carries.
    #
    # ``copy=True`` clones each segment independently rather than filling one shared buffer with
    # disjoint views into it, because a shared allocation would mean holding one segment pins the
    # whole buffer alive -- the opposite of what ``copy=True`` promises.
    segments = [
        array[begin:end] if end > begin else wp.empty(0, dtype=array.dtype, device=array.device)
        for begin, end in itertools.pairwise(bounds)
    ]
    return [wp.clone(segment) for segment in segments] if copy else segments

allclose 

allclose(
    a: array[Float] | array[vec3],
    b: array[Float] | array[vec3],
    *,
    rtol: float = 1e-05,
    atol: float = 1e-08,
) -> bool

Test whether two arrays are element-wise equal within a tolerance (numpy.allclose).

Reduces |a - b| <= atol + rtol * |b| (element-wise, and component-wise for wp.vec3) to a single Python bool on-device, without copying either array to the host. Matches the asymmetric numpy.allclose / torch.allclose tolerance convention.

Parameters:

Name Type Description Default
a array[Float] | array[vec3]

Length-n array of any float dtype (float16 / float32 / float64) or of wp.vec3, on the target device.

 required
b array[Float] | array[vec3]

Array of the same length and dtype as a.

 required
rtol float

Relative tolerance. Defaults to 1e-05. Converted to a's precision.

 1e-05
atol float

Absolute tolerance. Defaults to 1e-08. Converted to a's precision, so a float16 comparison cannot resolve a tolerance below its own epsilon.

 1e-08

Returns:

Type Description
bool

True when every element (every component, for wp.vec3) is within tolerance. True for empty inputs, following the numpy.allclose convention.

Raises:

Type Description
ValueError

If a and b have different lengths or dtypes.
RuntimeError

If a and b are not all on one device.
Source code in triwarp/array.py 
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
def allclose(
    a: wp.array[wp.Float] | wp.array[wp.vec3],
    b: wp.array[wp.Float] | wp.array[wp.vec3],
    *,
    rtol: float = 1e-05,
    atol: float = 1e-08,
) -> bool:
    """
    Test whether two arrays are element-wise equal within a tolerance (``numpy.allclose``).

    Reduces ``|a - b| <= atol + rtol * |b|`` (element-wise, and component-wise for ``wp.vec3``)
    to a single Python ``bool`` on-device, without copying either array to the host. Matches the
    asymmetric ``numpy.allclose`` / ``torch.allclose`` tolerance convention.

    Parameters
    ----------
    a
        Length-``n`` array of any float dtype (``float16`` / ``float32`` / ``float64``) or of
        ``wp.vec3``, on the target device.
    b
        Array of the same length and dtype as ``a``.
    rtol
        Relative tolerance. Defaults to ``1e-05``. Converted to ``a``'s precision.
    atol
        Absolute tolerance. Defaults to ``1e-08``. Converted to ``a``'s precision, so a
        ``float16`` comparison cannot resolve a tolerance below its own epsilon.

    Returns
    -------
    bool
        ``True`` when every element (every component, for ``wp.vec3``) is within tolerance.
        ``True`` for empty inputs, following the ``numpy.allclose`` convention.

    Raises
    ------
    ValueError
        If ``a`` and ``b`` have different lengths or dtypes.
    RuntimeError
        If ``a`` and ``b`` are not all on one device.
    """
    require_same_device(a=a, b=b)
    if a.dtype != b.dtype:
        raise ValueError(f"allclose requires matching dtypes, got {a.dtype} and {b.dtype}")
    n = int(a.shape[0])
    if n != int(b.shape[0]):
        raise ValueError(f"allclose requires equal lengths, got {n} and {b.shape[0]}")
    if n == 0:
        return True

    mask = wp.empty(n, dtype=wp.bool, device=a.device)
    if a.dtype == wp.vec3:
        wp.map(kernel_array.is_close_vec3, a, b, wp.float32(rtol), wp.float32(atol), out=mask)
    else:
        scalar = a.dtype
        wp.map(kernel_array.is_close_scalar, a, b, scalar(rtol), scalar(atol), out=mask)
    return bool(tw.reduce.all(mask))

sort_and_argsort 

sort_and_argsort(
    keys: array[Scalar], *, fill_value: int = -1
) -> tuple[wp.array[wp.Scalar], wp.array[wp.int32]]

Ascending sort of keys together with the permutation that produced it.

Both halves of numpy.sort and numpy.argsort at once: a radix sort produces the ordered keys as a side effect of computing the order, so returning only one of the two would throw work away. sorted_keys[i] == keys[order[i]].

Wraps warp.utils.radix_sort_pairs, which needs double-width scratch for both the keys and the payload; this allocates that scratch, seeds the payload with 0..n-1 and hands back length-n views of the sorted prefixes.

Parameters:

Name Type Description Default
keys array[Scalar]

Length-n sort keys (any radix-sortable scalar dtype; see sortable_dtype for which those are).

 required
fill_value int

Padding written into the upper half of the payload buffer, where the sort's scratch lives. Only matters to callers that read past n.

 -1

Returns:

Name Type Description
sorted_keys array

Length-n view of the ascending keys.
order array[int32]

Length-n view of the original index of each sorted key.
Notes

Both results are views into the scratch buffers, kept alive by the returned arrays. Clone them if they must outlive the caller's frame alongside another sort.

The sort is stable: equal keys keep their input order, so order is ascending within each run of duplicate keys. warp.utils.radix_sort_pairs documents this ("the sort is stable and operates in linear time"), it is inherent to its LSD radix passes, and it is verified on both devices against numpy.argsort(kind="stable"). Callers may rely on it -- split_batched does, to keep faces ascending within each component.

See Also

sort_rows sort_pair_indices sortable_dtype

Source code in triwarp/array.py 
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
def sort_and_argsort(
    keys: wp.array[wp.Scalar], *, fill_value: int = -1
) -> tuple[wp.array[wp.Scalar], wp.array[wp.int32]]:
    """
    Ascending sort of ``keys`` together with the permutation that produced it.

    Both halves of ``numpy.sort`` and ``numpy.argsort`` at once: a radix sort produces the ordered
    keys as a side effect of computing the order, so returning only one of the two would throw work
    away. ``sorted_keys[i] == keys[order[i]]``.

    Wraps ``warp.utils.radix_sort_pairs``, which needs double-width scratch for both the keys and
    the payload; this allocates that scratch, seeds the payload with ``0..n-1`` and hands back
    length-``n`` views of the sorted prefixes.

    Parameters
    ----------
    keys
        Length-``n`` sort keys (any radix-sortable scalar dtype; see
        [`sortable_dtype`][triwarp.typing.sortable_dtype] for which those are).
    fill_value
        Padding written into the upper half of the payload buffer, where the sort's scratch lives.
        Only matters to callers that read past ``n``.

    Returns
    -------
    sorted_keys : wp.array
        Length-``n`` view of the ascending keys.
    order : wp.array[wp.int32]
        Length-``n`` view of the original index of each sorted key.

    Notes
    -----
    Both results are **views** into the scratch buffers, kept alive by the returned arrays. Clone
    them if they must outlive the caller's frame alongside another sort.

    The sort is **stable**: equal keys keep their input order, so ``order`` is ascending within
    each run of duplicate keys. ``warp.utils.radix_sort_pairs`` documents this ("the sort is
    stable and operates in linear time"), it is inherent to its LSD radix passes, and it is
    verified on both devices against ``numpy.argsort(kind="stable")``. Callers may rely on it --
    [`split_batched`][triwarp.combine.split_batched] does, to keep faces ascending within each
    component.

    See Also
    --------
    [`sort_rows`][triwarp.array.sort_rows]
    [`sort_pair_indices`][triwarp.array.sort_pair_indices]
    [`sortable_dtype`][triwarp.typing.sortable_dtype]
    """
    device = keys.device
    n = int(keys.shape[0])
    if n == 0:
        return keys, wp.empty(0, dtype=wp.int32, device=device)
    keys_buffer = wp.empty(2 * n, dtype=keys.dtype, device=device)
    wp.copy(keys_buffer, keys, count=n)
    order_buffer = sort_pair_indices(n, fill_value, device)
    wp.utils.radix_sort_pairs(keys_buffer, order_buffer, count=n)
    return keys_buffer[:n], order_buffer[:n]

sort_rows 

sort_rows(data: Array2dInt32 | Array2dFloat32) -> None

Sort each row of a 2D array independently, in place, ascending.

Each row is treated as its own radix-sort segment, so rows are reordered internally but their relative row order is unaffected.

Parameters:

Name Type Description Default
data Array2dInt32 | Array2dFloat32

(n, w) int32 or float32 device array sorted in place, row by row.

 required

Raises:

Type Description
TypeError

If data is not int32 or float32.
Notes

Rows no wider than SORT_ROWS_INSERTION_MAX_COLS are sorted by a per-row insertion sort (one thread per row); wider rows fall back to a segmented radix sort. The narrow path is not a micro-optimization: segmented_sort_pairs pays a fixed cost per segment, which for rows this narrow dominates the comparison work by orders of magnitude.

The dtype restriction comes from that fallback -- warp.utils.segmented_sort_pairs takes int32 or float32 keys and nothing else -- and it is checked up front rather than left to whichever path a row width happens to select, so the accepted dtypes do not depend on w.

Source code in triwarp/array.py 
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
def sort_rows(data: twt.Array2dInt32 | twt.Array2dFloat32) -> None:
    """
    Sort each row of a 2D array independently, in place, ascending.

    Each row is treated as its own radix-sort segment, so rows are reordered internally
    but their relative row order is unaffected.

    Parameters
    ----------
    data
        ``(n, w)`` ``int32`` or ``float32`` device array sorted in place, row by row.

    Raises
    ------
    TypeError
        If ``data`` is not ``int32`` or ``float32``.

    Notes
    -----
    Rows no wider than ``SORT_ROWS_INSERTION_MAX_COLS`` are sorted by a per-row insertion sort (one
    thread per row); wider rows fall back to a segmented radix sort. The narrow path is not a
    micro-optimization: ``segmented_sort_pairs`` pays a fixed cost per *segment*, which for rows
    this narrow dominates the comparison work by orders of magnitude.

    The dtype restriction comes from that fallback -- ``warp.utils.segmented_sort_pairs`` takes
    ``int32`` or ``float32`` keys and nothing else -- and it is checked up front rather than left
    to whichever path a row width happens to select, so the accepted dtypes do not depend on ``w``.
    """
    # The narrow path's margin, for whoever considers deleting it: a fixed per-segment cost in
    # ``segmented_sort_pairs`` dominates the comparison work for rows this narrow, where a plain
    # compare-and-swap needs none of it.
    #
    # The dtype guard is here rather than at the wide branch because without it the two paths
    # disagree: the insertion kernel is generic, so a ``float64`` table used to sort silently at
    # ``w <= SORT_ROWS_INSERTION_MAX_COLS`` and raise ``RuntimeError: Unsupported data type:
    # float64`` from inside Warp one column later. Support that turns on the row width is worse
    # than no support.
    if data.dtype not in (wp.int32, wp.float32):
        raise TypeError(f"sort_rows requires an int32 or float32 array, got {data.dtype}")
    n = data.size
    n_rows, n_cols = int(data.shape[0]), int(data.shape[1])
    if n_rows == 0 or n_cols < 2:
        return
    if n_cols <= SORT_ROWS_INSERTION_MAX_COLS:
        wp.launch(
            kernel_array.SORT_ROWS_INSERTION[data.dtype],
            dim=n_rows,
            inputs=[data],
            device=data.device,
        )
        return

    data_buffer = wp.empty(n * 2, dtype=data.dtype, device=data.device)
    wp.copy(data_buffer, data, count=n)
    indices_buffer = sort_pair_indices(n, -1, data.device)
    segment_start_indices = arange(0, (n // n_cols + 1) * n_cols, n_cols, device=data.device)
    wp.utils.segmented_sort_pairs(
        data_buffer, indices_buffer, n, segment_start_indices=segment_start_indices
    )
    wp.copy(data, data_buffer, count=n)

triplet_buffers 

triplet_buffers(
    n_triplets: int, dtype: type, device: DeviceLike
) -> tuple[
    wp.array[wp.int32], wp.array[wp.int32], wp.array
]

Uninitialized (rows, cols, values) COO buffers for one bsr_from_triplets build.

Parameters:

Name Type Description Default
n_triplets int

Length of each of the three buffers: the number of (row, col, value) entries the writing kernel will emit, counting duplicates, since warp.sparse.bsr_from_triplets sums entries that land on the same position.

 required
dtype type

Element type of the value buffer. A scalar (wp.float32 / wp.float64) for a 1x1-block matrix, or a matrix type (wp.mat22d) for a block matrix.

 required
device DeviceLike

Warp device for all three buffers.

 required

Returns:

Type Description
(rows, cols, values)

Three length-n_triplets arrays on device. The index buffers are wp.int32; values takes dtype. All three are uninitialized -- the caller's kernel is expected to write every entry.
Notes

wp.empty rather than wp.zeros deliberately: a triplet writer fills all three buffers, so zeroing them first would be three wasted launches. A kernel that emits fewer than n_triplets entries must therefore write an explicit structural zero (typically a self-entry) rather than leave a slot untouched, which is also why every operator build in this package passes prune_numerical_zeros=False -- see the note on index_sparse, the one caller that prunes.

See Also

index_sparse

Source code in triwarp/array.py 
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
def triplet_buffers(
    n_triplets: int, dtype: type, device: wp.DeviceLike
) -> tuple[wp.array[wp.int32], wp.array[wp.int32], wp.array]:
    """
    Uninitialized ``(rows, cols, values)`` COO buffers for one ``bsr_from_triplets`` build.

    Parameters
    ----------
    n_triplets
        Length of each of the three buffers: the number of ``(row, col, value)`` entries the
        writing kernel will emit, counting duplicates, since ``warp.sparse.bsr_from_triplets``
        sums entries that land on the same position.
    dtype
        Element type of the value buffer. A scalar (``wp.float32`` / ``wp.float64``) for a
        1x1-block matrix, or a matrix type (``wp.mat22d``) for a block matrix.
    device
        Warp device for all three buffers.

    Returns
    -------
    rows, cols, values
        Three length-``n_triplets`` arrays on ``device``. The index buffers are ``wp.int32``;
        ``values`` takes ``dtype``. All three are **uninitialized** -- the caller's kernel is
        expected to write every entry.

    Notes
    -----
    ``wp.empty`` rather than ``wp.zeros`` deliberately: a triplet writer fills all three buffers,
    so zeroing them first would be three wasted launches. A kernel that emits *fewer* than
    ``n_triplets`` entries must therefore write an explicit structural zero (typically a
    self-entry) rather than leave a slot untouched, which is also why every operator build in
    this package passes ``prune_numerical_zeros=False`` -- see the note on
    [`index_sparse`][triwarp.array.index_sparse], the one caller that prunes.

    See Also
    --------
    [`index_sparse`][triwarp.array.index_sparse]
    """
    rows = wp.empty(n_triplets, dtype=wp.int32, device=device)
    cols = wp.empty(n_triplets, dtype=wp.int32, device=device)
    values = wp.empty(n_triplets, dtype=dtype, device=device)
    return rows, cols, values

empty_square_bsr 

empty_square_bsr(
    n_rows: int, dtype: type, device: DeviceLike
) -> wps.BsrMatrix[wp.float32]

Zero-nnz (n_rows, n_rows) operator, the n_faces == 0 return several assemblers share.

A square BSR matrix with no faces to build triplets from is still a valid, correctly-shaped operator -- just an empty one -- so laplacian.cotmatrix, laplacian.connection_laplacian, laplacian.graph_laplacian and energies's assembly wrappers all construct this exact matrix for their empty-mesh early return, and did so as four near-identical inline copies before this was factored out.

Parameters:

Name Type Description Default
n_rows int

Row and column count of the square matrix.

 required
dtype type

Element type of the (empty) value buffer -- a scalar for a 1x1-block matrix, or a matrix type (e.g. wp.mat22d) for a block matrix, exactly as triplet_buffers takes it.

 required
device DeviceLike

Warp device for the matrix.

 required

Returns:

Type Description
BsrMatrix

Square (n_rows, n_rows) matrix with zero stored entries, on device.
See Also

triplet_buffers

Source code in triwarp/array.py 
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
def empty_square_bsr(n_rows: int, dtype: type, device: wp.DeviceLike) -> wps.BsrMatrix[wp.float32]:
    """
    Zero-nnz ``(n_rows, n_rows)`` operator, the ``n_faces == 0`` return several assemblers share.

    A square BSR matrix with no faces to build triplets from is still a valid, correctly-shaped
    operator -- just an empty one -- so [`laplacian.cotmatrix`][triwarp.laplacian.cotmatrix],
    [`laplacian.connection_laplacian`][triwarp.laplacian.connection_laplacian],
    [`laplacian.graph_laplacian`][triwarp.laplacian.graph_laplacian] and
    [`energies`][triwarp.energies]'s assembly wrappers all construct this exact matrix for their
    empty-mesh early return, and did so as four near-identical inline copies before this was
    factored out.

    Parameters
    ----------
    n_rows
        Row and column count of the square matrix.
    dtype
        Element type of the (empty) value buffer -- a scalar for a 1x1-block matrix, or a matrix
        type (e.g. ``wp.mat22d``) for a block matrix, exactly as
        [`triplet_buffers`][triwarp.array.triplet_buffers] takes it.
    device
        Warp device for the matrix.

    Returns
    -------
    warp.sparse.BsrMatrix
        Square ``(n_rows, n_rows)`` matrix with zero stored entries, on ``device``.

    See Also
    --------
    [`triplet_buffers`][triwarp.array.triplet_buffers]
    """
    return wps.bsr_from_triplets(
        n_rows,
        n_rows,
        wp.empty(0, dtype=wp.int32, device=device),
        wp.empty(0, dtype=wp.int32, device=device),
        wp.empty(0, dtype=dtype, device=device),
        prune_numerical_zeros=False,
    )

index_sparse 

index_sparse(
    n_rows: int,
    indices: Array2dInt32,
    data: array[Scalar] | None = None,
    dtype: type[Scalar] | None = None,
    *,
    prune_numerical_zeros: bool = True,
) -> wps.BsrMatrix[wp.Scalar]

Build a sparse row/column incidence matrix from flat index columns.

This mirrors trimesh.geometry.index_sparse, but returns a warp.sparse.BsrMatrix in 1x1 BSR (CSR) form instead of scipy.sparse.coo_matrix.

Parameters:

Name Type Description Default
n_rows int

Number of matrix rows (e.g. vertex count). Matrix shape is (n_rows, len(indices)).

 required
indices Array2dInt32

Integer array of shape (m, d) — typically mesh.faces with three vertex indices per face.

 required
data array[Scalar] | None

Optional 1-D array of length m * d. If omitted, wp.ones is used; see dtype.

 None
dtype type[Scalar] | None

Scalar type for wp.ones when data is None (defaults to wp.float32 if dtype is None). When data and dtype are provided, the values of the matrix are cast to dtype.

 None
prune_numerical_zeros bool

Forwarded to warp.sparse.bsr_from_triplets.

 True

Returns:

Type Description
BsrMatrix

Sparse matrix with shape (n_rows, len(indices)) and 1x1 blocks.

Raises:

Type Description
ValueError

If data is given and its size differs from indices.size.
RuntimeError

If indices and data are not all on one device.
Source code in triwarp/array.py 
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
def index_sparse(
    n_rows: int,
    indices: twt.Array2dInt32,
    data: wp.array[wp.Scalar] | None = None,
    dtype: type[wp.Scalar] | None = None,
    *,
    prune_numerical_zeros: bool = True,
) -> wps.BsrMatrix[wp.Scalar]:
    """
    Build a sparse row/column incidence matrix from flat index columns.

    This mirrors ``trimesh.geometry.index_sparse``, but returns a ``warp.sparse.BsrMatrix``
    in 1x1 BSR (CSR) form instead of ``scipy.sparse.coo_matrix``.

    Parameters
    ----------
    n_rows
        Number of matrix rows (e.g. vertex count). Matrix shape is ``(n_rows, len(indices))``.
    indices
        Integer array of shape ``(m, d)`` — typically ``mesh.faces`` with three vertex indices
        per face.
    data
        Optional 1-D array of length ``m * d``. If omitted, ``wp.ones`` is used; see ``dtype``.
    dtype
        Scalar type for ``wp.ones`` when ``data`` is ``None`` (defaults to ``wp.float32`` if
        ``dtype`` is ``None``). When ``data`` and ``dtype`` are provided, the values of the
        matrix are cast to ``dtype``.
    prune_numerical_zeros
        Forwarded to ``warp.sparse.bsr_from_triplets``.

    Returns
    -------
    warp.sparse.BsrMatrix
        Sparse matrix with shape ``(n_rows, len(indices))`` and 1x1 blocks.

    Raises
    ------
    ValueError
        If ``data`` is given and its size differs from ``indices.size``.
    RuntimeError
        If ``indices`` and ``data`` are not all on one device.
    """
    require_same_device(indices=indices, data=data)
    prune_numerical_zeros = prune_numerical_zeros and data is not None
    if data is None:
        data = wp.ones(
            indices.size, dtype=dtype if dtype is not None else wp.float32, device=indices.device
        )
    else:
        if data.size != indices.size:
            raise ValueError(
                f"data must have the same size as indices, got {data.size} and {indices.size}"
            )
        if dtype is not None and data.dtype != dtype:
            data = astype(data, dtype)

    n_cols, n_repeats = indices.shape
    cols = arange_repeat(n_cols * n_repeats, n_repeats, indices.device)
    return wps.bsr_from_triplets(
        n_rows,
        indices.shape[0],
        indices.flatten(),
        cols,
        data,
        prune_numerical_zeros=prune_numerical_zeros,
    )

isin 

isin(
    elements: ArrayNd,
    test_elements: array[Int],
    *,
    max_index: int | None = None,
) -> wp.array[wp.bool]

Test whether each element appears in test_elements (numpy.isin for integers).

Works for every Warp integer dtype -- int8 through int64, uint8 through uint64 -- and for negative values. Both arrays must share one dtype.

Two strategies, chosen by the value span max - min + 1 taken over both inputs together. When the span is modest relative to len(test_elements), membership is a boolean lookup table indexed by value - min (fast for dense mesh indices). Otherwise test_elements is sorted and each query is a binary search, which keeps memory bounded when the values are sparse in their dtype.

Parameters:

Name Type Description Default
elements ArrayNd

Integer array of any rank on the target device. Membership is a per-element predicate, so the array is flattened, tested, and the result reshaped back; nothing in the two strategies looks at the shape.

 required
test_elements array[Int]

1D array of values to test membership against, of the same dtype as elements.

 required
max_index int | None

Optional exclusive upper bound on the values of both arrays, which must then be non-negative -- the same escape hatch, and the same shape of it, as hash_indices_rows' max_index. Supplying it pins the strategy to the direct-index table and skips the two triwarp.reduce.minmax reductions that would otherwise infer the span, and with them two host readbacks that serialise the device pipeline. Must be positive. Pass it wherever the bound is structural, as it is for any buffer of mesh vertex indices.

 None

Returns:

Type Description
array[bool]

Boolean array with the same shape as elements. All False when either input is empty.

Raises:

Type Description
TypeError

If either array is not an integer dtype, or the two dtypes differ.
ValueError

If max_index is not positive.
Warnings

A max_index smaller than the true maximum is a wrong answer, not a raise: a value at or above it reads as absent, however far above it lies and on either side of the comparison. Unlike hash_indices_rows' bound it is at least memory-safe, because the table lookup range-guards each slot.

See Also

indices_to_mask sortable_dtype numpy.isin

Notes

Dtypes narrower than four bytes are widened to int32 / uint32 (the sortable_dtype rule) before either strategy runs: Warp's radix sort does not accept them, and neither does the tiled min/max reduction the span needs.

There is no assume_unique, and it is worth saying why rather than leaving its absence to look like an omission: numpy.isin gains from one because the numpy.in1d sort path behind it calls numpy.unique on both arrays first, and neither strategy here dedups anything -- the table is an idempotent scatter and the binary search is over the sorted keys as given. Duplicates on either side are already free. What the two strategies do pay for is inferring the value span, which is why the guarantee this takes is a bound rather than a uniqueness claim.

Without max_index: two host readbacks, one min/max reduction per input, which is what selects the strategy and anchors the table.

Source code in triwarp/array.py 
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
def isin(
    elements: twt.ArrayNd, test_elements: wp.array[wp.Int], *, max_index: int | None = None
) -> wp.array[wp.bool]:
    """
    Test whether each element appears in ``test_elements`` (``numpy.isin`` for integers).

    Works for every Warp integer dtype -- ``int8`` through ``int64``, ``uint8`` through ``uint64``
    -- and for negative values. Both arrays must share one dtype.

    Two strategies, chosen by the value **span** ``max - min + 1`` taken over both inputs together.
    When the span is modest relative to ``len(test_elements)``, membership is a boolean lookup table
    indexed by ``value - min`` (fast for dense mesh indices). Otherwise ``test_elements`` is sorted
    and each query is a binary search, which keeps memory bounded when the values are sparse in
    their dtype.

    Parameters
    ----------
    elements
        Integer array of **any rank** on the target device. Membership is a per-element predicate,
        so the array is flattened, tested, and the result reshaped back; nothing in the two
        strategies looks at the shape.
    test_elements
        1D array of values to test membership against, of the same dtype as ``elements``.
    max_index
        Optional exclusive upper bound on the values of **both** arrays, which must then be
        non-negative -- the same escape hatch, and the same shape of it, as
        [`hash_indices_rows`][triwarp.grouping.hash_indices_rows]' ``max_index``. Supplying it
        pins the strategy to the direct-index table and skips the two ``triwarp.reduce.minmax``
        reductions that would otherwise infer the span, and with them two host readbacks that
        serialise the device pipeline. Must be positive. Pass it wherever the bound is structural,
        as it is for any buffer of mesh vertex indices.

    Returns
    -------
    wp.array[wp.bool]
        Boolean array with the same shape as ``elements``. All ``False`` when either
        input is empty.

    Raises
    ------
    TypeError
        If either array is not an integer dtype, or the two dtypes differ.
    ValueError
        If ``max_index`` is not positive.

    Warnings
    --------
    A ``max_index`` smaller than the true maximum is a wrong answer, not a raise: a value at or
    above it reads as absent, however far above it lies and on either side of the comparison.
    Unlike [`hash_indices_rows`][triwarp.grouping.hash_indices_rows]' bound it is at least
    memory-safe, because the table lookup range-guards each slot.

    See Also
    --------
    [`indices_to_mask`][triwarp.array.indices_to_mask]
    [`sortable_dtype`][triwarp.typing.sortable_dtype]
    [`numpy.isin`][]

    Notes
    -----
    Dtypes narrower than four bytes are widened to ``int32`` / ``uint32`` (the
    [`sortable_dtype`][triwarp.typing.sortable_dtype] rule) before either strategy runs: Warp's
    radix sort does not accept them, and neither does the tiled min/max reduction the span needs.

    There is no ``assume_unique``, and it is worth saying why rather than leaving its absence to
    look like an omission: [`numpy.isin`][] gains from one because the ``numpy.in1d`` sort path
    behind it calls [`numpy.unique`][] on both arrays first, and neither strategy here dedups
    anything -- the table is an idempotent scatter and the binary search is over the sorted keys as
    given. Duplicates on either side are already free. What the two
    strategies *do* pay for is inferring the value span, which is why the guarantee this takes is a
    bound rather than a uniqueness claim.

    Without ``max_index``: two host readbacks, one min/max reduction per input, which is what
    selects the strategy and anchors the table.
    """
    device = elements.device
    dtype = elements.dtype
    if dtype != test_elements.dtype:
        raise TypeError(
            f"isin requires one dtype for both arrays, got {dtype} and {test_elements.dtype}"
        )
    if not wp.types.type_is_int(dtype) or dtype == wp.bool:
        raise TypeError(f"isin requires an integer dtype, got {dtype}")

    if max_index is not None and max_index <= 0:
        raise ValueError(f"max_index must be positive, got {max_index}")

    k = int(test_elements.shape[0])
    if k == 0 or int(elements.size) == 0:
        return wp.zeros(elements.shape, dtype=wp.bool, device=device)

    is_flat = int(elements.ndim) == 1
    elements_flat = elements if is_flat else elements.flatten()
    # Widen sub-32-bit dtypes once, up front: neither ``reduce.minmax`` nor the radix sort accepts
    # them, and a widened span cannot overflow the type it is measured in (int8's span reaches 256).
    if wp.types.type_size_in_bytes(dtype) < 4:
        wide = twt.sortable_dtype(dtype)
        elements_flat = astype(elements_flat, wide)
        test_elements = astype(test_elements, wide)

    if max_index is not None:
        # A supplied bound is the whole point of the keyword: it is what the two reductions below
        # would have inferred, so it also settles the strategy -- a caller who knows the values are
        # dense mesh indices is describing exactly the table's best case, and honouring a
        # ``_ISIN_MASK_SIZE_FACTOR`` test here would spend the readback to reach the same branch.
        return _reshaped(_isin_lookup_mask(elements_flat, test_elements, max_index, 0), elements)

    lo_elements, hi_elements = tw.reduce.minmax(elements_flat)
    lo_test, hi_test = tw.reduce.minmax(test_elements)
    offset = min(int(lo_elements), int(lo_test))
    span = max(int(hi_elements), int(hi_test)) - offset + 1
    if span <= _ISIN_MASK_SIZE_FACTOR * k:
        out_flat = _isin_lookup_mask(elements_flat, test_elements, span, offset)
    else:
        out_flat = _isin_lookup_sorted(elements_flat, test_elements)

    return _reshaped(out_flat, elements)

flatnonzero 

flatnonzero(
    values: array[bool] | array[Scalar],
) -> wp.array[wp.int32]

Return the indices of the non-zero entries of a 1D array (numpy.flatnonzero).

Takes a boolean mask, which is the common case, or any scalar array — every non-zero value selects its index, exactly as numpy.flatnonzero does, so -2 and 3 both count and only 0 does not.

Parameters:

Name Type Description Default
values array[bool] | array[Scalar]

Length-n wp.bool mask, or a wp.int32 / float / other scalar array, on the target device.

 required

Returns:

Type Description
array[int32]

Selected indices on values.device, ascending. Empty when nothing is non-zero.

Raises:

Type Description
ValueError

If values is not rank-1.
See Also

indices_to_mask mask_to_compact_ranks numpy.flatnonzero

Source code in triwarp/array.py 
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
def flatnonzero(values: wp.array[wp.bool] | wp.array[wp.Scalar]) -> wp.array[wp.int32]:
    """
    Return the indices of the non-zero entries of a 1D array (``numpy.flatnonzero``).

    Takes a boolean mask, which is the common case, or any scalar array — every non-zero value
    selects its index, exactly as ``numpy.flatnonzero`` does, so ``-2`` and ``3`` both count and
    only ``0`` does not.

    Parameters
    ----------
    values
        Length-``n`` ``wp.bool`` mask, or a ``wp.int32`` / float / other scalar array, on the
        target device.

    Returns
    -------
    wp.array[wp.int32]
        Selected indices on ``values.device``, ascending. Empty when nothing is non-zero.

    Raises
    ------
    ValueError
        If ``values`` is not rank-1.

    See Also
    --------
    [`indices_to_mask`][triwarp.array.indices_to_mask]
    [`mask_to_compact_ranks`][triwarp.array.mask_to_compact_ranks]
    [`numpy.flatnonzero`][]
    """
    if int(values.ndim) != 1:
        raise ValueError(f"flatnonzero requires a 1D array, got ndim={values.ndim}")

    device = values.device
    n = int(values.shape[0])
    if n == 0:
        return wp.empty(0, dtype=wp.int32, device=device)

    flags = wp.empty(n, dtype=wp.int32, device=device)
    if values.dtype == wp.bool:
        # ``array_cast`` already yields exactly 0/1 from a mask, and ``wp.Scalar`` does not
        # instantiate for ``wp.bool`` anyway. For every other dtype the cast would copy the
        # *values*, and the scan below would then sum them instead of counting them.
        wp.utils.array_cast(values, flags)
    else:
        wp.map(kernel_array.nonzero_flag, values, out=flags)

    # Inclusive scan: the total is its last element, so one 4-byte tail read sizes the output
    # (the scatter kernel derives each exclusive position as inclusive[i] - 1).
    inclusive = wp.empty(n, dtype=wp.int32, device=device)
    wp.utils.array_scan(flags, out_array=inclusive, inclusive=True)
    n_out = int(read_scalar(inclusive))

    if n_out == 0:
        return wp.empty(0, dtype=wp.int32, device=device)

    out_indices = wp.empty(n_out, dtype=wp.int32, device=device)
    wp.launch(
        kernel_scatter.scatter_index_where,
        dim=n,
        inputs=[flags, inclusive, out_indices],
        device=device,
    )
    return out_indices

gather 

gather(
    src: array[DType] | ArrayNd, indices: array[int32]
) -> wp.array[DType] | twt.ArrayNd

Dense copy of src gathered along its first axis by indices (numpy.take).

Warp's src[indices] fancy indexing yields a warp.indexedarray view; this materializes a contiguous warp.array (performing the copy) so callers get a real array supporting .reshape and a stable return type. Works for rank-1 sources (src[indices]) and rank-2 row gather (src[indices, :]), with any scalar or vector dtype.

Only the first axis is indexable, which is what every caller in this package needs. To take a column, materialize it with wp.clone(src[:, k]) — a column view is strided, and Warp's fancy indexing silently ignores the stride of an index array.

Parameters:

Name Type Description Default
src array[DType] | ArrayNd

Rank-1 or rank-2 wp.array on the target device.

 required
indices array[int32]

1D wp.int32 array of indices into the first axis of src. Must be contiguous; see the warning above.

 required

Returns:

Type Description
array

Contiguous gathered copy on src.device with shape (len(indices), *src.shape[1:]) and the same dtype as src. Empty along the first axis when indices is empty.

Raises:

Type Description
RuntimeError

If src and indices are not all on one device.
See Also

index_sparse remap_indices The sentinel-preserving variant for index buffers that may carry -1 entries.

Source code in triwarp/array.py 
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
def gather(
    src: wp.array[DType] | twt.ArrayNd, indices: wp.array[wp.int32]
) -> wp.array[DType] | twt.ArrayNd:
    """
    Dense copy of ``src`` gathered along its first axis by ``indices`` (``numpy.take``).

    Warp's ``src[indices]`` fancy indexing yields a ``warp.indexedarray`` view; this
    materializes a contiguous ``warp.array`` (performing the copy) so callers get a real
    array supporting ``.reshape`` and a stable return type. Works for rank-1 sources
    (``src[indices]``) and rank-2 row gather (``src[indices, :]``), with any scalar or vector
    ``dtype``.

    Only the first axis is indexable, which is what every caller in this package needs. To take a
    *column*, materialize it with ``wp.clone(src[:, k])`` — a column view is strided, and Warp's
    fancy indexing silently ignores the stride of an index array.

    Parameters
    ----------
    src
        Rank-1 or rank-2 ``wp.array`` on the target device.
    indices
        1D ``wp.int32`` array of indices into the first axis of ``src``. **Must be contiguous**;
        see the warning above.

    Returns
    -------
    wp.array
        Contiguous gathered copy on ``src.device`` with shape
        ``(len(indices), *src.shape[1:])`` and the same ``dtype`` as ``src``. Empty along the
        first axis when ``indices`` is empty.

    Raises
    ------
    RuntimeError
        If ``src`` and ``indices`` are not all on one device.

    See Also
    --------
    [`index_sparse`][triwarp.array.index_sparse]
    [`remap_indices`][triwarp.array.remap_indices]
        The sentinel-preserving variant for index buffers that may carry ``-1`` entries.
    """
    require_same_device(src=src, indices=indices)
    k = int(indices.shape[0])
    out_shape = (k, *(int(dim) for dim in src.shape[1:]))
    out = wp.empty(out_shape, dtype=src.dtype, device=src.device)
    if k > 0:
        wp.copy(out, src[indices])
    return out

astype 

astype(values: ArrayNd, dtype: type) -> wp.array

Element-wise dtype conversion, shape and rank preserved (numpy.ndarray.astype).

The Python-scope counterpart of wp.cast, which exists only inside a kernel. Allocates a buffer of values' shape on values' device and fills it with warp.utils.array_cast -- the pair this replaces at twenty-odd call sites.

Parameters:

Name Type Description Default
values ArrayNd

Rank-1 or rank-2 Warp array of any scalar dtype array_cast accepts.

 required
dtype type

Target scalar dtype.

 required

Returns:

Type Description
array

A new array of values' shape on values' device, with element type dtype.

Raises:

Type Description
ValueError

If values is rank-2 and not contiguous, since the rank-2 path flattens.
Notes

warp.utils.array_cast's kernel is dest[i] = dest.dtype(src[i]), which is a scalar conversion -- handed a rank-2 array it fails to compile, because src[i] is a row. So rank-2 input is cast through paired flatten() views here rather than at each call site, which is what several of them were doing by hand.

The output keeps values' shape, so this is not a reinterpretation that changes rank: an (n, 3) float32 read as (n,) wp.vec3 is a different operation and still calls wp.utils.array_cast directly, as hash_rows and mean_vertex_normals do.

See Also

bitcast_to_int Reinterpret the bits rather than convert the value.

Source code in triwarp/array.py 
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
def astype(values: twt.ArrayNd, dtype: type) -> wp.array:
    """
    Element-wise dtype conversion, shape and rank preserved (``numpy.ndarray.astype``).

    The Python-scope counterpart of ``wp.cast``, which exists only inside a kernel. Allocates a
    buffer of ``values``' shape on ``values``' device and fills it with
    ``warp.utils.array_cast`` -- the pair this replaces at twenty-odd call sites.

    Parameters
    ----------
    values
        Rank-1 or rank-2 Warp array of any scalar dtype ``array_cast`` accepts.
    dtype
        Target scalar dtype.

    Returns
    -------
    wp.array
        A new array of ``values``' shape on ``values``' device, with element type ``dtype``.

    Raises
    ------
    ValueError
        If ``values`` is rank-2 and not contiguous, since the rank-2 path flattens.

    Notes
    -----
    ``warp.utils.array_cast``'s kernel is ``dest[i] = dest.dtype(src[i])``, which is a *scalar*
    conversion -- handed a rank-2 array it fails to compile, because ``src[i]`` is a row. So rank-2
    input is cast through paired ``flatten()`` views here rather than at each call site, which is
    what several of them were doing by hand.

    The output keeps ``values``' shape, so this is not a reinterpretation that changes rank: an
    ``(n, 3)`` ``float32`` read as ``(n,)`` ``wp.vec3`` is a different operation and still calls
    ``wp.utils.array_cast`` directly, as [`hash_rows`][triwarp.grouping.hash_rows] and
    [`mean_vertex_normals`][triwarp.vertices.mean_vertex_normals] do.

    See Also
    --------
    [`bitcast_to_int`][triwarp.array.bitcast_to_int]
        Reinterpret the *bits* rather than convert the value.
    """
    out = wp.empty(values.shape, dtype=dtype, device=values.device)
    if int(values.ndim) == 1:
        wp.utils.array_cast(values, out)
    else:
        if not values.is_contiguous:
            raise ValueError("astype requires a contiguous array for rank-2 input")
        wp.utils.array_cast(values.flatten(), out.flatten())
    return out

index_bound 

index_bound(indices: IntArray) -> int

Exclusive upper bound on an index buffer's values, as max(indices) + 1.

For a face or edge buffer this is the vertex count, and it follows libigl's F.maxCoeff() + 1 convention: one past the largest referenced index, so a mesh with trailing unreferenced vertices reports fewer than it has. Accepts any wp.int32 index buffer of any shape -- a length-3 * n_faces flat triangle buffer, an (n, 2) edge array -- and reads only the 4-byte maximum back to the host.

Named for the index buffer rather than for vertices because that is all it sees: it is index arithmetic, and nothing about it is geometric. Named a bound rather than a size because that is what a consumer wants it for -- to allocate a table the indices address, or to hand hash_indices_rows its radix -- and because a mesh with trailing unreferenced vertices has more of them than this reports.

Parameters:

Name Type Description Default
indices IntArray

A wp.int32 index buffer of any shape (e.g. a flat faces array or a (n, 2) edge array).

 required

Returns:

Type Description
int

max(indices) + 1, or 0 when indices is empty.
Source code in triwarp/array.py 
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
def index_bound(indices: twt.IntArray) -> int:
    """
    Exclusive upper bound on an index buffer's values, as ``max(indices) + 1``.

    For a face or edge buffer this is the vertex count, and it follows libigl's
    ``F.maxCoeff() + 1`` convention: one past the largest referenced index, so a mesh with trailing
    unreferenced vertices reports fewer than it has. Accepts any ``wp.int32`` index buffer of any
    shape -- a length-``3 * n_faces`` flat triangle buffer, an ``(n, 2)`` edge array -- and reads
    only the 4-byte maximum back to the host.

    Named for the index buffer rather than for vertices because that is all it sees: it is
    index arithmetic, and nothing about it is geometric. Named a *bound* rather than a size
    because that is what a consumer wants it for -- to allocate a table the indices address, or to
    hand [`hash_indices_rows`][triwarp.grouping.hash_indices_rows] its radix -- and because a mesh
    with trailing unreferenced vertices has more of them than this reports.

    Parameters
    ----------
    indices
        A ``wp.int32`` index buffer of any shape (e.g. a flat ``faces`` array or a ``(n, 2)``
        edge array).

    Returns
    -------
    int
        ``max(indices) + 1``, or ``0`` when ``indices`` is empty.
    """
    if int(indices.size) == 0:
        return 0
    # Device-side tiled max: only the 4-byte result crosses to the host, not the whole buffer.
    return int(tw.reduce.max(indices)) + 1

indices_to_mask 

indices_to_mask(
    indices: array[int32],
    n: int,
    *,
    device: DeviceLike = None,
) -> wp.array[wp.bool]

Boolean membership mask of length n marking each value in indices as True.

Wraps the mark_membership_mask scatter kernel: every indices[i] sets out_mask[indices[i]] = True. The inverse of flatnonzero.

Parameters:

Name Type Description Default
indices array[int32]

1D wp.int32 array of values in [0, n) to mark. May be empty.

 required
n int

Length of the returned mask.

 required
device DeviceLike

Target Warp device. Defaults to indices.device.

 None

Returns:

Type Description
array[bool]

Length-n mask, all False except at positions named by indices.
See Also

flatnonzero isin

Source code in triwarp/array.py 
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
def indices_to_mask(
    indices: wp.array[wp.int32], n: int, *, device: wp.DeviceLike = None
) -> wp.array[wp.bool]:
    """
    Boolean membership mask of length ``n`` marking each value in ``indices`` as ``True``.

    Wraps the ``mark_membership_mask`` scatter kernel: every ``indices[i]`` sets
    ``out_mask[indices[i]] = True``. The inverse of [`flatnonzero`][triwarp.array.flatnonzero].

    Parameters
    ----------
    indices
        1D ``wp.int32`` array of values in ``[0, n)`` to mark. May be empty.
    n
        Length of the returned mask.
    device
        Target Warp device. Defaults to ``indices.device``.

    Returns
    -------
    wp.array[wp.bool]
        Length-``n`` mask, all ``False`` except at positions named by ``indices``.

    See Also
    --------
    [`flatnonzero`][triwarp.array.flatnonzero]
    [`isin`][triwarp.array.isin]
    """
    device = device if device is not None else indices.device
    mask = wp.zeros(n, dtype=wp.bool, device=device)
    k = int(indices.shape[0])
    if k > 0:
        wp.launch(
            kernel_scatter.mark_membership_mask,
            dim=k,
            inputs=[indices, wp.int32(n), mask],
            device=device,
        )
    return mask

mask_to_compact_ranks 

mask_to_compact_ranks(
    mask: array[bool], *, invert: bool = False
) -> tuple[wp.array[wp.int32], int]

Compact index map over the True entries of a boolean mask, plus their count.

Not flatnonzero, though the two come from the same scan: flatnonzero returns which elements are selected (length count, values are positions into mask), while this returns where each element lands in the compacted numbering (length n, values are compact ranks). Renumbering consumers — the free/fixed degree-of-freedom partitions in triwarp.smoothing and triwarp.linalg — need this full-length scatter-side form, which flatnonzero cannot provide without an extra pass.

Parameters:

Name Type Description Default
mask array[bool]

Length-n wp.bool array.

 required
invert bool

When True, map the False entries instead. Useful for a free/fixed degree-of-freedom partition, where the mask marks the constrained entries and the compact map is wanted over the unconstrained complement (see free_partition).

 False

Returns:

Name Type Description
compact_ranks array[int32]

Length-n array on mask.device: an exclusive scan of the (optionally inverted) mask, so index_map[i] is the compact 0-based rank of element i among the selected entries at or before it (meaningful only where element i is itself selected).
count int

Total number of selected entries in mask.
Source code in triwarp/array.py 
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
def mask_to_compact_ranks(
    mask: wp.array[wp.bool], *, invert: bool = False
) -> tuple[wp.array[wp.int32], int]:
    """
    Compact index map over the ``True`` entries of a boolean mask, plus their count.

    Not [`flatnonzero`][triwarp.array.flatnonzero], though the two come from the same scan:
    ``flatnonzero`` returns *which* elements are selected (length ``count``, values are positions
    into ``mask``), while this returns *where each element lands* in the compacted numbering
    (length ``n``, values are compact ranks). Renumbering consumers — the free/fixed
    degree-of-freedom partitions in [`triwarp.smoothing`][triwarp.smoothing] and
    [`triwarp.linalg`][triwarp.linalg] — need this full-length scatter-side form, which
    ``flatnonzero`` cannot provide without an extra pass.

    Parameters
    ----------
    mask
        Length-``n`` ``wp.bool`` array.
    invert
        When ``True``, map the ``False`` entries instead. Useful for a free/fixed degree-of-freedom
        partition, where the mask marks the *constrained* entries and the compact map is wanted over
        the unconstrained complement (see [`free_partition`][triwarp.linalg.free_partition]).

    Returns
    -------
    compact_ranks : wp.array[wp.int32]
        Length-``n`` array on ``mask.device``: an exclusive scan of the (optionally inverted) mask,
        so ``index_map[i]`` is the compact 0-based rank of element ``i`` among the selected
        entries at or before it (meaningful only where element ``i`` is itself selected).
    count : int
        Total number of selected entries in ``mask``.
    """
    device = mask.device
    n = int(mask.shape[0])
    if n == 0:
        return wp.zeros(0, dtype=wp.int32, device=device), 0
    flags = wp.empty(n, dtype=wp.int32, device=device)
    if invert:
        wp.map(kernel_array.complement_flag, mask, out=flags)
    else:
        wp.utils.array_cast(mask, flags)
    return counts_to_offsets(flags)

counts_to_offsets 

counts_to_offsets(
    counts: array[int32], *, include_total: bool = False
) -> tuple[wp.array[wp.int32], int]

Exclusive prefix sum of counts, plus their total.

The CSR-building step that turns per-element counts into row starts. Done in one scan pass and one 4-byte host read: the scan runs inclusive into the tail of an n + 1 buffer whose leading zero is already in place, which makes the first n entries the exclusive sum and the last entry the total. The obvious spelling — one exclusive scan for the offsets and a second inclusive scan (or a reduce.sum) for the total — costs a second full pass over counts, and reading the total as inclusive.numpy()[-1] copies the whole array to the host to look at one element of it.

Parameters:

Name Type Description Default
counts array[int32]

Length-n wp.int32 per-element counts.

 required
include_total bool

Return the length-n + 1 CSR form, whose trailing element is total, instead of the length-n form. Free — that buffer is what gets built either way — and it is what warp.utils.segmented_sort_pairs and the other segment-bounds consumers want.

 False

Returns:

Name Type Description
offsets array[int32]

Exclusive prefix sum: length n by default, or n + 1 with offsets[n] == total when include_total is set. The default is a view into the n + 1 buffer, which the returned array keeps alive. Element i owns [offsets[i], offsets[i] + counts[i]).
total int

Sum of counts.
Notes

Two offsets conventions coexist in this package: the length-n form, with the total implicit, and the length-n + 1 form that stores it (halfedge.vertex_one_rings, geodesic_walk.trace_from_vertex, and every segmented_sort_pairs caller). Both come out of here, so no caller has to append the terminator afterwards.

This is for callers that want total, which it reads back unconditionally, and a host readback serialises the device pipeline. A caller that only needs the offsets and already knows its buffer size should keep the open-coded wp.zeros(n + 1) plus a scan into [1:], as halfedge.vertex_one_rings and adjacency.vertex_face_adjacency do: both size their payload from 3 * n_faces and would gain a synchronization they currently do not have.

See Also

flatnonzero mask_to_compact_ranks

Source code in triwarp/array.py 
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
def counts_to_offsets(
    counts: wp.array[wp.int32], *, include_total: bool = False
) -> tuple[wp.array[wp.int32], int]:
    """
    Exclusive prefix sum of ``counts``, plus their total.

    The CSR-building step that turns per-element counts into row starts. Done in **one** scan pass
    and one 4-byte host read: the scan runs *inclusive* into the tail of an ``n + 1`` buffer whose
    leading zero is already in place, which makes the first ``n`` entries the exclusive sum and the
    last entry the total. The obvious spelling — one exclusive scan for the offsets and a second
    inclusive scan (or a ``reduce.sum``) for the total — costs a second full pass over ``counts``,
    and reading the total as ``inclusive.numpy()[-1]`` copies the whole array to the host to look at
    one element of it.

    Parameters
    ----------
    counts
        Length-``n`` ``wp.int32`` per-element counts.
    include_total
        Return the length-``n + 1`` CSR form, whose trailing element is ``total``, instead of the
        length-``n`` form. Free — that buffer is what gets built either way — and it is what
        ``warp.utils.segmented_sort_pairs`` and the other segment-bounds consumers want.

    Returns
    -------
    offsets : wp.array[wp.int32]
        Exclusive prefix sum: length ``n`` by default, or ``n + 1`` with ``offsets[n] == total``
        when ``include_total`` is set. The default is a **view** into the ``n + 1`` buffer, which
        the returned array keeps alive. Element ``i`` owns ``[offsets[i], offsets[i] + counts[i])``.
    total : int
        Sum of ``counts``.

    Notes
    -----
    Two offsets conventions coexist in this package: the length-``n`` form, with the total
    implicit, and the length-``n + 1`` form that stores it (``halfedge.vertex_one_rings``,
    ``geodesic_walk.trace_from_vertex``, and every ``segmented_sort_pairs`` caller). Both come
    out of here, so no caller has to append the terminator afterwards.

    **This is for callers that want ``total``**, which it reads back unconditionally, and a host
    readback serialises the device pipeline. A caller that only needs the offsets and already knows
    its buffer size should keep the open-coded ``wp.zeros(n + 1)`` plus a scan into ``[1:]``, as
    ``halfedge.vertex_one_rings`` and ``adjacency.vertex_face_adjacency`` do: both size their
    payload from ``3 * n_faces`` and would gain a synchronization they currently do not have.

    See Also
    --------
    [`flatnonzero`][triwarp.array.flatnonzero]
    [`mask_to_compact_ranks`][triwarp.array.mask_to_compact_ranks]
    """
    # The unconditional readback of ``total`` is a host synchronization -- which is why the Notes
    # section advises open-coding the scan when only the offsets are wanted.
    n = int(counts.shape[0])
    device = counts.device
    if n == 0:
        return wp.zeros(1 if include_total else 0, dtype=wp.int32, device=device), 0
    # The leading zero from ``wp.zeros`` is the first exclusive offset; the inclusive scan fills the
    # rest, so ``buffer[n]`` is the total and ``buffer[:n]`` the exclusive offsets.
    buffer = wp.zeros(n + 1, dtype=wp.int32, device=device)
    wp.utils.array_scan(counts, out_array=buffer[1:], inclusive=True)
    return buffer if include_total else buffer[:n], int(read_scalar(buffer))

remap_indices 

remap_indices(
    indices: array[int32], remap: array[int32]
) -> wp.array[wp.int32]

Remap an index buffer through a lookup table, passing negative (sentinel) entries through.

Not a plain gather: the -1 slots that mark padded or removed entries in an index buffer must survive the remap unchanged, where a gather would read out of bounds on them. triwarp.repair relies on this — its functions preserve -1 face sentinels through vertex renumbering (see remove_unreferenced_vertices).

Parameters:

Name Type Description Default
indices array[int32]

1D wp.int32 array of indices into remap (e.g. a flat face buffer). Negative entries are passed through unchanged.

 required
remap array[int32]

1D wp.int32 lookup table (e.g. old-to-new vertex index map).

 required

Returns:

Type Description
array[int32]

Length len(indices) array on indices.device with out[i] = remap[indices[i]] for non-negative indices[i], and out[i] = indices[i] otherwise.

Raises:

Type Description
RuntimeError

If indices and remap are not all on one device.
See Also

gather The sentinel-free form: a dense first-axis gather for index buffers known to be in range.

Source code in triwarp/array.py 
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
def remap_indices(indices: wp.array[wp.int32], remap: wp.array[wp.int32]) -> wp.array[wp.int32]:
    """
    Remap an index buffer through a lookup table, passing negative (sentinel) entries through.

    Not a plain [`gather`][triwarp.array.gather]: the ``-1`` slots that mark padded or removed
    entries in an index buffer must survive the remap unchanged, where a gather would read out of
    bounds on them. [`triwarp.repair`][triwarp.repair] relies on this — its functions preserve
    ``-1`` face sentinels through vertex renumbering (see ``remove_unreferenced_vertices``).

    Parameters
    ----------
    indices
        1D ``wp.int32`` array of indices into ``remap`` (e.g. a flat face buffer). Negative
        entries are passed through unchanged.
    remap
        1D ``wp.int32`` lookup table (e.g. old-to-new vertex index map).

    Returns
    -------
    wp.array[wp.int32]
        Length ``len(indices)`` array on ``indices.device`` with ``out[i] = remap[indices[i]]``
        for non-negative ``indices[i]``, and ``out[i] = indices[i]`` otherwise.

    Raises
    ------
    RuntimeError
        If ``indices`` and ``remap`` are not all on one device.

    See Also
    --------
    [`gather`][triwarp.array.gather]
        The sentinel-free form: a dense first-axis gather for index buffers known to be in range.
    """
    require_same_device(indices=indices, remap=remap)
    n = int(indices.shape[0])
    device = indices.device
    if n == 0:
        return wp.empty(0, dtype=wp.int32, device=device)
    out = wp.empty(n, dtype=wp.int32, device=device)
    wp.launch(
        kernel_array.gather_1d_skip_negative, dim=n, inputs=[indices, remap, out], device=device
    )
    return out

trim_to_count 

trim_to_count(
    counter: array[int32], *buffers: array[DType]
) -> tuple[int, list[wp.array[DType]]]

Trim atomic-append output buffers to the number of elements actually written.

A kernel that emits an unpredictable number of results cannot size its output ahead of time. The usual pattern is to over-allocate the output buffers to a safe upper bound and have each thread claim its slots with wp.atomic_add on a shared length-1 counter, writing into buffer[slot]. After the launch, only the first n_out slots hold valid data and the tail is uninitialized, but n_out is known only on the device.

This finalizes that pattern: it reads the counter back to the host once, then copies the valid prefix buffer[:n_out] of each over-allocated buffer into a freshly allocated, exact-size array. Pass every buffer filled by the same counter in one call so they are all trimmed to a consistent length.

Parameters:

Name Type Description Default
counter array[int32]

Length-1 wp.int32 array holding the final atomic-append count on the target device.

 required
*buffers array[DType]

Over-allocated output buffers to trim, all indexed along their first axis by the same counter. Any rank and dtype; trailing dimensions are preserved.

 ()

Returns:

Name Type Description
n_out int

The counter value: the number of valid leading elements in each buffer.
trimmed list[array]

One contiguous (n_out, *buffer.shape[1:]) copy per input buffer, in order, each on its buffer's device.
Source code in triwarp/array.py 
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
def trim_to_count(
    counter: wp.array[wp.int32], *buffers: wp.array[DType]
) -> tuple[int, list[wp.array[DType]]]:
    """
    Trim atomic-append output buffers to the number of elements actually written.

    A kernel that emits an unpredictable number of results cannot size its output ahead of
    time. The usual pattern is to over-allocate the output buffers to a safe upper bound and
    have each thread claim its slots with ``wp.atomic_add`` on a shared length-1 ``counter``,
    writing into ``buffer[slot]``. After the launch, only the first ``n_out`` slots hold valid
    data and the tail is uninitialized, but ``n_out`` is known only on the device.

    This finalizes that pattern: it reads the counter back to the host once, then copies the
    valid prefix ``buffer[:n_out]`` of each over-allocated buffer into a freshly allocated,
    exact-size array. Pass every buffer filled by the same counter in one call so they are all
    trimmed to a consistent length.

    Parameters
    ----------
    counter
        Length-1 ``wp.int32`` array holding the final atomic-append count on the target device.
    *buffers
        Over-allocated output buffers to trim, all indexed along their first axis by the same
        counter. Any rank and ``dtype``; trailing dimensions are preserved.

    Returns
    -------
    n_out : int
        The counter value: the number of valid leading elements in each buffer.
    trimmed : list[wp.array]
        One contiguous ``(n_out, *buffer.shape[1:])`` copy per input buffer, in order, each on
        its buffer's device.
    """
    n_out = int(read_scalar(counter, 0))
    trimmed = []
    for buffer in buffers:
        out_shape = (n_out, *(int(dim) for dim in buffer.shape[1:]))
        out = wp.empty(out_shape, dtype=buffer.dtype, device=buffer.device)
        if n_out > 0:
            wp.copy(out, buffer[:n_out])
        trimmed.append(out)
    return n_out, trimmed

bitcast_to_int 

bitcast_to_int(
    data: array[Scalar], count: int | None = None
) -> wp.array[wp.int32] | wp.array[wp.int64]

Reinterpret an array's underlying bits as a same-width signed integer dtype.

32-bit-or-narrower dtypes (wp.int32, wp.uint32, wp.float32, and narrower) are reinterpreted as wp.int32; wider dtypes (wp.int64, wp.uint64, wp.float64) as wp.int64. Narrower-than-32-bit floating point values are first upcast to wp.float32 (a numeric cast, not a bit reinterpretation) so every dtype narrower than 32 bits shares one wp.int32 key space. Used by the hashing/uniqueness machinery (unique_1d) to give arbitrary scalar dtypes a common sortable, hashable integer key.

Parameters:

Name Type Description Default
data array[Scalar]

Rank-1 wp.array of any scalar dtype.

 required
count int | None

Output length. Defaults to data.shape[0]. When greater than the input length, the tail is left uninitialized (over-allocation for in-place radix-sort scratch).

 None

Returns:

Type Description
array[int32] | array[int64]

Bit-reinterpreted (or, for sub-32-bit floats, upcast-then-reinterpreted) copy of length count on data.device.
See Also

bitcast_from_int

Source code in triwarp/array.py 
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
def bitcast_to_int(
    data: wp.array[wp.Scalar], count: int | None = None
) -> wp.array[wp.int32] | wp.array[wp.int64]:
    """
    Reinterpret an array's underlying bits as a same-width signed integer dtype.

    32-bit-or-narrower dtypes (``wp.int32``, ``wp.uint32``, ``wp.float32``, and narrower) are
    reinterpreted as ``wp.int32``; wider dtypes (``wp.int64``, ``wp.uint64``, ``wp.float64``) as
    ``wp.int64``. Narrower-than-32-bit floating point values are first upcast to ``wp.float32``
    (a numeric cast, not a bit reinterpretation) so every dtype narrower than 32 bits shares one
    ``wp.int32`` key space. Used by the hashing/uniqueness machinery
    ([`unique_1d`][triwarp.grouping.unique_1d]) to give arbitrary scalar dtypes a common sortable,
    hashable integer key.

    Parameters
    ----------
    data
        Rank-1 ``wp.array`` of any scalar dtype.
    count
        Output length. Defaults to ``data.shape[0]``. When greater than the input length, the
        tail is left uninitialized (over-allocation for in-place radix-sort scratch).

    Returns
    -------
    wp.array[wp.int32] | wp.array[wp.int64]
        Bit-reinterpreted (or, for sub-32-bit floats, upcast-then-reinterpreted) copy of length
        ``count`` on ``data.device``.

    See Also
    --------
    [`bitcast_from_int`][triwarp.array.bitcast_from_int]
    """
    n_bits = wp.types.type_size_in_bytes(data.dtype) * 8
    n = data.shape[0]
    # ``is None``, not ``or``: zero is a legitimate length and ``or`` would read it as "not passed"
    # and hand back the whole buffer -- an ``n``-element array of stale bits where the caller asked
    # for an empty one. The signature already spells the distinction; only this line lost it.
    count = n if count is None else count
    copy_count = min(n, count)
    target = wp.int64 if n_bits > 32 else wp.int32

    # ``count=0`` reaches ``wp.copy`` and ``wp.utils.array_cast`` as *"copy the whole source"* --
    # a documented back-compatibility rule in Warp 1.17 (``if count == 0: count = src.size``), so
    # the zero that means "nothing" and the zero that means "everything" are the same argument.
    # Into a length-0 destination that is not even a clean refusal: it raises ``TypeError:
    # unsupported operand type(s) for +: 'NoneType' and 'int'`` from inside the copy. Returning the
    # empty allocation here is both the right answer and the only way to state it.
    if copy_count == 0:
        return wp.empty(count, dtype=target, device=data.device)

    if n_bits > 32:
        reinterpreted = wp.empty(count, dtype=wp.int64, device=data.device)
        wp.copy(reinterpreted, data, count=copy_count)
        return reinterpreted

    reinterpreted = wp.empty(count, dtype=wp.int32, device=data.device)
    if wp.types.type_is_float(data.dtype):
        src = data
        if n_bits < 32:
            src = wp.empty(copy_count, dtype=wp.float32, device=data.device)
            wp.utils.array_cast(data, src, count=copy_count)
        wp.copy(reinterpreted, src, count=copy_count)
    else:
        wp.utils.array_cast(data, reinterpreted, count=copy_count)
    return reinterpreted

bitcast_from_int 

bitcast_from_int(
    data: array[int32] | array[int64],
    dtype: type[Scalar],
    count: int | None = None,
) -> wp.array[wp.Scalar]

Inverse of bitcast_to_int: recover the original dtype.

Reinterprets (same-width) or upcasts-then-reinterprets (narrower target) the bits produced by bitcast_to_int back into dtype.

Parameters:

Name Type Description Default
data array[int32] | array[int64]

Rank-1 wp.array[wp.int32] or wp.array[wp.int64], typically the output of bitcast_to_int.

 required
dtype type[Scalar]

Target scalar dtype to reinterpret data as.

 required
count int | None

Output length. Defaults to data.shape[0]. When greater than the input length, the tail is left uninitialized.

 None

Returns:

Type Description
array[Scalar]

Array of dtype dtype and length count on data.device.
See Also

bitcast_to_int

Source code in triwarp/array.py 
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
def bitcast_from_int(
    data: wp.array[wp.int32] | wp.array[wp.int64], dtype: type[wp.Scalar], count: int | None = None
) -> wp.array[wp.Scalar]:
    """
    Inverse of [`bitcast_to_int`][triwarp.array.bitcast_to_int]: recover the original dtype.

    Reinterprets (same-width) or upcasts-then-reinterprets (narrower target) the bits produced
    by ``bitcast_to_int`` back into ``dtype``.

    Parameters
    ----------
    data
        Rank-1 ``wp.array[wp.int32]`` or ``wp.array[wp.int64]``, typically the output of
        [`bitcast_to_int`][triwarp.array.bitcast_to_int].
    dtype
        Target scalar dtype to reinterpret ``data`` as.
    count
        Output length. Defaults to ``data.shape[0]``. When greater than the input length, the
        tail is left uninitialized.

    Returns
    -------
    wp.array[wp.Scalar]
        Array of dtype ``dtype`` and length ``count`` on ``data.device``.

    See Also
    --------
    [`bitcast_to_int`][triwarp.array.bitcast_to_int]
    """
    n_bits = wp.types.type_size_in_bytes(data.dtype) * 8
    n_target_bits = wp.types.type_size_in_bytes(dtype) * 8
    n = data.shape[0]
    # ``is None`` rather than ``or``, and the zero-length short circuit, both for the reasons
    # recorded at ``bitcast_to_int``: Warp reads ``count=0`` as "copy everything".
    count = n if count is None else count
    copy_count = min(n, count)
    if copy_count == 0:
        return wp.empty(count, dtype=dtype, device=data.device)

    if n_bits == n_target_bits:
        reinterpreted = wp.empty(count, dtype=dtype, device=data.device)
        wp.copy(reinterpreted, data, count=copy_count)
        return reinterpreted

    if wp.types.type_is_float(dtype) and n_bits > n_target_bits:
        wide_dtype = getattr(wp, f"float{n_bits}")
        wide = wp.empty(count, dtype=wide_dtype, device=data.device)
        wp.copy(wide, data, count=copy_count)
        reinterpreted_casted = wp.empty(count, dtype=dtype, device=data.device)
        wp.utils.array_cast(wide, reinterpreted_casted, count=copy_count)
    else:
        reinterpreted_casted = wp.empty(count, dtype=dtype, device=data.device)
        wp.utils.array_cast(data, reinterpreted_casted, count=copy_count)
    return reinterpreted_casted