Replicate_sundin.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "id": "951a14df",
   "metadata": {},
   "source": [
    "---\n",
    "output-file: replicate_sundin.html\n",
    "title: Replication of Sundin et al, ICML 2019\n",
    "skip_exec: true\n",
    "---\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7cdc05b-2d0f-4a83-83de-f91a60e0ba43",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install git@https://github.com/cavan33/openbt_py.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "596aa338",
   "metadata": {},
   "outputs": [],
   "source": [
    "from asbe.base import *\n",
    "from asbe.models import *\n",
    "from asbe.estimators import *\n",
    "from econml.orf import DMLOrthoForest\n",
    "from econml.dml import CausalForestDML\n",
    "from openbt.openbt import OPENBT\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.linear_model import LinearRegression, LogisticRegression\n",
    "from sklearn.model_selection import train_test_split\n",
    "from copy import deepcopy\n",
    "import econml\n",
    "from sklearn.gaussian_process import GaussianProcessRegressor\n",
    "from sklearn.gaussian_process.kernels import ConstantKernel, RBF\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import pickle\n",
    "import scipy.linalg as la\n",
    "from scipy.stats import norm\n",
    "from xbart import XBART\n",
    "import GPy\n",
    "import pymc as pm\n",
    "import pymc_bart as pmb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d879dda",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BARTEstimator(BaseITEEstimator):\n",
    "    def fit(self, **kwargs):\n",
    "        self.model_bart = pm.Model()\n",
    "        X = np.concatenate((kwargs[\"X_training\"], kwargs[\"t_training\"].reshape((-1,1))),axis=1)            \n",
    "        with self.model_bart:\n",
    "            x_obs = pm.MutableData(\"X\", X)\n",
    "            self.model_bart.named_vars.pop(\"X\")\n",
    "            x_obs = pm.MutableData(\"X\", X)\n",
    "            μ = pmb.BART(\"μ\", X=x_obs, Y=kwargs[\"y_training\"], m=10)\n",
    "            y_pred = pm.Normal(\"y_pred\", mu=μ, observed=kwargs[\"y_training\"],\n",
    "                               shape=μ.shape)\n",
    "            self.trace = pm.sample(random_seed=1005)\n",
    "            \n",
    "    def predict(self, **kwargs):\n",
    "        X0 = np.concatenate((kwargs[\"X\"],\n",
    "                             np.zeros(kwargs[\"X\"].shape[0]).reshape((-1,1))),axis=1)\n",
    "        X1 = np.concatenate((kwargs[\"X\"],\n",
    "                             np.ones(kwargs[\"X\"].shape[0]).reshape((-1,1))),axis=1)\n",
    "        with self.model_bart:\n",
    "            pm.set_data({\"X\": X1})\n",
    "            p1 = pm.sample_posterior_predictive(self.trace)\n",
    "            pm.set_data({\"X\": X0})\n",
    "            p0 = pm.sample_posterior_predictive(self.trace)\n",
    "#             ite = p1[\"posterior_predictive\"][\"y_pred\"] - p0[\"posterior_predictive\"][\"y_pred\"]\n",
    "            p1_ypred = p1[\"posterior_predictive\"][\"y_pred\"].mean(axis=0)\n",
    "            p0_ypred = p0[\"posterior_predictive\"][\"y_pred\"].mean(axis=0)\n",
    "            ite = p1_ypred - p0_ypred\n",
    "            ite = ite.mean(axis=0).to_numpy() #across chains\n",
    "            if \"return_mean\" in kwargs:\n",
    "                if kwargs[\"return_mean\"]:\n",
    "                    out =  ite.mean(axis=0)\n",
    "            else:\n",
    "                out =  ite.T\n",
    "            if \"return_counterfactuals\" in kwargs:\n",
    "                if kwargs[\"return_counterfactuals\"]:\n",
    "                    out = ite.mean(axis=0),p1_ypred.mean().values,p1_ypred.std().values, p0_ypred.mean().values, p0_ypred.std().values\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c3b69ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def categorical2indicator(data, name, categorical_max=4):\n",
    "    '''\n",
    "    Transforms categorical variable with name 'name' form a data frame to indicator variables\n",
    "    \n",
    "    Taken from https://github.com/IirisSundin/active-learning-for-decision-making/blob/e0c83f58181f81da2f867da4c49f1333fa7d0ae6/src/util.py#L14\n",
    "    '''\n",
    "    values = data[name].values\n",
    "    values[values>= categorical_max] = categorical_max\n",
    "    uni = np.unique(values)\n",
    "    for i, value in enumerate(uni):\n",
    "        data[name+'.'+str(i)] = np.array((values==value), dtype=int)\n",
    "    data.drop(name, axis=1)\n",
    "    return data\n",
    "\n",
    "def preprocess(data, categorical_max=2):\n",
    "    '''\n",
    "    This function preprocesses the hill data \n",
    "    '''\n",
    "    #Normalization is enough here\n",
    "    data['bw'] = normalize(data['bw'])\n",
    "    data['nnhealth'] = normalize(data['nnhealth'])\n",
    "    data['preterm'] = normalize(data['preterm'])\n",
    "    #Taking logarithm does not harm here before normalizing, but might be unneccessary\n",
    "    data['b.head'] = normalize(np.log(data['b.head']))\n",
    "    data['momage'] = normalize(np.log(data['momage']))\n",
    "    \n",
    "    #Categorigal variables are made to indicators, could also be just normalized:\n",
    "    #Birth order is between 1\n",
    "    data = categorical2indicator(data, 'birth.o', categorical_max=categorical_max)\n",
    "    \n",
    "    #Everything else is binary, so processing doesn't really help, makes only understanding the results harder.\n",
    "    \n",
    "    #For some reason, indicator variable \"first\" is either 1 or 2, that is why we subtract 1 from it\n",
    "    data['first'] = data['first'] -1 \n",
    "    \n",
    "    return data\n",
    "\n",
    "def normalize(data):\n",
    "    '''\n",
    "    Transforms the data to zero mean unit variance\n",
    "    '''\n",
    "    return (data- np.mean(data))/np.std(data)\n",
    "\n",
    "def prepare_data(id_test, random_state, percent_for_train=0.1335):\n",
    "    inputs = pd.read_csv(\"/Users/zoltan/Desktop/ihdp_data/inputs.csv\")\n",
    "    #names = [\"treatment\", \"y_factual\", \"y_cfactual\", \"mu0\", \"mu1\"] + [f'x{x}' for x in range(25)])\n",
    "    inputs = preprocess(inputs, categorical_max=2)\n",
    "    N = inputs.shape[0]\n",
    "    N_train = int(np.ceil(N*percent_for_train))\n",
    "    D = inputs.columns.shape[0]\n",
    "    outcomes = pd.read_csv('/Users/zoltan/Desktop/ihdp_data/observed_outcomes.csv', sep=',')\n",
    "    potential_outcomes = pd.read_csv('/Users/zoltan/Desktop/ihdp_data/potential_outcomes.csv', sep=',')\n",
    "    potential_outcomes = potential_outcomes[['outcome_c0','outcome_c1']]\n",
    "    actions = inputs['treat']\n",
    "    pred_names = inputs.columns[1:D]\n",
    "    predictors = inputs[pred_names]\n",
    "    \n",
    "    outcomes = outcomes.values\n",
    "    potential_outcomes = potential_outcomes.values\n",
    "    actions = actions.values\n",
    "    predictors = predictors.values\n",
    "    \n",
    "    outcomes = np.zeros(actions.shape)\n",
    "    for i in range(actions.shape[0]):\n",
    "        outcomes[i] = potential_outcomes[i, actions[i]]\n",
    "    \n",
    "    np.random.seed(random_state)\n",
    "    ind = np.arange(N)\n",
    "    np.random.shuffle(ind)\n",
    "    outcomes = outcomes[ind]\n",
    "    potential_outcomes = potential_outcomes[ind,:]\n",
    "    actions = actions[ind]\n",
    "    predictors = predictors[ind,:]\n",
    "\n",
    "    \n",
    "    # Creating training data\n",
    "    X_test = np.copy(predictors[id_test,:])\n",
    "    y_test = np.copy(outcomes[id_test])\n",
    "    t_test = np.copy(actions[id_test])\n",
    "    ite_test = potential_outcomes[id_test, 1]  - potential_outcomes[id_test, 0]\n",
    "                     \n",
    "    #shuffle the rest again (not test sample)\n",
    "    np.random.seed(random_state)\n",
    "    \n",
    "    ind = np.arange(N-1)\n",
    "    np.random.shuffle(ind)\n",
    "    outcomes = outcomes[np.arange(N)!=id_test][ind]\n",
    "    potential_outcomes = potential_outcomes[np.arange(N)!=id_test][ind,:]\n",
    "    actions = actions[np.arange(N)!=id_test][ind]\n",
    "    predictors = predictors[np.arange(N)!=id_test,:][ind,:]  \n",
    "                     \n",
    "    X_train, y_train, t_train = predictors[:N_train,:], outcomes[:N_train], actions[:N_train]\n",
    "    y0_train, y1_train = potential_outcomes[:N_train, 0], potential_outcomes[:N_train, 1]\n",
    "    ite_train = y1_train - y0_train\n",
    "                     \n",
    "    X_pool, y_pool, t_pool = predictors[N_train:,:], outcomes[N_train:], actions[N_train:]\n",
    "    y0_pool, y1_pool = potential_outcomes[N_train:, 0], potential_outcomes[N_train:, 1]\n",
    "    ite_pool = y1_pool - y0_pool\n",
    "            \n",
    "#     for col in ['bw','nnhealth', 'preterm', 'b.head','momage']:\n",
    "#         scalers[col] = StandardScaler()\n",
    "#         X_train[col] = scalers[col].fit_transform(X_train[col].to_numpy().reshape(-1,1))\n",
    "#         X_pool[col] = scalers[col].transform(X_pool[col].to_numpy().reshape(-1,1))\n",
    "#         try:\n",
    "#             X_test[col] = scalers[col].transform(X_test[col].to_numpy().reshape(-1,1))\n",
    "#         except:\n",
    "#             X_test[col] = scalers[col].transform(X_test[col].reshape(-1,1))\n",
    "    #X_train, X_pool, y_train, y_pool = X[], X_pool, y_train, y_pool\n",
    "    ds = {\"X_training\": X_train,\n",
    "     \"y_training\": y_train,\n",
    "     \"t_training\": t_train,\n",
    "     \"ite_training\": ite_train,\n",
    "     \"X_pool\": X_pool, \n",
    "     \"y_pool\": y_pool,\n",
    "     \"t_pool\": t_pool,\n",
    "     \"y1_pool\": y1_pool,\n",
    "     \"y0_pool\": y0_pool,\n",
    "     \"ite_pool\": ite_pool,\n",
    "     \"X_test\": X_test.reshape((1,-1)),\n",
    "     \"y_test\": y_test,\n",
    "     \"t_test\": t_test,\n",
    "     \"ite_test\": ite_test.reshape((-1,1))\n",
    "     }\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b6bf979",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPyEstimator(BaseITEEstimator):\n",
    "    # https://github.com/IirisSundin/active-learning-for-decision-making/blob/master/src/gpmodel.py\n",
    "    def __init__(self,**kwds):\n",
    "        super().__init__(**kwds)\n",
    "        self.name = \"gpy\"\n",
    "        self.update = False\n",
    "    \n",
    "    def _create_model(self, x, y):\n",
    "        d = x.shape[1]\n",
    "        prior = GPy.core.parameterization.priors.Gamma(a=1.5,b=3.0)\n",
    "        kern = GPy.kern.RBF(input_dim=d, ARD=True)\n",
    "        kern.variance.set_prior(prior, warning=False)\n",
    "        kern.lengthscale.set_prior(prior, warning=False)\n",
    "        lik1 = GPy.likelihoods.Gaussian()\n",
    "        lik1.variance.set_prior(prior, warning=False)\n",
    "        lik_expert = GPy.likelihoods.Gaussian()\n",
    "        lik_expert.variance.set_prior(prior, warning=False)\n",
    "        lik = GPy.likelihoods.MixedNoise([lik1, lik_expert])\n",
    "        output_index = np.ones(x.shape[0])\n",
    "        model = GPy.core.GP(X =  x,\n",
    "                                 Y =  y.reshape(-1, 1), \n",
    "                                 kernel=kern,\n",
    "                                 likelihood=lik, \n",
    "                                 Y_metadata = {'output_index':output_index})\n",
    "        model.optimize()\n",
    "        return model\n",
    "        \n",
    "    def fit(self, **kwargs):\n",
    "        if self.update:\n",
    "            action = \"t\" if kwargs[\"t_training\"] == 1 else \"c\"\n",
    "            predictor = kwargs[\"X_training\"]\n",
    "            oracle_outcome = kwargs[\"y_training\"]\n",
    "            mcopy = {}\n",
    "            mcopy[\"c\"] = deepcopy(self.models[\"c\"])\n",
    "            mcopy[\"t\"] = deepcopy(self.models[\"t\"])\n",
    "            mcopy[action].Y_metadata['output_index'] = np.r_[mcopy[action].Y_metadata['output_index'], np.array([1])]\n",
    "            mcopy[action].set_XY(np.r_[mcopy[action].X, predictor], np.r_[mcopy[action].Y, oracle_outcome], )\n",
    "            mcopy[action].optimize()\n",
    "            self.updated_model = mcopy\n",
    "        else:\n",
    "            self.models = {}\n",
    "            Xt, yt = kwargs[\"X_training\"][(kwargs[\"t_training\"] == 1), :],\\\n",
    "                    kwargs[\"y_training\"][(kwargs[\"t_training\"] == 1)]\n",
    "            Xc, yc = kwargs[\"X_training\"][(kwargs[\"t_training\"] == 0), :],\\\n",
    "            kwargs[\"y_training\"][(kwargs[\"t_training\"] == 0)]\n",
    "            self.models[\"c\"] = self._create_model(Xc, yc)\n",
    "            self.models[\"t\"] = self._create_model(Xt, yt)\n",
    "        \n",
    "    def predict(self, X, **kwargs):\n",
    "        if len(X.shape) == 1:\n",
    "            X = X.reshape((1, -1))\n",
    "        p1 = self.models[\"t\"].posterior_samples_f(X, 100)\n",
    "        p0 = self.models[\"c\"].posterior_samples_f(X, 100)\n",
    "        ite = p1 - p0\n",
    "        ite = ite.squeeze(1)\n",
    "        if \"return_mean\" in kwargs:\n",
    "            if kwargs[\"return_mean\"]:\n",
    "                ite = self.models[\"t\"].predict_noiseless(X)[\n",
    "                    0] - self.models[\"c\"].predict_noiseless(X)[0]\n",
    "        if \"return_counterfactuals\" in kwargs:\n",
    "            if kwargs[\"return_counterfactuals\"]:\n",
    "                p1, p1s = self.models[\"t\"]._raw_predict(X)\n",
    "                p0, p0s = self.models[\"c\"]._raw_predict(X)\n",
    "                ite = (ite, p1, p0, p1s, p0s)\n",
    "        return ite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "465c9ded",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ExpectedReliability(BaseAcquisitionFunction): #'decerr'\n",
    "        '''\n",
    "        Uses Gauss-Hermite quadrature to compute expected type S error rate\n",
    "        '''\n",
    "        def error_rate(self, preds):\n",
    "            mu_tau = np.mean(preds[0])\n",
    "            sd_tau = np.std(preds[0])\n",
    "            alpha = norm.cdf(-np.abs(mu_tau)/sd_tau)\n",
    "            return alpha\n",
    "        \n",
    "        def calculate_metrics(self, model, dataset):\n",
    "            utilities = np.zeros(dataset[\"X_pool\"].shape[0])\n",
    "            reload = False\n",
    "            for n in range(dataset[\"X_pool\"].shape[0]):\n",
    "                x_star = dataset[\"X_pool\"][n,:].reshape(1,-1)\n",
    "                a_star = 1 - dataset[\"t_pool\"][n] # counterfactual action\n",
    "                decerr = 0.0\n",
    "                points, weights = np.polynomial.hermite.hermgauss(16) \n",
    "                for ii, yy in enumerate(points):\n",
    "                    preds_star = model.predict(X=x_star, return_counterfactuals=True)\n",
    "                    if a_star == 1:\n",
    "                        mu_star, S_star = preds_star[1], preds_star[3]\n",
    "                    else:\n",
    "                        mu_star, S_star = preds_star[2], preds_star[4]\n",
    "                    y_star = np.sqrt(2)*np.sqrt(S_star)*yy + mu_star #for substitution\n",
    "                    new_data = {\"X_training\": np.concatenate((dataset[\"X_training\"], x_star)),\n",
    "                               \"y_training\":  np.concatenate((dataset[\"y_training\"], \n",
    "                                                              y_star.reshape((1,)))),\n",
    "                                \"t_training\": np.concatenate((dataset[\"t_training\"], \n",
    "                                                              a_star.reshape((1,))))\n",
    "                               }\n",
    "                    if model.name == \"gpy\":\n",
    "                        model.update = True\n",
    "                        model.fit(X_training=x_star,\n",
    "                                        y_training=y_star,\n",
    "                                        t_training=a_star.reshape((1,)))\n",
    "                        gpy_new = GPyEstimator()\n",
    "                        gpy_new.models = model.updated_model\n",
    "                        preds_next = gpy_new.predict(X = dataset[\"X_test\"],\n",
    "                                         return_counterfactuals=True)\n",
    "                        model.update = False\n",
    "                    else:\n",
    "                        try:\n",
    "                            new_model = deepcopy(model)\n",
    "                            new_model.fit(**new_data)\n",
    "                            preds_next = new_model.predict(X = dataset[\"X_test\"],\n",
    "                                         return_counterfactuals=True)\n",
    "                        except TypeError:\n",
    "                            nm = XBARTEstimator(model = XBART(num_sweeps=20), two_model = False)\n",
    "                            nm.fit(**new_data)\n",
    "                            preds_next = nm.predict(X = dataset[\"X_test\"],\n",
    "                                         return_counterfactuals=True)\n",
    "                    util = 1-self.error_rate(preds_next)\n",
    "                    decerr += weights[ii] * 1/np.sqrt(np.pi) * util\n",
    "                utilities[n] = decerr\n",
    "                print(f\"{n} is done\")\n",
    "            return decerr\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb21ba10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# class ExpectedTargetedIG(BaseAcquisitionFunction):\n",
    "#     def mvn_KL(self, mu1, S1, mu2, S2):\n",
    "#         '''\n",
    "#         KL divergence between two multivariate normals\n",
    "#         '''\n",
    "#         if type(mu1)!=np.ndarray:\n",
    "#             d = 1\n",
    "#         else:\n",
    "#             d = mu1.shape[0]\n",
    "#         KL = 0.5*(np.log(la.det(S2)/la.det(S1)) + np.trace(la.inv(S2).dot(S1)) + (mu1 - mu2).T.dot(la.inv(S2).dot((mu1 - mu2))) - d)\n",
    "#         return KL\n",
    "    \n",
    "#     def calculate_metrics(self, model, dataset):\n",
    "#         '''\n",
    "#         Uses Gauss-Hermite quadrature to compute expected information gain in decision\n",
    "#         '''\n",
    "#         ite_pred, p1, p0, p1s, p0s = model.predict(X = dataset[\"X_test\"],\n",
    "#                                          return_counterfactuals=True)\n",
    "#         ypred = np.where(p1 > p0, p1, p0)\n",
    "#         Spred = np.where(p1 > p0, p1s, p0s)\n",
    "#         N = dataset[\"X_pool\"].shape[0]\n",
    "#         # d = self._decide(preds)\n",
    "#         # ypred, Spred = preds[d]\n",
    "#         X_STAR = dataset[\"X_pool\"]\n",
    "#         # x_star = self.dat['predictors'][n,:].reshape(1,-1)\n",
    "#         a_star = 1 - dataset[\"t_pool\"] # counterfactual action\n",
    "\n",
    "#         points, weights = np.polynomial.hermite.hermgauss(32) \n",
    "#         scores = []\n",
    "#         NUM_SIMULATIONS = 10\n",
    "#         sims = np.random.choice(N, NUM_SIMULATIONS, replace=False)\n",
    "#         for i in sims:\n",
    "#             x_star = X_STAR[i, :].reshape((1, -1))\n",
    "#             preds_star = model.predict(x_star, return_counterfactuals=True)\n",
    "#             targig = 0.0\n",
    "#             for ii, yy in enumerate(points):\n",
    "#                 if dataset[\"t_pool\"][i] == 1:\n",
    "#                     mu_star, S_star = preds_star[1], preds_star[3]\n",
    "#                 else:\n",
    "#                     mu_star, S_star = preds_star[2], preds_star[4]\n",
    "#                 y_star = np.sqrt(2)*np.sqrt(S_star)*yy + mu_star #for substitution\n",
    "#                 #try:\n",
    "#                 new_model = deepcopy(model)\n",
    "#                 new_data = {\"X_training\": np.concatenate((dataset[\"X_training\"], x_star)),\n",
    "#                            \"y_training\":  np.concatenate((dataset[\"y_training\"], \n",
    "#                                                           dataset[\"y_pool\"][i].reshape((1,)))),\n",
    "#                             \"t_training\": np.concatenate((dataset[\"t_training\"], \n",
    "#                                                           dataset[\"t_pool\"][i].reshape((1,))))\n",
    "#                            }\n",
    "#                 new_model.fit(**new_data)\n",
    "#                 preds_next = new_model.predict(X = dataset[\"X_test\"],\n",
    "#                                  return_counterfactuals=True)\n",
    "#                 #model_star = self.update(a_star, x_star, y_star)\n",
    "#                 # preds_next = self.predict(new_predictors, model_star)\n",
    "#                 #    preds_next = self.predict(new_predictors)\n",
    "#                 D_KL = 0\n",
    "#                 for d in range(2):\n",
    "#                     ypred_next, Spred_next = preds_next[d+1], preds_next[d+3]\n",
    "#                     D_KL += self.mvn_KL(ypred_next, Spred_next, ypred, Spred)\n",
    "#                 targig += weights[ii] * 1/np.sqrt(np.pi) * D_KL\n",
    "#             scores += [targig]\n",
    "#         scs = np.zeros(N)\n",
    "#         scs[sims] = np.array(scores).ravel()\n",
    "#         return scs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca501e89",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7102c7d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class XBARTEstimator(BaseITEEstimator):\n",
    "    def __init__(self, model, two_model=False):\n",
    "        super().__init__(model = model, two_model = two_model,dataset=None, ps_model=None)\n",
    "        \n",
    "    def predict(self, **kwargs):\n",
    "        X0 = np.concatenate((kwargs[\"X\"],\n",
    "                             np.zeros(kwargs[\"X\"].shape[0]).reshape((-1,1))),axis=1)\n",
    "        X1 = np.concatenate((kwargs[\"X\"],\n",
    "                             np.ones(kwargs[\"X\"].shape[0]).reshape((-1,1))),axis=1)\n",
    "        if \"return_mean\" in kwargs:\n",
    "            if kwargs[\"return_mean\"] is True:\n",
    "                out = self.model.predict(X1) - self.model.predict(X0)\n",
    "        else:\n",
    "            out = self.model.predict(X1, return_mean=False) - self.model.predict(X0, return_mean=False)\n",
    "        if \"return_counterfactuals\" in kwargs:\n",
    "            if kwargs[\"return_counterfactuals\"] is True:\n",
    "                p1 = self.model.predict(X1, return_mean=False)\n",
    "                p0 = self.model.predict(X0, return_mean=False)\n",
    "                return (out, np.mean(p1), np.mean(p0), np.std(p1), np.std(p0))\n",
    "        return out\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1459914",
   "metadata": {},
   "outputs": [],
   "source": [
    "ests = {#\"gpy\":GPyEstimator()\n",
    "#         \"gpy\":GPyEstimator(),\n",
    "#          \"openbt\": OPENBTITEEstimator(model=openbt.OPENBT(model=\"bart\"),two_model=True),\n",
    "#          \"bart\":BARTEstimator(),\n",
    "         \"xbart_optimized\":XBARTEstimator(model=XBART( num_sweeps= 50,\n",
    "                                            num_trees=400,\n",
    "                                            beta=3,alpha=0.6))\n",
    "}\n",
    "ests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dbc0709",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_results = {}\n",
    "for key, value in ests.items():\n",
    "    res = {}\n",
    "    for i in range(1):\n",
    "        for d in range(747):\n",
    "            ds = prepare_data(d, i)\n",
    "            asl = BaseActiveLearner(estimator = value,\n",
    "                                     acquisition_function=[RandomAcquisitionFunction(name=\"random\",\n",
    "                                                                                     method = \"top\" ),\n",
    "                                                           EMCMAcquisitionFunction(name=\"emcm\",\n",
    "                                                                                   method = \"top\"), \n",
    "                                                           UncertaintyAcquisitionFunction(name=\"unc\"),\n",
    "#                                                            ExpectedReliability(name=\"er\", method=\"top\")\n",
    "                                                         ],\n",
    "                                    assignment_function=BaseAssignmentFunction(),\n",
    "                                    stopping_function = None,\n",
    "                                    dataset=ds,\n",
    "                                    al_steps=6)\n",
    "            _ = asl.simulate(no_query=1, metric=[\"PEHE\", \"decision\"])\n",
    "            res[f\"{key}_{i}_{d}\"] =  pd.DataFrame(asl.simulation_results)\n",
    "            res[f\"{key}_{i}_{d}\"][\"sim\"] = i\n",
    "            res[f\"{key}_{i}_{d}\"][\"data\"] = d\n",
    "#             with open(f\"/Users/zoltan/Desktop/{key}_res.pickle\", \"wb\") as output_file:\n",
    "#                 pickle.dump(res, output_file)\n",
    "#             print(res)\n",
    "#             print(f\"D {d} is done\")\n",
    "#             temp_df = create_table(res)\n",
    "#             temp_df.to_csv(\"/Users/zoltan/Desktop/asbe_ihdp_results_xbart.csv\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}