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 (TypeVar(Array, bound= Array)) – input array of values.

  • test_elements (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:

TypeVar(Array, bound= 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]]