Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

To array and to sparse functions are ready for linearoperator #368

Open
wants to merge 9 commits into
base: devel
Choose a base branch
from
359 changes: 339 additions & 20 deletions psydac/linalg/basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@
from abc import ABC, abstractmethod
from scipy.sparse import coo_matrix
import numpy as np
import itertools
from scipy import sparse
from psydac.linalg.kernels.toarray_kernels import write_out_stencil_dense_3D, write_out_stencil_dense_2D, write_out_stencil_dense_1D, write_out_block_dense_3D, write_out_block_dense_2D, write_out_block_dense_1D

__all__ = ('VectorSpace', 'Vector', 'LinearOperator', 'ZeroOperator', 'IdentityOperator', 'ScaledLinearOperator',
'SumLinearOperator', 'ComposedLinearOperator', 'PowerLinearOperator', 'InverseLinearOperator', 'LinearSolver')
Expand Down Expand Up @@ -220,15 +223,346 @@ def codomain(self):
@abstractmethod
def dtype(self):
pass

def __tosparse_array(self, out=None, is_sparse=False):
"""
Transforms the linear operator into a matrix, which is either stored in dense or sparse format.

@abstractmethod
Parameters
----------
out : Numpy.ndarray, optional
If given, the output will be written in-place into this array.
is_sparse : bool, optional
If set to True the method returns the matrix as a Scipy sparse matrix, if set to false
it returns the full matrix as a Numpy.ndarray

Returns
-------
out : Numpy.ndarray or scipy.sparse.csr.csr_matrix
The matrix form of the linear operator. If ran in parallel each rank gets the full
matrix representation of the linear operator.
"""
# v will be the unit vector with which we compute Av = ith column of A.
v = self.domain.zeros()
# We define a temporal vector
tmp2 = self.codomain.zeros()

#We need to determine if we are a blockvector or a stencilvector but we are not able to use
#the BlockVectorSpace and StencilVectorSpace classes in here. So we check if domain has the spaces
#attribute in which case the domain would be a BlockVectorSpace. If that is not the case we check
#if the domain has the cart atrribute, in which case it will be a StencilVectorSpace.
if hasattr(self.domain, 'spaces'):
BoS = "b"
elif hasattr(self.domain, 'cart'):
BoS = "s"
else:
raise Exception(
'The domain of the LinearOperator must be a BlockVectorSpace or a StencilVectorSpace.')

#We also need to know if the codomain is a StencilVectorSpace or a BlockVectorSpace
if hasattr(self.codomain, 'spaces'):
BoS2 = "b"
elif hasattr(self.codomain, 'cart'):
BoS2 = "s"
else:
raise Exception(
'The codomain of the LinearOperator must be a BlockVectorSpace or a StencilVectorSpace.')

if BoS == "b":
comm = self.domain.spaces[0].cart.comm
elif BoS == "s":
comm = self.domain.cart.comm
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If it is not a Blockvector or a Stencilvector then comm is not defined. I would add an assert in the beginning.

rank = comm.Get_rank()
size = comm.Get_size()

if (is_sparse == False):
if out is None:
# We declare the matrix form of our linear operator
out = np.zeros(
[self.codomain.dimension, self.domain.dimension], dtype=self.dtype)
else:
assert isinstance(out, np.ndarray)
assert out.shape[0] == self.codomain.dimension
assert out.shape[1] == self.domain.dimension
else:
if out is not None:
raise Exception(
'If is_sparse is True then out must be set to None.')
numrows = self.codomain.dimension
numcols = self.domain.dimension
# We define a list to store the non-zero data, a list to sotre the row index of said data and a list to store the column index.
data = []
row = []
colarr = []

# V is either a BlockVector or a StencilVector depending on the domain of the linear operator.
if BoS == "b":
# we collect all starts and ends in two big lists
starts = [vi.starts for vi in v]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it is possible to use the same code for both cases BoS == 'b' and BoS == 's' with some small changes, for instance:
starts = [vi.starts] if BoS == "s" else [vi.starts for vi in v]

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Although there are some similarities between the code for the BlockVectorSpaces and the StencilVectorSpaces, there are also a sizeable number of differences. Therefore, getting rid of the big if(BoS =="b")- elif(BoS =="s") would result in having to write a myriad of more tiny if-elif pairs. Which would negatively affect the readability of the code.

ends = [vi.ends for vi in v]
# We collect the dimension of the BlockVector
npts = [sp.npts for sp in self.domain.spaces]
# We get the number of space we have
nsp = len(self.domain.spaces)
# We get the number of dimensions each space has.
ndim = [sp.ndim for sp in self.domain.spaces]
elif BoS == "s":
# We get the start and endpoint for each sublist in v
starts = [v.starts]
ends = [v.ends]
# We get the dimensions of the StencilVector
npts = [self.domain.npts]
# We get the number of space we have
nsp = 1
# We get the number of dimensions the StencilVectorSpace has.
ndim = [self.domain.ndim]

# First each rank is going to need to know the starts and ends of all other ranks
startsarr = np.array([starts[i][j] for i in range(nsp)
for j in range(ndim[i])], dtype=int)

endsarr = np.array([ends[i][j] for i in range(nsp)
for j in range(ndim[i])], dtype=int)

# Create an array to store gathered data from all ranks
allstarts = np.empty(size * len(startsarr), dtype=int)

# Use Allgather to gather 'starts' from all ranks into 'allstarts'
comm.Allgather(startsarr, allstarts)

# Reshape 'allstarts' to have 9 columns and 'size' rows
allstarts = allstarts.reshape((size, len(startsarr)))

# Create an array to store gathered data from all ranks
allends = np.empty(size * len(endsarr), dtype=int)

# Use Allgather to gather 'ends' from all ranks into 'allends'
comm.Allgather(endsarr, allends)

# Reshape 'allends' to have 9 columns and 'size' rows
allends = allends.reshape((size, len(endsarr)))


# Before we begin computing the dot products we need to know which entries of the output vector tmp2 belong to our rank.
if BoS2 == "s":
# We get the start and endpoint for each sublist in tmp2
starts2 = tmp2.starts
ends2 = tmp2.ends
# We get the dimensions of the StencilVector
npts2 = np.array(self.codomain.npts)
# We get the number of space we have
nsp2 = 1
# We get the number of dimensions the StencilVectorSpace has.
ndim2 = self.codomain.ndim
#We build our ranges of iteration
if (is_sparse == False):
itterables2 = []
for ii in range(ndim2):
itterables2.append([starts2[ii], ends2[ii]+1])
#itterables2.append(range(starts2[ii], ends2[ii]+1))
itterables2 = np.array(itterables2)
#We also get the StencilVector's pads
pds = np.array(tmp2.pads)
else:
itterables2 = []
for ii in range(ndim2):
itterables2.append(
range(starts2[ii], ends2[ii]+1))

elif BoS2 == "b":
# we collect all starts and ends in two big lists
starts2 = [vi.starts for vi in tmp2]
ends2 = [vi.ends for vi in tmp2]
# We collect the dimension of the BlockVector
npts2 = np.array([sp.npts for sp in self.codomain.spaces])
# We get the number of space we have
nsp2 = len(self.codomain.spaces)
# We get the number of dimensions each space has.
ndim2 = [sp.ndim for sp in self.codomain.spaces]
if (is_sparse == False):
#We also get the BlockVector's pads
pds = np.array([vi.pads for vi in tmp2])
#We build the range of iteration
itterables2 = []
# since the size of npts changes denpending on h we need to compute a starting point for
# our row index
spoint2 = 0
spoint2list = [np.int64(spoint2)]
for hh in range(nsp2):
itterables2aux = []
for ii in range(ndim2[hh]):
itterables2aux.append(
[starts2[hh][ii], ends2[hh][ii]+1])
itterables2.append(itterables2aux)
cummulative2 = 1
for ii in range(ndim2[hh]):
cummulative2 *= npts2[hh][ii]
spoint2 += cummulative2
spoint2list.append(spoint2)
else:
itterables2 = []
# since the size of npts changes denpending on h we need to compute a starting point for
# our row index
spoint2 = 0
spoint2list = [spoint2]
for hh in range(nsp2):
itterables2aux = []
for ii in range(ndim2[hh]):
itterables2aux.append(
range(starts2[hh][ii], ends2[hh][ii]+1))
itterables2.append(itterables2aux)
cummulative2 = 1
for ii in range(ndim2[hh]):
cummulative2 *= npts2[hh][ii]
spoint2 += cummulative2
spoint2list.append(spoint2)


currentrank = 0
# Each rank will take care of setting to 1 each one of its entries while all other entries remain zero.
while (currentrank < size):
# since the size of npts changes denpending on h we need to compute a starting point for
# our column index
spoint = 0
npredim = 0
# We iterate over the stencil vectors inside the BlockVector
for h in range(nsp):
itterables = []
for i in range(ndim[h]):
itterables.append(
range(allstarts[currentrank][i+npredim], allends[currentrank][i+npredim]+1))
# We iterate over all the entries that belong to rank number currentrank
for i in itertools.product(*itterables):

#########################################
if BoS == "b":
if (rank == currentrank):
v[h][i] = 1.0
v[h].update_ghost_regions()
elif BoS == "s":
if (rank == currentrank):
v[i] = 1.0
v.update_ghost_regions()
#########################################

# Compute dot product with the linear operator.
self.dot(v, out=tmp2)
# Compute to which column this iteration belongs
col = spoint
col += np.ravel_multi_index(i, npts[h])

# Case in which tmp2 is a StencilVector
if BoS2 == "s":
if is_sparse == False:
#We iterate over the entries of tmp2 that belong to our rank
#The pyccel kernels are tantamount to this for loop.
#for ii in itertools.product(*itterables2):
#if (tmp2[ii] != 0):
#out[np.ravel_multi_index(
#ii, npts2), col] = tmp2[ii]
if (ndim2 == 3):
write_out_stencil_dense_3D(itterables2, tmp2._data, out, npts2, col, pds)
elif (ndim2 == 2):
write_out_stencil_dense_2D(itterables2, tmp2._data, out, npts2, col, pds)
elif (ndim2 == 1):
write_out_stencil_dense_1D(itterables2, tmp2._data, out, npts2, col, pds)
else:
raise Exception("The codomain dimension must be 3, 2 or 1.")

else:
#We iterate over the entries of tmp2 that belong to our rank
for ii in itertools.product(*itterables2):
if (tmp2[ii] != 0):
data.append(tmp2[ii])
colarr.append(col)
row.append(
np.ravel_multi_index(ii, npts2))
elif BoS2 =="b":
# We iterate over the stencil vectors inside the BlockVector
for hh in range(nsp2):


if is_sparse == False:
itterables2aux = np.array(itterables2[hh])
# We iterate over all the tmp2 entries that belong to rank number currentrank
#for ii in itertools.product(*itterables2aux):
#if (tmp2[hh][ii] != 0):
#out[spoint2list[hh]+np.ravel_multi_index(
#ii, npts2[hh]), col] = tmp2[hh][ii]
if (ndim2[hh] == 3):
write_out_block_dense_3D(itterables2aux, tmp2[hh]._data, out, npts2[hh], col, pds[hh], spoint2list[hh])
elif (ndim2[hh] == 2):
write_out_block_dense_2D(itterables2aux, tmp2[hh]._data, out, npts2[hh], col, pds[hh], spoint2list[hh])
elif (ndim2[hh] == 1):
write_out_block_dense_1D(itterables2aux, tmp2[hh]._data, out, npts2[hh], col, pds[hh], spoint2list[hh])
else:
raise Exception("The codomain dimension must be 3, 2 or 1.")
else:
itterables2aux = itterables2[hh]
for ii in itertools.product(*itterables2aux):
if (tmp2[hh][ii] != 0):
data.append(tmp2[hh][ii])
colarr.append(col)
row.append(
spoint2list[hh]+np.ravel_multi_index(ii, npts2[hh]))
#################################
if BoS == "b":
if (rank == currentrank):
v[h][i] = 0.0
v[h].update_ghost_regions()
elif BoS == "s":
if (rank == currentrank):
v[i] = 0.0
v.update_ghost_regions()
##################################
cummulative = 1
for i in range(ndim[h]):
cummulative *= npts[h][i]
spoint += cummulative
npredim += ndim[h]
currentrank += 1

if is_sparse == False:
return out
else:
return sparse.csr_matrix((data, (row, colarr)), shape=(numrows, numcols))


# Function that returns the local matrix corresponding to the linear operator. Returns a scipy.sparse.csr.csr_matrix.
def tosparse(self):
pass
"""
Transforms the linear operator into a matrix, which is stored in sparse csr format.

@abstractmethod
def toarray(self):
pass
Returns
-------
out : Numpy.ndarray or scipy.sparse.csr.csr_matrix
The matrix form of the linear operator. If ran in parallel each rank gets the local
matrix representation of the linear operator.
"""
return self.__tosparse_array(is_sparse=True)


# Function that returns the matrix corresponding to the linear operator. Returns a numpy array.
def toarray(self, out=None):
"""
Transforms the linear operator into a matrix, which is stored in dense format.

Parameters
----------
out : Numpy.ndarray, optional
If given, the output will be written in-place into this array.

Returns
-------
out : Numpy.ndarray
The matrix form of the linear operator. If ran in parallel each rank gets the local
matrix representation of the linear operator.
"""
return self.__tosparse_array(out=out, is_sparse=False)




@abstractmethod
def dot(self, v, out=None):
""" Apply linear operator to Vector v. Result is written to Vector out, if provided."""
Expand Down Expand Up @@ -748,9 +1082,6 @@ def multiplicants(self):
def dtype(self):
return None

def toarray(self):
raise NotImplementedError('toarray() is not defined for ComposedLinearOperators.')

def tosparse(self):
mats = [M.tosparse() for M in self._multiplicants]
M = mats[0]
Expand Down Expand Up @@ -849,12 +1180,6 @@ def operator(self):
def factorial(self):
return self._factorial

def toarray(self):
raise NotImplementedError('toarray() is not defined for PowerLinearOperators.')

def tosparse(self):
raise NotImplementedError('tosparse() is not defined for PowerLinearOperators.')

def transpose(self, conjugate=False):
return PowerLinearOperator(domain=self._codomain, codomain=self._domain, A=self._operator.transpose(conjugate=conjugate), n=self._factorial)

Expand Down Expand Up @@ -905,12 +1230,6 @@ def linop(self):
def options(self):
return self._options

def toarray(self):
raise NotImplementedError('toarray() is not defined for InverseLinearOperators.')

def tosparse(self):
raise NotImplementedError('tosparse() is not defined for InverseLinearOperators.')

def get_info(self):
return self._info

Expand Down
Loading