<html><head><meta http-equiv="content-type" content="text/html; charset=utf-8"></head><body style="overflow-wrap: break-word; -webkit-nbsp-mode: space; line-break: after-white-space;">I am trying to multiply a sequential PETsc matrix with an mpi PETSc matrix in parallel. The final step is to concatenate the product matrix, which is a local sequential PETSc matrix that is different for every proc, so that I get the full mpi matrix as a result. This has proven to work, but setting the values one by one using a loop is very inefficient and slow.<div><br></div><div>In the following MFE, I am trying to make this concatenation more efficient by setting the values in batches. However it doesn’t work and I am wondering why:</div><div><br></div><div><div style="color: rgb(29, 33, 38); background-color: rgb(255, 255, 255); font-family: Menlo, Monaco, "Courier New", monospace; line-height: 18px; white-space: pre;"><div><span style="color: rgb(10, 66, 131);">"""Experimenting with PETSc mat-mat multiplication"""</span></div><br><div><span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">time</span></div><br><div><span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">numpy</span> <span style="color: rgb(168, 3, 20);">as</span> <span style="color: rgb(72, 31, 148);">np</span></div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">colorama</span> <span style="color: rgb(168, 3, 20);">import</span> Fore</div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">firedrake</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(13, 92, 181);">COMM_SELF</span>, <span style="color: rgb(13, 92, 181);">COMM_WORLD</span></div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">firedrake</span>.<span style="color: rgb(72, 31, 148);">petsc</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">PETSc</span></div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">mpi4py</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">MPI</span></div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">numpy</span>.<span style="color: rgb(72, 31, 148);">testing</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">assert_array_almost_equal</span></div><br><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">utilities</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">Print</span></div><br><div>size <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span>.size</div><div>rank <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span>.rank</div><div style="color: rgb(29, 33, 38); background-color: rgb(255, 255, 255); font-family: Menlo, Monaco, "Courier New", monospace; line-height: 18px; white-space: pre;"><br></div><div style="color: rgb(29, 33, 38); background-color: rgb(255, 255, 255); font-family: Menlo, Monaco, "Courier New", monospace; line-height: 18px; white-space: pre;"><div style="line-height: 18px;"><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix</span>(<span style="color: rgb(166, 61, 0);">input_array</span>, <span style="color: rgb(166, 61, 0);">sparse</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">True</span>):</div><div>    <span style="color: rgb(10, 66, 131);">"""Create a PETSc matrix from an input_array</span></div><br><div><span style="color: rgb(10, 66, 131);">    Args:</span></div><div><span style="color: rgb(10, 66, 131);">        input_array (np array): Input array</span></div><div><span style="color: rgb(10, 66, 131);">        partition_like (PETSc mat, optional): Petsc matrix. Defaults to None.</span></div><div><span style="color: rgb(10, 66, 131);">        sparse (bool, optional): Toggle for sparese or dense. Defaults to True.</span></div><br><div><span style="color: rgb(10, 66, 131);">    Returns:</span></div><div><span style="color: rgb(10, 66, 131);">        PETSc mat: PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);">    """</span></div><div>    <span style="color: rgb(164, 176, 188);"># Check if input_array is 1D and reshape if necessary</span></div><div>    <span style="color: rgb(168, 3, 20);">assert</span> <span style="color: rgb(72, 31, 148);">len</span>(<span style="color: rgb(166, 61, 0);">input_array</span>.shape) <span style="color: rgb(168, 3, 20);">==</span> <span style="color: rgb(13, 92, 181);">2</span>, <span style="color: rgb(10, 66, 131);">"Input array should be 2-dimensional"</span></div><div>    global_rows, global_cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">input_array</span>.shape</div><div><br></div><div>    size <span style="color: rgb(168, 3, 20);">=</span> ((<span style="color: rgb(13, 92, 181);">None</span>, global_rows), (global_cols, global_cols))</div><br><div>    <span style="color: rgb(164, 176, 188);"># Create a sparse or dense matrix based on the 'sparse' argument</span></div><div>    <span style="color: rgb(168, 3, 20);">if</span> <span style="color: rgb(166, 61, 0);">sparse</span>:</div><div>        matrix <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createAIJ(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>size, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">COMM_WORLD</span>)</div><div>    <span style="color: rgb(168, 3, 20);">else</span>:</div><div>        matrix <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createDense(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>size, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">COMM_WORLD</span>)</div><div>    matrix.setUp()</div><br><div>    local_rows_start, local_rows_end <span style="color: rgb(168, 3, 20);">=</span> matrix.getOwnershipRange()</div><br><div>    <span style="color: rgb(168, 3, 20);">for</span> counter, i <span style="color: rgb(168, 3, 20);">in</span> <span style="color: rgb(72, 31, 148);">enumerate</span>(<span style="color: rgb(72, 31, 148);">range</span>(local_rows_start, local_rows_end)):</div><div>        <span style="color: rgb(164, 176, 188);"># Calculate the correct row in the array for the current process</span></div><div>        row_in_array <span style="color: rgb(168, 3, 20);">=</span> counter <span style="color: rgb(168, 3, 20);">+</span> local_rows_start</div><div>        matrix.setValues(</div><div>            i, <span style="color: rgb(72, 31, 148);">range</span>(global_cols), <span style="color: rgb(166, 61, 0);">input_array</span>[row_in_array, :], <span style="color: rgb(166, 61, 0);">addv</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">False</span></div><div>        )</div><br><div>    <span style="color: rgb(164, 176, 188);"># Assembly the matrix to compute the final structure</span></div><div>    matrix.assemblyBegin()</div><div>    matrix.assemblyEnd()</div><br><div>    <span style="color: rgb(168, 3, 20);">return</span> matrix</div></div></div><div style="color: rgb(29, 33, 38); background-color: rgb(255, 255, 255); font-family: Menlo, Monaco, "Courier New", monospace; line-height: 18px; white-space: pre;"><br></div><div style="color: rgb(29, 33, 38); background-color: rgb(255, 255, 255); font-family: Menlo, Monaco, "Courier New", monospace; line-height: 18px; white-space: pre;"><div style="line-height: 18px;"><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">get_local_submatrix</span>(<span style="color: rgb(166, 61, 0);">A</span>):</div><div>    <span style="color: rgb(10, 66, 131);">"""Get the local submatrix of A</span></div><br><div><span style="color: rgb(10, 66, 131);">    Args:</span></div><div><span style="color: rgb(10, 66, 131);">        A (mpi PETSc mat): partitioned PETSc matrix</span></div><br><div><span style="color: rgb(10, 66, 131);">    Returns:</span></div><div><span style="color: rgb(10, 66, 131);">        seq mat: PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);">    """</span></div><div>    local_rows_start, local_rows_end <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A</span>.getOwnershipRange()</div><div>    local_rows <span style="color: rgb(168, 3, 20);">=</span> local_rows_end <span style="color: rgb(168, 3, 20);">-</span> local_rows_start</div><div>    comm <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A</span>.getComm()</div><div>    rows <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.IS().createStride(</div><div>        local_rows, <span style="color: rgb(166, 61, 0);">first</span><span style="color: rgb(168, 3, 20);">=</span>local_rows_start, <span style="color: rgb(166, 61, 0);">step</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">1</span>, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span>comm</div><div>    )</div><div>    _, k <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A</span>.getSize()  <span style="color: rgb(164, 176, 188);"># Get the number of columns (k) from A's size</span></div><div>    cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.IS().createStride(k, <span style="color: rgb(166, 61, 0);">first</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">0</span>, <span style="color: rgb(166, 61, 0);">step</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">1</span>, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span>comm)</div><br><div>    <span style="color: rgb(164, 176, 188);"># Getting the local submatrix</span></div><div>    <span style="color: rgb(164, 176, 188);"># </span><span style="color: rgb(168, 3, 20);">TODO</span><span style="color: rgb(164, 176, 188);">: To be replaced by MatMPIAIJGetLocalMat() in the future (see petsc-users mailing list). There is a missing petsc4py binding, need to add it myself (and please create a merge request)</span></div><div>    A_local <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A</span>.createSubMatrices(rows, cols)[<span style="color: rgb(13, 92, 181);">0</span>]</div><div>    <span style="color: rgb(168, 3, 20);">return</span> A_local</div></div></div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix_seq</span>(<span style="color: rgb(166, 61, 0);">input_array</span>):</div><div>    <span style="color: rgb(10, 66, 131);">"""Building a sequential PETSc matrix from an array</span></div><br><div><span style="color: rgb(10, 66, 131);">    Args:</span></div><div><span style="color: rgb(10, 66, 131);">        input_array (np array): Input array</span></div><br><div><span style="color: rgb(10, 66, 131);">    Returns:</span></div><div><span style="color: rgb(10, 66, 131);">        seq mat: PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);">    """</span></div><div>    <span style="color: rgb(168, 3, 20);">assert</span> <span style="color: rgb(72, 31, 148);">len</span>(<span style="color: rgb(166, 61, 0);">input_array</span>.shape) <span style="color: rgb(168, 3, 20);">==</span> <span style="color: rgb(13, 92, 181);">2</span></div><br><div>    m, n <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">input_array</span>.shape</div><div>    matrix <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createAIJ(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>(m, n), <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">COMM_SELF</span>)</div><div>    matrix.setUp()</div><br><div>    matrix.setValues(<span style="color: rgb(72, 31, 148);">range</span>(m), <span style="color: rgb(72, 31, 148);">range</span>(n), <span style="color: rgb(166, 61, 0);">input_array</span>, <span style="color: rgb(166, 61, 0);">addv</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">False</span>)</div><br><div>    <span style="color: rgb(164, 176, 188);"># Assembly the matrix to compute the final structure</span></div><div>    matrix.assemblyBegin()</div><div>    matrix.assemblyEnd()</div><br><div>    <span style="color: rgb(168, 3, 20);">return</span> matrix</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">multiply_matrices_seq</span>(<span style="color: rgb(166, 61, 0);">A_seq</span>, <span style="color: rgb(166, 61, 0);">B_seq</span>):</div><div>    <span style="color: rgb(10, 66, 131);">"""Multiply 2 sequential matrices</span></div><br><div><span style="color: rgb(10, 66, 131);">    Args:</span></div><div><span style="color: rgb(10, 66, 131);">        A_seq (seqaij): local submatrix of A</span></div><div><span style="color: rgb(10, 66, 131);">        B_seq (seqaij): sequential matrix B</span></div><br><div><span style="color: rgb(10, 66, 131);">    Returns:</span></div><div><span style="color: rgb(10, 66, 131);">        seq mat: PETSc matrix that is the product of A_seq and B_seq</span></div><div><span style="color: rgb(10, 66, 131);">    """</span></div><div>    _, A_seq_cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A_seq</span>.getSize()</div><div>    B_seq_rows, _ <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">B_seq</span>.getSize()</div><div>    <span style="color: rgb(168, 3, 20);">assert</span> (</div><div>        A_seq_cols <span style="color: rgb(168, 3, 20);">==</span> B_seq_rows</div><div>    ), <span style="color: rgb(168, 3, 20);">f</span><span style="color: rgb(10, 66, 131);">"Incompatible matrix sizes for multiplication: </span><span style="color: rgb(13, 92, 181);">{</span>A_seq_cols<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);"> != </span><span style="color: rgb(13, 92, 181);">{</span>B_seq_rows<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);">"</span></div><div>    C_local <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A_seq</span>.matMult(<span style="color: rgb(166, 61, 0);">B_seq</span>)</div><div>    <span style="color: rgb(168, 3, 20);">return</span> C_local</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">concatenate_local_to_global_matrix</span>(<span style="color: rgb(166, 61, 0);">local_matrix</span>, <span style="color: rgb(166, 61, 0);">mat_type</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">None</span>):</div><div>    <span style="color: rgb(10, 66, 131);">"""Create the global matrix C from the local submatrix local_matrix</span></div><br><div><span style="color: rgb(10, 66, 131);">    Args:</span></div><div><span style="color: rgb(10, 66, 131);">        local_matrix (seqaij): local submatrix of global_matrix</span></div><div><span style="color: rgb(10, 66, 131);">        partition_like (mpiaij): partitioned PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);">        mat_type (str): type of the global matrix. Defaults to None. If None, the type of local_matrix is used.</span></div><br><div><span style="color: rgb(10, 66, 131);">    Returns:</span></div><div><span style="color: rgb(10, 66, 131);">        mpi PETSc mat: partitioned PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);">    """</span></div><div>    local_matrix_rows, local_matrix_cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">local_matrix</span>.getSize()</div><div>    global_rows <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span>.allreduce(local_matrix_rows, <span style="color: rgb(166, 61, 0);">op</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(72, 31, 148);">MPI</span>.SUM)</div><br><div>    <span style="color: rgb(164, 176, 188);"># Determine the local portion of the vector</span></div><div>    size <span style="color: rgb(168, 3, 20);">=</span> ((<span style="color: rgb(13, 92, 181);">None</span>, global_rows), (local_matrix_cols, local_matrix_cols))</div><br><div>    <span style="color: rgb(168, 3, 20);">if</span> <span style="color: rgb(166, 61, 0);">mat_type</span> <span style="color: rgb(168, 3, 20);">is</span> <span style="color: rgb(13, 92, 181);">None</span>:</div><div>        <span style="color: rgb(166, 61, 0);">mat_type</span> <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">local_matrix</span>.getType()</div><br><div>    <span style="color: rgb(168, 3, 20);">if</span> <span style="color: rgb(10, 66, 131);">"dense"</span> <span style="color: rgb(168, 3, 20);">in</span> <span style="color: rgb(166, 61, 0);">mat_type</span>:</div><div>        sparse <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">False</span></div><div>    <span style="color: rgb(168, 3, 20);">else</span>:</div><div>        sparse <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">True</span></div><br><div>    <span style="color: rgb(168, 3, 20);">if</span> sparse:</div><div>        global_matrix <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createAIJ(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>size, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">COMM_WORLD</span>)</div><div>    <span style="color: rgb(168, 3, 20);">else</span>:</div><div>        global_matrix <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createDense(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>size, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">COMM_WORLD</span>)</div><div>    global_matrix.setUp()</div><br><div>    <span style="color: rgb(164, 176, 188);"># The exscan operation is used to get the starting global row for each process.</span></div><div>    <span style="color: rgb(164, 176, 188);"># The result of the exclusive scan is the sum of the local rows from previous ranks.</span></div><div>    global_row_start <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span>.exscan(local_matrix_rows, <span style="color: rgb(166, 61, 0);">op</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(72, 31, 148);">MPI</span>.SUM)</div><div>    <span style="color: rgb(168, 3, 20);">if</span> rank <span style="color: rgb(168, 3, 20);">==</span> <span style="color: rgb(13, 92, 181);">0</span>:</div><div>        global_row_start <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">0</span></div><br><div>    concatenate_start <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">time</span>.<span style="color: rgb(72, 31, 148);">time</span>()</div><div style="color: rgb(29, 33, 38); background-color: rgb(255, 255, 255); font-family: Menlo, Monaco, "Courier New", monospace; line-height: 18px; white-space: pre;"><br></div><div style="color: rgb(29, 33, 38); background-color: rgb(255, 255, 255); font-family: Menlo, Monaco, "Courier New", monospace; line-height: 18px; white-space: pre;"><div>    <span style="color: rgb(164, 176, 188);"># </span>This works but is very inefficient</div></div><div>    <span style="color: rgb(164, 176, 188);"># for i in range(local_matrix_rows):</span></div><div>    <span style="color: rgb(164, 176, 188);">#     cols, values = local_matrix.getRow(i)</span></div><div>    <span style="color: rgb(164, 176, 188);">#     global_row = i + global_row_start</span></div><div>    <span style="color: rgb(164, 176, 188);">#     global_matrix.setValues(global_row, cols, values)</span></div><br><div>    all_cols <span style="color: rgb(168, 3, 20);">=</span> []</div><div>    all_values <span style="color: rgb(168, 3, 20);">=</span> []</div><div>    all_global_rows <span style="color: rgb(168, 3, 20);">=</span> [i <span style="color: rgb(168, 3, 20);">+</span> global_row_start <span style="color: rgb(168, 3, 20);">for</span> i <span style="color: rgb(168, 3, 20);">in</span> <span style="color: rgb(72, 31, 148);">range</span>(local_matrix_rows)]</div><br><div>    <span style="color: rgb(168, 3, 20);">for</span> i <span style="color: rgb(168, 3, 20);">in</span> <span style="color: rgb(72, 31, 148);">range</span>(<span style="color: rgb(72, 31, 148);">len</span>(all_global_rows)):</div><div>        cols, values <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">local_matrix</span>.getRow(i)</div><div>        <span style="color: rgb(164, 176, 188);"># print(f"cols: {cols}, values: {values}")</span></div><div>        all_cols.<span style="color: rgb(72, 31, 148);">append</span>(cols)</div><div>        all_values.<span style="color: rgb(72, 31, 148);">append</span>(values)</div><br><div>    <span style="color: rgb(168, 3, 20);">for</span> j <span style="color: rgb(168, 3, 20);">in</span> <span style="color: rgb(72, 31, 148);">range</span>(local_matrix_cols):</div><div>        columns <span style="color: rgb(168, 3, 20);">=</span> [all_cols[i][j] <span style="color: rgb(168, 3, 20);">for</span> i <span style="color: rgb(168, 3, 20);">in</span> <span style="color: rgb(72, 31, 148);">range</span>(<span style="color: rgb(72, 31, 148);">len</span>(all_cols))]</div><div>        values <span style="color: rgb(168, 3, 20);">=</span> [all_values[i][j] <span style="color: rgb(168, 3, 20);">for</span> i <span style="color: rgb(168, 3, 20);">in</span> <span style="color: rgb(72, 31, 148);">range</span>(<span style="color: rgb(72, 31, 148);">len</span>(all_values))]</div><br><div>        global_matrix.setValues(all_global_rows, columns, values)</div><br><div>    concatenate_end <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">time</span>.<span style="color: rgb(72, 31, 148);">time</span>()</div><div>    <span style="color: rgb(72, 31, 148);">Print</span>(</div><div>        <span style="color: rgb(168, 3, 20);">f</span><span style="color: rgb(10, 66, 131);">"  -Setting values: </span><span style="color: rgb(13, 92, 181);">{</span>concatenate_end <span style="color: rgb(168, 3, 20);">-</span> concatenate_start<span style="color: rgb(168, 3, 20);">: 2.2f</span><span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);"> s"</span>,</div><div>        Fore.GREEN,</div><div>    )</div><br><div>    global_matrix.assemblyBegin()</div><div>    global_matrix.assemblyEnd()</div><br><div>    <span style="color: rgb(168, 3, 20);">return</span> global_matrix</div><br><br><div><span style="color: rgb(164, 176, 188);"># --------------------------------------------</span></div><div><span style="color: rgb(164, 176, 188);"># EXP: Multiplication of an mpi PETSc matrix with a sequential PETSc matrix</span></div><div><span style="color: rgb(164, 176, 188);">#  C = A * B</span></div><div><span style="color: rgb(164, 176, 188);"># [m x k] = [m x k] * [k x k]</span></div><div><span style="color: rgb(164, 176, 188);"># --------------------------------------------</span></div><br><div>m, k <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">11</span>, <span style="color: rgb(13, 92, 181);">7</span></div><div><span style="color: rgb(164, 176, 188);"># Generate the random numpy matrices</span></div><div><span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">random</span>.seed(<span style="color: rgb(13, 92, 181);">0</span>)  <span style="color: rgb(164, 176, 188);"># sets the seed to 0</span></div><div>A_np <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">random</span>.randint(<span style="color: rgb(166, 61, 0);">low</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">0</span>, <span style="color: rgb(166, 61, 0);">high</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">6</span>, <span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>(m, k))</div><div>B_np <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">random</span>.randint(<span style="color: rgb(166, 61, 0);">low</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">0</span>, <span style="color: rgb(166, 61, 0);">high</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">6</span>, <span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>(k, k))</div><br><div><span style="color: rgb(164, 176, 188);"># Create B as a sequential matrix on each process</span></div><div>B_seq <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix_seq</span>(B_np)</div><div><br></div><div><span style="color: rgb(13, 92, 181);">A</span> <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix</span>(A_np)</div><div><br></div><div><span style="color: rgb(164, 176, 188);"># Getting the correct local submatrix to be multiplied by B_seq</span></div><div>A_local <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">get_local_submatrix</span>(<span style="color: rgb(13, 92, 181);">A</span>)</div><br><div><span style="color: rgb(164, 176, 188);"># Multiplication of 2 sequential matrices</span></div><div>C_local <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">multiply_matrices_seq</span>(A_local, B_seq)</div><br><div><span style="color: rgb(164, 176, 188);"># Creating the global C matrix</span></div><div><span style="color: rgb(13, 92, 181);">C</span> <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">concatenate_local_to_global_matrix</span>(C_local) <span style="color: rgb(168, 3, 20);">if</span> size <span style="color: rgb(168, 3, 20);">></span> <span style="color: rgb(13, 92, 181);">1</span> <span style="color: rgb(168, 3, 20);">else</span> C_local</div><div><br></div><div><span style="color: rgb(164, 176, 188);"># --------------------------------------------</span></div><div><span style="color: rgb(164, 176, 188);"># TEST: Multiplication of 2 numpy matrices</span></div><div><span style="color: rgb(164, 176, 188);"># --------------------------------------------</span></div><div>AB_np <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">dot</span>(A_np, B_np)</div><div><span style="color: rgb(72, 31, 148);">Print</span>(<span style="color: rgb(168, 3, 20);">f</span><span style="color: rgb(10, 66, 131);">"MATRIX AB_np [</span><span style="color: rgb(13, 92, 181);">{</span>AB_np.shape[<span style="color: rgb(13, 92, 181);">0</span>]<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);">x</span><span style="color: rgb(13, 92, 181);">{</span>AB_np.shape[<span style="color: rgb(13, 92, 181);">1</span>]<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);">]"</span>)</div><div><span style="color: rgb(72, 31, 148);">Print</span>(<span style="color: rgb(168, 3, 20);">f</span><span style="color: rgb(10, 66, 131);">"</span><span style="color: rgb(13, 92, 181);">{</span>AB_np<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);">"</span>)</div><br><div><span style="color: rgb(164, 176, 188);"># Get the local values from C</span></div><div>local_rows_start, local_rows_end <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">C</span>.getOwnershipRange()</div><div>C_local <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">C</span>.getValues(<span style="color: rgb(72, 31, 148);">range</span>(local_rows_start, local_rows_end), <span style="color: rgb(72, 31, 148);">range</span>(k))</div><br><div><span style="color: rgb(164, 176, 188);"># Assert the correctness of the multiplication for the local subset</span></div><div><span style="color: rgb(72, 31, 148);">assert_array_almost_equal</span>(C_local, AB_np[local_rows_start:local_rows_end, :], <span style="color: rgb(166, 61, 0);">decimal</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">5</span>)</div><br></div><div><br></div><div><br></div><div>Any idea how to fix this?</div><div><br></div><div>Thanks,</div><div>Thanos</div><div><br></div></div></body></html>