text:
import torch
import scipy as sp
import networkx as nx

EPS = 1e-5

if torch.cuda.is_available():
    torch.set_default_tensor_type(torch.cuda.FloatTensor)

def get_laplacian_matrix(adj, normalize_L):
    r"""
    Get the Laplacian/normalized Laplacian matrices from a batch of adjacency matrices
    Parameters
    --------------
        adj: tensor(..., N, N)
            Batches of symmetric adjacency matrices

        normalize_L: bool
            Whether to normalize the Laplacian matrix
            If `False`, then `L = D - A`
            If `True`, then `L = D^-1 (D - A)`
    Returns
    -------------
        L: tensor(..., N, N)
            Resulting Laplacian matrix
    """

    # Apply the equation L = D - A
    N = adj.shape[-1]
    arr = torch.arange(N)
    L = -adj
    D = torch.sum(adj, dim=-1)
    L[..., arr, arr] = D

    # Normalize by the degree : L = D^-1 (D - A)
    if normalize_L:
        Dinv = torch.zeros_like(L)
        Dinv[..., arr, arr] = D ** -1
        L = torch.matmul(Dinv, L)

    return L

def get_k_lowest_eig(adj, k):
    r"""
    Compute the k-lowest eigenvectors of the Laplacian matrix
    for each connected components of the graph. If there are disconnected
    graphs, then the first k eigenvectors are computed for each sub-graph
    separately.
    Parameters
    --------------
        adj: tensor(..., N, N)
            Batches of symmetric adjacency matrices
        k: int
            Compute the k-th smallest eigenvectors and eigenvalues.

        normalize_L: bool
            Whether to normalize the Laplacian matrix
            If `False`, then `L = D - A`
            If `True`, then `L = D^-1 (D - A)`
    Returns
    -------------
        eigvec: tensor(..., N, k)
            Resulting k-lowest eigenvectors of the Laplacian matrix of each sub-graph,
            with the same batching as the `adj` tensor.
            The dim==-1 represents the k-th vectors.
            The dim==-2 represents the N elements of each vector.
            If the a given graph is disconnected, it will give the first ``k`` eigenvector
            of each sub-graph, and will force the first eigenvector to be 0-vectors.
            If there are ``m`` eigenvectors for a given sub-graph, with ``m < k``, it will
            return 0-vectors for all eigenvectors ``> m``
    """

    # Reshape as a 3D tensor for easier looping along batches
    device = adj.device
    shape = list(adj.shape)
    if adj.ndim == 2:
        adj = adj.unsqueeze(0)
    elif adj.ndim > 3:
        adj = adj.view(-1, shape[-2], shape[-1])

    L = get_laplacian_matrix(adj, normalize_L=False)

    # Compute and sort the eigenvectors

    eigval_all, eigvec_all = torch.symeig(L.cpu(), eigenvectors=True)
    eigval_all = eigval_all.to(device)
    eigvec_all = eigvec_all.to(device)
    sort_idx = torch.argsort(eigval_all.abs(), dim=-1, descending=False)
    sort_idx_vec = sort_idx.unsqueeze(-2).expand(eigvec_all.shape)
    eigval_sort = torch.gather(eigval_all, dim=-1, index=sort_idx)
    eigvec_sort = torch.gather(eigvec_all, dim=-1, index=sort_idx_vec)

    k_lowest_eigvec = []

    # Loop each graph to detect if some of them are disconnected. If they are disconnected,
    # then modify the eigenvectors such that the lowest k eigenvectors are returned for
    # each sub-graph
    for ii in range(adj.shape[0]):
        this_eigval = eigval_sort[ii]
        num_connected = torch.sum(this_eigval.abs() < EPS)

        # If there is a single connected graph, then return the k lowest eigen functions
        if num_connected <= 1:
            this_eigvec = eigvec_sort[ii, :, :k]
            if k > this_eigvec.shape[-1]:
                temp_eigvec = torch.zeros(this_eigvec.shape[0], k)
                temp_eigvec[:, :k] = this_eigvec
                this_eigvec = temp_eigvec
            k_lowest_eigvec.append(this_eigvec)

        # Otherwise, return the k lowest eigen functions for each sub-graph
        elif num_connected > 1:
            eigvec0 = eigvec_sort[ii, :, :num_connected]
            unique_idx = torch.zeros(1)
            factor = 100

            # Use the eigenvectors with 0 eigenvalues to find the unique sub-graphs
            # And loop to make sure the number of detected sub-graphs is consistent with the
            # Number of connected sub-graphs.
            while (max(unique_idx) + 1) != num_connected:
                eigvec0_round = torch.round(eigvec0 / (factor * EPS))
                _, unique_idx = torch.unique(eigvec0_round, return_inverse=True, dim=0)

                if (max(unique_idx) + 1) < num_connected:
                    factor = (factor / 2)
                elif (max(unique_idx) + 1) > num_connected:
                    factor = (factor * 3)

            # Find the eigenvectors associated to each sub-graph
            sub_graph_factors = torch.zeros(num_connected, len(this_eigval))
            for sub_ii in range(num_connected):
                sub_idx = torch.where(unique_idx == sub_ii)[0]
                sub_graph_factors[sub_ii, :] = torch.mean(torch.abs(eigvec_sort[ii, sub_idx, :]), dim=-2)
            max_idx = torch.argmax(sub_graph_factors, dim=0)[num_connected:]

            # Concatenate the k lowest eigenvectors of each sub-graph
            this_k_lowest_eigvec = torch.zeros(len(this_eigval), k)
            for sub_ii in range(num_connected):
                sub_idx = torch.where(unique_idx == sub_ii)[0]
                k_lowest_idx = torch.where(max_idx == sub_ii)[0][:k-1] + num_connected
                for kk_enum, kk in enumerate(k_lowest_idx):
                    this_k_lowest_eigvec[sub_idx, kk_enum+1] = eigvec_sort[ii, sub_idx, kk]

            k_lowest_eigvec.append(this_k_lowest_eigvec)

    # Stack and Reshape to match the input batch shape
    k_lowest_eigvec = torch.stack(k_lowest_eigvec, dim=0).view(*(shape[:-2] + [-1, k]))

    return k_lowest_eigvec

def get_k_lowest_eig_old(adj, k, normalize_L):
    r"""
    Compute the k-lowest eigenvalues and eigenvectors of the Laplacian matrix
    for each connected components of the graph.
    Parameters
    --------------
        adj: tensor(..., N, N)
            Batches of symmetric adjacency matrices
        k: int
            Compute the k-th smallest eigenvectors and eigenvalues.

        normalize_L: bool
            Whether to normalize the Laplacian matrix
            If `False`, then `L = D - A`
            If `True`, then `L = D^-1 (D - A)`
    Returns
    -------------
        eigval: tensor(..., k)
            Resulting k-lowest eigenvalues of the Laplacian matrix,
            with the same batching as the `adj` tensor
        eigvec: tensor(..., N, k)
            Resulting k-lowest eigenvectors of the Laplacian matrix,
            with the same batching as the `adj` tensor.
            The dim==-1 represents the k-th vectors.
            The dim==-2 represents the N elements of each vector.
    """

    device = adj.device
    # Reshape as a 3D tensor for easier looping along batches
    shape = list(adj.shape)
    if adj.ndim == 2:
        adj = adj.unsqueeze(0)
    elif adj.ndim > 3:
        adj = adj.view(-1, shape[-2], shape[-1])

    L = get_laplacian_matrix(adj, normalize_L=False)

    eigvec = torch.zeros(adj.shape[0],adj.shape[-1],k).to(device)
    for ii in range(adj.shape[0]):
        #get singular values decomposition
        u,s,v = torch.svd(L[ii])
        #get 0 singular values positions in svd
        null_mask = (s <= EPS)
        #get basis of 0 singular space
        zero_eig_vect = torch.transpose(v[null_mask],0,1)
        assert zero_eig_vect.shape[0] == adj.shape[-1]

        components = []
        for jj in range(zero_eig_vect.shape[-1]):
            #get the indices of the connected components by getting nonzero entries in the
            #0 singular vector associated to the component
            compo = torch.nonzero(zero_eig_vect[:,jj]).squeeze()
            components.append(compo)

        #get the eigenvalues and eigenvectors of each components
        ii_eigvec = torch.zeros(adj.shape[-1],k).to(device)
        for jj, compo in enumerate(components):
            #select the sub matrix corresponding to the component
            sub_adj = adj[ii,:,:][compo].T[compo].T

            #if a component has only one node, all eigen vectors are 0
            if sub_adj.shape==torch.Size([]):
                compo_eigvec = torch.zeros(1,k)
                ii_eigvec[compo,:] = compo_eigvec
            else:
                dim = sub_adj.shape[0]
                #if the components Laplacian has less eigenvectors than k, add zero vectors in padding
                if k > dim:
                    compo_eigval, compo_eigvec = get_k_lowest_eig(sub_adj, dim, normalize_L)
                    compo_eigvec = torch.cat((compo_eigvec, torch.zeros(1, dim, k-dim, device=device)), dim=2)

                else:
                    compo_eigval, compo_eigvec = get_k_lowest_eig_connected(sub_adj, k, normalize_L)

                #place the k eigenvectors of each component in a vector using the indices of each nodes
                #of the components nodes
                ii_eigvec[compo,:] = compo_eigvec
        eigvec[ii,:,:]=ii_eigvec

    #makes sense for connected graphs but not for disconnected ones so it is left empty. Remove later
    eigval = []

    return eigval, eigvec

def get_k_lowest_eig_connected_old(adj, k, normalize_L):
    r"""
    Compute the k-lowest eigenvalues and eigenvectors of the Laplacian matrix of a connected graph.
    Parameters
    --------------
        adj: tensor(..., N, N)
            Batches of symmetric adjacency matrices
        k: int
            Compute the k-th smallest eigenvectors and eigenvalues.

        normalize_L: bool
            Whether to normalize the Laplacian matrix
            If `False`, then `L = D - A`
            If `True`, then `L = D^-1 (D - A)`
    Returns
    -------------
        eigval: tensor(..., k)
            Resulting k-lowest eigenvalues of the Laplacian matrix,
            with the same batching as the `adj` tensor
        eigvec: tensor(..., N, k)
            Resulting k-lowest eigenvectors of the Laplacian matrix,
            with the same batching as the `adj` tensor.
            The dim==-1 represents the k-th vectors.
            The dim==-2 represents the N elements of each vector.
    """

    # Reshape as a 3D tensor for easier looping along batches
    device = adj.device
    shape = list(adj.shape)
    if adj.ndim == 2:
        adj = adj.unsqueeze(0)
    elif adj.ndim > 3:
        adj = adj.view(-1, shape[-2], shape[-1])

    L = get_laplacian_matrix(adj, normalize_L=normalize_L)
    eigvec = []
    eigval = []

    # Compute all the k-th smallest eigenvectors and eigenvalues
    for ii in range(L.shape[0]):
        Laplacian = sp.sparse.csr_matrix(L[ii, :, :].clone().detach().cpu().numpy())
        w, v = sp.sparse.linalg.eigsh(Laplacian, k=k, which='SM', tol=0)
        eigval.append(torch.Tensor(w).to(device))
        eigvec.append(torch.Tensor(v).to(device))

    # Stack and Reshape to match the input batch shape
    eigval = torch.stack(eigval, dim=0).view(*(shape[:-2] + [k]))
    eigvec = torch.stack(eigvec, dim=0).view(*(shape[:-2] + [-1, k]))

    return eigval, eigvec

def get_adjacency_from_gradient_of_features(adj, features, normalization='none', add_diag=True, absolute_adj=False):
    r"""
    Computes an adjacency matrix from the gradient of the ``features`` on the graph described
    by the adjacency matrix ``adj``, with different options of normalization.
    The gradient is a function on the edges, computed from these node features.
    For 2 nodes ``n_i, n_j`` with node features ``f_i, f_j`` and an edge weight of ``w_ij``,
    the gradient on edge ``e_ij`` is given by ``e_ij = w_ij * (f_j - f_i)``.
    If ``X`` is a function of the nodes, then ``matmul(grad_adj, X)`` behaves as the derivative of
    ``X`` in the direction of the gradient of ``features``.
    Parameters
    --------------
        adj: tensor(..., N, N)
            Batches of adjacency matrices of graphs
        features: tensor(..., N)
            Batches of features on the nodes. These features are used to compute the gradient
        normalization: str, Optional
            Normalization strategy for the ``grad_adj`` matrix.
            - 'none': No normalization is applied

            - 'row-abs': Normalize such that the absolute sum of each row of ``grad_adj`` is 1.
              All values below the global variable ``EPS=1e-5`` will be set to 0.
            - 'in-out-field': Normalize such that the sum of the norm of the input and output field
              is 1. The input field is the negative values in the ``grad_adj``, while the
              output field is the positive values. The norm is computed as the root of the squared sum.
              All values below the global variable ``EPS=1e-5`` will be set to 0.
              For example, if there are no positive values above EPS for a specific node, then
              the output field of that specific node will be set to 0.
            (Default='none')

        add_diag: bool, Optional
            Whether to add a diagonal element such as each row of ``grad_adj`` sums to 0.
            The diagonal is added after the normalization.

            - If True, ``matmul(grad_adj, X)`` will behave like a forward/backward derivative
              at the local minima/maxima from the features function.
            - If False, ``matmul(grad_adj, X)`` will behave similarily to a center derivative
              with a zero-padding added before the local minima/maxima
            (Default=True)
        absolute_adj: bool, Optional
            Return the absolute value of ``grad_adj``.
            ``matmul(grad_adj, X)`` becomes a directional smoothing instead of a directional
            derivative.
            (Default=False)
    Returns
    -------------
        grad_adj: tensor(..., N, N)
            Batches of adjacency matrices representing the gradient of the features on the graph.

    """
    device = adj.device
    # Compute the gradient of the features
    grad_adj = (adj * (features.unsqueeze(-2) - features.unsqueeze(-1) + EPS))

    # Don't normalize
    if (normalization is None) or (normalization.lower() == 'none'):
        pass

    # Normalize by the sum of each row
    elif normalization.lower() == 'row-abs':
        grad_norm = torch.abs(grad_adj)
        grad_norm *= grad_norm > EPS
        grad_adj = grad_adj / (torch.sum(grad_norm, dim=-1, keepdims=True) + EPS)

    # Normalize by the input and output vector field
    elif normalization.lower() == 'in-out-field':
        pos_adj = grad_adj * (grad_adj > EPS)
        neg_adj = grad_adj * (grad_adj < -EPS)
        out_field = torch.sqrt(torch.sum(pos_adj ** 2, dim=-1, keepdims=True)) + EPS
        in_field = torch.sqrt(torch.sum(neg_adj ** 2, dim=-1, keepdims=True)) + EPS
        in_out_field = out_field + in_field
        grad_adj = (pos_adj / in_out_field) + (neg_adj / in_out_field)

    else:
        raise ValueError('Unsupported normalization option `{}`')

    # Add an element on the diagonal to make each row sum to 0
    if add_diag:
        eye = torch.eye(adj.shape[-1]).to(device)[(None, )*(adj.ndim-2)] # Unsqueeze multiple times
        grad_adj = grad_adj - (eye * torch.sum(grad_adj, dim=-1, keepdims=True))

    # Compute the absolute value
    if absolute_adj:
        grad_adj = torch.abs(grad_adj)

    return grad_adj

def get_eig_adjacency(adj, eig_idx, normalization='none', add_diag=True, absolute_adj=False, normalize_L=False, eig_acos=True):
    r"""
    Computes an adjacency matrix from the gradient of the eigenvectors of the graph Laplacian,
    with the graph being described by the adjacency matrix ``adj``.
    The gradient is a function on the edges, computed from these node features.
    For 2 nodes ``n_i, n_j`` with node features ``f_i, f_j`` and an edge weight of ``w_ij``,
    the gradient on edge ``e_ij`` is given by ``e_ij = w_ij * (f_j - f_i)``.
    If ``X`` is a function of the nodes, then ``matmul(grad_adj, X)`` behaves as the derivative of
    ``X`` in the direction of the gradient of ``features``.
    Parameters
    --------------
        adj: tensor(..., N, N)
            Batches of adjacency matrices of graphs. It is used to compute the Laplacian matrix,
            and the eigenvectors of the Laplacian.
        eig_idx: int, iterator(int)
            Indexes of the eigenvectors to compute, sorted by smallest eigenvalues.
            The current function is efficient for low indexes.
            e.g. ``eig_idx=[1, 2, 4]`` will compute the eigenvectors ``[0, 1, 2, 3, 4]``, but will
            then ignore the eigenvectors 0 and 3.
            If ``0`` is in the ``eig_idx``, then the normalized adjacency matrix is returned for
            that specific index, no matter the normalization chosen.
        normalization: str, Optional
            Normalization strategy for the ``grad_adj`` matrix. It does not apply to eig_idx==0.
            - 'none': No normalization is applied

            - 'row-abs': Normalize such that the absolute sum of each row of ``grad_adj`` is 1.
              All values below the global variable ``EPS=1e-5`` will be set to 0.
            - 'in-out-field': Normalize such that the the norm of the input and output field
              are 1. The input field is the negative values in the ``grad_adj``, while the
              output field is the positive values. The norm is computed as the root of the squared sum.
              All values below the global variable ``EPS=1e-5`` will be set to 0.
              For example, if there are no positive values above EPS for a specific node, then
              the output field of that specific node will be set to 0.
            (Default='none')

        add_diag: bool, Optional
            Whether to add a diagonal element such as each row of ``grad_adj`` sums to 0.
            The diagonal is added after the normalization. It does not apply to eig_idx==0.

            - If True, ``matmul(grad_adj, X)`` will behave like a forward/backward derivative
              at the local minima/maxima from the features function.
            - If False, ``matmul(grad_adj, X)`` will behave similarily to a center derivative
              with a zero-padding added before the local minima/maxima
            (Default=True)
        absolute_adj: bool, Optional
            Return the absolute value of ``grad_adj``.
            ``matmul(grad_adj, X)`` becomes a directional smoothing instead of a directional
            derivative.
            (Default=False)
        normalize_L: bool, Optional
            Whether to normalize the Laplacian matrix
            If `False`, then `L = D - A`
            If `True`, then `L = D^-1 (D - A)`
            (Default=False)

        eig_acos: bool, Optional
            Whether to compute the arcosine of the eigenvectors instead of the eigenvector.
            This will sort-of `linearize` the eigenvector. If ``normalize=='in-out-field'``,
            this parameter shouldn't change anything, appart from numerical error. For other
            normalization, it will have an important impact, especially near the borders
            of the graph.
    Returns
    -------------
        eig_adj: dict(tensor(..., N, N))
            Dictionary of Batches of adjacency matrices representing the gradient of
            the eigenvectors on the graph. Each key corresponds to a specific ``eig_idx``,
            and the value corresponds to the associated batch of eigen-adjacencies.

    """

    # Convert the eig_idx into a list
    try:
        eig_idx = list(eig_idx)
    except:
        eig_idx = [eig_idx]

    # Generate an adjacency matrix for each eigenvectors from `eig_idx`
    eigvec = get_k_lowest_eig(adj, max(eig_idx)+1, normalize_L=normalize_L)
    eig_adj = {}
    for ii in eig_idx:
        if ii != 0:
            this_eigvec = eigvec[..., ii]
            if eig_acos:
                this_eigvec = torch.acos(this_eigvec / torch.max(torch.abs(this_eigvec)))
            w_adj = get_adjacency_from_gradient_of_features(adj, features=this_eigvec,
                        normalization=normalization, add_diag=add_diag, absolute_adj=absolute_adj)
            eig_adj[ii] = w_adj

        else:
            eig_adj[ii] = adj / (torch.sum(adj.abs(), dim=-1, keepdims=True) + EPS)

    return eig_adj

def aggregate_sum(X, adj, self_loop=False, device='cpu', avg_d=None):
    """
    Aggregate each node by the sum of it's neighbours, weighted by the edge between the nodes.
    Parameters
    --------------
        X: tensor(..., N, N, Din)
            Input feature tensor to aggregate.
        adj: tensor(..., N, N)
            Batches of adjacency matrices of graphs.
    Returns
    -------------
        X_sum: tensor(..., N, N, Din)
            Aggregation results

    """

    if self_loop:  # add self connections
        N = adj.shape[-1]
        adj = adj + torch.eye(N, device=device).unsqueeze(0)

    X_sum = torch.sum(torch.mul(X, adj.unsqueeze(-1)), dim=-2)
    return X_sum

def aggregate_eigs(X, adj, eig_idx, normalization='none', add_diag=True, agg_type='derivative', normalize_L=False,
                    eig_acos=True, self_loop=False, device='cpu', avg_d=None):
    r"""
    Aggregate the features ``X`` in the direction of the gradient of the Laplacian
    eigenvectors ``eig_idx`` such as the indexes are sorted by lowest eigenvalues.
    Parameters
    --------------
        X: tensor(..., N, N, Din)
            Input feature tensor to aggregate.
        adj: tensor(..., N, N)
            Batches of adjacency matrices of graphs. It is used to compute the Laplacian matrix,
            and the eigenvectors of the Laplacian.
        eig_idx: int, iterator(int)
            Indexes of the eigenvectors to compute, sorted by smallest eigenvalues.
            The current function is efficient for low indexes.
            e.g. ``eig_idx=[1, 2, 4]`` will compute the eigenvectors ``[0, 1, 2, 3, 4]``, but will
            then ignore the eigenvectors 0 and 3.
            If ``0`` is in the ``eig_idx``, then a mean aggregation is done for that specific index.
        normalization: str, Optional
            Normalization strategy for the ``grad_adj`` matrix.
            - 'none': No normalization is applied

            - 'row-abs': Normalize such that the absolute sum of each row of ``grad_adj`` is 1.
              All values below the global variable ``EPS=1e-5`` will be set to 0.
            - 'in-out-field': Normalize such that the the norm of the input and output field
              are 1. The input field is the negative values in the ``grad_adj``, while the
              output field is the positive values. The norm is computed as the root of the squared sum.
              All values below the global variable ``EPS=1e-5`` will be set to 0.
              For example, if there are no positive values above EPS for a specific node, then
              the output field of that specific node will be set to 0.
            (Default='none')

        add_diag: bool, Optional
            Whether to add a diagonal element such as each row of ``grad_adj`` sums to 0.
            The diagonal is added after the normalization.

            - If True, ``matmul(grad_adj, X)`` will behave like a forward/backward derivative
              at the local minima/maxima from the features function.
            - If False, ``matmul(grad_adj, X)`` will behave similarily to a center derivative
              with a zero-padding added before the local minima/maxima
            (Default=True)
        agg_type: str, Optional
            Aggregation options
            - 'derivative': Use the eigen-adjacency to compute the directional derivative
              of the signal, then compute the absolute value of the derivative.
            - 'smoothing': Use the absolute value of the eigen-adjacency to compute the
              directional smoothing of the signal.
            - 'both': Use both methods.

            (Default='derivative')
        normalize_L: bool, Optional
            Whether to normalize the Laplacian matrix
            If `False`, then `L = D - A`
            If `True`, then `L = D^-1 (D - A)`
            (Default=False)

        eig_acos: bool, Optional
            Whether to compute the arcosine of the eigenvectors instead of the eigenvector.
            This will sort-of `linearize` the eigenvector. If ``normalize=='in-out-field'``,
            this parameter shouldn't change anything, appart from numerical error. For other
            normalization, it will have an important impact, especially near the borders
            of the graph.
        self_loop: bool, Optional
            Whether to add a self-loop in the aggregation
            (Default=False)
        device:
            Torch device used to run the computation
            (Default='cpu')
        avg_d: Ignored
            Ignored parameter. Here for compatibility
    Returns
    -------------
        X_agg: tensor(..., N, N, len(eig_idx)*Din) or tensor(..., N, N, len(eig_idx)*2*Din)
            Aggregation results for each eigenvector, concatenated alongside the last dimension
            The last dimension is of size ``len(eig_idx)*2*Din`` when ``agg_type=='both'``.

    """

    agg_type = agg_type.lower()
    if agg_type not in ['derivative', 'smoothing', 'both']:
        raise ValueError(f'Unknown agg_type "{agg_type}"')

    # Get dictionary of adjacency matrices
    adj_dict = get_eig_adjacency(adj, eig_idx, normalization=normalization, add_diag=add_diag,
                                absolute_adj=False, normalize_L=normalize_L, eig_acos=eig_acos)

    X_agg = []
    for ii, this_adj in adj_dict.items():
        # Compute derivative
        if ((agg_type == 'derivative') or (agg_type == 'both')) and (ii != 0):
            this_agg = aggregate_sum(X, this_adj, self_loop=self_loop, device=device)
            X_agg.append(this_agg)

        # Compute smoothing
        if (agg_type == 'smoothing') or (agg_type == 'both') or (ii == 0):
            this_agg = aggregate_sum(X, this_adj.abs(), self_loop=self_loop, device=device)
            X_agg.append(this_agg)

    return torch.cat(X_agg, dim=-1)

max_stars_repo_name: xinzhou-ntu/DGNs

max_issues_repo_name: xinzhou-ntu/DGNs

max_forks_repo_name: xinzhou-ntu/DGNs

reason: import scipy,import networkx


text:
# get overview of data
# read data from given files, and produce a reduced image of each data file

import argparse
import logging
import math
import os
import shutil
import subprocess
from typing import Any, Dict, List, Optional

import h5py
import nibabel as nib
import numpy as np
from PIL import Image
from tqdm import tqdm

from AssistedVolumeSegmentation.common import (
    default_seg_file,
    get_file_list,
    get_full_path,
    get_source_data_path,
    init_logging,
    load_config,
    overview_bound_size,
    write_annot_file,
)

def reduce_source_data(
    config: Dict[str, Any],
    bound_size: List[int],
    subdir_num: Optional[int],
    launch_editor: bool,
    reduce_annotation: bool = False,
    source_format: Optional[str] = None,
    data_path: Optional[str] = None,
    write_data_path: Optional[str] = None,
):
    """
    Read source data, from either a stack of tiff files from the given dir or an HDF5 data file, and resize the
    source volume data to the given size.

    :param Dict[str, Any] config: Configuration map
    :param List[int] bound_size: Maximum size of the volume, as (height, width, layers)
    :param Optional[int] subdir_num: Number of subdirectory data to read
    :param bool launch_editor: If true, launch slicer editor using the produced overview data
    :param bool reduce_annotation: If true, perform reduction of a generated segmentation.  This requires source_format,
           data_path and write_data_path to be defined
    :param Optional[str] source_format: Format of source data (read from config if not specified)
    :param Optional[str] data_path: Path of source data (read from config if not specified)
    :param Optional[str] write_data_path: Path to write reduced data to (read from config if not specified)
    :return: Array of source images reduced to given size
    """
    if source_format is None:
        source_format = config["source_data_format"][subdir_num]
    read_tiff_stack = source_format == "tiff-stack"

    # select source data path according to index in config file. assume source data path is an array
    if data_path is None:
        assert subdir_num is not None
        data_path = get_source_data_path(config, subdir_num)
    if write_data_path is None:
        write_data_path = get_full_path(
            config, subdir_num, "overview_reduced_data"
        )
    assert source_format is not None and data_path is not None
    if reduce_annotation:
        logging.info(
            "Reducing generated annotation, reading from source %s, format %s, writing to %s, bound size %s"
            % (
                data_path,
                source_format,
                write_data_path,
                bound_size,
            )
        )
    else:
        write_coverage_path = get_full_path(
            config, subdir_num, "overview_coverage"
        )
        logging.info(
            "Producing overview, reading from source %s, format %s, writing to %s, %s, bound size %s"
            % (
                data_path,
                source_format,
                write_data_path,
                write_coverage_path,
                bound_size,
            )
        )

    if read_tiff_stack:
        # get the list of files
        file_list = get_file_list(data_path)
        input_stacks = len(file_list)

        # read the first image to get dimensions.  assume images are the same size
        first_img = Image.open(file_list[0])
        input_w, input_h = first_img.size  # (width, height)
        input_dims = (input_h, input_w)
        input_count = input_stacks
        iterate_dim = 2
    else:
        # open as hdf5 file
        h5_file = h5py.File(data_path, "r")
        h5_data_key = config["source_hdf5_dataset_name"]
        h5_data = h5_file[h5_data_key]

        # with H5 data, iterate over the first dimension (x), which is different to when reading from image stacks.
        # it is much faster to slice on the initial dimension(s) than on later ones (eg z).
        # find dimensions of data.  if 4 dims assume format is (layers, x, y, z), and only read from first layer.  don't try and
        # index the layer at the start as it would load the whole array in memory, perform indexing with each iteration
        if h5_data.ndim == 4:
            input_size = h5_data.shape[1:]
        else:
            input_size = h5_data.shape
        input_count = input_size[0]
        input_dims = input_size[1:3]
        iterate_dim = 0

    # find the downsample factor to fit within the given bounds.  only allow reducing size
    ratios = [min(1.0, x / y) for x, y in zip(bound_size, input_dims)]
    reduce_ratio = np.min(ratios)

    # find target image size
    target_imgs = math.ceil(input_count * reduce_ratio)
    target_dims = tuple([math.ceil(x * reduce_ratio) for x in input_dims])

    def find_frame_score(input_idx, div_steps, out_slice):
        lower_bound = max(input_idx, div_steps[out_slice])
        upper_bound = min(input_idx + 1, div_steps[out_slice + 1])
        return upper_bound - lower_bound

    # read slices in turn
    logging.info(
        "found %d images(slices), input dim (w %d, h %d), target dim %s, %d slices"
        % (input_count, input_dims[1], input_dims[0], target_dims, target_imgs)
    )
    input_slices = {}
    division_steps = np.arange(target_imgs + 1) * (1.0 / reduce_ratio)
    output_slice = 0
    range_start = 0
    range_end = math.ceil(division_steps[1])
    output_slices = []
    for count in tqdm(range(input_count)):
        if read_tiff_stack:
            this_file = file_list[count]
            this_img = Image.open(this_file)
        else:
            if h5_data.ndim == 4:
                this_img = Image.fromarray(h5_data[0, count])
            else:
                this_img = Image.fromarray(h5_data[count])

        resized_img = this_img.resize(target_dims, Image.NEAREST)
        input_slices[count] = resized_img

        # check if we have enough slices to output the next slice
        if count >= range_end - 1:
            slices = np.stack(
                [
                    np.array(input_slices[x])
                    for x in range(range_start, range_end)
                ]
            )  # (slices, width, height)
            slice_weights = np.array(
                [
                    find_frame_score(x, division_steps, output_slice)
                    for x in range(range_start, range_end)
                ]
            )  # (slices,)

            reduced_slice = (slices * slice_weights[:, None, None]).sum(
                axis=0
            ) / slice_weights.sum()
            output_slices.append(reduced_slice)

            # update output slice number and input slice range
            output_slice += 1
            if output_slice < target_imgs:
                range_start = math.floor(division_steps[output_slice])
                range_end = math.ceil(division_steps[output_slice + 1])

                # remove stored input slices outside of range
                remove_input_frames = [
                    x
                    for x in input_slices.keys()
                    if x < range_start or x > range_end
                ]
                for x in remove_input_frames:
                    input_slices.pop(x)

    # write output slices
    output_array = np.stack(output_slices, axis=iterate_dim)
    if reduce_annotation:
        # discretise results and write in annotation format
        discrete_array = (output_array > 0.5).astype("int")

        # define segment fields.  only a single segment is given as output in current method
        sample_fields = {
            "Segment0_Color": "0.525666 0.813434 0.324",
            "Segment0_ColorAutoGenerated": "1",
            "Segment0_Extent": "0 %d 0 %d 0 %d"
            % tuple((np.array(output_array.shape) - 1).tolist()),
            "Segment0_ID": "Generated_0",
            "Segment0_LabelValue": "1",
            "Segment0_Layer": "0",
            "Segment0_Name": "Generated_0",
            "Segment0_NameAutoGenerated": "1",
            "Segment0_Tags": "Segmentation.Status:inprogress|TerminologyEntry:Segmentation category and type - 3D Slicer General Anatomy list~SRT^T-D0050^Tissue~SRT^T-D0050^Tissue~^^~Anatomic codes - DICOM master list~^^~^^|",
        }

        scales = np.identity(3, dtype="float")
        write_annot_file(
            index_name=None,
            write_path=None,
            annot_data=discrete_array,
            annot_fields=sample_fields,
            scales=scales,
            annot_write_path=write_data_path,
        )
        logging.info(
            "Created reduced segmentation in: %s" % (write_data_path,)
        )
        launch_args = [write_data_path]

    else:
        # write out reduced volume data
        # create Nifti1 format object from array
        # define affine transform matrix representing scaling transform
        aff_matrix = np.identity(4, dtype="float")
        np.fill_diagonal(aff_matrix, [1.0 / reduce_ratio] * 3)
        nifti_object = nib.Nifti1Image(output_array, aff_matrix)
        overviews_path = get_full_path(config, subdir_num, "overviews_path")
        os.makedirs(overviews_path, exist_ok=True)
        nib.save(nifti_object, write_data_path)

        # create sample coverage segmentation file
        if not os.path.exists(write_coverage_path):
            script_path = os.path.dirname(os.path.abspath(__file__))
            default_file = os.path.join(
                script_path, os.pardir, default_seg_file
            )
            shutil.copyfile(default_file, write_coverage_path)
            logging.info(
                "Copying initial coverage annotation from %s to %s"
                % (default_file, write_coverage_path)
            )

        logging.info(
            "Created overview data in: %s , and empty coverage annotation in: %s"
            % (write_data_path, write_coverage_path)
        )
        launch_args = [write_data_path, write_coverage_path]

    args = [config["slicer_path"]] + launch_args
    launch_command = " ".join(args)
    if launch_editor:
        logging.info("Launching Slicer editor with arguments: %s" % args)
        subprocess.run(launch_command, shell=True)
    else:
        logging.info(
            "You can open this in Slicer manually, or launch it with: %s"
            % (launch_command,)
        )

def main():
    init_logging()
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-c", "--config_file", help="Project config file", required=True
    )
    parser.add_argument(
        "-s", "--subdir", help="Data subdirectory number", required=False
    )
    parser.add_argument(
        "-l",
        "--launch",
        help="Launch Slicer to edit piece",
        action="store_true",
    )
    parser.add_argument(
        "-g",
        "--generated_data",
        help="Generated annotation to reduce",
        required=False,
    )
    parser.add_argument(
        "-o",
        "--generated_data_output",
        help="Output file for reducing generated data",
        required=False,
    )

    args = parser.parse_args()
    config = load_config(args.config_file)

    if args.generated_data is None:
        # produce overview, read source data from tiff files in given directory or an HDF5 file,
        # and downsample to given size
        if not isinstance(args.subdir, str) or not args.subdir.isnumeric():
            raise RuntimeError("Subdir should be a number")
        subdir_num = int(args.subdir)
        reduce_source_data(
            config,
            overview_bound_size,
            subdir_num,
            args.launch,
        )
    else:
        # perform reduction of generated annotation data
        assert args.generated_data_output is not None
        if os.path.exists(args.generated_data_output):
            raise RuntimeError(
                "Generated data output already exists: %s"
                % args.generated_data_output
            )

        reduce_source_data(
            config,
            overview_bound_size,
            None,
            args.launch,
            reduce_annotation=True,
            source_format="hdf5",
            data_path=args.generated_data,
            write_data_path=args.generated_data_output,
        )

if __name__ == "__main__":
    main()

ext: py

reason: import numpy


text:
import scipy.sparse as sp
import numpy as np
from ...survey import BaseTimeRx
from ...utils import Zero

class Pressure(BaseTimeRx):
    """Richards Receiver Object"""

    def __call__(self, U, simulation):
        P = self.getP(simulation.mesh, simulation.time_mesh)
        u = np.concatenate(U)
        return P * u

    def deriv(self, U, simulation, du_dm_v=None, v=None, adjoint=False):
        P = self.getP(simulation.mesh, simulation.time_mesh)
        if not adjoint:
            return P * du_dm_v  # + 0 for dRx_dm contribution
        if v is None:
            raise Exception("v must be provided if computing adjoint")
        return P.T * v, Zero()

class Saturation(BaseTimeRx):
    """Richards Receiver Object"""

    def __call__(self, U, simulation):
        # The water retention curve model should have been updated in the prob
        P = self.getP(simulation.mesh, simulation.time_mesh)
        usat = np.concatenate([simulation.water_retention(ui) for ui in U])
        return P * usat

    def deriv(self, U, simulation, du_dm_v=None, v=None, adjoint=False):
        # The water retention curve model should have been updated in the prob

        P = self.getP(simulation.mesh, simulation.time_mesh)
        dT_du = sp.block_diag([simulation.water_retention.derivU(ui) for ui in U])

        if simulation.water_retention.needs_model:
            dT_dm = sp.vstack([simulation.water_retention.derivM(ui) for ui in U])
        else:
            dT_dm = Zero()

        if v is None and not adjoint:
            # this is called by the fullJ in the problem
            return P * (dT_du * du_dm_v) + P * dT_dm
        if not adjoint:
            return P * (dT_du * du_dm_v) + P * (dT_dm * v)

        # for the adjoint return both parts of the sum separately
        if v is None:
            raise Exception("v must be provided if computing adjoint")
        PTv = P.T * v
        return dT_du.T * PTv, dT_dm.T * PTv

max_stars_repo_name: ElliotCheung/simpeg

max_issues_repo_name: ElliotCheung/simpeg

max_forks_repo_name: ElliotCheung/simpeg

reason: import numpy,import scipy


text:
from Stream import Stream
from Operators import lf, stream_agent
from examples_element_wrapper import print_stream
import numpy as np
import random

x = Stream('x')
z = Stream('z')

print_stream(x)
print_stream(z)

def average(a, state):
        n, cum = state
        b = np.array(a)
        b[0] += cum
        b = np.cumsum(b)
        n_array = np.arange(n+1, n+len(b)+1, 1)
        avg = b/np.rint(n_array)
        state = (n_array[-1], b[-1])
        return (avg, state)
# Initial state is (0, 0.0).
lf(x, z, average, (0, 0.0))

#x.extend(range(10))
x.extend([random.randint(1, 4) for _ in range(10)])

max_stars_repo_name: AnomalyInc/StreamPy

max_issues_repo_name: StreamPy/StreamPy

max_forks_repo_name: StreamPy/StreamPy


text:
"""
Handlers for keys related to set membership: integer, rational, etc.
"""
from sympy.assumptions import Q, ask
from sympy.assumptions.handlers import CommonHandler

class AskIntegerHandler(CommonHandler):
    """
    Handler for key 'integer'
    Test that an expression belongs to the field of integer numbers
    """

    @staticmethod
    def _number(expr, assumptions):
        # helper method
        if expr.as_real_imag()[1] == 0:
            return expr.evalf(1) == expr
        return False

    @staticmethod
    def Add(expr, assumptions):
        """
        Integer + Integer       -> Integer
        Integer + !Integer      -> !Integer
        !Integer + !Integer -> ?
        """
        if expr.is_number:
            return AskIntegerHandler._number(expr, assumptions)
        return test_closed_group(expr, assumptions, 'integer')

    @staticmethod
    def Mul(expr, assumptions):
        """
        Integer*Integer      -> Integer
        Integer*Irrational   -> !Integer
        Odd/Even             -> !Integer
        Integer*Rational     -> ?
        """
        if expr.is_number:
            return AskIntegerHandler._number(expr, assumptions)
        _output = True
        for arg in expr.args:
            if not ask(arg, Q.integer, assumptions):
                if arg.is_Rational:
                    if arg.q == 2:
                        return ask(2*expr, Q.even, assumptions)
                    if ~(arg.q & 1):
                        return None
                elif ask(arg, Q.irrational, assumptions):
                    if _output:
                        _output = False
                    else:
                        return
                else:
                    return
        else:
            return _output

    Pow = Add

    @staticmethod
    def int(expr, assumptions):
        return True

    @staticmethod
    def Integer(expr, assumptions):
        return True

    @staticmethod
    def Rational(expr, assumptions):
        # rationals with denominator one get
        # evaluated to Integers
        return False

    @staticmethod
    def Real(expr, assumptions):
        return int(expr) == expr

    @staticmethod
    def Pi(expr, assumptions):
        return False

    @staticmethod
    def Exp1(expr, assumptions):
        return False

    @staticmethod
    def Infinity(expr, assumptions):
        return False

    @staticmethod
    def NegativeInfinity(expr, assumptions):
        return False

    @staticmethod
    def ImaginaryUnit(expr, assumptions):
        return False

    @staticmethod
    def Abs(expr, assumptions):
        return ask(expr.args[0], Q.integer, assumptions)

class AskRationalHandler(CommonHandler):
    """
    Handler for key 'rational'
    Test that an expression belongs to the field of rational numbers
    """

    @staticmethod
    def Add(expr, assumptions):
        """
        Rational + Rational     -> Rational
        Rational + !Rational    -> !Rational
        !Rational + !Rational   -> ?
        """
        if expr.is_number:
            if expr.as_real_imag()[1]:
                return False
        return test_closed_group(expr, assumptions, 'rational')

    Mul = Add

    @staticmethod
    def Pow(expr, assumptions):
        """
        Rational ** Integer      -> Rational
        Irrational ** Rational   -> Irrational
        Rational ** Irrational   -> ?
        """
        if ask(expr.exp, Q.integer, assumptions):
            return ask(expr.base, Q.rational, assumptions)
        elif ask(expr.exp, Q.rational, assumptions):
            if ask(expr.base, Q.prime, assumptions):
                return False

    @staticmethod
    def Rational(expr, assumptions):
        return True

    @staticmethod
    def Real(expr, assumptions):
        # it's finite-precission
        return True

    @staticmethod
    def ImaginaryUnit(expr, assumptions):
        return False

    @staticmethod
    def Infinity(expr, assumptions):
        return False

    @staticmethod
    def NegativeInfinity(expr, assumptions):
        return False

    @staticmethod
    def Pi(expr, assumptions):
        return False

    @staticmethod
    def Exp1(expr, assumptions):
        return False

class AskIrrationalHandler(CommonHandler):

    @staticmethod
    def Basic(expr, assumptions):
        _real = ask(expr, Q.real, assumptions)
        if _real:
            _rational = ask(expr, Q.rational, assumptions)
            if _rational is None: return None
            return not _rational
        else: return _real

class AskRealHandler(CommonHandler):
    """
    Handler for key 'real'
    Test that an expression belongs to the field of real numbers
    """

    @staticmethod
    def _number(expr, assumptions):
        return not expr.as_real_imag()[1]

    @staticmethod
    def Add(expr, assumptions):
        """
        Real + Real              -> Real
        Real + (Complex & !Real) -> !Real
        """
        if expr.is_number:
            return AskRealHandler._number(expr, assumptions)
        return test_closed_group(expr, assumptions, 'real')

    @staticmethod
    def Mul(expr, assumptions):
        """
        Real*Real               -> Real
        Real*Imaginary          -> !Real
        Imaginary*Imaginary     -> Real
        """
        if expr.is_number:
            return AskRealHandler._number(expr, assumptions)
        result = True
        for arg in expr.args:
            if ask(arg, Q.real, assumptions):
                pass
            elif ask(arg, Q.imaginary, assumptions):
                result = result ^ True
            else:
                break
        else:
            return result

    @staticmethod
    def Pow(expr, assumptions):
        """
        Real**Integer         -> Real
        Positive**Real        -> Real
        Real**(Integer/Even)  -> Real if base is nonnegative
        Real**(Integer/Odd)   -> Real
        """
        if expr.is_number:
            return AskRealHandler._number(expr, assumptions)
        if ask(expr.base, Q.real, assumptions):
            if ask(expr.exp, Q.integer, assumptions):
                return True
            elif expr.exp.is_Rational:
                if (expr.exp.q % 2 == 0):
                    return ask(expr.base, Q.real, assumptions) and \
                       not ask(expr.base, Q.negative, assumptions)
                else: return True
            elif ask(expr.exp, Q.real, assumptions):
                if ask(expr.base, Q.positive, assumptions):
                    return True

    @staticmethod
    def Rational(expr, assumptions):
        return True

    @staticmethod
    def Real(expr, assumptions):
        return True

    @staticmethod
    def Pi(expr, assumptions):
        return True

    @staticmethod
    def Exp1(expr, assumptions):
        return True

    @staticmethod
    def Abs(expr, assumptions):
        return True

    @staticmethod
    def re(expr, assumptions):
        return True

    im = re

    @staticmethod
    def ImaginaryUnit(expr, assumptions):
        return False

    @staticmethod
    def Infinity(expr, assumptions):
        return False

    @staticmethod
    def NegativeInfinity(expr, assumptions):
        return False

    @staticmethod
    def sin(expr, assumptions):
        if ask(expr.args[0], Q.real, assumptions):
            return True

    cos, exp = sin, sin

class AskExtendedRealHandler(AskRealHandler):
    """
    Handler for key 'extended_real'
    Test that an expression belongs to the field of extended real numbers,
    that is real numbers union {Infinity, -Infinity}
    """

    @staticmethod
    def Add(expr, assumptions):
        return test_closed_group(expr, assumptions, 'extended_real')

    Mul, Pow = Add, Add

    @staticmethod
    def Infinity(expr, assumptions):
        return True

    @staticmethod
    def NegativeInfinity(expr, assumptions):
        return True

class AskComplexHandler(CommonHandler):
    """
    Handler for key 'complex'
    Test that an expression belongs to the field of complex numbers
    """

    @staticmethod
    def Add(expr, assumptions):
        return test_closed_group(expr, assumptions, 'complex')

    Mul, Pow = Add, Add

    @staticmethod
    def Number(expr, assumptions):
        return True

    @staticmethod
    def NumberSymbol(expr, assumptions):
        return True

    @staticmethod
    def Abs(expr, assumptions):
        return True

    @staticmethod
    def ImaginaryUnit(expr, assumptions):
        return True

    @staticmethod
    def Infinity(expr, assumptions):
        return False

    @staticmethod
    def NegativeInfinity(expr, assumptions):
        return False

    sin, cos, exp, re, im = [Abs]*5 # they are all complex functions

class AskImaginaryHandler(CommonHandler):
    """
    Handler for key 'imaginary'
    Test that an expression belongs to the field of imaginary numbers,
    that is, numbers in the form x*I, where x is real
    """

    @staticmethod
    def _number(expr, assumptions):
        # helper method
        return not expr.as_real_imag()[0]

    @staticmethod
    def Add(expr, assumptions):
        """
        Imaginary + Imaginary -> Imaginary
        Imaginary + Complex   -> ?
        Imaginary + Real      -> !Imaginary
        """
        if expr.is_number:
            return AskImaginaryHandler._number(expr, assumptions)
        reals = 0
        for arg in expr.args:
            if ask(arg, Q.imaginary, assumptions):
                pass
            elif ask(arg, Q.real, assumptions):
                reals += 1
            else:
                break
        else:
            if reals == 0:
                return True
            if reals == 1 or (len(expr.args) == reals):
                # two reals could sum 0 thus giving an imaginary
                return False

    @staticmethod
    def Mul(expr, assumptions):
        """
        Real*Imaginary      -> Imaginary
        Imaginary*Imaginary -> Real
        """
        if expr.is_number:
            return AskImaginaryHandler._number(expr, assumptions)
        result = False
        reals = 0
        for arg in expr.args:
            if ask(arg, Q.imaginary, assumptions):
                result = result ^ True
            elif not ask(arg, Q.real, assumptions):
                break
        else:
            if reals == len(expr.args):
                return False
            return result

    Pow = Add

    @staticmethod
    def Number(expr, assumptions):
        return not (expr.as_real_imag()[1] == 0)

    NumberSymbol = Number

    @staticmethod
    def ImaginaryUnit(expr, assumptions):
        return True

class AskAlgebraicHandler(CommonHandler):
    """Handler for 'algebraic' key. """

    @staticmethod
    def Add(expr, assumptions):
        return test_closed_group(expr, assumptions, 'algebraic')

    @staticmethod
    def Mul(expr, assumptions):
        return test_closed_group(expr, assumptions, 'algebraic')

    @staticmethod
    def Pow(expr, assumptions):
        return expr.exp.is_Rational and ask(expr.base, 'algebraic', assumptions)

    @staticmethod
    def Number(expr, assumptions):
        return False

    @staticmethod
    def Rational(expr, assumptions):
        return expr.q != 0

    @staticmethod
    def ImaginaryUnit(expr, assumptions):
        return True

    @staticmethod
    def AlgebraicNumber(expr, assumptions):
        return True

#### Helper methods

def test_closed_group(expr, assumptions, key):
    """
    Test for membership in a group with respect
    to the current operation
    """
    result = True
    for arg in expr.args:
        _out = ask(arg, key, assumptions)
        if _out is None: break
        elif _out is False:
            if result: result = False
            else: break
    else:
        return result

max_stars_repo_name: benjaminmcdonald/sympy

max_issues_repo_name: benjaminmcdonald/sympy

max_forks_repo_name: benjaminmcdonald/sympy

reason: from sympy


text:
import cv2 as cv
import numpy as np

#STEP BY STEP:
#
#   1. Creating self.image object:
#               *Open a file
#               *Convert to grayscale
#               *Resize
#   2. Denoise
#   3. Adaptive thresholding
#   4. Eliminate grid lines
#   5. Finding numbers
#   6. Create numpy object

class SudokuExtraction(object):
    def __init__(self, image_path):
        super().__init__()
        self.image_path = image_path
        self.image = cv.imread(image_path)

    def find_sudoku_grid(self):
        gray = cv.cvtColor(self.image,cv.COLOR_BGR2GRAY)
        processed_temp_image = cv.GaussianBlur(gray, (7,7), 0)
        processed_temp_image = cv.adaptiveThreshold(processed_temp_image,
                                                        255,
                                                        cv.ADAPTIVE_THRESH_GAUSSIAN_C,
                                                        cv.THRESH_BINARY_INV,
                                                        11,
                                                        2)
        countours, _ = cv.findContours(processed_temp_image, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
        if not countours:
            return False # no contour found

        sudoku_grid = sorted(countours, key=cv.contourArea, reverse=True)[0]
        epsilon = 0.1*cv.arcLength(sudoku_grid,True)

        self.grid_edges = cv.approxPolyDP(sudoku_grid, epsilon,True)
        self.sudoku_grid = cv.drawContours(self.image, [self.grid_edges], -1, (255,255,255), 3)
        return True

    def rotate_image(self, angle):
        image_center = tuple(np.array(self.image.shape[1::-1]) / 2)
        rot_mat = cv.getRotationMatrix2D(image_center, angle, 1.0)
        return cv.warpAffine(self.image, rot_mat, self.image.shape[1::-1], flags=cv.INTER_LINEAR)

    @property
    def longest_side(self):
        return max(
                    [
                        cv.norm(self.grid_edges[0][0], self.grid_edges[1][0]), # top left vs top right
                        cv.norm(self.grid_edges[1][0], self.grid_edges[2][0]), # top right vs bottom right
                        cv.norm(self.grid_edges[2][0], self.grid_edges[3][0]), # bottom right vs bottom left
                        cv.norm(self.grid_edges[3][0], self.grid_edges[0][0])  # bottom left vs top left
                    ]
                )

    def zoom_sudoku(self):
        src = np.array([self.grid_edges[1][0],
                            self.grid_edges[0][0],
                            self.grid_edges[3][0],
                            self.grid_edges[2][0]], dtype=np.float32)

        dst = np.array([[0, 0],
                            [self.longest_side - 1, 0],
                            [self.longest_side - 1, self.longest_side - 1],
                            [0, self.longest_side - 1]], dtype=np.float32)

        tranformed_perspective = cv.getPerspectiveTransform(src, dst)
        transformed_image = cv.warpPerspective(self.image,
                                                tranformed_perspective,
                                                (int(self.longest_side), int(self.longest_side)))
        self.sudoku_grid = self.rotate_image(90) # TODO: check if it works

    def get_numbers(self):
        processed_image = cv.cvtColor(self.sudoku_grid,cv.COLOR_BGR2GRAY)
        cv.bitwise_not(processed_image, processed_image)
        for i in range(9):
            for j in range(9):
                width = int(processed_image.shape[0]/9)
                height = int(processed_image.shape[1]/9)
                yield processed_image[i*width:i*width + width,j*height:j*height + height]

class NumbersProcessing:
    # Crop the image
    UP = 0.15
    LEFT = 0.2
    # Thresholds
    NUMBER_THREASHOLD = 150
    MIN_NUMBER_SIZE = 75

    def __init__(self, numbers, resolution):
        self.numbers = numbers
        self.resolution = resolution

    def _find_center(self, contours):
        M = cv.moments(sorted(contours, key=cv.contourArea, reverse=True)[0])
        return int(M['m10']/M['m00']), int(M['m01']/M['m00']) # cx, cy

    def _center_number(self, number):
        contours, _ = cv.findContours(number, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
        cx, cy = self._find_center(contours)
        number = cv.drawContours(np.zeros((cy*2, cx*2)), contours, -1, (255,255,255))
        number = cv.fillPoly(number,pts=contours, color=(255,255,255))
        number = cv.resize(number, (self.resolution,self.resolution))
        return cv.erode(number, (3,3))

    def preprocess_numbers(self):
        for num in self.numbers:
            num = num[int(num.shape[0]*self.UP):num.shape[0],int(num.shape[1]*self.LEFT):num.shape[1]]
            num = cv.resize(num, (60,60))
            _,num = cv.threshold(num,self.NUMBER_THREASHOLD,255,cv.THRESH_BINARY)

            if cv.countNonZero(num) > self.MIN_NUMBER_SIZE:
                yield self._center_number(number=num)
            else:
                yield None


text:
import numpy as np

# Define theta generator function for Adam
# [Kingma, D. P., & Ba, J. L. (2015). Adam: A Method for Stochastic Optimization. International Conference on Learning Representations.](https://arxiv.org/pdf/1412.6980v8.pdf)
def adapt(params, grad_fun):

    # Define theta generator
    def theta_gen_adam(theta):
        # Initialize hyperparameters, moment and iteration variables
        alpha = params['alpha']
        beta1 = params['beta1']
        beta2 = params['beta2']
        epsilon = params['epsilon']
        moment1 = np.zeros(theta.shape[0])
        moment2 = np.zeros(theta.shape[0])
        t = 1

        # Generate new theta
        while True:
            # Get gradient
            gradient = grad_fun(theta)

            # Update moment estimates
            moment1 = beta1 * moment1 + (1 - beta1) * gradient
            moment2 = beta2 * moment2 + (1 - beta2) * np.square(gradient)

            # Yield adapted gradient
            theta = ( theta - alpha * (1 - beta2**t)**0.5 / (1 - beta1**t) *
                        moment1 / (epsilon + np.sqrt(moment2)) )
            yield theta
            t += 1

    return theta_gen_adam

    # Less efficient calculation of theta that is easier to follow
    # moment1_hat = moment1 / (1 - beta1**t)
    # moment2_hat = moment2 / (1 - beta2**t)
    # theta = theta - alpha * moment1_hat / (epsilon + np.sqrt(moment2_hat))

max_stars_repo_name: octaviansoldea/pysgd

max_issues_repo_name: octaviansoldea/pysgd

max_forks_repo_name: octaviansoldea/pysgd


text:
"""
Tests for snn.Weight.
"""
import unittest
from unit_tests import ModuleTest
import numpy as np
from spikey.snn import weight

def weight_generator(shape):
    matrix = np.random.uniform(0, 9999, size=shape)
    return np.ma.array(matrix, mask=(matrix == 0), fill_value=0)

class TestWeight(unittest.TestCase, ModuleTest):
    """
    Tests for snn.Weight.
    """

    TYPES = [weight.Manual, weight.Random]
    BASE_CONFIG = {
        "n_inputs": 10,
        "n_neurons": 10,
        "max_weight": 1,
        "force_unidirectional": False,
        "matrix": weight_generator((10 + 10, 10)),
        "weight_generator": weight_generator,
        "matrix_mask": None,
    }

    def _check_matrix_types(self, weights):
        self.assertIsInstance(weights._matrix, np.ndarray)
        self.assertIsInstance(weights.matrix, np.ndarray)
        self.assertTrue(hasattr(weights._matrix, "mask"))
        self.assertTrue(not hasattr(weights.matrix, "mask"))
        self.assertEqual(weights._matrix.fill_value, 0)

    def _assert_clipped(self, weights):
        self.assertTrue(
            np.all(weights.matrix >= 0)
            & np.all(weights.matrix <= self.BASE_CONFIG["max_weight"])
        )

    @ModuleTest.run_all_types
    def test_arith(self):
        w_shape = (
            self.BASE_CONFIG["n_inputs"] + self.BASE_CONFIG["n_neurons"],
            self.BASE_CONFIG["n_neurons"],
        )

        weights = self.get_obj(
            max_weight=2, matrix=np.ma.array(np.ones(w_shape)), weight_generator=np.ones
        )
        self.assertTrue(np.mean(weights._matrix.mask) <= 0.05)
        weights += np.ones((w_shape[0], 1))
        self.assertTrue(np.mean(weights._matrix == 2) >= 0.95)
        self.assertTrue(np.mean(weights.matrix == 2) >= 0.95)

        weights = self.get_obj(
            max_weight=2, matrix=np.zeros(w_shape), weight_generator=np.zeros
        )
        self.assertTrue(np.all(weights._matrix.mask))
        weights += np.ones((w_shape[0], 1))
        self.assertTrue(np.all(weights.matrix == 0))

    @ModuleTest.run_all_types
    def test_usage(self):
        weights = self.get_obj()
        self._check_matrix_types(weights)
        self._assert_clipped(weights)

        spike_shape = (self.BASE_CONFIG["n_inputs"] + self.BASE_CONFIG["n_neurons"], 1)
        weights * np.random.uniform(0, 1, spike_shape)
        weights / np.random.uniform(0, 1, spike_shape)
        weights + np.random.uniform(0, 1, spike_shape)
        weights - np.random.uniform(0, 1, spike_shape)

        weights *= np.random.uniform(0, 1, spike_shape)
        weights /= np.random.uniform(0, 1, spike_shape)
        weights += np.random.uniform(0, 1, spike_shape)
        weights -= np.random.uniform(0, 1, spike_shape)

        weights.matrix
        self._check_matrix_types(weights)
        self._check_matrix_types(weights.copy())

        layers = [np.ones((10, 5)), np.ones((4, 4)), np.ones((1, 1))]
        weights = self.get_obj(**{"matrix": layers, "matrix_mask": layers})
        self._check_matrix_types(weights)
        self._assert_clipped(weights)

if __name__ == "__main__":
    unittest.main()

max_stars_repo_name: SpikeyCNS/spikey

max_issues_repo_name: SpikeyCNS/spikey

max_forks_repo_name: SpikeyCNS/spikey


text:
from OpenGL.GL import *
from PIL import Image
import numpy

def load_texture(path):
    texture = glGenTextures(1)
    glBindTexture(GL_TEXTURE_2D, texture)
    # Set the texture wrapping parameters
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
    # Set texture filtering parameters
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    # load image
    image = Image.open(path)
    img_data = numpy.array(list(image.getdata()), numpy.uint8)
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, image.width, image.height, 0, GL_RGB, GL_UNSIGNED_BYTE, img_data)
    return texture

max_stars_repo_name: geraldzakwan/ITBMaps

max_stars_repo_licenses: MIT WTFPL Unlicense

max_issues_repo_name: geraldzakwan/ITBMaps

max_issues_repo_licenses: MIT WTFPL Unlicense

max_forks_repo_name: geraldzakwan/ITBMaps

max_forks_repo_licenses: MIT WTFPL Unlicense


text:
#!/usr/bin/env python
"""
This file is part of the WorkMATe source code, provided as supplemental material
for the article:
    "Flexible Working Memory through Selective Gating and Attentional Tagging"
Wouter Kruijne, Sander M Bohte, Pieter R Roelfsema, Christian N L Olivers

Please see the README file for additional details.
Questions or comments regarding this file can be addressed to w.kruijne@vu.nl
-------------------
This file implements four functions for individual runs for the four tasks
tasks. The random seeds that each run get here yield 'illustrative'
convergence times.
"""
#%% Import modules
import numpy as np
import matplotlib.pyplot as plt

import world
import workmate_match
from m_runner import matchnet_regression, mean_std

#%% Time window constants

trial_print = 150 #printing of trials
max_time = 50000  #total number of training trials

#%% DMS task constants

n_switches = 5 # i.e. 6 stimulus sets, 3 stimuli per set
n_stim = 3
conv_crit = 0.85

#%% Run a trial
#np.random.seed(20) # a randnum [0- 1e10] > Use the np.random.seed NOT random.seed
#These are two different generators!
ncore = 4

def _run_trial(agent, env,log_q=False, **kwargs):
    """
    Generic function to run a trial:
    agent   : the agent being run
    env     : the environment implementing the task
    env     : if log_q=True, function returns both the reward and the output
                values in the q-layer

    Additional kwargs can be passed to be set in the environment (not used here)
    """
    if not kwargs is None:
        env.__dict__.update(kwargs)
    # init trial results:
    totrew = 0.0
    qlog = []
    match_per_trial = []
    tt_per_trial = []

    # step until obs == 'RESET', collect reward and q-values
    while True:
        #r, matches, tt_match = agent.step()
        r, matches, tt = agent.step()
        totrew += r
        if matches != None and tt != None:
            match_per_trial.append(matches)
            tt_per_trial.append(tt)
        if log_q:
            print((agent.x))
            qlog.append(agent.q)
        if agent.obs == 'RESET':
            break

    # format result:
    #print(match_per_trial)
    #print(tt_per_trial)
    retval = (totrew, np.array(qlog)) if log_q else totrew
    match_info = [match_per_trial, tt_per_trial]
    return retval, match_info

#%% Print information from trial to analyse the learning

def watch_trial(agent, env):
    """
    Function to watch a trial giving a specific agent and environment
    Returns the inputs, q-values, actions and rewards for each time step
    Returns the total reward received

    """
    #This is a function that prints the agents strategy during a trial
    #This will give insight into its functioning

    #For each time step, I want to track
    xlog = []       #inputs
    qlog = []       #q-values
    zlog = []       #actions
    rlog = []       #rewards
    tlog = []
    envtlog = []
    obslog = []
    #trial_typelog = []
    totrew = 0      #total reward

    #step until obs == 'RESET', collect reward and q-values
    while True:

        totrew += agent.step()
        #print(agent.x)
        obslog.append(agent.obs)
        xlog.append(agent.x)
        qlog.append(agent.q)
        zlog.append(agent.z)
        rlog.append(agent.r)
        tlog.append(agent.t)
        envtlog.append(env.t)
        if agent.obs == 'RESET':
            trial_type = env.trial_type
            break

    retval = (totrew, obslog, trial_type, np.array(xlog), np.array(qlog), np.array(zlog), np.array(rlog), np.array(tlog), np.array(envtlog))
    return retval

#def run_dms(W_hx_pre, W_qh_pre, pre_stim, viz_converged=True, agent_type = 'W', print_all = 'off'):
def run_dms_match(matchnet, viz_converged=True, print_all = 'off', stim_type = 'alpha'):

    #stim_type can be
    ###'alpha' > original small binary vectors
    ###'cifar' > learned latent vector representation CIFAR-10

    #np.random.seed(20)
    #n_switches = 5 # i.e. 6 stimulus sets, 3 stimuli per set

    #Initialise agent and task
    dms = world.DMS(n_stim=n_stim, n_switches=n_switches, stim_type = stim_type)
    agent = workmate_match.WorkMATe(dms, nblocks=2)
    agent.matchnet = matchnet
    print(agent)

    ### Initialize training:
    # buffer to store 'the moment of switching' (=convergence)
    iswi = np.zeros((n_switches+1))
    # buffer for last 500 trials:
    saved_p = np.zeros(1)
    total_buff = np.zeros(500)
    # buffer for performance immediately after a switch
    swiperf = np.nan * np.ones((n_switches+1, total_buff.size))

    #Buffers for matchnet:
    x2_buff = np.zeros((agent.nblocks,1))
    tt_buff = np.zeros((agent.nblocks,1))
    performance_all_blocks = np.zeros(agent.nblocks)
    perf_match = [] #performance match_net

    # counters
    last=False
    i = 0
    i_ = 0
    i_l = 0

    aa = True
    while aa:
        # run trial, get performance
        r, match_info = _run_trial(agent, dms)
        # increase i
        i += 1
        # was the trial correct?
        corr = ( r >= dms.bigreward )
        total_buff[0] = corr
        total_buff = np.roll(total_buff, 1)

        #buffers for match info
        x2_buff = np.append(x2_buff, np.array(match_info[0]).transpose(), axis=1)
        tt_buff = np.append(tt_buff, np.array(match_info[1]).transpose(), axis=1)

        #if the past 100 trials were 85% correct, set is 'learned'
        if np.mean(total_buff[:100]) >= conv_crit and last == False:
            print('Convergence at {}\tSwitch to set {}'.format(i,dms.setnr+1))
            iswi[dms.setnr] = i
            # if criterion reached in less than 500 trials,
            # 'performanc post-switch' hasn't been logged yet -- do that now,
            # using only the trials with this set:
            if i < i_ + 500:
                swiperf[dms.setnr, :(i- i_)] = total_buff[:(i - i_)] # leaves nans for the rest of performance

            if dms.setnr == n_switches:
                dms.stimuli = dms.STIM
                last = True

            dms.switch_set()
            total_buff *= 0 # reset performance buffer
            i_ = i

        if last == True:
            i_l = i_l + 1

        if i_l >= 2000:
            break

        # @ iswi + 500: store post-switch performance:
        if i == i_ + 500:
            swiperf[dms.setnr, :] = total_buff

        # print progress:
        if i % trial_print == 0:
            #get matchnet performance
            for ff in range(agent.nblocks):
                performance = matchnet_regression(x2_buff[ff,:], tt_buff[ff,:])
                performance_all_blocks[ff] = performance.copy()
            perf_match.append(performance_all_blocks.copy())
            #reset buffers to zero
            x2_buff = np.zeros((agent.nblocks,1))
            tt_buff = np.zeros((agent.nblocks,1))

            #print performance in a list
            if print_all == 'on':
                print(i, '\t', np.mean(total_buff))
            saved_p = np.append(saved_p,np.mean(total_buff))

        if i >= max_time:
            aa = False
    print('Loop ended')

    return (dms, saved_p, iswi, agent, perf_match)

#%% Visualisation phase 2: Evaluation example

def plot_dms(saved_p,iswi):
    plt.rc('axes', labelsize=28)
    plt.rc('xtick', labelsize=24)    # fontsize of the tick labels
    plt.rc('ytick', labelsize=24)

    #print DMS performance in a graph
    arr = [i*trial_print for i in range(0,len(saved_p))]
    #get the markers on the line
    a = (np.floor(iswi/trial_print)).astype(int)    #find point closest below
    b = (np.ceil(iswi/trial_print)).astype(int)     #find point closest above

    iswi_p = np.zeros(len(iswi))
    for x in range(len(iswi)):                      #for all but the last one..
        if b[x] == len(saved_p):
            b[x] = b[x]                          #if its the last trial, set b[x] back to floor
        diff = saved_p[a[x] - 1] - saved_p[b[x] - 1]        #find difference in performance
        if diff < 0:                                #is b is bigger than a
           iswi_p[x] = b[x] - 1                          #choose b
        else:
           iswi_p[x] = a[x] - 1                         #if a is bigger choose a

    iswi_p = iswi_p.astype(int)

    fig, ax = plt.subplots(figsize=(10,8))
    #add switch positions to graph
    #iswi_index_insert = list(np.ceil(np.array(iswi)/trial_print) + range(n_switches+1))
    saved_p = list(saved_p)
    for i, value in enumerate(iswi):
        iswi_index_insert = np.ceil(value/trial_print) + i
        saved_p.insert(int(iswi_index_insert), conv_crit)
        arr.insert(int(iswi_index_insert), value)

    xposition = iswi
    for xc in xposition:
        plt.axvline(x=xc, color='k', linestyle='--', linewidth=4)

    #plt.plot(iswi-trial_print,saved_p[iswi_p],'b*', markersize = 10)
    plt.plot(arr,saved_p,'r', linewidth=4)
    plt.xlabel('Trials')
    plt.ylabel('Ratio full task reward')
    plt.title('DMS task performance', fontsize=34)
    plt.grid()
    plt.show()
    return arr, saved_p

def plot_perf_DMS(iswi_buff):
    #Exclude non-converged agents
    iswi_conv = np.array([x for x in iswi_buff if np.all(x) == True])
    #How many excluded?
    nr_conv = len(iswi_conv)
    #Not cumulative
    iswi_diff = np.c_[iswi_conv[:,0], np.diff(iswi_conv)]
    #Average + std for each level
    iswi_avg, iswi_std = mean_std(iswi_diff)
    #Plot
    fig, ax = plt.subplots(figsize=(10,7))
    #ax.plot(np.arange(6)+1, iswi_avg, 'o')
    ax.set_axisbelow(True)
    plt.grid()
    plt.bar(np.arange(6)+1, iswi_avg/1000, color='navy', alpha=0.8)
    caplines = plt.errorbar(np.arange(6)+1, iswi_avg/1000, yerr=iswi_std/1000, ls='None', lolims=True, capsize=10, color='k')
    caplines[0].set_marker('_')

    plt.ylim([0,(iswi_avg[0]+iswi_std[0]+1000)/1000])
    plt.title('DMS trials until set convergence', fontsize=34)
    plt.ylabel('Trials (x1000)')
    plt.xlabel('DMS stimulus sets')
    return nr_conv

#%%
if __name__ == '__main__':
    run_dms_match(seed=1)


text:
import os
from datetime import datetime
import numpy as np
import openpyxl as xl

from profile_ import Profile
from variable_mgr import Date

def ExportWriter(writer):
    try:
        writer.save()
    except PermissionError:
        print('Unable to export, the file is in use')

def WriteProfile(ws, profile, variables, entity_names, phaser=False):
    n = profile.time().size

    if phaser:
        headers = ['Date', 'Oil rate (stb/day)', 'Total Gas Rate (kscf/day)', 'Water Rate (stb/day)',
                   'Lift Gas Rate (kscf/day)', 'Gas Injection Rate (kscf/day)', 'Water Injection Rate(stb/day)']

        unit_conversions = np.full(len(variables), 1.e3)

    else:
        headers = ['{} [{}]'.format(v.GetMenuLabel(), v.GetUnit()) for v in variables]
        unit_conversions = np.full(len(variables), 1.)

    # write headers
    for j, _ in enumerate(variables):
        ws.cell(row=1, column=j+2).value = headers[j]

    # write name
    ws.cell(row=1, column=1).value = 'EntityName'
    for i in range(n):
        ws.cell(row=i+2, column=1).value = entity_names[i]

    # write variables
    for j, variable in enumerate(variables):
        array = profile.Get(variable.GetId())

        if isinstance(variable, Date):
            array = array.astype(datetime)
        else:
            array *= unit_conversions[j]

        for i in range(n):
            ws.cell(row=i+2, column=j + 2).value = array[i]

def GetPath(directory, name):
    # pre-append the date
    file_name = datetime.today().strftime('%Y-%m-%d') + '_' + name

    path = os.path.join(directory, file_name)
    if path[-5:] != '.xlsx':
        path += '.xlsx'

    return path

def ToExcel(cases, items, simulations, variables, directory, dateline=None, phaser=False):
    # items is a {{[]}} list with items {file_name: {sheet_name: [entities]}}

    # write to excel
    for case in cases:

        for simulation in simulations:

            for file_name, sheets in items.items():

                wb = xl.Workbook()

                for sheet_name, entities in sheets.items():

                    names = []
                    profiles = Profile()

                    for entity in entities:

                        profile = entity.GetSimulationResult(simulation).GetProfile(idx=case)

                        if dateline is not None:
                            profile = profile.resample(dateline)

                        profiles.stack(profile)
                        names += [entity.GetName()] * profile.time().size

                    # write to sheet
                    ws = wb.create_sheet(title=sheet_name)
                    WriteProfile(ws, profiles, variables, names, phaser=phaser)

                # delete default sheet
                wb.remove_sheet(wb.get_sheet_by_name('Sheet'))

                case_name = 'LOW' if case == 0 else 'MID' if case == 1 else 'HIGH'
                path = GetPath(directory, '{}_{}_{}'.format(simulation.GetName(), file_name, case_name))
                wb.save(path)

max_stars_repo_name: FrederikLehn/alveus

max_issues_repo_name: FrederikLehn/alveus

max_forks_repo_name: FrederikLehn/alveus


text:
from abc import abstractmethod

import ipywidgets as widgets
import matplotlib.pyplot as plt
import numpy as np
from nwbwidgets import base
from plotly import graph_objects as go
from pynwb.behavior import SpatialSeries, BehavioralEvents
from pynwb import TimeSeries

from .utils.timeseries import (
    get_timeseries_tt,
    get_timeseries_in_units,
    timeseries_time_to_ind,
)
from .timeseries import (
    AlignMultiTraceTimeSeriesByTrialsConstant,
    AlignMultiTraceTimeSeriesByTrialsVariable,
    AbstractTraceWidget,
    SingleTracePlotlyWidget,
    SeparateTracesPlotlyWidget,
)

from .base import lazy_tabs

from .controllers import StartAndDurationController

def show_behavioral_events(beh_events: BehavioralEvents, neurodata_vis_spec: dict):
    return base.dict2accordion(
        beh_events.time_series, neurodata_vis_spec, ls="", marker="|"
    )

def show_spatial_series(node: SpatialSeries, **kwargs):
    data, unit = get_timeseries_in_units(node)
    tt = get_timeseries_tt(node)

    if len(data.shape) == 1:
        fig, ax = plt.subplots()
        ax.plot(tt, data, **kwargs)
        ax.set_xlabel("t (sec)")
        if unit:
            ax.set_xlabel("x ({})".format(unit))
        else:
            ax.set_xlabel("x")
        ax.set_ylabel("x")

    elif data.shape[1] == 2:
        fig, ax = plt.subplots()
        ax.plot(data[:, 0], data[:, 1], **kwargs)
        if unit:
            ax.set_xlabel("x ({})".format(unit))
            ax.set_ylabel("y ({})".format(unit))
        else:
            ax.set_xlabel("x")
            ax.set_ylabel("y")
        ax.axis("equal")

    elif data.shape[1] == 3:
        import ipyvolume.pylab as p3

        fig = p3.figure()
        p3.scatter(data[:, 0], data[:, 1], data[:, 2], **kwargs)
        p3.xlim(np.min(data[:, 0]), np.max(data[:, 0]))
        p3.ylim(np.min(data[:, 1]), np.max(data[:, 1]))
        p3.zlim(np.min(data[:, 2]), np.max(data[:, 2]))

    else:
        raise NotImplementedError

    return fig

def route_spatial_series(spatial_series, **kwargs):
    if len(spatial_series.data.shape) == 1:
        dict_ = {
            "over time": SingleTracePlotlyWidget,
            "trial aligned": trial_align_spatial_series,
        }
    elif spatial_series.data.shape[1] == 2:
        dict_ = {
            "over time": SeparateTracesPlotlyWidget,
            "trace": SpatialSeriesTraceWidget2D,
            "trial aligned": trial_align_spatial_series,
        }
    elif spatial_series.data.shape[1] == 3:
        dict_ = {
            "over time": SeparateTracesPlotlyWidget,
            "trace": SpatialSeriesTraceWidget3D,
            "trial aligned": trial_align_spatial_series,
        }
    else:
        return widgets.HTML("Unsupported number of dimensions.")
    return lazy_tabs(dict_, spatial_series)

class SpatialSeriesTraceWidget(AbstractTraceWidget):
    @abstractmethod
    def plot_data(self, data, units, tt):
        return

    @abstractmethod
    def update_plot(self, data, tt):
        return

    def __init__(
        self,
        timeseries: TimeSeries,
        foreign_time_window_controller: StartAndDurationController = None,
        **kwargs,
    ):
        super().__init__(
            timeseries=timeseries,
            foreign_time_window_controller=foreign_time_window_controller,
            **kwargs,
        )

    def set_out_fig(self):
        timeseries = self.controls["timeseries"].value
        time_window = self.controls["time_window"].value

        istart = timeseries_time_to_ind(timeseries, time_window[0])
        istop = timeseries_time_to_ind(timeseries, time_window[1])
        data, units = get_timeseries_in_units(timeseries, istart, istop)

        tt = get_timeseries_tt(timeseries, istart, istop)

        self.plot_data(data, units, tt)

        def on_change(change):
            time_window = self.controls["time_window"].value
            istart = timeseries_time_to_ind(timeseries, time_window[0])
            istop = timeseries_time_to_ind(timeseries, time_window[1])
            data, units = get_timeseries_in_units(timeseries, istart, istop)
            tt = get_timeseries_tt(timeseries, istart, istop)
            self.update_plot(data, tt)

        self.controls["time_window"].observe(on_change)

class SpatialSeriesTraceWidget2D(SpatialSeriesTraceWidget):
    def plot_data(self, data, units, tt):
        if units is None:
            units = "no units"
        self.out_fig = go.FigureWidget(
            data=go.Scattergl(
                x=list(data[:, 0]),
                y=list(data[:, 1]),
                mode="markers+lines",
                marker_color=tt,
                marker_colorscale="Viridis",
                marker_colorbar=dict(thickness=20, title="time (s)"),
                marker_size=5,
            )
        )

        self.out_fig.update_layout(
            title=self.timeseries.name,
            xaxis_title=f"x ({units})",
            yaxis_title=f"y ({units})",
        )

    def update_plot(self, data, tt):
        self.out_fig.data[0].x = list(data[:, 0])
        self.out_fig.data[0].y = list(data[:, 1])
        self.out_fig.update_traces(marker_color=list(tt))

class SpatialSeriesTraceWidget3D(SpatialSeriesTraceWidget):
    def plot_data(self, data, units, tt):
        if units is None:
            units = "no units"

        self.out_fig = go.FigureWidget(
            data=go.Scatter3d(
                x=data[:, 0],
                y=data[:, 1],
                z=data[:, 2],
                mode="markers+lines",
                marker_color=tt,
                marker_colorscale="Viridis",
                marker_colorbar=dict(thickness=20, title="time (s)"),
                marker_size=5,
            )
        )

        self.out_fig.update_layout(
            scene=dict(
                xaxis_title=f"x ({units})",
                yaxis_title=f"y ({units})",
                zaxis_title=f"x ({units})",
            )
        )

    def update_plot(self, data, tt):
        self.out_fig.data[0].x = list(data[:, 0])
        self.out_fig.data[0].y = list(data[:, 1])
        self.out_fig.data[0].z = list(data[:, 2])
        self.out_fig.update_traces(marker_color=list(tt))

def trial_align_spatial_series(spatial_series, trials=None):
    options = [("x", 0), ("y", 1), ("z", 2)][: spatial_series.data.shape[1]]
    if spatial_series.rate is None:
        return AlignMultiTraceTimeSeriesByTrialsVariable(
            time_series=spatial_series,
            trials=trials,
            trace_controller_kwargs=dict(options=options),
        )
    else:
        return AlignMultiTraceTimeSeriesByTrialsConstant(
            time_series=spatial_series,
            trials=trials,
            trace_controller_kwargs=dict(options=options),
        )

max_issues_repo_name: catalystneuro/nwb-jupyter-widgets

max_forks_repo_name: catalystneuro/nwb-jupyter-widgets


text:
import numpy as np
import matplotlib.pyplot as plt

# factor for 90% coverage with 90% confidence using Normal distribution
# with 10 samples from table XII in [1]
# [1] Montgomery, D. C., & Runger, G. C. (2014). Applied statistics and
# probability for engineers. Sixth edition. John Wiley & Sons.
k = 2.535

run_times = np.load('6_break_times.npy')
n = np.load('n.npy')
run_times1 = np.load('20_break_times.npy')

run_times_means = run_times.mean(axis=2)
run_times_stds = run_times.std(axis=2, ddof=1)
run_times_means1 = run_times1.mean(axis=2)
run_times_stds1 = run_times1.std(axis=2, ddof=1)

plt.figure()
plt.title('AMD FX-8350 + Nvidia TitanXp: 6 line segments')
plt.grid()
plt.errorbar(n, run_times_means[0], yerr=k*run_times_stds[0], capsize=2.0,
             label='npumpy.linalg.lstsq')
plt.errorbar(n, run_times_means[1], yerr=k*run_times_stds[1], capsize=2.0,
             label='cupy.linalg.lstsq')
plt.errorbar(n, run_times_means[2], yerr=k*run_times_stds[2], capsize=2.0,
             label='tf.linalg.lstsq (CPU)')
# plt.errorbar(n, run_times_means[3], yerr=k*run_times_stds[3], capsize=2.0,
#              label='tf.linalg.lstsq fast=False (CPU)')

plt.xlabel('Number of data points')
plt.ylabel('Run time (seconds, Lower is better)')
plt.semilogx()
plt.semilogy()
plt.legend()
plt.savefig('fx_six_breaks.png', bbox_inches='tight')

print('cupy 1e7 time faster', run_times_means[0][-1]/run_times_means[1][-1])

plt.figure()
plt.title('AMD FX-8350 + Nvidia TitanXp: 20 line segments')
plt.grid()
plt.errorbar(n, run_times_means1[0], yerr=k*run_times_stds1[0], capsize=2.0,\
             label='npumpy.linalg.lstsq')
plt.errorbar(n, run_times_means1[1], yerr=k*run_times_stds1[1], capsize=2.0,
             label='cupy.linalg.lstsq')
plt.errorbar(n, run_times_means1[2], yerr=k*run_times_stds1[2], capsize=2.0,
             label='tf.linalg.lstsq (CPU)')
# plt.errorbar(n, run_times_means1[3], yerr=k*run_times_stds1[3], capsize=2.0,
#              label='tf.linalg.lstsq fast=False (CPU)')

plt.xlabel('Number of data points')
plt.ylabel('Run time (seconds, Lower is better)')
plt.semilogx()
plt.semilogy()
plt.legend()
plt.savefig('fx_twenty_breaks.png', bbox_inches='tight')

print('cupy 1e7 time faster', run_times_means1[0][-1]/run_times_means1[1][-1])

plt.show()