text:
# Based on code of Leland McInnes
#
# License: BSD 3 clause

from numba import jit, prange
import numpy as np
from sklearn.base import BaseEstimator
from sklearn.utils import check_array

from sklearn.cluster import KMeans
from sklearn.neighbors import NearestNeighbors
from scipy.spatial import Delaunay

from scipy.sparse import csr_matrix, eye
from scipy.sparse.linalg import spsolve as lin_solve

SMOOTH_K_TOLERANCE = 1e-5
MIN_K_DIST_SCALE = 1e-3
NPY_INFINITY = np.inf

def euclid(x_,r_,s_,sigma_,nn_,a=1,b=1, y_=None, H_=None):
    n_execs, n_dim = nn_.shape[0], s_.shape[1]

    if y_ is None:
        y_ = np.zeros( ( n_execs,n_dim ) , dtype=np.float32 )
    else:
        for k in range(n_execs):
            for j in range(n_dim):
                y_[k,j] = 0
    if H_ is None:
        H_ = np.zeros( (n_execs) , dtype=np.float32 )
    else:
        for k in range(n_execs):
            H_[k] = 0

    return euclid__(x_,r_,s_,sigma_,nn_,a,b, y_, H_)

@jit(nopython=True, fastmath=True, parallel=True)
def euclid__(x_,r_,s_,sigma_,nn_,a,b, y_, H_):
    n_execs, n_samps, m_dim, n_dim = nn_.shape[0], nn_.shape[1], r_.shape[1], s_.shape[1]

    for k in range(n_execs):
        w_ = np.zeros( (n_samps) , dtype=np.float32 )
        w_i_div_sigma_i_sum = 0
        for samp in range(n_samps):
            i = nn_[k,samp]
            dist = 0
            for j in range(m_dim):
                dist += (x_[k,j]-r_[i,j])**2
            w_[samp] = 1/(1+a*dist**b)
            w_i_div_sigma_i = w_[samp] / sigma_[i]
            w_i_div_sigma_i_sum += w_i_div_sigma_i
            for j in range(n_dim):
                y_[k,j] += w_i_div_sigma_i*s_[i,j]
        for j in range(n_dim):
            y_[k,j] /= w_i_div_sigma_i_sum
        for samp in range(n_samps):
            dist = 0
            for j in range(n_dim):
                dist += (y_[k,j]-s_[i,j])**2
            logv_i = -dist/sigma_[samp]
            w_i = w_[samp]
            H_[k] += w_i*logv_i + (1-w_i)*(1-np.exp(logv_i))
    return y_, H_

@jit(nopython=True, fastmath=True, parallel=True)
def weight__(x_,r_,sigma_,nn_,a,b):
    n_execs, n_samps, m_dim = nn_.shape[0], nn_.shape[1], r_.shape[1]

    w_ = np.zeros( (n_execs,n_samps) , dtype=np.float32 )
    for k in range(n_execs):
        w_i_div_sigma_i_sum = 0
        for samp in range(n_samps):
            i = nn_[k,samp]
            dist = 0
            for j in range(m_dim):
                dist += (x_[k,j]-r_[i,j])**2
            w_ki = 1/(1+a*dist**b) / sigma_[i]
            w_i_div_sigma_i_sum += w_ki
            w_[k,samp] = w_ki
        for samp in range(n_samps):
            w_[k,samp] /= w_i_div_sigma_i_sum
    return w_

def compute_H(x_,y_,r_,s_,sigma_,nn_=None,a=1,b=1, H_=None,att_=None,rep_=None, compute_H_only=True):
    n_execs = x_.shape[0]
    if nn_ is None:
        n_samps = r_.shape[0]
        use_all_data = True
        nn_ = np.empty( (0,0), dtype=np.int )
    else:
        n_samps = nn_.shape[1]
        use_all_data = False
    if compute_H_only:
        if H_ is None:
            H_ = np.empty( (n_execs) , dtype=np.float32 )
        if att_ is None:
            att_ = np.empty( (0) , dtype=np.float32 )
        if rep_ is None:
            rep_ = np.empty( (0) , dtype=np.float32 )
    else:
        if H_ is None:
            H_ = np.empty( (0) , dtype=np.float32 )
        if att_ is None:
            att_ = np.empty( (n_execs) , dtype=np.float32 )
        if rep_ is None:
            rep_ = np.empty( (n_execs) , dtype=np.float32 )

    res = compute_H__(x_,y_,r_,s_,sigma_,n_samps,use_all_data,nn_,a,b, H_,att_,rep_, compute_H_only)

    if compute_H_only:
        return res[0]
    else:
        return res[1],res[2]

@jit(nopython=True, fastmath=True, parallel=True)
def compute_H__(x_,y_,r_,s_,sigma_,n_samps,use_all_data,nn_,a,b, H_,att_,rep_, compute_H_only):
    n_execs, m_dim, n_dim = x_.shape[0], r_.shape[1], s_.shape[1]

    for k in range(n_execs):
        att_k, rep_k = 0,0
        for samp in range(n_samps):
            if use_all_data:
                i = samp
            else:
                i = nn_[k,samp]

            dist = 0
            for j in range(m_dim):
                dist += (x_[k,j]-r_[i,j])**2
            w_i = 1/(1+a*dist**b)

            dist = 0
            for j in range(n_dim):
                dist += (y_[k,j]-s_[i,j])**2
            logv_i = -dist/sigma_[samp]

            att_k, rep_k = att_k+w_i*logv_i , rep_k+(1-w_i)*(1-np.exp(logv_i))
        if compute_H_only:
            H_[k] = att_k+rep_k
        else:
            att_[k], rep_[k] = att_k, rep_k

    return  H_, att_, rep_

@jit(nopython=True, fastmath=True, parallel=True)
def smooth(y_,H_,nn_,itrs,in_place=True):
    n_execs, n_dim, n_samps = y_.shape[0], y_.shape[1], nn_.shape[1]

    H_min, H_max = H_.min(), H_.max()
    H_norm = (H_.copy()-H_min)/(H_max-H_min)

    if in_place:
        for itr in range(itrs):
            perm = np.random.permutation(n_execs)
            for kp in range(n_execs):
                k = perm[kp]
                H_norm_k = H_norm[k]
                one_H_norm_k = 1-H_norm_k
                for j in range(n_dim):
                    ava_ji = 0
                    for i in range(n_samps):
                        ava_ji += y_[nn_[k,i],j]
                    y_[k,j] = one_H_norm_k*y_[k,j] + H_norm_k*ava_ji/n_samps
        return y_
    else:
        y_a,y_b = y_, np.zeros( y_.shape , dtype=np.float32 )
        for itr in range(itrs):
            for k in prange(n_execs):
                H_norm_k = H_norm[k]
                one_H_norm_k = 1-H_norm_k
                for j in range(n_dim):
                    ava_ji = 0
                    for i in range(n_samps):
                        ava_ji += y_a[nn_[k,i],j]
                    y_b[k,j] = one_H_norm_k*y_a[k,j] + H_norm_k*ava_ji/n_samps
            y_c = y_b
            y_b = y_a
            y_a = y_c
        return y_a

#@jit(parallel=True)
def select_neighbors(candidates):
        n_sets, n_select, n_neighbors = candidates.shape[0], candidates.shape[1], candidates.shape[2]

        result = np.zeros( (n_sets, n_neighbors) , dtype=np.int32 )
        k = np.zeros( n_select , dtype=np.int32 )
        for i in range(n_sets):
            cand = set()
            j, l = 0, 0
            for j_ in range(n_select):
                k[j_] = 0
            while l < n_neighbors:
                nn_id = candidates[i,j,k[j]]
                k[j] += 1
                if not(nn_id in cand):
                    cand.add(nn_id)
                    result[i,l] = nn_id
                    l = l+1
                    j = (j+1)%n_select
        return result

@jit(nopython=True, fastmath=True, parallel=True)
def dist_each(X,Ys,nn):
    n_execs, n_samps, n_dim = nn.shape[0], nn.shape[1], Ys.shape[1]
    dists = np.zeros( (n_execs,n_samps) , dtype=np.float32 )
    for i in range(n_execs):
        x_i = X[i]
        for samp in range(n_samps):
            #k = nn[i,samp]
            y_k = Ys[nn[i,samp]]
            dist = 0
            for d in range(n_dim):
                dist += (x_i[d]-y_k[d])**2
            dists[i,samp] = dist
    return dists

@jit(nopython=True, parallel=True, fastmath=True)
def compute_sigma(distances, k, n_iter=64, bandwidth=1.0):
    """Compute a continuous version of the distance to the kth nearest
    neighbor. That is, this is similar to knn-distance but allows continuous
    k values rather than requiring an integral k. In esscence we are simply
    computing the distance such that the cardinality of fuzzy set we generate
    is k.

    Parameters
    ----------
    distances: array of shape (n_samples, n_neighbors)
        Distances to nearest neighbors for each samples. Each row should be a
        sorted list of distances to a given samples nearest neighbors.

    k: float
        The number of nearest neighbors to approximate for.

    n_iter: int (optional, default 64)
        We need to binary search for the correct distance value. This is the
        max number of iterations to use in such a search.

    local_connectivity: int (optional, default 1)
        The local connectivity required -- i.e. the number of nearest
        neighbors that should be assumed to be connected at a local level.
        The higher this value the more connected the manifold becomes
        locally. In practice this should be not more than the local intrinsic
        dimension of the manifold.UMAP

    bandwidth: float (optional, default 1)
        The target bandwidth of the kernel, larger values will produce
        larger return values.

    Returns
    -------
    knn_dist: array of shape (n_samples,)
        The distance to kth nearest neighbor, as suitably approximated.

    nn_dist: array of shape (n_samples,)
        The distance to the 1st nearest neighbor for each point.
    """
    target = np.log2(k) * bandwidth
    result = np.zeros(distances.shape[0], dtype=np.float32)

    for i in range(distances.shape[0]):
        lo = 0.0
        hi = NPY_INFINITY
        mid = 1.0

        for n in range(n_iter):

            psum = 0.0
            for j in range(1, distances.shape[1]):
                d = distances[i, j]
                if d > 0:
                    psum += np.exp(-(d / mid))
                else:
                    psum += 1.0

            if np.fabs(psum - target) < SMOOTH_K_TOLERANCE:
                break

            if psum > target:
                hi = mid
                mid = (lo + hi) / 2.0
            else:
                lo = mid
                if hi == NPY_INFINITY:
                    mid *= 2
                else:
                    mid = (lo + hi) / 2.0

        result[i] = mid

        if result[i] < MIN_K_DIST_SCALE * np.mean(distances):
             result[i] = MIN_K_DIST_SCALE * np.mean(distances)

    return result

def find_grid_nn(X,n_neighbors):
    return NearestNeighbors().fit(X).kneighbors(X, n_neighbors=n_neighbors, return_distance=False)

class StochasticEmbedding(BaseEstimator):

    def __init__(
        self,
        n_neighbors=20,
        n_smoothing_neighbors=None,
        n_smoothing_epochs=25,
        n_centroids=None,
        border_min_dist=1,
        a=None,
        b=None,
        verbose=False,
    ):

        self.n_neighbors = n_neighbors
        self.n_centroids = n_centroids
        self.n_smoothing_neighbors = n_smoothing_neighbors
        self.n_smoothing_epochs = n_smoothing_epochs
        self.border_min_dist = border_min_dist
        self.verbose = verbose
        self.a = a
        self.b = b

    def _validate_parameters(self):
        if self.n_neighbors < 2:
            raise ValueError("n_neighbors must be greater than 2")
        if (self.n_centroids is not None) and ((not isinstance(self.n_centroids, int)) or self.n_centroids <= 0):
            raise ValueError("n_centroids must be a positive integer or None for self estimid")
        if self.border_min_dist <= 0:
            raise ValueError("border_min_dist must be greater then zero")
        if (self.n_smoothing_neighbors is not None) and ((not isinstance(self.n_smoothing_neighbors, int)) or self.n_smoothing_neighbors < 2):
            raise ValueError("n_smoothing_neighbors must be a positive integer greater then 2 or None for self estimid")

    def _fit(self, X, Y, lab=None, direct_adaption=True, eta=0.1, max_itr=500, F=None):
        """Fit X into an embedded space.

        Optionally use y for supervised dimension reduction.

        Parameters
        ----------
        X : array, shape (n_samples, n_features) or (n_samples, n_samples)
            If the metric is 'precomputed' X must be a square distance
            matrix. Otherwise it contains a sample per row. If the method
            is 'exact', X may be a sparse matrix of type 'csr', 'csc'
            or 'coo'.

        y : array, shape (n_samples)
            A target array for supervised dimension reduction. How this is
            handled is determined by parameters UMAP was instantiated with.
            The relevant attributes are ``target_metric`` and
            ``target_metric_kwds``.
        """
        X = check_array(X, dtype=np.float32, accept_sparse="csr")
        Y = check_array(Y, dtype=np.float32, accept_sparse="csr")

        if len(Y.shape) != 2:
            raise ValueError("Shape missmatch, expected (samples,feature)")
        if len(X.shape) != 2:
            raise ValueError("Shape missmatch, expected (samples,components)")

        if X.shape[0] != Y.shape[0]:
            raise ValueError("Shape missmatch, sample")

        #self.s_ = Y
        self.r_ = X

        # Handle all the optional arguments, setting default
        self._validate_parameters()

        if self.a is None or self.b is None:
            self._a, self._b = np.float32(1), np.float32(1)
        else:
            self._a, self._b = np.float32(self.a), np.float32(self.b)

        if self.n_centroids is None:
            if lab is None:
                n_centroids = min(int(np.ceil( Y.shape[0]**0.5 )),Y.shape[0])
            else:
                n = np.unique(lab).shape[0]
                n_centroids = min(int(np.ceil( (Y.shape[0]*n)**0.5 )),Y.shape[0])
        else:
            n_centroids = self.n_centroids
        if self.n_smoothing_neighbors is None:
            self._n_smoothing_neighbors = 2**X.shape[1]+1
        else:
            self._n_smoothing_neighbors = self.n_smoothing_neighbors

        if self.verbose:
            print(str(self))

        if X.shape[0] <= self.n_neighbors:
            #warn(
            #    "n_neighbors is larger than the dataset size; truncating to "
            #    "X.shape[0] - 1"
            #)
            self._n_neighbors = X.shape[0] - 1
        else:
            self._n_neighbors = self.n_neighbors

        if self.verbose:
            print("compute triangolation")
        self.min_,self.max_ = np.min(X,axis=0)-self.border_min_dist,np.max(X,axis=0)+self.border_min_dist
        if lab is None:
            centroieds = KMeans(n_clusters=n_centroids,n_jobs=-1).fit(X).cluster_centers_
        else:
            n = lab.max()+1
            labY = np.zeros( (X.shape[0],n) )
            for i in range(lab.shape[0]):
                labY[i,lab[i]] = 10
            centroieds = KMeans(n_clusters=n_centroids,n_jobs=-1).fit(X,y=lab).cluster_centers_
        border = np.array([[self.min_[0],self.min_[1]],[self.max_[0],self.min_[1]],[self.min_[0],self.max_[1]],[self.max_[0],self.max_[1]]])
        self._centroieds = np.concatenate((border,centroieds), axis=0)

        if self.verbose:
            print("precompute neighborhood")
        neighbors_search = NearestNeighbors().fit(X)
        self._centroied_neighbors = neighbors_search.kneighbors(self._centroieds, n_neighbors=self._n_neighbors, return_distance=False)
        self._triangolation = Delaunay(self._centroieds)

        if self.verbose:
            print("compute sigma")
        """Ynn = neighbors_search.kneighbors(X, n_neighbors=self._n_neighbors, return_distance=False)[:,1:]
        Ydist = dist_each(Y,Y,Ynn)
        self.sigma_ = compute_sigma(Ydist, self._n_neighbors)"""
        self.sigma_ = np.ones(X.shape[0])

        if not(direct_adaption):
            self.s_ = Y
        else:
            if self.verbose:
                print("compute s")
            nn = self.find_nn(self.r_)
            w = weight__(self.r_,self.r_,self.sigma_,nn,self._a,self._b)

            #lambd = 0.1
            Wp = ([],([],[]))
            #diag = set([i for i in range(nn.shape[0])])
            for i in range(nn.shape[0]):
                for j in range(nn.shape[1]):
                    Wp[1][0].append(i)
                    Wp[1][1].append(nn[i,j])
                    Wp[0].append(w[i,j])
                    #if i == nn[i,j]:
                    #    diag.remove(i)
            #for i in diag:
            #    Wp[1][0].append(i)
            #    Wp[1][1].append(i)
            #    Wp[0].append(0)
            W = csr_matrix( Wp, ( self.r_.shape[0], self.r_.shape[0] ), dtype=np.float32)
            W = W.toarray()

            # out = W(x) @ S = (S^T @ W(x)^T)^T
            # Y = W(R) @ S => Y^T = S^T @ W(R)^T
            #              => S = (S^T)^T := (W^T @ W)^-1 @ (W^T @ Y)
            #              => (W^T @ W) S = W^T @ Y

            #W = W + eye(W.shape[0])
            #Wt = W.transpose()
            #G = Wt @ W + lambd * eye(W.shape[0])

            #print( W.toarray(), W.toarray().sum(axis=1), (W.transpose() @ W).toarray() )
            #self.s_ = np.array( lin_solve(W.transpose() @ W , W.transpose() @ Y), dtype=np.float32 )

            """GT,YT = W.transpose(),Y.transpose()
            print(GT.shape)
            GTI = np.linalg.pinv(GT)
            ST = YT @ GTI
            self.s_ = ST.transpose()"""

            """S = np.copy(Y)
            eta = 0.01
            for i in range(1000):
                diff = Y - W @ S
                if i % 10 == 0:
                    err = np.linalg.norm(diff, axis=1)
                    print(err.min(),err.mean(),err.max()," ",err.var(),"   ",diff.shape,W.shape)
                S = S + eta * (W @ diff)
            self.s_ = S"""

            S = np.copy(Y)
            if F is None:
                euc_itr = max_itr
            else:
                euc_itr = int(max_itr/2)

            for i in range(euc_itr):
                E = (W @ S - Y)
                #if i % 50 == 0:
                #    print("eucl err ", np.linalg.norm(E, axis=1).mean())
                S = S - eta * W.transpose() @ E

            if not(F is None):
                buf = np.empty( (X.shape[0], Y.shape[1]) )
                for i in range(max_itr):
                    E = W @ S - Y
                    err = 0
                    for k in range(X.shape[0]):
                        nsqerr = np.dot(F[k],E[k])
                        err += np.inner(nsqerr,nsqerr)**0.5
                        buf[k] = np.dot(F[k].transpose(), nsqerr)
                    err /= X.shape[0]
                    #if i % 50 == 0:
                    #    print("eucl err ", np.linalg.norm((W @ S - Y), axis=1).mean(), ", fish err ", err)
                    grad = W.transpose() @ buf

                    eta_a, eta_b = 0, 0
                    for k in range(X.shape[0]):
                        """print(grad.shape)
                        print(W.shape)
                        print(W[k].shape)
                        print(F[k].shape)
                        print(np.dot(grad.transpose(),W[k].reshape([-1])).shape)
                        print("Ek", E[k].shape) """
                        d = np.dot(F[k],np.dot(grad.transpose(),W[k].reshape([-1])))
                        e = np.dot(F[k],            E[k] )
                        eta_a += np.inner(e,d)
                        eta_b += np.inner(d,d)

                    S = S - eta_a/eta_b * grad

            self.s_ = S
            """err = np.linalg.norm(euclid(X,self.r_,Y,self.sigma_,nn,self._a,self._b )[0] - Y, axis=1)
            print(err.min(),err.mean(),err.max()," ",err.var())
            err = np.linalg.norm(euclid(X,self.r_,self.s_,self.sigma_,nn,self._a,self._b )[0] - Y, axis=1)
            print(err.min(),err.mean(),err.max()," ",err.var())"""

        return self

    def find_nn(self, X):
        return select_neighbors(self._centroied_neighbors[self._triangolation.simplices[self._triangolation.find_simplex(X)]])

    def _transform(self,x,return_H=None):
        if self.verbose:
            print("compute initial embedding")
        nn = self.find_nn(x)
        y,H = euclid(x,self.r_,self.s_,self.sigma_,nn,self._a,self._b )
        if self.verbose:
            print("iterative smoothing")
        y = smooth(y,H,find_grid_nn(x, self._n_smoothing_neighbors),self.n_smoothing_epochs,in_place=False)
        if return_H is None:
            return y
        else:
            if self.verbose:
                print("compute H")
            if return_H is "fast":
                return y, compute_H(x,y,self.r_,self.s_,self.sigma_,nn_=nn,a=self._a,b=self._b,H_=H)
            elif return_H is "all":
                return y, compute_H(x,y,self.r_,self.s_,self.sigma_,nn_=None,a=self._a,b=self._b,H_=H)
            else:
                raise ValueError("Unexpected return_H value, possible values are None, 'fast', 'all'")

reason: import numpy,from scipy,from numba


text:
from collections import OrderedDict
import math
import torch
import torch.nn.functional as F
import torch.nn as nn
from torch.utils import model_zoo
import copy
import numpy as np
from torchvision import utils
from .modules import *

class model(nn.Module):
    def __init__(self, Encoder, num_features, block_channel):
        super(model, self).__init__()
        self.E = Encoder
        self.MFF = MFF(block_channel)
        self.R = R(block_channel)

    def forward(self, x):
        x_block1, x_block2, x_block3, x_block4 = self.E(x)
        x_mff = self.MFF(x_block1, x_block2, x_block3, x_block4,[x_block1.size(2)*2,x_block1.size(3)*2])
        out = self.R(x_mff)
        out = F.upsample(out, size=[228,304], mode='bilinear', align_corners=True)

        return out

ext: py

reason: import numpy


text:
# Copyright 2016 United States Government as represented by the Administrator
# of the National Aeronautics and Space Administration. All Rights Reserved.
#
# Portion of this code is Copyright Geoscience Australia, Licensed under the
# Apache License, Version 2.0 (the "License"); you may not use this file
# except in compliance with the License. You may obtain a copy of the License
# at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# The CEOS 2 platform is licensed under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0.
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

import gdal, osr
import collections
import gc
import numpy as np
import xarray as xr
from datetime import datetime
import collections
from collections import OrderedDict
import hdmedians as hd

import datacube
from . import dc_utilities as utilities
from .dc_utilities import create_default_clean_mask
import hdmedians as hd

"""
Utility Functions
"""

def convert_to_dtype(data, dtype):
    """
    A utility function converting xarray, pandas, or NumPy data to a given dtype.

    Parameters
    ----------
    data: xarray.Dataset, xarray.DataArray, pandas.Series, pandas.DataFrame,
             or numpy.ndarray
    dtype: str or numpy.dtype
        A string denoting a Python datatype name (e.g. int, float) or a NumPy dtype (e.g.
        np.int16, np.float32) to convert the data to.
    """
    if dtype is None: # Don't convert the data type.
        return data
    return data.astype(dtype)

"""
Compositing Functions
"""

def create_mosaic(dataset_in, clean_mask=None, no_data=-9999, dtype=None, intermediate_product=None, **kwargs):
    """
    Creates a most-recent-to-oldest mosaic of the input dataset.

    Parameters
    ----------
    dataset_in: xarray.Dataset
        A dataset retrieved from the Data Cube; should contain:
        coordinates: time, latitude, longitude
        variables: variables to be mosaicked (e.g. red, green, and blue bands)
    clean_mask: np.ndarray
        An ndarray of the same shape as `dataset_in` - specifying which values to mask out.
        If no clean mask is specified, then all values are kept during compositing.
    no_data: int or float
        The no data value.
    dtype: str or numpy.dtype
        A string denoting a Python datatype name (e.g. int, float) or a NumPy dtype (e.g.
        np.int16, np.float32) to convert the data to.

    Returns
    -------
    dataset_out: xarray.Dataset
        Compositited data with the format:
        coordinates: latitude, longitude
        variables: same as dataset_in
    """
    dataset_in = dataset_in.copy(deep=True)

    # Default to masking nothing.
    if clean_mask is None:
        clean_mask = create_default_clean_mask(dataset_in)

    # Mask data with clean_mask. All values where clean_mask==False are set to no_data.
    for key in list(dataset_in.data_vars):
        dataset_in[key].values[np.invert(clean_mask)] = no_data
    if intermediate_product is not None:
        dataset_out = intermediate_product.copy(deep=True)
    else:
        dataset_out = None
    time_slices = reversed(
        range(len(dataset_in.time))) if 'reverse_time' in kwargs and kwargs['reverse_time'] else range(
            len(dataset_in.time))
    for index in time_slices:
        dataset_slice = dataset_in.isel(time=index).drop('time')
        if dataset_out is None:
            dataset_out = dataset_slice.copy(deep=True)
            utilities.clear_attrs(dataset_out)
        else:
            for key in list(dataset_in.data_vars):
                dataset_out[key].values[dataset_out[key].values == -9999] = dataset_slice[key].values[dataset_out[key]
                                                                                .values == -9999]
                dataset_out[key].attrs = OrderedDict()
    utilities.nan_to_num(dataset_out, no_data)
    if dtype is not None:
        convert_to_dtype(dataset_out, dtype)
    return dataset_out

def create_mean_mosaic(dataset_in, clean_mask=None, no_data=-9999, dtype=None, **kwargs):
    """
    Method for calculating the mean pixel value for a given dataset.

    Parameters
    ----------
    dataset_in: xarray.Dataset
        A dataset retrieved from the Data Cube; should contain:
        coordinates: time, latitude, longitude
        variables: variables to be mosaicked (e.g. red, green, and blue bands)
    clean_mask: np.ndarray
        An ndarray of the same shape as `dataset_in` - specifying which values to mask out.
        If no clean mask is specified, then all values are kept during compositing.
    no_data: int or float
        The no data value.
    dtype: str or numpy.dtype
        A string denoting a Python datatype name (e.g. int, float) or a NumPy dtype (e.g.
        np.int16, np.float32) to convert the data to.

    Returns
    -------
    dataset_out: xarray.Dataset
        Compositited data with the format:
        coordinates: latitude, longitude
        variables: same as dataset_in
    """
    # Default to masking nothing.
    if clean_mask is None:
        clean_mask = create_default_clean_mask(dataset_in)

    dataset_in_filtered = dataset_in.where((dataset_in != no_data) & (clean_mask))
    dataset_out = dataset_in_filtered.mean(dim='time', skipna=True, keep_attrs=False)
    utilities.nan_to_num(dataset_out, no_data)
    #manually clear out dates/timestamps/sats.. median won't produce meaningful reslts for these.
    for key in ['timestamp', 'date', 'satellite']:
        if key in dataset_out:
            dataset_out[key].values[::] = no_data
    utilities.nan_to_num(dataset_out, no_data)
    if dtype is not None:
        convert_to_dtype(dataset_out, dtype)
    return dataset_out

def create_median_mosaic(dataset_in, clean_mask=None, no_data=-9999, dtype=None, **kwargs):
    """
    Method for calculating the median pixel value for a given dataset.

    Parameters
    ----------
    dataset_in: xarray.Dataset
        A dataset retrieved from the Data Cube; should contain:
        coordinates: time, latitude, longitude
        variables: variables to be mosaicked (e.g. red, green, and blue bands)
    clean_mask: np.ndarray
        An ndarray of the same shape as `dataset_in` - specifying which values to mask out.
        If no clean mask is specified, then all values are kept during compositing.
    no_data: int or float
        The no data value.
    dtype: str or numpy.dtype
        A string denoting a Python datatype name (e.g. int, float) or a NumPy dtype (e.g.
        np.int16, np.float32) to convert the data to.

    Returns
    -------
    dataset_out: xarray.Dataset
        Compositited data with the format:
        coordinates: latitude, longitude
        variables: same as dataset_in
    """
    # Default to masking nothing.
    if clean_mask is None:
        clean_mask = create_default_clean_mask(dataset_in)

    dataset_in_filtered = dataset_in.where((dataset_in != no_data) & (clean_mask))
    dataset_out = dataset_in_filtered.median(dim='time', skipna=True, keep_attrs=False)
    utilities.nan_to_num(dataset_out, no_data)
    #manually clear out dates/timestamps/sats.. median won't produce meaningful reslts for these.
    for key in ['timestamp', 'date', 'satellite']:
        if key in dataset_out:
            dataset_out[key].values[::] = no_data
    utilities.nan_to_num(dataset_out, no_data)
    if dtype is not None:
        convert_to_dtype(dataset_out, dtype)
    return dataset_out

def create_max_ndvi_mosaic(dataset_in, clean_mask=None, no_data=-9999, dtype=None, intermediate_product=None, **kwargs):
    """
    Method for calculating the pixel value for the max ndvi value.

    Parameters
    ----------
    dataset_in: xarray.Dataset
        A dataset retrieved from the Data Cube; should contain:
        coordinates: time, latitude, longitude
        variables: variables to be mosaicked (e.g. red, green, and blue bands)
    clean_mask: np.ndarray
        An ndarray of the same shape as `dataset_in` - specifying which values to mask out.
        If no clean mask is specified, then all values are kept during compositing.
    no_data: int or float
        The no data value.
    dtype: str or numpy.dtype
        A string denoting a Python datatype name (e.g. int, float) or a NumPy dtype (e.g.
        np.int16, np.float32) to convert the data to.

    Returns
    -------
    dataset_out: xarray.Dataset
        Compositited data with the format:
        coordinates: latitude, longitude
        variables: same as dataset_in
    """
    dataset_in = dataset_in.copy(deep=True)

    # Default to masking nothing.
    if clean_mask is None:
        clean_mask = create_default_clean_mask(dataset_in)

    for key in list(dataset_in.data_vars):
        dataset_in[key].values[np.invert(clean_mask)] = no_data

    if intermediate_product is not None:
        dataset_out = intermediate_product.copy(deep=True)
    else:
        dataset_out = None

    time_slices = range(len(dataset_in.time))
    for timeslice in time_slices:
        dataset_slice = dataset_in.isel(time=timeslice).drop('time')
        ndvi = (dataset_slice.nir - dataset_slice.red) / (dataset_slice.nir + dataset_slice.red)
        ndvi.values[np.invert(clean_mask)[timeslice, ::]] = -1000000000
        dataset_slice['ndvi'] = ndvi
        if dataset_out is None:
            dataset_out = dataset_slice.copy(deep=True)
            utilities.clear_attrs(dataset_out)
        else:
            for key in list(dataset_slice.data_vars):
                dataset_out[key].values[dataset_slice.ndvi.values >
                                        dataset_out.ndvi.values] = dataset_slice[key].values[dataset_slice.ndvi.values >
                                                                                dataset_out.ndvi.values]
    utilities.nan_to_num(dataset_out, no_data)
    if dtype is not None:
        convert_to_dtype(dataset_out, dtype)
    return dataset_out

def create_min_ndvi_mosaic(dataset_in, clean_mask=None, no_data=-9999, dtype=None, intermediate_product=None, **kwargs):
    """
    Method for calculating the pixel value for the min ndvi value.

    Parameters
    ----------
    dataset_in: xarray.Dataset
        A dataset retrieved from the Data Cube; should contain:
        coordinates: time, latitude, longitude
        variables: variables to be mosaicked (e.g. red, green, and blue bands)
    clean_mask: np.ndarray
        An ndarray of the same shape as `dataset_in` - specifying which values to mask out.
        If no clean mask is specified, then all values are kept during compositing.
    no_data: int or float
        The no data value.
    dtype: str or numpy.dtype
        A string denoting a Python datatype name (e.g. int, float) or a NumPy dtype (e.g.
        np.int16, np.float32) to convert the data to.

    Returns
    -------
    dataset_out: xarray.Dataset
        Compositited data with the format:
        coordinates: latitude, longitude
        variables: same as dataset_in
    """
    dataset_in = dataset_in.copy(deep=True)

    # Default to masking nothing.
    if clean_mask is None:
        clean_mask = create_default_clean_mask(dataset_in)

    for key in list(dataset_in.data_vars):
        dataset_in[key].values[np.invert(clean_mask)] = no_data

    if intermediate_product is not None:
        dataset_out = intermediate_product.copy(deep=True)
    else:
        dataset_out = None

    time_slices = range(len(dataset_in.time))
    for timeslice in time_slices:
        dataset_slice = dataset_in.isel(time=timeslice).drop('time')
        ndvi = (dataset_slice.nir - dataset_slice.red) / (dataset_slice.nir + dataset_slice.red)
        ndvi.values[np.invert(clean_mask)[timeslice, ::]] = 1000000000
        dataset_slice['ndvi'] = ndvi
        if dataset_out is None:
            dataset_out = dataset_slice.copy(deep=True)
            utilities.clear_attrs(dataset_out)
        else:
            for key in list(dataset_slice.data_vars):
                dataset_out[key].values[dataset_slice.ndvi.values <
                                        dataset_out.ndvi.values] = dataset_slice[key].values[dataset_slice.ndvi.values <
                                                                                dataset_out.ndvi.values]
    utilities.nan_to_num(dataset_out, no_data)
    if dtype is not None:
        convert_to_dtype(dataset_out, dtype)
    return dataset_out

def unpack_bits(land_cover_endcoding, data_array, cover_type):
    """
	Description:
		Unpack bits for end of ls7 and ls8 functions
	-----
	Input:
		land_cover_encoding(dict hash table) land cover endcoding provided by ls7 or ls8
        data_array( xarray DataArray)
        cover_type(String) type of cover
	Output:
        unpacked DataArray
	"""
    boolean_mask = np.isin(data_array.values, land_cover_endcoding[cover_type])
    return xr.DataArray(boolean_mask.astype(np.int8),
                        coords = data_array.coords,
                        dims = data_array.dims,
                        name = cover_type + "_mask",
                        attrs = data_array.attrs)

def ls8_unpack_qa( data_array , cover_type):

    land_cover_endcoding = dict( fill         =[1] ,
                                 clear        =[322, 386, 834, 898, 1346],
                                 water        =[324, 388, 836, 900, 1348],
                                 shadow       =[328, 392, 840, 904, 1350],
                                 snow         =[336, 368, 400, 432, 848, 880, 812, 944, 1352],
                                 cloud        =[352, 368, 416, 432, 848, 880, 912, 944, 1352],
                                 low_conf_cl  =[322, 324, 328, 336, 352, 368, 834, 836, 840, 848, 864, 880],
                                 med_conf_cl  =[386, 388, 392, 400, 416, 432, 898, 900, 904, 928, 944],
                                 high_conf_cl =[480, 992],
                                 low_conf_cir =[322, 324, 328, 336, 352, 368, 386, 388, 392, 400, 416, 432, 480],
                                 high_conf_cir=[834, 836, 840, 848, 864, 880, 898, 900, 904, 912, 928, 944],
                                 terrain_occ  =[1346,1348, 1350, 1352]
                               )
    return unpack_bits(land_cover_endcoding, data_array, cover_type)

def ls7_unpack_qa( data_array , cover_type):

    land_cover_endcoding = dict( fill     =  [1],
                                 clear    =  [66,  130],
                                 water    =  [68,  132],
                                 shadow   =  [72,  136],
                                 snow     =  [80,  112, 144, 176],
                                 cloud    =  [96,  112, 160, 176, 224],
                                 low_conf =  [66,  68,  72,  80,  96,  112],
                                 med_conf =  [130, 132, 136, 144, 160, 176],
                                 high_conf=  [224]
                               )
    return unpack_bits(land_cover_endcoding, data_array, cover_type)

def ls5_unpack_qa( data_array , cover_type):

    land_cover_endcoding = dict( fill     =  [1],
                                 clear    =  [66,  130],
                                 water    =  [68,  132],
                                 shadow   =  [72,  136],
                                 snow     =  [80,  112, 144, 176],
                                 cloud    =  [96,  112, 160, 176, 224],
                                 low_conf =  [66,  68,  72,  80,  96,  112],
                                 med_conf =  [130, 132, 136, 144, 160, 176],
                                 high_conf=  [224]
                               )
    return unpack_bits(land_cover_endcoding, data_array, cover_type)

def create_hdmedians_multiple_band_mosaic(dataset_in,
                                          clean_mask=None,
                                          no_data=-9999,
                                          dtype=None,
                                          intermediate_product=None,
                                          operation="median",
                                          **kwargs):
    """
    Parameters
    ----------
    dataset_in: xarray.Dataset
        A dataset retrieved from the Data Cube; should contain:
        coordinates: time, latitude, longitude
        variables: variables to be mosaicked (e.g. red, green, and blue bands)
    clean_mask: np.ndarray
        An ndarray of the same shape as `dataset_in` - specifying which values to mask out.
        If no clean mask is specified, then all values are kept during compositing.
    no_data: int or float
        The no data value.
    dtype: str or numpy.dtype
        A string denoting a Python datatype name (e.g. int, float) or a NumPy dtype (e.g.
        np.int16, np.float32) to convert the data to.
    operation: str in ['median', 'medoid']

    Returns
    -------
    dataset_out: xarray.Dataset
        Compositited data with the format:
        coordinates: latitude, longitude
        variables: same as dataset_in
    """
    # Default to masking nothing.
    if clean_mask is None:
        clean_mask = create_default_clean_mask(dataset_in)
    assert operation in ['median', 'medoid'], "Only median and medoid operations are supported."

    dataset_in_filtered = dataset_in.where((dataset_in != no_data) & (clean_mask))

    band_list = list(dataset_in_filtered.data_vars)
    arrays = [dataset_in_filtered[band] for band in band_list]
    stacked_data = np.stack(arrays)
    bands_shape, time_slices_shape, lat_shape, lon_shape = stacked_data.shape[0], \
    stacked_data.shape[1], stacked_data.shape[2], stacked_data.shape[3]
    # Reshape to remove lat/lon
    reshaped_stack = stacked_data.reshape(bands_shape, time_slices_shape,
                                          lat_shape * lon_shape)
    # Build zeroes array across time slices.
    hdmedians_result = np.zeros((bands_shape, lat_shape * lon_shape))

    for x in range(reshaped_stack.shape[2]):
        try:
            hdmedians_result[:, x] = hd.nangeomedian(
                reshaped_stack[:, :, x], axis=1) if operation == "median" else hd.nanmedoid(
                    reshaped_stack[:, :, x], axis=1)
        except ValueError:
            no_data_pixel_stack = reshaped_stack[:, :, x]
            no_data_pixel_stack[np.isnan(no_data_pixel_stack)] = no_data
            hdmedians_result[:, x] = np.full((bands_shape), no_data) if operation == "median" else hd.nanmedoid(
                no_data_pixel_stack, axis=1)
    output_dict = {
        value: (('latitude', 'longitude'), hdmedians_result[index, :].reshape(lat_shape, lon_shape))
        for index, value in enumerate(band_list)
    }
    dataset_out = xr.Dataset(output_dict,
                             coords={'latitude': dataset_in['latitude'], 'longitude': dataset_in['longitude']},
                             attrs = dataset_in.attrs)
    utilities.nan_to_num(dataset_out, no_data)
    if dtype is not None:
        convert_to_dtype(dataset_out, dtype)
    return dataset_out

def create_hdmedians_multiple_band_mosaic(dataset_in,
                                          clean_mask=None,
                                          no_data=-9999,
                                          intermediate_product=None,
                                          operation="median",
                                          **kwargs):
    """
    Description:
    Calculates the medoid using a multi-band processing method.
    -----
    Input:
    dataset_in (xarray dataset) - the set of data with clouds and no data removed.
    Optional Inputs:
    no_data (int/float) - no data value.
    """

    assert clean_mask is not None, "A boolean mask for clean_mask must be supplied."
    assert operation in ['median', 'medoid'], "Only median and medoid operations are supported."

    dataset_in_filtered = dataset_in.where((dataset_in != no_data) & (clean_mask))

    band_list = list(dataset_in_filtered.data_vars)
    arrays = [dataset_in_filtered[band] for band in band_list]

    stacked_data = np.stack(arrays)
    bands_shape, time_slices_shape, lat_shape, lon_shape = stacked_data.shape[0], stacked_data.shape[
        1], stacked_data.shape[2], stacked_data.shape[3]

    reshaped_stack = stacked_data.reshape(bands_shape, time_slices_shape,
                                          lat_shape * lon_shape)  # Reshape to remove lat/lon
    hdmedians_result = np.zeros((bands_shape, lat_shape * lon_shape))  # Build zeroes array across time slices.

    for x in range(reshaped_stack.shape[2]):
        try:
            hdmedians_result[:, x] = hd.nangeomedian(
                reshaped_stack[:, :, x], axis=1) if operation == "median" else hd.nanmedoid(
                    reshaped_stack[:, :, x], axis=1)
        except ValueError:
            no_data_pixel_stack = reshaped_stack[:, :, x]
            no_data_pixel_stack[np.isnan(no_data_pixel_stack)] = no_data
            hdmedians_result[:, x] = np.full((bands_shape), no_data) if operation == "median" else hd.nanmedoid(
                no_data_pixel_stack, axis=1)

    output_dict = {
        value: (('latitude', 'longitude'), hdmedians_result[index, :].reshape(lat_shape, lon_shape))
        for index, value in enumerate(band_list)
    }

    dataset_out = xr.Dataset(
        output_dict, coords={'latitude': dataset_in['latitude'],
                             'longitude': dataset_in['longitude']})
    utilities.nan_to_num(dataset_out, no_data)
    return dataset_out


text:
import os

import edward as ed
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

from mcmc.mcmc import run_experiment, compare_vae_hmc_loss
from mcmc.util import plot

mnist = input_data.read_data_sets('./MNIST_data', one_hot=True)
mb_size = 64
z_dim = 100
X_dim = mnist.train.images.shape[1]
y_dim = mnist.train.labels.shape[1]
h_dim = 128
c = 0
lr = 1e-3

def xavier_init(size):
    in_dim = size[0]
    xavier_stddev = 1. / tf.sqrt(in_dim / 2.)
    return tf.random_normal(shape=size, stddev=xavier_stddev)

# =============================== Q(z|X) Encoder ======================================

Q_W1 = tf.Variable(xavier_init([X_dim, h_dim]))
Q_b1 = tf.Variable(tf.zeros(shape=[h_dim]))

Q_W2_mu = tf.Variable(xavier_init([h_dim, z_dim]))
Q_b2_mu = tf.Variable(tf.zeros(shape=[z_dim]))

Q_W2_sigma = tf.Variable(xavier_init([h_dim, z_dim]))
Q_b2_sigma = tf.Variable(tf.zeros(shape=[z_dim]))

def Q(X, reuse=True):
    with tf.variable_scope('encoder', reuse=reuse):
        h = tf.nn.relu(tf.matmul(X, Q_W1) + Q_b1)
        z_mu = tf.matmul(h, Q_W2_mu) + Q_b2_mu
        z_logvar = tf.matmul(h, Q_W2_sigma) + Q_b2_sigma
        return z_mu, z_logvar

# =============================== Sampling Helper ======================================
def sample_z(mu, log_var):
    eps = tf.random_normal(shape=tf.shape(mu))
    return mu + tf.exp(log_var / 2) * eps

# =============================== P(X|z) Decoder ======================================

P_W1 = tf.Variable(xavier_init([z_dim, h_dim]))
P_b1 = tf.Variable(tf.zeros(shape=[h_dim]))

P_W2 = tf.Variable(xavier_init([h_dim, X_dim]))
P_b2 = tf.Variable(tf.zeros(shape=[X_dim]))

def P(z, reuse=True):
    with tf.variable_scope('decoder', reuse=reuse):
        h = tf.nn.relu(tf.matmul(z, P_W1) + P_b1)
        logits = tf.matmul(h, P_W2) + P_b2
        prob = tf.nn.sigmoid(logits)
        return prob, logits

# =============================== DEFINE LOSS ====================================
X = tf.placeholder(tf.float32, shape=[None, X_dim])
z = tf.placeholder(tf.float32, shape=[None, z_dim])

z_mu, z_logvar = Q(X)
z_sample = sample_z(z_mu, z_logvar)
_, logits = P(z_sample)

# Sampling from random z
X_samples, _ = P(z)

# E[log P(X|z)]
recon_loss = tf.reduce_sum(tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=X), 1)
# D_KL(Q(z|X) || P(z)); calculate in closed form as both dist. are Gaussian
kl_loss = 0.5 * tf.reduce_sum(tf.exp(z_logvar) + z_mu**2 - 1. - z_logvar, 1)
# VAE loss
vae_loss = tf.reduce_mean(recon_loss + kl_loss)

solver = tf.train.AdamOptimizer().minimize(vae_loss)

saver = tf.train.Saver()

# sess = tf.Session()
sess = ed.get_session() # need to make sure tf and edward share the global session
sess.run(tf.global_variables_initializer())

if not os.path.exists('out/'):
    os.makedirs('out/')

# =============================== TRAINING ====================================
i = 0
max_iter = 15000

for it in range(max_iter):
    X_mb, _ = mnist.train.next_batch(mb_size)

    _, loss = sess.run([solver, vae_loss], feed_dict={X: X_mb})

    if it % 1000 == 0:
        print('Iter: {}'.format(it))
        print('Loss: {:.4}'. format(loss))
        print()

        samples = sess.run(X_samples, feed_dict={z: np.random.randn(16, z_dim)})

        fig = plot(samples)
        plt.savefig('./out/{}.png'.format(str(i).zfill(3)), bbox_inches='tight')
        i += 1
        plt.close(fig)

save_path = saver.save(sess, "./models/vae_model.ckpt")
print("Model saved in file: %s" % save_path)

# =============================== INFERENCE ====================================
inference_batch_size = 1
x_gt, _ = mnist.train.next_batch(inference_batch_size)
plot(x_gt)

config = {
    'model': 'hmc',
    'inference_batch_size': 1,
    'T': 10000,
    'img_dim': 28,
    'step_size': None,
    'leapfrog_steps': None,
    'friction': None,
    'z_dim': 100,
    'likelihood_variance': 0.4
}

qz, qz_kept = run_experiment(P, Q, x_gt, config)
compare_vae_hmc_loss(P, Q, x_gt, qz_kept, num_samples=100)

max_stars_repo_name: shenkev/Exact-Inference-VAE-Robustness

max_issues_repo_name: shenkev/Exact-Inference-VAE-Robustness

max_forks_repo_name: shenkev/Exact-Inference-VAE-Robustness


text:
from numpy import linalg
import numpy as np
from scipy.spatial import ConvexHull
import mpl_toolkits.mplot3d as a3
import matplotlib.pyplot as plt
import pypoman as ppm
from shapely.geometry.polygon import Polygon
from models.auv import *
from models.hovercraft import *
from util import *

def plot_env(map_limits, Obstacles):
    fig = plt.figure()
    if len(map_limits) == 2:
        # Configure
        xmin, xmax = map_limits[0]
        ymin, ymax = map_limits[1]
        plt.xlim(xmin-2, xmax+2)
        plt.ylim(ymin-2, ymax+2)
        # # viz obstacles
        for A, b in Obstacles:
            poly = Polygon(ppm.duality.compute_polytope_vertices(A, b))
            x, y = poly.exterior.xy
            plt.fill(x, y, facecolor='r', alpha=0.3)
        ax = fig.gca()
        plt.gca().set_aspect('equal', adjustable='box')
        plt.axis('off')
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
        plt.margins(0, 0)

    elif len(map_limits) == 3:
        ax = fig.add_subplot(111, projection='3d')
        xmin, xmax = map_limits[0]
        ymin, ymax = map_limits[1]
        zmin, zmax = map_limits[2]
        ax.set_xlim3d(xmin, xmax)
        ax.set_ylim3d(ymin, ymax)
        ax.set_zlim3d(zmin, zmax)
        # set_equal TODO
        plt.axis('off')
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
        ax.get_zaxis().set_visible(False)
        plt.margins(0, 0, 0)
        for obs in Obstacles:
            A, b = obs[0], obs[1]
            if len(A) == 4:
                A = np.array(A.tolist() + [[0, 0, -1], [0, 0, 1]])
                b = np.array(b.tolist() + [[-map_limits[2][0]], [map_limits[2][1]]])
            plot_polytope_3d(A, b, ax=ax, color='yellow', trans=0.2)

    return fig, ax

def plot_goals(Goals):
    True
    for A, b in Goals:
        poly = Polygon(ppm.duality.compute_polytope_vertices(A, b))
        x, y = poly.exterior.xy
        plt.fill(x, y, facecolor='g', alpha=0.3)

def plot_thetas(thetas):
    for x, y in thetas:
        A, b = make_rectangle_center(x, y, 5, 5)
        poly = Polygon(ppm.duality.compute_polytope_vertices(A, b))
        x, y = poly.exterior.xy
        plt.fill(x, y, facecolor='blue', alpha=0.3)

def extract_paths(models, ma_thetas, ma_segs):
    # viz initial states and trajectories
    agent_num = len(models)
    paths = []
    for idx in range(agent_num):
        theta = ma_thetas[idx]
        segs = ma_segs[idx]

        if len(theta) == 2:
            ref_x, ref_y, ref_theta, tru_x, tru_y, tru_theta, times =  [], [], [], [], [], [], []
            q0 = theta + [0]
            for seg in segs:
                t, qref, uref = seg
                ref_x = ref_x + [qref[i][0] for i in range(len(qref))]
                ref_y = ref_y + [qref[i][1] for i in range(len(qref))]
                ref_theta = ref_theta + [qref[i][2] for i in range(len(qref))]
                times = times + [t[i] for i in range(len(t))]

                run = models[idx].run_model
                q = run(q0, t, qref, uref)
                q0 = q[-1]
                tru_x = tru_x + [q[i][0] for i in range(len(q)-1)]
                tru_y = tru_y + [q[i][1] for i in range(len(q)-1)]
                tru_theta = tru_theta + [q[i][2] for i in range(len(q)-1)]
            paths.append([ref_x, ref_y, ref_theta, tru_x, tru_y, tru_theta, times])

        elif len(theta) == 3:

            if isinstance(models[0], AUV):
                q0 = theta + [0, 0, 0]
                ref_x, ref_y, ref_z, times = [], [], [], []
                tru_x, tru_y, tru_z, tru_thi, tru_theta, tru_psi = [], [], [], [], [], []
                for seg in segs:
                    t, qref, uref = seg
                    ref_x = ref_x + [qref[i][0] for i in range(len(qref))]
                    ref_y = ref_y + [qref[i][1] for i in range(len(qref))]
                    ref_z = ref_z + [qref[i][2] for i in range(len(qref))]
                    times = times + [t[i] for i in range(len(t))]

                    run = models[idx].run_model
                    q = run(q0, t, qref, uref)
                    q0 = q[-1]
                    tru_x = tru_x + [q[i][0] for i in range(len(q)-1)]
                    tru_y = tru_y + [q[i][1] for i in range(len(q)-1)]
                    tru_z = tru_z + [q[i][2] for i in range(len(q)-1)]
                    tru_thi = tru_thi + [q[i][3] for i in range(len(q)-1)]
                    tru_theta = tru_theta + [q[i][4] for i in range(len(q)-1)]
                    tru_psi = tru_psi + [q[i][5] for i in range(len(q)-1)]

                paths.append([ref_x, ref_y, ref_z,
                              tru_x, tru_y, tru_z, tru_thi, tru_theta, tru_psi, times])

            elif isinstance(models[0], Hovercraft):
                q0 = theta + [0, 0, 0]
                ref_x, ref_y, ref_z, tru_x, tru_y, tru_z, times = [], [], [], [], [], [], []

                for seg in segs:
                    t, qref, uref = seg
                    ref_x = ref_x + [qref[i][0] for i in range(len(qref))]
                    ref_y = ref_y + [qref[i][1] for i in range(len(qref))]
                    ref_z = ref_z + [qref[i][2] for i in range(len(qref))]
                    times = times + [t[i] for i in range(len(t))]

                    run = models[idx].run_model
                    q = run(q0, t, qref, uref)
                    q0 = q[-1]
                    tru_x = tru_x + [q[i][0] for i in range(len(q)-1)]
                    tru_y = tru_y + [q[i][1] for i in range(len(q)-1)]
                    tru_z = tru_z + [q[i][2] for i in range(len(q)-1)]
                paths.append([ref_x, ref_y, ref_z, tru_x, tru_y, tru_z, times])

    return paths

class Faces():
    def __init__(self,tri, sig_dig=12, method="convexhull"):
        self.method=method
        self.tri = np.around(np.array(tri), sig_dig)
        self.grpinx = list(range(len(tri)))
        norms = np.around([self.norm(s) for s in self.tri], sig_dig)
        _, self.inv = np.unique(norms,return_inverse=True, axis=0)

    def norm(self,sq):
        cr = np.cross(sq[2]-sq[0],sq[1]-sq[0])
        return np.abs(cr/np.linalg.norm(cr))

    def isneighbor(self, tr1,tr2):
        a = np.concatenate((tr1,tr2), axis=0)
        return len(a) == len(np.unique(a, axis=0))+2

    def order(self, v):
        if len(v) <= 3:
            return v
        v = np.unique(v, axis=0)
        n = self.norm(v[:3])
        y = np.cross(n,v[1]-v[0])
        y = y/np.linalg.norm(y)
        c = np.dot(v, np.c_[v[1]-v[0],y])
        if self.method == "convexhull":
            h = ConvexHull(c)
            return v[h.vertices]
        else:
            mean = np.mean(c,axis=0)
            d = c-mean
            s = np.arctan2(d[:,0], d[:,1])
            return v[np.argsort(s)]

    def simplify(self):
        for i, tri1 in enumerate(self.tri):
            for j,tri2 in enumerate(self.tri):
                if j > i:
                    if self.isneighbor(tri1,tri2) and \
                       self.inv[i]==self.inv[j]:
                        self.grpinx[j] = self.grpinx[i]
        groups = []
        for i in np.unique(self.grpinx):
            u = self.tri[self.grpinx == i]
            u = np.concatenate([d for d in u])
            u = self.order(u)
            groups.append(u)
        return groups

def plot_polytope_3d(A, b, ax = None, color = 'red', trans = 0.2):
    verts = np.array(ppm.compute_polytope_vertices(A, b))
    # compute the triangles that make up the convex hull of the data points
    hull = ConvexHull(verts)
    triangles = [verts[s] for s in hull.simplices]
    # combine co-planar triangles into a single face
    faces = Faces(triangles, sig_dig=1).simplify()
    # viz
    if ax == None:
        ax = a3.Axes3D(plt.figure())

    pc = a3.art3d.Poly3DCollection(faces,
                                   facecolor=color,
                                   edgecolor="k", alpha=trans)
    ax.add_collection3d(pc)
    # define view
    yllim, ytlim = ax.get_ylim()
    xllim, xtlim = ax.get_xlim()
    zllim, ztlim = ax.get_zlim()
    x = verts[:,0]
    x = np.append(x, [xllim, xtlim])
    y = verts[:,1]
    y = np.append(y, [yllim, ytlim])
    z = verts[:,2]
    z = np.append(z, [zllim, ztlim])
    ax.set_xlim(np.min(x)-1, np.max(x)+1)
    ax.set_ylim(np.min(y)-1, np.max(y)+1)
    ax.set_zlim(np.min(z)-1, np.max(z)+1)

reason: import numpy,from numpy,from scipy


text:
import cv2
import numpy as np

num_down =2
num_bilateral =7

# img_rgb = cv2.imread("avatar.png")
cap = cv2.VideoCapture(0)

while(True):
    # Capture frame-by-frame
    ret, frame = cap.read()

    # Our operations on the frame come here
    img_rgb= cv2.resize(frame, (900,900))

    img_color = img_rgb

    for _ in range(num_down):
        img_color = cv2.pyrDown(img_color)

    for _ in range(num_bilateral):
        img_color = cv2.bilateralFilter(img_color, d=9, sigmaColor=9,sigmaSpace=7)

    for _ in range(num_down):
        img_color = cv2.pyrUp(img_color)

    img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY)
    img_blur = cv2.medianBlur(img_gray, 7)

    img_edge = cv2.adaptiveThreshold(img_blur, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, blockSize=9, C=2)

    img_edge = cv2.cvtColor(img_edge, cv2.COLOR_GRAY2RGB)
    img_cartoon = cv2.bitwise_and(img_color, img_edge)

    stack = np.hstack([img_rgb, img_cartoon])
    cv2.imshow('Stacked Images', stack)
    # Display the resulting frame
    # cv2.imshow('frame',gray)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# When everything done, release the capture
cap.release()
cv2.destroyAllWindows()


text:
import time
import fire
import kitti_common as kitti
from eval import get_official_eval_result, get_coco_eval_result
import pdb
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pickle
import os

def _read_imageset_file(path):
    with open(path, 'r') as f:
        lines = f.readlines()
    return [int(line) for line in lines]

def evaluate(label_path,
             result_path,
             save_path,
             current_class=['Car', 'Pedestrian', 'Cyclist'],
             coco=False,
             score_thresh=-1):
    class_to_name = {0: 'Car',
                     1: 'Pedestrian',
                     2: 'Cyclist',
                     3: 'DontCare'}

    gt_annos = kitti.get_label_annos(label_path)
    dt_annos = kitti.get_label_annos(result_path)
    # visualize(gt_annos, dt_annos)
    print(len(dt_annos))
    if score_thresh > 0:
        dt_annos = kitti.filter_annos_low_score(dt_annos, score_thresh)

    if coco:
        print(get_coco_eval_result(dt_annos, dt_annos, current_class))
    else:
        result_str, _ = get_official_eval_result(gt_annos, dt_annos, current_class, class_to_name)
        print(result_str)
        with open(save_path, 'w+') as f:
            f.write("\n")
            f.write(result_str)

def analyze(front_path, frontleft_path, left_path, save_path):

    gt_front = kitti.get_label_annos(front_path)
    gt_frontleft = kitti.get_label_annos(frontleft_path)
    gt_left = kitti.get_label_annos(left_path)
    gt_dataset = [gt_front, gt_frontleft, gt_left]

    annos_name = ["Front", "FrontLeft", "Left"]
    split = "left_pred_"
    for i in range(0, 3):
        count = {'Car': 0, 'Pedestrian': 0, 'Cyclist': 0, 'DontCare': 0}
        ranges = {'Car': [], 'Pedestrian': [], 'Cyclist': [], 'DontCare': []}
        rot_y = {'Car': [], 'Pedestrian': [], 'Cyclist': [], 'DontCare': []}
        for frame in gt_dataset[i]:
            for j in range(0, len(frame['name'])):
                cls = frame['name'][j]
                distance = frame['location'][j][2]
                ranges[cls].append(distance)
                ry = frame['rotation_y'][j] * (180 / np.pi)
                rot_y[cls].append(ry)
                count[cls] += 1
        print("Count of classes:" + str(count) + '\n')
        with open(save_path + split + "classes.txt", 'a+') as f:
            f.write(annos_name[i] + ' ' + str(count) + '\n')
        sns_plot = sns.distplot(ranges["Car"], color="skyblue", label="Car").set_title(
            annos_name[i] + " Car Distances")
        fig = sns_plot.get_figure()
        fig.savefig("output/eval/dataset_vis/" + split + annos_name[i] + "_car.png")
        plt.clf()

        sns_plot_ped = sns.distplot(ranges["Pedestrian"], color="red", label="Pedestrian").set_title(
            annos_name[i] + " Pedestrian Distances")
        fig_ped = sns_plot_ped.get_figure()
        fig_ped.savefig("output/eval/dataset_vis/" + split + annos_name[i] + "_ped.png")
        plt.clf()

        sns_plot_cyc = sns.distplot(ranges["Cyclist"], color="teal", label="Cyclist").set_title(
            annos_name[i] + " Cyclist Distances")
        fig_cyc = sns_plot_cyc.get_figure()
        fig_cyc.savefig("output/eval/dataset_vis/" + split + annos_name[i] + "_cyc.png")
        plt.clf()

        sns_plot = sns.distplot(rot_y["Car"], color="skyblue", label="Car").set_title(annos_name[i] + " Car Rotation")
        fig = sns_plot.get_figure()
        fig.savefig("output/eval/dataset_vis/" + split + "ry_" + annos_name[i] + "_car.png")
        plt.clf()

        sns_plot_ped = sns.distplot(rot_y["Pedestrian"], color="red", label="Pedestrian").set_title(
            annos_name[i] + " Pedestrian Rotation")
        fig_ped = sns_plot_ped.get_figure()
        fig_ped.savefig("output/eval/dataset_vis/" + split + "ry_" + annos_name[i] + "_ped.png")
        plt.clf()

        sns_plot_cyc = sns.distplot(rot_y["Cyclist"], color="teal", label="Cyclist").set_title(
            annos_name[i] + " Cyclist Rotation")
        fig_cyc = sns_plot_cyc.get_figure()
        fig_cyc.savefig("output/eval/dataset_vis/" + split + "ry_" + annos_name[i] + "_cyc.png")
        plt.clf()
    #print("Count of ranges:" + str(ranges)+'\n')
    #print("Count of ry:" + str(rot_y)+'\n')

def training_analyze(front_path, frontleft_path, left_path):

    gt_front = kitti.get_label_annos(front_path)
    gt_frontleft = kitti.get_label_annos(frontleft_path)
    gt_left = kitti.get_label_annos(left_path)
    gt_dataset = [gt_front, gt_frontleft, gt_left]

    x = ['Car', 'Pedestrian', 'Cyclist', 'DontCare']
    annos_name = ["Front", "FrontLeft", "Left"]
    count = {"Front": [0, 0, 0, 0], "FrontLeft": [0, 0, 0, 0], "Left": [0, 0, 0, 0]}

    for i in range(0, 3):
        view = annos_name[i]
        for frame in gt_dataset[i]:
            view_count = count[view]
            for j in range(0, len(frame['name'])):
                cls = frame['name'][j]
                view_count[x.index(cls)] += 1

    # set width of bar
    barWidth = 0.25

    # set height of bar
    bars1 = count["Front"]
    bars2 = count["FrontLeft"]
    bars3 = count["Left"]

    # Set position of bar on X axis
    r1 = np.arange(len(bars1))
    r2 = [x + barWidth for x in r1]
    r3 = [x + barWidth for x in r2]

    # Make the plot
    plt.bar(r1, bars1, color='#7f6d5f', width=barWidth, edgecolor='white', label='Front')
    plt.bar(r2, bars2, color='#557f2d', width=barWidth, edgecolor='white', label='FrontLeft')
    plt.bar(r3, bars3, color='#2d7f5e', width=barWidth, edgecolor='white', label='Left')

    # Add xticks on the middle of the group bars
    plt.xlabel('group', fontweight='bold')
    plt.xticks([r + barWidth for r in range(len(bars1))], ['Car', 'Pedestrian', 'Cyclist', 'Sign'])

    plt.title("Class Count Across Camera Views")
    # Create legend & Show graphic
    plt.legend()
    plt.savefig("output/eval/dataset_vis/train_class_count.png")

def training_density_analysis(front_path, frontleft_path, left_path):

    gt_front = kitti.get_label_annos(front_path)
    gt_frontleft = kitti.get_label_annos(frontleft_path)
    gt_left = kitti.get_label_annos(left_path)
    gt_dataset = [gt_front, gt_frontleft, gt_left]

    x = ['Car', 'Pedestrian', 'Cyclist', 'DontCare']
    annos_name = ["Front", "FrontLeft", "Left"]
    count = {"Front": [0, 0, 0, 0], "FrontLeft": [0, 0, 0, 0], "Left": [0, 0, 0, 0]}

    total_images = 0
    total_car = 0
    for i in range(0, 1):
        view = annos_name[i]
        for frame in gt_dataset[i]:
            view_count = count[view]
            if len(frame['name']) != 0:
            	total_images += 1
            for j in range(0, len(frame['name'])):
                cls = frame['name'][j]
                if (cls) == 'Car' or (cls) == 'Pedestrian' or (cls) == 'Cyclist' :
                	total_car += 1
    print(str(total_car/total_images))

    # set width of bar
    barWidth = 0.25

    # set height of bar
    bars1 = count["Front"]
    bars2 = count["FrontLeft"]
    bars3 = count["Left"]

    # Set position of bar on X axis
    r1 = np.arange(len(bars1))
    r2 = [x + barWidth for x in r1]
    r3 = [x + barWidth for x in r2]

    # Make the plot
    plt.bar(r1, bars1, color='#7f6d5f', width=barWidth, edgecolor='white', label='Front')
    plt.bar(r2, bars2, color='#557f2d', width=barWidth, edgecolor='white', label='FrontLeft')
    plt.bar(r3, bars3, color='#2d7f5e', width=barWidth, edgecolor='white', label='Left')

    # Add xticks on the middle of the group bars
    plt.xlabel('group', fontweight='bold')
    plt.xticks([r + barWidth for r in range(len(bars1))], ['Car', 'Pedestrian', 'Cyclist', 'Sign'])

    plt.title("Class Count Across Camera Views")
    # Create legend & Show graphic
    plt.legend()
    plt.savefig("output/eval/dataset_vis/train_class_count.png")

def visualize(gt_annos, dt_annos):
    sns.set()
    ranges = {'Car': [], 'Pedestrian': [], 'Cyclist': [], 'DontCare': []}
    for frame in gt_annos:
        for i in range(0, len(frame['location'])):
            distance = frame['location'][i][2]
            cls = frame['name'][i]
            ranges[cls].append(distance)

    sns_plot = sns.distplot(ranges["Car"], color="skyblue", label="Car")
    fig = sns_plot.get_figure()
    fig.savefig("output/eval/dataset_vis/val_car.png")
    plt.clf()

    sns_plot_ped = sns.distplot(ranges["Pedestrian"], color="red", label="Pedestrian")
    fig_ped = sns_plot_ped.get_figure()
    fig_ped.savefig("output/eval/dataset_vis/val_ped.png")
    plt.clf()

    sns_plot_cyc = sns.distplot(ranges["Cyclist"], color="teal", label="Cyclist")
    fig_cyc = sns_plot_cyc.get_figure()
    fig_cyc.savefig("output/eval/dataset_vis/val_cyc.png")
    plt.clf()

    ranges = {'Car': [], 'Pedestrian': [], 'Cyclist': [], 'DontCare': []}
    for frame in dt_annos:
        for i in range(0, len(frame['location'])):
            distance = frame['location'][i][2]
            cls = frame['name'][i]
            ranges[cls].append(distance)

    sns_plot = sns.distplot(ranges["Car"], color="skyblue", label="Car")
    fig = sns_plot.get_figure()
    fig.savefig("output/eval/dataset_vis/pred_car.png")
    plt.clf()

    sns_plot_ped = sns.distplot(ranges["Pedestrian"], color="red", label="Pedestrian")
    fig_ped = sns_plot_ped.get_figure()
    fig_ped.savefig("output/eval/dataset_vis/pred_ped.png")
    plt.clf()

    sns_plot_cyc = sns.distplot(ranges["Cyclist"], color="teal", label="Cyclist")
    fig_cyc = sns_plot_cyc.get_figure()
    fig_cyc.savefig("output/eval/dataset_vis/pred_cyc.png")
    breakpoint()

def analyze_tp():
    training_views = ['front', 'frontleft', 'left']
    eval_views = ['front', 'frontleft', 'left']
    for train in training_views:
        for evalu in eval_views:
            name = 'output/eval/detections/' + train + '_' + evalu + '.pkl'
            with open(name, 'rb') as f:
                tp_annos = pickle.load(f)

                fig, axes = plt.subplots(ncols=2)
                ranges = [x[2] for x in tp_annos['locs']]
                sns.distplot(ranges, color="skyblue", ax=axes[0]).set_title("Distance from car")
                ry = [x * (180 / np.pi) for x in tp_annos["ry"]]
                sns.distplot(ry, color="red", ax=axes[1]).set_title("Rotation_y of detections")
                fig.savefig("output/eval/detections/" + train + '_' + evalu + '.png')
                plt.clf()

                ry = frame['rotation_y'][j] * (180 / np.pi)
                rot_y[cls].append(ry)
                count[cls] += 1
        print("Count of classes:" + str(count) + '\n')
        with open(save_path + split + "classes.txt", 'a+') as f:
            f.write(annos_name[i] + ' ' + str(count) + '\n')
        sns_plot = sns.distplot(ranges["Car"], color="skyblue", label="Car").set_title(
            annos_name[i] + " Car Distances")
        fig = sns_plot.get_figure()
        fig.savefig("output/eval/dataset_vis/" + split + annos_name[i] + "_car.png")
        plt.clf()

        sns_plot_ped = sns.distplot(ranges["Pedestrian"], color="red", label="Pedestrian").set_title(
            annos_name[i] + " Pedestrian Distances")
        fig_ped = sns_plot_ped.get_figure()
        fig_ped.savefig("output/eval/dataset_vis/" + split + annos_name[i] + "_ped.png")
        plt.clf()

        sns_plot_cyc = sns.distplot(ranges["Cyclist"], color="teal", label="Cyclist").set_title(
            annos_name[i] + " Cyclist Distances")
        fig_cyc = sns_plot_cyc.get_figure()
        fig_cyc.savefig("output/eval/dataset_vis/" + split + annos_name[i] + "_cyc.png")
        plt.clf()

        sns_plot = sns.distplot(rot_y["Car"], color="skyblue", label="Car").set_title(annos_name[i] + " Car Rotation")
        fig = sns_plot.get_figure()
        fig.savefig("output/eval/dataset_vis/" + split + "ry_" + annos_name[i] + "_car.png")
        plt.clf()

        sns_plot_ped = sns.distplot(rot_y["Pedestrian"], color="red", label="Pedestrian").set_title(
            annos_name[i] + " Pedestrian Rotation")
        fig_ped = sns_plot_ped.get_figure()
        fig_ped.savefig("output/eval/dataset_vis/" + split + "ry_" + annos_name[i] + "_ped.png")
        plt.clf()

        sns_plot_cyc = sns.distplot(rot_y["Cyclist"], color="teal", label="Cyclist").set_title(
            annos_name[i] + " Cyclist Rotation")
        fig_cyc = sns_plot_cyc.get_figure()
        fig_cyc.savefig("output/eval/dataset_vis/" + split + "ry_" + annos_name[i] + "_cyc.png")
        plt.clf()
    #print("Count of ranges:" + str(ranges)+'\n')
    #print("Count of ry:" + str(rot_y)+'\n')

def analyze_velocity():
    CALIB_PATH = "data/waymo/training/img_calib"
    path, dirs, files = next(os.walk(CALIB_PATH))
    files.sort()
    velocity = []
    angular_velocity = []
    hs_count = 0
    for file in files:
        f = open(CALIB_PATH + '/' + file)
        lines = f.readlines()
        v0 = lines[5].split(' ')[1:]
        cam0_vel = np.array([float(x) for x in v0])
        cam0_v = np.linalg.norm(cam0_vel[:3])
        cam0_w = np.linalg.norm(cam0_vel[3:])
        velocity.append(cam0_v*3.6)
        angular_velocity.append(cam0_w)
        if cam0_v*3.6 > 40:
        	hs_count +=1

    print("Percentage is: " + str(hs_count/len(files)))
    sns_plot_cyc = sns.distplot(velocity, color="teal", label="Velocity").set_title("Velocity distribution in Training Set")
    fig_cyc = sns_plot_cyc.get_figure()
    fig_cyc.savefig("output/eval/dataset_vis/velocity.png")
    plt.clf()

    sns_plot_cyc = sns.distplot(angular_velocity, color="teal", label="Angular Velocity").set_title("Ang Velocity Distribution in Training Set")
    fig_cyc = sns_plot_cyc.get_figure()
    fig_cyc.savefig("output/eval/dataset_vis/ang_velocity.png")
    plt.clf()

if __name__ == '__main__':
    fire.Fire()


text:
# The MIT License (MIT)

# Copyright (c) 2020-2021 CoML

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Rachid RIAD, Hadrien TITEUX, Léopold FAVRE

import logging
from typing import Union, Iterable, Callable

import numpy as np
import numpy.random
from pyannote.core import Segment
from sortedcontainers import SortedSet

from .continuum import Annotator, Continuum

class CorpusShufflingTool:
    """
    Corpus shuffling tool as detailed in section 6.3 of the gamma paper
    (https://www.aclweb.org/anthology/J15-3003.pdf#page=30).
    """
    SHIFT_FACTOR = 2
    SPLIT_FACTOR = 2.5
    FALSE_POS_FACTOR = 1

    def __init__(self,
                 magnitude: float,
                 reference_continuum: Continuum,
                 categories: Iterable[str] = None):
        """
        Parameters
        ----------
        magnitude:
            magnitude m of the cst (cf gamma paper)
        reference_continuum:
            this continuum will serve as reference for the tweaks made by the corpus shuffling tool.
        categories:
            this is used to consider additionnal categories when shuffling the corpus, in the eventuality that the
            reference continuum does not contain any unit of a possible category.
        """
        self.magnitude: float = magnitude
        reference_annotators = reference_continuum.annotators
        if len(reference_annotators) > 1:
            logging.warning("Warning : a reference continuum with multiple annotators was given to the CST, so "
                            "its first annotator in alphabetical order will be used as reference.")
        self._reference_annotator: Annotator = reference_annotators[0]
        self._reference_continuum: Continuum = reference_continuum
        self._categories: SortedSet = self._reference_continuum.categories
        if categories is not None:
            for category in categories:
                self._categories.add(category)

    def corpus_from_reference(self, new_annotators: Union[int, Iterable[Annotator]]):
        # TODO: add docstring
        continuum = Continuum()
        continuum._categories = self._reference_continuum.categories
        continuum.bound_inf, continuum.bound_sup = self._reference_continuum.bounds
        if isinstance(new_annotators, int):
            new_annotators = [f"annotator_{i}" for i in range(new_annotators)]
        for unit in self._reference_continuum.iter_annotator(self._reference_annotator):
            for new_annotator in new_annotators:
                continuum.add(new_annotator,
                              Segment(unit.segment.start, unit.segment.end),
                              unit.annotation)
        return continuum

    def shift_shuffle(self, continuum: Continuum) -> None:
        """
        Tweaks the given continuum by shifting the ends of each segment, with uniformly distributed values
        of bounds proportionnal to the magnitude of the CST and the length of the segment.
        """
        shift_max = self.magnitude * self.SHIFT_FACTOR * \
            self._reference_continuum.avg_length_unit
        for annotator in continuum.annotators:
            for unit in continuum[annotator]:
                continuum.remove(annotator, unit)
                start_seg, end_seg = 0.0, 0.0
                while start_seg >= end_seg:
                    start_seg = unit.segment.start + np.random.uniform(-1, 1) * shift_max
                    end_seg = unit.segment.end + np.random.uniform(-1, 1) * shift_max
                continuum.add(annotator, Segment(start_seg, end_seg), unit.annotation)

    def false_neg_shuffle(self, continuum: Continuum) -> None:
        """
        Tweaks the continuum by randomly removing units ("false negatives").
        Every unit (for each annotator) have a probability equal to the magnitude of being removed.
        If this probability is one, a single random unit (for each annotator) will be left alone.
        """
        for annotator in continuum.annotators:
            security = np.random.choice(continuum._annotations[annotator])
            # security : if an annotator doesnt have any annotations gamma cant be computed.
            for unit in list(continuum[annotator]):
                if np.random.random() < self.magnitude:
                    continuum.remove(annotator, unit)
            if len(continuum._annotations[annotator]) == 0:
                continuum.add(annotator, security.segment, security.annotation)

    def false_pos_shuffle(self, continuum: Continuum) -> None:
        """
        Tweaks the continuum by randomly adding "false positive" units.
        The number of added units per annotator is constant & proportionnal to the magnitude of the CST.
        The chosen category is random and depends on the probability of occurence of the category in the reference.
        The length of the segment is random (normal distribution) based on the average and standard deviation
        of those of the reference.
        """
        ref_units = self._reference_continuum[self._reference_annotator]
        avg_dur = np.average([unit.segment.end - unit.segment.start for unit in ref_units])
        var_dur = np.std([unit.segment.end - unit.segment.start for unit in ref_units])
        category_weights = self._reference_continuum.category_weights
        bounds_inf, bounds_sup = self._reference_continuum.bound_inf, self._reference_continuum.bound_sup
        for annotator in continuum.annotators:
            for _ in range(int(self.magnitude * self.FALSE_POS_FACTOR * len(self._reference_continuum))):
                # a random unit is generated from a (all random) central point, duration, and category
                category = np.random.choice(category_weights.keys(), p=category_weights.values())
                center = np.random.uniform(bounds_inf, bounds_sup)
                duration = abs(np.random.normal(avg_dur, var_dur))
                continuum.add(annotator,
                              Segment(center - duration / 2, center + duration / 2),
                              annotation=category)

    def category_shuffle(self, continuum: Continuum,
                         overlapping_fun: Callable[[str, str], float] = None,
                         prevalence: bool = False):
        """
        Shuffles the categories of the annotations in the given continuum using the process described in
        section 3.3.5 of https://hal.archives-ouvertes.fr/hal-00769639/.
        Parameters
        ----------
        overlapping_fun:
            gives the "categorical distance" between two annotations, which is taken into account when provided.
            (the lower the distance between categories, the higher the chance one will be changed into the other).
        prevalence:
            specify whether or not to consider the proportion of presence of each category in the reference.
        """
        category_weights = self._reference_continuum.category_weights
        # matrix "A"
        nb_categories = len(category_weights)
        prob_matrix = np.eye(nb_categories)
        # matrix "B or C"
        if prevalence:
            sec_matrix = np.ones(nb_categories) / nb_categories
        else:
            sec_matrix = np.array([list(category_weights.values())] * nb_categories)

        categories = list(category_weights.keys())
        if overlapping_fun is None:
            # this formula was deduced from the graphs.
            prob_matrix = prob_matrix * (1 - self.magnitude ** 2) + sec_matrix * self.magnitude ** 2
        else:
            overlapping_matrix = np.zeros((len(categories), len(categories)))
            for id1, cat1 in enumerate(categories):
                sum_line = 0
                for id2, cat2 in enumerate(categories):
                    elem = overlapping_fun(cat1, cat2)
                    sum_line += elem
                    overlapping_matrix[id1, id2] = elem
                overlapping_matrix[id1] /= sum_line
            # this formula was also deduced from the graphs.
            prob_matrix = (prob_matrix * (1 - self.magnitude)
                           + sec_matrix * self.magnitude ** 3
                           + overlapping_matrix * (self.magnitude - self.magnitude ** 3)
                           )
        for annotator in continuum.annotators:
            for unit in list(continuum[annotator]):
                continuum.remove(annotator, unit)
                try:
                    new_category = np.random.choice(categories, p=prob_matrix[category_weights.index(unit.annotation)])
                except ValueError as e:
                    raise e
                continuum.add(annotator, Segment(unit.segment.start, unit.segment.end), new_category)
                del unit

    def splits_shuffle(self, continuum: Continuum):
        """
        Tweak the continuum by randomly splitting segments.
        Number of splits per annotator is constant & proportionnal to the magnitude of the CST
        and the number of units in the reference.
        A splitted segment can be re-splitted.
        """
        for _ in range(int(self.magnitude *
                           self.SPLIT_FACTOR *
                           self._reference_continuum.avg_num_annotations_per_annotator)):
            for annotator in continuum.annotators:
                units = continuum._annotations[annotator]
                to_split = units.pop(numpy.random.randint(0, len(units)))
                security = (to_split.segment.end - to_split.segment.start) * 0.01
                cut = numpy.random.uniform(to_split.segment.start + security, to_split.segment.end)

                try:
                    continuum.add(annotator, Segment(cut, to_split.segment.end), to_split.annotation)
                    continuum.add(annotator, Segment(to_split.segment.start, cut), to_split.annotation)
                except ValueError:
                    continuum.add(annotator, to_split.segment, to_split.annotation)
                    continuum.add(annotator, to_split.segment, to_split.annotation)

    def corpus_shuffle(self,
                       annotators: Union[int, Iterable[str]],
                       shift: bool = False,
                       false_pos: bool = False,
                       false_neg: bool = False,
                       split: bool = False,
                       cat_shuffle: bool = False,
                       include_ref: bool = False
                       ) -> Continuum:
        """
        Generates a new shuffled corpus with the provided (or generated) reference annotation set,
        using the method described in 6.3 of the gamma paper, https://www.aclweb.org/anthology/J15-3003.pdf#page=30
        (and missing elements described in another article : https://hal.archives-ouvertes.fr/hal-00769639/).
        """
        continuum = self.corpus_from_reference(annotators)
        if shift:
            self.shift_shuffle(continuum)
        if false_pos:
            self.false_pos_shuffle(continuum)
        if false_neg:
            self.false_neg_shuffle(continuum)
        if cat_shuffle:
            self.category_shuffle(continuum)
        if split:
            self.splits_shuffle(continuum)
        if include_ref:
            assert self._reference_annotator not in continuum.annotators, \
                "Reference annotator can't be included as " \
                "an annotator with the same name is in the " \
                "generated corpus."
            for unit in self._reference_continuum[next(iter(self._reference_continuum.annotators))]:
                continuum.add(self._reference_annotator, unit.segment, unit.annotation)
        return continuum

max_stars_repo_name: valentinoli/pygamma-agreement

max_issues_repo_name: valentinoli/pygamma-agreement

max_forks_repo_name: valentinoli/pygamma-agreement