1""" Extra constraints based on e.g. surface properties of the star, or a binary companion. """
  2
  3import logging
  4import sys
  5from functools import partial
  6from pathlib import Path
  7
  8import numpy as np
  9import pandas as pd
 10
 11logger = logging.getLogger("logger.ac")
 12
 13
 14################################################################################
 15def surface_constraint(
 16    merit_values_file,
 17    observations_file=None,
 18    nsigma=3,
 19    constraint_companion=None,
 20    isocloud_grid_summary=None,
 21    surface_grid_file=None,
 22    free_parameters=["Z", "M", "logD", "aov", "fov", "Xc"],
 23    evolution_parameter="Xc",
 24    evolution_step=-1e-2,
 25):
 26    """
 27    Enforce an n-sigma constraint on the models based on the surface observables.
 28    Save this as a file with prefix indicating how many sigma the error box was.
 29
 30    Parameters
 31    ----------
 32    merit_values_file: string
 33        Path to the hdf5 files with the merit function values and the surface info of the models in the grid.
 34    observations_file: string
 35        Path to the tsv file with observations, with a column for each observable and each set of errors.
 36        Column names specify the observable, and "_err" suffix denotes that it's the error.
 37    nsigma: int
 38        How many sigma you want to make the interval to accept models.
 39    constraint_companion: dict
 40        Information on the companion star. Set to None to model single stars,
 41        or provide this to include binary constraints using isochrone-clouds.
 42    isocloud_grid_summary: dict
 43        Nested dictionary, the keys at its two levels are metallicity and mass.
 44        Holds the surface properties of the grid for the isochrone-cloud modelling per combination of metallicity-mass.
 45    surface_grid_file: string
 46        File with the surface properties and ages of the model-grid.
 47    free_parameters: list of strings
 48        List of all the parameters varied in the model grid.
 49    evolution_parameter: string
 50        Name of the parameter that is used to track the evolutionary steps of the model.
 51    evolution_step: float
 52        Change in the evolutionary parameter from one step to the next (negative if quantity decreases, e.g. central hydrogen content Xc)
 53    """
 54    obs_dataframe = pd.read_table(observations_file, sep="\s+", header=0, index_col="index")
 55    dataframe_theory = pd.read_hdf(merit_values_file)
 56
 57    if "Teff" in obs_dataframe.columns:
 58        dataframe_theory = dataframe_theory[
 59            dataframe_theory.logTeff
 60            < np.log10(obs_dataframe["Teff"].iloc[0] + nsigma * obs_dataframe["Teff_err"].iloc[0])
 61        ]
 62        dataframe_theory = dataframe_theory[
 63            dataframe_theory.logTeff
 64            > np.log10(obs_dataframe["Teff"].iloc[0] - nsigma * obs_dataframe["Teff_err"].iloc[0])
 65        ]
 66    if "logg" in obs_dataframe.columns:
 67        dataframe_theory = dataframe_theory[
 68            dataframe_theory.logg < obs_dataframe["logg"].iloc[0] + nsigma * obs_dataframe["logg_err"].iloc[0]
 69        ]
 70        dataframe_theory = dataframe_theory[
 71            dataframe_theory.logg > obs_dataframe["logg"].iloc[0] - nsigma * obs_dataframe["logg_err"].iloc[0]
 72        ]
 73    if "logL" in obs_dataframe.columns:
 74        dataframe_theory = dataframe_theory[
 75            dataframe_theory.logL < obs_dataframe["logL"].iloc[0] + nsigma * obs_dataframe["logL_err"].iloc[0]
 76        ]
 77        dataframe_theory = dataframe_theory[
 78            dataframe_theory.logL > obs_dataframe["logL"].iloc[0] - nsigma * obs_dataframe["logL_err"].iloc[0]
 79        ]
 80
 81    if constraint_companion is not None:
 82        if (isocloud_grid_summary is None) or (surface_grid_file is None):
 83            logger.error(
 84                "Please supply a summary file for the isocloud grid and a path to the file with the grid surface properties and ages."
 85            )
 86            sys.exit()
 87
 88        surface_grid_dataframe = pd.read_hdf(surface_grid_file)
 89
 90        func = partial(
 91            enforce_binary_constraints,
 92            constraint_companion=constraint_companion,
 93            isocloud_grid_summary=isocloud_grid_summary,
 94            nsigma=nsigma,
 95            surface_grid_dataframe=surface_grid_dataframe,
 96            free_parameters=free_parameters,
 97            evolution_parameter=evolution_parameter,
 98            evolution_step=evolution_step,
 99        )
100        indices_to_drop = dataframe_theory.apply(func, axis=1)
101        for index_to_drop in indices_to_drop:
102            # Check if index_to_drop equals itself to filter out NaN values
103            if (index_to_drop is not None) and (index_to_drop == index_to_drop):
104                dataframe_theory.drop(index_to_drop, inplace=True)
105
106    output_file = f"{nsigma}sigmaBox_{merit_values_file}"
107    Path(output_file).parent.mkdir(parents=True, exist_ok=True)
108    dataframe_theory.to_hdf(path_or_buf=output_file, key="surface_constrained_models", format="table", mode="w")
109
110
111################################################################################
112def get_age(
113    model, df, free_parameters=["Z", "M", "logD", "aov", "fov", "Xc"], evolution_parameter="Xc", evolution_step=-1e-2
114):
115    """
116    Get the age of the models one step older and younger than the provided model.
117
118    Parameters
119    ----------
120    model: pandas series
121        Parameters of the model.
122    df: pandas dataFrame
123        Dataframe with the model parameters and age (and surface info) of the theoretical models.
124    free_parameters: list of strings
125        List of all the parameters varied in the model grid.
126    evolution_parameter: string
127        Name of the parameter that is used to track the evolutionary steps of the model.
128    evolution_step: float
129        Change in the evolutionary parameter from one step to the next (negative if quantity decreases, e.g. central hydrogen content Xc)
130
131    Returns
132    ----------
133    min_age, max_age: tuple of integers
134        Age of the model one step younger and older than the provided model,
135        these are the minimum and maximum age to accept models in the isochrone-cloud.
136    """
137
138    # copy to prevent deletion in the list outside this function
139    params = list(free_parameters)
140    params.remove(evolution_parameter)
141    for param in params:
142        df = df.loc[np.isclose(getattr(df, param), getattr(model, param))]
143
144    model_evolution_attr = getattr(model, evolution_parameter)
145    grid_evolution_attr = getattr(df, evolution_parameter)
146
147    if abs(model_evolution_attr - max(grid_evolution_attr)) < abs(0.5 * evolution_step):
148        min_age = 0
149        max_age = int((df.loc[np.isclose(grid_evolution_attr, model_evolution_attr + evolution_step)].age).iloc[0])
150    elif abs(model_evolution_attr - min(grid_evolution_attr)) < abs(0.5 * evolution_step):
151        min_age = int((df.loc[np.isclose(grid_evolution_attr, model_evolution_attr - evolution_step)].age).iloc[0])
152        age = int((df.loc[np.isclose(grid_evolution_attr, model_evolution_attr)].age).iloc[0])
153        max_age = age + age - min_age
154    else:
155        min_age = int((df.loc[np.isclose(grid_evolution_attr, model_evolution_attr - evolution_step)].age).iloc[0])
156        max_age = int((df.loc[np.isclose(grid_evolution_attr, model_evolution_attr + evolution_step)].age).iloc[0])
157    return min_age, max_age
158
159
160################################################################################
161def enforce_binary_constraints(
162    df_theory_row,
163    constraint_companion=None,
164    isocloud_grid_summary=None,
165    nsigma=3,
166    surface_grid_dataframe=None,
167    free_parameters=["Z", "M", "logD", "aov", "fov", "Xc"],
168    evolution_parameter="Xc",
169    evolution_step=-1e-2,
170):
171    """
172    Enforce an n-sigma constraint on the models based on
173    spectroscopic observations of the binary companion employing isochrone-clouds.
174    Assumes the same metallicity 'Z' for both primary and secondary,
175    masses 'M' compatible with observed mass ratio 'q', and ages similar within 1 grid step.
176
177    Parameters
178    ----------
179    df_theory_row: tuple, made of (int, pandas series)
180        tuple returned from pandas.iterrows(), first tuple entry is the row index of the pandas dataFrame
181        second tuple entry is a pandas series, containing a row from the pandas dataFrame.
182        (This row holds model parameters, the merit function value, and surface properties.)
183    constraint_companion: dict
184        Information on the companion star, including surface parameters, mass ratio (q), the errors,
185        and a boolean indicating whether the primary or secondary star is assumed pulsating and hence being modelled.
186    isocloud_grid_summary: dict
187        Nested dictionary, the keys at its two levels are metallicity and mass.
188        Holds the surface properties of the grid for the isochrone-cloud modelling per combination of metallicity-mass.
189    nsigma: int
190        How many sigma you want to make the interval to accept models.
191    surface_grid_dataframe: pandas DataFrame
192        DataFrame with the surface properties and ages of the model-grid.
193    free_parameters: list of strings
194        List of all the parameters varied in the model grid.
195    evolution_parameter: string
196        Name of the parameter that is used to track the evolutionary steps of the model.
197    evolution_step: float
198        Change in the evolutionary parameter from one step to the next (negative if quantity decreases, e.g. central hydrogen content Xc)
199
200    Returns
201    ----------
202    index: int or None
203        Index of the dataframe that needs to be removed if binary constraints do not allow the model to remain.
204        Returns None if the binary constraints do not discard the model.
205    """
206    model = df_theory_row
207    min_age, max_age = get_age(
208        model,
209        surface_grid_dataframe,
210        free_parameters=free_parameters,
211        evolution_parameter=evolution_parameter,
212        evolution_step=evolution_step,
213    )
214    q = constraint_companion["q"]
215    q_err = constraint_companion["q_err"]
216    if constraint_companion["primary_pulsates"]:
217        m2_min = round(model.M * (q - q_err), 1)
218        m2_max = round(model.M * (q + q_err), 1)
219    else:
220        m2_min = round(model.M / (q + q_err), 1)
221        m2_max = round(model.M / (q - q_err), 1)
222
223    isocloud_dict = isocloud_grid_summary[f"{model.Z}"]
224    for key_mass, df in zip(isocloud_dict.keys(), isocloud_dict.values()):
225        # Convert from string to float for the comparisons
226        key_mass = float(key_mass)
227        if key_mass < m2_min or key_mass > m2_max:
228            # Only keep models that fall within mass range
229            continue
230        else:
231            df = df[(df.star_age < max_age) & (df.star_age > min_age)]
232
233            if df.shape[0] == 0:
234                continue
235
236            # Check for all provided constraints if the track passes through the uncertainty region
237            if constraint_companion["Teff"] is not None:
238                df = df[
239                    df.log_Teff < np.log10(constraint_companion["Teff"] + nsigma * constraint_companion["Teff_err"])
240                ]
241                df = df[
242                    df.log_Teff > np.log10(constraint_companion["Teff"] - nsigma * constraint_companion["Teff_err"])
243                ]
244            if constraint_companion["logg"] is not None:
245                df = df[df.log_g < constraint_companion["logg"] + nsigma * constraint_companion["logg_err"]]
246                df = df[df.log_g > constraint_companion["logg"] - nsigma * constraint_companion["logg_err"]]
247            if constraint_companion["logL"] is not None:
248                df = df[df.log_L < constraint_companion["logL"] + nsigma * constraint_companion["logL_err"]]
249                df = df[df.log_L > constraint_companion["logL"] - nsigma * constraint_companion["logL_err"]]
250            if df.shape[0] > 0:
251                # If some models fall within the constraints, return None to not remove the model.
252                return None
253
254    return model.name