# Recurrent neural networks

We show how to implement RNNs from scratch.
Based on sec 8.5 of http://d2l.ai/chapter_recurrent-neural-networks/rnn-scratch.html.

```python
!pip install -q flax
```

         |████████████████████████████████| 184 kB 4.2 MB/s
         |████████████████████████████████| 136 kB 47.1 MB/s
         |████████████████████████████████| 72 kB 777 kB/s

```python
import jax.numpy as jnp
import matplotlib.pyplot as plt
import math
from IPython import display

import jax
import flax.linen as nn
from flax import jax_utils
import optax

import collections
import re
import random
import os
import requests
import hashlib
import time
import functools

random.seed(0)
rng = jax.random.PRNGKey(0)

!mkdir figures # for saving plots
```

# Data

 As data, we use the book "The Time Machine" by H G Wells,
preprocessed using the code in [this colab](https://colab.research.google.com/github/probml/pyprobml/blob/master/notebooks/text_preproc_torch.ipynb).

```python
class SeqDataLoader:
    """An iterator to load sequence data."""
    def __init__(self, batch_size, num_steps, use_random_iter, max_tokens):
        if use_random_iter:
            self.data_iter_fn = seq_data_iter_random
        else:
            self.data_iter_fn = seq_data_iter_sequential
        self.corpus, self.vocab = load_corpus_time_machine(max_tokens)
        self.batch_size, self.num_steps = batch_size, num_steps

    def __iter__(self):
        return self.data_iter_fn(self.corpus, self.batch_size, self.num_steps)

class Vocab:
    """Vocabulary for text."""
    def __init__(self, tokens=None, min_freq=0, reserved_tokens=None):
        if tokens is None:
            tokens = []
        if reserved_tokens is None:
            reserved_tokens = []
        # Sort according to frequencies
        counter = count_corpus(tokens)
        self.token_freqs = sorted(counter.items(), key=lambda x: x[1],
                                  reverse=True)
        # The index for the unknown token is 0
        self.unk, uniq_tokens = 0, ['<unk>'] + reserved_tokens
        uniq_tokens += [
            token for token, freq in self.token_freqs
            if freq >= min_freq and token not in uniq_tokens]
        self.idx_to_token, self.token_to_idx = [], dict()
        for token in uniq_tokens:
            self.idx_to_token.append(token)
            self.token_to_idx[token] = len(self.idx_to_token) - 1

    def __len__(self):
        return len(self.idx_to_token)

    def __getitem__(self, tokens):
        if not isinstance(tokens, (list, tuple)):
            return self.token_to_idx.get(tokens, self.unk)
        return [self.__getitem__(token) for token in tokens]

    def to_tokens(self, indices):
        if not isinstance(indices, (list, tuple)):
            return self.idx_to_token[indices]
        return [self.idx_to_token[index] for index in indices]
```

```python
def tokenize(lines, token='word'):
    """Split text lines into word or character tokens."""
    if token == 'word':
        return [line.split() for line in lines]
    elif token == 'char':
        return [list(line) for line in lines]
    else:
        print('ERROR: unknown token type: ' + token)

def count_corpus(tokens):
    """Count token frequencies."""
    # Here `tokens` is a 1D list or 2D list
    if len(tokens) == 0 or isinstance(tokens[0], list):
        # Flatten a list of token lists into a list of tokens
        tokens = [token for line in tokens for token in line]
    return collections.Counter(tokens)

def seq_data_iter_random(corpus, batch_size, num_steps):
    """Generate a minibatch of subsequences using random sampling."""
    # Start with a random offset (inclusive of `num_steps - 1`) to partition a
    # sequence
    corpus = corpus[random.randint(0, num_steps - 1):]
    # Subtract 1 since we need to account for labels
    num_subseqs = (len(corpus) - 1) // num_steps
    # The starting indices for subsequences of length `num_steps`
    initial_indices = list(range(0, num_subseqs * num_steps, num_steps))
    # In random sampling, the subsequences from two adjacent random
    # minibatches during iteration are not necessarily adjacent on the
    # original sequence
    random.shuffle(initial_indices)

    def data(pos):
        # Return a sequence of length `num_steps` starting from `pos`
        return corpus[pos:pos + num_steps]

    num_batches = num_subseqs // batch_size
    for i in range(0, batch_size * num_batches, batch_size):
        # Here, `initial_indices` contains randomized starting indices for
        # subsequences
        initial_indices_per_batch = initial_indices[i:i + batch_size]
        X = [data(j) for j in initial_indices_per_batch]
        Y = [data(j + 1) for j in initial_indices_per_batch]
        yield jnp.array(X), jnp.array(Y)

def seq_data_iter_sequential(corpus, batch_size, num_steps):
    """Generate a minibatch of subsequences using sequential partitioning."""
    # Start with a random offset to partition a sequence
    offset = random.randint(0, num_steps)
    num_tokens = ((len(corpus) - offset - 1) // batch_size) * batch_size
    Xs = jnp.array(corpus[offset:offset + num_tokens])
    Ys = jnp.array(corpus[offset + 1:offset + 1 + num_tokens])
    Xs, Ys = Xs.reshape(batch_size, -1), Ys.reshape(batch_size, -1)
    num_batches = Xs.shape[1] // num_steps
    for i in range(0, num_steps * num_batches, num_steps):
        X = Xs[:, i:i + num_steps]
        Y = Ys[:, i:i + num_steps]
        yield X, Y
```

```python
def download(name, cache_dir=os.path.join('..', 'data')):
    """Download a file inserted into DATA_HUB, return the local filename."""
    assert name in DATA_HUB, f"{name} does not exist in {DATA_HUB}."
    url, sha1_hash = DATA_HUB[name]
    os.makedirs(cache_dir, exist_ok=True)
    fname = os.path.join(cache_dir, url.split('/')[-1])
    if os.path.exists(fname):
        sha1 = hashlib.sha1()
        with open(fname, 'rb') as f:
            while True:
                data = f.read(1048576)
                if not data:
                    break
                sha1.update(data)
        if sha1.hexdigest() == sha1_hash:
            return fname  # Hit cache
    print(f'Downloading {fname} from {url}...')
    r = requests.get(url, stream=True, verify=True)
    with open(fname, 'wb') as f:
        f.write(r.content)
    return fname

def read_time_machine():
    """Load the time machine dataset into a list of text lines."""
    with open(download('time_machine'), 'r') as f:
        lines = f.readlines()
    return [re.sub('[^A-Za-z]+', ' ', line).strip().lower() for line in lines]

def load_corpus_time_machine(max_tokens=-1):
    """Return token indices and the vocabulary of the time machine dataset."""
    lines = read_time_machine()
    tokens = tokenize(lines, 'char')
    vocab = Vocab(tokens)
    # Since each text line in the time machine dataset is not necessarily a
    # sentence or a paragraph, flatten all the text lines into a single list
    corpus = [vocab[token] for line in tokens for token in line]
    if max_tokens > 0:
        corpus = corpus[:max_tokens]
    return corpus, vocab

def load_data_time_machine(batch_size, num_steps, use_random_iter=False,
                           max_tokens=10000):
    """Return the iterator and the vocabulary of the time machine dataset."""
    data_iter = SeqDataLoader(batch_size, num_steps, use_random_iter,
                              max_tokens)
    return data_iter, data_iter.vocab
```

```python
DATA_HUB = dict()
DATA_URL = 'http://d2l-data.s3-accelerate.amazonaws.com/'
DATA_HUB['time_machine'] = (DATA_URL + 'timemachine.txt',
                                '090b5e7e70c295757f55df93cb0a180b9691891a')

batch_size, num_steps = 32, 35
train_iter, vocab = load_data_time_machine(batch_size, num_steps)
```

    Downloading ../data/timemachine.txt from http://d2l-data.s3-accelerate.amazonaws.com/timemachine.txt...

# Model

We fit an unconditional RNN, for language modeling (ie. not vec2seq or seq2seq. Following the D2L notation, the model has the form
$$\begin{align}
H_t &= \phi(X_t W_{xh} + H_{t-1}  W_{hh} + b_h) \\
O_t &= H_t W_{hq} + b_q
\end{align}
$$
where $X_t$ is the $(n,d)$ matrix of (one-hot) inputs
(for batch size $n$ and vocabulary size $d$),
$H_t$ is the $(n,h)$ matrix of hidden states
(for $h$ hidden states),
and $O_t$ is the $(n,q)$ matrix of output logits
(for $q$ output labels, often $q=d$).

```python
# Create the initial parameters
def get_params(vocab_size, num_hiddens, init_rng):
    num_inputs = num_outputs = vocab_size

    def normal(shape, rng):
        return jax.random.normal(rng, shape=shape) * 0.01

    hidden_rng, out_rng = jax.random.split(init_rng)
    # Hidden layer parameters
    W_xh = normal((num_inputs, num_hiddens), hidden_rng)
    W_hh = normal((num_hiddens, num_hiddens), hidden_rng)
    b_h = jnp.zeros(num_hiddens)
    # Output layer parameters
    W_hq = normal((num_hiddens, num_outputs), out_rng)
    b_q = jnp.zeros(num_outputs)
    params = [W_xh, W_hh, b_h, W_hq, b_q]
    return params
```

```python
# Create the initial state
# We assume this is a tuple of one element (later we will use longer tuples)
def init_rnn_state(batch_size, num_hiddens):
    return (jnp.zeros((batch_size, num_hiddens)),)
```

```python
# Forward function.
# Input sequence is (T,B,V), where T is length of the sequence, B is batch size, V is vocab size.
# We iterate over each time step, and process the batch (for that timestep in parallel).
# Output sequence is (T*B, V), since we concatenate all the time steps.
# We also return the final state, so we can process the next subsequence.
@jax.jit
def rnn(params, state, inputs):
    # Here `inputs` shape: (`num_steps`, `batch_size`, `vocab_size`)
    W_xh, W_hh, b_h, W_hq, b_q = params
    H, = state
    outputs = []
    # Shape of `X`: (`batch_size`, `vocab_size`)
    for X in inputs:
        H = jnp.tanh((X @ W_xh) + (H @ W_hh) + b_h)
        Y = H @ W_hq + b_q
        outputs.append(Y)
    return jnp.concatenate(outputs, axis=0), (H,)
```

```python
# Make the model class
# Input X to apply is (B,T) matrix of integers (from vocab encoding).
# We transpose this to (T,B) then one-hot encode to (T,B,V), where V is vocab.
# The result is passed to the forward function.
# (We define the forward function as an argument, so we can change it later.)

class RNNModelScratch:
    """A RNN Model implemented from scratch."""
    def __init__(self, vocab_size, num_hiddens, get_params, init_state,
                 forward_fn):
        self.vocab_size, self.num_hiddens = vocab_size, num_hiddens
        self.init_state, self.get_params = init_state, get_params
        self.forward_fn = forward_fn

    def apply(self, params, state, X):
        X = jax.nn.one_hot(X.T, num_classes=self.vocab_size)
        return self.forward_fn(params, state, X)

    def begin_state(self, batch_size):
        return self.init_state(batch_size, self.num_hiddens)

    def init_params(self, rng):
        return self.get_params(self.vocab_size, self.num_hiddens, rng)
```

```python
num_hiddens = 512
net = RNNModelScratch(len(vocab), num_hiddens, get_params, init_rnn_state, rnn)
```

```python
X = jnp.arange(10).reshape((2, 5)) # batch 2, sequence length is 5
params = net.init_params(rng)
state = net.begin_state(X.shape[0])
print(len(state)) # length 1
print(state[0].shape) # (2,512)

Y, new_state = net.apply(params, state, X)
print(len(vocab)) # 28
print(Y.shape) # (2x5, 28)
```

    1
    (2, 512)
    28
    (10, 28)

# Prediction (generation)

We pass in an initial prefix string, that is not generated. This is used to "warm-up" the hidden state. Specifically, we update the hidden state given the observed prefix, but don't generate anything. After that, for each of the T steps, we compute the (1,V) output tensor, pick the argmax index, and append it to the output. Finally, we convert the to indices to readable token sequence of size (1,T). (Note that this is a greedy, deterministic procedure.)

```python
def predict(prefix, num_preds, net, params, vocab):
    """Generate new characters following the `prefix`."""
    state = net.begin_state(batch_size=1)
    outputs = [vocab[prefix[0]]]
    get_input = lambda: jnp.array([outputs[-1]]).reshape((1, 1))
    for y in prefix[1:]:  # Warm-up period
        _, state = net.apply(params, state, get_input())
        outputs.append(vocab[y])
    for _ in range(num_preds):  # Predict `num_preds` steps
        y, state = net.apply(params, state, get_input())
        y = y.reshape(-1, y.shape[-1])
        outputs.append(int(y.argmax(axis=1).reshape(1)))
    return ''.join([vocab.idx_to_token[i] for i in outputs])
```

```python
# sample 10 characters after the prefix.
# since the model is untrained, the results will be garbage.
predict('time traveller ', 10, net, params, vocab)
```

    'time traveller cncncncncn'

# Training

To ensure the gradient doesn't blow up when doing backpropagation through many layers, we use gradient clipping, which corresponds to the update
$$
g := \min(1, \theta /||g||) g
$$
where $\theta$ is the scaling parameter, and $g$ is the gradient vector.

```python
@jax.jit
def grad_clipping(grads, theta):
    """Clip the gradient."""
    def grad_update(grads):
      return jax.tree_map(lambda g: g * theta / norm, grads)

    norm = jnp.sqrt(sum(jax.tree_util.tree_leaves(jax.tree_map(
        lambda x: jnp.sum(x ** 2), grads))))
    # Update gradient if norm > theta
    # This is jax.jit compatible
    grads = jax.lax.cond(norm > theta, grad_update, lambda g: g, grads)
    return grads
```

The training step is fairly standard, except for the use of gradient clipping, and the issue of the hidden state.
If the data iterator uses random ordering of the sequences, we need to initialize the hidden state for each minibatch. However, if the data iterator uses sequential ordering, we only initialize the hidden state at the very beginning of the process. In the latter case, the hidden state will depend on the value at the previous minibatch.

The state vector may be a tensor or a tuple, depending on what kind of RNN we are using. In addition, the parameter updater can be an optax optimizer, or a simpler custom sgd optimizer.

```python
class Animator:
    """For plotting data in animation."""
    def __init__(self, xlabel=None, ylabel=None, legend=None, xlim=None,
                 ylim=None, xscale='linear', yscale='linear',
                 fmts=('-', 'm--', 'g-.', 'r:'), nrows=1, ncols=1,
                 figsize=(3.5, 2.5)):
        # Incrementally plot multiple lines
        if legend is None:
            legend = []
        display.set_matplotlib_formats('svg')
        self.fig, self.axes = plt.subplots(nrows, ncols, figsize=figsize)
        if nrows * ncols == 1:
            self.axes = [self.axes,]
        # Use a lambda function to capture arguments
        self.config_axes = lambda: set_axes(self.axes[
            0], xlabel, ylabel, xlim, ylim, xscale, yscale, legend)
        self.X, self.Y, self.fmts = None, None, fmts

    def add(self, x, y):
        # Add multiple data points into the figure
        if not hasattr(y, "__len__"):
            y = [y]
        n = len(y)
        if not hasattr(x, "__len__"):
            x = [x] * n
        if not self.X:
            self.X = [[] for _ in range(n)]
        if not self.Y:
            self.Y = [[] for _ in range(n)]
        for i, (a, b) in enumerate(zip(x, y)):
            if a is not None and b is not None:
                self.X[i].append(a)
                self.Y[i].append(b)
        self.axes[0].cla()
        for x, y, fmt in zip(self.X, self.Y, self.fmts):
            self.axes[0].plot(x, y, fmt)
        self.config_axes()
        display.display(self.fig)
        display.clear_output(wait=True)

class Timer:
    """Record multiple running times."""
    def __init__(self):
        self.times = []
        self.start()

    def start(self):
        """Start the timer."""
        self.tik = time.time()

    def stop(self):
        """Stop the timer and record the time in a list."""
        self.times.append(time.time() - self.tik)
        return self.times[-1]

    def avg(self):
        """Return the average time."""
        return sum(self.times) / len(self.times)

    def sum(self):
        """Return the sum of time."""
        return sum(self.times)

    def cumsum(self):
        """Return the accumulated time."""
        return jnp.array(self.times).cumsum().tolist()

class Accumulator:
    """For accumulating sums over `n` variables."""
    def __init__(self, n):
        self.data = [0.0] * n

    def add(self, *args):
        self.data = [a + float(b) for a, b in zip(self.data, args)]

    def reset(self):
        self.data = [0.0] * len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

def set_axes(axes, xlabel, ylabel, xlim, ylim, xscale, yscale, legend):
    """Set the axes for matplotlib."""
    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)
    axes.set_xscale(xscale)
    axes.set_yscale(yscale)
    axes.set_xlim(xlim)
    axes.set_ylim(ylim)
    if legend:
        axes.legend(legend)
    axes.grid()

@jax.jit
def sgd(params, grads, lr, batch_size):
    """Minibatch stochastic gradient descent."""
    params = jax.tree_multimap(lambda p, g: p - lr * g / batch_size, params,
                               grads)
    return params
```

```python
@jax.jit
def train_step(apply_fn, loss_fn, params, state, X, Y):
      def loss(params, state, X, Y):
        y = Y.T.reshape(-1) #(B,T) -> (T,B)
        y_hat, state = apply_fn(params, state, X)
        y_hat = y_hat.reshape(-1, y_hat.shape[-1])
        y_one_hot = jax.nn.one_hot(y, num_classes=y_hat.shape[-1])
        return loss_fn(y_hat, y_one_hot).mean(), state

      grad_fn = jax.value_and_grad(loss, has_aux=True)
      (l, state), grads = grad_fn(params, state, X, Y)
      grads = grad_clipping(grads, 1)
      return l, state, grads
```

```python
def train_epoch(net, params, train_iter, loss, updater, use_random_iter):
    state, timer = None, Timer()
    metric = Accumulator(2)  # Sum of training loss, no. of tokens
    if isinstance(updater, optax.GradientTransformation):
      updater_state = updater.init(params)
    # Convert to jax Partial functions for jax.jit compatibility
    apply_fn = jax.tree_util.Partial(net.apply)
    loss_fn = jax.tree_util.Partial(loss)
    for X, Y in train_iter:
        if state is None or use_random_iter:
            # Initialize `state` when either it is the first iteration or
            # using random sampling
            state = net.begin_state(batch_size=X.shape[0])
        l, state, grads = train_step(apply_fn, loss_fn, params, state, X, Y)
        if isinstance(updater, optax.GradientTransformation):
            updates, updater_state = updater.update(grads, updater_state)
            params = optax.apply_updates(params, updates)
        else:
            # batch_size=1 since the `mean` function has been invoked
            params = updater(params, grads, batch_size=1)
        metric.add(l * Y.size, Y.size)
    return params, math.exp(metric[0] / metric[1]), metric[1] / timer.stop()
```

The main training function is fairly standard.
The loss function is per-symbol cross-entropy, $-\log q(x_t)$, where $q$ is the model prediction from the RNN. Since we compute the average loss across time steps within a batch, we are computing $-\frac{1}{T} \sum_{t=1}^T \log p(x_t|x_{1:t-1})$. The exponential of this is the perplexity (ppl). We plot this metric during training, since it is independent of document length. In addition, we print the MAP sequence prediction following the suffix 'time traveller', to get a sense of what the model is doing.

```python
def train(net, params, train_iter, vocab, lr, num_epochs,
          use_random_iter=False):
    loss = optax.softmax_cross_entropy
    animator = Animator(xlabel='epoch', ylabel='perplexity',
                            legend=['train'], xlim=[10, num_epochs])
    # Initialize
    if isinstance(net, nn.Module):
        updater = optax.sgd(lr)
    else:
        updater = lambda params, grads, batch_size: sgd(params, grads, lr,
                                                        batch_size)
    num_preds = 50
    predict_ = lambda prefix: predict(prefix, num_preds, net, params, vocab)
    # Train and predict
    for epoch in range(num_epochs):
        params, ppl, speed = train_epoch(net, params, train_iter, loss, updater,
                                     use_random_iter)
        if (epoch + 1) % 10 == 0:
            # Prediction takes time on the flax model
            # print(predict_('time traveller'))
            animator.add(epoch + 1, [ppl])
    device = jax.default_backend()
    print(f'perplexity {ppl:.1f}, {speed:.1f} tokens/sec on {device}')
    print(predict_('time traveller'))
    print(predict_('traveller'))
    return params
```

```python
random.seed(0)
num_epochs, lr = 500, 1
num_hiddens = 512
net = RNNModelScratch(len(vocab), num_hiddens, get_params, init_rnn_state, rnn)
params = net.init_params(rng)
params = train(net, params, train_iter, vocab, lr, num_epochs)
```

    perplexity 1.0, 34154.9 tokens/sec on gpu
    time travelleryou can show black is white by argument said filby
    travelleryou can show black is white by argument said filby

```python
num_preds = 100
predict_ = lambda prefix: predict(prefix, num_preds, net, params, vocab)

print(predict_('time traveller'))
print(predict_('the'))
```

    time travelleryou can show black is white by argument said filby but you willnever convince mepossibly not said th
    the got hege tare ve you in a moment weincline to overlook this fact there are really four dimensionsth

```python
num_preds = 500
predict_= lambda prefix: predict(prefix, num_preds, net, params, vocab)

print(predict_('the'))
```

    the got hege tare ve you in a moment weincline to overlook this fact there are really four dimensionsthree which we call the three planes of space and a fourth timethere is however a tendency to draw an unreal distinction betweenthe former three dimensions and the latter bectursat dithe thicg to expect us to begin uponsaid filby an argumentative person with red hairi do not mean to ask you to accept anything without reasonableground for it you will soon admit as much as i need from you youknow of c

# Creating a Flax module

We now show how to use create an RNN as a module, which is faster than our pure Python implementation.

While Flax has cells for more advanced recurrent models, it does not have a basic RNNCell. Therefore, we create an RNNCell similar to those defined in `flax.linen.recurrent` [here](https://flax.readthedocs.io/en/latest/_modules/flax/linen/recurrent.html).

```python
from typing import Any, Callable, Tuple

PRNGKey = Any
Shape = Tuple[int]
Dtype = Any
Array = Any

class RNNCell(nn.recurrent.RNNCellBase):
    """RNN Cell."""
    activation_fn: Callable[..., Any] = nn.activation.tanh
    kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = (
        nn.linear.default_kernel_init)
    recurrent_kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = (
        nn.initializers.orthogonal())
    bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = nn.initializers.zeros

    @nn.compact
    def __call__(self, carry, inputs):
        """RNN Cell.

        Args:
            carry: the hidden state of the RNN cell,
                initialized using `RNNCell.initialize_carry`.
            inputs: an ndarray with the input for the current time step.
                All dimensions except the final are considered batch dimensions.

        Returns:
            A tuple with the new carry and the output.
        """
        h = carry
        hidden_features = h.shape[-1]
        # Dense layer applied to the previous state
        dense_h = functools.partial(nn.Dense,
                          features=hidden_features,
                          use_bias=False,
                          kernel_init=self.recurrent_kernel_init,
                          bias_init=self.bias_init)
        # Dense layer applied to the input, i
        dense_i = functools.partial(nn.Dense,
                          features=hidden_features,
                          use_bias=True,
                          kernel_init=self.kernel_init,
                          bias_init=self.bias_init)
        new_h = self.activation_fn(dense_i()(inputs) + dense_h()(h))
        return new_h, new_h

    @staticmethod
    def initialize_carry(rng, batch_dims, size,
                         init_fn=nn.initializers.zeros):
        """Initialize the RNN cell carry.
        Args:
            rng: random number generator passed to the init_fn.
            batch_dims: a tuple providing the shape of the batch dimensions.
            size: the size or number of features of the memory.
            init_fn: initializer function for the carry.
        Returns:
            An initialized carry for the given RNN cell.
        """
        mem_shape = batch_dims + (size,)
        return init_fn(rng, mem_shape)
```

Now, we create an RNN module to call the RNNCell for each step.

```python
class RNN(nn.Module):
    @functools.partial(
        nn.transforms.scan,
        variable_broadcast='params',
        in_axes=0, out_axes=0,
        split_rngs={'params': False})
    @nn.compact
    def __call__(self, state, x):
        return RNNCell()(state, x)

    @staticmethod
    def initialize_carry(rng, batch_dims, size):
        return RNNCell.initialize_carry(rng, batch_dims, size)
```

```python
num_hiddens = 256
rnn_layer = RNN()

batch_size, num_steps = 32, 35
state = rnn_layer.initialize_carry(rng, (batch_size,), num_hiddens)
state.shape
```

    (32, 256)

Now we update the state with a random one-hot array of inputs.

```python
X = jax.random.normal(rng, shape=(num_steps, batch_size, len(vocab)))
params = rnn_layer.init(rng, state, X)
state_new, Y = rnn_layer.apply(params, state, X)
Y.shape, state_new.shape
```

    ((35, 32, 256), (32, 256))

Now we define our model. It consists of an RNN Layer followed by a dense layer.

```python
class RNNModel(nn.Module):
    """The RNN model."""
    rnn: nn.Module
    vocab_size: int
    num_hiddens: int
    bidirectional: bool = False

    def setup(self):
        # If the RNN is bidirectional (to be introduced later),
        # `num_directions` should be 2, else it should be 1.
        if not self.bidirectional:
            self.num_directions = 1
        else:
            self.num_directions = 2

    @nn.compact
    def __call__(self, state, inputs):
        X = jax.nn.one_hot(inputs.T, num_classes=self.vocab_size)
        state, Y = self.rnn(state, X)
        output = nn.Dense(self.vocab_size)(Y)
        return output, state

    def begin_state(self, batch_size=1):
        # Use fixed random key since default state init fn is just `zeros`.
        return self.rnn.initialize_carry(jax.random.PRNGKey(0), (batch_size,),
                                         num_hiddens)
```

Test the untrained model.

```python
net = RNNModel(rnn=rnn_layer, vocab_size=len(vocab), num_hiddens=num_hiddens)
params = net.init(rng, state, jnp.ones([batch_size, num_steps]))
predict('time traveller', 50, net, params, vocab)
```

    'time travellerfvytfjwjsyijjtjtmfjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjj'

Train it. The results are similar to the 'from scratch' implementation, but much faster.

```python
num_epochs, lr = 500, 1
params = train(net, params, train_iter, vocab, lr, num_epochs)
```

    perplexity 1.3, 42947.2 tokens/sec on gpu
    time traveller fricif ffor d inou with rave yor some malk has i
    traveller after the payser whis the fopmer d on anlithith m


```python
from google.colab import drive
drive.mount('/content/drive')
```

    Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3aietf%3awg%3aoauth%3a2.0%3aoob&response_type=code&scope=email%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdocs.test%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive.photos.readonly%20https%3a%2f%2fwww.googleapis.com%2fauth%2fpeopleapi.readonly

    Enter your authorization code:
    ··········
    Mounted at /content/drive

# Neuromatch Academy: Week 3, Day 1, Tutorial 1
# Real Neurons: The Leaky Integrate-and-Fire (LIF) Neuron Model
__Content creators:__ Qinglong Gu, Songtin Li, John Murray, Richard Naud, Arvind Kumar

__Content reviewers:__  Maryam Vaziri-Pashkam, Ella Batty, Lorenzo Fontolan, Richard Gao, Matthew Krause, Spiros Chavlis, Michael Waskom

---
# Tutorial Objectives
This is Tutorial 1 of a series on implementing realistic neuron models. In this tutorial, we will build up a leaky integrate-and-fire (LIF) neuron model and study its dynamics in response to various types of inputs. In particular, we are going to write a few lines of code to:

- simulate the LIF neuron model

- drive the LIF neuron with external inputs, such as direct currents, Gaussian white noise, and Poisson spike trains, etc.

- study how different inputs affect the  LIF neuron's output (firing rate and spike time irregularity)

Here, we will especially emphasize identifying conditions (input statistics) under which a neuron can spike at low firing rates and in an irregular manner. The reason for focusing on this is that in most cases, neocortical neurons spike in an irregular manner.

---
# Setup

```python
# Imports
import numpy as np
import matplotlib.pyplot as plt
```

```python
# @title Figure Settings
import ipywidgets as widgets  # interactive display
%config InlineBackend.figure_format = 'retina'
# use NMA plot style
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
my_layout = widgets.Layout()
```

```python
# @title Helper functions

def plot_volt_trace(pars, v, sp):
  """
  Plot trajetory of membrane potential for a single neuron

  Expects:
  pars   : parameter dictionary
  v      : volt trajetory
  sp     : spike train

  Returns:
  figure of the membrane potential trajetory for a single neuron
  """

  V_th = pars['V_th']
  dt, range_t = pars['dt'], pars['range_t']
  if sp.size:
    sp_num = (sp / dt).astype(int) - 1
    v[sp_num] += 20  # draw nicer spikes

  plt.plot(pars['range_t'], v, 'b')
  plt.axhline(V_th, 0, 1, color='k', ls='--')
  plt.xlabel('Time (ms)')
  plt.ylabel('V (mV)')
  plt.legend(['Membrane\npotential', r'Threshold V$_{\mathrm{th}}$'],
             loc=[1.05, 0.75])
  plt.ylim([-80, -40])

def plot_GWN(pars, I_GWN):
  """
  Args:
    pars  : parameter dictionary
    I_GWN : Gaussian white noise input

  Returns:
    figure of the gaussian white noise input
  """

  plt.figure(figsize=(12, 4))
  plt.subplot(121)
  plt.plot(pars['range_t'][::3], I_GWN[::3], 'b')
  plt.xlabel('Time (ms)')
  plt.ylabel(r'$I_{GWN}$ (pA)')
  plt.subplot(122)
  plot_volt_trace(pars, v, sp)
  plt.tight_layout()

def my_hists(isi1, isi2, cv1, cv2, sigma1, sigma2):
  """
  Args:
    isi1 : vector with inter-spike intervals
    isi2 : vector with inter-spike intervals
    cv1  : coefficient of variation for isi1
    cv2  : coefficient of variation for isi2

  Returns:
    figure with two histograms, isi1, isi2

  """
  plt.figure(figsize=(11, 4))
  my_bins = np.linspace(10, 30, 20)
  plt.subplot(121)
  plt.hist(isi1, bins=my_bins, color='b', alpha=0.5)
  plt.xlabel('ISI (ms)')
  plt.ylabel('count')
  plt.title(r'$\sigma_{GWN}=$%.1f, CV$_{\mathrm{isi}}$=%.3f' % (sigma1, cv1))

  plt.subplot(122)
  plt.hist(isi2, bins=my_bins, color='b', alpha=0.5)
  plt.xlabel('ISI (ms)')
  plt.ylabel('count')
  plt.title(r'$\sigma_{GWN}=$%.1f, CV$_{\mathrm{isi}}$=%.3f' % (sigma2, cv2))
  plt.tight_layout()
  plt.show()

# the function plot the raster of the Poisson spike train
def my_raster_Poisson(range_t, spike_train, n):
  """
  Generates poisson trains

  Args:
    range_t     : time sequence
    spike_train : binary spike trains, with shape (N, Lt)
    n           : number of Poisson trains plot

  Returns:
    Raster plot of the spike train
  """

  # find the number of all the spike trains
  N = spike_train.shape[0]

  # n should smaller than N:
  if n > N:
    print('The number n exceeds the size of spike trains')
    print('The number n is set to be the size of spike trains')
    n = N

  # plot rater
  i = 0
  while i < n:
    if spike_train[i, :].sum() > 0.:
      t_sp = range_t[spike_train[i, :] > 0.5]  # spike times
      plt.plot(t_sp, i * np.ones(len(t_sp)), 'k|', ms=10, markeredgewidth=2)
    i += 1
  plt.xlim([range_t[0], range_t[-1]])
  plt.ylim([-0.5, n + 0.5])
  plt.xlabel('Time (ms)', fontsize=12)
  plt.ylabel('Neuron ID', fontsize=12)
```

---
# Section 1: The Leaky Integrate-and-Fire (LIF) model

```python
#@title Video 1: Reduced Neuron Models
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='rSExvwCVRYg', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=rSExvwCVRYg

## Implementation of an LIF neuron model

Now, it's your turn to implement one of the simplest mathematical model of a neuron: the leaky integrate-and-fire (LIF) model. The basic idea of LIF neuron was proposed in 1907 by Louis Édouard Lapicque, long before we understood the electrophysiology of a neuron (see a translation of [Lapicque's paper](https://pubmed.ncbi.nlm.nih.gov/17968583/) ). More details of the model can be found in the book [**Theoretical neuroscience**](http://www.gatsby.ucl.ac.uk/~dayan/book/) by Peter Dayan and Laurence F. Abbott.

The subthreshold membrane potential dynamics of a LIF neuron is described by

\begin{eqnarray}
C_m\frac{dV}{dt} = -g_L(V-E_L) + I,\quad (1)
\end{eqnarray}

where $C_m$ is the membrane capacitance, $V$ is the membrane potential, $g_L$ is the leak conductance ($g_L = 1/R$, the inverse of the leak resistance $R$ mentioned in previous tutorials), $E_L$ is the resting potential, and $I$ is the external input current.

Dividing both sides of the above equation by $g_L$ gives

\begin{align}
\tau_m\frac{dV}{dt} = -(V-E_L) + \frac{I}{g_L}\,,\quad (2)
\end{align}

where the $\tau_m$ is membrane time constant and is defined as $\tau_m=C_m/g_L$.

You might wonder why dividing capacitance by conductance gives units of time! Find out yourself why this is true.

Below, we will use Eqn.(2) to simulate LIF neuron dynamics.

If $I$ is sufficiently strong such that $V$ reaches a certain threshold value $V_{\rm th}$, $V$ is reset to a reset potential $V_{\rm reset}< V_{\rm th}$, and voltage is clamped to $V_{\rm reset}$ for $\tau_{\rm ref}$ ms, mimicking the refractoriness of the neuron during an action potential:

\begin{eqnarray}
\mathrm{if}\quad V(t_{\text{sp}})\geq V_{\rm th}&:& V(t)=V_{\rm reset} \text{  for } t\in(t_{\text{sp}}, t_{\text{sp}} + \tau_{\text{ref}}]
\end{eqnarray}
where $t_{\rm sp}$ is the spike time when $V(t)$ just exceeded $V_{\rm th}$.

(__Note__: in the lecture slides, $\theta$ corresponds to the threshold voltage $V_{th}$, and $\Delta$ corresponds to the refractory time $\tau_{\rm ref}$.)

Thus, the LIF model captures the facts that a neuron:
- performs spatial and temporal integration of synaptic inputs
- generates a spike when the voltage reaches a certain threshold
- goes refractory during the action potential
- has a leaky membrane

The LIF model assumes that the spatial and temporal integration of inputs is linear. Also, membrane potential dynamics close to the spike threshold are much slower in LIF neurons than in real neurons.

## Exercise 1: Python code to simulate the LIF neuron

We now write Python code to calculate Eqn. (2) and simulate the LIF neuron dynamics. We will use the Euler method, which you saw in the linear systems case last week,  to numerically integrate Eq 2.

The cell below initializes a dictionary that stores parameters of the LIF neuron model and the simulation scheme. You can use `pars=default_pars(T=simulation_time, dt=time_step)` to get the parameters (you can try to print the dictionary `pars`). Note that, `simulation_time` and `time_step` have the unit `ms`. In addition, you can add the value to a new parameter by `pars['New_param'] = value`.

```python
# @title

# @markdown Execute this code to initialize the default parameters

def default_pars(**kwargs):
  pars = {}

  # typical neuron parameters#
  pars['V_th'] = -55.     # spike threshold [mV]
  pars['V_reset'] = -75.  # reset potential [mV]
  pars['tau_m'] = 10.     # membrane time constant [ms]
  pars['g_L'] = 10.       # leak conductance [nS]
  pars['V_init'] = -75.   # initial potential [mV]
  pars['E_L'] = -75.      # leak reversal potential [mV]
  pars['tref'] = 2.       # refractory time (ms)

  # simulation parameters #
  pars['T'] = 400.  # Total duration of simulation [ms]
  pars['dt'] = .1   # Simulation time step [ms]

  # external parameters if any #
  for k in kwargs:
    pars[k] = kwargs[k]

  pars['range_t'] = np.arange(0, pars['T'], pars['dt'])  # Vector of discretized time points [ms]

  return pars

pars = default_pars()
```

The cell below defines the function to simulate the LIF neuron when receiving external current inputs. You can use `v, sp = run_LIF(pars, Iinj)` to get the membrane potential (`v`) and spike train (`sp`) given the dictionary `pars` and input current `Iinj`.

```python
def run_LIF(pars, Iinj, stop=False):
  """
  Simulate the LIF dynamics with external input current

  Args:
    pars       : parameter dictionary
    Iinj       : input current [pA]. The injected current here can be a value
                 or an array
    stop       : boolean. If True, use a current pulse

  Returns:
    rec_v      : membrane potential
    rec_sp     : spike times
  """

  # Set parameters
  V_th, V_reset = pars['V_th'], pars['V_reset']
  tau_m, g_L = pars['tau_m'], pars['g_L']
  V_init, E_L = pars['V_init'], pars['E_L']
  dt, range_t = pars['dt'], pars['range_t']
  Lt = range_t.size
  tref = pars['tref']

  # Initialize voltage and current
  v = np.zeros(Lt)
  v[0] = V_init
  Iinj = Iinj * np.ones(Lt)
  if stop:  # set end of current to 0 if current pulse
    Iinj[:int(len(Iinj) / 2) - 1000] = 0
    Iinj[int(len(Iinj) / 2) + 1000:] = 0
  tr = 0.  # the count for refractory duration

  # Simulate the LIF dynamics
  rec_spikes = []  # record spike times
  for it in range(Lt - 1):
    if tr > 0:  # check for refractoriness
      v[it] = V_reset
      tr = tr - 1
    elif v[it] >= V_th:  # reset voltage and record spike event
      rec_spikes.append(it)
      v[it] = V_reset
      tr = tref / dt

    ########################################################################
    ## TODO for students: compute the membrane potential v, spike train sp #
    # Fill out function and remove
    #raise NotImplementedError('Student Exercise: calculate the dv/dt and the update step!')
    ########################################################################
    # calculate the increment of the membrane potential
    dv = ( -(v[it]-E_L) + Iinj[it] /g_L)*(dt/tau_m)

    # update the membrane potential
    v[it+1] = v[it] + dv

  rec_spikes = np.array(rec_spikes) * dt

  return v, rec_spikes

pars = default_pars(T=500)
# Uncomment below to test your function
v, sp = run_LIF(pars, Iinj=100, stop=True)
plot_volt_trace(pars, v, sp)
plt.ylim([-80, -60])
plt.show()
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_RealNeurons/solutions/W3D1_Tutorial1_Solution_c779cd88.py)

*Example output:*

---
# Section 2: Response of an LIF model to different types of input currents
In the following section, we will learn how to inject direct current and white noise to study the response of an LIF neuron.

```python
#@title Video 2: Response of the LIF neuron to different inputs
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='preNGdab7Kk', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=preNGdab7Kk

## Section 2.1: Direct current (DC)

### LIF neuron driven by constant current
Investigate the voltage response of the LIF neuron when receiving a DC input of 300 pA by `run_LIF` function.

```python
pars = default_pars(T=100)  # get the parameters

# Run the model to obtain v and sp
v, sp = run_LIF(pars, Iinj=300)

plot_volt_trace(pars, v, sp)
plt.show()
```

In the plot above, you see the membrane potential of an LIF neuron. You may notice that the neuron generates a spike. But this is just a cosmetic spike only for illustration purposes. In an LIF neuron, we only need to keep track of times when the neuron hit the threshold so the postsynaptic neurons can be informed of the spike.

### Interactive Demo: Parameter exploration of DC input amplitude
Here's an interactive demo that shows how the LIF neuron behavior changes for DC input with different amplitudes.

How much DC is needed to reach the threshold (rheobase current)? How does the membrane time constant affect the frequency of the neuron?

```python
# @title
# @markdown Make sure you execute this cell to enable the widget!
my_layout.width = '450px'
@widgets.interact(
    I_dc=widgets.FloatSlider(50., min=0., max=300., step=10.,
                               layout=my_layout),
    tau_m=widgets.FloatSlider(10., min=2., max=20., step=2.,
                                layout=my_layout)
)

def diff_DC(I_dc=200., tau_m=10.):
  pars = default_pars(T=100.)
  pars['tau_m'] = tau_m
  v, sp = run_LIF(pars, Iinj=I_dc)
  plot_volt_trace(pars, v, sp)
  plt.show()
```

    interactive(children=(FloatSlider(value=50.0, description='I_dc', layout=Layout(width='450px'), max=300.0, ste…

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_RealNeurons/solutions/W3D1_Tutorial1_Solution_6babe2e2.py)

## Section 2.2: Gaussian white noise (GWN) current

Given the noisy nature of neuronal activity _in vivo_, neurons usually receive complex, time-varying inputs.

To mimic this, we will now investigate the neuronal response when the LIF neuron receives Gaussian white noise $\xi(t)$ with mean

\begin{eqnarray}
E[\xi(t)]=\mu=0,
\end{eqnarray}

and autocovariance
\begin{eqnarray}
E[\xi(t)\xi(t+\tau)]=\sigma_\xi^2 \delta(\tau)
\end{eqnarray}

Note that the GWN has zero mean, that is, it describes only the fluctuations of the input received by a neuron. We can thus modify our definition of GWN to have a nonzero mean value $\mu$ that equals the DC input, since this is the average input into the cell. The cell below defines the modified gaussian white noise currents with nonzero mean $\mu$.

```python
#@title
#@markdown Execute this cell to get function to generate GWN

def my_GWN(pars, mu, sig, myseed=False):
  """
  Function that generates Gaussian white noise input

  Args:
    pars       : parameter dictionary
    mu         : noise baseline (mean)
    sig        : noise amplitute (standard deviation)
    myseed     : random seed. int or boolean
                 the same seed will give the same
                 random number sequence

  Returns:
    I          : Gaussian white noise input
  """

  # Retrieve simulation parameters
  dt, range_t = pars['dt'], pars['range_t']
  Lt = range_t.size

  # set random seed
  # you can fix the seed of the random number generator so that the results
  # are reliable however, when you want to generate multiple realization
  # make sure that you change the seed for each new realization.
  if myseed:
      np.random.seed(seed=myseed)
  else:
      np.random.seed()

  # generate GWN
  # we divide here by 1000 to convert units to sec.
  I_gwn = mu + sig * np.random.randn(Lt) / np.sqrt(dt / 1000.)

  return I_gwn
```

### Exercise 2: LIF neuron driven by GWN

You can generate a noisy input with `my_GWN(pars, mu, sig, myseed=False)`. Here, $\mu=250$ and $\sigma=5$. Note that fixing the value of the random seed (e.g., `myseed=2020`) will allow you to obtain the same result every time you run this.

```python
pars = default_pars(T=100.)
sig_gwn = 5.
mu_gwn = 250.
#######################################################
## TODO for students: compute the noisy input I_GWN   #
# and then compute v, sp with inject I_GWN into LIF   #
#######################################################

# Calculate the GWN current
I_GWN = my_GWN(pars, mu_gwn, sig_gwn, myseed=2020)

# Run the model and calculate the v and the sp
v, sp = run_LIF(pars, I_GWN)

# Uncomment these lines after obtaining v, sp
plot_GWN(pars, I_GWN)
plt.show()
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_RealNeurons/solutions/W3D1_Tutorial1_Solution_14e10901.py)

*Example output:*

### Interactive Demo: LIF neuron Explorer for noisy input

The mean of the GWN is the amplitude of DC. Indeed, when $\sigma = 0$, GWN is just a DC.

So the question arises how does $\sigma$ of the GWN affect the spiking behavior of the neuron. For instance we may want to know
- how does the minimum input (i.e. $\mu$) needed to make a neuron spike change with increase in $\sigma$
- how does the spike regularity change with increase in $\sigma$

To get an intuition about these questions you can use the following interactive demo that shows how the LIF neuron behavior changes for noisy input with different amplitudes (the mean $\mu$) and fluctuation sizes ($\sigma$).

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

my_layout.width = '450px'
@widgets.interact(
    mu_gwn=widgets.FloatSlider(200., min=100., max=300., step=5.,
                               layout=my_layout),
    sig_gwn=widgets.FloatSlider(2.5, min=0., max=5., step=.5,
                                layout=my_layout)
)

def diff_GWN_to_LIF(mu_gwn, sig_gwn):
  pars = default_pars(T=100.)
  I_GWN = my_GWN(pars, mu=mu_gwn, sig=sig_gwn)
  v, sp = run_LIF(pars, Iinj=I_GWN)
  plt.figure(figsize=(12, 4))
  plt.subplot(121)
  plt.plot(pars['range_t'][::3], I_GWN[::3], 'b')
  plt.xlabel('Time (ms)')
  plt.ylabel(r'$I_{GWN}$ (pA)')
  plt.subplot(122)
  plot_volt_trace(pars, v, sp)
  plt.tight_layout()
  plt.show()
```

    interactive(children=(FloatSlider(value=200.0, description='mu_gwn', layout=Layout(width='450px'), max=300.0, …

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_RealNeurons/solutions/W3D1_Tutorial1_Solution_fb89860a.py)

## Think!
- As we increase the input average ($\mu$) or the input fluctuation ($\sigma$), the spike count changes. How much can we increase the spike count, and what might be the relationship between GWN mean/std or DC value and spike count?

- We have seen above that when we inject DC, the neuron spikes in a regular manner (clock like), and this regularity is reduced when GWN is injected. The question is, how irregular can we make the neurons spiking by changing the parameters of the GWN?

---
# Section 3: Firing rate and spike time irregularity
When we plot the output firing rate as a function of GWN mean or DC value, it is called the input-output transfer function of the neuron (so simply F-I curve).

Spike regularity can be quantified as the **coefficient of variance (CV) of the inter-spike-interval (ISI)**:
\begin{align}
\text{CV}_{\text{ISI}} = \frac{std(\text{ISI})}{mean(\text{ISI})}
\end{align}

A Poisson train is an example of high irregularity, in which $\textbf{CV}_{\textbf{ISI}} \textbf{= 1}$. And for a clocklike (regular) process we have $\textbf{CV}_{\textbf{ISI}} \textbf{= 0}$ because of **std(ISI)=0**.

## Interactive Demo: F-I Explorer for different `sig_gwn`

How does the F-I curve of the LIF neuron change as we increase the $\sigma$ of the GWN? We can already expect that the F-I curve will be stochastic and the results will vary from one trial to another. But will there be any other change compared to the F-I curved measured using DC?

Here's an interactive demo that shows how the F-I curve of a LIF neuron changes for different levels of fluctuation $\sigma$.

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

my_layout.width = '450px'
@widgets.interact(
    sig_gwn=widgets.FloatSlider(3.0, min=0., max=6., step=0.5,
                                layout=my_layout)
)

def diff_std_affect_fI(sig_gwn):
  pars = default_pars(T=1000.)
  I_mean = np.arange(100., 400., 10.)
  spk_count = np.zeros(len(I_mean))
  spk_count_dc = np.zeros(len(I_mean))

  for idx in range(len(I_mean)):
      I_GWN = my_GWN(pars, mu=I_mean[idx], sig=sig_gwn, myseed=2020)
      v, rec_spikes = run_LIF(pars, Iinj=I_GWN)
      v_dc, rec_sp_dc = run_LIF(pars, Iinj=I_mean[idx])
      spk_count[idx] = len(rec_spikes)
      spk_count_dc[idx] = len(rec_sp_dc)

  # Plot the F-I curve i.e. Output firing rate as a function of input mean.
  plt.figure()
  plt.plot(I_mean, spk_count, 'k',
           label=r'$\sigma_{\mathrm{GWN}}=%.2f$' % sig_gwn)
  plt.plot(I_mean, spk_count_dc, 'k--', alpha=0.5, lw=4, dashes=(2, 2),
           label='DC input')
  plt.ylabel('Spike count')
  plt.xlabel('Average injected current (pA)')
  plt.legend(loc='best')
  plt.show()
```

    interactive(children=(FloatSlider(value=3.0, description='sig_gwn', layout=Layout(width='450px'), max=6.0, ste…

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_RealNeurons/solutions/W3D1_Tutorial1_Solution_9a4f261d.py)

### Exercise 3: Compute $CV_{ISI}$ values

As shown above, the F-I curve becomes smoother while increasing the amplitude of the fluctuation ($\sigma$). In addition, the fluctuation can also change the irregularity of the spikes. Let's investigate the effect of $\mu=250$ with $\sigma=0.5$ vs $\sigma=3$.

Fill in the code below to compute ISI, then plot the histogram of the ISI and compute the $CV_{ISI}$. Note that, you can use `np.diff` to calculate ISI.

```python
def isi_cv_LIF(spike_times):
  """
  Calculates the inter-spike intervals (isi) and
  the coefficient of variation (cv) for a given spike_train

  Args:
    spike_times : (n, ) vector with the spike times (ndarray)

  Returns:
    isi         : (n-1,) vector with the inter-spike intervals (ms)
    cv          : coefficient of variation of isi (float)

  """
  ########################################################################
  ## TODO for students: compute the membrane potential v, spike train sp #
  # Fill out function and remove
  #raise NotImplementedError('Student Exercise: calculate the isi and the cv!')
  ########################################################################
  if len(spike_times) >= 2:
    # Compute isi
    isi = np.diff(spike_times)
    # Compute cv
    cv = np.std(isi)/np.mean(isi)
  else:
    isi = np.nan
    cv = np.nan

  return isi, cv

pars = default_pars(T=1000.)
mu_gwn = 250
sig_gwn1 = 0.5
sig_gwn2 = 3.0

I_GWN1 = my_GWN(pars, mu=mu_gwn, sig=sig_gwn1, myseed=2020)
_, sp1 = run_LIF(pars, Iinj=I_GWN1)

I_GWN2 = my_GWN(pars, mu=mu_gwn, sig=sig_gwn2, myseed=2020)
_, sp2 = run_LIF(pars, Iinj=I_GWN2)

# Uncomment to check your function
isi1, cv1 = isi_cv_LIF(sp1)
isi2, cv2 = isi_cv_LIF(sp2)

my_hists(isi1, isi2, cv1, cv2, sig_gwn1, sig_gwn2)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_RealNeurons/solutions/W3D1_Tutorial1_Solution_c9232818.py)

*Example output:*

## Interactive Demo: Spike irregularity explorer for different `sig_gwn`

In the above illustration, we see that the CV of inter-spike-interval (ISI) distribution depends on $\sigma$ of GWN. What about the mean of GWN, should that also affect the CV$_{\rm ISI}$? If yes, how? Does the efficacy of $\sigma$ in increasing the CV$_{\rm ISI}$ depend on $\mu$?

In the following interactive demo, you will examine how different levels of fluctuation $\sigma$ affect the CVs for different average injected currents ($\mu$).

```python
#@title

#@markdown Make sure you execute this cell to enable the widget!

my_layout.width = '450px'
@widgets.interact(
    sig_gwn=widgets.FloatSlider(0.0, min=0., max=10.,
                                step=0.5, layout=my_layout)
)

def diff_std_affect_fI(sig_gwn):
  pars = default_pars(T=1000.)
  I_mean = np.arange(100., 400., 20)
  spk_count = np.zeros(len(I_mean))
  cv_isi = np.empty(len(I_mean))

  for idx in range(len(I_mean)):
    I_GWN = my_GWN(pars, mu=I_mean[idx], sig=sig_gwn)
    v, rec_spikes = run_LIF(pars, Iinj=I_GWN)
    spk_count[idx] = len(rec_spikes)
    if len(rec_spikes) > 3:
      isi = np.diff(rec_spikes)
      cv_isi[idx] = np.std(isi) / np.mean(isi)

  # Plot the F-I curve i.e. Output firing rate as a function of input mean.
  plt.figure()
  plt.plot(I_mean[spk_count > 5], cv_isi[spk_count > 5], 'bo', alpha=0.5)
  plt.xlabel('Average injected current (pA)')
  plt.ylabel(r'Spike irregularity ($\mathrm{CV}_\mathrm{ISI}$)')
  plt.ylim(-0.1, 1.5)
  plt.grid(True)
  plt.show()
```

    interactive(children=(FloatSlider(value=0.0, description='sig_gwn', layout=Layout(width='450px'), max=10.0, st…

### Try to answer the following:
- Does the standard deviation of the injected current affect the F-I curve in any qualitative manner?
- Why does increasing the mean of GWN reduce the $CV_{ISI}$?
- If you plot spike count (or rate) vs. $CV_{ISI}$, should there be a relationship between the two? Try out yourself.

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_RealNeurons/solutions/W3D1_Tutorial1_Solution_7356e407.py)

---
# Section 4: Generation of Poisson type spike trains

*In the next tutorials, we will often use Poisson type spike train to explore properties of neurons and synapses. Therefore, it is good to know how to generate Poisson type spike trains.*

Mathematically, a spike train is a Point Process. One of the simplest models of a sequence of presynaptic pulse inputs is the Poisson process. We know that given temporal integration and refractoriness, neurons cannot behave as a Poisson Process, and Gamma Process gives a better approximation (*find out what might be the difference in the two processes*).
Here, however, we will assume that the incoming spikes are following Poisson statistics. A question arises about how to simulate a Poisson process. The generation of the Poisson process can be realized by at least two following ways:

- By definition, for a Poisson process with rate $\lambda$, the probability of finding one event in the time window with a sufficiently small length $\Delta t$ is $P(N = 1) = \lambda \Delta t$. Therefore, in each time window, we generate a uniformly distributed random variable $r \in [0,1]$ and generate a Poisson event when $r <\lambda \Delta t$. This method allows us to generate Poisson distributed spikes in an online manner.

- The interval $t_{k+1}-t_{k}$ between two Poisson events with rate $\lambda$ follows the exponential distribution, i.e., $P(t_{k+1}-t_{k}<t) = 1 - e^{\lambda t}$. Therefore, we only need to generate a set of exponentially distributed variables $\{s_k\}$ to obtain the timing of Poisson events $t_{k+1}=t_{k}+s_{k}$. In this method, we need to generate all future spikes at once.

Below, we use the first method in a function `Poisson_generator`, which takes arguments `(pars, rate, n, myseed)`.

```python
# @title
# @markdown Execute this cell to get a Poisson_generator function

def Poisson_generator(pars, rate, n, myseed=False):
  """
  Generates poisson trains

  Args:
    pars            : parameter dictionary
    rate            : noise amplitute [Hz]
    n               : number of Poisson trains
    myseed          : random seed. int or boolean

  Returns:
    pre_spike_train : spike train matrix, ith row represents whether
                      there is a spike in ith spike train over time
                      (1 if spike, 0 otherwise)
  """

  # Retrieve simulation parameters
  dt, range_t = pars['dt'], pars['range_t']
  Lt = range_t.size

  # set random seed
  if myseed:
      np.random.seed(seed=myseed)
  else:
      np.random.seed()

  # generate uniformly distributed random variables
  u_rand = np.random.rand(n, Lt)

  # generate Poisson train
  poisson_train = 1. * (u_rand < rate * (dt / 1000.))

  return poisson_train

```

```python
# we can use Poisson_generator to mimic presynaptic spike trains
pars = default_pars()
pre_spike_train = Poisson_generator(pars, rate=10, n=100, myseed=2020)
my_raster_Poisson(pars['range_t'], pre_spike_train, 100)

```

How do we make sure that the above spike trains are following Poisson statistics?
A Poisson process must have the following properties:
- The ratio of the mean and variance of spike count is 1
- Inter-spike-intervals are exponentially distributed
- Spike times are irregular i.e. $CV_{\rm ISI} = 1$
- Adjacent spike intervals are independent of each other.

---
# Summary

Congratulations! You've just built a leaky integrate-and-fire (LIF) neuron model from scratch, and studied its dynamics in response to various types of inputs, having:

- simulated the LIF neuron model

- driven the LIF neuron with external inputs, such as direct current, Gaussian white noise, and Poisson spike trains, etc.

- studied how different inputs affect the  LIF neuron's output (firing rate and spike time irregularity),

with a special focus on low rate and irregular firing regime to mimc real cortical neurons. The next tutorial will look at how spiking statistics may be influenced by a neuron's input statistics.

However, if you have extra time, follow the section below to explore a different type of noise input.

---
# Bonus 1: Orenstein-Uhlenbeck Process

When a neuron receives spiking input, the synaptic current is Shot Noise -- which is a kind of colored noise and the spectrum of the noise determined by the synaptic kernel time constant. That is, a neuron is driven by **colored noise** and not GWN.

We can model colored noise using the Ohrenstein-Uhlenbeck process - filtered white noise.

## Ornstein-Uhlenbeck (OU) current

We next study if the input current is temporally correlated and is modeled as an Ornstein-Uhlenbeck process $\eta(t)$, i.e., low-pass filtered GWN with a time constant $\tau_{\eta}$:

$$\tau_\eta \frac{d}{dt}\eta(t) = \mu-\eta(t) + \sigma_\eta\sqrt{2\tau_\eta}\xi(t).$$

**Hint:** An OU process as defined above has

$$E[\eta(t)]=\mu$$

and autocovariance

$$[\eta(t)\eta(t+\tau)]=\sigma_\eta^2e^{-|t-\tau|/\tau_\eta},$$

which can be used to check your code.

```python
# @title `my_OU(pars, mu, sig, myseed)`
# @markdown Ececute this cell to enable the OU process

def my_OU(pars, mu, sig, myseed=False):
  """
  Function that produces Ornstein-Uhlenbeck input

  Args:
    pars       : parameter dictionary
    sig        : noise amplitute
    myseed     : random seed. int or boolean

  Returns:
    I_ou       : Ornstein-Uhlenbeck input current
  """

  # Retrieve simulation parameters
  dt, range_t = pars['dt'], pars['range_t']
  Lt = range_t.size
  tau_ou = pars['tau_ou']  # [ms]

  # set random seed
  if myseed:
    np.random.seed(seed=myseed)
  else:
    np.random.seed()

  # Initialize
  noise = np.random.randn(Lt)
  I_ou = np.zeros(Lt)
  I_ou[0] = noise[0] * sig

  # generate OU
  for it in range(Lt-1):
    I_ou[it+1] = I_ou[it] + (dt / tau_ou) * (mu - I_ou[it]) + np.sqrt(2 * dt / tau_ou) * sig * noise[it + 1]

  return I_ou
```

### Interactive Demo: LIF Explorer with OU input
In the following, we will check how a neuron responds to a noisy current that follows the statistics of an OU process.

```python
# @title

# @markdown Remember to enable the widget by running the cell!

my_layout.width = '450px'
@widgets.interact(
  tau_ou=widgets.FloatSlider(10.0, min=5., max=20.,
                             step=2.5, layout=my_layout),
  sig_ou=widgets.FloatSlider(10.0, min=5., max=40.,
                             step=2.5, layout=my_layout),
  mu_ou=widgets.FloatSlider(190.0, min=180., max=220.,
                            step=2.5, layout=my_layout)
)

def LIF_with_OU(tau_ou=10., sig_ou=40., mu_ou=200.):
  pars = default_pars(T=1000.)
  pars['tau_ou'] = tau_ou  # [ms]

  I_ou = my_OU(pars, mu_ou, sig_ou)

  v, sp = run_LIF(pars, Iinj=I_ou)

  plt.figure(figsize=(12, 4))
  plt.subplot(121)
  plt.plot(pars['range_t'], I_ou, 'b', lw=1.0)
  plt.xlabel('Time (ms)')
  plt.ylabel(r'$I_{\mathrm{OU}}$ (pA)')
  plt.subplot(122)
  plot_volt_trace(pars, v, sp)
  plt.tight_layout()
  plt.show()
```

## Think!
- How does the OU type input change neuron responsiveness?
- What do you think will happen to the spike pattern and rate if you increased or decreased the time constant of the OU process?

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D1_RealNeurons/solutions/W3D1_Tutorial1_Solution_13568d3a.py)

---
# Bonus 2: Generalized Integrate-and-Fire models

LIF model is not the only abstraction of real neurons. If you want to learn about more realistic types of neuronal models, watch the Bonus Video!

```python
#@title Video 3 (Bonus): Extensions to Integrate-and-Fire models
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='G0b6wLhuQxE', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=G0b6wLhuQxE


# Constraining Cosmological Parameters by Using Type Ia Supernova Data

## 1. Data Preparation

In this notebook, we will use type Ia supernova data from https://sne.space/.

```python
# import necessary modules
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from scipy import integrate, optimize

%matplotlib inline

# set basic figure parameters
import matplotlib as mpl
mpl_param = {'figure.figsize'   : [8.0, 6.0],
             'savefig.dpi'      : 100,
             'axes.titlesize'   : 'xx-large',
             'axes.labelsize'   : 'xx-large',
             'text.usetex'      : False,
             'font.family'      : 'serif'}
mpl.rcParams.update(mpl_param)

# =============================================================================
# Data cleaning
# =============================================================================

# read json data
data = pd.read_json('supernova.json')

# define list of strings to be removed
bad_string = '(host|spectroscopic|heliocentric|cmb|photometric|cluster|,)'

# define list of columns to be removed
bad_column = ['z', 'dL (Mpc)', 'mmax', 'Mmax']

# replace z and dL value with valid float numbers
# nan means Not A Number
for i in bad_column:
    data[i] = data[i].str.replace(bad_string, '')
    data[i].loc[data[i] == ''] = np.nan
    data[i] = data[i].astype(float)

# sort data by redshift value
data = data.sort_values(by=['z'])

# redshift cut until z~2
data = data.loc[data['z'] <= 2]

data.head()
```

    /home/irhamta/anaconda3/lib/python3.6/site-packages/pandas/core/indexing.py:189: SettingWithCopyWarning:
    A value is trying to be set on a copy of a slice from a DataFrame

    See the caveats in the documentation: http://pandas.pydata.org/pandas-docs/stable/indexing.html#indexing-view-versus-copy
      self._setitem_with_indexer(indexer, value)

      Airmass
      Aliases
      Alt. (°)
      Azi. (°)
      Data
      Dec.
      Disc. Date
      Discoverer
      E(B-V)
      ...
      Radio
      References
      Spec.
      Type
      Vsky
      X-ray
      dL (Mpc)
      mmax
      v☉
      z

      244

      1.180
      AAVSO 1221+15
      57.897
      187.928

      +15:03:17.1 15.05475°
      1984/01/04 MJD: 45703.5
      Kimeridze, Rosino,Kimeridze,
      0.0284
      ...
      0
      1986Afz....25..279K1989A&A...220...83B2012A&A....

      Ia-HV Ia,I
      4.679

      17.00000
      12.10
      -256 host
      -0.000920

      2273

      -15.862
      74.072

      -01:07:08.9 -1.11914°
      2006/10/02 MJD: 54010.5

      0.1783
      ...

      2014arXiv1401.3317S2012ApJ...755...61S2008AJ.....

      Ia

      1061.80000
      20.87
      56200
      0.000000

      11076

      LMCSNR J053402-695503, MCSNR J0534-6955
      -49.212
      210.072

      -69:55:03 -69.91750°

      0.8733
      ...

      2016A&A...585A.162M2016A&A...594A..13P2011ApJ....

      Ia

      0.05001
      NaN
      3.388
      0.000011

      11075

      LMCSNR J051935-690209, MCSNR J0519-6902
      -50.415
      211.568

      -69:02:09 -69.03583°

      0.7637
      ...

      2016A&A...585A.162M2016A&A...594A..13P2011ApJ....

      Ia

      0.05001
      NaN
      3.388
      0.000011

      11078

      LMCSNR J053418-703326, MCSNR J0534-7033
      -49.242
      209.093

      -70:33:26 -70.55722°

      0.2237
      ...

      2016A&A...585A.162M2016A&A...594A..13P2011ApJ....

      Ia

      0.05001
      NaN
      3.388
      0.000011

5 rows × 31 columns

```python
# plot redshift vs distance modulus
plt.plot(data['z'], data['dL (Mpc)'], 'b.')
plt.xlabel('redshift $(z)$')
plt.ylabel('Luminosity Distance (MPC)')
```

## 2. Luminosity Distance

Luminosity distance ($d_L$) for various cosmological models are usually formulated as

\begin{equation}
    d_L(z) = \frac{c (1+z_2)}{H_0}
        \times \int_{z_1}^{z_2}{
            \frac{dz'}{\sqrt{
                \Omega_{\rm m}(1+z')^3 + \Omega_{\rm r}(1+z')^4 + \Omega_{\rm de}(1+z')^{3(1+w_{\rm de})
               }}}}
\end{equation}

Here, we consider a spatially flat universe ($\Omega_{\rm m} + \Omega_\Lambda = 1$, $w_{\rm de} = -1$). Therefore, that we can write it as following codes.

```python
# define some basic constants
c = 3 * 1e5 # km/s
H0 = 67.7 #km / (Mpc s)
Omega_m = 0.307
Omega_r = 0 * 1e-5 # too small
Omega_lambda = 1 - Omega_m

# make luminosity distance function
def lum_dist (z, Omega_m, Omega_lambda, H0):

    Omega_r = 0 * 1e-5 # too small

    # integration part
    # integration is calculated from redshift=0 to redshift=z
    fn = lambda z: (Omega_r*(1+z)**4. \
                    + Omega_m*(1+z)**3 \
                    + Omega_lambda\
                   )**-0.5

    # return array values
    return c*(1+z)/H0 * np.asarray([integrate.quad(fn, 0, _z)[0] for _z in z])
```

For a quick insight, we can do the least-square fitting to obtain cosmological parameters.

```python
# remove NaN values
data_good = data[['z', 'dL (Mpc)']].dropna()
data_good = data_good.sample(n=500)

# guess initial parameters
# Omega_m, Omega_lambda, H0, respectively
initial_param = np.array([0.3, 0.7, 70])

# least-square fitting
opt_param, cov = optimize.curve_fit(lum_dist,
                                    data_good['z'].values,
                                    data_good['dL (Mpc)'].values,
                                    p0=initial_param)
err_param = np.sqrt(np.diag(cov))
```

    /home/irhamta/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:15: RuntimeWarning: invalid value encountered in double_scalars
      from ipykernel import kernelapp as app
    /home/irhamta/anaconda3/lib/python3.6/site-packages/scipy/integrate/quadpack.py:385: IntegrationWarning: The occurrence of roundoff error is detected, which prevents
      the requested tolerance from being achieved.  The error may be
      underestimated.
      warnings.warn(msg, IntegrationWarning)

```python
# =============================================================================
# Plot the result
# =============================================================================

plt.figure()
plt.plot(data['z'], data['dL (Mpc)'], 'b.', label='Data')
plt.plot(data['z'], lum_dist(data['z'], *opt_param),
         'g-', label='Fitted')
plt.ylabel('Distance (Mpc)')
plt.xlabel('Redshift')
plt.legend()
plt.show()

print ('======================================')
print ('Least-Square Fitting Final parameter:')
print ('Omega_m = %.2f (%f)' %(opt_param[0], err_param[0]))
print ('Omega_lambda = %.2f (%f)' %(opt_param[1], err_param[1]))
print ('H0 = %.2f (%f)' %(opt_param[2], err_param[2]))
print ('======================================')
```

## 3. Maximum Likelihood Fitting

We will use maximization of the likelihood function to constrain and compare the models. First, we calculate $d_{L}^{\rm obs}$ and $d_{L}^{\rm th}$. Then, the join likelihood function for all parameters, based on a flat Bayesian prior, is

\begin{equation}
    \mathcal{L} = \prod_{i} \exp{\bigg[-\frac{\chi^2}{2} \bigg]}
\end{equation}

where for each measurement

\begin{equation}
    \chi^2_i = \frac{(d_{L_i}^{\rm obs} - d_{L_i}^{\rm th})^2}{\sigma^2_{d_{L_i}}}
\end{equation}

In python, those equations can be written as below.

```python
# =============================================================================
# Maximum likelihood fitting
# =============================================================================

# define likelihood function as in Equation 11 in Leaf et al. (2018)
def lnlike(theta, X, y, yerr):

    Omega_m, Omega_lambda, H0 = theta
    model = lum_dist(X, Omega_m, Omega_lambda, H0)

    # chi-square
    chi2 = ((y-model)**2)/yerr**2

    return np.sum( np.exp(-chi2/2) )

X = data_good['z'].values
y = data_good['dL (Mpc)'].values
yerr = 0.05 # assuming there is uniform error about 5%
```

```python
from scipy import optimize

# optimize module minimizes functions whereas we would like to maximize the likelihood
# that's why I put the minus(-) sign
nll = lambda *args: -lnlike(*args)
#result = optimize.minimize(nll, initial_param, args=(X, y, yerr))
#m_ml, b_ml, h0_ml, wde_ml = result["x"]

print ('======================================')
#print ('Maximum Likelihood Result')
#print ('Omega_m = %.2f (%.2f)' %(m_ml, 0))
#print ('Omega_lambda = %.2f (%.2f)' %(b_ml, 0))
#print ('H0 = %.2f (%.2f)' %(h0_ml, 0))
#print ('wde = %.2f (%.2f)' %(wde_ml, 0))
print ('======================================\n')
```

    ======================================
    ======================================

## 4. Markov Chain Monte Carlo Fitting

In order to improve our fitting accuracy, we will use emcee, which is pure-Python implementation of Goodman & Weare's Affine Invariant Markov chain Monte Carlo (MCMC) Ensemble sampler. These pages (http://emcee.readthedocs.io/en/stable/) will show you how to use it. We would like to marginalize over some "nuisance parameters" and find an estimate of the posterior probability function (the distribution of parameters that is consistent with dataset) for others.

We will define uniform (so-called "uninformative") priors on initial parameters ($\Omega_{\rm m}$, $\Omega_{\Lambda}$, and $H_0$). Then, combining this with the definition of lnlike from above, the full log-probability function can be calculated.

```python
# =============================================================================
# MCMC fitting
# see http://dfm.io/emcee/current/user/line/ for the detail
# =============================================================================

# define prior
def lnprior(theta):
    Omega_m, Omega_lambda, H0 = theta

    if 0.95 <= Omega_m + Omega_lambda <= 1.05 \
    and 0 < Omega_m < 1.5 \
    and 0 < Omega_lambda < 1.5 \
    and 60 < H0 < 80:
        return 0
    return -np.inf

# define the full probability
def lnprob(theta, X, y, yerr):
    lp = lnprior(theta)
    if not np.isfinite(lp):
        return -np.inf
    return lp + lnlike(theta, X, y, yerr)
```

```python
ndim, nwalkers = 3, 1000

pos = [initial_param + 1e-2*np.random.randn(ndim) for i in range(nwalkers)]

import emcee
import sys

sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(X, y, yerr), threads=3)

nsteps = 500
width = 30

print ('running MCMC.....')
for i, result in enumerate(sampler.sample(pos, iterations=nsteps)):
    n = int((width+1) * float(i) / nsteps)
    sys.stdout.write("\r[{0}{1}]".format('#' * n, ' ' * (width - n)))
sys.stdout.write("\n")

samples = sampler.chain[:, 50:, :].reshape((-1, ndim))
```

    running MCMC.....
    [##############################]

```python
import corner
fig = corner.corner(samples, labels=["$\Omega_m$", "$\Omega_\Lambda$", "$H_0$"],
                      truths=[Omega_m, Omega_lambda, H0])

plt.savefig('result/supernova.png')
plt.show()

```

```python
m_mcmc, b_mcmc, h0_mcmc = map(lambda v: (v[1], v[2]-v[1], v[1]-v[0]),
                                   zip(*np.percentile(samples, [16, 50, 84],
                                                      axis=0)))

print ('============================================================================')
print ('MCMC Result')
print ('Omega_m = ', m_mcmc)
print ('Omega_lambda = ', b_mcmc)
print ('H0 = ', h0_mcmc)
print ('============================================================================')

output_data = pd.DataFrame({'omega_m': samples[:, 0],
                            'omega_l': samples[:, 1],
                            'h0'     : samples[:, 2]})
output_data.to_csv('result/output_supernova.csv', index=False)
```

    ============================================================================
    MCMC Result
    Omega_m =  (0.30166214003237946, 0.010601759941790578, 0.006677495582416815)
    Omega_lambda =  (0.6749308524546039, 0.021825491761950255, 0.012705856920557035)
    H0 =  (68.6734042618426, 0.6701585025228525, 0.6150923190433701)
    ============================================================================

```python

```


# Autograd (3)：梯度下降法解量子化学 RHF 自洽场能量

> 创建时间：2019-12-17

这一篇文档将会简单地介绍量子化学计算中经常使用的 RHF 方法，在自动求导下的一种应用。

RHF (Restricted Hartree-Fock) 方法是量化计算的基本方法；对于更高精度的分子体系或晶体等周期性体系的计算，譬如 DFT、Post-HF 方法如 CC 和 CI，通常都是以 HF 方法为原型开发。但 RHF 方法的求解方程通常都是通过自洽场方法 (self-consistent field, SCF) 迭代计算；一种经典的迭代计算算法是 DIIS (见 [简单理解 SCF 中的 DIIS](../../QC_Notes/SCF_Series/diis_comprehen.ipynb))。关于 SCF 求解 RHF 的入门文档可以参考 Szabo, Ostlund [^Szabo-Ostlund.Dover.1996]。关于自洽场求解 RHF 的程序相关问题，可以参考 [pyxdh 文档](https://py-xdh.readthedocs.io/zh_CN/latest/qcbasic/basic_rhf.html)。

而事实上，更经典但现在不太用于 RHF 方程求解的方法是将其化为局域凸优化问题。尽管通常这类凸优化问题属于类牛顿法的范畴，但如果将 Hessian 矩阵始终视为单位矩阵的常数倍，这就化为了普通的梯度下降问题，也恰好是自动求导可以解决的问题范畴。这篇文档就用 PyTorch 简单地实现这样一个功能。

在继续这篇文档之前，我们先引入程序库：

```python
import torch
import math
import numpy as np
import scipy
from pyscf import gto, scf
import warnings

torch.set_printoptions(precision=5, sci_mode=False, linewidth=120)
np.set_printoptions(precision=5, suppress=True, linewidth=120)
warnings.filterwarnings("ignore")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
```

:::{attention}

这篇文档是找了一个作者稍微熟悉的优化问题，用自动求导来解决。但这不是通常处理这类问题的做法，况且并不高效。

同时需要指出的是，一般来说 RHF 并不是一个凸优化问题；但在比较合理的初猜下，可以将问题 **近似地** 看成凸优化问题。后面尽管默认了凸优化这个条件，但许多现实的 RHF 收敛问题实际上来源于非凸优化的复杂性。作者也不是搞优化问题的，因此很多术语可能用得不对。

因此，这篇文档是 **娱乐向** 的文档。

:::

## 分子体系的定义

### 分子体系能量

我们使用的分子体系是不对称的双氧水分子，基组为 6-31G，用 PySCF 进行计算。分子体系的总能量为 `energy_tot` $E_\mathrm{tot}$ -150.5850338 Hartree。这篇文档的目标就是计算得到这个数值。

```python
mol = gto.Mole()
mol.atom = """
O  0.0  0.0  0.0
O  0.0  0.0  1.5
H  1.0  0.0  0.0
H  0.0  0.7  1.0
"""
mol.basis = "6-31G"
mol.verbose = 0
mol.build()
```

    <pyscf.gto.mole.Mole at 0x7fa30ea25160>

```python
scf_eng = scf.RHF(mol).run()
energy_tot = scf_eng.e_tot
energy_tot
```

    -150.5850337808369

### 恒定物理量与矩阵

这里列举的是随分子不同而不同的量。这些量可能计算耗时小，也可能是计算过程中的重要构成部分；但它们与凸优化问题的变量无关，在这篇文档中姑且称为恒定量。

分子体系总能量中，一部分能量是无需通过量化方法，而只需要经典物理的库伦排斥力计算得到。这部分能量是原子核排斥能 `energy_nuc` $E_\mathrm{nuc}$ 37.8846744 Hartree。为了简化叙述，我们就直接借用了 PySCF 的程序的结果。

```python
energy_nuc = scf_eng.energy_nuc()
energy_nuc
```

    37.88467440864127

下面是原子轨道角标的矩阵或张量。无电子相互作用的 Hamiltonian 积分 (Hamiltonian Core 积分) $h_{\mu \nu}$ 记为 `H`，重叠积分 $S_{\mu \nu}$ 记为 `S`，双电子积分 $(\mu \nu | \kappa \lambda)$ 记为 `eri_ao`。

```python
T = mol.intor("int1e_kin")
Vnuc = mol.intor("int1e_nuc")
H = T + Vnuc
S = mol.intor("int1e_ovlp")
eri_ao = mol.intor("int2e")
```

上述矩阵或张量都是通过 PySCF 给出的；PySCF 使用的默认引擎是 numpy。我们需要将它们转为 PyTorch 引擎的矩阵或张量。

```python
H = torch.tensor(H, device=device)
S = torch.tensor(S, device=device)
eri_ao = torch.tensor(eri_ao, device=device)
```

### 凸优化自变量与密度矩阵

在凸优化问题中，RHF 问题可以看作是对下述原子轨道空间下 $\mathbf{X}$ 矩阵 (或写作矩阵元的形式，$X_{\mu \nu}$) 作为自变量的函数优化问题 (Helgaker et al. [^Helgaker-Jorgensen.Wiley.2013], eq 10.7.69，但这篇文档对密度矩阵 $R_{\mu \nu}$ 的定义恰好是 Helgaker 课本中的两倍)：

$$
E_\mathrm{tot} [\mathbf{X}] = \mathrm{tr} (\mathbf{R}[\mathbf{X}]) \mathbf{h} + \frac{1}{4} \mathrm{tr} (\mathbf{R}[\mathbf{X}] \mathbf{G} \mathbf{R}[\mathbf{X}]) + E_\mathrm{nuc}
$$

其中，根据 Helgaker, eq 10.7.32，有

$$
\mathbf{R}[\mathbf{X}] = \exp (- \mathbf{X} \mathbf{S}) \mathbf{R}_0 \exp (\mathbf{S} \mathbf{X})
$$

关于这些公式的意义与计算过程，我们后面还会有更详细的描述。这里需要留意的是，$\mathbf{G}$ 是一个四维张量；其它地方的记号与矩阵乘法没有区别。

我们现在需要知道的是，其一，$E_\mathrm{tot} [\mathbf{X}]$ 是我们的凸优化目标；对应于统计或机器学习中的概念，就类似于损失函数 (loss function)。其二，尽管 $\mathbf{X}$ 是我们的学习目标，但真正有物理意义的、可以被实验观测的量是 $\mathbf{R}[\mathbf{X}]$；该量的意义是电子云密度 (简称密度)。换一种说法的话，RHF 方法作为量子化学方法的核心目标是给出基态的电子云密度；有了电子云密度之后，许多其它的分子性质是可以从中导出的。因此，对于 RHF 方法而言，其电子云密度的价值近乎等价于作为 Schrodinger 方程的求解目的的波函数。

方才我们已经用 PySCF 计算了分子的 RHF 的能量。PySCF 也可以给出分子基态下的 RHF 密度 `D` $D_{\mu \nu}$：

```python
D = scf_eng.make_rdm1()
D = torch.tensor(D, device=device)
```

对于 RHF 问题而言，我们会说，$\mathbf{X}$ 取到函数 $E_\mathrm{tot} [\mathbf{X}]$ 处于最小值的极小值几乎等价于问题

$$
\mathbf{D} = \mathbf{R} [\mathbf{X}]
$$

我们等下会验证这个结论。

需要补充的是，PySCF 给出电子密度的方式与上述过程完全不相同。在一般的自洽场迭代过程中，电子态密度是由分子轨道系数得到的：

$$
\mathbf{D} = 2 \mathbf{C}_\mathrm{occ} \mathbf{C}_\mathrm{occ}^\mathrm{T}
$$

或用 Einstein Summation Convention 的语言，

$$
D_\mathrm{\mu \nu} = 2 C_{\mu i} C_{\nu i}
$$

这与 $\mathbf{D} = \mathbf{R} [\mathbf{X}]$ 看似是完全无关的公式，但至少从结果上，两者确实相同。

## 损失函数的定义

### PyTorch 矩阵幂运算

这一段只是讨论一个技术细节。尽管 PyTorch 功能确实很强大，但一个比较尴尬的地方是 PyTorch 不支持不少需要一些数学技巧才能保证效率与稳定性的算法。其中一个算法是矩阵幂 (作者 Bing 了一刻钟发现实在找不到，于是作出了断言，但高 zi 贵 bi 的作者是不会提 issue 或者 stackoverflow 的 >.>)。

为了解决这样一个量化的问题，我们无奈需要自己造轮子，目标是只要能用就行。

下面的程序 `calc_exp` 是根据矩阵幂的定义给出的：

$$
\exp(\mathbf{X}) = \sum_{k=0}^{K} \frac{\mathbf{X}^k}{k!}
$$

其中用到了矩阵指数运算 $\mathbf{X}^k$ 的函数 `mat_power`；该函数通过递归定义。

我们知道，这种函数的求和上界应当是 $K \rightarrow \infty$；但实际的程序不能允许这种操作，必须让 $K$ 作有界的截断。这种有界的截断的判据是，若通过截断计算得到的 $\exp(\mathbf{X})$ 和 numpy 给出的不截断的 $\exp(\mathbf{X})$ 近乎相等 (`np.allclose` 的判标下)，那么就允许截断。

```python
def mat_power(X, order):
    if order > 0:
        return X @ mat_power(X, order - 1)
    return torch.eye(X.shape[0], dtype=X.dtype, device=X.device)

def calc_exp(X, debug=False):
    target = scipy.linalg.expm(X.cpu().clone().detach().numpy())
    result = mat_power(X, 0) + mat_power(X, 1)
    order = 1
    while not np.allclose(result.cpu().clone().detach().numpy(), target):
        order += 1
        result += mat_power(X, order) / math.factorial(order)
        if debug: print(np.linalg.norm(result.cpu().clone().detach().numpy() - target))
    if debug: print("Current order is", order)
    return result
```

读者可以尝试几个矩阵幂运算的例子。但需要留意，这个程序的鲁棒性显然是不高的；遇到上述定义下 $\mathbf{X}$ 超出收敛域的情形下，这个程序就失效了。

### 分子能量函数 $E_\mathrm{tot} [\mathbf{X}]$

下面我们回到分子能量函数 (损失函数) $E_\mathrm{tot} [\mathbf{X}]$ 的定义中。事实上让程序来做这件事其实是相当容易的，只需要几行代码就能搞定：

```python
def calc_RX(X, S, R0):
    return calc_exp(-X @ S) @ R0 @ calc_exp(S @ X)
```

```python
def calc_EX(X, S, R0, H, eri_ao, energy_nuc):
    RX = calc_RX(X, S, R0)
    return (
        torch.einsum("uv, uv ->", RX, H)
        + 0.5 * torch.einsum("uv, uvkl, kl ->", RX, eri_ao, RX)
        - 0.25 * torch.einsum("uv, ukvl, kl ->", RX, eri_ao, RX)
        + energy_nuc
    )
```

但我们需要对其中的过程作一些说明。首先是关于密度的导出式 `calc_RX`：

$$
\mathbf{R}[\mathbf{X}] = \exp (- \mathbf{X} \mathbf{S}) \mathbf{R}_0 \exp (\mathbf{S} \mathbf{X})
$$

这里的 $\mathbf{S}$ 是重叠积分。$\mathbf{R}_0$ 原则上是任意的、满足下述三式的矩阵 (Helgaker, eq 10.7.25-27)：

$$
\begin{align}
\mathbf{R}_0^\mathrm{T} &= \mathbf{R}_0 \\
\mathrm{tr} (\mathbf{R}_0 \mathbf{S}) &= N \\
\mathbf{R}_0 \mathbf{S} \mathbf{R}_0 &= \mathbf{R}_0
\end{align}
$$

上述三式对应的性质分别称为对称性、电子数守恒、基于重叠积分 $\mathbf{S}$ 的幂等性。其中，$N$ 代表体系的电子数。

但实际上，从数值稳定性的角度上讲，我们不希望让 $\mathbf{X}$ 太大以至于超出收敛域；因此 $\mathbf{R}_0$ 总是需要尽可能接近真实密度 $\mathbf{D}$。在这篇文档中，我们称 $\mathbf{R}_0$ 为密度初猜。

同时，$\mathbf{X}$ 是一种旋转矩阵；它必须要满足反对称性，即 $\mathbf{X}^\mathrm{T} = - \mathbf{X}$。这是相当重要且有用的性质，但我们这里不作更多分析。

而对于能量函数 `calc_EX`：

$$
E_\mathrm{tot} [\mathbf{X}] = \mathrm{tr} (\mathbf{R}[\mathbf{X}]) \mathbf{h} + \frac{1}{4} \mathrm{tr} (\mathbf{R}[\mathbf{X}] \mathbf{G} \mathbf{R}[\mathbf{X}]) + E_\mathrm{nuc}
$$

由于其中的四脚标张量的存在稍有麻烦，在程序编写便利、以及借用强大的 `einsum` 功能的角度上，我们采用 Einstein Summation Convention 可以将上式写作

$$
E_\mathrm{tot} [\mathbf{X}] = R_{\mu \nu} [\mathbf{X}] h_{\mu \nu} + \frac{1}{2} R_{\mu \nu} [\mathbf{X}] (\mu \nu | \kappa \lambda) R_{\kappa \lambda} [\mathbf{X}] - \frac{1}{4} R_{\mu \nu} [\mathbf{X}] (\mu \kappa | \nu \lambda) R_{\kappa \lambda} [\mathbf{X}] + E_\mathrm{nuc}
$$

各个矩阵元都已经在上文中被定义了，因此程序的书写也并不困难。

### 最小值的极小点分析

这一段我们分析极小点附近的性质。我们之前提到过，$\mathbf{R}_0$ 总是需要尽可能接近真实密度 $\mathbf{D}$。如果这两者相等是什么情形？

现在我们假定 $E_\mathrm{tot} [\mathbf{X}]$ 已经取到了处于最小值的极小点。我们说起过这种情况下与 $\mathbf{D} = \mathbf{R} [\mathbf{X}]$ 几乎等价。但我们并没有验证过这个结论。下面我们就来验证之。

若同时有 $\mathbf{R}_0 = \mathbf{D}$ 与 $\mathbf{D} = \mathbf{R} [\mathbf{X}]$，那么

$$
\exp (- \mathbf{X} \mathbf{S}) \mathbf{D} \exp (\mathbf{S} \mathbf{X}) = \mathbf{D}
$$

一个显然满足上式的 $\mathbf{X}$ 是零矩阵。在这种情况下，我们应当会期待分子体系总能量 `energy_autograd` 就是使用现成量化软件给出的总能量 `energy_tot`：

```python
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = D.clone()
```

```python
energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
float(energy_autograd)
```

    -150.58503378083833

并且，我们会期待 $\mathbf{X}$ 确实取到了极小值，通过自动求导得到的导数是不会发生任何变化的零矩阵：

```python
energy_autograd.backward()
gX = X.grad
```

```python
torch.norm(gX)
```

    tensor(    0.00000, device='cuda:0', dtype=torch.float64)

从这个意义上，我们可以说，这一套工作流程是基本合理的。

### 真实密度初猜的分析

但上面一段是假设我们已经知道体系密度的理想情况。现实情况是，我们并不知道体系的真正密度，从而需要求自变量 $\mathbf{X}$ 的结果。

现实中有许多给出密度初猜 $\mathbf{R}_0$ 的方式。其中最简单的密度初猜是将 Fock 矩阵当作 Hamiltonian Core 矩阵，进行对角化后得到的密度矩阵。在 PySCF 中，它通过 `init_guess_by_1e` 方法给出。同时，初始的自变量 (或称旋转矩阵) $\mathbf{X}$ 定为零矩阵。

```python
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = torch.tensor(scf_eng.init_guess_by_1e(), device=S.device)
```

以这个初始密度代入能量的计算公式中，会得到一个比刚才的 -150 Hartree 高出不少的能量：

```python
energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
energy_autograd
```

    tensor(-140.54132, device='cuda:0', dtype=torch.float64, grad_fn=<AddBackward0>)

得到这样的结果也是情理之中：这完全满足变分原理的要求，即任意非基态的波函数 (在 RHF 中，也可以称为其对应的基态密度 $\mathbf{D}$)，那么其给出的能量会高于最低能量。这也满足优化问题的要求，即任何偏离作为最小值的极小点附近的变量 $\mathbf{X}$，一定给出高于最小值的结果来。

下面我们来考虑关于 $\mathbf{X}$ 的梯度。首先，由于 $\mathbf{X}$ 在最小值附近，我们假设其附近是凸性的，因此其梯度必然不为零：

```python
energy_autograd.backward()
gX = X.grad
torch.norm(gX)
```

    tensor(6.70115, device='cuda:0', dtype=torch.float64)

其次，我们说起过，$\mathbf{X}$ 具有反对称的性质。如果让能量函数 $E_\mathrm{tot} [\mathbf{X}]$ 作反向传播，那么 $\mathbf{X}$ 的梯度也应当是满足这种反对称性质。

```python
torch.allclose(-gX.T, gX)
```

    True

## RHF 的梯度下降求取

### 学习率递降 class `Scheduler`

这其实是一个非常细节的问题，但却需要不少代码来解决。

我们都知道在深度学习中，梯度下降算法与学习率递降都是非常有必要的。但由于我们需要在训练的过程中始终保证作为自变量的 $\mathbf{X}$ 具有反对称性，因此它不能像 PyTorch 深度网络中参数一样，使用较为自由的优化方式；而是需要在优化 (机器学习中类似于训练) 过程中始终重新将 $\mathbf{X}$ 反对称化。

因此，这里暂时不使用 PyTorch 提供的 optim 类的 `Adam` 或 `ReduceLROnPlateau` 类来解决问题，而需要手动造与学习率有关的轮子；至于梯度下降法只能用 Naive 的版本，而不使用 Adam 等更高级的版本。至于是不是一定不能用 PyTorch 提供的自动化工具，作者暂时还不清楚。

```python
class Sheduler:

    def __init__(self, init_lr=0.05, min_lr=1e-6, factor=0.95, patience=3, stop_train=30, stop_thresh=1e-6, debug=True):
        self.lr = init_lr
        self.min_lr = min_lr
        self.factor = factor
        self.patience = patience
        self.stop_train = stop_train
        self.stop_thresh = stop_thresh
        self.debug = debug

        self.loss = 1e10
        self.epoch = 0

        self.flag_min = False
        self.end_train = False
        self.loss_list = []

    def step(self, loss, epoch_id=-1):
        self.epoch += 1
        self.loss_list.append(float(loss))
        if loss < self.loss:
            self.loss = loss
            self.epoch = 0
        if not self.flag_min:
            if self.epoch >= self.patience:
                self.lr *= self.factor
                self.epoch = 0
                print("learning rate decrease to {:7.4f} on epoch {:5d}.".format(self.lr, epoch)) if self.debug else None
            if self.lr < self.min_lr:
                self.flag_min = True
                print("Hit minimum learning rate.") if self.debug else None
        if self.flag_min:
            if self.epoch >= self.stop_train:
                print("End learning at epoch {:5d} since {:5d} epoch loss is larger than lowest loss.".format(epoch, self.stop_train)) if self.debug else None
                self.end_train = True
        if np.all(np.abs(np.array(self.loss_list[-self.stop_train:]) - float(self.loss)) < self.stop_thresh) and len(self.loss_list) > 30:
            print("End learning at epoch {:5d} since last {:5d} epoch loss is allclose to lowest loss".format(epoch, self.stop_train) +
                  " within threshold {:7.2e}.".format(self.stop_thresh)) if self.debug else None
            self.end_train = True
```

这个类的优化过程稍微有些一点点复杂 (但比起现成的程序已经简单太多)，一些简单的说明是：

- 优化过程的学习率从 `init_lr` (默认 0.05) 渐渐降低为 `min_lr` (默认 1e-6)，每次降低 `factor` (默认 0.95) 倍，但降到最低后就不再下降；
- 学习率降低的判据是，如果连续 `patience` (默认 3) 次的损失函数大于最低的损失函数值，就降低学习率
- 上述过程比通常深度学习中的要求高很多；这是因为通常的深度学习不是凸优化问题，在损失函数势能面上游走不见得是坏事；但 RHF 问题我们假设在合理的密度初猜下是凸优化，损失函数不应该经常上升；
- 停止优化的第一判据是，若学习率 `self.lr` 降到最低后连续 `stop_train` (默认 30) 次损失函数的值大于最低的那一次，那么停止优化；
- 停止优化的第二判据是，若最后连续 `stop_train` (默认 30) 次损失函数都与最低的那一次之间差不超过 `stop_thresh` (默认 1e-6) Hartree 大小，那么停止优化。

### 优化过程

有了上面文字的铺垫，后面的优化过程若读者有任何一点深度学习的脚本经验，应该是非常容易理解的。这里用到的初始条件是

- 密度初猜 $\mathbf{R}_0$：采用 Hamiltonian Core 作为 Fock 矩阵对角化所得密度
- 初始自变量 $\mathbf{X}$：采用零矩阵

```python
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = torch.tensor(scf_eng.init_guess_by_1e(), device=S.device)
```

随后学习率从 0.04 开始递降，进行损失函数 (即能量函数 $E_\mathrm{tot} [\mathbf{X}]$) 的优化。

```python
sheduler = Sheduler(init_lr=0.04)

for epoch in range(0, 5000):
    # Calculate loss function
    energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
    energy_autograd.backward()
    # Make optimization step
    t = X - X.grad * sheduler.lr
    # Force anti-symmetrize X
    X = ((t - t.T) / 2).clone().detach().requires_grad_(True)
    # Update learning rate or stop training
    sheduler.step(energy_autograd, epoch)
    if sheduler.end_train:
        print("End at epoch {:5d}.".format(epoch))
        break
print(float(energy_autograd))
```

    learning rate decrease to  0.0380 on epoch     7.
    learning rate decrease to  0.0361 on epoch    10.
    learning rate decrease to  0.0343 on epoch   182.
    learning rate decrease to  0.0326 on epoch   196.
    End learning at epoch   859 since last    30 epoch loss is allclose to lowest loss within threshold 1.00e-06.
    End at epoch   859.
    -150.5850332855298

这就完成了借助于自动求导的 RHF 的能量计算问题！我们可以拿它与 PySCF 作为量化程序计算得到的能量作比较：

```python
energy_tot
```

    -150.5850337808369

## 小结

我们在这份文档中简单地讨论了在给定电子积分的情况下，使用 PyTorch 的自动求导功能，以 Naive 梯度下降法给出 RHF 能量。这样的程序既可以在 CPU 下运行，也能在 GPU 下运行。

简单的代码总结可以是，在定义了与量子化学问题无关的矩阵指数 `mat_power`、矩阵幂 `calc_exp`、学习率递降管理器 `Sheduler` 后，我们可以进行 RHF 的能量计算，其代码也非常简单。

**分子定义**

```python
mol = gto.Mole()
mol.atom = """
O  0.0  0.0  0.0
O  0.0  0.0  1.5
H  1.0  0.0  0.0
H  0.0  0.7  1.0
"""
mol.basis = "6-31G"
mol.verbose = 0
mol.build()
```

    <pyscf.gto.mole.Mole at 0x7fa2800cb908>

**电子积分定义**

```python
H = torch.tensor(mol.intor("int1e_kin") + mol.intor("int1e_nuc"), device=device)
S = torch.tensor(mol.intor("int1e_ovlp"), device=device)
eri_ao = torch.tensor(mol.intor("int2e"), device=device)
```

**核库伦排斥能定义**

```python
natm = mol.natm
Z_A, A_t = mol.atom_charges(), mol.atom_coords()
r_AB = np.linalg.norm(A_t[:, None, :] - A_t[None, :, :], axis=-1)
r_AB += np.diag(np.ones(natm) * np.inf)
energy_nuc = 0.5 * (Z_A[None, :] * Z_A[:, None] / r_AB).sum()
```

**密度矩阵与能量损失函数定义**

```python
def calc_RX(X, S, R0):
    return calc_exp(-X @ S) @ R0 @ calc_exp(S @ X)
```

```python
def calc_EX(X, S, R0, H, eri_ao, energy_nuc):
    RX = calc_RX(X, S, R0)
    return (
        torch.einsum("uv, uv ->", RX, H)
        + 0.5 * torch.einsum("uv, uvkl, kl ->", RX, eri_ao, RX)
        - 0.25 * torch.einsum("uv, ukvl, kl ->", RX, eri_ao, RX)
        + energy_nuc
    )
```

**密度初猜与自变量定义**

```python
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = torch.tensor(scf_eng.init_guess_by_1e(), device=S.device)
```

**梯度下降得到 RHF 能量**

```python
sheduler = Sheduler(init_lr=0.04, debug=False)

for epoch in range(0, 5000):
    energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
    energy_autograd.backward()
    t = X - X.grad * sheduler.lr
    X = ((t - t.T) / 2).clone().detach().requires_grad_(True)
    sheduler.step(energy_autograd, epoch)
    if sheduler.end_train: break
```

**最终能量**

```python
float(energy_autograd)
```

    -150.5850332855298

[^Szabo-Ostlund.Dover.1996]: Szabo, A.; Ostlund, N. S. *Modern Quantum Chemistry: Introduction to Advanced Electronic Structure Theory*; Dover Publications, 1996.

[^Helgaker-Jorgensen.Wiley.2013]: Helgaker, T.; Olsen, J.; Jorgensen, P. *Molecular Electronic-Structure Theory*; Wiley-Blackwell, 2013.