text:
import torch as th
import torch.nn as nn
import torch.nn.functional as F
import pandas as pd
import numpy as np
from dgl.nn import MaxPooling, AvgPooling, GlobalAttentionPooling
from dgl import function as fn
from dgl.nn.pytorch.softmax import edge_softmax
from dgllife.model import GATPredictor, GIN, load_pretrained, GAT, WLN

"""
- in general, for DGLLife, GNN is just a graph layer while GNNPredictor includes a graph layer, a readout layer, and an mlp layer
- 4 types of pretrained models: moleculenet, generative_models, property_prediction, reaction
- property_prediction are GIN supervised (contextpred,infomax,edgepred,masking) and AttentiveFP, they do not contain reaodut and output MLP
- moleculenet models are a variety of architectures fine-tuned on moleculenet benchmark tasks (Tox21, PCBA, etc), the do contain readout and output MLP
- reaction and generative_models are not relevant
"""

def get_model(args, data_feats):

	model_type = args["model_type"]

	if model_type == "gat_p":
		model = GATP(data_feats["n_feats"],**args)

	elif model_type == "gin":
		model = GIN1(data_feats["n_feats"],data_feats["e_feats"],**args)

	elif model_type == "gat":
		model = GAT1(data_feats["n_feats"], data_feats["grover_fp_feats"], **args)

	elif model_type == "mlp_fingerprint":
		model = MLP(data_feats["finger_feats"], data_feats["grover_fp_feats"], **args)

	elif model_type == "gin_pt1":
		pretrain_name = "gin_supervised_contextpred_Tox21"
		model = GINPT(
			data_feats["n_feats"],
			data_feats["e_feats"],
			pretrain_name,
			data_feats["grover_fp_feats"],
			**args)
	elif model_type == "gin_pt2":
		pretrain_name = "gin_supervised_infomax_Tox21"
		model = GINPT(
			data_feats["n_feats"],
			data_feats["e_feats"],
			pretrain_name,
			data_feats["grover_fp_feats"],
			**args)
	elif model_type == "gin_pt3":
		pretrain_name = "gin_supervised_infomax_ToxCast"
		model = GINPT(
			data_feats["n_feats"],
			data_feats["e_feats"],
			pretrain_name,
			data_feats["grover_fp_feats"],
			**args)
	elif model_type == "gin_pt4":
		pretrain_name = "gin_supervised_infomax"
		model = GINPT(
			data_feats["n_feats"],
			data_feats["e_feats"],
			pretrain_name,
			data_feats["grover_fp_feats"],
			**args)
	elif model_type == "cnn":
		model = CNN(data_feats["c_feats"], data_feats["l_feats"],**args)
	elif model_type == "grover_pt":
		model = GROVERPT(data_feats["grover_fp_feats"],**args)
	elif model_type == "wln":
		model = WLN1(data_feats["n_feats"],data_feats["e_feats"],**args)
	else:
		print(model_type)
		raise ValueError(f"invalid model_type {model_type}")
	return model

class Identity(nn.Module):
	"""
	for deleting layers in pretrained models
	from https://discuss.pytorch.org/t/how-to-delete-layer-in-pretrained-model/17648
	"""

	def __init__(self):
		super(Identity,self).__init__()

	def forward(self,x):
		return x

class GNIdentity(Identity):
	"""
	identity for operation that takes graph, node attributes (for example, graph pool layer)
	"""

	def forward(self,g,nh):
		return nh

class GNEIdentity(Identity):
	"""
	identity for operation that takes graph, node, edge attributes (for example. message passing layer)
	"""

	def forward(self,g,nh,eh):
		return nh

class LinearBlock(nn.Module):
	""" basic block in an MLP, with dropout and batch norm """

	def __init__(self, in_feats, out_feats, dropout=0.1):
		super(LinearBlock, self).__init__()
		self.linear = nn.Linear(in_feats, out_feats)
		self.bn = nn.BatchNorm1d(out_feats)
		self.dropout = nn.Dropout(dropout)

	def forward(self, x):
		return self.bn(self.dropout(F.relu(self.linear(x))))

class GATP(nn.Module):
	""" silly wrapper around GATPredictor1 """

	def __init__(self,n_feats,**kwargs):

		super(GATP,self).__init__()
		model_d = {}
		model_d['in_feats'] = n_feats
		model_d['hidden_feats'] = [kwargs['gnn_hidden_feats'] for x in range(kwargs['gnn_num_layers'])]
		model_d['num_heads'] = [kwargs['gat_num_heads'] for x in range(kwargs['gnn_num_layers'])]
		model_d['feat_drops'] = [kwargs['gnn_dropout'] for x in range(kwargs['gnn_num_layers'])]
		model_d['attn_drops'] = [kwargs['gnn_dropout'] for x in range(kwargs['gnn_num_layers'])]
		model_d["alphas"] = [0.2 for _ in range(kwargs['gnn_num_layers'])]
		# residuals list of bool for number of layers
		model_d['residuals'] = [True for _ in range(kwargs['gnn_num_layers'])]
		"""
		The way to aggregate multi-head attention results for each GAT layer, which can be either
		'flatten' for concatenating all-head results or 'mean' for averaging all-head results.
		``agg_modes[i]`` gives the way to aggregate multi-head attention results for the i-th
		GAT layer. ``len(agg_modes)`` equals the number of GAT layers. By default, we flatten
		multi-head results for intermediate GAT layers and compute mean of multi-head results
		for the last GAT layer.
		"""
		model_d['agg_modes'] = ['flatten' for _ in range(kwargs['gnn_num_layers'] - 1)] + ['mean']
		# Apply Elu to all the layers besides the last one
		model_d['activations'] = [nn.ELU(model_d['alphas'][i]) for i in range(kwargs['gnn_num_layers'] - 1)] + [None]
		model_d['n_tasks'] = kwargs['out_feats']
		model_d['predictor_hidden_feats'] = kwargs['ff_hidden_feats']
		model_d['predictor_dropout'] = kwargs['ff_dropout']
		print(model_d)
		self.gnn_p = GATPredictor(**model_d)

	def forward(self,input_d):
		""" this is compatible with canonical and pretrafsefin featurization """

		graph = input_d["graph"]
		node_h = graph.ndata["h"]
		return self.gnn_p(graph,node_h)

class MLP(nn.Module):
	def __init__(self, n_feats, grover_fp_feats, **kwargs):

		super(MLP, self).__init__()
		self.n_feats = n_feats
		self.grover_fp_feats = grover_fp_feats # can be 0
		for k,v in kwargs.items():
			setattr(self, k, v)

		self.ff_layers = construct_mlp(
			self.n_feats + self.grover_fp_feats,
			self.out_feats,
			self.ff_hidden_feats,
			self.ff_num_layers,
			self.ff_dropout
		)

	def forward(self, input_d):
		fingerprints = input_d["fingerprints"]

		if self.grover_fp in ['base', 'large']:
			output = self.ff_layers(th.cat([fingerprints, input_d["grover_fp"]], dim=1))
		else:
			output = self.ff_layers(fingerprints)

		return output

class GAT1(nn.Module):

	def __init__(self, n_feats, grover_fp_feats, **kwargs):

		super(GAT1,self).__init__()
		self.n_feats = n_feats
		self.grover_fp_feats = grover_fp_feats # can be 0
		for k,v in kwargs.items():
			setattr(self,k,v)
		self.gnn_layers = GAT(
			in_feats=self.n_feats,
			hidden_feats=[self.gnn_hidden_feats for i in range(self.gnn_num_layers)],
			num_heads=[self.gat_num_heads for i in range(self.gnn_num_layers)],
			feat_drops=[self.gnn_dropout for i in range(self.gnn_num_layers)],
			attn_drops=[self.gnn_dropout for i in range(self.gnn_num_layers)],
			alphas=[0.2 for i in range(self.gnn_num_layers)],
			residuals=[True for i in range(self.gnn_num_layers)],
			agg_modes=["flatten" for i in range(self.gnn_num_layers-1)] + ["mean"],
			activations=[nn.ELU(0.2) for i in range(self.gnn_num_layers-1)] + [None]
		)
		self.pool = construct_pool(self.gnn_hidden_feats,self.gnn_pool_type)
		if not self.separate_heads:
			self.ff_layers = construct_mlp(
				self.gnn_hidden_feats + self.grover_fp_feats,
				self.out_feats,
				self.ff_hidden_feats,
				self.ff_num_layers,
				self.ff_dropout
			)
		else:
			assert self.out_feats % 2 == 0
			self.ff_layers1 = construct_mlp(
				self.gnn_hidden_feats + self.grover_fp_feats,
				int(self.out_feats / 2),
				self.ff_hidden_feats,
				self.ff_num_layers,
				self.ff_dropout
			)
			self.ff_layers2 = construct_mlp(
				self.gnn_hidden_feats + self.grover_fp_feats,
				int(self.out_feats/2),
				self.ff_hidden_feats,
				self.ff_num_layers,
				self.ff_dropout
			)

	def forward(self, input_d):

		graph = input_d["graph"]
		node_h = graph.ndata["h"]
		node_h = self.gnn_layers(graph, node_h)
		global_h = self.pool(graph, node_h)

		if not self.separate_heads:
			if self.grover_fp in ['base', 'large']:
				output = self.ff_layers(th.cat([global_h, input_d["grover_fp"]], dim=1))
			else:
				output = self.ff_layers(global_h)

		else:
			if self.grover_fp in ['base', 'large']:
				output1 = self.ff_layers1(th.cat([global_h, input_d["grover_fp"]], dim=1))
				output2 = self.ff_layers2(th.cat([global_h, input_d["grover_fp"]], dim=1))
				output = th.cat([output1, output2], dim=1)
			else:
				output1 = self.ff_layers1(global_h)
				output2 = self.ff_layers2(global_h)
				output = th.cat([output1, output2], dim=1)

		return output

	def get_attention(self, graph, layer_idx=0):

		assert not self.training
		with th.no_grad():
			# run the model forward until reaching layer layer_idx
			graph = graph.local_var()
			feat = graph.ndata["h"]
			skip_gats = self.gnn_layers.gnn_layers[:layer_idx]
			for skip_gat in skip_gats:
				feat = skip_gat(graph,feat)
			# get attention weights of target GAT layer
			targ_gat = self.gnn_layers.gnn_layers[layer_idx].gat_conv
			h_src = h_dst = targ_gat.feat_drop(feat)
			feat_src = feat_dst = targ_gat.fc(h_src).view(
				-1, targ_gat._num_heads, targ_gat._out_feats)
			# NOTE: GAT paper uses "first concatenation then linear projection"
			# to compute attention scores, while ours is "first projection then
			# addition", the two approaches are mathematically equivalent:
			# We decompose the weight vector a mentioned in the paper into
			# [a_l || a_r], then
			# a^T [Wh_i || Wh_j] = a_l Wh_i + a_r Wh_j
			# Our implementation is much efficient because we do not need to
			# save [Wh_i || Wh_j] on edges, which is not memory-efficient. Plus,
			# addition could be optimized with DGL's built-in function u_add_v,
			# which further speeds up computation and saves memory footprint.
			el = (feat_src * targ_gat.attn_l).sum(dim=-1).unsqueeze(-1)
			er = (feat_dst * targ_gat.attn_r).sum(dim=-1).unsqueeze(-1)
			graph.srcdata.update({'ft': feat_src, 'el': el})
			graph.dstdata.update({'er': er})
			# compute edge attention, el and er are a_l Wh_i and a_r Wh_j respectively.
			graph.apply_edges(fn.u_add_v('el', 'er', 'e'))
			e = targ_gat.leaky_relu(graph.edata.pop('e'))
			# compute softmax
			e_attn = targ_gat.attn_drop(edge_softmax(graph, e))
		# average attention over bidirectional edges (not vectorized)
		attn_dict = {}
		e_attn = e_attn.cpu()
		e_tups = zip(graph.edges()[0],graph.edges()[1])
		for e_idx, e_tup in enumerate(e_tups):
			if e_tup[0] <= e_tup[1]:
				e_tup = (int(e_tup[0]),int(e_tup[1]))
			else:
				e_tup = (int(e_tup[1]),int(e_tup[0]))
			if e_tup in attn_dict:
				attn_dict[e_tup] += 0.5*e_attn[e_idx].squeeze(1)
			else:
				attn_dict[e_tup] = 0.5*e_attn[e_idx].squeeze(1)
		return attn_dict

class GIN1(nn.Module):
	""" GIN model """

	def __init__(self,n_feats,e_feats,**kwargs):

		super(GIN1,self).__init__()
		self.n_feats = n_feats
		self.e_feats = e_feats
		for k,v in kwargs.items():
			setattr(self,k,v)
		self.gnn_layers = GIN(
			num_node_emb_list=self.n_feats,
			num_edge_emb_list=self.e_feats,
			num_layers=self.gnn_num_layers,
			emb_dim=self.gnn_hidden_feats,
			dropout=self.gnn_dropout
		)
		self.pool = construct_pool(self.gnn_hidden_feats,self.gnn_pool_type)
		self.ff_layers = construct_mlp(self.gnn_hidden_feats,self.out_feats,self.ff_hidden_feats,self.ff_num_layers,self.ff_dropout)

	def forward(self,input_d):

		graph = input_d["graph"]
		ndata_tensors = [ncol.data for ncol in graph.ndata.values()]
		edata_tensors = [ecol.data for ecol in graph.edata.values()]
		# import pdb; pdb.set_trace()
		node_h = self.gnn_layers(graph,ndata_tensors,edata_tensors)
		global_h = self.pool(graph,node_h)
		output = self.ff_layers(global_h)
		return output

class WLN1(nn.Module):
	""" WLN model """

	def __init__(self,n_feats,e_feats,**kwargs):

		super(WLN1,self).__init__()
		self.n_feats = n_feats
		self.e_feats = e_feats
		for k,v in kwargs.items():
			setattr(self,k,v)
		self.gnn_layers = WLN(
			self.n_feats,
			self.e_feats,
			self.gnn_hidden_feats,
			self.gnn_num_layers
		)
		self.pool = construct_pool(self.gnn_hidden_feats,self.gnn_pool_type)
		self.ff_layers = construct_mlp(self.gnn_hidden_feats,self.out_feats,self.ff_hidden_feats,self.ff_num_layers,self.ff_dropout)

	def forward(self,input_d):

		graph = input_d["graph"]
		node_h = graph.ndata["h"]
		edge_h = graph.edata["e"]
		# import pdb; pdb.set_trace()
		node_h = self.gnn_layers(graph,node_h,edge_h)
		global_h = self.pool(graph,node_h)
		output = self.ff_layers(global_h)
		return output

class GINPT(nn.Module):
	""" gin_supervised_contextpred_Tox21 """

	def __init__(self, n_feats, e_feats, pretrain_name, grover_fp_feats, **kwargs):

		super(GINPT,self).__init__()
		self.grover_fp_feats = grover_fp_feats # can be 0
		self.n_feats = n_feats
		self.e_feats = e_feats
		self.pretrain_name = pretrain_name
		for k,v in kwargs.items():
			setattr(self,k,v)
		assert len(n_feats) == 2 and n_feats[0] == 120 and n_feats[1] == 3, n_feats
		assert len(e_feats) == 2 and e_feats[0] == 6 and e_feats[1] == 3, e_feats
		self.gnn_layers = load_pretrained(self.pretrain_name)
		# remove the output layers
		if hasattr(self.gnn_layers,"readout"):
			self.gnn_layers.readout = GNIdentity()
		if hasattr(self.gnn_layers,"predict"):
			self.gnn_layers.predict = Identity()
		# fix pretrained weights
		if self.fix_pt_weights:
			# note: this is not the same as putting it in eval() mode
			# batchnorm/dropout still works as it does in train() mode
			for param in self.gnn_layers.parameters():
				param.requires_grad = False
		if self.pretrain_name in ["gin_supervised_contextpred_Tox21", "gin_supervised_infomax_Tox21"]:
			self.gnn_out_dim = 1800
		elif self.pretrain_name in ["gin_supervised_infomax_ToxCast", "gin_supervised_infomax"]:
			self.gnn_out_dim = 300
		else:
			raise NotImplementedError
		self.pool = construct_pool(self.gnn_out_dim,self.gnn_pool_type)
		if not self.separate_heads:
			self.ff_layers = construct_mlp(
				self.gnn_out_dim + self.grover_fp_feats,
				self.out_feats,
				self.ff_hidden_feats,
				self.ff_num_layers,
				self.ff_dropout
			)
		else:
			assert self.out_feats % 2 == 0
			self.ff_layers1 = construct_mlp(
				self.gnn_out_dim + self.grover_fp_feats,
				int(self.out_feats / 2),
				self.ff_hidden_feats,
				self.ff_num_layers,
				self.ff_dropout
			)
			self.ff_layers2 = construct_mlp(
				self.gnn_out_dim + self.grover_fp_feats,
				int(self.out_feats/2),
				self.ff_hidden_feats,
				self.ff_num_layers,
				self.ff_dropout
			)

	def forward(self,input_d):
		""" assumes graph is using pretrain featurization """

		graph = input_d["graph"]
		ndata_tensors = [ncol.data for ncol in graph.ndata.values()]
		edata_tensors = [ecol.data for ecol in graph.edata.values()]
		node_h = self.gnn_layers(graph, ndata_tensors, edata_tensors)
		global_h = self.pool(graph, node_h)

		if not self.separate_heads:
			if self.grover_fp in ['base', 'large']:
				output = self.ff_layers(th.cat([global_h, input_d["grover_fp"]], dim=1))
			else:
				output = self.ff_layers(global_h)

		else:
			if self.grover_fp in ['base', 'large']:
				output1 = self.ff_layers1(th.cat([global_h, input_d["grover_fp"]], dim=1))
				output2 = self.ff_layers2(th.cat([global_h, input_d["grover_fp"]], dim=1))
				output = th.cat([output1, output2], dim=1)
			else:
				output1 = self.ff_layers1(global_h)
				output2 = self.ff_layers2(global_h)
				output = th.cat([output1, output2], dim=1)

		return output

class GROVERPT(nn.Module):

	def __init__(self,grover_fp_feats,**kwargs):
		""" just an MLP on the grover embedding """

		super(GROVERPT,self).__init__()
		self.grover_fp_feats = grover_fp_feats
		assert self.grover_fp_feats > 0, self.grover_fp_feats
		for k,v in kwargs.items():
			setattr(self,k,v)
		self.ff_layers = construct_mlp(self.grover_fp_feats,self.out_feats,self.ff_hidden_feats,self.ff_num_layers,self.ff_dropout)

	def forward(self,input_d):
		""" ignores graph """

		grover_fp = input_d["grover_fp"]
		output = self.ff_layers(grover_fp)
		return output

# class GINPT1(nn.Module):
# 	""" gin_supervised_contextpred_Tox21 """

# 	def __init__(self,n_feats,e_feats,**kwargs):

# 		super(GINPT1,self).__init__()
# 		self.n_feats = n_feats
# 		self.e_feats = e_feats
# 		for k,v in kwargs.items():
# 			setattr(self,k,v)
# 		assert len(n_feats) == 2 and n_feats[0] == 120 and n_feats[1] == 3, n_feats
# 		assert len(e_feats) == 2 and e_feats[0] == 6 and e_feats[1] == 3, e_feats
# 		self.gnn_layers = load_pretrained("gin_supervised_contextpred_Tox21")
# 		# remove the output layers
# 		self.gnn_layers.readout = GNIdentity()
# 		self.gnn_layers.predict = Identity()
# 		# fix pretrained weights
# 		if self.fix_pt_weights:
# 			# note: this is not the same as putting it in eval() mode
# 			# batchnorm/dropout still works as it does in train() mode
# 			for param in self.gnn_layers.parameters():
# 				param.requires_grad = False
# 		self.gnn_out_dim = 1800
# 		self.pool = construct_pool(self.gnn_out_dim,self.gnn_pool_type)
# 		self.ff_layers = construct_mlp(self.gnn_out_dim,self.out_feats,self.ff_hidden_feats,self.ff_num_layers,self.ff_dropout)

# 	def forward(self,graph):
# 		""" assumes graph is using pretrain featurization """

# 		ndata_tensors = [ncol.data for ncol in graph.ndata.values()]
# 		edata_tensors = [ecol.data for ecol in graph.edata.values()]
# 		# import pdb; pdb.set_trace()
# 		node_h = self.gnn_layers(graph,ndata_tensors,edata_tensors)
# 		global_h = self.pool(graph,node_h)
# 		output = self.ff_layers(global_h)
# 		return output

class Conv1dBlock(nn.Module):

	def __init__(self,in_channels,out_channels,kernel_size=5,stride=1,pool_type="avg",pool_size=4,batch_norm=True,dropout=0.1,activation="relu"):
		"""
		padding is always the same
		order of ops:
		1. conv
		2. batch norm (if applicable)
		3. activation
		4. pool (if applicable)
		5. dropout (if applicable)
		"""

		super(Conv1dBlock,self).__init__()
		self.in_channels = in_channels
		self.out_channels = out_channels
		ops_list = []
		assert kernel_size % 2 == 1, kernel_size
		padding = kernel_size // 2
		ops_list.append(nn.Conv1d(in_channels,out_channels,kernel_size,stride=stride,padding=padding))
		if batch_norm:
			ops_list.append(nn.BatchNorm1d(out_channels))
		if activation == "relu":
			ops_list.append(nn.ReLU())
		else:
			raise ValueError
		if pool_type == "max":
			ops_list.append(nn.MaxPool1d(pool_size,ceil_mode=True))
		elif pool_type == "avg":
			ops_list.append(nn.AvgPool1d(pool_size,ceil_mode=True))
		else:
			assert pool_type == "none", pool_type
		if dropout > 0.:
			ops_list.append(nn.Dropout(dropout))
		self.ops = nn.Sequential(*ops_list)

	def forward(self,b_x):

		assert b_x.ndim == 3, b_x.ndim
		assert b_x.shape[1] == self.in_channels, b_x.shape
		b_y = self.ops(b_x)
		return b_y

class CNN(nn.Module):

	def __init__(self,c_feats,l_feats,**kwargs):

		super(CNN,self).__init__()
		self.c_feats = c_feats
		self.l_feats = l_feats
		for k,v in kwargs.items():
			setattr(self,k,v)

		conv_hp = dict(
			kernel_size=self.conv_kernel_size,
			stride=self.conv_stride,
			pool_type=self.conv_pool_type,
			pool_size=self.conv_pool_size,
			dropout=self.ff_dropout,
			activation="relu"
		)

		def compute_conv_out_length(in_length,kernel_size,stride):
			padding = kernel_size // 2
			dilation = 1
			out_length = np.floor((in_length + 2*padding - dilation * (kernel_size-1) - 1) / stride + 1)
			return int(out_length)

		def compute_pool_out_length(in_length,kernel_size):
			if self.conv_pool_type == "none":
				return in_length
			else:
				padding = 0
				dilation = 1
				stride = kernel_size
				out_length = np.ceil((in_length + 2*padding - dilation * (kernel_size-1) - 1) / stride + 1)
				return int(out_length)

		conv_list = []
		conv_list.append(Conv1dBlock(self.c_feats,self.conv_num_kernels,**conv_hp))
		conv_out_length = compute_conv_out_length(self.l_feats,self.conv_kernel_size,self.conv_stride)
		pool_out_length = compute_pool_out_length(conv_out_length,self.conv_pool_size)
		# print(0,conv_out_length,pool_out_length)
		for i in range(self.conv_num_layers-1):
			conv_list.append(Conv1dBlock(self.conv_num_kernels,self.conv_num_kernels,**conv_hp))
			conv_out_length = compute_conv_out_length(pool_out_length,self.conv_kernel_size,self.conv_stride)
			pool_out_length = compute_pool_out_length(conv_out_length,self.conv_pool_size)
			# print(i+1,conv_out_length,pool_out_length)
		self.conv_layers = nn.Sequential(*conv_list)

		self.ff_layers = construct_mlp(pool_out_length*self.conv_num_kernels,self.out_feats,self.ff_hidden_feats,self.ff_num_layers,self.ff_dropout)

	def forward(self, input_d):

		seq = input_d["seq"]
		# assumes seq is padded
		batch_size = seq.shape[0]
		ch = self.conv_layers(seq)
		ch = ch.view(batch_size,-1)
		fh = self.ff_layers(ch)
		return fh

def construct_mlp(in_dim, out_dim, hidden_dim, hidden_layers, dropout, different_heads=False):
	"""
	total number of layers = hidden_layers + 1 (the + 1 is for the output linear)
	no activation/batch norm/dropout on output layer
	"""

	assert hidden_layers >= 1, hidden_layers
	mlp_list = []
	mlp_list.append(LinearBlock(in_dim,hidden_dim,dropout=dropout))
	for i in range(hidden_layers-1):
		mlp_list.append(LinearBlock(hidden_dim,hidden_dim,dropout=dropout))
	mlp_list.append(nn.Linear(hidden_dim,out_dim))
	mlp = nn.Sequential(*mlp_list)
	return mlp

def construct_pool(node_dim,pool_type):

	if pool_type == "max":
		return MaxPooling()
	elif pool_type == "avg":
		return AvgPooling()
	elif pool_type == "attn":
		pool_gate_nn = nn.Linear(node_dim, 1)
		return GlobalAttentionPooling(pool_gate_nn)

def test_forward_backward(model_type):

	from data_loader import get_datasets
	from arguments import get_args
	args = get_args()
	print(model_type)
	args["model_type"] = model_type
	if model_type == "grover_pt":
		args["grover_fp"] = "base"
	ds = get_datasets(args)[0]
	data_feats_d = ds.get_data_feats()
	model = get_model(args, data_feats_d)
	collate_fn = ds.get_collate_fn()
	d0 = ds[0]
	d1 = ds[1]
	d2 = ds[2]
	bd = collate_fn([d0, d1, d2])
	# print(g)
	o = model(bd)
	print(o)
	assert pd.Series(bd["mut_label"]).isnull().sum() == 0
	l = th.mean((bd["mut_label"]-o[:, 0])**2)
	print(l)
	l.backward()
	import ipdb;
	ipdb.set_trace()

def test_attention_viz():

	from data_loader import get_datasets
	from arguments import get_args
	from plot_utils import viz_attention
	args = get_args()
	args["model_type"] = "gat"
	ds = get_datasets(args)[0]
	data_feats_d = ds.get_data_feats()
	model = get_model(args, data_feats_d)
	collate_fn = ds.get_collate_fn()
	d0 = ds[0]
	model.eval()
	attn_dict = model.get_attention(d0["graph"])
	print(attn_dict)
	im = viz_attention(smiles,attn_dict)
	im.save("test.png")

if __name__ == "__main__":

	# test_attention_viz()
	test_forward_backward("mlp_fingerprint")
	test_forward_backward("grover_pt")
	test_forward_backward("gat")
	# test_forward_backward("gin_pt1")
	test_forward_backward("gin_pt2")
	# test_forward_backward("gin_pt3")
	# test_forward_backward("gin_pt4")
	test_forward_backward("cnn")
	test_forward_backward("wln")

max_stars_repo_name: bowang-lab/CONCERTO

max_issues_repo_name: bowang-lab/CONCERTO

max_forks_repo_name: bowang-lab/CONCERTO


text:
# Databricks notebook source
# MAGIC %md
# MAGIC ScaDaMaLe Course [site](https://lamastex.github.io/scalable-data-science/sds/3/x/) and [book](https://lamastex.github.io/ScaDaMaLe/index.html)
# MAGIC
# MAGIC This is a 2019-2021 augmentation and update of [Adam Breindel](https://www.linkedin.com/in/adbreind)'s initial notebooks.
# MAGIC
# MAGIC _Thanks to [Christian von Koch](https://www.linkedin.com/in/christianvonkoch/) and [William Anzén](https://www.linkedin.com/in/william-anz%C3%A9n-b52003199/) for their contributions towards making these materials Spark 3.0.1 and Python 3+ compliant._

# COMMAND ----------

# MAGIC %md
# MAGIC # Convolutional Neural Networks
# MAGIC ## aka CNN, ConvNet

# COMMAND ----------

# MAGIC %md
# MAGIC As a baseline, let's start a lab running with what we already know.
# MAGIC
# MAGIC We'll take our deep feed-forward multilayer perceptron network, with ReLU activations and reasonable initializations, and apply it to learning the MNIST digits.
# MAGIC
# MAGIC The main part of the code looks like the following (full code you can run is in the next cell):
# MAGIC
# MAGIC ```
# MAGIC # imports, setup, load data sets
# MAGIC
# MAGIC model = Sequential()
# MAGIC model.add(Dense(20, input_dim=784, kernel_initializer='normal', activation='relu'))
# MAGIC model.add(Dense(15, kernel_initializer='normal', activation='relu'))
# MAGIC model.add(Dense(10, kernel_initializer='normal', activation='softmax'))
# MAGIC model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['categorical_accuracy'])
# MAGIC
# MAGIC categorical_labels = to_categorical(y_train, num_classes=10)
# MAGIC
# MAGIC history = model.fit(X_train, categorical_labels, epochs=100, batch_size=100)
# MAGIC
# MAGIC # print metrics, plot errors
# MAGIC ```
# MAGIC
# MAGIC Note the changes, which are largely about building a classifier instead of a regression model:
# MAGIC * Output layer has one neuron per category, with softmax activation
# MAGIC * __Loss function is cross-entropy loss__
# MAGIC * Accuracy metric is categorical accuracy

# COMMAND ----------

# MAGIC %md
# MAGIC Let's hold pointers into wikipedia for these new concepts.

# COMMAND ----------

# MAGIC %scala
# MAGIC //This allows easy embedding of publicly available information into any other notebook
# MAGIC //Example usage:
# MAGIC // displayHTML(frameIt("https://en.wikipedia.org/wiki/Latent_Dirichlet_allocation#Topics_in_LDA",250))
# MAGIC def frameIt( u:String, h:Int ) : String = {
# MAGIC       """
# MAGIC

# MAGIC
# MAGIC       Fallback link for browsers that, unlikely, don't support frames
# MAGIC
# MAGIC

# MAGIC """
# MAGIC    }
# MAGIC displayHTML(frameIt("https://en.wikipedia.org/wiki/Cross_entropy#Cross-entropy_error_function_and_logistic_regression",500))

# COMMAND ----------

# MAGIC %scala
# MAGIC displayHTML(frameIt("https://en.wikipedia.org/wiki/Softmax_function",380))

# COMMAND ----------

# MAGIC %md
# MAGIC The following is from: [https://www.quora.com/How-does-Keras-calculate-accuracy](https://www.quora.com/How-does-Keras-calculate-accuracy).
# MAGIC
# MAGIC **Categorical accuracy:**
# MAGIC
# MAGIC ```%python
# MAGIC def categorical_accuracy(y_true, y_pred):
# MAGIC  return K.cast(K.equal(K.argmax(y_true, axis=-1),
# MAGIC  K.argmax(y_pred, axis=-1)),
# MAGIC  K.floatx())
# MAGIC  ```
# MAGIC
# MAGIC > `K.argmax(y_true)` takes the highest value to be the prediction and matches against the comparative set.

# COMMAND ----------

# MAGIC %md
# MAGIC Watch (1:39)
# MAGIC * [![Udacity: Deep Learning by Vincent Vanhoucke - Cross-entropy](http://img.youtube.com/vi/tRsSi_sqXjI/0.jpg)](https://www.youtube.com/watch?v=tRsSi_sqXjI)
# MAGIC
# MAGIC Watch (1:54)
# MAGIC * [![Udacity: Deep Learning by Vincent Vanhoucke - Minimizing Cross-entropy](http://img.youtube.com/vi/x449QQDhMDE/0.jpg)](https://www.youtube.com/watch?v=x449QQDhMDE)

# COMMAND ----------

from keras.models import Sequential
from keras.layers import Dense
from keras.utils import to_categorical
import sklearn.datasets
import datetime
import matplotlib.pyplot as plt
import numpy as np

train_libsvm = "/dbfs/databricks-datasets/mnist-digits/data-001/mnist-digits-train.txt"
test_libsvm = "/dbfs/databricks-datasets/mnist-digits/data-001/mnist-digits-test.txt"

X_train, y_train = sklearn.datasets.load_svmlight_file(train_libsvm, n_features=784)
X_train = X_train.toarray()

X_test, y_test = sklearn.datasets.load_svmlight_file(test_libsvm, n_features=784)
X_test = X_test.toarray()

model = Sequential()
model.add(Dense(20, input_dim=784, kernel_initializer='normal', activation='relu'))
model.add(Dense(15, kernel_initializer='normal', activation='relu'))
model.add(Dense(10, kernel_initializer='normal', activation='softmax'))
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['categorical_accuracy'])

categorical_labels = to_categorical(y_train, num_classes=10)
start = datetime.datetime.today()

history = model.fit(X_train, categorical_labels, epochs=40, batch_size=100, validation_split=0.1, verbose=2)

scores = model.evaluate(X_test, to_categorical(y_test, num_classes=10))

print
for i in range(len(model.metrics_names)):
	print("%s: %f" % (model.metrics_names[i], scores[i]))

print ("Start: " + str(start))
end = datetime.datetime.today()
print ("End: " + str(end))
print ("Elapse: " + str(end-start))

# COMMAND ----------

# MAGIC %md
# MAGIC after about a minute we have:
# MAGIC
# MAGIC ```
# MAGIC ...
# MAGIC
# MAGIC Epoch 40/40
# MAGIC 1s - loss: 0.0610 - categorical_accuracy: 0.9809 - val_loss: 0.1918 - val_categorical_accuracy: 0.9583
# MAGIC
# MAGIC ...
# MAGIC
# MAGIC loss: 0.216120
# MAGIC
# MAGIC categorical_accuracy: 0.955000
# MAGIC
# MAGIC Start: 2017-12-06 07:35:33.948102
# MAGIC
# MAGIC End: 2017-12-06 07:36:27.046130
# MAGIC
# MAGIC Elapse: 0:00:53.098028
# MAGIC ```

# COMMAND ----------

import matplotlib.pyplot as plt

fig, ax = plt.subplots()
fig.set_size_inches((5,5))
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'val'], loc='upper left')
display(fig)

# COMMAND ----------

# MAGIC %md
# MAGIC What are the big takeaways from this experiment?
# MAGIC
# MAGIC 1. We get pretty impressive "apparent error" accuracy right from the start! A small network gets us to training accuracy 97% by epoch 20
# MAGIC 2. The model *appears* to continue to learn if we let it run, although it does slow down and oscillate a bit.
# MAGIC 3. Our test accuracy is about 95% after 5 epochs and never gets better ... it gets worse!
# MAGIC 4. Therefore, we are overfitting very quickly... most of the "training" turns out to be a waste.
# MAGIC 5. For what it's worth, we get 95% accuracy without much work.
# MAGIC
# MAGIC This is not terrible compared to other, non-neural-network approaches to the problem. After all, we could probably tweak this a bit and do even better.
# MAGIC
# MAGIC But we talked about using deep learning to solve "95%" problems or "98%" problems ... where one error in 20, or 50 simply won't work. If we can get to "multiple nines" of accuracy, then we can do things like automate mail sorting and translation, create cars that react properly (all the time) to street signs, and control systems for robots or drones that function autonomously.
# MAGIC
# MAGIC Try two more experiments (try them separately):
# MAGIC 1. Add a third, hidden layer.
# MAGIC 2. Increase the size of the hidden layers.
# MAGIC
# MAGIC Adding another layer slows things down a little (why?) but doesn't seem to make a difference in accuracy.
# MAGIC
# MAGIC Adding a lot more neurons into the first topology slows things down significantly -- 10x as many neurons, and only a marginal increase in accuracy. Notice also (in the plot) that the learning clearly degrades after epoch 50 or so.
# MAGIC
# MAGIC ... We need a new approach!
# MAGIC
# MAGIC ---
# MAGIC
# MAGIC ... let's think about this:
# MAGIC
# MAGIC ### What is layer 2 learning from layer 1? Combinations of pixels
# MAGIC
# MAGIC #### Combinations of pixels contain information but...
# MAGIC
# MAGIC There are a lot of them (combinations) and they are "fragile"
# MAGIC
# MAGIC In fact, in our last experiment, we basically built a model that memorizes a bunch of "magic" pixel combinations.
# MAGIC
# MAGIC What might be a better way to build features?
# MAGIC
# MAGIC * When humans perform this task, we look not at arbitrary pixel combinations, but certain geometric patterns -- lines, curves, loops.
# MAGIC * These features are made up of combinations of pixels, but they are far from arbitrary
# MAGIC * We identify these features regardless of translation, rotation, etc.
# MAGIC
# MAGIC Is there a way to get the network to do the same thing?
# MAGIC
# MAGIC I.e., in layer one, identify pixels. Then in layer 2+, identify abstractions over pixels that are translation-invariant 2-D shapes?
# MAGIC
# MAGIC We could look at where a "filter" that represents one of these features (e.g., and edge) matches the image.
# MAGIC
# MAGIC How would this work?
# MAGIC
# MAGIC ### Convolution
# MAGIC
# MAGIC Convolution in the general mathematical sense is define as follows:
# MAGIC
# MAGIC
# MAGIC
# MAGIC The convolution we deal with in deep learning is a simplified case. We want to compare two signals. Here are two visualizations, courtesy of Wikipedia, that help communicate how convolution emphasizes features:
# MAGIC
# MAGIC
# MAGIC
# MAGIC ---
# MAGIC
# MAGIC #### Here's an animation (where we change \\({\tau}\\))
# MAGIC
# MAGIC
# MAGIC __In one sense, the convolution captures and quantifies the pattern matching over space__
# MAGIC
# MAGIC If we perform this in two dimensions, we can achieve effects like highlighting edges:
# MAGIC
# MAGIC
# MAGIC
# MAGIC The matrix here, also called a convolution kernel, is one of the functions we are convolving. Other convolution kernels can blur, "sharpen," etc.
# MAGIC
# MAGIC ### So we'll drop in a number of convolution kernels, and the network will learn where to use them? Nope. Better than that.
# MAGIC
# MAGIC ## We'll program in the *idea* of discrete convolution, and the network will learn what kernels extract meaningful features!
# MAGIC
# MAGIC The values in a (fixed-size) convolution kernel matrix will be variables in our deep learning model. Although inuitively it seems like it would be hard to learn useful params, in fact, since those variables are used repeatedly across the image data, it "focuses" the error on a smallish number of parameters with a lot of influence -- so it should be vastly *less* expensive to train than just a huge fully connected layer like we discussed above.
# MAGIC
# MAGIC This idea was developed in the late 1980s, and by 1989, Yann LeCun (at AT&T/Bell Labs) had built a practical high-accuracy system (used in the 1990s for processing handwritten checks and mail).
# MAGIC
# MAGIC __How do we hook this into our neural networks?__
# MAGIC
# MAGIC * First, we can preserve the geometric properties of our data by "shaping" the vectors as 2D instead of 1D.
# MAGIC
# MAGIC * Then we'll create a layer whose value is not just activation applied to weighted sum of inputs, but instead it's the result of a dot-product (element-wise multiply and sum) between the kernel and a patch of the input vector (image).
# MAGIC     * This value will be our "pre-activation" and optionally feed into an activation function (or "detector")
# MAGIC
# MAGIC
# MAGIC
# MAGIC
# MAGIC If we perform this operation at lots of positions over the image, we'll get lots of outputs, as many as one for every input pixel.
# MAGIC
# MAGIC
# MAGIC
# MAGIC
# MAGIC * So we'll add another layer that "picks" the highest convolution pattern match from nearby pixels, which
# MAGIC     * makes our pattern match a little bit translation invariant (a fuzzy location match)
# MAGIC     * reduces the number of outputs significantly
# MAGIC * This layer is commonly called a pooling layer, and if we pick the "maximum match" then it's a "max pooling" layer.
# MAGIC
# MAGIC
# MAGIC
# MAGIC __The end result is that the kernel or filter together with max pooling creates a value in a subsequent layer which represents the appearance of a pattern in a local area in a prior layer.__
# MAGIC
# MAGIC __Again, the network will be given a number of "slots" for these filters and will learn (by minimizing error) what filter values produce meaningful features. This is the key insight into how modern image-recognition networks are able to generalize -- i.e., learn to tell 6s from 7s or cats from dogs.__
# MAGIC
# MAGIC
# MAGIC
# MAGIC ## Ok, let's build our first ConvNet:
# MAGIC
# MAGIC First, we want to explicity shape our data into a 2-D configuration. We'll end up with a 4-D tensor where the first dimension is the training examples, then each example is 28x28 pixels, and we'll explicitly say it's 1-layer deep. (Why? with color images, we typically process over 3 or 4 channels in this last dimension)
# MAGIC
# MAGIC A step by step animation follows:
# MAGIC * http://cs231n.github.io/assets/conv-demo/index.html

# COMMAND ----------

train_libsvm = "/dbfs/databricks-datasets/mnist-digits/data-001/mnist-digits-train.txt"
test_libsvm = "/dbfs/databricks-datasets/mnist-digits/data-001/mnist-digits-test.txt"

X_train, y_train = sklearn.datasets.load_svmlight_file(train_libsvm, n_features=784)
X_train = X_train.toarray()

X_test, y_test = sklearn.datasets.load_svmlight_file(test_libsvm, n_features=784)
X_test = X_test.toarray()

X_train = X_train.reshape( (X_train.shape[0], 28, 28, 1) )
X_train = X_train.astype('float32')
X_train /= 255
y_train = to_categorical(y_train, num_classes=10)

X_test = X_test.reshape( (X_test.shape[0], 28, 28, 1) )
X_test = X_test.astype('float32')
X_test /= 255
y_test = to_categorical(y_test, num_classes=10)

# COMMAND ----------

# MAGIC %md
# MAGIC Now the model:

# COMMAND ----------

from keras.layers import Dense, Dropout, Activation, Flatten, Conv2D, MaxPooling2D

model = Sequential()

model.add(Conv2D(8, # number of kernels
				(4, 4), # kernel size
                padding='valid', # no padding; output will be smaller than input
                input_shape=(28, 28, 1)))

model.add(Activation('relu'))

model.add(MaxPooling2D(pool_size=(2,2)))

model.add(Flatten())
model.add(Dense(128))
model.add(Activation('relu')) # alternative syntax for applying activation

model.add(Dense(10))
model.add(Activation('softmax'))

model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# COMMAND ----------

# MAGIC %md
# MAGIC ... and the training loop and output:

# COMMAND ----------

start = datetime.datetime.today()

history = model.fit(X_train, y_train, batch_size=128, epochs=8, verbose=2, validation_split=0.1)

scores = model.evaluate(X_test, y_test, verbose=1)

print
for i in range(len(model.metrics_names)):
	print("%s: %f" % (model.metrics_names[i], scores[i]))

# COMMAND ----------

fig, ax = plt.subplots()
fig.set_size_inches((5,5))
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'val'], loc='upper left')
display(fig)

# COMMAND ----------

# MAGIC %md
# MAGIC ### Our MNIST ConvNet
# MAGIC
# MAGIC In our first convolutional MNIST experiment, we get to almost 99% validation accuracy in just a few epochs (a minutes or so on CPU)!
# MAGIC
# MAGIC The training accuracy is effectively 100%, though, so we've almost completely overfit (i.e., memorized the training data) by this point and need to do a little work if we want to keep learning.
# MAGIC
# MAGIC Let's add another convolutional layer:

# COMMAND ----------

model = Sequential()

model.add(Conv2D(8, # number of kernels
						(4, 4), # kernel size
                        padding='valid',
                        input_shape=(28, 28, 1)))

model.add(Activation('relu'))

model.add(Conv2D(8, (4, 4)))
model.add(Activation('relu'))

model.add(MaxPooling2D(pool_size=(2,2)))

model.add(Flatten())
model.add(Dense(128))
model.add(Activation('relu'))

model.add(Dense(10))
model.add(Activation('softmax'))

model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

history = model.fit(X_train, y_train, batch_size=128, epochs=15, verbose=2, validation_split=0.1)

scores = model.evaluate(X_test, y_test, verbose=1)

print
for i in range(len(model.metrics_names)):
	print("%s: %f" % (model.metrics_names[i], scores[i]))

# COMMAND ----------

# MAGIC %md
# MAGIC While that's running, let's look at a number of "famous" convolutional networks!
# MAGIC
# MAGIC ### LeNet (Yann LeCun, 1998)
# MAGIC
# MAGIC
# MAGIC
# MAGIC

# COMMAND ----------

# MAGIC %md

# COMMAND ----------

# MAGIC %md
# MAGIC ### AlexNet (2012)
# MAGIC
# MAGIC
# MAGIC
# MAGIC

# COMMAND ----------

# MAGIC %md
# MAGIC ### Back to our labs: Still Overfitting
# MAGIC
# MAGIC We're making progress on our test error -- about 99% -- but just a bit for all the additional time, due to the network overfitting the data.
# MAGIC
# MAGIC There are a variety of techniques we can take to counter this -- forms of regularization.
# MAGIC
# MAGIC Let's try a relatively simple solution solution that works surprisingly well: add a pair of `Dropout` filters, a layer that randomly omits a fraction of neurons from each training batch (thus exposing each neuron to only part of the training data).
# MAGIC
# MAGIC We'll add more convolution kernels but shrink them to 3x3 as well.

# COMMAND ----------

model = Sequential()

model.add(Conv2D(32, # number of kernels
						(3, 3), # kernel size
                        padding='valid',
                        input_shape=(28, 28, 1)))

model.add(Activation('relu'))

model.add(Conv2D(32, (3, 3)))
model.add(Activation('relu'))

model.add(MaxPooling2D(pool_size=(2,2)))

model.add(Dropout(rate=1-0.25)) # <- regularize, new parameter rate added (rate=1-keep_prob)
model.add(Flatten())
model.add(Dense(128))
model.add(Activation('relu'))

model.add(Dropout(rate=1-0.5)) # <-regularize, new parameter rate added (rate=1-keep_prob)
model.add(Dense(10))
model.add(Activation('softmax'))

model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
history = model.fit(X_train, y_train, batch_size=128, epochs=15, verbose=2)

scores = model.evaluate(X_test, y_test, verbose=2)

print
for i in range(len(model.metrics_names)):
	print("%s: %f" % (model.metrics_names[i], scores[i]))

# COMMAND ----------

# MAGIC %md
# MAGIC While that's running, let's look at some more recent ConvNet architectures:
# MAGIC
# MAGIC ### VGG16 (2014)
# MAGIC
# MAGIC

# COMMAND ----------

# MAGIC %md
# MAGIC ### GoogLeNet (2014)
# MAGIC
# MAGIC
# MAGIC
# MAGIC *"Inception" layer: parallel convolutions at different resolutions*
# MAGIC
# MAGIC ### Residual Networks (2015-)
# MAGIC
# MAGIC Skip layers to improve training (error propagation). Residual layers learn from details at multiple previous layers.
# MAGIC
# MAGIC

# COMMAND ----------

# MAGIC %md
# MAGIC ---
# MAGIC
# MAGIC > __ASIDE: Atrous / Dilated Convolutions__
# MAGIC
# MAGIC > An atrous or dilated convolution is a convolution filter with "holes" in it. Effectively, it is a way to enlarge the filter spatially while not adding as many parameters or attending to every element in the input.
# MAGIC
# MAGIC > Why? Covering a larger input volume allows recognizing coarser-grained patterns; restricting the number of parameters is a way of regularizing or constraining the capacity of the model, making training easier.
# MAGIC
# MAGIC ---

# COMMAND ----------

# MAGIC %md
# MAGIC ## *Lab Wrapup*
# MAGIC
# MAGIC From the last lab, you should have a test accuracy of over 99.1%
# MAGIC
# MAGIC For one more activity, try changing the optimizer to old-school "sgd" -- just to see how far we've come with these modern gradient descent techniques in the last few years.
# MAGIC
# MAGIC Accuracy will end up noticeably worse ... about 96-97% test accuracy. Two key takeaways:
# MAGIC
# MAGIC * Without a good optimizer, even a very powerful network design may not achieve results
# MAGIC * In fact, we could replace the word "optimizer" there with
# MAGIC     * initialization
# MAGIC     * activation
# MAGIC     * regularization
# MAGIC     * (etc.)
# MAGIC * All of these elements we've been working with operate together in a complex way to determine final performance

# COMMAND ----------

# MAGIC %md
# MAGIC Of course this world evolves fast - see the new kid in the CNN block -- **capsule networks**
# MAGIC
# MAGIC > Hinton: “The pooling operation used in convolutional neural networks is a big mistake and the fact that it works so well is a disaster.”
# MAGIC
# MAGIC Well worth the 8 minute read:
# MAGIC * [https://medium.com/ai%C2%B3-theory-practice-business/understanding-hintons-capsule-networks-part-i-intuition-b4b559d1159b](https://medium.com/ai%C2%B3-theory-practice-business/understanding-hintons-capsule-networks-part-i-intuition-b4b559d1159b)
# MAGIC
# MAGIC To understand deeper:
# MAGIC * original paper: [https://arxiv.org/abs/1710.09829](https://arxiv.org/abs/1710.09829)
# MAGIC
# MAGIC [Keras capsule network example](https://keras.io/examples/cifar10_cnn_capsule/)

# COMMAND ----------

# MAGIC %md
# MAGIC # More resources
# MAGIC
# MAGIC  - http://www.wildml.com/2015/12/implementing-a-cnn-for-text-classification-in-tensorflow/
# MAGIC  - https://openai.com/

# COMMAND ----------

ext: py

max_stars_repo_name: r-e-x-a-g-o-n/scalable-data-science

max_issues_repo_name: r-e-x-a-g-o-n/scalable-data-science

max_forks_repo_name: r-e-x-a-g-o-n/scalable-data-science

reason: import numpy


text:
import numpy as np
import random

from boxenv import *
from agent import *

COND = 'GC'
STATE_DIM = 2
DIM = STATE_DIM

policy_function = GaussianPolicyFunction(2 * STATE_DIM, 2)
policy = GaussianPolicy()

# initial training task list
TASKS = [(0.5, 0.8), (-0.5, 0.8)]
# TASKS = [(1., 0.8), (0.33, 0.8), (-0.33, 0.8), (-1, 0.8)]

# create a stationary test task list
rng = np.random.default_rng(1)
TEST_TASKS = rng.random((10,2))
TEST_TASKS = TEST_TASKS * 2 - 1
TEST_TASKS[:,1] = 0.8

print('Testing zero-shot generalization on tasks ')
print(TEST_TASKS)

def compute_rewards(s1, s2, g):
    """
    input: s1 - state before action
           s2 - state after action
    rewards based on proximity to each goal
    """
    dist1 = np.linalg.norm(s1 - g)
    dist2 = np.linalg.norm(s2 - g)
    r = dist1 - dist2
    return r

rewards = []
for SEED in [123, 456, 789]:
    np.random.seed(SEED)
    random.seed(SEED)
    torch.manual_seed(SEED)

    box = BoxWorld()
    NB_TASKS = TEST_TASKS.shape[0]

    policy_function.load_state_dict(torch.load(
        'models/{}/{}/policy_seed{}.pth'.format(COND, len(TASKS), SEED)))
    rewards.append([])
    for gid in range(NB_TASKS):
        rewards[-1].append([])
        for i in range(10):
            s = box.reset()
            done = False
            states = []
            rewards[-1][gid].append(0)
            while not done:
                states.append(s)
                s = torch.Tensor(np.concatenate((s, TEST_TASKS[gid])))
                # get action and logprobs
                mu, sigma = policy_function(s)
                unscaled_action, logprob, entropy = policy.forward(mu, sigma)
                # scale action to environment limits
                a = box.scale_action(unscaled_action.detach().numpy())
                # step the environment
                s, _, done = box.step(a)
                r = compute_rewards(states[-1], s, TEST_TASKS[gid])
                rewards[-1][gid][-1] += r
rewards = np.stack(rewards)
print(rewards.mean(-1).mean(-1))
print(rewards.mean())
print(np.std(rewards.mean(-1).mean(-1)))

max_stars_repo_name: himanshusahni/task-biased-url

max_issues_repo_name: himanshusahni/task-biased-url

max_forks_repo_name: himanshusahni/task-biased-url


text:
from lega.legendre_basis import mass_matrix, stiffness_matrix
from scipy.linalg import eigh
from math import sqrt
import numpy as np

n = 10

# These are matrices correspoding to L_0=1 and so on. As such the A matrix is
# not positive definite and the space spanned by L_0 is the kernel
M = mass_matrix(n)
A = stiffness_matrix(n)

# So we leave consider only basis L_1 ... and solve a well posed problem
lmbda1, V1 = eigh(A.toarray()[1:, 1:], M.toarray()[1:, 1:])

# And now add the zero into the spectrum along with its eigenvector [1, 0, ....]
lmbda = np.concatenate([np.array([0]), lmbda1])
V = np.eye(n)/sqrt(2.)   # Note that the column/eigenvector is normalized
V[1:, 1:] = V1

# Let check we have not messed up anything. Can expect exact
for a, v in zip(lmbda, V.T):
    assert np.linalg.norm(A.dot(v) - a*M.dot(v))/n < 1E-12

assert np.linalg.norm(V.T.dot(M.dot(V)) - np.eye(n))/n < 1E-12
assert np.linalg.norm(V.T.dot(A.dot(V)) - np.diag(lmbda))/n < 1E-12

# Greate so now we know how to deal with the Neumann problem

max_stars_repo_name: MiroK/lega

max_issues_repo_name: MiroK/lega

max_forks_repo_name: MiroK/lega

reason: import numpy,from scipy


text:
from typing import Mapping, Set, Tuple, Sequence, AbstractSet, Union
from abc import abstractmethod
from numbers import Real

import numpy as np

from qupulse import ChannelID
from qupulse.comparable import Comparable
from qupulse.utils.types import SingletonABCMeta

class Transformation(Comparable):
    _identity_singleton = None
    """Transforms numeric time-voltage values for multiple channels to other time-voltage values. The number and names
     of input and output channels might differ."""

    @abstractmethod
    def __call__(self, time: Union[np.ndarray, float],
                 data: Mapping[ChannelID, Union[np.ndarray, float]]) -> Mapping[ChannelID, Union[np.ndarray, float]]:
        """Apply transformation to data
        Args:
            time:
            data:

        Returns:
            transformed: A DataFrame that has been transformed with index == output_channels
        """

    @abstractmethod
    def get_output_channels(self, input_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        """Return the channel identifiers"""

    @abstractmethod
    def get_input_channels(self, output_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        """Channels that are required for getting data for the requested output channel"""

    def chain(self, next_transformation: 'Transformation') -> 'Transformation':
        if next_transformation is IdentityTransformation():
            return self
        else:
            return chain_transformations(self, next_transformation)

    def is_constant_invariant(self):
        """Signals if the transformation always maps constants to constants."""
        return False

class IdentityTransformation(Transformation, metaclass=SingletonABCMeta):
    def __call__(self, time: Union[np.ndarray, float],
                 data: Mapping[ChannelID, Union[np.ndarray, float]]) -> Mapping[ChannelID, Union[np.ndarray, float]]:
        return data

    def get_output_channels(self, input_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        return input_channels

    @property
    def compare_key(self) -> None:
        return None

    def get_input_channels(self, output_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        return output_channels

    def chain(self, next_transformation: Transformation) -> Transformation:
        return next_transformation

    def __repr__(self):
        return 'IdentityTransformation()'

    def is_constant_invariant(self):
        """Signals if the transformation always maps constants to constants."""
        return True

class ChainedTransformation(Transformation):
    def __init__(self, *transformations: Transformation):
        self._transformations = transformations

    @property
    def transformations(self) -> Tuple[Transformation, ...]:
        return self._transformations

    def get_output_channels(self, input_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        for transformation in self._transformations:
            input_channels = transformation.get_output_channels(input_channels)
        return input_channels

    def get_input_channels(self, output_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        for transformation in reversed(self._transformations):
            output_channels = transformation.get_input_channels(output_channels)
        return output_channels

    def __call__(self, time: Union[np.ndarray, float],
                 data: Mapping[ChannelID, Union[np.ndarray, float]]) -> Mapping[ChannelID, Union[np.ndarray, float]]:
        for transformation in self._transformations:
            data = transformation(time, data)
        return data

    @property
    def compare_key(self) -> Tuple[Transformation, ...]:
        return self._transformations

    def chain(self, next_transformation) -> Transformation:
        return chain_transformations(*self.transformations, next_transformation)

    def __repr__(self):
        return 'ChainedTransformation%r' % (self._transformations,)

    def is_constant_invariant(self):
        """Signals if the transformation always maps constants to constants."""
        return all(trafo.is_constant_invariant() for trafo in self._transformations)

class LinearTransformation(Transformation):
    def __init__(self,
                 transformation_matrix: np.ndarray,
                 input_channels: Sequence[ChannelID],
                 output_channels: Sequence[ChannelID]):
        """

        Args:
            transformation_matrix: Matrix describing the transformation with shape (output_channels, input_channels)
            input_channels: Channel ids of the columns
            output_channels: Channel ids of the rows
        """
        transformation_matrix = np.asarray(transformation_matrix)

        if transformation_matrix.shape != (len(output_channels), len(input_channels)):
            raise ValueError('Shape of transformation matrix does not match to the given channels')

        output_sorter = np.argsort(output_channels)
        transformation_matrix = transformation_matrix[output_sorter, :]

        input_sorter = np.argsort(input_channels)
        transformation_matrix = transformation_matrix[:, input_sorter]

        self._matrix = transformation_matrix
        self._input_channels = tuple(sorted(input_channels))
        self._output_channels = tuple(sorted(output_channels))
        self._input_channels_set = frozenset(self._input_channels)
        self._output_channels_set = frozenset(self._output_channels)

    def __call__(self, time: Union[np.ndarray, float],
                 data: Mapping[ChannelID, Union[np.ndarray, float]]) -> Mapping[ChannelID, Union[np.ndarray, float]]:
        data_out = {forwarded_channel: data[forwarded_channel]
                    for forwarded_channel in set(data.keys()).difference(self._input_channels)}

        if len(data_out) == len(data):
            # only forwarded data
            return data_out

        try:
            data_in = np.stack(data[in_channel] for in_channel in self._input_channels)
        except KeyError as error:
            raise KeyError('Invalid input channels', set(data.keys()), set(self._input_channels)) from error

        transformed_data = self._matrix @ data_in

        for idx, out_channel in enumerate(self._output_channels):
            data_out[out_channel] = transformed_data[idx, ...]

        return data_out

    def get_output_channels(self, input_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        if not input_channels >= self._input_channels_set:
            # input_channels is not a superset of the required input channels
            raise KeyError('Invalid input channels', input_channels, self._input_channels_set)

        return (input_channels - self._input_channels_set) | self._output_channels_set

    def get_input_channels(self, output_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        forwarded = output_channels - self._output_channels_set
        if not forwarded.isdisjoint(self._input_channels):
            raise KeyError('Is input channel', forwarded & self._input_channels_set)
        elif output_channels.isdisjoint(self._output_channels):
            return output_channels
        else:
            return forwarded | self._input_channels_set

    @property
    def compare_key(self) -> Tuple[Tuple[ChannelID], Tuple[ChannelID], bytes]:
        return self._input_channels, self._output_channels, self._matrix.tobytes()

    def __repr__(self):
        return ('LinearTransformation('
                'transformation_matrix={transformation_matrix},'
                'input_channels={input_channels},'
                'output_channels={output_channels})').format(transformation_matrix=self._matrix.tolist(),
                                                             input_channels=self._input_channels,
                                                             output_channels=self._output_channels)

    def is_constant_invariant(self):
        """Signals if the transformation always maps constants to constants."""
        return True

class OffsetTransformation(Transformation):
    def __init__(self, offsets: Mapping[ChannelID, Real]):
        """Adds an offset to each channel specified in offsets.

        Channels not in offsets are forewarded

        Args:
            offsets: Channel -> offset mapping
        """
        self._offsets = dict(offsets.items())

    def __call__(self, time: Union[np.ndarray, float],
                 data: Mapping[ChannelID, Union[np.ndarray, float]]) -> Mapping[ChannelID, Union[np.ndarray, float]]:
        return {channel: channel_values + self._offsets[channel] if channel in self._offsets else channel_values
                for channel, channel_values in data.items()}

    def get_input_channels(self, output_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        return output_channels

    def get_output_channels(self, input_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        return input_channels

    @property
    def compare_key(self) -> frozenset:
        return frozenset(self._offsets.items())

    def __repr__(self):
        return 'OffsetTransformation(%r)' % self._offsets

    def is_constant_invariant(self):
        """Signals if the transformation always maps constants to constants."""
        return True

class ScalingTransformation(Transformation):
    def __init__(self, factors: Mapping[ChannelID, Real]):
        self._factors = dict(factors.items())

    def __call__(self, time: Union[np.ndarray, float],
                 data: Mapping[ChannelID, Union[np.ndarray, float]]) -> Mapping[ChannelID, Union[np.ndarray, float]]:
        return {channel: channel_values * self._factors[channel] if channel in self._factors else channel_values
                for channel, channel_values in data.items()}

    def get_input_channels(self, output_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        return output_channels

    def get_output_channels(self, input_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        return input_channels

    @property
    def compare_key(self) -> frozenset:
        return frozenset(self._factors.items())

    def __repr__(self):
        return 'ScalingTransformation(%r)' % self._factors

    def is_constant_invariant(self):
        """Signals if the transformation always maps constants to constants."""
        return True

try:
    import pandas

    def linear_transformation_from_pandas(transformation: pandas.DataFrame) -> LinearTransformation:
        """ Creates a LinearTransformation object out of a pandas data frame.

        Args:
            transformation (pandas.DataFrame): The pandas.DataFrame object out of which a LinearTransformation will be formed.

        Returns:
            the created LinearTransformation instance
        """
        return LinearTransformation(transformation.values, transformation.columns, transformation.index)

    LinearTransformation.from_pandas = linear_transformation_from_pandas
except ImportError:
    pass

class ParallelConstantChannelTransformation(Transformation):
    def __init__(self, channels: Mapping[ChannelID, Real]):
        """Set channel values to given values regardless their former existence

        Args:
            channels: Channels present in this map are set to the given value.
        """
        self._channels = {channel: float(value)
                          for channel, value in channels.items()}

    def __call__(self, time: Union[np.ndarray, float],
                 data: Mapping[ChannelID, Union[np.ndarray, float]]) -> Mapping[ChannelID, Union[np.ndarray, float]]:
        overwritten = {channel: np.full_like(time, fill_value=value, dtype=float)
                       for channel, value in self._channels.items()}
        return {**data, **overwritten}

    @property
    def compare_key(self) -> Tuple[Tuple[ChannelID, float], ...]:
        return tuple(sorted(self._channels.items()))

    def get_input_channels(self, output_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        return output_channels - self._channels.keys()

    def get_output_channels(self, input_channels: AbstractSet[ChannelID]) -> AbstractSet[ChannelID]:
        return input_channels | self._channels.keys()

    def __repr__(self):
        return 'ParallelConstantChannelTransformation(%r)' % self._channels

    def is_constant_invariant(self):
        """Signals if the transformation always maps constants to constants."""
        return True

def chain_transformations(*transformations: Transformation) -> Transformation:
    parsed_transformations = []
    for transformation in transformations:
        if transformation is IdentityTransformation() or transformation is None:
            pass
        elif isinstance(transformation, ChainedTransformation):
            parsed_transformations.extend(transformation.transformations)
        else:
            parsed_transformations.append(transformation)
    if len(parsed_transformations) == 0:
        return IdentityTransformation()
    elif len(parsed_transformations) == 1:
        return parsed_transformations[0]
    else:
        return ChainedTransformation(*parsed_transformations)

max_stars_repo_name: eendebakpt/qupulse

max_issues_repo_name: eendebakpt/qupulse

max_forks_repo_name: eendebakpt/qupulse


text:
# -*- coding: utf-8 -*-
"""
This module contains all classes and functions specific for processing GLS.
"""
import pandas as pd
import numpy as np
import scipy
import scipy.sparse as sps
import scipy.sparse.linalg as spsl
import sandy

__author__ = "Aitor Bengoechea"
__all__ = [
        "gls_update",
        "_y_calc",
        "chi_individual",
        "chi_diag",
        "chi_square",
        "ishikawa_factor",
        "constrained_gls_update"
        ]

x_prior = [1, 2, 3]
y_extra = pd.Series([2, 3, 4], index=[1, 2, 3])
S = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
Vx_prior = [[0, 0, 0], [0, 3, 0], [0, 0, 8]]
Vy_extra = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
N_e = 1

def gls_update(x_prior, S, Vx_prior, Vy_extra, y_extra, sparse=False,
               threshold=None):
    """
    Perform GlS update for a given variances, vectors and sensitivity.
    .. math::
        $$
        x_{post} = x_{prior} + V_{x_{prior}}\cdot S.T \cdot \left(S\cdot V_{x_{prior}}\cdot S.T + V_{y_{extra}}\right)^{-1} \cdot \left(y_{extra} - y_{calc}\right)
        $$

    Parameters
    ----------
    x_prior: 1D iterable
        Vector in which we are going to apply GLS (MX1)
    Vx_prior : 2D iterable
        2D covariance matrix of x_prior (MXN).
    Vy_extra : 2D iterable
        2D covariance matrix for y_extra (MXN).
    S : 2D iterable
        2D sensitivity of the model y=f(x) (MXN).
    y_extra : 1D iterable
        1D extra info on output (NX1)
    sparse : `bool`, optional
        Option to use sparse matrix for calculations. The default is False.
    threshold : `int`, optional
        Thereshold to avoid numerical fluctuations. The default is None.

    Returns
    -------
    `pd.Series`
        GLS apply to a vector x_prior given S, Vx_prior, Vy_extra, y_calc
        and y_extra.

    Example
    -------
    >>> S = [[1, 2], [3, 4]]
    >>> y = pd.Series([1, 1])
    >>> Vx = sandy.CategoryCov.from_var([1, 1]).data
    >>> Vy = pd.DataFrame([[1, 0], [0, 1]], index=[1, 2], columns=[1, 2])
    >>> x = [1, 1]
    >>> x_p = [2, 2]
    >>> gls_update(y, S, Vx, Vy, x_p)
    0   2.00000e-01
    1   4.85714e-01
    dtype: float64

    >>> gls_update(y, S, Vx, Vy, x_p, sparse=True)
    0   2.00000e-01
    1   4.85714e-01
    dtype: float64
    """
    # Model calculus:
    x_prior_ = pd.Series(x_prior)
    S_ = pd.DataFrame(S).reindex(columns=x_prior_.index)
    y_calc_ = _y_calc(x_prior, S, sparse=sparse)
    y_extra_ = pd.Series(y_extra)
    y_calc_ = y_calc_.reindex(y_extra_.index)
    S_ = S_.reindex(index=y_extra_.index)
    # Data in a appropriate format
    delta = y_extra_ - y_calc_
    Vx_prior_ = sandy.CategoryCov(Vx_prior)
    # GLS update
    A = Vx_prior_._gls_general_sensitivity(S_, Vy_extra, sparse=sparse,
                                           threshold=threshold).values
    if sparse:
        A = sps.csr_matrix(A)
        index = x_prior_.index
        x_prior_ = x_prior_.values
        x_post = x_prior_ + A.dot(delta)
        x_post = pd.Series(x_post, index=index)
    else:
        x_post = x_prior_ + A.dot(delta)
    return x_post

def _y_calc(x_prior, S, sparse=False):
    """
    Perform model calculation in GLS.

    Parameters
    ----------
    x_prior: 1D iterable
        Vector in which we are going to apply GLS (MX1).
    S : 2D iterable
        2D sensitivity of the model y=f(x) (MXN).
    sparse : `bool`, optional
        Option to use sparse matrix for calculations. The default is False.

    Returns
    -------
    y_calc : `pd.Series`
        1D calculated output using S.dot(x_prior), e.g. calculated CFY

    Example
    -------
    S square matrix:
    >>> _y_calc(x_prior, S)
    0    1
    1    2
    2    3
    dtype: int64

    Different number of row and columns in S:
    >>> S = [[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 1]]
    >>> _y_calc(x_prior, S)
    0    1
    1    2
    2    3
    3    6
    dtype: int64
    >>> _y_calc(x_prior, S, sparse=True)
    0    1
    1    2
    2    3
    3    6
    dtype: int64
    """
    S_ = pd.DataFrame(S)
    x_prior_ = pd.Series(x_prior).reindex(S_.columns).fillna(0)
    if sparse:
        index = S_.index
        S_ = sps.csr_matrix(S_.values)
        y_calc = S_.dot(x_prior_.values)
        y_calc = pd.Series(y_calc, index=index)
    else:
        y_calc = S_.dot(x_prior_)
    return y_calc

def chi_individual(x_prior, S, Vx_prior, Vy_extra, y_extra, sparse=False):
    """
    Function to calculate individual chi-value measured in sigmas according to
    https://www.oecd-nea.org/jcms/pl_19760/intermediate-report-on-methods-and-approaches-to-provide-feedback-from-nuclear-and-covariance-data-adjustment-for-improvement-of-nuclear-data-files
    (page 9, equation (4.2))

    Parameters
    ----------
    x_prior: 1D iterable
        Vector in which we are going to apply GLS (MX1)
    S : 2D iterable
        2D sensitivity of the model y=f(x) (MXN).
    Vx_prior : 2D iterable
        2D covariance matrix of x_prior (MXN).
    Vy_extra : 2D iterable
        2D covariance matrix for y_extra (MXN).
    y_extra : 1D iterable
        1D extra info on output (NX1).
    sparse : `bool`, optional
        Option to use sparse matrix for calculations. The default is False.

    Returns
    -------
    `pd.Series`
        individual chi-value measured in sigmas.

    Example
    -------
    >>> chi_individual(x_prior, S, Vx_prior, Vy_extra, y_extra)
    1   1.00000e+00
    2   5.00000e-01
    3   3.33333e-01
    dtype: float64
    >>> chi_individual(x_prior, S, Vx_prior, Vy_extra, y_extra, sparse=True)
    1   1.00000e+00
    2   5.00000e-01
    3   3.33333e-01
    dtype: float64
    """
    Vx_prior_ = sandy.CategoryCov(Vx_prior)
    G = Vx_prior_._gls_G(S, Vy_extra, sparse=sparse)
    G = np.sqrt(np.diag(G))
    y_calc_ = _y_calc(x_prior, S, sparse=sparse).values
    y_extra_ = pd.Series(y_extra)
    delta = np.abs(y_extra_.values - y_calc_)
    return pd.Series(delta / G, index=y_extra_.index)

def chi_diag(x_prior, S, Vx_prior, Vy_extra, y_extra, sparse=False):
    """
    Function to calculate diagonal chi-value measured in sigmas
    $\chi_{ind,i}$>>1 according to
    https://www.oecd-nea.org/jcms/pl_19760/intermediate-report-on-methods-and-approaches-to-provide-feedback-from-nuclear-and-covariance-data-adjustment-for-improvement-of-nuclear-data-files
    (page 9, equation (4.3))

    Parameters
    ----------
    x_prior: 1D iterable
        Vector in which we are going to apply GLS (MX1)
    S : 2D iterable
        2D sensitivity of the model y=f(x) (MXN).
    Vx_prior : 2D iterable
        2D covariance matrix of x_prior (MXN).
    Vy_extra : 2D iterable
        2D covariance matrix for y_extra (MXN).
    y_extra : 1D iterable
        1D extra info on output (NX1)
    sparse : `bool`, optional
        Option to use sparse matrix for calculations. The default is False.

    Returns
    -------
    `pd.Series`
        diagonal chi-value measured in sigmas $\chi_{ind,i}$>>1

    Example
    -------
    >>> chi_diag(x_prior, S, Vx_prior, Vy_extra, y_extra)
    1   1.00000e+00
    2   2.00000e+00
    3   3.00000e+00
    dtype: float64
    >>> chi_diag(x_prior, S, Vx_prior, Vy_extra, y_extra, sparse=True)
    1   1.00000e+00
    2   2.00000e+00
    3   3.00000e+00
    dtype: float64
    """
    Vx_prior_ = sandy.CategoryCov(Vx_prior)
    G_inv = Vx_prior_._gls_G_inv(S, Vy_extra, sparse=sparse).values
    G_inv = np.sqrt(np.diag(G_inv))
    y_calc_ = _y_calc(x_prior, S, sparse=sparse).values
    y_extra_ = pd.Series(y_extra)
    delta = np.abs(y_extra_.values - y_calc_)
    return pd.Series(delta / G_inv, index=y_extra_.index)

def chi_square(x_prior, S, Vx_prior, Vy_extra, y_extra, N_e, sparse=False):
    """
    Function to calculate contribution to chi-square value according to
    https://www.oecd-nea.org/jcms/pl_19760/intermediate-report-on-methods-and-approaches-to-provide-feedback-from-nuclear-and-covariance-data-adjustment-for-improvement-of-nuclear-data-files
    (page 10, equation (4.4))

    Parameters
    ----------
    x_prior: 1D iterable
        Vector in which we are going to apply GLS (MX1)
    S : 2D iterable
        2D sensitivity of the model y=f(x) (MXN).
    Vx_prior : 2D iterable
        2D covariance matrix of x_prior (MXN).
    Vy_extra : 2D iterable
        2D covariance matrix for y_extra (MXN).
    y_extra : 1D iterable
        1D extra info on output (NX1)
    N_e : `int`
        Number of experimental values used in adjustment.
    sparse : `bool`, optional
        Option to use sparse matrix for calculations. The default is False.

    Returns
    -------
    `pd.Series`
        contribution to chi-square value

    Example
    -------
    >>> chi_square(x_prior, S, Vx_prior, Vy_extra, y_extra, N_e)
    1   1.00000e+00
    2   2.50000e-01
    3   1.11111e-01
    dtype: float64
    >>> chi_square(x_prior, S, Vx_prior, Vy_extra, y_extra, N_e, sparse=True)
    1   1.00000e+00
    2   2.50000e-01
    3   1.11111e-01
    dtype: float64
    """
    Vx_prior_ = sandy.CategoryCov(Vx_prior)
    G_inv = Vx_prior_._gls_G_inv(S, Vy_extra, sparse=sparse).values
    y_calc_ = _y_calc(x_prior, S, sparse=sparse).values
    y_extra_ = pd.Series(y_extra)
    delta = y_extra_.values - y_calc_
    chi_square = delta.T.dot(G_inv) * delta / N_e
    return pd.Series(chi_square, index=y_extra_.index)

def ishikawa_factor(S, Vx_prior, Vy_extra, sparse=False):
    """
    Function to obtain Ishikawa factor according to
    https://www.oecd-nea.org/jcms/pl_19760/intermediate-report-on-methods-and-approaches-to-provide-feedback-from-nuclear-and-covariance-data-adjustment-for-improvement-of-nuclear-data-files
    (page 10, equation (4.5))

    Parameters
    ----------
    S : 2D iterable
        2D sensitivity of the model y=f(x) (MXN).
    Vx_prior : 2D iterable
        2D covariance matrix of x_prior (MXN).
    Vy_extra : 2D iterable
        2D covariance matrix for y_extra (MXN).
    sparse : `bool`, optional
        Option to use sparse matrix for calculations. The default is False.

    Returns
    -------
    `pd.Series`
        Ishikawa factor.

    Example
    -------
    >>> ishikawa_factor(S, Vx_prior, Vy_extra)
    0   0.00000e+00
    1   3.00000e+00
    2   8.00000e+00
    dtype: float64
    >>> ishikawa_factor(S, Vx_prior, Vy_extra, sparse=True)
    0   0.00000e+00
    1   3.00000e+00
    2   8.00000e+00
    dtype: float64
    """
    Vx_prior_ = sandy.CategoryCov(Vx_prior)
    Vy_calc = Vx_prior_._gls_Vy_calc(S, sparse=sparse)
    Vy_values = np.diag(Vy_calc)
    Vy_extra_ = np.diag(pd.DataFrame(Vy_extra).values)
    index = pd.DataFrame(Vy_extra).index
    return pd.Series(Vy_values / Vy_extra_, index=index)

def constrained_gls_update(x_prior, S, Vx_prior, sparse=True, threshold=None):
    """
    Perform Constrained Least-Squares update for a given variances, vectors
    and sensitivity:
    .. math::
        $$
        x_{post} = x_{prior} + \left(S.T \cdot x_{prior} - x_{prior} \cdot S.T\right) \cdot \left(S\cdot V_{x_{prior}}\cdot S.T + V_{y_{extra}}\right)^{-1} \cdot S \cdot V_{x_{prior}}
        $$

    Parameters
    ----------
    x_prior: 1D iterable
        Vector in which we are going to apply GLS (MX1)
    S : 2D iterable
        2D sensitivity of the model y=f(x) (MXN).
    Vx_prior : 2D iterable
        2D covariance matrix of x_prior (MXN).
    sparse : `bool`, optional
        Option to use sparse matrix for calculations. The default is False.
    threshold : `int`, optional
            Thereshold to avoid numerical fluctuations. The default is None.

    Returns
    -------
    x_post : `pandas.Series`
        Constrained Least-Squares apply to a vector x_prior given S and
        Vx_prior.

    Example
    -------
    >>> S = np.array([[1, 2], [3, 4]])
    >>> Vx_prior = sandy.CategoryCov.from_var([1, 1]).data
    >>> x_prior = np.array([1, 2])
    >>> constrained_gls_update(x_prior, S, Vx_prior)
    0   -4.00000e+00
    1    5.50000e+00
    dtype: float64
    >>> constrained_gls_update(x_prior, S, Vx_prior, sparse=True)
    0   -4.00000e+00
    1    5.50000e+00
    dtype: float64
    """
    x_prior_ = pd.Series(x_prior)
    S_ = pd.DataFrame(S).reindex(columns=x_prior_.index)
    # Data in a appropriate format:
    delta = _y_calc(x_prior, S_.T, sparse=sparse) - x_prior.dot(S_.T)
    delta = delta.reindex(x_prior_.index)
    Vx_prior_ = sandy.CategoryCov(Vx_prior)
    # Constrained Least-Square sensitivity
    A = Vx_prior_._constrained_gls_sensitivity(S, sparse=sparse,
                                               threshold=threshold).values
    x_post = x_prior_ + delta.dot(A)
    return x_post

reason: import numpy,import scipy


text:
import numpy as np
import tensorflow as tf
from PIL import Image
import os
import matplotlib.pyplot as plt

input_dir = input('input data directory please: ') # input_dir/*.png ...
output_dir = os.path.join(input_dir, "tf_records_train")
if not os.path.isdir(output_dir):
    os.makedirs(output_dir)

def _int64_feature(value):
    """
    generate int64 feature.
    """
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def _bytes_feature(value):
    """
    generate byte feature.
    """
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def createRecord(imageDir):
    """
    create TFRecord data.

    Arguments:
    imageDir -- image directory.
    Return: none.
    """
    writer = tf.python_io.TFRecordWriter(os.path.join(imageDir, "tf_records_train/train"))
    # classNames = ["cat", "dog", "horse"]

    # for classIndex, className in enumerate(classNames):
    #     print "class name = ",className
    #     currentClassDir = os.path.join(imageDir,className)
    #     print "current dir = ",currentClassDir
    count = 0
    for index, imageName in enumerate(os.listdir(imageDir)):
        if not os.path.isdir(os.path.join(imageDir,imageName)):
            image = Image.open(os.path.join(imageDir,imageName))
            img_shape=np.shape(image)
            print(img_shape)
            if img_shape[-1] == 3:
                image_raw = image.tobytes() # convert image to binary format
                print ('nice, ', index, imageName)

                example = tf.train.Example(features = tf.train.Features(feature = {
                "image/height": _int64_feature(img_shape[0]),
                "image/width": _int64_feature(img_shape[1]),
                "image/encoded": _bytes_feature(image_raw),
                }))
                writer.write(example.SerializeToString())
            count += 1
    writer.close()

createRecord(input_dir)

max_stars_repo_name: drboog/FPK

max_issues_repo_name: drboog/FPK

max_forks_repo_name: drboog/FPK


text:
import scipy.sparse as sp
import scipy.sparse.linalg as la
import numpy as np
import math as m
import time
import sys
import os
import pandas as pd
import statsmodels.formula.api as sm

resultDir = os.environ.get('RESULTS')
if resultDir == None :
    print ("WARNING! $RESULTS not set! Attempt to write results will fail!\n")

botConc = float(sys.argv[1])
topConc = float(sys.argv[2])
l = float(sys.argv[3])
L = np.uint32(int(sys.argv[4]))
numVecs = int(sys.argv[5])
boundMult = float(sys.argv[6])
tolerance = float(sys.argv[7])
numTimeSlices = int(sys.argv[8])
fileInfo = sys.argv[9]

resultsPlace = resultDir+"/"+fileInfo+"/"

if not os.path.exists(resultsPlace):
    os.makedirs(resultsPlace)

with open(resultsPlace+'settings', 'w') as f:
    f.write('BotConcentration = ' + str(botConc) +'\n')
    f.write('TopConcentration = ' + str(topConc) +'\n')
    f.write('Lambda = ' + str(l) +'\n')
    f.write('SysSize = ' + str(L) +'\n')
    f.write('NumVecs = ' + str(numVecs)+'\n')
    f.write('BoundMult = ' + str(boundMult)+'\n')
    f.write('Tolerance = ' +str(tolerance)+'\n')

N = np.uint32(2**(L+4))

rateMatrix = sp.lil_matrix((N, N), dtype = np.float64)
densityMatrix = sp.lil_matrix((L+4, N), dtype = np.float64)
currentMatrix = sp.lil_matrix((L+3, N), dtype = np.float64)

topIncRate = (1.0+l)*boundMult*m.sqrt(l*topConc/(1.0-topConc))
topOutRate = (1.0+l)*boundMult*m.sqrt(l*(1.0-topConc)/topConc)
botIncRate = (1.0+l)*boundMult*m.sqrt(l*botConc/(1.0-botConc))
botOutRate = (1.0+l)*boundMult*m.sqrt(l*(1.0-botConc)/botConc)

t0 = time.clock()
for i in range(0, N):
    state = format(i, '032b')
    totLeakage = 0.0
    for position in range(32 - 4 - L, 32-2 - L):
        if state[position] == '1':
            newState = state[:position]+'0'+state[(position+1):]
#            print state
#            print newState
            j = np.uint32(int(newState, 2))
#            print format(j, '032b')+"\n"
            rateMatrix[j, i] += botOutRate
            totLeakage += botOutRate
            densityMatrix[position-(32-4-L), i] = 1
        else:
            newState = state[:position]+'1'+state[(position+1):]
#            print state
#            print newState
            j = np.uint32(int(newState, 2))
#            print format(j, '032b')+"\n"
            rateMatrix[j, i] += botIncRate
            totLeakage += botIncRate
    for position in range(32-3 - L , 32-1):
        if state[position] == '1':
            densityMatrix[position-(32-4-L), i] = 1
            if state[position-1] == '0':
                newState = state[:(position-1)]+'1'+'0'+state[(position+1):]
#                print state
#                print newState
                j = np.uint32(int(newState, 2))
#                print format(j, '032b')+"\n"
                if state[position+1] == '0':
                    rateMatrix[j, i] += 1.0
                    currentMatrix[position - (32-3-L), i] -= 1.0
                    totLeakage += 1.0
                else:
                    rateMatrix[j, i] += l
                    currentMatrix[position - (32-3-L), i] -= l
                    totLeakage += l
            if state[position+1] == '0':
                newState = state[:(position)]+'0'+'1'+state[(position+2):]
#                print state
#                print newState
                j = np.uint32(int(newState, 2))
#                print format(j, '032b')+"\n"
                if state[position-1] == '0':
                    rateMatrix[j, i] += 1.0
                    currentMatrix[position + 1 - (32-3-L), i] += 1.0
                    totLeakage += 1.0
                else:
                    rateMatrix[j, i] += l
                    currentMatrix[position + 1 - (32-3-L), i] += l
                    totLeakage += l
    for position in range(32 - 2, 32):
        if state[position] == '1':
            densityMatrix[position-(32-4-L), i] = 1
            newState = state[:position]+'0'+state[(position+1):]
#            print state
#            print newState
            j = np.uint32(int(newState, 2))
#            print format(j, '032b')+"\n"
            rateMatrix[j, i] += topOutRate
            totLeakage += topOutRate
        else:
            newState = state[:position]+'1'+state[(position+1):]
#            print state
#            print newState
            j = np.uint32(int(newState, 2))
#            print format(j, '032b')+"\n"
            rateMatrix[j, i] += topIncRate
            totLeakage += topIncRate
    rateMatrix[i, i] -= totLeakage

#print("RateMatrix created.")

cscRateMatrix = rateMatrix.tocsc()
cscDensityMatrix = densityMatrix.tocsc()
cscCurrentMatrix = currentMatrix.tocsc()

#print("RateMatrix reformatted.")

valsLR, vecsLR = la.eigs(cscRateMatrix, k=numVecs, tol=tolerance, which='LR', maxiter=100*N)

errs = []
for index in range(0, numVecs):
    overallMass = np.sum(vecsLR[:, index])
    vecsLR[:, index] = np.sign(overallMass)*vecsLR[:, index]/(np.linalg.norm(vecsLR[:, index], 1))
    errs.append(2.0*np.linalg.norm(cscRateMatrix.dot(vecsLR[:, index])-valsLR[index]*vecsLR[:, index], 1)/(np.linalg.norm(cscRateMatrix.dot(vecsLR[:, index]), 1)+np.abs(valsLR[index])*np.linalg.norm(vecsLR[:, index], 1)))

t1 = time.clock()

print(str(t1-t0)+"s for computation of this session's eigenpairs.\n")

print("Done messing around, now for time-dependence:\n")

avDens = cscDensityMatrix.dot(vecsLR)

groundState = vecsLR[:, 0]
position = 32 - 4 - L + (int(L/2) + 2)
#groundState = groundState[:position]+'1'+state[(position+1):]
for i in range(0, N):
    state = format(i, '032b')
    if state[position]=='0':
        groundState[i] = 0
groundState = groundState/(np.linalg.norm(groundState, 1))
stateTimeSeries = la.expm_multiply(cscRateMatrix, groundState, start=0.0, num=numTimeSlices, stop=10.0, endpoint=True)
densTimeSeries = cscDensityMatrix.dot(np.transpose(stateTimeSeries))
entropySeries = []
timeSeries = []
autoCorr = []
trivWeights = []
corrInfo = []
with open(resultsPlace+'autoCorr.dat', 'w') as f:
    for i in range(int(numTimeSlices/2), numTimeSlices):
        time = i*10.0/float(numTimeSlices-1)
        f.write(str(time)+" "+str(m.log(abs(np.real(densTimeSeries[(int(L/2) + 2)][i] - avDens[(int(L/2) + 2)][0]))))+"\n")
        timeSeries.append(i*10.0/float(numTimeSlices-1))
        autoCorr.append(m.log(abs(np.real(densTimeSeries[(int(L/2) + 2)][i] - avDens[(int(L/2) + 2)][0]))))
        trivWeights.append(1.0)
        corrInfo.append(str(i*10.0/float(numTimeSlices-1))+" "+str(m.log(abs(np.real(densTimeSeries[(int(L/2) + 2)][i] - avDens[(int(L/2) + 2)][0]))))+"\n")

y_list = timeSeries
x_list = autoCorr
y_err = trivWeights

# put x and y into a pandas DataFrame, and the weights into a Series
ws = pd.DataFrame({
    'x': x_list,
    'y': y_list
})
weights = pd.Series(trivWeights)

wls_fit = sm.wls('x ~ y', data=ws, weights=1.0 / ((weights)**2)).fit()
#ols_fit = sm.ols('x ~ y', data=ws).fit()

#print avDens
#print("\nThe mean current should be:\n")
avCurr = cscCurrentMatrix.dot(vecsLR)
#print avCurr

with open(resultsPlace+'eigenvalues.dat', 'w') as f:
    for eig in valsLR:
        f.write(str(np.real(eig))+'\n')

#with open(resultsPlace+'fullEigenvalues.dat', 'w') as f:
#    for eig in vals:
#        f.write(str(eig)+'\n')

with open(resultsPlace+'fullEigenvalues.dat', 'w') as f:
    for eig in valsLR:
        f.write(str(eig)+'\n')
#    for eig in valsSR:
#        f.write(str(eig)+'\n')
#    for eig in valsLI:
#        f.write(str(eig)+'\n')
#    for eig in valsSI:
#        f.write(str(eig)+'\n')

for index in range(0, 1):
    with open(resultsPlace+'densVec'+str(index)+'.dat', 'w') as f:
        for position in range(0, L+4):
            f.write(str(np.real(avDens[position, index]))+'\n')

#for index in range(0, numVecs):
#    with open(resultsPlace+'fullDensVec'+str(index)+'.dat', 'w') as f:
#        for position in range(0, L+4):
#            f.write(str(avDens[position, index])+' ')

for index in range(0, 1):
    with open(resultsPlace+'currVec'+str(index)+'.dat', 'w') as f:
        for position in range(0, L+3):
            f.write(str(np.real(avCurr[position, index]))+'\n')

#for index in range(0, numVecs):
#    with open(resultsPlace+'fullCurrVec'+str(index)+'.dat', 'w') as f:
#        for position in range(0, L+3):
#            f.write(str(avCurr[position, index])+' ')

with open(resultsPlace+'eigenErrs.dat', 'w') as f:
    for err in errs:
        f.write(str(err)+'\n')

entropy = 0.0
for i in range(0, N):
    temp = vecsLR[i, 0]
    if temp > 0.0:
        entropy += - temp*m.log(2, temp)

with open(resultsPlace+'groundEntropy.dat', 'w') as f:
    f.write(str(np.real(entropy))+'\n')

#with open(resultsPlace+"autoCorrRegression.dat", 'w') as f:
#    f.writelines([str(wls_fit.summary())+"\n", str(wls_fit.params[0])+" "+str(wls_fit.bse[0])+"\n", str(wls_fit.params[1])+" "+str(wls_fit.bse[1])+"\n"])

with open(resultsPlace+"autoCorrTime.dat", 'w') as f:
    f.write(str(-1.0/wls_fit.params[1])+" "+str(abs(wls_fit.bse[1]/wls_fit.params[1])))

max_stars_repo_name: joshuahellier/PhDStuff

max_issues_repo_name: joshuahellier/PhDStuff

max_forks_repo_name: joshuahellier/PhDStuff

reason: import numpy,import scipy,import statsmodels


text:
#!/usr/bin/env python
import argparse
import subprocess
import re
import csv
import numpy
import hgvs
from hgvs import *

import hgvs.dataproviders.uta
import hgvs.assemblymapper
import hgvs.config
from hgvs import variantmapper

#from hgvs.variantmapper import  AssemblyMapper

##This is version 2, will now do the lookup in HGVS lands, using local HGVS
#    # after install    #    #
#export HGVS_SEQREPO_DIR=/usr/local/share/seqrepo/2018-11-26
#export UTA_DB_URL=postgresql://anonymous@localhost:15032/uta/uta_20170117
##

parser = argparse.ArgumentParser(description='given a list of hgvss\'/rsids from snpeff and VEP, see if the mapping is correct')
#parser.add_argument("--mut", help="Batch output",required=True)
parser.add_argument("--vg", help="the initial file, chr\tpos\tHGVS",required=True)
parser.add_argument("--rsid", help="if the file (--vg) is rsids, use this flag.",default = False)
args = parser.parse_args()

#mutfi = args.mut
vgrfi = args.vg
vgr = csv.DictReader(open(vgrfi,'r'),delimiter='\t')
#muts = open(mutfi,'r')

database = hgvs.dataproviders.uta.connect("postgresql://anonymous@localhost:15032/uta/uta_20170117")

##-----------defs--------##

def checkMatch(inc_chr,inc_pos,mapd_chr,mapd_pos):#newtranslated_pos,dictionary

    found = False

    inc_val = inc_chr + ":" + inc_pos
    chk_val = "chr" + mapd_chr + ":" + mapd_pos

    if chk_val == inc_val:
    	found =True

    return found

#------------main--------##

#blah, turn the hgvs into a dict

varmapper = hgvs.assemblymapper.AssemblyMapper(database, assembly_name='GRCh37', alt_aln_method='splign')

for hg in vgr:

    storelocale = {}
    isitgood = None

    varparse = hgvs.parser.Parser()

    try:
        vartomap = varparse.parse_hgvs_variant(hg['hgvs'])
        mappedvar = varmapper.c_to_g(vartomap)

    except (IndexError, hgvs.exceptions.HGVSDataNotAvailableError, hgvs.exceptions.HGVSInvalidIntervalError,hgvs.exceptions.HGVSParseError, hgvs.exceptions.HGVSInvalidVariantError) as e:

        print("NO_MAPPING\t" + str(hg['chr']) + "\t" + str(hg['pos']) + "\t" +  str(hg['hgvs']) )
        mappedvar = None

    if re.match('NC_',str(mappedvar)):
        m = re.match('NC_0{1,6}(\d+)\.(\d+):g.(\d+)(.*)',str(mappedvar))
        storelocale[int(m.group(2))] = m.group(3) #loads the version, so that it can be sorted on.
        #print(m.group(3) + "\t" + m.group(2) + "\t" + m.group(1))
        #print("match " + m.group(2) )
        #print(str(len(storelocale.keys())) + "\t" +  str(storelocale.keys()))

    if len(storelocale.keys()) == 1:
        storek =  storelocale.keys()#this statement, will autosort??

        isitgood = checkMatch(hg['chr'],hg['pos'],m.group(1),m.group(3))

        if isitgood == True:
            print("MATCH_FOUND\t" + hg['hgvs'] + "\t" + hg['chr'] + "\t" + hg['pos']  + "\tchr" + m.group(1) + "\t" + storelocale[storek[0]])
            mappedvar = None
        else:
            print("INCORRECT_POSITION\t" + hg['hgvs'] + "\t" + hg['chr'] + "\t" + hg['pos']  + "\t" + str(mappedvar) )
            mappedvar = None


text:
# coding: utf-8

""" Read in the VL2 stream data, output a JSON file for Three.js """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys
import glob
import json
import logging

# Third-party
import numpy as np

# project
from streams.util import project_root

# Create logger
logger = logging.getLogger(__name__)

def main(N=None):
    path = os.path.join(project_root, "data", "VL2 Streams")
    all_filenames = glob.glob(os.path.join(path, "*.dat"))

    xyz = dict()
    for ii,filename in enumerate(all_filenames):
        with open(filename) as f:
            d = np.loadtxt(f)
            if N is not None:
                xyz[str(ii)] = dict(data=d[1:N+1,:3].tolist())
            else:
                xyz[str(ii)] = dict(data=d[1:,:3].tolist())

    output_file = os.path.join(path,
                    "{}_vl2.json".format(N if N is not None else 'all'))
    with open(output_file, 'w') as f:
        f.write(json.dumps(xyz))

if __name__ == "__main__":
    main(N=5000)

max_stars_repo_name: adrn/streams

max_issues_repo_name: adrn/streams

max_forks_repo_name: adrn/streams