Implement Generator In Rust Using Genawaiter Crate

Iterator

In programming languages, an Iterator refers to a type that sequentially provides a specific sequence. For example, the following code, which provides values in the given array (or vector) in order of index, is also using an iterator.

// Rust code
let list = vec![0, 1, 2, 3];
for v in list{
    // v = 0, 1, 2, 3
}

# Python code
list = [0, 1, 2, 3]
for v in list:
  # v = 0, 1, 2, 3

Similar to Python, in Rust, any struct can be used as an Iterator by defining the Iterator trait. The Iterator trait has many functions, but when defining a new struct with the Iterator trait, only the type of the value that the iterator outputs, Item, and the function next that returns the next value need to be defined, and the rest is automatically completed.

impl Iterator for MyStruct{
  type Item;

  fn next(&mut self) -> Option<Self::Item>
}

Generator

There are many ways to create an Iterator, but there is a method that is similar to yield in other languages but not in stable Rust. This is a Generator. A Generator is an iterator that sequentially returns values derived from a specific operation. Although the Iterator trait provided by default can only be used if there is a way to infer the ‘next state’ from the ‘current state’, Generators can output more diverse values because they can hold all previous information.

The following python code defines a generator that outputs the Fibonacci numbers from 0 to num-1.

def getFibonnaciSeries(num):
    c1, c2 = 0, 1
    count = 0
    while count < num:
        yield c1
        c3 = c1 + c2
        c1 = c2
        c2 = c3
        count += 1
for i in getFibonnaciSeries(7):
    print(i)

Of course, in Rust, we can provide the same functionality by storing c1, c2, count in a struct.

struct FibIterator{
  c1 : usize,
  c2 : usize,
  count : usize,
}

impl FibIterator{
  fn new() -> Self{
    Self{ c1 : 0, c2 : 1, count : 0 }
  }
}

impl Iterator for FibIterator{
  type Item = usize;

  fn next(&mut self) -> Option<Self::Item>{
    if count == num{
      return None;
    } else {
      let temp = c1;
      c2 = c1 + c2;
      c1 = c2;
      count += 1;
      Some(temp);
    }
  }
}

However, due to the disadvantage that the number of fields in the iterator increases as the generator’s structure becomes more complex, it is not possible to change all generators into iterators of other structures. However, Rust does not provide generators in stable mode. So, without using generators, we often use various other methods to implement similar functionality. I couldn’t figure out exactly why generator and yield are not used in stable mode, but I think it might be because of memory safety issues. Generators are also known as “resumable functions”, meaning that when a function is running and returns a value, it stops, and when the function is “resumed”, it starts from the calculation after it stopped. In other words, generators can be understood as a structure that stops and waits to be called again until it is called again from the outside. If a generator uses variables defined outside, the variable becomes a ‘shared mutable reference’, and since Rust strictly limits this case, it is not used.

So, in the current Rust version, generators only exist as nightly features, and to use them in stable mode, we need to use the genawaiter crate. genawaiter crate implements generator functionality using await / async. In this document, I will explain how to simplify the Hamiltonian configuration using genawaiter in the crate that uses generators to diagonalize the Hamiltonian of a spin chain.

System instruction

Spin states

First, the spin states of the particles in the system can be either up spin or down spin in the $z$ direction, and if up spin is represented by 1 and down spin by 0, the state of a system with $N$ particles can be represented by an $N$-bit number. \begin{equation} \ket{26} = \ket{11010} \end{equation} Also, the spin of the particles can be represented by bit flip.

Hamiltonian

Next, let’s look at the Hamiltonian of the spin-XXZ model to understand the Hamiltonian configuration. \begin{equation} H_0 = - \frac{J}{2} \sum_{n=0}^{L - 1}[\sigma_n^x \sigma^x_{n+1} + \sigma_n^y \sigma^y_{n+1} + \Delta_0 \sigma_n^z \sigma^z_{n+1}] \end{equation} Here, $\sigma^x_n, \sigma^y_n, \sigma^z_n$ represent the $x, y, z$ spin of the $n$-th particle, respectively. $\sigma^x_n$ and $\sigma^y_n$ are written using the raising and lowering operators $\sigma_n^{\pm} = (\sigma_n^x \pm i \sigma_n^y) / 2$ as follows. \begin{equation} H_0 = -J\sum_{n=0}^{L - 1} [\sigma_n^+ \sigma_{n+1}^- + \sigma_n^- \sigma_{n+1}^+ + \frac{\Delta_0}{2} \sigma_n^z \sigma^z_{n+1} ] \end{equation} This means that the first two terms mean that there is an energy of $-J$ when the $z$ spin of the $n$-th particle and the $n+1$-th particle are different, and the quantum state of the system changes to the state where the spin of the $n$-th particle and the $n+1$-th particle are exchanged by the Hamiltonian. The last term means that if the $z$ spin of two adjacent particles is the same, $J \Delta_0 / 2$ is added, and if it is different, $J \Delta_0 / 2$ is added, and in this case, the quantum state of the system does not change. If we explain it using bit, the state where 3 out of 5 particles are up spin $\ket{11010}$ changes as follows when the Hamiltonian is applied.

\begin{align} H_0 \ket{11010} = &- \frac{3}{2} \Delta_0 J\ket{11010} - J\ket{10110} -J \ket{11100} \nonumber \newline & -J \ket{11001} -J \ket{01011} \end{align}

In the matrix representation of $H_0$, the $i,j$ element is given by $\bra{i}H_0\ket{j}$, and to obtain this, it is necessary to output $i$ for which $H_{0, ij}$ is not zero when the number $j$ is given. This is the point where we want to use generators.

Generator implementation

Let’s create a generator that outputs $i$ and $H_{0, ij}$ values continuously when the number $j$ is given.

// src/hamiltonian/mod.rs

use genawaiter::{sync::gen, yield_};

pub struct PeriodicNearestXXZ{
    pub delta_x : f64,
    pub delta_z : f64,
}

impl PeriodicNearestXXZ{
    pub fn new(delta_x : f64, delta_z : f64) -> Self{
        Self{
            delta_x,
            delta_z,
        }
    }

    pub fn apply_to<'a, S>(&'a self, state : &'a S) -> impl Iterator<Item = (usize, f64)> + 'a
        where S : State {
        gen!({
            let num = state.rep();
            let mut sum = 0f64;
            for ((i, si), (j, sj)) in state.periodic_pair_enumerate(){
                if si == sj{
                    sum -= self.delta_z / 2f64;
                } else {
                    sum += self.delta_z / 2f64;
                    let flipped = bit_flip_unsafe(num, i, j);
                    yield_!((flipped, -self.delta_x));
                }
            }

            yield_!((num, sum));
        }).into_iter()
    }
}

Let’s take a look at each part.

use genawaiter::{sync::gen, yield_};

We will use the gen macro and yield_ macro from the genawaiter crate.

pub struct PeriodicNearestXXZ{
    pub delta_x : f64,
    pub delta_z : f64,
}

First, we define a struct to store the properties of the Hamiltonian. Although we only need delta_z for the XXZ model, we added delta_x and delta_z just in case.

pub fn new(delta_x : f64, delta_z : f64) -> Self{
    Self{
        delta_x,
        delta_z,
    }
}

This is a function to create a struct, and

pub fn apply_to<'a, S>(&'a self, state : &'a S) -> impl Iterator<Item = (usize, f64)> + 'a
    where S : State{
    // Something
}

This is the function to apply the Hamiltonian to a specific state. In the crate, there is a State trait, and we add State to the trait bound to allow any struct to be the target. The more important thing is lifetime. We will only receive references to self and state, and since the generator’s nature is that the iterator is not immediately consumed, the lifetime of the reference to self and state needs to be longer than the lifetime of the iterator being output. To do this, we specify the lifetime.

gen!({
    // Something
}).into_iter()

The gen macro converts a code block into a generator, and the into_iter() function at the end converts it into an iterator in the form that is convenient to use in our code.

let num = state.rep();
let mut sum = 0f64;
for ((i, si), (j, sj)) in state.periodic_pair_enumerate(){
    if si == sj{
        sum -= self.delta_z / 2f64;
    } else {
        sum += self.delta_z / 2f64;
        let flipped = bit_flip_unsafe(num, i, j);
        yield_!((flipped, -self.delta_x));
    }
}
yield_!((num, sum));

In the code block, we receive the representation number of state into num, and iterate over the bit numbers. The iterator that performs this function is periodic_pair_enumerate(). It outputs the $i$-th bit $s_i$ and the $j$-th bit $s_j$, but since the bits are adjacent anyway, the reason for dividing by $i,j$ is that the system is periodic, so $(L, 0)$ pair also exists. If $s_i$ and $s_j$ are different, we yield the bit flipped number and the energy value. And the value corresponding to the diagonal term is added up and finally the representation number of the state itself and sum are returned and the iterator ends.

In this way, the code can be used as follows in the actual calculation code.

let xxz = PeriodicNearestXXZ::new(1f64, delta);
let mut hamiltonian : Array2<Complex64> = Array2::zeros((n, n));

for (rep, value) in xxz.apply_to(state){
    hamiltonian[[rep, state.rep()]] += value;
}

Of course, if the code block inside the generator is not very complicated, it may not be necessary to use the generator. However, since it is also written while studying, and it is easier to read the code in a simpler context than the iterator, it may be useful.