text:
#******************************************************************************
#
# tempoGAN: A Temporally Coherent, Volumetric GAN for Super-resolution Fluid Flow
# Copyright 2018 You Xie, Erik Franz, Mengyu Chu, Nils Thuerey
#
# This program is free software, distributed under the terms of the
# Apache License, Version 2.0
# http://www.apache.org/licenses/LICENSE-2.0
#
#******************************************************************************

import os, math
import shutil, sys
from random import seed, random, randrange
import uniio
import numpy as np
import scipy.misc
import scipy.ndimage
import imageio

# check whether matplotlib is available to generate vector/quiver plots
import imp
try:
	imp.find_module('matplotlib')
	import matplotlib.pyplot
	found_matplotlib = True
except ImportError:
    found_matplotlib = False
#import matplotlib.pyplot as plt

# global channel keys, have to be one char
C_KEY_DEFAULT = 'd'
C_KEY_VELOCITY = 'v'
C_KEY_VORTICITY = 'x'
C_KEY_POSITION = 'p'

DATA_KEY_LOW = 0
DATA_KEY_HIGH= 1

#keys for augmentation operations
AOPS_KEY_ROTATE = 'rot'
AOPS_KEY_SCALE = 'scale'
AOPS_KEY_ROT90 = 'rot90'
AOPS_KEY_FLIP = 'flip'

seed( 42 )

# default channel layouts
C_LAYOUT = {
	'dens':C_KEY_DEFAULT,
	'dens_vel':'d,vx,vy,vz'
	}

class TileCreator(object):

	def __init__(self, tileSizeLow, simSizeLow=64, upres=2, dim=2, dim_t=1, overlapping=0, densityMinimum=0.02, premadeTiles=False, partTrain=0.8, partTest=0.2, partVal=0, channelLayout_low=C_LAYOUT['dens_vel'], channelLayout_high=C_LAYOUT['dens'], highIsLabel=False, loadPN=False, padding=0):
		'''
			tileSizeLow, simSizeLow: int, [int,int] if 2D, [int,int,int]
			channelLayout: 'key,key,...'
				the keys are NOT case sensitive and leading and trailing whitespace characters are REMOVED.
				key:
					default:	d
					velocity:	v[label](x|y|z)
						label can be arbitrary or empty,
						key must be unique and x,y must exist while z is optional in 2D, x,y,z must exist in 3D.
						if x does not exist y,z will be ignored (treaded as 'd').
					rest is not yet supported

			premadeTiles: cut regular tiles when loading data, can't use data augmentation
			part(Train|Test|Val): relative size of the different data sets
			highIsLabel: high data is not augmented
			loadHigh:
			simPath: path to the uni simulation files
			loadPath: packed simulations are stored here
		'''

		# DATA DIMENSION
		self.dim_t = dim_t # same for hi_res or low_res
		if dim!=2 and dim!=3:
			self.TCError('Data dimension must be 2 or 3.')
		self.dim = dim
		# TILE SIZE
		if np.isscalar(tileSizeLow):
			self.tileSizeLow = [tileSizeLow, tileSizeLow, tileSizeLow]
		elif len(tileSizeLow)==2 and self.dim==2:
			self.tileSizeLow = [1]+tileSizeLow
		elif len(tileSizeLow)==3:
			self.tileSizeLow = tileSizeLow
		else:
			self.TCError('Tile size mismatch.')
		self.tileSizeLow = np.asarray(self.tileSizeLow)
		#SIM SIZE
		if np.isscalar(simSizeLow):
			self.simSizeLow = [simSizeLow, simSizeLow, simSizeLow]
		elif len(simSizeLow)==2 and self.dim==2:
			self.simSizeLow = [1]+simSizeLow
		elif len(simSizeLow)==3:
			self.simSizeLow = simSizeLow
		else:
			self.TCError('Simulation size mismatch.')
		self.simSizeLow = np.asarray(self.simSizeLow)

		if upres < 1:
			self.TCError('Upres must be at least 1.')
		self.upres = upres
		if not highIsLabel:
			self.tileSizeHigh = self.tileSizeLow*upres
			self.simSizeHigh = self.simSizeLow*upres
		else:
			self.tileSizeHigh = np.asarray([1])
			self.simSizeHigh = np.asarray([1])

		if self.dim==2:
			self.tileSizeLow[0]=1
			self.tileSizeHigh[0]=1
			self.simSizeLow[0]=1
			self.simSizeHigh[0]=1
		if np.less(self.simSizeLow, self.tileSizeLow).any():
			self.TCError('Tile size {} can not be larger than sim size {}, {}.'.format(self.tileSizeLow,self.simSizeLow))

		if densityMinimum<0.:
			self.TCError('densityMinimum can not be negative.')
		self.densityMinimum = densityMinimum
		self.premadeTiles = premadeTiles
		self.useDataAug = False

		#CHANNELS
		self.c_lists = {}
		self.c_low, self.c_lists[DATA_KEY_LOW] = self.parseChannels(channelLayout_low)
		self.c_high, self.c_lists[DATA_KEY_HIGH] = self.parseChannels(channelLayout_high)

		# print info
		print('\n')
		print('Dimension: {}, time dimension: {}'.format(self.dim,self.dim_t))
		print('Low-res data:')
		print('  channel layout: {}'.format(self.c_low))
		print('  default channels: {}'.format(self.c_lists[DATA_KEY_LOW][C_KEY_DEFAULT]))
		if len(self.c_lists[DATA_KEY_LOW][C_KEY_VELOCITY])>0:
			print('  velocity channels: {}'.format(self.c_lists[DATA_KEY_LOW][C_KEY_VELOCITY]))
		if len(self.c_lists[DATA_KEY_LOW][C_KEY_VORTICITY])>0:
			print('  vorticity channels: {}'.format(self.c_lists[DATA_KEY_LOW][C_KEY_VORTICITY]))
		print('High-res data:')
		if highIsLabel:
			print('  is Label')
		print('  channel layout: {}'.format(self.c_high))
		print('  default channels: {}'.format(self.c_lists[DATA_KEY_HIGH][C_KEY_DEFAULT]))
		if len(self.c_lists[DATA_KEY_HIGH][C_KEY_VELOCITY])>0:
			print('  velocity channels: {}'.format(self.c_lists[DATA_KEY_HIGH][C_KEY_VELOCITY]))
		if len(self.c_lists[DATA_KEY_HIGH][C_KEY_VORTICITY])>0:
			print('  vorticity channels: {}'.format(self.c_lists[DATA_KEY_HIGH][C_KEY_VORTICITY]))
		#self.channels=len(self.c)

		self.data_flags = {
			DATA_KEY_LOW:{
				'isLabel':False,
				'channels':len(self.c_low),
				C_KEY_VELOCITY:len(self.c_lists[DATA_KEY_LOW][C_KEY_VELOCITY])>0,
				C_KEY_VORTICITY:len(self.c_lists[DATA_KEY_LOW][C_KEY_VORTICITY])>0,
				C_KEY_POSITION:False
			},
			DATA_KEY_HIGH:{
				'isLabel':highIsLabel,
				'channels':len(self.c_high),
				C_KEY_VELOCITY:len(self.c_lists[DATA_KEY_HIGH][C_KEY_VELOCITY])>0,
				C_KEY_VORTICITY:len(self.c_lists[DATA_KEY_HIGH][C_KEY_VORTICITY])>0,
				C_KEY_POSITION:False
			}
		}
		if loadPN:
			self.TCError('prev and next tiles not supported.')
		self.hasPN = loadPN
		self.padding=padding

		#if self.hasPN:
		#[z,y,x, velocities an/or position if enabled (density,vel,vel,vel, pos, pos [,pos])]

		#DATA SHAPES
		self.tile_shape_low = np.append(self.tileSizeLow, [self.data_flags[DATA_KEY_LOW]['channels']])
		self.frame_shape_low = np.append(self.simSizeLow, [self.data_flags[DATA_KEY_LOW]['channels']])
		if not self.data_flags[DATA_KEY_HIGH]['isLabel']:
			self.tile_shape_high = np.append(self.tileSizeHigh, [self.data_flags[DATA_KEY_HIGH]['channels']])
			self.frame_shape_high = np.append(self.simSizeHigh, [self.data_flags[DATA_KEY_HIGH]['channels']])
		else:
			self.tile_shape_high = self.tileSizeHigh[:]
			self.frame_shape_high = self.simSizeHigh[:]

		self.densityThreshold = (self.densityMinimum * self.tile_shape_low[0] * self.tile_shape_low[1] * self.tile_shape_low[2])

		self.data = {
			DATA_KEY_LOW:[],
			DATA_KEY_HIGH:[]
		}

		all=partTrain+partTest+partVal
		self.part_train=partTrain/all
		self.part_test=partTest/all
		self.part_validation=partVal/all

	def initDataAugmentation(self, rot=2, minScale=0.85, maxScale=1.15 ,flip=True):
		'''
			set up data augmentation
			rot: 1: 90 degree rotations; 2: full rotation; else: nop rotation
			Scale: if both 1 disable scaling
		'''
		self.useDataAug = True

		"""
			specify the special augmentation operation needed for some channel types here
			will only be applyed if the specified channel type is in the data
			** Tempo Datum may have multiple channels as coherent frames, [batch, z, y, x, t*channels]
			** They are reshaped first before these aops, [batch, z, y, x, t, channels], and then reshape back
			** Because of this extra time dimention, all aops can only do isolate calculations, for e.g., value scaling,
			** Any calculation relay on neighborhood will be wrong, for e.g., spacial scaling (zoom).
		"""
		self.aops = {
			DATA_KEY_LOW:{
				AOPS_KEY_ROTATE:{
					C_KEY_VELOCITY:self.rotateVelocities,
					C_KEY_VORTICITY:self.rotateVelocities
				},
				AOPS_KEY_SCALE:{
					C_KEY_VELOCITY:self.scaleVelocities,
					C_KEY_VORTICITY:self.scaleVelocities
				},
				AOPS_KEY_ROT90:{
					C_KEY_VELOCITY:self.rotate90Velocities,
					C_KEY_VORTICITY:self.rotate90Velocities
				},
				AOPS_KEY_FLIP:{
					C_KEY_VELOCITY:self.flipVelocities,
					C_KEY_VORTICITY:self.flipVelocities
				}

			},
			DATA_KEY_HIGH:{
				AOPS_KEY_ROTATE:{
					C_KEY_VELOCITY:self.rotateVelocities,
					C_KEY_VORTICITY:self.rotateVelocities
				},
				AOPS_KEY_SCALE:{
					C_KEY_VELOCITY:self.scaleVelocities,
					C_KEY_VORTICITY:self.scaleVelocities
				},
				AOPS_KEY_ROT90:{
					C_KEY_VELOCITY:self.rotate90Velocities,
					C_KEY_VORTICITY:self.rotate90Velocities
				},
				AOPS_KEY_FLIP:{
					C_KEY_VELOCITY:self.flipVelocities,
					C_KEY_VORTICITY:self.flipVelocities
				}
			}
		}

		msg = 'data augmentation: '

		if rot==2:
			self.do_rotation = True
			self.do_rot90 = False
			msg += 'rotation, '
		elif rot==1:
			self.do_rotation = False
			self.do_rot90 = True
			msg += 'rot90, '
			z=(2,1)
			nz=(1,2)
			x=(1,0)
			y=(0,2)
			nx=(0,1)
			ny=(2,0)
			# thanks to http://www.euclideanspace.com/maths/discrete/groups/categorise/finite/cube/
			self.cube_rot = {2: [[],[z],[z,z],[nz]], 3: [[],[x],[y],[x,x],[x,y],[y,x],[y,y],[nx],[x,x,y],[x,y,x],[x,y,y],[y,x,x],[y,y,x],[ny],[nx,y],[x,x,y,x],[x,x,y,y],[x,y,x,x],[x,ny],[y,nx],[ny,x],[nx,y,x],[x,y,nx],[x,ny,x]]}
		else:
			self.do_rotation = False
			self.do_rot90 = False

		self.scaleFactor = [minScale, maxScale]
		if (self.scaleFactor[0]==1 and self.scaleFactor[1]==1):
			self.do_scaling = False
		else:
			self.do_scaling = True
			msg += 'scaling, '

		self.do_flip = flip
		if self.do_flip:
			msg += 'flip'
		msg += '.'
		print(msg)
		self.interpolation_order = 1
		self.fill_mode = 'constant'

	def addData(self, low, high):
		'''
			add data, tiles if premadeTiles, frames otherwise.
			low, high: list of or single 3D data np arrays
		'''
		# check data shape
		low = np.asarray(low)
		high = np.asarray(high)

		if not self.data_flags[DATA_KEY_HIGH]['isLabel']:
			if len(low.shape)!=len(high.shape): #high-low mismatch
				self.TCError('Data shape mismatch. Dimensions: {} low vs {} high. Dimensions must match or use highIsLabel.'.format(len(low.shape),len(high.shape)) )
		if not (len(low.shape)==4 or len(low.shape)==5): #not single frame or sequence of frames
			self.TCError('Input must be single 3D data or sequence of 3D data. Format: ([batch,] z, y, x, channels). For 2D use z=1.')

		if (low.shape[-1]!=(self.dim_t * self.data_flags[DATA_KEY_LOW]['channels'])):
			self.TCError('Dim_t ({}) * Channels ({}, {}) configured for LOW-res data don\'t match channels ({}) of input data.'.format(self.dim_t, self.data_flags[DATA_KEY_LOW]['channels'], self.c_low,  low.shape[-1]) )
		if not self.data_flags[DATA_KEY_HIGH]['isLabel']:
			if (high.shape[-1]!=(self.dim_t * self.data_flags[DATA_KEY_HIGH]['channels'])):
				self.TCError('Dim_t ({}) * Channels ({}, {}) configured for HIGH-res data don\'t match channels ({}) of input data.'.format(self.dim_t, self.data_flags[DATA_KEY_HIGH]['channels'], self.c_high, high.shape[-1]) )

		low_shape = low.shape
		high_shape = high.shape
		if len(low.shape)==5: #sequence
			if low.shape[0]!=high.shape[0]: #check amount
				self.TCError('Unequal amount of low ({}) and high ({}) data.'.format(low.shape[1], high.shape[1]))
			# get single data shape
			low_shape = low_shape[1:]
			if not self.data_flags[DATA_KEY_HIGH]['isLabel']:
				high_shape = high_shape[1:]
			else: high_shape = [1]
		else: #single
			low = [low]
			high = [high]

		if self.premadeTiles:
			if not (self.dim_t == 1):
				self.TCError('Currently, Dim_t = {} > 1 is not supported by premade tiles'.format(self.dim_t))
			if not np.array_equal(low_shape, self.tile_shape_low) or not np.array_equal(high_shape,self.tile_shape_high):
				self.TCError('Tile shape mismatch: is - specified\n\tlow: {} - {}\n\thigh {} - {}'.format(low_shape, self.tile_shape_low, high_shape,self.tile_shape_high))
		else:
			single_frame_low_shape = list(low_shape)
			single_frame_high_shape = list(high_shape)
			single_frame_low_shape[-1] = low_shape[-1] // self.dim_t
			if not self.data_flags[DATA_KEY_HIGH]['isLabel']:
				single_frame_high_shape[-1] = high_shape[-1] // self.dim_t

			if not np.array_equal(single_frame_low_shape, self.frame_shape_low) or not np.array_equal(single_frame_high_shape,self.frame_shape_high):
				self.TCError('Frame shape mismatch: is - specified\n\tlow: {} - {}\n\thigh: {} - {}, given dim_t as {}'.format(single_frame_low_shape, self.frame_shape_low, single_frame_high_shape,self.frame_shape_high, self.dim_t))

		self.data[DATA_KEY_LOW].extend(low)
		self.data[DATA_KEY_HIGH].extend(high)

		print('\n')
		print('Added {} datasets. Total: {}'.format(low.shape[0], len(self.data[DATA_KEY_LOW])))
		self.splitSets()

	def splitSets(self):
		'''
			calculate the set borders for training, testing and validation set
		'''
		length = len(self.data[DATA_KEY_LOW])

		end_train = int( length * self.part_train )
		end_test = end_train + int( length * self.part_test )
		#just store the borders of the different sets to avoid data duplication
		self.setBorders = [end_train, end_test, length]

		print('Training set: {}'.format(self.setBorders[0]))
		print('Testing set:  {}'.format(self.setBorders[1]-self.setBorders[0]))
		print('Validation set:  {}'.format(self.setBorders[2]-self.setBorders[1]))

	def clearData(self):
		'''
			clears the data buffer
		'''
		self.data = {
			DATA_KEY_LOW:[],
			DATA_KEY_HIGH:[]
		}

	def createTiles(self, data, tileShape, strides=-1):
		'''
			create tiles from a single frame. fixed, regular pattern
			strides: <=0 or tileShape is normal, otherwise create overlapping tiles
		'''
		dataShape = data.shape #2D sim: [1,res,res,channels]
		pad = [self.padding,self.padding,self.padding,0]
		if np.isscalar(strides):
			if strides <= 0:
				strides = tileShape
			else:
				strides = [strides,strides,strides]
		if dataShape[0]<=1:
			pad[0] = 0
			strides[0] = 1
		channels = dataShape[3]
		noTiles = [ (dataShape[0]-tileShape[0])//strides[0]+1, (dataShape[1]-tileShape[1])//strides[1]+1, (dataShape[2]-tileShape[2])//strides[2]+1 ]
		tiles = []

		for tileZ in range(0, noTiles[0]):
			for tileY in range(0, noTiles[1]):
				for tileX in range(0, noTiles[2]):
					idx_from=[tileZ*strides[0], tileY*strides[1], tileX*strides[2]]
					idx_to=[idx_from[0]+tileShape[0], idx_from[1]+tileShape[1], idx_from[2]+tileShape[2]]
					currTile=data[ idx_from[0]:idx_to[0], idx_from[1]:idx_to[1], idx_from[2]:idx_to[2], :]
					if self.padding > 0:
						currTile = np.pad(currTile, pad, 'edge')

					tiles.append(currTile)
		return np.array(tiles)

	def cutTile(self, data, tileShape, offset=[0,0,0]):
		'''
			cut a tile of with shape and offset
		'''
		offset = np.asarray(offset)
		tileShape = np.asarray(tileShape)
		tileShape[-1] = data.shape[-1]
		if np.less(data.shape[:3], tileShape[:3]+offset[:3]).any():
			self.TCError('Can\'t cut tile with shape {} and offset{} from data with shape {}.'.format(tileShape, offset, data.shape))

		tile = data[offset[0]:offset[0]+tileShape[0], offset[1]:offset[1]+tileShape[1], offset[2]:offset[2]+tileShape[2], :]

		if not np.array_equal(tile.shape,tileShape):
			self.TCError('Wrong tile shape after cutting. is: {}. goal: {}.'.format(tile.shape,tileShape))
		return tile

#####################################################################################
# batch creation
#####################################################################################

	def selectRandomTiles(self, selectionSize, isTraining=True, augment=False, tile_t = 1):
		'''
			main method to create baches
			Return:
				shape: [selectionSize, z, y, x, channels * tile_t]
				if 2D z = 1
				channels: density, [vel x, vel y, vel z], [pos x, pox y, pos z]
		'''
		if isTraining:
			if self.setBorders[0]<1:
				self.TCError('no training data.')
		else:
			if (self.setBorders[1] - self.setBorders[0])<1:
				self.TCError('no test data.')
		if(tile_t > self.dim_t):
			self.TCError('not enough coherent frames. Requested {}, available {}'.format(tile_t, self.dim_t))
		batch_low = []
		batch_high = []
		for i in range(selectionSize):
			if augment and self.useDataAug: #data augmentation
				low, high = self.generateTile(isTraining, tile_t)
			else: #cut random tile without augmentation
				low, high = self.getRandomDatum(isTraining, tile_t)
				if not self.premadeTiles:
					low, high = self.getRandomTile(low, high)
			batch_low.append(low)
			batch_high.append(high)

		return np.asarray(batch_low), np.asarray(batch_high)

	def generateTile(self, isTraining=True, tile_t = 1):
		'''
			generates a random low-high pair of tiles (data augmentation)
		'''
		# get a frame, is a copy to avoid transormations affecting the raw dataset
		data = {}
		data[DATA_KEY_LOW], data[DATA_KEY_HIGH] = self.getRandomDatum(isTraining, tile_t)

		if not self.premadeTiles:
			#cut a tile for faster transformation
			if self.do_scaling or self.do_rotation:
				factor = 1
				if self.do_rotation: # or self.do_scaling:
					factor*=1.5 # scaling: to avoid size errors caused by rounding
				if self.do_scaling:
					scaleFactor = np.random.uniform(self.scaleFactor[0], self.scaleFactor[1])
					factor/= scaleFactor
				tileShapeLow = np.ceil(self.tile_shape_low*factor)
				if self.dim==2:
					tileShapeLow[0] = 1
				data[DATA_KEY_LOW], data[DATA_KEY_HIGH] = self.getRandomTile(data[DATA_KEY_LOW], data[DATA_KEY_HIGH], tileShapeLow.astype(int))

			#random scaling, changes resolution
			if self.do_scaling:
				data = self.scale(data, scaleFactor)

			bounds = np.zeros(4)

			#rotate
			if self.do_rotation:
				bounds = np.array(data[DATA_KEY_LOW].shape)*0.16 #bounds applied on all sides, 1.5*(1-2*0.16)~1
				data = self.rotate(data)

			#get a tile
			data[DATA_KEY_LOW], data[DATA_KEY_HIGH] = self.getRandomTile(data[DATA_KEY_LOW], data[DATA_KEY_HIGH], bounds=bounds) #includes "shifting"

		if self.do_rot90:
			rot = np.random.choice(self.cube_rot[self.dim])
			for axis in rot:
				data = self.rotate90(data, axis)

		#flip once
		if self.do_flip:
			axis = np.random.choice(4)
			if axis < 3: # axis < self.dim
				data = self.flip(data, [axis])

		# check tile size
		target_shape_low = np.copy(self.tile_shape_low)
		target_shape_high = np.copy(self.tile_shape_high)
		target_shape_low[-1] *= tile_t
		target_shape_high[-1] *= tile_t

		if not np.array_equal(data[DATA_KEY_LOW].shape,target_shape_low) or (not np.array_equal(data[DATA_KEY_HIGH].shape,target_shape_high) and not self.data_flags[DATA_KEY_HIGH]['isLabel']):
			self.TCError('Wrong tile shape after data augmentation. is: {},{}. goal: {},{}.'.format(data[DATA_KEY_LOW].shape, data[DATA_KEY_HIGH].shape, target_shape_low, target_shape_high))

		return data[DATA_KEY_LOW], data[DATA_KEY_HIGH]

	def getRandomDatum(self, isTraining=True, tile_t = 1):
		'''returns a copy of a random frame'''
		if isTraining:
			randNo = randrange(0, self.setBorders[0])
		else:
			randNo = randrange(self.setBorders[0], self.setBorders[1])
		randFrame = 0
		if tile_t<self.dim_t:
			randFrame = randrange(0, self.dim_t - tile_t)
		else:
			tile_t = self.dim_t

		return self.getDatum(randNo*self.dim_t+randFrame, tile_t)

	def getDatum(self, index, tile_t = 1):
		'''returns a copy of the indicated frame or tile'''
		begin_ch = 0
		if(self.dim_t > 1):
			begin_ch = (index % self.dim_t) * self.tile_shape_low[-1]
		end_ch = begin_ch + tile_t * self.tile_shape_low[-1]
		begin_ch_y = 0
		if(self.dim_t > 1):
			begin_ch_y = (index % self.dim_t) * self.tile_shape_high[-1]
		end_c_h_y = begin_ch_y + tile_t * self.tile_shape_high[-1]

		if not self.data_flags[DATA_KEY_HIGH]['isLabel']:
			return np.copy(self.data[DATA_KEY_LOW][index//self.dim_t][:,:,:,begin_ch:end_ch]), np.copy(self.data[DATA_KEY_HIGH][index//self.dim_t][:,:,:,begin_ch_y:end_c_h_y])
		else:
			return np.copy(self.data[DATA_KEY_LOW][index//self.dim_t][:,:,:,begin_ch:end_ch]), np.copy(self.data[DATA_KEY_HIGH][index//self.dim_t])

	def getRandomTile(self, low, high, tileShapeLow=None, bounds=[0,0,0,0]): #bounds to avoid mirrored parts
		'''
			cut a random tile (low and high) from a given frame, considers densityMinimum
			bounds: ignore edges of frames, used to discard mirrored parts after rotation
		'''

		if tileShapeLow is None:
			tileShapeLow = np.copy(self.tile_shape_low) # use copy is very important!!!
		tileShapeHigh = tileShapeLow*self.upres

		frameShapeLow = np.asarray(low.shape)
		if len(low.shape)!=4 or len(tileShapeLow)!=4:
			self.TCError('Data shape mismatch.')
		if  len(high.shape)!=4 and not self.data_flags[DATA_KEY_HIGH]['isLabel']:
			self.TCError('Data shape mismatch.')

		start = np.ceil(bounds)
		end = frameShapeLow - tileShapeLow + np.ones(4) - start

		offset_up = np.array([self.upres, self.upres, self.upres])

		if self.dim==2:
			start[0] = 0
			end[0] = 1
			offset_up[0] = 1
			tileShapeHigh[0] = 1

		# check if possible to cut tile
		if np.amin((end-start)[:3]) < 0:
			self.TCError('Can\'t cut tile {} from frame {} with bounds {}.'.format(tileShapeLow, frameShapeLow, start))

		# cut tile
		hasMinDensity = False
		i = 1
		while (not hasMinDensity) and i<20:
			offset = np.asarray([randrange(start[0], end[0]), randrange(start[1], end[1]), randrange(start[2], end[2])])
			lowTile = self.cutTile(low, tileShapeLow, offset)
			offset *= offset_up
			if not self.data_flags[DATA_KEY_HIGH]['isLabel']:
				highTile = self.cutTile(high, tileShapeHigh, offset)
			else:
				highTile = high
			hasMinDensity = self.hasMinDensity(lowTile)
			i+=1
		return lowTile, highTile

#####################################################################################
# AUGMENTATION
#####################################################################################

	def special_aug(self, data, ops_key, param):
		"""
			wrapper to call the augmentation operations specified in self.aops in initAugmentation
		"""
		for data_key in data:
			if self.data_flags[data_key]['isLabel']: continue
			orig_shape = data[data_key].shape
			tile_t = orig_shape[-1] // self.data_flags[data_key]['channels']
			data_array = data[data_key]
			if(tile_t > 1): data_array = data[data_key].reshape( (-1, tile_t, self.data_flags[data_key]['channels']) )
			for c_key, op in self.aops[data_key][ops_key].items():
				if self.data_flags[data_key][c_key]:
					data_array = op(data_array, self.c_lists[data_key][c_key], param)
			if (tile_t > 1): data[data_key] = data_array.reshape(orig_shape)
		return data

	def rotate(self, data):
		'''
			random uniform rotation of low and high data of a given frame
		'''
		#check if single frame

		#2D:
		if self.dim==2:
			theta = np.pi * np.random.uniform(0, 2)
			rotation_matrix = np.array([[1, 0, 0, 0 ],
										[0, np.cos(theta), -np.sin(theta), 0],
										[0, np.sin(theta), np.cos(theta) , 0],
										[0, 0, 0, 1]  ])

		#3D:
		elif self.dim==3:
			# random uniform rotation in 3D
			quat = np.random.normal(size=4)
			quat/= np.linalg.norm(quat)

			q = np.outer(quat, quat)*2
			rotation_matrix = np.array([[1-q[2, 2]-q[3, 3],   q[1, 2]-q[3, 0],   q[1, 3]+q[2, 0], 0],
										[  q[1, 2]+q[3, 0], 1-q[1, 1]-q[3, 3],   q[2, 3]-q[1, 0], 0],
										[  q[1, 3]-q[2, 0],   q[2, 3]+q[1, 0], 1-q[1, 1]-q[2, 2], 0],
										[                0,                 0,                 0, 1]])

		data = self.special_aug(data, AOPS_KEY_ROTATE, rotation_matrix)

		for data_key in data:
			if not self.data_flags[data_key]['isLabel']:
				data[data_key] = self.applyTransform(data[data_key], rotation_matrix.T)

		return data

	def rotate_simple(self, low, high, angle):
		'''
			use a different method for rotation. about 30-40% faster than with rotation matrix, but only one axis.
		'''
		if len(low.shape)!=4 or len(high.shape)!=4:
			self.TCError('Data shape mismatch.')
		#test rot around z (axis order z,y,x,c)
		low = scipy.ndimage.rotate(low, angle, [1,2] , reshape=False, order=self.interpolation_order, mode=self.fill_mode, cval=1.0)
		high = scipy.ndimage.rotate(high, angle, [1,2] , reshape=False, order=self.interpolation_order, mode=self.fill_mode, cval=1.0)
		return low, high

	def rotateVelocities(self, datum, c_list, rotationMatrix):
		'''
			rotate vel vectors (channel 1-3)
		'''

		rotation3 = rotationMatrix[:3, :3]
		rotation2 = rotationMatrix[1:3, 1:3]
		channels = np.split(datum, datum.shape[-1], -1)
		for v in c_list:
			if len(v) == 3: # currently always ends here!! even for 2D, #z,y,x to match rotation matrix
				vel = np.stack([channels[v[2]].flatten(),channels[v[1]].flatten(),channels[v[0]].flatten()])
				vel = rotation3.dot(vel)
				channels[v[2]] = np.reshape(vel[0], channels[v[2]].shape)
				channels[v[1]] = np.reshape(vel[1], channels[v[1]].shape)
				channels[v[0]] = np.reshape(vel[2], channels[v[0]].shape)
			if len(v) == 2:
				vel = np.concatenate([channels[v[1]],channels[v[0]]], -1) #y,x to match rotation matrix
				shape = vel.shape
				vel = np.reshape(vel, (-1, 2))
				vel = np.reshape(rotation2.dot(vel.T).T, shape)
				vel = np.split(vel, 2, -1)
				channels[v[1]] = vel[0]
				channels[v[0]] = vel[1]

		return np.concatenate(channels, -1)

	def rotate90(self, data, axes):
		'''
			rotate the frame by 90 degrees from the first axis counterclockwise to the second
			axes: 2 int, from axis to axis; see np.rot90
				0,1,2 -> z,y,x
		'''
		if len(axes)!=2:
			self.TCError('need 2 axes for rotate90.')

		for data_key in data:
			if not self.data_flags[data_key]['isLabel']:
				data[data_key] = np.rot90(data[data_key], axes=axes)

		data = self.special_aug(data, AOPS_KEY_ROT90, axes)

		return data

	def rotate90Velocities(self, datum, c_list, axes):
		if len(axes)!=2:
			self.TCError('need 2 axes for rotate90.')

		channels = np.split(datum, datum.shape[-1], -1)
		for v in c_list: #axes z,y,x -> vel x,y,z: 0,1,2 -> 2,1,0
			channels[v[-axes[0]+2]], channels[v[-axes[1]+2]] = -channels[v[-axes[1]+2]], channels[v[-axes[0]+2]]

		return np.concatenate(channels, -1)

	def flip(self, data, axes, isFrame=True): #axes=list, flip multiple at once
		'''
			flip low and high data (single frame/tile) along the specified axes
			low, high: data format: (z,x,y,c)
			axes: list of axis indices 0,1,2-> z,y,x
		'''
		# axis: 0,1,2 -> z,y,x

		if not isFrame:
			axes = np.asarray(axes) + np.ones(axes.shape)

		#flip tiles/frames
		for axis in axes:
			for data_key in data:
				if not self.data_flags[data_key]['isLabel']:
					data[data_key] = np.flip(data[data_key], axis)

		data = self.special_aug(data, AOPS_KEY_FLIP, axes)

		return data

	def flipVelocities(self, datum, c_list, axes):
		'''
			flip velocity vectors along the specified axes
			low: data with velocity to flip (4 channels: d,vx,vy,vz)
			axes: list of axis indices 0,1,2-> z,y,x
		'''

		# !axis order: data z,y,x
		channels = np.split(datum, datum.shape[-1], -1)
		for v in c_list: # x,y,[z], 2,1,0
			if 2 in axes: # flip vel x
				channels[v[0]] *= (-1)
			if 1 in axes:
				channels[v[1]] *= (-1)
			if 0 in axes and len(v)==3:
				channels[v[2]] *= (-1)

		return np.concatenate(channels, -1)

	def scale(self, data, factor):
		'''
			 changes frame resolution to "round((factor) * (original resolution))"
		'''
		# only same factor in every dim for now. how would it affect vel scaling?
		# check for 2D

		scale = [factor, factor, factor, 1] #single frame
		if self.dim==2:
			scale[0] = 1

		# to ensure high/low ration stays the same
		scale = np.round(np.array(data[DATA_KEY_LOW].shape) * scale )/np.array(data[DATA_KEY_LOW].shape)
		if len(data[DATA_KEY_LOW].shape)==5: #frame sequence
			scale = np.append([1],scale)

		#apply transform
		#low = self.applyTransform(low, zoom_matrix)
		#high = self.applyTransform(high, zoom_matrix)

		#changes the size of the frame. should work well with getRandomTile(), no bounds needed
		for data_key in data:
			if not self.data_flags[data_key]['isLabel']:
				data[data_key]  = scipy.ndimage.zoom( data[data_key], scale, order=self.interpolation_order, mode=self.fill_mode, cval=0.0)

		#necessary?
		data = self.special_aug(data, AOPS_KEY_SCALE, factor)

		return data

	def scaleVelocities(self, datum, c_list, factor):
		#scale vel? vel*=factor
		channels = np.split(datum, datum.shape[-1], -1)
		for v in c_list: # x,y,[z]; 2,1,0
			channels[v[0]] *= factor
			channels[v[1]] *= factor
			if len(v)==3:
				channels[v[2]] *= factor

		return np.concatenate(channels, -1)

	def applyTransform(self, data, transform_matrix, data_dim=3):
		# change axis order from z,y,x to x,y,z? (invert axis order +channel)
		if len(data.shape)!=4:
			self.TCError('Data shape mismatch.')

		#set transform to center; from fluiddatagenerator.py
		offset = np.array(data.shape) / 2 - np.array([0.5, 0.5, 0.5, 0])
		offset_matrix = np.array([[1, 0, 0, offset[0]], [0, 1, 0, offset[1]], [0, 0, 1, offset[2]], [0, 0, 0, 1]])
		reset_matrix  = np.array([[1, 0, 0,-offset[0]], [0, 1, 0,-offset[1]], [0, 0, 1,-offset[2]], [0, 0, 0, 1]])
		transform_matrix = np.dot(np.dot(offset_matrix, transform_matrix), reset_matrix)

		data = np.rollaxis(data, 3, 0) #channel to front
		channel_data = [scipy.ndimage.interpolation.affine_transform(
			channel,
			transform_matrix[:data_dim,:data_dim],
			transform_matrix[:data_dim, data_dim],
			order=self.interpolation_order,
			mode=self.fill_mode,
			cval=0.) for channel in data]
		data = np.stack(channel_data, axis=-1) # stack axis=-1 ?\
		return data

#####################################################################################
# HELPER METHODS
#####################################################################################

	def concatTiles(self, tiles, frameShape ,tileBorder=[0,0,0,0]):
		'''
			build a frame by concatenation of the given tiles.
			tiles: numpy array of same shaped tiles [batch,z,y,x,c]
			frameShape: the shape of the frame in tiles [z,y,x]
			tileBorder: cut off borders of the tiles. [z,y,x,c]
		'''
		if len(tiles.shape)!=5 or len(frameShape)!=3 or len(tileBorder)!=4:
			self.TCError('Data shape mismatch.')
		tiles_in_frame = frameShape[0]*frameShape[1]*frameShape[2]
		if tiles_in_frame != len(tiles):
			self.TCError('given tiles do not match required tiles.')

		# cut borders
		tileBorder = np.asarray(tileBorder)
		if np.less(np.zeros(4),tileBorder).any():
			tileShape = tiles.shape[1:] - 2*tileBorder
			tiles_cut = []
			for tile in tiles:
				tiles_cut.append(self.cutTile(tile, tileShape, tileBorder))
			tiles = tiles_cut

		#combine tiles to image
		frame = []
		for z in range(frameShape[0]):
			frame_slices = []
			for y in range(frameShape[1]):
				offset=z*frameShape[1]*frameShape[2] + y*frameShape[2]
				frame_slices.append(np.concatenate(tiles[offset:offset+frameShape[2]],axis=2)) #combine x
			frame.append(np.concatenate(frame_slices, axis=1)) #combine y
		frame = np.concatenate(frame, axis=0) #combine z

		return frame

	def hasMinDensity(self, tile):
		return self.getTileDensity(tile) >= (self.densityMinimum * tile.shape[0] * tile.shape[1] * tile.shape[2])

	def getTileDensity(self, tile):
		if self.data_flags[DATA_KEY_LOW]['channels'] > 1:
			tile = np.split(tile, [1], axis=-1)[0]
		return tile.sum( dtype=np.float64 )

	def getFrameTiles(self, index):
		''' returns the frame as tiles'''
		low, high = self.getDatum(index)
		return self.createTiles(low, self.tile_shape_low), self.createTiles(high, self.tile_shape_high)

#####################################################################################
# CHANNEL PARSING
#####################################################################################

	def parseChannels(self, channelString):
		''' arbitrary channel structure from string, expand if necessary. USE GLOBAL KEYS ^
			'd': default/ density; data that needs no special operations during augmentation
			'v[label](x|y|z)': vector/velocity; is transformed according to the augmentation
		'''
		#need this for low and high, +high only labels
		c = channelString.lower().split(',')
		for i in range(len(c)):
			c[i] = c[i].strip()

		c_types = {
			C_KEY_DEFAULT:[], #list of indices of default channels. e.g. for normal sim data [0]
			C_KEY_VELOCITY:[], #list of ordered (x,y,z; chnage to z,y,x to match data?) triples of velocity sets. e.g. for normal sim data [[1,2,3]]
			C_KEY_VORTICITY:[]
			}

		self.parse = {
			C_KEY_DEFAULT:self.parseCDefault,
			C_KEY_VELOCITY:self.parseCVelocity,
			C_KEY_VORTICITY:self.parseCVorticity
			}

		for i in range(len(c)):
			if len(c[i])==0: # check empty key
				self.TCError('empty channel key.'.format(i))
			try:
				self.parse[c[i][0]](c, i, c_types)
			except KeyError:
				self.TCError('channel {}: unknown channel key \"{}\".'.format(i, c[i]))

		# TODO check unused channels here
		return c, c_types

	def parseCDefault(self, c, i, c_types):
		# check length
		if c[i]=='d':
			c_types[C_KEY_DEFAULT].append(i)
		else:
			self.TCError('channel {}: unknown channel key \"{}\".'.format(i, c[i]))

	def parseCVector(self, c, i, c_types, c_key, c_name='vector'):
		# c_key[label](x|y|z)
		if c[i][-1] == 'x' or c[i][-1] == 'y' or c[i][-1] == 'z':
			label = c[i][1:-1] #can be empty

			#get matching keys
			v_x = c_key+label+'x'
			v_y = c_key+label+'y'
			v_z = c_key+label+'z'

			#check for duplicates
			if c.count(v_x)>1:
				self.TCError('Duplicate {} ({}) x-channel with label \"{}\": {}. Vector keys must be unique.'.format(c_name, c_key, label, v_x))
			if c.count(v_y)>1:
				self.TCError('Duplicate {} ({}) y-channel with label \"{}\": {}. Vector keys must be unique.'.format(c_name, c_key, label, v_y))
			if c.count(v_z)>1:
				self.TCError('Duplicate {} ({}) z-channel with label \"{}\": {}. Vector keys must be unique.'.format(c_name, c_key, label, v_z))

			#check missing
			if c.count(v_x)==0:
				self.TCError('Missing {} ({}) x-channel with label \"{}\": {}'.format(c_name, c_key, label, v_x))
			if c.count(v_y)==0:
				self.TCError('Missing {} ({}) y-channel with label \"{}\": {}'.format(c_name, c_key, label, v_y))
			if self.dim==3 and c.count(v_z)==0:
				self.TCError('Missing {} ({}) z-channel with label \"{}\": {}'.format(c_name, c_key, label, v_z))

			if c[i][-1] == 'x':
				if(c.count(v_z)==0 and self.dim==2):
					c_types[C_KEY_VELOCITY].append([c.index(v_x),c.index(v_y)])
				else:
					c_types[C_KEY_VELOCITY].append([c.index(v_x),c.index(v_y),c.index(v_z)])

		# check wrong suffix
		else:
			self.TCError('Channel {}, \"{}\": unknown {} ({}) channel suffix \"{}\". Valid suffixes are \"x\", \"y\", \"z\".'.format(i, c[i], c_name, c_key, c[i][-1]))

	def parseCVelocity(self, c, i, c_types):
		# C_KEY_VELOCITY[label](x|y|z)
		self.parseCVector(c, i, c_types, C_KEY_VELOCITY, 'velociy')

	def parseCVorticity(self, c, i, c_types):
		# C_KEY_VELOCITY[label](x|y|z)
		self.parseCVector(c, i, c_types, C_KEY_VELOCITY, 'vorticity')

#####################################################################################
# ERROR HANDLING
#####################################################################################

	def TCError(self, msg):
		raise TilecreatorError(msg)

class TilecreatorError(Exception):
	''' Tilecreator errors '''

#####################################################################################
# IMAGE OUTPUT
#####################################################################################

# save summary images of all channels in a batch generated by the tile creator
# projects 3D data onto different axes, data has to be B-ZYX-C format
batchCounterGlob=0
def savePngsBatch(low,high, TC, path, batchCounter=-1, save_vels=False, dscale=1., vscale=1.):
	global batchCounterGlob
	if(low.shape[1]==1):
		dim=2
	else:
		dim=3

	# figure out good tile size, and project all axes for 3D
	batch = low.shape[0]
	tileX = 4 if batch>=4 else batch
	if batch%tileX != 0: tileX = batch
	# file names
	if batchCounter < 0:
		batchCounter = batchCounterGlob
		batchCounterGlob += 1
	path = path+"batch{:04d}_".format(batchCounter)

	# show scalar channels, in tiled images
	aNames = ["xy_","xz_","yz_"]
	for axis in range(1 if dim==2 else 3):
		suff = aNames[axis]

		if dim == 3:
			highD = np.average(high, axis=axis+1) * dscale
			lowD  = np.average(low, axis=axis+1) * dscale
		if dim == 2:
			highD = high*brightness ; lowD = low * dscale
			lowD.shape = (batch, tll, tll, cl)
			highD.shape = (batch, tlh, tlh, ch)

		# note - outputs all channels as images, also vel channels...
		clout = np.arange(low.shape[4])
		savePngsGrayscale(lowD, path+'low_'+suff, tiles_in_image=[batch//tileX,tileX], channels=clout )
		chout = np.arange(high.shape[4])
		savePngsGrayscale(tiles=highD, path=path+'high_'+suff, imageCounter=0, tiles_in_image=[batch//tileX,tileX], channels=chout )

	# plot velocities , for individual samples
	if save_vels:
		for i in range(low.shape[0]):
			saveVelChannels(low[i], TC.c_lists[DATA_KEY_LOW][C_KEY_VELOCITY], path=path+'low_vel_i{:02d}_'.format(i), name="", scale=vscale )
		for i in range(high.shape[0]):
			saveVelChannels(high[i], TC.c_lists[DATA_KEY_HIGH][C_KEY_VELOCITY], path=path+'high_vel_i{:02d}_'.format(i), name="", scale=vscale )

# simpler function to output multiple tiles into grayscale pngs
def savePngsGrayscale(tiles, path, imageCounter=0, tiles_in_image=[1,1], channels=[0], save_gif=False, plot_vel_x_y=False, save_rgb=None, rgb_interval=[-1,1]):
	'''
		tiles_in_image: (y,x)
		tiles: shape: (tile,y,x,c)
	'''
	tilesInImage = tiles_in_image[0]*tiles_in_image[1]
	if len(tiles)%tilesInImage!=0:
		print('ERROR: number of tiles does not match tiles per image')
		return
	tiles = np.asarray(tiles)
	noImages = len(tiles)//tilesInImage
	if save_gif:
		gif=[]

	for image in range(noImages):
		img = []
		#combine tiles to image
		for y in range(tiles_in_image[0]):
			offset=image*tilesInImage + y*tiles_in_image[1]
			img.append(np.concatenate(tiles[offset:offset+tiles_in_image[1]],axis=1)) #combine x
		img = np.concatenate(img, axis=0) #combine y
		# move channels to first dim.
		img_c = np.rollaxis(img, -1, 0)
		if len(img_c)>1 and (plot_vel_x_y or save_rgb!=None):
			if plot_vel_x_y: saveVel(img, path, imageCounter+image)
			if save_rgb!=None: saveRGBChannels(img,path, save_rgb,value_interval=rgb_interval, imageCounter=imageCounter+image)
		if len(channels) == 1:
			scipy.misc.toimage(img_c[channels[0]], cmin=0.0, cmax=1.0).save(path + 'img_{:04d}.png'.format(imageCounter*noImages+image))
		else:
			for i in channels:
				scipy.misc.toimage(img_c[i], cmin=0.0, cmax=1.0).save(path + 'img_{:04d}_c{:04d}.png'.format(imageCounter*noImages+image, i))

# store velocity as quiver plot
def saveVel(tile, path, imageCounter=0, name='vel-x-y'):
	# origin is in upper left corner, transform acordingly
	y, x = np.mgrid[-tile.shape[0]:0, 0:tile.shape[1]]
	vx = None; vy = None
	if tile.shape[-1]==4:
		d, vx, vy, vz = np.split(tile, 4, -1)
	elif tile.shape[-1]==2:
		vx, vy  = np.split(tile, 2, -1)
	else:
		print('ERROR: unknown nr of channels for vel input '+format(tile.shape))
	vx = vx[::-1, ...]
	vy = vy[::-1, ...]
	if found_matplotlib:
		matplotlib.pyplot.quiver(x,y,vx.flatten(),vy.flatten(), units = 'xy', scale = 1)
		matplotlib.pyplot.axis('equal')
		matplotlib.pyplot.savefig(path + '{}_{:04d}.png'.format(name,imageCounter))
		matplotlib.pyplot.clf()

# save velocity channels from the tilecreator with multiple axis projections (uses saveVel)
def saveVelChannels(data, c_idx, path, average=False, scale=1.0, normalize=True, name=''):
	channels = np.split(data, data.shape[-1], -1)

	vpath = path
	vcnt = 0
	for v in c_idx:
		if(len(c_idx)>1):
			vpath = path + "vc{}".format(vcnt)
		vcnt += 1
		# compute scale factor
		vscale = scale
		if normalize:
			vavg = np.concatenate( [ channels[v[0]],channels[v[1]] ], -1)
			if(len(v))>2: vavg = np.concatenate( [ vavg, channels[v[2]] ],-1)
			vscale *= (1./(np.max( vavg )+1e-10)) # normalize

		vavg = np.concatenate( [ channels[v[0]],channels[v[1]] ] , -1)
		vavg = np.average(vavg, axis=0)
		vavg *= vscale
		saveVel(vavg, path=vpath, name='_xy' )

		if(len(v))>2: # also output xz,yz
			vavg = np.concatenate( [ channels[v[0]],channels[v[2]] ] , -1)
			vavg = np.average(vavg, axis=1)
			vavg *= vscale
			saveVel(vavg, path=vpath, name='_xz' )

			vavg = np.concatenate( [ channels[v[1]],channels[v[2]] ] , -1)
			vavg = np.average(vavg, axis=2)
			vavg *= vscale
			saveVel(vavg, path=vpath, name='_yz' )

def saveRGBChannels(data, path, channel_list, imageCounter=0, value_interval=[-1,1]):
	"""
		data: shape[y,x,c]
		channels: list of triples of channel ids saved as RGB image
	"""

	cmin = value_interval[0]
	cmax = value_interval[1]
	num_channels = data.shape[-1]

	channels = np.split(data, num_channels, -1)
	for i in channel_list:
		if len(i)==2:
			img = np.concatenate([channels[i[0]], channels[i[1]], np.ones_like(channels[i[0]])*cmin], -1)
		else:
			img = np.concatenate([channels[i[0]], channels[i[1]], channels[i[2]]], -1)
		scipy.misc.toimage(img, cmin=-1.0, cmax=1.0).save(path + 'img_rgb_{:04d}.png'.format(imageCounter))

def save3DasUni(tiles, path, motherUniPath, imageCounter=0, tiles_in_image=[1,1]):
	'''
		tiles_in_image: (y,x)
		tiles: shape: (image,y,x,c)
	'''

	tilesInImage = tiles_in_image[0]*tiles_in_image[1]
	if len(tiles)%tilesInImage!=0:
		print('ERROR: number of tiles does not match tiles per image')
		return
	tiles = np.asarray(tiles)
	noImages = len(tiles)//tilesInImage

	tiles = np.reshape(tiles,(len(tiles), tiles.shape[1], tiles.shape[2], tiles.shape[-1]))

	#only save the average
	img_all = []
	for image in range(noImages):
		img = []
		#combine tiles to image
		for y in range(tiles_in_image[0]):
			offset=( image) * tilesInImage + (y)*tiles_in_image[1]
			img.append(np.concatenate(tiles[offset:offset+tiles_in_image[1]],axis=2)) #combine y
		img = np.array(img)
		img = np.concatenate(img, axis=1) #combine x
		img = np.array(img)
		# move channels to first dim.
		img_c = np.rollaxis(img, 0, 0)
		img_all.append(img_c)
	img_all = np.array(img_all)
	img_all = np.concatenate(img_all, axis=0)
	img_all = np.array(img_all)
	TDarrayToUni(img_all, path + 'source_{:04d}.uni'.format(imageCounter), motherUniPath, img_all.shape[0], img_all.shape[1], img_all.shape[2])

def TDarrayToUni(input, savePath, motherUniPath, imageHeight, imageWidth, imageDepth, is_vel=False):
	head, _ = uniio.readUni(motherUniPath)
	head['dimX'] = imageWidth
	head['dimY'] = imageHeight
	head['dimZ'] = imageDepth

	if not is_vel:
		fixedArray = np.zeros((imageHeight, imageWidth, imageDepth), dtype='f')
		for x in range(0, imageHeight):
			for y in range(0, imageWidth):
				for z in range(0, imageDepth):
					fixedArray[x][y][z] = input[imageDepth - 1 - z][y][(imageHeight - 1) - x]
	else:
		fixedArray = np.zeros((imageHeight, imageWidth, imageDepth, 3), dtype='f')
		for x in range(0, imageHeight):
			for y in range(0, imageWidth):
				for z in range(0, imageDepth):
					fixedArray[x][y][z] = input[imageDepth - 1 - z][y][(imageHeight - 1) - x]

	uniio.writeUni(savePath, head, fixedArray)

# ******************************************************************************
# faster functions, batch operations
#

# grid interpolation method, order: only linear tested
# for velocity, macgridbatch.shape should be [b,z,y,x,3] (in 2D, should be [b,1,ny,nx,3])
# for density , macgridsource.shape should be [z,y,x,1] (in 2D, should be [1,ny,nx,1])
def gridInterpolBatch(macgridbatch, targetshape, order=1):
	assert (targetshape[-1] == macgridbatch.shape[-1])  # no interpolation between channels
	assert (len(targetshape) == 5 and len(macgridbatch.shape) == 5)
	dim = 3
	if (macgridbatch.shape[1] == 1 and targetshape[1] == 1):  dim = 2

	x_ = np.linspace(0.5, targetshape[3] - 0.5, targetshape[3])
	y_ = np.linspace(0.5, targetshape[2] - 0.5, targetshape[2])
	z_ = np.linspace(0.5, targetshape[1] - 0.5, targetshape[1])
	c_ = np.linspace(0, targetshape[4] - 1, targetshape[4])  # no interpolation between channels
	b_ = np.linspace(0, targetshape[0] - 1, targetshape[0])  # no interpolation between batches

	b, z, y, x, c = np.meshgrid(b_, z_, y_, x_, c_, indexing='ij')

	# scale
	fx = float(macgridbatch.shape[3]) / targetshape[3]
	fy = float(macgridbatch.shape[2]) / targetshape[2]
	fz = float(macgridbatch.shape[1]) / targetshape[1]

	mactargetbatch = scipy.ndimage.map_coordinates(macgridbatch, [b, z * fz, y * fy, x * fx, c], order=order, mode='reflect')

	return mactargetbatch;

# macgrid_batch shape b, z, y, x, 3
# return a matrix in size [b,z,y,x,3] ( 2D: [b,y,x,2]), last channel in z-y-x order!(y-x order for 2D)
def getMACGridCenteredBatch(macgrid_batch, is3D):
	_bn, _zn, _yn, _xn, _cn = macgrid_batch.shape

	valid_idx = list(range(1, _xn))
	valid_idx.append(_xn - 1)
	add_x = macgrid_batch.take(valid_idx, axis=3)[:, :, :, :, 0]  # shape, [b,z,y,x]
	valid_idx = list(range(1, _yn))
	valid_idx.append(_yn - 1)
	add_y = macgrid_batch.take(valid_idx, axis=2)[:, :, :, :, 1]  # shape, [b,z,y,x]

	add_y = add_y.reshape([_bn, _zn, _yn, _xn, 1])
	add_x = add_x.reshape([_bn, _zn, _yn, _xn, 1])
	if (is3D):
		valid_idx = list(range(1, _zn))
		valid_idx.append(_zn - 1)
		add_z = macgrid_batch.take(valid_idx, axis=1)[:, :, :, :, 2]  # shape, [b,z,y,x]
		add_z = add_z.reshape([_bn, _zn, _yn, _xn, 1])
		resultgrid = 0.5 * (macgrid_batch[:, :, :, :, ::-1] + np.concatenate((add_z, add_y, add_x), axis=-1))
		return resultgrid.reshape([_bn, _zn, _yn, _xn, 3])

	resultgrid = 0.5 * (macgrid_batch[:, :, :, :, -2::-1] + np.concatenate((add_y, add_x), axis=4))
	return resultgrid.reshape([_bn, _yn, _xn, 2])

# macgrid_batch shape b, z, y, x, 3 ( b,1,y,x,3 for 2D )
# return the re-sampling positions as a matrix, in size of [b,z,y,x,3] ( 2D: [b,y,x,2])
def getSemiLagrPosBatch(macgrid_batch, dt, cube_len_output=-1):  # check interpolation later
	assert (len(macgrid_batch.shape) == 5)
	_bn, _zn, _yn, _xn, _cn = macgrid_batch.shape
	assert (_cn == 3)
	is3D = (_zn > 1)
	if (cube_len_output == -1): cube_len_output = _xn
	factor = float(_xn) / cube_len_output
	x_ = np.linspace(0.5, int(_xn / factor + 0.5) - 0.5, int(_xn / factor + 0.5))
	y_ = np.linspace(0.5, int(_yn / factor + 0.5) - 0.5, int(_yn / factor + 0.5))
	interp_shape = [_bn, int(_zn / factor + 0.5), int(_yn / factor + 0.5), int(_xn / factor + 0.5), 3]
	if (not is3D): interp_shape[1] = 1

	if (is3D):
		z_ = np.linspace(0.5, int(_zn / factor + 0.5) - 0.5, int(_zn / factor + 0.5))
		z, y, x = np.meshgrid(z_, y_, x_, indexing='ij')
		posArray = np.stack((z, y, x), axis=-1)  # shape, z,y,x,3
		tarshape = [1, int(_zn / factor + 0.5), int(_yn / factor + 0.5), int(_xn / factor + 0.5), 3]
	else:
		y, x = np.meshgrid(y_, x_, indexing='ij')
		posArray = np.stack((y, x), axis=-1)  # shape, y,x,2
		tarshape = [1, int(_yn / factor + 0.5), int(_xn / factor + 0.5), 2]
	posArray = posArray.reshape(tarshape)
	if (cube_len_output == _xn):
		return (posArray - getMACGridCenteredBatch(macgrid_batch, is3D) * dt)
	# interpolate first

	inter_mac_batch = gridInterpolBatch(macgrid_batch, interp_shape, 1)
	inter_mac_batch = getMACGridCenteredBatch(inter_mac_batch, is3D) / factor
	return (posArray - (inter_mac_batch) * dt)

# error! muss not shuffle for fluid data loader!
def selectRandomTempoTiles(self, selectionSize, isTraining=True, augment=False, n_t=3, dt=0.5, adv_flag = 1.0):
	'''
		main method to create coherent baches
		Return:
			shape: [n_t * selectionSize//n_t, z, y, x, channels]
			if 2D z = 1
			channels: density, [vel x, vel y, vel z], [pos x, pox y, pos z]
	'''
	batch_sz = int( max( 1, selectionSize // n_t) )
	batch_low, batch_high = self.selectRandomTiles(batch_sz, isTraining, augment, tile_t = n_t)
	real_batch_sz = batch_sz * n_t
	ori_input_shape = batch_low.reshape((batch_sz, self.tileSizeLow[0], self.tileSizeLow[1], self.tileSizeLow[2], n_t, -1))
	ori_input_shape = np.transpose(ori_input_shape, (0,4,1,2,3,5))
	ori_input_shape = ori_input_shape.reshape((real_batch_sz, self.tileSizeLow[0], self.tileSizeLow[1], self.tileSizeLow[2], -1))
	vel_pos_high_inter = None
	if adv_flag:
		# TODO check velocity channels and 3D
		macgrid_input = ori_input_shape[:, :, :, :, self.c_lists[DATA_KEY_LOW][C_KEY_VELOCITY][0]]
		macgrid_input = macgrid_input.reshape( (real_batch_sz, self.tileSizeLow[0], self.tileSizeLow[1], self.tileSizeLow[2], 3))
		dtArray = np.array([i * dt for i in range(n_t // 2, -n_t // 2, -1)] * batch_sz, dtype=np.float32)
		if (self.dim == 2):
			dtArray = dtArray.reshape((-1, 1, 1, 1))
		else:
			dtArray = dtArray.reshape((-1, 1, 1, 1, 1))

		vel_pos_high_inter = getSemiLagrPosBatch(macgrid_input, dtArray, self.tileSizeHigh[1]).reshape((real_batch_sz, -1))

	# return reshape_input, selectedOutputs, vel_pos_high_inter
	batch_high = batch_high.reshape((batch_sz, self.tileSizeHigh[0], self.tileSizeHigh[1], self.tileSizeHigh[2], n_t, -1))
	batch_high = np.transpose(batch_high, (0,4,1,2,3,5))
	return ori_input_shape.reshape((real_batch_sz, -1)), batch_high.reshape((real_batch_sz, -1)), vel_pos_high_inter

TileCreator.selectRandomTempoTiles = selectRandomTempoTiles

def pngs_to_gif(path, start_idx=0, end_idx=199, step=1, fps=20, mask="img_%04d.png"):
	print("creating gif from {} to {} with {} fps".format(mask % start_idx, mask % end_idx, fps))
	with imageio.get_writer(path + 'step%d.gif'%step, mode='I', fps=fps) as writer:
		for i in range(start_idx - 1, end_idx, step):
			image = imageio.imread(path+(mask % (i+1)))
			writer.append_data(image)

ext: py

reason: import numpy,import scipy


text:
import os
import cv2
import numpy as np
from PIL  import Image

recognizer = cv2.face.LBPHFaceRecognizer_create()
path='dataSet'

def getImagesWithID(path):
    imagepaths=[os.path.join(path,f) for f in os.listdir(path)]
    faces=[]
    IDs=[]
    for imagepath in imagepaths:
        faceImg=Image.open(imagepath).convert('L');
        faceNp=np.array(faceImg,'uint8')
        ID=int(os.path.split(imagepath)[-1].split('.')[1])
        faces.append(faceNp)
        print (ID)
        IDs.append(ID)
        cv2.imshow("training",faceNp)
        cv2.waitKey(10)
    return np.array(IDs), faces

IDs,faces=getImagesWithID(path)
recognizer.train(faces,np.array(IDs))
recognizer.save('recognizer/TrainingData.yml')
cv2.destroyAllWindows()

reason: import numpy


text:
import os
from typing import List, Union
import h5py
import numpy as np

def merge_hdf5_files(
    new_filename: str,
    log_paths: List[str],
    file_ids: Union[None, List[str]] = None,
    delete_files: bool = False,
) -> None:
    """Merges a set of hdf5 files into a new hdf5 file with more groups."""
    file_to = h5py.File(new_filename, "w")
    for i, log_p in enumerate(log_paths):
        file_from = h5py.File(log_p, "r")
        datasets = get_datasets("/", file_from)
        if file_ids is None:
            write_data_to_file(file_to, file_from, datasets)
        else:
            # Maintain unique config id even if they have same random seed
            write_data_to_file(file_to, file_from, datasets, file_ids[i])
        file_from.close()

        # Delete individual log file if desired
        if delete_files:
            os.remove(log_p)
    file_to.close()

def get_datasets(key: str, archive: h5py.File):
    """Collects different paths to datasets in recursive fashion."""
    if key[-1] != "/":
        key += "/"
    out = []
    for name in archive[key]:
        path = key + name
        if isinstance(archive[path], h5py.Dataset):
            out += [path]
        else:
            out += get_datasets(path, archive)
    return out

def write_data_to_file(
    file_to: h5py.File,
    file_from: h5py.File,
    datasets: List[str],
    file_id: Union[str, None] = None,
):
    """Writes the datasets from-to file."""
    # get the group-names from the lists of datasets
    groups = list(set([i[::-1].split("/", 1)[1][::-1] for i in datasets]))
    if file_id is None:
        groups = [i for i in groups if len(i) > 0]
    else:
        groups = [i[0] + file_id + "/" + i[1:] for i in groups if len(i) > 0]

    # sort groups based on depth
    idx = np.argsort(np.array([len(i.split("/")) for i in groups]))
    groups = [groups[i] for i in idx]

    # create all groups that contain dataset that will be copied
    for group in groups:
        file_to.create_group(group)

    # copy datasets
    for path in datasets:
        # - get group name // - minimum group name // - copy data
        group = path[::-1].split("/", 1)[1][::-1]
        if len(group) == 0:
            group = "/"
        if file_id is not None:
            group_to_index = group[0] + file_id + "/" + group[1:]
        else:
            group_to_index = group
        file_from.copy(path, file_to[group_to_index])

    file_from.close()

max_stars_repo_name: rameshKrSah/mle-logging

max_issues_repo_name: rameshKrSah/mle-logging

max_forks_repo_name: rameshKrSah/mle-logging


text:
import sys
import numpy as np
from icatt._dopri import ffi, lib
import icatt.kepler as kepler
import icatt.elements as elements
import time
from numba import cfunc, types, farray

def benchmark(times):
    y = np.array([8.59072560e+02, -4.13720368e+03, 5.29556871e+03, 7.37289205e+00, 2.08223573e+00, 4.39999794e-01])
    y0 = y.copy()
    x = 0.0
    mu = 3.986004418e5
    rpar = np.array([mu])
    el = elements.elements(y[0:3], y[3:], mu)
    tp = kepler.period(el[0], mu)
    dopri(numba_gravity.cffi, x, y, tp, rpar)

    best = np.inf
    worst = -np.inf
    total = 0.0
    for _ in range(times):
        t0 = time.clock()
        dopri(numba_gravity.cffi, x, y, tp, rpar)
        t1 = time.clock()
        x = 0.0
        y = y0.copy()
        current = t1 - t0
        if current < best:
            best = current
        if current > worst:
            worst = current
        total += current
    print("[",total/times,",",best,",",worst,"]")

c_sig = types.void(
    types.CPointer(types.intc),
    types.CPointer(types.double),
    types.CPointer(types.double),
    types.CPointer(types.double),
    types.CPointer(types.double),
    types.CPointer(types.intc),
)

@cfunc(c_sig)
def numba_gravity(_n, _x, _y, _f, _rpar, _ipar):
    n = _n[0]
    x = _x[0]
    y = farray(_y, 6)
    f = farray(_f, 6)
    mu = _rpar[0]

    r = np.sqrt(y[0]*y[0]+y[1]*y[1]+y[2]*y[2])
    r3 = r*r*r
    f[0] = y[3]
    f[1] = y[4]
    f[2] = y[5]
    f[3] = -mu*y[0]/r3
    f[4] = -mu*y[1]/r3
    f[5] = -mu*y[2]/r3

@ffi.def_extern()
def solout(nr, xold, x, y, n, con, icomp, nd, rpar, ipar, irtrn, xout):
    pass

def dopri(func, x, y, xend, rpar, reltol=1e-6, abstol=1e-8):
    n = len(y)
    lwork = 11*n + 8*n + 21
    liwork = n + 21
    lwork_ = ffi.new("int *", lwork)
    liwork_ = ffi.new("int *", liwork)
    work = np.zeros(lwork)
    iwork = np.zeros(liwork, dtype=np.int32)
    rtol = np.array([reltol])
    atol = np.array([abstol])

    _n  = ffi.new("int *", n)
    _x = ffi.new('double *', x)
    _xend = ffi.new('double *', xend)
    _iout = ffi.new("int *", 0)
    _idid = ffi.new("int *", 0)
    _itol = ffi.new("int *", 0)
    _lwork = ffi.new("int *", lwork)
    _liwork = ffi.new("int *", liwork)
    _ipar = ffi.new("int []", [])
    _y = ffi.cast('double *', y.ctypes.data)
    _work = ffi.cast('double *', work.ctypes.data)
    _iwork = ffi.cast('int *', iwork.ctypes.data)
    _rtol = ffi.cast('double *', rtol.ctypes.data)
    _atol = ffi.cast('double *', atol.ctypes.data)
    _rpar = ffi.cast('double *', rpar.ctypes.data)
    lib.c_dop853(_n, func, _x, _y, _xend, _rtol, _atol, _itol, lib.solout, _iout, _work, _lwork, _iwork, _liwork, _rpar, _ipar, _idid)

max_issues_repo_name: OpenAstrodynamics/benchmarks

max_forks_repo_name: OpenAstrodynamics/benchmarks

reason: import numpy,from numba


text:
import numpy

import cupy
from cupy import cusparse
from cupy.cuda import cusolver
from cupy.cuda import device
from cupy.cuda import runtime
from cupy.linalg import _util
from cupyx.scipy import sparse

import warnings
try:
    import scipy.sparse
    import scipy.sparse.linalg
    scipy_available = True
except ImportError:
    scipy_available = False

def lsqr(A, b):
    """Solves linear system with QR decomposition.

    Find the solution to a large, sparse, linear system of equations.
    The function solves ``Ax = b``. Given two-dimensional matrix ``A`` is
    decomposed into ``Q * R``.

    Args:
        A (cupy.ndarray or cupyx.scipy.sparse.csr_matrix): The input matrix
            with dimension ``(N, N)``
        b (cupy.ndarray): Right-hand side vector.

    Returns:
        tuple:
            Its length must be ten. It has same type elements
            as SciPy. Only the first element, the solution vector ``x``, is
            available and other elements are expressed as ``None`` because
            the implementation of cuSOLVER is different from the one of SciPy.
            You can easily calculate the fourth element by ``norm(b - Ax)``
            and the ninth element by ``norm(x)``.

    .. seealso:: :func:`scipy.sparse.linalg.lsqr`
    """
    if runtime.is_hip:
        raise RuntimeError('HIP does not support lsqr')
    if not sparse.isspmatrix_csr(A):
        A = sparse.csr_matrix(A)
    _util._assert_nd_squareness(A)
    _util._assert_cupy_array(b)
    m = A.shape[0]
    if b.ndim != 1 or len(b) != m:
        raise ValueError('b must be 1-d array whose size is same as A')

    # Cast to float32 or float64
    if A.dtype == 'f' or A.dtype == 'd':
        dtype = A.dtype
    else:
        dtype = numpy.promote_types(A.dtype, 'f')

    handle = device.get_cusolver_sp_handle()
    nnz = A.nnz
    tol = 1.0
    reorder = 1
    x = cupy.empty(m, dtype=dtype)
    singularity = numpy.empty(1, numpy.int32)

    if dtype == 'f':
        csrlsvqr = cusolver.scsrlsvqr
    else:
        csrlsvqr = cusolver.dcsrlsvqr
    csrlsvqr(
        handle, m, nnz, A._descr.descriptor, A.data.data.ptr,
        A.indptr.data.ptr, A.indices.data.ptr, b.data.ptr, tol, reorder,
        x.data.ptr, singularity.ctypes.data)

    # The return type of SciPy is always float64. Therefore, x must be casted.
    x = x.astype(numpy.float64)
    ret = (x, None, None, None, None, None, None, None, None, None)
    return ret

def spsolve_triangular(A, b, lower=True, overwrite_A=False, overwrite_b=False,
                       unit_diagonal=False):
    """Solves a sparse triangular system ``A x = b``.

    Args:
        A (cupyx.scipy.sparse.spmatrix):
            Sparse matrix with dimension ``(M, M)``.
        b (cupy.ndarray):
            Dense vector or matrix with dimension ``(M)`` or ``(M, K)``.
        lower (bool):
            Whether ``A`` is a lower or upper trinagular matrix.
            If True, it is lower triangular, otherwise, upper triangular.
        overwrite_A (bool):
            (not supported)
        overwrite_b (bool):
            Allows overwriting data in ``b``.
        unit_diagonal (bool):
            If True, diagonal elements of ``A`` are assumed to be 1 and will
            not be referencec.

    Returns:
        cupy.ndarray:
            Solution to the system ``A x = b``. The shape is the same as ``b``.
    """
    if not cusparse.check_availability('csrsm2'):
        raise NotImplementedError

    if not sparse.isspmatrix(A):
        raise TypeError('A must be cupyx.scipy.sparse.spmatrix')
    if not isinstance(b, cupy.ndarray):
        raise TypeError('b must be cupy.ndarray')
    if A.shape[0] != A.shape[1]:
        raise ValueError('A must be a square matrix (A.shape: {})'.
                         format(A.shape))
    if b.ndim not in [1, 2]:
        raise ValueError('b must be 1D or 2D array (b.shape: {})'.
                         format(b.shape))
    if A.shape[0] != b.shape[0]:
        raise ValueError('The size of dimensions of A must be equal to the '
                         'size of the first dimension of b '
                         '(A.shape: {}, b.shape: {})'.format(A.shape, b.shape))
    if A.dtype.char not in 'fdFD':
        raise TypeError('unsupported dtype (actual: {})'.format(A.dtype))

    if not (sparse.isspmatrix_csr(A) or sparse.isspmatrix_csc(A)):
        warnings.warn('CSR or CSC format is required. Converting to CSR '
                      'format.', sparse.SparseEfficiencyWarning)
        A = A.tocsr()
    A.sum_duplicates()

    if (overwrite_b and A.dtype == b.dtype and
            (b._c_contiguous or b._f_contiguous)):
        x = b
    else:
        x = b.astype(A.dtype, copy=True)

    cusparse.csrsm2(A, x, lower=lower, unit_diag=unit_diagonal)

    if x.dtype.char in 'fF':
        # Note: This is for compatibility with SciPy.
        dtype = numpy.promote_types(x.dtype, 'float64')
        x = x.astype(dtype)
    return x

def spsolve(A, b):
    """Solves a sparse linear system ``A x = b``

    Args:
        A (cupyx.scipy.sparse.spmatrix):
            Sparse matrix with dimension ``(M, M)``.
        b (cupy.ndarray):
            Dense vector or matrix with dimension ``(M)`` or ``(M, 1)``.

    Returns:
        cupy.ndarray:
            Solution to the system ``A x = b``.
    """
    if not cupy.cusolver.check_availability('csrlsvqr'):
        raise NotImplementedError
    if not sparse.isspmatrix(A):
        raise TypeError('A must be cupyx.scipy.sparse.spmatrix')
    if not isinstance(b, cupy.ndarray):
        raise TypeError('b must be cupy.ndarray')
    if A.shape[0] != A.shape[1]:
        raise ValueError('A must be a square matrix (A.shape: {})'.
                         format(A.shape))
    if not (b.ndim == 1 or (b.ndim == 2 and b.shape[1] == 1)):
        raise ValueError('Invalid b.shape (b.shape: {})'.format(b.shape))
    if A.shape[0] != b.shape[0]:
        raise ValueError('matrix dimension mismatch (A.shape: {}, b.shape: {})'
                         .format(A.shape, b.shape))

    if not sparse.isspmatrix_csr(A):
        warnings.warn('CSR format is required. Converting to CSR format.',
                      sparse.SparseEfficiencyWarning)
        A = A.tocsr()
    A.sum_duplicates()
    b = b.astype(A.dtype, copy=False).ravel()

    return cupy.cusolver.csrlsvqr(A, b)

class SuperLU():

    def __init__(self, obj):
        """LU factorization of a sparse matrix.

        Args:
            obj (scipy.sparse.linalg.SuperLU): LU factorization of a sparse
                matrix, computed by `scipy.sparse.linalg.splu`, etc.
        """
        if not scipy_available:
            raise RuntimeError('scipy is not available')
        if not isinstance(obj, scipy.sparse.linalg.SuperLU):
            raise TypeError('obj must be scipy.sparse.linalg.SuperLU')

        self.shape = obj.shape
        self.nnz = obj.nnz
        self.perm_r = cupy.array(obj.perm_r)
        self.perm_c = cupy.array(obj.perm_c)
        self.L = sparse.csr_matrix(obj.L.tocsr())
        self.U = sparse.csr_matrix(obj.U.tocsr())

        self._perm_r_rev = cupy.argsort(self.perm_r)
        self._perm_c_rev = cupy.argsort(self.perm_c)

    def solve(self, rhs, trans='N'):
        """Solves linear system of equations with one or several right-hand sides.

        Args:
            rhs (cupy.ndarray): Right-hand side(s) of equation with dimension
                ``(M)`` or ``(M, K)``.
            trans (str): 'N', 'T' or 'H'.
                'N': Solves ``A * x = rhs``.
                'T': Solves ``A.T * x = rhs``.
                'H': Solves ``A.conj().T * x = rhs``.

        Returns:
            cupy.ndarray:
                Solution vector(s)
        """
        if not isinstance(rhs, cupy.ndarray):
            raise TypeError('ojb must be cupy.ndarray')
        if rhs.ndim not in (1, 2):
            raise ValueError('rhs.ndim must be 1 or 2 (actual: {})'.
                             format(rhs.ndim))
        if rhs.shape[0] != self.shape[0]:
            raise ValueError('shape mismatch (self.shape: {}, rhs.shape: {})'
                             .format(self.shape, rhs.shape))
        if trans not in ('N', 'T', 'H'):
            raise ValueError('trans must be \'N\', \'T\', or \'H\'')
        if not cusparse.check_availability('csrsm2'):
            raise NotImplementedError

        x = rhs.astype(self.L.dtype)
        if trans == 'N':
            if self.perm_r is not None:
                x = x[self._perm_r_rev]
            cusparse.csrsm2(self.L, x, lower=True, transa=trans)
            cusparse.csrsm2(self.U, x, lower=False, transa=trans)
            if self.perm_c is not None:
                x = x[self.perm_c]
        else:
            if self.perm_c is not None:
                x = x[self._perm_c_rev]
            cusparse.csrsm2(self.U, x, lower=False, transa=trans)
            cusparse.csrsm2(self.L, x, lower=True, transa=trans)
            if self.perm_r is not None:
                x = x[self.perm_r]

        if not x._f_contiguous:
            # For compatibility with SciPy
            x = x.copy(order='F')
        return x

class CusparseLU(SuperLU):

    def __init__(self, a):
        """Incomplete LU factorization of a sparse matrix.

        Args:
            a (cupyx.scipy.sparse.csr_matrix): Incomplete LU factorization of a
                sparse matrix, computed by `cusparse.csrilu02`.
        """
        if not scipy_available:
            raise RuntimeError('scipy is not available')
        if not sparse.isspmatrix_csr(a):
            raise TypeError('a must be cupyx.scipy.sparse.csr_matrix')

        self.shape = a.shape
        self.nnz = a.nnz
        self.perm_r = None
        self.perm_c = None
        # TODO(anaruse): Computes tril and triu on GPU
        a = a.get()
        al = scipy.sparse.tril(a)
        al.setdiag(1.0)
        au = scipy.sparse.triu(a)
        self.L = sparse.csr_matrix(al.tocsr())
        self.U = sparse.csr_matrix(au.tocsr())

def factorized(A):
    """Return a function for solving a sparse linear system, with A pre-factorized.

    Args:
        A (cupyx.scipy.sparse.spmatrix): Sparse matrix to factorize.

    Returns:
        callable: a function to solve the linear system of equations given in
        ``A``.

    Note:
        This function computes LU decomposition of a sparse matrix on the CPU
        using `scipy.sparse.linalg.splu`. Therefore, LU decomposition is not
        accelerated on the GPU. On the other hand, the computation of solving
        linear equations using the method returned by this function is
        performed on the GPU.

    .. seealso:: :func:`scipy.sparse.linalg.factorized`
    """
    return splu(A).solve

def splu(A, permc_spec=None, diag_pivot_thresh=None, relax=None,
         panel_size=None, options={}):
    """Computes the LU decomposition of a sparse square matrix.

    Args:
        A (cupyx.scipy.sparse.spmatrix): Sparse matrix to factorize.
        permc_spec (str): (For further augments, see
            :func:`scipy.sparse.linalg.splu`)
        diag_pivot_thresh (float):
        relax (int):
        panel_size (int):
        options (dict):

    Returns:
        cupyx.scipy.sparse.linalg.SuperLU:
            Object which has a ``solve`` method.

    Note:
        This function LU-decomposes a sparse matrix on the CPU using
        `scipy.sparse.linalg.splu`. Therefore, LU decomposition is not
        accelerated on the GPU. On the other hand, the computation of solving
        linear equations using the ``solve`` method, which this function
        returns, is performed on the GPU.

    .. seealso:: :func:`scipy.sparse.linalg.splu`
    """
    if not scipy_available:
        raise RuntimeError('scipy is not available')
    if not sparse.isspmatrix(A):
        raise TypeError('A must be cupyx.scipy.sparse.spmatrix')
    if A.shape[0] != A.shape[1]:
        raise ValueError('A must be a square matrix (A.shape: {})'
                         .format(A.shape))
    if A.dtype.char not in 'fdFD':
        raise TypeError('Invalid dtype (actual: {})'.format(A.dtype))

    a = A.get().tocsc()
    a_inv = scipy.sparse.linalg.splu(
        a, permc_spec=permc_spec, diag_pivot_thresh=diag_pivot_thresh,
        relax=relax, panel_size=panel_size, options=options)
    return SuperLU(a_inv)

def spilu(A, drop_tol=None, fill_factor=None, drop_rule=None,
          permc_spec=None, diag_pivot_thresh=None, relax=None,
          panel_size=None, options={}):
    """Computes the incomplete LU decomposition of a sparse square matrix.

    Args:
        A (cupyx.scipy.sparse.spmatrix): Sparse matrix to factorize.
        drop_tol (float): (For further augments, see
            :func:`scipy.sparse.linalg.spilu`)
        fill_factor (float):
        drop_rule (str):
        permc_spec (str):
        diag_pivot_thresh (float):
        relax (int):
        panel_size (int):
        options (dict):

    Returns:
        cupyx.scipy.sparse.linalg.SuperLU:
            Object which has a ``solve`` method.

    Note:
        This function computes incomplete LU decomposition of a sparse matrix
        on the CPU using `scipy.sparse.linalg.spilu` (unless you set
        ``fill_factor`` to ``1``). Therefore, incomplete LU decomposition is
        not accelerated on the GPU. On the other hand, the computation of
        solving linear equations using the ``solve`` method, which this
        function returns, is performed on the GPU.

        If you set ``fill_factor`` to ``1``, this function computes incomplete
        LU decomposition on the GPU, but without fill-in or pivoting.

    .. seealso:: :func:`scipy.sparse.linalg.spilu`
    """
    if not scipy_available:
        raise RuntimeError('scipy is not available')
    if not sparse.isspmatrix(A):
        raise TypeError('A must be cupyx.scipy.sparse.spmatrix')
    if A.shape[0] != A.shape[1]:
        raise ValueError('A must be a square matrix (A.shape: {})'
                         .format(A.shape))
    if A.dtype.char not in 'fdFD':
        raise TypeError('Invalid dtype (actual: {})'.format(A.dtype))

    if fill_factor == 1:
        # Computes ILU(0) on the GPU using cuSparse functions
        if not sparse.isspmatrix_csr(A):
            a = A.tocsr()
        else:
            a = A.copy()
        cusparse.csrilu02(a)
        return CusparseLU(a)

    a = A.get().tocsc()
    a_inv = scipy.sparse.linalg.spilu(
        a, fill_factor=fill_factor, drop_tol=drop_tol, drop_rule=drop_rule,
        permc_spec=permc_spec, diag_pivot_thresh=diag_pivot_thresh,
        relax=relax, panel_size=panel_size, options=options)
    return SuperLU(a_inv)

max_stars_repo_name: Palash-Vishnani/cupy

max_issues_repo_name: Palash-Vishnani/cupy

max_forks_repo_name: Palash-Vishnani/cupy

reason: import numpy,import scipy,import cupy,from cupy