Repository: Aarhus-Psychiatry-Research/timeseriesflattener

Issue:
feat: take multiple features as long format

Discussion:

Patch:
diff --git a/src/timeseriesflattener/feature_spec_objects.py b/src/timeseriesflattener/feature_spec_objects.py
index 0ae0db2..e739fa9 100644
--- a/src/timeseriesflattener/feature_spec_objects.py
+++ b/src/timeseriesflattener/feature_spec_objects.py
@@ -12,7 +12,7 @@ from pydantic import BaseModel as PydanticBaseModel
 from pydantic import Extra
 
 from timeseriesflattener.resolve_multiple_functions import resolve_multiple_fns
-from timeseriesflattener.utils import data_loaders
+from timeseriesflattener.utils import data_loaders, split_df_dict
 
 log = logging.getLogger(__name__)
 
@@ -40,18 +40,12 @@ def in_dict_and_not_none(d: dict, key: str) -> bool:
     return key in d and d[key] is not None
 
 
-def resolve_values_df(data: dict[str, Any]):
-    """Resolve the values_df attribute to a dataframe."""
-    if "values_loader" not in data and "values_df" not in data:
-        raise ValueError("Either values_loader or a dataframe must be specified.")
-
-    if in_dict_and_not_none(d=data, key="values_loader") and in_dict_and_not_none(
-        key="values_df",
-        d=data,
-    ):
-        raise ValueError("Only one of values_loader or df can be specified.")
-
-    if "values_df" not in data or data["values_df"] is None:
+def resolve_from_dict_or_registry(data: dict[str, Any]):
+    """Resolve values_df from a dictionary or registry."""
+    if "values_name" in data and data["values_name"] is not None:
+        data["values_df"] = split_df_dict.get(data["values_name"])
+        data["feature_name"] = data["values_name"]
+    else:
         if isinstance(data["values_loader"], str):
             data["feature_name"] = data["values_loader"]
             data["values_loader"] = data_loaders.get(data["values_loader"])
@@ -66,6 +60,28 @@ def resolve_values_df(data: dict[str, Any]):
                 feature_name=data["feature_name"],
             )
 
+
+def resolve_values_df(data: dict[str, Any]):
+    """Resolve the values_df attribute to a dataframe."""
+    if not any(key in data for key in ["values_loader", "values_name", "values_df"]):
+        raise ValueError(
+            "Either values_loader or a dictionary containing dataframes or a single dataframe must be specified.",
+        )
+
+    if (
+        sum(
+            in_dict_and_not_none(data, key)
+            for key in ["values_loader", "values_name", "values_df"]
+        )
+        > 1
+    ):
+        raise ValueError(
+            "Only one of values_loader or values_name or df can be specified.",
+        )
+
+    if "values_df" not in data or data["values_df"] is None:
+        resolve_from_dict_or_registry(data)
+
     if not isinstance(data["values_df"], pd.DataFrame):
         raise ValueError("values_df must be or resolve to a pandas DataFrame.")
 
@@ -115,6 +131,10 @@ class AnySpec(BaseModel):
     # Loader for the df. Tries to resolve from the resolve_multiple_nfs registry,
     # then calls the function which should return a dataframe.
 
+    values_name: Optional[str] = None
+    # A string that corresponds to a key in a dictionary of multiple dataframes that
+    # correspods to a name of a type of values.
+
     loader_kwargs: Optional[dict[str, Any]] = None
     # Optional kwargs for the values_loader
 
@@ -187,7 +207,7 @@ class StaticSpec(AnySpec):
 
 
 class TemporalSpec(AnySpec):
-    """The minimum specification required for all collapsed time series
+    """The minimum specification required for all collapsed time series.
 
     (temporal features), whether looking ahead or behind.
 
@@ -301,7 +321,7 @@ class PredictorSpec(TemporalSpec):
 
 
 class OutcomeSpec(TemporalSpec):
-    """Specification for a single predictor, where the df has been resolved."""
+    """Specification for a single outcome, where the df has been resolved."""
 
     prefix: str = "outc"
 
@@ -357,17 +377,19 @@ class OutcomeSpec(TemporalSpec):
 
 
 class MinGroupSpec(BaseModel):
-    """Minimum specification for a group of features, whether they're looking
-
-    ahead or behind.
+    """Minimum specification for a group of features, whether they're looking ahead or behind.
 
     Used to generate combinations of features.
     """
 
-    values_loader: list[str]
-    # Loader for the df. Tries to resolve from the resolve_multiple_nfs registry,
+    values_loader: Optional[list[str]] = None
+    # Loader for the df. Tries to resolve from the data_loaders registry,
     # then calls the function which should return a dataframe.
 
+    values_name: Optional[list[str]] = None
+    # List of strings that corresponds to a key in a dictionary of multiple dataframes
+    # that correspods to a name of a type of values.
+
     values_df: Optional[pd.DataFrame] = None
     # Dataframe with the values.
 
@@ -387,12 +409,9 @@ class MinGroupSpec(BaseModel):
     # If NaN is higher than this in the input dataframe during resolution, raise an error.
 
     prefix: Optional[str] = None
-    # Prefix for the column name. Overrides the default prefix for the feature type.
-
-    def __init__(self, **data):
-        super().__init__(**data)
 
-        # Check that all passed loaders are valid
+    def _check_loaders_are_valid(self):
+        """Check that all loaders can be resolved from the data_loaders catalogue."""
         invalid_loaders = list(
             set(self.values_loader) - set(data_loaders.get_all().keys()),
         )
@@ -415,6 +434,15 @@ class MinGroupSpec(BaseModel):
                 f"""Available loaders:{nl}{avail_loaders_str}""",
             )
 
+    # Prefix for the column name. Overrides the default prefix for the feature type.
+
+    def __init__(self, **data):
+        super().__init__(**data)
+
+        # Check that all passed loaders are valid
+        if self.values_loader is not None:
+            self._check_loaders_are_valid()
+
         if self.output_col_name_override:
             input_col_name = (
                 "value"
diff --git a/src/timeseriesflattener/utils.py b/src/timeseriesflattener/utils.py
index a602b18..2ef4adb 100644
--- a/src/timeseriesflattener/utils.py
+++ b/src/timeseriesflattener/utils.py
@@ -14,10 +14,56 @@ import catalogue
 import pandas as pd
 
 data_loaders = catalogue.create("timeseriesflattener", "data_loaders")
+split_df_dict = {}
+
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
 
 
+def split_df_and_register_to_dict(
+    df: pd.DataFrame,
+    id_col_name: str = "dw_ek_borger",
+    timestamp_col_name: str = "timestamp",
+    value_col_name: str = "value",
+    value_names_col_name: str = "value_names",
+):
+    """Split a df with multiple different value types into dataframes only containing values for each specific value type.
+
+    Registers the seperated dfs in the df_dict.
+
+    Args:
+        df (pd.DataFrame): A dataframe in long format containing the values to be grouped into a catalogue.
+        id_col_name (str): Name of the column containing the patient id for each value. Defaults to "dw_ek_borger".
+        timestamp_col_name (str): Name of the column containing the timestamp for each value. Defaults to "timestamp".
+        value_col_name (str): Name of the column containing the value for each value. Defaults to "values".
+        value_names_col_name (str): Name of the column containing the names of the different types of values for each value. Defaults to "value_names".
+    """
+    passed_columns = [
+        id_col_name,
+        timestamp_col_name,
+        value_col_name,
+        value_names_col_name,
+    ]
+    missing_columns = [col for col in passed_columns if col not in list(df.columns)]
+
+    # If any of the required columns is missing, raise an error
+    if len(missing_columns) > 0:
+        raise ValueError(
+            f"The following required column(s) is/are missing from the input dataframe: {missing_columns}. Available columns are {df.columns}.",
+        )
+
+    # Get the unique value names from the dataframe
+    value_names = df[value_names_col_name].unique()
+
+    for value_name in value_names:
+
+        value_df = df[df[value_names_col_name] == value_name][
+            [id_col_name, timestamp_col_name, value_col_name]
+        ]
+
+        split_df_dict[value_name] = value_df
+
+
 def format_dict_for_printing(d: dict) -> str:
     """Format a dictionary for printing. Removes extra apostrophes, formats
 


Test patch:
diff --git a/src/timeseriesflattener/testing/utils_for_testing.py b/src/timeseriesflattener/testing/utils_for_testing.py
index 0970c9d..0b08985 100644
--- a/src/timeseriesflattener/testing/utils_for_testing.py
+++ b/src/timeseriesflattener/testing/utils_for_testing.py
@@ -15,6 +15,7 @@ from timeseriesflattener.feature_spec_objects import AnySpec
 from timeseriesflattener.testing.load_synth_data import (
     load_synth_outcome,
     load_synth_prediction_times,
+    synth_predictor_binary,
 )
 from timeseriesflattener.utils import data_loaders
 
@@ -155,4 +156,22 @@ def synth_predictor():
 @pytest.fixture(scope="function")
 def synth_outcome():
     """Load the synth outcome times."""
-    return load_synth_outcome(n_rows=1_000)
+    return load_synth_outcome()
+
+
+@pytest.fixture(scope="function")
+def long_df():
+    """Create a long df."""
+    synth_df = synth_predictor_binary()
+    synth_df = synth_df.rename(columns={"value": "value_name_1"})
+    synth_df["value_name_2"] = synth_df["value_name_1"]
+
+    df = pd.melt(
+        synth_df,
+        id_vars=["dw_ek_borger", "timestamp"],
+        value_vars=["value_name_1", "value_name_2"],
+        var_name="value_names",
+        value_name="value",
+    )
+
+    return df
diff --git a/tests/test_timeseriesflattener/test_feature_spec_objects.py b/tests/test_timeseriesflattener/test_feature_spec_objects.py
index 3f0bf25..d04028d 100644
--- a/tests/test_timeseriesflattener/test_feature_spec_objects.py
+++ b/tests/test_timeseriesflattener/test_feature_spec_objects.py
@@ -4,14 +4,14 @@ import pytest
 
 from timeseriesflattener.feature_spec_objects import (
     AnySpec,
+    PredictorGroupSpec,
     check_that_col_names_in_kwargs_exist_in_df,
 )
 from timeseriesflattener.testing.load_synth_data import (  # pylint: disable=unused-import
     load_synth_predictor_float,
 )
-
-# Avoid ruff removing as unused
-used_loaders = [load_synth_predictor_float]
+from timeseriesflattener.testing.utils_for_testing import long_df
+from timeseriesflattener.utils import split_df_and_register_to_dict
 
 
 def test_anyspec_init():
@@ -38,6 +38,17 @@ def test_loader_kwargs():
     assert len(spec.values_df) == 10
 
 
+def test_invalid_multiple_data_args():
+    """Test that error is raised if multiple data args are passed."""
+
+    with pytest.raises(ValueError, match=r".*nly one of.*"):
+        AnySpec(
+            values_loader="synth_predictor_float",
+            values_name="synth_data",
+            prefix="test",
+        )
+
+
 def test_anyspec_incorrect_values_loader_str():
     """Raise error if values loader is not a key in registry."""
     with pytest.raises(ValueError, match=r".*in registry.*"):
@@ -57,3 +68,21 @@ def test_that_col_names_in_kwargs_exist_in_df():
     data = {"col_name_1": "A", "col_name_2": "D", "values_df": df}
     with pytest.raises(ValueError, match="D is not in df"):
         check_that_col_names_in_kwargs_exist_in_df(data=data, df=df)
+
+
+def test_create_combinations_while_resolving_from_registry(long_df: pd.DataFrame):
+    """Test that split_df_and_register_to_dict resolves correctly when multiple dataframes are fetched."""
+
+    split_df_and_register_to_dict(df=long_df)
+
+    group_spec = PredictorGroupSpec(
+        values_name=[
+            "value_name_1",
+            "value_name_2",
+        ],
+        resolve_multiple_fn=["mean"],
+        lookbehind_days=[30],
+        fallback=[0],
+    ).create_combinations()
+
+    assert len(group_spec) == 2
diff --git a/tests/test_timeseriesflattener/test_utils.py b/tests/test_timeseriesflattener/test_utils.py
new file mode 100644
index 0000000..a7ce88d
--- /dev/null
+++ b/tests/test_timeseriesflattener/test_utils.py
@@ -0,0 +1,32 @@
+"""Test that feature spec objects work as intended."""
+import pandas as pd
+import pytest
+
+from timeseriesflattener.feature_spec_objects import AnySpec
+from timeseriesflattener.testing.load_synth_data import synth_predictor_binary  # noqa
+from timeseriesflattener.testing.utils_for_testing import long_df
+from timeseriesflattener.utils import split_df_and_register_to_dict, split_df_dict
+
+
+def test_split_df_and_register_in_dict(long_df: pd.DataFrame):
+    """Test that the split_df_and_register_to_dict function works as intended."""
+
+    split_df_and_register_to_dict(df=long_df)
+
+    assert len(split_df_dict) == 2
+    assert split_df_dict["value_name_1"].shape == (10000, 3)
+    assert split_df_dict["value_name_2"].shape == (10000, 3)
+
+
+def test_resolve_from_df_dict(long_df: pd.DataFrame):
+    """Test that a split_df_and_register_to_dict resolves from the  correctly."""
+
+    split_df_and_register_to_dict(df=long_df)
+
+    spec = AnySpec(
+        values_name="value_name_1",
+        feature_name="test",
+        prefix="test",
+    )
+
+    assert len(spec.values_df) == 10000