# Notebook from ShepherdCode/ShepherdML
Path: Workshop/MLP_309.ipynb

# MLP 309
_____no_output_____

NC_FILENAME='ncRNA.gc34.unprocessed.fasta'
PC_FILENAME='pcRNA.gc34.unprocessed.fasta'
DATAPATH=""
try:
    from google.colab import drive
    IN_COLAB = True
    PATH='/content/drive/'
    drive.mount(PATH)
    DATAPATH=PATH+'My Drive/data/'  # must end in "/"
    NC_FILENAME = DATAPATH+NC_FILENAME
    PC_FILENAME = DATAPATH+PC_FILENAME
except:
    IN_COLAB = False
    DATAPATH=""

EPOCHS=200
SPLITS=5
K=3
VOCABULARY_SIZE=4**K+1   # e.g. K=3 => 64 DNA K-mers + 'NNN'
EMBED_DIMEN=16
FILENAME='MLP305'   # load not save!
NEURONS=16Drive already mounted at /content/drive/; to attempt to forcibly remount, call drive.mount("/content/drive/", force_remount=True).
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import ShuffleSplit
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import RepeatedKFold
from sklearn.model_selection import StratifiedKFold
import tensorflow as tf
from tensorflow import keras
from keras.wrappers.scikit_learn import KerasRegressor
from keras.models import Sequential
from keras.layers import Bidirectional
from keras.layers import GRU
from keras.layers import Dense
from keras.layers import LayerNormalization
import time
dt='float32'
tf.keras.backend.set_floatx(dt)_____no_output_____

## Build model_____no_output_____

def compile_model(model):
    adam_default_learn_rate = 0.001
    schedule = tf.keras.optimizers.schedules.ExponentialDecay(
        initial_learning_rate = adam_default_learn_rate*10,
        #decay_steps=100000, decay_rate=0.96, staircase=True)
        decay_steps=10000, decay_rate=0.99, staircase=True)
    # learn rate = initial_learning_rate * decay_rate ^ (step / decay_steps)
    opt = tf.keras.optimizers.Adam(learning_rate=schedule)
    bc=tf.keras.losses.BinaryCrossentropy(from_logits=False)
    print("COMPILE...")
    model.compile(loss=bc, optimizer=opt, metrics=["accuracy"])
    print("...COMPILED")
    return model

def build_model(maxlen):
    act="elu"
    #embed_layer  = keras.layers.Embedding(
    #    VOCABULARY_SIZE,EMBED_DIMEN,input_length=maxlen);
    dense1_layer = keras.layers.Dense(NEURONS, activation=act,dtype=dt,
                                      input_dim=VOCABULARY_SIZE)
    dense2_layer = keras.layers.Dense(NEURONS, activation=act,dtype=dt)
    #dense3_layer = keras.layers.Dense(NEURONS, activation=act,dtype=dt)
    output_layer = keras.layers.Dense(1,  activation="sigmoid",dtype=dt)
    mlp = keras.models.Sequential()
    #mlp.add(embed_layer)
    mlp.add(dense1_layer)
    mlp.add(dense2_layer)
    #mlp.add(dense3_layer)
    mlp.add(output_layer)
    mlpc = compile_model(mlp)
    return mlpc_____no_output_____

## Load and partition sequences_____no_output_____

# Assume file was preprocessed to contain one line per seq.
# Prefer Pandas dataframe but df does not support append.
# For conversion to tensor, must avoid python lists.
def load_fasta(filename,label):
    DEFLINE='>'
    labels=[]
    seqs=[]
    lens=[]
    nums=[]
    num=0
    with open (filename,'r') as infile:
        for line in infile:
            if line[0]!=DEFLINE:
                seq=line.rstrip()
                num += 1   # first seqnum is 1
                seqlen=len(seq)
                nums.append(num)
                labels.append(label)
                seqs.append(seq)
                lens.append(seqlen)
    df1=pd.DataFrame(nums,columns=['seqnum'])
    df2=pd.DataFrame(labels,columns=['class'])
    df3=pd.DataFrame(seqs,columns=['sequence'])
    df4=pd.DataFrame(lens,columns=['seqlen'])
    df=pd.concat((df1,df2,df3,df4),axis=1)
    return df

def separate_X_and_y(data):
    y=   data[['class']].copy()
    X=   data.drop(columns=['class','seqnum','seqlen'])
    return (X,y)

_____no_output_____

## Make K-mers_____no_output_____

def make_kmer_table(K):
    npad='N'*K
    shorter_kmers=['']
    for i in range(K):
        longer_kmers=[]
        for mer in shorter_kmers:
            longer_kmers.append(mer+'A')
            longer_kmers.append(mer+'C')
            longer_kmers.append(mer+'G')
            longer_kmers.append(mer+'T')
        shorter_kmers = longer_kmers
    all_kmers = shorter_kmers
    kmer_dict = {}
    kmer_dict[npad]=0
    value=1
    for mer in all_kmers:
        kmer_dict[mer]=value
        value += 1
    return kmer_dict

KMER_TABLE=make_kmer_table(K)

def strings_to_vectors(data,uniform_len):
    all_seqs=[]
    for seq in data['sequence']:
        i=0
        seqlen=len(seq)
        kmers=[]
        while i < seqlen-K+1 -1:  # stop at minus one for spaced seed
            #kmer=seq[i:i+2]+seq[i+3:i+5]    # SPACED SEED 2/1/2 for K=4
            kmer=seq[i:i+K]
            i += 1
            value=KMER_TABLE[kmer]
            kmers.append(value)
        pad_val=0
        while i < uniform_len:
            kmers.append(pad_val)
            i += 1
        all_seqs.append(kmers)
    pd2d=pd.DataFrame(all_seqs)
    return pd2d   # return 2D dataframe, uniform dimensions_____no_output_____def make_kmers(MAXLEN,train_set):
    (X_train_all,y_train_all)=separate_X_and_y(train_set)
    X_train_kmers=strings_to_vectors(X_train_all,MAXLEN)
    # From pandas dataframe to numpy to list to numpy
    num_seqs=len(X_train_kmers)
    tmp_seqs=[]
    for i in range(num_seqs):
        kmer_sequence=X_train_kmers.iloc[i]
        tmp_seqs.append(kmer_sequence)
    X_train_kmers=np.array(tmp_seqs)
    tmp_seqs=None
    labels=y_train_all.to_numpy()
    return (X_train_kmers,labels)_____no_output_____def make_frequencies(Xin):
    Xout=[]
    VOCABULARY_SIZE= 4**K + 1  # plus one for 'NNN'
    for seq in Xin:
        freqs =[0] * VOCABULARY_SIZE
        total = 0
        for kmerval in seq:
            freqs[kmerval] += 1
            total += 1
        for c in range(VOCABULARY_SIZE):
            freqs[c] = freqs[c]/total
        Xout.append(freqs)
    Xnum = np.asarray(Xout)
    return (Xnum)
def make_slice(data_set,min_len,max_len):
    slice = data_set.query('seqlen <= '+str(max_len)+' & seqlen>= '+str(min_len))
    return slice_____no_output_____

## Cross validation_____no_output_____

from sklearn.metrics import confusion_matrix
def do_evaluation(X,y):
    cv_scores = []
    fold=1
    while fold<=5:
        X_valid=X
        y_valid=y
        bestname=DATAPATH+FILENAME+".cv."+str(fold)+".best"
        best_model=keras.models.load_model(bestname)
        scores = best_model.evaluate(X_valid, y_valid, verbose=0)
        #print("%s: %.2f%%" % (best_model.metrics_names[1], scores[1]*100))
        print("%.2f%% "%(scores[1]*100), end=' ')
        cv_scores.append(scores[1] * 100)

        y_pred = best_model.predict(X_valid)
        y_pred = [1 if n >= 0.5 else 0 for n in y_pred]
        cm=confusion_matrix(y_valid,y_pred)
        print("confusion")
        print(cm)

        fold += 1
    print()
    print("mean %.2f%% (+/- %.2f%%)" % (np.mean(cv_scores), np.std(cv_scores)))_____no_output_____

## Train on RNA lengths 200-1Kb_____no_output_____

MINLEN=200
MAXLEN=1000
print("Load data from files.")
nc_seq=load_fasta(NC_FILENAME,0)
pc_seq=load_fasta(PC_FILENAME,1)
train_set=pd.concat((nc_seq,pc_seq),axis=0)
nc_seq=None
pc_seq=None
print("Ready: train_set")
print ("Data prep")
subset=make_slice(train_set,MINLEN,MAXLEN)# One array to two: X and y
print ("Data reshape")
(X_train,y_train)=make_kmers(MAXLEN,subset)
X_train=make_frequencies(X_train)Load data from files.
Ready: train_set
Data prep
Data reshape
_____no_output_____print ("Evaluation")
do_evaluation(X_train,y_train)
print ("Done")Evaluation
75.92%  confusion
[[21322  2625]
 [12013 24832]]
74.80%  confusion
[[21876  2071]
 [13249 23596]]
75.63%  confusion
[[21727  2220]
 [12592 24253]]
75.56%  confusion
[[21575  2372]
 [12484 24361]]
76.19%  confusion
[[21214  2733]
 [11739 25106]]

mean 75.62% (+/- 0.47%)
Done
_____no_output__________no_output_____


# Notebook from lhqing/snmCAT-seq_integration
Path: input/prepare.ipynb

import anndata
import xarray as xr
import pandas as pd
import scipy.sparse as ss
import h5py_____no_output_____use_genes1 = pd.read_csv('Seurat_HVGs.txt', index_col=0).index
use_genes1 = use_genes1.str.replace('-', '_')
use_genes2 = pd.read_csv('SCF_hvgs.txt', index_col=0).index

use_genes = use_genes1 | use_genes2
use_genes.name = 'gene'_____no_output_____mch_meta = pd.read_csv('human_frontal_mch_metadata.tsv.gz', sep='\t', index_col=0)
mch_meta['cell_id'] = mch_meta.index.str[:-4]

rna_meta = pd.read_csv('human_frontal_rna_metadata.tsv.gz', sep='\t', index_col=0)
rna_meta['cell_id'] = mch_meta.index.str[:-4]_____no_output_____

## RNA_____no_output_____

mcds = xr.open_mfdataset('/home/hanliu/project/human_multiomic/mCTN/dataset/snmCT_*.withRNA.mcds',
                         concat_dim='cell', combine='nested').squeeze()
rna_data = mcds['rna_da'].load().sel(cell=rna_meta['cell_id'].values, gene=use_genes).to_pandas()_____no_output_____rna_adata = anndata.AnnData(X=rna_data.values,
                            obs=rna_meta,
                            var=pd.DataFrame([], index=rna_data.columns))
rna_adata.write_h5ad('rna_adata.h5ad', compression=True)
rna_adata... storing 'Technology' as categorical
... storing 'Brain Region' as categorical
... storing 'AD Index' as categorical
... storing 'ClusterAnno' as categorical
... storing 'MajorType' as categorical
... storing 'SubMarker' as categorical
... storing 'SubClusterAnno' as categorical
... storing 'MajorMarker' as categorical
... storing 'Individual ID' as categorical
... storing 'Individual Label' as categorical
... storing 'm3c_cluster' as categorical

## mC Frac_____no_output_____

mcds = xr.open_mfdataset(
    '/home/hanliu/project/human_multiomic/mCTN/dataset/ForSeurat/gene_frac/snmCT_*.withRNA.mcdsgene_da_rate.nc',
    concat_dim='cell',
    combine='nested').squeeze()
mch_data = mcds.sel(mc_type='HCHN')['gene_da_rate'].load().sel(
    cell=mch_meta['cell_id'].values, gene=mcds.get_index('gene') & use_genes).to_pandas()_____no_output_____mch_adata = anndata.AnnData(X=mch_data.values,
                            obs=mch_meta,
                            var=pd.DataFrame([], index=mch_data.columns))
mch_adata.write_h5ad('mch_adata.h5ad', compression=True)
mch_adata... storing 'Technology' as categorical
... storing 'Brain Region' as categorical
... storing 'AD Index' as categorical
... storing 'ClusterAnno' as categorical
... storing 'MajorType' as categorical
... storing 'SubMarker' as categorical
... storing 'SubClusterAnno' as categorical
... storing 'MajorMarker' as categorical
... storing 'Individual ID' as categorical
... storing 'Individual Label' as categorical
... storing 'm3c_cluster' as categorical

## Prepare LIGER_____no_output_____

use_genes = pd.read_csv('../input_archive/human_frontal_mch_hvfeatures.reverse.tsv.gz',
                        index_col=0, sep='\t').index
use_genes.name = 'gene'
use_genes_____no_output_____rna_adata = anndata.read_h5ad('rna_adata.h5ad')
mc_adata = anndata.read_h5ad('mch_adata.h5ad')_____no_output_____use_genes = use_genes & rna_adata.var_names & mc_adata.var_names
use_genes_____no_output_____rna_adata = rna_adata[:, use_genes].copy()
csr_x = ss.csr_matrix(rna_adata.X)
with h5py.File('rna.h5', 'w') as f:
    f['matrix/barcodes'] = rna_adata.obs_names.values
    f['matrix/data'] = csr_x.data
    f['matrix/indices'] = csr_x.indices
    f['matrix/indptr'] = csr_x.indptr
    f['matrix/features/name'] = rna_adata.var_names.values_____no_output_____mc_adata = mc_adata[:, use_genes].copy()
mc_adata.X = mc_adata.X.max() - mc_adata.X

csr_x = ss.csr_matrix(mc_adata.X)
with h5py.File('mch.h5', 'w') as f:
    f['matrix/barcodes'] = mc_adata.obs_names.values
    f['matrix/data'] = csr_x.data
    f['matrix/indices'] = csr_x.indices
    f['matrix/indptr'] = csr_x.indptr
    f['matrix/features/name'] = mc_adata.var_names.values_____no_output_____


# Notebook from carushi/Catactor
Path: example/mini_catactor_tutorial.ipynb

# Example 1: Computation of meta-analytic marker gene enrichment using mini_catactor
* Use 5 marker gene sets as default (SF, SC, CU, TA, TN + SM=SF for neuronal subtypes)
* Input: scanpy object (row=cell, column=gene name)
* Output: scanpy object, pickled file, or csv file_____no_output_____

import pickle
import scanpy as sc
import sys
from Catactor import mini_catactor
scanpy_object_path = './scanpy_obj.pyn'
with open(scanpy_object_path, "rb") as f:
     adata = pickle.load(f)

# tsne_1 and tsne_2 columns are used as a dimension reduction space
dimension='tsne_'
# Return a scanpy object with new columns for signal enrichment
new_adata_with_signal = mini_catactor.run_mini_catactor(adata, plot_gene=True, dimension=dimension)
print(new_adata_with_signal.obs)
# Write to a pickle object
new_adata_with_signal = mini_catactor.run_mini_catactor(adata, output_ext='pyn', dimension=dimension)
# Write to a csv file
new_adata_with_signal = mini_catactor.run_mini_catactor(adata, output_ext='csv', dimension=dimension)

# Plot meta-analytic marker signals
new_adata_with_signal = mini_catactor.run_mini_catactor(adata, plot=True, dimension=dimension)

# Plot each single marker gene enrichment
new_adata_with_signal = mini_catactor.run_mini_catactor(adata, plot_gene=True, dimension=dimension)
_____no_output_____

# Example 2: Computation of pseudo-bulk profiles for each cluster or top cells showing marker signal enrichment
* Input: scanpy object (row=cell, column=genomic bin)
* Additional input: scanpy object (row=cell, column=gene name)
* Output: 4 csv files for pseudo-bulk profiles_____no_output_____

from Catactor import pseudo_bulk
scanpy_bin_object_path = './scanpy_bin_obj.pyn'
with open(scanpy_bin_object_path, "rb") as f:
     pdata = pickle.load(f)

# Pseudo-bulk profiles for each cluster
pseudo_bulk.run_average_profiling(pdata, cluster=['cluster'])

# Pseudo-bulk profiles for top 500 cells in terms of the marker gene signal computed from reference
pseudo_bulk.run_average_profiling(pdata, reference=scanpy_obj_path)

# Add equally-divided genomic bin index (global_index_1000, global_index_2000, ...)
pseudo_bulk.genome_binning(pdata.var)

# Matrix conversion
cell_group, cell_original = 'global_index', 'global_index'
gene_group, gene_original = 'global_index_5000', 'global_index_1000'
new_pdata = pseudo_bulk.convert_row_and_column(pdata, gene_group, gene_original, cell_group, cell_original)
print(new_pdata)
_____no_output_____


# Notebook from google-research/proteinfer
Path: colabs/Price.ipynb

!git clone https://github.com/google-research/proteinfer
!pip install -q -r proteinfer/requirements.txt
%cd proteinfer
Cloning into 'proteinfer'...
remote: Enumerating objects: 612, done.
remote: Counting objects: 100% (335/335), done.
remote: Compressing objects: 100% (286/286), done.
remote: Total 612 (delta 113), reused 160 (delta 31), pack-reused 277
Receiving objects: 100% (612/612), 88.74 MiB | 26.54 MiB/s, done.
Resolving deltas: 100% (252/252), done.
     |████████████████████████████████| 99 kB 5.4 MB/s
     |████████████████████████████████| 2.3 MB 37.1 MB/s
     |████████████████████████████████| 10.8 MB 51.6 MB/s
     |████████████████████████████████| 2.8 MB 50.4 MB/s
     |████████████████████████████████| 50 kB 8.1 MB/s
     |████████████████████████████████| 59 kB 7.0 MB/s
     |████████████████████████████████| 89 kB 10.8 MB/s
     |████████████████████████████████| 56 kB 6.0 MB/s
     |████████████████████████████████| 17.3 MB 52 kB/s
     |████████████████████████████████| 10.5 MB 25.5 MB/s
     |████████████████████████████████| 107 kB 68.4 MB/s
     |████████████████████████████████| 13.1 MB 209 kB/s
     |████████████████████████████████| 4.4 MB 53.1 MB/s
     |████████████████████████████████| 226 kB 74.6 MB/s
     |████████████████████████████████| 6.8 MB 21.4 MB/s
     |████████████████████████████████| 110.5 MB 1.2 MB/s
     |████████████████████████████████| 411.0 MB 25 kB/s
     |████████████████████████████████| 503 kB 67.1 MB/s
     |████████████████████████████████| 103 kB 80.1 MB/s
     |████████████████████████████████| 45 kB 4.9 MB/s
     |████████████████████████████████| 328 kB 59.6 MB/s
     |████████████████████████████████| 9.5 MB 17.7 MB/s
     |████████████████████████████████| 63 kB 2.2 MB/s
     |████████████████████████████████| 3.8 MB 55.4 MB/s
  Building wheel for absl-py (setup.py) ... done
  Building wheel for gast (setup.py) ... done
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
xarray 0.18.2 requires numpy>=1.17, but you have numpy 1.16.2 which is incompatible.
tensorflow-probability 0.13.0 requires gast>=0.3.2, but you have gast 0.2.2 which is incompatible.
tensorflow-metadata 1.2.0 requires absl-py<0.13,>=0.9, but you have absl-py 0.7.1 which is incompatible.
spacy 2.2.4 requires tqdm<5.0.0,>=4.38.0, but you have tqdm 4.28.1 which is incompatible.
pyerfa 2.0.0 requires numpy>=1.17, but you have numpy 1.16.2 which is incompatible.
pyarrow 3.0.0 requires numpy>=1.16.6, but you have numpy 1.16.2 which is incompatible.
panel 0.12.1 requires tqdm>=4.48.0, but you have tqdm 4.28.1 which is incompatible.
kapre 0.3.5 requires numpy>=1.18.5, but you have numpy 1.16.2 which is incompatible.
kapre 0.3.5 requires tensorflow>=2.0.0, but you have tensorflow 1.15.4 which is incompatible.
jaxlib 0.1.70+cuda110 requires numpy>=1.18, but you have numpy 1.16.2 which is incompatible.
jax 0.2.19 requires numpy>=1.18, but you have numpy 1.16.2 which is incompatible.
google-colab 1.0.0 requires astor~=0.8.1, but you have astor 0.7.1 which is incompatible.
google-colab 1.0.0 requires six~=1.15.0, but you have six 1.12.0 which is incompatible.
google-api-python-client 1.12.8 requires six<2dev,>=1.13.0, but you have six 1.12.0 which is incompatible.
google-api-core 1.26.3 requires six>=1.13.0, but you have six 1.12.0 which is incompatible.
fbprophet 0.7.1 requires tqdm>=4.36.1, but you have tqdm 4.28.1 which is incompatible.
datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.
cupy-cuda101 9.1.0 requires numpy>=1.17, but you have numpy 1.16.2 which is incompatible.
astropy 4.3.1 requires numpy>=1.17, but you have numpy 1.16.2 which is incompatible.
albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible.
/content/proteinfer
!cd proteinfer
!ls/bin/bash: line 0: cd: proteinfer: No such file or directory
baseline_utils.py	  install_models.py	   protein_model_test.py
baseline_utils_test.py	  LICENSE		   README.md
colab_evaluation.py	  misc			   requirements.txt
colab_evaluation_test.py  parenthood_bin.py	   testdata
colabs			  parenthood_lib.py	   test_util.py
CONTRIBUTING.md		  parenthood_lib_test.py   test_util_test.py
evaluation.py		  protein_dataset.py	   train.py
evaluation_test.py	  protein_dataset_test.py  train_test.py
hparams_sets.py		  proteinfer.py		   utils.py
inference.py		  proteinfer_test.py	   utils_test.py
inference_test.py	  protein_model.py
import inference
import pandas as pd
import evaluation
import parenthood_lib

import re_____no_output_____

# Parse Price's ground truth labels_____no_output_____

!wget https://storage.googleapis.com/brain-genomics-public/research/proteins/proteinfer/colab_support/price_enzymes.tsv--2021-09-13 02:01:43--  https://storage.googleapis.com/brain-genomics-public/research/proteins/proteinfer/colab_support/price_enzymes.tsv
Resolving storage.googleapis.com (storage.googleapis.com)... 142.250.141.128, 142.251.2.128, 2607:f8b0:4023:c0d::80, ...
Connecting to storage.googleapis.com (storage.googleapis.com)|142.250.141.128|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 400317 (391K) [application/octet-stream]
Saving to: ‘price_enzymes.tsv’

price_enzymes.tsv     0%[                    ]       0  --.-KB/s
price_enzymes.tsv   100%[===================>] 390.93K  --.-KB/s    in 0.004s

2021-09-13 02:01:43 (93.7 MB/s) - ‘price_enzymes.tsv’ saved [400317/400317]

with open('price_enzymes.tsv') as f:
  price_enzymes = pd.read_csv(f, sep='\t')_____no_output_____price_enzymes_w_annotation = price_enzymes[price_enzymes.new_annotation.apply(lambda x: x.count('EC ') == 1)].copy()_____no_output_____

# Predict for these sequences_____no_output_____

!wget https://storage.googleapis.com/brain-genomics-public/research/proteins/proteinfer/models/zipped_models/noxpnd_cnn_swissprot_ec_random_swiss-cnn_for_swissprot_ec_random-13685140.tar.gz--2021-09-13 02:01:43--  https://storage.googleapis.com/brain-genomics-public/research/proteins/proteinfer/models/zipped_models/noxpnd_cnn_swissprot_ec_random_swiss-cnn_for_swissprot_ec_random-13685140.tar.gz
Resolving storage.googleapis.com (storage.googleapis.com)... 142.250.141.128, 142.251.2.128, 2607:f8b0:4023:c0b::80, ...
Connecting to storage.googleapis.com (storage.googleapis.com)|142.250.141.128|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 134391480 (128M) [application/octet-stream]
Saving to: ‘noxpnd_cnn_swissprot_ec_random_swiss-cnn_for_swissprot_ec_random-13685140.tar.gz’

noxpnd_cnn_swisspro 100%[===================>] 128.17M   111MB/s    in 1.2s

2021-09-13 02:01:45 (111 MB/s) - ‘noxpnd_cnn_swissprot_ec_random_swiss-cnn_for_swissprot_ec_random-13685140.tar.gz’ saved [134391480/134391480]

!tar xzf noxpnd_cnn_swissprot_ec_random_swiss-cnn_for_swissprot_ec_random-13685140.tar.gz_____no_output_____inferrer = inference.Inferrer('noxpnd_cnn_swissprot_ec_random_swiss-cnn_for_swissprot_ec_random-13685140', use_tqdm=True)WARNING: Logging before flag parsing goes to stderr.
W0913 02:01:47.257570 140413780555648 deprecation.py:323] From /usr/local/lib/python3.7/dist-packages/tensorflow_core/python/ops/ragged/ragged_tensor.py:1586: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.where in 2.0, which has the same broadcast rule as np.where
predictions = inferrer.get_activations(list(price_enzymes_w_annotation.protein_sequence))Annotating batches of sequences: 100%|██████████| 3/3 [01:08<00:00, 26.95s/it]
# Turn prediction array into actual labels by using the vocab.
cnn_label_vocab = inferrer.get_variable('label_vocab:0').astype(str)
price_enzymes_w_annotation['predicted_label'] = evaluation.filter_predictions_to_above_threshold(predictions, 0.148222, cnn_label_vocab) # Threshold is that which maximizes f1 score._____no_output_____!wget https://storage.googleapis.com/brain-genomics-public/research/proteins/proteinfer/colab_support/parenthood.json.gz--2021-09-13 02:02:56--  https://storage.googleapis.com/brain-genomics-public/research/proteins/proteinfer/colab_support/parenthood.json.gz
Resolving storage.googleapis.com (storage.googleapis.com)... 142.250.141.128, 142.251.2.128, 74.125.137.128, ...
Connecting to storage.googleapis.com (storage.googleapis.com)|142.250.141.128|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 992539 (969K) [application/octet-stream]
Saving to: ‘parenthood.json.gz’

parenthood.json.gz  100%[===================>] 969.28K  --.-KB/s    in 0.007s

2021-09-13 02:02:57 (138 MB/s) - ‘parenthood.json.gz’ saved [992539/992539]

# Find the most specific label(s) given to each protein.
label_normalizer = parenthood_lib.get_applicable_label_dict('parenthood.json.gz')
price_enzymes_w_annotation = parenthood_lib.filter_labels_to_most_specific(price_enzymes_w_annotation, label_normalizer)_____no_output_____to_examine = price_enzymes_w_annotation.copy()

# Drop labels where we have > 1 prediction for ease of comparison.
to_examine = to_examine[(to_examine.predicted_label.apply(len) == 1)].copy()

# Format predicted and true label identically.
to_examine['true_label'] = to_examine.new_annotation.apply(lambda x: re.findall(r'\(EC (.{1,3}\..{1,3}\..{1,3}\..{1,3})\)', x)[0])
to_examine['predicted_label'] = to_examine.predicted_label.apply(lambda x: list(x)[0].replace('EC:', ''))_____no_output_____

# Compute agreement and plot_____no_output_____

performance_by_level = pd.DataFrame([
    evaluation.ec_agreement_for_level(to_examine, 1),
    evaluation.ec_agreement_for_level(to_examine, 2),
    evaluation.ec_agreement_for_level(to_examine, 3),
    evaluation.ec_agreement_for_level(to_examine, 4)
], columns=['Correct', 'Incorrect', 'Not called'], index=range(1, 5))[['Correct','Not called',"Incorrect"]]_____no_output_____print(performance_by_level)   Correct  Not called  Incorrect
1      103           0         12
2       91          14         10
3       75          28         12
4       13          79         18
array = performance_by_level.to_numpy()

array/array.sum(axis=1,keepdims=True)_____no_output_____performance_by_level['position']=performance_by_level.index
melted=pd.melt(performance_by_level,'position')_____no_output_____from plotnine import *
import matplotlib as mpl
mpl.rcParams['figure.dpi']= 800_____no_output_____
melted['variable'] = melted['variable'].astype('category').cat.reorder_categories(["Incorrect","Not called","Correct"], ordered=True )

(ggplot(melted,aes(x="position",fill="variable",y="value"))+
geom_col(color="black", position="fill")+
scale_fill_manual(["red","lightgray","green"])+ theme_classic()
)
_____no_output__________no_output_____


# Notebook from ricardo-hdz/trading_ml
Path: early_rl.ipynb

# Early Reinforcement Learning

With the advances of modern computing power, the study of Reinforcement Learning is having a heyday. Machines are now able to learn complex tasks once thought to be solely in the domain of humans, from conrolling the [heating and cooling in massive data centers](https://www.technologyreview.com/s/611902/google-just-gave-control-over-data-center-cooling-to-an-ai/) to beating [grandmasters at Starcraft](https://storage.googleapis.com/deepmind-media/research/alphastar/AlphaStar_unformatted.pdf). As magnificent as it may seem today, it had humble roots many decades ago. Seeing how far it's come, it's a wonder to see how far it will go!

Let's take a step back in time to see how these early algorithms developed. Many of these algorithms make sense given the context of when they were created. Challenge yourself and see if you can come up with the same strategies given the right problem. Ok! Time to cozy up for a story.

This is the hero of our story, the gumdrop emoji. It was enjoying a cool winter day building a snowman when suddely, it slipped and fell on a frozen lake of death.

The lake can be thought of as a 4 x 4 grid where the gumdrop can move left (0), down (1), right (2) and up (3). Unfortunately, this frozen lake of death has holes of death where if the gumdrop enters that square, it will fall in and meet an untimely demise. To make matters worse, the lake is surrounded by icy boulders that if the gumdrop attempts to climb, will have it slip back into its original position. Thankfully, at the bottom right of the lake is a safe ramp that leads to a nice warm cup of hot cocoa.

## Set Up

We can try and save the gumdrop ourselves! This is a common game people begin their Reinforcement Learning journey with, and is included in the OpenAI's python package [Gym](https://gym.openai.com/) and is aptly named [FrozenLake-v0](https://gym.openai.com/envs/FrozenLake-v0/) ([code](https://github.com/openai/gym/blob/master/gym/envs/toy_text/frozen_lake.py)). No time to waste, let's get the environment up and running. Run the below to install the needed libraries if they are not installed already._____no_output_____

!python3 -m pip freeze | grep gym || sudo python3 -m pip install gym==0.12.5Collecting gym==0.12.5
  Downloading gym-0.12.5.tar.gz (1.5 MB)
     |████████████████████████████████| 1.5 MB 4.6 MB/s eta 0:00:01
Requirement already satisfied: scipy in /usr/local/lib/python3.5/dist-packages (from gym==0.12.5) (1.4.1)
Requirement already satisfied: numpy>=1.10.4 in /usr/local/lib/python3.5/dist-packages (from gym==0.12.5) (1.18.1)
Requirement already satisfied: six in /usr/local/lib/python3.5/dist-packages (from gym==0.12.5) (1.14.0)
Collecting pyglet>=1.2.0
  Downloading pyglet-1.5.0-py2.py3-none-any.whl (1.0 MB)
     |████████████████████████████████| 1.0 MB 19.3 MB/s eta 0:00:01
Requirement already satisfied: future in /usr/local/lib/python3.5/dist-packages (from pyglet>=1.2.0->gym==0.12.5) (0.18.2)
Building wheels for collected packages: gym
  Building wheel for gym (setup.py) ... done
  Created wheel for gym: filename=gym-0.12.5-py3-none-any.whl size=1614824 sha256=1f68f160e37e965330f2c0d1957d13e2123cbdf2fe2e5b666621cfc9c6401615
  Stored in directory: /root/.cache/pip/wheels/97/b1/1c/da7e0452ed21171c1f59fac9eec5d2d1723c671c69b4439ab9
Successfully built gym
Installing collected packages: pyglet, gym
Successfully installed gym-0.12.5 pyglet-1.5.0

###### Note: Restart the kernel if the above libraries needed to be installed

There are [four methods from Gym](http://gym.openai.com/docs/) that are going to be useful to us in order to save the gumdrop.
* `make` allows us to build the environment or game that we can pass actions to
* `reset` will reset an environment to it's starting configuration and return the state of the player
* `render` displays the environment for human eyes
* `step` takes an action and returns the player's next state.

Let's make, reset, and render the game. The output is an ANSI string with the following characters:
* `S` for starting point
* `F` for frozen
* `H` for hole
* `G` for goal
* A red square indicates the current position_____no_output_____

import gym
import numpy as np
import random
env = gym.make('FrozenLake-v0', is_slippery=False)
state = env.reset()
env.render()
SFFF
FHFH
FFFH
HFFG

If we print the state we'll get `0`. This is telling us which square we're in. Each square is labeled from `0` to `15` from left to right, top to bottom, like this:

| | | | |
|-|-|-|-|
|0|1|2|3|
|4|5|6|7|
|8|9|10|11|
|12|13|14|15|
_____no_output_____

print(state)0

We can make a simple print function to let us know whether it's game won, game over, or game on._____no_output_____

def print_state(state, done):
    statement = "Still Alive!"
    if done:
        statement = "Cocoa Time!" if state == 15 else "Game Over!"
    print(state, "-", statement)_____no_output_____

We can control the gumdrop ourselves with the `step` method. Run the below cell over and over again trying to move from the starting position to the goal. Good luck!_____no_output_____

#0 left
#1 down
#2 right
#3 up

# Uncomment to reset the game
#env.reset()
action = 2  # Change me, please!
state, _, done, _ = env.step(action)
env.render()
print_state(state, done)  (Right)
SFFF
FHFH
FFFH
HFFG
15 - Cocoa Time!

Were you able to reach the hot chocolate? If so, great job! There are multiple paths through the maze. One solution is `[1, 1, 2, 2, 1, 2]`. Let's loop through our actions in order to get used to interacting with the environment programmatically._____no_output_____

def play_game(actions):
    state = env.reset()
    step = 0
    done = False

    while not done and step < len(actions):
        action = actions[step]
        state, _, done, _ = env.step(action)
        env.render()
        step += 1
        print_state(state, done)

actions = [1, 1, 2, 2, 1, 2]  # Replace with your favorite path.
play_game(actions)  (Down)
SFFF
FHFH
FFFH
HFFG
4 - Still Alive!
  (Down)
SFFF
FHFH
FFFH
HFFG
8 - Still Alive!
  (Right)
SFFF
FHFH
FFFH
HFFG
9 - Still Alive!
  (Right)
SFFF
FHFH
FFFH
HFFG
10 - Still Alive!
  (Down)
SFFF
FHFH
FFFH
HFFG
14 - Still Alive!
  (Right)
SFFF
FHFH
FFFH
HFFG
15 - Cocoa Time!

Nice, so we know how to get through the maze, but how do we teach that to the gumdrop? It's just some bytes in an android phone. It doesn't have our human insight.

We could give it our list of actions directly, but then it would be copying us and not really learning. This was a tricky one to the mathematicians and computer scientists origionally trying to solve this problem. How do we teach a machine to do this without human insight?

## Value Iteration

Let's turn the clock back on our time machines to 1957 to meet Mr. [Richard Bellman](https://en.wikipedia.org/wiki/Richard_E._Bellman). Bellman started his academic career in mathematics, but due to World War II, left his postgraduate studies at John Hopkins to teach electronics as part of the war effort (as chronicled by J. J. O'Connor and E. F. Robertson [here](https://www-history.mcs.st-andrews.ac.uk/Biographies/Bellman.html)). When the war was over, and it came time for him to focus on his next area of research, he became fascinated with [Dynamic Programming](https://en.wikipedia.org/wiki/Dynamic_programming): the idea of breaking a problem down into sub-problems and using recursion to solve the larger problem.

Eventually, his research landed him on [Markov Decision Processes](https://en.wikipedia.org/wiki/Markov_decision_process). These processes are a graphical way of representing how to make a decision based on a current state. States are connected to other states with positive and negative rewards that can be picked up along the way.

Sound familiar at all? Perhaps our Frozen Lake?

In the lake case, each cell is a state. The `H`s and the `G` are a special type of state called a "Terminal State", meaning they can be entered, but they have no leaving connections. What of rewards? Let's say the value of losing our life is the negative opposite of getting to the goal and staying alive. Thus, we can assign the reward of entering a death hole as -1, and the reward of escaping as +1.

Bellman's first breakthrough with this type of problem is now known as Value Iteration ([his original paper](http://www.iumj.indiana.edu/IUMJ/FULLTEXT/1957/6/56038)). He introduced a variable, gamma (γ), to represent discounted future rewards. He also introduced a function of policy (π) that takes a state (s), and outputs corresponding suggested action (a). The goal is to find the value of a state (V), given the rewards that occur when following an action in a particular state (R).

Gamma, the discount, is the key ingredient here. If my time steps were in days, and my gamma was .9, `$100` would be worth `$100` to me today, `$90` tomorrow, `$81` the day after, and so on. Putting this all together, we get the Bellman Equation

source: [Wikipedia](https://en.wikipedia.org/wiki/Bellman_equation)

In other words, the value of our current state, `current_values`, is equal to the discount times the value of the next state, `next_values`, given the policy the agent will follow. For now, we'll have our agent assume a greedy policy: it will move towards the state with the highest calculated value. If you're wondering what P is, don't worry, we'll get to that later.

Let's program it out and see it in action! We'll set up an array representing the lake with -1 as the holes, and 1 as the goal. Then, we'll set up an array of zeros to start our iteration._____no_output_____

LAKE = np.array([[0,  0,  0,  0],
                 [0, -1,  0, -1],
                 [0,  0,  0, -1],
                 [-1, 0,  0,  1]])
LAKE_WIDTH = len(LAKE[0])
LAKE_HEIGHT = len(LAKE)

DISCOUNT = .9  # Change me to be a value between 0 and 1.
current_values = np.zeros_like(LAKE)_____no_output_____

The Gym environment class has a handy property for finding the number of states in an environment called `observation_space`. In our case, there a 16 integer states, so it will label it as "Discrete". Similarly, `action_space` will tell us how many actions are abailable to the agent.

Let's take advantage of these to make our code portable between different lakes sizes._____no_output_____

print("env.observation_space -", env.observation_space)
print("env.observation_space.n -", env.observation_space.n)
print("env.action_space -", env.action_space)
print("env.action_space.n -", env.action_space.n)

STATE_SPACE = env.observation_space.n
ACTION_SPACE = env.action_space.n
STATE_RANGE = range(STATE_SPACE)
ACTION_RANGE = range(ACTION_SPACE)env.observation_space - Discrete(16)
env.observation_space.n - 16
env.action_space - Discrete(4)
env.action_space.n - 4

We'll need some sort of function to figure out what the best neighboring cell is. The below function take's a cell of the lake, and looks at the current value mapping (to be called with `current_values`, and see's what the value of the adjacent state is corresponding to the given `action`._____no_output_____

def get_neighbor_value(state_x, state_y, values, action):
    """Returns the value of a state's neighbor.

    Args:
        state_x (int): The state's horizontal position, 0 is the lake's left.
        state_y (int): The state's vertical position, 0 is the lake's top.
        values (float array): The current iteration's state values.
        policy (int): Which action to check the value for.

    Returns:
        The corresponding action's value.
    """
    left = [state_y, state_x-1]
    down = [state_y+1, state_x]
    right = [state_y, state_x+1]
    up = [state_y-1, state_x]
    actions = [left, down, right, up]

    direction = actions[action]
    check_x = direction[1]
    check_y = direction[0]

    is_boulder = check_y < 0 or check_y >= LAKE_HEIGHT \
        or check_x < 0 or check_x >= LAKE_WIDTH

    value = values[state_y, state_x]
    if not is_boulder:
        value = values[check_y, check_x]

    return value_____no_output_____

But this doesn't find the best action, and the gumdrop is going to need that if it wants to greedily get off the lake. The `get_max_neighbor` function we've defined below takes a number corresponding to a cell as `state_number` and the same value mapping as `get_neighbor_value`._____no_output_____

def get_state_coordinates(state_number):
    state_x = state_number % LAKE_WIDTH
    state_y = state_number // LAKE_HEIGHT
    return state_x, state_y

def get_max_neighbor(state_number, values):
    """Finds the maximum valued neighbor for a given state.

    Args:
        state_number (int): the state to find the max neighbor for
        state_values (float array): the respective value of each state for
            each cell of the lake.

    Returns:
        max_value (float): the value of the maximum neighbor.
        policy (int): the action to take to move towards the maximum neighbor.
    """
    state_x, state_y = get_state_coordinates(state_number)

    # No policy or best value yet
    best_policy = -1
    max_value = -np.inf

    # If the cell has something other than 0, it's a terminal state.
    if LAKE[state_y, state_x]:
        return LAKE[state_y, state_x], best_policy

    for action in ACTION_RANGE:
        neighbor_value = get_neighbor_value(state_x, state_y, values, action)
        if neighbor_value > max_value:
            max_value = neighbor_value
            best_policy = action

    return max_value, best_policy_____no_output_____

Now, let's write our value iteration code. We'll write a function that codes out one step of the iteration by cheacking each state and finding its maximum neighbor. The values will be reshaped so that it's in the form of the lake, but the policy will stay as a list of ints. This way, when Gym returns a state, all we need to do is look at the corresponding index in the policy list to tell our agent where to go._____no_output_____

def iterate_value(current_values):
    """Finds the future state values for an array of current states.

    Args:
        current_values (int array): the value of current states.

    Returns:
        next_values (int array): The value of states based on future states.
        next_policies (int array): The recommended action to take in a state.
    """
    next_values = []
    next_policies = []

    for state in STATE_RANGE:
        value, policy = get_max_neighbor(state, current_values)
        next_values.append(value)
        next_policies.append(policy)

    next_values = np.array(next_values).reshape((LAKE_HEIGHT, LAKE_WIDTH))
    return next_values, next_policies

next_values, next_policies = iterate_value(current_values)_____no_output_____

This is what our values look like after one step. Right now, it just looks like the lake. That's because we started with an array of zeros for `current_values`, and the terminal states of the lake were loaded in._____no_output_____

next_values_____no_output_____

And this is what our policy looks like reshaped into the form of the lake. The `-1`'s are terminal states. Right now, the agent will move left in any non-terminal state, because it sees all of those states as equal. Remember, if the gumdrop is along the leftmost side of the lake, and tries to move left, it will slip on a boulder and return to the same position._____no_output_____

np.array(next_policies).reshape((LAKE_HEIGHT ,LAKE_WIDTH))_____no_output_____

There's one last step to apply the Bellman Equation, the `discount`! We'll multiply our next states by the `discount` and set that to our `current_values`. One loop done!_____no_output_____

current_values = DISCOUNT * next_values
current_values_____no_output_____

Run the below cell over and over again to see how our values change with each iteration. It should be complete after six iterations when the values no longer change. The policy will also change as the values are updated._____no_output_____

next_values, next_policies = iterate_value(current_values)
print("Value")
print(next_values)
print("Policy")
print(np.array(next_policies).reshape((4,4)))
current_values = DISCOUNT * next_values_____no_output_____

Have a completed policy? Let's see it in action! We'll update our `play_game` function to instead take our list of policies. That way, we can start in a random position and still get to the end._____no_output_____

def play_game(policy):
    state = env.reset()
    step = 0
    done = False

    while not done:
        action = policy[state]  # This line is new.
        state, _, done, _ = env.step(action)
        env.render()
        step += 1
        print_state(state, done)

play_game(next_policies)_____no_output_____

Phew! Good job, team! The gumdrop made it out alive. So what became of our gumdrop hero? Well, the nect day, it was making another snowman and fell onto an even more slippery and deadly lake. Doh! Turns out this story is part of a trilogy. Feel free to move onto the next section after your own sip of cocoa, coffee, tea, or poison of choice.

## Policy Iteration

You may have noticed that the first lake was built with the parameter `is_slippery=False`. This time, we're going to switch it to `True`._____no_output_____

env = gym.make('FrozenLake-v0', is_slippery=True)
state = env.reset()
env.render()
SFFF
FHFH
FFFH
HFFG

Hmm, looks the same as before. Let's try applying our old policy and see what happens._____no_output_____

play_game(next_policies)_____no_output_____

Was there a game over? There's a small chance that the gumdrop made it to the end, but it's much more likely that it accidentally slipped and fell into a hole. Oh no! We can try repeatedly testing the above code cell over and over again, but it might take a while. In fact, this is a similar rodeblock Bellman and his colleagues faced.

How efficient is Value Iteration? On our modern machines, this algorithm ran fairly quickly, but back in 1960, that wasn't the case. Let's say our lake is a long straight line like this:

| | | | | | | |
|-|-|-|-|-|-|-|
|S|F|F|F|F|F|H|

This is the worse case scenario for value iteration. In each iteration, we look at every state (s) and each action per state (a), so one step of value iteration is O(s*a). In the case of our lake line, each iteration only updates one cell. In other words, the value iteration step needs to be run `s` times. In total, that's O(s2a).

Back in 1960, that was computationally heavy, and so [Ronald Howard](https://en.wikipedia.org/wiki/Ronald_A._Howard) developed an alteration of Value Iteration that mildly sacrificed mathematical accuaracy for speed.

Here's the strategy: it was observed that the optimal policy often converged before value iteration was complete. To take advantage of this, we'll start with random policy. When we iterate over our values, we'll use this policy instead of trying to find the maximum neighbor. This has been coded out in `find_future_values` below._____no_output_____

def find_future_values(current_values, current_policies):
    """Finds the next set of future values based on the current policy."""
    next_values = []

    for state in STATE_RANGE:
        current_policy = current_policies[state]
        state_x, state_y = get_state_coordinates(state)

        # If the cell has something other than 0, it's a terminal state.
        value = LAKE[state_y, state_x]
        if not value:
            value = get_neighbor_value(
                state_x, state_y, current_values, current_policy)
        next_values.append(value)

    return np.array(next_values).reshape((LAKE_HEIGHT, LAKE_WIDTH))_____no_output_____

After we've calculated our new values, then we'll updated the policy (and not the values) based on the maximum neighbor. If there's no change in the policy, then we're done. The below is very similar to our `get_max_neighbor` function. Can you see the differences?_____no_output_____

def find_best_policy(next_values):
    """Finds the best policy given a value mapping."""
    next_policies = []
    for state in STATE_RANGE:
        state_x, state_y = get_state_coordinates(state)

        # No policy or best value yet
        max_value = -np.inf
        best_policy = -1

        if not LAKE[state_y, state_x]:
            for policy in ACTION_RANGE:
                neighbor_value = get_neighbor_value(
                    state_x, state_y, next_values, policy)
                if neighbor_value > max_value:
                    max_value = neighbor_value
                    best_policy = policy

        next_policies.append(best_policy)
    return next_policies_____no_output_____

To complete the Policy Iteration algorithm, we'll combine the two functions above. Conceptually, we'll be alternating between updating our value function and updating our policy function._____no_output_____

def iterate_policy(current_values, current_policies):
    """Finds the future state values for an array of current states.

    Args:
        current_values (int array): the value of current states.
        current_policies (int array): a list where each cell is the recommended
            action for the state matching its index.

    Returns:
        next_values (int array): The value of states based on future states.
        next_policies (int array): The recommended action to take in a state.
    """
    next_values = find_future_values(current_values, current_policies)
    next_policies = find_best_policy(next_values)
    return next_values, next_policies_____no_output_____

Next, let's modify the `get_neighbor_value` function to now include the slippery ice. Remember the `P` in the Bellman Equation above? It stands for the probabilty of ending up in a new state given the current state and action taken. That is, we'll take a weighted sum of the values of all possible states based on our chances to be in those states.

How does the physics of the slippery ice work? For this lake, whenever the gumdrop tries to move in a particular direction, there are three possible positions that it could end up with. It could move where it was intending to go, but it could also end up to the left or right of the direction it was facing. For instance, if it wanted to move right, it could end up on the square above or below it! This is depicted below, with the yellow squares being potential positions after attempting to move right.

Each of these has an equal probability chance of happening. So since there are three outcomes, they each have about a 33% chance to happen. What happens if we slip in the direction of a boulder? No problem, we'll just end up not moving anywhere. Let's make a function to find what our possible locations could be given a policy and state coordinates._____no_output_____

def get_locations(state_x, state_y, policy):
    left = [state_y, state_x-1]
    down = [state_y+1, state_x]
    right = [state_y, state_x+1]
    up = [state_y-1, state_x]
    directions = [left, down, right, up]
    num_actions = len(directions)

    gumdrop_right = (policy - 1) % num_actions
    gumdrop_left = (policy + 1) % num_actions
    locations = [gumdrop_left, policy, gumdrop_right]
    return [directions[location] for location in locations]_____no_output_____

Then, we can add it to `get_neighbor_value` to find the weighted value of all the possible states the gumdrop can end up in._____no_output_____

def get_neighbor_value(state_x, state_y, values, policy):
    """Returns the value of a state's neighbor.

    Args:
        state_x (int): The state's horizontal position, 0 is the lake's left.
        state_y (int): The state's vertical position, 0 is the lake's top.
        values (float array): The current iteration's state values.
        policy (int): Which action to check the value for.

    Returns:
        The corresponding action's value.
    """
    locations = get_locations(state_x, state_y, policy)
    location_chance = 1.0 / len(locations)
    total_value = 0

    for location in locations:
        check_x = location[1]
        check_y = location[0]

        is_boulder = check_y < 0 or check_y >= LAKE_HEIGHT \
            or check_x < 0 or check_x >= LAKE_WIDTH

        value = values[state_y, state_x]
        if not is_boulder:
            value = values[check_y, check_x]
        total_value += location_chance * value

    return total_value_____no_output_____

For Policy Iteration, we'll start off with a random policy if only because the Gumdrop doesn't know any better yet. We'll reset our current values while we're at it._____no_output_____

current_values = np.zeros_like(LAKE)
policies = np.random.choice(ACTION_RANGE, size=STATE_SPACE)
np.array(policies).reshape((4,4))_____no_output_____

As before with Value Iteration, run the cell below multiple until the policy no longer changes. It should only take 2-3 clicks compared to Value Iteration's 6._____no_output_____

next_values, policies = iterate_policy(current_values, policies)
print("Value")
print(next_values)
print("Policy")
print(np.array(policies).reshape((4,4)))
current_values = DISCOUNT * next_values_____no_output_____

Hmm, does this work? Let's see! Run the cell below to watch the gumdrop slip its way to victory._____no_output_____

play_game(policies)_____no_output_____

So what was the learned strategy here? The gumdrop learned to hug the left wall of boulders until it was down far enough to make a break for the exit. Instead of heading directly for it though, it took advantage of actions that did not have a hole of death in them. Patience is a virtue!

We promised this story was a trilogy, and yes, the next day, the gumdrop fell upon a frozen lake yet again.

## Q Learning
Value Iteration and Policy Iteration are great techniques, but what if we don't know  how big the lake is? With real world problems, not knowing how many potential states are can be a definite possibility.

Enter [Chris Watkins](http://www.cs.rhul.ac.uk/~chrisw/). Inspired by how animals learn with delayed rewards, he came up with the idea of [Q Learning](http://www.cs.rhul.ac.uk/~chrisw/new_thesis.pdf) as an evolution of [Richard Sutton's](https://en.wikipedia.org/wiki/Richard_S._Sutton) [Temporal Difference Learning](https://en.wikipedia.org/wiki/Temporal_difference_learning). Watkins noticed that animals learn from positive and negative rewards, and that they often make mistakes in order to optimize a skill.

From this emerged the idea of a Q table. In the lake case, it would look something like this.

| |Left|Down|Right|Up|
|-|-|-|-|-|
|0| | | | |
|1| | | | |
|...| | | | |

Here's the strategy: our agent will explore the environment. As the agent observes new states, we'll add more rows to our table. Whenever it moves from one state to the next, we'll update the cell corresponding to the old state based on the Bellman Equation. The agent doesn't need to know what the probabilities are between transitions. It'll learn the value of these as it experiments.

For Q learning, this works by looking at the row that corresponds the agent's current state. Then, we'll select the action with the highest value. There are multiple ways to initialize the Q-table, but for us, we'll start with all zeros. In that case, when selecting the best action, we'll randomly select between tied max values. If we don't, the agent will favor certain actions which will limit its exploration.

To be able to hande an unknown number of states, we'll initialize our q_table as one row to represent our initial state. Then, we'll make a dictionary to map new states to rows in the table._____no_output_____

new_row = np.zeros((1, env.action_space.n))
q_table = np.copy(new_row)
q_map = {0: 0}

def print_q(q_table, q_map):
    print("mapping")
    print(q_map)
    print("q_table")
    print(q_table)

print_q(q_table, q_map)_____no_output_____

Our new `get_action` function will help us read the `q_table` and find the best action.

First, we'll give the agent the ability to act randomly as opposed to choosing the best known action. This gives it the ability to explore and find new situations. This is done with a random chance to act randomly. So random!

When the Gumdrop chooses not to act randomly, it will instead act based on the best action recorded in the `q_table`. Numpy's [argwhere](https://docs.scipy.org/doc/numpy/reference/generated/numpy.argwhere.html) is used to find the indexes with the maximum value in the q-table row corresponding to our current state. Since numpy is often used with higher dimensional data, each index is returned as a list of ints. Our indexes are really one dimensional since we're just looking within a single row, so we'll use [np.squeeze](https://docs.scipy.org/doc/numpy/reference/generated/numpy.squeeze.html) to remove the extra brackets. To randomly select from the indexes, we'll use [np.random.choice](https://docs.scipy.org/doc/numpy-1.14.1/reference/generated/numpy.random.choice.html)._____no_output_____

def get_action(q_map, q_table, state_row, random_rate):
    """Find max-valued actions and randomly select from them."""
    if random.random() < random_rate:
        return random.randint(0, ACTION_SPACE-1)

    action_values = q_table[state_row]
    max_indexes = np.argwhere(action_values == action_values.max())
    max_indexes = np.squeeze(max_indexes, axis=-1)
    action = np.random.choice(max_indexes)
    return action_____no_output_____

Here, we'll define how the `q_table` gets updated. We'll apply the Bellman Equation as before, but since there is so much luck involved between slipping and random actions, we'll update our `q_table` as a weighted average between the `old_value` we're updating and the `future_value` based on the best action in the next state. That way, there's a little bit of memory between old and new experiences._____no_output_____

def update_q(q_table, new_state_row, reward, old_value):
    """Returns an updated Q-value based on the Bellman Equation."""
    learning_rate = .1  # Change to be between 0 and 1.
    future_value = reward + DISCOUNT * np.max(q_table[new_state_row])
    return old_value + learning_rate * (future_value - old_value)_____no_output_____

We'll update our `play_game` function to take our table and mapping, and at the end, we'll return any updates to them. Once we observe new states, we'll check our mapping and add then to the table if space isn't allocated for them already.

Finally, for every `state` - `action` - `new-state` transition, we'll update the cell in `q_table` that corresponds to the `state` and `action` with the Bellman Equation.

There's a little secret to solving this lake problem, and that's to have a small negative reward when moving between states. Otherwise, the grumdrop will become too afraid of slipping in a death hole to explore out of what is thought to be safe positions._____no_output_____

def play_game(q_table, q_map, random_rate, render=False):
    state = env.reset()
    step = 0
    done = False

    while not done:
        state_row = q_map[state]
        action = get_action(q_map, q_table, state_row, random_rate)
        new_state, _, done, _ = env.step(action)

        #Add new state to table and mapping if it isn't there already.
        if new_state not in q_map:
            q_map[new_state] = len(q_table)
            q_table = np.append(q_table, new_row, axis=0)
        new_state_row = q_map[new_state]

        reward = -.01  #Encourage exploration.
        if done:
            reward = 1 if new_state == 15 else -1
        current_q = q_table[state_row, action]
        q_table[state_row, action] = update_q(
            q_table, new_state_row, reward, current_q)

        step += 1
        if render:
            env.render()
            print_state(new_state, done)
        state = new_state

    return q_table, q_map_____no_output_____

Ok, time to shine, gumdrop emoji! Let's do one simulation and see what happens._____no_output_____

# Run to refresh the q_table.
random_rate = 1
q_table = np.copy(new_row)
q_map = {0: 0}_____no_output_____q_table, q_map = play_game(q_table, q_map, random_rate, render=True)
print_q(q_table, q_map)_____no_output_____

Unless the gumdrop was incredibly lucky, chances were, it fell in some death water. Q-learning is markedly different from Value Iteration or Policy Iteration in that it attempts to simulate how an animal learns in unknown situations. Since the layout of the lake is unknown to the Gumdrop, it doesn't know which states are death holes, and which ones are safe. Because of this, it's going to make many mistakes before it can start making successes.

Feel free to run the above cell multiple times to see how the gumdrop steps through trial and error. When you're ready, run the below cell to have the gumdrop play 1000 times._____no_output_____

for _ in range(1000):
    q_table, q_map = play_game(q_table, q_map, random_rate)
    random_rate = random_rate * .99
print_q(q_table, q_map)
random_rate_____no_output_____

Cats have nine lives, our Gumdrop lived a thousand! Moment of truth. Can it get out of the lake now that it matters?_____no_output_____

q_table, q_map = play_game(q_table, q_map, 0, render=True)_____no_output_____

Third time's the charm!

Each of these techniques has their pros and cons. For instance, while Value Iteration is the mathematically correct solution, it's not as time efficient at Policy Iteration or as flexible as Q-Learning.

| |Value Iteration|Policy Iteration|Q Tables|
|-|-|-|-|
|Avoids locally optimal routes|✓|x|x|
|On-policy (greedy)|✓|✓|x|
|Model Free|x|x|✓|
|Most time efficient|x|✓|x|

Congratulations on making it through to the end. Now if you ever fall on a Frozen Lake, you'll have may different ways to calculate your survival. The gumdrop thanks you!

_____no_output_____Copyright 2019 Google Inc.
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License._____no_output_____


# Notebook from DavidAssaraf106/2021-CS109B
Path: docs/lectures/lecture20/notebook/gru_scaffold-varhar.ipynb

## Title :
Vanishing Gradients

## Description :
The goal of this exercise is to understand the vanishing gradient problem in training RNNs and using various methods to improve training.

In order to do this exercise, we will use the IMDB movie review dataset to perform sentiment analysis.

Your final comparison for the trace plot may look something like this:

## Instructions:
- Read the IMDB dataset from the helper code given.
- Take a quick look at your training inputs and labels.
- Pad the values to a fix number `max_words` in-order to have sequences of the same size.
- First post pad the inputs with `padding='post'` i.e sequences smaller than `max_words` will be followed by zero padding.
- Build, compile and fit a Vanilla RNN and evaluate it on the test set.
- Now refit the model, but this time post pad the inputs with `padding='pre'`
- Again evaluate the model performance on the test set.
- Finally, rebuild a model with the Gated Recurrent Unit and fit and evaluate on the training and test set respectively.
- Compare the performance of all three models similar to the table above.

## Hints:

tensorflow.keras.preprocessing.sequence.pad_sequences() Pad the sequences to same length - pre/post._____no_output_____## Vanilla RNNs

We will use Vanilla Recurrent Neural Networks to perform sentiment analysis in `tensorflow.keras` using imdb movie reviews dataset.
The dataset is a subset consisting of 10,000 reviews in the training and test set each. [see here for more info](https://www.tensorflow.org/datasets/catalog/imdb_reviews)_____no_output_____

# Import required libraries
from tensorflow.keras.datasets import imdb
from tensorflow.keras.preprocessing import sequence
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Embedding, SimpleRNN, Dense,GRU
import pickle
import numpy as np
import matplotlib.pyplot as plt
from prettytable import PrettyTable
from pprint import pprint_____no_output_____# We fix a vocabulary size of 5000
# Use the code below to call a small subset of the imdb dataset
# We keep the vocabulary size fixed because it was used to curate the sub-dataset
vocabulary_size = 5000
with open('imdb_mini.pkl','rb') as f:
    X_train, y_train, X_test, y_test = pickle.load(f)
_____no_output_____

#### Inspect a sample review and its label_____no_output_____

# Run the code below to see the first tokenized review and the label
print('---review---')
print(X_train[0])
print('---label---')
print(y_train[0])---review---
[1, 48, 2498, 2, 16, 4, 4, 1554, 149, 14, 22, 95, 198, 51, 29, 62, 2039, 46, 11, 189, 10, 10, 146, 806, 1588, 21, 2, 195, 75, 69, 79, 3514, 4, 1122, 292, 2, 5, 150, 14, 803, 4, 4351, 57, 45, 24, 618, 6, 2, 15, 289, 7, 4, 2, 323, 2, 2, 5, 4, 85, 381, 160, 87, 698, 284, 4, 293, 1141, 2, 11, 148, 2, 9, 4, 2, 7, 4, 108, 36, 1173, 93, 8, 4171, 363, 36, 71, 2812, 631, 108, 19, 6, 955, 1382, 841, 15, 43, 566, 30, 2, 11, 4, 2997, 4, 681, 9, 1215, 5, 51, 128, 96, 8, 2330, 6, 1215, 22, 74, 8, 2, 12, 19, 6, 1034, 42, 60, 6, 755, 10, 10, 2, 2, 69, 6, 1211, 2790, 159, 12, 60, 569, 4, 1320, 2, 2, 9, 60, 53, 2, 15, 85, 301, 1215, 108, 12, 2, 512, 1089, 3167, 6, 1242, 156, 4, 228, 603, 270, 1328, 15, 2, 2786, 5, 32, 4, 537, 3392, 4, 22, 10, 10, 38, 133, 266, 6, 1034, 92, 3250, 57, 2, 71, 2, 11, 4, 231, 7, 14, 1034, 49, 678, 1409, 7, 4, 65, 887, 8, 30, 2, 18, 4, 682, 2997, 2, 2, 2, 21, 198, 43, 44, 4, 226, 863, 38, 75, 202, 4, 1002, 6, 2, 405, 2, 501, 601, 19, 1465, 228, 2183, 18, 4, 706, 2629, 95, 19, 57, 2, 1699, 2, 23, 4, 1111, 15, 2580, 4, 2, 46, 21, 11, 4, 147, 182, 14, 586, 593, 1788, 43, 92, 140, 1012, 202, 90, 6, 541, 4304, 18, 3666, 247, 74, 4, 2, 7, 4580, 5, 25, 28, 4, 1034, 2, 1523, 151, 218, 12, 10, 10, 45, 43, 15, 12, 16, 32, 2984, 23, 19, 6, 2, 4, 403, 2, 71, 331, 2, 220, 1671, 23, 50, 16, 57, 281, 7, 1830, 23, 4, 1111, 57, 2, 7, 513, 8, 1277, 129, 2, 43, 6, 171, 4165, 2, 44, 6, 2585, 5, 15, 16, 12, 2, 16, 43, 616, 34, 24, 743, 46, 101, 2, 33, 32, 5, 1935, 16, 3506, 8, 387, 41, 79, 245, 19, 12, 54, 29, 435, 83, 4, 73, 25, 43, 697, 29, 62, 79, 2897, 12, 4, 881, 16, 2, 32, 4, 96, 8, 4, 130, 5, 25, 43, 473, 12, 8, 2, 56, 5, 130, 4, 2, 16, 427, 642, 5, 161, 124, 54, 8, 570, 10, 10, 15, 277, 9, 242, 4, 118, 96, 8, 2, 4, 1474, 200, 4, 107, 31, 630, 11, 4, 91, 307, 2, 103, 4, 91, 2993, 251, 4, 85, 630, 19, 6, 1208, 365, 1260, 12, 32, 8, 4, 2, 552, 1174, 10, 10, 13, 447, 4, 204, 21, 435, 8, 4, 438, 19, 35, 911, 330, 5, 16, 2229, 8, 67, 4, 22, 13, 317, 2, 11, 4, 1857, 15, 14, 22, 80, 242, 130, 56, 4158, 6, 2, 1198, 64, 14, 58, 2, 1551, 1437]
---label---
0
# You can get the word2id mapping by
# using the imdb.get_word_index() function

word2id = imdb.get_word_index()
# We need to adjust the mapping by 3 because of tensorflow.keras preprocessing
# more here: https://stackoverflow.com/questions/42821330/restore-original-text-from-keras-s-imdb-dataset
word2id = {k:(v+3) for k,v in word2id.items()}
word2id["<PAD>"] = 0
word2id["<START>"] = 1
word2id["<UNK>"] = 2
word2id["<UNUSED>"] = 3

# Reversing the key,value pair will give the id2word
id2word = {i: word for word, i in word2id.items()}
pprint('---review with words---')
pprint(" ".join([id2word[i] for i in X_train[0]]))
pprint('---label---')
pprint(y_train[0]);'---review with words---'
("<START> if edward <UNK> was the the flicks watching this film then that's "
 "what he would scream out in horror br br i'm sorry folks but <UNK> enough we "
 "had get carter the italian job <UNK> and now this what's the similarities no "
 "it's not exactly a <UNK> that three of the <UNK> star <UNK> <UNK> and the "
 'other stars another great british actor the main common <UNK> in those <UNK> '
 "is the <UNK> of the films they weren't made to impress hollywood they were "
 'quirky english films with a unique charm atmosphere that just cannot be '
 '<UNK> in the usa the word is cult and what better way to destroy a cult film '
 'than to <UNK> it with a remake or even a sequel br br <UNK> <UNK> had a '
 'tough task before it even hit the road <UNK> <UNK> is even more <UNK> that '
 'other said cult films it <UNK> genre intelligent scripts a grade actors the '
 'music score set pieces that <UNK> description and all the stories '
 "surrounding the film br br so here comes a remake don't worry no <UNK> were "
 '<UNK> in the making of this remake some major aspects of the story needed to '
 "be <UNK> for the modern usa <UNK> <UNK> <UNK> but that's just about the "
 'whole premise so we give the cop a <UNK> style <UNK> past complete with '
 'shock music flashbacks for the cheap scares then with no <UNK> phone <UNK> '
 "on the island that sorts the <UNK> out but in the real world this wouldn't "
 "happen cops just don't go missing give him a blood link for motivation "
 'rather than the <UNK> of beliefs and you have the remake <UNK> thin though '
 "isn't it br br it's just that it was all laid on with a <UNK> the name <UNK> "
 'were simply <UNK> almost carry on there was no sense of community on the '
 'island no <UNK> of town to catch your <UNK> just a few houses <UNK> about a '
 'forest and that was it <UNK> was just annoying by not giving out any <UNK> '
 'at all and cage was useless to let her get away with it when he went into '
 'the well you just knew he would get locked it the screenplay was <UNK> all '
 'the way to the end and you just wanted it to <UNK> up and end the <UNK> was '
 "absolutely hilarious and didn't know when to stop br br that ending is "
 'probably the best way to <UNK> the difference between the two one ends in '
 'the most beautiful <UNK> after the most horrific day the other ends with a '
 'post production explain it all to the <UNK> type conclusion br br i loved '
 'the original but went to the cinema with an open mind and was excited to see '
 'the film i left <UNK> in the knowledge that this film will probably end up '
 'beneath a <UNK> somewhere only this time <UNK> forgotten forever')
'---label---'
0

Maximum review length and minimum review length_____no_output_____### ⏸ The `<UNK>` tag is reserved for?

#### A. Special characters
#### B. Out of vocabulary words
#### C. Start of sentence
#### D. End of sentence_____no_output_____

### edTest(test_chow1) ###
# Submit an answer choice as a string below (eg. if you choose option C, put 'C')
answer1 = '___'_____no_output_____# For training we need our sequences to be of fixed length, but the reviews
# are of different sizes
print(f'Maximum review length: {len(max([i for i in X_train]+[i for i in X_test], key=len))}')
print(f'Minimum review length: {len(min([i for i in X_train]+[i for i in X_test], key=len))}')Maximum review length: 2494
Minimum review length: 7
# We also create two indices for short and long reviews
# we will use this later
idx_short = [i for i,val in enumerate(X_train) if len(val)<100]
idx_long = [i for i,val in enumerate(X_train) if len(val)>500]_____no_output_____

### Pad sequences

In order to feed this data into our RNN, all input documents must have the same length. We will limit the maximum review length to max_words by truncating longer reviews and padding shorter reviews. We can accomplish this using the pad_sequences() function in `tensorflow.keras`. For now, set max_words to 500._____no_output_____### ⏸ If we use *post-padding* on a sequence, the new sequence is"

#### A. "Padded with zeros before the start of original sequence"
#### B. "Padded with zeros after the end of the original sequence"
#### C. "Padded with ones before the start of original sequence"
#### D. "Padded with ones after the end of the original sequence"_____no_output_____

### edTest(test_chow2) ###
# Submit an answer choice as a string below (eg. if you choose option C, put 'C')
answer2 = '___'_____no_output_____# We will clip large reviews and pad smaller reviews to 500 words
max_words = 500
# We can pad the smaller sequences with 0s before, or after.
# This choice can severely affect network performance
# In the first case we, will pad after the sequence
# Please utilize sequence.pad_sequences()
postpad_X_train = ___
postpad_X_test = ________no_output_____

### RNN model for sentiment analysis

We build the model architecture in the code cell below. We have imported some layers from `tensorflow.keras` that you might need but feel free to use any other layers / transformations you like.

Remember that our input is a sequence of words (technically, integer word IDs) of maximum length = max_words, and our output is a binary sentiment label (0 or 1)._____no_output_____

def model_maker(summary=True,gru=False):
# One layer RNN model with 32 rnn cells
    embedding_size=32
    model=Sequential()
    model.add(Embedding(vocabulary_size, embedding_size, input_length=max_words))

    # We can specify if we want the GRU cell or the vanilla RNN cell
    if gru:
        model.add(GRU(32))
    else:
        model.add(SimpleRNN(32))
    model.add(Dense(1, activation='sigmoid'))
    if summary:
        print(model.summary())

    # model compile step
    model.compile(loss='binary_crossentropy',
                 optimizer='adam',
                 metrics=['accuracy'])
    return model_____no_output_____

## Trace Plot analysis
We expect the postpadded model to perform train slowly because of vanishing gradients.
Let us investigate the cause by training two models,
- One with shorter reviews (that were post padded) using `idx_short`
- The other with longer reviews (that were truncated) using `idx_long`_____no_output_____

# We build two new models with vanilla RNNs
model_short = model_maker(summary=False)
model_long = model_maker(summary=False)_____no_output_____# First we train `model_short` with short reviews
# set epochs to 10

X_short = ___
y_short = ___
epochs = ___
history_short = model_short.fit(___, ___, epochs=epochs,batch_size=640,verbose=0);_____no_output_____# Then we train `model_long` with short reviews

X_long = ___
y_long = ___
history_long = model_long.fit(___,___, epochs=epochs,batch_size=640,verbose=0);_____no_output_____### edTest(test_chow2_1) ###
X_short_shape, y_short_shape = X_short.shape, y_short.shape
X_long_shape, y_long_shape = X_long.shape, y_long.shape
print(X_short_shape, y_short_shape,X_long_shape, y_long_shape)_____no_output_____# Helper function to plot the data
# Plot the MSE of the model
plt.rcParams["figure.figsize"] = (8,6)
plt.title("Padding='post'",fontsize=20)
plt.semilogy(history_short.history['loss'], label='Shorter reviews', color='#FF9A98', linewidth=3)
plt.semilogy(history_long.history['loss'],  label='Longer reviews', color='#75B594', linewidth=3)
plt.legend()

# Set the axes labels
plt.xlabel('Epochs',fontsize=14)
plt.xticks(range(1,epochs,4))
plt.ylabel('MSE Loss',fontsize=14)
plt.legend(fontsize=14)
plt.show()_____no_output_____

## Pre-padding sequences

As we can see, the vanishing gradient problem is real and can severely affect the training of our network. The short review network negligibly trains, whereas the longer review model trains very well.
To counter this, we will now *pre-pad* the shorter sequences._____no_output_____

# We can pre-pad by using `sequence.pad_sequences` with `padding='pre'`
max_words = 500
prepad_X_train = ___
prepad_X_test = ________no_output_____

### Trace Plot - Take 2

Again, we investigate the trace plots for the two categories, but this time, with pre-padding_____no_output_____

# Reinitializing the models for the two categories
model_short = model_maker(summary=False)
model_long = model_maker(summary=False)_____no_output_____# Again we train `model_short` with short reviews

X_short = ___
y_short = ___
epochs = 10
history_short = model_short.fit(___, ___, epochs=epochs,batch_size=640,verbose=0);_____no_output_____# Then we train `model_long` with short reviews

X_long = ___
y_long = ___
history_long = model_long.fit(___, ___, epochs=epochs,batch_size=640,verbose=0);_____no_output_____# Helper function to plot the data
# Plot the MSE of the model
plt.rcParams["figure.figsize"] = (8,6)
plt.title("Padding='pre'",fontsize=20)
plt.semilogy(history_short.history['loss'], label='Shorter reviews', color='#FF9A98', linewidth=3)
plt.semilogy(history_long.history['loss'],  label='Longer reviews', color='#75B594', linewidth=3)
plt.legend()

# Set the axes labels
plt.xlabel('Epochs',fontsize=14)
plt.xticks(range(1,epochs,4))
plt.ylabel('MSE Loss',fontsize=14)
plt.legend(fontsize=14)
plt.show()_____no_output_____

### 🍲 Further improvements

We solved the vanishing gradient problem by *pre-padding* the sequences, but what other design choices can help you improve performance?_____no_output_____

### edTest(test_chow3) ###
# Submit your answer as a string below
answer3 = '___'_____no_output_____


# Notebook from tomdbar/rb-cqed
Path: jupyter_notebooks/notebooks-single/py_pulses/.ipynb_checkpoints/vStirap_simple_atom-checkpoint.ipynb

<code>
%matplotlib inline
#%matplotlib notebook

import matplotlib.pyplot as plt
import numpy as np
import scipy.constants as const
import time
import csv
from itertools import chain

from qutip import *

np.set_printoptions(threshold=np.inf)_____no_output_____'''
Define all of our relevant basis.
'''
i=np.complex(0,1)

'''
Rotation from Lab -> Atom: R_AL: {|+>,|->} -> {|H>,|V>}

R_AL.|+>_A = |+>_L --> R_AL.(1 0)^tr = (1  i)^tr / sqrt(2)
R_AL.|->_A = |->_L --> R_AL.(0 1)^tr = (1 -i)^tr / sqrt(2)
'''
R_AL = np.sqrt(1/2) * np.matrix([[1, i],
                                 [i, 1]])

'''
Rotation from Lab -> Cavity: R_LA: {|H>,|V>} -> {|X>,|Y>}
'''
# Option 1: Measured parameters for our cavity.
#alpha_CL, phi1_CL, phi2_CL = 0.459, -64.9*np.pi/180, -40.1*np.pi/180

## Option 2: Circular cavity eigenmodes.
## Note phi2=90 <--> |X>=|+>, phi2=-90 <--> |X>=|->
#alpha_CL, phi1_CL, phi2_CL = 1/np.sqrt(2), 0*np.pi/180, 90*np.pi/180

## Option 3: Linear cavity eigenmodes.
alpha_CL, phi1_CL, phi2_CL = 1, 0*np.pi/180, 0*np.pi/180

beta_CL = np.sqrt(1-alpha_CL**2)

R_CL = np.matrix([[alpha_CL* np.exp(i*phi1_CL), -beta_CL * np.exp(-i*phi2_CL)],
                  [beta_CL * np.exp(i*phi2_CL),  alpha_CL* np.exp(-i*phi1_CL)]])
R_LC = R_CL.getH()

'''
Rotation from Atom -> Cavity: R_LA: {|+>,|->} -> {|X>,|Y>}
'''
R_AC = R_LC*R_AL
R_CA = R_AC.getH()

'''
alpha, phi for writing the atomic basis in terms of the cavity basis
'''
alpha = np.abs(R_AC[0, 0])
beta = np.sqrt(1-alpha**2)
phi1, phi2 = np.angle(R_AC[0, 0]), np.angle(R_AC[1, 0])

args_system = dict([('alpha',alpha), ('beta',beta), ('phi1',phi1), ('phi2',phi2), ('i',i)])_____no_output_____'''
Coupling parameters.

   kappa:  Decay of the electric field out of the cavity.
   gamma:  Decay of the atomic amplitude.
   deltaP: Splitting of the cavity polarisation eigenmodes.
'''
kappa = 2 * 2.*np.pi
gamma = 0 * 2.*np.pi
deltaP = 20 * 2.*np.pi

# State the detuning of the Raman resonance
deltaOffsetVStirap = 0 * 2.*np.pi

# The laser and cavity detuning for photon production
deltaL = deltaOffsetVStirap
deltaC = deltaOffsetVStirap

deltaC -= deltaP/2 # Puts the |X> cavity mode in Raman resonance
#deltaC += deltaP/2 # Puts the |Y> cavity mode on Raman resonance

# Atom-cavity coupling rate
g = 4 * (2*np.pi)_____no_output_____'''
Here we perform the set up required before we can define and simulate the
Hamiltonian of the system.  This entails:
    1. Defining the size of our Hilbert space and how the states with it are
      labelled.
    2. Creating shorthand functions for the 'ket'/'bras' of the state vectors.
    3. Creating the operators who's expectation values we will wish to track
      through the simulations.  These are the number operators and for the
      population of each atomic state and each cavity mode in various
      polarisation bases (the cavity basis, {|H>,|V>} modes, {|+>,|->} modes
      etc.).
    4. Creating the collapse operators to account for spontaneous decay in our
      system.
'''

# Set where to truncate Fock states 0,1,2..,N-1 for cavity modes.
N = 2
cavStates = list(range(N))

# A dictionary of the atomic states.
atomStates = {
    "x0":0,
    "gM":1, "gP":2,
    "d":3 # d state is all excess spont. emm. will be dumped.
}
M = len(atomStates)

'''
State definitions - pre-calculated for speed.
'''
def ket(atom, cavX, cavY):
    return tensor(basis(M, atomStates[atom]), basis(N, cavX), basis(N, cavY))

def bra(atom, cavX, cavY):
    return ket(atom, cavX, cavY).dag()

kets,bras = {},{}
ketbras = {}
import itertools
s=[ list(atomStates), [0, 1], [0,1] ]
states = list(map(list, list(itertools.product(*s))))
for state in states:
    kets[str(state)] = ket(*state)
    bras[str(state)] = bra(*state)

for x in list(map(list, list(itertools.product(*[states,states])))):
    ketbras[str(x)] = ket(*x[0])*bra(*x[1])

'''
Shorthand definitions for coding Hamiltonians.
'''
i=np.complex(0,1)

def ePit(x,t):
    return np.exp(i*t*x)
def ePmit(x,t):
    return np.exp(-i*t*x)
def kb(x,y):
    return ketbras[str([x,y])]

'''
Create the operators that give the population of the atomic states.
'''
aDict = {}
def createStateOp(s):
    try:
        aOp = kb([s,0,0],[s,0,0]) + kb([s,1,0],[s,1,0]) + kb([s,0,1],[s,0,1])+ kb([s,1,1],[s,1,1])
        aDict[s]=aOp
    except KeyError:
        aOp =  None
    return aOp

[ad, agM, agP, ax0] = [createStateOp(s) for s in ["d", "gM", "gP", "x0"]]

'''
Create the photon number operators.

   anX, anY - photon number operator for cavity modes X, Y.

   anRotP_fast, anRotM_fast - photon number operators for different polarisation
        bases. They basis they correspond to is defined by the parameters of the
        rotation matrix from the new basis to the cavity basis:
               R = [[alpha* np.exp(i*phi1), -beta * np.exp(-i*phi2)],
                   [beta * np.exp(i*phi2),  alpha* np.exp(-i*phi1)]].
        '_fast' denotes that these functions return the appropriate operators
        at a given time, t, with all time-independent calculations already
        performed to maximise the speed.
'''
# Create the photon operators
aX = tensor(qeye(M), destroy(N), qeye(N))
aY = tensor(qeye(M), qeye(N), destroy(N))
anX = aX.dag()*aX
anY = aY.dag()*aY

# Pre-compute time independent terms to speed up calcualtions.
allAtomicStates = list(atomStates)

an_fast_1  = sum(map(lambda s: kb([s,1,0],[s,1,0]) + kb([s,1,1],[s,1,1]),allAtomicStates))
an_fast_2  = sum(map(lambda s: kb([s,0,1],[s,0,1]) + kb([s,1,1],[s,1,1]),allAtomicStates))
an_fast_3  = sum(map(lambda s: kb([s,0,1],[s,1,0]),allAtomicStates))
an_fast_4  = sum(map(lambda s: kb([s,1,0],[s,0,1]),allAtomicStates))

def anRotP_fast(t, alpha=alpha, phi1=phi1, phi2=phi2):
    beta = np.sqrt(1-alpha**2)
    delta_phi = phi2 - phi1
    return \
        (alpha**2 * an_fast_1 + beta**2 * an_fast_2) + \
        alpha*beta * (
            np.exp(-i*deltaP*t) * np.exp(i*delta_phi) * an_fast_3 + \
            np.exp(i*deltaP*t) * np.exp(-i*delta_phi) * an_fast_4
        )

def anRotM_fast(t, alpha=alpha, phi1=phi1, phi2=phi2):
    beta = np.sqrt(1-alpha**2)
    delta_phi = phi2 - phi1
    return \
        (alpha**2 * an_fast_2 + beta**2 * an_fast_1) - \
        alpha*beta * (
            np.exp(-i*deltaP*t) * np.exp(i*delta_phi) * an_fast_3 + \
            np.exp(i*deltaP*t) * np.exp(-i*delta_phi) * an_fast_4
        )

'''
Takes a series of times, t_series = [t0,t1,...], and returns
    [ [anRotP_fast[t0],anRotP_fast[t1],...],
      [anRotM_fast[t0],anRotM_fast[t1],...] ]
with the fewest possible calcualtions.  It is approx. twice as fast as calling
anRotP_fast and andRotM_fast independently for every time in t_series.
'''
def anRot_series_fast(t_series, alpha=alpha, phi1=phi1, phi2=phi2):
    beta = np.sqrt(1-alpha**2)
    delta_phi = phi2 - phi1

    an0P = (alpha**2 * an_fast_1 + beta**2 * an_fast_2)
    an0M = (alpha**2 * an_fast_2 + beta**2 * an_fast_1)

    an1s = [alpha*beta * (
                np.exp(-i*deltaP*t) * np.exp(i*delta_phi) * an_fast_3 + \
                np.exp(i*deltaP*t) * np.exp(-i*delta_phi) * an_fast_4
                )
            for t in t_series]

    anRots = [[an0P+an1,an0M-an1] for an1 in an1s]

    return [list(i) for i in zip(*anRots)]

'''
The c_op_list is the collapse operators of the system. Namely
  - the rate of photon decay from the cavity
  - spontaneous decay of the excited atomic states
'''
# Define collapse operators
c_op_list = []

# Cavity decay rate
c_op_list.append(np.sqrt(2*kappa) * aX)
c_op_list.append(np.sqrt(2*kappa) * aY)

spontEmmChannels = [
                    # |F',mF'> --> |F=1,mF=-1>
                    ('gM','x0',0),
                    ('gP','x0',0),
                    ('d','x0',1)
                    ]
#
# spontEmmChannels = []

spontDecayOps = []

for x in spontEmmChannels:
    try:
        spontDecayOps.append(x[2] * np.sqrt(2*gamma) *
                     tensor(
                         basis(M, atomStates[x[0]]) * basis(M, atomStates[x[1]]).dag(), qeye(N), qeye(N)))
    except KeyError:
        pass

c_op_list += spontDecayOps

sigma_spontDecayOp = sum([x.dag()*x for x in spontDecayOps])_____no_output_____'''
Define the pump pulse parameters.  The pump pulse by default has a
sin**2 amplitude profile.

Parameters:
    lengthStirap - length of the pump pulse
    wStirap - pi / lengthStirap
    lengthSim - length of the simulation (as the photon can continue to decay
                out of the cavity after the pump pulse is over, we typically
                use lengthSim > lengthStirap).
    OmegaStirap - the peak Rabi frequency of the pump pulse.  Note that, again,
                  this does not include the angualr dependence of the
                  transition so the Rabi frequency on a given transition is
                        A * OmegaStirap
                where,
                    A: Angular dependence of the considered transition.
                e.g. A=np.sqrt(5/24) for F=1,mF=+/-1 <--> F=1,mF=1 in zero
                magnetic field.
'''
lengthStirap= 1
lengthSim = 1.2
OmegaStirap = 2* 2 * np.pi

wStirap = np.pi / lengthStirap
t,tStep = np.linspace(0, lengthSim, 251, retstep=True)

Omega = lambda t, A=OmegaStirap, w=wStirap: np.piecewise(t, [t<lengthStirap], [A*np.sin(w*t)**2,0])

args_stirap = dict([('wStirap',wStirap), ('lengthStirap',lengthStirap), ('lengthSim',lengthSim)])

f, a1 = plt.subplots(figsize=(3, 2))
a1.plot(t, [Omega(x)/(2*np.pi) for x in t], 'b')
f.patch.set_facecolor('white')_____no_output_____'''
Create couplings for the Hamiltonian
'''

'''
Create a laser coupling.

Parameters:
    Omega - The peak rabi frequency of the pump pulse.
    g - The ground atomic atomic level.
    x - The excited atomic level.
    omegaL - The detuning of the pump laser.
    deltaM - The angular momentum change from g --> x.  This is ignored but included
             for consistancy with the cavityCoupling function.
    args_list - A dictionary of arguments for the qutip simulation.
    pulseShape - The shape of the pump pulse.

Returns:
    (List of cython-ready Hamiltonian terms,
     args_list with relevant parameters added)
'''
def laserCoupling(Omega,g,x,omegaL,deltaM,args_list,pulseShape='np.sin(w*t)**2'):
     omegaL_lab = 'omegaL_{0}{1}'.format(g,x)
     args_list[omegaL_lab] = omegaL
     return (
         [
         [ -(Omega/2)*(
                ( kb([g,0,0],[x,0,0]) + kb([g,0,1],[x,0,1]) + kb([g,1,0],[x,1,0]) + kb([g,1,1],[x,1,1]) ) +
                ( kb([x,0,0],[g,0,0]) + kb([x,0,1],[g,0,1]) + kb([x,1,0],[g,1,0]) + kb([x,1,1],[g,1,1]) )
        ),'{0} * np.cos({1}*t)'.format(pulseShape,omegaL_lab)],
        [ i*(Omega/2)*(
                ( kb([x,0,0],[g,0,0]) + kb([x,0,1],[g,0,1]) + kb([x,1,0],[g,1,0]) + kb([x,1,1],[g,1,1]) ) -
                ( kb([g,0,0],[x,0,0]) - kb([g,0,1],[x,0,1]) - kb([g,1,0],[x,1,0]) - kb([g,1,1],[x,1,1]) )
        ),'{0} * np.sin({1}*t)'.format(pulseShape,omegaL_lab)]
         ],
     args_list
     )

'''
Create a cavity coupling.

Parameters:
    g0 - The atom-cavity coupling rate.
    g - The ground atomic atomic level.
    x - The excited atomic level.
    omegaC - The detuning of the cavity resonance.
    deltaM - The angular momentum change from g --> x.
    args_list - A dictionary of arguments for the qutip simulation.

Returns:
    (List of cython-ready Hamiltonian terms,
     args_list with relevant parameters added)
'''
def cavityCoupling(g0,g,x,omegaC,deltaM,args_list):
    omegaC_X = omegaC + deltaP/2
    omegaC_Y = omegaC - deltaP/2
    omegaC_X_lab = 'omegaC_X_{0}{1}'.format(g,x)
    omegaC_Y_lab = 'omegaC_Y_{0}{1}'.format(g,x)

    args_list[omegaC_X_lab] = omegaC_X
    args_list[omegaC_Y_lab] = omegaC_Y

    if deltaM==1:
        H_coupling = (
            [
                [ -g0*alpha*(
                    kb([g,1,0],[x,0,0])+ kb([g,1,1],[x,0,1]) + kb([x,0,0],[g,1,0]) + kb([x,0,1],[g,1,1])
                ),'np.cos({0}*t + phi1)'.format(omegaC_X_lab)],

                [ -i*g0*alpha*(
                    kb([g,1,0],[x,0,0])+ kb([g,1,1],[x,0,1]) - kb([x,0,0],[g,1,0]) - kb([x,0,1],[g,1,1])
                ),'np.sin({0}*t + phi1)'.format(omegaC_X_lab)],

                [ -g0*beta*(
                    kb([g,0,1],[x,0,0]) + kb([g,1,1],[x,1,0]) + kb([x,0,0],[g,0,1]) + kb([x,1,0],[g,1,1])
                ),'np.cos({0}*t + phi2)'.format(omegaC_Y_lab)],

                [ -i*g0*beta*(
                    kb([g,0,1],[x,0,0]) + kb([g,1,1],[x,1,0]) - kb([x,0,0],[g,0,1]) - kb([x,1,0],[g,1,1])
                ),'np.sin({0}*t + phi2)'.format(omegaC_Y_lab)]
            ],
            args_list
        )
    elif deltaM==-1:
        H_coupling = (
            [
                [ -g0*alpha*(
                    kb([g,0,1],[x,0,0])+ kb([g,1,1],[x,1,0]) + kb([x,0,0],[g,0,1]) + kb([x,1,0],[g,1,1])
                ),'np.cos({0}*t - phi1)'.format(omegaC_Y_lab)],

                [ -i*g0*alpha*(
                    kb([g,0,1],[x,0,0])+ kb([g,1,1],[x,1,0]) - kb([x,0,0],[g,0,1]) - kb([x,1,0],[g,1,1])
                ),'np.sin({0}*t - phi1)'.format(omegaC_Y_lab)],

                [ g0*beta*(
                    kb([g,1,0],[x,0,0]) + kb([g,1,1],[x,0,1]) + kb([x,0,0],[g,1,0]) + kb([x,0,1],[g,1,1])
                ),'np.cos({0}*t - phi2)'.format(omegaC_X_lab)],

                [ i*g0*beta*(
                    kb([g,1,0],[x,0,0]) + kb([g,1,1],[x,0,1]) - kb([x,0,0],[g,1,0]) - kb([x,0,1],[g,1,1])
                ),'np.sin({0}*t - phi2)'.format(omegaC_X_lab)]
            ],
            args_list
        )
    else:
        raise Exception("deltaM must be +/-1")

    return H_coupling

'''
Get the laser and cavity couplings for between
all configured levels with angular momentum changes
of +/-1.

Parameters:
    delta - The detuning of the coupling field.
'''
def getCouplings_SigmaPlus(delta):
    return [
    ('gM', 'x0', delta, 1)
    ]

def getCouplings_SigmaMinus(delta):
    return [
    ('gP', 'x0', delta, -1)
    ]

'''
Define the couplings used here.
'''
# Setting the cavity to coupling the transition with deltaM=-1 from g->x will produce sigma_- photons
cavityCouplings = getCouplings_SigmaMinus(deltaC)
laserCouplings  = getCouplings_SigmaPlus(deltaL)
#laserCouplings = []

# Setting the cavity to coupling the transition with deltaM=+1 from g->x will produce sigma_+ photons
# cavityCouplings = getCouplings_SigmaPlus(deltaC)
# laserCouplings  = getCouplings_SigmaMinus(deltaL)

'''
Convert a list of couplings as returned by getCouplings...() to a list
of time dependent Hamiltonian terms as required by qutip.mesolve(...).

Needs: couplings:  the list of couplings as returned by getCouplings...()
       rabiFreq:   the intensity of the couplings (with angular dependence factored out)
       pulseShape: the shape of the pulse (in the string format for mesolve compilation
                   into Cython.
                   If None the coupling is presumed constant and a cavityCoupling(...)
                   is set, otherwise laserCoupling(...) is used.

Returns: hams: the list of Hamiltonian terms
         args_hams: the dictionary of args needed by qutip.mesolve
'''
def couplingsToHamiltonians(couplings, rabiFreq, pulseShape=None):
    hams, args_hams = [], dict()
    for x in couplings:
        # Check if this is a coupling between configured states.
        if  x[0] in atomStates and x[1] in atomStates:
            if pulseShape != None:
                ham, args_ham = laserCoupling(rabiFreq, x[0], x[1], x[2], x[3], args_hams, pulseShape)
            else:
                ham, args_ham = cavityCoupling(rabiFreq, x[0], x[1], x[2], x[3], args_hams)
            hams.append(ham)
    return list(chain(*hams)), args_hams

'''
Create Hamiltonian terms for configured couplings.
'''
hams_cavity, args_hams_cavity = couplingsToHamiltonians(cavityCouplings, g)
hams_laser, args_hams_laser = couplingsToHamiltonians(laserCouplings, OmegaStirap,
                                                 'np.piecewise(t, [t<lengthStirap], [np.sin(wStirap*t)**2,0])')

'''
The full list of Hamiltonian terms and arguments.
'''
H_Stirap = list(chain(*[hams_cavity, hams_laser]))
args_hams_Stirap = {**args_hams_cavity,**args_hams_laser,
                    **args_stirap,**args_system}_____no_output_____'''
When repeatedly simulating a system where only the time-dependent variables,
or initial state change, it is possible to reuse the Hamiltonian data stored
in QuTiP and there by avoid spending time needlessly preparing the Hamiltonian
and collapse terms for simulation. To turn on the the reuse features, we must
pass a qutip.Options object with the rhs_reuse flag turned on, or we can
pre-compile the Hamiltonian and pass it explicitly to the solver.
'''

# A list of Hamiltonians to pre-compile in the form:
#    (Cythonied Hamiltonian, arguement dictionary, compiled file name)
H_list = [
    (H_Stirap, args_hams_Stirap, 'H_Stirap')
]

H_rhs_compiled=[]

for H, args, label in H_list:
    t_start = time.time()
    rhs_generate(H, c_op_list, args=args, name=label, cleanup=False)
    H_rhs_compiled.append(H)
    print('Hamiltonian \'{0}\': states pre-computed in in {1} seconds'.format(label, time.time()-t_start))Hamiltonian 'H_Stirap': states pre-computed in in 0.016782045364379883 seconds
'''
Perform simulation
'''
t,tStep = np.linspace(0, lengthSim, 501, retstep=True)

psi0 = ket(*['gM',0,0]) # Initial state of the system

#psi0=(ket(*['gP',1,0]) - i*ket(*['gP',0,1]))/np.sqrt(2)

opts = Options(rhs_reuse=False, rhs_filename='H_Stirap')

t_start = time.time()

output = mesolve(H_Stirap, psi0, t, c_op_list, [], args=args_hams_Stirap, options=opts, progress_bar=None)

print('Simulation with {0} timesteps completed in {1} seconds'.format(
t.size,np.round(time.time()-t_start,3) ))Simulation with 501 timesteps completed in 0.266 seconds
'''
Handle simulation results.  Here we produce plots showing the
photon emission in three polarisation bases:
    - circular {|+>,|->},
    - cavity {|X>,|Y>},
    - linear {|H>,|V>},
and the evolution of the atomic state.
'''

output_states = output.states

'''
Photon number operators for calculating the population in the
circular and linear bases.
'''
#anP_t = [anRotP_fast(time, alpha=alpha, phi1=phi1, phi2=phi2) for time in t]
#anM_t = [anRotM_fast(time, alpha=alpha, phi1=phi1, phi2=phi2) for time in t]

print(alpha,phi1,phi2)

anP_t,anM_t = anRot_series_fast(t, alpha=alpha, phi1=phi1, phi2=phi2)

alpha_LC = np.clip(np.abs(R_LC[0, 0]),0,1)
phi1_LC, phi2_LC = np.angle(R_LC[0, 0]), np.angle(R_LC[1, 0])

# anH_t = [anRotP_fast(time, alpha=alpha_LC, phi1=phi1_LC, phi2=phi2_LC) for time in t]
# anV_t = [anRotM_fast(time, alpha=alpha_LC, phi1=phi1_LC, phi2=phi2_LC) for time in t]
anH_t,anV_t = anRot_series_fast(t, alpha=alpha_LC, phi1=phi1_LC, phi2=phi2_LC)

'''
Calculate expectation values of photon number operators.
'''
#####################
# Slow but explicit #
#####################
# exp_anP = np.abs(np.array([(x[0]*x[1]).tr() for x in zip(output_states, anP_t)]))
# exp_anM = np.abs(np.array([(x[0]*x[1]).tr() for x in zip(output_states, anM_t)]))
#
# exp_anH = np.abs(np.array([(x[0]*x[1]).tr() for x in zip(output_states, anH_t)]))
# exp_anV = np.abs(np.array([(x[0]*x[1]).tr() for x in zip(output_states, anV_t)]))
#
#exp_anX = np.abs(np.array([(x*anX).tr() for x in output_states]))
#exp_anY = np.abs(np.array([(x*anY).tr() for x in output_states]))
#
#exp_spontDecay = np.abs( np.array([(x*sigma_spontDecayOp).tr() for x in output_states]) )

###############
# > 3x Faster #
###############
exp_anP, exp_anM, exp_anH, exp_anV =\
    np.abs(np.array(
        [expect(list(an_list), state) for an_list,state in zip(zip(anP_t,anM_t,anH_t,anV_t),output_states)]
    )).T

exp_anX,exp_anY,exp_spontDecay =\
    np.abs( expect([anX,anY,sigma_spontDecayOp], output_states) )

# Total photonic population of each mode.
n_ph = np.trapz(2*kappa*(exp_anP+exp_anM), dx=tStep)

n_P = np.trapz(2*kappa*(exp_anP), dx=tStep)
n_M = np.trapz(2*kappa*(exp_anM), dx=tStep)
n_X = np.trapz(2*kappa*(exp_anX), dx=tStep)
n_Y = np.trapz(2*kappa*(exp_anY), dx=tStep)
n_H = np.trapz(2*kappa*(exp_anH), dx=tStep)
n_V = np.trapz(2*kappa*(exp_anV), dx=tStep)

# Total spontaneous emission.
n_spont = np.trapz(exp_spontDecay, dx=tStep)

# Atomic populations
# exp_agM = np.abs(np.array([(x*agM).tr() for x in output_states]) )
# exp_agP = np.abs( np.array([(x*agP).tr() for x in output_states]) )
# exp_ax0 = np.abs( np.array([(x*ax0).tr() for x in output_states]) )
# exp_ad = np.abs( np.array([(x*ad).tr() for x in output_states]) )
exp_agM,exp_agP,exp_ax0,exp_ad =\
    np.abs( expect([agM,agP,ax0,ad], output_states) )

'''
Summaries
'''
plt.rcParams['text.usetex'] = True

print('Photon:', np.round(n_ph,3))
print('Total spontaneous emission:', np.round(n_spont,3))
print('Photon / Total spontaneous emission:', np.round(n_ph/n_spont,3))

print('Photon + / Photon - / frac_+: {0} / {1} = {2}'.format(*[np.round(n,3) for n in [n_P,n_M,n_P/n_ph]]))
print('Photon X / Photon Y / frac_X: {0} / {1} = {2}'.format(*[np.round(n,3) for n in [n_X,n_Y,n_X/n_ph]]))
print('Photon H / Photon V / frac_H: {0} / {1} = {2}'.format(*[np.round(n,3) for n in [n_H,n_V,n_H/n_ph]]))

print('Final pops: |1,-1>,|1,1>:', np.round(exp_agM[-1],3),np.round(exp_agP[-1],3))

# Plot the results
f, (a1, a2, a3, a4) = plt.subplots(4, 1, sharex=True, figsize=(12, 11/.7))

a1.plot(t, 2*kappa * exp_anP, 'b', label='$+ (\sigma^{+})$')
a1.plot(t, 2*kappa * exp_anM, 'g', label='$- (\sigma^{-})$')
a1.set_ylabel('Cavity emission rate, $1/\mu s$')
a1.legend(loc=2)

a2.plot(t, 2*kappa * exp_anX, 'b', label='$X$')
a2.plot(t, 2*kappa * exp_anY, 'g', label='$Y$')
a2.set_ylabel('Cavity emission rate, $1/\mu s$')
a2.legend(loc=2)

a3.plot(t, 2*kappa * exp_anH, 'b', label='$H$')
a3.plot(t, 2*kappa * exp_anV, 'g', label='$V$')
a3.set_ylabel('Cavity emission rate, $1/\mu s$')
a3.legend(loc=2)

for a in [a1,a2,a3]:
    aB = a.twinx()
    aB.plot(t, [Omega(x)/(2*np.pi) for x in t], '--k', label='$\Omega(t)$')
    aB.legend(loc=0)

a4.plot(t, exp_agM, 'b', label='$gM$')
a4.plot(t, exp_agP, 'g', label='$gP$')
a4.plot(t, exp_ax0, '--r', label='$x0$')
a4.plot(t, exp_ad, '--y', label='$d$')
a4.legend(loc=2)0.707106781187 0.0 1.57079632679
Photon: 0.692
Total spontaneous emission: 0.0
Photon / Total spontaneous emission: inf
Photon + / Photon - / frac_+: 0.341 / 0.351 = 0.493
Photon X / Photon Y / frac_X: 0.685 / 0.008 = 0.989
Photon H / Photon V / frac_H: 0.685 / 0.008 = 0.989
Final pops: |1,-1>,|1,1>: 0.307 0.692
'''
Export photon shapes for outside use
'''
info_str="g{0}_k{1}_gamma{2}_deltaP{3}_Omega{4}_deltaC{5}".format(*[int(np.round(x/(2.*np.pi))) for x in
                                                                        [CGg1Mx1 * d * coupling_factor,
                                                                         kappa, gamma, deltaP, OmegaStirap * CGg1Mx1, deltaC]
                                                                                ])

export_targets = [
                  (t, "time"),
                  (list(map(float,2*kappa * exp_anR)), "emissionRateR"),
                  (list(map(float,2*kappa * exp_anL)), "emissionRateL"),
                  (list(map(float,2*kappa * exp_anH)), "emissionRateH"),
                  (list(map(float,2*kappa * exp_anV)), "emissionRateV"),
                  ]

export_path='/data/18-07-08/3lvlSystem/wavepackets/{0}/'.format(info_str)

export_dir_local = '..' + export_path
export_dir_aldaq = '/Volumes/KuhnGroup/Tom/Python/STIRAP modelling_Zeeman Scheme with Birefringence' + export_path
#export_dir = '/Volumes/KuhnGroup/Tom/Python/STIRAP modelling_Zeeman Scheme with Birefringence\
#              /data/18-06-13/Omega{0}gBar{1}deltaZ{2}deltaP{3}/gM_gP'.format(*[int(np.round(x/(2.*np.pi))) for x in
#                                                                             [OmegaStirap,d*coupling_factor,deltaZ,deltaP] ])

for export_dir in [export_dir_local,export_dir_aldaq]:

    if not os.path.exists(os.path.dirname(export_dir)):
         try:
             os.makedirs(os.path.dirname(export_dir))
             print('created dir: ', export_dir)
         except OSError as exc: # Guard against race condition
             if exc.errno != exc.errno.EEXIST:
                 raise
         except Exception:
             pass

    for data, fname in export_targets:
        with open(os.path.join(export_dir, fname + '.csv'), 'w+') as file:
            wr = csv.writer(file)
            wr.writerow(data)
        print (os.path.join(export_dir, fname + '.csv'))_____no_output_____
</code>