<commit_before>import numpy as np
import random
import itertools
import statsmodels.api as sm
from scipy.stats import norm
from scipy.optimize import fmin_bfgs
from math import sqrt
import cvxpy as cvx
from sklearn import linear_model


class CausalModel(object):


	def __init__(self, Y, D, X):

		self._Y, self._D, self._X = Y, D, X
		self._N, self._K = self._X.shape
		self._treated, self._controls = np.nonzero(D)[0], np.nonzero(D==0)[0]
		self._Y_t, self._Y_c = self._Y[self._treated], self._Y[self._controls]
		self._X_t, self._X_c = self._X[self._treated], self._X[self._controls]
		self._N_t = np.count_nonzero(D)
		self._N_c = self._N - self._N_t


	def normalized_diff(self):

		"""
		Computes normalized difference in covariates for assessing balance.

		Normalized difference is the difference in group means, scaled by the
		square root of the average of the two within-group variances. Large
		values indicate that simple linear adjustment methods may not be adequate
		for removing biases that are associated with differences in covariates.

		Unlike t-statistic, normalized differences do not, in expectation,
		increase with sample size, and thus is more appropriate for assessing
		balance.
		"""

		return (self._X_t.mean(0) - self._X_c.mean(0)) / np.sqrt((self._X_t.var(0) +
		                                                         self._X_c.var(0))/2)


	def _sigmoid(self, x):

		return 1/(1+np.exp(-x))


	def _log1exp(self, x):

		return np.log(1 + np.exp(-x))


	def _neg_loglike(self, beta, X_t, X_c):

		return self._log1exp(X_t.dot(beta)).sum() + \
		       self._log1exp(-X_c.dot(beta)).sum()


	def _neg_gradient(self, beta, X_t, X_c):

		return (self._sigmoid(X_c.dot(beta))*X_c.T).sum(1) - \
		       (self._sigmoid(-X_t.dot(beta))*X_t.T).sum(1)


	def _pscore(self, X):

		X_t = X[self._treated]
		X_c = X[self._controls]
		K = X.shape[1]

		neg_loglike = lambda x: self._neg_loglike(x, X_t, X_c)
		neg_gradient = lambda x: self._neg_gradient(x, X_t, X_c)

		logit = fmin_bfgs(neg_loglike, np.zeros(K), neg_gradient, full_output=True, disp=False)

		beta = logit[0]
		loglike = -logit[1]
		fitted = np.empty(self._N)
		fitted[self._treated] = self._sigmoid(X_t.dot(beta))
		fitted[self._controls] = self._sigmoid(X_c.dot(beta))

		return (beta, loglike, fitted)


	def _form_matrix(self, const, lin, qua):

		mat = np.empty((self._N, const+len(lin)+len(qua)))

		current_col = 0
		if const:
			mat[:, current_col] = 1
			current_col += 1
		if lin:
			mat[:, current_col:current_col+len(lin)] = self._X[:, lin]
			current_col += len(lin)
		for term in qua:
			mat[:, current_col] = self._X[:, term[0]] * self._X[:, term[1]]
			current_col += 1

		return mat


	def pscore(self, const=True, lin=None, qua=[]):

		if lin is None:
			lin = xrange(self._X.shape[1])
		if qua:
			qua = list(itertools.combinations_with_replacement(qua, 2))

		return self._pscore(self._form_matrix(const, lin, qua))
		

	def _pscore_select(self, const, X_cur, X_pot, crit, X_lin=[]):
	
		if not X_pot:
			return X_cur

		if not X_lin:  # X_lin is empty, so linear terms not yet decided
			ll_null = self._pscore(self._form_matrix(const, X_cur, []))[1]
		else:  # X_lin is not empty, so linear terms are already fixed
			ll_null = self._pscore(self._form_matrix(const, X_lin, X_cur))[1]

		lr = np.empty(len(X_pot))
		if not X_lin:
			for i in xrange(len(X_pot)):
				lr[i] = 2*(self._pscore(self._form_matrix(const, X_cur+[X_pot[i]], []))[1] - ll_null)
		else:
			for i in xrange(len(X_pot)):
				lr[i] = 2*(self._pscore(self._form_matrix(const, X_lin, X_cur+[X_pot[i]]))[1] - ll_null)

		argmax = np.argmax(lr)
		if lr[argmax] < crit:
			return X_cur
		else:
			new_term = X_pot.pop(argmax)
			return self._pscore_select(const, X_cur+[new_term], X_pot, crit, X_lin)


	def pscore_select(self, const=True, X_B=[], C_lin=1, C_qua=2.71):<commit_msg>documentation for pscore shenanigans<commit_after>import numpy as np
import random
import itertools
import statsmodels.api as sm
from scipy.stats import norm
from scipy.optimize import fmin_bfgs
from math import sqrt
import cvxpy as cvx
from sklearn import linear_model


class CausalModel(object):


	def __init__(self, Y, D, X):

		self._Y, self._D, self._X = Y, D, X
		self._N, self._K = self._X.shape
		self._treated, self._controls = np.nonzero(D)[0], np.nonzero(D==0)[0]
		self._Y_t, self._Y_c = self._Y[self._treated], self._Y[self._controls]
		self._X_t, self._X_c = self._X[self._treated], self._X[self._controls]
		self._N_t = np.count_nonzero(D)
		self._N_c = self._N - self._N_t


	def normalized_diff(self):

		"""
		Computes normalized difference in covariates for assessing balance.

		Normalized difference is the difference in group means, scaled by the
		square root of the average of the two within-group variances. Large
		values indicate that simple linear adjustment methods may not be adequate
		for removing biases that are associated with differences in covariates.

		Unlike t-statistic, normalized differences do not, in expectation,
		increase with sample size, and thus is more appropriate for assessing
		balance.

		Returns
		-------
			Vector of normalized differences.
		"""

		return (self._X_t.mean(0) - self._X_c.mean(0)) / np.sqrt((self._X_t.var(0) +
		                                                         self._X_c.var(0))/2)


	def _sigmoid(self, x):
	
		"""
		Computes 1/(1+exp(-x)) for input x, to be used in maximum likelihood
		estimation of propensity score.

		Arguments
		---------
			x: array-like

		Returns
		-------
			Vector or scalar 1/(1+exp(-x)), depending on input x.
		"""

		return 1/(1+np.exp(-x))


	def _log1exp(self, x):

		"""
		Computes log(1+exp(-x)) for input x, to be used in maximum likelihood
		estimation of propensity score.

		Arguments
		---------
			x: array-like

		Returns
		-------
			Vector or scalar log(1+exp(-x)), depending on input x.
		"""

		return np.log(1 + np.exp(-x))


	def _neg_loglike(self, beta, X_t, X_c):

		"""
		Computes the negative of the log likelihood function for logit, to be used
		in maximum likelihood estimation of propensity score. Negative because SciPy
		optimizier does minimization only.

		Arguments
		---------
			beta: array-like
				Logisitic regression parameters to maximize over.
			X_t: array-like
				Covariate matrix of the treated units.
			X_c: array-like
				Covariate matrix of the control units.

		Returns
		-------
			Negative log likelihood evaluated at input values.
		"""

		return self._log1exp(X_t.dot(beta)).sum() + \
		       self._log1exp(-X_c.dot(beta)).sum()


	def _neg_gradient(self, beta, X_t, X_c):

		"""
		Computes the negative of the gradient of the log likelihood function for
		logit, to be used in maximum likelihood estimation of propensity score.
		Negative because SciPy optimizier does minimization only.

		Arguments
		---------
			beta: array-like
				Logisitic regression parameters to maximize over.
			X_t: array-like
				Covariate matrix of the treated units.
			X_c: array-like
				Covariate matrix of the control units.

		Returns
		-------
			Negative gradient of log likelihood function evaluated at input values.
		"""

		return (self._sigmoid(X_c.dot(beta))*X_c.T).sum(1) - \
		       (self._sigmoid(-X_t.dot(beta))*X_t.T).sum(1)


	def _pscore(self, X):

		"""
		Estimates via logit the propensity score based on input covariate matrix X.

		Arguments
		---------
			X: array-like
				Covariate matrix to estimate propensity score on.

		Returns
		-------
			beta: array-like
				Estimated logistic regression coefficients.
			loglike: scalar
				Maximized log-likelihood value.
			fitted: array-like
				Estimated propensity scores for each unit.
		"""

		X_t = X[self._treated]
		X_c = X[self._controls]
		K = X.shape[1]

		neg_loglike = lambda x: self._neg_loglike(x, X_t, X_c)
		neg_gradient = lambda x: self._neg_gradient(x, X_t, X_c)

		logit = fmin_bfgs(neg_loglike, np.zeros(K), neg_gradient, full_output=True, disp=False)

		beta = logit[0]
		loglike = -logit[1]
		fitted = np.empty(self._N)
		fitted[self._treated] = self._sigmoid(X_t.dot(beta))
		fitted[self._controls] = self._sigmoid(X_c.dot(beta))

		return (beta, loglike, fitted)


	def _form_matrix(self, const, lin, qua):

		"""
		Forms covariate matrix for use in propensity score estimation, based on
		requirements on constant term, linear terms, and quadratic terms.

		Arguments
		---------
			const: Boolean
				Includes a column of one's if True.
			lin: list
				Column numbers of the base self._X covariate matrix
				to include linearly.
			qua: list
				Tuples indicating which columns of the base self._X
				covariate matrix to multiply and include. E.g.,
				[(0,0), (1,2)] indicates squaring the 0th column and
				including the product of the 1st and 2nd columns.

		Returns
		-------
			mat: array-like
				Covariate matrix formed based on requirements on
				constant, linear, and quadratic terms.
		"""

		mat = np.empty((self._N, const+len(lin)+len(qua)))

		current_col = 0
		if const:
			mat[:, current_col] = 1
			current_col += 1
		if lin:
			mat[:, current_col:current_col+len(lin)] = self._X[:, lin]
			current_col += len(lin)
		for term in qua:
			mat[:, current_col] = self._X[:, term[0]] * self._X[:, term[1]]
			current_col += 1

		return mat


	def pscore(self, const=True, lin=None, qua=[]):

		"""
		Estimates via logit the propensity score based on requirements on
		constant term, linear terms, and quadratic terms.

		Arguments
		---------
			const: Boolean
				Includes a column of one's if True. Defaults to
				True.
			lin: list
				Column numbers of the base covariate matrix X
				to include linearly. Defaults to using the
				whole covariate matrix.
			qua: list
				Column numbers of the base covariate matrix X
				to include quadratic. E.g., [0,2] will include
				squares of the 0th and 2nd columns, and the product
				of these two columns. Default is not include any
				quadratic terms.

		Returns
		-------
			beta: array-like
				Estimated logistic regression coefficients.
			loglike: scalar
				Maximized log-likelihood value.
			fitted: array-like
				Estimated propensity scores for each unit.
		"""

		if lin is None:
			lin = xrange(self._X.shape[1])
		if qua:
			qua = list(itertools.combinations_with_replacement(qua, 2))

		return self._pscore(self._form_matrix(const, lin, qua))
		

	def _pscore_select(self, const, X_cur, X_pot, crit, X_lin=[]):
	
		"""
		Estimates via logit the propensity score using Imbens and Rubin's
		covariate selection algorithm.

		Arguments
		---------
			const: Boolean
				Includes a column of one's if True. 
			X_cur: list
				List containing terms that are currently included
				in the logistic regression.
			X_pot: list
				List containing candidate terms to be iterated through.
			crit: scalar
				Critical value used in likelihood ratio test to decide
				whether candidate terms should be included.
			X_lin: list
				List containing linear terms that have been decided on.
				If non-empty, then X_cur and X_pot should be containing
				candidate quadratic terms. If empty, then those two
				matrices should be containing candidate linear terms.

		Returns
		-------
			List containing terms that the algorithm has settled on for inclusion.
		"""

		if not X_pot:
			return X_cur

		if not X_lin:  # X_lin is empty, so linear terms not yet decided
			ll_null = self._pscore(self._form_matrix(const, X_cur, []))[1]
		else:  # X_lin is not empty, so linear terms are already fixed
			ll_null = self._pscore(self._form_matrix(const, X_lin, X_cur))[1]

		lr = np.empty(len(X_pot))
		if not X_lin:
			for i in xrange(len(X_pot)):
				lr[i] = 2*(self._pscore(self._form_matrix(const, X_cur+[X_pot[i]], []))[1] - ll_null)
		else:
			for i in xrange(len(X_pot)):
				lr[i] = 2*(self._pscore(self._form_matrix(const, X_lin, X_cur+[X_pot[i]]))[1] - ll_null)

		argmax = np.argmax(lr)
		if lr[argmax] < crit:
			return X_cur
		else:
			new_term = X_pot.pop(argmax)
			return self._pscore_select(const, X_cur+[new_term], X_pot, crit, X_lin)


	def pscore_select(self, const=True, X_B=[], C_lin=1, C_qua=2.71):

		"""
		Estimates via logit the propensity score using Imbens and Rubin's
		covariate selection algorithm.

		Arguments
		---------
			const: Boolean
				Includes a column of one's if True. Defaults to
				True.
			X_B: list
				Column numbers of the base covariate matrix X
				that should be included as linear terms
				regardless. Defaults to empty list, meaning
				every column of X is subjected to the selection
				algorithm.
			C_lin: scalar
				Critical value used in likelihood ratio test to decide
				whether candidate linear terms should be included.
				Defaults to 1 as in Imbens (2014).
			C_qua: scalar
				Critical value used in likelihood ratio test to decide
				whether candidate quadratic terms should be included.
				Defaults to 2.71 as in Imbens (2014).

		Returns
		-------
			beta: array-like
				Estimated logistic regression coefficients.
			loglike: scalar
				Maximized log-likelihood value.
			fitted: array-like
				Estimated propensity scores for each unit.

		References
		----------
			<NAME>. & <NAME>. (2015). Causal Inference in Statistics,
				Social, and Biomedical Sciences: An Introduction.
			<NAME>. (2014). Matching Methods in Practice: Three Examples.
		"""
{'fraction_non_alphanumeric': 0.11342453397474443, 'fraction_numerical': 0.008343355381840048, 'mean_word_length': 3.172154280338664, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1390299', 'n_tokens_mistral': 5264, 'n_tokens_neox': 4427, 'n_words': 1408}