Meta learners

causalpy/__init__.py

@@ -3,7 +3,9 @@
 import causalpy.pymc_experiments
 import causalpy.pymc_models
 import causalpy.skl_experiments
+import causalpy.skl_meta_learners
 import causalpy.skl_models
+import causalpy.summary
 from causalpy.version import __version__
 
 from .data import load_data

causalpy/pymc_models.py

@@ -4,7 +4,9 @@
 import numpy as np
 import pandas as pd
 import pymc as pm
+import pymc_bart as pmb
 from arviz import r2_score
+from pymc.distributions.distribution import DistributionMeta
 
 
 class ModelBuilder(pm.Model):
@@ -113,3 +115,120 @@ def build_model(self, X, y, coords):
             sigma = pm.HalfNormal("sigma", 1)
             mu = pm.Deterministic("mu", pm.math.dot(X, beta), dims="obs_ind")
             pm.Normal("y_hat", mu, sigma, observed=y, dims="obs_ind")
+
+
+class BARTRegressor(ModelBuilder):
+    """
+    Class for building BART based regressors for meta-learners.
+
+    Parameters
+    ----------
+    m : int.
+        Number of trees to fit.
+    sigma : float.
+        Prior standard deviation.
+    sample_kwargs : dict.
+        Keyword arguments for sampler.
+    """
+
+    def __init__(
+        self,
+        m: int = 20,
+        sigma: float = 1.0,
+        sample_kwargs: Optional[dict[str, Any]] = None,
+    ):
+        self.m = m
+        self.sigma = sigma
+        super().__init__(sample_kwargs)
+
+    def build_model(self, X, y, coords=None):
+        with self:
+            self.add_coords(coords)
+            X_ = pm.MutableData("X", X, dims=["obs_ind", "coeffs"])
+            mu = pmb.BART("mu", X_, y, m=self.m, dims="obs_ind")
+            pm.Normal("y_hat", mu=mu, sigma=self.sigma, observed=y, dims="obs_ind")
+
+
+class BARTClassifier(ModelBuilder):
+    """
+    Class for building BART based models for meta-learners.
+
+    Parameters
+    ----------
+    m : int.
+        Number of trees to fit.
+    sample_kwargs : dict.
+        Keyword arguments for sampler.
+    """
+
+    def __init__(
+        self,
+        m: int = 20,
+        sample_kwargs: Optional[dict[str, Any]] = None,
+    ):
+        self.m = m
+        super().__init__(sample_kwargs)
+
+    def build_model(self, X, y, coords=None):
+        with self:
+            self.add_coords(coords)
+            X_ = pm.MutableData("X", X, dims=["obs_ind", "coeffs"])
+            mu_ = pmb.BART("mu_", X_, y, m=self.m, dims="obs_ind")
+            mu = pm.Deterministic("mu", pm.math.sigmoid(mu_), dims="obs_ind")
+            pm.Bernoulli("y_hat", mu, observed=y, dims="obs_ind")
+
+
+class LogisticRegression(ModelBuilder):
+    """
+    Custom PyMC model for logistic regression.
+
+    Parameters
+    ----------
+    coeff_distribution :    PyMC distribution.
+                Prior distribution of coefficient vector.
+    distribution_kwargs :   dict.
+                Keyword arguments for prior distribution.
+    sample_kwargs :   dict.
+                Keyword arguments for sampler.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import pymc as pm
+    >>> from causalpy.pymc_models import LogisticRegression
+    >>>
+    >>> X = np.random.rand(10, 10)
+    >>> y = np.random.rand(10)
+    >>> m = LogisticRegression(
+    >>>         coeff_distribution=pm.Cauchy,
+    >>>         coeff_distribution_kwargs={"alpha": 0, "beta": 1}
+    >>> )
+    >>>
+    >>> m.fit(X, y)
+    """
+
+    def __init__(
+        self,
+        sample_kwargs: Optional[dict[str, Any]] = None,
+        coeff_distribution: DistributionMeta = pm.Normal,
+        coeff_distribution_kwargs: Optional[dict[str, Any]] = None,
+    ):
+        self.coeff_distribution = coeff_distribution
+        if coeff_distribution_kwargs is None:
+            self.coeff_distribution_kwargs = {"mu": 0, "sigma": 50}
+        else:
+            self.coeff_distribution_kwargs = coeff_distribution_kwargs
+
+        super().__init__(sample_kwargs)
+
+    def build_model(self, X, y, coords) -> None:
+        with self:
+            self.add_coords(coords)
+            X_ = pm.MutableData("X", X, dims=["obs_ind", "coeffs"])
+            beta = self.coeff_distribution(
+                "beta", dims="coeffs", **self.coeff_distribution_kwargs
+            )
+            mu = pm.Deterministic(
+                "mu", pm.math.sigmoid(pm.math.dot(X_, beta)), dims="obs_ind"
+            )
+            pm.Bernoulli("y_hat", mu, observed=y, dims="obs_ind")

causalpy/summary.py

@@ -0,0 +1,245 @@
+"Summary objects."
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+def html_header(columns: list[str], colspan: int = 1) -> str:
+    """
+    Returns HTML code for table header.
+
+    Parameters
+    ----------
+    columns :   list[str].
+        List containing column names.
+    colspan :   int.
+        Column span.
+    """
+    string = (
+        f'<th colspan={colspan} style="text-align: center"> '
+        + (f'</th> <th style="text-align: center" colspan={colspan}>'.join(columns))
+        + "</th>"
+    )
+    return '<tr style="text-align: center">' + string + "</tr>"
+
+
+def html_rows(index: str, values: list, colspan: int = 1) -> str:
+    """
+    Returns HTML code for table rows.
+
+    Parameters
+    ----------
+    items :   dict.
+        The keys of items is the index-set of the rows to be added. The values are
+        lists containing the values to be added.
+    colspan :   int.
+        Column span.
+    """
+    string = ""
+    string += f'<th style="text-align: left"> {index} </th>'
+
+    if not isinstance(values, list):
+        values = [values]
+
+    values = map(str, values)
+    string += (
+        f'<td colspan={colspan} style="text-align: center">'
+        + ('</td> <td style="text-align: center">'.join(values))
+        + "</td>"
+    )
+
+    string = f"<tr> {string} </tr>"
+
+    return string
+
+
+def str_header(columns: list[str], length: int) -> str:
+    """
+    Returns table header string.
+
+    Parameters
+    ----------
+    columns :   list[str].
+        List containing column names.
+    length :   int.
+        Length of box containing header.
+    """
+    # If first column is empty, it's box should not be displayed
+    first_col_empty = columns[0] == ""
+
+    if first_col_empty:
+        columns = columns[1:]
+
+    n_cols = len(columns)
+
+    top = f"{length * '═'}╦"
+    bot = f"{length * '═'}╩"
+
+    spaces = first_col_empty * (length + 1) * " "
+
+    return f"""\
+    {spaces}╔{(n_cols - 1) * top}{length * "═"}╗
+    {spaces}║{"║".join(map(lambda x: x.center(length), columns))}║
+    ┏{length * '━'}╚{(n_cols - 1) * bot}{length * "═"}╝\
+    """
+
+
+def str_row(index: str, values: str, length: int, row_type: str = "inner") -> str:
+    """
+    Returns string table row.
+
+    Parameters
+    ----------
+    index : str.
+        Index of the row.
+    values :    int.
+        Values of the row.
+    length :    int.
+        Length of boxes.
+    row_type :  str.
+        One of "inner", "first", "last".
+    """
+    n_cols = len(values)
+    values = map(str, values)
+
+    s = ""
+
+    if row_type == "first":
+        s += f"""    ┏{length * '━'}┱"""
+        s += f"{(n_cols - 1) * (length * '─' + '┬')}"
+        s += f"{(length * '─' + '┐')}"
+        s += "\n"
+
+    s += f"""\
+    ┃{index.center(length)}┃{"│".join(map(lambda x: x.center(length), values))}│
+    """
+    if row_type == "last":
+        s += f"┗{length * '━'}┹{(n_cols - 1) * (length * '─' + '┴')}{length * '─'}┘"
+    else:
+        s += f"┣{length * '━'}╋{(n_cols - 1) * (length * '─' + '┼')}{length * '─'}┤"
+
+    return s
+
+
+def str_title(title: str):
+    """
+    Returns title in a box.
+
+    Parameters
+    ----------
+    title : str.
+        String to return in a box.
+    """
+    length = len(title)
+    return f"""\
+    ╔{(length + 2) * '═'}╗
+    ║ {title} ║
+    ╚{(length + 2) * '═'}╝
+    """
+
+
+@dataclass
+class Record:
+    """
+    Class representing either a header or a row of a table.
+
+    Parameters
+    ----------
+    record_type : str.
+        Either 'header', 'title' or 'row'.
+    values : list.
+        List of values in record.
+    colspan : int.
+        Column span.
+    """
+
+    record_type: str
+    values: list
+    colspan: int
+    index: Optional[str] = None
+
+    def to_html(self):
+        "Returns HTML code for record."
+        if self.record_type in ["header", "title"]:
+            return html_header(self.values, self.colspan)
+        else:
+            return html_rows(self.index, self.values, self.colspan)
+
+
+class Summary:
+    """
+    Base summary class.
+    """
+
+    def __init__(self):
+        self.records = []
+
+    def add_header(self, col_names, colspan) -> None:
+        "Adds a header."
+        self.records.append(Record("header", col_names, colspan))
+
+    def add_row(self, index: str, values: list, colspan: int) -> None:
+        "Adds a row."
+        self.records.append(Record("row", values, colspan, index=index))
+
+    def add_title(self, title) -> None:
+        "Adds title."
+        self.records.append(Record("title", title, 1))
+
+    def get_longest_length(self) -> int:
+        """
+        Returns the lenght of the longest item currently in the table.
+        """
+        non_titles = [r for r in self.records if not r.record_type == "title"]
+
+        def length_of_items(L):
+            return [len(str(x)) for x in L]
+
+        vals = [length_of_items(r.values) for r in non_titles]
+        longest_value_length = max([max(x) for x in vals])
+        indx = [len(str(x.index)) for x in self.records]
+        longest_index_length = max(indx)
+        return max(longest_index_length, longest_value_length)
+
+    def to_string(self) -> str:
+        """
+        Return string summary table in string form.
+        """
+        type_of_next_row = "first"
+        length = self.get_longest_length()
+        s = ""
+
+        for i, x in enumerate(self.records):
+            if i == len(self.records) - 1:
+                type_of_next_row = "last"
+            elif self.records[i + 1].record_type == "title":
+                type_of_next_row = "last"
+
+            if x.record_type == "header":
+                s += str_header(x.values, length)
+                type_of_next_row = "inner"
+            elif x.record_type == "title":
+                s += str_title(x.values[0])
+                type_of_next_row = "first"
+            else:
+                s += str_row(x.index, x.values, length, type_of_next_row)
+                type_of_next_row = "inner"
+            s += "\n"
+        return s
+
+    def to_html(self) -> str:
+        """
+        Returns HTML code for summary table.
+        """
+        s = ""
+
+        for x in self.records:
+            s += x.to_html()
+
+        return "<table>" + s + "</table>"
+
+    def _repr_html_(self) -> str:
+        return self.to_html()
+
+    def __repr__(self) -> str:
+        return self.to_string()

causalpy/utils.py

@@ -1,5 +1,18 @@
 import pandas as pd
 
+from causalpy.pymc_models import ModelBuilder
+
+
+def _fit(model, X, y, coords):
+    """
+    Fits model to X, y, where model is either a sklearn model or a ModelBuilder
+    instance. In the later case it passes coords, in the first case coords is ignored.
+    """
+    if isinstance(model, ModelBuilder):
+        model.fit(X, y, coords)
+    else:
+        model.fit(X, y)
+
 
 def _is_variable_dummy_coded(series: pd.Series) -> bool:
     """Check if a data in the provided Series is dummy coded. It should be 0 or 1

pyproject.toml

@@ -35,10 +35,11 @@ dependencies = [
     "pandas",
     "patsy",
     "pymc>=5.0.0",
+    "pymc-bart",
     "scikit-learn>=1",
     "scipy",
     "seaborn>=0.11.2",
-    "xarray>=v2022.11.0",
+    "xarray>=v2022.11.0"
 ]
 
 # List additional groups of dependencies here (e.g. development dependencies). Users