From ace0b5276824663f00e3e495aec4937e1d3051c9 Mon Sep 17 00:00:00 2001 From: "lennart.carstens-behrens" Date: Tue, 13 May 2025 10:23:51 +0200 Subject: [PATCH] added from_path to InitialEstimates --- pytmle/estimates.py | 46 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/pytmle/estimates.py b/pytmle/estimates.py index 42eb4d7..ff74d99 100644 --- a/pytmle/estimates.py +++ b/pytmle/estimates.py @@ -1,5 +1,6 @@ import numpy as np import pandas as pd +from pathlib import Path from dataclasses import dataclass, field from typing import Optional, List, Union @@ -50,6 +51,51 @@ def __setattr__(self, name, value): self._check_compatibility(value, check_width=True) super().__setattr__(name, value) + @classmethod + def from_path(cls, data: pd.DataFrame, path: str | Path): + if type(path) != Path: + path = Path(path) + surv_0_df = pd.read_csv(path / "surv_0.csv") + surv_1_df = pd.read_csv(path / "surv_1.csv") + cens_surv_0_df = pd.read_csv(path / "cens_surv_0.csv") + cens_surv_1_df = pd.read_csv(path / "cens_surv_1.csv") + haz1_0_df = pd.read_csv(path / "haz1_0.csv") + haz1_1_df = pd.read_csv(path / "haz1_1.csv") + haz2_0_df = pd.read_csv(path / "haz2_0.csv") + haz2_1_df = pd.read_csv(path / "haz2_1.csv") + prop_0_df = pd.read_csv(path / "prop_0.csv", header=None) + prop_1_df = pd.read_csv(path / "prop_1.csv", header=None) + + # List of dataframes to check + dataframes = [surv_0_df, surv_1_df, cens_surv_0_df, cens_surv_1_df, haz1_0_df, haz1_1_df, haz2_0_df, haz2_1_df] + + # Check if all dataframes have the same columns + all_col_equal = all(df.columns.equals(dataframes[0].columns) for df in dataframes) + + assert all_col_equal, "All columns must be equal." + + dataframes += [data, prop_0_df, prop_1_df] + + # Check if all dataframes have the same index + all_indices_equal = all(df.index.equals(dataframes[0].index) for df in dataframes) + + assert all_indices_equal, "All indices must be equal" + + return { + 0: cls(times = surv_0_df.columns.astype(float), + g_star_obs= 1 - data["chemo"].values, + propensity_scores=prop_0_df.values.squeeze(), + hazards=np.stack([haz1_0_df.values, haz2_0_df.values], axis=-1), + event_free_survival_function=surv_0_df.values, + censoring_survival_function=cens_surv_0_df.values), + 1: cls(times = surv_1_df.columns.astype(float), + g_star_obs= data["chemo"].values, + propensity_scores=prop_1_df.values.squeeze(), + hazards=np.stack([haz1_1_df.values, haz2_1_df.values], axis=-1), + event_free_survival_function=surv_1_df.values, + censoring_survival_function=cens_surv_1_df.values), + } + def _check_compatibility(self, new_element, check_width): # check that all given estimates have the same length (first dimension size) if self._length is None: