text:
import torch.nn.functional as F
import torch.nn as nn
import torch
from utils.FaFModel import FaFNet,FaFMIMONet_512_16_16, FaFMIMONet_512_16_16_KD, FaFMIMONet_256_32_32, FaFMIMONet_256_32_32_KD, FaFMIMONet_128_64_64, \
	FaFMIMONet_128_64_64_KD, FaFMIMONet_64_128_128,FaFMIMONet_32_256_256, FaFMIMONet_layer_3_and_4, FeatEncoder,FaFMGDA
from utils.detection_util import *
from utils.min_norm_solvers import MinNormSolver
import numpy
import matplotlib.pyplot as plt
from data.obj_util import coor_to_vis

class FaFModule(object):
	def __init__(self, model,config,optimizer,criterion):
		self.MGDA = config.MGDA
		if self.MGDA:
			self.encoder = model[0]
			self.head = model[1]
			self.optimizer_encoder = optimizer[0]
			self.optimizer_head = optimizer[1]
			self.scheduler_encoder = torch.optim.lr_scheduler.MultiStepLR(self.optimizer_encoder, milestones=[50, 100, 150, 200], gamma=0.5)
			self.scheduler_head = torch.optim.lr_scheduler.MultiStepLR(self.optimizer_head, milestones=[50, 100, 150, 200], gamma=0.5)
			self.MGDA = config.MGDA
		else:
			self.model = model
			self.optimizer = optimizer
			self.scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[50, 100, 150, 200], gamma=0.5)
		self.criterion = criterion#{'cls_loss','loc_loss'}

		self.out_seq_len = config.pred_len
		self.category_num = config.category_num
		self.code_size = config.box_code_size
		self.loss_scale = None

		self.code_type = config.code_type
		self.loss_type=  config.loss_type
		self.pred_len = config.pred_len
		self.only_det = config.only_det
		if self.code_type in ['corner_1','corner_2','corner_3']:
			self.alpha = 1.
		elif self.code_type == 'faf':
			if self.loss_type == 'corner_loss':
				self.alpha= 1.
				if not self.only_det:
					self.alpha = 1.
			else:
				self.alpha = 0.1
		self.config = config

	def resume(self,path):
		def map_func(storage, location):
			return storage.cuda()

		if os.path.isfile(path):
			if rank == 0:
				print("=> loading checkpoint '{}'".format(path))

			checkpoint = torch.load(path, map_location=map_func)
			self.model.load_state_dict(checkpoint['state_dict'], strict=False)

			ckpt_keys = set(checkpoint['state_dict'].keys())
			own_keys = set(model.state_dict().keys())
			missing_keys = own_keys - ckpt_keys
			for k in missing_keys:
				print('caution: missing keys from checkpoint {}: {}'.format(path, k))
		else:
				print("=> no checkpoint found at '{}'".format(path))

	def corner_loss(self,anchors,reg_loss_mask,reg_targets,pred_result):
		N = pred_result.shape[0]
		anchors = anchors.unsqueeze(-2).expand(anchors.shape[0],anchors.shape[1],anchors.shape[2],anchors.shape[3],reg_loss_mask.shape[-1],anchors.shape[-1])
		assigned_anchor = anchors[reg_loss_mask]
		assigned_target = reg_targets[reg_loss_mask]
		assigned_pred = pred_result[reg_loss_mask]
		#print(assigned_anchor.shape,assigned_pred.shape,assigned_target.shape)
		#exit()
		pred_decode = bev_box_decode_torch(assigned_pred,assigned_anchor)
		target_decode = bev_box_decode_torch(assigned_target,assigned_anchor)
		pred_corners = center_to_corner_box2d_torch(pred_decode[...,:2],pred_decode[...,2:4],pred_decode[...,4:])
		target_corners = center_to_corner_box2d_torch(target_decode[...,:2],target_decode[...,2:4],target_decode[...,4:])
		loss_loc = torch.sum(torch.norm(pred_corners-target_corners,dim=-1)) / N

		return loss_loc

	def loss_calculator(self,result,anchors,reg_loss_mask,reg_targets,labels,N,motion_labels = None,motion_mask=None):
		loss_num =0
		# calculate loss
		weights = torch.Tensor([0.005, 1.0, 1.0, 1.0, 1.0]).cuda().double()
		loss_cls = torch.sum(self.criterion['cls'](result['cls'],labels)) /N
		loss_num += 1
		#loss_loc = torch.sum(self.criterion['loc'](result['loc'],reg_targets,mask = reg_loss_mask)) / N

		#Motion state
		if not motion_labels is None:
			loss_motion = torch.sum(self.criterion['cls'](result['state'],motion_labels)) /N
			loss_num += 1

		loss_mask_num = torch.nonzero(reg_loss_mask.contiguous().view(-1,reg_loss_mask.shape[-1])).size(0)
		#print(loss_mask_num)
		#print(torch.sum(reg_targets[:,:,:,:,0][reg_loss_mask[:,:,:,:,2]]))

		if self.code_type in ['corner_1','corner_2','corner_3']:
				target = reg_targets[reg_loss_mask].reshape(-1,5,2)
				flip_target = torch.stack([target[:,0],target[:,3],target[:,4],target[:,1],target[:,2]],dim=-2)
				pred = result['loc'][reg_loss_mask].reshape(-1,5,2)
				t = torch.sum(torch.norm(pred-target,dim=-1),dim=-1)
				f = torch.sum(torch.norm(pred-flip_target,dim=-1),dim=-1)
				loss_loc = torch.sum(torch.min(t,f)) / N

		elif self.code_type == 'faf':
			if self.loss_type == 'corner_loss':
				if self.only_det:
					loss_loc = self.corner_loss(anchors,reg_loss_mask,reg_targets,result['loc'])
					loss_num += 1
				elif self.config.pred_type in ['motion','center']:

					###only center/motion for pred

					loss_loc_1 = self.corner_loss(anchors,reg_loss_mask[...,0][...,[0]],reg_targets[...,[0],:],result['loc'][...,[0],:])
					pred_reg_loss_mask = reg_loss_mask[...,1:,:]
					if self.config.motion_state:
						pred_reg_loss_mask = motion_mask #mask out static object
					loss_loc_2 = F.smooth_l1_loss(result['loc'][...,1:,:][pred_reg_loss_mask],reg_targets[...,1:,:][pred_reg_loss_mask])
					loss_loc = loss_loc_1 + loss_loc_2
					loss_num += 2

				###corners for pred
				else:
					loss_loc = self.corner_loss(anchors,reg_loss_mask,reg_targets,result['loc'])
					loss_num += 1
			else:

				loss_loc = F.smooth_l1_loss(result['loc'][reg_loss_mask],reg_targets[reg_loss_mask])
				loss_num += 1

		if self.loss_scale is not None:
			if len(self.loss_scale)==4:
				loss = self.loss_scale[0]*loss_cls + self.loss_scale[1]*loss_loc_1 + self.loss_scale[2]*loss_loc_2 + self.loss_scale[3]*loss_motion
			elif len(self.loss_scale)==3:
				loss = self.loss_scale[0]*loss_cls + self.loss_scale[1]*loss_loc_1 + self.loss_scale[2]*loss_loc_2
			else:
				loss = self.loss_scale[0]*loss_cls + self.loss_scale[1]*loss_loc
		elif not motion_labels is None:
			loss = loss_cls + loss_loc + loss_motion
		else:
			loss = loss_cls + loss_loc

		if loss_num == 2:
			return (loss_num,loss, loss_cls,loss_loc)
		elif loss_num == 3:
			return (loss_num,loss, loss_cls,loss_loc_1,loss_loc_2)
		elif loss_num == 4:
			return (loss_num,loss, loss_cls,loss_loc_1,loss_loc_2,loss_motion)

	def step(self,data,batch_size, center_agent, forcast_num):

		bev_seq = data['bev_seq']
		labels = data['labels']
		reg_targets = data['reg_targets']
		reg_loss_mask = data['reg_loss_mask']
		anchors = data['anchors']
		vis_maps = data['vis_maps']
		trans_matrices = data['trans_matrices']
		num_agent = data['num_agent']
		newest_time = torch.zeros((batch_size, 1))
		delta_t = []
		for iter in range(batch_size):
			newest_time[iter] = int(data['file_name'][center_agent[iter]][iter].split(',')[-1].strip('/').split('_')[-1])
			delta_t_temp = []
			for iter_num in range(5):
				delta_t_temp.append(newest_time[iter] - int(data['file_name'][iter_num][iter].split(',')[-1].strip('/').split('_')[-1]))
			delta_t.append(delta_t_temp)
		# with torch.autograd.set_detect_anomaly(True):
		if self.MGDA:
			self.loss_scale = self.cal_loss_scale(data)
			x = self.encoder(bev_seq)
			result = self.head(x)
		else:
			result = self.model(bev_seq, trans_matrices, num_agent, batch_size=batch_size, center_agent = center_agent, delta_t = delta_t)

		# labels = labels[:,-1]
		# anchors = anchors[:, -1]
		# reg_loss_mask = reg_loss_mask[:, -1]
		# reg_targets = reg_targets[:,-1]
		labels = labels.view(result['cls'].shape[0],-1,result['cls'].shape[-1])

		N = bev_seq.shape[0]

		loss_collect = self.loss_calculator(result,anchors,reg_loss_mask,reg_targets,labels,N)
		loss_num = loss_collect[0]
		if loss_num == 3:
			loss_num,loss, loss_cls,loss_loc_1,loss_loc_2 = loss_collect
		elif loss_num ==2:
			loss_num,loss, loss_cls,loss_loc = loss_collect
		elif loss_num == 4:
			loss_num,loss, loss_cls,loss_loc_1,loss_loc_2,loss_motion = loss_collect

		if self.MGDA:
			self.optimizer_encoder.zero_grad()
			self.optimizer_head.zero_grad()
			loss.backward()
			self.optimizer_encoder.step()
			self.optimizer_head.step()
		else:
			self.optimizer.zero_grad()
			loss.backward()
			self.optimizer.step()

		if self.config.pred_type in ['motion','center'] and not self.only_det:
			if self.config.motion_state:
				return loss.item(),loss_cls.item(),loss_loc_1.item(),loss_loc_2.item(), loss_motion.item()
			else:
				return loss.item(),loss_cls.item(),loss_loc_1.item(),loss_loc_2.item()
		else:
			return loss.item(),loss_cls.item(),loss_loc.item()

	def predict(self,data,validation=True):

		bev_seq = data['bev_seq']
		vis_maps = data['vis_maps']

		if self.MGDA:
			x = self.encoder(bev_seq)
			result = self.head(x)
		else:
			result = self.model(bev_seq,vis=vis_maps)

		N = bev_seq.shape[0]

		if validation:
			labels = data['labels']
			anchors = data['anchors']
			reg_targets = data['reg_targets']
			reg_loss_mask = data['reg_loss_mask']
			motion_labels = None
			motion_mask = None

			labels = labels.view(result['cls'].shape[0],-1,result['cls'].shape[-1])
			if self.config.motion_state:
				motion_labels = data['motion_label']
				motion_mask = data['motion_mask']
				motion_labels = motion_labels.view(result['state'].shape[0],-1,result['state'].shape[-1])
			N = bev_seq.shape[0]

			loss_collect = self.loss_calculator(result,anchors,reg_loss_mask,reg_targets,labels,N,motion_labels,motion_mask)
			loss_num = loss_collect[0]
			if loss_num == 3:
				loss_num,loss, loss_cls,loss_loc_1,loss_loc_2 = loss_collect
			elif loss_num ==2:
				loss_num,loss, loss_cls,loss_loc = loss_collect
			elif loss_num == 4:
				loss_num,loss, loss_cls,loss_loc_1,loss_loc_2,loss_motion = loss_collect

			batch_box_preds = result['loc']
			batch_cls_preds = result['cls']

			if self.config.motion_state:
				batch_motion_preds = result['state']
			else:
				batch_motion_preds = None

			if not self.only_det:
				if self.config.pred_type == 'center':
					batch_box_preds[:,:,:,:,1:,2:] = batch_box_preds[:,:,:,:,[0],2:]

		class_selected = apply_nms_det(batch_box_preds, batch_cls_preds,anchors,self.code_type,self.config,batch_motion_preds)
		#class_selected = None
		if validation:
			if self.config.pred_type in ['motion','center'] and not self.only_det:
				if self.config.motion_state:
					return loss.item(),loss_cls.item(),loss_loc_1.item(),loss_loc_2.item(), loss_motion.item(),class_selected
				else:
					return loss.item(),loss_cls.item(),loss_loc_1.item(),loss_loc_2.item(),class_selected
			else:
				return loss.item(),loss_cls.item(),loss_loc.item(),class_selected
		else:
			return class_selected

	def predict_all(self,data,batch_size,validation=True, center_agent = 0):
		NUM_AGENT = 5
		bev_seq = data['bev_seq']
		vis_maps = data['vis_maps']
		trans_matrices = data['trans_matrices']
		num_agent_tensor = data['num_agent']
		num_sensor = num_agent_tensor[0, 0]

		if self.MGDA:
			x = self.encoder(bev_seq)
			result = self.head(x)
		else:
			with torch.no_grad():
				result= self.model(bev_seq, trans_matrices, num_agent_tensor, batch_size=batch_size, center_agent = center_agent)
			# result = self.model(bev_seq,vis=vis_maps,training=False)
        #
		N = bev_seq.shape[0]

		if validation:
			labels = data['labels']
			anchors = data['anchors']
			reg_targets = data['reg_targets']
			reg_loss_mask = data['reg_loss_mask']
			motion_labels = None
			motion_mask = None

			labels = labels.view(result['cls'].shape[0],-1,result['cls'].shape[-1])

			if self.config.motion_state:
				motion_labels = data['motion_label']
				motion_mask = data['motion_mask']
				motion_labels = motion_labels.view(result['state'].shape[0],-1,result['state'].shape[-1])
			N = bev_seq.shape[0]

			loss_collect = self.loss_calculator(result,anchors,reg_loss_mask,reg_targets,labels,N,motion_labels,motion_mask)
			loss_num = loss_collect[0]
			if loss_num == 3:
				loss_num,loss, loss_cls,loss_loc_1,loss_loc_2 = loss_collect
			elif loss_num ==2:
				loss_num,loss, loss_cls,loss_loc = loss_collect
			elif loss_num == 4:
				loss_num,loss, loss_cls,loss_loc_1,loss_loc_2,loss_motion = loss_collect

		seq_results = [[] for i in range(NUM_AGENT)]
		# global_points = [[] for i in range(num_sensor)]
		# cls_preds = [[] for i in range(num_sensor)]

		for k in range(NUM_AGENT):
			bev_seq = torch.unsqueeze(data['bev_seq'][k, :, :, :, :], 0)

			if torch.nonzero(bev_seq).shape[0] == 0:
				seq_results[k] = []
			else:
				batch_box_preds = torch.unsqueeze(result['loc'][k, :, :, :, :, :],0)
				batch_cls_preds = torch.unsqueeze(result['cls'][k, :, :], 0)
				anchors = torch.unsqueeze(data['anchors'][k, :, :, :, :],0)
				batch_motion_preds = None

				if not self.only_det:
					if self.config.pred_type == 'center':
						batch_box_preds[:,:,:,:,1:,2:] = batch_box_preds[:,:,:,:,[0],2:]

				class_selected = apply_nms_det(batch_box_preds, batch_cls_preds,anchors,self.code_type,self.config,batch_motion_preds)
				seq_results[k] = class_selected

				# global_points[k], cls_preds[k] = apply_box_global_transform(trans_matrices_map[k],batch_box_preds,batch_cls_preds,anchors,self.code_type,self.config,batch_motion_preds)

		# all_points_scene = numpy.concatenate(tuple(global_points), 0)
		# cls_preds_scene = torch.cat(tuple(cls_preds), 0)
		# class_selected_global = apply_nms_global_scene(all_points_scene, cls_preds_scene)

		if validation:
			return loss.item(),loss_cls.item(),loss_loc.item(),seq_results
		else:
			return seq_results

	def predict_all_with_box_com(self, data, trans_matrices_map, validation=True):
		NUM_AGENT = 5
		bev_seq = data['bev_seq']
		vis_maps = data['vis_maps']
		trans_matrices = data['trans_matrices']
		num_agent_tensor = data['num_agent']
		num_sensor = num_agent_tensor[0, 0]

		if self.MGDA:
			x = self.encoder(bev_seq)
			result = self.head(x)
		else:
			result = self.model(bev_seq, trans_matrices, num_agent_tensor, batch_size=1)

		N = bev_seq.shape[0]

		if validation:
			labels = data['labels']
			anchors = data['anchors']
			reg_targets = data['reg_targets']
			reg_loss_mask = data['reg_loss_mask']
			motion_labels = None
			motion_mask = None

			labels = labels.view(result['cls'].shape[0], -1, result['cls'].shape[-1])

			if self.config.motion_state:
				motion_labels = data['motion_label']
				motion_mask = data['motion_mask']
				motion_labels = motion_labels.view(result['state'].shape[0], -1, result['state'].shape[-1])
			N = bev_seq.shape[0]

			loss_collect = self.loss_calculator(result, anchors, reg_loss_mask, reg_targets, labels, N, motion_labels,
												motion_mask)
			loss_num = loss_collect[0]
			if loss_num == 3:
				loss_num, loss, loss_cls, loss_loc_1, loss_loc_2 = loss_collect
			elif loss_num == 2:
				loss_num, loss, loss_cls, loss_loc = loss_collect
			elif loss_num == 4:
				loss_num, loss, loss_cls, loss_loc_1, loss_loc_2, loss_motion = loss_collect

		seq_results = [[] for i in range(NUM_AGENT)]
		local_results_wo_local_nms = [[] for i in range(NUM_AGENT)]
		local_results_af_local_nms = [[] for i in range(NUM_AGENT)]

		global_points = [[] for i in range(num_sensor)]
		cls_preds = [[] for i in range(num_sensor)]
		global_boxes_af_localnms = [[] for i in range(num_sensor)]
		box_scores_af_localnms = [[] for i in range(num_sensor)]

		forward_message_size = 0
		forward_message_size_two_nms = 0

		for k in range(NUM_AGENT):
			bev_seq = torch.unsqueeze(data['bev_seq'][k, :, :, :, :], 0)

			if torch.nonzero(bev_seq).shape[0] == 0:
				seq_results[k] = []
			else:
				batch_box_preds = torch.unsqueeze(result['loc'][k, :, :, :, :, :], 0)
				batch_cls_preds = torch.unsqueeze(result['cls'][k, :, :], 0)
				anchors = torch.unsqueeze(data['anchors'][k, :, :, :, :], 0)

				if self.config.motion_state:
					batch_motion_preds = result['state']
				else:
					batch_motion_preds = None

				if not self.only_det:
					if self.config.pred_type == 'center':
						batch_box_preds[:, :, :, :, 1:, 2:] = batch_box_preds[:, :, :, :, [0], 2:]

				class_selected, box_scores_pred_cls = apply_nms_det(batch_box_preds, batch_cls_preds, anchors,
																	self.code_type, self.config, batch_motion_preds)

				# transform all the boxes before local nms to the global coordinate
				# global_points[k], cls_preds[k] = apply_box_global_transform(trans_matrices_map[k], batch_box_preds,
				#                                                            batch_cls_preds, anchors, self.code_type,
				#                                                            self.config, batch_motion_preds)

				# transform the boxes after local nms to the global coordinate
				global_boxes_af_localnms[k], box_scores_af_localnms[k] = apply_box_global_transform_af_localnms(
					trans_matrices_map[k], class_selected, box_scores_pred_cls)
				# print(cls_preds[k].shape, box_scores_af_localnms[k].shape)

				forward_message_size = forward_message_size + 256 * 256 * 6 * 4 * 2
				forward_message_size_two_nms = forward_message_size_two_nms + global_boxes_af_localnms[k].shape[
					0] * 4 * 2

		# global results with one NMS
		# all_points_scene = numpy.concatenate(tuple(global_points), 0)
		# cls_preds_scene = torch.cat(tuple(cls_preds), 0)
		# class_selected_global = apply_nms_global_scene(all_points_scene, cls_preds_scene)

		# global results with two NMS
		global_boxes_af_local_nms = numpy.concatenate(tuple(global_boxes_af_localnms), 0)
		box_scores_af_local_nms = torch.cat(tuple(box_scores_af_localnms), 0)
		class_selected_global_af_local_nms = apply_nms_global_scene(global_boxes_af_local_nms, box_scores_af_local_nms)

		# transform the consensus global boxes to local agents (two NMS)
		back_message_size_two_nms = 0
		for k in range(num_sensor):
			local_results_af_local_nms[k], ms = apply_box_local_transform(class_selected_global_af_local_nms,
																		  trans_matrices_map[k])
			back_message_size_two_nms = back_message_size_two_nms + ms

		sample_bandwidth_two_nms = forward_message_size_two_nms + back_message_size_two_nms

		# transform the consensus global boxes to local agents (One NMS)
		# back_message_size = 0
		# for k in range(num_sensor):
		#    local_results_wo_local_nms[k], ms = apply_box_local_transform(class_selected_global, trans_matrices_map[k])
		#    back_message_size = back_message_size + ms

		# sample_bandwidth = forward_message_size + back_message_size

		return loss.item(), loss_cls.item(), loss_loc.item(), local_results_af_local_nms, class_selected_global_af_local_nms, sample_bandwidth_two_nms

	def cal_loss_scale(self,data):
		bev_seq = data['bev_seq']
		labels = data['labels']
		reg_targets = data['reg_targets']
		reg_loss_mask = data['reg_loss_mask']
		anchors = data['anchors']
		motion_labels = None
		motion_mask = None

		with torch.no_grad():
			shared_feats = self.encoder(bev_seq)
		shared_feats_tensor = shared_feats.clone().detach().requires_grad_(True)
		result = self.head(shared_feats_tensor)
		if self.config.motion_state:
			motion_labels = data['motion_label']
			motion_mask = data['motion_mask']
			motion_labels = motion_labels.view(result['state'].shape[0],-1,result['state'].shape[-1])
		self.optimizer_encoder.zero_grad()
		self.optimizer_head.zero_grad()
		grads = {}
		labels = labels.view(result['cls'].shape[0],-1,result['cls'].shape[-1])
		N = bev_seq.shape[0]

		# calculate loss
		grad_len = 0

		'''
		Classification Loss
		'''
		loss_cls = self.alpha*torch.sum(self.criterion['cls'](result['cls'],labels)) /N
		#loss_loc = torch.sum(self.criterion['loc'](result['loc'],reg_targets,mask = reg_loss_mask)) / N
		self.optimizer_encoder.zero_grad()
		self.optimizer_head.zero_grad()

		loss_cls.backward(retain_graph=True)
		grads[0] = []
		grads[0].append(shared_feats_tensor.grad.data.clone().detach())
		shared_feats_tensor.grad.data.zero_()
		grad_len += 1

		'''
		Localization Loss
		'''
		loc_scale = False
		loss_mask_num = torch.nonzero(reg_loss_mask.view(-1,reg_loss_mask.shape[-1])).size(0)

		if self.code_type in ['corner_1','corner_2','corner_3']:
				target = reg_targets[reg_loss_mask].reshape(-1,5,2)
				flip_target = torch.stack([target[:,0],target[:,3],target[:,4],target[:,1],target[:,2]],dim=-2)
				pred = result['loc'][reg_loss_mask].reshape(-1,5,2)
				t = torch.sum(torch.norm(pred-target,dim=-1),dim=-1)
				f = torch.sum(torch.norm(pred-flip_target,dim=-1),dim=-1)
				loss_loc = torch.sum(torch.min(t,f)) / N

		elif self.code_type == 'faf':
			if self.loss_type == 'corner_loss':
				if self.only_det:
					loss_loc = self.corner_loss(anchors,reg_loss_mask,reg_targets,result['loc'])
				elif self.config.pred_type in ['motion','center']:

					###only center/motion for pred

					loss_loc_1 = self.corner_loss(anchors,reg_loss_mask[...,0][...,[0]],reg_targets[...,[0],:],result['loc'][...,[0],:])
					pred_reg_loss_mask = reg_loss_mask[...,1:,:]
					if self.config.motion_state:
						pred_reg_loss_mask = motion_mask #mask out static object
					loss_loc_2 = F.smooth_l1_loss(result['loc'][...,1:,:][pred_reg_loss_mask],reg_targets[...,1:,:][pred_reg_loss_mask])

					self.optimizer_encoder.zero_grad()
					self.optimizer_head.zero_grad()

					loss_loc_1.backward(retain_graph=True)
					grads[1] = []
					grads[1].append(shared_feats_tensor.grad.data.clone().detach())
					shared_feats_tensor.grad.data.zero_()

					self.optimizer_encoder.zero_grad()
					self.optimizer_head.zero_grad()

					loss_loc_2.backward(retain_graph=True)
					grads[2] = []
					grads[2].append(shared_feats_tensor.grad.data.clone().detach())
					shared_feats_tensor.grad.data.zero_()
					loc_scale = True
					grad_len += 2

				###corners for pred
				else:
					loss_loc = self.corner_loss(anchors,reg_loss_mask,reg_targets,result['loc'])
			else:

				loss_loc = F.smooth_l1_loss(result['loc'][reg_loss_mask],reg_targets[reg_loss_mask])

			if not loc_scale:
				grad_len += 1
				self.optimizer_encoder.zero_grad()
				self.optimizer_head.zero_grad()
				loss_loc.backward(retain_graph=True)
				grads[1] = []
				grads[1].append(shared_feats_tensor.grad.data.clone().detach())
				shared_feats_tensor.grad.data.zero_()

		'''
		Motion state Loss
		'''
		if self.config.motion_state:
			loss_motion = torch.sum(self.criterion['cls'](result['state'],motion_labels)) /N

			self.optimizer_encoder.zero_grad()
			self.optimizer_head.zero_grad()

			loss_motion.backward(retain_graph=True)
			grads[3] = []
			grads[3].append(shared_feats_tensor.grad.data.clone().detach())
			shared_feats_tensor.grad.data.zero_()
			grad_len += 1

		# ---------------------------------------------------------------------
		# -- Frank-Wolfe iteration to compute scales.
		scale = np.zeros(grad_len, dtype=np.float32)
		sol, min_norm = MinNormSolver.find_min_norm_element([grads[t] for t in range(grad_len)])
		for i in range(grad_len):
			scale[i] = float(sol[i])

		#print(scale)
		return scale

class FaFModuleKD(object):
	def __init__(self, model, teacher, config, optimizer, criterion):
		self.MGDA = config.MGDA
		if self.MGDA:
			self.encoder = model[0]
			self.head = model[1]
			self.optimizer_encoder = optimizer[0]
			self.optimizer_head = optimizer[1]
			self.scheduler_encoder = torch.optim.lr_scheduler.MultiStepLR(self.optimizer_encoder,
																		  milestones=[50, 100, 150, 200], gamma=0.5)
			self.scheduler_head = torch.optim.lr_scheduler.MultiStepLR(self.optimizer_head,
																	   milestones=[50, 100, 150, 200], gamma=0.5)
			self.MGDA = config.MGDA
		else:
			self.model = model
			self.teacher = teacher

			for k, v in self.teacher.named_parameters():
				v.requires_grad = False

			self.optimizer = optimizer
			self.scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[50, 100, 150, 200], gamma=0.5)
		self.criterion = criterion  # {'cls_loss','loc_loss'}

		self.out_seq_len = config.pred_len
		self.category_num = config.category_num
		self.code_size = config.box_code_size
		self.loss_scale = None

		self.code_type = config.code_type
		self.loss_type = config.loss_type
		self.pred_len = config.pred_len
		self.only_det = config.only_det
		if self.code_type in ['corner_1', 'corner_2', 'corner_3']:
			self.alpha = 1.
		elif self.code_type == 'faf':
			if self.loss_type == 'corner_loss':
				self.alpha = 1.
				if not self.only_det:
					self.alpha = 1.
			else:
				self.alpha = 0.1
		self.config = config

	def resume(self, path):
		def map_func(storage, location):
			return storage.cuda()

		if os.path.isfile(path):
			if rank == 0:
				print("=> loading checkpoint '{}'".format(path))

			checkpoint = torch.load(path, map_location=map_func)
			self.model.load_state_dict(checkpoint['state_dict'], strict=False)

			ckpt_keys = set(checkpoint['state_dict'].keys())
			own_keys = set(model.state_dict().keys())
			missing_keys = own_keys - ckpt_keys
			for k in missing_keys:
				print('caution: missing keys from checkpoint {}: {}'.format(path, k))
		else:
			print("=> no checkpoint found at '{}'".format(path))

	def corner_loss(self, anchors, reg_loss_mask, reg_targets, pred_result):
		N = pred_result.shape[0]
		anchors = anchors.unsqueeze(-2).expand(anchors.shape[0], anchors.shape[1], anchors.shape[2], anchors.shape[3],
											   reg_loss_mask.shape[-1], anchors.shape[-1])
		assigned_anchor = anchors[reg_loss_mask]
		assigned_target = reg_targets[reg_loss_mask]
		assigned_pred = pred_result[reg_loss_mask]
		# print(assigned_anchor.shape,assigned_pred.shape,assigned_target.shape)
		# exit()
		pred_decode = bev_box_decode_torch(assigned_pred, assigned_anchor)
		target_decode = bev_box_decode_torch(assigned_target, assigned_anchor)
		pred_corners = center_to_corner_box2d_torch(pred_decode[..., :2], pred_decode[..., 2:4], pred_decode[..., 4:])
		target_corners = center_to_corner_box2d_torch(target_decode[..., :2], target_decode[..., 2:4],
													  target_decode[..., 4:])
		loss_loc = torch.sum(torch.norm(pred_corners - target_corners, dim=-1)) / N

		return loss_loc

	def loss_calculator(self, result, anchors, reg_loss_mask, reg_targets, labels, N, motion_labels=None,
						motion_mask=None):
		loss_num = 0
		# calculate loss
		weights = torch.Tensor([0.005, 1.0, 1.0, 1.0, 1.0]).cuda().double()
		loss_cls = torch.sum(self.criterion['cls'](result['cls'], labels)) / N
		loss_num += 1
		# loss_loc = torch.sum(self.criterion['loc'](result['loc'],reg_targets,mask = reg_loss_mask)) / N

		# Motion state
		if not motion_labels is None:
			loss_motion = torch.sum(self.criterion['cls'](result['state'], motion_labels)) / N
			loss_num += 1

		loss_mask_num = torch.nonzero(reg_loss_mask.view(-1, reg_loss_mask.shape[-1])).size(0)
		# print(loss_mask_num)
		# print(torch.sum(reg_targets[:,:,:,:,0][reg_loss_mask[:,:,:,:,2]]))

		if self.code_type in ['corner_1', 'corner_2', 'corner_3']:
			target = reg_targets[reg_loss_mask].reshape(-1, 5, 2)
			flip_target = torch.stack([target[:, 0], target[:, 3], target[:, 4], target[:, 1], target[:, 2]], dim=-2)
			pred = result['loc'][reg_loss_mask].reshape(-1, 5, 2)
			t = torch.sum(torch.norm(pred - target, dim=-1), dim=-1)
			f = torch.sum(torch.norm(pred - flip_target, dim=-1), dim=-1)
			loss_loc = torch.sum(torch.min(t, f)) / N

		elif self.code_type == 'faf':
			if self.loss_type == 'corner_loss':
				if self.only_det:
					loss_loc = self.corner_loss(anchors, reg_loss_mask, reg_targets, result['loc'])
					loss_num += 1
				elif self.config.pred_type in ['motion', 'center']:

					###only center/motion for pred

					loss_loc_1 = self.corner_loss(anchors, reg_loss_mask[..., 0][..., [0]], reg_targets[..., [0], :],
												  result['loc'][..., [0], :])
					pred_reg_loss_mask = reg_loss_mask[..., 1:, :]
					if self.config.motion_state:
						pred_reg_loss_mask = motion_mask  # mask out static object
					loss_loc_2 = F.smooth_l1_loss(result['loc'][..., 1:, :][pred_reg_loss_mask],
												  reg_targets[..., 1:, :][pred_reg_loss_mask])
					loss_loc = loss_loc_1 + loss_loc_2
					loss_num += 2

				###corners for pred
				else:
					loss_loc = self.corner_loss(anchors, reg_loss_mask, reg_targets, result['loc'])
					loss_num += 1
			else:

				loss_loc = F.smooth_l1_loss(result['loc'][reg_loss_mask], reg_targets[reg_loss_mask])
				loss_num += 1

		if self.loss_scale is not None:
			if len(self.loss_scale) == 4:
				loss = self.loss_scale[0] * loss_cls + self.loss_scale[1] * loss_loc_1 + self.loss_scale[
					2] * loss_loc_2 + self.loss_scale[3] * loss_motion
			elif len(self.loss_scale) == 3:
				loss = self.loss_scale[0] * loss_cls + self.loss_scale[1] * loss_loc_1 + self.loss_scale[2] * loss_loc_2
			else:
				loss = self.loss_scale[0] * loss_cls + self.loss_scale[1] * loss_loc
		elif not motion_labels is None:
			loss = loss_cls + loss_loc + loss_motion
		else:
			loss = loss_cls + loss_loc

		if loss_num == 2:
			return (loss_num, loss, loss_cls, loss_loc)
		elif loss_num == 3:
			return (loss_num, loss, loss_cls, loss_loc_1, loss_loc_2)
		elif loss_num == 4:
			return (loss_num, loss, loss_cls, loss_loc_1, loss_loc_2, loss_motion)

	def step(self, data, batch_size):
		bev_seq = data['bev_seq']
		bev_seq_teacher = data['bev_seq_teacher']
		kd_weight = data['kd_weight']
		layer = data['layer']
		self.teacher.eval()

		labels = data['labels']
		reg_targets = data['reg_targets']
		reg_loss_mask = data['reg_loss_mask']
		anchors = data['anchors']
		vis_maps = data['vis_maps']
		trans_matrices = data['trans_matrices']
		num_agent = data['num_agent']

		# with torch.autograd.set_detect_anomaly(True):
		if self.MGDA:
			self.loss_scale = self.cal_loss_scale(data)
			x = self.encoder(bev_seq)
			result = self.head(x)
		else:
			result, x_8, x_7, x_6, x_5, x_3 = self.model(bev_seq, trans_matrices, num_agent, batch_size=batch_size)

		x_8_teacher, x_7_teacher, x_6_teacher, x_5_teacher,x_3_teacher = self.teacher(bev_seq_teacher, vis=vis_maps)
		kl_loss_mean = nn.KLDivLoss(size_average=True, reduce=True)

		# size: x_8: 32*256*256, x_7: 64*128*128, x_8: 128*64*64, x_8: 256*32*32
		# if layer==4:
		# 	kd_loss = kd_weight * torch.sum(torch.norm((x_8.reshape(5 * batch_size, -1) - x_8_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size) + \
		# 			  kd_weight * torch.sum(torch.norm((x_7.reshape(5 * batch_size, -1) - x_7_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size) + \
		# 			  kd_weight * torch.sum(torch.norm((x_6.reshape(5 * batch_size, -1) - x_6_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size) + \
		# 			  kd_weight * torch.sum(torch.norm((x_5.reshape(5 * batch_size, -1) - x_5_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size)
		# elif layer==3:
		# 	kd_loss = kd_weight * torch.sum(torch.norm((x_8.reshape(5 * batch_size, -1) - x_8_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size) + \
		# 			  kd_weight * torch.sum(torch.norm((x_7.reshape(5 * batch_size, -1) - x_7_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size) + \
		# 			  kd_weight * torch.sum(torch.norm((x_6.reshape(5 * batch_size, -1) - x_6_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size) + \
		# 			  kd_weight * torch.sum(torch.norm((x_5.reshape(5 * batch_size, -1) - x_5_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size)

		# elif layer == 2:
		# 	kd_loss = kd_weight * torch.sum(torch.norm((x_8.reshape(5 * batch_size, -1) - x_8_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size) + \
		# 			  kd_weight * torch.sum(torch.norm((x_7.reshape(5 * batch_size, -1) - x_7_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size) + \
		# 			  kd_weight * torch.sum(torch.norm((x_6.reshape(5 * batch_size, -1) - x_6_teacher.reshape(5 * batch_size, -1)), dim=1)) / (5 * batch_size)

		if layer==4:

			target_x8 = x_8_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*256*256, -1)
			student_x8 = x_8.permute(0, 2, 3, 1).reshape(5 *batch_size*256*256, -1)
			kd_loss_x8 = kl_loss_mean(F.log_softmax(student_x8, dim=1), F.softmax(target_x8, dim=1))

			target_x7 = x_7_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*128*128, -1)
			student_x7 = x_7.permute(0, 2, 3, 1).reshape(5 *batch_size*128*128, -1)
			kd_loss_x7 = kl_loss_mean(F.log_softmax(student_x7, dim=1), F.softmax(target_x7, dim=1))

			target_x6 = x_6_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*64*64, -1)
			student_x6 = x_6.permute(0, 2, 3, 1).reshape(5 *batch_size*64*64, -1)
			kd_loss_x6 = kl_loss_mean(F.log_softmax(student_x6, dim=1), F.softmax(target_x6, dim=1))

			target_x5 = x_5_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*32*32, -1)
			student_x5 = x_5.permute(0, 2, 3, 1).reshape(5 *batch_size*32*32, -1)
			kd_loss_x5 = kl_loss_mean(F.log_softmax(student_x5, dim=1), F.softmax(target_x5, dim=1))

			kd_loss = kd_weight * (kd_loss_x8 + kd_loss_x7 + kd_loss_x6 + kd_loss_x5)

		elif layer==3:
			target_x8 = x_8_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*256*256, -1)
			student_x8 = x_8.permute(0, 2, 3, 1).reshape(5 *batch_size*256*256, -1)
			kd_loss_x8 = kl_loss_mean(F.log_softmax(student_x8, dim=1), F.softmax(target_x8, dim=1))

			target_x7 = x_7_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*128*128, -1)
			student_x7 = x_7.permute(0, 2, 3, 1).reshape(5 *batch_size*128*128, -1)
			kd_loss_x7 = kl_loss_mean(F.log_softmax(student_x7, dim=1), F.softmax(target_x7, dim=1))

			target_x6 = x_6_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*64*64, -1)
			student_x6 = x_6.permute(0, 2, 3, 1).reshape(5 *batch_size*64*64, -1)
			kd_loss_x6 = kl_loss_mean(F.log_softmax(student_x6, dim=1), F.softmax(target_x6, dim=1))

			target_x5 = x_5_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*32*32, -1)
			student_x5 = x_5.permute(0, 2, 3, 1).reshape(5 *batch_size*32*32, -1)
			kd_loss_x5 = kl_loss_mean(F.log_softmax(student_x5, dim=1), F.softmax(target_x5, dim=1))

			target_x3 = x_3_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*32*32, -1)
			student_x3 = x_3.permute(0, 2, 3, 1).reshape(5 *batch_size*32*32, -1)
			kd_loss_x3 = kl_loss_mean(F.log_softmax(student_x3, dim=1), F.softmax(target_x3, dim=1))

			kd_loss = kd_weight * (kd_loss_x8 + kd_loss_x7 + kd_loss_x6 + kd_loss_x5 + kd_loss_x3)

		elif layer==2:
			target_x8 = x_8_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*256*256, -1)
			student_x8 = x_8.permute(0, 2, 3, 1).reshape(5 *batch_size*256*256, -1)
			kd_loss_x8 = kl_loss_mean(F.log_softmax(student_x8, dim=1), F.softmax(target_x8, dim=1))

			target_x7 = x_7_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*128*128, -1)
			student_x7 = x_7.permute(0, 2, 3, 1).reshape(5 *batch_size*128*128, -1)
			kd_loss_x7 = kl_loss_mean(F.log_softmax(student_x7, dim=1), F.softmax(target_x7, dim=1))

			target_x6 = x_6_teacher.permute(0, 2, 3, 1).reshape(5 *batch_size*64*64, -1)
			student_x6 = x_6.permute(0, 2, 3, 1).reshape(5 *batch_size*64*64, -1)
			kd_loss_x6 = kl_loss_mean(F.log_softmax(student_x6, dim=1), F.softmax(target_x6, dim=1))

			kd_loss = kd_weight * (kd_loss_x8 + kd_loss_x7 + kd_loss_x6)

		labels = labels.view(result['cls'].shape[0], -1, result['cls'].shape[-1])
		N = bev_seq.shape[0]
		loss_collect = self.loss_calculator(result, anchors, reg_loss_mask, reg_targets, labels, N)

		# -------- for debugging teacher model---------#
		# loss_collect_teacher = self.loss_calculator(result_teacher,anchors,reg_loss_mask,reg_targets,labels,N)
		# loss_num, loss, loss_cls, loss_loc = loss_collect_teacher
		# print(loss, loss_cls, loss_loc)

		loss_num = loss_collect[0]
		if loss_num == 3:
			loss_num, loss, loss_cls, loss_loc_1, loss_loc_2 = loss_collect
		elif loss_num == 2:
			loss_num, loss, loss_cls, loss_loc = loss_collect
		elif loss_num == 4:
			loss_num, loss, loss_cls, loss_loc_1, loss_loc_2, loss_motion = loss_collect

		loss = loss + kd_loss
		print(kd_loss.item())

		if self.MGDA:
			self.optimizer_encoder.zero_grad()
			self.optimizer_head.zero_grad()
			loss.backward()
			self.optimizer_encoder.step()
			self.optimizer_head.step()
		else:
			self.optimizer.zero_grad()
			loss.backward()
			self.optimizer.step()

		if self.config.pred_type in ['motion', 'center'] and not self.only_det:
			if self.config.motion_state:
				return loss.item(), loss_cls.item(), loss_loc_1.item(), loss_loc_2.item(), loss_motion.item()
			else:
				return loss.item(), loss_cls.item(), loss_loc_1.item(), loss_loc_2.item()
		else:
			return loss.item(), loss_cls.item(), loss_loc.item(), kd_loss.item()

	def predict(self, data, validation=True):

		bev_seq = data['bev_seq']
		vis_maps = data['vis_maps']

		if self.MGDA:
			x = self.encoder(bev_seq)
			result = self.head(x)
		else:
			result = self.model(bev_seq, vis=vis_maps)

		N = bev_seq.shape[0]

		if validation:
			labels = data['labels']
			anchors = data['anchors']
			reg_targets = data['reg_targets']
			reg_loss_mask = data['reg_loss_mask']
			motion_labels = None
			motion_mask = None

			labels = labels.view(result['cls'].shape[0], -1, result['cls'].shape[-1])
			if self.config.motion_state:
				motion_labels = data['motion_label']
				motion_mask = data['motion_mask']
				motion_labels = motion_labels.view(result['state'].shape[0], -1, result['state'].shape[-1])
			N = bev_seq.shape[0]

			loss_collect = self.loss_calculator(result, anchors, reg_loss_mask, reg_targets, labels, N, motion_labels,
												motion_mask)
			loss_num = loss_collect[0]
			if loss_num == 3:
				loss_num, loss, loss_cls, loss_loc_1, loss_loc_2 = loss_collect
			elif loss_num == 2:
				loss_num, loss, loss_cls, loss_loc = loss_collect
			elif loss_num == 4:
				loss_num, loss, loss_cls, loss_loc_1, loss_loc_2, loss_motion = loss_collect

			batch_box_preds = result['loc']
			batch_cls_preds = result['cls']

			if self.config.motion_state:
				batch_motion_preds = result['state']
			else:
				batch_motion_preds = None

			if not self.only_det:
				if self.config.pred_type == 'center':
					batch_box_preds[:, :, :, :, 1:, 2:] = batch_box_preds[:, :, :, :, [0], 2:]

		class_selected = apply_nms_det(batch_box_preds, batch_cls_preds, anchors, self.code_type, self.config,
									   batch_motion_preds)
		# class_selected = None
		if validation:
			if self.config.pred_type in ['motion', 'center'] and not self.only_det:
				if self.config.motion_state:
					return loss.item(), loss_cls.item(), loss_loc_1.item(), loss_loc_2.item(), loss_motion.item(), class_selected
				else:
					return loss.item(), loss_cls.item(), loss_loc_1.item(), loss_loc_2.item(), class_selected
			else:
				return loss.item(), loss_cls.item(), loss_loc.item(), class_selected
		else:
			return class_selected

	def predict_all(self, data, batch_size, validation=True):
		NUM_AGENT = 5
		bev_seq = data['bev_seq']
		vis_maps = data['vis_maps']
		trans_matrices = data['trans_matrices']
		num_agent_tensor = data['num_agent']
		num_sensor = num_agent_tensor[0, 0]

		if self.MGDA:
			x = self.encoder(bev_seq)
			result = self.head(x)
		else:
			with torch.no_grad():
				result, x_8, x_7, x_6, x_5, x_3 = self.model(bev_seq, trans_matrices, num_agent_tensor, batch_size=batch_size)
		# result = self.model(bev_seq,vis=vis_maps,training=False)
		#
		N = bev_seq.shape[0]

		if validation:
			labels = data['labels']
			anchors = data['anchors']
			reg_targets = data['reg_targets']
			reg_loss_mask = data['reg_loss_mask']
			motion_labels = None
			motion_mask = None

			labels = labels.view(result['cls'].shape[0], -1, result['cls'].shape[-1])

			if self.config.motion_state:
				motion_labels = data['motion_label']
				motion_mask = data['motion_mask']
				motion_labels = motion_labels.view(result['state'].shape[0], -1, result['state'].shape[-1])
			N = bev_seq.shape[0]

			loss_collect = self.loss_calculator(result, anchors, reg_loss_mask, reg_targets, labels, N, motion_labels,
												motion_mask)
			loss_num = loss_collect[0]
			if loss_num == 3:
				loss_num, loss, loss_cls, loss_loc_1, loss_loc_2 = loss_collect
			elif loss_num == 2:
				loss_num, loss, loss_cls, loss_loc = loss_collect
			elif loss_num == 4:
				loss_num, loss, loss_cls, loss_loc_1, loss_loc_2, loss_motion = loss_collect

		seq_results = [[] for i in range(NUM_AGENT)]
		# global_points = [[] for i in range(num_sensor)]
		# cls_preds = [[] for i in range(num_sensor)]

		for k in range(NUM_AGENT):
			bev_seq = torch.unsqueeze(data['bev_seq'][k, :, :, :, :], 0)

			if torch.nonzero(bev_seq).shape[0] == 0:
				seq_results[k] = []
			else:
				batch_box_preds = torch.unsqueeze(result['loc'][k, :, :, :, :, :], 0)
				batch_cls_preds = torch.unsqueeze(result['cls'][k, :, :], 0)
				anchors = torch.unsqueeze(data['anchors'][k, :, :, :, :], 0)
				batch_motion_preds = None

				if not self.only_det:
					if self.config.pred_type == 'center':
						batch_box_preds[:, :, :, :, 1:, 2:] = batch_box_preds[:, :, :, :, [0], 2:]

				class_selected = apply_nms_det(batch_box_preds, batch_cls_preds, anchors, self.code_type, self.config,
											   batch_motion_preds)
				seq_results[k] = class_selected

		# global_points[k], cls_preds[k] = apply_box_global_transform(trans_matrices_map[k],batch_box_preds,batch_cls_preds,anchors,self.code_type,self.config,batch_motion_preds)

		# all_points_scene = numpy.concatenate(tuple(global_points), 0)
		# cls_preds_scene = torch.cat(tuple(cls_preds), 0)
		# class_selected_global = apply_nms_global_scene(all_points_scene, cls_preds_scene)

		if validation:
			return loss.item(), loss_cls.item(), loss_loc.item(), seq_results
		else:
			return seq_results

	def predict_all_with_box_com(self, data, trans_matrices_map, validation=True):
		NUM_AGENT = 5
		bev_seq = data['bev_seq']
		vis_maps = data['vis_maps']
		trans_matrices = data['trans_matrices']
		num_agent_tensor = data['num_agent']
		num_sensor = num_agent_tensor[0, 0]

		if self.MGDA:
			x = self.encoder(bev_seq)
			result = self.head(x)
		else:
			result = self.model(bev_seq, trans_matrices, num_agent_tensor, batch_size=1)

		N = bev_seq.shape[0]

		if validation:
			labels = data['labels']
			anchors = data['anchors']
			reg_targets = data['reg_targets']
			reg_loss_mask = data['reg_loss_mask']
			motion_labels = None
			motion_mask = None

			labels = labels.view(result['cls'].shape[0], -1, result['cls'].shape[-1])

			if self.config.motion_state:
				motion_labels = data['motion_label']
				motion_mask = data['motion_mask']
				motion_labels = motion_labels.view(result['state'].shape[0], -1, result['state'].shape[-1])
			N = bev_seq.shape[0]

			loss_collect = self.loss_calculator(result, anchors, reg_loss_mask, reg_targets, labels, N, motion_labels,
												motion_mask)
			loss_num = loss_collect[0]
			if loss_num == 3:
				loss_num, loss, loss_cls, loss_loc_1, loss_loc_2 = loss_collect
			elif loss_num == 2:
				loss_num, loss, loss_cls, loss_loc = loss_collect
			elif loss_num == 4:
				loss_num, loss, loss_cls, loss_loc_1, loss_loc_2, loss_motion = loss_collect

		seq_results = [[] for i in range(NUM_AGENT)]
		local_results_wo_local_nms = [[] for i in range(NUM_AGENT)]
		local_results_af_local_nms = [[] for i in range(NUM_AGENT)]

		global_points = [[] for i in range(num_sensor)]
		cls_preds = [[] for i in range(num_sensor)]
		global_boxes_af_localnms = [[] for i in range(num_sensor)]
		box_scores_af_localnms = [[] for i in range(num_sensor)]

		forward_message_size = 0
		forward_message_size_two_nms = 0

		for k in range(NUM_AGENT):
			bev_seq = torch.unsqueeze(data['bev_seq'][k, :, :, :, :], 0)

			if torch.nonzero(bev_seq).shape[0] == 0:
				seq_results[k] = []
			else:
				batch_box_preds = torch.unsqueeze(result['loc'][k, :, :, :, :, :], 0)
				batch_cls_preds = torch.unsqueeze(result['cls'][k, :, :], 0)
				anchors = torch.unsqueeze(data['anchors'][k, :, :, :, :], 0)

				if self.config.motion_state:
					batch_motion_preds = result['state']
				else:
					batch_motion_preds = None

				if not self.only_det:
					if self.config.pred_type == 'center':
						batch_box_preds[:, :, :, :, 1:, 2:] = batch_box_preds[:, :, :, :, [0], 2:]

				class_selected, box_scores_pred_cls = apply_nms_det(batch_box_preds, batch_cls_preds, anchors,
																	self.code_type, self.config, batch_motion_preds)

				# transform all the boxes before local nms to the global coordinate
				# global_points[k], cls_preds[k] = apply_box_global_transform(trans_matrices_map[k], batch_box_preds,
				#                                                            batch_cls_preds, anchors, self.code_type,
				#                                                            self.config, batch_motion_preds)

				# transform the boxes after local nms to the global coordinate
				global_boxes_af_localnms[k], box_scores_af_localnms[k] = apply_box_global_transform_af_localnms(
					trans_matrices_map[k], class_selected, box_scores_pred_cls)
				# print(cls_preds[k].shape, box_scores_af_localnms[k].shape)

				forward_message_size = forward_message_size + 256 * 256 * 6 * 4 * 2
				forward_message_size_two_nms = forward_message_size_two_nms + global_boxes_af_localnms[k].shape[
					0] * 4 * 2

		# global results with one NMS
		# all_points_scene = numpy.concatenate(tuple(global_points), 0)
		# cls_preds_scene = torch.cat(tuple(cls_preds), 0)
		# class_selected_global = apply_nms_global_scene(all_points_scene, cls_preds_scene)

		# global results with two NMS
		global_boxes_af_local_nms = numpy.concatenate(tuple(global_boxes_af_localnms), 0)
		box_scores_af_local_nms = torch.cat(tuple(box_scores_af_localnms), 0)
		class_selected_global_af_local_nms = apply_nms_global_scene(global_boxes_af_local_nms, box_scores_af_local_nms)

		# transform the consensus global boxes to local agents (two NMS)
		back_message_size_two_nms = 0
		for k in range(num_sensor):
			local_results_af_local_nms[k], ms = apply_box_local_transform(class_selected_global_af_local_nms,
																		  trans_matrices_map[k])
			back_message_size_two_nms = back_message_size_two_nms + ms

		sample_bandwidth_two_nms = forward_message_size_two_nms + back_message_size_two_nms

		# transform the consensus global boxes to local agents (One NMS)
		# back_message_size = 0
		# for k in range(num_sensor):
		#    local_results_wo_local_nms[k], ms = apply_box_local_transform(class_selected_global, trans_matrices_map[k])
		#    back_message_size = back_message_size + ms

		# sample_bandwidth = forward_message_size + back_message_size

		return loss.item(), loss_cls.item(), loss_loc.item(), local_results_af_local_nms, class_selected_global_af_local_nms, sample_bandwidth_two_nms

	def cal_loss_scale(self, data):
		bev_seq = data['bev_seq']
		labels = data['labels']
		reg_targets = data['reg_targets']
		reg_loss_mask = data['reg_loss_mask']
		anchors = data['anchors']
		motion_labels = None
		motion_mask = None

		with torch.no_grad():
			shared_feats = self.encoder(bev_seq)
		shared_feats_tensor = shared_feats.clone().detach().requires_grad_(True)
		result = self.head(shared_feats_tensor)
		if self.config.motion_state:
			motion_labels = data['motion_label']
			motion_mask = data['motion_mask']
			motion_labels = motion_labels.view(result['state'].shape[0], -1, result['state'].shape[-1])
		self.optimizer_encoder.zero_grad()
		self.optimizer_head.zero_grad()
		grads = {}
		labels = labels.view(result['cls'].shape[0], -1, result['cls'].shape[-1])
		N = bev_seq.shape[0]

		# calculate loss
		grad_len = 0

		'''
		Classification Loss
		'''
		loss_cls = self.alpha * torch.sum(self.criterion['cls'](result['cls'], labels)) / N
		# loss_loc = torch.sum(self.criterion['loc'](result['loc'],reg_targets,mask = reg_loss_mask)) / N
		self.optimizer_encoder.zero_grad()
		self.optimizer_head.zero_grad()

		loss_cls.backward(retain_graph=True)
		grads[0] = []
		grads[0].append(shared_feats_tensor.grad.data.clone().detach())
		shared_feats_tensor.grad.data.zero_()
		grad_len += 1

		'''
		Localization Loss
		'''
		loc_scale = False
		loss_mask_num = torch.nonzero(reg_loss_mask.view(-1, reg_loss_mask.shape[-1])).size(0)

		if self.code_type in ['corner_1', 'corner_2', 'corner_3']:
			target = reg_targets[reg_loss_mask].reshape(-1, 5, 2)
			flip_target = torch.stack([target[:, 0], target[:, 3], target[:, 4], target[:, 1], target[:, 2]], dim=-2)
			pred = result['loc'][reg_loss_mask].reshape(-1, 5, 2)
			t = torch.sum(torch.norm(pred - target, dim=-1), dim=-1)
			f = torch.sum(torch.norm(pred - flip_target, dim=-1), dim=-1)
			loss_loc = torch.sum(torch.min(t, f)) / N

		elif self.code_type == 'faf':
			if self.loss_type == 'corner_loss':
				if self.only_det:
					loss_loc = self.corner_loss(anchors, reg_loss_mask, reg_targets, result['loc'])
				elif self.config.pred_type in ['motion', 'center']:

					###only center/motion for pred

					loss_loc_1 = self.corner_loss(anchors, reg_loss_mask[..., 0][..., [0]], reg_targets[..., [0], :],
												  result['loc'][..., [0], :])
					pred_reg_loss_mask = reg_loss_mask[..., 1:, :]
					if self.config.motion_state:
						pred_reg_loss_mask = motion_mask  # mask out static object
					loss_loc_2 = F.smooth_l1_loss(result['loc'][..., 1:, :][pred_reg_loss_mask],
												  reg_targets[..., 1:, :][pred_reg_loss_mask])

					self.optimizer_encoder.zero_grad()
					self.optimizer_head.zero_grad()

					loss_loc_1.backward(retain_graph=True)
					grads[1] = []
					grads[1].append(shared_feats_tensor.grad.data.clone().detach())
					shared_feats_tensor.grad.data.zero_()

					self.optimizer_encoder.zero_grad()
					self.optimizer_head.zero_grad()

					loss_loc_2.backward(retain_graph=True)
					grads[2] = []
					grads[2].append(shared_feats_tensor.grad.data.clone().detach())
					shared_feats_tensor.grad.data.zero_()
					loc_scale = True
					grad_len += 2

				###corners for pred
				else:
					loss_loc = self.corner_loss(anchors, reg_loss_mask, reg_targets, result['loc'])
			else:

				loss_loc = F.smooth_l1_loss(result['loc'][reg_loss_mask], reg_targets[reg_loss_mask])

			if not loc_scale:
				grad_len += 1
				self.optimizer_encoder.zero_grad()
				self.optimizer_head.zero_grad()
				loss_loc.backward(retain_graph=True)
				grads[1] = []
				grads[1].append(shared_feats_tensor.grad.data.clone().detach())
				shared_feats_tensor.grad.data.zero_()

		'''
		Motion state Loss
		'''
		if self.config.motion_state:
			loss_motion = torch.sum(self.criterion['cls'](result['state'], motion_labels)) / N

			self.optimizer_encoder.zero_grad()
			self.optimizer_head.zero_grad()

			loss_motion.backward(retain_graph=True)
			grads[3] = []
			grads[3].append(shared_feats_tensor.grad.data.clone().detach())
			shared_feats_tensor.grad.data.zero_()
			grad_len += 1

		# ---------------------------------------------------------------------
		# -- Frank-Wolfe iteration to compute scales.
		scale = np.zeros(grad_len, dtype=np.float32)
		sol, min_norm = MinNormSolver.find_min_norm_element([grads[t] for t in range(grad_len)])
		for i in range(grad_len):
			scale[i] = float(sol[i])

		# print(scale)
		return scale

max_stars_repo_name: Chezacar/CollaborationWithLatency

max_issues_repo_name: Chezacar/CollaborationWithLatency

max_forks_repo_name: Chezacar/CollaborationWithLatency


text:
#!/usr/bin/env python
"""
find contours/curves in an image, a contour is a curve of points, with no
gaps in the curve. Contours are extremely useful for such
things as shape approximation and analysis.

Parameters

Usage

Example
    $ python <scriptname>.py --image ../img/<filename>.png

## Explain

"""
import numpy as np
import argparse
import cv2

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("-i", "--image", required=True, help="Path to the image")
    args = vars(ap.parse_args())

    image = cv2.imread(args["image"])
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (11,11), 0)
    cv2.imshow("Image", image)

    edged = cv2.Canny(blurred, 30, 150)
    cv2.imshow("Edges", edged)

    (cnts, _) = cv2.findContours(edged.copy(),
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE
        )
    # image is destroyed, therefore .copy method
    # returns a 2-tuple, image, contours and hierarchy of contours
    # param1: type of contours - retrieve only the outermost contours/EXTERNAL
        # RETR_LIST will return all contours; RETR_COMP and RETR_TREE can return hierarchy
    # param2: how to approximate contours; simple compress horizontal, vertical and diagonal - save computation and memory
    print("[INFO] Detected {} coins in this image".format(len(cnts)))

    coins = image.copy()
    cv2.drawContours(coins, cnts, -1, (0, 255, 0), 2)
    # param2: list of contours
    # param3: which contours to draw, -1 all
    # param4: color of contour
    # param5: thickness
    cv2.imshow("Coins", coins)
    cv2.waitKey(0)

    # crop individual coins
    for (i, c) in enumerate(cnts):
        (x, y, w, h) = cv2.boundingRect(c)# finds enclosing region of contour

        print("[INFO] #{}".format(i + 1))
        coin = image[y:y + h, x:x + w]# slice/crop region
        cv2.imshow("Coin", coin)

        mask = np.zeros(image.shape[:2], dtype="uint8")
        ((centerX, centerY), radius) = cv2.minEnclosingCircle(c)# fits circle to contour

        cv2.circle(mask, (int(centerX), int(centerY)), int(radius), 255, -1)# create mask
        mask = mask[y:y + h, x:x + w]# slice/crop mask

        cv2.imshow("Masked coin", cv2.bitwise_and(coin, coin, mask = mask))
        # remove background
        cv2.waitKey(0)

if __name__=="__main__":
    main()

ext: py

max_stars_repo_name: knielbo/cds-visual

max_issues_repo_name: knielbo/cds-visual

max_forks_repo_name: knielbo/cds-visual

reason: import numpy


text:
# coding=utf-8
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import torch
import os
import pickle
import torchvision.models as models

import math
import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
# a =  torch.randn((1,2,9,9))
# print a
#
# b = torch.nn.functional.adaptive_max_pool2d(a, 3)
# print b
# b = Variable(torch.ones((3,2)), requires_grad=False)
# b[b<0] = 0
# print(b)
# model = torch.load('/home/rgh/PycharmProjects/pytorch-faster-rcnn/data/imagenet_weights/vgg16-00b39a1b.pth')
# print(model.keys())
# model = models.vgg16()
#
# b = torch.ones((3,2))
# c = b[:, 0]
# # c = b.new(b.size(0), 1) #.zero_()
# print(b)
# print(c.size())
# clss = torch.from_numpy(np.array([0,1,2]))
# bbox_targets = clss.new(clss.numel(), 4 * 3).zero_()
# bbox_target_data = torch.from_numpy(np.array([[-1,-2,-3,-4],[-1,-2,-3,-4],[-1,-2,-3,-4]]))
# inds = torch.from_numpy(np.array([1,2]))
#
# clss = clss[inds].contiguous().view(-1,1)  # (n, 1)
# print(clss)
# dim1_inds = inds.unsqueeze(1).expand(inds.size(0), 4)  # (n, 4) [ [1,1,1,1],[3,3,3,3],[4,4,4,4] ]
# print(dim1_inds)
# dim2_inds = torch.cat([4*clss, 4*clss+1, 4*clss+2, 4*clss+3], 1).long()
# print(dim2_inds)
# bbox_targets[dim1_inds, dim2_inds] = bbox_target_data[inds][:, 0:]
#
# scores = np.array([[1,2],[3,4],[5,6]])
# mask_target_channel = scores.argmax(1)
# print(mask_target_channel.shape)
# img = np.array([0.2,0.5])
# cv2.imwrite('/media/rgh/rgh-data/PycharmProjects/cvpr2018/temp/mask/test.png',img)
# result = np.zeros((3,3))
#
# mask_a = np.array([[1,0],[3,4]])
# index = result[0:2,0:2]==0
# result[0:2,0:2][index] = mask_a[index]
#
# mask_b = np.array([[100,101],[102,103]])
# index = result[0:2,0:2]==0
# result[0:2,0:2][index] = mask_b[index]
# print(result)
# print(bbox_targets)
# a = torch.from_numpy(np.array([[1,2,3],[4,8,6]]))
# b = a[0:1,:]
# print(b)
# b = a.new(2, 1).zero_()
# for i in range(2):
#     b[i,:] = a[i,1]
# b[0][0] = 10
# print(a)
# x = np.ones((3,2))
# y = np.ones(3)
# print(x+y)
# overlaps = torch.from_numpy(np.array([[1,2,3],[4,8,6]]))
# gt_boxes = torch.from_numpy(np.array([[0,0,0,0,-1],[0,0,0,0,-2],[0,0,0,0,-3]]))
# max_overlaps, gt_assignment = overlaps.max(1)
# print(max_overlaps,gt_assignment)
# labels = gt_boxes[gt_assignment, [4]]
# print(labels)
# for key, value in dict(model.named_parameters()).items():
#     print(key)
# obj2 = pickle.load(open(os.path.join('/home/rgh/PycharmProjects/pytorch-faster-rcnn/output/vgg16/lip_val/global_cat/vgg16_faster_rcnn_iter_70000',
#                        'hair_pr.pkl'), 'rb'))
# print(obj2)

# a = torch.ones((256,512,37,39))
#
# b = torch.zeros((256,37,39))
# print(b.size())
# a = Variable(a)
# print(a.size())
# b = Variable(b)
# print(a*b)
# all_boxes = [[[3] for _ in range(2)]
#          for _ in range(3)]
# all_boxes[0][0] = 10
# print(all_boxes)
# print(all_boxes[0][0][1])
# img = cv2.imread('/home/rgh/Pictures/temp/t.jpg')
# img = np.array(img)
# img = np.array([img])
# b = torch.from_numpy(img)
# b = b.float()
# b = b.permute(0, 3, 1, 2)
# print(b.size())
# # b has the size (1, 3, 360, 640)
# flow = torch.zeros(1, 420, 1002 , 2)
# b = Variable(b)
# flow = Variable(flow)
# out = F.grid_sample(b, flow)
# print(out.data.shape)
# img = out.data.numpy()
# print(img.shape)
# img = img.squeeze().transpose(1,2,0)
# cv2.imshow('g',img)
# cv2.waitKey()

# grads = {}
# def save_grad(name):
#     def hook(grad):
#         grads[name] = grad
#     return hook

# x = Variable(torch.ones((2,2)), requires_grad=True)
#
# z = x * 2
# z[0,0] = 0
# z.register_hook(save_grad('z'))
# t = z+2
#
# t.backward(torch.Tensor([[10]]))
# print(x.grad,grads['z'])

# x = Variable(torch.ones(2, 2), requires_grad = True)
# w = Variable(torch.ones(2, 2), requires_grad = True)
# z = x*w -100
# loss = z.sum()
# loss.backward()  # loss.backward(torch.Tensor([1]))
#
# print(x.grad)
# print(w.grad)
# grads = {}
# def save_grad(name):
#     def hook(grad):
#         grads[name] = grad
#     return hook
#
# x = Variable(torch.ones((2,2)), requires_grad=True)
# y = Variable(torch.ones((1,2)), requires_grad=True)
# t = Variable(torch.ones((1,2)), requires_grad=True)
# t = x[0,:] + y
# t.register_hook(save_grad('t'))
# z = torch.cat((y, t),1)
# z = torch.cat((y, z),1)
# z.backward(torch.Tensor([[1, 2, 3, 4, 1,1]]))
# print(z, x.grad, y.grad, t.grad, grads['t'])
#
# z = x[0,0]
# z.backward(torch.Tensor([[10]]))
# print(z, x.grad, y.grad, t.grad)
# a = torch.ones((3,2))
# a[0,0] = 3
# print(a==3)

# a = np.array([[1,2],[3,4]])
# b = np.zeros((2,2,3))
# mask = a == 1
# c = b[mask]
# c[0][0] = -100
# print (c)
# print (b)
# d = a[mask]
# d [0] = -100
# print (d)
# print (a)
# a = np.array([[1,2],[3,4]])
# mask = a == 1
# b = a[mask]
# b = -100
# print (b)
# print (a)
# a[mask] = -100
# print (a)
# img_names = open('/media/rgh/rgh-data/Dataset/Lip_t_d_zdf/train_full.txt', 'r')
# total = []
# for name in img_names:
#     name = name.strip()
#     if name not in total:
#         total.append(name)
#     else:
#         print(name)
# print (len(total))
class Vars(object):

    def __init__(self):
        self.count = 0
        self.defs = {}
        self.lookup = {}

    def add(self, *v):
        name = self.lookup.get(v, None)
        print (v)
        if name is None:
            if v[0] == '+':
                if v[1] == 0:
                    return v[2]
                elif v[2] == 0:
                    return v[1]
            elif v[0] == '*':
                if v[1] == 1:
                    return v[2]
                elif v[2] == 1:
                    return v[1]
                elif v[1] == 0:
                    return 0
                elif v[2] == 0:
                    return 0

            self.count += 1
            name = "v" + str(self.count)
            self.defs[name] = v
            self.lookup[v] = name

        return name

    def __getitem__(self, name):
        return self.defs[name]

    def __iter__(self):
        return self.defs.iteritems()

def diff(vars, acc, v, w):
    if v == w:
        return acc

    v = vars[v]
    if v[0] == 'input':
        return 0
    elif v[0] == "sin":
        return diff(vars, vars.add("*", acc, vars.add("cos", v[1])), v[1], w)
    elif v[0] == '+':
        gx = diff(vars, acc, v[1], w)
        gy = diff(vars, acc, v[2], w)
        return vars.add("+", gx, gy)
    elif v[0] == '*':
        gx = diff(vars, vars.add("*", v[2], acc), v[1], w)
        gy = diff(vars, vars.add("*", v[1], acc), v[2], w)
        return vars.add("+", gx, gy)

    raise NotImplementedError

def autodiff(vars, v, *wrt):
    return tuple(diff(vars, 1, v, w) for w in wrt)

# z = (sin x) + (x * y)

vars = Vars()

x = vars.add("input",1)
#a= raw_input()
y = vars.add("input",2)
z = vars.add("+", vars.add("*",x,y),vars.add("sin",x))

a= raw_input()
print (autodiff(vars, z, x, y))

for k, v in vars:
    print (k, v)

max_stars_repo_name: Sundrops/pytorch-faster-rcnn

max_issues_repo_name: Sundrops/pytorch-faster-rcnn

max_forks_repo_name: Sundrops/pytorch-faster-rcnn


text:
import cocos.device
import cocos.numerics as cn
import numpy as np
import pytest

test_data = [np.array([[1, 2, 3],
                       [4, 5, 6],
                       [7, 8, 20]],
                      dtype=np.int32),
             np.array([[0.2, 1.0, 0.5],
                       [0.4, 0.5, 0.6],
                       [0.7, 0.2, 0.25]],
                      dtype=np.float32),
             np.array([[0.5, 2.3, 3.1],
                       [4, 5.5, 6],
                       [7 - 9j, 8 + 1j, 2 + 10j]],
                      dtype=np.complex64),
             np.array([[[1.0, 2], [3, 4]],
                       [[5, 6], [7, 8]]],
                      dtype=np.float32)]

test_axes = (0, 1)
test_offset = (-2, -1, 0, 1, 2)

@pytest.mark.parametrize("A_numpy", test_data)
@pytest.mark.parametrize("axis", test_axes)
@pytest.mark.parametrize("shift", test_offset)
def test_roll(A_numpy, axis, shift):
    print(f'axis={axis}, shift={shift}')
    cocos.device.init()
    A_cocos = cn.array(A_numpy)

    B_numpy = np.roll(A_numpy, shift=shift, axis=axis)
    B_cocos = cn.roll(A_cocos, shift=shift, axis=axis)
    assert np.allclose(B_cocos, B_numpy)

max_stars_repo_name: michaelnowotny/cocos

max_issues_repo_name: michaelnowotny/cocos

max_forks_repo_name: michaelnowotny/cocos


text:
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Test code for pooling"""
import numpy as np
import pytest
import sys

import tvm
from tvm import topi
from tvm import te
import tvm.topi.testing
from tvm.topi.utils import get_const_tuple

from ..conftest import requires_hexagon_toolchain

class TestAdaptivePool:
    dshape, out_size, pool_type, layout = tvm.testing.parameters(
        ((1, 3, 112, 112), (1, 1), "max", "NCHW"),
        ((1, 3, 112, 112), (1, 1), "avg", "NCHW"),
        ((1, 14, 56, 78), (34, 13), "max", "NCHW"),
        ((1, 5, 46, 97), (4, 96), "avg", "NCHW"),
        ((1, 112, 112, 3), (1, 1), "max", "NHWC"),
        ((1, 5, 46, 97), (4, 96), "avg", "NHWC"),
        ((1, 16, 32, 32, 32), (1, 1, 1), "max", "NCDHW"),
        ((1, 16, 32, 32, 32), (1, 1, 1), "avg", "NCDHW"),
        ((1, 16, 32, 32, 32), (2, 2, 2), "avg", "NCDHW"),
        (
            (1, 16, 64, 32, 32),
            (7, 8, 9),
            "avg",
            "NCDHW",
        ),
        (
            (1, 16, 64, 32, 32),
            (8, 16, 16),
            "avg",
            "NCDHW",
        ),
        ((1, 16, 32, 32, 32), (1, 1, 1), "avg", "NDHWC"),
        ((1, 16, 32, 32, 32), (2, 2, 2), "max", "NDHWC"),
        ((1, 16, 32, 32, 32), (2, 4, 4), "max", "NDHWC"),
    )

    @requires_hexagon_toolchain
    def test_adaptive_pool(self, hexagon_session, dshape, out_size, pool_type, layout):
        dtype = "float32"
        np_data = np.random.uniform(low=0, high=255, size=dshape).astype(dtype)
        np_out = tvm.topi.testing.adaptive_pool(np_data, out_size, pool_type, layout)
        oshape = np_out.shape

        data = te.placeholder(dshape, name="data", dtype=dtype)
        if len(out_size) == 2:
            out = topi.nn.adaptive_pool(data, out_size, pool_type, layout)
        else:
            assert len(out_size) == 3
            out = topi.nn.adaptive_pool3d(data, out_size, pool_type, layout)

        target_hexagon = tvm.target.hexagon("v68")
        with tvm.target.Target(target_hexagon):
            fschedule = topi.hexagon.schedule_adaptive_pool
            s = fschedule(out)

        func = tvm.build(
            s,
            [data, out],
            tvm.target.Target(target_hexagon, host=target_hexagon),
            name="adaptive-pool",
        )
        mod = hexagon_session.load_module(func)

        dev = hexagon_session.device
        a = tvm.nd.array(np_data, dev)
        b = tvm.nd.array(np.zeros(get_const_tuple(oshape), dtype=out.dtype), dev)
        mod["adaptive-pool"](a, b)

        tvm.testing.assert_allclose(b.numpy(), np_out, rtol=4e-5, atol=1e-6)

def verify_poolnd(
    hexagon_session,
    n,
    input_shape,
    kernel,
    stride,
    dilation,
    padding,
    pool_type,
    ceil_mode,
    count_include_pad=True,
    layout="NCW",
):
    A = te.placeholder(input_shape, name="A")

    if n == 1:
        B = topi.nn.pool1d(
            A,
            kernel=kernel,
            stride=stride,
            dilation=dilation,
            padding=padding,
            pool_type=pool_type,
            ceil_mode=ceil_mode,
            layout=layout,
            count_include_pad=count_include_pad,
        )
    elif n == 2:
        B = topi.nn.pool2d(
            A,
            kernel=kernel,
            stride=stride,
            dilation=dilation,
            padding=padding,
            pool_type=pool_type,
            ceil_mode=ceil_mode,
            layout=layout,
            count_include_pad=count_include_pad,
        )
    elif n == 3:
        B = topi.nn.pool3d(
            A,
            kernel=kernel,
            stride=stride,
            dilation=dilation,
            padding=padding,
            pool_type=pool_type,
            ceil_mode=ceil_mode,
            layout=layout,
            count_include_pad=count_include_pad,
        )
    else:
        raise ValueError(f"PoolND only supports n=1, 2, 3 got n={n}")

    B = topi.nn.relu(B)
    dtype = A.dtype
    output_shape = [int(i) for i in B.shape]

    input_np = np.random.uniform(low=0.001, size=input_shape).astype(dtype)

    padding_before = padding[:n]
    padding_after = padding[n:]
    ref_np = tvm.topi.testing.poolnd_python(
        input_np,
        kernel,
        stride,
        dilation,
        padding_before,
        padding_after,
        pool_type,
        count_include_pad,
        ceil_mode,
        layout=layout,
    )

    np.testing.assert_equal(tuple(output_shape), tuple(ref_np.shape))

    target_hexagon = tvm.target.hexagon("v68")
    with tvm.target.Target(target_hexagon):
        fschedule = topi.hexagon.schedule_pool
        s = fschedule(B, layout)

    func = tvm.build(s, [A, B], tvm.target.Target(target_hexagon, host=target_hexagon), name="pool")
    mod = hexagon_session.load_module(func)

    dev = hexagon_session.device
    a = tvm.nd.array(input_np, dev)
    b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=dtype), dev)
    mod["pool"](a, b)

    tvm.testing.assert_allclose(b.numpy(), ref_np, rtol=1e-5)

class TestPool1D:
    (
        input_shape,
        kernel,
        stride,
        dilation,
        padding,
        pool_type,
        ceil_mode,
        count_include_pad,
        layout,
    ) = tvm.testing.parameters(
        ([1, 16, 32], [2], [2], [1], [0, 0], "avg", False, True, "NCW"),
        ([1, 16, 31], [3], [3], [1], [1, 2], "avg", False, True, "NCW"),
        ([1, 16, 32], [2], [2], [1], [1, 2], "avg", False, False, "NCW"),
        ([1, 16, 31], [4], [4], [1], [3, 3], "avg", False, False, "NCW"),
        ([1, 16, 31], [4], [4], [1], [0, 0], "avg", False, False, "NCW"),
        ([1, 16, 32], [2], [2], [1], [0, 0], "max", False, True, "NCW"),
        ([1, 16, 31], [3], [3], [1], [2, 1], "max", False, True, "NCW"),
        ([1, 16, 31], [3], [3], [1], [2, 1], "max", True, True, "NCW"),
        ([1, 16, 31], [3], [3], [1], [2, 5], "avg", False, True, "NCW"),
        ([1, 16, 32], [2], [2], [1], [0, 3], "avg", False, False, "NCW"),
        ([1, 16, 31], [3], [3], [1], [1, 4], "max", False, True, "NCW"),
        ([1, 16, 31], [3], [3], [1], [3, 0], "max", True, True, "NCW"),
        # Test non-1 dilations
        ([1, 16, 31], [3], [3], [2], [2, 5], "avg", False, True, "NCW"),
        ([1, 16, 32], [2], [2], [3], [0, 3], "avg", False, False, "NCW"),
        ([1, 16, 31], [3], [3], [2], [1, 4], "max", False, True, "NCW"),
        ([1, 16, 31], [3], [3], [3], [3, 0], "max", True, True, "NCW"),
        # Test Channel last
        ([1, 32, 16], [2], [2], [1], [0, 0], "avg", False, True, "NWC"),
        ([1, 31, 16], [3], [3], [1], [1, 2], "avg", False, True, "NWC"),
        ([1, 32, 16], [2], [2], [1], [1, 2], "avg", False, False, "NWC"),
        ([1, 31, 16], [4], [4], [1], [3, 3], "avg", False, False, "NWC"),
        ([1, 31, 16], [4], [4], [1], [0, 0], "avg", False, False, "NWC"),
        ([1, 32, 16], [2], [2], [1], [0, 0], "max", False, True, "NWC"),
        ([1, 31, 16], [3], [3], [1], [2, 1], "max", False, True, "NWC"),
        ([1, 31, 16], [3], [3], [1], [2, 1], "max", True, True, "NWC"),
        ([1, 31, 16], [3], [3], [1], [2, 5], "avg", False, True, "NWC"),
        ([1, 31, 16], [2], [2], [1], [0, 3], "avg", False, False, "NWC"),
        ([1, 31, 16], [3], [3], [1], [1, 4], "max", False, True, "NWC"),
        ([1, 31, 16], [3], [3], [1], [3, 0], "max", True, True, "NWC"),
        ([1, 31, 16], [3], [3], [2], [2, 5], "avg", False, True, "NWC"),
        ([1, 32, 16], [2], [2], [3], [0, 3], "avg", False, False, "NWC"),
        ([1, 31, 16], [3], [3], [2], [1, 4], "max", False, True, "NWC"),
        ([1, 31, 16], [3], [3], [3], [3, 0], "max", True, True, "NWC"),
    )

    @requires_hexagon_toolchain
    def test_pool1d(
        self,
        hexagon_session,
        input_shape,
        kernel,
        stride,
        dilation,
        padding,
        pool_type,
        ceil_mode,
        count_include_pad,
        layout,
    ):
        verify_poolnd(
            hexagon_session,
            1,
            input_shape,
            kernel,
            stride,
            dilation,
            padding,
            pool_type,
            ceil_mode,
            count_include_pad,
            layout,
        )

class TestPool2D:
    (
        input_shape,
        kernel,
        stride,
        dilation,
        padding,
        pool_type,
        ceil_mode,
        count_include_pad,
        layout,
    ) = tvm.testing.parameters(
        ([1, 16, 32, 32], [2, 2], [2, 2], [1, 1], [0, 0, 0, 0], "avg", False, True, "NCHW"),
        ([1, 16, 31, 31], [3, 3], [3, 3], [1, 1], [1, 2, 1, 2], "avg", False, True, "NCHW"),
        ([1, 16, 32, 32], [2, 2], [2, 2], [1, 1], [1, 2, 1, 2], "avg", False, False, "NCHW"),
        ([1, 16, 31, 31], [4, 4], [4, 4], [1, 1], [3, 3, 3, 3], "avg", False, False, "NCHW"),
        ([1, 16, 31, 31], [4, 4], [4, 4], [1, 1], [0, 0, 0, 0], "avg", False, False, "NCHW"),
        ([1, 16, 32, 32], [2, 3], [2, 2], [1, 1], [0, 0, 0, 0], "max", False, True, "NCHW"),
        ([1, 16, 31, 31], [3, 3], [3, 3], [1, 1], [2, 1, 2, 1], "max", False, True, "NCHW"),
        ([1, 16, 31, 31], [3, 3], [3, 3], [1, 1], [2, 1, 2, 1], "max", True, True, "NCHW"),
        ([1, 16, 31, 31], [3, 3], [3, 3], [1, 1], [2, 1, 0, 3], "avg", False, True, "NCHW"),
        ([1, 16, 32, 32], [2, 3], [2, 2], [1, 1], [0, 3, 2, 1], "avg", False, False, "NCHW"),
        ([1, 16, 31, 31], [3, 3], [3, 3], [1, 1], [1, 0, 3, 2], "max", False, True, "NCHW"),
        ([1, 16, 31, 31], [3, 3], [3, 3], [1, 1], [3, 2, 1, 0], "max", True, True, "NCHW"),
        # Test non-1 dilations
        ([1, 16, 31, 31], [3, 3], [3, 3], [2, 1], [2, 1, 0, 3], "avg", False, True, "NCHW"),
        ([1, 16, 32, 32], [2, 3], [2, 2], [2, 3], [0, 3, 2, 1], "avg", False, False, "NCHW"),
        ([1, 16, 31, 31], [3, 3], [3, 3], [3, 3], [1, 0, 3, 2], "max", False, True, "NCHW"),
        ([1, 16, 31, 31], [3, 3], [3, 3], [2, 2], [3, 2, 1, 0], "max", True, True, "NCHW"),
        # Test channel last
        ([1, 32, 32, 16], [2, 2], [2, 2], [1, 1], [0, 0, 0, 0], "avg", False, True, "NHWC"),
        ([1, 31, 31, 16], [3, 3], [3, 3], [1, 1], [1, 2, 1, 2], "avg", False, True, "NHWC"),
        ([1, 32, 32, 16], [2, 2], [2, 2], [1, 1], [1, 2, 1, 2], "avg", False, False, "NHWC"),
        ([1, 31, 31, 16], [4, 4], [4, 4], [1, 1], [3, 3, 3, 3], "avg", False, False, "NHWC"),
        ([1, 31, 31, 16], [4, 4], [4, 4], [1, 1], [0, 0, 0, 0], "avg", False, False, "NHWC"),
        ([1, 32, 32, 16], [2, 3], [2, 2], [1, 1], [0, 0, 0, 0], "max", False, True, "NHWC"),
        ([1, 31, 31, 16], [3, 3], [3, 3], [1, 1], [2, 1, 2, 1], "max", False, True, "NHWC"),
        ([1, 31, 31, 16], [3, 3], [3, 3], [1, 1], [2, 1, 2, 1], "max", True, True, "NHWC"),
        ([1, 31, 31, 16], [3, 3], [3, 3], [1, 1], [2, 1, 0, 3], "avg", False, True, "NHWC"),
        ([1, 32, 32, 16], [2, 3], [2, 2], [1, 1], [0, 3, 2, 1], "avg", False, False, "NHWC"),
        ([1, 31, 31, 16], [3, 3], [3, 3], [1, 1], [1, 0, 3, 2], "max", False, True, "NHWC"),
        ([1, 31, 31, 16], [3, 3], [3, 3], [1, 1], [3, 2, 1, 0], "max", True, True, "NHWC"),
        ([1, 31, 31, 16], [3, 3], [3, 3], [2, 1], [2, 1, 0, 3], "avg", False, True, "NHWC"),
        ([1, 32, 32, 16], [2, 3], [2, 2], [2, 3], [0, 3, 2, 1], "avg", False, False, "NHWC"),
        ([1, 31, 31, 16], [3, 3], [3, 3], [3, 3], [1, 0, 3, 2], "max", False, True, "NHWC"),
        ([1, 31, 31, 16], [3, 3], [3, 3], [2, 2], [3, 2, 1, 0], "max", True, True, "NHWC"),
    )

    @requires_hexagon_toolchain
    def test_pool2d(
        self,
        hexagon_session,
        input_shape,
        kernel,
        stride,
        dilation,
        padding,
        pool_type,
        ceil_mode,
        count_include_pad,
        layout,
    ):
        verify_poolnd(
            hexagon_session,
            2,
            input_shape,
            kernel,
            stride,
            dilation,
            padding,
            pool_type,
            ceil_mode,
            count_include_pad,
            layout,
        )

class TestPool3D:
    (
        input_shape,
        kernel,
        stride,
        dilation,
        padding,
        pool_type,
        ceil_mode,
        count_include_pad,
        layout,
    ) = tvm.testing.parameters(
        (
            [1, 16, 32, 32, 32],
            [2, 2, 2],
            [2, 2, 2],
            [1, 1, 1],
            [0, 0, 0, 0, 0, 0],
            "avg",
            False,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [1, 1, 2, 2, 2, 1],
            "avg",
            False,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 32, 32, 32],
            [2, 2, 2],
            [2, 2, 2],
            [1, 1, 1],
            [1, 1, 2, 2, 2, 1],
            "avg",
            False,
            False,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [4, 4, 4],
            [4, 4, 4],
            [1, 1, 1],
            [3, 3, 3, 3, 3, 3],
            "avg",
            False,
            False,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [4, 4, 4],
            [4, 4, 4],
            [1, 1, 1],
            [0, 0, 0, 0, 0, 0],
            "avg",
            False,
            False,
            "NCDHW",
        ),
        (
            [1, 16, 32, 32, 32],
            [2, 2, 2],
            [2, 2, 2],
            [1, 1, 1],
            [0, 0, 0, 0, 0, 0],
            "max",
            False,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [2, 2, 1, 1, 1, 2],
            "max",
            False,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [2, 2, 1, 1, 1, 2],
            "max",
            True,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [2, 1, 0, 5, 4, 3],
            "avg",
            False,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 32, 32, 32],
            [2, 2, 2],
            [2, 2, 2],
            [1, 1, 1],
            [0, 5, 4, 3, 2, 1],
            "avg",
            False,
            False,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [1, 0, 5, 4, 3, 2],
            "max",
            False,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [3, 2, 1, 0, 5, 4],
            "max",
            True,
            True,
            "NCDHW",
        ),
        # Test non-1 dilation
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [3, 3, 3],
            [2, 1, 0, 5, 4, 3],
            "avg",
            False,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 32, 32, 32],
            [2, 2, 2],
            [2, 2, 2],
            [2, 2, 2],
            [0, 5, 4, 3, 2, 1],
            "avg",
            False,
            False,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [2, 1, 3],
            [1, 0, 5, 4, 3, 2],
            "max",
            False,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [2, 2, 3],
            [3, 2, 1, 0, 5, 4],
            "max",
            True,
            True,
            "NCDHW",
        ),
        # Test channel last layouts
        (
            [1, 32, 32, 32, 16],
            [2, 2, 2],
            [2, 2, 2],
            [1, 1, 1],
            [0, 0, 0, 0, 0, 0],
            "avg",
            False,
            True,
            "NDHWC",
        ),
        (
            [1, 31, 31, 31, 16],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [1, 1, 2, 2, 2, 1],
            "avg",
            False,
            True,
            "NDHWC",
        ),
        (
            [1, 32, 32, 32, 16],
            [2, 2, 2],
            [2, 2, 2],
            [1, 1, 1],
            [1, 1, 2, 2, 2, 1],
            "avg",
            False,
            False,
            "NDHWC",
        ),
        (
            [1, 31, 31, 31, 16],
            [4, 4, 4],
            [4, 4, 4],
            [1, 1, 1],
            [3, 3, 3, 3, 3, 3],
            "avg",
            False,
            False,
            "NDHWC",
        ),
        (
            [1, 31, 31, 31, 16],
            [4, 4, 4],
            [4, 4, 4],
            [1, 1, 1],
            [0, 0, 0, 0, 0, 0],
            "avg",
            False,
            False,
            "NDHWC",
        ),
        (
            [1, 32, 32, 32, 16],
            [2, 2, 2],
            [2, 2, 2],
            [1, 1, 1],
            [0, 0, 0, 0, 0, 0],
            "max",
            False,
            True,
            "NDHWC",
        ),
        (
            [1, 31, 31, 31, 16],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [2, 2, 1, 1, 1, 2],
            "max",
            False,
            True,
            "NDHWC",
        ),
        (
            [1, 31, 31, 31, 16],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [2, 2, 1, 1, 1, 2],
            "max",
            True,
            True,
            "NDHWC",
        ),
        (
            [1, 31, 31, 31, 16],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [2, 1, 0, 5, 4, 3],
            "avg",
            False,
            True,
            "NDHWC",
        ),
        (
            [1, 32, 32, 32, 16],
            [2, 2, 2],
            [2, 2, 2],
            [1, 1, 1],
            [0, 5, 4, 3, 2, 1],
            "avg",
            False,
            False,
            "NDHWC",
        ),
        (
            [1, 31, 31, 31, 16],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [1, 0, 5, 4, 3, 2],
            "max",
            False,
            True,
            "NDHWC",
        ),
        (
            [1, 31, 31, 31, 16],
            [3, 3, 3],
            [3, 3, 3],
            [1, 1, 1],
            [3, 2, 1, 0, 5, 4],
            "max",
            True,
            True,
            "NDHWC",
        ),
        # Test non-1 dilation
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [3, 3, 3],
            [2, 1, 0, 5, 4, 3],
            "avg",
            False,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 32, 32, 32],
            [2, 2, 2],
            [2, 2, 2],
            [2, 2, 2],
            [0, 5, 4, 3, 2, 1],
            "avg",
            False,
            False,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [2, 1, 3],
            [1, 0, 5, 4, 3, 2],
            "max",
            False,
            True,
            "NCDHW",
        ),
        (
            [1, 16, 31, 31, 31],
            [3, 3, 3],
            [3, 3, 3],
            [2, 2, 3],
            [3, 2, 1, 0, 5, 4],
            "max",
            True,
            True,
            "NCDHW",
        ),
    )

    @requires_hexagon_toolchain
    def test_pool3d(
        self,
        hexagon_session,
        input_shape,
        kernel,
        stride,
        dilation,
        padding,
        pool_type,
        ceil_mode,
        count_include_pad,
        layout,
    ):
        verify_poolnd(
            hexagon_session,
            3,
            input_shape,
            kernel,
            stride,
            dilation,
            padding,
            pool_type,
            ceil_mode,
            count_include_pad,
            layout,
        )

if __name__ == "__main__":
    sys.exit(pytest.main(sys.argv))

max_stars_repo_name: pfk-beta/tvm

max_issues_repo_name: pfk-beta/tvm

max_forks_repo_name: pfk-beta/tvm


text:
from __future__ import print_function
import pickle
import os.path
import sys
from googleapiclient.discovery import build
from google_auth_oauthlib.flow import InstalledAppFlow
from google.auth.transport.requests import Request
import cvxpy as cp
import numpy as np

# If modifying these scopes, delete the file token.pickle.
SCOPES = ['https://www.googleapis.com/auth/spreadsheets.readonly']

# The ID and range of a sample spreadsheet.
FIRSTNAME_RANGE_NAME = 'H3:BH3'
LASTNAME_RANGE_NAME = 'H2:BH2'
STUDENTTYPE_RANGE_NAME = 'H1:BH1'
GAAP_RANGE_NAME = 'D4:D16'
RANKING_RANGE_NAME = 'H4:BH16'
GAAP_COUNT = 13
STUDENT_COUNT = 53

EXTERNES_EXPERIMENTES = [5, 41, 49] # Index des externes ayant plus de 3 ans d'expérience pro
WOMEN = [0, 1, 2, 3, 4, 5, 8, 13, 16, 17, 18, 19, 20, 21, 23, 28, 31, 36, 38, 40, 42, 43, 44, 47, 49] # Index des femmes dans le tableau
MEN = [6, 7, 9, 10, 11, 12, 14, 15, 22, 24, 25, 26, 27, 29, 30, 32, 33, 34, 35, 37, 39, 41, 45, 46, 48, 50, 51, 52] # Index des hommes dans le tableau

def solve(firstnames, lastnames, studenttypes, gaaps, rankings):
    for i in range(GAAP_COUNT):
        for j in range(STUDENT_COUNT):
            if len(rankings[i]) <= j:
                rankings[i].append(1000)
            elif rankings[i][j] == "":
                rankings[i][j] = 1000
            elif rankings[i][j] != "1" and studenttypes[j] == "exterieur": # Ajout d'un coefficient 15 sur les choix secondaires des externes pour prioriser leur premier choix
                rankings[i][j] = 15*int(rankings[i][j])
            else:
                rankings[i][j] = int(rankings[i][j])

    rankings = np.array(rankings)
    groupmax = np.array([5]*GAAP_COUNT)
    groupmin = np.array([4]*GAAP_COUNT)
    externemax = np.array([2]*GAAP_COUNT)
    experimentesmin = np.array([1]*GAAP_COUNT)
    womenmin = np.array([1]*GAAP_COUNT)
    menmin = np.array([1]*GAAP_COUNT)

    selection = cp.Variable(shape=rankings.shape,boolean=True)

    group_constraint_1 = cp.sum(selection,axis=1) <= groupmax

    group_constraint_2 = cp.sum(selection,axis=1) >= groupmin

    externe_filter = np.zeros((STUDENT_COUNT, STUDENT_COUNT))
    for i in range(STUDENT_COUNT):
        if studenttypes[i] == "exterieur":
            externe_filter[i][i] = 1
    externe_constraint = cp.sum((selection @ externe_filter),axis=1) <= externemax

    experimentes_filter = np.zeros((STUDENT_COUNT, STUDENT_COUNT))
    for i in range(STUDENT_COUNT):
        if studenttypes[i] == "IVP" or "Interne" in studenttypes[i] or i in EXTERNES_EXPERIMENTES:
            experimentes_filter[i][i] = 1
    experimentes_constraint = cp.sum((selection @ experimentes_filter),axis=1) >= experimentesmin

    women_filter = np.zeros((STUDENT_COUNT, STUDENT_COUNT))
    for i in range(STUDENT_COUNT):
        if i in WOMEN:
            women_filter[i][i] = 1
    women_constraint = cp.sum((selection @ women_filter),axis=1) >= womenmin

    men_filter = np.zeros((STUDENT_COUNT, STUDENT_COUNT))
    for i in range(STUDENT_COUNT):
        if i in MEN:
            men_filter[i][i] = 1
    men_constraint = cp.sum((selection @ men_filter),axis=1) >= menmin

    assignment_constraint = cp.sum(selection,axis=0) == 1

    cost = cp.sum(cp.multiply(rankings,selection))

    constraints = [group_constraint_1, group_constraint_2, assignment_constraint, externe_constraint, experimentes_constraint, women_constraint, men_constraint]

    assign_prob = cp.Problem(cp.Minimize(cost),constraints)

    assign_prob.solve(solver=cp.GLPK_MI)

    for i in range(GAAP_COUNT):
        print("GAAP " + str(i+1) + " - " + gaaps[i][0])
        for j in range(STUDENT_COUNT):
            if selection.value[i][j] == 1:
                print(firstnames[j] + " " + lastnames[j])

        print()

    selected_rankings = np.multiply(rankings, selection.value)
    unique, counts = np.unique(selected_rankings, return_counts=True)
    occurences = dict(zip(unique, counts))
    for i in range(1, GAAP_COUNT + 1):
        print("Nombre de choix #" + str(i) + " satisfaits : " + str(occurences.get(i, 0) + occurences.get(i*15, 0)))

    print("Nombre de choix non satisfaits : " + str(occurences.get(1000, 0)))

def main():
    """Shows basic usage of the Sheets API.
    Prints values from a sample spreadsheet.
    """
    creds = None
    # The file token.pickle stores the user's access and refresh tokens, and is
    # created automatically when the authorization flow completes for the first
    # time.
    if os.path.exists('token.pickle'):
        with open('token.pickle', 'rb') as token:
            creds = pickle.load(token)
    # If there are no (valid) credentials available, let the user log in.
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            flow = InstalledAppFlow.from_client_secrets_file(
                'credentials.json', SCOPES)
            creds = flow.run_local_server(port=0)
        # Save the credentials for the next run
        with open('token.pickle', 'wb') as token:
            pickle.dump(creds, token)

    service = build('sheets', 'v4', credentials=creds)

    # Call the Sheets API
    sheet = service.spreadsheets()

    spreadsheet = sys.argv[1]

    firstname_result = sheet.values().get(spreadsheetId=spreadsheet,
                                range=FIRSTNAME_RANGE_NAME).execute()
    firstnames = firstname_result.get('values', [])

    lastname_result = sheet.values().get(spreadsheetId=spreadsheet,
                                range=LASTNAME_RANGE_NAME).execute()
    lastnames = lastname_result.get('values', [])

    studenttype_result = sheet.values().get(spreadsheetId=spreadsheet,
                                range=STUDENTTYPE_RANGE_NAME).execute()
    studenttypes = studenttype_result.get('values', [])

    gaap_result = sheet.values().get(spreadsheetId=spreadsheet,
                                range=GAAP_RANGE_NAME).execute()
    gaaps = gaap_result.get('values', [])

    ranking_result = sheet.values().get(spreadsheetId=spreadsheet,
                                range=RANKING_RANGE_NAME).execute()
    rankings = ranking_result.get('values', [])

    solve(firstnames[0], lastnames[0], studenttypes[0], gaaps, rankings)

if __name__ == '__main__':
    main()

max_stars_repo_name: tbourvon/GAAP

max_issues_repo_name: tbourvon/GAAP

max_forks_repo_name: tbourvon/GAAP

reason: import numpy,import cvxpy


text:
# -*- coding: utf-8 -*-
# Copyright (c) 2015-2018, Exa Analytics Development Team
# Distributed under the terms of the Apache License 2.0
import six
#import numpy as np
import pandas as pd
from io import StringIO
#from exa import Series, TypedMeta
from exa import TypedMeta
from exatomic.core import Editor, Tensor

class Meta(TypedMeta):
    tensor = Tensor

class RTensor(six.with_metaclass(Meta, Editor)):
    """
    This is a simple script to read a rank-2 tensor file with frame, label and atom index
    labels. The format for such a file is,

    0: f=** l=** a=**
    1: xx   xy   xz
    2: yx   yy   yz
    3: zx   zy   zz
    4:
    5: Same as above for a second tensor

    """
## Must make this into a class that looks like the XYZ and Cube
#           classes. Must have something like parse_tensor.
#           Then on the Tensor class there should be something that can
#           be activated to find the eigenvalues and eigenvectors of the
#           matrix to plot the basis vectors.
#           Look at untitled1.ipynb for more info.
    _to_universe = Editor.to_universe

    def to_universe(self):
        raise NotImplementedError("Tensor file format has no atom table")

    def parse_tensor(self):
        df = pd.read_csv(StringIO(str(self)), delim_whitespace=True, header=None,
                            skip_blank_lines=False)
        #print(df)
        try:
            i=0
            data = ''
            while True:
                a = df.loc[[i*5],:].values[0]
                labels = []
                for lbl in a:
                    d = lbl.split('=')
                    labels.append(d[1])
                cols = ['xx','xy','xz','yx','yy','yz','zx','zy','zz']
                af = pd.DataFrame([df.loc[[i*5+1,i*5+2,i*5+3],:].unstack().values], \
                                            columns=cols)
                af['frame'] = labels[0] if labels[0] != '' else 0
                af['label'] = labels[1] if labels[1] != '' else None
                af['atom'] = labels[2] if labels[0] != '' else 0
                if i >= 1:
                    data = pd.concat([data,af],keys=[o for o in range(i+1)])
                    #data = data.append(af)
                    print('tens.py--------')
                    print(data)
                    print('---------------')
                else:
                    data = af
                i+=1
        except:
            print('tens.py--------')
            print("Reached EOF reading {} tensor".format(i))
            print(data)
            print('---------------')
        self.tensor = data

#    @classmethod
#    def from_universe(cls):
#        pass

max_stars_repo_name: tjduigna/exatomic

max_issues_repo_name: tjduigna/exatomic

max_forks_repo_name: tjduigna/exatomic


text:
# -*- coding: utf-8 -*-
"""
GenZ pilot analysis script.

"""
import mne

from functools import partial

import numpy as np
from scipy import linalg
import itertools

from .io.pick import pick_types, pick_channels, pick_channels_regexp
from .io.constants import FIFF
from .io.ctf.trans import _make_ctf_coord_trans_set
from .forward import (_magnetic_dipole_field_vec, _create_meg_coils,
                      _concatenate_coils)
from .cov import make_ad_hoc_cov, compute_whitener
from .transforms import (apply_trans, invert_transform, _angle_between_quats,
                         quat_to_rot, rot_to_quat)
from .utils import verbose, logger, use_log_level, _check_fname, warn

from mne.chpi import (_get_hpi_info, _get_hpi_initial_fit, _magnetic_dipole_objective, _fit_magnetic_dipole,
_chpi_objective, _unit_quat_constraint, _fit_chpi_quat, _setup_hpi_struct, _fit_cHPI_amplitudes, _fit_device_hpi_positions)

# coil 1 = 83 hz
# coil 2 = 143 hz
# coil 3 = 203 hz
# coil 4 = 263 hz
# coil 5 = 323 hz

# define variables we need

coil_choices = ['one', 'two', 'three', 'four', 'five']

coil_placement = ['back', 'front']

positions = ['control', 'down', 'fardown', 'forward']

s_freq = [1000, 2000]

lowpass = [330, 660]

dir_path = '/home/nordme/data/cHPI_test'

files = []

# read in the raw file

raw = mne.io.read_raw_fif('path')

info = mne.io.read_info(raw)

# report Elekta's measurements from digitizing
# use the results from _get_hpi_initial_fit and raw_info to make sure they match

hpi_dig_locations = mne.chpi._get_hpi_initial_fit(info)
print('File: %s Initial coil locations: %s' % (file, hpi_dig_locations))
print('raw_info: %s' % raw.info['hpi_meas'])

# report mne-python calculations for coil location and coil distances
# use the results from _fit_device_hpi_positions

# report discrepencies between original and measured, plus measured fit quality

max_stars_repo_path: working file.py

max_issues_repo_path: working file.py

max_forks_repo_path: working file.py

reason: import numpy,from scipy


text:
from abc import ABC, abstractmethod
from dataclasses import dataclass
from pathlib import Path
from typing import Dict

import numpy as np
import pandas as pd
import pyopencl as cl
import xarray as xr

from ..enums.forcings import Forcing

@dataclass
class KernelConfig(ABC):
    pass

class Kernel(ABC):
    def __init__(
        self,
        forcing_data: Dict[Forcing, xr.Dataset],
        p0: xr.Dataset,
        advect_time: pd.DatetimeIndex,
        save_every: int,
        config: KernelConfig,
        context: cl.Context,
    ):
        """
        :param forcing_data: dict containing forcing datasets
        :param p0: initial state of particles
        :param advect_time: the timeseries which the kernel advects on
        :param save_every: number of timesteps between each writing of particle state
        :param config: dataclass with any extra settings needed by subclass implementation
        :param context: PyopenCL context for executing OpenCL programs
        """
        # save some arguments for creating output dataset
        self.p0 = p0
        self.out_time = advect_time[::save_every][1:]
        self.advect_time = advect_time

        # advection time parameters
        self.start_time = np.float64(advect_time[0].timestamp())
        self.dt = np.float64(pd.Timedelta(advect_time.freq).total_seconds())
        self.ntimesteps = np.uint32(len(advect_time) - 1)  # initial position given
        self.save_every = np.uint32(save_every)

        # create opencl objects
        self.context = context
        self.queue = cl.CommandQueue(context)
        self.cl_kernel = cl.Program(context, self._kernel_source_code).build(
            options=["-I", str(self._model_core_path)]
        )

        # some handy timers
        self.data_load_time = 0
        self.buf_time = 0
        self.kernel_time = 0

    @property
    def _model_core_path(self):
        return Path(__file__).parent.parent / "model_core"

    @property
    @abstractmethod
    def _kernel_source_code(self) -> str:
        pass

    @abstractmethod
    def execute(self) -> xr.Dataset:
        pass

    @abstractmethod
    def get_memory_footprint(self) -> dict:
        pass

    @abstractmethod
    def get_data_loading_time(self) -> float:
        pass

    @abstractmethod
    def get_buffer_transfer_time(self) -> float:
        pass

    @abstractmethod
    def get_kernel_execution_time(self) -> float:
        pass

max_stars_repo_name: john-science/ADVECTOR

max_issues_repo_name: TheOceanCleanupAlgorithms/ADVECT

max_forks_repo_name: TheOceanCleanupAlgorithms/ADVECT


text:
import numpy as np

H, N = map(int, input().split())
magic = [[0, 0] for i in range(N)]
for i in range(N):
  magic[i] = list(map(int, input().split()))

magic = np.sort(magic, axis = 0)
magic = np.sort(magic, axis = 1)
magic = tolist(magic) # Convert to list type