Repository: Aarhus-Psychiatry-Research/timeseriesflattener

Issue:
Refactor: Main class TSFLattener, add df_getter method

Discussion:

Patch:
diff --git a/src/timeseriesflattener/__init__.py b/src/timeseriesflattener/__init__.py
index c4eae1f..1516176 100644
--- a/src/timeseriesflattener/__init__.py
+++ b/src/timeseriesflattener/__init__.py
@@ -1,2 +1,2 @@
 """Init timeseriesflattener."""
-from .flattened_dataset import FlattenedDataset
+from .flattened_dataset import TimeseriesFlattener
diff --git a/src/timeseriesflattener/flattened_dataset.py b/src/timeseriesflattener/flattened_dataset.py
index 165a39e..cc200ac 100644
--- a/src/timeseriesflattener/flattened_dataset.py
+++ b/src/timeseriesflattener/flattened_dataset.py
@@ -31,7 +31,7 @@ from timeseriesflattener.utils import load_dataset_from_file, write_df_to_file
 ProgressBar().register()
 
 
-class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
+class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
     """Turn a set of time-series into tabular prediction-time data.
 
     Attributes:
@@ -111,18 +111,18 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
         if "value" in prediction_times_df.columns:
             prediction_times_df.drop("value", axis=1, inplace=True)
 
-        self.df = prediction_times_df
+        self._df = prediction_times_df
 
         ValidateInitFlattenedDataset(
-            df=self.df,
+            df=self._df,
             timestamp_col_name=self.timestamp_col_name,
             id_col_name=self.id_col_name,
         ).validate_dataset()
 
         # Create pred_time_uuid_columne
-        self.df[self.pred_time_uuid_col_name] = self.df[self.id_col_name].astype(
+        self._df[self.pred_time_uuid_col_name] = self._df[self.id_col_name].astype(
             str,
-        ) + self.df[self.timestamp_col_name].dt.strftime("-%Y-%m-%d-%H-%M-%S")
+        ) + self._df[self.timestamp_col_name].dt.strftime("-%Y-%m-%d-%H-%M-%S")
 
     def _load_most_recent_df_matching_pattern(
         self,
@@ -179,7 +179,7 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
 
         # Expand fallback column
         df = pd.merge(
-            left=self.df[self.pred_time_uuid_col_name],
+            left=self._df[self.pred_time_uuid_col_name],
             right=df,
             how="left",
             on=self.pred_time_uuid_col_name,
@@ -245,7 +245,7 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
         else:
             raise ValueError(f"Unknown output_spec type {type(output_spec)}")
 
-        df = FlattenedDataset._drop_records_outside_interval_days(
+        df = TimeseriesFlattener._drop_records_outside_interval_days(
             df,
             direction=direction,
             interval_days=output_spec.interval_days,
@@ -254,7 +254,7 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
         )
 
         # Add back prediction times that don't have a value, and fill them with fallback
-        df = FlattenedDataset._add_back_prediction_times_without_value(
+        df = TimeseriesFlattener._add_back_prediction_times_without_value(
             df=df,
             pred_times_with_uuid=prediction_times_with_uuid_df,
             pred_time_uuid_colname=pred_time_uuid_col_name,
@@ -262,7 +262,7 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
 
         df["timestamp_val"].replace({output_spec.fallback: pd.NaT}, inplace=True)
 
-        df = FlattenedDataset._resolve_multiple_values_within_interval_days(
+        df = TimeseriesFlattener._resolve_multiple_values_within_interval_days(
             resolve_multiple=output_spec.resolve_multiple_fn,
             df=df,
             pred_time_uuid_colname=pred_time_uuid_col_name,
@@ -308,10 +308,10 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
                     f"{value_col_str[20]}, {n_trials}: Generated_df was all fallback values, regenerating",
                 )
 
-            n_to_generate = int(min(n_to_generate, len(self.df)))
+            n_to_generate = int(min(n_to_generate, len(self._df)))
 
             generated_df = self._flatten_temporal_values_to_df(
-                prediction_times_with_uuid_df=self.df.sample(
+                prediction_times_with_uuid_df=self._df.sample(
                     n=n_to_generate,
                     replace=False,
                 ),
@@ -470,7 +470,7 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
             msg.info("No cache dir specified, not attempting load")
 
         df = self._flatten_temporal_values_to_df(
-            prediction_times_with_uuid_df=self.df[
+            prediction_times_with_uuid_df=self._df[
                 [
                     self.pred_time_uuid_col_name,
                     self.id_col_name,
@@ -553,7 +553,7 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
         msg.info(f"Concatenation took {round(end_time - start_time, 3)} seconds")
 
         msg.info("Merging with original df")
-        self.df = self.df.merge(right=new_features, on=self.pred_time_uuid_col_name)
+        self._df = self._df.merge(right=new_features, on=self.pred_time_uuid_col_name)
 
     def add_temporal_predictors_from_pred_specs(  # pylint: disable=too-many-branches
         self,
@@ -621,14 +621,16 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
 
         data_of_birth_col_name = f"{output_prefix}_{input_date_of_birth_col_name}"
 
-        self.df[output_age_col_name] = (
-            (self.df[self.timestamp_col_name] - self.df[data_of_birth_col_name]).dt.days
+        self._df[output_age_col_name] = (
+            (
+                self._df[self.timestamp_col_name] - self._df[data_of_birth_col_name]
+            ).dt.days
             / (365.25)
         ).round(2)
 
         if birth_year_as_predictor:
             # Convert datetime to year
-            self.df["pred_birth_year"] = self.df[data_of_birth_col_name].dt.year
+            self._df["pred_birth_year"] = self._df[data_of_birth_col_name].dt.year
 
     def add_static_info(
         self,
@@ -676,8 +678,8 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
             },
         )
 
-        self.df = pd.merge(
-            self.df,
+        self._df = pd.merge(
+            self._df,
             df,
             how="left",
             on=self.id_col_name,
@@ -697,7 +699,7 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
         outcome_timestamp_col_name = f"{self.timestamp_col_name}_outcome"
 
         df = pd.merge(
-            self.df,
+            self._df,
             outcome_spec.values_df,
             how="left",
             on=self.id_col_name,
@@ -727,7 +729,7 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
 
         df.drop(["value"], axis=1, inplace=True)
 
-        self.df = df
+        self._df = df
 
     def add_temporal_outcome(
         self,
@@ -781,8 +783,8 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
                 f"{self.timestamp_col_name} is of type {timestamp_col_type}, not 'Timestamp' from Pandas. Will cause problems. Convert before initialising FlattenedDataset.",
             )
 
-        df = FlattenedDataset._flatten_temporal_values_to_df(
-            prediction_times_with_uuid_df=self.df[
+        df = TimeseriesFlattener._flatten_temporal_values_to_df(
+            prediction_times_with_uuid_df=self._df[
                 [
                     self.id_col_name,
                     self.timestamp_col_name,
@@ -795,7 +797,7 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
             pred_time_uuid_col_name=self.pred_time_uuid_col_name,
         )
 
-        self.df = self.df.merge(
+        self._df = self._df.merge(
             right=df,
             on=self.pred_time_uuid_col_name,
             validate="1:1",
@@ -911,3 +913,11 @@ class FlattenedDataset:  # pylint: disable=too-many-instance-attributes
             ["is_in_interval", "time_from_pred_to_val_in_days"],
             axis=1,
         )
+
+    def get_df(self) -> DataFrame:
+        """Get the flattened dataframe.
+
+        Returns:
+            DataFrame: Flattened dataframe.
+        """
+        return self._df


Test patch:
diff --git a/src/timeseriesflattener/testing/utils_for_testing.py b/src/timeseriesflattener/testing/utils_for_testing.py
index d9fa75f..13d28bc 100644
--- a/src/timeseriesflattener/testing/utils_for_testing.py
+++ b/src/timeseriesflattener/testing/utils_for_testing.py
@@ -10,7 +10,7 @@ import pytest
 from pandas import DataFrame
 from pandas.testing import assert_series_equal
 
-from timeseriesflattener import FlattenedDataset
+from timeseriesflattener import TimeseriesFlattener
 from timeseriesflattener.feature_spec_objects import TemporalSpec
 from timeseriesflattener.testing.load_synth_data import (
     load_synth_outcome,
@@ -86,7 +86,7 @@ def assert_flattened_data_as_expected(
     if isinstance(prediction_times_df, str):
         prediction_times_df = str_to_df(prediction_times_df)
 
-    flattened_ds = FlattenedDataset(
+    flattened_ds = TimeseriesFlattener(
         prediction_times_df=prediction_times_df,
         n_workers=4,
     )
@@ -96,12 +96,12 @@ def assert_flattened_data_as_expected(
     if expected_df:
         for col in expected_df.columns:
             assert_series_equal(
-                left=flattened_ds.df[col],
+                left=flattened_ds.get_df()[col],
                 right=expected_df[col],
                 check_dtype=False,
             )
     elif expected_values:
-        output = flattened_ds.df[output_spec.get_col_str()].values.tolist()
+        output = flattened_ds.get_df()[output_spec.get_col_str()].values.tolist()
         expected = list(expected_values)
 
         for i, expected_val in enumerate(expected):
diff --git a/tests/test_timeseriesflattener/test_add_values.py b/tests/test_timeseriesflattener/test_add_values.py
index 60d6e5a..04e829f 100644
--- a/tests/test_timeseriesflattener/test_add_values.py
+++ b/tests/test_timeseriesflattener/test_add_values.py
@@ -6,7 +6,7 @@ import numpy as np
 import pandas as pd
 import pytest
 
-from timeseriesflattener import FlattenedDataset
+from timeseriesflattener import TimeseriesFlattener
 from timeseriesflattener.feature_spec_objects import AnySpec, OutcomeSpec, PredictorSpec
 from timeseriesflattener.testing.utils_for_testing import (
     assert_flattened_data_as_expected,
@@ -203,7 +203,7 @@ def test_static_predictor():
                         1,1994-12-31 00:00:01
                         """
 
-    dataset = FlattenedDataset(prediction_times_df=str_to_df(prediction_times_df))
+    dataset = TimeseriesFlattener(prediction_times_df=str_to_df(prediction_times_df))
     dataset.add_static_info(
         static_spec=AnySpec(
             values_df=str_to_df(static_predictor),
@@ -224,7 +224,7 @@ def test_static_predictor():
     )
 
     pd.testing.assert_series_equal(
-        left=dataset.df[output_col_name].reset_index(drop=True),
+        left=dataset.get_df()[output_col_name].reset_index(drop=True),
         right=expected_values[output_col_name].reset_index(drop=True),
         check_dtype=False,
     )
@@ -240,7 +240,7 @@ def test_add_age():
                         1,1994-12-31 00:00:00
                         """
 
-    dataset = FlattenedDataset(prediction_times_df=str_to_df(prediction_times_df))
+    dataset = TimeseriesFlattener(prediction_times_df=str_to_df(prediction_times_df))
 
     output_prefix = "eval"
 
@@ -262,7 +262,7 @@ def test_add_age():
     )
 
     pd.testing.assert_series_equal(
-        left=dataset.df["eval_age_in_years"].reset_index(drop=True),
+        left=dataset.get_df()["eval_age_in_years"].reset_index(drop=True),
         right=expected_values[f"{output_prefix}_age_in_years"].reset_index(drop=True),
         check_dtype=False,
     )
@@ -278,7 +278,7 @@ def test_add_age_error():
                         1,94-12-31 00:00:00
                         """
 
-    dataset = FlattenedDataset(prediction_times_df=str_to_df(prediction_times_df))
+    dataset = TimeseriesFlattener(prediction_times_df=str_to_df(prediction_times_df))
 
     with pytest.raises(ValueError):
         dataset.add_age_and_birth_year(
@@ -311,7 +311,7 @@ def test_incident_outcome_removing_prediction_times():
     event_times_df = str_to_df(event_times_str)
     expected_df = str_to_df(expected_df_str)
 
-    flattened_dataset = FlattenedDataset(
+    flattened_dataset = TimeseriesFlattener(
         prediction_times_df=prediction_times_df,
         timestamp_col_name="timestamp",
         id_col_name="dw_ek_borger",
@@ -329,7 +329,7 @@ def test_incident_outcome_removing_prediction_times():
         ),
     )
 
-    outcome_df = flattened_dataset.df.reset_index(drop=True)
+    outcome_df = flattened_dataset.get_df().reset_index(drop=True)
 
     for col in expected_df.columns:
         pd.testing.assert_series_equal(
@@ -375,7 +375,7 @@ def test_add_multiple_static_predictors():
     birthdates_df = str_to_df(birthdates_df_str)
     male_df = str_to_df(male_df_str)
 
-    flattened_dataset = FlattenedDataset(
+    flattened_dataset = TimeseriesFlattener(
         prediction_times_df=prediction_times_df,
         timestamp_col_name="timestamp",
         id_col_name="dw_ek_borger",
@@ -403,7 +403,7 @@ def test_add_multiple_static_predictors():
         static_spec=AnySpec(values_df=male_df, feature_name="male", prefix="pred"),
     )
 
-    outcome_df = flattened_dataset.df
+    outcome_df = flattened_dataset.get_df()
 
     for col in (
         "dw_ek_borger",
@@ -446,7 +446,7 @@ def test_add_temporal_predictors_then_temporal_outcome():
     event_times_df = str_to_df(event_times_str)
     expected_df = str_to_df(expected_df_str)
 
-    flattened_dataset = FlattenedDataset(
+    flattened_dataset = TimeseriesFlattener(
         prediction_times_df=prediction_times_df,
         timestamp_col_name="timestamp",
         id_col_name="dw_ek_borger",
@@ -477,7 +477,7 @@ def test_add_temporal_predictors_then_temporal_outcome():
         ),
     )
 
-    outcome_df = flattened_dataset.df.set_index("dw_ek_borger").sort_index()
+    outcome_df = flattened_dataset.get_df().set_index("dw_ek_borger").sort_index()
     expected_df = expected_df.set_index("dw_ek_borger").sort_index()
 
     for col in expected_df.columns:
@@ -508,7 +508,7 @@ def test_add_temporal_incident_binary_outcome():
     event_times_df = str_to_df(event_times_str)
     expected_df = str_to_df(expected_df_str)
 
-    flattened_dataset = FlattenedDataset(
+    flattened_dataset = TimeseriesFlattener(
         prediction_times_df=prediction_times_df,
         timestamp_col_name="timestamp",
         id_col_name="dw_ek_borger",
@@ -526,7 +526,7 @@ def test_add_temporal_incident_binary_outcome():
         ),
     )
 
-    outcome_df = flattened_dataset.df
+    outcome_df = flattened_dataset.get_df()
 
     for col in [c for c in expected_df.columns if "outc" in c]:
         for df in (outcome_df, expected_df):
diff --git a/tests/test_timeseriesflattener/test_errors.py b/tests/test_timeseriesflattener/test_errors.py
index 6dc2f6a..8287c5e 100644
--- a/tests/test_timeseriesflattener/test_errors.py
+++ b/tests/test_timeseriesflattener/test_errors.py
@@ -3,7 +3,7 @@
 import pytest
 
 from timeseriesflattener.feature_spec_objects import PredictorSpec
-from timeseriesflattener.flattened_dataset import FlattenedDataset
+from timeseriesflattener.flattened_dataset import TimeseriesFlattener
 from timeseriesflattener.testing.utils_for_testing import (
     str_to_df,  # pylint: disable=import-error
 )
@@ -19,7 +19,7 @@ def test_col_does_not_exist_in_prediction_times():
     prediction_times_df = str_to_df(prediction_times_str)
 
     with pytest.raises(ValueError):
-        FlattenedDataset(  # noqa
+        TimeseriesFlattener(  # noqa
             prediction_times_df=prediction_times_df,
             timestamp_col_name="timestamp",
             id_col_name="dw_ek_borger",
@@ -39,7 +39,7 @@ def test_col_does_not_exist():
     prediction_times_df = str_to_df(prediction_times_str)
     event_times_df = str_to_df(event_times_str)
 
-    flattened_df = FlattenedDataset(
+    flattened_df = TimeseriesFlattener(
         prediction_times_df=prediction_times_df,
         timestamp_col_name="timestamp",
         id_col_name="dw_ek_borger",
@@ -64,6 +64,6 @@ def test_duplicate_prediction_times():
                                 1,2021-11-31 00:00:00
                                 """
 
-        FlattenedDataset(
+        TimeseriesFlattener(
             prediction_times_df=str_to_df(prediction_times_df_str),
         )
diff --git a/tests/test_timeseriesflattener/test_flattened_dataset_generation.py b/tests/test_timeseriesflattener/test_flattened_dataset_generation.py
index 359975d..7c1ad59 100644
--- a/tests/test_timeseriesflattener/test_flattened_dataset_generation.py
+++ b/tests/test_timeseriesflattener/test_flattened_dataset_generation.py
@@ -15,7 +15,7 @@ from timeseriesflattener.feature_spec_objects import (
     PredictorSpec,
     TemporalSpec,
 )
-from timeseriesflattener.flattened_dataset import FlattenedDataset
+from timeseriesflattener.flattened_dataset import TimeseriesFlattener
 from timeseriesflattener.testing.load_synth_data import (
     load_synth_prediction_times,
     synth_predictor_binary,
@@ -94,7 +94,7 @@ def create_flattened_df(
     prediction_times_df: pd.DataFrame,
 ):
     """Create a dataset df for testing."""
-    flat_ds = FlattenedDataset(
+    flat_ds = TimeseriesFlattener(
         prediction_times_df=prediction_times_df,
         n_workers=1,
         feature_cache_dir=cache_dir,
@@ -104,7 +104,7 @@ def create_flattened_df(
         predictor_specs=predictor_specs,
     )
 
-    return flat_ds.df
+    return flat_ds.get_df()
 
 
 @pytest.mark.parametrize(