A Python package for PSO-based parameter optimization in galaxy evolution modeling using the SAGE semi-analytic model.
- Features
- Requirements
- Installation
- Package Structure
- Usage
- Search Space Configuration
- Constraints System
- Simulation Support
- Data Files and Formats
- Diagnostics and Output
- Testing
- License
- Particle Swarm Optimization (PSO) for SAGE parameter calibration
- Multiple constraint types: SMF, BHMF, BHBM, CSFRDH, HIMF, H2MF, MZR, SHMR, SMD, FICS, FICS_Mvir, MLF
- Red/blue galaxy stellar mass function discrimination
- Multi-simulation support: miniUchuu, miniMillennium, MTNG
- Automatic CSV data generation from SAGE HDF5 output
- HPC/SLURM integration for parallel execution
- Diagnostic plots, animations, and parameter uncertainty analysis
- Python 3.8+
- SAGE binary (user-provided)
- SAGE input
.parfile (user-provided) - Merger tree age list file (simulation-specific)
Python dependencies:
- numpy
- pandas
- matplotlib
- scipy
- h5py
Install with:
pip install -r requirements.txtgit clone https://github.com/yourusername/sage-pso.git
cd sage-pso
pip install -r requirements.txtSAGE-PSO/
├── main.py # Entry point: SAGE execution, CSV generation, PSO orchestration
├── space.txt # Parameter search space specification
├── requirements.txt # Python dependencies
├── run_pso.sh # Single PSO run script
│
├── src/
│ ├── pso.py # PSO algorithm implementation
│ ├── constraints.py # Constraint definitions and data loading
│ ├── analysis.py # Statistical tests (chi-squared, Student's t)
│ ├── execution.py # SAGE binary execution, SLURM job submission
│ ├── diagnostics.py # Post-PSO plots and animations
│ ├── simulation_config.py # Simulation parameters and snapshot mappings
│ ├── routines.py # HDF5 reading and data extraction
│ ├── common.py # Utility functions
│ ├── redshift_utils.py # Redshift/snapshot conversion
│ └── pso_uncertainty.py # Parameter uncertainty analysis
│
├── run_types/
│ ├── run_multiple_pso.sh # Sequential multiple PSO runs
│ ├── run_multiple_pso_slurm.sh # SLURM parallel PSO runs
│ ├── submit_pso_array.sh # SLURM array job submission
│ ├── analyze_pso_array.sh # Analyze array job results
│ ├── analyze_multiple_pso.py # Multi-run analysis script
│ └── make_comparison_plot.sh # Generate comparison plots
│
├── tests/
│ ├── test_constraint_data.py # Constraint data loading tests
│ ├── test_pso_benchmarks.py # PSO algorithm validation
│ ├── quick_pso_test.py # Minimal sanity check
│ └── visual_pso_test.py # Visual convergence test
│
└── data/
└── (constraint observational data files)
- CSV Data Check: At startup, required
sage_*.csvfiles are checked in the output directory - Automatic Generation: Missing CSVs are regenerated from SAGE HDF5 output
- PSO Execution: The optimizer runs using constraints and configuration
- Diagnostics: Plots and statistics are generated after PSO completes
| Argument | Description |
|---|---|
-c, --config |
Path to SAGE input .par file |
-b, --sage-binary |
Path to SAGE binary |
| Argument | Default | Description |
|---|---|---|
-o, --outdir |
. |
Output directory |
-v, --subvolumes |
0 |
Subvolumes to process |
-k, --keep |
off | Keep temporary output files |
-sn, --snapshot |
auto | Snapshot numbers to analyze |
| Argument | Default | Description |
|---|---|---|
--sim |
0 |
Simulation type: 0=miniUchuu, 1=miniMillennium, 2=MTNG |
--boxsize |
sim-specific | Simulation box size in Mpc/h |
--vol-frac |
1.0 |
Volume fraction of simulation box |
--age-alist-file |
sim-specific | Path to merger tree age list file |
--Omega0 |
sim-specific | Matter density parameter |
--h0 |
sim-specific | Hubble parameter (H0/100) |
| Argument | Default | Description |
|---|---|---|
-s, --swarm-size |
10 + 2*sqrt(D) |
Number of particles |
-m, --max-iterations |
20 |
Maximum iterations |
-S, --space-file |
space.txt |
Search space specification |
Box size, cosmology (h0, Omega0), the processed volume fraction and the
scale-factor list are read automatically — the first three and the volume
fraction from the SAGE output's Header/Simulation and Header/Runtime, the
scale-factor list from the .par file's FileWithSnapList (falling back to the
output header). --boxsize, --h0, --Omega0, --vol-frac and
--age-alist-file remain as overrides, but a value that contradicts the SAGE
output is refused rather than applied: a wrong box size rescales every
volume-dependent constraint by (ratio)³, silently.
--sim still has to be set. It selects the snapshot-to-redshift map, which the
output does not record: 0 for the 50-snapshot Uchuu grid (miniUchuu and
microUchuu both), 1 for the 64-snapshot Millennium grid, 2 for MTNG.
| -t, --stat-test | student-t | Statistical test (student-t, chi2) |
| -x, --constraints | BHMF,SMF_z0,BHBM | Constraints to use |
| -csv, --csv-output | none | Save results to CSV |
| -r, --random-seed | random | Seed for reproducibility |
| --omega | 0.729 | PSO inertia weight |
| --phip | 1.49445 | Cognitive parameter |
| --phig | 1.49445 | Social parameter |
| Argument | Default | Description |
|---|---|---|
-H, --hpc-mode |
off | Enable HPC mode |
-C, --cpus |
1 |
CPUs per SAGE instance |
-M, --memory |
1500m |
Memory per instance |
-N, --nodes |
auto | Number of nodes |
-a, --account |
none | SLURM account |
-q, --queue |
none | SLURM queue |
-w, --walltime |
1:00:00 |
Walltime per job |
-u, --username |
none | SLURM username |
python main.py \
-b ./sage \
-c ./input/millennium.par \
-o ./output \
--sim 1 \
-x "SMF_z0(8-11)*5,BHMF_z0,BHBM" \
-s 20 \
-m 30 \
-csv results.csvThe space.txt file defines the parameter search space:
SfrEfficiency,eSFR,1,0.01,0.1
FeedbackReheatingEpsilon,eReheat,0,0.0,6.0
FeedbackEjectionEfficiency,eEject,0,0.1,1.0
ReIncorporationFactor,eReinc,0,0.05,0.3
RadioModeEfficiency,eRadio,1,0.001,1.0
QuasarModeEfficiency,eQuasar,1,0.001,0.5
BlackHoleGrowthRate,eBHgrowth,1,0.0001,0.5
Format: ParameterName,Label,Sampling,LowerBound,UpperBound
Sampling:0= linear,1= logarithmic,2= integer switch
| Test | Behaviour |
|---|---|
chi2 |
Gaussian. Most sensitive when the errors are trustworthy; an outlier costs its distance squared. |
student-t |
Heavy-tailed, degrees of freedom estimated from the residuals. Tends to chi2 as the fit improves. |
huber |
Quadratic within 2σ, linear beyond. Same as chi2 for good points, no dof to estimate. |
cauchy |
Heaviest tail (Student-t with ν=1). For when some observations are simply wrong. |
abs |
Sum of |residual|/σ. The total reads as "typically N sigma out". |
cash |
Poisson/Cash (1979) C statistic on galaxy counts, for the mass functions. Falls back to chi2 for constraints that are not histograms. |
cash treats the model as the Poisson realisation and the observation as the
underlying rate — the right way round, since a survey covers far more volume
than a mini-box, so the simulation is the noisy side. It needs the constraint's
bin_width to recover counts from log₁₀φ, which the mass functions set and the
relations (BHBM, MZR, SHMR, FICS, MLF) do not.
When to use cash: it is the only test that handles sparse and empty bins
correctly. At the massive end a bin holds 1–4 galaxies, where the 68% Poisson
interval on log₁₀φ is wildly asymmetric ([-3.0, +0.3] dex for N=1) and no
symmetric error bar represents it. But it weights every bin by its object
count, so the well-populated low-mass end (10⁴ galaxies, φ known to 1%)
dominates — and there the real uncertainty is systematic (completeness, IMF,
apertures), not shot noise. Note the model's Poisson error is already folded
into chi2 via hist_smf_err; what chi2 gets wrong is only the asymmetry at
N < 5 and the empty-bin sentinel. Prefer cash when the massive end is what
you care about, and do not mix it with non-counting constraints in one run
without rebalancing weights — the mass functions will otherwise outweigh them
by a factor of ~100.
All are reduced by the point count per constraint, then combined by relative
weight, so they are comparable within a fixed -x string but not across
different ones.
SAGE reads some parameters as integers — mode selectors such as
DynamicDisruptionSplit, ConcentrationOn, FeedbackFreeModeOn,
SFprescription. Its parameter reader parses these with strtol and aborts
the whole run if anything is left over, so a continuous PSO value like 1.37
kills the optimisation rather than just scoring badly.
Sampling code 2 makes such a parameter safe to optimise: the particle position
stays continuous, but the value is rounded to the nearest level and clamped to
the declared bounds before it is written to the parameter file, and the tracks,
CSV and best-fit report the level that SAGE actually ran.
DynamicDisruptionSplit,split,2,0,2
ConcentrationOn,cOn,2,0,3
SfrEfficiency,aSF,1,0.005,0.2
Bounds for a switch must be whole numbers, and must lie inside the range SAGE
accepts for that parameter (SAGE validates switches and aborts on an
out-of-range value). load_space warns if a known SAGE integer parameter is
declared with code 0/1, or if a continuous parameter is declared with 2.
Two things to keep in mind when doing this:
- The objective is piecewise constant along a switch axis, so PSO gets no
gradient information there and particles that round to the same level run
identical SAGE configurations. For a switch with only two or three levels,
running one campaign per level is more informative — you get a fit for each
mechanism to compare, rather than one winner and no information about the
rest. Switches with many levels (
FeedbackFreeModeOnhas eight) are where co-fitting earns its place. - Sampling is over
[lb-0.5, ub+0.5]so every level gets equal probability; sampling[lb, ub]and rounding would give the two end levels half the width of the interior ones.
| Constraint | Description |
|---|---|
SMF_z0, SMF_z05, SMF_z10, SMF_z20, SMF_z30, SMF_z40 |
Stellar Mass Function at z=0, 0.5, 1, 2, 3, 4 |
SMF_Red_z0, SMF_Blue_z0 |
Red (quiescent) and blue (star-forming) galaxy SMF at z=0 |
BHMF_z0, BHMF_z10 |
Black Hole Mass Function |
BHBM |
Black Hole - Bulge Mass relation |
CSFRDH |
Cosmic Star Formation Rate Density History |
HIMF |
HI Mass Function |
H2MF |
H2 Mass Function |
MZR |
Mass-Metallicity Relation |
SHMR |
Stellar-Halo Mass Relation |
SMD |
Stellar Mass Density history |
FICS |
Intracluster-star mass fraction, f_ICS = m_ICS / M_*,halo, from z=0 to z=2 |
FICS_Mvir |
Intracluster-star mass fraction vs host halo mass at z=0 (Contini 2021) |
MLF |
Galactic-wind mass-loading factor vs circular velocity at z=0 (requires FIREmodeOn=1) |
-x "SMF_z0(8-11)*5,BHMF_z0*10,BHBM"- Domain restriction:
SMF_z0(8-11)limits to log(M/M☉) = 8-11 - Weighting:
BHMF_z0*10applies weight of 10
- 50 snapshots (0-49), snapshot 49 ≈ z=0
- Box size: 400 Mpc/h
- Cosmology: h=0.6774, Ω₀=0.3089
- 64 snapshots (0-63), snapshot 63 = z=0
- Box size: 62.5 Mpc/h
- Cosmology: h=0.73, Ω₀=0.25
- 100 snapshots (0-99), snapshot 99 = z=0
- Box size: 500 Mpc/h
- Cosmology: h=0.6774, Ω₀=0.3089
Each simulation has its own snapshot-to-redshift mapping defined in src/simulation_config.py.
SAGE output is automatically converted to CSV files:
| File | Contents |
|---|---|
sage_smf_all_redshifts.csv |
Stellar Mass Function |
sage_smf_red_all_redshifts.csv |
Red galaxy SMF |
sage_smf_blue_all_redshifts.csv |
Blue galaxy SMF |
sage_bhmf_all_redshifts.csv |
Black Hole Mass Function |
sage_bhbm_all_redshifts.csv |
BHBM relation (median, std, counts) |
sage_halostellar_all_redshifts.csv |
Halo-Stellar mass relation |
sage_himf_all_redshifts.csv |
HI Mass Function |
sage_h2mf_all_redshifts.csv |
H2 Mass Function |
sage_mzr_all_redshifts.csv |
Mass-Metallicity Relation |
sage_history.csv |
Cosmic history (CSFRDH, SMD, FICS) |
Files are tab-separated with no headers.
After PSO completion:
sage_pso.log- Run logtracks/track_*_pos.npy,tracks/track_*_fx.npy- Particle trajectories- Parameter evolution plots
- Likelihood curves
- Swarm movement visualizations
- Pairplots and KDE distributions
- Constraint comparison grids
- GIF animations of swarm evolution (optional)
Run tests from the project root:
# Constraint data loading
python tests/test_constraint_data.py
# PSO algorithm benchmarks
python tests/test_pso_benchmarks.py --test all
# Quick sanity check
python tests/quick_pso_test.py
# Visual convergence test
python tests/visual_pso_test.pyMIT