isin¶

atlas_schema.isin(element, test_elements, axis=-1)[source]¶

Find test_elements in element. Similar in API as numpy.isin().

Calculates element in test_elements, broadcasting over element elements only. Returns a boolean array of the same shape as element that is True where an element of element is in test_elements and False otherwise.

This works by first transforming test_elements to an array with one more dimension than the element, placing the test_elements at axis, and then doing a comparison.

Parameters:

element (dask_awkward.Array or ak.Array) – input array of values.
test_elements (dask_awkward.Array or ak.Array) – one-dimensional set of values against which to test each value of element.
axis (int) – the axis along which the comparison is performed

Returns:

result of comparison for test_elements in element

Return type:

dask_awkward.Array or ak.Array

Example

>>> import awkward as ak
>>> import atlas_schema as ats
>>> truth_origins = ak.Array([[1, 2, 3], [4], [5, 6, 7], [1]])
>>> prompt_origins = ak.Array([1, 2, 7])
>>> ats.isin(truth_origins, prompt_origins).to_list()
[[True, True, False], [False], [False, False, True], [True]]