Usage#

If all you want to do is to predict red giant granulation parameters using the built-in neural network model, you can just use predict()

import pandas as pd
X = pd.DataFrame({
    "M" : [1.38], # solar units
    "R" : [11.14], # solar units
    "Teff" : [4864.98], # Kelvin
    "FeH" : [-0.13], # metallicity
    "KepMag" : [12.55], # apparent magnitude in Kepler band
    "phase" : [2] # 0: any, 1: red giant branch, 2: red clump/helium burning
})

X
M R Teff FeH KepMag phase
0 1.38 11.14 4864.98 -0.13 12.55 2
from grannules import predict
y = predict(X, to_df = True)

# ppm^2/uHz, ppm^2/uHz, s, unitless
y
/home/docs/checkouts/readthedocs.org/user_builds/grannules/envs/latest/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[2], line 2
      1 from grannules import predict
----> 2 y = predict(X, to_df = True)
      4 # ppm^2/uHz, ppm^2/uHz, s, unitless
      5 y

File ~/checkouts/readthedocs.org/user_builds/grannules/envs/latest/lib/python3.12/site-packages/grannules/neural_net.py:608, in predict(X, to_df, *args, **kwargs)
    584 def predict(
    585         X: pd.DataFrame, to_df: bool = False, 
    586         *args, **kwargs
    587     ) -> np.ndarray | pd.DataFrame:
    588     r"""Predicts the parameters :math:`H,\, P,\, \tau,` and 
    589     :math:`\alpha` for red giant stars using a pre-trained neural network.
    590 
   (...)    606     :rtype: numpy.ndarray | pandas.DataFrame
    607     """
--> 608     predictor = NNPredictor.get_default_predictor(*args, **kwargs)
    609     return predictor.predict(X, to_df)

File ~/checkouts/readthedocs.org/user_builds/grannules/envs/latest/lib/python3.12/site-packages/grannules/neural_net.py:145, in NNPredictor.get_default_predictor(cls, *args, **kwargs)
    139 """Loads a pre-trained NNPredictor singleton.
    140 
    141 :return: A pre-trained NNPredictor
    142 :rtype: NNPredictor
    143 """
    144 if cls._default_predictor is None:
--> 145     cls._default_predictor = cls._default_from_serialize(*args, **kwargs)
    147 return cls._default_predictor

File ~/checkouts/readthedocs.org/user_builds/grannules/envs/latest/lib/python3.12/site-packages/grannules/neural_net.py:582, in NNPredictor._default_from_serialize(cls, path)
    577 @classmethod
    578 def _default_from_serialize(
    579         cls,
    580         path = files(__name__) / "data/default-serialized"
    581 ):
--> 582     return cls.deserialize(path)

File ~/checkouts/readthedocs.org/user_builds/grannules/envs/latest/lib/python3.12/site-packages/grannules/neural_net.py:560, in NNPredictor.deserialize(cls, path)
    551 blank_state = train_state.TrainState.create(
    552     apply_fn=model.apply,
    553     params=model.init(
   (...)    557         tx = optax.adam(1e-3)
    558 )
    559 with open(state_path, "rb") as f:
--> 560     state_dict = pickle.load(f)
    562 state = from_state_dict(blank_state, state_dict)
    564 with open(data_transform_path, "rb") as f:

File ~/checkouts/readthedocs.org/user_builds/grannules/envs/latest/lib/python3.12/site-packages/jax/_src/array.py:126, in _reconstruct_array(fun, args, arr_state, aval_state)
    124 np_value.__setstate__(arr_state)
    125 jnp_value = api.device_put(np_value)
--> 126 jnp_value.aval = jnp_value.aval.update(**aval_state)
    127 return jnp_value

File ~/checkouts/readthedocs.org/user_builds/grannules/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:2217, in ShapedArray.update(self, shape, dtype, weak_type, **kwargs)
   2215 if 'memory_space' not in kwargs:
   2216   kwargs['memory_space'] = self.memory_space
-> 2217 return ShapedArray(shape, dtype, weak_type, **kwargs)

TypeError: ShapedArray.__init__() got an unexpected keyword argument 'named_shape'

Tada! If you don’t have a GPU, jax will print out a warning. This isn’t a huge problem unless you’re trying to train a new model in which case this will be extremely slow.

grannules also has helper methods to evaluate the model power spectrum over a frequency array in utils.psd.

from grannules.utils.psd import PSD, nu_max
import numpy as np
import holoviews as hv
hv.extension("bokeh")

nu = np.logspace(
    start = np.log10(10),
    stop = np.log10(300),
    num = 200
)
nm = nu_max(X["M"], X["R"], X["Teff"]).values
p = PSD(nu, nm, y["H"], y["P"], y["tau"], y["alpha"])[0]

nn_only = hv.Curve(
    data = (nu, p),
    kdims = "Frequency (uHz)", vdims="Power ppm^2/uHz",
    label = "Neural Network Prediction"
).opts(
    width = 600, height = 600,
    logx = True, logy = True
)
nn_only

There is also functionality to use scaling relations on \(\nu_\mathrm{max}\), with values from de Assis Peralta et al. 2018 and from in-house fitting in utils.scalingrelations

from grannules.utils.scalingrelations import SRPredictor

sr_predictor = SRPredictor()
sr_y = sr_predictor.predict(nm, [2])
sr_p = PSD(nu, nm, sr_y["H"], sr_y["P"], sr_y["tau"], sr_y["alpha"])[0]

sr_only = hv.Curve(
    data = (nu, sr_p),
    kdims = "Frequency (uHz)", vdims="Power ppm^2/uHz",
    label = "Scaling Relation Prediction"
).opts(
    width = 600, height = 600,
    logx = True, logy = True
)
sr_only * nn_only

Better yet, let’s our predictions with actual data from Kepler. We’ll use KIC 757137, an arbitrary red giant.

from grannules.utils.scalingrelations import compare_psd
compare_psd(
    M = 1.55,
    R = 13.26,
    Teff = 4751,
    FeH = -0.08,
    phase = 1,
    KepMag = 9.196,
    KIC = 757137
)
Searching...
Downloading...
Processing lightcurve...
0    18802.64906
Name: P, dtype: float64

Given a new data set, grannules also has functionality to handily train a new model, albeit using a more limited set of network structures (see NNPredictor, utils.model, and utils.datatransform)