Sparse matrix testing and added features (Bears-R-Us#3745)

* Add func to convert sparse matrix to pdarrays

We convert sparse matrices to a 3 tuple of pdarrays with row, cols, and
vals

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

* switch to using list of pdarrays

We use the MsgTuple.fromResponses function to return a list of multiple
pdarrays instead of a tupl, this minimzes the comm between the server
and client.

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

* Add fill_vals to populate a sparse array with values

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

* Add tests for sparse matrix features

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

* Add sparse_test to pytest, module in ServerModule

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

* Add sparse matrix compat module for 2.0

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

* Add missing newlines at the end of files

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

* List public object in sparrayclass with __all__

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

* Comment out non distributed version of mult

Comment out the non ditributed version of the sparse matrix
multiplication and also reduce the problem size in the test to make sure
we don't run out of memory

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

* Changes made based on feedback

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

---------

Signed-off-by: Shreyas Khandekar <60454060+ShreyasKhandekar@users.noreply.github.com>

ServerModules.cfg

@@ -30,6 +30,7 @@ SegmentedMsg
 SequenceMsg
 SetMsg
 SortMsg
+SparseMatrixMsg
 StatsMsg
 TimeClassMsg
 TransferMsg

arkouda/sparrayclass.py

@@ -1,17 +1,23 @@
 from __future__ import annotations
 
 import builtins
-from typing import Optional, Sequence, Union
+from typing import Optional, Sequence, Union, cast
 
 import numpy as np
 from typeguard import typechecked
 
 from arkouda.client import generic_msg
-from arkouda.dtypes import dtype, int_scalars
+from arkouda.dtypes import NumericDTypes, dtype, int_scalars
 from arkouda.logger import getArkoudaLogger
+from arkouda.pdarrayclass import create_pdarrays, pdarray
 
 logger = getArkoudaLogger(name="sparrayclass")
 
+__all__ = [
+ "sparray",
+ "create_sparray",
+]
+
 
 class sparray:
  """
@@ -94,6 +100,37 @@ def __str__(self): # This won't work out of the box for sparrays need to add th
  # print("Called repr")
  # return generic_msg(cmd="repr", args={"array": self, "printThresh": sparrayIterThresh})
 
+ """
+ Converts the sparse matrix to a list of 3 pdarrays (rows, cols, vals)
+
+ Returns
+ -------
+ List[ak.pdarray]
+ A list of 3 pdarrays which contain the row indices, the column indices,
+ and the values at the respective indices within the sparse matrix.
+ Examples
+ --------
+ >>> a = ak.random_sparse_matrix(100,0.2,"CSR");
+ >>> a.to_pdarray()
+ [array([1 1 1 ... 100 100 100]), array([17 21 29 ... 75 77 85]), array([0 0 0 ... 0 0 0])]
+ """
+
+ @typechecked
+ def to_pdarray(self):
+ dtype = self.dtype
+ dtype_name = cast(np.dtype, dtype).name
+ # check dtype for error
+ if dtype_name not in NumericDTypes:
+ raise TypeError(f"unsupported dtype {dtype}")
+ responseArrays = generic_msg(cmd="sparse_to_pdarrays", args={"matrix": self})
+ array_list = create_pdarrays(responseArrays)
+ return array_list
+
+ """"""
+
+ def fill_vals(self, a: pdarray):
+ generic_msg(cmd="fill_sparse_vals", args={"matrix": self, "vals": a})
+
 
 # creates sparray object
 # only after:

pytest.ini

@@ -44,6 +44,7 @@ testpaths =
  tests/series_test.py
  tests/setops_test.py
  tests/sort_test.py
+ tests/sparse_test.py
  tests/stats_test.py
  tests/string_test.py
  tests/symbol_table_test.py

src/SparseMatrix.chpl

@@ -1,17 +1,57 @@
 module SparseMatrix {
 
  public use SpsMatUtil;
+ use ArkoudaSparseMatrixCompat;
+
+
+ // Quick and dirty, not permanent
+ proc fillSparseMatrix(ref spsMat, const A: [?D] ?eltType) throws {
+ if A.rank != 1 then
+ throw getErrorWithContext(
+ msg="fill vals requires a 1D array; got a %iD array".format(A.rank),
+ lineNumber=getLineNumber(),
+ routineName=getRoutineName(),
+ moduleName=getModuleName(),
+ errorClass="IllegalArgumentError"
+ );
+ if A.size != spsMat.domain.getNNZ() then
+ throw getErrorWithContext(
+ msg="fill vals requires an array of the same size as the sparse matrix; got %i elements, expected %i".format(A.size, spsMat.domain.getNNZ()),
+ lineNumber=getLineNumber(),
+ routineName=getRoutineName(),
+ moduleName=getModuleName(),
+ errorClass="IllegalArgumentError"
+ );
+ if eltType != spsMat.eltType then
+ throw getErrorWithContext(
+ msg="fill vals requires an array of the same type as the sparse matrix; got %s, expected %s".format(eltType, spsMat.eltType),
+ lineNumber=getLineNumber(),
+ routineName=getRoutineName(),
+ moduleName=getModuleName(),
+ errorClass="IllegalArgumentError"
+ );
+ for((i,j), idx) in zip(spsMat.domain,A.domain) {
+ spsMat[i,j] = A[idx];
+ }
+ }
 
+ proc sparseMatToPdarray(const ref spsMat, ref rows, ref cols, ref vals){
+ for((i,j), idx) in zip(spsMat.domain,0..) {
+ rows[idx] = i;
+ cols[idx] = j;
+ vals[idx] = spsMat[i, j];
+ }
+ }
  // sparse, outer, matrix-matrix multiplication algorithm; A is assumed
  // CSC and B CSR
- proc sparseMatMatMult(A, B) {
- var spsData: sparseMatDat;
+ // proc sparseMatMatMult(A, B) {
+ //  var spsData: sparseMatDat;
 
- sparseMatMatMult(A, B, spsData);
+ //  sparseMatMatMult(A, B, spsData);
 
- var C = makeSparseMat(A.domain.parentDom, spsData);
- return C;
- }
+ //  var C = makeSparseMat(A.domain.parentDom, spsData);
+ //  return C;
+ // }
 
  // This version forms the guts of the above and permits a running set
  // of nonzeroes to be passed in and updated rather than assuming that

src/SparseMatrixMsg.chpl

@@ -53,7 +53,6 @@ module SparseMatrixMsg {
  return new MsgTuple(errorMsg, MsgType.ERROR);
  }
  }
-
  }
 
 
@@ -80,9 +79,78 @@ module SparseMatrixMsg {
  }
 
 
+ proc sparseMatrixtoPdarray(cmd: string, msgArgs: borrowed MessageArgs, st: borrowed SymTab): MsgTuple throws {
+ var repMsg: string; // response message with the details of the new arr
+
+ var gEnt = getGenericSparseArrayEntry(msgArgs.getValueOf("matrix"), st);
+
+ var size = gEnt.nnz;
+ var rows = makeDistArray(size, int);
+ var cols = makeDistArray(size, int);
+ var vals = makeDistArray(size, int);
+
+ if gEnt.layoutStr=="CSC" {
+ // Hardcode for int right now
+ var sparrayEntry = gEnt.toSparseSymEntry(int, dimensions=2, layout.CSC);
+ sparseMatToPdarray(sparrayEntry.a, rows, cols, vals);
+ } else if gEnt.layoutStr=="CSR" {
+ // Hardcode for int right now
+ var sparrayEntry = gEnt.toSparseSymEntry(int, dimensions=2, layout.CSR);
+ sparseMatToPdarray(sparrayEntry.a, rows, cols, vals);
+ } else {
+ throw getErrorWithContext(
+ msg="unsupported layout for sparse matrix: %s".format(gEnt.layoutStr),
+ lineNumber=getLineNumber(),
+ routineName=getRoutineName(),
+ moduleName=getModuleName(),
+ errorClass="NotImplementedError"
+ );
+ }
+
+ var responses: [0..2] MsgTuple;
+ responses[0] = st.insert(createSymEntry(rows));
+ responses[1] = st.insert(createSymEntry(cols));
+ responses[2] = st.insert(createSymEntry(vals));
+ sparseLogger.debug(getModuleName(),getRoutineName(),getLineNumber(), "Converted sparse matrix to pdarray");
+ return MsgTuple.fromResponses(responses);
+ }
+
+
+ proc fillSparseMatrixMsg(cmd: string, msgArgs: borrowed MessageArgs, st: borrowed SymTab): MsgTuple throws {
+ var repMsg: string; // response message with the details of the new arr
+
+ var gEnt = getGenericSparseArrayEntry(msgArgs.getValueOf("matrix"), st);
+ var gEntVals: borrowed GenSymEntry = getGenericTypedArrayEntry(msgArgs.getValueOf("vals"), st);
+
+ //Hardcode int for now
+ var vals = toSymEntry(gEntVals,int);
+ if gEnt.layoutStr=="CSC" {
+ // Hardcode for int right now
+ var sparrayEntry = gEnt.toSparseSymEntry(int, dimensions=2, layout.CSC);
+ fillSparseMatrix(sparrayEntry.a, vals.a);
+ } else if gEnt.layoutStr=="CSR" {
+ // Hardcode for int right now
+ var sparrayEntry = gEnt.toSparseSymEntry(int, dimensions=2, layout.CSR);
+ fillSparseMatrix(sparrayEntry.a, vals.a);
+ } else {
+ throw getErrorWithContext(
+ msg="unsupported layout for sparse matrix: %s".format(gEnt.layoutStr),
+ lineNumber=getLineNumber(),
+ routineName=getRoutineName(),
+ moduleName=getModuleName(),
+ errorClass="NotImplementedError"
+ );
+ }
+ sparseLogger.debug(getModuleName(),getRoutineName(),getLineNumber(), "Filled sparse Array with values");
+ return MsgTuple.success();
+ }
+
+
 
  use CommandMap;
  registerFunction("random_sparse_matrix", randomSparseMatrixMsg, getModuleName());
  registerFunction("sparse_matrix_matrix_mult", sparseMatrixMatrixMultMsg, getModuleName());
+ registerFunction("sparse_to_pdarrays", sparseMatrixtoPdarray, getModuleName());
+ registerFunction("fill_sparse_vals", fillSparseMatrixMsg, getModuleName());
 
-}
+}

src/compat/eq-20/ArkoudaSparseMatrixCompat.chpl

@@ -0,0 +1,86 @@
+module ArkoudaSparseMatrixCompat {
+ use SparseBlockDist;
+
+ proc SparseBlockDom.setLocalSubdomain(locIndices, loc: locale = here) {
+ if loc != here then
+ halt("setLocalSubdomain() doesn't currently support remote updates");
+ ref myBlock = this.myLocDom!.mySparseBlock;
+ if myBlock.type != locIndices.type then
+ compilerError("setLocalSubdomain() expects its argument to be of type ",
+ myBlock.type:string);
+ else
+ myBlock = locIndices;
+ }
+
+ proc SparseBlockArr.getLocalSubarray(localeRow, localeCol) const ref {
+ return this.locArr[localeRow, localeCol]!.myElems;
+ }
+
+ proc SparseBlockArr.getLocalSubarray(localeIdx) const ref {
+ return this.locArr[localeIdx]!.myElems;
+ }
+
+ proc SparseBlockArr.setLocalSubarray(locNonzeroes, loc: locale = here) {
+ if loc != here then
+ halt("setLocalSubarray() doesn't currently support remote updates");
+ ref myBlock = this.myLocArr!.myElems;
+ if myBlock.type != locNonzeroes.type then
+ compilerError("setLocalSubarray() expects its argument to be of type ",
+ myBlock.type:string);
+ else
+ myBlock.data = locNonzeroes.data;
+ }
+
+ proc SparseBlockDom.dsiTargetLocales() const ref {
+ return dist.targetLocales;
+ }
+
+ proc SparseBlockArr.dsiTargetLocales() const ref {
+ return dom.dsiTargetLocales();
+ }
+
+ use LayoutCS;
+
+ proc CSDom.rows() {
+ return this.rowRange;
+ }
+
+ proc CSDom.cols() {
+ return this.colRange;
+ }
+
+ @chpldoc.nodoc
+ iter CSDom.uidsInRowCol(rc) {
+ for uid in startIdx[rc]..<startIdx[rc+1] do
+ yield uid;
+ }
+
+ proc CSArr.rows() {
+ return this.dom.rows();
+ }
+
+ proc CSArr.cols() {
+ return this.dom.cols();
+ }
+
+ @chpldoc.nodoc
+ iter CSArr.indsAndVals(rc) {
+ ref dom = this.dom;
+ for uid in dom.uidsInRowCol(rc) do
+ yield (dom.idx[uid], this.data[uid]);
+ }
+
+ iter CSArr.colsAndVals(r) {
+ if this.dom.compressRows == false then
+ compilerError("Can't (efficiently) iterate over rows using a CSC layout");
+ for colVal in indsAndVals(r) do
+ yield colVal;
+ }
+
+ iter CSArr.rowsAndVals(c) {
+ if this.dom.compressRows == true then
+ compilerError("Can't (efficiently) iterate over columns using a CSR layout");
+ for rowVal in indsAndVals(c) do
+ yield rowVal;
+ }
+}

src/compat/eq-21/ArkoudaSparseMatrixCompat.chpl

@@ -0,0 +1 @@
+module ArkoudaSparseMatrixCompat {}

src/compat/ge-22/ArkoudaSparseMatrixCompat.chpl

@@ -0,0 +1 @@
+module ArkoudaSparseMatrixCompat {}