diffusion.hpp

Go to the documentation of this file.

/// @file pde/diffusion.hpp
/// @brief Diffusion operators and implicit system builders for 2D grids.
#pragma once
 
#include "core/policy.hpp"
#include "core/vector.hpp"
#include "fields/grid2d.hpp"
#include "linalg/solvers/cg.hpp"
#include "linalg/solvers/linear_solver.hpp"
#include "linalg/sparse/sparse.hpp"
#include "pde/stencil.hpp"
 
namespace num::pde {
 
inline void diffusion_step_2d(Vector& u, int N, double coeff, Backend b = best_backend) {
    Vector lap(u.size());
    laplacian_stencil_2d_periodic(u, lap, N);
    axpy(coeff, lap, u, b);
}
 
inline void diffusion_step_2d_dirichlet(Vector& u,
                                        int N,
                                        double coeff,
                                        Backend b = best_backend) {
    Vector lap(u.size());
    laplacian_stencil_2d(u, lap, N);
    axpy(coeff, lap, u, b);
}
 
inline void diffusion_step_2d_4th_dirichlet(Vector& u,
                                            int N,
                                            double coeff,
                                            Backend b = best_backend) {
    Vector lap(u.size());
    laplacian_stencil_2d_4th(u, lap, N);
    axpy(coeff, lap, u, b);
}
 
inline void diffusion_step_2d_dirichlet(ScalarField2D& g,
                                        double coeff,
                                        Backend b = best_backend) {
    diffusion_step_2d_dirichlet(g.vec(), g.N(), coeff, b);
}
 
inline void diffusion_step_2d_4th_dirichlet(ScalarField2D& g,
                                            double coeff,
                                            Backend b = best_backend) {
    diffusion_step_2d_4th_dirichlet(g.vec(), g.N(), coeff, b);
}
 
inline SparseMatrix laplacian_sparse_2d(int N) {
    const int n = N * N;
    std::vector<idx> rows, cols;
    std::vector<real> vals;
    rows.reserve(5 * n);
    cols.reserve(5 * n);
    vals.reserve(5 * n);
    for (int i = 0; i < N; ++i) {
        for (int j = 0; j < N; ++j) {
            int k = i * N + j;
            rows.push_back(k);
            cols.push_back(k);
            vals.push_back(-4.0);
            if (i > 0) {
                rows.push_back(k);
                cols.push_back((i - 1) * N + j);
                vals.push_back(1.0);
            }
            if (i < N - 1) {
                rows.push_back(k);
                cols.push_back((i + 1) * N + j);
                vals.push_back(1.0);
            }
            if (j > 0) {
                rows.push_back(k);
                cols.push_back(i * N + (j - 1));
                vals.push_back(1.0);
            }
            if (j < N - 1) {
                rows.push_back(k);
                cols.push_back(i * N + (j + 1));
                vals.push_back(1.0);
            }
        }
    }
    return SparseMatrix::from_triplets(n, n, rows, cols, vals);
}
 
inline SparseMatrix backward_euler_matrix(int N, double coeff) {
    const int n = N * N;
    std::vector<idx> rows, cols;
    std::vector<real> vals;
    rows.reserve(5 * n);
    cols.reserve(5 * n);
    vals.reserve(5 * n);
    for (int i = 0; i < N; ++i) {
        for (int j = 0; j < N; ++j) {
            int k = i * N + j;
            rows.push_back(k);
            cols.push_back(k);
            vals.push_back(1.0 + 4.0 * coeff);
            if (i > 0) {
                rows.push_back(k);
                cols.push_back((i - 1) * N + j);
                vals.push_back(-coeff);
            }
            if (i < N - 1) {
                rows.push_back(k);
                cols.push_back((i + 1) * N + j);
                vals.push_back(-coeff);
            }
            if (j > 0) {
                rows.push_back(k);
                cols.push_back(i * N + (j - 1));
                vals.push_back(-coeff);
            }
            if (j < N - 1) {
                rows.push_back(k);
                cols.push_back(i * N + (j + 1));
                vals.push_back(-coeff);
            }
        }
    }
    return SparseMatrix::from_triplets(n, n, rows, cols, vals);
}
 
inline SparseMatrix backward_euler_matrix(const Grid2D& grid, double coeff) {
    return backward_euler_matrix(grid.N, coeff);
}
 
inline LinearSolver make_cg_solver(const SparseMatrix& A, real tol = 1e-6) {
    return [&A, tol](const Vector& rhs, Vector& x) {
        operators::SparseOp op(A);
        return cg(op, rhs, x, tol);
    };
}
 
} // namespace num::pde
 
namespace num {
using pde::make_cg_solver;
}