text:
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.ticker import ScalarFormatter
import seaborn as sns
import numpy
import pandas
import copy
import os
import plotly.graph_objs as go

from .. import Dataset, MSDataset, NMRDataset
from ..enumerations import VariableType, SampleType, AssayRole
from ..utilities import rsd
from ._plotVariableScatter import plotVariableScatter

def plotRSDs(dataset, featureName='Feature Name', ratio=False, logx=True, xlim=None, withExclusions=True, sortOrder=True, savePath=None, color=None, featName=False, hLines=None, figureFormat='png', dpi=72, figureSize=(11,7)):
	"""
	plotRSDs(dataset, ratio=False, savePath=None, color=None \*\*kwargs)

	Visualise analytical *versus* biological variance.

	Plot RSDs calculated in study-reference samples (analytical variance), versus those calculated in study samples (biological variance). RSDs can be visualised either in absolute terms, or as a ratio to analytical variation (*ratio=*\ ``True``).

	:py:func:`plotRSDs` requires that the dataset have at least two samples with the :py:attr:`~nPYc.enumerations.AssayRole.PrecisionReference` :term:`assay role`, if present, RSDs calculated on independent sets of :py:attr:`~nPYc.enumerations.AssayRole.PrecisionReference` samples will also be plotted.

	:param Dataset dataset: Dataset object to plot, the object must have greater that one 'Study Sample' and 'Study-Reference Sample' defined
	:param bool ratio: If ``True`` plot the ratio of analytical variance to biological variance instead of raw values
	:param str featureName: featureMetadata column name by which to label features
	:param bool logx: If ``True`` plot RSDs on a log10 scaled axis
	:param xlim: Tuple of (min, max) RSD values to plot
	:type xlim: None or tuple(float, float)
	:param hLines: None or list of y positions at which to plot an horizontal line. Features are positioned from 1 to nFeat
	:type hLines: None or list
	:param savePath: If ``None`` plot interactively, otherwise save the figure to the path specified
	:type savePath: None or str
	:param color: Allows the default colour pallet to be overridden
	:type color: None or seaborn.palettes._ColorPalette
	:param bool featName: If ``True`` y-axis label is the feature Name, if ``False`` features are numbered.
	"""
	rsdTable = _plotRSDsHelper(dataset, featureName=featureName, ratio=ratio, withExclusions=withExclusions, sortOrder=sortOrder)

	# if RSD not able to be calculated for some features - rsdTable size will be less than dataset.featureMetadata
	if hLines is not None:
		if  dataset.featureMetadata.shape[0] != rsdTable.shape[0]:
			temp = [x for x in rsdTable['Feature Name'].values.tolist() if x in dataset.featureMetadata[featureName][dataset.featureMetadata['Passing Selection'] == False].values.tolist()]
			hLines = [len(temp)]

	# Plot
	if xlim:
		xLim = xlim
	else:
		minRSD = numpy.min(rsdTable[rsdTable.columns[1:]].values)
		maxRSD = numpy.max(rsdTable[rsdTable.columns[1:]].values)
		xLim = (minRSD, maxRSD)

	if logx:
		xlab = 'RSD (%)'
	else:
		xlab = 'RSD (%)'

	# Add Feature Name if required
	if featName:
		rsdTable['yName'] = rsdTable['Feature Name']
		ylab = 'Feature Name'
	else:
		ylab = 'Feature Number'

	plotVariableScatter(rsdTable, logX=logx, xLim=xLim, xLabel=xlab, yLabel=ylab, sampletypeColor=True, hLines=hLines, vLines=None, savePath=savePath, figureFormat=figureFormat, dpi=dpi, figureSize=figureSize)

def plotRSDsInteractive(dataset, featureName='Feature Name', ratio=False, logx=True):
	"""
	Plotly-based interactive version of :py:func:`plotRSDs`

	Visualise analytical *versus* biological variance.

	Plot RSDs calculated in study-reference samples (analytical variance), versus those calculated in study samples (biological variance). RSDs can be visualised either in absolute terms, or as a ratio to analytical variation (*ratio=*\ ``True``).

	:py:func:`plotRSDsInteractive` requires that the dataset have at least two samples with the :py:attr:`~nPYc.enumerations.AssayRole.PrecisionReference` :term:`assay role`, if present, RSDs calculated on independent sets of :py:attr:`~nPYc.enumerations.AssayRole.PrecisionReference` samples will also be plotted.

	:param Dataset dataset: Dataset object to plot, the object must have greater that one 'Study Sample' and 'Study-Reference Sample' defined
	:param str featureName: featureMetadata column name by which to label features
	:param bool ratio: If ``True`` plot the ratio of analytical variance to biological variance instead of raw values
	:param bool logx: If ``True`` plot RSDs on a log10 scaled axis

	"""
	rsdTable = _plotRSDsHelper(dataset, featureName=featureName, ratio=ratio)

	reversedIndex =  numpy.arange(len(rsdTable)-1,-1, -1)
	data = []
	if SampleType.StudySample in rsdTable.columns:
		studySamples = go.Scatter(
			x = rsdTable[SampleType.StudySample].values,
			y = reversedIndex,
			mode = 'markers',
			text = rsdTable['Feature Name'],
			name = 'Study Sample',
			marker = dict(
				color = 'rgba(89, 117, 164, .8)',
			),
			hoverinfo = 'x+text',
		)
		data.append(studySamples)

	if SampleType.ExternalReference in rsdTable.columns:
		externalRef = go.Scatter(
			x = rsdTable[SampleType.ExternalReference].values,
			y = reversedIndex,
			mode = 'markers',
			text = rsdTable['Feature Name'],
			name = 'Long-Term Reference',
			marker = dict(
				color = 'rgba(181, 93, 96, .8)',
			),
			hoverinfo = 'x+text',
		)
		data.append(externalRef)

	if SampleType.StudyPool in rsdTable.columns:
		studyPool = go.Scatter(
			x = rsdTable[SampleType.StudyPool].values,
			y = reversedIndex,
			mode = 'markers',
			text = rsdTable['Feature Name'],
			name = 'Study Reference',
			 marker = dict(
				color = 'rgba(95, 158, 110, .8)',
			),
			hoverinfo = 'x+text',
		)
		data.append(studyPool)

	if logx:
		xaxis = dict(
					type='log',
					title='RSD (%)',
					autorange=True
					)
	else:
		xaxis = dict(
					title='RSD (%)'
					)

	layout = go.Layout(
				title='Feature RSDs',
				legend=dict(
					orientation="h"
				),
				hovermode = "closest",
				yaxis=dict(
						title='Feature Number'
					),
				xaxis=xaxis
				)
	figure = go.Figure(data=data, layout=layout)

	return figure

def _plotRSDsHelper(dataset, featureName='Feature Name', ratio=False, withExclusions=False, sortOrder=True):

	if not dataset.VariableType == VariableType.Discrete:
		raise ValueError('Only datasets with discreetly sampled variables are supported.')

	if sum(dataset.sampleMetadata.loc[dataset.sampleMask, 'SampleType'].values == SampleType.StudySample) <= 2:
		raise ValueError('More than two Study Samples must be defined to calculate biological RSDs.')

	## Calculate RSD for every SampleType with enough PrecisionReference samples.
	rsdVal = dict()

	precRefMask = dataset.sampleMetadata.loc[:, 'AssayRole'].values == AssayRole.PrecisionReference
	precRefMask = numpy.logical_and(precRefMask, dataset.sampleMask)
	sTypes = list(set(dataset.sampleMetadata.loc[precRefMask, 'SampleType'].values))

	if withExclusions:
		rsdVal['Feature Name'] = dataset.featureMetadata.loc[dataset.featureMask, featureName].values
		rsdVal[SampleType.StudyPool] = dataset.rsdSP[dataset.featureMask]
		ssMask = (dataset.sampleMetadata['SampleType'].values == SampleType.StudySample) & dataset.sampleMask
		rsdList = rsd(dataset.intensityData[ssMask, :])
		rsdVal[SampleType.StudySample] = rsdList[dataset.featureMask]
	else:
		rsdVal['Feature Name'] = dataset.featureMetadata.loc[:, featureName].values
		rsdVal[SampleType.StudyPool] = dataset.rsdSP
		ssMask = (dataset.sampleMetadata['SampleType'].values == SampleType.StudySample) & dataset.sampleMask
		rsdList = rsd(dataset.intensityData[ssMask, :])
		rsdVal[SampleType.StudySample] = rsdList

	# Only keep features with finite values for SP and SS
	finiteMask = (rsdVal[SampleType.StudyPool] < numpy.finfo(numpy.float64).max)
	finiteMask = finiteMask & (rsdVal[SampleType.StudySample] < numpy.finfo(numpy.float64).max)

	for sType in sTypes:
		if not sTypes == SampleType.StudyPool:
			sTypeMask = dataset.sampleMetadata.loc[:, 'SampleType'].values == sType
			# precRefMask limits to Precision Reference and dataset.sampleMask
			sTypeMask = numpy.logical_and(sTypeMask, precRefMask)

			# minimum 3 points needed
			if sum(sTypeMask) >= 3:
				rsdList = rsd(dataset.intensityData[sTypeMask, :])
				if withExclusions:
					rsdVal[sType] = rsdList[dataset.featureMask]
				else:
					rsdVal[sType] = rsdList
				finiteMask = finiteMask & (rsdVal[sType] < numpy.finfo(numpy.float64).max)

	## apply finiteMask
	for sType in rsdVal.keys():
		rsdVal[sType] = rsdVal[sType][finiteMask]

	if ratio:
		rsdSP = copy.deepcopy(rsdVal[SampleType.StudyPool])
		for sType in sTypes:
			rsdVal[sType] = numpy.divide(rsdVal[sType], rsdSP)

	# reorder from largest to smallest RSD in Study Pool
	if sortOrder:
		sortIndex = reversed(numpy.argsort(rsdVal[SampleType.StudyPool]))
		rsdTable = pandas.DataFrame(rsdVal).reindex(sortIndex)
		rsdTable.reset_index(drop=True, inplace=True)
	else:
		rsdTable = pandas.DataFrame(rsdVal)

	return rsdTable

max_stars_repo_name: ghaggart/nPYc-Toolbox

max_issues_repo_name: ghaggart/nPYc-Toolbox

max_forks_repo_name: ghaggart/nPYc-Toolbox


text:
import numpy as np
import scipy.io.wavfile as wav
import torch

def get_mask_from_lengths(lengths):
    max_len = torch.max(lengths).item()
    ids = torch.arange(0, max_len, out=torch.cuda.LongTensor(max_len))
    mask = (ids < lengths.unsqueeze(1)).bool()
    return mask

def load_wav_to_torch(full_path):
    sampling_rate, sound = wav.read(full_path)

    if len(sound.shape) > 1:
        if sound.shape[1] == 1:
            sound = sound.squeeze()
        else:
            sound = sound.mean(axis=1)  # multiple channels, average

    return torch.FloatTensor(sound.astype(np.float32)), sampling_rate

def load_filepaths_and_text(filename, split="|"):
    with open(filename, encoding='utf-8') as f:
        filepaths_and_text = [line.strip().split(split) for line in f]
    return filepaths_and_text

def files_to_list(filename):
    """
    Takes a text file of filenames and makes a list of filenames
    """
    with open(filename, encoding='utf-8') as f:
        files = f.readlines()

    files = [f.rstrip() for f in files]
    return files

def to_gpu(x):
    x = x.contiguous()

    if torch.cuda.is_available():
        x = x.cuda(non_blocking=True)
    return torch.autograd.Variable(x)

ext: py

max_stars_repo_name: Rasmusafj/mellotron

max_issues_repo_name: Rasmusafj/mellotron

max_forks_repo_name: Rasmusafj/mellotron

reason: import numpy,import scipy


text:
from netCDF4 import Dataset
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
import numpy as np
from mpl_toolkits.axes_grid1 import make_axes_locatable

#-------------------------------------------------------------------------------

def plot_testcase_diff(testDir, baseDir, strainType, filenameOut):

    if (strainType not in ["varvar","weakvar"]):
        raise RuntimeError("Unknown strainType: %s" %(strainType))

    iTime = -1

    # base dir
    filein = Dataset("./%s/output.2000.nc" %(baseDir),"r")

    uVelocityBase = filein.variables["uVelocity"][iTime,:]
    vVelocityBase = filein.variables["vVelocity"][iTime,:]

    stressDivergenceUBase = filein.variables["stressDivergenceU"][iTime,:]
    stressDivergenceVBase = filein.variables["stressDivergenceV"][iTime,:]

    if (strainType == "varvar"):

        principalStress1Base = filein.variables["principalStress1Var"][iTime,:,2]
        principalStress2Base = filein.variables["principalStress2Var"][iTime,:,2]

        strain11varBase = filein.variables["strain11var"][iTime,:]
        strain22varBase = filein.variables["strain22var"][iTime,:]
        strain12varBase = filein.variables["strain12var"][iTime,:]

    filein.close()

    # test dir
    filein = Dataset("./%s/output.2000.nc" %(testDir),"r")

    nCells = len(filein.dimensions["nCells"])
    nVertices = len(filein.dimensions["nVertices"])
    vertexDegree = len(filein.dimensions["vertexDegree"])
    maxEdges = len(filein.dimensions["maxEdges"])

    cellsOnVertex = filein.variables["cellsOnVertex"][:]
    cellsOnVertex -= 1

    verticesOnCell = filein.variables["verticesOnCell"][:]
    verticesOnCell -= 1

    edgesOnCell = filein.variables["edgesOnCell"][:]
    edgesOnCell -= 1

    xVertex = filein.variables["xVertex"][:]
    yVertex = filein.variables["yVertex"][:]

    xCell = filein.variables["xCell"][:]
    yCell = filein.variables["yCell"][:]

    xEdge = filein.variables["xEdge"][:]
    yEdge = filein.variables["yEdge"][:]

    uVelocityTest = filein.variables["uVelocity"][iTime,:]
    vVelocityTest = filein.variables["vVelocity"][iTime,:]

    stressDivergenceUTest = filein.variables["stressDivergenceU"][iTime,:]
    stressDivergenceVTest = filein.variables["stressDivergenceV"][iTime,:]

    if (strainType == "varvar"):

        principalStress1 = filein.variables["principalStress1Var"][iTime,:,2]
        principalStress2 = filein.variables["principalStress2Var"][iTime,:,2]

        strain11varTest = filein.variables["strain11var"][iTime,:]
        strain22varTest = filein.variables["strain22var"][iTime,:]
        strain12varTest = filein.variables["strain12var"][iTime,:]

    elif (strainType == "weakvar"):

        strain11weakTest = filein.variables["strain11weak"][iTime,:]
        strain22weakTest = filein.variables["strain22weak"][iTime,:]
        strain12weakTest = filein.variables["strain12weak"][iTime,:]

        principalStress1 = filein.variables["principalStress1Weak"][iTime,:]
        principalStress2 = filein.variables["principalStress2Weak"][iTime,:]

    filein.close()

    # differences
    uVelocity = uVelocityTest[:]# - uVelocityBase[:]
    vVelocity = vVelocityTest[:]# - vVelocityBase[:]

    stressDivergenceU = stressDivergenceUTest[:]# - stressDivergenceUBase[:]
    stressDivergenceV = stressDivergenceVTest[:]# - stressDivergenceVBase[:]

    if (strainType == "varvar"):
        strain11var = strain11varTest[:] - strain11varBase[:]
        strain22var = strain22varTest[:] - strain22varBase[:]
        strain12var = strain12varTest[:] - strain12varBase[:]
    elif (strainType == "weakvar"):
        strain11weak = strain11weakTest[:]
        strain22weak = strain22weakTest[:]
        strain12weak = strain12weakTest[:]

    xmin = np.amin(xVertex) - 10000.0
    xmax = np.amax(xVertex) + 10000.0
    ymin = np.amin(yVertex) - 10000.0
    ymax = np.amax(yVertex) + 10000.0

    # get vertex patches list
    patchesVertex = []
    deletedVertices = []
    for iVertex in range(0,nVertices):
        vertices = []
        useVertex = True
        for iCellOnVertex in range(0,vertexDegree):
            iCell = cellsOnVertex[iVertex,iCellOnVertex]
            if (iCell != -1):
                vertices.append((xCell[iCell],yCell[iCell]))
            else:
                useVertex = False
        if (useVertex):
            patchesVertex.append(Polygon(vertices,True))
        else:
            deletedVertices.append(iVertex)
    deletedVertices = np.array(deletedVertices)

    # get strain patches list
    patchesStrain = []
    for iCell in range(0,nCells):
        for iVertexOnCell in range(0,maxEdges):
            iVertex = verticesOnCell[iCell,iVertexOnCell]
            iEdgeOnCell1 = (iVertexOnCell - 1)
            if (iEdgeOnCell1 < 0): iEdgeOnCell1 += maxEdges
            iEdgeOnCell2 = iVertexOnCell
            #iEdgeOnCell1 = iVertexOnCell
            #iEdgeOnCell2 = iVertexOnCell + 1
            #if (iEdgeOnCell2 > maxEdges-1): iEdgeOnCell2 -= maxEdges
            iEdge1 = edgesOnCell[iCell,iEdgeOnCell1]
            iEdge2 = edgesOnCell[iCell,iEdgeOnCell2]
            vertices = [(xCell[iCell],yCell[iCell]),
                        (xEdge[iEdge1],yEdge[iEdge1]),
                        (xVertex[iVertex],yVertex[iVertex]),
                        (xEdge[iEdge2],yEdge[iEdge2])]
            patchesStrain.append(Polygon(vertices,True))

    # patches cell
    patchesCell = []
    for iCell in range(0,nCells):
        vertices = []
        for iVertexOnCell in range(0,maxEdges):
            iVertex = verticesOnCell[iCell,iVertexOnCell]
            vertices.append((xVertex[iVertex],yVertex[iVertex]))
        patchesCell.append(Polygon(vertices,True))

    uVelocity = np.delete(uVelocity, deletedVertices)
    vVelocity = np.delete(vVelocity, deletedVertices)
    stressDivergenceU = np.delete(stressDivergenceU, deletedVertices)
    stressDivergenceV = np.delete(stressDivergenceV, deletedVertices)
    principalStress1 = np.where(principalStress1 < 1e20, principalStress1, 0.0)
    principalStress2 = np.where(principalStress2 < 1e20, principalStress2, 0.0)

    # patch collections
    pcUVelocity = PatchCollection(patchesVertex, cmap=plt.get_cmap("jet"))
    pcUVelocity.set_array(uVelocity)

    pcVVelocity = PatchCollection(patchesVertex, cmap=plt.get_cmap("jet"))
    pcVVelocity.set_array(vVelocity)

    pcStressDivergenceU = PatchCollection(patchesVertex, cmap=plt.get_cmap("jet"))
    pcStressDivergenceU.set_array(stressDivergenceU)

    pcStressDivergenceV = PatchCollection(patchesVertex, cmap=plt.get_cmap("jet"))
    pcStressDivergenceV.set_array(stressDivergenceV)

    strainMin = None#-5e-7
    strainMax = None# 5e-7

    if (strainType == "varvar"):

        pcStrain11 = PatchCollection(patchesStrain, cmap=plt.get_cmap("jet"))
        pcStrain11.set_array(strain11var.flatten())
        pcStrain11.set_clim([strainMin, strainMax])

        pcStrain22 = PatchCollection(patchesStrain, cmap=plt.get_cmap("jet"))
        pcStrain22.set_array(strain22var.flatten())
        pcStrain22.set_clim([strainMin, strainMax])

        pcStrain12 = PatchCollection(patchesStrain, cmap=plt.get_cmap("jet"))
        pcStrain12.set_array(strain12var.flatten())
        pcStrain12.set_clim([strainMin, strainMax])

    elif (strainType == "weakvar"):

        pcStrain11 = PatchCollection(patchesCell, cmap=plt.get_cmap("jet"))
        pcStrain11.set_array(strain11weak.flatten())
        pcStrain11.set_clim([strainMin, strainMax])

        pcStrain22 = PatchCollection(patchesCell, cmap=plt.get_cmap("jet"))
        pcStrain22.set_array(strain22weak.flatten())
        pcStrain22.set_clim([strainMin, strainMax])

        pcStrain12 = PatchCollection(patchesCell, cmap=plt.get_cmap("jet"))
        pcStrain12.set_array(strain12weak.flatten())
        pcStrain12.set_clim([strainMin, strainMax])

    # plot
    fig, axes = plt.subplots(2,4,figsize=(10,5))

    axes[0,0].add_collection(pcUVelocity)
    axes[0,0].set_xlim((xmin,xmax))
    axes[0,0].set_ylim((ymin,ymax))
    axes[0,0].set_aspect('equal')
    axes[0,0].set_xticks([])
    axes[0,0].set_yticks([])
    axes[0,0].set_title("uVelocity")
    divider = make_axes_locatable(axes[0,0])
    cax = divider.append_axes('right', size='5%', pad=0.05)
    fig.colorbar(pcUVelocity, cax=cax)

    axes[0,1].add_collection(pcVVelocity)
    axes[0,1].set_xlim((xmin,xmax))
    axes[0,1].set_ylim((ymin,ymax))
    axes[0,1].set_aspect('equal')
    axes[0,1].set_xticks([])
    axes[0,1].set_yticks([])
    axes[0,1].set_title("vVelocity")
    divider = make_axes_locatable(axes[0,1])
    cax = divider.append_axes('right', size='5%', pad=0.05)
    fig.colorbar(pcVVelocity, cax=cax)

    axes[0,2].add_collection(pcStressDivergenceU)
    axes[0,2].set_xlim((xmin,xmax))
    axes[0,2].set_ylim((ymin,ymax))
    axes[0,2].set_aspect('equal')
    axes[0,2].set_xticks([])
    axes[0,2].set_yticks([])
    axes[0,2].set_title("stressDivU")
    divider = make_axes_locatable(axes[0,2])
    cax = divider.append_axes('right', size='5%', pad=0.05)
    fig.colorbar(pcStressDivergenceU, cax=cax)

    axes[0,3].add_collection(pcStressDivergenceV)
    axes[0,3].set_xlim((xmin,xmax))
    axes[0,3].set_ylim((ymin,ymax))
    axes[0,3].set_aspect('equal')
    axes[0,3].set_xticks([])
    axes[0,3].set_yticks([])
    axes[0,3].set_title("stressDivV")
    divider = make_axes_locatable(axes[0,3])
    cax = divider.append_axes('right', size='5%', pad=0.05)
    fig.colorbar(pcStressDivergenceV, cax=cax)

    axes[1,0].add_collection(pcStrain11)
    axes[1,0].set_xlim((xmin,xmax))
    axes[1,0].set_ylim((ymin,ymax))
    axes[1,0].set_aspect('equal')
    axes[1,0].set_xticks([])
    axes[1,0].set_yticks([])
    axes[1,0].set_title("strain11")
    divider = make_axes_locatable(axes[1,0])
    cax = divider.append_axes('right', size='5%', pad=0.05)
    fig.colorbar(pcStrain11, cax=cax)

    axes[1,1].add_collection(pcStrain22)
    axes[1,1].set_xlim((xmin,xmax))
    axes[1,1].set_ylim((ymin,ymax))
    axes[1,1].set_aspect('equal')
    axes[1,1].set_xticks([])
    axes[1,1].set_yticks([])
    axes[1,1].set_title("strain22")
    divider = make_axes_locatable(axes[1,1])
    cax = divider.append_axes('right', size='5%', pad=0.05)
    fig.colorbar(pcStrain22, cax=cax)

    axes[1,2].add_collection(pcStrain12)
    axes[1,2].set_xlim((xmin,xmax))
    axes[1,2].set_ylim((ymin,ymax))
    axes[1,2].set_aspect('equal')
    axes[1,2].set_xticks([])
    axes[1,2].set_yticks([])
    axes[1,2].set_title("strain12")
    divider = make_axes_locatable(axes[1,2])
    cax = divider.append_axes('right', size='5%', pad=0.05)
    fig.colorbar(pcStrain12, cax=cax)

    axes[1,3].scatter(principalStress1, principalStress2, s=1, c="k")
    axes[1,3].set_aspect('equal')
    axes[1,3].set_xticks([])
    axes[1,3].set_yticks([])

    #axes[1,1].axis('off')

    plt.tight_layout()
    plt.savefig(filenameOut,dpi=1200)

#-------------------------------------------------------------------------------

def plot_testcase():

    plot_testcase_diff("output_hex_pwl_0082x0094_120",
                       "output_hex_wachspress_0082x0094_120",
                       "varvar",
                       "square_pwl_wachs_diff.png")

    plot_testcase_diff("output_hex_pwlavg_0082x0094_120",
                       "output_hex_wachsavg_0082x0094_120",
                       "varvar",
                       "square_pwlavg_wachsavg_diff.png")

    plot_testcase_diff("output_hex_weak_0082x0094_120",
                       "output_hex_wachspress_0082x0094_120",
                       "weakvar",
                       "square_weak_wachs_diff.png")

    plot_testcase_diff("output_hex_weakwachs_0082x0094_120",
                       "output_hex_wachspress_0082x0094_120",
                       "varvar",
                       "square_weakwachs_wachs_diff.png")

#-------------------------------------------------------------------------------

if __name__ == "__main__":

    plot_testcase()

max_stars_repo_licenses: zlib-acknowledgement FTL RSA-MD

max_issues_repo_licenses: zlib-acknowledgement FTL RSA-MD

max_forks_repo_licenses: zlib-acknowledgement FTL RSA-MD

reason: import numpy


text:
"""
This file contains a usage script, intended to test using interface.
Developed by Taras Kucherenko (tarask@kth.se)
"""

import train as tr
import utils.data as dt
import utils.flags as fl
from learn_dataset_encoding import create_nn, prepare_motion_data

import numpy as np

import sys

DATA_DIR = sys.argv[1]
TEST_FILE = sys.argv[2]
OUTPUT_FILE = sys.argv[3]

if __name__ == '__main__':

    # Get the data
    Y_train_normalized, Y_train, Y_test_normalized, Y_test, Y_dev_normalized, max_val, mean_pose  = prepare_motion_data(DATA_DIR)

    # Train the network
    nn = create_nn(Y_train_normalized, Y_dev_normalized, max_val, mean_pose, restoring=True)

    # Read the encoding
    encoding = np.loadtxt(TEST_FILE)

    print(encoding.shape)

    # Decode it
    decoding = tr.decode(nn, encoding)

    print(decoding.shape)

    np.save(OUTPUT_FILE, decoding)

    # Close Tf session
    nn.session.close()


text:
import os

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import linear_sum_assignment
from scipy.spatial import distance

import data_generator
from model import Generator, Critic
from reg_losses import get_regularization_term

slim = tf.contrib.slim

__eval_step_list__ = [10, 50, 100, 250, 500, 1000, 2500, 5000, 10000, 15000, 20000]

flags = tf.app.flags
flags.DEFINE_integer("n_epoch", 20000, "Epoch to train [20000]")
flags.DEFINE_integer("n_batch_size", 256, "Batch size to train [256]")
flags.DEFINE_integer("latent_dimensionality", 2, "Dimensionality of the latent variables [2]")

"""
During training, 10 critic updates are performed for every generator update,
except for the first 25 generator updates,
where the critic is updated 100 times for each generator update
in order to get closer to the optimal critic in the beginning of training.
"""

flags.DEFINE_integer("begining_init_step", 25, "[25]")
flags.DEFINE_integer("n_c_iters_under_begining_init_step", 100, "[100]")
flags.DEFINE_integer("n_c_iters_over_begining_init_step", 10, "[10]")
flags.DEFINE_integer("interval_record_earth_mover", 10, "[10]")

flags.DEFINE_float("learning_rate", 5e-5, "Learning rate of optimizer [5e-5]")
flags.DEFINE_float("Lambda", 5., "Weights for critics' regularization term [5]")
flags.DEFINE_string("Regularization_type", "LP", "[no_reg, no_reg_but_clipping, LP, GP]")
flags.DEFINE_string("Purturbation_type", "dragan_only_training",
                    "[no_purf, wgan_gp, dragan_only_training, dragan_both]")
flags.DEFINE_string("dataset", 'GeneratorSwissRoll',
                    "Which dataset is used? [GeneratorGaussians8, GeneratorGaussians25, GeneratorSwissRoll]")

flags.DEFINE_string("critic_variable_scope_name", "Critic", "[Critic]")
flags.DEFINE_string("generator_variable_scope_name", "Generator", "Generator")

flags.DEFINE_bool("emd_records", False, "Whether EMD is recorded. (It takes some time...)[True, False]")
FLAGS = flags.FLAGS

class Trainer(object):
    def __init__(self):
        self.dataset_generator = None
        self.real_input = None

        self.z = None

        self.generator = None
        self.critic_x = None
        self.critic_gz = None

        self.g_loss = None
        self.c_negative_loss = None
        self.c_regularization_loss = None
        self.c_loss = None
        self.c_clipping = None
        self.x_hat = None

        self.ckpt_dir = None
        self.summary_writer = None
        self.c_summary_op = None
        self.g_summary_op = None
        self.emd_placeholder = None
        self.emd_summary = None

        self.saver = None

        self.step = None

        self.sess = None
        self.step_inc = None
        self.g_opt = None
        self.c_opt = None

        self.g_update_fetch_dict = None
        self.c_update_fetch_dict = None
        self.c_feed_dict = None

        self.coord = None
        self.threads = None

        self.define_dataset()
        self.define_latent()
        self.define_model()
        self.define_loss()
        self.define_optim()
        self.define_writer_and_summary()
        self.define_saver()
        self.initialize_session_and_etc()
        self.define_feed_and_fetch()

    def define_dataset(self):
        self.dataset_generator = iter(getattr(data_generator, FLAGS.dataset)(FLAGS.n_batch_size))
        self.real_input = tf.placeholder(tf.float32, shape=(None, 2))

    def define_latent(self):
        self.z = tf.random_normal([FLAGS.n_batch_size, FLAGS.latent_dimensionality], mean=0.0, stddev=1.0, name='z')

    def define_model(self):
        self.generator = Generator(self.z,
                                   variable_scope_name=FLAGS.generator_variable_scope_name)
        self.critic_x = Critic(self.real_input,
                               variable_scope_name=FLAGS.critic_variable_scope_name)
        self.critic_gz = Critic(self.generator.output_tensor,
                                variable_scope_name=FLAGS.critic_variable_scope_name,
                                reuse=True)

    def define_loss(self):
        self.g_loss = -tf.reduce_mean(self.critic_gz.output_tensor)
        self.c_negative_loss = -self.g_loss - tf.reduce_mean(self.critic_x.output_tensor)
        if FLAGS.Regularization_type == 'no_reg_but_clipping' or \
           FLAGS.Regularization_type == 'no_reg':
            self.c_regularization_loss = tf.Variable(0., trainable=False)
        else:
            self.c_regularization_loss, self.x_hat = get_regularization_term(
                                                training_samples=self.real_input,
                                                generated_samples=self.generator.output_tensor,
                                                reg_type=FLAGS.Regularization_type,
                                                per_type=FLAGS.Purturbation_type,
                                                critic_variable_scope_name=FLAGS.critic_variable_scope_name
                                                )

        self.c_loss = self.c_negative_loss + FLAGS.Lambda * self.c_regularization_loss

    def define_optim(self):
        self.step = tf.Variable(0, name='step', trainable=False)
        self.step_inc = tf.assign(self.step, self.step + 1)

        optimizer = tf.train.RMSPropOptimizer(FLAGS.learning_rate)

        self.g_opt = optimizer.minimize(self.g_loss, var_list=self.generator.var_list)
        self.c_opt = optimizer.minimize(self.c_loss, var_list=self.critic_x.var_list)

        with tf.control_dependencies([self.c_opt]):
            if FLAGS.Regularization_type == 'no_reg_but_clipping':
                self.c_clipping = [p.assign(tf.clip_by_value(p, -0.01, 0.01)) for p in self.critic_x.var_list]
            else:
                self.c_clipping = [tf.no_op()]

    def define_writer_and_summary(self):
        self.ckpt_dir = ''.join(['ckpts/',
                                 FLAGS.dataset+'_',
                                 FLAGS.Regularization_type+'_',
                                 FLAGS.Purturbation_type+'_',
                                 str(FLAGS.Lambda)+'_',
                                 str(FLAGS.emd_records),
                                 '/'])

        if not os.path.exists(self.ckpt_dir):
            os.makedirs(self.ckpt_dir)

        self.summary_writer = tf.summary.FileWriter(self.ckpt_dir)

        self.c_summary_op = tf.summary.merge([
            tf.summary.scalar('loss/c', self.c_loss),
            tf.summary.scalar('loss/c_negative_loss', self.c_negative_loss),
            tf.summary.scalar('loss/c_regularization_loss', self.c_regularization_loss)
        ])
        self.g_summary_op = tf.summary.merge([
            tf.summary.scalar('loss/g', self.g_loss)
        ])

        self.emd_placeholder = tf.placeholder(tf.float32, shape=())
        self.emd_summary = tf.summary.scalar('EMD', self.emd_placeholder)

    def define_saver(self):
        self.saver = tf.train.Saver()

    def initialize_session_and_etc(self):
        gpu_options = tf.GPUOptions(allow_growth=True)
        sess_config = tf.ConfigProto(allow_soft_placement=True,
                                     gpu_options=gpu_options)
        self.sess = tf.Session(config=sess_config)

        self.sess.run(tf.local_variables_initializer())
        self.sess.run(tf.global_variables_initializer())

        self.coord = tf.train.Coordinator()
        self.threads = tf.train.start_queue_runners(sess=self.sess, coord=self.coord)

    def define_feed_and_fetch(self):
        self.g_update_fetch_dict = {
            "opt": self.g_opt,
            "z": self.z,
            "G_z": self.generator.output_tensor,
            "loss": self.g_loss,
            'summary': self.g_summary_op,
            "step": self.step
        }

        self.c_update_fetch_dict = {
            'gradient_clipping': self.c_clipping,
            "x": self.real_input,
            "G_z": self.generator.output_tensor,
            "loss": self.c_loss,
            "negative_loss": self.c_negative_loss,
            "regularization_loss": self.c_regularization_loss,
            'summary': self.c_summary_op,
            "step": self.step
        }

        self.c_feed_dict = {
            self.real_input: None
        }

    def draw_level_sets(self, step,
                        x_min=-2.5, x_max=2.5,
                        y_min=-2.5, y_max=2.5,
                        n_batch=2):
        x = np.linspace(x_min, x_max, 200)
        y = np.linspace(y_min, y_max, 200)

        x, y = np.meshgrid(x, y)
        grid_pts = np.stack([x.flatten(), y.flatten()], axis=1)

        z = self.sess.run(self.critic_x.output_tensor, feed_dict={self.real_input: grid_pts})
        z = np.reshape(z, (200, 200))

        plt.contour(x, y, z, 30, cmap='copper')

        real = list()
        fake = list()
        perturbated = list()
        for i in range(n_batch):
            __real__ = next(self.dataset_generator)

            __fake__, __perturbated__ = \
                self.sess.run([self.generator.output_tensor, self.x_hat], feed_dict={self.real_input: __real__})
            real.append(__real__)
            fake.append(__fake__)
            perturbated.append(__perturbated__)

        real = np.vstack(real)
        fake = np.vstack(fake)
        perturbated = np.vstack(perturbated)

        plt.scatter(perturbated[:, 0], perturbated[:, 1], c='r', s=1)
        plt.scatter(real[:, 0], real[:, 1], c='y', s=2)
        plt.scatter(fake[:, 0], fake[:, 1], c='g', s=2)

        plt.savefig(self.ckpt_dir+str(step)+'.png')
        plt.clf()

    def estimate_earth_mover_distance(self, step, n_batch=2):
        real = list()
        fake = list()

        for i in range(n_batch):
            __real__ = next(self.dataset_generator)
            __fake__ = self.sess.run(self.generator.output_tensor)

            real.append(__real__)
            fake.append(__fake__)

        real = np.vstack(real)
        fake = np.vstack(fake)

        cost_matrix = distance.cdist(fake, real, 'euclidean')

        row_ind, col_ind = linear_sum_assignment(cost_matrix)
        linear_sum = cost_matrix[row_ind, col_ind].sum()
        emd = linear_sum/real.shape[0]

        emd_fetch = self.sess.run(self.emd_summary, feed_dict={self.emd_placeholder: emd})
        self.summary_writer.add_summary(emd_fetch, step)
        self.summary_writer.flush()

    def train(self):
        try:
            c_fetch_dict = None
            print("[.] Learning Start...")
            step = 0
            while not self.coord.should_stop():
                if step > FLAGS.n_epoch:
                    break

                self.c_feed_dict[self.real_input] = next(self.dataset_generator)
                step = self.sess.run(self.step)

                n_c_iters = (FLAGS.n_c_iters_under_begining_init_step
                             if step < FLAGS.begining_init_step
                             else FLAGS.n_c_iters_over_begining_init_step)
                for _ in range(n_c_iters):
                    c_fetch_dict = self.sess.run(self.c_update_fetch_dict,
                                                 feed_dict=self.c_feed_dict)

                g_fetch_dict = self.sess.run(self.g_update_fetch_dict)

                self.summary_writer.add_summary(c_fetch_dict["summary"], c_fetch_dict["step"])
                self.summary_writer.add_summary(g_fetch_dict["summary"], g_fetch_dict["step"])
                self.summary_writer.flush()

                if step in __eval_step_list__:
                    self.draw_level_sets(step)

                if FLAGS.emd_records and step % FLAGS.interval_record_earth_mover == 0 and step != 0:
                    self.estimate_earth_mover_distance(step)

                self.sess.run(self.step_inc)

        except KeyboardInterrupt:
            print("Interrupted")
            self.coord.request_stop()
        finally:
            self.saver.save(self.sess, self.ckpt_dir)
            print('Stop')
            self.coord.request_stop()
            self.coord.join(self.threads)

if __name__ == '__main__':
    trainer = Trainer()
    trainer.train()
    trainer.sess.close()

max_stars_repo_name: mikigom/WGAN-LP-tensorflow

max_issues_repo_name: mikigom/WGAN-LP-tensorflow

max_forks_repo_name: mikigom/WGAN-LP-tensorflow

reason: import numpy,from scipy


text:
import sys
import argparse
import numpy as np

from doom.actions import ActionBuilder
from doom.game import Game
from doom.utils import register_game_args
from src.utils import bool_flag, register_model_args
from src.agent import Agent

import tensorflow as tf

def main():
    parser = argparse.ArgumentParser(description='LUBAN runner')
    register_model_args(parser)
    register_game_args(parser)
    params, unparsed = parser.parse_known_args(sys.argv)
    params.game_variables = [('health', 101), ('sel_ammo', 301)]

    params.frame_skip = 3
    params.eval_time = 900

    sess = tf.Session()
    agent = Agent(sess, params)
    agent.load("./checkpoint")

    action_builder = ActionBuilder(params)
    game = Game(
        scenario='full_deathmatch',
        action_builder=action_builder,
        score_variable='USER2',
        freedoom=True,
        screen_resolution='RES_400X225',
        screen_format='CRCGCB',
        use_screen_buffer=True,
        use_depth_buffer=False,
        labels_mapping="",
        game_features="target,enemy",
        mode=('SPECTATOR' if params.human_player else 'PLAYER'),
        render_hud=False,
        render_crosshair=True,
        render_weapon=True,
        freelook=params.freelook,
        visible=0,
        n_bots=4,
        use_scripted_marines=True
    )
    game.start(map_id=3)
    game.init_bots_health(100)

    n_iter = 0
    last_states = []
    # eval_time: seconds, 35: fps, eval_time*35: total frames
    while n_iter * params.frame_skip < params.eval_time * 35:
        n_iter += 1

        if game.is_player_dead():
            game.respawn_player()
            # TODO: add reset method to network
            # network.reset()

         # observe the game state / select the next action
        game.observe_state(params, last_states)
        # single frame
        screen1 = last_states[11].screen.reshape((1, 1, 60, 108, 3))
        screen1 = screen1.astype(np.float32) / 255.
        # two frame
        screen2 = np.array([last_states[i+10].screen.transpose((1, 2, 0)) for i in range(2)]).reshape((1, 2, 60, 108, 3))
        screen2 = screen2.astype(np.float32) / 255.
        # four frame
        screen4 = np.array([last_states[i+8].screen.transpose((1, 2, 0)) for i in range(4)]).reshape((1, 4, 60, 108, 3))
        screen4 = screen4.astype(np.float32) / 255.
        # eight frame
        screen8 = np.array([last_states[i+4].screen.transpose((1, 2, 0)) for i in range(8)]).reshape((1, 8, 60, 108, 3))
        screen8 = screen8.astype(np.float32) / 255.
        # eight frame
        screen12 = np.array([last_states[i].screen.transpose((1, 2, 0)) for i in range(12)]).reshape((1, 12, 60, 108, 3))
        screen12 = screen12.astype(np.float32) / 255.
        action, output_gf = agent.choose_action(screen2)
        '''
        if action == 0:
            action = 5
        elif action == 1:
            action = 7
        elif action == 2:
            action = 9
        elif action == 3:
            action = 2
        elif action == 4:
            action = 4
        else:
            action = 0
        '''
        print("target: %.4f, enemy: %.4f" % (output_gf[0][0], output_gf[0][1]))
        #action = np.random.randint(0, 29)
        print("health: %d\tsel_ammo: %d\taction: %d" % (last_states[3].variables[0], last_states[3].variables[1], action))

        sleep = 0.01
        game.make_action(action, frame_skip=params.frame_skip, sleep=sleep)

    game.close()
    game.print_statistics()

if __name__ == '__main__':
    main()

max_stars_repo_name: luxinglong/ViZDoom-SL

max_issues_repo_name: luxinglong/ViZDoom-SL

max_forks_repo_name: luxinglong/ViZDoom-SL


text:
import numpy as np
import pytest
from omlt import OffsetScaling
from omlt.scaling import convert_to_dict

def test_convert_to_dict():
    x = ['a', 'b']
    x = convert_to_dict(x)
    assert sorted(x.keys()) == [0, 1]
    assert x[0] == 'a'
    assert x[1] == 'b'

    x = {2: 'a', 1: 'b'}
    x = convert_to_dict(x)
    assert sorted(x.keys()) == [1, 2]
    assert x[2] == 'a'
    assert x[1] == 'b'

def test_offset_scaling():
    x_offset = [42, 65]
    x_factor = [1975, 1964]
    y_offset = [-4, 2, 1.784]
    y_factor = [2, 1.5, 1.3]

    x = {0: 1, 1: 2}
    y = {0: -1, 1: 2, 2: 3}
    x_scal = {k: (v - x_offset[k]) / x_factor[k] for (k, v) in x.items()}
    y_scal = {k: (v - y_offset[k]) / y_factor[k] for (k, v) in y.items()}

    scaling = OffsetScaling(
        offset_inputs=x_offset,
        factor_inputs=x_factor,
        offset_outputs=y_offset,
        factor_outputs=y_factor,
    )

    test_x_scal = scaling.get_scaled_input_expressions(x)
    test_x = scaling.get_unscaled_input_expressions(test_x_scal)
    test_y = scaling.get_unscaled_output_expressions(y_scal)
    test_y_scal = scaling.get_scaled_output_expressions(y)
    np.testing.assert_almost_equal(list(test_x_scal.values()), list(x_scal.values()))
    np.testing.assert_almost_equal(list(test_x.values()), list(x.values()))
    np.testing.assert_almost_equal(list(test_y.values()), list(y.values()))
    np.testing.assert_almost_equal(list(test_y_scal.values()), list(y_scal.values()))

def test_incorrect_keys():
    x_offset = {1: 42, 42: 65}
    x_factor = {1: 1975, 42: 1964}
    y_offset = {7: -4, 9: 2, 11: 1.784}
    y_factor = {7: 2, 9: 1.5, 11: 1.3}

    x = {1: 1, 42: 2}
    y = {7: -1, 9: 2, 11: 3}
    x_scal = {k: (v - x_offset[k]) / x_factor[k] for (k, v) in x.items()}
    y_scal = {k: (v - y_offset[k]) / y_factor[k] for (k, v) in y.items()}

    scaling = OffsetScaling(
        offset_inputs=x_offset,
        factor_inputs=x_factor,
        offset_outputs=y_offset,
        factor_outputs=y_factor,
    )

    test_x_scal = scaling.get_scaled_input_expressions(x)
    test_y_unscal = scaling.get_unscaled_output_expressions(y_scal)
    np.testing.assert_almost_equal(list(test_x_scal.values()), list(x_scal.values()))
    np.testing.assert_almost_equal(list(test_y_unscal.values()), list(y.values()))

    x = {1: 42, 2: 65}
    with pytest.raises(ValueError) as excinfo:
        test_x_scal = scaling.get_scaled_input_expressions(x)
    expected_msg = "get_scaled_input_expressions called with input_vars that " \
        "do not have the same indices as offset_inputs or factor_inputs.\nKeys " \
        "in input_vars: [1, 2].\nKeys in offset_inputs: [1, 42].\nKeys in "\
        "offset_factor: [1, 42]."
    assert str(excinfo.value) == expected_msg

    y = {7: -1, 19: 2, 11: 3}
    with pytest.raises(ValueError) as excinfo:
        test_y_scal = scaling.get_scaled_output_expressions(y)
    expected_msg = "get_scaled_output_expressions called with output_vars that " \
        "do not have the same indices as offset_outputs or factor_outputs.\nKeys " \
        "in output_vars: [7, 11, 19]\nKeys in offset_outputs: [7, 9, 11]\nKeys in " \
        "offset_factor: [7, 9, 11]"
    assert str(excinfo.value) == expected_msg

    x_scal = {1: 42, 2: 65}
    with pytest.raises(ValueError) as excinfo:
        test_x_unscal = scaling.get_unscaled_input_expressions(x_scal)
    expected_msg = "get_scaled_input_expressions called with input_vars that " \
        "do not have the same indices as offset_inputs or factor_inputs.\nKeys " \
        "in input_vars: [1, 2]\nKeys in offset_inputs: [1, 42]\nKeys in " \
        "offset_factor: [1, 42]"
    assert str(excinfo.value) == expected_msg

    y_scal = {7: -1, 8: 2, 11: 3}
    with pytest.raises(ValueError) as excinfo:
        test_y_unscal = scaling.get_unscaled_output_expressions(y_scal)
    expected_msg = "get_scaled_output_expressions called with output_vars that do " \
        "not have the same indices as offset_outputs or factor_outputs.\nKeys in " \
        "output_vars: [7, 8, 11]\nKeys in offset_outputs: [7, 9, 11]\nKeys in " \
        "offset_factor: [7, 9, 11]"
    assert str(excinfo.value) == expected_msg

def test_negative_offsets():
    x_offset = [42, 65]
    x_factor = [-1975, 1964]
    y_offset = [-4, 2, 1.784]
    y_factor = [2, 1.5, 1.3]

    with pytest.raises(ValueError) as excinfo:
        scaling = OffsetScaling(
            offset_inputs=x_offset,
            factor_inputs=x_factor,
            offset_outputs=y_offset,
            factor_outputs=y_factor,
        )
    assert str(excinfo.value) == "OffsetScaling only accepts positive values" \
       " for factor_inputs. Negative value found at" \
       " index 0."

    x_offset = [42, 65]
    x_factor = [1975, 1964]
    y_offset = [-4, 2, 1.784]
    y_factor = [2, -1.5, 1.3]

    with pytest.raises(ValueError) as excinfo:
        scaling = OffsetScaling(
            offset_inputs=x_offset,
            factor_inputs=x_factor,
            offset_outputs=y_offset,
            factor_outputs=y_factor,
        )
    assert str(excinfo.value) == "OffsetScaling only accepts positive values" \
       " for factor_outputs. Negative value found at" \
       " index 1."

if __name__ == '__main__':
    test_incorrect_keys()

max_stars_repo_name: fracek/OMLT

max_issues_repo_name: fracek/OMLT

max_forks_repo_name: fracek/OMLT


text:
import os
import csv
import cv2
import numpy as np # for np.array() np.append()
from datetime import datetime # print timestamps

# for loss history visualization image
import matplotlib; matplotlib.use('agg')
import matplotlib.pyplot as plt

from scipy import ndimage

# list of common functions
# from commonFunctions_vxx import get_log_path
# from commonFunctions_vxx import get_lines_logfile
# from commonFunctions_vxx import get_info_from_lines
# from commonFunctions_vxx import get_info_from_logfile
# from commonFunctions_vxx import flip_horizontally
# from commonFunctions_vxx import visualize_loss_history
# from commonFunctions_vxx import RGB2YUV
# from commonFunctions_vxx import print_info
# from commonFunctions_vxx import

# History
# v01 : Start
# v02 : add nb_images to read parameter
# v03 : from scipy import ndimage, due to cv2.imread will get images in BGR format, while drive.py uses RGB. In the video above one way you could keep the same image formatting is to do "image = ndimage.imread(current_path)"
# v04 : use left + right images to augment data + measurements extrapolations
# v05 : add Generator Function + modify all other functions whenever necessary to use generator function ...
# v06 : Re-start from v04 aas fit_generator and need to add generator obsolete.
#       Latest Keras.Model.fit integrates a generator in itself.
#        ie v06 : visualize loss history
# v07 : For nvidia model, convert RGB to YUV
# v08 : add print_info() to print debug/progress info
# v12 : Try to avoid list to numpy conversion, taking few minutes, start with numpy image array straight from start
#       But failed. Need to use and adapt to Generator

driving_log_file = 'driving_log.csv'

# Select right sample data folder whether in GPU mode or not
# check if ./data/driving_log.csv exists otherwise select
# simulationData/001_1stTrackSampleDrivingData/
def get_log_path() :
    if os.path.exists("./data/" + driving_log_file) :
        return("./data/")
    else :
        return("./simulationData/001_1stTrackSampleDrivingData/")

def get_lines_logfile() :
    l = []
    with open (get_log_path() + driving_log_file ) as csv_file :
        reader = csv.reader(csv_file)
        for line in reader :
            l.append(line)
    return l

def get_info_from_lines(l,leftright_steer_corr,nb_images=None) :
    imgs = []
    meas = []

    log_path = get_log_path()
    first_im=ndimage.imread(log_path + l[1][0].strip())
    row,col,channel = first_im.shape
    print('image shape : {}'.format(first_im.shape))
    np_imgs = None
    nb = 2*len(l)
    print('nb : {}'.format(nb))
    np_imgs = np.zeros((nb,row,col,channel))

    print('Size np_imgs : {}'.format(np_imgs.shape))

    for line in l[1:nb_images] :
        #image = cv2.imread(log_path + line[0].strip())
        for i in range(3) :         # center image, left , right images
            image = ndimage.imread(log_path + line[i].strip())
            imgs.append(image)
        measurement = float(line[3]) # center image
        meas.append(measurement)
        measurement += leftright_steer_corr # left image
        meas.append(measurement)
        measurement -= leftright_steer_corr # right image
        meas.append(measurement)
    return imgs,meas

def get_info_from_logfile(leftright_steer_correction,nb_images=None) :
    lines = get_lines_logfile()
    return get_info_from_lines(lines,leftright_steer_correction,nb_images)

def flip_horizontally(img,meas):
    aug_img, aug_meas = [],[]
    for i,m in zip(img,meas) :
        aug_img.append(cv2.flip(i,1))
        aug_meas.append(m*(-1.0))
    return aug_img,aug_meas

def visualize_loss_history(history) :
    ### plot the training and validation loss for each epoch
    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('model mean squared error loss')
    plt.ylabel('mean squared error loss')
    plt.xlabel('epoch')
    plt.legend(['training set', 'validation set'], loc='upper right')
    # plt.show()
    plt.savefig('lossHistory.png')

def RGB2YUV(im):
    yuv = []
    for i in im :
        yuv.append(cv2.cvtColor(i, cv2.COLOR_RGB2YUV))
    return yuv

def print_info(info):
    now = datetime.now()
    infotime = now.strftime("%H:%M:%S")
    # can not use f-string due to GPU python version v3.5.2
    print('{}. Time : {}'.format(info,infotime))


text:
from __future__ import print_function
import argparse
from engine.displaced_people_branch.single_image_inference_hra_2class import single_img_HRA_inference
from keras.preprocessing import image
from scipy.misc import imread
from matplotlib import pyplot as plt
from inference.displacenet_single_image_inference_unified import displaceNet_inference

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--img_path", type = str,help = 'path of the input image')
    parser.add_argument("--hra_model_backend_name", type = str,help = 'One of `VGG16`, `VGG19`, `ResNet50` or `VGG16_Places365`')
    parser.add_argument("--emotic_model_backend_name", type=str, help='One of `VGG16`, `VGG19` or `ResNet50`')
    parser.add_argument("--nb_of_conv_layers_to_fine_tune", type=int, help='indicates the number of convolutional layers that were fine-tuned during training')

    args = parser.parse_args()
    return args

args = get_args()

raw_preds, sole_classifier_overlayed_text, top_1_predicted_label = single_img_HRA_inference(img_path=args.img_path,
                                                                                violation_class='dp',
                                                                                model_backend_name=args.hra_model_backend_name,
                                                                                nb_of_conv_layers_to_fine_tune=args.nb_of_conv_layers_to_fine_tune)

img = image.load_img(args.img_path, target_size=(224, 224))
print ('Vanilla CNN prediction: ', raw_preds[0])

emotic_model_a_backend_name = 'VGG16'
emotic_model_b_backend_name = None
emotic_model_c_backend_name = None

final_preds = displaceNet_inference(args.img_path,
                                    emotic_model_a_backend_name=args.emotic_model_backend_name,
                                    emotic_model_b_backend_name=None,
                                    emotic_model_c_backend_name=None,
                                    hra_model_backend_name=args.hra_model_backend_name,
                                    nb_of_fine_tuned_conv_layers=args.nb_of_conv_layers_to_fine_tune,
                                    violation_class='dp')

print('DisplaceNet prediction: ', final_preds)

numpy_img_path = imread(args.img_path)
plt.figure(figsize=(10, 12))
plt.imshow(numpy_img_path)

current_axis = plt.gca()

# configure colours for bounding box and text
violation_bounding_box_colour_rgbvar = (255, 3, 62)
violation_bounding_box_colour_rgbvar2 = [x / 255.0 for x in violation_bounding_box_colour_rgbvar]

no_violation_bounding_box_colour_rgbvar = (34, 139, 34)
no_violation_bounding_box_colour_rgbvar2 = [x / 255.0 for x in no_violation_bounding_box_colour_rgbvar]

abusenet_overlayed_text = str(final_preds[0][0][1]) + ' (' + str(round(final_preds[0][0][2], 2)) + ')'

# print (abusenet_overlayed_text)

abusenet_overlayed_text = 'DisplaceNet: '+abusenet_overlayed_text
sole_classifier_overlayed_text = 'Vanilla CNN: '+ sole_classifier_overlayed_text

if final_preds[0][0][1] == 'displaced_populations':
    current_axis.text(0, -28, abusenet_overlayed_text, size='x-large', color='white',
                      bbox={'facecolor': violation_bounding_box_colour_rgbvar2, 'alpha': 1.0})

    if top_1_predicted_label == 'displaced_populations':
        current_axis.text(0, -7, sole_classifier_overlayed_text, size='x-large', color='white',
                          bbox={'facecolor': violation_bounding_box_colour_rgbvar2, 'alpha': 1.0})
    else:
        current_axis.text(0, -7, sole_classifier_overlayed_text, size='x-large', color='white',
                          bbox={'facecolor': no_violation_bounding_box_colour_rgbvar2, 'alpha': 1.0})

elif final_preds[0][0][1] == 'no_displaced_populations':
    current_axis.text(0, -45, abusenet_overlayed_text, size='x-large', color='white',
                      bbox={'facecolor': no_violation_bounding_box_colour_rgbvar2, 'alpha': 1.0})

    if top_1_predicted_label == 'displaced_populations':
        current_axis.text(0, -7, sole_classifier_overlayed_text, size='x-large', color='white',
                          bbox={'facecolor': violation_bounding_box_colour_rgbvar2, 'alpha': 1.0})
    else:
        current_axis.text(0, -7, sole_classifier_overlayed_text, size='x-large', color='white',
                          bbox={'facecolor': no_violation_bounding_box_colour_rgbvar2, 'alpha': 1.0})

plt.axis('off')
plt.show()

max_stars_repo_name: GKalliatakis/DisplaceNet

max_issues_repo_name: GKalliatakis/DisplaceNet

max_forks_repo_name: GKalliatakis/DisplaceNet

reason: from scipy


text:
import networkx as nx
import os
import pickle as pkl
import time
import pdb
# import utils.data_utils as data_utils
import data_utils
from itertools import combinations

def hyperbolicity_sample(G):
    curr_time = time.time()
    #sp = dict(nx.shortest_path_length(G, weight=None))

    num_samples = 500000
    hyps = []
    # for i in range(num_samples):
    #   node_tuple = np.random.choice(G.nodes(), 4, replace=False)

    comb = combinations(G.nodes(), 4)
    for node_tuple in comb:
        # print(node_tuple)
        s = []
        #d01 = sp[node_tuple[0]][node_tuple[1]] if node_tuple[0] in sp and node_tuple[1] in sp[node_tuple[0]] else 0
        #d23 = sp[node_tuple[2]][node_tuple[3]] if node_tuple[2] in sp and node_tuple[3] in sp[node_tuple[2]] else 0
        #d02 = sp[node_tuple[0]][node_tuple[2]] if node_tuple[0] in sp and node_tuple[2] in sp[node_tuple[0]] else 0
        #d13 = sp[node_tuple[1]][node_tuple[3]] if node_tuple[1] in sp and node_tuple[3] in sp[node_tuple[1]] else 0
        #d03 = sp[node_tuple[0]][node_tuple[3]] if node_tuple[0] in sp and node_tuple[3] in sp[node_tuple[0]] else 0
        #d12 = sp[node_tuple[1]][node_tuple[2]] if node_tuple[1] in sp and node_tuple[2] in sp[node_tuple[1]] else 0
        try:
            d01 = nx.shortest_path_length(G, source=node_tuple[0], target=node_tuple[1], weight=None)
            d23 = nx.shortest_path_length(G, source=node_tuple[2], target=node_tuple[3], weight=None)
            d02 = nx.shortest_path_length(G, source=node_tuple[0], target=node_tuple[2], weight=None)
            d13 = nx.shortest_path_length(G, source=node_tuple[1], target=node_tuple[3], weight=None)
            d03 = nx.shortest_path_length(G, source=node_tuple[0], target=node_tuple[3], weight=None)
            d12 = nx.shortest_path_length(G, source=node_tuple[1], target=node_tuple[2], weight=None)
            s.append(d01 + d23)
            s.append(d02 + d13)
            s.append(d03 + d12)
            s.sort()
            hyps.append((s[-1] - s[-2]) / 2.)
            # print('get')
            if len(hyps) > 1000:
                break
        except Exception as e:
            continue
    #print('Time for hyp: ', time.time() - curr_time)
    # try:
    return max(hyps)

    # except:
    #     return -1

def hyp_ppi():
    data_path = 'data/human_ppi'
    dataset_str = 'human_ppi'
    graph = pkl.load(open(os.path.join(data_path, dataset_str + '.p'), 'rb'))
    print('Computing hyperbolicity', graph.number_of_nodes(), graph.number_of_edges())
    hyp = hyperbolicity_sample(graph)
    print('Hyp: ', hyp)

def hyp_airport():
    data_path = 'data/airport'
    dataset_str = 'airport'
    graph = pkl.load(open(os.path.join(data_path, dataset_str + '.p'), 'rb'))
    print('Computing hyperbolicity', graph.number_of_nodes(), graph.number_of_edges())
    hyp = hyperbolicity_sample(graph)
    print('Hyp: ', hyp)

def hyp_syn():
    data_path = 'data/synthetic_tree_disease_propagation'
    dataset_str = 'synthetic_tree_disease_propagation'
    adj, feat, label = data_utils.load_synthetic_data(dataset_str, False, data_path)
    graph = nx.from_scipy_sparse_matrix(adj)
    print('Computing hyperbolicity', graph.number_of_nodes(), graph.number_of_edges())
    hyp = hyperbolicity_sample(graph)
    print('Hyp: ', hyp)

def hyp_cora():
    # data_path = './data/cora'
    # dataset_str = 'cora'
    # output = data_utils.load_citation_data(dataset_str, False, data_path)
    # adj = output[0]
    # graph = nx.from_scipy_sparse_matrix(adj)
    # print('Computing hyperbolicity', graph.number_of_nodes(), graph.number_of_edges())
    # # hyp = hyperbolicity_sample(graph)
    #
    # pdb.set_trace()
    # print('Hyp: ', hyp)
    import pandas as pd
    data_dir = os.path.expanduser("./data/cora2")
    edgelist = pd.read_csv(os.path.join(data_dir, "cora.cites"), sep='\t', header=None, names=["target", "source"])
    edgelist = edgelist.values
    G = nx.DiGraph()
    for i in range(edgelist.shape[0]):
        G.add_edge(edgelist[i, 0], edgelist[i, 1])

    # largest = max(nx.strongly_connected_components(G), key=len)
    hyp = hyperbolicity_sample(G)
    print(hyp)

def read_triple(file_path, entity2id, relation2id):
    '''
    Read triples and map them into ids.
    '''
    triples = []
    with open(file_path) as fin:
        for line in fin:
            h, r, t = line.strip().split('\t')
            triples.append((entity2id[h], relation2id[r], entity2id[t]))
    return triples

def hyp_general(dataset='wn18'):
    print('dataset: %s' % dataset)
    data_dir = './data/%s' % dataset
    with open(os.path.join(data_dir, 'entities.dict')) as fin:
        entity2id = dict()
        for line in fin:
            eid, entity = line.strip().split('\t')
            entity2id[entity] = int(eid)

    with open(os.path.join(data_dir, 'relations.dict')) as fin:
        relation2id = dict()
        for line in fin:
            rid, relation = line.strip().split('\t')
            relation2id[relation] = int(rid)
    train_triples = read_triple(os.path.join(data_dir, 'train.txt'), entity2id, relation2id)
    valid_triples = read_triple(os.path.join(data_dir, 'valid.txt'), entity2id, relation2id)
    test_triples = read_triple(os.path.join(data_dir, 'test.txt'), entity2id, relation2id)
    all_true_triples = train_triples + valid_triples + test_triples
    print(len(all_true_triples))
    from sage.graphs.hyperbolicity import hyperbolicity
    G = Graph(loops=True)

    for i in all_true_triples:
        G.add_edge(i[0], i[2])

    if G.loops is not None:
        # print(G.loops())
        for i in G.loops():
            G.delete_edge(i[0], i[1])
    if G.connected_components_number()>1:
        G = G.connected_components_subgraphs()[0]
    pdb.set_trace()
    h, _, _ = hyperbolicity(G)
    print('| entire G |  %5d  |  %5d  | %d  |' % (G.order(), G.size(), h))

    for relation in range(0, len(relation2id)):
        try:
            G = Graph(loops=True)
            for i in all_true_triples:
                if i[1] == relation:
                    G.add_edge(i[0], i[2])

            if G.loops is not None:
                # print(G.loops())
                for i in G.loops():
                    G.delete_edge(i[0], i[1])
            if G.connected_components_number() > 1:
                G = G.connected_components_subgraphs()[0]
            if dataset == 'FB15k':
                if G.order() > 500:
                    h, _, _ = hyperbolicity(G)
                    print('|  %d  |  %5d  |  %5d  | %d  |' %(relation, G.order(), G.size(), h))
            else:
                h, _, _ = hyperbolicity(G)
                print('|  %d  |  %5d  |  %5d  | %d  |' % (relation, G.order(), G.size(), h))
        except:
            print('error encountered')
            continue

# entire graph
#     G = nx.DiGraph()
#     # for i in range(len(entity2id)):
#     #     G.add_node(i)
#     for triple in all_true_triples:
#         G.add_edge(triple[0], triple[2])
#     # pdb.set_trace()
#
#     hyp = hyperbolicity_sample(G)
#     print('| entire G |  %5d  |  %5d  | %d  |' % (G.number_of_nodes(), G.number_of_edges(), hyp))
#
#     for relation in range(0, len(relation2id)):
#         # print('------- %d --------' % relation)
#         G = nx.DiGraph()
#         node_set = []
#         for triple in all_true_triples:
#             if triple[1] == relation:
#                 node_set.append(triple[0])
#                 node_set.append(triple[2])
#
#                 G.add_edge(triple[0], triple[2])
#         for i in set(node_set):
#             G.add_node(i)
#         # if G.number_of_nodes() > 500:
#         hyp = hyperbolicity_sample(G)
#         print('|  %d  |  %5d  |  %5d  | %d  |' %(relation, G.number_of_nodes(), G.number_of_edges(), hyp))

if __name__ == '__main__':
    #hyp_airport()
    # hyp_syn()
    #hyp_ppi()
    hyp_general('FB15k')
    # hyp_cora()
    # hyp_circle()

max_stars_repo_name: honglin-chen/KnowledgeGraphEmbedding

max_issues_repo_name: honglin-chen/KnowledgeGraphEmbedding

max_forks_repo_name: honglin-chen/KnowledgeGraphEmbedding

reason: from sage,import networkx


text:
#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Mon Dec  9 15:41:14 2019

@author: kako
"""

#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Mon Dec  2 16:02:09 2019

@author: kako
"""
import numpy as np
import matplotlib.pyplot as plt

folder3000 = '/home/kako/run_dirs/gempic/strongA/weibel_191203/pyt/'
folder = '/home/kako/run_dirs/gempic/strongA/iaw_191213/'

# Reference solution
#ref = np.loadtxt(folder3000+'disgrad_f3_s2_thdiag.dat')
ref = np.loadtxt('/home/kako/run_dirs/gempic/strongA/weibel_191203/ref/strong_weibel_dg_f3_s3_thdiag.dat')
ref2 = np.loadtxt('/home/kako/run_dirs/gempic/strongA/iaw_191213/symplectic_f-1_s2_thdiag.dat')

sdeg_vec = [1, 2, 3, 4 ]
dfem_vec = [1,2,3]#[-1,1,2,3,4,5]
propagators = ['disgrad', 'symplectic']
propagators = ['symplectic']
def check_conservation(data ):
    print('Conservation Gauss:', np.max(np.abs(data[:,11]-data[0,11])))
    print('Conservation energy:', np.max(np.abs(data[:,8]-data[0,8])))
    print('Momentum error', np.max(np.abs(data[:,9])), np.max(np.abs(data[:,10])))
    #plt.semilogy(data[:,0],data[:,3])

#dfem_vec = [3,4]
#sdeg_vec = [3]
#dfem_vec = [1,2,3,4]
#propagators = ['symplectic']
#sdeg_vec = [3]
#dfem_vec = [3]
ran = range(0,40001)
for sdeg in sdeg_vec :
    for propagator in propagators :
        for dfem in dfem_vec :
            prefix = propagator +'_f'+str(dfem)+'_s' + str(sdeg)
            data = np.loadtxt(folder + prefix +'_thdiag.dat')
            plt.semilogy(data[ran,0],data[ran,6])
            print(prefix)
            #check_conservation(data)
       # plt.semilogy(ref[ran,0],ref[ran,3])#,ref2[ran,0],ref2[ran,3])
        plt.semilogy(ref2[ran,0],ref2[ran,6])
        plt.legend(dfem_vec)
        plt.title(str(sdeg)+' '+propagator)
        plt.show()

err = np.zeros([6,4])
ind = 0
for dfem in dfem_vec :
    for propagator in propagators :
        ins = 0
        for sdeg in sdeg_vec :
            prefix = propagator +'_f'+str(dfem)+'_s' + str(sdeg)
            data = np.loadtxt(folder + prefix +'_thdiag.dat')
            plt.semilogy(data[ran,0],data[ran,1])
            print(prefix, np.max(np.abs(data[:,6]-ref2[:,6])))
            err[ind,ins] = np.sqrt(np.sum((data[:,1]-ref2[:,1])**2))#np.max(np.abs(data[:,1]-ref2[:,1]))
            ins = ins+1
            #check_conservation(data)
       # plt.semilogy(ref[ran,0],ref[ran,3])#,ref2[ran,0],ref2[ran,3])
        #plt.semilogy(ref2[ran,0],ref2[ran,6])
        plt.legend(sdeg_vec)
        plt.title(str(dfem)+' '+propagator)
        plt.show()
    ind = ind +1

max_stars_repo_name: utastudents/selalib

max_issues_repo_name: utastudents/selalib

max_forks_repo_name: utastudents/selalib


text:
import itertools
from typing import List, Dict, Tuple, Optional

import numpy as np

from remat.core.dfgraph import DFGraph
from remat.core.schedule import OperatorEvaluation, AllocateRegister, DeallocateRegister, Schedule, SchedulerAuxData
from remat.core.utils.timer import Timer

class ScheduleBuilder:
    def __init__(self, g, verbosity: int = 2):
        self.max_ram = 0
        self.total_cpu = 0
        self.g = g
        self.schedule: Schedule = []
        self.live_registers: Dict[int, int] = {}
        self.next_free_register_id = 0
        self.verbosity = verbosity
        self.ram_timeline: List[int] = []

    def is_op_cached(self, op_id: int):
        return op_id in self.live_registers.keys()

    def allocate_register(self, op_id: int):
        """
        Schedule a register allocation
        :param op_id: ID for operation whose output will be stored in this register,
        :return: the newly allocated register ID
        """
        if op_id in self.live_registers.keys():
            if self.verbosity >= 2:
                print("WARNING! Double allocating output register for op #{}, skipping allocation to reuse reg #{}"
                      .format(op_id, self.live_registers[op_id]))
            return self.live_registers[op_id]
        reg = AllocateRegister(self.next_free_register_id, op_id, self.g.cost_ram[op_id])
        self.live_registers[op_id] = reg.register_id
        self.schedule.append(reg)
        self.next_free_register_id += 1
        self.max_ram = max(self.max_ram, self.current_mem())
        self.ram_timeline.append(self.current_mem())
        return reg.register_id

    def run_operator(self, op_id: int, update_aux_vars: bool):
        debug_str = "Dependency not fulfilled for op #{}, ops in ram now are {} but I need {}".format(
            op_id, set(self.live_registers.keys()), self.g.predecessors(op_id))
        assert all([pred in self.live_registers.keys() for pred in self.g.predecessors(op_id)]), debug_str
        out_reg = self.allocate_register(op_id)
        in_regs = {pred_id: self.live_registers[pred_id] for pred_id in self.g.predecessors(op_id)}
        eval_op = OperatorEvaluation(op_id, in_regs, out_reg, self.g.cost_cpu[op_id],
                                     update_aux_vars=update_aux_vars, is_backwards=op_id not in self.g.vfwd)
        self.schedule.append(eval_op)
        self.total_cpu += self.g.cost_cpu[op_id]
        self.ram_timeline.append(self.current_mem())

    def deallocate_register(self, op_id: int):
        """
        Schedule a register deallocation
        :param op_id: ID for operation whose output will be stored in this register
        """
        if op_id not in self.live_registers.keys():
            print("WARNING! Double free output register for op #{}".format(op_id))
        reg_id = self.live_registers.pop(op_id)
        self.schedule.append(DeallocateRegister(op_id, reg_id))
        self.ram_timeline.append(self.current_mem())

    def current_mem(self):
        return sum(map(self.g.cost_ram.get, self.live_registers.keys()))

def schedule_from_rs(g: DFGraph, r: np.ndarray, s: np.ndarray) -> Tuple[Optional[Schedule], Optional[SchedulerAuxData]]:
    if r is None or s is None:
        return None, None  # infeasible
    T = g.size

    def _used_after(t_, u_, i_):
        """Returns True if v_u is used after v_i in stage t"""
        is_retained_snapshot = t_ < T - 1 and s[t_ + 1, u_] == 1
        is_used_by_successor = not all([r[t_, v] == 0 or v <= i_ for v in g.successors(u_)])
        return is_retained_snapshot or is_used_by_successor

    with Timer('schedule_rs_matrix') as schedule_timer:
        # compute last usage to determine whether to update auxiliary variables
        last_used = {i: max([t for t in range(T) if r[t, i] == 1]) for i in range(T)}
        mem_usage = np.zeros((T, T), dtype=np.int)
        sb = ScheduleBuilder(g, verbosity=1)
        for t in range(T):
            # Free unused checkpoints
            for i in filter(lambda x: sb.is_op_cached(x), range(T)):
                if not _used_after(t, i, i):
                    sb.deallocate_register(i)

            for i in range(T):
                if r[t, i] == 1:
                    sb.run_operator(i, last_used[i] == t)
                mem_usage[t, i] = sb.current_mem() + g.cost_ram_fixed

                # Free memory
                for u in filter(lambda x: sb.is_op_cached(x), itertools.chain(g.predecessors(i), [i])):
                    if not _used_after(t, u, i):
                        sb.deallocate_register(u)
        total_ram = sb.max_ram + g.cost_ram_fixed
        ram_timeline = [mem + g.cost_ram_fixed for mem in sb.ram_timeline]

    return sb.schedule, SchedulerAuxData(R=r, S=s, cpu=sb.total_cpu, peak_ram=total_ram,
                                         activation_ram=sb.max_ram, mem_grid=mem_usage,
                                         mem_timeline=ram_timeline, schedule_time_s=schedule_timer.elapsed)

max_stars_repo_name: uwsampl/dtr-prototype

max_issues_repo_name: merrymercy/dtr-prototype

max_forks_repo_name: uwsampl/dtr-prototype


text:
from collections.abc import Mapping
from dataclasses import fields, is_dataclass
from functools import lru_cache
import logging
from multiprocessing.context import BaseContext
from pathlib import Path
import platform
import subprocess
from typing import (
    Any,
    Callable,
    Dict,
    Final,
    Iterator,
    List,
    Literal,
    NamedTuple,
    Optional,
    Sequence,
    Tuple,
    Type,
    TypeVar,
    Union,
    cast,
    get_args,
    overload,
)
from zipfile import BadZipFile

from PIL import Image
import albumentations as A  # type: ignore
import cv2
import numpy as np
import numpy.typing as npt
from ranzen.misc import gcopy
from ranzen.torch.data import Subset, prop_random_split
import torch
from torch import Tensor
from torch._six import string_classes
from torch.utils.data import ConcatDataset
from torch.utils.data._utils.collate import (
    default_collate_err_msg_format,
    np_str_obj_array_pattern,
)
from torch.utils.data.dataloader import DataLoader, _worker_init_fn_t
from torch.utils.data.sampler import Sampler
from torchvision.datasets.utils import (  # type: ignore
    _detect_file_type,
    download_url,
    extract_archive,
)
from torchvision.transforms import functional as TF  # type: ignore
from typing_extensions import TypeAlias, TypeGuard

from conduit.data.datasets.base import CdtDataset
from conduit.data.structures import (
    BinarySample,
    Dataset,
    LoadedData,
    NamedSample,
    PseudoCdtDataset,
    RawImage,
    SampleBase,
    SizedDataset,
    TernarySample,
    TrainTestSplit,
)

__all__ = [
    "AlbumentationsTform",
    "AudioTform",
    "CdtDataLoader",
    "GdriveFileInfo",
    "ImageLoadingBackend",
    "ImageTform",
    "PillowTform",
    "UrlFileInfo",
    "apply_image_transform",
    "cdt_collate",
    "check_integrity",
    "download_from_gdrive",
    "download_from_url",
    "extract_base_dataset",
    "extract_labels_from_dataset",
    "get_group_ids",
    "img_to_tensor",
    "infer_al_backend",
    "infer_il_backend",
    "infer_sample_cls",
    "is_tensor_list",
    "load_image",
    "make_subset",
    "stratified_split",
]

ImageLoadingBackend: TypeAlias = Literal["opencv", "pillow"]

@overload
def load_image(
    filepath: Union[Path, str], *, backend: Literal["opencv"] = ...
) -> npt.NDArray[np.integer]:
    ...

@overload
def load_image(filepath: Union[Path, str], *, backend: Literal["pillow"] = ...) -> Image.Image:
    ...

def load_image(filepath: Union[Path, str], *, backend: ImageLoadingBackend = "opencv") -> RawImage:
    """Load an image from disk using the requested backend.

    :param: The path of the image-file to be loaded.
    :param backend: Backed to use for loading the image: either 'opencv' or 'pillow'.

    :returns: The loaded image file as a numpy array if 'opencv' was the selected backend
    and a PIL image otherwise.
    """
    if backend == "opencv":
        if isinstance(filepath, Path):
            # cv2 can only read string filepaths
            filepath = str(filepath)
        image = cv2.imread(filepath)  # type: ignore
        if image is None:
            raise OSError(f"Image-file could not be read from location '{filepath}'")
        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # type: ignore
    return Image.open(filepath)

AlbumentationsTform: TypeAlias = Union[A.Compose, A.BasicTransform]
PillowTform: TypeAlias = Callable[[Image.Image], Any]
ImageTform: TypeAlias = Union[AlbumentationsTform, PillowTform]

def infer_il_backend(transform: Optional[ImageTform]) -> ImageLoadingBackend:
    """Infer which image-loading backend to use based on the type of the image-transform.

    :param transform: The image transform from which to infer the image-loading backend.
    If the transform is derived from Albumentations, then 'opencv' will be selected as the
    backend, else 'pillow' will be selected.

    :returns: The backend to load images with based on the supplied image-transform: either
    'opencv' or 'pillow'.
    """
    # Default to openccv is transform is None as numpy arrays are generally
    # more tractable
    if transform is None or isinstance(transform, get_args(AlbumentationsTform)):
        return "opencv"
    return "pillow"

def apply_image_transform(
    image: RawImage, *, transform: Optional[ImageTform]
) -> Union[RawImage, Tensor]:
    image_ = image
    if transform is not None:
        if isinstance(transform, (A.Compose, A.BasicTransform)):
            if isinstance(image, Image.Image):
                image = np.array(image)
            image_ = transform(image=image)["image"]
        else:
            if isinstance(image, np.ndarray):
                image = Image.fromarray(image)
            image_ = transform(image)
    return image_

def img_to_tensor(img: Union[Image.Image, np.ndarray]) -> Tensor:
    if isinstance(img, Image.Image):
        return TF.pil_to_tensor(img)
    return torch.from_numpy(
        np.moveaxis(img / (255.0 if img.dtype == np.uint8 else 1), -1, 0).astype(np.float32)
    )

AudioLoadingBackend: TypeAlias = Literal["sox_io", "soundfile"]

def infer_al_backend() -> AudioLoadingBackend:
    """Infer which audio-loading backend to use based on the operating system."""
    soundfile: Final = "soundfile"
    sox: Final = "sox_io"
    return soundfile if platform.system() == "Windows" else sox

AudioTform: TypeAlias = Callable[[Tensor], Tensor]

def apply_audio_transform(waveform: Tensor, *, transform: Optional[AudioTform]) -> Tensor:
    return waveform if transform is None else transform(waveform)

@overload
def extract_base_dataset(
    dataset: Dataset, *, return_subset_indices: Literal[True] = ...
) -> Tuple[Dataset, Union[Tensor, slice]]:
    ...

@overload
def extract_base_dataset(
    dataset: Dataset, *, return_subset_indices: Literal[False] = ...
) -> Dataset:
    ...

def extract_base_dataset(
    dataset: Dataset, *, return_subset_indices: bool = True
) -> Union[Dataset, Tuple[Dataset, Union[Tensor, slice]]]:
    """Extract the innermost dataset of a nesting of datasets.

    Nested datasets are inferred based on the existence of a 'dataset'
    attribute and the base dataset is extracted by recursive application
    of this rule.

    :param dataset: The dataset from which to extract the base dataset.

    :param return_subset_indices: Whether to return the indices from which
    the overall subset of the dataset was created (works for multiple levels of
    subsetting).

    :returns: The base dataset, which may be the original dataset if one does not
    exist or cannot be determined.
    """

    def _closure(
        dataset: Dataset, rel_indices_ls: Optional[List[List[int]]] = None
    ) -> Union[Dataset, Tuple[Dataset, Union[Tensor, slice]]]:
        if rel_indices_ls is None:
            rel_indices_ls = []
        if hasattr(dataset, "dataset"):
            if isinstance(dataset, Subset):
                rel_indices_ls.append(list(dataset.indices))
            return _closure(dataset.dataset, rel_indices_ls)  # type: ignore
        if return_subset_indices:
            if rel_indices_ls:
                abs_indices = torch.as_tensor(rel_indices_ls.pop(), dtype=torch.long)
                for indices in rel_indices_ls[::-1]:
                    abs_indices = abs_indices[indices]
            else:
                abs_indices = slice(None)
            return dataset, abs_indices
        return dataset

    return _closure(dataset)

@lru_cache(typed=True)
def extract_labels_from_dataset(
    dataset: PseudoCdtDataset,
) -> Tuple[Optional[Tensor], Optional[Tensor]]:
    """Attempt to extract s/y labels from a dataset."""
    base_dataset = dataset

    def _closure(dataset: Dataset) -> Tuple[Optional[Tensor], Optional[Tensor]]:
        base_dataset, indices = extract_base_dataset(dataset=dataset, return_subset_indices=True)
        _s = None
        _y = None
        if getattr(base_dataset, "s", None) is not None:
            _s = base_dataset.s[indices]  # type: ignore
        if getattr(base_dataset, "y", None) is not None:
            _y = base_dataset.y[indices]  # type: ignore

        _s = torch.from_numpy(_s) if isinstance(_s, np.ndarray) else _s
        _y = torch.from_numpy(_y) if isinstance(_y, np.ndarray) else _y

        return _s, _y

    if isinstance(base_dataset, (ConcatDataset)):
        s_all_ls: List[Tensor] = []
        y_all_ls: List[Tensor] = []
        for _dataset in base_dataset.datasets:
            s, y = _closure(_dataset)
            if s is not None:
                s_all_ls.append(s)
            if y is not None:
                s_all_ls.append(y)
        s_all = torch.cat(s_all_ls, dim=0) if s_all_ls else None
        y_all = torch.cat(y_all_ls, dim=0) if y_all_ls else None
    else:
        s_all, y_all = _closure(base_dataset)
    return s_all, y_all

def get_group_ids(dataset: Dataset) -> Tensor:
    s_all, y_all = extract_labels_from_dataset(dataset)
    # group_ids: Optional[Tensor] = None
    if s_all is None:
        if y_all is None:
            raise ValueError(
                "Unable to compute group ids for dataset because no labels could be extracted."
            )
        group_ids = y_all
    elif y_all is None:
        group_ids = s_all
    else:
        group_ids = (y_all * len(s_all.unique()) + s_all).squeeze()
    return group_ids.long()

def compute_instance_weights(dataset: Dataset, *, upweight: bool = False) -> Tensor:
    group_ids = get_group_ids(dataset)
    _, inv_indexes, counts = group_ids.unique(return_inverse=True, return_counts=True)
    # Upweight samples according to the cardinality of their intersectional group
    if upweight:
        group_weights = len(group_ids) / counts
    # Downweight samples according to the cardinality of their intersectional group
    # - this approach should be preferred due to being more numerically stable
    # (very small counts can lead to very large weighted loss values when upweighting)
    else:
        group_weights = 1 - (counts / len(group_ids))
    return group_weights[inv_indexes]

PCD = TypeVar("PCD", bound=PseudoCdtDataset)

def make_subset(
    dataset: Union[PCD, Subset[PCD]],
    *,
    indices: Optional[Union[List[int], npt.NDArray[np.uint64], Tensor, slice]],
    deep: bool = False,
) -> PCD:
    """Create a subset of the dataset from the given indices.

    :param indices: The sample-indices from which to create the subset.
    In the case of being a numpy array or tensor, said array or tensor
    must be 0- or 1-dimensional.

    :param deep: Whether to create a copy of the underlying dataset as
    a basis for the subset. If False then the data of the subset will be
    a view of original dataset's data.

    :returns: A subset of the dataset from the given indices.
    """
    if isinstance(indices, (np.ndarray, Tensor)):
        if indices.ndim > 1:
            raise ValueError("If 'indices' is an array it must be a 0- or 1-dimensional.")
        indices = cast(List[int], indices.tolist())

    current_indices = None
    if isinstance(dataset, Subset):
        base_dataset, current_indices = extract_base_dataset(dataset, return_subset_indices=True)
        if not isinstance(base_dataset, CdtDataset):
            raise TypeError(
                f"Subsets can only be created from {CdtDataset.__name__} instances or PyTorch "
                "Subsets of them."
            )
        base_dataset = cast(PCD, base_dataset)

        if isinstance(current_indices, Tensor):
            current_indices = current_indices.tolist()
    else:
        base_dataset = dataset
    subset = gcopy(base_dataset, deep=deep)

    def _subset_from_indices(_dataset: PCD, _indices: Union[List[int], slice]) -> PCD:
        _dataset.x = _dataset.x[_indices]
        if _dataset.y is not None:
            _dataset.y = _dataset.y[_indices]
        if _dataset.s is not None:
            _dataset.s = _dataset.s[_indices]
        return _dataset

    if current_indices is not None:
        subset = _subset_from_indices(_dataset=subset, _indices=current_indices)
    if indices is not None:
        subset = _subset_from_indices(_dataset=subset, _indices=indices)

    return subset

def infer_sample_cls(
    sample: Union[List[LoadedData], Tuple[LoadedData, ...], Dict[str, LoadedData], LoadedData]
) -> Type[NamedSample]:
    """ "Attempt to infer the appropriate sample class based on the length of the input."""
    if not isinstance(sample, (list, tuple, dict)) or (len(sample) == 1):
        return NamedSample
    elif len(sample) == 2:
        return BinarySample
    elif len(sample) == 3:
        return TernarySample
    else:
        raise ValueError("Only items with 3 or fewer elements can be cast to 'Sample' instances.")

class cdt_collate:
    def __init__(
        self, cast_to_sample: bool = True, *, converter: Optional[Union[Type[Any], Callable]] = None
    ) -> None:
        self.cast_to_sample = cast_to_sample
        self.converter = converter

    def _collate(self, batch: Any) -> Any:
        elem = batch[0]
        elem_type = type(elem)
        if isinstance(elem, Tensor):
            out = None
            if torch.utils.data.get_worker_info() is not None:  # type: ignore
                # If we're in a background process, concatenate directly into a
                # shared memory tensor to avoid an extra copy
                numel = sum(x.numel() for x in batch)
                storage = elem.storage()._new_shared(numel)
                out = elem.new(storage).resize_(len(batch), *list(elem.size()))
            ndims = elem.dim()
            if (ndims > 0) and ((ndims % 2) == 0):
                return torch.cat(batch, dim=0, out=out)  # type: ignore
            return torch.stack(batch, dim=0, out=out)  # type: ignore
        elif (
            elem_type.__module__ == "numpy"
            and elem_type.__name__ != "str_"
            and elem_type.__name__ != "string_"
        ):
            elem = batch[0]
            if elem_type.__name__ == "ndarray":
                # array of string classes and object
                if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
                    raise TypeError(default_collate_err_msg_format.format(elem.dtype))
                return self._collate([torch.as_tensor(b) for b in batch])
        elif isinstance(elem, float):
            return torch.tensor(batch, dtype=torch.float64)
        elif isinstance(elem, int):
            return torch.tensor(batch)
        elif isinstance(elem, string_classes):
            return batch
        elif isinstance(elem, Mapping):
            return {key: self._collate([d[key] for d in batch]) for key in elem}
        elif isinstance(elem, tuple) and hasattr(elem, "_fields"):  # namedtuple
            return elem_type(**(self._collate(samples) for samples in zip(*batch)))
        elif is_dataclass(elem):  # dataclass
            return elem_type(
                **{
                    field.name: self._collate([getattr(d, field.name) for d in batch])
                    for field in fields(elem)
                }
            )
        elif isinstance(elem, (tuple, list)):
            transposed = zip(*batch)
            return [self._collate(samples) for samples in transposed]
        raise TypeError(default_collate_err_msg_format.format(elem_type))

    def __call__(self, batch: Any) -> Any:
        collated_batch = self._collate(batch=batch)
        if self.converter is not None:
            collated_batch = self.converter(collated_batch)
        if self.cast_to_sample and (not isinstance(collated_batch, SampleBase)):
            if isinstance(collated_batch, Tensor):
                collated_batch = NamedSample(x=collated_batch)
            elif isinstance(collated_batch, (tuple, list, dict)):
                sample_cls = infer_sample_cls(collated_batch)
                if not isinstance(collated_batch, dict):
                    collated_batch = dict(zip(["x", "y", "s"], collated_batch))
                collated_batch = sample_cls(**collated_batch)
            else:
                raise ValueError(
                    f"batch of type '{type(collated_batch)}' could not be automatically cast to a "
                    "'Sample' instance. Batch must be of type 'dict', 'tuple', or 'list'."
                )
        return collated_batch

I = TypeVar("I", bound=NamedSample)

class CdtDataLoader(DataLoader[I]):
    def __init__(
        self,
        dataset: SizedDataset[I],
        *,
        batch_size: Optional[int] = 1,
        shuffle: bool = False,
        sampler: Optional[Sampler[int]] = None,
        batch_sampler: Optional[Sampler[Sequence[int]]] = None,
        num_workers: int = 0,
        pin_memory: bool = False,
        drop_last: bool = False,
        timeout: float = 0,
        worker_init_fn: Optional[_worker_init_fn_t] = None,
        multiprocessing_context: Optional[Union[BaseContext, str]] = None,
        generator: Optional[torch.Generator] = None,
        prefetch_factor: int = 2,
        persistent_workers: bool = False,
        cast_to_sample: bool = True,
        converter: Optional[Union[Type[Any], Callable]] = None,
    ) -> None:
        super().__init__(
            dataset,  # type: ignore
            batch_size=batch_size,
            shuffle=shuffle,
            sampler=sampler,
            batch_sampler=batch_sampler,
            num_workers=num_workers,
            collate_fn=cdt_collate(cast_to_sample=cast_to_sample, converter=converter),
            pin_memory=pin_memory,
            drop_last=drop_last,
            timeout=timeout,
            worker_init_fn=worker_init_fn,
            multiprocessing_context=multiprocessing_context,
            generator=generator,
            prefetch_factor=prefetch_factor,
            persistent_workers=persistent_workers,
        )

    def __iter__(self) -> Iterator[I]:
        return super().__iter__()

def check_integrity(*, filepath: Path, md5: Optional[str]) -> None:
    from torchvision.datasets.utils import check_integrity  # type: ignore

    ext = filepath.suffix
    if ext not in [".zip", ".7z"] and check_integrity(fpath=str(filepath), md5=md5):
        raise RuntimeError('Dataset corrupted; try deleting it and redownloading it.')

class UrlFileInfo(NamedTuple):
    name: str
    url: str
    md5: Optional[str] = None

def download_from_url(
    *,
    file_info: Union[UrlFileInfo, List[UrlFileInfo]],
    root: Union[Path, str],
    logger: Optional[logging.Logger] = None,
    remove_finished: bool = True,
) -> None:

    logger = logging.getLogger(__name__) if logger is None else logger
    file_info_ls = file_info if isinstance(file_info, list) else [file_info]
    if not isinstance(root, Path):
        root = Path(root).expanduser()
    # Create the specified root directory if it doesn't already exist
    root.mkdir(parents=True, exist_ok=True)

    for info in file_info_ls:
        filepath = root / info.name

        filepath_str = str(filepath)
        suffix = _detect_file_type(filepath_str)[0]
        extracted_filepath = Path(filepath_str.split(suffix)[0])

        if extracted_filepath.exists():
            logger.info(f"File '{info.name}' already downloaded and extracted.")
        else:
            if filepath.exists():
                logger.info(f"File '{info.name}' already downloaded.")
            else:
                logger.info(f"Downloading file '{info.name}' from address '{info.url}'.")
                download_url(url=info.url, filename=info.name, root=str(root), md5=info.md5)

            logger.info(f"Extracting '{filepath.resolve()}' to '{root.resolve()}'")
            try:
                extract_archive(
                    from_path=str(filepath),
                    to_path=str(extracted_filepath),
                    remove_finished=remove_finished,
                )
            # Fall back on using jar to unzip the archive
            except BadZipFile:
                try:
                    subprocess.run(["jar", "-xvf", str(filepath)], check=True, cwd=root)
                except subprocess.CalledProcessError:
                    logger.info(
                        "Attempted to fall back on using Java to extract malformed .zip file; "
                        "however, there was a problem. Try redownloading the zip file or "
                        "checking that Java has been properly added to your system variables."
                    )

class GdriveFileInfo(NamedTuple):
    name: str
    id: str
    md5: Optional[str] = None

def download_from_gdrive(
    *,
    file_info: Union[GdriveFileInfo, List[GdriveFileInfo]],
    root: Union[Path, str],
    logger: Optional[logging.Logger] = None,
) -> None:
    """Attempt to download data if files cannot be found in the root directory."""

    logger = logging.getLogger(__name__) if logger is None else logger

    file_info_ls = file_info if isinstance(file_info, list) else [file_info]
    if not isinstance(root, Path):
        root = Path(root).expanduser()
    # Create the specified root directory if it doesn't already exist
    root.mkdir(parents=True, exist_ok=True)

    for info in file_info_ls:
        filepath = root / info.name
        if filepath.exists():
            logger.info(f"File '{info.name}' already downloaded.")
        else:
            import gdown  # type: ignore

            logger.info(f"Downloading file '{info.name}' from Google Drive.")
            gdown.cached_download(
                url=f"https://drive.google.com/uc?id={info.id}",
                path=str(filepath),
                quiet=False,
                md5=info.md5,
            )
        if filepath.suffix == ".zip":
            if filepath.with_suffix("").exists():
                logger.info(f"File '{info.name}' already unzipped.")
            else:
                check_integrity(filepath=filepath, md5=info.md5)
                # ------------------------------ Unzip the data ------------------------------
                import zipfile

                logger.info(f"Unzipping '{filepath.resolve()}'; this could take a while.")
                with zipfile.ZipFile(filepath, "r") as fhandle:
                    fhandle.extractall(str(root))

def random_split(
    dataset: PCD,
    props: Union[Sequence[float], float],
    deep: bool = False,
    seed: Optional[int] = None,
) -> List[PCD]:
    """Randomly split the dataset into subsets according to the given proportions.

    :param props: The fractional size of each subset into which to randomly split the data.
    Elements must be non-negative and sum to 1 or less; if less then the size of the final
    split will be computed by complement.

    :param deep: Whether to create a copy of the underlying dataset as
    a basis for the random subsets. If False then the data of the subsets will be
    views of original dataset's data.

    :param seed: PRNG seed to use for sampling.

    :returns: Random subsets of the data of the requested proportions.
    """
    split_indices = prop_random_split(dataset=dataset, props=props, as_indices=True, seed=seed)
    splits = [make_subset(dataset, indices=indices, deep=deep) for indices in split_indices]
    return splits

def stratified_split(
    dataset: PCD,
    *,
    default_train_prop: float,
    train_props: Optional[Dict[int, Union[Dict[int, float], float]]] = None,
    seed: Optional[int] = None,
) -> TrainTestSplit[PCD]:
    """Splits the data into train/test sets conditional on super- and sub-class labels.

    :param default_train_prop: Proportion of samples for a given to sample for
    the training set for those y-s combinations not specified in ``train_props``.

    :param train_props: Proportion of each superclass-subclass combination to sample for
    the training set. Keys correspond to the superclass while values can be either a float,
    in which case the proportion is applied to the superclass as a whole, or a dict, in which
    case the sampling is applied only to the subclasses of the superclass given by the keys.
    If ``None`` then the function reduces to a simple random split of the data.

    :param seed: PRNG seed to use for sampling.

    :returns: Train-test split.
    """
    if dataset.y is None:
        raise TypeError(
            f"Dataset of type {dataset.__class__.__name__} has no superclass labels to use "
            "for stratification."
        )
    train_props = {} if train_props is None else train_props
    # Initialise the random-number generator
    generator = torch.default_generator if seed is None else torch.Generator().manual_seed(seed)

    group_ids = get_group_ids(dataset)
    y_unique = dataset.y.unique()
    groups, id_counts = group_ids.unique(return_counts=True)
    card_s = None if dataset.s is None else len(dataset.s.unique())
    ncols = 1 if card_s is None else card_s
    group_train_props = dict.fromkeys(groups.tolist(), default_train_prop)

    if train_props is not None:
        for superclass, value in train_props.items():
            # Apply the same splitting proportion to the entire superclass
            if superclass not in y_unique:
                raise ValueError(
                    f"No samples belonging to superclass 'y={superclass}' exist in the dataset of "
                    f"type {dataset.__class__.__name__}."
                )
            if isinstance(value, float):
                if not 0 <= value <= 1:
                    raise ValueError(
                        "All splitting proportions speicfied in 'train_props' must be in the "
                        "range [0, 1]."
                    )
                if card_s is None:
                    group_train_props[superclass] = value
                else:
                    superclass_props = dict.fromkeys(
                        range(superclass * card_s, (superclass + 1) * card_s),
                        value,
                    )
                    group_train_props.update(superclass_props)
            # Specifying proportions at the superclass/subclass level, rather than superclass-wide
            else:
                for subclass, train_prop in value.items():
                    if not 0 <= train_prop <= 1:
                        raise ValueError(
                            "All splitting proportions specified in 'train_props' must be in the "
                            "range [0, 1]."
                        )
                    group_id = superclass * ncols + subclass
                    if group_id not in groups:
                        raise ValueError(
                            f"No samples belonging to the subset '(y={superclass}', s={subclass})' "
                            f"exist in the dataset of type {dataset.__class__.__name__}."
                        )
                    group_train_props[group_id] = train_prop

    # Shuffle the samples before sampling
    perm_inds = torch.randperm(len(group_ids), generator=generator)
    group_ids_perm = group_ids[perm_inds]

    sort_inds = group_ids_perm.sort(dim=0, stable=True).indices
    thresholds = cast(
        Tensor, (torch.as_tensor(tuple(group_train_props.values())) * id_counts).round().long()
    )
    thresholds = torch.stack([thresholds, id_counts], dim=-1)
    thresholds[1:] += id_counts.cumsum(0)[:-1].unsqueeze(-1)

    train_test_inds = sort_inds.tensor_split(thresholds.flatten()[:-1], dim=0)
    train_inds = perm_inds[torch.cat(train_test_inds[0::2])]
    test_inds = perm_inds[torch.cat(train_test_inds[1::2])]

    train_data = make_subset(dataset=dataset, indices=train_inds)
    test_data = make_subset(dataset=dataset, indices=test_inds)

    return TrainTestSplit(train=train_data, test=test_data)

def is_tensor_list(ls: List[Any]) -> TypeGuard[List[Tensor]]:
    return isinstance(ls[0], Tensor)

max_stars_repo_name: wearepal/conduit

max_issues_repo_name: wearepal/conduit

max_forks_repo_name: wearepal/conduit


text:
from os import read
import streamlit as st
import pandas as pd
import numpy as np
import sys
sys.path.append("./")
# 自作のやつら
from pytorch_utils import OnnxPredictor
from goo_lab.goo_apis import Goo

uploaded_file = st.file_uploader("ファイルアップロード", type='csv')

def load_model(model_path:str = "./onnx_model/epoch=9-valid_loss=0.1356-valid_acc=0.9745_quant.onnx"):
    predictor = OnnxPredictor(model_path=model_path, device="cpu")
    return predictor

model = load_model()

if uploaded_file is not None:
    df = pd.read_csv(uploaded_file)
    l = []
    for text in df["tweet"]:
        # st.write(text)
        # st.write(np.argmax(model.predict(text)[0][0]))
        score = model.predict(text)[0][0][1]
        l.append(score)
        # st.write(score)
    df["score"] = l
    st.write(df.query("score > 0.9")[["username", "tweet", "score"]])
    st.write(df.query("score <= 0.9")[["username", "tweet", "score"]])

    app_id = "e8be92a5e7fbf6a4b60bb8ff34cbdbf551e65a626b32090fe095864a7f2565e3"
    st.write(df["tweet"][0])
    st.write(df["tweet"][1])
    st.write(Goo(app_id=app_id, request_id="record001").textpair(text1=df["tweet"][0], text2=df["tweet"][1]).text)