Merge remote-tracking branch 'upstream/master'

fealpy/experimental/fem/block_form.py

@@ -8,7 +8,6 @@
 from ..sparse import COOTensor, CSRTensor
 
 class BlockForm(Form):
-    #为什么有这一行
     _M = None
 
     def __init__(self, blocks:List):
@@ -37,9 +36,11 @@ def assembly(self, format='coo'):
         row_offset = bm.concatenate((bm.array([0]),row_offset))
         col_offset = bm.concatenate((bm.array([0]),col_offset))
 
-        # type类型处理
-        indices = bm.empty((2, 0), dtype=bm.int32)
-        values = bm.empty((0,), dtype=bm.float64)
+        for j in range(self.ncols):
+            block = self.blocks[0][j]
+            if block is not None:
+                indices = bm.empty((2, 0), dtype=block.space.mesh.itype)
+                values = bm.empty((0,), dtype=block.space.mesh.ftype)
         sparse_shape = self.shape
 
         for i in range(self.nrows):
@@ -48,7 +49,7 @@ def assembly(self, format='coo'):
                 if block is None:
                     continue
                 ## 不用assemble的方法
-                block_matrix = block.assembly()
+                block_matrix = block.assembly().tocoo()
                 block_indices = block_matrix.indices() + bm.array([[row_offset[i]], [col_offset[j]]])
                 block_values = block_matrix.values() 
                 indices = bm.concatenate((indices, block_indices), axis=1)
@@ -72,16 +73,18 @@ def __matmul__(self, u: TensorLike):
         kwargs = bm.context(u)
         v = bm.zeros_like(u, **kwargs)
 
-        row_offset = bm.cumsum(bm.max(self.block_shape[...,0],axis = 1))
-        col_offset = bm.cumsum(bm.max(self.block_shape[...,1],axis = 0))
-        row_offset = bm.concatenate(([0],row_offset))
-        col_offset = bm.concatenate(([0],col_offset))
+        row_offset = bm.cumsum(bm.max(self.block_shape[...,0],axis = 1), axis=0)
+        col_offset = bm.cumsum(bm.max(self.block_shape[...,1],axis = 0), axis=0)
+        row_offset = bm.concatenate((bm.array([0]),row_offset))
+        col_offset = bm.concatenate((bm.array([0]),col_offset))
         for i in range(self.nrows):
             for j in range(self.ncols):
                 block = self.blocks[i][j]
                 if block is None:
                     continue
-                v[row_offset[i]:row_offset[i+1]] += block @ u[row_offset[i]:row_offset[i+1]] 
+                v = bm.index_add(v, bm.arange(row_offset[i],row_offset[i+1]), 
+                                 block @ u[row_offset[j]:row_offset[j+1]] )
+                #v[row_offset[i]:row_offset[i+1]] += block @ u[row_offset[j]:row_offset[j+1]] 
         return v
 
 

fealpy/experimental/sparse/coo_tensor.py

@@ -202,32 +202,6 @@ def coalesce(self, accumulate: bool=True) -> 'COOTensor':
 
         return COOTensor(new_indices, new_values, self.sparse_shape, is_coalesced=True)
 
-        # unique_indices, inverse_indices = bm.unique(
-        #     self._indices, return_inverse=True, axis=1
-        # )
-
-        # if self._values is not None:
-        #     value_shape = self.dense_shape + (unique_indices.shape[-1], )
-        #     new_values = bm.zeros(value_shape, **self.values_context())
-        #     new_values = bm.index_add(new_values, inverse_indices, self._values, axis=-1)
-
-        #     return COOTensor(
-        #         unique_indices, new_values, self.sparse_shape, is_coalesced=True
-        #     )
-
-        # else:
-        #     if accumulate:
-        #         kwargs = bm.context(self._indices)
-        #         ones = bm.ones((self.nnz, ), **kwargs)
-        #         new_values = bm.zeros((unique_indices.shape[-1], ), **kwargs)
-        #         new_values = bm.index_add(new_values, inverse_indices, ones, axis=-1)
-        #     else:
-        #         new_values = None
-
-        #     return COOTensor(
-        #         unique_indices, new_values, self.sparse_shape, is_coalesced=True
-        #     )
-
     @overload
     def reshape(self, shape: Size, /) -> 'COOTensor': ...
     @overload
@@ -265,9 +239,24 @@ def T(self):
                              f"but got {self.sparse_ndim}")
         return COOTensor(new_indices, self._values, shape)
 
-    def tril(self, k: int=0) -> 'COOTensor':
-        indices, values = tril_coo(self._indices, self._values, k)
-        return COOTensor(indices, values, self._spshape)
+    def partial(self, index: Union[TensorLike, slice], /):
+        new_indices = bm.copy(self.indices()[:, index])
+        new_values = self.values()
+
+        if new_values is not None:
+            new_values = bm.copy(new_values[..., index])
+
+        return COOTensor(new_indices, new_values, self._spshape)
+
+    def tril(self, k: int = 0) -> 'COOTensor':
+        indices = self.indices()
+        tril_loc = (indices[-2] + k) >= indices[-1]
+        return self.partial(tril_loc)
+
+    def triu(self, k: int = 0) -> 'COOTensor':
+        indices = self.indices()
+        triu_loc = (indices[-1] - k) >= indices[-2]
+        return self.partial(triu_loc)
 
     @classmethod
     def concat(cls, coo_tensors: Sequence['COOTensor'], /, *, axis: int=0) -> 'COOTensor':

fealpy/experimental/sparse/csr_tensor.py

@@ -91,6 +91,15 @@ def crow(self) -> TensorLike:
         """Return the row location of non-zero elements."""
         return self._crow
 
+    def row(self) -> TensorLike:
+        """Generate the row id of non-zero elements."""
+        crow = self.crow()
+        n_row = crow.shape[0] - 1
+        return bm.repeat(
+            bm.arange(n_row, dtype=crow.dtype, device=bm.get_device(crow)),
+            crow[1:] - crow[:-1]
+        )
+
     def col(self) -> TensorLike:
         """Return the column of non-zero elements."""
         return self._col
@@ -194,6 +203,35 @@ def ravel(self) -> 'CSRTensor':
     def flatten(self) -> 'CSRTensor':
         pass
 
+    @property
+    def T(self):
+        raise NotImplementedError
+
+    def partial(self, index: Union[TensorLike, slice]):
+        crow = self.crow()
+        ZERO = bm.zeros([1], dtype=crow.dtype, device=bm.get_device(crow))
+        new_col = bm.copy(self.col()[..., index])
+        is_selected = bm.zeros((self.nnz,), dtype=bm.bool, device=bm.get_device(new_col))
+        is_selected = bm.set_at(is_selected, index, True)
+        selected_cum = bm.concat([ZERO, bm.cumsum(is_selected, axis=0)], axis=0)
+        new_nnz_per_row = selected_cum[crow[1:]] - selected_cum[crow[:-1]]
+        new_crow = bm.concat([ZERO, bm.cumsum(new_nnz_per_row, axis=0)], axis=0)
+
+        new_values = self.values()
+
+        if new_values is not None:
+            new_values = bm.copy(new_values[..., index])
+
+        return CSRTensor(new_crow, new_col, new_values, self._spshape)
+
+    def tril(self, k: int = 0) -> 'CSRTensor':
+        tril_loc = (self.row() + k) >= self.col()
+        return self.partial(tril_loc)
+
+    def triu(self, k: int = 0) -> 'CSRTensor':
+        tril_loc = (self.col() - k) >= self.row()
+        return self.partial(tril_loc)
+
     ### 6. Arithmetic Operations ###
     def neg(self) -> 'CSRTensor':
         """Negation of the CSR tensor. Returns self if values is None."""

fealpy/experimental/tests/fem/test_block_form.py

@@ -13,7 +13,7 @@
 from fealpy.experimental.tests.fem.block_form_data import *
 
 class TestBlockForm:
-    @pytest.mark.parametrize("backend", ['pytorch', 'numpy'])
+    @pytest.mark.parametrize("backend", ['numpy','pytorch'])
     @pytest.mark.parametrize("data", diag_diffusion)
     def test_diag_diffusion(self, backend, data): 
         bm.set_backend(backend)
@@ -29,41 +29,20 @@ def test_diag_diffusion(self, backend, data):
 
         blockform = BlockForm([[bform0,None],[None,bform1]])
 
+        a = bm.arange(blockform.shape[1], dtype=mesh.ftype)
+        vector = blockform @ a
+
         matrix = blockform.assembly() 
-        true_matrix = COOTensor(bm.array(data["indices"]), bm.array(data["values"]), blockform.shape)
-
+        true_matrix = COOTensor(bm.array(data["indices"],dtype=mesh.itype), bm.array(data["values"], dtype=mesh.ftype), blockform.shape) 
+
+        true_vector = true_matrix.toarray() @ a
+
         np.testing.assert_array_almost_equal(matrix.toarray(), true_matrix.toarray(), 
                                      err_msg=f" `blockform` function is not equal to real result in backend {backend}")
+        np.testing.assert_array_almost_equal(vector, true_vector, 
+                                     err_msg=f" `blockform __mult__` function is not equal to real result in backend {backend}")
 
 if __name__ == "__main__":
-    pytest.main(['./test_block_form.py', '-k', 'test_diag_diffusion'])
-
-
-'''
-bm.set_backend('numpy')
-#bm.set_backend('pytorch')
-
-mesh = TriangleMesh.from_box(box=[0, 1, 0, 1], nx=1, ny=1)
-space = LagrangeFESpace(mesh, p=1)
-space1 = LagrangeFESpace(mesh, p=2)
-
-bform0 = BilinearForm(space)
-bform0.add_integrator(ScalarDiffusionIntegrator())
-A1 = bform0.assembly()
-bform1 = BilinearForm(space1)
-bform1.add_integrator(ScalarDiffusionIntegrator())
-A2 = bform1.assembly()
+    pytest.main(['./test_block_form.py', '-sk', 'test_diag_diffusion'])
 
-from scipy.sparse import coo_array, bmat
-def coo(A):
-    data = A._values
-    indices = A._indices
-    return coo_array((data, indices))
 
-A = bmat([[coo(A1), None],
-        [None, coo(A2)]], format='coo')
-A = COOTensor(bm.stack([A.row,A.col],axis=0), A.data, spshape=A.shape)
-print(A.indices())
-print(A.values())
-print(A)
-'''

fealpy/experimental/tests/sparse/test_csr_tensor.py

@@ -56,6 +56,24 @@ def test_to_dense(backend):
                     [1.22, 0, 0]]], dtype=bm.float64)
     )
 
+@pytest.mark.parametrize("backend", ALL_BACKENDS)
+def test_tril(backend):
+    bm.set_backend(backend)
+    crow = bm.tensor([0, 0, 2, 3, 4, 4])
+    col = bm.tensor([1, 2, 3, 2])
+    values = bm.tensor([4, 3, 1, 2], dtype=bm.float32)
+    spshape = (5, 4)
+    csr_tensor = CSRTensor(crow, col, values, spshape)
+    tril_tensor = csr_tensor.tril(k=0)
+
+    expected_crow = bm.tensor([0, 0, 1, 1, 2, 2])
+    expected_col = bm.tensor([1, 2])
+    expected_values = bm.tensor([4, 2], dtype=bm.float32)
+
+    assert bm.all(bm.equal(tril_tensor.crow(), expected_crow))
+    assert bm.all(bm.equal(tril_tensor.col(), expected_col))
+    assert bm.allclose(tril_tensor.values(), expected_values)
+
 def create_csr_tensor(crow, col, values, shape):
     return CSRTensor(crow=crow, col=col, values=values, spshape=shape)