# Scale bijectors and LinearOperator

This reading is an introduction to scale bijectors, as well as the `LinearOperator` class, which can be used with them.

```python
!pip install tensorflow=='2.2.0'
```

    Collecting tensorflow==2.2.0
      Downloading tensorflow-2.2.0-cp37-cp37m-manylinux2010_x86_64.whl (516.2 MB)
         |████████████████████████████████| 516.2 MB 3.8 kB/s
    Requirement already satisfied: astunparse==1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (1.6.3)
    Collecting tensorflow-estimator<2.3.0,>=2.2.0
      Downloading tensorflow_estimator-2.2.0-py2.py3-none-any.whl (454 kB)
         |████████████████████████████████| 454 kB 50.5 MB/s
    Requirement already satisfied: protobuf>=3.8.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (3.17.3)
    Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (0.37.1)
    Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (1.1.0)
    Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (1.15.0)
    Requirement already satisfied: numpy<2.0,>=1.16.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (1.19.5)
    Requirement already satisfied: google-pasta>=0.1.8 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (0.2.0)
    Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (3.3.0)
    Collecting tensorboard<2.3.0,>=2.2.0
      Downloading tensorboard-2.2.2-py3-none-any.whl (3.0 MB)
         |████████████████████████████████| 3.0 MB 43.2 MB/s
    Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (1.13.3)
    Requirement already satisfied: grpcio>=1.8.6 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (1.43.0)
    Collecting h5py<2.11.0,>=2.10.0
      Downloading h5py-2.10.0-cp37-cp37m-manylinux1_x86_64.whl (2.9 MB)
         |████████████████████████████████| 2.9 MB 25.9 MB/s
    Requirement already satisfied: keras-preprocessing>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (1.1.2)
    Requirement already satisfied: scipy==1.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (1.4.1)
    Requirement already satisfied: absl-py>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow==2.2.0) (1.0.0)
    Collecting gast==0.3.3
      Downloading gast-0.3.3-py2.py3-none-any.whl (9.7 kB)
    Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (1.8.1)
    Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (2.23.0)
    Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (1.35.0)
    Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (0.4.6)
    Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (57.4.0)
    Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (1.0.1)
    Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (3.3.6)
    Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (4.2.4)
    Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (4.8)
    Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (0.2.8)
    Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (1.3.1)
    Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (4.10.1)
    Requirement already satisfied: typing-extensions>=3.6.4 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (3.10.0.2)
    Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (3.7.0)
    Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (0.4.8)
    Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (3.0.4)
    Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (1.24.3)
    Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (2.10)
    Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (2021.10.8)
    Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.3.0,>=2.2.0->tensorflow==2.2.0) (3.2.0)
    Installing collected packages: tensorflow-estimator, tensorboard, h5py, gast, tensorflow
      Attempting uninstall: tensorflow-estimator
        Found existing installation: tensorflow-estimator 2.7.0
        Uninstalling tensorflow-estimator-2.7.0:
          Successfully uninstalled tensorflow-estimator-2.7.0
      Attempting uninstall: tensorboard
        Found existing installation: tensorboard 2.7.0
        Uninstalling tensorboard-2.7.0:
          Successfully uninstalled tensorboard-2.7.0
      Attempting uninstall: h5py
        Found existing installation: h5py 3.1.0
        Uninstalling h5py-3.1.0:
          Successfully uninstalled h5py-3.1.0
      Attempting uninstall: gast
        Found existing installation: gast 0.4.0
        Uninstalling gast-0.4.0:
          Successfully uninstalled gast-0.4.0
      Attempting uninstall: tensorflow
        Found existing installation: tensorflow 2.7.0
        Uninstalling tensorflow-2.7.0:
          Successfully uninstalled tensorflow-2.7.0
    Successfully installed gast-0.3.3 h5py-2.10.0 tensorboard-2.2.2 tensorflow-2.2.0 tensorflow-estimator-2.2.0

```python
!pip install tensorflow_probability=='0.10.0'
```

    Collecting tensorflow_probability==0.10.0
      Downloading tensorflow_probability-0.10.0-py2.py3-none-any.whl (3.5 MB)
         |████████████████████████████████| 3.5 MB 5.2 MB/s
    Requirement already satisfied: gast>=0.3.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow_probability==0.10.0) (0.3.3)
    Requirement already satisfied: numpy>=1.13.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow_probability==0.10.0) (1.19.5)
    Requirement already satisfied: cloudpickle>=1.2.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow_probability==0.10.0) (1.3.0)
    Requirement already satisfied: six>=1.10.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow_probability==0.10.0) (1.15.0)
    Requirement already satisfied: decorator in /usr/local/lib/python3.7/dist-packages (from tensorflow_probability==0.10.0) (4.4.2)
    Installing collected packages: tensorflow-probability
      Attempting uninstall: tensorflow-probability
        Found existing installation: tensorflow-probability 0.15.0
        Uninstalling tensorflow-probability-0.15.0:
          Successfully uninstalled tensorflow-probability-0.15.0
    Successfully installed tensorflow-probability-0.10.0

```python
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
print("TF version:", tf.__version__)
print("TFP version:", tfp.__version__)

from IPython.display import Image
```

    TF version: 2.2.0
    TFP version: 0.10.0

## Introduction

You have now seen how bijectors can be used to transform tensors and tensor spaces. Until now, you've only seen this in the scalar case, where the bijector acts on a single value. When the tensors you fed into the bijectors had multiple components, the bijector acted on each component individually by applying batch operations to scalar values. For probability distributions, this corresponds to a scalar event space.

However, bijectors can also act on higher-dimensional space. You've seen, for example, the multivariate normal distribution, for which samples are tensors with more than one component. You'll need higher-dimensional bijectors to work with such distributions. In this reading, you'll see how bijectors can be used to generalise scale transformations to higher dimensions. You'll also see the `LinearOperator` class, which you can use to construct highly general scale bijectors. In this reading, you'll walk through the code, and we'll use figure examples to demonstrate these transformations.

This reading contains many images, as this allows you to visualise how a space is transformed. For this reason, the examples are limited to two dimensions, since these allow easy plots. However, these ideas generalise naturally to higher dimensions. Let's start by creating a point that is randomly distributed across the unit square $[0, 1] \times [0, 1]$:

```python
# Create the base distribution and a single sample

uniform = tfd.Uniform(low=[0.0, 0.0], high=[1.0, 1.0], name='uniform2d')
x = uniform.sample()
x
```

    <tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.4135183, 0.5548992], dtype=float32)>

We will be applying linear transformations to this data. To get a feel for how these transformations work, we show ten example sample points, and plot them, as well as the domain of the underlying distribution:

```python
# Run this cell to download and view a figure to show example data points

!wget -q -O x.png --no-check-certificate "https://docs.google.com/uc?export=download&id=1DLqzh7xcjM7BS3C_QmgeF1xET2sXgMG0"
Image("x.png", width=500)
```

Each of the ten points is hence represented by a two-dimensional vector. Let $\mathbf{x} = [x_1, x_2]^T$ be one of these points. Then scale bijectors are linear transformations of $\mathbf{x}$, which can be represented by a $2 \times 2$ matrix $B$. The forward bijection to $\mathbf{y} = [y_1, y_2]^T$ is

$$
\mathbf{y}
=
\begin{bmatrix}
y_1 \\ y_2
\end{bmatrix}
= B \mathbf{x}
= \begin{bmatrix}
b_{11} & b_{12} \\
b_{21} & b_{22} \\
\end{bmatrix}
\begin{bmatrix}
x_1 \\ x_2
\end{bmatrix}
$$

This is important to remember: any two-dimensional scale bijector can be represented by a $2 \times 2$ matrix. For this reason, we'll sometimes use the term "matrix" to refer to the bijector itself. You'll be seeing how these points and domain are transformed under different bijectors  in two dimensions.

## The `ScaleMatvec` bijectors

### The `ScaleMatvecDiag` bijector

We'll start with a simple scale bijector created using the `ScaleMatvecDiag` class:

```python
# Create the ScaleMatvecDiag bijector

bijector = tfb.ScaleMatvecDiag(scale_diag=[1.5, -0.5])
```

which creates a bijector represented by the diagonal matrix
$$ B =
\begin{bmatrix}
1.5 & 0 \\
0 & -0.5 \\
\end{bmatrix}.
$$

We can apply this to the data using `y = bijector(x)` for each of the ten points. This transforms the data as follows:

```python
# Run this cell to download and view a figure to illustrate the transformation

!wget -q -O diag.png --no-check-certificate "https://docs.google.com/uc?export=download&id=1sgfZ_Qzdd2v7CErP2zIk04p6R6hUW7RR"
Image("diag.png", width=500)
```

You can see what happened here: the first coordinate is multiplied by 1.5 while the second is multipled by -0.5, flipping it through the horizontal axis.

```python
# Apply the bijector to the sample point

y = bijector(x)
y
```

    <tf.Tensor: shape=(2,), dtype=float32, numpy=array([ 0.62027746, -0.2774496 ], dtype=float32)>

### The `ScaleMatvecTriL` bijector

In the previous example, the bijector matrix was diagonal, which essentially performs an independent scale operation on each of the two dimensions. The domain under the bijection remains rectangular. However, not all scale tarnsformations have to be like this. With a non-diagonal matrix, the domain will transform to a quadrilateral. One way to do this is by using the `tfb.ScaleMatvecTriL` class, which implements a bijection based on a lower-triangular matrix. For example, to implement the lower-triangular matrix
$$ B =
\begin{bmatrix}
-1 & 0 \\
-1 & -1 \\
\end{bmatrix}
$$
you can use the `tfb.ScaleMatvecTriL` bijector as follows:

```python
# Create the ScaleMatvecTriL bijector

bijector = tfb.ScaleMatvecTriL(scale_tril=[[-1., 0.],
                                           [-1., -1.]])
```

```python
# Apply the bijector to the sample x

y = bijector(x)
y
```

    <tf.Tensor: shape=(2,), dtype=float32, numpy=array([-0.4135183, -0.9684175], dtype=float32)>

A graphical overview of this change is:

```python
# Run this cell to download and view a figure to illustrate the transformation

!wget -q -O lower_triangular.png --no-check-certificate "https://docs.google.com/uc?export=download&id=1eMYwPzMVpmt1FYscplu7RRn1S4gmFo5B"
Image("lower_triangular.png", width=500)
```

## Inverse and composition

Scale transformations always map the point $[0, 0]$ to itself and are only one particular class of bijectors. As you saw before, you can create more complicated bijections by composing one with another. This works just like you would expect. For example, you can compose a scale transformation with a shift to the left (by one unit) as follows:

```python
# Create a scale and shift bijector

scale_bijector = tfb.ScaleMatvecTriL(scale_tril=[[-1., 0.],
                                                 [-1., -1.]])
shift_bijector = tfb.Shift([-1., 0.])
bijector = shift_bijector(scale_bijector)
```

```python
# Apply the bijector to the sample x

y = bijector(x)
y
```

    <tf.Tensor: shape=(2,), dtype=float32, numpy=array([-1.4135183, -0.9684175], dtype=float32)>

which has the expected result:

```python
# Run this cell to download and view a figure to illustrate the transformation

!wget -q -O scale_and_shift.png --no-check-certificate "https://docs.google.com/uc?export=download&id=1iucwJlG2ropvJOkRfBMgEpuFNpYa_JH6"
Image("scale_and_shift.png", width=500)
```

Furthermore, bijectors are always invertible (with just a few special cases, see e.g. [`Absolute Value`](https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/AbsoluteValue)), and these scale transformations are no exception. For example, running

```python
# Apply the inverse transformation to the image of x

bijector = tfb.ScaleMatvecTriL(scale_tril=[[-1., 0.],
                                           [-1., -1.]])
y = bijector.inverse(bijector(x))
```

recovers `x`:

```python
# Run this cell to download and view a figure to illustrate the transformation

!wget -q -O inverse.png --no-check-certificate "https://docs.google.com/uc?export=download&id=1CHCkSfz6EnOYiZaw6vGZ_s6BzyP1NK1X"
Image("inverse.png", width=500)
```

so that the original and transformed data is the same.

```python
# Check that all y and x values are the same

tf.reduce_all(y == x)
```

    <tf.Tensor: shape=(), dtype=bool, numpy=True>

## The `LinearOperator` class and `ScaleMatvecLinearOperator` bijector

The examples you just saw used the `ScaleMatvecDiag` and `ScaleMatvecTriL` bijectors, whose transformations can be represented by diagonal and lower-triangular matrices respectively. These are convenient since it's easy to check whether such matrices are invertible (a requirement for a bijector). However, this comes at a cost of generality: there are acceptable bijectors whose matrices are not diagonal or lower-triangular. To construct these more general bijectors, you can use the `ScaleMatvecLinearOperator` class, which operates on instances of `tf.linalg.LinearOperator`.

The `LinearOperator` is a class that allows the creation and manipulation of linear operators in TensorFlow. It's rare to call the class directly, but its subclasses represent many of the common linear operators. It's programmed in a way to have computational advantages when working with big linear operators, although we won't discuss these here. What matters now is that we can use these linear operators to define bijectors using the `ScaleMatvecLinearOperator` class. Let's see how this works.

### The `LinearOperatorDiag` class

First, let's use this framework to recreate our first bijector, represented by the diagonal matrix

$$ B =
\begin{bmatrix}
1.5 & 0 \\
0 & -0.5 \\
\end{bmatrix}.
$$

You can do this using the `ScaleMatvecLinearOperator` as follows. First, we'll create the linear operator that represents the scale transformation using

```python
scale = tf.linalg.LinearOperatorDiag(diag=[1.5, -0.5])
```

where `LinearOperatorDiag` is one of the subclasses of `LinearOperator`. As the name suggests, it implements a diagonal matrix. We then use this to create the bijector using the `tfb.ScaleMatvecLinearOperator`:

```python
# Create the ScaleMatvecLinearOperator bijector

bijector = tfb.ScaleMatvecLinearOperator(scale)
```

This bijector is the same as the first one above:

```python
# Run this cell to download and view a figure to illustrate the transformation

!wget -q -O linear_operator_diag.png --no-check-certificate "https://docs.google.com/uc?export=download&id=1KaCJl28Thp6NjxspG3pq251vDJrmDd97"
Image("linear_operator_diag.png", width=500)
```

```python
# Apply the bijector to the sample x

y = bijector(x)
y
```

    <tf.Tensor: shape=(2,), dtype=float32, numpy=array([ 0.62027746, -0.2774496 ], dtype=float32)>

### The `LinearOperatorFullMatrix` class

We can also use this framework to create a bijector represented by a custom matrix. Suppose we have the matrix

$$ B =
\begin{bmatrix}
0.5 & 1.5 \\
1.5 & 0.5 \\
\end{bmatrix}
$$

which is neither diagonal nor lower-triangular. We can implement a bijector for it using the `ScaleMatvecLinearOperator` class by using another subclass of `LinearOperator`, namely the `LinearOperatorFullMatrix`, as follows:

```python
# Create a ScaleMatvecLinearOperator bijector

B = [[0.5, 1.5],
     [1.5, 0.5]]
scale = tf.linalg.LinearOperatorFullMatrix(matrix=B)
bijector = tfb.ScaleMatvecLinearOperator(scale)
```

which leads to the following transformation:

```python
# Run this cell to download and view a figure to illustrate the transformation

!wget -q -O linear_operator_full.png --no-check-certificate "https://docs.google.com/uc?export=download&id=1Zk5lp7-VTwmX5r0yPAqVGGzWIgYTjJIJ"
Image("linear_operator_full.png", width=500)
```

```python
# Apply the bijector to the sample x

y = bijector(x)
y
```

    <tf.Tensor: shape=(2,), dtype=float32, numpy=array([1.039108 , 0.8977271], dtype=float32)>

### Batch operations and broadcasting

As you've seen before, it's important to be very careful with shapes in TensorFlow Probability. That's because there are three possible components to a shape: the event shape (dimensionality of the random variable), sample shape (dimensionality of the samples drawn) and batch shape (multiple distributions can be considered in one object). This subtlety is especially important for bijectors, but can be harnassed to make powerful, and very computationally efficient, transformations of spaces. Let's examine this a little bit in this section.

In the previous examples, we applied a bijector to a two-dimensional data point $\mathbf{x}$ to create a two-dimensional data point $\mathbf{y}$. This was done using $\mathbf{y} = B \mathbf{x}$ where $B$ is the $2 \times 2$ matrix that represents the scale bijector. This is simply matrix multiplication. To implement this, we created a tensor `x` with `x.shape == [2]` and a bijector using a matrix of shape `B.shape == [2, 2]`. This generalises straightforwardly to higher dimensions: if $\mathbf{x}$ is $n$-dimensional, the bijection matrix must be of shape $n \times n$ for some $n>0$. In this case, $\mathbf{y}$ is $n$-dimensional.

But what if you wanted to apply the same bijection to ten $\mathbf{x}$ values at once? You can then arrange all these samples into a single tensor `x` with `x.shape == [10, 2]` and create a bijector as usual, with a matrix of shape `B.shape == [2, 2]`.

```python
# Create 10 samples from the uniform distribution

x = uniform.sample(10)
x
```

    <tf.Tensor: shape=(10, 2), dtype=float32, numpy=
    array([[0.29726505, 0.4838661 ],
           [0.07308066, 0.7001848 ],
           [0.05950093, 0.7649404 ],
           [0.06935227, 0.10923171],
           [0.33912015, 0.45880222],
           [0.51004636, 0.72990084],
           [0.67883885, 0.5524392 ],
           [0.6320689 , 0.9856808 ],
           [0.9687066 , 0.80901885],
           [0.87916744, 0.99455404]], dtype=float32)>

```python
# Recreate the diagonal matrix transformation with LinearOperatorDiag

scale = tf.linalg.LinearOperatorDiag(diag=[1.5, -0.5])
scale.to_dense()
```

    <tf.Tensor: shape=(2, 2), dtype=float32, numpy=
    array([[ 1.5,  0. ],
           [ 0. , -0.5]], dtype=float32)>

```python
# Create the ScaleMatvecLinearOperator bijector

bijector = tfb.ScaleMatvecLinearOperator(scale)
```

```python
# Apply the bijector to the 10 samples

y = bijector(x)
y
```

    <tf.Tensor: shape=(10, 2), dtype=float32, numpy=
    array([[ 0.44589758, -0.24193305],
           [ 0.10962099, -0.3500924 ],
           [ 0.0892514 , -0.3824702 ],
           [ 0.1040284 , -0.05461586],
           [ 0.5086802 , -0.22940111],
           [ 0.76506954, -0.36495042],
           [ 1.0182583 , -0.2762196 ],
           [ 0.9481033 , -0.4928404 ],
           [ 1.4530599 , -0.40450943],
           [ 1.3187511 , -0.49727702]], dtype=float32)>

This gives us the same plot we had before:

```python
# Run this cell to download and view a figure to illustrate the transformation

!wget -q -O diag.png --no-check-certificate "https://docs.google.com/uc?export=download&id=1sgfZ_Qzdd2v7CErP2zIk04p6R6hUW7RR"
Image("diag.png", width=500)
```

For matrix multiplication to work, we need `B.shape[-1] == x.shape[-1]`, and the output tensor has last dimension `y.shape[-1] == B.shape[-2]`. For invertibility, we also need the matrix `B` to be square. Any dimensions except for the last one on `x` become sample/batch dimensions: the operation is broadcast across these dimensions as we are used to. It's probably easiest to understand through a table of values, where `s`, `b`, `m`, and `n` are positive integers and `m != n`:

| `B.shape` | `x.shape` | `y.shape` |
| ----- | ----- | ----- |
| `(2, 2)` | `(2)` | `(2)` |
| `(n, n)` | `(m)` | `ERROR` |
| `(n, n)` | `(n)` | `(n)` |
| `(n, n)` | `(s, n)` | `(s, n)` |
| `(b, n, n)` | `(n)` | `(b, n)` |
| `(b, n, n)` | `(b, n)` | `(b, n)` |
| `(b, n, n)` | `(s, 1, n)` | `(s, b, n)` |

These rules and the ability to broadcast make batch operations easy.

We can also easily apply multiple bijectors. Suppose we want to apply both these bijectors:

$$
\begin{align}
B_1 =
\begin{bmatrix}
1 & 0 \\
0 & -1 \\
\end{bmatrix}
& \qquad
B_2 =
\begin{bmatrix}
-1 & 0 \\
0 & 1 \\
\end{bmatrix}.
\end{align}
$$

We can do this using the batched bijector

```python
# Create a batched ScaleMatvecLinearOperator bijector

diag = tf.stack((tf.constant([1, -1.]),
                 tf.constant([-1, 1.])))  # (2, 2)
scale = tf.linalg.LinearOperatorDiag(diag=diag)  # (2, 2, 2)
bijector = tfb.ScaleMatvecLinearOperator(scale=scale)
```

and we can broadcast the samples across both bijectors in the batch, as well as broadcasting the bijectors across all samples. For this, we need to include a batch dimension in the samples Tensor.

```python
# Add a singleton batch dimension to x

x = tf.expand_dims(x, axis=1)
x.shape
```

    TensorShape([10, 1, 2])

```python
# Apply the batched bijector to x

y = bijector(x)
y.shape  # (S, B, E) shape semantics
```

    TensorShape([10, 2, 2])

which gives two batches of forward values for each sample:

```python
# Run this cell to download and view a figure to illustrate the transformation

!wget -q -O linear_operator_batch.png --no-check-certificate "https://docs.google.com/uc?export=download&id=1obgl3sOIYsH_ijxxkhgBu4miBxq23fny"
Image("linear_operator_batch.png", width=500)
```

## Conclusion

In this reading, you saw how to construct scale bijectors in two dimensions using the various `ScaleMatvec` classes. You also had a quick introduction to the general `LinearOperators` class and some of its subclasses. Finally, you saw how batching makes large computations clean and efficient. Be careful to keep track of the tensor shapes, as broadcasting and the difference between batch shapes and event shapes makes errors easy. Finally, note that these bijectors are still amenable to composition (via `Chain` or simply feeding one into another) and inversion, which retains the same syntax you're used to. Enjoy using this powerful tool!

### Further reading and resources

* `ScaleMatvec` bijectors:
  * https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/ScaleMatvecDiag\n",
  * https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/ScaleMatvecLinearOperator\n",
  * https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/ScaleMatvecLU\n",
  * https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/ScaleMatvecTriL\n",
* `LinearOperator` class (see also subclasses)
  * https://www.tensorflow.org/api_docs/python/tf/linalg/LinearOperator

Week3/Scale bijectors and LinearOperator.ipynb


# FFM234, Klassisk fysik och vektorfält - Veckans tal

**[Christian Forssén](http://fy.chalmers.se/subatom/tsp/), Institutionen för  fysik, Chalmers**

Date: **Aug 10, 2019**

## Kurvintegral längs komplicerad ellips

Beräkna integralen

$$
\begin{equation}
  \oint_\Gamma \vec{F} \cdot \mbox{d}\vec{r},
\label{_auto1} \tag{1}
\end{equation}
$$

där

$$
\begin{equation}
  \vec{F} = \left[x^2-a\left(y+z\right)\right]\hat{x} + \left(y^2-az\right)
\hat{y} + \left[z^2-a\left(x+y\right)\right] \hat{z},
\label{_auto2} \tag{2}
\end{equation}
$$

och $\Gamma$ är den kurva som utgör skärningen mellan cylindern

$$
\begin{equation}
 \left(x-a\right)^2 +y^2 = a^2,\quad z \ge 0,
\label{_auto3} \tag{3}
\end{equation}
$$

och sfären

$$
\begin{equation}
  x^2 + y^2 + z^2 = R^2, \quad R> 2a,
\label{_auto4} \tag{4}
\end{equation}
$$

där $a$ är en konstant med dimensionen längd.

**Answer.**
$\pi a^3$.

**Solution.**
Vi kan först konstatera att skärningen mellan cylinder och sfär är en ellips vars exakta form är något komplicerad att fastställa.  Eftersom kurvan $\Gamma$ är en sluten kurva är det lockande att använda Stokes sats, så vi beräknar rotationen

$$
\vec{\nabla} \times \vec{F} =
  \left|\begin{array}{ccc}
\hat{x} & \hat{y} & \hat{z} \\
\frac{\partial}{\partial x} & \frac{\partial}{\partial y} &
\frac{\partial}{\partial z} \\
x^2-a\left(y+z\right) & y^2-az & z^2-a\left(x+y\right) \\
\end{array}
\right| \nonumber
$$

$$
\begin{equation}
= \left(-a+a\right) \hat{x} + \left(-a+a\right) \hat{y} + a\hat{z}
= a\hat{z}.
\label{_auto5} \tag{5}
\end{equation}
$$

Alltså är rotationen av $\vec{F}$ en rent vertikal vektor.

Vi kan nu använda Stokes sats

$$
\begin{equation}
  \oint_\Gamma \vec{F} \cdot \mbox{d}\vec{r} = \int_S \vec{\nabla} \times \vec{F} \cdot
\mbox{d}\vec{S}.
\label{_auto6} \tag{6}
\end{equation}
$$

Lägg märke till att ytan skall orienteras så att den följer högerhandsregeln.  Detta betyder att om vi följer kurvan $\Gamma$ moturs så skall normalen $\hat{n}$ till $S$ peka uppåt.

$$
\begin{equation}
  \int_S \vec{\nabla} \times \vec{F} \cdot \mbox{d}\vec{S} = \int_S a \hat{z} \cdot
\hat{n} \mbox{d}S = a \int_S \hat{z} \cdot \hat{n} \mbox{d}S.
\label{_auto7} \tag{7}
\end{equation}
$$

Skalärprodukten i den sista integralen betyder att vi projicerar ner arean $S$ på ett plan vinkelrät mot $\hat{z}$, det vill säga på $xy$-planet.  I detta planet är skärningen cylinderns tvärsnittsyta, en cirkel med radien $a$, och integralen blir cirkelarean $\pi a^2$. Alltså blir integralen till slut

$$
\begin{equation}
  \oint_\Gamma \vec{F} \cdot \mbox{d}\vec{r} = a \pi a^2 = \pi a^3.
\label{_auto8} \tag{8}
\end{equation}
$$


Table of Contents

1  Test with interactions from Banchi et al..

2  Reproduce the same parameters and gate, using the same topology but not knowing the correct parameters:

```python
%load_ext autoreload
%autoreload 2
%matplotlib notebook
import glob
from IPython.display import display
from IPython.core.debugger import set_trace
from collections import OrderedDict
import itertools
import os
import sys
import pickle
import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
import scipy
import sympy
import pandas as pd

# import plotly
# import plotly.plotly as py
# import plotly.graph_objs as go
# plotly.offline.init_notebook_mode(connected=True)
# import cufflinks

import qutip
import theano
import theano.tensor as T

src_dir = os.path.join(os.getcwd(), os.pardir, os.pardir, 'src')
sys.path.append(src_dir)

import qubit_network.net_analysis_tools as nat
import qubit_network.utils
from qubit_network.utils import (chop, complex2bigreal, bigreal2complex,
                                 bigreal2qobj, pauli_product, normalize_phase)
from qubit_network.QubitNetwork import pauli_product
from qubit_network.model import QubitNetworkGateModel
from qubit_network.Optimizer import Optimizer
from qubit_network.net_analysis_tools import NetDataFile, NetsDataFolder
```

# Test with interactions from [*Banchi et al.*](https://arxiv.org/abs/1509.04298).
The Fredkin gate is indeed reproduced with good fidelity, and regardless of the initial ancilla state.

```python
def build_Banchietal_model():
    J01xx = 13.6
    J02xx = J01xx
    J12xx = -4.712
    J12yy = J12zz = J12xx
    J13xx = 8.4
    J23xx = J13xx
    J01zz = 11.15
    J02zz = J01zz
    h3x = 1.025
    h0z = np.pi
    interactions_values = OrderedDict({
        (1, 1, 0, 0): J01xx / 4,
        (1, 0, 1, 0): J02xx / 4,
        (0, 1, 1, 0): J12xx / 4,
        (0, 2, 2, 0): J12yy / 4,
        (0, 3, 3, 0): J12zz / 4,
        (0, 1, 0, 1): J13xx / 4,
        (0, 0, 1, 1): J23xx / 4,
        (3, 3, 0, 0): J01zz / 4,
        (3, 0, 3, 0): J02zz / 4,
        (3, 0, 0, 0): h0z / 2,
        (0, 0, 0, 1): h3x / 2
    })
    interactions_used = [i[0] for i in interactions_values.items()]
    initial_values = [i[1] for i in interactions_values.items()]
    # build sympy parameters
    free_pars = []
    for interaction_string in interactions_used:
        interaction_string = ''.join(str(char) for char in interaction_string)
        free_pars.append(sympy.Symbol('J' + interaction_string))
    # build sympy model of active interactions
    fredkin_model = pauli_product(0, 0, 0, 0) * 0
    for par, par_symbol, val in zip(interactions_used, free_pars, initial_values):
        fredkin_model += par_symbol * pauli_product(*par)
    # build initial values dict
    initial_values_dict = {symb: val for symb, val in zip(free_pars, initial_values)}
    # fredkin_model
    model = QubitNetworkGateModel(
        sympy_expr=fredkin_model,
        num_system_qubits=3,
        initial_values=initial_values_dict,
        target_gate=qutip.fredkin()
    )
    return model
```

The above already gives pretty good results, with a fidelity ~$0.99999$:

```python
%matplotlib inline
model = build_Banchietal_model()
display(model.fidelity_test(n_samples=100))
nat.plot_gate(model, permutation=[0, 1, 2, 3])
```

```python
gate = chop(normalize_phase(model.get_current_gate().ptrace([0, 1, 2])), eps=1e-2)
gate / gate[0, 0]
```

Quantum object: dims = [[2, 2, 2], [2, 2, 2]], shape = (8, 8), type = oper, isherm = True\begin{equation*}\left(\begin{array}{*{11}c}1.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 1.000 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 1.000 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 1.000 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 1.000 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.000 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.000 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.000\\\end{array}\right)\end{equation*}

A few steps of optimisation improve the fidelity to 1 up to numerical precision:

```python
optimizer = Optimizer(
    net=model,
    learning_rate=.1,
    decay_rate=.005,
    n_epochs=500,
    batch_size=2,
    target_gate=qutip.fredkin(),
    training_dataset_size=200,
    test_dataset_size=100,
    sgd_method='momentum'
)
optimizer.run()
```

```python
display(model.fidelity_test(n_samples=100))
chop(normalize_phase(model.get_current_gate().ptrace([0, 1, 2])))
```

    1.0

Quantum object: dims = [[2, 2, 2], [2, 2, 2]], shape = (8, 8), type = oper, isherm = True\begin{equation*}\left(\begin{array}{*{11}c}1.743 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 1.743 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 1.743 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 1.743 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 1.743 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.743 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.743 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.743\\\end{array}\right)\end{equation*}

```python
optimizer.save_results('../../data/nets/fredkin_Banchietal.pickle')
```

    Successfully saved to ../../data/nets/fredkin_Banchietal.pickle

Using the same identical topology as per the paper, and the same set of initial interactions, the perfect fidelity is reproduced without problems (also, the net is saved into `data/nets/fredkin_Banchietal.pickle`):

Load the above net from file, and analyze the corresponding matrix.

Tracing over the ancilla, multiplying the whole matrix with a global phase, and chopping away the small matrix elements, we get a perfect Fredkin over the system qubits.
Moreover, the unitary obtained is a tensor product of a unitary over the system qubits (a Fredkin gate), and a unitary over the ancilla qubit:
$$ \mathcal U_{} = \mathcal U_{fredkin} \otimes \mathcal V.$$
We can test this in a few ways:
1. Calling $\rho_{out}$ the density matrix after partial tracing the ancilla, we check that $\operatorname{Tr}(\rho_{out}^2) = 1$, that is, that the resulting state is pure.
2. Now that we know that $\mathcal U$ is a tensor product, we can trace over the ancilla to gate (up to a constant factor), the Fredkin gate.
3. Knowing that the associated fidelity is 1, and hence that the gate acting over the system is a Fredkin gate, we can obtain the gate acting on the ancilla by tracing $\mathcal U$ over the system, and dividing the result by the trace of the Fredkin gate (that is, 6).

Point **(1)** is necessarily verified if the fidelity is 1.
Let us see points **(2)** and **(3)** in action:

```python
qn.trace_ancillae_and_normalize(net)
```

Quantum object: dims = [[2, 2, 2], [2, 2, 2]], shape = (8, 8), type = oper, isherm = True\begin{equation*}\left(\begin{array}{*{11}c}1.743 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 1.743 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 1.743 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 1.743 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 1.743 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.743 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.743 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.743\\\end{array}\right)\end{equation*}

```python
gate_ancilla = net.get_current_gate().ptrace([3]) / 6
qn.normalize_phase(gate)
```

Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = False\begin{equation*}\left(\begin{array}{*{11}c}0.872 & -0.490j\\-0.490j & 0.872\\\end{array}\right)\end{equation*}

# Reproduce the same parameters and gate, using the same topology but not knowing the correct parameters:

If we start with random initial parameters, with only `xx`, `yy`, `zz`, and the two self-interaction used in the paper, we don't get very far:

```python
# set up a set of interaction mimicking those of Banchi et al.
interactions = []
for pair in pairs:
    interactions.append((pair, 'xx'))
    interactions.append((pair, 'yy'))
    interactions.append((pair, 'zz'))

interactions.append((0, 'z'))
interactions.append((3, 'x'))

# define the network
net = qn.QubitNetwork(
    num_qubits=4,
    system_qubits=3,
    interactions=interactions,
    ancillae_state=qutip.basis(2, 0)
)
net.J.set_value(np.ones_like(net.J.get_value()))

# train the network
qn.sgd_optimization(
    net=net,
    learning_rate=1,
    n_epochs=200,
    batch_size=10,
    target_gate=qutip.fredkin(),
    training_dataset_size=100,
    test_dataset_size=1000,
    decay_rate=.01
)
```

    Generating training data...
    Building the model...
    Let's roll!

    <IPython.core.display.Javascript object>

    Finished training
    Final fidelity: 0.7304610421775868

Same thing if we use the same topology as per the paper, including the constraints on the various parameters:

```python
# set up a set of interaction mimicking those of Banchi et al.
net_topology = OrderedDict()
net_topology[((0, 1), 'xx')] = '1'
net_topology[((0, 2), 'xx')] = '1'
net_topology[((1, 2), 'xx')] = '2'
net_topology[((1, 2), 'yy')] = '2'
net_topology[((1, 2), 'zz')] = '2'
net_topology[((1, 3), 'xx')] = '3'
net_topology[((2, 3), 'xx')] = '3'
net_topology[((0, 1), 'zz')] = '4'
net_topology[((0, 2), 'zz')] = '4'
net_topology[(3, 'x')] = '5'
net_topology[(0, 'z')] = '6'

# define the network
net = qn.QubitNetwork(
    num_qubits=4,
    system_qubits=3,
    ancillae_state=qutip.basis(2, 0),
    net_topology=net_topology,
    target_gate=qutip.fredkin()
)
# net.J.set_value(np.array([13.6 / 4, -4.712 / 4, 8.4 / 4, 11.15 / 4, 1.025 / 2, 3.141 / 2]))

# train the network
qn.sgd_optimization(
    net=net,
    learning_rate=.1,
    n_epochs=200,
    batch_size=10,
    target_gate=qutip.fredkin(),
    training_dataset_size=100,
    test_dataset_size=1000,
    decay_rate=.01
)
```

    Generating training data...
    Building the model...
    Let's roll!

    <IPython.core.display.Javascript object>

    Finished training
    Final fidelity: 0.6036661701260336

*However*, setting the initial values of every interaction parameter to `10`, still with the same topology as per the paper, the network converges *extremely quickly* to fidelity 1.

The resulting interactions are different than those in the paper, but they produce the same unitary evolution over the all qubit network:

```python
# define the network
net = qn.QubitNetwork(
    num_qubits=4,
    system_qubits=3,
    ancillae_state=qutip.basis(2, 0),
    net_topology=net_topology,
    target_gate=qutip.fredkin()
)
net.J.set_value(np.array([10., 10., 10., 10., 10., 10.]))
# net.J.set_value(np.array([13.6 / 4, -4.712 / 4, 8.4 / 4, 11.15 / 4, 1.025 / 2, 3.141 / 2]))

# train the network
qn.sgd_optimization(
    net=net,
    learning_rate=.1,
    n_epochs=40,
    batch_size=10,
    target_gate=qutip.fredkin(),
    training_dataset_size=100,
    test_dataset_size=1000,
    decay_rate=.01,
    saveafter_file='../data/nets/fredkin_3q+1a_paperTopology_1fid.pickle'
)
```

    Generating training data...
    Building the model...
    Let's roll!

    <IPython.core.display.Javascript object>

    Finished training
    Final fidelity: 0.9999999999999274

The interactions *look* very different than those in the paper:

```python
outdict = net.get_interactions_with_Js(renormalize_parameters=True)
print_OrderedDict(outdict)
```

    {
        "(((0, 1), 'xx'), ((0, 2), 'xx'))": 38.95448261505736,
        "(((1, 2), 'xx'), ((1, 2), 'yy'), ((1, 2), 'zz'))": 39.26990817184952,
        "(((1, 3), 'xx'), ((2, 3), 'xx'))": 36.62955921252687,
        "(((0, 1), 'zz'), ((0, 2), 'zz'))": 42.93821433320934,
        "((3, 'x'),)": 19.999921102204446,
        "((0, 'z'),)": 21.991147620075512
    }

***However***, they produce basically the same evolution, with the difference of the unitary acting on the ancilla being different:

```python
gate_paper = chop(qn.load_network_from_file('../data/nets/fredkin_Banchietal.pickle').get_current_gate())
gate_new = chop(qn.load_network_from_file('../data/nets/fredkin_3q+1a_paperTopology_1fid.pickle').get_current_gate())

gate1 = qn.trace_ancillae_and_normalize(gate_paper, 3)
gate2 = qn.trace_ancillae_and_normalize(gate_new, 3)

display(gate1)
display(gate2)

display(qn.normalize_phase(gate_paper.ptrace([3]) / 6))
display(qn.normalize_phase(gate_new.ptrace([3]) / 6))
```

Quantum object: dims = [[2, 2, 2], [2, 2, 2]], shape = (8, 8), type = oper, isherm = True\begin{equation*}\left(\begin{array}{*{11}c}1.743 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 1.743 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 1.743 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 1.743 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 1.743 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.743 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.743 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.743\\\end{array}\right)\end{equation*}

Quantum object: dims = [[2, 2, 2], [2, 2, 2]], shape = (8, 8), type = oper, isherm = True\begin{equation*}\left(\begin{array}{*{11}c}1.678 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 1.678 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 1.678 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 1.678 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 1.678 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.678 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.678 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.678\\\end{array}\right)\end{equation*}

Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = False\begin{equation*}\left(\begin{array}{*{11}c}0.872 & -0.490j\\-0.490j & 0.872\\\end{array}\right)\end{equation*}

Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = False\begin{equation*}\left(\begin{array}{*{11}c}0.839 & -0.544j\\-0.544j & 0.839\\\end{array}\right)\end{equation*}


```python
from IPython.display import HTML, Latex
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from ipywidgets import IntProgress
```

# What's the Optimal Size of the House?

---

The House of Representatives has stood at 435 members for nearly 100 years as the population of the country has tripled. How much larger should it be to minimize the difference in how many people each Member represents?

To find out, let's run a computational experiment. The proof-of-concept for these functions, tested against the official Congressional tally, is available in the `Apportionment.ipynb` file in this repo or on [Google CoLab](https://colab.research.google.com/drive/1bBHxmUGtqs2kIH-myie5_bOZgFaT-hPT?usp=sharing).

## Loading the Data

Congress reapportions the House every 10 years after the decennial Census figures are released. The [GitHub repo for this demo](https://github.com/TimeMagazineLabs/CongressionalApportionment) has data files from Census.gov for the 2010 and 2020 decennial Census counts as well as the official apportionments so that we can check our work. Here's what the data looks like:

```python
dataByDecade = {}

for year in range(1910, 2021, 10):
  data = pd.read_csv('https://raw.githubusercontent.com/TimeMagazineLabs/CongressionalApportionment/main/data/apportionment_%s.csv' % year, dtype={'State': 'string', 'Abbr': 'string', 'Reps': 'Int64'})
  data['Per_Rep'] = data['Population'] / data['Reps']
  dataByDecade['y' + str(year)] = data
```

```python
dataByDecade['y2020'].head()
```

      State
      Abbr
      Year
      Population
      Reps
      Per_Rep

      0
      Alabama
      AL
      2020
      5030053
      7
      718579.000000

      1
      Alaska
      AK
      2020
      736081
      1
      736081.000000

      2
      Arizona
      AZ
      2020
      7158923
      9
      795435.888889

      3
      Arkansas
      AR
      2020
      3013756
      4
      753439.000000

      4
      California
      CA
      2020
      39576757
      52
      761091.480769

As we can easily see, the number of people represented by each Member of Congress has ballooned drastically in the past 100 years

```python
reps1920 = dataByDecade['y1920'][dataByDecade['y1920']['Reps'] > 0]
reps2020 = dataByDecade['y2020'][dataByDecade['y2020']['Reps'] > 0]

print('In 1920, the %s House Members for %d states represented between %s and %s people' % (
  sum(reps1920['Reps']), reps1920.shape[0], '{0:,}'.format(int(min(reps1920['Per_Rep']))), '{0:,}'.format(int(max(reps1920['Per_Rep'])))
))

print('In 2020, the %s House Members for %d states represented between %s and %s people' % (
  sum(reps2020['Reps']), reps2020.shape[0], '{0:,}'.format(int(min(reps2020['Per_Rep']))), '{0:,}'.format(int(max(reps2020['Per_Rep'])))
))
```

    In 1920, the 435 House Members for 48 states represented between 77,407 and 360,350 people
    In 2020, the 435 House Members for 50 states represented between 542,703 and 990,837 people

## How Apportionment Works: The Equal Proportions Method

Since 1940, Congress [has used the "Equal Proportions Method"](https://www.census.gov/topics/public-sector/congressional-apportionment/about/computing.html) to apportion the 435 House seats to each state.

First, per the Constitution, every state gets one seat. For the remaining 385, each state's "priority" is measured by a simple formula that divides its population, as of the most recent decennial Census, by the square root of the product of the number of seats it currently has and that number plus 1 (the [geometric mean](https://mathworld.wolfram.com/GeometricMean.html)):
\begin{equation}
priority = \frac{P}{\sqrt{n*(n+1)}}
\end{equation}
For each seat, one at a time, the state with the highest priority is awarded that seat. Then its priority is recalculated and it moves further back in the line. Let's write a Python function to calculate the priority of a state, representated as a row in the table (aka "DataFrame") above. (We'll be adding a `RepsCalculated` column to compare to the official tally.)

```python
def EqualProportionsMethod(st):
  reps = st["RepsCalculated"]
  priority = st["Population"] / np.sqrt(reps * (reps + 1))
  return priority
```

## Recalculating the apportionments

As covered in detail in the previous Notebook, we can then write a function to apportion any number of seats we like:

```python
def addNextSeat(df, ORDER=[]):
  indexNext = df['Priority'].idxmax() # The index of the row with the highest priority
  # Add a seat to the state in that row
  df.loc[indexNext,'RepsCalculated'] += 1
  # Recompute the priority for this state
  df.loc[indexNext,'Priority'] = EqualProportionsMethod(df.loc[indexNext])
  # Add the state to the ORDER list
  ORDER.append(df.loc[indexNext,'Abbr'])

def calculateApportionment(popData=dataByDecade['y2020'], ignoreStates=['District of Columbia', 'Puerto Rico'], TOTAL_SEATS=435):
  ORDER_SEATS = []

  data_filtered = popData[~popData['State'].isin(ignoreStates)]
  data_filtered.reset_index(drop=True)

  test_apportionment = data_filtered.copy()
  test_apportionment['RepsCalculated'] = 1
  test_apportionment["Priority"] = test_apportionment.apply(EqualProportionsMethod, 1)

  SEATS_LEFT = TOTAL_SEATS - test_apportionment.shape[0] # This is the number of rows, which will account for any hypothetical states we add
  while SEATS_LEFT > 0:
    addNextSeat(test_apportionment, ORDER_SEATS)
    SEATS_LEFT -= 1

  test_apportionment['Per_RepCalculated'] = test_apportionment['Population'] / test_apportionment['RepsCalculated']

  return test_apportionment, ORDER_SEATS
```

```python
test_2020, order_2020 = calculateApportionment()
test_2020.head()
```

      State
      Abbr
      Year
      Population
      Reps
      Per_Rep
      RepsCalculated
      Priority
      Per_RepCalculated

      0
      Alabama
      AL
      2020
      5030053
      7
      718579.000000
      7
      672169.105833
      718579.000000

      1
      Alaska
      AK
      2020
      736081
      1
      736081.000000
      1
      520487.866603
      736081.000000

      2
      Arizona
      AZ
      2020
      7158923
      9
      795435.888889
      9
      754616.742459
      795435.888889

      3
      Arkansas
      AR
      2020
      3013756
      4
      753439.000000
      4
      673896.328360
      753439.000000

      4
      California
      CA
      2020
      39576757
      52
      761091.480769
      52
      753877.180693
      761091.480769

Just to reiterate, this matches the official apportionment, which the final seat going to Minnesota:

```python
error = np.abs(test_2020['Reps'] - test_2020['RepsCalculated'])
print('Total difference from actual apportionment: %d' % max(error))
print('Final seat went to %s' % order_2020[-1])
```

    Total difference from actual apportionment: 0
    Final seat went to MN

Great, it works! But how well does the current apportionment represent every American equally? Since we have the number of residents per Representatives for each state, one way to measure this is the "Coefficient of Variation," a simple statistical method of normalizing the standard deviation by dividing it by the average value:

```python
def coeff_of_variance(vals):
  return np.std(vals) / np.mean(vals)

coeff_of_variance(test_2020['Per_Rep'])
```

    0.10299165948743454

Okay, so the current variation is about 0.1, which could be better with more seats. Since we made the number of seats a variable, let's see if it improves if we add 25 seats:

```python
test_460, order_460 = calculateApportionment(TOTAL_SEATS=460)
print(order_460[-1])
coeff_of_variance(test_460['Per_RepCalculated'])
```

    MO

    0.09990156250807743

Yep, a House of Representatives of 460 seats would have a lower variance in how equally each American is represented. Is there an optimal number? Let's just try every value from 435 to 1,305, three times that.

While we could run the `calculateApportionment` function 970 times, this will take around 10 minutes--believe me, I tried! But that would be wasteful, because each iteration is recomputing the same order of seats and just adding another. So we can save a lot of time by running this once, for 1,035 seats, and then apply the order to however many seats we want, like so:

```python
apportionment_1305, ORDER_1305 = calculateApportionment(TOTAL_SEATS=435*3)
```

Now we just need a function to take a subset of that grand order, up to however many seats we want, and dole them out:

```python
def computeSeats(popData=dataByDecade['y2020'], ignoreStates=['District of Columbia', 'Puerto Rico'], TOTAL_SEATS=435):
  # Remember to subtract 50 from each state since we apportion one automatically
  orderX = ORDER_1305[:TOTAL_SEATS - 50]

  # Here's a little trick: Assuming that every state gets at least one seat in the 1,305 scenario, we can use `set` to reduce to unique abbreviations:
  abbrs = set(ORDER_1305)
  # print(len(abbrs)) # will be 50

  # Here's another neat trick for tallying the number of duplicate items in a list
  # h/t https://stackoverflow.com/questions/42317134/how-to-get-a-tally-of-each-letter-in-a-string
  # This will assign 1 to each state even if it never shows up in the order (e.g. AK in the 435 scenario)
  seats = { st: 1 + orderX.count(st) for st in abbrs }

  # convert the dictionary to a list of tuples
  seats = [ (st,reps) for st,reps in seats.items() ]

  # then convert it to a dataframe
  df = pd.DataFrame(seats, columns = ['Abbr', 'RepsCalculated'])
  df = df.astype({ 'Abbr': 'string', 'RepsCalculated': 'Int64' })

  # Again make a copy of the population dataframe, then join them
  data_filtered = popData[~popData['State'].isin(ignoreStates)]
  data_filtered.reset_index(drop=True)

  apportionmentX = data_filtered.copy()
  apportionmentX = apportionmentX.set_index("Abbr").join(df.set_index("Abbr"))
  apportionmentX['Per_RepCalculated'] = apportionmentX['Population'] / apportionmentX['RepsCalculated']

  # print(len(orderX), len(ORDER_1035))
  return apportionmentX, orderX
```

Again, let's run a sanity test to make sure this method computes correctly for the real apportionment

```python
test_2020B, order_2020B = computeSeats(TOTAL_SEATS=435)
test_2020B.head()

error = np.abs(test_2020B['Reps'] - test_2020B['RepsCalculated'])
print('Total difference from actual apportionment: %d' % max(error))
print('Final seat went to %s' % order_2020B[-1])
```

    Total difference from actual apportionment: 0
    Final seat went to MN

Once again, looks good! Let's try it for all sizes:

```python
SIZES = list(range(435,435*3))
TRIALS = []

progressBar = IntProgress(min=435, max=435*3) # instantiate the bar
display(progressBar) # display the bar

for s in SIZES:
  apportionmentN, orderN = computeSeats(TOTAL_SEATS=s)
  TRIALS.append({
      'size': s,
      'apportionment': apportionmentN,
      'order': orderN,
      'score': coeff_of_variance(apportionmentN['Per_RepCalculated'])
  })
  progressBar.value += 1
```

    IntProgress(value=435, max=1305, min=435)

Trust me, that went WAY faster than reculating the apportionment from scratch every time. How did we do?

```python
SCORES = [x['score'] for x in TRIALS]

fig = plt.figure(figsize=(12, 8), frameon=True);
plt.title("Variance by Size of House")
plt.xlabel('Number of seats')
plt.ylabel('Covariance in People / Representative')
ax = fig.gca();
ax.plot(SIZES, SCORES, 'o', markersize=2)
```

This is curious! While the variation, or "score," goes down generally as the size of the House increases, as we'd expect, there are clearly ideal values, as we saw last time [I ran this experiment](https://time.com/5423623/house-representatives-number-seats/). Let's identify them:

```python
results = pd.DataFrame({
    'size': SIZES,
    'score': SCORES
})
results['delta'] = results.score -  results.score.shift(1)
idealValues = results.sort_values('delta').head(15).sort_values('score')
idealValues
```

      size
      score
      delta

      824
      1259
      0.042905
      -0.002023

      673
      1108
      0.045743
      -0.005190

      482
      917
      0.049247
      -0.004085

      413
      848
      0.065905
      -0.002544

      399
      834
      0.070985
      -0.001775

      384
      819
      0.075262
      -0.002355

      376
      811
      0.077571
      -0.009300

      208
      643
      0.080377
      -0.001516

      206
      641
      0.081955
      -0.010556

      315
      750
      0.082174
      -0.001641

      293
      728
      0.082856
      -0.007041

      91
      526
      0.083756
      -0.006935

      38
      473
      0.091049
      -0.008361

      166
      601
      0.092744
      -0.007077

      8
      443
      0.101340
      -0.001685

```python
dropPoints = list(idealValues['size'][1:])
dropValues = list(idealValues['score'][1:])

fig = plt.figure(figsize=(12, 8), frameon=True);
plt.title("Variance by Size of House")
plt.xlabel('Number of seats')
plt.ylabel('Covariance in People / Representative')
ax = fig.gca();
ax.plot(SIZES, SCORES, 'o', markersize=2)
ax.plot(dropPoints, dropValues, 'o', markersize=8)
```

Let's see if there are any interesting properties, like the smallest number of reps per state, the minimum, maximum and mean number of people per representative, and a few other items:

```python
testPoints = [x for x in TRIALS if x['size'] in dropPoints]
SCENARIOS = []

for trial in testPoints:
  vals = trial['apportionment']['Per_RepCalculated']
  seats = trial['apportionment']['RepsCalculated'].tolist()

  SCENARIOS.append({
    'size': trial['size'],
    'score': trial['score'],
    'perRepMin': min(vals),
    'perRepMax': max(vals),
    'perRepMean': np.mean(vals),
    'minSeats': min(seats),
    'minSeatCount': seats.count(min(seats)),
    'lastSeat': trial['order'][-1]
  })
```

```python
df = pd.DataFrame.from_dict(SCENARIOS)
df.sort_values('size')
```

      size
      score
      perRepMin
      perRepMax
      perRepMean
      minSeats
      minSeatCount
      lastSeat

      0
      443
      0.101340
      542703.500000
      990837.000000
      740876.872452
      1
      6
      ID

      1
      473
      0.091049
      495418.500000
      887770.000000
      692717.863082
      1
      5
      DE

      2
      526
      0.083756
      443885.000000
      779702.000000
      629202.604718
      1
      4
      SD

      3
      601
      0.092744
      389851.000000
      736081.000000
      549025.108946
      1
      3
      ND

      4
      641
      0.081955
      368040.500000
      643503.000000
      511273.805858
      1
      2
      AK

      5
      643
      0.080377
      368040.500000
      643503.000000
      508031.535423
      1
      2
      WV

      6
      728
      0.082856
      321751.500000
      577719.000000
      456041.125297
      1
      1
      VT

      7
      750
      0.082174
      321751.500000
      577719.000000
      440050.345793
      1
      1
      MT

      8
      811
      0.077571
      288859.500000
      495418.500000
      405550.329794
      2
      6
      WY

      9
      819
      0.075262
      288859.500000
      459696.333333
      400053.037279
      2
      5
      DE

      10
      834
      0.070985
      288859.500000
      454527.333333
      390546.797816
      2
      5
      NH

      11
      848
      0.065905
      288859.500000
      443885.000000
      383700.929714
      2
      5
      ME

      12
      917
      0.049247
      288859.500000
      392666.600000
      358310.239965
      2
      4
      SD

      13
      1108
      0.045743
      245360.333333
      330279.000000
      295759.126984
      2
      2
      AK

As we see, there are many optimal values that get increasingly more equal as the size grows. **The value 917 has a dramatic drop to a covariance below .05**, and increasing the size beyong that doesn't improve the score that much.

But if that's too big a House for your liking, even 473 is a huge improvement.

```python
apportionment917, _ = calculateApportionment(TOTAL_SEATS=917)
apportionment917[['State', 'Population', 'Reps', 'Per_Rep', 'RepsCalculated', 	'Per_RepCalculated']]
```

      State
      Population
      Reps
      Per_Rep
      RepsCalculated
      Per_RepCalculated

      0
      Alabama
      5030053
      7
      718579.000000
      14
      359289.500000

      1
      Alaska
      736081
      1
      736081.000000
      2
      368040.500000

      2
      Arizona
      7158923
      9
      795435.888889
      20
      357946.150000

      3
      Arkansas
      3013756
      4
      753439.000000
      8
      376719.500000

      4
      California
      39576757
      52
      761091.480769
      109
      363089.513761

      5
      Colorado
      5782171
      8
      722771.375000
      16
      361385.687500

      6
      Connecticut
      3608298
      5
      721659.600000
      10
      360829.800000

      7
      Delaware
      990837
      1
      990837.000000
      3
      330279.000000

      9
      Florida
      21570527
      28
      770375.964286
      60
      359508.783333

      10
      Georgia
      10725274
      14
      766091.000000
      30
      357509.133333

      11
      Hawaii
      1460137
      2
      730068.500000
      4
      365034.250000

      12
      Idaho
      1841377
      2
      920688.500000
      5
      368275.400000

      13
      Illinois
      12822739
      17
      754278.764706
      35
      366363.971429

      14
      Indiana
      6790280
      9
      754475.555556
      19
      357383.157895

      15
      Iowa
      3192406
      4
      798101.500000
      9
      354711.777778

      16
      Kansas
      2940865
      4
      735216.250000
      8
      367608.125000

      17
      Kentucky
      4509342
      6
      751557.000000
      12
      375778.500000

      18
      Louisiana
      4661468
      6
      776911.333333
      13
      358574.461538

      19
      Maine
      1363582
      2
      681791.000000
      4
      340895.500000

      20
      Maryland
      6185278
      8
      773159.750000
      17
      363839.882353

      21
      Massachusetts
      7033469
      9
      781496.555556
      19
      370182.578947

      22
      Michigan
      10084442
      13
      775726.307692
      28
      360158.642857

      23
      Minnesota
      5709752
      8
      713719.000000
      16
      356859.500000

      24
      Mississippi
      2963914
      4
      740978.500000
      8
      370489.250000

      25
      Missouri
      6160281
      8
      770035.125000
      17
      362369.470588

      26
      Montana
      1085407
      2
      542703.500000
      3
      361802.333333

      27
      Nebraska
      1963333
      3
      654444.333333
      5
      392666.600000

      28
      Nevada
      3108462
      4
      777115.500000
      9
      345384.666667

      29
      New Hampshire
      1379089
      2
      689544.500000
      4
      344772.250000

      30
      New Jersey
      9294493
      12
      774541.083333
      26
      357480.500000

      31
      New Mexico
      2120220
      3
      706740.000000
      6
      353370.000000

      32
      New York
      20215751
      26
      777528.884615
      56
      360995.553571

      33
      North Carolina
      10453948
      14
      746710.571429
      29
      360480.965517

      34
      North Dakota
      779702
      1
      779702.000000
      2
      389851.000000

      35
      Ohio
      11808848
      15
      787256.533333
      33
      357843.878788

      36
      Oklahoma
      3963516
      5
      792703.200000
      11
      360319.636364

      37
      Oregon
      4241500
      6
      706916.666667
      12
      353458.333333

      38
      Pennsylvania
      13011844
      17
      765402.588235
      36
      361440.111111

      40
      Rhode Island
      1098163
      2
      549081.500000
      3
      366054.333333

      41
      South Carolina
      5124712
      7
      732101.714286
      14
      366050.857143

      42
      South Dakota
      887770
      1
      887770.000000
      3
      295923.333333

      43
      Tennessee
      6916897
      9
      768544.111111
      19
      364047.210526

      44
      Texas
      29183290
      38
      767981.315789
      81
      360287.530864

      45
      Utah
      3275252
      4
      818813.000000
      9
      363916.888889

      46
      Vermont
      643503
      1
      643503.000000
      2
      321751.500000

      47
      Virginia
      8654542
      11
      786776.545455
      24
      360605.916667

      48
      Washington
      7715946
      10
      771594.600000
      21
      367426.000000

      49
      West Virginia
      1795045
      2
      897522.500000
      5
      359009.000000

      50
      Wisconsin
      5897473
      8
      737184.125000
      16
      368592.062500

      51
      Wyoming
      577719
      1
      577719.000000
      2
      288859.500000

```python
apportionment473, _ = calculateApportionment(TOTAL_SEATS=473)
apportionment473[['State', 'Population', 'Reps', 'Per_Rep', 'RepsCalculated', 	'Per_RepCalculated']]
```

      State
      Population
      Reps
      Per_Rep
      RepsCalculated
      Per_RepCalculated

      0
      Alabama
      5030053
      7
      718579.000000
      7
      718579.000000

      1
      Alaska
      736081
      1
      736081.000000
      1
      736081.000000

      2
      Arizona
      7158923
      9
      795435.888889
      10
      715892.300000

      3
      Arkansas
      3013756
      4
      753439.000000
      4
      753439.000000

      4
      California
      39576757
      52
      761091.480769
      56
      706727.803571

      5
      Colorado
      5782171
      8
      722771.375000
      8
      722771.375000

      6
      Connecticut
      3608298
      5
      721659.600000
      5
      721659.600000

      7
      Delaware
      990837
      1
      990837.000000
      2
      495418.500000

      9
      Florida
      21570527
      28
      770375.964286
      31
      695823.451613

      10
      Georgia
      10725274
      14
      766091.000000
      15
      715018.266667

      11
      Hawaii
      1460137
      2
      730068.500000
      2
      730068.500000

      12
      Idaho
      1841377
      2
      920688.500000
      3
      613792.333333

      13
      Illinois
      12822739
      17
      754278.764706
      18
      712374.388889

      14
      Indiana
      6790280
      9
      754475.555556
      10
      679028.000000

      15
      Iowa
      3192406
      4
      798101.500000
      5
      638481.200000

      16
      Kansas
      2940865
      4
      735216.250000
      4
      735216.250000

      17
      Kentucky
      4509342
      6
      751557.000000
      6
      751557.000000

      18
      Louisiana
      4661468
      6
      776911.333333
      7
      665924.000000

      19
      Maine
      1363582
      2
      681791.000000
      2
      681791.000000

      20
      Maryland
      6185278
      8
      773159.750000
      9
      687253.111111

      21
      Massachusetts
      7033469
      9
      781496.555556
      10
      703346.900000

      22
      Michigan
      10084442
      13
      775726.307692
      14
      720317.285714

      23
      Minnesota
      5709752
      8
      713719.000000
      8
      713719.000000

      24
      Mississippi
      2963914
      4
      740978.500000
      4
      740978.500000

      25
      Missouri
      6160281
      8
      770035.125000
      9
      684475.666667

      26
      Montana
      1085407
      2
      542703.500000
      2
      542703.500000

      27
      Nebraska
      1963333
      3
      654444.333333
      3
      654444.333333

      28
      Nevada
      3108462
      4
      777115.500000
      4
      777115.500000

      29
      New Hampshire
      1379089
      2
      689544.500000
      2
      689544.500000

      30
      New Jersey
      9294493
      12
      774541.083333
      13
      714961.000000

      31
      New Mexico
      2120220
      3
      706740.000000
      3
      706740.000000

      32
      New York
      20215751
      26
      777528.884615
      29
      697094.862069

      33
      North Carolina
      10453948
      14
      746710.571429
      15
      696929.866667

      34
      North Dakota
      779702
      1
      779702.000000
      1
      779702.000000

      35
      Ohio
      11808848
      15
      787256.533333
      17
      694638.117647

      36
      Oklahoma
      3963516
      5
      792703.200000
      6
      660586.000000

      37
      Oregon
      4241500
      6
      706916.666667
      6
      706916.666667

      38
      Pennsylvania
      13011844
      17
      765402.588235
      19
      684833.894737

      40
      Rhode Island
      1098163
      2
      549081.500000
      2
      549081.500000

      41
      South Carolina
      5124712
      7
      732101.714286
      7
      732101.714286

      42
      South Dakota
      887770
      1
      887770.000000
      1
      887770.000000

      43
      Tennessee
      6916897
      9
      768544.111111
      10
      691689.700000

      44
      Texas
      29183290
      38
      767981.315789
      42
      694840.238095

      45
      Utah
      3275252
      4
      818813.000000
      5
      655050.400000

      46
      Vermont
      643503
      1
      643503.000000
      1
      643503.000000

      47
      Virginia
      8654542
      11
      786776.545455
      12
      721211.833333

      48
      Washington
      7715946
      10
      771594.600000
      11
      701449.636364

      49
      West Virginia
      1795045
      2
      897522.500000
      3
      598348.333333

      50
      Wisconsin
      5897473
      8
      737184.125000
      8
      737184.125000

      51
      Wyoming
      577719
      1
      577719.000000
      1
      577719.000000

# Next Up:

What if there's a better method?