CFDManager/docs/theory/solver_2x_sdf/solver.py

# Оркестратор модели 2×+SDF: держит persistent-буферы, собирает один шаг из компонентов и гоняет цикл.
# Один шаг = замкнутая step() над фикс. буферами без host→device (готова к захвату в CUDA-граф).
import math
import backend as B
import lattice as L
import geometry as G
import equilibrium as E
import collision as Coll
import streaming as Strm
import boundary as Bnd
import forces as Fz
import forces_stress as FzS
import amr as Amr

cp = B.cp; xp = B.xp; DTYPE = B.DTYPE; to_cpu = B.to_cpu
CYa = L.CYa

def _capture(step_fn):
    """Захват одного шага в CUDA-граф. При сбое гарантированно выводит поток из режима записи."""
    cp.cuda.Device().synchronize()
    s = cp.cuda.Stream(non_blocking=True)
    with s:
        s.begin_capture()
        try:
            step_fn()
        except Exception:
            try: s.end_capture()
            except Exception: pass
            raise
        return s.end_capture()

class Solver:
    def __init__(self, cfg):
        self.cfg = cfg
        self.collide = Coll.collide        # KBC-N1 (единственный оператор столкновения)
        self.force = Fz.force_gmem         # галилей-инвариантный обмен импульсом
        self.geo = G.build(cfg)
        self.P1 = self.geo.P1
        self.beta1 = cfg.beta1
        if cfg.sponge_len:
            # губка: τ(x) = τ0 + (mult−1)·ν·s(x)/cs², s — smoothstep 0→1 в последних sponge_len
            # столбцах. β становится ПОЛЕМ (Nx,) → broadcast в kbc_collide; патч L1 губки не видит
            # (гарантировано assert'ом в config), там остаётся скалярный beta1.
            xs = cp.arange(cfg.Nx, dtype=DTYPE)
            s = cp.clip((xs - (cfg.Nx - 1 - cfg.sponge_len)) / float(cfg.sponge_len), 0.0, 1.0)
            s = s*s*(3.0 - 2.0*s)
            tau_x = cfg.tau0 + (cfg.sponge_nu_mult - 1.0) * cfg.nu / (1.0/3.0) * s
            self.beta0 = (1.0 / (2.0 * tau_x)).astype(DTYPE)[None, :]   # (1,Nx) → broadcast по y
        else:
            self.beta0 = cfg.beta0
        self.R01 = cfg.R01; self.Rf01 = cfg.Rf01
        self.refine = cfg.refine
        # ЗОНД скорости в следе: физическая точка (px,py) в координатах L0.
        # Берём её с ТОНКОЙ сетки L1 (меньше численного размытия вихрей → чище St / rms u_y),
        # если точка внутри патча; иначе fallback на L0. Индексы фиксированы → graph-safe.
        self.probe_l1 = (cfg.ax1 <= cfg.px <= cfg.bx1) and (cfg.ay1 <= cfg.py <= cfg.by1)
        if self.probe_l1:
            self.pfx = (cfg.px - cfg.ax1) * cfg.refine
            self.pfy = (cfg.py - cfg.ay1) * cfg.refine
        else:
            self.pfx = cfg.px; self.pfy = cfg.py
        # ГУ (SDF/Bouzidi) на L0 (цилиндр+стенки) и L1 (цилиндр)
        self.BC0 = Bnd.build_bc(self.geo.solid0, self.geo.phi0, self.geo.cyl0)
        self.BC1 = Bnd.build_bc(self.geo.solid1, self.geo.phi1, self.geo.solid1)
        self.fl1 = self.geo.fl1
        # persistent-состояние
        Ny, Nx = cfg.Ny, cfg.Nx; Nfy, Nfx = self.P1["Nfy"], self.P1["Nfx"]
        self.F0 = E.feq(xp.ones((Ny, Nx), DTYPE), xp.zeros((2, Ny, Nx), DTYPE)).copy()
        self.F1 = E.feq(xp.ones((Nfy, Nfx), DTYPE), xp.zeros((2, Nfy, Nfx), DTYPE)).copy()
        self.onecol = xp.ones((Ny, 1), DTYPE)
        self.uin = xp.zeros((2, Ny, 1), DTYPE)               # вход; обновляется СНАРУЖИ захвата
        self.Fx_s = xp.zeros((), DTYPE); self.Fy_s = xp.zeros((), DTYPE); self.uy_s = xp.zeros((), DTYPE)
        self.Tz_s = xp.zeros((), DTYPE)
        # координаты узлов L1 относительно центра тела — плечи торка (прекомпьют вне захвата)
        ccx1 = (cfg.cx - cfg.ax1) * cfg.refine; ccy1 = (cfg.cy - cfg.ay1) * cfg.refine
        Yf, Xf = cp.meshgrid(cp.arange(Nfy), cp.arange(Nfx), indexing="ij")
        self.rxy1 = ((Xf - ccx1).astype(DTYPE), (Yf - ccy1).astype(DTYPE))
        # контур кросс-чека силы σ·n (forces_stress; считается редко, вне CUDA-графа)
        self.sprobe = FzS.build_probe(cfg) if cfg.stress_every else None
        # разгон входа на device: U*smoothstep(t/ramp)
        ts = cp.minimum(cp.arange(cfg.steps + 1) / cfg.ramp, 1.0)
        self.ramp_u = (cfg.U * (ts*ts*(3 - 2*ts))).astype(DTYPE)
        # стартовое возмущение: поперечный sin-импульс u_y входа в окне [ramp, ramp+pert_dur] (триггер схода)
        uyp = cp.zeros(cfg.steps + 1, DTYPE)
        if cfg.pert_dur > 0 and cfg.pert_amp != 0.0:
            k = cp.arange(cfg.steps + 1)
            inw = (k >= cfg.ramp) & (k < cfg.ramp + cfg.pert_dur)
            ph = (k - cfg.ramp).astype(DTYPE) / float(cfg.pert_dur)
            uyp = cp.where(inw, (cfg.pert_amp * cfg.U) * cp.sin(math.pi * ph), 0.0).astype(DTYPE)
        self.ramp_uy = uyp
        # ⟨ρ⟩ — диагностика дрейфа массы. Считать ТОЛЬКО по жидкости: узлы внутри цилиндра —
        # фиктивные (их популяции не имеют физического смысла, ГУ их не трогает), а их плотность
        # уезжает независимо. При D=12 на 120×60 это 1.6% узлов и смещение ⟨ρ⟩ на −5.4e-4 —
        # ровно тот разряд, в котором ⟨ρ⟩ и цитируется в отчётах (1.0000 / 1.001).
        self.fluid0 = (~self.geo.solid0).astype(DTYPE)
        self._ncells = float(self.fluid0.sum())

    def step(self):
        F0 = self.F0; F1 = self.F1; P1 = self.P1
        feq = E.feq; macros = E.macros; stream = Strm.stream
        # ----- уровень 0 -----
        rho, u0 = macros(F0)
        pre = F0.copy()
        post0 = self.collide(F0, feq(rho, u0), self.beta0)
        f0 = stream(post0); f0 = Bnd.apply_bc(f0, post0, self.BC0)
        # порядок важен: стенки снимают y-заворот roll, затем Zou-He закрывает x-заворот НА ВЕСЬ
        # столбец, включая угловые узлы (иначе вход и выход связаны напрямую — см. boundary.py)
        f0 = Bnd.free_slip_walls(f0)                            # верх/низ — скользящие стенки (specular)
        f0 = Bnd.channel_bc(f0, self.uin, outlet_uy=self.cfg.outlet_uy)  # Zou-He вход/выход
        # ----- уровень 1: refine подшагов по времени с временно́й интерполяцией ghost -----
        g1o = Amr.ghost(pre, P1, self.R01); g1n = Amr.ghost(f0, P1, self.R01)
        f1 = F1
        # сила/момент: меряем на КАЖДОМ подшаге L1 и усредняем (анти-алиасинг; раньше — мгновенная
        # на последнем подшаге). Средний Cd не меняется, ряды становятся глаже.
        fxa = B.ZERO; fya = B.ZERO; tza = B.ZERO
        for s1 in range(self.refine):
            r1, u1 = macros(f1); post1 = self.collide(f1, feq(r1, u1), self.beta1)
            f1 = stream(post1); f1 = Bnd.apply_bc(f1, post1, self.BC1)
            fx, fy, tz = self.force(post1, f1, self.BC1, self.rxy1)
            fxa = fxa + fx; fya = fya + fy; tza = tza + tz
            w = (s1 + 1) / self.refine                         # временна́я интерполяция ghost: o→n
            f1 = Amr.fill(f1, (1 - w)*g1o + w*g1n)
        self.Fx_s[...] = fxa / self.refine; self.Fy_s[...] = fya / self.refine
        self.Tz_s[...] = tza / self.refine
        f0 = Amr.restrict(f1, f0, P1, self.Rf01, self.fl1)
        F1[...] = f1
        # зонд скорости в следе (с тонкой сетки L1, если точка внутри патча) + запись состояния
        psrc = f1 if self.probe_l1 else f0
        col = psrc[:, self.pfy, self.pfx]; self.uy_s[...] = (CYa*col).sum() / col.sum()
        F0[...] = f0

    def run(self, progress=None):
        cfg = self.cfg
        Fx = xp.zeros(cfg.steps + 1, DTYPE); Fy = xp.zeros(cfg.steps + 1, DTYPE); uy = xp.zeros(cfg.steps + 1, DTYPE)
        Tz = xp.zeros(cfg.steps + 1, DTYPE)           # момент (торк) вокруг центра тела
        rho = xp.zeros(cfg.steps + 1, DTYPE)          # ⟨ρ⟩(t) — диагностика дрейфа массы
        # кросс-чек σ·n: редкая выборка (каждые stress_every шагов), вне CUDA-графа
        nst = (cfg.steps // cfg.stress_every + 1) if cfg.stress_every else 0
        Fx_st = xp.zeros(nst, DTYPE); Fy_st = xp.zeros(nst, DTYPE); Tz_st = xp.zeros(nst, DTYPE)
        kst = 0
        frames = []
        graph = None; graph_failed = False
        def snap(): return (self.F0.copy(), self.F1.copy())
        def restore(s): self.F0[...] = s[0]; self.F1[...] = s[1]
        def maxdiff(s): return max(float(xp.abs(s[0] - self.F0).max()), float(xp.abs(s[1] - self.F1).max()))
        for t in range(cfg.steps + 1):
            self.uin[0, :, 0] = self.ramp_u[t]
            self.uin[1, :, 0] = self.ramp_uy[t]       # поперечный импульс-триггер (0 вне окна возмущения)
            did = False
            if t == 1 and cfg.use_cuda_graph and not graph_failed:
                # захват + РАНТАЙМ-САМОПРОВЕРКА (граф vs eager на одном шаге); при расхождении — откат
                try:
                    graph = _capture(self.step)
                    before = snap(); graph.launch(); gres = snap()
                    restore(before); self.step()
                    d = maxdiff(gres)
                    if d > 1e-3:
                        print(f"\n[graph] self-check разошёлся (Δ={d:.2e}) — отключаю графы, работаю eager")
                        graph = None; graph_failed = True
                    else:
                        print(f"\n[graph] self-check OK (Δ={d:.2e}) — CUDA graphs активны")
                    did = True
                except Exception as e:
                    print(f"\n[graph] недоступны ({type(e).__name__}: {e}) — работаю eager")
                    import traceback as _tb; _tb.print_exc()
                    graph = None; graph_failed = True
                    try: cp.cuda.Device().synchronize()
                    except Exception: pass
            if not did:
                if graph is not None and t >= 1:
                    graph.launch()
                else:
                    self.step()
            Fx[t] = self.Fx_s; Fy[t] = self.Fy_s; Tz[t] = self.Tz_s; uy[t] = self.uy_s
            rho[t] = (self.F0.sum(0) * self.fluid0).sum() / self._ncells   # ⟨ρ⟩ по жидкости (дрейф массы)
            if cfg.stress_every and t % cfg.stress_every == 0:
                # eager-вычисление на текущем F1 после launch/step — упорядочено на стриме,
                # сам захваченный граф не трогаем (паттерн как у rho[t] выше)
                sfx, sfy, stz = FzS.force_stress(self.F1, self.sprobe, cfg.tau1)
                Fx_st[kst] = sfx; Fy_st[kst] = sfy; Tz_st[kst] = stz; kst += 1
            if progress is not None:
                progress(t, cfg.steps)
            if t % 500 == 0 and not bool(xp.isfinite(self.F0).all()):
                print("BLEW UP", t)
                Fx = Fx[:t]; Fy = Fy[:t]; Tz = Tz[:t]; uy = uy[:t]; rho = rho[:t]
                Fx_st = Fx_st[:kst]; Fy_st = Fy_st[:kst]; Tz_st = Tz_st[:kst]
                break
            if cfg.frame_every and t % cfg.frame_every == 0:
                frames.append((t, to_cpu(E.macros(self.F0)[1]), to_cpu(E.macros(self.F1)[1])))
        return dict(Fx=Fx, Fy=Fy, Tz=Tz, uy=uy, rho=rho, frames=frames, DL=cfg.DL,
                    Fx_st=Fx_st, Fy_st=Fy_st, Tz_st=Tz_st, stress_every=cfg.stress_every)