<filename>agent.py
# -*- coding: utf-8 -*-
"""
Created on Sun Feb 28 11:15:15 2021

@author: <NAME>
"""
from numpy import random
from sympy import hyper

from agent_control import AgentControl
from replay_buffer import ReplayBuffer
from collections import namedtuple
import time
import numpy as np
import math
import telegram_bot as tg

class Agent:

    Transition = namedtuple('Transition',('state', 'action', 'next_state', 'reward', 'done'), rename = False) # 'rename' means not to overwrite invalid field

    def __init__(self, env, hyperparameters, device, writer, max_games, tg_bot):
        self.eps_start = hyperparameters['eps_start']
        self.eps_end = hyperparameters['eps_end']
        self.eps_decay = hyperparameters['eps_decay']
        self.epsilon = hyperparameters['eps_start']
        self.n_iter_update_nn = hyperparameters['n_iter_update_nn']
        self.max_games = max_games
        self.tg_bot = tg_bot
        self.env = env

        self.agent_control = AgentControl(env, device, hyperparameters['learning_rate'], hyperparameters['gamma'], hyperparameters['multi_step'], hyperparameters['double_dqn'], hyperparameters['dueling'])
        self.replay_buffer = ReplayBuffer(hyperparameters['buffer_size'], hyperparameters['buffer_minimum'], hyperparameters['multi_step'], hyperparameters['gamma'])
        self.summary_writer = writer

        self.num_iterations = 0
        self.total_reward = 0
        self.num_games = 0
        self.total_loss = []
        self.ts_frame = 0
        self.ts = time.time()
        self.birth_time = time.time()
        self.rewards = []

        if self.tg_bot:
            tg.welcome_msg(hyperparameters['multi_step'], hyperparameters['double_dqn'], hyperparameters['dueling'])

    def select_greedy_action(self, obs):
        # Give current state to the control who will pass it to NN which will
        # return all actions and the control will take max and return it here
        return self.agent_control.select_greedy_action(obs)

    def select_eps_greedy_action(self, obs):
        rand_num = random.rand()
        if self.epsilon > rand_num:
            # Select random action - explore
            return self.env.action_space.sample()
        else:
            # Select best action
            return self.select_greedy_action(obs)

    def add_to_buffer(self, obs, action, new_obs, reward, done):
        transition = self.Transition(state = obs, action = action, next_state = new_obs, reward = reward, done = done)
        self.replay_buffer.append(transition)
        self.num_iterations = self.num_iterations + 1
        if self.epsilon > self.eps_end:
            self.epsilon = self.eps_start - self.num_iterations / self.eps_decay
        self.total_reward = self.total_reward + reward

    def sample_and_improve(self, batch_size):
        # If buffer is big enough
        if len(self.replay_buffer.buffer) > self.replay_buffer.minimum:
            # Sample batch_size number of transitions from buffer B
            mini_batch = self.replay_buffer.sample(batch_size)
            # Calculate loss and improve NN
            loss = self.agent_control.improve(mini_batch)
            # So we can calculate mean of all loss during one game
            self.total_loss.append(loss)

        if ( self.num_iterations % self.n_iter_update_nn) == 0:
            self.agent_control.update_target_nn()

    def reset_parameters(self):
        self.rewards.append(self.total_reward)
        self.total_reward = 0
        self.num_games = self.num_games + 1
        self.total_loss = []

    def print_info(self):
        # print(self.num_iterations, self.ts_frame, time.time(), self.ts)
        fps = (self.num_iterations-self.ts_frame)/(time.time()-self.ts)
        print('%d %d rew:%d mean_rew:%.2f fps:%d, eps:%.2f, loss:%.4f' % (self.num_iterations, self.num_games, self.total_reward, np.mean(self.rewards[-40:]), fps, self.epsilon, np.mean(self.total_loss)))
        self.ts_frame = self.num_iterations
        self.ts = time.time()

        if self.summary_writer != None:
            self.summary_writer.add_scalar('reward', self.total_reward, self.num_games)
            self.summary_writer.add_scalar('mean_reward', np.mean(self.rewards[-40:]), self.num_games)
            self.summary_writer.add_scalar('10_mean_reward', np.mean(self.rewards[-10:]), self.num_games)
            self.summary_writer.add_scalar('esilon', self.epsilon, self.num_games)
            self.summary_writer.add_scalar('loss', np.mean(self.total_loss), self.num_games)

        if self.tg_bot:
            if (self.num_games % 10) == 0:
                tg.info_msg(self.num_games+1, self.max_games, np.mean(self.rewards[-40:]), np.mean(self.total_loss))
            if self.num_games == (self.max_games - 1):
                tg.end_msg(time.time() - self.birth_time)


import random
from fractions import Fraction

MAX_HASH = pow(2, 32)
FILES_TO_SIMULATE = 100000
MAX_SERVERS = 10

def generate_fraction(servs):
    res = [[(0, Fraction(1, 1))]]  # server index, fraction
    for i in range(1, servs):
        rn = []
        fraction_to_subtract = {}
        for si in range(0, i):
            fraction_to_subtract[si] = Fraction(1, i * (i + 1))
        # print('to subtract per server:', fraction_to_subtract)
        # take a part from each previous server.
        for piece in res[i - 1]:

            # print('starting:', piece, end='\t')
            si = piece[0]
            fr = piece[1]

            if (not fr > fraction_to_subtract[si]):
                # print('giving it all to nfew server')
                rn.append([i, fr])
                fraction_to_subtract[si] -= fr
            else:
                if si % 2:  # makes similar parts sit together
                    # rewriting old one but now its fraction decreases
                    rn.append([si, fr - fraction_to_subtract[si]])
                    # adding new one
                    rn.append([i, fraction_to_subtract[si]])
                else:
                    rn.append([i, fraction_to_subtract[si]])
                    rn.append([si, fr - fraction_to_subtract[si]])
                fraction_to_subtract[si] = Fraction(0, 1)

        # connect pieces siting next to each other
        r = []
        p = [-1, Fraction(1, 1)]
        # print('rn:', rn)
        for j in rn:
            if j[0] != p[0]:
                r.append(j)
            else:
                r[-1] = ([j[0], j[1] + p[1]])
            p = j
        # print(r)
        res.append(r)
    return res

fractions = generate_fraction(MAX_SERVERS)
ranges = []
for s, f in enumerate(fractions):
    ul = 0
    rang = []
    for i in f:
        ul += i[1]
        rang.append([i[0], ul])
    print(rang)
    ranges.append(rang)
    print('servers:', s + 1, '\tpieces:', len(f), '\nfractions:', f, '\nranges:', rang)

def algo1(fm, servers):
    return fm % servers

def algo2(fm, servers):
    fr = (MAX_HASH - 1) / servers
    return int(fm / fr)

def algo3(fm, servers):
    ff = Fraction(fm, MAX_HASH)
    ul = 0
    # print(ranges[servers])
    for i in ranges[servers]:
        if ff < i[1]:
            # print(ff, i, float(ff))
            return i[0]

print("=======================================")

matches = [0] * (MAX_SERVERS - 1)
for files in range(FILES_TO_SIMULATE):
    fm = random.randint(0, MAX_HASH)
    endup = []
    for servers in range(0, MAX_SERVERS):
        s = algo3(fm, servers)
        endup.append(s)
    # print('-' * 30)
    # print(fm, endup)

    l = endup[0]
    for i, e in enumerate(endup[1:]):
        if e == l:
            matches[i] += 1
        l = e
print('=' * 30)
print('reused from total of', FILES_TO_SIMULATE, 'files')
print(matches)


<reponame>waitong94/smarties
#!/usr/bin/env python3
#
#  smarties
#  Copyright (c) 2018 CSE-Lab, ETH Zurich, Switzerland. All rights reserved.
#  Distributed under the terms of the MIT license.
#
#  Created by <NAME> (<EMAIL>).
#
import sys
import numpy as np
import matplotlib.pyplot as plt
from numpy import random
from scipy.stats.stats import pearsonr

XCOL = int(sys.argv[1])
YCOL = int(sys.argv[2])
NFILES = len(range(3, len(sys.argv)))

for j in range(3, len(sys.argv)):
  DATA = np.fromfile(sys.argv[j], dtype=np.float32)
  print (DATA.size // 4)
  DATA = DATA.reshape([DATA.size // 4, 4])
  start = 3000000
  print(np.corrcoef(DATA[start:, XCOL], DATA[start:, YCOL]))
  plt.semilogx(DATA[start:, XCOL], DATA[start:, YCOL], '.')
plt.show()


<reponame>jonepatr/lets_face_it<filename>code/glow_pytorch/glow/modules.py
import numpy as np
import scipy.linalg
import torch
import torch.nn as nn
import torch.nn.functional as F

from glow_pytorch.glow import thops

class ActNorm2d(nn.Module):
    """
    Activation Normalization
    Initialize the bias and scale with a given minibatch,
    so that the output per-channel have zero mean and unit variance for that.

    After initialization, `bias` and `logs` will be trained as parameters.
    """

    def __init__(self, num_features, scale=1.0):
        super().__init__()
        # register mean and scale
        size = [1, num_features]
        self.register_parameter("bias", nn.Parameter(torch.zeros(*size)))
        self.register_parameter("logs", nn.Parameter(torch.zeros(*size)))
        self.num_features = num_features
        self.scale = float(scale)
        self.inited = False

    def _check_input_dim(self, input):
        return NotImplemented

    def initialize_parameters(self, input):

        if not self.training:
            return
        assert input.device == self.bias.device
        with torch.no_grad():
            bias = thops.mean(input.clone(), dim=0, keepdim=True) * -1.0
            vars = thops.mean((input.clone() + bias) ** 2, dim=0, keepdim=True)
            logs = torch.log(self.scale / (torch.sqrt(vars) + 1e-6))
            self.bias.data.copy_(bias.data)
            self.logs.data.copy_(logs.data)
            self.inited = True

    def _center(self, input, reverse=False):
        if not reverse:
            return input + self.bias
        else:
            return input - self.bias

    def _scale(self, input, logdet=None, reverse=False):
        logs = self.logs
        if not reverse:
            input = input * torch.exp(logs)
        else:
            input = input * torch.exp(-logs)
        if logdet is not None:
            """
            logs is log_std of `mean of channels`
            so we need to multiply on the channel length
            """
            dlogdet = thops.sum(logs) * input.size(1)
            if reverse:
                dlogdet *= -1
            logdet = logdet + dlogdet
        return input, logdet

    def forward(self, input, logdet=None, reverse=False):
        if not self.inited:
            self.initialize_parameters(input)
        # no need to permute dims as old version
        if not reverse:
            # center and scale
            input = self._center(input, reverse)
            input, logdet = self._scale(input, logdet, reverse)
        else:
            # scale and center
            input, logdet = self._scale(input, logdet, reverse)
            input = self._center(input, reverse)
        return input, logdet

class LinearZeros(nn.Linear):
    def __init__(self, in_channels, out_channels, logscale_factor=3):
        super().__init__(in_channels, out_channels)
        self.logscale_factor = logscale_factor
        # set logs parameter
        self.register_parameter("logs", nn.Parameter(torch.zeros(out_channels)))
        # init
        self.weight.data.zero_()
        self.bias.data.zero_()

    def forward(self, input):
        output = super().forward(input)
        return output * torch.exp(self.logs * self.logscale_factor)

class Permute2d(nn.Module):
    def __init__(self, num_channels, shuffle):
        super().__init__()
        self.num_channels = num_channels
        self.indices = np.arange(self.num_channels - 1).astype(np.long)
        self.indices_inverse = np.zeros((self.num_channels), dtype=np.long)
        for i in range(self.num_channels):
            self.indices_inverse[self.indices[i]] = i
        if shuffle:
            self.reset_indices()

    def reset_indices(self):
        np.random.shuffle(self.indices)
        for i in range(self.num_channels):
            self.indices_inverse[self.indices[i]] = i

    def forward(self, input, reverse=False):
        assert len(input.size()) == 4
        if not reverse:
            return input[:, self.indices]
        else:
            return input[:, self.indices_inverse]

class InvertibleConv1x1(nn.Module):
    def __init__(self, num_channels, LU_decomposed=False):
        super().__init__()
        w_shape = [num_channels, num_channels]
        w_init = np.linalg.qr(np.random.randn(*w_shape))[0].astype(np.float32)
        if not LU_decomposed:
            # Sample a random orthogonal matrix:
            self.register_parameter("weight", nn.Parameter(torch.Tensor(w_init)))
        else:
            np_p, np_l, np_u = scipy.linalg.lu(w_init)
            np_s = np.diag(np_u)
            np_sign_s = np.sign(np_s)
            np_log_s = np.log(np.abs(np_s))
            np_u = np.triu(np_u, k=1)
            l_mask = np.tril(np.ones(w_shape, dtype=np.float32), -1)
            eye = np.eye(*w_shape, dtype=np.float32)

            self.register_buffer("p", torch.Tensor(np_p.astype(np.float32)))
            self.register_buffer("sign_s", torch.Tensor(np_sign_s.astype(np.float32)))
            self.l = nn.Parameter(torch.Tensor(np_l.astype(np.float32)))
            self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(np.float32)))
            self.u = nn.Parameter(torch.Tensor(np_u.astype(np.float32)))
            self.l_mask = torch.Tensor(l_mask)
            self.eye = torch.Tensor(eye)
        self.w_shape = w_shape
        self.LU = LU_decomposed

    def get_weight(self, input, reverse):
        w_shape = self.w_shape
        if not self.LU:
            dlogdet = torch.slogdet(self.weight)[1] * input.size(1)
            if not reverse:
                weight = self.weight.view(w_shape[0], w_shape[1])
            else:
                weight = (
                    torch.inverse(self.weight.double())
                    .float()
                    .view(w_shape[0], w_shape[1])
                )
            return weight, dlogdet
        else:
            self.p = self.p.to(input.device)
            self.sign_s = self.sign_s.to(input.device)
            self.l_mask = self.l_mask.to(input.device)
            self.eye = self.eye.to(input.device)
            l = self.l * self.l_mask + self.eye
            u = self.u * self.l_mask.transpose(0, 1).contiguous() + torch.diag(
                self.sign_s * torch.exp(self.log_s)
            )
            dlogdet = thops.sum(self.log_s) * input.size(1)
            if not reverse:
                w = torch.matmul(self.p, torch.matmul(l, u))
            else:
                l = torch.inverse(l.double()).float()
                u = torch.inverse(u.double()).float()
                w = torch.matmul(u, torch.matmul(l, self.p.inverse()))
            return w.view(w_shape[0], w_shape[1]), dlogdet

    def forward(self, input, logdet=None, reverse=False):
        """
        log-det = log|abs(|W|)| * pixels
        """
        weight, dlogdet = self.get_weight(input, reverse)
        if not reverse:
            z = torch.matmul(input, weight)
            if logdet is not None:
                logdet = logdet + dlogdet
            return z, logdet
        else:
            z = torch.matmul(input, weight)
            if logdet is not None:
                logdet = logdet - dlogdet
            return z, logdet

class GaussianDiag:
    Log2PI = float(np.log(2 * np.pi))

    @staticmethod
    def likelihood_simplified(x):
        """
        lnL = -1/2 * { ln|Var| + ((X - Mu)^T)(Var^-1)(X - Mu) + kln(2*PI) }
              k = 1 (Independent)
              Var = logs ** 2
        """
        return -0.5 * ((x ** 2) + GaussianDiag.Log2PI)

    @staticmethod
    def logp_simplified(x):
        likelihood = GaussianDiag.likelihood_simplified(x)
        return torch.sum(likelihood, dim=1)

    @staticmethod
    def likelihood(mean, logs, x):
        """
        lnL = -1/2 * { ln|Var| + ((X - Mu)^T)(Var^-1)(X - Mu) + kln(2*PI) }
              k = 1 (Independent)
              Var = logs ** 2
        """
        return -0.5 * (
            logs * 2.0 + ((x - mean) ** 2) / torch.exp(logs * 2.0) + GaussianDiag.Log2PI
        )

    @staticmethod
    def logp(mean, logs, x):
        likelihood = GaussianDiag.likelihood(mean, logs, x)
        return thops.sum(likelihood, dim=[1])

    @staticmethod
    def sample(output_shape, eps_std=1):
        return torch.normal(
            mean=torch.zeros_like(output_shape),
            std=torch.ones_like(output_shape) * eps_std,
        )


<reponame>whigg/SeaIceVariability
"""
Script plots ranking on meshgrid plot for sea ice volume

Notes
-----
    Source : http://psc.apl.washington.edu/zhang/IDAO/data_piomas.html
    Author : <NAME>
    Date   : 23 September 2016
"""

### Import modules
import numpy as np
import matplotlib.pyplot as plt
import datetime
import read_SeaIceThick_PIOMAS as CT
import read_SeaIceConc_PIOMAS as CC
import calc_PiomasArea as CA
import scipy.stats as sts

### Define directories
directorydata = '/home/zlabe/Surtsey/seaice_obs/PIOMAS/'
directoryfigure = '/home/zlabe/Desktop/'

### Define time
now = datetime.datetime.now()
currentmn = str(now.month)
currentdy = str(now.day)
currentyr = str(now.year)
currenttime = currentmn + '_' + currentdy + '_' + currentyr
titletime = currentmn + '/' + currentdy + '/' + currentyr
print '\n' '----Calculate sea ice volume - %s----' % titletime

### Alott time series
yearmin = 1979
yearmax = 2015
years = np.arange(yearmin,yearmax+1,1)
months = [r'Jan',r'Feb',r'Mar',r'Apr',r'May',r'Jun',r'Jul',r'Aug',
          r'Sep',r'Oct',r'Nov',r'Dec']

### Call functions
#lats,lons,sit = CT.readPiomas(directorydata,years,0.15)
#lats,lons,sic = CC.readPiomas(directorydata,years,0.01)
#area = CA.readPiomasArea(directorydata)

### Calculate siv per year
def sivYear(sit,sic,area,conc):
    """
    Function calculates time series of sea ice volume per YEAR

    Parameters
    ----------
    sit : 4d array [year,month,lat,lon]
        sea ice thickness (m)
    sic : 4d array [year,month,lat,lon]
        sea ice concentration (fraction, 0-1)
    area : 2d array [lat,lon]
        area of grid cell in PIOMAS (km^2)
    conc : boolean
        True or False ----> turning on/off using sea ice concentration

    Returns
    -------
    sivyr : 2d array [year,month]
        sea ice volume

    Usage
    -----
    siv_yr = sivYear(sit,sic,area,conc)
    """

    print '\n>>> Using sivYear function!'

    print 'Calculating sea ice volume'
    if conc == True:
        siv = sit*sic*area
    elif conc == False:
        siv = sit*area
    else:
        RuntimeError('Did not use correct argument for function!')

    ### Take temporal average of year
    siv = np.squeeze(np.apply_over_axes(np.nansum,
                                        siv[:,:,:,:],(2,3)))

    ### Correct units
    siv = siv/10**6

    print '*Completed: Calculated sea ice volume per year!'
    return siv

###########################################################################
###########################################################################
###########################################################################
### Call function
sivyr = sivYear(sit,sic,area,False)
sivyr[np.where(sivyr==0.0)]=np.nan

sivyr = sivyr.transpose()

### Try ranking
rank = np.empty(sivyr.shape)
for i in xrange(sivyr.shape[0]):
    rank[i,:] = sts.rankdata(sivyr[i,:],method='min')

rank = np.flipud(rank)

### Call parameters
plt.rcParams['text.usetex']=True
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.sans-serif'] = 'Avant Garde'
plt.rcParams['xtick.direction'] = 'out'
plt.rcParams['xtick.major.width'] = 1.3
plt.rcParams['xtick.major.size'] = 4

### Plot first meshgrid
fig = plt.figure()
ax = plt.subplot(111)

ax.spines['top'].set_color('none')
ax.spines['right'].set_color('none')
ax.spines['bottom'].set_color('none')
ax.spines['left'].set_color('none')
ax.get_xaxis().set_tick_params(direction='out', width=1.5,length=5)

plt.tick_params(
    axis='x',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='on',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='on')
plt.tick_params(
    axis='y',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    left='off',      # ticks along the bottom edge are off
    right='off',         # ticks along the top edge are off
    labelleft='on')

cs = plt.pcolormesh(rank,shading='faceted',edgecolor='w',
                    linewidth=0.3,vmin=1,vmax=38)

cs.set_cmap('viridis')

cbar = plt.colorbar(cs,orientation='horizontal')
cbar.set_ticks([])
cbar.ax.invert_xaxis()
cbar.set_label(r'\textbf{Sea Ice Volume Rank by Year}')

ylabels = ['D','N','O','S','A','J','J','M','A','M','F','J']
plt.yticks(np.arange(0.5,12.5,1),ylabels,ha='center')
yax = ax.get_yaxis()
yax.set_tick_params(pad=7)
plt.xticks(np.arange(0.5,37.5,3),map(str,np.arange(1979,2016,3)))
plt.xlim([0,37])

plt.text(-3,-5.3,r'Highest rank')
plt.text(34,-5.3,r'Lowest rank')

for i in xrange(rank.shape[0]):
    for j in xrange(rank.shape[1]):
        plt.text(j+0.5,i+0.5,'%s' % int(rank[i,j]),fontsize=6,
                 color='w',va='center',ha='center')

### Save figure
plt.savefig(directoryfigure + 'siv_ranks.png',dpi=300)

###########################################################################
###########################################################################
###########################################################################
### Calculate climo
climyr = np.where((years>=1981) & (years<2010))[0]
sivave = np.nanmean(sivyr[:,climyr],axis=1)

anoms = sivyr.transpose() - sivave

### Plot second meshgrid
fig = plt.figure()
ax = plt.subplot(111)

ax.spines['top'].set_color('none')
ax.spines['right'].set_color('none')
ax.spines['bottom'].set_color('none')
ax.spines['left'].set_color('none')
ax.get_xaxis().set_tick_params(direction='out', width=1.3,length=5)

plt.tick_params(
    axis='x',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='on',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='on')
plt.tick_params(
    axis='y',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    left='off',      # ticks along the bottom edge are off
    right='off',         # ticks along the top edge are off
    labelleft='on')

cs = plt.pcolormesh(np.flipud(anoms.transpose()),shading='faceted',
                    edgecolor='w',linewidth=0.3,clim=np.arange(-10,11,0.2),
                    vmin=-10,vmax=10)

for i in xrange(rank.shape[0]):
    for j in xrange(rank.shape[1]):
        plt.text(j+0.5,i+0.5,'%s' % int(rank[i,j]),fontsize=6,
                 color='k',va='center',ha='center')

cs.set_cmap('RdBu')

cbar = plt.colorbar(cs,orientation='horizontal',shrink=0.7)
cbar.set_ticks(np.arange(-10,11,2))
cbar.set_ticklabels(map(str,np.arange(-10,11,2)))
cbar.set_label(r'\textbf{Sea Ice Volume Anomalies ($\times$1000\ km${^3}$)}',
                         color='k')

#cbar().tick_params(axis='y', direction='out')
#cbarxtks = plt.getp(cbar.ax.axes, 'xticklines')
#plt.setp(cbarxtks, direction='out')

ylabels = ['D','N','O','S','A','J','J','M','A','M','F','J']
plt.yticks(np.arange(0.5,12.5,1),ylabels,ha='center')
yax = ax.get_yaxis()
yax.set_tick_params(pad=7)
plt.xticks(np.arange(0.5,37.5,3),map(str,np.arange(1979,2016,3)))
plt.xlim([0,37])

### Save figure
plt.savefig(directoryfigure + 'siv_anoms.png',dpi=300)


<gh_stars>0
#==============================
# Just started working on that
#==============================
from scipy import interpolate
import pickle
from os import path
import pandas as pd
import numpy as np

class Interpolation():
    def __init__(self):
        self.model_path = path.join('models', 'interpolation.pk')

    def train(self, coor, vals):
        """Training Interpolation model and saving it to pickle format.

        Parameters
        ----------
        coor : ndarray of floats
            Data point coordinates
        vals : ndarray of float
            Data values.
        """
        interp = interpolate.LinearNDInterpolator(coor, vals)
        with open(self.model_path, 'wb') as f:
            pickle.dump(interp, f)

    def predict_grid(self, x_min, x_max, y_min, y_max, step=0.01):
        """Calculating interpolation for the grid created by given x_min_max
           and y_min_max by given step.

        Parameters
        ----------
        x_min : float
            Min grid value on x coordinate
        x_max : float
            Max grid value on x coordinate
        y_min : float
            Min grid value on y coordinate
        y_max : float
            Max grid value on y coordinate
        step : float, optional
            How dense should the grid be, by default 0.01

        Returns
        -------
        df
            Returns pandas dataframe, with cols: x, y
        """
        with open(self.model_path, 'rb') as f:
            interp = pickle.load(f)
        grid = self._generate_grid(x_min, x_max, y_min, y_max, step)
        print('Starting Prediction')
        grid.loc[:, "val"] = grid.apply(lambda x: interp(x[0], x[1]), axis=1)
        return grid

    def _generate_grid(self, x_min, x_max, y_min, y_max, step=0.01):
        x = np.arange(x_min, x_max, step)
        y = np.arange(y_min, y_max, step)
        xv, yv = np.meshgrid(x, y, sparse=False, indexing='ij')
        grid = np.transpose(
            np.array([np.matrix.flatten(xv), np.matrix.flatten(yv)]))
        grid = pd.DataFrame(grid, columns=['x', 'y'])
        return grid

# %% Tests
Tests = False
if Tests is True:
    from src.data.createDB import (Database, db_backend, db_name)
    dbms = Database(db_backend, db_name)
    query = """WITH ranked AS(
        SELECT m.*, ROW_NUMBER() OVER(PARTITION BY id ORDER BY time DESC) as rm
        FROM `airly_raw_data` as m
    ), locs as (SELECT id, lat, lng, city, street FROM 'airly_locs')
    SELECT r.*, l.lat, l.lng
    FROM ranked r LEFT JOIN locs as l on l.id = r.id where rm = 1
    """
    data = dbms.pd_readSQL(sql=query)
    inter = Interpolation()
    inter.train(data.loc[:, ['lat', 'lng']].values, data.loc[:, 'pm10'])
    data_pred = inter.predict_grid(49.8, 50.24, 19.6, 20.25, 0.001)
    print(data_pred.head())


<gh_stars>0
from sympy import *

u = Matrix([4, 3])
v = Matrix([1, 1])
s = u.dot(v)
print(s)


<filename>qmla/exploration_strategies/genetic_algorithms/genetic_exploration_strategy.py
import numpy as np
import itertools
import sys
import os
import random
import copy
import scipy
import time

import pandas as pd
import sklearn as skl

import matplotlib
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import seaborn as sns

try:
    from lfig import LatexFigure
except:
    from qmla.shared_functionality.latex_figure import LatexFigure
from qmla.exploration_strategies import exploration_strategy
import qmla.shared_functionality.probe_set_generation
import qmla.model_building_utilities

import qmla.shared_functionality.genetic_algorithm

__all__ = [
    "Genetic",
    "GeneticTest",
    "GeneticAlgorithmQMLAFullyConnectedLikewisePauliTerms",
]

def hamming_distance(str1, str2):
    return sum(c1 != c2 for c1, c2 in zip(str1, str2))

class Genetic(exploration_strategy.ExplorationStrategy):
    r"""
    Exploration Strategy where the model search is mediated through a genetic algorithm.
    Genetic algorithm is implemented through :class:`qmla.GeneticAlgorithmQMLA`.
    This forms the base class for genetic algorithm applications within QMLA.

    :param str exploration_rules: name of exploration strategy used
    :param list genes: terms which are permitted in the model search,
        which become genes in the chromomsomes of the genetic algorithm
    :param str true_model: name of the target model.

    """

    def __init__(self, exploration_rules, genes, true_model, **kwargs):
        super().__init__(exploration_rules=exploration_rules, **kwargs)

        self.genes = genes
        self.true_model = true_model
        self.ratings_class = qmla.shared_functionality.rating_system.ModifiedEloRating(
            initial_rating=1000, k_const=30
        )  # for use when ranking/rating models

        self.branch_champion_selection_stratgey = "fitness"  # 'ratings'
        self.fitness_method = "elo_rating"
        self.prune_completed_initially = True
        self.prune_complete = True
        self.fitness_by_f_score = pd.DataFrame()
        self.fitness_df = pd.DataFrame()
        self.num_sites = qmla.model_building_utilities.get_num_qubits(self.true_model)
        self.num_probes = 50
        self.max_num_qubits = 7
        self.hypothetical_final_generation = False

        self.qhl_models = [
            "pauliSet_1J2_zJz_d3+pauliSet_1J3_yJy_d3+pauliSet_1J3_zJz_d3+pauliSet_2J3_xJx_d3+pauliSet_2J3_zJz_d3",
            "pauliSet_1J3_yJy_d3+pauliSet_1J3_zJz_d3+pauliSet_2J3_xJx_d3+pauliSet_2J3_zJz_d3",
            "pauliSet_1J2_zJz_d3+pauliSet_1J3_zJz_d3+pauliSet_2J3_xJx_d3+pauliSet_2J3_zJz_d3",
        ]
        self.spawn_step = 0  # 1st generation's ID

        self.mutation_probability = 0.1

        if "log_file" not in kwargs:
            kwargs["log_file"] = self.log_file

        self.genetic_algorithm = (
            qmla.shared_functionality.genetic_algorithm.GeneticAlgorithmQMLA(
                genes=genes,
                num_sites=self.num_sites,
                true_model=self.true_model,
                mutation_probability=self.mutation_probability,
                **kwargs,
            )
        )

        self.true_chromosome = self.genetic_algorithm.true_chromosome
        self.true_chromosome_string = self.genetic_algorithm.true_chromosome_string

        self.num_possible_models = 2 ** len(self.true_chromosome)

        self.max_num_probe_qubits = self.num_sites

        # default test - 32 generations x 16 starters
        self.max_spawn_depth = 24
        self.initial_num_models = 16
        self.initial_models = self.genetic_algorithm.random_initial_models(
            num_models=self.initial_num_models
        )
        self.model_f_scores = {}
        self.model_points_at_step = {}
        self.generation_model_rankings = {}
        self.models_ranked_by_fitness = {}
        self.model_fitness_by_generation = {}
        self.fitness_correlations = {}

        self.tree_completed_initially = False
        self.max_num_models_by_shape = {
            self.num_sites: (self.initial_num_models * self.max_spawn_depth) / 10,
            "other": 0,
        }
        self.num_processes_to_parallelise_over = self.initial_num_models

        self.max_time_to_consider = 15
        self.min_param = 0.35
        self.max_param = 0.65

        self.fitness_mechanism_names = {
            "f_score": r"$F_1$",
            "hamming_distance": r"$H$",
            "inverse_ll": r"$g^L$",
            "inverse_ll_sq": r"$-\frac{1}{L^2}$",
            "akaike_info_criterion": r"$\frac{1}{AIC}$",
            "aic_sq": r"$\frac{1}{AIC^2}$",
            "aicc": r"$\frac{1}{AICc}$",
            "aicc_sq": r"$g^{A}$",
            "bayesian_info_criterion": r"$\frac{1}{BIC}$",
            "bic_sq": r"$g^{B}$",
            "akaike_weight": r"$w_{A}$",
            "bayes_weight": r"$w_{B}$",
            "mean_residuals": r"$r_{\mu}$",
            "mean_residuals_sq": r"$r_{\mu}^2$",
            "rs_mean": r"$1-\overline{r}$",
            "rs_median": r"$1-\tilde{r}$",
            "rs_mean_sq": r"$g^{r}$",  # r"$(1-\overline{r})^2$",
            "rs_median_sq": r"$(1-\tilde{r})^2$",
            "bf_points": r"$g^{p}$",
            "bf_rank": r"$g^{R}$",
            "elo_rating": r"$g^{E}$",
        }

    def nominate_champions(self):
        r"""Choose model with highest fitness on final generation"""

        self.champion_model = self.models_ranked_by_fitness[self.spawn_step][0]
        self.log_print(
            [
                "Final generation:",
                self.spawn_step,
                "\nModel rankings on final generation:",
                self.models_ranked_by_fitness[self.spawn_step],
                "\nChampion:",
                self.champion_model,
            ]
        )

        return [self.champion_model]

    def analyse_generation(self, model_points, model_names_ids, **kwargs):
        r"""
        Following a complete generation of the genetic algorithm,
        perform all necessary processing to enable construction of next set of models.

        :param dict model_points: the number of Bayes factor comparisons for which each candidate
            within the generation was deemed superior against a contemporary model
        :param dict model_names_ids: mapping between models' names and their IDs from the QMLA environment;
            this enables analaysing further data passed from QMLA within kwargs.
        """

        self.spawn_step += 1

        self.log_print(["Analysing generation at spawn step ", self.spawn_step])
        self.log_print(["model names ids:", model_names_ids])
        self.model_points_at_step[self.spawn_step] = model_points

        # model_names_ids = model_names_ids
        sum_wins = sum(list(model_points.values()))
        if sum_wins == 0:
            sum_wins = 1  # TODO hack to get over some times passing empty dict from update_branch -- find a better way
        model_ids = list(model_points.keys())

        # model rankings  by number of wins
        ranked_model_list = sorted(model_points, key=model_points.get, reverse=True)
        ranked_models_by_name = [model_names_ids[m] for m in ranked_model_list]
        self.log_print(
            [
                "Ranked models:",
                ranked_model_list,
                "\n Names:",
                ranked_models_by_name,
                "\n with fitnesses:",
            ]
        )

        self.generation_model_rankings[self.spawn_step] = ranked_models_by_name
        rankings = list(range(1, len(ranked_model_list) + 1))
        rankings.reverse()
        num_points = sum(rankings)  # number of points to distribute
        ranking_points = list(
            zip(ranked_models_by_name, [r / num_points for r in rankings])
        )
        ranking_points = dict(ranking_points)

        # Model ratings  (Elo ratings)
        precomputed_ratings = self.ratings_class.get_ratings(list(model_points.keys()))
        original_ratings_by_name = {
            model_names_ids[m]: precomputed_ratings[m] for m in model_ids
        }
        min_rating = min(original_ratings_by_name.values())
        ratings_by_name = {
            m: original_ratings_by_name[m] - min_rating
            for m in original_ratings_by_name
        }
        self.log_print(["Rating (as fraction of starting rating):\n", ratings_by_name])
        sum_ratings = np.sum(list(ratings_by_name.values()))
        model_elo_ratings = {
            m: ratings_by_name[m] / sum_ratings for m in ratings_by_name
        }

        # New dictionaries which can be used as fitnesses:
        model_f_scores = {"fitness_type": "f_score"}
        model_hamming_distances = {"fitness_type": "hamming_distance"}
        model_number_wins = {"fitness_type": "number_wins"}
        model_win_ratio = {"fitness_type": "win_ratio"}
        mean_residuals = {"fitness_type": "mean_residuals"}
        log_likelihoods = {"fitness_type": "log_likelihoods"}

        # Alter finished dicts also useable as fitness
        # log_likelihoods['fitness_type'] = 'log_likelihoods'
        model_elo_ratings["fitness_type"] = "elo_ratings"
        ranking_points["fitness_type"] = "ranking"

        model_instances = [self.tree.model_storage_instances[m] for m in model_ids]
        aic_values = {
            model.model_id: model.akaike_info_criterion for model in model_instances
        }
        aicc_values = {
            model.model_id: model.akaike_info_criterion_c for model in model_instances
        }
        min_aicc = min(aicc_values.values())
        self.log_print(
            ["At generation {}, AIC of models: {}".format(self.spawn_step, aic_values)]
        )

        # store info on each model for analysis
        for m in model_ids:
            # Access the model storage instance and retrieve some attributes from there
            model_storage_instance = self.tree.model_storage_instances[m]
            self.log_print(["Model storage instance:", model_storage_instance])
            mod = model_storage_instance.model_name
            model_number_wins[mod] = model_points[m]
            hamming_dist = self.hamming_distance_model_comparison(
                test_model=mod
            )  # for fitness use 1/H
            model_hamming_distances[mod] = (
                self.genetic_algorithm.num_terms - hamming_dist
            ) / self.genetic_algorithm.num_terms
            model_f_scores[mod] = np.round(
                self.f_score_model_comparison(test_model=mod), 2
            )  # TODO get from model instance
            self.model_f_scores[m] = model_f_scores[mod]
            model_win_ratio[mod] = model_number_wins[mod] / sum_wins

            # store scores for offline analysis
            this_model_fitnesses = {
                # When adding a new fitness fnc -- add a name in self.fitness_mechanism_names
                "model": mod,
                "model_id": m,
                "generation": self.spawn_step,
                # absolute metrics (not available in real experiments)
                "f_score": model_f_scores[mod],
                "hamming_distance": model_hamming_distances[mod],
                # from storage instance
                # 'eval_log_likelihood' : model_storage_instance.evaluation_log_likelihood,
                "inverse_ll": -1 / model_storage_instance.evaluation_log_likelihood,
                "inverse_ll_sq": (-1 / model_storage_instance.evaluation_log_likelihood)
                ** 2,
                "akaike_info_criterion": 1
                / model_storage_instance.akaike_info_criterion,
                "aicc": 1 / model_storage_instance.akaike_info_criterion_c,
                "aic_sq": (1 / model_storage_instance.akaike_info_criterion) ** 2,
                "aicc_sq": (1 / model_storage_instance.akaike_info_criterion_c) ** 2,
                "bayesian_info_criterion": (
                    1 / model_storage_instance.bayesian_info_criterion
                ),
                "bic_sq": (1 / model_storage_instance.bayesian_info_criterion) ** 2,
                "akaike_weight": np.e
                ** ((min_aicc - model_storage_instance.akaike_info_criterion_c) / 2),
                "bayes_weight": np.e
                ** (-1 * model_storage_instance.bayesian_info_criterion / 2),
                "mean_residuals": 1 - model_storage_instance.evaluation_mean_pr0_diff,
                "mean_residuals_sq": (
                    1 - model_storage_instance.evaluation_mean_pr0_diff
                )
                ** 2,
                "rs_mean": 1
                - model_storage_instance.evaluation_residual_squares["mean"],
                "rs_median": 1
                - model_storage_instance.evaluation_residual_squares["median"],
                "rs_mean_sq": (
                    1 - model_storage_instance.evaluation_residual_squares["mean"]
                )
                ** 2,
                "rs_median_sq": (
                    1 - model_storage_instance.evaluation_residual_squares["median"]
                )
                ** 2,
                # relative to other models in this branch
                "bf_points": model_win_ratio[mod],
                "bf_rank": ranking_points[mod],
                "elo_rating": model_elo_ratings[mod],
                # 'original_elo_rating' : original_ratings_by_name[mod],
            }

            self.fitness_by_f_score = self.fitness_by_f_score.append(
                pd.Series(this_model_fitnesses), ignore_index=True
            )

            recorded_fitness_types = list(
                this_model_fitnesses.keys()
                - [
                    "model",
                    "model_id",
                    "generation",
                    "hamming_distance",
                ]
            )
            for f in recorded_fitness_types:
                try:
                    new_entry = pd.Series(
                        {
                            "generation": this_model_fitnesses["generation"],
                            "f_score": this_model_fitnesses["f_score"],
                            "fitness": this_model_fitnesses[f],
                            "fitness_type": f,
                            "fitness_type_name": self.fitness_mechanism_names[f],
                            "active_fitness_method": self.fitness_method == f,
                        }
                    )
                    self.fitness_df = self.fitness_df.append(
                        new_entry, ignore_index=True
                    )
                except:
                    self.log_print(
                        [
                            "fitness name keys:",
                            list(self.fitness_mechanism_names.keys())
                            # "f={}; type name = {}".format(f, self.fitness_mechanism_names[f])
                        ]
                    )
                    raise

        # Extract fitness specified by user (exploration strategy's fitness_method attribute)
        # to use for generating models within genetic algorithm
        fitnesses = self.fitness_by_f_score[
            self.fitness_by_f_score.generation == self.spawn_step
        ][["model", self.fitness_method]]

        genetic_algorithm_fitnesses = dict(
            zip(fitnesses["model"], fitnesses[self.fitness_method])
        )

        self.log_print(
            [
                "fitness method:{} => Fitnesses={}".format(
                    self.fitness_method, genetic_algorithm_fitnesses
                )
            ]
        )
        self.models_ranked_by_fitness[self.spawn_step] = sorted(
            genetic_algorithm_fitnesses,
            key=genetic_algorithm_fitnesses.get,
            reverse=True,
        )
        self.model_fitness_by_generation[self.spawn_step] = genetic_algorithm_fitnesses

        self.genetic_algorithm.consolidate_generation(
            model_fitnesses=genetic_algorithm_fitnesses
        )

        # return genetic_algorithm_fitnesses
        return self.models_ranked_by_fitness[self.spawn_step]

    def generate_models(self, model_list, **kwargs):
        r"""
        Model generation using genetic algorithm.

        Follows rules of :meth:`~qmla.exploration_strategies.ExplorationStrategy.generate_models`.
        """

        # Analysis of the previous generation is called by the exploration strategy tree.
        genetic_algorithm_fitnesses = self.model_fitness_by_generation[self.spawn_step]

        self.log_print(
            [
                "Spawn step:",
                self.spawn_step,
            ]
        )

        # Spawn models from genetic algorithm
        new_models = self.genetic_algorithm.genetic_algorithm_step(
            model_fitnesses=genetic_algorithm_fitnesses,
            num_pairs_to_sample=self.initial_num_models
            / 2,  # for every pair, 2 chromosomes proposed
        )

        return new_models

    def finalise_model_learning(self, **kwargs):
        return

    def hamming_distance_model_comparison(
        self,
        test_model,
        target_model=None,
    ):
        r"""
        Compare test_model with target_model by Hamming distance
        """

        if target_model is None:
            target_model = self.true_chromosome_string
        else:
            target_model = self.genetic_algorithm.chromosome_string(
                self.genetic_algorithm.map_model_to_chromosome(target_model)
            )
        test_model = self.genetic_algorithm.chromosome_string(
            self.genetic_algorithm.map_model_to_chromosome(test_model)
        )

        h = sum(c1 != c2 for c1, c2 in zip(test_model, target_model))
        return h

    def f_score_model_comparison(
        self,
        test_model,
        target_model=None,
        beta=1,
    ):
        r"""
        Get F score of candidate model, measure of overlap between the terms of the candidate and target model

        :param str test_model: name of candidate model
        :param str target_model: name of target model, if None, assumed that target is self.true_model
        :param float beta: relative importance of precision to sensitivity. in general this is F-beta score,
            usually beta = 1
        """
        if target_model is None:
            target_model = self.true_model

        true_set = set(
            self.latex_name(mod)
            for mod in qmla.model_building_utilities.get_constituent_names_from_name(
                target_model
            )
        )
        terms = [
            self.latex_name(term)
            for term in qmla.model_building_utilities.get_constituent_names_from_name(
                test_model
            )
        ]
        learned_set = set(sorted(terms))

        total_positives = len(true_set)
        true_positives = len(true_set.intersection(learned_set))
        false_positives = len(learned_set - true_set)
        false_negatives = len(true_set - learned_set)
        precision = true_positives / (true_positives + false_positives)
        sensitivity = true_positives / total_positives
        try:
            f_score = (1 + beta ** 2) * (
                (precision * sensitivity) / (beta ** 2 * precision + sensitivity)
            )
        except BaseException:
            # both precision and sensitivity=0 as true_positives=0
            f_score = 0
        return f_score

    def f_score_from_chromosome_string(
        self,
        chromosome,
    ):
        r"""
        F1 score between chromosome and true model
        """

        mod = np.array([int(a) for a in list(chromosome)])

        try:
            f = skl.metrics.f1_score(mod, self.true_chromosome)
            return f
        except:
            self.log_print(
                [
                    "F score from chromosome {} with mod {} not working against true chrom {}".format(
                        mod, chromosome, self.true_chromosome
                    )
                ]
            )
            raise

    def exploration_strategy_finalise(self):
        r"""
        Genetic algorithm specific version of :meth:`qmla.ExplorationStrategy.exploration_strategy_finalise`.
        """

        # hypothetical generation_models
        if self.hypothetical_final_generation:
            # TODO this will cause a crash in QHL mode since.
            # in general this should be turned off so not worth a large fix
            self.log_print(["Running hypothetical step to get some models"])
            hypothetical_models = self.genetic_algorithm.genetic_algorithm_step(
                model_fitnesses=self.model_fitness_by_generation[self.spawn_step - 1],
                num_pairs_to_sample=self.initial_num_models
                / 2,  # for every pair, 2 chromosomes proposed
            )
            self.log_print(["hypothetical generation models:", hypothetical_models])

        self.storage.fitness_correlations = self.fitness_correlations
        self.storage.fitness_by_f_score = self.fitness_by_f_score
        self.storage.fitness_df = self.fitness_df
        self.storage.true_model_chromosome = self.true_chromosome_string
        self.storage.ratings_df = self.ratings_class.ratings_df
        gene_pool = self.genetic_algorithm.gene_pool
        gene_pool["objective_function"] = self.fitness_mechanism_names[
            self.fitness_method
        ]
        self.storage.gene_pool = gene_pool
        birth_register = self.genetic_algorithm.birth_register
        birth_register["objective_function"] = self.fitness_mechanism_names[
            self.fitness_method
        ]
        birth_register["max_time_considered"] = self.max_time_to_consider
        self.storage.birth_register = birth_register
        self.storage.ratings = self.ratings_class.all_ratings

        chromosomes = sorted(
            list(set(self.genetic_algorithm.previously_considered_chromosomes))
        )
        self.unique_chromosomes = pd.DataFrame(
            columns=[
                "chromosome",
                "numeric_chromosome",
                "f_score",
                "num_terms",
                "hamming_distance",
            ]
        )
        for c in chromosomes:
            hamming_dist = self.hamming_distance_model_comparison(
                test_model=self.genetic_algorithm.map_chromosome_to_model(c)
            )  # for fitness use 1/H

            chrom_data = pd.Series(
                {
                    "chromosome": str(c),
                    "numeric_chromosome": int(c, 2),
                    "num_terms": self.genetic_algorithm.num_terms,
                    "hamming_distance": hamming_dist,
                    "f_score": np.round(self.f_score_from_chromosome_string(c), 3),
                }
            )
            self.unique_chromosomes.loc[len(self.unique_chromosomes)] = chrom_data
        self.log_print(["self.unique_chromosomes:\n", self.unique_chromosomes])
        self.storage.unique_chromosomes = self.unique_chromosomes

        dud_chromosome = str("1" + "0" * self.genetic_algorithm.num_terms)
        if dud_chromosome in chromosomes:
            self.log_print(
                [
                    "{} in previous chromosomes:\n{}".format(
                        dud_chromosome,
                        self.genetic_algorithm.previously_considered_chromosomes,
                    )
                ]
            )
        chromosome_numbers = sorted([int(c, 2) for c in chromosomes])
        # self.exploration_strategy_specific_data_to_store['chromosomes_tested'] = chromosome_numbers
        try:
            f_scores = []
            for c in chromosomes:
                try:
                    f_scores.append(np.round(self.f_score_from_chromosome_string(c), 3))
                except:
                    self.log_print(
                        ["Could not compute f score for chromosome: {}".format(c)]
                    )
            # self.exploration_strategy_specific_data_to_store['f_score_tested_models' ] = f_scores
        except:
            self.log_print(
                [
                    "Could not compute f score for chromosome list: {}".format(
                        chromosomes
                    )
                ]
            )
            pass

        self.storage.chromosomes_tested = chromosome_numbers
        self.storage.f_score_tested_models = f_scores

    def check_tree_completed(self, spawn_step, **kwargs):
        r"""
        Genetic algorithm specific version of :meth:`qmla.ExplorationStrategy.check_tree_completed`.
        """

        if self.spawn_step == self.max_spawn_depth:
            self.log_print(["Terminating at spawn depth ", self.spawn_step])
            return True
        elif self.genetic_algorithm.best_model_unchanged:
            self.champion_determined = True
            self.champion_model = (
                self.genetic_algorithm.most_elite_models_by_generation[
                    self.genetic_algorithm.genetic_generation - 1
                ]
            )

            self.log_print(
                [
                    "Terminating search early (after {} generations) b/c elite model unchanged in {} generations.".format(
                        self.genetic_algorithm.genetic_generation,
                        self.genetic_algorithm.unchanged_elite_num_generations_cutoff,
                    ),
                    "\nDeclaring champion:",
                    self.champion_model,
                ]
            )
            # check if elite model hasn't changed in last N generations
            return True
        else:
            self.log_print(["Elite models changed recently; continuing search."])
            return False

    def check_tree_pruned(self, **kwargs):
        r"""
        Genetic algorithm specific version of :meth:`qmla.ExplorationStrategy.check_tree_pruned`.
        """
        # no pruning for GA, winner is champion of final branch
        return True

    def set_specific_plots(self, **kwargs):
        r"""
        Genetic algorithm specific version of :meth:`qmla.ExplorationStrategy.set_specific_plots`.
        """

        self.plot_methods_by_level = {
            1: [],
            2: [
                self._plot_correlation_fitness_with_f_score,
                self._plot_fitness_v_fscore_by_generation,
                self.__plot_gene_pool_progression,
            ],
            3: [
                self._plot_fitness_v_fscore,
                self._plot_fitness_v_generation,
            ],
            4: [
                self._plot_model_ratings,
                self._plot_gene_pool,
            ],
            5: [self.plot_generational_metrics, self._plot_selection_probabilities],
            6: [],
        }

        # Plots that need arguments so are called individually
        if self.plot_level >= 2:
            try:
                self.ratings_class.plot_models_ratings_against_generation(
                    f_scores=self.model_f_scores,
                    save_directory=self.save_directory,
                    f_score_cmap=self.f_score_cmap,
                    figure_format=self.figure_format,
                )
            except Exception as e:
                self.log_print(
                    [
                        "plot failed plot_models_ratings_against_generation with error ",
                        e,
                    ]
                )

            try:
                self.ratings_class.plot_rating_progress_single_model(
                    target_model_id=champion_model_id,
                    save_to_file=os.path.join(
                        self.save_directory, "ratings_progress_champion.png"
                    ),
                )
                if true_model_id != -1 and true_model_id != champion_model_id:
                    self.ratings_class.plot_rating_progress_single_model(
                        target_model_id=true_model_id,
                        save_to_file=os.path.join(
                            save_directory, "ratings_progress_true_model.png"
                        ),
                    )
            except Exception as e:
                self.log_print(
                    ["plot failed plot_rating_progress_single_model with error ", e]
                )

    def _plot_correlation_fitness_with_f_score(
        self,
        save_to_file=None,
    ):
        r"""
        Show how the fitness of models at each generation progress in terms of F score.
        """

        plt.clf()
        correlations = pd.DataFrame(columns=["Generation", "Method", "Correlation"])
        fitness_types_to_ignore = ["f_score", "hamming_distance"]
        for t in self.fitness_df.fitness_type.unique():
            if t not in fitness_types_to_ignore:
                this_fitness_type = self.fitness_df[
                    self.fitness_df["fitness_type"] == t
                ]

                for g in this_fitness_type.generation.unique():
                    this_type_this_gen = this_fitness_type[
                        this_fitness_type.generation == g
                    ]

                    corr = this_type_this_gen["f_score"].corr(
                        this_type_this_gen["fitness"]
                    )
                    cov = this_type_this_gen["f_score"].cov(
                        this_type_this_gen["fitness"]
                    )

                    corr = {
                        "Generation": g,
                        "Method": self.fitness_mechanism_names[t],
                        # 'Method' : t,
                        "Correlation": corr,
                        "Covariance": cov,
                    }
                    correlations = correlations.append(
                        pd.Series(corr), ignore_index=True
                    )

        self.fitness_correlations = correlations
        self.log_print(["fitness correlations:\n", self.fitness_correlations])
        fig, ax = plt.subplots(figsize=(15, 10))

        if len(correlations.Generation.unique()) == 1:
            sns.scatterplot(
                y="Correlation",
                x="Generation",
                # style= 'Method',
                hue="Method",
                data=correlations,
                ax=ax,
                # markers = ['*', 'X', '<', '^'],
            )
        else:
            sns.lineplot(
                y="Correlation",
                x="Generation",
                # style= 'Method',
                hue="Method",
                data=correlations,
                ax=ax,
                markers=["*", "X", "<", "^"],
            )
        ax.axhline(0, ls="--", c="k")

        if save_to_file is None:
            save_to_file = os.path.join(
                self.save_directory,
                "correlations_bw_fitness_and_f_score.png".format(self.qmla_id),
            )

        plt.savefig(save_to_file)

    def _plot_fitness_v_generation(self, save_to_file=None):
        r"""
        Plot progression of fitness against generations of the genetic algorithm.
        """
        import matplotlib.pyplot as plt
        import seaborn as sns

        plt.clf()
        fig, ax = plt.subplots()
        sns.set(rc={"figure.figsize": (11.7, 8.27)})

        cmap = sns.cubehelix_palette(dark=0.3, light=0.8, as_cmap=True)
        sns.boxplot(
            x="generation",
            y="fitness",
            data=self.fitness_df[
                # self.fitness_df['fitness_type'] == 'model_hamming_distances'
                self.fitness_df["active_fitness_method"]
                == True
            ],
            ax=ax,
        )
        ax.legend(loc="lower right")
        ax.set_xlabel("Generation")
        ax.set_ylabel("Fitness")
        ax.set_title("Fitness method: {}".format(self.fitness_method))
        # ax.set_xlim((0,1))
        if save_to_file is None:
            save_to_file = os.path.join(
                self.save_directory,
                "fitness_v_generation.{}".format(self.figure_format),
            )

        plt.savefig(save_to_file)

    def _plot_fitness_v_fscore_by_generation(
        self,
    ):
        r"""
        Plot fitness vs f score throughout generations of the genetic algorithm.
        """

        plt.clf()
        sanity_check_df = self.fitness_df[
            (self.fitness_df["fitness_type"] == "f_score")
            | (self.fitness_df["fitness_type"] == "model_hamming_distances")
        ]
        candidate_fitnesses = self.fitness_df[
            (self.fitness_df["fitness_type"] == "elo_rating")
            | (self.fitness_df["fitness_type"] == "ranking")
            | (self.fitness_df["fitness_type"] == "model_win_ratio")
        ]

        g = sns.FacetGrid(
            candidate_fitnesses,
            row="generation",
            hue="fitness_type",
            hue_kws=dict(marker=["x", "+", "*"]),
            # col_wrap=5,
            xlim=(-0.1, 1.1),
            # ylim=(0,1),
            size=4,
            aspect=2,
        )
        g = g.map(plt.scatter, "f_score", "fitness").add_legend()

        save_to_file = os.path.join(
            self.save_directory, "fitness_types.{}".format(self.figure_format)
        )
        plt.savefig(save_to_file)

    def _plot_model_ratings(
        self,
    ):
        r"""
        Plot ratings of models on all generations, as determined by the RatingSystem
        """

        plt.clf()
        ratings = self.ratings_class.all_ratings
        generations = [int(g) for g in ratings.generation.unique()]
        num_generations = len(generations)

        lf = LatexFigure(use_gridspec=True, gridspec_layout=(num_generations, 1))

        # TODO : unique linestyle and colour combo for each model ID and tracks across subplots
        ratings["Model ID"] = ratings["model_id"]

        for gen in generations:
            ax = lf.new_axis()

            this_gen_ratings = ratings[ratings.generation == gen]
            colours = {
                m: self.f_score_cmap(self.model_f_scores[m])
                for m in this_gen_ratings["model_id"]
            }
            sns.lineplot(
                x="idx",
                y="rating",
                hue=r"Model ID",
                hue_order=sorted(this_gen_ratings.model_id.unique()),
                data=this_gen_ratings,
                ax=ax,
                legend="full",
                palette=colours,
            )

            ax.set_title("Generation {}".format(gen), pad=-15)
            ax.set_xlabel("")
            ax.set_ylabel("Elo rating")
            ax.legend(bbox_to_anchor=(1, 1))

        save_to_file = os.path.join(self.save_directory, "ratings".format(self.qmla_id))

        lf.save(save_to_file, file_format=self.figure_format)

    def _plot_fitness_v_fscore(self):
        r"""
        Plot fitness against f score
        """

        plt.clf()
        fig, ax = plt.subplots()
        sns.set(rc={"figure.figsize": (11.7, 8.27)})

        cmap = sns.cubehelix_palette(dark=0.3, light=0.8, as_cmap=True)
        sns.scatterplot(
            x="f_score",
            y="elo_rating",
            # hue='generation',
            # palette = cmap,
            label="Rating",
            data=self.fitness_by_f_score,
            ax=ax,
        )

        sns.scatterplot(
            x="f_score",
            y="win_ratio",
            # hue='generation',
            # palette = cmap,
            label="Win ratio",
            data=self.fitness_by_f_score,
            ax=ax,
        )

        ax.legend(loc="lower right")
        ax.set_xlabel("F score")
        ax.set_ylabel("Fitness (as probability)")
        # bplot.set_ylim((0,1))
        ax.set_xlim((-0.05, 1.05))
        save_to_file = os.path.join(
            self.save_directory, "fitness_v_fscore.png".format(self.qmla_id)
        )

        ax.figure.savefig(save_to_file)

    def _plot_gene_pool(self):
        r"""
        Show the F scores of all models in all generations
        """
        ga = self.genetic_algorithm

        plt.clf()
        fig, axes = plt.subplots(
            figsize=(10, 8),
            constrained_layout=True,
        )

        gs = GridSpec(nrows=2, ncols=1, height_ratios=[7, 1])
        label_fontsize = 10
        # TODO get f score cmap from exploration strategy
        # f_score_cmap = matplotlib.colors.ListedColormap(["sienna", "red", "darkorange", "gold", "blue"])
        f_score_cmap = self.f_score_cmap

        # Bar plots for probability of gene being selected, coloured by f score
        ax = fig.add_subplot(gs[0, 0])

        generations = list(sorted(ga.gene_pool.generation.unique()))
        probability_grouped_by_f_by_generation = {
            g: {
                f: ga.gene_pool[
                    (ga.gene_pool.f_score == f) & (ga.gene_pool.generation == g)
                ].probability.sum()
                for f in ga.gene_pool.f_score.unique()
            }
            for g in generations
        }
        probability_grouped_by_f_by_generation = pd.DataFrame(
            probability_grouped_by_f_by_generation
        ).T

        sorted_f_scores = list(sorted(ga.gene_pool.f_score.unique()))
        below = [0] * len(generations)
        for f in sorted_f_scores[:]:
            probs_this_f = list(probability_grouped_by_f_by_generation[f])
            ax.bar(
                generations,
                probs_this_f,
                color=f_score_cmap(f),
                bottom=below,
                edgecolor=["black"] * len(generations),
            )

            below = [b + p for b, p in zip(below, probs_this_f)]
        ax.set_xticks(generations)
        ax.set_ylabel("Probability", fontsize=label_fontsize)
        ax.set_xlabel("Generation", fontsize=label_fontsize)
        ax.set_title("Gene pool", fontsize=label_fontsize)

        # Colour bar
        ax = fig.add_subplot(gs[1, 0])
        sm = plt.cm.ScalarMappable(
            cmap=f_score_cmap, norm=plt.Normalize(vmin=0, vmax=1)
        )
        sm.set_array(np.linspace(0, 1, 100))
        plt.colorbar(sm, cax=ax, orientation="horizontal")
        ax.set_xlabel("F-score", fontsize=label_fontsize)

        # Save figure
        save_to_file = os.path.join(
            self.save_directory, "gene_pool.{}".format(self.figure_format)
        )

        fig.savefig(save_to_file)

    def _plot_selection_probabilities(self):
        r"""
        Plot pie charts of the selection probabilities of prospective parents at each generation.
        Models are signified by their F score.
        """
        generations = sorted(self.genetic_algorithm.gene_pool.generation.unique())
        self.log_print(["[_plot_selection_probabilities] generations:", generations])
        lf = LatexFigure(auto_gridspec=len(generations))

        for g in generations:
            ax = lf.new_axis()
            this_gen_genes = self.genetic_algorithm.gene_pool[
                self.genetic_algorithm.gene_pool.generation == g
            ]
            f_scores = this_gen_genes.f_score
            colours = [self.f_score_cmap(f) for f in f_scores]
            probabilities = this_gen_genes.probability

            ax.pie(
                probabilities,
                colors=colours,
                radius=2,
            )

        save_to_file = os.path.join(self.save_directory, "selection_probabilities")
        lf.save(save_to_file, figure_format=self.figure_format)

    def plot_generational_metrics(self):
        r"""
        Show various metrics across all generations
        """
        fig, axes = plt.subplots(figsize=(15, 10), constrained_layout=True)
        gs = GridSpec(
            nrows=2,
            ncols=1,
        )

        ax = fig.add_subplot(gs[0, 0])
        sns.boxplot(y="f_score", x="generation", data=self.fitness_by_f_score, ax=ax)
        ax.set_ylabel("F-score")
        ax.set_xlabel("Generation")
        ax.set_title("F score")
        ax.set_ylim(0, 1)
        ax.legend()

        ax = fig.add_subplot(gs[1, 0])
        sns.boxplot(
            y="log_likelihood", x="generation", data=self.fitness_by_f_score, ax=ax
        )
        ax.set_ylabel("log-likelihood")
        ax.set_xlabel("Generation")
        ax.set_title("Evaluation log likeihood")
        ax.legend()

        # Save figure
        save_to_file = os.path.join(self.save_directory, "generation_progress.png")

        fig.savefig(save_to_file)

    def __plot_gene_pool_progression(
        self,
    ):
        r"""
        Succinct representation of the progression of gene pool with respect to F score.
        """
        lf = LatexFigure()
        ax = lf.new_axis()
        gene_pool = self.genetic_algorithm.gene_pool
        gene_pool.sort_values("f_score", inplace=True, ascending=False)

        self.gene_pool_progression(
            gene_pool=gene_pool,
            ax=ax,
            f_score_cmap=self.f_score_cmap,
        )
        lf.save(
            save_to_file=os.path.join(self.save_directory, "gene_pool_progression"),
            file_format=self.figure_format,
        )

    @staticmethod
    def gene_pool_progression(
        gene_pool, ax, f_score_cmap=None, draw_cbar=True, cbar_ax=None
    ):
        r"""
        Method for plotting succinct summary of progression of gene pool with respect to F score.
        """
        if f_score_cmap is None:
            f_score_cmap = matplotlib.cm.RdBu
        num_models_per_generation = len(gene_pool[gene_pool.generation == 1])
        num_generations = gene_pool.generation.nunique()
        f_scores_of_gene_pool = np.empty((num_models_per_generation, num_generations))
        for g in gene_pool.generation.unique():

            f_scores_by_gen = gene_pool[gene_pool.generation == g].f_score

            f_scores_of_gene_pool[:, g - 1] = f_scores_by_gen

        sns.heatmap(
            f_scores_of_gene_pool,
            cmap=f_score_cmap,
            vmin=0,
            vmax=1,
            ax=ax,
            cbar=draw_cbar,
            cbar_kws=dict(
                label=r"$F_1$-score",
                aspect=25,
                ticks=[0, 0.5, 1],
            ),
        )
        ax.set_yticks([])
        xtick_pos = range(5, num_generations + 1, 5)
        ax.set_xticks([g - 0.5 for g in xtick_pos])
        ax.set_xticklabels(xtick_pos)
        ax.set_xlabel("Generation")

        if cbar_ax is not None:
            cbar = ax.collections[0].colorbar
            cbar.ax.set_ylabel(r"$F_1$", rotation=0, labelpad=10)  # if F horizontal
            cbar.ax.yaxis.set_label_position(
                "right",
            )
            cbar.ax.tick_params(labelleft=True, labelright=False)

class GeneticTest(Genetic):
    r"""
    Exactly as the genetic exploration strategy, but small depth to test quickly.

    """

    def __init__(self, exploration_rules, **kwargs):
        true_model = "pauliSet_1J2_zJz_d4+pauliSet_1J3_zJz_d4+pauliSet_2J3_zJz_d4+pauliSet_2J4_zJz_d4+pauliSet_3J4_zJz_d4"
        self.true_model = qmla.model_building_utilities.alph(true_model)
        num_sites = qmla.model_building_utilities.get_num_qubits(true_model)
        terms = []
        for i in range(1, 1 + num_sites):
            for j in range(i + 1, 1 + num_sites):
                for t in ["x", "y", "z"]:
                    new_term = "pauliSet_{i}J{j}_{o}J{o}_d{N}".format(
                        i=i,
                        j=j,
                        o=t,
                        N=num_sites,
                    )
                    terms.append(new_term)

        super().__init__(
            exploration_rules=exploration_rules,
            genes=terms,
            true_model=self.true_model,
            **kwargs,
        )
        self.max_spawn_depth = 2
        self.max_num_probe_qubits = self.num_sites
        self.initial_num_models = 6
        self.initial_models = self.genetic_algorithm.random_initial_models(
            num_models=self.initial_num_models
        )
        self.tree_completed_initially = False
        self.max_num_models_by_shape = {
            self.num_sites: (self.initial_num_models * self.max_spawn_depth) / 10,
            "other": 0,
        }
        self.num_processes_to_parallelise_over = self.initial_num_models

class GeneticAlgorithmQMLAFullyConnectedLikewisePauliTerms(Genetic):
    r"""
    Exact structure of :class:`~qmla.Genetic`, where the avaiable terms
    are assumed to follow conventional pauliSet format, and all sites are connected.
    e.g. terms of the form:
    pauliSet_1J2_xJx_d2, pauliSet_1J2_yJy_d2, pauliSet_1J2_zJz_d2,
    """

    def __init__(
        self,
        exploration_rules,
        true_model,
        num_sites=None,
        base_terms=["x", "y", "z"],
        **kwargs
    ):
        if num_sites is None:
            num_sites = qmla.model_building_utilities.get_num_qubits(true_model)
        terms = []
        for i in range(1, 1 + num_sites):
            for j in range(i + 1, 1 + num_sites):
                for t in base_terms:
                    new_term = "pauliSet_{i}J{j}_{o}J{o}_d{N}".format(
                        i=i,
                        j=j,
                        o=t,
                        N=num_sites,
                    )
                    terms.append(new_term)

        super().__init__(
            exploration_rules=exploration_rules,
            genes=terms,
            true_model=true_model,
            **kwargs,
        )


<filename>A6/A6Part4.py<gh_stars>0
import os
import sys
import numpy as np
import math
from scipy.signal import get_window
import matplotlib.pyplot as plt

sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../../software/models/'))
import utilFunctions as UF
import harmonicModel as HM
import dftModel as DFT
import stft

eps = np.finfo(float).eps

"""
A6Part4 - Improving the implementation of the two way mismatch f0 estimation algorithm

Improve the performance of the current implementation of the two way mismatch algorithm in sms-tools
used for fundamental frequency estimation. This is an optional open question and will not contribute
towards the final grade. There is no definite answer for this question. Its main purpose is to
understand the limitations of the current implementations of the TWM algorithm and to come up with
some community driven solutions based on collective thinking.

In this question you will directly modify the core functions that implement the TWM algorithm in
sms-tools. To assist you with this task, we have copied all the needed functions into this python
file. Hence, you just need to modify the functions in this file and not anywhere else.

Estimating fundamental frequency from an audio signal is still a challenging and unsolved problem
to a large extent. By this time you might have also realized that many times the performance of the
TWM f0 estimation algorithm falls short of the expectations. There can be a systematic explanation
for the scenarios where TWM fails for specific categories or characteristics of the sounds. Some of
the known scenarios where the current implementation of the TWM algorithm fails to estimate a correct
fundamental frequency are:

1) Missing fundamental frequency: For many sounds the fundamental frequency component is very low and
therefore during the spectral peak picking step we do not obtain any peak corresponding to the f0.
Since the TWM algorithm implemented in sms-tools considers only the detected spectral peaks as the
f0 candidates, we do not get any candidate corresponding to the f0. This causes f0 estimation to fail.
For example, such a scenario is encountered in low pitched vocal sounds.

2) Pseudo-harmonicity in the sound. Many instruments such as piano exhibit some deviation from perfect
harmonicity wherein their harmonic partials are not perfectly located at integral multiples of the
fundamental frequency. Since the TWM algorithm computes error function assuming that the harmonic
locations are at integral multiples, its performance is poorer when such deviations exist.

In this question we propose to work on these two scenarios. Go to freesound and download sound examples
of low pitched vocal sounds and of piano. Run current implementation of TMW to identify the limitations
and propose improvements to the code in order to obtain better f0 estimation for those two particular
scenarios.

The core TWM algorithm is implemented in the function TWM_p(), which takes in an array of f0 candidates
and detect the candidate that has the lowest error. TWM_p() is called by f0Twm(), which generates
f0 candidates (f0c = np.argwhere((pfreq>minf0) & (pfreq<maxf0))[:,0]). This function also implements
a memory based prunning of the f0 candidates. If the f0 contour is found to be stable (no drastic
transitions across frames) then only the f0 candidates close to the stable f0 value are retained.
f0Twm() is called for every audio frame by f0Detection().

You can use computeAndPlotF0(), which calls f0Detection() for estimating f0 for every audio frame.
In addition, it also plots the f0 contour on the top of the spectrogram. If you set plot=1, it shows
the plot, plot=2 saves the plot as can be seen in the code.

Once you implement your proposed enhancement, discuss and share your ideas on the discussion forum
assigned for A6Part4 - https://class.coursera.org/audio-001/forum/list?forum_id=10026. Along with the
text you should include 2 plots showing the f0 contour before and after your changes. Use the same
values of the analysis parameters while showing the improvement in the performance. in the discussion,
also include a link to the sound in freesound.

TIP: An identified limitation of the current implementation for the case of low vocal sounds is that
it can only find f0 if there is a peak present in the magnitude spectrum. A possible improvement is
to generate additional f0 candidates from the identified peaks. Another identified limitation for
the case of piano sounds is the assumption of perfect harmonicity. For these sounds you can think
of modifying the generation of the ideal harmonic series that is computed in the code, incorporating
the typical deviation from harmonicity encountered in piano sounds.

NOTE: Before you start making changes in the TWM implementation make sure you have reached the best
possible performance that can be achieved by tuning the analysis parameters. If the analysis parameters
are inappropriately set, it is not completely meaningful to just improve the TWM implementation.

To maintain the integrity if the sms-tools package for future assignments, please make changes only
to the functions in this file and not the other files in sms-tools.
"""

def computeAndPlotF0(inputFile = '../../sounds/piano.wav'):
    """
    Function to estimate fundamental frequency (f0) in an audio signal using TWM.
    Input:
        inputFile (string): wav file including the path
    """
    window='hamming'
    M=2048
    N=2048
    H=256
    f0et=5.0
    t=-80
    minf0=100
    maxf0=300

    fs, x = UF.wavread(inputFile)                               #reading inputFile
    w  = get_window(window, M)                                  #obtaining analysis window
    f0 = f0Detection(x, fs, w, N, H, t, minf0, maxf0, f0et)  #estimating F0

    ## Code for plotting the f0 contour on top of the spectrogram
    # frequency range to plot
    maxplotfreq = 500.0
    fontSize = 16
    plot = 1

    fig = plt.figure()
    ax = fig.add_subplot(111)

    mX, pX = stft.stftAnal(x, w, N, H)                      #using same params as used for analysis
    mX = np.transpose(mX[:,:int(N*(maxplotfreq/fs))+1])

    timeStamps = np.arange(mX.shape[1])*H/float(fs)
    binFreqs = np.arange(mX.shape[0])*fs/float(N)

    plt.pcolormesh(timeStamps, binFreqs, mX)
    plt.plot(timeStamps, f0, color = 'k', linewidth=1.5)

    plt.autoscale(tight=True)
    plt.ylabel('Frequency (Hz)', fontsize = fontSize)
    plt.xlabel('Time (s)', fontsize = fontSize)
    plt.legend(('f0',))

    xLim = ax.get_xlim()
    yLim = ax.get_ylim()
    ax.set_aspect((xLim[1]-xLim[0])/(2.0*(yLim[1]-yLim[0])))

    if plot == 1:
        plt.autoscale(tight=True)
        plt.show()
    elif plot == 2:                   #you can save the plot too!
        fig.tight_layout()
        fig.savefig('f0_over_Spectrogram.png', dpi=150, bbox_inches='tight')

def f0Detection(x, fs, w, N, H, t, minf0, maxf0, f0et):
    """
    Fundamental frequency detection of a sound using twm algorithm
    x: input sound; fs: sampling rate; w: analysis window;
    N: FFT size; t: threshold in negative dB,
    minf0: minimum f0 frequency in Hz, maxf0: maximim f0 frequency in Hz,
    f0et: error threshold in the f0 detection (ex: 5),
    returns f0: fundamental frequency
    """
    if (minf0 < 0):                                            # raise exception if minf0 is smaller than 0
        raise ValueError("Minumum fundamental frequency (minf0) smaller than 0")

    if (maxf0 >= 10000):                                       # raise exception if maxf0 is bigger than fs/2
        raise ValueError("Maximum fundamental frequency (maxf0) bigger than 10000Hz")

    if (H <= 0):                                               # raise error if hop size 0 or negative
        raise ValueError("Hop size (H) smaller or equal to 0")

    hN = N/2                                                   # size of positive spectrum
    hM1 = int(math.floor((w.size+1)/2))                        # half analysis window size by rounding
    hM2 = int(math.floor(w.size/2))                            # half analysis window size by floor
    x = np.append(np.zeros(hM2),x)                             # add zeros at beginning to center first window at sample 0
    x = np.append(x,np.zeros(hM1))                             # add zeros at the end to analyze last sample
    pin = hM1                                                  # init sound pointer in middle of anal window
    pend = x.size - hM1                                        # last sample to start a frame
    fftbuffer = np.zeros(N)                                    # initialize buffer for FFT
    w = w / sum(w)                                             # normalize analysis window
    f0 = []                                                    # initialize f0 output
    f0t = 0                                                    # initialize f0 track
    f0stable = 0                                               # initialize f0 stable
    while pin<pend:
        x1 = x[pin-hM1:pin+hM2]                                  # select frame
        mX, pX = DFT.dftAnal(x1, w, N)                           # compute dft
        ploc = UF.peakDetection(mX, t)                           # detect peak locations
        iploc, ipmag, ipphase = UF.peakInterp(mX, pX, ploc)      # refine peak values
        ipfreq = fs * iploc/N                                    # convert locations to Hez
        f0t = f0Twm(ipfreq, ipmag, f0et, minf0, maxf0, f0stable)  # find f0
        if ((f0stable==0)&(f0t>0)) \
                or ((f0stable>0)&(np.abs(f0stable-f0t)<f0stable/5.0)):
            f0stable = f0t                                         # consider a stable f0 if it is close to the previous one
        else:
            f0stable = 0
        f0 = np.append(f0, f0t)                                  # add f0 to output array
        pin += H                                                 # advance sound pointer
    return f0

def f0Twm(pfreq, pmag, ef0max, minf0, maxf0, f0t=0):
    """
    Function that wraps the f0 detection function TWM, selecting the possible f0 candidates
    and calling the function TWM with them
    pfreq, pmag: peak frequencies and magnitudes,
    ef0max: maximum error allowed, minf0, maxf0: minimum  and maximum f0
    f0t: f0 of previous frame if stable
    returns f0: fundamental frequency in Hz
    """
    if (minf0 < 0):                                  # raise exception if minf0 is smaller than 0
        raise ValueError("Minumum fundamental frequency (minf0) smaller than 0")

    if (maxf0 >= 10000):                             # raise exception if maxf0 is bigger than 10000Hz
        raise ValueError("Maximum fundamental frequency (maxf0) bigger than 10000Hz")

    if (pfreq.size < 3) & (f0t == 0):                # return 0 if less than 3 peaks and not previous f0
        return 0

    f0c = np.argwhere((pfreq>minf0) & (pfreq<maxf0))[:,0] # use only peaks within given range
    if (f0c.size == 0):                              # return 0 if no peaks within range
        return 0
    f0cf = pfreq[f0c]                                # frequencies of peak candidates
    f0cm = pmag[f0c]                                 # magnitude of peak candidates

    if f0t>0:                                        # if stable f0 in previous frame
        shortlist = np.argwhere(np.abs(f0cf-f0t)<f0t/2.0)[:,0]   # use only peaks close to it
        maxc = np.argmax(f0cm)
        maxcfd = f0cf[maxc]%f0t
        if maxcfd > f0t/2:
            maxcfd = f0t - maxcfd
        if (maxc not in shortlist) and (maxcfd>(f0t/4)): # or the maximum magnitude peak is not a harmonic
            shortlist = np.append(maxc, shortlist)
        f0cf = f0cf[shortlist]                         # frequencies of candidates

    if (f0cf.size == 0):                             # return 0 if no peak candidates
        return 0

    f0, f0error = TWM_p(pfreq, pmag, f0cf)        # call the TWM function with peak candidates

    if (f0>0) and (f0error<ef0max):                  # accept and return f0 if below max error allowed
        return f0
    else:
        return 0

def TWM_p(pfreq, pmag, f0c):
    """
    Two-way mismatch algorithm for f0 detection (by Beauchamp&Maher)
    [better to use the C version of this function: UF_C.twm]
    pfreq, pmag: peak frequencies in Hz and magnitudes,
    f0c: frequencies of f0 candidates
    returns f0, f0Error: fundamental frequency detected and its error
    """

    p = 0.5                                          # weighting by frequency value
    q = 1.4                                          # weighting related to magnitude of peaks
    r = 0.5                                          # scaling related to magnitude of peaks
    rho = 0.33                                       # weighting of MP error
    Amax = max(pmag)                                 # maximum peak magnitude
    maxnpeaks = 10                                   # maximum number of peaks used
    harmonic = np.matrix(f0c)
    ErrorPM = np.zeros(harmonic.size)                # initialize PM errors
    MaxNPM = min(maxnpeaks, pfreq.size)
    for i in range(0, MaxNPM) :                      # predicted to measured mismatch error
        difmatrixPM = harmonic.T * np.ones(pfreq.size)
        difmatrixPM = abs(difmatrixPM - np.ones((harmonic.size, 1))*pfreq)
        FreqDistance = np.amin(difmatrixPM, axis=1)    # minimum along rows
        peakloc = np.argmin(difmatrixPM, axis=1)
        Ponddif = np.array(FreqDistance) * (np.array(harmonic.T)**(-p))
        PeakMag = pmag[peakloc]
        MagFactor = 10**((PeakMag-Amax)/20)
        ErrorPM = ErrorPM + (Ponddif + MagFactor*(q*Ponddif-r)).T
        harmonic = harmonic+f0c

    ErrorMP = np.zeros(harmonic.size)                # initialize MP errors
    MaxNMP = min(maxnpeaks, pfreq.size)
    for i in range(0, f0c.size) :                    # measured to predicted mismatch error
        nharm = np.round(pfreq[:MaxNMP]/f0c[i])
        nharm = (nharm>=1)*nharm + (nharm<1)
        FreqDistance = abs(pfreq[:MaxNMP] - nharm*f0c[i])
        Ponddif = FreqDistance * (pfreq[:MaxNMP]**(-p))
        PeakMag = pmag[:MaxNMP]
        MagFactor = 10**((PeakMag-Amax)/20)
        ErrorMP[i] = sum(MagFactor * (Ponddif + MagFactor*(q*Ponddif-r)))

    Error = (ErrorPM[0]/MaxNPM) + (rho*ErrorMP/MaxNMP)  # total error
    f0index = np.argmin(Error)                       # get the smallest error
    f0 = f0c[f0index]                                # f0 with the smallest error

    return f0, Error[f0index]


from scipy.ndimage.fourier import fourier_gaussian
import numpy
import scipy

print('numpy version:', numpy.__version__)
print('scipy version:', scipy.__version__)

print('fourier_gaussian:', fourier_gaussian(numpy.ones([1, 1]), 1))


import sys

from matplotlib import axes
from matplotlib.figure import Figure
from PyQt5 import QtCore, QtWidgets
from scipy import constants as Cn

from Solver import Projectile
from ui import Ui

# Constants

#yapf: disable
t           : float     = 0
dt          : float     = 0.005
A           : float     = 4e-5
B           : float     = 0

g           : float     = Cn.g
v           : float     = 0

Projectiles : dict      = {}
lineList    : dict      = {}

max_xrange  : list      = [0, 1]
max_yrange  : list      = [0, 1]
max_time    : float     = 0
bgCache                 = None

fig         : Figure
ax          : axes
#yapf: enable

def PrepareFigure():
    global fig, ax

    fig = MyUI.canvas.fig
    ax = fig.gca()

    ax.cla()
    ax.yaxis.grid(color='gainsboro', linestyle='dotted', linewidth=1.5)
    ax.xaxis.grid(color='gainsboro', linestyle='dotted', linewidth=0.8)
    ax.axhline(0, linestyle='dotted', color='grey')
    ax.axvline(0, linestyle='dotted', color='grey')

    ax.set_title('Projectile Trajectory')
    ax.set_xlabel(r'$x$', loc='center')
    ax.set_ylabel(r'$y$', loc='center', rotation=0)

    # fig.tight_layout(h_pad=2, w_pad=1)
    fig.tight_layout()

def LoadData():
    global dt, A, B, Projectiles

    dt, A, B, Prjtls = MyUI.GetData().values()

    Projectiles = {}

    for id, Prjctl in Prjtls.items():
        projectile = Projectile(dt, A, B, Prjctl['m'], Prjctl['x_i'], Prjctl['y_i'], Prjctl['v'], Prjctl['ang'], g)
        Projectiles[id] = projectile

def ProcessData():
    global Projectiles, max_xrange, max_yrange, max_time

    max_xrange = [0, 1]
    max_yrange = [0, 1]
    max_time = 0.0

    for Prjctl in Projectiles.values():
        Prjctl: Projectile

        Prjctl.Solve()

        max_time = max(max_time, Prjctl.getMaxTime())

        xr = Prjctl.getMaxXRange()

        yr = Prjctl.getMaxYRange()

        max_xrange[0] = min(xr[0], max_xrange[0])
        max_xrange[1] = max(xr[1], max_xrange[1])

        max_yrange[0] = min(yr[0], max_yrange[0])
        max_yrange[1] = max(yr[1], max_yrange[1])

    MyUI.SetMaxTime(max_t=max_time)
    set_range()

def set_range():
    global max_xrange, max_yrange

    # set aspect

    # padding
    d = (max_xrange[0] - max_xrange[1]) * 0.08
    max_xrange[0] = max_xrange[0] + d
    max_xrange[1] = max_xrange[1] - d

    d = (max_yrange[0] - max_yrange[1]) * 0.08
    max_yrange[0] = max_yrange[0] + d
    max_yrange[1] = max_yrange[1] - d

def plot_points(animated: bool = False):
    global lineList, Projectiles

    for id, Prjctl in Projectiles.items():
        Prjctl: Projectile

        x, y = Prjctl.getResults(slice(0, t))

        lineList[id] = ax.plot(x,
                               y,
                               marker='o',
                               markersize=4,
                               markevery=[-1],
                               label=str(Prjctl.angle),
                               animated=animated)

def update_points(blit: bool = False):
    global t

    for id, Prjctl in Projectiles.items():
        Prjctl: Projectile
        x, y = Prjctl.getResults(slice(0, t))
        lineList[id][0].set_xdata(x)
        lineList[id][0].set_ydata(y)

    if blit == True:
        MyUI.canvas.restore_region(bgCache)
        for id, Prjctl in Projectiles.items():
            Prjctl: Projectile
            ax.draw_artist(lineList[id][0])
        MyUI.canvas.blit(ax.bbox)
    else:
        MyUI.canvas.draw()

def cache_bg():
    global bgCache

    bgCache = MyUI.canvas.copy_from_bbox(ax.bbox)

def RedrawPlots():
    global lineList, ax, max_xrange, max_yrange, t, fig

    t = 0

    LoadData()
    ProcessData()

    RefreshPlots()

    ax.set_xlim(tuple(max_xrange))
    ax.set_ylim(tuple(max_yrange))

    ax.legend()

    MyUI.canvas.draw()

    cache_bg()

def RefreshPlots(animated: bool = False):
    global lineList

    for id, line in lineList.items():
        line[0].remove()

    lineList = {}
    ax.set_prop_cycle(None)

    plot_points(animated=animated)

    MyUI.canvas.draw()

def startAnimation():
    RefreshPlots(animated=True)
    cache_bg()

def stopAnimation():
    RefreshPlots()

def animate(t_new: float):
    global dt, t

    t = t_new

    update_points(blit=True)

def windowResize():
    global fig

    fig.tight_layout()

if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)

    MyUI = Ui()

    PrepareFigure()

    MyUI.Redraw.connect(RedrawPlots)
    MyUI.blit_start.connect(startAnimation)
    MyUI.blit_stop.connect(stopAnimation)
    MyUI.update_plot.connect(animate)
    MyUI.resizeSignal.connect(windowResize)

    # MyUI.showMaximized()
    MyUI.show()

    windowResize()
    RedrawPlots()

    app.exec_()