SSP interpolation#
Load SSP Grid#
# NBVAL_SKIP
from rubix.spectra.ssp.templates import BruzualCharlot2003
print(BruzualCharlot2003)
2025-11-10 17:13:20,829 - rubix - INFO -
___ __ _____ _____ __
/ _ \/ / / / _ )/ _/ |/_/
/ , _/ /_/ / _ |/ /_> <
/_/|_|\____/____/___/_/|_|
2025-11-10 17:13:20,831 - rubix - INFO - Rubix version: 0.0.post626+g42b4b7505.d20251110
2025-11-10 17:13:20,831 - rubix - INFO - JAX version: 0.7.2
2025-11-10 17:13:20,909 - rubix - INFO - Running on [CpuDevice(id=0)] devices
2025-11-10 17:13:20,911 - rubix - WARNING - python-fsps is not installed. Please install it to use this function. Install using pip install fsps and check the installation page: https://dfm.io/python-fsps/current/installation/ for more details. Especially, make sure to set all necessary environment variables.
HDF5SSPGrid(age=Array([ 0. , 5.100002 , 5.1500006, 5.1999993, 5.25 ,
5.3000016, 5.350002 , 5.4000006, 5.4500012, 5.500002 ,
5.550002 , 5.600002 , 5.6500025, 5.700002 , 5.750002 ,
5.8000026, 5.850003 , 5.900003 , 5.950003 , 6. ,
6.0200005, 6.040001 , 6.0599985, 6.0799985, 6.100002 ,
6.120001 , 6.1399984, 6.16 , 6.18 , 6.1999993,
6.2200007, 6.24 , 6.2599998, 6.2799997, 6.2999997,
6.3199987, 6.3399997, 6.3600006, 6.3799996, 6.3999987,
6.4200006, 6.44 , 6.4599996, 6.4799995, 6.499999 ,
6.52 , 6.539999 , 6.56 , 6.5799994, 6.6 ,
6.6199994, 6.6399994, 6.66 , 6.679999 , 6.699999 ,
6.72 , 6.7399993, 6.7599993, 6.7799997, 6.799999 ,
6.819999 , 6.839999 , 6.8599997, 6.879999 , 6.899999 ,
6.919999 , 6.939999 , 6.959999 , 6.9799986, 6.999999 ,
7.0200005, 7.040001 , 7.0599985, 7.0799985, 7.099998 ,
7.119998 , 7.1399984, 7.16 , 7.18 , 7.1999993,
7.2199984, 7.24 , 7.2599998, 7.2799997, 7.2999997,
7.3199987, 7.3399997, 7.3599987, 7.3799996, 7.3999987,
7.4199986, 7.4399986, 7.462398 , 7.4771214, 7.4913616,
7.50515 , 7.518514 , 7.531479 , 7.544068 , 7.5563025,
7.5682015, 7.5797834, 7.5910645, 7.60206 , 7.628389 ,
7.6532125, 7.6766934, 7.69897 , 7.7201595, 7.7403626,
7.7565446, 7.806545 , 7.8565454, 7.906545 , 7.9565454,
8.006543 , 8.056546 , 8.1065445, 8.156547 , 8.206545 ,
8.256547 , 8.306547 , 8.356546 , 8.406547 , 8.456547 ,
8.506547 , 8.556547 , 8.606546 , 8.656548 , 8.706548 ,
8.756548 , 8.806548 , 8.856548 , 8.9065485, 8.956549 ,
9.006547 , 9.05655 , 9.106548 , 9.156549 , 9.206551 ,
9.225309 , 9.230449 , 9.255273 , 9.278753 , 9.30103 ,
9.322219 , 9.3424225, 9.361728 , 9.380211 , 9.39794 ,
9.414973 , 9.439333 , 9.477121 , 9.511884 , 9.544068 ,
9.574031 , 9.60206 , 9.628389 , 9.653213 , 9.676694 ,
9.69897 , 9.72016 , 9.740363 , 9.759667 , 9.7781515,
9.79588 , 9.812913 , 9.829304 , 9.8450985, 9.860338 ,
9.875061 , 9.889301 , 9.90309 , 9.916454 , 9.929419 ,
9.942008 , 9.954243 , 9.966142 , 9.977724 , 9.989004 ,
10. , 10.010724 , 10.02119 , 10.031408 , 10.041392 ,
10.051152 , 10.060698 , 10.070038 , 10.079182 , 10.088136 ,
10.09691 , 10.10551 , 10.113943 , 10.122216 , 10.130334 ,
10.138303 , 10.146128 , 10.153815 , 10.161368 , 10.168792 ,
10.176091 , 10.1832695, 10.190331 , 10.197281 , 10.20412 ,
10.210854 , 10.2174835, 10.224015 , 10.230449 , 10.236789 ,
10.243038 , 10.249198 , 10.255273 , 10.261263 , 10.267172 ,
10.273002 , 10.278753 , 10.2844305, 10.290034 , 10.2955675,
10.30103 ], dtype=float32), metallicity=Array([0.0001, 0.0004, 0.004 , 0.008 , 0.02 , 0.05 ], dtype=float32), wavelength=Array([ 91., 94., 96., 98., 100., 102., 104., 106.,
108., 110., 114., 118., 121., 125., 127., 128.,
131., 132., 134., 137., 140., 143., 147., 151.,
155., 159., 162., 166., 170., 173., 177., 180.,
182., 186., 191., 194., 198., 202., 205., 210.,
216., 220., 223., 227., 230., 234., 240., 246.,
252., 257., 260., 264., 269., 274., 279., 284.,
290., 296., 301., 308., 318., 328., 338., 348.,
357., 366., 375., 385., 395., 405., 414., 422.,
430., 441., 451., 460., 470., 480., 490., 500.,
506., 512., 520., 530., 540., 550., 560., 570.,
580., 590., 600., 610., 620., 630., 640., 650.,
658., 665., 675., 685., 695., 705., 716., 726.,
735., 745., 755., 765., 775., 785., 795., 805.,
815., 825., 835., 845., 855., 865., 875., 885.,
895., 905., 915., 925., 935., 945., 955., 965.,
975., 985., 995., 1005., 1015., 1025., 1035., 1045.,
1055., 1065., 1075., 1085., 1095., 1105., 1115., 1125.,
1135., 1145., 1155., 1165., 1175., 1185., 1195., 1205.,
1215., 1225., 1235., 1245., 1255., 1265., 1275., 1285.,
1295., 1305., 1315., 1325., 1335., 1345., 1355., 1365.,
1375., 1385., 1395., 1405., 1415., 1425., 1435., 1442.,
1447., 1455., 1465., 1475., 1485., 1495., 1505., 1512.,
1517., 1525., 1535., 1545., 1555., 1565., 1575., 1585.,
1595., 1605., 1615., 1625., 1635., 1645., 1655., 1665.,
1672., 1677., 1685., 1695., 1705., 1715., 1725., 1735.,
1745., 1755., 1765., 1775., 1785., 1795., 1805., 1815.,
1825., 1835., 1845., 1855., 1865., 1875., 1885., 1895.,
1905., 1915., 1925., 1935., 1945., 1955., 1967., 1976.,
1984., 1995., 2005., 2015., 2025., 2035., 2045., 2055.,
2065., 2074., 2078., 2085., 2095., 2105., 2115., 2125.,
2135., 2145., 2155., 2165., 2175., 2185., 2195., 2205.,
2215., 2225., 2235., 2245., 2255., 2265., 2275., 2285.,
2295., 2305., 2315., 2325., 2335., 2345., 2355., 2365.,
2375., 2385., 2395., 2405., 2415., 2425., 2435., 2445.,
2455., 2465., 2475., 2485., 2495., 2505., 2513., 2518.,
2525., 2535., 2545., 2555., 2565., 2575., 2585., 2595.,
2605., 2615., 2625., 2635., 2645., 2655., 2665., 2675.,
2685., 2695., 2705., 2715., 2725., 2735., 2745., 2755.,
2765., 2775., 2785., 2795., 2805., 2815., 2825., 2835.,
2845., 2855., 2865., 2875., 2885., 2895., 2910., 2930.,
2950., 2970., 2990., 3010., 3030., 3050., 3070., 3090.,
3110., 3130., 3150., 3170., 3190., 3210., 3230., 3250.,
3270., 3290., 3310., 3330., 3350., 3370., 3390., 3410.,
3430., 3450., 3470., 3490., 3510., 3530., 3550., 3570.,
3590., 3610., 3630., 3640., 3650., 3670., 3690., 3710.,
3730., 3750., 3770., 3790., 3810., 3830., 3850., 3870.,
3890., 3910., 3930., 3950., 3970., 3990., 4010., 4030.,
4050., 4070., 4090., 4110., 4130., 4150., 4170., 4190.,
4210., 4230., 4250., 4270., 4290., 4310., 4330., 4350.,
4370., 4390., 4410., 4430., 4450., 4470., 4490., 4510.,
4530., 4550., 4570., 4590., 4610., 4630., 4650., 4670.,
4690., 4710., 4730., 4750., 4770., 4790., 4810., 4830.,
4850., 4870., 4890., 4910., 4930., 4950., 4970., 4990.,
5010., 5030., 5050., 5070., 5090., 5110., 5130., 5150.,
5170., 5190., 5210., 5230., 5250., 5270., 5290., 5310.,
5330., 5350., 5370., 5390., 5410., 5430., 5450., 5470.,
5490., 5510., 5530., 5550., 5570., 5590., 5610., 5630.,
5650., 5670., 5690., 5710., 5730., 5750., 5770., 5790.,
5810., 5830., 5850., 5870., 5890., 5910., 5930., 5950.,
5970., 5990., 6010., 6030., 6050., 6070., 6090., 6110.,
6130., 6150., 6170., 6190., 6210., 6230., 6250., 6270.,
6290., 6310., 6330., 6350., 6370., 6390., 6410., 6430.,
6450., 6470., 6490., 6510., 6530., 6550., 6570., 6590.,
6610., 6630., 6650., 6670., 6690., 6710., 6730., 6750.,
6770., 6790., 6810., 6830., 6850., 6870., 6890., 6910.,
6930., 6950., 6970., 6990., 7010., 7030., 7050., 7070.,
7090., 7110., 7130., 7150., 7170., 7190., 7210., 7230.,
7250., 7270., 7290., 7310., 7330., 7350., 7370., 7390.,
7410., 7430., 7450., 7470., 7490., 7510., 7530., 7550.,
7570., 7590., 7610., 7630., 7650., 7670., 7690., 7710.,
7730., 7750., 7770., 7790., 7810., 7830., 7850., 7870.,
7890., 7910., 7930., 7950., 7970., 7990., 8010., 8030.,
8050., 8070., 8090., 8110., 8130., 8150., 8170., 8190.,
8210., 8230., 8250., 8270., 8290., 8310., 8330., 8350.,
8370., 8390., 8410., 8430., 8450., 8470., 8490., 8510.,
8530., 8550., 8570., 8590., 8610., 8630., 8650., 8670.,
8690., 8710., 8730., 8750., 8770., 8790., 8810., 8830.,
8850., 8870., 8890., 8910., 8930., 8950., 8970., 8990.,
9010., 9030., 9050., 9070., 9090., 9110., 9130., 9150.,
9170., 9190., 9210., 9230., 9250., 9270., 9290., 9310.,
9330., 9350., 9370., 9390., 9410., 9430., 9450., 9470.,
9490., 9510., 9530., 9550., 9570., 9590., 9610., 9630.,
9650., 9670., 9690., 9710., 9730., 9750., 9770., 9790.,
9810., 9830., 9850., 9870., 9890., 9910., 9930., 9950.,
9970., 9990., 10025., 10075., 10125., 10175., 10225., 10275.,
10325., 10375., 10425., 10475., 10525., 10575., 10625., 10675.,
10725., 10775., 10825., 10875., 10925., 10975., 11025., 11075.,
11125., 11175., 11225., 11275., 11325., 11375., 11425., 11475.,
11525., 11575., 11625., 11675., 11725., 11775., 11825., 11875.,
11925., 11975., 12025., 12075., 12125., 12175., 12225., 12275.,
12325., 12375., 12425., 12475., 12525., 12575., 12625., 12675.,
12725., 12775., 12825., 12875., 12925., 12975., 13025., 13075.,
13125., 13175., 13225., 13275., 13325., 13375., 13425., 13475.,
13525., 13575., 13625., 13675., 13725., 13775., 13825., 13875.,
13925., 13975., 14025., 14075., 14125., 14175., 14225., 14275.,
14325., 14375., 14425., 14475., 14525., 14570., 14620., 14675.,
14725., 14775., 14825., 14875., 14925., 14975., 15025., 15075.,
15125., 15175., 15225., 15275., 15325., 15375., 15425., 15475.,
15525., 15575., 15625., 15675., 15725., 15775., 15825., 15875.,
15925., 15975., 16050., 16150., 16250., 16350., 16450., 16550.,
16650., 16750., 16850., 16950., 17050., 17150., 17250., 17350.,
17450., 17550., 17650., 17750., 17850., 17950., 18050., 18150.,
18250., 18350., 18450., 18550., 18650., 18750., 18850., 18950.,
19050., 19150., 19250., 19350., 19450., 19550., 19650., 19750.,
19850., 19950.], dtype=float32), flux=Array([[[9.08833684e-08, 1.93420703e-07, 3.10973348e-07, ...,
1.92249590e-05, 1.88633931e-05, 1.85086974e-05],
[9.08833684e-08, 1.93420703e-07, 3.10973348e-07, ...,
1.92249590e-05, 1.88633931e-05, 1.85086974e-05],
[9.08833684e-08, 1.93420703e-07, 3.10973348e-07, ...,
1.92249590e-05, 1.88633931e-05, 1.85086974e-05],
...,
[5.92562333e-10, 8.93100538e-10, 1.15493171e-09, ...,
2.39835890e-06, 2.35784546e-06, 2.32140042e-06],
[5.92806859e-10, 8.92882435e-10, 1.15413190e-09, ...,
2.37455151e-06, 2.33498645e-06, 2.29807620e-06],
[5.95643035e-10, 8.97048713e-10, 1.15942633e-09, ...,
2.35168159e-06, 2.31248464e-06, 2.27596547e-06]],
[[2.11160405e-08, 4.68378190e-08, 7.72740307e-08, ...,
2.08794318e-05, 2.04886637e-05, 2.01090988e-05],
[2.11160405e-08, 4.68378190e-08, 7.72740307e-08, ...,
2.08794318e-05, 2.04886637e-05, 2.01090988e-05],
[2.11160405e-08, 4.68378190e-08, 7.72740307e-08, ...,
2.08794318e-05, 2.04886637e-05, 2.01090988e-05],
...,
[5.63963209e-10, 8.50090109e-10, 1.09938125e-09, ...,
2.57541342e-06, 2.53532630e-06, 2.49656500e-06],
[5.59437219e-10, 8.43146331e-10, 1.09030318e-09, ...,
2.55510099e-06, 2.51477172e-06, 2.47722096e-06],
[5.78517234e-10, 8.71934414e-10, 1.12751075e-09, ...,
2.53303801e-06, 2.49305162e-06, 2.45587876e-06]],
[[1.11427291e-10, 2.75856810e-10, 4.93186603e-10, ...,
3.00550819e-05, 2.95078007e-05, 2.89541367e-05],
[1.11427291e-10, 2.75856810e-10, 4.93186603e-10, ...,
3.00550819e-05, 2.95078007e-05, 2.89541367e-05],
[1.11427291e-10, 2.75856810e-10, 4.93186603e-10, ...,
3.00550819e-05, 2.95078007e-05, 2.89541367e-05],
...,
[1.51815840e-08, 1.92815222e-08, 2.29955877e-08, ...,
3.14909880e-06, 3.10474729e-06, 3.06152378e-06],
[1.55623212e-08, 1.97692778e-08, 2.35827819e-08, ...,
3.12075917e-06, 3.07683240e-06, 3.03407387e-06],
[1.56620601e-08, 1.98958627e-08, 2.37337012e-08, ...,
3.10205382e-06, 3.05840922e-06, 3.01598016e-06]],
[[6.33916183e-11, 1.56637481e-10, 2.80225038e-10, ...,
3.40314473e-05, 3.34144715e-05, 3.28001406e-05],
[6.33916183e-11, 1.56637481e-10, 2.80225038e-10, ...,
3.40314473e-05, 3.34144715e-05, 3.28001406e-05],
[6.33916183e-11, 1.56637481e-10, 2.80225038e-10, ...,
3.40314473e-05, 3.34144715e-05, 3.28001406e-05],
...,
[1.13446195e-08, 1.44345762e-08, 1.72374950e-08, ...,
3.58108127e-06, 3.53232667e-06, 3.49160928e-06],
[1.14191590e-08, 1.45293875e-08, 1.73506933e-08, ...,
3.54622898e-06, 3.49792595e-06, 3.45767330e-06],
[1.14927898e-08, 1.46229295e-08, 1.74622912e-08, ...,
3.51071185e-06, 3.46286311e-06, 3.42306453e-06]],
[[1.03717389e-14, 2.60376945e-14, 6.23507932e-14, ...,
4.28130661e-05, 4.20417018e-05, 4.12843074e-05],
[1.03717389e-14, 2.60376945e-14, 6.23507932e-14, ...,
4.28130661e-05, 4.20417018e-05, 4.12843074e-05],
[1.03717389e-14, 2.60376945e-14, 6.23507932e-14, ...,
4.28130661e-05, 4.20417018e-05, 4.12843074e-05],
...,
[2.74051143e-10, 4.33427960e-10, 5.86995785e-10, ...,
3.62579908e-06, 3.56578244e-06, 3.53157429e-06],
[2.80006740e-10, 4.42861414e-10, 5.99826022e-10, ...,
3.59876890e-06, 3.53911469e-06, 3.50530217e-06],
[2.81731083e-10, 4.45578630e-10, 6.03499362e-10, ...,
3.57047224e-06, 3.51121457e-06, 3.47779246e-06]],
[[2.64753693e-18, 8.02830980e-18, 2.30857457e-17, ...,
5.49388205e-05, 5.39541179e-05, 5.29583958e-05],
[2.64753693e-18, 8.02830980e-18, 2.30857457e-17, ...,
5.49388205e-05, 5.39541179e-05, 5.29583958e-05],
[2.69226858e-18, 8.17344360e-18, 2.35313512e-17, ...,
5.90876080e-05, 5.80271771e-05, 5.69552649e-05],
...,
[2.86055124e-10, 4.52389348e-10, 6.12669249e-10, ...,
3.57395697e-06, 3.51914946e-06, 3.49452603e-06],
[2.92348756e-10, 4.62365729e-10, 6.26242114e-10, ...,
3.54419944e-06, 3.48981166e-06, 3.46525371e-06],
[2.94150426e-10, 4.65220779e-10, 6.30102970e-10, ...,
3.51500717e-06, 3.46103275e-06, 3.43656484e-06]]], dtype=float32))
SSP lookup#
# NBVAL_SKIP
import matplotlib.pyplot as plt
from rubix.spectra.ssp.templates import BruzualCharlot2003
from jax import jit
ssp = BruzualCharlot2003
wave = ssp.wavelength
age_index = 0
met_index = 3
target_age = ssp.age[age_index] + 0.5*(ssp.age[age_index+1] - ssp.age[age_index])
print(target_age)
target_met = ssp.metallicity[met_index] + 0.5*(ssp.metallicity[met_index+1] - ssp.metallicity[met_index])
lookup = ssp.get_lookup_interpolation()
spec_calc = lookup(target_met, target_age)
spec_true = ssp.flux[met_index, age_index, :]
plt.plot(wave, spec_calc, label='calc')
plt.plot(wave, spec_true, label='true')
plt.legend()
plt.yscale('log')
2.550001
# NBVAL_SKIP
# Check if it works with jit
spec_calc = jit(lookup)(target_met, target_age)
plt.plot(wave, spec_calc, label='calc jit')
[<matplotlib.lines.Line2D at 0x73621018c140>]
# NBVAL_SKIP
from rubix.utils import load_galaxy_data
data, units = load_galaxy_data("output/rubix_galaxy.h5")
mass = data["particle_data"]["stars"]["mass"]
metallicity = data["particle_data"]["stars"]["metallicity"]
age = data["particle_data"]["stars"]["age"]
VMAP#
Vmap the lookup over the stellar particles
# NBVAL_SKIP
# Calculate spectra with vmap
from jax import vmap
lookup = ssp.get_lookup_interpolation()
subset = 1000
# Use only subset because it is too big to fit into gpu memory
met_subset = metallicity[:subset]
age_subset = age[:subset]
# Clip the metallicity and age values to the range of the SSP
met_subset = met_subset.clip(min(ssp.metallicity), max(ssp.metallicity))
age_subset = age_subset.clip(min(ssp.age), max(ssp.age))
spec_calc = vmap(lookup)(met_subset, age_subset)
spec_calc.shape
(1000, 842)
# NBVAL_SKIP
# check if it contains nan values
import jax.numpy as jnp
jnp.isnan(spec_calc).any()
Array(False, dtype=bool)
Use configuration to load lookup function#
#NBVAL_SKIP
config ={ "ssp": {
"template": {
"name": "BruzualCharlot2003"
},
"method": "cubic"
}
}
# NBVAL_SKIP
from rubix.core.ssp import get_lookup_interpolation
lookup = get_lookup_interpolation(config)
2025-11-10 17:13:25,095 - rubix - DEBUG - Using method defined in config: cubic
# NBVAL_SKIP
# Check how many particles are outside the range of the SSP
import numpy as np
np.sum(metallicity < ssp.metallicity[0]), np.sum(metallicity > ssp.metallicity[-1])
(Array(695, dtype=int32), Array(35669, dtype=int32))
# NBVAL_SKIP
np.sum(age < ssp.age[0]), np.sum(age > ssp.age[-1])
(Array(0, dtype=int32), Array(45331, dtype=int32))
# NBVAL_SKIP
len(metallicity)
643940
# NBVAL_SKIP
# clip the metallicity and age values to the range of the SSP
met_subset = met_subset.clip(min(ssp.metallicity), max(ssp.metallicity))
age_subset = age_subset.clip(min(ssp.age), max(ssp.age))
lookup(met_subset, age_subset)
Array([[ 0.0000000e+00, 2.7755576e-17, 5.5511151e-17, ...,
4.5619352e-05, 4.4804754e-05, 4.3980177e-05],
[ 0.0000000e+00, 2.7755576e-17, 5.5511151e-17, ...,
4.5822202e-05, 4.5003901e-05, 4.4175598e-05],
[ 0.0000000e+00, 2.7755576e-17, 5.5511151e-17, ...,
4.8169481e-05, 4.7308338e-05, 4.6436926e-05],
...,
[ 0.0000000e+00, 2.7755576e-17, 5.5511151e-17, ...,
4.6039131e-05, 4.5216868e-05, 4.4384582e-05],
[-5.2917185e-12, -1.3075374e-11, -2.3385127e-11, ...,
4.5421279e-05, 4.4608027e-05, 4.3794011e-05],
[ 0.0000000e+00, 2.7755576e-17, 5.5511151e-17, ...,
4.4646134e-05, 4.3849301e-05, 4.3042593e-05]], dtype=float32)