refactor: TreatmentEffect builder pattern

.vscode/settings.json

@@ -20,7 +20,7 @@
   "python.terminal.activateEnvironment": true,
   "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python",
   "python.terminal.activateEnvInCurrentTerminal": true,
-  "autoDocstring.docstringFormat": "sphinx",
+  "autoDocstring.docstringFormat": "google",
   "python.analysis.typeCheckingMode": "strict",
   "[python]": {
     "editor.defaultFormatter": "charliermarsh.ruff"

medmodels/__init__.py

@@ -1,3 +1,4 @@
 from medmodels.medrecord import MedRecord
+from medmodels.treatment_effect_estimation import TreatmentEffect
 
-__all__ = [MedRecord]
+__all__ = [MedRecord, TreatmentEffect]

medmodels/treatment_effect_estimation/__init__.py

@@ -0,0 +1,3 @@
+from medmodels.treatment_effect_estimation.treatment_effect import TreatmentEffect
+
+__all__ = ["TreatmentEffect"]

medmodels/treatment_effect_estimation/analysis_modules/adjust.py

@@ -1,13 +1,59 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional, Set, Literal, Tuple
+from typing_extensions import TypeAlias
+
+from medmodels import MedRecord
+from medmodels.medrecord.types import (
+    NodeIndex,
+)
 
 if TYPE_CHECKING:
     from medmodels.treatment_effect_estimation.treatment_effect import TreatmentEffect
 
 
+MatchingMethod: TypeAlias = Literal["propensity", "nearest_neighbors"]
+
+
 class Adjust:
     _treatment_effect: TreatmentEffect
 
     def __init__(self, treatment_effect: TreatmentEffect) -> None:
         self._treatment_effect = treatment_effect
+
+    def _apply_matching(
+        self,
+        method: Optional[MatchingMethod],
+        medrecord: MedRecord,
+        treatment_all: Set[NodeIndex],
+        control_true: Set[NodeIndex],
+        control_false: Set[NodeIndex],
+    ) -> Tuple[Set[NodeIndex], Set[NodeIndex]]:
+        """
+        Update the treatment effect object with the matched controls.
+
+        Args:
+            medrecord (MedRecord): The MedRecord object containing the data.
+            treatment_all (Set[NodeIndex]): The set of all patients in the treatment
+                group.
+            control_true (Set[NodeIndex]): The set of patients in the control group with
+                the outcome of interest.
+            control_false (Set[NodeIndex]): The set of patients in the control group
+                without the outcome of interest.
+            control_false (Set[NodeIndex]): The set of patients in the control group
+                without the outcome of interest.
+
+        Returns:
+            Tuple[Set[NodeIndex], Set[NodeIndex]]: The updated control_true and
+                control_false sets after matching.
+        """
+        if method is None:
+            return control_true, control_false
+
+        # If it is not None, apply the matching method
+        method_function = getattr(self, method)
+        control_true, control_false = method_function(
+            medrecord, treatment_all, control_true, control_false
+        )
+
+        return control_true, control_false