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

# # Matrix Basis Tutorial
# 
# Consider the space of density matrices corresponding to a Hilbert space $\mathcal{H}$ of dimension $d$.  The basis used for this Hilbert-Schmidt space, $B(\mathcal{H})$, can be any set of $d\times d$ matrices which span the density matrix space.  pyGSTi supports arbitrary bases by deriving from the `pygsti.tools.Basis` class, and constains built-in support for the following basis sets:
# 
# - the matrix unit, or "standard" basis, consisting of the matrices with a single unit (1.0) element and otherwise zero.  This basis is selected by passing `"std"` to appropriate function arguments.
# - the Pauli-product basis, consisting of tensor products of the four Pauli matrices {I, X, Y, Z} normalized so that $Tr(B_i B_j) = \delta_{ij}$.  All of these matrices are Hermitian, so that Hilbert-Schmidt vectors and matrices are real when this basis is used. This basis can only be used when the $d = 4^i$ for integer $i$, and is selected using the string `"pp"`.
# - the Gell-Mann basis, consisting of the normalized Gell-Mann matrices (see Wikipedia if you don't know what these are).  Similar to the Pauli-product case, these matrices are also Hermitian, so that Hilbert-Schmidt vectors and matrices are real when this basis is used.  Unlike the Pauli-product case, since Gell-Mann matrices are well defined in any dimension, the Gell-Mann basis is *not* restricted to cases when $d=4^i$.  This basis is selected using the string `"gm"`.
# - a special basis of $3 \times 3$ matricies designed for Qutrit systems formed by taking the symmetric subspace of a 2-qubit system.  This basis is selected using the string `"qt"`. 
# 
# Numerous functions within pyGSTi require knowledge of what Hilbert-Schmidt basis is being used.  The `pygsti.objects.Basis` object encapsulates a basis, and is the most flexible way of specifying a basis in pyGSTi.  Alternatively, many functions also accept the short strings `"std"`, `"gm"`, `"pp"`, and `"qt"` to select one of the standard bases.  In this tutorial, we'll demonstrate how to create a `Basis` object and use it and related functions to obtain and change the basis of the gate matrices and SPAM vectors stored in a `GateSet`.

# The most straightforward way to create a `Basis` object is to provide its short name and dimension:

# In[1]:


from pygsti import Basis
pp  = Basis('pp',  2)
std = Basis('std', 2)
gm  = Basis('gm',  2)
qt  = Basis('qt',  3) # qt must be dim 3
bases = [pp, std, gm, qt]


# Each of the `pp`, `std`, and `gm` bases created will have $4$ $2x2$ matrices each. The `qt` basis has $9$ $3x3$ matrices instead:

# In[2]:


for basis in bases:
    print('{} basis (dim {}):'.format(basis.name, basis.dim.dmDim))
    print('{} elements:'.format(len(basis)))
    for element in basis:
        print(element)


# However, custom bases can be easily created by supplying matrices:

# In[3]:


import numpy as np
std2x2Matrices = [
        np.array([[1, 0],
                  [0, 0]]),
        np.array([[0, 1],
                  [0, 0]]),
        np.array([[0, 0],
                  [1, 0]]),
        np.array([[0, 0],
                  [0, 1]])]

alt_standard = Basis(matrices=std2x2Matrices,
                     name='std',
                     longname='Standard')
print(alt_standard)


# More complex bases can be created by chaining other bases together along the diagonal. For example, a composition of the $2x2$ `std` basis with the $1x1$ `std` basis leads to a basis with state vectors of length $5$, or $5x5$ matrices:

# In[4]:


comp = Basis('std', [2, 1])
comp = Basis([('std', 2), ('std', 1)])
comp = Basis([std, ('std', 1)])
comp = Basis([Basis('std', 2), Basis('std', 1)])
print(comp)
for element in comp:
    print(element)


# ### Basis usage
# Once created, bases are used to manipulate matrices and vectors within pygsti:

# In[5]:


from pygsti.tools import change_basis, flexible_change_basis

mx = np.array([[1, 0, 0, 1],
               [0, 1, 2, 0],
               [0, 2, 1, 0],
               [1, 0, 0, 1]])

change_basis(mx, 'std', 'gm') # shortname lookup
change_basis(mx, std, gm)     # object only
change_basis(mx, std, 'gm')   # combination


# Composite bases can be converted between expanded and contracted forms:

# In[6]:


mxInStdBasis = np.array([[1,0,0,2],
                         [0,0,0,0],
                         [0,0,0,0],
                         [3,0,0,4]],'d')

begin = Basis('std', 2)
end   = Basis('std', [1,1])
mxInReducedBasis = flexible_change_basis(mxInStdBasis, begin, end)
print(mxInReducedBasis)
original         = flexible_change_basis(mxInReducedBasis, end, begin)


# In[ ]: