Various changes

Author: connorjward

pyop3/__init__.py

@@ -34,7 +34,7 @@
     Subset,
     TabulatedMapComponent,
 )
-from pyop3.itree.tree import ScalarIndex
+from pyop3.itree.tree import ScalarIndex, as_index_forest
 from pyop3.lang import (  # noqa: F401
     INC,
     MAX_RW,

pyop3/array/base.py

@@ -1,10 +1,10 @@
 import abc
 
-from pyop3.lang import KernelArgument, ReplaceAssignment
+from pyop3.lang import FunctionArgument, ReplaceAssignment
 from pyop3.utils import UniqueNameGenerator
 
 
-class Array(KernelArgument, abc.ABC):
+class Array(FunctionArgument, abc.ABC):
     _prefix = "array"
     _name_generator = UniqueNameGenerator()
 

pyop3/array/harray.py

@@ -139,19 +139,23 @@ def __init__(
 
             if data is not None:
                 data = np.asarray(data, dtype=dtype)
-                shape = data.shape
-            else:
-                shape = axes.global_size
+
+                # always deal with flattened data
+                if len(data.shape) > 1:
+                    data = data.flatten()
+                if data.size != axes.unindexed.global_size:
+                    raise ValueError("Data shape does not match axes")
 
             # IndexedAxisTrees do not currently have SFs, so create a dummy one here
             if isinstance(axes, AxisTree):
                 sf = axes.sf
             else:
                 assert isinstance(axes, IndexedAxisTree)
-                sf = serial_forest(axes.global_size)
+                # not sure this is the right thing to do
+                sf = serial_forest(axes.unindexed.global_size)
 
             data = DistributedBuffer(
-                shape,
+                axes.unindexed.global_size,  # not a useful property anymore
                 sf,
                 dtype,
                 name=self.name,
@@ -165,6 +169,8 @@ def __init__(
         # TODO This attr really belongs to the buffer not the array
         self.constant = constant
 
+        # self._cache = {}
+
     def __str__(self):
         return self.name
 
@@ -177,14 +183,20 @@ def getitem(self, indices, *, strict=False):
         if indices is Ellipsis:
             return self
 
+        # key = (indices, strict)
+        # if key in self._cache:
+        #     return self._cache[key]
+
         index_forest = as_index_forest(indices, axes=self.axes, strict=strict)
         if index_forest.keys() == {pmap()}:
             index_tree = index_forest[pmap()]
             indexed_axes = index_axes(index_tree, pmap(), self.axes)
             axes = compose_axes(indexed_axes, self.axes)
-            return HierarchicalArray(
+            dat = HierarchicalArray(
                 axes, data=self.buffer, max_value=self.max_value, name=self.name
             )
+            # self._cache[key] = dat
+            return dat
 
         array_per_context = {}
         for loop_context, index_tree in index_forest.items():
@@ -194,7 +206,9 @@ def getitem(self, indices, *, strict=False):
                 axes, data=self.buffer, name=self.name, max_value=self.max_value
             )
 
-        return ContextSensitiveMultiArray(array_per_context)
+        dat = ContextSensitiveMultiArray(array_per_context)
+        # self._cache[key] = dat
+        return dat
 
     # Since __getitem__ is implemented, this class is implicitly considered
     # to be iterable (which it's not). This avoids some confusing behaviour.
@@ -218,16 +232,16 @@ def data(self):
     @property
     def data_rw(self):
         self._check_no_copy_access()
-        return self.buffer.data_rw[self._buffer_indices]
+        return self.buffer.data_rw[self.axes._buffer_indices]
 
     @property
     def data_ro(self):
-        if not isinstance(self._buffer_indices, slice):
+        if not isinstance(self.axes._buffer_indices, slice):
             warning(
                 "Read-only access to the array is provided with a copy, "
                 "consider avoiding if possible."
             )
-        return self.buffer.data_ro[self._buffer_indices]
+        return self.buffer.data_ro[self.axes._buffer_indices]
 
     @property
     def data_wo(self):
@@ -239,7 +253,7 @@ def data_wo(self):
         can be dropped.
         """
         self._check_no_copy_access()
-        return self.buffer.data_wo[self._buffer_indices]
+        return self.buffer.data_wo[self.axes._buffer_indices]
 
     @property
     @deprecated(".data_rw_with_halos")
@@ -249,16 +263,16 @@ def data_with_halos(self):
     @property
     def data_rw_with_halos(self):
         self._check_no_copy_access()
-        return self.buffer.data_rw[self._buffer_indices_ghost]
+        return self.buffer.data_rw[self.axes._buffer_indices_ghost]
 
     @property
     def data_ro_with_halos(self):
-        if not isinstance(self._buffer_indices_ghost, slice):
+        if not isinstance(self.axes._buffer_indices_ghost, slice):
             warning(
                 "Read-only access to the array is provided with a copy, "
                 "consider avoiding if possible."
             )
-        return self.buffer.data_ro[self._buffer_indices_ghost]
+        return self.buffer.data_ro[self.axes._buffer_indices_ghost]
 
     @property
     def data_wo_with_halos(self):
@@ -270,54 +284,10 @@ def data_wo_with_halos(self):
         can be dropped.
         """
         self._check_no_copy_access()
-        return self.buffer.data_wo[self._buffer_indices_ghost]
-
-    @cached_property
-    def _buffer_indices(self):
-        return self._collect_buffer_indices(ghost=False)
-
-    @cached_property
-    def _buffer_indices_ghost(self):
-        return self._collect_buffer_indices(ghost=True)
-
-    def _collect_buffer_indices(self, *, ghost: bool):
-        # TODO: This method is inefficient as for affine things we still tabulate
-        # everything first. It would be best to inspect index_exprs to determine
-        # if a slice is sufficient, but this is hard.
-        # TODO: This should be more widely cached, don't want to tabulate more often
-        # than required.
-
-        size = self.axes.size if ghost else self.axes.owned.size
-        assert size > 0
-
-        indices = np.full(size, -1, dtype=IntType)
-        # TODO: Handle any outer loops.
-        # TODO: Generate code for this.
-        for i, p in enumerate(self.axes.iter()):
-            indices[i] = self.axes.offset(p.source_exprs, p.source_path)
-        debug_assert(lambda: (indices >= 0).all())
-
-        # The packed indices are collected component-by-component so, for
-        # numbered multi-component axes, they are not in ascending order.
-        # We sort them so we can test for "affine-ness".
-        indices.sort()
-
-        # See if we can represent these indices as a slice. This is important
-        # because slices enable no-copy access to the array.
-        steps = np.unique(indices[1:] - indices[:-1])
-        if len(steps) == 0:
-            start = just_one(indices)
-            return slice(start, start + 1, 1)
-        elif len(steps) == 1:
-            start = indices[0]
-            stop = indices[-1] + 1
-            (step,) = steps
-            return slice(start, stop, step)
-        else:
-            return indices
+        return self.buffer.data_wo[self.axes._buffer_indices_ghost]
 
     def _check_no_copy_access(self):
-        if not isinstance(self._buffer_indices, slice):
+        if not isinstance(self.axes._buffer_indices, slice):
             raise FancyIndexWriteException(
                 "Writing to the array directly is not supported for "
                 "non-trivially indexed (i.e. sliced) arrays."
@@ -541,7 +511,8 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
 
-# Now ContextSensitiveDat
+# NOTE: I think I can probably get rid of this class and wrap the
+# context-sensitivity inside the axis tree.
 class ContextSensitiveMultiArray(Array, ContextSensitive):
     def __init__(self, arrays):
         name = single_valued(a.name for a in arrays.values())
@@ -596,6 +567,11 @@ def __getitem__(self, indices) -> ContextSensitiveMultiArray:
     def buffer(self):
         return self._shared_attr("buffer")
 
+    # this is really nasty, but need to know if wrapping a Mat
+    @property
+    def mat(self):
+        return self._shared_attr("mat")
+
     @property
     def dtype(self):
         return self._shared_attr("dtype")

pyop3/array/petsc.py

@@ -104,15 +104,7 @@ def __init__(
         self.mat_type = mat_type
         self.mat = mat
 
-        # TODO: delete
-        # self.rtarget_paths = rtarget_paths
-        # self.rindex_exprs = rindex_exprs
-        # self.orig_raxes = orig_raxes
-        # self.router_loops = router_loops
-        # self.ctarget_paths = ctarget_paths
-        # self.cindex_exprs = cindex_exprs
-        # self.orig_caxes = orig_caxes
-        # self.couter_loops = couter_loops
+        # self._cache = {}
 
     def __getitem__(self, indices):
         return self.getitem(indices, strict=False)
@@ -122,6 +114,11 @@ def __getitem__(self, indices):
     __iter__ = None
 
     def getitem(self, indices, *, strict=False):
+        # does not work as indices may not be hashable, parse first?
+        # cache_key = (indices, strict)
+        # if cache_key in self._cache:
+        #     return self._cache[cache_key]
+
         if len(indices) != 2:
             raise ValueError
 
@@ -174,13 +171,15 @@ def getitem(self, indices, *, strict=False):
             indexed_caxes = index_axes(ctree, pmap(), self.caxes)
             caxes = compose_axes(indexed_caxes, self.caxes)
 
-            return type(self)(
+            mat = type(self)(
                 raxes,
                 caxes,
                 mat_type=self.mat_type,
                 mat=self.mat,
                 name=self.name,
             )
+            # self._cache[cache_key] = mat
+            return mat
 
         # Otherwise we are context-sensitive
         arrays = {}
@@ -202,7 +201,9 @@ def getitem(self, indices, *, strict=False):
                 name=self.name,
             )
         # But this is now a PetscMat...
-        return ContextSensitiveMultiArray(arrays)
+        mat = ContextSensitiveMultiArray(arrays)
+        # self._cache[cache_key] = mat
+        return mat
 
     # like Dat, bad name? handle?
     @property
@@ -305,7 +306,9 @@ def _iter_nest_labels(
                 yield (rlabel_acc_, clabel_acc_)
 
     @cached_property
+    @PETSc.Log.EventDecorator()
     def maps(self):
+        print("HIT!")
         from pyop3.axtree.layout import my_product
 
         # TODO: Don't think these need to be lists here.
@@ -342,11 +345,11 @@ def maps(self):
                     ].items()
                     if ax == cfield_axis.label
                 )
-                orig_caxes = AxisTree(self.orig_caxes[cfield].node_map)
+                orig_caxes = AxisTree(self.caxes.unindexed[cfield].node_map)
                 orig_caxess = [orig_caxes]
                 dropped_ckeys = {cfield_axis.label}
             else:
-                orig_caxess = [self.orig_caxes]
+                orig_caxess = [self.caxes.unindexed]
                 dropped_ckeys = set()
         else:
             orig_raxess = [self.raxes.unindexed]
@@ -426,6 +429,22 @@ def rmap(self):
     def cmap(self):
         return self.maps[1]
 
+    @cached_property
+    def row_lgmap_dat(self):
+        if self.nested or self.mat_type == "baij":
+            raise NotImplementedError("Use a smaller set of axes here")
+        return HierarchicalArray(self.raxes, data=self.raxes.unindexed.global_numbering)
+
+    @cached_property
+    def column_lgmap_dat(self):
+        if self.nested or self.mat_type == "baij":
+            raise NotImplementedError("Use a smaller set of axes here")
+        return HierarchicalArray(self.caxes, data=self.caxes.unindexed.global_numbering)
+
+    @cached_property
+    def comm(self):
+        return single_valued([self.raxes.comm, self.caxes.comm])
+
     @property
     def shape(self):
         return (self.raxes.size, self.caxes.size)
@@ -447,15 +466,6 @@ def axes(self):
     @classmethod
     def _make_mat(cls, raxes, caxes, mat_type):
         if isinstance(mat_type, collections.abc.Mapping):
-            # if strictly_all(c.unit for c in raxes.root.components):
-            #     riter = tuple((c.label, raxes[c.label]) for c in raxes.root.components)
-            # else:
-            #     riter = [(None, raxes)]
-            # if strictly_all(c.unit for c in caxes.root.components):
-            #     citer = tuple((c.label, caxes[c.label]) for c in caxes.root.components)
-            # else:
-            #     citer = [(None, caxes)]
-
             # TODO: This is very ugly
             rsize = max(x or 0 for x, _ in mat_type.keys()) + 1
             csize = max(y or 0 for _, y in mat_type.keys()) + 1
@@ -482,10 +492,6 @@ def kernel_dtype(self):
 
 
 class Sparsity(AbstractMat):
-    # def __init__(self, *args, **kwargs):
-    #     super().__init__(*args, **kwargs)
-    #     self._lazy_template = None
-
     def materialize(self) -> PETSc.Mat:
         if not hasattr(self, "_lazy_template"):
             self.assemble()
@@ -495,6 +501,7 @@ def materialize(self) -> PETSc.Mat:
             # template.preallocateWithMatPreallocator(self.mat)
             # We can safely set these options since by using a sparsity we
             # are asserting that we know where the non-zeros are going.
+            # NOTE: These may already get set by PETSc.
             template.setOption(PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR, True)
             template.setOption(PETSc.Mat.Option.IGNORE_ZERO_ENTRIES, True)
 
@@ -537,15 +544,12 @@ def _make_monolithic_mat(cls, raxes, caxes, mat_type: str):
             mat = PETSc.Mat().create(comm)
             mat.setType(PETSc.Mat.Type.PREALLOCATOR)
 
-            #     breakpoint()
-            # else:
             # None is for the global size, PETSc will figure it out for us
             sizes = ((raxes.owned.size, None), (caxes.owned.size, None))
-
             mat.setSizes(sizes)
 
-            rlgmap = PETSc.LGMap().create(raxes.global_numbering(), comm=comm)
-            clgmap = PETSc.LGMap().create(caxes.global_numbering(), comm=comm)
+            rlgmap = PETSc.LGMap().create(raxes.global_numbering, comm=comm)
+            clgmap = PETSc.LGMap().create(caxes.global_numbering, comm=comm)
             mat.setLGMap(rlgmap, clgmap)
 
         mat.setUp()
@@ -613,8 +617,8 @@ def _make_monolithic_mat(cls, raxes, caxes, mat_type: str):
             sizes = ((raxes.owned.size, None), (caxes.owned.size, None))
             mat.setSizes(sizes)
 
-            rlgmap = PETSc.LGMap().create(raxes.global_numbering(), comm=comm)
-            clgmap = PETSc.LGMap().create(caxes.global_numbering(), comm=comm)
+            rlgmap = PETSc.LGMap().create(raxes.global_numbering, comm=comm)
+            clgmap = PETSc.LGMap().create(caxes.global_numbering, comm=comm)
             mat.setLGMap(rlgmap, clgmap)
 
         mat.setUp()