Commit d71028a7 authored by Gaetano Scandariato's avatar Gaetano Scandariato
Browse files

fix: float32 ensemble warm-up (bounded ESS slices, NaN-free eclipse mask) and... 

fix: float32 ensemble warm-up (bounded ESS slices, NaN-free eclipse mask) and mixed precision for NUTS with narrow priors

Co-Authored-By: default avatarClaude Opus 5.5 <noreply@anthropic.com>
parent c898ae77

CHANGELOG.md

+14 −0
Changes for CHANGELOG.md: 14 added lines, 0 removed lines.
Original line number Diff line number Diff line
@@ -72,6 +72,20 @@ generate it or edit an existing one. 
Requirements: Python ≥ 3.10, JAX, NumPyro ≥ 0.22, BlackJAX ≥ 1.6.2,

jaxoplanet, tinygp; see `docs/dependencies.md`.



### Fixed before the release (2026-10-03)



- **Phase curves in float32:** the eclipse mask divided the planet's emission by its day-side amplitude; when a

  sampler drove the amplitude to its lower bound the product with 1e-6 was a denormal, flushed to zero, and the

  model became NaN. The mask is now computed with a unit-amplitude surface, without division.

- **Ensemble (ESS) warmup from the prior:** the first slices of a poor fit are enormous, and walkers flew to

  unconstrained values of 1e3–1e5, where a float32 sigmoid sits exactly on the prior bound; the warmup ended with

  stuck walkers of infinite energy or NaN positions. The ESS walkers are now confined to |z| ≤ 15 on parameters with

  interval priors (`ensemble_kernels.BoundedESS`; the excluded prior mass is 3e-7 of the interval at each end).

- **NUTS with narrow priors:** `use_float64: auto` now selects the mixed precision for Branch 2 when a free

  parameter's prior is narrower than 1e-4 of its value (`precision.NUTS_RELATIVE_WIDTH_LIMIT`).

- These were found by the paper's NUTS benchmarks: until then the default route had run only on models with GPs,

  which are always in float64.



### Upgrading from 3.x



These notes matter only for configurations or results made with the

docs/config/sampling.md

+7 −1
Changes for docs/config/sampling.md: 7 added lines, 1 removed line.
Original line number Diff line number Diff line
@@ -309,7 +309,13 @@ For a $D$-dimensional posterior: 
  - Gaussian-process detrending (`GP` or `multi_GP`);

  - a transit epoch fitted directly (`tau`) at a large absolute value

    (|T0| > 1000 d, e.g. a BJD near 2 460 000): float32 spaces numbers there

    by 0.25 d, so the epoch itself would be quantised.

    by 0.25 d, so the epoch itself would be quantised;

  - with `production_sampler: nuts` or `auto` (Branch 2), a free parameter

    whose normal or uniform prior is narrower than 1e-4 of its value, typically

    a period fitted directly: its posterior spans only a few hundred float32

    values, and the small leapfrog steps of NUTS round to zero, so the chains

    stop moving (on K2-140, R-hat 75 in float32 and 1.01 with this rule).

    Nested sampling moves by about the posterior width and is not affected.



  Otherwise it uses float32. The folded times passed to the physical models

  are at most half a period, so their float32 precision scales with the

paper_benchmarks/k2140_juliet/runs/nuts_cpu24/PyAETNAouts/k2140_results.json

deleted100644 → 0
+0 −399
Changes for paper_benchmarks/k2140_juliet/runs/nuts_cpu24/PyAETNAouts/k2140_results.json: 0 added lines, 399 removed lines.
Original line number Diff line number Diff line 
{

 "format": "pyaetna-results",

 "format_version": 1,

 "config_name": "k2140",

 "run": {

  "sampler": "ess",

  "device": "cpu",

  "precision": "float32",

  "num_chains": 60,

  "num_draws_per_chain": 2000,

  "num_warmup": 2000,

  "sampling_elapsed_s": 855.8341546058655,

  "random_seed": 1,

  "timestamp_utc": "2026-10-02T14:07:21Z"

 },

 "provenance": {

  "pyaetna_version": "3.8.0",

  "git_commit": "7de93d64bae90e17f88ba7e025f8aebd625f2bf2",

  "git_dirty": false,

  "git_diff_sha256": null,

  "python": "3.14.4",

  "platform": "Linux-7.0.0-34-generic-x86_64-with-glibc2.43",

  "packages": {

   "jax": "0.11.0",

   "jaxlib": "0.11.0",

   "jax-cuda12-plugin": "0.11.0",

   "nvidia-cuda-runtime-cu12": "12.9.79",

   "numpyro": "0.22.0",

   "blackjax": "1.6.2",

   "tinygp": "0.3.1",

   "jaxoplanet": "0.1.0",

   "numpy": "2.5.1"

  },

  "jax_platforms": "cpu",

  "gpu": [

   {

    "name": "NVIDIA GeForce RTX 5070 Ti Laptop GPU",

    "driver": "595.91.07",

    "memory": "12227 MiB"

   }

  ],

  "compilation_cache_entries": 165855

 },

 "parameters": {

  "impact_b": {

   "mean": 0.1675270050764084,

   "std": 0.16985861957073212,

   "ess": 72623.15923079789,

   "rhat": 1.0013524293899536,

   "ess_per_sec": 84.85657979406395

  },

  "kRV_b": {

   "mean": 96.72758483886719,

   "std": 30.648242950439453,

   "ess": 3763.127756899081,

   "rhat": 1.0194811820983887,

   "ess_per_sec": 4.397029186842984

  },

  "log_jitter_CORALIE": {

   "mean": 1.569044589996338,

   "std": 2.2733895778656006,

   "ess": 869.3828799587207,

   "rhat": 1.0897291898727417,

   "ess_per_sec": 1.0158310173529996

  },

  "log_jitter_FIES": {

   "mean": 1.8005396127700806,

   "std": 2.2552735805511475,

   "ess": 886.3653007079738,

   "rhat": 1.0773344039916992,

   "ess_per_sec": 1.0356741384271684

  },

  "log_jitter_HARPS": {

   "mean": 2.8003151416778564,

   "std": 1.6617658138275146,

   "ess": 2217.97573520677,

   "rhat": 1.0309522151947021,

   "ess_per_sec": 2.5915952562423814

  },

  "log_jitter_K2": {

   "mean": -7.852592945098877,

   "std": 0.06052347645163536,

   "ess": 3763.3111208460487,

   "rhat": 1.016507625579834,

   "ess_per_sec": 4.3972434385715236

  },

  "log_jitter_LCOGT": {

   "mean": -10.956531524658203,

   "std": 2.2593882083892822,

   "ess": 13072.704199939073,

   "rhat": 1.0062683820724487,

   "ess_per_sec": 15.274810113133897

  },

  "norm_CORALIE": {

   "mean": 1.2163112163543701,

   "std": 0.025345368310809135,

   "ess": 55488.47777722019,

   "rhat": 1.0018681287765503,

   "ess_per_sec": 64.83554959637492

  },

  "norm_FIES": {

   "mean": 1.1474117040634155,

   "std": 0.031486768275499344,

   "ess": 1380.6143642391173,

   "rhat": 1.0506521463394165,

   "ess_per_sec": 1.6131797928477476

  },

  "norm_HARPS": {

   "mean": 1.237528681755066,

   "std": 0.04433504492044449,

   "ess": 14152.814913938064,

   "rhat": 1.0061594247817993,

   "ess_per_sec": 16.53686621148675

  },

  "norm_K2": {

   "mean": 1.000032901763916,

   "std": 1.442806842533173e-05,

   "ess": 2182.7619284220527,

   "rhat": 1.014286994934082,

   "ess_per_sec": 2.5504496597559525

  },

  "norm_LCOGT": {

   "mean": 0.9997591972351074,

   "std": 0.000783976458478719,

   "ess": 72329.4110153168,

   "rhat": 1.0018879175186157,

   "ess_per_sec": 84.5133494919076

  },

  "period_b": {

   "mean": 6.569289684295654,

   "std": 5.848598448210396e-05,

   "ess": 54058.4011002558,

   "rhat": 1.0025402307510376,

   "ess_per_sec": 63.16457553058413

  },

  "q1_K2": {

   "mean": 0.25799116492271423,

   "std": 0.24276341497898102,

   "ess": 12655.570167694768,

   "rhat": 1.0076960325241089,

   "ess_per_sec": 14.787409569466174

  },

  "q1_LCOGT": {

   "mean": 0.5747063755989075,

   "std": 0.2835370600223541,

   "ess": 2018.1976816333677,

   "rhat": 1.0387287139892578,

   "ess_per_sec": 2.35816445367596

  },

  "q2_K2": {

   "mean": 0.566866397857666,

   "std": 0.2815672755241394,

   "ess": 43517.80618972963,

   "rhat": 1.0004806518554688,

   "ess_per_sec": 50.84841023874625

  },

  "q2_LCOGT": {

   "mean": 0.28667953610420227,

   "std": 0.2777012884616852,

   "ess": 1186.9567249639767,

   "rhat": 1.0646941661834717,

   "ess_per_sec": 1.386900392530609

  },

  "rho_star": {

   "mean": 1.4903908967971802,

   "std": 0.21488428115844727,

   "ess": 116477.99667040345,

   "rhat": 1.000196099281311,

   "ess_per_sec": 136.0987944259419

  },

  "rp_by_rs_b": {

   "mean": 0.11544886976480484,

   "std": 0.0022013464476913214,

   "ess": 18659.8428215672,

   "rhat": 1.0002360343933105,

   "ess_per_sec": 21.80310603537499

  },

  "tau_b": {

   "mean": 0.0007005384541116655,

   "std": 0.00042484357254579663,

   "ess": 4709.36190617044,

   "rhat": 1.0150859355926514,

   "ess_per_sec": 5.502657122090701

  }

 },

 "fixed_params": {

  "norm_mode_K2": "multiplicative",

  "norm_mode_LCOGT": "multiplicative",

  "norm_mode_CORALIE": "additive",

  "norm_mode_FIES": "additive",

  "norm_mode_HARPS": "additive",

  "q1_CORALIE": 0.5,

  "q1_FIES": 0.5,

  "q1_HARPS": 0.5,

  "q2_CORALIE": 0.1,

  "q2_FIES": 0.1,

  "q2_HARPS": 0.1,

  "inc_star": 1.5707963267948966,

  "ecos_b": 0.0,

  "esin_b": 0.0,

  "amp_night_b": 0.0,

  "amp_day_minus_amp_night_b": 0.0,

  "amp_beam_b": 0.0,

  "amp_ellip_b": 0.0,

  "phoffset_b": 0.0,

  "tidallag_b": 0.0,

  "lambda_spin_b": 0.0

 },

 "map": null,

 "config": {

  "data": {

   "file_path": "k2140c.dat",

   "timeZeroPoint": 2457588.284,

   "working_dir": "./"

  },

  "instruments": [

   "K2",

   "LCOGT",

   "CORALIE",

   "FIES",

   "HARPS"

  ],

  "data_type": [

   "LC",

   "LC",

   "RV",

   "RV",

   "RV"

  ],

  "sampling": {

   "chain_energy_threshold": 5.0,

   "checkpoint_interval": 500,

   "internal_warmup_if_nuts": 500,

   "max_tree_depth": 10,

   "num_chains_if_nuts": 4,

   "num_samples_if_ensemble": 2000,

   "num_samples_if_nuts": 2000,

   "num_warmup": 2000,

   "production_sampler": "auto",

   "random_seed": 1,

   "reset": true,

   "target_accept_prob": 0.9,

   "use_float64": false,

   "walkers_per_free_parameter": 3,

   "warmup_sampler": "ess"

  },

  "detrending": {

   "normalization": [

    "per_instrument",

    [

     "normal",

     1.0,

     0.1

    ],

    [

     "normal",

     1.0,

     0.1

    ],

    [

     "normal",

     1.22,

     0.05

    ],

    [

     "normal",

     1.215,

     0.05

    ],

    [

     "normal",

     1.24,

     0.05

    ]

   ],

   "normalization_mode": [

    "per_instrument",

    "multiplicative",

    "multiplicative",

    "additive",

    "additive",

    "additive"

   ]

  },

  "log_jitter": [

   "per_instrument",

   [

    "uniform",

    -13.82,

    -7.6

   ],

   [

    "uniform",

    -13.82,

    -5.3

   ],

   [

    "uniform",

    -2.3,

    4.61

   ],

   [

    "uniform",

    -2.3,

    4.61

   ],

   [

    "uniform",

    -2.3,

    4.61

   ]

  ],

  "planet_b": {

   "amp_beam": 0.0,

   "amp_day_minus_amp_night": 0.0,

   "amp_ellip": 0.0,

   "amp_night": 0.0,

   "ecos": 0.0,

   "esin": 0.0,

   "impact": [

    "global",

    "uniform",

    0.0,

    1.2

   ],

   "kRV": [

    "global",

    "uniform",

    0.0,

    1000.0

   ],

   "period": [

    "global",

    "normal",

    6.5693,

    0.0001

   ],

   "phoffset": 0.0,

   "rp_by_rs": [

    "global",

    "uniform",

    0.01,

    0.5

   ],

   "tau": [

    "global",

    "normal",

    0.0,

    0.001

   ],

   "tidallag": 0.0

  },

  "reports": {

   "discard_extra_warmup": 0,

   "run": false,

   "trace_max_points": 500

  },

  "star": {

   "q1": [

    "per_instrument",

    [

     "uniform",

     0,

     1

    ],

    [

     "uniform",

     0,

     1

    ],

    0.5,

    0.5,

    0.5

   ],

   "q2": [

    "per_instrument",

    [

     "uniform",

     0,

     1

    ],

    [

     "uniform",

     0,

     1

    ],

    0.1,

    0.1,

    0.1

   ],

   "rho_star": [

    "global",

    "uniform",

    0.01,

    12.0

   ]

  }

 }

}
 
 No newline at end of file

paper_benchmarks/k2140_juliet/runs/nuts_cpu24/PyAETNAouts/k2140_run.log

deleted100644 → 0
+0 −628

File deleted.

Preview size limit exceeded, changes collapsed.

paper_benchmarks/k2140_juliet/runs/nuts_cpu24/PyAETNAouts/k2140_samples.npz

deleted100644 → 0
−7.23 MiB

File deleted.

View file