text:
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

class MigrVisualizer():
    def __init__(self, d_row, d_col):
        """initial vis. TopLeft is origin point.

        :param d_row: grid height
        :type d_row: int
        :param d_col: grid width
        :type d_col: grid width
        """
        self.d_row = d_row
        self.d_col = d_col
        self.states_num = d_row * d_col  # total states

    def potential_heatmap(self, data, **kwargs):
        """plot potential heatmap and save png.
        2 reserved keys:

        * ``title`` figure title
        * ``path`` figure path prefix

        :param data: i-th row represents the potential of d state
        :type data: ndarray
        """
        for i in range(self.states_num):
            distribution = data[i, :].reshape(self.d_row, self.d_col)
            axes = sns.heatmap(distribution)
            axes.set_title(f"{self.ind2rowcol(i)}_{kwargs['title']}")
            fig = axes.get_figure()
            fig.savefig(f"{kwargs['path']}_{i}_step.png")
            plt.close(fig)

    # FIXME: CHANGE API
    def visualize_location(self, xx, yy, xy_size, **kwargs):
        """plot grid location distribution. Origin point is in DownLeft.
        3 reservered key in kwargs:
        * ``fig_name``: savefig
        * ``xlabel``: xlabel
        * ``ylabel``

        :param xx: scatter plot x
        :type xx: ndarray
        :param yy: scatter plot y
        :type yy: ndarray
        :param xy_size: at position (x,y) the number of particles
        :type xy_size: ndarray
        """
        cmap = sns.cubehelix_palette(dark=.3, light=.8, as_cmap=True)
        # FIXME: remove the comments. use the quickplot to refactor code here
        with sns.axes_style("whitegrid"):
            axes = sns.scatterplot(x=xx, y=yy, size=xy_size,
                                   sizes=(0, np.sqrt(np.max(xy_size)) /
                                          np.sqrt(np.sum(xy_size)) * 5000),
                                   palette=cmap)

            axes.set_xlim([0 - 1, self.d_col + 1])
            axes.set_ylim([0 - 1, self.d_row + 1])

            axes.set_xticks(np.arange(1, self.d_col + 1, 3))
            axes.set_xticks(np.arange(1, self.d_col + 1), minor=True)
            axes.set_yticks(np.arange(1, self.d_row + 1, 3))
            axes.set_yticks(np.arange(1, self.d_col + 1), minor=True)
            axes.set_yticklabels([])
            axes.set_xticklabels([])

            if "ylabel" in kwargs:
                axes.set_ylabel(kwargs['ylabel'])
            if "xlabel" in kwargs:
                axes.set_xlabel(kwargs['xlabel'])

            axes.legend_.remove()
            axes.grid(which='both')
            axes.grid(which='minor', alpha=0.2)
            axes.grid(which='major', alpha=0.5)
            axes.yaxis.label.set_size(32)
            axes.xaxis.label.set_size(32)
            plt.savefig(kwargs["fig_name"], pad_inches=-4)
            plt.close()

    def migration(self, data, **kwargs):
        """draw migration figure, heat map distribution
        3 reservered key in kwargs:
        * ``fig_name``: savefig
        * ``xlabel``: xlabel
        * ``ylabel``

        :param data: [i,j] record i-th particle position at j timestamp
        """
        _, time_length = data.shape
        for i in range(time_length):
            locations = data[:, i]  # i timestamp distribution
            bins, _ = np.histogram(locations, np.arange(self.states_num + 1))

            if 'ylabel' in kwargs:
                kwargs['ylabel'] = f"t={i}"
            kwargs['fig_name'] = f"{kwargs['path']}_{i}_step.png"

            self.visualize_map_bins(bins, **kwargs)

    def visualize_map_bins(self, bins, **kwargs):
        """converts the statistics bins data to the map figure.
        3 reservered key in kwargs:
        * ``fig_name``: savefig
        * ``xlabel``: xlabel
        * ``ylabel``

        :param bins: every element represents how many particles in the cell
        :type bins: list or ndarray
        """
        xx = []
        yy = []
        xy_cnt = []
        for xy, cnt in enumerate(bins):
            if cnt > 0:
                row, col = self.ind2rowcol(xy)
                xx.append(col)
                yy.append(row)
                xy_cnt.append(int(cnt))
        self.visualize_location(xx, yy, xy_cnt, **kwargs)

    def ind2rowcol(self, index):
        index = np.array(index).astype(np.int64)
        row = (index / self.d_col).astype(np.int64)
        col = index % self.d_col
        return row, col

max_stars_repo_name: zqsh/cbp

max_issues_repo_name: zqsh/cbp

max_forks_repo_name: zqsh/cbp


text:
import nextnanopy as nn
from nextnanopy.utils.misc import mkdir_if_not_exist
import sys,os
import numpy as np
import matplotlib.pyplot as plt

#import config_nextnano     # This should be your default configuration.
#import config_nextnano_temp # This could be a modified configuration file.
# config file is stored in C:\Users\<User>\.nextnanopy-config

#FigFormat = '.pdf'
FigFormat = '.svg'
#FigFormat = '.jpg'
#FigFormat = '.png'

#++++++++++++++++++++++++++++++++++++++++++++++
# These lines have to be adjusted by the user.
#++++++++++++++++++++++++++++++++++++++++++++++
#================================
# Specify software product here!
#================================
software = 'nextnano++'
#software = 'nextnano3'
#software = 'nextnano.NEGF'
#software = 'nextnano.MSB'
#===========================

folder_examples_nnp    = r'C:\Program Files\nextnano\2021_08_06\Sample files\nextnano++ sample files'
folder_examples_nn3    = r'C:\Program Files\nextnano\2021_08_06\Sample files\nextnano3 sample files'
folder_examples_nnNEGF = r'D:\nextnano.NEGF\nextnanoNEGF_2020_11_16\nextnano.NEGF sample files'
folder_examples_nnMSB  = r'D:\nextnano.MSB\nextnano.MSB_2017_12_20\nextnano\2017_12_20\nextnano.MSB sample files'

#++++++++++++++++++++++++++++++++++++++++++++++
# These lines have to be adjusted by the user.
#++++++++++++++++++++++++++++++++++++++++++++++
#folder_examples_nn3 = r'N:\users\nextnano\nextnano GmbH - Tutorials\Tutorials\2D The CBR method (Transmission)'
#===========================
if(software=="nextnano++"):
   subfolder = ''
  #  subfolder = r'\Quantum Mechanics examples'
elif(software=="nextnano3"):
   subfolder = ''
  #  subfolder = r'\Quantum Mechanics examples'
elif(software=="nextnano.NEGF"):
    subfolder = ''
elif(software=="nextnano.MSB"):
    subfolder = ''
#===========================

software_short_nnp    = '_nnp'
software_short_nn3    = '_nn3'
software_short_nnNEGF = '_nnNEGF'
software_short_nnMSB  = '_nnMSB'

#===========================
if(software=="nextnano++"):
    software_short = software_short_nnp
elif(software=="nextnano3"):
    software_short = software_short_nn3
elif(software=="nextnano.NEGF"):
    software_short = software_short_nnNEGF
elif(software=="nextnano.MSB"):
    software_short = software_short_nnMSB
#===========================

#===========================
if(software=="nextnano++"):
    folder_examples = folder_examples_nnp + subfolder   # nextnano++
elif(software=="nextnano3"):
    folder_examples = folder_examples_nn3 + subfolder   # nextnano3
elif(software=="nextnano.NEGF"):
    folder_examples = folder_examples_nnNEGF + subfolder # nextnano.NEGF
elif(software=="nextnano.MSB"):
    folder_examples = folder_examples_nnMSB + subfolder # nextnano.MSB
#===========================

#================================================================
# Define output folders. If they do not exist, they are created.
#================================================================
folder_output        = nn.config.config[software]['outputdirectory']
folder_output_python = os.path.join(folder_output,r'nextnanopy')
mkdir_if_not_exist(folder_output)
mkdir_if_not_exist(folder_output_python)

#--------------------------------------------------------
# Specify input file without file extension '.in'/.'xml'
#--------------------------------------------------------
#my_input_file_no_extension_nnp = r'Jogai_AlGaNGaN_FET_JAP2003_GaNcap_Fig6Fig5_1D_nnp'
#my_input_file_no_extension_nn3 = r'Jogai_AlGaNGaN_FET_JAP2003_GaNcap_Fig6Fig5_1D_nn3'
#my_input_file_no_extension_nn3 = r'2D_CBR_MamaluySabathilJAP2003_AharonovBohm'
#my_input_file_no_extension_nn3 = r'2D_CBR_square'
my_input_file_no_extension_nnp = r'LaserDiode_InGaAs_1D_qm_nnp_nnmat_AvsAscii_one_file'
my_input_file_no_extension_nn3 = r'Transmission_Double_Barrier_1D_nn3'
my_input_file_no_extension_nnNEGF = r'THz_QCL_GaAs_AlGaAs_Fathololoumi_OptExpress2012_10K-FAST'
my_input_file_no_extension_nnMSB = r'1D_Transmission_DoubleBarrier_CBR_paper_MSB'

#===========================
if(software=="nextnano++"):
    my_input_file_no_extension = my_input_file_no_extension_nnp
elif(software=="nextnano3"):
    my_input_file_no_extension = my_input_file_no_extension_nn3
elif(software=="nextnano.NEGF"):
    my_input_file_no_extension = my_input_file_no_extension_nnNEGF
elif(software=="nextnano.MSB"):
    my_input_file_no_extension = my_input_file_no_extension_nnMSB
#===========================

if(software=="nextnano.NEGF" or software=="nextnano.MSB"):
    FileExtension = '.xml' # for nextnano.NEGF and nextnano.MSB
else:
    FileExtension = '.in'  # for nextnano++ and nextnano3

my_input_file = my_input_file_no_extension+FileExtension

# plt.ion() # interactive mode

print(f"starting nextnano...")
input_file_name = os.path.join(folder_examples,my_input_file)
input_file = nn.InputFile(r'C:\Users\naoki.mitsui\Documents\nextnano\My input files\laser_diode_tutorial\LaserDiode_InGaAs_1D_qm_nnp_nnmat_AvsAscii_one_file.in')

print(f'')
print(f'=====================================')
############################
# Execute nextnano software
############################
# input_file.execute() # Put line into comment if you only want to to post-processing of results
print(f'=====================================')

############################
# plot energy-resolved density + conduction bandedge (cb)
############################
ListOfValues = ['00000', '00001', '00002', '00003', '00004', '00005', '00006', '00007', '00008', '00009', '00010']
#ListOfValues = ['00001', '00008']
bias = 0
count = 0
inner_count = 0

fig_several = plt.figure()

for i in ListOfValues:
    bias = round(bias, 1)
    #===========================
    # Conduction and valence band edge
    #===========================
    file_be = os.path.join(folder_output,my_input_file_no_extension+r'\bias_'+i+r'\bandedges.dat')

    print(f"Read in file:")
    print(file_be)
    df_be = nn.DataFile(file_be,product=software)

    #===========================
    # Energy-resolved electron density
    #===========================
    file_elec = os.path.join(folder_output,my_input_file_no_extension+r'\bias_'+i+r'\electron_density_vs_energy.fld')
    #  file = os.path.join(folder_output,input_file_name_variable+r'\Results'+r'\LocalDOS_sg1_deg1.vtr')
    # Caution! Specify "format2D=AvsAscii_one_file" in the input file. Otherwise this line causes error due to the binary data reading.
    datafile_2d = nn.DataFile(file_elec,product=software)

    x=datafile_2d.coords['x']
    y=datafile_2d.coords['y']
    z_elec=datafile_2d.variables[0]

    fig, ax = plt.subplots(1)
    pcolor = ax.pcolormesh(x.value,y.value,z_elec.value.T)
    cbar = fig.colorbar(pcolor)
    cbar.set_label('Electron density [cm$^{-3}$eV$^{-1}$]')
    # conduction bandedge plot
    ax.plot(df_be.coords[0].value,df_be.variables[0].value, color='white', linestyle='-')
    # HH valence bandedge plot
    ax.plot(df_be.coords[0].value,df_be.variables[1].value, color='white', linestyle='-')
    # electron Fermi level plot
    ax.plot(df_be.coords[0].value,df_be.variables[4].value, color='red',  linestyle='dotted', linewidth=1.8)
    # hole Fermi level plot
    ax.plot(df_be.coords[0].value,df_be.variables[5].value, color='green',  linestyle='dotted', linewidth=1.8)

    ax.set_xlabel(f"{x.name} [{x.unit}]")
    ax.set_ylabel('Energy [eV]')
    ax.set_title('Energy resolved electron density at '+str(bias)+' V')
    fig.tight_layout()
    # plt.show()
    fig.savefig(file_elec+'.jpg', facecolor='white')

    #===========================
    # Energy-resolved hole density
    #===========================
    file_hole = os.path.join(folder_output,my_input_file_no_extension+r'\bias_'+i+r'\hole_density_vs_energy.fld')
    #  file = os.path.join(folder_output,input_file_name_variable+r'\Results'+r'\LocalDOS_sg1_deg1.vtr')
    # Caution! Specify "format2D=AvsAscii_one_file" in the input file. Otherwise this line causes error due to the binary data reading.
    datafile_2d = nn.DataFile(file_hole,product=software)

    x=datafile_2d.coords['x']
    y=datafile_2d.coords['y']
    z_hole=datafile_2d.variables[0]

    fig, ax = plt.subplots(1)
    pcolor = ax.pcolormesh(x.value,y.value,z_hole.value.T, vmin=0.00, vmax=2.00*10**15)
    cbar = fig.colorbar(pcolor)
    cbar.set_label('Hole density [cm$^{-3}$eV$^{-1}$]')
    # conduction bandedge plot
    ax.plot(df_be.coords[0].value,df_be.variables[0].value, color='white', linestyle='-')
    # HH valence bandedge plot
    ax.plot(df_be.coords[0].value,df_be.variables[1].value, color='white', linestyle='-')
    # electron Fermi level plot
    ax.plot(df_be.coords[0].value,df_be.variables[4].value, color='red',  linestyle='dotted', linewidth=1.8)
    # hole Fermi level plot
    ax.plot(df_be.coords[0].value,df_be.variables[5].value, color='green',  linestyle='dotted', linewidth=1.8)

    ax.set_xlabel(f"{x.name} [{x.unit}]")
    ax.set_ylabel('Energy [eV]')
    ax.set_title('Energy resolved hole density at '+str(bias)+' V')
    fig.tight_layout()
    # plt.show()
    fig.savefig(file_hole+'.jpg', facecolor='white')

    #===========================
    # Energy-resolved electron and hole density
    #===========================
    z_elec_plus_hole = z_elec.value + z_hole.value

    fig, ax = plt.subplots(1)
    pcolor = ax.pcolormesh(x.value,y.value,z_elec_plus_hole.T, vmin=0.00, vmax=2.00*10**15)
    cbar = fig.colorbar(pcolor)
    cbar.set_label('Electron and hole densities [cm$^{-3}$eV$^{-1}$]')
    # conduction bandedge plot
    ax.plot(df_be.coords[0].value,df_be.variables[0].value, color='white', linestyle='-')
    # HH valence bandedge plot
    ax.plot(df_be.coords[0].value,df_be.variables[1].value, color='white', linestyle='-')
    # electron Fermi level plot
    ax.plot(df_be.coords[0].value,df_be.variables[4].value, color='red',  linestyle='dotted', linewidth=1.8)
    # hole Fermi level plot
    ax.plot(df_be.coords[0].value,df_be.variables[5].value, color='green',  linestyle='dotted', linewidth=1.8)

    ax.set_xlabel(f"{x.name} [{x.unit}]")
    ax.set_ylabel('Energy [eV]')
    ax.set_title('Energy resolved densities of hole and electron at '+str(bias)+' V', y=1.035)
    fig.tight_layout()
    # plt.show()
    elec_hole_img = os.path.join(folder_output,my_input_file_no_extension+r'\bias_'+i+r'\electron_hole_density_vs_energy.jpg')
    fig.savefig(elec_hole_img, dpi=300, facecolor='white')

    #===========================
    # Energy-resolved electron and hole densities in several biases
    #===========================
    if count in [2,4,8,10]:
        ax = fig_several.add_subplot(2,2,inner_count+1)

        pcolor = ax.pcolormesh(x.value,y.value,z_elec_plus_hole.T, vmin=0.00, vmax=2.00*10**15)
        #cbar = fig.colorbar(pcolor)
        #cbar.set_label('Electron and hole densities [cm$^{-3}$eV$^{-1}$]')
        # conduction bandedge plot
        ax.plot(df_be.coords[0].value,df_be.variables[0].value, color='white', linestyle='-')
        # HH valence bandedge plot
        ax.plot(df_be.coords[0].value,df_be.variables[1].value, color='white', linestyle='-')
        # electron Fermi level plot
        ax.plot(df_be.coords[0].value,df_be.variables[4].value, color='red',  linestyle='dotted', linewidth=1.8)
        # hole Fermi level plot
        ax.plot(df_be.coords[0].value,df_be.variables[5].value, color='green',  linestyle='dotted', linewidth=1.8)

        ax.set_xlabel(f"{x.name} [{x.unit}]", fontsize=8)
        ax.set_ylabel('Energy [eV]', fontsize=8)
        ax.tick_params(axis='x', labelsize=8)
        ax.tick_params(axis='y', labelsize=8)
        ax.set_title('Bias = '+str(bias)+' V', y=1.035, fontsize=8)

        inner_count += 1

    bias += 0.1
    count+=1

#cbar = fig_several.colorbar(pcolor)
#cbar.set_label('Electron and hole densities [cm$^{-3}$eV$^{-1}$]')
fig_several.suptitle('Energy resolved densities of electron and hole')
fig_several.tight_layout()
elec_hole_img_several = os.path.join(folder_output,my_input_file_no_extension+r'\bias_'+i+r'\electron_hole_density_vs_energy_several.jpg')
fig_several.savefig(elec_hole_img_several, dpi=300, facecolor='white')
plt.show()

print(f'=====================================')
print(f'Done nextnanopy.')
print(f'=====================================')

ext: py

max_stars_repo_name: nextnanopy/nextnanopy

max_issues_repo_name: nextnanopy/nextnanopy

max_forks_repo_name: nextnanopy/nextnanopy

reason: import numpy


text:
#=============================================================================#
#                                                                             #
# Copyright (c) 2005 - 2015 G-Truc Creation (www.g-truc.net)                  #
#                                                                             #
# This module is a derivative work of the OpenGL Mathematics                  #
# (glm.g-truc.net) gtc/matrix_transform module.                               #
#                                                                             #
# Permission is hereby granted, free of charge, to any person obtaining a     #
# copy of this software and associated documentation files (the "Software"),  #
# to deal in the Software without restriction, including without limitation   #
# the rights to use, copy, modify, merge, publish, distribute, sublicense,    #
# and/or sell copies of the Software, and to permit persons to whom the       #
# Software is furnished to do so, subject to the following conditions:        #
#                                                                             #
# The above copyright notice and this permission notice shall be included in  #
# all copies or substantial portions of the Software.                         #
#                                                                             #
# Restrictions:                                                               #
#		By making use of the Software for military purposes, you choose to make   #
#		a Bunny unhappy.                                                          #
#                                                                             #
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR  #
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,    #
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE #
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER      #
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING     #
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER         #
# DEALINGS IN THE SOFTWARE.                                                   #
#                                                                             #
#=============================================================================#

from math import sin, cos, tan
from numpy import array, identity, zeros, ones, copyto
from numpy.linalg import norm

def m_translate(m, v):
	"""
	Builds a translation 4 * 4 matrix created from a vector of 3 components.

	\param m  Input matrix multiplied by this translation matrix.
	\param v  Coordinates of a translation vector.
	"""
	t = array(m)
	t[3] = m[0]*v[0] + m[1]*v[1] + m[2]*v[2] + m[3]
	return t

def m_translate_in_place(m, v):
	"""
	Translates a 4 * 4 matrix created from a vector of 3 components.

	\param m  Input matrix multiplied by this translation matrix.
	\param v  Coordinates of a translation vector.
	"""
	m[3] = m[0]*v[0] + m[1]*v[1] + m[2]*v[2] + m[3]

def m_rotate_axis(m, theta, axis):
	"""
	Builds a rotation 4 * 4 matrix created from an axis vector and an angle.

	\param m      Input matrix multiplied by this rotation matrix.
	\param theta (rad) Rotation angle.
	\param axis   Rotation axis, recommended to be normalized.
	"""
	c = cos(theta)
	s = sin(theta)

	v = array(axis)
	v /= norm(v)

	tmp = v * (1-c)

	rotation = identity(3)

	rotation[0,0] = tmp[0]*v[0] + c
	rotation[0,1] = tmp[0]*v[1] + s*v[2]
	rotation[0,2] = tmp[0]*v[2] - s*v[1]

	rotation[1,0] = tmp[1]*v[0] - s*v[2]
	rotation[1,1] = tmp[1]*v[1] + c
	rotation[1,2] = tmp[1]*v[2] + s*v[0]

	rotation[2,0] = tmp[2]*v[0] + s*v[1]
	rotation[2,1] = tmp[2]*v[1] - s*v[0]
	rotation[2,2] = tmp[2]*v[2] + c

	r = m.copy()

	r[0] = m[0]*rotation[0,0] + m[1]*rotation[0,1] + m[2]*rotation[0,2]
	r[1] = m[0]*rotation[1,0] + m[1]*rotation[1,1] + m[2]*rotation[1,2]
	r[2] = m[0]*rotation[2,0] + m[1]*rotation[2,1] + m[2]*rotation[2,2]

	return r

def m_rotate(m, r):
	"""
	http://www.opengl-tutorial.org/assets/faq_quaternions/index.html
	"""
	A = cos(r[0])
	B = sin(r[0])
	C = cos(r[1])
	D = sin(r[1])
	E = cos(r[2])
	F = sin(r[2])

	AD = A*D
	BD = B*D

	mat = array([[  C*E      ,  -C*F      ,    D,  0.],
	             [ BD*E + A*F, -BD*F + A*E, -B*C,  0.],
	             [-AD*E + B*F,  AD*F + B*E,  A*C,  0.],
	             [         0.,          0.,   0.,  1.]])

	return mat.dot(m)

def m_rotate_in_place(m, r):
	"""
	http://www.opengl-tutorial.org/assets/faq_quaternions/index.html
	"""
	A = cos(r[0])
	B = sin(r[0])
	C = cos(r[1])
	D = sin(r[1])
	E = cos(r[2])
	F = sin(r[2])

	AD = A*D
	BD = B*D

	mat = array([[  C*E      ,  -C*F      ,    D,  0.],
	             [ BD*E + A*F, -BD*F + A*E, -B*C,  0.],
	             [-AD*E + B*F,  AD*F + B*E,  A*C,  0.],
	             [         0.,          0.,   0.,  1.]])

	copyto(m, mat.dot(m))

def m_scale(m, v):
	"""
	Builds a scale 4 * 4 matrix created from 3 scalars.

	\param m  Input matrix multiplied by this scale matrix.
	\param v  Ratio of scaling for each axis.
	"""
	s = array(m)
	s[0] = m[0]*v[0]
	s[1] = m[1]*v[1]
	s[2] = m[2]*v[2]
	s[3] = m[3]
	return s

def m_scale_in_place(m, v):
	"""
	Builds a scale 4 * 4 matrix created from 3 scalars.

	\param m  Input matrix multiplied by this scale matrix.
	\param v  Ratio of scaling for each axis.
	"""
	m[0] *= v[0]
	m[1] *= v[1]
	m[2] *= v[2]

def m_perspective(fov, aspect, min_z, max_z):
	"""
	Creates a matrix for a symetric perspective-view frustum based on the default
	handedness.

	\param fov    (rad) Specifies the field of view angle in the y direction.
	\param aspect  Specifies the aspect ratio that determines the field of view
                 in the x direction. The aspect ratio is the ratio of x (width)
	               to y (height).
	\param min_z   Specifies the distance from the viewer to the near clipping
	               plane (always positive).
	\param max_z   Specifies the distance from the viewer to the far clipping
	               plane (always positive).
	"""

	if aspect == 0.:
		raise ValueError("Invalid aspect ratio.")
	if min_z > max_z:
		raise ValueError("Minimum draw distance cannot be greater than maximum.")

	f = tan(fov/2.)

	m = zeros((4,4))

	m[0,0] = 1./(aspect*f)
	m[1,1] = 1./f
	m[2,2] = -(max_z+min_z)/(max_z-min_z)
	m[2,3] = -1.
	m[3,2] = -(2.*max_z*min_z)/(max_z-min_z)

	return m

def m_orthographic(left, right, bottom, top, min_z, max_z):
	"""
	Creates a matrix for an orthographic parallel viewing volume, using
	right-handed coordinates. 	The near and far clip planes correspond to z
	normalized device coordinates of -1 and +1 respectively.

	"""

	m = identity(4)

	m[0,0] = 2./(right-left)
	m[1,1] = 2./(top-bottom)
	m[2,2] = -2./(max_z-min_z)
	m[3,0] = -(right+left)/(right-left)
	m[3,1] = -(top+bottom)/(top-bottom)
	m[3,2] = -(max_z+min_z)/(max_z-min_z)

	return m

max_stars_repo_name: RobinSatterthwaite/opengl-experiments

max_issues_repo_name: RobinSatterthwaite/opengl-experiments

max_forks_repo_name: RobinSatterthwaite/opengl-experiments

reason: from numpy


text:
#Metro Solver
#Given a metro system, this program will use Q Learning to determine what the
#best route between a given start point and end point is, given the assumption that the time value of each trip between any two stops is equal.

import numpy as np
import random
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation
from keras.optimizers import RMSprop
from IPython.display import clear_output

#Number of metro stations
numStations = 10 #this is for the small set
#numStations = 62 #this is for the london underground

goal = 8
#def our states, we have numStations possible states with each station
#having 1 - numStations possible actions
#row being station A row2 being at station B, row3 being at station C, and row 4 being at station D. A 0 means no direct connection, 1 means connection, 10 means action results in getting to the target station
#for now, the target station is D. We will start at A
def initState():
    state = np.zeros((2,numStations))
    state[0,3] = 1 #place player at station n
    state[1,goal] = 1 #place goal
    return state

def initStateC(start,goal):
    state = np.zeros((2,numStations))
    state[0,start] = 1 #place player at station 1 for simplicity
    state[1,goal] = 1 #place goal
    return state

def initStateRand():
    state = np.zeros((2,numStations))
    station = random.randint(0,numStations-1)
    while station == goal:
        station = random.randint(0,numStations-1)
    state[0,station] = 1
    state[1,goal] = 1
    return state

def initRewardM(targetStation):
    '''
    rM = np.array((   [ -1,  -10,   -1,  -1],
                      [-10,  10,  -10,  -1],
                      [ -1,  -10,   -1,  -1],
                      [ -1,  10,   -1,  -1]))
    '''
    ''' a hand encoded matrix

    #             1   2   3    4   5   6   7   8   9   10
    rM = np.array(( [ -1, -1,-10,-10,-10,-10, -1,-10,-10,-10],  #1
                    [ -1, -1,-10,-10,-10,-10,-10, -1,-10,-10],  #2
                    [-10,-10, -1, -1, -1, 10,-10,-10,-10,-10],  #3
                    [-10,-10, -1, -1,-10,-10, -1,-10, -1,-10],  #4
                    [-10,-10, -1,-10, -1,-10,-10,-10, -1,-10],  #5
                    [-10,-10, -1,-10,-10, 10,-10, -1,-10,-10],  #6
                    [ -1,-10,-10, -1,-10,-10, -1,-10,-10,-10],  #7
                    [-10, -1,-10,-10,-10, 10,-10, -1,-10,-10],  #8
                    [-10,-10,-10,-10, -1,-10,-10,-10, -1, -1],  #9
                    [-10,-10,-10,-10,-10,-10,-10,-10, -1, -1],  #10
                ))
    '''
    #ok but let's actually use a csv file, so it can be larger and more
    #editable
    #rM = np.genfromtxt('londonUnderground.csv',delimiter=',')
    #rM = np.genfromtxt('small.csv',delimiter=',')
    rM = np.genfromtxt('medium.csv',delimiter=',')

    #now, for every -1 in the target station column, change it to a 10

    col = rM[:,targetStation]
    row = rM[targetStation,:]

    for i in range (0,len(col)):
         if col[i] == -1:
            col[i] = 10
         if row[i] == -1:
            row[i] = 10
    return rM

def getCurrentStation(state):
    for i in range(0,numStations):
        if(state[0,i] == 1):
            return i;

def isValidMove(state,action,rM):
    #remember the shape of our state means the top row is
    #just the
    curStation = getCurrentStation(state)
    if(rM[curStation,action] == -10):
        return False
    else:
        return True

def makeMove(state, action, rM):
    #we will take our action, which is just an int of the station
    #we have randomly chosen to move to, check if we can move there
    #and then if so, update the one hot encoded 0 row of our state
    #to reflect our current position
    #check to see if we can make the move
    if isValidMove(state,action, rM):
        #reset our current station
        state[0] = np.zeros(numStations)
        state[0,action] = 1
    return state

def getReward(state,rM):
    print("getting reward for: \n",state)
    station = getCurrentStation(state);
    print("Station: ",station);
    r = rM[station,station];
    print("Reward: ",r)
    return r
 #   curStation = getCurrentStation(state)
 #   print("getting reward for: ")
 #  print("state: ", state, "location: ",curStation)
 #   r = state[curStation+1,curStation]
 #   if r == 10:
 #       return 10
 #   else:
 #       return -1

st = initStateC(5,3)
rLL = initRewardM(3)

getReward(st,rLL)

def getRandomAction(state,rM):
    #get all items in the row that are
    index = getCurrentStation(state)
    actionSpace = []
    possible = rM[:,index]
    for i in range(0,len(possible)):
        if(possible[i] != -10 and i != index):
            actionSpace.append(i)
    return random.choice(actionSpace)

def getBestValidAction(state,prev,qval,rM):
    print("*********************")
    index = getCurrentStation(state)
    print("Getting best valid action for: ",index)
    actionSpace = []
    possible = rM[index,:]
    print("possible: ",possible)
    for i in range(0,len(possible)):
        if(possible[i] != -10 and i != index and i != prev):
            actionSpace.append(i)
    if(len(actionSpace) < 1):
        #print("=//////////////////////")
        print("Problem isolated node")
        #print("=//////////////////////")
        for i in range(0,len(possible)):
            if(possible[i] != -10 and i != index):
                actionSpace.append(i)

    print("Action Space: ",actionSpace)
    #Ok, we have our action space of possible nodes to move to
    #now, we need to get the Q value for each one, and return
    #what move to make based on that

    best = 0 #arbitrarily small number
    bestQ = 0
    print("comp to QVAL: ",qval)
    for i in range(0, len(actionSpace)):
        print("---Action Space Val: ",actionSpace[i])
        print("---QVal: ",qval[0:,actionSpace[i]])
        if(qval[0:,actionSpace[i]] > bestQ or bestQ == 0):
            best = actionSpace[i]
            bestQ = qval[0:,best]
    print("Best: ",best)
    print("Best QVal: ",qval[0:,best])
    print("********************")
    return best

model = Sequential()
model.add(Dense(220, init='lecun_uniform', input_shape=(numStations*2,)))
model.add(Activation('relu'))
#model.add(Dropout(0.2)) I'm not using dropout, but maybe you wanna give it a try?

model.add(Dense(150, init='lecun_uniform'))
model.add(Activation('relu'))
#model.add(Dropout(0.2))

model.add(Dense(numStations, init='lecun_uniform'))
model.add(Activation('linear')) #linear output so we can have range of real-valued outputs

rms = RMSprop()
model.compile(loss='mse', optimizer=rms)

model.compile(loss='mse', optimizer=rms)#reset weights of neural network
epochs = 100
gamma = 0.975
epsilon = 1
batchSize = 40
buffer = 80
replay = []
h = 0
tTot = 0
for i in range(epochs):
    print("\n")
    print("Init New Game \n")
    rM = initRewardM(goal)
    print("RM::::\n")
    print(rM)
    #state = initState() #using the deterministic
    state = initStateRand() #using the random start
    prev = getCurrentStation(state)
    status = 1
    #while game still in progress
    t = 0
    randMoves = 0
    #print("&&&&&&&&&&&&&&&&&&&&&&&&&&&&&")
    #print("&&&&&&&&&&&&&&&&&&&&&&&&&&&&&")
    #print("&&&&     GAME ",i,"        &&")
    while(status == 1):
        #We are in state S
        #Let's run our Q function on S to get Q values for all possible actions
        #print("+++++++++++++++++++++++++++++++++++++++++")
        qval = model.predict(state.reshape(1,numStations*2), batch_size=1)
        if (random.random() < epsilon): #choose random action
            #action = np.random.randint(0,numStations)
            #action = getBestValidAction(state,rM)
            randMoves += 1
            action = getRandomAction(state,rM)
        else: #choose best action from Q(s,a) values
            #get indexs of possible actions
            #select max of those
            action = getBestValidAction(state,prev,qval,rM)
            #action = (np.argmax(qval))
        #Take action, observe new state S'
        print("take action: ",action)

        prev = getCurrentStation(state)
        new_state = makeMove(state, action, rM)
        reward = getReward(new_state, rM)

        #Experience replay storage
        if (len(replay) < buffer): #if buffer not filled, add to it
            replay.append((state, action, reward, new_state))
        else: #if buffer full, overwrite old values
            if (h < (buffer-1)):
                h += 1
            else:
                h = 0
            replay[h] = (state, action, reward, new_state)
            #randomly sample our experience replay memory
            minibatch = random.sample(replay, batchSize)
            X_train = []
            y_train = []
            for memory in minibatch:
                #Get max_Q(S',a)
                old_state, action, reward, new_state = memory
                old_qval = model.predict(old_state.reshape(1,numStations*2), batch_size=1)
                newQ = model.predict(new_state.reshape(1,numStations*2), batch_size=1)
                maxQ = np.max(newQ)
                y = np.zeros((1,numStations))
                y[:] = old_qval[:]
                if reward == -1: #non-terminal state
                    update = (reward + (gamma * maxQ))
                else: #terminal state
                    update = reward
                y[0][action] = update
                X_train.append(old_state.reshape(numStations*2,))
                y_train.append(y.reshape(numStations,))

            X_train = np.array(X_train)
            y_train = np.array(y_train)
            print("Game #: %s" % (i,))
            model.fit(X_train, y_train, batch_size=batchSize, nb_epoch=1, verbose=1)
            state = new_state
            reward = getReward(state,rM)
            print("---_----__Reward Check: ",reward)
            print(state)

        t += 1
        if reward == 10: # or reward == 10.0: #if reached terminal state, update game status
            status = 0
            print("VICTORY")
            print("Num Turns: ",t)
            print("Rand Moves: ",randMoves)
            tTot += t
        clear_output(wait=True)
    if epsilon > 0.1: #decrement epsilon over time
        epsilon -= (1/epochs)

print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
print("t total: ",tTot)
print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")
print("%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%")

losess = 0

wins = 0

def testAlgo():
    i = 0
    state = initState()
    state = initStateRand()

    print("Initial State:")
    print(state)

    prev = getCurrentStation(state)
    status = 1
    #while game still in progress
    while(status == 1):
        #print("((((((((()))))))")
        #print("((((((((()))))))")
        #print("((((((((()))))))")
        #print("      Turn: ",i)
        qval = model.predict(state.reshape(1,numStations*2), batch_size=1)
        print(qval)
        #action = (np.argmax(qval)) #take action with highest Q-value
        action = getBestValidAction(state,prev,qval,rM)

        #take action with highest Q-value

        prev = getCurrentStation(state)
        print('Move #: %s; Taking action: %s' % (i, action))
        state = makeMove(state, action, rM)
        print(state)
        reward = getReward(state,rM)
        #reward = rM[getCurrentStation(state), getCurrentStation(state)]
        print("got reward: ",reward)
        if reward == 10:
            status = 0
            print("Reward: %s" % (reward,))
            print ("___________VICTORY_________")
            return 1
        i += 1 #If we're taking more than 10 actions, just stop, we probably can't win this game
        if (i > 60):
            print("Game lost; too many moves.")
            print("______FAILURE_____");
            return 0
            break
total = 0
numTests = 100
for i in range(0,numTests):
    total += testAlgo()

print("Wins: ",total)
print("Losses: ",numTests-total)

max_stars_repo_name: rugggg/metroSolv

max_issues_repo_name: rugggg/metroSolv

max_forks_repo_name: rugggg/metroSolv


text:
# Souza & Ramos da Silva,

# Ocean-Land Atmosphere Model (OLAM) performance for major extreme
#   meteorological events near the coastal region of southern Brazil,

# Climate Research, in revision 2020
"""
Created on Sat Feb  6 17:34:55 2021

Script for analysing SLp and Wind Speed data

@author: Danilo
"""
import statistics_Danilo as st
from prepare_data import (regridSLPWind, GetSLPWindData, regrid)
from stats_accprec_time_evo import (mean_corrl, DataToDataFrame)
from scipy import stats
#
import numpy as np
import pandas as pd
import csv
# plotting packages
import pylab as pl
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import matplotlib.gridspec as gridspec
import cartopy.feature as cfeature
import matplotlib.ticker as mticker
import cmocean
from matplotlib import colors
from matplotlib.colors import LinearSegmentedColormap
# ----------
def Get_SLP_Wind_At_EvPeak(event, re):
    olam,obs = GetSLPWindData(event)[0],GetSLPWindData(event)[1]
    # dates corresponding to each event peak time
    dates = ['2018-01-11T18:00:00.000000000',
            '2017-06-04T21:00:00.000000000',
            '2011-09-08T09:00:00.000000000',
            '2011-05-28T06:00:00.000000000',
            '2011-01-21T21:00:00.000000000',
            '2010-04-10T09:00:00.000000000',
            '2008-11-22T12:00:00.000000000',
            '2008-05-03T21:00:00.000000000',
            '2007-07-28T12:00:00.000000000',
            '2006-09-03T18:00:00.000000000',
            '2005-08-10T09:00:00.000000000',
            '2004-03-28T06:00:00.000000000']
    obs = obs.sel(time=dates[event-1])
    olam = olam.sel(time=dates[event-1]).assign(lev=obs.lev)
    # regrid data, if asked
    if re == 'y':
        olam = regridSLPWind(obs.lon,obs.lat,olam,event)
    return obs, olam

# ----------
# params for plotting
# --
# bounding box
min_lon, max_lon, min_lat, max_lat = -54, -45, -34, -26
# --
# Make state boundaries feature
states_provinces = cfeature.NaturalEarthFeature(category='cultural',
                                                name='admin_1_states_provinces_lines',
                                                scale='50m', facecolor='none')
# Make country borders feature
country_borders = cfeature.NaturalEarthFeature(category='cultural',
                                               name='admin_0_countries',
                                               scale='50m', facecolor='none')
# --
def plot_background(ax):
    ax.set_extent([min_lon, max_lon, min_lat, max_lat])
    ax.coastlines('50m', edgecolor='black', linewidth=0.5)
    ax.add_feature(states_provinces, edgecolor='black', linewidth=0.5)
    ax.add_feature(country_borders, edgecolor='black', linewidth=0.5)
    return ax

def plot_slp_panel():
    '''
    Plot SLP panels for each event peak time
    '''
    cmap = cmocean.cm.balance
    # figure params
    fig = plt.figure(figsize=(10,15) , constrained_layout=False)
    gs1 = gridspec.GridSpec(6, 2, hspace=0.25, wspace=0.15, left=0.01, right=0.45)
    gs2 = gridspec.GridSpec(6, 2, hspace=0.25, wspace=0.15, left=0.50, right=0.94)
    axs = []
    datacrs = ccrs.PlateCarree()
    #
    ev = 1
    panel1 = 0
    panel2 = 0
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
    for i in range (1,25):
        # get data
        if i % 2 != 0:
            tmp = Get_SLP_Wind_At_EvPeak(ev,'n')
            obs = tmp[0].transpose('lat','lon','lev')
            olam = tmp[1].transpose('lat','lon', 'lev')
            # reanalysis data
            lonsre, latsre = obs.lon, obs.lat
            obs_slp = obs.SLP/100
            ure, vre = obs.U, obs.V
            # olam data
            lonso, latso = olam.lon, olam.lat
            olam_slp = olam.sslp/100
            uo, vo = olam.uwnd, olam.vwnd
            # adjust SLP limts accordingly to each event data
            minslp = round(np.amin([int(np.amin(obs_slp.values)), int(np.amin(olam_slp.values))])-5,-1)
            maxslp = round(np.amax([int(np.amax(obs_slp.values)), int(np.amax(olam_slp.values))])+5,-1)
            if minslp > 1014:
                minslp = 1010
            if maxslp < 1025:
                maxslp = 1030
            norm = colors.DivergingNorm(vmin=minslp, vcenter=1014, vmax=maxslp)
        # figure
        if ev % 2 != 0:
            panel1 += 1
            axs.append(fig.add_subplot(gs1[panel1 - 1], projection=datacrs))
        if ev % 2 == 0:
            panel2 += 1
            axs.append(fig.add_subplot(gs2[panel2 - 1], projection=datacrs))
        ax1 = axs[-1]
        axs.append(ax1)
        # plot reanalysis
        if i % 2 != 0:
            cf1 = ax1.contourf(lonsre, latsre, obs_slp, norm=norm,
                               cmap=cmap)
            ax1.contour(lonsre, latsre, obs_slp, cf1.levels,colors='grey', linewidths=1)
            qv = ax1.quiver(lonsre[::2],latsre[::2],
                       ure.sel(lev=1000)[::2,::2],vre.sel(lev=1000)[::2,::2],
                       color= 'k')
            ax1.text(-53,-27.5,str(ev), fontsize = 18, bbox=props)
            ax1.quiverkey(qv,1.1, 1.07, 20, r'$10 \frac{m}{s}$', labelpos = 'E',
                           coordinates='axes', labelsep = 0.05,
                           fontproperties={'size': 14, 'weight': 'bold'})
            if ev < 3:
                ax1.text(-51,-25.8,'OBS.', fontsize=16)
        # plot OLAM
        else:
            cf2 = ax1.contourf(lonso, latso, olam_slp, norm=norm,
                               cmap=cmap)
            ax1.contour(lonso, latso, olam_slp, cf2.levels,colors='grey', linewidths=1)
            ax1.quiver(lonso[::20],latso[::20],uo[::20,::20],vo[::20,::20],
                       color= 'k')
            # colorbar
            pos = ax1.get_position()
            cbar_ax = fig.add_axes([pos.x1+0.005, pos.y0, 0.01, pos.height])
            cbar = plt.colorbar(cf1, cax=cbar_ax, orientation='vertical')
            cbar.ax.tick_params(labelsize=12)
            for label in cbar.ax.xaxis.get_ticklabels()[::1]:
                label.set_visible(False)
            if ev < 3:
                ax1.text(-51,-25.8,'OLAM.', fontsize=16)
            ev += 1
        # map cosmedics
        plot_background(ax1)
        gl = ax1.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
                      linewidth=1, color='gray', alpha=0.5,linestyle='--')
        gl.xlabels_top = False
        gl.ylabels_left = False
        if i % 2 == 0:
            gl.ylabels_right = False
        gl.xlocator = mticker.FixedLocator(range(-54,-40,2))
        gl.xlabel_style = {'size': 13, 'color': 'gray'}
        gl.ylabel_style = {'size': 14, 'color': 'gray'}
        ax1.outline_patch.set_edgecolor('gray')

    pl.savefig('./figures/slp_wind/slp_wind_panel.jpg', format='jpg')
    pl.savefig('./figures/slp_wind/slp_wind_panel.eps', format='eps', dpi=300)

def plot_wind_panel():
    '''
    Plot Wind Speed panels for each event peak time
    '''
    # create colormap
    col_hcl = [
     [0.9921568627450981, 0.6588235294117647, 0.7058823529411765],
     [0.9294117647058824, 0.4392156862745098, 0.6627450980392157],
     [0.8, 0.16470588235294117, 0.6470588235294118],
     [0.5294117647058824, 0.058823529411764705, 0.5254901960784314],
     [0.36470588235294116,0.1568627450980392, 0.39215686274509803],
     [0.3215686274509804, 0.2549019607843137, 0.4549019607843137],
     [0.1843137254901961, 0.4627450980392157, 0.5725490196078431],
     [0.0, 0.5843137254901961, 0.6862745098039216],
     [0.09411764705882353, 0.7411764705882353, 0.6901960784313725],
     [0.9450980392156862, 0.9450980392156862, 0.9450980392156862]
     ]
    col_hcl.reverse()
    cmap = LinearSegmentedColormap.from_list(
        'MyMap', col_hcl, N=20)
    cmap.set_under('white')
    # figure params
    fig = plt.figure(figsize=(10,15) , constrained_layout=False)
    gs1 = gridspec.GridSpec(6, 2, hspace=0.25, wspace=0.15, left=0.01, right=0.45)
    gs2 = gridspec.GridSpec(6, 2, hspace=0.25, wspace=0.15, left=0.50, right=0.94)
    axs = []
    datacrs = ccrs.PlateCarree()
    #
    ev = 1
    panel1 = 0
    panel2 = 0
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
    for i in range (1,25):
        # get data
        if i % 2 != 0:
            tmp = Get_SLP_Wind_At_EvPeak(ev,'n')
            obs = tmp[0].transpose('lat','lon','lev')
            olam = tmp[1].transpose('lat','lon', 'lev')
            # reanalysis data
            lonsre, latsre = obs.lon, obs.lat
            ure, vre = obs.U, obs.V
            wsre = np.sqrt(ure**2 + vre**2)
            wsre = wsre.sel(lev=1000)
            # olam data
            lonso, latso = olam.lon, olam.lat
            uo, vo = olam.uwnd, olam.vwnd
            wso = np.sqrt(uo**2 + vo**2)
            maxws = round(np.amax([int(np.amax(wso.values)), int(np.nanmax(wsre.values))]),-1)
            norm = colors.Normalize(vmin=0, vmax=maxws)
        # figure
        if ev % 2 != 0:
            panel1 += 1
            axs.append(fig.add_subplot(gs1[panel1 - 1], projection=datacrs))
        if ev % 2 == 0:
            panel2 += 1
            axs.append(fig.add_subplot(gs2[panel2 - 1], projection=datacrs))
        ax1 = axs[-1]
        axs.append(ax1)
        # plot reanalysis data
        if i % 2 != 0:
            cf1 = ax1.contourf(lonsre, latsre, wsre, norm=norm,
                               cmap=cmap)
            ax1.contour(lonsre, latsre, wsre, cf1.levels,colors='grey', linewidths=1)
            ax1.quiver(lonsre[::2],latsre[::2],
                       ure.sel(lev=1000)[::2,::2],vre.sel(lev=1000)[::2,::2],
                       color= 'k')
            ax1.text(-53,-27.5,str(ev), fontsize = 18, bbox=props)
            if ev < 3:
                ax1.text(-51,-25.8,'OBS.', fontsize=16)
        # plot OLAM data
        else:
            cf2 = ax1.contourf(lonso, latso, wso, norm=norm,
                               cmap=cmap)
            ax1.contour(lonso, latso, wso, cf2.levels,colors='grey', linewidths=1)
            ax1.quiver(lonso[::20],latso[::20],uo[::20,::20],vo[::20,::20],
                       color= 'k')
            # colorbar
            pos = ax1.get_position()
            cbar_ax = fig.add_axes([pos.x1+0.005, pos.y0, 0.01, pos.height])
            cbar = plt.colorbar(cf1, cax=cbar_ax, orientation='vertical')
            cbar.ax.tick_params(labelsize=12)
            for label in cbar.ax.xaxis.get_ticklabels()[::5]:
                label.set_visible(False)
            if ev < 3:
                ax1.text(-51,-25.8,'OLAM.', fontsize=16)
            ev += 1
        # map cosmedics
        plot_background(ax1)
        gl = ax1.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
                      linewidth=1, color='gray', alpha=0.5,linestyle='--')
        gl.xlabels_top = False
        gl.ylabels_left = False
        if i % 2 == 0:
            gl.ylabels_right = False
        gl.xlocator = mticker.FixedLocator(range(-54,-40,2))
        gl.xlabel_style = {'size': 13, 'color': 'gray'}
        gl.ylabel_style = {'size': 14, 'color': 'gray'}
        ax1.outline_patch.set_edgecolor('gray')

    pl.savefig('./figures/slp_wind/windspeed_panel.jpg', format='jpg')
    pl.savefig('./figures/slp_wind/windspeed_panel.eps', format='eps', dpi=300)

def obs_model_panel_slp():
    '''
    Make scater plot of observed  x simulated data,
        and also plot a linear regression line
    '''
    # fig params
    fig = plt.figure(figsize=(15,15))
    gs1 = gridspec.GridSpec(4, 3, hspace=0.2, wspace=0.4)
    axs = []
    cmap = cmocean.cm.balance
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
    for i in range(1,13):
        # fig
        axs.append(fig.add_subplot(gs1[i - 1]))
        ax1 = axs[-1]
        # get data
        tmp = Get_SLP_Wind_At_EvPeak(i,'y')
        obs, olam = tmp[0].transpose('lat','lev','lon'),tmp[1].transpose('lat','lev','lon')
        obs, olam = obs.SLP/100, olam.sslp/100
        # This first method returns the regression data
        #   in a comprehensive way, but not best for plot
        B0, B1, reg_line = st.linear_regression(obs.values, olam.values)
        R = st.Scorr(obs, olam)[0]
        text = ''' R^2: {}
        y = {} + {}X'''.format(round(R**2, 2),
                               round(B0, 2),
                               round(B1, 2))
        # This method is not great for presenting data
        #   but it is easier for plotting
        gradient, intercept, r_value, p_value, std_err = \
            stats.linregress((olam.values.ravel()),
                             (obs.values.ravel()))
        # Adjust min and max plot limits
        minslp = round(np.amin([int(np.amin(obs.values)), int(np.amin(olam.values))])-5,-1)
        maxslp = round(np.amax([int(np.amax(obs.values)), int(np.amax(olam.values))])+5,-1)
        if minslp > 1014:
            minslp = 1010
        if maxslp < 1025:
            maxslp = 1030
        x1=np.linspace(minslp,maxslp,500)
        y1=gradient*x1+intercept
        norm = colors.DivergingNorm(vmin=minslp, vcenter=1014, vmax=maxslp)
        # plot
        ax1.scatter(olam,obs,c=obs,cmap=cmap,norm=norm,alpha=0.85)
        ax1.plot(x1,y1,"k")
        # plot cosmedics
        if i > 9:
            ax1.set_xlabel('Olam', fontsize = 18)
        if i  == 1 or i == 4 or i == 7 or i == 10:
            ax1.set_ylabel('Reanalysis',fontsize = 18)
        ax1.text(0.1,0.85, str(i), fontsize = 18, transform=ax1.transAxes, bbox=props)
        ax1.text(0.3,0.1, s=text, fontsize = 12, transform=ax1.transAxes, bbox=props)
        plt.setp(ax1.xaxis.get_majorticklabels(), rotation=20)
        ax1.set_xlim(minslp,maxslp)
        ax1.set_ylim(minslp,maxslp)
        ax1.set_aspect('equal', 'datalim')
    pl.savefig('./figures/slp_wind/obs_x_olam_slp_panel.jpg', format='jpg')
    pl.savefig('./figures/slp_wind/obs_x_olam_slp_panel.eps', format='eps', dpi=300)

def obs_model_panel_wind():

    # create colormap
    col_hcl = [
     [0.9921568627450981, 0.6588235294117647, 0.7058823529411765],
     [0.9294117647058824, 0.4392156862745098, 0.6627450980392157],
     [0.8, 0.16470588235294117, 0.6470588235294118],
     [0.5294117647058824, 0.058823529411764705, 0.5254901960784314],
     [0.36470588235294116,0.1568627450980392, 0.39215686274509803],
     [0.3215686274509804, 0.2549019607843137, 0.4549019607843137],
     [0.1843137254901961, 0.4627450980392157, 0.5725490196078431],
     [0.0, 0.5843137254901961, 0.6862745098039216],
     [0.09411764705882353, 0.7411764705882353, 0.6901960784313725],
     [0.9450980392156862, 0.9450980392156862, 0.9450980392156862]
     ]
    col_hcl.reverse()
    cmap = LinearSegmentedColormap.from_list(
        'MyMap', col_hcl, N=20)
    # Fig params
    fig = plt.figure(figsize=(15,15))
    gs1 = gridspec.GridSpec(4, 3, hspace=0.2, wspace=0.4)
    axs = []
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
    for i in range(1,13):
        # fig
        axs.append(fig.add_subplot(gs1[i - 1]))
        ax1 = axs[-1]
        # get data
        tmp = Get_SLP_Wind_At_EvPeak(i,'y')
        obs, olam = tmp[0].transpose('lat','lev','lon'),tmp[1].transpose('lat','lev','lon')
        obs = np.sqrt(obs.U**2 + obs.V**2).sel(lev=1000)
        olam = np.sqrt(olam.uwnd**2 + olam.vwnd**2)
        # data limits
        minws = round(np.amin([int(np.nanmin(obs.values)), int(np.nanmin(olam.values))])-5,-1)
        maxws = round(np.amax([int(np.nanmax(obs.values)), int(np.nanmax(olam.values))])+5,-1)
        # This first method returns the regression data
        #   in a comprehensive way, but not best for plot
        mask = ~np.isnan(olam.values) & ~np.isnan(obs.values)
        B0, B1, reg_line = st.linear_regression(obs.values[mask], olam.values[mask])
        R = st.Scorr(obs, olam)[0]
        text = ''' R^2: {}
        y = {} + {}X'''.format(round(R**2, 2),
                               round(B0, 2),
                               round(B1, 2))
        # This method is not great for presenting data
        #   but it is easier for plotting
        mask = ~np.isnan(olam.values.ravel()) & ~np.isnan(obs.values.ravel())
        gradient, intercept, r_value, p_value, std_err = \
            stats.linregress(np.log(olam.values.ravel()[mask]),
                             np.log(obs.values.ravel()[mask]))
        x1=np.linspace(minws,maxws,500)
        y1=gradient*x1+intercept
        # plot
        ax1.scatter(np.log(olam),np.log(obs),c=obs,cmap=cmap,alpha=0.85)
        ax1.loglog(x1,y1,"k")
        # cosmedics
        if i > 9:
            ax1.set_xlabel('Olam', fontsize = 18)
        if i  == 1 or i == 4 or i == 7 or i == 10:
            ax1.set_ylabel('Reanalysis',fontsize = 18)
        ax1.text(0.1,0.85, str(i), fontsize = 18, transform=ax1.transAxes, bbox=props)
        ax1.text(0.3,0.1, s=text, fontsize = 12, transform=ax1.transAxes, bbox=props)
        ax1.set_xlim(minws,maxws)
        ax1.set_ylim(minws,maxws)
        ax1.set_aspect('equal', 'datalim')
    pl.savefig('./figures/slp_wind/obs_x_olam_wind_panel.jpg', format='jpg')
    pl.savefig('./figures/slp_wind/obs_x_olam_wind_panel.eps', format='eps', dpi=300)

def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100):
    new_cmap = colors.LinearSegmentedColormap.from_list(
        'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval, b=maxval),
        cmap(np.linspace(minval, maxval, n)))
    return new_cmap

def export_stats_SLP():
    '''
    Export statistics to csv file
    '''
    # save data into array
    arr = []
    for event in range(1,13):
        print('------------------------------------------------')
        print('making statistics for event '+str(event)+'...')
        ev = []
        # get data
        obs, olam = Get_SLP_Wind_At_EvPeak(event,'y')
        obs, olam = obs.SLP/100, olam.sslp/100
        # append stats to array
        ev.append(str(event))
        ev.append(np.amax(obs).values)
        ev.append(np.amin(obs).values)
        ev.append(np.std(obs).values)
        ev.append(np.amax(olam).values)
        ev.append(np.amin(olam).values)
        ev.append(np.std(olam).values)
        ev.append(st.BIAS(obs, olam))
        ev.append(st.MAE(obs, olam))
        ev.append(st.MSE(obs, olam))
        ev.append(st.MSE_diss(obs, olam))
        ev.append(st.MSE_disp(obs, olam))
        ev.append(st.Scorr(obs, olam)[0])
        ev.append(st.Scorr(obs, olam)[1])
        ev.append(st.Scorr(obs, olam)[2])
        ev.append(st.RMSE(obs, olam))
        ev.append(st.RMSE_bias(obs, olam))
        ev.append(np.std(st.di_acc(obs, olam)).values)
        ev.append(st.S_sqr(obs, olam))
        ev.append(st.ConcordanceIndex(obs, olam))
        ev.append(st.D_pielke(obs, olam))
        arr.append(ev)
    # write array into csv
    with open('SLP_validation_stats.csv','w') as f:
        writer = csv.writer(f)
        writer.writerow(['Event', 'max_obs', 'min_obs','obs_std','max_model',
                         'min_model', 'model_std', 'Bias',
                         'MAE','MSEtot','MSE_diss','MSE_disp',
                         'Corr_Pearson', 'Corr_Spearmann', 'Corr_Kendall',
                         'RMSE','RMSE_bias','diff_std','S_sqr','ConcordanceIndex',
                         'D_pielke'])
        writer.writerows(arr)

def Create_df(obs,olam):
    '''
    Create a Pandas dataframe for handling the data.
      (workaround for computating correlations with NaN values)
    '''
    data_obs = []
    for i in obs.values:
            for j in i:
                data_obs.append(j)
    data_olam = []
    for i in olam.values:
            for j in i:
                data_olam.append(j)
    diff = []
    for i in range(len(data_obs)):
            diff.append(data_obs[i]-data_olam[i])
    df = pd.DataFrame({'Obs': data_obs,
                       'olam': data_olam,
                       'diff': diff})
    return df

def export_stats_wind():
    '''
    Export statistics to csv file
    '''
    # save data into array
    arr = []
    for event in range(1,13):
        print('------------------------------------------------')
        print('making statistics for event '+str(event)+'...')
        ev = []
        # get data
        obs, olam = Get_SLP_Wind_At_EvPeak(event,'y')
        obs = np.sqrt(obs.U**2 + obs.V**2).sel(lev=1000)
        olam = np.sqrt(olam.uwnd**2 + olam.vwnd**2)
        # pass data to data frame
        df = Create_df(obs,olam)
        # append stats to array
        mean_model = float(olam.mean('lon').mean('lat').values)
        mean_obs = float(obs.mean('lon').mean('lat').values)
        ev.append(str(event))
        ev.append(round(mean_obs))
        ev.append(np.amax(obs).values)
        ev.append(np.std(obs).values)
        ev.append(round(mean_model))
        ev.append(np.amax(olam).values)
        ev.append(np.std(olam).values)
        ev.append(round(mean_obs - mean_model))
        ev.append(st.BIAS(obs, olam))
        ev.append(st.MAE(obs, olam))
        ev.append(st.MSE(obs, olam))
        ev.append(st.MSE_diss(obs, olam))
        ev.append(st.MSE_disp(obs, olam))
        ev.append(df.corr('pearson')['Obs']['olam'])
        ev.append(df.corr('spearman')['Obs']['olam'])
        ev.append(df.corr('kendall')['Obs']['olam'])
        ev.append(st.RMSE(obs, olam))
        ev.append(st.RMSE_bias(obs, olam))
        ev.append(np.std(st.di_acc(obs, olam)).values)
        ev.append(st.S_sqr(obs, olam))
        ev.append(st.ConcordanceIndex(obs, olam))
        ev.append(st.D_pielke(obs, olam))
        arr.append(ev)
    # write array into csv
    with open('wind_validation_stats.csv','w') as f:
        writer = csv.writer(f)
        writer.writerow(['Event', 'mean_obs','max_obs','obs_std','mean_model',
                         'max_model','model_std', 'Deltamean', 'Bias',
                         'MAE','MSEtot','MSE_diss','MSE_disp',
                         'Corr_Pearson', 'Corr_Spearmann', 'Corr_Kendall',
                         'RMSE','RMSE_bias','diff_std','S_sqr','ConcordanceIndex',
                         'D_pielke'])
        writer.writerows(arr)

def CorrlEachTime(event):
    '''
    Get spatial correlation values of wind Speed and SLP
        for each Data time step (3h).
    '''
    # getdata
    tmp = GetSLPWindData(event)
    olam,obs = tmp[0],tmp[1]
    # arrays for storing the data
    corrl_ws = []
    corrl_slp = []
    # loop through time steps
    for t in olam.time:
        # data for each time step
        obs_ = obs.sel(time=t)
        olam_ = olam.sel(time=t).assign(lev=obs.lev)
        olam_ = regridSLPWind(obs.lon,obs.lat,olam_,event)
        # calculate wind speed
        obs_ws = np.sqrt(obs_.U**2 + obs_.V**2).sel(lev=1000)
        olam_ws = np.sqrt(olam_.uwnd**2 + olam_.vwnd**2)
        # Pass data to Dataframe for handling NaNs
        df = Create_df(obs_ws,olam_ws)
        corrl_ws.append(df.corr('pearson')['Obs']['olam'])
        corrl_slp.append(st.Scorr(obs_.SLP, olam_.sslp)[0])
    return corrl_ws, corrl_slp, olam.time

def corrl_time_evo(prec):
    '''
    Creates a figure showing temporal evolution of correlation values
        for each event.
    Parameters
    ----------
    prec : float ('y' or 'n')
        Decides if plot correlation for accumulated preciptation data
    '''
    # Fig params
    fig = plt.figure(figsize=(15,15))
    gs1 = gridspec.GridSpec(4, 3, hspace=0.225, wspace=0.1)
    axs = []
    # box for plotting texts
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
    for i in range(1,13):
        # fig
        axs.append(fig.add_subplot(gs1[i - 1]))
        ax1 = axs[-1]
        # get data
        tmp = CorrlEachTime(i)
        corrl_ws, corrl_slp = tmp[0], tmp[1]
        time = tmp[2]
        # plot
        ax1.plot(time,corrl_ws,c='#836953',marker='o',
                 alpha=0.85,label='Wind Spd.', linewidth=2)
        ax1.plot(time,corrl_slp,c='#028e2c',marker='v',
                 alpha=0.85,  label='SLP', linewidth=2)
        # plot prec if requested
        if prec == 'y':
            corrl_prec = mean_corrl(i)[2]
            ax1.plot(time,corrl_prec,c='#54718F',marker='s',
                 alpha=0.85,  label='Acc. Prec.', linewidth=2)
        # figure cosmedics
        if i  == 1 or i == 4 or i == 7 or i == 10:
            ax1.set_ylabel('Corrl. I.',fontsize = 16)
        if i == 1:
            ax1.legend(fontsize=16)
        ax1.text(0.8,0.1, str(i), fontsize = 18, transform=ax1.transAxes, bbox=props)
        plt.setp(ax1.xaxis.get_majorticklabels(), rotation=20)
        ax1.set_ylim(0,1)
        ax1.grid(linestyle='--',linewidth=0.25,c='gray')
    if prec == 'n':
        pl.savefig('./figures/slp_wind/time_corrl.jpg', format='jpg')
        pl.savefig('./figures/slp_wind/time_corrl.eps', format='eps', dpi=300)
    elif prec == 'y':
        pl.savefig('./figures/slp_wind/time_corrl_slp_wind_prec.jpg', format='jpg')
        pl.savefig('./figures/slp_wind/time_corrl_slp_wind_prec.eps', format='eps', dpi=300)

def TestMannWithneyU_DailySLP():
    '''
    Performs a Mann Whitnney U test for daily means of SLP data

    Returns
    -------
    Dictionary containing test results for each event and each day
    '''
    # Dict
    MWUT = {}
    for i in range(1,13):
        # Create list for each event
        MWUT['Event '+str(i)] = []
        # Get data
        tmp = GetSLPWindData(i)
        olam,obs = tmp[0],tmp[1]
        # Get daily means
        olam = olam.sslp.resample(time='1D').mean('time')
        obs = obs.SLP.resample(time='1D').mean('time')
        # Daily corrl
        for d in olam.time:
            stat, p = stats.mannwhitneyu(obs.sel(time=d).values.ravel(), olam.sel(time=d).values.ravel())
            alpha = 0.05
            if p > alpha:
                MWUT['Event '+str(i)].append('Same distribution (fail to reject H0)')
            else:
                	MWUT['Event '+str(i)].append('Different distribution (reject H0)')
    return MWUT

def TestMannWithneyU_allSLP():
    '''
    Performs a Mann Whitnney U test for SLP data, for each time step

    Returns
    -------
    Dictionary containing test results for each event and each time step
    '''
    # Dict
    MWUT = {}
    for i in range(1,13):
        # Create list for each event
        MWUT['Event '+str(i)] = []
         # Get data
        tmp = GetSLPWindData(i)
        olam,obs = tmp[0].sslp,tmp[1].SLP
        # Daily corrl
        for d in olam.time:
            stat, p = stats.mannwhitneyu(obs.sel(time=d).values.ravel(), olam.sel(time=d).values.ravel())
            alpha = 0.05
            if p > alpha:
                MWUT['Event '+str(i)].append('Same distribution (fail to reject H0)')
            else:
                	MWUT['Event '+str(i)].append('Different distribution (reject H0)')
    return MWUT

def PlotBoxPLotSLP(which):
    '''
    Create Boxplots for daily mean SLP
    '''
    # Figure params
    fig = plt.figure(figsize=(15,15))
    gs1 = gridspec.GridSpec(4, 3, hspace=0.125, wspace=0.3)
    axs = []
    # colors for plots
    c1 ='#0077b6'
    c2 = '#ff6961'
    # box for plotting texts
    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
    for i in range(1,13):
        # get data
        tmp = GetSLPWindData(i)
        olam,obs = tmp[0],tmp[1]
        olam = olam.resample(time='1D').mean('time')
        obs = obs.resample(time='1D').mean('time')
        if which == 'slp':
            obs_data, olam_data = obs.SLP/100, olam.sslp/100
        elif which == 'wind':
            obs_data = np.sqrt(obs.U**2 + obs.V**2)
            olam_data = np.sqrt(olam.uwnd**2 + olam.vwnd**2)
        # pass data to Pandas Dataframe
        olam_df = DataToDataFrame(olam_data)
        obs_df = DataToDataFrame(obs_data)
        # mask NaNs in wind data
        if which == 'wind':
            mask = ~np.isnan(obs_df.values)
            obs_df = obs_df[mask]
        # Daily corrl
        olam_datar = regrid(obs.lon,obs.lat,olam_data,i)
        corrl = []
        for d in olam.time:
            corrl.append(st.Scorr(obs_data.sel(time=d), olam_datar.sel(time=d))[0])
            for ct in range(3):
                corrl.append(np.nan)
        xs = range(len(corrl))
        s1 = pd.Series(corrl, index=xs)
#       # figure
        axs.append(fig.add_subplot(gs1[i - 1]))
        ax1 = axs[-1]
        axs.append(ax1)
        ax2 = axs[-1].twinx()
        # ticks for the x-axis
        ticks = np.arange(0,len(obs_df.columns))
        # plot reanalysis
        ax1.boxplot(obs_df, labels=obs_df.columns,
        positions=np.array(ticks)*4-1, widths=1.6,
        showfliers=False,
        notch=True, patch_artist=True,
        boxprops=dict(facecolor=c1),
        medianprops=dict(color='k', linewidth=2),
        flierprops=dict(color='k'))
        # plot olam
        ax1.boxplot(olam_df, labels=olam_df.columns,
        positions=np.array(ticks)*4+1,widths=1.6,
        showfliers=False,
        notch=True, patch_artist=True,
        boxprops=dict(facecolor=c2),
        medianprops=dict(color='k', linewidth=2),
        flierprops=dict(color='k'))
        # plot correlation values
        ax2.plot(s1.dropna(),c='#028e2c',marker='v',
                 alpha=0.85,  label='Corrl.', linewidth=2)
        ax2.set_ylim(0,1)
        ax2.grid(linewidth=0.25,c='gray')
        # create fake labels
        if i == 1:
            plt.plot([], c=c1, label='Reanalysis')
            plt.plot([], c=c2, label='OLAM')
            plt.legend(fontsize=14,loc='upper right')
        # ticks
        plt.xticks(np.arange(0, (len(ticks)) * 4, 4), range(1,len(obs_df.columns)+1))
        # cosmedics
        ax1.text(.05,0.05, str(i), fontsize = 14,
                 transform=ax1.transAxes, bbox=props)
        ax1.tick_params(labelsize=12)
        ax2.tick_params(labelsize=12)
        plt.xlim(-3, (len(ticks)*4)+1)
        if i > 9:
            ax1.set_xlabel('Time (days)', fontsize = 16)
        if i  == 1 or i == 4 or i == 7 or i == 10:
            if which == 'slp':
                ax1.set_ylabel('SLP (hPa)',fontsize = 16)
            elif which == 'wind':
                ax1.set_ylabel('Wind Speed (m s-1)',fontsize = 16)
        if i  == 3 or i == 6 or i == 9 or i == 12:
                ax2.set_ylabel('Corrl I.',fontsize = 16)
    pl.savefig('./figures/slp_wind/daily_boxplot'+str(which)+'.jpg', format='jpg')
    pl.savefig('./figures/slp_wind/daily_boxplot'+str(which)+'.eps', format='eps', dpi=300)

# ----------
if __name__ == "__main__":
    plot_slp_panel()
    plot_wind_panel()
    obs_model_panel_slp()
    obs_model_panel_wind()
    export_stats_SLP()
    export_stats_wind()
    corrl_time_evo('n')
    corrl_time_evo('y')
    PlotBoxPLotSLP('slp')
    PlotBoxPLotSLP('wind')

reason: import numpy,from scipy


text:
#!/usr/bin/env python
# coding: utf-8

# # Covid19 Vaccine Analysis

# Many vaccines have been introduced so far to fight covid-19. No vaccine has guaranteed 100% accuracy so far, but most manufacturing companies claim their vaccine is not 100% accurate, but still, it will save your life by giving you immunity.
#
# Thus, each country tries to vaccinate a large part of its population so as not to depend on a single vaccine. That’s I am going to analysis in this project, which is how many vaccines each country is using to fight covid-19. In the section below, I have made my project on Covid-19 vaccines analysis with Python.

# ### Importing the necessary Python libraries and the dataset

# In[3]:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
data = pd.read_csv("country_vaccinations.csv")
data.head()

# ### Exploring this data before we start analyzing the vaccines taken by countries

# In[4]:

data.describe()

# In[5]:

pd.to_datetime(data.date)

# In[6]:

data.country.value_counts()

# The United Kingdom is made up of England, Scotland, Wales, and Northern Ireland. But in the above data, these countries are mentioned separately with the same values as in the United Kingdom. So this may be an error while recording this data. So for fixing this error:

# In[7]:

data = data[data.country.apply(lambda x: x not in ["England", "Scotland", "Wales", "Northern Ireland"])]
data.country.value_counts()

# ### Exploring the vaccines available in this dataset

# In[8]:

data.vaccines.value_counts()

# So we have almost all the Covid-19 vaccines available in this dataset. Now I will create a new DataFrame by only selecting the vaccine and the country columns to explore which vaccine is taken by which country:

# In[9]:

df = data[['vaccines', 'country']]
df.head()

# #### Now let’s see how many countries are taking each of the vaccines mentioned in this data:

# In[10]:

dict_ = {}
for i in df.vaccines.unique():
    dict_[i] = [df['country'][j] for j in df[df['vaccines'] == i].index]

vaccines = {}
for key, value in dict_.items():
    vaccines[key] = set(value)
for i,j in vaccines.items():
    print(f'{i}:>>{j}')

# ### Now let’s visualize this data to have a look at what combination of vaccines every country is using:

# In[11]:

import plotly.express as px
import plotly.offline as py

vaccine_map = px.choropleth(data, locations = 'iso_code', color = 'vaccines')
vaccine_map.update_layout(height = 300, margin = {'r':0, 't':0, 'l':0, 'b':0})
vaccine_map.show()

# In[ ]:


text:
import qimpy as qp
import numpy as np
from scipy.special import sph_harm
from typing import Sequence, Any, List, Tuple

def get_harmonics_ref(l_max: int, r: np.ndarray) -> np.ndarray:
    """Reference real solid harmonics based on SciPy spherical harmonics."""
    rMag = np.linalg.norm(r, axis=-1)
    theta = np.arccos(r[..., 2] / rMag)
    phi = np.arctan2(r[..., 1], r[..., 0])
    phi += np.where(phi < 0.0, 2 * np.pi, 0)
    results = []
    for l in range(l_max + 1):
        result = np.zeros((2 * l + 1,) + r.shape[:-1])
        for m in range(0, l + 1):
            ylm = ((-1) ** m) * (rMag ** l) * sph_harm(m, l, phi, theta)
            if m == 0:
                result[l] = ylm.real
            else:
                result[l + m] = np.sqrt(2) * ylm.real
                result[l - m] = np.sqrt(2) * ylm.imag
        results.append(result)
    return np.concatenate(results, axis=0)

def get_lm(l_max: int) -> List[Tuple[int, int]]:
    """Get list of all (l,m) in order up to (and including) l_max"""
    return [(l, m) for l in range(l_max + 1) for m in range(-l, l + 1)]

def format_array(array: Sequence[Any], fmt: str) -> str:
    """Convert `array` to string with format `fmt` for each entry."""
    return "[" + ", ".join(fmt.format(elem) for elem in array) + "]"

def generate_harmonic_coefficients(l_max_hlf: int) -> None:
    """Generate tables of recursion coefficients for computing real
    solid harmonics up to l_max = 2 * l_max_hlf, as well as tables of
    product coefficients (Clebsch-Gordon coefficients) for real solid
    harmonics up to order l_max_hlf. Print results formatted as Python
    code that can be pasted into _spherical_harmonics_data.py."""
    l_max = 2 * l_max_hlf
    qp.log.info(
        "from typing import List, Tuple, Dict\n\n"
        f"L_MAX: int = {l_max}  # Maximum l for harmonics\n"
        f"L_MAX_HLF: int = {l_max_hlf}  # Maximum l for products"
    )
    # Calculate solid harmonics on a mesh covering unit cube:
    grids1d = 3 * (np.linspace(-1.0, 1.0, 2 * l_max),)  # avoids zero
    r = np.array(np.meshgrid(*grids1d)).reshape(3, -1).T
    r_sq = (r ** 2).sum(axis=-1)
    ylm = get_harmonics_ref(l_max, r)
    # Calculate recursion coefficients:
    ERR_TOL = 1e-14
    COEFF_TOL = 1e-8
    qp.log.info(
        "CooIndices = Tuple[List[int], List[int], List[float]]\n\n"
        "# Recursion coefficients for computing real harmonics at l>1\n"
        "# from products of those at l = 1 and l-1. The integers index\n"
        "# a sparse matrix with (2l+1) rows and 3*(2l-1) columns.\n"
        "YLM_RECUR: List[CooIndices] = ["
    )
    Y_00 = np.sqrt(0.25 / np.pi)
    Y_1m_prefac = np.sqrt(0.75 / np.pi)
    qp.log.info(f"    ([], [], [{Y_00:.16f}]), ([], [], [{Y_1m_prefac:.16f}]),")
    for l in range(2, l_max + 1):
        l_minus_1_slice = slice((l - 1) ** 2, l ** 2)
        y_product = ylm[l_minus_1_slice, None, :] * ylm[None, 1:4, :]
        y_product = y_product.reshape((2 * l - 1) * 3, -1)
        index_row = []
        index_col = []
        values = []
        for m in range(-l, l + 1):
            # List pairs of m at l = 1 and l-1 that can add up to m:
            m_pairs_all = set(
                [
                    (sign * m + dsign * dm, dm)
                    for sign in (-1, 1)
                    for dsign in (-1, 1)
                    for dm in (-1, 0, 1)
                ]
            )
            m_pairs = [m_pair for m_pair in m_pairs_all if abs(m_pair[0]) < l]
            m_pair_indices = [3 * (l - 1 + m) + (1 + dm) for m, dm in m_pairs]
            # Solve for coefficients of the linear combination:
            for n_sel in range(1, len(m_pair_indices) + 1):
                # Try increasing numbers till we get one:
                y_product_allowed = y_product[m_pair_indices[:n_sel]]
                y_target = ylm[l * (l + 1) + m]
                coeff = np.linalg.lstsq(y_product_allowed.T, y_target, rcond=None)[0]
                residual = np.dot(coeff, y_product_allowed) - y_target
                err = np.linalg.norm(residual) / np.linalg.norm(y_target)
                if err < ERR_TOL:
                    break
            assert err < ERR_TOL
            # Select non-zero coefficients to form product expansion:
            sel = np.where(np.abs(coeff) > COEFF_TOL * np.linalg.norm(coeff))[0]
            indices = np.array(m_pair_indices)[sel]
            coeff = coeff[sel]
            # Sort by index and add to lists for current l:
            sort_index = indices.argsort()
            index_row += [l + m] * len(sort_index)
            index_col += list(indices[sort_index])
            values += list(coeff[sort_index])
        # Format as python code:
        qp.log.info(
            f"    ("
            f"{format_array(index_row, '{:d}')}, "
            f"{format_array(index_col, '{:d}')}, "
            f"{format_array(values, '{:.16f}')}),"
        )
    qp.log.info("]\n")
    # Calculate Clebsch-Gordon coefficients:
    lm_hlf = get_lm(l_max_hlf)
    qp.log.info(
        "# Clebsch-Gordon coefficients for products of real harmonics.\n"
        "# The integer indices correspond to l*(l+1)+m for each (l,m).\n"
        "YLM_PROD: Dict[Tuple[int, int],"
        " Tuple[List[int], List[float]]] = {"
    )
    for ilm1, (l1, m1) in enumerate(lm_hlf):
        for ilm2, (l2, m2) in enumerate(lm_hlf[: ilm1 + 1]):
            # List (l,m) pairs allowed by angular momentum addition rules:
            m_allowed = {m1 + m2, m1 - m2, m2 - m1, -(m1 + m2)}
            l_allowed = range(l1 - l2, l1 + l2 + 1, 2)
            lm_all = np.array(
                [(l, m) for l in l_allowed for m in m_allowed if (abs(m) <= l)]
            )
            l_all = lm_all[:, 0]
            m_all = lm_all[:, 1]
            ilm = l_all * (l_all + 1) + m_all  # flattened index
            # Solve for coefficients of the linear combination:
            y_product = ylm[ilm1] * ylm[ilm2]
            y_terms = ylm[ilm] * (r_sq[None, :] ** ((l1 + l2 - l_all) // 2)[:, None])
            results = np.linalg.lstsq(y_terms.T, y_product, rcond=None)
            coeff = results[0]
            err = np.sqrt(results[1][0]) / np.linalg.norm(y_product)
            assert err < ERR_TOL
            # Select non-zero coefficients to form product expansion:
            sel = np.where(np.abs(coeff) > COEFF_TOL * np.linalg.norm(coeff))[0]
            ilm = ilm[sel]
            coeff = coeff[sel]
            # Sort by (l,m):
            sort_index = ilm.argsort()
            ilm = ilm[sort_index]
            coeff = coeff[sort_index]
            # Format as python code:
            qp.log.info(
                f"    ({ilm1}, {ilm2}): ("
                f"{format_array(ilm, '{:d}')}, "
                f"{format_array(coeff, '{:.16f}')}),"
            )
    qp.log.info("}")

def main():
    qp.rc.init()
    assert qp.rc.n_procs == 1  # no MPI
    qp.utils.log_config()  # after rc to suppress header messages
    generate_harmonic_coefficients(l_max_hlf=3)

if __name__ == "__main__":
    main()