```python
import pandas as pd
import ms3
from  ms3.utils import *
import os
from ms3 import Score
import matplotlib.pyplot as plt
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
from scipy import stats
import pretty_midi
import statistics
from scipy.spatial import distance
import torch.nn
from torch import nn
from torch.optim import SGD
# include directory
home_dir = '/Users/Pushkar/PycharmProjects/Tresillo'
```

# Towards a Rhythmical Definition of the Tressilio beat and a Tracing of it in Popular Music

## 1) Introduction & Research Question

### 1.1) Research question:
"Can we compute to which extent a defined rhythm, which we refer to as Tresillo rhythm, is used in a given pop song and if so can we measure the intensity of Tresillo rhythm use in top-20-billboard songs over the past twenty years?"

**Discussion:**
In our project we would like to discuss the use of a rhythm, which we refer to as 'Tresillo rhythm', in the popular music of the last 20 years. We define this rhythm in our project, given secondary literature, and thus obtain a precise notation and formulization of the Tresillo rhythm. Given this definition we can then compute the similarity between the Tresillo rhythm and the rhythm of a given pop song. Thus, we hope to obtain a similarity coefficient which measures how similar the rythm of a given pop song is to our self defined Tresillo rythm. Given the computed similarity coefficients, we hope to measure the use of the Tresillo rhythm in the top 20 billboard songs of the past 20 years (1999-2019).

### 1.2) Assumptions

- Pop songs found in the billboard charts mostly come with a simple melodic and rythmic structure without a lot of variation troughout the song. We assume that for most pop songs present in the billboard charts, one can identify one dominant rythm per song.
- For our presntation of songs in rythm space, we also assume that the majority of the songs in the billboard charts is in the time signature of 4/4. Our data exploration justifies this assumption.(There are 420 songs with 4/4, 6 with 3/4 and 16 with changing time signature).
- We assume that a time sample longer than 30 seconds is suitable to present the main rythm of a pop song. This assumptions is related to the difficultys of finding puplic available midi data for recent pop music. While we couldn't find a data source providing the midi data of the full songs, we managed to find midi data for shorter samples of the songs.

```python
#Some numbers regarding our song length assumption, not adressed in the former notebooks
list_midis = os.listdir(r"C:/Users/Florian/Documents/GitHub/Tresillo/dataset/project_midi/billboard")
length_midis = []
for el in list_midis:
    midi_data = pretty_midi.PrettyMIDI("C:/Users/Florian/Documents/GitHub/Tresillo/dataset/project_midi/billboard" +'/'+el)
    length_midis.append(midi_data.get_end_time())

average_midi_length = (sum(length_midis)/len(length_midis))
plt.figure(1)
plt.hist(length_midis, 20)
midi_length_standart_dev = statistics.stdev(length_midis)

short_midis = []
for songs in length_midis:
    if songs < 60:
        short_midis.append(songs)

plt.figure(2)
plt.hist(short_midis, 20)

counter_30 = 0
counter_40 = 0
counter_30 = sum (i < 30 for i in length_midis)
counter_40 = sum (i < 40 for i in length_midis)
print("Samples shorter than 30 sec: " , counter_30)
print("Samples shorter 40 sec: "  , counter_40)
print ("Average sample length" , average_midi_length)
print("Standart deveiation sample length" , midi_length_standart_dev)
```

### 1.3) Data Representation
Initially our data is represented in the MIDI file format. The representation of music in the MIDI format has the advantage, that often several voices of different instruments are represented in such files. In contrast to musescore, where often only the voice of one instrument (mostly piano) is notated.
However, to obtain a list of onsets of every musical event in a given song, we have to convert our MIDI (.midi) files to Musescore (.mscx) files.
To convert and further analyze our files, we will use the [ms3](https://pypi.org/project/ms3/) python library.
To convert a directory of .midi files to .mscx files we use following command:

```python
path_midi = '/home/nulpe/Desktop/Tresillo/dataset/project_midi/tresillo/'
target = '/home/nulpe/Desktop/Tresillo/dataset/project_mscx/mscx_tresillos_billboard/'

dir_list = os.listdir(path_midi)

for el in dir_list:
    convert(path_midi+el, target+el[:-4]+'.mscx', MS='musescore3')
```

**TA instructions**: Precision of Research Question

- State the final version of your research question as you understand it now.

-  Make all assumptions and hypotheses explicit.

- In case Milestone 2 did not include the final data representation that you are analyzing, present the representation you finally used and the necessary steps to get to it.

**Thoughts Aurel:**
 - Reasearch Question: Move to a fuzzy definition of Tresillio-ness: E.g.:
    Are pop songs increasingly using a rythm pattern which is similar to a rythm pattern which we reffer to as 'Tresillio rythm'?

 - Assumptions are very important for them. Here we have to note our definition(s) of the Tresillio rythm pattern and how we derive them (incl sources). Furthermore, we have to discuss cases where there are Rythms which are similar to the Tresillio rythm but not equivalent (e.g.: reggaeton) and how we deal with them computationally.

 - Here we have to discuss the conversion of our MIDI files to the musescore3 file format. Furthermore, we have to discuss what the musescore3 format offers us, and why it is the better choice for our analysis.

  **Toughts Florian:**
    -Assumptions:

    -30 second piece of pop song is suitable to identfy the main rythm

    -Main rythm can be identified by counting the onsets

    -the great majority of pop songs comes in 4/4 (not shure if we really need that assumption)

## 2) Methods

### 2.1) Definition of the Tresillo rhythm
We notate a clean version of a rythm wich we from now on use as our defintion for the clean tresillo rythm, in the following context also called vanilla tresillo.

The rhythm pattern consists of a dotted eighth note, followed by a sixteenth note, an eighth rest and an eighth note and is repeated two times in a 4/4 bar. The rytm pattern is beeing used as a own rythm or as a rythmic part of a more complex rythm, for example the "clave" pattern or the ragaetton rythm Floyd, 1999).
By notating the rythm in MuseScore 3 and saving it as a .mscx file, we can use it our data processing pipeline described in the following sections.

### 2.2) Rhythm histograms and vectors

To be able to measure the similarity between rhythm we must have a clear definition and thus following representation of rhythm. In general, one can define rhythm as "a series of onsets and durations of musical events.” (Rohrmeier, 2020). In our specific case however, we are interested in the dominat and repeating rhythm of a given song. Therefore we prefer a narrower definition of rhythm as “repeated recurrences in alternate heavy and light beats” (Chin and Wu, 1992). To furthermore simplify our data, we assume that the main rhythm of a song can be defined by the onsets  of its musical events (notes).
To obtain a representation of the dominate rhythm of a song, we preceed to aggregate all musical onset to one bar. Collapsing all musical onsets to one bar and thus obtaining onset 'histograms' is a common pratice and has been used beside others to analyze western classical music (Palmer and Krumhansl, 1990) and american folk music (Huron and Ommen, 2006).
With the onset histogram of a song we can compile a n dimensional vector for each song, which we refer to as a 'rhythm vector'. Given that the number of songs with meters other than 4/4 is negible, we only consider songs with a 4/4 meter in our analysis.  Given that we only consider songs with 4/4 meters, we obtain for each song a 16 dimensional vector.

### 2.3) Evaluating different similarity measures
In this section we try out different similarity measures to choose the best one. A good similarity measure would have a high similarity for all songs which have a $tresillo^+$ pattern and a low similarity for songs which do not have such a pattern.
To measure this, we use 'Similarity Goodness' $S^*$, which is the ratio of mean similarity in songs which have tresillo and mean similarity of songs which don't.
\begin{equation}
S^* = \frac{\frac{\sum{\text{similarity(Songs with tresillo)}}}{n}}{\frac{\sum{\text{similarity(Songs without tresillo)}}}{m}}
\end{equation}
`n` and `m` denote the number of songs in each of respective categories.

In each of the subsections we calculate relavant statistics by dividing the similarity methods based on the following techniques
* Template based center vs Centroid computed from known tresillo songs
* Cosine similarity vs Euclidean Distance
* Parameterized vs non parameterized distance function

Each section contains the mean similarity/distance and it's standard deviation.

At the end of the section we compare the $S^*$ (Model Goodness) for all of the above subsections.

_NOTE: Both the datasets used to evaluate $S^*$ are from the validation/test set and not used as part of training/finding the centroid._

```python
# Collapse all voices into a single voice by taking the mean over the respective beat.
def collapse_normalize_vectors(df_rythm_vectors):
    rythm_vector_collaped = df_rythm_vectors.groupby(['song_artist']).agg(['sum'])
    rythm_vector_collaped.columns = rythm_vector_collaped.columns.droplevel(1)
    rythm_vector_collaped = rythm_vector_collaped.drop(['instrument', 'level_1'],axis=1)
    rythm_vector_collaped["sum"] = rythm_vector_collaped.sum(axis=1)
    rythm_vector_collaped = rythm_vector_collaped.loc[:,"0":"15"].div(rythm_vector_collaped["sum"], axis=0)
    return rythm_vector_collaped
```

```python
df_tresillo_not_billb = pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_tresillo_not_billboard.csv')
df_not_tre_validation = pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_not_tresillo_validation.csv')
df_rythm_vectors_train = pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_no_tresillos_4_4.csv')
df_rythm_vectors = pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_tresillos_billboard.csv')

tresillo_train = collapse_normalize_vectors(df_tresillo_not_billb)
tresillo_test = collapse_normalize_vectors(df_rythm_vectors)
df_not_tre_validation = collapse_normalize_vectors(df_not_tre_validation)
non_tresillio_vectors = collapse_normalize_vectors(df_rythm_vectors_train)

np_tresillo_train = tresillo_train.to_numpy()
np_tresillo_test = tresillo_test.to_numpy()
np_not_tre_validation = df_not_tre_validation.to_numpy()
np_non_tresillo = non_tresillio_vectors.to_numpy()

df_synt_tresillo = pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_tresillio.csv')
df_synt_tresillo = df_synt_tresillo[df_synt_tresillo['song_artist']!='Raggetone'].drop(['instrument','level_1'], axis=1)
df_synt_tresillo.index = df_synt_tresillo.song_artist
df_synt_tresillo.drop(['song_artist'], axis=1)
df_synt_tresillo["sum"] = df_synt_tresillo.sum(axis=1)
df_synt_tresillo = df_synt_tresillo.loc[:,"0":"15"].div(df_synt_tresillo["sum"], axis=0)
vanilla_tresillo_vector = np.asarray(df_synt_tresillo.loc['Vanilla_Tresillo',"0":"15"]).reshape(-1, 1).T
```

#### 2.3.1) Rhythm similarity measured with cosine simularity, using tresillo template as centre.
Given the 16 dimensional rhythm vectors we obtain following the methode described above, we can now compute simple similarity metrics.
In rhythm analysis a common similarity metric which is used to calculate the similarity between two rhythm vectors is the cosine distance (see: Panteli et al., 2014; Parry and Essa, 2003).  The cosine similarity metric is scale invariante, which as such is interesting for rhythm similarity given that thus only relative frequencies of onsets are important and not absolute frequencies.
The cosine distance between two vectors A and B is defined as following:
\begin{equation}
\cos ({\bf A},{\bf B})= {{\bf A}*{\bf B} \over \|{\bf A}\| \|{\bf B}\|} = \frac{ \sum_{i=1}^{n}{{\bf A}_i{\bf B}_i} }{ \sqrt{\sum_{i=1}^{n}{({\bf A}_i)^2}} \sqrt{\sum_{i=1}^{n}{({\bf B}_i)^2}} }
\end{equation}

Given the definition of the cosine similarity we can now compute the similarity between our self defined Tresillo rhythm and the billboard songs.
First, however we will validate this similarity metric by testing it on our self compiled list of songs which do comprise a Tresillo rhythm and songs which do not comprise a Tresillo rhythm. We then compute the mean 'Tresillo-ness' (similarity to Tresillo rhythm) of both samples. By employing the Bootstrapping method we can also obtain a measurement of uncertainty, as provided by 2.5% and 97.5% confidence intervals.

```python
model_goodness_1 = 1
```

#### 2.3.2) Rhythm similarity measured with inverse euclidean distance, using tresillo template as centre.

In this section the similarity will be calculated by calculating the Euclidean distance between two points. We will go with the convention of "high value -> high similarity" and hence use the inverse distance.

```python
model_goodness_2 = 1
```

#### 2.3.3) Rhythm similarity measured with cosine simularity, using the centroid of tresillo songs as centre.

```python
centroid = tresillo_train.sum(axis=0)/tresillo_train.shape[0]
centroid = np.array(centroid).reshape(1,-1)
```

```python
sim_present_train = cosine_similarity(tresillo_train, centroid)
```

Similarity in test set. Tresillo Present

```python
sim_present = cosine_similarity(np_tresillo_test, centroid)
print(f"Mean Similarity +- Standard Dev: {round(np.mean(sim_present), 4)} +- {round(np.std(sim_present), 4)}")
```

    Mean Similarity +- Standard Dev: 0.9433 +- 0.0236

Similarity for songs without tresillo

```python
sim_not_present = cosine_similarity(np_non_tresillo, np.array(centroid).reshape(1,-1))
print(f"Mean Similarity +- Standard Dev: {round(np.mean(sim_not_present), 4)} +- {round(np.std(sim_not_present), 4)}")
```

    Mean Similarity +- Standard Dev: 0.924 +- 0.025

$S^*$ Model Goodness

```python
model_goodness_3 = round(np.mean(sim_present) / np.mean(sim_not_present), 5)
model_goodness_3
```

    1.02092

#### 2.3.4) Rhythm similarity measured with inverse euclidean distance, using the centroid of tresillo songs as centre.

Training Tresillo songs mean distance from the centroid

```python
dst = [distance.euclidean(point, centroid) for point in np_tresillo_train]
print(f"Mean Similarity +- Standard Dev: {round(np.mean(dst), 4)} +- {round(np.std(dst), 4)}")
```

    Mean Similarity +- Standard Dev: 0.1123 +- 0.0362

Test Tresillo songs mean distance from the centroid

```python
dst_test = [distance.euclidean(point, centroid) for point in np_tresillo_test]
print(f"Mean Similarity +- Standard Dev: {round(np.mean(dst_test), 4)} +- {round(np.std(dst_test), 4)}")
```

    Mean Similarity +- Standard Dev: 0.1116 +- 0.028

Not Tresillo songs mean distance from the centroid

```python
dst_test_nt_tresillo = [distance.euclidean(point, centroid) for point in np_non_tresillo]
print(f"Mean Similarity +- Standard Dev: {round(np.mean(dst_test_nt_tresillo), 4)} +- {round(np.std(dst_test_nt_tresillo), 4)}")
```

    Mean Similarity +- Standard Dev: 0.1384 +- 0.0266

$S^*$ Model Goodness for inverse distance

```python
# Since bigger distance implies low similarity
model_goodness_4 = round(1 / (np.mean(dst_test) / np.mean(dst_test_nt_tresillo)), 4)
model_goodness_4
```

    1.2392

#### 2.3.5) Rhythm similarity measured with parameterized cosine simularity, using the centroid of tresillo songs as centre.

```python
class ParameterizedDistance(nn.Module):
    def __init__(self, theta, device, distance_function):
        super().__init__()
        self.device = device
        if len(theta.shape) == 1:
            theta = theta.reshape(1, -1)
        self.theta = torch.nn.Parameter(torch.from_numpy(theta), requires_grad=True).to(device)
        self.distance_function = distance_function

    def forward(self, tresillo_vectors, not_tresillo_vectors, vanila_tresillo_vector):
        assert isinstance(tresillo_vectors, np.ndarray)
        assert isinstance(not_tresillo_vectors, np.ndarray)
        assert isinstance(vanila_tresillo_vector, np.ndarray)
        assert self.theta.shape[1] == tresillo_vectors.shape[1] == not_tresillo_vectors.shape[1] == \
               vanila_tresillo_vector.shape[1]

        not_tresillo_vectors = torch.from_numpy(not_tresillo_vectors).to(self.device)
        tresillo_vectors = torch.from_numpy(tresillo_vectors).to(self.device)
        vanila_tresillo_vector = torch.from_numpy(vanila_tresillo_vector).to(self.device)

        parameterized_vector_not_tresillo = self.theta * not_tresillo_vectors
        parameterized_vector_tresillo = self.theta * tresillo_vectors
        parameterized_vector_vanilla = self.theta * vanila_tresillo_vector

        cosine_similarity_not_t = torch.mean(
            self.distance_function(parameterized_vector_not_tresillo, parameterized_vector_vanilla))
        cosine_similarity_t = torch.mean(
            self.distance_function(parameterized_vector_tresillo, parameterized_vector_vanilla))

        assert cosine_similarity_t.cpu().detach().numpy() != 0, "0 Similarity between Tresillo set and Vanilla-Tresillo Beat"

        return cosine_similarity_not_t / cosine_similarity_t

    def similarity(self, x, y):
        assert isinstance(x, np.ndarray)
        assert isinstance(y, np.ndarray)
        assert self.theta.shape[1] == x.shape[1] == y.shape[1]

        x = torch.from_numpy(x).to(self.device)
        y = torch.from_numpy(y).to(self.device)

        x = self.theta * x
        y = self.theta * y
        return self.distance_function(x, y)
```

```python
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ParameterizedDistance(np.random.rand(1, 16), device, torch.cosine_similarity)
optim = SGD(model.parameters(), lr=1e-2, weight_decay=1e-4)
optim.zero_grad()
model.train()
prev_validation_ratio = float('inf')
for i in range(100):
    loss = model(np_tresillo_train, np_non_tresillo, centroid)
    loss.backward()
    optim.step()
    with torch.no_grad():
        validation_ratio = model(np_tresillo_test, np_not_tre_validation, centroid)
        if validation_ratio > prev_validation_ratio:
            break
        prev_validation_ratio = validation_ratio
```

```python
model_goodness_5 = 1/validation_ratio
model_goodness_5 = model_goodness_5.cpu().detach().item()
model_goodness_5
```

    1.040989750239579

```python
similarity = model.similarity(np_tresillo_test, vanilla_tresillo_vector).cpu().detach().numpy()
print(f"Mean Similarity +- Standard Dev: {round(np.mean(similarity), 4)} +- {round(np.std(similarity), 4)}")
```

    Mean Similarity +- Standard Dev: 0.7023 +- 0.182

```python
similarity_not_tre = model.similarity(np_not_tre_validation, vanilla_tresillo_vector).cpu().detach().numpy()
print(f"Mean Similarity +- Standard Dev: {round(np.mean(similarity_not_tre), 4)} +- {round(np.std(similarity_not_tre), 4)}")
```

    Mean Similarity +- Standard Dev: 0.4663 +- 0.0735

```python
model.theta.cpu().detach().numpy()
```

    array([[0.72976688, 0.69940337, 0.65299366, 0.92980249, 0.62651748,
            0.77724899, 0.03002122, 0.50504894, 0.15016441, 0.72617848,
            0.66931015, 0.43817627, 0.8624732 , 0.53485404, 0.09186637,
            0.22150938]])

#### 2.3.6) Rhythm similarity measured with parameterized cosine simularity, using tresillo template as centre.

```python
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ParameterizedDistance(np.random.rand(1, 16), device, torch.cosine_similarity)
optim = SGD(model.parameters(), lr=1e-3, weight_decay=1e-4)
optim.zero_grad()
model.train()
prev_validation_ratio = float('inf')
for i in range(200):
    loss = model(np_tresillo_train, np_non_tresillo, vanilla_tresillo_vector)
    loss.backward()
    optim.step()
    with torch.no_grad():
        validation_ratio = model(np_tresillo_test, np_not_tre_validation, vanilla_tresillo_vector)
        if validation_ratio > prev_validation_ratio:
            break
        prev_validation_ratio = validation_ratio
```

```python
model_goodness_6 = 1/validation_ratio
model_goodness_6 = model_goodness_6.cpu().detach().item()
model_goodness_6
```

    2.4386116929263775

```python
similarity = model.similarity(np_tresillo_test, vanilla_tresillo_vector).cpu().detach().numpy()
print(f"Mean Similarity +- Standard Dev: {round(np.mean(similarity), 4)} +- {round(np.std(similarity), 4)}")
```

    Mean Similarity +- Standard Dev: 0.5919 +- 0.2708

```python
similarity_not_tre = model.similarity(np_not_tre_validation, vanilla_tresillo_vector).cpu().detach().numpy()
print(f"Mean Similarity +- Standard Dev: {round(np.mean(similarity_not_tre), 4)} +- {round(np.std(similarity_not_tre), 4)}")
```

    Mean Similarity +- Standard Dev: 0.2427 +- 0.1563

```python
model.theta.cpu().detach().numpy()
```

    array([[-0.12385222,  0.67118366,  0.64942834,  1.06137132,  1.18011032,
             0.10959411,  0.0499411 ,  0.4357717 , -0.12146487,  0.68657846,
             0.58290389,  1.05967711,  1.0022932 ,  0.85658089,  0.12920415,
             0.37872593]])

The parameters of the similarity function hint that it is better to not consider the value of 8th beat. It gives high importance to 3rd, 5th, 11th and 12th beat out of which 3rd and 11th are part of the tresillo peaks.

```python
import plotly.express as px
df_goodness = pd.DataFrame(columns=["Cosine with template", "Euclidean with template", "Cosine with centroid", "Euclidean with template",
                                    "Parameterized with template", "Parameterized with centroid"])
df_goodness.append({"Cosine with template": model_goodness_1, "Euclidean with template": model_goodness_2,
                    "Cosine with centroid": model_goodness_3, "Euclidean with template": model_goodness_4,
                    "Parameterized with template": model_goodness_5, "Euclidean with template": model_goodness_6}, ignore_index=True)
fig = px.bar(data_canada, x='year', y='pop')
fig.show()
```

```python

```

### 2.4) Reducing noise: (pushkar)

**Thoughts Pushkar**:
* I don't think this is needed if the results from the previous section are reasonable.
* This was a reserved approach in mind incase the parameterized similarity was still not good enough.
* As per our research question, Clustering is not `required`, as we have already created a method with reasonable confidence, which does the task. We can have a section at the end, if needed, in which we can do such exploratory analysis.
* Many of the points in this section should be covered/answered from the previous section.

**TA instructions**:
- How did you deal with the problems you mentioned in Milestone 2?

- Which methods did you use to obtain the final results? Provide self-contained explanations and make sure to cite relevant literature where appropriate.

- Explain your core calculations using equations.

- Do not describe all the methods you tried out but only those that lead to the results: the final analysis is not an exploratory analysis anymore.

- Specify any adjustments you made to pre-existing methods

**Thoughts Aurel**:
 -  Talk about how we got to the bar representation of our music. Furthermore, also discuss how we get to a 'perfect Tresillio Histogram'

 - Here I propose we try out several things and compare the results of several methods:

  a) A first big topics is how we define the perfect Tresillo:
   1. Given predefined rythm patterns (by Florian)
   2. Given songs with high Tresillio-ness
    - All instruments collapsed
    - Only key instruments
    - Certain instruments

  a) The second big question is how do we measure Tresillio-ness in the pop songs, I would suggest three approaches. All aproaches require that we first obtain 16 dimensional Rythm vectors of each song of interest:
   1. Compare our vanilla self defined Tresillio rythm vector with all vectors of our pop songs. Measure distance or simularity with some commonly used distance measure in the literature
   2. Very similar to 1) but this time use the Tresillio rythm vector as defined by our songs
   3. Prior clustering of the rythm vectors. Obtaining centroid and measuring with it Tresillio ness in the charts (method as proposed by Pushkar)

- Equations should be included in the prior part
- Discuss critically any outliers, problems and limitations of our methodology. Extra focus on the question how we deal with related but not the same rythm (e.g.: Reggaeton)

**Toughts Florian**
- Following the last paper discussion one measuurement we could use is the Cosin distance
- Not sure about an instrument selection. We can mention that a reduction seems not favorable as the tresillo rythm  is presented with different instruments throughout the dataset

## 3) Final Results

### 3.1) Onset histograms and rhythm vectors
In this first part we will use onset histograms to compute rhythm vectors.
To obtain the onset histogram of a given song, we use the notes representation provided by the [ms3](https://pypi.org/project/ms3/) libary and colapse all musical onsets to one bar. In the example below we will compile the histograms for our self defined 'Vanilla Tresillo' and for the example song 'shape of you' by Ed Sheran. Then we will proceed to compute the rhythm vectors for both songs

```python
#paths to both examples
shape_of_you = home_dir+'dataset/project_mscx/mscx_tresillos_billboard/Shape of you-Ed Sheran.mscx'
vanilla_tresillo = home_dir+'dataset/project_mscx/mscx_tresillos/Vanilla_Tresillo.mscx'

# get the note scores of both examples
df_shape_of_you = Score(shape_of_you).mscx.notes
df_vanilla_tresillo = Score(vanilla_tresillo).mscx.notes

#calculate quarter note position for each note
df_shape_of_you['quarter_beats'] = (df_shape_of_you.mc_onset*16).astype('int32')
df_vanilla_tresillo['quarter_beats'] = (df_vanilla_tresillo.mc_onset*16).astype('int32')

fig, ax = plt.subplots(1,2, figsize=(12,3))

ax[0].hist(df_shape_of_you['quarter_beats'], bins=16)
ax[1].hist(df_vanilla_tresillo['quarter_beats'], bins=16)
ax[0].set_xlabel('quarter_beats')
ax[1].set_xlabel('quarter_beats')
ax[0].xaxis.set_ticks(np.arange(0, 16, 1))
ax[1].xaxis.set_ticks(np.arange(0, 16, 1))
ax[0].set_ylabel('count')
ax[0].set_title('Shape of you')
ax[1].set_title('Vanilla Tresillo')

#mention somewhere that we are working with quarter notes
```

In a next step we want to compile rhythm of each song given this notion of histograms. E.g: every dimension incorporates the absolute frequency of onsets on one given quarter note. This is done as follows:

```python
rhythm_vector_shape_you = df_shape_of_you.groupby(['quarter_beats'])['mn'].agg(['count'])
rhythm_vector_shape_you  = rhythm_vector_shape_you.reindex(list(range(0,16)),fill_value=0).T
rhythm_vector_shape_you
```

<div>
<style scoped>
    .dataframe tbody tr th:only-of-type {
        vertical-align: middle;
    }

    .dataframe tbody tr th {
        vertical-align: top;
    }

    .dataframe thead th {
        text-align: right;
    }
</style>
<table border="1" class="dataframe">
  <thead>
    <tr style="text-align: right;">
      <th>quarter_beats</th>
      <th>0</th>
      <th>1</th>
      <th>2</th>
      <th>3</th>
      <th>4</th>
      <th>5</th>
      <th>6</th>
      <th>7</th>
      <th>8</th>
      <th>9</th>
      <th>10</th>
      <th>11</th>
      <th>12</th>
      <th>13</th>
      <th>14</th>
      <th>15</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <th>count</th>
      <td>206</td>
      <td>33</td>
      <td>44</td>
      <td>119</td>
      <td>72</td>
      <td>37</td>
      <td>169</td>
      <td>12</td>
      <td>200</td>
      <td>29</td>
      <td>75</td>
      <td>112</td>
      <td>75</td>
      <td>27</td>
      <td>167</td>
      <td>10</td>
    </tr>
  </tbody>
</table>
</div>

In the assumption that the rhythm vectors of distinct voices within a song might include information we want to preserve, we compiled one rhythm vector per instrument in a song as follows:

```python
# Define instruments
shape_of_you_score = Score(shape_of_you)
instrument_dict = {}
for key in shape_of_you_score.mscx.metadata['parts']:
    for staff in shape_of_you_score.mscx.metadata['parts'][key].keys():
            instrument_dict[staff] = key

#staff to voice/instruments
df_shape_of_you['instrument'] = [instrument_dict[el] if el in instrument_dict else 'na' for el in df_shape_of_you.staff]

#compute rhythm vectors per voice
rhythm_vector_shape_you_instruments = df_shape_of_you.groupby(['instrument','quarter_beats'])['mn'].agg(['count'])
rhythm_vector_shape_you_instruments = rhythm_vector_shape_you_instruments.groupby(level=0).apply(lambda x: x.reset_index(level = 0).drop(['instrument'],axis=1).reindex(list(range(0,16)),fill_value=0).T)
rhythm_vector_shape_you_instruments = rhythm_vector_shape_you_instruments.reset_index()
rhythm_vector_shape_you_instruments
```

<div>
<style scoped>
    .dataframe tbody tr th:only-of-type {
        vertical-align: middle;
    }

    .dataframe tbody tr th {
        vertical-align: top;
    }

    .dataframe thead th {
        text-align: right;
    }
</style>
<table border="1" class="dataframe">
  <thead>
    <tr style="text-align: right;">
      <th>quarter_beats</th>
      <th>instrument</th>
      <th>level_1</th>
      <th>0</th>
      <th>1</th>
      <th>2</th>
      <th>3</th>
      <th>4</th>
      <th>5</th>
      <th>6</th>
      <th>7</th>
      <th>8</th>
      <th>9</th>
      <th>10</th>
      <th>11</th>
      <th>12</th>
      <th>13</th>
      <th>14</th>
      <th>15</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <th>0</th>
      <td>Grand Piano</td>
      <td>count</td>
      <td>66</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
      <td>66</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
    </tr>
    <tr>
      <th>1</th>
      <td>Marimba, untitled</td>
      <td>count</td>
      <td>66</td>
      <td>0</td>
      <td>0</td>
      <td>53</td>
      <td>14</td>
      <td>0</td>
      <td>63</td>
      <td>0</td>
      <td>68</td>
      <td>0</td>
      <td>0</td>
      <td>57</td>
      <td>17</td>
      <td>0</td>
      <td>63</td>
      <td>0</td>
    </tr>
    <tr>
      <th>2</th>
      <td>Melodic Drum</td>
      <td>count</td>
      <td>0</td>
      <td>0</td>
      <td>33</td>
      <td>36</td>
      <td>13</td>
      <td>19</td>
      <td>39</td>
      <td>0</td>
      <td>0</td>
      <td>0</td>
      <td>50</td>
      <td>28</td>
      <td>10</td>
      <td>19</td>
      <td>38</td>
      <td>0</td>
    </tr>
    <tr>
      <th>3</th>
      <td>Overdrive Gtr</td>
      <td>count</td>
      <td>6</td>
      <td>2</td>
      <td>1</td>
      <td>1</td>
      <td>4</td>
      <td>4</td>
      <td>6</td>
      <td>2</td>
      <td>6</td>
      <td>1</td>
      <td>6</td>
      <td>2</td>
      <td>6</td>
      <td>1</td>
      <td>6</td>
      <td>0</td>
    </tr>
    <tr>
      <th>4</th>
      <td>Percussion</td>
      <td>count</td>
      <td>24</td>
      <td>23</td>
      <td>0</td>
      <td>2</td>
      <td>19</td>
      <td>2</td>
      <td>22</td>
      <td>2</td>
      <td>21</td>
      <td>21</td>
      <td>1</td>
      <td>2</td>
      <td>19</td>
      <td>2</td>
      <td>23</td>
      <td>2</td>
    </tr>
    <tr>
      <th>5</th>
      <td>Tenor Sax</td>
      <td>count</td>
      <td>20</td>
      <td>8</td>
      <td>10</td>
      <td>8</td>
      <td>16</td>
      <td>12</td>
      <td>16</td>
      <td>8</td>
      <td>16</td>
      <td>7</td>
      <td>18</td>
      <td>6</td>
      <td>15</td>
      <td>5</td>
      <td>14</td>
      <td>8</td>
    </tr>
    <tr>
      <th>6</th>
      <td>Woodblock</td>
      <td>count</td>
      <td>2</td>
      <td>0</td>
      <td>0</td>
      <td>2</td>
      <td>0</td>
      <td>0</td>
      <td>2</td>
      <td>0</td>
      <td>2</td>
      <td>0</td>
      <td>0</td>
      <td>1</td>
      <td>2</td>
      <td>0</td>
      <td>2</td>
      <td>0</td>
    </tr>
    <tr>
      <th>7</th>
      <td>Xylophone</td>
      <td>count</td>
      <td>22</td>
      <td>0</td>
      <td>0</td>
      <td>17</td>
      <td>6</td>
      <td>0</td>
      <td>21</td>
      <td>0</td>
      <td>21</td>
      <td>0</td>
      <td>0</td>
      <td>16</td>
      <td>6</td>
      <td>0</td>
      <td>21</td>
      <td>0</td>
    </tr>
  </tbody>
</table>
</div>

If we want to compile rhythm vectors (per voice) for all mscx files in one directory we can use following loop:

```python
def rythm_vectors(in_dir, out_dir):
    list_sheet_music = os.listdir(in_dir)
    df_rythm_vectors =[]

    for idx, el in enumerate(list_sheet_music):
        if el[-4:] == 'mscx':

            #Get notes with onsets
            s = Score(dir_sheet_music+el)
            df = s.mscx.notes

            # Define instruments
            instrument_dict = {}
            for key in s.mscx.metadata['parts']:
                for staff in s.mscx.metadata['parts'][key].keys():
                    instrument_dict[staff] = key

            #staff to instruments
            df['instrument'] = [instrument_dict[el] if el in instrument_dict else 'na' for el in df.staff]

            # define quarter beat
            df['quarter_beats'] = (df.mc_onset*16).astype('int32')

            #make rythm matrix & data frame
            df_histogram = df.groupby(['instrument','quarter_beats'])['mn'].agg(['count'])
            df_histogram = df_histogram.groupby(level=0).apply(lambda x: x.reset_index(level = 0).drop(['instrument'],axis=1).reindex(list(range(0,16)),fill_value=0).T)
            df_histogram = df_histogram.reset_index()

            df_histogram.insert(loc=0, column='song_artist', value=el[:-5])

            #concat to big rythm vector df
            if len(df_rythm_vectors) == 0: df_rythm_vectors = df_histogram

            df_rythm_vectors = pd.concat([df_rythm_vectors,df_histogram], axis=0)

    df_rythm_vectors.to_csv(out_dir, index = False)

dir_sheet_music = home_dir + '/dataset/project_mscx/mscx_billboard/'
out_dir = home_dir + '/dataset/rythm_vectors/rythm_vectors_billboard.csv'
rythm_vectors(dir_sheet_music, out_dir)
```

### 3.2) Tresilio-ness with Cosine Similarity
In a first naive analysis we will employ the cosine similarity measurement to assess if the vector of a given song is similar to our defined Tresillo rhythm. We will need some helper functions.
Following two helper functions will help us to collapse all rhythms of all instruments to one rhythm vector per song.
The second function calculates the cosine similarity between a pandas data frame of rhythm vectors and one single rhythm vector.
The third function calculates the 2.5% and the 97.5% confidence intervals of the distribution of a mean of a given data set. This function will allow us to assess, how big the uncerntaingty is in our data set and if the means of two distributions are indeed significantly different.

```python
def calc_cosine_sim(rythm_vectors, tresillo_vector):
    rythm_vectors['cosine_sim_tresillo'] = cosine_similarity(rythm_vectors.loc[:,"0":"15"],tresillo_vector)
    return rythm_vectors

def bootstrap_CI(data, nbr_draws):
    means = np.zeros(nbr_draws)
    data = np.array(data)

    for n in range(nbr_draws):
        indices = np.random.randint(0, len(data), len(data))
        data_tmp = data[indices]
        means[n] = np.nanmean(data_tmp)
    return [np.nanpercentile(means, 2.5),np.nanpercentile(means, 97.5)]
```

Now lets read in the rhythm vector of our vanilla Tresillo.

```python
df_synt_tresillo = pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_tresillio.csv') # read in all defined 'tresillos' and variations
df_vanilla_tresillo = df_synt_tresillo[df_synt_tresillo['song_artist']=='Vanilla_Tresillo'].loc[: ,"0":"15"] # only use the 16d vector of our vanilla_tresillo
vector_vanilla_tresillo = np.asarray(df_vanilla_tresillo)
vector_vanilla_tresillo
```

    array([[20,  0,  0, 20,  0,  0, 20,  0, 20,  0,  0, 20,  0,  0, 20,  0]])

We will now calculate the cosine similarity between our vanilla Tresillo vector and between a hand selected set of Tresillo songs.
This set of songs has been selected by us ourself and none of those 'validation' songs  are included in the billboard data sets.

```python

tresillo_test_set_vectors = pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_tresillo_not_billboard.csv')
#tresillo_test_set_vectors = pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_tresillos_billboard.csv')
tresillo_test_set_vectors = collapse_normalize_vectors(tresillo_test_set_vectors)
tresillo_test_set_vectors = calc_cosine_sim(tresillo_test_set_vectors, vector_vanilla_tresillo)
print(tresillo_test_set_vectors['cosine_sim_tresillo'])
```

    song_artist
    Attention Charlie Puth                     0.613390
    Chandelier Sia                             0.605698
    Cold Water Major Lazer                     0.896750
    Hips Don't Lie Shakira                     0.819003
    I Don't Care Ed Sheeran & Justin Bieber    0.876542
    One Dance - Drake & Wizkid, Kyla Reid      0.717454
    Rockabye Baby Clean Bandit                 0.880677
    Sorry Justin Bieber                        0.853424
    Titanium David Guetta                      0.652046
    Name: cosine_sim_tresillo, dtype: float64

Let's calculate the mean Tresillo-ness and also the 2.5% and 97.5% confidence intervals of the mean as obtained by bootstrapping.

```python
print('mean Tresillo-ness in the test set: ', tresillo_test_set_vectors['cosine_sim_tresillo'].mean())
lower_tresillo_ci, upper_tresillo_ci = bootstrap_CI(tresillo_test_set_vectors['cosine_sim_tresillo'], 100)
print('tresillo upper and lower ci on 100 draws: ', lower_tresillo_ci, upper_tresillo_ci)
```

    mean Tresillo-ness in the test set:  0.7683315382006799
    tresillo upper and lower ci on 100 draws:  0.6808770780048731 0.8481138158323536

Let's do exact same thing, but with a comparison data set of songs which we kno to not include any Tresillo songs.
Also calculating the mean and the confidence intervals of the not Tresillo songs.

```python
non_tresillio_vectors =  pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_not_tresillo_validation.csv')
#non_tresillio_vectors = pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_no_tresillos_4_4.csv')

non_tresillio_vectors = collapse_normalize_vectors(non_tresillio_vectors)
non_tresillio_vectors = calc_cosine_sim(non_tresillio_vectors, vector_vanilla_tresillo)
print(non_tresillio_vectors['cosine_sim_tresillo'])

print('mean Tresillo-ness in the test set: ', non_tresillio_vectors['cosine_sim_tresillo'].mean())
lower_non_tresillo_ci, upper_non_tresillo_ci = bootstrap_CI(non_tresillio_vectors['cosine_sim_tresillo'], 100)
print('tresillo upper and lower ci on 100 draws: ', lower_non_tresillo_ci, upper_non_tresillo_ci)
```

    song_artist
    Dress You Up           0.666450
    Fear of the Dark       0.689443
    Folsom Prison Blues    0.657568
    Immigrant Song         0.584275
    Jailhouse Rock         0.593940
    One Step Closer        0.523399
    Sweet Child O'Mine     0.724700
    Sweet Home Alabama     0.534338
    Name: cosine_sim_tresillo, dtype: float64
    mean Tresillo-ness in the test set:  0.6217641548092555
    tresillo upper and lower ci on 100 draws:  0.5792063648094466 0.6717705312519476

Already looking at the means of the two samples and their confidence intervalls, they seem to be significantly different.
However let's also calculate the t-test statistics to ensure that the distributions are actually different.

```python
ttest_pvalue = stats.ttest_ind(tresillo_test_set_vectors['cosine_sim_tresillo'], non_tresillio_vectors['cosine_sim_tresillo']).pvalue
print('p value that the means are the same: ', ttest_pvalue)
```

    p value that the means are the same:  0.009618772484088895

**Discussion:**
Looking at the the Tresillio cosine similarity metrics of both data sets and comparing the distributions of the means whit each others (confidence intervals and t-test), we see that their mean Tresillo-ness is indeed significantly diffferent.
However, we also see that 1) this Tresillo-ness measurement seems to be quite noisy 2) The distribution of the two samples are not as far apart as we would like them to be.
The noisiness of this metric, can be especially seen, if one looks at the individual cosine Tresillo-ness values of the songs. Here we see that there is great variance in Tresillo-ness in the hand selected Tresillo data set. Values can be as small as 0.6, which corresponds to the value of a none Tresillo song.
Furthermore, it seems that also songs wich we classified to have no Tresillo, have a cosine Tresillo-ness up to 0.72.
In general we can state that the cosine similarity methods seems to identify a ceratin Tresillo-ness, however it is questionable how robust this method is.

### 3.3) Tresilio-ness over Time

Now that we explored several ways to compute Tresillo-ness, we can calculate the Tresilo-ness of ou billboard data set.
First let us load the rhythm vectors of our billboard songs and let's also merge it to the metadata of the billboard songs.
The metadata of the billboard songs includes time signatures with which we can Tresillo-ness over time.

```python
df_billboard = pd.read_csv(home_dir + '/dataset/rythm_vectors/rythm_vectors_billboard.csv')
df_billboard_meta = pd.read_csv(home_dir + '/dataset/billboard_data_sets/billboard_1999-2019_unique_top_20_selection.csv')

#calculate cosine simularity
df_billboard = collapse_normalize_vectors(df_billboard)
df_billboard_sim = calc_cosine_sim(df_billboard, vector_vanilla_tresillo)

#prepare for merge with meta data
df_billboard_sim['song_artist'] = df_billboard_sim.index
df_billboard_sim['song'] = df_billboard_sim.song_artist.apply(lambda x: x.split('_')[0])
df_billboard_sim['artist'] = df_billboard_sim.song_artist.apply(lambda x: x.split('_')[1][:-1])
df_billboard_sim = df_billboard_sim.drop(['song_artist'], axis=1)

#merge data frames
df_billboard_merged = df_billboard_sim.merge(df_billboard_meta, left_on=['song','artist'], right_on=['Name', 'Artists'], how='left')
df_billboard_reduced = df_bilboard_merged[['Name', 'Artists', 'Peak.position', 'Week', 'Genre', 'cosine_sim_tresillo']]
#let's look at some songs with very high tresillo-ness
df_billboard_reduced[df_billboard_reduced['cosine_sim_tresillo']>0.85]

```

<div>
<style scoped>
    .dataframe tbody tr th:only-of-type {
        vertical-align: middle;
    }

    .dataframe tbody tr th {
        vertical-align: top;
    }

    .dataframe thead th {
        text-align: right;
    }
</style>
<table border="1" class="dataframe">
  <thead>
    <tr style="text-align: right;">
      <th></th>
      <th>Name</th>
      <th>Artists</th>
      <th>Peak.position</th>
      <th>Week</th>
      <th>Genre</th>
      <th>cosine_sim_tresillo</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <th>89</th>
      <td>Cheap Thrills</td>
      <td>Sia</td>
      <td>19.0</td>
      <td>2016-06-18</td>
      <td>Jamaica,Remix,Australia,Rap,Synth-Pop,Pop</td>
      <td>0.908945</td>
    </tr>
    <tr>
      <th>293</th>
      <td>Rather Be</td>
      <td>Clean Bandit</td>
      <td>19.0</td>
      <td>2014-08-08</td>
      <td>Electro-Pop,Deep House,Eurodance,Dance-Pop,Cla...</td>
      <td>0.861075</td>
    </tr>
    <tr>
      <th>319</th>
      <td>Shape Of You</td>
      <td>Ed Sheeran</td>
      <td>1.0</td>
      <td>2017-01-28</td>
      <td>Tropical House,Dancehall,Pop,UK</td>
      <td>0.915565</td>
    </tr>
    <tr>
      <th>387</th>
      <td>Treat You Better</td>
      <td>Shawn Mendes</td>
      <td>20.0</td>
      <td>2016-07-30</td>
      <td>Canada,Pop</td>
      <td>0.883772</td>
    </tr>
    <tr>
      <th>389</th>
      <td>Turn Me On</td>
      <td>Kevin Lyttle</td>
      <td>20.0</td>
      <td>2004-06-23</td>
      <td>Pop</td>
      <td>0.937243</td>
    </tr>
    <tr>
      <th>423</th>
      <td>Where Have You Been</td>
      <td>Rihanna</td>
      <td>18.0</td>
      <td>2012-05-25</td>
      <td>House,Electro-Pop,Pop</td>
      <td>0.947448</td>
    </tr>
  </tbody>
</table>
</div>

Let's plot weekly cosine Tresillo-ness

```python
df_bilboard_merged_weekly = df_bilboard_merged.groupby(['Week'])['cosine_sim_tresillo'].agg(['mean']).reset_index()
plt.plot(pd.to_datetime(df_bilboard_merged_weekly.Week), df_bilboard_merged_weekly['mean'])
plt.title('Weekly Cosine Tresillo-ness in the Top 20 Billboards')
plt.xlabel('Time')
plt.ylabel('Cosine Tresillo-ness')
```

The time trend above is very noisy and hard to read. By calculating a 4 weeks moveing average we might get smoother results.

```python
df_bilboard_merged_weekly['rolling_mean'] = df_bilboard_merged_weekly.iloc[:,1].rolling(window=4).mean()

plt.plot(pd.to_datetime(df_bilboard_merged_weekly.Week), df_bilboard_merged_weekly['rolling_mean'])
plt.title('4 Weeks Moving Average Cosine Tresillo-ness in the Top 20 Billboards')
plt.xlabel('Time')
plt.ylabel('Cosine Tresillo-ness')
```

```python

```

```python

```

**TA instructions**:
- Present your results in relation to your research question.
- Present them in a logical order that does not have to be the order in which you achieved them.

**Thoughts Aurel:**

 4. Discussion of the clustering method, either k-mean clustering or something we dont have to set the cluster number
 5. Finding Tresillio-ness in the pop charts with all three methods a) Vanilla Tresillio-ness b) Tresillio songs vector c)

## 4) Outlook on final interpretation

Points to discuss as stated by TAs:
- Interpreting your results is the final step that you will do in preparing Milestone 4 (your presentations). Please end your submission by giving a first,preliminary outlook on this final step: what aspects of your results do you find interesting with respect to your hypotheses and previous literature? What do you think might be the main points to elaborate upon in the discussion?

```python

```

```python

```

## References

- Chin,  F.  and  Wu,  S.  (1992).   An  efficient  algorithm  for  rhythm-finding.Computer MusicJournal, 16(2):35–44.
- Dixon,  S.,  Gouyon,  F.,  Widmer,  G.,  et  al.  (2004).   Towards  characterisation  of  music  viarhythmic patterns.  InISMIR.
- Huron, D. and Ommen, A. (2006).  An empirical study of syncopation in american popularmusic, 1890–1939.Music Theory Spectrum, 28(2):211–231.
- Palmer, C. and Krumhansl, C. L. (1990).  Mental representations for musical meter.Journalof Experimental Psychology: Human Perception and Performance, 16(4):728.
- Panteli,  M.,  Bogaards,  N.,  Honingh,  A.  K.,  et  al.  (2014).   Modeling  rhythm  similarity  forelectronic dance music.  InISMIR, pages 537–542.
- Parry, M. and Essa, I. (2003).  Rhythmic similarity through elaboration.
- Pohle,  T.,  Schnitzer,  D.,  Schedl,  M.,  Knees,  P.,  and  Widmer,  G.  (2009).   On  rhythm  andgeneral music similarity.  InISMIR, pages 525–530.
- Rohrmeier, M. (2020).  Towards a formalization of musical rhythm.  InProc. of the 21st Int.Society for Music Information Retrieval Conf
- Floyd, Samuel A. "Black music in the circum-Caribbean." American Music (1999): 5-7.

```python

```


Chapter 1: Restricted Hartree Fock (RHF) for Simple Atom (He and Be) using  Slater Type Orbital (STO)
    with  Double Zeta Basis in Python

Integral calculated by sympy

    SIMUC PROJECT

    By Jinze (Richard) Xue   /
    Source code

1. It will be better to understand if you have read Levine, Quantum Chemistry, 7th Edition, chapter 11 and 14.
There is an example calculation of Helium atom at page 412-414.

1. For well-organized and documented python file, please find in github: https://github.com/yueyericardo/simuc/tree/master/notebooks/pchem/hartree-fock/sto

1. This notebook is self-contained.
However, if you wanna try in other place, you should put `hf.py` file in the same directory where you wanna import it.

### Chapters
- Chapter 1: Restricted Hartree Fock (RHF) for Simple Atom (He and Be) using Slater Type Orbital (STO) with Double Zeta Basis in Python
- Chapter 2: Slater Type Orbital (STO) VS Gaussian Type Orbital (GTO) 【WIP】
  Interactive tool 【WIP】
- Chapter 3: Restricted Hartree Fock (RHF) for Simple diatomic molecule (H2 and HeH+) using Gaussian Type Orbital (GTO) with STO-3G Basis in Python  【Code is on [github](https://github.com/yueyericardo/simuc/tree/master/notebooks/pchem/hartree-fock/gto)】【Document WIP】
- Chapter 4: Restricted Hartree Fock (RHF) for Polyatomic Molecule using Simple and Powerful Package: PSI4【WIP】

### Overview
1. [Introduction](#1.-Introduction)
1. [Born–Oppenheimer approximation](#2.-Born–Oppenheimer-approximation)
2. [Hartree fock approximation](#3.-Hartree-fock-approximation)
3. [Why build matrix and how to solve secular equation?](#4.-Why-build-matrix-and-how-to-solve-secular-equation?)
4. [How to build matrix](#5.-How-to-build-matrix)
5. [Secular Equation](#6.-Secular-Equation)
6. [Total Energy](#7.-Total-Energy)
7. [Utils](#8.-Utils)
8. [Run Hartree Fock](#9.-Run-Hartree-Fock)
9. [Test](#10.-Test)
10. [Excise - Plot the charge density of orbitals](#11.-Excise---Plot-the-charge-density-of-orbitals)
10. [Limitations](#12.-Limitations)
11. [Reference](#13.-Reference)

---

### 1. Introduction
In this notebook, we will use Helium Atom as an example to explain Restricted Hartree Fock (RHF) method.
(Restricted means only for closed shell molecule, orbitals are either doubly occupied or empty)

Helium has two electrons on 1s orbital.

One-electron wavefunction of 1s orbital could be written as $^*$
$$\chi_{1 \mathrm{s}}^{\mathrm{STO}}=\left(\frac{\zeta^{3}}{\pi}\right)^{1 / 2} \exp (-\zeta r)$$

     [*]
    For simplicity, there is no angular part since it's s orbital.

    $\zeta$ is a constant related to the effective charge of the nucleus, the nuclear charge being partly shielded by electrons.

With double zeta, we are saying we will buil two $\chi_{1 \mathrm{s}}^{\mathrm{STO}}$ function. And the final 1s orbital of Helium will be a linear combination of these two STO.  $$\phi_1 = c_{11} \chi_1 + c_{21} \chi_2  \quad \text{(occupied orbital)}\;\; \tag {Eq1}$$ $$\phi_2 = c_{12} \chi_1 + c_{22} \chi_2 \quad \text{(unoccupied orbital)}$$

The goal of hartree fock is using Self-Consistent Field method (SCF) to optimize coefficients to get close to the real wavefunction.

For Helium, from reference [[1]](https://www.sciencedirect.com/science/article/pii/S0092640X74800161), two double zeta are 1.45363, 2.91093.
We could build these two STO with each zeta using sympy like below.

```python
import hf
import sympy as sp
import numpy as np
import scipy.linalg
from sympy import oo
from sympy import diff
import time
import matplotlib
from matplotlib import pyplot as plt
from IPython.display import Math
sp.init_printing()
%matplotlib inline
```

```python
r, r1, r2, zeta = sp.symbols("r, r1, r2, zeta")
n = sp.Symbol('n', integer=True)
```

```python
def STO(zeta, n, r=r):
    """
    Define a Slater Type Orbital function using sympy.

    INPUT:
    zeta: zeta for the STO.
    n: principle quantum number for the STO.
    """
    f = r ** (n - 1) * sp.exp(-zeta * r)
    # normalization
    N = sp.sqrt(1 / sp.integrate(4 * sp.pi * f * f * r * r, (r, 0, +oo)))
    return N * f
```

```python
f1s_1 = hf.STO(zeta=1.45363, n=1)
f1s_2 = hf.STO(zeta=2.91093, n=1)

display(Math('$\chi_1 :'))
display(f1s_1)
display(Math('$\chi_2 :'))
display(f1s_2)
```

 Note:  More detail about this will be covered at part 2, but I want to point out the main idea at beginning. So it could remind you what's our final goal, when you feel distracted.

By using Roothan equations below , which could be solved self-consistently for the orbital coefficient matrix **C** and orbital energy eigenvalues $\epsilon_i$ by iterations, we could finally
1. get close to the real wavefunction (using improved Coefficients for Eq1).
2. get close to correct Molecular orbital energies (using $\epsilon_i$), to get the real total energy

$${\mathbf  {F}}{\mathbf  {C}}={\mathbf  {S}}{\mathbf  {C}}{\mathbf  {\epsilon }}$$

     Note:
    F (Fock matrix), S (Overlap matrix) are inputs.

    S matrix is fixed, F matrix is changing every iteration because of the improved C.

    C (Coefficient matrix) and $\epsilon_i$ (eigenvalues) are results.

---

### 2. Born–Oppenheimer approximation
Using Born–Oppenheimer approximation, molecular Hamiltonian could be expressed as

$${\displaystyle H=H_{\text{e}}+T_{\text{n}}}$$

     Note:
    e (electron), n (nuclear)

Where

$${\displaystyle H_{\text{e}}=-\sum _{i}{{\frac {1}{2}}\nabla _{i}^{2}}-\sum _{i,\alpha}{\frac {Z_{\alpha}}{r_{i\alpha}}}+\sum _{i>j}{\frac {1}{r_{ij}}}+\sum _{\beta>\alpha}{\frac {Z_{\alpha}Z_{\beta}}{R_{\alpha \beta}}}\quad {\text{and}}\quad T_{\text{n}}=-\sum _{\alpha}{{\frac {1}{2M_{\alpha}}}\nabla _{\alpha}^{2}}}$$

     Note:
    i, j (electron), $\alpha$, $\beta$ (nuclear)

$H_e$:
1. sum of kinetic-energy operators for each electron
2. sum of nuclear–electronic Coulombic attraction terms
3. sum of electron-electron repulsion energy
4. sum of nuclear-nuclear repulsion energy ($V_{NN}$)

$T_n$:
1. sum of kinetic-energy operators for each nuclear

Classically, during the time of a cycle of electronic motion, the change in nuclear configuration is negligible. Thus, considering the nuclei as fixed, we omit the nuclear kinetic-energy terms $T_n$. So commonly when we say hartree fock energy, we are only talking about $H_e$ term.

$${\displaystyle H_{\text{e}}=-\sum _{i}{{\frac {1}{2}}\nabla _{i}^{2}}-\sum _{i,\alpha}{\frac {Z_{\alpha}}{r_{i\alpha}}}+\sum_{i} \sum_{j>i}{\frac {1}{r_{ij}}} + V_{NN}}$$

The first 3 terms together is purely electronic Hamiltonian.

---

### 3. Hartree fock approximation

Because of the inter-electronic repulsion term $\frac{1}{r_{i j}}$, the Schrödinger equation for a molecule wavefunction is not separable. So the true wave function cannot be written as the product of n one-electron functions.

The essence of hartree-fock approximation is to treat electron-electron repulsion in an average way, so this complicated many-electron problem could be solved as one-electron problem.

Then the molecular wavefunction could be written as a product of all one-electron wavefunctions.
The functions chosen to represent each electron is based on the hydrogen-like atomic wavefunction. ( Note: This is the reason why we use Slater type Orbital (STO), because it's transformed from exact wavefunction for an electron around a hydrogen atom)

$$
\Psi({r_1 r_2})=\phi_{1}\left(\boldsymbol{r}_{1}\right) \phi_{2}\left(\boldsymbol{r}_{2}\right)
$$

However this product does not satisfy antisymmetric requirements (which means if you swap electrons the sign of the wavefunction should invert). This problem can be overcome by taking a linear combination of both products:
$$
\begin{aligned} \Psi\left(\mathbf{r}_{1}, \mathbf{r}_{2}\right) &=\frac{1}{\sqrt{2}}\left\{\phi_{1}\left(\mathbf{r}_{1}\right) \phi_{2}\left(\mathbf{r}_{2}\right)-\phi_{1}\left(\mathbf{r}_{2}\right) \phi_{2}\left(\mathbf{r}_{1}\right)\right\} \\ &=\frac{1}{\sqrt{2}}\left|\begin{array}{ll}{\phi_{1}\left(\mathbf{r}_{1}\right)} & {\phi_{2}\left(\mathbf{r}_{1}\right)} \\ {\phi_{1}\left(\mathbf{r}_{2}\right)} & {\phi_{2}\left(\mathbf{r}_{2}\right)}\end{array}\right| \end{aligned}
$$

In this way, the Schrödinger equation would then be separated into n one-electron hydrogenlike equations.

Suppose there is an operator called Fock operator $\hat F$, the eigenvalue corresponding to $\hat F$ on a one-electron wavefunction is the energy related to this electron.

Each electron energy $\varepsilon_{i}$ will include
1. kinetic-energy for this electron
2. sum of coulombic attraction between this electron and all nuclears
3. the potential of this electron interacting with an averaged distribution of other electrons, (which is calculated by treating all of the other electrons within the molecule as a smooth distribution of negative charge, and this is the major simplification inherent in the Hartree–Fock method).
(What does this means? explained at below )

 Note:   If we take$\sum_{i}^{n} {\varepsilon_{i}}$, we will count each interelectronic repulsion twice, which needs to be subtracted when calculating total energy of molecule.

What's the meaing of the potential of this electron interacting with an averaged distribution of other electrons?
suppose we want to find the electron-electron repulsion potential ($\text{Vee}$) of electron 1 with electron 2.

$$\text{Vee}_{1 2}=\left\langle\phi (1)^*|\frac{1}{r_{1 2}} | \phi (1)\right\rangle$$
This is not solvable, because we don't know the location of electron (2).

And hartree fock simplify this to:
(recall that $\left|\phi(2)\right|^{2}$ is the probability density of electron (2)

$$\text{Vee}_{1 2}=\left\langle\phi (1)^*|\;\;\; \int \frac{\left|\phi(2)\right|^{2}}{r_{12}} d v_{2} \;\;\; | \phi (1)\right\rangle \tag{Eq2}$$

By using the probability density of electron (2) $\left|\phi(2)\right|^{2}$,

$$\text{infinitesimal charge density} * \text{infinitesimal volume} = \text{infinitesimal charge}$$

and integrate over all space, we could get the repulsion energy above.
It could be also rewritten as, which is more commonly used:

$$
\text{Vee}_{1 2}= \int \int \frac{\left|\phi(2)\right|^{2} \phi (1)^* \phi (1)}{r_{12}} d v_{2}d v_{1}
$$
$$
= \int \int \frac{\phi(2)^*\phi(2) \phi (1)^* \phi (1)}{r_{12}} d v_{2}d v_{1}
$$

Fock operator

The operator corresponding to one electron energy $\varepsilon_{i}$ is fock operator. (For the restricted case which assumes closed-shell orbitals and single- determinantal wavefunctions)

$$\hat{F}(i)=-\frac{1}{2} \nabla_{i}^{2}-\sum_{\alpha} \frac{Z_{\alpha}}{r_{1 \alpha}} + \sum_{j=1}^{n / 2}\left[2 \hat{J}_{j}(i)-\hat{K}_{j}(i)\right]$$

where:

- ${\displaystyle {\hat {F}}(i)}$ is the Fock operator for the i-th electron in the system,
- $-\frac{1}{2} \nabla_{i}^{2}-\sum_{\alpha} \frac{Z_{\alpha}}{r_{1 \alpha}} $ are kinetics energy and sum of nuclear-electron attraction respectively.
These two terms are often considered as the core terms, and refered as $\hat{H}^{\mathrm{core}}_i$

So Fock operator could be rewritten as
$$\hat{F}(i)=\hat{H}_{core}(i)+\sum_{j=1}^{n / 2}\left[2 \hat{J}_{j}(i)-\hat{K}_{j}(i)\right]$$

- ${\displaystyle n}$ is the number of electrons and ${\displaystyle {\frac {n}{2}}}$ is the number of occupied orbitals in the closed-shell system,
- ${\displaystyle {\hat {J}}_{j}(i)}$ is the Coulomb operator, defining the repulsive force between the j-th and i-th electrons in the system,  (explained at Eq2)
$$
\hat{J}_{j}(1) f(1) =f(1) \int\left|\phi_{j}(2)\right|^{2} \frac{1}{r_{12}} d v_{2}
$$

     Note:
    ${\displaystyle f(1)}$, ${\displaystyle f(2)}$ are the one-electron wavefunctions acted upon by the exchange operator as functions of the electron positions,
 and ${\displaystyle \phi _{j}(1)}$ and ${\displaystyle \phi _{j}(2)}$ are the one-electron wavefunction of the jth electron as functions of the positions of the electrons.

- ${\displaystyle {\hat {K}}_{j}(i)}$ is the exchange operator, defining the quantum effect produced by exchanging two electrons.

$$\hat{K}_{j}(1) f(1) =\phi_{j}(1) \int \frac{\phi_{j}^{*}(2) f(2)}{r_{12}} d v_{2}$$

The Coulomb operator is multiplied by two since there are two electrons in each occupied orbital. The exchange operator is not multiplied by two since it has a non-zero result only for electrons which have the same spin as the i-th electron.

Roothan equation

The eigenvalue corresponding to the fock operator is the molecular orbital energies, electrons will occupy start from the lowest level, since we are dealing closed shell molecule, each orbital will has 2 or 0 electrons.

The Roothaan equations are a representation of the Hartree–Fock equation in a non orthonormal basis set (molecular orbitals represented by linear combination of atomic orbitals), which could be solved in matrix way that computer is good at.

$${\mathbf  {F}}{\mathbf  {C}}={\mathbf  {S}}{\mathbf  {C}}{\mathbf  {\epsilon }}$$

     Note:
    F (Fock matrix), S (Overlap matrix) are inputs.

    S matrix is fixed, F matrix is changing every iteration because of the improved C.

    C (Coefficient matrix) and $\epsilon_i$ (eigenvalues) are results.

Why iterations?

Since the Fock operator depends on the orbitals used to construct the corresponding Fock matrix, the eigenfunctions of the Fock operator are in turn new orbitals, which can be used to construct a new Fock operator. In this way, the Hartree–Fock orbitals are optimized iteratively until the change in total electronic energy falls below a predefined threshold. In this way, a set of self-consistent one-electron orbitals is calculated. The Hartree–Fock electronic wave function is then the Slater determinant constructed from these orbitals. Following the basic postulates of quantum mechanics, the Hartree–Fock wave function can then be used to compute any desired chemical or physical property within the framework of the Hartree–Fock method and the approximations employed.

---

### 4. Why build matrix and how to solve secular equation?
Because in this way, computer could solve it efficiently.

     Why build matrix? Click to uncollapse
(This part is from Calculating Orbital Energies and Expansion Coefficients - Chemistry LibreTexts

Solving secular equation is actually calculating orbital energies and coefficients based on **variation principle**, which states that any approximate wavefunction must have a higher energy than the true wavefunction.
(This part is from [Calculating Orbital Energies and Expansion Coefficients - Chemistry LibreTexts](https://chem.libretexts.org/Bookshelves/Physical_and_Theoretical_Chemistry_Textbook_Maps/Book%3A_Symmetry_(Vallance)/20%3A_Calculating_Orbital_Energies_and_Expansion_Coefficients) .)

Let's ignore the coulomb $\hat J$ and exchange $\hat K$ operator in fock operator, what's left is one electron Hamiltonian core operator $\hat H_{core}$. So how to calculate eigenvalue of Hamiltonian core operator and it's corresponding wavefunction of Helium with double zeta basis?
(Fock operator could be used as the same way.)

$$
E=\frac{\langle\phi|\hat{H}| \phi\rangle}{\langle\phi | \phi\rangle}
$$
Where
$$\phi = c_{1} \chi_1 + c_{2} \chi_2 \quad \text{(unnormalized)} $$
Plug in $\phi$ and expand

$$
\begin{aligned} E &=\frac{\left\langle c_{1} \chi_{1}+c_{2} \chi_{2}|\hat{H}| c_{1} \chi_{1}+c_{2} \chi_{2}\right\rangle}{\left\langle c_{1} \chi_{1}+c_{2} \chi_{2} | c_{1} \chi_{1}+c_{2} \chi_{2}\right\rangle} \\ &=\frac{\left\langle c_{1} \chi_{1}|\hat{H}| c_{1} \chi_{1}\right\rangle+\left\langle c_{1} \chi_{1}|\hat{H}| c_{2} \chi_{2}\right\rangle+\left\langle c_{2} \chi_{2}|\hat{H}| c_{1} \chi_{1}\right\rangle+\left\langle c_{2} \chi_{2}|\hat{H}| c_{2} \chi_{2}\right\rangle}{\left\langle c_{1} \chi_{1}| c_{1}\chi_{1}\right\rangle+\left\langle c_{1} \chi_{1}| c_{2} \chi_{2}\right\rangle+\left\langle c_{2} \chi_{2} | c_{1} \chi_{1}\right\rangle+\left\langle c_{2} \chi_{2} | c_{2} \chi_{2}\right\rangle} \\ &=\frac{c_{1}^{2}\left\langle\chi_{1}|\hat{H}| \chi_{1}\right\rangle+ c_{1} c_{2}\left\langle\chi_{1}|\hat{H}| \chi_{2}\right\rangle+ c_{2} c_{1}\left\langle\chi_{2}|\hat{H}| \chi_{1}\right\rangle+ c_{2}^{2}\left\langle\chi_{2}|\hat{H}| \chi_{2}\right\rangle}{c_{1}^{2}\left\langle\chi_{1} | \chi_{1}\right\rangle+ c_{1} c_{2}\left\langle\chi_{1} | \chi_{2}\right\rangle+ c_{2} c_{1}\left\langle\chi_{2} | \chi_{1}\right\rangle+ c_{2}^{2}\left\langle  \chi_{2}|\chi_{2}\right\rangle} \end{aligned}
$$

If define
$$
H_{i j}=\left\langle\chi_{i}|\hat{H}| \chi_{j}\right\rangle
\quad and \quad
S_{i j}=\left\langle\chi_{i} | \chi_{j}\right\rangle
\tag {Eq3}
$$

     Note: This is where H matrix (same to Fock matrix) and S matrix comes from! We will come build these matrix at part 5.

and note that  $H_{ij}  =  H_{j}$  and  $S_{ij}  =  S_{ji}$ ,

$$
E=\frac{c_{1}^{2} H_{11}+2 c_{1} c_{2} H_{12}+c_{2}^{2} H_{22}}{c_{1}^{2} S_{11}+2 c_{1} c_{2} S_{12}+c_{2}^{2} S_{22}}
$$

$$
E\left(c_{1}^{2} S_{11}+2 c_{1} c_{2} S_{12}+c_{2}^{2} S_{22}\right)=c_{1}^{2} H_{11}+2 c_{1} c_{2} H_{12}+c_{2}^{2} H_{22}
$$

To minimize the energy with respect to  c1  and  c2,  we require

$$
\frac{\partial E}{\partial c_{1}}=0
\quad and \quad
\frac{\partial E}{\partial c_{2}}=0
$$

If we differentiate the above equation through separately by  c1  and  c2  and apply this condition, we will end up with two equations in the two unknowns  c1  and  c2 , which we can solve to determine the coefficients and the energy.

$$
\begin{array}{l}{E\left(2 c_{1} S_{11}+2 c_{2} S_{12}\right)=2 c_{1} H_{11}+2 c_{2} H_{12}} \\ {E\left(2 c_{1} S_{12}+2 c_{2} S_{22}\right)=2 c_{1} H_{12}+2 c_{2} H_{22}}\end{array}
$$

These are normally rewritten slightly, in the form
$$
\begin{array}{l}{c_{1}\left(H_{11}-E S_{11}\right)+c_{2}\left(H_{12}-E S_{12}\right)=0} \\ {c_{1}\left(H_{12}-E S_{12}\right)+c_{2}\left(H_{22}-E S_{22}\right)=0}\end{array} \tag{Eq4}
$$
Write this in matrix form gives
$$
\left(\begin{array}{cc}{H_{11}-E S_{11}} & {H_{12}-E S_{12}} \\ {H_{12}-E S_{12}} & {H_{22}-E S_{22}}\end{array}\right)\left(\begin{array}{c}{c_{1}} \\ {c_{2}}\end{array}\right)=\left(\begin{array}{l}{0} \\ {0}\end{array}\right)
$$
For the equations to have a solution, the determinant of the matrix must be equal to zero. Which means
$$
\left(H_{11}-E\right)\left(H_{22}-E\right)-\left(H_{12}-E S_{12}\right)^{2}=0
$$
Now, there is only one unkown variable E, solve the equation will commonly give us two E (eigenvalue).
And put this two E back to Eq4 will give us two set of (c1, c2), which corresponding to
$$\phi_1 = c_{11} \chi_1 + c_{21} \chi_2 \quad (\varepsilon_{1})$$
$$\phi_2 = c_{12} \chi_1 + c_{22} \chi_2 \quad (\varepsilon_{2})$$

Thanks for powerful scipy, secular equation could be solved simply by calling
`eigenvalue, C = scipy.linalg.eigh(H, S)`

```python
# H and S are calculated from next part
H = np.array([[-1.85073991, -1.88346692],
              [-1.88346692, -1.58510327]])
S = np.array([[1.        , 0.83752358],
              [0.83752358, 1.        ]])

e, Co = scipy.linalg.eigh(H, S)

print(e)
print(Co)
```

    [-1.97961968  1.03859384]
    [[-0.66167682  1.70635833]
     [-0.37818627 -1.79065634]]

Let's test if whether the eigenvalue and coefficient satisfy Eq4.
Note: `[-0.66167682, -0.37818627]` is the eigenvector corresponding to eigenvalue `-1.97961968`
`[1.70635833, -1.79065634]` is the eigenvector corresponding to eigenvalue `1.03859384`
The result below $e^{-17}$ and $e^{-16}$ are so close to 0, there is no much difference with 0.

```python
tmp1 = Co[0, 0] * (H[0, 0] - e[0] * S[0, 0]) + Co[1, 0] * (H[0, 1] - e[0] * S[0, 1])
tmp2 = Co[0, 1] * (H[0, 1] - e[1] * S[0, 1]) + Co[1, 1] * (H[1, 1] - e[1] * S[1, 1])
print(tmp1)
print(tmp2)
```

    2.7755575615628914e-17
    8.881784197001252e-16

### 5. How to build matrix

Note:   To avoid confusion, letters `r`, `s`, `t`, `u` are used to label matrix element and the basis functions $\chi$, and the letters `i`, `j` are used to label the MOs $\phi$.

Back to Eq3, $H_{rs}$ in H matrix and $S_{rs}$ in S matrix are defined as
$$
H_{r s}=\left\langle\chi_{r}|\hat{H}| \chi_{s}\right\rangle
\quad and \quad
S_{r s}=\left\langle\chi_{r} | \chi_{s}\right\rangle
$$

(a) Hamiltonian core matrix H

Where Hamiltonian core operator $\hat H$ is
$$\hat{H}^{\mathrm{core}}_i \equiv-\frac{1}{2} \nabla_{i}^{2}-\sum_{\alpha} \frac{Z_{\alpha}}{r_{i \alpha}}$$

$$H_{r s}=\left\langle\chi_{r}|\hat{H}| \chi_{s}\right\rangle
\quad $$

$$=\int_{0}^\infty \chi_r \hat{H} \chi_s \; 4\pi r^2dr$$

$$= \int_{0}^\infty \chi_r ((-\dfrac{1}{2}) \nabla^2 -\sum_{\alpha} \frac{Z_{\alpha}}{r_{i \alpha}})\chi_s \; 4\pi r^2 dr$$

Where $
{\text { Laplace operator: } \nabla^{2}} {=\frac{1}{r^{2}} \frac{\partial}{\partial r}\left(r^{2} \frac{\partial f}{\partial r}\right)=\frac{1}{r} \frac{\partial^{2}}{\partial r^{2}}(r f)}$
(For wavefunction which only has radial part.)

And since we are calculating atom, there is only one nuclear $\alpha$

$$\therefore H_{rs}= \int_{0}^\infty \chi_r ((-\dfrac{1}{2}) \dfrac{1}{r} \dfrac{\partial}{\partial r} \dfrac{\partial}{\partial r} r \chi_s - \dfrac{Z_{\alpha}}{r} \chi_s )4\pi r^2 dr$$

Let's write this in sympy.
For easy reading, the code below seperate equation above into
$ H_{rs}= \int_{0}^\infty \chi_r (T - V )4\pi r^2 dr$

```python
def H_int(fr, fs, Z):
    """
    Compute H_core integral between two STO functions.
    H_core = electron kinetics energy + electron nuclear potential energy

    INPUT:
    Z: Nuclear charge
    """
    T = - ((1 / 2) * (1 / r) * diff(diff(r * fs, r), r))
    V = - (Z / r) * fs
    return sp.integrate(fr * (T + V) * 4 * sp.pi * r * r, (r, 0, +oo))
```

`H_int` function calculate $H_{rs}$ element in H matrix. To build H matrix, we just need to go over `r` from (1, 2, 3 ... num_bfs), and go over `s` also from (1, 2, 3 ... num_bfs), where `num_bfs` is number of basis functions, we will see this more at below.

```python
def H_matrix(bfs, Z):
    """
    Compute the core hamiltonian matrix H.
    H_core = electron kinetics energy + electron nuclear potential energy

    INPUT:
        bfs: basis functions
        Z: nuclear charge
    OUTPUT:
        H: core hamiltonian matrix
    """
    num_bfs = len(bfs)
    H = np.zeros((num_bfs, num_bfs))

    for r in range(num_bfs):
        for s in range(num_bfs):
            H[r, s] = H_int(bfs[r], bfs[s], Z)

    return H
```

(b) Overlap matrix S

If you understand H matrix, S (Overlap) matrix will be very easy.

$$S_{rs} = \int_{0}^\infty \chi_r^* \chi_s  \; 4 \pi r^2dr$$

```python
def S_int(fr, fs):
    """
    Compute overlap integral between two STO functions.
    """
    return sp.integrate(fr * fs * 4 * sp.pi * r * r, (r, 0, +oo))
```

```python
def S_matrix(bfs):
    """
    Compute overlap matrix S.

    INPUT:
        fs: basis functions
    OUTPUT:
        S: Overlap matrix
    """
    num_bfs = len(bfs)
    S = np.zeros((num_bfs, num_bfs))

    for r in range(num_bfs):
        for s in range(num_bfs):
            S[r, s] = S_int(bfs[r], bfs[s])

    return S
```

(c) Fock matrix F

Fock matrix F = H matrix + G matrix

Fock matrix could be generated similarly like H, where Fock operator is
$$\hat{F}(i)=\hat{H}_{core}(i)+\sum_{j=1}^{n / 2}\left[2 \hat{J}_{j}(i)-\hat{K}_{j}(i)\right]$$

The second term is often denoted as $\hat G$.

$$\hat{F}(i)=\hat{H}_{core}(i)+\hat G$$

Because we already have **H matrix** above, so we only need to build **G matrix**, add them together will generate **Fock matrix**.

G matrix

The defination of $J$ (Coulomb operator) and $K$ (exchange operator) are. (1) (2) below simple means they are different electron.

$$
\begin{aligned} \hat{J}_{j}(1) f(1) &=f(1) \int\left|\phi_{j}(2)\right|^{2} \frac{1}{r_{12}} d v_{2} \\ \hat{K}_{j}(1) f(1) &=\phi_{j}(1) \int \frac{\phi_{j}^{*}(2) f(2)}{r_{12}} d v_{2} \end{aligned}
$$

     Note:
    ${\displaystyle f(1)}$, ${\displaystyle f(2)}$ are the one-electron wavefunctions acted upon by the exchange operator as functions of the electron positions,
 and ${\displaystyle \phi _{j}(1)}$ and ${\displaystyle \phi _{j}(2)}$ are the one-electron wavefunction of the jth electron as functions of the positions of the electrons.

Let's see the result first, we will walk through an example ($G_{12}$ of helium) below explain how to get here. (Basically, it's just expand $\phi_i$ into linear combination of the basis functions.)
$$
\left\langle\chi_{r}(1) | \hat{J}_{j}(1) \chi_{s}(1)\right\rangle=\sum_{t} \sum_{u} c_{t j}^{*} c_{u j} \iint \frac{\chi_{r}^{*}(1) \chi_{s}(1) \chi_{t}^{*}(2) \chi_{u}(2)}{r_{12}} d v_{1} d v_{2}
$$
$$
\left\langle\chi_{r}(1) | \hat{K}_{j}(1) \chi_{s}(1)\right\rangle=\sum_{t} \sum_{u} c_{t j}^{*} c_{u j} \iint \frac{\chi_{r}^{*}(1) \chi_{u}(1) \chi_{t}^{*}(2) \chi_{s}(2)}{r_{12}} d v_{1} d v_{2}
$$
If define

$$
(r s | t u) \equiv \iint \frac{\chi_{r}^{*}(1) \chi_{s}(1) \chi_{t}^{*}(2) \chi_{u}(2)}{r_{12}} d v_{1} d v_{2}
$$

It could be rewritten as
$$
\left\langle\chi_{r}(1) | \hat{J}_{j}(1) \chi_{s}(1)\right\rangle=\sum_{t=1}^{b} \sum_{u=1}^{b} c_{t j}^{*} c_{u j}(r s | t u)
$$

$$
\left\langle\chi_{r}(1) | \hat{K}_{j}(1) \chi_{s}(1)\right\rangle=\sum_{t=1}^{b} \sum_{u=1}^{b} c_{t j}^{*} c_{u j}(r u | t s)
$$
And final $G_{rs}$ will be
$$
G_{r s}=\sum_{t=1}^{b} \sum_{u=1}^{b} \sum_{j=1}^{n / 2} c_{t j}^{*} c_{u j}[2(r s | t u)-(r u | t s)]
$$

     Example:

Let's take  $G_{12}$ element in G matrix (Helium atom) as an example to see how this comes out. (click to uncollapse)

$$G_{r s}=\left\langle\chi_{r}|\hat{G}| \chi_{s}\right\rangle$$
$$G_{1 2}=\left\langle\chi_{1}|\hat{G}| \chi_{2}\right\rangle$$

$$G_{1 2}= \sum_{j=1}^{n / 2} \left\langle\chi_{1}|2 \hat{J}_{j}(i)-\hat{K}_{j}(i)| \chi_{2}\right\rangle$$

For helium atom, num of electron (n) is 2, there is only one 1s orbital.

$$G_{1 2}= \left\langle\chi_{1}|2 \hat{J}_{j}(i)-\hat{K}_{j}(i)| \chi_{2}\right\rangle$$
$$G_{1 2}= 2\left\langle\chi_{1}| \hat{J}_{j}(i)| \chi_{2}\right\rangle - \left\langle\chi_{1}|\hat{K}_{j}(i)| \chi_{2}\right\rangle$$

Let's do $\left\langle\chi_{1}| \hat{J}_{j}(i)| \chi_{2}\right\rangle$ first

$$\left\langle\chi_{1}(1)| \hat{J}_{1}| \chi_{2}(1)\right\rangle $$
$$= \left\langle\chi_{1}(1)| \;\;\; \int \frac{\left|\phi_1(2)\right|^{2}}{r_{12}} d v_{2} \;\;\;  | \chi_{2}(1)\right\rangle $$
$$
= \int \int \frac{\phi_1(2)^*\phi_1(2) \chi (1)^* \chi (1)}{r_{12}} d v_{2}d v_{1}
$$

$$
= \int \int \frac{[c_{11}\chi_{1}^*(2) + c_{21}\chi_{2}^*(2)][c_{11}\chi_{1}(2) + c_{21}\chi_{2}(2)] \chi (1)^* \chi (1)}{r_{12}} d v_{2}d v_{1}
$$

$$
= \int \int \frac{c_{11}\chi_{1}^*(2)c_{11}\chi_{1}(2) \chi (1)^* \chi (1)}{r_{12}} d v_{2}d v_{1}
+ \int \int \frac{c_{11}\chi_{1}^*(2)c_{21}\chi_{2}(2) \chi (1)^* \chi (1)}{r_{12}} d v_{2}d v_{1}
$$
$$
+ \int \int \frac{c_{21}\chi_{2}^*(2)c_{11}\chi_{1}(2) \chi (1)^* \chi (1)}{r_{12}} d v_{2}d v_{1}
+ \int \int \frac{c_{21}\chi_{2}^*(2)c_{21}\chi_{2}(2) \chi (1)^* \chi (1)}{r_{12}} d v_{2}d v_{1}
$$

Use
$
(r s | t u) \equiv \iint \frac{\chi_{r}^{*}(1) \chi_{s}(1) \chi_{t}^{*}(2) \chi_{u}(2)}{r_{12}} d v_{1} d v_{2}
$ to simplify the representation.

$$
\left\langle\chi_{1}| \hat{J}_{j}(i)| \chi_{2}\right\rangle
= c_{11}c_{11}(1 2 | 1 1) + c_{11}c_{21}(1 2 | 1 2) + c_{21}c_{11}(1 2 | 2 1) + c_{21}c_{21}(1 2 | 2 2)
$$
$$
=\sum_{t=1}^{2} \sum_{u=1}^{2} \sum_{j=1}^{2 / 2} c_{t j}^{*} c_{u j} (1 \;2 | t \;u)
$$

Similar $\left\langle\chi_{1}| \hat{K}_{j}(i)| \chi_{2}\right\rangle$ could also be inferred,
$$
\left\langle\chi_{1}| \hat{K}_{j}(i)| \chi_{2}\right\rangle
= c_{11}c_{11}(1 1 | 1 2) + c_{11}c_{21}(1 2 | 1 2) + c_{21}c_{11}(1 1 | 2 2) + c_{21}c_{21}(1 2 | 2 2)
$$
$$
=\sum_{t=1}^{2} \sum_{u=1}^{2} \sum_{j=1}^{2 / 2} c_{t j}^{*} c_{u j} (1 \;u | t \;2)
$$

So $G_{1 2}$ is
$$G_{1 2}= \left\langle\chi_{1}|2 \hat{J}_{j}(i)-\hat{K}_{j}(i)| \chi_{2}\right\rangle$$
$$
=\sum_{t=1}^{2} \sum_{u=1}^{2} \sum_{j=1}^{2 / 2} c_{t j}^{*} c_{u j}[2(1 \;2 | t \;u)-(1\; u | t \;2)]
$$

Build G matrix in python

recall $G_{rs}$ is
$$
G_{r s}=\sum_{t=1}^{b} \sum_{u=1}^{b} \sum_{j=1}^{n / 2} c_{t j}^{*} c_{u j}[2(r s | t u)-(r u | t s)]
$$

If define $P_{rs}$ (density matrix) as
$$
P_{t u} \equiv 2 \sum_{j=1}^{n / 2} c_{t j}^{*} c_{u j}, \quad t=1,2, \ldots, b, \quad u=1,2, \ldots, b
$$

Then $G_{rs}$ could be simplified again as
$$
G_{r s}=\sum_{t=1}^{b} \sum_{u=1}^{b} P_{t u}\left[(r s | t u)-\frac{1}{2}(r u | t s)\right]
$$

If we built a matrix **R** which is 4 dimension, and save all the posible $(r s | t u)$, then we could build G matrix like below.

```python
def G_matrix(P, R):
    """
    Compute G matrix.
    G =  coulombic repulsion energy + exchange energy

    INPUT:
        P: density matrix
        R: electron repulsion matrix
    OUTPUT:
        G: repulsion matrix
    """
    num_bfs = P.shape[0]
    G = np.zeros((num_bfs, num_bfs))

    for r in range(num_bfs):
        for s in range(num_bfs):
            g = 0
            for t in range(num_bfs):
                for u in range(num_bfs):
                    int1 = R[r, s, t, u]
                    int2 = R[r, u, t, s]
                    g += P[t, u] * (int1 - 0.5 * int2)
            G[r, s] = g

    return G
```

P matrix

And also P (density) matrix
$$
P_{t u} \equiv 2 \sum_{j=1}^{n / 2} c_{t j}^{*} c_{u j}, \quad t=1,2, \ldots, b, \quad u=1,2, \ldots, b
$$

```python
def P_matrix(Co, N):
    """
    Compute density matrix P.

    INPUT:
        Co: coefficents matrix
        N: num of electrons
    OUTPUT:
        P: repulsion matrix
    """
    P = np.zeros([Co.shape[0], Co.shape[0]])

    for t in range(Co.shape[0]):
        for u in range(Co.shape[0]):
            for j in range(int(N/2)):
                P[t, u] += 2 * Co[t, j] * Co[u, j]
    return P
```

R matrix

R matrix (4 dimenstion), which is the most computation-expensive part in hartree-fock method. (In principle, there are a lot of elements in R matrix are equal by symmetry, but for code simplicity, we didn't implemented this.)

```python
def R_matrix(bfs):
    """
    Compute the electron repulsion integral matrix R.

    INPUT:
        fs: basis functions
    OUTPUT:
        R: repulsion matrix
    """
    start = time.time()
    num_bfs = len(bfs)
    R = np.zeros((num_bfs, num_bfs, num_bfs, num_bfs))

    for r in range(num_bfs):
        for s in range(num_bfs):
            for t in range(num_bfs):
                for u in range(num_bfs):
                    R[r, s, t, u] = R_int([bfs[r], bfs[s], bfs[t], bfs[u]])

    stop = time.time()
    print('time Repu: {:.1f} s'.format(stop-start))
    return R
```

Recall the element $R_{rstu}$ in R matrix is
$$
(r s | t u) \equiv \iint \frac{\chi_{r}^{*}(1) \chi_{s}(1) \chi_{t}^{*}(2) \chi_{u}(2)}{r_{12}} d v_{1} d v_{2}
$$

For 1s or 2s orbital which only has radial part, this could be calculated approximately in sympy. However, for the orbitals which has angular part, how to solve two-electron Repulsion integral of Slater Type Orbital (STO) is still slow and difficult. This is the main reason why Gaussian Type Orbital (GTO) is more frequently used, and we will try GTO in next chapter.

$$(rs|tu)  = \int_{0}^\infty \int_{0}^\infty \dfrac{\chi_r^*(1) \chi_s(1) \chi_t^*(2) \chi_u(2)}{r_{12}} \; 4 \pi r_1^2dr_1\; 4 \pi r_2^2dr_2  $$

$$(rs|tu) = \int_{0}^\infty \chi_r^*(1) \chi_s(1)  \; 4 \pi r_1^2dr_1\int_{0}^\infty \frac{ \chi_t^*(2) \chi_u(2)}{r_{12}}\; 4 \pi r_2^2dr_2  $$

$r_{12}$ here is tricky to deal with, from problem 9.14 in quantum_chemistry by levine

$$(rs|tu) = \int_{0}^\infty \chi_r^*(1) \chi_s(1)  \; 4 \pi r_1^2dr_1\int_{0}^\infty \frac{ \chi_t^*(2) \chi_u(2)}{r_{>}}\; 4 \pi r_2^2dr_2  $$

     Note:
    $r_{>}$ is the larger one between $r1$ and $r2$

$$(rs|tu) = \int_{0}^\infty \chi_r^*(1) \chi_s(1)  \; 4 \pi r_1^2dr_1(\int_{0}^{r_1} \frac{ \chi_t^*(2) \chi_u(2)}{r_{1}}\; 4 \pi r_2^2dr_2 + \int_{r_1}^\infty \frac{ \chi_t^*(2) \chi_u(2)}{r_{2}}\; 4 \pi r_2^2dr_2)  $$

Let$ \; B= \int_{0}^{r_1} \frac{ \chi_t^*(2) \chi_u(2)}{r_{1}}\; 4 \pi r_2^2dr_2 + \int_{r_1}^\infty \frac{ \chi_t^*(2) \chi_u(2)}{r_{2}}\; 4 \pi r_2^2dr_2$

$$(rs|tu) = \int_{0}^\infty \chi_r^*(1) \chi_s(1) B \; 4 \pi  r_1^2 dr_1  $$

So the element $R_{rstu}$ in R matrix could be calculated by

```python
def R_int(four_bfs):
    """
    Compute electron-electron repulsion integral.

    INPUT:
    four_bfs: an array contain 4 basis functions
    """
    f1, f2, f3, f4 = four_bfs

    f1 = f1.subs(r, r1)
    f2 = f2.subs(r, r1)
    f3 = f3.subs(r, r2)
    f4 = f4.subs(r, r2)

    B = (1 / r1) * sp.integrate(f3 * f4 * 4 * sp.pi * r2 * r2, (r2, 0, r1)) + sp.integrate((1 / r2) * f3 * f4 * 4 * sp.pi * r2 * r2, (r2, r1, +oo))
    return sp.integrate(f1 * f2 * 4 * sp.pi * r1 * r1 * B, (r1, 0, +oo))
```

This is basically all the matrixes we will need to use, and we already solved the most difficult part of hartree fock!!

### 6. Secular Equation

**Fock matrix** is the sum of **H matrix** and **G matrix**.
F = H + G

Then, Roothan equation could be solved simply by calling
`eigenvalue, C = scipy.linalg.eigh(F, S)`
eigenvalues are MO orbital energies.

$${\mathbf  {F}}{\mathbf  {C}}={\mathbf  {S}}{\mathbf  {C}}{\mathbf  {\epsilon }}$$

     Note:
    F (Fock matrix), S (Overlap matrix) are inputs.

    S matrix is fixed, F matrix is changing every iteration because of the improved C.

    C (Coefficient matrix) and $\epsilon_i$ (eigenvalues) are results.

We could wrap `scipy.linalg.eigh` into a function

```python
def secular_eqn(F, S):
    """
    Slove secular equation, return the MO energies (eigenvalue) and improved coeffients (eigenvector)

    INPUT:
        F: fock matrix or h_core matrix
        S: overlap integral
    OUTPUT:
        ei: eigenvalue
        C: eigenvector
    """
    ei, C = scipy.linalg.eigh(F, S)
    return ei, C
```

### 7. Total Energy
Finally, hartree fock total energy is
$$
E_{\mathrm{HF}}=2 \sum_{i=1}^{n / 2} \varepsilon_{i}-\sum_{i=1}^{n / 2} \sum_{j=1}^{n / 2}\left(2 J_{i j}-K_{i j}\right)+V_{N N}
$$

     Note:
    i, j are electron (only r, s, t, u are used for matrix element)

Where every term is
1. Energy of all the electrons (in closed shell condition, every orbital is doubly-occupied)
2. Recall that when we build F operator and calculat electron energy, electron-electron repulsion is counted twice when you sum all the electrons energy. This is the remove the extra count.
3. Nuclear-Nuclear repulsion term, which is not included in F operator. (because we are calculating Atom, this term is always 0)

half sum of the all the electrons $({n / 2})$ could also be written as
$$
\sum_{i=1}^{n / 2} \varepsilon_{i}=\sum_{i=1}^{n / 2} H_{i i}^{\mathrm{core}}+\sum_{i=1}^{n / 2} \sum_{j=1}^{n / 2}\left(2 J_{i j}-K_{i j}\right)
$$

So $E_{\mathrm{HF}}$ could be simplifed as
$$
E_{\mathrm{HF}}=\sum_{i=1}^{n / 2} \varepsilon_{i}+\sum_{i=1}^{n / 2} H_{i i}^{\mathrm{core}}+V_{N N}
$$

     Note:
    i, j are electron (only r, s, t, u are used for matrix element)

    H here is not H matrix, which is $
H_{i i}^{\text {core }}=\left\langle\phi_{i}\left|\hat{H}^{\text {core }}\right| \phi_{i}\right\rangle
$

Expand $\phi_{i}$ into basis functions and simplify, we could calculate total energy using the matrix we have built before.
$$
E_{\mathrm{HF}}=\sum_{i=1}^{n / 2} \varepsilon_{i}+\frac{1}{2} \sum_{r=1}^{b} \sum_{s=1}^{b} P_{r s} H_{r s}^{\mathrm{core}}+V_{N N}
$$

```python
def energy_tot(e, N, P, H, Vnn=0):
    """
    Compute the total energy.

    INPUT:
    e: MO energies
    N: num of electrons
    P: density matrix
    H: h_core matrix
    Vnn: nuclear nuclear repulsion energy, for atom is 0
    """
    e_tot = 0

    for i in range(int(N/2)):
        e_tot += e[i].real

    e_tot = e_tot + 0.5 * (P * H).sum() + Vnn
    return e_tot
```

### 8. Utils

And also some utils function to print information and compare our result with reference

```python
def print_info(S, H, e, Co, P, hf_e, start, stop, delta_e=0, verbose=False):
    """
    Print information while doing SCF interations.
    """
    if(verbose):
        # overlap
        print('Overlap:')
        print(S)

        # hamiltonian
        print('Core hamiltonian:')
        print(H)

        # Co
        print('Coefficients:')
        print(Co)

        # density
        print('Density matrix:')
        print(P)

        # MOs
        print('MO energies:')
        message = ', '
        m_list = ['e{} = {:0.3f}'.format(i+1, x) for i, x in enumerate(e)]
        message = message.join(m_list)
        print(message)

    print('HF energy: {:0.5f} (hartree) = {:0.5f} (eV)'.format(hf_e, hf_e*27.211))
    if delta_e != 0:
        print('dE       : {:.2e}'.format(delta_e))
    print('time used: {:.1f} s'.format(stop-start))

def compare(cal, ref, tol=1.0e-4):
    """
    Compare calculated result with reference data.
    """
    delta = np.abs(ref - cal)
    if delta < tol:
        message = '\33[32m' + 'PASSED' + '\x1b[0m'
    else:
        message = '\033[91m' + 'FAILED' + '\033[0m'
    print('-' * 32, message, '-' * 33)
    print('cal: {:.7f}, ref: {:.7f}\n\n'.format(cal, ref))
```

### 9. Run Hartree Fock

Steps to run hartree fock

1. Initialization
  - Let Fock matrix = H_core matrix, without considering electron repulsion
  - Solve secular equation with H and S to get initial Co (means initial guessed molecular orbitals) and build inital P (density) matrix
  - Prepare Repulsion matrix R (take time)
2. Iteration
  - Using P matrix and R matrix to calculate G matrix
  - F matrix = H matrix + G matrix
  - Solve secular equation with F and S to get improved Co (means improved molecular orbitals)
  - Using improved Co to build improved P matrix
  - check whether converged (the change of total energy smaller than converge requirement)

```python
def run_hf(bfs, Z):
    """
    Run restricted hartree fock for a single atom.

    INPUT:
        bfs: basis functions
        Z: nuclear charge of the atom
    """
    print('------------------------------', "Initialization", '------------------------------')
    print('-------------------------', "Ignore repulsion integral", '------------------------')
    N = Z  # num of electron = nuclear charege (since it's atom)
    start = time.time()

    # initialization
    H = H_matrix(bfs, Z)
    S = S_matrix(bfs)
    e, Co = secular_eqn(H, S)
    P = P_matrix(Co, N)
    Vnn = 0  # A single atom does not have nuclear repulsion
    hf_e = energy_tot(e, N, P, H, Vnn)

    stop = time.time()
    print_info(S, H, e, Co, P, hf_e, start, stop, verbose=verbose)
    print('-----------', "Caculating Electron Repulsion Integral (takes time)", '------------')
    R = R_matrix(bfs)
    delta_e = 1
    ITER = 0
    previous_e = hf_e

    # Iterations
    while(delta_e > E_conv and ITER < MAXITER):
        print('------------------------------', "Iteration", ITER + 1, '------------------------------')
        start = time.time()

        # important scf steps
        G = G_matrix(P, R)
        F = H + G
        e, Co = secular_eqn(F, S)
        P = P_matrix(Co, N)
        hf_e = energy_tot(e, N, P, H, Vnn)

        delta_e = np.abs(hf_e - previous_e)
        previous_e = hf_e
        ITER += 1
        stop = time.time()
        print_info(S, H, e, Co, P, hf_e, start, stop, delta_e, verbose=verbose)

    return hf_e
```

Set converge criterion

```python
MAXITER = 40     # Maximum SCF iterations
E_conv = 1.0e-6  # Energy convergence criterion
verbose = False  # whether to print matrix information while iterating
```

### 10. Test

Run hartree fock for Helium

```python
def test1():
    # Use 2 Slater Type ourbital to represent Helium 1s orbital.
    # The final Helium 1s orbital is a linear combination of these two STO.
    f1s_1 = STO(zeta=1.45363, n=1)
    f1s_2 = STO(zeta=2.91093, n=1)

    # all basis functions
    fs = [f1s_1, f1s_2]

    #  nuclear charge of He
    Z = 2

    # run hartree fock
    hf_e = run_hf(fs, Z)

    # compare result with reference
    ref_hf_e = -2.8616726
    compare(hf_e, ref_hf_e)
```

```python
test1()
```

    ------------------------------ Initialization ------------------------------
    ------------------------- Ignore repulsion integral ------------------------
    HF energy: -3.95924 (hartree) = -107.73486 (eV)
    time used: 0.2 s
    ----------- Caculating Electron Repulsion Integral (takes time) ------------
    time Repu: 3.2 s
    ------------------------------ Iteration 1 ------------------------------
    HF energy: -2.78457 (hartree) = -75.77099 (eV)
    dE       : 1.17e+00
    time used: 0.0 s
    ------------------------------ Iteration 2 ------------------------------
    HF energy: -2.85860 (hartree) = -77.78532 (eV)
    dE       : 7.40e-02
    time used: 0.0 s
    ------------------------------ Iteration 3 ------------------------------
    HF energy: -2.86152 (hartree) = -77.86490 (eV)
    dE       : 2.92e-03
    time used: 0.0 s
    ------------------------------ Iteration 4 ------------------------------
    HF energy: -2.86167 (hartree) = -77.86877 (eV)
    dE       : 1.42e-04
    time used: 0.0 s
    ------------------------------ Iteration 5 ------------------------------
    HF energy: -2.86167 (hartree) = -77.86896 (eV)
    dE       : 7.01e-06
    time used: 0.0 s
    ------------------------------ Iteration 6 ------------------------------
    HF energy: -2.86167 (hartree) = -77.86897 (eV)
    dE       : 3.45e-07
    time used: 0.0 s
    -------------------------------- PASSED ---------------------------------
    cal: -2.8616726, ref: -2.8616726

Run hartree fock for Beryllium

```python
def test2():
    """
    Test of Be (1s, 2s)
    """
    # Use 2 STO to represent Be 1s orbital and another 2 STO for 2s orbital
    # The final 1s orbital is a linear combination of these 4 STO.
    # Same for 2s orbital.
    f1s_1 = STO(zeta=5.59108, n=1)
    f1s_2 = STO(zeta=3.35538, n=1)
    f2s_1 = STO(zeta=1.01122, n=2)
    f2s_2 = STO(zeta=0.61000, n=2)

    # all basis functions
    fs = [f1s_1, f1s_2, f2s_1, f2s_2]

    # nuclear charge of Be
    Z = 4

    # run hartree fock
    hf_e = run_hf(fs, Z)

    # compare result with reference
    ref_hf_e = -14.572369
    compare(hf_e, ref_hf_e)
```

```python
test2()
```

    ------------------------------ Initialization ------------------------------
    ------------------------- Ignore repulsion integral ------------------------
    HF energy: -19.51846 (hartree) = -531.11686 (eV)
    time used: 1.0 s
    ----------- Caculating Electron Repulsion Integral (takes time) ------------
    time Repu: 56.9 s
    ------------------------------ Iteration 1 ------------------------------
    HF energy: -14.28744 (hartree) = -388.77544 (eV)
    dE       : 5.23e+00
    time used: 0.0 s
    ------------------------------ Iteration 2 ------------------------------
    HF energy: -14.50425 (hartree) = -394.67518 (eV)
    dE       : 2.17e-01
    time used: 0.0 s
    ------------------------------ Iteration 3 ------------------------------
    HF energy: -14.55105 (hartree) = -395.94867 (eV)
    dE       : 4.68e-02
    time used: 0.0 s
    ------------------------------ Iteration 4 ------------------------------
    HF energy: -14.56616 (hartree) = -396.35976 (eV)
    dE       : 1.51e-02
    time used: 0.0 s
    ------------------------------ Iteration 5 ------------------------------
    HF energy: -14.57061 (hartree) = -396.48075 (eV)
    dE       : 4.45e-03
    time used: 0.0 s
    ------------------------------ Iteration 6 ------------------------------
    HF energy: -14.57187 (hartree) = -396.51521 (eV)
    dE       : 1.27e-03
    time used: 0.0 s
    ------------------------------ Iteration 7 ------------------------------
    HF energy: -14.57223 (hartree) = -396.52493 (eV)
    dE       : 3.57e-04
    time used: 0.0 s
    ------------------------------ Iteration 8 ------------------------------
    HF energy: -14.57233 (hartree) = -396.52766 (eV)
    dE       : 1.01e-04
    time used: 0.0 s
    ------------------------------ Iteration 9 ------------------------------
    HF energy: -14.57236 (hartree) = -396.52843 (eV)
    dE       : 2.83e-05
    time used: 0.0 s
    ------------------------------ Iteration 10 ------------------------------
    HF energy: -14.57237 (hartree) = -396.52865 (eV)
    dE       : 7.94e-06
    time used: 0.0 s
    ------------------------------ Iteration 11 ------------------------------
    HF energy: -14.57237 (hartree) = -396.52871 (eV)
    dE       : 2.23e-06
    time used: 0.0 s
    ------------------------------ Iteration 12 ------------------------------
    HF energy: -14.57237 (hartree) = -396.52872 (eV)
    dE       : 6.27e-07
    time used: 0.0 s
    -------------------------------- PASSED ---------------------------------
    cal: -14.5723687, ref: -14.5723690

### 11. Excise - Plot the charge density of orbitals
Write a function based on the code of `run_hf()`, plot the charge density $4 \pi r^2 |\phi_i|^2$ of all the orbitals (final iteration).

For example, Beryllium:
Note: You only need to plot the final one (most bottom)
Hint see below.

Hint: convert a sympy expression to python function by using [sp.lambdify()](https://docs.sympy.org/latest/modules/numeric-computation.html#lambdify)
For example, to plot $\chi_1$, which is the 1st wavefunction of Beryllium Basis set.

```python
f1s_1 = STO(zeta=5.59108, n=1)
display(f1s_1)
f = sp.lambdify(r, f1s_1, "numpy")
```

```python
x = np.linspace(0, 5, 300)
y = f(x)
plt.figure(figsize=(16,2))
plt.plot(x, y, label='$\chi_{}$'.format(1))
plt.legend()
plt.xlim(0, x[-1])
plt.xlabel('r')
plt.ylabel('$\chi_1$')
plt.title('Wavefunction of $\chi_1$')
plt.show()
```

Code below is example solution

```python
def run_hf(bfs, Z):
    """
    Run restricted hartree fock for a single atom.

    INPUT:
        bfs: basis functions
        Z: nuclear charge of the atom
    """
    print('------------------------------', "Initialization", '------------------------------')
    print('-------------------------', "Ignore repulsion integral", '------------------------')
    N = Z  # num of electron = nuclear charege (since it's atom)
    start = time.time()

    # initialization
    H = H_matrix(bfs, Z)
    S = S_matrix(bfs)
    e, Co = secular_eqn(H, S)
    P = P_matrix(Co, N)
    Vnn = 0  # A single atom does not have nuclear repulsion
    hf_e = energy_tot(e, N, P, H, Vnn)

    stop = time.time()
    print_info(S, H, e, Co, P, hf_e, start, stop, verbose=verbose)
    print('-----------', "Caculating Electron Repulsion Integral (takes time)", '------------')
    R = R_matrix(bfs)
    delta_e = 1
    ITER = 0
    previous_e = hf_e

    densities = [] # [[d1, d2], [d1, d2], [d1, d2]]
    # plot
    x = np.linspace(0, 5, 300)
    tmp_density = get_density(bfs, Co, x)
    densities.append(tmp_density)

    # Iterations
    while(delta_e > E_conv and ITER < MAXITER):
        print('------------------------------', "Iteration", ITER + 1, '------------------------------')
        start = time.time()

        # important scf steps
        G = G_matrix(P, R)
        F = H + G
        e, Co = secular_eqn(F, S)
        P = P_matrix(Co, N)
        hf_e = energy_tot(e, N, P, H, Vnn)

        delta_e = np.abs(hf_e - previous_e)
        previous_e = hf_e
        ITER += 1
        stop = time.time()
        print_info(S, H, e, Co, P, hf_e, start, stop, delta_e, verbose=verbose)

        # plot
        tmp_density = get_density(bfs, Co, x)
        densities.append(tmp_density)

    plot_density(densities, x)

    return hf_e
```

```python
def get_density(bfs, Co, x):
    r = sp.Symbol('r')
    density = []

    # all orbitals
    for i in range(Co.shape[0]):
        tmp_orbital = 0
        # all basis functions
        for j, f in enumerate(bfs):
            tmp_orbital += Co[j][i] * bfs[j]
        tmp_d_function = tmp_orbital * tmp_orbital * r * r * 4 * np.pi
        tmp_f = sp.lambdify(r, tmp_d_function, "numpy")
        tmp_d = tmp_f(x)
        density.append(tmp_d)
    return density
```

```python
def plot_density(densities, x):
    num_orbitals = len(densities[0])
    num_iterations = len(densities)

    for i in range(num_orbitals):
        plt.figure(figsize=(16,2))
        for j, d in enumerate(densities):
            if j == 0 or j == (num_iterations - 1):
                plt.plot(x, d[i], label='iteration {}'.format(j+1))
            else:
                plt.plot(x, d[i])
        plt.legend()
        plt.xlim(0, x[-1])
        plt.xlabel('r')
        plt.ylabel('$4 \pi r^2 |\phi_{}|^2$'.format(i+1))
        plt.title('charge density of $\phi_{}$'.format(i+1))
        plt.show()

    plt.figure(figsize=(16,2))
    for i, d in enumerate(densities[-1]):
        plt.plot(x, d, label='$\phi_{}$'.format(i+1))
    plt.legend()
    plt.xlim(0, x[-1])
    plt.xlabel('r')
    plt.ylabel('$4 \pi r^2 |\phi_i|^2$')
    plt.title('charge density of all orbitals $\phi_i$ (final iteration)')
    plt.show()
```

```python
test1()
```

```python
test2()
```

### 12. Limitations
Limitations of this implementation:
1. Because of the coordinate system (Spherical coordinate system), this implementation can only deal with atom.
1. Because angular part is not included in current STO function, only atoms which only have s orbital could be represented by STO.
1. For restricted hartree fock, could only run closed shell atom.
1. Integral calculated by sympy, which is slow but easy to understand.

### 13. Reference
[1] Levine, Quantum Chemistry, 7th Edition, chapter 14
[2] Wikipedia
[Hartree–Fock method](https://en.wikipedia.org/wiki/Hartree%E2%80%93Fock_method#Hartree%E2%80%93Fock_algorithm),
[Fock matrix](https://en.wikipedia.org/wiki/Fock_matrix),
[Roothaan equations](https://en.wikipedia.org/wiki/Roothaan_equations),
[Coulomb operator](https://en.wikipedia.org/wiki/Coulomb_operator),
[Exchange operator](https://en.wikipedia.org/wiki/Exchange_operator)
[3] Clementi, Enrico, and Carla Roetti. [Roothaan-Hartree-Fock atomic wavefunctions: Basis functions and their coefficients for ground and certain excited states of neutral and ionized atoms, Z≤ 54.](https://www.sciencedirect.com/science/article/pii/S0092640X74800161) Atomic data and nuclear data tables 14.3-4 (1974): 177-478.
[4] Acosta C R. [Restricted closed shell Hartree Fock Roothaan matrix method applied to Helium atom using Mathematica[J].](https://files.eric.ed.gov/fulltext/EJ1051495.pdf) European Journal of Physics Education, 2017, 5(1): 1-14.
[5] Jacob Martin, [Simple Quantum Chemistry: Hartree-Fock in Python](http://nznano.blogspot.com/2018/03/simple-quantum-chemistry-hartree-fock.html), 2018
[6] Claire Vallance, [Calculating Orbital Energies and Expansion Coefficients](https://chem.libretexts.org/Bookshelves/Physical_and_Theoretical_Chemistry_Textbook_Maps/Book%3A_Symmetry_(Vallance)/20%3A_Calculating_Orbital_Energies_and_Expansion_Coefficients), 2019