<reponame>irxat/geoist<filename>geoist/inversion/pfmodel_ts.py
import pathlib
from datetime import datetime
from functools import wraps
import numpy as np
from scipy import linalg as splin
from scipy import sparse as spsparse
from scipy.optimize import minimize
import pandas as pd

from geoist import gridder
from geoist.pfm import prism,tesseroid
from geoist.inversion.mesh import PrismMesh,TesseroidMesh
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import ImageGrid

print_level = -1 # control indentation of prints.
last_print_level = -2

# A helper decorator print time consumption of f.
def timeit(f):
    @wraps(f)
    def wrap(*args,**kwargs):
        global print_level
        global last_print_level
        print_level += 1
        if print_level == last_print_level:
            print('')
        print(' '*4*print_level+'calling {}'.format(f.__name__))
        st = datetime.now()
        res = f(*args,**kwargs)
        ed = datetime.now()
        print(' '*4*print_level+'{} completed in {}'.format(f.__name__,ed-st))
        last_print_level = print_level
        print_level -= 1
        return res
    return wrap

class SmoothOperator:
    def __init__(self):
        self.axis = {'tzyx':{'x':-1,'y':-2,'z':-3,'t':-4},
                     'txyz':{'t':-4,'x':-3,'y':-2,'z':-1}}

    def derivation(self,v,component='dx',array_order='tzyx'):
        for axis_i in component[1:]:
            slices = [slice(None)]*v.ndim
            slices[self.axis[array_order][axis_i]] = slice(-1,None,-1)
            v = np.diff(v[tuple(slices)],axis=self.axis[array_order][axis_i])[tuple(slices)]
        return v

    def rderivation(self,v,component='dx',array_order='tzyx'):
        for axis_i in component[-1:0:-1]:
            slices = [slice(None)]*v.ndim
            slices[self.axis[array_order][axis_i]] = 0
            shape = list(v.shape)
            shape[self.axis[array_order][axis_i]] = 1
            prepend=np.zeros_like(v[tuple(slices)].reshape(tuple(shape)))
            append=np.zeros_like(v[tuple(slices)].reshape(tuple(shape)))
            v = np.diff(v,
                       axis=self.axis[array_order][axis_i],
                       prepend=prepend,
                       append=append)
        return v

class AbicLSQOperatorMS:
    '''An operator doing matrix vector multiplication. The matrix is:
        $\alpha_g G^TG + \sum \alpha_i B_i^TB_i$. Where $\alpha$'s are
        weights, $G$ is kernel matrix, $B_i$'s are smooth matrix.
    Args:
        kernel_matrix (list of ndarray): kernel matrix of each survey
        ns (int): number of surveys
        nzyx (tuple of int): number of cells along z-axis, y-axis and x-axis
        smmoth_components (list of str): which components should be smoothed, acceptable
            string could be 'dx','dxx','dy','dyy','dxy','dt',...
        weights (dict): name and values of weights.
    '''
    def __init__(self,
                 kernel_matrices,
                 ns,
                 nzyx,
                 smooth_components=set(),
                 weights=None):
        self.kernel_matrices = kernel_matrices
        self.ns = ns
        self.nz,self.ny,self.nx = nzyx
        self.shape = (self.ns*self.ny*self.nx,self.ns*self.ny*self.nx)
        self.weights = weights
        self.smooth_components = smooth_components
        self.smop = SmoothOperator()
        if self.weights is None:
            self.weights = {'bound':1,'obs':1,'dx':1,'dy':1,'dt':1, 'refer':1}

    def matvec(self,v):
        tmp = np.zeros_like(v)
        v = v.reshape(self.ns,-1)
        for i,k in enumerate(self.kernel_matrices):
            tmp[i*self.ny*self.nx:(i+1)*self.ny*self.nx] = self.weights['obs'] * (k.T @ (k @ v[i]))
        if 'refer' in self.weights.keys():
            tmp += self.weights['refer']*v.ravel()
        for key in self.smooth_components:
            tmp2 = v.reshape(self.ns,self.nz,self.ny,self.nx)
            tmp2 = self.smop.derivation(tmp2,component=key)
            tmp2 = self.smop.rderivation(tmp2,component=key)
            tmp += self.weights[key]*tmp2.ravel()
        return tmp

class InvModelTS:
    '''Multiple survey gravity inversion use Abic
    Args:
        nyx (tuple of int): number of cells along y-axis and x-axis
        smooth_components (list of str): which components should be smoothed, acceptable
            string could be 'dx','dxx','dy','dyy','dxy','dt',...
        weights (dict): name and values of weights.
        optimize_weights (list of str): specify which weight should be optimized by Abic.
        source_volume (tuple of float): dimension of the underlying gravity source.
        margin (tuple of float): margin size around the source volume.
        cell_type (str): either 'prism' or 'Tesseroid'
        data_dir (str): a folder where observation data resides
    Attributes:
        ns (int): number of surveys
        nx,ny,nz (int): number of cells along x-axis, y-axis and z-axis
        source_volume (tuple of float): dimension of the underlying gravity source
        margin (tuple of float): margin size around the source volume.
        cell_type (str): either 'prism' or 'Tesseroid'
        smooth_components (list of str): which components should be smoothed
        weights (dict): name and values of weights.
        optimize_weights (list of str): specify which weight should be optimized by Abic.
        data_dir (str): a folder where observation data resides
        abic_val (float): abic value
        log_total_det_val (float): log determinate value of the total matrix
        log_obs_det_val (float): log determinate value of the observation matrix
        min_u_val (float): minimum of U
        min_density (float): minimum density
        max_density (float): maximum density
        kernel_matrices (list of ndarray): kernels of each survey.
        rhs (ndarray): right hand side when solve min_u
        solution (ndarray): model densityies solve min_u
        abic_log (dict): save the weights and abic value during calculation
        orig_data (DataFrame): original data loaded from files.
        mesh (Mesh): mesh of the underlying density model.
        R (float): radius of the Earth (m)
    '''
    def __init__(self,
                 nyx=(20,20),
                 smooth_components=None,
                 weights=None,
                 optimize_weights=None,
                 source_volume=None,
                 margin=(0,0,0,0),
                 cell_type='prism',
                 data_dir='./data'):
        self.ny,self.nx = nyx
        self.nz = 1
        self.data_dir = pathlib.Path(data_dir)
        self.source_volume = source_volume
        self.margin= margin
        self.kernel_op = None
        self.smooth_components = smooth_components
        self._weights = weights
        self.optimize_weights = optimize_weights
        self.cell_type = cell_type.lower()
        self.smop = SmoothOperator()
        self.kernel_matrices = []
        self.abic_val = 0
        self.log_total_det_val = 0
        self.log_prior_det_val = 0
        self.log_obs_det_val = 0
        self.min_u_val = 0
        self.min_density = -1.0e4
        self.max_density = 1.0e4
        self._btb_exist = False
        self._gtg_exist = False
        self.abic_log = {'weights':[],'abic_val':[]}
        self._abic_iter = 0
        self.R = 6371000
        self._keymap = None
        self.orig_data = None

    def load_data(self,pattern='*.txt',names=['lon','lat','g'],**kwargs):
        ''' load all data inside data dir, result saved in self.orig_data.
        Args:
            pattern (str): filename pattern.
            names (list of str): name of the observation must be 'g', coordinates
                must be 'lon','lat' or 'x','y'
            kwargs : same as pd.read_csv
        '''
        data_files = sorted(self.data_dir.glob(pattern))
        self.ns = len(data_files)
        if not kwargs:
            kwargs = dict()
            kwargs['delim_whitespace'] = True
        if self.cell_type == 'prism':
            self._keymap = ['x','y']
        elif self.cell_type == 'tesseroid':
            self._keymap = ['lon','lat']
        orig_data = []
        for i,data_file in enumerate(data_files):
            df = pd.read_csv(data_file,names=names,**kwargs)
            i_survey = np.zeros(len(df),dtype=int) + i
            df['i_survey'] = i_survey
            orig_data.append(df)
        self.orig_data = pd.concat(orig_data)
        self.orig_data['z'] = 0.0

    def set_refer(self,refer_density):
        self.refer_density = refer_density

    def deg2xy(self):
        dlon = self.orig_data.groupby('i_survey')['lon'].apply(lambda x: x-x.mean())
        dlat = self.orig_data.groupby('i_survey')['lat'].apply(lambda x: x-x.mean())
        x = dlat*self.R*np.pi/180.
        y = dlon*self.R*np.cos(self.orig_data['lat']*np.pi/180.)*np.pi/180.
        self.orig_data['x'] = x
        self.orig_data['y'] = y

    @property
    def weights(self):
        return self._weights
    @weights.setter
    def weights(self,values):
        self._weights = values
        if not self.kernel_op is None:
            self.kernel_op.weights = self._weights
    @property
    def nx(self):
        return self._nx
    @nx.setter
    def nx(self,value):
        self._nx = value
        self._btb_exist = False
        self._gtg_exist = False

    @property
    def ny(self):
        return self._ny
    @ny.setter
    def ny(self,value):
        self._ny = value
        self._btb_exist = False
        self._gtg_exist = False

    @property
    def smooth_components(self):
        return self._smooth_components
    @smooth_components.setter
    def smooth_components(self,values):
        self._smooth_components = values
        if not self.kernel_op is None:
            self.kernel_op.smooth_components = self._smooth_components

    def _gen_source_volume(self,source_volume=None,margin=None):
        if source_volume is None:
            source_volume = list(self.source_volume)
        fun_list = [min,max,min,max]
        key_ind = [self._keymap[0],self._keymap[0],self._keymap[1],self._keymap[1]]
        for i in range(4):
            if source_volume[i] is None:
                source_volume[i] = fun_list[i](self.orig_data[key_ind[i]])
                source_volume[i] += margin[i]*(-1)**(i+1)
        self.source_volume = source_volume

    def gen_mesh(self,source_volume=None,margin=None):
        shape = (self.nz, self.ny, self.nx)
        if margin is None:
            margin = self.margin
        self._gen_source_volume(source_volume,margin)
        if self.cell_type =='prism':
            self.mesh = PrismMesh(self.source_volume, shape)
        elif self.cell_type =='tesseroid':
            self.mesh = TesseroidMesh(self.source_volume, shape)
        else:
            raise ValueError('cell_type must be \'prism\' or \'tesseroid\'!!')
        density = np.ones(shape)*1.0e3
        self.mesh.addprop('density', density.ravel())

    def gen_obs_grid(self,height=-1):
        """ generate obs grid
         """
        obs_area = (self.source_volume[0]+0.5*self.mesh.dims[0],
                         self.source_volume[1]-0.5*self.mesh.dims[0],
                         self.source_volume[2]+0.5*self.mesh.dims[1],
                         self.source_volume[3]-0.5*self.mesh.dims[1])
        obs_shape = (self.nx, self.ny)
        return gridder.regular(obs_area, obs_shape, z=height)

    def gen_kernel(self):
        self.kernel_matrices = []
        if self.cell_type == 'prism':
            for i in range(self.ns):
                xp = self.orig_data[self.orig_data['i_survey']==i]['x'].values
                yp = self.orig_data[self.orig_data['i_survey']==i]['y'].values
                zp = self.orig_data[self.orig_data['i_survey']==i]['z'].values
                kernel0 = np.zeros((len(xp),self.ny*self.nx))
                for j,cell in enumerate(self.mesh):
                    kernel0[:,j] = prism.gz(xp,yp,zp,[cell])
                self.kernel_matrices.append(np.array(kernel0))
        elif self.cell_type == 'tesseroid':
            for i in range(self.ns):
                xp = self.orig_data[self.orig_data['i_survey']==i]['lon'].values
                yp = self.orig_data[self.orig_data['i_survey']==i]['lat'].values
                zp = self.orig_data[self.orig_data['i_survey']==i]['z'].values
                kernel0 = np.zeros((len(xp),self.ny*self.nx))
                for j,cell in enumerate(self.mesh):
                    kernel0[:,j] = tesseroid.gz(xp,yp,zp,[cell])
                self.kernel_matrices.append(np.array(kernel0))
        else:
            raise ValueError('cell_type must be \'prism\' or \'tesseroid\'!!')
        self.kernel_op = AbicLSQOperatorMS(self.kernel_matrices,
                                           ns=self.ns,
                                           nzyx=(self.nz,self.ny,self.nx),
                                         smooth_components=self._smooth_components,
                                         weights=self._weights)
        self._gtg_exist = False

    def _diagvec(self,vec=None,diag=None):
        if vec.ndim == 1:
            return vec * diag
        else:
            return  vec * diag.reshape(1,-1)

    @timeit
    def forward(self,model_density=None):
        obs_g = []
        if model_density is None:
            x = self.solution.reshape(self.ns,-1)
        else:
            x = model_density.reshape(self.ns,-1)
        for i,k in enumerate(self.kernel_matrices):
            obs_g.append((k @ x[i]).ravel())
        return np.hstack(obs_g)

    def _gen_rhs(self):
        self.rhs = np.zeros(self.ns*self.ny*self.nx)
        for i,k in enumerate(self.kernel_matrices):
            s = i*self.ny*self.nx
            e = (i+1)*self.ny*self.nx
            g = self.orig_data[self.orig_data['i_survey']==i]['g'].values
            self.rhs[s:e] = self._weights['obs'] * (k.T @ g)
        if 'refer' in self._weights.keys():
            self.rhs += (self._weights['refer']*self.refer_density.ravel())

    def _gen_btb(self):
        if self._btb_exist:
            return
        self._btb = dict()
        for key in self.smooth_components:
            if 't' in key:
                tmp = np.eye(self.ns*self.nx*self.ny).reshape(-1,self.ns,self.nz,self.ny,self.nx)
                nrow = self.ns*self.ny*self.nx
            else:
                tmp = np.eye(self.nx*self.ny).reshape(-1,self.ny,self.nx)
                nrow = self.nx*self.ny
            self._btb[key] = self.smop.rderivation(self.smop.derivation(tmp,component=key),
                                                   component=key).reshape(nrow,-1)
        self._btb_exist = True

    def _gen_gtg(self):
        if self._gtg_exist:
            return
        self._gtg = []
        for k in self.kernel_matrices:
            self._gtg.append(k.T @ k)
        self._gtg_exist = True

    @timeit
    def do_linear_solve(self):
        self.do_linear_solve_quiet()

    def do_linear_solve_quiet(self):
        self._gen_rhs()
        self.solution = spsparse.linalg.cg(self.kernel_op,self.rhs,tol=1.0e-5)[0]

    @timeit
    def calc_u(self,solved=False,x=None):
        return self.calc_u_quiet(solved,x)

    @timeit
    def calc_min_u(self,solved=False,x=None):
        return self.calc_u_quiet(solved,x)

    def calc_u_quiet(self,solved=False,x=None):
        if x is None:
            if not solved:
                self.do_linear_solve_quiet()
            x = self.solution
        v = x.reshape(self.ns,-1)
        self.min_u_val = 0.
        for i,k in enumerate(self.kernel_matrices):
            g = self.orig_data[self.orig_data['i_survey']==i]['g'].values
            self.min_u_val += self._weights['obs']*np.linalg.norm(k @ v[i] - g)**2
        for key in self._smooth_components:
            tmp2 = self.smop.derivation(v.reshape(self.ns,self.nz,self.ny,self.nx),
                                        component=key)
            self.min_u_val += self._weights[key]*np.linalg.norm(tmp2.ravel())**2
        if 'refer' in self._weights.keys():
            v = x - self.refer_density.ravel()
            self.min_u_val += self._weights['refer'] * np.linalg.norm(v)**2
        return self.min_u_val

    def calc_res(self):
        self.residuals = dict()
        self.stds = dict()
        x = self.solution.reshape(self.ns,-1)
        res = []
        for i,k in enumerate(self.kernel_matrices):
            g = self.orig_data[self.orig_data['i_survey']==i]['g'].values
            res.append((k@x[i] - g).ravel())
        res = np.hstack(res)
        self.residuals['obs'] = np.linalg.norm(res)**2
        self.stds['obs'] = np.std(res)
        for key in self._smooth_components:
            tmp2 = self.solution.reshape(self.ns,self.nz,self._ny,self._nx)
            tmp2 = self.smop.derivation(tmp2,component=key)
            self.residuals[key] = np.linalg.norm(tmp2.ravel())**2
            self.stds[key] = np.std(tmp2.ravel())
        if 'refer' in self._weights.keys():
            self.residuals['refer'] = []
            self.stds['refer'] = []
            self.residuals['refer'].append(np.linalg.norm(self.solution.ravel()-self.refer_density.ravel())**2)
            self.stds['refer'].append(np.std(self.solution.ravel()-self.refer_density.ravel()))

    def calc_log_prior_total_det_quiet(self,precision=1.0e-6):
        self._gen_gtg()
        self._gen_btb()
        self.log_prior_det_val = 0
        self.log_total_det_val = 0
        prior_eigs = np.zeros(self._nx*self._ny*self.nz)
        total_eigs = np.zeros(self._nx*self._ny*self.nz)
        tmp_mat = np.zeros((self.ns*self.nx*self.ny,self.ns*self.nx*self.ny))
        for key in self.smooth_components:
            if 't' in key:
                tmp_mat += self.weights[key] * self._btb[key]
            else:
                for i in range(self.ns):
                    tmp_mat[i*self.nx*self.ny:(i+1)*self.nx*self.ny,
                            i*self.nx*self.ny:(i+1)*self.nx*self.ny] += self.weights[key] * self._btb[key]
        prior_eigs = np.linalg.eigvalsh(tmp_mat)
        self.log_prior_det_val = sum(np.log(prior_eigs[prior_eigs > precision]))
        for i,k in enumerate(self.kernel_matrices):
            tmp_mat[i*self.nx*self.ny:(i+1)*self.nx*self.ny,
                    i*self.nx*self.ny:(i+1)*self.nx*self.ny] += self.weights['obs'] * self._gtg[i]
        if 'refer' in self._weights.keys():
            tmp_mat += self._weights['refer']*np.eye(self.nx*self.ny*self.ns)
        total_eigs = np.linalg.eigvalsh(tmp_mat)
        self.log_total_det_val = sum(np.log(total_eigs[total_eigs > precision]))
        self.eigs = {'prior':prior_eigs,'total':total_eigs}

    @timeit
    def calc_log_prior_total_det(self,precision=1.0e-6):
        return self.calc_log_prior_total_det_quiet(precision)

    def calc_log_obs_det_quiet(self):
        if 'refer' in self._weights.keys():
            self.log_obs_det_val = (self.nx*self.ny*self.ns*np.log(self.weights['refer'])
                                    +np.log(self._weights['obs'])*len(self.orig_data))
        else:
            self.log_obs_det_val = np.log(self._weights['obs'])*len(self.orig_data)
        return self.log_obs_det_val

    @timeit
    def calc_log_obs_det(self):
        return self.calc_log_obs_det_quiet()

    @timeit
    def calc_abic(self):
        '''-log_prior_det_value+log_total_det-log_obs_det+min_u'''
        self.calc_log_prior_total_det()
        self.calc_u()
        self.calc_log_obs_det()
        self.abic_val = (self.log_total_det_val
                         + self.min_u_val
                         - self.log_prior_det_val
                         - self.log_obs_det_val)
        self.abic_log['weights'].append(self._weights.copy())
        self.abic_log['abic_val'].append(self.abic_val)
        return self.abic_val

    def calc_abic_quiet(self,precision=1.0e-6):
        '''-log_prior_det_value+log_total_det-log_obs_det+min_u'''
        self.calc_log_prior_total_det_quiet(precision)
        self.calc_u_quiet()
        self.calc_log_obs_det_quiet()
        self.abic_val = (self.log_total_det_val
                         + self.min_u_val
                         - self.log_prior_det_val
                         - self.log_obs_det_val)
        self.abic_log['weights'].append(self._weights.copy())
        self.abic_log['abic_val'].append(self.abic_val)
        if self._abic_iter % 10 == 0:
            print('abic value is:{}'.format(self.abic_val))
        self._abic_iter += 1
        return self.abic_val

    def _abic_optimize_exp(self,precision=1.0e-6):
        def abic_target(x):
            for i,key in enumerate(self.optimize_weights):
                self._weights[key] = np.exp(x[i])
            return self.calc_abic_quiet(precision=1.0e-6)
        x0 = np.zeros(len(self.optimize_weights))
        for i,key in enumerate(self.optimize_weights):
            x0[i] = np.log(self._weights[key])
        self.abic_optimize_summary = minimize(abic_target,
                                              x0,
                                              method='Nelder-Mead')

    @timeit
    def abic_optimize(self):
        self._abic_iter = 0
        self._abic_optimize_exp(precision=1.0e-6)

    @timeit
    def para_grad(self,x):
        pass

    def u_bound(self):
        pass

    def print_summary(self):
        print('abic values:{}'.format(self.abic_val))
        print('log total det:{}'.format(self.log_total_det_val))
        print('log prior det:{}'.format(self.log_prior_det_val))
        print('log obs det:{}'.format(self.log_obs_det_val))
        print('min u:{}'.format(self.min_u_val))
        print('std:',end=' ')
        print(self.stds)
        print('1/var:',end=' ')
        print({k:1./v**2 for k,v in self.stds.items()})
        print('norms:',end=' ')
        print(self.residuals)

    def plot_density(self,density=None,surveys=None,fname=None):
        if surveys is None:
            surveys = range(self.ns)
        fig = plt.figure(figsize=(10, 10))
        nrows = int(np.ceil(np.sqrt(len(surveys))))
        grid = ImageGrid(fig, 111,
                        nrows_ncols=(nrows, nrows),
                        axes_pad=0.05,
                        cbar_mode='single',
                        cbar_location='right',
                        cbar_pad=0.1
                        )
        if density is None:
            x = self.solution.reshape(self.ns,self.ny,self.nx)
        else:
            x = density.reshape(self.ns,self.ny,self.nx)
        if self.cell_type == 'prism':
            #rint(x.shape)
            x = np.transpose(x,axes=[0,2,1]) #axis //chenshi
        for ind,i_survey in enumerate(surveys):
            grid[ind].set_axis_off()
            im = grid[ind].imshow(x[i_survey],origin='lower')
        cbar = grid.cbar_axes[0].colorbar(im)
        if fname is None:
            plt.show()
        else:
            plt.savefig(fname,dpi=150)

    def plot_field(self,field=None,surveys=None,fname=None,plot_station=True):
        if surveys is None:
            surveys = range(self.ns)
        if field is None:
            obs_g = self.orig_data['g']
        else:
            obs_g = pd.Series(field,index=self.orig_data.index)
        fig = plt.figure(figsize=(10, 10))
        if self.cell_type == 'prism':
            axis_order = ['y','x']
        elif self.cell_type == 'tesseroid':
            axis_order = ['lon','lat']
        nrows = int(np.ceil(np.sqrt(len(surveys))))
        grid = ImageGrid(fig, 111,
                        nrows_ncols=(nrows, nrows),
                        axes_pad=0.05,
                        cbar_mode='single',
                        cbar_location='right',
                        cbar_pad=0.1
                        )
        for ind,i_survey in enumerate(surveys):
            grid[ind].set_axis_off()
            tmp = self.orig_data[self.orig_data['i_survey']==ind]
            x = tmp[axis_order[0]].values
            y = tmp[axis_order[1]].values
            g = obs_g[self.orig_data['i_survey']==ind].values
            im = grid[ind].tricontourf(x, y, g, 20)
            if plot_station:
                im2 = grid[ind].scatter(x,y)
        cbar = grid.cbar_axes[0].colorbar(im)
        if fname is None:
            plt.show()
        else:
            plt.savefig(fname,dpi=150)


import numpy as np
import scipy.sparse.csgraph

import dijkstra

def test_dijkstra_against_scipy():
    W = np.random.random((100, 100))
    W += W.transpose()
    W[W<1.0] = np.inf
    W = scipy.sparse.csr_matrix(W)

    for seed in [0,1,2]:
        result0 = dijkstra.dijkstra(W, seed)
        result1 = scipy.sparse.csgraph.dijkstra(W, indices = [seed])[0]
        assert all(result0 == result1)


from cmath import rect, phase
from math import radians, degrees
import cv2
import numpy as np

from dyslexia import preprocessing

def mean_angle(deg):
    if len(deg) == 0:
        return 0

    return degrees(phase(sum(rect(1, radians(d)) for d in deg) / len(deg)))

def reject_outliers(values, m=2):
    values = np.array(values)

    return values[abs(values - mean_angle(values)) < m * np.std(values)]

def compute_rotation_angle(image):
    """Find best text angle direction

    It's recommended that you use theses functions before using this one:

    # ``remove_shadow(image)``
    # ``image_to_gray(image, threshold=True)``
    """
    gray = 255 - image

    # Apply threshold to binarize image
    thresh = cv2.threshold(gray, 0, 255,
                           cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
    thresh = 255 - thresh

    # apply close to connect the white areas
    # kernel = np.ones((5, 5), np.uint8)
    # morph = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
    # kernel = np.ones((1, 9), np.uint8)
    # morph = cv2.morphologyEx(morph, cv2.MORPH_CLOSE, kernel)
    # kernel = np.ones((9, 1), np.uint8)
    # morph = cv2.morphologyEx(morph, cv2.MORPH_CLOSE, kernel)

    # Use HoughLines method to find lines
    lines = cv2.HoughLinesP(thresh,
                            rho=1,
                            theta=1 * np.pi / 180,
                            threshold=100,
                            minLineLength=100,
                            maxLineGap=50)

    if lines is None:
        return 0
    if len(lines) == 0:
        return 0

    lines = np.reshape(lines, (lines.shape[0], 4))

    angles = []
    for x1, y1, x2, y2 in lines:  #[:1000]:
        # get angle
        angle = np.arctan2((y2 - y1), (x2 - x1))

        # To degrees
        angle = angle * 180 / np.pi

        # Since we don't know where is the top of
        # the picture then we just want the direction
        if angle < 0:
            angle = -angle

        # Multiply angle by 2 so that mean on degrees work
        # for line direction only
        angles.append(angle * 2)

    # Reject outliers (like paper border lines)
    angles = reject_outliers(angles, m=1)

    # Divides mean by 2 to retrieves direction
    angle = mean_angle(angles) / 2

    return angle

def find_best_rotation_angle(image: np.ndarray, threshold=1):
    """Find best text angle direction

    It's recommended that you use theses functions before using this one:

    # ``remove_shadow(image)``
    # ``image_to_gray(image, threshold=True)``

    Apply this algorithm to find the best angle.

    #. use ``compute_rotation_angle()`` to get first angle detected
    #. if angle is close to 180 or 0 then it's already good
    #. else rotated image by angle
    #. Use this image to use again ``compute_rotation_angle()``
    #. if angle is close to 180 or 0 then use this angle
    #. do the same for but with minus first angle found
    #. If neither angles are close enough then use the closest from 180 or 0

    Parameters
    ----------
    image : np.ndarray
        Source image
    threshold : int, default 1
        Distance in degree from horizontal line (180 or 0 degrees)
        to consider an angle to be good

    Returns
    -------
    number:
        Best rotation angle for this image
    """
    angle = compute_rotation_angle(image)

    if abs(angle - 90) > 90 - threshold:
        return angle

    img_rotated = preprocessing.rotate_img(image, angle)
    new_angle1 = compute_rotation_angle(img_rotated)

    if abs(new_angle1 - 90) > 90 - threshold:
        return angle

    img_rotated = preprocessing.rotate_img(image, -angle)
    new_angle2 = compute_rotation_angle(img_rotated)

    if abs(new_angle2 - 90) > 90 - threshold:
        return -angle

    if abs(new_angle1 - 90) > abs(new_angle2 - 90):
        return angle

    return -angle


<gh_stars>1-10
# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------

"""Utilities to train a surrogate model from teacher."""

import numpy as np
from scipy.sparse import issparse, isspmatrix_csr, vstack as sparse_vstack

def _soft_logit(values, clip_val=5):
    """Compute a soft logit on an iterable by bounding outputs to a min/max value.

    :param values: Iterable of numeric values to logit and clip.
    :type values: iter
    :param clip_val: Clipping threshold for logit output.
    :type clip_val: Union[Int, Float]
    """
    new_values = np.log(values / (1 - values))
    return np.clip(new_values, -clip_val, clip_val)

def _model_distill(teacher_model_predict_fn, uninitialized_surrogate_model, data, original_training_data,
                   explainable_model_args):
    """Teach a surrogate model to mimic a teacher model.

    :param teacher_model_predict_fn: Blackbox model's prediction function.
    :type teacher_model_predict_fn: function
    :param uninitialized_surrogate_model: Uninitialized model used to distill blackbox.
    :type uninitialized_surrogate_model: uninitialized model
    :param data: Representative data (or training data) to train distilled model.
    :type data: numpy.ndarray
    :param original_training_data: Representative data (or training data) to get predictions from teacher model.
    :type original_training_data: numpy.ndarray
    :param explainable_model_args: An optional map of arguments to pass to the explainable model
        for initialization.
    :type explainable_model_args: dict
    """
    # For regression, teacher_y is a real value whereas for classification it is a probability between 0 and 1
    teacher_y = teacher_model_predict_fn(original_training_data)
    multiclass = False
    training_labels = None
    is_classifier = len(teacher_y.shape) == 2
    # If the predict_proba function returned one column but this is a classifier, modify to [1-p, p]
    if is_classifier and teacher_y.shape[1] == 1:
        teacher_y = np.column_stack((1 - teacher_y, teacher_y))
    if is_classifier and teacher_y.shape[1] > 2:
        # If more than two classes, use multiclass surrogate
        multiclass = True
        # For multiclass case, we need to train on the class label
        training_labels = np.argmax(teacher_y, axis=1)
        unique_labels = set(np.unique(training_labels))
        if len(unique_labels) < teacher_y.shape[1]:
            # Get the missing labels
            missing_labels = set(range(teacher_y.shape[1])).difference(unique_labels)
            # Append some rows with the missing labels
            for missing_label in missing_labels:
                # Find max prob for missing label
                max_row_index = np.argmax(teacher_y[:, missing_label])
                # Append the extra label to data and y value
                training_labels = np.append(training_labels, missing_label)
                if issparse(data) and not isspmatrix_csr(data):
                    data = data.tocsr()
                vstack = sparse_vstack if issparse(data) else np.vstack
                data = vstack([data, data[max_row_index:max_row_index + 1, :]])
        surrogate_model = uninitialized_surrogate_model(multiclass=multiclass,
                                                        **explainable_model_args)
    else:
        surrogate_model = uninitialized_surrogate_model(**explainable_model_args)
    if is_classifier and teacher_y.shape[1] == 2:
        # Make sure output has only 1 dimension
        teacher_y = teacher_y[:, 1]
        # Transform to logit space and fit regression
        surrogate_model.fit(data, _soft_logit(teacher_y))
    else:
        # Use hard labels for regression or multiclass case
        if training_labels is None:
            training_labels = teacher_y
        surrogate_model.fit(data, training_labels)
    return surrogate_model


# -*- coding: utf-8 -*-
"""
Created on Wed Apr 28 14:40:35 2021

@author: utente
"""
#Functions that calculate the variance and the entropy, both phenotypic and genotypic

from random import sample, uniform
from scipy.spatial.distance import euclidean
import math
import numpy as np

def phen_variance(pop):
            #Calculating the avg fitness of the population
            avg_fitness = sum([i.fitness for i in pop.individuals])/pop.size
            #calculating the variance over the population
            return sum([(i.fitness-avg_fitness)**2 for i in pop.individuals]) / (pop.size- 1)

def gen_variance(pop):
            #selecting a random individual to be the origin
            origin = sample(pop.individuals, 1)[0]
            #calculating the distances of each point to the distance
            distances = [ sum([ np.linalg.norm(ind[i] -  origin[i]) for i in range(len(ind)) ]) / len(ind) for ind in pop.individuals]
            #calculating the average distance over the population
            avg_distance = sum(distances) / pop.size
            #calculating the variance over the population
            return sum([(distance-avg_distance)**2 for distance in distances]) / (pop.size- 1)

def phen_entropy(pop):
            #Calculating the fitnesses of the population
            fitnesses = [i.fitness for i in pop.individuals]
            #calculating the entropy over the population
            return sum([ fitnesses.count(fitness) / len(fitnesses) * math.log(fitnesses.count(fitness) / len(fitnesses),  2) for fitness in fitnesses])

def gen_entropy(pop):
            #selecting a random individual to be the origin
            origin = sample(pop.individuals, 1)[0]
            #calculating the distances of each point to the distance
            distances = [ sum([ np.linalg.norm(ind[i] - origin[i]) for i in range(len(ind)) ]) / len(ind) for ind in pop.individuals]
            #calculating the variance over the population
            return sum([ distances.count(distance) / len(distances) * math.log(distances.count(distance) / len(distances), 2) for distance in distances])

def fs(pop):
            #building the distance-half matrix
            #starting by initializing the matrix with zeros
            distance_matrix = np.zeros((len(pop), len(pop)))
            #iterating over the individuals and calculating the euclidean distances
            for i in range(pop.size):
                for j in range(i, pop.size):
                    distance_matrix[i,j] = sum([ np.linalg.norm(pop.individuals[i][index] - pop.individuals[j][index]) for index in range(4) ]) / 4
            #normalizing distances in [0,1] and reverting them, so the result is big if the distance was small and viceversa
            max_ = distance_matrix.max()
            min_ = distance_matrix.min()
            distance_matrix = 1 - (distance_matrix - min_)/(max_ - min_)
            #defining a sharing coeffient for each element in the population
            sharing_coeff = []
            for i in range(pop.size):
                sharing_coeff.append( sum(distance_matrix[i] + distance_matrix.T[i]) )
            #updating the individuals fitness
            for i, individual in enumerate(pop.individuals):

                if pop.optim == 'max':

                    individual.fitness = individual.fitness / sharing_coeff[i]

                else:

                    individual.fitness = individual.fitness * int(sharing_coeff[i])

def mo_selection(pop):
    '''
    Multiobjective optimization selection process, can be used only if the fitness function is a tuple.

    '''
    #copying all individuals in a list s
    s = [(pop.individuals.index(i), i.fitness) for i in pop.individuals]
    #Initializing the flag dictionary, the flag and the individuals
    flag = {}
    i = 1
    fit = s[0]
    fit2 = s[1]

    if pop.optim == 'max':

        #keep iterating while the set is not empty
        while len(s) > 0:
            #checking if the individual is dominated
            if fit[1][0] < fit2[1][0] and fit[1][1] < fit2[1][1]:
                #checking if we already compare all individuals
                if s.index(fit) == len(s) - 1:
                    #reinitializing the loop and increasing the flag
                    fit = s[0]
                    fit2 = s[1]
                    i += 1
                #if it is checking for another individual
                fit = s[s.index(fit) + 1 ]

            #otherwise
            else:
                #checking if  we already check for all individuals
                if fit2 == s[-1]:
                    #then the individual is not-dominated, saving it in the flag dict with the right flag
                    flag[str(fit[0])] = i
                    #removing it from the list
                    s.remove(fit)
                    #checking how many individual are left
                    if len(s) == 1:
                        flag[str(s[0][0])] = i + 1
                        break
                    elif len(s) == 0:
                        break
                    else:
                        #reinitializing the individuals
                        fit = s[0]
                        fit2 = s[1]
                #otherwise
                else:
                    #changing the individual to check for
                    fit2 = s[s.index(fit2) + 1]

    else:

        #keep iterating while the set is not empty
        while len(s) > 0:
            #checking if the individual is dominated
            if fit[1][0] > fit2[1][0] and fit[1][1] > fit2[1][1]:
                #checking if we already compare all individuals
                if s.index(fit) == len(s) - 1:
                    #reinitializing the loop and increasing the flag
                    fit = s[0]
                    fit2 = s[1]
                    i += 1
                #if it is checking for another individual
                fit = s[s.index(fit) + 1 ]

            #otherwise
            else:
                #checking if  we already check for all individuals
                if fit2 == s[-1]:
                    #then the individual is not-dominated, saving it in the flag dict with the right flag
                    flag[str(fit[0])] = i
                    #removing it from the list
                    s.remove(fit)
                    #checking how many individual are left
                    if len(s) == 1:
                        flag[str(s[0][0])] = i + 1
                        break
                    elif len(s) == 0:
                        break
                    else:
                        #reinitializing the individuals
                        fit = s[0]
                        fit2 = s[1]
                #otherwise
                else:
                    #changing the individual to check for
                    fit2 = s[s.index(fit2) + 1]

    #assinging to each ind the probability of being picked of 1 - ind.flag / sum(flag)
    tot_flag = sum(flag.values())
    prob = {}
    for ind in list(flag.keys()):
        prob[ind] = 1 - flag[ind]/tot_flag

    #sorting the probabilities dictionary
    sorted_prob = {}
    sorted_keys = sorted(prob, key=prob.get)
    for w in sorted_keys:
        sorted_prob[w] = prob[w]
    #selecting the two parents
    spin = uniform(0, sum(sorted_prob.values()))
    position = 0
    for individual in sorted_prob.keys():
        position += sorted_prob[individual]
        if position > spin:
            index1 = individual
            break

    spin = uniform(0, sum(sorted_prob.values()))
    position = 0
    # Find individual in the position of the spin
    for individual in sorted_prob.keys():
        position += sorted_prob[individual]
        if position > spin:
            index2 = individual
            break

    return pop.individuals[int(index1)], pop.individuals[int(index2)]


<gh_stars>0
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: <NAME>
# License: MIT

"""
Mie theory functions
"""

import numpy as np
import scipy.special as sp

def sj(n, q):
    return sp.spherical_jn(n, q)

def dsj(n, q):
    return sp.spherical_jn(n, q, derivative=True)

def sy(n, q):
    return sp.spherical_yn(n, q)

def dsy(n, q):
    return sp.spherical_yn(n, q, derivative=True)

def sh1(n, q):
    return sj(n, q) + 1j * sy(n, q)

def dsh1(n, q):
    return dsj(n, q) + 1j * dsy(n, q)

def fdot(n, q, f, df):
    return 1 / q * (f(n, q) + q * df(n, q))

def sjdot(n, q):
    return fdot(n, q, sj, dsj)

def sh1dot(n, q):
    return fdot(n, q, sh1, dsh1)

def get_cross_sections_analytical(k0, a, eps_sphere=4, eps_bg=1, Nmax=25):
    k1 = k0 * (eps_bg.conjugate()) ** 0.5
    k2 = k0 * (eps_sphere.conjugate()) ** 0.5
    chi = k2 / k1

    def coeffs(n):
        q1 = sjdot(n, k1 * a) * sj(n, k2 * a) - chi * sj(n, k1 * a) * sjdot(n, k2 * a)
        q2 = sh1dot(n, k1 * a) * sj(n, k2 * a) - chi * sh1(n, k1 * a) * sjdot(n, k2 * a)

        c_over_a = -q1 / q2

        q1 = sj(n, k1 * a) * sjdot(n, k2 * a) - chi * sjdot(n, k1 * a) * sj(n, k2 * a)
        q2 = sh1(n, k1 * a) * sjdot(n, k2 * a) - chi * sh1dot(n, k1 * a) * sj(n, k2 * a)

        d_over_b = -q1 / q2
        return c_over_a, d_over_b

    Cs, Ce = 0, 0

    for n in range(1, Nmax):
        A, B = coeffs(n)
        # print(np.mean((np.abs(A)**2 + np.abs(B)**2)))
        Cs += (2 * n + 1) * (np.abs(A) ** 2 + np.abs(B) ** 2)
        Ce += -(2 * n + 1) * ((A) + (B)).real

    Cs *= 2 * np.pi / k1 ** 2
    Ce *= 2 * np.pi / k1 ** 2
    Ca = Ce - Cs
    return Cs, Ce, Ca


<filename>interpolate/pozo_0.5m_to_5m_interpolation_canopy.py
import numpy as np
from scipy.interpolate import griddata
from matplotlib import pyplot as pl

method = 'linear'
#method = 'cubic'

z2 = np.load('pozo_0.5m_dem_mean_cl2.npy')
z5 = np.load('pozo_0.5m_dem_mean_cl5.npy')
z5[np.isnan(z5)] = z2[np.isnan(z5)]

# input coords
w = 0.5
xb = np.arange(0, 500+w, w)
yb = np.arange(0, 500+w, w)
xc = xb[:-1] + w/2.
yc = yb[:-1] + w/2.
x, y = np.meshgrid(xc, yc)
x2, y2, z2 = x[~np.isnan(z2)], y[~np.isnan(z2)], z2[~np.isnan(z2)]
x5, y5, z5 = x[~np.isnan(z5)], y[~np.isnan(z5)], z5[~np.isnan(z5)]

# output coords
d = 5
Xb = np.arange(0, 500+d, d)
Yb = np.arange(0, 500+d, d)
Xc = Xb[:-1] + d/2.
Yc = Yb[:-1] + d/2.
X, Y = np.meshgrid(Xc, Yc)
Z2 = griddata((x2, y2), z2, (X, Y), method = method)
Z5 = griddata((x5, y5), z5, (X, Y), method = method)

pl.title('Canopy height [m]')
pl.pcolormesh(Xb, Xb, np.ma.masked_invalid(Z5-Z2),
        cmap = pl.cm.seismic,
        vmin = -15, vmax = 15,
        )
pl.axes().set_aspect('equal')
pl.colorbar()
pl.show()


import matplotlib.pyplot as plt
import altair as alt
alt.data_transformers.disable_max_rows()
import pandas as pd
import numpy as np
import os
from scipy.ndimage import gaussian_filter

def smooth_chrom(xs=[], ys=[], smooth_factor=1, source=None, filename=None, save_as=None):
    '''
    Function to smooth chromatogram from MS data.
    Values can either be read from

    :param xs: (array-like) x values from chromatogram
    :param ys: (array-like) y values from chromatogram
    :param source: (string) dictates from where the function should read data.
        None --> implies data is passed as an argument
        'clip' --> implies pandas should read data from clipboard
        'excel' --> implies pandas should read data from excel file
    '''
    if source is None:
        asrt_text = 'If no source is provided, data arrays must be passed as arguments'
        assert xs != [] and ys != [], asrt_text
    elif source=='clip':
        df = pd.read_clipboard()
    elif source=='excel':
        df = pd.read_excel(filename)
    else:
        assert xs is not None, 'Found no valid X values'
        assert ys is not None, 'Found no valid Y values'
    xs = df.iloc[3:, 0].astype(float)
    ys = df.iloc[3:, 1].astype(float)
    ys = gaussian_filter(ys, smooth_factor)
    plt.plot(xs, ys)
    plt.fill_between(xs, ys, alpha=0.3)
    if save_as:
        plt.savefig(save_as)

def get_files(directory='.', exts=['.']):
    '''
    Function that searches the defined directory and reutrns list
    of all files with the specified extension or ending.

    :param directory: (str) raw string of directory to be searched
    :param exts: (list) list of extensions or endings to be returned
    '''
    all_files = []
    for root, _, files in os.walk(directory, topdown=True):
        if exts == ['.']:
            all_files.extend([os.path.join(root, name) for name in files])
        for name in files:
            file_path = os.path.join(root, name)
            for ext in exts:
                if file_path.endswith(ext):
                    all_files.append(file_path)
    return all_files

def find_nearest(array, value):
    '''
    Function that searches an array and returns the value nearest to the
    one passed.

    :param array: (array-like) array to be searched
    :param value: (int or float) experimental value
    '''
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return array[idx]

def plot_ms2_data(xs, ys, peptide, frag_dict, mods=None, show_error=False, tolerance=25):
    '''
    Function to return altair plot of identified fragments for a theoretical
    peptide.

    :param xs: (array) x/time data
    :param ys: (array) intensity data
    :param peptide: (string) peptide sequence
    :param frag_dict: (dict) output returned from data_processing.fragments func
    '''
    # handle -endian order of data
    data = [xs, ys]
    for i, array in enumerate(data):
        if array.dtype.byteorder == '>': # is big-endian
            data[i] = data[i].byteswap().newbyteorder()
    xs, ys = data

    df = pd.DataFrame({
        'x':xs, 'y':ys,
        'fragment':[None]*len(xs),
        'label':['']*len(xs)
    })

    if mods is not None:
        assert isinstance(mods, dict), 'modifications must enter as dictionary'
        for k in mods:
            if not isinstance(mods[k], list):
                frag_dict[k] = list([mods[k]])
            else:
                frag_dict[k] = mods[k]

    dom = {
        'b': '#3b5bad',
        'y': '#c42e23',
        'HexHexNAc': '#3d8f2e',
        'NeuAc': '#3d8f2e',
        'NeuAc-18': '#3d8f2e',
        'HexNac': '#3d8f2e',
        'HexNac-18': '#3d8f2e',
        'HexNac-36': '#3d8f2e',
        'HexNac-fg': '#3d8f2e',
        'Hex': '#3d8f2e',
        'Hex-18': '#3d8f2e',
    }

    err_mass, err_dist, err_kind= [], [], []

    for k, v in frag_dict.items():
        for frag in v:
            nearest = find_nearest(df.x, frag)
            error = mass_error(frag, nearest)
            if abs(error) <= tolerance:
                err_mass.append(nearest)
                err_dist.append(error)
                err_kind.append(k)
                df.loc[(df.x==nearest), 'fragment'] = k
                if k in ['b']:
                    df.loc[(df.x==nearest), 'label'] = k+f'{v.index(frag)+1}'
                elif k in ['y']:
                    df.loc[(df.x==nearest), 'label'] = k+f'{len(v) - v.index(frag)+1}'
                else:
                    df.loc[(df.x==nearest), 'label'] = k

    df.dropna(inplace=True)
    df.loc[:, 'y'] = df.y / np.max(df.y) * 100
    df['label position'] = df.y + 5

    bars = alt.Chart(df).mark_bar(size=2).encode(
        x=alt.X('x', title='m/z', axis=alt.Axis(grid=False)),
        y=alt.Y('y', title='Relative Abundance',
                axis=alt.Axis(grid=False, tickCount=1),
                scale=alt.Scale(domain=(0, 100))),
        color=alt.Color('fragment', scale=alt.Scale(domain=list(dom.keys()),
                        range=list(dom.values())), legend=None)
    ).properties(
        title=peptide,
        width=600
    )

    text = alt.Chart(df).mark_text().encode(
        y=alt.Y('label position'),
        x=alt.X('x'),
        text='label'
    )
    chart = alt.vconcat()
    chart &= alt.layer(bars, text)
    if show_error:
        err_df = pd.DataFrame({
            'mass':err_mass,
            'error':err_dist,
            'kind':err_kind
        })
        dots = alt.Chart(err_df).mark_circle().encode(
            x=alt.X('mass:Q', title='m/z', axis=alt.Axis(grid=False)),
            y=alt.Y('error:Q', title='error (ppm)', axis=alt.Axis(grid=True, tickCount=3),
                    scale=alt.Scale(domain=(-tolerance, tolerance))),
            color=alt.Color('kind:O', scale=alt.Scale(domain=list(dom.keys()), range=list(dom.values()))
            )).properties(height=100)

        line = alt.Chart(pd.DataFrame({'y': [0]})).mark_rule(
            strokeDash=[10, 10]).encode(y='y')

        chart &= (dots + line)
    return chart.configure_view(
        strokeWidth=0
    )

def mass_error(measured, exact):
    '''
    Returns mass error between measured and theoretical values.
    '''
    dif = measured - exact
    quo = dif / exact
    return quo * 10**6

modifications = {
    'oxonium':{
        'Hex-36': 127.06,
        'HexNAc-fg': 138.05,
        'Hex': 163.06,
        'HexNAc-36': 168.09,
        'HexNAc-18': 186.09,
        'HexNAc': 204.09,
        'NeuAc-18': 274.09,
        'NeuAc': 292.08,
        'HexHexNAc': 366.14,
    }
}


<reponame>fetchai/agents-tac<filename>data/experiments/analysis/stats.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ------------------------------------------------------------------------------
#
#   Copyright 2018-2020 Fetch.AI Limited
#
#   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.
#
# ------------------------------------------------------------------------------

"""Quick stats from data."""

import numpy as np
import pandas as pd
from scipy import stats  # type: ignore
import seaborn as sns  # type: ignore
import matplotlib.pyplot as plt  # type: ignore

sns.set(color_codes=True)

def analysis_scores_initial_and_final_vs_equilibrium():
    """Analyse the initial and final vs equilibrium scores for baseline agents only."""
    # load scores
    print("---------------\nAnalysis 1:")
    df1 = pd.read_csv("scores_final_ex2.csv")
    df1 = df1.drop(["Unnamed: 0"], axis=1)
    df2 = pd.read_csv("scores_final_ex3.csv")
    df2 = df2.drop(["Unnamed: 0"], axis=1)
    df_final = pd.concat([df1, df2], ignore_index=True, sort=False)
    df_final.reset_index(inplace=True)

    df3 = pd.read_csv("scores_equilibrium_ex2.csv")
    df3 = df3.drop(["Unnamed: 0"], axis=1)
    df4 = pd.read_csv("scores_equilibrium_ex3.csv")
    df4 = df4.drop(["Unnamed: 0"], axis=1)
    df_eq = pd.concat([df3, df4], ignore_index=True, sort=False)
    df_eq.reset_index(inplace=True)

    df5 = pd.read_csv("scores_initial_ex2.csv")
    df5 = df5.drop(["Unnamed: 0"], axis=1)
    df6 = pd.read_csv("scores_initial_ex3.csv")
    df6 = df6.drop(["Unnamed: 0"], axis=1)
    df_initial = pd.concat([df5, df6], ignore_index=True, sort=False)
    df_initial.reset_index(inplace=True)

    df_in_v_eq = df_initial.subtract(df_eq)
    df_fi_v_eq = df_final.subtract(df_eq)
    print("> Initial vs equilibrium and final vs equilibrium scores (only baseline)...")
    # sample:
    print(
        "Number agents: {}, number games: {}".format(
            df_in_v_eq.shape[1], df_in_v_eq.shape[0]
        )
    )
    # population: all games with the same config
    in_v_eq = np.array([])
    _in = np.array([])
    fi_v_eq = np.array([])
    _fi = np.array([])

    for column_name in df_in_v_eq.columns:
        if column_name == "index":
            continue
        elif column_name[-3:] == "_wm":
            continue
        else:
            in_v_eq = np.append(in_v_eq, df_in_v_eq[column_name].to_numpy())
            _in = np.append(_in, df_initial[column_name].to_numpy())

    for column_name in df_fi_v_eq.columns:
        if column_name == "index":
            continue
        elif column_name[-3:] == "_wm":
            continue
        else:
            fi_v_eq = np.append(fi_v_eq, df_fi_v_eq[column_name].to_numpy())
            _fi = np.append(_fi, df_final[column_name].to_numpy())

    # scipy.stats.ttest_ind(cat1['values'], cat2['values'], equal_var=False)
    # https://en.wikipedia.org/wiki/Welch%27s_t-test
    # https://stats.stackexchange.com/questions/305/when-conducting-a-t-test-why-would-one-prefer-to-assume-or-test-for-equal-vari
    alpha = 0.0001
    H_0 = "H_0: (baseline) agents mean of diff in final scores to equilibrium scores <= (baseline) agents mean of diff in initial scores to equilibrium scores (one-sided t-test)"
    print(H_0)
    result = stats.ttest_ind(in_v_eq, fi_v_eq)
    p_value = result.pvalue / 2.0
    reject = p_value < alpha
    print("One-sided t-test p-value: {}, reject H_0: {}".format(p_value, reject))
    means = [np.mean(fi_v_eq), np.mean(in_v_eq)]
    print("Mean: fi_v_eq={}, in_v_eq={}".format(means[0], means[1]), "score bin")
    _plot(
        in_v_eq,
        fi_v_eq,
        "scores_in_d_eq_vs_fi_d_eq.png",
        "score bin",
        label_one="Difference btw initial and equilibrium scores",
        label_two="Difference btw final and equilibrium scores",
    )

def analysis_scores_final_min_initial():
    """Analyse the final minus initial scores for both agent groups."""
    # load scores
    print("---------------\nAnalysis 2:")
    df1 = pd.read_csv("scores_final_ex2.csv")
    df1 = df1.drop(["Unnamed: 0"], axis=1)
    df2 = pd.read_csv("scores_final_ex3.csv")
    df2 = df2.drop(["Unnamed: 0"], axis=1)
    df_final = pd.concat([df1, df2], ignore_index=True, sort=False)
    df_final.reset_index(inplace=True)

    df3 = pd.read_csv("scores_initial_ex2.csv")
    df3 = df3.drop(["Unnamed: 0"], axis=1)
    df4 = pd.read_csv("scores_initial_ex3.csv")
    df4 = df4.drop(["Unnamed: 0"], axis=1)
    df_initial = pd.concat([df3, df4], ignore_index=True, sort=False)
    df_initial.reset_index(inplace=True)

    df = df_final.subtract(df_initial)

    print("> Initial minus final scores...")
    # sample:
    print("Number agents: {}, number games: {}".format(df.shape[1], df.shape[0]))
    # population: all games with the same config
    baseline = np.array([])
    w_model = np.array([])

    for column_name in df.columns:
        if column_name == "index":
            continue
        elif column_name[-3:] == "_wm":
            w_model = np.append(w_model, df[column_name].to_numpy())
        else:
            baseline = np.append(baseline, df[column_name].to_numpy())

    # scipy.stats.ttest_ind(cat1['values'], cat2['values'], equal_var=False)
    # https://en.wikipedia.org/wiki/Welch%27s_t-test
    # https://stats.stackexchange.com/questions/305/when-conducting-a-t-test-why-would-one-prefer-to-assume-or-test-for-equal-vari
    alpha = 0.0001
    H_0 = "H_0: world modelling agents diff mean scores <= baseline agents diff mean scores (one-sided t-test)"
    print(H_0)
    result = stats.ttest_ind(baseline, w_model)
    p_value = result.pvalue / 2.0
    reject = p_value < alpha
    print("One-sided t-test p-value: {}, reject H_0: {}".format(p_value, reject))
    means = [np.mean(w_model), np.mean(baseline)]
    print("Mean: w_model={}, baseline={}".format(means[0], means[1]))
    # sns.distplot(w_model)
    # plt.show()
    _plot(baseline, w_model, "scores_final_min_initial.png", "score bin")

def analysis_scores():
    """Analyse the final scores."""
    # load scores
    print("---------------\nAnalysis 3:")
    df1 = pd.read_csv("scores_final_ex2.csv")
    df1 = df1.drop(["Unnamed: 0"], axis=1)
    df2 = pd.read_csv("scores_final_ex3.csv")
    df2 = df2.drop(["Unnamed: 0"], axis=1)
    df = pd.concat([df1, df2], ignore_index=True, sort=False)
    df.reset_index(inplace=True)

    print("> Final scores...")
    # sample:
    print("Number agents: {}, number games: {}".format(df.shape[1], df.shape[0]))
    # population: all games with the same config
    baseline = np.array([])
    w_model = np.array([])

    for column_name in df.columns:
        if column_name == "index":
            continue
        elif column_name[-3:] == "_wm":
            w_model = np.append(w_model, df[column_name].to_numpy())
        else:
            baseline = np.append(baseline, df[column_name].to_numpy())

    # scipy.stats.ttest_ind(cat1['values'], cat2['values'], equal_var=False)
    # https://en.wikipedia.org/wiki/Welch%27s_t-test
    # https://stats.stackexchange.com/questions/305/when-conducting-a-t-test-why-would-one-prefer-to-assume-or-test-for-equal-vari
    alpha = 0.0001
    H_0 = "H_0: world modelling agents mean scores <= baseline agents mean scores (one-sided t-test)"
    print(H_0)
    result = stats.ttest_ind(baseline, w_model)
    p_value = result.pvalue / 2.0
    reject = p_value < alpha
    print("One-sided t-test p-value: {}, reject H_0: {}".format(p_value, reject))
    means = [np.mean(w_model), np.mean(baseline)]
    print("Mean: w_model={}, baseline={}".format(means[0], means[1]))
    # sns.distplot(w_model)
    # plt.show()
    _plot(baseline, w_model, "scores.png", "score", extra = 2)

def analysis_txs():
    """Analyse the txs."""
    # load scores
    for typ in ["seller", "buyer"]:
        print("---------------\nAnalysis of {} txs:".format(typ))
        df1 = pd.read_csv("transactions_{}_ex2.csv".format(typ))
        df1 = df1.drop(["Unnamed: 0"], axis=1)
        df2 = pd.read_csv("transactions_{}_ex3.csv".format(typ))
        df2 = df2.drop(["Unnamed: 0"], axis=1)
        df = pd.concat([df1, df2], ignore_index=True, sort=False)
        df.reset_index(inplace=True)
        # sample:
        print("Number agents: {}, number games: {}".format(df.shape[1], df.shape[0]))
        # population: all games with the same config
        baseline = np.array([])
        w_model = np.array([])

        for column_name in df.columns:
            if column_name == "index":
                continue
            elif column_name[-3:] == "_wm":
                w_model = np.append(w_model, df[column_name].to_numpy())
            else:
                baseline = np.append(baseline, df[column_name].to_numpy())

        # scipy.stats.ttest_ind(cat1['values'], cat2['values'], equal_var=False)
        # https://en.wikipedia.org/wiki/Welch%27s_t-test
        # https://stats.stackexchange.com/questions/305/when-conducting-a-t-test-why-would-one-prefer-to-assume-or-test-for-equal-vari
        alpha = 0.0001
        H_0 = "H_0: world modelling agents mean txs >= baseline agents mean txs (one-sided t-test)"
        print(H_0)
        # Alternate: world modelling agents mean txs < baseline agents mean txs
        result = stats.ttest_ind(baseline, w_model)
        p_value = result.pvalue / 2.0
        reject = p_value < alpha
        print(
            "One-sided t-test p-value: {}, reject H_0 (world modelling agents mean txs >= baseline agents mean txs): {}".format(
                p_value, reject
            )
        )
        means = [np.mean(w_model), np.mean(baseline)]
        print("Mean: w_model={}, baseline={}".format(means[0], means[1]))
        # sns.distplot(w_model)
        # plt.show()
        _plot(baseline, w_model, "txs_{}.png".format(typ), "no. of transactions", extra = 2)

def analysis_prices():
    """Analyse the prices."""
    print("---------------\nAnalysis of prices:")
    df1 = pd.read_csv("prices_ex2.csv", dtype=np.float64)
    df1 = df1.drop(["Unnamed: 0"], axis=1)
    df2 = pd.read_csv("prices_ex3.csv", dtype=np.float64)
    df2 = df2.drop(["Unnamed: 0"], axis=1)
    df = pd.concat([df1, df2], ignore_index=True, sort=False)
    df.reset_index(inplace=True)
    # sample:
    print("Number observations: {}".format(df.shape[0]))
    # population: all games with the same config
    baseline = df["baseline"].to_numpy()
    w_model = df["w_model"].to_numpy()
    w_model = w_model[~np.isnan(w_model)]

    # scipy.stats.ttest_ind(cat1['values'], cat2['values'], equal_var=False)
    # https://en.wikipedia.org/wiki/Welch%27s_t-test
    # https://stats.stackexchange.com/questions/305/when-conducting-a-t-test-why-would-one-prefer-to-assume-or-test-for-equal-vari
    alpha = 0.0001
    H_0 = "H_0: world modelling agents mean prices <= baseline agents mean prices (one-sided t-test)"
    print(H_0)
    result = stats.ttest_ind(baseline, w_model)
    p_value = result.pvalue / 2.0
    reject = p_value < alpha
    print(
        "One-sided t-test p-value: {}, reject H_0 (world modelling agents mean prices <= baseline agents mean prices): {}".format(
            p_value, reject
        )
    )
    means = [np.mean(w_model), np.mean(baseline)]
    print("Mean: w_model={}, baseline={}".format(means[0], means[1]))
    # sns.distplot(w_model)
    # plt.show()
    _plot(baseline, w_model, "prices.png", "price", extra=0.2)

def _plot(baseline, w_model, file, x_label, label_one="baseline", label_two="w_model", cut_at_zero = True, extra = 10):
    """Plot helpert function."""
    min_bin = min(baseline.min(), w_model.min()) - extra
    if cut_at_zero:
        min_bin = max(min_bin, 0)
    max_bin = max(baseline.max(), w_model.max()) + extra
    bins = np.linspace(min_bin, max_bin, 100)
    plt.hist(baseline, bins, alpha=0.5, label=label_one)
    plt.hist(w_model, bins, alpha=0.5, label=label_two)
    plt.legend(loc="upper right")
    plt.ylabel('count')
    plt.xlabel(x_label)
    # plt.show()
    plt.savefig(file)
    plt.clf()

if __name__ == "__main__":
    analysis_scores_initial_and_final_vs_equilibrium()
    analysis_scores_final_min_initial()
    analysis_scores()
    analysis_txs()
    analysis_prices()


"""
Parametric spline interpolator

Useful when you need to interpolate a curve of values, like points along a flux surface or a chord
"""
import warnings
import numpy as np
from scipy.interpolate import splprep, splev

class ParametricSpline:
    """
    A wrapper class around slprep and splev from scipy.interpolate
    Uses cubic spline interpolation
    """

    def __init__(
        self,
        sample_points: list,
        t_points: np.array = None,
        smoothing: float = 0.0,
        periodic: bool = False,
        order: int = 3,
    ):
        """
        Calculates the knots and coefficients for a parametric cubic spline interpolator

        sample_points should be a list of 1D arrays, i.e. [x_points, y_points, ...]
        If t_points is given, then the arrays given in sample points should be parametrised by t

        If smoothing if > 0.0, then the input points will be smoothed. (Typically smoothing required << 1, recommend
        to check the fit)

        Order should be less than the number of sample points - 1
        """

        sample_length = None
        for sample_array in sample_points:
            if sample_length is None:
                # Use the first array to set the length
                sample_length = len(sample_array)
            else:
                assert len(sample_array) == sample_length

        assert sample_length > order, (
            f"Not enough sample points ({sample_length}) for an order ({order}) "
            "ParametricSpline"
        )

        if order < 3 and smoothing > 0.0:
            warnings.warn(
                UserWarning(
                    "Should not use smoothing for order < 3 in ParametricSpline"
                )
            )

        if t_points is None:
            tck, self.t_points = splprep(
                sample_points, s=smoothing, per=(1 if periodic else 0), k=order
            )
        else:
            tck, self.t_points = splprep(
                sample_points,
                u=t_points,
                s=smoothing,
                per=(1 if periodic else 0),
                k=order,
            )

        self.t_min, self.t_max = np.min(self.t_points), np.max(self.t_points)

        self.knots, self.coeffs, self.order = tck

    def __call__(self, t_evaluations: np.array):
        """
        Returns the spline evaluations at points given by t_evaluations
        N.b. if no t_points are provided to init, then the given sample points are assumed to be parametrised between 0 and 1
        """
        # Interpolation only!
        assert (
            t_evaluations.min() >= self.t_min and t_evaluations.max() <= self.t_max
        ), f"Requested points in the range\
            {t_evaluations.min()}, {t_evaluations.max()}, which is outside the interval {self.t_min}, {self.t_max}"

        return splev(t_evaluations, (self.knots, self.coeffs, self.order))


import torch
from collections import OrderedDict
from .base_model import BaseModel
from . import networks
from typing import Union
from util import util
import numpy as np
from . import saw_utils
from skimage.transform import resize
from scipy.ndimage.filters import maximum_filter
from scipy.ndimage.measurements import label
import cv2
import json

class BrightestCasModel(BaseModel):
    """ This class implements the pix2pix model, for learning a mapping from input images to output images given paired data.

    The model training requires '--dataset_mode aligned' dataset.
    By default, it uses a '--netG unet256' U-Net generator,
    a '--netD basic' discriminator (PatchGAN),
    and a '--gan_mode' vanilla GAN loss (the cross-entropy objective used in the orignal GAN paper).

    pix2pix paper: https://arxiv.org/pdf/1611.07004.pdf
    """
    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.

        Returns:
            the modified parser.

        For pix2pix, we do not use image buffer
        The training objective is: GAN Loss + lambda_L1 * ||G(A)-B||_1
        By default, we use vanilla GAN loss, UNet with batchnorm, and aligned datasets.
        """
        # changing the default values to match the pix2pix paper (https://phillipi.github.io/pix2pix/)
        parser.set_defaults(norm='batch', netG='unet_256', dataset_mode='aligned')
        if is_train:
            parser.set_defaults(pool_size=0, gan_mode='vanilla')
            parser.add_argument('--lambda_S', type=float, default=1.0, help='weight for Shading loss')
            parser.add_argument('--lambda_AL', type=float, default=1.0, help='weight for Reflection loss')
            parser.add_argument('--lambda_BA', type=float, default=1.0, help='weight for Brightest area loss')
            parser.add_argument('--lambda_BP', type=float, default=1.0, help='weight for Brightest pixel loss')
            parser.add_argument('--lambda_BC', type=float, default=1.0, help='weight for Brightest coordinate loss')
        parser.add_argument('--cat_AL', action='store_true', help='Concat AL')
        parser.add_argument('--cat_In', action='store_true', help='Concat Input')
        parser.add_argument('--cat_In_AL', action='store_true', help='Concat Input and AL')

        return parser

    def __init__(self, opt):
        """Initialize the pix2pix class.

        Parameters:
            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        BaseModel.__init__(self, opt)

        self.loss_names = ['G_AL', 'G_SH', 'G_BA', 'G_BP', 'G_BC']
        self.visual_names = ['input', 'pr_BA', 'pr_BA2', 'gt_BA', 'pr_BP', 'pr_BP2', 'gt_BP', 'pr_AL', 'gt_AL', 'pr_SH', 'gt_SH', 'mask']

        self.model_names = ['G1', 'G2', 'G3']

        self.netG1 = networks.define_G(opt.input_nc, 3, opt.ngf, 'unet_256_multi', opt.norm,
                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
        self.netG2 = networks.define_G(opt.input_nc, 1, opt.ngf, 'resnet_9blocks_multi', opt.norm,
                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)

        g3_input_nc = opt.input_nc
        if opt.cat_AL or opt.cat_In:
            g3_input_nc = g3_input_nc + 3
        if opt.cat_In_AL:
            g3_input_nc = g3_input_nc + 6
        self.netG3 = networks.define_G(g3_input_nc, 1, opt.ngf, 'resnet_9blocks_multi', opt.norm,
                                        not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)
        if self.isTrain:
            # define loss functions
            self.criterionR = torch.nn.MSELoss()
            self.criterionS = torch.nn.MSELoss()
            self.criterionBA = torch.nn.MSELoss()
            self.criterionBP = torch.nn.MSELoss()
            self.criterionBC = torch.nn.MSELoss()
            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            self.optimizer_G1 = torch.optim.Adam(self.netG1.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizer_G2 = torch.optim.Adam(self.netG2.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizer_G3 = torch.optim.Adam(self.netG3.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizers.append(self.optimizer_G1)
            self.optimizers.append(self.optimizer_G2)
            self.optimizers.append(self.optimizer_G3)

    def set_input(self, input):
        self.input = torch.squeeze(input['A'],0).to(self.device) # [bn, 3, 256, 256]
        self.image_paths = input['A_paths']
        self.gt_AL = torch.squeeze(input['gt_AL'],0).to(self.device) # [bn, 3, 256, 256]
        self.gt_SH = torch.squeeze(input['gt_SH'],0).to(self.device) # [bn, 3, 256, 256]
        self.mask = torch.squeeze(input['mask'],0).to(self.device) # [bn, 1, 256, 256]
        self.gt_BA = torch.squeeze(input['gt_BA'],0).to(self.device) # [bn, 1, 256, 256]
        self.gt_BP = torch.squeeze(input['gt_BP'],0).to(self.device) # [bn, 1, 256, 256]
        self.gt_BC = input['gt_BC'].to(self.device)

    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
        pr_SH, pr_AL, color = self.netG1(self.input)  # G(A)
        self.pr_AL = pr_AL
        pr_SH = pr_SH.repeat(1, 3, 1, 1)
        pr_SH = pr_SH * 0.5 + 0.5
        color = torch.unsqueeze(torch.unsqueeze(color, 2), 3)
        self.pr_SH = pr_SH * color
        self.pr_SH = self.pr_SH * 2.0 - 1.0
        self.pr_BC, self.pr_BA, self.pr_BP = self.netG2(self.input)

        if self.opt.cat_AL:
            g3_input = torch.cat((self.pr_SH, self.pr_AL), 1)
        elif self.opt.cat_In:
            g3_input = torch.cat((self.pr_SH, self.input), 1)
        elif self.opt.cat_In_AL:
            g3_input = torch.cat((self.pr_SH, self.pr_AL), 1)
            g3_input = torch.cat((g3_input, self.input), 1)
        else:
            g3_input = self.pr_SH

        self.pr_BC2, self.pr_BA2, self.pr_BP2 = self.netG3(g3_input)

    def backward_G(self):
        """Calculate GAN and L1 loss for the generator"""
        mask = self.mask*0.5 + 0.5
        gt_BC = self.gt_BC[:,:,:2]
        condition = int(self.gt_BC[:, 0, 2].item())
        bc_num = int(self.gt_BC[:, 0, 3].item())

        self.loss_G_AL = self.criterionR(self.pr_AL*mask, self.gt_AL*mask) * self.opt.lambda_AL
        self.loss_G_SH = self.criterionS(self.pr_SH*mask, self.gt_SH*mask) * self.opt.lambda_S
        self.loss_G_BA = self.criterionBA(self.pr_BA*mask, self.gt_BA*mask) * self.opt.lambda_BA
        self.loss_G_BP = self.criterionBP(self.pr_BP*mask, self.gt_BP*mask) * self.opt.lambda_BP
        self.loss_G_BA2 = self.criterionBA(self.pr_BA2*mask, self.gt_BA*mask) * self.opt.lambda_BA
        self.loss_G_BP2 = self.criterionBP(self.pr_BP2*mask, self.gt_BP*mask) * self.opt.lambda_BP

        self.loss_G = self.loss_G_AL + self.loss_G_SH + self.loss_G_BA + self.loss_G_BP + self.loss_G_BA2 + self.loss_G_BP2
        if condition==1:
            self.loss_G_BC = self.criterionBC(self.pr_BC, gt_BC.squeeze(1)) * self.opt.lambda_BC
            self.loss_G_BC2 = self.criterionBC(self.pr_BC2, gt_BC.squeeze(1)) * self.opt.lambda_BC
            self.loss_G += self.loss_G_BC + self.loss_G_BC2
        # else:
        elif condition==2:
            loss_G_BC = util.min_loss_BC(self.pr_BC, gt_BC, bc_num, self.criterionBC)
            loss_G_BC2 = util.min_loss_BC(self.pr_BC2, gt_BC, bc_num, self.criterionBC)

            self.loss_G_BC = loss_G_BC * self.opt.lambda_BC
            self.loss_G_BC2 = loss_G_BC2 * self.opt.lambda_BC
            self.loss_G += self.loss_G_BC + self.loss_G_BC2
        else:
            print('Pass loss_G_BC because condition is {}'.format(condition))

        self.loss_G.backward()

    def optimize_parameters(self):
        # with torch.autograd.set_detect_anomaly(True):
        self.forward()                   # compute fake images: G(A)
        self.optimizer_G1.zero_grad()        # set G's gradients to zero
        self.optimizer_G2.zero_grad()        # set G's gradients to zero
        self.optimizer_G3.zero_grad()        # set G's gradients to zero
        self.backward_G()                   # calculate graidents for G
        self.optimizer_G3.step()             # udpate G's weights
        self.optimizer_G1.step()             # udpate G's weights
        self.optimizer_G2.step()             # udpate G's weights

    def get_current_visuals(self):
        """Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""
        visual_ret = OrderedDict()
        for name in self.visual_names:
            if isinstance(name, str):
                visual_ret[name] = getattr(self, name)
        visual_ret['pr_BP_BC'] = util.get_current_BC(self.pr_BC, self.pr_BP, self.opt)
        visual_ret['pr_BP_BC2'] = util.get_current_BC(self.pr_BC2, self.pr_BP2, self.opt)
        visual_ret['pr_BP_BP'] = util.get_current_BP(self.pr_BP, self.opt)
        visual_ret['pr_BP_BP2'] = util.get_current_BP(self.pr_BP2, self.opt)
        return visual_ret

    def eval_label(self):
        label = ['idx', 'condition', 'gt_BC_num']
        label += self.label_base()['BC'] + self.label_sh()['BC'] + self.label_pr()['BC'] + self.label_pr(True, '2')['BC']
        label += self.label_base()['bcDist'] + self.label_sh()['bcDist'] + self.label_pr()['bcDist'] + self.label_pr(True, '2')['dict_BC']
        label += self.label_base()['baMSE'] + self.label_sh()['baMSE'] + self.label_pr()['baMSE'] + self.label_pr(True, '2')['baMSE']
        label += self.label_base()['bpMSE'] + self.label_sh()['bpMSE'] + self.label_pr()['bpMSE'] + self.label_pr(True, '2')['bpMSE']
        label += self.label_sh()['shEval']

        return label

    def eval_brightest_pixel(self, idx=0):
        with torch.no_grad():
            self.forward()
            self.compute_visuals()

        res_base = self.eval_bp_base(self.mask, self.gt_BA, self.gt_BP, self.gt_BC, self.input)
        res_sh = self.eval_bp_sh(self.mask, self.gt_BA, self.gt_BP, self.gt_BC, self.pr_SH)
        res_sh.update(self.eval_sh(self.mask, self.gt_SH, self.pr_SH))
        res_pr = self.eval_bp_pr(self.mask, self.gt_BA, self.gt_BP, self.gt_BC, self.pr_BA, self.pr_BP, self.pr_BC, '')
        res_pr2 = self.eval_bp_pr(self.mask, self.gt_BA, self.gt_BP, self.gt_BC, self.pr_BA2, self.pr_BP2, self.pr_BC2, '2')

        result = [idx]
        label = self.eval_label(idx)
        for l in label:
            if l in res_base:
                result.append(res_base[l])
            if l in res_pr:
                result.append(res_pr[l])
            if l in res_pr2:
                result.append(res_pr2[l])
        return list(result)


from . import select
from . import utils
from scipy import sparse

import numpy as np
import pandas as pd
import scipy.signal

def library_size(data):
    """Measure the library size of each cell.

    Parameters
    ----------
    data : array-like, shape=[n_samples, n_features]
        Input data

    Returns
    -------
    library_size : list-like, shape=[n_samples]
        Sum over all genes for each cell
    """
    library_size = utils.matrix_sum(data, axis=1)
    if isinstance(library_size, pd.Series):
        library_size.name = "library_size"
    return library_size

def gene_set_expression(
    data,
    genes=None,
    library_size_normalize=False,
    starts_with=None,
    ends_with=None,
    exact_word=None,
    regex=None,
):
    """Measure the expression of a set of genes in each cell.

    Parameters
    ----------
    data : array-like, shape=[n_samples, n_features]
        Input data
    genes : list-like, shape<=[n_features], optional (default: None)
        Integer column indices or string gene names included in gene set
    library_size_normalize : bool, optional (default: False)
        Divide gene set expression by library size
    starts_with : str or None, optional (default: None)
        If not None, select genes that start with this prefix
    ends_with : str or None, optional (default: None)
        If not None, select genes that end with this suffix
    exact_word : str, list-like or None, optional (default: None)
        If not None, select genes that contain this exact word.
    regex : str or None, optional (default: None)
        If not None, select genes that match this regular expression

    Returns
    -------
    gene_set_expression : list-like, shape=[n_samples]
        Sum over genes for each cell
    """
    if library_size_normalize:
        from .normalize import library_size_normalize

        data = library_size_normalize(data)
    gene_data = select.select_cols(
        data,
        idx=genes,
        starts_with=starts_with,
        ends_with=ends_with,
        exact_word=exact_word,
        regex=regex,
    )
    if len(gene_data.shape) > 1:
        gene_set_expression = library_size(gene_data)
    else:
        gene_set_expression = gene_data
    if isinstance(gene_set_expression, pd.Series):
        gene_set_expression.name = "expression"
    return gene_set_expression

def gene_variability(data, kernel_size=0.005, smooth=5, return_means=False):
    """Measure the variability of each gene in a dataset.

    Variability is computed as the deviation from
    the rolling median of the mean-variance curve

    Parameters
    ----------
    data : array-like, shape=[n_samples, n_features]
        Input data
    kernel_size : float or int, optional (default: 0.005)
        Width of rolling median window. If a float between 0 and 1, the width is given
        by kernel_size * data.shape[1]. Otherwise should be an odd integer
    smooth : int, optional (default: 5)
        Amount of smoothing to apply to the median filter
    return_means : boolean, optional (default: False)
        If True, return the gene means

    Returns
    -------
    variability : list-like, shape=[n_samples]
        Variability for each gene
    """
    columns = data.columns if isinstance(data, pd.DataFrame) else None
    data = utils.to_array_or_spmatrix(data)
    if isinstance(data, sparse.dia_matrix):
        data = data.tocsc()
    data_std = utils.matrix_std(data, axis=0) ** 2
    data_mean = utils.toarray(data.mean(axis=0)).flatten()

    if kernel_size < 1:
        kernel_size = 2 * (int(kernel_size * len(data_std)) // 2) + 1

    order = np.argsort(data_mean)
    data_std_med = np.empty_like(data_std)
    data_std_order = data_std[order]
    # handle overhang with reflection
    data_std_order = np.r_[
        data_std_order[kernel_size::-1],
        data_std_order,
        data_std_order[:-kernel_size:-1],
    ]
    medfilt = scipy.signal.medfilt(data_std_order, kernel_size=kernel_size)[
        kernel_size:-kernel_size
    ]

    # apply a little smoothing
    for i in range(smooth):
        medfilt = np.r_[(medfilt[1:] + medfilt[:-1]) / 2, medfilt[-1]]

    data_std_med[order] = medfilt
    result = data_std - data_std_med

    if columns is not None:
        result = pd.Series(result, index=columns, name="variability")
        data_mean = pd.Series(data_mean, index=columns, name="mean")
    if return_means:
        result = result, data_mean
    return result

def gene_capture_count(data, cutoff=0):
    """Measure the number of cells in which each gene has non-negligible counts.

    Parameters
    ----------
    data : array-like, shape=[n_samples, n_features]
        Input data
    cutoff : float, optional (default: 0)
        Number of counts above which expression is deemed non-negligible

    Returns
    -------
    capture-count : list-like, shape=[m_features]
        Capture count for each gene
    """
    gene_sums = np.array(utils.matrix_sum(data > cutoff, axis=0)).reshape(-1)
    if isinstance(data, pd.DataFrame):
        gene_sums = pd.Series(gene_sums, index=data.columns, name="capture_count")
    return gene_sums


<gh_stars>1-10
#!/usr/bin/env python3
"""
I/O and related utils.

Attributes
----------
EXT_1D : list
    List of supported TXT/1D file extensions, in lower case.
EXT_MAT : list
    List of supported matlab file extensions, in lower case.
EXT_NIFTI : list
    List of supported nifti file extensions, in lower case.
EXT_XLS : list
    List of supported XLS-like file extensions, in lower case.
LGR
    Logger
"""

import logging

from os import makedirs
from os.path import exists, join

import numpy as np

from nigsp.utils import change_var_type

EXT_1D = ['.txt', '.csv', '.tsv', '.1d', '.par', '.tsv.gz', '.csv.gz']
EXT_XLS = ['.xls']
EXT_MAT = ['.mat']
EXT_NIFTI = ['.nii', '.nii.gz']
EXT_ALL = EXT_1D + EXT_XLS + EXT_MAT + EXT_NIFTI

EXT_DICT = {'1D': EXT_1D, 'xls': EXT_XLS, 'mat': EXT_MAT, 'nifti': EXT_NIFTI}

LGR = logging.getLogger(__name__)

def check_ext(all_ext, fname, scan=False, remove=False):
    """
    Check which extension a file has, and possibly remove it.

    Parameters
    ----------
    all_ext : list
        All possible extensions to check within
    fname : str or os.PathLikeLike
        The filename to check
    scan : bool, optional
        Scan the given path to see if there is a file with that extension
        If True and no path declared, check if fname has a path, if not scan '.'
        If False, don't scan any folder
    remove : bool, optional
        Remove the extention from fname if it has one

    Returns
    -------
    obj_return : Uses a list to return variable amount of options.
        has_ext : boolean
            True if the extension is found, false otherwise
        fname : str or os.PathLike
            If `remove` is True, return (extensionless) fname
        ext : str
            If both `remove` and `has_ext` are True, returns also found extension
    """
    has_ext = False
    all_ext = change_var_type(all_ext, list, stop=False, silent=True)
    for ext in all_ext:
        if fname.lower().endswith(ext):
            has_ext = True
            LGR.debug(f'{fname} ends with extension {ext}')
            break

    if not has_ext and scan:
        for ext in all_ext:
            if exists(f'{fname}{ext}'):
                fname = f'{fname}{ext}'
                LGR.warning(f'Found {fname}{ext}, using it as input henceforth')
                has_ext = True
                break

    obj_return = [has_ext]

    if remove:
        if has_ext:
            obj_return += [fname[:-len(ext)], ext]  # case insensitive solution
        else:
            obj_return += [fname, '']
    else:
        obj_return += [fname]

    return obj_return[:]

def check_nifti_dim(fname, data, dim=4):
    """
    Check number of dimensions in nifti file.

    Parameters
    ----------
    fname : str
        The name of the file representing `data`
    data : numpy.ndarray
        The data which dimensionality needs to be checked
    dim : int, optional
        The amount of dimensions expected/desired in the data.

    Returns
    -------
    numpy.ndarray
        If `data.ndim` = `dim`, returns data.

    Raises
    ------
    ValueError
        If `data` has different dimensions than `dim`
    """
    data = data.squeeze()

    if data.ndim != dim:
        raise ValueError(f'A {dim}D nifti file is required, but {fname} is '
                         f'{data.ndim}D. Please check the input file.')

    return data

def check_mtx_dim(fname, data, shape=None):
    """
    Check dimensions of a matrix.

    Parameters
    ----------
    fname : str
        The name of the file representing `data`
    data : np.ndarray
        The data which dimensionality needs to be checked
    shape : None, 'square', or 'rectangle'}, str, optional
        Shape of matrix, if empty, skip shape check

    Returns
    -------
    np.ndarray
        If `data.ndim` = 2, returns data.
        If `data.ndim` = 1 and `shape` == 'rectangle',
        Returns data with added empty axis.

    Raises
    ------
    NotImplementedError
        If `data` has more than 3 dimensions.
        If `shape` is not None but `data` is 3D.
    ValueError
        If `data` is empty
        If `shape` == 'square' and `data` dimensions have different lenghts.
    """
    data = data.squeeze()
    LGR.info('Checking data shape.')

    if data.shape[0] == 0:
        raise ValueError(f'{fname} is empty!')
    if data.ndim > 3:
        raise NotImplementedError('Only matrices up to 3D are supported, but '
                                  f'given matrix is {data.ndim}D.')
    if shape is not None:
        if data.ndim > 2:
            raise NotImplementedError('Cannot check shape of 3D matrix.')
        if data.ndim == 1 and shape == 'rectangle':
            data = data[..., np.newaxis]
            LGR.warning(f'Rectangular matrix required, but {fname} is a vector. '
                        'Adding empty dimension.')
        if shape == 'square' and data.shape[0] != data.shape[1]:
            raise ValueError(f'Square matrix required, but {fname} matrix has '
                             f'shape {data.shape}.')

    return data

def load_nifti_get_mask(fname, is_mask=False, ndim=4):
    """
    Load a nifti file and returns its data, its image, and a 3d mask.

    Parameters
    ----------
    fname : str
        The filename to read in
    is_mask : bool, optional
        If the file contains a mask.
        Default: False
    ndim : int
        The number of dimensions expected in the mask

    Returns
    -------
    data : numpy.ndarray
        Data from nifti file.
    mask : numpy.ndarray
        If `is_mask` is False, numpy.ndarray of one dimension less than data,
        in which any element that has at least a value different from zero
        in the last dimension of `data` is True.
        If `is_mask` is True, mask is a boolean representation of data.
    img : nib.img
        Image object from nibabel.

    """
    try:
        import nibabel as nib
    except ImportError:
        raise ImportError('nibabel is required to import nifti files. '
                          'Please see install instructions.')

    LGR.info(f'Loading {fname}.')
    img = nib.load(fname)
    data = img.get_fdata()
    data = check_nifti_dim(fname, data, dim=ndim)

    if is_mask:
        mask = (data != 0)
        LGR.info(f'{fname} loaded as mask.')
    else:
        mask = data.any(axis=-1).squeeze()
        LGR.info(f'Data loaded from {fname}.')

    return data, mask, img

def load_txt(fname, shape=None):
    """
    Read files in textual format.

    Parameters
    ----------
    fname : str or os.PathLike
        Path to the txt file
    shape : None, 'square', or 'rectangle', optional
        Shape of matrix, if empty, skip check

    Returns
    -------
    mtx : numpy.ndarray
        Data matrix

    See also
    --------
    check_mtx_dim
    """
    LGR.info(f'Loading {fname}.')

    _, _, ext = check_ext(EXT_1D, fname, scan=True, remove=True)

    if ext in ['.csv', '.csv.gz']:
        delimiter = ','
    elif ext in ['.tsv', '.tsv.gz']:
        delimiter = '\t'
    elif ext in ['.txt', '.1d', '.par']:
        delimiter = ' '
    else:
        delimiter = None

    mtx = np.genfromtxt(fname, delimiter=delimiter)

    mtx = check_mtx_dim(fname, mtx, shape)

    return mtx

def load_mat(fname, shape=None):
    """
    Read files in matlab format.

    Assumes the existence of a matrix/vector in the mat file, rendered as
    a numpy.ndarray. If there is more than a marix, the one with the largest
    size will be selected.

    Parameters
    ----------
    fname : str or os.PathLike
        Path to the mat file
    shape : None, 'square', or 'rectangle'}, str, optional
        Shape of matrix, if empty, skip check

    Returns
    -------
    mtx : numpy.ndarray
        Data matrix

    Notes
    -----
    Requires module pymatreader to work

    See also
    --------
    check_mtx_dim

    Raises
    ------
    EOFError
        If the mat file does not contain matrix or vectors.
    ImportError
        If pymatreader is not installed or can't be read.
    """
    try:
        from pymatreader import read_mat
    except ImportError:
        raise ImportError('pymatreader is required to import mat files. '
                          'Please see install instructions.')

    LGR.info(f'Loading {fname}.')
    data = read_mat(fname)

    data_keys = []
    for k in data.keys():
        # Check data key only if it's not hidden
        # (skip '__header__', '__version__', '__global__')
        if '__' not in k:
            LGR.info(f'Checking {fname} key {str(k)} content for data (float array/matrices in MATLAB).')
            if type(data[k]) is np.ndarray:
                data_keys.append(k)

    if len(data_keys) < 1:
        raise EOFError(f'{fname} does not seem to contain a numeric matrix.')
    elif len(data_keys) > 1:
        LGR.warning('Found multiple possible arrays to load. '
                    'Selecting the biggest (highest pythonic size).')

    key = data_keys[0]
    for k in data_keys[1:]:
        if data[k].size > data[key].size:
            key = k

    LGR.info(f'Selected data from MATLAB variable {key}')
    mtx = data[key]
    mtx = check_mtx_dim(fname, mtx, shape)

    return mtx

def load_xls(fname, shape=''):
    """
    Read files in xls format.

    Parameters
    ----------
    fname : str or os.PathLike
        Path to the mat file
    shape : None, 'square', or 'rectangle'}, str, optional
        Shape of matrix, if empty, skip check

    Notes
    -----
    Requires module _ to work

    See also
    --------
    check_mtx_dim

    No Longer Returned
    ------------------
    mtx : numpy.ndarray
        Data matrix

    Raises
    ------
    NotImplementedError
        Spreadheet loading is not implemented yet.
    """
    raise NotImplementedError('Spreadsheet loading is not implemented yet')

def export_nifti(data, img, fname):
    """
    Export a nifti file.

    Parameters
    ----------
    data : numpy.ndarray
        Data to be exported
    img : nib.img
        Nibabel image object
    fname : str or os.PathLike
        Name of the output file
    """
    try:
        import nibabel as nib
    except ImportError:
        raise ImportError('nibabel is required to export nifti files. '
                          'Please see install instructions.')

    for e in EXT_NIFTI:
        has_ext, fname, ext = check_ext(e, fname, remove=True)
        if has_ext:
            break

    if ext == '':
        ext = '.nii.gz'

    LGR.info(f'Exporting nifti data into {fname}{ext}.')
    out_img = nib.Nifti1Image(data, img.affine, img.header)
    out_img.to_filename(f'{fname}{ext}')

    return 0

def export_txt(data, fname, ext=None):
    """
    Export data into a text-like or mat file.

    Parameters
    ----------
    data : np.ndarray
        Data to be exported.
    fname : str or os.PathLike
        Name of the output file.
    ext : str or None, optional
        Selected extension for export.

    Returns
    -------
    0
        On a successful run
    """
    if ext.lower() in ['.csv', '.csv.gz', '', None]:
        delimiter = ','
    elif ext.lower() in ['.tsv', '.tsv.gz']:
        delimiter = '\t'
    elif ext.lower() in ['.txt', '.1d', '.par']:
        delimiter = ' '
    else:
        delimiter = None

    if data.ndim < 3:
        np.savetxt(f'{fname}{ext}', data, fmt='%.6f', delimiter=delimiter)
    elif data.ndim == 3:
        makedirs(fname, exist_ok=True)
        for i in range(data.shape[-1]):
            np.savetxt(join(fname, f'{i:03d}{ext}'), data[:, :, i], fmt='%.6f',
                       delimiter=delimiter)

    return 0

def export_mtx(data, fname, ext=None):
    """
    Export data into a text-like or mat file.

    Parameters
    ----------
    data : np.ndarray
        Data to be exported.
    fname : str or os.PathLike
        Name of the output file.
    ext : str or None, optional
        Selected extension for export.

    Notes
    -----
    Requires module scipy to export in .mat format.
    (Will require other modules to export in XLS-like format)

    Raises
    ------
    BrokenPipeError
        If somewhat an extension that is not supported passes all checks.
        (This should never happen)
    ImportError
        If scipy is not installed or cannot be found.
    NotImplementedError
        Spreadheet output is not implemented yet.

    Returns
    -------
    0
        On a successful run
    """
    if ext is None:
        # Check if extension was provided in fname.
        for e in EXT_ALL:
            has_ext, fname, ext = check_ext(e, fname, remove=True)
            if has_ext:
                break
    elif ext.lower() not in EXT_ALL:
        # Check if extension is supported.
        ext = None

    if ext in [None, '']:
        LGR.warning('Extension not specified, or specified extension not '
                    'supported. Forcing export in CSV format.')
        ext = '.csv'
    elif ext.lower() in EXT_NIFTI:
        LGR.warning('Found nifti extension, exporting data in .1D instead')
        ext = '.1D'

    LGR.info(f'Exporting data into {fname}{ext}.')
    if ext.lower() in EXT_MAT:
        try:
            import scipy
        except ImportError:
            raise ImportError('To export .mat files, scipy is required. '
                              'Please install it.')
        scipy.io.savemat(f'{fname}{ext}', {'data': data})
    elif ext.lower() in EXT_XLS:
        raise NotImplementedError('Spreadsheet output is not implemented yet')
    elif ext.lower() in EXT_1D:
        export_txt(data, fname, ext)
    else:
        raise BrokenPipeError(f'This should not have happened: {ext} was the '
                              'selected extension.')

    return 0

"""
Copyright 2022, <NAME>.

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.
"""


# Copyright 2020 Google LLC
#
# 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.

"""Test `BrownianPath`.

The suite tests both running on CPU and CUDA (if available).
"""
import sys

sys.path = sys.path[1:]  # A hack so that we always import the installed library.

import math
import numpy as np
import numpy.random as npr
import torch
from scipy.stats import norm, kstest

import torchsde
import pytest

torch.manual_seed(1147481649)
torch.set_default_dtype(torch.float64)

D = 3
BATCH_SIZE = 131072
REPS = 3
ALPHA = 0.00001

devices = [cpu, gpu] = [torch.device('cpu'), torch.device('cuda')]

def _setup(device):
    t0, t1 = torch.tensor([0., 1.], device=device)
    w0, w1 = torch.randn([2, BATCH_SIZE, D], device=device)
    t = torch.rand([], device=device)
    bm = torchsde.BrownianPath(t0=t0, w0=w0)
    return t, bm

@pytest.mark.parametrize("device", devices)
def test_basic(device):
    if device == gpu and not torch.cuda.is_available():
        pytest.skip(msg="CUDA not available.")

    t, bm = _setup(device)
    sample = bm(t)
    assert sample.size() == (BATCH_SIZE, D)

@pytest.mark.parametrize("device", devices)
def test_determinism(device):
    if device == gpu and not torch.cuda.is_available():
        pytest.skip(msg="CUDA not available.")

    t, bm = _setup(device)
    vals = [bm(t) for _ in range(REPS)]
    for val in vals[1:]:
        assert torch.allclose(val, vals[0])

@pytest.mark.parametrize("device", devices)
def test_normality(device):
    if device == gpu and not torch.cuda.is_available():
        pytest.skip(msg="CUDA not available.")

    t0_, t1_ = 0.0, 1.0
    eps = 1e-2
    for _ in range(REPS):
        w0_ = npr.randn() * math.sqrt(t1_)
        w0 = torch.tensor(w0_, device=device).repeat(BATCH_SIZE)

        bm = torchsde.BrownianPath(t0=t0_, w0=w0)  # noqa

        w1_ = bm(t1_).cpu().numpy()

        t_ = npr.uniform(low=t0_ + eps, high=t1_ - eps)  # Avoid sampling too close to the boundary.
        samples_ = bm(t_).cpu().numpy()

        # True expected mean from Brownian bridge.
        mean_ = ((t1_ - t_) * w0_ + (t_ - t0_) * w1_) / (t1_ - t0_)
        std_ = math.sqrt((t1_ - t_) * (t_ - t0_) / (t1_ - t0_))
        ref_dist = norm(loc=np.zeros_like(mean_), scale=np.ones_like(std_))

        _, pval = kstest((samples_ - mean_) / std_, ref_dist.cdf)
        assert pval >= ALPHA


<reponame>codersUP/MO-Labs
import numpy as np

from scipy.optimize import minimize

def rosen(x):

    """The Rosenbrock function"""

    return sum(100.0*(x[1:]-x[:-1]**2.0)**2.0 + (1-x[:-1])**2.0)

def exercise_2(n:int):
    # x0 = np.array([0 for i in range(n)])
    x0 = np.array([1.3, 0.7, 0.8, 1.9, 1.2]) # Este es el de la documentación
    # x0 = np.array([20, 12, 57, 9, 76])

    res = minimize(rosen, x0, method='nelder-mead',

               options={'xatol': 1e-8, 'disp': True})
    print(res.x)

exercise_2(5)