# Operações com *DataFrames*

Como dissemos anterioremente, o *DataFrame* é a segunda estrutura basilar do *pandas*. Um *DataFrame*:
- é uma tabela, ou seja, é bidimensional;
- tem cada coluna formada como uma *Series* do *pandas*;
- pode ter *Series* contendo tipos de dado diferentes.

import numpy as np
import pandas as pd

## Criação de um *DataFrame*

O método padrão para criarmos um *DataFrame* é através de uma função com mesmo nome.

```python
df_exemplo = pd.DataFrame(dados_de_interesse, index = indice_de_interesse,
                          columns = colunas_de_interesse)
```

Ao criar um *DataFrame*, podemos informar
- `index`: rótulos para as linhas (atributos *index* das *Series*).
- `columns`: rótulos para as colunas (atributos *name* das *Series*).

No _template_, `dados_de_interesse` pode ser

* um dicionário de:
  * *arrays* unidimensionais do *numpy*;
  * listas;
  * dicionários;
  * *Series* do *pandas*.
* um *array* bidimensional do *numpy*;
* uma *Series* do *Pandas*;
* outro *DataFrame*.

### *DataFrame* a partir de dicionários de *Series*

Neste método de criação, as *Series* do dicionário não precisam possuir o mesmo número de elementos. O *index* do *DataFrame* será dado pela **união** dos *index* de todas as *Series* contidas no dicionário.

Exemplo:

serie_Idade = pd.Series({'Ana':20, 'João': 19, 'Maria': 21, 'Pedro': 22}, name="Idade")

serie_Peso = pd.Series({'Ana':55, 'João': 80, 'Maria': 62, 'Pedro': 67, 'Túlio': 73}, name="Peso")

serie_Altura = pd.Series({'Ana':162, 'João': 178, 'Maria': 162, 'Pedro': 165, 'Túlio': 171}, name="Altura")

dicionario_series_exemplo = {'Idade': serie_Idade, 'Peso': serie_Peso, 'Altura': serie_Altura}

df_dict_series = pd.DataFrame(dicionario_series_exemplo)

df_dict_series

Compare este resultado com a criação de uma planilha pelos métodos usuais. Veja que há muita flexibilidade para criarmos ou modificarmos uma tabela.

Vejamos exemplos sobre como acessar intervalos de dados na tabela.

pd.DataFrame(dicionario_series_exemplo, index=['Ana','Maria'])

pd.DataFrame(dicionario_series_exemplo, index=['Ana','Maria'], columns=['Peso','Altura'])

Neste exemplo, adicionamos a coluna `IMC`, ainda sem valores calculados.

pd.DataFrame(dicionario_series_exemplo, index=['Ana','Maria','Paula'],
             columns=['Peso','Altura','IMC'])

df_exemplo_IMC = pd.DataFrame(dicionario_series_exemplo,
             columns=['Peso','Altura','IMC'])

Agora, mostramos como os valores do IMC podem ser calculados diretamente por computação vetorizada sobre as *Series*.

df_exemplo_IMC['IMC']=round(df_exemplo_IMC['Peso']/(df_exemplo_IMC['Altura']/100)**2,2)

df_exemplo_IMC

### *DataFrame* a partir de dicionários de listas ou *arrays* do *numpy*

Neste método de criação, os *arrays* ou as listas **devem** possuir o mesmo comprimento. Se o *index* não for informado, o *index* será dado de forma similar ao do objeto tipo *Series*.

Exemplo com dicionário de listas:

dicionario_lista_exemplo = {'Idade': [20,19,21,22,20],
                            'Peso': [55,80,62,67,73],
                            'Altura': [162,178,162,165,171]}

pd.DataFrame(dicionario_lista_exemplo)

Mais exemplos:

pd.DataFrame(dicionario_lista_exemplo, index=['Ana','João','Maria','Pedro','Túlio'])

Exemplos com dicionário de *arrays* do *numpy*:

dicionario_array_exemplo = {'Idade': np.array([20,19,21,22,20]),
                            'Peso': np.array([55,80,62,67,73]),
                            'Altura': np.array([162,178,162,165,171])}

pd.DataFrame(dicionario_array_exemplo)

Mais exemplos:

pd.DataFrame(dicionario_array_exemplo, index=['Ana','João','Maria','Pedro','Túlio'])

### *DataFrame* a partir de uma *Series* do *pandas*

Neste caso, o *DataFrame* terá o mesmo *index* que a *Series* do *pandas* e apenas uma coluna.

series_exemplo = pd.Series({'Ana':20, 'João': 19, 'Maria': 21, 'Pedro': 22, 'Túlio': 20})

pd.DataFrame(series_exemplo)

Caso a *Series* possua um atributo `name` especificado, este será o nome da coluna do *DataFrame*.

series_exemplo_Idade = pd.Series({'Ana':20, 'João': 19, 'Maria': 21, 'Pedro': 22, 'Túlio': 20}, name="Idade")

pd.DataFrame(series_exemplo_Idade)

### *DataFrame* a partir de lista de *Series* do *pandas*

Neste caso, a entrada dos dados da lista no *DataFrame* será feita por linha.

pd.DataFrame([serie_Peso, serie_Altura, serie_Idade])

Podemos corrigir a orientação usando o método `transpose`.

pd.DataFrame([serie_Peso, serie_Altura, serie_Idade]).transpose()

### *DataFrame* a partir de arquivos

Para criar um *DataFrame* a partir de um arquivo, precisamos de funções do tipo `pd.read_FORMATO`, onde `FORMATO` indica o formato a ser importado sob o pressuposto de que a biblioteca *pandas* foi devidamente importada com `pd`.

Os formatos mais comuns são:

* *csv* (comma-separated values),
* *xls* ou *xlsx* (formatos do Microsoft Excel),
* *hdf5* (comumente utilizado em *big data*),
* *json* (comumente utilizado em páginas da internet).

As funções para leitura correspondentes são:
* `pd.read_csv`,
* `pd.read_excel`,
* `pd.read_hdf`,
* `pd.read_json`,

respectivamente.

De todas elas, a função mais utilizada é `read_csv`. Ela possui vários argumentos. Vejamos os mais utilizados:

* `file_path_or_buffer`: o endereço do arquivo a ser lido. Pode ser um endereço da internet.
* `sep`: o separador entre as entradas de dados. O separador padrão é `,`.
* `index_col`: a coluna que deve ser usada para formar o *index*. O padrão é `None`. Porém pode ser alterado para outro. Um separador comumente encontrado é o `\t` (TAB).
* `names`: nomes das colunas a serem usadas. O padrão é `None`.
* `header`: número da linha que servirá como nome para as colunas. O padrão é `infer` (ou seja, tenta deduzir automaticamente). Se os nomes das colunas forem passados através do `names`, então `header` será automaticamente considerado como `None`.

**Exemplo:** considere o arquivo `data/exemplo_data.csv` contendo:

```
,coluna_1,coluna_2
2020-01-01,-0.4160923582996922,1.8103644347460834
2020-01-02,-0.1379696602473578,2.5785204825192785
2020-01-03,0.5758273450544708,0.06086648807755068
2020-01-04,-0.017367186564883633,1.2995865328684455
2020-01-05,1.3842792448510655,-0.3817320973859929
2020-01-06,0.5497056238566345,-1.308789022968975
2020-01-07,-0.2822962331437976,-1.6889791765925102
2020-01-08,-0.9897300598660013,-0.028120707936426497
2020-01-09,0.27558240737928663,-0.1776585993494299
2020-01-10,0.6851316082235455,0.5025348904591399
```

Para ler o arquivo acima basta fazer:

df_exemplo_0 = pd.read_csv('data/exemplo_data.csv')

df_exemplo_0

No exemplo anterior, as colunas receberam nomes corretamentes exceto pela primeira coluna que gostaríamos de considerar como *index*. Neste caso fazemos:

df_exemplo = pd.read_csv('data/exemplo_data.csv', index_col=0)

df_exemplo

#### O método *head* do *DataFrame*

O método `head`, sem argumento, permite que visualizemos as 5 primeiras linhas do *DataFrame*.

df_exemplo.head()

Se for passado um argumento com valor `n`, as `n` primeiras linhas são impressas.

df_exemplo.head(2)

df_exemplo.head(7)

#### O método `tail` do *DataFrame*

O método `tail`, sem argumento, retorna as últimas 5 linhas do *DataFrame*.

df_exemplo.tail()

Se for passado um argumento com valor `n`, as `n` últimas linhas são impressas.

df_exemplo.tail(2)

df_exemplo.tail(7)

## Atributos de *Series* e *DataFrames*

Atributos comumente usados para *Series* e *DataFrames* são:

* `shape`: fornece as dimensões do objeto em questão (*Series* ou *DataFrame*) em formato consistente com o atributo `shape` de um *array* do *numpy*.
* `index`: fornece o índice do objeto. No caso do *DataFrame* são os rótulos das linhas.
* `columns`: fornece as colunas (apenas disponível para *DataFrames*)

Exemplo:

df_exemplo.shape

serie_1 = pd.Series([1,2,3,4,5])

serie_1.shape

df_exemplo.index

serie_1.index

df_exemplo.columns

Se quisermos obter os dados contidos nos *index* ou nas *Series* podemos utilizar a propriedade `.array`.

serie_1.index.array

df_exemplo.columns.array

Se o interesse for obter os dados como um `array` do *numpy*, devemos utilizar o método `.to_numpy()`.

Exemplo:

serie_1.index.to_numpy()

df_exemplo.columns.to_numpy()

O método `.to_numpy()` também está disponível em *DataFrames*:

df_exemplo.to_numpy()

A função do *numpy* `asarray()` é compatível com *index*, *columns* e *DataFrames* do *pandas*:

np.asarray(df_exemplo.index)

np.asarray(df_exemplo.columns)

np.asarray(df_exemplo)

### Informações sobre as colunas de um *DataFrame*

Para obtermos uma breve descrição sobre as colunas de um *DataFrame* utilizamos o método `info`.

Exemplo:

df_exemplo.info()

## Criando arquivos a partir de *DataFrames*

Para criar arquivos a partir de *DataFrames*, basta utilizar os métodos do tipo `pd.to_FORMATO`, onde `FORMATO` indica o formato a ser exportado e supondo que a biblioteca *pandas* foi importada com `pd`.

Com relação aos tipos de arquivo anteriores, os métodos para exportação correspondentes são:
* `.to_csv` ('endereço_do_arquivo'),
* `.to_excel` ('endereço_do_arquivo'),
* `.to_hdf` ('endereço_do_arquivo'),
* `.to_json`('endereço_do_arquivo'),

onde `endereço_do_arquivo` é uma `str` que contém o endereço do arquivo a ser exportado.

Exemplo:

Para exportar para o arquivo `exemplo_novo.csv`, utilizaremos o método `.to_csv` ao *DataFrame* `df_exemplo`:

df_exemplo.to_csv('data/exemplo_novo.csv')

### Exemplo COVID-19 PB

Dados diários de COVID-19 do estado da Paraíba:

*Fonte: https://superset.plataformatarget.com.br/superset/dashboard/microdados/*

dados_covid_PB = pd.read_csv('https://superset.plataformatarget.com.br/superset/explore_json/?form_data=%7B%22slice_id%22%3A1550%7D&csv=true',
                             sep=',', index_col=0)

dados_covid_PB.info()

dados_covid_PB.head()

dados_covid_PB.tail()

dados_covid_PB['estado'] = 'PB'

dados_covid_PB.head()

dados_covid_PB.to_csv('data/dadoscovidpb.csv')

### Índices dos valores máximos ou mínimos

Os métodos `idxmin()` e `idxmax()` retornam o *index* cuja entrada fornece o valor mínimo ou máximo da *Series* ou *DataFrame*. Se houver múltiplas ocorrências de mínimos ou máximos, o método retorna a primeira ocorrência.

Vamos recriar um *DataFrame* genérico.

serie_Idade = pd.Series({'Ana':20, 'João': 19, 'Maria': 21, 'Pedro': 22, 'Túlio': 20}, name="Idade")
serie_Peso = pd.Series({'Ana':55, 'João': 80, 'Maria': 62, 'Pedro': 67, 'Túlio': 73}, name="Peso")
serie_Altura = pd.Series({'Ana':162, 'João': 178, 'Maria': 162, 'Pedro': 165, 'Túlio': 171}, name="Altura")

dicionario_series_exemplo = {'Idade': serie_Idade, 'Peso': serie_Peso, 'Altura': serie_Altura}

df_dict_series = pd.DataFrame(dicionario_series_exemplo)

df_dict_series

Assim, podemos localizar quem possui menores idade, peso e altura.

df_dict_series.idxmin()

De igual forma, localizamos quem possui maiores idade, peso e altura.

df_dict_series.idxmax()

**Exemplo:** Aplicaremos as funções `idxmin()` e `idxmax()` aos dados do arquivo `data/exemplo_data.csv` para localizar entradas de interesse.

df_exemplo = pd.read_csv('data/exemplo_data.csv', index_col=0); df_exemplo

df_exemplo = pd.DataFrame(df_exemplo, columns=['coluna_1','coluna_2','coluna_3'])

Inserimos uma terceira coluna com dados fictícios.

df_exemplo['coluna_3'] = pd.Series([1,2,3,4,5,6,7,8,np.nan,np.nan],index=df_exemplo.index)
df_exemplo

Os *index* correspondentes aos menores e maiores valores são datas, evidentemente.

df_exemplo.idxmin()

df_exemplo.idxmax()

### Reindexação de *DataFrames*

No *pandas*, o método `reindex`

* reordena o *DataFrame* de acordo com o conjunto de rótulos inserido como argumento;
* insere valores faltantes caso um rótulo do novo *index* não tenha valor atribuído no conjunto de dados;
* remove valores correspondentes a rótulos que não estão presentes no novo *index*.

Exemplos:

df_dict_series.reindex(index=['Victor', 'Túlio', 'Pedro', 'João'], columns=['Altura','Peso','IMC'])

## Remoção de linhas ou colunas de um *DataFrame*

Para remover linhas ou colunas de um *DataFrame* do *pandas* podemos utilizar o método `drop`. O argumento `axis` identifica o eixo de remoção: `axis=0`, que é o padrão, indica a remoção de uma ou mais linhas; `axis=1` indica a remoção de uma ou mais colunas.

Nos exemplos que segue, note que novos *DataFrames* são obtidos a partir de `df_dict_series` sem que este seja sobrescrito.

df_dict_series.drop('Túlio') # axis=0 implícito

df_dict_series.drop(['Ana','Maria'], axis=0)

df_dict_series.drop(['Idade'], axis=1)

### Renomeando *index* e *columns*

O método `rename` retorna uma cópia na qual o *index* (no caso de *Series* e *DataFrames*) e *columns* (no caso de *DataFrames*) foram renomeados. O método aceita como entrada um dicionário, uma *Series* do *pandas* ou uma função.

Exemplo:

serie_exemplo = pd.Series([1,2,3], index=['a','b','c'])

serie_exemplo

serie_exemplo.rename({'a':'abacaxi', 'b':'banana', 'c': 'cebola'})

Exemplo:

df_dict_series

df_dict_series.rename(index = {'Ana':'a', 'João':'j', 'Maria':'m', 'Pedro':'p','Túlio':'t'},
                     columns = {'Idade':'I', 'Peso':'P','Altura':'A'})

No próximo exemplo, usamos uma *Series* para renomear os rótulos.

indice_novo = pd.Series({'Ana':'a', 'João':'j', 'Maria':'m', 'Pedro':'p','Túlio':'t'})

df_dict_series.rename(index = indice_novo)

Neste exemplo, usamos a função `str.upper` (altera a `str` para "todas maiúsculas") para renomear colunas.

df_dict_series.rename(columns=str.upper)

## Ordenação de *Series* e *DataFrames*

É possível ordenar ambos pelos rótulos do *index* (para tanto é necessário que eles sejam ordenáveis) ou por valores nas colunas.

O método `sort_index` ordena a *Series* ou o *DataFrame* pelo *index*. O método `sort_values` ordena a *Series* ou o *DataFrame* pelos valores (escolhendo uma ou mais colunas no caso de *DataFrames*). No caso do *DataFrame*, o argumento `by` é necessário para indicar qual(is) coluna(s) será(ão) utilizada(s) como base para a ordenação.

Exemplos:

serie_desordenada = pd.Series({'Maria': 21, 'Pedro': 22, 'Túlio': 20, 'João': 19, 'Ana':20});
serie_desordenada

serie_desordenada.sort_index() # ordenação alfabética

Mais exemplos:

df_desordenado = df_dict_series.reindex(index=['Pedro','Maria','Ana','Túlio','João'])

df_desordenado

df_desordenado.sort_index()

Mais exemplos:

serie_desordenada.sort_values()

df_desordenado.sort_values(by=['Altura']) # ordena por 'Altura'

No caso de "empate", podemos utilizar outra coluna para desempatar.

df_desordenado.sort_values(by=['Altura','Peso']) # usa a coluna 'Peso' para desempatar

Os métodos `sort_index` e `sort_values` admitem o argumento opcional `ascending`, que permite inverter a ordenação caso tenha valor `False`.

df_desordenado.sort_index(ascending=False)

df_desordenado.sort_values(by=['Idade'], ascending=False)

## Comparação de *Series* e *DataFrames*

*Series* e *DataFrames* possuem os seguintes métodos de comparação lógica:

- `eq` (igual);
- `ne` (diferente);
- `lt` (menor do que);
- `gt` (maior do que);
- `le` (menor ou igual a);
- `ge` (maior ou igual a)

que permitem a utilização dos operadores binários `==`, `!=`, `<`, `>`, `<=` e `>=`, respectivamente. As comparações são realizadas em cada entrada da *Series* ou do *DataFrame*.

**Observação**: Para que esses métodos sejam aplicados, todos os objetos presentes nas colunas do *DataFrame* devem ser de mesma natureza. Por exemplo, se um *DataFrame* possui algumas colunas numéricas e outras com *strings*, ao realizar uma comparação do tipo `> 1`, um erro ocorrerá, pois o *pandas* tentará comparar objetos do tipo `int` com objetos do tipo `str`, assim gerando uma incompatibilidade.

Exemplos:

serie_exemplo

serie_exemplo == 2

De outra forma:

serie_exemplo.eq(2)

serie_exemplo > 1

Ou, na forma funcional:

serie_exemplo.gt(1)

df_exemplo > 1

**Importante:** Ao comparar *np.nan*, o resultado tipicamente é falso:

np.nan == np.nan

np.nan > np.nan

np.nan >= np.nan

Só é verdadeiro para indicar que é diferente:

np.nan != np.nan

Neste sentido, podemos ter tabelas iguais sem que a comparação usual funcione:

# 'copy', como o nome sugere, gera uma cópia do DataFrame
df_exemplo_2 = df_exemplo.copy()

(df_exemplo == df_exemplo_2)

O motivo de haver entradas como `False` ainda que `df_exemplo_2` seja uma cópia exata de `df_exemplo` é a presença do `np.nan`. Neste caso, devemos utilizar o método `equals` para realizar a comparação.

df_exemplo.equals(df_exemplo_2)

## Os métodos `any`, `all` e a propriedade `empty`

O método `any` é aplicado a entradas booleanas (verdadeiras ou falsas) e retorna *verdadeiro* se existir alguma entrada verdadeira, ou *falso*, se todas forem falsas. O método `all` é aplicado a entradas booleanas e retorna *verdadeiro* se todas as entradas forem verdadeiras, ou *falso*,  se houver pelo menos uma entrada falsa. A propriedade `empty` retorna *verdadeiro* se a *Series* ou o *DataFrame* estiver vazio, ou *falso* caso contrário.

Exemplos:

serie_exemplo

serie_exemplo_2 = serie_exemplo-2;
serie_exemplo_2

(serie_exemplo_2 > 0).any()

(serie_exemplo > 1).all()

Este exemplo reproduz um valor `False` único.

(df_exemplo == df_exemplo_2).all().all()

serie_exemplo.empty

Mais exemplos:

(df_exemplo == df_exemplo_2).any()

df_exemplo.empty

df_vazio = pd.DataFrame()

df_vazio.empty

## Seleção de colunas de um *DataFrame*

Para selecionar colunas de um *DataFrame*, basta aplicar *colchetes* a uma lista contendo os nomes das colunas de interesse.

No exemplo abaixo, temos um *DataFrame* contendo as colunas `'Idade'`, `'Peso'` e `'Altura'`. Selecionaremos `'Peso'` e `'Altura'`, apenas.

df_dict_series[['Peso','Altura']]

Se quisermos selecionar apenas uma coluna, não há necessidade de inserir uma lista. Basta utilizar o nome da coluna:

df_dict_series['Peso']

Para remover colunas, podemos utilizar o método `drop`.

df_dict_series.drop(['Peso','Altura'], axis=1)

### Criação de novas colunas a partir de colunas existentes

Um método eficiente para criar novas colunas a partir de colunas já existentes é `eval`. Neste método, podemos utilizar como argumento uma *string* contendo uma expressão matemática envolvendo nomes de colunas do *DataFrame*.

Como exemplo, vamos ver como calcular o IMC no *DataFrame* anterior:

df_dict_series.eval('Peso/(Altura/100)**2')

Se quisermos obter um *DataFrame* contendo o IMC como uma nova coluna, podemos utilizar o método `assign` (sem modificar o *DataFrame* original):

df_dict_series.assign(IMC=round(df_dict_series.eval('Peso/(Altura/100)**2'),2))

df_dict_series # não modificado

Se quisermos modificar o *DataFrame* para incluir a coluna IMC fazemos:

df_dict_series['IMC']=round(df_dict_series.eval('Peso/(Altura/100)**2'),2)

df_dict_series # modificado "in-place"

## Seleção de linhas de um *DataFrame*

Podemos selecionar linhas de um *DataFrame* de diversas formas diferentes. Veremos agora algumas dessas formas.

Diferentemente da forma de selecionar colunas, para selecionar diretamente linhas de um *DataFrame* devemos utilizar o método `loc` (fornecendo o *index*, isto é, o rótulo da linha) ou o `iloc` (fornecendo a posição da linha):

Trabalharemos a seguir com um banco de dados atualizado sobre a COVID-19. Para tanto, importaremos o módulo `datetime` que nos auxiliará com datas.

import datetime

dados_covid_PB = pd.read_csv('https://superset.plataformatarget.com.br/superset/explore_json/?form_data=%7B%22slice_id%22%3A1550%7D&csv=true',
                             sep=',', index_col=0)

# busca o banco na data D-1, visto que a atualização
# ocorre em D
ontem = (datetime.date.today() - datetime.timedelta(days=1)).strftime('%Y-%m-%d')

dados_covid_PB.head(1)

Podemos ver as informações de um único dia como argumento. Para tanto, excluímos a coluna `'Letalidade'` (valor não inteiro) e convertemos o restante para valores inteiros:

dados_covid_PB.loc[ontem].drop('Letalidade').astype('int')

Podemos selecionar um intervalo de datas como argumento (excluindo letalidade):

dados_covid_PB.index = pd.to_datetime(dados_covid_PB.index) # Convertendo o index de string para data
dados_covid_PB.loc[pd.date_range('2021-02-01',periods=5,freq="D")].drop('Letalidade',axis=1)
                #função pd.date_range é muito útil para criar índices a partir de datas.

Podemos colocar uma lista como argumento:

dados_covid_PB.loc[pd.to_datetime(['2021-01-01','2021-02-01'])]

Vamos agora examinar os dados da posição 100 (novamente excluindo a coluna letalidade e convertendo para inteiro):

dados_covid_PB.iloc[100].drop('Letalidade').astype('int')

Podemos colocar um intervalo como argumento:

dados_covid_PB.iloc[50:55].drop('Letalidade', axis=1).astype('int')

### Seleção de colunas com `loc` e `iloc`

Podemos selecionar colunas utilizando os métodos `loc` e `iloc` utilizando um argumento adicional.

dados_covid_PB.loc[:,['casosNovos','obitosNovos']]

dados_covid_PB.iloc[:,4:6] # fatiamento na coluna

### Seleção de linhas e colunas específicas com `loc` e `iloc`

Usando o mesmo princípio de *fatiamento* aplicado a *arrays* do numpy, podemos selecionar linhas e colunas em um intervalo específico de forma a obter uma subtabela.

dados_covid_PB.iloc[95:100,4:6]

Neste exemplo um pouco mais complexo, buscamos casos novos e óbitos novos em um período específico e ordenamos a tabela da data mais recente para a mais antiga.

dados_covid_PB.loc[pd.date_range('2020-04-06','2020-04-10'),['casosNovos','obitosNovos']].sort_index(ascending=False)

Suponha que o peso de Ana foi medido corretamente, mas registrado de maneira errônea no *DataFrame* `df_dict_series` como 55.

df_dict_series

Supondo que, na realidade, o valor é 65, alteramos a entrada específica com um simples `loc`. Em seguida, atualizamos a tabela.

df_dict_series.loc['Ana','Peso'] = 65

df_dict_series = df_dict_series.assign(IMC=round(df_dict_series.eval('Peso/(Altura/100)**2'),2)) # O IMC mudou

df_dict_series

### Seleção de linhas através de critérios lógicos ou funções

Vamos selecionar quais os dias em que houve mais de 40 mortes registradas:

dados_covid_PB.loc[dados_covid_PB['obitosNovos']>40]

Selecionando os dias com mais de 25 óbitos e mais de 1500 casos novos:

dados_covid_PB.loc[(dados_covid_PB.obitosNovos > 25) & (dados_covid_PB.casosNovos>1500)]

**Obs**.: Note que podemos utilizar o nome da coluna como um atributo.

Vamos inserir uma coluna sobrenome no `df_dict_series`:

df_dict_series['Sobrenome'] = ['Silva', 'PraDo', 'Sales', 'MachadO', 'Coutinho']
df_dict_series

Vamos encontrar as linhas cujo sobrenome termina em "do". Para tanto, note que a função abaixo retorna `True` se o final é "do" e `False` caso contrário.

```python
def verifica_final_do(palavra):
    return palavra.lower()[-2:] == 'do'
```
**Obs**.: Note que convertemos tudo para minúsculo.

Agora vamos utilizar essa função para alcançar nosso objetivo:

# 'map' aplica a função lambda a cada elemento da *Series*
df_dict_series['Sobrenome'].map(lambda palavra: palavra.lower()[-2:]=='do')

# procurando no df inteiro
df_dict_series.loc[df_dict_series['Sobrenome'].map(lambda palavra: palavra.lower()[-2:]=='do')]

Vamos selecionar as linhas do mês 2 (fevereiro) usando `index.month`:

dados_covid_PB.loc[dados_covid_PB.index.month==2].head()

### Seleção de linhas com o método *query*

Similarmente ao método `eval`, ao utilizarmos `query`, podemos criar expressões lógicas a partir de nomes das colunas do *DataFrame*.

Assim, podemos reescrever o código

```python
dados_covid_PB.loc[(dados_covid_PB.obitosNovos>25) &
                   (dados_covid_PB.casosNovos>1500)]
```
como

dados_covid_PB.query('obitosNovos>25 and casosNovos>1500') # note que 'and' é usado em vez de '&'


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on May 31 2019
@author: Nathan de Lara <ndelara@enst.fr>
"""
import numpy as np

def top_k(scores: np.ndarray, k: int = 1):
    """Index of the k elements with highest value.

    Parameters
    ----------
    scores : np.ndarray
        Array of values.
    k : int
        Number of elements to return.

    Examples
    --------
    >>> scores = np.array([0, 1, 0, 0.5])
    >>> top_k(scores, k=2)
    array([1, 3])

    Notes
    -----
    This is a basic implementation that sorts the entire array to find its top k elements.
    """
    return np.argsort(-scores)[:k]


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

# In[ ]:

# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load in

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the "../input/" directory.
# For example, running this (by clicking run or pressing Shift+Enter) will list the files in the input directory

from subprocess import check_output
print(check_output(["ls", "../input"]).decode("utf8"))

# Any results you write to the current directory are saved as output.

# Read in the file (King's County housing prices) and summarize the information

df = pd.read_csv('../input/kc_house_data.csv')

df.info()

# In[ ]:

df.describe() # Look at a description of the features

# In[ ]:

df.columns # What are the features?

# In[ ]:

feat_null = [feat for feat in df.columns if df[feat].isnull().sum()!=0]
feat_null # No Null values
# Delve deeper later

# In[ ]:


# 1. Start Python and check versions

# Check the versions of libraries
# Python version
import sys
print('Python: {}'.format(sys.version))
# scipy
import scipy
print('scipy: {}'.format(scipy.__version__))
# numpy
import numpy
print('numpy: {}'.format(numpy.__version__))
# matplotlib
import matplotlib
print('matplotlib: {}'.format(matplotlib.__version__))
# pandas
import pandas
print('pandas: {}'.format(pandas.__version__))
# scikit-learn
import sklearn
print('sklearn: {}'.format(sklearn.__version__))

# FUTURE: Develop into a function
#def lib_review(libref):
#    import libref
#    print('libref: {}'.format(libref.__version__))
#    return librefout

#lib_review(scipy)


# coding: utf-8

# ># IOOS System Test: Rapid Deployment gages from USGS
#
# ###Can we obtain water level data from the rapid deployment gages, deployment for Hurricane Irene?
# This notebook is based on IOOS System Test: Inundation
#
# #### Methodology:
#
# * USGS gage data (ASCII Text Files) obtained from http://ga.water.usgs.gov/flood/hurricane/irene/sites/datafiles/ and 90% of data zipped up for use in the notebook (large files were removed for efficiency)
# * Station data Automatically gets unzipped
# * Process data files (195 stations), and store data in dictionary for access
# * Plot Water level data for the New Jersey area
# * Plot Barometric pressure data for the New Jersey area
# * Plot Gage Locations on a map, overlaid with the Hurricane Irene track line
# * Plot time series of maximum waterlevel
# * Plot locations of maximum waterlevel

# ### import required libraries

# In[1]:

import re
import os
import csv
import uuid
import zipfile

from datetime import datetime

import folium
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.dates as md
import matplotlib.pyplot as plt

from IPython.display import HTML, Javascript, display

from utilities import css_styles, inline_map
css_styles()

#
Extract Data  - Does the data dir exist, if not extract it

# In[2]:

def unzip(source_filename, dest_dir):
    with zipfile.ZipFile(source_filename) as zf:
        zf.extractall(dest_dir)

# In[3]:

if os.path.isdir("data_files"):
    pass
else:
    print("Data Dir does not exist... Extracting.")
    unzip('sample_data_files.zip', os.getcwd())

# In[4]:

files = os.listdir('data_files')
print("Water Level Files: %s" % len(files))

#
Process Data  - Read the data files and create a dict of the fields

# In[5]:

def parse_metadata(fname):
    meta_data = {}
    non_decimal = re.compile(r'[^\d.]+')
    fields = {'Sensor location latitude': 'lat',
              'Sensor location longitude': 'lon',
              'Site id =': 'name',
              'Sensor elevation above NAVD 88 =': 'elevation',
              'Barometric sensor site (source of bp) =': 'bp_source',
              'Lowest recordable water elevation is': 'lowest_wl'}
    with open(os.path.join('data_files', fname)) as f:
        content = f.readlines()
        for k, ln in enumerate(content):
            content[k] = ln.strip()
            if content[k].startswith('#'):
                for fd in fields:
                    if fd in content[k]:
                        if fields[fd] == 'name':
                            meta_data[fields[fd]] = content[k].split(fd)[-1]
                        else:
                            val = (content[k].split(fd)[-1])
                            val = float(non_decimal.sub('', val))
                            meta_data[fields[fd]] = val
                        if fields[fd] == 'lon':
                            meta_data[fields[fd]] = -meta_data[fields[fd]]
    return meta_data

# In[6]:

divid = str(uuid.uuid4())

pb = HTML("""

""" % divid)

display(pb)
full_data = {}
for count, fname in enumerate(files):
    meta_data = parse_metadata(fname)
    kw = dict(parse_dates=True, sep='\t', skiprows=29, index_col=0)
    actual_data = pd.read_csv(os.path.join('data_files', fname), **kw)
    full_data[fname] = {'meta': meta_data,
                        'data': actual_data}

    percent_complete = ((float(count+1) / float(len(files))) * 100.)
    display(Javascript("$('div#%s').width('%i%%')" %
                       (divid, int(percent_complete))))

# #### Show the available fields from the processed data files

# In[7]:

print("Data Fields: {}, {}, {}".format(actual_data.index.name,
                                       *actual_data.columns))

# # Remove 'Sensor elevation above NAVD 88 (ft)'

# In[8]:

for key, value in full_data.iteritems():
    offset = float(value['meta']['elevation'])
    value['data']['elevation'] -= offset

# ## Plot all Water Level data in the NJ area

# In[9]:

fig, ax = plt.subplots(figsize=(16, 3))

fig.suptitle('Water Elevation', fontsize=14)

for key, value in full_data.iteritems():
    try:
        if 'SSS-NJ' in key:
            df = value['data']
            ax.plot(df.index, df['elevation'])
            ax.set_xlabel('Date', fontsize=14)
            ax.set_ylabel('Elevation (ft)', fontsize=14)
    except Exception as e:
        print(e)

# ## Plot all Pressure data in the NJ area

# In[10]:

fig, ax = plt.subplots(figsize=(16, 3))

fig.suptitle('nearest_barometric_sensor_psi', fontsize=14)

for key, value in full_data.iteritems():
    try:
        if 'SSS-NJ' in key:
            df = value['data']
            ax.plot(df.index, df['nearest_barometric_sensor_psi'])
            ax.set_xlabel('Date', fontsize=14)
            ax.set_ylabel('Nearest barometric sensor (psi)', fontsize=14)
    except Exception as e:
        print(e)

# ## Map the Gage locations

# In[11]:

map = folium.Map(width=800, height=500, location=[30, -73], zoom_start=4)

# Generate the color map for the storms.
color_list = {"Tropical Storm": '#4AD200',
              "Category 1 Hurricane": '#CFD900',
              "Category 2 Hurricane": '#E16400',
              "Category 3 Hurricane": '#ff0000'}

# Add the track line.
with open('track.csv', 'rb') as csvfile:
    spamreader = csv.reader(csvfile, delimiter=',', quotechar='|')
    for row in spamreader:
        lon, lat = row[3], row[2]
        popup = "{} : {}
 {}".format(row[0], row[1], row[6])
        map.circle_marker([lat, lon], popup=popup,
                          fill_color=color_list[row[6]],
                          radius=10000, line_color='#000000')

# Add the station.
for st in full_data:
    lat = full_data[st]['meta']['lat']
    lon = full_data[st]['meta']['lon']
    map.simple_marker([lat, lon], popup=st, clustered_marker=True)

map.add_layers_to_map()
inline_map(map)

# ## Generate Plot Of Maximum Water Levels from each gage

# In[12]:

dt, dv = [], []

fig, ax = plt.subplots(figsize=(16, 3))
fig.suptitle('Max Water Level (ft), 2011', fontsize=14)

for key, value in full_data.iteritems():
    df = value['data']
    z = df['elevation'].values

    idx = np.argmax(z)
    val = z[idx]
    t = df.index[idx]

    dt.append(t)
    dv.append(val)

    data_dict = {'elevation': dv,
                 'dates': dt}

    df = pd.DataFrame(data=data_dict, index=dt, columns=['elevation', 'dates'])

    ax.scatter(df.index, df['elevation'])
    ax.set_xlabel('Date', fontsize=14)
    ax.set_ylabel('Water Level (ft)', fontsize=14)

ax.set_ylim(0, 20)
ax.set_xlim(md.date2num(datetime(2011, 8, 25, 20)),
            md.date2num(datetime(2011, 8, 30)))
ax.xaxis.set_major_formatter(md.DateFormatter('%B,%d\n%H:%M'))

# ## Generate Plot of maximum water level and its location

# In[13]:

mpl.rcParams['legend.fontsize'] = 10

x, y, zz, bpz = [], [], [], []

fig, ax = plt.subplots(figsize=(10, 10))

for key, value in full_data.iteritems():
    df = value['data']
    z = df['elevation'].values
    bp = df['nearest_barometric_sensor_psi'].values
    lon = value['meta']['lon']
    lat = value['meta']['lat']

    idx = np.argmax(z)
    bpz.append(bp[idx])
    zz.append(z[idx])
    x.append(lon)
    y.append(lat)

bpz = np.array(bpz) * 10.
pts = ax.scatter(x, y, c=zz, s=bpz)
ax.set_xlabel('Lon')
ax.set_ylabel('Lat')
title = ("Plot Showing Locations of Maximum Water Level\n"
         "Color coded by Maximum water level (ft)\n"
         "Sized by barometric pressure (psi)")
ax.set_title(title)
cb = fig.colorbar(pts)


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

# In[ ]:

import dask
import numpy as np
import numpy.typing as npt
from typing import *
import pandas as pd
import dask.dataframe as dd

# In[ ]:

url = "https://gender-pay-gap.service.gov.uk/viewing/download-data/2021"
#tag::ex_load_1kb[]
many_chunks = dd.read_csv(url, blocksize="1kb")
many_chunks.index
#end::ex_load_1kb[]

# In[ ]:

#tag::ex_load_uk_gender_pay_gap_infered[]
df = dd.read_csv(
    "https://gender-pay-gap.service.gov.uk/viewing/download-data/2021")
#end::ex_load_uk_gender_pay_gap_infered[]

# In[ ]:

# The df.compute() is not needed for the load, but because Dask is lazy we need to trigger a compute
# for Dask to evaluate the DataFrame and notice the error.
try:
    df.compute() # Observe the failure
except Exception as e:
        print(e)
# - CompanyNumber
#  ValueError("invalid literal for int() with base 10: 'SC312912'")
#end::ex_load_uk_gender_pay_gap_infered

# In[ ]:

#tag::ex_load_uk_gender_pay_gap[]
df = dd.read_csv(
    "https://gender-pay-gap.service.gov.uk/viewing/download-data/2021",
    dtype={'CompanyNumber': 'str', 'DiffMeanHourlyPercent': 'float64'})
#end::ex_load_uk_gender_pay_gap[]

# In[ ]:

#tag::csv_gender_pay_gap_with_full_inference[]
df = dd.read_csv(
    "https://gender-pay-gap.service.gov.uk/viewing/download-data/2021",
    sample=256000000000000000000000000000000000000000000) # size in bytes to sample
# One day this should work, but for now it does not and we get the same error as if we had not sampled the entire CSV file
# df.compute()
#end::csv_gender_pay_gap_with_full_inference[]

# In[ ]:

from fsspec.registry import known_implementations
known_implementations

# In[ ]:

#tag::filna_ex[]
def fillna(df):
    return df.fillna(value={"PostCode": "UNKNOWN"}).fillna(value=0)

new_df = df.map_partitions(fillna)
# Since there could be an NA in the index clear the partition / division information
new_df.clear_divisions()
#end::filna_ex[]

# In[ ]:

new_df.compute()

# In[ ]:

narrow_df = new_df[["PostCode", "EmployerSize", "DiffMeanHourlyPercent"]]

# In[ ]:

grouped_df = narrow_df.groupby("PostCode")

# In[ ]:

alt_grouped_df = new_df.groupby(["PostCode", "SicCodes"])
alt_grouped_df.sum().head(2)

# In[ ]:

avg_by_postalcode = grouped_df.mean()

# In[ ]:

avg_by_postalcode.compute()

# In[ ]:

ops_by_postcalcode = narrow_df.set_index("PostCode", npartitions=10)
len(list(ops_by_postcalcode.partitions))

# In[ ]:

# Le sad, you can see this doesn't actually respect the partition size of one byte.
dask.visualize(narrow_df.set_index("PostCode", npartitions="auto", partition_size=1))

# In[ ]:

indexed = narrow_df.set_index("PostCode")
#tag::repartition[]
reparted = indexed.repartition(partition_size="20kb")
#end::repartition[]
dask.visualize(narrow_df.set_index("PostCode").repartition(partition_size="20kb"))

# In[ ]:

dask.visualize(ops_by_postcalcode)

# In[ ]:

fast_grouped_df = ops_by_postcalcode.groupby("PostCode")
fast_grouped_df.mean().compute()

# In[ ]:

# Kind of hacky string munging to get a median-ish to weight our values.
def update_empsize_to_median(df):
    def to_median(value):
        if " to " in value:
            f , t = value.replace(",", "").split(" to ")
            return (int(f) + int(t)) / 2.0
        elif "Less than" in value:
            return 100
        else:
            return 10000
    df["EmployerSize"] = df["EmployerSize"].apply(to_median)
    return df

df_with_median_emp_size = narrow_df.map_partitions(update_empsize_to_median)

# In[ ]:

# In[ ]:

df_with_median_emp_size.head(1)

# In[ ]:

def join_emp_with_diff(df):
    # In practice life would be easier if we multiplied these together but to illustrate
    # the custom aggregate we'll make this a tuple for now
    df["empsize_diff"] = list(df[["EmployerSize", "DiffMeanHourlyPercent"]].to_records(index=False))
    return df
df_diff_with_emp_size = df_with_median_emp_size.map_partitions(
    join_emp_with_diff)
df_diff_with_emp_size.head(1)

# In[ ]:

#tag::custom_agg[]
# Write a custom weighted mean, we get either a DataFrameGroupBy with multiple columns or SeriesGroupBy for each chunk
def process_chunk(chunk):
    def weighted_func(df):
        return (df["EmployerSize"] * df["DiffMeanHourlyPercent"]).sum()
    return (chunk.apply(weighted_func), chunk.sum()["EmployerSize"])

def agg(total, weights):
    return (total.sum(), weights.sum())

def finalize(total, weights):
    return total / weights

weighted_mean = dd.Aggregation(
    name='weighted_mean',
    chunk=process_chunk,
    agg=agg,
    finalize=finalize)

aggregated = df_diff_with_emp_size.groupby("PostCode")["EmployerSize", "DiffMeanHourlyPercent"].agg(weighted_mean)
#end::custom_agg[]
j = aggregated.head(4)
j

# In[ ]:

#tag::custom_agg_hyperloglog[]
# Wrap Dask's hyperloglog in dd.Aggregation

from dask.dataframe import hyperloglog

approx_unique = dd.Aggregation(
    name='aprox_unique',
    chunk=hyperloglog.compute_hll_array,
    agg=hyperloglog.reduce_state,
    finalize=hyperloglog.estimate_count)

aggregated = df_diff_with_emp_size.groupby("PostCode")["EmployerSize", "DiffMeanHourlyPercent"].agg(weighted_mean)
#end::custom_agg_hyperloglog[]
j = aggregated.head(4)
j

# In[ ]:

aggregated = new_df.groupby("PostCode")["EmployerId"].apply(lambda g: list(g))
aggregated.head(4)

# In[ ]:

# For loading data example the note here is that whatever params we pass through read_x
# if not consumed by dask (e.g. blocksize is used by Dask),
# More generally all of Dask's DataFrame functions follow this pattern.
sf_covid_df = dd.read_csv("https://data.sfgov.org/api/views/gqw3-444p/rows.csv?accessType=DOWNLOAD", blocksize=None, dtype={
    'pct_tot_new_cases': 'float64',
    'pct_tot_new_cases_7_day_avg': 'float64',
    'new_case_rate': 'float64',
    'new_case_rate_7_day_avg': 'float64',
    'new_cases_7_day_avg': 'float64'}, parse_dates=['specimen_collection_date'], infer_datetime_format=True)

# In[ ]:

sf_covid_df.columns

# In[ ]:

sf_covid_df.head(10)

# In[ ]:

#tag::compute_entire_max_mean[]
dask.compute(
    sf_covid_df[["new_cases"]].max(),
    sf_covid_df[["new_cases"]].mean()
)
#end::compute_entire_max_mean[]

# In[ ]:

#tag::static_group[]
raw_grouped = sf_covid_df.groupby(lambda x: 0)
#end::static_group[]

#tag::max_mean[]
dask.compute(
    raw_grouped[["new_cases"]].max(),
    raw_grouped[["new_cases"]].mean())
#end::max_mean[]

# In[ ]:

# Drop columns & rows we don't care about before repartitioning
#tag::index_covid_data[]
mini_sf_covid_df = sf_covid_df[sf_covid_df['vaccination_status'] == 'All'][['specimen_collection_date', 'new_cases']]
#end::index_covid_data[]

# In[ ]:

mini_sf_covid_df.index

# In[ ]:

indexed_df = mini_sf_covid_df.set_index('specimen_collection_date', npartitions=5)
indexed_df.head(1)

# In[ ]:

from datetime import datetime

#tag::set_index_with_rolling_window[]
divisions = pd.date_range(start="2021-01-01", end=datetime.today(), freq='7D').tolist()
partitioned_df_as_part_of_set_index = mini_sf_covid_df.set_index(
    'specimen_collection_date', divisions=divisions)
#end::set_index_with_rolling_window[]

# In[ ]:

partitioned_df_as_part_of_set_index.divisions

# In[ ]:

len(list(indexed_df.partitions))

# In[ ]:

# Repartition on 14 day window
partitioned_df = indexed_df.repartition(freq='14D', force=True)

# In[ ]:

indexed_df.divisions

# In[ ]:

partitioned_df.divisions

# In[ ]:

# Rolling average with time delta
#tag::rolling_date_ex[]
rolling_avg = partitioned_df.map_overlap(lambda df: df.rolling('5D').mean(), pd.Timedelta('5D'), 0)
#end::rolling_date_ex[]

# In[ ]:

rolling_avg.compute()

# In[ ]:


# Copyright 2018-2020 Xanadu Quantum Technologies 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.
"""
Unit tests for the :mod:`pennylane.circuit_graph.to_openqasm()` method.
"""
# pylint: disable=no-self-use,too-many-arguments,protected-access
from textwrap import dedent

import numpy as np
import pytest

import pennylane as qml
from pennylane import CircuitGraph
from pennylane.wires import Wires

class TestToQasmUnitTests:
    """Unit tests for the to_openqasm() method"""

    def test_empty_circuit(self):
        """Test that an empty circuit graph is properly
        serialized into an empty QASM program."""
        circuit = CircuitGraph([], {}, Wires([]))
        res = circuit.to_openqasm()
        expected = 'OPENQASM 2.0;\ninclude "qelib1.inc";\n'
        assert res == expected

    def test_native_qasm_gates(self):
        """Test that a circuit containing solely native QASM
        gates is properly serialized."""
        ops = [
            qml.RX(0.43, wires=0),
            qml.RY(0.35, wires=1),
            qml.RZ(0.35, wires=2),
            qml.CNOT(wires=[0, 1]),
            qml.Hadamard(wires=2),
            qml.CNOT(wires=[2, 0]),
            qml.PauliX(wires=1),
        ]

        circuit = CircuitGraph(ops, {}, Wires([0, 1, 2]))
        res = circuit.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[3];
            creg c[3];
            rx(0.43) q[0];
            ry(0.35) q[1];
            rz(0.35) q[2];
            cx q[0],q[1];
            h q[2];
            cx q[2],q[0];
            x q[1];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            """
        )

        assert res == expected

    def test_native_inverse_gates(self):
        """Test that a circuit containing inverse gates that are supported
        natively by QASM, such as sdg, are correctly serialized."""
        ops = [
            qml.S(wires=0),
            qml.S(wires=0).inv(),
            qml.T(wires=0),
            qml.T(wires=0).inv(),
        ]

        circuit = CircuitGraph(ops, {}, Wires([0]))
        res = circuit.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[1];
            creg c[1];
            s q[0];
            sdg q[0];
            t q[0];
            tdg q[0];
            measure q[0] -> c[0];
            """
        )

        assert res == expected

    def test_unused_wires(self):
        """Test that unused wires are correctly taken into account"""
        ops = [
            qml.Hadamard(wires=4),
            qml.CNOT(wires=[1, 0]),
        ]

        circuit = CircuitGraph(ops, {}, Wires([0, 1, 2, 3, 4]))
        res = circuit.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[5];
            creg c[5];
            h q[4];
            cx q[1],q[0];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            measure q[3] -> c[3];
            measure q[4] -> c[4];
            """
        )

        assert res == expected

    def test_rotation_gate_decomposition(self):
        """Test that gates not natively supported by QASM, such as the
        rotation gate, are correctly decomposed and serialized."""
        ops1 = [qml.Rot(0.3, 0.1, 0.2, wires=1)]
        circuit1 = CircuitGraph(ops1, {}, Wires([0, 1]))
        qasm1 = circuit1.to_openqasm()

        ops2 = qml.Rot.decomposition(0.3, 0.1, 0.2, wires=1)
        circuit2 = CircuitGraph(ops2, {}, Wires([0, 1]))
        qasm2 = circuit2.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[2];
            creg c[2];
            rz(0.3) q[1];
            ry(0.1) q[1];
            rz(0.2) q[1];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            """
        )

        assert qasm1 == expected
        assert qasm1 == qasm2

    def test_state_initialization_decomposition(self):
        """Test that the Mottonen state prepration decomposition
        is correctly applied."""
        psi = np.array([1, -1, -1, 1]) / np.sqrt(4)

        ops1 = [qml.QubitStateVector(psi, wires=[0, 1])]
        circuit1 = CircuitGraph(ops1, {}, Wires([0, 1]))
        qasm1 = circuit1.to_openqasm()

        ops2 = qml.QubitStateVector.decomposition(psi, wires=[0, 1])
        circuit2 = CircuitGraph(ops2, {}, Wires([0, 1]))
        qasm2 = circuit2.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[2];
            creg c[2];
            ry(1.5707963267948968) q[0];
            ry(1.5707963267948968) q[1];
            cx q[0],q[1];
            cx q[0],q[1];
            cx q[0],q[1];
            rz(3.141592653589793) q[1];
            cx q[0],q[1];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            """
        )

        assert qasm1 == expected
        assert qasm1 == qasm2

    def test_basis_state_initialization_decomposition(self):
        """Test that the basis state preparation decomposition

        is correctly applied."""
        basis_state = np.array([1, 0, 1, 1])

        ops1 = [qml.BasisState(basis_state, wires=[0, 1, 2, 3])]
        circuit1 = CircuitGraph(ops1, {}, Wires([0, 1, 2, 3]))
        qasm1 = circuit1.to_openqasm()

        ops2 = qml.BasisState.decomposition(basis_state, wires=[0, 1, 2, 3])
        circuit2 = CircuitGraph(ops2, {}, Wires([0, 1, 2, 3]))
        qasm2 = circuit2.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[4];
            creg c[4];
            x q[0];
            x q[2];
            x q[3];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            measure q[3] -> c[3];
            """
        )

        assert qasm1 == expected
        assert qasm1 == qasm2

    def test_unsupported_gate(self):
        """Test an exception is raised if an unsupported operation is
        applied."""
        U = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
        ops = [qml.S(wires=0), qml.QubitUnitary(U, wires=[0, 1])]

        circuit = CircuitGraph(ops, {}, Wires([0, 1]))

        with pytest.raises(
            ValueError, match="QubitUnitary not supported by the QASM serializer"
        ):
            res = circuit.to_openqasm()

    def test_rotations(self):
        """Test that observable rotations are correctly applied."""
        ops = [
            qml.Hadamard(wires=0),
            qml.CNOT(wires=[0, 1]),
            qml.expval(qml.PauliX(0)),
            qml.expval(qml.PauliZ(1)),
            qml.expval(qml.Hadamard(2)),
        ]

        circuit = CircuitGraph(ops, {}, Wires([0, 1, 2]))
        res = circuit.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[3];
            creg c[3];
            h q[0];
            cx q[0],q[1];
            h q[0];
            ry(-0.7853981633974483) q[2];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            """
        )

        assert res == expected

        ops2 = circuit.operations + circuit.diagonalizing_gates
        circuit2 = CircuitGraph(ops2, {}, Wires([0, 1, 2]))
        qasm2 = circuit2.to_openqasm()

        assert res == qasm2

class TestQNodeQasmIntegrationTests:
    """Test that the QASM serialization works correctly
    when circuits are created via QNodes."""

    def test_empty_circuit(self):
        """Test that an empty QNode is properly
        serialized into an empty QASM program."""
        dev = qml.device("default.qubit", wires=1)

        @qml.qnode(dev)
        def qnode():
            return qml.expval(qml.PauliZ(0))

        # construct the qnode circuit
        qnode()

        res = qnode.qtape.graph.to_openqasm()
        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[1];
            creg c[1];
            measure q[0] -> c[0];
            """
        )
        assert res == expected

    def test_native_qasm_gates(self):
        """Test that a QNode containing solely native QASM
        gates is properly serialized."""
        dev = qml.device("default.qubit", wires=3)

        @qml.qnode(dev)
        def qnode():
            qml.RX(0.43, wires=0)
            qml.RY(0.35, wires=1)
            qml.RZ(0.35, wires=2)
            qml.CNOT(wires=[0, 1])
            qml.Hadamard(wires=2)
            qml.CNOT(wires=[2, 0])
            qml.PauliX(wires=1)
            return qml.expval(qml.PauliZ(0))

        # construct the qnode circuit
        qnode()
        res = qnode.qtape.graph.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[3];
            creg c[3];
            rx(0.43) q[0];
            ry(0.35) q[1];
            rz(0.35) q[2];
            cx q[0],q[1];
            h q[2];
            cx q[2],q[0];
            x q[1];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            """
        )

        assert res == expected

    def test_parametrized_native_qasm_gates(self):
        """Test that a QNode containing solely native QASM
        gates, as well as input parameters, is properly serialized.
        In addition, double check the serialization changes as parameters
        are changed."""
        dev = qml.device("default.qubit", wires=3)

        @qml.qnode(dev)
        def qnode(x, y):
            qml.RX(x, wires=0)
            qml.RY(y[0], wires=1)
            qml.RZ(y[1], wires=2)
            qml.CNOT(wires=[0, 1])
            qml.Hadamard(wires=2)
            qml.CNOT(wires=[2, 0])
            qml.PauliX(wires=1)
            return qml.expval(qml.PauliZ(0))

        # execute the QNode with parameters, and serialize
        params = np.array([0.5, [0.2, 0.1]])
        qnode(*params)

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[3];
            creg c[3];
            rx(0.5) q[0];
            ry(0.2) q[1];
            rz(0.1) q[2];
            cx q[0],q[1];
            h q[2];
            cx q[2],q[0];
            x q[1];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            """
        )

        res = qnode.qtape.graph.to_openqasm()
        assert res == expected

        # execute the QNode with new parameters, and serialize again
        params = np.array([0.1, [0.3, 0.2]])
        qnode(*params)

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[3];
            creg c[3];
            rx(0.1) q[0];
            ry(0.3) q[1];
            rz(0.2) q[2];
            cx q[0],q[1];
            h q[2];
            cx q[2],q[0];
            x q[1];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            """
        )

        res = qnode.qtape.graph.to_openqasm()
        assert res == expected

    def test_native_inverse_gates(self):
        """Test that a QNode containing inverse gates that are supported
        natively by QASM, such as sdg, are correctly serialized."""
        dev = qml.device("default.qubit", wires=1)

        @qml.qnode(dev)
        def qnode():
            qml.S(wires=0)
            qml.S(wires=0).inv()
            qml.T(wires=0)
            qml.T(wires=0).inv()
            return qml.expval(qml.PauliZ(0))

        # construct the qnode circuit
        qnode()
        res = qnode.qtape.graph.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[1];
            creg c[1];
            s q[0];
            sdg q[0];
            t q[0];
            tdg q[0];
            measure q[0] -> c[0];
            """
        )

        assert res == expected

    def test_unused_wires(self):
        """Test that unused wires are correctly taken into account"""
        dev = qml.device("default.qubit", wires=5)

        @qml.qnode(dev)
        def qnode():
            qml.Hadamard(wires=4)
            qml.CNOT(wires=[1, 0])
            return qml.expval(qml.PauliZ(0))

        # construct the qnode circuit
        qnode()
        res = qnode.qtape.graph.to_openqasm(wires=dev.wires)

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[5];
            creg c[5];
            h q[4];
            cx q[1],q[0];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            measure q[3] -> c[3];
            measure q[4] -> c[4];
            """
        )

        assert res == expected

    def test_rotation_gate_decomposition(self):
        """Test that gates not natively supported by QASM, such as the
        rotation gate, are correctly decomposed and serialized."""
        dev = qml.device("default.qubit", wires=2)

        @qml.qnode(dev)
        def qnode():
            qml.Rot(0.3, 0.1, 0.2, wires=1)
            return qml.expval(qml.PauliZ(0))

        # construct the qnode circuit
        qnode()
        res = qnode.qtape.graph.to_openqasm(wires=dev.wires)

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[2];
            creg c[2];
            rz(0.3) q[1];
            ry(0.1) q[1];
            rz(0.2) q[1];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            """
        )

        assert res == expected

    def test_state_initialization_decomposition(self):
        """Test that the Mottonen state prepration decomposition
        is correctly applied."""
        dev = qml.device("default.qubit", wires=2)

        @qml.qnode(dev)
        def qnode(state=None):
            qml.QubitStateVector(state, wires=[0, 1])
            return qml.expval(qml.PauliZ(0))

        # construct the qnode circuit
        qnode(state=np.array([1, -1, -1, 1]) / np.sqrt(4))
        res = qnode.qtape.graph.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[2];
            creg c[2];
            ry(1.5707963267948968) q[0];
            ry(1.5707963267948968) q[1];
            cx q[0],q[1];
            cx q[0],q[1];
            cx q[0],q[1];
            rz(3.141592653589793) q[1];
            cx q[0],q[1];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            """
        )

        assert res == expected

    def test_basis_state_initialization_decomposition(self):
        """Test that the basis state prepration decomposition
        is correctly applied."""
        dev = qml.device("default.qubit", wires=4)

        @qml.qnode(dev)
        def qnode(state=None):
            qml.BasisState(state, wires=[0, 1, 2, 3])
            return qml.expval(qml.PauliZ(0))

        # construct the qnode circuit
        qnode(state=np.array([1, 0, 1, 1]))
        res = qnode.qtape.graph.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[4];
            creg c[4];
            x q[0];
            x q[2];
            x q[3];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            measure q[3] -> c[3];
            """
        )

        assert res == expected

    def test_unsupported_gate(self):
        """Test an exception is raised if an unsupported operation is
        applied."""
        U = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
        dev = qml.device("default.qubit", wires=1)

        @qml.qnode(dev)
        def qnode():
            qml.S(wires=0)
            qml.QubitUnitary(U, wires=0)
            return qml.expval(qml.PauliZ(0))

        qnode()

        with pytest.raises(
            ValueError, match="QubitUnitary not supported by the QASM serializer"
        ):
            qnode.qtape.graph.to_openqasm()

    def test_rotations(self):
        """Test that observable rotations are correctly applied."""
        dev = qml.device("default.qubit", wires=3)

        @qml.qnode(dev)
        def qnode():
            qml.Hadamard(wires=0)
            qml.CNOT(wires=[0, 1])
            return [
                qml.expval(qml.PauliX(0)),
                qml.expval(qml.PauliZ(1)),
                qml.expval(qml.Hadamard(2)),
            ]

        qnode()
        res = qnode.qtape.graph.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[3];
            creg c[3];
            h q[0];
            cx q[0],q[1];
            h q[0];
            ry(-0.7853981633974483) q[2];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            """
        )

        assert res == expected

    def test_wires(self):
        """Test that the QASM serializer correctly integrates with the new wires class."""
        dev = qml.device("default.qubit", wires=["a", "b", "c"])

        @qml.qnode(dev)
        def qnode():
            qml.Hadamard(wires="a")
            qml.CNOT(wires=["b", "a"])
            return [
                qml.expval(qml.PauliX("c")),
                qml.expval(qml.PauliZ("a")),
                qml.expval(qml.Hadamard("b")),
            ]

        qnode()
        res = qnode.qtape.graph.to_openqasm()

        expected = dedent(
            """\
            OPENQASM 2.0;
            include "qelib1.inc";
            qreg q[3];
            creg c[3];
            h q[0];
            cx q[1],q[0];
            h q[2];
            ry(-0.7853981633974483) q[1];
            measure q[0] -> c[0];
            measure q[1] -> c[1];
            measure q[2] -> c[2];
            """
        )

        assert res == expected

class TestQASMConformanceTests:
    """Conformance tests to ensure that the CircuitGraph
    serialized QASM conforms to the QASM standard as implemented
    by Qiskit. Note that this test class requires Qiskit and
    PennyLane-Qiskit as a dependency."""

    @pytest.fixture
    def check_dependencies(self):
        self.qiskit = pytest.importorskip("qiskit", minversion="0.14.1")
        pl_qiskit = pytest.importorskip("pennylane_qiskit")

    def test_agrees_qiskit_plugin(self, check_dependencies):
        """Test that the QASM generated by the CircuitGraph agrees
        with the QASM generated by the PennyLane-Qiskit plugin."""
        dev = qml.device("qiskit.basicaer", wires=3)

        @qml.qnode(dev)
        def qnode(x):
            qml.Hadamard(wires=0)
            qml.RY(x[1], wires=0)
            qml.CNOT(wires=[0, 1])
            qml.RX(x[0], wires=1)
            return [
                qml.expval(qml.PauliX(0)),
                qml.expval(qml.PauliZ(1)),
                qml.expval(qml.Hadamard(2)),
            ]

        qnode([0.1, 0.2])
        res = qnode.qtape.graph.to_openqasm()

        # Note: Qiskit hardcodes in pi as a QASM constant.
        # Here, we replace it with its numerical value.
        expected = dev._circuit.qasm().replace("pi/4", str(np.pi / 4))

        assert res == expected

    def test_basis_state_agrees_qiskit_plugin(self, check_dependencies):
        """Test that the basis state prepration QASM agrees
        with that generated by the PennyLane-Qiskit plugin. This is
        a useful test to ensure that we are using the correct qubit
        ordering convention."""
        dev = qml.device("qiskit.basicaer", wires=4)

        @qml.qnode(dev)
        def qnode(state=None):
            qml.BasisState(state, wires=[0, 1, 2, 3])
            return qml.expval(qml.PauliZ(0))

        # construct the qnode circuit
        qnode(state=np.array([1, 0, 1, 1]))
        res = qnode.qtape.graph.to_openqasm(wires=dev.wires)
        expected = dev._circuit.qasm()

        assert res == expected

    def test_qiskit_load_generated_qasm(self, check_dependencies):
        """Test that the QASM generated by the CircuitGraph
        corresponds to valid QASM, that can be loaded by Qiskit."""
        dev = qml.device("default.qubit", wires=3)

        @qml.qnode(dev)
        def qnode(x):
            qml.Hadamard(wires=0)
            qml.RY(x[1], wires=0)
            qml.CNOT(wires=[0, 1])
            qml.RX(x[0], wires=1)
            return [
                qml.expval(qml.PauliX(0)),
                qml.expval(qml.PauliZ(1)),
                qml.expval(qml.Hadamard(2)),
            ]

        params = [0.1, 0.2]
        qnode(params)
        qasm = qnode.qtape.graph.to_openqasm()
        qc = self.qiskit.QuantumCircuit.from_qasm_str(qasm)

        gates = [g for g, _, _ in qc.data]

        for idx, g in enumerate(gates):
            # attach a wires attribute to each gate, containing
            # a list of wire integers it acts on, so we can assert
            # correctness below.
            g.wires = [q.index for q in qc.data[idx][1]]

        # operations
        assert gates[0].name == "h"
        assert gates[0].wires == [0]

        assert gates[1].name == "ry"
        assert gates[1].wires == [0]
        assert gates[1].params == [params[1]]

        assert gates[2].name == "cx"
        assert gates[2].wires == [0, 1]

        assert gates[4].name == "rx"
        assert gates[4].wires == [1]
        assert gates[4].params == [params[0]]

        # rotations
        assert gates[3].name == "h"
        assert gates[3].wires == [0]

        assert gates[5].name == "ry"
        assert gates[5].wires == [2]
        assert gates[5].params == [-np.pi / 4]


# This tests the compilation and execution of the source code generated with
# utilities.codegen. The compilation takes place in a temporary directory that
# is removed after the test. By default the test directory is always removed,
# but this behavior can be changed by setting the environment variable
# SYMPY_TEST_CLEAN_TEMP to:
#   export SYMPY_TEST_CLEAN_TEMP=always   : the default behavior.
#   export SYMPY_TEST_CLEAN_TEMP=success  : only remove the directories of working tests.
#   export SYMPY_TEST_CLEAN_TEMP=never    : never remove the directories with the test code.
# When a directory is not removed, the necessary information is printed on
# screen to find the files that belong to the (failed) tests. If a test does
# not fail, py.test captures all the output and you will not see the directories
# corresponding to the successful tests. Use the --nocapture option to see all
# the output.

# All tests below have a counterpart in utilities/test/test_codegen.py. In the
# latter file, the resulting code is compared with predefined strings, without
# compilation or execution.

# All the generated Fortran code should conform with the Fortran 95 standard,
# and all the generated C code should be ANSI C, which facilitates the
# incorporation in various projects. The tests below assume that the binary cc
# is somewhere in the path and that it can compile ANSI C code.

from __future__ import print_function

from sympy.abc import x, y, z
from sympy.utilities.pytest import skip
from sympy.utilities.codegen import codegen, make_routine, get_code_generator
import sys
import os
import tempfile
import subprocess

# templates for the main program that will test the generated code.

main_template = {}
main_template['F95'] = """
program main
  include "codegen.h"
  integer :: result;
  result = 0

  %(statements)s

  call exit(result)
end program
"""

main_template['C'] = """
#include "codegen.h"
#include <stdio.h>
#include <math.h>

int main() {
  int result = 0;

  %(statements)s

  return result;
}
"""

# templates for the numerical tests

numerical_test_template = {}
numerical_test_template['C'] = """
  if (fabs(%(call)s)>%(threshold)s) {
    printf("Numerical validation failed: %(call)s=%%e threshold=%(threshold)s\\n", %(call)s);
    result = -1;
  }
"""

numerical_test_template['F95'] = """
  if (abs(%(call)s)>%(threshold)s) then
    write(6,"('Numerical validation failed:')")
    write(6,"('%(call)s=',e15.5,'threshold=',e15.5)") %(call)s, %(threshold)s
    result = -1;
  end if
"""
# command sequences for supported compilers

compile_commands = {}
compile_commands['cc'] = [
    "cc -c codegen.c -o codegen.o",
    "cc -c main.c -o main.o",
    "cc main.o codegen.o -lm -o test.exe"
]

compile_commands['gfortran'] = [
    "gfortran -c codegen.f90 -o codegen.o",
    "gfortran -ffree-line-length-none -c main.f90 -o main.o",
    "gfortran main.o codegen.o -o test.exe"
]

compile_commands['g95'] = [
    "g95 -c codegen.f90 -o codegen.o",
    "g95 -ffree-line-length-huge -c main.f90 -o main.o",
    "g95 main.o codegen.o -o test.exe"
]

compile_commands['ifort'] = [
    "ifort -c codegen.f90 -o codegen.o",
    "ifort -c main.f90 -o main.o",
    "ifort main.o codegen.o -o test.exe"
]

combinations_lang_compiler = [
    ('C', 'cc'),
    ('F95', 'ifort'),
    ('F95', 'gfortran'),
    ('F95', 'g95')
]

def try_run(commands):
    """Run a series of commands and only return True if all ran fine."""
    null = open(os.devnull, 'w')
    for command in commands:
        retcode = subprocess.call(command, stdout=null, shell=True,
                stderr=subprocess.STDOUT)
        if retcode != 0:
            return False
    return True

def run_test(label, routines, numerical_tests, language, commands, friendly=True):
    """A driver for the codegen tests.

       This driver assumes that a compiler ifort is present in the PATH and that
       ifort is (at least) a Fortran 90 compiler. The generated code is written in
       a temporary directory, together with a main program that validates the
       generated code. The test passes when the compilation and the validation
       run correctly.
    """

    # Check input arguments before touching the file system
    language = language.upper()
    assert language in main_template
    assert language in numerical_test_template

    # Check that evironment variable makes sense
    clean = os.getenv('SYMPY_TEST_CLEAN_TEMP', 'always').lower()
    if clean not in ('always', 'success', 'never'):
        raise ValueError("SYMPY_TEST_CLEAN_TEMP must be one of the following: 'always', 'success' or 'never'.")

    # Do all the magic to compile, run and validate the test code
    # 1) prepare the temporary working directory, switch to that dir
    work = tempfile.mkdtemp("_sympy_%s_test" % language, "%s_" % label)
    oldwork = os.getcwd()
    os.chdir(work)

    # 2) write the generated code
    if friendly:
        # interpret the routines as a name_expr list and call the friendly
        # function codegen
        codegen(routines, language, "codegen", to_files=True)
    else:
        code_gen = get_code_generator(language, "codegen")
        code_gen.write(routines, "codegen", to_files=True)

    # 3) write a simple main program that links to the generated code, and that
    #    includes the numerical tests
    test_strings = []
    for fn_name, args, expected, threshold in numerical_tests:
        call_string = "%s(%s)-(%s)" % (
            fn_name, ",".join(str(arg) for arg in args), expected)
        if language == "F95":
            call_string = fortranize_double_constants(call_string)
            threshold = fortranize_double_constants(str(threshold))
        test_strings.append(numerical_test_template[language] % {
            "call": call_string,
            "threshold": threshold,
        })

    if language == "F95":
        f_name = "main.f90"
    elif language == "C":
        f_name = "main.c"
    else:
        raise NotImplemented(
            "FIXME: filename extension unknown for language: %s" % language)

    with open(f_name, "w") as f:
        f.write(
            main_template[language] % {'statements': "".join(test_strings)})

    # 4) Compile and link
    compiled = try_run(commands)

    # 5) Run if compiled
    if compiled:
        executed = try_run(["./test.exe"])
    else:
        executed = False

    # 6) Clean up stuff
    if clean == 'always' or (clean == 'success' and compiled and executed):
        def safe_remove(filename):
            if os.path.isfile(filename):
                os.remove(filename)
        safe_remove("codegen.f90")
        safe_remove("codegen.c")
        safe_remove("codegen.h")
        safe_remove("codegen.o")
        safe_remove("main.f90")
        safe_remove("main.c")
        safe_remove("main.o")
        safe_remove("test.exe")
        os.chdir(oldwork)
        os.rmdir(work)
    else:
        print("TEST NOT REMOVED: %s" % work, file=sys.stderr)
        os.chdir(oldwork)

    # 7) Do the assertions in the end
    assert compiled, "failed to compile %s code with:\n%s" % (
        language, "\n".join(commands))
    assert executed, "failed to execute %s code from:\n%s" % (
        language, "\n".join(commands))

def fortranize_double_constants(code_string):
    """
    Replaces every literal float with literal doubles
    """
    import re
    pattern_exp = re.compile('\d+(\.)?\d*[eE]-?\d+')
    pattern_float = re.compile('\d+\.\d*(?!\d*d)')

    def subs_exp(matchobj):
        return re.sub('[eE]', 'd', matchobj.group(0))

    def subs_float(matchobj):
        return "%sd0" % matchobj.group(0)

    code_string = pattern_exp.sub(subs_exp, code_string)
    code_string = pattern_float.sub(subs_float, code_string)

    return code_string

def is_feasible(language, commands):
    # This test should always work, otherwise the compiler is not present.
    routine = make_routine("test", x)
    numerical_tests = [
        ("test", ( 1.0,), 1.0, 1e-15),
        ("test", (-1.0,), -1.0, 1e-15),
    ]
    try:
        run_test("is_feasible", [routine], numerical_tests, language, commands,
                friendly=False)
        return True
    except AssertionError:
        return False

valid_lang_commands = []
invalid_lang_compilers = []
for lang, compiler in combinations_lang_compiler:
    commands = compile_commands[compiler]
    if is_feasible(lang, commands):
        valid_lang_commands.append((lang, commands))
    else:
        invalid_lang_compilers.append((lang, compiler))

# We test all language-compiler combinations, just to report what is skipped

def test_C_cc():
    if ("C", 'cc') in invalid_lang_compilers:
        skip("`cc' command didn't work as expected")

def test_F95_ifort():
    if ("F95", 'ifort') in invalid_lang_compilers:
        skip("`ifort' command didn't work as expected")

def test_F95_gfortran():
    if ("F95", 'gfortran') in invalid_lang_compilers:
        skip("`gfortran' command didn't work as expected")

def test_F95_g95():
    if ("F95", 'g95') in invalid_lang_compilers:
        skip("`g95' command didn't work as expected")

# Here comes the actual tests

def test_basic_codegen():
    numerical_tests = [
        ("test", (1.0, 6.0, 3.0), 21.0, 1e-15),
        ("test", (-1.0, 2.0, -2.5), -2.5, 1e-15),
    ]
    name_expr = [("test", (x + y)*z)]
    for lang, commands in valid_lang_commands:
        run_test("basic_codegen", name_expr, numerical_tests, lang, commands)

def test_intrinsic_math1_codegen():
    # not included: log10
    from sympy import acos, asin, atan, ceiling, cos, cosh, floor, log, ln, \
        sin, sinh, sqrt, tan, tanh, N
    name_expr = [
        ("test_fabs", abs(x)),
        ("test_acos", acos(x)),
        ("test_asin", asin(x)),
        ("test_atan", atan(x)),
        ("test_cos", cos(x)),
        ("test_cosh", cosh(x)),
        ("test_log", log(x)),
        ("test_ln", ln(x)),
        ("test_sin", sin(x)),
        ("test_sinh", sinh(x)),
        ("test_sqrt", sqrt(x)),
        ("test_tan", tan(x)),
        ("test_tanh", tanh(x)),
    ]
    numerical_tests = []
    for name, expr in name_expr:
        for xval in 0.2, 0.5, 0.8:
            expected = N(expr.subs(x, xval))
            numerical_tests.append((name, (xval,), expected, 1e-14))
    for lang, commands in valid_lang_commands:
        if lang == "C":
            name_expr_C = [("test_floor", floor(x)), ("test_ceil", ceiling(x))]
        else:
            name_expr_C = []
        run_test("intrinsic_math1", name_expr + name_expr_C,
                 numerical_tests, lang, commands)

def test_instrinsic_math2_codegen():
    # not included: frexp, ldexp, modf, fmod
    from sympy import atan2, N
    name_expr = [
        ("test_atan2", atan2(x, y)),
        ("test_pow", x**y),
    ]
    numerical_tests = []
    for name, expr in name_expr:
        for xval, yval in (0.2, 1.3), (0.5, -0.2), (0.8, 0.8):
            expected = N(expr.subs(x, xval).subs(y, yval))
            numerical_tests.append((name, (xval, yval), expected, 1e-14))
    for lang, commands in valid_lang_commands:
        run_test("intrinsic_math2", name_expr, numerical_tests, lang, commands)

def test_complicated_codegen():
    from sympy import sin, cos, tan, N
    name_expr = [
        ("test1", ((sin(x) + cos(y) + tan(z))**7).expand()),
        ("test2", cos(cos(cos(cos(cos(cos(cos(cos(x + y + z))))))))),
    ]
    numerical_tests = []
    for name, expr in name_expr:
        for xval, yval, zval in (0.2, 1.3, -0.3), (0.5, -0.2, 0.0), (0.8, 2.1, 0.8):
            expected = N(expr.subs(x, xval).subs(y, yval).subs(z, zval))
            numerical_tests.append((name, (xval, yval, zval), expected, 1e-12))
    for lang, commands in valid_lang_commands:
        run_test(
            "complicated_codegen", name_expr, numerical_tests, lang, commands)


import numpy as np

a = np.arange(10)
print(a)
# [0 1 2 3 4 5 6 7 8 9]

a_0 = a[:6]
print(a_0)
# [0 1 2 3 4 5]

a_1 = a_0.reshape(2, 3)
print(a_1)
# [[0 1 2]
#  [3 4 5]]

print(a_0.base)
# [0 1 2 3 4 5 6 7 8 9]

print(a_1.base)
# [0 1 2 3 4 5 6 7 8 9]

a_copy = a.copy()
print(a_copy)
# [0 1 2 3 4 5 6 7 8 9]

print(a_copy.base)
# None

print(a.base)
# None

print(a_0.base is None)
# False

print(a_copy.base is None)
# True

print(a.base is None)
# True

print(a_0.base is a)
# True

print(a_0.base is a_1.base)
# True


# -*- coding: utf-8 -*-
"""
Created on Mon May 29 14:13:00 2017
@author: azkei
After the Data is in a DataFrame format we can start Manipulating Data.
The three phases of data manipulation are:
    1. Data Preparation
    2. Data Transformation
    3. Data Aggregation

Data preparation procedures:
    1. Loading
    2. Assembling
        a. Merging
        b. Concatenating
        c. Combining
    3. Reshaping (pivoting)
    4. Removing
"""
# Assembling can be done in different ways:
#   Merging: using pandas.merge() function
#   Concatenating: using the pandas.concat()
#   Combining: using pandas.DataFrame.combine_first()

# 1. Merging - SQL join
import numpy as np
import pandas as pd
frame1 = pd.DataFrame({'id':['ball','pencil','pen','mug','ashtray'],
                       'price':[12.33,11.44,33.21,13.23,33.62]})
frame1
frame2 = pd.DataFrame({'id':['pencil','pencil','ball','pen'],
                       'color':['white','red','red','black']})
frame2
# Carry out merge function
pd.merge(frame1,frame2)
# As we can see as a result, the returned DataFrame consists of all the rows that
# have an ID in common between the DataFrame. In addition the common column
# that both have been added.

# Specifying which column to be the base for merging.
frame1 = pd.DataFrame({'id':['ball','pencil','pen','mug','ashtray'],
                       'color':['white','red','red','black','green'],
                       'brand':['OMG','ABC','ABC','POD','POD']})
frame1
frame2 = pd.DataFrame({'id':['pencil','pencil','ball','pen'],
                       'brand':['OMG','POD','ABC','POD']})
frame2
# In this case we have 2 DataFrames having columns with the same name.
# So if we merge, we do not get any results
pd.merge(frame1,frame2)
# So it is necessary to explicitly define the criterion of merging using on option
pd.merge(frame1,frame2,on='id')
pd.merge(frame1,frame2,on='brand')
# The results vary considerably depending on the criteria for merging.
# A potential proble can arise where the column names you want to merge
# do not have the same name
# to fix this we use left_on, right_on
frame2 = frame2.rename(columns={'id':'sid'})
frame2
pd.merge(frame1,frame2,left_on='id',right_on='sid')
# By Default the merge() function perorms an inner join.
# The keys in the result are the result in the intersection.
# Using the how option, you can perform, left join, right join and outer joins
frame2 = frame2.rename(columns={'sid':'id'})
# Default inner join
pd.merge(frame1,frame2,on='id')
# Full outer join
pd.merge(frame1,frame2,on='id',how='outer')
# Left join
pd.merge(frame1,frame2,on='id',how='left')
# Right join
pd.merge(frame1,frame2,on='id',how='right')
# To make a merge of multiple keys, just add a list to the on option
pd.merge(frame1,frame2,on=['id','brand'],how='outer')

# 2. Merging on Index - using index as a criterion for merging
# Using right_index, left_index
pd.merge(frame1,frame2,right_index=True, left_index=True)
# DataFrame objects have a join() function when you want to do merging by index

# This wont work as the index names are the same
frame1.join(frame2)
# rename
frame2 = frame2.rename(columns={'id':'id1','brand':'brand2'})
frame1.join(frame2)
# As you can see there are some values corresponding to a frame that have NaN

pandas3 - Data Manipulation/pandas1 - Data Preparation - Merging.py


# Unit 3 Sprint 1 Module 1

import numpy as np
import pandas as pd

#def null_count(df): Check a dataframe for nulls and return the
#       number of missing values.
#   returns total sum of null values from given DataFrame
def null_count(df):
    return df.isnull().sum().sum()

# Create a Train/Test split function for a dataframe and returns
# both the Training and Testing sets. Frac referes to the precent of
# data you would like to set aside for training.
def train_test_split(df, frac):
    n_df = len(df)
    if n_df == 0:
        raise ValueError('At least one array required as input')

    train = df.head(int(len(df)* frac))
    test = df.tail(int(len(df)* (1-frac)))
    return train, test


import pandas as pd
import numpy as np

"""
    How to create a new dataframe using Pandas

"""
filename_two = "heart.csv"

df = pd.read_csv(filename_two)

# Using pandas.
df = pd.DataFrame(data={"A": [1, 2, 3], "B": ["Sam", "Alex", "John"]})

print(df.memory_usage())


"""
Module file containing functions that allow to reproduce FIG. 2 of the article

    Interaction of "Solitons" in a Collisionless Plasma and the Recurrence of Initial States,
    N. J. Zabusky and M. D. Kruskal,
    Phys. Rev. Lett. 15, 240 (1965)

This module was prepared as a part of the scientific short course

  A brief guide to publication-ready scientific figures using Python's matplotlib

held during the 2020 seminar week of the Ultrafast Laser Laboratory
at Institute of Quantum Optics at Leibniz University Hannover.

Author: O. Melchert
Date: 2020-09-09
"""
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

__author__ = 'Oliver Melchert'
__date__ = '2020-09-09'

def fetch_data(path):
    """fetch data

    Reads in data from file in numpy npz-format

    Args:
      path (str): path to npz-file

    Returns: (x, t, uxt)
      x (1D array): x samples
      t (1D array): t samples
      uxt (2D array): wave profile u(x,t)
    """
    dat = np.load(path)
    return dat['x'], dat['t'], dat['uxt']

def set_style():
    """set figure style

    Function that customizes the default style to be conform with the Physical
    Review style and notation guide [1]. For instructions on how to set the
    default style using style sheets see [2].

    Notes:
    - main font size is chosen as 8pt, matching the fontsize of figure captions
    - fontsize of legends and auxiliary text labels are set to 6pt, exceeding
      the minimally required pointsize of 4.25 pt. (1.5 mm)
    - default rc (rc = "run commands", i.e. startup information) settings are
      changed dynamically
    - the custom font-scheme 'type2' depends on your latex installation and
      is not guaranteed to run on your specific system

    Refs:
      [1] https://journals.aps.org/prl/authors
      [2] https://matplotlib.org/3.3.1/tutorials/introductory/customizing.html
    """

    fig_width_1col = 3.4        # figure width in inch
    fig_width_2col = 7.0        # figure width in inch
    fig_aspect_ratio = 1.45     # width to height aspect ratio
    font_size = 8               # font size in pt
    font_size_small = 6         # font size in pt
    font_scheme = 'type2'       # options:
                                #   None    - default matplotlib fonts
                                #   'type1' - text: Helvetica, math: Computer modern
                                #   'type2' - text: Helvetica, math: Helvetica

    mpl.rcParams['figure.figsize'] = fig_width_1col, fig_aspect_ratio*fig_width_1col
    mpl.rcParams['axes.labelsize'] = font_size
    mpl.rcParams['font.size'] = font_size
    mpl.rcParams['legend.fontsize'] = font_size_small
    mpl.rcParams['xtick.labelsize'] = font_size
    mpl.rcParams['ytick.labelsize'] = font_size
    mpl.rcParams['xtick.direction'] = 'out'
    mpl.rcParams['ytick.direction'] = 'out'
    mpl.rcParams['lines.linewidth'] = 1.0
    mpl.rcParams['axes.linewidth'] =  0.5

    if font_scheme == 'type1':
        mpl.rcParams['text.usetex'] = True
        mpl.rcParams['font.family'] = 'sans-serif'
        mpl.rcParams['font.sans-serif'] = 'Helvetica'
        mpl.rcParams['mathtext.fontset'] = 'cm'

    if font_scheme == 'type2':
        mpl.rcParams['text.usetex'] = True
        mpl.rcParams['text.latex.preamble'] = [
           r'\usepackage{siunitx}',
           r'\sisetup{detect-all}',
           r'\usepackage{helvet}',
           r'\usepackage{sansmath}',
           r'\sansmath'
        ]

def set_circle(ax, x0, y0, label):
    """set circle

    Function that generates a circle with text-label at its center.
    For more options on scatter plots, see [1], for more options on
    setting text, see [2]

    Refs:
      [1] https://matplotlib.org/3.3.1/api/_as_gen/matplotlib.pyplot.scatter.html
      [2] https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.text.html

    Args:
      ax (object): figure part for which the labeled circle is intended
      x0 (float): x-position for center of circle
      y0 (float): y-position for center of circle
      label (str): text that should be displayed within circle
    """
    # -- scatter plot consisting of a single object
    ax.scatter(x0, y0, s=60, linewidth=0.75, facecolors='white', edgecolor='black', zorder=10)
    # -- place text at the center of the scatter plot object
    ax.text(x0, y0, label, backgroundcolor='none', ha='center', va='center', color='black', zorder=11, fontsize=6)

def set_colorbar(fig, img, ax):
    """set colorbar

    Function that generates a custom colorbar. For more options on
    colorbars and tick parameters see Refs. [1,2]

    Refs:
      [1] https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.colorbar.html
      [2] https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.tick_params.html

    Args:
      fig (object): main figure object
      img (object): image that should be described
      ax (object): axes object with position info for colorbar placement
    """
    # -- extract position information for colorbar placement
    refPos = ax.get_position()
    x0, y0, w, h = refPos.x0, refPos.y0, refPos.width, refPos.height
    # -- set new axes as reference for colorbar
    colorbar_axis = fig.add_axes([x0, y0 + 1.015*h, w, 0.0175*h])

    # -- set custom colorbar
    colorbar = fig.colorbar(img,        # image described by colorbar
            cax = colorbar_axis,        # reference axex
            orientation = 'horizontal', # colorbar orientation
            extend = 'both'             # ends with out-of range values
            )
    colorbar.ax.tick_params(
            color = 'k',                # tick color
            labelcolor = 'k',           # label color
            bottom = False,             # no ticks at bottom
            labelbottom = False,        # no labels at bottom
            labeltop = True,            # labels on top
            top = True,                 # ticks on top
            direction = 'out',          # place ticks outside
            length = 3,                 # tick length in pts.
            labelsize = 8.,             # tick font in pts.
            pad = 0.                    # tick-to-label distance in pts.
            )
    colorbar.set_ticks((-2, -1, 0, 1, 2))
    colorbar.ax.set_title(r"Real-valued field $u(x,t)$",
                            fontsize=8., y=2.5)

def set_colorbar_small(fig, img, ax):
    """set small colorbar

    Function that generates a custom colorbar. For more options on
    colorbars and tick parameters see Refs. [1,2]

    Refs:
      [1] https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.colorbar.html
      [2] https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.tick_params.html

    Args:
      fig (object): main figure object
      img (object): image that should be described
      ax (object): axes object with position info for colorbar placement
    """
    # -- extract position information for colorbar placement
    refPos = ax.get_position()
    x0, y0, w, h = refPos.x0, refPos.y0, refPos.width, refPos.height
    # -- set new axes as reference for colorbar
    colorbar_axes = fig.add_axes([x0 + 0.875*w, y0 + 0.83*h, 0.03*w, 0.15*h])

    # -- set custom colorbar
    colorbar = fig.colorbar(img,        # image described by colorbar
            cax = colorbar_axes,        # reference axex
            orientation = 'vertical', # colorbar orientation
            extend = 'both'             # ends with out-of range values
            )
    colorbar.ax.tick_params(
            color = 'k',                # tick color
            labelcolor = 'k',           # label color
            right=False,                # no right ticks
            labelright=False,           # no right labels
            left=True,                  # ticks left
            labelleft=True,             # labels left
            direction = 'out',          # place ticks outside
            length = 3,                 # tick length in pts.
            labelsize = 8.,             # tick font in pts.
            pad = 0.                    # tick-to-label distance in pts.
            )
    colorbar.set_ticks((-2, -1, 0, 1, 2))
    colorbar.ax.set_ylabel(r"Field $u(x,t)$", fontsize = 8.)

    # -- enhance contrast by placeing patch under colorbar
    import matplotlib.patches as mpatches
    rectangle = mpatches.Rectangle(
            (-3., -0.1),            # bottom left corner coords.
             6.65, 1.2,             # width, height
            facecolor = 'white',    # path color
            edgecolor = 'none',     # disable frame
            alpha = 0.7,            # set transparencty
            transform = colorbar_axes.transAxes # use axes coords.
            )
    ax.add_patch(rectangle)

def save_figure(fig_format = None, fig_name = 'test'):
    """ save figure

    Function that saves figure or shows interactive plot

    Note:
    - if no valid option is provided, an interactive plot is shown

    Args:
      fig_format (str): format to save figure in (options: png, pdf, svg)
      fig_name (str): name for figure (default: 'test')
    """
    if fig_format == 'png':
        plt.savefig(fig_name+'.png', format='png', dpi=600)
    elif fig_format == 'pdf':
        plt.savefig(fig_name+'.pdf', format='pdf', dpi=600)
    elif fig_format == 'svg':
        plt.savefig(fig_name+'.svg', format='svg')
    else:
        plt.show()

def generate_figure(x, t, uxt, fig_format=None, fig_name='fig02'):
    """generate figure

    Function generating a figure reproducing FIG. 2 of [1].

    Refs:
      [1] Interaction of "Solitons" in a Collisionless Plasma and the Recurrence of Initial States,
          N. J. Zabusky and M. D. Kruskal,
          Phys. Rev. Lett. 15, 240 (1965)

    Args:
      x (1D array): x samples
      t (1D array): t samples
      uxt (2D array): wave profile u(x,t)
      fig_format (str): format for output figure
                        (choices: png, pdf, svg; default: interactive figure)
      fig_name (str): name for output figure wihtout suffix (default='fig02')
    """

    # (1) SET A STYLE THAT FITS THE TARGET JOURNAL
    set_style()

    # (2) SET FIGURE LAYOUT
    fig = plt.figure()
    plt.subplots_adjust(left = 0.17,    # pos. of subplots left border
                        bottom = 0.065, # pos. of subplots bottom border
                        right = 0.96,   # pos. of subplots right border
                        top = 0.91,     # pos. of subplots top border
                        wspace = 0.3,   # horizontal space between supblots
                        hspace = 0.05   # vertical space between subplots
                        )
    gs00 = GridSpec(nrows = 1, ncols = 6)   # set geometry of subplot grid
    ax01 = fig.add_subplot(gs00[0, 0:5])
    ax02 = fig.add_subplot(gs00[0, 5])

    # (3.1) SUBPLOT 1 - SET AXES CONTENTS
    tR = 30.4/np.pi
    t = t/tR
    img = ax01.pcolorfast(x, t, uxt[:-1,:-1],
                          vmin=-2., vmax=2.,                    # set color range
                          cmap = mpl.cm.get_cmap('coolwarm')    # set colormap
                          )
    set_colorbar(fig, img, ax01)

    # -- add auxiliary horizontal lines
    for t0 in [0.5,1./3,1./4,1./5,1./6]:
        ax01.axhline(t0, color = 'black', dashes = [2,2], linewidth = 1)

    # -- add dashed circles
    ax01.scatter(0.45, 0.5,     # position in data coordinates
                s = 1400,       # symbol size in pts.-squared
                fc = 'None',    # facecolor
                ec = 'black',   # edgecolor
                lw = 1,         # linewidth
                ls = '--',      # linestyle
                zorder = 100    # layering order
                )
    ax01.scatter(1.90, 0.4,  s=600,  lw=1, fc='None', ec='black', ls='--', zorder=100)
    ax01.scatter(1.67, 1./3, s=600,  lw=1, fc='None', ec='black', ls='--', zorder=100)
    ax01.scatter(1.23, 1./4, s=400,  lw=1, fc='None', ec='black', ls='--', zorder=100)
    ax01.scatter(0.98, 1./5, s=400,  lw=1, fc='None', ec='black', ls='--', zorder=100)
    ax01.scatter(0.83, 1./6, s=200,  lw=1, fc='None', ec='black', ls='--', zorder=100)

    # -- add numbered circles
    for idx, x0 in enumerate([0.59, 0.33, 0.1, 1.86, 1.64, 1.43, 1.23, 1.02, 0.82]):
      set_circle(ax01, x0, 0.11, r"$%d$"%(idx+1))

    # (4.1) SUBPLOT 1 - SET AXIS DETAILS
    # -- customize x-axis
    x_lim = (0,2)
    x_ticks = (0,0.6,1.2,1.8)
    ax01.tick_params(axis='x', length=3., pad=2, right=False)
    ax01.set_xlim(x_lim)
    ax01.set_xticks(x_ticks)
    ax01.set_xlabel(r"Normalized distance $x$")

    # -- customize y-axis
    y_lim = (0.1,0.6)
    y_ticks = (0.1,1./6, 1./5,1./4,0.3,1./3,0.4,0.5)
    y_ticklabels = (r'$0.1\,T_{\mathrm{R}}$', r'$T_{\mathrm{R}}/6$',
                    r'$T_{\mathrm{R}}/5$', r'$T_{\mathrm{R}}/4$',
                    r'$0.3\,T_{\mathrm{R}}$', r'$T_{\mathrm{R}}/3$',
                    r'$0.4\,T_{\mathrm{R}}$', r'$0.5\,T_{\mathrm{R}}$')
    ax01.tick_params(axis='y', length=3., pad=2, top=False)
    ax01.set_ylim(y_lim)
    ax01.set_yticks(y_ticks)
    ax01.set_yticklabels(y_ticklabels)
    ax01.set_ylabel(r"Normalized time $t$")

    # (3.2) SUBPLOT 2 - SET AXES CONTENTS

    # -- trace path of 1st soliton in subfigure 1
    x_idx = np.zeros(len(t), dtype=int)
    a = np.zeros(len(t)); a[0]=1
    mask = np.zeros(len(x)); mask[x<0.2]=1; mask[x>1.8]=1
    for i in range(1,len(t)):
       x_idx[i] = np.argmax(uxt[i]*mask)
       a[i] = np.real(uxt[i,x_idx[i]])
       s = np.min( [ np.abs(x_idx[i]-x_idx[i-1]), np.abs(x_idx[i]-x_idx[-1]+len(x))    ]  )
       mask = np.roll(mask, s )

    ax02.plot(a, t, color='black')

    # -- highlight traced path of 1st soliton in subfigure 1
    ax01.plot(x[x_idx][::20],t[::20],color='white', linewidth=0, marker='.', markersize=2)

    # -- add auxiliary horizontal lines
    for t0 in [0.5,1./3,1./4,1./5,1./6]:
        ax02.axhline(y=t0, xmin=0., xmax=a[np.argmin(np.abs(t-t0))]/4.,  color='black', dashes=[2,2], linewidth=1)

    # (4.2) SUBPLOT 2 - SET AXIS DETAILS
    x_lim = (0,4)
    x_ticks = (0,2,4)
    ax02.tick_params(axis='x',length=3.,pad=1,top=False)
    ax02.set_xlim(x_lim)
    ax02.set_xticks(x_ticks)
    ax02.set_xlabel(r'$A_{\mathrm{S1}}$')

    ax02.tick_params(axis='y',length=3.,pad=1,labelleft=False,right=False)
    ax02.set_ylim(y_lim)
    ax02.set_yticks(y_ticks)

    ax02.spines['right'].set_visible(False)
    ax02.spines['top'].set_visible(False)

    ax02.text(1., 0.01, r'Amplitude of soliton no. 1',
            horizontalalignment='center', verticalalignment='bottom',
            rotation='vertical', transform=ax02.transAxes)

    # -- add subfigure labels
    ax01.text(-0.02, 1.01, r'(a)', fontsize=8,
            horizontalalignment='right', verticalalignment='bottom',
            transform=ax01.transAxes)
    ax02.text(0.0, 1.01, r'(b)', fontsize=8,
            horizontalalignment='left', verticalalignment='bottom',
            transform=ax02.transAxes)

    # (6) SAVE FIGURE
    save_figure(fig_format, fig_name)

def main():
    x, t, uxt = fetch_data('KdV_raw_data.npz')
    generate_figure(x, t, uxt, fig_format='png', fig_name='fig02')

if __name__ == '__main__':
    main()