schur_test_diag_mat.cpp

Go to the documentation of this file.

/**
 * @file schur_test_diag_mat.cpp
 * @example mofem/atom_tests/schur_test_diag_mat.cpp
 * @brief Test block matrix and Schur complement matrix.
 *
 * @copyright Copyright (c) 2024
 *
 */
 
#include <MoFEM.hpp>
using namespace MoFEM;
 
static char help[] = "...\n\n";
 
template <int DIM> struct ElementsAndOps {};
 
template <> struct ElementsAndOps<2> {
  using DomainEle = PipelineManager::FaceEle;
};
 
template <> struct ElementsAndOps<3> {
  using DomainEle = VolumeElementForcesAndSourcesCore;
};
 
//! [Define dimension]
constexpr int SPACE_DIM = 3;
constexpr int FIELD_DIM = SPACE_DIM;
 
constexpr AssemblyType A =
    AssemblyType::BLOCK_MAT; //< selected assembly type
constexpr IntegrationType I =
    IntegrationType::GAUSS; //< selected integration type
 
using EntData = EntitiesFieldData::EntData;
using DomainEle = ElementsAndOps<SPACE_DIM>::DomainEle;
using DomainEleOp = DomainEle::UserDataOperator;
 
SmartPetscObj<Mat> petsc_mat;
SmartPetscObj<Mat> block_mat;
 
template <>
typename MoFEM::OpBaseImpl<PETSC, DomainEleOp>::MatSetValuesHook
    MoFEM::OpBaseImpl<PETSC, DomainEleOp>::matSetValuesHook =
        [](ForcesAndSourcesCore::UserDataOperator *op_ptr,
           const EntitiesFieldData::EntData &row_data,
           const EntitiesFieldData::EntData &col_data, MatrixDouble &m) {
          return MatSetValues<AssemblyTypeSelector<PETSC>>(
              petsc_mat, row_data, col_data, m, ADD_VALUES);
        };
 
template <>
typename MoFEM::OpBaseImpl<BLOCK_MAT, DomainEleOp>::MatSetValuesHook
    MoFEM::OpBaseImpl<BLOCK_MAT, DomainEleOp>::matSetValuesHook =
        [](ForcesAndSourcesCore::UserDataOperator *op_ptr,
           const EntitiesFieldData::EntData &row_data,
           const EntitiesFieldData::EntData &col_data, MatrixDouble &m) {
          return MatSetValues<AssemblyTypeSelector<BLOCK_MAT>>(
              block_mat, row_data, col_data, m, ADD_VALUES);
        };
 
template <>
typename MoFEM::OpBaseImpl<BLOCK_SCHUR, DomainEleOp>::MatSetValuesHook
    MoFEM::OpBaseImpl<BLOCK_SCHUR, DomainEleOp>::matSetValuesHook =
        [](ForcesAndSourcesCore::UserDataOperator *op_ptr,
           const EntitiesFieldData::EntData &row_data,
           const EntitiesFieldData::EntData &col_data, MatrixDouble &m) {
          return MatSetValues<AssemblyTypeSelector<BLOCK_SCHUR>>(
              block_mat, row_data, col_data, m, ADD_VALUES);
        };
 
template <>
typename MoFEM::OpBaseImpl<BLOCK_PRECONDITIONER_SCHUR,
                           DomainEleOp>::MatSetValuesHook
    MoFEM::OpBaseImpl<BLOCK_PRECONDITIONER_SCHUR,
                      DomainEleOp>::matSetValuesHook =
        [](ForcesAndSourcesCore::UserDataOperator *op_ptr,
           const EntitiesFieldData::EntData &row_data,
           const EntitiesFieldData::EntData &col_data, MatrixDouble &m) {
          return MatSetValues<AssemblyTypeSelector<BLOCK_PRECONDITIONER_SCHUR>>(
              block_mat, row_data, col_data, m, ADD_VALUES);
        };
 
constexpr bool debug = false;
 
int main(int argc, char *argv[]) {
 
  // initialize petsc
  MoFEM::Core::Initialize(&argc, &argv, (char *)0, help);
 
  try {
 
    // Create MoAB database
    moab::Core moab_core;
    moab::Interface &moab = moab_core;
 
    // Create MoFEM database and link it to MoAB
    MoFEM::Core mofem_core(moab);
    MoFEM::Interface &m_field = mofem_core;
 
    // Register DM Manager
    DMType dm_name = "DMMOFEM";
    CHKERR DMRegister_MoFEM(dm_name);
 
    // Add logging channel for example
    auto core_log = logging::core::get();
    core_log->add_sink(
        LogManager::createSink(LogManager::getStrmWorld(), "Timeline"));
    LogManager::setLog("Timeline");
    MOFEM_LOG_TAG("Timeline", "Timeline");
 
    // Simple interface
    auto simple = m_field.getInterface<Simple>();
 
    // get options from command line
    CHKERR simple->getOptions();
    // load mesh file
    CHKERR simple->loadFile();
 
    CHKERR simple->addDomainField("V", H1, AINSWORTH_LEGENDRE_BASE, FIELD_DIM);
    CHKERR simple->addDomainField("T", L2, AINSWORTH_LEGENDRE_BASE, FIELD_DIM);
    CHKERR simple->addDomainField("S", L2, AINSWORTH_LEGENDRE_BASE, FIELD_DIM);
    CHKERR simple->addDomainField("O", L2, AINSWORTH_LEGENDRE_BASE, FIELD_DIM);
 
    // set fields order, i.e. for most first cases order is sufficient.
    int order = 3;
    CHKERR PetscOptionsGetInt(PETSC_NULLPTR, "", "-order", &order, PETSC_NULLPTR);
    CHKERR simple->setFieldOrder("V", order);
    CHKERR simple->setFieldOrder("T", order);
    CHKERR simple->setFieldOrder("S", order - 1);
    CHKERR simple->setFieldOrder("O", order - 2);
 
    // setup problem
    CHKERR simple->setUp();
 
    // If block "VOL" exist, first index is removed from "T" field. That test
    // blocks handling when some DOFs are removed.
    auto bc_mng = m_field.getInterface<BcManager>();
    CHKERR bc_mng->removeBlockDOFsOnEntities(simple->getProblemName(), "VOL",
                                             "T", 0, 1);
 
    auto schur_dm = createDM(m_field.get_comm(), "DMMOFEM");
    CHKERR DMMoFEMCreateSubDM(schur_dm, simple->getDM(), "SCHUR");
    CHKERR DMMoFEMSetSquareProblem(schur_dm, PETSC_TRUE);
    CHKERR DMMoFEMAddElement(schur_dm, simple->getDomainFEName());
    CHKERR DMMoFEMAddSubFieldRow(schur_dm, "V");
    CHKERR DMMoFEMAddSubFieldCol(schur_dm, "V");
    CHKERR DMSetUp(schur_dm);
 
    auto block_dm = createDM(m_field.get_comm(), "DMMOFEM");
    CHKERR DMMoFEMCreateSubDM(block_dm, simple->getDM(), "BLOCK");
    CHKERR DMMoFEMSetSquareProblem(block_dm, PETSC_TRUE);
    CHKERR DMMoFEMAddElement(block_dm, simple->getDomainFEName());
    CHKERR DMMoFEMAddSubFieldRow(block_dm, "T");
    CHKERR DMMoFEMAddSubFieldCol(block_dm, "T");
    CHKERR DMMoFEMAddSubFieldRow(block_dm, "S");
    CHKERR DMMoFEMAddSubFieldCol(block_dm, "S");
    CHKERR DMMoFEMAddSubFieldRow(block_dm, "O");
    CHKERR DMMoFEMAddSubFieldCol(block_dm, "O");
    CHKERR DMSetUp(block_dm);
 
    petsc_mat = createDMMatrix(simple->getDM());
    auto S = createDMMatrix(schur_dm);
 
    auto shell_data =
        createBlockMat(simple->getDM(),
 
                       createBlockMatStructure(simple->getDM(),
 
                                               {{simple->getDomainFEName(),
 
                                                 {{"V", "V"},
                                                  {"T", "T"},
                                                  {"S", "S"},
                                                  {"O", "O"},
                                                  {"V", "T"},
                                                  {"T", "V"},
                                                  {"S", "T"},
                                                  {"T", "S"},
                                                  {"T", "O"},
                                                  {"O", "T"},
                                                  {"S", "O"},
                                                  {"O", "S"}
 
                                                 }}}
 
                                               )
 
        );
 
    auto [mat, block_data_ptr] = shell_data;
    block_mat = mat;
 
    std::vector<std::string> fields{"T", "S", "O"};
 
    auto [nested_mat, nested_data_ptr] = createSchurNestedMatrix(
 
        createSchurNestedMatrixStruture(
 
            {schur_dm, block_dm}, block_data_ptr,
 
            fields, {nullptr, nullptr, nullptr}, true
 
            )
 
    );
 
    using OpMassPETSCAssemble = FormsIntegrators<DomainEleOp>::Assembly<
        PETSC>::BiLinearForm<GAUSS>::OpMass<1, FIELD_DIM>;
    using OpMassBlockAssemble = FormsIntegrators<DomainEleOp>::Assembly<
        BLOCK_SCHUR>::BiLinearForm<GAUSS>::OpMass<1, FIELD_DIM>;
    using OpMassBlockPreconditionerAssemble =
        FormsIntegrators<DomainEleOp>::Assembly<BLOCK_PRECONDITIONER_SCHUR>::
            BiLinearForm<GAUSS>::OpMass<1, FIELD_DIM>;
 
    // get operators tester
    auto pip_mng = m_field.getInterface<PipelineManager>(); // get interface to
                                                        // pipeline manager
 
    auto close_zero = [](double, double, double) { return 1; };
    auto beta = [](double, double, double) { return -1./2; };
    auto gamma = [](double, double, double) { return -1. / 4; };
 
    pip_mng->setDomainLhsIntegrationRule(
        [](int, int, int o) { return 2 * o; });
    auto &pip_lhs = pip_mng->getOpDomainLhsPipeline();
 
    pip_lhs.push_back(new OpMassPETSCAssemble("V", "V"));
    // pip_lhs.push_back(new OpMassPETSCAssemble("T", "T"));
    pip_lhs.push_back(new OpMassPETSCAssemble("V", "T"));
    pip_lhs.push_back(new OpMassPETSCAssemble("T", "V"));
    pip_lhs.push_back(new OpMassPETSCAssemble("S", "S", close_zero));
    pip_lhs.push_back(new OpMassPETSCAssemble("S", "T", beta));
    pip_lhs.push_back(new OpMassPETSCAssemble("T", "S", beta));
    pip_lhs.push_back(new OpMassPETSCAssemble("O", "O", close_zero));
    pip_lhs.push_back(new OpMassPETSCAssemble("T", "O", beta));
    pip_lhs.push_back(new OpMassPETSCAssemble("O", "T", beta));
    pip_lhs.push_back(new OpMassPETSCAssemble("S", "O", gamma));
    pip_lhs.push_back(new OpMassPETSCAssemble("O", "S", beta));
 
    pip_lhs.push_back(createOpSchurAssembleBegin());
    pip_lhs.push_back(new OpMassBlockAssemble("V", "V"));
    // pip_lhs.push_back(new OpMassBlockAssemble("T", "T"));
    pip_lhs.push_back(new OpMassBlockAssemble("V", "T"));
    pip_lhs.push_back(new OpMassBlockAssemble("T", "V"));
    pip_lhs.push_back(new OpMassBlockAssemble("S", "S", close_zero));
    pip_lhs.push_back(new OpMassBlockAssemble("S", "T", beta));
    pip_lhs.push_back(new OpMassBlockAssemble("T", "S", beta));
    pip_lhs.push_back(new OpMassBlockAssemble("O", "O", close_zero));
    pip_lhs.push_back(new OpMassBlockAssemble("T", "O", beta));
    pip_lhs.push_back(new OpMassBlockAssemble("O", "T", beta));
    pip_lhs.push_back(new OpMassBlockAssemble("S", "O", gamma));
    pip_lhs.push_back(new OpMassBlockAssemble("O", "S", beta));
    pip_lhs.push_back(new OpMassBlockPreconditionerAssemble("T", "T"));
 
    auto schur_is = getDMSubData(schur_dm)->getSmartRowIs();
    auto ao_up = createAOMappingIS(schur_is, PETSC_NULLPTR);
 
    pip_lhs.push_back(createOpSchurAssembleEnd(
 
        fields, {nullptr, nullptr, nullptr}, ao_up, S, true, true)
 
    );
 
    {
      MOFEM_LOG_CHANNEL("Timeline");
      MOFEM_LOG_TAG("Timeline", "timer");
      BOOST_LOG_SCOPED_THREAD_ATTR("Timeline", attrs::timer());
      MOFEM_LOG("Timeline", Sev::inform) << "Assemble start";
      CHKERR DMoFEMLoopFiniteElements(simple->getDM(),
                                      simple->getDomainFEName(),
                                      pip_mng->getDomainLhsFE());
      MOFEM_LOG("Timeline", Sev::inform) << "Assemble end";
    }
    {
      MOFEM_LOG_CHANNEL("Timeline");
      MOFEM_LOG_TAG("Timeline", "timer");
      BOOST_LOG_SCOPED_THREAD_ATTR("Timeline", attrs::timer());
      MOFEM_LOG("Timeline", Sev::inform) << "Mat assemble start";
      CHKERR MatAssemblyBegin(petsc_mat, MAT_FINAL_ASSEMBLY);
      CHKERR MatAssemblyEnd(petsc_mat, MAT_FINAL_ASSEMBLY);
      MOFEM_LOG("Timeline", Sev::inform) << "Mat assemble end";
    }
 
    auto get_random_vector = [&](auto dm) {
      auto v = createDMVector(dm);
      PetscRandom rctx;
      PetscRandomCreate(PETSC_COMM_WORLD, &rctx);
      CHK_MOAB_THROW(VecSetRandom(v, rctx), "generate rand vector");
      PetscRandomDestroy(&rctx);
      return v;
    };
    auto v = get_random_vector(simple->getDM());
 
    auto y_petsc = createDMVector(simple->getDM());
    auto y_block = createDMVector(simple->getDM());
 
    auto test = [](auto msg, auto y, double norm0) {
      MoFEMFunctionBegin;
 
    double eps = 1e-10;
    CHKERR PetscOptionsGetScalar(PETSC_NULLPTR, "", "-eps", &eps, PETSC_NULLPTR);
 
      PetscReal norm;
      CHKERR VecNorm(y, NORM_2, &norm);
      norm = norm / norm0;
      MOFEM_LOG_CHANNEL("WORLD");
      MOFEM_TAG_AND_LOG("WORLD", Sev::inform, "TestBlockMat")
          << msg << ": norm of difference: " << norm;
      if (norm > eps || std::isnan(norm) || std::isinf(norm)) {
        SETERRQ(PETSC_COMM_WORLD, 1, "norm of difference is too big");
      }
      MoFEMFunctionReturn(0);
    };
 
    std::vector<int> zero_rows_and_cols = {
        0, 1, 10, 20,
        500}; // not to remove dofs for TENSOR field, inverse will not work
    CHKERR MatZeroRowsColumns(petsc_mat, zero_rows_and_cols.size(),
                              &*zero_rows_and_cols.begin(), 1, PETSC_NULLPTR,
                              PETSC_NULLPTR);
    CHKERR MatZeroRowsColumns(block_mat, zero_rows_and_cols.size(),
                              &*zero_rows_and_cols.begin(), 1, PETSC_NULLPTR,
                              PETSC_NULLPTR);
 
    // test mult
    {
      MOFEM_LOG_CHANNEL("Timeline");
      MOFEM_LOG_TAG("Timeline", "timer");
      BOOST_LOG_SCOPED_THREAD_ATTR("Timeline", attrs::timer());
      MOFEM_LOG("Timeline", Sev::inform)
          << "MatMult(petsc_mat, v, y_petsc) star";
      CHKERR MatMult(petsc_mat, v, y_petsc);
      MOFEM_LOG("Timeline", Sev::inform)
          << "MatMult(petsc_mat, v, y_petsc) end";
    }
    double nrm0;
    CHKERR VecNorm(y_petsc, NORM_2, &nrm0);
 
    {
      MOFEM_LOG_CHANNEL("Timeline");
      MOFEM_LOG_TAG("Timeline", "timer");
      BOOST_LOG_SCOPED_THREAD_ATTR("Timeline", attrs::timer());
      MOFEM_LOG("Timeline", Sev::inform)
          << "MatMult(block_mat, v, y_block) star";
      CHKERR MatMult(block_mat, v, y_block);
      MOFEM_LOG("Timeline", Sev::inform)
          << "MatMult(block_mat, v, y_block) end";
    }
    CHKERR VecAXPY(y_petsc, -1.0, y_block);
    CHKERR test("mult", y_petsc, nrm0);
 
    // tets mult add
 
    MOFEM_LOG_CHANNEL("Timeline");
    MOFEM_LOG_TAG("Timeline", "timer");
 
    CHKERR MatMult(petsc_mat, v, y_petsc);
    CHKERR MatMult(block_mat, v, y_block);
    CHKERR MatMultAdd(petsc_mat, v, y_petsc, y_petsc);
    CHKERR MatMultAdd(block_mat, v, y_block, y_block);
    CHKERR VecAXPY(y_petsc, -1.0, y_block);
    CHKERR test("mult add", y_petsc, nrm0);
 
    // test mult transpose
    CHKERR MatMultTranspose(petsc_mat, v, y_petsc);
    CHKERR MatMultTranspose(block_mat, v, y_block);
    CHKERR VecAXPY(y_petsc, -1.0, y_block);
    CHKERR test("mult transpose", y_petsc, nrm0);
 
    // tets mult add transpose
    CHKERR MatMultTranspose(petsc_mat, v, y_petsc);
    CHKERR MatMultTranspose(block_mat, v, y_block);
    CHKERR MatMultTransposeAdd(petsc_mat, v, y_petsc, y_petsc);
    CHKERR MatMultTransposeAdd(block_mat, v, y_block, y_block);
    CHKERR VecAXPY(y_petsc, -1.0, y_block);
    CHKERR test("mult add transpose", y_petsc, nrm0);
    
    // test schur nested mat with preconditioner
    CHKERR schurSwitchPreconditioner(std::get<1>(*nested_data_ptr)[3]);
    auto y_nested = createDMVector(simple->getDM());
    CHKERR MatMult(petsc_mat, v, y_petsc);
    CHKERR MatMult(nested_mat, v, y_nested);
    CHKERR VecAXPY(y_petsc, -1.0, y_nested);
    CHKERR test("mult nested", y_petsc, nrm0);
 
    CHKERR schurSwitchPreconditioner(std::get<1>(*nested_data_ptr)[3]);
    auto diag_mat = std::get<0>(*nested_data_ptr)[3];
    auto diag_block_x = get_random_vector(block_dm);
    auto diag_block_f = createDMVector(block_dm);
    auto block_solved_x = createDMVector(block_dm);
    CHKERR MatMult(diag_mat, diag_block_x, diag_block_f);
    
    // That is if one like to use MatSolve directly, not though PC, as it is
    // below 
    CHKERR MatSolve(diag_mat, diag_block_f, block_solved_x);
    auto diag_block_f_test = createDMVector(block_dm);
    CHKERR MatMult(diag_mat, block_solved_x, diag_block_f_test);
    CHKERR VecAXPY(diag_block_f_test, -1.0, diag_block_f); 
    CHKERR test("diag solve", diag_block_f_test, nrm0);
    // Transpose solve test
    CHKERR MatSolveTranspose(diag_mat, diag_block_f, block_solved_x);
    CHKERR MatMultTranspose(diag_mat, block_solved_x, diag_block_f_test);
    CHKERR VecAXPY(diag_block_f_test, -1.0, diag_block_f); 
    CHKERR test("diag transpose solve", diag_block_f_test, nrm0);
 
    // set matrix type to shell, set data
    CHKERR DMSetMatType(block_dm, MATSHELL);
    CHKERR DMMoFEMSetBlocMatData(block_dm, std::get<1>(*nested_data_ptr)[3]);
    // set empty operator, since block data are already calculated
    CHKERR DMKSPSetComputeOperators(
        block_dm,
        [](KSP, Mat, Mat, void *) {
          MOFEM_LOG("WORLD", Sev::inform) << "empty operator";
          return 0;
        },
        nullptr);
 
    auto ksp = createKSP(m_field.get_comm());
    CHKERR KSPSetDM(ksp, block_dm);
    CHKERR KSPSetFromOptions(ksp);
 
    // set preconditioner to block mat
    auto get_pc = [](auto ksp) {
      PC pc_raw;
      CHKERR KSPGetPC(ksp, &pc_raw);
      return SmartPetscObj<PC>(pc_raw, true); // bump reference
    };
    CHKERR setSchurA00MatSolvePC(get_pc(ksp));
    CHKERR KSPSetUp(ksp);
 
    CHKERR VecZeroEntries(block_solved_x);
    CHKERR KSPSolve(ksp, diag_block_f, block_solved_x);
 
    CHKERR MatMult(diag_mat, block_solved_x, diag_block_f_test);
    CHKERR VecAXPY(diag_block_f_test, -1.0, diag_block_f);
    CHKERR test("diag solve", diag_block_f_test, nrm0);
 
    // solve transpose system
    CHKERR VecZeroEntries(block_solved_x);
    CHKERR KSPSolveTranspose(ksp, diag_block_f, block_solved_x);
    CHKERR MatMultTranspose(diag_mat, block_solved_x, diag_block_f_test);
    CHKERR VecAXPY(diag_block_f_test, -1.0, diag_block_f);
    CHKERR test("diag transpose solve", diag_block_f_test, nrm0);
 
    if (m_field.get_comm_rank() == 0) {
      CHKERR schurSaveBlockMesh(block_data_ptr, "block_mesh.vtk");
    }
 
    petsc_mat.reset();
    block_mat.reset();
  }
  CATCH_ERRORS;
 
  // finish work cleaning memory, getting statistics, etc.
  MoFEM::Core::Finalize();
 
  return 0;
}