From d5816104cd96f527a31658fbaca910cfecc887f9 Mon Sep 17 00:00:00 2001 From: Saurabh Shukla <136088049+ss2098@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:10:59 -0800 Subject: [PATCH] Update adapt cost microbenchmark for current APIs --- scripts/adapt_cost_microbench.py | 190 +++++++++++++++++++++++-------- 1 file changed, 143 insertions(+), 47 deletions(-) diff --git a/scripts/adapt_cost_microbench.py b/scripts/adapt_cost_microbench.py index dfe435a7e..69c7a49ff 100644 --- a/scripts/adapt_cost_microbench.py +++ b/scripts/adapt_cost_microbench.py @@ -3,110 +3,206 @@ res-16 Ra=1e5, same setup as the saturation runner. Breaks the adaptation into metric / mover / remap+restokes. """ + from __future__ import annotations + import time import numpy as np import sympy import underworld3 as uw -from underworld3.meshing import ( - smooth_mesh_interior, metric_density_from_gradient) + +from underworld3.meshing import metric_density_from_gradient + RA, RES, r_inner, r_o = 1.0e5, 16, 0.5, 1.0 +DT_SAFETY = 0.1 + + +def scalar_dt(value): + """Make estimate_dt output safe for AdvDiffusionSLCN.solve().""" + try: + dt = float(value) + except TypeError: + arr = np.asarray(value, dtype=float) + dt = float(np.nanmin(arr)) + + if not np.isfinite(dt) or dt <= 0.0: + raise ValueError(f"Bad timestep from estimate_dt(): {value!r}") + + return DT_SAFETY * dt + + +def check_snes(system, label): + """Stop immediately if a solve diverged, so this benchmark is honest.""" + reason = system.snes.getConvergedReason() + if reason < 0: + raise RuntimeError(f"{label} diverged: SNES reason={reason}") + return reason + + +m = uw.meshing.Annulus( + radiusOuter=r_o, + radiusInner=r_inner, + cellSize=1.0 / RES, + qdegree=3, +) -m = uw.meshing.Annulus(radiusOuter=r_o, radiusInner=r_inner, - cellSize=1.0 / RES, qdegree=3) r, th = m.CoordinateSystem.R -v = uw.discretisation.MeshVariable("V", m, vtype=uw.VarType.VECTOR, - degree=2, continuous=True) -P = uw.discretisation.MeshVariable("P", m, vtype=uw.VarType.SCALAR, - degree=1, continuous=True) -T = uw.discretisation.MeshVariable("T", m, vtype=uw.VarType.SCALAR, - degree=3, continuous=True) + +v = uw.discretisation.MeshVariable( + "V", + m, + vtype=uw.VarType.VECTOR, + degree=2, + continuous=True, +) + +P = uw.discretisation.MeshVariable( + "P", + m, + vtype=uw.VarType.SCALAR, + degree=1, + continuous=True, +) + +T = uw.discretisation.MeshVariable( + "T", + m, + vtype=uw.VarType.SCALAR, + degree=3, + continuous=True, +) + stokes = uw.systems.Stokes(m, velocityField=v, pressureField=P) stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 stokes.tolerance = 1.0e-5 stokes.penalty = 0.0 + ur = m.CoordinateSystem.unit_e_0 + stokes.add_essential_bc((0.0, 0.0), m.boundaries.Lower.name) -stokes.add_natural_bc(1.0e6 * v.sym.dot(ur) * ur, - m.boundaries.Upper.name) -stokes.bodyforce = RA * (T.sym[0] - (r_o - r) / (r_o - r_inner)) * ur -adv = uw.systems.AdvDiffusionSLCN(m, u_Field=T, V_fn=v.sym, - verbose=False, theta=0.5, - monotone_mode="clamp") +stokes.add_essential_bc((0.0, 0.0), m.boundaries.Upper.name) + +stokes.bodyforce = RA * ( + T.sym[0] - (r_o - r) / (r_o - r_inner) +) * ur + +adv = uw.systems.AdvDiffusionSLCN( + m, + u_Field=T, + V_fn=v.sym, + verbose=False, + theta=0.5, + monotone_mode="clamp", +) + adv.constitutive_model = uw.constitutive_models.DiffusionModel adv.constitutive_model.Parameters.diffusivity = 1.0 adv.tolerance = 1.0e-4 + adv.add_dirichlet_bc(1.0, m.boundaries.Lower.name) adv.add_dirichlet_bc(0.0, m.boundaries.Upper.name) -init_t = (0.01 * sympy.sin(5.0 * th) - * sympy.sin(np.pi * (r - r_inner) / (r_o - r_inner)) - + (r_o - r) / (r_o - r_inner)) -T.data[...] = np.asarray(uw.function.evaluate( - init_t, T.coords)).reshape(-1, 1) + +init_t = ( + 0.01 + * sympy.sin(5.0 * th) + * sympy.sin(np.pi * (r - r_inner) / (r_o - r_inner)) + + (r_o - r) / (r_o - r_inner) +) + +T.data[...] = np.asarray( + uw.function.evaluate(init_t, T.coords) +).reshape(-1, 1) + X0 = np.asarray(m.X.coords).copy() X0_Tx = np.asarray(T.coords).copy() + stokes.solve(zero_init_guess=True) +check_snes(stokes, "initial Stokes solve") -# warm a few steps so the field is representative (plumes growing) +# Warm a few steps so the field is representative. for _ in range(8): - dt = adv.estimate_dt() + dt = scalar_dt(adv.estimate_dt()) adv.solve(timestep=dt, zero_init_guess=False) - stokes.solve(zero_init_guess=False) + check_snes(adv, "warmup AdvDiffusion solve") + + stokes.solve(zero_init_guess=True) + check_snes(stokes, "warmup Stokes solve") -# (1) time plain (adv+stokes) steps +# Time plain adv+stokes steps. N = 8 t0 = time.perf_counter() + for _ in range(N): - dt = adv.estimate_dt() + dt = scalar_dt(adv.estimate_dt()) adv.solve(timestep=dt, zero_init_guess=False) - stokes.solve(zero_init_guess=False) + check_snes(adv, "timing AdvDiffusion solve") + + stokes.solve(zero_init_guess=True) + check_snes(stokes, "timing Stokes solve") + t_step = (time.perf_counter() - t0) / N print(f"plain (adv+stokes) step : {t_step:6.3f} s (mean of {N})") -# (2) time ONE pristine adaptation, broken down -X_prev = np.asarray(m.X.coords).copy() -T_prev = np.asarray(T.data).copy() - +# Time ONE pristine adaptation, broken down. ta = time.perf_counter() vals0 = np.asarray(uw.function.evaluate(T.sym[0], X0_Tx)).reshape(-1) -m._deform_mesh(X0); T.data[:, 0] = vals0 +m.deform(X0) +T.data[:, 0] = vals0 t_remap_in = time.perf_counter() - ta tb = time.perf_counter() rho = metric_density_from_gradient(m, T, amp=8.0, name="mb") t_metric = time.perf_counter() - tb -X0c = np.asarray(m.X.coords).copy(); T0 = np.asarray(T.data).copy() +X0c = np.asarray(m.X.coords).copy() +T0 = np.asarray(T.data).copy() + tc = time.perf_counter() -smooth_mesh_interior(m, metric=rho, method="anisotropic", - method_kwargs=dict(aniso_cap=2.0, relax=0.2, - n_outer=8)) +uw.meshing.node_redistribution(m, rho) t_mover = time.perf_counter() - tc new_X = np.asarray(m.X.coords).copy() new_Tx = np.asarray(T.coords).copy() + td = time.perf_counter() -m._deform_mesh(X0c); T.data[...] = T0 +m.deform(X0c) +T.data[...] = T0 + valsN = np.asarray(uw.function.evaluate(T.sym[0], new_Tx)).reshape(-1) -m._deform_mesh(new_X); T.data[:, 0] = valsN + +m.deform(new_X) +T.data[:, 0] = valsN t_remap_out = time.perf_counter() - td te = time.perf_counter() -stokes.solve(zero_init_guess=False) +stokes.solve(zero_init_guess=True) +check_snes(stokes, "post-adaptation Stokes solve") t_restokes = time.perf_counter() - te -t_adapt = (t_remap_in + t_metric + t_mover + t_remap_out - + t_restokes) +t_adapt = ( + t_remap_in + + t_metric + + t_mover + + t_remap_out + + t_restokes +) + print(f" remap-in (eval+deform) : {t_remap_in:6.3f} s") print(f" metric (grad projection): {t_metric:6.3f} s") -print(f" MOVER (anisotropic) : {t_mover:6.3f} s") +print(f" MOVER (redistribution) : {t_mover:6.3f} s") print(f" remap-out (eval+deform x2) : {t_remap_out:6.3f} s") print(f" re-stokes : {t_restokes:6.3f} s") print(f"ONE pristine adaptation TOTAL: {t_adapt:6.3f} s") -print(f"\nratio adaptation / (adv+stokes step) = " - f"{t_adapt / t_step:5.1f}x") -print(f"amortised over adapt-every-5: +{t_adapt/(5*t_step):.1f}x " - f"work vs a non-adaptive res-16 run " - f"(5 steps cost {5*t_step:.2f}s + 1 adapt {t_adapt:.2f}s)") + +print( + f"\nratio adaptation / (adv+stokes step) = " + f"{t_adapt / t_step:5.1f}x" +) + +print( + f"amortised over adapt-every-5: +{t_adapt / (5 * t_step):.1f}x " + f"work vs a non-adaptive res-16 run " + f"(5 steps cost {5 * t_step:.2f}s + 1 adapt {t_adapt:.2f}s)" +)