Dataset for MSMS binary: Millions of Main-Sequence Binary Stars from Gaia BP/RP Spectra

We present the main-sequence binary (MSMS) Catalog derived from Gaia Data Release 3 BP/RP (XP) spectra. Leveraging the exceptional quality of low-resolution XP spectra, we develop a forward modeling approach that maps stellar parameters of stellar mass and photometric metallcity to XP spectra through a neural network. Our methodology identifies binary systems through statistical comparison of single- and binary-star model fits, enabling detection of binaries with mass ratios between 0.4 and 0.8 and flux ratios larger than 0.1. From an initial sample of 35 million stars within 1 kpc, we identify 14 million binary candidates and define a high-confidence "golden sample" of 1 million binary systems. 

This code demonstrates how to generate emulated spectra using trained neural network models for both single and binary star systems. It shows how to load the necessary models and scalers, create forward models with user-specified parameters (mass, metallicity, extinction, and mass ratio for binaries), and convert the model outputs to physical flux units. The code is designed to work with Gaia XP spectral data and includes functionality to visualize the differences between single and binary star spectra, simulate component contributions in binary systems.

import numpy as np
import torch
import pickle
import os
from yutu import YuTuqPredictor, NNPredictor
from nnm import NNModel

def load_ms_models(model_dir, device=None):
"""
Load MS models for both single star and binary star prediction.

Parameters:
model_dir (str): Directory containing model files
device (torch.device, optional): Device to load models to

Returns:
tuple: (single_ms, binary_msms, flux_scaler)
- single_ms: NNPredictor for single star spectra
- binary_msms: YuTuqPredictor for binary star spectra
- flux_scaler: Scaler for flux normalization
"""
if device is None:
device = torch.device("cpu")

# Load the flux scaler
flux_scaler_path = os.path.join(model_dir, 'flux_scaler_v250203.pkl')
with open(flux_scaler_path, 'rb') as f:
flux_scaler = pickle.load(f)

# Load the main model containing the neural network weights
n_label = 2
n_pixel = 61

nnm_mm = NNModel(n_label=n_label, n_pixel=n_pixel,
n_hidden=64, n_layer=5, drop_rate=1e-6, wave=None)
nnm_mm.load_state_dict(torch.load(os.path.join(model_dir, "nnm_25022_mm2xp.pt"),
map_location=device))

# Extract weights and biases from the model
w = [nnm_mm.state_dict()[k].cpu().numpy() for k in nnm_mm.state_dict().keys() if "weight" in k]
b = [nnm_mm.state_dict()[k].cpu().numpy()[:,None] for k in nnm_mm.state_dict().keys() if "bias" in k]
alpha = 0.01 if nnm_mm.activation == "relu" else 1.0

xmin = nnm_mm.xmin
xmax = nnm_mm.xmax

ext_curve_file = os.path.join(model_dir, "../data/extinction_curve.npy")

# Create predictors
single_ms = NNPredictor(w, b, alpha, xmin, xmax, flux_scaler=flux_scaler,
n_label=n_label, n_pixel=n_pixel,
ext_curve_file=ext_curve_file)

binary_msms = YuTuqPredictor(w, b, alpha, xmin, xmax, flux_scaler=flux_scaler,
n_label=n_label, n_pixel=n_pixel,
ext_curve_file=ext_curve_file)

return single_ms, binary_msms, flux_scaler

def forward_model_spectra(mass, mh, extinction=0.0, mass_ratio=None, single_ms=None,
binary_msms=None, flux_scaler=None):
"""
Generate forward model spectra for single or binary stars.

Parameters:
mass (float): Mass of the primary star in solar masses
mh (float): Metallicity [Fe/H] of the star(s)
extinction (float, optional): Extinction value (E(B-V))
mass_ratio (float, optional): Mass ratio (M2/M1) if modeling a binary
single_ms (NNPredictor): Single star spectrum predictor
binary_msms (YuTuqPredictor): Binary star spectrum predictor
flux_scaler: Scaler to transform spectra back to physical scale

Returns:
tuple: (wavelengths, flux)
- wavelengths (ndarray): Wavelength grid in nm
- flux (ndarray): Predicted flux, with extinction applied if specified
- For binaries, also returns the individual component spectra
"""
# Ensure the predictors are provided
if single_ms is None or binary_msms is None or flux_scaler is None:
raise ValueError("Spectrum predictors and flux_scaler must be provided")

# Set up wavelength grid
wavelengths = np.linspace(380, 930, 61) # XP wavelength grid

# Generate single star spectrum
if mass_ratio is None:
# Single star model
params = np.array([mass, mh, extinction])
normalized_flux = single_ms.predict_one_spectrum_ext(params)

# Convert back to physical flux scale
physical_flux = 10**(flux_scaler.inverse_transform(normalized_flux.reshape(1, -1)).flatten())

return wavelengths, physical_flux

else:
# Binary star model
binary_params = np.array([mass, mh, mass_ratio, extinction])
normalized_flux_binary = binary_msms.predict_one_binary_spectrum(binary_params, use_ext_predict=True)

# Also generate individual component spectra
mass_secondary = mass * mass_ratio
primary_params = np.array([mass, mh, extinction])
secondary_params = np.array([mass_secondary, mh, extinction])

normalized_flux_primary = single_ms.predict_one_spectrum_ext(primary_params)
normalized_flux_secondary = single_ms.predict_one_spectrum_ext(secondary_params)

# Convert back to physical flux scale
physical_flux_binary = 10**(flux_scaler.inverse_transform(normalized_flux_binary.reshape(1, -1)).flatten())
physical_flux_primary = 10**(flux_scaler.inverse_transform(normalized_flux_primary.reshape(1, -1)).flatten())
physical_flux_secondary = 10**(flux_scaler.inverse_transform(normalized_flux_secondary.reshape(1, -1)).flatten())

return wavelengths, physical_flux_binary, physical_flux_primary, physical_flux_secondary

# Example usage:
if __name__ == "__main__":
# Load models
model_dir = "./model/"
single_ms, binary_msms, flux_scaler = load_ms_models(model_dir)

# Example 1: Generate single star spectrum
wave, flux_single = forward_model_spectra(
mass=0.8,
mh=-0.5,
extinction=0.1,
single_ms=single_ms,
binary_msms=binary_msms,
flux_scaler=flux_scaler
)

# Example 2: Generate binary star spectrum
wave, flux_binary, flux_primary, flux_secondary = forward_model_spectra(
mass=1.0,
mh=0.0,
extinction=0.05,
mass_ratio=0.7,
single_ms=single_ms,
binary_msms=binary_msms,
flux_scaler=flux_scaler
)