<commit_before>
import numpy as np
import random
import pandas as pd
from scipy.stats import norm
from sklearn import linear_model

def EstimateWeights(X_c, X_t):

	"""
	Function that estimates synthetic control weights for each treated unit.

	For each treated unit, we find w that minimizes:
		|| w' X_c - X_t ||^2,

	Since the number of controls should always be greater than the number of
	covariates, the linear system is underdetermined, and least squares	is
	used to obtain a solution w.

	Args:
		X_c = N_c-by-k matrix of control units
		X_t = N_t-by-k matrix of treated units

	Returns:
		W_hat = N_t-by-N_c matrix of estimated weights
	"""

	N_c, N_t = len(X_c), len(X_t)

	W_hat = np.zeros(shape=(N_t, N_c))  # matrix of weight estimates

	for i in xrange(N_t):
		W_hat[i, ] = np.linalg.lstsq(X_c.T, X_t[i, ])[0]

	return W_hat

def TreatmentEffects(Y_c, Y_t, W, return_individual=False):

	"""
	Function that estimates individual and average treatment effects
	via synthetic control method, using provided synthetic control weights.

	Args:
		Y_c = N_c-dimensional array of control unit outcomes
		Y_t = N_t-dimensional array of treated unit outcomes
		W = N_t-by-N_c matrix of synthetic control weights

	Returns:
		ITT = N_t-dimensional array of estimated individual treatment effects
		      for the treated sample
		ATT = Scalar value of estaimted average treatment effect on the treated
	"""

	ITT = Y_t - np.dot(W, Y_c)
	ATT = ITT.mean()

	if return_individual:
		return ATT, ITT
	else:
		return ATT

def SyntheticEstimates(Y, D, X, higher_moments=False):

	"""
	Function that computes ATT estimates from observerd data using
	synthetic control method.

	If higher_moments is true, then matching with be done on higher
	moments of X as well. For now this only includes absolute first
	moments, and second moments.
	"""

	control = (D == 0)  # Boolean index of control units
	treated = (D == 1)  # Boolean index of treated units
	if higher_moments:
		X_control = np.hstack((X[control], abs(X[control]), X[control]**2))
		X_treated = np.hstack((X[treated], abs(X[treated]), X[treated]**2))
		W_hat = EstimateWeights(X_control, X_treated)
	else:
		W_hat = EstimateWeights(X[control], X[treated])  # synthetic control weights

	return TreatmentEffects(Y[control], Y[treated], W_hat)

def MatchingWithReplacement(Y, D, X):

	X_c, X_t, Y_c, Y_t = X[D==0], X[D==1], Y[D==0], Y[D==1]

	N_t = len(X_t)

	ITT = np.zeros(N_t)

	for i in xrange(N_t):
		match_index = ((X_c - X_t[i])**2).sum(1).argmin()
		ITT[i] = Y_t[i] - Y_c[match_index]

	return ITT.mean()

def MatchingWithoutReplacement(Y, D, X):

	if 2*sum(D) > len(D):  # if N_t > N_c
		print 'Not enough control units, matching with replacement instead.'
		return MatchingWithReplacement(Y, D, X)

	pscore = sm.Logit(D, X).fit(disp=False).predict()  # estimate pscore with logit
	order = np.argsort(pscore)[::-1]  # sort pscore in descending order, get index
	Y, D, X = Y[order], D[order], X[order]  # any way to do it more elegantly?

	X_c, X_t, Y_c, Y_t = X[D==0], X[D==1], Y[D==0], Y[D==1]

	N_c, N_t = len(X_c), len(X_t)

	ITT = np.zeros(N_t)

	for i in xrange(N_t):
		match_index = ((X_c - X_t[i])**2).sum(1).argmin()
		ITT[i] = Y_t[i] - Y_c[match_index]
		X_c = np.delete(X_c, match_index, axis=0)
		Y_c = np.delete(Y_c, match_index, axis=0)

	return ITT.mean()

def MatchingEstimates(Y, D, X, with_replacement=True):

	"""
	Function that estimates the average treatment effect for the treated (ATT)
	using matching.

	Can specify whether matching is done with or without replacement. Swtiches
	to matching with replacement if number of control units is less than
	number of treated units.

	Args:
		Y = N-dimensional array of observed outcomes
		D = N-dimensional array of treatment indicator; 1=treated, 0=control
		X = N-by-k matrix of covariates

	Returns:
		ATT estimate; matching done with or without replacement as specified
	"""

	if with_replacement:
		return MatchingWithReplacement(Y, D, X)
	else:
		return MatchingWithoutReplacement(Y, D, X)

def OLSEstimates(Y, D, X):

	"""
	Function that estimates average treatment effect for the treated (ATT)
	using OLS, which is consistent when the true model is linear.

	Args:
		Y = N-dimensional array of observed outcomes
		D = N-dimensional array of treatment indicator; 1=treated, 0=control
		X = N-by-k matrix of covariates
	"""

	k = X.shape[1]
	D = D.reshape((D.size, 1))  # convert D into N-by-1 vector
	dX = X - X.mean(0)  # demean covariates
	DdX = D * dX
	# construct design matrix; no constant term as it will be added by sklearn
	W = np.column_stack((D, dX, DdX))

	# use sklearn to run regression
	reg = linear_model.LinearRegression()
	reg.fit(W, Y)

	# for derivation of this estimator, see my notes on Treatment Effects
	return reg.coef_[0] + np.dot((DdX.sum(0) / D.sum()), reg.coef_[-k:])

class parameters:

	"""
	Class object that stores the parameter values for use in the baseline model.

	See SimulateData function for model description.

	Args:
		N = Sample size (control + treated units) to generate
		k = Number of covariates
	"""

	def __init__(self, N=500, k=3):  # set initial parameter values
		self.N = N  # sample size (control + treated units)
		self.k = k  # number of covariates

		self.delta = 3
		self.beta = np.ones(k)
		self.theta = np.ones(k)
		self.mu = np.zeros(k)
		self.Sigma = np.identity(k)
		self.Gamma = np.identity(2)

def SimulateData(para=parameters(), nonlinear=False, return_counterfactual=False):

	"""
	Function that generates data according to one of two simple models that
	satisfies the Unconfoundedness assumption.

	The covariates and error terms are generated according to
		X ~ N(mu, Sigma), epsilon ~ N(0, Gamma).
	The counterfactual outcomes are generated by
		Y_0 = X*beta + epsilon_0,
		Y_1 = delta + X*(beta+theta) + epsilon_1,
	if the nonlinear Boolean is False, or by
		Y_0 = sum_of_abs(X) + epsilon_0,
		Y_1 = sum_of_squares(X) + epsilon_1,
	if the nonlinear Boolean is True.

	Selection is done according to the following propensity score function:
		P(D=1|X) = Phi(X*beta),
	where Phi is the standard normal CDF.

	Args:
		para = Model parameter values supplied by parameter class object
		nonlinear = Boolean indicating whether the data generating model should
		            be linear or not
		return_counterfactual = Boolean indicating whether to return vectors of
		                        counterfactual outcomes

	Returns:
		Y = N-dimensional array of observed outcomes
		D = N-dimensional array of treatment indicator; 1=treated, 0=control
		X = N-by-k matrix of covariates
		Y_0 = N-dimensional array of non-treated outcomes
		Y_1 = N-dimensional array of treated outcomes
	"""

	k = len(para.mu)

	X = np.random.multivariate_normal(mean=para.mu, cov=para.Sigma,
	                                  size=para.N)
	epsilon = np.random.multivariate_normal(mean=np.zeros(2), cov=para.Gamma,
	                                        size=para.N)

	Xbeta = np.dot(X, para.beta)

	pscore = norm.cdf(Xbeta)
	# for each p in pscore, generate Bernoulli rv with success probability p
	D = np.array([np.random.binomial(1, p, size=1) for p in pscore]).flatten()

	if nonlinear:
		Y_0 = abs(X).sum(1) + epsilon[:, 0]
		Y_1 = (X**2).sum(1) + epsilon[:, 1]
	else:
		Y_0 = Xbeta + epsilon[:, 0]
		Y_1 = para.delta + np.dot(X, para.beta+para.theta) + epsilon[:, 1]

	Y = (1 - D) * Y_0 + D * Y_1  # compute observed outcome

	if return_counterfactual:
		return Y, D, X, Y_0, Y_1
	else:
		return Y, D, X

def MonteCarlo(B=500, para=parameters(), nonlinear=False, print_progress=True):

	"""
	Function that returns ATT estimates using synthetic control, matching,
	and OLS computed over B repetitions on simulated data.

	See SimulateData function for data generating process.

	Args:
		B = Number of Monte Carlo simulations to perform
		para = Model parameter values supplied by parameter class object
		nonlinear = Boolean indicating whether the data generating model should
		            be linear or not
		print_progress = Boolean indicating whether to dispaly current number
		                 of completed repetitions (every 10 repetitions)

	Returns:
		ATT_true = Actual average treatment effect for the treated
		ATT_syn = Estimated ATT using synthetic controls
		ATT_match = Estimated ATT using matching
		ATT_ols = Estimated ATT using OLS
	"""

	ATT_true = np.zeros(B)
	ATT_syn = np.zeros(B)
	ATT_match = np.zeros(B)
	ATT_ols = np.zeros(B)

	for i in xrange(B):

		Y, D, X, Y_0, Y_1 = SimulateData(para, nonlinear, True)

		ATT_true[i] = (Y_1[D==1]-Y_0[D==1]).mean()
		ATT_syn[i] = SyntheticEstimates(Y, D, X, higher_moments=nonlinear)
		ATT_match[i] = MatchingEstimates(Y, D, X, with_replacement=False)
		ATT_ols[i] = OLSEstimates(Y, D, X)

		if print_progress and (i+1) % 10 == 0:
			print i+1, 'out of', B, 'simulations completed.'

	return ATT_true, ATT_syn, ATT_match, ATT_ols

def CalculateMSE(B=500, para=parameters(), nonlinear=False):

	"""
	Function that calcuates MSE after performing Monte Carlo simulations.

	Args:
		B = Number of Monte Carlo simulations to perform
		para = Model parameter values supplied by parameter class object
		nonlinear = Boolean indicating whether the data generating model should
		            be linear or not

	Returns:
		MSE_syn = Estimated MSE using synthetic control estimates
		MSE_match = Estimated MSE using matching estimates
		MSE_ols = Estimated MSE using OLS estimates
	"""

	ATT_true, ATT_syn, ATT_match, ATT_ols = MonteCarlo(B, nonlinear=nonlinear)

	MSE_syn = ((ATT_syn - ATT_true)**2).mean()
	MSE_match = ((ATT_match - ATT_true)**2).mean()
	MSE_ols = ((ATT_ols - ATT_true)**2).mean()

	return MSE_syn, MSE_match, MSE_ols

def UseLalondeData():

	"""
	Utility function that reads and applies synthetic control
	estimator on Lalonde's NSW experimental data.

	Note that because of complete randomization, the mean
	difference identifies both the ATE and ATT.
	"""

	url = 'http://www.stanford.edu/~lwong1/data.csv'
	lalonde = pd.read_csv(url)  # read CSV data from url

	covariate_list = ['age', 'education', 'black', 'hispanic', 'married',
	                  'nodegree', 're74', 're75', 'u74', 'u75']

	# don't know how to not convert to array first
	Y = np.array(lalonde['re78'])
	D = np.array(lalonde['treat'])
	X = np.array(lalonde[covariate_list])

	ATT_syn = SyntheticEstimates(Y, D, X)
	ATT_match = MatchingEstimates(Y, D, X)
	ATT_ols = OLSEstimates(Y, D, X)

	print "Using Lalonde's National Supported Work (NSW) experimental data..."
	print 'Estimated ATT using synthetic control estimator:', ATT_syn
	print 'Estimated ATT using matching (with replacement) estimator:', ATT_match
	print 'Estimated ATT using OLS estimator:', ATT_ols
	print 'Mean difference:', Y[D==1].mean() - Y[D==0].mean()

def main():

	#UseLalondeData()

	B = 500
	print '\n' + 'Performing Monte Carlo simulations with B =', B
	MSE_syn, MSE_match, MSE_ols = CalculateMSE(B, nonlinear=True)
	print 'Estimated Mean Squared Error for matching estimator:', MSE_match
	print 'Estimated Mean Squared Error for synthetic control estimator:', MSE_syn
	print 'Estimated Mean Squared Error for OLS estimator:', MSE_ols

if __name__ == '__main__':
	main()
<commit_msg>added option for Mahalanobis metric, but no documentation yet<commit_after>
import numpy as np
import random
import pandas as pd
from scipy.stats import norm
from sklearn import linear_model

def EstimateWeights(X_c, X_t):

	"""
	Function that estimates synthetic control weights for each treated unit.

	For each treated unit, we find w that minimizes:
		|| w' X_c - X_t ||^2,

	Since the number of controls should always be greater than the number of
	covariates, the linear system is underdetermined, and least squares	is
	used to obtain a solution w.

	Args:
		X_c = N_c-by-k matrix of control units
		X_t = N_t-by-k matrix of treated units

	Returns:
		W_hat = N_t-by-N_c matrix of estimated weights
	"""

	N_c, N_t = len(X_c), len(X_t)

	W_hat = np.zeros(shape=(N_t, N_c))  # matrix of weight estimates

	for i in xrange(N_t):
		W_hat[i, ] = np.linalg.lstsq(X_c.T, X_t[i, ])[0]

	return W_hat

def TreatmentEffects(Y_c, Y_t, W, return_individual=False):

	"""
	Function that estimates individual and average treatment effects
	via synthetic control method, using provided synthetic control weights.

	Args:
		Y_c = N_c-dimensional array of control unit outcomes
		Y_t = N_t-dimensional array of treated unit outcomes
		W = N_t-by-N_c matrix of synthetic control weights

	Returns:
		ITT = N_t-dimensional array of estimated individual treatment effects
		      for the treated sample
		ATT = Scalar value of estaimted average treatment effect on the treated
	"""

	ITT = Y_t - np.dot(W, Y_c)
	ATT = ITT.mean()

	if return_individual:
		return ATT, ITT
	else:
		return ATT

def SyntheticEstimates(Y, D, X, higher_moments=False):

	"""
	Function that computes ATT estimates from observerd data using
	synthetic control method.

	If higher_moments is true, then matching with be done on higher
	moments of X as well. For now this only includes absolute first
	moments, and second moments.
	"""

	control = (D == 0)  # Boolean index of control units
	treated = (D == 1)  # Boolean index of treated units
	if higher_moments:
		X_control = np.hstack((X[control], abs(X[control]), X[control]**2))
		X_treated = np.hstack((X[treated], abs(X[treated]), X[treated]**2))
		W_hat = EstimateWeights(X_control, X_treated)
	else:
		W_hat = EstimateWeights(X[control], X[treated])  # synthetic control weights

	return TreatmentEffects(Y[control], Y[treated], W_hat)

def MatchingWithReplacement(Y, D, X, mahalanobis=True):

	X_c, X_t, Y_c, Y_t = X[D==0], X[D==1], Y[D==0], Y[D==1]

	N_t = len(X_t)

	ITT = np.zeros(N_t)

	for i in xrange(N_t):
		x = X_c - X_t[i]
		if mahalanobis:
			V = np.cov(X, rowvar=0)
			match_index = (x.dot(np.linalg.inv(V))*x).sum(axis=1).argmin()
		else:
			match_index = (x**2).sum(axis=1).argmin()
		ITT[i] = Y_t[i] - Y_c[match_index]

	return ITT.mean()  # return ITT.std() as well for standard errors

def MatchingWithoutReplacement(Y, D, X, mahalanobis=True):

	if 2*sum(D) > len(D):  # if N_t > N_c
		print 'Not enough control units, matching with replacement instead.'
		return MatchingWithReplacement(Y, D, X, mahalanobis)

	pscore = sm.Logit(D, X).fit(disp=False).predict()  # estimate pscore with logit
	order = np.argsort(pscore)[::-1]  # sort pscore in descending order, get index
	Y, D, X = Y[order], D[order], X[order]  # any way to do it more elegantly?

	X_c, X_t, Y_c, Y_t = X[D==0], X[D==1], Y[D==0], Y[D==1]

	N_c, N_t = len(X_c), len(X_t)

	ITT = np.zeros(N_t)

	for i in xrange(N_t):
		x = X_c - X_t[i]
		if mahalanobis:
			V = np.cov(X, rowvar=0)
			match_index = (x.dot(np.linalg.inv(V))*x).sum(axis=1).argmin()
		else:
			match_index = (x**2).sum(axis=1).argmin()
		ITT[i] = Y_t[i] - Y_c[match_index]
		X_c = np.delete(X_c, match_index, axis=0)
		Y_c = np.delete(Y_c, match_index, axis=0)

	return ITT.mean()  # return ITT.std() as well for standard errors

def MatchingEstimates(Y, D, X, with_replacement=True):

	"""
	Function that estimates the average treatment effect for the treated (ATT)
	using matching.

	Can specify whether matching is done with or without replacement. Swtiches
	to matching with replacement if number of control units is less than
	number of treated units.

	Args:
		Y = N-dimensional array of observed outcomes
		D = N-dimensional array of treatment indicator; 1=treated, 0=control
		X = N-by-k matrix of covariates

	Returns:
		ATT estimate; matching done with or without replacement as specified
	"""

	if with_replacement:
		return MatchingWithReplacement(Y, D, X)
	else:
		return MatchingWithoutReplacement(Y, D, X)

def OLSEstimates(Y, D, X):

	"""
	Function that estimates average treatment effect for the treated (ATT)
	using OLS, which is consistent when the true model is linear.

	Args:
		Y = N-dimensional array of observed outcomes
		D = N-dimensional array of treatment indicator; 1=treated, 0=control
		X = N-by-k matrix of covariates
	"""

	k = X.shape[1]
	D = D.reshape((D.size, 1))  # convert D into N-by-1 vector
	dX = X - X.mean(0)  # demean covariates
	DdX = D * dX
	# construct design matrix; no constant term as it will be added by sklearn
	W = np.column_stack((D, dX, DdX))

	# use sklearn to run regression
	reg = linear_model.LinearRegression()
	reg.fit(W, Y)

	# for derivation of this estimator, see my notes on Treatment Effects
	return reg.coef_[0] + np.dot((DdX.sum(0) / D.sum()), reg.coef_[-k:])

class parameters:

	"""
	Class object that stores the parameter values for use in the baseline model.

	See SimulateData function for model description.

	Args:
		N = Sample size (control + treated units) to generate
		k = Number of covariates
	"""

	def __init__(self, N=500, k=3):  # set initial parameter values
		self.N = N  # sample size (control + treated units)
		self.k = k  # number of covariates

		self.delta = 3
		self.beta = np.ones(k)
		self.theta = np.ones(k)
		self.mu = np.zeros(k)
		self.Sigma = np.identity(k)
		self.Gamma = np.identity(2)

def SimulateData(para=parameters(), nonlinear=False, return_counterfactual=False):

	"""
	Function that generates data according to one of two simple models that
	satisfies the Unconfoundedness assumption.

	The covariates and error terms are generated according to
		X ~ N(mu, Sigma), epsilon ~ N(0, Gamma).
	The counterfactual outcomes are generated by
		Y_0 = X*beta + epsilon_0,
		Y_1 = delta + X*(beta+theta) + epsilon_1,
	if the nonlinear Boolean is False, or by
		Y_0 = sum_of_abs(X) + epsilon_0,
		Y_1 = sum_of_squares(X) + epsilon_1,
	if the nonlinear Boolean is True.

	Selection is done according to the following propensity score function:
		P(D=1|X) = Phi(X*beta),
	where Phi is the standard normal CDF.

	Args:
		para = Model parameter values supplied by parameter class object
		nonlinear = Boolean indicating whether the data generating model should
		            be linear or not
		return_counterfactual = Boolean indicating whether to return vectors of
		                        counterfactual outcomes

	Returns:
		Y = N-dimensional array of observed outcomes
		D = N-dimensional array of treatment indicator; 1=treated, 0=control
		X = N-by-k matrix of covariates
		Y_0 = N-dimensional array of non-treated outcomes
		Y_1 = N-dimensional array of treated outcomes
	"""

	k = len(para.mu)

	X = np.random.multivariate_normal(mean=para.mu, cov=para.Sigma,
	                                  size=para.N)
	epsilon = np.random.multivariate_normal(mean=np.zeros(2), cov=para.Gamma,
	                                        size=para.N)

	Xbeta = np.dot(X, para.beta)

	pscore = norm.cdf(Xbeta)
	# for each p in pscore, generate Bernoulli rv with success probability p
	D = np.array([np.random.binomial(1, p, size=1) for p in pscore]).flatten()

	if nonlinear:
		Y_0 = abs(X).sum(1) + epsilon[:, 0]
		Y_1 = (X**2).sum(1) + epsilon[:, 1]
	else:
		Y_0 = Xbeta + epsilon[:, 0]
		Y_1 = para.delta + np.dot(X, para.beta+para.theta) + epsilon[:, 1]

	Y = (1 - D) * Y_0 + D * Y_1  # compute observed outcome

	if return_counterfactual:
		return Y, D, X, Y_0, Y_1
	else:
		return Y, D, X

def MonteCarlo(B=500, para=parameters(), nonlinear=False, print_progress=True):

	"""
	Function that returns ATT estimates using synthetic control, matching,
	and OLS computed over B repetitions on simulated data.

	See SimulateData function for data generating process.

	Args:
		B = Number of Monte Carlo simulations to perform
		para = Model parameter values supplied by parameter class object
		nonlinear = Boolean indicating whether the data generating model should
		            be linear or not
		print_progress = Boolean indicating whether to dispaly current number
		                 of completed repetitions (every 10 repetitions)

	Returns:
		ATT_true = Actual average treatment effect for the treated
		ATT_syn = Estimated ATT using synthetic controls
		ATT_match = Estimated ATT using matching
		ATT_ols = Estimated ATT using OLS
	"""

	ATT_true = np.zeros(B)
	ATT_syn = np.zeros(B)
	ATT_match = np.zeros(B)
	ATT_ols = np.zeros(B)

	for i in xrange(B):

		Y, D, X, Y_0, Y_1 = SimulateData(para, nonlinear, True)

		ATT_true[i] = (Y_1[D==1]-Y_0[D==1]).mean()
		ATT_syn[i] = SyntheticEstimates(Y, D, X, higher_moments=nonlinear)
		ATT_match[i] = MatchingEstimates(Y, D, X, with_replacement=False)
		ATT_ols[i] = OLSEstimates(Y, D, X)

		if print_progress and (i+1) % 10 == 0:
			print i+1, 'out of', B, 'simulations completed.'

	return ATT_true, ATT_syn, ATT_match, ATT_ols

def CalculateMSE(B=500, para=parameters(), nonlinear=False):

	"""
	Function that calcuates MSE after performing Monte Carlo simulations.

	Args:
		B = Number of Monte Carlo simulations to perform
		para = Model parameter values supplied by parameter class object
		nonlinear = Boolean indicating whether the data generating model should
		            be linear or not

	Returns:
		MSE_syn = Estimated MSE using synthetic control estimates
		MSE_match = Estimated MSE using matching estimates
		MSE_ols = Estimated MSE using OLS estimates
	"""

	ATT_true, ATT_syn, ATT_match, ATT_ols = MonteCarlo(B, nonlinear=nonlinear)

	MSE_syn = ((ATT_syn - ATT_true)**2).mean()
	MSE_match = ((ATT_match - ATT_true)**2).mean()
	MSE_ols = ((ATT_ols - ATT_true)**2).mean()

	return MSE_syn, MSE_match, MSE_ols

def UseLalondeData():

	"""
	Utility function that reads and applies synthetic control
	estimator on Lalonde's NSW experimental data.

	Note that because of complete randomization, the mean
	difference identifies both the ATE and ATT.
	"""

	url = 'http://www.stanford.edu/~lwong1/data.csv'
	lalonde = pd.read_csv(url)  # read CSV data from url

	covariate_list = ['age', 'education', 'black', 'hispanic', 'married',
	                  'nodegree', 're74', 're75', 'u74', 'u75']

	# don't know how to not convert to array first
	Y = np.array(lalonde['re78'])
	D = np.array(lalonde['treat'])
	X = np.array(lalonde[covariate_list])

	ATT_syn = SyntheticEstimates(Y, D, X)
	ATT_match = MatchingEstimates(Y, D, X)
	ATT_ols = OLSEstimates(Y, D, X)

	print "Using Lalonde's National Supported Work (NSW) experimental data..."
	print 'Estimated ATT using synthetic control estimator:', ATT_syn
	print 'Estimated ATT using matching (with replacement) estimator:', ATT_match
	print 'Estimated ATT using OLS estimator:', ATT_ols
	print 'Mean difference:', Y[D==1].mean() - Y[D==0].mean()

def main():

	#UseLalondeData()

	B = 500
	print '\n' + 'Performing Monte Carlo simulations with B =', B
	MSE_syn, MSE_match, MSE_ols = CalculateMSE(B, nonlinear=True)
	print 'Estimated Mean Squared Error for matching estimator:', MSE_match
	print 'Estimated Mean Squared Error for synthetic control estimator:', MSE_syn
	print 'Estimated Mean Squared Error for OLS estimator:', MSE_ols

if __name__ == '__main__':
	main()
<|endoftext|>


<commit_before>'''
Validates a user entered pattern against STIXPattern grammar.
'''

from __future__ import print_function
import argparse

from antlr4 import InputStream
from antlr4 import CommonTokenStream
from antlr4.error.ErrorListener import ErrorListener
import six

from .grammars.STIXPatternLexer import STIXPatternLexer
from .grammars.STIXPatternParser import STIXPatternParser
from .grammars.STIXPatternListener import STIXPatternListener

class STIXPatternErrorListener(ErrorListener):
    '''
    Modifies ErrorListener to collect error message and set flag to False when
    invalid pattern is encountered.
    '''
    def __init__(self):
        super(STIXPatternErrorListener, self).__init__()
        self.err_strings = []

    def syntaxError(self, recognizer, offendingSymbol, line, column, msg, e):
        self.err_strings.append("FAIL: Error found at line %d:%d. %s" %
                                (line, column, msg))

def run_validator(pattern):
    '''
    Validates a pattern against the STIX Pattern grammar.  Error messages are
    returned in a list.  The test passed if the returned list is empty.
    '''

    start = ''
    if isinstance(pattern, six.string_types):
        start = pattern[:2]
        pattern = InputStream(pattern)

    if not start:
        start = pattern.readline()[:2]
        pattern.seek(0)

    parseErrListener = STIXPatternErrorListener()

    lexer = STIXPatternLexer(pattern)
    # it always adds a console listener by default... remove it.
    lexer.removeErrorListeners()

    stream = CommonTokenStream(lexer)

    parser = STIXPatternParser(stream)
    parser.buildParseTrees = False
    # it always adds a console listener by default... remove it.
    parser.removeErrorListeners()
    parser.addErrorListener(parseErrListener)

    # To improve error messages, replace "<INVALID>" in the literal
    # names with symbolic names.  This is a hack, but seemed like
    # the simplest workaround.
    for i, lit_name in enumerate(parser.literalNames):
        if lit_name == u"<INVALID>":
            parser.literalNames[i] = parser.symbolicNames[i]

    parser.pattern()

    # replace with easier-to-understand error message
    if not (start[0] == '[' or start == '(['):
        parseErrListener.err_strings[0] = "FAIL: Error found at line 1:0. " \
                                          "input is missing square brackets"

    return parseErrListener.err_strings

def validate(user_input, ret_errs=False, print_errs=False):
    '''
    Wrapper for run_validator function that returns True if the user_input
    contains a valid STIX pattern or False otherwise. The error messages may
    also be returned or printed based upon the ret_errs and print_errs arg
    values.
    '''

    errs = run_validator(user_input)
    passed = len(errs) == 0

    if print_errs:
        for err in errs:
            print(err)

    if ret_errs:
        return passed, errs

    return passed

def main():
    '''
    Continues to validate patterns until it encounters EOF within a pattern
    file or Ctrl-C is pressed by the user.
    '''
    parser = argparse.ArgumentParser(description='Validate STIX Patterns.')
    parser.add_argument('-f', '--file',
                        help="Specify this arg to read patterns from a file.",
                        type=argparse.FileType("r"))
    args = parser.parse_args()

    pass_count = fail_count = 0

    # I tried using a generator (where each iteration would run raw_input()),
    # but raw_input()'s behavior seems to change when called from within a
    # generator: I only get one line, then the generator completes!  I don't
    # know why behavior changes...
    import functools
    if args.file:
        nextpattern = args.file.readline
    else:
        nextpattern = functools.partial(raw_input, "Enter a pattern to validate: ")

    try:
        while True:
            pattern = nextpattern()
            if not pattern:
                break
            tests_passed, err_strings = validate(pattern, True)

            if tests_passed:
                print("\nPASS: %s" % pattern)
                pass_count += 1

            else:
                for err in err_strings:
                    print(err, '\n')
                fail_count += 1
    except (EOFError, KeyboardInterrupt):
        pass
    finally:
        if args.file:
            args.file.close()

    print("\nPASSED:", pass_count, " patterns")
    print("FAILED:", fail_count, " patterns")

if __name__ == '__main__':
    main()
<commit_msg>Make interactive mode work on Python 3<commit_after>'''
Validates a user entered pattern against STIXPattern grammar.
'''

from __future__ import print_function
import argparse

from antlr4 import InputStream
from antlr4 import CommonTokenStream
from antlr4.error.ErrorListener import ErrorListener
import six

from .grammars.STIXPatternLexer import STIXPatternLexer
from .grammars.STIXPatternParser import STIXPatternParser
from .grammars.STIXPatternListener import STIXPatternListener

class STIXPatternErrorListener(ErrorListener):
    '''
    Modifies ErrorListener to collect error message and set flag to False when
    invalid pattern is encountered.
    '''
    def __init__(self):
        super(STIXPatternErrorListener, self).__init__()
        self.err_strings = []

    def syntaxError(self, recognizer, offendingSymbol, line, column, msg, e):
        self.err_strings.append("FAIL: Error found at line %d:%d. %s" %
                                (line, column, msg))

def run_validator(pattern):
    '''
    Validates a pattern against the STIX Pattern grammar.  Error messages are
    returned in a list.  The test passed if the returned list is empty.
    '''

    start = ''
    if isinstance(pattern, six.string_types):
        start = pattern[:2]
        pattern = InputStream(pattern)

    if not start:
        start = pattern.readline()[:2]
        pattern.seek(0)

    parseErrListener = STIXPatternErrorListener()

    lexer = STIXPatternLexer(pattern)
    # it always adds a console listener by default... remove it.
    lexer.removeErrorListeners()

    stream = CommonTokenStream(lexer)

    parser = STIXPatternParser(stream)
    parser.buildParseTrees = False
    # it always adds a console listener by default... remove it.
    parser.removeErrorListeners()
    parser.addErrorListener(parseErrListener)

    # To improve error messages, replace "<INVALID>" in the literal
    # names with symbolic names.  This is a hack, but seemed like
    # the simplest workaround.
    for i, lit_name in enumerate(parser.literalNames):
        if lit_name == u"<INVALID>":
            parser.literalNames[i] = parser.symbolicNames[i]

    parser.pattern()

    # replace with easier-to-understand error message
    if not (start[0] == '[' or start == '(['):
        parseErrListener.err_strings[0] = "FAIL: Error found at line 1:0. " \
                                          "input is missing square brackets"

    return parseErrListener.err_strings

def validate(user_input, ret_errs=False, print_errs=False):
    '''
    Wrapper for run_validator function that returns True if the user_input
    contains a valid STIX pattern or False otherwise. The error messages may
    also be returned or printed based upon the ret_errs and print_errs arg
    values.
    '''

    errs = run_validator(user_input)
    passed = len(errs) == 0

    if print_errs:
        for err in errs:
            print(err)

    if ret_errs:
        return passed, errs

    return passed

def main():
    '''
    Continues to validate patterns until it encounters EOF within a pattern
    file or Ctrl-C is pressed by the user.
    '''
    parser = argparse.ArgumentParser(description='Validate STIX Patterns.')
    parser.add_argument('-f', '--file',
                        help="Specify this arg to read patterns from a file.",
                        type=argparse.FileType("r"))
    args = parser.parse_args()

    pass_count = fail_count = 0

    # I tried using a generator (where each iteration would run raw_input()),
    # but raw_input()'s behavior seems to change when called from within a
    # generator: I only get one line, then the generator completes!  I don't
    # know why behavior changes...
    import functools
    if args.file:
        nextpattern = args.file.readline
    else:
        nextpattern = functools.partial(six.moves.input, "Enter a pattern to validate: ")

    try:
        while True:
            pattern = nextpattern()
            if not pattern:
                break
            tests_passed, err_strings = validate(pattern, True)

            if tests_passed:
                print("\nPASS: %s" % pattern)
                pass_count += 1

            else:
                for err in err_strings:
                    print(err, '\n')
                fail_count += 1
    except (EOFError, KeyboardInterrupt):
        pass
    finally:
        if args.file:
            args.file.close()

    print("\nPASSED:", pass_count, " patterns")
    print("FAILED:", fail_count, " patterns")

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>"""Test I/O."""
import logging

from tensorpac.io import (set_log_level, is_pandas_installed,
                          is_statsmodels_installed, progress_bar)

logger = logging.getLogger('tensorpac')

levels = ['profiler', 'debug', 'info', 'warning', 'error', 'critical']

class TestIO(object):
    """Test statistical functions."""

    def test_log_level(self):
        """Test setting the log level."""
        for l in levels:
            set_log_level(l)
        set_log_level(True)
        set_log_level(False)

    def test_logger(self):
        """Test logger levels."""
        logger.profiler("profiler")
        logger.debug("debug")
        logger.info("info")
        logger.warning("warning")
        logger.critical("critical")

    def test_progress_bar(self):
        """Test progress bar."""
        progress_bar(5, 10)

    def test_dependance(self):
        """Test dependancies."""
        is_statsmodels_installed()
        is_pandas_installed()
<commit_msg>test match<commit_after>"""Test I/O."""
import logging

from tensorpac.io import (set_log_level, is_pandas_installed,
                          is_statsmodels_installed, progress_bar)

logger = logging.getLogger('tensorpac')

levels = ['profiler', 'debug', 'info', 'warning', 'error', 'critical']

class TestIO(object):
    """Test statistical functions."""

    def test_log_level(self):
        """Test setting the log level."""
        for l in levels:
            set_log_level(l)
        set_log_level(False)
        set_log_level(True)
        set_log_level(match="ok")
        logger.info("show me ok")
        logger.info("show me")

    def test_logger(self):
        """Test logger levels."""
        set_log_level("profiler")
        logger.profiler("profiler")
        logger.debug("debug")
        logger.info("info")
        logger.warning("warning")
        logger.critical("critical")

    def test_progress_bar(self):
        """Test progress bar."""
        progress_bar(5, 10)

    def test_dependance(self):
        """Test dependancies."""
        is_statsmodels_installed()
        is_pandas_installed()
<|endoftext|>


<commit_before>from struct import pack, unpack, calcsize
from collections import OrderedDict

def getbytes(s, n):
    b = s.read(n)
    assert len(b) == n, "Unexpected EOF"
    return b

def getbyte(s):
    return getbytes(s, 1)

class Seek(object):
    def __init__(self, s, *args, **kwargs):
        self.old_pos = None
        self.s = s
        self.args = args
        self.kwargs = kwargs

    def __enter__(self):
        self.old_pos = self.s.tell()
        self.s.seek(*self.args, **self.kwargs)

    def __exit__(self, exc_type, exc_value, traceback):
        self.s.seek(self.old_pos)

class FakeWriteStream(object):
    def __init__(self, offset):
        self.offset = offset

    def seek(self, offset):
        self.offset = offset

    def tell(self):
        return self.offset

    def write(self, data):
        self.offset += len(data)

class BaseField(object):
    def unpack(self, s):
        self.data = self.unpack_data(s)

    def unpack_data(self, s):
        raise notImplementedError

    def pack(self, s):
        self.pack_data(s, self.data)

    def pack_data(self, s, data):
        raise NotImplementedError

    def full_pack(self, s):
        new_data = self.data
        while True:
            old_data = new_data
            self.pack(FakeWriteStream(s.tell()))
            new_data = self.data
            if old_data == new_data:
                break
        self.pack(s)

    def serialize(self):
        return self.data

class Struct(BaseField):
    def add_field(self, name, f):
        assert name not in self.field, name
        self.field[name] = f
        input_type, v = self.input
        if input_type == "data":
            f.data = v.get(name, None)
        elif input_type == "stream":
            f.unpack(v)
        else:
            assert False, input_type
        return f

    def F(self, name, f):
        return self.add_field(name, f)

    def unpack(self, s):
        self.field = OrderedDict()
        self.input = ("stream", s)
        self.fields()
        del self.input

    def pack(self, s):
        for name, f in self.field.iteritems():
            f.pack(s)

    @property
    def data(self):
        data = OrderedDict()
        for k, v in self.field.iteritems():
            data[k] = v.data
        return data

    @data.setter
    def data(self, v):
        self.field = OrderedDict()
        self.input = ("data", v)
        self.fields()
        del self.input

    def serialize(self):
        data = OrderedDict()
        for k, v in self.field.iteritems():
            data[k] = v.serialize()
        return data

    def fields(self):
        raise NotImplementedError

class Magic(BaseField):
    def __init__(self, magic):
        self.magic = magic

    def unpack(self, s):
        data = getbytes(s, len(self.magic))
        assert data == self.magic

    def pack(self, s):
        s.write(self.magic)

    @property
    def data(self):
        return self.magic

    @data.setter
    def data(self, v):
        assert v == self.magic or v is None, v

class Format(BaseField):
    def __init__(self, fmt):
        if fmt[0] in "@=<>!":
            bosa = fmt[0]
            fmt = fmt[1:]
        else:
            bosa = "<"
        self.bosa = bosa
        self.fmt = fmt
        self.single = len(fmt) == 1

    def unpack_data(self, s):
        fmt = self.bosa + self.fmt
        size = calcsize(fmt)
        b = getbytes(s, size)
        data = unpack(fmt, b)
        if self.single:
            assert len(data) == 1
            data = data[0]
        return data

    def pack_data(self, s, data):
        if self.single:
            data = (data,)
        s.write(pack(self.fmt, *data))

class BaseArray(BaseField):
    def __init__(self, field_maker=None, field_function=None):
        if field_function is None:
            field_function = lambda i, f: field_maker()
        self.field_fun = field_function

    def unpack(self, s):
        self.field = [self.field_fun(i, self) for i in xrange(self.size)]
        for f in self.field:
            f.unpack(s)

    def pack(self, s):
        for f in self.field:
            f.pack(s)

    @property
    def data(self):
        return [f.data for f in self.field]

    @data.setter
    def data(self, v):
        self.field = [self.field_fun(i, self) for i in xrange(len(v))]
        for f, fv in zip(self.field, v):
            f.data = fv

    def serialize(self):
        return [f.serialize() for f in self.field]

class Array(BaseArray):
    def __init__(self, size, *args, **kwargs):
        self.size = size
        BaseArray.__init__(self, *args, **kwargs)

class PrefixedArray(BaseArray):
    def __init__(self, prefix_field, *args, **kwargs):
        self.prefix_field = prefix_field
        BaseArray.__init__(self, *args, **kwargs)

    @property
    def size(self):
        return self.prefix_field.data

    def unpack(self, s):
        self.prefix_field.unpack(s)
        BaseArray.unpack(self, s)

    def pack(self, s):
        self.prefix_field.data = len(self.field)
        self.prefix_field.pack(s)
        BaseArray.pack(self, s)

class BaseBlob(BaseField):
    def unpack_data(self, s):
        return getbytes(s, self.size)

    def pack_data(self, s, data):
        s.write(data)

class Blob(BaseBlob):
    def __init__(self, size):
        self.size = size

class PrefixedBlob(BaseBlob):
    def __init__(self, prefix_field, *args, **kwargs):
        self.prefix_field = prefix_field
        BaseBlob.__init__(self, *args, **kwargs)

    @property
    def size(self):
        return self.prefix_field.data

    def unpack(self, s):
        self.prefix_field.unpack(s)
        BaseBlob.unpack(self, s)

    def pack(self, s):
        self.prefix_field.data = len(self.field)
        self.prefix_field.pack(s)
        BaseBlob.pack(self, s)

class String(BaseField):
    def unpack_data(self, s):
        lc = []
        c = getbyte(s)
        while c != "\0":
            lc.append(c)
            c = getbyte(s)
        return "".join(lc)

    def pack_data(self, s, data):
        s.write(data)
        s.write('\0')

class FixedString(BaseField):
    def __init__(self, size):
        self.size = size

    def unpack_data(self, s):
        data = getbytes(s, self.size)
        data = data.rstrip("\0")
        return data

    def pack_data(self, s, data):
        data = data.ljust(self.size, "\0")
        s.write(data)

class Index(BaseField):
    def __init__(self, array, index_field):
        self.array = array
        self.index_field = index_field

    def unpack_data(self, s):
        self.index_field.unpack(s)
        return self.array.field[self.index_field.data].data

    def pack_data(self, s, data):
        if data not in self.array.data:
            self.array.data = self.array.data + [data]
        self.index_field.data = self.array.data.index(data)
        self.index_field.pack(s)

class Offset(BaseField):
    def unpack_data(self, s):
        return s.tell()

    def pack_data(self, s, data):
        self.data = s.tell()

class Pointer(BaseField):
    def __init__(self, offset, field):
        self.offset = offset
        self.field = field

    def unpack(self, s):
        with Seek(s, self.offset):
            self.field.unpack(s)

    @property
    def data(self):
        return self.field.data

    @data.setter
    def data(self, v):
        self.field.data = v
<commit_msg>name in full_pack<commit_after>from struct import pack, unpack, calcsize
from collections import OrderedDict

def getbytes(s, n):
    b = s.read(n)
    assert len(b) == n, "Unexpected EOF"
    return b

def getbyte(s):
    return getbytes(s, 1)

class Seek(object):
    def __init__(self, s, *args, **kwargs):
        self.old_pos = None
        self.s = s
        self.args = args
        self.kwargs = kwargs

    def __enter__(self):
        self.old_pos = self.s.tell()
        self.s.seek(*self.args, **self.kwargs)

    def __exit__(self, exc_type, exc_value, traceback):
        self.s.seek(self.old_pos)

class FakeWriteStream(object):
    def __init__(self, offset, name):
        self.offset = offset
        self.name = name

    def seek(self, offset):
        self.offset = offset

    def tell(self):
        return self.offset

    def write(self, data):
        self.offset += len(data)

class BaseField(object):
    def unpack(self, s):
        self.data = self.unpack_data(s)

    def unpack_data(self, s):
        raise notImplementedError

    def pack(self, s):
        self.pack_data(s, self.data)

    def pack_data(self, s, data):
        raise NotImplementedError

    def full_pack(self, s):
        new_data = self.data
        while True:
            old_data = new_data
            self.pack(FakeWriteStream(s.tell(), s.name))
            new_data = self.data
            if old_data == new_data:
                break
        self.pack(s)

    def serialize(self):
        return self.data

class Struct(BaseField):
    def add_field(self, name, f):
        assert name not in self.field, name
        self.field[name] = f
        input_type, v = self.input
        if input_type == "data":
            f.data = v.get(name, None)
        elif input_type == "stream":
            f.unpack(v)
        else:
            assert False, input_type
        return f

    def F(self, name, f):
        return self.add_field(name, f)

    def unpack(self, s):
        self.field = OrderedDict()
        self.input = ("stream", s)
        self.fields()
        del self.input

    def pack(self, s):
        for name, f in self.field.iteritems():
            f.pack(s)

    @property
    def data(self):
        data = OrderedDict()
        for k, v in self.field.iteritems():
            data[k] = v.data
        return data

    @data.setter
    def data(self, v):
        self.field = OrderedDict()
        self.input = ("data", v)
        self.fields()
        del self.input

    def serialize(self):
        data = OrderedDict()
        for k, v in self.field.iteritems():
            data[k] = v.serialize()
        return data

    def fields(self):
        raise NotImplementedError

class Magic(BaseField):
    def __init__(self, magic):
        self.magic = magic

    def unpack(self, s):
        data = getbytes(s, len(self.magic))
        assert data == self.magic

    def pack(self, s):
        s.write(self.magic)

    @property
    def data(self):
        return self.magic

    @data.setter
    def data(self, v):
        assert v == self.magic or v is None, v

class Format(BaseField):
    def __init__(self, fmt):
        if fmt[0] in "@=<>!":
            bosa = fmt[0]
            fmt = fmt[1:]
        else:
            bosa = "<"
        self.bosa = bosa
        self.fmt = fmt
        self.single = len(fmt) == 1

    def unpack_data(self, s):
        fmt = self.bosa + self.fmt
        size = calcsize(fmt)
        b = getbytes(s, size)
        data = unpack(fmt, b)
        if self.single:
            assert len(data) == 1
            data = data[0]
        return data

    def pack_data(self, s, data):
        if self.single:
            data = (data,)
        s.write(pack(self.fmt, *data))

class BaseArray(BaseField):
    def __init__(self, field_maker=None, field_function=None):
        if field_function is None:
            field_function = lambda i, f: field_maker()
        self.field_fun = field_function

    def unpack(self, s):
        self.field = [self.field_fun(i, self) for i in xrange(self.size)]
        for f in self.field:
            f.unpack(s)

    def pack(self, s):
        for f in self.field:
            f.pack(s)

    @property
    def data(self):
        return [f.data for f in self.field]

    @data.setter
    def data(self, v):
        self.field = [self.field_fun(i, self) for i in xrange(len(v))]
        for f, fv in zip(self.field, v):
            f.data = fv

    def serialize(self):
        return [f.serialize() for f in self.field]

class Array(BaseArray):
    def __init__(self, size, *args, **kwargs):
        self.size = size
        BaseArray.__init__(self, *args, **kwargs)

class PrefixedArray(BaseArray):
    def __init__(self, prefix_field, *args, **kwargs):
        self.prefix_field = prefix_field
        BaseArray.__init__(self, *args, **kwargs)

    @property
    def size(self):
        return self.prefix_field.data

    def unpack(self, s):
        self.prefix_field.unpack(s)
        BaseArray.unpack(self, s)

    def pack(self, s):
        self.prefix_field.data = len(self.field)
        self.prefix_field.pack(s)
        BaseArray.pack(self, s)

class BaseBlob(BaseField):
    def unpack_data(self, s):
        return getbytes(s, self.size)

    def pack_data(self, s, data):
        s.write(data)

class Blob(BaseBlob):
    def __init__(self, size):
        self.size = size

class PrefixedBlob(BaseBlob):
    def __init__(self, prefix_field, *args, **kwargs):
        self.prefix_field = prefix_field
        BaseBlob.__init__(self, *args, **kwargs)

    @property
    def size(self):
        return self.prefix_field.data

    def unpack(self, s):
        self.prefix_field.unpack(s)
        BaseBlob.unpack(self, s)

    def pack(self, s):
        self.prefix_field.data = len(self.field)
        self.prefix_field.pack(s)
        BaseBlob.pack(self, s)

class String(BaseField):
    def unpack_data(self, s):
        lc = []
        c = getbyte(s)
        while c != "\0":
            lc.append(c)
            c = getbyte(s)
        return "".join(lc)

    def pack_data(self, s, data):
        s.write(data)
        s.write('\0')

class FixedString(BaseField):
    def __init__(self, size):
        self.size = size

    def unpack_data(self, s):
        data = getbytes(s, self.size)
        data = data.rstrip("\0")
        return data

    def pack_data(self, s, data):
        data = data.ljust(self.size, "\0")
        s.write(data)

class Index(BaseField):
    def __init__(self, array, index_field):
        self.array = array
        self.index_field = index_field

    def unpack_data(self, s):
        self.index_field.unpack(s)
        return self.array.field[self.index_field.data].data

    def pack_data(self, s, data):
        if data not in self.array.data:
            self.array.data = self.array.data + [data]
        self.index_field.data = self.array.data.index(data)
        self.index_field.pack(s)

class Offset(BaseField):
    def unpack_data(self, s):
        return s.tell()

    def pack_data(self, s, data):
        self.data = s.tell()

class Pointer(BaseField):
    def __init__(self, offset, field):
        self.offset = offset
        self.field = field

    def unpack(self, s):
        with Seek(s, self.offset):
            self.field.unpack(s)

    @property
    def data(self):
        return self.field.data

    @data.setter
    def data(self, v):
        self.field.data = v
<|endoftext|>


<commit_before>import csv
from contextlib import contextmanager
import datetime
import os
import shutil
import tempfile
from zipfile import ZipFile, ZIP_DEFLATED

from ichnaea.backup.s3 import S3Backend, compute_hash
from ichnaea.models import (
    CellMeasure,
    MEASURE_TYPE,
    MeasureBlock,
    WifiMeasure,
)
from ichnaea.tasks import DatabaseTask
from ichnaea.worker import celery

@contextmanager
def selfdestruct_tempdir(s3_key):
    """
    We need two temp directories to do this properly.

    The base_path is a temp directory that holds all the content that
    will go into our zip file. This is effectively a working
    directory that gets immediately deleted once the zip file is
    ready.

    The zip_path is the filename of the zip file that will get
    uploaded into S3. It is the responsibility of the caller of
    selfdestruct_tempdir to remove the zip_path and parent directory.
    """
    short_name = os.path.split(s3_key)[-1]

    base_path = tempfile.mkdtemp()
    s3_path = os.path.join(tempfile.mkdtemp(), short_name)
    try:
        zip_path = os.path.join(base_path, s3_path)
        yield base_path, zip_path
    finally:
        try:
            z = ZipFile(zip_path, "w", ZIP_DEFLATED)
            try:
                for root, dirs, files in os.walk(base_path):
                    for fn in files:
                        absfn = os.path.join(root, fn)
                        zip_fn = absfn[len(base_path) + len(os.sep):]
                        z.write(absfn, zip_fn)
            finally:
                z.close()
        finally:
            shutil.rmtree(base_path)

@celery.task(base=DatabaseTask, bind=True)
def write_cellmeasure_s3_backups(self, cleanup_zip=True):
    measure_type = MEASURE_TYPE['cell']
    zip_prefix = 'CellMeasure'
    csv_name = 'cell_measure.csv'
    measure_cls = CellMeasure
    return write_measure_s3_backups(self,
                                    measure_type,
                                    zip_prefix,
                                    csv_name,
                                    measure_cls,
                                    cleanup_zip)

@celery.task(base=DatabaseTask, bind=True)
def write_wifimeasure_s3_backups(self, cleanup_zip=True):
    measure_type = MEASURE_TYPE['wifi']
    zip_prefix = 'WifiMeasure'
    csv_name = 'wifi_measure.csv'
    measure_cls = WifiMeasure
    return write_measure_s3_backups(self,
                                    measure_type,
                                    zip_prefix,
                                    csv_name,
                                    measure_cls,
                                    cleanup_zip)

def write_measure_s3_backups(self, measure_type,
                             zip_prefix, csv_name,
                             measure_cls, cleanup_zip):
    """
    Iterate over each of the measure block records that aren't
    backed up yet and back them up.

    Assume that this is running in a single task.
    """

    zips = []
    utcnow = datetime.datetime.utcnow()
    s3_backend = S3Backend(
        self.app.s3_settings['backup_bucket'],
        self.app.s3_settings['backup_prefix'],
        self.heka_client)

    with self.db_session() as session:
        query = session.query(MeasureBlock).filter(
            MeasureBlock.measure_type == measure_type).filter(
            MeasureBlock.archive_date.is_(None)).order_by(
            MeasureBlock.end_id)
        for cmb in query.all():
            cmb.s3_key = '%s/%s_%d_%d.zip' % (
                utcnow.strftime("%Y%m"),
                zip_prefix,
                cmb.start_id,
                cmb.end_id)

            with selfdestruct_tempdir(cmb.s3_key) as (tmp_path, zip_path):
                rset = session.execute("select * from alembic_version")
                rev = rset.first()[0]
                with open(os.path.join(tmp_path,
                                       'alembic_revision.txt'), 'w') as f:
                    f.write('%s\n' % rev)

                cm_fname = os.path.join(tmp_path, csv_name)

                cm_query = session.query(measure_cls).filter(
                    measure_cls.id >= cmb.start_id).filter(
                    measure_cls.id <= cmb.end_id)

                col_names = None
                with open(cm_fname, 'w') as f:
                    csv_out = csv.writer(f, dialect='excel')
                    for i, row in enumerate(cm_query.all()):
                        if i == 0:
                            col_names = [c.name for c in row.__table__.columns]
                            csv_out.writerow(col_names)
                            pass
                        data_row = [getattr(row, cname) for cname in col_names]
                        csv_out.writerow(data_row)

            cmb.archive_sha = compute_hash(zip_path)

            try:
                if not s3_backend.backup_archive(cmb.s3_key, zip_path):
                    continue
                self.heka_client.incr('s3.backup.%s' % measure_type,
                                      (cmb.end_id - cmb.start_id))
            finally:
                if cleanup_zip:
                    if os.path.exists(zip_path):
                        zip_dir, zip_file = os.path.split(zip_path)
                        if os.path.exists(zip_dir):
                            shutil.rmtree(zip_dir)
                else:
                    zips.append(zip_path)

            session.add(cmb)
            session.commit()
    return zips

def schedule_measure_archival(self, measure_type, measure_cls, batch=100):
    blocks = []
    with self.db_session() as session:
        query = session.query(MeasureBlock.end_id).filter(
            MeasureBlock.measure_type == measure_type).order_by(
            MeasureBlock.end_id.desc())
        record = query.first()
        if record:
            min_id = record[0]
        else:
            query = session.query(measure_cls.id).order_by(
                measure_cls.id.asc())
            record = query.first()
            min_id = record[0]

        query = session.query(measure_cls.id).order_by(
            measure_cls.id.desc())
        record = query.first()

        # We're using half-open ranges, so we need to bump the max_id
        max_id = record[0] + 1

        if max_id - min_id < batch - 1:
            # Not enough to fill a block
            return blocks

        this_max_id = min_id + batch

        while (this_max_id - min_id) == batch:
            cm_blk = MeasureBlock(start_id=min_id,
                                  end_id=this_max_id,
                                  measure_type=measure_type)
            blocks.append((cm_blk.start_id, cm_blk.end_id))
            session.add(cm_blk)

            min_id = this_max_id
            this_max_id = min(batch + this_max_id, max_id)
        session.commit()
    return blocks

@celery.task(base=DatabaseTask, bind=True)
def schedule_cellmeasure_archival(self, batch=100):
    return schedule_measure_archival(
        self, MEASURE_TYPE['cell'], CellMeasure, batch)

@celery.task(base=DatabaseTask, bind=True)
def schedule_wifimeasure_archival(self, batch=100):
    return schedule_measure_archival(
        self, MEASURE_TYPE['wifi'], WifiMeasure, batch)

def delete_measure_records(self, measure_type, measure_cls, cleanup_zip):
    s3_backend = S3Backend(
        self.app.s3_settings['backup_bucket'],
        self.app.s3_settings['backup_prefix'],
        self.heka_client)

    with self.db_session() as session:
        query = session.query(MeasureBlock).filter(
            MeasureBlock.measure_type == measure_type).filter(
            MeasureBlock.s3_key.isnot(None)).filter(
            MeasureBlock.archive_date.is_(None))
        for cmb in query.all():
            expected_sha = cmb.archive_sha
            if s3_backend.check_archive(expected_sha, cmb.s3_key):
                q = session.query(measure_cls).filter(
                    measure_cls.id >= cmb.start_id,
                    measure_cls.id <= cmb.end_id)
                q.delete()
                session.commit()

@celery.task(base=DatabaseTask, bind=True)
def delete_cellmeasure_records(self, cleanup_zip=True):
    return delete_measure_records(
        self, MEASURE_TYPE['cell'], CellMeasure, cleanup_zip)

@celery.task(base=DatabaseTask, bind=True)
def delete_wifimeasure_records(self, cleanup_zip=True):
    return delete_measure_records(
        self, MEASURE_TYPE['wifi'], WifiMeasure, cleanup_zip)
<commit_msg>Gracefully handle empty measure tables archival tasks.<commit_after>import csv
from contextlib import contextmanager
import datetime
import os
import shutil
import tempfile
from zipfile import ZipFile, ZIP_DEFLATED

from ichnaea.backup.s3 import S3Backend, compute_hash
from ichnaea.models import (
    CellMeasure,
    MEASURE_TYPE,
    MeasureBlock,
    WifiMeasure,
)
from ichnaea.tasks import DatabaseTask
from ichnaea.worker import celery

@contextmanager
def selfdestruct_tempdir(s3_key):
    """
    We need two temp directories to do this properly.

    The base_path is a temp directory that holds all the content that
    will go into our zip file. This is effectively a working
    directory that gets immediately deleted once the zip file is
    ready.

    The zip_path is the filename of the zip file that will get
    uploaded into S3. It is the responsibility of the caller of
    selfdestruct_tempdir to remove the zip_path and parent directory.
    """
    short_name = os.path.split(s3_key)[-1]

    base_path = tempfile.mkdtemp()
    s3_path = os.path.join(tempfile.mkdtemp(), short_name)
    try:
        zip_path = os.path.join(base_path, s3_path)
        yield base_path, zip_path
    finally:
        try:
            z = ZipFile(zip_path, "w", ZIP_DEFLATED)
            try:
                for root, dirs, files in os.walk(base_path):
                    for fn in files:
                        absfn = os.path.join(root, fn)
                        zip_fn = absfn[len(base_path) + len(os.sep):]
                        z.write(absfn, zip_fn)
            finally:
                z.close()
        finally:
            shutil.rmtree(base_path)

@celery.task(base=DatabaseTask, bind=True)
def write_cellmeasure_s3_backups(self, cleanup_zip=True):
    measure_type = MEASURE_TYPE['cell']
    zip_prefix = 'CellMeasure'
    csv_name = 'cell_measure.csv'
    measure_cls = CellMeasure
    return write_measure_s3_backups(self,
                                    measure_type,
                                    zip_prefix,
                                    csv_name,
                                    measure_cls,
                                    cleanup_zip)

@celery.task(base=DatabaseTask, bind=True)
def write_wifimeasure_s3_backups(self, cleanup_zip=True):
    measure_type = MEASURE_TYPE['wifi']
    zip_prefix = 'WifiMeasure'
    csv_name = 'wifi_measure.csv'
    measure_cls = WifiMeasure
    return write_measure_s3_backups(self,
                                    measure_type,
                                    zip_prefix,
                                    csv_name,
                                    measure_cls,
                                    cleanup_zip)

def write_measure_s3_backups(self, measure_type,
                             zip_prefix, csv_name,
                             measure_cls, cleanup_zip):
    """
    Iterate over each of the measure block records that aren't
    backed up yet and back them up.

    Assume that this is running in a single task.
    """

    zips = []
    utcnow = datetime.datetime.utcnow()
    s3_backend = S3Backend(
        self.app.s3_settings['backup_bucket'],
        self.app.s3_settings['backup_prefix'],
        self.heka_client)

    with self.db_session() as session:
        query = session.query(MeasureBlock).filter(
            MeasureBlock.measure_type == measure_type).filter(
            MeasureBlock.archive_date.is_(None)).order_by(
            MeasureBlock.end_id)
        for cmb in query.all():
            cmb.s3_key = '%s/%s_%d_%d.zip' % (
                utcnow.strftime("%Y%m"),
                zip_prefix,
                cmb.start_id,
                cmb.end_id)

            with selfdestruct_tempdir(cmb.s3_key) as (tmp_path, zip_path):
                rset = session.execute("select * from alembic_version")
                rev = rset.first()[0]
                with open(os.path.join(tmp_path,
                                       'alembic_revision.txt'), 'w') as f:
                    f.write('%s\n' % rev)

                cm_fname = os.path.join(tmp_path, csv_name)

                cm_query = session.query(measure_cls).filter(
                    measure_cls.id >= cmb.start_id).filter(
                    measure_cls.id <= cmb.end_id)

                col_names = None
                with open(cm_fname, 'w') as f:
                    csv_out = csv.writer(f, dialect='excel')
                    for i, row in enumerate(cm_query.all()):
                        if i == 0:
                            col_names = [c.name for c in row.__table__.columns]
                            csv_out.writerow(col_names)
                            pass
                        data_row = [getattr(row, cname) for cname in col_names]
                        csv_out.writerow(data_row)

            cmb.archive_sha = compute_hash(zip_path)

            try:
                if not s3_backend.backup_archive(cmb.s3_key, zip_path):
                    continue
                self.heka_client.incr('s3.backup.%s' % measure_type,
                                      (cmb.end_id - cmb.start_id))
            finally:
                if cleanup_zip:
                    if os.path.exists(zip_path):
                        zip_dir, zip_file = os.path.split(zip_path)
                        if os.path.exists(zip_dir):
                            shutil.rmtree(zip_dir)
                else:
                    zips.append(zip_path)

            session.add(cmb)
            session.commit()
    return zips

def schedule_measure_archival(self, measure_type, measure_cls, batch=100):
    blocks = []
    with self.db_session() as session:
        query = session.query(MeasureBlock.end_id).filter(
            MeasureBlock.measure_type == measure_type).order_by(
            MeasureBlock.end_id.desc())
        record = query.first()
        if record:
            min_id = record[0]
        else:
            query = session.query(measure_cls.id).order_by(
                measure_cls.id.asc())
            record = query.first()
            if record:
                min_id = record[0]
            else:
                min_id = 0

        query = session.query(measure_cls.id).order_by(
            measure_cls.id.desc())
        record = query.first()
        if not record:
            return blocks

        # We're using half-open ranges, so we need to bump the max_id
        max_id = record[0] + 1

        if max_id - min_id < batch - 1:
            # Not enough to fill a block
            return blocks

        this_max_id = min_id + batch

        while (this_max_id - min_id) == batch:
            cm_blk = MeasureBlock(start_id=min_id,
                                  end_id=this_max_id,
                                  measure_type=measure_type)
            blocks.append((cm_blk.start_id, cm_blk.end_id))
            session.add(cm_blk)

            min_id = this_max_id
            this_max_id = min(batch + this_max_id, max_id)
        session.commit()
    return blocks

@celery.task(base=DatabaseTask, bind=True)
def schedule_cellmeasure_archival(self, batch=100):
    return schedule_measure_archival(
        self, MEASURE_TYPE['cell'], CellMeasure, batch)

@celery.task(base=DatabaseTask, bind=True)
def schedule_wifimeasure_archival(self, batch=100):
    return schedule_measure_archival(
        self, MEASURE_TYPE['wifi'], WifiMeasure, batch)

def delete_measure_records(self, measure_type, measure_cls, cleanup_zip):
    s3_backend = S3Backend(
        self.app.s3_settings['backup_bucket'],
        self.app.s3_settings['backup_prefix'],
        self.heka_client)

    with self.db_session() as session:
        query = session.query(MeasureBlock).filter(
            MeasureBlock.measure_type == measure_type).filter(
            MeasureBlock.s3_key.isnot(None)).filter(
            MeasureBlock.archive_date.is_(None))
        for cmb in query.all():
            expected_sha = cmb.archive_sha
            if s3_backend.check_archive(expected_sha, cmb.s3_key):
                q = session.query(measure_cls).filter(
                    measure_cls.id >= cmb.start_id,
                    measure_cls.id <= cmb.end_id)
                q.delete()
                session.commit()

@celery.task(base=DatabaseTask, bind=True)
def delete_cellmeasure_records(self, cleanup_zip=True):
    return delete_measure_records(
        self, MEASURE_TYPE['cell'], CellMeasure, cleanup_zip)

@celery.task(base=DatabaseTask, bind=True)
def delete_wifimeasure_records(self, cleanup_zip=True):
    return delete_measure_records(
        self, MEASURE_TYPE['wifi'], WifiMeasure, cleanup_zip)
<|endoftext|>


<commit_before>"""
Use the DBLP dataset to recommend experts. Find the optimal parameters.
"""
import gc
import os
import numpy
import logging
import sys
import argparse
from exp.influence2.GraphRanker import GraphRanker
from exp.influence2.RankAggregator import RankAggregator
from exp.influence2.ArnetMinerDataset import ArnetMinerDataset
from apgl.util.Latex import Latex
from apgl.util.Evaluator import Evaluator

logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
numpy.set_printoptions(suppress=True, precision=3, linewidth=160)
numpy.random.seed(21)

parser = argparse.ArgumentParser(description='Run reputation evaluation experiments')
parser.add_argument("-r", "--runLDA", action="store_true", help="Run Latent Dirchlet Allocation")
args = parser.parse_args()

averagePrecisionN = 20
similarityCutoff = 0.30
ns = numpy.arange(5, 55, 5)
runLSI = not args.runLDA

dataset = ArnetMinerDataset(runLSI=runLSI)
#dataset.dataFilename = dataset.dataDir + "DBLP-citation-100000.txt"

#dataset.dataFilename = dataset.dataDir + "DBLP-citation-1000000.txt"
#dataset.dataFilename = dataset.dataDir + "DBLP-citation-5000000.txt"
#dataset.dataFilename = dataset.dataDir + "DBLP-citation-7000000.txt"
dataset.dataFilename = dataset.dataDir + "DBLP-citation-Feb21.txt"
dataset.ks = [100, 200, 300, 400, 500, 600]
dataset.minDfs = [0.05, 0.01, 0.005, 0.001]
dataset.overwriteGraph = True
dataset.overwriteModel = True

dataset.overwriteVectoriser = True
dataset.vectoriseDocuments()
dataset.overwriteVectoriser = False

dataset.modelSelection()
dataset.overwriteModel = False

for field in dataset.fields:
    logging.debug("Field = " + field)
    dataset.learnModel()
    relevantExperts = dataset.findSimilarDocuments(field)

    graph, authorIndexer = dataset.coauthorsGraph(field, relevantExperts)
    expertMatches = dataset.matchExperts(relevantExperts, dataset.testExpertDict[field])

    expertMatchesInds = authorIndexer.translate(expertMatches)
    relevantAuthorInds = authorIndexer.translate(relevantExperts)
    assert (numpy.array(relevantAuthorInds) < len(relevantAuthorInds)).all()

    if len(expertMatches) != 0:
        #First compute graph properties
        computeInfluence = False
        graphRanker = GraphRanker(k=100, numRuns=100, computeInfluence=computeInfluence, p=0.05, trainExpertsIdList=expertMatchesInds)
        outputLists = graphRanker.vertexRankings(graph, relevantAuthorInds, [relevantAuthorInds])
        itemList = RankAggregator.generateItemList(outputLists)
        #methodNames = graphRanker.getNames()

        numMethods = len(outputLists)
        precisions = numpy.zeros((len(ns), numMethods))
        averagePrecisions = numpy.zeros(numMethods)

        for i, n in enumerate(ns):
            for j in range(len(outputLists)):
                precisions[i, j] = Evaluator.precisionFromIndLists(expertMatchesInds, outputLists[j][0:n])

        for j in range(len(outputLists)):
            averagePrecisions[j] = Evaluator.averagePrecisionFromLists(expertMatchesInds, outputLists[j][0:averagePrecisionN], averagePrecisionN)

        precisions2 = numpy.c_[numpy.array(ns), precisions]

        logging.debug(Latex.array2DToRows(precisions2))
        logging.debug(Latex.array1DToRow(averagePrecisions))

        resultsFilename = dataset.getResultsDir(field) + "precisions.npz"
        numpy.savez(resultsFilename, precisions, averagePrecisions)
        logging.debug("Saved precisions as " + resultsFilename)

logging.debug("All done!")
<commit_msg>Param tweak <commit_after>"""
Use the DBLP dataset to recommend experts. Find the optimal parameters.
"""
import gc
import os
import numpy
import logging
import sys
import argparse
from exp.influence2.GraphRanker import GraphRanker
from exp.influence2.RankAggregator import RankAggregator
from exp.influence2.ArnetMinerDataset import ArnetMinerDataset
from apgl.util.Latex import Latex
from apgl.util.Evaluator import Evaluator

logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
numpy.set_printoptions(suppress=True, precision=3, linewidth=160)
numpy.random.seed(21)

parser = argparse.ArgumentParser(description='Run reputation evaluation experiments')
parser.add_argument("-r", "--runLDA", action="store_true", help="Run Latent Dirchlet Allocation")
args = parser.parse_args()

averagePrecisionN = 20
similarityCutoff = 0.30
ns = numpy.arange(5, 55, 5)
runLSI = not args.runLDA

dataset = ArnetMinerDataset(runLSI=runLSI)
#dataset.dataFilename = dataset.dataDir + "DBLP-citation-100000.txt"

#dataset.dataFilename = dataset.dataDir + "DBLP-citation-1000000.txt"
#dataset.dataFilename = dataset.dataDir + "DBLP-citation-5000000.txt"
#dataset.dataFilename = dataset.dataDir + "DBLP-citation-7000000.txt"
dataset.dataFilename = dataset.dataDir + "DBLP-citation-Feb21.txt"
dataset.ks = [100, 200, 300, 400, 500, 600]
dataset.minDfs = [0.2, 0.1, 0.01, 0.001, 0.0001]
dataset.overwriteGraph = True
dataset.overwriteModel = True

dataset.overwriteVectoriser = True
dataset.vectoriseDocuments()
dataset.overwriteVectoriser = False

dataset.modelSelection()
dataset.overwriteModel = False

for field in dataset.fields:
    logging.debug("Field = " + field)
    dataset.learnModel()
    relevantExperts = dataset.findSimilarDocuments(field)

    graph, authorIndexer = dataset.coauthorsGraph(field, relevantExperts)
    expertMatches = dataset.matchExperts(relevantExperts, dataset.testExpertDict[field])

    expertMatchesInds = authorIndexer.translate(expertMatches)
    relevantAuthorInds = authorIndexer.translate(relevantExperts)
    assert (numpy.array(relevantAuthorInds) < len(relevantAuthorInds)).all()

    if len(expertMatches) != 0:
        #First compute graph properties
        computeInfluence = False
        graphRanker = GraphRanker(k=100, numRuns=100, computeInfluence=computeInfluence, p=0.05, trainExpertsIdList=expertMatchesInds)
        outputLists = graphRanker.vertexRankings(graph, relevantAuthorInds, [relevantAuthorInds])
        itemList = RankAggregator.generateItemList(outputLists)
        #methodNames = graphRanker.getNames()

        numMethods = len(outputLists)
        precisions = numpy.zeros((len(ns), numMethods))
        averagePrecisions = numpy.zeros(numMethods)

        for i, n in enumerate(ns):
            for j in range(len(outputLists)):
                precisions[i, j] = Evaluator.precisionFromIndLists(expertMatchesInds, outputLists[j][0:n])

        for j in range(len(outputLists)):
            averagePrecisions[j] = Evaluator.averagePrecisionFromLists(expertMatchesInds, outputLists[j][0:averagePrecisionN], averagePrecisionN)

        precisions2 = numpy.c_[numpy.array(ns), precisions]

        logging.debug(Latex.array2DToRows(precisions2))
        logging.debug(Latex.array1DToRow(averagePrecisions))

        resultsFilename = dataset.getResultsDir(field) + "precisions.npz"
        numpy.savez(resultsFilename, precisions, averagePrecisions)
        logging.debug("Saved precisions as " + resultsFilename)

logging.debug("All done!")
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tqdm import tqdm
import pandas as pd
import tensorflow as tf

# Connect to Google's BigQuery
from google.cloud import bigquery
client = bigquery.Client(project="wikidetox-viz")

FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_integer("zones_num", 10, "Number of zones to be used for sampling.")
tf.app.flags.DEFINE_string("output_sample_path", None, "Path where the output sample CSV should be written.")
tf.app.flags.DEFINE_string("output_short_sample_path", None, "Path where the small (100) output sample CSV should be written.")

def query(low_tox=0, high_tox=1, rand=.5, limit=1000, type="ADDITION", max_txt=1000):
    """
    Google Big Query sampling data of specific types
    :param low_tox: comments should be higher than that
    :param high_tox: comments should be lower than that
    :param rand: random selection (p=50%)
    :param limit: fetch some
    :param type: kind of comments (wikipedia specific)
    :param max_txt: max length of comment
    :return: dataframe with results
    """
    query = """
    SELECT *
        FROM `wikiconv_v2.en_20180701_external`
        WHERE length(cleaned_content)>1 AND length(cleaned_content)<{max_txt} AND replyTo_id!="null" AND type="{typ}" AND toxicity>{low} AND toxicity<{high} AND RAND()<{ran}
        LIMIT {lim}
    """.format(low=low_tox, high=high_tox, ran=rand, lim=limit, typ=type, max_txt=max_txt)
    query_job = client.query(query)
    results = query_job.result()
    return results.to_dataframe()

def sample(toxic_zones_num=10):
    """
    Create a sample of comments using toxic zones and a random batch.
    :param toxic_zones_num:
    :return: toxic_zones_num samples + one with randomly selected comments
    """
    # Create random sample (initialization)
    table = query()
    table["toxic_zone"] = [-1] * table.shape[0]
    # Add samples following a toxicity-preselection
    for i in tqdm(range(toxic_zones_num)):
        preselected = query(low_tox=i / float(toxic_zones_num), high_tox=(i + 1) / float(toxic_zones_num))
        preselected["toxic_zone"] = [i] * preselected.shape[0]
        # Stack to initialized/updated table
        table = pd.concat([table, preselected], axis=0)
    # Shuffle
    table = table.sample(frac=1).reset_index(drop=True)
    return table

if __name__ == "__main__":
    table = sample(FLAGS.zones_num)
    table.to_csv(FLAGS.output_sample_path, index=False) # e.g., "sample.11-toxic-zones.addition.csv"
    table[:100].to_csv(FLAGS.output_sample_path, index=False) # e.g., table[:100].to_csv(, index=False)

<commit_msg>Addressing comments by Lucas: flags updated with more variables random selection in query fixed licence<commit_after># coding=utf-8
# Copyright 2018 The Conversation-AI.github.io Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Preselect Wikipedia comments.

A library to preselect wikipedia comments based on their toxicity score
among others aspects.

The initial goal was to create a sample of comments to use for annotation;
the sample was stratified to ensure existence of toxic comments. If one
uses 3 zones, then [0,1] will be divided in [0,0.33), [0.33,0.66) and
[0.66,1], so that the comments are scored from all three zones.

The BigQuery table queried to create this sample contains 240M rows (332Gb).

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tqdm import tqdm
import pandas as pd
import tensorflow as tf
from google.cloud import bigquery

FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_integer("limit", 100, "Number of comments per zone (zones_num*limit + 1).")
tf.app.flags.DEFINE_integer("randomly_selected_limit", 100, "Number of comments selected at random, not based on toxicity.")
tf.app.flags.DEFINE_string("comment_type", "ADDITION", "The type of comment ('ADDITION', 'DELETION', 'CREATION', 'MODIFICATION').")
tf.app.flags.DEFINE_integer("max_txt_len", 1000, "Maximum length for comments of the sample.")
tf.app.flags.DEFINE_string("bgq_project_name", None, "The name of the Google BigQuery project.")
tf.app.flags.DEFINE_integer("zones_num", 10, "Number of zones to be used for sampling.")
tf.app.flags.DEFINE_string("output_csv_path", None, "Path where the output sample CSV should be written.")
tf.app.flags.DEFINE_string("output_json_path", None, "Path where the output sample JSON should be written.")

def connect_to_gbq(gbq_project_name):
    """
    Connect to Google BigQuery
    :param gbq_project_name: project name
    :return:
    """
    client = bigquery.Client(project=gbq_project_name)
    return client

def query(low_tox, high_tox, limit, max_txt_len, type):
    """
    Google Big Query sampling data of specific types
    :param low_tox: comments should be higher than that
    :param high_tox: comments should be lower than that
    :param limit: fetch some
    :param type: kind of comments (wikipedia specific)
    :param max_txt_len: max length of comment
    :return: dataframe with results
    """
    assert low_tox>=0 and high_tox<=1
    query = """
    SELECT *, rand() as r
        FROM `wikiconv_v2.en_20180701_external`
        WHERE length(cleaned_content)>1 AND length(cleaned_content)<{max_txt_len} AND replyTo_id!="null" AND type="{typ}" AND toxicity>{low} AND toxicity<{high}
        order by r
        LIMIT {lim}
    """.format(low=low_tox, high=high_tox, lim=limit, typ=type, max_txt_len=max_txt_len)
    query_job = client.query(query)
    results = query_job.result()
    return results.to_dataframe()

def sample(toxic_zones_num, max_txt_len, type, limit, randomly_selected_limit):
    """
    Create a sample of comments using toxic zones and a random batch.
    :param toxic_zones_num:
    :return: toxic_zones_num samples + one with randomly selected comments
    """
    # Create random sample (initialization)
    table = query(low_tox=0, high_tox=1, max_txt_len=max_txt_len, limit=randomly_selected_limit, type=type)
    table["toxic_zone"] = [-1] * table.shape[0]
    # Add samples following a toxicity-preselection
    for i in tqdm(range(toxic_zones_num)):
        preselected = query(low_tox=float(i)/toxic_zones_num, high_tox=float(i+1)/toxic_zones_num, max_txt_len=max_txt_len, limit=limit, type=type)
        preselected["toxic_zone"] = [i] * preselected.shape[0]
        # Stack to initialized/updated table
        table = pd.concat([table, preselected], axis=0)
    # Shuffle
    table = table.sample(frac=1).reset_index(drop=True)
    return table

if __name__ == "__main__":
    assert FLAGS.bgq_project_name is not None
    assert (FLAGS.output_json_path is None and FLAGS.output_csv_path is not None) or (FLAGS.output_json_path is not None and FLAGS.output_csv_path is None)
    assert FLAGS.comment_type in {'ADDITION', 'DELETION', 'CREATION', 'MODIFICATION'}
    # connect to gbq, sample and save
    client = connect_to_gbq(FLAGS.bgq_project_name)
    table = sample(toxic_zones_num=FLAGS.zones_num, max_txt_len=FLAGS.max_txt_len, limit=FLAGS.limit, type=FLAGS.comment_type, randomly_selected_limit=FLAGS.randomly_selected_limit)
    if FLAGS.output_csv_path is not None:
        # e.g., "sample.11-toxic-zones.addition.csv"
        table.to_csv(FLAGS.output_csv_path, index=False)
    elif FLAGS.output_json_path is not None:
        # e.g., "sample.11-toxic-zones.addition.json"
        table.to_json(path_or_buf=FLAGS.output_json_path, orient="records", lines=True)
    else:
        print ("Not implemented yet.")<|endoftext|>


<commit_before>"""The entity endpoint"""

from flask import jsonify, json, g, request, url_for
from flask.views import MethodView
from mongoengine import ValidationError

from tentd.flask import Blueprint
from tentd.control import follow
from tentd.utils.exceptions import APIException, APIBadRequest
from tentd.documents.entity import Entity, Follower, Post

entity = Blueprint('entity', __name__, url_prefix='/<string:entity>')

@entity.url_value_preprocessor
def fetch_entity(endpoint, values):
    """Replace `entity` (which is a string) with the actuall entity"""
    values['entity'] = Entity.objects.get_or_404(name=values['entity'])

@entity.route('')
def link(entity):
    link = '<{url}>; rel="https://tent.io/rels/profile"'.format(
        url=url_for('entity.profile', entity=entity.name, _external=True))

    resp = jsonify(entity.to_json())
    resp.headers['Link'] = link

    return resp

@entity.route('/profile')
def profile(entity):
    """Return the info types belonging to the entity"""
    return jsonify({p.schema: p.to_json() for p in entity.profiles})

@entity.route('/followers', methods=['POST'])
def followers(entity):
    """Starts following a user, defined by the post data"""
    try:
        post_data = json.loads(request.data)
    except json.JSONDecodeError:
        raise APIBadRequest()

    if not post_data:
        raise APIBadRequest("No POST data.")

    follower = follow.start_following(entity, post_data)
    return jsonify(follower.to_json())

@entity.route_class('/followers/<string:follower_id>')
class FollowerView(MethodView):
    endpoint = 'follower'

    def get(self, entity, follower_id):
        """Returns the json representation of a follower"""
        return jsonify(entity.followers.get_or_404(id=follower_id).to_json())

    def put(self, entity, follower_id):
        try:
            post_data = json.loads(request.data)
        except json.JSONDecodeError:
            raise JSONBadRequest()
        updated_follower = follow.update_follower(entity, follower_id, post_data)
        return jsonify(updated_follower.to_json())

    def delete(self, entity, follower_id):
        try:
            follow.stop_following(entity, follower_id)
            return '', 200
        except ValidationError:
            raise APIBadRequest("The given follower id was invalid")

@entity.route('/notification', methods=['GET'])
def get_notification(entity):
    """ Alerts of a notification """
    return '', 200

@entity.route('/posts', methods=['GET', 'POST'])
def get_posts(entity):
    """ Returns all public posts. Can be scoped. """
    # TODO Filter to public posts only when that's included.
    # TODO Apply other filters included as part of the GET.
    #      We'll need to think about how we handle these filters as this is
    #      potential unsanatised input from the user and is therefore vunerable
    #      to SQL injection attacks.
    if request.method == 'GET':
        posts=[post.to_json() for post in entity.posts]
        print posts
        return jsonify(posts), 200
    if request.method == 'POST':
        try:
            data = json.loads(request.data)
        except json.JSONDecodeError:
            raise APIBadRequest()
        post = Post()
        post.entity = entity
        post.schema = data['schema']
        post.content = data['content']

        # TODO add in something to this affect:
        # for follower in entity.followers:
        #     follower.notify(post)
        post.save()
        return post.to_json(), 200

@entity.route('/posts/<string:post_id>', methods=['GET', 'PUT', 'DELETE'])
class PostsView:

    def get(self, entity, post_id):
        return entity.posts.get_or_404(id=post_id).to_json(), 200
    def put(self, entity, post_id):
        post = entity.posts.get_or_404(id=post_id)
        try:
            post_data = json.loads(request.data)
        except json.JSONDecodeError:
            raise APIBadRequest()

        #TODO Perform the update

        post.save()
        return post.to_json(), 200
    def delete(self, entity, post_id):
        post = entity.posts.get_or_404(id=post_id)
        #TODO Notify?
        return post.delete(), 200
<commit_msg>Fixed up the APIBadRequest to have error messages.<commit_after>"""The entity endpoint"""

from flask import jsonify, json, g, request, url_for
from flask.views import MethodView
from mongoengine import ValidationError

from tentd.flask import Blueprint
from tentd.control import follow
from tentd.utils.exceptions import APIException, APIBadRequest
from tentd.documents.entity import Entity, Follower, Post

entity = Blueprint('entity', __name__, url_prefix='/<string:entity>')

@entity.url_value_preprocessor
def fetch_entity(endpoint, values):
    """Replace `entity` (which is a string) with the actuall entity"""
    values['entity'] = Entity.objects.get_or_404(name=values['entity'])

@entity.route('')
def link(entity):
    link = '<{url}>; rel="https://tent.io/rels/profile"'.format(
        url=url_for('entity.profile', entity=entity.name, _external=True))

    resp = jsonify(entity.to_json())
    resp.headers['Link'] = link

    return resp

@entity.route('/profile')
def profile(entity):
    """Return the info types belonging to the entity"""
    return jsonify({p.schema: p.to_json() for p in entity.profiles})

@entity.route('/followers', methods=['POST'])
def followers(entity):
    """Starts following a user, defined by the post data"""
    try:
        post_data = json.loads(request.data)
    except json.JSONDecodeError as e:
        raise APIBadRequest(str(e))

    if not post_data:
        raise APIBadRequest("No POST data.")

    follower = follow.start_following(entity, post_data)
    return jsonify(follower.to_json())

@entity.route_class('/followers/<string:follower_id>')
class FollowerView(MethodView):
    endpoint = 'follower'

    def get(self, entity, follower_id):
        """Returns the json representation of a follower"""
        return jsonify(entity.followers.get_or_404(id=follower_id).to_json())

    def put(self, entity, follower_id):
        try:
            post_data = json.loads(request.data)
        except json.JSONDecodeError as e:
            raise APIBadRequest(str(e))
        updated_follower = follow.update_follower(entity, follower_id, post_data)
        return jsonify(updated_follower.to_json())

    def delete(self, entity, follower_id):
        try:
            follow.stop_following(entity, follower_id)
            return '', 200
        except ValidationError:
            raise APIBadRequest("The given follower id was invalid")

@entity.route('/notification', methods=['GET'])
def get_notification(entity):
    """ Alerts of a notification """
    return '', 200

@entity.route('/posts', methods=['GET', 'POST'])
def get_posts(entity):
    """ Returns all public posts. Can be scoped. """
    # TODO Filter to public posts only when that's included.
    # TODO Apply other filters included as part of the GET.
    #      We'll need to think about how we handle these filters as this is
    #      potential unsanatised input from the user and is therefore vunerable
    #      to SQL injection attacks.
    if request.method == 'GET':
        posts=[post.to_json() for post in entity.posts]
        return jsonify(posts), 200
    if request.method == 'POST':
        try:
            data = json.loads(request.data)
        except json.JSONDecodeError as e:
            raise APIBadRequest(str(e))
        post = Post()
        post.entity = entity
        post.schema = data['schema']
        post.content = data['content']

        # TODO add in something to this affect:
        # for follower in entity.followers:
        #     follower.notify(post)
        post.save()
        return post.to_json(), 200

@entity.route('/posts/<string:post_id>', methods=['GET', 'PUT', 'DELETE'])
class PostsView:

    def get(self, entity, post_id):
        return entity.posts.get_or_404(id=post_id).to_json(), 200
    def put(self, entity, post_id):
        post = entity.posts.get_or_404(id=post_id)
        try:
            post_data = json.loads(request.data)
        except json.JSONDecodeError as e:
            raise APIBadRequest(str(e))

        #TODO Perform the update

        post.save()
        return post.to_json(), 200
    def delete(self, entity, post_id):
        post = entity.posts.get_or_404(id=post_id)
        #TODO Notify?
        return post.delete(), 200
<|endoftext|>