text:
import unittest
import shutil
from pathlib import Path

import numpy as np

import pytest

from spikeinterface import extract_waveforms, WaveformExtractor
from spikeinterface.extractors import toy_example

from spikeinterface.toolkit.postprocessing import WaveformPrincipalComponent, compute_principal_components

def setup_module():
    for folder in ('toy_rec_1seg', 'toy_sorting_1seg', 'toy_waveforms_1seg',
                   'toy_rec_2seg', 'toy_sorting_2seg', 'toy_waveforms_2seg'):
        if Path(folder).is_dir():
            shutil.rmtree(folder)

    recording, sorting = toy_example(num_segments=2, num_units=10)
    recording = recording.save(folder='toy_rec_2seg')
    sorting = sorting.save(folder='toy_sorting_2seg')
    we = extract_waveforms(recording, sorting, 'toy_waveforms_2seg',
                           ms_before=3., ms_after=4., max_spikes_per_unit=500,
                           n_jobs=1, chunk_size=30000)

    recording, sorting = toy_example(num_segments=1, num_units=10, num_channels=12)
    recording = recording.save(folder='toy_rec_1seg')
    sorting = sorting.save(folder='toy_sorting_1seg')
    we = extract_waveforms(recording, sorting, 'toy_waveforms_1seg',
                           ms_before=3., ms_after=4., max_spikes_per_unit=500,
                           n_jobs=1, chunk_size=30000)

def test_WaveformPrincipalComponent():
    we = WaveformExtractor.load_from_folder('toy_waveforms_2seg')
    unit_ids = we.sorting.unit_ids
    num_channels = we.recording.get_num_channels()
    pc = WaveformPrincipalComponent(we)

    for mode in ('by_channel_local', 'by_channel_global'):
        pc.set_params(n_components=5, mode=mode)
        print(pc)
        pc.run()
        for i, unit_id in enumerate(unit_ids):
            comp = pc.get_components(unit_id)
            # print(comp.shape)
            assert comp.shape[1:] == (5, 4)

        # import matplotlib.pyplot as plt
        # cmap = plt.get_cmap('jet', len(unit_ids))
        # fig, axs = plt.subplots(ncols=num_channels)
        # for i, unit_id in enumerate(unit_ids):
        # comp = pca.get_components(unit_id)
        # print(comp.shape)
        # for chan_ind in range(num_channels):
        # ax = axs[chan_ind]
        # ax.scatter(comp[:, 0, chan_ind], comp[:, 1, chan_ind], color=cmap(i))
        # plt.show()

    for mode in ('concatenated',):
        pc.set_params(n_components=5, mode=mode)
        print(pc)
        pc.run()
        for i, unit_id in enumerate(unit_ids):
            comp = pc.get_components(unit_id)
            assert comp.shape[1] == 5
            # print(comp.shape)

    all_labels, all_components = pc.get_all_components()

    # import matplotlib.pyplot as plt
    # cmap = plt.get_cmap('jet', len(unit_ids))
    # fig, ax = plt.subplots()
    # for i, unit_id in enumerate(unit_ids):
    # comp = pca.get_components(unit_id)
    # print(comp.shape)
    # ax.scatter(comp[:, 0], comp[:, 1], color=cmap(i))
    # plt.show()

def test_compute_principal_components_for_all_spikes():
    we = WaveformExtractor.load_from_folder('toy_waveforms_1seg')
    pc = compute_principal_components(we, load_if_exists=True)
    print(pc)

    pc_file1 = 'all_pc1.npy'
    pc.run_for_all_spikes(pc_file1, max_channels_per_template=7, chunk_size=10000, n_jobs=1)
    all_pc1 = np.load(pc_file1)

    pc_file2 = 'all_pc2.npy'
    pc.run_for_all_spikes(pc_file2, max_channels_per_template=7, chunk_size=10000, n_jobs=2)
    all_pc2 = np.load(pc_file2)

    assert np.array_equal(all_pc1, all_pc2)

if __name__ == '__main__':
    setup_module()
    test_compute_principal_components_for_all_spikes()

max_stars_repo_name: marcbue/spikeinterface

max_issues_repo_name: marcbue/spikeinterface

max_forks_repo_name: marcbue/spikeinterface


text:
#dependencies
#Anaconda 3.6
#Pytorch
#Tensorboardx
#OpenCV python
import torch
import torch.nn as nn
import os
from torch.autograd import Variable
from skimage import io, transform
import numpy as np
import cv2
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import re
from tensorboardX import SummaryWriter
import time
from Neural_Network_Class import conv_deconv #Class where the network is defined

writer = SummaryWriter()

class ImageDataset(Dataset): #Defining the class to load datasets

    def __init__(self,input_dir,output_dir,train=True,transform=None):
        self.input_dir=input_dir
        self.output_dir=output_dir
        self.transform=transform
        self.train=train

    def __len__ (self):
        if self.train:
            return len(os.listdir(self.input_dir))-50 #I have kept size of testing data to be 50
        else:
            return 50

    def __getitem__(self,idx):
        if self.train:
            idx=idx+1
        else:
            idx=idx+1+len(os.listdir(self.input_dir))-50
        input_image=io.imread(self.input_dir+"/"+str(idx)+".jpg").transpose((2, 0, 1)) #The convolution function in pytorch expects data in
        output_image=io.imread(self.output_dir+"/"+str(idx)+".jpg").transpose((2, 0, 1))#format (N,C,H,W) N is batch size , C are channels
                                                                                # H is height and W is width. here we convert image from
                                                                                #(H,W,C) to (C,H,W)

        sample = {'input_image': input_image, 'output_image': output_image}

        if self.transform:
            sample= self.transform(sample)

        return sample

train_dataset=ImageDataset(input_dir="input",output_dir="output") #Training Dataset
test_dataset=ImageDataset(input_dir="input",output_dir="output",train=False) #Testing Dataset
batch_size = 10 #mini-batch size
n_iters = 10000 #total iterations
num_epochs = n_iters / (len(train_dataset) / batch_size)
num_epochs = int(num_epochs)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)

model=conv_deconv() # Neural network model object
if torch.cuda.is_available(): #use gpu if available
	model.cuda()               # move model to gpu

iter=0
iter_new=0
check=os.listdir("checkpoints") #checking if checkpoints exist to resume training
if len(check):
    check.sort(key=lambda x:int((x.split('_')[2]).split('.')[0]))
    model=torch.load("checkpoints/"+check[-1])
    iter=int(re.findall(r'\d+',check[-1])[0])
    iter_new=iter
    print("Resuming from iteration " + str(iter))
    #os.system('python visualise.py')

criterion=nn.MSELoss()  #Loss Class

learning_rate = 0.001
optimizer = torch.optim.Adam(model.parameters(),lr=learning_rate) #optimizer class
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)# this will decrease the learning rate by factor of 0.1
                                                                              # https://discuss.pytorch.org/t/can-t-import-torch-optim-lr-scheduler/5138/6
beg=time.time() #time at the beginning of training
print("Training Started!")
for epoch in range(num_epochs):
    print("\nEPOCH " +str(epoch+1)+" of "+str(num_epochs)+"\n")
    for i,datapoint in enumerate(train_loader):
        datapoint['input_image']=datapoint['input_image'].type(torch.FloatTensor) #typecasting to FloatTensor as it is compatible with CUDA
        datapoint['output_image']=datapoint['output_image'].type(torch.FloatTensor)
        if torch.cuda.is_available(): #move to gpu if available
                input_image = Variable(datapoint['input_image'].cuda()) #Converting a Torch Tensor to Autograd Variable
                output_image = Variable(datapoint['output_image'].cuda())
        else:
                input_image = Variable(datapoint['input_image'])
                output_image = Variable(datapoint['output_image'])

        optimizer.zero_grad()  #https://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903/3
        outputs = model(input_image)
        loss = criterion(outputs, output_image)
        loss.backward() #Backprop
        optimizer.step()    #Weight update
        writer.add_scalar('Training Loss',loss.data[0]/10, iter)
        iter=iter+1
        if iter % 10 == 0 or iter==1:
            # Calculate Accuracy
            test_loss = 0
            total = 0
            # Iterate through test dataset
            for j,datapoint_1 in enumerate(test_loader): #for testing
                datapoint_1['input_image']=datapoint_1['input_image'].type(torch.FloatTensor)
                datapoint_1['output_image']=datapoint_1['output_image'].type(torch.FloatTensor)

                if torch.cuda.is_available():
                    input_image_1 = Variable(datapoint_1['input_image'].cuda())
                    output_image_1 = Variable(datapoint_1['output_image'].cuda())
                else:
                    input_image_1 = Variable(datapoint_1['input_image'])
                    output_image_1 = Variable(datapoint_1['output_image'])

                # Forward pass only to get logits/output
                outputs = model(input_image_1)
                test_loss += criterion(outputs, output_image_1).data[0]
                total+=datapoint_1['output_image'].size(0)
            test_loss=test_loss/total   #sum of test loss for all test cases/total cases
            writer.add_scalar('Test Loss',test_loss, iter)
            # Print Loss
            time_since_beg=(time.time()-beg)/60
            print('Iteration: {}. Loss: {}. Test Loss: {}. Time(mins) {}'.format(iter, loss.data[0]/10, test_loss,time_since_beg))
        if iter % 500 ==0:
            torch.save(model,'checkpoints/model_iter_'+str(iter)+'.pt')
            print("model saved at iteration : "+str(iter))
            writer.export_scalars_to_json("graphs/all_scalars_"+str(iter_new)+".json") #saving loss vs iteration data to be used by visualise.py
    scheduler.step()
writer.close()

#decrease learning rate

ext: py

reason: import numpy


text:
import numpy as np
import os
import warnings
from sklearn import decomposition
import pandas as pd
from lcc.utils.helpers import check_depth
from matplotlib import pyplot as plt

def plotProbabSpace(star_filter, plot_ranges=None, opt="show",
                    path=".", file_name="params_space.png", N=400,
                    title="Params space", x_lab="", y_lab="",
                    searched_coords=None, contamination_coords=None, OVERLAY=0.6):
    """
    Plot params space

    Parameters
    ----------
    star_filter : StarsFilter object
        Trained stars filter object

    plot_ranges : tuple, list
        List of ranges. For example: [(1,10), (20,50)] - for 2D plot

    opt : str
        Option whether save/show/return

    title : str
        Title of the plot

    path : str
        Path to the output file location

    file_name : str
        Name of the file

    OVERLAY : float
        Percentage overlay of borders despite of data ranges

    Returns
    -------
        None
    """
    contamination_coords = contamination_coords if contamination_coords is not None else []
    searched_coords = searched_coords if searched_coords is not None else []

    if not len(searched_coords) and not len(contamination_coords) and hasattr(star_filter,
                                                                              "searched_coords") and hasattr(
            star_filter, "others_coords"):
        searched_coords = star_filter.searched_coords.values
        contamination_coords = star_filter.others_coords.values

    dim = len(star_filter.searched_coords.columns)

    if isinstance(searched_coords, pd.DataFrame) and isinstance(contamination_coords, pd.DataFrame):
        searched_coords = searched_coords.values.tolist()
        contamination_coords = contamination_coords.values.tolist()

    if not plot_ranges:
        plot_ranges = []
        trained_coo = np.array(
            star_filter.searched_coords.values.tolist() + star_filter.others_coords.values.tolist()).T
        for i in range(dim):
            rang = [np.min(trained_coo[i]), np.max(trained_coo[i])]
            overl = abs(rang[0] - rang[1]) * OVERLAY
            plot_ranges.append([rang[0] - overl, rang[1] + overl])

    if dim == 1:
        if not x_lab and not y_lab:
            x_lab = star_filter.descriptors[0].LABEL
            y_lab = "Probability"
        plt_data = plot1DProbabSpace(
            star_filter, plot_ranges, N,
            searched_coords=searched_coords,
            contaminatiom_coords=contamination_coords)
    elif dim == 2:
        if not x_lab and not y_lab:
            if len(star_filter.descriptors) == 2:
                x_lab = star_filter.descriptors[0].LABEL
                y_lab = star_filter.descriptors[1].LABEL
            else:
                labels = []
                for desc in star_filter.descriptors:
                    if hasattr(desc.LABEL, "__iter__"):
                        labels += desc.LABEL
                    else:
                        labels.append(desc.LABEL)
                if len(labels) == 2:
                    x_lab = labels[0]
                    y_lab = labels[1]
                else:
                    x_lab = ", ".join(labels)
                    y_lab = ""
        plt_data = plot2DProbabSpace(star_filter, plot_ranges, N,
                                     searched_coords=searched_coords,
                                     contaminatiom_coords=contamination_coords)

    else:
        plt_data = plotNDProbabSpace(star_filter, plot_ranges, N,
                                  searched_coords=searched_coords,
                                  contaminatiom_coords=contamination_coords
                                  )
        # return np.array([[], [], []])

    plt.xlabel(str(x_lab))
    plt.ylabel(str(y_lab))
    plt.title(str(title))

    if opt == "show":
        plt.show()
    elif opt == "save":
        plt.savefig(os.path.join(path, file_name))
    elif opt == "return":
        return plt_data

def plotNDProbabSpace(star_filter, plot_ranges, N, searched_coords=[],
                      contaminatiom_coords=[]):
    """
    Plot N dim probability space

    Parameters
    ----------
    star_filter : StarsFilter object
        Trained stars filter

    plot_ranges : iterable
        Ranges (max/min) for all axis

    N : int
        Number of points per axis

    searched_coords : list, iterable
        List of coordinates of searched objects

    contaminatiom_coords : list, iterable
        List of coordinates of contamination objects

    Returns
    -------
    tuple
        x, y, Z
    """
    OVERLAY = 0.4
    ns = len(searched_coords)
    if not isinstance(searched_coords, list):
        try:
            searched_coords = searched_coords.tolist()
            contaminatiom_coords = contaminatiom_coords.tolist()
        except:
            pass

    coords = searched_coords + contaminatiom_coords

    if not coords:
        try:
            coords = star_filter.searched_coords.values.tolist() + star_filter.others_coords.values.tolist()

        except:
            try:
                coords = star_filter.searched_coords.tolist() + star_filter.others_coords.tolist()
            except:
                coords = star_filter.searched_coords + star_filter.others_coords

    pca = decomposition.PCA(n_components=2)
    pca.fit(coords)
    red_coords = pca.transform(coords)

    xmax, ymax = np.max(red_coords, axis=0).tolist()
    xmin, ymin = np.min(red_coords, axis=0).tolist()

    xwidth = xmax - xmin
    ywidth = ymax - ymin

    x = np.linspace(xmin - xwidth*OVERLAY, xmax + xwidth*OVERLAY, N)
    y = np.linspace(ymin - ywidth*OVERLAY, ymax + ywidth*OVERLAY, N)
    X, Y = np.meshgrid(x, y)

    to_transf = np.c_[X.ravel(), Y.ravel()]
    back_transf_data = pca.inverse_transform(to_transf)

    z = np.array(star_filter.evaluateCoordinates(back_transf_data))
    Z = z.reshape(X.shape)

    plt.xlim(xmin, xmax)
    plt.ylim(ymin, ymax)

    plt.pcolor(X, Y, Z)
    plt.colorbar()

    if coords:
        s = np.array(red_coords[:ns]).T
        c = np.array(red_coords[ns:]).T
        plt.plot(s[0], s[1], "m*", label="Searched objects", markersize=17)
        plt.plot(
            c[0], c[1], "k*", label="Contamination objects", markersize=17)
        plt.legend()

    return x, y, Z, pca

def plot2DProbabSpace(star_filter, plot_ranges, N, searched_coords=[],
                      contaminatiom_coords=[]):
    """
    Plot probability space

    Parameters
    ----------
    star_filter : StarsFilter object
        Trained stars filter

    plot_ranges : iterable
        Ranges (max/min) for all axis

    N : int
        Number of points per axis

    searched_coords : list, iterable
        List of coordinates of searched objects

    contaminatiom_coords : list, iterable
        List of coordinates of contamination objects

    Returns
    -------
    tuple
        x, y, Z
    """
    if check_depth(plot_ranges, 1, ifnotraise=False):
        plot_ranges = [plot_ranges, plot_ranges]
    x = np.linspace(plot_ranges[0][0], plot_ranges[0][1], N)
    y = np.linspace(plot_ranges[1][0], plot_ranges[1][1], N)
    X, Y = np.meshgrid(x, y)

    z = np.array(star_filter.evaluateCoordinates(np.c_[X.ravel(), Y.ravel()]))
    Z = z.reshape(X.shape)

    plt.xlim(plot_ranges[0][0], plot_ranges[0][1])
    plt.ylim(plot_ranges[1][0], plot_ranges[1][1])

    plt.pcolor(X, Y, Z)
    plt.colorbar()

    if len(searched_coords) or len(contaminatiom_coords):
        s = np.array(searched_coords).T
        c = np.array(contaminatiom_coords).T
        plt.plot(s[0], s[1], "m*", label="Searched objects", markersize=17)
        plt.plot(
            c[0], c[1], "k*", label="Contamination objects", markersize=17)
        plt.legend()

    return x, y, Z

def plot1DProbabSpace(star_filter, plot_ranges, N,
                      searched_coords=[], contaminatiom_coords=[]):
    """
    Plot probability space

    Parameters
    ----------
    star_filter : StarsFilter object
        Trained stars filter

    plot_ranges : iterable
        Ranges (max/min) for all axis

    N : int
        Number of points per axis

    searched_coords : list, iterable
        List of coordinates of searched objects

    contaminatiom_coords : list, iterable
        List of coordinates of contamination objects

    Returns
    -------
    tuple
        x, y
    """

    if check_depth(plot_ranges, 2, ifnotraise=False):
        plot_ranges = plot_ranges[0]
    x = np.linspace(plot_ranges[0], plot_ranges[1])
    y = star_filter.evaluateCoordinates([[y] for y in x])

    plt.plot(x, y, linewidth=3)

    if len(searched_coords) or len(contaminatiom_coords):
        s = [qq[0] for qq in searched_coords]
        c = [qq[0] for qq in contaminatiom_coords]
        s_weights = np.ones_like(s) / len(s)
        c_weights = np.ones_like(c) / len(c)
        plt.hist(s, bins=x,
                 histtype='bar', weights=s_weights,
                 label="Searched objects")
        plt.hist(c, bins=x,
                 histtype='bar', weights=c_weights,
                 label="Contamination objects")
        plt.legend()

    return x, np.array(y)

def plotHist(searched_coo, cont_coo, labels=[], bins=None, save_path=None,
             file_name="hist.png"):
    """
    Plot histogram

    Parameters
    ----------
    searched_coo : iterable
        Coordinates of searched objects to plot the histogram

    cont_coo : iterable
        Coordinates of contamination objects to plot the histogram

    labels : list, tuple of str
        Labels for axis

    save_path : str, NoneType
        Path to the folder where plots are saved if not None, else
        plots are showed immediately

    bins : int, NoneType
        Number of bins for histogram

    file_name : str
        Name of the plot file

    Returns
    -------
    None
    """
    x = np.array(searched_coo).T
    y = np.array(cont_coo).T

    if len(x) != len(y):
        raise Exception(
            "Dimension of both searched and contamination sample have to be the same.\nGot: %i, %i" % (len(x), len(y)))
    if len(x) != len(labels):
        warnings.warn(
            "Dimension of the dimension of train sample and labels have to be the same.\nGot: %i, %i" % (len(x), len(labels)))
        labels = ["" for _ in x]

    for x_param, y_param, lab in zip(x, y, labels):
        plt.clf()

        if not bins:
            x_bins = 1 + 3.32 * np.log10(len(x_param))
            y_bins = 1 + 3.32 * np.log10(len(y_param))
        else:
            x_bins = bins
            y_bins = bins

        x_weights = np.ones_like(x_param) / len(x_param)
        y_weights = np.ones_like(y_param) / len(y_param)

        plt.hist(x_param, bins=int(x_bins), weights=x_weights,
                 histtype='bar', color="crimson",
                 label="Searched objects")
        plt.hist(y_param, bins=int(y_bins), weights=y_weights,
                 label="Others")
        plt.title("Distribution of the parameters coordinates")

        plt.xlabel(lab)
        plt.ylabel("Normalized counts")

        plt.legend()
        if save_path:
            plt.savefig(os.path.join(
                save_path, file_name + "_hist_%s.png" % (lab.replace(" ", "_"))))
        else:
            plt.show()

def plotUnsupProbabSpace(coords, decider, opt="show", N=100):
    if len(coords) and len(coords[0]) == 2:
        return plot2DUnsupProbabSpace(coords, decider, opt, N)
    elif len(coords) and len(coords[0]) == 1:
        return plot1DUnsupProbabSpace(coords, decider, opt, N)
    else:
        return plotNDUnsupProbabSpace(coords, decider, opt, N)

def plot2DUnsupProbabSpace(coords, decider, opt="show", N=50):
    OVERLAY = 0.2

    x_min, x_max = coords[:, 0].min(), coords[:, 0].max()
    y_min, y_max = coords[:, 1].min(), coords[:, 1].max()
    xo = (x_max - x_min) * OVERLAY
    yo = (y_max - y_min) * OVERLAY
    x, y = np.linspace(
        x_min - xo, x_max + xo, N), np.linspace(y_min - yo, y_max + yo, N)
    xx, yy = np.meshgrid(x, y)

    # Obtain labels for each point in mesh. Use last trained model.
    Z = decider.evaluate(np.c_[xx.ravel(), yy.ravel()])

    # Put the result into a color plot
    Z = Z.reshape(xx.shape)

    # Plot the centroids as a white X
    centroids = decider.classifier.cluster_centers_

    if opt == "show":
        plt.figure(1)
        plt.clf()
        plt.imshow(Z, interpolation='nearest',
                   extent=(xx.min(), xx.max(), yy.min(), yy.max()),
                   cmap=plt.cm.Paired,
                   aspect='auto', origin='lower')

        plt.plot(coords[:, 0], coords[:, 1], 'k.', markersize=2)
        plt.scatter(centroids[:, 0], centroids[:, 1],
                    marker='x', s=169, linewidths=3,
                    color='w', zorder=10)
        plt.title('')
        plt.xlim(x_min, x_max)
        plt.ylim(y_min, y_max)
        plt.xticks(())
        plt.yticks(())

        plt.show()

    return x, y, Z, centroids

def plot1DUnsupProbabSpace(coords, decider, opt, N):
    OVERLAY = 0.2

    x_min, x_max = coords[:, 0].min(), coords[:, 0].max()
    xo = (x_max - x_min) * OVERLAY
    x = np.linspace(x_min - xo, x_max + xo, N)

    y = decider.evaluate([[xx] for xx in x])
    centroids = decider.classifier.cluster_centers_
    return x, y, centroids

def plotNDUnsupProbabSpace(coords, decider, opt="show", N=8):
    """

    """
    OVERLAY = 0.2

    if not isinstance(coords, list):
        coords = list(coords)

    pca = decomposition.PCA(n_components=2)
    pca.fit(coords)
    red_coords = pca.transform(coords)

    xmax, ymax = np.max(red_coords, axis=0).tolist()
    xmin, ymin = np.min(red_coords, axis=0).tolist()

    xwidth = xmax-xmin
    ywidth = ymax - ymin

    x = np.linspace(xmin - xwidth*OVERLAY, xmax + ywidth*OVERLAY, N)
    y = np.linspace(ymin - ywidth*OVERLAY, ymax + ywidth*OVERLAY, N)
    X, Y = np.meshgrid(x, y)

    to_transf = np.c_[X.ravel(), Y.ravel()]
    back_transf_data = pca.inverse_transform(to_transf)

    z = np.array(decider.evaluate(back_transf_data))
    Z = z.reshape(X.shape)

    # Plot the centroids as a white X
    centroids = pca.transform(decider.classifier.cluster_centers_)

    plt.xlim(xmin, xmax)
    plt.ylim(ymin, ymax)

    plt.pcolor(X, Y, Z)
    plt.colorbar()

    if opt == "show":
        plt.figure(1)
        plt.clf()
        plt.imshow(Z, interpolation='nearest',
                   extent=(xmin, xmax, ymin, ymax),
                   cmap=plt.cm.Paired,
                   aspect='auto', origin='lower')

        plt.plot(red_coords[:, 0], red_coords[:, 1], 'k.', markersize=2)
        plt.scatter(centroids[:, 0], centroids[:, 1],
                    marker='x', s=169, linewidths=3,
                    color='w', zorder=10)
        plt.title('')
        plt.xlim(xmin, xmax)
        plt.ylim(ymin, ymax)
        plt.xticks(())
        plt.yticks(())

        plt.show()

    return x, y, Z, centroids, red_coords


text:
import math as m
import numpy as np
from scipy.integrate import ode
import plotly.graph_objects as go
import datetime

# Planetary Body & Orbit Initializer
class Orbit:
    def __init__(self, body, state, time_span, dt, TLES=False):

        if len(state) == 2:
            self.body = body
            self.r_vec = state[0]
            self.v_vec = state[1]
            self.time_span = time_span
            self.dt = dt

        elif TLES:
            self.filename = state

            with open(self.filename, 'r') as t:
                lines = t.readlines()

            line_Id = 0
            tle = []
            self.tles = []

            for line in lines:
                line_Id += 1
                tle.append(line)
                if line_Id == 3:
                    self.tles.append(tle)
                    tle = []
                    line_Id = 0
                else:
                    pass

            self.body = body
            self.filename = state
            self.time_span = time_span
            self.dt = dt

            self.r_vecs = []
            self.v_vecs = []

            self.tle_sat_names = []

            for tle_ in self.tles:
                TLE_into_COES = self.TLE_coes(tle_)

                TLECOES_into_rv = self.COES_rv(TLE_into_COES[0], TLE=True)

                tle_r_vec_ = []
                tle_v_vec_ = []

                TLE_r_vec = TLECOES_into_rv[0]
                TLE_v_vec = TLECOES_into_rv[1]

                for n in TLE_r_vec:
                    tle_r_vec_.append(n)

                for n in TLE_v_vec:
                    tle_v_vec_.append(n)

                self.r_vecs.append(tle_r_vec_)
                self.v_vecs.append(tle_v_vec_)

                self.tle_sat_names.append(TLE_into_COES[1])

            tle_v_vec_Id = 0
            self.tle_rv_vecs = []

            for tle_r_vec in self.r_vecs:
                self.tle_rv_vecs.append([tle_r_vec, self.v_vecs[tle_v_vec_Id]])
                tle_v_vec_Id += 1

        else:
            self.body = body
            self.time_span = time_span
            self.dt = dt

            COES_ = self.COES_rv(state)

            COES_r_vec = COES_[0]
            COES_v_vec = COES_[1]

            self.r_vec = []
            self.v_vec = []

            for n in COES_r_vec:
                self.r_vec.append(n)

            for n in COES_v_vec:
                self.v_vec.append(n)

    def diffrential_q(self, t, y, mu):
        rx, ry, rz, vx, vy, vz = y
        r = np.array([rx, ry, rz])

        norm_r = np.linalg.norm(r)

        ax, ay, az = -r * self.body['Mu']/norm_r**3

        return [vx, vy, vz, ax, ay, az]

    def Propagate_Orbit(self, TLES=False):
        if TLES:

            self.tle_orbits = []

            for tle_rv_vec in self.tle_rv_vecs:

                self.r_vec = tle_rv_vec[0]
                self.v_vec = tle_rv_vec[1]

                n_steps = int(np.ceil(self.time_span/self.dt))

                ys = np.zeros((n_steps, 6))
                ts = np.zeros((n_steps, 1))

                y0 = self.r_vec + self.v_vec
                ys[0] = np.array(y0)
                step = 1

                solver = ode(self.diffrential_q)
                solver.set_integrator('lsoda')
                solver.set_initial_value(y0, 0)
                solver.set_f_params(self.body['Mu'])

                while solver.successful() and step < n_steps:
                    solver.integrate(solver.t + self.dt)
                    ts[step] = solver.t
                    ys[step] = solver.y
                    step += 1

                rs = ys[:, :3]
                vs = ys[:, 3:]

                print(rs)

                rxs = []
                rys = []
                rzs = []
                for r in rs:
                    rxs.append(r[0])
                    rys.append(r[1])
                    rzs.append(r[2])

                self.tle_orbits.append([rxs, rys, rzs])

            return self.tle_orbits

        else:
            n_steps = int(np.ceil(self.time_span/self.dt))

            ys = np.zeros((n_steps, 6))
            ts = np.zeros((n_steps, 1))

            y0 = self.r_vec + self.v_vec
            ys[0] = np.array(y0)
            step = 1

            solver = ode(self.diffrential_q)
            solver.set_integrator('lsoda')
            solver.set_initial_value(y0, 0)
            solver.set_f_params(self.body['Mu'])

            while solver.successful() and step < n_steps:
                solver.integrate(solver.t + self.dt)
                ts[step] = solver.t
                ys[step] = solver.y
                step += 1

            rs = ys[:, :3]
            vs = ys[:, 3:]

            print(rs)

            rxs = []
            rys = []
            rzs = []
            for r in rs:
                rxs.append(r[0])
                rys.append(r[1])
                rzs.append(r[2])
            return [rxs, rys, rzs]

    def COES_rv(self, COES, deg=False, TLE=False):
        if deg:
            self.A, self.Ecc, self.I, self.T, self.P, self.R = COES
            self.I = self.I * (np.pi / 180)
            self.T = self.T * (np.pi / 180)
            self.P = self.P * (np.pi / 180)
            self.R = self.R * (np.pi / 180)
        else:
            self.A, self.Ecc, self.I, self.T, self.P, self.R = COES

        self.E = self.ecc_anomaly([self.T, self.Ecc], 'T&E')

        norm_r = self.A * (1 - self.Ecc ** 2) / (1 + self.Ecc * np.cos(self.T))

        r_perif = norm_r * np.array([m.cos(self.T), m.sin(self.T), 0])
        v_perif = m.sqrt(self.body['Mu'] * self.A) / norm_r * \
            np.array([-m.sin(self.E), m.cos(self.E)
                      * m.sqrt(1 - self.Ecc ** 2), 0])

        perif_eci = np.transpose(self.eci_perif(self.R, self.P, self.I))

        COES_r_vec = np.dot(perif_eci, r_perif)
        COES_v_vec = np.dot(perif_eci, v_perif)

        return [COES_r_vec, COES_v_vec]

    def eci_perif(self, R, P, I):
        row0 = [-m.sin(R) * m.cos(I) * m.sin(P) + m.cos(R) * m.cos(P), m.cos(R)
                * m.cos(I) * m.sin(P) + m.sin(R) * m.cos(P), m.sin(I) * m.sin(P)]
        row1 = [-m.sin(R) * m.cos(I) * m.cos(P) - m.cos(R) * m.sin(P), m.cos(R)
                * m.cos(I) * m.cos(P) - m.sin(R) * m.sin(P), m.sin(I) * m.cos(P)]
        row2 = [m.sin(R) * m.sin(I), -m.cos(R) * m.sin(I), m.cos(I)]
        return np.array([row0, row1, row2])

    def ecc_anomaly(self, vars_, method, tol=1e-8):
        if method == 'Newton':
            M, E = vars_
            if M < np.pi / 2.0:
                E0 = M + E / 2.0
            else:
                E0 = M - E
            for n in range(200):
                ratio = (E0 - E * np.sin(E0) - M) / (1 - E * np.cos(E0))
                if abs(ratio) < tol:
                    if n == 0:
                        return E0
                    else:
                        return E1
                else:
                    E1 = E0 - ratio
                    E0 = E1
            return False

        elif method == 'T&E':
            T, E = vars_
            return 2 * m.atan(m.sqrt((1-E) / (1+E)) * m.tan(T/2.0))

    def TLE_coes(self, tle):
        tle_sat = tle[0].strip()
        line1 = tle[1].strip().split()
        line2 = tle[2].strip().split()

        Epoch = line1[3]
        yy, mm, dd, h = self.calc_Epoch(Epoch)

        I = float(line2[2]) * (np.pi / 180)
        R = float(line2[3]) * (np.pi / 180)
        E_ = line2[4]
        Ecc = float('0.' + E_)
        P = float(line2[5]) * (np.pi / 180)
        M = float(line2[6]) * (np.pi / 180)
        M_motion = float(line2[7])
        T_P = 1 / M_motion * 24 * 3600
        A = (T_P ** 2 * self.body['Mu'] / 4.0 /
             np.pi ** 2) ** (1.0 / 3.0)

        self.E_A = self.ecc_anomaly(
            [M, Ecc], 'Newton')

        T = self.true_anomaly([self.E_A, Ecc])

        r_mag = A * (1 - Ecc * np.cos(self.E_A))

        return [A, Ecc, I, T, P, R], tle_sat

    def calc_Epoch(self, Epoch):
        YY = int('20' + Epoch[:2])

        Epoch = Epoch[2:].split('.')

        dd = int(Epoch[0]) - 1

        h = float('0.' + Epoch[1]) * 24.0

        date = datetime.date(YY, 1, 1) + datetime.timedelta(dd)

        MM = float(date.month)
        DD = float(date.day)

        return YY, MM, DD, h

    def true_anomaly(self, EccEA):
        E_A, Ecc = EccEA
        return 2 * np.arctan(np.sqrt((1 + Ecc) / (1 - Ecc)) * np.tan(E_A / 2.0))

    def Names(self):
        return self.tle_sat_names

    def spheres(self, radius, clr, clr_gradient=None, dist=0):
        theta = np.linspace(0, 2 * np.pi, 100)
        phi = np.linspace(0, np.pi, 100)

        x_ = dist + radius * np.outer(np.cos(theta), np.sin(phi))
        y_ = radius * np.outer(np.sin(theta), np.sin(phi))
        z_ = radius * np.outer(np.ones(100), np.cos(phi))

        if clr_gradient == None:
            clr_gradient = clr
        else:
            pass

        trace = go.Surface(x=x_, y=y_, z=z_,
                           colorscale=[
                               [0, clr], [1, clr_gradient]],
                           showlegend=True,
                           name=self.body['Name']
                           )
        trace.update(showscale=False)

        return trace

max_stars_repo_name: lalithu/EllipsePy

max_issues_repo_name: lalithu/EllipsePy

max_forks_repo_name: lalithu/EllipsePy

reason: import numpy,from scipy


text:
import numpy as np

from mps_tomo.utils import kron_embed, pauli_group

def pauli_proj(dens_mat, pauli_prod):
    return np.einsum("ij,ji", dens_mat, pauli_prod) * pauli_prod / np.size(dens_mat, 0)

def R_hat(sigmas, num_qubits):
    return sum(kron_embed(*i, num_qubits) for i in enumerate(sigmas))

def iteration(k, sigmas, num_qubits, max_its=100, delta=0.1):
    dim = 2 ** num_qubits
    paulis = pauli_group(k)

    R = R_hat(sigmas, num_qubits)

    Y = np.zeros((dim, dim), dtype=np.complex128)
    for _ in range(max_its):
        w, v = np.linalg.eigh(Y)
        y_val, y_vec = w[-1], v[:, -1]

        X = np.zeros_like(Y)
        for qubit in range(num_qubits - k + 1):
            y_vec_reshaped = np.reshape(y_vec, (2 ** qubit, 2 ** k, -1))
            for p in paulis:
                X += np.dot(
                    y_vec.conj(), np.ravel(np.einsum("ij,ajb->aib", p, y_vec_reshaped))
                ) * kron_embed(qubit, p, num_qubits)
        X *= y_val / dim

        Y += delta * (R - X)

    return y_vec

max_stars_repo_name: AidanGG/mps-tomo

max_issues_repo_name: AidanGG/mps-tomo

max_forks_repo_name: AidanGG/mps-tomo


text:
import csv, os, time
import numpy as np
import pandas as pd
import pandasql as ps
import time

filename = os.path.realpath("../plural_sent_pairs.csv")

df = pd.read_csv(filename)

query = """
select *
from (select distinct
    t1.b as a, t1.a as b, t2.b as c, t2.a as d, t1.category, 'plural|from-single' as subcategory
from
    df t1
inner join df t2
    on t1.b != t2.b
    and t1.category == t2.category
limit 3000000)
union
select *
from (select distinct
    t1.a, t1.b, t2.a as c, t2.b as d, t1.category, 'plural|to-' || t1.subcategory as subcategory
from
    df t1
inner join df t2
    on t1.a != t2.a
    and t1.category == t2.category
    and t1.subcategory == t2.subcategory
limit 3000000);
"""
print("running query")

t1 = time.perf_counter()
result_df = ps.sqldf(query)
t2 = time.perf_counter()

result_df.to_csv("../processed_plural_to_singular_pairs.csv", index=False)

max_stars_repo_name: alexwarstadt/phrase-analogies-large-vae

max_issues_repo_name: alexwarstadt/phrase-analogies-large-vae

max_forks_repo_name: alexwarstadt/phrase-analogies-large-vae


text:
#!/usr/bin/env python

from __future__ import print_function, division

import matplotlib
matplotlib.use('Qt4Agg')

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

import argparse
import numpy
import os
import re

def numeric(x):
	try: return int(x)
	except ValueError: pass

	try: return float(x)
	except ValueError: pass

	if '%' in x:
		try: return float(x[:-1])/100
		except ValueError: pass

	return None

class AdHocFigure(object):
	def __init__(self, fig=None, force_dims=None, x=None, y=None, z=None, xlabel=None, ylabel=None, zlabel=None, xlim=None, ylim=None, zlim=None, title=None):
		self.fig = plt.figure() if fig is None else fig
		self.dims = force_dims
		self.title = title

		self.x = x
		self.y = y
		self.z = z
		self.xlabel = xlabel
		self.ylabel = ylabel
		self.zlabel = zlabel
		self.xlim = xlim
		self.ylim = ylim
		self.zlim = zlim

		self.X = []
		self.Y = []
		self.Z = []

	def parse_tsv(self, infile):
		autolabel = True if (self.xlabel is None and self.ylabel is None and self.zlabel is None) else False
		autocol = True if (self.x is None and self.y is None and self.z is None) else False

		if self.dims is None:
			autodims = True if (self.x is None and self.y is None and self.z is None) else False
		else: autodims = False

		for l in infile:
			if not l.strip(): continue
			elif l.lstrip().startswith('#'):
				if autolabel:
					if '\t' in l: sl = l[1:].split('\t')
					#elif re.split(', *', l)
					elif re.findall(', *', l): sl = l[1:].split(', *')

					else: sl = l[1:].split('\t')
					try: self.xlabel = sl.pop(0)
					except IndexError: pass
					try: self.ylabel = sl.pop(0)
					except IndexError: pass
					try: self.zlabel = sl.pop(0)
					except IndexError: pass
			else:
				if autocol:
					sl = l.split('\t')
					done = [False, False, False]
					for rawval in sl:
						val = numeric(rawval)
						if val is None: continue

						if not done[0]:
							self.X.append(val)
							done[0] = True
						elif not done[1]:
							self.Y.append(val)
							done[1] = True
						elif not done[2]:
							self.Z.append(val)
							done[2] = True
				else:
					sl = l.split('\t')
					try:
						if self.x is not None:
							val = numeric(sl[self.x])
							#TODO: do some validation or something
							self.X.append(val)
						if self.y is not None:
							val = numeric(sl[self.y])
							#TODO: do some validation or something
							self.Y.append(val)
						if self.z is not None:
							val = numeric(sl[self.z])
							#TODO: do some validation or something
							self.Z.append(val)
					except IndexError: pass #maybe this shouldn't happen
		if self.xlabel is None: self.xlabel = ''
		if self.ylabel is None: self.ylabel = ''
		if self.zlabel is None: self.zlabel = ''

	def plot(self):
		if self.title is not None: self.add_title()
		truedims = sum([bool(self.X), bool(self.Y), bool(self.Z)])
		if self.dims is None: self.dims = truedims
		print(self.dims, truedims)

		if self.dims < truedims:
			if self.dims == 1:
				if truedims == 2: stuff = zip([211, 212], [self.X, self.Y], [self.xlim, self.ylim], [self.xlabel, self.ylabel])
				else: stuff = zip([221, 222, 223], [self.X, self.Y, self.Z], [self.xlim, self.ylim, self.zlim], [self.xlabel, self.ylabel, self.zlabel])
				#TODO: generalize to 4+ series
				for subplotspec, series, lim, label in stuff:
					ax = self.fig.add_subplot(subplotspec)
					hist = ax.hist(series, bins='fd')
					if lim is not None: ax.set_xlim(lim)
					if label is not None: ax.set_xlabel(label)

				ax = self.fig.add_subplot(224, projection='3d')
				scatter = ax.plot(self.X, self.Y, self.Z, lw=0, marker='o')
				if self.xlim is not None: ax.set_xlim(self.xlim)
				if self.ylim is not None: ax.set_ylim(self.ylim)
				if self.zlim is not None: ax.set_zlim(self.zlim)

				if self.xlabel is not None: ax.set_xlabel(self.xlabel)
				if self.ylabel is not None: ax.set_ylabel(self.ylabel)
				if self.zlabel is not None: ax.set_zlabel(self.zlabel)

			elif self.dims == 2:
				stuff = zip([221, 222, 223],
					[[self.X, self.Z], [self.Y, self.Z], [self.X, self.Y]],
					[[self.xlim, self.zlim], [self.ylim, self.zlim], [self.xlim, self.ylim]],
					[[self.xlabel, self.zlabel], [self.ylabel, self.zlabel], [self.xlabel, self.ylabel]],
				)
				for subplotspec, series, lims, labels in stuff:
					ax = self.fig.add_subplot(subplotspec)
					scatter = ax.plot(series[0], series[1], lw=0, marker='o')
					if lims[0] is not None: ax.set_xlim(lims[0])
					if lims[1] is not None: ax.set_ylim(lims[1])
					if labels[0] is not None: ax.set_xlabel(labels[0])
					if labels[1] is not None: ax.set_ylabel(labels[1])
				ax = self.fig.add_subplot(224, projection='3d')
				scatter = ax.plot(self.X, self.Y, self.Z, lw=0, marker='o')
				if self.xlim is not None: ax.set_xlim(self.xlim)
				if self.ylim is not None: ax.set_ylim(self.ylim)
				if self.zlim is not None: ax.set_zlim(self.zlim)

				if self.xlabel is not None: ax.set_xlabel(self.xlabel)
				if self.ylabel is not None: ax.set_ylabel(self.ylabel)
				if self.zlabel is not None: ax.set_zlabel(self.zlabel)

		#elif self.dims == truedims:
		else:
			if truedims == 1:
				ax = self.fig.add_subplot(111)
				hist = ax.hist(self.X, bins='fd')
				if self.xlim is not None: ax.set_xlim(self.xlim)
			elif truedims == 2:
				ax = self.fig.add_subplot(111)
				scatter = ax.plot(self.X, self.Y, marker='o', lw=0)
				if self.xlim is not None: ax.set_xlim(self.xlim)
				if self.ylim is not None: ax.set_ylim(self.ylim)
			elif truedims == 3:
				ax = self.fig.add_subplot(111, projection='3d')
				scatter = ax.plot(self.X, self.Y, self.Z, marker='o', lw=0)
				if self.xlim is not None: ax.set_xlim(self.xlim)
				if self.ylim is not None: ax.set_ylim(self.ylim)
				if self.zlim is not None: ax.set_zlim(self.zlim)

				if self.xlabel is not None: ax.set_xlabel(self.xlabel)
				if self.ylabel is not None: ax.set_ylabel(self.ylabel)
				if self.zlabel is not None: ax.set_zlabel(self.zlabel)
	def add_title(self, title=None):
		if title is None: title = self.title
		ax = self.fig.add_subplot(111)
		#ax.set_xticks([])
		ax.tick_params(labelcolor=(0,0,0,0), color=(0,0,0,0))
		ax.axes.get_yaxis().set_visible(False)
		ax.set_frame_on(False)
		ax.set_title(title)
	def show(self):
		self.fig.set_tight_layout(True)
		plt.show()

	def save(self, outfile, dpi=300):
		self.fig.savefig(outfile, dpi=dpi)

if __name__ == '__main__':
	parser = argparse.ArgumentParser()

	parser.add_argument('-o', '--outfile', help='Where to save the figure')
	parser.add_argument('infile', help='TSV file to read data from')
	parser.add_argument('--force-dims', default=None, type=int, help='Force the resulting plot to be this many dimensions. Can be less than or equal to the number of data series.')
	parser.add_argument('--title', default=None, help='Figure title')

	parser.add_argument('-x', default=None, type=int, help='(0-indexed) column with x values')
	parser.add_argument('-y', default=None, type=int, help='(0-indexed) column with y values')
	parser.add_argument('-z', default=None, type=int, help='(0-indexed) column with z values')
	parser.add_argument('--xlabel', default=None, help='Label for x axis')
	parser.add_argument('--ylabel', default=None, help='Label for y axis')
	parser.add_argument('--zlabel', default=None, help='Label for z axis')

	parser.add_argument('--xlim', nargs=2, type=float, default=None, help='Limits for x axis')
	parser.add_argument('--ylim', nargs=2, type=float, default=None, help='Limits for y axis')
	parser.add_argument('--zlim', nargs=2, type=float, default=None, help='Limits for z axis')

	args = parser.parse_args()

	ahfig = AdHocFigure(force_dims=args.force_dims, x=args.x, y=args.y, z=args.z, xlabel=args.xlabel, ylabel=args.ylabel, zlabel=args.zlabel, xlim=args.xlim, ylim=args.ylim, zlim=args.zlim, title=args.title)
	with open(args.infile) as f: ahfig.parse_tsv(f)
	ahfig.plot()
	ahfig.show()
	if args.outfile: ahfig.save(args.outfile)

max_stars_repo_name: SaierLaboratory/biotools

max_issues_repo_name: SaierLaboratory/biotools

max_forks_repo_name: SaierLaboratory/biotools


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Modifications for Guinet et al.

import io
import warnings
import numpy as np
import argparse
from utils import *
from query_aux import *

#Disable warnings for Meta-features
warnings.filterwarnings("ignore")

# to use bool for parsing
def str2bool(v):
    """Parse String to bool
    Args:
        v: String or Bool
    Returns:
        bool
    Raises:
        ArgumentTypeError: If v is not a String nor a bool
    """
    if isinstance(v, bool):
        return v
    if v.lower() in ("yes", "true", "t", "y", "1"):
        return True
    elif v.lower() in ("no", "false", "f", "n", "0"):
        return False
    else:
        raise argparse.ArgumentTypeError("Boolean value expected.")

parser = argparse.ArgumentParser(description="Extraction of queries simplified")
parser.add_argument(
    "--src_emb", type=str, default="", help="Load source embeddings for training"
)
parser.add_argument(
    "--tgt_emb", type=str, default="", help="Load target embeddings for validation"
)
parser.add_argument(
    "--filename", type=str, default="", help="Filename of lightsvm files extracted"
)
parser.add_argument(
    "--center", action="store_true", help="whether to center embeddings or not"
)
parser.add_argument(
    "--dico", type=str, default="", help="Dictionary for query extraction"
)
parser.add_argument("--maxload", type=int, default=200000)

parser.add_argument(
    "--query_relevance_type",
    type=str,
    default="",
    help="Type of query relevance: binary or continuous",
)
parser.add_argument("--query_size", type=int, default=10, help="Size of the query")
parser.add_argument(
    "--add_csls_coord",
    type=str2bool,
    default=True,
    help="Whether to add to query coord CSLS distance",
)
parser.add_argument(
    "--k_csls",
    type=int,
    default=10,
    help="Number of coord in query for CSLS distance (from 0 to k)",
)
parser.add_argument(
    "--testing_query",
    type=str2bool,
    default=False,
    help="Whether to impose the ground truth traduction presence in the query",
)
parser.add_argument(
    "--add_word_coord",
    type=str2bool,
    default=False,
    help="Whether to add to query coord word embedding",
)
parser.add_argument(
    "--discard_empty_query",
    type=str2bool,
    default=False,
    help="Whether to remove query without the right traduction or not",
)
parser.add_argument(
    "--use_csls",
    type=str2bool,
    default=False,
    help="Whether to use CSLS distance or CosineSim",
)
parser.add_argument(
    "--add_query_coord",
    type=str2bool,
    default=False,
    help="Whether to add to query coord query word embedding",
)

parser.add_argument(
    "--add_meta_features",
    type=str2bool,
    default=True,
    help="Whether to add to meta-features of the 2 clouds (source and target)",
)

parser.add_argument(
    "--center_meta_features",
    type=str2bool,
    default=True,
    help="Whether to add to center the meta-features of the target clouds",
)

parser.add_argument(
    "--nn_size_meta_features",
    type=int,
    default=10,
    help="Number of neighbors to use when computing meta-features",
)

params = parser.parse_args()

###### MAIN ######

query_extractor = (
    compute_binary_distance
    if params.query_relevance_type == "binary"
    else compute_embedding_distance
)

print("Extraction of queries alignment on %s" % params.dico)

words_tgt, x_tgt = load_vectors(
    params.tgt_emb, maxload=params.maxload, center=params.center
)
words_src, x_src = load_vectors(
    params.src_emb, maxload=params.maxload, center=params.center
)

print("Loading and extracting data")
src2tgt, lexicon_size = load_lexicon(params.dico, words_src, words_tgt)

query_extractor(
    x_src,
    x_tgt,
    params.filename,
    src2tgt,
    add_csls_coord=params.add_csls_coord,
    k_csls=params.k_csls,
    testing_query=params.testing_query,
    discard_empty_query=params.discard_empty_query,
    add_word_coord=params.add_word_coord,
    add_query_coord=params.add_query_coord,
    add_meta_features=params.add_meta_features,
    center_meta_features=params.center_meta_features,
    nn_size_meta_features=params.nn_size_meta_features,
    query_size=params.query_size,
    use_csls=params.use_csls
)
print("Query file extracted")

max_stars_repo_name: Gguinet/semisupervised-alignment

max_issues_repo_name: Gguinet/semisupervised-alignment

max_forks_repo_name: Gguinet/semisupervised-alignment


text:
import pandas as pd
import os
import csv
import os
from glob import glob
from skimage.io import imread, imsave
import numpy as np

from glob import glob
import os

path_to_train = '../../data/train/'
glob_train_imgs = os.path.join(path_to_train, '*_sat.jpg')
glob_train_masks = os.path.join(path_to_train, '*_msk.png')

train_img_paths = glob(glob_train_imgs)
train_mask_paths = glob(glob_train_masks)
print(train_img_paths[:10])
print(train_mask_paths[:10])

from skimage.io import imread
from skimage.transform import resize
from skimage.color import rgb2gray

def get_img_id(img_path):
    img_basename = os.path.basename(img_path)
    img_id = os.path.splitext(img_basename)[0][:-len('_sat')]
    return img_id

def img_gen(img_paths, img_size=(512, 512)):
    # Iterate over all image paths
    for img_path in img_paths:
        img_id = get_img_id(img_path)
        mask_path = os.path.join(path_to_train, img_id + '_msk.png')

        img = imread(img_path) / 255
        mask = rgb2gray(imread(mask_path))

        #     img = resize(img, img_size, preserve_range = True)
        #     mask = resize(mask, img_size, mode='constant', preserve_range = True)

        mask = (mask >= 0.5).astype(float)
        mask = np.reshape(mask, (512, 512, 1))
        yield img, mask

def get_non_outlier_data(train_img_paths):
    train_path_without_outlier = []
    for index, image_path in enumerate(train_img_paths):
        if index % 500 == 0:
            print(index)
        img_id = get_img_id(image_path)
        mask_path = os.path.join('../../data/train/', img_id + '_msk.png')
        mask = rgb2gray(imread(mask_path))
        if len(np.where(mask.flatten() != 0)[0]) != 0:
            train_path_without_outlier.append(image_path)
    return train_path_without_outlier

train_img_paths = get_non_outlier_data(train_img_paths)
import numpy as np

# Keras takes its input in batches
# (i.e. a batch size of 32 would correspond to 32 images and 32 masks from the generator)

# The generator should run forever
def image_batch_generator(img_paths, batchsize=32):
    while True:
        ig = img_gen(img_paths)
        batch_img, batch_mask = [], []

        for img, mask in ig:
            # Add the image and mask to the batch
            batch_img.append(img)
            batch_mask.append(mask)
            # If we've reached our batchsize, yield the batch and reset
            if len(batch_img) == batchsize:
                yield np.stack(batch_img, axis=0), np.stack(batch_mask, axis=0)
                batch_img, batch_mask = [], []

        # If we have an nonempty batch left, yield it out and reset
        if len(batch_img) != 0:
            yield np.stack(batch_img, axis=0), np.stack(batch_mask, axis=0)
            batch_img, batch_mask = [], []

from sklearn.model_selection import train_test_split

BATCHSIZE = 1

# Split the data into a train and validation set
train_img_paths, val_img_paths = train_test_split(train_img_paths, test_size=0.1,random_state=42)
print(len(train_img_paths))
print(len(val_img_paths))

# Create the train and validation generators
traingen = image_batch_generator(train_img_paths, batchsize=BATCHSIZE)
valgen = image_batch_generator(val_img_paths, batchsize=BATCHSIZE)
print(next(traingen)[0].shape)
def calc_steps(data_len, batchsize):
    return (data_len + batchsize - 1) // batchsize

# Calculate the steps per epoch
train_steps = calc_steps(len(train_img_paths), BATCHSIZE)
val_steps = calc_steps(len(val_img_paths), BATCHSIZE)
print("train steps {}".format(train_steps))
print("val steps {}".format(val_steps))

import os
import sys
import random
import warnings

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt

from tqdm import tqdm
from itertools import chain
from skimage.io import imread, imshow, imread_collection, concatenate_images
from skimage.transform import resize
from skimage.morphology import label

from keras.models import Model, load_model
from keras.layers import Input, BatchNormalization
from keras.layers.core import Dropout, Lambda
from keras.layers.convolutional import Conv2D, Conv2DTranspose
from keras.layers.pooling import MaxPooling2D
from keras.layers.merge import concatenate
from keras.callbacks import *
from keras import backend as K

import tensorflow as tf

# Build U-Net model
inputs = Input((512, 512, 3))
s = inputs

c0 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(inputs)
c0 = Dropout(0.2)(c0)
c0 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c0)
p0 = MaxPooling2D((2, 2))(c0)

c1 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(p0)
c1 = Dropout(0.2)(c1)
c1 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c1)
p1 = MaxPooling2D((2, 2))(c1)

c2 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(p1)
c2 = Dropout(0.2)(c2)
c2 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c2)
p2 = MaxPooling2D((2, 2))(c2)

c3 = Conv2D(256, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(p2)
c3 = Dropout(0.2)(c3)
c3 = Conv2D(256, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c3)
p3 = MaxPooling2D((2, 2))(c3)

c4 = Conv2D(512, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(p3)
c4 = Dropout(0.3)(c4)
c4 = Conv2D(512, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c4)
p4 = MaxPooling2D(pool_size=(2, 2))(c4)

c5 = Conv2D(1024, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(p4)
c5 = Dropout(0.4)(c5)
c5 = Conv2D(1024, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c5)

u6 = Conv2DTranspose(512, (2, 2), strides=(2, 2), padding='same')(c5)
u6 = concatenate([u6, c4])
c6 = Conv2D(512, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(u6)
c6 = Dropout(0.3)(c6)
c6 = Conv2D(512, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c6)

u7 = Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(c6)
u7 = concatenate([u7, c3])
c7 = Conv2D(256, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(u7)
c7 = Dropout(0.2)(c7)
c7 = Conv2D(256, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c7)

u8 = Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(c7)
u8 = concatenate([u8, c2])
c8 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(u8)
c8 = Dropout(0.1)(c8)
c8 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c8)

u9 = Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(c8)
u9 = concatenate([u9, c1])
c9 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(u9)
c9 = Dropout(0.1)(c9)
c9 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c9)

u10 = Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(c9)
u10 = concatenate([u10, c0], axis=3)
c10 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(u10)
c10 = Dropout(0.1)(c10)
c10 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c10)

outputs = Conv2D(1, (1, 1), activation='sigmoid')(c10)

import keras.backend as K
from keras.optimizers import Adam
from keras.losses import binary_crossentropy

smooth = 1e-9

# This is the competition metric implemented using Keras
def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred = K.cast(y_pred, 'float32')
    y_pred_f = K.cast(K.greater(K.flatten(y_pred), 0.5), 'float32')
    intersection = y_true_f * y_pred_f
    score = 2. * (K.sum(intersection) + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)
    return score

# We'll construct a Keras Loss that incorporates the DICE score
def dice_loss(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return 1. - (2. * intersection + 1.) / (K.sum(y_true_f) + K.sum(y_pred_f) + 1.)

def bce_dice_loss(y_true, y_pred):
    return 0.5 * binary_crossentropy(y_true, y_pred) + dice_loss(y_true, y_pred)

def ln_dice(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred = K.cast(y_pred, 'float32')
    y_pred_f = K.cast(K.greater(K.flatten(y_pred), 0.5), 'float32')
    intersection = y_true_f * y_pred_f
    score = 2. * (K.sum(intersection) + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)
    return K.log(score)

def new_bce_dice_loss(y_true, y_pred):
    return binary_crossentropy(y_true, y_pred) - ln_dice(y_true, y_pred)

model = Model(inputs=[inputs], outputs=[outputs])
model.compile(Adam(1e-4), loss=bce_dice_loss, metrics=[dice_coef])
model.summary()
MODEL_NAME = 'nonblack_13layer-{epoch:02d}-{val_loss:.4f}-id7.h5'
import math
checkpointer = ModelCheckpoint(MODEL_NAME, verbose=1,monitor='val_loss', save_best_only=True)
def step_decay(epoch):
   initial_lrate = 1e-4
   drop = 0.5
   epochs_drop = 15.0
   lrate = initial_lrate * math.pow(drop,
           math.floor((1+epoch)/epochs_drop))
   return lrate
lrate = LearningRateScheduler(step_decay)
EPOCH_NUM = 60
# Train the model
history = model.fit_generator(
    traingen,
    steps_per_epoch=train_steps,
    epochs=EPOCH_NUM, # Change this to a larger number to train for longer
    validation_data=valgen,
    validation_steps=val_steps,
    verbose=1,
    max_queue_size=5,  # Change this number based on memory restrictions
    callbacks = [checkpointer,lrate]
)


text:
import numpy as np
import torch
from torch.autograd import Variable
from torch.nn import functional as F
from scipy import sparse
from sklearn import metrics
from scipy.sparse import csr_matrix

from utils import load_data

class ProblemLosses:
    @staticmethod
    def multilabel_classification(preds, targets):
        return F.multilabel_soft_margin_loss(preds, targets)

    @staticmethod
    def classification(preds, targets):
        return F.cross_entropy(preds, targets)

    @staticmethod
    def regression_mae(preds, targets):
        return F.l1_loss(preds, targets)

class ProblemMetrics:
    @staticmethod
    def multilabel_classification(y_true, y_pred):
        y_pred = (y_pred > 0).astype(int)
        return {
            "micro" : float(metrics.f1_score(y_true, y_pred, average="micro")),
            "macro" : float(metrics.f1_score(y_true, y_pred, average="macro")),
        }

    @staticmethod
    def classification(y_true, y_pred):
        y_pred = np.argmax(y_pred, axis=1)
        return {
            "micro" : float(metrics.f1_score(y_true, y_pred, average="micro")),
            "macro" : float(metrics.f1_score(y_true, y_pred, average="macro")),
        }

    @staticmethod
    def regression_mae(y_true, y_pred):
        return float(np.abs(y_true - y_pred).mean())

class Problem(object):
    def __init__(self, file_path, task, cuda=True, batch_size=100, max_degree=25):

        self.task = task
        self.cuda = False
        self.loss_fn = getattr(ProblemLosses, self.task)
        self.metric_fn = getattr(ProblemMetrics, self.task)

        print('Problem: loading started')
        self.max_degree = max_degree
        self.batch_size = batch_size

    def construct_adj(self):
        raise NotImplementedError()

    def construct_test_adj(self):
        adj = len(self.id2idx)*np.ones((len(self.id2idx)+1, self.max_degree))
        for nodeid in self.G.nodes():
            neighbors = np.array([self.id2idx[neighbor]
                for neighbor in self.G.neighbors(nodeid)])
            if len(neighbors) == 0:
                continue
            if len(neighbors) > self.max_degree:
                neighbors = np.random.choice(neighbors, self.max_degree, replace=False)
            elif len(neighbors) < self.max_degree:
                neighbors = np.random.choice(neighbors, self.max_degree, replace=True)
            adj[self.id2idx[nodeid], :] = neighbors
        return adj

    def _to_torch(self):
        if not sparse.issparse(self.val_adj):
            self.val_adj = Variable(torch.LongTensor(self.val_adj))
            self.train_adj = Variable(torch.LongTensor(self.train_adj))
            if self.cuda:
                self.adj = self.val_adj.cuda()
                self.train_adj = self.train_adj.cuda()

        if self.features is not None:
            self.features = Variable(torch.FloatTensor(self.features))
            if self.cuda:
                self.features = self.features.cuda()

    def _batch_to_torch(self, mids, targets):
        """ convert batch to torch """
        mids = Variable(torch.LongTensor(mids))
        # print(targets)
        if self.task == 'multilabel_classification':
            targets = Variable(torch.FloatTensor(targets))
        elif self.task == 'classification':
            targets = Variable(torch.LongTensor(targets))
        elif 'regression' in self.task:
            targets = Variable(torch.FloatTensor(targets))
        else:
            raise Exception('NodeDataLoader: unknown task: %s' % self.task)

        if self.cuda:
            mids, targets = mids.cuda(), targets.cuda()

        return mids, targets

    def iterate(self, mode, batch_size=512, shuffle=False):
        nodes = self.nodes[mode]
        idx = np.arange(nodes.shape[0])
        if shuffle:
            idx = np.random.permutation(idx)

        n_chunks = idx.shape[0] // batch_size + 1
        """
        chunk: batch id
        chunkd id: 1, 2,..., n
        """
        for chunk_id, chunk in enumerate(np.array_split(idx, n_chunks)):
            mids = nodes[chunk]
            targets = np.concatenate(self.targets[mids]).reshape(len(mids), -1)
            mids, targets = self._batch_to_torch(mids, targets)
            yield mids, targets, chunk_id / n_chunks

class NodeProblem(Problem):
    def __init__(self, file_path, task, cuda=True, batch_size=100, max_degree=25):
        self.G, self.id2idx, self.id2target, self.features = load_data(file_path)
        self.max_degree = max_degree
        self.batch_num = 0

        Problem.__init__(self, file_path, task, cuda=True, batch_size=100, max_degree=25)

        self.targets = self.id2target.values

        self.n_dim = self.features.shape[1] if self.features is not None else None

        if isinstance(list(self.id2target.values[0]), list):
            self.n_class = len(list(self.id2target.values)[0])
        else:
            self.n_class = len(set(self.id2target.values))

        self.train_adj, self.deg = self.construct_adj()
        self.val_adj = self.construct_test_adj()

        self.n_node   = self.train_adj.shape[0]
        self.val_nodes = [n for n in self.G.nodes() if self.G.node[n]['val']]
        self.test_nodes = [n for n in self.G.nodes() if self.G.node[n]['test']]

        self.no_train_nodes_set = set(self.val_nodes + self.test_nodes)
        self.train_nodes = set(self.G.nodes()).difference(self.no_train_nodes_set)

        # don't train on nodes that only have edges to test set
        self.train_nodes = [n for n in self.train_nodes if self.deg[self.id2idx[n]] > 0]

        self.nodes = {
            'train': np.array(self.train_nodes),
            'val': np.array(self.val_nodes),
            'test': np.array(self.test_nodes),
        }

        print(self.nodes['train'].shape)
        print(self.nodes['val'].shape)
        print(self.nodes['test'].shape)

        self._to_torch()

        print('NodeProblem: loading finished')

    def construct_adj(self):
        adj = len(self.id2idx)*np.ones((len(self.id2idx)+1, self.max_degree))
        deg = np.zeros((len(self.id2idx),))

        for nodeid in self.G.nodes():
            if self.G.node[nodeid]['test'] or self.G.node[nodeid]['val']:
                continue
            neighbors = np.array([self.id2idx[neighbor]
                for neighbor in self.G.neighbors(nodeid)
                if (not self.G[nodeid][neighbor]['train_removed'])])
            deg[self.id2idx[nodeid]] = len(neighbors)
            if len(neighbors) == 0:
                continue
            if len(neighbors) > self.max_degree:
                neighbors = np.random.choice(neighbors, self.max_degree, replace=False)
            elif len(neighbors) < self.max_degree:
                neighbors = np.random.choice(neighbors, self.max_degree, replace=True)
            adj[self.id2idx[nodeid], :] = neighbors
        return adj, deg


text:
import os
import argparse
from collections import defaultdict
from multiprocessing import cpu_count
import glob

import numpy as np
import torch
import torch.nn.functional as F
from trimesh import load_mesh
from torch.utils.data import DataLoader
from torch.utils.data.dataloader import default_collate
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, RichModelSummary, RichProgressBar
from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
import wandb

from src import config, data
from src.common import compute_iou

class LitConvOnet(pl.LightningModule):
    def __init__(self,
                 cfg,
                 train_dataset,
                 val_dataset,
                 visualize: bool = True,
                 vis_every_n_evals: int = 1):
        super().__init__()
        self.cfg = cfg
        self.model = config.get_model(cfg, train_dataset)
        self.val_dataset = val_dataset
        self.threshold = float(cfg["test"]["threshold"])
        self.n_evals = 0
        self.vis_every_n_evals = vis_every_n_evals
        self.visualize = visualize
        self.save_hyperparameters(cfg)

        if visualize:
            out_dir = os.path.abspath(cfg["training"]["out_dir"])
            self.vis_path = os.path.join(out_dir, "vis")
            os.makedirs(self.vis_path, exist_ok=True)
            self.generator = config.get_generator(self.model, self.cfg, device=None)
            self.vis_data = dict()
            self.model_counter = defaultdict(int)
            self.vis_this_eval = False

    def generate_mesh(self, inputs):
        if not isinstance(inputs, torch.Tensor):
            inputs = torch.from_numpy(inputs)
        if inputs.shape[0] != 1:
            inputs = torch.unsqueeze(inputs, dim=0)
        return self.generator.generate_mesh({"inputs": inputs}, return_stats=False)

    def visualize_data(self, batch, batch_idx):
        if self.n_evals == 0:
            for idx, inputs in zip(batch.get("idx"), batch.get("inputs")):
                model_dict = self.val_dataset.get_model_dict(idx)
                category_id = model_dict.get("category")

                instance = self.model_counter[category_id]
                if instance < self.cfg["generation"]["vis_n_outputs"]:
                    category_name = self.val_dataset.metadata[category_id].get("name")
                    category_name = category_name.split(',')[0]

                    if batch_idx in self.vis_data:
                        self.vis_data[batch_idx].append({"inputs": inputs,
                                                         "category_name": category_name,
                                                         "instance": instance})
                    else:
                        self.vis_data[batch_idx] = [{"inputs": inputs,
                                                     "category_name": category_name,
                                                     "instance": instance}]

                    mesh = self.generate_mesh(inputs)
                    if len(mesh.faces) > 0 and len(mesh.vertices) > 0:
                        mesh.export(os.path.join(self.vis_path, f"{self.global_step}_{category_name}_{instance}.obj"))
                self.model_counter[category_id] += 1
        elif batch_idx in self.vis_data:
            for batch_data in self.vis_data[batch_idx]:
                inputs = batch_data["inputs"]
                category_name = batch_data["category_name"]
                instance = batch_data["instance"]

                mesh = self.generate_mesh(inputs)
                if len(mesh.faces) > 0 and len(mesh.vertices) > 0:
                    mesh.export(os.path.join(self.vis_path, f"{self.global_step}_{category_name}_{instance}.obj"))

    def forward(self, x):
        return self.model(**x)

    def on_fit_start(self):
        if self.visualize:
            self.generator.device = self.device

    def on_train_start(self):
        self.n_evals = 0
        if self.visualize:
            self.vis_data = dict()
            self.model_counter = defaultdict(int)
            meshes = glob.glob(os.path.join(self.vis_path, f"{self.global_step}*.obj"))
            for mesh in meshes:
                os.remove(mesh)

    def training_step(self, batch, batch_idx):
        logits = self({"points": batch.get("points"), "inputs": batch.get("inputs")}).logits
        occ = batch.get("points.occ")
        loss = F.binary_cross_entropy_with_logits(logits, occ, reduction="none").sum(-1).mean()
        self.log("loss", loss, batch_size=self.cfg["training"]["batch_size"])
        return loss

    def on_validation_start(self):
        self.vis_this_eval = False
        if self.visualize and self.n_evals % self.vis_every_n_evals == 0:
            self.vis_this_eval = True

    def validation_step(self, batch, batch_idx):
        out = self({"points": batch.get("points_iou"), "inputs": batch.get("inputs")})
        occ_iou = batch.get("points_iou.occ")
        loss = F.binary_cross_entropy_with_logits(out.logits, occ_iou, reduction="none").sum(-1).mean()
        self.log("val_loss", loss, prog_bar=True, batch_size=self.cfg["training"]["batch_size_val"])

        occ_iou_np = (occ_iou >= 0.5).cpu().numpy()
        occ_iou_hat_np = (out.probs >= self.threshold).cpu().numpy()
        iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
        self.log("val_iou", iou, prog_bar=True, batch_size=self.cfg["training"]["batch_size_val"])

        if self.vis_this_eval:
            self.visualize_data(batch, batch_idx)
        return loss

    def on_validation_end(self):
        self.n_evals += 1

    def test_step(self, batch, batch_idx):
        out = self({"points": batch.get("points_iou"), "inputs": batch.get("inputs")})
        occ_iou = batch.get("points_iou.occ")
        loss = F.binary_cross_entropy_with_logits(out.logits, occ_iou, reduction="none").sum(-1).mean()
        self.log("test_loss", loss, prog_bar=True, batch_size=self.cfg["training"]["batch_size_test"])

        occ_iou_np = (occ_iou >= 0.5).cpu().numpy()
        occ_iou_hat_np = (out.probs >= self.threshold).cpu().numpy()
        iou = compute_iou(occ_iou_np, occ_iou_hat_np).mean()
        self.log("test_iou", iou, prog_bar=True, batch_size=self.cfg["training"]["batch_size_test"])
        return loss

    def predict_step(self, batch, batch_idx, dataloader_idx=None):
        pass

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-4)
        return optimizer

class LitDataModule(pl.LightningDataModule):
    def __init__(self,
                 cfg,
                 num_workers: int = cpu_count(),
                 prefetch_factor: int = 2,
                 pin_memory: bool = False,
                 seed: int = 0):
        super().__init__()
        self.cfg = cfg
        self.train = None
        self.val = None
        self.test = None
        self.num_workers = num_workers
        self.prefetch_factor = prefetch_factor
        self.pin_memory = pin_memory
        self.generator = torch.Generator()
        self.generator.manual_seed(seed)
        if self.cfg["data"]["pointcloud_n"] and not self.cfg["data"]["voxelize"]:
            self.collate_fn = default_collate
        else:
            self.collate_fn = data.heterogeneous_batching

    def setup(self, stage=None):
        self.train = config.get_dataset("train", self.cfg)
        self.val = config.get_dataset("val", self.cfg, return_idx=True)
        self.test = config.get_dataset("test", self.cfg, return_idx=True)

    def train_dataloader(self):
        return torch.utils.data.DataLoader(self.train,
                                           batch_size=self.cfg["training"]["batch_size"],
                                           num_workers=self.num_workers,
                                           collate_fn=self.collate_fn,
                                           shuffle=True,
                                           pin_memory=self.pin_memory,
                                           worker_init_fn=data.worker_init_reset_seed,
                                           generator=self.generator,
                                           prefetch_factor=self.prefetch_factor,
                                           persistent_workers=True)

    def val_dataloader(self):
        return torch.utils.data.DataLoader(self.val,
                                           batch_size=self.cfg["training"]["batch_size_val"],
                                           num_workers=self.num_workers,
                                           collate_fn=self.collate_fn,
                                           prefetch_factor=self.prefetch_factor,
                                           pin_memory=self.pin_memory,
                                           worker_init_fn=data.worker_init_reset_seed,
                                           generator=self.generator,
                                           persistent_workers=True)

    def test_dataloader(self):
        return torch.utils.data.DataLoader(self.test,
                                           batch_size=self.cfg["training"]["batch_size_test"],
                                           num_workers=self.num_workers,
                                           prefetch_factor=self.prefetch_factor,
                                           pin_memory=self.pin_memory,
                                           worker_init_fn=data.worker_init_reset_seed,
                                           generator=self.generator)

class WandbVisualizeCallback(pl.Callback):
    def on_validation_epoch_end(self, trainer, pl_module):
        if pl_module.vis_this_eval:
            vis_data = pl_module.vis_data
            vis_path = pl_module.vis_path
            paths = list()
            names = list()
            instances = list()
            for batch_idx in vis_data.keys():
                for batch_data in vis_data[batch_idx]:
                    category_name = batch_data["category_name"]
                    instance = batch_data["instance"]
                    mesh_path = os.path.join(vis_path, f"{trainer.global_step}_{category_name}_{instance}.obj")
                    if os.path.isfile(mesh_path):
                        paths.append(mesh_path)
                        names.append(category_name)
                        instances.append(instance)

            if paths:
                trainer.logger.experiment.log({
                    "vis/completion": [wandb.Object3D(p,
                                                      caption=f"{n}_{i}") for p, n, i in zip(paths, names, instances)],
                    "global_step": trainer.global_step
                })

def main():
    parser = argparse.ArgumentParser(description="Train a 3D reconstruction model.")
    parser.add_argument("config", type=str, help="Path to config file.")
    parser.add_argument("--weights", type=str, help="Path to pre-trained weights.")
    parser.add_argument("--checkpoint", type=str, help="Path to PyTorch Lightning checkpoint.")
    parser.add_argument("--resume", action="store_true", help="Resume training instead of starting from scratch.")
    parser.add_argument("--early_stopping", type=str, default="none", choices=["none", "val_iou", "val_loss"],
                        help="Terminate if validation loss or iou stops improving.")
    parser.add_argument("--auto_lr", action="store_true", help="Tune learning rate automatically.")
    parser.add_argument("--auto_batch_size", action="store_true", help="Tune batch size automatically.")
    parser.add_argument("--test", action="store_true", help="Runs evaluation on the test set after training.")
    parser.add_argument("--wandb", action="store_true", help="Use the weights & biases logger.")
    parser.add_argument("--name", type=str, help="The name of the train run.")
    parser.add_argument("--project", type=str, help="The name of project.")
    parser.add_argument("--id", type=str, help="The weights & biases run id.")
    parser.add_argument("--offline", action="store_true", help="Log offline.")
    parser.add_argument("--precision", type=int, default=32, help="Train with float16, 32 or 64 precision.")
    parser.add_argument("--profile", action="store_true", help="Profile code to find bottlenecks.")
    parser.add_argument("--no_progress", action="store_true", help="Disable progress bar.")
    parser.add_argument("--pin_memory", action="store_true", help="Enable GPU memory pinning in data loaders.")
    parser.add_argument("--verbose", action="store_true", help="Enable verbose output.")
    parser.add_argument("--visualize", action="store_true", help="Generate some validation meshes.")
    parser.add_argument("--seed", type=int, default=0, help="Set random seed.")
    parser.add_argument("--num_workers", type=int, default=cpu_count(),
                        help="Number of workers to spawn for data loading.")
    parser.add_argument("--prefetch_factor", type=int, default=2, help="Number of batches to prefetch per worker.")
    args = parser.parse_args()

    data.seed_all_rng(args.seed)

    cfg = config.load_config(args.config, "configs/default.yaml")
    batch_size = cfg["training"]["batch_size"]
    model_selection_metric = cfg["training"]["model_selection_metric"]
    model_selection_mode = cfg["training"]["model_selection_mode"]
    out_dir = os.path.abspath(cfg["training"]["out_dir"])
    max_iter = cfg["training"]["max_iter"]
    validate_every = cfg["training"]["validate_every"]
    visualize_every = cfg["training"]["visualize_every"]
    print_every = cfg["training"]["print_every"]

    os.makedirs(out_dir, exist_ok=True)
    split_out_dir = out_dir.split('/')
    save_dir = '/'.join(split_out_dir[:-1])
    name = args.name if args.name else split_out_dir[-1]
    project = args.project if args.project else split_out_dir[-2]

    num_workers = args.num_workers
    if num_workers == -1:
        num_workers = cpu_count()
    dataset = LitDataModule(cfg, num_workers, args.prefetch_factor, args.pin_memory, args.seed)
    dataset.setup()

    n_train_batches = int(np.ceil(len(dataset.train) / batch_size))
    eval_every_n_epochs = np.round(validate_every / n_train_batches, 1)
    vis_every_n_evals = int(np.ceil(visualize_every / validate_every))
    max_steps = max_iter
    max_epochs = max_steps // n_train_batches
    print(f"Number of train batches:", n_train_batches)
    print(f"Evaluating every {validate_every} steps")
    print(f"Evaluating every {eval_every_n_epochs} epochs")
    if args.visualize:
        print(f"Visualizing every {vis_every_n_evals} evaluations")
    print(f"Training for max. {max_steps} steps")
    print(f"Training for max. {max_epochs} epochs")
    print(f"Logging every {print_every} steps")

    conv_onet = LitConvOnet(cfg, dataset.train, dataset.val, args.visualize, vis_every_n_evals)
    if not args.resume:
        if args.weights:
            state_dict = torch.load(os.path.abspath(args.weights))["model"]
            conv_onet.model.load_state_dict(state_dict)
        if args.checkpoint:
            conv_onet = LitConvOnet.load_from_checkpoint(os.path.abspath(args.checkpoint),
                                                         cfg=cfg,
                                                         train_dataset=dataset.train,
                                                         val_dataset=dataset.val,
                                                         vis_every_n_evals=vis_every_n_evals)
    callbacks = [ModelCheckpoint(filename="{epoch}-{step}-{val_loss:.2f}-{val_iou:.2f}",
                                 monitor=f"val_{model_selection_metric}",
                                 verbose=args.verbose,
                                 save_last=True,
                                 mode="max" if "max" in model_selection_mode else "min"),
                 RichModelSummary()]
    if not args.no_progress:
        callbacks.append(RichProgressBar())
    patience = max_epochs // int(np.ceil(eval_every_n_epochs)) // 10
    if args.early_stopping != "none" and patience >= 3:
        print(f"Will terminate after {patience} evaluations without improvement")
        metric = args.early_stopping
        callbacks.append(EarlyStopping(monitor=metric,
                                       mode="max" if metric == "val_iou" else "min",
                                       patience=patience,
                                       verbose=args.verbose))

    if args.wandb:
        logger = WandbLogger(name=name,
                             save_dir='/'.join(split_out_dir[:-2]),
                             offline=args.offline,
                             id=args.id,
                             project=project)
        logger.watch(conv_onet)
        if args.visualize:
            callbacks.append(WandbVisualizeCallback())
    else:
        logger = TensorBoardLogger(save_dir=save_dir, name=name)
    logger.log_hyperparams(cfg)

    trainer = pl.Trainer(logger=logger,
                         default_root_dir=out_dir,
                         callbacks=callbacks,
                         gpus=1,
                         enable_progress_bar=not args.no_progress,
                         accumulate_grad_batches=32 // batch_size if batch_size < 32 else None,
                         max_epochs=max_epochs,
                         max_steps=max_steps,
                         check_val_every_n_epoch=int(np.ceil(eval_every_n_epochs)),
                         val_check_interval=validate_every if eval_every_n_epochs < 1.0 else 1.0,
                         log_every_n_steps=print_every,
                         precision=args.precision,
                         profiler="simple" if args.profile else None,
                         auto_lr_find=args.auto_lr,
                         auto_scale_batch_size=args.auto_batch_size)
    trainer.fit(conv_onet,
                dataset,
                ckpt_path=os.path.abspath(args.checkpoint) if args.checkpoint and args.resume else None)
    print(f"Best validation {model_selection_metric}: {trainer.checkpoint_callback.best_model_score:.2f}")

    state_dict = torch.load(trainer.checkpoint_callback.best_model_path)["state_dict"]
    model_best = state_dict.copy()
    for key, value in state_dict.items():
        model_best[key.replace("model.", "")] = model_best.pop(key)
    torch.save({"model": model_best}, os.path.join(out_dir, "model_best.pt"))

    if args.test:
        trainer.test(datamodule=dataset)

if __name__ == "__main__":
    main()


text:
from .base import _get_response

from .. import auc
from .. import roc_curve
from .._base import _check_pos_label_consistency

from ...utils import check_matplotlib_support, deprecated

class RocCurveDisplay:
    """ROC Curve visualization.

    It is recommend to use
    :func:`~sklearn.metrics.RocCurveDisplay.from_estimator` or
    :func:`~sklearn.metrics.RocCurveDisplay.from_predictions` to create
    a :class:`~sklearn.metrics.RocCurveDisplay`. All parameters are
    stored as attributes.

    Read more in the :ref:`User Guide <visualizations>`.

    Parameters
    ----------
    fpr : ndarray
        False positive rate.

    tpr : ndarray
        True positive rate.

    roc_auc : float, default=None
        Area under ROC curve. If None, the roc_auc score is not shown.

    estimator_name : str, default=None
        Name of estimator. If None, the estimator name is not shown.

    pos_label : str or int, default=None
        The class considered as the positive class when computing the roc auc
        metrics. By default, `estimators.classes_[1]` is considered
        as the positive class.

        .. versionadded:: 0.24

    Attributes
    ----------
    line_ : matplotlib Artist
        ROC Curve.

    ax_ : matplotlib Axes
        Axes with ROC Curve.

    figure_ : matplotlib Figure
        Figure containing the curve.

    See Also
    --------
    roc_curve : Compute Receiver operating characteristic (ROC) curve.
    RocCurveDisplay.from_estimator : Plot Receiver Operating Characteristic
        (ROC) curve given an estimator and some data.
    RocCurveDisplay.from_predictions : Plot Receiver Operating Characteristic
        (ROC) curve given the true and predicted values.
    roc_auc_score : Compute the area under the ROC curve.

    Examples
    --------
    >>> import matplotlib.pyplot as plt
    >>> import numpy as np
    >>> from sklearn import metrics
    >>> y = np.array([0, 0, 1, 1])
    >>> pred = np.array([0.1, 0.4, 0.35, 0.8])
    >>> fpr, tpr, thresholds = metrics.roc_curve(y, pred)
    >>> roc_auc = metrics.auc(fpr, tpr)
    >>> display = metrics.RocCurveDisplay(fpr=fpr, tpr=tpr, roc_auc=roc_auc,
    ...                                   estimator_name='example estimator')
    >>> display.plot()
    <...>
    >>> plt.show()
    """

    def __init__(self, *, fpr, tpr, roc_auc=None, estimator_name=None, pos_label=None):
        self.estimator_name = estimator_name
        self.fpr = fpr
        self.tpr = tpr
        self.roc_auc = roc_auc
        self.pos_label = pos_label

    def plot(self, ax=None, *, name=None, **kwargs):
        """Plot visualization

        Extra keyword arguments will be passed to matplotlib's ``plot``.

        Parameters
        ----------
        ax : matplotlib axes, default=None
            Axes object to plot on. If `None`, a new figure and axes is
            created.

        name : str, default=None
            Name of ROC Curve for labeling. If `None`, use `estimator_name` if
            not `None`, otherwise no labeling is shown.

        Returns
        -------
        display : :class:`~sklearn.metrics.plot.RocCurveDisplay`
            Object that stores computed values.
        """
        check_matplotlib_support("RocCurveDisplay.plot")

        name = self.estimator_name if name is None else name

        line_kwargs = {}
        if self.roc_auc is not None and name is not None:
            line_kwargs["label"] = f"{name} (AUC = {self.roc_auc:0.2f})"
        elif self.roc_auc is not None:
            line_kwargs["label"] = f"AUC = {self.roc_auc:0.2f}"
        elif name is not None:
            line_kwargs["label"] = name

        line_kwargs.update(**kwargs)

        import matplotlib.pyplot as plt

        if ax is None:
            fig, ax = plt.subplots()

        (self.line_,) = ax.plot(self.fpr, self.tpr, **line_kwargs)
        info_pos_label = (
            f" (Positive label: {self.pos_label})" if self.pos_label is not None else ""
        )

        xlabel = "False Positive Rate" + info_pos_label
        ylabel = "True Positive Rate" + info_pos_label
        ax.set(xlabel=xlabel, ylabel=ylabel)

        if "label" in line_kwargs:
            ax.legend(loc="lower right")

        self.ax_ = ax
        self.figure_ = ax.figure
        return self

    @classmethod
    def from_estimator(
        cls,
        estimator,
        X,
        y,
        *,
        sample_weight=None,
        drop_intermediate=True,
        response_method="auto",
        pos_label=None,
        name=None,
        ax=None,
        **kwargs,
    ):
        """Create a ROC Curve display from an estimator.

        Parameters
        ----------
        estimator : estimator instance
            Fitted classifier or a fitted :class:`~sklearn.pipeline.Pipeline`
            in which the last estimator is a classifier.

        X : {array-like, sparse matrix} of shape (n_samples, n_features)
            Input values.

        y : array-like of shape (n_samples,)
            Target values.

        sample_weight : array-like of shape (n_samples,), default=None
            Sample weights.

        drop_intermediate : bool, default=True
            Whether to drop some suboptimal thresholds which would not appear
            on a plotted ROC curve. This is useful in order to create lighter
            ROC curves.

        response_method : {'predict_proba', 'decision_function', 'auto'} \
                default='auto'
            Specifies whether to use :term:`predict_proba` or
            :term:`decision_function` as the target response. If set to 'auto',
            :term:`predict_proba` is tried first and if it does not exist
            :term:`decision_function` is tried next.

        pos_label : str or int, default=None
            The class considered as the positive class when computing the roc auc
            metrics. By default, `estimators.classes_[1]` is considered
            as the positive class.

        name : str, default=None
            Name of ROC Curve for labeling. If `None`, use the name of the
            estimator.

        ax : matplotlib axes, default=None
            Axes object to plot on. If `None`, a new figure and axes is created.

        **kwargs : dict
            Keyword arguments to be passed to matplotlib's `plot`.

        Returns
        -------
        display : :class:`~sklearn.metrics.plot.RocCurveDisplay`
            The ROC Curve display.

        See Also
        --------
        roc_curve : Compute Receiver operating characteristic (ROC) curve.
        RocCurveDisplay.from_predictions : ROC Curve visualization given the
            probabilities of scores of a classifier.
        roc_auc_score : Compute the area under the ROC curve.

        Examples
        --------
        >>> import matplotlib.pyplot as plt
        >>> from sklearn.datasets import make_classification
        >>> from sklearn.metrics import RocCurveDisplay
        >>> from sklearn.model_selection import train_test_split
        >>> from sklearn.svm import SVC
        >>> X, y = make_classification(random_state=0)
        >>> X_train, X_test, y_train, y_test = train_test_split(
        ...     X, y, random_state=0)
        >>> clf = SVC(random_state=0).fit(X_train, y_train)
        >>> RocCurveDisplay.from_estimator(
        ...    clf, X_test, y_test)
        <...>
        >>> plt.show()
        """
        check_matplotlib_support(f"{cls.__name__}.from_estimator")

        name = estimator.__class__.__name__ if name is None else name

        y_pred, pos_label = _get_response(
            X,
            estimator,
            response_method=response_method,
            pos_label=pos_label,
        )

        return cls.from_predictions(
            y_true=y,
            y_pred=y_pred,
            sample_weight=sample_weight,
            drop_intermediate=drop_intermediate,
            name=name,
            ax=ax,
            pos_label=pos_label,
            **kwargs,
        )

    @classmethod
    def from_predictions(
        cls,
        y_true,
        y_pred,
        *,
        sample_weight=None,
        drop_intermediate=True,
        pos_label=None,
        name=None,
        ax=None,
        **kwargs,
    ):
        """Plot ROC curve given the true and predicted values.

        Read more in the :ref:`User Guide <visualizations>`.

        .. versionadded:: 1.0

        Parameters
        ----------
        y_true : array-like of shape (n_samples,)
            True labels.

        y_pred : array-like of shape (n_samples,)
            Target scores, can either be probability estimates of the positive
            class, confidence values, or non-thresholded measure of decisions
            (as returned by “decision_function” on some classifiers).

        sample_weight : array-like of shape (n_samples,), default=None
            Sample weights.

        drop_intermediate : bool, default=True
            Whether to drop some suboptimal thresholds which would not appear
            on a plotted ROC curve. This is useful in order to create lighter
            ROC curves.

        pos_label : str or int, default=None
            The label of the positive class. When `pos_label=None`, if `y_true`
            is in {-1, 1} or {0, 1}, `pos_label` is set to 1, otherwise an
            error will be raised.

        name : str, default=None
            Name of ROC curve for labeling. If `None`, name will be set to
            `"Classifier"`.

        ax : matplotlib axes, default=None
            Axes object to plot on. If `None`, a new figure and axes is
            created.

        **kwargs : dict
            Additional keywords arguments passed to matplotlib `plot` function.

        Returns
        -------
        display : :class:`~sklearn.metrics.DetCurveDisplay`
            Object that stores computed values.

        See Also
        --------
        roc_curve : Compute Receiver operating characteristic (ROC) curve.
        RocCurveDisplay.from_estimator : ROC Curve visualization given an
            estimator and some data.
        roc_auc_score : Compute the area under the ROC curve.

        Examples
        --------
        >>> import matplotlib.pyplot as plt
        >>> from sklearn.datasets import make_classification
        >>> from sklearn.metrics import RocCurveDisplay
        >>> from sklearn.model_selection import train_test_split
        >>> from sklearn.svm import SVC
        >>> X, y = make_classification(random_state=0)
        >>> X_train, X_test, y_train, y_test = train_test_split(
        ...     X, y, random_state=0)
        >>> clf = SVC(random_state=0).fit(X_train, y_train)
        >>> y_pred = clf.decision_function(X_test)
        >>> RocCurveDisplay.from_predictions(
        ...    y_test, y_pred)
        <...>
        >>> plt.show()
        """
        check_matplotlib_support(f"{cls.__name__}.from_predictions")

        fpr, tpr, _ = roc_curve(
            y_true,
            y_pred,
            pos_label=pos_label,
            sample_weight=sample_weight,
            drop_intermediate=drop_intermediate,
        )
        roc_auc = auc(fpr, tpr)

        name = "Classifier" if name is None else name
        pos_label = _check_pos_label_consistency(pos_label, y_true)

        viz = RocCurveDisplay(
            fpr=fpr, tpr=tpr, roc_auc=roc_auc, estimator_name=name, pos_label=pos_label
        )

        return viz.plot(ax=ax, name=name, **kwargs)

@deprecated(
    "Function :func:`plot_roc_curve` is deprecated in 1.0 and will be "
    "removed in 1.2. Use one of the class methods: "
    ":meth:`sklearn.metric.RocCurveDisplay.from_predictions` or "
    ":meth:`sklearn.metric.RocCurveDisplay.from_estimator`."
)
def plot_roc_curve(
    estimator,
    X,
    y,
    *,
    sample_weight=None,
    drop_intermediate=True,
    response_method="auto",
    name=None,
    ax=None,
    pos_label=None,
    **kwargs,
):
    """Plot Receiver operating characteristic (ROC) curve.

    Extra keyword arguments will be passed to matplotlib's `plot`.

    Read more in the :ref:`User Guide <visualizations>`.

    Parameters
    ----------
    estimator : estimator instance
        Fitted classifier or a fitted :class:`~sklearn.pipeline.Pipeline`
        in which the last estimator is a classifier.

    X : {array-like, sparse matrix} of shape (n_samples, n_features)
        Input values.

    y : array-like of shape (n_samples,)
        Target values.

    sample_weight : array-like of shape (n_samples,), default=None
        Sample weights.

    drop_intermediate : bool, default=True
        Whether to drop some suboptimal thresholds which would not appear
        on a plotted ROC curve. This is useful in order to create lighter
        ROC curves.

    response_method : {'predict_proba', 'decision_function', 'auto'} \
            default='auto'
        Specifies whether to use :term:`predict_proba` or
        :term:`decision_function` as the target response. If set to 'auto',
        :term:`predict_proba` is tried first and if it does not exist
        :term:`decision_function` is tried next.

    name : str, default=None
        Name of ROC Curve for labeling. If `None`, use the name of the
        estimator.

    ax : matplotlib axes, default=None
        Axes object to plot on. If `None`, a new figure and axes is created.

    pos_label : str or int, default=None
        The class considered as the positive class when computing the roc auc
        metrics. By default, `estimators.classes_[1]` is considered
        as the positive class.

    **kwargs : dict
        Additional keywords arguments passed to matplotlib `plot` function.

        .. versionadded:: 0.24

    Returns
    -------
    display : :class:`~sklearn.metrics.RocCurveDisplay`
        Object that stores computed values.

    See Also
    --------
    roc_curve : Compute Receiver operating characteristic (ROC) curve.
    RocCurveDisplay.from_estimator : ROC Curve visualization given an estimator
        and some data.
    RocCurveDisplay.from_predictions : ROC Curve visualisation given the
        true and predicted values.
    roc_auc_score : Compute the area under the ROC curve.

    Examples
    --------
    >>> import matplotlib.pyplot as plt
    >>> from sklearn import datasets, metrics, model_selection, svm
    >>> X, y = datasets.make_classification(random_state=0)
    >>> X_train, X_test, y_train, y_test = model_selection.train_test_split(
    ...     X, y, random_state=0)
    >>> clf = svm.SVC(random_state=0)
    >>> clf.fit(X_train, y_train)
    SVC(random_state=0)
    >>> metrics.plot_roc_curve(clf, X_test, y_test) # doctest: +SKIP
    <...>
    >>> plt.show()
    """
    check_matplotlib_support("plot_roc_curve")

    y_pred, pos_label = _get_response(
        X, estimator, response_method, pos_label=pos_label
    )

    fpr, tpr, _ = roc_curve(
        y,
        y_pred,
        pos_label=pos_label,
        sample_weight=sample_weight,
        drop_intermediate=drop_intermediate,
    )
    roc_auc = auc(fpr, tpr)

    name = estimator.__class__.__name__ if name is None else name

    viz = RocCurveDisplay(
        fpr=fpr, tpr=tpr, roc_auc=roc_auc, estimator_name=name, pos_label=pos_label
    )

    return viz.plot(ax=ax, name=name, **kwargs)