diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index b3317f4f..48cac627 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -34,6 +34,4 @@ jobs: pip install . - name: Test with pytest run: | - pytest tests/test_sfh.py tests/test_pipeline_module.py tests/test_run.py - # make test-units - # make test-notebooks + pytest tests/test_* diff --git a/README.md b/README.md index 16de27b0..82606ad9 100644 --- a/README.md +++ b/README.md @@ -3,41 +3,126 @@ [![Documentation Status](https://readthedocs.org/projects/besta/badge/?version=latest)](https://besta.readthedocs.io/en/latest/?badge=latest) [![test](https://github.com/PabloCorcho/besta/actions/workflows/test.yml/badge.svg?branch=main)](https://github.com/PabloCorcho/besta/actions/workflows/test.yml) -BESTA is a Python package for inferring stellar population properties from spectroscopic and/or photometric data using Bayesian inference and Monte Carlo methods. +BESTA is a Python library for Bayesian inference of galaxy stellar population properties from spectroscopic and/or photometric observations. -Full documentation is available at [besta.readthedocs.io](https://besta.readthedocs.io). +It provides a practical interface to build fitting workflows on top of: -## Framework +- [Population Synthesis Toolkit (PST)](https://github.com/paranoya/population-synthesis-toolkit) for stellar population modeling. +- [CosmoSIS](https://cosmosis.readthedocs.io) for parameter estimation and sampling. -BESTA is built on top of two core libraries: +Full user and API documentation: [besta.readthedocs.io](https://besta.readthedocs.io) -- **[Population Synthesis Toolkit (PST)](https://github.com/paranoya/population-synthesis-toolkit)** — flexible stellar population synthesis models. -- **[CosmoSIS](https://cosmosis.readthedocs.io)** — modular parameter estimation via Monte Carlo sampling. +## What BESTA Can Do -## Features +BESTA supports end-to-end stellar population fitting workflows, including: -### Pipeline modules +- Spectroscopic and photometric inference workflows. +- Joint fitting of stellar population and kinematics. +- Configurable star formation history (SFH) parameterizations. +- Grid-based and emulator-based SFH inference options. +- Sequential and batch execution of fitting pipelines. +- Post-processing and visualization utilities for fitted solutions. -Ready-to-use CosmoSIS pipeline modules are provided in `besta.pipeline_modules`: +For detailed pipeline setup, module-level configuration, and examples, see the online documentation: -| Module | Description | -|---|---| -| `GalaxySpectraModule` | Full spectroscopic SED fitting | -| `FullSpectralFitModule` | Joint stellar population and kinematics fit from spectra | -| `GalaxyPhotometryModule` | Broadband photometric SED fitting | -| `SFHPhotometryGridModule` | SFH inference from photometry via model grid | -| `SFHPhotometryEmulatorModule` | SFH inference from photometry via emulator | +- [Quick guide](https://besta.readthedocs.io/en/latest/quick_guide.html) +- [Configuration reference](https://besta.readthedocs.io/en/latest/configuration.html) +- [Pipeline manager docs](https://besta.readthedocs.io/en/latest/manager.html) +- [Tutorials](https://besta.readthedocs.io/en/latest/index.html) -### Star formation history models +## Installation -Analytic: `ExponentialSFH`, `DelayedTauSFH`, `DelayedTauQuenchedSFH`, `LogNormalSFH`, `LogNormalQuenchedSFH` +### Python version -Piece-wise: `FixedTimeSFH`, `FixedTime_sSFR_SFH`, `FixedMassFracSFH` +- `Python >= 3.10` -## Installation +### Install from PyPI ```bash pip install besta ``` -Requires Python ≥ 3.10. See the [installation guide](https://besta.readthedocs.io) for CosmoSIS setup instructions. +### Install from source + +Depending on your platform and scientific Python installation, you may also need system libraries typically used by MPI and linear algebra backends (for example OpenMPI, BLAS/LAPACK, and a Fortran compiler). + +```bash +git clone https://github.com/PabloCorcho/besta.git +cd besta +python -m pip install -r requirements.txt +python -m pip install . +``` + +### Ubuntu/Debian example: + +```bash +sudo apt update +sudo apt install -y \ + gfortran \ + liblapack-dev \ + libopenblas-dev \ + openmpi-bin \ + openmpi-common \ + libopenmpi-dev \ + libgtk2.0-dev +``` + +using conda + +```bash +conda install -c conda-forge \ + gcc_linux-64 \ + gxx_linux-64 \ + gfortran_linux-64 \ + openmpi \ + mpi4py \ + gsl \ + cfitsio \ + fftw \ + lapack \ + openblas \ + git \ + make \ + cmake \ + pkg-config +``` + +Package names can vary by distribution and version. If your environment already provides BLAS/LAPACK and MPI through conda, you may not need to install all system-level packages. + +## Contributing + +Contributions are welcome, including bug fixes, new features, tests, docs, and tutorials. + +### Report issues + +- Open issues or feature requests at: [github.com/PabloCorcho/besta/issues](https://github.com/PabloCorcho/besta/issues) + +### Contribute code + +1. Fork the repository. +2. Create a feature branch. +3. Implement your changes and add or update tests. +4. Update documentation when behavior or APIs change. +5. Open a pull request with a clear summary and motivation. + +### Development quick start + +```bash +git clone https://github.com/PabloCorcho/besta.git +cd besta +python -m pip install -r requirements.txt +python -m pip install -e . +pytest -q +``` + +## Citation + +If BESTA contributes to your research, please cite the project and acknowledge the software in your publication. + +## License + +BSD 3-Clause. + +## Contact + +For questions, please send an email to p.corcho.caballero@rug.nl diff --git a/docs/requirements.txt b/docs/requirements.txt index 5e7477f8..1cc263d7 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -4,7 +4,7 @@ sphinx>=5.0.0 nbsphinx>=0.9.3 cosmosis astropy>=5.0.0 -numpy<2.0 +numpy matplotlib>=3.7.0 extinction>=0.4 scipy>=1.10 diff --git a/docs/source/configuration.rst b/docs/source/configuration.rst index bcabc127..186be9cd 100644 --- a/docs/source/configuration.rst +++ b/docs/source/configuration.rst @@ -269,6 +269,9 @@ Photometry-only fit [runtime] sampler = maxlike + [maxlike] + method = Powell + [output] filename = ./photometry_fit format = text @@ -297,6 +300,13 @@ Single-spectrum fit [runtime] sampler = maxlike emcee + [maxlike] + method = Powell + + [emcee] + nwalkers = 16 + nsteps = 500 + [output] filename = ./spectral_fit format = text @@ -334,6 +344,11 @@ spectra of the same galaxy simultaneously with shared parameters: from besta.pipeline_modules.full_spectral_fit import FullSpectralFitModule cfg = { + + "runtime": {"sampler": "maxlike emcee"}, + "maxlike": {"method": "Nelder-Mead", "tolerance": 1e-3, "maxiter": 3000}, + "emcee": {"walkers": 32, "samples": 100, "nsteps": 100}, + "output": {"filename": "./fit_all", "format": "text"}, "pipeline": { "modules": "FullSpectralFit_blue FullSpectralFit_red", diff --git a/docs/source/grid.rst b/docs/source/grid.rst index cc5a720c..2f1439b4 100644 --- a/docs/source/grid.rst +++ b/docs/source/grid.rst @@ -4,40 +4,52 @@ Grid-Based Inference ==================== BESTA includes a direct **model-grid inference** workflow that does not require -CosmoSIS samplers. This mode is useful when you already have a finite model -library and want fast posterior evaluation over that library. - -Why use the grid module instead of CosmoSIS sampling? ------------------------------------------------------ - -Use the grid module when: - -- Your model space is already discretized (for example: precomputed SED/SSP libraries). -- You want robust, repeatable inference with no MCMC tuning. -- You need high throughput over many objects using candidate selection and parallel workers. -- You want direct control over priors/likelihoods in Python. - -Prefer the CosmoSIS pipeline (see :ref:`pipeline_manager`) when: - -- Your parameter space is continuous and not naturally represented by a fixed grid. -- You need sampler diagnostics and chain-level convergence checks. -- You rely on existing module wiring in the CosmoSIS runtime. - +CosmoSIS samplers. This mode is particularly useful when you already have a finite model +library and want fast posterior evaluation over that library or the volume of data very large. Core Concepts ------------- +The model-grid workflow is built around several key modules/classes: + - :class:`besta.grid.grid.ModelGrid`: container for model observables/targets and metadata. - :class:`besta.grid.grid.GridFitter`: computes posterior weights over grid models. - :mod:`besta.grid.prob`: prior and likelihood building blocks. - :mod:`besta.grid.binning`: candidate selectors to avoid evaluating the full grid for every object. +Working with ModelGrids +----------------------- + +A :class:`~besta.grid.grid.ModelGrid` is a container of a grid of models split into ``observables`` and ``targets`` (see example below). They support multiple I/O formats: + +- FITS tables: :meth:`~besta.grid.grid.ModelGrid.from_fits_table`, + ``ModelGrid.to_fits_table(...)`` +- HDF5: :meth:`~besta.grid.grid.ModelGrid.from_hdf5`, + :meth:`~besta.grid.grid.ModelGrid.to_hdf5` +- Pickle: :meth:`~besta.grid.grid.ModelGrid.from_pickle`, + :meth:`~besta.grid.grid.ModelGrid.to_pickle` +- automatic loader: :meth:`~besta.grid.grid.ModelGrid.load_auto` + +Priors and Likelihoods +---------------------- + +The grid module is fully Bayesian; you choose the ingredients from +:mod:`besta.grid.prob`: + +- priors: :class:`~besta.grid.prob.FlatPrior`, :class:`~besta.grid.prob.CompositePrior`, + :class:`~besta.grid.prob.ObservableDependentPrior`, and others, +- likelihoods: :class:`~besta.grid.prob.GaussianProductLikelihood`, + :class:`~besta.grid.prob.CompositeLikelihood`, etc. + +This gives a similar statistical structure to sampling methods, but evaluated +directly on a finite model set instead of drawing chains. + Minimal Workflow ---------------- 1. Build or load a :class:`~besta.grid.grid.ModelGrid`. -2. Create a :class:`~besta.grid.grid.GridFitter` with a likelihood and prior. +2. Create a :class:`~besta.grid.grid.GridFitter` with a likelihood and a set of priors. 3. Evaluate posterior summaries for one object, or run :meth:`~besta.grid.grid.GridFitter.fit_batch` for many. Example (single-object posterior on one target): @@ -48,12 +60,12 @@ Example (single-object posterior on one target): from besta.grid import ModelGrid, GridFitter from besta.grid.prob import GaussianProductLikelihood, FlatPrior - # N models, P observables, Q targets + # N models, 3 observables, 4 targets grid = ModelGrid( - observables=obs_models, # shape (N, P) - targets=target_models, # shape (N, Q) + observables=obs_models, # shape (N, 3) + targets=target_models, # shape (N, 3) observable_names=["mag_g", "mag_r", "mag_i"], - target_names=["logM", "age", "Z", "z"], + target_names=["logM", "age", "Z", "z"], # (stellar mass, mean age, metals, redshift) ) fitter = GridFitter( @@ -63,23 +75,24 @@ Example (single-object posterior on one target): use_standardised=True, ) - x = np.array([22.1, 21.5, 21.2]) # observed data (P,) - sx = np.array([0.03, 0.03, 0.04]) # observational errors (P,) + x = np.array([22.1, 21.5, 21.2]) # observed data (3,) + sigma_x = np.array([0.03, 0.03, 0.04]) # observational errors (3,) bins = np.linspace(7.0, 12.0, 101) # bins for logM + # Estimate the marginal stellar mass posterior PDF post_logM, centers = fitter.posterior_over_target( x_native=x, - sigma_native=sx, - target_col="logM", + sigma_native=sigma_x, + target_col="logM", bins=bins, ) -Batch Inference (many objects) ------------------------------- +Batch Inference +---------------- -For catalogs, use :meth:`~besta.grid.grid.GridFitter.fit_batch`. -It supports: +For large samples, use :meth:`~besta.grid.grid.GridFitter.fit_batch`. +This method supports: - optional candidate selection (`binner=`), - thread/process parallelism (`n_jobs`, `backend`), @@ -87,11 +100,12 @@ It supports: - optional per-target summary statistics (`stats_for`, `stats_bins`), - optional HDF5 output (`output_hdf5_path`). + .. code-block:: python from besta.grid.binning import KDTreeBinner - binner = KDTreeBinner(dims=[0, 1, 2]).fit(grid) # select candidates in observable space + binner = KDTreeBinner(dims=[0, 1, 2]).fit(grid) # Use the first three observable columns results = fitter.fit_batch( X_native=X_catalog, # shape (M, P) @@ -99,7 +113,7 @@ It supports: binner=binner, n_jobs=8, backend="thread", - stats_for=["logM", "age"], + stats_for=["logM", "age"], # the posterior statistics for these two quantities stats_bins=[np.linspace(7, 12, 120), np.linspace(0, 14, 120)], output_hdf5_path="grid_fit_results.h5", output_hdf5_group="/run1", @@ -111,39 +125,15 @@ It supports: pass -Priors and Likelihoods ----------------------- +Candidate selection +^^^^^^^^^^^^^^^^^^^ -The grid module is fully Bayesian; you choose the ingredients from -:mod:`besta.grid.prob`: +This is an essential feature when it comes to using large high-dimensional model grids and large datasets. The binners act as model cadidate selectors, rather than using the entire grid on each evaluation, based on the observables of the input sources. This reduces significantly the number of posterior evaluations (sometimes by orders of magnitude). -- priors: :class:`~besta.grid.prob.FlatPrior`, :class:`~besta.grid.prob.CompositePrior`, - :class:`~besta.grid.prob.ObservableDependentPrior`, and others, -- likelihoods: :class:`~besta.grid.prob.GaussianProductLikelihood`, - :class:`~besta.grid.prob.CompositeLikelihood`, etc. +Posterior analysis +^^^^^^^^^^^^^^^^^^ -This gives a similar statistical structure to sampling methods, but evaluated -directly on a finite model set instead of drawing chains. - - -Input/Output and Reproducibility --------------------------------- - -:class:`~besta.grid.grid.ModelGrid` supports multiple formats: - -- FITS tables: :meth:`~besta.grid.grid.ModelGrid.from_fits_table`, - ``ModelGrid.to_fits_table(...)`` -- HDF5: :meth:`~besta.grid.grid.ModelGrid.from_hdf5`, - :meth:`~besta.grid.grid.ModelGrid.to_hdf5` -- Pickle: :meth:`~besta.grid.grid.ModelGrid.from_pickle`, - :meth:`~besta.grid.grid.ModelGrid.to_pickle` -- automatic loader: :meth:`~besta.grid.grid.ModelGrid.load_auto` +BESTA includes some built-in tools for post-processing the posterior PDF and estimate several key quantities such as percentiles, MAP, or covariance matrix, per source. -Practical Tips --------------- -- Start with `FlatPrior + GaussianProductLikelihood` as a baseline. -- If the grid is large, use a binner to cut candidate counts before posterior evaluation. -- Use `return_mode="iter"` for low-memory streaming over large catalogs. -- Write batch outputs to HDF5 for reproducible downstream post-processing. diff --git a/docs/source/manager.rst b/docs/source/manager.rst index c879fdee..cf8b7270 100644 --- a/docs/source/manager.rst +++ b/docs/source/manager.rst @@ -70,6 +70,69 @@ Plotting results If ``plot_result=True``, the best-fit spectra/photometry are plotted for each module using :meth:`besta.pipeline_modules.base_module.BaseModule.plot_solution` and saved alongside the output text files. Each module is re-instantiated from the ``.ini`` file to rebuild the model before plotting. +Running Independent Pipelines In Parallel +----------------------------------------- + +For independent runs (for example, fitting many galaxies with the same workflow or using IFU observations), use :class:`besta.pipeline.BatchPipeline`. + +Each element of ``pipeline_configuration_list`` is one full ``MainPipeline`` input (i.e., a list of sub-pipeline configuration dictionaries): + +.. code-block:: python + + from besta.pipeline import BatchPipeline + + # Two independent jobs, each one contains a single-stage MainPipeline. + job_a = [config_a] + job_b = [config_b] + + batch = BatchPipeline( + pipeline_configuration_list=[job_a, job_b], + n_jobs_parallel=2, + ) + + # Returns one status code per independent job (0 means success). + results = batch.run_all_pipelines(plot_result=False) + +You can also instantiate the internal ``MainPipeline`` objects without running: + +.. code-block:: python + + pipelines = batch.build_pipelines() + +Parameter sweeps with ``from_running_parameters`` +------------------------------------------------- + +A very common use case is when all runs share a base configuration. It is possible to create a batch from parameter updates: + +.. code-block:: python + + from besta.pipeline import BatchPipeline + + base_config = { + # full BESTA configuration dict + } + + # Parameters that differ between runs. + running_parameters = [ + {"FullSpectralFit": {"redshift": 0.10}}, + {"FullSpectralFit": {"redshift": 0.12}}, + {"FullSpectralFit": {"redshift": 0.14}}, + ] + + batch = BatchPipeline.from_running_parameters( + pipeline_configuration=base_config, + running_parameters=running_parameters, + n_jobs_parallel=3, + ) + + results = batch.run_all_pipelines() + +Notes: + +- ``running_parameters`` must be a list of dictionaries. +- Each parameter set is deep-copied from the base configuration, so updates from one run do not leak into the others. +- ``BatchPipeline`` returns statuses in the same order as the input jobs. + Short checklist --------------- diff --git a/requirements.txt b/requirements.txt index 86b5974b..d72048ff 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,8 @@ cosmosis>=3.12 population-synthesis-toolkit -numpy<2.0 +numpy matplotlib>=3.9 scipy>=1.13 astropy psutil -numba \ No newline at end of file +numba diff --git a/src/besta/grid/binning.py b/src/besta/grid/binning.py index b2c1d46b..b31d6864 100644 --- a/src/besta/grid/binning.py +++ b/src/besta/grid/binning.py @@ -6,6 +6,7 @@ from itertools import product import json import os +import re import numpy as np import matplotlib.pyplot as plt @@ -124,6 +125,32 @@ def _chunk_ranges(n: int, batch_size: Optional[int]): s = e +def _tuple_key_to_str(key: Sequence[Any]) -> str: + """Serialize tuple-like keys using plain Python ints for stable JSON keys.""" + values = [int(v) for v in key] + if not values: + return "()" + if len(values) == 1: + return f"({values[0]},)" + return "(" + ", ".join(str(v) for v in values) + ")" + + +def _parse_tuple_key_str(key_str: str) -> Tuple[int, ...]: + """Parse tuple-like keys from legacy and NumPy scalar string representations.""" + k_clean = key_str.strip().strip("()") + if k_clean == "": + return tuple() + parts = [p.strip() for p in k_clean.split(",") if p.strip() != ""] + out: List[int] = [] + for part in parts: + nums = re.findall(r"[-+]?\d+", part) + if not nums: + raise ValueError(f"Could not parse tuple key component: {part!r}") + # Use the last integer token so tokens like 'np.int64(2)' map to 2. + out.append(int(nums[-1])) + return tuple(out) + + def _sigma_to_space(sig_native: np.ndarray, T: dict) -> np.ndarray: """Map 1-sigma vector from native into the transform space described by T.""" mode = (T.get("mode") or "none").lower() @@ -692,7 +719,7 @@ def save(self, path: str) -> None: }, "_edges": [[e.tolist() for e in lev] for lev in self._edges], "_layers": { - str(li): {str(k): v.tolist() for k, v in layer.items()} + str(li): {_tuple_key_to_str(k): v.tolist() for k, v in layer.items()} for li, layer in enumerate(self._layers) }, "_level_bin_size": self._level_bin_size.tolist(), @@ -725,10 +752,7 @@ def load(cls, path: str) -> "RectBinner": layer_d = d["_layers"][str(li)] layer: Dict[Tuple[int, ...], np.ndarray] = {} for k_str, v in layer_d.items(): - # Robust tuple parsing from the str(key) form - k_clean = k_str.strip().strip("()") - parts = [p.strip() for p in k_clean.split(",") if p.strip() != ""] - key = tuple(int(p) for p in parts) + key = _parse_tuple_key_str(k_str) layer[key] = np.asarray(v, dtype=np.int64) layers.append(layer) obj._layers = layers @@ -1032,7 +1056,7 @@ def save(self, path: str) -> None: "_cell_width": self._cell_width.tolist(), "_coord_min": self._coord_min.tolist(), "_coord_max": self._coord_max.tolist(), - "_cells": {str(k): v.tolist() for k, v in self._cells.items()}, + "_cells": {_tuple_key_to_str(k): v.tolist() for k, v in self._cells.items()}, } with open(path, "w") as f: json.dump(blob, f) @@ -1066,9 +1090,7 @@ def load(cls, path: str) -> "HashedGridBinner": obj._coord_max = np.asarray(d["_coord_max"], dtype=np.int64) cells = {} for k_str, v in d["_cells"].items(): - k_clean = k_str.strip().strip("()") - parts = [p.strip() for p in k_clean.split(",") if p.strip() != ""] - key = tuple(int(p) for p in parts) + key = _parse_tuple_key_str(k_str) cells[key] = np.asarray(v, dtype=np.int64) obj._cells = cells return obj @@ -1373,11 +1395,7 @@ def dims(self) -> List[int]: @staticmethod def _parse_tuple_key(key_str: str) -> Tuple[int, ...]: - k_clean = key_str.strip().strip("()") - if k_clean == "": - return tuple() - parts = [p.strip() for p in k_clean.split(",") if p.strip() != ""] - return tuple(int(p) for p in parts) + return _parse_tuple_key_str(key_str) @staticmethod def _coords_for_edges(X: np.ndarray, edges_per_dim: List[np.ndarray]) -> np.ndarray: @@ -1481,7 +1499,7 @@ def _pack_binner(b: BaseBinner) -> Dict[str, Any]: "_cell_width": None if b._cell_width is None else b._cell_width.tolist(), "_coord_min": None if b._coord_min is None else b._coord_min.tolist(), "_coord_max": None if b._coord_max is None else b._coord_max.tolist(), - "_cells": {str(k): v.tolist() for k, v in b._cells.items()}, + "_cells": {_tuple_key_to_str(k): v.tolist() for k, v in b._cells.items()}, } if isinstance(b, RectBinner): @@ -1505,7 +1523,7 @@ def _pack_binner(b: BaseBinner) -> Dict[str, Any]: }, "_edges": [[e.tolist() for e in lev] for lev in b._edges], "_layers": { - str(li): {str(k): v.tolist() for k, v in layer.items()} + str(li): {_tuple_key_to_str(k): v.tolist() for k, v in layer.items()} for li, layer in enumerate(b._layers) }, "_level_bin_size": b._level_bin_size.tolist(), @@ -1785,12 +1803,16 @@ def save(self, path: str) -> None: "final_target_k": self.final_target_k, "use_global_fallback": self.use_global_fallback, "_primary_edges": [ed.tolist() for ed in self._primary_edges], - "_primary_cells": {str(k): v.tolist() for k, v in self._primary_cells.items()}, + "_primary_cells": { + _tuple_key_to_str(k): v.tolist() for k, v in self._primary_cells.items() + }, "_secondary_models": { - str(k): self._pack_binner(v) for k, v in self._secondary_models.items() + _tuple_key_to_str(k): self._pack_binner(v) + for k, v in self._secondary_models.items() }, "_secondary_global_index": { - str(k): v.tolist() for k, v in self._secondary_global_index.items() + _tuple_key_to_str(k): v.tolist() + for k, v in self._secondary_global_index.items() }, "_fallback_secondary": None if self._fallback_secondary is None diff --git a/src/besta/grid/grid.py b/src/besta/grid/grid.py index 6fd16dbe..766e154c 100644 --- a/src/besta/grid/grid.py +++ b/src/besta/grid/grid.py @@ -1,8 +1,5 @@ """ Model grid container and fitting machinery. - -This module defines the ModelGrid class, which stores a grid of models with -their parameters and provides methods for fitting and evaluating these models. """ from __future__ import annotations @@ -26,7 +23,7 @@ from scipy.spatial import cKDTree from scipy.stats import gaussian_kde -from besta.grid.prob import ( +from .prob import ( Prior, FlatPrior, ObservableDependentPrior, @@ -34,17 +31,14 @@ GaussianProductLikelihood, posterior_over_models as posterior_over_models_fn, ) +from .transforms import LinearStandardiser + from besta.postprocess import ( enclosed_fraction_map, - pit_from_discrete_posterior, pdf_stats, - photoz_metrics, weighted_quantiles, ) -from .transforms import LinearStandardiser - -from besta.utils import available_memory_bytes from besta.logging import get_logger os.environ.setdefault("OMP_NUM_THREADS", "1") @@ -53,8 +47,10 @@ logger = get_logger(__name__) +# ------------- Helper methods --------------- def _guess_slices(n_objects, n_observables, n_jobs, tasks_per_worker=6): + """Helper function for guessing the amount of parallel fit slices.""" # fewer, larger slices when P is large base_tasks = n_jobs * tasks_per_worker scale = max(1, n_observables // 8) @@ -322,6 +318,7 @@ class ModelGrid: default=None, init=False, repr=False ) + # Model interpolation _kdtree: Optional["cKDTree"] = field(default=None, init=False, repr=False) _kdtree_standardized: Optional[bool] = field(default=None, init=False, repr=False) @@ -611,7 +608,7 @@ def get_kdtree(self, *, standardize: bool = True): self._kdtree_standardized = standardize return self._kdtree - def in_boundaries(self, targets_query: np.ndarray) -> np.ndarray: + def _in_boundaries(self, targets_query: np.ndarray) -> np.ndarray: """ Check which target queries are within the model grid boundaries. @@ -650,11 +647,17 @@ def interpolate_observables( ridge: float = 1e-8, # Tikhonov regularization for stability ) -> np.ndarray: """ - Interpolate observables for arbitrary target values using cached KDTree. + Interpolate observables for arbitrary target values using a KDTree. - mode="nearest" : nearest-neighbour lookup - mode="idw" : inverse-distance weighted KNN - mode="local_linear": weighted local affine fit (exact for linear functions) + Parameters + ---------- + targets_query: np.nddarray + Target values to interpolate. + method: str + Optional interpolation method. Currently the methods supported are: + - "nearest" (select nearest-neighbour candidate) + - "idw" (inverse-distance weighted KNN) + - "local_linear" (weighted local affine fit, only exact for linear functions) """ tq = np.atleast_2d(np.asarray(targets_query, dtype=float)) if tq.shape[1] != self.n_targets: @@ -670,7 +673,7 @@ def interpolate_observables( f"Unknown mode={mode!r} (use 'nearest', 'idw' or 'local_linear')." ) - bad_q = ~np.isfinite(tq).all(axis=1) | ~self.in_boundaries(tq) + bad_q = ~np.isfinite(tq).all(axis=1) | ~self._in_boundaries(tq) out = np.full((tq.shape[0], self.n_observables), fill_value, dtype=float) if bad_q.all(): return out diff --git a/src/besta/grid/prob.py b/src/besta/grid/prob.py index 968addb8..20306727 100644 --- a/src/besta/grid/prob.py +++ b/src/besta/grid/prob.py @@ -5,20 +5,13 @@ from dataclasses import dataclass from typing import Optional, Sequence, Tuple, List import warnings +from numba import njit, prange import numpy as np from besta.logging import get_logger logger = get_logger(__name__) -try: - from numba import njit, prange - - NUMBA_OK = True -except Exception: - NUMBA_OK = False - logger.warning("numba could not be imported") - # ------------------------------- utilities ------------------------------- @@ -90,7 +83,7 @@ def _std_norm_cdf(x: np.ndarray) -> np.ndarray: class Prior(ABC): """ - Abstract prior interface over model targets. + Prior base class. A prior returns log p(theta) for each model row. It may depend on specific target columns (e.g., redshift) and optionally on other @@ -308,7 +301,7 @@ class EmpiricalHistogramPrior1D(Prior): Parameters ---------- target_col : int - Index of the target column to build the prior on (e.g., redshift). + Index of the target column to build the prior. edges : ndarray, shape (K+1,) Histogram bin edges. Must cover the support of the target. density_floor : float, optional @@ -334,14 +327,19 @@ def fit_from_targets( ------- self : EmpiricalHistogramPrior1D """ + + logger.debug("Fitting EmpiricalHistogramPrior 1D") + # Select the target column t = targets[:, self.target_col] + # Compute the histogram hist, _ = np.histogram(t, bins=self.edges, weights=weights, density=False) - mass = hist.astype(float) - mass = ( - mass / np.sum(mass) - if np.sum(mass) > 0 - else np.full_like(mass, 1.0 / mass.size) - ) + mass = hist * np.diff(self.edges) + # Normalise histogram + norm = np.sum(mass) + if norm > 0: + mass /= norm + else: + mass = np.full_like(mass, 1.0 / mass.size) mass = np.clip(mass, self.density_floor, None) self._logp_per_bin = np.log(mass) return self @@ -350,6 +348,7 @@ def log_prob_for_models(self, targets: np.ndarray) -> np.ndarray: if not hasattr(self, "_logp_per_bin"): raise RuntimeError("Prior not fitted. Call fit_from_targets first.") t = targets[:, self.target_col] + # Bin targets using the pre-defined bins j = np.digitize(t, self.edges) - 1 j = np.clip(j, 0, self._logp_per_bin.size - 1) return self._logp_per_bin[j] @@ -358,17 +357,17 @@ def log_prob_for_models(self, targets: np.ndarray) -> np.ndarray: @dataclass class EmpiricalFlatteningPriorND(Prior): """ - Empirical flattening prior over several target columns. + Empirical flattening prior over arbitrary target columns. This prior uses the model grid itself to estimate the (possibly - non-flat) distribution of a set of parameters and builds a prior + non-uniform) distribution of a set of parameters and builds a prior that counteracts those inhomogeneities. Two modes are provided: - 'factorised': build 1-D histograms for each column separately and form a product prior over dimensions. This approximately flattens - the *marginal* distributions of those parameters. + the marginal distributions of those parameters. - 'joint': build a joint N-D histogram over all selected columns and assign prior mass proportional to 1 / N_k for each occupied @@ -433,41 +432,41 @@ def fit_from_targets( ------- self : EmpiricalFlatteningPriorND """ + logger.debug("Fitting EmpiricalHistogramPriorND") t = targets[:, self.target_cols] # (N, D) + # Grid dimensions D = t.shape[1] if weights is not None and weights.shape[0] != t.shape[0]: raise ValueError("weights must have shape (N,) if provided.") if self.mode == "factorised": + logger.debug("Using 'factorised' mode (per-dim prior)") # One histogram per dimension, store log inverse-mass per bin log_inv_mass_list = [] for d in range(D): edges = self.edges_list[d] + # get all parameter values td = t[:, d] - + # compute histogram counts, _ = np.histogram(td, bins=edges, weights=weights, density=False) - counts = counts.astype(float) - total = np.sum(counts) if total <= 0: raise RuntimeError( - f"No models in any bin for dimension {d}; cannot fit prior." + f"No models in any user-provided bin for dimension {d}; cannot fit prior." ) - + # Clip prior to prevent zero division counts = np.clip(counts, self.count_floor, None) - # Define per-bin mass proportional to 1 / counts + # Define per-bin prior mass proportional to 1 / counts inv_counts = 1.0 / counts inv_counts /= np.sum(inv_counts) - - # Store log(mass_d per bin) or directly log(1/count_d) up to a constant - # For our purpose, log prior for a model in bin j_d is sum_d log(inv_counts_d[j_d]) log_inv_mass_list.append(np.log(inv_counts)) self._log_inv_mass_list = log_inv_mass_list else: # mode == 'joint' + logger.debug("Using 'joint' mode (multi-dim prior)") # Build joint N-D histogram bin_indices = [] bin_sizes = [] @@ -480,12 +479,12 @@ def fit_from_targets( bin_indices.append(j) bin_sizes.append(edges.size - 1) - bin_indices = np.stack(bin_indices, axis=0) # (D, N) + bin_indices = np.stack(bin_indices, axis=0) # Flatten to 1-D indices for bincount linear_indices = np.ravel_multi_index( bin_indices, dims=tuple(bin_sizes) - ) # (N,) + ) counts_flat = np.bincount( linear_indices, @@ -561,7 +560,7 @@ def log_prob_for_models(self, targets: np.ndarray) -> np.ndarray: class ObservableDependentPrior(Prior): - """TODO""" + """Base class for priors that depend on observables.""" def fit_from_grid(self): raise NotImplementedError() @@ -570,14 +569,16 @@ def fit_from_grid(self): @dataclass class MagDependentRedshiftPrior(ObservableDependentPrior): """ - Magnitude-dependent redshift prior p(z | m) from a 2-D histogram. + Magnitude-dependent redshift prior, i.e. likelihood + of an object with a magnitude ``m`` being detected at + redshift ``z``. Parameters ---------- z_col : int Index of redshift in targets. mag_observable_index : int - Index of magnitude in observables (e.g., VIS magnitude column). + Index of magnitude in observables. z_edges : ndarray Bin edges in redshift. m_edges : ndarray @@ -598,6 +599,7 @@ class MagDependentRedshiftPrior(ObservableDependentPrior): z_edges: np.ndarray m_edges: np.ndarray density_floor: float = 1e-12 + # TODO: allow for optional user-provided prior def fit_from_grid( self, @@ -606,7 +608,7 @@ def fit_from_grid( weights: Optional[np.ndarray] = None, ) -> "MagDependentRedshiftPrior": """ - Fit conditional histogram from the model grid. + Fit conditional histogram from input dataset. Parameters ---------- @@ -623,19 +625,19 @@ def fit_from_grid( H, z_edges, m_edges = np.histogram2d( z, m, bins=[self.z_edges, self.m_edges], weights=weights ) - # normalise each magnitude column to sum 1 over z + # compute the conditional distribution colsum = H.sum(axis=0, keepdims=True) colsum[colsum == 0] = 1.0 - P = H / colsum - P = np.clip(P, self.density_floor, None) - self._logP_z_given_m = np.log(P) # shape (Kz, Km) + p_z_given_m = H / colsum + p_z_given_m = np.clip(p_z_given_m, self.density_floor, None) + self._logP_z_given_m = np.log(p_z_given_m) return self def log_prob_for_models( self, targets: np.ndarray, observables: Optional[np.ndarray] = None ) -> np.ndarray: """ - Evaluate log p(z | m) per model row. + Evaluate :math:`log p(z | m)` per model row. Parameters ---------- @@ -653,6 +655,7 @@ def log_prob_for_models( raise ValueError("observables must be provided to evaluate p(z|m)") z = targets[:, self.z_col] m = observables[:, self.mag_observable_index] + # Interpolate input magnitudes iz = np.clip( np.digitize(z, self.z_edges) - 1, 0, self._logP_z_given_m.shape[0] - 1 ) @@ -662,29 +665,6 @@ def log_prob_for_models( return self._logP_z_given_m[iz, im] -class HierarchicalPrior(Prior): - """Abstract base class for priors controlled by learnable hyperparameters.""" - - def __init__(self, hyperparams: dict): - self.hyperparams = hyperparams - - def update_hyperparams(self, new_values: dict) -> None: - self.hyperparams.update(new_values) - - @abstractmethod - def log_prob_for_models(self, targets: np.ndarray, **kwargs) -> np.ndarray: - pass - - @abstractmethod - def fit_from_data( - self, - targets: np.ndarray, - observables: np.ndarray, - weights: np.ndarray | None = None, - ) -> None: - pass - - @dataclass class CompositePrior(Prior): """ @@ -693,9 +673,9 @@ class CompositePrior(Prior): The total log prior is defined as a weighted sum of component log priors: - log p_total(model) = sum_i w_i * log p_i(model) + :math:`\log p_{total}(model) = sum_i w_i * \log p_i(model)` - where each p_i is a Prior that does *not* depend on observables. + where each p_i is a Prior that does not depend on observables. Parameters ---------- @@ -715,6 +695,7 @@ class CompositePrior(Prior): def __post_init__(self): self.priors = list(self.priors) + logger.debug(f"Setting up CompositePrior with {len(self.priors)} priors") if not self.priors: raise ValueError("CompositePrior requires at least one component prior.") @@ -726,6 +707,7 @@ def __post_init__(self): ) if self.weights is not None: + logger.debug("Using user-provided relative prior weights") if len(self.weights) != len(self.priors): raise ValueError( "weights must have the same length as priors " @@ -771,20 +753,22 @@ def log_prob_for_models(self, targets: np.ndarray) -> np.ndarray: warnings.warn( "All weights were zero in CompositePrior; returning flat prior." ) + logger.warning( + "All weights were zero in CompositePrior; returning flat prior." + ) return np.zeros(N, dtype=float) - return logp_total @dataclass class ObservableCompositePrior(ObservableDependentPrior): """ - Composite prior combining Priors, including observable-dependent ones. + Same as :class:`CompositePrior`, but including :class:`ObservableDependentPrior`. The total log prior is defined as a weighted sum of component log priors: - log p_total(model) = sum_i w_i * log p_i(model) + :math:`\log p_{total}(model) = sum_i w_i * \log p_i(model)` Components can be: - Plain Prior (target-only), evaluated as @@ -819,12 +803,15 @@ class ObservableCompositePrior(ObservableDependentPrior): def __post_init__(self): self.priors = list(self.priors) + logger.debug(f"Setting up CompositePrior with {len(self.priors)} priors") + if not self.priors: raise ValueError( "ObservableCompositePrior requires at least one component prior." ) if self.weights is not None: + logger.debug("Using user-provided relative prior weights") if len(self.weights) != len(self.priors): raise ValueError( "weights must have the same length as priors " @@ -899,7 +886,7 @@ def log_prob_for_models( class Likelihood(ABC): """ - Abstract likelihood interface p(x | model). + Base likelihood class representing :math:`p(x | model)`. Methods ------- @@ -935,10 +922,6 @@ class GaussianProductLikelihood(Likelihood): """ Independent per-dimension Gaussian product likelihood. - The likelihood is proportional to the product over j of - N(x_j | X_ij, h_j^2), where h_j is derived from sigma_native with - an optional floor. - Parameters ---------- bandwidth_floor : float, optional @@ -953,8 +936,8 @@ class GaussianProductLikelihood(Likelihood): def log_likelihood( self, x_native: np.ndarray, sigma_native: np.ndarray, X_models: np.ndarray ) -> np.ndarray: + # truncate prior h = np.maximum(self.scale * sigma_native, self.bandwidth_floor) - # broadcast to (Nc, P) diff = X_models - x_native[None, :] var = h[None, :] ** 2 # sum of 1-D logpdfs @@ -965,66 +948,33 @@ def log_likelihood( @dataclass -class CensoredSizeLikelihood(Likelihood): - """ - Photometry-only Gaussian product with a left-censored size factor. +class SplitGaussianProductLikelihood(Likelihood): + """Independent per-dimension split Gaussian likelihood - This is useful when the apparent size is below a reliability floor - (e.g., PSF or measurement threshold). The photometric part is a - Gaussian product over selected photometry indices. The size part - adds a log CDF factor log Phi((s_min - s_model) / h_s), where s is - log10(Re) and h_s is derived from sigma_native[size_index]. + The likelihood along each dimension is given by - Parameters - ---------- - phot_indices : Sequence[int] - Indices of observable columns to include in the Gaussian product - (typically colours and anchor magnitude). - size_index : int - Index of the size observable column (e.g., log10(Re)). - s_min : float - Left-censoring threshold in the same units as the size observable. - bandwidth_floor : float, optional - Minimum bandwidth per dimension in native units. Default 0.0. - scale : float, optional - Multiplicative scale applied to sigma_native. Default 1.0. - """ + .. math: - phot_indices: Sequence[int] - size_index: int - s_min: float - bandwidth_floor: float = 0.0 - scale: float = 1.0 + \mathcal{L} = \mathcal{N}(x | \mu, \sigma_L),\, if x\leq\mu\\ + \mathcal{L} = \mathcal{N}(x | \mu, \sigma_R),\, if x>\mu + + and the total likelihood is the product along all dimensions. + """ + bandwith_floor: float = 0.0 + scale_left: float = 1.0 + scale_right: float = 1.0 def log_likelihood( self, x_native: np.ndarray, sigma_native: np.ndarray, X_models: np.ndarray ) -> np.ndarray: - # Photometry part - phot_idx = np.asarray(self.phot_indices, dtype=int) - x_ph = x_native[phot_idx] - sig_ph = np.maximum(self.scale * sigma_native[phot_idx], self.bandwidth_floor) - Xm_ph = X_models[:, phot_idx] - diff = Xm_ph - x_ph[None, :] - var = sig_ph[None, :] ** 2 - logL_ph = -0.5 * ( - np.sum(np.log(2.0 * np.pi * var), axis=1) + np.sum(diff**2 / var, axis=1) - ) - - # Censored size factor: log Phi((s_min - s_model)/h_s) - h_s = max(self.scale * sigma_native[self.size_index], self.bandwidth_floor) - s_model = X_models[:, self.size_index] - z = (self.s_min - s_model) / h_s - # avoid log(0) - cdf = np.clip(_std_norm_cdf(z), 1e-300, 1.0) - logL_sz = np.log(cdf) - - return logL_ph + logL_sz + # truncate prior + raise NotImplementedError("Class not implemented") @dataclass class CompositeLikelihood(Likelihood): """ - Sum of multiple likelihood terms (log-likelihoods add). + Sum of multiple likelihood terms. Parameters ---------- @@ -1111,7 +1061,7 @@ def posterior_over_models( return w -# Numba-dedicated likelihood +# Numba-dedicated likelihood for increased performance @njit(parallel=True, fastmath=True, cache=True) @@ -1122,17 +1072,14 @@ def _quadform_diag_parallel(X, x, h): for i in prange(N): s = 0.0 Xi = X[i] - # unrolled-style simple loop lets numba vectorise well for j in range(P): d = (Xi[j] - x[j]) * invh[j] s += d * d out[i] = s - return out # squared Mahalanobis with diagonal covariance - + return out @njit(parallel=True, fastmath=True, cache=True) def _loglike_gaussprod_diag(X, x, h): - # log L_i = -0.5 * sum_j ((X_ij - x_j)/h_j)^2 (constants drop) q = _quadform_diag_parallel(X, x, h) return -0.5 * q @@ -1154,10 +1101,7 @@ def __init__( self.prefer_batch = prefer_batch def log_likelihood(self, x_native, sigma_native, X_models): - if not NUMBA_OK: - return super().log_likelihood(x_native, sigma_native, X_models) - - # Expect C-contiguous float64 for best performance + # C-contiguous float64 for best performance (https://stackoverflow.com/questions/67784563/how-to-make-two-arrays-contiguous-so-that-numba-can-speed-up-np-dot) x = np.ascontiguousarray(x_native, dtype=np.float64) X = np.ascontiguousarray(X_models, dtype=np.float64) sigma = np.ascontiguousarray(sigma_native, dtype=np.float64) diff --git a/src/besta/grid/transforms.py b/src/besta/grid/transforms.py index 961b7710..eda31c07 100644 --- a/src/besta/grid/transforms.py +++ b/src/besta/grid/transforms.py @@ -1,8 +1,6 @@ -"""Lightweight transforms used by model grids and emulators.""" - -# pst/transforms.py -# -*- coding: utf-8 -*- - +"""Transforms used by model grids and emulators.""" +# +# from __future__ import annotations from dataclasses import dataclass @@ -14,24 +12,20 @@ @dataclass class LinearStandardiser: """ - Simple per-dimension standardisation: x_std = (x - mean) / sd. - - Notes - ----- - - Uses ddof=0 by default (population std), matching your ModelGrid. - - If sd == 0, it is set to 1 to avoid division by zero. + Per-dimension linear standardisation: x_std = (x - mean) / sd. """ mean: Optional[np.ndarray] = None sd: Optional[np.ndarray] = None def fit(self, X: np.ndarray, *, ddof: int = 0) -> "LinearStandardiser": + """Fit the standariser parameters.""" X = np.asarray(X) mu = np.nanmean(X, axis=0) - sd = np.nanstd(X, axis=0, ddof=ddof) - sd = np.where(sd == 0.0, 1.0, sd) + std = np.nanstd(X, axis=0, ddof=ddof) + std = np.where(std == 0.0, 1.0, std) self.mean = mu - self.sd = sd + self.sd = std return self @property @@ -64,32 +58,6 @@ def from_dict(cls, d: Dict[str, Any]) -> "LinearStandardiser": sd = None if d.get("sd") is None else np.asarray(d["sd"], dtype=float) return cls(mean=mean, sd=sd) +# TODO: vmin/vmax linear transform -@dataclass -class MagTransform: - """ - Magnitude transform. - - mag = -2.5 log10(flux) + zero_point - flux = 10^((zero_point - mag)/2.5) - """ - zero_point: float = 0.0 - eps: float = 1e-30 - - def flux_to_mag(self, flux: np.ndarray) -> np.ndarray: - f = np.maximum(np.asarray(flux), self.eps) - return (-2.5 * np.log10(f)) + self.zero_point - - def mag_to_flux(self, mag: np.ndarray) -> np.ndarray: - m = np.asarray(mag) - return np.power(10.0, (self.zero_point - m) / 2.5) - - def to_dict(self) -> Dict[str, Any]: - return {"zero_point": float(self.zero_point), "eps": float(self.eps)} - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> "MagTransform": - return cls( - zero_point=float(d.get("zero_point", 0.0)), eps=float(d.get("eps", 1e-30)) - ) diff --git a/src/besta/io.py b/src/besta/io.py index 95362e7a..34a3b905 100644 --- a/src/besta/io.py +++ b/src/besta/io.py @@ -5,6 +5,7 @@ import importlib.util import sys from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple import numpy as np @@ -14,7 +15,6 @@ from cosmosis.datablock import DataBlock, SectionOptions from astropy.table import Table -from besta import pipeline_modules from besta.logging import get_logger, setup_logging from besta.utils import expand_env_vars @@ -148,9 +148,9 @@ def _parse_group(token: str): def _parse_scalar(token: str): low = token.lower() - if low in {"true", "yes", "on"}: + if low in {"t", "true", "yes", "on"}: return True - if low in {"false", "no", "off"}: + if low in {"f", "false", "no", "off"}: return False if low in {"none", "null"}: return "none" @@ -329,7 +329,7 @@ def make_values_file(config, overwrite=True, values_sec="values"): make_ini_file(values_filename, config[values_sec], ignore_sec=None) @expand_env_vars() -def read_results_file(path): +def read_results_file(path, delimiter="\t"): """Read the results produced during a CosmoSIS run. Parameters @@ -343,13 +343,21 @@ def read_results_file(path): Table containing the results. """ with open(path, "r", encoding="utf-8") as f: - header = f.readline().strip("#") - columns = header.replace("\n", "").split("\t") + header = f.readline() + if not header.startswith("#"): + raise ValueError("Expected first line header starting with '#'.") + + columns = [col.strip().lower() for col in header.strip("# \n").split(delimiter)] matrix = np.atleast_2d(np.loadtxt(path)) table = Table() - if matrix.size > 1: - for ith, c in enumerate(columns): - table.add_column(matrix.T[ith], name=c.lower()) + if matrix.size <= 1: + return table + if matrix.shape[1] != len(columns): + raise ValueError( + f"Data has {matrix.shape[1]} columns but header lists {len(columns)}." + ) + for ith, name in enumerate(columns): + table[name] = matrix[:, ith] return table def load_class_from_path(file_path, class_name): @@ -369,6 +377,78 @@ def load_class_from_path(file_path, class_name): return getattr(module, class_name) +# Table operations + +def parse_table_format(path): + """Parse the format of a table file based on its extension. + + Parameters + ---------- + path : str + Path to the table file. + + Returns + ------- + format : str + Format of the table file (e.g., "csv", "fits", "ascii"). + """ + extension = os.path.splitext(path)[1].lower() + if extension in [".csv"]: + format = "csv" + elif extension in [".fits"]: + format = "fits" + elif extension in [".txt", ".dat"]: + format = "ascii" + else: + logger.warning(f"Could not guess file format from extension '{extension}'. Defaulting to ASCII.") + format = "ascii" + + return format + +def burn_table(table, nwalkers: int, burn_in: int) -> Table: + """Discard the first `burn_in` samples per walker from the table.""" + nrows = len(table) + expected = nwalkers * burn_in + if nrows < expected: + raise ValueError(f"Not enough rows in table ({nrows}) for burn_in={burn_in} and nwalkers={nwalkers} (expected at least {expected}).") + # Keep rows after burn-in for each walker + mask = np.ones(nrows, dtype=bool) + for w in range(nwalkers): + start = w * burn_in + end = (w + 1) * burn_in + mask[start:end] = False + return table[mask] + +def _select_parameter_keys( + table: Table, + *, + parameter_prefix: str = "--", + parameter_keys: Optional[Sequence[str]] = None, +) -> List[str]: + if parameter_keys is not None: + keys = list(parameter_keys) + else: + keys = [k for k in table.colnames if parameter_prefix in k] + if len(keys) == 0: + raise ValueError("No parameter keys found/selected.") + return keys + +def _split_param_key(key: str, prefix: str = "--") -> Tuple[str, str]: + """ + Split a parameter key into (section, name) using the delimiter/prefix. + + If the key cannot be split, returns ("", key). + """ + if prefix in key: + sect, name = key.split(prefix, 1) + return sect, name + return "", key + + +################### +# Main reader class +################### + class Reader(object): r"""CosmoSIS run results reader. @@ -459,7 +539,7 @@ def ini_values(self, value): @property def values_file(self) -> str: """Path to the CosmoSIS (prior) values configuration file.""" - return getattr(self, "_ini_file", None) + return getattr(self, "_values_file", None) @values_file.setter def values_file(self, value): @@ -497,10 +577,6 @@ def get_module(self, module_name): module = load_class_from_path(self.ini[module_name]["file"], "module") logger.debug(f"Loaded module {module_name} from {self.ini[module_name]['file']}") return module(self.ini, alias=module_name) - # if not hasattr(pipeline_modules, module_class): - # raise ValueError( - # f"Module class {module_class} not found in besta.pipeline_modules.") - # return getattr(pipeline_modules, module_class)(options) #TODO: deprecate @property @@ -574,12 +650,24 @@ def __init__(self, ini_file=None, results_file=None): overwrite=logging_overwrite, console=logging_console) - def load_results(self): - """Load the cosmosis run results associated to the ``ini`` file.""" + def load_results(self, burn_in=0, nwalkers=1): + """Load the cosmosis run results associated to the ``ini`` file. + + Parameters + ---------- + burn_in : int, optional + Number of initial samples to discard per walker. Default is 0 (no burn-in). + nwalkers : int, optional + Number of walkers used during the sampling. Required if burn_in > 0. Default is 1. + """ path = self.ini["output"]["filename"] if ".txt" not in path: path += ".txt" self.results_table = read_results_file(path) + if burn_in > 0: + self.results_table = burn_table( + self.results_table, + nwalkers=self.ini["pipeline"].get("nwalkers", nwalkers), burn_in=burn_in) def get_maxlike_solution(self, log_prob="post", as_datablock=False, **kwargs): @@ -645,7 +733,8 @@ def get_top_frac_solutions(self, frac=1, log_prob="post", as_datablock=False, -------- :func:`solution_to_datablock` """ - assert frac > 0 and frac <= 100, "Fraction must be in (0, 100]" + if not (0 < frac <= 100): + raise ValueError("Fraction must be in (0, 100].") good_sample = np.isfinite(self.results_table[log_prob]) tab = self.results_table[good_sample] post_sort = np.argsort(tab[log_prob]) diff --git a/src/besta/kinematics.py b/src/besta/kinematics.py index e6099ea8..de927388 100644 --- a/src/besta/kinematics.py +++ b/src/besta/kinematics.py @@ -1,358 +1,431 @@ -"""This module contains the tools for modelling kinematic effects on spectra.""" +"""Kinematic convolution utilities for spectral modeling. + +This module provides pixel-space LOSVD kernel classes and convolution helpers +used by BESTA spectral fitting modules. Kernels are designed to be reusable +across likelihood calls via lightweight in-memory caching. + +Conventions +----------- +- Velocities and dispersions are provided in physical units and converted to + pixel units using ``velocity_scale``. +- Convolutions act on the last axis of the input arrays. +- Kernels are normalized to unit sum before convolution. +""" import numpy as np import re from scipy.signal import fftconvolve from scipy.special import erf from scipy import sparse -from astropy.modeling import Fittable1DModel -from astropy.modeling.models import Gaussian1D, Hermite1D -from astropy.convolution.kernels import Model1DKernel - from astropy import units as u -from astropy.convolution import convolve, convolve_fft -from besta import spectrum from besta import config as CONFIG +from besta.logging import get_logger + +logger = get_logger(__name__) -# TODO: implement split Gaussian model -# This model is much more stable and never -# produces negative densities +SQRT2 = np.sqrt(2) +SQRT2PI = np.sqrt(2 * np.pi) +CACHE_PIX_DECIMALS = 3 +CACHE_NMODELS = 256 +DELTA_KERNEL_ATOL = 1e-3 -class GaussHermite(Fittable1DModel): - """Gauss-Hermite model.""" +class LOSVDPixelKernel: + """Line-of-sight velocity distribution kernel. - _param_names = () + Attributes + ---------- + velocity_scale : float + Velocity step represented by one pixel (same units as LOS velocities). + kernel_weight : np.ndarray or None + Normalized kernel weights sampled on the pixel grid. + edge_pixels : int + Number of edge pixels likely affected by convolution artifacts. + skip_convolution : bool + If ``True``, convolution is treated as identity (delta-like kernel). + """ - def __init__(self, order, *args, **kwargs): - self._order = int(order) - if self._order < 3: - self._order = 0 + @property + def kernel_weight(self): + """Kernel weights.""" + return self._kernel_weight + + @kernel_weight.setter + def kernel_weight(self, value): + if value is not None: + value = np.asarray(value, dtype=float) + norm = np.sum(value) + + if norm > 0: + # Check if all the weight is on a single pixel (delta kernel) + if np.isclose(norm, value.max(), atol=DELTA_KERNEL_ATOL): + self.skip_convolution = True + else: + self.skip_convolution = False + + self._kernel_weight = value / norm + + else: + logger.warning("Kernel weights sum to zero; using unnormalized values.") + raise ValueError("Kernel weights sum to zero; cannot normalize.") + # self.skip_convolution = False + # self._kernel_weight = value - self._gaussian = Gaussian1D() - # Hermite series - if self._order: - self._hermite = Hermite1D(self._order) + @property + def size(self): + """Kernel size in pixels.""" + if self.kernel_weight is not None: + return self.kernel_weight.size else: - self._hermite = None + return 0 - self._param_names = self._generate_coeff_names() - super(GaussHermite, self).__init__(*args, **kwargs) + def __init__(self, velocity_scale): + """Initialize a pixel-space LOSVD kernel container. - def _generate_coeff_names(self): - names = list(self._gaussian.param_names) # Gaussian parameters - names += ["h{}".format(i) for i in range(3, self._order + 1)] # Hermite coeffs + Parameters + ---------- + velocity_scale : float + Velocity step represented by one spectral pixel. + """ + logger.debug("Initializing LOSVDPixelKernel with velocity_scale=%s", velocity_scale) + self.velocity_scale = velocity_scale + self._kernel_weight = None + self.skip_convolution = False + self.edge_pixels = 0 + self._cache = {} - def _hi_order(self, name): - # One could store the compiled regex, but it will crash the deepcopy: - # "cannot deepcopy this pattern object" + def _cached_kernel(self, key, build_kernel): + """Populate ``kernel_weight`` from cache or from a builder callback. - match = re.match("h(?P\d+)", name) # h3, h4, etc. - order = int(match.groupdict()["order"]) if match else 0 + Parameters + ---------- + key : hashable + Cache key describing kernel parameters. + build_kernel : Callable[[], np.ndarray] + Callback used to build the kernel when ``key`` is absent. + """ + kernel = self._cache.get(key) + if kernel is None: + kernel = build_kernel() + self._cache[key] = kernel + # Keep cache bounded to avoid unbounded memory growth in long chains. + if len(self._cache) > CACHE_NMODELS: + self._cache.pop(next(iter(self._cache))) + self.kernel_weight = kernel + + def convolve(self, spectra): + """Convolve the input spectra with the LOSVD kernel.""" + if self.kernel_weight is not None: + if self.skip_convolution: + return spectra + if np.ndim(spectra) == 1: + return fftconvolve(spectra, self.kernel_weight, mode="same") + + kernel = self.kernel_weight.reshape((1,) * (np.ndim(spectra) - 1) + (-1,)) + return fftconvolve(spectra, kernel, mode="same", axes=-1) + else: + raise ValueError("Kernel weights are not set.") - return order + def get_percentile_pixel(self, percentile): + """Get a percentile location in pixel units relative to kernel center. - def _generate_coeff_names(self): - names = list(self._gaussian.param_names) # Gaussian parameters - names += ["h{}".format(i) for i in range(3, self._order + 1)] # Hermite coeffs + Parameters + ---------- + percentile : float + Desired percentile (between 0 and 100). - return tuple(names) + Returns + ------- + pixel_offset : float + Pixel offset relative to the central kernel pixel. + """ + if self.kernel_weight is None: + raise ValueError("Kernel weights are not set.") + if not (0.0 <= percentile <= 100.0): + raise ValueError("percentile must be in [0, 100].") + + cumulative = np.cumsum(self.kernel_weight) + # Interpolate on bin edges so symmetric kernels yield zero-centered + # median offsets instead of the half-pixel bias from center-grid CDF. + cdf_edges = np.concatenate(([0.0], cumulative)) + pixel_edges = np.arange(len(self.kernel_weight) + 1, dtype=float) - 0.5 + pixel = np.interp(percentile / 100.0, cdf_edges, pixel_edges) + center = 0.5 * (len(self.kernel_weight) - 1) + return pixel - center + + def get_percentile_velocity(self, percentile): + """Get a percentile location in velocity units relative to kernel center.""" + return self.get_percentile_pixel(percentile) * self.velocity_scale + + def parse_parameters(self, datablock): + """Read kinematic parameters from a DataBlock and set kernel weights. + + Notes + ----- + Subclasses must implement this method according to their parameterization. + """ + raise NotImplementedError("This method should be implemented by subclasses to parse parameters from the kernel model.") - def __getattr__(self, attr): - if attr[0] == "_": - super(GaussHermite, self).__getattr__(attr) - elif attr in self._gaussian.param_names: - return self._gaussian.__getattribute__(attr) - elif self._order and self._hi_order(attr) >= 3: - return self._hermite.__getattribute__(attr.replace("h", "c")) - else: - super(GaussHermite, self).__getattr__(attr) - - def __setattr__(self, attr, value): - if attr[0] == "_": - super(GaussHermite, self).__setattr__(attr, value) - elif attr in self._gaussian.param_names: - self._gaussian.__setattr__(attr, value) - elif self._order and self._hi_order(attr) >= 3: - self._hermite.__setattr__(attr.replace("h", "c"), value) - else: - super(GaussHermite, self).__setattr__(attr, value) - @property - def param_names(self): - """Tuple of Gaussian and Hermite coefficient parameter names.""" - return self._param_names +class GaussianPixelKernel(LOSVDPixelKernel): + """Single-Gaussian LOSVD kernel in pixel space.""" - def evaluate(self, x, *params): - """Evaluate the Gauss-Hermite profile. + def __init__(self, velocity_scale, sigma_truncation=5.0): + """Create a Gaussian LOSVD kernel. Parameters ---------- - x : array_like - Coordinate values where the profile is evaluated. - *params - Gaussian parameters followed by Hermite coefficients. + velocity_scale : float + Velocity step represented by one spectral pixel. + sigma_truncation : float, optional + Kernel half-width in units of sigma when building finite support. + """ + super().__init__(velocity_scale) + self.sigma_truncation = float(sigma_truncation) - Returns - ------- - array_like - Profile values at ``x``. + def parse_parameters(self, datablock): + vel = datablock["kinematics", "los_vel"] + sigma = datablock["kinematics", "los_sigma"] + self.set_parameters(vel, sigma) + + def set_parameters(self, vel, sigma): + """Set Gaussian LOSVD parameters and build/cache kernel weights. + + Parameters + ---------- + vel : float + Mean LOS velocity. + sigma : float + LOS velocity dispersion. """ - a, m, s = params[:3] # amplitude, mean, stddev - f = self._gaussian.evaluate(x, a, m, s) - if self._order: - f *= 1 + self._hermite.evaluate((x - m) / s, 0, 0, 0, *params[3:]) - - return f - - -# TODO : remove and homogeneize -def losvd(vel_pixel, sigma_pixel, h3=0, h4=0): - """Evaluate a Gauss-Hermite line-of-sight velocity distribution kernel.""" - - y = vel_pixel / sigma_pixel - g = ( - np.exp(-(y**2) / 2) - / sigma_pixel - / np.sqrt(2 * np.pi) - * ( - 1 - + h3 * (y * (2 * y**2 - 3) / np.sqrt(3)) # H3 - + h4 * ((4 * (y**2 - 3) * y**2 + 3) / np.sqrt(24)) # H4 - ) - ) - return g + sigma_pixel = sigma / self.velocity_scale + vel_pixel = vel / self.velocity_scale + if sigma_pixel <= 0: + self.edge_pixels = 0 + self.kernel_weight = np.array([1.0], dtype=float) + return -def get_losvd_kernel(kernel_model, x_size): - """Create a ``Model1DKernel`` from an input ``Model``. + half_width = max(1, int(np.ceil(self.sigma_truncation * sigma_pixel + np.abs(vel_pixel)))) + self.edge_pixels = int(np.ceil(self.sigma_truncation * sigma_pixel)) + key = (round(vel_pixel, CACHE_PIX_DECIMALS), round(sigma_pixel, CACHE_PIX_DECIMALS), half_width) - Parameters - ---------- - kernel_model : :class:`astropy.models.FittableModel` - Model used to build the kernel. - x_size : int - Kernel size + def build_kernel(): + x_edges = np.arange(-half_width - 0.5, half_width + 1.5, 1.0) + cmf = self.__cumulative_distribution(x_edges, sigma_pixel, vel_pixel) + return cmf[1:] - cmf[:-1] - Returns - ------- - kernel : :class:`Model1DKernel` - Kernel model - """ - ker = Model1DKernel(kernel_model, x_size=x_size, mode="integrate") - return ker + self._cached_kernel(key, build_kernel) + def __cumulative_distribution(self, x, sigma_pixel, vel_pixel): + return 0.5 * (1 + erf((x - vel_pixel) / (sigma_pixel * SQRT2))) -def convolve_spectra_with_kernel(spectra, kernel, use_fft=True): - """Convolve an input spectra with a given kernel. - Parameters - ---------- - spectra : np.ndarray - Target spectra to convolve with the kernel. - kernel : :class:`Model1DKernel` - Convolution kernel. +class SplitGaussianPixelKernel(LOSVDPixelKernel): + """Split-Gaussian LOSVD kernel in pixel space. - Returns - ------- - convolved_spectra : np.ndarray - Spectra convolved with the input kernel. + This kernel uses different velocity dispersions on blue and red sides + around the LOS velocity centroid. """ - # TODO: use np.convolve to increase performance when using - # synthetic observations that do not contain nan - try: - if use_fft: - return convolve_fft( - spectra, kernel, boundary="fill", fill_value=0.0, normalize_kernel=True - ) + def __init__(self, velocity_scale, sigma_truncation=5.0): + super().__init__(velocity_scale) + self.sigma_truncation = float(sigma_truncation) - else: - return convolve( - spectra, kernel, boundary="fill", fill_value=0.0, normalize_kernel=True - ) + def parse_parameters(self, datablock): + vel = datablock["kinematics", "los_vel"] + sigma_blue = datablock["kinematics", "los_sigma_blue"] + sigma_red = datablock["kinematics", "los_sigma_red"] + self.set_parameters(vel, sigma_blue, sigma_red) - except ValueError as e: - logging.error(f"Error during convolution: {e}") - return np.full(spectra.size, np.nan) - + def set_parameters(self, vel, sigma_blue, sigma_red): + """Set split-Gaussian LOSVD parameters and build/cache kernel weights. -def convolve_ssp_with_lsf(ssp, lsf_sigma_pixels): - """Convolve a given SSP model with an LSF. - - Parameters - ---------- - ssp : pst.SSP.SSPBase - lsf_sigma_pixels : float, np.ndarray - Gaussian standard deviation of the LSF. + Parameters + ---------- + vel : float + Mean LOS velocity. + sigma_blue : float + Dispersion used for pixels blueward of the centroid. + sigma_red : float + Dispersion used for pixels redward of the centroid. + """ + sigma_blue_px = sigma_blue / self.velocity_scale + sigma_red_px = sigma_red / self.velocity_scale + vel_pixel = vel / self.velocity_scale + + if sigma_blue_px <= 0 or sigma_red_px <= 0: + self.edge_pixels = 0 + self.kernel_weight = np.array([1.0], dtype=float) + return + + sigma_max = max(sigma_blue_px, sigma_red_px) + half_width = max(1, int(np.ceil(self.sigma_truncation * sigma_max + np.abs(vel_pixel)))) + self.edge_pixels = int(np.ceil(self.sigma_truncation * sigma_max)) + key = ( + round(vel_pixel, CACHE_PIX_DECIMALS), + round(sigma_blue_px, CACHE_PIX_DECIMALS), + round(sigma_red_px, CACHE_PIX_DECIMALS), + half_width, + ) + + def build_kernel(): + x = np.arange(-half_width, half_width + 1, dtype=float) - vel_pixel + sigma = np.where(x < 0.0, sigma_blue_px, sigma_red_px) + w = np.exp(-0.5 * (x / sigma) ** 2) / (sigma * SQRT2PI) + return np.where(np.isfinite(w), w, 0.0) + + self._cached_kernel(key, build_kernel) + + +class GaussHermitePixelKernel(LOSVDPixelKernel): + """Gauss-Hermite LOSVD kernel in pixel space. + + The profile is controlled by mean velocity, velocity dispersion and + optional third/fourth-order Hermite moments ``h3`` and ``h4``. """ - # Account for LSF variability - if lsf_sigma_pixels.size == ssp.wavelength.size: - ssp.L_lambda = np.array([ - fftconvolve( - ssp.L_lambda.value, - losvd(np.arange(-CONFIG.kinematics["lsf_sigma_truncation"] * sigma, - CONFIG.kinematics["lsf_sigma_truncation"] * sigma), - sigma_pixel=sigma)[np.newaxis, np.newaxis], - mode="same", axes=2)[:, :, ith] - for ith, sigma in enumerate(lsf_sigma_pixels)]) * ssp.L_lambda.unit - # Constant LSF - elif np.atleast_1d(lsf_sigma_pixels).size == 1: - ssp.L_lambda = fftconvolve( - ssp.L_lambda.value, - losvd(np.arange(-CONFIG.kinematics["lsf_sigma_truncation"] * lsf_sigma_pixels, - CONFIG.kinematics["lsf_sigma_truncation"] * lsf_sigma_pixels), - sigma_pixel=lsf_sigma_pixels)[np.newaxis, np.newaxis], - mode="same", axes=2) * ssp.L_lambda.unit - else: - raise ArithmeticError("Dimensions of SSP and LSF do not match") - -def convolve_ssp(module_config, los_sigma, los_vel, los_h3=0.0, los_h4=0.0): - """Convolve SSP spectra stored in a module configuration with LOS kinematics.""" - - velscale = module_config["velscale"] - extra_pixels = module_config["extra_pixels"] - ssp_sed = module_config["ssp_sed"] - flux = module_config["flux"] - if los_sigma <= 0: - raise ValueError("los_sigma must be positive for convolution.") - if np.abs(los_h3) > 0.5 or np.abs(los_h4) > 0.5: - raise ValueError("Gauss-Hermite coefficients h3/h4 are out of bounds (|h|<=0.5).") - # Kinematics - sigma_pixel = los_sigma / velscale - veloffset_pixel = los_vel / velscale - x = ( - np.arange( - -CONFIG.kinematics["lsf_sigma_truncation"] * sigma_pixel, - CONFIG.kinematics["lsf_sigma_truncation"] * sigma_pixel, + + def __init__(self, velocity_scale, sigma_truncation=5.0): + super().__init__(velocity_scale) + self.sigma_truncation = float(sigma_truncation) + + def parse_parameters(self, datablock): + self.set_parameters( + float(datablock["kinematics", "los_vel"]), + float(datablock["kinematics", "los_sigma"]), + h3=float(datablock["kinematics", "los_h3"]), + h4=float(datablock["kinematics", "los_h4"]), ) - - veloffset_pixel - ) - losvd_kernel = losvd(x, sigma_pixel=sigma_pixel, h3=los_h3, h4=los_h4) - sed = fftconvolve(ssp_sed, np.atleast_2d(losvd_kernel), mode="same", axes=1) - # Rebin model spectra to observed grid - sed = sed[:, extra_pixels:-extra_pixels] - ### Mask pixels at the edges with artifacts produced by the convolution - mask = np.ones_like(flux, dtype=bool) - mask[: int(CONFIG.kinematics["lsf_sigma_truncation"] * sigma_pixel)] = False - mask[-int(CONFIG.kinematics["lsf_sigma_truncation"] * sigma_pixel) :] = False - return sed, mask - - -def convolve_ssp_model(module_config, los_sigma, los_vel, h3=0.0, h4=0.0): - """Convolve an SSP model instance in place with LOS kinematics.""" - - velscale = module_config["velscale"] - extra_pixels = int(module_config["extra_pixels"]) - ssp = module_config["ssp_model"] - wl = module_config["wavelength"] - if los_sigma <= 0: - raise ValueError("los_sigma must be positive for convolution.") - if np.abs(h3) > 0.5 or np.abs(h4) > 0.5: - raise ValueError("Gauss-Hermite coefficients h3/h4 are out of bounds (|h|<=0.5).") - # Kinematics - sigma_pixel = los_sigma / velscale - veloffset_pixel = los_vel / velscale - x = ( - np.arange( - -CONFIG.kinematics["lsf_sigma_truncation"] * sigma_pixel, - CONFIG.kinematics["lsf_sigma_truncation"] * sigma_pixel, + + def set_parameters(self, vel, sigma, h3=0.0, h4=0.0): + """Set Gauss-Hermite LOSVD parameters and build/cache kernel weights. + + Parameters + ---------- + vel : float + Mean LOS velocity. + sigma : float + LOS velocity dispersion. + h3 : float, optional + Third-order Gauss-Hermite coefficient. + h4 : float, optional + Fourth-order Gauss-Hermite coefficient. + """ + sigma_pixel = float(sigma) / self.velocity_scale + vel_pixel = float(vel) / self.velocity_scale + + if sigma_pixel <= 0: + self.edge_pixels = 0 + self.kernel_weight = np.array([1.0], dtype=float) + return + + half_width = max( + 1, + int(np.ceil(self.sigma_truncation * sigma_pixel + np.abs(vel_pixel))), ) - - veloffset_pixel - ) - losvd_kernel = losvd(x, sigma_pixel=sigma_pixel, h3=h3, h4=h4) - ssp.L_lambda = ( - fftconvolve( - ssp.L_lambda.value, - losvd_kernel[np.newaxis, np.newaxis], - mode="same", - axes=2, + self.edge_pixels = int(np.ceil(self.sigma_truncation * sigma_pixel)) + key = ( + round(vel_pixel, CACHE_PIX_DECIMALS), + round(sigma_pixel, CACHE_PIX_DECIMALS), + round(float(h3), CACHE_PIX_DECIMALS), + round(float(h4), CACHE_PIX_DECIMALS), + half_width, ) - * ssp.L_lambda.unit - ) - # Rebin model spectra to observed grid - pixels = slice(extra_pixels, -extra_pixels) - new_sed = ssp.L_lambda[:, :, pixels] - ssp.L_lambda = new_sed - if not isinstance(wl, u.Quantity): - ssp.wavelength = wl * ssp.wavelength.unit - else: - ssp.wavelength = wl - ### Mask pixels at the edges with artifacts produced by the convolution - mask = np.ones(wl.size, dtype=bool) - mask[: int(CONFIG.kinematics["lsf_sigma_truncation"] * sigma_pixel)] = False - mask[-int(CONFIG.kinematics["lsf_sigma_truncation"] * sigma_pixel) :] = False - return ssp, mask + + def build_kernel(): + x = np.arange(-half_width, half_width + 1, dtype=float) - vel_pixel + w = self.__losvd(x, sigma_pixel=sigma_pixel, h3=float(h3), h4=float(h4)) + w = np.where(np.isfinite(w), w, 0.0) + if np.sum(w) <= 0: + w = np.exp(-0.5 * (x / sigma_pixel) ** 2) / (sigma_pixel * SQRT2PI) + return w + + self._cached_kernel(key, build_kernel) + + def __losvd(self, vel_pixel, sigma_pixel, h3=0, h4=0): + """Evaluate a Gauss-Hermite line-of-sight velocity distribution kernel.""" + + y = vel_pixel / sigma_pixel + + g = (np.exp(-(y**2) / 2) / sigma_pixel / SQRT2PI + * ( + 1 + + h3 * (y * (2 * y**2 - 3) / np.sqrt(3)) + + h4 * ((4 * (y**2 - 3) * y**2 + 3) / np.sqrt(24)) + ) + ) + return g + +class PieceWisePixelKernel(LOSVDPixelKernel): + """Piecewise-constant LOSVD kernel defined in velocity bins. + + The kernel is specified by sampled bin weights in velocity space and then + interpolated onto the module pixel grid. + """ + + def __init__(self, velocity_scale, velocity_bin_size, velocity_min, velocity_max): + """Create a piecewise LOSVD kernel parameterization. + + Parameters + ---------- + velocity_scale : float + Velocity step represented by one spectral pixel. + velocity_bin_size : float + Width of input velocity bins. + velocity_min : float + Lower bound of the piecewise velocity grid. + velocity_max : float + Upper bound of the piecewise velocity grid. + """ + super().__init__(velocity_scale) + self.velocity_bin_size = velocity_bin_size + if velocity_min >= velocity_max: + raise ValueError("velocity_min must be less than velocity_max.") + self.velocity_bin_edges = np.arange( + velocity_min, velocity_max + velocity_bin_size, velocity_bin_size) + # Cache for querying the datablock + self.bin_ids = np.arange(0, self.velocity_bin_edges.size - 1, 1) + pixel_min = velocity_min / self.velocity_scale + pixel_max = velocity_max / self.velocity_scale + # Ensure kernel array is odd, symmetric and covers both edges + edges = np.abs([pixel_min, pixel_max]).max() + self.x_pixel_edges = np.arange(-edges - 0.5, edges + 1.5, 1.0) + self.x_vel_edges = self.x_pixel_edges * self.velocity_scale + + def parse_parameters(self, datablock): + weights = np.asarray([datablock["kinematics", f"vel_bin_{ith}"] for ith in self.bin_ids], dtype=float) + # resample into the velocity scale pixel grid + cum_kernel = np.cumsum(weights) + cum_kernel = np.insert(cum_kernel, 0, 0.0) # add zero at the beginning for interpolation + cum_kernel = np.interp( + self.x_vel_edges, + self.velocity_bin_edges, + cum_kernel, + left=0.0, + right=cum_kernel[-1], + ) + self.edge_pixels = int(np.ceil(np.max(np.abs(self.x_pixel_edges)))) + self.kernel_weight = np.diff(cum_kernel) + def normal_cdf(x, mu=0.0, sigma=1.0): - """Normal cumulative density function.""" - return (1.0 + erf((x - mu) / sigma / np.sqrt(2.0))) / 2.0 + """Normal cumulative density function. -# def convolve_variable_gaussian_kernel(spectra, sigma_pixel, -# kappa_sigma_thresh=3.0): -# """Convolve an input spectra with a Gaussian kernel of varying width. - -# Parameters -# ---------- -# spectra : :class:`np.ndarray` or :class:`u.Quantity` -# N-dimensional spectra. The last dimension must correspond to the -# wavelength axis. -# sigma_pixel : :class:`np.ndarray` -# Value of the standard deviation of the Gaussian LSF for each spectral -# resolution element, expressed in pixel units. -# kappa_sigma_thresh : float, optional -# Threshold in units of the Gaussian sigma. If a **single pixel** contains -# at least ``kappa_sigma_thresh`` sigmas of the LSF, that pixel is left -# untouched (i.e. the convolution kernel is clamped to a delta function -# at that pixel). Default is ``3.0``. - -# Returns -# ------- -# convolved_spectra : :class:`np.ndarray` or :class:`u.Quantity` -# A convolved version of ``spectra``. -# """ -# # Convert sigma_pixel to a plain ndarray (it may be a Quantity or list) -# sigma_pixel = np.asarray(sigma_pixel, dtype=float) -# n_pix = sigma_pixel.size - -# # Identify pixels where the LSF is effectively contained within a single pixel: -# # 0.5 pixel half-width contains >= kappa_sigma_thresh sigmas. -# # 0.5 / sigma >= kappa -> sigma <= 0.5 / kappa -# small_lsf_mask = (0.5 / sigma_pixel) >= kappa_sigma_thresh -# # Guard against sigma == 0 as well (also treated as delta kernel) -# small_lsf_mask |= (sigma_pixel <= 0.0) - -# # Use a "safe" sigma for computing the CDF to avoid division issues; -# # rows that will be clamped later can use any finite sigma. -# safe_sigma = sigma_pixel.copy() -# safe_sigma[small_lsf_mask] = 1.0 - -# # mean_pixel has shape (n_pix, n_pix + 1) and represents bin edges -# mean_pixel = (np.arange(-0.5, n_pix + 0.5, 1)[np.newaxis, :] -# - np.arange(0, n_pix, 1)[:, np.newaxis]) - -# cmf = normal_cdf(mean_pixel / safe_sigma[:, np.newaxis]) -# weights = cmf[:, 1:] - cmf[:, :-1] # (n_pix, n_pix) - -# # Normalise each row -# weights /= np.sum(weights, axis=1)[:, np.newaxis] - -# # Clamp rows where the LSF is narrower than a pixel: -# # replace the whole row by a delta kernel. -# if np.any(small_lsf_mask): -# weights[small_lsf_mask, :] = 0.0 -# idx = np.nonzero(small_lsf_mask)[0] -# # Put the delta on the diagonal: output[i] = input[i] -# weights[idx, idx] = 1.0 - -# # Broadcast to match the input spectra shape (last axis is spectral axis) -# extra_dim = spectra.ndim - 1 -# if extra_dim > 0: -# axis = [dim for dim in np.arange(0, extra_dim, 1)] -# weights = np.expand_dims(weights, axis=axis) - -# # Perform the (variable) convolution along the last axis -# return np.sum(np.expand_dims(spectra, axis=-1) * weights, axis=-1) + Parameters + ---------- + x : array-like + Evaluation points. + mu : float, optional + Mean of the normal distribution. + sigma : float, optional + Standard deviation of the normal distribution. + + Returns + ------- + np.ndarray + CDF values evaluated at ``x``. + """ + return (1.0 + erf((x - mu) / sigma / np.sqrt(2.0))) / 2.0 def convolve_variable_gaussian_kernel( spectra, diff --git a/src/besta/pipeline.py b/src/besta/pipeline.py index 49cd7014..a0e33692 100644 --- a/src/besta/pipeline.py +++ b/src/besta/pipeline.py @@ -4,8 +4,11 @@ import os import subprocess +import copy import numpy as np +from typing import List, Dict + from matplotlib import pyplot as plt from cosmosis import DataBlock @@ -102,15 +105,16 @@ def execute_pipeline( ) io.make_ini_file(ini_filename, config) else: - assert os.path.isfile(os.path.expandvars(ini_filename) - ), f"{os.path.expandvars(ini_filename)} not found" + ini_filename = os.path.expandvars(ini_filename) + if not os.path.isfile(ini_filename): + raise FileNotFoundError(f"{ini_filename} not found") if ini_values_filename is None: io.make_values_file(config) else: - assert os.path.isfile( - ini_values_filename - ), f"{ini_values_filename} not found" + ini_values_filename = os.path.expandvars(ini_values_filename) + if not os.path.isfile(ini_values_filename): + raise FileNotFoundError(f"{ini_values_filename} not found") config["pipeline"]["values"] = ini_values_filename if n_cores == -1: @@ -141,11 +145,16 @@ def execute_all(self, plot_result=False): ): if prev_solution is not None: logger.info("Updating configuration file with previous run results") + module_name = subpipe_config["pipeline"]["modules"].replace(",", " ").split()[0] + if module_name not in subpipe_config: + raise KeyError( + f"Module '{module_name}' not found in subpipeline configuration." + ) # Update the input values - subpipe_config[subpipe_config["pipeline"]["modules"]].update( + subpipe_config[module_name].update( (k, v) for k, v in prev_solution.items() - if k in subpipe_config[subpipe_config["pipeline"]["modules"]] + if k in subpipe_config[module_name] ) # Execute sub-pipepline ini_filename = self.execute_pipeline( @@ -183,3 +192,231 @@ def execute_all(self, plot_result=False): pipeline_module.plot_solution(solution_datablock, figname=figname) return 0 + + +# Run multiple independent pipelines in parallel +def _run_main_pipeline_job(job): + """Worker helper used by multiprocessing to run one MainPipeline.""" + pipeline = MainPipeline( + job["pipeline_config"], + n_cores_list=job["n_cores_list"], + ini_files=job["ini_files"], + ini_values_files=job["ini_values_files"], + ) + status = pipeline.execute_all(plot_result=job["plot_result"]) + return {"index": job["index"], "status": status} + + +class BatchPipeline(object): + """Manager for running multiple independent MainPipeline instances in parallel. + + Parameters + ---------- + pipeline_configuration_list : list[list[dict]] + Each entry is the configuration list required to initialize a MainPipeline. + n_cores_list : list, optional + Per-pipeline n_cores lists to pass to each MainPipeline. + ini_files : list, optional + Per-pipeline ini_files lists to pass to each MainPipeline. + ini_values_files : list, optional + Per-pipeline ini_values_files lists to pass to each MainPipeline. + n_jobs_parallel : int, optional + Maximum number of parallel processes. + """ + + def __init__( + self, + *, + pipeline_configuration_list: List[List[Dict]], + n_cores_list=None, + ini_files=None, + ini_values_files=None, + n_jobs_parallel=None, + ): + # Store the list of independent pipeline configurations. + self.all_pipelines_config = pipeline_configuration_list + + n_pipelines = len(self.all_pipelines_config) + + self.all_n_cores_list = ( + [None] * n_pipelines if n_cores_list is None else n_cores_list + ) + self.all_ini_files = [None] * n_pipelines if ini_files is None else ini_files + self.all_ini_values_files = ( + [None] * n_pipelines if ini_values_files is None else ini_values_files + ) + self.n_jobs_parallel = n_jobs_parallel + + self._validate_inputs() + + def _validate_inputs(self): + """Validate top-level list lengths.""" + n_pipelines = len(self.all_pipelines_config) + for attr_name in ( + "all_n_cores_list", + "all_ini_files", + "all_ini_values_files", + ): + value = getattr(self, attr_name) + if len(value) != n_pipelines: + raise ValueError( + f"{attr_name} length ({len(value)}) must match " + f"number of pipelines ({n_pipelines})" + ) + + def cpu_info(self): + """Log available CPU cores and number of independent pipelines.""" + n_pipelines = len(self.all_pipelines_config) + n_cores = os.cpu_count() or 1 + logger.info(f"Available CPU cores: {n_cores}") + logger.info(f"Number of pipelines to run: {n_pipelines}") + + def build_pipelines(self): + """Build all MainPipeline objects without executing them.""" + pipelines = [] + for pipeline_config, n_cores, ini_file, ini_values_file in zip( + self.all_pipelines_config, + self.all_n_cores_list, + self.all_ini_files, + self.all_ini_values_files, + ): + pipelines.append( + MainPipeline( + pipeline_config, + n_cores_list=n_cores, + ini_files=ini_file, + ini_values_files=ini_values_file, + ) + ) + return pipelines + + def _build_jobs(self, plot_result=False): + jobs = [] + for index, (pipeline_config, n_cores, ini_file, ini_values_file) in enumerate( + zip( + self.all_pipelines_config, + self.all_n_cores_list, + self.all_ini_files, + self.all_ini_values_files, + ) + ): + jobs.append( + { + "index": index, + "pipeline_config": pipeline_config, + "n_cores_list": n_cores, + "ini_files": ini_file, + "ini_values_files": ini_values_file, + "plot_result": plot_result, + } + ) + return jobs + + def run_single_pipeline(self, index, plot_result=False): + """Run one independent pipeline by index.""" + jobs = self._build_jobs(plot_result=plot_result) + if index < 0 or index >= len(jobs): + raise IndexError(f"Pipeline index {index} out of range") + return _run_main_pipeline_job(jobs[index])["status"] + + def run_all_pipelines(self, plot_result=False): + """Run all configured MainPipeline instances in parallel.""" + from multiprocessing import get_context + + self.cpu_info() + jobs = self._build_jobs(plot_result=plot_result) + if not jobs: + logger.info("No pipelines to run") + return [] + + max_jobs = os.cpu_count() or 1 + if self.n_jobs_parallel is None: + n_jobs = min(len(jobs), max_jobs) + else: + n_jobs = max(1, min(self.n_jobs_parallel, len(jobs), max_jobs)) + + logger.info(f"Running {n_jobs} pipelines in parallel") + + if n_jobs == 1: + raw_results = [_run_main_pipeline_job(job) for job in jobs] + else: + with get_context("spawn").Pool(processes=n_jobs) as pool: + raw_results = pool.map(_run_main_pipeline_job, jobs) + + raw_results.sort(key=lambda x: x["index"]) + results = [item["status"] for item in raw_results] + + n_failed = sum(status != 0 for status in results) + if n_failed > 0: + logger.error(f"{n_failed}/{len(results)} pipelines failed") + else: + logger.info("All pipelines executed successfully") + + return results + + @classmethod + def from_running_parameters( + cls, + pipeline_configuration: Dict, + running_parameters: List[Dict], + **kwargs, + ): + """Factory method to create a BatchPipeline from a single pipeline configuration and multiple running parameters. + + Parameters + ---------- + pipeline_configuration : dict + Base configuration for the pipeline, which will be updated with each set of running parameters. + running_parameters : list[dict] + List of dictionaries containing the parameters to update in the base configuration for each independent pipeline run. + **kwargs + Additional keyword arguments to pass to the BatchPipeline constructor (e.g., n_cores_list + or ini_files). These will be applied to all pipelines created from the running parameters. + + Returns + ------- + BatchPipeline + An instance of BatchPipeline configured to run multiple pipelines based on the provided configuration and parameters. + + Example + ------- + >>> from besta.pipeline import BatchPipeline + >>> base_config = { + ... "pipeline": {"modules": "FullSpectralFit", "values": "./values.ini"}, + ... "output": {"filename": "./fit_result"}, + ... "FullSpectralFit": {"file": "/path/to/module.py", "redshift": 0.1}, + ... } + >>> running_parameters = [ + ... {"FullSpectralFit": {"redshift": 0.10}, "output": {"filename": "./fit_z010"}}, + ... {"FullSpectralFit": {"redshift": 0.12}, "output": {"filename": "./fit_z012"}}, + ... ] + >>> batch = BatchPipeline.from_running_parameters( + ... pipeline_configuration=base_config, + ... running_parameters=running_parameters, + ... n_jobs_parallel=2, + ... ) + >>> results = batch.run_all_pipelines() + >>> len(results) + 2 + """ + if not isinstance(running_parameters, list): + raise TypeError("running_parameters must be a list of dictionaries") + + pipeline_config_list = [] + + # Recursively merge nested dictionaries from one parameter set. + def recursive_update(d, u): + for k, v in u.items(): + if isinstance(v, dict): + d[k] = recursive_update(d.get(k, {}), v) + else: + d[k] = v + return d + + for params in running_parameters: + if not isinstance(params, dict): + raise TypeError("Each running parameter entry must be a dictionary") + config_copy = copy.deepcopy(pipeline_configuration) + config_copy = recursive_update(config_copy, params) + pipeline_config_list.append([config_copy]) # Wrap in list for MainPipeline + return cls(pipeline_configuration_list=pipeline_config_list, **kwargs) \ No newline at end of file diff --git a/src/besta/pipeline_modules/base_module.py b/src/besta/pipeline_modules/base_module.py index 1f188055..5d2776c2 100644 --- a/src/besta/pipeline_modules/base_module.py +++ b/src/besta/pipeline_modules/base_module.py @@ -18,7 +18,6 @@ from besta.visualization import draw_dict_in_axes import numpy as np -from scipy.stats import norm from sklearn.decomposition import NMF from astropy import units as u from astropy.io import ascii @@ -38,6 +37,7 @@ from besta import io from besta import utils from besta.grid import ModelGrid +from . import likelihoods from besta.config import cosmology, memory from besta.logging import get_logger, setup_logging @@ -50,7 +50,10 @@ def _log(*args): class BaseModule(ClassModule): """BESTA Pipeline module base class.""" - def __init__(self, options, *, alias=None): + _default_flux_units = "1e-16 erg / (s cm2 Angstrom)" + _default_luminosity_units = "1e-16 erg / (s Angstrom)" + + def __init__(self, options, *, alias=None, likelihood_kind=None, likelihood_method=None): """ Set up the CosmoSIS module. @@ -101,6 +104,30 @@ def __init__(self, options, *, alias=None): self.like_name = self.name + "_like" _log("Setting module likelihood name to default: ", self.like_name) + if likelihood_kind is None: + name = (self.__class__.__name__ + " " + self.name).lower() + if "photometry" in name: + likelihood_kind = "photometry" + else: + likelihood_kind = "spectra" + + if options.has_value("likelihood_kind"): + likelihood_kind = options.get_string("likelihood_kind", default=likelihood_kind) + if options.has_value("likelihood_method"): + likelihood_method = options.get_string("likelihood_method", default="auto") + + self.likelihood_kind = str(likelihood_kind).strip().lower() + self.likelihood_method = str(likelihood_method or "auto").strip().lower() + + if self.likelihood_kind == "photometry": + self.log_like = likelihoods.make_photometry_loglike(self.likelihood_method) + elif self.likelihood_kind == "spectra": + self.log_like = likelihoods.make_spectra_loglike(self.likelihood_method) + else: + raise ValueError( + f"Unknown likelihood_kind={self.likelihood_kind!r}; expected 'spectra' or 'photometry'." + ) + @abstractmethod def make_observable(self, *args, **kwargs): """Create an observable from an input set of model parameters.""" @@ -162,10 +189,10 @@ def prepare_ssp_model(self, options, normalize=False, velocity_buffer=800.0): during convolution. The buffer is applied to both sides of the SSP spectra. """ - _log("\n-> Configuring SSP model") + _log("Configuring SSP model") if options.has_value("SSPModelFromPickle"): - _log("\n-> Loading preconfigured SSP model from pickle") + _log("Loading preconfigured SSP model from pickle") if not os.path.isfile( os.path.expandvars(options["SSPModelFromPickle"])): raise FileNotFoundError( @@ -186,7 +213,7 @@ def prepare_ssp_model(self, options, normalize=False, velocity_buffer=800.0): extra_offset_pixel = int(velocity_buffer / velscale) self.config["velscale"] = velscale self.config["extra_pixels"] = extra_offset_pixel - _log("-> Configuration done.") + _log("Configuration done.") return ssp_name = options["SSPModel"] @@ -260,22 +287,32 @@ def prepare_ssp_model(self, options, normalize=False, velocity_buffer=800.0): # Convolve with instrumental LSF if "lsf" in self.config: _log("Convolving SSP model with instrumental LSF") - inst_lsf = np.interp(ssp.wavelength, self.config["wavelength"], - self.config["lsf"]) + # A given rest-frame wavelength is observed at a redshifted wavelength + # so the instrumental LSF in the rest frame is given by interpolating + # the input LSF at the observed wavelength + inst_lsf = np.interp( + ssp.wavelength, + self.config["wavelength"] * (1 + self.config["redshift"]), + self.config["lsf"]) if options.has_value("SSPLSF"): _log("Including SSP resolution") + # Load SSP LSF in the rest frame ssp_lsf_wl, ssp_lsf_fwhm = np.loadtxt( os.path.expandvars(options["SSPLSF"]), unpack=True, usecols=(0, 1)) - ssp_lsf_fwhm = np.interp(ssp.wavelength, - ssp_lsf_wl << u.AA, ssp_lsf_fwhm) + ssp_lsf_fwhm = np.interp(ssp.wavelength, ssp_lsf_wl << u.AA, + ssp_lsf_fwhm) else: ssp_lsf_fwhm = np.zeros(ssp.wavelength.size, dtype=float) # Assume both LSF are Gaussian effective_lsf_disp = (inst_lsf / 2.355)**2 - (ssp_lsf_fwhm / 2.355)**2 if (effective_lsf_disp < 0).any(): + logger.error("Effective LSF dispersion has negative values. Check the input instrumental and SSP LSFs.") + logger.debug("Instrumental LSF (FWHM): %s", inst_lsf) + logger.debug("SSP LSF (FWHM): %s", ssp_lsf_fwhm) + logger.debug("Effective LSF dispersion: %s", effective_lsf_disp) raise ValueError("Effective SSP LSF cannot be negative!" + "SSP models do not have enough resolution") effective_lsf = np.sqrt(effective_lsf_disp) @@ -315,7 +352,7 @@ def prepare_ssp_model(self, options, normalize=False, velocity_buffer=800.0): if options.has_value("SaveSSPModel"): _log("Saving SSP model to ", options["SaveSSPModel"]) ssp.to_pickle(os.path.expandvars(options["SaveSSPModel"])) - _log("-> Configuration done.") + _log("Configuration done.") return def prepare_extinction_law(self, options): @@ -332,7 +369,7 @@ def prepare_extinction_law(self, options): _log("Extinction law: ", ext_law) # TODO: add more extinction laws self.config["extinction_law"] = dust.DustScreen(ext_law) - _log("-> Configuration is done.") + _log("Configuration is done.") def prepare_sfh_model(self, options): """Prepare the SFH model. @@ -342,7 +379,7 @@ def prepare_sfh_model(self, options): options : :class:`DataBlock` Input options to initialise the model. """ - _log("\n-> Configuring SFH model") + _log("Configuring SFH model") sfh_model_name = options["SFHModel"] sfh_args = [] sfh_kwargs = {} @@ -378,103 +415,7 @@ def prepare_sfh_model(self, options): sfh_model = getattr(sfh, sfh_model_name) sfh_model = sfh_model(*sfh_args, **sfh_kwargs, **self.config) self.config["sfh_model"] = sfh_model - _log("-> Configuration done") - - def log_like(self, data, model, var, weights=None, is_upper=None, is_lower=None, include_norm=True): - """Compute log-likelihood between data and model. - - Parameters - ---------- - data : np.ndarray - For detections: measured values. - For limits: the limit value (upper or lower). - model : np.ndarray - Model prediction for each datum. - var : np.ndarray - Data variance. Must match shape of data/model. - weights : np.ndarray, optional - Data weights. If provided, returns weighted mean log-likelihood. - is_upper : np.ndarray[bool], optional - Mask for upper limits (x < data). - is_lower : np.ndarray[bool], optional - Mask for lower limits (x > data). - include_norm : bool, optional, default=True - If True, include Gaussian normalization terms for detections: - -0.5*log(2*pi*var). - - Returns - ------- - loglike : float - Total (or weighted-mean) log-likelihood. - """ - if data.shape != model.shape or data.shape != var.shape: - raise ValueError("data, model, var must have the same shape (var is per-datum variance).") - if np.any(var <= 0): - raise ValueError("All var entries must be > 0 (variance).") - - if is_upper is None: - is_upper = np.zeros_like(data, dtype=bool) - else: - is_upper = np.asarray(is_upper, dtype=bool) - - if is_lower is None: - is_lower = np.zeros_like(data, dtype=bool) - else: - is_lower = np.asarray(is_lower, dtype=bool) - - if is_upper.shape != data.shape or is_lower.shape != data.shape: - raise ValueError("is_upper and is_lower must have the same shape as data/model.") - if np.any(is_upper & is_lower): - raise ValueError("A data point cannot be both an upper and a lower limit.") - - if weights is None: - weights = np.ones_like(data, dtype=float) - normalize = False - else: - weights = np.asarray(weights, dtype=float) - if weights.shape != data.shape: - raise ValueError("weights must have the same shape as data/model.") - if np.any(weights < 0): - raise ValueError("weights must be non-negative.") - normalize = True - - # TODO: to avoid this step the pipeline should store sigma - sigma = np.sqrt(var) - - logp = np.empty_like(data, dtype=float) - det = ~(is_upper | is_lower) - - # For detections: Gaussian logpdf - if np.any(det): - if include_norm: - logp[det] = norm.logpdf(data[det], loc=model[det], scale=sigma[det]) - else: - z = (data[det] - model[det]) / sigma[det] - logp[det] = -0.5 * z**2 - - # Upper limits: P(x < L) - if np.any(is_upper): - z_u = (data[is_upper] - model[is_upper]) / sigma[is_upper] - logp[is_upper] = norm.logcdf(z_u) - - # Lower limits: P(x > L) = 1 - P(x < L) - if np.any(is_lower): - z_l = (data[is_lower] - model[is_lower]) / sigma[is_lower] - logp[is_lower] = norm.logsf(z_l) # 1 - logcdf - - # User-provided weighted mean - if normalize: - wsum = np.sum(weights) - if wsum <= 0: - if np.all(weights == 0): - raise ValueError("All weights are zero.") - else: - raise ValueError("Sum of weights must be > 0 for normalization.") - return np.sum(logp * weights) / wsum - - return np.sum(logp * weights) - - + _log("Configuration done") class SpectraFitModule(BaseModule): """Base class for spectral fitting modules in BESTA.""" @@ -489,13 +430,17 @@ def prepare_observed_spectra( normalize : bool, optional If ``True``, normalizes the spectra using the given wavelength range. """ - _log("\n-> Configuring input observed spectra") + _log("Configuring input observed spectra") filename = os.path.expandvars(options["inputSpectrum"]) # Read wavelength and spectra _log("Loading observed spectra from input file: ", filename) wavelength, flux, error = np.loadtxt(filename, unpack=True) - _log("Wavelength coverage: ", wavelength[[0, -1]]) - _log("Size: ", wavelength.size) + + if options.get_bool("is_variance", default=False): + error = np.sqrt(error) + elif options.get_bool("is_inverse_variance", default=False): + error = np.divide(1.0, error**0.5, out=error, where=error > 0) + error[error <= 0] = np.inf # Convert units if needed if options.has_value("wlUnits"): @@ -507,21 +452,21 @@ def prepare_observed_spectra( wl_units = u.angstrom if options.has_value("fluxUnits"): - _log("Converting flux units to 1e-16 erg/s/cm^2/Angstrom") + _log(f"Converting flux units to {self._default_flux_units}") flux_units = u.Unit(options["fluxUnits"]) - flux = (flux << flux_units).to( - "1e-16 erg / (s cm2 Angstrom)").value + flux = (flux << flux_units).to(self._default_flux_units).value error = (error << flux_units).to( - "1e-16 erg / (s cm2 Angstrom)").value + self._default_flux_units).value else: - _log("Assuming input flux units are in 1e-16 erg/s/cm^2/Angstrom") - flux_units = u.Unit("1e-16 erg / (s cm2 Angstrom)") + _log(f"Assuming input flux units are in {self._default_flux_units}") + flux_units = u.Unit(self._default_flux_units) # Wavelength range to include in the fit if options.has_value("wlRange"): wl_range = (np.asarray(options["wlRange"]) << wl_units ).to("Angstrom").value else: + _log("No input wavelength range provided; using full wavelength coverage") wl_range = wavelength[[0, -1]] # Wavelength range to renormalize the spectra if options.has_value("wlNormRange"): @@ -546,6 +491,28 @@ def prepare_observed_spectra( if weights.size != flux.size: raise ValueError( "Input mask size does not match the input spectrum size.") + + # Check for negative weights and wrong errors + if np.any(weights < 0): + raise ValueError("Input weights contain negative values.") + if np.any(error <= 0): + raise ValueError("Input errors contain negative values.") + + # non-finite numbers handling + mask_non_finite = options.get_bool("mask_non_finite", default=True) + if not np.isfinite(error).all() and mask_non_finite: + logger.warning("Input error array contains non-finite values." + "Setting weights to zero for those pixels.") + weights[~np.isfinite(error)] = 0.0 + # Set non-finite errors to the median of finite errors to avoid issues during interpolation or convolution + error[~np.isfinite(error)] = np.nanmedian(error[np.isfinite(error)]) + if not np.isfinite(flux).all() and mask_non_finite: + logger.warning("Input flux array contains non-finite values." + "Setting weights to zero for those pixels.") + weights[~np.isfinite(flux)] = 0.0 + # Set non-finite fluxes to the median of finite fluxes to avoid issues during interpolation or convolution + flux[~np.isfinite(flux)] = np.nanmedian(flux[np.isfinite(flux)]) + # Load the instrumental LSF if options.has_value("lsf"): lsf_wl, lsf_fwhm = np.loadtxt(os.path.expandvars(options["lsf"]), @@ -638,7 +605,21 @@ def prepare_observed_spectra( flux = flux_conserving_interpolation(ln_wave, np.log(wavelength), flux) cov = flux_conserving_interpolation(ln_wave, np.log(wavelength), cov) - weights = np.interp(ln_wave, np.log(wavelength), weights) + weights = np.interp(ln_wave, np.log(wavelength), weights).clip(0, None) + bad_cov = cov <= 0 + bad_flux = ~np.isfinite(flux) + if np.any(bad_cov): + _log( + f"Warning: some interpolated covariance values ({np.count_nonzero(bad_cov)}) are non-positive.", + "Setting weights to zero for those pixels.") + weights[bad_cov] = 0.0 + + if np.any(bad_flux): + _log( + f"Warning: some interpolated flux values ({np.count_nonzero(bad_flux)}) are non-finite.", + "Setting weights to zero for those pixels.") + weights[bad_flux] = 0.0 + instrumental_lsf = np.interp(ln_wave, np.log(wavelength), instrumental_lsf) new_wavelength = np.exp(ln_wave) @@ -669,6 +650,7 @@ def prepare_observed_spectra( self.config["flux"] = flux self.config["var"] = cov + self.config["ivar"] = np.divide(1.0, cov, out=np.zeros_like(cov), where=cov > 0) self.config["redshift"] = redshift self.config["wlUnits"] = wl_units self.config["fluxUnits"] = flux_units @@ -680,7 +662,10 @@ def prepare_observed_spectra( if not (instrumental_lsf == 0).all(): self.config["lsf"] = instrumental_lsf - _log("-> Configuration done.") + if options.get_bool("use_features", default=True): + self.get_feature_weights(options) + self.config["weights"] *= self.config["feature_weights"] + _log("Configuration done.") def prepare_galaxy(self, options): """Build and configure a :class:`pst.galaxy.GalaxySED` model. @@ -751,7 +736,7 @@ def prepare_legendre_polynomials(self, options): options : :class:`DataBlock` Input options to initialise the model. """ - _log("\n-> Configuring multiplicative polynomial") + _log("Configuring multiplicative polynomial") if options.has_value("legendre_deg"): kwargs = {} if options.has_value("legendre_bounds"): @@ -766,7 +751,88 @@ def prepare_legendre_polynomials(self, options): self.config["wavelength"], options["legendre_deg"], **kwargs) else: _log(f"Not using multiplicative Legendre polynomials") - _log("-> Configuration done") + _log("Configuration done") + + def prepare_losvd_kernel(self, options): + """Prepare the LOSVD convolution kernel. + + Parameters + ---------- + options : :class:`DataBlock` + Input options to initialise the model. + """ + _log("Configuring LOSVD convolution kernel") + # Get the velocity scale from options or previously prepared config. + if options.has_value("velscale"): + velocity_scale = options["velscale"] + self.config["velscale"] = velocity_scale + elif "velscale" in self.config: + velocity_scale = self.config["velscale"] + else: + raise ValueError("LOSVD convolution requires a defined velocity scale (velscale).") + + _log("Pixel velocity scale for LOSVD convolution: ", + velocity_scale, " km/s per pixel") + + kernel_name = options.get_string("losvd_kernel", default="gauss-hermite") + kernel_name = kernel_name.strip().lower().replace("_", "-") + _log("LOSVD kernel name: ", kernel_name) + + if kernel_name == "gaussian": + sigma_truncation = options.get_double("sigma_truncation", default=5.0) + self.config["losvd_kernel"] = kinematics.GaussianPixelKernel( + velocity_scale=velocity_scale, + sigma_truncation=sigma_truncation, + ) + + elif kernel_name in {"gauss-hermite", "gausshermite"}: + sigma_truncation = options.get_double("sigma_truncation", default=5.0) + self.config["losvd_kernel"] = kinematics.GaussHermitePixelKernel( + velocity_scale=velocity_scale, + sigma_truncation=sigma_truncation, + ) + + elif kernel_name in {"split-gaussian", "splitgaussian"}: + sigma_truncation = options.get_double("sigma_truncation", default=5.0) + self.config["losvd_kernel"] = kinematics.SplitGaussianPixelKernel( + velocity_scale=velocity_scale, + sigma_truncation=sigma_truncation, + ) + + elif kernel_name in {"piece-wise", "piecewise"}: + velocity_bin_size = options.get_int("velocity_bin_size", default=10) + velocity_min = options.get_double("velocity_min", default=-500.0) + velocity_max = options.get_double("velocity_max", default=500.0) + self.config["losvd_kernel"] = kinematics.PieceWisePixelKernel( + velocity_scale=velocity_scale, + velocity_bin_size=velocity_bin_size, + velocity_min=velocity_min, + velocity_max=velocity_max, + ) + else: + raise ValueError(f"Unsupported LOSVD kernel: {kernel_name}") + _log("Configuration done") + + def get_feature_weights(self, options): + """TODO""" + logger.info("Computing feature weights from input spectra") + # Estimate the continuum + continuum, continuum_err = spectrum.estimate_continuum( + self.config["wavelength"].to_value("AA"), + self.config["flux"], + err=self.config["var"]**0.5, + weights=self.config["weights"], + knot_spacing=options.get_double("continuum_knot_spacing", default=200.0), + sigma_clip=options.get_double("continuum_sigma_clip", default=3.0), + ) + self.config["continuum"] = continuum + self.config["continuum_err"] = continuum_err + # Favour features over/under continuum + weight_powlaw = options.get_double("feature_weight_powlaw", 2.0) + w = (np.abs(self.config["flux"] - continuum) / continuum_err)**weight_powlaw + w = np.where(np.isfinite(w), w, 0.0) + w /= w.max() + self.config["feature_weights"] = w def measure_emission_lines(self, solution: DataBlock, **kwargs): """Measure emission line fluxes and EWs from the best-fit solution. @@ -802,13 +868,14 @@ def measure_emission_lines(self, solution: DataBlock, **kwargs): def plot_solution(self, solution: DataBlock, figname=None, plot_lines=True): """Plot the fit.""" flux_model = self.make_observable(solution, parse=True) + continuum_model = self.config.get("continuum") + continuum_model_err = self.config.get("continuum_err") + if isinstance(flux_model, tuple): weights = flux_model[1] flux_model = flux_model[0] else: weights = np.ones_like(flux_model) - # Include input weights - weights *= self.config["weights"] # Grab the solution values (visualuzation purpose only) sol_keys = solution.keys() @@ -829,6 +896,14 @@ def plot_solution(self, solution: DataBlock, figname=None, plot_lines=True): # Display the information ax = axs[0, 1] # Pixel masking information + like_weights = self.config["weights"] + if "sweep_weights" in self.config: + sweep_weights = self.config["sweep_weights"] + weights *= sweep_weights + + like_eff_pixels = np.sum(like_weights > 0) + + mask_info = {"Total pixels": self.config["flux"].size, "Masked pixels (w=0)": np.sum(weights <= 0), " - Telluric abs.": np.sum( @@ -887,6 +962,19 @@ def plot_solution(self, solution: DataBlock, figname=None, plot_lines=True): # Plot model ax.plot(self.config["wavelength"], flux_model, c="b", label="Model", lw=0.7) + if continuum_model is not None: + ax.plot(self.config["wavelength"], continuum_model, c="cornflowerblue", + label="Continuum", lw=0.7) + if continuum_model_err is not None: + ax.fill_between( + self.config["wavelength"].value, + continuum_model - continuum_model_err, + continuum_model + continuum_model_err, + color="cornflowerblue", + alpha=0.4, + label="Continuum error" + ) + # Plot residuals residuals = flux_model - self.config["flux"] ax.plot( @@ -953,15 +1041,14 @@ def plot_solution(self, solution: DataBlock, figname=None, plot_lines=True): ax.set_xlim(self.config["wavelength"].value[[0, -1]]) # Plot chi2 good_pixels = weights > 0 - chi2 = (flux_model - self.config["flux"]) ** 2 / self.config["var"] + chi2 = (flux_model - self.config["flux"]) ** 2 * self.config["ivar"] mean_chi2 = np.nanmean(chi2[good_pixels]) median_chi2 = np.nanmedian(chi2[good_pixels]) nmad_chi2 = 1.4826 * np.nanmedian( np.abs(chi2[good_pixels] - median_chi2)) loglike = self.log_like(self.config["flux"][good_pixels], flux_model[good_pixels], - self.config["var"][good_pixels], - weights=weights[good_pixels]) + self.config["ivar"][good_pixels] * weights[good_pixels]) ax = axs[1, 0] ax.plot(self.config["wavelength"], chi2, c="k", lw=0.7) ax.grid(visible=True) @@ -969,6 +1056,12 @@ def plot_solution(self, solution: DataBlock, figname=None, plot_lines=True): ax.set_yscale("symlog", linthresh=1.0) ax.set_xlabel("Wavelength (AA)") + twax = ax.twinx() + twax.fill_between( + np.array(self.config["wavelength"]), 0.0, weights, + color="lime", alpha=0.2, label="Weights") + twax.set_ylabel("Weight") + ax = axs[1, 1] ax.hist( chi2, @@ -1006,7 +1099,7 @@ def prepare_observed_photometry(self, options: SectionOptions): ---------- options : :class:`DataBlock` """ - _log("\n-> Configuring photometric data") + _log("Configuring photometric data") photometry_file = os.path.expandvars(options["inputPhotometry"]) # Read the data @@ -1058,7 +1151,7 @@ def prepare_observed_photometry(self, options: SectionOptions): redshift = options.get_double("redshift", default=0.0) self.config["redshift"] = redshift _log("Source redshift: ", redshift) - _log("-> Configuration done.") + _log("Configuration done.") def prepare_galaxy(self, options): """Build and configure a :class:`pst.galaxy.GalaxySED` model for photometry. @@ -1124,6 +1217,7 @@ def prepare_galaxy(self, options): z_obs = self.config.get("redshift", 0.0) + #TODO: make sure this is well documented if options.get_bool("logwave", False): target_wl = np.geomspace(min_wl.to_value("AA") / (1 + z_obs), max_wl.to_value("AA"), 3000) << u.AA @@ -1244,12 +1338,12 @@ def plot_solution(self, solution: DataBlock, figname=None): ax.set_xlim(min_wl.to_value(u.AA) * 0.8, max_wl.to_value(u.AA) * 1.2) # Flux density per wavelength unit ax = axs[1, 0] - flam = (full_spec * u.Unit("uJy")).to("1e-16 erg / (s cm**2 AA)", u.spectral_density(self.config["galaxy"].target_wavelength)) + flam = (full_spec * u.Unit("uJy")).to(self._default_flux_units, u.spectral_density(self.config["galaxy"].target_wavelength)) ax.plot( self.config["galaxy"].target_wavelength.to_value("AA"), flam, color="k", alpha=0.4) - ax.set_ylabel("Flux density (1e-16 erg / (s cm**2 AA))") + ax.set_ylabel(f"Flux density ({self._default_flux_units})") # chi2 as function of wavelength chi2 = (flux_model - self.config["photometry_flux"]) ** 2 / self.config["photometry_flux_var"] ax = axs[2, 0] @@ -1355,7 +1449,7 @@ def prepare_grid_model(self, options): options : :class:`DataBlock` Input options to initialise the model. """ - logger.info("-> Configuring model grid") + logger.info("Configuring model grid") if not options.has_value("modelGridFile"): raise ValueError("No input model grid file provided.") grid_file = os.path.expandvars(options["modelGridFile"]) @@ -1385,7 +1479,7 @@ def prepare_grid_model(self, options): self.config["knn"] = options["knn"] else: self.config["knn"] = int(4 * model_grid.n_targets) - logger.info("-> Configuration done.") + logger.info("Configuration done.") class EmulatorMixin: @@ -1408,7 +1502,7 @@ def prepare_emulator(self, options): except ImportError: raise ImportError("joblib is required to load ML emulators." "Please install joblib and try again.") - logger.info("-> Configuring ML emulator") + logger.info("Configuring ML emulator") if not options.has_value("emulatorFile"): raise ValueError("No input emulator file provided.") emulator_file = os.path.expandvars(options["emulatorFile"]) @@ -1418,4 +1512,4 @@ def prepare_emulator(self, options): logger.info("Reading ML emulator...") ml_emulator = joblib.load(emulator_file) self.config["ml_emulator"] = ml_emulator - logger.info("-> Configuration done.") + logger.info("Configuration done.") diff --git a/src/besta/pipeline_modules/full_spectral_fit.py b/src/besta/pipeline_modules/full_spectral_fit.py index fbec41e6..1426c411 100644 --- a/src/besta/pipeline_modules/full_spectral_fit.py +++ b/src/besta/pipeline_modules/full_spectral_fit.py @@ -27,7 +27,7 @@ def __init__(self, options, **kwargs): **kwargs : dict Extra keyword arguments forwarded to ``SpectraFitModule``. """ - super().__init__(options, **kwargs) + super().__init__(options, likelihood_kind="spectra", **kwargs) options = self.parse_options(options) # Check for the necessary options and prepare the models @@ -39,6 +39,8 @@ def __init__(self, options, **kwargs): self.prepare_sfh_model(options) self.prepare_extinction_law(options) self.prepare_legendre_polynomials(options) + self.prepare_losvd_kernel(options) + self._losvd_kernel = self.config["losvd_kernel"] @spectrum.legendre_decorator def make_observable(self, block, parse=False): @@ -50,34 +52,17 @@ def make_observable(self, block, parse=False): luminosity_model = sfh_model.model.compute_SED( self.config["ssp_model"], t_obs=sfh_model.today, allow_negative=False ) - flux_model = 1e10 * luminosity_model.to_value("1e-16 erg / (s Angstrom)" + flux_model = 1e10 * luminosity_model.to_value(self._default_luminosity_units ) / self.config["dl_sq"] # Kinematics - velscale = self.config["velscale"] - # Kinematics - sigma_pixel = block["kinematics", "los_sigma"] / velscale - veloffset_pixel = block["kinematics", "los_vel"] / velscale - # Build the kernel. TOO SLOW? Initialise only once? - kernel_model = kinematics.GaussHermite( - 4, - mean=veloffset_pixel, - stddev=sigma_pixel, - h3=block["kinematics", "los_h3"], - h4=block["kinematics", "los_h4"], - ) - kernel_n_pixel = 10 * np.clip(int(np.round(np.abs(veloffset_pixel) + sigma_pixel)), 1, - None) + 1 - kernel = kinematics.get_losvd_kernel( - kernel_model, - x_size=kernel_n_pixel - ) + self._losvd_kernel.parse_parameters(block) # Perform the convolution - flux_model = kinematics.convolve_spectra_with_kernel(flux_model, kernel) + flux_model = self._losvd_kernel.convolve(flux_model) # Track those pixels at the edges mask = flux_model > 0 - mask[: int(10 * sigma_pixel)] = False - mask[-int(10 * sigma_pixel) :] = False + mask[:self._losvd_kernel.size // 2] = False + mask[-self._losvd_kernel.size // 2:] = False # Sample to observed resolution extra_pixels = self.config["extra_pixels"] pixels = slice(extra_pixels, -extra_pixels) @@ -87,7 +72,7 @@ def make_observable(self, block, parse=False): # Apply dust extinction dust_model = self.config["extinction_law"] flux_model = dust_model.apply_extinction( - self.config["wavelength"], flux_model, a_v=block["dust.extinction", "a_v"] + self.config["wavelength"], flux_model, a_v=block["dust.attenuation", "a_v"] ).value weights = self.config["weights"] * mask @@ -107,18 +92,19 @@ def execute(self, block): if not valid: # To track invalid samples users can set debug=T # logger.warning("Invalid sample") + logger.debug("Invalid sample: %s", block) block[section_names.likelihoods, self.like_name] = -1e20 * penalty block["extra", "stellar_mass"] = np.nan return 0 # Obtain parameters from setup - cov = self.config["var"] flux_model, weights = self.make_observable(block) # Calculate likelihood-value of the fit good_pixels = weights > 0 like = self.log_like(self.config["flux"][good_pixels], flux_model[good_pixels], - cov[good_pixels], - weights=weights[good_pixels]) + self.config["ivar"][good_pixels] * weights[good_pixels]) + # To make it compatible with photometric likelihoods + like /= np.sum(weights[good_pixels]) # Final posterior for sampling block[section_names.likelihoods, self.like_name] = like return 0 diff --git a/src/besta/pipeline_modules/galaxy_photometry.py b/src/besta/pipeline_modules/galaxy_photometry.py index 46d0f159..4de5343a 100644 --- a/src/besta/pipeline_modules/galaxy_photometry.py +++ b/src/besta/pipeline_modules/galaxy_photometry.py @@ -24,7 +24,7 @@ class GalaxyPhotometryModule(PhotometryFitModule): def __init__(self, options, **kwargs): """Set up the module from a CosmoSIS configuration block.""" - super().__init__(options, **kwargs) + super().__init__(options, likelihood_kind="photometry", **kwargs) options = self.parse_options(options) self.prepare_observed_photometry(options) self.prepare_galaxy(options) diff --git a/src/besta/pipeline_modules/galaxy_spectra.py b/src/besta/pipeline_modules/galaxy_spectra.py index fb3c564b..e17e1e9a 100644 --- a/src/besta/pipeline_modules/galaxy_spectra.py +++ b/src/besta/pipeline_modules/galaxy_spectra.py @@ -19,11 +19,13 @@ class GalaxySpectraModule(SpectraFitModule): def __init__(self, options, **kwargs): """Set up the module from a CosmoSIS configuration block.""" - super().__init__(options, **kwargs) + super().__init__(options, likelihood_kind="spectra", **kwargs) options = self.parse_options(options) self.prepare_observed_spectra(options) self.prepare_galaxy(options) self.prepare_legendre_polynomials(options) + self.prepare_losvd_kernel(options) + self._losvd_kernel = self.config["losvd_kernel"] # Set parameters fixed in this module self.config["galaxy"].redshift.fixed = True @@ -45,32 +47,15 @@ def make_observable(self, block, parse=False): galaxy.update_parameters(parameters, strict=False) # Synthesis flux_model = 1e10 * galaxy.emission_spectrum( - to_obs_frame=False).to_value("1e-16 erg / (s Angstrom)") / self.config["dl_sq"] - - # Kinematics #TODO: this should be done by PST stars.kinematics - velscale = self.config["velscale"] - sigma_pixel = block["kinematics", "los_sigma"] / velscale - veloffset_pixel = block["kinematics", "los_vel"] / velscale - - kernel_model = kinematics.GaussHermite( - 4, - mean=veloffset_pixel, - stddev=sigma_pixel, - h3=block["kinematics", "los_h3"], - h4=block["kinematics", "los_h4"], - ) - kernel_n_pixel = 10 * np.clip(int(np.round(np.abs(veloffset_pixel) + sigma_pixel)), 1, - None) + 1 - kernel = kinematics.get_losvd_kernel( - kernel_model, - x_size=kernel_n_pixel - ) + to_obs_frame=False).to_value(self._default_luminosity_units) / self.config["dl_sq"] + + self._losvd_kernel.parse_parameters(block) # Perform the convolution - flux_model = kinematics.convolve_spectra_with_kernel(flux_model, kernel) + flux_model = self._losvd_kernel.convolve(flux_model) # Track those pixels at the edges mask = flux_model > 0 - mask[: int(10 * sigma_pixel)] = False - mask[-int(10 * sigma_pixel) :] = False + mask[:self._losvd_kernel.size // 2] = False + mask[-self._losvd_kernel.size // 2:] = False # Sample to observed resolution extra_pixels = self.config["extra_pixels"] pixels = slice(extra_pixels, -extra_pixels) @@ -98,14 +83,14 @@ def execute(self, block): block["extra", "stellar_mass"] = np.nan return 0 # Obtain parameters from setup - cov = self.config["var"] flux_model, weights = self.make_observable(block) # Calculate likelihood-value of the fit good_pixels = weights > 0 like = self.log_like(self.config["flux"][good_pixels], flux_model[good_pixels], - cov[good_pixels], - weights=weights[good_pixels]) + self.config["ivar"][good_pixels] * weights[good_pixels]) + # To make it compatible with photometric likelihoods + like /= np.sum(weights[good_pixels]) # Final posterior for sampling block[section_names.likelihoods, self.like_name] = like return 0 diff --git a/src/besta/pipeline_modules/likelihoods.py b/src/besta/pipeline_modules/likelihoods.py new file mode 100644 index 00000000..f37aa137 --- /dev/null +++ b/src/besta/pipeline_modules/likelihoods.py @@ -0,0 +1,223 @@ +"""Likelihood helpers for BESTA pipeline modules. + +The module exposes fast, backend-selectable likelihood functions for spectra and +photometry. Spectral modules use an inverse-variance Gaussian likelihood with no +limit handling. Photometry modules use the more general likelihood that can +handle optional upper/lower limits, with a fast no-limits path when requested. +""" + +from __future__ import annotations + +import math +from typing import Callable + +import numpy as np +from scipy.stats import norm + +try: + from numba import njit + + NUMBA_AVAILABLE = True +except ImportError: # pragma: no cover - optional dependency guard + NUMBA_AVAILABLE = False + njit = None # type: ignore[assignment] + + +_LOG_2PI = math.log(2.0 * math.pi) + + +def _validate_spectra_inputs(data, model, ivar): + if data.shape != model.shape or data.shape != ivar.shape: + raise ValueError("data, model, ivar must have the same shape (ivar is per-datum inverse variance).") + if np.any(ivar < 0): + raise ValueError("All ivar entries must be >= 0 (inverse variance).") + + +def _validate_photometry_inputs(data, model, var, weights, is_upper, is_lower): + if data.shape != model.shape or data.shape != var.shape: + raise ValueError("data, model, var must have the same shape (var is per-datum variance).") + if np.any(var <= 0): + raise ValueError("All var entries must be > 0 (variance).") + + if is_upper is None: + is_upper = np.zeros_like(data, dtype=bool) + else: + is_upper = np.asarray(is_upper, dtype=bool) + + if is_lower is None: + is_lower = np.zeros_like(data, dtype=bool) + else: + is_lower = np.asarray(is_lower, dtype=bool) + + if is_upper.shape != data.shape or is_lower.shape != data.shape: + raise ValueError("is_upper and is_lower must have the same shape as data/model.") + if np.any(is_upper & is_lower): + raise ValueError("A data point cannot be both an upper and a lower limit.") + + if weights is None: + weights = np.ones_like(data, dtype=float) + normalize = False + else: + weights = np.asarray(weights, dtype=float) + if weights.shape != data.shape: + raise ValueError("weights must have the same shape as data/model.") + if np.any(weights < 0): + raise ValueError("weights must be non-negative.") + normalize = True + + return weights, is_upper, is_lower, normalize + + +def spectra_loglike_numpy(data, model, ivar): + """Inverse-variance Gaussian log-likelihood for spectra.""" + _validate_spectra_inputs(data, model, ivar) + chi2 = np.sum(ivar * (data - model) ** 2) + return -0.5 * chi2 + + +if NUMBA_AVAILABLE: + + @njit(cache=True, fastmath=False) + def spectra_loglike_numba(data, model, ivar): + chi2 = 0.0 + for i in range(data.size): + residual = data[i] - model[i] + chi2 += ivar[i] * residual * residual + return -0.5 * chi2 + + + @njit(cache=True, fastmath=False) + def photometry_loglike_numba_no_limits(data, model, var, weights, normalize, include_norm): + total = 0.0 + wsum = 0.0 + for i in range(data.size): + sigma = math.sqrt(var[i]) + z = (data[i] - model[i]) / sigma + logp = -0.5 * z * z + if include_norm: + logp += -0.5 * (_LOG_2PI + math.log(var[i])) + w = weights[i] + total += logp * w + if normalize: + wsum += w + if normalize: + return total / wsum + return total + +else: # pragma: no cover - exercised only when numba is missing + + def spectra_loglike_numba(*args, **kwargs): + raise ImportError("numba is not available") + + def photometry_loglike_numba_no_limits(*args, **kwargs): + raise ImportError("numba is not available") + + +def photometry_loglike_numpy( + data, + model, + var, + weights=None, + is_upper=None, + is_lower=None, + include_norm=True, +): + """Gaussian photometry likelihood with optional upper/lower limits.""" + weights, is_upper, is_lower, normalize = _validate_photometry_inputs( + data, model, var, weights, is_upper, is_lower + ) + + if not np.any(is_upper | is_lower): + sigma = np.sqrt(var) + logp = -0.5 * ((data - model) / sigma) ** 2 + if include_norm: + logp = logp - 0.5 * (np.log(2.0 * np.pi) + np.log(var)) + if normalize: + wsum = np.sum(weights) + if wsum <= 0: + raise ValueError("Sum of weights must be > 0 for normalization.") + return np.sum(logp * weights) / wsum + return np.sum(logp * weights) + + sigma = np.sqrt(var) + logp = np.empty_like(data, dtype=float) + det = ~(is_upper | is_lower) + + if np.any(det): + if include_norm: + logp[det] = norm.logpdf(data[det], loc=model[det], scale=sigma[det]) + else: + z = (data[det] - model[det]) / sigma[det] + logp[det] = -0.5 * z**2 + + if np.any(is_upper): + z_u = (data[is_upper] - model[is_upper]) / sigma[is_upper] + logp[is_upper] = norm.logcdf(z_u) + + if np.any(is_lower): + z_l = (data[is_lower] - model[is_lower]) / sigma[is_lower] + logp[is_lower] = norm.logsf(z_l) + + if normalize: + wsum = np.sum(weights) + if wsum <= 0: + raise ValueError("Sum of weights must be > 0 for normalization.") + return np.sum(logp * weights) / wsum + + return np.sum(logp * weights) + + +def make_spectra_loglike(method: str = "auto") -> Callable: + """Return a spectra likelihood implementation for the requested backend.""" + method = (method or "auto").strip().lower() + if method in {"auto", "numba"} and NUMBA_AVAILABLE: + return spectra_loglike_numba + if method in {"auto", "numpy"}: + return spectra_loglike_numpy + raise ValueError("Valid spectra likelihood methods are: auto, numpy, numba") + + +def make_photometry_loglike(method: str = "auto") -> Callable: + """Return a photometry likelihood implementation for the requested backend.""" + method = (method or "auto").strip().lower() + + if method in {"auto", "numba"} and NUMBA_AVAILABLE: + + def _numba_or_numpy( + data, + model, + var, + weights=None, + is_upper=None, + is_lower=None, + include_norm=True, + ): + if is_upper is None and is_lower is None: + weights_arr = np.ones_like(data, dtype=float) if weights is None else np.asarray(weights, dtype=float) + normalize = weights is not None + if normalize and np.any(weights_arr < 0): + raise ValueError("weights must be non-negative.") + return photometry_loglike_numba_no_limits( + np.asarray(data, dtype=float), + np.asarray(model, dtype=float), + np.asarray(var, dtype=float), + weights_arr, + normalize, + include_norm, + ) + return photometry_loglike_numpy( + data, + model, + var, + weights=weights, + is_upper=is_upper, + is_lower=is_lower, + include_norm=include_norm, + ) + + return _numba_or_numpy + + if method in {"auto", "numpy"}: + return photometry_loglike_numpy + + raise ValueError("Valid photometry likelihood methods are: auto, numpy, numba") diff --git a/src/besta/pipeline_modules/sfh_photometry_emu.py b/src/besta/pipeline_modules/sfh_photometry_emu.py index 154f813b..ac1c6206 100644 --- a/src/besta/pipeline_modules/sfh_photometry_emu.py +++ b/src/besta/pipeline_modules/sfh_photometry_emu.py @@ -20,7 +20,7 @@ class SFHPhotometryEmulatorModule(PhotometryFitModule, EmulatorMixin): def __init__(self, options): """Set up the module from a CosmoSIS configuration block.""" - super().__init__(options) + super().__init__(options, likelihood_kind="photometry") options = self.parse_options(options) # Pipeline values file self.prepare_observed_photometry(options) diff --git a/src/besta/pipeline_modules/sfh_photometry_grid.py b/src/besta/pipeline_modules/sfh_photometry_grid.py index 8cf0d7bd..ad48c4b6 100644 --- a/src/besta/pipeline_modules/sfh_photometry_grid.py +++ b/src/besta/pipeline_modules/sfh_photometry_grid.py @@ -20,7 +20,7 @@ class SFHPhotometryGridModule(PhotometryFitModule, GridFitMixin): def __init__(self, options): """Set up the module from a CosmoSIS configuration block.""" - super().__init__(options) + super().__init__(options, likelihood_kind="photometry") options = self.parse_options(options) # Pipeline values file self.prepare_observed_photometry(options) diff --git a/src/besta/pipeline_modules/spectra_redshift_fit.py b/src/besta/pipeline_modules/spectra_redshift_fit.py index 91ea9e8f..df3aece4 100644 --- a/src/besta/pipeline_modules/spectra_redshift_fit.py +++ b/src/besta/pipeline_modules/spectra_redshift_fit.py @@ -2,15 +2,103 @@ import os import numpy as np +from astropy.table import Table from besta.pipeline_modules.base_module import SpectraFitModule from cosmosis.datablock import names as section_names from cosmosis.datablock import SectionOptions from besta import spectrum from besta.logging import get_logger - +from besta.io import parse_table_format +from numba import njit, prange logger = get_logger(__name__) + +@njit(parallel=True, fastmath=False) +def compute_redshift_chi2_from_slices( + flux_model, + target_flux, + candidate_weights, + slice_starts, + slice_stops, + good_idx, + ): + """Compute the chi2 values for all redshift steps defined by the model slices. + + Description + ----------- + + The returned scales are the best-fit normalization factors for each redshift + step, which can be used to compute the best-fit model fluxes if needed. + + Parameters + ---------- + flux_model : array-like + The model flux values. + target_flux : array-like + The target flux values. + candidate_weights : array-like + The weights for each pixel. + slice_starts : array-like + The starting indices for each redshift slice. + slice_stops : array-like + The stopping indices for each redshift slice. + good_idx : array-like + The indices of the good pixels. + + Returns + ------- + z_chi2 : array-like + The chi2 values for each redshift step. + z_scales : array-like + The best-fit normalization factors for each redshift step. + """ + n_z = len(slice_starts) + + z_chi2 = np.full(n_z, np.inf) + z_scales = np.full(n_z, np.nan) + + # Loop over all redshift steps in parallel + for i in prange(n_z): + start = slice_starts[i] + + numerator = 0.0 + denominator = 0.0 + + for jj in range(len(good_idx)): + k = start + good_idx[jj] + + f = flux_model[k] + t = target_flux[jj] + w = candidate_weights[jj] + + numerator += w * f * t + denominator += w * f * f + + if denominator <= 0.0: + continue + + scale = numerator / denominator + z_scales[i] = scale + + chi2 = 0.0 + + for jj in range(len(good_idx)): + k = start + good_idx[jj] + + f = flux_model[k] + t = target_flux[jj] + w = candidate_weights[jj] + + residual = scale * f - t + chi2 += w * residual * residual + + # Normalise chi2 by the number of good pixels + if len(good_idx) > 0: + z_chi2[i] = chi2 / len(good_idx) + + return z_chi2, z_scales + class SpectraRedshiftFitModule(SpectraFitModule): """Fit stellar populations and kinematics directly from galaxy spectra.""" @@ -27,7 +115,7 @@ def __init__(self, options, **kwargs): **kwargs : dict Extra keyword arguments forwarded to ``SpectraFitModule``. """ - super().__init__(options, **kwargs) + super().__init__(options, likelihood_kind="spectra", **kwargs) options = self.parse_options(options) if options.has_value("redshift") and options["redshift"] != 0: @@ -114,6 +202,7 @@ def __init__(self, options, **kwargs): ) self.config["model_slices"] = model_slices self.config["model_start"] = np.array([s.start for s in model_slices]) + self.config["model_stop"] = np.array([slc.stop for slc in model_slices], dtype=np.int64) # Likelihood values of all redshift stepsg. slice_redshifts = observed_wavelength[0] / model_wavelength[self.config["model_start"]] - 1.0 @@ -129,7 +218,10 @@ def __init__(self, options, **kwargs): # Check if the output file already exists to avoid overwriting previous results. if os.path.exists(self.z_loglike_path): logger.info("Loading existing redshift log-likelihood profile from file.") - z, loglike = np.loadtxt(self.z_loglike_path, unpack=True) + + format = parse_table_format(self.z_loglike_path) + t = Table.read(self.z_loglike_path, format=format) + z, loglike = t["redshift"].value, t["log_likelihood"].value if np.array_equal(z, slice_redshifts): self.z_loglike = loglike else: @@ -139,36 +231,27 @@ def __init__(self, options, **kwargs): f"Model wavelength range in rest frame: {self.config['ssp_model'].wavelength[model_start].to_value('AA'):.1f} - {self.config['ssp_model'].wavelength[model_stop-1].to_value('AA'):.1f} AA") logger.info(f"Number of redshift steps: {len(model_slices)} (z={z_max:.1f} to {z_min:.1f})") - w = np.ones_like(self.config["weights"]) - if options.has_value("use_features"): - print(options["use_features"]) - continuum, continuum_err = spectrum.estimate_continuum( - self.config["wavelength"].to_value("AA"), - self.config["flux"], - err=self.config["var"]**0.5, - weights=self.config["weights"], - knot_spacing=options.get_double("continuum_knot_spacing", default=200.0), - sigma_clip=options.get_double("continuum_sigma_clip", default=3.0), - ) - self.config["continuum"] = continuum - self.config["continuum_err"] = continuum_err - # Favour features over/under continuum - w = (np.abs(self.config["flux"] - continuum) / continuum_err)**2 - w = np.where(np.isfinite(w), w, 0.0) - w_sum = np.nansum(w) - if w_sum <= 0: - raise ValueError("Feature-based weights sum to zero; cannot perform redshift fit.") - w /= w_sum - logger.info("Using feature-based weights for redshift fitting.") + if options.get_bool("use_features", default=False): + self.get_feature_weights(options) else: logger.info("Using original weights for redshift fitting.") - self.config["sweep_weights"] = w * self.config["weights"] / self.config["norm_obs_var"] + w = self.config.get("feature_weights", + np.ones_like(self.config["flux"], dtype=np.float32)) + self.config["sweep_weights"] = np.where( + (w > 0) & np.isfinite(self.config["norm_obs_var"]), + w * self.config["weights"] / self.config["norm_obs_var"], + 0.0 + ) + # This are used for the likelihood in combination with the variance + self.config["weights_orig"] = self.config["weights"].copy() self.config["weights"] *= w + self.config["good"] = self.config["sweep_weights"] > 0 + self.config["good_idx"] = np.flatnonzero(self.config["good"]).astype(np.int64) @spectrum.legendre_decorator - def make_observable(self, block, parse=False): + def make_observable(self, block, parse=False, ): """Create the spectra model from the input parameters""" # Stellar population synthesis sfh_model = self.config["sfh_model"] @@ -178,45 +261,34 @@ def make_observable(self, block, parse=False): # Here we compute the luminosity by we call it flux and rescale later flux_model = sfh_model.model.compute_SED( self.config["ssp_model"], t_obs=sfh_model.today, allow_negative=False - ).to_value("1e-16 erg / (s Angstrom)") + ).to_value(self._default_luminosity_units) # Apply dust extinction dust_model = self.config["extinction_law"] if dust_model is not None: flux_model = dust_model.apply_extinction( self.config["ssp_model"].wavelength, flux_model, - a_v=block["dust.extinction", "a_v"] + a_v=block["dust_attenuation", "a_v"] ).value w = self.config["sweep_weights"] - mask = w > 0 + good = self.config["good"] + good_idx = self.config["good_idx"] + candidate_weights = w[good] + slc_starts = self.config["model_start"] + slc_stops = self.config["model_stop"] + norm_obs_flux = self.config["norm_obs_flux"] - z_chi2 = np.full(len(self.config["model_slices"]), np.inf) - z_scales = np.full(len(self.config["model_slices"]), np.nan) - - # Sweep over all target slices using the same weighted least-squares - # scale that is applied to the selected model below. - for i, slc in enumerate(self.config["model_slices"]): - candidate_flux = flux_model[slc] - good = ( - mask - & np.isfinite(candidate_flux) - & np.isfinite(norm_obs_flux) - & np.isfinite(w) - ) - if not np.any(good): - continue - candidate_flux = candidate_flux[good] - candidate_weights = w[good] - target_flux = norm_obs_flux[good] - denominator = np.nansum(candidate_weights * candidate_flux**2) - if denominator <= 0: - continue - scale = np.nansum(candidate_weights * candidate_flux * target_flux) / denominator - z_scales[i] = scale - z_chi2[i] = np.nansum( - candidate_weights * (candidate_flux * scale - target_flux) ** 2 - ) + target_flux = norm_obs_flux[good] + + z_chi2, z_scales = compute_redshift_chi2_from_slices( + flux_model, + target_flux, + candidate_weights, + slc_starts, + slc_stops, + good_idx, + ) # Keep track of the likelihood values for all redshift steps. self.z_loglike = np.maximum(self.z_loglike, -0.5 * z_chi2) @@ -233,14 +305,23 @@ def make_observable(self, block, parse=False): z_best, z_chi2[best_fit_slice_index], ) - # Re-scale model flux to match the observed flux level, and convert to physical units for - # luminosity distance calculation. The stellar mass is then inferred from the normalization. - # dl_sq = cosmology.luminosity_distance(z_best).to_value("cm")**2 - # flux_model /= 4 * np.pi * dl_sq - n_pix = self.config["flux"].size - normalization = z_scales[best_fit_slice_index] * self.config["norm_obs_flux_scale"] + + # TODO: I am not sure if this will bias the likelihood + # Use the original weights to estimate the normalization + w = self.config["weights_orig"] + candidate_flux = flux_model[best_fit_index : best_fit_index + w.size] + denominator = np.nansum(w[good] * candidate_flux[good]**2) + + if denominator <= 0: + logger.warning( + f"Denominator for redshift step {best_fit_slice_index} is non-positive; skipping this step." + ) + return np.full_like(candidate_flux, np.nan), w + + scale = np.nansum(w[good] * candidate_flux[good] * target_flux) / denominator + normalization = scale * self.config["norm_obs_flux_scale"] #block["extra", "stellar_mass"] = np.log10(normalization) + 10 - return flux_model[best_fit_index : best_fit_index + n_pix] * normalization, self.config["weights"] + return candidate_flux * normalization, self.config["weights"] def execute(self, block): """Function executed by sampler @@ -255,28 +336,31 @@ def execute(self, block): block["extra", "stellar_mass"] = np.nan return 0 # Obtain parameters from setup - var = self.config["var"] flux_model, weights = self.make_observable(block) # Calculate likelihood-value of the fit good_pixels = weights > 0 like = self.log_like(self.config["flux"][good_pixels], flux_model[good_pixels], - var[good_pixels], - weights=weights[good_pixels]) + self.config["ivar"][good_pixels] * weights[good_pixels]) + # To make it compatible with photometric likelihoods + like /= np.sum(weights[good_pixels]) # Final posterior for sampling block[section_names.likelihoods, self.like_name] = like return 0 def cleanup(self): - """Persist the redshift likelihood profile if requested.""" + """Save the redshift likelihood profile if requested.""" if self.save_z_loglike: logger.info(f"Saving redshift log-likelihood profile to {self.z_loglike_path}") - np.savetxt( - self.z_loglike_path, - np.column_stack( - (self.config["slice_redshifts"], self.z_loglike)), - header="redshift log_likelihood") + t = Table( + [self.config["slice_redshifts"], self.z_loglike], + names=["redshift", "log_likelihood"], + meta={"description": "Redshift log-likelihood profile from SpectraRedshiftFitModule"} + ) + # Guess the format from the file extension + format = parse_table_format(self.z_loglike_path) + t.write(self.z_loglike_path, format=format, overwrite=True) def setup(options): diff --git a/src/besta/postprocess.py b/src/besta/postprocess.py index 37e45792..125b94da 100644 --- a/src/besta/postprocess.py +++ b/src/besta/postprocess.py @@ -23,11 +23,15 @@ from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple import numpy as np +from matplotlib import pyplot as plt +from scipy import stats +from scipy.optimize import minimize + from astropy.io import fits from astropy.table import Table, Column from astropy import units as u -from matplotlib import pyplot as plt -from scipy import stats + +from besta import io from besta.logging import get_logger logger = get_logger(__name__) @@ -346,24 +350,24 @@ def histogram_pdf_1d( # Convert probability per bin -> density with np.errstate(divide="ignore", invalid="ignore"): pdf = hist / dx - centers = 0.5 * (edges[:-1] + edges[1:]) # Ensure integrates to 1 (numerical) integral = np.nansum(pdf * dx) if integral > 0: pdf /= integral - return centers, pdf + return edges, pdf def kde_pdf_1d( x: np.ndarray, weights: np.ndarray, - grid: np.ndarray, + bins: np.ndarray, ) -> np.ndarray: """ KDE PDF on a provided grid; returns NaNs on failure. """ x = _as_float_array(x).ravel() w = _as_float_array(weights).ravel() - g = _as_float_array(grid).ravel() + edges = _as_float_array(bins).ravel() + g = 0.5 * (edges[:-1] + edges[1:]) mask = np.isfinite(x) & np.isfinite(w) x = x[mask] w = w[mask] @@ -446,61 +450,253 @@ def kde_or_hist_pdf_2d( Z /= integral return xc, yc, Z +def pit_from_pdf(x_edges, pdf, x_true): + """ + Compute PIT value for a true value given a PDF defined by edges and values. + + Parameters + ---------- + x_edges : array shape (N+1,) bin edges + pdf : array shape (N,) PDF values for each bin, normalized to integrate to 1 + x_true : scalar true value + + Returns + ------- + pit : scalar in [0,1] representing the cumulative probability up to ``x_true`` + """ + x_edges = _as_float_array(x_edges).ravel() + pdf = _as_float_array(pdf).ravel() + if x_edges.size != pdf.size + 1: + raise ValueError("x_edges must have one more element than pdf.") + if not np.isfinite(x_true): + raise ValueError("x_true must be finite.") + dx = np.diff(x_edges) + cdf = np.cumsum(pdf * dx) + if not np.isclose(cdf[-1], 1.0): + raise ValueError("pdf must be normalized to integrate to 1.") + + return np.interp(x_true, x_edges, np.r_[0.0, cdf]) + +def pdf_stats(edges: np.ndarray, pdf: np.ndarray, + quantiles=None, + find_multimodal: bool = False) -> dict: + """ + Compute summary stats from a discrete pdf over bin centers. + + Parameters + ---------- + edges : ndarray, shape (K,) + Bin edges. + pdf : ndarray, shape (K,) + PDF defined by the edges. + quantiles : tuple of float, optional + Quantiles to report (default 16, 50, 84 percent). + find_multimodal : bool, optional + If True, return rough modes by local-maximum search. + + Returns + ------- + stats : dict + Keys: mean, std, map, q, lo68, hi68, modes (optional). + """ + # Ensure normalization + norm = np.sum(pdf * np.diff(edges)) + pdf /= norm if norm > 0 else 1.0 + centers = 0.5 * (edges[:-1] + edges[1:]) + # Trapecium integrals + mean = np.sum(pdf * centers * np.diff(edges)) + var = np.sum(pdf * (centers - mean)**2 * np.diff(edges)) + std = var ** 0.5 + k_map = np.argmax(pdf * np.diff(edges)) + v_map = centers[k_map] + + cdf = np.cumsum(pdf * np.diff(edges)) + cdf = np.insert(cdf, 0, 0) + + if quantiles is not None: + qs = np.array(quantiles, float) + qvals = np.interp(qs, cdf, edges, left=edges[0], right=edges[-1]) + else: + qvals = None + + median = np.interp(0.5, cdf, edges, left=edges[0], right=edges[-1]) + lo68 = np.interp(0.16, cdf, edges, left=edges[0], right=edges[-1]) + hi68 = np.interp(0.84, cdf, edges, left=edges[0], right=edges[-1]) + + out = {"mean": mean, "std": std, "map": v_map, "q": qvals, + "median": median, "lo68": lo68, "hi68": hi68} + + if find_multimodal: + modes = [] + for i in range(1, len(pdf) - 1): + if pdf[i] > pdf[i - 1] and pdf[i] > pdf[i + 1]: + modes.append(centers[i]) + if not modes: + modes = [v_map] + out["modes"] = np.asarray(modes) + return out + # ----------------------------------------------------------------------------- -# I/O helpers +# Autocorrelation # ----------------------------------------------------------------------------- -def read_results_file(path: str, *, delimiter: str = "\t") -> Table: + +def next_pow_two(n): + i = 1 + while i < n: + i <<= 1 + return i + + +def autocorr_func_1d(x, norm=True): """ - Read a CosmoSIS-style text results file: - - First line is a commented header starting with '#' - - Columns are delimiter-separated - - Remaining lines numeric + Estimate the 1D autocorrelation function using FFT. + + Parameters + ---------- + x : array_like, shape (n_sample,) + One chain, e.g. one walker for one parameter. + norm : bool + If True, normalize so acf[0] = 1. - Returns an Astropy Table with lowercase column names. + Returns + ------- + acf : ndarray, shape (n_sample,) + Autocorrelation function. """ - with open(path, "r", encoding="utf-8") as f: - header = f.readline() - if not header.startswith("#"): - raise ValueError("Expected first line header starting with '#'.") - - columns = header.strip("# \n").split(delimiter) - matrix = np.atleast_2d(np.loadtxt(path)) - tab = Table() - if matrix.size <= 1: - return tab - if matrix.shape[1] != len(columns): - raise ValueError( - f"Data has {matrix.shape[1]} columns but header lists {len(columns)}." - ) - for i, c in enumerate(columns): - tab[c.strip().lower()] = matrix[:, i] - return tab + x = np.asarray(x, dtype=float) -def _select_parameter_keys( - table: Table, - *, - parameter_prefix: str = "--", - parameter_keys: Optional[Sequence[str]] = None, -) -> List[str]: - if parameter_keys is not None: - keys = list(parameter_keys) - else: - keys = [k for k in table.colnames if parameter_prefix in k] - if len(keys) == 0: - raise ValueError("No parameter keys found/selected.") - return keys + if x.ndim != 1: + raise ValueError("x must be one-dimensional") + + if x.size < 2: + raise ValueError("x must contain at least two samples") + + if not np.all(np.isfinite(x)): + raise ValueError("x contains NaN or infinite values") + + x = x - np.mean(x) + + if np.allclose(x, 0.0): + raise ValueError("cannot compute autocorrelation of a constant chain") -def _split_param_key(key: str, prefix: str = "--") -> Tuple[str, str]: + n = next_pow_two(len(x)) + + f = np.fft.fft(x, n=2 * n) + acf = np.fft.ifft(f * np.conjugate(f))[: len(x)].real + + if norm: + acf /= acf[0] + + return acf + + +def auto_window(taus, c=5.0): """ - Split a parameter key into (section, name) using the delimiter/prefix. + Automated windowing criterion. - If the key cannot be split, returns ("", key). + Finds the first lag M where M >= c * tau(M). """ - if prefix in key: - sect, name = key.split(prefix, 1) - return sect, name - return "", key + taus = np.asarray(taus, dtype=float) + + m = np.arange(len(taus)) < c * taus + + if np.any(~m): + return np.argmin(m) + + return len(taus) - 1 + +def autocorr_time_one_dim(y, c=5.0): + """ + Estimate autocorrelation time for one parameter. + + Parameters + ---------- + y : ndarray, shape (n_walker, n_sample) + Chains for one parameter. + + Returns + ------- + tau : float + Integrated autocorrelation time. + """ + y = np.asarray(y, dtype=float) + + if y.ndim != 2: + raise ValueError("y must have shape (n_walker, n_sample)") + + n_walker, n_sample = y.shape + + acf = np.zeros(n_sample) + + for walker in range(n_walker): + acf += autocorr_func_1d(y[walker]) + + acf /= n_walker + + taus = 2.0 * np.cumsum(acf) - 1.0 + + window = auto_window(taus, c=c) + + return taus[window] + +def autocorr_time_chain(chain, c=5.0): + """ + Estimate autocorrelation time for each dimension. + + Parameters + ---------- + chain : ndarray, shape (n_dim, n_walker, n_sample) + MCMC chain. + + Returns + ------- + taus : ndarray, shape (n_dim,) + Autocorrelation time for each parameter. + """ + chain = np.asarray(chain, dtype=float) + + if chain.ndim != 3: + raise ValueError("chain must have shape (n_dim, n_walker, n_sample)") + + n_dim, n_walker, n_sample = chain.shape + + taus = np.empty(n_dim) + + for dim in range(n_dim): + taus[dim] = autocorr_time_one_dim(chain[dim], c=c) + + return taus + +# Following the suggestion from Goodman & Weare (2010) +def autocorr_gw2010(y, c=5.0): + f = autocorr_func_1d(np.mean(y, axis=0)) + taus = 2.0 * np.cumsum(f) - 1.0 + window = auto_window(taus, c) + return taus[window] + + +def autocorr_new(y, c=5.0): + f = np.zeros(y.shape[1]) + # Loop over all variables + for variable in y: + f += autocorr_func_1d(variable) + f /= y.shape[0] + taus = 2.0 * np.cumsum(f) - 1.0 + window = auto_window(taus, c) + return taus[window] + +# ----------------------------------------------------------------------------- +# I/O helpers +# ----------------------------------------------------------------------------- +def flat_chain_to_walkers(data, nwalkers, nsamples): + nrows = data.shape[0] + expected = nwalkers * nsamples + if nrows != expected: + raise ValueError( + f"Cannot reshape chain with {nrows} rows into (nsteps={nsamples}, nwalkers={nwalkers})." + ) + return data.reshape((nsamples, nwalkers, -1)) def _logsumexp_weighted(a: np.ndarray, w: np.ndarray) -> float: a = _as_float_array(a).ravel() @@ -753,6 +949,11 @@ def _tolist(a): } return out + @classmethod + def from_json(cls): + #TODO + pass + def write_json(self, path: str, *, overwrite: bool = True, indent: int = 2) -> str: """Write a JSON summary file.""" if (not overwrite) and os.path.exists(path): @@ -829,15 +1030,27 @@ def to_fits(self) -> fits.HDUList: hdus.append(fits.BinTableHDU(t_pct, name="PERCENTILES", header=pct_hdr)) - # PDF1D table: store grid/pdf/kde per parameter as separate columns + # PDF1D table: store edges/pdf/kde per parameter as separate columns t_pdf1 = Table() + t_pdf1_edges = Table() for name, d in self.pdf_1d.items(): - t_pdf1[f"{name}_x"] = _as_float_array(d["grid"]) - t_pdf1[f"{name}_pdf"] = _as_float_array(d["hist_pdf"]) - t_pdf1[f"{name}_kde"] = _as_float_array(d.get("kde_pdf", np.full_like(d["grid"], np.nan))) + edges = _as_float_array(d["edges"]) + centers = _as_float_array(d.get("grid", 0.5 * (edges[:-1] + edges[1:]))) + hist_pdf = _as_float_array(d["hist_pdf"]) + kde_pdf = _as_float_array(d.get("kde_pdf", np.full_like(hist_pdf, np.nan))) + + # FITS bin table columns must have consistent lengths across rows. + # Store centers/PDF/KDE in PDF1D (all length = n_bins). + t_pdf1[f"{name}_x"] = centers + t_pdf1[f"{name}_pdf"] = hist_pdf + t_pdf1[f"{name}_kde"] = kde_pdf + # Store raw histogram edges separately (length = n_bins + 1). + t_pdf1_edges[f"{name}_edges"] = edges if len(t_pdf1.colnames) > 0: hdus.append(fits.BinTableHDU(t_pdf1, name="PDF1D")) + if len(t_pdf1_edges.colnames) > 0: + hdus.append(fits.BinTableHDU(t_pdf1_edges, name="PDF1D_EDGES")) # PDF2D images for (n0, n1), d in self.pdf_2d.items(): @@ -857,9 +1070,10 @@ def write_fits(self, path: str, *, overwrite: bool = True) -> str: hdul.writeto(path, overwrite=overwrite) return path - # ------------------------------------------------------------------------- - # Plot helpers (optional) - # ------------------------------------------------------------------------- + @classmethod + def from_fits(cls): + #TODO + pass def plot_1d_pdfs( self, @@ -953,15 +1167,7 @@ def corner_plot( dpi: int = 200, show: bool = False, ) -> str: - """ - Lightweight corner plot without external dependencies. - - Diagonal: 1D hist PDFs (weighted) - Off-diagonal: scatter of (subsampled) points colored by weight rank (simple) - - For serious usage, consider adding an optional dependency later (corner/arviz), - but this is a decent built-in baseline. - """ + """Build a corner plot""" npar = len(self.parameter_names) if npar == 0 or self.samples.size == 0: raise ValueError("No samples available to plot.") @@ -977,21 +1183,25 @@ def corner_plot( S = self.samples[:, idx] w = self.weights[idx] - fig, axes = plt.subplots(npar, npar, figsize=(2.2 * npar, 2.2 * npar), constrained_layout=True) + fig, axes = plt.subplots( + npar, npar, figsize=(2.2 * npar, 2.2 * npar), + sharex="col", + constrained_layout=True) for i in range(npar): for j in range(npar): ax = axes[i, j] if i == j: x = self.samples[i, :] - xc, pdf = histogram_pdf_1d(x, self.weights, bins=bins) + edges, pdf = histogram_pdf_1d(x, self.weights, bins=bins) + xc = 0.5 * (edges[:-1] + edges[1:]) ax.plot(xc, pdf, lw=1.2) # mark mean/MAP ax.axvline(self.mean[i], lw=1.0, alpha=0.8) ax.axvline(self.map[i], lw=1.0, alpha=0.8) ax.set_yticks([]) elif i > j: - ax.scatter(S[j, :], S[i, :], s=2, alpha=0.25) + ax.hist2d(S[j, :], S[i, :], weights=w) else: ax.axis("off") @@ -1021,6 +1231,8 @@ def summarize_results( *, output_fits: Optional[str] = None, output_json: Optional[str] = None, + nwalkers: Optional[int] = 1, + burn_in: int = 0, parameter_prefix: str = "--", posterior_key: str = "post", parameter_keys: Optional[Sequence[str]] = None, @@ -1079,10 +1291,14 @@ def summarize_results( ------- ResultsSummary """ + if burn_in > 0: + # discard the first burn_in samples per walker; assumes samples are ordered as (walker0, walker1, ..., walkerN, walker0, ...) + table = io.burn_table(table, nwalkers=nwalkers, burn_in=burn_in) + if posterior_key not in table.colnames: raise KeyError(f"posterior_key='{posterior_key}' not in table.") - keys = _select_parameter_keys(table, parameter_prefix=parameter_prefix, parameter_keys=parameter_keys) + keys = io._select_parameter_keys(table, parameter_prefix=parameter_prefix, parameter_keys=parameter_keys) # Extract and filter samples logpost_all = _as_float_array(table[posterior_key]) # Finite mask across posterior and all selected parameters @@ -1118,7 +1334,7 @@ def summarize_results( sections = [] names = [] for k in keys: - sect, nm = _split_param_key(k, prefix=parameter_prefix) + sect, nm = io._split_param_key(k, prefix=parameter_prefix) sections.append(sect) names.append(nm) @@ -1158,10 +1374,11 @@ def summarize_results( pdf1d: Dict[str, Dict[str, np.ndarray]] = {} if compute_1d: for i, nm in enumerate(names): - grid, hist_pdf = histogram_pdf_1d(samples[i, :], w, bins=pdf_bins_1d) - d = {"grid": grid, "hist_pdf": hist_pdf} + edges, hist_pdf = histogram_pdf_1d(samples[i, :], w, bins=pdf_bins_1d) + centers = 0.5 * (edges[:-1] + edges[1:]) + d = {"grid": centers, "edges": edges, "hist_pdf": hist_pdf} if kde_1d: - d["kde_pdf"] = kde_pdf_1d(samples[i, :], w, grid) + d["kde_pdf"] = kde_pdf_1d(samples[i, :], w, edges) pdf1d[nm] = d # 2D PDFs @@ -1250,7 +1467,7 @@ def summarize_results_file( **kwargs, ) -> ResultsSummary: """Read a results file and summarize it (passes kwargs to summarize_results).""" - tab = read_results_file(results_path, delimiter=delimiter) + tab = io.read_results_file(results_path, delimiter=delimiter) return summarize_results(tab, output_fits=output_fits, output_json=output_json, **kwargs) @@ -1360,100 +1577,6 @@ def weighted_quantiles(x: np.ndarray, w: np.ndarray, cdf = cdf / cdf[-1] return np.interp(qs, cdf, x) -def pdf_stats(edges: np.ndarray, pdf: np.ndarray, - quantiles=None, - find_multimodal: bool = False) -> dict: - """ - Compute summary stats from a discrete pdf over bin centers. - - Parameters - ---------- - edges : ndarray, shape (K,) - Bin edges. - pdf : ndarray, shape (K,) - PDF defined by the edges. - quantiles : tuple of float, optional - Quantiles to report (default 16, 50, 84 percent). - find_multimodal : bool, optional - If True, return rough modes by local-maximum search. - - Returns - ------- - stats : dict - Keys: mean, std, map, q, lo68, hi68, modes (optional). - """ - # Ensure normalization - norm = np.sum(pdf * np.diff(edges)) - pdf /= norm if norm > 0 else 1.0 - centers = 0.5 * (edges[:-1] + edges[1:]) - # Trapecium integrals - mean = np.sum(pdf * centers * np.diff(edges)) - var = np.sum(pdf * (centers - mean)**2 * np.diff(edges)) - std = var ** 0.5 - k_map = np.argmax(pdf * np.diff(edges)) - v_map = centers[k_map] - - cdf = np.cumsum(pdf * np.diff(edges)) - cdf = np.insert(cdf, 0, 0) - - if quantiles is not None: - qs = np.array(quantiles, float) - qvals = np.interp(qs, cdf, edges, left=edges[0], right=edges[-1]) - else: - qvals = None - - median = np.interp(0.5, cdf, edges, left=edges[0], right=edges[-1]) - lo68 = np.interp(0.16, cdf, edges, left=edges[0], right=edges[-1]) - hi68 = np.interp(0.84, cdf, edges, left=edges[0], right=edges[-1]) - - out = {"mean": mean, "std": std, "map": v_map, "q": qvals, - "median": median, "lo68": lo68, "hi68": hi68} - - if find_multimodal: - modes = [] - for i in range(1, len(pdf) - 1): - if pdf[i] > pdf[i - 1] and pdf[i] > pdf[i + 1]: - modes.append(centers[i]) - if not modes: - modes = [v_map] - out["modes"] = np.asarray(modes) - return out - - -def pit_from_discrete_posterior(z_true: np.ndarray, - posts: np.ndarray, - z_edges: np.ndarray) -> np.ndarray: - """ - Probability Integral Transform for discrete posteriors on bins. - - Assumes uniform density within each bin for within-bin interpolation. - - Parameters - ---------- - z_true : ndarray, shape (N,) - True values. - posts : ndarray, shape (N, K) - Row-normalised posteriors over K bins. - z_edges : ndarray, shape (K+1,) - Bin edges. - - Returns - ------- - pit : ndarray, shape (N,) - PIT values in [0, 1]. - """ - N, K = posts.shape - assert K == len(z_edges) - 1 - cdf_bins = np.cumsum(posts, axis=1) - j = np.clip(np.digitize(z_true, z_edges) - 1, 0, K - 1) - idx = np.arange(N) - below = np.where(j > 0, cdf_bins[idx, j - 1], 0.0) - widths = z_edges[1:] - z_edges[:-1] - frac = np.clip((z_true - z_edges[j]) / widths[j], 0.0, 1.0) - pit = below + posts[idx, j] * frac - return np.clip(pit, 0.0, 1.0) - - def photoz_metrics(z_true: np.ndarray, z_est: np.ndarray) -> dict: """ Standard photo-z metrics using delta z over 1+z. @@ -1471,6 +1594,135 @@ def photoz_metrics(z_true: np.ndarray, z_est: np.ndarray) -> dict: return {"bias": float(med), "nmad": float(nmad), "outlier": outlier, "rmse": rmse} +def specz_posterior(path: str, pct_val=[0.16, 0.5, 0.84]) -> dict: + """Load spectral redshift posterior from a file. + + Parameters + ---------- + path : str + Path to the posterior file. + pct_val : list of float + Percentiles to compute (default: 16, 50, 84). + + Returns + ------- + results : dict + Keys: pct, mean, var, modes, mode_loglike, mode_log_amplitude. + """ + z, loglike = np.loadtxt(path, dtype=np.float64, skiprows=1, unpack=True) + + # sort by z + idx = np.argsort(z) + z = z[idx] + loglike = loglike[idx] + # Renormalize to get a proper PDF + loglike -= np.nanmax(loglike) + like = np.exp(loglike) + pdf = like / np.trapz(like, z) + + pct = weighted_quantile(z, like, pct_val) + mean = weighted_mean(z, like) + var = weighted_covariance(z[None, :], like, unbiased=False)[0, 0] + # Analyze multimodality (simple local maxima) + modes = [] + modes_loglike = [] + modes_idx = [] # list of list of indices corresponding to modes + modes_pct = [] + modes_mean = [] + modes_var = [] + modes_log_amplitude = [] + modes_evidence = [] + idx_cont = [] # to track indices contributing to modes for continuum estimation + + # Step 1: find local minima + for i in range(1, len(z) - 1): + if loglike[i] < loglike[i - 1] and loglike[i] < loglike[i + 1]: + idx_cont.append(i) + + # Step 2: characterise local maxima and assign mode indices + start = 0 + for i in range(len(idx_cont) + 1): # add end index to capture last segment + if i < len(idx_cont): + segment_idx = range(start, idx_cont[i]) + else: + segment_idx = range(start, len(z)) + if len(segment_idx) == 0: + continue + # Find local maximum in this segment + seg_loglike = loglike[segment_idx] + max_idx_in_seg = np.argmax(seg_loglike) + global_idx = segment_idx[max_idx_in_seg] + modes.append(z[global_idx]) + modes_loglike.append(loglike[global_idx]) + modes_idx.append(list(segment_idx)) + + # mode quantities + mode_like = np.exp(seg_loglike - loglike[global_idx]) + mode_mean = weighted_mean(z[segment_idx], mode_like) + mode_var = weighted_covariance( + z[segment_idx][None, :], mode_like, unbiased=False + )[0, 0] + mode_pct = weighted_quantile(z[segment_idx], mode_like, pct_val) + modes_mean.append(mode_mean) + modes_var.append(mode_var) + modes_pct.append(mode_pct) + # mode loglike amplitude above local continuum + left_cont = loglike[segment_idx[0]] if segment_idx[0] > 0 else loglike[0] + right_cont = loglike[segment_idx[-1]] if segment_idx[-1] < len(z) - 1 else loglike[-1] + + cont_loglike = np.interp(z[global_idx], [z[segment_idx[0]], z[segment_idx[-1]]], [left_cont, right_cont]) + mode_log_amplitude = loglike[global_idx] - cont_loglike + modes_log_amplitude.append(mode_log_amplitude) + + # mode evidence + mode_evidence = np.trapz(pdf[segment_idx], z[segment_idx]) + modes_evidence.append(mode_evidence) + + if i < len(idx_cont): + start = idx_cont[i] + 1 + + if modes: + modes = np.array(modes) + modes_loglike = np.array(modes_loglike) + modes_log_amplitude = np.array(modes_log_amplitude) + modes_evidence = np.array(modes_evidence) + modes_pct = np.array(modes_pct) + modes_mean = np.array(modes_mean) + modes_var = np.array(modes_var) + + # from matplotlib import pyplot as plt + # plt.figure() + # plt.plot(z, loglike, label="loglike") + # plt.scatter(modes, modes_loglike, c=np.log(modes_evidence), label="modes") + # plt.colorbar() + # plt.legend() + else: + modes = np.array([z[np.argmax(loglike)]]) + modes_loglike = np.array([np.max(loglike)]) + modes_log_amplitude = np.array([0.0]) + modes_evidence = np.array([np.trapz(np.exp(loglike), z)]) + modes_evidence_contsub = np.array([0.0]) + modes_pct = np.array([0.0]) + modes_mean = np.array([0.0]) + modes_var = np.array([0.0]) + modes_idx = [list(range(len(z)))] + + results = { + "pct": pct, + "mean": mean, + "var": var, + "modes": modes, + "mode_loglike": modes_loglike, + "mode_log_amplitude": modes_log_amplitude, + "mode_evidence": modes_evidence, + "mode_pct": modes_pct, + "mode_mean": modes_mean, + "mode_var": modes_var, + "mode_indices": modes_idx, + } + return results + + def plot_chains(table, truth_values=None, output_dir=None, posterior_key="post"): """Make trace plots from an astropy Table containing chain results. diff --git a/src/besta/samplers/__init__.py b/src/besta/samplers/__init__.py new file mode 100644 index 00000000..dbc3b464 --- /dev/null +++ b/src/besta/samplers/__init__.py @@ -0,0 +1,5 @@ +"""Custom sampler implementations shipped with BESTA.""" + +from .pymc_sampler import PymcSampler, PyMCSampler + +__all__ = ["PymcSampler", "PyMCSampler"] diff --git a/src/besta/samplers/pymc_sampler.py b/src/besta/samplers/pymc_sampler.py new file mode 100644 index 00000000..42192d8b --- /dev/null +++ b/src/besta/samplers/pymc_sampler.py @@ -0,0 +1,209 @@ +"""PyMC sampler compatible with modern PyMC versions. + +This module is designed to be loaded by CosmoSIS using the runtime option +`import_samplers = /path/to/pymc_sampler.py`. +""" + +from __future__ import annotations + +from typing import Tuple + +import numpy as np + +from cosmosis.runtime import logs +from cosmosis.samplers import ParallelSampler + + +def _build_logpost_op(pipeline): + """Create a PyTensor Op that evaluates the CosmoSIS log-posterior.""" + try: + from pytensor import wrap_py as _wrap_py + except ImportError: # pragma: no cover - compatibility with older PyTensor + from pytensor.compile.ops import as_op as _wrap_py + import pytensor.tensor as pt + + @_wrap_py(itypes=[pt.dvector], otypes=[pt.dscalar]) + def logpost_from_unit_cube(theta_unit): + theta_unit = np.asarray(theta_unit, dtype=float) + if theta_unit.ndim != 1: + return np.array(-np.inf, dtype=np.float64) + if not np.all(np.isfinite(theta_unit)): + return np.array(-np.inf, dtype=np.float64) + + try: + params = pipeline.denormalize_vector(theta_unit) + except Exception: + return np.array(-np.inf, dtype=np.float64) + + result = pipeline.run_results(params) + post = float(result.post) + if not np.isfinite(post): + return np.array(-np.inf, dtype=np.float64) + return np.array(post, dtype=np.float64) + + return logpost_from_unit_cube + + +class PymcSampler(ParallelSampler): + """Modern PyMC sampler for CosmoSIS pipelines. + + Notes + ----- + - Uses a Metropolis step method to avoid requiring gradients through the + black-box pipeline likelihood. + - Samples in normalized unit-cube coordinates and denormalizes via + `pipeline.denormalize_vector` when evaluating the posterior. + """ + + sampler_outputs = [("prior", float), ("post", float), ("like", float)] + parallel_output = False + supports_smp = False + + def _make_step_method(self): + """Build the configured PyMC step method.""" + name = self.step_method_name + + if name == "metropolis": + return self.pm.Metropolis() + if name == "demetropolis": + return self.pm.DEMetropolis() + if name == "demetropolisz": + return self.pm.DEMetropolisZ() + if name == "slice": + return self.pm.Slice() + + if name in {"nuts", "hmc", "hamiltonianmc"}: + raise ValueError( + f"step_method={name!r} requires gradients, but this sampler wraps " + "a black-box CosmoSIS likelihood and does not provide gradients. " + "Use one of: metropolis, demetropolis, demetropolisz, slice." + ) + + raise ValueError( + f"Unknown step_method={name!r}. " + "Valid options are: metropolis, demetropolis, demetropolisz, slice." + ) + + def config(self): + try: + import pymc as pm + except ImportError as exc: + raise ImportError( + "PyMC is required for besta.samplers.PymcSampler. " + "Install it in your active environment (for example: pip install pymc)." + ) from exc + + self.pm = pm + self._model = None + self._done = False + + self.ndim = len(self.pipeline.varied_params) + self.samples = max(1, self.read_ini("samples", int, 1000)) + self.nsteps = max(1, self.read_ini("nsteps", int, self.samples)) + self.chains = max(1, self.read_ini("chains", int, 1)) + self.target_accept = float(self.read_ini("target_accept", float, 0.8)) + self.progressbar = bool(self.read_ini("progressbar", bool, False)) + self.step_method_name = self.read_ini("step_method", str, "demetropolisz").strip().lower() + + fburn = self.read_ini("burn_fraction", float, 0.0) + if 0.0 <= fburn < 1.0: + self.nburn = int(fburn * self.samples) + else: + self.nburn = max(0, int(fburn)) + + seed_raw = self.read_ini("seed", int, -1) + self.random_seed = None if seed_raw < 0 else seed_raw + + self._logpost_op = _build_logpost_op(self.pipeline) + + with self.pm.Model() as model: + theta_unit = self.pm.Uniform( + "theta_unit", lower=0.0, upper=1.0, shape=(self.ndim,) + ) + self.pm.Potential("cosmosis_logpost", self._logpost_op(theta_unit)) + self._model = model + + def _evaluate_samples(self, theta_unit: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Evaluate pipeline diagnostics for sampled points and stream to output.""" + traces = [] + posts = [] + + for u in theta_unit: + try: + params = np.asarray(self.pipeline.denormalize_vector(u), dtype=float) + except Exception: + continue + + result = self.pipeline.run_results(params) + prior = float(result.prior) + post = float(result.post) + like = float(result.like) + + if not np.isfinite(post): + continue + + extra = np.asarray(result.extra) + self.output.parameters(params, extra, prior, post, like) + + traces.append(params) + posts.append(post) + + if traces: + traces_arr = np.asarray(traces, dtype=float) + posts_arr = np.asarray(posts, dtype=float) + self.distribution_hints.set_from_sample(traces_arr, posts_arr) + return traces_arr, posts_arr + + return np.empty((0, self.ndim), dtype=float), np.empty((0,), dtype=float) + + def execute(self): + if self._done: + return + + draws = max(1, self.samples) + tune = max(0, self.nburn) + + logs.overview( + f"Running PyMC with draws={draws}, tune={tune}, chains={self.chains}, " + f"nsteps={self.nsteps}" + ) + + with self._model: + idata = self.pm.sample( + draws=draws, + tune=tune, + chains=self.chains, + cores=1, + random_seed=self.random_seed, + progressbar=self.progressbar, + compute_convergence_checks=False, + discard_tuned_samples=True, + return_inferencedata=True, + step=self._make_step_method(), + ) + + theta = np.asarray(idata.posterior["theta_unit"]) + theta = theta.reshape((-1, self.ndim)) + + traces, posts = self._evaluate_samples(theta) + self.num_samples = traces.shape[0] + self._done = True + + if posts.size: + logs.overview( + f"Done PyMC sampling with {self.num_samples} samples. " + f"max(post)={posts.max():.6g}" + ) + else: + logs.overview("Done PyMC sampling but no finite-posterior samples were recorded.") + + def worker(self): + while not self.is_converged(): + self.execute() + + def is_converged(self): + return self._done + + +# legacy alias +PyMCSampler = PymcSampler diff --git a/src/besta/sfh.py b/src/besta/sfh.py index 56872c60..43749122 100644 --- a/src/besta/sfh.py +++ b/src/besta/sfh.py @@ -97,9 +97,21 @@ def make_ini(self, ini_file): Path to the output .ini file. """ logger.info("Making ini file: %s", ini_file) + + free_params = self.free_params.copy() + + if self.use_transforms: + logger.info("transforming default SFH values into latent variables") + if getattr(self, "sfh_bin_keys", None): + sfh_values = self.to_latent([free_params[k][1] for k in self.sfh_bin_keys]) + for key, val in zip(self.sfh_bin_keys, sfh_values): + free_params[key] = [-5, val, 5] + with open(ini_file, "w", encoding="utf-8") as file: + file.write(f"; Default prior file for SFH model: {str(self.__class__)}\n") + file.write(f"; use_transforms: {str(self.use_transforms)}\n") file.write(f"[{self.sect_name}]\n") - for key, val in self.free_params.items(): + for key, val in free_params.items(): if len(val) > 1: file.write(f"{key} = {val[0]} {val[1]} {val[2]}\n") else: @@ -169,15 +181,13 @@ class FixedTimeSFH(ZPowerLawMixin, SFHBase, PieceWiseSFHMixin): The SFH of a galaxy is modelled as a stepwise function where the free parameters correspond to the mass fraction formed on each bin. - Upon initializaiton, the free parameters are set between -8 to 0 in terms - of log(M/Msun). The starting point corresponds to the mass fraction formed - assuming a constant star formation history. + The default starting point for uniform priors corresponds to the mass + fraction formed assuming a constant star formation history. Attributes ---------- lookback_time : astropy.units.Quantity Lookback time bin edges. - """ def __init__(self, lookback_time_bins, *args, **kwargs): @@ -301,12 +311,15 @@ def __init__(self, lookback_time, *args, **kwargs): # Maximum value of the sSFR max_logssfr = np.log10(1 / lbt) self.max_ssfr_logyr[ith] = max_logssfr - self.free_params[k] = [ - -14.0, - np.log10(1 / self.today.to_value("yr")), - max_logssfr, - ] - + if not self.use_transforms: + self.free_params[k] = [ + -14.0, + np.log10(1 / self.today.to_value("yr")), + max_logssfr, + ] + else: + self.free_params[k] = [-10.0, 0.0, 10.0] + # log(tau1 / tau2) where tau1 > tau2 self.delta_logtau = - np.diff(np.log10(self.lookback_time.to_value("yr"))) @@ -395,11 +408,18 @@ def __init__(self, mass_fraction, *args, **kwargs): for frc in mass_fraction: k = f"t_at_frac_{frc:.4f}" self.sfh_bin_keys.append(k) - self.free_params[k] = [ - 1e-3, - frc * self.today.to_value("Gyr"), - self.today.to_value("Gyr") * 0.999, - ] + if not self.use_transforms: + self.free_params[k] = [ + 1e-3, + frc * self.today.to_value("Gyr"), + self.today.to_value("Gyr") * 0.999, + ] + else: + self.free_params[k] = [-10.0, 0.0, 10.0] + + if self.use_transforms: + logger.info("Using transforms to enforce monotonicity and bounds on time bins.") + mass_fraction = mass_fraction[:-1] self.model = cem.TabularMassFracCEM( mass_frac=mass_fraction, @@ -408,40 +428,42 @@ def __init__(self, mass_fraction, *args, **kwargs): mass_today=1 << u.Msun, ism_metallicity_today=kwargs.get("ism_metallicity_today", 0.02) << u.dimensionless_unscaled, - alpha_powerlaw=kwargs.get("alpha", 0.0), + alpha_powerlaw=kwargs.get("alpha_powerlaw", 0.0), ) def parse_datablock(self, datablock: DataBlock): """Update the fixed-mass-fraction SFH model from a CosmoSIS DataBlock.""" times = self.get_sfh_parameters_array(datablock) - if self.use_transforms: - # Enforce strictly increasing times within [0, today] - deltas = np.exp(times) # positive - times = np.cumsum(deltas) - # TEMPFIX: - times = times / times[-1] * self.today.to_value("Gyr") * 0.999 - # TODO: this enforces that mass_fraction[-1] occurs at present time - # which is unphysical. This transformed sampling should therefore - # include and additional fraction (*mass_frac, today). - # Ensure monotonically increasing and always smaller than the age of the Universe - delta_t = times[1:] - times[:-1] - if (delta_t <= 0).any() or times[-1] >= self.today.to_value("Gyr"): - return 0, 1 + np.abs(delta_t[delta_t < 0].sum()) - # Update the mass of the tabular model - self.model.times = times << u.Gyr self.model.alpha_powerlaw = datablock[self.sect_name, "alpha_powerlaw"] self.model.ism_metallicity_today = ( datablock[self.sect_name, "ism_metallicity_today"] << u.dimensionless_unscaled ) + + if self.use_transforms: + times = self.to_physical(times) + times = np.insert(times, 0, 0) + # Bypass the setter and avoid the insert + self.model._times = Parameter( + times << u.Gyr, + fixed=False, + doc="Observing-time SFH anchors", + ) + # Ensure monotonically increasing and always smaller than the age of the Universe + else: + delta_t = times[1:] - times[:-1] + if (delta_t <= 0).any() or times[-1] >= self.today.to_value("Gyr"): + return 0, 1 + np.abs(delta_t[delta_t < 0].sum()) + # Update the mass of the tabular model + self.model.times = times << u.Gyr return 1, None def to_physical(self, latent): """Map unconstrained latents to strictly increasing times (Gyr).""" if self.use_transforms: - deltas = np.exp(latent) - times = np.cumsum(deltas) - times = times / times[-1] * self.today.to_value("Gyr") + # Softmax transform + time_frac = _softmax(latent) + times = np.cumsum(time_frac) * self.today.to_value("Gyr") return times return np.asarray(latent, dtype=float) diff --git a/src/besta/spectrum.py b/src/besta/spectrum.py index 937012e2..671cd245 100644 --- a/src/besta/spectrum.py +++ b/src/besta/spectrum.py @@ -27,6 +27,69 @@ logger = get_logger(__name__) +def vacuum_to_air_wavelength(wavelength_vacuum): + """Convert vacuum wavelengths to air wavelengths Morton (1991, ApJS, 77, 119). + + Parameters + ---------- + wavelength_vacuum : array-like + Wavelengths in vacuum (in Angstrom). + + Returns + ------- + array-like + Wavelengths in air (in Angstrom). + """ + if isinstance(wavelength_vacuum, u.Quantity): + wl = wavelength_vacuum.to_value(u.AA) + unit = u.AA + else: + wl = np.asanyarray(wavelength_vacuum) + unit = 1 + + wl_air = wl / (1.0 + 2.735182e-4 + 131.4182 / wl**2 + 2.76249e8 / wl**4) + return wl_air * unit + +def wavelength_offset_to_velocity(wavelength, rest_wavelength): + """Convert a wavelength offset to a velocity in km/s. + + Parameters + ---------- + wavelength : array-like + Observed wavelength (in Angstrom). + rest_wavelength : array-like + Reference rest wavelength (in Angstrom). + + Returns + ------- + array-like + Velocity offset in km/s. If input wavelengths are :class:`astropy.units.Quantity`, + the output will be a Quantity in km/s. Otherwise, it will be a dimensionless array. + """ + # Preserve units if input is Quantity, otherwise return dimensionless array + if isinstance(wavelength, u.Quantity): + return (constants.c * (wavelength - rest_wavelength) / rest_wavelength).to(u.km / u.s) + return constants.c.to_value(u.km / u.s) * (wavelength - rest_wavelength) / rest_wavelength + +def wavelength_dispersion_to_velocity_dispersion(wavelength_sigma, rest_wavelength): + """Convert a wavelength dispersion (sigma) to a velocity dispersion in km/s. + + Parameters + ---------- + wavelength_sigma : array-like + Wavelength dispersion (sigma) in Angstrom. + rest_wavelength : array-like + Reference rest wavelength in Angstrom. + + Returns + ------- + array-like + Velocity dispersion in km/s. If input is Quantity, output will be Quantity in km/s. Otherwise, dimensionless array. + """ + if isinstance(wavelength_sigma, u.Quantity): + return (constants.c * wavelength_sigma / rest_wavelength).to(u.km / u.s) + return constants.c.to_value(u.km / u.s) * wavelength_sigma / rest_wavelength + def get_legendre_polynomial_array( wavelength, order, bounds=None, scale=None, clip_first_zero=True ): @@ -245,6 +308,14 @@ class EmissionLine: flag: int = 0 metadata: dict = field(default_factory=dict) + def velocity_offset(self, observed_wavelength, redshift=0.0): + """Compute velocity offset of the line given an observed wavelength.""" + return wavelength_offset_to_velocity(observed_wavelength, self.rest_wavelength * (1 + redshift)) + + def velocity_dispersion(self, wavelength_sigma, redshift=0.0): + """Compute velocity dispersion of the line given a wavelength sigma.""" + return wavelength_dispersion_to_velocity_dispersion( + wavelength_sigma, self.rest_wavelength * (1 + redshift)) class EmissionLineList: """Collection of emission lines with helper methods.""" @@ -532,6 +603,9 @@ def from_file(cls, filename: str, **table_kwargs) -> EmissionLineList: def get_default_emission_lines(): """Return the default optical emission-line list. + Lines are expressed in air wavelengths, and the default half-widths (in AA) are + nominal values in Angstrom for masking purposes. + Returns ------- EmissionLineList @@ -583,7 +657,16 @@ def get_default_sky_emission_lines(): def _parse_lines_param_decorator(func): - """Decorator to parse the `lines` parameter.""" + """Decorator to parse the `lines` parameter in a function. + + Description + ----------- + + This decorator is designed to be applied to functions that accept a `lines` + parameter, which can be provided in multiple formats. The decorator will + handle the parsing of the `lines` parameter and convert it into a standardized + format (an instance of `EmissionLineList`) before passing it to the decorated function. + """ @wraps(func) def wrapper(*args, **kwargs): @@ -1238,7 +1321,26 @@ def _gaussian_fit(self, wl, flux, err, weights, line_id, line_mask): ) def fit_line(self, line_id: int) -> dict: - """Fit a Gaussian to the specified line_id and return fit parameters.""" + """Fit a Gaussian to the specified line_id and return fit parameters. + + Parameters + ---------- + line_id : int + The integer ID of the line segment to fit (corresponding to the segmentation map). + + Returns + ------- + dict + A dictionary containing the fit parameters and metadata for the line: + - id: line_id + - line_flux: integrated flux of the fitted Gaussian + - line_flux_err: uncertainty of the line flux + - center: fitted central wavelength of the line + - sigma: fitted Gaussian sigma (line width in wavelength units) + - npixels: number of valid pixels used in the fit + - flag: fit quality flag (0=good fit, 1=no valid pixels, 2=not enough + pixels for fit, 3=fit failed, used MLE estimates) + """ mask = self.get_line_mask(line_id) return self._gaussian_fit( wl=self.wavelength, @@ -1250,7 +1352,31 @@ def fit_line(self, line_id: int) -> dict: ) def fit_all_lines(self) -> Table: - """Fit all lines in the segmentation map and return a table of results.""" + """Fit all lines in the segmentation map and return a table of results. + + Description + ----------- + This method iterates over all unique line IDs in the segmentation map, + fits a Gaussian profile to each line segment using the continuum-subtracted flux, + and compiles the fit parameters into an Astropy Table. + + It also constructs an EmissionLineList with the fitted line parameters + and builds a composite emission line spectrum (``self.eline_flux``) by + summing the fitted Gaussians for all lines. + + Returns + ------- + Table + An Astropy Table containing the fit parameters for each line segment, with columns: + - id: line_id + - line_flux: integrated flux of the fitted Gaussian + - line_flux_err: uncertainty of the line flux + - center: fitted central wavelength of the line + - sigma: fitted Gaussian sigma (line width in wavelength units) + - npixels: number of valid pixels used in the fit + - flag: fit quality flag (0=good fit, 1=no valid pixels, 2=not enough + pixels for fit, 3=fit failed, used MLE estimates) + """ output_table = Table( names=[ "id", @@ -1267,6 +1393,7 @@ def fit_all_lines(self) -> Table: "flags": "0=good fit, 1=no valid pixels, 2=not enough pixels for fit, 3=fit failed, used MLE estimates", }, ) + measured_lines = [] for line_id in range(1, self.nlines + 1): fit_params = self.fit_line(line_id) @@ -1308,8 +1435,10 @@ def _watershed_1d( """ 1D watershed segmentation for deblending overlapping emission lines. - Starting from labeled seed regions (markers), flood outward following - descending signal gradient until regions meet or the mask boundary is reached. + This function performs a simple 1D watershed flood-fill. It starts from + already-labeled seed pixels in markers and expands those labels into nearby + unlabeled pixels inside mask, always filling the currently highest-signal + available pixel first. Parameters ---------- @@ -1333,33 +1462,33 @@ def _watershed_1d( in_queue = np.zeros(len(signal), dtype=bool) # Priority queue: (-signal_value, pixel_index, label) - # Negative because heapq is a min-heap; we want to process highest signal first + # to use heapq (which is a min-heap, for smallest values are procesed first). heap = [] - - # Seed the queue with all boundary pixels of each marker region + # First seed the queue with all boundary pixels of each marker region for i in range(len(signal)): if labels[i] > 0 and mask[i]: for di in (-1, 1): nb = i + di + # Pixel nb is currently unlabeled, but it touches label[i], + # so it is eligible for flooding if ( 0 <= nb < len(signal) and mask[nb] and labels[nb] == 0 and not in_queue[nb] ): + # Add to the queue heappush(heap, (-signal[nb], nb, labels[i])) in_queue[nb] = True + # Flood in descending signal order while heap: neg_val, idx, lbl = heappop(heap) - # Skip if already labeled by a previous (higher-priority) flood if labels[idx] != 0: continue - labels[idx] = lbl - - # Expand to unlabeled neighbours + # Expand again to unlabeled neighbours for di in (-1, 1): nb = idx + di if ( @@ -1368,9 +1497,9 @@ def _watershed_1d( and labels[nb] == 0 and not in_queue[nb] ): + # Flood this neighbour next, with the same label heappush(heap, (-signal[nb], nb, lbl)) in_queue[nb] = True - return labels @@ -1571,11 +1700,18 @@ def find_emission_lines( if lines is not None: # Match detected lines to input line list based on proximity of centers + logger.info( + "Matching detected lines to input line list with redshift z=%.3f", + redshift) matched_line_names = [] + matched_line_velocities = [] + matched_line_velocity_dispersion = [] for row in output_table: line_center = row["center"] if not np.isfinite(line_center): matched_line_names.append("unknown") + matched_line_velocities.append(np.nan) + matched_line_velocity_dispersion.append(np.nan) continue matched_line = lines.get_closest_line(line_center, redshift=redshift) tolerance = ( @@ -1589,15 +1725,31 @@ def find_emission_lines( < tolerance ): matched_line_names.append(matched_line.name) + matched_line_velocities.append( + matched_line.velocity_offset(line_center, redshift=redshift)) + matched_line_velocity_dispersion.append( + matched_line.velocity_dispersion(row["sigma"], redshift=redshift)) else: matched_line_names.append("unknown") + matched_line_velocities.append(np.nan) + matched_line_velocity_dispersion.append(np.nan) + output_table["line_name"] = matched_line_names + output_table["velocity"] = matched_line_velocities + output_table["velocity_dispersion"] = matched_line_velocity_dispersion + if line_segm_map.lines is not None: line_segm_map.lines = line_segm_map.lines.with_names(matched_line_names) if to_rest_frame: output_table["center"] = output_table["center"] / (1 + redshift) output_table["sigma"] = output_table["sigma"] / (1 + redshift) + if lines is not None: + output_table["velocity"] = output_table["velocity"] + output_table["velocity_dispersion"] = output_table["velocity_dispersion"] + output_table.meta["rest_frame"] = True + output_table.meta["redshift"] = redshift + if line_segm_map.lines is not None: line_segm_map.lines = EmissionLineList( [ diff --git a/tests/_test_pymc_sampler.py b/tests/_test_pymc_sampler.py new file mode 100644 index 00000000..9a0b04f4 --- /dev/null +++ b/tests/_test_pymc_sampler.py @@ -0,0 +1,250 @@ +from __future__ import annotations + +import shutil +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +from cosmosis.output import InMemoryOutput + +from besta import io +from besta.pipeline import MainPipeline +from besta.samplers.pymc_sampler import PymcSampler + + +def _write_ini(path: Path, text: str) -> Path: + path.write_text(text, encoding="utf-8") + return path + + +def test_pymc_sampler_smoke_with_dummy_pipeline(tmp_path): + class DummyPipeline: + def __init__(self): + self.varied_params = [object(), object()] + self.modules = [] + + def output_names(self): + # Varied params + one extra output. + return ["p0", "p1", "extra0"] + + def denormalize_vector(self, u): + # Map unit-cube to a bounded physical space. + return 4.0 * np.asarray(u, dtype=float) - 2.0 + + def run_results(self, p): + p = np.asarray(p, dtype=float) + post = -0.5 * float(np.dot(p, p)) + prior = -0.1 * float(np.dot(p, p)) + like = post - prior + extra = np.array([float(p.sum())], dtype=float) + return SimpleNamespace(post=post, prior=prior, like=like, extra=extra) + + ini_path = _write_ini( + tmp_path / "pymc_smoke.ini", + """ +[runtime] +sampler = pymc + +[pymc] +samples = 24 +nsteps = 8 +burn_fraction = 0.0 +chains = 1 +progressbar = F +seed = 11 +""".strip(), + ) + + output = InMemoryOutput() + sampler = PymcSampler(str(ini_path), DummyPipeline(), output) + sampler.config() + sampler.execute() + + assert sampler.is_converged() + assert sampler.num_samples > 0 + assert len(output.rows) == sampler.num_samples + + col_names = [c[0] for c in output.columns] + assert "prior" in col_names + assert "post" in col_names + assert "like" in col_names + + +def test_pymc_sampler_accepts_slice_step_method(tmp_path): + class DummyPipeline: + def __init__(self): + self.varied_params = [object(), object()] + self.modules = [] + + def output_names(self): + return ["p0", "p1", "extra0"] + + def denormalize_vector(self, u): + return 4.0 * np.asarray(u, dtype=float) - 2.0 + + def run_results(self, p): + p = np.asarray(p, dtype=float) + post = -0.5 * float(np.dot(p, p)) + prior = -0.1 * float(np.dot(p, p)) + like = post - prior + extra = np.array([float(p.sum())], dtype=float) + return SimpleNamespace(post=post, prior=prior, like=like, extra=extra) + + ini_path = _write_ini( + tmp_path / "pymc_slice.ini", + """ +[runtime] +sampler = pymc + +[pymc] +samples = 12 +nsteps = 4 +burn_fraction = 0.0 +chains = 1 +progressbar = F +seed = 5 +step_method = slice +""".strip(), + ) + + output = InMemoryOutput() + sampler = PymcSampler(str(ini_path), DummyPipeline(), output) + sampler.config() + sampler.execute() + + assert sampler.is_converged() + assert sampler.num_samples > 0 + + +def test_pymc_sampler_rejects_gradient_step_methods(tmp_path): + class DummyPipeline: + def __init__(self): + self.varied_params = [object(), object()] + self.modules = [] + + def output_names(self): + return ["p0", "p1", "extra0"] + + def denormalize_vector(self, u): + return 4.0 * np.asarray(u, dtype=float) - 2.0 + + def run_results(self, p): + p = np.asarray(p, dtype=float) + post = -0.5 * float(np.dot(p, p)) + prior = -0.1 * float(np.dot(p, p)) + like = post - prior + extra = np.array([float(p.sum())], dtype=float) + return SimpleNamespace(post=post, prior=prior, like=like, extra=extra) + + ini_path = _write_ini( + tmp_path / "pymc_nuts.ini", + """ +[runtime] +sampler = pymc + +[pymc] +samples = 6 +nsteps = 3 +burn_fraction = 0.0 +chains = 1 +progressbar = F +seed = 1 +step_method = nuts +""".strip(), + ) + + output = InMemoryOutput() + sampler = PymcSampler(str(ini_path), DummyPipeline(), output) + sampler.config() + with pytest.raises(ValueError, match="requires gradients"): + sampler.execute() + + +def test_pymc_sampler_full_besta_run(tmp_path): + cosmosis_exe = shutil.which("cosmosis") + if cosmosis_exe is None: + pytest.skip("cosmosis executable not available in PATH") + + module_path = tmp_path / "dummy_like_module.py" + module_path.write_text( + """ +def setup(options): + return {} + + +def execute(block, config): + p1 = block['parameters', 'p1'] + p2 = block['parameters', 'p2'] + like = -0.5 * (p1**2 + p2**2) + block['extra', 'sum_p'] = p1 + p2 + block['likelihoods', 'dummy_like'] = like + return 0 + + +def cleanup(config): + return 0 +""".strip(), + encoding="utf-8", + ) + + output_root = tmp_path / "pymc_full_run" + values_path = tmp_path / "values.ini" + values_path.write_text( + """ +[parameters] +p1 = -2.0 0.0 2.0 +p2 = -2.0 0.0 2.0 +""".strip(), + encoding="utf-8", + ) + + sampler_path = Path(__file__).resolve().parents[1] / "src" / "besta" / "samplers" / "pymc_sampler.py" + + config = { + "runtime": { + "sampler": "pymc", + "import_samplers": str(sampler_path), + }, + "pymc": { + "samples": 16, + "nsteps": 8, + "burn_fraction": 0.0, + "chains": 1, + "progressbar": False, + "seed": 3, + }, + "output": { + "filename": str(output_root), + "format": "text", + }, + "pipeline": { + "modules": "DummyLike", + "values": str(values_path), + "likelihoods": "dummy", + "quiet": "T", + "debug": "T", + "extra_output": "extra/sum_p", + }, + "DummyLike": { + "file": str(module_path), + }, + } + + runner = MainPipeline( + [config], + n_cores_list=[1], + ini_values_files=[str(values_path)], + ) + status = runner.execute_all(plot_result=False) + assert status == 0 + + results_file = Path(str(output_root) + ".txt") + assert results_file.exists() + + table = io.read_results_file(str(results_file)) + assert len(table) > 0 + assert "post" in table.colnames + assert "prior" in table.colnames + assert "like" in table.colnames diff --git a/tests/test_fitting_modules.py b/tests/test_fitting_modules.py index b6808082..7d183681 100644 --- a/tests/test_fitting_modules.py +++ b/tests/test_fitting_modules.py @@ -60,7 +60,7 @@ def test_full_spectral_fit_make_observable(tmp_path): spec = make_dummy_spectrum(tmp_path) block = DataBlock() all_params = { - "dust.extinction": {"a_v": 0.0}, + "dust_attenuation": {"a_v": 0.0}, "kinematics": { "los_vel": 0.0, "los_sigma": 100.0, @@ -95,3 +95,9 @@ def test_full_spectral_fit_make_observable(tmp_path): flux_model, weights = mod.make_observable(block) assert flux_model.shape == mod.config["flux"].shape assert weights.shape == mod.config["flux"].shape + +if __name__ == "__main__": + import sys + import unittest + + unittest.main(argv=[sys.argv[0]]) \ No newline at end of file diff --git a/tests/test_grid.py b/tests/test_grid.py index 31ca4854..65561943 100644 --- a/tests/test_grid.py +++ b/tests/test_grid.py @@ -3,7 +3,7 @@ from besta.grid.grid import ModelGrid, GridFitter, _truncate_posterior_mass from besta.grid.binning import RectBinner, KDTreeBinner, HashedGridBinner, NestedBinner -from besta.grid.transforms import LinearStandardiser, MagTransform +from besta.grid.transforms import LinearStandardiser from besta.grid.prob import ( FlatPrior, GaussianPrior1D, @@ -18,7 +18,6 @@ CompositePrior, ObservableCompositePrior, GaussianProductLikelihood, - CensoredSizeLikelihood, CompositeLikelihood, NumbaGaussianProductLikelihood, posterior_over_models, @@ -181,8 +180,7 @@ def test_modelgrid_hdf5_and_fits_roundtrip(tmp_path): np.testing.assert_allclose(g_fits.observables, grid.observables) np.testing.assert_allclose(g_fits.targets, grid.targets) - -def test_linear_standardiser_and_mag_transform_roundtrip(): +def test_linear_standardiser_roundtrip(): X = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) std = LinearStandardiser().fit(X) Xz = std.transform(X) @@ -192,12 +190,6 @@ def test_linear_standardiser_and_mag_transform_roundtrip(): std2 = LinearStandardiser.from_dict(state) np.testing.assert_allclose(std2.transform(X), Xz) - mt = MagTransform(zero_point=25.0) - flux = np.array([1e-1, 1.0, 10.0]) - mag = mt.flux_to_mag(flux) - np.testing.assert_allclose(mt.mag_to_flux(mag), flux) - - def test_rect_binner_fit_candidates_and_persistence(tmp_path): grid = _make_grid(n_side=7) rb = RectBinner( @@ -398,15 +390,6 @@ def test_likelihoods_and_posterior_helper_normalization(): assert ll.shape == (grid.n_models,) assert np.argmax(ll) == 10 - cl = CensoredSizeLikelihood( - phot_indices=[0, 1], - size_index=2, - s_min=float(grid.observables[:, 2].mean()), - bandwidth_floor=1e-3, - ) - ll2 = cl.log_likelihood(x, sig, grid.observables) - assert ll2.shape == (grid.n_models,) - comp = CompositeLikelihood([gp, gp]) llc = comp.log_likelihood(x, sig, grid.observables) np.testing.assert_allclose(llc, 2.0 * ll) @@ -539,3 +522,6 @@ def test_gridfitter_fit_batch_dry_run(): out = fitter.fit_batch(X, S, dry_run=True, return_mode="list") assert out == [] + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/test_kinematics.py b/tests/test_kinematics.py new file mode 100644 index 00000000..f35d6c8c --- /dev/null +++ b/tests/test_kinematics.py @@ -0,0 +1,95 @@ +import unittest + +import numpy as np +from astropy import units as u + +from besta import kinematics + + +class TestKinematics(unittest.TestCase): + def test_gaussian_kernel_normalization_and_centered_percentile(self): + kernel = kinematics.GaussianPixelKernel(velocity_scale=100.0, sigma_truncation=5.0) + kernel.set_parameters(vel=0.0, sigma=200.0) + + self.assertIsNotNone(kernel.kernel_weight) + self.assertTrue(np.isclose(np.sum(kernel.kernel_weight), 1.0, atol=1e-12)) + self.assertTrue(kernel.size % 2 == 1) + + # Percentiles are now returned relative to the central pixel. + self.assertTrue(np.isclose(kernel.get_percentile_pixel(50.0), 0.0, atol=1e-2)) + self.assertTrue(np.isclose(kernel.get_percentile_velocity(50.0), 0.0, atol=1.0)) + + def test_gaussian_delta_kernel_skips_convolution(self): + kernel = kinematics.GaussianPixelKernel(velocity_scale=100.0) + kernel.set_parameters(vel=0.0, sigma=0.0) + + x = np.linspace(0.0, 1.0, 64) + y = kernel.convolve(x) + + self.assertTrue(kernel.skip_convolution) + self.assertTrue(np.allclose(y, x, atol=0.0, rtol=0.0)) + + def test_gausshermite_parse_parameters_from_datablock(self): + block = { + ("kinematics", "los_vel"): 30.0, + ("kinematics", "los_sigma"): 150.0, + ("kinematics", "los_h3"): 0.05, + ("kinematics", "los_h4"): -0.02, + } + + kernel = kinematics.GaussHermitePixelKernel(velocity_scale=100.0) + kernel.parse_parameters(block) + + self.assertIsNotNone(kernel.kernel_weight) + self.assertTrue(np.isclose(np.sum(kernel.kernel_weight), 1.0, atol=1e-10)) + self.assertGreater(kernel.edge_pixels, 0) + + spec = np.sin(np.linspace(0.0, 2 * np.pi, 128)) + conv = kernel.convolve(spec) + self.assertEqual(conv.shape, spec.shape) + self.assertTrue(np.isfinite(conv).all()) + + def test_piecewise_kernel_parse_and_centered_percentile(self): + kernel = kinematics.PieceWisePixelKernel( + velocity_scale=100.0, + velocity_bin_size=100.0, + velocity_min=-200.0, + velocity_max=200.0, + ) + + block = {} + # Symmetric piecewise weights around zero velocity. + block["kinematics", "vel_bin_0"] = 0.1 + block["kinematics", "vel_bin_1"] = 0.4 + block["kinematics", "vel_bin_2"] = 0.4 + block["kinematics", "vel_bin_3"] = 0.1 + + kernel.parse_parameters(block) + + self.assertTrue(np.isclose(np.sum(kernel.kernel_weight), 1.0, atol=1e-10)) + self.assertTrue(np.isclose(kernel.get_percentile_velocity(50.0), 0.0, atol=100.0)) + + def test_convolve_variable_gaussian_kernel_identity_limit(self): + rng = np.random.default_rng(42) + spec = rng.normal(size=(3, 128)) + + # Tiny sigma means all rows are clamped to delta kernels. + sigma_pixel = np.full(spec.shape[-1], 0.01) + out = kinematics.convolve_variable_gaussian_kernel(spec, sigma_pixel) + + self.assertEqual(out.shape, spec.shape) + self.assertTrue(np.allclose(out, spec, atol=1e-12, rtol=0.0)) + + def test_convolve_variable_gaussian_kernel_preserves_units(self): + spec = np.ones((2, 64)) * u.Unit("erg / (s cm2 Angstrom)") + sigma_pixel = np.full(64, 0.8) + + out = kinematics.convolve_variable_gaussian_kernel(spec, sigma_pixel) + + self.assertTrue(hasattr(out, "unit")) + self.assertEqual(out.unit, spec.unit) + self.assertEqual(out.shape, spec.shape) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_likelihoods.py b/tests/test_likelihoods.py new file mode 100644 index 00000000..99b0ae69 --- /dev/null +++ b/tests/test_likelihoods.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +from time import perf_counter + +import numpy as np +import pytest + +from besta.pipeline_modules import likelihoods +from besta.pipeline_modules.base_module import BaseModule + + +class _DummySpectraModule(BaseModule): + name = "DummySpectra" + + def __init__(self, options): + super().__init__(options, likelihood_kind="spectra") + + def make_observable(self, *args, **kwargs): + raise NotImplementedError + + def execute(self, *args, **kwargs): + raise NotImplementedError + + def plot_solution(self, *args, **kwargs): + raise NotImplementedError + + +class _DummyPhotometryModule(BaseModule): + name = "DummyPhotometry" + + def __init__(self, options): + super().__init__(options, likelihood_kind="photometry") + + def make_observable(self, *args, **kwargs): + raise NotImplementedError + + def execute(self, *args, **kwargs): + raise NotImplementedError + + def plot_solution(self, *args, **kwargs): + raise NotImplementedError + + +@pytest.mark.skipif(not likelihoods.NUMBA_AVAILABLE, reason="numba is not available") +def test_likelihood_method_selection_from_init(): + spectra_mod = _DummySpectraModule({"DummySpectra": {"likelihood_method": "numba"}}) + photometry_mod = _DummyPhotometryModule({"DummyPhotometry": {"likelihood_method": "numba"}}) + + assert spectra_mod.likelihood_kind == "spectra" + assert photometry_mod.likelihood_kind == "photometry" + assert spectra_mod.likelihood_method == "numba" + assert photometry_mod.likelihood_method == "numba" + assert spectra_mod.log_like is likelihoods.spectra_loglike_numba + assert photometry_mod.log_like is not None + + +def _benchmark(func, *args, repeats: int = 200): + start = perf_counter() + value = None + for _ in range(repeats): + value = func(*args) + return perf_counter() - start, value + + +@pytest.mark.skipif(not likelihoods.NUMBA_AVAILABLE, reason="numba is not available") +def test_likelihood_backends_match_and_report_performance(capsys): + rng = np.random.default_rng(123) + + # Spectra: inverse-variance Gaussian with no limits. + spectra_data = rng.normal(size=20000) + spectra_model = spectra_data + rng.normal(scale=0.05, size=20000) + spectra_ivar = np.full(20000, 25.0) + + spectra_numpy = likelihoods.spectra_loglike_numpy + spectra_numba = likelihoods.make_spectra_loglike("numba") + spectra_numba(spectra_data, spectra_model, spectra_ivar) # warmup compile + t_np, v_np = _benchmark(spectra_numpy, spectra_data, spectra_model, spectra_ivar) + t_nb, v_nb = _benchmark(spectra_numba, spectra_data, spectra_model, spectra_ivar) + assert np.isclose(v_np, v_nb) + + # Photometry: run both the no-limits fast path and the limits-aware path. + photo_data = rng.normal(size=2000) + photo_model = photo_data + rng.normal(scale=0.05, size=2000) + photo_var = np.full(2000, 0.04) + photo_weights = rng.uniform(0.2, 1.0, size=2000) + no_limits = (None, None) + upper = np.zeros(2000, dtype=bool) + upper[::73] = True + lower = np.zeros(2000, dtype=bool) + lower[::89] = True + lower[::73] = False + + photo_numpy = likelihoods.photometry_loglike_numpy + photo_numba = likelihoods.make_photometry_loglike("numba") + photo_numba(photo_data, photo_model, photo_var, photo_weights, *no_limits) # warmup compile + t_pn, v_pn = _benchmark(photo_numpy, photo_data, photo_model, photo_var, photo_weights, *no_limits) + t_pj, v_pj = _benchmark(photo_numba, photo_data, photo_model, photo_var, photo_weights, *no_limits) + assert np.isclose(v_pn, v_pj) + + t_pl, v_pl = _benchmark( + photo_numpy, + photo_data, + photo_model, + photo_var, + photo_weights, + upper, + lower, + ) + t_pjl, v_pjl = _benchmark( + photo_numba, + photo_data, + photo_model, + photo_var, + photo_weights, + upper, + lower, + ) + assert np.isclose(v_pl, v_pjl) + + print( + "Likelihood timing report (seconds over repeated calls):\n" + f" spectra numpy={t_np:.6f} numba={t_nb:.6f} speedup={t_np / t_nb:.2f}x\n" + f" photo(no limits) numpy={t_pn:.6f} numba={t_pj:.6f} speedup={t_pn / t_pj:.2f}x\n" + f" photo(limits) numpy={t_pl:.6f} numba={t_pjl:.6f} speedup={t_pl / t_pjl:.2f}x" + ) + + captured = capsys.readouterr() + assert "Likelihood timing report" in captured.out diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index dd23f90d..b41ed78d 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -1,12 +1,8 @@ -import os -import numpy as np import pytest -from astropy.table import Table - -from cosmosis import DataBlock from besta import io -from besta.pipeline import MainPipeline +from besta.pipeline import MainPipeline, BatchPipeline +import besta.pipeline as pipeline_module def test_pipeline_execute_all_stops_on_failure(tmp_path, monkeypatch): @@ -26,6 +22,109 @@ def run_command(self, command): monkeypatch.setattr(io, "make_values_file", lambda *args, **kwargs: None) assert pipe.execute_all() == 1 + +def test_batch_pipeline_validate_input_lengths(): + configs = [[{"pipeline": {"modules": "Dummy"}, "Dummy": {}, "output": {"filename": "x"}}]] + with pytest.raises(ValueError): + BatchPipeline( + pipeline_configuration_list=configs, + n_cores_list=[1, 2], + ) + + +def test_batch_from_running_parameters_deepcopy_isolated_configs(): + base_config = { + "pipeline": {"modules": "Dummy", "values": "values.ini"}, + "output": {"filename": "base"}, + "Dummy": { + "file": "dummy.py", + "nested": {"alpha": 1, "beta": 2}, + }, + } + running_parameters = [ + {"Dummy": {"nested": {"alpha": 10}}}, + {"Dummy": {"nested": {"alpha": 20}}}, + ] + + batch = BatchPipeline.from_running_parameters( + pipeline_configuration=base_config, + running_parameters=running_parameters, + ) + + cfg_a = batch.all_pipelines_config[0][0] + cfg_b = batch.all_pipelines_config[1][0] + + assert cfg_a["Dummy"]["nested"]["alpha"] == 10 + assert cfg_b["Dummy"]["nested"]["alpha"] == 20 + assert cfg_a["Dummy"]["nested"] is not cfg_b["Dummy"]["nested"] + assert base_config["Dummy"]["nested"]["alpha"] == 1 + + +def test_batch_from_running_parameters_rejects_invalid_inputs(): + base_config = { + "pipeline": {"modules": "Dummy", "values": "values.ini"}, + "output": {"filename": "base"}, + "Dummy": {"file": "dummy.py"}, + } + + with pytest.raises(TypeError): + BatchPipeline.from_running_parameters( + pipeline_configuration=base_config, + running_parameters={"Dummy": {"x": 1}}, + ) + + with pytest.raises(TypeError): + BatchPipeline.from_running_parameters( + pipeline_configuration=base_config, + running_parameters=[{"Dummy": {"x": 1}}, "bad-entry"], + ) + + +def test_batch_from_running_parameters_calls_keyword_only_constructor(monkeypatch): + base_config = { + "pipeline": {"modules": "Dummy", "values": "values.ini"}, + "output": {"filename": "base"}, + "Dummy": {"file": "dummy.py"}, + } + + captured = {} + + def fake_init(self, *, pipeline_configuration_list, **kwargs): + captured["pipeline_configuration_list"] = pipeline_configuration_list + captured["kwargs"] = kwargs + + monkeypatch.setattr(BatchPipeline, "__init__", fake_init) + + BatchPipeline.from_running_parameters( + pipeline_configuration=base_config, + running_parameters=[{"Dummy": {"x": 1}}], + n_jobs_parallel=2, + ) + + assert "pipeline_configuration_list" in captured + assert len(captured["pipeline_configuration_list"]) == 1 + assert captured["kwargs"]["n_jobs_parallel"] == 2 + + +def test_batch_run_all_pipelines_sequential_order(monkeypatch): + def fake_worker(job): + return {"index": job["index"], "status": 100 + job["index"]} + + monkeypatch.setattr(pipeline_module, "_run_main_pipeline_job", fake_worker) + + configs = [ + [{"pipeline": {"modules": "Dummy"}, "Dummy": {}, "output": {"filename": "a"}}], + [{"pipeline": {"modules": "Dummy"}, "Dummy": {}, "output": {"filename": "b"}}], + [{"pipeline": {"modules": "Dummy"}, "Dummy": {}, "output": {"filename": "c"}}], + ] + + batch = BatchPipeline( + pipeline_configuration_list=configs, + n_jobs_parallel=1, + ) + + assert batch.run_all_pipelines() == [100, 101, 102] + if __name__ == "__main__": from besta.logging import setup_logging setup_logging() diff --git a/tests/test_pipeline_module.py b/tests/test_pipeline_module.py index 36dc7c32..6e66f62b 100644 --- a/tests/test_pipeline_module.py +++ b/tests/test_pipeline_module.py @@ -3,6 +3,7 @@ import os import numpy as np +from requests.exceptions import ConnectTimeout from cosmosis import DataBlock from besta.pipeline_modules import FullSpectralFitModule, GalaxySpectraModule, GalaxyPhotometryModule @@ -59,7 +60,7 @@ def test_full_spectral_fit(self): }} block = DataBlock() - block['dust.extinction', 'a_v'] = 0 + block['dust_attenuation', 'a_v'] = 0 block['kinematics', 'los_vel'] = 0 block['kinematics', 'los_sigma'] = 100. block['kinematics', 'los_h3'] = 0 @@ -91,7 +92,7 @@ def test_galaxy_spectra(self): }} block = DataBlock() - block['dust.extinction', 'a_v'] = 0 + block['dust_attenuation', 'a_v'] = 0 block['kinematics', 'los_vel'] = 0 block['kinematics', 'los_sigma'] = 100. block['kinematics', 'los_h3'] = 0 @@ -123,7 +124,7 @@ def test_galaxy_photometry(self): }} block = DataBlock() - block['dust.extinction', 'a_v'] = 0 + block['dust_attenuation', 'a_v'] = 0 block['kinematics', 'los_vel'] = 0 block['kinematics', 'los_sigma'] = 100. block['kinematics', 'los_h3'] = 0 @@ -132,11 +133,15 @@ def test_galaxy_photometry(self): block['stars.sfh', 'alpha_powerlaw'] = 1 block['stars.sfh', 'ism_metallicity_today'] = 0.02 - module = module(config) + try: + module = module(config) + except ConnectTimeout as exc: + print("SVO filter query probably failed: skipping") + return self.assertFalse(module.execute(block)) flux_model = module.make_observable(block, parse=True) self.assertTrue(np.isfinite(flux_model).all()) print("Module successfully executed") if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py index b7fcbcbe..17dacfd8 100644 --- a/tests/test_postprocessing.py +++ b/tests/test_postprocessing.py @@ -7,11 +7,10 @@ from astropy.table import Table from besta.postprocess import ( - read_results_file, summarize_results, ResultsSummary, ) - +from besta import io class TestPostprocessing(unittest.TestCase): """ @@ -85,23 +84,9 @@ def _get_results_path(self) -> str: return self.data_file return self._make_synthetic_results_file() - def test_read_results_file(self): - path = self._get_results_path() - table = read_results_file(path) - - self.assertIsInstance(table, Table) - self.assertGreater(len(table), 0, "Results table is empty") - - # Basic column expectations - self.assertIn("post", table.colnames, "Missing 'post' column") - # The synthetic file has prior; a real sfh.txt might not. - # So only assert prior if present in the file. - # Parameter columns: at least one 'section--param' should exist - self.assertTrue(any("--" in c for c in table.colnames), "No parameter columns found (expected '--' delimiter).") - def test_summarize_results_and_exports(self): path = self._get_results_path() - table = read_results_file(path) + table = io.read_results_file(path) with tempfile.TemporaryDirectory() as tmp: out_fits = os.path.join(tmp, "besta_summary.fits") diff --git a/tests/test_run.py b/tests/test_run.py index cc7c515e..27bcd56d 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -33,7 +33,7 @@ def setUpClass(cls): [ssp.wavelength, np.random.normal(sed, sed * 0.01), sed * 0.01]).T) # Create values file - text = """[dust.extinction] + text = """[dust_attenuation] a_v = 0 0 1 [stars.sfh] alpha_powerlaw = 0 1 10 @@ -111,6 +111,7 @@ def test_fit(self): "SFHModel": "ExponentialSFH", "velscale": 50.0, "ExtinctionLaw": "ccm89", + "use_features": "T", }} t0 = time() @@ -140,4 +141,4 @@ def test_postprocess(self): results = summarize_results(results.results_table) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/wip_test_pymc_sampler.py b/tests/wip_test_pymc_sampler.py new file mode 100644 index 00000000..9a0b04f4 --- /dev/null +++ b/tests/wip_test_pymc_sampler.py @@ -0,0 +1,250 @@ +from __future__ import annotations + +import shutil +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +from cosmosis.output import InMemoryOutput + +from besta import io +from besta.pipeline import MainPipeline +from besta.samplers.pymc_sampler import PymcSampler + + +def _write_ini(path: Path, text: str) -> Path: + path.write_text(text, encoding="utf-8") + return path + + +def test_pymc_sampler_smoke_with_dummy_pipeline(tmp_path): + class DummyPipeline: + def __init__(self): + self.varied_params = [object(), object()] + self.modules = [] + + def output_names(self): + # Varied params + one extra output. + return ["p0", "p1", "extra0"] + + def denormalize_vector(self, u): + # Map unit-cube to a bounded physical space. + return 4.0 * np.asarray(u, dtype=float) - 2.0 + + def run_results(self, p): + p = np.asarray(p, dtype=float) + post = -0.5 * float(np.dot(p, p)) + prior = -0.1 * float(np.dot(p, p)) + like = post - prior + extra = np.array([float(p.sum())], dtype=float) + return SimpleNamespace(post=post, prior=prior, like=like, extra=extra) + + ini_path = _write_ini( + tmp_path / "pymc_smoke.ini", + """ +[runtime] +sampler = pymc + +[pymc] +samples = 24 +nsteps = 8 +burn_fraction = 0.0 +chains = 1 +progressbar = F +seed = 11 +""".strip(), + ) + + output = InMemoryOutput() + sampler = PymcSampler(str(ini_path), DummyPipeline(), output) + sampler.config() + sampler.execute() + + assert sampler.is_converged() + assert sampler.num_samples > 0 + assert len(output.rows) == sampler.num_samples + + col_names = [c[0] for c in output.columns] + assert "prior" in col_names + assert "post" in col_names + assert "like" in col_names + + +def test_pymc_sampler_accepts_slice_step_method(tmp_path): + class DummyPipeline: + def __init__(self): + self.varied_params = [object(), object()] + self.modules = [] + + def output_names(self): + return ["p0", "p1", "extra0"] + + def denormalize_vector(self, u): + return 4.0 * np.asarray(u, dtype=float) - 2.0 + + def run_results(self, p): + p = np.asarray(p, dtype=float) + post = -0.5 * float(np.dot(p, p)) + prior = -0.1 * float(np.dot(p, p)) + like = post - prior + extra = np.array([float(p.sum())], dtype=float) + return SimpleNamespace(post=post, prior=prior, like=like, extra=extra) + + ini_path = _write_ini( + tmp_path / "pymc_slice.ini", + """ +[runtime] +sampler = pymc + +[pymc] +samples = 12 +nsteps = 4 +burn_fraction = 0.0 +chains = 1 +progressbar = F +seed = 5 +step_method = slice +""".strip(), + ) + + output = InMemoryOutput() + sampler = PymcSampler(str(ini_path), DummyPipeline(), output) + sampler.config() + sampler.execute() + + assert sampler.is_converged() + assert sampler.num_samples > 0 + + +def test_pymc_sampler_rejects_gradient_step_methods(tmp_path): + class DummyPipeline: + def __init__(self): + self.varied_params = [object(), object()] + self.modules = [] + + def output_names(self): + return ["p0", "p1", "extra0"] + + def denormalize_vector(self, u): + return 4.0 * np.asarray(u, dtype=float) - 2.0 + + def run_results(self, p): + p = np.asarray(p, dtype=float) + post = -0.5 * float(np.dot(p, p)) + prior = -0.1 * float(np.dot(p, p)) + like = post - prior + extra = np.array([float(p.sum())], dtype=float) + return SimpleNamespace(post=post, prior=prior, like=like, extra=extra) + + ini_path = _write_ini( + tmp_path / "pymc_nuts.ini", + """ +[runtime] +sampler = pymc + +[pymc] +samples = 6 +nsteps = 3 +burn_fraction = 0.0 +chains = 1 +progressbar = F +seed = 1 +step_method = nuts +""".strip(), + ) + + output = InMemoryOutput() + sampler = PymcSampler(str(ini_path), DummyPipeline(), output) + sampler.config() + with pytest.raises(ValueError, match="requires gradients"): + sampler.execute() + + +def test_pymc_sampler_full_besta_run(tmp_path): + cosmosis_exe = shutil.which("cosmosis") + if cosmosis_exe is None: + pytest.skip("cosmosis executable not available in PATH") + + module_path = tmp_path / "dummy_like_module.py" + module_path.write_text( + """ +def setup(options): + return {} + + +def execute(block, config): + p1 = block['parameters', 'p1'] + p2 = block['parameters', 'p2'] + like = -0.5 * (p1**2 + p2**2) + block['extra', 'sum_p'] = p1 + p2 + block['likelihoods', 'dummy_like'] = like + return 0 + + +def cleanup(config): + return 0 +""".strip(), + encoding="utf-8", + ) + + output_root = tmp_path / "pymc_full_run" + values_path = tmp_path / "values.ini" + values_path.write_text( + """ +[parameters] +p1 = -2.0 0.0 2.0 +p2 = -2.0 0.0 2.0 +""".strip(), + encoding="utf-8", + ) + + sampler_path = Path(__file__).resolve().parents[1] / "src" / "besta" / "samplers" / "pymc_sampler.py" + + config = { + "runtime": { + "sampler": "pymc", + "import_samplers": str(sampler_path), + }, + "pymc": { + "samples": 16, + "nsteps": 8, + "burn_fraction": 0.0, + "chains": 1, + "progressbar": False, + "seed": 3, + }, + "output": { + "filename": str(output_root), + "format": "text", + }, + "pipeline": { + "modules": "DummyLike", + "values": str(values_path), + "likelihoods": "dummy", + "quiet": "T", + "debug": "T", + "extra_output": "extra/sum_p", + }, + "DummyLike": { + "file": str(module_path), + }, + } + + runner = MainPipeline( + [config], + n_cores_list=[1], + ini_values_files=[str(values_path)], + ) + status = runner.execute_all(plot_result=False) + assert status == 0 + + results_file = Path(str(output_root) + ".txt") + assert results_file.exists() + + table = io.read_results_file(str(results_file)) + assert len(table) > 0 + assert "post" in table.colnames + assert "prior" in table.colnames + assert "like" in table.colnames diff --git a/tutorials/fit_redshift/fit_jwst_redshift.ipynb b/tutorials/fit_redshift/fit_jwst_redshift.ipynb index 30f23dcf..623ebf60 100644 --- a/tutorials/fit_redshift/fit_jwst_redshift.ipynb +++ b/tutorials/fit_redshift/fit_jwst_redshift.ipynb @@ -153,7 +153,9 @@ "\n", "If you want to infer the redshift from more than one spectrum, it is highly recommended to first pre-build the SSP stellar template to the desired resolution.\n", "\n", - "In this case, we will resample the Bruzual & Charlote 2003 (version updated in 2016) SSP models to a resolution of constant velocity spanning the entire range available. The SSP grid can be saved as a pickle file, that can be loaded on each run, skipping the interpolation step." + "In this case, we will resample the Bruzual & Charlote 2003 (version updated in 2016) SSP models to a resolution of constant velocity spanning the entire range available. The SSP grid can be saved as a pickle file, that can be loaded on each run, skipping the interpolation step.\n", + "\n", + "**Security disclaimer**: Python pickle files can execute arbitrary code when loaded. Only load .pkl files that you generated yourself or received from a trusted source. Never load pickle files from untrusted or unverified locations." ] }, { @@ -194,7 +196,7 @@ "- The search interval is controlled by `z_min` and `z_max`.\n", "- Keep the observed spectrum in the observed frame; the module shifts the model internally during the scan.\n", "\n", - "In this tutorial we also enable `use_features=\"T\"`, which estimates a smooth continuum and feature-sensitive weights. This usually improves robustness when broad-band continuum shape mismatches dominate over line information." + "In this tutorial we also enable `use_features=\"T\"`, which estimates a smooth continuum and feature-sensitive weights. This usually improves robustness when broad-band continuum shape mismatches dominate over line information. In poor SNR conditions, it is recommended to disable it." ] }, { @@ -215,8 +217,6 @@ "metadata": {}, "outputs": [], "source": [ - "# -- Configuration -------------------------------------------------------------\n", - "\n", "output_root = tutorial_dir / \"results.jwst_redshift_maxlike.txt\"\n", "values_path = tutorial_dir / \"tutorial_values.ini\"\n", "\n", @@ -309,7 +309,7 @@ " f.write(\"quenching_time = 13.6 13.8 13.8\\n\")\n", " f.write(\"alpha_powerlaw = 1\\n\")\n", " f.write(\"ism_metallicity_today = 0.005 0.02 0.05\\n\")\n", - " f.write(\"[dust.extinction]\\n\")\n", + " f.write(\"[dust_attenuation]\\n\")\n", " f.write(\"a_v = 0.0 0.1 1.0\\n\")\n" ] }, @@ -521,7 +521,7 @@ ], "metadata": { "kernelspec": { - "display_name": "cosmo-env", + "display_name": "besta", "language": "python", "name": "python3" }, @@ -535,7 +535,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.10" + "version": "3.14.4" } }, "nbformat": 4, diff --git a/tutorials/fit_sdss_spectra/sdss_spectra.ipynb b/tutorials/fit_sdss_spectra/sdss_spectra.ipynb index 8c9f34d0..929df920 100644 --- a/tutorials/fit_sdss_spectra/sdss_spectra.ipynb +++ b/tutorials/fit_sdss_spectra/sdss_spectra.ipynb @@ -54,8 +54,8 @@ "from astropy.io import fits\n", "\n", "hdul = fits.open(\n", - " \"spec-2240-53823-0178.fits\" # Quiescent galaxy\n", - " # \"spec-1058-52520-0595.fits\" # Star-forming galaxy\n", + " # \"spec-2240-53823-0178.fits\" # Quiescent galaxy\n", + " \"spec-1058-52520-0595.fits\" # Star-forming galaxy\n", " # \"spec-1678-53433-0425.fits\" # Starburst galaxy\n", " )" ] @@ -191,9 +191,11 @@ "metadata": {}, "outputs": [], "source": [ - "sdss_lsf_file = \"lsf_SDSS.dat\"\n", + "lsf_pixels = hdul[1].data[\"wdisp\"]\n", + "sdss_lsf_wl = lsf_pixels * np.gradient(wl.value)\n", + "sdss_lsf_fwhm = wl\n", + "# convert to FWHM in Angstrom at each wavelength accounting for each pixels varying dispersion\n", "emiles_lsf_file = \"e-miles_spectral_resolution.dat\"\n", - "sdss_lsf_wl, sdss_lsf_fwhm = np.loadtxt(sdss_lsf_file, unpack=True)\n", "emiles_lsf_wl, emiles_lsf_fwhm = np.loadtxt(emiles_lsf_file, unpack=True, usecols=(0, 1))\n", "\n", "wl_min = 3850\n", @@ -404,7 +406,7 @@ "metadata": {}, "outputs": [], "source": [ - "new_sections = \"\"\"[dust.extinction]\n", + "new_sections = \"\"\"[dust.attenuation]\n", "a_v = 0 0.3 2.0\n", "\"\"\"\n", "\n", @@ -547,28 +549,20 @@ "\n", " # Sampling strategy\n", "\n", - " \"runtime\": {\"sampler\": \"gridmax emcee\"},\n", + " \"runtime\": {\"sampler\": \"maxlike emcee\"},\n", "\n", " # Optimizer based on scipy\n", " \"maxlike\": {\n", " \"method\": \"Nelder-Mead\",\n", - " \"tolerance\": 1e-3,\n", + " \"tolerance\": 1e-6,\n", " \"maxiter\": 5000,\n", " \"repeats\": 5,\n", " \"start_method\": \"prior\",\n", " },\n", "\n", - " # Optimizer that accepts parallel processing\n", - " \"gridmax\": {\"nsteps\": 10, \"tolerance\": 0.1, \"output_ini\": \"gridmax_output.txt\",\n", - " \"max_iterations\": 100},\n", - "\n", " # MCMC sampler configuration\n", " \"emcee\": {\"walkers\": 64, \"samples\": 100, \"nsteps\": 100},\n", - "\n", - "\n", - " \"polychord\": {\"max_iterations\": 1000, \"live_points\": 200, \"feedback\": 1,\n", - " \"polychord_outfile_root\": \"./polychord_output\"}, \n", - " \n", + " \n", " # Output configuration\n", " \"output\": {\"filename\": \"results_sdss_tutorial.txt\", \"format\": \"text\"},\n", "\n", @@ -587,6 +581,7 @@ " \"FullSpectralFit\": {\n", " \"file\": FullSpectralFitModule.get_path(),\n", " \"logging_console\": \"T\",\n", + " # \"logging_level\": \"DEBUG\",\n", " \"inputSpectrum\": spectra_file,\n", " \"fluxUnits\": \"'erg/ (s cm^2 Angstrom)'\",\n", " \"wlRange\": [wl_min, wl_max],\n", @@ -765,12 +760,26 @@ "\n", "summary = summarize_results(reader.results_table, compute_2d=True,\n", " parameter_key_pairs=[(\"kinematics--los_vel\", \"kinematics--los_sigma\"),\n", - " (\"dust.extinction--a_v\", \"stars.sfh--t_at_frac_0.9900\")])\n", + " (\"dust.attenuation--a_v\", \"stars.sfh--t_at_frac_0.9900\")])\n", "\n", "summary.corner_plot(outpath=\"./corner_plot.png\")\n", "summary.plot_2d_pdfs(outdir=\"./2d_pdfs\")" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "b19408ad", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Percentiles: \", summary.percentiles)\n", + "\n", + "for param, pct in zip(summary.parameter_keys, summary.percentiles_values):\n", + " print(f\"{param}: {pct}\")\n", + " " + ] + }, { "cell_type": "code", "execution_count": null, @@ -779,7 +788,7 @@ "outputs": [], "source": [ "# disable logging for the module to avoid cluttering the output\n", - "# reader.ini[\"FullSpectralFit\"][\"logging_level\"] = \"ERROR\"\n", + "reader.ini[\"FullSpectralFit\"][\"logging_level\"] = \"ERROR\"\n", "\n", "# load the module to access its methods for post-processing\n", "module = reader.get_module(\"FullSpectralFit\")" @@ -969,7 +978,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "besta", "language": "python", "name": "python3" }, @@ -983,7 +992,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.10" + "version": "3.14.4" } }, "nbformat": 4,