added code

code/utils/dataset_creation.py

@@ -4,6 +4,7 @@ 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
 
@@ -115,11 +116,11 @@ def get_feature_values(feature_type: str | None,
     return feature_values
 
 
-def train_scalers(feature_type: str,
-                  feature_name: str,
-                  scaler_type,
-                  sample,
-                  env) -> dict:
+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()
@@ -147,7 +148,33 @@ def train_scalers(feature_type: str,
     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:
@@ -161,12 +188,13 @@ def scale_item(data: dict,
     for feature_set in data_format:
         for feature_name in data_format[feature_set]:
             feature_values = data[feature_set][feature_name]
-            if scalers[feature_name] is not None:
+            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[feature_name].transform(
+                    scaled_feature_values = scalers[subset_feature_name].transform(
                         feature_values.reshape(-1, 1)).flatten()
                 else:
-                    scaled_feature_values = scalers[feature_name].transform(
+                    scaled_feature_values = scalers[subset_feature_name].transform(
                         np.array(feature_values).reshape(-1, 1)).flatten()
             else:
                 scaled_feature_values = feature_values
@@ -177,13 +205,15 @@ def scale_item(data: dict,
 
 def inverse_scale_feature(input_feature: np.ndarray | int | float | torch.Tensor,
                           feature_names: np.ndarray | list | str,
-                          scalers: dict) -> np.ndarray:
+                          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
@@ -196,7 +226,7 @@ def inverse_scale_feature(input_feature: np.ndarray | int | float | 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[feature_names[0]].inverse_transform([[input_feature]])[0][0]
+        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)
@@ -210,7 +240,8 @@ def inverse_scale_feature(input_feature: np.ndarray | int | float | torch.Tensor
 
     input_scaled = input_feature.copy()
     for i in range(input_dim):
-        input_scaled[:, i] = scalers[feature_names[i]].inverse_transform(np.array([input_scaled[:, i]]))
+        input_scaled[:, i] = scalers[f"{subset_name}_{feature_names[i]}"].inverse_transform(
+            np.array([input_scaled[:, i]]))
 
     return input_scaled