#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
The Jonker-Volgenant algorithm for finding the maximum assignment.

<NAME>, University of Texas at Austin
last modified 10/23/2018

This is a direct adaptation of the Pascal code from
"A Shortest Augmenting Path Algorithm for Dense and Sparse Linear Assignment Problems"
<NAME> and <NAME>, Computing 1987

the __main__ code at the bottom tests this implementation, comparing it to
Scipy's linear_sum_assignment function. You'll need to have scipy in your
distribution to run this file on its own, but not to import it in other files.
"""
import numpy as np
import numba as nb

inf= 1e9 # inf is a suitably large number

@nb.jit(nb.f8(nb.f8[:,:], nb.i8[:], nb.i8[:], nb.f8[:],
              nb.i8[:], nb.i8, nb.i8[:], nb.i8, nb.f8[:], nb.i8[:]), nopython=True)
def SSP(c, x, y, v, rows2use, nrows2use, cols2use, ncols2use, d, pred):
    """ solves full 2D assignment problem
    c: matrix
    x: column indices that match to row, or -1 if row is missing
    y: match indices for column
    v: column reductions
    rows2use, nrows2use: rows in rows2use[:nrows2use] are considered part of the problem
    cols2use, ncols2use: " "
    d, pred: workspace for SSP, remember costs and path backwards for each column
    returns cost of assignment
    """
    C = 0.

    # basic column reduction - basically running some rows in a convenient order
    nrows = nrows2use
    for ri in xrange(nrows2use-1,-1,-1):
        i = rows2use[ri]
        j = np.argmin(c[i,:])
        if c[i,j] < 0 and y[j] == -1:
            x[i] = j
            y[j] = i
            C += c[i,j]
            nrows -= 1
            rows2use[ri] = rows2use[nrows]
            rows2use[nrows] = i

    for i1 in rows2use[:nrows]:
        d[:] = c[i1,:] - v
        pred[:] = i1
        minmissi = i1
        minmissval = 0.
        ncolsunused = ncols2use
        emergcounter = 0
        while True:
            emergcounter += 1
            assert emergcounter < 2000
            minval = minmissval
            minj = -1
            mincolidx = 0
            for colidx, j in enumerate(cols2use[:ncolsunused]):
                dj = d[j]
                if dj < minval:
                    minj = j
                    minval = dj
                    mincolidx = colidx
            j = minj
            if j == -1:
                break # hit unmatched row
            i = y[j]
            if i == -1:
                break # hit unmatched column
            # this column should no longer be considered
            v[j] += minval
            ncolsunused -= 1
            cols2use[mincolidx] = cols2use[ncolsunused]
            cols2use[ncolsunused] = j
            # update distances to other columns
            u1 = c[i,j] - v[j]
            if -u1 < minmissval:
                # this row is the closest to missing
                minmissi = i
                minmissval = -u1
            for j in cols2use[:ncolsunused]:
                dj = c[i,j] - v[j] - u1
                if dj < d[j]:
                    d[j] = dj
                    pred[j] = i

        # augment
        # travel back through shortest path to find matches
        if j==-1:
            i = minmissi
            j = x[i]
            x[i] = -1
        emergcounter = 0
        while i != i1:
            emergcounter += 1
            assert emergcounter < 2000
            i = pred[j]
            y[j] = i
            k = j
            j = x[i]
            x[i] = k
        # updating of column prices
        for j in cols2use[ncolsunused:ncols2use]:
            v[j] -= minval
        C += minval
    return C

@nb.jit(nb.f8(nb.f8[:,:], nb.i8[:], nb.i8[:], nb.f8[:], nb.i8[:],
              nb.i8, nb.i8[:], nb.i8, nb.f8[:], nb.i8[:],
              nb.i8, nb.i8, nb.b1[:], nb.b1, nb.f8), nopython=True)
def spStep(c, x, y, v, rows2use, nrows2use, cols2use, ncols2use, d, pred,
           i1, j1, eliminate_els, eliminate_miss, cost_bound):
    """ solves Murty subproblem given solution to originating problem
    same inputs as SSP and also:
    i1, j1 = row and column that are now unassigned
    eliminate_els = boolean array, whether matching a column with i1 is prohibited
    eliminate_miss = whether i1 is prohibited to miss
    cost_bound = function will stop early and return inf if the solution is known
                    to be above this bound
    returns cost of shortest path, a.k.a. this solution's cost minus original solution's
    """

    if j1>=0:
        u0 = c[i1,j1]-v[j1] # not necessary to get solution, but gives accurate cost
    else:
        u0 = 0.
    pred[:] = i1
    ncols = ncols2use
    for j in cols2use[:ncols]:
        d[j] = inf if eliminate_els[j] else c[i1,j] - v[j] - u0
    minmissj = -1
    minmissi = i1
    minmissval = inf if eliminate_miss else -u0
    miss_unused = True
    missing_from_row = False
    missing_cost = 0. # this is a dual cost on auxiliary columns
    emergcounter = 0
    while True:
        emergcounter += 1
        assert emergcounter < 2000
        minval = minmissval
        minj = -2
        minjcol = -1
        for jcol, j in enumerate(cols2use[:ncols]):
            dj = d[j]
            if dj < minval:
                minj = j
                minval = dj
                minjcol = jcol
        if minval > cost_bound: return inf # that's all it takes for early stopping!
        j = minj
        if j==j1: break
        if j == -2:
            if not miss_unused: # if you got here again, costs must be really high
                return inf
            # entry to missing zone: row was matched but is now missing
            missing=True
            missing_from_row = True
        else:
            i = y[j]
            # this column should no lonber be considered
            ncols -= 1
            cols2use[minjcol] = cols2use[ncols]
            cols2use[ncols] = j
            if i==-1:
                # entry to missing zone: col was missing but is now matched
                if miss_unused:
                    minmissj = j
                    missing=True
                    missing_from_row = False
                else:
                    # already covered the missing zone, this is a dead end
                    continue
            else:
                missing=False
        if missing:
            if j1 == -1:
                j=-1
                break
            miss_unused = False
            missing_cost = minval
            minmissval = inf
            u1 = -minval
            # exit from missing zone: row that was missing is matched
            for i in rows2use[:nrows2use]:
                if x[i]==-1:
                    for j in cols2use[:ncols]:
                        dj = c[i,j]-v[j]-u1
                        if dj < d[j]:
                            d[j] = dj
                            pred[j] = i
            # exit from missing zone: col that was matched is missing
            for j in cols2use[:ncols]:
                if y[j] >= 0:
                    dj = -v[j]-u1
                    if dj < d[j]:
                        d[j] = dj
                        pred[j] = -1
        else:
            u1 = c[i,j]-v[j]-minval
            if miss_unused and -u1<minmissval:
                minmissi = i
                minmissval = -u1
            for j in cols2use[:ncols]:
                dj = c[i,j]-v[j]-u1
                if dj < d[j]:
                    d[j] = dj
                    pred[j] = i

    # augment
    # updating of column prices
    v[cols2use[ncols:ncols2use]] += d[cols2use[ncols:ncols2use]] - minval
    if not miss_unused:
        v[cols2use[:ncols2use]] += minval - missing_cost
    # travel back through shortest path to find matches
    i = i1+1 # any number that isn't i1
    emergcounter = 0
    while i != i1:
        emergcounter += 1
        assert emergcounter < 2000
        if j == -1:
            # exit from missing zone: row was missing but is now matched
            i = -1
        else:
            i = pred[j]
            y[j] = i
        if i == -1:
            # exit from missing zone: column j was matched but is now missing
            if missing_from_row:
                # entry to missing zone: row was matched but is now missing
                i = minmissi
                j = x[i]
                x[i] = -1
            else:
                # entry to missing zone: col was missing but is now matched
                j = minmissj
        else:
            k = j
            j = x[i]
            x[i] = k
    v[y==-1] = 0.
    return minval

if __name__ == '__main__':
    """
    create a random matrix
    try assignment, check for equality
    """
    from scipy.optimize import linear_sum_assignment

    m=10
    n=20
#    P = np.random.exponential(size=(n,m))
#    mX = np.random.exponential(size=(n,))
#    mY = np.random.exponential(size=(m,))
    P = np.random.rand(m,n)
    mX = np.random.rand(m)
    mY = np.random.rand(n)

    # make full square version, use standard code
    c1 = np.zeros((m+n,m+n))
    c1[:m,:n] = P
    c1[:m,n:] = 1e4
    c1[range(m),range(n,m+n)] = mX
    c1[m:,:n] = 1e4
    c1[range(m,m+n), range(n)] = mY
    sol = linear_sum_assignment(c1)
    x1 = np.array(sol[1][:m])
    x1[x1>=n] = -1
    y1 = np.arange(n)
    for k,j in enumerate(sol[1]):
        j = sol[1][k]
        if j < n:
            if k < m:
                y1[j] = k
            else:
                y1[j] = -1
    print x1
    print y1

    y = np.zeros(n, dtype=int) - 1
    x = np.zeros(m, dtype=int) - 1
    v = np.zeros(n)
    c2 = P - mX[:,None] - mY[None,:]
    rows2use = np.arange(m)
    cols2use = np.arange(n)
    d = np.zeros(n)
    pred = np.zeros(n, dtype=int)
    SSP(c2, x, y, v, rows2use, cols2use, d, pred)
    print x
    print y

    v += mY
    u = mX.copy()
    xmatch = x>=0
    xmis = xmatch==False
    ymis = y==-1
    u[xmatch] = P[xmatch,x[xmatch]] - v[x[xmatch]]
    u2 = np.append(u, np.zeros(n))
    v2 = np.append(v, np.zeros(m))
    x2 = np.append(x, y+n)
    x2[np.where(x==-1)[0]] = np.where(x==-1)[0]+n
    x2[np.where(y==-1)[0]+m] = np.where(y==-1)[0]
    slack = c1 - u2[:,None] - v2
    assert np.min(slack) > -1e-8
    assert all(slack[range(m+n), x2] < 1e-8)
    assert np.min(v[ymis]) >= -1e-8 if any(ymis) else True


"""
BSD 3-Clause License

Copyright (c) 2020, Cyber Security Research Centre Limited
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this
   list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its
   contributors may be used to endorse or promote products derived from
   this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
from tick.hawkes import SimuHawkesExpKernels, SimuHawkesMulti, HawkesExpKern, SimuHawkes, HawkesEM, HawkesKernelTimeFunc, HawkesBasisKernels
from tick.plot import plot_hawkes_kernels, plot_point_process
import numpy as np
import matplotlib.pyplot as plt
from HelperFunctions import ConcatClusters, ComputeClusterLengths
from tick_goodness_of_fit import plot_resid, goodness_of_fit_par
from scipy import integrate
import math

#contains various functions used to train or simulate models related to hawkes processes
#some maybe broken, but the ones used in the notebooks work

#takes a list of inner cluster timestamps
#attemps to train an exp model to produce more clusters
def TrainInnerClusterExp(clusters, num_decays=2000, decay_low=-10, decay_high=10):
    data = ConcatClusters(clusters, 0)
    best_score = -1e100
    #decays for multiple dimention process
    #update this to have different decays for each process
    #num_decays = 2000
    #print(f"Total decay combinations = {num_decays*num_decays*num_decays*num_decays}")
    decay_candidates = np.logspace(decay_low, decay_high, num_decays, dtype=np.dtype('d'))
    print(f"Training on {len(clusters)} clusters")
    print(f"Decay Range: {decay_candidates[0]} -> {decay_candidates[-1]}")
    best_decay = decay_candidates[0]
    score_list = np.zeros(num_decays)

    #x*e^(-xt)
    l = 0
    floaterrors = 0
    baseline_errors = 0
    for i, decay in enumerate(decay_candidates):
        decay = decay * np.ones((2,2))
        try:
            #might need a hyperbolic kernel?
            #it seems to get too excited and decays too slowly
            #only small decay values seem to make sense
            learner = HawkesExpKern(decay, penalty='l2', C=1000, max_iter=1000, solver='agd', tol=1e-3)#, max_iter=1000, tol=1e-5) #gofit='likelihood'
            ###Error functions
            #l1 - has 0 step errors
            #l2 - runs, but the results do not look good, heavily favours higher decay values that produce nonsense graphs
            #elasticnet (elastic_net_ratio, def 0.95) - values closer to 0 work better (since it uses l2) otherwise it produces step errors. Still similar to l2.
            #nuclear - basically the same
            #none - how can you have no penalty function?
            ###solvers
            #agd - all penalties favour super high decays, basicaly wants random event generation
            #gd - basically the same
            #bfgs - does weird things, but is quick
            #svrg

            learner.fit(data, start=learner.coeffs)

            """cluster_num = 0
            for cluster in clusters:
                if (cluster_num % 100 == 0):
                    #print out training progress
                    s = f"It: {i}, Decay: {decay[0]}, Cluster: {cluster_num}"
                    print(f"\r{' '*l}\r", end='')
                    print(f"It: {i}, Decay: {decay[0]}, Cluster: {cluster_num}", end='', flush=True)
                    l = len(s)
                learner.fit(cluster, start=learner.coeffs)
                cluster_num += 1"""
            hawkes_score = learner.score()
            #print(hawkes_score)
            #print(f"Coeffs: {learner.coeffs}")

            #ensure there is a non-0 baseline
            numb = 0
            for b in learner.baseline:
                if (b > 0):
                    numb += 1
            if (numb == 0):
                baseline_errors += 1
                continue

            #record the score for plotting
            score_list[i] = hawkes_score

            #record the best
            if (hawkes_score > best_score):
                best_score = hawkes_score
                best_learner = learner
                best_decay = decay

            step = 0.01
            #residuals = goodness_of_fit_par(learner,data,step,integrate.simps)
            #plot_resid(residuals,2,1)

        except ZeroDivisionError:
            #print("float error");
            floaterrors += 1
            continue;

    #create a score plot
    plt.plot(decay_candidates, score_list)
    plt.xscale('log')
    plt.yscale('log')
    plt.title('decay Scores')
    plt.grid(True)
    plt.show()

    print(f"\nTraining Done")
    print(f"Float Errors: {floaterrors} ({100/num_decays*floaterrors}%)")
    print(f"Baseline Errors: {baseline_errors} ({100/num_decays*baseline_errors}%)")
    print(f"==========\nSuccessful Results: {num_decays - floaterrors - baseline_errors} ({100/num_decays*(num_decays - floaterrors - baseline_errors)}%)\n==========\n")

    print(f"\nBest Score: {best_score}")
    print(f"Best Decay: {best_decay}")
    plot_hawkes_kernels(best_learner)

    print(f"Adjacency: {best_learner.adjacency}")
    print(f"Baseline: {best_learner.baseline}")
    print(f"Coeffs: {best_learner.coeffs}")

    #return best_learner.adjacency, best_learner.baseline, best_decay
    return best_learner, best_decay

def TrainSeriesExp(series, num_decays=2000, decay_low=-10, decay_high=10):
    best_score = -1e100
    #decays for multiple dimention process
    #update this to have different decays for each process
    #num_decays = 2000
    #print(f"Total decay combinations = {num_decays*num_decays*num_decays*num_decays}")
    decay_candidates = np.logspace(decay_low, decay_high, num_decays, dtype=np.dtype('d'))
    print(f"Decay Range: {decay_candidates[0]} -> {decay_candidates[-1]}")
    best_decay = decay_candidates[0]
    score_list = np.zeros(num_decays)

    #x*e^(-xt)
    l = 0
    floaterrors = 0
    baseline_errors = 0
    for i, decay in enumerate(decay_candidates):
        #decay = decay * np.ones((2,2))
        try:
            #might need a hyperbolic kernel?
            #it seems to get too excited and decays too slowly
            #only small decay values seem to make sense
            learner = HawkesExpKern(decay, penalty='l2', C=1e-3, max_iter=1000, solver='agd', tol=1e-5)#, max_iter=1000, tol=1e-5) #gofit='likelihood'
            ###Error functions
            #l1 - has 0 step errors
            #l2 - runs, but the results do not look good, heavily favours higher decay values that produce nonsence graphs
            #elasticnet (elastic_net_ratio, def 0.95) - values closer to 0 work better (since it uses l2) otherwise it produces step errors. Still similar to l2.
            #nuclear - basically the same
            #none - how can you have no penalty function?
            ###solvers
            #agd - all penalties favour super high decays, basicaly wants random event generation
            #gd - basically the same
            #bfgs - does weird things, but is quick
            #svrg

            learner.fit([series])

            """cluster_num = 0
            for cluster in clusters:
                if (cluster_num % 100 == 0):
                    #print out training progress
                    s = f"It: {i}, Decay: {decay[0]}, Cluster: {cluster_num}"
                    print(f"\r{' '*l}\r", end='')
                    print(f"It: {i}, Decay: {decay[0]}, Cluster: {cluster_num}", end='', flush=True)
                    l = len(s)
                learner.fit(cluster, start=learner.coeffs)
                cluster_num += 1"""
            hawkes_score = learner.score()
            #print(hawkes_score)
            #print(f"Coeffs: {learner.coeffs}")

            #ensure there is a non-0 baseline
            numb = 0
            for b in learner.baseline:
                if (b > 0):
                    numb += 1
            if (numb == 0):
                baseline_errors += 1
                continue

            #record the score for plotting
            score_list[i] = hawkes_score

            #record the best
            if (hawkes_score > best_score):
                best_score = hawkes_score
                best_learner = learner
                best_decay = decay
        except ZeroDivisionError:
            #print("float error");
            floaterrors += 1
            continue;

    #create a score plot
    plt.plot(decay_candidates, score_list)
    plt.xscale('log')
    plt.yscale('log')
    plt.title('decay Scores')
    plt.grid(True)
    plt.show()

    print(f"\nTraining Done")
    print(f"Float Errors: {floaterrors} ({100/num_decays*floaterrors}%)")
    print(f"Baseline Errors: {baseline_errors} ({100/num_decays*baseline_errors}%)")
    print(f"==========\nSuccessful Results: {num_decays - floaterrors - baseline_errors} ({100/num_decays*(num_decays - floaterrors - baseline_errors)}%)\n==========\n")

    print(f"\nBest Score: {best_score}")
    print(f"Best Decay: {best_decay}")
    plot_hawkes_kernels(best_learner)

    print(f"Adjacency: {best_learner.adjacency}")
    print(f"Baseline: {best_learner.baseline}")
    print(f"Coeffs: {best_learner.coeffs}")

    #return best_learner.adjacency, best_learner.baseline, best_decay
    return best_learner, best_decay

#careful, setting the time to be too high will eat up all available memory
def SimulateExp(baseline, adjacency, decays, time):

    hawkes = SimuHawkesExpKernels(adjacency=adjacency, decays=decays, baseline=baseline, verbose=False)

    hawkes.end_time = time
    dt = 0.001 #millisecond granularity
    #hawkes.track_intensity(dt)
    print(f"Starting sim")
    hawkes.simulate()
    timestamps = hawkes.timestamps
    l = 0
    for series in timestamps:
        l += len(series)
    print(f"Simulated {l} points")
    return hawkes.timestamps

def TrainInnerClusterEM(clusters, k_time = 1, k_size = 100):
    #merge all the clusters, the learner seems to only be able to fit a single long datastream
    num_clusters = len(clusters)
    data = ConcatClusters(clusters, 0)

    #kernel size is the granularity
    #kernel support is something... (is it the size of each step?)
    em_learner = HawkesEM(kernel_support=k_time, kernel_size=k_size, n_threads=8, verbose=True, tol=1e-5, max_iter=1000)
    em_learner.fit(data)

    """#train the em learner on each cluster
    cluster_num = 0
    for cluster in clusters:
        if (cluster_num % 10 == 0):
            #print out training progress
            s = f"Cluster: {cluster_num}/{num_clusters}"
            print(f"\r{' '*l}\r", end='')
            print(f"Cluster: {cluster_num}/{num_clusters}", end='', flush=True)
            l = len(s)
        print(em_learner.baseline)
        print(em_learner.kernel)
        print("==========")
        if (cluster_num == 0):
            em_learner.fit(cluster)
        else:
            em_learner.fit(cluster, baseline_start=em_learner.baseline, kernel_start=em_learner.kernel)
        cluster_num += 1"""
    #maybe add variation in kernel sie later?
    #use em_learner.score() to evaluate goodness
    print(f"\nEM Score: {em_learner.score()}")
    fig = plot_hawkes_kernels(em_learner) #TODO, remove this?

    t = np.linspace(0, k_time, endpoint=False, num=k_size)

    m = []
    for i in range(2):
        for j in range(2):
            m.append(max(em_learner.kernel[i][j]))
    #normalise to make a proper hawkes process
    spectral_radius = max(m)
    if (spectral_radius < 1):
        spectral_radius = 1

    #create a 2x2 array of time func kernels
    k = [[],[]]
    for i in range(2):
        for j in range(2):
            k[i].append(HawkesKernelTimeFunc(t_values=t, y_values=em_learner.kernel[i][j]/np.linalg.norm(em_learner.kernel[i][j])))

    #return k, em_learner.baseline #the kernel, baseline
    return em_learner

def SimulateEM(kernel, baseline, time=600):
    sim_em = SimuHawkes(kernels=kernel, baseline=baseline, verbose=False, end_time=time)

    dt = 0.001 #millisecond granularity
    #sim_em.track_intensity(dt)
    sim_em.simulate()

    timestamps = sim_em.timestamps
    l = 0
    for series in timestamps:
        l += len(series)
    print(f"Simulated {l} points")
    return sim_em.timestamps

def TrainInnerClusterBasis(clusters, k_time=1, k_size=100, num_kernels=2):
    num_clusters = len(clusters)
    #data = ConcatClusters(clusters, 0)
    l = 0

    basis_learner = HawkesBasisKernels(kernel_support=k_time, kernel_size=k_size, n_basis=num_kernels, C=1e-3, n_threads=8, verbose=False, ode_tol=1e-5, max_iter=1000)
    #train the basis learner on each cluster
    cluster_num = 0
    for cluster in clusters:
        if (cluster_num % 10 == 0):
            #print out training progress
            s = f"Cluster: {cluster_num}/{num_clusters}"
            print(f"\r{' '*l}\r", end='')
            print(f"Cluster: {cluster_num}/{num_clusters}", end='', flush=True)
            l = len(s)
        if (cluster_num == 0):
            basis_learner.fit(cluster)
        else:
            basis_learner.fit(cluster, baseline_start=basis_learner.baseline, amplitudes_start=basis_learner.amplitudes, basis_kernels_start=basis_learner.basis_kernels)
        cluster_num += 1

    #kernel size is the granularity
    #kernel support is something... (is it the size of each step?)
    #basis_learner = HawkesBasisKernels(kernel_support=k_time, kernel_size=k_size, n_basis=num_kernels, C=1e-3, n_threads=8, verbose=True, ode_tol=1e-5, max_iter=1000)
    #basis_learner.fit(data)
    #maybe add variation in kernel sie later?
    #use em_learner.score() to evaluate goodness
    #print(f"\nEM Score: {basis_learner.score()}")

    #TODO, remove this?
    fig = plot_hawkes_kernels(basis_learner)

    print(basis_learner.basis_kernels)
    print(basis_learner.amplitudes)
    print(basis_learner.baseline)
    return None

def SimulateBasis(kernel, baseline, time=600):
    sim_em = SimuHawkes(kernels=kernel, baseline=baseline, verbose=False, end_time=time)

    dt = 0.001 #millisecond granularity
    sim_em.track_intensity(dt)
    sim_em.simulate()

    timestamps = sim_em.timestamps
    l = 0
    for series in timestamps:
        l += len(series)
    print(f"Simulated {l} points")
    return sim_em.timestamps

#takes clusters of timestamps
#trains using clusters instead of concatenating them
def TrainInnerTimestampsExp(clusters, num_decays=2000, decay_low=-10, decay_high=10, max_iterations=100, tolerance=1e-5, e=10):
    cat_clusters = ConcatClusters(clusters, 0)
    best_score = -1e100
    print(f"Training on {len(clusters)} clusters")
    unique_decays = int(num_decays**(1.0/4))
    num_decays = unique_decays**4
    decay_candidates = np.logspace(decay_low, decay_high, unique_decays, dtype=np.dtype('d'))
    print(f"Decay Range: {decay_candidates[0]} -> {decay_candidates[-1]}")
    print(f"{unique_decays} unique decays. {num_decays} total")
    best_decay = None
    score_list = np.zeros(num_decays)

    #x*e^(-xt)
    l = 0
    floaterrors = 0
    baseline_errors = 0
    for i in range(num_decays):
        decay = np.ones((2,2))
        decay[0][0] = decay_candidates[int(i/(unique_decays**3))%unique_decays]
        decay[0][1] = decay_candidates[int(i/(unique_decays**2))%unique_decays]
        decay[1][0] = decay_candidates[int(i/(unique_decays**1))%unique_decays]
        decay[1][1] = decay_candidates[int(i)%unique_decays]
        prev_score = float('-inf')
        #print(decay)
        try:
            learner = HawkesExpKern(decay, penalty='l2', C=e, max_iter=1, solver='agd', tol=1e-5)

            #do the learning loop
            #need a stopping point
            for i in range(max_iterations):
                for cluster in clusters:
                    learner.fit(cluster, start=learner.coeffs)
                it_score = learner.score()
                print(f"It: {i}, Score: {it_score}")
                if (it_score <= prev_score + tolerance):
                    #barely changed
                    break
                prev_score = it_score
            hawkes_score = learner.score(events=clusters)

            #ensure there is a non-0 baseline
            numb = 0
            for b in learner.baseline:
                if (b > 0):
                    numb += 1
            if (numb == 0):
                baseline_errors += 1
                continue

            #record the score for plotting
            score_list[i] = hawkes_score

            #record the best
            if (hawkes_score > best_score):
                best_score = hawkes_score
                best_learner = learner
                best_decay = decay

        except ZeroDivisionError:
            #print("float error");
            floaterrors += 1
            continue;

    #create a score plot
    plt.plot(score_list)
    plt.xscale('log')
    plt.yscale('log')
    plt.title('decay Scores')
    plt.grid(True)
    plt.show()

    print(f"\nTraining Done")
    print(f"Float Errors: {floaterrors} ({100/num_decays*floaterrors}%)")
    print(f"Baseline Errors: {baseline_errors} ({100/num_decays*baseline_errors}%)")
    print(f"==========\nSuccessful Results: {num_decays - floaterrors - baseline_errors} ({100/num_decays*(num_decays - floaterrors - baseline_errors)}%)\n==========\n")

    print(f"\nBest Score: {best_score}")
    print(f"Best Decay: {best_decay}")
    plot_hawkes_kernels(best_learner)

    print(f"Adjacency: {best_learner.adjacency}")
    print(f"Baseline: {best_learner.baseline}")
    print(f"Coeffs: {best_learner.coeffs}")

    #activate this for residuals (Warning, it is REALLLLLLLLLLY SLOOOOOOOOOOOOW)
    step = 0.1
    residuals = goodness_of_fit_par(best_learner,cat_clusters,step,integrate.simps)
    plot_resid(residuals,2,1)

    return best_learner.adjacency, best_learner.baseline, best_decay

#"you just put the clusters into it and it does it (tm)"
#takes clusters of timestamps
#trains using clusters instead of concatenating them
def TrainInnerTimestampsExp2(clusters, num_decays=2000, decay_low=-10, decay_high=10, e=10):
    best_score = -1e100
    print(f"Training on {len(clusters)} clusters")
    unique_decays = int(num_decays**(1.0/4))
    num_decays = unique_decays**4
    decay_candidates = np.logspace(decay_low, decay_high, unique_decays, dtype=np.dtype('d'))
    print(f"Decay Range: {decay_candidates[0]} -> {decay_candidates[-1]}")
    print(f"{unique_decays} unique decays. {num_decays} total")
    best_decay = None
    score_list = np.zeros(num_decays)

    #x*e^(-xt)
    l = 0
    floaterrors = 0
    baseline_errors = 0
    for i in range(num_decays):
        s = f"Decay {i} ({format(100/num_decays*i, '.2f')}% done)"
        l = len(s)
        #print(f"{' '*l}\r", end="", flush=True)
        print(f"{' '*l}\r{s}\r", end='', flush=True)
        decay = np.ones((2,2))
        decay[0][0] = decay_candidates[int(i/(unique_decays**3))%unique_decays]
        decay[0][1] = decay_candidates[int(i/(unique_decays**2))%unique_decays]
        decay[1][0] = decay_candidates[int(i/(unique_decays**1))%unique_decays]
        decay[1][1] = decay_candidates[int(i)%unique_decays]
        prev_score = float('-inf')
        #print(decay)
        try:
            learner = HawkesExpKern(decay, penalty='l2', C=e, max_iter=1000, solver='agd', tol=1e-5)
            learner.fit(clusters)
            hawkes_score = learner.score()

            #ensure there is a non-0 baseline
            numb = 0
            for b in learner.baseline:
                if (b > 0):
                    numb += 1
            if (numb == 0):
                baseline_errors += 1
                continue

            #record the score for plotting
            score_list[i] = hawkes_score

            #record the best
            if (hawkes_score > best_score):
                best_score = hawkes_score
                best_learner = learner
                best_decay = decay

        except ZeroDivisionError:
            #print("float error");
            floaterrors += 1
            continue;

    #create a score plot
    plt.plot(score_list)
    plt.xscale('log')
    plt.yscale('log')
    plt.title('decay Scores')
    plt.grid(True)
    plt.show()

    print(f"\nTraining Done")
    print(f"Float Errors: {floaterrors} ({100/num_decays*floaterrors}%)")
    print(f"Baseline Errors: {baseline_errors} ({100/num_decays*baseline_errors}%)")
    print(f"==========\nSuccessful Results: {num_decays - floaterrors - baseline_errors} ({100/num_decays*(num_decays - floaterrors - baseline_errors)}%)\n==========\n")

    print(f"\nBest Score: {best_score}")
    print(f"Best Decay: {best_decay}")
    plot_hawkes_kernels(best_learner)

    print(f"Adjacency: {best_learner.adjacency}")
    print(f"Baseline: {best_learner.baseline}")
    print(f"Coeffs: {best_learner.coeffs}")

    #activate this for residuals (Warning, it is REALLLLLLLLLLY SLOOOOOOOOOOOOW)
    cat_clusters = ConcatClusters(clusters, 0)
    step = 0.1
    residuals = goodness_of_fit_par(best_learner,cat_clusters,step,integrate.simps)
    plot_resid(residuals,2,1)

    return best_learner.adjacency, best_learner.baseline, best_decay

def SimExp(baseline, adjacency, decays, num_clusters, data):
    hawkes = SimuHawkesExpKernels(adjacency=adjacency, decays=decays, baseline=baseline, verbose=False)

    #dt = 0.001 #millisecond granularity
    #hawkes.track_intensity(dt) # turning this on will eat up memory

    #need to compute and draw from the cluster length distrbution from the original data
    cluster_lengths = ComputeClusterLengths(data)

    multi = SimuHawkesMulti(hawkes, n_simulations=num_clusters)

    multi.end_time = np.random.choice(cluster_lengths, size=num_clusters, replace=True)
    multi.simulate()
    sim_inner_timestamps = multi.timestamps

    l = 0
    for realisation in sim_inner_timestamps:
        for series in realisation:
            l += len(series)
    print(f"Simulated {l} points")
    return sim_inner_timestamps

#k_time is the timescale on which a model should be formed
#k_size is how many points it should consider on this timescale
def TrainEM(times, k_time = 1, k_size = 100):

    #kernel size is the granularity
    #kernel support is the length of time to support
    em_learner = HawkesEM(kernel_support=k_time, kernel_size=k_size, n_threads=4, verbose=True, tol=1e-5, max_iter=250)
    em_learner.fit(times)

    #fig = plot_hawkes_kernels(em_learner) #TODO, remove this?

    return em_learner

def SimEM(smodel, time=600):
    sim_em = SimuHawkes(kernels=smodel.time_kernel, baseline=smodel.baseline, verbose=False, end_time=time)
    sim_em.simulate()
    return sim_em.timestamps

class SavedModel:

    def __init__(self, model, k_time, k_size, is_em=True, decay=0):
        self.baseline = model.baseline
        self.k_time = k_time
        self.k_size = k_size
        self.time_kernel = None
        self.intensity = 0
        self.decay = decay
        self.is_em = is_em
        if (is_em):
            self.kernel = model.kernel
            self.MakeTimeKernel(self.kernel)
            self.n_realizations = model.n_realizations
            self.n_nodes = model.n_nodes
        else:
            self.intensity = model.adjacency
            self.coeffs = model.coeffs

    def MakeTimeKernel(self, kernel):
        #check the spectral radius
        m = []
        for process in kernel:
            for intensity in process:
                m.append(max(intensity))
        self.spectral_radius = max(m)
        if (self.spectral_radius > 1):
            #oh no
            self.time_kernel = None
            return

        #create a time kernel for simulation
        t = np.linspace(0, self.k_time, endpoint=False, num=self.k_size)
        k = []
        for process in kernel:
            for intensity in process:
                k.append(HawkesKernelTimeFunc(t_values=t, y_values=intensity))

        k = np.reshape(k, kernel.shape[:2])
        self.time_kernel = k

    def Save(self, fname):
        np.savez(fname, savedmodel=self)


# -- coding: utf-8 --

# Copyright 2018 <NAME> <<EMAIL>>

import scipy.optimize as opt
"""
Provides some useful mathematical functions which are not present in numpy.
"""

import numpy as np

def prod(x):
    p = 1
    for k in x:
        p *= k
    return p

def perm(x):
    """
    permutation with repetition of n factors
    x = [n1,n2,n3,...nk]
    n = n1+n2+...+nk
    return n!/(n1!*n2*...*nk!)
    """
    n = sum(x)
    num = list(range(2,n+1))
    denom = list(range(2,x[0]+1))
    for k in x[1:]:
        denom += list(range(2,k+1))
    for x in num:
        if x in denom:
            num.remove(x)
            denom.remove(x)
    return prod(num)/prod(denom)

def closest_arg(array, value):
    return np.argmin(np.abs(array-value))

def strictly_positify(x):
    """
    Make the result strictly positive by setting the minimum value to
    the lower allowed float value
    """
    return np.fmax(x, np.finfo(x.dtype).eps)

def positify(x):
    """
    Set to zero all negative values
    """
    return np.fmax(0, x)

def clip01(x):
    """
    clip data x between 0 and 1
    """
    return np.fmax(np.fmin(x, 1), 0)

def fact(x):
    """
    Return the factors of an integer as a list.
    Warning: This function is not a factorial!
    """
    if x < 0 or type(x) is not int:
        raise ValueError("input must be a positive integer")
    if x < 2:
        return x
    f = []
    i = 2
    while True:
        while x % i == 0:
            f.append(i)
            x /= i
        i += 1
        if x == 1:
            return f

def moving_average(x, N):
    assert len(x) > N
    c = np.cumsum(x)
    return (c[N:]-c[:-N])/N

def butter_lowpass(cutOff, fs, order=5):
    import scipy.signal
    nyq = 0.5 * fs
    normalCutoff = cutOff / nyq
    b, a = scipy.signal.butter(order, normalCutoff, btype='low', analog=True)
    return b, a

def butter_lowpass_filter(data, cutOff, fs, order=4):
    import scipy.signal
    b, a = butter_lowpass(cutOff, fs, order=order)
    y = scipy.signal.lfilter(b, a, data)
    return y

def Gauss(x, x0, s, amp=None, **kargs):
    if 'A' in kargs:
        from warnings import warn
        warn("Parameter A is deprecated. Please use amp in order to set the amplitude!")
        amp = kargs.pop('A')
    elif 'Amp' in kargs:
        from warnings import warn
        warn("Parameter Amp is deprecated. Please use amp in order to set the amplitude!")
        amp = kargs.pop('Amp')
    R = np.exp(-(x-x0)**2/(2*s**2))
    if amp is None:
        R /= (s*np.sqrt(2*np.pi))
    else:
        R *= amp
    R[s==0] = (x[s==0]==x0)*1.0
    return R

def Lorentz(x, x0, gamma, amp=None, **kargs):
    if 'A' in kargs:
        from warnings import warn
        warn("Parameter A is deprecated. Please use amp in order to set the amplitude!")
        Amp = kargs['A']
    elif 'Amp' in kargs:
        from warnings import warn
        warn("Parameter Amp is deprecated. Please use amp in order to set the amplitude!")
        amp = kargs.pop('Amp')
    R = 1/((x-x0)**2+(.5*gamma)**2)
    if amp is None:
        return .5*gamma*R/np.pi
    return amp*R*(.5*gamma)**2

def CDF(x,mu,sig, amp=1, lg=0, **kargs):
    if 'Amp' in kargs:
        from warnings import warn
        warn("Parameter Amp is deprecated. Please use amp in order to set the amplitude!")
        amp = kargs.pop('Amp')
    from scipy.special import erf
    g = sig*np.sqrt(2*np.log(2))
    return amp*lg*(.5+np.arctan2(x-mu,g)/np.pi)+(1-lg)*amp*.5*(1+erf((x-mu)/(sig*np.sqrt(2))))

def LG(x, x0, sig=None, amp=None, lg=.5, asym=1, FWHM=None, **kargs):
    if 'Amp' in kargs:
        from warnings import warn
        warn("Parameter Amp is deprecated. Please use amp in order to set the amplitude!")
        amp = kargs.pop('Amp')

    assert sig is not None or FWHM is not None

    if FWHM is None:
        FWHM = 2*np.sqrt(2*np.log(2))*sig
    if sig is None:
        sig = FWHM/(2*np.sqrt(2*np.log(2)))
    Y = (1-lg)*Gauss(x,x0,sig,amp=amp)+lg*Lorentz(x,x0,FWHM,amp=amp)
    if asym!=1:
        Yr = (1-lg)*Gauss(x, x0, sig*asym, amp=amp)+lg*Lorentz(x, x0, FWHM*asym, amp=amp)
        Y[x>x0] = Yr[x>x0]
    return Y

def logistic(x, lower=0, upper=1, growth=1, x0=0, nu=1, C=1):
    return lower+(upper-lower)/(C+np.exp(-growth*(x-x0)))**(1/nu)

def fitCDF1line(A):
    line = np.zeros(A.shape[1])
    for x in range(A.shape[1]):
        popt, pcov = opt.curve_fit(CDF,
            np.arange(A.shape[0]),
            A[:,x],
            (A.shape[0]/2, 1, np.max(A[:,x])),
            bounds=(0,(A.shape[1],np.inf,np.inf))
            )
        line[x] = popt[0]
    return line

def FT(x, ufunc=np.real, real=False):
    """
    shortcut for 1D/2D Fourier Transform (real and centered)
    """
    assert isinstance(x, np.ndarray)
    if len(x.shape) == 1:
        if real:
            F = np.fft.rfft(x)
        else:
            F = np.fft.fft(x)
    elif len(x.shape) == 2:
        if real:
            F = np.fft.rfft2(x)
        else:
            F = np.fft.fft2(x)
    else:
        raise TypeError("The array should be 1D or 2D")
    return ufunc(np.fft.fftshift(F))

def binning(data, N=2, axis=0, ufunc=np.sum):
    w = int(np.floor(data.shape[axis]/N))
    r = np.copy(data)
    size = list(data.shape)
    size[axis] = w
    size = size[:axis+1]+[N]+size[axis+1:]
    r.resize(size)
    return ufunc(r, axis=axis+1)

def stat_info(data):
    import matplotlib.pyplot as plt
    D = np.ravel(data)
    U = np.unique(D)
    if len(U)>1:
        sep = np.min(U[1:]-U[:-1])
        N = min(100, int(np.ceil((np.max(D)-np.min(D))/sep)))
    else:
        N = 1

    mean = np.mean(D)
    std = np.std(D)

    fig, ax = plt.subplots(2,1,figsize=(21,4))
    ax[0].boxplot(D, 0, 'ro', 0);
    ax[1].hist(D, N, density=True);
    ax[1].axvline(mean, color='r', label='mean')
    ax[1].axvline(mean+std, color='r', linestyle='--', label='1$\\sigma$')
    ax[1].axvline(mean-std, color='r', linestyle='--', label='1$\\sigma$')
    if mean-2*std >= U[0]:
        ax[1].axvline(mean-2*std, color='r', linestyle=':', label='2$\\sigma$')
    if mean+2*std <= U[-1]:
        ax[1].axvline(mean+2*std, color='r', linestyle=':', label='2$\\sigma$')
    ax[1].legend();
    print("Stats")
    print("\tAverage:", mean)
    print("\tStandard-deviation:", std)
    print("\tMinimum:", np.min(D))
    print("\tQ1:", np.percentile(D, 25))
    print("\tMedian:", np.percentile(D, 50))
    print("\tQ3:", np.percentile(D, 75))
    print("\tMaximum:", np.max(D))

def ellipse(a,b,phi):
    """
    return the x,y coordinates of an ellipse with major axis=a and minor axis=b at angle phi(radians)
    """
    r = np.zeros(phi.shape) # if a & b is zero, result is zero
    m = np.logical_and(a!=0,b!=0)
    r[m] = a*b/np.sqrt((b*np.cos(phi[m]))**2+(a*np.sin(phi[m]))**2)
    return r

def asymm_ellipse(left, right, upper, lower, phi):
    phi = np.divmod(phi+2*np.pi,2*np.pi)[1] # Be sure phi ∈ [0,2π]
    b = np.where(phi<=np.pi, upper, lower)
    a = np.where(np.logical_or(phi<=np.pi/2, phi>=3*np.pi/2), right, left)
    m = np.logical_and(a!=0,b!=0)
    r = np.zeros(phi.shape) # if a & b is zero, result is zero
    r[m] = a[m]*b[m]/np.sqrt((b[m]*np.cos(phi[m]))**2+(a[m]*np.sin(phi[m]))**2)
    return r

def LG2D(XY, amplitude=1, angle=0, sig_x=10, sig_y=10, x0=None, y0=None, LG_x=0, LG_y=0, assym_x=1, assym_y=1, bg=0):
    """
    Return a 2D Lorentz-Gauss.
    XY: (X,Y) tuple
    A: amplitude
    a: angle
    sx: sigma for x-axis
    sy: sigma fdor y-axis
    x0,y0 : center coordinates of the peak
    lgx, lgy: Lorentz-Gauss proportion (for x,y axis)
    assym_x, assym_y: The assymetry in sig_x/sig_y for the left/right or upper/lower part of the curve
    """
    if x0 is None:
        x0 = XY[0][0,XY[0].shape[1]//2]
    if y0 is None:
        y0 = XY[1][XY[1].shape[0]//2,0]
    X1 = (XY[0]-x0)*np.cos(angle) - (XY[1]-y0)*np.sin(angle)
    Y1 = (XY[0]-x0)*np.sin(angle) + (XY[1]-y0)*np.cos(angle)

    R1 = np.sqrt(X1**2+Y1**2)
    angle = np.arctan2(Y1,X1)
    sig = asymm_ellipse(sig_x, sig_x*assym_x, sig_y, sig_y*assym_y, angle)
    gamma = np.sqrt(2*np.log(2))*sig
    LG = ellipse(LG_x+1, LG_y+1, angle)-1

    Gxy = Gauss(R1, 0, sig, 1)
    Lxy = 1/((R1/gamma)**2+1)

    f = (1-LG)*Gxy+LG*Lxy
    out = bg+amplitude*f
    return out

def LG2Da(XY, amplitude=1, angle=0, sigN=10, sigS=None, sigE=10, sigW=None, x0=None, y0=None, LGN=0, LGS=None, LGE=0, LGW=None, bg=0):
    if x0 is None:
        x0 = XY[0][0,XY[0].shape[1]//2]
    if y0 is None:
        y0 = XY[1][XY[1].shape[0]//2,0]
    if sigS is None:
        sigS = sigN
    if sigW is None:
        sigW = sigE
    if LGS is None:
        LGS = LGN
    if LGW is None:
        LGW = LGE

    X1 = (XY[0]-x0)*np.cos(angle) - (XY[1]-y0)*np.sin(angle)
    Y1 = (XY[0]-x0)*np.sin(angle) + (XY[1]-y0)*np.cos(angle)

    R1 = np.sqrt(X1**2+Y1**2)
    angle = np.arctan2(Y1, X1)
    sig = asymm_ellipse(sigW, sigE, sigN, sigS, angle)
    gamma = np.sqrt(2*np.log(2))*sig # HFHM
    LG = asymm_ellipse(LGW+1, LGE+1, LGN+1, LGS+1, angle)-1

    Gxy = Gauss(R1, 0, sig, 1)
    Lxy = 1/((R1/gamma)**2+1)

    f = (1-LG)*Gxy+LG*Lxy
    out = bg+amplitude*f
    return out

def MaxwellBoltzmann(E, T):
    from . import constants as const
    return 2*const.qe*np.sqrt(E/np.pi)*np.exp(-E/(const.kb*T))/(const.kb*T)**1.5

def Voigt(x, x0, sig, gamma, A=1):
    import scipy
    L = Lorentz(x, x0, gamma, A=1)
    G = Gauss(x, x0, sig)
    out = scipy.signal.convolve(L, G, 'same')
    out /= np.max(out)
    return A*out


<reponame>AhmetCanSolak/aydin
import math
import time
from random import randrange
from typing import Optional, Callable

import numpy
from numba import jit, prange, vectorize, float32
from numpy import absolute
from numpy.typing import ArrayLike
from scipy.ndimage import sobel, gaussian_filter

from aydin.util.edge_filter.fast_edge_filter import fast_edge_filter
from aydin.util.log.log import lprint, lsection

def representative_crop(
    image: ArrayLike,
    mode: str = 'contrast',
    crop_size: Optional[int] = None,
    min_length: int = 32,
    smoothing_sigma: int = 0.5,
    equal_sides: bool = False,
    favour_odd_lengths: bool = False,
    search_mode: str = 'random',
    granularity_factor: int = 4,
    random_search_mode_num_crops: int = 1512,
    min_num_crops: int = 512,
    timeout_in_seconds: float = 2,
    return_slice: bool = False,
    display_crop: bool = False,
    std_fun: Callable = None,  # numpy.std
):
    """Extract a representative crop from the image. Searches for the crop of given
    (approximate) size with highest score. The score is simply the sum of sobel
    magnitudes (~tenengrad) which is a good metric for estimating how much interesting
    content each crop contains. Empirically, this highly correlates with where
    I (<NAME>) tend to look at in images.

    Parameters
    ----------
    image : ArrayLike
        Image to extract representative crop from

    mode : str
        Metric for picking crop. Can be : 'contrast' (fastest), 'sobel', 'sobelmin',
        'sobelmax' We recommend 'contrast'.

    crop_size : int
        Crop size in voxels. Default (None) is 32000.

    min_length : int
        Crop axis lengths cannot be smaller than this number.

    smoothing_sigma : int
        Sigma value for Gaussian filter smoothing to achieve some crude denoising and thus
        make it a bit easier to estimate the score per crop.

    equal_sides : bool
        When True the crop will have all its sides of equal size (square, cube, ...)

    favour_odd_lengths : bool
        If possible favours crops that have odd shape lengths.

    search_mode: bool
        Search mode for best crops. Can be 'random' or 'systematic'. In
        random mode we pick random crops, in systematic mode we check every
        possible strided crop.

    granularity_factor: int
        Granularity of search. higher values correspond to more overlap between candidate crops.

    random_search_mode_num_crops: int
        Number of crops to check in 'random' search mode.

    min_num_crops : int
        Min number of crops to examine.

    timeout_in_seconds: float
        Maximum amount of time in seconds that this function should run for.
        This avoids excessive computation for very large images.

    return_slice : bool
        If True the slice is returned too:

    display_crop: bool
        Displays crop, for debugging purposes...

    Returns
    -------
    Most representative crop, and if return_slice is True the actual slice object too.

    """

    # Debug:
    # _fast_std.parallel_diagnostics(level=4)

    # Std function:
    if std_fun is None:
        std_fun = _fast_std

    # Compile numba functions:
    # _fast_std(numpy.ones(100, dtype=numpy.float32))
    # _normalise(numpy.ones(100, dtype=numpy.float32))

    # Start time:
    start_time = time.time()

    with lsection(
        f"Cropping image of size: {image.shape} with at most {crop_size} voxels and mode {mode}"
    ):

        # save reference to original image:
        original_image = image

        with lsection("Cast and normalise image..."):
            # Cast, if needed:
            image = image.astype(numpy.float32, copy=False)
            # Normalise:
            # image = _normalise(image)

        # Apply filter:
        with lsection(f"Apply cropping filter to image of shape: {image.shape}"):

            # Smoothing:
            sigma = tuple(
                (smoothing_sigma if s > min_length else 0 for s in image.shape)
            )
            image = gaussian_filter(image, sigma=sigma)

            if mode == 'contrast':
                pass
            elif mode == 'sobelfast':
                image = _sobel_fast(image)
            elif mode == 'sobel':
                image = _sobel_magnitude(image)
            elif mode == 'sobelmin':
                image = _sobel_minimum(image)
            elif mode == 'sobelmax':
                image = _sobel_maximum(image)
            else:
                raise ValueError(f"Unknown mode: {mode}")
        #
        # import napari
        # with napari.gui_qt():
        #     viewer = napari.Viewer()
        #     viewer.add_image(image, name='image')
        #     viewer.add_image(filtered_image, name='filtered_image')

        # To speed up cropping we apply a divide and conquer to 'pre-crop' the image recusively:
        # if precropping:
        #     image = _precrop(image,
        #                      crop_size=crop_size,
        #                      min_length=min_length,
        #                      mode=mode,
        #                      smoothing_size=smoothing_size)

        # Number of voxels in image:
        num_voxels = image.size

        # Default number of voxels:
        if crop_size is None:
            crop_size = 32000

        # Ratio by which to crop to achieve max num voxels:
        ratio = (crop_size / num_voxels) ** (1 / image.ndim)

        if ratio >= 1:
            # If the image is small enough no point in getting a crop!
            if return_slice:
                return image, (slice(None),) * image.ndim
            else:
                return image

        # cropped shape:
        cropped_shape = tuple(
            min(max(min_length, int(s * ratio)), s) for s in image.shape
        )

        # If the crop size is still too big, we adjust that. This happens because
        # we cannot crop dimensions that are too small, leading to an
        # underestimation of the ratio.
        for tries in range(8):
            # First let's figure out the current crop size:
            current_crop_size = math.prod(cropped_shape)

            # we check if it is ok, or too large:
            if current_crop_size < 1.05 * crop_size:
                # we are ok if the crop size is within 5% of the desired size.
                break

            # If too large we compute the ratio by which to adjust it:
            ratio = (crop_size / current_crop_size) ** (1 / image.ndim)

            # we compute a new crop shape:
            cropped_shape = tuple(
                min(max(min_length, int(s * ratio)), s) for s in cropped_shape
            )

        # Favour odd lengths if requested:
        if favour_odd_lengths:
            cropped_shape = tuple((s // 2) * 2 + 1 for s in cropped_shape)

        # We enforce equal sides if requested:
        if equal_sides:
            min_length = min(cropped_shape)
            cropped_shape = tuple((min_length,) * image.ndim)

        # range for translation:
        translation_range = tuple(s - cs for s, cs in zip(image.shape, cropped_shape))

        # We loop through a number of crops and keep the one wit the best score:
        best_score = -1
        best_slice = None
        best_crop = None

        # Instead of searching for all possible crops, we take into
        # account the size of the crops to define a 'granularity' (
        # stride) of the translations used for search:

        granularity = tuple(cs // granularity_factor for cs in cropped_shape)

        if search_mode == 'random':

            # We make sure that the number of crops is not too large given
            # the relative size of the crop versus whole image:
            random_search_mode_num_crops = min(
                random_search_mode_num_crops,
                (granularity_factor**image.ndim) * int(image.size / crop_size),
            )

            for i in range(random_search_mode_num_crops):

                # translation:
                translation = tuple(
                    (randrange(0, max(1, (s - cs) // g)) * g if cs != s else 0)
                    for s, cs, g in zip(image.shape, cropped_shape, granularity)
                )

                # function to get crop slice:
                def _crop_slice(translation, cropped_shape, downscale: int = 1):
                    return tuple(
                        slice(t, t + s, downscale)
                        for t, s in zip(translation, cropped_shape)
                    )

                # slice object for cropping:
                crop_slice = _crop_slice(
                    translation, cropped_shape, 2 if image.size > 1e8 else 1
                )

                # extract crop:
                crop = image[crop_slice]

                score = std_fun(crop)

                # slice object for the actual crop:
                crop_slice = _crop_slice(translation, cropped_shape, 1)

                # update best score and image:
                if score > best_score and not math.isinf(score):
                    best_score = score
                    best_slice = crop_slice

                    # We make sure to have the full and original crop!
                    best_crop = original_image[best_slice]

                if i >= min_num_crops and time.time() > start_time + timeout_in_seconds:
                    lprint(
                        f"Interrupting crop search because of timeout after {i} crops examined!"
                    )
                    break

        elif search_mode == 'systematic':

            # grid for translations:
            translation_indices = tuple(
                max(1, int(granularity_factor * r / cs))
                for r, cs in zip(translation_range, cropped_shape)
            )

            for i, index in enumerate(numpy.ndindex(translation_indices)):

                # print(
                #     f"i={i}, index={index}, translation_indices={translation_indices}"
                # )

                # translation:
                translation = tuple(
                    int(j * cs / granularity_factor)
                    for j, cs in zip(index, cropped_shape)
                )

                # slice object for cropping:
                crop_slice = tuple(
                    slice(t, t + s) for t, s in zip(translation, cropped_shape)
                )

                # extract crop:
                crop = image[crop_slice]

                score = std_fun(crop)

                # update best score and image:
                if score > best_score and not math.isinf(score):
                    best_score = score
                    best_slice = crop_slice

                    # We make sure to have the full and original crop!
                    best_crop = original_image[best_slice]

                if i >= min_num_crops and time.time() > start_time + timeout_in_seconds:
                    lprint(
                        f"Interrupting crop search because of timeout after {i} crops examined!"
                    )
                    break
        else:
            raise ValueError(f"Unsupported search mode: {search_mode}")

        if display_crop:

            import napari

            viewer = napari.Viewer()
            viewer.add_image(image.squeeze(), name='image')
            viewer.add_image(best_crop.squeeze(), name='best_crop')
            napari.run()

        #     print(_fast_std.signatures)
        #     for sig in _fast_std.signatures:
        #         print(_fast_std.inspect_asm(sig))`

        if return_slice:
            return best_crop, best_slice
        else:
            return best_crop

def _sobel_magnitude(image):
    magnitude = numpy.zeros_like(image)
    for axis in range(image.ndim):
        if image.shape[axis] < 32:
            continue
        magnitude += sobel(image, axis=axis) ** 2
    return numpy.sqrt(magnitude)

def _sobel_minimum(image):
    minimum = None
    for axis in range(image.ndim):
        if image.shape[axis] < 32:
            continue
        sobel_image = absolute(sobel(image, axis=axis))
        if minimum is None:
            minimum = sobel_image
        else:
            minimum = numpy.minimum(minimum, sobel_image)
    if minimum is None:
        return image
    else:
        return minimum

def _sobel_maximum(image):
    maximum = None
    for axis in range(image.ndim):
        if image.shape[axis] < 32:
            continue
        sobel_image = absolute(fast_edge_filter(image, axis=axis))
        if maximum is None:
            maximum = sobel_image
        else:
            maximum = numpy.maximum(maximum, sobel_image)
    if maximum is None:
        return image
    else:
        return maximum

def _sobel_fast(image):

    longest_axis = max(image.shape)
    axis = image.shape.index(longest_axis)

    return absolute(fast_edge_filter(image, axis=axis))

@jit(nopython=True, parallel=True, fastmath=True)
def _normalise(image):

    # Normalise:
    image_min = float32(image.min())
    image_max = float32(image.max())
    if image_max - image_min > 0:
        return _rescale(image, image_min, image_max)
    else:
        return image - image_min

@vectorize([float32(float32, float32, float32)])
def _rescale(x, min_value, max_value):
    return (x - min_value) / (max_value - min_value)

@jit(nopython=True, parallel=True, fastmath=True)
def _fast_std(image: ArrayLike, workers=16, decimation=1):

    array = image.ravel()
    length = array.size
    num_chunks = workers
    chunk_length = (length // num_chunks) + num_chunks

    std = 0.0

    for c in prange(num_chunks):

        start = c * chunk_length
        stop = (c + 1) * chunk_length
        if stop >= length:
            stop = length

        sub_array = array[start:stop:decimation]
        chunk_std = numpy.std(sub_array)
        std = max(std, chunk_std)

    return std


<reponame>maartenbreddels/mab<filename>mab/gd/jeans/jeansmodel2.py
from numpy import *
import numpy
import mab.gd.logging as logging
import mab.parallelize
from scipy import special
import pickle
import mab.utils.numpy
import os
import sys
from kaplot import *

store_cache = True
logger = logging.getLogger("gd.jeans.jeansmodel")

def h3(x):
	return x**3-3*x

def h4(x):
	return x**4-6*x**2+3
def h6(x):
	return x**6-15*x**4+45*x**2-15

fac3 = 1*2*3
fac4 = fac3*4
fac6 = fac3 * 4 * 5 * 6

def makeresize(parameterset, i, j):
	return (parameterset.ranges_org[i][0], parameterset.ranges_org[j][0]), (parameterset.ranges_org[i][1], parameterset.ranges_org[j][1])

class JeansModel(object):
	def __init__(self, modelpath, modelsetname, parameterset, projection_matrix, observation, use_fourth_moment=False, vmean=0, filters=[], postfix=""):
		self.modelpath = modelpath
		self.parameterset = parameterset
		self.projection_matrix = projection_matrix
		self.observation = observation
		self.schwsetname = modelsetname #"biased_10k_jeans2"
		self.vmean = vmean
		self.postfix = postfix
		self.filters  = filters

		u1=-3.
		u2=3.
		length=250
		self.u = (arange(length)+0.5)/(length+0.) * (u2-u1) + u1
		self.r = 10**self.u
		self.du = (u2-u1)/length

		self.dimension = self.parameterset.dimension

		#self.schwsetname = "biased_jeans2"
		self.dirname = os.path.join(self.modelpath, "jeans", self.schwsetname)

		#self.probability_range = parameterset.ranges_org
		#self.labels = parameterset.paramnames
		#self.paramlabels = parameterset.paramlabels

		self.use_m4 = use_fourth_moment

		self.modelcache = {}
		resultdir = os.path.join(self.dirname, "results")
		if self.use_m4:
			self.modelcachefilename = os.path.join(resultdir, "modelcache_m4_d%d%s_mod.pickle" % (self.parameterset.dimension, self.postfix))
		else:
			self.modelcachefilename = os.path.join(resultdir, "modelcache_d%d%s.pickle" % (self.parameterset.dimension, self.postfix))

		self.init()
	def save(self, iteration=0):
		resultdir = os.path.join(self.dirname, "results")
		logger.info("ensuring directory exists: %s" % resultdir)
		if not os.path.exists(resultdir):
			os.makedirs(resultdir)

		filename = os.path.join(resultdir, "probability_grid_d%d%s%03d.npy" % (self.parameterset.dimension, self.postfix, iteration))
		logger.info("storing probability grid: %s" % filename)
		numpy.save(filename, self.probability_grid)
		if store_cache:
			logger.info("storing modelcache: %s" % self.modelcachefilename)
			fp = open(self.modelcachefilename, "wb")
			pickle.dump(self.modelcache, fp)
			fp.close()

		filename = os.path.join(resultdir, "parameter_points_d%d%s%03d.npy" % (self.parameterset.dimension, self.postfix, iteration))
		logger.info("storing parameter_points: %s" % filename)
		numpy.save(filename, self.parameter_points)

		#numpy.save(filename, self.modelcache)

#self.paramvalues

		return

		filename = os.path.join(resultdir, "anisotropy_grid_%03d.npy" % iteration)
		logger.info("storing anisotropy grid: %s" % filename)
		numpy.save(filename, self.anisotropy_grid)

		filename = os.path.join(resultdir, "mass_enclosed_grid_%03d.npy" % iteration)
		logger.info("storing mass_enclosed grid: %s" % filename)
		numpy.save(filename, self.mass_enclosed_grid)

		filename = os.path.join(resultdir, "logslope_grid_%03d.npy" % iteration)
		logger.info("storing logslope grid: %s" % filename)
		numpy.save(filename, self.logslope_grid)

	def load(self, iteration=0):
		resultdir = os.path.join(self.dirname, "results")

		#filename = os.path.join(resultdir, "probability_grid_%03d.npy" % iteration)
		filename = os.path.join(resultdir, "probability_grid_d%d%s%03d.npy" % (self.parameterset.dimension, self.postfix, iteration))
		logger.info("loading probability grid: %s" % filename)
		self.probability_grid = numpy.load(filename)

		if os.path.exists(self.modelcachefilename):
			logger.info("loading modelcache: %s" % self.modelcachefilename)
			#self.modelcache = numpy.load(filename)
			fp = open(self.modelcachefilename, "rb")
			self.modelcache = pickle.load(fp)
			fp.close()

		#filename = os.path.join(resultdir, "parameter_points_%03d.npy" % iteration)
		filename = os.path.join(resultdir, "parameter_points_d%d%s%03d.npy" % (self.parameterset.dimension, self.postfix, iteration))
		logger.info("loading parameter_points: %s" % filename)
		self.parameter_points = numpy.load(filename)

		return

		filename = os.path.join(resultdir, "anisotropy_grid_%03d.npy" % iteration)
		logger.info("loading anisotropy grid: %s" % filename)
		self.anisotropy_grid = numpy.load(filename)

		filename = os.path.join(resultdir, "mass_enclosed_grid_%03d.npy" % iteration)
		logger.info("loading mass_enclosed grid: %s" % filename)
		self.mass_enclosed_grid = numpy.load(filename)

		filename = os.path.join(resultdir, "logslope_grid_%03d.npy" % iteration)
		logger.info("loading logslope grid: %s" % filename)
		self.logslope_grid = numpy.load(filename)

	def _collect(self, iteration=0, cores=4):
		w = 80

		self.parameter_points = transpose(array(self.parameterset.points))

		#mab.gd.gdfast.ProfilerStart("sparse")
		self.probability_grid = mab.utils.numpy.mmapzeros((w,) * self.parameterset.dimension)
		make_linear = False
		if make_linear:
			self.parameterset.make_linear()
		if self.parameterset.dimension == 1:
			for i in range(w):
				x1, x2 = self.parameterset.ranges_org[0]
				x = x1 + (x2-x1) * i / (w-1.)
				self.probability_grid[i] = self.parameterset((x))
		if self.parameterset.dimension == 2:
			for i in range(w):
				x1, x2 = self.parameterset.ranges_org[0]
				x = x1 + (x2-x1) * i / (w-1.)
				for j in range(w):
					y1, y2 = self.parameterset.ranges_org[1]
					y = y1 + (y2-y1) * j / (w-1.)
					#print self.parameterset((x, y))
					self.probability_grid[i,j] = self.parameterset((x, y))
		if self.parameterset.dimension == 3:
			x1, x2 = 0., 1.#self.parameterset.ranges_org[0]
			y1, y2 = 0., 1.#self.parameterset.ranges_org[1]
			z1, z2 = 0., 1.#self.parameterset.ranges_org[2]
			logger.info("interpolating grid")
			@mab.parallelize.parallelize(cores=cores, info=info)
			def do(i):
				x = x1 + (x2-x1) * i / (w-1.)
				for j in range(w):
					y = y1 + (y2-y1) * j / (w-1.)
					for k in range(w):
						z = z1 + (z2-z1) * k / (w-1.)
						self.probability_grid[i,j,k] = self.parameterset.eval_normalized((x, y, z))
			xindices = range(w)
			do(xindices)
		if self.parameterset.dimension == 4:
			x1, x2 = self.parameterset.ranges_org[0]
			y1, y2 = self.parameterset.ranges_org[1]
			z1, z2 = self.parameterset.ranges_org[2]
			v1, v2 = self.parameterset.ranges_org[3]
			logger.info("interpolating grid")
			#info = False
			@mab.parallelize.parallelize(cores=cores, info=info)
			def do(ij):
					i = ij%w
					j = ij/w
					x = x1 + (x2-x1) * i / (w-1.)
					y = y1 + (y2-y1) * j / (w-1.)
					for k in range(w):
						z = z1 + (z2-z1) * k / (w-1.)
						for l in range(w):
							v = v1 + (v2-v1) * l / (w-1.)
							#print (x, y, z, v), self.parameterset((x, y, z, v))
							self.probability_grid[i,j,k,l] = self.parameterset((x, y, z, v))
			xindices = range(w**2)
			do(xindices)
		if self.parameterset.dimension == 5:
			x1, x2 = self.parameterset.ranges_org[0]
			y1, y2 = self.parameterset.ranges_org[1]
			z1, z2 = self.parameterset.ranges_org[2]
			a1, a2 = self.parameterset.ranges_org[3]
			b1, b2 = self.parameterset.ranges_org[4]
			logger.info("interpolating grid")
			#info = False
			@mab.parallelize.parallelize(cores=cores, info=info)
			def do(ijk):
					i = ijk%w
					j = (ijk/w)%w
					k = (ijk/w)/w
					x = x1 + (x2-x1) * i / (w-1.)
					y = y1 + (y2-y1) * j / (w-1.)
					z = z1 + (z2-z1) * k / (w-1.)
					for l in range(w):
						a = a1 + (a2-a1) * l / (w-1.)
						for m in range(w):
							b = b1 + (b2-b1) * m / (w-1.)
							#print (x, y, z, v), self.parameterset((x, y, z, v))
							self.probability_grid[i,j,k,l,m] = self.parameterset((x, y, z, a, b))
			xindices = range(w**3)
			do(xindices)
		#mab.gd.gdfast.ProfilerStop()
		if not make_linear:
			#print self.probability_grid
			self.probability_grid -= self.probability_grid.max()
			#print self.probability_grid
			self.probability_grid = exp(self.probability_grid)
			mask = numpy.isnan(self.probability_grid)
			self.probability_grid[mask] = 0
			#print self.probability_grid

	def init(self, iteration=0):
		self.parameterset.init()
		#self.parameterset.load(iteration=iteration)
		self.projection_matrix.load()
		self.observation.load()
		self.aperture = self.projection_matrix.gridR

		self.M = self.projection_matrix.matrices

		#print self.Mr.shape
		#print self.projection_matrix.gridR.radial_surface_densities.shape

		light_model = self.projection_matrix.light_model
		self.R = R = light_model.arcsec_to_kpc(self.projection_matrix.gridR.aperture_rcenters)
		dR = light_model.arcsec_to_kpc(self.projection_matrix.gridR.aperture_rcenters[1] - self.projection_matrix.gridR.aperture_rcenters[0])

		total_mass = light_model.total_mass()
		self.massr = light_model.densityr(self.r)/total_mass * self.r**3 * 4 * pi * log(10) * self.du
		#massR = tensordot(self.Ms, self.massr, axes=[(1,), (0,)])
		self.massR = (2*pi*R*dR)*light_model.densityR(R,M=1.)
		print "sum", sum(self.massr)

		for i in range(self.M.shape[0]):
			self.M[i] = (self.M[i].T /self.massR).T
			self.M[i] = (self.M[i] *self.massr)

		self.Ms = self.M[0]
		self.Mr = self.M[0] - self.M[1]
		self.Mt = self.M[1]
		#self.Mr = (self.Mr.T /self.massR).T
		#self.Mt = (self.Mt.T /self.massR).T
		#self.Mr = (self.Mr *self.massr)
		#self.Mt = (self.Mt *self.massr)

		stars = self.observation.stars
		for filter in self.filters:
			stars = filter(stars)

		logger.info("using %d stars/observations" % len(stars))
		#stars_inrange = stars.filter(lambda star: self.storage_2d.aperture.inrange(star.xi, star.eta))
		#print stars
		stars_inrange = []
		for i in range(len(stars)):
			star = stars[i]
			#print star
			if self.aperture.inrange(stars.xi[i], stars.eta[i]):
				stars_inrange.append(star)
		use_numpy = not isinstance(stars, mab.cvsfile.CsvObject)
		if use_numpy:
			stars_inrange = array(stars_inrange, dtype=stars.dtype).view(recarray)
		else:
			stars_inrange = mab.cvsfile.CsvObject(stars_inrange)
		#import pdb;pdb.set_trace()
		logger.info("stars in aperture range     : %d" % len(stars_inrange))
		logger.info("stars outside aperture range: %d" % (len(stars)-len(stars_inrange)))
		stars = stars_inrange
		aperture_indices = []
		for i in range(len(stars)):
			aperture_indices.append(self.aperture.aperture.findindex(stars.xi[i], stars.eta[i]))
		self.aperture_indices = array(aperture_indices)
		#print self.aperture_indices

		real_data = True #False
		if not real_data:
			sigma_v = 2.01
			numpy.random.seed(8)
			for i in range(len(stars)):
				stars.vlos[i] = stars.vlos_true[i] + numpy.random.normal(0, sigma_v)
				#star.v_index = int(((star.vlos+vmax)/(2*vmax)) * self.storage_2d.Nv);
				#outlier = True
				#for losvd in losvds:
				#	if losvd[star.v_index, star.aperture_index] != 0:
				#		outlier = False
				#		break
				#star.is_outlier = outlier

				#print star.aperture_index
			#indices = array([star.aperture_index for star in stars])
		self.vlos = stars.vlos - self.vmean #array([star.vlos for star in stars])
		self.vlos_sigma = stars.e_vlos #array([star.e_vlos for star in stars])
		print self.aperture_indices.min(), self.aperture_indices.max()
		#print self.storage_2d.aperture
		#moment0 = self.storage_2d.aperture.radial_surface_densities
		#moment0 = self.aperture.radial_surface_densities
		#print vlos_sigma
		#print indices[argsort(indices)]
		#print stars.rc[argsort(indices)]

		if os.path.exists(self.modelcachefilename):
			logger.info("loading modelcache: %s" % self.modelcachefilename)
			fp = open(self.modelcachefilename, "rb")
			self.modelcache = pickle.load(fp)
			fp.close()
			#self.modelcache = numpy.load(filename)
			logger.info("model cache contains %d models" % len(self.modelcache))

	def test(self, init, scope):
		self.scope = scope
		mozaic(2,1,box)

		scope.flush()
		scope.init()
		scope["dm_density_twoslope.rs"] = 0.5
		scope["dm_density_twoslope.alpha"] = 1
		#scope["jeans.beta"] = -0.5
		jeans = scope["jeans"]
		self.drawsigmar(jeans, color="black")

		scope.flush()
		scope.init()
		scope["dm_density_twoslope.rs"] = 10**-0.85
		scope["dm_density_twoslope.alpha"] = -0.56
		#scope["jeans.beta"] = -0.5
		jeans = scope["jeans"]
		self.drawsigmar(jeans, color="green", linestyle="dot")

		scope.flush()
		scope.init()
		scope["dm_density_twoslope.rs"] = 10**-0.57
		scope["dm_density_twoslope.alpha"] = 0.37
		#scope["jeans.beta"] = -0.5
		jeans = scope["jeans"]
		self.drawsigmar(jeans, color="red", linestyle="dash")

		scope.flush()
		scope.init()
		scope["dm_density_twoslope.rs"] = 10**-0.392
		scope["dm_density_twoslope.alpha"] = 0.860
		#scope["jeans.beta"] = -0.5
		jeans = scope["jeans"]
		self.drawsigmar(jeans, color="blue", linestyle="dot")

		select(0, 0)
		ylim(0, 15)
		select(1, 0)
		ylim(0, 15)
		draw()

	def drawsigmar(self, jeans, **kwargs):
		m2 = array([jeans.m2(k) for k in self.r])
		sigmar = sqrt(m2)
		beta = jeans.beta
		m2_los = tensordot(self.M[0]-beta*self.M[1], m2, axes=[(1,), (0,)])
		sigma_los = m2_los**0.5
		select(0, 0)
		graph(self.u, sigmar, **kwargs)
		select(1, 0)
		graph(self.R, sigma_los, **kwargs)

	def run(self, args, opts, scope, info=True):
		self.scope = scope
		cores=opts.cores
		#if 1:
		logm = 7.91
		if 0:
			scope.flush()
			scope.init()
			logprior = 0
			paramnames = ['dm_density_twoslope.rs', 'dm_density_twoslope.alpha']#, 'jeans.beta']
			values = [10**-0.5, 0.54, -0.5]
			values = [3.16227766017, -0.25]

			for name, value in zip(paramnames, values):
				scope[name] = value
			#logprior = log(paramvalue.p_prior)
			#scope["jeans.beta"] = -0.5
			jeans = scope["jeans"]
			print self.run_one(0, test=True)
			sys.exit(0)
			#print "beta = ", jeans.beta
			beta = jeans.beta
			logr = arange(-2, 2, 0.1)
			r = 10**logr
			m2 = numpy.array([jeans.m2(k) for k in r])
			m4 = numpy.array([jeans.m4(k) for k in r])
			gamma = m4/m2**2
			graph(logr, gamma)
			draw()
			sys.exit(0)

		if 0:
			if 1:
				i = 0
				logm += 0.01
				logrs = -0.4796
				self.parameterset.paramvalues[i].values_org = list(self.parameterset.paramvalues[i].values_org)
				self.parameterset.paramvalues[i].values_trans = list(self.parameterset.paramvalues[i].values_trans)
				self.parameterset.paramvalues[i].values_org[0] = logm
				self.parameterset.paramvalues[i].values_trans[0] = 10**logm
				self.parameterset.paramvalues[i].values_org[1] = logrs
				self.parameterset.paramvalues[i].values_trans[1] = 10**logrs
				print self.run_one(0)
				logrs = -0.4795
				self.parameterset.paramvalues[i].values_org = list(self.parameterset.paramvalues[i].values_org)
				self.parameterset.paramvalues[i].values_trans = list(self.parameterset.paramvalues[i].values_trans)
				self.parameterset.paramvalues[i].values_org[0] = logm
				self.parameterset.paramvalues[i].values_trans[0] = 10**logm
				self.parameterset.paramvalues[i].values_org[1] = logrs
				self.parameterset.paramvalues[i].values_trans[1] = 10**logrs
				print self.run_one(0)
				sys.exit(0)
			while 1:
				logm += 0.01
				print "run, logm=", logm
				if 0:
					for i, p in enumerate(self.parameterset.paramvalues):
						v = p.values_org
						if (v[0] >= 7.97) and (v[0] < 7.98) and (v[1] > -0.52) and (v[1] < -0.25):
							print p.modelname, v
							print self.modelcache[p.modelname]

				#logrss = arange(-0.5, -0.45, 0.01/10)
				#logrss = arange(-0.55, -0.45, 0.01/100)
				logrss = arange(-0.4798, -0.4794, 0.01/1000)
				#logrss = arange(-0.5, 0.0, 0.01/4)
				logps = []
				for i, logrs in enumerate(logrss):
					#logm = 8.0 #7.9749999999999996
					self.parameterset.paramvalues[i].values_org = list(self.parameterset.paramvalues[i].values_org)
					self.parameterset.paramvalues[i].values_trans = list(self.parameterset.paramvalues[i].values_trans)
					self.parameterset.paramvalues[i].values_org[0] = logm
					self.parameterset.paramvalues[i].values_trans[0] = 10**logm
					self.parameterset.paramvalues[i].values_org[1] = logrs
					self.parameterset.paramvalues[i].values_trans[1] = 10**logrs
				@mab.parallelize.parallelize(cores=cores, info=info)
				def wrap(index):
					return self.run_one(index)
				logps = wrap(range(len(logrss)))
				print logps
				#logp = self.run_one(0)
				#print logrs, logp
				#	logps.append(logps)
				logps = array(logps)
				logps -= logps.max()
				ps = exp(logps)
				clear()
				box()
				graph(logrss, ps)
				draw()

			sys.exit(0)
		def do():
			self.parameterset.valuemap = {}
			all_indices = range(len(self.parameterset.parameter_values))
			# skip cached models
			known_indices = [i for i in all_indices if self.parameterset.parameter_values[i].name in self.modelcache]
			unknown_indices = [i for i in all_indices if self.parameterset.parameter_values[i].name not in self.modelcache]
			if 0:
				self.run_one(0)
			else:
				@mab.parallelize.parallelize(cores=cores, info=info)
				def wrap(index):
					return self.run_one(index)
				if cores == 1:
					for i in unknown_indices:
						print i
						print self.run_one(i)
				else:
					logps_solved = wrap(unknown_indices)
					logps_solved = array(logps_solved)
					logps = zeros(len(all_indices))
					logps[unknown_indices] = logps_solved
					logps[known_indices] = [self.modelcache[self.parameterset.parameter_values[i].name] for i in known_indices]
					for i in unknown_indices:
						self.modelcache[self.parameterset.parameter_values[i].name] = logps[i]
					for i in range(len(self.parameterset.parameter_values)):
						parametervalue = self.parameterset.parameter_values[i]
						self.parameterset.valuemap[tuple(parametervalue.values)] = logps[i]
					filename = os.path.join(self.modelpath, "logps.npy")
					save(filename, logps)
					resultdir = os.path.join(self.dirname, "results")
					if store_cache:
						dirname = os.path.dirname(self.modelcachefilename)
						if not os.path.exists(dirname):
							os.makedirs(dirname)
						logger.info("storing modelcache: %s (%d models)" % (self.modelcachefilename, len(self.modelcache)))
						fp = open(self.modelcachefilename, "wb")
						pickle.dump(self.modelcache, fp)
						fp.close()

					print "logps", logps
				#self.parameterset.feed(logps)

		#print self.parameterset.paramnames
		do()
		if 0:
			for i in range(self.parameterset.max_iterations):
				#print "iterate"
				logger.info("iteration %d (out of %i)" % (i+1,self.parameterset.max_iterations))
				self.parameterset.iterate(i,cores=cores)
				do()
				#self.parameterset.refresh()

		self.parameterset.grid()
		#self._collect(iteration=0, cores=cores)

	def run_one(self, index, test=False):
		scope = self.scope
		if not test:
			#paramnames = self.parameterset.parameter_values
			paramvalues = self.parameterset.parameter_values[index]
			#paramprior = self.parameterset.parampriors
			#values = paramvalue.values
			scope.flush()
			scope.re_readfiles()
			scope.init()
			logprior = 0
			#for name, value in zip(paramnames, values):
			#for paramvalues in self.parameterset.parameter_values[index]
			for name, value in paramvalues.items:
				#paramvalues.name, paramvalues.value
				#print name, value
				scope[name] = value
			logprior = log(paramvalues.p_prior)
			#scope["jeans.beta"] = -0.5
		jeans = scope["jeans"]
		#print "beta = ", jeans.beta
		beta = jeans.beta(self.r)

		#jeans = jeans.fast()
		#sigmar = jeans.sigmar(self.r)
		#print jeans.m2(280.0)
		#print [(k,jeans.m2(k)) for k in self.r]
		#dsa

		sigmar = array([jeans.m2(k)**0.5 for k in self.r])
		#if any(isnan(sigmar)):
		sigmar[isnan(sigmar)] = 0
		#if any(isnan(sigmar)):
		#	import pdb;
		#	pdb.set_trace()
		#print sigmar
		#g = 1 - 2 * self.beta * R**2/r**2 + self.beta*(1+self.beta)/2*R**4/r**4

		use_m4 = self.use_m4
		if use_m4:
			m4 = array([jeans.m4(k) for k in self.r])
			kappa4 = m4##/ sigmar**4

		varr = sigmar**2# * self.massr
		vart =(1-beta)*varr

		varr_los = tensordot(self.Mr, varr, axes=[(1,), (0,)])
		vart_los = tensordot(self.Mt, vart, axes=[(1,), (0,)])
		#light_model = jeans.light_model
		#varR_los = (varr_los + vart_los)
		m2 = sigmar**2
		m2_los = tensordot(self.M[0]-beta*self.M[1], m2, axes=[(1,), (0,)])
		if use_m4:
			#g = 1 - 2 * self.beta * R**2/r**2 + self.beta*(1+self.beta)/2*R**4/r**4
			m4_los = tensordot(self.M[0]-2*beta*self.M[1]+beta*(1+beta)/2*self.M[3], m4, axes=[(1,), (0,)])
		# m2*(1-beta*
		varR_los = m2_los
		m2_los = varR_los[self.aperture_indices] + self.vlos_sigma**2
		if use_m4:
			m4_los = m4_los[self.aperture_indices] + 3*self.vlos_sigma**4
			kappa4 = m4_los
		sigma_los = sqrt(m2_los)
		#print sigma_los
		#print sigmasqs
		#print sigmasqs.shape
		logps = sum(- self.vlos**2/(2*m2_los) - log(sqrt(m2_los*2*pi)))
		if use_m4:
			u = (kappa4/sigma_los**4-3)
			g = m4_los/sigma_los**4

			if 0:
				a1, a2, a3, b1, b2,b3 = [-0.10971413, -0.03265063,  0.02697869,  -0.52812814, -0.05336269, -0.0592689] # unnormalized
				gamma = g
				mask = gamma < 3
				#print "%e" % gamma
				sigma = sigma_los
				sigma_mod = sigma * 1.0
				sigma_mod[mask] = sigma[mask] * (1 + (3-gamma[mask])**2*a1 + (3-gamma[mask])**3*a2 + (3-gamma[mask])**4*a3)
				kappa2 = sigma**2
				kappa4 = gamma * kappa2**2
				kappa4_mod = kappa4 * 1.
				kappa4_mod[mask] = kappa4[mask] * (1 + (3-gamma[mask])**2*b1 + (3-gamma[mask])**3*b2 + (3-gamma[mask])**4*b3)
				gamma_mod = gamma * 1.
				gamma_mod[mask] = kappa4_mod[mask]/kappa2[mask]**2
				u = gamma_mod * 1.-3
				sigma_los = sigma_mod

			#kappa4 = kappa4_mod
			#print list(g)
			#print
			#print g.min(), g.max(), mean(g)
			#u = minimum(u, 2)
			#u = maximum(u, -1)
			#u = maximum(u, -0.25)
			#u = minimum(u, 0.25)
			kappa4 = (u+3)*sigma_los**4
			s4 = (kappa4 - 3 * sigma_los**4)/sigma_los**6
			b1 = 0 #self.sigma*(1./fac3 * self.s3 * h3(x/self.sigma))
			b2 = sigma_los**2 * (1./fac4 * s4 * h4(self.vlos/sigma_los))
			#print list(g)
			#print b2.min(), b2.max(), mean(b2)
			extra = b2
			#mask = (1+b2)>0
			mask = u > 0.
			mask = b2 > 0.
			z = b2
			z[mask] = log(1+b2[mask])

			if 0:
				u = (kappa4/sigma_los**4-3)
				#print b2[mask].max(), u.min(), u.max()
				assert numpy.all(1+b2[mask]) > 0, zip(b2[mask], u[mask], s4[mask])
				oldb2 = b2 * 1
				b2[mask] = log(1+b2[mask])
				#print b2.min(), b2.max(), mean(b2)
				nans = numpy.isnan(b2)
				if isnan(sum(b2)) and not test:
					for name, value in zip(paramnames, values):
						print name, value

				#assert not isnan(sum(b2)), zip(b2[nans], oldb2[nans], u[nans], s4[nans], (self.vlos/sigma_los)[nans], self.vlos[nans], sigma_los[nans])
				print logps, sum(b2)
			else:
				assert not isnan(sum(z)), zip(b2[nans], oldb2[nans], u[nans], s4[nans], (self.vlos/sigma_los)[nans], self.vlos[nans], sigma_los[nans])
			r = logps + sum(z)
			import pdb;
			#pdb.set_trace()
			return r + logprior
		else:
			return logps + logprior

		if 0:
			from kaplot import *
			box()
			graph(sqrt(varR_los))
			sigmar = jeans.sigma_los(self.R)
			graph(sigmar, color="green")
			draw()
		if 0:
			R = light_model.arcsec_to_kpc(self.projection_matrix.gridR.aperture_rcenters)
			dR = light_model.arcsec_to_kpc(self.projection_matrix.gridR.aperture_rcenters[1] - self.projection_matrix.gridR.aperture_rcenters[0])
			varR_los = (varr_los + vart_los)#/self.massR # /(2*pi*R*dR)*(jeans.light_model.densityR(R)/1e6)#massR*(2*pi*R*dR)# * 10
			mozaic(2,2,box)
			area = 1 #R / jeans.light_model.densityR(R,M=1.) # * self.projection_matrix.gridR.radial_surface_densities# / 100
			sigmar = jeans.sigma_los(R)
			graph(sigmar, color="green")
			graph(sqrt(varR_los*area))
			print sigmar.shape, varR_los.shape
			graph(sigmar**2/(varR_los*area), color="red")
			select(0,1)
			graph(self.u, jeans.light_model.densityr(self.r))
			select(1,1)
			graph(self.projection_matrix.gridR.radial_surface_densities)
			select(1,0)
			graph(self.massR/(2*pi*R*dR), color="red")

			graph(jeans.light_model.densityR(R)/1e6, color="green", linestyle="dot")
			draw()
		#print u


<gh_stars>0
import numpy
from pylab import *
from scipy.interpolate import interp1d

d1,e1,ee1,f1,ef1=numpy.loadtxt("full.txt",unpack=True)
f1=-f1*31.6e-15
inds=argsort(d1)
d1=d1[inds]
f1=f1[inds]

d2,e2,ee2,f2,ef2=numpy.loadtxt("PEC.txt",unpack=True)
f2=-f2*31.6e-15
inds=argsort(d2)
d2=d2[inds]
f2=f2[inds]

d3,e3,ee3,f3,ef3=numpy.loadtxt("temp.txt",unpack=True)
f3=-f3*31.6e-15
inds=argsort(d3)
d3=d3[inds]
f3=f3[inds]

datafile="../../Mathematica/calculated_vals.tsv"
dist,fpfa,fnaive,fright,ftemp=numpy.loadtxt(datafile,unpack=True)
dist=dist*1e6

PFA_datafile="../../Mathematica/calculated_pfa_vals.tsv"
dist2,fpfa2,fnaive2,fright2,ftemp2=numpy.loadtxt(PFA_datafile,unpack=True)
dist2=dist2*1e6

plot(d2,f2,label="PEC")
plot(d1,f1,label="SiO2/Au")
plot(d3,f3,label="PEC T=300")

plot(dist,fpfa,label="PFA",linestyle='dashed')
plot(dist,fright,label="SiO2/Au",linestyle='dashed')
plot(dist,ftemp,label="SiO2/Au T=300",linestyle='dashed')
xscale('log')
yscale('log')
xlabel('Distance (microns)')
ylabel('Force (N)')
title('Analytical (Dashed) v Numerical (Solid) Calculations')
legend()
show()
#savefig('analytic_v_numerical')

clf()
iPFA = interp1d(dist,fpfa)
rPFA=iPFA(d2)/f2
iPFA2 = interp1d(dist2,fpfa2)
rPFA2=iPFA2(d2)/f2
plot(d2,rPFA,label="Gradient Expansion PFA")
plot(d2,rPFA2,label="Normal PFA")
xscale('log')
yscale('log')
xlabel('Distance (Microns)')
ylabel('(PFA/BEM) Force Ratio')
title('Comparion between Calculations, grid=1 micron')
legend()
savefig("pfa_v_pec.png")

clf()
tag,xi,e,f=numpy.loadtxt("mesh_byXi.txt",unpack=True)
f=-f*31.6e-15
xis=numpy.unique(xi)
for ix in range(0,len(xis)):
    inds = numpy.where(xi == xis[ix])
    xplot=tag[inds]
    yplot=f[inds]
    plot(xplot,yplot,label=str(xis[ix]))

xscale('log')
yscale('log')
xlim(0.5,5.0)
legend(title='Int. Frequency',loc='lower left')
xlabel('Mesh length scale (microns)')
ylabel('Force Integrand Value')
title('Convergence Tests @ 1 micron')
savefig('convergence.png')


<reponame>kjanjua26/CHIPathy<filename>generatedata/aug.py<gh_stars>1-10
"""
Augment the images in the directory and generate multiple images from that \
    to increase the numnber of available data.
"""
import cv2
import imgaug.augmenters as iaa
import matplotlib.pyplot as plt
from scipy import ndimage
import numpy as np
import argparse
from glob import glob
from tqdm import tqdm

parser = argparse.ArgumentParser(description="Augment the images in the directory and generate multiple images.",
                                formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument('--src', '-s', type=str, required=True,
                    help='The source directory containing the images.')
parser.add_argument('--dst', '-d', type=str, required=True,
                    help='The destination directory to save the images.')
parser.add_argument('--techniques', '-t', type=str,
                    help='The techniques to apply to augment the images.\nBy default all the techniques will be applied.')
args = parser.parse_args()

class Augment:
    def __init__(self):
        self.input_dir = args.src
        self.output_dir = args.dst
        self.techniques = args.techniques
        self.zoom_factors = [1.15, 1.25, 1.35, 1.45, 1.50, 1.60, 1.70, 1.80, 1.90]
        self.severity = [1, 2, 3, 4, 5]
        self.count = 0

        if self.techniques is not None:
            print(f'The augmentations to apply are: {[x for x in self.techniques]}')
        else:
            print("Applying all the augmentations.")

    def crop_image_from_gray(self, img, tol=10):
        dims = img.ndim

        if dims == 3:
            gray_image = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        elif dims == 2:
            gray_image = img

        mask = gray_image > tol

        # find start crop point
        objs = ndimage.find_objects(mask)
        height_offset = objs[0][0].start
        width_offset = objs[0][1].start
        ix = np.ix_(mask.any(1),mask.any(0))

        if dims == 2:
            return img[ix], height_offset, width_offset

        check_shape = img[:,:,0][ix].shape[0]
        if (check_shape == 0): # image is too dark so that we crop out everything,
            return img, 0, 0 # return original image

        img1 = img[:,:,0][ix]
        img2 = img[:,:,1][ix]
        img3 = img[:,:,2][ix]
        img = np.stack([img1, img2, img3],axis=-1)

        return img, height_offset, width_offset

    def circle_crop(self, img, sigmaX=50):
        """
        Create circular crop around image centre
        """
        img, height_offset, width_offset = self.crop_image_from_gray(img)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        img = cv2.addWeighted(img,4, cv2.GaussianBlur(img, (0,0), sigmaX), -4, 128)

        return img, height_offset, width_offset

    def flip_image(self, img, mode='h'):
        """
        Flips the image horizontally or vertically depending on the mode.
        """

        if mode == 'h':
            aug = iaa.Fliplr()
        else:
            aug = iaa.Flipud()

        flipped_img = aug(image=img)
        return flipped_img

    def compression_blur(self, img, severity):
        """
        Blur in an image caused due to the JPEG compression.
        """
        aug = iaa.imgcorruptlike.JpegCompression(severity=severity)
        return aug(image=img)

    def cv2_clipped_zoom(self, img, zoom_factor):
        """
        Center zoom in/out of the given image and returning an enlarged/shrinked view of
        the image without changing dimensions
        Args:
            img : Image array
            zoom_factor : amount of zoom as a ratio (0 to Inf)

        Ftn Taken From: https://stackoverflow.com/a/48097478/6735773
        """
        height, width = img.shape[:2] # It's also the final desired shape
        new_height, new_width = int(height * zoom_factor), int(width * zoom_factor)

        ### Crop only the part that will remain in the result (more efficient)
        # Centered bbox of the final desired size in resized (larger/smaller) image coordinates
        y1, x1 = max(0, new_height - height) // 2, max(0, new_width - width) // 2
        y2, x2 = y1 + height, x1 + width
        bbox = np.array([y1,x1,y2,x2])
        # Map back to original image coordinates
        bbox = (bbox / zoom_factor).astype(np.int)
        y1, x1, y2, x2 = bbox
        cropped_img = img[y1:y2, x1:x2]

        # Handle padding when downscaling
        resize_height, resize_width = min(new_height, height), min(new_width, width)
        pad_height1, pad_width1 = (height - resize_height) // 2, (width - resize_width) //2
        pad_height2, pad_width2 = (height - resize_height) - pad_height1, (width - resize_width) - pad_width1
        pad_spec = [(pad_height1, pad_height2), (pad_width1, pad_width2)] + [(0,0)] * (img.ndim - 2)

        result = cv2.resize(cropped_img, (resize_width, resize_height))
        result = np.pad(result, pad_spec, mode='constant')
        assert result.shape[0] == height and result.shape[1] == width
        return result

    def main(self):
        """
        The main caller function of the augmentation file.
        """

        exts = glob(self.input_dir + "*.jpeg") + glob(self.input_dir + "*.png") + \
                glob(self.input_dir + "*.jpg") + glob(self.input_dir + "*.tif")
        list_of_images = [x for x in exts]
        list_of_transformed_images = []

        for i, img_path in enumerate(list_of_images):
            print(f"Processing {i}/{len(list_of_images)}")
            img = cv2.imread(img_path)

            print("Applying the base transformation: CropnBlur.")
            img_t, height_offset, width_offset = self.circle_crop(img, 50)
            list_of_transformed_images.append(img_t)

            print("Applying the horizontal/vertical flips.")
            img_t_fliplr = self.flip_image(img_t)
            list_of_transformed_images.append(img_t_fliplr)

            img_t_flipud = self.flip_image(img_t, mode='v')
            list_of_transformed_images.append(img_t_flipud)

            print("Applying zoom transformation.")
            for zoom_factor in self.zoom_factors:
                zoomed_img = self.cv2_clipped_zoom(img_t, zoom_factor)
                list_of_transformed_images.append(zoomed_img)

            print("Applying JPEG compression blur.")
            for sev in self.severity:
                compressed_img_t = self.compression_blur(img_t, sev)
                compressed_img_fliplr = self.compression_blur(img_t_fliplr, sev)
                compressed_img_flipud = self.compression_blur(img_t_flipud, sev)

                list_of_transformed_images.append(compressed_img_t)
                list_of_transformed_images.append(compressed_img_fliplr)
                list_of_transformed_images.append(compressed_img_flipud)

        for transformed_img in tqdm(list_of_transformed_images):
            self.count += 1
            transformed_img = cv2.cvtColor(transformed_img, cv2.COLOR_BGR2RGB)
            cv2.imwrite(self.output_dir + f"{self.count}.jpg", transformed_img)

if __name__ == "__main__":
    ag = Augment()
    ag.main()


<reponame>sumau/tick<gh_stars>100-1000
# License: BSD 3 clause

import unittest

import numpy as np
from scipy.sparse import csr_matrix

from tick.linear_model import SimuLinReg
from tick.base_model.tests.generalized_linear_model import TestGLM
from tick.robust import ModelLinRegWithIntercepts
from tick.simulation import weights_sparse_gauss

class Test(TestGLM):
    def test_ModelLinRegWithInterceptsWithGlobalIntercept(self):
        """...Numerical consistency check of loss and gradient for linear
        regression with sample intercepts and a global intercept
        """
        np.random.seed(12)
        n_samples, n_features = 200, 5
        w0 = np.random.randn(n_features)
        intercept0 = 50 * weights_sparse_gauss(n_weights=n_samples, nnz=30)
        c0 = None
        X, y = SimuLinReg(w0, c0, n_samples=n_samples, verbose=False,
                          seed=2038).simulate()
        # Add gross outliers to the labels
        y += intercept0
        X_spars = csr_matrix(X)
        model = ModelLinRegWithIntercepts(fit_intercept=False).fit(X, y)
        model_spars = ModelLinRegWithIntercepts(fit_intercept=False) \
            .fit(X_spars, y)
        self.run_test_for_glm(model, model_spars)
        self._test_glm_intercept_vs_hardcoded_intercept(model)

        self.assertAlmostEqual(model.get_lip_mean(), 6.324960325598532)
        self.assertAlmostEqual(model.get_lip_max(), 30.277118951892113)
        self.assertAlmostEqual(model.get_lip_mean(),
                               model_spars.get_lip_mean())
        self.assertAlmostEqual(model.get_lip_max(), model_spars.get_lip_max())
        self.assertAlmostEqual(model.get_lip_best(), 2.7217793249045439)

    def test_ModelLinRegWithInterceptsWithoutGlobalIntercept(self):
        """...Numerical consistency check of loss and gradient for linear
        regression with sample intercepts and no global intercept
        """
        np.random.seed(12)
        n_samples, n_features = 200, 5
        w0 = np.random.randn(n_features)
        intercept0 = 50 * weights_sparse_gauss(n_weights=n_samples, nnz=30)
        c0 = None
        X, y = SimuLinReg(w0, c0, n_samples=n_samples, verbose=False,
                          seed=2038).simulate()
        # Add gross outliers to the labels
        y += intercept0
        X_spars = csr_matrix(X)
        model = ModelLinRegWithIntercepts(fit_intercept=True).fit(X, y)
        model_spars = ModelLinRegWithIntercepts(fit_intercept=True) \
            .fit(X_spars, y)
        self.run_test_for_glm(model, model_spars)

        self.assertAlmostEqual(model.get_lip_mean(), 7.324960325598536)
        self.assertAlmostEqual(model.get_lip_max(), 31.277118951892113)
        self.assertAlmostEqual(model.get_lip_mean(),
                               model_spars.get_lip_mean())
        self.assertAlmostEqual(model.get_lip_max(), model_spars.get_lip_max())
        self.assertAlmostEqual(model.get_lip_best(), 2.7267793249045438)

    def test_ModelLinRegWithInterceptsWithoutGlobalInterceptExtras(self):
        """...Extra tests for linear regression with sample intercepts and not
        global intercept, check gradient wrt homemade gradient
        """
        np.random.seed(12)
        n_samples, n_features = 200, 5
        w0 = np.random.randn(n_features)
        intercept0 = 50 * weights_sparse_gauss(n_weights=n_samples, nnz=30)
        c0 = None
        X, y = SimuLinReg(w0, c0, n_samples=n_samples, verbose=False,
                          seed=2038).simulate()
        # Add gross outliers to the labels
        y += intercept0
        model = ModelLinRegWithIntercepts(fit_intercept=False).fit(X, y)
        coeffs = np.random.randn(model.n_coeffs)
        grad1 = model.grad(coeffs)
        X2 = np.hstack((X, np.identity(n_samples)))
        grad2 = X2.T.dot(X2.dot(coeffs) - y) / n_samples
        np.testing.assert_almost_equal(grad1, grad2, decimal=10)

    def test_ModelLinRegWithInterceptsWithGlobalInterceptExtras(self):
        """...Extra tests for linear regression with sample intercepts and
        global intercept, check gradient wrt homemade gradient
        """
        np.random.seed(12)
        n_samples, n_features = 200, 5
        w0 = np.random.randn(n_features)
        intercept0 = 50 * weights_sparse_gauss(n_weights=n_samples, nnz=30)
        c0 = -1.
        X, y = SimuLinReg(w0, c0, n_samples=n_samples, verbose=False,
                          seed=2038).simulate()
        # Add gross outliers to the labels
        y += intercept0
        model = ModelLinRegWithIntercepts(fit_intercept=True).fit(X, y)
        coeffs = np.random.randn(model.n_coeffs)
        grad1 = model.grad(coeffs)
        X2 = np.hstack((X, np.ones((n_samples, 1)), np.identity(n_samples)))
        grad2 = X2.T.dot(X2.dot(coeffs) - y) / n_samples
        np.testing.assert_almost_equal(grad1, grad2, decimal=10)

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