diff --git a/scripts/hykkt/hykkt_matrix_splitter.py b/scripts/hykkt/hykkt_matrix_splitter.py new file mode 100644 index 000000000..f65111697 --- /dev/null +++ b/scripts/hykkt/hykkt_matrix_splitter.py @@ -0,0 +1,158 @@ +''' +Split HyKKT matrix into matrix blocks (D_d, D_s, J, J_d) and RHS vector into vector blocks (r_x, r_s, r_y, r_yd). + +Usage: python ./hykkt_matrix_splitter.py [HyKKT matrix local path/URL] [HyKKT RHS local path/URL] [output directory] + +Notes: +Matrix and RHS must be .mtx files if using local path. They can be .mtx.gz files if using URLs. +It doesn't matter if output directory has "/" at the end. + +Might not work for very specific edge cases, such as when there is a very long diagonal of -1's outside the -I block. +''' + +import sys +import os +import requests +import numpy as np +import scipy.io + +if len(sys.argv) == 1: + print("Arguments: matrix path/URL, RHS path/URL, output directory. Files can be .mtx or .mtx.gz.") + sys.exit() + +try: + import cupy as cp + import cupyx.scipy.sparse as cpsparse + cp.cuda.Device(0).use() + cp.cuda.runtime.setDevice(0) +except Exception: + import scipy.sparse as cpsparse + cp = np + +output_dir = sys.argv[3].rstrip("/") + +if sys.argv[1].startswith("http"): + print("Downloading matrix file.") + matrix_url = sys.argv[1] + matrix_file = os.path.basename(matrix_url) + if (matrix_file.endswith(".mtx.gz")): + matrix_name = matrix_file[:-7] + else: + matrix_name = os.path.splitext(matrix_file)[0] + matrix_path = f"{output_dir}/{matrix_file}" + with open(matrix_path, "wb") as file: + file.write(requests.get(matrix_url).content) +else: + matrix_path = sys.argv[1] + matrix_file = os.path.basename(matrix_path) + matrix_name = os.path.splitext(matrix_file)[0] + +if sys.argv[2].startswith("http"): + print("Downloading rhs file.") + rhs_url = sys.argv[2] + rhs_file = os.path.basename(rhs_url) + if rhs_file.endswith(".mtx.gz"): + rhs_name = rhs_file[:-7] + else: + rhs_name = os.path.splitext(rhs_file)[0] + rhs_path = f"{output_dir}/{rhs_file}" + with open(rhs_path, "wb") as file: + file.write(requests.get(rhs_url).content) +else: + rhs_path = sys.argv[2] + rhs_file = os.path.basename(rhs_path) + rhs_name = os.path.splitext(rhs_file)[0] + +print("Splitting matrix.") + +A_cpu = scipy.io.mmread(matrix_path).tocoo() +n = A_cpu.shape[0] + +A_row = cp.asarray(A_cpu.row) +A_col = cp.asarray(A_cpu.col) +A_val = cp.asarray(A_cpu.data) + +mo_mask = (A_val == -1) +mo_row = A_row[mo_mask] +mo_col = A_col[mo_mask] + +offsets = mo_row - mo_col +mI_mask = (offsets > 0) +offsets = offsets[mI_mask] +assert len(offsets) > 0, "Not a valid HyKKT matrix!" + +unique_offsets, counts = cp.unique(offsets, return_counts=True) +most_common_offset = unique_offsets[cp.argmax(counts)] + +mI_block_col = mo_col[mI_mask] +mI_block_col = cp.sort(mI_block_col[offsets == most_common_offset]) +assert mI_block_col[-1] - mI_block_col[0] == len(mI_block_col) - 1, "Stray -1!" + +nx = mI_block_col[0] +md = len(mI_block_col) +mc = (n - nx - md - md).item() + +print(f"n_x={nx}, m_d={md}, m_c={mc}") + +R0_mask = A_row < nx +R1_mask = ~R0_mask & (A_row < nx + md) +R2_mask = ~R0_mask & ~R1_mask & (A_row < nx + md + mc) +R3_mask = ~R0_mask & ~R1_mask & ~R2_mask + +C0_mask = A_col < nx +C1_mask = ~C0_mask & (A_col < nx + md) +C2_mask = ~C0_mask & ~C1_mask & (A_col < nx + md + mc) +C3_mask = ~C0_mask & ~C1_mask & ~C2_mask + +H_mask = R0_mask & C0_mask +Dd_mask = R1_mask & C1_mask +J_mask = R2_mask & C0_mask +Jd_mask = R3_mask & C0_mask + +R0_end = nx +R1_end = nx + md +R2_end = nx + md + mc +C0_end = nx + +H_block = cpsparse.coo_matrix((A_val[H_mask], (A_row[H_mask], A_col[H_mask])), shape=(nx, nx)) +Dd_block = cpsparse.coo_matrix((A_val[Dd_mask], (A_row[Dd_mask] - R0_end, A_col[Dd_mask] - C0_end)), shape=(md, md)) +J_block = cpsparse.coo_matrix((A_val[J_mask], (A_row[J_mask] - R1_end, A_col[J_mask])), shape=(mc, nx)) +Jd_block = cpsparse.coo_matrix((A_val[Jd_mask], (A_row[Jd_mask] - R2_end, A_col[Jd_mask])), shape=(md, nx)) + +print(f"H nnz={H_block.nnz}, D_s nnz={Dd_block.nnz}, J nnz={J_block.nnz}, Jd nnz={Jd_block.nnz}") + +if cp == np: + scipy.io.mmwrite(f"{output_dir}/block_H_{matrix_name}.mtx", H_block) + scipy.io.mmwrite(f"{output_dir}/block_Dd_{matrix_name}.mtx", Dd_block) + scipy.io.mmwrite(f"{output_dir}/block_J_{matrix_name}.mtx", J_block) + scipy.io.mmwrite(f"{output_dir}/block_Jd_{matrix_name}.mtx", Jd_block) +else: + scipy.io.mmwrite(f"{output_dir}/block_H_{matrix_name}.mtx", H_block.get()) + scipy.io.mmwrite(f"{output_dir}/block_Dd_{matrix_name}.mtx", Dd_block.get()) + scipy.io.mmwrite(f"{output_dir}/block_J_{matrix_name}.mtx", J_block.get()) + scipy.io.mmwrite(f"{output_dir}/block_Jd_{matrix_name}.mtx", Jd_block.get()) + +rhs = cp.asarray(scipy.io.mmread(rhs_path)).flatten() + +rx = rhs[: R0_end] +rs = rhs[R0_end : R1_end] +ry = rhs[R1_end : R2_end] +ryd = rhs[R2_end :] + +if cp == np: + scipy.io.mmwrite(f"{output_dir}/block_rx_{rhs_name}.mtx", rx.reshape(-1, 1)) + scipy.io.mmwrite(f"{output_dir}/block_rs_{rhs_name}.mtx", rs.reshape(-1, 1)) + scipy.io.mmwrite(f"{output_dir}/block_ry_{rhs_name}.mtx", ry.reshape(-1, 1)) + scipy.io.mmwrite(f"{output_dir}/block_ryd_{rhs_name}.mtx", ryd.reshape(-1, 1)) +else: + scipy.io.mmwrite(f"{output_dir}/block_rx_{rhs_name}.mtx", rx.get().reshape(-1, 1)) + scipy.io.mmwrite(f"{output_dir}/block_rs_{rhs_name}.mtx", rs.get().reshape(-1, 1)) + scipy.io.mmwrite(f"{output_dir}/block_ry_{rhs_name}.mtx", ry.get().reshape(-1, 1)) + scipy.io.mmwrite(f"{output_dir}/block_ryd_{rhs_name}.mtx", ryd.get().reshape(-1, 1)) + +if sys.argv[1].startswith("http"): + os.remove(matrix_path) +if sys.argv[2].startswith("http"): + os.remove(rhs_path) + +print("Done.") diff --git a/scripts/hykkt/requirements.txt b/scripts/hykkt/requirements.txt new file mode 100644 index 000000000..6969f73b5 --- /dev/null +++ b/scripts/hykkt/requirements.txt @@ -0,0 +1,3 @@ +numpy +cupy +scipy