diff --git a/pyttb/export_data.py b/pyttb/export_data.py index 520e89a2..09dd3d36 100644 --- a/pyttb/export_data.py +++ b/pyttb/export_data.py @@ -109,12 +109,7 @@ def export_sparse_array( """Export sparse array data in coordinate format.""" if not fmt_data: fmt_data = "%.16e" - # TODO: looping through all values may take a long time, can this be more efficient? - for i in range(A.nnz): - # 0-based indexing in package, 1-based indexing in file - subs = A.subs[i, :] + index_base - subs.tofile(fp, sep=" ", format="%d") - print(end=" ", file=fp) - val = A.vals[i][0] - val.tofile(fp, sep=" ", format=fmt_data) - print(file=fp) + # 0-based indexing in package, 1-based indexing in file + subs = A.subs + index_base + vals = A.vals[:, 0].reshape(-1, 1) + np.savetxt(fp, np.hstack((subs, vals)), fmt="%d " * subs.shape[1] + fmt_data) diff --git a/pyttb/import_data.py b/pyttb/import_data.py index 7c73c26d..c31cc167 100644 --- a/pyttb/import_data.py +++ b/pyttb/import_data.py @@ -99,12 +99,13 @@ def import_sparse_array( fp: TextIO, n: int, nz: int, index_base: int = 1 ) -> tuple[np.ndarray, np.ndarray]: """Extract sparse data subs and vals from coordinate format data.""" - subs = np.zeros((nz, n), dtype="int64") - vals = np.zeros((nz, 1)) - for k in range(nz): - line = fp.readline().strip().split(" ") - subs[k, :] = [np.int64(i) - index_base for i in line[:-1]] - vals[k, 0] = line[-1] + data = np.loadtxt(fp) + subs = data[:, :-1].astype("int64") - index_base + vals = data[:, -1].reshape(-1, 1) + if subs.shape[0] != nz: + raise ValueError("Imported nonzeros are not of expected size") + if subs.shape[1] != n: + raise ValueError("Imported tensor is not of expected shape") return subs, vals