<commit_before># clean, reroot, ladderize newick tree
# output to tree.json
import numpy as np
import time
from collections import defaultdict
from matplotlib import pyplot as plt
from itertools import izip

class HI_tree(object):

	def __init__(self, HI_fname = 'source-data/HI_titers.txt',**kwargs):
		self.HI, tmp = self.read_HI_titers(HI_fname)

	def read_HI_titers(self, fname):
		strains = set()
		measurements = {}
		with open(fname, 'r') as infile:
			for line in infile:
				entries = line.strip().split()
				test, ref, val = entries[0], entries[1], map(float, entries[2:])
				if len(val):
					measurements[(test, ref)] = val
					strains.update([test, ref])
				else:
					print line.strip()
		return measurements, strains

	def normalize_HI(self):
		consensus_func = np.mean
		self.HI_normalized = {}
		sera = set()
		HI_strains = set()
		for (test, ref), val in self.HI.iteritems():
			if test.lower() in self.node_lookup and ref.lower() in self.node_lookup:
				HI_strains.add(test.lower())
				HI_strains.add(ref.lower())
				if (ref,ref) in self.HI:
					sera.add(ref)
					normalized_val = consensus_func(np.log2(self.HI[(ref, ref)])) - consensus_func(np.log2(val))
					self.HI_normalized[(test, ref)] = normalized_val
		self.sera = list(sera)
		self.HI_strains = list(HI_strains)

	def add_mutations(self):
		'''
		add amino acid mutations to the tree
		'''
		self.tree.seed_node.mutations= ''
		for node in self.tree.postorder_node_iter():
			if node is not self.tree.seed_node:
				node.mutations = [a+str(pos+1)+b for pos, (a,b) in 
								enumerate(izip(node.parent_node.aa_seq, node.aa_seq)) if a!=b]

	def mark_HI_splits(self):
		# flag all branches on the tree with HI_strain = True if they lead to strain with titer data
		for leaf in self.tree.leaf_iter():
			if leaf.strain.lower() in self.HI_strains:
				leaf.serum = leaf.strain.lower() in self.sera
				leaf.HI_info= 1
			else:
				leaf.serum, leaf.HI_info=False, 0

		for node in self.tree.postorder_node_iter():
			if not node.is_leaf():
				node.HI_info = sum([c.HI_info for c in node.child_nodes()])

		# combine sets of branches that span identical sets of HI measurements
		self.HI_split_count = 0  # HI measurment split counter
		self.HI_split_to_branch = defaultdict(list)
		for node in self.tree.preorder_node_iter():
			node.dHI, node.cHI, node.constraints =0, 0, 0
			if node.is_internal(): node.serum=False
			if node.HI_info>1:
				node.HI_branch_index = self.HI_split_count
				self.HI_split_to_branch[node.HI_branch_index].append(node)
				if sum([c.HI_info for c in node.child_nodes()])>1:
					self.HI_split_count+=1

		print "# of reference strains:",len(self.sera), "# of branches with HI constraint", self.HI_split_count

	def get_path_no_terminals(self, v1, v2):
		if v1 in self.node_lookup and v2 in self.node_lookup:
			p1 = [self.node_lookup[v1]]
			p2 = [self.node_lookup[v2]]
			for tmp_p in [p1,p2]:
				while tmp_p[-1].parent_node != self.tree.seed_node:
					tmp_p.append(tmp_p[-1].parent_node)
				tmp_p.append(self.tree.seed_node)
				tmp_p.reverse()

			for pi, (tmp_v1, tmp_v2) in enumerate(izip(p1,p2)):
				if tmp_v1!=tmp_v2:
					break
			path = p1[pi:-1] + p2[pi:-1]
		else:
			path = None
		return path

	def make_treegraph(self):
		tree_graph = []
		HI_dist = []
		n_params = self.HI_split_count + len(self.sera) + len(self.HI_strains)
		for (test, ref), val in self.train_HI.iteritems():
			if not np.isnan(val):
				try:
					if ref in self.node_lookup and test in self.node_lookup:
						path = self.get_path_no_terminals(test, ref)
						tmp = np.zeros(n_params)
						branches = np.unique([c.HI_branch_index for c in path if hasattr(c, 'HI_branch_index')])
						if len(branches): tmp[branches] = 1
						tmp[self.HI_split_count+self.sera.index(ref)] = 1
						tree_graph.append(tmp)
						tmp[self.HI_split_count+self.sera.index(ref)] = 1
						tmp[self.HI_split_count+len(self.sera)+self.HI_strains.index(test)] = 1
						HI_dist.append(val)
				except:
					import pdb; pdb.set_trace()
					print test, ref, "ERROR"

		self.HI_dist =  np.array(HI_dist)
		self.tree_graph= np.array(tree_graph)
		print "Found", self.tree_graph.shape, "measurements x parameters"

	def fit_func(self):
		return np.mean( (self.HI_dist - np.dot(self.tree_graph, self.params))**2 )

	def fit_l1reg(self):
		from cvxopt import matrix
		from l1regls import l1regls
		A = matrix(self.tree_graph)
		b = matrix(self.HI_dist)
		return np.array([x for x in l1regls(A/np.sqrt(self.lam),b/np.sqrt(self.lam))])

	def fit_nnls(self):
		from scipy.optimize import nnls
		return nnls(self.tree_graph, self.HI_dist)[0]

	def fit_nnl2reg(self):
		from cvxopt import matrix, solvers
		n_params = self.tree_graph.shape[1]
		P = matrix(np.dot(self.tree_graph.T, self.tree_graph) + self.lam*np.eye(n_params))
		q = matrix( -np.dot( self.HI_dist, self.tree_graph))
		h = matrix(np.zeros(n_params)) # Gw <=h
		G = matrix(-np.eye(n_params))
		W = solvers.qp(P,q,G,h)
		return np.array([x for x in W['x']])

	def fit_nnl1reg(self):
		from cvxopt import matrix, solvers
		n_params = self.tree_graph.shape[1]
		HI_sc = self.HI_split_count
		n_sera = len(self.sera)
		n_v = len(self.HI_strains)
		P1 = np.zeros((n_params+HI_sc,n_params+HI_sc))
		P1[:n_params, :n_params] = np.dot(self.tree_graph.T, self.tree_graph)
		for ii in xrange(HI_sc, n_params):
			P1[ii,ii]+=self.lam
		P = matrix(P1)

		q1 = np.zeros(n_params+HI_sc)
		q1[:n_params] = -np.dot( self.HI_dist, self.tree_graph)
		q1[n_params:] = self.lam
		q = matrix(q1)

		h = matrix(np.zeros(2*HI_sc)) 	# Gw <=h
		G1 = np.zeros((2*HI_sc,n_params+HI_sc))
		G1[:HI_sc, :HI_sc] = -np.eye(HI_sc)
		G1[HI_sc:, :HI_sc] = np.eye(HI_sc)
		G1[HI_sc:, n_params:] = -np.eye(HI_sc)
		G = matrix(G1)
		W = solvers.qp(P,q,G,h)
		return np.array([x for x in W['x']])[:n_params]

	def map_HI_to_tree(self, training_fraction = 1.0, method = 'nnls', lam=10, cutoff_date = None):
		self.normalize_HI()
		self.add_mutations()
		self.mark_HI_splits()
		self.lam = lam
		if training_fraction<1.0:
			self.test_HI, self.train_HI = {}, {}
			for key, val in self.HI_normalized.iteritems():
				if np.random.uniform()>training_fraction:
					self.test_HI[key]=val
				else:
					self.train_HI[key]=val
		else:
			self.train_HI = self.HI_normalized

		if cutoff_date is not None:
			self.train_HI = {key:val for key,val in self.train_HI.iteritems()
							if self.node_lookup[key[0]].num_date<cutoff_date and 
							   self.node_lookup[key[1]].num_date<cutoff_date}

		self.make_treegraph()
		if method=='l1reg':  # l1 regularized fit, no constraint on sign of effect
			self.params = self.fit_l1reg()
		elif method=='nnls':  # non-negative least square, not regularized
			self.params = self.fit_nnls()
		elif method=='nnl2reg':	# non-negative L2 norm regularized fit
			self.params = self.fit_nnl2reg()
		elif method=='nnl1reg':  # non-negative fit, branch terms L1 regularized, avidity terms L2 regularized
			self.params = self.fit_nnl1reg()
<commit_msg>added comments and a relaxatio of the HI -> tree map to avoid shortening of dHI terms due to the regularization.<commit_after># clean, reroot, ladderize newick tree
# output to tree.json
import numpy as np
import time
from collections import defaultdict
from matplotlib import pyplot as plt
from itertools import izip

class HI_tree(object):

	def __init__(self, HI_fname = 'source-data/HI_titers.txt',**kwargs):
		self.HI, tmp = self.read_HI_titers(HI_fname)

	def read_HI_titers(self, fname):
		strains = set()
		measurements = {}
		with open(fname, 'r') as infile:
			for line in infile:
				entries = line.strip().split()
				test, ref, val = entries[0], entries[1], map(float, entries[2:])
				if len(val):
					measurements[(test, ref)] = val
					strains.update([test, ref])
				else:
					print line.strip()
		return measurements, strains

	def normalize_HI(self):
		consensus_func = np.mean
		self.HI_normalized = {}
		sera = set()
		HI_strains = set()
		for (test, ref), val in self.HI.iteritems():
			if test.lower() in self.node_lookup and ref.lower() in self.node_lookup:
				HI_strains.add(test.lower())
				HI_strains.add(ref.lower())
				if (ref,ref) in self.HI:
					sera.add(ref)
					normalized_val = consensus_func(np.log2(self.HI[(ref, ref)])) - consensus_func(np.log2(val))
					self.HI_normalized[(test, ref)] = normalized_val
		self.sera = list(sera)
		self.HI_strains = list(HI_strains)

	def add_mutations(self):
		'''
		add amino acid mutations to the tree
		'''
		self.tree.seed_node.mutations= ''
		for node in self.tree.postorder_node_iter():
			if node is not self.tree.seed_node:
				node.mutations = [a+str(pos+1)+b for pos, (a,b) in 
								enumerate(izip(node.parent_node.aa_seq, node.aa_seq)) if a!=b]

	def mark_HI_splits(self):
		# flag all branches on the tree with HI_strain = True if they lead to strain with titer data
		for leaf in self.tree.leaf_iter():
			if leaf.strain.lower() in self.HI_strains:
				leaf.serum = leaf.strain.lower() in self.sera
				leaf.HI_info= 1
			else:
				leaf.serum, leaf.HI_info=False, 0

		for node in self.tree.postorder_node_iter():
			if not node.is_leaf():
				node.HI_info = sum([c.HI_info for c in node.child_nodes()])

		# combine sets of branches that span identical sets of HI measurements
		self.HI_split_count = 0  # HI measurment split counter
		self.HI_split_to_branch = defaultdict(list)
		for node in self.tree.preorder_node_iter():
			node.dHI, node.cHI, node.constraints =0, 0, 0
			if node.is_internal(): node.serum=False
			if node.HI_info>1:
				node.HI_branch_index = self.HI_split_count
				self.HI_split_to_branch[node.HI_branch_index].append(node)
				if sum([c.HI_info for c in node.child_nodes()])>1:
					self.HI_split_count+=1

		print "# of reference strains:",len(self.sera), "# of branches with HI constraint", self.HI_split_count

	def get_path_no_terminals(self, v1, v2):
		if v1 in self.node_lookup and v2 in self.node_lookup:
			p1 = [self.node_lookup[v1]]
			p2 = [self.node_lookup[v2]]
			for tmp_p in [p1,p2]:
				while tmp_p[-1].parent_node != self.tree.seed_node:
					tmp_p.append(tmp_p[-1].parent_node)
				tmp_p.append(self.tree.seed_node)
				tmp_p.reverse()

			for pi, (tmp_v1, tmp_v2) in enumerate(izip(p1,p2)):
				if tmp_v1!=tmp_v2:
					break
			path = p1[pi:-1] + p2[pi:-1]
		else:
			path = None
		return path

	def make_treegraph(self):
		tree_graph = []
		HI_dist = []
		n_params = self.HI_split_count + len(self.sera) + len(self.HI_strains)
		for (test, ref), val in self.train_HI.iteritems():
			if not np.isnan(val):
				try:
					if ref in self.node_lookup and test in self.node_lookup:
						path = self.get_path_no_terminals(test, ref)
						tmp = np.zeros(n_params)
						branches = np.unique([c.HI_branch_index for c in path if hasattr(c, 'HI_branch_index')])
						if len(branches): tmp[branches] = 1
						tmp[self.HI_split_count+self.sera.index(ref)] = 1
						tree_graph.append(tmp)
						tmp[self.HI_split_count+self.sera.index(ref)] = 1
						tmp[self.HI_split_count+len(self.sera)+self.HI_strains.index(test)] = 1
						HI_dist.append(val)
				except:
					import pdb; pdb.set_trace()
					print test, ref, "ERROR"

		self.HI_dist =  np.array(HI_dist)
		self.tree_graph= np.array(tree_graph)
		print "Found", self.tree_graph.shape, "measurements x parameters"

	def fit_func(self):
		return np.mean( (self.HI_dist - np.dot(self.tree_graph, self.params))**2 )

	def fit_l1reg(self):
		from cvxopt import matrix
		from l1regls import l1regls
		A = matrix(self.tree_graph)
		b = matrix(self.HI_dist)
		return np.array([x for x in l1regls(A/np.sqrt(self.lam),b/np.sqrt(self.lam))])

	def fit_nnls(self):
		from scipy.optimize import nnls
		return nnls(self.tree_graph, self.HI_dist)[0]

	def fit_nnl2reg(self):
		from cvxopt import matrix, solvers
		n_params = self.tree_graph.shape[1]
		P = matrix(np.dot(self.tree_graph.T, self.tree_graph) + self.lam*np.eye(n_params))
		q = matrix( -np.dot( self.HI_dist, self.tree_graph))
		h = matrix(np.zeros(n_params)) # Gw <=h
		G = matrix(-np.eye(n_params))
		W = solvers.qp(P,q,G,h)
		return np.array([x for x in W['x']])

	def fit_nnl1reg(self):
		from cvxopt import matrix, solvers
		n_params = self.tree_graph.shape[1]
		HI_sc = self.HI_split_count
		n_sera = len(self.sera)
		n_v = len(self.HI_strains)

		# set up the quadratic matrix containing the deviation term (linear xterm below)
		# and the l2-regulatization of the avidities and potencies
		P1 = np.zeros((n_params+HI_sc,n_params+HI_sc))
		P1[:n_params, :n_params] = np.dot(self.tree_graph.T, self.tree_graph)
		for ii in xrange(HI_sc, n_params):
			P1[ii,ii]+=self.lam
		P = matrix(P1)

		# set up cost for auxillary parameter and the linear cross-term
		q1 = np.zeros(n_params+HI_sc)
		q1[:n_params] = -np.dot( self.HI_dist, self.tree_graph)
		q1[n_params:] = self.lam
		q = matrix(q1)

		# set up linear constraint matrix to enforce positivity of the
		# dHIs and bounding of dHI by the auxillary parameter
		h = matrix(np.zeros(2*HI_sc)) 	# Gw <=h
		G1 = np.zeros((2*HI_sc,n_params+HI_sc))
		G1[:HI_sc, :HI_sc] = -np.eye(HI_sc)
		G1[HI_sc:, :HI_sc] = np.eye(HI_sc)
		G1[HI_sc:, n_params:] = -np.eye(HI_sc)
		G = matrix(G1)
		W = solvers.qp(P,q,G,h)
		sol = np.array([x for x in W['x']])[:n_params]

		# redo the linear cost relaxing terms that seem to be relevant to avoid 
		# compression of the fit. 0.2 seems to be a good cut-off, linear tune to zero
		q1[n_params:] = self.lam*(1-5.0*np.minimum(0.2,sol[:HI_sc]))
		q = matrix(q1)
		W = solvers.qp(P,q,G,h)
		sol = np.array([x for x in W['x']])[:n_params]

		return sol

	def map_HI_to_tree(self, training_fraction = 1.0, method = 'nnls', lam=10, cutoff_date = None):
		self.normalize_HI()
		self.add_mutations()
		self.mark_HI_splits()
		self.lam = lam
		if training_fraction<1.0:
			self.test_HI, self.train_HI = {}, {}
			for key, val in self.HI_normalized.iteritems():
				if np.random.uniform()>training_fraction:
					self.test_HI[key]=val
				else:
					self.train_HI[key]=val
		else:
			self.train_HI = self.HI_normalized

		if cutoff_date is not None:
			self.train_HI = {key:val for key,val in self.train_HI.iteritems()
							if self.node_lookup[key[0]].num_date<cutoff_date and 
							   self.node_lookup[key[1]].num_date<cutoff_date}

		self.make_treegraph()
		if method=='l1reg':  # l1 regularized fit, no constraint on sign of effect
			self.params = self.fit_l1reg()
		elif method=='nnls':  # non-negative least square, not regularized
			self.params = self.fit_nnls()
		elif method=='nnl2reg':	# non-negative L2 norm regularized fit
			self.params = self.fit_nnl2reg()
		elif method=='nnl1reg':  # non-negative fit, branch terms L1 regularized, avidity terms L2 regularized
			self.params = self.fit_nnl1reg()

{'fraction_non_alphanumeric': 0.13360984565173048, 'fraction_numerical': 0.01210308016590739, 'mean_word_length': 3.825459317585302, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3478655', 'n_tokens_mistral': 6512, 'n_tokens_neox': 5699, 'n_words': 1282}