Repository navigation
Smooth the stagnant-lid kernels and validate rheology inputs - #85
Merged
Merged
Conversation
Verify near-zero surface pressures within round-off tolerance clamp to zero across rheology kernels, while negative pressures beyond tolerance raise ValueError.
Allow pressures down to minus one Pascal in input validation to accept surface round-off from atmospheric and interior solvers. Clamp negative values within this tolerance to zero Pascal before evaluating Arrhenius enthalpy, yield stress, and stagnant lid properties.
Assert that negative driving stress in compute_effective_viscosity lid mode raises ValueError on eager NumPy and returns NaN under JAX JIT.
Change driving stress validation rule from finite to nonneg in compute_effective_viscosity lid mode. Driving stress represents a physical stress magnitude that must be non-negative.
Assert that compute_arrhenius_viscosity produces strictly positive results without floating point underflow to zero when solid viscosity is small.
Adjust the lower clip bound in compute_arrhenius_viscosity dynamically when solid viscosity is small. This prevents the product of solid viscosity and exponent from underflowing to zero.
Assert that compute_yield_stress, compute_strain_rate_local, compute_mlt_top_slope, and viscous_mixing_length_factor reject out of bounds or non real arguments on NumPy and evaluate to NaN under JAX JIT.
Add input validation and JAX NaN masking to compute_yield_stress, compute_strain_rate_local, compute_mlt_top_slope, and viscous_mixing_length_factor. Update test asserting invalid negative pressure to expect ValueError.
Assert that nonneg_or_posinf exists in validation rules and wrap the infinite yield stress call in test_limit_infinite_yield_stress to convert mutant failures into test assertions.
Rename the validation rule key not_neginf to nonneg_or_posinf across rheology modules. This accurately reflects that the rule admits non-negative finite values and positive infinity.
Add fixture based test asserting bitwise output equivalence against origin main across all rheology kernels on a valid input grid. Add Hypothesis property tests verifying input validation boundaries across kernel functions.
…rge input checks Add reference links to the solid-state convection explanation document in docstrings for SolidRheologyParams and compute_stagnant_lid_state. Consolidate input validation checks in compute_arrhenius_viscosity into a single call.
Expand Hypothesis test coverage across pressure roundoff bitwise parity, Arrhenius underflow lower bounds, and input validation for helpers. Assert ValueError on eager NumPy with matching field names, NaN return under JAX JIT compilation, and finite outputs on valid domain draws.
Remove duplicate checks in compute_arrhenius_viscosity that repeat compute_arrhenius_enthalpy. Merge scalar and array branches in input validation, inline activation parameter in enthalpy calculation, and remove redundant temperature and hydration guards in Arrhenius viscosity.
Assert bitwise output parity when scalar parameters in rheology kernels are supplied as single precision float32 values compared with standard float values.
…ence temperature Ensure scalar parameters passed as float32 in compute_arrhenius_enthalpy and compute_arrhenius_viscosity are explicitly cast to float64, maintaining bitwise parity against standard float inputs under NumPy 2 promotion rules.
Remove the redundant t_sol_0 keyword argument from compute_stagnant_lid_state so the zero-pressure solidus reference is taken from params.t_sol_0. Update compute_solidus_reference_values.py, checkpoint_profile_175kyr.npz, and the solid-state convection documentation to evaluate the PALEOS MgSiO3 solidus (0.8 * T_liq) and report all linear-fallback crossings over 0 to 135 GPa.
Commit tests/reference/build_checkpoint_profile_175kyr.py to rebuild and verify tests/reference/checkpoint_profile_175kyr.npz bitwise from the 175 kyr snapshot and the PALEOS MgSiO3 melting curve, and document which arrays and slices are run values versus PALEOS-mapped values in every test that loads the fixture. Extend tools/verification/compute_solidus_reference_values.py with T_sol(1 GPa), maximum relative deviation, and the Chebyshev minimax and ordinary least-squares linear fits over 0 to 135 GPa, sync docs/Explanations/solid_state_convection.md and SolidRheologyParams to those tool outputs, document the measured dw_weight_df_max regression bound, and remove the duplicate t_sol_0 range check in compute_stagnant_lid_state.
Refer to the 175 kyr fixture on main in the module docstring of tests/reference/build_checkpoint_profile_175kyr.py while keeping MAIN_BASE_COMMIT in the code.
Test the log10 tanh yield blend in lid mode, the softplus viscosity ceiling, the smooth lid-presence and convective-activity weights, the smooth temperature-gradient floor and lid-mask cap, the smooth maximum in boundary-layer weights, the sign-changing total-flux denominator floor, the C2 quintic solidus melt ramp, the local strain-rate floor and softmin cap, the subthreshold convective-flux transition center, and the runtime weight and viscosity assertions.
Replace hard piecewise switches and kinks in aragog.rheology and aragog.rheology_lid with smooth gates: add a log10 softplus viscosity ceiling (width 0.05) and non-real scalar/array input rejection in compute_arrhenius_viscosity, smooth log10 tanh and softplus yield closures plus a 1e-100 strain-rate regularizer in compute_effective_viscosity, a Boltzmann-weighted smooth maximum in _boundary_weights, a compactly supported C2 quintic solidus melt ramp on [0.008, 0.012] in _melt_ramp, smooth w_no_lid gating and scale_cell/scale_mask softplus floors/caps, sign-change total-flux floor protection, and fixed-reference sub-threshold fallback weighting in compute_stagnant_lid_state.
Compare eta_i, T_i, and mu_i from compute_stagnant_lid_state against reference literals from main across interior_flux_fraction in {0.05, 0.10, 0.20, 0.50} for uniform convective flux fractions in {0.0, 6.00843e-5, 1.0e-4, 1.0e-3, 0.03} and on the 175 kyr profile fixture to a relative tolerance of 1e-12 outside the transition bands.
…tl/ssc-wp2-kernels
…l/ssc-wp2-kernels
Evaluate w_weight inside _boundary_weights from w_conv_int using params.interior_flux_fraction and _smooth_max rather than a fixed 0.05 reference threshold, restore the W_Ti_tilde fallback and has_lid temperature difference, and update the solid-state convection documentation and subthreshold transition test.
In the merged kernel branch, compute_stagnant_lid_state takes the zero-pressure reference solidus temperature solely from the parameters object. Remove the unused t_sol_0 key from the valid stagnant lid inputs helper dictionary in test_rheology_hypothesis.
Assert gated d_lid and tau_d outputs at W_sum = 0.75e-6 and 1.25e-6 with interior_flux_fraction = 0.01 so w_act_gate stays active while w_sum_gate transitions across 1e-6 with width 2.5e-7.
Combine the PALEOS solidus calibration and domain/underflow guards on tl/ssc-wp2-kernels with the smooth stagnant-lid and viscosity-ceiling closures from tl/ssc-wp2-smooth-lid, preserving both sides' behaviour across all seven conflict hunks.
Gate the direct vs * exp(exponent) branch on (d_cap < -40.0) & (exponent <= max_exponent) so small viscosity_solid values evaluate 10**log10_clipped above exponent == 700 instead of plateauing at vs * exp(700) until d_cap == -40.0.
Combine the two-argument _compute_f_tot_safe(F_tot, xp) call from origin/tl/ssc-wp2-kernels with the xp.where(f_raw >= 1.0, 1.0, f_raw) flux fraction cap in compute_stagnant_lid_state.
Extend test_arrhenius_viscosity_grad_finite_at_masked_invalid_nodes to check that jax.grad of a masked log10 viscosity sum is zero rather than NaN when masked temperature, water_prefactor, or arrhenius_t_ref values are zero, negative, or NaN.
Replace xp.where(x <= 0.0, 1.0, x) with xp.where(x > 0.0, x, 1.0) for t_safe, t_ref_safe, and wp_pos in compute_arrhenius_viscosity so masked NaN values map to 1.0 in the unmasked branch and keep jax.grad finite while invalid outputs remain NaN.
Add tests verifying that None inputs for theta, radii, r_in, r_out, mixing_length, and viscous_velocity raise ValueError. Add test for viscous mixing length factor with r_in equal to r_out when radii equals r_in. Add test verifying that invalid domain configurations evaluate to NaN under JAX eager and jit execution.
Domain checks enforce that radii lie within [r_in, r_out] with r_in strictly less than r_out. Because boundary slopes are positive, viscous mixing length factors are strictly non-negative throughout the physical domain. The outer maximum clamp is redundant dead code and a non-differentiable switch.
…idus Extend the masked invalid-node gradient test to cover pressure at -5 Pa, NaN, and infinity and non-positive or non-finite solidus temperatures with gradients checked for temperature and all scalar parameters.
Replace invalid pressure, solidus temperature, and enthalpy values with safe finite operands before the exponential in compute_arrhenius_enthalpy and compute_arrhenius_viscosity while keeping the output NaN mask unchanged, so jax.grad stays finite at masked nodes.
Check r_in and r_out bounds when both are finite and evaluate radius range checks on finite radii only. Adopt repo mesh rule max(1 m, 1e-9 span) for radius tolerance. Validate None inputs for eta_diff and required array fields in stagnant lid state. Use 1e-300 floor for log_eff in lid viscosity blend. Rename scalar validator to _require_scalar and assert f_raw at saturation test node. Update documentation and docstrings.
Catch unexpected exceptions and assert exception type and message for domain tolerance and stagnant lid input validations to ensure clean assertion failures under mutation. Pin effective viscosity output for sub-1e-30 diffusion viscosities to verify the 1e-300 floor in the lid viscosity blend.
Rely on IEEE 754 comparison behavior where comparisons with NaN evaluate to False to avoid filtering finite radii explicitly during mixing length domain checks.
Verify that compute_arrhenius_viscosity returns the log10 viscosity ceiling on both NumPy and JAX when the activation enthalpy overflows float64 to positive infinity.
Guard activation enthalpy with xp.isnan rather than xp.isfinite before computing the Arrhenius exponent so that an activation enthalpy overflow to positive infinity reaches the log10 viscosity ceiling while NaN enthalpies at invalid masked nodes remain replaced by zero for finite reverse-mode gradients.
Catch Exception and assert both ValueError type and the required field name in the message for every None test case in rheology kernel and hypothesis suites, calling pytest.fail only if no exception is raised.
Capture the exception in a local variable and assert both ValueError type and the required field name in the message for all None test cases in rheology kernel and hypothesis suites. This ensures that a missing check fails on an AssertionError rather than an unasserted fail.
There was a problem hiding this comment.
🟡 Changes recommended
The local viscosity closure has a zero-strain discontinuity, and the reference tool can mislabel a nonoptimal fit as minimax.
2 open issues
What changed in this PR
This PR prepares Aragog’s standalone solid-rheology kernels for future solid-state convection integration. The kernels are not yet called by the solver, so solver output remains unchanged.
Changes:
- Smooth viscosity and stagnant-lid transitions and expand input validation.
- Update PALEOS solidus comparisons and checkpoint fixture reconstruction.
- Extend regression, property-based, and NumPy/JAX parity checks.
| File | Description |
|---|---|
tools/verification/compute_solidus_reference_values.py |
Adds PALEOS comparisons, line fits, and reporting of all crossings. |
tests/test_rheology_parity.py |
Updates parity expectations and checks fixture consistency. |
tests/test_rheology_params_one_owner.py |
Checks parameter ownership, viscosity ceilings, and documentation. |
tests/test_rheology_hypothesis.py |
Adds validation, bounds, and smooth-transition properties. |
tests/test_foley_bercovici_lid_stress.py |
Updates stress regressions and smooth-gate expectations. |
tests/reference/build_rheology_kernel_reference_outputs.py |
Builds and verifies baseline kernel reference outputs. |
tests/reference/build_checkpoint_profile_175kyr.py |
Reconstructs and verifies the PALEOS-mapped checkpoint fixture. |
src/aragog/rheology.py |
Adds input checks, pressure tolerance, and smooth viscosity ceilings. |
src/aragog/rheology_lid.py |
Smooths lid diagnostics and viscosity closures; centralizes solidus parameters. |
Review effort: Balanced
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #85 +/- ##
==========================================
+ Coverage 97.96% 98.00% +0.03%
==========================================
Files 44 44
Lines 6882 7010 +128
Branches 763 765 +2
==========================================
+ Hits 6742 6870 +128
Misses 67 67
Partials 73 73
Flags with carried forward coverage won't be shown. Click here to find out more.
|
Assert that compute_solidus_reference_values evaluates the Chebyshev minimax linear fit for a non-convex table via linear programming, achieving the 375 K maximum error line instead of the shifted secant error of 450 K.
Evaluate the Chebyshev minimax linear fit in compute_solidus_reference_values using scipy.optimize.linprog over a dense grid with table nodes included. This replaces the shifted secant approximation, which is minimax only for concave or convex tables. The default PALEOS output numbers are unchanged to the printed digits, and non-convex tables achieve the true Chebyshev minimax error.
Parse the printed Chebyshev minimax error from the output and evaluate the reference errors for the Copilot line and the shifted secant directly over the table points, asserting that the minimax error is strictly smaller than both.
… eta_diff Assert that compute_effective_viscosity in local mode is continuous across strain_rate = 0, 1e-300, and 1e-200 when eta_diff exceeds tau_y / (2 * LOCAL_STRAIN_RATE_FLOOR) on both NumPy and JAX, and that its JAX derivative at strain_rate = 0 is zero.
Remove (sr <= 0.0) from is_unyielded in compute_effective_viscosity so that strain_rate = 0 evaluates the continuous sr -> 0+ limit softmin(eta_diff, tau_y_fin / (2 * LOCAL_STRAIN_RATE_FLOOR)) while keeping infinite yield stress unyielded.
Update test_compute_solidus_reference_values_crossings_and_paleos to assert all distinct alternation peaks and parse Chebyshev minimax line coefficients, error bounds, and peak pressures with tolerances.
Group dense grid points within 1e-6 K of the Chebyshev minimax error into distinct peak pressures and format all peak locations in increasing pressure order.
Remove exact string assertions that include peak location lists in solidus reference tests and parse Chebyshev line coefficients, maximum deviations, and peak locations with numerical tolerances.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This updates the standalone solid-rheology kernels in
aragog.rheologyandaragog.rheology_lidahead of wiring them into the solver. No solver path calls these kernels yet (onlyconfig/phases.pyimports the parameter class), so the solver output is unchanged. This is the first step towards the solid-state convection work of FormingWorlds/PROTEUS#915; the wiring follows in a separate PR.tools/verification/compute_solidus_reference_values.py, which now reports every crossing and fits the minimax line by linear programming). The linear form stays a fallback for table-free benchmarks. The 175 kyr fixture takes its temperature and melt fraction from one PALEOS melting curve, andcompute_stagnant_lid_statereadst_sol_0fromparamsonly.ValueErrorwith the field name on the NumPy path (NaNon the JAX path, andNaNin givesNaNout per node). This includes a stated round-off tolerance of 1 Pa for negative surface pressure, a non-negative driving stress in'lid'mode, andNonefor a required input.compute_effective_viscosity, smooth lid weights incompute_stagnant_lid_state, a softplus viscosity ceiling incompute_arrhenius_viscosity, a smooth maximum and a softplus total-flux floor in the boundary weights, and a compact-support melt ramp. Outside each transition band the output equals main.docs/Explanations/solid_state_convection.mdlists the forms and the hard switches that stay.+inf) and returns the viscosity ceiling.Validation of changes
macOS (Apple silicon), Python 3.12, tip
231c395a1c09a2c1f7d1fd54ca9a554b7626eaac,ARAGOG_TEST_EOS_DIRset:pytest -m "unit and not slow" -o "addopts="in one process: 2127 passed, 397 deselected.tests/test_rheology_*.py,tests/test_foley_bercovici_lid_stress.py,tests/test_mlt_calibration.py): 569 passed.test_rheology_kernel_reference_grid_parityholds the outputs of every changed kernel equal to main on a fixed input grid atrtol=1e-12.coverage run: all 241 changed lines inrheology.pyandrheology_lid.pyrun (rheology.py99.4 %,rheology_lid.py100 %); codecov/patch and codecov/project pass on the tip.isfinitefailstest_arrhenius_viscosity_ceiling_when_enthalpy_overflows, and treating a zero strain rate as unyielded again failstest_local_mode_zero_strain_rate_continuity_at_high_viscosity.ruff check,ruff format --check,tools/validate_test_structure.shandtools/agents/check_agents_md.pypass.No coupled run and no validation case here: with no solver caller the kernels cannot change a run. The perturbed starts and the validation cases come with the wiring PR.
Checklist