Files
temperature-based-fertility…/code/utils/dataset_creation.py
T
2025-09-10 10:37:55 +02:00

267 lines
10 KiB
Python

import os
import lmdb
import numpy as np
import pandas as pd
import torch
from sklearn.base import BaseEstimator
from sklearn.preprocessing import StandardScaler, MinMaxScaler
import joblib
from utils.lmdb_utils import load_from_lmdb, get_lmdb_keys
def get_features(cycle: dict,
feature_config: dict) -> dict:
"""
Computes the feature for a given cycle defined by the feature_config
:param cycle: cycle data as dictionary
:param feature_config: feature config as dictionary
:return: dict with feature according feature config
"""
features = {}
feature_sets = feature_config["feature_sets"]
for feature_set in feature_sets:
for feature_def in feature_config[feature_set]:
if feature_set not in features:
features[feature_set] = {}
# skip features that are marked as ignored
if "ignored_features" in feature_config and feature_def["name"] in feature_config["ignored_features"]:
continue
feature_return = feature_def["fn"](cycle=cycle)
if isinstance(feature_return, dict):
if len(feature_return) > 1:
for feature_name, feature in feature_return.items():
features[feature_set][f"{feature_def['name']}_{feature_name}"] = feature
else:
features[feature_set][feature_def["name"]] = list(feature_return.values())[0]
else:
features[feature_set][feature_def["name"]] = feature_return
# check, if all features have same length
if feature_set in features and len(features[feature_set]) > 0 and isinstance(
list(features[feature_set].values())[0], np.ndarray):
feature_lengths = [len(x) for x in features[feature_set].values()]
if len(set(feature_lengths)) > 1:
raise ValueError(f"Feature set {feature_set} has features of different lengths: {feature_lengths}")
return features
def save_scalers(scalers: dict,
scaler_dir: str):
if not os.path.exists(scaler_dir):
os.makedirs(scaler_dir)
for feature_type in scalers:
for feature_name in scalers[feature_type]:
joblib.dump(scalers[feature_type][feature_name],
f"{scaler_dir}/{feature_name}.pkl")
def load_scalers(scaler_dir: str) -> dict:
scalers = dict()
for scaler_file in os.listdir(scaler_dir):
if scaler_file.endswith(".pkl"):
feature_name = scaler_file.replace(".pkl", "")
scalers[feature_name] = joblib.load(f"{scaler_dir}/{scaler_file}")
return scalers
def get_scalers_for_model(model_configuration: dict):
"""
Get the scalers for the features of a model configuration
Args:
model_configuration: configuration of the model
Returns:
scalers: scalers for the model
"""
scaler_dir = os.path.join(model_configuration["feature_config"]["dataset_dir"], "scalers")
scalers = load_scalers(scaler_dir)
return scalers
def get_feature_values(feature_type: str | None,
feature_name: str,
env: lmdb.Environment,
keys: list = None) -> list:
"""
Get the values of a feature from the lmdb dataset for all keys
:param feature_type: type of feature, can be None, then the first feature with the given name will be used
:param feature_name: feature name to extract
:param env: lmdb environment
:param keys: keys to extract the feature from, if None, all keys will be used / fetched from the database
:return: list with feature values
"""
if keys is None:
keys = get_lmdb_keys(env)
feature_values = []
for key in keys:
data = load_from_lmdb(env, key)
if feature_type is None:
for feature_type, feature_data in data.items():
if feature_name in feature_data:
feature_values.append(feature_data[feature_name])
break
else:
for current_feature_name, feature_data in data[feature_type].items():
# catch sub features that have been prefixed with the feature name
if current_feature_name.startswith(feature_name):
feature_values.append(feature_data)
break
return feature_values
def _train_scalers(feature_type: str,
feature_name: str,
scaler_type,
sample,
env) -> dict:
scalers = dict()
all_features = sample[feature_type]
individual_feature_names = list()
for individual_feature_name in all_features:
if individual_feature_name.startswith(feature_name) and not individual_feature_name.endswith("_scaled"):
individual_feature_names.append(individual_feature_name)
for individual_feature_name in individual_feature_names:
feature_values = get_feature_values(feature_type, individual_feature_name, env)
if len(feature_values) == 0:
raise ValueError(f"No feature values found for feature {individual_feature_name}")
if isinstance(feature_values[0], list) or isinstance(feature_values[0], np.ndarray):
features_reshaped = np.concatenate(feature_values).reshape(-1, 1)
else:
features_reshaped = np.array(feature_values).reshape(-1, 1)
del feature_values
if scaler_type is not None:
scaler = scaler_type()
scaler.fit(features_reshaped)
else:
scaler = None
scalers[individual_feature_name] = scaler
return scalers
def train_scalers(feature_type: str,
feature_name: str,
feature_config: dict,
env) -> BaseEstimator | None:
feature_values = get_feature_values(feature_type, feature_name, env)
if len(feature_values) == 0:
raise ValueError(f"No feature values found for feature {feature_name}")
if isinstance(feature_values[0], list) or isinstance(feature_values[0], np.ndarray):
features_reshaped = np.concatenate(feature_values).reshape(-1, 1)
else:
features_reshaped = np.array(feature_values).reshape(-1, 1)
del feature_values
current_feature_config = feature_config[feature_type][feature_name]
scaler_type = current_feature_config.get("scaler", None)
if scaler_type is not None:
scaler = scaler_type()
scaler.fit(features_reshaped)
else:
scaler = None
return scaler
def scale_item(data: dict,
subset_name: str,
scalers: dict):
data_format = dict()
for feature_set in data:
if feature_set not in data_format:
data_format[feature_set] = dict()
for feature_name in data[feature_set]:
if feature_name.endswith("_scaled"):
continue
data_format[feature_set][feature_name] = data[feature_set][feature_name]
for feature_set in data_format:
for feature_name in data_format[feature_set]:
feature_values = data[feature_set][feature_name]
subset_feature_name = f"{subset_name}_{feature_name}"
if subset_feature_name in scalers and scalers[subset_feature_name] is not None:
if isinstance(feature_values, list) or isinstance(feature_values, np.ndarray):
scaled_feature_values = scalers[subset_feature_name].transform(
feature_values.reshape(-1, 1)).flatten()
else:
scaled_feature_values = scalers[subset_feature_name].transform(
np.array(feature_values).reshape(-1, 1)).flatten()
else:
scaled_feature_values = feature_values
data[feature_set][f"{feature_name}_scaled"] = scaled_feature_values
return data
def inverse_scale_feature(input_feature: np.ndarray | int | float | torch.Tensor,
feature_names: np.ndarray | list | str,
scalers: dict,
subset_name: str = "train") -> np.ndarray:
"""
Inverse scales the input features, can handle single and multi feature input
Args:
input_feature: input feature as a numpy array
feature_names: names of input features as reference for scalers
scalers: dict of feature scalers
subset_name: subset name, either train, val or test for proper inverse scaling
Returns:
scaled input features as numpy array
"""
if isinstance(feature_names, str):
feature_names = [feature_names]
if isinstance(input_feature, np.ndarray) or isinstance(input_feature, torch.Tensor):
input_dim = input_feature.shape[1] if len(input_feature.shape) > 1 else 1
elif isinstance(input_feature, float) or isinstance(input_feature, int):
# handle case, where input is scalar
input_scaled = scalers[f"{subset_name}_{feature_names[0]}"].inverse_transform([[input_feature]])[0][0]
return input_scaled
elif isinstance(input_feature, list):
input_feature = np.array(input_feature)
input_dim = input_feature.shape[1] if len(input_feature.shape) > 1 else 1
else:
raise ValueError(f"Unsupported input type: {type(input_feature)}")
# reshape, if input is one dimensional
if input_dim == 1:
input_feature = input_feature.reshape(-1, 1)
input_scaled = input_feature.copy()
for i in range(input_dim):
input_scaled[:, i] = scalers[f"{subset_name}_{feature_names[i]}"].inverse_transform(
np.array([input_scaled[:, i]]))
return input_scaled
def combine_features(cycles: list[dict], feature_config: dict) -> dict:
combined_features = dict()
for feature_set in feature_config["feature_sets"]:
combined_features[feature_set] = dict()
if feature_set not in cycles[0]:
continue
features_in_set = cycles[0][feature_set].keys()
for feature in features_in_set:
feature_values = [cycles[i][feature_set][feature] for i in range(len(cycles))]
if isinstance(feature_values[0], dict):
for key in feature_values[0].keys():
feature_array = np.concatenate([feature_values[i][key] for i in range(len(feature_values))])
combined_features[feature_set][key] = feature_array
else:
feature_array = np.concatenate(feature_values)
combined_features[feature_set][feature] = feature_array
return combined_features