(weight, edges, landmarks, output_path)
| 50 | |
| 51 | # create matrix function |
| 52 | def create_matrix(weight, edges, landmarks, output_path): |
| 53 | |
| 54 | # TODO: implement sparse Q matrix construction |
| 55 | DEFINE_DENSE = 1 |
| 56 | USE_GPU = 0 |
| 57 | |
| 58 | if USE_GPU: |
| 59 | print("TODO: add cupy and deal with CUDA OUT OF MEMORY") |
| 60 | # mempool = cupy.get_default_memory_pool() |
| 61 | # pinned_mempool = cupy.get_default_pinned_memory_pool() |
| 62 | |
| 63 | M = int(edges[:,1].max()) |
| 64 | N = int(edges[:,0].max()) |
| 65 | print(f"M: {M}, N: {N}") |
| 66 | n_observation = edges.shape[0] |
| 67 | V3 = coo_matrix((weight, (edges[:, 0]-1, edges[:, 1]-1)), shape=(N, M)).tocsr() |
| 68 | |
| 69 | observation_index = coo_matrix((np.arange(1, len(edges) + 1), (edges[:, 0]-1, edges[:, 1]-1)), shape=(N, M)).tocsr() |
| 70 | |
| 71 | Q2 = diags(V3.sum(axis=1).A1, format='csr') |
| 72 | Q3 = diags(V3.sum(axis=0).A1, format='csr') |
| 73 | |
| 74 | # Initialize sparse matrices |
| 75 | Q1 = np.zeros((3 * N, 3 * N)) |
| 76 | V1 = np.zeros((3 * N, N)) |
| 77 | V2 = np.zeros((3 * N, M)) |
| 78 | |
| 79 | # Prepare lists to collect the indices and data for assembly |
| 80 | Q1_data = [] |
| 81 | V1_data = [] |
| 82 | V2_data = [] |
| 83 | |
| 84 | print("Start reading data") |
| 85 | |
| 86 | with ProcessPoolExecutor() as executor: |
| 87 | # Run the process_observation function in parallel for each i |
| 88 | futures = [executor.submit(process_observation, i, observation_index.getrow(i).indices, landmarks[observation_index.getrow(i).data-1,:], weight[observation_index.getrow(i).data-1]) for i in range(N)] |
| 89 | |
| 90 | for future in futures: |
| 91 | |
| 92 | i, Q1_block, V1_block, V2_block, ind_l = future.result() |
| 93 | |
| 94 | # Collect data and indices for Q1 |
| 95 | row_indices_Q1 = np.repeat(np.arange(3*i, 3*i+3), 3) |
| 96 | col_indices_Q1 = np.tile(np.arange(3*i, 3*i+3), 3) |
| 97 | Q1_values = Q1_block.flatten() |
| 98 | Q1_data.append((Q1_values, row_indices_Q1, col_indices_Q1)) |
| 99 | |
| 100 | # Collect data and indices for V1 |
| 101 | row_indices_V1 = np.arange(3*i, 3*i+3) |
| 102 | col_indices_V1 = np.array([i]*3) |
| 103 | V1_values = V1_block.flatten() |
| 104 | V1_data.append((V1_values, row_indices_V1, col_indices_V1)) |
| 105 | |
| 106 | # Collect data and indices for V2 |
| 107 | rows_V2 = np.repeat(np.arange(3*i, 3*i+3), len(ind_l)) |
| 108 | cols_V2 = np.tile(ind_l, 3) |
| 109 | V2_values = V2_block.flatten() |
no test coverage detected