catenets.models.jax.transformation_utils module

Utils for transformations

_get_transformation_function(transformation_name: str) → Any

Get transformation function associated with a name

aipw_te_transformation(y: numpy.ndarray, w: numpy.ndarray, p: Optional[numpy.ndarray], mu_0: numpy.ndarray, mu_1: numpy.ndarray) → numpy.ndarray

Transforms data to efficient influence function pseudo-outcome for CATE estimation

Parameters

  • y (array-like of shape (n_samples,) or (n_samples, )) – The observed outcome variable

  • w (array-like of shape (n_samples,)) – The observed treatment indicator

  • p (array-like of shape (n_samples,)) – The treatment propensity, estimated or known. Can be None, then p=0.5 is assumed

  • mu_0 (array-like of shape (n_samples,)) – Estimated or known potential outcome mean of the control group

  • mu_1 (array-like of shape (n_samples,)) – Estimated or known potential outcome mean of the treatment group

Returns

EIF transformation for CATE

Return type

d_hat

ht_te_transformation(y: numpy.ndarray, w: numpy.ndarray, p: Optional[numpy.ndarray] = None, mu_0: Optional[numpy.ndarray] = None, mu_1: Optional[numpy.ndarray] = None) → numpy.ndarray

Transform data to Horvitz-Thompson transformation for CATE

Parameters

  • y (array-like of shape (n_samples,) or (n_samples, )) – The observed outcome variable

  • w (array-like of shape (n_samples,)) – The observed treatment indicator

  • p (array-like of shape (n_samples,)) – The treatment propensity, estimated or known. Can be None, then p=0.5 is assumed

  • mu_0 (array-like of shape (n_samples,)) – Placeholder, not used. Estimated or known potential outcome mean of the control group

  • mu_1 (array-like of shape (n_samples,)) – Placerholder, not used. Estimated or known potential outcome mean of the treatment group

Returns

res – Horvitz-Thompson transformed data

Return type

array-like of shape (n_samples,)

ra_te_transformation(y: numpy.ndarray, w: numpy.ndarray, p: Optional[numpy.ndarray], mu_0: numpy.ndarray, mu_1: numpy.ndarray) → numpy.ndarray

Transform data to regression adjustment for CATE

Parameters

  • y (array-like of shape (n_samples,) or (n_samples, )) – The observed outcome variable

  • w (array-like of shape (n_samples,)) – The observed treatment indicator

  • p (array-like of shape (n_samples,)) – Placeholder, not used. The treatment propensity, estimated or known.

  • mu_0 (array-like of shape (n_samples,)) – Estimated or known potential outcome mean of the control group

  • mu_1 (array-like of shape (n_samples,)) – Estimated or known potential outcome mean of the treatment group

Returns

res – Regression adjusted transformation

Return type

array-like of shape (n_samples,)