Adding functions to create and manipulate sparse matrices

Author: pavanky

arrayfire/__init__.py

@@ -73,6 +73,7 @@
 from .interop    import *
 from .timer      import *
 from .random     import *
+from .sparse     import *
 
 # do not export default modules as part of arrayfire
 del ct

arrayfire/sparse.py

@@ -0,0 +1,275 @@
+#######################################################
+# Copyright (c) 2015, ArrayFire
+# All rights reserved.
+#
+# This file is distributed under 3-clause BSD license.
+# The complete license agreement can be obtained at:
+# http://arrayfire.com/licenses/BSD-3-Clause
+########################################################
+
+"""
+Functions to create and manipulate sparse matrices.
+"""
+
+from .library import *
+from .array import *
+import numbers
+from .interop import to_array
+
+__to_sparse_enum = [STORAGE.DENSE,
+                    STORAGE.CSR,
+                    STORAGE.CSC,
+                    STORAGE.COO]
+
+
+def sparse(values, row_idx, col_idx, nrows, ncols, storage = STORAGE.CSR):
+    """
+    Create a sparse matrix from it's constituent parts.
+
+    Parameters
+    ----------
+
+    values : af.Array.
+          - Contains the non zero elements of the sparse array.
+
+    row_idx : af.Array.
+          - Contains row indices of the sparse array.
+
+    col_idx : af.Array.
+          - Contains column indices of the sparse array.
+
+    nrows   : int.
+          - specifies the number of rows in sparse matrix.
+
+    ncols   : int.
+          - specifies the number of columns in sparse matrix.
+
+    storage : optional: arrayfire.STORAGE. default: arrayfire.STORAGE.CSR.
+          - Can be one of arrayfire.STORAGE.CSR, arrayfire.STORAGE.COO.
+
+    Returns
+    -------
+
+    A sparse matrix.
+    """
+    assert(isinstance(values, Array))
+    assert(isinstance(row_idx, Array))
+    assert(isinstance(col_idx, Array))
+    out = Array()
+    safe_call(backend.get().af_create_sparse_array(ct.pointer(out.arr), c_dim_t(nrows), c_dim_t(ncols),
+                                                   values.arr, row_idx.arr, col_idx.arr, storage.value))
+    return out
+
+def sparse_from_host(values, row_idx, col_idx, nrows, ncols, storage = STORAGE.CSR):
+    """
+    Create a sparse matrix from it's constituent parts.
+
+    Parameters
+    ----------
+
+    values : Any datatype that can be converted to array.
+          - Contains the non zero elements of the sparse array.
+
+    row_idx : Any datatype that can be converted to array.
+          - Contains row indices of the sparse array.
+
+    col_idx : Any datatype that can be converted to array.
+          - Contains column indices of the sparse array.
+
+    nrows   : int.
+          - specifies the number of rows in sparse matrix.
+
+    ncols   : int.
+          - specifies the number of columns in sparse matrix.
+
+    storage : optional: arrayfire.STORAGE. default: arrayfire.STORAGE.CSR.
+          - Can be one of arrayfire.STORAGE.CSR, arrayfire.STORAGE.COO.
+
+    Returns
+    -------
+
+    A sparse matrix.
+    """
+    return sparse(to_array(values), to_array(row_idx), to_array(col_idx), nrows, ncols, storage)
+
+def sparse_from_dense(dense, storage = STORAGE.CSR):
+    """
+    Create a sparse matrix from a dense matrix.
+
+    Parameters
+    ----------
+
+    dense : af.Array.
+          - A dense matrix.
+
+    storage : optional: arrayfire.STORAGE. default: arrayfire.STORAGE.CSR.
+          - Can be one of arrayfire.STORAGE.CSR, arrayfire.STORAGE.COO.
+
+    Returns
+    -------
+
+    A sparse matrix.
+    """
+    assert(isinstance(dense, Array))
+    out = Array()
+    safe_call(backend.get().af_create_sparse_array_from_dense(ct.pointer(out.arr), dense.arr, storage.value))
+    return out
+
+def sparse_to_dense(sparse):
+    """
+    Create a dense matrix from a sparse matrix.
+
+    Parameters
+    ----------
+
+    sparse : af.Array.
+          - A sparse matrix.
+
+    Returns
+    -------
+
+    A dense matrix.
+    """
+    out = Array()
+    safe_call(backend.get().af_sparse_to_dense(ct.pointer(out.arr), sparse.arr))
+    return out
+
+def sparse_get_info(sparse):
+    """
+    Get the constituent arrays and storage info from a sparse matrix.
+
+    Parameters
+    ----------
+
+    sparse : af.Array.
+           - A sparse matrix.
+
+    Returns
+    --------
+    (values, row_idx, col_idx, storage) where
+    values : arrayfire.Array containing non zero elements from sparse matrix
+    row_idx : arrayfire.Array containing the row indices
+    col_idx : arrayfire.Array containing the column indices
+    storage : sparse storage
+    """
+    values = Array()
+    row_idx = Array()
+    col_idx = Array()
+    stype = ct.c_int(0)
+    safe_call(backend.get().af_sparse_get_info(ct.pointer(values.arr), ct.pointer(row_idx.arr),
+                                               ct.pointer(col_idx.arr), ct.pointer(stype),
+                                               sparse.arr))
+    return (values, row_idx, col_idx, __to_sparse_enum[stype.value])
+
+def sparse_get_values(sparse):
+    """
+    Get the non zero values from sparse matrix.
+
+    Parameters
+    ----------
+
+    sparse : af.Array.
+           - A sparse matrix.
+
+    Returns
+    --------
+    arrayfire array containing the non zero elements.
+
+    """
+    values = Array()
+    safe_call(backend.get().af_sparse_get_values(ct.pointer(values.arr), sparse.arr))
+    return values
+
+def sparse_get_row_idx(sparse):
+    """
+    Get the row indices from sparse matrix.
+
+    Parameters
+    ----------
+
+    sparse : af.Array.
+           - A sparse matrix.
+
+    Returns
+    --------
+    arrayfire array containing the non zero elements.
+
+    """
+    row_idx = Array()
+    safe_call(backend.get().af_sparse_get_row_idx(ct.pointer(row_idx.arr), sparse.arr))
+    return row_idx
+
+def sparse_get_col_idx(sparse):
+    """
+    Get the column indices from sparse matrix.
+
+    Parameters
+    ----------
+
+    sparse : af.Array.
+           - A sparse matrix.
+
+    Returns
+    --------
+    arrayfire array containing the non zero elements.
+
+    """
+    col_idx = Array()
+    safe_call(backend.get().af_sparse_get_col_idx(ct.pointer(col_idx.arr), sparse.arr))
+    return col_idx
+
+def sparse_get_nnz(sparse):
+    """
+    Get the column indices from sparse matrix.
+
+    Parameters
+    ----------
+
+    sparse : af.Array.
+           - A sparse matrix.
+
+    Returns
+    --------
+    Number of non zero elements in the sparse matrix.
+
+    """
+    nnz = c_dim_t(0)
+    safe_call(backend.get().af_sparse_get_nnz(ct.pointer(nnz), sparse.arr))
+    return nnz.value
+
+def sparse_get_storage(sparse):
+    """
+    Get the column indices from sparse matrix.
+
+    Parameters
+    ----------
+
+    sparse : af.Array.
+           - A sparse matrix.
+
+    Returns
+    --------
+    Number of non zero elements in the sparse matrix.
+
+    """
+    storage = ct.c_int(0)
+    safe_call(backend.get().af_sparse_get_storage(ct.pointer(storage), sparse.arr))
+    return __to_sparse_enum[storage.value]
+
+def sparse_convert_to(sparse, storage):
+    """
+    Convert sparse matrix from one format to another.
+
+    Parameters
+    ----------
+
+    storage : arrayfire.STORAGE.
+
+    Returns
+    -------
+
+    Sparse matrix converted to the appropriate type.
+    """
+    out = Array()
+    safe_call(backend.get().af_sparse_convert_to(ct.pointer(out.arr), sparse.arr, storage.value))
+    return out

arrayfire/tests/simple/__init__.py

@@ -19,4 +19,5 @@
 from .signal import *
 from .statistics import *
 from .random import *
+from .sparse import *
 from ._util import tests

arrayfire/tests/simple/sparse.py

@@ -0,0 +1,28 @@
+#!/usr/bin/python
+#######################################################
+# Copyright (c) 2015, ArrayFire
+# All rights reserved.
+#
+# This file is distributed under 3-clause BSD license.
+# The complete license agreement can be obtained at:
+# http://arrayfire.com/licenses/BSD-3-Clause
+########################################################
+
+import arrayfire as af
+from . import _util
+
+def simple_sparse(verbose=False):
+    display_func = _util.display_func(verbose)
+    print_func   = _util.print_func(verbose)
+
+    dd = af.randu(5, 5)
+    ds = dd * (dd > 0.5)
+    sp = af.sparse_from_dense(ds)
+    display_func(af.sparse_get_info(sp))
+    display_func(af.sparse_get_values(sp))
+    display_func(af.sparse_get_row_idx(sp))
+    display_func(af.sparse_get_col_idx(sp))
+    print_func(af.sparse_get_nnz(sp))
+    print_func(af.sparse_get_storage(sp))
+
+_util.tests['sparse'] = simple_sparse