Merge pull request #23 from mtsokol/finch-random

API: Add `finch.random` function

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "finch-tensor"
-version = "0.1.6"
+version = "0.1.7"
 description = ""
 authors = ["Willow Ahrens <willow.marie.ahrens@gmail.com>"]
 readme = "README.md"

src/finch/__init__.py

@@ -14,7 +14,7 @@
 from .tensor import (
     Tensor,
     astype,
-    fsprand,
+    random,
     permute_dims,
     multiply,
     sum,
@@ -64,7 +64,7 @@
     "Storage",
     "DenseStorage",
     "astype",
-    "fsprand",
+    "random",
     "permute_dims",
     "int_",
     "int8",

src/finch/julia.py

@@ -7,3 +7,4 @@
 from juliacall import Main as jl  # noqa
 
 jl.seval("using Finch")
+jl.seval("using Random: default_rng")

src/finch/tensor.py

@@ -396,7 +396,15 @@ def construct_csf(cls, arg: TupleOf3Arrays, shape: tuple[int, ...]) -> "Tensor":
         return Tensor(cls.construct_csf_jl_object(arg, shape))
 
 
-def fsprand(*args):
+def random(shape, density=0.01, random_state=None):
+    args = [*shape, density]
+    if random_state is not None:
+        if isinstance(random_state, np.random.Generator):
+            seed = random_state.integers(np.iinfo(np.int32).max)
+        else:
+            seed = random_state
+        rng = jl.default_rng(seed)
+        args = [rng] + args
     return Tensor(jl.fsprand(*args))
 
 

tests/test_sparse.py

@@ -136,3 +136,11 @@ def test_astype(arr3d, order):
 
     with pytest.raises(ValueError, match="Unable to avoid a copy while casting in no-copy mode."):
         finch.astype(arr_finch, finch.float64, copy=False)
+
+
+@pytest.mark.parametrize("random_state", [42, np.random.default_rng(42)])
+def test_random(random_state):
+    result = finch.random((10, 20, 30), density=0.0, random_state=random_state)
+    expected = sparse.random((10, 20, 30), density=0.0, random_state=random_state)
+
+    assert_equal(result.todense(), expected.todense())