Inference: config validation and failure handling

.gitignore

@@ -342,3 +342,10 @@ src/proteus/_version.py
 # Local scratch directories (not part of the project)
 /.playwright-mcp/
 /platon/
+
+# output logs
+*.out
+
+# test inference tomls
+input/test_infer_proteus/
+input/inference/test_infer

docs/How-to/inference.md

@@ -32,6 +32,7 @@ The system performs Bayesian optimization to infer planetary formation parameter
     | `async_BO.py`      | Parallel BO implementation                |
     | `BO.py`            | Single BO step implementation             |
     | `objective.py`     | PROTEUS interface and objective function  |
+    | `failures.py`      | Functions for handling failing simulations |
     | `plot.py`          | Visualization utilities                   |
     | `utils.py`         | Helper functions for inference scheme     |
     | `gen_D_init.py`    | Generate initial data                     |
@@ -122,6 +123,7 @@ The system generates several outputs in:
 - `logs.csv`: Detailed logs of each BO step
 - `Ts.csv`: Timestamps for performance analysis
 - `init.csv`: Data used as an initial guess for starting the optimisation
+- `failures.csv`: One row per simulation that failed or was excluded, written only when there is at least one (see [Failed and excluded simulations](#failed-and-excluded-simulations))
 
 ### Plots
 The BO scheme will generate many plots upon completion.
@@ -144,10 +146,15 @@ Plots prefixed with `result_` show the results of the optimisation.
 
 ### Results Summary
 The system prints the final results including:
+
 - Best found parameters
 - Corresponding simulated observables
 - Comparison with target observables
 
+### Failed and excluded simulations
+
+During the inference run, some PROTEUS simulations might crash or fail, or stop on a status that is excluded in the inference configuration (e.g. maximum runtime reached). The run carries on when there are failures unless `abort_on_failure` is set to `true` in the inference config. At the end of the study all failures are written to `failures.csv` in the output folder, and summarised. 
+
 ## Customization
 
 ### Adding New Parameters

docs/Reference/config/atmosphere.md

@@ -25,6 +25,7 @@ parameter selects the surface energy balance scheme (`mixed_layer`,
 | `module` | str | `"agni"` | Which atmosphere module to use. Choices: `"dummy"`, `"agni"`, `"janus"`. |
 | `spectral_group` | str | `"Honeyside"` | Spectral file group defining gas opacities. See https://proteus-framework.org/SOCRATES/Reference/proteus_spectral_file_reference.html. |
 | `spectral_bands` | str | `"48"` | Number of wavenumber bands in k-table. |
+| `spectral_cache` | str or none | `none` | Folder in which to reuse prepared spectral files across runs that share a stellar spectrum. None disables the cache and every run builds its own. |
 | `num_levels` | int | `50` | Number of vertical atmosphere levels. Must be >= 15. |
 | `p_top` | float | `1e-06` | Top-of-atmosphere pressure \[bar\]. Must be > 0. |
 | `p_obs` | float | `0.02` | Observation pressure level \[bar\] (transit radius). Must be > 0. |

docs/Reference/config/config_schema.json

@@ -8673,6 +8673,22 @@
       "group": null,
       "group_qualifier": null
     },
+    {
+      "path": "atmos_clim.spectral_cache",
+      "toml_section": "atmos_clim",
+      "class": "AtmosClim",
+      "type": "str or none",
+      "accepts_none": true,
+      "default": "none",
+      "choices": null,
+      "bounds": null,
+      "description": "Folder in which to reuse prepared spectral files across runs that share a stellar spectrum. None disables the cache and every run builds its own.",
+      "doc_source": "attributes",
+      "group_order": 0,
+      "group_position": 3,
+      "group": null,
+      "group_qualifier": null
+    },
     {
       "path": "atmos_clim.num_levels",
       "toml_section": "atmos_clim",
@@ -8690,7 +8706,7 @@
       "description": "Number of vertical atmosphere levels.",
       "doc_source": "attributes",
       "group_order": 0,
-      "group_position": 3,
+      "group_position": 4,
       "group": null,
       "group_qualifier": null
     },
@@ -8711,7 +8727,7 @@
       "description": "Top-of-atmosphere pressure [bar].",
       "doc_source": "attributes",
       "group_order": 0,
-      "group_position": 4,
+      "group_position": 5,
       "group": null,
       "group_qualifier": null
     },
@@ -8732,7 +8748,7 @@
       "description": "Observation pressure level [bar] (transit radius).",
       "doc_source": "attributes",
       "group_order": 0,
-      "group_position": 5,
+      "group_position": 6,
       "group": null,
       "group_qualifier": null
     },
@@ -8748,7 +8764,7 @@
       "description": "Gas overlap method. Choices: 'ro', 'rorr', 'ee'.",
       "doc_source": "attributes",
       "group_order": 0,
-      "group_position": 6,
+      "group_position": 7,
       "group": null,
       "group_qualifier": null
     },

input/all_options.toml

@@ -637,6 +637,7 @@ config_version = "3.0"
     # Grid and spectral setup (shared by agni + janus)
     spectral_group            = "Honeyside"  # opacity k-table set; see docs/assets/spectral_files.pdf
     spectral_bands            = "48"         # wavenumber bands in k-table
+    spectral_cache            = "none"       # reuse prepared spectral files across runs with the same star; none = off
     num_levels                = 50           # vertical atmosphere levels (min 15)
     p_top                     = 1.0e-6       # top-of-atmosphere pressure [bar]
     p_obs                     = 0.02         # observation pressure level [bar] (transit radius)

input/inference/example.infer.toml

@@ -10,16 +10,19 @@ logging = "INFO"
 output = "bayesopt_infer_se_mat12_lei_5_100/"
 
 # Path to base (reference) config file relative to PROTEUS root folder
-# ref_config = "output/bayesopt_se/init_coupler.toml"
 ref_config = "input/inference/example.toml"
 
 # Method for initialising the inference scheme (one of these must be 'none')
 init_samps = 3        # Number of random samples if starting from scratch.
 init_grid  = 'none' # grid_demo/'   # Path pre-computed grid (relative to PROTEUS output folder)
 
-# Excluded exit codes (in addition to errors)
+# Completion codes to exclude from the fit.
 failure_codes = [11, ] # solidified (10), escaped (15), max_runtime (11)
 
+# Stop the whole inference run at the first simulation that fails, instead of scoring
+# it as a poor sample and carrying on. Off by default.
+abort_on_failure = false
+
 # Parameters for Bayesian optimisation
 n_workers  = 5        # Number of parallel workers
 kernel     = "MAT1/2"    # Kernel type for GP, "RBF" | "MAT1/2" | "MAT3/2" | "MAT5/2"

src/proteus/atmos_clim/agni.py

@@ -12,6 +12,7 @@
 from scipy.interpolate import PchipInterpolator
 
 from proteus.atmos_clim.common import clip_radius_to_hill, get_oarr_from_parr, get_spfile_path
+from proteus.atmos_clim.spectral_cache import cache_key, seed_from_cache, store_in_cache
 from proteus.utils.constants import gas_list, noble_gases
 from proteus.utils.helper import (
     UpdateStatusfile,
@@ -484,6 +485,18 @@ def init_agni_atmos(dirs: dict, config: Config, hf_row: dict):
     # bypass the glob entirely so a missing or empty `data/*.sflux` directory
     # is not a precondition for those modes.
 
+    # Fast I/O folder. Decided before the spectral file, because AGNI writes the
+    # prepared runtime.sf pair here and so this is where the cache reads from.
+    if (config.atmos_clim.agni.verbosity >= 2) or (config.params.out.logging == 'DEBUG'):
+        io_dir = dirs['output']
+    else:
+        io_dir = create_tmp_folder()
+    log.info(f'Temporary-file working dir: {io_dir}')
+
+    # Set when this run built a prepared spectral file that the cache does not
+    # yet hold, so it can be stored once the build is known to have succeeded.
+    cache_store_key = None
+
     # Spectral file path provided?
     if config.atmos_clim.agni.spectral_file is not None:
         # Grey gas?
@@ -527,12 +540,21 @@ def init_agni_atmos(dirs: dict, config: Config, hf_row: dict):
         input_sf = get_spfile_path(dirs['fwl'], config)
         input_star = sflux_path
 
-    # Fast I/O folder
-    if (config.atmos_clim.agni.verbosity >= 2) or (config.params.out.logging == 'DEBUG'):
-        io_dir = dirs['output']
-    else:
-        io_dir = create_tmp_folder()
-    log.info(f'Temporary-file working dir: {io_dir}')
+        # Reuse a cached file built earlier from this base file and this stellar
+        # spectrum, and skip the insertion.
+        if config.atmos_clim.spectral_cache:
+            key = cache_key(
+                input_sf,
+                sflux_path,
+                config.atmos_clim.spectral_group,
+                config.atmos_clim.spectral_bands,
+            )
+            if seed_from_cache(config.atmos_clim.spectral_cache, key, io_dir):
+                log.debug('Reusing prepared spectral file from cache')
+                input_sf = os.path.join(io_dir, 'runtime.sf')
+                input_star = ''
+            else:
+                cache_store_key = key
 
     # composition
     vol_dict = _construct_voldict(config, hf_row, dirs)
@@ -665,6 +687,10 @@ def init_agni_atmos(dirs: dict, config: Config, hf_row: dict):
     # Confirm the live Atmos_t contains every field that PROTEUS expects
     _check_agni_schema(atmos, dirs)
 
+    # Stored spectral file is now valid, so store it in the cache if requested.
+    if cache_store_key:
+        store_in_cache(config.atmos_clim.spectral_cache, cache_store_key, io_dir)
+
     # Set temperature profile from old NetCDF if it exists
     nc_files = glob.glob(os.path.join(dirs['output'], 'data', '*_atm.nc'))
     if len(nc_files) > 0:

src/proteus/atmos_clim/spectral_cache.py

@@ -0,0 +1,154 @@
+"""Reuse of prepared spectral files across runs that share a stellar spectrum.
+
+A run's `runtime.sf` is the base spectral file from FWL_DATA with that run's
+stellar spectrum inserted. This module keeps one copy per distinct input set so the
+second and later runs copy it instead.
+
+The cache is seeded into the run's output folder as `runtime.sf`, which is
+where the atmosphere wrapper already expects to manage it.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import os
+import shutil
+from pathlib import Path
+
+log = logging.getLogger('fwl.' + __name__)
+
+# Prepared spectral files come in a pair, and the existence check in the AGNI
+# module only tests the first. Seeding one without the other would leave the
+# module using a file whose companion is absent.
+SPECTRAL_SUFFIXES = ('', '_k')
+
+# Bytes read per chunk when fingerprinting a file.
+_CHUNK = 1 << 20
+
+
+def _file_digest(path: Path) -> str:
+    """Hash a file's contents."""
+    digest = hashlib.sha256()
+    with open(path, 'rb') as f:
+        while chunk := f.read(_CHUNK):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def cache_key(base_sf: Path | str, star_spectrum: Path | str, group: str, bands: str) -> str:
+    """Name the cache entry for a prepared spectral file. The stellar spectrum is hashed by content.
+
+    Parameters
+    ----------
+    - base_sf (Path | str): Base spectral file, from FWL_DATA.
+    - star_spectrum (Path | str): Stellar spectrum (`.sflux`) to be inserted.
+    - group (str): Spectral file group.
+    - bands (str): Number of wavenumber bands.
+
+    Returns
+    ----------
+    - str: Hex digest naming this combination.
+    """
+    base = Path(base_sf)
+    stat = base.stat()
+    parts = (
+        group,
+        bands,
+        base.name,
+        str(stat.st_size),
+        str(int(stat.st_mtime)),
+        _file_digest(Path(star_spectrum)),
+    )
+    return hashlib.sha256('\0'.join(parts).encode()).hexdigest()[:32]
+
+
+def _entry_paths(cache_dir: Path | str, key: str) -> list[Path]:
+    """Paths of the cached pair for one key."""
+    return [Path(cache_dir) / f'{key}.sf{suffix}' for suffix in SPECTRAL_SUFFIXES]
+
+
+def _output_paths(out_dir: Path | str) -> list[Path]:
+    """Paths of the prepared pair inside a run's output folder."""
+    return [Path(out_dir) / f'runtime.sf{suffix}' for suffix in SPECTRAL_SUFFIXES]
+
+
+def seed_from_cache(cache_dir: Path | str, key: str, out_dir: Path | str) -> bool:
+    """Copy a cached spectral file pair into a run's output folder.
+
+    Parameters
+    ----------
+    - cache_dir (Path | str): Folder holding cached entries.
+    - key (str): Key from `cache_key`.
+    - out_dir (Path | str): The run's output folder.
+
+    Returns
+    ----------
+    - bool: True if the run can now use `runtime.sf` as-is. False on a miss, or
+      on any error: a cache that cannot be read is a slower run, never a failed
+      one, so the caller falls back to building the file.
+    """
+    entries = _entry_paths(cache_dir, key)
+    if not all(entry.is_file() for entry in entries):
+        return False
+
+    targets = _output_paths(out_dir)
+    try:
+        Path(out_dir).mkdir(parents=True, exist_ok=True)
+        for entry, target in zip(entries, targets):
+            shutil.copyfile(entry, target)
+    except OSError as err:
+        log.warning(f'Could not seed spectral file from cache: {err}')
+        # A half-copied pair would be read as a prepared file. Clear it so the
+        # run rebuilds from scratch instead of starting from a truncated file.
+        for target in targets:
+            try:
+                target.unlink(missing_ok=True)
+            except OSError as cleanup_err:
+                log.debug(f'Could not remove partial spectral file {target}: {cleanup_err}')
+        return False
+
+    log.debug(f'Seeded spectral file from cache entry {key}')
+    return True
+
+
+def store_in_cache(cache_dir: Path | str, key: str, out_dir: Path | str) -> bool:
+    """Copy a run's prepared spectral file pair into the cache.
+
+    Written to a temporary name and renamed into place, so concurrent runs
+    racing to populate the same key never expose a half-written entry.
+
+    Parameters
+    ----------
+    - cache_dir (Path | str): Folder holding cached entries.
+    - key (str): Key from `cache_key`.
+    - out_dir (Path | str): The run's output folder, holding the prepared pair.
+
+    Returns
+    ----------
+    - bool: True if the entry was stored. False if the run produced no prepared
+      file, or the cache could not be written; neither is a fault of the run.
+    """
+    sources = _output_paths(out_dir)
+    if not all(source.is_file() for source in sources):
+        return False
+
+    try:
+        Path(cache_dir).mkdir(parents=True, exist_ok=True)
+        for source, entry in zip(sources, _entry_paths(cache_dir, key)):
+            tmp = entry.with_name(f'{entry.name}.{os.getpid()}.tmp')
+            try:
+                shutil.copyfile(source, tmp)
+                os.replace(tmp, entry)
+            except Exception:
+                try:
+                    tmp.unlink(missing_ok=True)
+                except OSError:
+                    pass
+                raise
+    except OSError as err:
+        log.warning(f'Could not store spectral file in cache: {err}')
+        return False
+
+    log.debug(f'Stored spectral file in cache as entry {key}')
+    return True