diff --git a/cpp/demo/hyperelasticity/main.cpp b/cpp/demo/hyperelasticity/main.cpp index 836ea23c0c..24a906107b 100644 --- a/cpp/demo/hyperelasticity/main.cpp +++ b/cpp/demo/hyperelasticity/main.cpp @@ -83,6 +83,229 @@ using namespace dolfinx; using T = PetscScalar; using U = typename dolfinx::scalar_value_t; +std::shared_ptr> create_mesh() +{ + return std::make_shared>(mesh::create_box( + MPI_COMM_WORLD, {{{0.0, 0.0, 0.0}, {1.0, 1.0, 1.0}}}, {10, 10, 10}, + mesh::CellType::tetrahedron, graph::partition_graph)); +} + +std::shared_ptr> +create_function_space(const std::shared_ptr>& mesh) +{ + auto element = basix::create_element( + basix::element::family::P, basix::cell::type::tetrahedron, 1, + basix::element::lagrange_variant::unset, + basix::element::dpc_variant::unset, false); + + return std::make_shared>(fem::create_functionspace( + mesh, std::make_shared>( + element, mesh->geometry().dim(), std::vector{3}))); +} + +std::shared_ptr> +create_rotation(const std::shared_ptr>& V) +{ + auto u_rotation = std::make_shared>(V); + u_rotation->interpolate( + [](auto x) -> std::pair, std::vector> + { + constexpr U scale = 0.005; + + // Center of rotation + constexpr U x1_c = 0.5; + constexpr U x2_c = 0.5; + + // Large angle of rotation (60 degrees) + constexpr U theta = std::numbers::pi / 3; + + // New coordinates + std::vector fdata(3 * x.extent(1), 0); + md::mdspan> f( + fdata.data(), 3, x.extent(1)); + for (std::size_t p = 0; p < x.extent(1); ++p) + { + U x1 = x(1, p); + U x2 = x(2, p); + f(1, p) = scale + * (x1_c + (x1 - x1_c) * std::cos(theta) + - (x2 - x2_c) * std::sin(theta) - x1); + f(2, p) = scale + * (x2_c + (x1 - x1_c) * std::sin(theta) + + (x2 - x2_c) * std::cos(theta) - x2); + } + + return {std::move(fdata), std::vector{3, x.extent(1)}}; + }); + + return u_rotation; +} + +std::vector> +create_boundary_conditions(const std::shared_ptr>& V, + const std::shared_ptr>& u_rotation) +{ + auto bdofs_left = fem::locate_dofs_geometrical( + *V, + [](auto x) + { + constexpr U eps = 1.0e-6; + std::vector marker(x.extent(1), false); + for (std::size_t p = 0; p < x.extent(1); ++p) + { + if (std::abs(x(0, p)) < eps) + marker[p] = true; + } + return marker; + }); + auto bdofs_right = fem::locate_dofs_geometrical( + *V, + [](auto x) + { + constexpr U eps = 1.0e-6; + std::vector marker(x.extent(1), false); + for (std::size_t p = 0; p < x.extent(1); ++p) + { + if (std::abs(x(0, p) - 1) < eps) + marker[p] = true; + } + return marker; + }); + + std::vector> bcs; + bcs.reserve(2); + bcs.emplace_back(std::vector{0, 0, 0}, bdofs_left, V); + bcs.emplace_back(u_rotation, bdofs_right); + return bcs; +} + +void post_process(const std::shared_ptr>& mesh, + const std::shared_ptr>& u) +{ + // Compute Cauchy stress. Construct appropriate Basix element for + // stress. + fem::Expression sigma_expression = fem::create_expression( + *expression_hyperelasticity_sigma, {{"u", u}}, {}, {}); + + constexpr auto family = basix::element::family::P; + auto cell_type = mesh::cell_type_to_basix_type(mesh->topology()->cell_type()); + constexpr int k = 0; + constexpr bool discontinuous = true; + basix::FiniteElement S_element = basix::create_element( + family, cell_type, k, basix::element::lagrange_variant::unset, + basix::element::dpc_variant::unset, discontinuous); + auto S = std::make_shared>(fem::create_functionspace( + mesh, + std::make_shared>(S_element, mesh->geometry().dim(), + std::vector{3, 3}))); + + fem::Function sigma(S); + sigma.name = "cauchy_stress"; + fem::interpolate(sigma, sigma_expression); + + // Save solution in VTK format + io::VTKFile file_u(mesh->comm(), "u.pvd", "w"); + file_u.write({*u}, 0); + + // Save Cauchy stress in XDMF format + io::XDMFFile file_sigma(mesh->comm(), "sigma.xdmf", "w"); + file_sigma.write_mesh(*mesh); + file_sigma.write_function(sigma, 0); +} + +void solve_hyperelasticity() +{ + // Inside `solve_hyperelasticity`, we begin by defining a tetrahedral + // mesh of the domain and the function space on this mesh. Here, we + // choose to create a unit cube mesh with 25 ( = 24 + 1) vertices in + // one direction and 17 ( = 16 + 1) vertices in the other two + // directions. With this mesh, we initialize the (finite element) + // function space defined by the generated code. + + // Create mesh and define function space + auto mesh = create_mesh(); + auto V = create_function_space(mesh); + + auto B = std::make_shared>(std::vector{0, 0, 0}); + auto traction = std::make_shared>(std::vector{0, 0, 0}); + + // Define solution function + auto u = std::make_shared>(V); + fem::Form a + = fem::create_form(*form_hyperelasticity_J_form, {V, V}, {{"u", u}}, + {{"B", B}, {"T", traction}}, {}, {}); + fem::Form L + = fem::create_form(*form_hyperelasticity_F_form, {V}, {{"u", u}}, + {{"B", B}, {"T", traction}}, {}, {}); + + auto u_rotation = create_rotation(V); + + // Create Dirichlet boundary conditions + auto bcs = create_boundary_conditions(V, u_rotation); + + // `A_layout` and `b_layout` set the layout of the Jacobian and + // residual that the solver works with. `u_vec` shares data with the + // degrees-of-freedom of `u`, and holds the initial guess on entry + // to the solve and the solution on return. + la::petsc::Matrix A_layout(fem::petsc::create_matrix(a, "aij"), false); + la::petsc::Vector b_layout( + la::petsc::create_vector(*V->dofmap()->index_map, + V->dofmap()->index_map_bs()), + false); + la::petsc::Vector u_vec(la::petsc::create_vector_wrap(*u->x()), false); + std::vector>> bcs_ref( + bcs.begin(), bcs.end()); + + // Create the solver, and attach the residual and Jacobian assembly. + // Each callback assembles at the point `x` into the `b` or `Jmat` + // it is passed, which may or may not be `b_layout.vec()` or + // `A_layout.mat()`: a line search, for instance, evaluates the + // residual in a work vector duplicated from `b_layout.vec()`. + nls::petsc::SNESSolver solver(mesh->comm()); + solver.set_F([&L, &a, &bcs_ref, &u](const Vec x, Vec b) + { fem::petsc::assemble_residual(x, b, L, a, bcs_ref, *u); }, + b_layout.vec()); + solver.set_J( + [&a, &bcs_ref, &u](const Vec x, Mat Jmat, Mat) + { fem::petsc::assemble_jacobian(x, Jmat, nullptr, a, bcs_ref, *u); }, + A_layout.mat()); + + // Begin configuring the solver through the PETSc options database. The + // Newton update is solved for with a direct LU solver, and a failure to + // converge raises an error rather than being reported by the return value. + const U tol = 10 * std::numeric_limits::epsilon(); + common::petsc::set_option("hyperelasticity_ksp_type", "preonly"); + common::petsc::set_option("hyperelasticity_pc_type", "lu"); + common::petsc::set_option("hyperelasticity_snes_rtol", tol); + common::petsc::set_option("hyperelasticity_snes_atol", tol); + common::petsc::set_option("hyperelasticity_snes_error_if_not_converged"); + + solver.set_options_prefix("hyperelasticity_"); + solver.set_from_options(); + + if (solver.solve(u_vec.vec()) < 0) + throw std::runtime_error("SNES solver did not converge."); + common::petsc::check( + VecGhostUpdateBegin(u_vec.vec(), INSERT_VALUES, SCATTER_FORWARD), + "VecGhostUpdateBegin"); + common::petsc::check( + VecGhostUpdateEnd(u_vec.vec(), INSERT_VALUES, SCATTER_FORWARD), + "VecGhostUpdateEnd"); + + // The SNES object is available for anything the solver does not + // wrap, here the number of Newton and linear solver iterations + PetscInt niter = 0; + common::petsc::check(SNESGetIterationNumber(solver.snes(), &niter), + "SNESGetIterationNumber"); + PetscInt lin_iter = 0; + common::petsc::check(SNESGetLinearSolveIterations(solver.snes(), &lin_iter), + "SNESGetLinearSolveIterations"); + std::cout << "Number of Newton iterations: " << niter << std::endl; + std::cout << "Number of linear solver iterations: " << lin_iter << std::endl; + + post_process(mesh, u); +} + int main(int argc, char* argv[]) { init_logging(argc, argv); @@ -94,198 +317,8 @@ int main(int argc, char* argv[]) std::string fmt = std::format("[%Y-%m-%d %H:%M:%S.%e] [RANK {}] [%l] %v", mpi_rank); spdlog::set_pattern(fmt); - { - // Inside the `main` function, we begin by defining a tetrahedral - // mesh of the domain and the function space on this mesh. Here, we - // choose to create a unit cube mesh with 25 ( = 24 + 1) vertices in - // one direction and 17 ( = 16 + 1) vertices in the other two - // directions. With this mesh, we initialize the (finite element) - // function space defined by the generated code. - - // Create mesh and define function space - auto mesh = std::make_shared>(mesh::create_box( - MPI_COMM_WORLD, {{{0.0, 0.0, 0.0}, {1.0, 1.0, 1.0}}}, {10, 10, 10}, - mesh::CellType::tetrahedron, graph::partition_graph)); - - auto element = basix::create_element( - basix::element::family::P, basix::cell::type::tetrahedron, 1, - basix::element::lagrange_variant::unset, - basix::element::dpc_variant::unset, false); - - auto V - = std::make_shared>(fem::create_functionspace( - mesh, - std::make_shared>( - element, mesh->geometry().dim(), std::vector{3}))); - auto B = std::make_shared>(std::vector{0, 0, 0}); - auto traction = std::make_shared>(std::vector{0, 0, 0}); - - // Define solution function - auto u = std::make_shared>(V); - fem::Form a - = fem::create_form(*form_hyperelasticity_J_form, {V, V}, {{"u", u}}, - {{"B", B}, {"T", traction}}, {}, {}); - fem::Form L - = fem::create_form(*form_hyperelasticity_F_form, {V}, {{"u", u}}, - {{"B", B}, {"T", traction}}, {}, {}); - - auto u_rotation = std::make_shared>(V); - u_rotation->interpolate( - [](auto x) -> std::pair, std::vector> - { - constexpr U scale = 0.005; - - // Center of rotation - constexpr U x1_c = 0.5; - constexpr U x2_c = 0.5; - - // Large angle of rotation (60 degrees) - constexpr U theta = std::numbers::pi / 3; - - // New coordinates - std::vector fdata(3 * x.extent(1), 0); - md::mdspan> f( - fdata.data(), 3, x.extent(1)); - for (std::size_t p = 0; p < x.extent(1); ++p) - { - U x1 = x(1, p); - U x2 = x(2, p); - f(1, p) = scale - * (x1_c + (x1 - x1_c) * std::cos(theta) - - (x2 - x2_c) * std::sin(theta) - x1); - f(2, p) = scale - * (x2_c + (x1 - x1_c) * std::sin(theta) - + (x2 - x2_c) * std::cos(theta) - x2); - } - - return {std::move(fdata), std::vector{3, x.extent(1)}}; - }); - - // Create Dirichlet boundary conditions - auto bdofs_left = fem::locate_dofs_geometrical( - *V, - [](auto x) - { - constexpr U eps = 1.0e-6; - std::vector marker(x.extent(1), false); - for (std::size_t p = 0; p < x.extent(1); ++p) - { - if (std::abs(x(0, p)) < eps) - marker[p] = true; - } - return marker; - }); - auto bdofs_right = fem::locate_dofs_geometrical( - *V, - [](auto x) - { - constexpr U eps = 1.0e-6; - std::vector marker(x.extent(1), false); - for (std::size_t p = 0; p < x.extent(1); ++p) - { - if (std::abs(x(0, p) - 1) < eps) - marker[p] = true; - } - return marker; - }); - std::vector bcs - = {fem::DirichletBC(std::vector{0, 0, 0}, bdofs_left, V), - fem::DirichletBC(u_rotation, bdofs_right)}; - - // `A_layout` and `b_layout` set the layout of the Jacobian and - // residual that the solver works with. `u_vec` shares data with the - // degrees-of-freedom of `u`, and holds the initial guess on entry - // to the solve and the solution on return. - la::petsc::Matrix A_layout(fem::petsc::create_matrix(a, "aij"), false); - la::petsc::Vector b_layout( - la::petsc::create_vector(*V->dofmap()->index_map, - V->dofmap()->index_map_bs()), - false); - la::petsc::Vector u_vec(la::petsc::create_vector_wrap(*u->x()), false); - std::vector>> bcs_ref( - bcs.begin(), bcs.end()); - - // Create the solver, and attach the residual and Jacobian assembly. - // Each callback assembles at the point `x` into the `b` or `Jmat` - // it is passed, which may or may not be `b_layout.vec()` or - // `A_layout.mat()`: a line search, for instance, evaluates the - // residual in a work vector duplicated from `b_layout.vec()`. - nls::petsc::SNESSolver solver(mesh->comm()); - solver.set_F([&L, &a, &bcs_ref, &u](const Vec x, Vec b) - { fem::petsc::assemble_residual(x, b, L, a, bcs_ref, *u); }, - b_layout.vec()); - solver.set_J( - [&a, &bcs_ref, &u](const Vec x, Mat Jmat, Mat) - { fem::petsc::assemble_jacobian(x, Jmat, nullptr, a, bcs_ref, *u); }, - A_layout.mat()); - - // Begin configuring the solver through the PETSc options database. The - // Newton update is solved for with a direct LU solver, and a failure to - // converge raises an error rather than being reported by the return value. - const U tol = 10 * std::numeric_limits::epsilon(); - common::petsc::set_option("hyperelasticity_ksp_type", "preonly"); - common::petsc::set_option("hyperelasticity_pc_type", "lu"); - common::petsc::set_option("hyperelasticity_snes_rtol", tol); - common::petsc::set_option("hyperelasticity_snes_atol", tol); - common::petsc::set_option("hyperelasticity_snes_error_if_not_converged"); - - solver.set_options_prefix("hyperelasticity_"); - solver.set_from_options(); - - if (solver.solve(u_vec.vec()) < 0) - throw std::runtime_error("SNES solver did not converge."); - common::petsc::check( - VecGhostUpdateBegin(u_vec.vec(), INSERT_VALUES, SCATTER_FORWARD), - "VecGhostUpdateBegin"); - common::petsc::check( - VecGhostUpdateEnd(u_vec.vec(), INSERT_VALUES, SCATTER_FORWARD), - "VecGhostUpdateEnd"); - - // The SNES object is available for anything the solver does not - // wrap, here the number of Newton and linear solver iterations - PetscInt niter = 0; - common::petsc::check(SNESGetIterationNumber(solver.snes(), &niter), - "SNESGetIterationNumber"); - PetscInt lin_iter = 0; - common::petsc::check(SNESGetLinearSolveIterations(solver.snes(), &lin_iter), - "SNESGetLinearSolveIterations"); - std::cout << "Number of Newton iterations: " << niter << std::endl; - std::cout << "Number of linear solver iterations: " << lin_iter - << std::endl; - - // Compute Cauchy stress. Construct appropriate Basix element for - // stress. - fem::Expression sigma_expression = fem::create_expression( - *expression_hyperelasticity_sigma, {{"u", u}}, {}, {}); - - constexpr auto family = basix::element::family::P; - auto cell_type - = mesh::cell_type_to_basix_type(mesh->topology()->cell_type()); - constexpr int k = 0; - constexpr bool discontinuous = true; - basix::FiniteElement S_element = basix::create_element( - family, cell_type, k, basix::element::lagrange_variant::unset, - basix::element::dpc_variant::unset, discontinuous); - auto S - = std::make_shared>(fem::create_functionspace( - mesh, std::make_shared>( - S_element, mesh->geometry().dim(), - std::vector{3, 3}))); - - fem::Function sigma(S); - sigma.name = "cauchy_stress"; - fem::interpolate(sigma, sigma_expression); - - // Save solution in VTK format - io::VTKFile file_u(mesh->comm(), "u.pvd", "w"); - file_u.write({*u}, 0); - - // Save Cauchy stress in XDMF format - io::XDMFFile file_sigma(mesh->comm(), "sigma.xdmf", "w"); - file_sigma.write_mesh(*mesh); - file_sigma.write_function(sigma, 0); - } + solve_hyperelasticity(); common::petsc::check(PetscFinalize(), "PetscFinalize"); return 0;