emulator¶
Neural-network warm-start package for Kurucz stellar atmospheres.
The emulator replaces the traditional Fortran READ DECK6 starting guess with a
fast, pre-trained MLP that predicts atmospheric structure (RHOX, T, P, XNE,
ABROSS, ACCRAD) from 4D stellar parameters (Teff, logg, [Fe/H], [α/Fe]). This
allows atlas_py to converge in a single outer iteration rather than starting
from a grey approximation.
Most users invoke the emulator indirectly through
pykurucz.emulator_warmstart_atm().
See also:
Atmosphere Emulator¶
Atmosphere emulator for Kurucz stellar atmospheres.
This module provides a neural network-based emulator that predicts atmospheric structure from 4D stellar parameters (Teff, logg, [Fe/H], [α/Fe]).
Classes¶
AtmosphereEmulator(model, normalizer, default_tau_grid=None, device='cpu')
Emulator for Kurucz stellar atmosphere models.
This class provides an interface to predict atmospheric structure based on stellar parameters and optical depth using a pre-trained neural network model.
Input: 4D stellar parameters (Teff, logg, [Fe/H], [α/Fe]) Output: 80 depth points × 6 quantities (RHOX, T, P, XNE, ABROSS, ACCRAD)
Initialize the emulator with a pre-trained model and normalization helper.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model
|
Module
|
Pre-trained neural network model |
required |
normalizer
|
NormalizationHelper
|
Normalization helper object |
required |
default_tau_grid
|
Tensor
|
Default optical depth grid |
None
|
device
|
str
|
Device to run the model on ('cpu' or 'cuda') |
'cpu'
|
Methods:¶
predict(teff, logg, feh, afe, tau_grid=None)
Predict atmospheric structure for given stellar parameters.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
teff
|
float
|
Effective temperature in K |
required |
logg
|
float
|
Surface gravity (log10 cgs) |
required |
feh
|
float
|
Metallicity [Fe/H] |
required |
afe
|
float
|
Alpha enhancement [α/Fe] |
required |
tau_grid
|
array - like
|
Optical depth grid (80 points) If None, uses the default grid |
None
|
Returns:
| Name | Type | Description |
|---|---|---|
dict |
Dictionary containing atmospheric parameters for each depth point Keys: 'RHOX', 'T', 'P', 'XNE', 'ABROSS', 'ACCRAD', 'TAU' Each is a numpy array of shape (80,) |
predict_atmosphere_data(teff, logg, feh, afe, vturb=2.0, tau_grid=None)
Predict the 80x9 atmosphere data array for use in .atm files.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
teff
|
float
|
Effective temperature in K |
required |
logg
|
float
|
Surface gravity (log10 cgs) |
required |
feh
|
float
|
Metallicity [Fe/H] |
required |
afe
|
float
|
Alpha enhancement [α/Fe] |
required |
vturb
|
float
|
Microturbulent velocity in km/s (default: 2.0) |
2.0
|
tau_grid
|
array - like
|
TAU grid from closest atmosphere. IMPORTANT: Must use TAU from a real atmosphere for accurate predictions. The MLP was trained with TAU calculated from each atmosphere, not a fixed grid. |
None
|
Returns:
| Type | Description |
|---|---|
|
numpy.ndarray: 80x9 atmosphere data array with columns: RHOX, T, P, XNE, ABROSS, ACCRAD, VTURB, FLXCNV, VCONV |
Functions:¶
load_emulator(weights_path=None, norm_params_path=None, device='cpu')
Load the pre-trained atmosphere emulator.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
weights_path
|
str or Path
|
Path to model weights file. If None, uses bundled weights. |
None
|
norm_params_path
|
str or Path
|
Path to normalization parameters. If None, uses bundled params. |
None
|
device
|
str
|
Device to load the model on ('cpu' or 'cuda') |
'cpu'
|
Returns:
| Name | Type | Description |
|---|---|---|
AtmosphereEmulator |
Initialized emulator object |
Model¶
Neural network model for Kurucz stellar atmosphere prediction.
This model takes 4D stellar parameters (Teff, logg, [Fe/H], [α/Fe]) plus optical depth tau values and predicts 6D atmospheric structure (RHOX, T, P, XNE, ABROSS, ACCRAD).
Classes¶
StellarParamEncoder(input_dim=4, embed_dim=128)
Bases: Module
Encoder for global stellar parameters (Teff, logg, [M/H], [alpha/Fe])
TauPositionEncoder(embed_dim=64, depth_points=80)
Bases: Module
Position encoder for tau values at each depth point
AtmosphereNetMLPtau(output_size=6, depth_points=80, stellar_embed_dim=128, tau_embed_dim=64)
Bases: Module
MLP-based atmosphere model with separate encoders for stellar params and tau.
Input: 4 stellar parameters + 80 tau values Output: 80 depth points × 6 atmospheric quantities
Normalization¶
Normalization utilities for the atmosphere emulator.
Classes¶
NormalizationHelper(norm_params)
Helper class for normalizing and denormalizing data using saved parameters.
Initialize with normalization parameters.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
norm_params
|
dict
|
Normalization parameters |
required |
Methods:¶
normalize(param_name, data) -> torch.Tensor
Normalize data to [-1, 1] range with optional log transform.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
param_name
|
str
|
Name of the parameter to normalize |
required |
data
|
Tensor
|
Data to normalize |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
torch.Tensor: Normalized data in [-1, 1] range |
denormalize(param_name, normalized_data) -> torch.Tensor
Denormalize data from [-1, 1] range back to original scale.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
param_name
|
str
|
Name of the parameter to denormalize |
required |
normalized_data
|
Tensor
|
Normalized data to convert back |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
torch.Tensor: Denormalized data with gradients preserved |
Functions:¶
load_norm_params(norm_params_path) -> dict
Load normalization parameters from a standalone file.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
norm_params_path
|
str
|
Path to the normalization parameters file |
required |
Returns:
| Name | Type | Description |
|---|---|---|
dict |
dict
|
Normalization parameters |