diff --git a/include/tensor_solver/AdamsBashforthMoulton.h b/include/tensor_solver/AdamsBashforthMoulton.h index 7517bb1b..1420db26 100644 --- a/include/tensor_solver/AdamsBashforthMoulton.h +++ b/include/tensor_solver/AdamsBashforthMoulton.h @@ -23,6 +23,11 @@ class AdamsBashforthMoulton : public SplitOperatorBase virtual void computeBuffer() override; protected: + /// implicit solve mode: diagonal (element-wise divide) or matrix (linear solve) + MooseEnum _implicit_mode; + /// optional full linear operator matrix buffers (row-major) + std::vector> _linear_matrix; + unsigned int _substeps; std::size_t _predictor_order; std::size_t _corrector_order; diff --git a/src/tensor_solver/AdamsBashforthMoulton.C b/src/tensor_solver/AdamsBashforthMoulton.C index 49e7d22c..b7bd0ce1 100644 --- a/src/tensor_solver/AdamsBashforthMoulton.C +++ b/src/tensor_solver/AdamsBashforthMoulton.C @@ -45,11 +45,21 @@ AdamsBashforthMoulton::validParams() "corrector_steps", 0, "Number the Adams-Moulton corrector steps to take (one is usually sufficient)."); + + MooseEnum implicit_mode("diagonal matrix", "diagonal"); + params.addParam("implicit_mode", + implicit_mode, + "Implicit treatment of the linear term (diagonal or full matrix)."); + params.addParam>( + "linear_matrix", + {}, + "Row-major list of buffers defining the full linear operator matrix for implicit solves."); return params; } AdamsBashforthMoulton::AdamsBashforthMoulton(const InputParameters & parameters) : SplitOperatorBase(parameters), + _implicit_mode(getParam("implicit_mode")), _substeps(getParam("substeps")), _predictor_order(getParam("predictor_order") - 1), _corrector_order(getParam("corrector_order") - 1), @@ -58,6 +68,18 @@ AdamsBashforthMoulton::AdamsBashforthMoulton(const InputParameters & parameters) _sub_time(_tensor_problem.subTime()) { getVariables(_predictor_order); + + if (_implicit_mode == MooseEnum("matrix")) + { + const auto names = getParam>("linear_matrix"); + const auto n = _variables.size(); + if (names.size() != n * n) + paramError("linear_matrix", "The 'linear_matrix' parameter must contain n*n entries."); + _linear_matrix.resize(n, std::vector(n)); + for (std::size_t i = 0; i < n; ++i) + for (std::size_t j = 0; j < n; ++j) + _linear_matrix[i][j] = &getInputBufferByName(names[i * n + j]); + } } void @@ -92,30 +114,71 @@ AdamsBashforthMoulton::computeBuffer() // re-evaluate the solve compute _compute->computeBuffer(); forwardBuffers(); + if (_implicit_mode == 1) + { + const auto n = _variables.size(); + std::vector rhs(n); + for (const auto k : index_range(_variables)) + { + const auto & reciprocal_buffer = _variables[k]._reciprocal_buffer; + const auto & nonlinear_reciprocal = _variables[k]._nonlinear_reciprocal; + const auto & old_nonlinear_reciprocal = _variables[k]._old_nonlinear_reciprocal; + const auto n_old = old_nonlinear_reciprocal.size(); + const auto order = + std::min(substep < _predictor_order && dt_changed ? 0 : n_old, _predictor_order); + ubar = reciprocal_buffer + (_sub_dt * beta[order][0]) * nonlinear_reciprocal; + for (const auto i : make_range(order)) + ubar += (_sub_dt * beta[order][i + 1]) * old_nonlinear_reciprocal[i]; + rhs[k] = ubar; + } - // Adams-Bashforth predictor on all variables - for (auto & [u, - reciprocal_buffer, - linear_reciprocal, - nonlinear_reciprocal, - old_nonlinear_reciprocal] : _variables) + std::vector rows(n); + for (std::size_t i = 0; i < n; ++i) + { + std::vector row(n); + for (std::size_t j = 0; j < n; ++j) + row[j] = *(_linear_matrix[i][j]); + rows[i] = torch::stack(row); + } + auto L = torch::stack(rows); + auto I = torch::eye(n, L.options()); + for (int d = 2; d < L.dim(); ++d) + I = I.unsqueeze(-1); + auto M = I - _sub_dt * L; + auto B = torch::stack(rhs); + auto M_flat = M.view({n, n, -1}).permute({2, 0, 1}); + auto B_flat = B.view({n, -1}).permute({1, 0}); + auto U_flat = torch::linalg::solve(M_flat, B_flat); + auto U = U_flat.permute({1, 0}).view_as(B); + for (const auto k : index_range(_variables)) + _variables[k]._buffer = _domain.ifft(U[k]); + } + else { - const auto n_old = old_nonlinear_reciprocal.size(); + // Adams-Bashforth predictor on all variables + for (auto & [u, + reciprocal_buffer, + linear_reciprocal, + nonlinear_reciprocal, + old_nonlinear_reciprocal] : _variables) + { + const auto n_old = old_nonlinear_reciprocal.size(); - // Order is what the user requested, or what the available history allows for. - // If dt changes between steps, we start at first order again - const auto order = - std::min(substep < _predictor_order && dt_changed ? 0 : n_old, _predictor_order); + // Order is what the user requested, or what the available history allows for. + // If dt changes between steps, we start at first order again + const auto order = + std::min(substep < _predictor_order && dt_changed ? 0 : n_old, _predictor_order); - // Adams-Bashforth - ubar = reciprocal_buffer + (_sub_dt * beta[order][0]) * nonlinear_reciprocal; - for (const auto i : make_range(order)) - ubar += (_sub_dt * beta[order][i + 1]) * old_nonlinear_reciprocal[i]; + // Adams-Bashforth + ubar = reciprocal_buffer + (_sub_dt * beta[order][0]) * nonlinear_reciprocal; + for (const auto i : make_range(order)) + ubar += (_sub_dt * beta[order][i + 1]) * old_nonlinear_reciprocal[i]; - if (linear_reciprocal) - ubar /= (1.0 - _sub_dt * *linear_reciprocal); + if (linear_reciprocal) + ubar /= (1.0 - _sub_dt * *linear_reciprocal); - u = _domain.ifft(ubar); + u = _domain.ifft(ubar); + } } // AB: y[n+1] = y[n] + dt * f(y[n]) @@ -150,32 +213,82 @@ AdamsBashforthMoulton::computeBuffer() _compute->computeBuffer(); forwardBuffers(); - for (const auto k : index_range(_variables)) + if (_implicit_mode == 1) { - auto & u = _variables[k]._buffer; - const auto * linear_reciprocal = _variables[k]._linear_reciprocal; - const auto & nonlinear_reciprocal_pred = _variables[k]._nonlinear_reciprocal; - const auto & old_nonlinear_reciprocal = _variables[k]._old_nonlinear_reciprocal; - - const auto n_old = old_nonlinear_reciprocal.size(); - const auto order = - std::min(substep < _corrector_order && dt_changed ? 1 : n_old + 1, _corrector_order); - if (order == 0) - continue; - - ubar = ubar_n[k] + (_sub_dt * alpha[order][0]) * nonlinear_reciprocal_pred; + const auto n = _variables.size(); + std::vector rhs(n); + for (const auto k : index_range(_variables)) + { + const auto & nonlinear_reciprocal_pred = _variables[k]._nonlinear_reciprocal; + const auto & old_nonlinear_reciprocal = _variables[k]._old_nonlinear_reciprocal; + const auto n_old = old_nonlinear_reciprocal.size(); + const auto order = std::min(substep < _corrector_order && dt_changed ? 1 : n_old + 1, + _corrector_order); + if (order == 0) + { + rhs[k] = ubar_n[k]; + continue; + } + ubar = ubar_n[k] + (_sub_dt * alpha[order][0]) * nonlinear_reciprocal_pred; + if (order > 0) + { + ubar += (_sub_dt * alpha[order][1]) * N_n[k]; + for (const auto i : make_range(order - 1)) + ubar += (_sub_dt * alpha[order][i + 2]) * old_nonlinear_reciprocal[i]; + } + rhs[k] = ubar; + } - if (order > 0) + std::vector rows(n); + for (std::size_t i = 0; i < n; ++i) { - ubar += (_sub_dt * alpha[order][1]) * N_n[k]; - for (const auto i : make_range(order - 1)) - ubar += (_sub_dt * alpha[order][i + 2]) * old_nonlinear_reciprocal[i]; + std::vector row(n); + for (std::size_t j2 = 0; j2 < n; ++j2) + row[j2] = *(_linear_matrix[i][j2]); + rows[i] = torch::stack(row); } + auto L = torch::stack(rows); + auto I = torch::eye(n, L.options()); + for (int d = 2; d < L.dim(); ++d) + I = I.unsqueeze(-1); + auto M = I - _sub_dt * L; + auto B = torch::stack(rhs); + auto M_flat = M.view({n, n, -1}).permute({2, 0, 1}); + auto B_flat = B.view({n, -1}).permute({1, 0}); + auto U_flat = torch::linalg::solve(M_flat, B_flat); + auto U = U_flat.permute({1, 0}).view_as(B); + for (const auto k : index_range(_variables)) + _variables[k]._buffer = _domain.ifft(U[k]); + } + else + { + for (const auto k : index_range(_variables)) + { + auto & u = _variables[k]._buffer; + const auto * linear_reciprocal = _variables[k]._linear_reciprocal; + const auto & nonlinear_reciprocal_pred = _variables[k]._nonlinear_reciprocal; + const auto & old_nonlinear_reciprocal = _variables[k]._old_nonlinear_reciprocal; + + const auto n_old = old_nonlinear_reciprocal.size(); + const auto order = std::min(substep < _corrector_order && dt_changed ? 1 : n_old + 1, + _corrector_order); + if (order == 0) + continue; - if (linear_reciprocal) - ubar /= (1.0 - _sub_dt * *linear_reciprocal); + ubar = ubar_n[k] + (_sub_dt * alpha[order][0]) * nonlinear_reciprocal_pred; - u = _domain.ifft(ubar); + if (order > 0) + { + ubar += (_sub_dt * alpha[order][1]) * N_n[k]; + for (const auto i : make_range(order - 1)) + ubar += (_sub_dt * alpha[order][i + 2]) * old_nonlinear_reciprocal[i]; + } + + if (linear_reciprocal) + ubar /= (1.0 - _sub_dt * *linear_reciprocal); + + u = _domain.ifft(ubar); + } } } }