# wavelet library

def wavelet(Y,dt,pad=0.,dj=0.25,s0=-1,J1=-1,mother="MORLET",param=-1):
    """
This function is the translation of wavelet.m by Torrence and Compo

import wave_bases from wave_bases.py

The following is the original comment in wavelet.m

#WAVELET  1D Wavelet transform with optional singificance testing
%
%   [WAVE,PERIOD,SCALE,COI] = wavelet(Y,DT,PAD,DJ,S0,J1,MOTHER,PARAM)
%
%   Computes the wavelet transform of the vector Y (length N),
%   with sampling rate DT.
%
%   By default, the Morlet wavelet (k0=6) is used.
%   The wavelet basis is normalized to have total energy=1 at all scales.
%
%
% INPUTS:
%
%    Y = the time series of length N.
%    DT = amount of time between each Y value, i.e. the sampling time.
%
% OUTPUTS:
%
%    WAVE is the WAVELET transform of Y. This is a complex array
%    of dimensions (N,J1+1). FLOAT(WAVE) gives the WAVELET amplitude,
%    ATAN(IMAGINARY(WAVE),FLOAT(WAVE) gives the WAVELET phase.
%    The WAVELET power spectrum is ABS(WAVE)^2.
%    Its units are sigma^2 (the time series variance).
%
%
% OPTIONAL INPUTS:
%
% *** Note *** setting any of the following to -1 will cause the default
%               value to be used.
%
%    PAD = if set to 1 (default is 0), pad time series with enough zeroes to get
%         N up to the next higher power of 2. This prevents wraparound
%         from the end of the time series to the beginning, and also
%         speeds up the FFT's used to do the wavelet transform.
%         This will not eliminate all edge effects (see COI below).
%
%    DJ = the spacing between discrete scales. Default is 0.25.
%         A smaller # will give better scale resolution, but be slower to plot.
%
%    S0 = the smallest scale of the wavelet.  Default is 2*DT.
%
%    J1 = the # of scales minus one. Scales range from S0 up to S0*2^(J1*DJ),
%        to give a total of (J1+1) scales. Default is J1 = (LOG2(N DT/S0))/DJ.
%
%    MOTHER = the mother wavelet function.
%             The choices are 'MORLET', 'PAUL', or 'DOG'
%
%    PARAM = the mother wavelet parameter.
%            For 'MORLET' this is k0 (wavenumber), default is 6.
%            For 'PAUL' this is m (order), default is 4.
%            For 'DOG' this is m (m-th derivative), default is 2.
%
%
% OPTIONAL OUTPUTS:
%
%    PERIOD = the vector of "Fourier" periods (in time units) that corresponds
%           to the SCALEs.
%
%    SCALE = the vector of scale indices, given by S0*2^(j*DJ), j=0...J1
%            where J1+1 is the total # of scales.
%
%    COI = if specified, then return the Cone-of-Influence, which is a vector
%        of N points that contains the maximum period of useful information
%        at that particular time.
%        Periods greater than this are subject to edge effects.
%        This can be used to plot COI lines on a contour plot by doing:
%
%              contour(time,log(period),log(power))
%              plot(time,log(coi),'k')
%
%----------------------------------------------------------------------------
%   Copyright (C) 1995-2004, Christopher Torrence and Gilbert P. Compo
%
%   This software may be used, copied, or redistributed as long as it is not
%   sold and this copyright notice is reproduced on each copy made. This
%   routine is provided as is without any express or implied warranties
%   whatsoever.
%
% Notice: Please acknowledge the use of the above software in any publications:
%    ``Wavelet software was provided by C. Torrence and G. Compo,
%      and is available at URL: http://paos.colorado.edu/research/wavelets/''.
%
% Reference: Torrence, C. and G. P. Compo, 1998: A Practical Guide to
%            Wavelet Analysis. <I>Bull. Amer. Meteor. Soc.</I>, 79, 61-78.
%
% Please send a copy of such publications to either C. Torrence or G. Compo:
%  Dr. Christopher Torrence               Dr. Gilbert P. Compo
%  Research Systems, Inc.                 Climate Diagnostics Center
%  4990 Pearl East Circle                 325 Broadway R/CDC1
%  Boulder, CO 80301, USA                 Boulder, CO 80305-3328, USA
%  E-mail: chris[AT]rsinc[DOT]com         E-mail: compo[AT]colorado[DOT]edu
%----------------------------------------------------------------------------"""
    #modules
    import numpy as np

    #set default
    n1 = len(Y)
    if (s0 == -1): s0=2.*dt
    if (dj == -1): dj = 1./4.
    if (J1 == -1): J1=np.fix((np.log(n1*dt/s0)/np.log(2))/dj)
    if (mother == -1): mother = 'MORLET'
    #print "s0=",s0
    #print "J1=",J1

    #....construct time series to analyze, pad if necessary
    x = Y - np.mean(Y);
    if (pad == 1):
        base2 = np.fix(np.log(n1)/np.log(2) + 0.4999)   # power of 2 nearest to N
        temp=np.zeros((2**(base2+1)-n1,))
        x=np.concatenate((x,temp))

    n = len(x)

    #....construct wavenumber array used in transform [Eqn(5)]
    k = np.arange(1,np.fix(n/2)+1)
    k = k*(2.*np.pi)/(n*dt)
    k = np.concatenate((np.zeros((1,)),k, -k[-2::-1]));

    #....compute FFT of the (padded) time series
    f = np.fft.fft(x)    # [Eqn(3)]

    #....construct SCALE array & empty PERIOD & WAVE arrays
    scale=np.array([s0*2**(i*dj) for i in range(0,int(J1)+1)])
    period = scale.copy()
    wave = np.zeros((int(J1)+1,n),dtype=np.complex)  # define the wavelet array  # make it complex
    # loop through all scales and compute transform
    for a1 in range(0,int(J1)+1):
        daughter,fourier_factor,coi,dofmin=wave_bases(mother,k,scale[a1],param)
        wave[a1,:] = np.fft.ifft(f*daughter)  # wavelet transform[Eqn(4)]
    period = fourier_factor*scale
    coi=coi*dt*np.concatenate(([1.E-5],np.arange(1.,(n1+1.)/2.-1),np.flipud(np.arange(1,n1/2.)),[1.E-5])) # COI [Sec.3g]
    wave = wave[:,:n1]  # get rid of padding before returning
    return wave,period,scale,coi
    # end of code

def wave_bases(mother,k,scale,param):
    """
    This is translation of wave_bases.m by Torrence and Gilbert P. Compo

    The folloing is the original README

%    WAVE_BASES  1D Wavelet functions Morlet, Paul, or DOG
%
%  [DAUGHTER,FOURIER_FACTOR,COI,DOFMIN] = ...
%      wave_bases(MOTHER,K,SCALE,PARAM);
%
%   Computes the wavelet function as a function of Fourier frequency,
%   used for the wavelet transform in Fourier space.
%   (This program is called automatically by WAVELET)
%
% INPUTS:
%
%    MOTHER = a string, equal to 'MORLET' or 'PAUL' or 'DOG'
%    K = a vector, the Fourier frequencies at which to calculate the wavelet
%    SCALE = a number, the wavelet scale
%    PARAM = the nondimensional parameter for the wavelet function
%
% OUTPUTS:
%
%    DAUGHTER = a vector, the wavelet function
%    FOURIER_FACTOR = the ratio of Fourier period to scale
%    COI = a number, the cone-of-influence size at the scale
%    DOFMIN = a number, degrees of freedom for each point in the wavelet power
%             (either 2 for Morlet and Paul, or 1 for the DOG)
%
%----------------------------------------------------------------------------
%   Copyright (C) 1995-1998, Christopher Torrence and Gilbert P. Compo
%   University of Colorado, Program in Atmospheric and Oceanic Sciences.
%   This software may be used, copied, or redistributed as long as it is not
%   sold and this copyright notice is reproduced on each copy made.  This
%   routine is provided as is without any express or implied warranties
%   whatsoever.
%----------------------------------------------------------------------------
    """
    #import modules
    import numpy as np

    #
    mother = mother.upper()
    n = len(k)
    # define Heaviside step function
    def ksign(x):
        y=np.zeros_like(x)
        y[x>0]=1
        return y
    #
    if mother=='MORLET':  #-----------------------------------  Morlet
        if (param == -1): param = 6.
        k0 = param
        expnt = -(scale*k - k0)**2/2. *ksign(k)
        norm = np.sqrt(scale*k[1])*(np.pi**(-0.25))*np.sqrt(n)    # total energy=N   [Eqn(7)]
        daughter = norm*np.exp(expnt)
        daughter = daughter*ksign(k)  # Heaviside step function
        fourier_factor = (4.*np.pi)/(k0 + np.sqrt(2. + k0**2)) # Scale-->Fourier [Sec.3h]
        coi = fourier_factor/np.sqrt(2)            # Cone-of-influence [Sec.3g]
        dofmin = 2.                          # Degrees of freedom
    elif mother=='PAUL': #--------------------------------  Paul
        if (param == -1): param = 4.
        m = param
        expnt = -(scale*k)*ksign(k)
        norm = np.sqrt(scale*k[1])*(2.**m/np.sqrt(m*np.prod(np.arange(2,2*m))))*np.sqrt(n)
        daughter = norm*((scale*k)**m)*np.exp(expnt)
        daughter = daughter*ksign(k)      # Heaviside step function
        fourier_factor = 4*np.pi/(2.*m+1.)
        coi = fourier_factor*np.sqrt(2)
        dofmin = 2.
    elif mother=='DOG':  #--------------------------------  DOG
        if (param == -1): param = 2.
        m = param
        expnt = -(scale*k)**2 / 2.0
        from scipy.special import gamma
        norm = np.sqrt(scale*k[1]/gamma(m+0.5))*np.sqrt(n)
        daughter = -norm*(1j**m)*((scale*k)**m)*np.exp(expnt);
        fourier_factor = 2.*np.pi*np.sqrt(2./(2.*m+1.))
        coi = fourier_factor/np.sqrt(2)
        dofmin = 1.
    else:
        raise Exception("Mother must be one of MORLET,PAUL,DOG")

    return daughter,fourier_factor,coi,dofmin

    # end of code

def wave_signif(Y,dt,scale1,sigtest=-1,lag1=-1,siglvl=-1,dof=-1,mother=-1,param=-1):
    """
This function is the translation of wave_signif.m by Torrence and Compo
use scipy function "chi2" instead of  chisquare_inv

The following is the original comment in wave_signif.m

%WAVE_SIGNIF  Significance testing for the 1D Wavelet transform WAVELET
%
%   [SIGNIF,FFT_THEOR] = ...
%      wave_signif(Y,DT,SCALE,SIGTEST,LAG1,SIGLVL,DOF,MOTHER,PARAM)
%
% INPUTS:
%
%    Y = the time series, or, the VARIANCE of the time series.
%        (If this is a single number, it is assumed to be the variance...)
%    DT = amount of time between each Y value, i.e. the sampling time.
%    SCALE = the vector of scale indices, from previous call to WAVELET.
%
%
% OUTPUTS:
%
%    SIGNIF = significance levels as a function of SCALE
%    FFT_THEOR = output theoretical red-noise spectrum as fn of PERIOD
%
%
% OPTIONAL INPUTS:
% *** Note *** setting any of the following to -1 will cause the default
%               value to be used.
%
%    SIGTEST = 0, 1, or 2.    If omitted, then assume 0.
%
%         If 0 (the default), then just do a regular chi-square test,
%             i.e. Eqn (18) from Torrence & Compo.
%         If 1, then do a "time-average" test, i.e. Eqn (23).
%             In this case, DOF should be set to NA, the number
%             of local wavelet spectra that were averaged together.
%             For the Global Wavelet Spectrum, this would be NA=N,
%             where N is the number of points in your time series.
%         If 2, then do a "scale-average" test, i.e. Eqns (25)-(28).
%             In this case, DOF should be set to a
%             two-element vector [S1,S2], which gives the scale
%             range that was averaged together.
%             e.g. if one scale-averaged scales between 2 and 8,
%             then DOF=[2,8].
%
%    LAG1 = LAG 1 Autocorrelation, used for SIGNIF levels. Default is 0.0
%
%    SIGLVL = significance level to use. Default is 0.95
%
%    DOF = degrees-of-freedom for signif test.
%         IF SIGTEST=0, then (automatically) DOF = 2 (or 1 for MOTHER='DOG')
%         IF SIGTEST=1, then DOF = NA, the number of times averaged together.
%         IF SIGTEST=2, then DOF = [S1,S2], the range of scales averaged.
%
%       Note: IF SIGTEST=1, then DOF can be a vector (same length as SCALEs),
%            in which case NA is assumed to vary with SCALE.
%            This allows one to average different numbers of times
%            together at different scales, or to take into account
%            things like the Cone of Influence.
%            See discussion following Eqn (23) in Torrence & Compo.
%
%
%----------------------------------------------------------------------------
%   Copyright (C) 1995-1998, Christopher Torrence and Gilbert P. Compo
%   University of Colorado, Program in Atmospheric and Oceanic Sciences.
%   This software may be used, copied, or redistributed as long as it is not
%   sold and this copyright notice is reproduced on each copy made.  This
%   routine is provided as is without any express or implied warranties
%   whatsoever.
%----------------------------------------------------------------------------
    """
    from scipy.stats import chi2
    import numpy as np

    try:
        n1=len(Y)
    except:
        n1=1
    J1 = len(scale1) - 1
    scale = scale1
    s0 = np.min(scale)
    dj = np.log(scale[1]/scale[0])/np.log(2.)

    if (n1 == 1):
        variance = Y
    else:
        variance = np.std(Y)**2

    if (sigtest == -1): sigtest = 0
    if (lag1 == -1): lag1 = 0.0
    if (siglvl == -1): siglvl = 0.95
    if (mother == -1): mother = 'MORLET'

    mother = mother.upper()

    # get the appropriate parameters [see Table(2)]
    if (mother=='MORLET'):  #----------------------------------  Morlet
        if (param == -1): param = 6.
        k0 = param
        fourier_factor = (4.*np.pi)/(k0 + np.sqrt(2. + k0**2)) # Scale-->Fourier [Sec.3h]
        empir = [2.,-1,-1,-1]
        if (k0 == 6): empir[1:4]=[0.776,2.32,0.60]
    elif (mother=='PAUL'):  #--------------------------------  Paul
        if (param == -1): param = 4.
        m = param
        fourier_factor = 4.*np.pi/(2.*m+1.)
        empir = [2.,-1,-1,-1]
        if (m == 4): empir[1:4]=[1.132,1.17,1.5]
    elif (mother=='DOG'):  #---------------------------------  DOG
        if (param == -1): param = 2.
        m = param
        fourier_factor = 2.*np.pi*np.sqrt(2./(2.*m+1.))
        empir = [1.,-1,-1,-1]
        if (m == 2): empir[1:4] = [3.541,1.43,1.4]
        if (m == 6): empir[1:4] = [1.966,1.37,0.97]
    else:
        raise Exception("Mother must be one of MORLET,PAUL,DOG")

    period = scale*fourier_factor
    dofmin = empir[0]     # Degrees of freedom with no smoothing
    Cdelta = empir[1]     # reconstruction factor
    gamma_fac = empir[2]  # time-decorrelation factor
    dj0 = empir[3]        # scale-decorrelation factor

    freq = dt / period   # normalized frequency
    fft_theor = (1.-lag1**2) / (1.-2.*lag1*np.cos(freq*2.*np.pi)+lag1**2)  # [Eqn(16)]
    fft_theor = variance*fft_theor  # include time-series variance
    signif = fft_theor
    try:
        test=len(dof)
    except:
        if (dof == -1):
            dof = dofmin
        else:
            pass
    #
    if (sigtest == 0):    # no smoothing, DOF=dofmin [Sec.4]
        dof = dofmin
        chisquare = chi2.ppf(siglvl,dof)/dof
        signif = fft_theor*chisquare   # [Eqn(18)]
    elif (sigtest == 1):  # time-averaged significance
        try:
            test=len(dof)
        except:
            dof=np.zeros((J1+1,))+dof
        truncate = dof < 1
        dof[truncate] = 1.
        dof = dofmin*np.sqrt(1. + (dof*dt/gamma_fac / scale)**2 )   # [Eqn(23)]
        truncate = dof < dofmin
        dof[truncate] = dofmin   # minimum DOF is dofmin
        for a1 in range(J1+1):
            chisquare = chi2.ppf(siglvl,dof[a1])/dof[a1]
            signif[a1] = fft_theor[a1]*chisquare
    elif (sigtest == 2):  # time-averaged significance
        if not (len(dof) == 2):
            raise Exception("DOF must be set to [S1,S2], the range of scale-averages'")
        if (Cdelta == -1):
            raise Exception('Cdelta & dj0 not defined for '+mother+' with param = '+str(param))
        s1 = dof[0];
        s2 = dof[1];
        avg = (scale >= s1) & (scale <= s2)  # scales between S1 & S2
        navg=np.sum(avg)
        if navg==0:
            raise Exception('No valid scales between '+str(s1)+' and '+str(s2))
        Savg = 1./np.sum(1 / scale[avg])    # [Eqn(25)]
        Smid = np.exp((np.log(s1)+np.log(s2))/2.)     # power-of-two midpoint
        dof = (dofmin*navg*Savg/Smid)*np.sqrt(1. + (navg*dj/dj0)**2)  # [Eqn(28)]
        fft_theor = Savg*np.sum(fft_theor[avg] / scale[avg])  # [Eqn(27)]#
        chisquare = chi2.ppf(siglvl,dof)/dof
        signif = (dj*dt/Cdelta/Savg)*fft_theor*chisquare    # [Eqn(26)]
    else:
        raise Exception('sigtest must be either 0, 1, or 2')

    return signif,fft_theor

# end of code

def wavelet_inverse(wave, scale, dt, dj=0.25, mother="MORLET",param=-1):
    """Inverse continuous wavelet transform
    Torrence and Compo (1998), eq. (11)

    INPUTS
        waves (array like):
          WAVE is the WAVELET transform. This is a complex array.

        scale (array like):
           the vector of scale indices
        dt (float) :
            amount of time between each original value, i.e. the sampling time.
        dj (float, optional) :
            the spacing between discrete scales. Default is 0.25.
           A smaller # will give better scale resolution, but be slower to plot.
        mother (string, optional) :
            the mother wavelet function.
             The choices are 'MORLET', 'PAUL', or 'DOG'
         PARAM = the mother wavelet parameter.
            For 'MORLET' this is k0 (wavenumber), default is 6.
            For 'PAUL' this is m (order), default is 4.
            For 'DOG' this is m (m-th derivative), default is 2.

    OUTPUTS
        iwave (array like) :
            Inverse wavelet transform.
    """
    import numpy as np

    j1, n = wave.shape
    J1 = len(scale)
    if not j1 == J1:
        print j1,n,J1
        raise Exception("Input array are inconsistent")
    sj = np.dot(scale.reshape(len(scale),1),np.ones((1,n)))
    #
    mother = mother.upper()

    # psi0 comes from Table 1,2 Torrence and Compo (1998)
    # Cdelta comes from Table 2 Torrence and Compo (1998)
    if mother=='MORLET':  #-----------------------------------  Morlet
        if (param == -1): param = 6.
        psi0=np.pi**(-0.25)
        if param==6.:
            Cdelta = 0.776
    elif mother=='PAUL': #--------------------------------  Paul
        if (param == -1): param = 4.
        m = param
        psi0=np.real(2.**m*1j**m*np.prod(np.arange(2, m + 1))/np.sqrt(np.pi*np.prod(np.arange(2,2*m+1)))*(1**(-(m+1))))
        if m==4.:
           Cdelta = 1.132
    elif mother=='DOG':  #--------------------------------  DOG
        if (param == -1): param = 2.
        m = param
        from scipy.special import gamma
        from numpy.lib.polynomial import polyval
        from scipy.special.orthogonal import hermitenorm
        p = hermitenorm(m)
        psi0=(-1)**(m+1)/np.sqrt(gamma(m+0.5))*polyval(p, 0)
        print psi0
        if m==2.:
            Cdelta=3.541
        if m==6.:
            Cdelta=1.966
    else:
        raise Exception("Mother must be one of MORLET,PAUL,DOG")

    #eq. (11) in Torrence and Compo (1998)
    iwave = dj * np.sqrt(dt) / Cdelta /psi0 * (np.real(wave) / sj**0.5).sum(axis=0)
    return iwave


#ref: Sreenivas Sarwar Anik

# Filters work by convolution with a moving window called a kernel.
#Convolution is nothing but multiplication of two arrays of different sizes.
#The image will be of one size and the kernel with be of a different size,
# #usually much smaller than image
# The input pixel is at the centre of the kernel.
# The convolution is performed by sliding the kernel over the image,
# $usually from top left of image.
# Linear filters and non-linear filters.
# Gaussian is an example of linear filter.
#Non-linear filters preserve edges.
#Median filter is an example of non-linear filter.
#The algorithm selects the median value of all the pixels in the selected window
#NLM: https://scikit-image.org/docs/dev/auto_examples/filters/plot_nonlocal_means.html

### What is gaussian Kernel #####
import numpy
from matplotlib import pyplot as plt

def gaussian_kernel(size, size_y=None):
    size = int(size)
    if not size_y:
        size_y = size
    else:
        size_y = int(size_y)
    x, y = numpy.mgrid[-size:size+1, -size_y:size_y+1]
    g = numpy.exp(-(x**2/float(size)+y**2/float(size_y)))
    return g / g.sum()

gaussian_kernel_array = gaussian_kernel(3)
print(gaussian_kernel_array)
plt.imshow(gaussian_kernel_array, cmap=plt.get_cmap('jet'), interpolation='nearest')
plt.colorbar()
plt.show()

############################ Denoising filters ###############
from skimage.restoration import denoise_nl_means, estimate_sigma
from skimage import img_as_ubyte, img_as_float
from matplotlib import pyplot as plt
from skimage import io
import numpy as np

img = img_as_float(io.imread("images/denoising/noisy_img.jpg"))
#Need to convert to float as we will be doing math on the array

from scipy import ndimage as nd
gaussian_img = nd.gaussian_filter(img, sigma=3)
plt.imsave("images/gaussian.jpg", gaussian_img)

median_img = nd.median_filter(img, size=3)
plt.imsave("images/median.jpg", median_img)

gaussian_img = nd.gaussian_filter(img, sigma=3)
plt.imsave("images/gaussian.jpg", gaussian_img)

##### NLM#####

sigma_est = np.mean(estimate_sigma(img, multichannel=True))

patch_kw = dict(patch_size=5,
                patch_distance=3,
                multichannel=True)

denoise_img = denoise_nl_means(img, h=1.15 * sigma_est, fast_mode=False,
                               patch_size=5, patch_distance=3, multichannel=True)
"""
denoise_img = denoise_nl_means(img, h=1.15 * sigma_est, fast_mode=False,
                           **patch_kw)
"""
denoise_img_as_8byte = img_as_ubyte(denoise_img)

plt.imshow(denoise_img)
#plt.imshow(denoise_img_as_8byte, cmap=plt.cm.gray, interpolation='nearest')
plt.imsave("images/NLM.jpg",denoise_img)


import numpy as np
import taylor as ta

def f(x):
    # vector function depending nonlinearly on input vector x
    if (x.shape[0] != 3):
        print(f'ERROR: Jacobian example requires vector of length 3')
    f    = np.empty(4)
    f[0] = x[0]*np.power(x[2],2)
    f[1] = np.sin(x[1]*x[2])
    f[2] = np.exp(np.pi*x[0]) + 10
    f[3] = np.sum(x)
    return f

def J_f_exact(x):
    # computes the exact Jacobian of f(x)
    J = np.empty((4,3))
    # set Jacobian row by row
    J[0,:] = np.array([np.power(x[2],2), 0.0, 2*x[0]*x[2]])
    J[1,:] = np.array([0.0, np.cos(x[1]*x[2])*x[2],
                       np.cos(x[1]*x[2])*x[1]])
    J[2,:] = np.array([np.exp(np.pi*x[0])*np.pi,0.0,0.0])
    J[3,:] = np.array([1.0,1.0,1.0])
    return J

if (__name__ == "__main__"):
    print()
    print(f'--------------------- JACOBIAN TEST ---------------------')
    print(f'    f(x) = [ x_0*(x_2)^2,')
    print(f'             sin(x_1*x_2),')
    print(f'             exp(pi*x_0) + 10,')
    print(f'             x_1 + x_2 + x_3 ]\n')

    # set current point x and compute exact Jacobian at point
    x       = np.array([1.5,20,2.718])
    J_exact = J_f_exact(x)

    # print x and exact Jacobian at x
    print(f'x : {x}\n')
    print(f'J_f(x) (exact) :\n{J_exact}\n')

    # compute Jacobian of nonlinear f with respect to x numerically
    deriv_nonlinear_vec = ta.diff(f,x,1)
    print(f'J_f(x) (approximate) :\n{deriv_nonlinear_vec}\n')


import numpy as np
import math

# -------------------------------------------------------------------------
'''
    Problem 0: Graph Clustering (Spectral Clustering)
    In this problem, you will implement a version of the spectral clustering method to cluster the nodes in a graph into two groups.
    You could test the correctness of your code by typing `nosetests -v test0.py` in the terminal.

    Notations:
            ---------- input data ------------------------
            n: the number of nodes in the network, an integer scalar.
            A:  the adjacency matrix, a float numpy matrix of shape n by n.
                If there is a link between node i an node j, then A[i][j] = A[j][i] = 1.
            ---------- computed data ------------------------
            D:  the degree matrix, a numpy float matrix of shape n by n.
                All off-diagonal elements are 0. Each diagonal element represents the degree of the node (number of links).
            L:  the Laplacian matrix, a numpy float matrix of shape n by n.
                L = D-A
            E:  the eigen vectors of matrix L, a numpy float matrix of shape n by n. Each column of E corresponds to an eigen vector of matrix L.
            v:  the eigen values of matrix L, a numpy float array of length n. Each element corresponds to an eigen value of matrix L. E and v are paired.
            e2:  the eigen vector corresponding to the smallest non-zero eigen value, a numpy float matrix of shape (n by 1).
            tol: the tolerance threshold for eigen values, a float scalar. A very small positive value. If an eigen value is smaller than tol, then we consider the eigen value as being 0.
            x:  the binary vector of shape (n by 1), a numpy float vector of (0/1) values.
                It indicates a binary partition on the graph, such as [1.,1.,1., 0.,0.,0.].
            --------------------------------------------------
'''

# --------------------------
def compute_D(A):
    '''
        Compute the degree matrix D.
        Input:
            A:  the adjacency matrix, a float numpy matrix of shape n by n. Here n is the number of nodes in the network.
                If there is a link between node i an node j, then A[i][j] = A[j][i] = 1.
        Output:
            D:  the degree matrix, a numpy float matrix of shape n by n.
                All off-diagonal elements are 0. Each diagonal element represents the degree of the node (number of links).
    '''

    #########################################
    ## INSERT YOUR CODE HERE
    D = np.zeros(A.shape)
    # degree vector of the nodes

    for i in range(A.shape[0]):
        for j in range(A.shape[1]):
            if A[i, j] == 1:
                D[i, i] += 1

    # diagonal matrix
    D = np.asmatrix(D)
    #########################################
    return D

# --------------------------
def compute_L(D, A):
    '''
        Compute the Laplacian matrix L.
        Input:
            D:  the degree matrix, a numpy float matrix of shape n by n.
                All off-diagonal elements are 0. Each diagonal element represents the degree of the node (number of links).
            A:  the adjacency matrix, a float numpy matrix of shape n by n. Here n is the number of nodes in the network.
                If there is a link between node i an node j, then A[i][j] = A[j][i] = 1.
        Output:
            L:  the Laplacian matrix, a numpy float matrix of shape n by n.
    '''

    #########################################
    ## INSERT YOUR CODE HERE
    L = D - A
    #########################################
    return L

# --------------------------
def compute_eigen_pairs(L):
    '''
        Compute the eigen vectors and eigen values of L.
        Input:
            L:  the Laplacian matrix, a numpy float matrix of shape n by n.
        Output:
            E:  the eigen vectors of matrix L, a numpy float matrix of shape p by p. Each column of E corresponds to an eigen vector of matrix L.
            v:  the eigen values of matrix L, a numpy float array of length p. Each element corresponds to an eigen value of matrix L. E and v are paired.
        Hint: you could use np.linalg.eig() to compute the eigen vectors of a matrix.
        Note: in the returned eigen vectors "eig_vecs", each column (not row) is considered as an eigen vector.
    '''

    #########################################
    ## INSERT YOUR CODE HERE
    v, E = np.linalg.eig(L)
    #########################################
    return E, v

# --------------------------
def find_e2(E, v, tol=1e-4):
    '''
        find the eigen vector that corresponds to the smallest non-zeros eigen values.
        For example, if the eigen values are [0.5, 0.000001, 2.], the first eigen vector is the answer.
        For example, if the eigen values are [2., 0.000001, 0.3], the last eigen vector is the answer.
        For example, if the eigen values are [0.00003, 0.000001, 0.3], the last eigen vector is the answer.
        Input:
            E:  the eigen vectors of matrix L, a numpy float matrix of shape n by n. Each column of E corresponds to an eigen vector of matrix L.
            v:  the eigen values of matrix L, a numpy float array of length n. Each element corresponds to an eigen value of matrix L. E and v are paired.
            tol: the tolerance threshold for eigen values, a float scalar. A very small positive value. If an eigen value is smaller than tol, then we consider the eigen value as being 0.
        Output:
            e2:  the eigen vector corresponding to the smallest non-zero eigen value, a numpy float matrix of shape (n by 1).
    '''

    #########################################
    ## INSERT YOUR CODE HERE
    i = 0
    e2 = np.asmatrix(np.zeros(E.shape[0]))
    # sort the eigen values
    idx = np.argsort(v)
    while v[idx[i]] < tol:
        i += 1

    e2 = E[:, idx[i]]

    #########################################
    return e2

# --------------------------
def compute_x(e2):
    '''
        Compute the partition on the graph from the thresholding an eigen vector with 0 threshold.
        Input:
            e2:  the eigen vector corresponding to the smallest non-zero eigen value, a numpy float matrix of shape (n by 1).
        Output:
            x:  the binary vector of shape (n by 1), a numpy float vector of (0/1) values.
                It indicates a binary partition on the graph, such as [1.,1.,1., 0.,0.,0.].
    '''

    #########################################
    ## INSERT YOUR CODE HERE
    x = e2
    threshold = 0
    x[x < threshold] = 0
    x[x > threshold] = 1
    #########################################
    return x

# --------------------------
def spectral_clustering(A):
    '''
        Spectral clustering of a graph.
        Input:
            A:  the adjacency matrix, a float numpy matrix of shape n by n. Here n is the number of nodes in the network.
                If there is a link between node i an node j, then A[i][j] = A[j][i] = 1.
        Output:
            x:  the binary vector of shape (n by 1), a numpy float vector of (0/1) values.
                It indicates a binary partition on the graph, such as [1.,1.,1., 0.,0.,0.].
        Note: you cannot use any existing python package for spectral clustering, such as scikit-learn.
        Hint: x is related to the eigen vector of L with the smallest positive eigen values.
              For example, if the eigen vector is [0.2,-0.1, -0.2], the values larger than zero will be 1, so x=[1,0,0] in this example.
              Note: if the eigen value is small enough (say smaller than 0.000001), we can treat it as being zero.
    (5 points)
    '''

    #########################################
    ## INSERT YOUR CODE HERE

    # compute degree matrix
    D = compute_D(A)
    # compute laplacian matrix
    L = compute_L(D, A)
    # compute eigen pairs of L
    E, v = compute_eigen_pairs(L)
    # find the eigen vector with the smallest non-zero eigen value
    e2 = find_e2(E, v, tol=1e-6)
    # compute the graph partition
    x = compute_x(e2)
    #########################################
    return x


"""
Lattice calcs

Authors/Modifications:
-----------------------
* Tom Trainor (tptrainor@alaska.edu)
* convert matlab to python, Kunal Tanwar
* Convert classes, TPT

Notes on Convention and Generalized Transforms
----------------------------------------------
The set of basis vectors that defines the real lattice are
covarient quantities, while the set of indicies that define
real space vectors are contravarient:
                 |a|
    v = (x,y,z)* |b| = x*a + y*b + z*c
                 |c|
If A is an arbitrary matrix that transforms the contravarient
vector indicies, its assumed that the indicies left multiply
    (x',y',z') = (x,y,z) * A

Note that the usual convention is to express covarient quantities
as column vectors and contravarient quantities as row vectors
(subscripts and superscripts respectivley in tensor notation)
However, in the notes below we'll just take 1-D arrays to be
column vectors to simplify the notation.  Therefore, if A is defined
as above, the matrix vector multiplications will need to be changed to:
    |x'|                 |x|
    |y'| = transpose(A)* |y|
    |z'|                 |z|
We can express the same in numpy given a vector v and matrix A
    dot(v,A) = dot(transpose(A),v)
where v is multiplied as a row vector on the lhs and a column
vector on the rhs; ie the transpose of v bewteen column and
row vecor is implied by its position.

Assume there is a matrix that transforms the covarient basis
(via rotation and/or dilation) to a new set of basis vectors:
    |a'|      |a|        a' = (F11*a + F12*b + F13*c)
    |b'| = F* |b|   or   b' = (F21*a + F22*b + F23*c)
    |c'|      |c|        c' = (F31*a + F32*b + F33*c)
with a well defined inverse:
    |a|           |a'|
    |b| = inv(F)* |b'|
    |c|           |c'|

The covarient/contravarient relationship bewteen basis vectors
and the vector indicies implies that the indicies of a (stationary)
vector in the new (') basis are obtained via the inverse of the
F matrix:
    |x'|      |x|
    |y'| = M* |y|
    |z'|      |z|
where
    M = transpose(inv(F)) = inv(transpose(F))

The reciprocal relation is:
    |x|      |x'|
    |y| = N* |y'|
    |z|      |z'|
where
    N = inv(M) = transpose(F)

With respect to the real lattice the recip lattice basis vectors
are contravarient and the vector indicies are covarient:
                     |h|
    vr = (ar,br,cr)* |k| = h*ar + k*br + l*cr
                     |l|

Given the generalized basis transform F above, the covarient recip
lattice indicies transform in the same way as the lattice.
Therefore:
     |h'|      |h|
     |k'| = F* |k|
     |l'|      |l|

The reciprocal relation is:
    |h|      |k'|
    |k| = G* |k'|
    |l|      |l'|
where
    G = inv(F) = transpose(M)

The metric tensor
------------------
The metric tensor is defined in terms of the vector dot product
for a general real space system
  dot(u,v) = u*g*v
where
       |a*a, a*b, a*c|
  g =  |b*a, b*b, b*c|
       |c*a, c*b, c*c|

Therefore, (with a,b,c, correspoding to vectors magnitudes):

g = [ [ a*a, a*b*cosd(gam), a*c*cosd(bet) ],
      [ b*a*cosd(gam), b*b, b*c*cosd(alp) ],
      [ c*a*cosd(bet), c*b*cosd(alp), c*c ] ]

The metric tensor may also be defined in terms of a lattice transform
ie the real space basis vectors can be expressed in terms of the
recip space basis vectors via
      |a|      |ar|
      |b| = g* |br|
      |c|      |cr|
therefore
      |ar|       |a|
      |br| = gr* |b|
      |cr|       |c|
where gr = inv(g).  Hence the recip lattice vectors
can be plotted in the real lattice via
      ar = gr11*a + gr12*b + gr13*c
      br = gr21*a + gr22*b + gr23*c
      cr = gr31*a + gr32*b + gr33*c
and visa versa

Using the above relations for lattice transforms the
indicies of real lattice vectors can be calc in the recip
lattice and visa-versa via:
   |h|                             |x|      |x|
   |k| = (x,y,z)*g = transpose(g)* |y| = M* |y|
   |l|                             |z|      |z|
and
   |x|                      |h|      |h|
   |y| = inv(transpose(g))* |k| = N* |k|
   |z|                      |l|      |l|

e.g. in above we take recip lattice as primed basis and
real lattice as the unprimed basis, therefore equate:
  F = gr
  G = inv(F) = inv(gr) = g
  M = transpose(inv(F)) = transpose(g)
  N = inv(M) = transpose(F) = inv(transpose(g)) = transpose(gr)

Note that g and gr are always symmetric, therefore
  g = transpose(g) and gr = transpose(gr).
  F = N = gr
  G = M = g

"""
##########################################################################

import numpy as num
from mathutil import cosd, sind, tand
from mathutil import arccosd, arcsind, arctand

##########################################################################
class Lattice:
    """
    Class that defines a lattice and various
    operations within the lattice

    """
    def __init__(self,a=10.,b=10.,c=10.,alpha=90.,beta=90.,gamma=90.,lam=1.5406):
        """
        Initialize

        Parameters:
        -----------
        * a,b,c in angstroms
        * alpha, beta, gamma in degrees,
        * lambda in angstroms (default lambda is Cu Ka1)
        """
        self.update(a=a,b=b,c=c,alpha=alpha,beta=beta,gamma=gamma,lam=lam)

    def __repr__(self,):
        """ display """
        lout = "a=%6.5f, b=%6.5f, c=%6.5f" % (self.a, self.b, self.c)
        lout = "%s, alpha=%6.5f,beta=%6.5f,gamma=%6.5f\n" % (lout,
                                                             self.alpha,
                                                             self.beta,
                                                             self.gamma)
        lout = "%sar=%6.5f, br=%6.5f, cr=%6.5f" % (lout,self.ar, self.br, self.cr)
        lout = "%s, alphar=%6.5f,betar=%6.5f,gammar=%6.5f\n" % (lout,
                                                                self.alphar,
                                                                self.betar,
                                                                self.gammar)
        lout = "%sDefault wavelength for angle calculations=%6.5f\n" % (lout,
                                                                      self.lam)
        return lout

    def update(self,a=None,b=None,c=None,alpha=None,beta=None,gamma=None,lam=None):
        """
        Update lattice parameters
        """
        if a != None:     self.a = float(a)
        if b != None:     self.b = float(b)
        if c != None:     self.c = float(c)
        if alpha != None: self.alpha = float(alpha)
        if beta  != None: self.beta  = float(beta)
        if gamma != None: self.gamma = float(gamma)
        if lam != None:   self.lam = float(lam)
        # update calc quantities
        self._calc_g()

    def cell(self):
        """
        Return array of real lattice cell parameters
        """
        return num.array([self.a,self.b,self.c,
                          self.alpha,self.beta,self.gamma],dtype=float)

    def rcell(self):
        """
        Return array of reciprocal lattice cell parameters
        """
        return num.array([self.ar,self.br,self.cr,
                          self.alphar,self.betar,self.gammar],dtype=float)

    def _calc_g(self):
        """
        Calculate the metric tensors and recip lattice params
        self.g  = real space metric tensor
        self.gr = recip space metric tensor
        """
        (a,b,c,alp,bet,gam) = self.cell()
        # real metric tensor
        self.g = num.array([ [ a*a, a*b*cosd(gam), a*c*cosd(bet) ],
                             [ b*a*cosd(gam), b*b, b*c*cosd(alp) ],
                             [ c*a*cosd(bet), c*b*cosd(alp), c*c ] ])
        # recip lattice metric tensor
        # and recip lattice params
        self.gr     = num.linalg.inv(self.g)
        self.ar     = num.sqrt(self.gr[0,0])
        self.br     = num.sqrt(self.gr[1,1])
        self.cr     = num.sqrt(self.gr[2,2])
        self.alphar = arccosd(self.gr[1,2]/(self.br*self.cr))
        self.betar  = arccosd(self.gr[0,2]/(self.ar*self.cr))
        self.gammar = arccosd(self.gr[0,1]/(self.ar*self.br))

    def vol(self,recip=False):
        """
        Calculate the cell volume.

        Parameters:
        -----------
        * If recip == False, this is the real space vol in ang**3
          If recip == True, this is the recip space vol in ang**-3

        Notes:
        ------
        V_real  = sqrt(determiant(g))
        V_recip = sqrt(determiant(inv(g)))

        """
        if recip == True: g = self.gr
        else: g = self.g
        det = num.linalg.det(g)
        if det > 0:
            return num.sqrt(det)
        else:
            return 0.0

    def dot(self,u,v,recip=False):
        """
        Calculate dot product of two vectors (u,v)

        Notes:
        ------
        * If u and v are real space vectors, use recip = False
        * If u and v are recip space vectors, use recip = True

        Note u and v are assumed to be normal numpy arrays
        (ie not matrix objects)

        The generalized dot product for real space vectors
        may be defined in terms of the metric tensor (g):
           dot(u,v) = u*g*v
        If u and v are recip lattice vectors replace g with
        gr = inv(g)
        """
        if recip == True: g = self.gr
        else: g = self.g
        u = num.array(u,dtype=float)
        v = num.array(v,dtype=float)
        return num.dot(u,num.dot(g,v))

    def mag(self,v,recip=False):
        """
        Calculate the norm of a vector (v)

        Notes:
        ------
        * If v is real space vector, use recip = False
        * If v is recip space vector, use recip = True

        Note v is assumed to be normal numpy array
        (ie not matrix objects)
        """
        m = num.sqrt(self.dot(v,v,recip=recip))
        return m

    def angle(self,u,v,recip=False):
        """
        Calculate the angle between two vectors (u,v)

        Notes:
        ------
        * If u and v are real space vectors, use recip = False
        * If u and v are recip space vectors, use recip = True

        Note u and v are assumed to be normal numpy arrays
        (ie not matrix objects)
        """
        uv = self.dot(u,v,recip=recip)
        um = self.mag(u,recip=recip)
        vm = self.mag(v,recip=recip)
        arg = uv/(um*vm)
        if num.fabs(arg) > 1.0:
            arg = arg / num.fabs(arg)
        alpha = arccosd(arg)
        return alpha

    def angle_rr(self,x,h):
        """
        Calculate the angle between a real vector x = [x,y,z]
        and a recip vector h = [h,k,l]
        """
        x = num.array(x,dtype=float)
        h = num.array(h,dtype=float)
        hx = num.sum(x*h)
        xm = self.mag(x,recip=False)
        hm = self.mag(h,recip=True)
        arg = hx/(hm*xm)
        if num.fabs(arg) > 1.0:
            arg = arg / num.fabs(arg)
        alpha = arccosd(arg)
        return alpha

    def d(self,hkl):
        """
        Calculate d space for given [h,k,l]
        """
        if len(hkl)!=3:
            print "Error need an array of [h,k,l]"
            return 0.
        h = self.mag(hkl,recip=True)
        if h == 0.:
            #print "Error [h,k,l] magnitude is zero:", hkl
            return 0.0
        d = 1./h
        return d

    def tth(self,hkl,lam=None):
        """
        Calculate 2Theta for given [h,k,l] and wavelength

        Notes:
        ------
        If lam is None the default lambda will be used (Cu Ka1)
        If you pass in lam, this will change the default for
        subsequent calls
        """
        if lam != None: self.lam = float(lam)
        d = self.d(hkl)
        if d == 0.: return 0.
        r = self.lam/(2.*d)
        if num.fabs(r) > 1.0:
            r = r/num.fabs(r)
        tth = 2.*arcsind(r)
        return tth

    def dvec(self,hkl):
        """
        Calculate the real space vector d
        which has a magnitude of the d spacing
        and is normal to the plane hkl = [h,k,l]
        """
        # convert hkl vector to real space indicies
        #hkl  = num.array(hkl,dtype=float)
        #dvec = num.dot(self.gr,hkl)
        dvec = self.recip_to_real(hkl)
        dspc = self.d(hkl)
        if dspc == 0: return num.array([0.,0.,0.])
        dvec = (dspc**2.)*dvec
        return dvec

    def recip_to_real(self,hkl):
        """
        Given recip vector hkl = [h,k,l] calculate the
        vectors indicies in the real lattice
        """
        hkl = num.array(hkl,dtype=float)
        #v   = num.dot(self.gr.transpose(),hkl)
        v   = num.dot(self.gr,hkl)
        return v

    def real_to_recip(self,v):
        """
        Given real vector v = [x,y,z] calculate the
        vectors indicies in the reciprocal lattice
        """
        v   = num.array(v,dtype=float)
        hkl = num.dot(v,self.g)
        return hkl

##########################################################################
class LatticeTransform:
    """
    Generalized lattice transformations
    """
    def __init__(self,lattice,Va=None,Vb=None,Vc=None,shift=None):
        """
        Initialize

        Parameters:
        ----------
        * lattice is a Lattice instance.
          All transforms are defined with respect to this
          lattice, ie this is the original unprimed lattice
        * Va, Vb, Vc are the lattice vectors that define the
          a',b',c' coordinate system (rotation and contraction)
        * shift is the lattice shift vector
        * (V's and shift are defined in the unprimed system)
        """
        self.lattice = lattice
        self.Va    = num.array([1.,0.,0.])
        self.Vb    = num.array([0.,1.,0.])
        self.Vc    = num.array([0.,0.,1.])
        self.shift = num.array([0.,0.,0.])
        self._update(Va=Va,Vb=Vb,Vc=Vc,shift=shift)

    def basis(self,Va=None,Vb=None,Vc=None,shift=None):
        """
        Define new basis vectors.

        Parameters:
        ----------
        * Va, Vb and Vc should define the a',b',c' lattice
          vectors of the new basis (rotation/dialation part).

        * Shift describes an origin shift of the new lattice
          (ie take the new basis defined by Va,Vb,Vc then
          apply the shift vector)

        Notes:
        ------
        All the vectors should be defined in terms of fractional
        coordinate indicies in the original basis, for example:
           a' = x1*a + y1*b + z1*c
           b' = x2*a + y2*b + z2*c
           c' = x3*a + y3*b + z3*c
           shift = x4*a + y4*b + z4*c
        and
           Va = [x1,y1,y1]
           Vb = [x2,y2,y2]
           Vc = [x3,y3,y3]
           shift = [x4,y4,z4]
        """
        self._update(Va=Va,Vb=Vb,Vc=Vc,shift=shift)

    def _update(self,Va=None,Vb=None,Vc=None,shift=None):
        """ update """
        if Va != None: self.Va = num.array(Va,dtype=float)
        if Vb != None: self.Vb = num.array(Vb,dtype=float)
        if Vc != None: self.Vc = num.array(Vc,dtype=float)
        if shift != None: self.shift = num.array(shift,dtype=float)
        F = [self.Va,self.Vb,self.Vc]
        self.F = num.array(F, dtype=float)
        self.G = num.linalg.inv(self.F)
        self.M = self.G.transpose()
        self.N = self.F.transpose()

    def cartesian(self,shift=[0.,0.,0.]):
        """
        Calculates a cartesian basis using:
          Va = a' is parallel to a
          Vb = b' is perpendicular to a' and in the a/b plane
          Vc = c' is perpendicular to the a'/c' plane
        A shift vector may be specified to shift the origin
        of the cartesian lattice relative to the original lattice
        origin (specify shift in fractional coordinates of
        the original lattice)
        """
        (a,b,c,alp,bet,gam) = self.lattice.cell()
        (ar,br,cr,alpr,betr,gamr) = self.lattice.rcell()
        Va = [1./a,                 0. ,             0.]
        Vb = [-1./(a*tand(gam)), 1./(b*sind(gam)),   0.]
        Vc = [ar*cosd(betr),     br*cosd(alpr),      cr]
        self.basis(Va=Va,Vb=Vb,Vc=Vc,shift=shift)

    def xp(self,x):
        """
        Given x = [x,y,z] in the original lattice
        compute the indicies of the vector in the new basis
           xp = M*(x - shift)
        """
        x  = num.array(x,dtype=float)
        if self.shift.sum() != 0.:
            x = x - self.shift
        xp = num.dot(self.M,x)
        return xp

    def x(self,xp):
        """
        Given xp = [x',y',z'] in the primed basis
        compute the indicies of the vector in the original basis
            x = N*(xp + shiftp) = N*xp + shift
        """
        xp  = num.array(xp,dtype=float)
        x = num.dot(self.N,xp)
        if self.shift.sum() != 0.:
            x = x + self.shift
        return x

    def hp(self,h):
        """
        Given h = [h,k,l] in the original (recip) basis
        compute the indicies of the vector in the primed
        (recip) basis
        """
        h  = num.array(h,dtype=float)
        return num.dot(self.F,h)

    def h(self,hp):
        """
        Given hp = [h',k',l'] in the primed (recip) basis
        compute the indicies of the vector in the original
        (recip) basis
        """
        hp  = num.array(hp,dtype=float)
        return num.dot(self.G,hp)

    def plat(self):
        """
        Return a Lattice instance for the primed basis
        """
        a = self.lattice.mag(self.Va)
        b = self.lattice.mag(self.Vb)
        c = self.lattice.mag(self.Vc)
        alp = self.lattice.angle(self.Vb,self.Vc)
        bet = self.lattice.angle(self.Va,self.Vc)
        gam = self.lattice.angle(self.Va,self.Vb)
        return Lattice(a,b,c,alp,bet,gam)

##########################################################################
##########################################################################
def test_lattice():
    # create a new lattice instance
    cell = Lattice(5.6,5.6,13,90,90,120)

    # show lattice parameters and real
    # recip lattice cell volumes
    print cell
    print cell.vol()
    print cell.vol(recip=True)

    # what the hkl = [0,0,1] dspace
    # and 2 theta value for lambda = 1 ang.
    print cell.d([0,0,1])
    print cell.tth([0,0,1],lam=1.)

    # create a real space vector perpendicular
    # to hkl = [0,0,1], with magnitude = dspace(001)
    dv = cell.dvec([0,0,1])
    print dv
    print cell.mag(dv)

    # compute the angle between some vectors
    # recip lattice [1,0,0] and [0,1,0]
    print cell.angle([1,0,0],[0,1,0],recip=True)
    # recip lattice [1,0,0] and [0,0,1]
    print cell.angle([1,0,0],[0,0,1],recip=True)
    # recip lattice [0,0,1] and [1,1,1]
    print cell.angle([0,0,1],[1,1,1],recip=True)

    # real lattice [1,0,0] and [0,1,0]
    print cell.angle([1,0,0],[0,1,0])
    # real lattice [1,0,0] and [0,0,1]
    print cell.angle([1,0,0],[0,0,1])
    # real lattice [0,0,1] and [1,1,1]
    print cell.angle([0,0,1],[1,1,1])

    # real [0,0,1] and recip [0,0,1]
    print cell.angle_rr([0,0,1],[0,0,1])
    # real [1,1,1] and recip [0,0,1]
    print cell.angle_rr([1,1,1],[0,0,1])
    # real dv and recip [0,0,1]
    print cell.angle_rr(dv,[0,0,1])
    # real dv and recip [1,1,0]
    print cell.angle_rr(dv,[1,1,0])

    return cell

def test_transform():
    # create a new hexagonal lattice instance
    cell = Lattice(5.6,5.6,13,90,90,120)

    # test lattice transform
    # below defines rhombohedral basis vectors
    # in terms of the heaxagonal basis vectors
    Va_rhom = [0.6667, 0.3333, 0.3333]
    Vb_rhom = [-0.3333, 0.3333, 0.3333]
    Vc_rhom = [-0.3333, -0.6667, 0.3333]
    t = LatticeTransform(cell,Va=Va_rhom,Vb=Vb_rhom,Vc=Vc_rhom)

    # create an instance of the rhombohedral lattice
    # shows the rhomb lattice params
    print t.plat()

    # given the hkl = [001] in hex calc hkl in rhom
    print t.hp([0,0,1])

    # given the hkl = [111] in rhom calc hkl in hex
    print t.h([1,1,1])

    # create a cartesian representation
    t.cartesian()
    cart = t.plat()
    print cart

    # convert a [1,1,0] vector in hex lattice to cartesian
    vc = t.xp([1,1,0])
    print vc

    # check that the vector is the same length
    print cell.mag([1,1,0])
    print cart.mag(vc)

    return t

##########################################################################
##########################################################################
if __name__ == "__main__":
    """
    test
    """
    test_lattice()
    test_transform()


import numpy as np

def congruent_linear_generator(a : int, m : int, c : int, x0 : int):
    """
    :def: Linear Congruent Generator Technique
    : Method to generate pseudo-randoms
    :
    :param a:  Random Parameter
    :type a:   int
    :param c:  Constant used to create bigger variation of generated numbers
    :type c:   int
    :param m:  Parameter which determines the biggest value generated
    :type m:   int
    :param x0: Determines initial seed and stop trigger
    :type x0:  int
    :return:   Pseudo-random list of numbers
    :rtype:    list[int]
    """
    return (a*x0 + c)%m

def give_me_n_random(RANGE:int, SEED:int, a:int=44, m:int=95784, c:int=4565):
    pseudo_random = [SEED]
    for _ in range(RANGE):
        pseudo_random.append(congruent_linear_generator(a, m, c, pseudo_random[-1]))

    MIN = min(pseudo_random)
    MAX = max(pseudo_random)

    def percentage(MIN, MAX, value):
        return round(((value - MIN)/(MAX - MIN)),2)

    return [percentage(MIN, MAX, x) for x in pseudo_random][1:]


import numpy as np
from nptyping import Array

def lms(d: Array, x: Array, w: Array, mu: float, _) -> Array:
    """
    lms Least-Mean-Square filtering.

    Args:
        d (Array): Desired values.
        x (Array): Input(recorded) values.
        w (Array): Filter-weight matrix.
        mu (float): Step size.
        _ ([dummy argument]): DO NOT SET ANY VALUE.

    Returns:
        Array: Modified value array of w.
    """
    N = len(x)

    y = np.zeros_like(d)
    e = np.zeros_like(d)

    # w_delta: 係数行列wの更新量
    w_delta = np.zeros_like(w)

    for k in range(N):
        y[k] = np.dot(w.T, x[k])
        e[k] = d[k] - y[k]
        # 2 * J = e ** 2
        diff_J_for_w_ = e[k] * x[k]
        w_delta = mu * diff_J_for_w_
    return w_delta

def nlms(d: Array, x: Array, w: Array, mu: float, lambda_: float = 1.0) -> Array:
    """
    Normalized-LMS filtering.

    Args:
        d (Array): Desired values.
        x (Array): Input(recorded) values.
        w (Array): Filter-weight matrix.
        mu (float): Step size.
        lambda_ (float, optional): Learning rate. Defaults to 1.0.

    Returns:
        Array: Modified value array of w.
    """

    N = len(x)

    y = np.zeros_like(d)
    e = np.zeros_like(d)

    # w_delta: 係数行列wの更新量
    w_delta = np.zeros_like(w)

    for k in range(N):
        y[k] = np.dot(w.T, x[k])
        e[k] = d[k] - y[k]
        # 2 * J = e ** 2 - sum(w[i] ** 2 for i in range(N))
        diff_J_for_w_ = e[k] * x[k]
        w_delta = mu / (np.dot(x[k], x[k]) + lambda_) * diff_J_for_w_

    return w_delta


### Example script showing how to perform a Total-Variation filtering with proxTV
import prox_tv as ptv
import numpy as np
import matplotlib.pyplot as plt
import time

def _blockysignal():
    """Generates a blocky signal for the demo"""
    N = 1000
    s = np.zeros((N,1))
    s[int(N/4):int(N/2)] = 1
    s[int(N/2):int(3*N/4)] = -1
    s[int(3*N/4):int(-N/8)] = 2
    return s

### TV-L1 filtering

# Generate impulse (blocky) signal
s = _blockysignal()

# Introduce noise
n = s + 0.5*np.random.rand(*np.shape(s))

# Filter using TV-L1
lam=20
print('Filtering signal with TV-L1...')
start = time.time()
f = ptv.tv1_1d(n,lam)
end = time.time()
print('Elapsed time ' + str(end-start))

# Plot results
plt.subplot(3, 1, 1)
plt.title('TVL1 filtering')
plt.plot(s)
plt.ylabel('Original')
plt.grid(True)

plt.subplot(3, 1, 2)
plt.plot(n)
plt.ylabel('Noisy')
plt.grid(True)

plt.subplot(3, 1, 3)
plt.plot(f)
plt.ylabel('Filtered')
plt.grid(True)

plt.show()

### TV-L2 filtering

# Generate sinusoidal signal
N = 1000
s = np.sin(np.arange(1,N+1)/10.0) + np.sin(np.arange(1,N+1)/100.0)

# Introduce noise
n = s + 0.5*np.random.randn(*np.shape(s))

# Filter using TV-L2
lam=100;
print('Filtering signal with TV-L2...')
start = time.time()
f = ptv.tv2_1d(n,lam);
end = time.time()
print('Elapsed time ' + str(end-start))

# Plot results
plt.subplot(3, 1, 1)
plt.title('TVL2 filtering')
plt.plot(s)
plt.ylabel('Original')
plt.grid(True)

plt.subplot(3, 1, 2)
plt.plot(n)
plt.ylabel('Noisy')
plt.grid(True)

plt.subplot(3, 1, 3)
plt.plot(f)
plt.ylabel('Filtered')
plt.grid(True)

plt.show()

### Weighted TV-L1 filtering

# Generate impulse (blocky) signal
s = _blockysignal()

# Introduce noise
n = s + 0.5*np.random.randn(*np.shape(s))

# Generate weights
lam = np.linspace(0,2,N-1)

# Filter using weighted TV-L1
print('Filtering signal with weighted TV-L1...')
start = time.time()
f = ptv.tv1w_1d(n, lam)
end = time.time()
print('Elapsed time ' + str(end-start))

# Plot results
plt.subplot(4, 1, 1)
plt.title('Weighted TVL1 filtering')
plt.plot(s)
plt.ylabel('Original')
plt.grid(True)

plt.subplot(4, 1, 2)
plt.plot(n)
plt.ylabel('Noisy')
plt.grid(True)

plt.subplot(4, 1, 3)
plt.plot(f)
plt.ylabel('Filtered')
plt.grid(True)

plt.subplot(4, 1, 4)
plt.fill_between(np.arange(1,N), 0, lam)
plt.ylabel('Weights')
plt.grid(True)

plt.show()


import numpy as np

operator = 'and'

#----------------

atributes = np.array([ [0, 0], [0, 1], [1, 0], [1, 1]])

if operator == 'and':
	labels = np.array([0, 0, 0, 1])
elif operator == 'or':
	labels = np.array([0, 1, 1, 1])
elif operator == 'xor':
	labels = np.array([0, 1, 1, 0])

#----------------

w = [+9, +9] #initial random values for weights

threshold = 5
alpha = 0.5 #learning rate
epoch = 1000 #learning time
#----------------

print("learning rate: ", alpha,", threshold: ", threshold)

for i in range(0, epoch):
	print("epoch ", i+1)
	global_delta = 0 #this variable is used to terminate the for loop if learning completed in early epoch
	for j in range(len(atributes)):

		actual = labels[j]

		sum = atributes[j][0]*w[0] + atributes[j][1]*w[1]

		if sum > threshold: #then fire
			predicted = 1
		else: #do not fire
			predicted = 0

		delta = actual - predicted
		global_delta = global_delta + abs(delta)

		#update weights with respect to the error
		for k in range(0, 2):
			w[k] = w[k] + delta * alpha

		print(atributes[j][0]," ", operator, " ", atributes[j][1], " -> actual: ", actual, ", predicted: ", predicted, " (w: ",w[0],")")

	if global_delta == 0:
		break

	print("------------------------------")


"""
An implementation of the cheeger cut for networkx graphs.
"""
import numpy as np
import scipy as sp
import scipy.sparse
import scipy.sparse.linalg
import networkx as nx

def sweep_set(A, v_2, degrees):
    """
    Given the adjacency matrix of a graph, and the second eigenvalue of the laplacian matrix, use the sweep set
    algorithm to find a sparse cut.

    :param A: The adjacency matrix of the graph to use.
    :param v_2: The second eigenvector of the laplacian matrix of the graph
    :param degrees: a list with the degrees of each vertex in the graph
    :return: The set of vertices corresponding to the optimal cut
    """

    # Calculate n here once
    n = A.shape[0]
    # Keep track of the best cut so far
    best_cut_index = None
    best_conductance = None

    # Keep track of the size of the set and the cut weight to make computing the conductance
    # straightforward
    total_volume = np.sum(degrees)
    set_volume = 0
    set_size = 0
    cut_weight = 0

    # Normalise v_2 with the degrees of each vertex
    D = sp.sparse.diags(degrees, 0)
    v_2 = D.power(-(1/2)).dot(v_2)

    # First, sort the vertices based on their value in the second eigenvector
    sorted_vertices = [i for i, v in sorted(enumerate(v_2), key=(lambda x: x[1]))]

    # Keep track of which edges to add/subtract from the cut each time
    x = np.ones(n)

    # Loop through the vertices in the graph
    for (i, v) in enumerate(sorted_vertices[:-1]):
        # Update the set size and cut weight
        set_volume += degrees[v]
        set_size += 1

        # From now on, edges to this vertex will be removed from the cut at each iteration.
        x[v] = -1

        additional_weight = A[v, :].dot(x)
        cut_weight += additional_weight

        if cut_weight < 0:
            raise Exception('Something went wrong in sweep set: conducatance negative!')
        # Calculate the conductance
        if min(set_volume, total_volume - set_volume) == 0:
            this_conductance = 1
        else:
            this_conductance = cut_weight / min(set_volume, total_volume - set_volume)

        # Check whether this conductance is the best
        if best_conductance is None or this_conductance < best_conductance:
            best_cut_index = i
            best_conductance = this_conductance

    # Return the best cut
    return sorted_vertices[:best_cut_index+1]

def cheeger_cut(G):
    """
    Given a networkx graph G, find the cheeger cut.

    :param G: The graph on which to operate
    :return: A set containing the vertices on one side of the cheeger cut
    """
    if G.number_of_nodes() == 0:
        raise Exception(f'Cheeger cut: Graph should not be empty!')

    if nx.is_connected(G) is False:
        return list(list(nx.connected_components(G))[0])

    if G.number_of_nodes() == 1:
        return list(G.nodes())

    if G.number_of_nodes() == 2:
        return [list(G.nodes)[0]]

    # Compute the key graph matrices
    adjacency_matrix = nx.adjacency_matrix(G, weight='weight')
    laplacian_matrix = nx.normalized_laplacian_matrix(G, weight='weight')
    graph_degrees = [t[1] for t in nx.degree(G, weight='weight')]

    # Compute the second smallest eigenvalue of the laplacian matrix
    eig_vals, eig_vecs = sp.sparse.linalg.eigsh(laplacian_matrix, which="SM", k=2)
    v_2 = eig_vecs[:, 1]

    # Perform the sweep set operation to find the sparsest cut
    S = sweep_set(adjacency_matrix, v_2, graph_degrees)
    nodes = list(G.nodes())
    return [nodes[i] for i in S]

import numpy,import scipy,import networkx


#!/usr/bin/env python

import numpy as np
from tqdm import tqdm

# we have only one neuron
NEURONS = 1

# learning rate
LEARNING_RATE = 0.01

# number of iterations
NUM_OF_ITERS = 100000

# small difference to be able to approximately calculcate derivative
EPSILON = 0.01

class TwoVariablesBooleanFunctionNeuralNetwork:
    def __init__(self, x, y):
        #  - n rows: the number of samples
        self.n = x.shape[0]

        #  - m columns: the number of input variables
        self.m = x.shape[1]

        # x is a np.array which is a n x m matrix. It has:
        #  - n rows: the number of samples
        #  - m columns: the number of input variables
        self.input = x

        # np.shape returns a tuple: (n,m)
        # W is a matrix of m x l:
        #  - m rows: the number of input variables
        #  - l columns: the number of neurons
        self.W = np.random.rand(self.m, NEURONS)

        # y and output are column vectors of size n
        self.y      = y
        self.output = np.zeros(self.y.shape)

    def feedforward(self):
        self.output = sigmoid(np.dot(self.input, self.W))

    def J(self, W):
        # here the np.dot multiplication will give back
        # a column vector, which we reshape (so that it will
        # be transposed to a row vector), and then squeeze (to
        # get rid of the extra dimension)
        g_xw = np.array(sigmoid(np.dot(self.input, W))).reshape(1,-1).squeeze()

        # calculate the difference between our hypothesis and
        # the desired values
        err_diff = self.y - g_xw

        # return the error as sum of squared difference
        return (1 / self.n) * np.sum(np.square(err_diff))

    def W_i_plus(self, i):
        W_copy = self.W.copy()
        W_copy[i][0] = W_copy[i][0] + EPSILON
        return W_copy

    def W_i_minus(self, i):
        W_copy = self.W.copy()
        W_copy[i][0] = W_copy[i][0] - EPSILON
        return W_copy

    def delta_J_i(self, i):
        return (
            (self.J(self.W_i_plus(i)) - self.J(self.W_i_minus(i))) / (2*EPSILON)
        )

    def backprop(self):
        delta_J = []

        for i in range(0,self.m):
            delta_J += [self.delta_J_i(i)]

        delta_J = np.array(delta_J).reshape(-1,1)

        self.W -= LEARNING_RATE * delta_J

def sigmoid(x):
    return 1/(1+np.exp(-x))

if __name__ == "__main__":
    x = np.array([[1,0,0],
                  [1,0,1],
                  [1,1,0],
                  [1,1,1]])

    y = np.array([0,0,0,1])

    nn = TwoVariablesBooleanFunctionNeuralNetwork(x,y)

    pbar = tqdm(total = NUM_OF_ITERS)

    for i in range(0, NUM_OF_ITERS):
        nn.feedforward()
        nn.backprop()
        pbar.update(1)

    print("Number of iterations: {}".format(NUM_OF_ITERS))
    print("Weights (W0, W1, ... Wm):\n{}".format(nn.W))
    print("Output (h_W(x)_0, h_W(x)_1, ... h_W(x)_n):\n{}".format(nn.output))


import numpy as np
import math
from scipy.interpolate import BSpline

def gen_BSplines(n, size_step, deg, shift=0):
    """
    Generates Bsplines on [1,n] where the knots are spaces by size_step and the polynomials are degree deg. Results are stored in a matrix D. Support of atoms has size (1+deg)*size_step. By default the grid for knots is [0,size_step, 2*size_step...] but it can be shifted.

    Parameters
    ----------
    n : int
        The splines are valued on [1,n]

    size-step : int
        The number of indices in [1,n] between each knot

    deg : int
        degree of the splines

    shift : int (default 0)
        Set a shift for the origin of the splines

    Returns
    -------
    D : numpy array
        A matrix containing the B-splines columnwise
    """

    n_over_m = size_step
    m = int(n/n_over_m + 2*(deg-1)+1) # adding left and right padding
    t = shift+np.linspace(-(deg-1)*n_over_m,n+(deg-1)*n_over_m,m)
    mySplines = BSpline(t,np.ones(n),deg)

    # todo: start before 0 to have nontrivial elements on the borders
    D = []
    for i in range(0,m-deg):
        Bspline = mySplines.basis_element(t[i:i+deg], extrapolate=False)
        D.append(Bspline(np.linspace(1,n,n)))
    D = np.array(D).T
    D[np.isnan(D)]=0
    return D

def sam(x,y):
    """
    Computes the Spectral Angular Mapper between two spectra x and y.

    Parameters
    ----------
    x : numpy array
        a spectra to compare with y

    y : numpy array
        a spectra to compare with x

    Returns
    -------
    out : float
        the SAM between x and y
    """

    return math.acos(np.dot(x,y)/np.linalg.norm(x)/np.linalg.norm(y))


## @ingroup Methods-Aerodynamics-Common-Gas_Dynamics
# oblique_shock.py
#
# Created:  Jan 2018, W. Maier
# Modified:
#

# ----------------------------------------------------------------------
#  Imports
# ----------------------------------------------------------------------

import numpy as np

# ----------------------------------------------------------------------
#  Oblique Shock Relations
# ----------------------------------------------------------------------

## @ingroup Methods-Aerodynamics-Common-Gas_Dynamics
def oblique_shock_relations(M0,gamma,theta,beta):
    """Computes flow quatities/ratios after undergoing
    an oblique shock

    Assumptions:
    None

    Source:
    Chapter 12 of:
    https://web.stanford.edu/~cantwell/AA210A_Course_Material/AA210A_Course_Notes/

    https://arc.aiaa.org/doi/pdf/10.2514/2.2349

    Inputs:
    Mach, M0                             [-]
    Isentropic Expansion Factor, gamma   [-]
    Flow turn angle, theta               [rad]
    Shock Angle, beta                    [rad]

    Outputs:
    Mach, M1                             [-]
    Static Temperature Ratio, Tr         [-]
    Static Pressure Ratio, Pr            [-]
    Stagnation Pressure Ratio, Ptr       [-]

    Properties Used:
    N/A
    """

    # determine normal component of machs
    M0_n = M0*np.sin(beta)
    M1_n = np.sqrt(((gamma-1.)*M0_n*M0_n+2.)/(2.*gamma*M0_n*M0_n-(gamma-1.)))

    # determine flow quantaties and ratios
    M1   = M1_n/np.sin(beta-theta)
    Pr   = (2.*gamma*M0_n*M0_n-(gamma-1.))/(gamma+1.)
    Tr   = Pr*((((gamma-1.)*M0_n*M0_n)+2.)/((gamma+1.)*M0_n*M0_n))
    Ptr  = ((((gamma+1.)*M0_n*M0_n)/((gamma-1.)*M0_n*M0_n+2))**(gamma/(gamma-1.)))*((gamma+1.)/(2.*gamma*M0_n*M0_n-(gamma-1.)))**(1./(gamma-1.))

    return M1,Pr,Tr,Ptr

def theta_beta_mach(M0,gamma,theta,n=0):
    """Computes shock angle of an oblique shock

        Assumptions:
        None

        Source:
        Chapter 12 of:
        https://web.stanford.edu/~cantwell/AA210A_Course_Material/AA210A_Course_Notes/

        Inputs:
        Mach, M0                             [-]
        Isentropic Expansion Factor, gamma   [-]
        Flow turn angle, theta               [rad]
        Strong Shock (0 = weak), delta       [-]

        Outputs:
        Shock Angle, Beta                    [rad]

        Properties Used:
        N/A
        """

    # Calculate wave angle
    mu   = np.arcsin(1./M0)
    c    = np.tan(mu)*np.tan(mu)

    # Calculate shock angle
    a    = ((gamma-1.)/2.+(gamma+1.)*c/2.)*np.tan(theta)
    b    = ((gamma+1.)/2.+(gamma+3.)*c/2.)*np.tan(theta)
    d    = np.sqrt(((4.*(1.-3.*a*b)**3.)/(27.*a*a*c+9.*a*b-2.)**2)-1.)

    beta =np.arctan((b+9.*a*c)/(2.*(1.-3.*a*b))-(d*(27.*a**2.*c+9.*a*b-2))/(6.*a*(1.-3.*a*b))*np.tan(n*np.pi/3.+1./3.*np.arctan(1./d)))

    return beta


#!/usr/bin/env python

from sympy import pi, cos, sin, sqrt, arg, exp, Abs, S
from sympy import integrate, Rational, nfloat, conjugate, Symbol
from sympy.functions.special.delta_functions import Piecewise

t = Symbol('t')

def Heaviside(arg):
    ''' Heaviside die ook op arg=0 gedefinieerd is '''
    return Piecewise((0, arg < 0), (1, arg >= 0))

# functies om de fourier-reeks coefficienten uit te rekenen
def an(f, T, n, omega0):
    return 2/T*integrate(f*cos(n*omega0*t), (t, 0, T))

def bn(f, T, n, omega0):
    return 2/T*integrate(f*sin(n*omega0*t), (t, 0, T))

def a0(f, T, omega0):
    return 1/T*integrate(f, (t, 0, T)).nsimplify()

def fourreeks(f, T, N):
    N += 1  # voor Maple
    a = (N-1)*[0]
    b = (N-1)*[0]
    omega0 = 2*pi/T
    som = a0(f, T, omega0)
    for n in range(1, N):
        a[n-1] = an(f, T, n, omega0)
        b[n-1] = bn(f, T, n, omega0)
        som += (a[n-1]*cos(n*omega0*t) + b[n-1]*sin(n*omega0*t)).nsimplify()
    return som, a, b

def spec(a, b, T):
    omega0 = 2*pi/T
    mag = [nfloat(sqrt(a[i]**2 + b[i]**2)) for i in range(len(a))]
    phi = len(a)*[0]
    for i in range(len(b)):
        if a[i] == S.Zero:
            phi[i] = 0
        else:
            phi[i] = nfloat(arg(a[i] - S.ImaginaryUnit*b[i]))
    mag.insert(0, a0(f, T, omega0))
    if mag[0] >= 0:
        phi.insert(0, 0)
    else:
        phi.insert(0, pi)
    return (mag, phi)

def periodiek(f, T, N):
    som = S(0)
    for n in range(N):
        som += (Heaviside(t-n*T) - Heaviside(t-(n+1)*T))*f.replace(t, t-n*T)
    return som

def plot_color(p):
    color = ['r', 'b', 'g', 'y', 'm']
    Nc = len(color)
    for i, pc in enumerate(p):
        pc.line_color = color[i % Nc]

if __name__ == '__main__':
    from sympy.plotting import plot
    import matplotlib.pyplot as plt

    f = Heaviside(t+1) - Heaviside(t-1)
    T = 4
    N = 5
    teind = 10
    som, a, b = fourreeks(f, T, N)
    mag, phase = spec(a, b, T)
    plot(som, (t, 0, teind))
    p = plot(periodiek(f, T, int(teind/T)+1), som, (t, 0, teind),
             adaptive=False, show=False)
    plot_color(p)
    p.show()
    phase, mag = spec(a, b, T)
    plt.figure()
    plt.subplot(2, 1, 1)
    plt.plot(range(len(phase)), phase)
    plt.subplot(2, 1, 2)
    plt.plot(range(len(mag)), mag)
    plt.show()


""" @file controls.py
"""

from scipy.linalg import solve_continuous_are as care
import numpy as np
import copy

###############################
## Control Policies
###############################
class ZeroPol:

    def __init__(self, du):
        self.du = du
        self.u = np.zeros((du))

    def evaluate(self, time, state):
        return self.u

class LinearFeedback: # Continuous Infinite-Horizon Linear Quadratic Regulator

    def __init__(self, A, B, Q, R):
        self.P = care(A, B, Q, R)
        self.K = np.linalg.solve(np.linalg.inv(R), np.dot(B.T, self.P))
        self.A = A
        self.B = B

        self.R = R
        self.Q = Q

    def evaluate(self, time, state):
        return -np.dot(self.K, state)

    def get_P(self):
        return self.P

    def get_Q(self):
        return self.Q

    def get_R(self):
        return self.R

class LinearFeedbackConstTracker:

    """ Class for linear quadratic tracker for constant state tracking
    """

    def __init__(self, A, B, Q, R, const, g=None):

        """ LinearFeedbackConstTracker constructor

        Input:
        - A:            linear time-invariant state matrix
        - B:            linear time-invariant input matrix
        - Q:            control weighting matrix for state
        - R:            control weighting matrix for control inputs
        - const:        constant state to track
        - g:            state to track in error dynamics system

        """

        self.A = A
        self.B = B
        self.Q = Q
        self.R = R

        self.P = care(A, B, Q, R)
        self.K = np.linalg.solve(R, np.dot(B.T, self.P))
        self.Bt = copy.deepcopy(B.T)
        self.RBt = np.dot(np.linalg.inv(R), self.Bt)
        self.BRBt = np.dot(B, self.RBt)
        if g is None:
            self.g = np.dot(A, const)
        else:
            self.g = copy.deepcopy(g)
        # print("g = ", self.g)

        self.p = -np.linalg.solve(A.T - np.dot(self.P, self.BRBt), np.dot(self.P, self.g))
        self.R = copy.deepcopy(R)
        self.const = copy.deepcopy(const)

        # Closed loop is
        # \dot{x} = A_cl x + g_cl
        self.Acl = A - np.dot(B, self.K)

        self.g_cl = np.dot(B, np.dot(self.K, const)) - np.dot(B, np.dot(np.linalg.inv(R), np.dot(B.T, self.p)))

        self.tracking = None

        # steady-state
        self.xss = np.dot(self.BRBt, self.p) - self.g
        self.xss = np.dot(np.linalg.inv(self.A - np.dot(self.BRBt.T, self.P)), self.xss)
        # self.xss = self.xss - self.const

        # steady-state optimal control
        # self.uss = self.evaluate(0, self.xss)
        self.uss = -np.dot(self.RBt, np.dot(self.P, self.xss) + self.p)

    def evaluate(self, time, state):
        # print("TIME: ", time, " STATE: ", state.T)

        return -np.dot(self.RBt, np.dot(self.P, state - self.const) + self.p)

    def get_closed_loop_A(self):
        return self.Acl

    def get_closed_loop_g(self):
        return self.g_cl

    def get_P(self):
        return self.P

    def get_Q(self):
        return self.Q

    def get_R(self):
        return self.R

    def get_p(self):
        return self.p

    def get_uss(self):
        return self.uss
        # return -np.dot(self.RBt, np.dot(self.P, self.xss) + self.p)

    def get_xss(self):
        return self.xss

    def set_const(self, time, target_id, const):
        self.tracking = target_id

        self.g = np.dot(self.A, const)

        self.p = -np.linalg.solve(self.A.T - np.dot(self.P, self.BRBt), np.dot(self.P, self.g))
        self.const = copy.deepcopy(const)

        # Closed loop is
        # \dot{x} = A_cl x + g_cl
        self.Acl = self.A - np.dot(self.B, self.K)

        self.g_cl = np.dot(self.B, np.dot(self.K, self.const)) - np.dot(self.B,
                np.dot(np.linalg.inv(self.R), np.dot(self.B.T, self.p)))

        self.tracking = None

        # steady-state
        self.xss = np.dot(self.BRBt, self.p) - self.g
        self.xss = np.dot(np.linalg.inv(self.A - np.dot(self.BRBt.T, self.P)), self.xss)

        self.uss = -np.dot(self.RBt, np.dot(self.P, self.xss) + self.p)

    def cost_to_go(self, time, state1, const_state):

        """ Computes cost_to_go for a agent/target to track a single agent/target
        """

        error_state = state1 - const_state

        g = np.dot(self.A, const_state)

        p = -np.linalg.solve(self.A.T - np.dot(self.P, self.BRBt), np.dot(self.P, g))

        g_cl = np.dot(self.B, np.dot(self.K, const_state)) - np.dot(self.B, np.dot(np.linalg.inv(self.R),
            np.dot(self.B.T, p)))

        self.tracking = None

        # steady-state
        xss = np.dot(self.BRBt, p) - g
        xss = np.dot(np.linalg.inv(self.A - np.dot(self.BRBt.T, self.P)), xss)

        # steady-state optimal control
        # self.uss = self.evaluate(0, self.xss)
        uss = -np.dot(self.RBt, np.dot(self.P, self.xss) + self.p)

        cost_to_go = np.dot(error_state, np.dot(self.P, error_state)) + 2*np.dot(p.T, error_state) -\
            np.dot(xss, np.dot(self.P, xss)) + 2*np.dot(p.T, xss)

        return cost_to_go

# IN DEVELOPMENT
class MinimumTimeIntercept(): # augmented proportional navigation - min time intercept with target accel

    """
    see bryson and ho (1975) 288 ch.9 eqn. 9.4.30
    """

    def __init__(self, time_final, dx):
        self.time_final = time_final
        self.dim_position = dx
        self.cp = 1
        self.ce = 3

    def get_final_time(self, state2):
        pass

    def track(self):
        pass

    def evaluate(self, time, state1, state2, feedforward=0):

        t = time
        tf = self.time_final

        ce = self.ce
        cp = self.cp

        # pursuer position
        rp = state1[:3]
        # pursuer velocity
        vp = state1[3:6]

        # evader position
        re = state2[:3]
        # evader velocity
        ve = state2[3:6]

        # proportional navigation guidance law
        u = -3/((1-ce/cp)*(tf-t)**2) * (rp-re + (vp-ve)*(tf-t) )

        return u

    def cost_to_go(self, time, state1, state2):
        # y1 = state1[:self.dim_position] - state2[:self.dim_position] # relative position
        # y2 = state1[self.dim_position:] - state2[self.dim_position:] # relative velocity
        # return  y1 + y2*(time_final-time) # miss distance
        time_to_go = self.time_final - time


'''
learn a little bit about markov chains

'''

import numpy as np
import argparse

mat = np.array([ # state transition matrix
    [2,6,2],
    [1,6,3],
    [2,7,1]
])/10
states = "sleep,run,icecream".split(',')

class MarkovSimple:
    def __init__(self,transitionmatrix:np.ndarray,statenames:list,initialstate=-1):
        ''' initialize markov chain. if no initial state given, one chosen at random'''
        self.m = transitionmatrix
        self.s = statenames
        r,c = self.m.shape
        assert r == c, "state transition matrix isn't square"
        assert r == len(self.s), "number of states doesn't match matrix"
        irow:np.ndarray = None
        for irow in self.m:
            assert np.isclose(irow.sum(),1.0),"row doesn't add up to 1: {}".format(irow)
        if(initialstate == -1):
            initialstate = np.random.choice(len(self.s))
        self.curr_state = initialstate
        self.hist = [self.curr_state]
        print('initial state:',self.s[self.curr_state])

    def next(self):
        ''' based on current state, go to next state'''
        irow = self.m[self.curr_state]
        self.curr_state = np.random.choice(len(self.s),p=irow)
        self.hist.append(self.curr_state)
        print('new state:',self.s[self.curr_state])

if(__name__ == '__main__'):
    # test out weather markov chain, using probabilities (not dice)
    mat = np.array([
        [0.7,0,0,0,0.09,0.075,0.09,0.045],
        [0,0.7,0,0,0.09,0.075,0.09,0.045],
        [0,0,0.7,0,0.09,0.075,0.09,0.045],
        [0,0,0,0.7,0.09,0.075,0.09,0.045],
        [0.075,0.075,0.075,0.075,0.21,0.175,0.21,0.105],
        [0.075,0.075,0.075,0.075,0.21,0.175,0.21,0.105],
        [0.075,0.075,0.075,0.075,0.21,0.175,0.21,0.105],
        [0.075,0.075,0.075,0.075,0.21,0.175,0.21,0.105],
    ])
    states = 'cool & foggy,clear & cool,clear & warm,cloudy & warm,clear & dry,rain showers,rain,storm'.split(',')
    am = MarkovSimple(mat, states)
    for i in range(10):
        am.next()

# eof


# -*- coding: utf-8 -*-
"""
Created on Sun Dec 24 21:32:08 2017

@author: chris
"""

import os
os.chdir("E:\\workspace\\project")

from mymodule.GMM import GMM
import numpy as np
from scipy import stats
from matplotlib import pyplot as plt

def mixture(x,p,mu,sigma):
    mix = p[0]*stats.multivariate_normal.pdf(x,mu[0],sigma[0])
    for i in range(1,len(p)):
        mix += p[i]*stats.multivariate_normal.pdf(x,mu[i],sigma[i])
    return mix

def scatter_contour(y,model,G,devide=True):
    right=0
    colors = ['tab:blue','tab:orange','tab:olive']
    fig = plt.figure()
    # scatter
    for i in range(G):
        left = right
        right += group[i]
        plt.scatter(y[left:right,0],y[left:right,1],alpha=0.6,color=colors[i])
    # contour
    if (devide):
        for i in range(model.g):
            A,B = np.meshgrid(np.arange(min(y[:,0]),max(y[:,0]),0.1),np.arange(min(y[:,1]),max(y[:,1]),0.1))
            tmp = np.append(A.reshape(A.shape[0]*A.shape[1],1),B.reshape(B.shape[0]*B.shape[1],1),1)
            C = stats.multivariate_normal.pdf(tmp,model.mu[i],model.sigma[i]).reshape(A.shape)
            plt.contour(A,B,C,5)
    else:
        A,B = np.meshgrid(np.arange(min(Y[:,0]),max(Y[:,0]),0.1),np.arange(min(Y[:,1]),max(Y[:,1]),0.1))
        tmp = np.append(A.reshape(A.shape[0]*A.shape[1],1),B.reshape(B.shape[0]*B.shape[1],1),1)
        C = mixture(tmp,model.p,model.mu,model.sigma).reshape(A.shape)
        plt.contour(A,B,C,10)
    return fig

G = 3
N = 1000
P = [0.5,
     0.3,
     0.2]
MU = [[-3,0],
      [0,0],
      [0,3]]

# Homo
SIGMA = [[[1,0.2],[0.2,1]],
         [[1,0.2],[0.2,1]],
         [[1,0.2],[0.2,1]]]

# single trial
np.random.seed(0)
group = np.random.multinomial(N,P)
Y = np.random.multivariate_normal(MU[0],SIGMA[0],group[0])
for i in range(1,len(P)):
    Y = np.append(Y,np.random.multivariate_normal(MU[i],SIGMA[i],group[i]),axis=0)

# train model
model = GMM(Y,G,True,1e-3)
# plot hetero_model_homo_data.png
fig = scatter_contour(Y,model,G,True)
fig.savefig("hetero_model_homo_data.png",dpi=120)
# plot hetero_model_homo_data_mixcontour.png
fig = scatter_contour(Y,model,G,False)
fig.savefig("hetero_model_homo_data_mixcontour.png",dpi=120)

# train model
model = GMM(Y,G,False,1e-3)
# plot homo_model_homo_data.png
fig = scatter_contour(Y,model,G,True)
fig.savefig("homo_model_homo_data.png",dpi=120)
# plot homo_model_homo_data_mixcontour.png
fig = scatter_contour(Y,model,G,False)
fig.savefig("homo_model_homo_data_mixcontour.png",dpi=120)


import numpy as np
import scipy
import scipy.stats as ss

def calc_2d_pdf(data, pdf_res = 0.3, minval = -10, maxval = 10):

    #x_bins = np.linspace(np.amin(data[0,0:]), np.amax(data[0,0:]), pdf_res)
    #y_bins = np.linspace(np.amin(data[1,0:]), np.amax(data[1,0:]), pdf_res)
    x_bins = np.arange(minval,maxval,pdf_res)
    y_bins = np.arange(minval,maxval,pdf_res)

    a,b,c = np.histogram2d(data[0,0:],data[1,0:],(x_bins,y_bins), density=True)

    return a

def integrate_2d(twod_data, minval = -10, maxval = 10, pdf_res = 0.3):

    xax_vals = np.arange(minval,maxval,pdf_res)[0:-1] + pdf_res/2
    temp = [np.trapz(y = twod_data[0:,i], x = xax_vals) for i in range(pdf1.shape[1])]
    int_val = np.trapz(y = temp, x = xax_vals)
    return int_val

#---------------------------------------------------------------------
#------------------------KL divergence module-------------------------
#---------------------------------------------------------------------

def calc_KL_divergence_2d(dat1, dat2):
    # assuming pdf1 is the 'true' pdf and pdf2 is sampled from it

    pdf1 = dat1.copy()#/np.sum(dat1)
    pdf2 = dat2.copy()#/np.sum(dat2)
    pdf1[pdf1 == 0] = 1e-10
    pdf2[pdf2 == 0] = 1e-10

    KLD = np.zeros_like(pdf1)
    for i in range(pdf1.shape[0]):
        for j in range(pdf1.shape[1]):

            KLD[i,j] = pdf1[i,j] * (np.log10(pdf1[i,j]/pdf2[i,j]))
#             print(pdf1[i,j], pdf2[i,j], (np.log10(pdf1[i,j]) - np.log10(pdf2[i,j])))
    return KLD

#---------------------------------------------------------------------
#------------------------KS test (2d) module--------------------------
#---------------------------------------------------------------------

def fhCounts(x,edge):
    # computes local CDF at a given point considering all possible axis orderings

    templist = [np.sum((x[0,0:] >= edge[0]) & (x[1,0:] >= edge[1])),
                np.sum((x[0,0:] <= edge[0]) & (x[1,0:] >= edge[1])),
                np.sum((x[0,0:] <= edge[0]) & (x[1,0:] <= edge[1])),
                np.sum((x[0,0:] >= edge[0]) & (x[1,0:] <= edge[1]))]
    return templist

def kstest_2d(dist1, dist2, alpha = 0.05):

    num1 = dist1.shape[1]
    num2 = dist2.shape[1]

    KSstat = -np.inf

    for iX in (np.arange(0,num1+num2)):

        if iX < num1:
            edge = dist1[0:,iX]
        else:
            edge = dist2[0:,iX-num1]

#         vfCDF1 = np.sum(fhCounts(dist1, edge)) / num1
#         vfCDF2 = np.sum(fhCounts(dist2, edge)) / num2

        vfCDF1 = np.array(fhCounts(dist1, edge)) / num1
        vfCDF2 = np.array(fhCounts(dist2, edge)) / num2

        vfThisKSTS = np.abs(vfCDF1 - vfCDF2)
        fKSTS = np.amax(vfThisKSTS)

        if (fKSTS > KSstat):
            KSstat = fKSTS
            #print(KSstat, vfCDF1, vfCDF2)

    # Peacock Z calculation and P estimation

    n =  num1 * num2 /(num1 + num2)
    Zn = np.sqrt(n) * KSstat;
    Zinf = Zn / (1 - 0.53 * n**(-0.9));
    pValue = 2 *np.exp(-2 * (Zinf - 0.5)**2);

    # Clip invalid values for P
    if pValue > 1.0:
        pValue = 1.0

#     H = (pValue <= alpha)

    return pValue, KSstat

#---------------------------------------------------------------------
#------------------------Hotelling T2 test module---------------------
#---------------------------------------------------------------------

def hotelling_T2(dist1, dist2):

    # Hotelling's T-Squared test for comparing d-dimensional data from two independent samples,
    # assuming normality w/ common covariance matrix

    if dist1.shape[0] != dist2.shape[0]:
        print('Error: ordered pair dimensions do not match')
    else:
        p = dist1.shape[0]

    n = dist1.shape[1]
    m = dist2.shape[1]

    n = n+m
    mux = np.mean(dist1,1)
    muy = np.mean(dist2,1)

    Sx = np.cov(dist1)
    Sy = np.cov(dist2)

    Su = (n*Sx + m*Sy) / (n-2) # unbiased estimate
    d = np.zeros((p,1))
    d[0:,0] = mux - muy

    D2 = np.matmul(d.T,np.matmul(np.linalg.inv(Su),d)) # reversing the formula for python's matrix handling
    #D2 = d.T*np.linalg.inv(Su)*d
    D2 = D2[0][0]
    T2 = ((n*m)/n)*D2
    F = T2 * (n-p-1) / ((n-2)*p);

    pval = 1 - scipy.stats.f.cdf(F,p,n-p-1);
    return pval, T2

#---------------------------------------------------------------------
#------------------------Hotelling T2 test module---------------------
#---------------------------------------------------------------------

def energy_statistics_test(dist1, dist2, alpha = 0.05, flag = 'szekely-rizzo', nboot = 100, replace = False, progbar = False):

    if dist1.shape[0] != dist2.shape[0]:
        print('Error: ordered pair dimensions do not match')

    n = dist1.shape[1]
    m = dist2.shape[1]

    pooled = np.hstack((dist1,dist2))

    e_n = energy(dist1, dist2, flag)
    e_n_boot = np.zeros((nboot, ))
    e_n_boot[0] = e_n

    if progbar == True:

        for i in tqdm(range(1,nboot)):

            if replace == False:
                ind = np.random.choice(n+m, size=(n+m,))
            else:
                ind = np.random.permutation(np.arange(n+m))
            e_n_boot[i] = energy(pooled[0:,ind[0:n]],pooled[0:,ind[n:]], flag)
    else:
        for i in (range(1,nboot)):

            if replace == False:
                ind = np.random.choice(n+m, size=(n+m,))
            else:
                ind = np.random.permutation(np.arange(n+m))
            e_n_boot[i] = energy(pooled[0:,ind[0:n]],pooled[0:,ind[n:]], flag)

    p = np.sum(e_n_boot>=e_n)/nboot

    return p, e_n, e_n_boot

def dist(x,y):
    dx = scipy.spatial.distance.pdist(x,'euclidean')
    dy = scipy.spatial.distance.pdist(y,'euclidean')
    dxy = scipy.spatial.distance.cdist(x.T,y.T,'euclidean')
    return dx,dy,dxy

def energy(x,y,flag):
    n = x.shape[1]
    m = y.shape[1]
    dx, dy, dxy = dist(x,y)

    if flag == 'aslan-zech':
        temp = -np.log(dxy)
        z = (1/(n*(n-1)))*np.sum(-np.log(dx)) + (1/(m*(m-1)))*np.sum(-np.log(dy)) - (1/(n*m))*np.sum(temp[(temp>-np.inf) & (temp<np.inf)])
    elif flag == 'szekely-rizzo':
        z = (2/(n*m))*np.sum(dxy) - (1/(n**2))*np.sum(2*dx) - (1/(m**2))*np.sum(2*dy)
        z = ((n*m)/(n+m))*z
    else:
        z = np.nan
        print('did not understand method input')

    return z


import numpy as np
from random import shuffle

def softmax_loss_naive(W, X, y, reg):
  """
  Softmax loss function, naive implementation (with loops)

  Inputs have dimension D, there are C classes, and we operate on minibatches
  of N examples.

  Inputs:
  - W: A numpy array of shape (D, C) containing weights.
  - X: A numpy array of shape (N, D) containing a minibatch of data.
  - y: A numpy array of shape (N,) containing training labels; y[i] = c means
    that X[i] has label c, where 0 <= c < C.
  - reg: (float) regularization strength

  Returns a tuple of:
  - loss as single float
  - gradient with respect to weights W; an array of same shape as W
  """
  # Initialize the loss and gradient to zero.
  loss = 0.0
  dW = np.zeros_like(W)

  #############################################################################
  # TODO: Compute the softmax loss and its gradient using explicit loops.     #
  # Store the loss in loss and the gradient in dW. If you are not careful     #
  # here, it is easy to run into numeric instability. Don't forget the        #
  # regularization!                                                           #
  #############################################################################
  #  Good article on softmax gradient derivation: https://madalinabuzau.github.io/2016/11/29/gradient-descent-on-a-softmax-cross-entropy-cost-function.html
  m = X.shape[0]
  for i in range(m):
        x_i = X[i]
        y_i = y[i]
        score = x_i.dot(W)
        stable_score = score - np.max(score)
        exp_sum = 0
        target = 0
        for i, score in enumerate(stable_score):
            exp_score = np.exp(score)
            exp_sum += exp_score
            if i == y_i:
                target = exp_score
        for i, score in enumerate(stable_score):
            exp_score = np.exp(score)
            scale = exp_score / exp_sum
            if i == y_i:
                dW[:, i] += ((scale - 1) * x_i).T
            else:
                dW[:, i] += (scale * x_i).T
        final_score = target / exp_sum
        loss += -np.log(final_score)

  loss /= m
  loss += reg * np.sum(W * W)
  dW /= m
  dW += 2 * reg * W
  #############################################################################
  #                          END OF YOUR CODE                                 #
  #############################################################################

  return loss, dW

def softmax_loss_vectorized(W, X, y, reg):
  """
  Softmax loss function, vectorized version.

  Inputs and outputs are the same as softmax_loss_naive.
  """
  # Initialize the loss and gradient to zero.
  loss = 0.0
  dW = np.zeros_like(W)

  #############################################################################
  # TODO: Compute the softmax loss and its gradient using no explicit loops.  #
  # Store the loss in loss and the gradient in dW. If you are not careful     #
  # here, it is easy to run into numeric instability. Don't forget the        #
  # regularization!                                                           #
  #############################################################################
  m = X.shape[0]

  # Loss Calculation
  logits = X.dot(W)
  stable_logits = logits - np.max(logits, axis=1, keepdims=True)
  scores = np.exp(stable_logits) / np.sum(np.exp(stable_logits), axis=1, keepdims=True)
  loss = np.sum(-np.log(scores[np.arange(m), y]))
  loss /= m
  loss += reg * np.sum(W * W)

  # Gradient Calculation
  scores[np.arange(m), y] -= 1
  dW = X.T.dot(scores)
  dW /= m
  dW += 2 * reg * W
  #############################################################################
  #                          END OF YOUR CODE                                 #
  #############################################################################

  return loss, dW


import math

import numpy as np

class ZeroVectorError(Exception): pass

class Vector3:

    def __init__(self, x=0, y=0, z=0):
        self.x = x
        self.y = y
        self.z = z

    @classmethod
    def unity(cls):
        return cls(1,1,1)

    @property
    def A(self):
        return np.array((self.x, self.y, self.z))

    def __add__(self, other):
        return Vector3(
            self.x + other.x,
            self.y + other.y,
            self.z + other.z
        )

    def __eq__(self, other):
        return self.x == other.x and self.y == other.y and self.z == other.z

    def __sub__(self, other):
        return self + -other

    def __mul__(self, s):
        return Vector3(self.x * s, self.y * s, self.z * s)

    def __rmul__(self, s):
        return self * s

    def __neg__(self):
        return Vector3(-self.x, -self.y, -self.z)

    def __repr__(self):
        return "Vector3([{} {} {}])".format(self.x, self.y, self.z)

    @classmethod
    def zero(cls):
        """returns Vector(0,0,0)"""
        return cls(0, 0, 0)

    def as_tuple(self):
        return self.x, self.y, self.z

    def length(self):
        """ returns the magnitude of this vector """
        return math.sqrt(self.x**2 + self.y**2 + self.z**2)

    def dot(self, other):
        return self.x * other.x + self.y * other.y + self.z * other.z

    def cross(self, other):
        """ cross product for left handed coordinate system """
        a = self
        b = other
        return -1 * Vector3(
            a.y * b.z - a.z * b.y,
            a.z * b.x - a.x * b.z,
            a.x * b.y - a.y * b.x
        )

    def angle_between(self, other):
        """ returns the angle between this vector and other """
        dot = self.normalize().dot(other.normalize())
        angle = math.acos(dot)  # / (self.length() * other.length()))
        return angle

    def normalize(self):
        """ returns a vector in the same direction with unit length
        In the degenerate case of a zero vector, raises"""

        length = self.length()
        if length == 0:
            raise ZeroVectorError("can't normalize 0 vector")
        return Vector3(self.x / length, self.y / length, self.z / length)

    def lerp(self, other, t):
        """ lerp between self and other
        when t=0, returns self
        when t=1 returns other
        """
        return self + t*(other - self)


"""
Solving the two-dimensional diffusion equation

Example acquired from https://scipython.com/book/chapter-7-matplotlib/examples/the-two-dimensional-diffusion-equation/
"""

import numpy as np
import matplotlib.pyplot as plt

class SolveDiffusion2D:
    def __init__(self):
        """
        Constructor of class SolveDiffusion2D
        """
        # plate size, mm
        self.w = None
        self.h = None

        # intervals in x-, y- directions, mm
        self.dx = None
        self.dy = None

        # Number of discrete mesh points in X and Y directions
        self.nx = None
        self.ny = None

        # Thermal diffusivity of steel, mm^2/s
        self.D = None

        # Initial cold temperature of square domain
        self.T_cold = None

        # Initial hot temperature of circular disc at the center
        self.T_hot = None

        # Timestep
        self.dt = None

    def initialize_domain(self, w=10., h=10., dx=0.1, dy=0.1):

        assert isinstance(w, float)
        assert isinstance(h, float)
        assert isinstance(dx, float)
        assert isinstance(dy, float)

        self.w = w
        self.h = h
        self.dx = dx
        self.dy = dy
        self.nx = int(w / dx)
        #self.nx = int(h / dx)
        self.ny = int(h / dy)

    def initialize_physical_parameters(self, d=4., T_cold=300., T_hot=700.):

        assert isinstance(d, float)
        assert isinstance(T_cold, float)
        assert isinstance(T_hot, float)

        self.D = d
        self.T_cold = T_cold
        self.T_hot = T_hot

        # Computing a stable time step
        dx2, dy2 = self.dx * self.dx, self.dy * self.dy
        self.dt = dx2 * dy2 / (2 * self.D * (dx2 + dy2))

        #self.dt = dx2 + dy2 / (2 * self.D * (dx2 + dy2))

        print("dt = {}".format(self.dt))

    def set_initial_condition(self):
        u = self.T_cold * np.ones((self.nx, self.ny))

        # Initial conditions - circle of radius r centred at (cx,cy) (mm)
        r, cx, cy = 2, 5, 5
        r2 = r ** 2
        for i in range(self.nx):
            for j in range(self.ny):
                p2 = (i * self.dx - cx) ** 2 + (j * self.dy - cy) ** 2
                if p2 < r2:
                #if p2 > r2:
                    u[i, j] = self.T_hot

        return u.copy()

    def do_timestep(self, u_nm1):
        u = u_nm1.copy()

        dx2 = self.dx * self.dx
        dy2 = self.dy * self.dy

        # Propagate with forward-difference in time, central-difference in space
        u[1:-1, 1:-1] = u_nm1[1:-1, 1:-1] + self.D * self.dt * (
                (u_nm1[2:, 1:-1] - 2 * u_nm1[1:-1, 1:-1] + u_nm1[:-2, 1:-1]) / dx2
                + (u_nm1[1:-1, 2:] - 2 * u_nm1[1:-1, 1:-1] + u_nm1[1:-1, :-2]) / dy2)

        return u.copy()

    def create_figure(self, fig, u, n, fignum):
        fignum += 1
        ax = fig.add_subplot(220 + fignum)
        im = ax.imshow(u.copy(), cmap=plt.get_cmap('hot'), vmin=self.T_cold, vmax=self.T_hot)
        ax.set_axis_off()
        ax.set_title('{:.1f} ms'.format(n * self.dt * 1000))

        return fignum, im

def output_figure(fig, im):
    fig.subplots_adjust(right=0.85)
    cbar_ax = fig.add_axes([0.9, 0.15, 0.03, 0.7])
    cbar_ax.set_xlabel('$T$ / K', labelpad=20)
    fig.colorbar(im, cax=cbar_ax)
    plt.show()

def main():
    DiffusionSolver = SolveDiffusion2D()

    DiffusionSolver.initialize_domain()

    DiffusionSolver.initialize_physical_parameters()

    u0 = DiffusionSolver.set_initial_condition()
    print(u0)

    # Number of timesteps
    nsteps = 101

    # Output 4 figures at these timesteps
    n_output = [0, 10, 50, 100]

    fig_counter = 0
    fig = plt.figure()

    im = None

    # Time loop
    for n in range(nsteps):
        u = DiffusionSolver.do_timestep(u0)

        # Create figure
        if n in n_output:
            fig_counter, im = DiffusionSolver.create_figure(fig, u, n, fig_counter)

        u0 = u.copy()

    # Plot output figures
    output_figure(fig, im)

if __name__ == "__main__":
    main()