import numpy as np
from adios2 import FileReader
from scipy.sparse import csr_matrix
[docs]
def sparse_matrix_reader(filename, var_names):
r"""
Read the sparse‑matrix variables stored in an ADIOS2 BP directory.
Parameters
----------
filename : str
Path to the ``*.bp`` directory.
var_names : list[str]
The six variable names in this exact order:
* **ncols_name** – number of columns
* **nrows_name** – number of rows
* **width_name** – max. non‑zeros per row
* **nelement_name** – 1‑D array of non‑zeros per row
* **col_indices_name** – 2‑D array of **1‑based** column indices
* **values_name** – 2‑D array of non‑zero values
Returns
-------
tuple
``(ncols, nrows, nelement, width, col_indices, values)``, where
``col_indices`` has been converted to **0‑based** indices.
"""
if len(var_names) != 6:
raise ValueError("var_names must be a list of 6 variable names.")
with FileReader(filename) as reader:
ncols = int(reader.read(var_names[0]).item())
nrows = int(reader.read(var_names[1]).item())
width = int(reader.read(var_names[2]).item())
nelement = reader.read(var_names[3]) # shape: (nrows,)
col_indices = reader.read(var_names[4]) # shape: (nrows, width)
values = reader.read(var_names[5]) # shape: (nrows, width)
# Convert 1-based indices to 0-based indices.
col_indices = col_indices - 1
return ncols, nrows, nelement, width, col_indices, values
[docs]
class SparseMatrixCSR:
"""
A class for storing a sparse matrix in CSR format (using SciPy's csr_matrix).
The constructor takes a reader function as input. This reader function should return
the matrix data as:
(ncols, nrows, nelement, width, col_indices, values)
where:
- ncols (int): Number of columns.
- nrows (int): Number of rows.
- nelement (np.ndarray): A 1D array (length nrows) with the number of nonzeros in each row.
- width (int): The maximal number of nonzeros in any row.
- col_indices (np.ndarray): A 2D array (nrows x width) with 0-based column indices.
- values (np.ndarray): A 2D array (nrows x width) with the nonzero values.
The constructor reformats this data and stores the matrix in SciPy's CSR format.
"""
def __init__(self, reader_function):
# Read the raw data using the provided reader function.
ncols, nrows, nelement, width, col_indices, values = reader_function()
# Ensure proper types.
ncols = int(ncols)
nrows = int(nrows)
nelement = np.array(nelement, dtype=int)
width = int(width)
col_indices = np.array(col_indices, dtype=int)
values = np.array(values, dtype=float)
# Build CSR components: data, indices, and indptr.
data_list = []
indices_list = []
indptr = [0]
for i in range(nrows):
count = int(nelement[i])
# Get the valid entries for this row.
row_indices = col_indices[i, :count]
row_values = values[i, :count]
data_list.extend(row_values.tolist())
indices_list.extend(row_indices.tolist())
indptr.append(indptr[-1] + count)
# Convert lists to numpy arrays.
data_array = np.array(data_list, dtype=float)
indices_array = np.array(indices_list, dtype=int)
indptr_array = np.array(indptr, dtype=int)
# Create the CSR matrix.
self.csr = csr_matrix((data_array, indices_array, indptr_array), shape=(nrows, ncols))
[docs]
def get_csr_matrix(self):
"""
Returns:
csr_matrix: The sparse matrix in CSR format.
"""
return self.csr
[docs]
def multiply_left(self, vector):
"""
Multiplies the CSR matrix from the left by a 1D NumPy array (i.e. computes vector * A).
Args:
vector (np.ndarray): 1D array with length equal to the number of rows of the matrix.
Returns:
np.ndarray: 1D array of length equal to the number of columns of the matrix.
Raises:
ValueError: if the length of vector does not match the number of rows.
"""
if vector.ndim != 1:
raise ValueError("Input vector must be 1D.")
if vector.shape[0] != self.csr.shape[1]:
raise ValueError("Length of vector must equal the number of columns of the matrix.")
# Multiply from the left (vector treated as a row vector).
return self.csr.dot(vector)
[docs]
def multiply_transpose_left(self, vector):
"""
Multiplies the transpose of the CSR matrix from the left by a 1D NumPy array
(i.e. computes vector * A^T).
Args:
vector (np.ndarray): 1D array with length equal to the number of columns of the matrix.
Returns:
np.ndarray: 1D array of length equal to the number of rows of the matrix.
Raises:
ValueError: if the length of vector does not match the number of columns.
"""
if vector.ndim != 1:
raise ValueError("Input vector must be 1D.")
if vector.shape[0] != self.csr.shape[0]:
raise ValueError("Length of vector must equal the number of rows of the matrix.")
# Multiply from the left with the transpose.
return self.csr.transpose().dot(vector)
# Example usage:
if __name__ == "__main__":
bp_filename = "xgc.cnv_to_surf.bp" # Replace with your actual BP file name.
# Define the variable names for the matrix in this BP file.
var_names = ["ncols", "nrows", "width", "nelement", "eindex", "value"]
# Create the SparseMatrixCSR object using a lambda that calls the sparse_matrix_reader with the file and var_names.
sparse_matrix_obj = SparseMatrixCSR(lambda: sparse_matrix_reader(bp_filename, var_names))
csr_mat = sparse_matrix_obj.get_csr_matrix()
# Print details of the CSR matrix.
print("CSR matrix data:", csr_mat.data)
print("CSR matrix indices:", csr_mat.indices)
print("CSR matrix indptr:", csr_mat.indptr)
print("CSR matrix shape:", csr_mat.shape)