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

@@ -0,0 +1,84 @@
+import abc
+import typing
+
+from .julia import jl
+from . import levels
+from . import utils
+from dataclasses import dataclass
+
+
+@dataclass
+class Format:
+    levels: tuple[levels.AbstractLevel, ...]
+    order: tuple[int, ...]
+    leaf: levels.AbstractLeafLevel
+
+    def __init__(
+        self,
+        *,
+        levels: tuple[levels.AbstractLevel, ...],
+        order: tuple[int, ...] | None,
+        leaf: levels.AbstractLeafLevel,
+    ) -> None:
+        if order is None:
+            order = tuple(range(len(levels)))
+
+        utils.check_valid_order(order, ndim=len(levels))
+        self.order = order
+        self.levels = levels
+        self.leaf = leaf
+
+    def _construct(self, *, fill_value, dtype: jl.DataType, data=None):
+        out_level = self.leaf._construct(dtype=dtype, fill_value=fill_value)
+        for level in reversed(self.levels):
+            out_level = level._construct(inner_level=out_level)
+
+        swizzle_args = map(lambda x: x + 1, reversed(self.order))
+        if data is None:
+            return jl.swizzle(jl.Tensor(out_level), *swizzle_args)
+
+        return jl.swizzle(jl.Tensor(out_level, data), *swizzle_args)
+
+class FlexibleFormat(abc.ABC):
+    def _construct(self, *, ndim: int, fill_value, dtype: jl.DataType, data=None):
+        return self._get_format(ndim)._construct(fill_value=fill_value, dtype=dtype, data=data)
+
+    @abc.abstractmethod
+    def _get_format(self, ndim: int, /) -> Format:
+        pass
+
+@dataclass
+class Dense(FlexibleFormat):
+    order: typing.Literal["C", "F"] | tuple[int, ...] = "C"
+    shape: tuple[int | None, ...] | None = None
+
+    def __post_init__(self) -> None:
+        if isinstance(self.order, tuple):
+            utils.check_valid_order(self.order)
+
+            if self.shape is not None and len(self.order) != len(self.shape):
+                raise ValueError(f"len(self.order) != len(self.shape), {self.order}, {self.shape}")
+
+    def _get_format(self, ndim: int) -> Format:
+        super()._get_format(ndim)
+        match self.order:
+            case "C":
+                order = tuple(range(ndim))
+            case "F":
+                order = tuple(reversed(range(ndim)))
+            case _:
+                order = self.order
+
+        utils.check_valid_order(order, ndim=ndim)
+        
+        shape = self.shape
+        if shape is None:
+            shape = (None,) * ndim
+        
+        if len(shape) != ndim:
+            raise ValueError(f"len(self.shape != ndim), {shape=}, {ndim=}")
+        
+        topological_shape = utils.get_topological_shape(shape, order=order)
+        lvls = tuple(levels.Dense(dim=dim) for dim in topological_shape)
+
+        return Format(levels=lvls, order=order, leaf=levels.Element())

src/finch/formats.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)