Synthesis of semi-resolved stellar populations

Synthesizing a semi-resolved stellar population model with FASTAR is rather straightforward. Semi-resolved models follow the principles of evolutionary stellar population synthesis and therefore they are defined by the age of the population, its metallicity, and IMF. There is, however, a fundamental difference with respect to standard SSP models, as the user must define the number of stars contributing to the model. Let’s have a look at a basic example

from matplotlib import pyplot as plt

import jax
from fastar import SemiresolvedSSPSynthesizer
from fastar.imf import kroupa

# Initialize the synthesizer with the desired IMF functional form
semi = SemiresolvedSSPSynthesizer(imf_function=kroupa)

# This randomizes the IMF sampling
key = jax.random.PRNGKey(0)

# Basic synthesis
age = 10  # e.g. 10 Gyr
met = 0  # Solar metallicity
nstars = 1e2  # Number of stars

semi_spec, semi_stellar_mass = semi.synthesize(
    age=age, met=met, num_stars=int(nstars), key=key
)

fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(semi.wave, semi_spec, color='xkcd:green')

ax.set_xlabel('Wavelength [Å]')
ax.set_ylabel(r'Flux [erg s$^{-1}$ cm$^{-2}$ $\AA^{-1}$]')
ax.set_title(f'Age: {age} Gyr, [M/H]: {met}')

fig.tight_layout()
plt.show()

Note that the synthesis of semi-resolved models has two outputs: the model spectrum and the total stellar mass of the model. This is because, contrary to standard SSP models (usually normalized to one solar mass at birth), the mass of the stellar population depends on the number of stars. In addition, the actual value of the stellar mass depends on the specific IMF sampling, and therefore two different model realizations will most likely have different stellar masses even if the age, metallicity, IMF and number of stars are the same.

By default, FASTAR outputs the total stellar mass but it can be called so it returns an array with the individual stellar masses of the stars in the model. This is controlled with the out_masses keywords:

semi_spec, semi_individual_stellar_masses = semi.synthesize(
    age=age, met=met, num_stars=int(nstars), key=key, out_masses=True
)

Semi-resolved SSP models with a large number of stars

The semi-resolved version of FASTAR explicitly calculates all the individual stellar spectra contributing to the integrated flux. While this has the advantage of exactly representing the composite flux of finite ensemble of stars with no additional simplification, the computation can be memory intensive.

To enable the computation of models with large \(N_{stars}\) without running into memory issues, FASTAR can internally split the model computation in smaller batches. The output is exactly the same.

# Generate the key
key = jax.random.PRNGKey(0)

# And then simply synthesize a semi-resolved model
age = 11.0
met = 0.0
nstars = 1e5

# Synthesize SSP
semi_model, semi_stellar_mass = semi.synthesize_large(
    age=age, met=met, num_stars=int(nstars), key=key, batch_size=int(1e4)
)

Creating a large set of semi-resolved models

Measuring stellar population properties in th semi-resolved regime presents some technical challenges due to the unavoidable stochasticity. Therefore, it is common that a large number of different realizations for the same age, metallicity and \(N_{stars}\) are needed in order to model observed data. This is the case for example if you need to generate a large training set for a SBI framework.

A possible way to generate a large model grid is as follows

import numpy as np
from tqdm import tqdm, trange


def synthesize_vmappable(age, met, key):
    return semi.synthesize(age=age, met=met, num_stars=num_stars, key=key)


# Let's define the age, met
age_min = 0.1  # Gyr
age_max = 14.0  # Gyr

met_min = -1.0  # dex
met_min = +0.2  # dex

# As explained below, it is much more efficient to quantify Nstars
log_Nstars_min = 1.5
log_Nstars_max = 2.5
n_bins = 5  # This is the number of bins in Nstar
Nstars_bins = (
    (10 ** np.linspace(log_Nstars_min, log_Nstars_max, n_bins)).round().astype(int)
)

samples_per_bin = int(100)  # This sets the number of model realizations per Nstar bin.
key = jax.random.PRNGKey(0)

param_list = []
spec_list = []

tqdm.disable = True
for num_stars in tqdm(Nstars_bins):
    key, subkey = jax.random.split(key)
    keys = jax.random.split(subkey, samples_per_bin)

    key, subkey = jax.random.split(key)
    ages = jax.random.uniform(
        subkey, (samples_per_bin,), minval=age_min, maxval=age_max
    )

    key, subkey = jax.random.split(key)
    mets = jax.random.uniform(
        subkey, (samples_per_bin,), minval=met_min, maxval=met_min
    )

    for i in trange(samples_per_bin, leave=False):
        age, met, k = float(ages[i]), float(mets[i]), keys[i]
        output = synthesize_vmappable(age, met, k)

        # Extract spectrum and stellar mass
        spectrum = output[0]
        stellar_mass = np.log10(output[1])

        # Save parameters and spectrum
        param_list.append([float(age), float(met), np.log10(num_stars), stellar_mass])
        spec_list.append(np.array(spectrum))

  0%|          | 0/5 [00:00<?, ?it/s]
  0%|          | 0/100 [00:00<?, ?it/s]
  1%|          | 1/100 [00:04<06:44,  4.08s/it]
 56%|█████▌    | 56/100 [00:04<00:02, 18.82it/s]
 20%|██        | 1/5 [00:05<00:21,  5.26s/it]
  0%|          | 0/100 [00:00<?, ?it/s]
  1%|          | 1/100 [00:03<05:13,  3.17s/it]
 47%|████▋     | 47/100 [00:03<00:02, 20.13it/s]
 94%|█████████▍| 94/100 [00:03<00:00, 46.17it/s]
 40%|████      | 2/5 [00:08<00:12,  4.16s/it]
  0%|          | 0/100 [00:00<?, ?it/s]
  1%|          | 1/100 [00:03<05:02,  3.05s/it]
 36%|███▌      | 36/100 [00:03<00:04, 15.95it/s]
 70%|███████   | 70/100 [00:03<00:00, 35.35it/s]
 60%|██████    | 3/5 [00:12<00:07,  3.79s/it]
  0%|          | 0/100 [00:00<?, ?it/s]
  1%|          | 1/100 [00:03<05:17,  3.21s/it]
 26%|██▌       | 26/100 [00:03<00:06, 10.93it/s]
 51%|█████     | 51/100 [00:03<00:01, 24.58it/s]
 75%|███████▌  | 75/100 [00:03<00:00, 40.84it/s]
 99%|█████████▉| 99/100 [00:03<00:00, 60.22it/s]
 80%|████████  | 4/5 [00:15<00:03,  3.73s/it]
  0%|          | 0/100 [00:00<?, ?it/s]
  1%|          | 1/100 [00:03<05:19,  3.23s/it]
 17%|█▋        | 17/100 [00:03<00:11,  7.06it/s]
 33%|███▎      | 33/100 [00:03<00:04, 15.77it/s]
 49%|████▉     | 49/100 [00:03<00:01, 26.61it/s]
 65%|██████▌   | 65/100 [00:03<00:00, 39.52it/s]
 81%|████████  | 81/100 [00:03<00:00, 54.13it/s]
 97%|█████████▋| 97/100 [00:03<00:00, 69.74it/s]
100%|██████████| 5/5 [00:19<00:00,  3.90s/it]

The code above synthesizes 500 semi-resolved models with random ages between 0.1 and 14 Gyr and random metallicities from [M/H]=-1 to [M/H]=0.2 and uses some of key functionalities of JAX, in particular vmap, to speed up the process.

Note that, while age and metallicity are randomly selected, \(N_{stars}\) is discretized. This is because FASTAR internally uses JAX just-in-time (JIT) to optimize the computation but \(N_{stars}\) cannot be jitted since it changes the internal dimension of the arrays. Randomly changing \(N_{stars}\) would also work, but it is much slower.

Optimizing model synthesis

The example above works, but it is not necessarily the most efficient way to generate a large set of models. We encourage users to explore the capabilities of FASTAR in combination with JAX to accelerate model synthesis.

If you find a clear way to optimize the use of FASTAR, please let us know.