CelerisLab/examples/fixed_patch_jet.py

"""Low-Re startup flow in an x-periodic confined domain; nozzle remains coarse.

From this checkout, after checking the selected physical GPU is idle:
  PYTHONPATH="$PWD/src" CUDA_VISIBLE_DEVICES=0 python examples/fixed_patch_jet.py --output /tmp/patch-jet

The fixed-amplitude start includes acoustic transients, not a developed jet.
Only native leaf cells are plotted. No interpolation, restart or flow oracle.
Scale estimates use the continuous peak A=.02, not inlet mean speed: W/A=600,
H/A=2250, W**2/nu=1440, H**2/nu=20250, where H=NY-3=45 between on-site planes.
An acoustic crossing (~78 steps) does not imply that transients have disappeared.
Mach is sampled over all leaves every 50 steps and is usually inlet dominated;
the discrete inlet peak is slightly below continuous A. This is not an accuracy test.
The 120s checks are cooperative, not an exact interruption deadline. main elapsed
includes construction, advancement and plotting; it is not pure CFD timing.
"""
import argparse
from dataclasses import replace
import json
import os
from pathlib import Path
import tempfile
import time

import numpy as np

import CelerisLab
from CelerisLab.config import load_lbm_config
from CelerisLab.cuda.context import CudaContext
from CelerisLab.lbm.equilibrium import compute_feq_d2q9, compute_macro_from_ddf
from CelerisLab.lbm.refinement import ACTIVE
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper

BOUNDS, AMPLITUDE, WIDTH = (16, 4, 32, 28), .02, 12
TIMES, MAX_MACH, MAX_SECONDS = (0, 300, 1200), .1, 120


def measure(arrays, masks, centers, initial_mass):
    fields = []
    for array, mask in zip(arrays, masks):
        if not np.isfinite(array[:, mask]).all():
            raise RuntimeError("Nonfinite leaf populations")
        field = np.full((*mask.shape, 3), np.nan)
        field[mask] = [compute_macro_from_ddf(f) for f in array[:, mask].T]
        fields.append(field)
    leaves = np.concatenate([f[m] for f, m in zip(fields, masks)])
    speed = np.hypot(leaves[:, 1], leaves[:, 2])
    mass = sum(area*a[:, m].sum(dtype=np.float64)
               for area, a, m in zip((1, .25), arrays, masks))
    record = dict(finite=bool(np.isfinite(leaves).all()), rho_min=float(leaves[:, 0].min()),
                  rho_max=float(leaves[:, 0].max()), speed_max=float(speed.max()),
                  sampled_max_Mach=float(speed.max()*np.sqrt(3)),
                  density_deviation=float(np.abs(leaves[:, 0]-1).max()),
                  discrete_leaf_mass=float(mass), mass_change=float(mass-initial_mass))
    if not record['finite'] or record['rho_min'] <= 0 or record['sampled_max_Mach'] > MAX_MACH:
        raise RuntimeError(f"Finite/positive-density/Ma<={MAX_MACH} guard failed: {record}")
    walls = fields[0][[1, -2]]
    record['macro_sum_rho_v'] = dict(zip(('Q_in', 'Q_out'),
                                       (walls[:, :, 0]*walls[:, :, 2]).sum(axis=1).tolist()))
    record['centerline'] = []
    for level, y in ((1, 16.25), (0, 32.5)):
        xs, ys = centers[level]
        columns = np.argsort(np.abs(xs-24))[:2]
        row = int(np.flatnonzero(ys == y)[0])
        values = fields[level][row, columns].mean(axis=0)
        record['centerline'].append(dict(level=level, x=24, y=y,
            sampled_x=sorted(xs[columns].tolist()), rule='arithmetic mean of two neighboring native columns',
            rho_ux_uy=values.tolist()))
    return fields, record


def draw(out, snapshots, masks, centers):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    from matplotlib.collections import PolyCollection

    def decorate(ax):
        ax.set(xlim=(0, 48), ylim=(1.5, 46.5), aspect='equal', xlabel='X', ylabel='Y')
        ax.plot([16, 32, 32, 16, 16], [4, 4, 28, 28, 4], color='tab:red', lw=1)
        ax.axhline(1.5, color='black', lw=1)
        ax.axhline(46.5, color='black', lw=1, ls='--')
        ax.plot([18, 30], [1.5, 1.5], color='tab:orange', lw=4, clip_on=False)

    fig, ax = plt.subplots(figsize=(6, 6), constrained_layout=True)
    for mask, (xs, ys), dx in zip(masks, centers, (1, .5)):
        yy, xx = np.meshgrid(ys, xs, indexing='ij')
        cells = np.stack((xx[mask], yy[mask]), axis=-1)
        vertices = cells[:, None, :] + dx*np.array([[-.5, -.5], [.5, -.5], [.5, .5], [-.5, .5]])
        ax.add_collection(PolyCollection(vertices, facecolors='none', edgecolors='0.55', linewidths=.3))
    decorate(ax)
    ax.set_title('Native leaf grid: coarse dx=1, fine dx=0.5\n'
                 'Red: internal patch; orange: coarse slot; dashed: reservoir\n'
                 'Flow may cross patch sides; patch does not contain all startup flow', fontsize=10)
    fig.savefig(out/'layout.png', dpi=150, bbox_inches='tight', pad_inches=.15)
    plt.close(fig)
    vmax = max(1., *(np.nanmax(np.hypot(f[..., 1], f[..., 2]))/AMPLITUDE
                    for fields in snapshots for f in fields))
    fig, axes = plt.subplots(1, 3, figsize=(12, 4.8), constrained_layout=True)
    for ax, step, fields in zip(axes, TIMES, snapshots):
        for field, mask, (xs, ys), dx in zip(fields, masks, centers, (1, .5)):
            speed = np.ma.masked_where(~mask, np.hypot(field[..., 1], field[..., 2])/AMPLITUDE)
            image = ax.pcolormesh(np.r_[xs-dx/2, xs[-1]+dx/2],
                                 np.r_[ys-dx/2, ys[-1]+dx/2], speed,
                                 shading='flat', cmap='viridis', vmin=0, vmax=vmax, rasterized=True)
            stride = int(4/dx)
            yy, xx = np.meshgrid(ys, xs, indexing='ij')
            selection = np.zeros_like(mask)
            selection[::stride, ::stride] = mask[::stride, ::stride]
            ax.quiver(xx[selection], yy[selection], field[..., 1][selection], field[..., 2][selection],
                      angles='xy', scale_units='xy', scale=AMPLITUDE/3, width=.005, color='white')
        decorate(ax)
        ax.set_title(f'coarse t={step}, fine steps={2*step}')
    fig.colorbar(image, ax=axes, label='speed / imposed amplitude')
    fig.suptitle('Low-Re startup flow in an x-periodic confined domain\n'
                 'Fixed-amplitude start includes acoustic transients; not steady/developed\n'
                 'Nozzle remains coarse; flow may cross patch sides (not fully enclosed)', fontsize=11)
    fig.savefig(out/'velocity.png', dpi=150)
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, required=True)
    out = parser.parse_args().output.resolve()
    root = Path(__file__).resolve().parents[1]
    if Path(CelerisLab.__file__).resolve() != root/'src/CelerisLab/__init__.py':
        raise RuntimeError(f'Set PYTHONPATH to {root / "src"}')
    out.mkdir(parents=True, exist_ok=False)
    cfg = replace(load_lbm_config(str(root/'src/CelerisLab/configs/config_lbm.json')),
                  nx=48, ny=48, viscosity=.1, velocity=0., rho=1., wall_slot_width=WIDTH,
                  wall_normal_velocity=AMPLITUDE, y_wall_bc='on_site_velocity', top_wall_bc='reservoir',
                  lattice_model='D2Q9', collision='SRT', streaming='double_buffer',
                  store_precision='FP32', ddf_shifting=False, les_enabled=False)
    report = dict(config=dict(nx=48, ny=48, bounds=BOUNDS, W=WIDTH, A=AMPLITUDE, nu=.1,
                              rho=1, tau_coarse=.8, tau_fine=1.1, coarse_boundary='slot_reservoir'),
                  guards=dict(check_every=50, positive_density=True, finite=True,
                              max_Mach=MAX_MACH, max_seconds=MAX_SECONDS,
                              timing='cooperative checks, not an exact interruption deadline'),
                  scale_estimates=dict(velocity_basis='continuous peak A, not inlet mean', H=cfg.ny-3,
                      W_over_A=WIDTH/AMPLITUDE, H_over_A=(cfg.ny-3)/AMPLITUDE,
                      W_squared_over_nu=WIDTH**2/cfg.viscosity,
                      H_squared_over_nu=(cfg.ny-3)**2/cfg.viscosity,
                      acoustic_crossing=(cfg.ny-3)*np.sqrt(3)),
                  notes='Node mass uses areas 1 and 1/4, not an exact physical volume integral. '
                        'Q_in/Q_out are bottom/top macro sum(rho*vy), not a closed population budget. '
                        'Mach maxima are sampled every 50 steps, not every-step peaks. '
                        'One acoustic crossing does not imply transient decay; no steady-state claim.',
                  samples=[], status='running', last_synchronized_coarse_step=0)
    start, run = time.monotonic(), None
    try:
        with CudaContext(0) as context, FixedPatchStepper(cfg, context=context,
                coarse_boundary='slot_reservoir', refinements=[BOUNDS]) as run:
            p = run.patch
            masks = [np.ones(p.coarse_shape, bool), p.fine_kind == ACTIVE]
            masks[0][[0, -1]] = False
            masks[0][4:28, 16:32] = False
            centers = [(origin[0]+dx*(np.arange(shape[1])+.5),
                        origin[1]+dx*(np.arange(shape[0])+.5))
                       for shape, origin, dx in ((p.coarse_shape, (0, 0), 1),
                                                 (p.fine_shape, p.fine_origin, .5))]
            initial = [np.broadcast_to(compute_feq_d2q9(1., 0., 0.)[:, None, None], (9, *m.shape)).copy()
                       for m in masks]
            for x in range(cfg.nx):
                v = AMPLITUDE*max(0., 1-((2*x-(cfg.nx-1))/WIDTH)**2)**2
                initial[0][:, 1, x] = compute_feq_d2q9(1., 0., v)
            run.initialize(*initial)
            mass0 = sum(area*a[:, m].sum(dtype=np.float64)
                        for area, a, m in zip((1, .25), initial, masks))
            snapshots = []
            for step in range(0, TIMES[-1]+1, 50):
                if time.monotonic()-start > MAX_SECONDS:
                    raise TimeoutError('120 second demonstration limit reached')
                if step:
                    run.step(50)
                report['last_synchronized_coarse_step'] = run.step_count
                if time.monotonic()-start > MAX_SECONDS:
                    raise TimeoutError('120 second demonstration limit reached after chunk')
                fields, sample = measure(run.download(), masks, centers, mass0)
                sample.update(t=step, coarse_steps=step, fine_steps=2*step)
                report['samples'].append(sample)
                print(json.dumps(sample), flush=True)
                if step in TIMES:
                    snapshots.append(fields)
        with tempfile.TemporaryDirectory(prefix='.plot-', dir=out) as cache:
            os.environ['MPLCONFIGDIR'] = cache
            draw(out, snapshots, masks, centers)
        if time.monotonic()-start > MAX_SECONDS:
            raise TimeoutError('120 second demonstration limit reached after plotting')
        report['status'] = 'complete'
    except BaseException as exc:
        report.update(status='failed', error=f'{type(exc).__name__}: {exc}',
                      last_synchronized_coarse_step=run.step_count if run is not None else 0)
        raise
    finally:
        report['elapsed_seconds'] = time.monotonic()-start
        report['elapsed_scope'] = 'main construction, advancement and plotting; not pure CFD time'
        (out/'summary.json').write_text(json.dumps(report, indent=2)+'\n')
        print(f'{report["status"]}: last synchronized coarse step '
              f'{report["last_synchronized_coarse_step"]}; output={out}', flush=True)


if __name__ == '__main__':
    main()