text:
import sys, os
from pathlib import Path
import subprocess
import pickle
from typing import Any
import zipfile
import h5py as hf
import numpy as np
import torch
from omnibelt import unspecified_argument, get_printer
import omnifig as fig
from torch.nn import functional as F

prt = get_printer(__name__)

try:
	from google.cloud import storage
except ImportError:
	prt.info('Failed to import gsutil (install with "pip install gsutil")')

try:
	import wget
except ImportError:
	prt.info('Failed to import wget (install with "pip install wget")')

from ... import util
from ...data import register_dataset, Deviced, Batchable, Splitable, ImageDataset, \
	Downloadable, Dataset, MissingDatasetError, Subset_Dataset

from .transforms import Cropped

def _rec_decode(obj):
	'''
	recursively convert bytes to str
	:param obj: root obj
	:return:
	'''
	if isinstance(obj, dict):
		return {_rec_decode(k):_rec_decode(v) for k,v in obj.items()}
	if isinstance(obj, list):
		return [_rec_decode(x) for x in obj]
	if isinstance(obj, tuple):
		return tuple(_rec_decode(x) for x in obj)
	if isinstance(obj, bytes):
		return obj.decode()
	return obj

class DisentanglementDataset(ImageDataset):

	def get_factor_sizes(self):
		raise NotImplementedError
	def get_factor_order(self):
		raise NotImplementedError

	def get_observations(self):
		raise NotImplementedError
	def get_labels(self):
		raise NotImplementedError
	def update_data(self, indices):
		raise NotImplementedError

@fig.AutoModifier('selected')
class Selected(Splitable):
	def __init__(self, A, ordered=unspecified_argument, unordered=unspecified_argument,
	             reject=unspecified_argument, invert=None, eval_name=None, **kwargs):

		if reject is unspecified_argument:
			reject = A.pull('reject', {})
		if ordered is unspecified_argument:
			ordered = A.pull('ordered', None)
		if unordered is unspecified_argument:
			unordered = A.pull('unordered', None)

		if invert is None:
			invert = A.pull('invert', False)

		if eval_name is None:
			eval_name = A.pull('eval-name', 'extra')

		if reject is not None or ordered is not None or unordered is not None:
			A.push('load-labels', True)

		self.reject = reject
		self.ordered = ordered
		self.unordered = unordered

		self.invert = invert

		self.eval_name = eval_name

		super().__init__(A, **kwargs)

	def _split_load(self, dataset):

		sizes = self.get_factor_sizes()
		flts = {}

		if self.reject is not None:
			flts.update({int(k): v for k, v in self.reject.items()})

		if self.ordered is not None:
			for idx, ratio in self.ordered.items():
				idx = int(idx)
				if idx not in flts and ratio != 0:
					N = sizes[idx]
					vals = np.arange(N)
					num = min(np.ceil(abs(ratio) * N), N - 1) if isinstance(ratio, float) \
						else max(1, min(abs(ratio), N - 1))
					num = int(num)
					vals = vals[:num] if ratio > 0 else vals[-num:]
					vals = vals.tolist()

					flts[idx] = vals

		if self.unordered is not None:
			raise NotImplementedError

		self.reject = flts

		if len(flts):
			dataset = self._select(dataset, flts)

		return super()._split_load(dataset)

	def _select(self, dataset, flts):
		lbls = self.get_labels()
		ok = None
		for idx, valid in flts.items():
			valid = torch.tensor(valid, device=lbls.device).unsqueeze(0)
			vals = lbls[:, idx:idx+1]
			for v in valid:
				s = vals.sub(v).eq(0).sum(-1)
				if ok is None:
					ok = s
				else:
					ok *= s
		ok = ok.bool()
		if not self.invert:
			ok = torch.logical_not(ok)

		inds = torch.arange(len(lbls), device=lbls.device)

		extra = Subset_Dataset(dataset, inds[torch.logical_not(ok)])
		self.register_mode(self.eval_name, extra)

		dataset = Subset_Dataset(dataset, inds[ok])
		self.register_mode(self.split_src, dataset)

		return dataset

@register_dataset('dsprites')
class dSprites(Deviced, Batchable, Downloadable, DisentanglementDataset):

	din = (1, 64, 64)
	dout = 5

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

		label_type = A.pull('label_type', None)

		din = A.pull('din', self.din)
		dout = A.pull('dout', din if label_type is None else self.dout)

		filename = A.pull('filename', 'dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz')

		assert label_type in {None, 'value', 'class'}, 'Unknown type of label: {}'.format(label_type)
		# assert images is not None or dataroot is not None, 'nothing to use/load'

		if dout is None and label_type is not None:
			dout = 5 if label_type == 'value' else 113

		super().__init__(A, din=din, dout=din if dout is None else dout, **kwargs)

		dataroot = self.root

		if dataroot is not None:
			root = os.path.join(dataroot, 'dsprites')
			path = os.path.join(root, filename)
			print('Loading dSprites dataset from disk: {}'.format(path))
			data = np.load(path, allow_pickle=True, encoding='bytes')

			self.meta = _rec_decode(data['metadata'][()])

			images = torch.from_numpy(data['imgs']).unsqueeze(1)

			self.register_buffer('images', images)

			if label_type is not None:
				if label_type == 'value':
					labels = torch.from_numpy(data['latents_values'][:,1:]).float()
				else:
					labels = torch.from_numpy(data['latents_classes'][:,1:]).int()
				self.register_buffer('labels', labels)

		self.labeled = hasattr(self, 'labels')
		self.root = root

	def get_raw_data(self):
		if self.labeled:
			return self.images, self.labels
		return self.images

	def __len__(self):
		return len(self.images)

	def __getitem__(self, item):
		imgs = self.images[item].float()
		if self.labeled:
			return imgs, self.labels[item]
		return imgs,

@register_dataset('3dshapes')
class Shapes3D(Deviced, Batchable, Downloadable, DisentanglementDataset):

	din = (3, 64, 64)
	dout = 6

	def __init__(self, A, mode=None, labeled=None, load_labels=None, din=None, dout=None, **kwargs):

		if mode is None:
			mode = A.pull('mode', 'full')
		if labeled is None:
			labeled = A.pull('labeled', False)
		if load_labels is None and not labeled:
			load_labels = A.pull('load-labels', False)
		load_labels = load_labels or labeled
		label_type = A.pull('label_type', 'class') if load_labels else None

		if din is None: din = A.pull('din', self.din)
		if dout is None: dout = A.pull('dout', self.dout if labeled else din)

		super().__init__(A, din=din, dout=dout, **kwargs)
		self.add_existing_modes('test')

		self.factor_order = ['floor_hue', 'wall_hue', 'object_hue', 'scale', 'shape', 'orientation']
		self.factor_sizes = [10, 10, 10, 8, 4, 15]
		self.factor_num_values = {'floor_hue': 10, 'wall_hue': 10, 'object_hue': 10,
		                          'scale': 8, 'shape': 4, 'orientation': 15}

		raw_mins = torch.tensor([  0.  ,   0.  ,   0.  ,   0.75,   0.  , -30.  ]).float()
		raw_maxs = torch.tensor([  0.9 ,  0.9 ,  0.9 ,  1.25,  3.  , 30.  ]).float()

		dataroot = self.root / '3dshapes'
		dataroot.mkdir(exist_ok=True)
		self.root = dataroot

		path = dataroot / '3dshapes.h5'
		if not path.is_file():
			download = A.pull('download', False)
			if download:
				self.download(A)
			else:
				raise MissingDatasetError('3dshapes')

		self.images, self.labels, self.indices = None, None, None
		with hf.File(path, 'r') as data:

			images = torch.from_numpy(data['images'][()])
			labels = torch.from_numpy(data['labels'][()]).float() if load_labels else None

			indices = None if mode == 'full' else \
				(data['test_idx'][()] if mode == 'test' else data['train_idx'][()])
			if indices is not None:
				indices = torch.from_numpy(indices).long()
				images = images[indices]
				if labels is not None:
					labels = labels[indices]
				self.register_buffer('indices', indices)
			else:
				prt.warning('using the full dataset (train+test)')

			images = images.permute(0, 3, 1, 2)#.float().div(255)
			self.register_buffer('images', images)

			if labels is not None:
				if label_type == 'class':
					nums = torch.tensor(self.factor_sizes).float() - 1
					labels -= raw_mins
					labels /= raw_maxs - raw_mins
					labels *= nums
					labels = labels.round().long()
				self.register_buffer('labels', labels)

		self.labeled = labeled
		self.label_type = label_type

	_source_url = 'gs://3d-shapes/3dshapes.h5'
	@classmethod
	def download(cls, A, dataroot=None, **kwargs):

		if dataroot is None:
			dataroot = util.get_data_dir(A) / '3dshapes'

		dest = dataroot / '3dshapes.h5'

		force_download = A.pull('force-download', False, silent=True)

		if not dest.is_file() or force_download:
			print(f'Downloading 3dshapes dataset to {str(dest)} ...', end='')
			subprocess.run(['gsutil', 'cp', cls._source_url, str(dest)])
			print(' done!')

		ratio = A.pull('separate-testset', 0.2)

		rng = np.random.RandomState(0)

		with hf.File(dest, 'r+') as f:
			if ratio is not None and ratio > 0 and 'train_idx' not in f:

				N, H, W, C = f['images'].shape

				test_N = int(N * ratio)

				order = rng.permutation(N)

				train_idx = order[:-test_N]
				train_idx.sort()
				test_idx = order[-test_N:]
				test_idx.sort()

				f.create_dataset('train_idx', data=train_idx)
				f.create_dataset('test_idx', data=test_idx)

	def get_factor_sizes(self):
		return self.factor_sizes
	def get_factor_order(self):
		return self.factor_order

	def get_labels(self):
		return self.labels

	def get_raw_data(self):
		if self.labeled:
			return self.images, self.labels
		return self.images

	def update_data(self, indices):
		self.images = self.images[indices]
		if self.labeled:
			self.labels = self.labels[indices]

	def __len__(self):
		return len(self.images)

	def __getitem__(self, item):
		# if isinstance(item, (np.ndarray, torch.Tensor)):
		# 	item = item.tolist()
		images = self.images[item].float().div(255)
		# if self.noise is not None:
		# 	images = images.add(torch.randn_like(images).mul(self.noise)).clamp(0,1)
		if self.labeled:
			labels = self.labels[item]
			# labels = torch.from_numpy(self.labels[item]).float()
			return images, labels
		return images,

@register_dataset('rfd')
class RFD(Downloadable, DisentanglementDataset):

	""" Implementing RA dataset for RFD based on .h5 storage"""
	shape = (3, 128,128)
	din = shape
	dout = 9
	raw_subfolders = {'train':"finger", 'test':'finger', 'heldout_test':"finger_heldout_colors",
	                  'real_test':"finger_real"}
	raw_files = {'train':8, 'test':2, 'heldout_test':1} #train, test, heldout_test files

	# expected folder structure
	# root (whatever this is it must be passed in the kwargs)
	# |-RFD
	# 	|-raw
	#		|- finger
	#			finger_info.npz
	#		|- finger_heldout
	#			finger_heldout_info.npz
	#		|- finger_real
	#			finger_real_info.npz
	#			finger_real_labels.npz
	#			finger_real_images.npz
	# 	|-processed
	# 		|-processed
	# 		Here go all the .h5 files

	def __init__(self,
				 A,
	             mode=None,
				 labeled=None,
	             din=None, dout=None,
	             **kwargs) -> None:

		if mode is None: mode = A.pull('mode', 'train')
		if labeled is None: labeled = A.pull('labeled', False)

		if din is None: din = A.pull('din', self.din)
		if dout is None: dout = A.pull('dout', self.dout if labeled else din)

		super().__init__(A, din=din, dout=din if dout is None else dout, **kwargs)
		self.root = self.root / 'RFD'
		self.add_existing_modes(*self.raw_subfolders.keys())

		self.partitions_dict = None
		self.real_set = None
		self.info = None

		self.labeled = labeled

		self.switch_to(mode)

	def switch_to(self, mode):
		if mode == 'real_test' and self.real_set is None:
			self.real_set = self.read_real_images_and_labels()
		elif mode != self.get_mode() or self.partitions_dict is None:
			self.partitions_dict = self.init_partitions()
		self.info = self.read_info()
		return super().switch_to(mode)

	def read_real_images_and_labels(self):
		""" Loading function only for real test dataset images"""
		print("====== Opening real RFD Dataset ======")
		images = np.load(self.raw_folder+"_images.npz", allow_pickle=True)["images"]
		labels = np.load(self.raw_folder+"_labels.npz", allow_pickle=True)["labels"]
		return (images, labels)

	def init_partitions(self):
		"""Initialising the .h5 files and their order.
		Note: when reading .h5 file it won't be loaded into memory until its
		entries are called."""
		print("====== Opening RFD Dataset ======")
		partitions_dict = {idx: self.open_partition(idx) for idx in range(self.num_files)}
		return partitions_dict

	def get_relative_index(self, index):
		""" Given absolute index it returns the partition index and the
		relative index inside the partition
		#TODO: make it work for multiple indices"""
		len_per_file = self.partitions_dict[0][1]
		partition_idx = index//len_per_file
		relative_index = index%len_per_file
		if index >= len_per_file*self.num_files:
			# contained in last partition
			relative_index = index - len_per_file*self.num_files
		return partition_idx, relative_index

	def open_partition(self, number):
		filename = f"RFD_{self.set_name}_{number}.h5"
		f = hf.File(str(self.root / 'processed' / filename), 'r')
		images = f['images']
		labels = f['labels']
		length = images.shape[0]
		return (images, labels), length

	def __getitem__(self, index: int) -> Any:
		if not isinstance(index, int):
			index = index.item()
		if self.get_mode() == 'real_test':
			imgs, lbls = self.real_set
			rel_idx = index
		else:
			par_idx, rel_idx = self.get_relative_index(index)
			imgs, lbls = self.partitions_dict[par_idx][0]
		img = torch.from_numpy(imgs[rel_idx]).float().div(255).permute(2,0,1)#rescaling the uint8 # from (H,W,C) to (C, H, W)

		# both the above are numpy arrays
		# so they need to be cast to torch tensors
		#Note: no rescaling needed as the imgs have already been
		# processed by the ToTensor transformation.
		# if self.transform is not None:
		# 	img = self.transform(img)
		if self.labeled:
			lbl = torch.from_numpy(lbls[rel_idx])
			# if self.target_transform is not None:
			# 	lbl = self.target_transform(lbl)
			return img, lbl
		return img

	def __len__(self) -> int:
		#TODO: see here too for the "batched" case
		if self.get_mode() in {'heldout_test', 'real_test'}: return self.size
		if self.get_mode() != 'test': return int(self.size*0.8)
		# only test is remaining
		return self.size -  int(self.size*0.8)

	def read_info(self):
		""" Opens the labels and info .npz files and stores them in the class"""
		# loading labels
		info = dict(np.load(self.raw_folder+"_info.npz", allow_pickle=True))
		self.size = info["dataset_size"].item()
		self.factor_order = list(info['factor_values'].item().keys())
		self.factor_sizes = info['num_factor_values']
		self.factor_num_values = {k:v for k,v in zip(self.factor_order, self.factor_sizes)}
		return info

	def get_factor_sizes(self):
		return self.factor_sizes

	def get_factor_order(self):
		return self.factor_order

	def get_observations(self):
		raise NotImplementedError

	def get_labels(self):
		raise NotImplementedError

	def update_data(self, indices):
		raise NotImplementedError

	def close_all(self):
		#TODO: use this when finished to close all the hdf5 files opened
		pass

	@property
	def raw_folder(self) -> str:
		prefix = self.raw_subfolders.get(self.get_mode(), 'finger')
		return str(self.root / 'raw' / prefix / prefix)

	@property
	def num_files(self) -> int:
		return self.raw_files.get(self.get_mode(), 1)

	@property
	def set_name(self) -> str:
		if self.get_mode() == 'test': return "test"
		if self.get_mode() == 'heldout_test': return "HC"
		return "train"

	# @property
	# def processed_folder(self) -> str:
	# 	return self.root + 'RFD/processed/'

@register_dataset('full-celeba')  # probably shouldnt be used
class FullCelebA(Downloadable, ImageDataset):  # TODO: automate downloading and formatting

	din = (3, 218, 178)

	ATTRIBUTES = ['5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes',
	              'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair',
	              'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin',
	              'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones',
	              'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard',
	              'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks',
	              'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings',
	              'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young', ]

	def __init__(self, A, resize=unspecified_argument, label_type=unspecified_argument, mode=None,
	             din=None, dout=None, **kwargs):

		if label_type is unspecified_argument:
			label_type = A.pull('label_type', None)

		if mode is None:
			mode = A.pull('mode', 'train')
		if resize is unspecified_argument:
			resize = A.pull('resize', (256, 256))

		if din is None:
			din = A.pull('din', self.din)

		if dout is None:
			if label_type is None:
				dout = din
			elif label_type == 'attr':
				dout = 40
			elif label_type == 'landmark':
				dout = 10
			elif label_type == 'identity':
				dout = 1
			else:
				raise Exception('unknown {}'.format(label_type))

			dout = A.pull('dout', dout)

		_labels = {
			'attr': 'attrs',
			'identity': 'identities',
			'landmark': 'landmarks',
		}

		if resize is not None:  # TODO: use Interpolated as modifier
			din = 3, *resize

		super().__init__(A, din=din, dout=dout, **kwargs)
		# self.add_existing_modes('val', 'test')
		self.root = Path(self.root) / 'celeba'
		dataroot = self.root
		name = 'celeba_test.h5' if mode == 'test' else 'celeba_train.h5'

		with hf.File(dataroot/name, 'r') as f:
			self.images = f['images'][()]  # encoded as str
			self.labels = f[_labels[label_type]][()] if label_type is not None else None

			self.attr_names = f.attrs['attr_names']
			self.landmark_names = f.attrs['landmark_names']

		self.resize = resize

	# google drive ids
	_google_drive_ids = {
		'list_eval_partition.txt': '0B7EVK8r0v71pY0NSMzRuSXJEVkk',
		'identity_CelebA.txt': '1_ee_0u7vcNLOfNLegJRHmolfH5ICW-XS',
		'list_attr_celeba.txt': '0B7EVK8r0v71pblRyaVFSWGxPY0U',
		'list_bbox_celeba.txt': '0B7EVK8r0v71pbThiMVRxWXZ4dU0',
		'list_landmarks_align_celeba.txt': '0B7EVK8r0v71pd0FJY3Blby1HUTQ',
	}
	_google_drive_image_id = '0B7EVK8r0v71pZjFTYXZWM3FlRnM'

	@classmethod
	def download(cls, A, dataroot=None, **kwargs):

		raise NotImplementedError('Unfortunately, automatically downloading CelebA is current not supported.')

		dataroot = util.get_data_dir(A)
		dataroot = dataroot / 'celeba'
		dataroot.mkdir(exist_ok=True)

		prt.warning(
			'Downloading CelebA doesn\'t seem to work currently (issues with google drive), instead it must be downloaded manually:'
			'\n1. create a directory "celeba" in the local_data directory'
			'\n2. download "img_align_celeba.zip" into that directory'
			'\n3. download label files: {}'
			'\n4. extract "img_align_celeba.zip" to a directory "img_align_celeba"'.format(
				', '.join(f'"{c}"' for c in cls._google_drive_ids.keys())))

		raise NotImplementedError

		imgdir = dataroot / 'img_align_celeba'

		pbar = A.pull('pbar', None)

		util.download_file_from_google_drive('0B7EVK8r0v71pOXBhSUdJWU1MYUk', dataroot, pbar=pbar)

		for name, ID in cls._google_drive_ids.items():
			path = dataroot / name
			if not path.exists():
				util.download_file_from_google_drive(ID, dataroot, name, pbar=pbar)

		if not imgdir.is_dir():

			imgpath = dataroot / 'img_align_celeba.zip'

			# download zip
			if not imgpath.exists():
				imgpath = util.download_file_from_google_drive(cls._google_drive_image_id, dataroot, pbar=pbar)

			# extract zip
			imgdir.mkdir(exist_ok=True)
			with zipfile.ZipFile(str(imgpath), 'r') as zip_ref:
				zip_ref.extractall(str(imgdir))

			# os.remove(str(imgpath))

		raise NotImplementedError

	def get_factor_sizes(self):
		return [2] * len(self.ATTRIBUTES)

	def get_factor_order(self):
		return self.ATTRIBUTES

	def get_attribute_key(self, idx):
		try:
			return self.ATTRIBUTES[idx]
		except IndexError:
			pass

	def __len__(self):
		return len(self.images)

	def __getitem__(self, item):

		img = torch.from_numpy(util.str_to_jpeg(self.images[item])).permute(2, 0, 1).float().div(255)

		if self.resize is not None:
			img = F.interpolate(img.unsqueeze(0), size=self.resize, mode='bilinear').squeeze(0)

		if self.labels is None:
			return img,

		lbl = torch.from_numpy(self.labels[item])

		return img, lbl

@register_dataset('celeba')
class CelebA(Cropped, FullCelebA):
	def __init__(self, A, resize=None, crop_size=128, **kwargs):
		# raise NotImplementedError('doesnt work since FullCelebA automatically resizes to (256,256)')
		crop_size = A.pull('crop_size', crop_size) # essentially sets this as the default

		super().__init__(A, resize=resize, crop_size=crop_size, **kwargs)

@register_dataset('mpi3d')
class MPI3D(Deviced, Batchable, Downloadable, DisentanglementDataset):

	din = (3, 64, 64)
	dout = 7

	def __init__(self, A, mode=None, fid_ident=None, **kwargs):

		if mode is None:
			mode = A.pull('mode', 'train')
		labeled = A.pull('labeled', False)

		din = A.pull('din', self.din)
		dout = A.pull('dout', self.dout if labeled else din)

		cat = A.pull('category', 'toy')

		assert cat in {'toy', 'sim', 'realistic', 'real', 'complex'}, 'invalid category: {}'.format(cat)
		if cat == 'sim':
			cat = 'realistic'

		super().__init__(A, din=din, dout=dout, fid_ident=cat, **kwargs)
		self.add_existing_modes('test')
		dataroot = self.root / 'mpi3d'
		self.root = dataroot

		self.factor_order = ['object_color', 'object_shape', 'object_size', 'camera_height', 'background_color',
		                     'horizonal_axis', 'vertical_axis']
		self.factor_sizes = [6,6,2,3,3,40,40]

		self.factor_values = {
			'object_color': ['white', 'green', 'red', 'blue', 'brown', 'olive'],
			'object_shape': ['cone', 'cube', 'cylinder', 'hexagonal', 'pyramid', 'sphere'],
			'object_size': ['small', 'large'],
			'camera_height': ['top', 'center', 'bottom'],
			'background_color': ['purple', 'sea_green', 'salmon'],
			'horizonal_axis': list(range(40)),
			'vertical_axis': list(range(40)),
		}

		if cat == 'complex':

			self.factor_sizes[0], self.factor_sizes[1] = 4, 4
			self.factor_values['object_shape'] = ['mug', 'ball', 'banana', 'cup']
			self.factor_values['object_color'] = ['yellow', 'green', 'olive', 'red']

		sizes = np.array(self.factor_sizes)

		flr = np.cumprod(sizes[::-1])[::-1]
		flr[:-1] = flr[1:]
		flr[-1] = 1

		self._sizes = torch.from_numpy(sizes.copy()).long()
		self._flr = torch.from_numpy(flr.copy()).long()

		self.labeled = labeled

		path = dataroot / f'mpi3d_{cat}.h5'

		if not path.is_file():
			download = A.pull('download', False)
			if download:
				self.download(A)
			else:
				raise MissingDatasetError(f'mpi3d-{cat}')

		with hf.File(path, 'r') as f:
			if mode == 'full':
				images = np.concatenate([f['train_images'][()], f['test_images'][()]])
				indices = np.concatenate([f['train_idx'], f['test_idx'][()]])
			elif mode == 'test':
				images = f['test_images'][()]
				indices = f['test_idx'][()]
			else:
				images = f['train_images'][()]
				indices = f['train_idx'][()]

		ordered = A.pull('ordered', mode == 'full')
		if ordered:
			self.register_buffer('sel_index', torch.from_numpy(indices.argsort()))
		else:
			self.sel_index = None

		self.register_buffer('images', torch.from_numpy(images).permute(0,3,1,2))
		self.register_buffer('indices', torch.from_numpy(indices))

	_source_url = {
		'toy': 'https://storage.googleapis.com/disentanglement_dataset/Final_Dataset/mpi3d_toy.npz',
		'realistic': 'https://storage.googleapis.com/disentanglement_dataset/Final_Dataset/mpi3d_realistic.npz',
		'real': 'https://storage.googleapis.com/disentanglement_dataset/Final_Dataset/mpi3d_real.npz',
	}
	@classmethod
	def download(cls, A, dataroot=None, cat=None, **kwargs):

		if dataroot is None:
			dataroot = util.get_data_dir(A)
			dataroot = dataroot / 'mpi3d'
		dataroot.mkdir(exist_ok=True)

		if cat is None:
			cat = A.pull('category', 'toy')

		assert cat in {'toy', 'sim', 'realistic', 'real'}, f'invalid category: {cat}'
		if cat == 'sim':
			cat = 'realistic'

		ratio = A.pull('separate-testset', 0.2)

		path = dataroot / f'mpi3d_{cat}.h5'

		if not path.exists():
			rawpath = dataroot / f'mpi3d_{cat}.npz'
			if not rawpath.exists():
				print(f'Downloading mpi3d-{cat}')
				wget.download(cls._source_url[cat], str(rawpath))

			print('Loading full dataset into memory to split train/test')
			full = np.load(rawpath)['images']

			N = len(full)

			test_N = int(N * ratio)

			rng = np.random.RandomState(0)
			order = rng.permutation(N)

			train_idx = order[:-test_N]
			train_idx.sort()
			test_idx = order[-test_N:]
			test_idx.sort()

			with hf.File(path, 'w') as f:
				f.create_dataset('train_idx', data=train_idx)
				f.create_dataset('train_images', data=full[train_idx])

				f.create_dataset('test_idx', data=test_idx)
				f.create_dataset('test_images', data=full[test_idx])

			os.remove(str(rawpath))

	def get_factor_sizes(self):
		return self.factor_sizes
	def get_factor_order(self):
		return self.factor_order

	def get_labels(self):
		return self.get_label(self.indices)

	def update_data(self, indices):
		self.images = self.images[indices]
		if self.labeled:
			self.indices = indices

	def get_label(self, inds):
		try:
			len(inds)
			inds = inds.view(-1,1)
		except TypeError:
			pass

		lvls = inds // self._flr
		labels = lvls % self._sizes

		return labels

	def __len__(self):
		return len(self.images)

	def __getitem__(self, idx):
		if self.sel_index is not None:
			idx = self.sel_index[idx]
		imgs = self.images[idx].float().div(255)
		if self.labeled:
			labels = self.get_label(self.indices[idx])
			return imgs, labels
		return imgs,

class OldFullCelebA(Downloadable, ImageDataset):  # TODO: automate downloading and formatting

	din = (3, 218, 178)

	ATTRIBUTES = ['5_o_Clock_Shadow', 'Arched_Eyebrows', 'Attractive', 'Bags_Under_Eyes',
	              'Bald', 'Bangs', 'Big_Lips', 'Big_Nose', 'Black_Hair', 'Blond_Hair',
	              'Blurry', 'Brown_Hair', 'Bushy_Eyebrows', 'Chubby', 'Double_Chin',
	              'Eyeglasses', 'Goatee', 'Gray_Hair', 'Heavy_Makeup', 'High_Cheekbones',
	              'Male', 'Mouth_Slightly_Open', 'Mustache', 'Narrow_Eyes', 'No_Beard',
	              'Oval_Face', 'Pale_Skin', 'Pointy_Nose', 'Receding_Hairline', 'Rosy_Cheeks',
	              'Sideburns', 'Smiling', 'Straight_Hair', 'Wavy_Hair', 'Wearing_Earrings',
	              'Wearing_Hat', 'Wearing_Lipstick', 'Wearing_Necklace', 'Wearing_Necktie', 'Young', ]

	def __init__(self, A, resize=unspecified_argument, label_type=unspecified_argument, mode=None, **kwargs):
		self.add_existing_modes('val', 'test')

		if label_type is unspecified_argument:
			label_type = A.pull('label_type', None)

		if mode is None:
			mode = A.pull('mode', 'train')
		if resize is unspecified_argument:
			resize = A.pull('resize', (256, 256))

		din = A.pull('din', self.din)

		if label_type is None:
			dout = din
		elif label_type == 'attr':
			dout = 40
		elif label_type == 'landmark':
			dout = 10
		elif label_type == 'identity':
			dout = 1
		else:
			raise Exception('unknown {}'.format(label_type))

		dout = A.pull('dout', dout)

		_labels = {
			'attr': 'attrs',
			'identity': 'identities',
			'landmark': 'landmarks',
		}

		if resize is not None:  # TODO: use Interpolated as modifier
			din = 3, *resize

		super().__init__(A, din=din, dout=dout, **kwargs)
		self.root = Path(self.root) / 'celeba'
		dataroot = self.root
		name = 'celeba_test.h5' if mode == 'test' else 'celeba_train.h5'

		with hf.File(os.path.join(dataroot, 'celeba', name), 'r') as f:
			self.images = f['images'][()]  # encoded as str
			self.labels = f[_labels[label_type]][()] if label_type is not None else None

			self.attr_names = f.attrs['attr_names']
			self.landmark_names = f.attrs['landmark_names']

		self.resize = resize

	# google drive ids
	_google_drive_ids = {
		'img_align_celeba.zip': '0B7EVK8r0v71pZjFTYXZWM3FlRnM',
		'list_eval_partition.txt': '0B7EVK8r0v71pY0NSMzRuSXJEVkk',
		'identity_CelebA.txt': '1_ee_0u7vcNLOfNLegJRHmolfH5ICW-XS',
		'list_attr_celeba.txt': '0B7EVK8r0v71pblRyaVFSWGxPY0U',
		'list_bbox_celeba.txt': '0B7EVK8r0v71pbThiMVRxWXZ4dU0',
		'list_landmarks_align_celeba.txt': '0B7EVK8r0v71pd0FJY3Blby1HUTQ',
	}

	def download(cls, A, dataroot=None, **kwargs):

		dataroot = util.get_data_dir(A)
		dataroot = dataroot / 'celeba'

		raise NotImplementedError

	def get_factor_sizes(self):
		return [2] * len(self.ATTRIBUTES)

	def get_factor_order(self):
		return self.ATTRIBUTES

	def get_attribute_key(self, idx):
		try:
			return self.ATTRIBUTES[idx]
		except IndexError:
			pass

	def __len__(self):
		return len(self.images)

	def __getitem__(self, item):

		img = torch.from_numpy(util.str_to_jpeg(self.images[item])).permute(2, 0, 1).float().div(255)

		if self.resize is not None:
			img = F.interpolate(img.unsqueeze(0), size=self.resize, mode='bilinear').squeeze(0)

		if self.labels is None:
			return img,

		lbl = torch.from_numpy(self.labels[item])

		return img, lbl

ext: py

max_stars_repo_name: GiuliaLanzillotta/omni-learn

max_issues_repo_name: GiuliaLanzillotta/omni-learn

max_forks_repo_name: GiuliaLanzillotta/omni-learn

reason: import numpy


text:
#!/usr/bin/env python
#
# Run me as follows:
# cd tests/
# nosetests

import glob
import json
import librosa
from nose.tools import nottest, eq_, raises, assert_equals
import numpy.testing as npt
import os

# Msaf imports
import msaf.featextract
from msaf.input_output import FileStruct

# Global vars
audio_file = os.path.join("data", "Jaillet64.wav")
# audio_file = os.path.join("../datasets/BeatlesTUT/audio", "01AHardDaysNight.wav")
sr = msaf.Anal.sample_rate
audio, fs = librosa.load(audio_file, sr=sr)
y_harmonic, y_percussive = librosa.effects.hpss(audio)

def test_compute_beats():
	beats_idx, beats_times = msaf.featextract.compute_beats(y_percussive)
	assert len(beats_idx) > 0
	assert len(beats_idx) == len(beats_times)

def test_compute_features():
	mfcc, hpcp, tonnetz, cqt, gmt = msaf.featextract.compute_features(audio, y_harmonic)
	assert mfcc.shape[1] == msaf.Anal.mfcc_coeff
	assert hpcp.shape[1] == 12
	assert tonnetz.shape[1] == 6
	# assert gmt.shape[1] == msaf.Anal.mfcc_coeff + 6
	assert_equals(mfcc.shape[0], hpcp.shape[0])
	assert_equals(hpcp.shape[0], tonnetz.shape[0])
	assert_equals(mfcc.shape[0], gmt.shape[0])

def test_save_features():
	# Read audio and compute features
	tmp_file = "temp.json"
	features = msaf.featextract.compute_features_for_audio_file(audio_file)
	msaf.featextract.save_features(tmp_file, features)

	# Check that new file exists
	assert os.path.isfile(tmp_file)

	# Check that the json file is actually readable
	with open(tmp_file, "r") as f:
		features = json.load(f)
	npt.assert_almost_equal(features["analysis"]["dur"], len(audio) / float(sr),
							decimal=1)

	# Clean up
	os.remove(tmp_file)

def test_compute_beat_sync_features():
	# Compute features
	features = msaf.featextract.compute_features_for_audio_file(audio_file)

	# Beat track
	beats_idx, beats_times = msaf.featextract.compute_beats(y_percussive,
															sr=sr)

	# Compute beat sync feats
	bs_mfcc, bs_hpcp, bs_tonnetz, bs_cqt, bs_gmt = \
		msaf.featextract.compute_beat_sync_features(features, beats_idx)
	assert_equals(bs_mfcc.shape[0],	 len(beats_idx) - 1)
	assert_equals(bs_mfcc.shape[0], bs_hpcp.shape[0])
	assert_equals(bs_hpcp.shape[0], bs_tonnetz.shape[0])
	assert_equals(bs_mfcc.shape[0], bs_gmt.shape[0])

def test_compute_features_for_audio_file():
	features = msaf.featextract.compute_features_for_audio_file(audio_file)
	keys = ["mfcc", "hpcp", "tonnetz", "cqt", "gmt", "beats_idx", "beats", "bs_mfcc",
			"bs_hpcp", "bs_tonnetz", "bs_cqt", "bs_gmt", "anal"]
	anal_keys = ["frame_rate", "hop_size", "mfcc_coeff", "sample_rate",
				 "window_type", "n_mels", "dur"]
	for key in keys:
		assert key in features.keys()
	for anal_key in anal_keys:
		assert anal_key in features["anal"].keys()

def test_compute_all_features():
	# Create file struct
	file_struct = FileStruct(audio_file)

	# Set output file
	feat_file = "tmp.json"
	beats_file = "beats.wav"
	file_struct.features_file = feat_file

	# Remove previously computed outputs if exist
	if os.path.isfile(feat_file):
		os.remove(feat_file)
	if os.path.isfile(beats_file):
		os.remove(beats_file)

	# Call main function
	msaf.featextract.compute_all_features(file_struct, sonify_beats=False,
										  overwrite=False)
	assert os.path.isfile(feat_file)

	# Call again main function (should do nothing, since feat_file exists)
	msaf.featextract.compute_all_features(file_struct, sonify_beats=False,
										  overwrite=False)
	assert os.path.isfile(feat_file)

	# Overwrite
	msaf.featextract.compute_all_features(file_struct, sonify_beats=False,
										  overwrite=True)
	assert os.path.isfile(feat_file)

	# Sonify
	msaf.featextract.compute_all_features(file_struct, sonify_beats=True,
										  overwrite=True, out_beats=beats_file)
	assert os.path.isfile(feat_file) and os.path.isfile(beats_file)

	# Clean up
	os.remove(feat_file)
	os.remove(beats_file)

def test_process():
	# Set output file
	feat_file = "tmp.json"
	beats_file = "beats.wav"

	# Remove previously computed outputs if exist
	if os.path.isfile(feat_file):
		os.remove(feat_file)
	if os.path.isfile(beats_file):
		os.remove(beats_file)

	# Call main function
	msaf.featextract.process(audio_file, out_file=feat_file)
	assert os.path.isfile(feat_file)

	# Call again main function (should do nothing, since feat_file exists)
	msaf.featextract.process(audio_file, out_file=feat_file)
	assert os.path.isfile(feat_file)

	# Overwrite
	msaf.featextract.process(audio_file, out_file=feat_file, overwrite=True)
	assert os.path.isfile(feat_file)

	# Sonify
	msaf.featextract.process(audio_file, out_file=feat_file, overwrite=True,
							 sonify_beats=True, out_beats=beats_file)
	assert os.path.isfile(feat_file) and os.path.isfile(beats_file)

	# Clean up
	os.remove(feat_file)
	os.remove(beats_file)

def main():
	test_compute_features()

if __name__ == '__main__':
	main()

max_stars_repo_name: m-tian/msaf-copy

max_issues_repo_name: m-tian/msaf-copy

max_forks_repo_name: m-tian/msaf-copy


text:
"""
    minimap.py
"""
from __future__ import absolute_import, division, print_function, unicode_literals
from collections import namedtuple
import logging

from OpenGL import GL

from PySide import QtCore, QtGui
import numpy

from mcedit2.rendering import compass
from mcedit2.rendering.scenegraph import scenenode, rendernode
from mcedit2.rendering.layers import Layer
from mcedit2.util.glutils import gl
from mcedit2.util.raycast import rayCastInBounds, MaxDistanceError
from mcedit2.worldview.worldview import WorldView
from mceditlib.geometry import Ray

log = logging.getLogger(__name__)

class LineSegment(namedtuple("LineSegment", "p1 p2")):
    def atHeight(self, y):
        p1 = self.p1
        p2 = self.p2
        if not (p1.y < y < p2.y or p1.y > y > p2.y):
            return None

        r = Ray.fromPoints(p1, p2)
        return r.atHeight(y)

class ViewCornersRenderNode(rendernode.RenderNode):

    #
    # Renders the intersection of a horizontal plane with the view frustum
    # We only check the four vertical segments (with respect to the view angle)
    # and the four segments pointing outward from the viewer.
    # The four outward segments are checked first, if all four intersect they are used
    # as the corners. If only two outwards segments intersect, the two furthest
    # verticals are intersected and used as the last two corners.
    # If no outward segments intersect (rare) then all four verticals are used.

    # 0: bottom left, near
    # 1: bottom left, far
    # 2: top left, near
    # 3: top left, far
    # 4: bottom right, near
    # 5: bottom right, far
    # 6: top right, near
    # 7: top right, far

    verticalIndices = [
        (1, 3),
        (5, 7),
        (0, 2),
        (4, 6),
    ]
    outwardIndices = [
        (0, 1),
        (2, 3),
        (4, 5),
        (6, 7),
    ]

    def drawSelf(self):
        if self.sceneNode.corners is None or self.sceneNode.dimension is None:
            return
        corners = self.sceneNode.corners

        outwardSegments = [LineSegment(corners[i], corners[j]) for i, j in self.outwardIndices]
        verticalSegments = [LineSegment(corners[i], corners[j]) for i, j in self.verticalIndices]
        points = []

        for segment in outwardSegments:
            p = segment.atHeight(self.sceneNode.planeHeight)
            if p is not None:
                points.append(p)

        if len(points) < 4:
            # only intersected two outward segments. check the far verticals.
            for segment in verticalSegments[:2]:
                r = Ray.fromPoints(*segment)
                points.append(r.atHeight(self.sceneNode.planeHeight))

        if len(points) < 4:
            # intersected zero outward segments!
            # rarely occurs, the near verticals are 1/10 of a block tall
            for segment in verticalSegments[2:]:
                r = Ray.fromPoints(*segment)
                points.append(r.atHeight(self.sceneNode.planeHeight))

        if len(points) < 4:
            return

        p1, p2, p3, p4 = points[:4]
        points = [p1, p2, p4, p3, p1]

        with gl.glPushAttrib(GL.GL_DEPTH_BUFFER_BIT, GL.GL_COLOR_BUFFER_BIT):
            GL.glDepthMask(False)
            GL.glEnable(GL.GL_BLEND)
            GL.glVertexPointer(3, GL.GL_FLOAT, 0, numpy.array(points).ravel())

            GL.glLineWidth(3.0)

            GL.glColor(1, 1, .1, 0.5)
            GL.glDisable(GL.GL_DEPTH_TEST)
            GL.glDrawArrays(GL.GL_LINE_STRIP, 0, len(points))

class ViewCornersNode(scenenode.Node):
    RenderNodeClass = ViewCornersRenderNode

    _corners = None

    @property
    def corners(self):
        return self._corners

    @corners.setter
    def corners(self, value):
        self._corners = value
        self.dirty = True

    _dimension = None

    @property
    def dimension(self):
        return self._dimension

    @dimension.setter
    def dimension(self, value):
        self._dimension = value
        self.dirty = True

    _planeHeight = None

    @property
    def planeHeight(self):
        return self._planeHeight

    @planeHeight.setter
    def planeHeight(self, value):
        self._planeHeight = value
        self.dirty = True

class MinimapWorldView(WorldView):
    minScale = 1.
    def __init__(self, *a, **kw):
        WorldView.__init__(self, *a, **kw)
        self.setSizePolicy(QtGui.QSizePolicy.Policy.Minimum, QtGui.QSizePolicy.Policy.Minimum)
        self.scale = 1.0
        self.worldScene.minlod = 2
        self.viewCornersNode = ViewCornersNode()
        self.viewCornersNode.dimension = self.dimension
        self.worldNode.addChild(self.viewCornersNode)

    def createWorldScene(self):
        scene = super(MinimapWorldView, self).createWorldScene()
        self.layerToggleGroup.setVisibleLayers([Layer.Blocks])
        return scene

    def createCompass(self):
        compassNode = compass.CompassNode(small=True)
        compassNode.yawPitch = 180, 0
        return compassNode

    def updateMatrices(self):
        w, h = self.width(), self.height()
        w *= self.scale
        h *= self.scale

        projection = QtGui.QMatrix4x4()
        projection.ortho(-w/2, w/2, -h/2, h/2, -1000, 2000)
        self.matrixState.projection = projection

        modelview = QtGui.QMatrix4x4()
        modelview.rotate(90., 1., 0., 0.)
        modelview.translate(-self.centerPoint[0], 0, -self.centerPoint[2])
        self.matrixState.modelview = modelview

    def currentViewMatrixChanged(self, currentView):
        self.viewCornersNode.corners = currentView.getViewCorners()
        try:
            targetPoint, face = rayCastInBounds(Ray(currentView.centerPoint, currentView.cameraVector), self.dimension, 100)
            if targetPoint is None:
                raise MaxDistanceError
            planeHeight = targetPoint.y

        except MaxDistanceError:
            planeDistance = 20
            planeHeight = (currentView.centerPoint + currentView.cameraVector * planeDistance).y

        self.viewCornersNode.planeHeight = planeHeight

    def zoom(self, scale, (mx, my)):

        # Get mouse position in world coordinates
        worldPos = self.unprojectAtHeight(self.width()/2, self.height()/2, 0)

        if scale != self.scale:
            self.scale = scale

            # Get the new position under the mouse, find its distance from the old position,
            # and shift the centerPoint by that amount.

            newWorldPos = self.unprojectAtHeight(self.width()/2, self.height()/2, 0)
            delta = newWorldPos - worldPos
            self.centerPoint = self.centerPoint - delta

    def sizeHint(self):
        return QtCore.QSize(192, 192)

    def mousePressEvent(self, event):
        event.ignore()

    def mouseMoveEvent(self, event):
        event.ignore()

    def mouseReleaseEvent(self, event):
        event.ignore()


text:
from torch.utils.data import DataLoader
from oil.utils.utils import cosLr, islice, FixedNumpySeed,FixedPytorchSeed
from trainer.utils import LoaderTo
from oil.tuning.study import train_trial
from oil.datasetup.datasets import split_dataset
from oil.tuning.args import argupdated_config
import logging
import emlp.nn
import emlp.groups
import objax
from trainer.hamiltonian_dynamics import DoubleSpringPendulum
from jax import vmap,jit
from functools import partial
import jax.numpy as jnp
import numpy as np
import pandas as pd
import jax

levels = {'critical': logging.CRITICAL,'error': logging.ERROR,
                    'warn': logging.WARNING,'warning': logging.WARNING,
                    'info': logging.INFO,'debug': logging.DEBUG}

from trainer.classifier import Regressor,Classifier

def _rollout(model,x0,u,ts):
    return model.rollout(x0,u,ts)

def rollout(model):
    return objax.Jit(vmap(partial(_rollout,model),(0,0,None)),model.vars())

def rel_err(a,b):
    return jnp.sqrt(((a-b)**2).mean())/(jnp.sqrt((a**2).mean())+jnp.sqrt((b**2).mean()))#

def log_rollout_error(roller,k,minibatch):
    (z0, u,ts), gt_zts = minibatch
    pred_zt,_ = roller(z0,u[:,:k],ts[0][:k])
    d = pred_zt.shape[-1]
    state_errs = vmap(vmap(rel_err))(pred_zt[:,:,:d//2],gt_zts[:,:k,:d//2]) # (bs,T,)
    clamped_errs = jax.lax.clamp(1e-7,state_errs,np.inf)
    log_geo_mean = jnp.log(clamped_errs).mean()
    return log_geo_mean

class RolloutTrainer(Regressor):
    def __init__(self,model,*args,**kwargs):
        super().__init__(model,*args,**kwargs)
        self.rollout = rollout(model)
        self.loss = objax.Jit(self.loss,model.vars())
        #self.model = objax.Jit(self.model)
        self.gradvals = objax.Jit(objax.GradValues(self.loss,model.vars()))#objax.Jit(objax.GradValues(fastloss,model.vars()),model.vars())
        #self.model.predict = objax.Jit(objax.ForceArgs(model.__call__,training=False),model.vars())

    # def mse(self,minibatch):
    #     (x0,u, ts), true_x = minibatch
    #     pred_x,reg = self.rollout(x0,u,ts[0])
    #     return jnp.mean((pred_x - true_x)**2)

    def loss(self, minibatch):
        """ Standard cross-entropy loss """
        (x0,u, ts), true_x = minibatch
        pred_x,reg = self.rollout(x0,u,ts[0])
        # for k, v in self.model.vars().items():
        #     print(k,'NODE' in k)
        #l2 = sum((v.value ** 2).sum() for k, v in self.model.model.vars().items() if k.endswith('w'))
        return jnp.mean(jnp.abs(pred_x - true_x))+.05*jnp.mean(reg)

    def metrics(self, loader):
        #l2 = sum((v.value ** 2).sum() for k, v in self.model.vars().items() if k.endswith('w'))
        def _metrics(minibatch):
            (x0,u, ts), true_x = minibatch
            pred_x,reg = self.rollout(x0,u,ts[0])
            mse = np.asarray(jnp.mean(jnp.abs(pred_x - true_x)))
            reg = np.asarray(jnp.mean(reg))
            return pd.Series({'MAE':mse,'reg':reg})
        return self.evalAverageMetrics(loader,_metrics)
        #mse = lambda mb: np.asarray(self.mse(mb))
        #return {"MSE": self.evalAverageMetrics(loader, mse)}

    def logStuff(self, step, minibatch=None):
        extra_metrics = {}
        loader = self.dataloaders['_test_episodes']
        for k in [10,30,100]:
            extra_metrics[f'rgre@{k}']=np.exp(self.evalAverageMetrics(loader,partial(log_rollout_error,self.rollout,k)))
        self.logger.add_scalars('metrics', extra_metrics, step)
        super().logStuff(step,minibatch)

from mujoco_models import DeltaNN
import mujoco_models
from datasets import MujocoRegression,MujocoRollouts

def makeTrainer(*,network=DeltaNN,num_epochs=500,ndata=50000,seed=2021,chunk_len=20,
                bs=500,lr=1e-3,device='cuda',split={'train':5/6,'test':1/6},env='HopperFull-v0',
                net_config={'num_layers':2,'ch':128,'r2':2,'r3':4},log_level='warn',
                trainer_config={'log_dir':None,'log_args':{'minPeriod':.05,'timeFrac':.2},},#'early_stop_metric':'val_MSE'},
                save=False,):

    logging.getLogger().setLevel(levels[log_level])
    # Prep the datasets splits, model, and dataloaders
    with FixedNumpySeed(seed),FixedPytorchSeed(seed):
        base_ds = MujocoRegression(N=ndata,env=env,chunk_len=chunk_len)
        datasets = split_dataset(base_ds,splits=split)
        datasets['_test_episodes'] = MujocoRollouts(N=100,env=env)
    if network != mujoco_models.RPPall:
        net_config.pop('r2')
        net_config.pop('r3')
    model = network(base_ds.xdim,base_ds.udim,**net_config)
    dataloaders = {k:LoaderTo(DataLoader(v,batch_size=min(bs,len(v)),shuffle=(k=='train'),
                num_workers=0,pin_memory=False)) for k,v in datasets.items()}
    dataloaders['Train'] = dataloaders['train']
    #equivariance_test(model,dataloaders['train'],net_config['group'])
    opt_constr = objax.optimizer.Adam
    lr_sched = lambda e: lr#*cosLr(num_epochs)(e)#*min(1,e/(num_epochs/10))
    return RolloutTrainer(model,dataloaders,opt_constr,lr_sched,**trainer_config)

if __name__ == "__main__":
    cfg = argupdated_config(makeTrainer.__kwdefaults__,namespace=(mujoco_models,))
    trainer = makeTrainer(**cfg)
    trainer.train(cfg['num_epochs'])
    # Trial = ode_trial(makeTrainer)
    # cfg,outcome = Trial()
    # print(outcome)

max_stars_repo_name: mfinzi/residual-pathway-priors

max_issues_repo_name: mfinzi/residual-pathway-priors

max_forks_repo_name: mfinzi/residual-pathway-priors

reason: import numpy,import jax,from jax


text:
import config
import numpy as np
import tensorflow as tf

def get_tp_fp_fn(a, b):
    a = np.equal(a, 1)
    not_a = np.logical_not(a)
    b = np.equal(b, 1)
    not_b = np.logical_not(b)

    tp = np.logical_and(a, b).sum().astype(np.float64)
    fp = np.logical_and(a, not_b).sum().astype(np.float64)
    fn = np.logical_and(not_a, b).sum().astype(np.float64)

    return tp, fp, fn

def jaccard(a, b):
    tp, fp, fn = get_tp_fp_fn(a, b)

    jaccard_coef = None
    if tp + fp + fn == 0:
        jaccard_coef = 1.
    else:
        jaccard_coef = tp / (tp + fp + fn)

    return jaccard_coef

def dice(a, b):
    tp, fp, fn = get_tp_fp_fn(a, b)

    dice_coef = None
    if tp + fp + fn == 0:
        dice_coef = 1.
    else:
        dice_coef = (2 * tp) / (2 * tp + fp + fn)

    return dice_coef

def jaccard_and_dice(preds, gts, thr):
    jaccard_and_dice_res = np.zeros((len(preds), 2))

    for idx, (pred, gt) in enumerate(zip(preds, gts)):
        pred = (pred >= thr)

        jaccard_coef = jaccard(pred, gt)
        dice_coef = dice(pred, gt)

        jaccard_and_dice_res[idx] = (jaccard_coef, dice_coef)

    return jaccard_and_dice_res

# tensorflow implementation (with thr)
def tf_get_tp_fp_fn(a_in, b_in):
    a = tf.greater_equal(a_in, config.PRED_THR)
    not_a = tf.logical_not(a)
    b = tf.greater_equal(b_in, config.PRED_THR)
    not_b = tf.logical_not(b)

    tp_and = tf.logical_and(a, b)
    tp_count = tf.count_nonzero(tp_and)
    tp = tf.cast(tp_count, tf.float64)

    fp_and = tf.logical_and(a, not_b)
    fp_count = tf.count_nonzero(fp_and)
    fp = tf.cast(fp_count, tf.float64)

    fn_and = tf.logical_and(not_a, b)
    fn_count = tf.count_nonzero(fn_and)
    fn = tf.cast(fn_count, tf.float64)

    return tp, fp, fn

def tf_jaccard(a, b):
    tp, fp, fn = tf_get_tp_fp_fn(a, b)
    jaccard_coef = tf.cond(tf.equal(tp + fp + fn, 0),
                           lambda: tf.constant(1, tf.float64),
                           lambda: tp / (tp + fp + fn))
    return jaccard_coef

def tf_dice(a, b):
    tp, fp, fn = tf_get_tp_fp_fn(a, b)
    dice_coef = tf.cond(tf.equal(tp + fp + fn, 0),
                        lambda: tf.constant(1, tf.float64),
                        lambda: (2 * tp) / (2 * tp + fp + fn))
    return dice_coef


text:
# ----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# ----------------------------------------------------------------------------

from __future__ import absolute_import, division, print_function

import numpy as np

from ._base import Ordination, OrdinationResults
from ._utils import svd_rank

class CA(Ordination):
    r"""Compute correspondence analysis, a multivariate statistical
    technique for ordination.

    In general, rows in the data table will correspond to sites and
    columns to species, but the method is symmetric. In order to
    measure the correspondence between rows and columns, the
    :math:`\chi^2` distance is used, and those distances are preserved
    in the transformed space. The :math:`\chi^2` distance doesn't take
    double zeros into account, and so it is expected to produce better
    ordination that PCA when the data has lots of zero values.

    It is related to Principal Component Analysis (PCA) but it should
    be preferred in the case of steep or long gradients, that is, when
    there are many zeros in the input data matrix.

    Parameters
    ----------
    X : array_like
        Contingency table. It can be applied to different kinds of
        data tables but data must be non-negative and dimensionally
        homogeneous (quantitative or binary).

    Notes
    -----
    The algorithm is based on [1]_, \S 9.4.1., and is expected to give
    the same results as ``cca(X)`` in R's package vegan.

    See Also
    --------
    CCA

    References
    ----------
    .. [1] Legendre P. and Legendre L. 1998. Numerical
       Ecology. Elsevier, Amsterdam.

    """
    short_method_name = 'CA'
    long_method_name = 'Canonical Analysis'

    def __init__(self, X, row_ids, column_ids):
        self.X = np.asarray(X, dtype=np.float64)
        self._ca()
        self.row_ids = row_ids
        self.column_ids = column_ids

    def _ca(self):
        X = self.X
        r, c = X.shape

        if X.min() < 0:
            raise ValueError("Input matrix elements must be non-negative.")

        # Step 1 (similar to Pearson chi-square statistic)
        grand_total = X.sum()
        Q = X / grand_total

        column_marginals = Q.sum(axis=0)
        row_marginals = Q.sum(axis=1)
        # Let's store them since they're needed to compute scores
        self.column_marginals = column_marginals
        self.row_marginals = row_marginals

        # Formula 9.32 in Lagrange & Lagrange (1998). Notice that it's
        # an scaled version of the contribution of each cell towards
        # Pearson chi-square statistic.
        expected = np.outer(row_marginals, column_marginals)
        Q_bar = (Q - expected) / np.sqrt(expected)  # Eq. 9.32

        # Step 2 (Singular Value Decomposition)
        U_hat, W, Ut = np.linalg.svd(Q_bar, full_matrices=False)
        # Due to the centering, there are at most min(r, c) - 1 non-zero
        # eigenvalues (which are all positive)
        rank = svd_rank(Q_bar.shape, W)
        assert rank <= min(r, c) - 1
        self.U_hat = U_hat[:, :rank]
        self.W = W[:rank]
        self.U = Ut[:rank].T

    def scores(self, scaling):
        r"""Compute site and species scores for different scalings.

        Parameters
        ----------
        scaling : int

            For a more detailed explanation of the interpretation, check
            Legendre & Legendre 1998, section 9.4.3. The notes that
            follow are quick recommendations.

            Scaling type 1 maintains :math:`\chi^2` distances between
            rows (sites): in the transformed space, the euclidean
            distances between rows are equal to the :math:`\chi^2`
            distances between rows in the original space. It should be
            used when studying the ordination of sites. Rows (sites)
            that are near a column (species) have high contributions
            from it.

            Scaling type 2 preserves :math:`\chi^2` distances between
            columns (species), so euclidean distance between columns
            after transformation is equal to :math:`\chi^2` distance
            between columns in the original space. It is best used
            when we are interested in the ordination of species. A
            column (species) that is next to a row (site) means that
            it is more abundant there.

            Other types of scalings are currently not implemented, as
            they're less used by ecologists (Legendre & Legendre 1998,
            p. 456).

            In general, species appearing far from the center of the
            biplot and far from its edges will probably exhibit better
            relationships than species either in the center (may be
            multimodal species, not related to the shown ordination
            axes...) or the edges (sparse species...).

        Returns
        -------
        OrdinationResults
            Object that stores the computed eigenvalues, the
            proportion explained by each of them (per unit),
            transformed coordinates, etc.

        See Also
        --------
        OrdinationResults
        """

        if scaling not in {1, 2}:
            raise NotImplementedError(
                "Scaling {0} not implemented.".format(scaling))
        # Both scalings are a bit intertwined, so we'll compute both and
        # then choose
        V = self.column_marginals[:, None]**-0.5 * self.U
        V_hat = self.row_marginals[:, None]**-0.5 * self.U_hat
        F = V_hat * self.W
        # According to Formula 9.43, this should hold
        # assert np.allclose(F, (row_marginals**-1)[:, None] * Q.dot(V))
        # but it doesn't (notice that W**2==Lambda):
        # (9.43a) F = V_hat W = D(p_i+)^{-1/2} U_hat W
        #           = D(p_i+)^{-1/2} Q_bar U W^{-1} W  (substituting 9.38)
        #           = D(p_i+)^{-1/2} Q_bar U
        # (9.43b) F = D(p_i+)^{-1} Q V
        #           = D(p_i+)^{-1} Q D(p_+j)^{-1/2} U  (substituting 9.41)
        #           = D(p_i+)^{-1/2} D(p_i+)^{-1/2} Q D(p_+j)^{-1/2} U
        #           = D(p_i+)^{-1/2} Q_tilde U         (using 9.40)
        # It holds if we replace Q in 9.43b with Q after centering, ie
        # assert np.allclose(
        #    F,
        #    (row_marginals**-1)[:, None] * (Q - expected).dot(V))
        # Comparing results with vegan and the examples in the book, 9.43a
        # is the right one. The same issue happens in 9.44, where also
        # 9.44a is the one that matches vegan's output.
        # (9.44a) F_hat = V W = D(p_+j)^{-1/2} U W
        #               = D(p_+j)^{-1/2} Q_bar' U_hat W^{-1} W (using 9.39)
        #               = D(p_+j)^{-1/2} Q_bar' U_hat
        # (9.44b) F_hat = D(p_+j)^{-1} Q' V_hat
        #               = D(p_+j)^{-1/2} Q_tilde' U_hat (using 9.40 and 9.42)
        F_hat = V * self.W

        # Eigenvalues
        eigvals = self.W**2

        # Species scores
        species_scores = [V, F_hat][scaling - 1]
        # Site scores (weighted averages of species scores)
        site_scores = [F, V_hat][scaling - 1]
        return OrdinationResults(eigvals=eigvals, species=species_scores,
                                 site=site_scores, site_ids=self.row_ids,
                                 species_ids=self.column_ids)

max_stars_repo_name: Kleptobismol/scikit-bio

max_issues_repo_name: Kleptobismol/scikit-bio

max_forks_repo_name: Kleptobismol/scikit-bio


text:
import torch
import torch.nn as nn
import torch.nn.functional as F
import constants

class WordEmbedding(nn.Module):

    # word_embeddings comes from numpy
    def __init__(self, word_embeddings):
        super(WordEmbedding, self).__init__()
        self.word_embedding = nn.Embedding(num_embeddings=word_embeddings.shape[0],
                                           embedding_dim=word_embeddings.shape[1])
        #Cast to float because the character embeding will be returned as a float, and we need to concatenate the two
        self.word_embedding.weight = nn.Parameter(torch.from_numpy(word_embeddings).float())

        # Only the unknown embedding requires grad
        self.word_embedding.weight.requires_grad = False
        #self.word_embedding.weight[constants.UNK_ID].requires_grad = True

        del word_embeddings

    def forward(self, input_context, input_question):

        context_word_emb = self.word_embedding(input_context)

        context_word_emb = F.dropout(context_word_emb, p=0.1, training=self.training)

        question_word_emb = self.word_embedding(input_question)

        question_word_emb = F.dropout(question_word_emb, p=0.1, training=self.training)

        return context_word_emb, question_word_emb

max_stars_repo_name: brinaseidel/squad

max_issues_repo_name: brinaseidel/squad

max_forks_repo_name: brinaseidel/squad

reason: from numpy


text:
# -*- coding: utf-8 -*-
# =============================================================================
# Created by: Natalie LaLuzerne
# =============================================================================
import numpy as np
from scipy import stats

# =============================================================================
#   location defines the datatype to represent the location of the tumor
# =============================================================================
class location:
    LEFT = -1
    RIGHT = 1
    NONE = 0

# =============================================================================
#   slope thresholds. derived from temperature_processing_investigation
#   slope falls under minSlopeThreshold: tumor on left
#   slope exceeds maxSlopeThreshold: tumor in right
#   slope within bounds of both thresholds: no tumor
# =============================================================================
minSlopeThreshold = -0.006495790
maxSlopeThreshold = 0.007409563

def temp_proc(bins):

    # =========================================================================
    #   initialize tumor location storage structure
    # =========================================================================
    locationList = []

    # =========================================================================
    #   Apply background filter
    # =========================================================================
    rowAvg = [np.mean(bins[rw]) for rw in range(bins.shape[0])]
    minAvg = min(rowAvg)
    idx = rowAvg.index(minAvg)
    std = np.std(bins[idx])
    bins = np.array([[x if x >= (minAvg - (2 * std)) else 0 for x in y] for y in bins])

    # =========================================================================
    #   Analyze row by row each of the temperature profiles
    # =========================================================================
    for row in reversed(range(bins.shape[0])):
        # =====================================================================
        #   get the current row to analyze
        # =====================================================================
        currentRow = np.array(bins[row])

        # =====================================================================
        #   calculate x and y coordinates of each temperature for y = mx + b
        # =====================================================================
        Y = [i for i in currentRow if i != 0]
        X = [i for i in range(len(Y))]

        # =====================================================================
        #   calculate slope
        # =====================================================================
        slope, intercept, r_value, p_value, std_err = stats.linregress(X, Y)

        # =====================================================================
        #   compare slope to thresholds and store location estimate
        # =====================================================================
        if(slope > maxSlopeThreshold):
            locationList.append(location.RIGHT)
        elif(slope < minSlopeThreshold):
            locationList.append(location.LEFT)
        else:
            locationList.append(location.NONE)

    # =========================================================================
    #   sum the location values in the list, return diagnosis
    #   sum > 0: majority of slopes were positive, tumor on right
    #   sum < 0: majority of slopes were negative, tumor on left
    #   sum = 0: slopes were balanced, no tumor
    # =========================================================================
    SUM = np.sum(locationList)
    if(SUM > 0):
        return([0,1]) #right
    elif(SUM < 0):
        return([1,0]) #left
    else:
        return([0,0]) #none

max_stars_repo_name: nlaluzerne/TIGRE

max_issues_repo_name: nlaluzerne/TIGRE

max_forks_repo_name: nlaluzerne/TIGRE

reason: import numpy,from scipy


text:
#
# Copyright (C) 2019 Databricks, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import datetime
from decimal import Decimal
from distutils.version import LooseVersion

import numpy as np
import pandas as pd

from databricks import koalas as ks
from databricks.koalas.testing.utils import ReusedSQLTestCase
from databricks.koalas.utils import name_like_string

class ReshapeTest(ReusedSQLTestCase):
    def test_get_dummies(self):
        for pdf_or_ps in [
            pd.Series([1, 1, 1, 2, 2, 1, 3, 4]),
            # pd.Series([1, 1, 1, 2, 2, 1, 3, 4], dtype='category'),
            # pd.Series(pd.Categorical([1, 1, 1, 2, 2, 1, 3, 4],
            #                          categories=[4, 3, 2, 1])),
            pd.DataFrame(
                {
                    "a": [1, 2, 3, 4, 4, 3, 2, 1],
                    # 'b': pd.Categorical(list('abcdabcd')),
                    "b": list("abcdabcd"),
                }
            ),
        ]:
            kdf_or_kser = ks.from_pandas(pdf_or_ps)

            self.assert_eq(ks.get_dummies(kdf_or_kser), pd.get_dummies(pdf_or_ps, dtype=np.int8))

        kser = ks.Series([1, 1, 1, 2, 2, 1, 3, 4])
        with self.assertRaisesRegex(
            NotImplementedError, "get_dummies currently does not support sparse"
        ):
            ks.get_dummies(kser, sparse=True)

    def test_get_dummies_object(self):
        pdf = pd.DataFrame(
            {
                "a": [1, 2, 3, 4, 4, 3, 2, 1],
                # 'a': pd.Categorical([1, 2, 3, 4, 4, 3, 2, 1]),
                "b": list("abcdabcd"),
                # 'c': pd.Categorical(list('abcdabcd')),
                "c": list("abcdabcd"),
            }
        )
        kdf = ks.from_pandas(pdf)

        # Explicitly exclude object columns
        self.assert_eq(
            ks.get_dummies(kdf, columns=["a", "c"]),
            pd.get_dummies(pdf, columns=["a", "c"], dtype=np.int8),
        )

        self.assert_eq(ks.get_dummies(kdf), pd.get_dummies(pdf, dtype=np.int8))
        self.assert_eq(ks.get_dummies(kdf.b), pd.get_dummies(pdf.b, dtype=np.int8))
        self.assert_eq(
            ks.get_dummies(kdf, columns=["b"]), pd.get_dummies(pdf, columns=["b"], dtype=np.int8)
        )

    def test_get_dummies_date_datetime(self):
        pdf = pd.DataFrame(
            {
                "d": [
                    datetime.date(2019, 1, 1),
                    datetime.date(2019, 1, 2),
                    datetime.date(2019, 1, 1),
                ],
                "dt": [
                    datetime.datetime(2019, 1, 1, 0, 0, 0),
                    datetime.datetime(2019, 1, 1, 0, 0, 1),
                    datetime.datetime(2019, 1, 1, 0, 0, 0),
                ],
            }
        )
        kdf = ks.from_pandas(pdf)

        self.assert_eq(ks.get_dummies(kdf), pd.get_dummies(pdf, dtype=np.int8))
        self.assert_eq(ks.get_dummies(kdf.d), pd.get_dummies(pdf.d, dtype=np.int8))
        self.assert_eq(ks.get_dummies(kdf.dt), pd.get_dummies(pdf.dt, dtype=np.int8))

    def test_get_dummies_boolean(self):
        pdf = pd.DataFrame({"b": [True, False, True]})
        kdf = ks.from_pandas(pdf)

        self.assert_eq(ks.get_dummies(kdf), pd.get_dummies(pdf, dtype=np.int8))
        self.assert_eq(ks.get_dummies(kdf.b), pd.get_dummies(pdf.b, dtype=np.int8))

    def test_get_dummies_decimal(self):
        pdf = pd.DataFrame({"d": [Decimal(1.0), Decimal(2.0), Decimal(1)]})
        kdf = ks.from_pandas(pdf)

        self.assert_eq(ks.get_dummies(kdf), pd.get_dummies(pdf, dtype=np.int8))
        self.assert_eq(ks.get_dummies(kdf.d), pd.get_dummies(pdf.d, dtype=np.int8), almost=True)

    def test_get_dummies_kwargs(self):
        # pser = pd.Series([1, 1, 1, 2, 2, 1, 3, 4], dtype='category')
        pser = pd.Series([1, 1, 1, 2, 2, 1, 3, 4])
        kser = ks.from_pandas(pser)
        self.assert_eq(
            ks.get_dummies(kser, prefix="X", prefix_sep="-"),
            pd.get_dummies(pser, prefix="X", prefix_sep="-", dtype=np.int8),
        )

        self.assert_eq(
            ks.get_dummies(kser, drop_first=True),
            pd.get_dummies(pser, drop_first=True, dtype=np.int8),
        )

        # nan
        # pser = pd.Series([1, 1, 1, 2, np.nan, 3, np.nan, 5], dtype='category')
        pser = pd.Series([1, 1, 1, 2, np.nan, 3, np.nan, 5])
        kser = ks.from_pandas(pser)
        self.assert_eq(ks.get_dummies(kser), pd.get_dummies(pser, dtype=np.int8), almost=True)

        # dummy_na
        self.assert_eq(
            ks.get_dummies(kser, dummy_na=True), pd.get_dummies(pser, dummy_na=True, dtype=np.int8)
        )

    def test_get_dummies_prefix(self):
        pdf = pd.DataFrame({"A": ["a", "b", "a"], "B": ["b", "a", "c"], "D": [0, 0, 1],})
        kdf = ks.from_pandas(pdf)

        self.assert_eq(
            ks.get_dummies(kdf, prefix=["foo", "bar"]),
            pd.get_dummies(pdf, prefix=["foo", "bar"], dtype=np.int8),
        )

        self.assert_eq(
            ks.get_dummies(kdf, prefix=["foo"], columns=["B"]),
            pd.get_dummies(pdf, prefix=["foo"], columns=["B"], dtype=np.int8),
        )

        self.assert_eq(
            ks.get_dummies(kdf, prefix={"A": "foo", "B": "bar"}),
            pd.get_dummies(pdf, prefix={"A": "foo", "B": "bar"}, dtype=np.int8),
        )

        self.assert_eq(
            ks.get_dummies(kdf, prefix={"B": "foo", "A": "bar"}),
            pd.get_dummies(pdf, prefix={"B": "foo", "A": "bar"}, dtype=np.int8),
        )

        self.assert_eq(
            ks.get_dummies(kdf, prefix={"A": "foo", "B": "bar"}, columns=["A", "B"]),
            pd.get_dummies(pdf, prefix={"A": "foo", "B": "bar"}, columns=["A", "B"], dtype=np.int8),
        )

        with self.assertRaisesRegex(NotImplementedError, "string types"):
            ks.get_dummies(kdf, prefix="foo")
        with self.assertRaisesRegex(ValueError, "Length of 'prefix' \\(1\\) .* \\(2\\)"):
            ks.get_dummies(kdf, prefix=["foo"])
        with self.assertRaisesRegex(ValueError, "Length of 'prefix' \\(2\\) .* \\(1\\)"):
            ks.get_dummies(kdf, prefix=["foo", "bar"], columns=["B"])

        pser = pd.Series([1, 1, 1, 2, 2, 1, 3, 4], name="A")
        kser = ks.from_pandas(pser)

        self.assert_eq(
            ks.get_dummies(kser, prefix="foo"), pd.get_dummies(pser, prefix="foo", dtype=np.int8)
        )

        # columns are ignored.
        self.assert_eq(
            ks.get_dummies(kser, prefix=["foo"], columns=["B"]),
            pd.get_dummies(pser, prefix=["foo"], columns=["B"], dtype=np.int8),
        )

    def test_get_dummies_dtype(self):
        pdf = pd.DataFrame(
            {
                # "A": pd.Categorical(['a', 'b', 'a'], categories=['a', 'b', 'c']),
                "A": ["a", "b", "a"],
                "B": [0, 0, 1],
            }
        )
        kdf = ks.from_pandas(pdf)

        if LooseVersion("0.23.0") <= LooseVersion(pd.__version__):
            exp = pd.get_dummies(pdf, dtype="float64")
        else:
            exp = pd.get_dummies(pdf)
            exp = exp.astype({"A_a": "float64", "A_b": "float64"})
        res = ks.get_dummies(kdf, dtype="float64")
        self.assert_eq(res, exp)

    def test_get_dummies_multiindex_columns(self):
        pdf = pd.DataFrame(
            {
                ("x", "a", "1"): [1, 2, 3, 4, 4, 3, 2, 1],
                ("x", "b", "2"): list("abcdabcd"),
                ("y", "c", "3"): list("abcdabcd"),
            }
        )
        kdf = ks.from_pandas(pdf)

        self.assert_eq(
            ks.get_dummies(kdf), pd.get_dummies(pdf, dtype=np.int8).rename(columns=name_like_string)
        )
        self.assert_eq(
            ks.get_dummies(kdf, columns=[("y", "c", "3"), ("x", "a", "1")]),
            pd.get_dummies(pdf, columns=[("y", "c", "3"), ("x", "a", "1")], dtype=np.int8).rename(
                columns=name_like_string
            ),
        )
        self.assert_eq(
            ks.get_dummies(kdf, columns=["x"]),
            pd.get_dummies(pdf, columns=["x"], dtype=np.int8).rename(columns=name_like_string),
        )
        self.assert_eq(
            ks.get_dummies(kdf, columns=("x", "a")),
            pd.get_dummies(pdf, columns=("x", "a"), dtype=np.int8).rename(columns=name_like_string),
        )
        self.assert_eq(
            ks.get_dummies(kdf, columns=["x"]),
            pd.get_dummies(pdf, columns=["x"], dtype=np.int8).rename(columns=name_like_string),
        )

        self.assertRaises(KeyError, lambda: ks.get_dummies(kdf, columns=["z"]))
        self.assertRaises(KeyError, lambda: ks.get_dummies(kdf, columns=("x", "c")))
        self.assertRaises(ValueError, lambda: ks.get_dummies(kdf, columns=[("x",), "c"]))
        self.assertRaises(TypeError, lambda: ks.get_dummies(kdf, columns="x"))

max_stars_repo_name: xinrong-databricks/koalas

max_issues_repo_name: xinrong-databricks/koalas

max_forks_repo_name: xinrong-databricks/koalas


text:
import tensorflow as tf
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
train_x, train_y = mnist.train.next_batch(10000)
test_x, test_y = mnist.test.next_batch(500)

K_neighbours = 8
x_train = tf.placeholder(tf.float32, [None, 784])
x_test = tf.placeholder(tf.float32, [784])

distance = tf.negative(tf.sqrt(tf.reduce_sum(tf.square(tf.subtract(x_train, x_test)), reduction_indices = 1)))

indices = tf.nn.top_k(distance, k=K_neighbours, sorted=False)

init = tf.global_variables_initializer()
correct_class = 0
num_steps = 2000

with tf.Session() as session:
    session.run(init)

    for i in range(len(test_x)):
        values_indices = session.run(indices, feed_dict={x_train: train_x, x_test: test_x[i, :]})

        # predicting label for test data
        counter = np.zeros(10)
        for j in range(K_neighbours):
            counter[np.argmax(train_y[values_indices.indices[j]])] += 1

        prediction = np.argmax(counter)

        if prediction == np.argmax(test_y[i]):
            correct_class += 1.0

    print('Accuracy:', (correct_class / len(test_x)) * 100, '%')

max_stars_repo_name: LITDataScience/image-classifier