final corrections in inverse_indices and tests

src/ragged/_spec_set_functions.py

@@ -60,9 +60,8 @@ def unique_all(x: array, /) -> tuple[array, array, array, array]:
         return ragged.array(x), ragged.array([0]), ragged.array([0]), ragged.array([1])
 
     x_flat=ak.ravel(x._impl)   
-    values, indices, counts = np.unique(x_flat.layout.data, return_index=True, return_counts=True)   
-    sorted_indices = np.argsort(indices)
-    inverse_indices = values[sorted_indices]
+    values, indices, inverse_indices, counts = np.unique(x_flat.layout.data, return_index=True, return_inverse=True, return_counts=True)
+
     return ragged.array(values),ragged.array(indices), ragged.array(inverse_indices), ragged.array(counts)
 
 
@@ -138,10 +137,7 @@ def unique_inverse(x: array, /) -> tuple[array, array]:
        return ragged.array(x), ragged.array([0])
 
     x_flat=ak.ravel(x._impl)
-    values, first_indices = np.unique(x_flat.layout.data, return_index=True)
-
-    sorted_indices = np.argsort(first_indices)
-    inverse_indices = first_indices[sorted_indices]
+    values, inverse_indices = np.unique(x_flat.layout.data, return_inverse=True)
 
     return ragged.array(values), ragged.array(inverse_indices)
 

tests/test_spec_set_functions.py → tests/test_55_spec_set_functions.py

@@ -18,13 +18,14 @@
 #    if not then add it to the output and repeat the step}
 # 4 once the cycle is over return an array of unique elements in the input array (the output must be of the same type as input array)
 
-#unique_values tests
+
 def test_existence():
     assert ragged.unique_all is not None
     assert ragged.unique_counts is not None
     assert ragged.unique_inverse is not None
     assert ragged.unique_values is not None
 
+#unique_values tests
 def test_can_take_none():
     assert ragged.unique_values(None)==None
 
@@ -37,28 +38,24 @@ def test_can_take_empty_arr():
         assert ragged.unique_values(ragged.array([]))
 
 def test_can_take_moredimensions():
-    assert ragged.unique_values(ragged.array([[1,2,3,4],[5,6]]))
+    with pytest.raises(ValueError):
+        assert ragged.unique_values(ragged.array([[1,2,3,4],[5,6]]))
 
 def test_can_take_1d_array():
-    assert ragged.unique_values(ragged.array([5,6,7,8,8,9,1,2,3,4,10,0,15,2]))==ragged.array([0,1,2,3,4,5,6,7,8,9,10,15])
-
-def test_can_take_awkward():
-    with pytest.raises(TypeError):
-        assert ragged.unique_values(ak.Array([1,2,3,4,[5,6,7]]))
+    arr=ragged.array([5,6,7,8,8,9,1,2,3,4,10,0,15,2])
+    expected_unique_values = ragged.array([0,1,2,3,4,5,6,7,8,9,10,15])
+    assert ak.to_list(ragged.unique_values(arr))==ak.to_list(expected_unique_values)
 
 
 #unique_counts tests
 def test_can_count_none():
-    assert ragged.unique_counts(None) is None
+    with pytest.raises(TypeError):
+        assert ragged.unique_counts(None) is None
 
 def test_can_count_list():
     with pytest.raises(TypeError):
         assert ragged.unique_counts([1,2,4,3,4,5,6,20]) is None
 
-def test_can_create_empty_arr():
-    with pytest.raises(TypeError):
-        arr = ragged.array([])
-
 def test_can_count_simple_array():
     arr = ragged.array([1, 2, 2, 3, 3, 3, 4, 4, 4, 4])
     expected_unique_values = ragged.array([1, 2, 3, 4])
@@ -75,10 +72,6 @@ def test_can_count_normal_array():
     assert ak.to_list(unique_values) == ak.to_list(expected_unique_values)
     assert ak.to_list(unique_counts) == ak.to_list(expected_counts)
 
-def test_can_count_awkward():
-    arr = ak.Array([[[1, 2, 2], 3, 3], [3, 4, 4, 4, 4]])
-    with pytest.raises(TypeError):
-        assert ragged.unique_counts(arr) is None
 
 def test_can_count_scalar():
     arr = ragged.array([5])
@@ -90,40 +83,29 @@ def test_can_count_scalar():
 
 #unique_inverse tests
 def test_can_take_none():
-    assert ragged.unique_inverse(None) is None
+    with pytest.raises(TypeError):
+        assert ragged.unique_inverse(None) is None
 
 def test_can_take_list():
     with pytest.raises(TypeError):
         assert ragged.unique_inverse([1,2,4,3,4,5,6,20]) is None
 
-def test_can_take_empty_arr():
-    arr = ragged.array([])
-    expected_unique_values = ragged.array([])
-    expected_unique_indices = ragged.array([])
-    unique_values, unique_indices = ragged.unique_inverse(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)
-
 def test_can_take_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_unique_indices = ragged.array([0, 1, 3, 6])
-    unique_values, unique_indices = ragged.unique_inverse(arr)
+    expected_inverse_indices = ragged.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3])
+    unique_values, inverse_indices = ragged.unique_inverse(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(inverse_indices) == ak.to_list(expected_inverse_indices)
 
 def test_can_take_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_unique_indices = ragged.array([0, 1, 3, 6])
-    unique_values, unique_indices = ragged.unique_inverse(arr)
+    expected_inverse_indices = ragged.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3])
+    unique_values, inverse_indices = ragged.unique_inverse(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(inverse_indices) == ak.to_list(expected_inverse_indices)
 
-def test_can_take_awkward():
-    arr = ragged.array([[[1, 2, 2], 3, 3], [3, 4, 4, 4, 4]])
-    with pytest.raises(TypeError):
-        assert ragged.unique_inverse(arr) is None
 
 def test_can_take_scalar():
     arr = ragged.array([5])
@@ -134,3 +116,46 @@ def test_can_take_scalar():
     assert ak.to_list(unique_indices) == ak.to_list(expected_unique_indices)
 
 #unique_all tests
+def test_can_all_none():
+    with pytest.raises(TypeError):
+        assert ragged.unique_all(None) is None
+
+def test_can_all_list():
+    with pytest.raises(TypeError):
+        assert ragged.unique_all([1,2,4,3,4,5,6,20]) is None
+
+def test_can_all_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_unique_indices = ragged.array([0, 1, 3, 6])
+    expected_unique_inverse = ragged.array([0, 1,1,2,2,2,3,3,3,3])
+    expected_unique_counts = ragged.array([1, 2, 3, 4])
+    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_normal_array():
+    arr = ragged.array([[2, 2, 2], [3], [3, 5], [4, 4, 4], [4]])
+    expected_unique_values = ragged.array([2, 3, 4, 5])
+    expected_unique_indices = ragged.array([0, 3, 6, 5])
+    expected_unique_inverse = ragged.array([0,0,0,1,1,3,2,2,2,2])
+    expected_unique_counts = ragged.array([3, 2, 4, 1])
+    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_scalar():
+    arr = ragged.array([5])
+    expected_unique_values = ragged.array([5])
+    expected_unique_indices = ragged.array([0])
+    expected_unique_inverse = ragged.array([0])
+    expected_unique_counts = ragged.array([1])
+    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)