Function bodies 112 total

                KOKKOS_LAMBDA(int i, int j, double& lsum) {
                    const double r = v_star(i, j)
                        - alpha * physics::laplacian(v_star, i, j, inv_dx2, inv_dy2)
                        - rhs(i, j);
                    lsum += r * r;
                }, res2);

def load_ke(directory):
    """Load kinetic energy time history from output_ke.csv.

    Parameters
    ----------
    directory : str
        Path to output directory containing ``output_ke.csv``.

    Returns
    -------
    dict
        Keys ``"step"``, ``"time"``, ``"kinetic_energy"`` as 1D arrays.
    """
    path = os.path.join(directory, "output_ke.csv")
    raw = np.loadtxt(path, delimiter=",", skiprows=1)
    return {
        "step": raw[:, 0].astype(int),
        "time": raw[:, 1],
        "kinetic_energy": raw[:, 2],
    }

def load_divergence(directory):
    """Load divergence L2-norm time history from output_div.csv.

    Parameters
    ----------
    directory : str
        Path to output directory containing ``output_div.csv``.

    Returns
    -------
    dict
        Keys ``"step"``, ``"time"``, ``"l2_divergence"`` as 1D arrays.
    """
    path = os.path.join(directory, "output_div.csv")
    raw = np.loadtxt(path, delimiter=",", skiprows=1)
    return {
        "step": raw[:, 0].astype(int),
        "time": raw[:, 1],
        "l2_divergence": raw[:, 2],
    }

def plot_kinetic_energy(ke_data, *, ax=None):
    """Plot kinetic energy vs time.

    Parameters
    ----------
    ke_data : dict
        Output of :func:`load_ke`.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on. Created if None.

    Returns
    -------
    fig, ax
    """
    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = ax.figure

    ax.plot(ke_data["time"], ke_data["kinetic_energy"], color="k")
    ax.set_xlabel(r"$t$")
    ax.set_ylabel(r"$E_k$")
    ax.set_title("Kinetic Energy")
    ax.ticklabel_format(axis="y", style="sci", scilimits=(-2, 3))

    return fig, ax

def plot_divergence(div_data, *, ax=None, log_scale=True):
    """Plot L2 divergence norm vs time.

    Parameters
    ----------
    div_data : dict
        Output of :func:`load_divergence`.
    ax : matplotlib.axes.Axes, optional
        Axes to plot on. Created if None.
    log_scale : bool
        Use log scale on y-axis (default True, since divergence is
        typically near machine epsilon).

    Returns
    -------
    fig, ax
    """
    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = ax.figure

    ax.plot(div_data["time"], div_data["l2_divergence"], color="k")
    ax.set_xlabel(r"$t$")
    ax.set_ylabel(r"$\| \nabla \cdot \mathbf{u} \|_2$")
    ax.set_title("Divergence")
    if log_scale:
        ax.set_yscale("log")

    return fig, ax

def plot_diagnostics(ke_data, div_data):
    """Two-panel figure with kinetic energy (top) and divergence (bottom).

    Parameters
    ----------
    ke_data : dict
        Output of :func:`load_ke`.
    div_data : dict
        Output of :func:`load_divergence`.

    Returns
    -------
    fig, (ax_ke, ax_div)
    """
    fig, (ax_ke, ax_div) = plt.subplots(2, 1, figsize=(3.5, 4.5), sharex=True)

    plot_kinetic_energy(ke_data, ax=ax_ke)
    ax_ke.set_xlabel("")  # shared x-axis, label only on bottom

    plot_divergence(div_data, ax=ax_div)

    fig.tight_layout()
    return fig, (ax_ke, ax_div)

def _contour_field(x, y, z, *, cmap, title, label, levels=N_LEVELS,
                   ax=None, vmin=None, vmax=None):
    """Filled contour plot with colorbar.

    Returns (fig, ax, contour_set).
    """
    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = ax.figure

    if isinstance(levels, int):
        lo = vmin if vmin is not None else z.min()
        hi = vmax if vmax is not None else z.max()
        if lo == hi:
            hi = lo + 1.0
        levels = np.linspace(lo, hi, levels)

    cs = ax.contourf(x, y, z, levels=levels, cmap=cmap)
    cb = fig.colorbar(cs, ax=ax, label=label, fraction=0.046, pad=0.04)
    cb.ax.tick_params(labelsize=8, direction="in")
    ax.set_xlabel(r"$x$")
    ax.set_ylabel(r"$y$")
    ax.set_title(title)
    ax.set_aspect("equal")

    return fig, ax, cs

def plot_pressure(data, *, ax=None, levels=N_LEVELS):
    """Filled contour of the pressure field.

    Uses a diverging colormap centered on zero.
    """
    p = data["p"]
    vabs = np.abs(p["val"]).max()
    if vabs == 0:
        vabs = 1.0

    return _contour_field(
        p["x"], p["y"], p["val"],
        cmap=CMAP_PRESSURE,
        title="Pressure",
        label=r"$p$",
        levels=levels,
        ax=ax,
        vmin=-vabs,
        vmax=vabs,
    )

def plot_velocity_magnitude(data, *, ax=None, levels=N_LEVELS):
    """Filled contour of velocity magnitude.

    Interpolates staggered u, v to cell centers before computing
    |V| = sqrt(u_cc^2 + v_cc^2).
    """
    u = data["u"]["val"]  # shape (ny, nx+1)
    v = data["v"]["val"]  # shape (ny+1, nx)

    u_cc = 0.5 * (u[:, :-1] + u[:, 1:])
    v_cc = 0.5 * (v[:-1, :] + v[1:, :])
    vmag = np.sqrt(u_cc**2 + v_cc**2)

    return _contour_field(
        data["p"]["x"], data["p"]["y"], vmag,
        cmap=CMAP_VELOCITY,
        title="Velocity Magnitude",
        label=r"$|\mathbf{V}|$",
        levels=levels,
        ax=ax,
        vmin=0,
    )

def _load_field(path):
    """Read a single CSV file and reshape into 2D arrays.

    The C++ writer loops i (outer) then j (inner), so the flat data
    reshapes to (ni, nj) in C-order.  We transpose to (nj, ni) so that
    the first axis is y and the second is x, matching matplotlib's
    contourf convention.
    """
    raw = np.loadtxt(path, delimiter=",", skiprows=1)
    i_col = raw[:, 0].astype(int)
    j_col = raw[:, 1].astype(int)

    ni = i_col.max() + 1
    nj = j_col.max() + 1

    x = raw[:, 2].reshape(ni, nj).T
    y = raw[:, 3].reshape(ni, nj).T
    val = raw[:, 4].reshape(ni, nj).T

    return {"x": x, "y": y, "val": val, "ni": ni, "nj": nj}

def _find_field_file(directory, suffix):
    """Find a CSV file ending with the given suffix in directory."""
    matches = glob.glob(os.path.join(directory, f"*_{suffix}.csv"))
    if len(matches) == 0:
        raise FileNotFoundError(
            f"No *_{suffix}.csv file found in {directory}")
    if len(matches) > 1:
        raise ValueError(
            f"Multiple *_{suffix}.csv files found in {directory}: {matches}")
    return matches[0]

def load_snapshot(directory):
    """Load pressure, u-velocity, and v-velocity fields from a directory.

    Parameters
    ----------
    directory : str
        Path to directory containing ``*_p.csv``, ``*_u.csv``,
        ``*_v.csv`` files.

    Returns
    -------
    dict
        Keys ``"p"``, ``"u"``, ``"v"`` each mapping to a dict with
        ``"x"``, ``"y"``, ``"val"`` (2D arrays) and ``"ni"``, ``"nj"``.
        Also includes ``"nx"`` and ``"ny"`` (pressure cell counts).
    """
    p = _load_field(_find_field_file(directory, "p"))
    u = _load_field(_find_field_file(directory, "u"))
    v = _load_field(_find_field_file(directory, "v"))

    return {"p": p, "u": u, "v": v, "nx": p["ni"], "ny": p["nj"]}

def main():
    parser = argparse.ArgumentParser(
        description="Plot NS solver output fields")
    parser.add_argument("run_dir",
                        help="Run directory (e.g. data/run_64_64_100_0.01)")
    parser.add_argument("--show", action="store_true",
                        help="Display plots interactively")
    parser.add_argument("--format", default="png",
                        choices=["png", "pdf", "svg"],
                        help="Output image format")
    args = parser.parse_args()

    output_dir = os.path.join(args.run_dir, "output")
    figures_dir = os.path.join(args.run_dir, "figures")
    os.makedirs(figures_dir, exist_ok=True)

    data = load_snapshot(output_dir)

    fig, _, _ = plot_pressure(data)
    fig.savefig(os.path.join(figures_dir, f"pressure.{args.format}"))

    fig, _, _ = plot_velocity_magnitude(data)
    fig.savefig(os.path.join(figures_dir, f"velocity_magnitude.{args.format}"))

    ke_data = load_ke(output_dir)
    div_data = load_

void CSVOutput::write(const SimState& s) const {
    Kokkos::fence();

    auto u_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, s.u);
    auto v_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, s.v);
    auto p_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, s.p);

    const auto& g = s.grid;

    // --- pressure (cell-centred) ---
    {
        std::ofstream f(this->filename + "_p.csv");
        if (!f.is_open())
            throw std::runtime_error("Cannot open output file: " + this->filename + "_p.csv");
        f << std::scientific << std::setprecision(10);
        f << "i,j,x,y,p\n";
        for (int i = g.p_i_begin(); i < g.p_i_end(); ++i)
            for (int j = g.p_j_begin(); j < g.p_j_end(); ++j) {
                double x = (i - g.ng + 0.5) * g.dx;
                double y = (j - g.ng + 0.5) * g.dy;
                f << (i - g.ng) << "," << (j - g.ng) << ","
                  << x << "," << y << "," << p_h(i, j) << "\n";
            }

void CSVOutput::write_kinetic_energy(const SimState& s) const {
    std::ofstream f(this->filename + "_ke.csv");
    if (!f.is_open())
        throw std::runtime_error("Cannot open output file: " + this->filename + "_ke.csv");
    f << std::scientific << std::setprecision(10);
    f << "step,time,kinetic_energy\n";
    for (int n = 0; n < s.step; ++n) {
        f << n << "," << s.time_history(n) << "," << s.ke_history(n) << "\n";
    }

    std::cout << "Kinetic energy history written to " << this->filename << "_ke.csv\n";
}

void CSVOutput::write_l2_divergence(const SimState& s) const {
    std::ofstream f(this->filename + "_div.csv");
    if (!f.is_open())
        throw std::runtime_error("Cannot open output file: " + this->filename + "_div.csv");
    f << std::scientific << std::setprecision(10);
    f << "step,time,l2_divergence\n";
    for (int n = 0; n < s.step; ++n) {
        f << n << "," << s.time_history(n) << "," << s.div_history(n) << "\n";
    }

    std::cout << "L2 divergence history written to " << this->filename << "_div.csv\n";
}

void compute_u_rhs(const SimState& s, double re,
                   Kokkos::View<double**> rhs_u) {
    // convective + diffusive terms for u-momentum
    // rhs_u(i,j) = -d(uu)/dx - d(uv)/dy + (1/Re)(d2u/dx2 + d2u/dy2)
    const double inv_dx = 1.0 / s.grid.dx;
    const double inv_dy = 1.0 / s.grid.dy;
    const double inv_dx2 = 1.0 / (s.grid.dx*s.grid.dx);
    const double inv_dy2 = 1.0 / (s.grid.dy*s.grid.dy);
    const double inv_re  = 1.0 / re;
    
    SimState::View2D u = s.u;
    SimState::View2D v = s.v;

    Kokkos::parallel_for("compute_u_rhs",
    Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
        {s.grid.u_i_begin(), s.grid.u_j_begin()},
        {s.grid.u_i_end(), s.grid.u_j_end()}),
    KOKKOS_LAMBDA(int i, int j) {
        const double u_east = avg_x(u, i, j);
        const double u_west = avg_x(u, i-1, j);
        const double duudx = (u_east*u_east - u_west*u_west) * inv_dx;

        const double u_north = avg_y(u, i, j);
        const double u_south = avg_y(u, i, j-1)

    KOKKOS_LAMBDA(int i, int j) {
        const double u_east = avg_x(u, i, j);
        const double u_west = avg_x(u, i-1, j);
        const double duudx = (u_east*u_east - u_west*u_west) * inv_dx;

        const double u_north = avg_y(u, i, j);
        const double u_south = avg_y(u, i, j-1);
        const double v_north = avg_x(v, i-1, j+1);
        const double v_south = avg_x(v, i-1, j);
        const double duvdy = (u_north*v_north - u_south * v_south) *inv_dy;

        const double viscous_term = inv_re * laplacian(u, i, j, inv_dx2, inv_dy2);

        rhs_u(i,j) = - duudx - duvdy + viscous_term;
    });

void compute_v_rhs(const SimState& s, double re,
                   Kokkos::View<double**> rhs_v) {
    // convective + diffusive terms for v-momentum
    // rhs_v(i,j) = -d(uv)/dx - d(vv)/dy + (1/Re)(d2v/dx2 + d2v/dy2)
    const double inv_dx  = 1.0 / s.grid.dx;
    const double inv_dy  = 1.0 / s.grid.dy;
    const double inv_dx2 = 1.0 / (s.grid.dx*s.grid.dx);
    const double inv_dy2 = 1.0 / (s.grid.dy*s.grid.dy);
    const double inv_re  = 1.0 / re;

    SimState::View2D u = s.u;
    SimState::View2D v = s.v;

    Kokkos::parallel_for("compute_v_rhs",
    Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
        {s.grid.v_i_begin(), s.grid.v_j_begin()},
        {s.grid.v_i_end(),   s.grid.v_j_end()}),
    KOKKOS_LAMBDA(int i, int j) {
        const double v_north = avg_y(v, i, j);
        const double v_south = avg_y(v, i, j-1);
        const double dvvdy = (v_north*v_north - v_south*v_south) * inv_dy;

        const double v_east = avg_x(v, i,   j);
        const double v_west = avg_x(v, i-

    KOKKOS_LAMBDA(int i, int j) {
        const double v_north = avg_y(v, i, j);
        const double v_south = avg_y(v, i, j-1);
        const double dvvdy = (v_north*v_north - v_south*v_south) * inv_dy;

        const double v_east = avg_x(v, i,   j);
        const double v_west = avg_x(v, i-1, j);
        const double u_east = avg_y(u, i+1, j-1);
        const double u_west = avg_y(u, i,   j-1);
        const double duvdx = (u_east*v_east - u_west*v_west) * inv_dx;

        const double viscous_term = inv_re * laplacian(v, i, j, inv_dx2, inv_dy2);

        rhs_v(i,j) = -duvdx - dvvdy + viscous_term;
    });

void compute_u_conv_rhs(const SimState& s,
                        Kokkos::View<double**> rhs_u) {
    const double inv_dx = 1.0 / s.grid.dx;
    const double inv_dy = 1.0 / s.grid.dy;

    SimState::View2D u = s.u;
    SimState::View2D v = s.v;

    Kokkos::parallel_for("compute_u_conv_rhs",
    Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
        {s.grid.u_i_begin(), s.grid.u_j_begin()},
        {s.grid.u_i_end(),   s.grid.u_j_end()}),
    KOKKOS_LAMBDA(int i, int j) {
        const double u_east  = avg_x(u, i,   j);
        const double u_west  = avg_x(u, i-1, j);
        const double duudx   = (u_east*u_east - u_west*u_west) * inv_dx;

        const double u_north = avg_y(u, i,   j);
        const double u_south = avg_y(u, i,   j-1);
        const double v_north = avg_x(v, i-1, j+1);
        const double v_south = avg_x(v, i-1, j);
        const double duvdy   = (u_north*v_north - u_south*v_south) * inv_dy;

        rhs_u(i, j) = -duudx - duvdy;
    });
}

    KOKKOS_LAMBDA(int i, int j) {
        const double u_east  = avg_x(u, i,   j);
        const double u_west  = avg_x(u, i-1, j);
        const double duudx   = (u_east*u_east - u_west*u_west) * inv_dx;

        const double u_north = avg_y(u, i,   j);
        const double u_south = avg_y(u, i,   j-1);
        const double v_north = avg_x(v, i-1, j+1);
        const double v_south = avg_x(v, i-1, j);
        const double duvdy   = (u_north*v_north - u_south*v_south) * inv_dy;

        rhs_u(i, j) = -duudx - duvdy;
    });

void compute_u_diff_rhs(const SimState& s, double re,
                        Kokkos::View<double**> rhs_u) {
    const double inv_dx2 = 1.0 / (s.grid.dx * s.grid.dx);
    const double inv_dy2 = 1.0 / (s.grid.dy * s.grid.dy);
    const double inv_re  = 1.0 / re;

    SimState::View2D u = s.u;

    Kokkos::parallel_for("compute_u_diff_rhs",
    Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
        {s.grid.u_i_begin(), s.grid.u_j_begin()},
        {s.grid.u_i_end(),   s.grid.u_j_end()}),
    KOKKOS_LAMBDA(int i, int j) {
        rhs_u(i, j) = inv_re * laplacian(u, i, j, inv_dx2, inv_dy2);
    });
}

    KOKKOS_LAMBDA(int i, int j) {
        rhs_u(i, j) = inv_re * laplacian(u, i, j, inv_dx2, inv_dy2);
    });

void compute_v_conv_rhs(const SimState& s,
                        Kokkos::View<double**> rhs_v) {
    const double inv_dx = 1.0 / s.grid.dx;
    const double inv_dy = 1.0 / s.grid.dy;

    SimState::View2D u = s.u;
    SimState::View2D v = s.v;

    Kokkos::parallel_for("compute_v_conv_rhs",
    Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
        {s.grid.v_i_begin(), s.grid.v_j_begin()},
        {s.grid.v_i_end(),   s.grid.v_j_end()}),
    KOKKOS_LAMBDA(int i, int j) {
        const double v_north = avg_y(v, i,   j);
        const double v_south = avg_y(v, i,   j-1);
        const double dvvdy   = (v_north*v_north - v_south*v_south) * inv_dy;

        const double v_east  = avg_x(v, i,   j);
        const double v_west  = avg_x(v, i-1, j);
        const double u_east  = avg_y(u, i+1, j-1);
        const double u_west  = avg_y(u, i,   j-1);
        const double duvdx   = (u_east*v_east - u_west*v_west) * inv_dx;

        rhs_v(i, j) = -duvdx - dvvdy;
    });
}

    KOKKOS_LAMBDA(int i, int j) {
        const double v_north = avg_y(v, i,   j);
        const double v_south = avg_y(v, i,   j-1);
        const double dvvdy   = (v_north*v_north - v_south*v_south) * inv_dy;

        const double v_east  = avg_x(v, i,   j);
        const double v_west  = avg_x(v, i-1, j);
        const double u_east  = avg_y(u, i+1, j-1);
        const double u_west  = avg_y(u, i,   j-1);
        const double duvdx   = (u_east*v_east - u_west*v_west) * inv_dx;

        rhs_v(i, j) = -duvdx - dvvdy;
    });

void compute_v_diff_rhs(const SimState& s, double re,
                        Kokkos::View<double**> rhs_v) {
    const double inv_dx2 = 1.0 / (s.grid.dx * s.grid.dx);
    const double inv_dy2 = 1.0 / (s.grid.dy * s.grid.dy);
    const double inv_re  = 1.0 / re;

    SimState::View2D v = s.v;

    Kokkos::parallel_for("compute_v_diff_rhs",
    Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
        {s.grid.v_i_begin(), s.grid.v_j_begin()},
        {s.grid.v_i_end(),   s.grid.v_j_end()}),
    KOKKOS_LAMBDA(int i, int j) {
        rhs_v(i, j) = inv_re * laplacian(v, i, j, inv_dx2, inv_dy2);
    });
}

    KOKKOS_LAMBDA(int i, int j) {
        rhs_v(i, j) = inv_re * laplacian(v, i, j, inv_dx2, inv_dy2);
    });

double compute_kinetic_energy(const SimState& s){
    SimState::View2D u = s.u;
    SimState::View2D v = s.v;

    double sum_u2 = 0.0;
    Kokkos::parallel_reduce("KE_u", 
        Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
            {s.grid.u_i_begin(), s.grid.u_j_begin()},
            {s.grid.u_i_end(), s.grid.u_j_end()}),
        KOKKOS_LAMBDA(int i, int j, double& lsum){
            lsum += u(i, j) * u(i, j);
        }, sum_u2);

    
    double sum_v2 = 0.0;
    Kokkos::parallel_reduce("KE_v", 
        Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
            {s.grid.v_i_begin(), s.grid.v_j_begin()},
            {s.grid.v_i_end(), s.grid.v_j_end()}),
        KOKKOS_LAMBDA(int i, int j, double& lsum){
            lsum += v(i, j) * v(i, j);
        }, sum_v2);
    
    return 0.5 * (sum_u2 + sum_v2) * s.grid.dx * s.grid.dy;
}

        KOKKOS_LAMBDA(int i, int j, double& lsum){
            lsum += u(i, j) * u(i, j);
        }, sum_u2);

        KOKKOS_LAMBDA(int i, int j, double& lsum){
            lsum += v(i, j) * v(i, j);
        }, sum_v2);

double compute_cfl_dt(const SimState& s, double cfl) {
    SimState::View2D u = s.u;
    SimState::View2D v = s.v;

    double u_max = 0.0;
    Kokkos::parallel_reduce("cfl_u",
        Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
            {s.grid.u_i_begin(), s.grid.u_j_begin()},
            {s.grid.u_i_end(),   s.grid.u_j_end()}),
        KOKKOS_LAMBDA(int i, int j, double& lmax) {
            const double val = Kokkos::fabs(u(i, j));
            if (val > lmax) lmax = val;
        },
        Kokkos::Max<double>(u_max));

    double v_max = 0.0;
    Kokkos::parallel_reduce("cfl_v",
        Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
            {s.grid.v_i_begin(), s.grid.v_j_begin()},
            {s.grid.v_i_end(),   s.grid.v_j_end()}),
        KOKKOS_LAMBDA(int i, int j, double& lmax) {
            const double val = Kokkos::fabs(v(i, j));
            if (val > lmax) lmax = val;
        },
        Kokkos::Max<double>(v_max));

    const double vel_max = Kokkos::max(u_max, v_max);
    con

        KOKKOS_LAMBDA(int i, int j, double& lmax) {
            const double val = Kokkos::fabs(u(i, j));
            if (val > lmax) lmax = val;
        },

        KOKKOS_LAMBDA(int i, int j, double& lmax) {
            const double val = Kokkos::fabs(v(i, j));
            if (val > lmax) lmax = val;
        },

double compute_l2_divergence(const SimState& s) {
    const double inv_dx = 1.0 / s.grid.dx;
    const double inv_dy = 1.0 / s.grid.dy;

    SimState::View2D u = s.u;
    SimState::View2D v = s.v;

    double sum_div2 = 0.0;
    Kokkos::parallel_reduce("L2_div",
        Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
            {s.grid.p_i_begin(), s.grid.p_j_begin()},
            {s.grid.p_i_end(),   s.grid.p_j_end()}),
        KOKKOS_LAMBDA(int i, int j, double& lsum) {
            double d = divergence(u, v, i, j, inv_dx, inv_dy);
            lsum += d * d;
        }, sum_div2);

    // RMS divergence: divide by cell count before taking root so the metric is
    // grid-size-independent and meaningful for mesh refinement studies.
    return Kokkos::sqrt(sum_div2 / (s.grid.nx * s.grid.ny));
}

        KOKKOS_LAMBDA(int i, int j, double& lsum) {
            double d = divergence(u, v, i, j, inv_dx, inv_dy);
            lsum += d * d;
        }, sum_div2);

void compute_pressure_rhs(const SimState& s,
                           Kokkos::View<double**> u_star,
                           Kokkos::View<double**> v_star,
                           double dt,
                           Kokkos::View<double**> rhs) {
    // divergence of predicted velocity / dt
    //   rhs(i,j) = (1/dt)( du_star/dx + dv_star/dy )
    const double inv_dx = 1.0 / s.grid.dx;
    const double inv_dy = 1.0 / s.grid.dy;
    const double inv_dt = 1.0 / dt;

    Kokkos::parallel_for("compute_pressure_rhs",
    Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
        {s.grid.p_i_begin(), s.grid.p_j_begin()},
        {s.grid.p_i_end(), s.grid.p_j_end()}),
    KOKKOS_LAMBDA(int i, int j) {
        const double dudx_star = (u_star(i+1,j) - u_star(i,j)) * inv_dx;
        const double dvdy_star = (v_star(i,j+1) - v_star(i,j)) * inv_dy;

        rhs(i,j) = (dudx_star + dvdy_star) * inv_dt;
    }
    );
}

    KOKKOS_LAMBDA(int i, int j) {
        const double dudx_star = (u_star(i+1,j) - u_star(i,j)) * inv_dx;
        const double dvdy_star = (v_star(i,j+1) - v_star(i,j)) * inv_dy;

        rhs(i,j) = (dudx_star + dvdy_star) * inv_dt;
    }

void correct_velocity(SimState& s,
                      Kokkos::View<double**> u_star,
                      Kokkos::View<double**> v_star,
                      double dt) {
    // subtract pressure gradient from predicted velocity
    //   u = u_star - dt * dp/dx
    //   v = v_star - dt * dp/dy

    const double inv_dx = 1.0 / s.grid.dx;
    const double inv_dy = 1.0 / s.grid.dy;

    SimState::View2D p = s.p;
    SimState::View2D u = s.u;
    SimState::View2D v = s.v;

    Kokkos::parallel_for(
        "correct_u",
        Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
            {s.grid.u_i_begin(), s.grid.u_j_begin()},
            {s.grid.u_i_end(),   s.grid.u_j_end()}),
        KOKKOS_LAMBDA(int i, int j) {
            const double dpdx = (p(i, j) - p(i - 1, j)) * inv_dx;
            u(i, j) = u_star(i, j) - dt * dpdx;
        });

    Kokkos::parallel_for(
        "correct_v",
        Kokkos::MDRangePolicy<Kokkos::Rank<2>>(
            {s.grid.v_i_begin(), s.grid.v_j_begin()},
    

        KOKKOS_LAMBDA(int i, int j) {
            const double dpdx = (p(i, j) - p(i - 1, j)) * inv_dx;
            u(i, j) = u_star(i, j) - dt * dpdx;
        });

        KOKKOS_LAMBDA(int i, int j) {
            const double dpdy = (p(i, j) - p(i, j - 1)) * inv_dy;
            v(i, j) = v_star(i, j) - dt * dpdy;
        });

static int ilog2(int n) {
    int k = 0;
    while (n > 1) { n >>= 1; ++k; }
    return k;
}

void PressureSolver::init(const MacGrid2D& grid) {
    if (initialized_) return;

    const int nx = grid.nx;
    const int ny = grid.ny;
    lx_ = grid.lx;
    ly_ = grid.ly;

    if (!is_power_of_2(nx) || !is_power_of_2(ny))
        throw std::invalid_argument(
            "PressureSolver: nx and ny must be powers of 2 for geometric multigrid");

    // Coarsest level has 2x2 cells (don't go to 1x1 — degenerate for periodic)
    const int num_levels = ilog2(std::min(nx, ny));  // e.g. 64 -> 6 levels (64..2)

    levels_.resize(num_levels);
    for (int l = 0; l < num_levels; ++l) {
        auto& lev = levels_[l];
        lev.nx = nx >> l;
        lev.ny = ny >> l;
        lev.dx = lx_ / lev.nx;
        lev.dy = ly_ / lev.ny;

        const std::string tag = "mg_l" + std::to_string(l) + "_";
        lev.phi = Kokkos::View<double**>(tag + "phi", lev.nx + 2, lev.ny + 2);
        lev.f   = Kokkos::View<double**>(tag + "f",   lev.nx + 2, lev.ny + 2);
        lev.r   = Kokkos::View<double

void PressureSolver::periodic_fill(Kokkos::View<double**> v,
                                   int nx, int ny) const {
    // x-direction ghosts
    Kokkos::parallel_for("mg_pfill_x", Kokkos::RangePolicy<>(1, ny + 1),
        KOKKOS_LAMBDA(int j) {
            v(0,      j) = v(nx, j);
            v(nx + 1, j) = v(1,  j);
        });

    // y-direction ghosts
    Kokkos::parallel_for("mg_pfill_y", Kokkos::RangePolicy<>(1, nx + 1),
        KOKKOS_LAMBDA(int i) {
            v(i, 0)      = v(i, ny);
            v(i, ny + 1) = v(i, 1);
        });

    // corners
    Kokkos::parallel_for("mg_pfill_c", Kokkos::RangePolicy<>(0, 1),
        KOKKOS_LAMBDA(int) {
            v(0,      0)      = v(nx, ny);
            v(nx + 1, 0)      = v(1,  ny);
            v(0,      ny + 1) = v(nx, 1);
            v(nx + 1, ny + 1) = v(1,  1);
        });
}

        KOKKOS_LAMBDA(int j) {
            v(0,      j) = v(nx, j);
            v(nx + 1, j) = v(1,  j);
        });

        KOKKOS_LAMBDA(int i) {
            v(i, 0)      = v(i, ny);
            v(i, ny + 1) = v(i, 1);
        });

        KOKKOS_LAMBDA(int) {
            v(0,      0)      = v(nx, ny);
            v(nx + 1, 0)      = v(1,  ny);
            v(0,      ny + 1) = v(nx, 1);
            v(nx + 1, ny + 1) = v(1,  1);
        });

void PressureSolver::smooth_rbgs(Level& lev, int sweeps) const {
    const int nx = lev.nx;
    const int ny = lev.ny;
    const double inv_dx2 = 1.0 / (lev.dx * lev.dx);
    const double inv_dy2 = 1.0 / (lev.dy * lev.dy);
    const double diag_inv = 1.0 / (2.0 * inv_dx2 + 2.0 * inv_dy2);

    auto phi = lev.phi;
    auto f   = lev.f;

    for (int s = 0; s < sweeps; ++s) {
        // --- Red pass: (i+j) % 2 == 0 ---
        periodic_fill(phi, nx, ny);
        Kokkos::parallel_for("mg_rbgs_red",
            Kokkos::MDRangePolicy<Kokkos::Rank<2>>({1, 1}, {nx + 1, ny + 1}),
            KOKKOS_LAMBDA(int i, int j) {
                if ((i + j) % 2 != 0) return;
                phi(i, j) = ((phi(i+1,j) + phi(i-1,j)) * inv_dx2
                           + (phi(i,j+1) + phi(i,j-1)) * inv_dy2
                           - f(i, j)) * diag_inv;
            });

        // --- Black pass: (i+j) % 2 == 1 ---
        periodic_fill(phi, nx, ny);
        Kokkos::parallel_for("mg_rbgs_black",
       

            KOKKOS_LAMBDA(int i, int j) {
                if ((i + j) % 2 != 0) return;
                phi(i, j) = ((phi(i+1,j) + phi(i-1,j)) * inv_dx2
                           + (phi(i,j+1) + phi(i,j-1)) * inv_dy2
                           - f(i, j)) * diag_inv;
            });

            KOKKOS_LAMBDA(int i, int j) {
                if ((i + j) % 2 != 1) return;
                phi(i, j) = ((phi(i+1,j) + phi(i-1,j)) * inv_dx2
                           + (phi(i,j+1) + phi(i,j-1)) * inv_dy2
                           - f(i, j)) * diag_inv;
            });