text:
#***********************************************************************************************
#                               MÓDULO MÉTODOS COMPUTACIONAIS
#***********************************************************************************************

header = '''
File name: MetodosComputacionais.py
Package folder: <nome da pasta do código>
Description: Arquivo com os algoritmos mais usados em cálculo numérico
Homepage: https://github.com/LeandroTeodoroRJ/CalculoComPython
Stable: Yes
Version: 1.0
Last Update: 27.11.20
Current: Yes
Maintainer: leandroteodoro.rj@gmail.com
Contributor(special thanks): Prof. Eric Amâncio (UNESA)
Depends: numpy, sympy, matplotlib
Architecture: PC-32, PC-64
Compile/Interpreter: Python 3.5.4
Access: Public
Changelog: No
Readme: No
Documents:	No
Links: No
Files:  * Características de funcionamento          Caracteristicas_e_Limitacoes.txt
        e limitações dos métodos compu-
        tacionais.
Other Notes: No
Code Structs Comments:
    FUNCTIONS:
        help(void): --Retorna instruções de ajuda do módulo
        exemplos(void): --Retorna exemplos de uso do módulo

    CLASSE BISSECAO
        init(self, equacao) :: equacao[Symbolic] -> [void]  --Instância um objeto com variável da equação em 'x'
        set_limites(self, inf, sup) :: inf[float], sup[float] -> [void]   --Define os Limites Superiores e Inferiores
        gera_grafico(self) :: [void] -> [void]  --Gera o gráfico da função dentro do intervalo
        checa_raiz(self) :: [void] -> [Bool]  --Checa para saber se houve raiz no intervalo
        monotona(self) :: [void] -> [void]  --Gera gráfico da função para saber se é monótona
        retorna_iteracao(self, tol) :: tol[Float] -> [Int]  --Retorna o provável número de iterações para a tolerância
        calcula_zero(self, tol, N) :: tol[Float], N[Int] -> [Float]  --Retorna o zero da função

    CLASSE NEWTON :: <- Ext classe bissecao
        calcula_zero(self, tol, N) :: tol[Float], N[Int] -> [Float]  --Retorna o zero da função pelo método de Newton
        retorna_iteracao(self, tol) :: --Método não implementado

    CLASSE JACOBI
        init(self, A, b, x0) :: A, b, x0[numpy array] -> [void]  --Cria o objeto com as respectivas matrizes de Jacobi
                                                                   Onde:
                                                                   A: Matriz dos coeficientes
                                                                   b: Matriz dos termos independentes
                                                                   x0: Matriz de iteração inicial
        calcula(self, TOL, N) :: TOL[float], N[int] -> [numpy array]  --Calcula o sist. linear onde TOL é a tolerância
                                                                        e N é o número máximo de iterações
        retorna_erro(self) :: [void] -> [float]  --Retorna o erro depois que calculado a resulução do sistema
        retorna_iteração(self) :: [void] -> [int] --Retorna o número de iterações que foi necessário
        pivotemanto_parcial(self) :: [void] -> [void]  --Executa o pivoteamento parcial que gera outro sistema linear
                                                         equivalente quando o original não converge para uma solução
        retorna_matriz_A(self) :: [void] -> [np.array]  --Retorna a matriz A
        retorna_matriz_b(self) :: [void] -> [np.array]  --Retorna a matriz b

    CLASSE GAUSSSEIDEL :: <- Ext classe Jacobi
        calcula(self, TOL, N) :: TOL[float], N[int] -> [numpy array]  --Calcula o sist. linear onde TOL é a tolerância
                                                                        e N é o número máximo de iterações pelo Método
                                                                        de Gauss-Seidel
    CLASSE DADOS
        init(self, ex, ey) :: ex[np.array], ey[np.array] -> [void]  --Cria um objeto a partir de pontos coletados, onde
                                                                      ex: São os ponto do eixo x, e
                                                                      ey: São os pontos do eixo y.
        set_limites(self, inf, sup) :: inf[float], sup[float] -> [void]   --Define os Limites Superiores e Inferiores,
                                                                            os pontos devem pertencer a lista ex.
        gera_grafico(self) :: [void] -> [void]  --Gera o gráfico da função dentro do intervalo
        integral(self) :: [void] -> [float]  --Retorna o valor da integral definida dentros dos limites especificados
        get_valores_x(self) :: [void] -> [numpy array]  --Retorna os valores do eixo X
        get_valores_y(self) :: [void] -> [numpy array]  --Retorna os valores do eixo Y
        legendas(self, titulo, label_x, label_y) :: titulo[string], label_x[string], label_y[string] --> [void]
                                                    -- Informa as legandas para o gráfico

    CLASSE INTERPOLAÇÃO POLINOMIAL
        init(self, objDados) :: objDados[Dados] -> [void]  --Gera um objeto do tipo Interpolação Polinomial
        regressao(self) :: [void] -> [void]  --Faz os cálculos necessários para a regressão
        get_poli(self) ::  [void]  -> [symbol]   --Retorna o polinômio em x
        matriz_dd(self) :: [void] -> [numpy array]  --Retorna matriz dos coeficientes do polinômio não simplificado
        poli_x(self)  :: [void]  -> [symbol]   --Retorna o polinômio simplificado em 'x'
        matriz(self) :: [void] -> [numpy array]  --Retorna matriz da Diferença Dividida
        gera_grafico(self) :: [void] -> [void]  --Gera o gráfico do polinômio

'''

def help():
    print(header)

#** MÓDULOS NECESSÁRIOS **
import numpy as np
import matplotlib.pyplot as plt
from sympy import *
from numpy import linalg

#Variáveis simbólicas
from sympy.abc import x

#** MÉTODO DA BISSEÇÃO **
class bissecao(object):        #Nome da classe(classe Pai)

#** Constructor ***
    def __init__(self, equacao):
        self._exp = equacao
        self._limite_sup = 0
        self._limite_inf = 0

# ** Destructor **
    def __del__(self):
        pass

#** Métodos de Classe **
    def set_limites(self, inf, sup):
        self._limite_inf = inf
        self._limite_sup = sup

    def gera_grafico(self):
        self.intervalo = np.linspace(self._limite_inf, self._limite_sup)  #Cria o intervalo de plotagem
        self._eixo_y = []   #Cria uma lista de valores para o eixo y de plotagem
        for i in self.intervalo:
            self._eixo_y.append(self._exp.subs(x, i).evalf())  #Resolve a função
        plt.plot(self.intervalo, self._eixo_y)
        plt.title("Gráfico da Função")
        plt.grid()
        plt.show()

    def checa_raiz(self):   #Checa se o intervalo passou pela raiz
        if ((np.sign(self._exp.subs(x, self._limite_inf).evalf())*np.sign(self._exp.subs(x, self._limite_sup).evalf())) < 0):
            return True
        else:
            return False

    def monotona(self):
        self.func_diff = diff(self._exp, x, 1)
        self.intervalo = np.linspace(self._limite_inf, self._limite_sup)  # Cria o intervalo de plotagem
        self._eixo_y = []  # Cria uma lista de valores para o eixo y de plotagem
        for i in self.intervalo:
            self._eixo_y.append(self.func_diff.subs(x, i).evalf())  # Resolve a função
        plt.plot(self.intervalo, self._eixo_y)
        plt.title("Gráfico Função Monótona")
        plt.grid()
        plt.show()

    def retorna_iteracao(self, tol):  #Retorna o número provável de iterações
        return floor(abs((log((self._limite_sup - self._limite_inf) / tol, 2)).evalf()))

    def calcula_zero(self, tol, N):
        self.i = 1  # Esse seria o início da iteração
        self.fa = self._exp.subs(x, self._limite_inf).evalf()  # O valor da variável f(x) em "a"
        while (self.i <= N):     # Definindo as repetições
            self.k = self._limite_inf + (self._limite_sup - self._limite_inf) / 2    #Ponto médio dos intervalo
            self.fk = self._exp.subs(x, self.k).evalf()
            if ((self.fk == 0) or (np.absolute(self._limite_sup - self._limite_inf) / 2) < tol): # Condições de parada
                return self.k
            self.i = self.i + 1  # Vamos iniciar mais uma iteração
            if (self.fa * self.fk > 0):  # Aqui estamos analisando se o zero da função está do lado
                # direito de 'k', caso ele esteja vamos arrastar o limite da
                # esquerda parao valor k. O limite do intervalo será entr 'k'
                # e 'b'.
                self._limite_inf = self.k
                self.fa = self.fk
            else:  # Em caso contrário, o intervalo o limite intervalo será entre 'a'
                # 'k'
                self._limite_sup = self.k
        raise NameError('Número máximo de iterações foi excedido!!!')

#** MÉTODO DE NEWTON-RAPHSON **
class Newton(bissecao):

# ** Métodos de Classe **
    def calcula_zero(self, tol, N):
        self._x0 = self._limite_sup
        self.func_diff = diff(self._exp, x, 1)
        #x1 = x0 - funcao(x0)/derivada(x0)
        self._x1 = self._x0 - (self._exp.subs(x, self._x0)/self.func_diff.subs(x, self._x0))
        self._i = 1
        while ((abs( self._x1 - self._x0) > tol) and (self._i <= N)):
            self._x0 = self._x1
            self._x1 = self._x0 - (self._exp.subs(x, self._x0) / self.func_diff.subs(x, self._x0))
            self._i = self._i + 1
        if (self._i > N):
            print('Nao houve convergencia!')
            return 0
        else:
            self._x1 = self._x1.evalf()  #Calcula a raíz
            return self._x1

    def retorna_iteracao(self, tol):  # Retorna o número provável de iterações
        print('Não implementado.')

# ** CLASSE DADOS **
class Dados(object):
    def __init__(self, ex, ey):
        self._ex = ex
        self._ey = ey
        self._ltitulo = 'Gráfico da Função'
        self._leixo_x = ''
        self._leixo_y = ''

    def __del__(self):
        pass

    def set_limites(self, inf, sup):
        self._limite_inf = inf
        self._limite_sup = sup

    def integral(self):    #Aproximação do retângulo
        self.i = self._ex.index(self._limite_inf)
        self._area = 0
        while self.i < (self._ex.index(self._limite_sup)):
            self._area = self._area + ((self._ex[self.i+1]-self._ex[self.i])*self._ey[self.i])
            self.i = self.i+1
        return self._area

    def gera_grafico(self):
        self.intervalo = np.linspace(self._limite_inf, self._limite_sup)  #Cria o intervalo de plotagem
        for self.i in self._ex:
            if self.i == self._limite_inf:
                self._menor = self._ex.index(self.i)
            if self.i == self._limite_sup:
                self._maior = self._ex.index(self.i)
        self.intervalo = self._ex[self._menor:self._maior+1]
        self._eixo_y = self._ey[self._menor:self._maior+1]
        plt.plot(self.intervalo, self._eixo_y)
        plt.title(self._ltitulo)
        plt.xlabel(self._leixo_x)
        plt.ylabel(self._leixo_y)
        plt.grid()
        plt.show()

    def get_valores_x(self):
        return self._ex

    def get_valores_y(self):
        return self._ey

    def legendas(self, titulo, label_x, label_y):
        self._ltitulo = titulo
        self._leixo_x = label_x
        self._leixo_y = label_y

#** MÉTODO DE JACOBI **
class Jacobi(object):

#** Constructor ***
    def __init__(self, A, b, x0):
        # Tratamento preliminar das matrizes
        self._A = A.astype('double')
        self._b = b.astype('double')
        self._x0 = x0.astype('double')
        self._erro = 0
        self._it = 0

# ** Destructor **
    def __del__(self):
        pass

#** Métodos de Classe **
    def calcula(self, TOL, N):
        self._n = np.shape(self._A)[0]  # Vai atribuir a incógnita "n" o valor do número de linhas
        # da matriz A.
        self._x = np.copy(self._x0)
        # Iterações de Jacobi
        while (self._it < N):
            self._it = self._it + 1
            for self._i in np.arange(self._n):
                self._x[self._i] = self._b[self._i]
                for self._j in np.concatenate((np.arange(0, self._i), np.arange(self._i + 1, self._n))):
                    self._x[self._i] -= self._A[self._i, self._j] * self._x0[self._j]
                self._x[self._i] /= self._A[self._i, self._i]
                # Tolerância
                if (np.linalg.norm(self._x - self._x0, np.inf) / np.linalg.norm(self._x0, np.inf) < TOL):
                    self._erro = np.linalg.norm(self._x - self._x0, np.inf) / np.linalg.norm(self._x0, np.inf)
                    return self._x
                self._x0 = np.copy(self._x)
        raise NameError('Número máximo de iterações foi ultrapassado!!')

    def retorna_erro(self):
        return self._erro

    def retorna_iteração(self):
        return self._it

    def pivoteamento_parcial(self):
        # Acesse as linhas
        for self._i in range(len(self._A)):
            # Verifique qual o maior pivô na na coluna de análise
            self._pivo = np.fabs(self._A[self._i][self._i])
            self._linhapivo = self._i
            for self._j in range(self._i + 1, len(self._A)):
                if np.fabs(self._A[self._j][self._i]) > self._pivo:
                    self._pivo = np.fabs(self._A[self._j][self._i])
                    self._linhapivo = self._j
            # Permute as linhas
            if self._linhapivo != self._i:
                self._linhaAuxiliar = self._A[self._i]
                self._A[self._i] = self._A[self._linhapivo]
                self._A[self._linhapivo] = self._linhaAuxiliar

                self._bAuxiliar = self._b[self._i]
                self._b[self._i] = self._b[self._linhapivo]
                self._b[self._linhapivo] = self._bAuxiliar
            # Eliminação Gaussiana
            for self._m in range(self._i + 1, len(self._A)):
                self._multiplicador = self._A[self._m][self._i] / self._A[self._i][self._i]
                for self._n in range(self._i, len(self._A)):
                    self._A[self._m][self._n] -= self._multiplicador * self._A[self._i][self._n]
                self._b[self._m] -= self._multiplicador * self._b[self._i]
        # Printar matriz escalonada e o vetor b escalonado
        self._A_resp = np.empty((len(self._A), len(self._A)))
        for self._k in range(len(self._A)):
            self._A_resp[self._k] = self._A[self._k]
        self._A = self._A_resp
        print('\n')

    def retorna_matriz_A(self):
        return self._A

    def retorna_matriz_b(self):
        return self._b

#** MÉTODO DE JACOBI **
class GaussSeidel(Jacobi):
    def calcula(self, TOL, N):
        self._n = np.shape(self._A)[0]  # Vai atribuir a incógnita "n" o valor do número de linhas
        # da matriz A.
        self._x = np.copy(self._x0)
        while (self._it < N):
            self._it = self._it + 1
            # Iterações de Gauss-Seidel
            for self._i in np.arange(self._n):
                self._x[self._i] = self._b[self._i]
                for self._j in np.concatenate((np.arange(0, self._i), np.arange(self._i + 1, self._n))):
                    self._x[self._i] -= self._A[self._i, self._j] * self._x[self._j]
                self._x[self._i] /= self._A[self._i, self._i]

            # tolerância
            if (np.linalg.norm(self._x - self._x0, np.inf) < TOL):
                self._E = np.linalg.norm(self._x - self._x0, np.inf)
                return self._x
            # preparando nova iteração
            self._x0 = np.copy(self._x)
        raise NameError('Número máximo de iterações ultrapassado!!!')

#** MÉTODO DE INTERPOLAÇÃO POLINOMIAL **
class InterpolacaoPolinomial(object):

#** Constructor ***
    def __init__(self, objDados):
        # Insira os dados de entrada
        self._xi = objDados.get_valores_x()
        self._fi = objDados.get_valores_y()

# ** Destructor **
    def __del__(self):
        pass

# ** Médodos de Classe **
    def regressao(self):
        # PROCEDIMIENTO
        # Tabela para armazenar os resultados do método
        self._titulo = ['i', 'xi', 'fi']

        self._n = len(self._xi)  # Comprimento da linha xi

        self._ki = np.arange(0, self._n, 1)  # Criando uma matriz linha com 4 elementos

        self._matriz = np.concatenate(([self._ki], [self._xi], [self._fi]), axis=0)  # Organizando todos os valore
        # em uma matriz, matriz linha
        # compõe uma linha.

        self._matriz = np.transpose(self._matriz)  # Transponha a matriz para obter o formato onde
        # onde "i", "xi" e "fi" cada um represente uma
        # coluna.

        self._dd = np.zeros(shape=(self._n, self._n), dtype=float)  # Criando uma matriz vazia a qual
        # atribuiremos valores futuros

        self._matriz = np.concatenate((self._matriz, self._dd), axis=1)  # Agora estamos agrupando uma matriz
        # quadrada vazia "n x n" a matriz
        # anterior. Como se trata junção de
        # matrizes axis = 1.

        # Calcula matriz, inicia-se na coluna 4 (lembre-se que 0 é a primeira coluna)
        # Esse valor precisa ser mudado quando estamos trabalhando com polinômio maior
        [self._n, self._m] = np.shape(self._matriz)
        self._coluna = self._n - 1
        self._j = 3
        while (self._j < self._m):
            # Adicionar título para cada coluna (utilizamos j - 2 pq iniciamos com
            # j = 3)
            self._titulo.append('self._dd' + str(self._j - 2))

            self._espassar = self._j - 2  # Servirá para adicionar os valores nas colunas vazias
            self._i = 0
            while (self._i < self._coluna):
                self._numerador = self._matriz[self._i + 1, self._j - 1] - self._matriz[self._i, self._j - 1]
                self._denominador = self._xi[self._i + self._espassar] - self._xi[self._i]
                self._matriz[self._i, self._j] = self._numerador / self._denominador
                self._i = self._i + 1
            self._coluna = self._coluna - 1
            self._j = self._j + 1

        # POLINOMIO com diferencias divididas
        # caso: pontos equidistantes no eixo x
        self._h = self._xi[1] - self._xi[0]

        self._dd = self._matriz[0, 3:]  # pega os valores da primeira linha e quarta coluna

        self._n = len(self._dd)

        # Expressão do polinômio com o Sympy
        #x = sym.Symbol('x')
        self._polinomio = self._fi[0]
        for self._j in range(1, self._n, 1):
            self._fator = self._dd[self._j - 1]
            self._termino = 1
            for self._k in range(0, self._j, 1):
                self._termino = self._termino * (x - self._xi[self._k])
            self._polinomio = self._polinomio + self._termino * self._fator
        # Simplifique multiplicando (x-self._xi)
        self._polisimple = self._polinomio.expand()

        # polinômio para avaliação numérica
        self._px = lambdify(x, self._polisimple)

        # Pontos para o gráfico
        self._amostra = 101
        self._a = np.min(self._xi)
        self._b = np.max(self._xi)
        self._pxi = np.linspace(self._a, self._b, self._amostra)
        self._pfi = self._px(self._pxi)

    def get_poli(self):
        return self._polinomio

    def poli_x(self):
        return self._polisimple

    def matriz_dd(self):
        #print([self._titulo])
        return self._dd

    def matriz(self):
        return self._matriz

    def gera_grafico(self):
        # Gráfico
        plt.plot(self._xi, self._fi, 'o', label='Pontos')
        ##for i in range(0,self._n,1):
        ##    plt.axvline(self._xi[i],ls='--', color='yellow')
        plt.plot(self._pxi, self._pfi, label='Interpolação Polinomial')
        plt.legend()
        plt.xlabel('xi')
        plt.ylabel('fi')
        plt.title('Interpolação Diferenças Divididas de Newton')
        plt.show()

extext = '''
# EXEMPLO DE USO DA BIBLIOTECA DE MÉTODOS COMPUTACIONAIS

import MetodosComputacionais as mc
import numpy as np
#import matplotlib.pyplot as plt
from sympy import *

#Variáveis simbólicas
from sympy.abc import x

#Confugurando o ambiente
init_printing()

#Método da Bisseção
#Passo 1 - Criando o gráfico da função
print("Exemplo 1")
func = exp(x)-x-2
pprint(func)
exemp1 = mc.bissecao(func)
exemp1.set_limites(-2, 0)
exemp1.gera_grafico()

#Passo 2 - Saber se o intervalo passou pela raiz
print("Passou pela raiz?")
print(exemp1.checa_raiz())

#Passo 3 - Saber se a função é monotona no intervalo
print("É monótona no intervalo?")
exemp1.monotona()

#Passo 4 - Retornar o múmero de iterações
print("Qual o provável número de iterações?")
print(exemp1.retorna_iteracao(0.1))

#Passo 5 - Calcular o Zero da função
print("Calculando o Zero.")
print(exemp1.calcula_zero(0.1, 5))

#MÉTODO DE JACOBI PARA SISTEMAS LINEARES
print('Resolução de Sistemas Lineares Pelo método de Jacobi')

Sistema em questão:
(I)     -3.x1+x2+x3=2
(II)    2.x1+5.x2+x3=5
(III)   2.x1+3.x2+7.x3=-17

#Declarando as matrizes
A = np.array ([[-3,1,1],[2,5,1],[2,3,7]])   #Matriz dos coeficientes
b = np.array ([[2],[5],[-17]])              #Matriz dos termos isolados
x0 = np.array ([[1],[1],[-1]])              #Vetor inicial (arbitrado)

#Declarando a tolerência e o número máximo de iterações
TOL = 10**-2
N = 50
ex1 = mc.Jacobi(A, b, x0)
print('Matriz resultado')
print(ex1.calcula(TOL, N))
print('Retornando o erro')
print(ex1.retorna_erro())
print('Retornando o número de iterações que foi necessário')
print(ex1.retorna_iteração())

#MÉTODO DE GAUSS SEIDEL PARA SISTEMAS LINEARES
A = np.array ([[10,2,1],[1,5,1],[2,3,10]])
b = np.array ([[7],[-8],[6]])
x0 = np.array([[1],[1],[1]])
TOL = 0.01
N = 50
ex1 = mc.GaussSeidel(A, b, x0)
print('Matriz resultado')
print(ex1.calcula(TOL, N))

# **** MÉTODO DA INTERPOLAÇÃO POLINOMIAL ****
import numpy as np
import MetodosComputacionais as mc

# Insirindo os dados de entrada (pontos)
xi = np.array([-1, 0, 1, 3])    #Pontos no eixo x
fi = np.array([3, 1, 3, 43])    #Pondos do eixo y (imagem)

pontos = mc.Dados(xi, fi)
polinomial = mc.InterpolacaoPolinomial(pontos)

polinomial.regressao()

# Dados de Saída
print('Matriz Diferenças Divididas:')
mtzdd = polinomial.matriz()
print('[[i, xi, fi, dd1, dd2, dd3, dd4]]')
print(mtzdd)
print(type(mtzdd))

print('')
print('Polinomio não simplificado: ')
poli = polinomial.get_poli()
print(poli)
print(type(poli))

print('')
print('Matriz dos coeficientes do polinômio não simplificado:')
print('dd: ')
mdd = polinomial.matriz_dd()
print(mdd)
print(type(mdd))

print('')
print('polinomio simplificado: ')
poli = polinomial.poli_x()
print(poli)
print(type(poli))

#Gera o gráfico do polinimio
polinomial.gera_grafico()

'''

def exemplos():         #Imprime os exemplos
    print(extext)

max_stars_repo_name: LeandroTeodoroRJ/CalculoComPython

max_issues_repo_name: LeandroTeodoroRJ/CalculoComPython

max_forks_repo_name: LeandroTeodoroRJ/CalculoComPython

reason: import numpy,from numpy,from sympy


text:
# pylint: disable=bad-continuation
"""
Defines the Mediation test class
"""

from pprint import pprint

import numpy as np
import pandas as pd
from pandas.api.types import is_float_dtype
from pygam import LinearGAM, s

from causal_curve.core import Core
from causal_curve.utils import rand_seed_wrapper

class Mediation(Core):
    """
    Given three continuous variables (a treatment or independent variable of interest,
    a potential mediator, and an outcome variable of interest), Mediation provides a method
    to determine the average direct and indirect effect.

    Parameters
    ----------

    treatment_grid_num: int, optional (default = 10)
        Takes the treatment, and creates a quantile-based grid across its values.
        For instance, if the number 6 is selected, this means the algorithm will only take
        the 6 treatment variable values at approximately the 0, 20, 40, 60, 80, and 100th
        percentiles to estimate the causal dose response curve.
        Higher value here means the final curve will be more finely estimated,
        but also increases computation time. Default is usually a reasonable number.

    lower_grid_constraint:  float, optional(default = 0.01)
        This adds an optional constraint of the lower side of the treatment grid.
        Sometimes data near the minimum values of the treatment are few in number
        and thus generate unstable estimates. By default, this clips the bottom 1 percentile
        or lower of treatment values. This can be as low as 0, indicating there is no
        lower limit to how much treatment data is considered.

    upper_grid_constraint: float, optional (default = 0.99)
        See above parameter. Just like above, but this is an upper constraint.
        By default, this clips the top 99th percentile or higher of treatment values.
        This can be as high as 1.0, indicating there is no upper limit to how much
        treatment data is considered.

    bootstrap_draws: int, optional (default = 500)
        Bootstrapping is used as part of the mediation test. The parameter determines
        the number of draws from the original data to create a single bootstrap replicate.

    bootstrap_replicates: int, optional (default = 100)
        Bootstrapping is used as part of the mediation test. The parameter determines
        the number of bootstrapping runs to perform / number of new datasets to create.

    spline_order: int, optional (default = 3)
        Order of the splines to use fitting the final GAM.
        Must be integer >= 1. Default value creates cubic splines.

    n_splines: int, optional (default = 5)
        Number of splines to use for the mediation and outcome GAMs.
        Must be integer >= 2. Must be non-negative.

    lambda_: int or float, optional (default = 0.5)
        Strength of smoothing penalty. Must be a positive float.
        Larger values enforce stronger smoothing.

    max_iter: int, optional (default = 100)
        Maximum number of iterations allowed for the maximum likelihood algo to converge.

    random_seed: int, optional (default = None)
        Sets the random seed.

    verbose: bool, optional (default = False)
        Determines whether the user will get verbose status updates.

    Attributes
    ----------

    grid_values: array of shape (treatment_grid_num, )
        The gridded values of the treatment variable. Equally spaced.

    Methods
    ----------
    fit: (self, T, M, y)
        Fits the trio of relevant variables using generalized additive models.

    calculate_effects: (self, ci)
        Calculates the average direct and indirect effects.

    Examples
    --------
    >>> from causal_curve import Mediation
    >>> med = Mediation(treatment_grid_num = 200, random_seed = 512)
    >>> med.fit(T = df['Treatment'], M = df['Mediator'], y = df['Outcome'])
    >>> med_results = med.calculate_effects(0.95)

    References
    ----------

    Imai K., Keele L., Tingley D. A General Approach to Causal Mediation Analysis. Psychological
    Methods. 15(4), 2010, pp.309–334.

    """

    def __init__(
        self,
        treatment_grid_num=10,
        lower_grid_constraint=0.01,
        upper_grid_constraint=0.99,
        bootstrap_draws=500,
        bootstrap_replicates=100,
        spline_order=3,
        n_splines=5,
        lambda_=0.5,
        max_iter=100,
        random_seed=None,
        verbose=False,
    ):

        self.treatment_grid_num = treatment_grid_num
        self.lower_grid_constraint = lower_grid_constraint
        self.upper_grid_constraint = upper_grid_constraint
        self.bootstrap_draws = bootstrap_draws
        self.bootstrap_replicates = bootstrap_replicates
        self.spline_order = spline_order
        self.n_splines = n_splines
        self.lambda_ = lambda_
        self.max_iter = max_iter
        self.random_seed = random_seed
        self.verbose = verbose

        # Validate the params
        self._validate_init_params()
        rand_seed_wrapper()

        if self.verbose:
            print("Using the following params for the mediation analysis:")
            pprint(self.get_params(), indent=4)

    def _validate_init_params(self):
        """
        Checks that the params used when instantiating mediation tool are formatted correctly
        """

        # Checks for treatment_grid_num
        if not isinstance(self.treatment_grid_num, int):
            raise TypeError(
                f"treatment_grid_num parameter must be an integer, "
                f"but found type {type(self.treatment_grid_num)}"
            )

        if (isinstance(self.treatment_grid_num, int)) and self.treatment_grid_num < 4:
            raise ValueError(
                f"treatment_grid_num parameter should be >= 4 so the internal models "
                f"have enough resolution, but found value {self.treatment_grid_num}"
            )

        if (isinstance(self.treatment_grid_num, int)) and self.treatment_grid_num > 100:
            raise ValueError(f"treatment_grid_num parameter is too high!")

        # Checks for lower_grid_constraint
        if not isinstance(self.lower_grid_constraint, float):
            raise TypeError(
                f"lower_grid_constraint parameter must be a float, "
                f"but found type {type(self.lower_grid_constraint)}"
            )

        if (
            isinstance(self.lower_grid_constraint, float)
        ) and self.lower_grid_constraint < 0:
            raise ValueError(
                f"lower_grid_constraint parameter cannot be < 0, "
                f"but found value {self.lower_grid_constraint}"
            )

        if (
            isinstance(self.lower_grid_constraint, float)
        ) and self.lower_grid_constraint >= 1.0:
            raise ValueError(
                f"lower_grid_constraint parameter cannot >= 1.0, "
                f"but found value {self.lower_grid_constraint}"
            )

        # Checks for upper_grid_constraint
        if not isinstance(self.upper_grid_constraint, float):
            raise TypeError(
                f"upper_grid_constraint parameter must be a float, "
                f"but found type {type(self.upper_grid_constraint)}"
            )

        if (
            isinstance(self.upper_grid_constraint, float)
        ) and self.upper_grid_constraint <= 0:
            raise ValueError(
                f"upper_grid_constraint parameter cannot be <= 0, "
                f"but found value {self.upper_grid_constraint}"
            )

        if (
            isinstance(self.upper_grid_constraint, float)
        ) and self.upper_grid_constraint > 1.0:
            raise ValueError(
                f"upper_grid_constraint parameter cannot > 1.0, "
                f"but found value {self.upper_grid_constraint}"
            )

        # Checks for bootstrap_draws
        if not isinstance(self.bootstrap_draws, int):
            raise TypeError(
                f"bootstrap_draws parameter must be a int, "
                f"but found type {type(self.bootstrap_draws)}"
            )

        if (isinstance(self.bootstrap_draws, int)) and self.bootstrap_draws < 100:
            raise ValueError(
                f"bootstrap_draws parameter cannot be < 100, "
                f"but found value {self.bootstrap_draws}"
            )

        if (isinstance(self.bootstrap_draws, int)) and self.bootstrap_draws > 500000:
            raise ValueError(
                f"bootstrap_draws parameter cannot > 500000, "
                f"but found value {self.bootstrap_draws}"
            )

        # Checks for bootstrap_replicates
        if not isinstance(self.bootstrap_replicates, int):
            raise TypeError(
                f"bootstrap_replicates parameter must be a int, "
                f"but found type {type(self.bootstrap_replicates)}"
            )

        if (
            isinstance(self.bootstrap_replicates, int)
        ) and self.bootstrap_replicates < 50:
            raise ValueError(
                f"bootstrap_replicates parameter cannot be < 50, "
                f"but found value {self.bootstrap_replicates}"
            )

        if (
            isinstance(self.bootstrap_replicates, int)
        ) and self.bootstrap_replicates > 100000:
            raise ValueError(
                f"bootstrap_replicates parameter cannot > 100000, "
                f"but found value {self.bootstrap_replicates}"
            )

        # Checks for lower_grid_constraint isn't higher than upper_grid_constraint
        if self.lower_grid_constraint >= self.upper_grid_constraint:
            raise ValueError(
                "lower_grid_constraint should be lower than upper_grid_constraint!"
            )

        # Checks for spline_order
        if not isinstance(self.spline_order, int):
            raise TypeError(
                f"spline_order parameter must be an integer, "
                f"but found type {type(self.spline_order)}"
            )

        if (isinstance(self.spline_order, int)) and self.spline_order < 3:
            raise ValueError(
                f"spline_order parameter should be >= 1, but found {self.spline_order}"
            )

        if (isinstance(self.spline_order, int)) and self.spline_order >= 30:
            raise ValueError(f"spline_order parameter is too high!")

        # Checks for n_splines
        if not isinstance(self.n_splines, int):
            raise TypeError(
                f"n_splines parameter must be an integer, but found type {type(self.n_splines)}"
            )

        if (isinstance(self.n_splines, int)) and self.n_splines < 2:
            raise ValueError(
                f"n_splines parameter should be >= 2, but found {self.n_splines}"
            )

        if (isinstance(self.n_splines, int)) and self.n_splines >= 100:
            raise ValueError(f"n_splines parameter is too high!")

        # Checks for lambda_
        if not isinstance(self.lambda_, (int, float)):
            raise TypeError(
                f"lambda_ parameter must be an int or float, but found type {type(self.lambda_)}"
            )

        if (isinstance(self.lambda_, (int, float))) and self.lambda_ <= 0:
            raise ValueError(
                f"lambda_ parameter should be >= 2, but found {self.lambda_}"
            )

        if (isinstance(self.lambda_, (int, float))) and self.lambda_ >= 1000:
            raise ValueError(f"lambda_ parameter is too high!")

        # Checks for max_iter
        if not isinstance(self.max_iter, int):
            raise TypeError(
                f"max_iter parameter must be an int, but found type {type(self.max_iter)}"
            )

        if (isinstance(self.max_iter, int)) and self.max_iter <= 10:
            raise ValueError(
                f"max_iter parameter is too low! Results won't be reliable!"
            )

        if (isinstance(self.max_iter, int)) and self.max_iter >= 1e6:
            raise ValueError(f"max_iter parameter is unnecessarily high!")

        # Checks for random_seed
        if not isinstance(self.random_seed, (int, type(None))):
            raise TypeError(
                f"random_seed parameter must be an int, but found type {type(self.random_seed)}"
            )

        if (isinstance(self.random_seed, int)) and self.random_seed < 0:
            raise ValueError(f"random_seed parameter must be > 0")

        # Checks for verbose
        if not isinstance(self.verbose, bool):
            raise TypeError(
                f"verbose parameter must be a boolean type, but found type {type(self.verbose)}"
            )

    def _validate_fit_data(self):
        """Verifies that T, M, and y are formatted the right way"""
        # Checks for T column
        if not is_float_dtype(self.T):
            raise TypeError(f"Treatment data must be of type float")

        # Checks for M column
        if not is_float_dtype(self.M):
            raise TypeError(f"Mediation data must be of type float")

        # Checks for Y column
        if not is_float_dtype(self.y):
            raise TypeError(f"Outcome data must be of type float")

    def _grid_values(self):
        """Produces initial grid values for the treatment variable"""
        return np.quantile(
            self.T,
            q=np.linspace(
                start=self.lower_grid_constraint,
                stop=self.upper_grid_constraint,
                num=self.treatment_grid_num,
            ),
        )

    def _collect_mean_t_levels(self):
        """Collects the mean treatment value within each treatment bucket in the grid_values"""

        t_bin_means = []

        for index, _ in enumerate(self.grid_values):
            if index == (len(self.grid_values) - 1):
                continue

            t_bin_means.append(
                self.T[
                    (
                        (self.T >= self.grid_values[index])
                        & (self.T <= self.grid_values[index + 1])
                    )
                ].mean()
            )

        return t_bin_means

    def fit(self, T, M, y):
        """Fits models so that mediation analysis can be run.
        For now, this only accepts pandas columns.

        Parameters
        ----------
        T: array-like, shape (n_samples,)
            A continuous treatment variable
        M: array-like, shape (n_samples,)
            A continuous mediation variable
        y: array-like, shape (n_samples,)
            A continuous outcome variable

        Returns
        ----------
        self : object

        """
        self.T = T.reset_index(drop=True, inplace=False)
        self.M = M.reset_index(drop=True, inplace=False)
        self.y = y.reset_index(drop=True, inplace=False)

        # Validate this input data
        self._validate_fit_data()

        self.n = len(y)

        # Create grid_values
        self.grid_values = self._grid_values()

        # Loop through the comparisons in the grid_values
        if self.verbose:
            print("Beginning main loop through treatment bins...")

        # Collect loop results in this list
        self.final_bootstrap_results = []

        # Begin main loop
        for index, _ in enumerate(self.grid_values):
            if index == 0:
                continue
            if self.verbose:
                print(
                    f"***** Starting iteration {index} of {len(self.grid_values) - 1} *****"
                )

            temp_low_treatment = self.grid_values[index - 1]
            temp_high_treatment = self.grid_values[index]

            bootstrap_results = self._bootstrap_analysis(
                temp_low_treatment, temp_high_treatment
            )

            self.final_bootstrap_results.append(bootstrap_results)

    def calculate_mediation(self, ci=0.95):
        """Conducts mediation analysis on the fit data

        Parameters
        ----------
        ci: float (default = 0.95)
            The desired bootstrap confidence interval to produce. Default value is 0.95,
            corresponding to 95% confidence intervals. bounded (0, 1.0).

        Returns
        ----------
        dataframe: Pandas dataframe
            Contains the estimate of the direct and indirect effects
            and the proportion of indirect effects across the treatment grid values.
            The bootstrap confidence interval that is returned might not be symmetric.

        self : object

        """

        # Collect effect results in these lists
        self.t_bin_means = self._collect_mean_t_levels()
        self.prop_direct_list = []
        self.prop_indirect_list = []
        general_indirect = []

        lower = (1 - ci) / 2
        upper = ci + lower

        # Calculate results for each treatment bin
        for index, _ in enumerate(self.grid_values):

            if index == (len(self.grid_values) - 1):
                continue

            temp_bootstrap_results = self.final_bootstrap_results[index]

            mean_results = {
                key: temp_bootstrap_results[key].mean()
                for key in temp_bootstrap_results
            }

            tau_coef = (
                mean_results["d1"]
                + mean_results["d0"]
                + mean_results["z1"]
                + mean_results["z0"]
            ) / 2
            n0 = mean_results["d0"] / tau_coef
            n1 = mean_results["d1"] / tau_coef
            n_avg = (n0 + n1) / 2

            tau_general = (
                temp_bootstrap_results["d1"]
                + temp_bootstrap_results["d0"]
                + temp_bootstrap_results["z1"]
                + temp_bootstrap_results["z0"]
            ) / 2
            nu_0_general = temp_bootstrap_results["d0"] / tau_general
            nu_1_general = temp_bootstrap_results["d1"] / tau_general
            nu_avg_general = (nu_0_general + nu_1_general) / 2

            self.prop_direct_list.append(1 - n_avg)
            self.prop_indirect_list.append(n_avg)
            general_indirect.append(nu_avg_general)

        general_indirect = pd.concat(general_indirect)

        # Bootstrap these general_indirect values
        bootstrap_overall_means = []
        for i in range(0, 1000):
            bootstrap_overall_means.append(
                general_indirect.sample(frac=0.25, replace=True).mean()
            )

        bootstrap_overall_means = np.array(bootstrap_overall_means)

        final_results = pd.DataFrame(
            {
                "Treatment_Value": self.t_bin_means,
                "Proportion_Direct_Effect": self.prop_direct_list,
                "Proportion_Indirect_Effect": self.prop_indirect_list,
            }
        ).round(4)

        # Clip Proportion_Direct_Effect and Proportion_Indirect_Effect
        final_results["Proportion_Direct_Effect"].clip(lower=0, upper=1.0, inplace=True)
        final_results["Proportion_Indirect_Effect"].clip(
            lower=0, upper=1.0, inplace=True
        )

        # Calculate overall, mean, indirect effect
        total_prop_mean = round(np.array(self.prop_indirect_list).mean(), 4)
        total_prop_lower = self._clip_negatives(
            round(np.percentile(bootstrap_overall_means, q=lower * 100), 4)
        )
        total_prop_upper = self._clip_negatives(
            round(np.percentile(bootstrap_overall_means, q=upper * 100), 4)
        )

        print(
            f"""\n\nMean indirect effect proportion:
            {total_prop_mean} ({total_prop_lower} - {total_prop_upper})
            """
        )
        return final_results

    def _clip_negatives(self, number):
        """Helper function to clip negative numbers to zero"""
        if number < 0:
            return 0
        return number

    def _bootstrap_analysis(self, temp_low_treatment, temp_high_treatment):
        """The top-level function used in the fitting method"""

        bootstrap_collection = []

        for i in range(0, self.bootstrap_replicates):
            # Create single bootstrap replicate
            temp_t, temp_m, temp_y = self._create_bootstrap_replicate()
            # Create the models from this
            temp_mediator_model, temp_outcome_model = self._fit_gams(
                temp_t, temp_m, temp_y
            )
            # Make mediator predictions
            predict_m1, predict_m0 = self._mediator_prediction(
                temp_mediator_model,
                temp_t,
                temp_m,
                temp_y,
                temp_low_treatment,
                temp_high_treatment,
            )
            # Make outcome predictions
            outcome_preds = self._outcome_prediction(
                temp_low_treatment,
                temp_high_treatment,
                predict_m1,
                predict_m0,
                temp_outcome_model,
            )
            # Collect the replicate results here
            bootstrap_collection.append(outcome_preds)

        # Convert this into a dataframe
        bootstrap_results = pd.DataFrame(bootstrap_collection)

        return bootstrap_results

    def _create_bootstrap_replicate(self):
        """Creates a single bootstrap replicate from the data"""
        temp_t = self.T.sample(n=self.bootstrap_draws, replace=True)
        temp_m = self.M.iloc[temp_t.index]
        temp_y = self.y.iloc[temp_t.index]

        return temp_t, temp_m, temp_y

    def _fit_gams(self, temp_t, temp_m, temp_y):
        """Fits the mediator and outcome GAMs"""
        temp_mediator_model = LinearGAM(
            s(0, n_splines=self.n_splines, spline_order=self.spline_order),
            fit_intercept=True,
            max_iter=self.max_iter,
            lam=self.lambda_,
        )
        temp_mediator_model.fit(temp_t, temp_m)

        temp_outcome_model = LinearGAM(
            s(0, n_splines=self.n_splines, spline_order=self.spline_order)
            + s(1, n_splines=self.n_splines, spline_order=self.spline_order),
            fit_intercept=True,
            max_iter=self.max_iter,
            lam=self.lambda_,
        )
        temp_outcome_model.fit(pd.concat([temp_t, temp_m], axis=1), temp_y)

        return temp_mediator_model, temp_outcome_model

    def _mediator_prediction(
        self,
        temp_mediator_model,
        temp_t,
        temp_m,
        temp_y,
        temp_low_treatment,
        temp_high_treatment,
    ):
        """Makes predictions based on the mediator models"""

        m1_mean = temp_mediator_model.predict(temp_high_treatment)[0]
        m0_mean = temp_mediator_model.predict(temp_low_treatment)[0]

        std_dev = (
            (temp_mediator_model.deviance_residuals(temp_t, temp_m) ** 2).sum()
        ) / (self.n - (len(temp_mediator_model.get_params()["terms"]._terms) + 1))

        est_error = np.random.normal(loc=0, scale=std_dev, size=self.n)

        predict_m1 = m1_mean + est_error
        predict_m0 = m0_mean + est_error

        return predict_m1, predict_m0

    def _outcome_prediction(
        self,
        temp_low_treatment,
        temp_high_treatment,
        predict_m1,
        predict_m0,
        temp_outcome_model,
    ):
        """Makes predictions based on the outcome models"""

        outcome_preds = {}

        inputs = [
            ["d1", temp_high_treatment, temp_high_treatment, predict_m1, predict_m0],
            ["d0", temp_low_treatment, temp_low_treatment, predict_m1, predict_m0],
            ["z1", temp_high_treatment, temp_low_treatment, predict_m1, predict_m1],
            ["z0", temp_high_treatment, temp_low_treatment, predict_m0, predict_m0],
        ]

        for input in inputs:

            # Set treatment values
            t_1 = input[1]
            t_0 = input[2]

            # Set mediator values
            m_1 = input[3]
            m_0 = input[4]

            pr_1 = temp_outcome_model.predict(
                np.column_stack((np.repeat(t_1, self.n), m_1))
            )

            pr_0 = temp_outcome_model.predict(
                np.column_stack((np.repeat(t_0, self.n), m_0))
            )

            outcome_preds[input[0]] = (pr_1 - pr_0).mean()

        return outcome_preds

ext: py

max_stars_repo_name: piaodangdang/causal-curve

max_issues_repo_name: piaodangdang/causal-curve

max_forks_repo_name: piaodangdang/causal-curve

reason: import numpy


text:
#Copyright (c) Soumith Chintala 2016, All rights reserved.

import torch
import torchvision
import time
import numpy as np

def nms(boxes, scores, iou_threshold):
    """
    Performs non-maximum suppression (NMS) on the boxes according
    to their intersection-over-union (IoU).

    NMS iteratively removes lower scoring boxes which have an
    IoU greater than iou_threshold with another (higher scoring)
    box.

    Parameters
    ----------
    boxes : Tensor[N, 4])
        boxes to perform NMS on. They
        are expected to be in (x1, y1, x2, y2) format
    scores : Tensor[N]
        scores for each one of the boxes
    iou_threshold : float
        discards all overlapping
        boxes with IoU < iou_threshold

    Returns
    -------
    keep : Tensor
        int64 tensor with the indices
        of the elements that have been kept
        by NMS, sorted in decreasing order of scores
    """
    return torchvision.ops.nms(boxes, scores, iou_threshold)

def batched_nms(boxes, scores, idxs, iou_threshold, use_soft_nms):
    """
    Performs non-maximum suppression in a batched fashion.

    Each index value correspond to a category, and NMS
    will not be applied between elements of different categories.

    Parameters
    ----------
    boxes : Tensor[N, 4]
        boxes where NMS will be performed. They
        are expected to be in (x1, y1, x2, y2) format
    scores : Tensor[N]
        scores for each one of the boxes
    idxs : Tensor[N]
        indices of the categories for each one of the boxes.
    iou_threshold : float
        discards all overlapping boxes
        with IoU < iou_threshold

    Returns
    -------
    keep : Tensor
        int64 tensor with the indices of
        the elements that have been kept by NMS, sorted
        in decreasing order of scores
    """
    if boxes.numel() == 0:
        return torch.empty((0,), dtype=torch.int64, device=boxes.device)
    # strategy: in order to perform NMS independently per class.
    # we add an offset to all the boxes. The offset is dependent
    # only on the class idx, and is large enough so that boxes
    # from different classes do not overlap

    # max_coordinate = boxes.max()
    # offsets = idxs.to(boxes) * (max_coordinate + 1)
    # boxes_for_nms = boxes + offsets[:, None]
    boxes_for_nms = boxes

    if not use_soft_nms :
        #keep = py_cpu_nms(boxes_for_nms, scores, iou_threshold)
        keep = nms(boxes_for_nms, scores, iou_threshold)
        return keep
    else :
        boxes, scores = py_cpu_soft_nms(boxes_for_nms, scores, iou_thr=iou_threshold)
        labels = torch.ones_like(scores).to(torch.device('cuda'))
        keep = torch.tensor([i for i in range(len(labels))]).to(torch.device('cuda'))
        return boxes, scores, labels, keep

def py_cpu_nms(dets, scores, thresh):
    """Pure Python NMS baseline."""

    dets = np.array(dets.cpu())
    scores = np.array(scores.cpu())
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]

    return keep

def py_cpu_soft_nms(dets, scores, method='gaussian', iou_thr=0.3, sigma=0.4, score_thr=0.05):

    if method not in ('linear', 'gaussian', 'greedy'):
        raise ValueError('method must be linear, gaussian or greedy')

    dets = np.array(dets.cpu())
    scores = np.array(scores.cpu())
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    # expand dets with areas, and the second dimension is
    # x1, y1, x2, y2, score, area

    dets = np.concatenate((dets, scores[:, None]), axis=1)
    dets = np.concatenate((dets, areas[:, None]), axis=1)

    retained_box = []

    while dets.size > 0:
        max_idx = np.argmax(dets[:, 4], axis=0)
        dets[[0, max_idx], :] = dets[[max_idx, 0], :]
        retained_box.append(dets[0, :-1])

        xx1 = np.maximum(dets[0, 0], dets[1:, 0])
        yy1 = np.maximum(dets[0, 1], dets[1:, 1])
        xx2 = np.minimum(dets[0, 2], dets[1:, 2])
        yy2 = np.minimum(dets[0, 3], dets[1:, 3])

        w = np.maximum(xx2 - xx1 + 1, 0.0)
        h = np.maximum(yy2 - yy1 + 1, 0.0)
        inter = w * h
        iou = inter / (dets[0, 5] + dets[1:, 5] - inter)

        if method == 'linear':
            weight = np.ones_like(iou)
            weight[iou > iou_thr] -= iou[iou > iou_thr]
        elif method == 'gaussian':
            weight = np.exp(-(iou * iou) / sigma)
        else:  # traditional nms
            weight = np.ones_like(iou)
            weight[iou > iou_thr] = 0

        dets[1:, 4] *= weight
        retained_idx = np.where(dets[1:, 4] >= score_thr)[0]
        dets = dets[retained_idx + 1, :]

    out = np.vstack(retained_box)
    boxes = torch.tensor(out[:,:-1]).to(torch.device('cuda'))
    scores = torch.tensor(out[:,-1]).to(torch.device('cuda'))

    return boxes, scores


text:
"""Estimate direct damages to physical assets exposed to hazards

"""
import sys
import os

import pandas as pd
import geopandas as gpd
from shapely import wkb
import numpy as np

from scripts.direct_damages.analysis_utils import *
from tqdm import tqdm
tqdm.pandas()

def main(damage_data_path, results_data_path):

    asset_data_details = pd.read_csv(os.path.join(damage_data_path,
                        "network_layers_hazard_intersections_details.csv"))

    param_values = pd.read_csv(os.path.join(results_data_path, "parameter_combinations.txt"), sep=" ")
    direct_damages_results = os.path.join(results_data_path,
                                        "direct_damages")

    for asset_info in asset_data_details.itertuples():
        asset_damages_results = os.path.join(
                                        results_data_path,
                                        "direct_damages",
                                        f"{asset_info.asset_gpkg}_{asset_info.asset_layer}"
                                    )
        for param in param_values.itertuples():
            damage_file = os.path.join(
                                asset_damages_results,
                                f"{asset_info.asset_gpkg}_{asset_info.asset_layer}_direct_damages_parameter_set_{param.parameter_set}.csv"
                                )
            if os.path.isfile(damage_file) is True:
                expected_damages = []
                df = pd.read_csv(damage_file).fillna(0)
                haz_rcp_epoch_conf_subs_model = list(set(df.set_index(["hazard","rcp","epoch","confidence","subsidence","model"]).index.values.tolist()))
                for i,(haz,rcp,epoch,conf,subs,model) in enumerate(haz_rcp_epoch_conf_subs_model):
                    damages = df[(df.hazard == haz) & (df.rcp == rcp) & (df.epoch == epoch) & (df.confidence == conf) & (df.subsidence == subs) & (df.model == model)]
                    damages['probability'] = 1.0/damages['rp']
                    index_columns = [c for c in damages.columns.values.tolist() if c not in [
                                                                        'rp',
                                                                        'probability',
                                                                        'direct_damage_cost',
                                                                        'exposure']
                                    ]

                    expected_damage_df = risks_pivot(damages,index_columns,'probability',
                                                'direct_damage_cost',None,'EAD',
                                                flood_protection=None)
                    expected_damages.append(expected_damage_df)

                    del expected_damage_df

                expected_damages = pd.concat(expected_damages,axis=0,ignore_index=True)
                asset_damages_results = os.path.join(direct_damages_results,f"{asset_info.asset_gpkg}_{asset_info.asset_layer}")
                if os.path.exists(asset_damages_results) == False:
                    os.mkdir(asset_damages_results)
                expected_damages.to_csv(os.path.join(asset_damages_results,
                            f"{asset_info.asset_gpkg}_{asset_info.asset_layer}_EAD_parameter_set_{param.parameter_set}.csv"),
                            index=False)
        print (f"* Done with {asset_info.asset_gpkg} {asset_info.asset_layer}")

max_stars_repo_name: itrcrisks/traction

max_issues_repo_name: itrcrisks/traction

max_forks_repo_name: itrcrisks/traction


text:
import os
import random
import argparse

import numpy as np
import pandas as pd

parser = argparse.ArgumentParser(description='Generate smaller NAS-Bench files')
parser.add_argument('--api_loc', default='/disk/scratch_ssd/nasbench201/NASBench_v1_1.pth',
                    type=str, help='path to API')
parser.add_argument('--save_loc', default='bench', type=str, help='folder to save results')
parser.add_argument('--arch_start', default=0, type=int)
parser.add_argument('--arch_end', default=15625, type=int)
parser.add_argument('--seed', default=42, type=int)

args = parser.parse_args()

import torch
import torch.nn as nn
from nas_201_api import NASBench201API as API

# Reproducibility
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)

ARCH_START, ARCH_END = args.arch_start, args.arch_end

api = API(args.api_loc)
seed = 1

df = []

for arch in range(ARCH_START, ARCH_END):
    info = api.query_by_index(arch)
    cellstr = info.arch_str

    cifar10_val = info.get_metrics('cifar10-valid', 'x-valid')['accuracy'] #validation when training on the training split of a train/val/test split

    ### CHOOSE WHICH CIFAR-10 TEST ACC YOU WANT
    cifar10_test = info.get_metrics('cifar10-valid', 'ori-test')['accuracy'] # test accuracy when trained on the training split of train/val/test split
    cifar10_test = info.get_metrics('cifar10', 'ori-test')['accuracy'] # test accuracy when trained on training+val on train/val/test split

    ### CIFAR-100
    cifar100_test = info.get_metrics('cifar100','x-test')['accuracy']
    cifar100_val  = info.get_metrics('cifar100','x-valid')['accuracy']

    ## ImageNet-16-120
    imagenet_test = info.get_metrics('ImageNet16-120', 'x-test')['accuracy']
    imagenet_val  = info.get_metrics('ImageNet16-120', 'x-valid')['accuracy']

    df.append([arch, cellstr, cifar10_test, cifar10_val, cifar100_test, cifar100_val, imagenet_test, imagenet_val])

df = pd.DataFrame(df, columns=['arch','cellstr','cifar10-test','cifar10-val', 'cifar100-test', 'cifar100-val', 'imagenet-test', 'imagenet-val'])

df.to_pickle('mini-bench-arch-cell-accs.pd')


text:
import numpy as np
import strlearn as sl
import hashlib

def post_prunings():
    return (True, False)

def thetas():
    return (0, 0.025, 0.05, 0.075, 0.1)

def weight_calculation_methods():
    return ("same_for_each", "kuncheva", "pta_related_to_whole", "bell_curve")

def aging_methods():
    return ("weights_proportional", "constant", "gaussian")

def base_estimators():
    from sklearn.naive_bayes import GaussianNB
    from sklearn.tree import DecisionTreeClassifier
    from sklearn.neural_network import MLPClassifier
    from sklearn.neighbors import KNeighborsClassifier
    from sklearn.svm import SVC
    from skmultiflow.trees import HoeffdingTreeClassifier

    return {
        "MLP": MLPClassifier(random_state=1410),
    }

def methods():
    from strlearn.ensembles import WAE
    from csm import BALS

    _methods = []
    for base_estimator in base_estimators():
        for post_pruning in post_prunings():
            for theta in thetas():
                for wcm in weight_calculation_methods():
                    for aging_method in aging_methods():
                        clf = BALS(WAE(
                            base_estimators()[base_estimator],
                            theta=theta,
                            post_pruning=post_pruning,
                            weight_calculation_method=wcm,
                            aging_method=aging_method,
                        ), random_state=1410)
                        _methods.append(clf)
    return _methods

def replications():
    return 5

def random_state():
    return 1410

def properties_var():
    """
    Stream properties:
    1. n_chunks
    2. chunk_size
    """
    return (200, 250)

def metrics():
    from sklearn.metrics import accuracy_score

    return {"accuracy": accuracy_score}

def weights_var():
    """
    Class imbalance:
    - tuple for dynamic (n-drifts, sigmoid, amplitude)
    - list for static (positive should be minor)
    """
    return [[0.5, 0.5]]

def y_flip_var():
    """
    Label noise:
    - float for global label noise
    - tuple for separate label noise for every class
    """
    return [0.01]

def drift_var():
    """
    Concept drift:
    1. n_drifts
    2. concept_sigmoid_spacing (None for sudden)
    3. incremental [True] or gradual [False]
    4. recurring [True] or non-recurring [False]
    """
    return [(10, None, False, False), (10, 5, False, False), (10, 5, True, False)]

def concept_var():
    """
    Concept:
    1. n_features
    2. n_informative
    3. n_redundant
    4. n_repeated
    4. n_clusters_per_class
    """
    return [(8, 8, 0, 0, 1)]

def streams(random_state):
    """
    Build stream dictionary
    """
    streams = {}
    n_chunks, chunk_size = properties_var()
    for drift in drift_var():
        n_drifts, concept_sigmoid_spacing, incremental, recurring = drift
        for concept in concept_var():
            (
                n_features,
                n_informative,
                n_redundant,
                n_repeated,
                n_clusters_per_class,
            ) = concept
            for weights in weights_var():
                for y_flip in y_flip_var():
                    stream = sl.streams.StreamGenerator(
                        incremental=incremental,
                        weights=weights,
                        recurring=recurring,
                        random_state=random_state,
                        concept_sigmoid_spacing=concept_sigmoid_spacing,
                        y_flip=y_flip,
                        n_drifts=n_drifts,
                        n_features=n_features,
                        n_informative=n_informative,
                        n_redundant=n_redundant,
                        n_repeated=n_repeated,
                        n_clusters_per_class=n_clusters_per_class,
                        n_chunks=n_chunks,
                        chunk_size=chunk_size,
                        class_sep=1,
                    )
                    streams.update({hash(stream_to_string(stream)): stream})
    return streams

def hash(string):
    return hashlib.md5(string.encode("utf-8")).hexdigest()

def stream_to_string(stream):
    return "dr-%i-%s-%i-%i-co-%i-%i-%i-%i-%i-pr-%i-%i-we-%s-yf-%s-rs-%i" % (
        stream.n_drifts,
        "N"
        if stream.concept_sigmoid_spacing is None
        else "%i" % stream.concept_sigmoid_spacing,
        stream.incremental,
        stream.recurring,
        stream.n_features,
        stream.n_informative,
        stream.n_redundant,
        stream.n_repeated,
        stream.n_clusters_per_class,
        stream.n_chunks,
        stream.chunk_size,
        "D-%i-%i-%.2f" % stream.weights
        if type(stream.weights) is tuple
        else ("S-%s" % "-".join(["%.2f" % s for s in stream.weights])),
        "DYNDYN" if type(stream.y_flip) is tuple else "%.2f" % stream.y_flip,
        stream.random_state,
    )