Skip to content

Add matrix-based implicit solve option to Adams-Bashforth-Moulton #76

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: devel
Choose a base branch
from
Open

Add matrix-based implicit solve option to Adams-Bashforth-Moulton #76

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions include/tensor_solver/AdamsBashforthMoulton.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::vector<const torch::Tensor *>> _linear_matrix;

unsigned int _substeps;
std::size_t _predictor_order;
std::size_t _corrector_order;
Expand Down
189 changes: 151 additions & 38 deletions src/tensor_solver/AdamsBashforthMoulton.C
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<MooseEnum>("implicit_mode",
implicit_mode,
"Implicit treatment of the linear term (diagonal or full matrix).");
params.addParam<std::vector<TensorInputBufferName>>(
"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<MooseEnum>("implicit_mode")),
_substeps(getParam<unsigned int>("substeps")),
_predictor_order(getParam<std::size_t>("predictor_order") - 1),
_corrector_order(getParam<std::size_t>("corrector_order") - 1),
Expand All @@ -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<std::vector<TensorInputBufferName>>("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<const torch::Tensor *>(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
Expand Down Expand Up @@ -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<torch::Tensor> 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<torch::Tensor> rows(n);
for (std::size_t i = 0; i < n; ++i)
{
std::vector<torch::Tensor> 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])
Expand Down Expand Up @@ -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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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);
}
}
}
}
Expand Down