collocated.hpp

Go to the documentation of this file.

 
#ifndef PECLET_VORO_FV_COLLOCATED_HPP
#define PECLET_VORO_FV_COLLOCATED_HPP
 
#include <functional>
#include <Kokkos_Core.hpp>
 
#include "peclet/voro/fv/covolume.hpp"
#include "peclet/voro/fv/mesh.hpp"
#include "peclet/voro/fv/operators.hpp"
 
namespace peclet::voro::fv {
 
template <class Real>
struct CollocatedNS {
  using Exec = peclet::core::ExecSpace;
  FaceMesh<Real> m;
  Real nu = 0;
  Real convScale = 1;         // 0 = Stokes (diagnostic)
  bool incremental = true;    // flow's incremental_: false = classical Chorin (φ = the pressure)
  bool skewCorrected = true;  // centroid-consistent constraint pair (see the header); MEASURED
                              // TGV order 2.10 / 2.08 on jittered / CVT meshes vs 1.72 / 1.17 plain
  DV<Real> U, uf, p, force;   // 3N, nFaces, N, 3N (optional)
  DV<Real> Uwall, ub;         // 3 × nBoundary wall velocity, nBoundary wall flux (C3)
  bool wallQuadratic = true;  // second-order wall gradient (wallGradientLS) vs two-point
  DV<Real> wg;                // 3 × nBoundary wall gradient
  // Distributed hooks (rung C5, see fv/distributed.hpp): ghost refresh of a cell field and the
  // global sum / max. Unset = single rank.
  std::function<void(const DV<Real>&, int)> exchange;
  std::function<Real(Real)> sum, max;
  // Implicit diffusion (flow's step structure): explicit convection, backward-Euler viscous solve
  // per component with the two-point Laplacian and the two-point wall term implicit and the
  // quadratic wall-gradient correction lagged (deferred correction keeps the matrix SPD), then the
  // approximate projection. First order in time, no diffusive Δt limit (Stokes: Δt free).
  bool implicitDiffusion = false;
  bool rotational =
      true;  // P += φ − ν div u* (Timmermans; flow's rotationalP_) with the implicit step
  PressureSolver<Real> visc;
  Real viscDt = -1;
  DV<Real> rhs3, uc, wcorr, aT, divStar;
  DV<Real> a, U1, U2, uf1, uf2, div, gphi, gp, phi, g9;  // workspace
 
  void setup(const FaceMesh<Real>& mesh, Real viscosity, bool amg = false) {
    m = mesh;
    nu = viscosity;
    const int N = m.nCells, F = m.nFaces, NC = m.nCombined > 0 ? m.nCombined : N;
    U = DV<Real>("co.U", 3 * NC);
    uf = DV<Real>("co.uf", F);
    p = DV<Real>("co.p", NC);
    a = DV<Real>("co.a", 3 * N);
    U1 = DV<Real>("co.U1", 3 * NC);
    U2 = DV<Real>("co.U2", 3 * NC);
    uf1 = DV<Real>("co.uf1", F);
    uf2 = DV<Real>("co.uf2", F);
    div = DV<Real>("co.div", N);
    gphi = DV<Real>("co.gphi", F);
    gp = DV<Real>("co.gp", 3 * N);
    phi = DV<Real>("co.phi", NC);
    g9 = DV<Real>("co.g9", 9 * NC);
    setWallVelocity(DV<Real>{});
    poisson.setup(m, amg);
  }
  template <class Ex>
  void setExchange(Ex& ex) {
    exchange = [&ex](const DV<Real>& f, int nc) { ex.exchange(f, nc); };
    sum = [&ex](Real v) { return ex.sum(v); };
    max = [&ex](Real v) { return ex.max(v); };
    poisson.setHooks(exchange, sum);
  }
  void setWallVelocity(const DV<Real>& Uw) {
    const int nB = m.nFaces - m.nInterior, nI = m.nInterior;
    Uwall = Uw;
    ub = DV<Real>("fv.ub", nB);
    if (Uw.extent(0) == 0)
      return;
    const FaceMesh<Real> mm = m;
    const DV<Real> bb = ub;
    Kokkos::parallel_for(
        "fv.ub", Kokkos::RangePolicy<Exec>(0, nB), KOKKOS_LAMBDA(const int b) {
          Real un = 0;
          for (int c = 0; c < 3; ++c)
            un += Uw(3 * b + c) * mm.normal(nI + b, c);
          bb(b) = un;
        });
  }
  void applyWallFlux(const DV<Real>& uf) const {
    const int nB = m.nFaces - m.nInterior, nI = m.nInterior;
    const DV<Real> bb = ub;
    Kokkos::parallel_for(
        "fv.wall", Kokkos::RangePolicy<Exec>(0, nB),
        KOKKOS_LAMBDA(const int b) { uf(nI + b) = bb(b); });
  }
  PressureSolver<Real> poisson;
 
  void toFaces(const DV<Real>& Uc, const DV<Real>& out) {
    if (skewCorrected) {
      vectorGreenGauss(m, Uc, g9, Uwall);
      if (exchange)
        exchange(g9, 9);
      projectToFaces(m, Uc, out, g9, ub);
    } else {
      projectToFaces(m, Uc, out, DV<Real>{}, ub);
    }
  }
  void cellGrad(const DV<Real>& q, const DV<Real>& out) {
    faceGradient(m, q, gphi);
    faceInterpTranspose(m, gphi, out, skewCorrected, g9, Uwall.extent(0) > 0, exchange);
  }
  void tendency(const DV<Real>& Uin, const DV<Real>& ufin, const DV<Real>& out) {
    if (wallQuadratic && m.nFaces > m.nInterior) {
      if (wg.extent(0) == 0)
        wg = DV<Real>("co.wg", 3 * (m.nFaces - m.nInterior));
      vectorGreenGauss(m, Uin, g9, Uwall);
      wallGradientLS(m, Uin, g9, Uwall, wg);
      cellTendency(m, ufin, Uin, nu, out, force, convScale, Uwall, wg);
    } else {
      cellTendency(m, ufin, Uin, nu, out, force, convScale, Uwall);
    }
    if (!incremental)
      return;
    const DV<Real> gc = gp;
    cellGrad(p, gc);
    Kokkos::parallel_for(
        "co.gradp", Kokkos::RangePolicy<Exec>(0, 3 * m.nCells),
        KOKKOS_LAMBDA(const int i) { out(i) -= gc(i); });
  }
  void project(const DV<Real>& Ustar, const DV<Real>& ufOut, Real dt,
               const DV<Real>& divOut = DV<Real>{}) {
    const DV<Real> dv = div, g = gphi, gc = gp, ph = phi, pp = p;
    if (exchange)
      exchange(Ustar, 3);  // the predictor's ghosts (interface faces read both cells)
    toFaces(Ustar, ufOut);
    divergence(m, ufOut, dv);
    if (divOut.extent(0) > 0)
      Kokkos::deep_copy(divOut, dv);  // div u* for the rotational pressure update
    const Real s = -Real(1) / dt;
    Kokkos::parallel_for(
        "co.rhsP", Kokkos::RangePolicy<Exec>(0, m.nCells),
        KOKKOS_LAMBDA(const int i) { dv(i) *= s; });
    poisson.solve(dv, ph);
    faceGradient(m, ph, g);
    Kokkos::parallel_for(
        "co.corrF", Kokkos::RangePolicy<Exec>(0, m.nFaces),
        KOKKOS_LAMBDA(const int f) { ufOut(f) -= dt * g(f); });
    cellGrad(ph, gc);
    Kokkos::parallel_for(
        "co.corrC", Kokkos::RangePolicy<Exec>(0, 3 * m.nCells),
        KOKKOS_LAMBDA(const int i) { Ustar(i) -= dt * gc(i); });
    if (incremental)
      Kokkos::parallel_for(
          "co.pinc", Kokkos::RangePolicy<Exec>(0, m.nCells),
          KOKKOS_LAMBDA(const int i) { pp(i) += ph(i); });
    else
      Kokkos::deep_copy(pp, ph);
    if (exchange) {
      exchange(Ustar, 3);  // corrected cells (the next tendency's neighbours)
      exchange(pp, 1);     // the incremental predictor's −∇P
    }
  }
  void initialize(const DV<Real>& Uin) {
    Kokkos::deep_copy(U, Uin);
    Kokkos::deep_copy(p, Real(0));
    project(U, uf, Real(1));
    Kokkos::deep_copy(p, Real(0));
  }
  void step(Real dt) {
    if (implicitDiffusion)
      stepImplicit(dt);
    else
      stepRK3(dt);
  }
  void stepImplicit(Real dt) {
    const int N = m.nCells, NC = m.nCombined > 0 ? m.nCombined : N;
    if (viscDt != dt) {
      visc.setupVelocity(m, nu, dt, poisson.useAmg);
      visc.tol = poisson.tol;
      if (exchange)
        visc.setHooks(exchange, sum);
      viscDt = dt;
      if (rhs3.extent(0) == 0) {
        rhs3 = DV<Real>("co.rhs3", 3 * N);
        uc = DV<Real>("co.uc", NC);
        wcorr = DV<Real>("co.wcorr", 3 * N);
        aT = DV<Real>("co.aT", 3 * N);
        divStar = DV<Real>("co.divStar", N);
      }
    }
    const DV<Real> UU = U, aa = a, R3 = rhs3, WC = wcorr, AT = aT, UC = uc, gc = gp;
    const FaceMesh<Real> mm = m;
    const Real nuL = nu;
    const bool hasWall = m.nFaces > m.nInterior;
    // explicit part: convection + force (nu = 0) − G Pⁿ
    cellTendency(m, uf, UU, Real(0), aa, force, convScale, Uwall);
    if (incremental) {
      cellGrad(p, gc);
      Kokkos::parallel_for(
          "co.igradp", Kokkos::RangePolicy<Exec>(0, 3 * N),
          KOKKOS_LAMBDA(const int i) { aa(i) -= gc(i); });
    }
    // lagged wall correction: quadratic − two-point viscous wall term at Uⁿ (interior cancels)
    if (hasWall && wallQuadratic) {
      if (wg.extent(0) == 0)
        wg = DV<Real>("co.wg", 3 * (m.nFaces - m.nInterior));
      vectorGreenGauss(m, UU, g9, Uwall);
      wallGradientLS(m, UU, g9, Uwall, wg);
      const DV<Real> zero("co.zero", m.nFaces);
      cellTendency(m, zero, UU, nuL, WC, DV<Real>{}, Real(0), Uwall, wg);
      cellTendency(m, zero, UU, nuL, AT, DV<Real>{}, Real(0), Uwall);
      Kokkos::parallel_for(
          "co.wcorr", Kokkos::RangePolicy<Exec>(0, 3 * N),
          KOKKOS_LAMBDA(const int i) { aa(i) += WC(i) - AT(i); });
    }
    // right-hand side: V(U/Δt + a) + ν Σ_wall (A/h_A) U_wall
    const DV<Real> Uw = Uwall;
    const bool hasW = Uw.extent(0) > 0;
    const int nI = m.nInterior;
    Kokkos::parallel_for(
        "co.rhs", Kokkos::RangePolicy<Exec>(0, N), KOKKOS_LAMBDA(const int i) {
          const Real V = mm.cellVolume(i);
          Real wall[3] = {0, 0, 0};
          if (hasW)
            for (int t = mm.cellFacesBegin(i); t < mm.cellFacesEnd(i); ++t) {
              const int f = mm.cellFace(t);
              if (f >= nI) {
                const Real w = nuL * mm.faceArea(f) / mm.faceHa(f);
                for (int c = 0; c < 3; ++c)
                  wall[c] += w * Uw(3 * (f - nI) + c);
              }
            }
          for (int c = 0; c < 3; ++c)
            R3(3 * i + c) = V * (UU(3 * i + c) / dt + aa(3 * i + c)) + wall[c];
        });
    // three scalar solves; the previous velocity component is the initial guess's replacement
    // (PressureSolver starts from zero — fine for CG)
    for (int c = 0; c < 3; ++c) {
      const DV<Real> rc("co.rc", N);
      Kokkos::parallel_for(
          "co.gather", Kokkos::RangePolicy<Exec>(0, N),
          KOKKOS_LAMBDA(const int i) { rc(i) = R3(3 * i + c); });
      visc.solve(rc, UC);
      Kokkos::parallel_for(
          "co.scatter", Kokkos::RangePolicy<Exec>(0, N),
          KOKKOS_LAMBDA(const int i) { UU(3 * i + c) = UC(i); });
    }
    // approximate projection (exchanges the predictor's ghosts itself)
    const DV<Real> dS = divStar;
    project(UU, uf, dt, rotational ? dS : DV<Real>{});
    if (rotational && incremental) {
      const DV<Real> pp = p;
      Kokkos::parallel_for(
          "co.rot", Kokkos::RangePolicy<Exec>(0, N),
          KOKKOS_LAMBDA(const int i) { pp(i) -= nuL * dS(i); });
      if (exchange)
        exchange(pp, 1);
    }
  }
  void stepRK3(Real dt) {
    const int N3 = 3 * m.nCells;
    const DV<Real> UU = U, aa = a, V1 = U1, V2 = U2;
    tendency(UU, uf, aa);
    Kokkos::parallel_for(
        "co.s1", Kokkos::RangePolicy<Exec>(0, N3),
        KOKKOS_LAMBDA(const int i) { V1(i) = UU(i) + dt * aa(i); });
    project(V1, uf1, dt);
    tendency(V1, uf1, aa);
    Kokkos::parallel_for(
        "co.s2", Kokkos::RangePolicy<Exec>(0, N3), KOKKOS_LAMBDA(const int i) {
          V2(i) = Real(0.75) * UU(i) + Real(0.25) * (V1(i) + dt * aa(i));
        });
    project(V2, uf2, Real(0.25) * dt);
    tendency(V2, uf2, aa);
    Kokkos::parallel_for(
        "co.s3", Kokkos::RangePolicy<Exec>(0, N3),
        KOKKOS_LAMBDA(const int i) { UU(i) = (UU(i) + Real(2) * (V2(i) + dt * aa(i))) / Real(3); });
    project(UU, uf, Real(2) / Real(3) * dt);
  }
  Real kineticEnergy() const {
    const FaceMesh<Real> mm = m;
    const DV<Real> UU = U;
    Real e = 0;
    Kokkos::parallel_reduce(
        "co.ke", Kokkos::RangePolicy<Exec>(0, m.nCells),
        KOKKOS_LAMBDA(const int i, Real& acc) {
          acc += mm.cellVolume(i) * (UU(3 * i) * UU(3 * i) + UU(3 * i + 1) * UU(3 * i + 1) +
                                     UU(3 * i + 2) * UU(3 * i + 2));
        },
        e);
    if (sum)
      e = sum(e);
    return Real(0.5) * e;
  }
  Real maxFaceDivergence() { return maxDivOf(uf); }
  Real maxCellDivergence() {
    toFaces(U, uf1);
    return maxDivOf(uf1);
  }
  Real maxDivOf(const DV<Real>& flux) {
    const DV<Real> dv = div;
    divergence(m, flux, dv);
    Real mx = 0;
    Kokkos::parallel_reduce(
        "co.maxdiv", Kokkos::RangePolicy<Exec>(0, m.nCells),
        KOKKOS_LAMBDA(const int i, Real& acc) {
          const Real v = Kokkos::fabs(dv(i));
          if (v > acc)
            acc = v;
        },
        Kokkos::Max<Real>(mx));
    return max ? max(mx) : mx;
  }
};
 
}  // namespace peclet::voro::fv
 
#endif  // PECLET_VORO_FV_COLLOCATED_HPP