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
)

我们发布了基于盖亚数据发布3（Gaia Data Release 3）BP/RP（XP）光谱构建的主序双星（main-sequence binary, MSMS）星表。依托低分辨率XP光谱的卓越性能，我们开发了一种正演建模（forward modeling）方法，通过神经网络将恒星质量、测光金属丰度等恒星参数映射至XP光谱。本研究方法通过单星与双星模型拟合结果的统计比对识别双星系统，可检测质量比介于0.4至0.8之间、流量比（flux ratios）大于0.1的双星系统。我们从1千秒差距（kpc）范围内的3500万颗恒星初始样本中，筛选出1400万颗双星候选体，并定义了包含100万颗双星系统的高置信度“黄金样本（golden sample）”。

本代码演示了如何利用训练好的神经网络模型生成单星与双星系统的模拟光谱。代码展示了如何加载所需模型与标准化器（scaler），基于用户指定参数（恒星质量、金属丰度、消光量，双星系统还需质量比）构建正演模型，并将模型输出转换为物理流量单位。本代码适配盖亚XP光谱数据，支持可视化单星与双星光谱的差异，以及模拟双星系统中各子星的光谱贡献。

# 导入NumPy数值计算库
import numpy as np
# 导入PyTorch深度学习框架
import torch
# 导入pickle序列化工具
import pickle
# 导入os系统交互模块
import os
# 从yutu模块导入YuTuq预测器与神经网络预测器
from yutu import YuTuqPredictor, NNPredictor
# 从nnm模块导入神经网络模型
from nnm import NNModel

def load_ms_models(model_dir, device=None):
"""
加载用于单星与双星光谱预测的主序星模型。

参数：
model_dir (str): 模型文件所在目录路径
device (torch.device, 可选): 加载模型的计算设备

返回值：
元组: (single_ms, binary_msms, flux_scaler)
- single_ms: 用于单星光谱预测的NNPredictor
- binary_msms: 用于双星光谱预测的YuTuqPredictor
- flux_scaler: 用于流量归一化的标准化器
"""
if device is None:
# 若未指定设备，则默认使用CPU
device = torch.device("cpu")

# 加载流量标准化器
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)

# 加载包含神经网络权重的主模型
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))

# 提取模型权重与偏置
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")

# 创建单星光谱预测器
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):
"""
为单星或双星系统生成正演模拟光谱。

参数：
mass (float): 主星质量，单位为太阳质量
mh (float): 恒星的金属丰度[Fe/H]
extinction (float, 可选): 消光量E(B-V)，默认值为0.0
mass_ratio (float, 可选): 双星系统的质量比M2/M1，仅双星模式需指定
single_ms (NNPredictor): 单星光谱预测器
binary_msms (YuTuqPredictor): 双星光谱预测器
flux_scaler: 用于将光谱转换回物理流量单位的标准化器

返回值：
元组: (wavelengths, flux)
- wavelengths (ndarray): 波长网格，单位为纳米nm
- flux (ndarray): 预测得到的流量，若指定消光量则已应用消光校正
若为双星系统，还会额外返回两个子星的单独光谱
"""
# 检查是否提供了必要的预测器与标准化器
if single_ms is None or binary_msms is None or flux_scaler is None:
raise ValueError("必须提供光谱预测器与流量标准化器")

# 设置XP光谱的波长网格
wavelengths = np.linspace(380, 930, 61)

# 单星模式
if mass_ratio is None:
params = np.array([mass, mh, extinction])
normalized_flux = single_ms.predict_one_spectrum_ext(params)

# 转换为物理流量单位
physical_flux = 10**(flux_scaler.inverse_transform(normalized_flux.reshape(1, -1)).flatten())

return wavelengths, physical_flux

else:
# 双星模式
binary_params = np.array([mass, mh, mass_ratio, extinction])
normalized_flux_binary = binary_msms.predict_one_binary_spectrum(binary_params, use_ext_predict=True)

# 计算两颗子星的单独质量
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)

# 转换为物理流量单位
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

# 示例用法
if __name__ == "__main__":
# 加载模型
model_dir = "./model/"
single_ms, binary_msms, flux_scaler = load_ms_models(model_dir)

# 示例1：生成单星光谱
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
)

# 示例2：生成双星光谱
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
)