adi.hpp

Go to the documentation of this file.

/// @file pde/adi.hpp
/// @brief Crank-Nicolson ADI solver for 2D parabolic equations via fiber sweeps.
///
/// CrankNicolsonADI encapsulates the prefactored complex tridiagonals for
/// Strang-split time stepping of equations of the form:
///
///   i ∂ψ/∂t = -(1/2) ∆ψ + V(x,y) ψ   (Schrödinger)
///   ∂u/∂t   = κ ∆u                      (complex-coefficient diffusion)
///
/// Each CN sub-step τ along one axis solves a 1D system per fiber:
///
///   (I - ia·Lap₁D) ψⁿ⁺¹ = (I + ia·Lap₁D) ψⁿ,   ia = i·τ/(4h²)
///
/// Two tridiagonals are prefactored once — for τ = dt/2 and τ = dt —
/// covering the full Strang splitting: sweep_x(dt/2) → sweep_y(dt) → sweep_x(dt/2).
///
/// Typical usage (TDSE Strang splitting):
/// @code
///   num::CrankNicolsonADI adi(N, dt, h);
///
///   // one full time step:
///   adi.sweep(psi, true,  dt * 0.5);  // x, half-step
///   adi.sweep(psi, false, dt);         // y, full-step
///   adi.sweep(psi, true,  dt * 0.5);  // x, half-step
/// @endcode
#pragma once
 
#include "pde/stencil.hpp"
#include "linalg/factorization/tridiag_complex.hpp"
#include "core/vector.hpp"
#include <vector>
#include <complex>
 
namespace num {
 
struct CrankNicolsonADI {
    int    N  = 0;
    double dt = 0.0;
    double h  = 0.0;
 
    CrankNicolsonADI() = default;
 
    /// Pre-factor both tridiagonals.
    /// @param N_   Interior grid points per axis
    /// @param dt_  Full timestep
    /// @param h_   Grid spacing
    CrankNicolsonADI(int N_, double dt_, double h_)
        : N(N_), dt(dt_), h(h_)
    {
        using cplx = std::complex<double>;
        // alpha = tau/(4h²).  LHS tridiagonal: a=c=-i*alpha, b=1+2i*alpha.
        auto factor = [&](double tau) {
            double alpha = tau / (4.0 * h * h);
            cplx a(0.0, -alpha);
            cplx b(1.0,  2.0 * alpha);
            ComplexTriDiag td;
            td.factor(N, a, b, a);
            return td;
        };
        td_half_ = factor(dt * 0.5);
        td_full_ = factor(dt);
    }
 
    /// Apply one CN sweep along x (col fibers, x_axis=true) or y (row fibers).
    /// Selects the half-step tridiagonal for tau < 0.75*dt, full-step otherwise.
    void sweep(CVector& psi, bool x_axis, double tau) const {
        using cplx = std::complex<double>;
        const ComplexTriDiag& td = (tau < dt * 0.75) ? td_half_ : td_full_;
        const cplx ia(0.0,  tau / (4.0 * h * h));
        const cplx diag(1.0, -2.0 * tau / (4.0 * h * h));
 
        auto apply = [&](std::vector<cplx>& fiber) {
            std::vector<cplx> rhs(N);
            for (int i = 0; i < N; ++i) {
                cplx prev = (i > 0)   ? fiber[i-1] : cplx{};
                cplx next = (i < N-1) ? fiber[i+1] : cplx{};
                rhs[i] = ia * prev + diag * fiber[i] + ia * next;
            }
            td.solve(rhs);
            fiber = std::move(rhs);
        };
 
        if (x_axis) col_fiber_sweep(psi, N, apply);
        else        row_fiber_sweep(psi, N, apply);
    }
 
private:
    ComplexTriDiag td_half_;
    ComplexTriDiag td_full_;
};
 
} // namespace num