WIP: Modify level construction.

Author: hameerabbasi

.github/workflows/publish.yml

@@ -1,4 +1,4 @@
-name: Publish 
+name: Publish
 on:
   workflow_dispatch:
 jobs:

.gitignore

@@ -50,6 +50,7 @@ coverage.xml
 .hypothesis/
 .pytest_cache/
 cover/
+junit/
 
 # Translations
 *.mo

poetry.lock

@@ -1,16 +1,3 @@
-from .levels import (
-    Dense,
-    Element,
-    Pattern,
-    SparseList,
-    SparseByteMap,
-    RepeatRLE,
-    SparseVBL,
-    SparseCOO,
-    SparseHash,
-    Storage,
-    DenseStorage,
-)
 from .tensor import (
     Tensor,
     astype,

src/finch/__init__.py

@@ -1,93 +1,70 @@
-import numpy as np
+import abc
+
 
 from .julia import jl
-from .typing import OrderType
+from . import dtypes
+from dataclasses import dataclass
 
 
-class _Display:
+class _Display(abc.ABC):
     def __repr__(self):
         return jl.sprint(jl.show, self._obj)
 
     def __str__(self):
         return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj)
 
 
-# LEVEL
+class AbstractLeafLevel(abc.ABC):
+    @abc.abstractmethod
+    def _construct(self, *, dtype, fill_value):
+        ...
+
 
-class AbstractLevel(_Display):
-    pass
+# LEVEL
+class AbstractLevel(abc.ABC):
+    @abc.abstractmethod
+    def _construct(self, *, inner_level):
+        ...
 
 
 # core levels
-
+@dataclass
 class Dense(AbstractLevel):
-    def __init__(self, lvl, shape=None):
-        args = [lvl._obj]
-        if shape is not None:
-            args.append(shape)
-        self._obj = jl.Dense(*args)
-
-
-class Element(AbstractLevel):
-    def __init__(self, fill_value, data=None):
-        args = [fill_value]
-        if data is not None:
-            args.append(data)
-        self._obj = jl.Element(*args)
-
-
-class Pattern(AbstractLevel):
-    def __init__(self):
-        self._obj = jl.Pattern()
-
-
-# advanced levels
-
-class SparseList(AbstractLevel):
-    def __init__(self, lvl):
-        self._obj = jl.SparseList(lvl._obj)
-
-
-class SparseByteMap(AbstractLevel):
-    def __init__(self, lvl):
-        self._obj = jl.SparseByteMap(lvl._obj)
-
+    dim: int | None = None
+    index_type: jl.DataType = dtypes.int64
 
-class RepeatRLE(AbstractLevel):
-    def __init__(self, lvl):
-        self._obj = jl.RepeatRLE(lvl._obj)
+    def _construct(self, *, inner_level) -> jl.Dense:
+        if self.dim is None:
+            return jl.Dense[self.index_type](inner_level)
 
+        return jl.Dense[self.index_type](inner_level, self.dim)
 
-class SparseVBL(AbstractLevel):
-    def __init__(self, lvl):
-        self._obj = jl.SparseVBL(lvl._obj)
 
+@dataclass
+class Element(AbstractLeafLevel):
+    def _construct(self, *, dtype: jl.DataType, fill_value) -> jl.Element:
+        return jl.Element[fill_value, dtype]()
 
-class SparseCOO(AbstractLevel):
-    def __init__(self, ndim, lvl):
-        self._obj = jl.SparseCOO[ndim](lvl._obj)
 
+@dataclass
+class Pattern(AbstractLeafLevel):
+    def _construct(self, *, dtype, fill_value) -> jl.Pattern:
+        from .dtypes import bool
 
-class SparseHash(AbstractLevel):
-    def __init__(self, ndim, lvl):
-        self._obj = jl.SparseHash[ndim](lvl._obj)
+        if dtype != bool:
+            raise TypeError("`Pattern` can only have `dtype=bool`.")
+        if dtype(fill_value) != dtype(False):
+            raise TypeError("`Pattern` can only have `fill_value=False`.")
 
+        return jl.Pattern()
 
-# STORAGE
 
-class Storage:
-    def __init__(self, levels_descr: AbstractLevel, order: OrderType = None):
-        self.levels_descr = levels_descr
-        self.order = order if order is not None else "C"
-
-    def __str__(self) -> str:
-        return f"Storage(lvl={str(self.levels_descr)}, order={self.order})"
-
-
-class DenseStorage(Storage):
-    def __init__(self, ndim: int, dtype: np.dtype, order: OrderType = None):
-        lvl = Element(np.int_(0).astype(dtype))
-        for _ in range(ndim):
-            lvl = Dense(lvl)
+# advanced levels
+@dataclass
+class SparseList(AbstractLevel):
+    index_type: jl.DataType = dtypes.int64
+    pos_type: jl.DataType = dtypes.uint64
+    crd_type: jl.DataType = dtypes.uint64
 
-        super().__init__(levels_descr=lvl, order=order)
+    def _construct(self, *, inner_level) -> jl.SparseList:
+        return jl.SparseList[self.index_type, self.pos_type, self.crd_type](inner_level)

src/finch/levels.py

@@ -1,10 +1,10 @@
-from typing import Any, Optional, Union
+from typing import Union
 
 import numpy as np
 from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple
 
 from .julia import jl
-from .levels import _Display, Dense, Element, Storage
+from .levels import _Display, Dense, Element
 from .typing import OrderType, JuliaObj, spmatrix, TupleOf3Arrays
 
 
@@ -57,12 +57,13 @@ class Tensor(_Display):
     array([[0, 1, 2],
            [3, 4, 5]])
     """
+
     row_major = "C"
     column_major = "F"
 
     def __init__(
         self,
-        obj: Union[np.ndarray, spmatrix, Storage, JuliaObj],
+        obj: Union[np.ndarray, spmatrix, JuliaObj],
         /,
         *,
         fill_value: np.number = 0.0,
@@ -74,7 +75,9 @@ def __init__(
             jl_data = self._from_numpy(obj, fill_value=fill_value)
             self._obj = jl_data
         elif isinstance(obj, Storage):  # from-storage constructor
-            order = self.preprocess_order(obj.order, self.get_lvl_ndim(obj.levels_descr._obj))
+            order = self.preprocess_order(
+                obj.order, self.get_lvl_ndim(obj.levels_descr._obj)
+            )
             self._obj = jl.swizzle(jl.Tensor(obj.levels_descr._obj), *order)
         elif jl.isa(obj, jl.Finch.SwizzleArray):  # raw-Julia-object constructor
             self._obj = obj
@@ -143,25 +146,22 @@ def _order(self) -> tuple[int, ...]:
         return jl.typeof(self._obj).parameters[1]
 
     @classmethod
-    def preprocess_order(
-        cls, order: OrderType, ndim: int
-    ) -> tuple[int, ...]:
+    def preprocess_order(cls, order: OrderType, ndim: int) -> tuple[int, ...]:
         if order == cls.column_major:
             permutation = tuple(range(1, ndim + 1))
         elif order == cls.row_major or order is None:
             permutation = tuple(range(1, ndim + 1)[::-1])
         elif isinstance(order, tuple):
             if min(order) == 0:
                 order = tuple(i + 1 for i in order)
-            if (
-                len(order) == ndim and
-                all([i in order for i in range(1, ndim + 1)])
-            ):
+            if len(order) == ndim and all([i in order for i in range(1, ndim + 1)]):
                 permutation = order
             else:
                 raise ValueError(f"Custom order is not a permutation: {order}.")
         else:
-            raise ValueError(f"order must be 'C', 'F' or a tuple, but is: {type(order)}.")
+            raise ValueError(
+                f"order must be 'C', 'F' or a tuple, but is: {type(order)}."
+            )
 
         return permutation
 
@@ -214,11 +214,11 @@ def permute_dims(self, axes: tuple[int, ...]) -> "Tensor":
         new_tensor = Tensor(new_obj)
         return new_tensor
 
-    def to_device(self, device: Storage) -> "Tensor":
+    def to_device(self, device) -> "Tensor":
         return Tensor(self._from_other_tensor(self, storage=device))
 
     @classmethod
-    def _from_other_tensor(cls, tensor: "Tensor", storage: Optional[Storage]) -> JuliaObj:
+    def _from_other_tensor(cls, tensor: "Tensor", storage) -> JuliaObj:
         order = cls.preprocess_order(storage.order, tensor.ndim)
         return jl.swizzle(
             jl.Tensor(storage.levels_descr._obj, tensor._obj.body), *order
@@ -239,7 +239,10 @@ def _from_numpy(cls, arr: np.ndarray, fill_value: np.number) -> JuliaObj:
     def _from_scipy_sparse(cls, x) -> JuliaObj:
         if x.format == "coo":
             return cls.construct_coo_jl_object(
-                coords=(x.col, x.row), data=x.data, shape=x.shape[::-1], order=Tensor.row_major
+                coords=(x.col, x.row),
+                data=x.data,
+                shape=x.shape[::-1],
+                order=Tensor.row_major,
             )
         elif x.format == "csc":
             return cls.construct_csc_jl_object(
@@ -255,7 +258,9 @@ def _from_scipy_sparse(cls, x) -> JuliaObj:
             raise ValueError(f"Unsupported SciPy format: {type(x)}")
 
     @classmethod
-    def construct_coo_jl_object(cls, coords, data, shape, order, fill_value=0.0) -> JuliaObj:
+    def construct_coo_jl_object(
+        cls, coords, data, shape, order, fill_value=0.0
+    ) -> JuliaObj:
         assert len(coords) == 2
         ndim = len(shape)
         order = cls.preprocess_order(order, ndim)
@@ -264,12 +269,18 @@ def construct_coo_jl_object(cls, coords, data, shape, order, fill_value=0.0) ->
         ptr = jl.Vector[jl.Int]([1, len(data) + 1])
         tbl = tuple(jl.PlusOneVector(arr) for arr in coords)
 
-        jl_data = jl.swizzle(jl.Tensor(jl.SparseCOO[ndim](lvl, shape, ptr, tbl)), *order)
+        jl_data = jl.swizzle(
+            jl.Tensor(jl.SparseCOO[ndim](lvl, shape, ptr, tbl)), *order
+        )
         return jl_data
 
     @classmethod
-    def construct_coo(cls, coords, data, shape, order=row_major, fill_value=0.0) -> "Tensor":
-        return Tensor(cls.construct_coo_jl_object(coords, data, shape, order, fill_value))
+    def construct_coo(
+        cls, coords, data, shape, order=row_major, fill_value=0.0
+    ) -> "Tensor":
+        return Tensor(
+            cls.construct_coo_jl_object(coords, data, shape, order, fill_value)
+        )
 
     @staticmethod
     def _construct_compressed2d_jl_object(
@@ -288,22 +299,27 @@ def _construct_compressed2d_jl_object(
 
         lvl = jl.Element(dtype(fill_value), data)
         jl_data = jl.swizzle(
-            jl.Tensor(jl.Dense(jl.SparseList(lvl, shape[0], indptr, indices), shape[1])), *order
+            jl.Tensor(
+                jl.Dense(jl.SparseList(lvl, shape[0], indptr, indices), shape[1])
+            ),
+            *order,
         )
         return jl_data
 
     @classmethod
-    def construct_csc_jl_object(cls, arg: TupleOf3Arrays, shape: tuple[int, ...]) -> JuliaObj:
-        return cls._construct_compressed2d_jl_object(
-            arg=arg, shape=shape, order=(1, 2)
-        )
+    def construct_csc_jl_object(
+        cls, arg: TupleOf3Arrays, shape: tuple[int, ...]
+    ) -> JuliaObj:
+        return cls._construct_compressed2d_jl_object(arg=arg, shape=shape, order=(1, 2))
 
     @classmethod
     def construct_csc(cls, arg: TupleOf3Arrays, shape: tuple[int, ...]) -> "Tensor":
         return Tensor(cls.construct_csc_jl_object(arg, shape))
 
     @classmethod
-    def construct_csr_jl_object(cls, arg: TupleOf3Arrays, shape: tuple[int, ...]) -> JuliaObj:
+    def construct_csr_jl_object(
+        cls, arg: TupleOf3Arrays, shape: tuple[int, ...]
+    ) -> JuliaObj:
         return cls._construct_compressed2d_jl_object(
             arg=arg, shape=shape[::-1], order=(2, 1)
         )
@@ -331,7 +347,9 @@ def construct_csf_jl_object(
         for size, indices, indptr in zip(shape[:-1], indices_list, indptr_list):
             lvl = jl.SparseList(lvl, size, indptr, indices)
 
-        jl_data = jl.swizzle(jl.Tensor(jl.Dense(lvl, shape[-1])), *range(1, len(shape) + 1))
+        jl_data = jl.swizzle(
+            jl.Tensor(jl.Dense(lvl, shape[-1])), *range(1, len(shape) + 1)
+        )
         return jl_data
 
     @classmethod
@@ -377,7 +395,9 @@ def _slice_plus_one(s: slice, size: int) -> range:
 
     if s.stop is not None:
         stop_offset = 2 if step < 0 else 0
-        stop = normalize_axis_index(s.stop, size) + stop_offset if s.stop < size else size
+        stop = (
+            normalize_axis_index(s.stop, size) + stop_offset if s.stop < size else size
+        )
     else:
         stop = stop_default
 
@@ -429,6 +449,7 @@ def _expand_ellipsis(key: tuple, shape: tuple[int, ...]) -> tuple:
         key = new_key
     return key
 
+
 def _add_missing_dims(key: tuple, shape: tuple[int, ...]) -> tuple:
     for i in range(len(key), len(shape)):
         key = key + (jl.range(start=1, stop=shape[i]),)