POSYDON-code / POSYDON

POSYDON is a next-generation single and binary-star population synthesis code incorporating full stellar structure and evolution modeling with the use of MESA.
BSD 3-Clause "New" or "Revised" License
25 stars 19 forks source link

IF interpolator `out_nan_keys` as an input #342

Open maxbriel opened 1 week ago

maxbriel commented 1 week ago

See Issue #341

Adds the out_nan_keys to the training parameters and as an input to the BaseIFInterpolator.__init__.

maxbriel commented 1 week ago

Test training:

from posydon.grids.psygrid import PSyGrid
import numpy as np
from posydon.interpolation.IF_interpolation import IFInterpolator
from posydon.grids.MODELS import MODELS

# example file
interp_path = './CO_HeMS_RLO_1e+00_Zsun.h5'

grid = PSyGrid(interp_path)
interpolation_method = 'linear_RLO'
grid_type = 'CO-HeMS'

method = interpolation_method
if '_RLO' in method:
    method = method.split('_RLO')[0]

interp_method = [method, method]
interp_classes = ["stable_MT", "unstable_MT"]
if '_RLO' not in interpolation_method:
    interp_method += [method]
    interp_classes += ["no_MT"]
if 'CO' not in grid_type:
    interp_method += [method]
    interp_classes += ["stable_reverse_MT"]

out_keys = [key for key in grid.final_values.dtype.names if (
                                key != "model_number" and
                                (type(grid.final_values[key][0]) != np.str_)
                                and any(~np.isnan(grid.final_values[key]))
                                and "MODEL" not in key)]

out_nan_keys = [key for key in grid.final_values.dtype.names if 
                         (key != "model_number" and
                         (type(grid.final_values[key][0]) != np.str_)
                         and all(np.isnan(grid.final_values[key]))
                         and "MODEL" not in key)
                    ]

c_keys = ['interpolation_class', 'S1_state', 'S2_state', 'mt_history']

interpolators = [
    {
        "interp_method": interp_method,
        "interp_classes": interp_classes,
        "out_keys": out_keys,
        "out_nan_keys" : out_nan_keys,
        "class_method": "kNN",
        "c_keys": c_keys,
        "c_key": "interpolation_class"
    },
]
interp = IFInterpolator(grid=grid, interpolators=interpolators)
interp.train()

for i in interp.interpolators:
    print("Out keys:", i.out_keys)
    print("NaN keys:", i.out_nan_keys)
maxbriel commented 1 week ago

One additional change here:

I've removed the c_keys for f'S{i}_{MODEL_NAME}_SN_type' and f'S{i}_{MODEL_NAME}_CO_type' from when we're training IFinterpolators for each MODEL_NAME core-collapse quantities.

The c_keys at line 401 already contains the different f'S{i}_{MODEL_NAME}_CO_type' and f'S{i}_{MODEL_NAME}_SN_type'.

Since the classifier is independent of the c_key class, there is no need to include it again, when we're creating the IFinterpolators for the core-collapse MODELS. These classifiers are already being trained. There's no difference in the training between them.

maxbriel commented 1 week ago

I've successfully ran a population with this. The SN_type + CO_state seems to match; also as an affect of the forced check in the evolution we implemented in PR #324. I did notice other issues but not related to this one.