catenets.models.torch.pseudo_outcome_nets module

class DRLearner(n_unit_in: int, binary_y: bool, po_estimator: Optional[Any] = None, te_estimator: Optional[Any] = None, n_folds: int = 2, n_layers_out: int = 2, n_layers_out_t: int = 2, n_units_out: int = 100, n_units_out_t: int = 100, n_units_out_prop: int = 100, n_layers_out_prop: int = 0, weight_decay: float = 0.0001, weight_decay_t: float = 0.0001, lr: float = 0.0001, lr_t: float = 0.0001, n_iter: int = 10000, batch_size: int = 100, val_split_prop: float = 0.3, n_iter_print: int = 50, seed: int = 42, nonlin: str = 'elu', weighting_strategy: Optional[str] = 'prop', patience: int = 10, n_iter_min: int = 200, batch_norm: bool = True, early_stopping: bool = True, dropout: bool = False, dropout_prob: float = 0.2)

Bases: catenets.models.torch.pseudo_outcome_nets.PseudoOutcomeLearner

DR-learner for CATE estimation, based on doubly robust AIPW pseudo-outcome

_backward_hooks: Dict[int, Callable]

_buffers: Dict[str, Optional[torch.Tensor]]

_first_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, fit_mask: torch.Tensor, pred_mask: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

_forward_hooks: Dict[int, Callable]

_forward_pre_hooks: Dict[int, Callable]

_is_full_backward_hook: Optional[bool]

_load_state_dict_post_hooks: Dict[int, Callable]

_load_state_dict_pre_hooks: Dict[int, Callable]

_modules: Dict[str, Optional[Module]]

_non_persistent_buffers_set: Set[str]

_parameters: Dict[str, Optional[torch.nn.parameter.Parameter]]

_second_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, p: torch.Tensor, mu_0: torch.Tensor, mu_1: torch.Tensor) → None

_state_dict_hooks: Dict[int, Callable]

training: bool

class PWLearner(n_unit_in: int, binary_y: bool, po_estimator: Optional[Any] = None, te_estimator: Optional[Any] = None, n_folds: int = 2, n_layers_out: int = 2, n_layers_out_t: int = 2, n_units_out: int = 100, n_units_out_t: int = 100, n_units_out_prop: int = 100, n_layers_out_prop: int = 0, weight_decay: float = 0.0001, weight_decay_t: float = 0.0001, lr: float = 0.0001, lr_t: float = 0.0001, n_iter: int = 10000, batch_size: int = 100, val_split_prop: float = 0.3, n_iter_print: int = 50, seed: int = 42, nonlin: str = 'elu', weighting_strategy: Optional[str] = 'prop', patience: int = 10, n_iter_min: int = 200, batch_norm: bool = True, early_stopping: bool = True, dropout: bool = False, dropout_prob: float = 0.2)

Bases: catenets.models.torch.pseudo_outcome_nets.PseudoOutcomeLearner

PW-learner for CATE estimation, based on singly robust Horvitz Thompson pseudo-outcome

_backward_hooks: Dict[int, Callable]

_buffers: Dict[str, Optional[torch.Tensor]]

_first_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, fit_mask: torch.Tensor, pred_mask: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

_forward_hooks: Dict[int, Callable]

_forward_pre_hooks: Dict[int, Callable]

_is_full_backward_hook: Optional[bool]

_load_state_dict_post_hooks: Dict[int, Callable]

_load_state_dict_pre_hooks: Dict[int, Callable]

_modules: Dict[str, Optional[Module]]

_non_persistent_buffers_set: Set[str]

_parameters: Dict[str, Optional[torch.nn.parameter.Parameter]]

_second_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, p: torch.Tensor, mu_0: torch.Tensor, mu_1: torch.Tensor) → None

_state_dict_hooks: Dict[int, Callable]

training: bool

class PseudoOutcomeLearner(n_unit_in: int, binary_y: bool, po_estimator: Optional[Any] = None, te_estimator: Optional[Any] = None, n_folds: int = 2, n_layers_out: int = 2, n_layers_out_t: int = 2, n_units_out: int = 100, n_units_out_t: int = 100, n_units_out_prop: int = 100, n_layers_out_prop: int = 0, weight_decay: float = 0.0001, weight_decay_t: float = 0.0001, lr: float = 0.0001, lr_t: float = 0.0001, n_iter: int = 10000, batch_size: int = 100, val_split_prop: float = 0.3, n_iter_print: int = 50, seed: int = 42, nonlin: str = 'elu', weighting_strategy: Optional[str] = 'prop', patience: int = 10, n_iter_min: int = 200, batch_norm: bool = True, early_stopping: bool = True, dropout: bool = False, dropout_prob: float = 0.2)

Bases: catenets.models.torch.base.BaseCATEEstimator

Class implements TwoStepLearners based on pseudo-outcome regression as discussed in Curth &vd Schaar (2021): RA-learner, PW-learner and DR-learner

Parameters

• n_unit_in (int) – Number of features

• binary_y (bool, default False) – Whether the outcome is binary

• po_estimator (sklearn/PyTorch model, default: None) – Custom potential outcome model. If this parameter is set, the rest of the parameters are ignored.

• te_estimator (sklearn/PyTorch model, default: None) – Custom treatment effects model. If this parameter is set, the rest of the parameters are ignored.

• n_folds (int, default 1) – Number of cross-fitting folds. If 1, no cross-fitting

• n_layers_out (int) – First stage Number of hypothesis layers (n_layers_out x n_units_out + 1 x Linear layer)

• n_units_out (int) – First stage Number of hidden units in each hypothesis layer

• n_layers_r (int) – Number of shared & private representation layers before hypothesis layers

• n_units_r (int) – Number of hidden units in representation shared before the hypothesis layers.

• n_layers_out_t (int) – Second stage Number of hypothesis layers (n_layers_out x n_units_out + 1 x Linear layer)

• n_units_out_t (int) – Second stage Number of hidden units in each hypothesis layer

• n_layers_out_prop (int) – Number of hypothesis layers for propensity score(n_layers_out x n_units_out + 1 x Dense layer)

• n_units_out_prop (int) – Number of hidden units in each propensity score hypothesis layer

• weight_decay (float) – First stage l2 (ridge) penalty

• weight_decay_t (float) – Second stage l2 (ridge) penalty

• lr (float) – First stage learning rate for optimizer

• lr – Second stage learning rate for optimizer

• n_iter (int) – Maximum number of iterations

• batch_size (int) – Batch size

• val_split_prop (float) – Proportion of samples used for validation split (can be 0)

• n_iter_print (int) – Number of iterations after which to print updates

• seed (int) – Seed used

• nonlin (string, default 'elu') – Nonlinearity to use in NN. Can be ‘elu’, ‘relu’, ‘selu’ or ‘leaky_relu’.

• weighting_strategy (str, default "prop") – Weighting strategy. Can be “prop” or “1-prop”.

• patience (int) – Number of iterations to wait before early stopping after decrease in validation loss

• n_iter_min (int) – Minimum number of iterations to go through before starting early stopping

_backward_hooks: Dict[int, Callable]

_buffers: Dict[str, Optional[torch.Tensor]]

abstract _first_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, fit_mask: torch.Tensor, pred_mask: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

_forward_hooks: Dict[int, Callable]

_forward_pre_hooks: Dict[int, Callable]

_generate_po_estimator(name: str = 'po_estimator') → torch.nn.modules.module.Module

_generate_propensity_estimator(name: str = 'propensity_estimator') → torch.nn.modules.module.Module

_generate_te_estimator(name: str = 'te_estimator') → torch.nn.modules.module.Module

_impute_pos(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, fit_mask: torch.Tensor, pred_mask: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

_impute_propensity(X: torch.Tensor, w: torch.Tensor, fit_mask: torch._VariableFunctionsClass.tensor, pred_mask: torch.Tensor) → torch.Tensor

_impute_unconditional_mean(X: torch.Tensor, y: torch.Tensor, fit_mask: torch.Tensor, pred_mask: torch.Tensor) → torch.Tensor

_is_full_backward_hook: Optional[bool]

_load_state_dict_post_hooks: Dict[int, Callable]

_load_state_dict_pre_hooks: Dict[int, Callable]

_modules: Dict[str, Optional[Module]]

_non_persistent_buffers_set: Set[str]

_parameters: Dict[str, Optional[torch.nn.parameter.Parameter]]

abstract _second_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, p: torch.Tensor, mu_0: torch.Tensor, mu_1: torch.Tensor) → None

_state_dict_hooks: Dict[int, Callable]

fit(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor) → catenets.models.torch.pseudo_outcome_nets.PseudoOutcomeLearner

Train treatment effects nets.

Parameters

• X (array-like of shape (n_samples, n_features)) – Train-sample features

• y (array-like of shape (n_samples,)) – Train-sample labels

• w (array-like of shape (n_samples,)) – Train-sample treatments

predict(X: torch.Tensor, return_po: bool = False, training: bool = False) → torch.Tensor

Predict treatment effects

Parameters

X (array-like of shape (n_samples, n_features)) – Test-sample features

Returns

te_est – Predicted treatment effects

Return type

array-like of shape (n_samples,)

training: bool

class RALearner(n_unit_in: int, binary_y: bool, po_estimator: Optional[Any] = None, te_estimator: Optional[Any] = None, n_folds: int = 2, n_layers_out: int = 2, n_layers_out_t: int = 2, n_units_out: int = 100, n_units_out_t: int = 100, n_units_out_prop: int = 100, n_layers_out_prop: int = 0, weight_decay: float = 0.0001, weight_decay_t: float = 0.0001, lr: float = 0.0001, lr_t: float = 0.0001, n_iter: int = 10000, batch_size: int = 100, val_split_prop: float = 0.3, n_iter_print: int = 50, seed: int = 42, nonlin: str = 'elu', weighting_strategy: Optional[str] = 'prop', patience: int = 10, n_iter_min: int = 200, batch_norm: bool = True, early_stopping: bool = True, dropout: bool = False, dropout_prob: float = 0.2)

Bases: catenets.models.torch.pseudo_outcome_nets.PseudoOutcomeLearner

RA-learner for CATE estimation, based on singly robust regression-adjusted pseudo-outcome

_backward_hooks: Dict[int, Callable]

_buffers: Dict[str, Optional[torch.Tensor]]

_first_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, fit_mask: torch.Tensor, pred_mask: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

_forward_hooks: Dict[int, Callable]

_forward_pre_hooks: Dict[int, Callable]

_is_full_backward_hook: Optional[bool]

_load_state_dict_post_hooks: Dict[int, Callable]

_load_state_dict_pre_hooks: Dict[int, Callable]

_modules: Dict[str, Optional[Module]]

_non_persistent_buffers_set: Set[str]

_parameters: Dict[str, Optional[torch.nn.parameter.Parameter]]

_second_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, p: torch.Tensor, mu_0: torch.Tensor, mu_1: torch.Tensor) → None

_state_dict_hooks: Dict[int, Callable]

training: bool

class RLearner(n_unit_in: int, binary_y: bool, po_estimator: Optional[Any] = None, te_estimator: Optional[Any] = None, n_folds: int = 2, n_layers_out: int = 2, n_layers_out_t: int = 2, n_units_out: int = 100, n_units_out_t: int = 100, n_units_out_prop: int = 100, n_layers_out_prop: int = 0, weight_decay: float = 0.0001, weight_decay_t: float = 0.0001, lr: float = 0.0001, lr_t: float = 0.0001, n_iter: int = 10000, batch_size: int = 100, val_split_prop: float = 0.3, n_iter_print: int = 50, seed: int = 42, nonlin: str = 'elu', weighting_strategy: Optional[str] = 'prop', patience: int = 10, n_iter_min: int = 200, batch_norm: bool = True, early_stopping: bool = True, dropout: bool = False, dropout_prob: float = 0.2)

Bases: catenets.models.torch.pseudo_outcome_nets.PseudoOutcomeLearner

R-learner for CATE estimation. Based on pseudo-outcome (Y-mu(x))/(w-pi(x)) and sample weight (w-pi(x))^2 – can only be implemented if .fit of te_estimator takes argument ‘sample_weight’.

_backward_hooks: Dict[int, Callable]

_buffers: Dict[str, Optional[torch.Tensor]]

_first_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, fit_mask: torch.Tensor, pred_mask: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

_forward_hooks: Dict[int, Callable]

_forward_pre_hooks: Dict[int, Callable]

_is_full_backward_hook: Optional[bool]

_load_state_dict_post_hooks: Dict[int, Callable]

_load_state_dict_pre_hooks: Dict[int, Callable]

_modules: Dict[str, Optional[Module]]

_non_persistent_buffers_set: Set[str]

_parameters: Dict[str, Optional[torch.nn.parameter.Parameter]]

_second_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, p: torch.Tensor, mu_0: torch.Tensor, mu_1: torch.Tensor) → None

_state_dict_hooks: Dict[int, Callable]

training: bool

class ULearner(n_unit_in: int, binary_y: bool, po_estimator: Optional[Any] = None, te_estimator: Optional[Any] = None, n_folds: int = 2, n_layers_out: int = 2, n_layers_out_t: int = 2, n_units_out: int = 100, n_units_out_t: int = 100, n_units_out_prop: int = 100, n_layers_out_prop: int = 0, weight_decay: float = 0.0001, weight_decay_t: float = 0.0001, lr: float = 0.0001, lr_t: float = 0.0001, n_iter: int = 10000, batch_size: int = 100, val_split_prop: float = 0.3, n_iter_print: int = 50, seed: int = 42, nonlin: str = 'elu', weighting_strategy: Optional[str] = 'prop', patience: int = 10, n_iter_min: int = 200, batch_norm: bool = True, early_stopping: bool = True, dropout: bool = False, dropout_prob: float = 0.2)

Bases: catenets.models.torch.pseudo_outcome_nets.PseudoOutcomeLearner

U-learner for CATE estimation. Based on pseudo-outcome (Y-mu(x))/(w-pi(x))

_backward_hooks: Dict[int, Callable]

_buffers: Dict[str, Optional[torch.Tensor]]

_first_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, fit_mask: torch.Tensor, pred_mask: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

_forward_hooks: Dict[int, Callable]

_forward_pre_hooks: Dict[int, Callable]

_is_full_backward_hook: Optional[bool]

_load_state_dict_post_hooks: Dict[int, Callable]

_load_state_dict_pre_hooks: Dict[int, Callable]

_modules: Dict[str, Optional[Module]]

_non_persistent_buffers_set: Set[str]

_parameters: Dict[str, Optional[torch.nn.parameter.Parameter]]

_second_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, p: torch.Tensor, mu_0: torch.Tensor, mu_1: torch.Tensor) → None

_state_dict_hooks: Dict[int, Callable]

training: bool

class XLearner(*args: Any, weighting_strategy: str = 'prop', **kwargs: Any)

Bases: catenets.models.torch.pseudo_outcome_nets.PseudoOutcomeLearner

X-learner for CATE estimation. Combines two CATE estimates via a weighting function g(x): tau(x) = g(x) tau_0(x) + (1-g(x)) tau_1(x)

_backward_hooks: Dict[int, Callable]

_buffers: Dict[str, Optional[torch.Tensor]]

_first_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, fit_mask: torch.Tensor, pred_mask: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

_forward_hooks: Dict[int, Callable]

_forward_pre_hooks: Dict[int, Callable]

_is_full_backward_hook: Optional[bool]

_load_state_dict_post_hooks: Dict[int, Callable]

_load_state_dict_pre_hooks: Dict[int, Callable]

_modules: Dict[str, Optional[Module]]

_non_persistent_buffers_set: Set[str]

_parameters: Dict[str, Optional[torch.nn.parameter.Parameter]]

_second_step(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor, p: torch.Tensor, mu_0: torch.Tensor, mu_1: torch.Tensor) → None

_state_dict_hooks: Dict[int, Callable]

predict(X: torch.Tensor, return_po: bool = False, training: bool = False) → torch.Tensor

Predict treatment effects

Parameters

• X (array-like of shape (n_samples, n_features)) – Test-sample features

• return_po (bool, default False) – Whether to return potential outcome predictions. Placeholder, can only accept False.

Returns

te_est – Predicted treatment effects

Return type

array-like of shape (n_samples,)

training: bool