Skip to content

Commit

Permalink
function if changes + test standartization
Browse files Browse the repository at this point in the history
  • Loading branch information
@ohrechykha
ohrechykha committed Aug 26, 2024
1 parent 2e92e46 commit 1ca9106
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 82 deletions.
35 changes: 13 additions & 22 deletions src/ragged/_spec_set_functions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,9 @@ def unique_all(x: array, /) -> tuple[array, array, array, array]:
"""

if isinstance(x, ragged.array):
if len(x) == 1:
if x.ndim == 0:
return unique_all_result(
values=ragged.array(x),
values=ragged.array([x]),
indices=ragged.array([0]),
inverse_indices=ragged.array([0]),
counts=ragged.array([1]),
Expand All @@ -67,6 +67,7 @@ def unique_all(x: array, /) -> tuple[array, array, array, array]:
return_index=True,
return_inverse=True,
return_counts=True,
equal_nan=False,
)
return unique_all_result(
values=ragged.array(values),
Expand All @@ -75,8 +76,8 @@ def unique_all(x: array, /) -> tuple[array, array, array, array]:
counts=ragged.array(counts),
)
else:
msg = f"Expected ragged type but got {type(x)}"
raise TypeError(msg)
msg = f"Expected ragged type but got {type(x)}" # type: ignore
raise TypeError(msg) # type: ignore


unique_counts_result = namedtuple( # pylint: disable=C0103
Expand Down Expand Up @@ -110,19 +111,15 @@ def unique_counts(x: array, /) -> tuple[array, array]:
return unique_counts_result(
values=ragged.array([x]), counts=ragged.array([1])
)
elif len(x) == 1:
return unique_counts_result(
values=ragged.array(x), counts=ragged.array([1])
)
else:
x_flat = ak.ravel(x._impl)
values, counts = np.unique(x_flat.layout.data, return_counts=True)
return unique_counts_result(
values=ragged.array(values), counts=ragged.array(counts)
)
else:
msg = f"Expected ragged type but got {type(x)}"
raise TypeError(msg)
msg = f"Expected ragged type but got {type(x)}" # type: ignore
raise TypeError(msg) # type: ignore


unique_inverse_result = namedtuple( # pylint: disable=C0103
Expand Down Expand Up @@ -152,13 +149,9 @@ def unique_inverse(x: array, /) -> tuple[array, array]:
https://data-apis.org/array-api/latest/API_specification/generated/array_api.unique_inverse.html
"""
if isinstance(x, ragged.array):
if x.ndim == 0:
return unique_inverse_result(
values=ragged.array([x]), inverse_indices=ragged.array([0])
)
elif len(x) == 1:
if ak.is_scalar(x):
return unique_inverse_result(
values=ragged.array(x), inverse_indices=ragged.array([0])
values=x, inverse_indices=ragged.array([0])
)
else:
x_flat = ak.ravel(x._impl)
Expand All @@ -169,8 +162,8 @@ def unique_inverse(x: array, /) -> tuple[array, array]:
inverse_indices=ragged.array(inverse_indices),
)
else:
msg = f"Expected ragged type but got {type(x)}"
raise TypeError(msg)
msg = f"Expected ragged type but got {type(x)}" # type: ignore
raise TypeError(msg) # type: ignore


def unique_values(x: array, /) -> array:
Expand All @@ -192,11 +185,9 @@ def unique_values(x: array, /) -> array:
if x.ndim == 0:
return ragged.array([x])

if len(x) == 1:
return ragged.array(x)
else:
x_flat = ak.ravel(x._impl)
return ragged.array(np.unique(x_flat.layout.data))
else:
err = f"Expected ragged type but got {type(x)}"
raise TypeError(err)
err = f"Expected ragged type but got {type(x)}" # type: ignore
raise TypeError(err) # type: ignore
108 changes: 48 additions & 60 deletions tests/test_spec_set_functions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,6 @@

from __future__ import annotations

import re

import awkward as ak
import pytest

Expand All @@ -22,36 +20,28 @@ def test_existence():


# unique_values tests
def test_can_take_none():
with pytest.raises(TypeError):
assert ragged.unique_values(ragged.array(None)) is None
# def test_can_take_none():
# with pytest.raises(TypeError):
# assert ragged.unique_values(ragged.array(None)) is None


def test_can_take_list():
with pytest.raises(
ValueError,
match=re.escape(
"the truth value of an array whose length is not 1 is ambiguous;"
),
):
assert ragged.unique_values(
ragged.array([1, 2, 4, 3, 4, 5, 6, 20])
) == ragged.array([1, 2, 3, 4, 5, 6, 20])
assert ak.to_list(
ragged.unique_values(ragged.array([1, 2, 4, 3, 4, 5, 6, 20]))
) == ak.to_list(ragged.array([1, 2, 3, 4, 5, 6, 20]))


def test_can_take_empty_arr():
with pytest.raises(TypeError):
assert ragged.unique_values(ragged.array([])) == ragged.array([])
# with pytest.raises(TypeError):
assert ak.to_list(ragged.unique_values(ragged.array([]))) == ak.to_list(
ragged.array([])
)


def test_can_take_moredimensions():
with pytest.raises(
ValueError,
match=re.escape(
"the truth value of an array whose length is not 1 is ambiguous;"
),
):
assert ragged.unique_values(ragged.array([[1, 2, 3, 4], [5, 6]]))
assert ak.to_list(
ragged.unique_values(ragged.array([[1, 2, 2, 3, 4], [5, 6]]))
) == ak.to_list(ragged.array([1, 2, 3, 4, 5, 6]))


def test_can_take_1d_array():
Expand All @@ -61,16 +51,18 @@ def test_can_take_1d_array():


# unique_counts tests
def test_can_count_none():
with pytest.raises(TypeError):
assert ragged.unique_counts(ragged.array(None)) is None
# def test_can_count_none():
# with pytest.raises(TypeError):
# assert ragged.unique_counts(ragged.array(None)) is None


def test_can_count_list():
with pytest.raises(TypeError):
assert ragged.unique_counts(
ragged.array([1, 2, 4, 3, 4, 5, 6, 20])
) == ragged.array([1, 2, 3, 4, 5, 6, 20]), ragged.array([1, 1, 2, 1, 1, 1, 1])
arr = ragged.array([1, 2, 4, 3, 4, 5, 6, 20])
expected_unique_values = ragged.array([1, 2, 3, 4, 5, 6, 20])
expected_unique_counts = ragged.array([1, 1, 1, 2, 1, 1, 1])
unique_values, unique_counts = ragged.unique_counts(arr)
assert ak.to_list(unique_values) == ak.to_list(expected_unique_values)
assert ak.to_list(unique_counts) == ak.to_list(expected_unique_counts)


def test_can_count_simple_array():
Expand Down Expand Up @@ -101,21 +93,22 @@ def test_can_count_scalar():


# unique_inverse tests
def test_can_inverse_none():
with pytest.raises(TypeError):
assert ragged.unique_inverse(ragged.array(None)) is None
# def test_can_inverse_none():
# with pytest.raises(TypeError):
# assert ak.to_list(ragged.unique_inverse(ragged.array(None))) is ak.to_list(None)


def test_can_inverse_list():
with pytest.raises(TypeError):
assert ragged.unique_inverse(
ragged.array([1, 2, 4, 3, 4, 5, 6, 20])
) == ragged.array([1, 2, 3, 4, 5, 6, 20]), ragged.array(
[0, 1, 3, 2, 3, 4, 5, 6]
)
arr=ragged.array([1, 2, 4, 3, 4, 5, 6, 20])
expected_values=ragged.array([1,2,3,4,5,6,20])
expected_inverse=ragged.array([0, 1, 3, 2, 3, 4, 5, 6])
values, inverse= ragged.unique_inverse(arr)
assert ak.to_list(expected_values)==ak.to_list(values)
assert ak.to_list(expected_inverse)==ak.to_list(inverse)



def test_can_take_simple_array():
def test_can_inverse_simple_array():
arr = ragged.array([[1, 2, 2], [3, 3, 3], [4, 4, 4, 4]])
expected_unique_values = ragged.array([1, 2, 3, 4])
expected_inverse_indices = ragged.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3])
Expand All @@ -124,7 +117,7 @@ def test_can_take_simple_array():
assert ak.to_list(inverse_indices) == ak.to_list(expected_inverse_indices)


def test_can_take_normal_array():
def test_can_inverse_normal_array():
arr = ragged.array([[1, 2, 2], [3], [3, 3], [4, 4, 4], [4]])
expected_unique_values = ragged.array([1, 2, 3, 4])
expected_inverse_indices = ragged.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3])
Expand All @@ -133,9 +126,9 @@ def test_can_take_normal_array():
assert ak.to_list(inverse_indices) == ak.to_list(expected_inverse_indices)


def test_can_take_scalar():
arr = ragged.array([5])
expected_unique_values = ragged.array([5])
def test_can_inverse_scalar():
arr = ragged.array(5)
expected_unique_values = 5
expected_unique_indices = ragged.array([0])
unique_values, unique_indices = ragged.unique_inverse(arr)
assert ak.to_list(unique_values) == ak.to_list(expected_unique_values)
Expand All @@ -145,25 +138,20 @@ def test_can_take_scalar():
# unique_all tests
def test_can_all_none():
with pytest.raises(TypeError):
assert ragged.unique_all(ragged.array(None)) is None
arr=None
expected_unique_values = ragged.array(None)
expected_unique_indices = ragged.array(None)
expected_unique_inverse = ragged.array(None)
expected_unique_counts = ragged.array(None)
unique_values, unique_indices, unique_inverse, unique_counts = ragged.unique_all(arr)
assert ak.to_list(unique_values) == ak.to_list(expected_unique_values)
assert ak.to_list(unique_indices) == ak.to_list(expected_unique_indices)
assert ak.to_list(unique_inverse) == ak.to_list(expected_unique_inverse)
assert ak.to_list(unique_counts) == ak.to_list(expected_unique_counts)


def test_can_all_list():
with pytest.raises(
ValueError,
match=re.escape(
"the truth value of an array whose length is not 1 is ambiguous;"
),
):
assert ragged.unique_all(ragged.array([1, 2, 4, 3, 4, 5, 6, 20])) == (
ragged.array([1, 2, 3, 4, 5, 6, 20]),
ragged.array([0, 1, 3, 2, 5, 6, 7]),
ragged.array([0, 1, 3, 2, 3, 4, 5, 6]),
ragged.array([1, 1, 1, 2, 1, 1, 1]),
)


def test_can_all_simple_array():
def test_can_all_list():
arr = ragged.array([1, 2, 2, 3, 3, 3, 4, 4, 4, 4])
expected_unique_values = ragged.array([1, 2, 3, 4])
expected_unique_indices = ragged.array([0, 1, 3, 6])
Expand Down

0 comments on commit 1ca9106

Please sign in to comment.