"""A minimal adaptive-mesh-refinement code for Burgers' equation in one dimension. Block-structured AMR in the style of Berger & Oliger (1984) and Berger & Colella (1989): a hierarchy of levels with refinement ratio R, subcycling in time, conservative restriction and prolongation, and refluxing at the coarse-fine faces. To keep the code short, each level stores one array over the whole domain plus a mask of the cells it actually covers. Production codes store only the refined patches. python amr1d.py # runs the example and prints a summary """ import numpy as np R = 2 # refinement ratio between levels NG = 2 # ghost cells at each end of a level NEST = 2 # a refined cell must be this many coarse cells inside its parent level NBUF = 2 # buffer cells added around flagged cells TOL = 0.08 # flag a cell when the jump to a neighbour exceeds this CFL = 0.8 XMIN, XMAX = 0.0, 3.0 def flux(u): return 0.5 * u * u def riemann(ul, ur): """Godunov flux for Burgers' equation, from the exact Riemann solution.""" fl, fr = flux(ul), flux(ur) shock = np.maximum(fl, fr) fan = np.where(ul > 0, fl, np.where(ur < 0, fr, 0.0)) return np.where(ul > ur, shock, fan) def minmod(a, b): return np.where(a * b > 0, np.sign(a) * np.minimum(abs(a), abs(b)), 0.0) class Level: def __init__(self, l, n0): self.l = l self.n = n0 * R**l self.h = (XMAX - XMIN) / self.n self.x = XMIN + (np.arange(self.n) + 0.5) * self.h self.u = np.zeros(self.n + 2 * NG) # cell averages, ghosts included self.u_old = self.u.copy() # state at the start of the current step self.own = np.zeros(self.n, bool) # cells covered by this level self.F = np.zeros(self.n + 1) # face fluxes from the last step self.Fsum = np.zeros(self.n + 1) # dt-weighted fluxes summed over substeps self.work = 0 # cell updates so far @property def cells(self): """The level's cells without the ghosts (a view, so writes go through).""" return self.u[NG:NG + self.n] # ------------------------------------------------------------ between levels def prolong(uc): """Interpolate a coarse array (ghosts included) onto the fine cells under it. Each coarse cell becomes R fine cells on a slope-limited line, so the average of the fine cells equals the coarse value and no new extrema are created. The result is aligned with the fine level's array. """ s = np.zeros_like(uc) s[1:-1] = minmod(uc[1:-1] - uc[:-2], uc[2:] - uc[1:-1]) offs = (np.arange(R) + 0.5) / R - 0.5 fine = (uc[:, None] + s[:, None] * offs).ravel() start = NG * R - NG return fine[start:start + (len(uc) - 2 * NG) * R + 2 * NG] def restrict(coarse, fine): """Overwrite coarse cells under the fine level with the fine averages.""" under = fine.own.reshape(-1, R).all(axis=1) avg = fine.cells.reshape(-1, R).mean(axis=1) coarse.cells[under] = avg[under] return under def fill_ghosts(lv, coarse, frac): """Give every cell the level doesn't own a value. On the base level that means the outflow boundary condition. On finer levels it means interpolating from the coarse level, taken at a fraction `frac` of the way through the coarse step. """ if coarse is None: lv.u[:NG] = lv.u[NG] lv.u[-NG:] = lv.u[-NG - 1] return uc = (1 - frac) * coarse.u_old + frac * coarse.u need = np.ones(lv.n + 2 * NG, bool) need[NG:-NG] = ~lv.own lv.u[need] = prolong(uc)[need] # ------------------------------------------------------------- time stepping def step(lv, dt): """One first-order Godunov step on the cells the level owns.""" u = lv.u lv.F = riemann(u[NG - 1:NG + lv.n], u[NG:NG + lv.n + 1]) lv.Fsum += dt * lv.F du = -dt / lv.h * (lv.F[1:] - lv.F[:-1]) lv.cells[lv.own] += du[lv.own] lv.work += int(lv.own.sum()) def reflux(coarse, fine, under, dt): """Make both sides of every coarse-fine face agree on the flux through it. The coarse cell next to the fine region was updated with the coarse flux. Replace that by the time average of the fine fluxes, which is what the fine cells on the other side of the face actually used. """ Feff = fine.Fsum[::R] / dt dF = dt / coarse.h * (coarse.F - Feff) starts = np.flatnonzero(under[1:] & ~under[:-1]) + 1 # faces where a fine region begins ends = np.flatnonzero(under[:-1] & ~under[1:]) + 1 # faces where one ends coarse.cells[starts - 1] += dF[starts] coarse.cells[ends] -= dF[ends] def advance(levels, l, dt, frac=0.0, do_reflux=True): """Advance level l by dt, then every finer level by R substeps each.""" lv = levels[l] fill_ghosts(lv, levels[l - 1] if l > 0 else None, frac) lv.u_old[:] = lv.u step(lv, dt) if l + 1 < len(levels): fine = levels[l + 1] fine.Fsum[:] = 0.0 for k in range(R): advance(levels, l + 1, dt / R, k / R, do_reflux) under = restrict(lv, fine) if do_reflux: reflux(lv, fine, under, dt) # ------------------------------------------------------------------ regridding def flag(lv): """Cells whose jump to either neighbour is larger than TOL.""" u = lv.u jump = np.maximum(abs(u[NG + 1:NG + lv.n + 1] - u[NG:NG + lv.n]), abs(u[NG:NG + lv.n] - u[NG - 1:NG + lv.n - 1])) return (jump > TOL) & lv.own def dilate(m, k): out = m.copy() for s in range(1, k + 1): out[s:] |= m[:-s] out[:-s] |= m[s:] return out def regrid(levels, l): """Rebuild level l+1 from the flags on level l.""" coarse, fine = levels[l], levels[l + 1] allowed = ~dilate(~coarse.own, NEST) # proper nesting want = dilate(flag(coarse), NBUF) & allowed fill_ghosts(fine, coarse, 1.0) # new cells get interpolated data fine.own = np.repeat(want, R) # ---------------------------------------------------------------------- driver def leaves(levels): """For each level, the cells that are not covered by a finer level.""" out = [] for l, lv in enumerate(levels): leaf = lv.own.copy() if l + 1 < len(levels): leaf &= ~levels[l + 1].own.reshape(-1, R).all(axis=1) out.append(leaf) return out def mass(levels): return sum(lv.h * lv.cells[leaf].sum() for lv, leaf in zip(levels, leaves(levels))) def mass_error(history): """Change in total mass that is not explained by flow through the ends.""" return history[:, 1] - history[0, 1] - history[:, 2] def run(u0, n0=64, lmax=3, T=3.0, do_reflux=True, snapshots=()): levels = [Level(l, n0) for l in range(lmax + 1)] base = levels[0] base.own[:] = True base.cells[:] = u0(base.x) fill_ghosts(base, None, 0.0) for l in range(lmax): regrid(levels, l) t, inflow, history, saved = 0.0, 0.0, [(0.0, mass(levels), 0.0)], {} pending = sorted(snapshots) def snapshot(): saved[pending.pop(0)] = [(lv.x.copy(), lv.cells.copy(), leaf) for lv, leaf in zip(levels, leaves(levels))] if pending and pending[0] == 0.0: snapshot() while t < T - 1e-12: umax = max(abs(lv.cells[lv.own]).max() for lv in levels if lv.own.any()) dt = min(CFL * base.h / umax, T - t) if pending: dt = min(dt, pending[0] - t) advance(levels, 0, dt, do_reflux=do_reflux) t += dt inflow += dt * (base.F[0] - base.F[-1]) # mass entering through the ends for l in range(lmax): regrid(levels, l) history.append((t, mass(levels), inflow)) if pending and abs(t - pending[0]) < 1e-12: snapshot() return levels, np.array(history), saved def bump(x): return np.exp(-0.5 * ((x - 0.5) / 0.15) ** 2) if __name__ == "__main__": levels, hist, _ = run(bump) work = sum(lv.work for lv in levels) print(f"levels: {len(levels)}, finest spacing {levels[-1].h:.4f}") print(f"cell updates: {work:,}") print(f"mass error: {mass_error(hist)[-1]:.1e}") uni, hist_u, _ = run(bump, n0=levels[-1].n, lmax=0) print(f"uniform grid at the same resolution: {uni[0].work:,} cell updates")