diff --git a/AGENTS.md b/AGENTS.md index 1d21371..b7798e7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -20,6 +20,9 @@ The Python distribution name is `robert-exoplanets`; avoid introducing packaging explicit commands such as `conda run -n robert-exoplanets python ...` and `conda run -n robert-exoplanets python -m pytest ...`; do not rely on the currently activated shell environment. +- Always use Git over SSH for GitHub operations, with remotes such as + `git@github.com:owner/repository.git`. Do not require or use the GitHub CLI + (`gh`); create commits locally and fetch or push them through the SSH remote. - Keep physics-facing APIs explicit and typed so later scientific implementations can replace stubs without changing user-facing examples. - Prefer small, well-tested modules over broad framework code. - Do not add numerical approximations that look like real retrieval physics unless they are clearly labeled as placeholders. diff --git a/configurations/README.md b/configurations/README.md index eba696b..9438f42 100644 --- a/configurations/README.md +++ b/configurations/README.md @@ -32,7 +32,8 @@ same prior in every sampler variant. The WASP-69b cloud-free native-mode and fixed-catalogue Mie cases also ship as an inference benchmark matrix. Each file extends the corresponding science -configuration, so only the inference engine and run/output identity change: +configuration, so only the inference engine and run identity change. Writable +directories are local defaults beside each run's `configuration.yaml`: | Workflow | Cloud-free YAML | Mie-catalogue YAML | | --- | --- | --- | @@ -68,5 +69,6 @@ calibration-sensitivity test. For new projects, start from `TEMPLATE_all_supported_options.yaml`. It groups the editable inputs into system, data, atmosphere, cloud, opacity/RT, priors, -sampler, and housekeeping sections. The active block is a valid cloud-free retrieval; -commented alternatives show the currently supported optional modes and priors. +sampler, plotting, and one top-level paths section. The active block is a valid +cloud-free retrieval; commented alternatives show the supported one-region, +diluted, two-region, chemistry, cloud, and prior modes. diff --git a/configurations/TEMPLATE_all_supported_options.yaml b/configurations/TEMPLATE_all_supported_options.yaml index fc19435..cc6bedf 100644 --- a/configurations/TEMPLATE_all_supported_options.yaml +++ b/configurations/TEMPLATE_all_supported_options.yaml @@ -8,6 +8,19 @@ schema_version: 2 +# ============================================================================= +# PATHS: ALL MACHINE-SPECIFIC LOCATIONS LIVE HERE. +# Relative paths resolve beside this YAML. ROBERT creates ./outputs, ./scratch, +# and ./opacity_cache automatically, so they do not need to be configured. +# Environment variables such as ${ROBERT_DATA_ROOT} are supported. +# ============================================================================= +paths: + project_directory: . + observations_directory: /path/to/observations + fastchem_directory: /path/to/fastchem + k_table_directory: /path/to/opacity + optical_constants_directory: /path/to/optical_constants + # ============================================================================= # RUN IDENTITY # ============================================================================= @@ -62,7 +75,7 @@ observations: # uncertainty_scale: 1.15 # Fixed 15% inflation. # uncertainty_scale_parameter: miri_error_scale # jitter_parameter: miri_jitter - # path is populated from housekeeping.observations_directory. + # path is populated from paths.observations_directory. # ============================================================================= # ATMOSPHERE PROPERTIES @@ -106,7 +119,7 @@ atmosphere: model: fastchem_equilibrium metallicity_parameter: metallicity carbon_to_oxygen_parameter: CtoO - # fastchem_path is populated from housekeeping.fastchem_directory. + # fastchem_path is populated from paths.fastchem_directory. species: - {label: H2O, fastchem_name: H2O1} - {label: CO2, fastchem_name: C1O2} @@ -165,7 +178,7 @@ clouds: # geometric_stddev: 1.0 # quadrature_points: 1 # multiple_scattering_backend: sh4 - # optical_constants_path is populated from housekeeping.optical_constants_directory. + # optical_constants_path is populated from paths.optical_constants_directory. # Shared emission/transmission retrieved refractive-index Mie cloud: # model: mie_direct_nk @@ -177,6 +190,35 @@ clouds: # quadrature_points: 1 # multiple_scattering_backend: sh4 +# ============================================================================= +# PROJECTED-DISK EMISSION +# The top-level atmosphere/clouds define the primary (hot) column. Regional +# overrides inherit every omitted field. Two-region and dilution modes require +# their fraction parameter in `parameters` with bounds inside [0, 1]. +# ============================================================================= +disk_emission: + model: one_region + # Diluted one-column alternative: + # model: diluted_one_region + # dilution_parameter: dayside_dilution + # Two-column alternative ("2tp" is accepted as an alias): + # model: two_region + # hot_fraction_parameter: hot_area_fraction + # hot_region: {} # Inherits the top-level atmosphere/clouds. + # cold_region: + # atmosphere: + # temperature: + # model: isothermal + # parameter_name: cold_temperature + # chemistry: # Omit to share top-level chemistry. + # model: free + # species: [H2O, CO2, CO, CH4, NH3, SO2] + # parameter_names: {H2O: cold_log_H2O, CO2: cold_log_CO2, CO: cold_log_CO, CH4: cold_log_CH4, NH3: cold_log_NH3, SO2: cold_log_SO2} + # background_species: [H2, He] + # background_fractions: [0.8547, 0.1453] + # clouds: + # model: none + # ============================================================================= # OPACITY AND RADIATIVE TRANSFER # ============================================================================= @@ -189,7 +231,7 @@ opacity: use_rebin: false remove_zeros: true g_points: 8 # Used by exomol_cross_section_hdf only. - # path and cache_directory are populated from housekeeping. + # path is populated from paths.k_table_directory; cache defaults locally. radiative_transfer: model: emission # Alternative: transmission @@ -356,6 +398,9 @@ plotting: image_format: png # png, pdf, or svg. dpi: 180 max_posterior_samples: 20000 # Plotting limit; statistics always use all samples. + posterior_predictive_samples: 200 # Forward calls for spectral and T-P envelopes. + posterior_predictive_seed: 0 + corner_max_parameters: 20 # Larger states retain marginals/correlation only. dataset_colors: # Optional dataset-name overrides. f322w2: "#20639b" f444w: "#ef5675" @@ -369,16 +414,3 @@ plotting: runtime: mpi_processes: auto - -# ============================================================================= -# HOUSEKEEPING: ALL MACHINE- AND PROJECT-SPECIFIC PATHS LIVE HERE. -# Relative paths resolve relative to this YAML file. -# ============================================================================= -housekeeping: - observations_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/wasp69b_schlawin2024 - fastchem_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/chemistry/fastchem - k_table_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ktables_exomol - optical_constants_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/optical_constants/exo_skryer - opacity_cache_directory: opacity_cache - output_directory: outputs - scratch_directory: scratch diff --git a/configurations/l98_59b_clr_transmission_multinest.yaml b/configurations/l98_59b_clr_transmission_multinest.yaml index 6525537..96645c4 100644 --- a/configurations/l98_59b_clr_transmission_multinest.yaml +++ b/configurations/l98_59b_clr_transmission_multinest.yaml @@ -1,5 +1,9 @@ schema_version: 2 +paths: + observations_directory: ../data/observations/l98_59b_bello_arufe2025 + k_table_directory: /Users/jaketaylor/Dropbox/ktables_exomol + run: name: l98-59b-clr-transmission-multinest description: >- @@ -25,7 +29,6 @@ bodies: observations: loader: bello_arufe2025_l9859b - path: ../data/observations/l98_59b_bello_arufe2025 datasets: [nrs1, nrs2] verify_checksum: true miri_offset_parameter: nrs2_offset @@ -56,10 +59,8 @@ clouds: opacity: format: exomol_kta - path: /Users/jaketaylor/Dropbox/ktables_exomol resolution: R1000 species: [SO2, H2S, CO2] - cache_directory: ../examples/outputs/l98_59b_clr_transmission/opacity_cache binning: num: 100 use_rebin: false @@ -102,8 +103,6 @@ sampler: seed: 2719 invalid_loglike_floor: -1.0e100 -outputs: - directory: ../examples/outputs/l98_59b_clr_transmission plotting: enabled: true @@ -124,4 +123,3 @@ plotting: runtime: mpi_processes: 3 - scratch_directory: ../examples/outputs/l98_59b_clr_transmission/scratch diff --git a/configurations/picaso_jwst_transmission_retrieval_multinest.yaml b/configurations/picaso_jwst_transmission_retrieval_multinest.yaml index ba6879c..22416b2 100644 --- a/configurations/picaso_jwst_transmission_retrieval_multinest.yaml +++ b/configurations/picaso_jwst_transmission_retrieval_multinest.yaml @@ -1,5 +1,9 @@ schema_version: 2 +paths: + observations_directory: ../examples/outputs/picaso_jwst_transmission_retrieval/picaso_jwst_asimov_observation.npz + k_table_directory: ../opacity_data/exomol_xsec + run: name: picaso-jwst-transmission-retrieval description: >- @@ -23,7 +27,6 @@ bodies: observations: loader: robert_npz - path: ../examples/outputs/picaso_jwst_transmission_retrieval/picaso_jwst_asimov_observation.npz datasets: [picaso_asimov] verify_checksum: false miri_offset_parameter: null @@ -67,10 +70,8 @@ clouds: opacity: format: exomol_cross_section_hdf - path: ../opacity_data/exomol_xsec resolution: R100 species: [H2O, CO, CO2, CH4] - cache_directory: ../examples/outputs/picaso_jwst_transmission_retrieval/opacity_cache binning: num: 100 use_rebin: false @@ -141,8 +142,6 @@ sampler: show_status: true seed: 2718 -outputs: - directory: ../examples/outputs/picaso_jwst_transmission_retrieval plotting: enabled: true @@ -166,4 +165,3 @@ plotting: runtime: mpi_processes: 2 - scratch_directory: ../examples/outputs/picaso_jwst_transmission_retrieval/scratch diff --git a/configurations/synthetic_six_molecule_cloudy_transmission_injection_recovery_multinest.yaml b/configurations/synthetic_six_molecule_cloudy_transmission_injection_recovery_multinest.yaml index b8571bd..c2c24c0 100644 --- a/configurations/synthetic_six_molecule_cloudy_transmission_injection_recovery_multinest.yaml +++ b/configurations/synthetic_six_molecule_cloudy_transmission_injection_recovery_multinest.yaml @@ -1,5 +1,9 @@ schema_version: 2 +paths: + observations_directory: ../examples/outputs/synthetic_six_molecule_cloudy_transmission_injection_recovery/synthetic_transmission_observation.npz + k_table_directory: ../opacity_data/exomol_xsec + run: name: synthetic-six-molecule-cloudy-transmission-injection-recovery description: >- @@ -21,7 +25,6 @@ bodies: observations: loader: robert_npz - path: ../examples/outputs/synthetic_six_molecule_cloudy_transmission_injection_recovery/synthetic_transmission_observation.npz datasets: [synthetic_transit] verify_checksum: false miri_offset_parameter: null @@ -67,10 +70,8 @@ clouds: opacity: format: exomol_cross_section_hdf - path: ../opacity_data/exomol_xsec resolution: R100 species: [H2O, CO, CO2, CH4, NH3, HCN] - cache_directory: ../examples/outputs/synthetic_six_molecule_cloudy_transmission_injection_recovery/opacity_cache binning: num: 100 use_rebin: false @@ -149,8 +150,6 @@ sampler: show_status: true seed: 2717 -outputs: - directory: ../examples/outputs/synthetic_six_molecule_cloudy_transmission_injection_recovery plotting: enabled: true @@ -176,4 +175,3 @@ plotting: runtime: mpi_processes: 2 - scratch_directory: ../examples/outputs/synthetic_six_molecule_cloudy_transmission_injection_recovery/scratch diff --git a/configurations/synthetic_six_molecule_transmission_injection_recovery_multinest.yaml b/configurations/synthetic_six_molecule_transmission_injection_recovery_multinest.yaml index aae194d..f124a7e 100644 --- a/configurations/synthetic_six_molecule_transmission_injection_recovery_multinest.yaml +++ b/configurations/synthetic_six_molecule_transmission_injection_recovery_multinest.yaml @@ -1,5 +1,9 @@ schema_version: 2 +paths: + observations_directory: ../examples/outputs/synthetic_six_molecule_transmission_injection_recovery/synthetic_transmission_observation.npz + k_table_directory: ../opacity_data/exomol_xsec + run: name: synthetic-six-molecule-transmission-injection-recovery description: >- @@ -21,7 +25,6 @@ bodies: observations: loader: robert_npz - path: ../examples/outputs/synthetic_six_molecule_transmission_injection_recovery/synthetic_transmission_observation.npz datasets: [synthetic_transit] verify_checksum: false miri_offset_parameter: null @@ -57,10 +60,8 @@ clouds: opacity: format: exomol_cross_section_hdf - path: ../opacity_data/exomol_xsec resolution: R100 species: [H2O, CO, CO2, CH4, NH3, HCN] - cache_directory: ../examples/outputs/synthetic_six_molecule_transmission_injection_recovery/opacity_cache binning: num: 100 use_rebin: false @@ -123,8 +124,6 @@ sampler: show_status: true seed: 2716 -outputs: - directory: ../examples/outputs/synthetic_six_molecule_transmission_injection_recovery plotting: enabled: true @@ -146,4 +145,3 @@ plotting: runtime: mpi_processes: 2 - scratch_directory: ../examples/outputs/synthetic_six_molecule_transmission_injection_recovery/scratch diff --git a/configurations/synthetic_transmission_injection_recovery_multinest.yaml b/configurations/synthetic_transmission_injection_recovery_multinest.yaml index e1dd1b2..020ee77 100644 --- a/configurations/synthetic_transmission_injection_recovery_multinest.yaml +++ b/configurations/synthetic_transmission_injection_recovery_multinest.yaml @@ -1,5 +1,9 @@ schema_version: 2 +paths: + observations_directory: ../examples/outputs/synthetic_transmission_injection_recovery/synthetic_transmission_observation.npz + k_table_directory: ../opacity_data/exomol_xsec + run: name: synthetic-transmission-injection-recovery description: >- @@ -21,7 +25,6 @@ bodies: observations: loader: robert_npz - path: ../examples/outputs/synthetic_transmission_injection_recovery/synthetic_transmission_observation.npz datasets: [synthetic_transit] verify_checksum: false miri_offset_parameter: null @@ -51,10 +54,8 @@ clouds: opacity: format: exomol_cross_section_hdf - path: ../opacity_data/exomol_xsec resolution: R100 species: [H2O] - cache_directory: ../examples/outputs/synthetic_transmission_injection_recovery/opacity_cache binning: num: 100 use_rebin: false @@ -97,8 +98,6 @@ sampler: show_status: true seed: 2715 -outputs: - directory: ../examples/outputs/synthetic_transmission_injection_recovery plotting: enabled: true @@ -115,4 +114,3 @@ plotting: runtime: mpi_processes: 2 - scratch_directory: ../examples/outputs/synthetic_transmission_injection_recovery/scratch diff --git a/configurations/wasp69b_cloud_free_R1000.yaml b/configurations/wasp69b_cloud_free_R1000.yaml index ec00c7d..349a11b 100644 --- a/configurations/wasp69b_cloud_free_R1000.yaml +++ b/configurations/wasp69b_cloud_free_R1000.yaml @@ -1,7 +1,13 @@ # ROBERT cloud-free emission retrieval: WASP-69b NIRCam + MIRI native modes. -# Paths below are explicit DiRAC paths. Copy this file for each science run. schema_version: 2 +# All input locations live here. Writable outputs, scratch files, and opacity +# cache default to ./outputs, ./scratch, and ./opacity_cache beside this YAML. +paths: + observations_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/wasp69b_schlawin2024 + fastchem_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/chemistry/fastchem + k_table_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ktables_exomol + run: name: wasp69b-cloud-free-native-modes-R1000 description: Cloud-free equilibrium-chemistry emission retrieval @@ -21,7 +27,6 @@ bodies: observations: loader: schlawin2024_wasp69b - path: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/wasp69b_schlawin2024 datasets: [f322w2, f444w, lrs] # Native NIRCam modes plus MIRI/LRS verify_checksum: true miri_offset_parameter: null @@ -36,7 +41,6 @@ atmosphere: internal_temperature_k: 100.0 chemistry: model: fastchem_equilibrium - fastchem_path: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/chemistry/fastchem metallicity_parameter: metallicity carbon_to_oxygen_parameter: CtoO # "label" is ROBERT's molecule name; "fastchem_name" is FastChem's name. @@ -57,10 +61,8 @@ clouds: opacity: format: exomol_kta - path: /lustre/dirac3/scratch/dp448/dc-tayl1/ktables_exomol resolution: R1000 # Change to R15000 to select that directory/file suffix species: [H2O, CO2, CO, CH4, NH3, SO2] - cache_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-cache/wasp69b_native_modes binning: num: 300 use_rebin: false @@ -120,9 +122,5 @@ sampler: show_status: true seed: 2712 -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_cloud_free_R1000_400live - runtime: mpi_processes: auto # Uses SLURM_NTASKS; one process in the terminal - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_cloud_free_R1000 diff --git a/configurations/wasp69b_cloud_free_native_isothermal_R1000.yaml b/configurations/wasp69b_cloud_free_native_isothermal_R1000.yaml index d57d556..20d35de 100644 --- a/configurations/wasp69b_cloud_free_native_isothermal_R1000.yaml +++ b/configurations/wasp69b_cloud_free_native_isothermal_R1000.yaml @@ -14,8 +14,3 @@ parameters: - {name: CtoO, label: C/O, prior: {type: uniform, lower: 0.0, upper: 1.0}, value: 0.5} - {name: log_SO2, label: log10(SO2 VMR), prior: {type: uniform, lower: -10.0, upper: -4.0}, value: -8.0} - {name: temperature, unit: K, prior: {type: uniform, lower: 800.0, upper: 2200.0}, value: 1500.0} - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_cloud_free_native_isothermal_R1000 -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_cloud_free_native_isothermal_R1000 diff --git a/configurations/wasp69b_cloud_free_native_pg14_R1000.yaml b/configurations/wasp69b_cloud_free_native_pg14_R1000.yaml index b32ab6f..2bf9c0c 100644 --- a/configurations/wasp69b_cloud_free_native_pg14_R1000.yaml +++ b/configurations/wasp69b_cloud_free_native_pg14_R1000.yaml @@ -6,8 +6,3 @@ run: plotting: enabled: true - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_cloud_free_native_pg14_R1000 -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_cloud_free_native_pg14_R1000 diff --git a/configurations/wasp69b_cloud_free_native_pg14_R1000_multinest.yaml b/configurations/wasp69b_cloud_free_native_pg14_R1000_multinest.yaml index d9b67a1..60f35c1 100644 --- a/configurations/wasp69b_cloud_free_native_pg14_R1000_multinest.yaml +++ b/configurations/wasp69b_cloud_free_native_pg14_R1000_multinest.yaml @@ -8,8 +8,3 @@ sampler: engine: multinest multinest_max_iterations: 0 resume: resume - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_cloud_free_native_pg14_R1000_multinest -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_cloud_free_native_pg14_R1000_multinest diff --git a/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation.yaml b/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation.yaml index 7f0baff..54296e2 100644 --- a/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation.yaml +++ b/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation.yaml @@ -7,8 +7,3 @@ run: sampler: engine: optimal_estimation oe_max_iterations: 8 - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation diff --git a/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation_to_multinest.yaml b/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation_to_multinest.yaml index 20d0008..3b6adeb 100644 --- a/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation_to_multinest.yaml +++ b/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation_to_multinest.yaml @@ -12,8 +12,3 @@ sampler: require_oe_convergence: true multinest_max_iterations: 0 resume: resume - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_cloud_free_native_pg14_R1000_oe_to_multinest -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_cloud_free_native_pg14_R1000_oe_to_multinest diff --git a/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation_to_ultranest.yaml b/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation_to_ultranest.yaml index 6ada2f4..d920dc4 100644 --- a/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation_to_ultranest.yaml +++ b/configurations/wasp69b_cloud_free_native_pg14_R1000_optimal_estimation_to_ultranest.yaml @@ -10,8 +10,3 @@ sampler: prior_sigma: 4.0 minimum_prior_fraction: 0.05 require_oe_convergence: true - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_cloud_free_native_pg14_R1000_oe_to_ultranest -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_cloud_free_native_pg14_R1000_oe_to_ultranest diff --git a/configurations/wasp69b_cloud_free_nircam_pg14_R1000.yaml b/configurations/wasp69b_cloud_free_nircam_pg14_R1000.yaml index b26fef4..f5a4b8d 100644 --- a/configurations/wasp69b_cloud_free_nircam_pg14_R1000.yaml +++ b/configurations/wasp69b_cloud_free_nircam_pg14_R1000.yaml @@ -6,8 +6,3 @@ run: observations: datasets: [f322w2, f444w] - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_cloud_free_nircam_pg14_R1000 -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_cloud_free_nircam_pg14_R1000 diff --git a/configurations/wasp69b_mie_catalog_layer_by_layer_R1000_optimal_estimation.yaml b/configurations/wasp69b_mie_catalog_layer_by_layer_R1000_optimal_estimation.yaml index e6944a7..bde7c2c 100644 --- a/configurations/wasp69b_mie_catalog_layer_by_layer_R1000_optimal_estimation.yaml +++ b/configurations/wasp69b_mie_catalog_layer_by_layer_R1000_optimal_estimation.yaml @@ -13,7 +13,10 @@ atmosphere: knot_pressure: [100, 79.2016405019, 62.728998582, 49.6823959473, 39.349272631, 31.1652694493, 24.6834046707, 19.5496614309, 15.4836525659, 12.2633068418, 9.71274019847, 7.69264957488, 6.09270466137, 4.82552204274, 3.82189262063, 3.02700165376, 2.3974349678, 1.89880782447, 1.50388694696, 1.19110313328, 0.94337322163, 0.747167067587, 0.591768574819, 0.468690419231, 0.371210500907, 0.294004806433, 0.23285662985, 0.184426270859, 0.146068632036, 0.115688752832, 0.0916273901189, 0.0725703961232, 0.0574769442484, 0.0455226827551, 0.0360547115425, 0.0285559230199, 0.0226167594922, 0.0179128445462, 0.0141872667412, 0.0112365480014, 0.00889953035289, 0.00704857403645, 0.00558258626886, 0.00442149990737, 0.00350190046143, 0.0027735626142, 0.00219670709079, 0.00173982805293, 0.00137797235983, 0.00109137671465, 0.00086438826206, 0.000684609683857, 0.00054222210065, 0.000429448798879, 0.000340130493828, 0.000269388930959, 0.00021336045265, 0.000168984978681, 0.000133838875317, 0.000106002584881, 8.39557862e-05, 6.64943599667e-05, 5.26646239348e-05, 4.17112461206e-05, 3.30359912013e-05, 2.61650469875e-05, 2.07231464522e-05, 1.64130719538e-05, 1.29994222441e-05, 1.02957556731e-05, 8.15440739519e-06, 6.4584244302e-06, 5.11517809929e-06, 4.05130496924e-06, 3.20869999737e-06, 2.5413430367e-06, 2.01278537585e-06, 1.59415903746e-06, 1.26260010987e-06, 1e-06] parameter_names: [T_00, T_01, T_02, T_03, T_04, T_05, T_06, T_07, T_08, T_09, T_10, T_11, T_12, T_13, T_14, T_15, T_16, T_17, T_18, T_19, T_20, T_21, T_22, T_23, T_24, T_25, T_26, T_27, T_28, T_29, T_30, T_31, T_32, T_33, T_34, T_35, T_36, T_37, T_38, T_39, T_40, T_41, T_42, T_43, T_44, T_45, T_46, T_47, T_48, T_49, T_50, T_51, T_52, T_53, T_54, T_55, T_56, T_57, T_58, T_59, T_60, T_61, T_62, T_63, T_64, T_65, T_66, T_67, T_68, T_69, T_70, T_71, T_72, T_73, T_74, T_75, T_76, T_77, T_78, T_79] pressure_unit: bar - extrapolation: raise + # The knots are layer centres, while radiative transfer also evaluates the + # two bounding pressure edges. Hold the endpoint temperatures constant + # across those outer half-layers. + extrapolation: clip parameters: # Parameters shared with the completed PG14 MultiNest run are initialized @@ -121,9 +124,6 @@ sampler: oe_temperature_prior_sigma_k: 250.0 oe_temperature_correlation_length_dex: 1.5 -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_mie_catalog_layer_by_layer_R1000_optimal_estimation runtime: mpi_processes: 1 - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_mie_catalog_layer_by_layer_R1000_optimal_estimation diff --git a/configurations/wasp69b_mie_catalog_pg14_R1000.yaml b/configurations/wasp69b_mie_catalog_pg14_R1000.yaml index 09de690..70f3dc7 100644 --- a/configurations/wasp69b_mie_catalog_pg14_R1000.yaml +++ b/configurations/wasp69b_mie_catalog_pg14_R1000.yaml @@ -1,5 +1,8 @@ extends: wasp69b_cloud_free_R1000.yaml +paths: + optical_constants_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/optical_constants/exo_skryer + run: name: wasp69b-mie-catalog-pg14-R1000 description: One-region MgSiO3 Mie cloud with fixed catalogue optical constants and PG14 temperature profile. @@ -13,7 +16,6 @@ observations: clouds: model: mie_catalog - optical_constants_path: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/optical_constants/exo_skryer material: MgSiO3 particle_density_kg_m3: 3200.0 geometric_stddev: 1.0 @@ -41,8 +43,3 @@ sampler: resume: resume show_status: false seed: 2712 - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_mie_catalog_pg14_R1000 -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_mie_catalog_pg14_R1000 diff --git a/configurations/wasp69b_mie_catalog_pg14_R1000_multinest.yaml b/configurations/wasp69b_mie_catalog_pg14_R1000_multinest.yaml index 06724c7..019947e 100644 --- a/configurations/wasp69b_mie_catalog_pg14_R1000_multinest.yaml +++ b/configurations/wasp69b_mie_catalog_pg14_R1000_multinest.yaml @@ -8,8 +8,3 @@ sampler: engine: multinest multinest_max_iterations: 0 resume: resume - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_mie_catalog_pg14_R1000_multinest -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_mie_catalog_pg14_R1000_multinest diff --git a/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation.yaml b/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation.yaml index f7c55dc..26570e8 100644 --- a/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation.yaml +++ b/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation.yaml @@ -7,8 +7,3 @@ run: sampler: engine: optimal_estimation oe_max_iterations: 8 - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_mie_catalog_pg14_R1000_optimal_estimation -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_mie_catalog_pg14_R1000_optimal_estimation diff --git a/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation_to_multinest.yaml b/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation_to_multinest.yaml index 00032fe..807ed4f 100644 --- a/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation_to_multinest.yaml +++ b/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation_to_multinest.yaml @@ -12,8 +12,3 @@ sampler: require_oe_convergence: true multinest_max_iterations: 0 resume: resume - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_mie_catalog_pg14_R1000_oe_to_multinest -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_mie_catalog_pg14_R1000_oe_to_multinest diff --git a/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation_to_ultranest.yaml b/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation_to_ultranest.yaml index 052c8a6..39f865f 100644 --- a/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation_to_ultranest.yaml +++ b/configurations/wasp69b_mie_catalog_pg14_R1000_optimal_estimation_to_ultranest.yaml @@ -10,8 +10,3 @@ sampler: prior_sigma: 4.0 minimum_prior_fraction: 0.05 require_oe_convergence: true - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_mie_catalog_pg14_R1000_oe_to_ultranest -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_mie_catalog_pg14_R1000_oe_to_ultranest diff --git a/configurations/wasp69b_mie_direct_nk_pg14_R1000.yaml b/configurations/wasp69b_mie_direct_nk_pg14_R1000.yaml index d308255..1992102 100644 --- a/configurations/wasp69b_mie_direct_nk_pg14_R1000.yaml +++ b/configurations/wasp69b_mie_direct_nk_pg14_R1000.yaml @@ -32,8 +32,3 @@ parameters: - {name: cloud_logk_3, prior: {type: uniform, lower: -6.0, upper: 0.0}} - {name: cloud_logk_4, prior: {type: uniform, lower: -6.0, upper: 0.0}} - {name: cloud_logk_5, prior: {type: uniform, lower: -6.0, upper: 0.0}} - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp69b_mie_direct_nk_pg14_R1000 -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp69b_mie_direct_nk_pg14_R1000 diff --git a/configurations/wasp80b_cloud_free_native_isothermal_R1000.yaml b/configurations/wasp80b_cloud_free_native_isothermal_R1000.yaml index 3a2d264..22199fe 100644 --- a/configurations/wasp80b_cloud_free_native_isothermal_R1000.yaml +++ b/configurations/wasp80b_cloud_free_native_isothermal_R1000.yaml @@ -13,8 +13,3 @@ parameters: - {name: metallicity, label: log10(Z/Z_sun), unit: dex, prior: {type: uniform, lower: -1.0, upper: 2.0}, value: 0.5} - {name: CtoO, label: C/O, prior: {type: uniform, lower: 0.0, upper: 1.0}, value: 0.5} - {name: temperature, unit: K, prior: {type: uniform, lower: 700.0, upper: 2000.0}, value: 1200.0} - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp80b_cloud_free_native_isothermal_R1000 -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp80b_cloud_free_native_isothermal_R1000 diff --git a/configurations/wasp80b_cloud_free_native_pg14_R1000.yaml b/configurations/wasp80b_cloud_free_native_pg14_R1000.yaml index 41005f1..70547db 100644 --- a/configurations/wasp80b_cloud_free_native_pg14_R1000.yaml +++ b/configurations/wasp80b_cloud_free_native_pg14_R1000.yaml @@ -1,5 +1,8 @@ extends: wasp69b_cloud_free_R1000.yaml +paths: + observations_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/wasp80b_wiser2025 + run: name: wasp80b-cloud-free-native-pg14-R1000 description: Cloud-free PG14 equilibrium chemistry; independent WASP-80b NIRCam and MIRI modes. @@ -19,7 +22,6 @@ bodies: observations: loader: wiser2025_wasp80b - path: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/wasp80b_wiser2025 datasets: [f322w2, f444w, lrs] atmosphere: @@ -37,7 +39,6 @@ atmosphere: opacity: species: [H2O, CO2, CO, CH4, NH3, HCN] - cache_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-cache/wasp80b_native_modes parameters: - {name: metallicity, label: log10(Z/Z_sun), unit: dex, prior: {type: uniform, lower: -1.0, upper: 2.0}, value: 0.5} @@ -47,8 +48,3 @@ parameters: - {name: gamma2, prior: {type: log_uniform, lower: 1.0e-4, upper: 100.0}, value: 0.1} - {name: T_irr, unit: K, prior: {type: uniform, lower: 800.0, upper: 2200.0}, value: 1500.0} - {name: alpha, prior: {type: uniform, lower: 0.0, upper: 1.0}, value: 0.5} - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp80b_cloud_free_native_pg14_R1000 -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp80b_cloud_free_native_pg14_R1000 diff --git a/configurations/wasp80b_cloud_free_nircam_pg14_R1000.yaml b/configurations/wasp80b_cloud_free_nircam_pg14_R1000.yaml index e7873f2..0a8f55d 100644 --- a/configurations/wasp80b_cloud_free_nircam_pg14_R1000.yaml +++ b/configurations/wasp80b_cloud_free_nircam_pg14_R1000.yaml @@ -6,8 +6,3 @@ run: observations: datasets: [f322w2, f444w] - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp80b_cloud_free_nircam_pg14_R1000 -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp80b_cloud_free_nircam_pg14_R1000 diff --git a/configurations/wasp80b_mie_catalog_pg14_R1000.yaml b/configurations/wasp80b_mie_catalog_pg14_R1000.yaml index 934d8fd..6fd76e4 100644 --- a/configurations/wasp80b_mie_catalog_pg14_R1000.yaml +++ b/configurations/wasp80b_mie_catalog_pg14_R1000.yaml @@ -1,5 +1,8 @@ extends: wasp69b_mie_catalog_pg14_R1000.yaml +paths: + observations_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/wasp80b_wiser2025 + run: name: wasp80b-mie-catalog-pg14-R1000 description: WASP-80b one-region MgSiO3 Mie cloud with fixed catalogue optical constants and PG14 temperature profile. @@ -19,7 +22,6 @@ bodies: observations: loader: wiser2025_wasp80b - path: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-code/data/wasp80b_wiser2025 datasets: [f322w2, f444w, lrs] miri_offset_parameter: null @@ -38,7 +40,6 @@ atmosphere: opacity: species: [H2O, CO2, CO, CH4, NH3, HCN] - cache_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-cache/wasp80b_native_modes parameters: - {name: metallicity, label: '[M/H]', unit: dex, prior: {type: uniform, lower: 0.0, upper: 2.0}} @@ -52,8 +53,3 @@ parameters: - {name: log_cloud_radius_micron, prior: {type: uniform, lower: -3.0, upper: 1.0}} - {name: log_cloud_top_pressure_bar, prior: {type: uniform, lower: -6.0, upper: -2.0}} - {name: log_cloud_base_pressure_bar, prior: {type: uniform, lower: -1.0, upper: 2.0}} - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp80b_mie_catalog_pg14_R1000 -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp80b_mie_catalog_pg14_R1000 diff --git a/configurations/wasp80b_mie_direct_nk_pg14_R1000.yaml b/configurations/wasp80b_mie_direct_nk_pg14_R1000.yaml index 865d4fa..1f8f6d0 100644 --- a/configurations/wasp80b_mie_direct_nk_pg14_R1000.yaml +++ b/configurations/wasp80b_mie_direct_nk_pg14_R1000.yaml @@ -31,8 +31,3 @@ parameters: - {name: cloud_logk_3, prior: {type: uniform, lower: -6.0, upper: 0.0}} - {name: cloud_logk_4, prior: {type: uniform, lower: -6.0, upper: 0.0}} - {name: cloud_logk_5, prior: {type: uniform, lower: -6.0, upper: 0.0}} - -outputs: - directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-runs/wasp80b_mie_direct_nk_pg14_R1000 -runtime: - scratch_directory: /lustre/dirac3/scratch/dp448/dc-tayl1/ROBERT-work/wasp80b_mie_direct_nk_pg14_R1000 diff --git a/docs/configuration.md b/docs/configuration.md index 11f0ece..c2dc723 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -54,6 +54,9 @@ sbatch submit.sbatch The major sections are intentionally explicit: +- `paths`: the one top-level block for machine-specific input locations; + writable `outputs/`, `scratch/`, and `opacity_cache/` directories default + beside the YAML file and are created by ROBERT; - `bodies`: planetary and stellar parameters; - `observations`: published-data loader or self-describing ROBERT NPZ, data path, and selected instrument datasets; @@ -62,6 +65,8 @@ The major sections are intentionally explicit: corresponding to each molecule; - `clouds`: cloud-free or configured emission-cloud treatment; configured transmission currently requires `none`; +- `disk_emission`: one region, a diluted region, or two independently + configurable regions with selective atmosphere/chemistry/cloud overrides; - `opacity`: KTA or ExoMolOP cross-section root, resolution label, selected molecules, target-bin preparation, and an external cache directory; - `radiative_transfer`: emission geometry/backend or transmission reference @@ -73,9 +78,64 @@ The major sections are intentionally explicit: - `plotting`: optional automatic retrieval/forward post-processing, image format, Matplotlib style, sampling display limit, colours, labels, and the optional PSIS leave-one-out diagnostic; -- `outputs`: a project directory outside the source checkout; and -- `runtime`: `auto` uses `SLURM_NTASKS` under Slurm and one process otherwise, - while `scratch_directory` controls runtime caches. +- `runtime`: `auto` uses `SLURM_NTASKS` under Slurm and one process otherwise. + +## Self-contained paths + +Normal YAML files do not need `outputs.directory`, `opacity.cache_directory`, +or `runtime.scratch_directory`. Given `/project/configuration.yaml`, ROBERT +automatically resolves and creates `/project/outputs`, `/project/opacity_cache`, +and `/project/scratch`. External inputs are collected once at the top: + +```yaml +schema_version: 2 +paths: + project_directory: . + observations_directory: ${ROBERT_DATA_ROOT}/wasp69b + fastchem_directory: ${ROBERT_DATA_ROOT}/fastchem + k_table_directory: ${ROBERT_OPACITY_ROOT} +``` + +Relative paths resolve from the YAML directory. Undefined environment +variables are rejected. Explicit legacy writable paths remain readable for +old configurations, but new configurations should rely on the local defaults. + +## Projected-disk regions + +The top-level `atmosphere` and `clouds` blocks define the primary/hot column. +A diluted model only adds its projected emitting fraction: + +```yaml +disk_emission: + model: diluted_one_region + dilution_parameter: dayside_dilution +``` + +A two-region model adds hot and cold fluxes after their independent radiative +transfer calculations. Each override inherits omitted components: + +```yaml +disk_emission: + model: two_region + hot_fraction_parameter: hot_area_fraction + hot_region: {} + cold_region: + atmosphere: + temperature: + model: isothermal + parameter_name: cold_temperature + chemistry: + model: free + species: [H2O, CO2, CO, CH4, NH3, SO2] + parameter_names: {H2O: cold_log_H2O, CO2: cold_log_CO2} + background_species: [H2, He] + clouds: + model: none +``` + +All parameters named by either region, plus the area/dilution fraction, must +appear in `parameters`. Fraction priors must remain within `[0, 1]` physically. +Two-region and diluted models are emission-only. ## Stellar spectra @@ -238,5 +298,5 @@ cd /scratch/dp448/dc-tayl1/my_project/ sbatch submit.sbatch ``` -After any failed multi-writer launch, select a new `outputs.directory` in YAML. +After any failed multi-writer launch, use a new self-contained run directory. An HDF5 checkpoint touched by independent writers must not be resumed. diff --git a/docs/postprocessing.md b/docs/postprocessing.md index 52a8ea6..fb895a0 100644 --- a/docs/postprocessing.md +++ b/docs/postprocessing.md @@ -20,6 +20,9 @@ plotting: image_format: png dpi: 180 max_posterior_samples: 20000 + posterior_predictive_samples: 200 + posterior_predictive_seed: 0 + corner_max_parameters: 20 dataset_colors: f322w2: "#20639b" f444w: "#ef5675" @@ -36,8 +39,10 @@ plotting: After successful inference, MPI rank 0 post-processes every completed phase. A hybrid run therefore receives separate OE and nested-sampling plot folders. -The numerical posterior summary always uses every sample; -`max_posterior_samples` limits only the samples rendered in histograms. +The numerical posterior summary always uses every sample. +`max_posterior_samples` limits rendered marginal/corner samples, while +`posterior_predictive_samples` controls the forward evaluations used for the +68% spectral and temperature-profile credible envelopes. The benchmark configurations `wasp69b_cloud_free_native_pg14_R1000.yaml` and @@ -57,10 +62,19 @@ ROBERT discovers each completed phase and writes beneath `outputs/plots/`: - `fit_statistics.json`; - `posterior_summary.json`; - `fit_spectrum_residuals.png`; +- `temperature_profiles.png` when the forward model exposes a T-P profile; +- `posterior_corner.png` for nested posteriors up to `corner_max_parameters`; - `posterior_marginals.png` or `optimal_estimation_parameters.png`; - `parameter_correlation.png`; and - `plot_manifest.json`. +The spectrum plot shows the observation-grid posterior model as open squares +and a 68% posterior envelope. For ExoMol correlated-k configurations, ROBERT +also evaluates and plots the median and 68% envelope on the exact native +opacity grid. Formats for which a native-grid correlation is not defined are +reported honestly and retain the observation-grid model rather than drawing +an interpolated curve and calling it native. + When `plotting.leave_one_out.enabled` is true for a nested-sampling result, ROBERT additionally writes `leave_one_out.json`, `leave_one_out_arrays.npz`, and `leave_one_out.png`. This PSIS-LOO diagnostic diff --git a/docs/theory/retrieval_workflows.md b/docs/theory/retrieval_workflows.md index 74c6492..68cdb7b 100644 --- a/docs/theory/retrieval_workflows.md +++ b/docs/theory/retrieval_workflows.md @@ -77,6 +77,22 @@ python run_oe_from_nested.py \ --oe-config /path/to/layer-oe-run/configuration.yaml ``` +Create the OE run directory from the supplied source YAML so that its +`configuration.yaml` is a complete, flattened configuration: + +```bash +python scripts/create_run_directory.py \ + --project-dir /path/to/project \ + --config configurations/wasp69b_mie_catalog_layer_by_layer_R1000_optimal_estimation.yaml +``` + +Do not put `extends: configuration.yaml` in the isolated `configuration.yaml`; +that makes the file its own parent. For a like-for-like PG14 refinement, the +standard MultiNest and optimal-estimation configurations can be supplied to +the deferred runner. Their common temperature parameters and posterior +covariance transfer directly, without the layer-temperature conversion or a +correlated spline prior. + Parameters shared by the two models inherit the MultiNest best fit and posterior covariance. The PG14 best-fit temperature profile is evaluated at all 80 pressure-layer centres to initialize the new OE temperature state. The @@ -84,7 +100,9 @@ temperature prior covariance is `S_ij = sigma_T^2 exp(-|log10(P_i/P_j)| / L)`, with the supplied configuration recording `sigma_T=250 K` and `L=1.5 dex`. This smooths departures from the MultiNest profile instead of treating adjacent layer temperatures as -independent. The exact prior state and covariance are saved in +independent. Because the retrieved knots are the layer centres, the two outer +pressure edges use clipped endpoint temperatures across their half-layers. +The exact prior state and covariance are saved in `multinest_to_oe_prior.npz`. MultiNest outputs remain unchanged. By default the handoff refuses a result that has not reported convergence; `--allow-unconverged` is available only for explicit diagnostic work. diff --git a/postprocess_retrieval.py b/postprocess_retrieval.py index 43f5ff2..052851b 100644 --- a/postprocess_retrieval.py +++ b/postprocess_retrieval.py @@ -6,7 +6,11 @@ import argparse from pathlib import Path -from robert_exoplanets.io.configured_tasks import build_problem, load_observations +from robert_exoplanets.io.configured_tasks import ( + build_native_emission_model, + build_problem, + load_observations, +) from robert_exoplanets.io.task_config import initialize_task_directories, load_task_config from robert_exoplanets.postprocessing import ( discover_retrieval_result_directories, @@ -50,7 +54,9 @@ def main() -> None: config = load_task_config(args.config) initialize_task_directories(config) - problem = build_problem(config, load_observations(config)) + observations = load_observations(config) + problem = build_problem(config, observations) + native_spectrum_model = build_native_emission_model(config, observations) result_dirs = tuple(args.result_dir or ()) or discover_retrieval_result_directories( config.outputs.directory ) @@ -79,6 +85,12 @@ def main() -> None: max_posterior_samples=( args.max_samples or config.plotting.max_posterior_samples ), + posterior_predictive_samples=( + config.plotting.posterior_predictive_samples + ), + posterior_predictive_seed=config.plotting.posterior_predictive_seed, + corner_max_parameters=config.plotting.corner_max_parameters, + native_spectrum_model=native_spectrum_model, ) print( f"{phase}: reduced chi-squared={diagnostics['reduced_chi_squared']}; " diff --git a/run_oe_from_nested.py b/run_oe_from_nested.py index 4bacc07..1eab5fa 100644 --- a/run_oe_from_nested.py +++ b/run_oe_from_nested.py @@ -60,7 +60,7 @@ def main() -> None: initialize_task_directories(oe_config) oe_problem = build_problem(oe_config, load_observations(oe_config)) - temperature_overrides = layer_temperature_overrides( + temperature_overrides = temperature_state_overrides( nested_config, nested_result.best_fit_parameters, oe_problem, @@ -71,10 +71,12 @@ def main() -> None: state_overrides=temperature_overrides, covariance_floor_fraction=args.covariance_floor_fraction, ) - prior_covariance, temperature_pressure = configured_temperature_prior_covariance( - oe_config, - oe_problem, - prior_covariance, + prior_covariance, temperature_parameter_names, temperature_pressure = ( + temperature_prior_covariance( + oe_config, + oe_problem, + prior_covariance, + ) ) prior_path = oe_config.outputs.directory / "multinest_to_oe_prior.npz" np.savez( @@ -82,12 +84,13 @@ def main() -> None: parameter_names=np.asarray(oe_problem.parameter_names), prior_state=prior_state, prior_covariance=prior_covariance, - temperature_parameter_names=np.asarray( - _spline_temperature_profile(oe_problem).parameter_names - ), + temperature_parameter_names=np.asarray(temperature_parameter_names), temperature_pressure_bar=temperature_pressure, ) - smoke = smoke_evaluation(oe_problem) + # Validate the state OE will actually use. A generic prior midpoint can be + # outside the valid model domain even when the transferred nested state is + # physically valid. + smoke = smoke_evaluation(oe_problem, prior_state) write_config_snapshot(oe_config, args.oe_config) (oe_config.outputs.directory / "smoke_evaluation.json").write_text( json.dumps(smoke, indent=2), encoding="utf-8" @@ -183,12 +186,58 @@ def layer_temperature_overrides( } +def temperature_state_overrides( + nested_config: TaskConfig, + nested_best_fit: Mapping[str, float], + oe_problem: object, +) -> dict[str, float]: + """Prepare temperature overrides only when OE changes to a spline state.""" + + target_profile = _atmosphere_builder(oe_problem).temperature_profile + if isinstance(target_profile, SplineTemperatureProfile): + return layer_temperature_overrides( + nested_config, + nested_best_fit, + oe_problem, + ) + source_profile = _temperature_profile(nested_config) + if ( + type(source_profile) is not type(target_profile) + or source_profile.required_parameters() != target_profile.required_parameters() + ): + raise RobertConfigError( + "the nested and OE temperature profiles must use the same " + "parameterization unless OE uses a retrieved spline profile" + ) + # Same-name parameters transfer directly in nested_posterior_oe_prior. + return {} + + +def temperature_prior_covariance( + oe_config: TaskConfig, + oe_problem: object, + prior_covariance: np.ndarray, +) -> tuple[np.ndarray, tuple[str, ...], np.ndarray]: + """Apply optional spline smoothing and describe the temperature prior block.""" + + profile = _atmosphere_builder(oe_problem).temperature_profile + parameter_names = tuple(profile.required_parameters()) + if oe_config.sampler.oe_temperature_prior_sigma_k is None: + return prior_covariance, parameter_names, np.asarray([], dtype=float) + covariance, pressure = configured_temperature_prior_covariance( + oe_config, + oe_problem, + prior_covariance, + ) + return covariance, parameter_names, pressure + + def _atmosphere_builder(oe_problem: object): forward_model = getattr(oe_problem, "forward_model", None) builder = getattr(forward_model, "atmosphere_builder", None) if builder is None: raise RobertConfigError( - "layer-temperature handoff requires an emission problem with a shared " + "deferred OE handoff requires an emission problem with a shared " "atmosphere builder" ) return builder diff --git a/scripts/create_run_directory.py b/scripts/create_run_directory.py index 7d3d732..1189f62 100644 --- a/scripts/create_run_directory.py +++ b/scripts/create_run_directory.py @@ -133,16 +133,48 @@ def create_run_directory(*, project_dir: Path, source_config: Path) -> Path: ): shutil.copy2(ROOT / filename, run_directory / filename) - generated = config.model_dump(mode="json") - generated["outputs"]["directory"] = str(run_directory / "outputs") - generated["opacity"]["cache_directory"] = str(run_directory / "opacity_cache") - generated["runtime"]["scratch_directory"] = str(run_directory / "scratch") - if generated.get("housekeeping") is not None: - generated["housekeeping"]["output_directory"] = str(run_directory / "outputs") - generated["housekeeping"]["opacity_cache_directory"] = str( - run_directory / "opacity_cache" + generated = config.model_dump(mode="json", exclude_none=True) + configured_paths = config.paths or config.housekeeping + paths = ( + {} + if configured_paths is None + else configured_paths.model_dump(mode="json", exclude_none=True) + ) + paths.update( + { + "project_directory": ".", + "observations_directory": str(config.observations.path), + "k_table_directory": str(config.opacity.path), + } + ) + if config.atmosphere.chemistry.model == "fastchem_equilibrium": + paths["fastchem_directory"] = str( + config.atmosphere.chemistry.fastchem_path ) - generated["housekeeping"]["scratch_directory"] = str(run_directory / "scratch") + generated["atmosphere"]["chemistry"].pop("fastchem_path", None) + if config.clouds.model == "mie_catalog": + paths["optical_constants_directory"] = str( + config.clouds.optical_constants_path + ) + generated["clouds"].pop("optical_constants_path", None) + for key in ( + "opacity_cache_directory", + "output_directory", + "scratch_directory", + ): + paths.pop(key, None) + generated["observations"].pop("path", None) + generated["opacity"].pop("path", None) + generated["opacity"].pop("cache_directory", None) + generated.pop("outputs", None) + generated["runtime"].pop("scratch_directory", None) + generated.pop("housekeeping", None) + generated.pop("paths", None) + generated = { + "schema_version": generated.pop("schema_version"), + "paths": paths, + **generated, + } execution_config = run_directory / "configuration.yaml" execution_config.write_text( yaml.safe_dump(generated, sort_keys=False), encoding="utf-8" diff --git a/src/robert_exoplanets/__init__.py b/src/robert_exoplanets/__init__.py index 5a4d163..dc733f9 100644 --- a/src/robert_exoplanets/__init__.py +++ b/src/robert_exoplanets/__init__.py @@ -56,6 +56,8 @@ ExoMolOpacitySamplingSource, GreyScatteringCloudConfig, MultiDatasetEmissionForwardModel, + MultiDatasetDilutedEmissionModel, + MultiDatasetTwoRegionEmissionModel, NativeSpectrumMultiDatasetForwardModel, NativeSpectrumMultiDatasetPrediction, ParameterizedEmissionFactoryConfig, @@ -339,6 +341,8 @@ "MadhusudhanSeager2009TemperatureProfile", "MeanMolecularWeightModel", "MultiDatasetEmissionForwardModel", + "MultiDatasetDilutedEmissionModel", + "MultiDatasetTwoRegionEmissionModel", "MultiDatasetGaussianLikelihood", "NativeSpectrumMultiDatasetForwardModel", "NativeSpectrumMultiDatasetPrediction", diff --git a/src/robert_exoplanets/forward/__init__.py b/src/robert_exoplanets/forward/__init__.py index eb20b03..1e74a3f 100644 --- a/src/robert_exoplanets/forward/__init__.py +++ b/src/robert_exoplanets/forward/__init__.py @@ -44,6 +44,8 @@ from .inhomogeneous import ( DilutedEmissionModel, DiskEmissionModelConfig, + MultiDatasetDilutedEmissionModel, + MultiDatasetTwoRegionEmissionModel, TwoRegionEmissionModel, build_disk_emission_model, ) @@ -67,6 +69,8 @@ "ExoMolOpacitySamplingSource", "GreyScatteringCloudConfig", "MultiDatasetEmissionForwardModel", + "MultiDatasetDilutedEmissionModel", + "MultiDatasetTwoRegionEmissionModel", "NativeSpectrumMultiDatasetForwardModel", "NativeSpectrumMultiDatasetPrediction", "ParameterizedEmissionFactoryConfig", diff --git a/src/robert_exoplanets/forward/inhomogeneous.py b/src/robert_exoplanets/forward/inhomogeneous.py index 661a29b..9fcc92a 100644 --- a/src/robert_exoplanets/forward/inhomogeneous.py +++ b/src/robert_exoplanets/forward/inhomogeneous.py @@ -10,6 +10,9 @@ from robert_exoplanets.core import RobertValidationError, Spectrum EmissionEvaluator = Callable[[Mapping[str, float]], Spectrum] +MultiDatasetEmissionEvaluator = Callable[ + [Mapping[str, float]], Mapping[str, Spectrum] +] def _parameter(parameters: Mapping[str, float], name: str) -> float: @@ -37,6 +40,49 @@ def _validate_compatible(left: Spectrum, right: Spectrum) -> None: raise RobertValidationError("regional spectra must share units and observable") +def _mixed_spectrum( + hot: Spectrum, + cold: Spectrum, + hot_fraction: float, + *, + hot_fraction_parameter: str, +) -> Spectrum: + _validate_compatible(hot, cold) + return Spectrum( + spectral_grid=hot.spectral_grid, + values=hot_fraction * hot.values + (1.0 - hot_fraction) * cold.values, + unit=hot.unit, + observable=hot.observable, + metadata={ + "disk_model": "two_region_areal_mixture", + "hot_fraction_parameter": hot_fraction_parameter, + "hot_area_fraction": f"{hot_fraction:.17g}", + "source_equation": "Schlawin_et_al_2024_equation_1", + }, + ) + + +def _diluted_spectrum( + spectrum: Spectrum, + dilution: float, + *, + dilution_parameter: str, +) -> Spectrum: + return Spectrum( + spectral_grid=spectrum.spectral_grid, + values=dilution * spectrum.values, + unit=spectrum.unit, + observable=spectrum.observable, + metadata={ + **dict(spectrum.metadata), + "disk_model": "diluted_single_region", + "dilution_parameter": dilution_parameter, + "dayside_dilution": f"{dilution:.17g}", + "dilution_definition": "fractional_projected_dayside_emitting_area", + }, + ) + + @dataclass(frozen=True) class TwoRegionEmissionModel: """Area-weight two independent dayside emission columns. @@ -62,18 +108,11 @@ def __call__(self, parameters: Mapping[str, float]) -> Spectrum: ) hot = self.hot_model(parameters) cold = self.cold_model(parameters) - _validate_compatible(hot, cold) - return Spectrum( - spectral_grid=hot.spectral_grid, - values=hot_fraction * hot.values + (1.0 - hot_fraction) * cold.values, - unit=hot.unit, - observable=hot.observable, - metadata={ - "disk_model": "two_region_areal_mixture", - "hot_fraction_parameter": self.hot_fraction_parameter, - "hot_area_fraction": f"{hot_fraction:.17g}", - "source_equation": "Schlawin_et_al_2024_equation_1", - }, + return _mixed_spectrum( + hot, + cold, + hot_fraction, + hot_fraction_parameter=self.hot_fraction_parameter, ) @@ -99,21 +138,76 @@ def __call__(self, parameters: Mapping[str, float]) -> Spectrum: self.dilution_parameter, ) spectrum = self.emission_model(parameters) - return Spectrum( - spectral_grid=spectrum.spectral_grid, - values=dilution * spectrum.values, - unit=spectrum.unit, - observable=spectrum.observable, - metadata={ - **dict(spectrum.metadata), - "disk_model": "diluted_single_region", - "dilution_parameter": self.dilution_parameter, - "dayside_dilution": f"{dilution:.17g}", - "dilution_definition": "fractional_projected_dayside_emitting_area", - }, + return _diluted_spectrum( + spectrum, + dilution, + dilution_parameter=self.dilution_parameter, ) +@dataclass(frozen=True) +class MultiDatasetTwoRegionEmissionModel: + """Area-weight two regional models that each return named datasets.""" + + hot_model: MultiDatasetEmissionEvaluator + cold_model: MultiDatasetEmissionEvaluator + hot_fraction_parameter: str = "hot_area_fraction" + + def __post_init__(self) -> None: + if not self.hot_fraction_parameter: + raise RobertValidationError("hot_fraction_parameter must not be empty") + + def __call__(self, parameters: Mapping[str, float]) -> Mapping[str, Spectrum]: + hot_fraction = _validate_fraction( + _parameter(parameters, self.hot_fraction_parameter), + self.hot_fraction_parameter, + ) + hot = dict(self.hot_model(parameters)) + cold = dict(self.cold_model(parameters)) + if not hot or set(hot) != set(cold): + raise RobertValidationError( + "regional multi-dataset spectra must have matching non-empty names" + ) + return { + name: _mixed_spectrum( + hot[name], + cold[name], + hot_fraction, + hot_fraction_parameter=self.hot_fraction_parameter, + ) + for name in hot + } + + +@dataclass(frozen=True) +class MultiDatasetDilutedEmissionModel: + """Dilute every named spectrum from one multi-dataset emission model.""" + + emission_model: MultiDatasetEmissionEvaluator + dilution_parameter: str = "dayside_dilution" + + def __post_init__(self) -> None: + if not self.dilution_parameter: + raise RobertValidationError("dilution_parameter must not be empty") + + def __call__(self, parameters: Mapping[str, float]) -> Mapping[str, Spectrum]: + dilution = _validate_fraction( + _parameter(parameters, self.dilution_parameter), + self.dilution_parameter, + ) + spectra = dict(self.emission_model(parameters)) + if not spectra: + raise RobertValidationError("multi-dataset emission model returned no spectra") + return { + name: _diluted_spectrum( + spectrum, + dilution, + dilution_parameter=self.dilution_parameter, + ) + for name, spectrum in spectra.items() + } + + @dataclass(frozen=True) class DiskEmissionModelConfig: """Planet-independent selection of projected-dayside model geometry.""" @@ -178,6 +272,9 @@ def build_disk_emission_model( "DilutedEmissionModel", "DiskEmissionModelConfig", "EmissionEvaluator", + "MultiDatasetDilutedEmissionModel", + "MultiDatasetEmissionEvaluator", + "MultiDatasetTwoRegionEmissionModel", "TwoRegionEmissionModel", "build_disk_emission_model", ] diff --git a/src/robert_exoplanets/io/configured_tasks.py b/src/robert_exoplanets/io/configured_tasks.py index 5435739..7bbacaf 100644 --- a/src/robert_exoplanets/io/configured_tasks.py +++ b/src/robert_exoplanets/io/configured_tasks.py @@ -26,15 +26,20 @@ TabulatedTemperatureProfile, ) from robert_exoplanets.bodies import Planet, Star -from robert_exoplanets.core import PressureGrid, RobertConfigError +from robert_exoplanets.core import PressureGrid, RobertConfigError, SpectralGrid from robert_exoplanets.forward import ( + DilutedEmissionModel, + MultiDatasetDilutedEmissionModel, + MultiDatasetTwoRegionEmissionModel, ParameterizedEmissionFactoryConfig, ParameterizedEmissionModelConfig, ParameterizedDeckHazeCloudModel, ParameterizedMieCloudModel, ParameterizedTransmissionFactoryConfig, ParameterizedTransmissionModelConfig, + TwoRegionEmissionModel, build_multi_dataset_emission_model, + build_parameterized_emission_model, build_parameterized_transmission_model, ) from robert_exoplanets.instruments import ObservationCollection, ObservationDataset @@ -60,7 +65,16 @@ OpticalConstantsCatalog, ) -from .task_config import ParameterConfig, TaskConfig, initialize_task_directories +from .task_config import ( + CloudsConfig, + ChemistryConfig, + ParameterConfig, + ResolvedRegionConfig, + TaskConfig, + TemperatureConfig, + configured_regions, + initialize_task_directories, +) from .l9859b import load_bello_arufe2025_l9859b from .wasp69b import load_schlawin2024_wasp69b from .wasp80b import load_wiser2025_wasp80b @@ -95,12 +109,19 @@ def mpi_processes(config: TaskConfig) -> int: def describe_config(config: TaskConfig) -> str: """Return a concise preflight summary suitable for terminal inspection.""" + regions = configured_regions(config) + atmosphere_description = "; ".join( + f"{region.name}={region.atmosphere.temperature.model}/" + f"{region.atmosphere.chemistry.model}/clouds:{region.clouds.model}" + for region in regions + ) return "\n".join( [ f"Run: {config.run.name}", f"Target: {config.bodies.planet.name} / {config.bodies.star.name}", f"Data: {config.observations.loader} [{', '.join(config.observations.datasets)}]", - f"Atmosphere: {config.atmosphere.temperature.model}, {config.atmosphere.chemistry.model}, clouds={config.clouds.model}", + f"Disk emission: {config.disk_emission.model}", + f"Atmosphere regions: {atmosphere_description}", f"Opacity: {config.opacity.resolution} [{', '.join(config.opacity.species)}]", ( f"Transmission: reference={config.radiative_transfer.reference_pressure_bar:g} bar, " @@ -404,11 +425,211 @@ def _parameters(items: tuple[ParameterConfig, ...]) -> RetrievalParameterSet: return RetrievalParameterSet(tuple(parameters)) +def _temperature_profile(config: TemperatureConfig, *, gravity: float): + if config.model == "parmentier_guillot_2014": + return ParmentierGuillot2014TemperatureProfile( + gravity=gravity, + internal_temperature=config.internal_temperature_k, + kappa_ir_parameter_name=config.kappa_ir_parameter, + gamma1_parameter_name=config.gamma1_parameter, + gamma2_parameter_name=config.gamma2_parameter, + irradiation_temperature_parameter_name=( + config.irradiation_temperature_parameter + ), + alpha_parameter_name=config.alpha_parameter, + ) + if config.model == "isothermal": + return IsothermalTemperatureProfile( + temperature=config.temperature_k, + parameter_name=config.parameter_name, + ) + if config.model == "tabulated": + return TabulatedTemperatureProfile.from_csv( + config.profile_path, + pressure_column=config.pressure_column, + temperature_column=config.temperature_column, + pressure_unit=config.pressure_unit, + extrapolation=config.extrapolation, + ) + if config.model == "madhusudhan_seager_2009": + return MadhusudhanSeager2009TemperatureProfile( + pressure_unit=config.pressure_unit, + reference_pressure=config.reference_pressure, + p1_parameter_name=config.p1_parameter, + p2_parameter_name=config.p2_parameter, + p3_parameter_name=config.p3_parameter, + t0_parameter_name=config.t0_parameter, + alpha1_parameter_name=config.alpha1_parameter, + alpha2_parameter_name=config.alpha2_parameter, + ) + return SplineTemperatureProfile( + knot_pressure=np.asarray(config.knot_pressure, dtype=float), + knot_temperature=( + None + if config.knot_temperature_k is None + else np.asarray(config.knot_temperature_k, dtype=float) + ), + parameter_names=config.parameter_names, + pressure_unit=config.pressure_unit, + extrapolation=config.extrapolation, + ) + + +def _chemistry_components(config: ChemistryConfig): + if config.model == "fastchem_equilibrium": + chemistry = FastChemEquilibriumChemistry( + fastchem_path=config.fastchem_path, + fastchem_species=tuple(item.fastchem_name for item in config.species), + labels=tuple(item.label for item in config.species), + metallicity_parameter_name=config.metallicity_parameter, + carbon_to_oxygen_parameter_name=config.carbon_to_oxygen_parameter, + constant_log10_vmr_parameters=( + config.constant_log10_vmr_parameters or {} + ), + ) + return ( + chemistry, + CompositionMeanMolecularWeight(normalization="raw_sum"), + (), + ) + if config.fill_background: + fractions = config.background_fractions + background = BackgroundGasMixture( + {name: 1.0 for name in config.background_species} + if fractions is None + else dict(zip(config.background_species, fractions, strict=True)) + ) + else: + background = None + chemistry = FreeChemistry( + active_species=config.species, + background=background, + fixed_mixing_ratios=config.fixed_mixing_ratios, + parameter_names=config.parameter_names, + parameter_mode=config.parameter_mode, + fill_background=config.fill_background, + excess_policy=config.excess_policy, + ) + mean_molecular_weight = CompositionMeanMolecularWeight( + normalization="require" if config.fill_background else "normalize", + molecular_mass_parameters=( + {} + if config.phantom_species is None + else { + config.phantom_species: config.phantom_mean_molecular_weight_parameter + } + ), + ) + opacity_free_species = ( + () if config.phantom_species is None else (config.phantom_species,) + ) + return chemistry, mean_molecular_weight, opacity_free_species + + +def _cloud_model(config: CloudsConfig): + if config.model == "none": + return None + if config.model == "deck_haze": + return ParameterizedDeckHazeCloudModel( + log10_cloud_top_pressure_bar_parameter=( + config.log10_cloud_top_pressure_bar_parameter + ), + log10_cloud_optical_depth_parameter=( + config.log10_cloud_optical_depth_parameter + ), + log10_haze_mass_extinction_parameter=( + config.log10_haze_mass_extinction_parameter + ), + haze_slope_parameter=config.haze_slope_parameter, + haze_reference_wavelength_micron=( + config.haze_reference_wavelength_micron + ), + deck_single_scattering_albedo=config.deck_single_scattering_albedo, + deck_asymmetry_factor=config.deck_asymmetry_factor, + haze_single_scattering_albedo=config.haze_single_scattering_albedo, + haze_asymmetry_factor=config.haze_asymmetry_factor, + multiple_scattering_backend=config.multiple_scattering_backend, + ) + if config.model == "mie_catalog": + wavelength = () + real_parameters = () + imaginary_parameters = () + fixed_index = OpticalConstantsCatalog(config.optical_constants_path).load( + config.material + ) + else: + wavelength = config.refractive_index_wavelength_micron + real_parameters = config.real_index_parameter_names + imaginary_parameters = config.log10_imaginary_index_parameter_names + fixed_index = None + return ParameterizedMieCloudModel( + refractive_index_wavelength_micron=wavelength, + real_index_parameter_names=real_parameters, + log10_imaginary_index_parameter_names=imaginary_parameters, + fixed_refractive_index=fixed_index, + log10_condensate_mass_fraction_parameter=config.log10_mass_fraction_parameter, + log10_effective_radius_micron_parameter=config.log10_radius_micron_parameter, + particle_density_kg_m3=config.particle_density_kg_m3, + geometric_stddev=config.geometric_stddev, + log10_cloud_top_pressure_bar_parameter=config.log10_top_pressure_bar_parameter, + log10_cloud_base_pressure_bar_parameter=config.log10_base_pressure_bar_parameter, + quadrature_points=config.quadrature_points, + refractive_index_extrapolation="raise", + multiple_scattering_backend=config.multiple_scattering_backend, + ) + + +def _pressure_grid(region: ResolvedRegionConfig, *, planet: Planet) -> PressureGrid: + pressure = region.atmosphere.pressure + return PressureGrid.from_log_centers( + pressure.bottom_bar, + pressure.top_bar, + n_layers=pressure.layers, + unit="bar", + name=f"{planet.name} {region.name} configured pressure grid", + ) + + +def _opacity_providers( + config: TaskConfig, + observations: ObservationCollection, + *, + planet: Planet, +) -> dict[str, CorrelatedKOpacityProvider]: + providers = {} + for dataset in observations.datasets: + tables = { + species: _load_cached_table(config, dataset.name, species) + for species in config.opacity.species + } + reference = next(iter(tables.values())) + for species, table in tuple(tables.items()): + if not np.allclose( + table.g_samples, reference.g_samples, rtol=0.0, atol=1.0e-8 + ): + raise ValueError(f"{species} uses a different correlated-k g grid") + if not np.allclose( + table.g_weights, reference.g_weights, rtol=0.0, atol=1.0e-8 + ): + raise ValueError(f"{species} uses different correlated-k weights") + tables[species] = replace( + table, + g_samples=reference.g_samples, + g_weights=reference.g_weights, + ) + providers[dataset.name] = CorrelatedKOpacityProvider( + tables, + name=f"{planet.name}-{dataset.name}-{config.opacity.resolution}-binned", + interpolation="log_pressure_temperature_log_k_clip", + ) + return providers + + def build_problem( config: TaskConfig, observations: ObservationCollection | None = None, ) -> MultiDatasetRetrievalProblem: - """Construct the typed emission problem selected by the YAML file.""" + """Construct the typed regional retrieval problem selected by the YAML file.""" observations = load_observations(config) if observations is None else observations planet, gravity = _planet(config) @@ -420,111 +641,6 @@ def build_problem( log_g_cgs=star_item.log_g_cgs, metallicity_dex=star_item.metallicity_dex, ) - pressure_item = config.atmosphere.pressure - pressure = PressureGrid.from_log_centers( - pressure_item.bottom_bar, - pressure_item.top_bar, - n_layers=pressure_item.layers, - unit="bar", - name=f"{planet.name} configured pressure grid", - ) - temperature_item = config.atmosphere.temperature - if temperature_item.model == "parmentier_guillot_2014": - temperature = ParmentierGuillot2014TemperatureProfile( - gravity=gravity, - internal_temperature=temperature_item.internal_temperature_k, - ) - elif temperature_item.model == "isothermal": - temperature = IsothermalTemperatureProfile( - temperature=temperature_item.temperature_k, - parameter_name=temperature_item.parameter_name, - ) - elif temperature_item.model == "tabulated": - temperature = TabulatedTemperatureProfile.from_csv( - temperature_item.profile_path, - pressure_column=temperature_item.pressure_column, - temperature_column=temperature_item.temperature_column, - pressure_unit=temperature_item.pressure_unit, - extrapolation=temperature_item.extrapolation, - ) - elif temperature_item.model == "madhusudhan_seager_2009": - temperature = MadhusudhanSeager2009TemperatureProfile( - pressure_unit=temperature_item.pressure_unit, - reference_pressure=temperature_item.reference_pressure, - p1_parameter_name=temperature_item.p1_parameter, - p2_parameter_name=temperature_item.p2_parameter, - p3_parameter_name=temperature_item.p3_parameter, - t0_parameter_name=temperature_item.t0_parameter, - alpha1_parameter_name=temperature_item.alpha1_parameter, - alpha2_parameter_name=temperature_item.alpha2_parameter, - ) - else: - temperature = SplineTemperatureProfile( - knot_pressure=np.asarray(temperature_item.knot_pressure, dtype=float), - knot_temperature=( - None - if temperature_item.knot_temperature_k is None - else np.asarray(temperature_item.knot_temperature_k, dtype=float) - ), - parameter_names=temperature_item.parameter_names, - pressure_unit=temperature_item.pressure_unit, - extrapolation=temperature_item.extrapolation, - ) - chemistry_item = config.atmosphere.chemistry - if chemistry_item.model == "fastchem_equilibrium": - chemistry = FastChemEquilibriumChemistry( - fastchem_path=chemistry_item.fastchem_path, - fastchem_species=tuple( - item.fastchem_name for item in chemistry_item.species - ), - labels=tuple(item.label for item in chemistry_item.species), - metallicity_parameter_name=chemistry_item.metallicity_parameter, - carbon_to_oxygen_parameter_name=chemistry_item.carbon_to_oxygen_parameter, - constant_log10_vmr_parameters=( - chemistry_item.constant_log10_vmr_parameters or {} - ), - ) - mean_molecular_weight = CompositionMeanMolecularWeight(normalization="raw_sum") - opacity_free_species: tuple[str, ...] = () - else: - if chemistry_item.fill_background: - fractions = chemistry_item.background_fractions - if fractions is None: - background = BackgroundGasMixture( - {name: 1.0 for name in chemistry_item.background_species} - ) - else: - background = BackgroundGasMixture( - dict(zip(chemistry_item.background_species, fractions, strict=True)) - ) - else: - background = None - chemistry = FreeChemistry( - active_species=chemistry_item.species, - background=background, - fixed_mixing_ratios=chemistry_item.fixed_mixing_ratios, - parameter_names=chemistry_item.parameter_names, - parameter_mode=chemistry_item.parameter_mode, - fill_background=chemistry_item.fill_background, - excess_policy=chemistry_item.excess_policy, - ) - mean_molecular_weight = CompositionMeanMolecularWeight( - normalization="require" if chemistry_item.fill_background else "normalize", - molecular_mass_parameters=( - {} - if chemistry_item.phantom_species is None - else { - chemistry_item.phantom_species: ( - chemistry_item.phantom_mean_molecular_weight_parameter - ) - } - ), - ) - opacity_free_species = ( - () - if chemistry_item.phantom_species is None - else (chemistry_item.phantom_species,) - ) geometry_item = config.radiative_transfer.geometry geometry = ( normal_emission_geometry() @@ -553,145 +669,105 @@ def build_problem( impact_quadrature_order=rt.impact_quadrature_order, metadata={"configured_model": "transmission"}, ) - configs = {} - cloud_models = {} - spectral_grids = {} - clouds = config.clouds - shared_cloud_model = None - if clouds.model == "deck_haze": - shared_cloud_model = ParameterizedDeckHazeCloudModel( - log10_cloud_top_pressure_bar_parameter=( - clouds.log10_cloud_top_pressure_bar_parameter - ), - log10_cloud_optical_depth_parameter=( - clouds.log10_cloud_optical_depth_parameter - ), - log10_haze_mass_extinction_parameter=( - clouds.log10_haze_mass_extinction_parameter - ), - haze_slope_parameter=clouds.haze_slope_parameter, - haze_reference_wavelength_micron=( - clouds.haze_reference_wavelength_micron - ), - deck_single_scattering_albedo=( - clouds.deck_single_scattering_albedo - ), - deck_asymmetry_factor=clouds.deck_asymmetry_factor, - haze_single_scattering_albedo=( - clouds.haze_single_scattering_albedo - ), - haze_asymmetry_factor=clouds.haze_asymmetry_factor, - multiple_scattering_backend=clouds.multiple_scattering_backend, - ) - if clouds.model == "mie_catalog": - shared_cloud_model = ParameterizedMieCloudModel( - refractive_index_wavelength_micron=(), - real_index_parameter_names=(), - log10_imaginary_index_parameter_names=(), - fixed_refractive_index=OpticalConstantsCatalog( - clouds.optical_constants_path - ).load(clouds.material), - log10_condensate_mass_fraction_parameter=clouds.log10_mass_fraction_parameter, - log10_effective_radius_micron_parameter=clouds.log10_radius_micron_parameter, - particle_density_kg_m3=clouds.particle_density_kg_m3, - geometric_stddev=clouds.geometric_stddev, - log10_cloud_top_pressure_bar_parameter=clouds.log10_top_pressure_bar_parameter, - log10_cloud_base_pressure_bar_parameter=clouds.log10_base_pressure_bar_parameter, - quadrature_points=clouds.quadrature_points, - refractive_index_extrapolation="raise", - multiple_scattering_backend=clouds.multiple_scattering_backend, - ) - elif clouds.model == "mie_direct_nk": - shared_cloud_model = ParameterizedMieCloudModel( - refractive_index_wavelength_micron=clouds.refractive_index_wavelength_micron, - real_index_parameter_names=clouds.real_index_parameter_names, - log10_imaginary_index_parameter_names=clouds.log10_imaginary_index_parameter_names, - log10_condensate_mass_fraction_parameter=clouds.log10_mass_fraction_parameter, - log10_effective_radius_micron_parameter=clouds.log10_radius_micron_parameter, - particle_density_kg_m3=clouds.particle_density_kg_m3, - geometric_stddev=clouds.geometric_stddev, - log10_cloud_top_pressure_bar_parameter=clouds.log10_top_pressure_bar_parameter, - log10_cloud_base_pressure_bar_parameter=clouds.log10_base_pressure_bar_parameter, - quadrature_points=clouds.quadrature_points, - refractive_index_extrapolation="raise", - multiple_scattering_backend=clouds.multiple_scattering_backend, + providers = _opacity_providers(config, observations, planet=planet) + spectral_grids = { + dataset.name: dataset.observation.spectral_grid + for dataset in observations.datasets + } + regions = configured_regions(config) + regional_models = {} + regional_opacity_ids = {} + for region in regions: + temperature = _temperature_profile( + region.atmosphere.temperature, + gravity=gravity, ) - for dataset in observations.datasets: - tables = { - species: _load_cached_table(config, dataset.name, species) - for species in config.opacity.species - } - reference = next(iter(tables.values())) - for species, table in tuple(tables.items()): - if not np.allclose( - table.g_samples, reference.g_samples, rtol=0.0, atol=1.0e-8 - ): - raise ValueError(f"{species} uses a different correlated-k g grid") - if not np.allclose( - table.g_weights, reference.g_weights, rtol=0.0, atol=1.0e-8 - ): - raise ValueError(f"{species} uses different correlated-k weights") - tables[species] = replace( - table, g_samples=reference.g_samples, g_weights=reference.g_weights - ) - provider = CorrelatedKOpacityProvider( - tables, - name=f"{planet.name}-{dataset.name}-{config.opacity.resolution}-binned", - interpolation="log_pressure_temperature_log_k_clip", + chemistry, mean_molecular_weight, opacity_free_species = ( + _chemistry_components(region.atmosphere.chemistry) ) - spectral_grids[dataset.name] = dataset.observation.spectral_grid - if rt.model == "transmission": - factory = ParameterizedTransmissionFactoryConfig( - planet=planet, - star=star, - temperature_profile=temperature, - chemistry_model=chemistry, - mean_molecular_weight_model=mean_molecular_weight, - opacity_free_species=opacity_free_species, - pressure_grid=pressure, - cia_table=cia, - opacity_source=provider, - opacity_binning=None, - model=model_config, - cloud_model=shared_cloud_model, - ) - cloud_models[dataset.name] = build_parameterized_transmission_model( - factory, - spectral_grid=dataset.observation.spectral_grid, + pressure = _pressure_grid(region, planet=planet) + cloud = _cloud_model(region.clouds) + if rt.model == "emission": + factories = { + dataset.name: ParameterizedEmissionFactoryConfig( + planet=planet, + star=star, + temperature_profile=temperature, + chemistry_model=chemistry, + mean_molecular_weight_model=mean_molecular_weight, + opacity_free_species=opacity_free_species, + pressure_grid=pressure, + cia_table=cia, + geometry=geometry, + opacity_source=providers[dataset.name], + opacity_binning=None, + model=model_config, + cloud_model=cloud, + ) + for dataset in observations.datasets + } + regional_model = build_multi_dataset_emission_model( + factories, + spectral_grids=spectral_grids, ) + regional_opacity_ids[region.name] = { + f"{dataset}:{key}": value + for dataset, model in regional_model.models.items() + for key, value in model.opacity_identifiers.items() + } else: - factory = ParameterizedEmissionFactoryConfig( - planet=planet, - star=star, - temperature_profile=temperature, - chemistry_model=chemistry, - mean_molecular_weight_model=mean_molecular_weight, - opacity_free_species=opacity_free_species, - pressure_grid=pressure, - cia_table=cia, - geometry=geometry, - opacity_source=provider, - opacity_binning=None, - model=model_config, - cloud_model=shared_cloud_model, - ) - configs[dataset.name] = factory - if rt.model == "emission": - forward_model = build_multi_dataset_emission_model( - configs, spectral_grids=spectral_grids + dataset_models = {} + for dataset in observations.datasets: + factory = ParameterizedTransmissionFactoryConfig( + planet=planet, + star=star, + temperature_profile=temperature, + chemistry_model=chemistry, + mean_molecular_weight_model=mean_molecular_weight, + opacity_free_species=opacity_free_species, + pressure_grid=pressure, + cia_table=cia, + opacity_source=providers[dataset.name], + opacity_binning=None, + model=model_config, + cloud_model=cloud, + ) + dataset_models[dataset.name] = build_parameterized_transmission_model( + factory, + spectral_grid=dataset.observation.spectral_grid, + ) + regional_model = NamedRegionalModels(dataset_models) + regional_opacity_ids[region.name] = { + f"{dataset}:{key}": value + for dataset, model in dataset_models.items() + for key, value in model.opacity_identifiers.items() + } + regional_models[region.name] = regional_model + + disk_mode = config.disk_emission.model + if disk_mode in {"two_region", "2tp"}: + forward_model = MultiDatasetTwoRegionEmissionModel( + regional_models["hot"], + regional_models["cold"], + hot_fraction_parameter=config.disk_emission.hot_fraction_parameter, ) opacity_ids = { - f"{dataset}:{key}": value - for dataset, model in forward_model.models.items() - for key, value in model.opacity_identifiers.items() + f"{region}:{key}": value + for region, identifiers in regional_opacity_ids.items() + for key, value in identifiers.items() } + canonical_disk_mode = "two_region" + elif disk_mode in {"diluted", "diluted_one_region"}: + forward_model = MultiDatasetDilutedEmissionModel( + regional_models["primary"], + dilution_parameter=config.disk_emission.dilution_parameter, + ) + opacity_ids = regional_opacity_ids["primary"] + canonical_disk_mode = "diluted_one_region" else: - forward_model = NamedRegionalModels(cloud_models) - opacity_ids = { - f"{dataset}:{key}": value - for dataset, model in cloud_models.items() - for key, value in model.opacity_identifiers.items() - } + forward_model = regional_models["primary"] + opacity_ids = regional_opacity_ids["primary"] + canonical_disk_mode = "one_region" return MultiDatasetRetrievalProblem( name=config.run.name, observations=observations, @@ -704,7 +780,10 @@ def build_problem( metadata={ "configuration_schema_version": str(config.schema_version), "opacity_resolution": config.opacity.resolution, - "cloud_model": config.clouds.model, + "cloud_model": ",".join( + f"{region.name}:{region.clouds.model}" for region in regions + ), + "disk_emission_model": canonical_disk_mode, "geometry": geometry_item.model, "radiative_transfer_model": rt.model, }, @@ -712,14 +791,147 @@ def build_problem( ) -def smoke_evaluation(problem: MultiDatasetRetrievalProblem) -> dict[str, float]: - theta = problem.prior_transform(np.full(problem.ndim, 0.5)) +def build_native_emission_model( + config: TaskConfig, + observations: ObservationCollection, +): + """Build an exact native-opacity-grid model for retrieval diagnostics. + + Native diagnostic spectra are currently available for ExoMol correlated-k + tables. Cross-section HDF inputs require a separate spectral correlation + choice and therefore deliberately fall back to observation-grid plotting. + """ + + if ( + config.radiative_transfer.model != "emission" + or config.opacity.format != "exomol_kta" + ): + return None + provider = CorrelatedKOpacityProvider.from_exomol_kta_directory( + config.opacity.path, + species=config.opacity.species, + resolution=config.opacity.resolution, + name=f"native-{config.opacity.resolution}", + interpolation="log_pressure_temperature_log_k_clip", + nonfinite_policy="floor", + ) + reference = provider.tables[config.opacity.species[0]] + common_wavenumber = np.asarray(reference.wavenumber_cm_inverse, dtype=float) + for species in config.opacity.species[1:]: + common_wavenumber = np.intersect1d( + common_wavenumber, + provider.tables[species].wavenumber_cm_inverse, + ) + observed_wavelength = np.concatenate( + [dataset.observation.wavelength for dataset in observations.datasets] + ) + native_wavelength = 10_000.0 / common_wavenumber + selected = ( + (native_wavelength >= float(np.min(observed_wavelength))) + & (native_wavelength <= float(np.max(observed_wavelength))) + ) + native_wavelength = np.sort(native_wavelength[selected]) + if native_wavelength.size < 2: + raise RobertConfigError( + "native opacity grid does not cover the configured observations" + ) + spectral_grid = SpectralGrid.from_array( + native_wavelength, + unit="micron", + name=f"{config.opacity.resolution} native opacity diagnostic grid", + ) + planet, gravity = _planet(config) + star_item = config.bodies.star + star = Star( + name=star_item.name, + radius_m=star_item.radius_m, + effective_temperature_k=star_item.effective_temperature_k, + log_g_cgs=star_item.log_g_cgs, + metallicity_dex=star_item.metallicity_dex, + ) + geometry_item = config.radiative_transfer.geometry + geometry = ( + normal_emission_geometry() + if geometry_item.model == "normal_emission" + else gauss_legendre_disk_geometry(geometry_item.points) + ) + rt = config.radiative_transfer + model_config = ParameterizedEmissionModelConfig( + opacity_species=config.opacity.species, + include_rayleigh=rt.include_rayleigh, + gas_combination=rt.gas_combination, + thermal_integration_backend=rt.thermal_integration_backend, + stellar_spectrum_model=star_item.spectrum_model, + metadata={ + "configured_geometry": geometry_item.model, + "diagnostic_grid": "native_opacity", + }, + ) + cia = load_nemesispy_cia_table() + regional_models = {} + for region in configured_regions(config): + chemistry, mean_molecular_weight, opacity_free_species = ( + _chemistry_components(region.atmosphere.chemistry) + ) + factory = ParameterizedEmissionFactoryConfig( + planet=planet, + star=star, + temperature_profile=_temperature_profile( + region.atmosphere.temperature, + gravity=gravity, + ), + chemistry_model=chemistry, + mean_molecular_weight_model=mean_molecular_weight, + opacity_free_species=opacity_free_species, + pressure_grid=_pressure_grid(region, planet=planet), + cia_table=cia, + geometry=geometry, + opacity_source=provider, + opacity_binning=None, + model=model_config, + cloud_model=_cloud_model(region.clouds), + ) + regional_models[region.name] = build_parameterized_emission_model( + factory, + spectral_grid=spectral_grid, + ) + disk = config.disk_emission + if disk.model in {"two_region", "2tp"}: + return TwoRegionEmissionModel( + regional_models["hot"], + regional_models["cold"], + hot_fraction_parameter=disk.hot_fraction_parameter, + ) + if disk.model in {"diluted", "diluted_one_region"}: + return DilutedEmissionModel( + regional_models["primary"], + dilution_parameter=disk.dilution_parameter, + ) + return regional_models["primary"] + + +def smoke_evaluation( + problem: MultiDatasetRetrievalProblem, + state: np.ndarray | None = None, +) -> dict[str, float]: + """Evaluate a representative state and expose any hidden model error.""" + + theta = ( + problem.prior_transform(np.full(problem.ndim, 0.5)) + if state is None + else np.asarray(state, dtype=float) + ) started = time.perf_counter() value = problem.log_likelihood_from_vector(theta) elapsed = time.perf_counter() - started if not np.isfinite(value) or value <= problem.invalid_loglike: + # log_likelihood_from_vector deliberately converts physical-domain + # errors to the sampler's invalid floor. Re-evaluate through the + # unguarded OE interface so setup-time smoke failures retain their + # actionable original exception. + problem.gaussian_inputs_from_vector(theta) raise RuntimeError( - "prior-midpoint smoke evaluation returned an invalid likelihood" + "smoke evaluation returned an invalid likelihood" ) return {"elapsed_seconds": elapsed, "log_likelihood": float(value)} @@ -1017,6 +1229,10 @@ def _postprocess_retrieval_outputs( if parameter.label is not None } parameter_labels.update(config.plotting.parameter_labels) + native_spectrum_model = build_native_emission_model( + config, + problem.observations, + ) for result_dir in discover_retrieval_result_directories(config.outputs.directory): postprocess_retrieval_output( problem, @@ -1028,6 +1244,12 @@ def _postprocess_retrieval_outputs( image_format=config.plotting.image_format, dpi=config.plotting.dpi, max_posterior_samples=config.plotting.max_posterior_samples, + posterior_predictive_samples=( + config.plotting.posterior_predictive_samples + ), + posterior_predictive_seed=config.plotting.posterior_predictive_seed, + corner_max_parameters=config.plotting.corner_max_parameters, + native_spectrum_model=native_spectrum_model, leave_one_out=config.plotting.leave_one_out.enabled, loo_max_posterior_draws=( config.plotting.leave_one_out.max_posterior_draws @@ -1044,6 +1266,7 @@ def _postprocess_retrieval_outputs( __all__ = [ + "build_native_emission_model", "build_problem", "describe_config", "load_observations", diff --git a/src/robert_exoplanets/io/task_config.py b/src/robert_exoplanets/io/task_config.py index e3a670f..9a263ff 100644 --- a/src/robert_exoplanets/io/task_config.py +++ b/src/robert_exoplanets/io/task_config.py @@ -100,6 +100,24 @@ def validate_order(self) -> "PressureConfig": class ParmentierGuillotTemperatureConfig(ConfigModel): model: Literal["parmentier_guillot_2014"] internal_temperature_k: float = Field(default=100.0, ge=0.0) + kappa_ir_parameter: str = Field(default="kappa_IR", min_length=1) + gamma1_parameter: str = Field(default="gamma1", min_length=1) + gamma2_parameter: str = Field(default="gamma2", min_length=1) + irradiation_temperature_parameter: str = Field(default="T_irr", min_length=1) + alpha_parameter: str = Field(default="alpha", min_length=1) + + @model_validator(mode="after") + def validate_parameter_names(self) -> "ParmentierGuillotTemperatureConfig": + names = ( + self.kappa_ir_parameter, + self.gamma1_parameter, + self.gamma2_parameter, + self.irradiation_temperature_parameter, + self.alpha_parameter, + ) + if len(set(names)) != len(names): + raise ValueError("Parmentier-Guillot parameter names must be unique") + return self class IsothermalTemperatureConfig(ConfigModel): @@ -393,6 +411,53 @@ def validate_node_names(self) -> "MieDirectNkCloudConfig": ] +class RegionalAtmosphereOverrideConfig(ConfigModel): + """Optional regional replacements for the shared atmospheric defaults.""" + + pressure: PressureConfig | None = None + temperature: TemperatureConfig | None = None + chemistry: ChemistryConfig | None = None + + +class RegionOverrideConfig(ConfigModel): + """Overrides applied to one projected-disk emission region.""" + + atmosphere: RegionalAtmosphereOverrideConfig = RegionalAtmosphereOverrideConfig() + clouds: CloudsConfig | None = None + + +class OneRegionDiskEmissionConfig(ConfigModel): + model: Literal["one_region"] = "one_region" + + +class DilutedDiskEmissionConfig(ConfigModel): + model: Literal["diluted_one_region", "diluted"] + dilution_parameter: str = Field(default="dayside_dilution", min_length=1) + + +class TwoRegionDiskEmissionConfig(ConfigModel): + model: Literal["two_region", "2tp"] + hot_fraction_parameter: str = Field(default="hot_area_fraction", min_length=1) + hot_region: RegionOverrideConfig = RegionOverrideConfig() + cold_region: RegionOverrideConfig + + +DiskEmissionConfig = Annotated[ + OneRegionDiskEmissionConfig + | DilutedDiskEmissionConfig + | TwoRegionDiskEmissionConfig, + Field(discriminator="model"), +] + + +class ResolvedRegionConfig(ConfigModel): + """Fully resolved configuration for one physical emission column.""" + + name: str = Field(min_length=1) + atmosphere: AtmosphereConfig + clouds: CloudsConfig + + class OpacityBinningConfig(ConfigModel): num: PositiveInt = 300 use_rebin: bool = False @@ -556,6 +621,9 @@ class PlottingConfig(ConfigModel): image_format: Literal["png", "pdf", "svg"] = "png" dpi: PositiveInt = 180 max_posterior_samples: PositiveInt = 20_000 + posterior_predictive_samples: PositiveInt = 200 + posterior_predictive_seed: NonNegativeInt = 0 + corner_max_parameters: PositiveInt = 20 dataset_colors: dict[str, str] = Field(default_factory=dict) parameter_labels: dict[str, str] = Field(default_factory=dict) leave_one_out: LeaveOneOutConfig = LeaveOneOutConfig() @@ -580,14 +648,15 @@ class RuntimeConfig(ConfigModel): class HousekeepingConfig(ConfigModel): - """Machine- and project-specific locations for a user-facing YAML file.""" - - observations_directory: Path - fastchem_directory: Path - k_table_directory: Path - opacity_cache_directory: Path - output_directory: Path - scratch_directory: Path + """Machine- and project-specific paths kept in one visible YAML block.""" + + project_directory: Path | None = None + observations_directory: Path | None = None + fastchem_directory: Path | None = None + k_table_directory: Path | None = None + opacity_cache_directory: Path | None = None + output_directory: Path | None = None + scratch_directory: Path | None = None optical_constants_directory: Path | None = None @@ -595,11 +664,13 @@ class TaskConfig(ConfigModel): """Complete schema-versioned retrieval/forward-model configuration.""" schema_version: Literal[2] + paths: HousekeepingConfig | None = None run: RunConfig bodies: BodiesConfig observations: ObservationsConfig atmosphere: AtmosphereConfig clouds: CloudsConfig = CloudFreeConfig() + disk_emission: DiskEmissionConfig = OneRegionDiskEmissionConfig() opacity: OpacityConfig radiative_transfer: RadiativeTransferConfig likelihood: LikelihoodConfig = LikelihoodConfig() @@ -608,79 +679,67 @@ class TaskConfig(ConfigModel): outputs: OutputsConfig plotting: PlottingConfig = PlottingConfig() runtime: RuntimeConfig + # Legacy name retained so existing configurations continue to load. New + # YAML should use the top-level ``paths`` block. housekeeping: HousekeepingConfig | None = None @model_validator(mode="after") def validate_cross_references(self) -> "TaskConfig": + if self.paths is not None and self.housekeeping is not None: + raise ValueError("configure paths or legacy housekeeping, not both") opacity = self.opacity.species - chemistry_config = self.atmosphere.chemistry - chemistry = ( - tuple(item.label for item in chemistry_config.species) - if chemistry_config.model == "fastchem_equilibrium" - else chemistry_config.species + chemistry_config.background_species - ) if len(set(opacity)) != len(opacity): raise ValueError("opacity.species contains duplicates") - if len(set(chemistry)) != len(chemistry): - raise ValueError("atmosphere.chemistry.species labels contain duplicates") - missing = sorted(set(opacity) - set(chemistry)) - if missing: - raise ValueError( - "opacity species missing from chemistry labels: " + ", ".join(missing) - ) + regions = configured_regions(self) + for region in regions: + chemistry = _chemistry_species(region.atmosphere.chemistry) + if len(set(chemistry)) != len(chemistry): + raise ValueError( + f"{region.name} chemistry species labels contain duplicates" + ) + missing = sorted(set(opacity) - set(chemistry)) + if missing: + raise ValueError( + f"opacity species missing from {region.name} chemistry labels: " + + ", ".join(missing) + ) names = tuple(item.name for item in self.parameters) if len(set(names)) != len(names): raise ValueError("parameter names must be unique") required: set[str] = set() - if chemistry_config.model == "fastchem_equilibrium": - required.update( - { - chemistry_config.metallicity_parameter, - chemistry_config.carbon_to_oxygen_parameter, - } - ) - required.update( - (chemistry_config.constant_log10_vmr_parameters or {}).values() - ) + for region in regions: + required.update(_required_chemistry_parameters(region.atmosphere.chemistry)) + required.update(_required_temperature_parameters(region.atmosphere.temperature)) + required.update(_required_cloud_parameters(region.clouds)) + disk_mode = _disk_emission_mode(self.disk_emission) + if disk_mode == "diluted_one_region": + required.add(self.disk_emission.dilution_parameter) + fraction_parameter = self.disk_emission.dilution_parameter + elif disk_mode == "two_region": + required.add(self.disk_emission.hot_fraction_parameter) + fraction_parameter = self.disk_emission.hot_fraction_parameter else: - required.update( - chemistry_config.parameter_names.get(species, species) - for species in chemistry_config.species - if species not in chemistry_config.fixed_mixing_ratios - ) - if chemistry_config.phantom_mean_molecular_weight_parameter is not None: - required.add( - chemistry_config.phantom_mean_molecular_weight_parameter - ) - temperature = self.atmosphere.temperature - if temperature.model == "parmentier_guillot_2014": - required.update({"kappa_IR", "gamma1", "gamma2", "T_irr", "alpha"}) - elif temperature.model == "isothermal" and temperature.temperature_k is None: - required.add(temperature.parameter_name) - elif temperature.model == "madhusudhan_seager_2009": - required.update( - { - temperature.p1_parameter, - temperature.p2_parameter, - temperature.p3_parameter, - temperature.t0_parameter, - temperature.alpha1_parameter, - temperature.alpha2_parameter, - } - ) - elif temperature.model == "spline" and temperature.knot_temperature_k is None: - required.update( - temperature.parameter_names - or tuple( - f"temperature_{index}" - for index in range(len(temperature.knot_pressure)) + fraction_parameter = None + if fraction_parameter is not None: + parameter_lookup = {item.name: item for item in self.parameters} + configured_fraction = parameter_lookup.get(fraction_parameter) + if configured_fraction is not None and ( + configured_fraction.prior.lower < 0.0 + or configured_fraction.prior.upper > 1.0 + ): + raise ValueError( + f"{fraction_parameter} prior must lie within [0, 1]" ) - ) if self.sampler.oe_temperature_prior_sigma_k is not None: if not self.sampler.engine.startswith("optimal_estimation"): raise ValueError( "correlated OE temperature priors require an optimal-estimation engine" ) + if disk_mode != "one_region": + raise ValueError( + "correlated OE temperature priors currently require one_region disk emission" + ) + temperature = regions[0].atmosphere.temperature if ( temperature.model != "spline" or temperature.knot_temperature_k is not None @@ -709,11 +768,16 @@ def validate_cross_references(self) -> "TaskConfig": if self.observations.miri_offset_parameter is not None: required.add(self.observations.miri_offset_parameter) radiative_transfer = self.radiative_transfer + if radiative_transfer.model == "transmission" and disk_mode != "one_region": + raise ValueError( + "diluted and two-region disk models require radiative_transfer.model=emission" + ) if radiative_transfer.model == "transmission": + pressure = regions[0].atmosphere.pressure if not ( - self.atmosphere.pressure.top_bar + pressure.top_bar <= radiative_transfer.reference_pressure_bar - <= self.atmosphere.pressure.bottom_bar + <= pressure.bottom_bar ): raise ValueError( "transmission reference_pressure_bar must lie within the pressure grid" @@ -726,34 +790,42 @@ def validate_cross_references(self) -> "TaskConfig": if item.prior.type == "centered_log_ratio" } if clr_parameters: - if chemistry_config.model != "free": - raise ValueError( - "centered_log_ratio priors require free chemistry" - ) - if chemistry_config.parameter_mode != "log10": - raise ValueError( - "centered_log_ratio priors require chemistry parameter_mode=log10" + matched_clr_parameters: set[str] = set() + for region in regions: + chemistry_config = region.atmosphere.chemistry + if chemistry_config.model != "free": + continue + chemistry_parameters = _retrieved_free_chemistry_parameters( + chemistry_config ) - if not chemistry_config.fill_background: - raise ValueError( - "centered_log_ratio priors require a background closure category" - ) - chemistry_parameters = { - chemistry_config.parameter_names.get(species, species) - for species in chemistry_config.species - if species not in chemistry_config.fixed_mixing_ratios - } - if set(clr_parameters) != chemistry_parameters: - raise ValueError( - "all and only retrieved free-chemistry abundances must use " - "centered_log_ratio priors" - ) - groups = { - prior.group or "composition" for prior in clr_parameters.values() - } - if len(groups) != 1: + regional_clr = chemistry_parameters.intersection(clr_parameters) + if not regional_clr: + continue + if chemistry_config.parameter_mode != "log10": + raise ValueError( + "centered_log_ratio priors require chemistry parameter_mode=log10" + ) + if not chemistry_config.fill_background: + raise ValueError( + "centered_log_ratio priors require a background closure category" + ) + if regional_clr != chemistry_parameters: + raise ValueError( + "all and only retrieved free-chemistry abundances must use " + "centered_log_ratio priors within each region" + ) + groups = { + clr_parameters[name].group or "composition" + for name in regional_clr + } + if len(groups) != 1: + raise ValueError( + "free-chemistry centered_log_ratio priors must share one group per region" + ) + matched_clr_parameters.update(regional_clr) + if matched_clr_parameters != set(clr_parameters): raise ValueError( - "free-chemistry centered_log_ratio priors must share one group" + "centered_log_ratio priors require matching free chemistry parameters" ) if self.sampler.engine == "optimal_estimation" or self.sampler.engine.startswith( "optimal_estimation_to_" @@ -767,36 +839,138 @@ def validate_cross_references(self) -> "TaskConfig": "required model parameters are missing: " + ", ".join(missing_parameters) ) - clouds = self.clouds - if clouds.model == "deck_haze": - cloud_parameters = { - clouds.log10_cloud_top_pressure_bar_parameter, - clouds.log10_cloud_optical_depth_parameter, - clouds.log10_haze_mass_extinction_parameter, - clouds.haze_slope_parameter, - } - elif clouds.model != "none": - cloud_parameters = { - clouds.log10_mass_fraction_parameter, - clouds.log10_radius_micron_parameter, - clouds.log10_top_pressure_bar_parameter, - clouds.log10_base_pressure_bar_parameter, - } - if clouds.model == "mie_direct_nk": - cloud_parameters.update(clouds.real_index_parameter_names) - cloud_parameters.update(clouds.log10_imaginary_index_parameter_names) - else: - cloud_parameters = set() - if cloud_parameters: - missing_cloud_parameters = sorted(cloud_parameters - set(names)) - if missing_cloud_parameters: - raise ValueError( - "required cloud parameters are missing: " - + ", ".join(missing_cloud_parameters) - ) return self +def _disk_emission_mode(config: DiskEmissionConfig) -> str: + aliases = { + "one_region": "one_region", + "diluted": "diluted_one_region", + "diluted_one_region": "diluted_one_region", + "2tp": "two_region", + "two_region": "two_region", + } + return aliases[config.model] + + +def configured_regions(config: TaskConfig) -> tuple[ResolvedRegionConfig, ...]: + """Resolve regional overrides against the top-level atmospheric defaults.""" + + mode = _disk_emission_mode(config.disk_emission) + if mode != "two_region": + return ( + ResolvedRegionConfig( + name="primary", + atmosphere=config.atmosphere, + clouds=config.clouds, + ), + ) + disk = config.disk_emission + return ( + _resolve_region("hot", config.atmosphere, config.clouds, disk.hot_region), + _resolve_region("cold", config.atmosphere, config.clouds, disk.cold_region), + ) + + +def _resolve_region( + name: str, + base_atmosphere: AtmosphereConfig, + base_clouds: CloudsConfig, + override: RegionOverrideConfig, +) -> ResolvedRegionConfig: + atmosphere_override = override.atmosphere + atmosphere = AtmosphereConfig( + pressure=atmosphere_override.pressure or base_atmosphere.pressure, + temperature=atmosphere_override.temperature or base_atmosphere.temperature, + chemistry=atmosphere_override.chemistry or base_atmosphere.chemistry, + ) + return ResolvedRegionConfig( + name=name, + atmosphere=atmosphere, + clouds=override.clouds or base_clouds, + ) + + +def _chemistry_species(config: ChemistryConfig) -> tuple[str, ...]: + if config.model == "fastchem_equilibrium": + return tuple(item.label for item in config.species) + return config.species + config.background_species + + +def _required_chemistry_parameters(config: ChemistryConfig) -> set[str]: + if config.model == "fastchem_equilibrium": + return { + config.metallicity_parameter, + config.carbon_to_oxygen_parameter, + *(config.constant_log10_vmr_parameters or {}).values(), + } + required = _retrieved_free_chemistry_parameters(config) + if config.phantom_mean_molecular_weight_parameter is not None: + required.add(config.phantom_mean_molecular_weight_parameter) + return required + + +def _retrieved_free_chemistry_parameters(config: FreeChemistryConfig) -> set[str]: + return { + config.parameter_names.get(species, species) + for species in config.species + if species not in config.fixed_mixing_ratios + } + + +def _required_temperature_parameters(config: TemperatureConfig) -> set[str]: + if config.model == "parmentier_guillot_2014": + return { + config.kappa_ir_parameter, + config.gamma1_parameter, + config.gamma2_parameter, + config.irradiation_temperature_parameter, + config.alpha_parameter, + } + if config.model == "isothermal": + return set() if config.temperature_k is not None else {config.parameter_name} + if config.model == "madhusudhan_seager_2009": + return { + config.p1_parameter, + config.p2_parameter, + config.p3_parameter, + config.t0_parameter, + config.alpha1_parameter, + config.alpha2_parameter, + } + if config.model == "spline" and config.knot_temperature_k is None: + return set( + config.parameter_names + or tuple( + f"temperature_{index}" + for index in range(len(config.knot_pressure)) + ) + ) + return set() + + +def _required_cloud_parameters(config: CloudsConfig) -> set[str]: + if config.model == "deck_haze": + return { + config.log10_cloud_top_pressure_bar_parameter, + config.log10_cloud_optical_depth_parameter, + config.log10_haze_mass_extinction_parameter, + config.haze_slope_parameter, + } + if config.model == "none": + return set() + required = { + config.log10_mass_fraction_parameter, + config.log10_radius_micron_parameter, + config.log10_top_pressure_bar_parameter, + config.log10_base_pressure_bar_parameter, + } + if config.model == "mie_direct_nk": + required.update(config.real_index_parameter_names) + required.update(config.log10_imaginary_index_parameter_names) + return required + + def load_task_config(path: str | Path) -> TaskConfig: """Load YAML, resolve its paths relative to the file, and validate it.""" @@ -804,7 +978,7 @@ def load_task_config(path: str | Path) -> TaskConfig: if not source.is_file(): raise FileNotFoundError(source) raw = _load_yaml_mapping(source) - _apply_housekeeping_paths(raw) + _apply_configured_paths(raw) for keys in ( ("observations", "path"), ("atmosphere", "chemistry", "fastchem_path"), @@ -812,6 +986,21 @@ def load_task_config(path: str | Path) -> TaskConfig: ("opacity", "path"), ("opacity", "cache_directory"), ("clouds", "optical_constants_path"), + ("disk_emission", "hot_region", "atmosphere", "chemistry", "fastchem_path"), + ("disk_emission", "hot_region", "atmosphere", "temperature", "profile_path"), + ("disk_emission", "hot_region", "clouds", "optical_constants_path"), + ("disk_emission", "cold_region", "atmosphere", "chemistry", "fastchem_path"), + ("disk_emission", "cold_region", "atmosphere", "temperature", "profile_path"), + ("disk_emission", "cold_region", "clouds", "optical_constants_path"), + ("paths", "project_directory"), + ("paths", "observations_directory"), + ("paths", "fastchem_directory"), + ("paths", "k_table_directory"), + ("paths", "opacity_cache_directory"), + ("paths", "output_directory"), + ("paths", "scratch_directory"), + ("paths", "optical_constants_directory"), + ("housekeeping", "project_directory"), ("housekeeping", "observations_directory"), ("housekeeping", "fastchem_directory"), ("housekeeping", "k_table_directory"), @@ -830,21 +1019,47 @@ def load_task_config(path: str | Path) -> TaskConfig: else: if not isinstance(section, dict) or keys[-1] not in section: continue - value = Path(section[keys[-1]]).expanduser() + if section[keys[-1]] is None: + continue + value = _expanded_path(section[keys[-1]]) section[keys[-1]] = str( value if value.is_absolute() else source.parent / value ) return TaskConfig.model_validate(raw) -def _apply_housekeeping_paths(raw: dict) -> None: - """Fill internal path fields from one readable user-facing path block.""" - - housekeeping = raw.get("housekeeping") - if housekeeping is None: - return - if not isinstance(housekeeping, dict): - raise ValueError("housekeeping must be a YAML mapping") +def _expanded_path(value: object) -> Path: + text = os.path.expandvars(str(value)) + if "$" in text: + raise ValueError(f"path contains an undefined environment variable: {value}") + return Path(text).expanduser() + + +def _apply_configured_paths(raw: dict) -> None: + """Fill component paths and local writable defaults from the top path block.""" + + if raw.get("paths") is not None and raw.get("housekeeping") is not None: + raise ValueError("configure paths or legacy housekeeping, not both") + path_config = raw.get("paths", raw.get("housekeeping", {})) + if path_config is None: + path_config = {} + if not isinstance(path_config, dict): + raise ValueError("paths must be a YAML mapping") + project_directory = path_config.get("project_directory") or "." + derived = { + "opacity_cache_directory": str( + Path(str(project_directory)) / "opacity_cache" + ), + "output_directory": str(Path(str(project_directory)) / "outputs"), + "scratch_directory": str(Path(str(project_directory)) / "scratch"), + } + for key, value in derived.items(): + if path_config.get(key) is None: + path_config[key] = value + if raw.get("paths") is not None: + raw["paths"] = path_config + elif raw.get("housekeeping") is not None: + raw["housekeeping"] = path_config mappings = ( (("observations", "path"), "observations_directory"), (("opacity", "path"), "k_table_directory"), @@ -853,7 +1068,7 @@ def _apply_housekeeping_paths(raw: dict) -> None: (("runtime", "scratch_directory"), "scratch_directory"), ) for keys, source_key in mappings: - if source_key not in housekeeping: + if path_config.get(source_key) is None: continue section = raw for key in keys[:-1]: @@ -862,23 +1077,39 @@ def _apply_housekeeping_paths(raw: dict) -> None: section = section.setdefault(key, {}) if not isinstance(section, dict): raise ValueError("configuration sections must be YAML mappings") - section.setdefault(keys[-1], housekeeping[source_key]) - chemistry = raw.get("atmosphere", {}).get("chemistry", {}) - if ( - isinstance(chemistry, dict) - and chemistry.get("model") == "fastchem_equilibrium" - and "fastchem_directory" in housekeeping - ): - chemistry.setdefault("fastchem_path", housekeeping["fastchem_directory"]) - clouds = raw.get("clouds") - if ( - isinstance(clouds, dict) - and clouds.get("model") == "mie_catalog" - and "optical_constants_directory" in housekeeping - ): - clouds.setdefault( - "optical_constants_path", housekeeping["optical_constants_directory"] - ) + section.setdefault(keys[-1], path_config[source_key]) + _fill_regional_input_paths(raw, path_config) + + +def _fill_regional_input_paths(raw: dict, path_config: dict) -> None: + atmosphere_cloud_pairs = [ + (raw.get("atmosphere", {}), raw.get("clouds")), + ] + disk = raw.get("disk_emission", {}) + if isinstance(disk, dict): + for name in ("hot_region", "cold_region"): + region = disk.get(name, {}) + if isinstance(region, dict): + atmosphere_cloud_pairs.append( + (region.get("atmosphere", {}), region.get("clouds")) + ) + for atmosphere, clouds in atmosphere_cloud_pairs: + if isinstance(atmosphere, dict): + chemistry = atmosphere.get("chemistry", {}) + if ( + isinstance(chemistry, dict) + and chemistry.get("model") == "fastchem_equilibrium" + and path_config.get("fastchem_directory") is not None + ): + chemistry.setdefault("fastchem_path", path_config["fastchem_directory"]) + if ( + isinstance(clouds, dict) + and clouds.get("model") == "mie_catalog" + and path_config.get("optical_constants_directory") is not None + ): + clouds.setdefault( + "optical_constants_path", path_config["optical_constants_directory"] + ) def _load_yaml_mapping(source: Path, ancestors: tuple[Path, ...] = ()) -> dict: @@ -901,6 +1132,11 @@ def _load_yaml_mapping(source: Path, ancestors: tuple[Path, ...] = ()) -> dict: if not parent.is_absolute() else parent.resolve() ) + if parent == source: + raise ValueError( + f"configuration extends itself: {source} declares extends: {extends!r}; " + "a self-contained run configuration must not contain an extends entry" + ) if not parent.is_file(): raise FileNotFoundError(parent) return _deep_merge(_load_yaml_mapping(parent, (*ancestors, source)), raw) diff --git a/src/robert_exoplanets/opacity/correlated_k.py b/src/robert_exoplanets/opacity/correlated_k.py index b61b3f7..b5f1afb 100644 --- a/src/robert_exoplanets/opacity/correlated_k.py +++ b/src/robert_exoplanets/opacity/correlated_k.py @@ -281,12 +281,6 @@ def from_exomol_cross_section_hdf( modify the source molecular cross sections. """ - try: - import h5py - except ImportError as exc: # pragma: no cover - dependency error path - raise RobertValidationError( - "loading ExoMol cross sections requires h5py" - ) from exc if ( isinstance(g_points, bool) or int(g_points) != g_points @@ -297,6 +291,12 @@ def from_exomol_cross_section_hdf( raise RobertValidationError( "ExoMol cross-section correlation requires spectral bin edges" ) + try: + import h5py + except ImportError as exc: # pragma: no cover - dependency error path + raise RobertValidationError( + "loading ExoMol cross sections requires h5py" + ) from exc source = Path(path).expanduser().resolve() edge_grid = SpectralGrid( values=spectral_grid.bin_edges, diff --git a/src/robert_exoplanets/postprocessing.py b/src/robert_exoplanets/postprocessing.py index d23b0f9..a3abcf6 100644 --- a/src/robert_exoplanets/postprocessing.py +++ b/src/robert_exoplanets/postprocessing.py @@ -52,6 +52,10 @@ def postprocess_retrieval_output( image_format: str = "png", dpi: int = 180, max_posterior_samples: int = 20_000, + posterior_predictive_samples: int = 200, + posterior_predictive_seed: int = 0, + corner_max_parameters: int = 20, + native_spectrum_model: object | None = None, leave_one_out: bool = False, loo_max_posterior_draws: int = 2_000, loo_seed: int = 0, @@ -62,12 +66,18 @@ def postprocess_retrieval_output( _validate_plot_options(style, image_format, dpi) if max_posterior_samples < 1: raise RobertValidationError("max_posterior_samples must be positive") + if posterior_predictive_samples < 1: + raise RobertValidationError("posterior_predictive_samples must be positive") result_path = Path(result_dir).expanduser() summary = _read_json(result_path / "result.json") arrays = _read_npz(result_path / "result_arrays.npz") names = tuple(str(name) for name in summary.get("parameter_names", ())) if not names: raise RobertDataError(f"retrieval result has no parameter names: {result_path}") + if names != problem.parameter_names: + raise RobertDataError( + "retrieval result parameter order does not match the configured problem" + ) best = _float_mapping(summary.get("best_fit_parameters"), "best_fit_parameters") spectra = problem.model_spectra(best) diagnostics = calculate_fit_statistics( @@ -96,6 +106,31 @@ def postprocess_retrieval_output( ) posterior = posterior_summary(names, arrays) diagnostics["posterior"] = posterior + posterior_draws = _posterior_parameter_draws( + names, + arrays, + maximum=posterior_predictive_samples, + seed=posterior_predictive_seed, + ) + bounds = np.asarray(problem.parameters.bounds, dtype=float) + posterior_draws = np.clip( + posterior_draws, + bounds[:, 0], + bounds[:, 1], + ) + spectral_quantiles = _posterior_spectral_quantiles(problem, posterior_draws) + native_quantiles = _posterior_native_spectral_quantiles( + problem, + posterior_draws, + native_spectrum_model, + ) + temperature_quantiles = _posterior_temperature_quantiles( + problem, + posterior_draws, + ) + diagnostics["posterior_predictive_draws"] = int(posterior_draws.shape[0]) + diagnostics["native_opacity_spectrum"] = native_quantiles is not None + diagnostics["temperature_profile_regions"] = tuple(temperature_quantiles) destination = Path(plot_dir).expanduser() destination.mkdir(parents=True, exist_ok=True) @@ -147,7 +182,16 @@ def postprocess_retrieval_output( dataset_colors=dataset_colors, style=style, dpi=dpi, + posterior_quantiles=spectral_quantiles, + native_quantiles=native_quantiles, ) + if temperature_quantiles: + _plot_temperature_profiles( + temperature_quantiles, + destination / f"temperature_profiles.{image_format}", + style=style, + dpi=dpi, + ) _plot_parameters( names, arrays, @@ -157,6 +201,8 @@ def postprocess_retrieval_output( image_format=image_format, dpi=dpi, max_samples=max_posterior_samples, + seed=posterior_predictive_seed, + corner_max_parameters=corner_max_parameters, ) _write_plot_manifest( destination, @@ -385,6 +431,173 @@ def weighted_quantile( return np.interp(probability, cumulative, sorted_data) +def _posterior_parameter_draws( + names: Sequence[str], + arrays: Mapping[str, np.ndarray], + *, + maximum: int, + seed: int, +) -> np.ndarray: + rng = np.random.default_rng(seed) + if "samples" in arrays: + samples = np.asarray(arrays["samples"], dtype=float) + if samples.ndim != 2 or samples.shape[1] != len(names): + raise RobertDataError("nested posterior samples do not match parameter names") + weights = _normalized_weights(arrays.get("weights"), samples.shape[0]) + count = min(maximum, max(1, samples.shape[0])) + indices = rng.choice(samples.shape[0], size=count, replace=True, p=weights) + return np.asarray(samples[indices], dtype=float) + if "state_vector" in arrays and "covariance" in arrays: + state = np.asarray(arrays["state_vector"], dtype=float) + covariance = np.asarray(arrays["covariance"], dtype=float) + return np.asarray( + rng.multivariate_normal( + state, + covariance, + size=maximum, + check_valid="raise", + ), + dtype=float, + ) + raise RobertDataError("result arrays contain neither nested samples nor OE state") + + +def _posterior_spectral_quantiles( + problem: MultiDatasetRetrievalProblem, + draws: np.ndarray, +) -> dict[str, np.ndarray]: + predictions: dict[str, list[np.ndarray]] = { + name: [] for name in problem.observations.names + } + for draw in draws: + spectra = problem.model_spectra(draw) + for name in predictions: + predictions[name].append(np.asarray(spectra[name].values, dtype=float)) + return { + name: np.quantile(np.stack(values), (0.16, 0.5, 0.84), axis=0) + for name, values in predictions.items() + } + + +def _posterior_native_spectral_quantiles( + problem: MultiDatasetRetrievalProblem, + draws: np.ndarray, + model: object | None, +) -> dict[str, np.ndarray | str] | None: + if model is None: + return None + spectra = [model(problem.parameter_mapping(draw)) for draw in draws] + if not spectra or any(not isinstance(spectrum, Spectrum) for spectrum in spectra): + raise RobertValidationError("native spectrum model must return Spectrum") + reference = spectra[0] + if any( + spectrum.unit != reference.unit + or spectrum.observable != reference.observable + or not np.array_equal( + spectrum.spectral_grid.values, + reference.spectral_grid.values, + ) + for spectrum in spectra[1:] + ): + raise RobertValidationError( + "native posterior spectra must share one grid, unit, and observable" + ) + return { + "wavelength": np.asarray(reference.spectral_grid.values, dtype=float), + "quantiles": np.quantile( + np.stack([spectrum.values for spectrum in spectra]), + (0.16, 0.5, 0.84), + axis=0, + ), + "unit": reference.unit, + } + + +def _posterior_temperature_quantiles( + problem: MultiDatasetRetrievalProblem, + draws: np.ndarray, +) -> dict[str, dict[str, np.ndarray]]: + builders = _named_atmosphere_builders(problem.forward_model) + output = {} + for name, builder in builders.items(): + profiles = [] + for draw in draws: + parameters = problem.parameter_mapping(draw) + profiles.append( + np.asarray( + builder.temperature_profile.evaluate( + parameters, + builder.pressure_grid, + ), + dtype=float, + ) + ) + output[name] = { + "pressure": np.asarray(builder.pressure_grid.centers, dtype=float), + "quantiles": np.quantile( + np.stack(profiles), + (0.16, 0.5, 0.84), + axis=0, + ), + } + return output + + +def _named_atmosphere_builders(model: object) -> dict[str, object]: + if hasattr(model, "hot_model") and hasattr(model, "cold_model"): + output = {} + for prefix, regional in ( + ("hot", model.hot_model), + ("cold", model.cold_model), + ): + nested = _named_atmosphere_builders(regional) + for name, builder in nested.items(): + output[prefix if name == "primary" else f"{prefix}_{name}"] = builder + return output + if hasattr(model, "emission_model"): + return _named_atmosphere_builders(model.emission_model) + builder = getattr(model, "atmosphere_builder", None) + if builder is not None: + return {"primary": builder} + models = getattr(model, "models", None) + if isinstance(models, Mapping) and models: + return _named_atmosphere_builders(next(iter(models.values()))) + return {} + + +def _plot_temperature_profiles( + profiles: Mapping[str, Mapping[str, np.ndarray]], + output: Path, + *, + style: str, + dpi: int, +) -> None: + plt = _pyplot() + with plt.style.context(style): + figure, axis = plt.subplots(figsize=(6.5, 7.0)) + for index, (name, values) in enumerate(profiles.items()): + pressure = values["pressure"] + quantiles = values["quantiles"] + color = DEFAULT_DATASET_COLORS[index % len(DEFAULT_DATASET_COLORS)] + axis.fill_betweenx( + pressure, + quantiles[0], + quantiles[2], + color=color, + alpha=0.22, + ) + axis.plot(quantiles[1], pressure, color=color, label=f"{name} median") + axis.set_yscale("log") + axis.invert_yaxis() + axis.set_xlabel("Temperature (K)") + axis.set_ylabel("Pressure (bar)") + axis.set_title("Posterior temperature-pressure profile (68% interval)") + axis.legend(fontsize=8) + figure.tight_layout() + figure.savefig(output, dpi=dpi, bbox_inches="tight") + plt.close(figure) + + def _plot_fit( problem: MultiDatasetRetrievalProblem, spectra: Mapping[str, Spectrum], @@ -395,6 +608,8 @@ def _plot_fit( dataset_colors: Mapping[str, str] | None, style: str, dpi: int, + posterior_quantiles: Mapping[str, np.ndarray] | None = None, + native_quantiles: Mapping[str, np.ndarray | str] | None = None, ) -> None: plt = _pyplot() colors = _dataset_color_mapping(problem, dataset_colors) @@ -410,6 +625,26 @@ def _plot_fit( gridspec_kw={"height_ratios": (3, 1)}, ) flux_label = "Model and observation" + if native_quantiles is not None: + native_wavelength = np.asarray(native_quantiles["wavelength"], dtype=float) + native_values = np.asarray(native_quantiles["quantiles"], dtype=float) + native_scale, flux_label = _plot_scale(str(native_quantiles["unit"])) + fit_axis.fill_between( + native_wavelength, + native_values[0] * native_scale, + native_values[2] * native_scale, + color="0.35", + alpha=0.18, + linewidth=0.0, + label="Native-grid 68% envelope", + ) + fit_axis.plot( + native_wavelength, + native_values[1] * native_scale, + color="0.15", + linewidth=1.15, + label="Native-opacity-grid median", + ) for dataset in problem.observations.datasets: name = dataset.name observation = dataset.observation @@ -431,7 +666,43 @@ def _plot_fit( alpha=0.75, label=observation.instrument or name, ) - fit_axis.plot(wavelength, model * scale, color=color, linewidth=1.6) + quantiles = ( + None + if posterior_quantiles is None + else posterior_quantiles.get(name) + ) + if quantiles is not None and quantiles.shape[1] == valid.size: + quantiles = quantiles[:, valid] + plotted_model = model + if quantiles is not None: + plotted_model = quantiles[1] + fit_axis.fill_between( + wavelength, + quantiles[0] * scale, + quantiles[2] * scale, + color=color, + alpha=0.2, + linewidth=0.0, + label=( + "68% posterior envelope" + if dataset is problem.observations.datasets[0] + else None + ), + ) + fit_axis.plot( + wavelength, + plotted_model * scale, + linestyle="none", + marker="s", + markersize=4.0, + markerfacecolor="none", + markeredgecolor=color, + label=( + "Observation-grid model" + if dataset is problem.observations.datasets[0] + else None + ), + ) residual_axis.plot( wavelength, (data - model) / uncertainty, @@ -465,6 +736,8 @@ def _plot_parameters( image_format: str, dpi: int, max_samples: int, + seed: int, + corner_max_parameters: int, ) -> None: plt = _pyplot() labels = {name: name for name in names} @@ -515,6 +788,20 @@ def _plot_parameters( style=style, dpi=dpi, ) + if len(names) <= corner_max_parameters: + _plot_corner( + _posterior_parameter_draws( + names, + arrays, + maximum=max_samples, + seed=seed, + ), + names, + labels, + output_dir / f"posterior_corner.{image_format}", + style=style, + dpi=dpi, + ) return if "state_vector" in arrays and "covariance" in arrays: state = np.asarray(arrays["state_vector"], dtype=float) @@ -545,6 +832,64 @@ def _plot_parameters( ) +def _plot_corner( + samples: np.ndarray, + names: Sequence[str], + labels: Mapping[str, str], + output: Path, + *, + style: str, + dpi: int, +) -> None: + plt = _pyplot() + count = len(names) + with plt.style.context(style): + figure, axes = plt.subplots( + count, + count, + figsize=(max(3.2, 2.15 * count), max(3.2, 2.15 * count)), + squeeze=False, + ) + for row in range(count): + for column in range(count): + axis = axes[row, column] + if row < column: + axis.set_visible(False) + continue + if row == column: + axis.hist( + samples[:, column], + bins=35, + color="#20639b", + alpha=0.8, + density=True, + ) + axis.set_yticks([]) + else: + axis.plot( + samples[:, column], + samples[:, row], + ".", + color="#20639b", + alpha=min(0.35, max(0.03, 150.0 / samples.shape[0])), + markersize=1.4, + rasterized=True, + ) + if row == count - 1: + axis.set_xlabel(labels[names[column]], fontsize=8) + else: + axis.set_xticklabels([]) + if column == 0 and row > 0: + axis.set_ylabel(labels[names[row]], fontsize=8) + elif column > 0: + axis.set_yticklabels([]) + axis.tick_params(labelsize=7) + figure.suptitle("Posterior corner plot", y=1.0) + figure.tight_layout() + figure.savefig(output, dpi=dpi, bbox_inches="tight") + plt.close(figure) + + def _plot_correlation( covariance: np.ndarray, names: Sequence[str], diff --git a/tests/test_configured_tasks.py b/tests/test_configured_tasks.py index 6ff095e..1804ab8 100644 --- a/tests/test_configured_tasks.py +++ b/tests/test_configured_tasks.py @@ -2,24 +2,64 @@ from __future__ import annotations +from copy import deepcopy from pathlib import Path from types import SimpleNamespace import numpy as np +import pytest from robert_exoplanets import ( CorrelatedKTable, + MultiDatasetTwoRegionEmissionModel, Observation, ObservationCollection, ObservationDataset, ) from robert_exoplanets.io import configured_tasks -from robert_exoplanets.io.task_config import load_task_config +from robert_exoplanets.io.task_config import TaskConfig, load_task_config ROOT = Path(__file__).resolve().parents[1] +def test_smoke_evaluation_can_validate_an_explicit_oe_state() -> None: + state = np.asarray([1.0, 2.0]) + + class StubProblem: + ndim = 2 + invalid_loglike = -1.0e100 + + def prior_transform(self, cube): + raise AssertionError("the midpoint must not be used") + + def log_likelihood_from_vector(self, vector): + np.testing.assert_array_equal(vector, state) + return -3.0 + + smoke = configured_tasks.smoke_evaluation(StubProblem(), state) + + assert smoke["log_likelihood"] == -3.0 + + +def test_smoke_evaluation_exposes_the_underlying_model_error() -> None: + class StubProblem: + ndim = 1 + invalid_loglike = -1.0e100 + + def prior_transform(self, cube): + return cube + + def log_likelihood_from_vector(self, vector): + return self.invalid_loglike + + def gaussian_inputs_from_vector(self, vector): + raise ValueError("temperature is outside opacity coverage") + + with pytest.raises(ValueError, match="outside opacity coverage"): + configured_tasks.smoke_evaluation(StubProblem()) + + def _table(species: str) -> CorrelatedKTable: return CorrelatedKTable( species=species, @@ -75,3 +115,84 @@ def capture_factories(factories, *, spectral_grids): assert problem.name == config.run.name assert captured_factories["f322w2"].opacity_free_species == () + + +def test_configured_two_region_problem_builds_independent_columns(monkeypatch) -> None: + base = load_task_config( + ROOT / "configurations" / "wasp69b_cloud_free_R1000.yaml" + ) + raw = deepcopy(base.model_dump(mode="python")) + species = tuple(raw["opacity"]["species"]) + raw["disk_emission"] = { + "model": "two_region", + "hot_fraction_parameter": "hot_area_fraction", + "cold_region": { + "atmosphere": { + "temperature": {"model": "isothermal", "temperature_k": 900.0}, + "chemistry": { + "model": "free", + "species": species, + "fixed_mixing_ratios": {name: 1.0e-8 for name in species}, + }, + }, + "clouds": {"model": "deck_haze"}, + }, + } + raw["parameters"] = ( + *raw["parameters"], + { + "name": "hot_area_fraction", + "prior": {"type": "uniform", "lower": 0.0, "upper": 1.0}, + }, + *( + { + "name": name, + "prior": {"type": "uniform", "lower": -12.0, "upper": 4.0}, + } + for name in ( + "log_cloud_top_pressure_bar", + "log_cloud_optical_depth", + "log_haze_mass_extinction", + "haze_slope", + ) + ), + ) + config = TaskConfig.model_validate(raw) + observation = Observation.from_arrays( + wavelength=[2.0, 2.5], + flux=[1.0e-3, 1.1e-3], + uncertainty=[1.0e-4, 1.0e-4], + instrument="NIRCam", + ) + observations = ObservationCollection( + (ObservationDataset("f322w2", observation),) + ) + captured = [] + + monkeypatch.setattr( + configured_tasks, + "_load_cached_table", + lambda _config, _dataset, item: _table(item), + ) + monkeypatch.setattr(configured_tasks, "load_nemesispy_cia_table", lambda: None) + + class StubRegionalModel: + models = {"f322w2": SimpleNamespace(opacity_identifiers={})} + + def __call__(self, parameters): + return {} + + def capture(factories, *, spectral_grids): + captured.append(factories["f322w2"]) + return StubRegionalModel() + + monkeypatch.setattr(configured_tasks, "build_multi_dataset_emission_model", capture) + + problem = configured_tasks.build_problem(config, observations) + + assert isinstance(problem.forward_model, MultiDatasetTwoRegionEmissionModel) + assert len(captured) == 2 + assert captured[0].chemistry_model.__class__.__name__ == "FastChemEquilibriumChemistry" + assert captured[0].cloud_model is None + assert captured[1].chemistry_model.__class__.__name__ == "FreeChemistry" + assert captured[1].cloud_model.__class__.__name__ == "ParameterizedDeckHazeCloudModel" diff --git a/tests/test_create_run_directory.py b/tests/test_create_run_directory.py index 2d4bf62..4c848e8 100644 --- a/tests/test_create_run_directory.py +++ b/tests/test_create_run_directory.py @@ -111,7 +111,7 @@ def test_create_run_directory_refuses_to_mix_runs(tmp_path: Path) -> None: create_run_directory(**kwargs) -def test_create_run_directory_updates_housekeeping_writable_paths( +def test_create_run_directory_uses_top_level_paths_and_local_writable_defaults( tmp_path: Path, ) -> None: run_directory = create_run_directory( @@ -120,9 +120,9 @@ def test_create_run_directory_updates_housekeeping_writable_paths( ) config = load_task_config(run_directory / "configuration.yaml") - assert config.housekeeping is not None - assert config.housekeeping.output_directory == run_directory / "outputs" - assert ( - config.housekeeping.opacity_cache_directory == run_directory / "opacity_cache" - ) - assert config.housekeeping.scratch_directory == run_directory / "scratch" + assert config.paths is not None + assert config.paths.project_directory == run_directory + assert config.housekeeping is None + assert config.outputs.directory == run_directory / "outputs" + assert config.opacity.cache_directory == run_directory / "opacity_cache" + assert config.runtime.scratch_directory == run_directory / "scratch" diff --git a/tests/test_deferred_oe.py b/tests/test_deferred_oe.py index 21425ea..7318497 100644 --- a/tests/test_deferred_oe.py +++ b/tests/test_deferred_oe.py @@ -10,6 +10,7 @@ from pydantic import ValidationError from robert_exoplanets.atmosphere import SplineTemperatureProfile +from robert_exoplanets.atmosphere.builder import _edge_evaluation_grid from robert_exoplanets.io.configured_tasks import ( configured_temperature_prior_covariance, describe_config, @@ -20,6 +21,8 @@ _pressure_grid, _temperature_profile, layer_temperature_overrides, + temperature_prior_covariance, + temperature_state_overrides, ) @@ -32,6 +35,14 @@ / "configurations" / "wasp69b_mie_catalog_layer_by_layer_R1000_optimal_estimation.yaml" ) +PG14_NESTED_CONFIG = ( + ROOT / "configurations" / "wasp69b_cloud_free_native_pg14_R1000_multinest.yaml" +) +PG14_OE_CONFIG = ( + ROOT + / "configurations" + / "wasp69b_cloud_free_native_pg14_R1000_optimal_estimation.yaml" +) def test_layer_by_layer_oe_configuration_matches_pressure_grid() -> None: @@ -44,6 +55,7 @@ def test_layer_by_layer_oe_configuration_matches_pressure_grid() -> None: assert config.sampler.oe_temperature_correlation_length_dex == 1.5 assert "correlated_temperature_prior=250 K/1.5 dex" in describe_config(config) assert temperature.model == "spline" + assert temperature.extrapolation == "clip" assert len(temperature.knot_pressure) == config.atmosphere.pressure.layers == 80 assert temperature.parameter_names == tuple(f"T_{index:02d}" for index in range(80)) np.testing.assert_allclose( @@ -63,6 +75,13 @@ def test_layer_by_layer_oe_configuration_matches_pressure_grid() -> None: "alpha", }.intersection(item.name for item in config.parameters) + profile = _temperature_profile(config) + edge_temperature = profile.evaluate( + {name: 1500.0 for name in profile.required_parameters()}, + _edge_evaluation_grid(_pressure_grid(config)), + ) + np.testing.assert_allclose(edge_temperature, 1500.0) + def test_pg14_best_fit_initializes_every_oe_temperature_layer() -> None: nested_config = load_task_config(NESTED_CONFIG) @@ -100,6 +119,42 @@ def test_pg14_best_fit_initializes_every_oe_temperature_layer() -> None: ) +def test_same_pg14_parameterization_transfers_without_temperature_overrides() -> None: + nested_config = load_task_config(PG14_NESTED_CONFIG) + oe_config = load_task_config(PG14_OE_CONFIG) + target_profile = _temperature_profile(oe_config) + problem = SimpleNamespace( + forward_model=SimpleNamespace( + atmosphere_builder=SimpleNamespace(temperature_profile=target_profile) + ) + ) + + overrides = temperature_state_overrides(nested_config, {}, problem) + + assert overrides == {} + + +def test_same_pg14_parameterization_keeps_nested_temperature_covariance() -> None: + oe_config = load_task_config(PG14_OE_CONFIG) + profile = _temperature_profile(oe_config) + problem = SimpleNamespace( + forward_model=SimpleNamespace( + atmosphere_builder=SimpleNamespace(temperature_profile=profile) + ) + ) + covariance = np.eye(len(oe_config.parameters)) + + actual, names, pressure = temperature_prior_covariance( + oe_config, + problem, + covariance, + ) + + assert actual is covariance + assert names == profile.required_parameters() + assert pressure.size == 0 + + def test_log_pressure_temperature_covariance_has_requested_smoothing() -> None: covariance = log_pressure_correlated_covariance( [1.0, 10.0**1.5, 1000.0], diff --git a/tests/test_inhomogeneous_emission.py b/tests/test_inhomogeneous_emission.py index 4315c8c..f1792ca 100644 --- a/tests/test_inhomogeneous_emission.py +++ b/tests/test_inhomogeneous_emission.py @@ -8,6 +8,8 @@ from robert_exoplanets import ( DilutedEmissionModel, DiskEmissionModelConfig, + MultiDatasetDilutedEmissionModel, + MultiDatasetTwoRegionEmissionModel, SpectralGrid, Spectrum, TwoRegionEmissionModel, @@ -77,3 +79,28 @@ def test_disk_model_configuration_selects_existing_regional_hardware() -> None: np.testing.assert_allclose(one_region({}).values, [10.0, 20.0]) np.testing.assert_allclose(diluted({"dayside_dilution": 0.5}).values, [5.0, 10.0]) np.testing.assert_allclose(two_region({"hot_area_fraction": 0.5}).values, [6.0, 12.0]) + + +def test_multi_dataset_two_region_model_mixes_each_named_spectrum() -> None: + def hot(parameters): + return {"nircam": _constant([10.0, 20.0])(parameters)} + + def cold(parameters): + return {"nircam": _constant([2.0, 4.0])(parameters)} + + model = MultiDatasetTwoRegionEmissionModel(hot, cold) + + spectra = model({"hot_area_fraction": 0.25}) + + np.testing.assert_allclose(spectra["nircam"].values, [4.0, 8.0]) + + +def test_multi_dataset_dilution_scales_each_named_spectrum() -> None: + def regional(parameters): + return {"miri": _constant([10.0, 20.0])(parameters)} + + model = MultiDatasetDilutedEmissionModel(regional) + + spectra = model({"dayside_dilution": 0.4}) + + np.testing.assert_allclose(spectra["miri"].values, [4.0, 8.0]) diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index b2be949..7ba4c20 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -6,10 +6,16 @@ from pathlib import Path import subprocess import sys +from types import SimpleNamespace import numpy as np -from robert_exoplanets import Observation, Spectrum +from robert_exoplanets import ( + IsothermalTemperatureProfile, + Observation, + PressureGrid, + Spectrum, +) from robert_exoplanets.instruments import ObservationCollection, ObservationDataset from robert_exoplanets.postprocessing import ( calculate_fit_statistics, @@ -43,16 +49,24 @@ def _problem() -> MultiDatasetRetrievalProblem: (ObservationDataset("synthetic", observation),) ) - def forward(parameters): - level = parameters["level"] - return { - "synthetic": Spectrum.from_arrays( - observation.wavelength, - level * observation.wavelength, - unit=observation.flux_unit, - observable=observation.observable, - ) - } + class Forward: + atmosphere_builder = SimpleNamespace( + temperature_profile=IsothermalTemperatureProfile( + parameter_name="level" + ), + pressure_grid=PressureGrid.logspace(1.0e-5, 10.0, 5, unit="bar"), + ) + + def __call__(self, parameters): + level = parameters["level"] + return { + "synthetic": Spectrum.from_arrays( + observation.wavelength, + level * observation.wavelength, + unit=observation.flux_unit, + observable=observation.observable, + ) + } return MultiDatasetRetrievalProblem( name="synthetic-postprocessing", @@ -60,7 +74,7 @@ def forward(parameters): parameters=RetrievalParameterSet( (RetrievalParameter("level", UniformPrior(0.5, 1.5)),) ), - forward_model=forward, + forward_model=Forward(), ) @@ -118,6 +132,12 @@ def test_retrieval_postprocessing_writes_statistics_and_plots(tmp_path: Path) -> problem, result_dir, plot_dir=plot_dir, + native_spectrum_model=lambda parameters: Spectrum.from_arrays( + np.linspace(1.0, 3.0, 9), + parameters["level"] * np.linspace(1.0, 3.0, 9), + unit="relative_flux", + observable="relative_flux", + ), ) assert diagnostics["reduced_chi_squared"] == 0.0 @@ -127,6 +147,10 @@ def test_retrieval_postprocessing_writes_statistics_and_plots(tmp_path: Path) -> assert (plot_dir / "fit_spectrum_residuals.png").is_file() assert (plot_dir / "posterior_marginals.png").is_file() assert (plot_dir / "parameter_correlation.png").is_file() + assert (plot_dir / "posterior_corner.png").is_file() + assert (plot_dir / "temperature_profiles.png").is_file() + assert diagnostics["posterior_predictive_draws"] == 3 + assert diagnostics["native_opacity_spectrum"] is True assert discover_retrieval_result_directories(tmp_path / "outputs") == (result_dir,) diff --git a/tests/test_task_config.py b/tests/test_task_config.py index 6181841..adf2605 100644 --- a/tests/test_task_config.py +++ b/tests/test_task_config.py @@ -7,6 +7,7 @@ import numpy as np import pytest +import yaml from pydantic import ValidationError from robert_exoplanets.io.configured_tasks import ( @@ -15,6 +16,7 @@ ) from robert_exoplanets.io.task_config import ( TaskConfig, + configured_regions, initialize_task_directories, load_task_config, ) @@ -68,6 +70,19 @@ def test_wasp69b_example_exposes_complete_native_mode_run() -> None: assert config.bodies.star.metallicity_dex == 0.0 +def test_yaml_rejects_direct_self_extension_with_actionable_message( + tmp_path: Path, +) -> None: + source = tmp_path / "configuration.yaml" + source.write_text("extends: configuration.yaml\n", encoding="utf-8") + + with pytest.raises( + ValueError, + match="configuration extends itself.*must not contain an extends entry", + ): + load_task_config(source) + + def test_yaml_can_select_blackbody_stellar_spectrum() -> None: config = load_task_config(EXAMPLE) raw = deepcopy(config.model_dump(mode="python")) @@ -78,6 +93,124 @@ def test_yaml_can_select_blackbody_stellar_spectrum() -> None: assert parsed.bodies.star.spectrum_model == "blackbody" +def _two_region_raw() -> dict: + config = load_task_config(EXAMPLE) + raw = deepcopy(config.model_dump(mode="python")) + cold_species = list(raw["opacity"]["species"]) + raw["disk_emission"] = { + "model": "two_region", + "hot_fraction_parameter": "hot_area_fraction", + "cold_region": { + "atmosphere": { + "temperature": {"model": "isothermal", "temperature_k": 900.0}, + "chemistry": { + "model": "free", + "species": cold_species, + "fixed_mixing_ratios": { + species: 1.0e-8 for species in cold_species + }, + "background_species": ["H2", "He"], + }, + }, + "clouds": { + "model": "deck_haze", + "log10_cloud_top_pressure_bar_parameter": "cold_cloud_top", + "log10_cloud_optical_depth_parameter": "cold_cloud_tau", + "log10_haze_mass_extinction_parameter": "cold_haze_extinction", + "haze_slope_parameter": "cold_haze_slope", + }, + }, + } + raw["parameters"] = [ + *raw["parameters"], + { + "name": "hot_area_fraction", + "prior": {"type": "uniform", "lower": 0.0, "upper": 1.0}, + }, + *( + { + "name": name, + "prior": {"type": "uniform", "lower": -12.0, "upper": 4.0}, + } + for name in ( + "cold_cloud_top", + "cold_cloud_tau", + "cold_haze_extinction", + "cold_haze_slope", + ) + ), + ] + return raw + + +def test_yaml_resolves_independent_two_region_physics() -> None: + parsed = TaskConfig.model_validate(_two_region_raw()) + + hot, cold = configured_regions(parsed) + + assert parsed.disk_emission.model == "two_region" + assert hot.atmosphere.temperature.model == "parmentier_guillot_2014" + assert hot.atmosphere.chemistry.model == "fastchem_equilibrium" + assert hot.clouds.model == "none" + assert cold.atmosphere.pressure is parsed.atmosphere.pressure + assert cold.atmosphere.temperature.model == "isothermal" + assert cold.atmosphere.chemistry.model == "free" + assert cold.clouds.model == "deck_haze" + + +def test_two_region_yaml_requires_the_area_fraction_parameter() -> None: + raw = _two_region_raw() + raw["parameters"] = [ + item for item in raw["parameters"] if item["name"] != "hot_area_fraction" + ] + + with pytest.raises(ValidationError, match="hot_area_fraction"): + TaskConfig.model_validate(raw) + + +def test_yaml_accepts_diluted_emission_and_rejects_it_for_transmission() -> None: + config = load_task_config(EXAMPLE) + raw = deepcopy(config.model_dump(mode="python")) + raw["disk_emission"] = { + "model": "diluted_one_region", + "dilution_parameter": "dayside_dilution", + } + raw["parameters"] = ( + *raw["parameters"], + { + "name": "dayside_dilution", + "prior": {"type": "uniform", "lower": 0.0, "upper": 1.0}, + }, + ) + + parsed = TaskConfig.model_validate(raw) + assert parsed.disk_emission.model == "diluted_one_region" + + raw["radiative_transfer"]["model"] = "transmission" + with pytest.raises(ValidationError, match="require.*emission"): + TaskConfig.model_validate(raw) + + +def test_yaml_defaults_writable_paths_to_the_configuration_directory( + tmp_path: Path, +) -> None: + config = load_task_config(EXAMPLE) + raw = deepcopy(config.model_dump(mode="json")) + raw.pop("paths", None) + raw.pop("housekeeping", None) + raw.pop("outputs") + raw["runtime"].pop("scratch_directory") + raw["opacity"].pop("cache_directory") + source = tmp_path / "configuration.yaml" + source.write_text(yaml.safe_dump(raw, sort_keys=False), encoding="utf-8") + + parsed = load_task_config(source) + + assert parsed.outputs.directory == tmp_path / "outputs" + assert parsed.runtime.scratch_directory == tmp_path / "scratch" + assert parsed.opacity.cache_directory == tmp_path / "opacity_cache" + + def test_yaml_configures_transmission_and_real_exomol_h2o() -> None: config = load_task_config(TRANSMISSION) @@ -481,20 +614,21 @@ def test_direct_nk_default_replaces_catalogue_cloud_fields() -> None: assert len(config.clouds.real_index_parameter_names) == 6 -def test_complete_template_uses_housekeeping_for_internal_paths() -> None: +def test_complete_template_uses_one_top_level_path_block() -> None: config = load_task_config(TEMPLATE) assert config.clouds.model == "none" - assert config.housekeeping is not None - assert config.observations.path == config.housekeeping.observations_directory + assert config.paths is not None + assert config.housekeeping is None + assert config.observations.path == config.paths.observations_directory assert ( config.atmosphere.chemistry.fastchem_path - == config.housekeeping.fastchem_directory + == config.paths.fastchem_directory ) - assert config.opacity.path == config.housekeeping.k_table_directory - assert config.opacity.cache_directory == config.housekeeping.opacity_cache_directory - assert config.outputs.directory == config.housekeeping.output_directory - assert config.runtime.scratch_directory == config.housekeeping.scratch_directory + assert config.opacity.path == config.paths.k_table_directory + assert config.opacity.cache_directory == ROOT / "configurations" / "opacity_cache" + assert config.outputs.directory == ROOT / "configurations" / "outputs" + assert config.runtime.scratch_directory == ROOT / "configurations" / "scratch" assert config.plotting.enabled is False assert config.sampler.max_calls is None assert config.plotting.dataset_colors["f322w2"] == "#20639b"