# Underwater Image Restoration By Blue-Green Channels Dehazing and Red Channel Correction

## 1. The Challenge and Proposed Solution
Capturing clear underwater images is challenging due to various reasons namely,

  - Physical properties of underwater environment

  - Effects of absorption and scattering

  - Varying attenuation of light in different wavelengths

Fig.1 -This figure plots the differeing light intensities in water [4]

An underwater image can be decomposed into the superposition of a *direct*, *forward-scattering* and a *back-scattering* component. The forward scattering causes blurring of the image while the back-scattering masks the details of the senarios.

With the help of Fig.1, we make the following obsevations:

 - Red light is more easily absorbed than blue or green light
 - The shorter wavelengths of green and blue light scatter more than the long wavelength of red light

In [4], the following method is proposed for de-hazing of underwater images - using a modified version the Dark Channel Prior method to first de-haze the blue and green channels followed by the Red channel correction using the *Gray World Assumption theory*.

### 1.1. Flowchart

Fig.2 -This figures depticts the flowchart for step-by-step implementation of the proposed method [4]

```python
from IPython.display import HTML
from IPython.display import display

tag = HTML('''
To show/hide this cell's raw code input, click <a href="javascript:code_toggle()">here</a>.''')
display(tag)
```

To show/hide this cell's raw code input, click here.

```python
import os
import datetime
import numpy as np
import cv2
import natsort
import matplotlib.pyplot as plt
import io, os, sys, types
from IPython import get_ipython
from nbformat import read
from IPython.core.interactiveshell import InteractiveShell

def find_notebook(fullname, path=None):
    """find a notebook, given its fully qualified name and an optional path

    This turns "foo.bar" into "foo/bar.ipynb"
    and tries turning "Foo_Bar" into "Foo Bar" if Foo_Bar
    does not exist.
    """
    name = fullname.rsplit('.', 1)[-1]
    if not path:
        path = ['']
    for d in path:
        nb_path = os.path.join(d, name + ".ipynb")
        if os.path.isfile(nb_path):
            return nb_path
        # let import Notebook_Name find "Notebook Name.ipynb"
        nb_path = nb_path.replace("_", " ")
        if os.path.isfile(nb_path):
            return nb_path

class NotebookLoader(object):
    """Module Loader for Jupyter Notebooks"""
    def __init__(self, path=None):
        self.shell = InteractiveShell.instance()
        self.path = path

    def load_module(self, fullname):
        """import a notebook as a module"""
        path = find_notebook(fullname, self.path)

        print ("importing Jupyter notebook from %s" % path)

        # load the notebook object
        with io.open(path, 'r', encoding='utf-8') as f:
            nb = read(f, 4)

        mod = types.ModuleType(fullname)
        mod.__file__ = path
        mod.__loader__ = self
        mod.__dict__['get_ipython'] = get_ipython
        sys.modules[fullname] = mod

        save_user_ns = self.shell.user_ns
        self.shell.user_ns = mod.__dict__

        try:
          for cell in nb.cells:
            if cell.cell_type == 'code':
                # transform the input to executable Python
                code = self.shell.input_transformer_manager.transform_cell(cell.source)
                # run the code in themodule
                exec(code, mod.__dict__)
        finally:
            self.shell.user_ns = save_user_ns
        return mod

class NotebookFinder(object):
    """Module finder that locates Jupyter Notebooks"""
    def __init__(self):
        self.loaders = {}

    def find_module(self, fullname, path=None):
        nb_path = find_notebook(fullname, path)
        if not nb_path:
            return

        key = path
        if path:
            # lists aren't hashable
            key = os.path.sep.join(path)

        if key not in self.loaders:
            self.loaders[key] = NotebookLoader(path)
        return self.loaders[key]

sys.meta_path.append(NotebookFinder())

# Importing the required libraries and supporting functions
from extras.DetermineDepth import determineDepth
from extras.TransmissionEstimation import getTransmission
from extras.getAdaptiveSceneRadiance import AdaptiveSceneRadiance
from extras.getAtomsphericLight import getAtomsphericLight
from extras.sceneRadianceGb import sceneRadianceGB
from extras.sceneRadianceR import sceneradiance
from extras.GuidedFilter import GuidedFilter
from extras.guidedfilter_he import guided_filter_he

```

    importing Jupyter notebook from C:\Users\Giridhar\Downloads\UnderwaterDeHaze_Lalitha\extras\DetermineDepth.ipynb
    importing Jupyter notebook from C:\Users\Giridhar\Downloads\UnderwaterDeHaze_Lalitha\extras\TransmissionEstimation.ipynb
    importing Jupyter notebook from C:\Users\Giridhar\Downloads\UnderwaterDeHaze_Lalitha\extras\getAdaptiveSceneRadiance.ipynb
    importing Jupyter notebook from C:\Users\Giridhar\Downloads\UnderwaterDeHaze_Lalitha\extras\getAtomsphericLight.ipynb
    importing Jupyter notebook from C:\Users\Giridhar\Downloads\UnderwaterDeHaze_Lalitha\extras\sceneRadianceGb.ipynb
    importing Jupyter notebook from C:\Users\Giridhar\Downloads\UnderwaterDeHaze_Lalitha\extras\sceneRadianceR.ipynb
    importing Jupyter notebook from C:\Users\Giridhar\Downloads\UnderwaterDeHaze_Lalitha\extras\GuidedFilter.ipynb
    importing Jupyter notebook from C:\Users\Giridhar\Downloads\UnderwaterDeHaze_Lalitha\extras\guidedfilter_he.ipynb

```python
tag = HTML('''
To show/hide this cell's raw code input, click <a href="javascript:code_toggle()">here</a>.''')
display(tag)

def  refinedtransmission(transmission, img):
    # This function refines the transmission by using a guided filter to remove the halo effects
    gimfiltR = 50  # The size of the radius when guiding the filter
    eps = 10 ** -3  #The value of epsilon during bootstrap filtering

    guided_filter = GuidedFilter(img, gimfiltR, eps)
    transmission[:,:,0] = guided_filter.filter(transmission[:,:,0])
    transmission[:,:,1] = guided_filter.filter(transmission[:,:,1])
    transmission = np.clip(transmission, 0.1, 0.9)

    return transmission
```

To show/hide this cell's raw code input, click here.

```python
tag = HTML('''
To show/hide this cell's raw code input, click <a href="javascript:code_toggle()">here</a>.''')
display(tag)

def adapexpMap(img, sceneRadiance, Lambda, blockSize):
    # Applied as the final step of the image correction process where the regions that are either too bright or too dark are
    # adjusted accordingly to provide a balanced scene
    minValue = 10 ** -2
    img = np.uint8(img)
    sceneRadiance = np.uint8(sceneRadiance)

    yj_cr_cb = cv2.cvtColor(sceneRadiance, cv2.COLOR_BGR2YCrCb)
    yi_cr_cb = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
    normyj_cr_cb = (yj_cr_cb - yj_cr_cb.min()) / (yj_cr_cb.max() - yj_cr_cb.min())
    normyi_cr_cb = (yi_cr_cb - yi_cr_cb.min()) / (yi_cr_cb.max() - yi_cr_cb.min())
    Yi = normyi_cr_cb[:, :, 0]
    Yj = normyj_cr_cb[:, :, 0]
    Yi = np.clip(Yi, minValue,1)
    Yj = np.clip(Yj, minValue,1)

    S = (Yj * Yi + 0.3 * Yi ** 2) / (Yj ** 2 + 0.3 * Yi ** 2)

    filR = 50  # The size of the radius when guiding the filter
    eps = 10 ** -3  # The value of epsilon during bootstrap filtering

    guided_filter = GuidedFilter(yi_cr_cb, filR, eps)

    refS = guided_filter.filter(S)

    S_three = np.zeros(img.shape)
    S_three[:, :, 0] = S_three[:, :, 1] = S_three[:, :, 2] = refS

    return S_three
```

To show/hide this cell's raw code input, click here.

## 2. Blue-Green Channels Dehazing

Now, the attenuation of the blue and green channels result from scattering which is similar to the type of attenuation seen in tradionally terrestrial images. Based on this intuition, we can then describe the underwater imaging model described as

$ \begin{align}
I^{c}(x) = J^{c}(x)t(x) + B^{c}(1-t(x)), c \in \{g,b\}
\end{align} $

Here,

    $I(x)$ : Observed image

    $J(x)$ : Restored image

    $B$ : Background light

    $t(x) \in [0,1]$ is the medium transmission map based on camera parameters

The Background light *B* can be estimated from the knowledge that the Red channel attuates much faster than the green and blue channels underwater. First, the largest difference *D(x)* is computed by finding the difference between the maximum intensity of the red channel and the maximum intensity of the green and blue channels.

$ \begin{align}
D(x) = \max_{ x \in \Omega,c \in r } I^{c}(x) - \max_{ x \in \Omega,c \in \{g,b\} } I^{c}(x)
\end{align} $

Here, $\Omega$ corresponds to a local patch in the image. From this, *B* is estimated as:

$ \begin{align}
B^{c} = avg(I^{c}(arg \min_{x} D(x))), c \in \{g,b\}
\end{align} $

### 2.1. Transmission Map for Blue and Green Channels
Based on the Rayleigh scattering theory the following assumptions are valid:
 - Blue and Green channels have equal attenuation underwater
 - Transmission map for a given patch is same for both channels

To find the transmission map we minimise the imaging model equation over $c$ and simplify based on Dark Channel Prior theory. Rearranging the resulting equation we get:

$ \begin{align}
t(x) = 1 - \min_{c \in \{g,b\}}(\min_{x\in\Omega}(\frac{I^{c}(x)}{B^{c}}))
\end{align} $

Fig.3- Medium Transmission map (a) Original underwater image (b) Coarse transmission map (c) Refined medium transmission maps [4]

Fig.3(b) shows *halo* effects and block artifacts seen in the coarse tranmission map. This is because the coarse transmission map is a rough estimate over a patch of image. To fine-tune the transmission map, the coarse estimates are refined with the help of the guided filter which results in a more clear *t(x)* as seen in Fig.3(c).

### 2.2. Restored haze-free Blue-Green Channels
The last step to the de-hazing process of the Blue-Green channels is restoring the respective channels after correction. This is done by a simple back-substitution of the estimated parameters into the original image equation as

$ \begin{align}
J^{c}(x) = \frac{I^{c}(x)-B^{c}}{t^{c}(x)}+B^{c}, c\in\{g,b\}
\end{align} $

where $J^{c}(x)$ represents the restored channel.

## 3. Red Channel Correction
It is not straightforward to calculate the absorption rate for red light. Therefore the red light correction is done based on the Gray-World assumption theory which states that the average value of object color in an ideal image is gray. The following assumption forms the basis for the red channel correction process:

 $ \begin{align}
(avgRr + avgBr +avgGr)/3 = 0.5
\end{align} $

Here, *avgR*,*avgB* and *avgG* are the normalized average values of the respective recorvered channels.
Then, the average value of the recovered red channel is estimated as:

 $ \begin{align}
avgRr = 1.5 - avgBr - avgGr
\end{align} $

The estimated *compensation coefficient*, $\delta$ is

$ \begin{align}
\delta = avgRr / avgR
\end{align} $

where,

$avgR$ : Normalised *average* value of original red channel

Finally, the recovered red channel is found as:

$ \begin{align}
Rrec = R .* \delta
\end{align} $

## 4. Adaptive Exposure Map Estimation
Some bright and dark parts of the image become too dark or bright after the dehazing and correction process. To restore the right amount of brightness/darkness an Adaptive Exposure Map is implemented on the image. This Adaptive Exposure Map $s(x)$ is obtained by solving an optimization problem.

$ \begin{align}
\min_{s}\sum_{x}\{[1-s(x)\frac{Y_{J(x)}}{Y_{I(x)}}]^{2} + \lambda[s(x) - 1]^{2}\} + \phi(s)
\end{align} $

Here,

$s(x)$: The adaptive exposure map

$Y_{J}$: Illumination intensity of restored image

$\lambda$: A tuning hyperparameter set to $0.3$

$\phi(.)$: Smoothness regularization

This optimization problem can be solved in two steps:
 - Solve $s(x)$ without the smoothness rgularization
 - Apply guided filter to smooth the solution

The output after optimization can be obtained as:

$ \begin{align}
OutputExp = J^{c}(x).*s(x), c\in\{r,g,b\}
\end{align} $

```python
from IPython.display import HTML
from IPython.display import display

tag = HTML('''
To show/hide this cell's raw code input, click <a href="javascript:code_toggle()">here</a>.''')
display(tag)

if __name__ == '__main__':
    pass
starttime = datetime.datetime.now() #Records the time taken to run the code

folder = (os.path.abspath("."))
path = folder + "/Datasets"
files = os.listdir(path)
files =  natsort.natsorted(files)

for i in range(len(files)):
    file = files[i]
    filepath = path + "/" + file
    prefix = file.split('.')[0]
    if os.path.isfile(filepath):
        img1 = cv2.imread(folder +'/Datasets/' + file)
        img = (img1 - img1.min()) / (img1.max() - img1.min()) * 255
        blockSize = 9
        largest_diff = determineDepth(img, blockSize)
        #print(largest_diff)
        AtomsphericLight, AtomsphericLightGB, AtomsphericLightRGB = getAtomsphericLight(largest_diff, img)
       # print('AtomsphericLightRGB',AtomsphericLightRGB)
        transmission = getTransmission(img, AtomsphericLightRGB, blockSize)
        transmission = refinedtransmission(transmission, img)

        sceneRadiance_GB = sceneRadianceGB(img, transmission, AtomsphericLightRGB)
        sceneRadiance = sceneradiance(img, sceneRadiance_GB)

        S_x = adapexpMap(img, sceneRadiance, Lambda=0.3, blockSize=blockSize)
        sceneRadiance_adap = AdaptiveSceneRadiance(sceneRadiance, S_x)

        sceneRadiance_adap = cv2.cvtColor(sceneRadiance_adap, cv2.COLOR_BGR2RGB)
        img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
        plt.subplot(1, 2, 1)
        plt.imshow(img1)
        plt.gca().set_title('Original')
        plt.subplot(1, 2, 2)
        plt.imshow(sceneRadiance_adap)
        plt.gca().set_title('Corrected')
        plt.show()

Endtime = datetime.datetime.now()
Time = Endtime - starttime
print('Time', Time)
```

## 5. Comparison to Other Methods
In order to further highleten the requirement for the three color channels to be given the treatment as discussed above for underwater images we compare the results of the Green-Blue channel dehazing, Red channel correction method with the dehazing algorithms proposed for traditionally teresstrial images.

Here, we see that the traditional method of Dark Channel Prior gives poor results when it comes to de-hazing and color correcting an underwater image. This once again affirms our observation that different datasets require different methods for de-hazing and color correction.

```python

```


# Efficient multigrid methods for high-order nonlinear solid mechanics on emerging architectures

## Rezgar Shakeri1,

Valeria Barra2, Natalie Beams3, Jed Brown4, Yohann Dudouit5, Karen Stengel4, Jeremy L. Thompson4

1 Department of Civil Engineering, CU Boulder
2 Department of Environmental Science and Engineering, California Institute of Technology
3 Innovative Computing Laboratory, University of Tennessee
4 Department of Computer Science, CU Boulder
5 Lawrence Livermore National Laboratory

Copper-Iterative, 07 Apr 2022

```python
from IPython.display import SVG, Video, HTML, IFrame
```

# Outline

### Why Matrix Free ?

### Model Equations

* Neo-Hookean hyperelasticity at finite strain
* Initial vs Current Configuration

### Matrix Free FE

* Math
* libCEED

### P-Multigrid

### Results

# Why Matrix Free ?

* Modern hardware does 10 flops per byte.
* Matrix-free methods store and move less data, compute faster.
* The Jacobian of a non-linear operator rapidly loses the sparsity and the number of nonzero entries per dof grows with $p^d$ as the order $p$ of basis function increases in $d$ dimension

# Model Equations

* Strong form
\begin{equation}
-\nabla_X \cdot \boldsymbol{P} - \rho_0 \boldsymbol{g} = 0
\end{equation}
where $\boldsymbol{P}=\boldsymbol{FS}$ and
$$
\color{blue}{
\boldsymbol{S}=\lambda \, \log J \, \boldsymbol{C}^{-1} + \mu \left( \boldsymbol{I} - \boldsymbol{C}^{-1} \right).}
$$

* Variational form ($\color{red}{\text{Initial}}$ configuration)
\begin{equation}
\int_{\Omega_0}{\nabla_X \boldsymbol{v} \colon \boldsymbol{FS}} \, dV
 - \int_{\Omega_0}{\boldsymbol{v} \cdot \rho_0 \boldsymbol{g}} \, dV
 = 0, \quad \forall \boldsymbol v \in \mathcal V,
\end{equation}

* Variational form ($\color{red}{\text{Current}}$ configuration)
\begin{equation}
\int_{\Omega_0}{\nabla_x \boldsymbol{v} \colon \boldsymbol{\tau}} \, dV
 - \int_{\Omega_0}{\boldsymbol{v} \cdot \rho_0 \boldsymbol{g}} \, dV
 = 0, \quad \forall \boldsymbol v \in \mathcal V,
\end{equation}
where
$$
\color{teal}{
\boldsymbol{\tau}=\lambda \, \log J \, \boldsymbol{I} + \, \mu \left( \boldsymbol{b}-\boldsymbol{I} \right).}
$$

### Matrix-free Finite Element Formulation

\begin{gather*}
    v^T F(u) \sim \int_\Omega v \cdot \color{olive}{f_0(u, \nabla u)} + \nabla v \!:\! \color{olive}{f_1(u, \nabla u)} \quad
    v^T J du \sim \int_\Omega \begin{bmatrix} v \\ \nabla v \end{bmatrix}^T \color{teal}{\begin{bmatrix} f_{0,0} & f_{0,1} \\ f_{1,0} & f_{1,1} \end{bmatrix}}
    \begin{bmatrix} du \\ \nabla du \end{bmatrix} \\
\end{gather*}

### [libCEED](https://libceed.readthedocs.io): Efficient Extensible Discretization

* $\mathcal P$: Process decomposition,

* $\mathcal E$ : Element restriction/assembly operator

* $B$ : Basis (DoFs-to-Qpts) evaluator,

* $D$ : Operator at quadrature point- Qfunction ($f_0, f_1, f_{0,0},...$)

## [libCEED](https://libceed.readthedocs.io): fast algebra for finite elements

* Single source vanilla C for QFunctions

* Backend plugins with run-time selection
    * e.g., `./bps -ceed /gpu/cuda`

* Same source code can call multiple CEEDs with different backends

* Open source (BSD-2 license) C library with Fortran, Python, Julia, and Rust interfaces

* Available via MFEM, PETSc, Nek5000

# P-Multigrid

* h-multigrid vs p-multigrid

* 2nd order Chebyshev/Jacobi iteration in the range $[0.1\lambda_{max}, 1.1\lambda_{max}]$
* $\lambda_{max}$ estimate eigenvalue of $\left(\text{diag}\boldsymbol{A}_f \right)^{-1}\boldsymbol{A}_f$
* AMG use only information of the matrix sparsity and its entries of the assembled operators

# Results

* Schwarz Primitive mesh, left wall fixed, thickness 0.2 and 0.05, $Q_1, Q_2, Q_3$ elements

# Assembled vs Matrix-Free on Lassen machine

* Matrix-free becomes more efficient as the order increases

* Both are latency-limited for smaller problem size (left side of figure)

* Even for linear element, efficiency of matrix-free grows as DoFs increases

# Total nonlinear solve efficiency on different machines

* $Q_3$ elements using matrix-free
* Newton-Krylov with p-MG preconditioning and Boomer-AMG coarse solve
* For example in 10 seconds total time, we can solve 0.6 MDoFs/s/GPU on Perlmutter, so the target problem would be scaled to about 6 MDoFs/GPU

# Flame Graph

```python
IFrame("schwarz-q2-cuda-flame.svg", width="2000", height="200")
```

# high-order FE in engineering problem

* mesh A

* mesh B

# Future works

## Buckling problem

* implement arc-length method

# ratel

[https://gitlab.com/micromorph/ratel](https://gitlab.com/micromorph/ratel)


# **CS224W - Colab 5**

In this Colab, we will shift our focus from homogenous graphs to heterogeneous graphs. Heterogeneous graphs extend the traditional homogenous graphs that we have been working with by incorperating different node and edge types. This additional information allows us to extend the graph neural nework models that we have worked with before. Namely, we can apply heterogenous message passing, where different message types now exist between different node and edge type relationships.

In this notebook, we will first learn how to transform NetworkX graphs into DeepSNAP representations. Then, we will dive deeper into how DeepSNAP stores and represents heterogeneous graphs as PyTorch Tensors.

Lastly, we will build our own heterogenous graph neural netowrk models using PyTorch Geometric and DeepSNAP. We will then apply our models for a node property prediction task; specifically, we will evaluate these models on the heterogeneous ACM node prediction dataset.

**Note**: Make sure to **sequentially run all the cells in each section**, so that the intermediate variables / packages will carry over to the next cell

Have fun and good luck on Colab 5 :)

# Device
You might need to use GPU for this Colab.

Please click `Runtime` and then `Change runtime type`. Then set the `hardware accelerator` to **GPU**.

# Installation

```python
# Install torch geometric
import os
if 'IS_GRADESCOPE_ENV' not in os.environ:
  !pip install torch-scatter -f https://data.pyg.org/whl/torch-1.10.0+cu111.html
  !pip install torch-sparse -f https://data.pyg.org/whl/torch-1.10.0+cu111.html
  !pip install torch-geometric
  !pip install -q git+https://github.com/snap-stanford/deepsnap.git
  !pip install -U -q PyDrive
```

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  !nvcc --version
  !python -c "import torch; print(torch.version.cuda)"
```

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  import torch
  print(torch.__version__)
  import torch_geometric
  print(torch_geometric.__version__)
```

# DeepSNAP Basics

In previous Colabs we used both of graph class (NetworkX) and tensor (PyG) representations of graphs separately. The graph class `nx.Graph` provides rich analysis and manipulation functionalities, such as the clustering coefficient and PageRank. To feed the graph into the model, we need to transform the graph into tensor representations including edge tensor `edge_index` and node attributes tensors `x` and `y`. But only using tensors (as the graphs formatted in PyG `datasets` and `data`) will make many graph manipulations and analysis less efficient and harder. So, in this Colab we will use DeepSNAP which combines both representations and offers a full pipeline for GNN training / validation / testing.

In general, [DeepSNAP](https://github.com/snap-stanford/deepsnap) is a Python library to assist efficient deep learning on graphs. DeepSNAP features in its support for flexible graph manipulation, standard pipeline, heterogeneous graphs and simple API.

1. DeepSNAP is easy to be used for the sophisticated graph manipulations, such as feature computation, pretraining, subgraph extraction etc. during/before the training.
2. In most frameworks, standard pipelines for node, edge, link, graph-level tasks under inductive or transductive settings are left to the user to code. In practice, there are additional design choices involved (such as how to split dataset for link prediction). DeepSNAP provides such a standard pipeline that greatly saves repetitive coding efforts, and enables fair comparision for models.
3. Many real-world graphs are heterogeneous graphs. But packages support for heterogeneous graphs, including data storage and flexible message passing, is lacking. DeepSNAP provides an efficient and flexible heterogeneous graph that supports both the node and edge heterogeneity.

[DeepSNAP](https://github.com/snap-stanford/deepsnap) is a newly released project and it is still under development. If you find any bugs or have any improvement ideas, feel free to raise issues or create pull requests on the GitHub directly :)

In this Colab, we will focus on learning using Heterogeneous Graphs. Not many libraries are able to handle heterogeneous graphs, but DeepSNAP handles them quite elegantly, which is why we're introducing it here!

# 1) DeepSNAP Heterogeneous Graph

First, we will explore how to transform a NetworkX graph into the format supported by DeepSNAP.

In DeepSNAP we have three levels of attributes. We can have **node level** attributes including `node_feature` and `node_label`. The other two levels of attributes are graph and edge attributes. The usage is similar to the node level one except that the feature becomes `edge_feature` or `graph_feature` and label becomes `edge_label` or `graph_label` etc.

DeepSNAP extends its traditional graph representation to include heterogeneous graphs by including the following graph property features:
* `node_feature`: The feature of each node (`torch.tensor`)
* `edge_feature`: The feature of each edge (`torch.tensor`)
* `node_label`: The label of each node (`int`)
* `node_type`: The node type of each node (`string`)
* `edge_type`: The edge type of each edge (`string`)

where the key **new** features we add are `node_type` and `edge_type`, which enables us to perform heterogenous message passing.

For this first question we will work with the familiar [karate club graph](https://networkx.github.io/documentation/stable/auto_examples/graph/plot_karate_club.html) seen in Colab 1. To start, since each node in the graph belongs to one of two clubs (club "Mr. Hi" or club "Officer"), we will treat the club as the `node_type`. The code below demonstrates how to differentiate the nodes in the NetworkX graph.

```python
import networkx as nx
from networkx.algorithms.community import greedy_modularity_communities
import matplotlib.pyplot as plt
import copy

if 'IS_GRADESCOPE_ENV' not in os.environ:
  from pylab import show
  G = nx.karate_club_graph()
  community_map = {}
  for node in G.nodes(data=True):
    if node[1]["club"] == "Mr. Hi":
      community_map[node[0]] = 0
    else:
      community_map[node[0]] = 1
  node_color = []
  color_map = {0: 0, 1: 1}
  node_color = [color_map[community_map[node]] for node in G.nodes()]
  pos = nx.spring_layout(G)
  plt.figure(figsize=(7, 7))
  nx.draw(G, pos=pos, cmap=plt.get_cmap('coolwarm'), node_color=node_color)
  show()
```

### Question 1.1: Assigning Node Type and Node Features

Using the `community_map` dictionary and graph `G` from above, add node attributes `node_type` and `node_label` to the graph G. Namely, for `node_type` assign nodes in the "Mr. Hi" club to a node type `n0` and nodes in club "Officer" a node type `n1`. Note: the node type should be a `string` property.

Then for `node_label`, assign nodes in "Mr. Hi" club to a `node_label` `0` and nodes in club "Officer" a `node_label` of `1`.

Lastly, assign every node the *tensor* feature vector $[1, 1, 1, 1, 1]$.

**Hint**: Look at the NetworkX function `nx.classes.function.set_node_attributes`.

**Note**: This question is not specifically graded but is important for later questions.

```python
import torch

def assign_node_types(G, community_map):
  # TODO: Implement a function that takes in a NetworkX graph
  # G and community map assignment (mapping node id --> 0/1 label)
  # and adds 'node_type' as a node_attribute in G.

  ############# Your code here ############
  ## (~2 line of code)
  ## Note
  ## 1. Look up NetworkX `nx.classes.function.set_node_attributes`
  ## 2. Look above for the two node type values!

  pass

  #########################################

def assign_node_labels(G, community_map):
  # TODO: Implement a function that takes in a NetworkX graph
  # G and community map assignment (mapping node id --> 0/1 label)
  # and adds 'node_label' as a node_attribute in G.

  ############# Your code here ############
  ## (~2 line of code)
  ## Note
  ## 1. Look up NetworkX `nx.classes.function.set_node_attributes`

  pass

  #########################################

def assign_node_features(G):
  # TODO: Implement a function that takes in a NetworkX graph
  # G and adds 'node_feature' as a node_attribute in G. Each node
  # in the graph has the same feature vector - a torchtensor with
  # data [1., 1., 1., 1., 1.]

  ############# Your code here ############
  ## (~2 line of code)
  ## Note
  ## 1. Look up NetworkX `nx.classes.function.set_node_attributes`

  pass

  #########################################

if 'IS_GRADESCOPE_ENV' not in os.environ:
  assign_node_types(G, community_map)
  assign_node_labels(G, community_map)
  assign_node_features(G)

  # Explore node properties for the node with id: 20
  node_id = 20
  print (f"Node {node_id} has properties:", G.nodes(data=True)[node_id])
```

### Question 1.2: Assigning Edge Types

Next, we will assign three different `edge_types`:
* Edges within club "Mr. Hi": `e0`
* Edges within club "Officer": `e1`
* Edges between the two clubs: `e2`

**Hint**: Use the `community_map` from before and `nx.classes.function.set_edge_attributes`

```python
def assign_edge_types(G, community_map):
  # TODO: Implement a function that takes in a NetworkX graph
  # G and community map assignment (mapping node id --> 0/1 label)
  # and adds 'edge_type' as a edge_attribute in G.

  ############# Your code here ############
  ## (~5 line of code)
  ## Note
  ## 1. Create an edge assignment dict following rules above
  ## 2. Look up NetworkX `nx.classes.function.set_edge_attributes`

  pass

  #########################################

if 'IS_GRADESCOPE_ENV' not in os.environ:
  assign_edge_types(G, community_map)

  # Explore edge properties for a sampled edge and check the corresponding
  # node types
  edge_idx = 15
  n1 = 0
  n2 = 31
  edge = list(G.edges(data=True))[edge_idx]
  print (f"Edge ({edge[0]}, {edge[1]}) has properties:", edge[2])
  print (f"Node {n1} has properties:", G.nodes(data=True)[n1])
  print (f"Node {n2} has properties:", G.nodes(data=True)[n2])
```

## Heterogeneous Graph Visualization

Now we can visualize the Heterogeneous Graph we have generated.

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  edge_color = {}
  for edge in G.edges():
    n1, n2 = edge
    edge_color[edge] = community_map[n1] if community_map[n1] == community_map[n2] else 2
    if community_map[n1] == community_map[n2] and community_map[n1] == 0:
      edge_color[edge] = 'blue'
    elif community_map[n1] == community_map[n2] and community_map[n1] == 1:
      edge_color[edge] = 'red'
    else:
      edge_color[edge] = 'green'

  G_orig = copy.deepcopy(G)
  nx.classes.function.set_edge_attributes(G, edge_color, name='color')
  colors = nx.get_edge_attributes(G,'color').values()
  labels = nx.get_node_attributes(G, 'node_type')
  plt.figure(figsize=(8, 8))
  nx.draw(G, pos=pos, cmap=plt.get_cmap('coolwarm'), node_color=node_color, edge_color=colors, labels=labels, font_color='white')
  show()
```

where we differentiate edges within each clubs (2 types) and edges between the two clubs (1 type). Different types of nodes and edges are visualized in different colors. The NetworkX object `G` in following code can be transformed into `deepsnap.hetero_graph.HeteroGraph` directly.

## Transforming to DeepSNAP representation

We will now work through transforming the NetworkX object `G` into a `deepsnap.hetero_graph.HeteroGraph`.

```python
from deepsnap.hetero_graph import HeteroGraph

if 'IS_GRADESCOPE_ENV' not in os.environ:
  hete = HeteroGraph(G_orig)
```

## Question 1.3: How many nodes are of each type (10 Points)

```python
def get_nodes_per_type(hete):
  # TODO: Implement a function that takes a DeepSNAP dataset object
  # and return the number of nodes per `node_type`.

  num_nodes_n0 = 0
  num_nodes_n1 = 0

  ############# Your code here ############
  ## (~2 line of code)
  ## Note
  ## 1. Colab autocomplete functionality might be useful.

  pass

  #########################################

  return num_nodes_n0, num_nodes_n1

if 'IS_GRADESCOPE_ENV' not in os.environ:
  num_nodes_n0, num_nodes_n1 = get_nodes_per_type(hete)
  print("Node type n0 has {} nodes".format(num_nodes_n0))
  print("Node type n1 has {} nodes".format(num_nodes_n1))
```

## Question 1.4: Message Types - How many edges are of each message type (10 Points)

When working with heterogenous graphs, as we have discussed before, we now work with heterogenous message types (i.e. different message types for the different `node_type` and `edge_type` combinations). For example, an edge of type `e0` connecting two nodes in club "Mr. HI" would have a message type of (`n0`, `e0`, `n0`). In this problem we will analyze how many edges in our graph are of each message type.

**Hint**: If you want to learn more about what the different message types are try the call `hete.message_types`

```python
def get_num_message_edges(hete):
  # TODO: Implement this function that takes a DeepSNAP dataset object
  # and return the number of edges for each message type.
  # You should return a list of tuples as
  # (message_type, num_edge)

  message_type_edges = []

  ############# Your code here ############
  ## (~2 line of code)
  ## Note
  ## 1. Colab autocomplete functionality might be useful.

  pass

  #########################################

  return message_type_edges

if 'IS_GRADESCOPE_ENV' not in os.environ:
  message_type_edges = get_num_message_edges(hete)
  for (message_type, num_edges) in message_type_edges:
    print("Message type {} has {} edges".format(message_type, num_edges))
```

## Question 1.5: Dataset Splitting - How many nodes are in each dataset split? (10 Points)

DeepSNAP has built in Dataset creation and splitting methods for heterogeneous graphs. Here we will create train, validation, and test datasets for a node prediction task and inspect the resulting subgraphs. Specifically, write a function that computes the number of nodes with a known label in each dataset split.

```python
from deepsnap.dataset import GraphDataset

def compute_dataset_split_counts(datasets):
  # TODO: Implement a function that takes a dict of datasets in the form
  # {'train': dataset_train, 'val': dataset_val, 'test': dataset_test}
  # and returns a dict mapping dataset names to the number of labeled
  # nodes used for supervision in that respective dataset.

  data_set_splits = {}

  ############# Your code here ############
  ## (~3 line of code)
  ## Note
  ## 1. The DeepSNAP `node_label_index` dictionary will be helpful.
  ## 2. Remember to count both node_types
  ## 3. Remember each dataset only has one graph that we need to access
  ##    (i.e. dataset[0])

  pass

  #########################################

  return data_set_splits

if 'IS_GRADESCOPE_ENV' not in os.environ:
  dataset = GraphDataset([hete], task='node')
  # Splitting the dataset
  dataset_train, dataset_val, dataset_test = dataset.split(transductive=True, split_ratio=[0.4, 0.3, 0.3])
  datasets = {'train': dataset_train, 'val': dataset_val, 'test': dataset_test}

  data_set_splits = compute_dataset_split_counts(datasets)
  for dataset_name, num_nodes in data_set_splits.items():
    print("{} dataset has {} nodes".format(dataset_name, num_nodes))
```

## DeepSNAP Dataset Visualization

We can now visualize the different nodes and edges used in each graph dataset split.

```python
from deepsnap.dataset import GraphDataset

if 'IS_GRADESCOPE_ENV' not in os.environ:
  dataset = GraphDataset([hete], task='node')
  # Splitting the dataset
  dataset_train, dataset_val, dataset_test = dataset.split(transductive=True, split_ratio=[0.4, 0.3, 0.3])
  titles = ['Train', 'Validation', 'Test']

  for i, dataset in enumerate([dataset_train, dataset_val, dataset_test]):
    n0 = hete._convert_to_graph_index(dataset[0].node_label_index['n0'], 'n0').tolist()
    n1 = hete._convert_to_graph_index(dataset[0].node_label_index['n1'], 'n1').tolist()

    plt.figure(figsize=(7, 7))
    plt.title(titles[i])
    nx.draw(G_orig, pos=pos, node_color="grey", edge_color=colors, labels=labels, font_color='white')
    nx.draw_networkx_nodes(G_orig.subgraph(n0), pos=pos, node_color="blue")
    nx.draw_networkx_nodes(G_orig.subgraph(n1), pos=pos, node_color="red")
    show()
```

# 2) Heterogeneous Graph Node Property Prediction

Now, we will use PyTorch Geometric and DeepSNAP to implement a GNN model for heterogeneous graph node property prediction (node classification). We will draw upon our understanding of heterogeneous graphs from lecture and previous work in implementing GNN layers using PyG (introduced in Colab 3).

First let's take a look at the general structure of a heterogeneous GNN layer by working through an example:

Let's assume we have a graph $G$, which contains two node types $a$ and $b$, and three message types $m_1=(a, r_1, a)$, $m_2=(a, r_2, b)$ and $m_3=(a, r_3, b)$. Note: during message passing we view each message as (src, relation, dst), where messages "flow" from src to dst node types. For example, during message passing, updating node type $b$ relies on two different message types $m_2$ and $m_3$.

When applying message passing in heterogenous graphs, we seperately apply message passing over each message type. Therefore, for the graph $G$, a heterogeneous GNN layer contains three seperate Heterogeneous Message Passing layers (`HeteroGNNConv` in this Colab), where each `HeteroGNNConv` layer performs message passing and aggregation with respect to *only one message type*. Since a message type is viewed as (src, relation, dst) and messages "flow" from src to dst, each `HeteroGNNConv` layer only computes embeddings for the *dst* nodes of a given message type. For example, the `HeteroGNNConv` layer for message type $m_2$ outputs updated embedding representations *only* for node's with type b.

---

An overview of the heterogeneous layer we will create is shown below:

where we highlight the following notation:

- $H_a^{(l)[m_1]}$ is the intermediate matrix of of node embeddings for node type $a$, generated by the $l$th `HeteroGNNConv` layer for message type $m_1$.
- $H_a^{(l)}$ is the matrix with current embeddings for nodes of type $a$ after the $l$th layer of our Heterogeneous GNN model. Note that these embeddings can rely on one or more intermediate `HeteroGNNConv` layer embeddings(i.e. $H_b^{(l)}$ combines $H_b^{(l)[m_2]}$ and $H_b^{(l)[m_3]}$).

Since each `HeteroGNNConv` is only applied over a single message type, we additionally define a Heterogeneous GNN Wrapper layer (`HeteroGNNWrapperConv`). This wrapper manages and combines the output of each `HeteroGNNConv` layer in order to generate the complete updated node embeddings for each node type in layer $l$ of our model. More specifically, the $l^{th}$ `HeteroGNNWrapperConv` layer takes as input the node embeddings computed for each message type and node type (e.g. $H_b^{(l)[m_2]}$ and $H_b^{(l)[m_3]}$) and aggregates across message types with the same $dst$ node type. The resulting output of the $l^{th}$ `HeteroGNNWrapperConv` layer is the updated embedding matrix $H_i^{(l)}$ for each node type i.

Continuing on our example above, to compute the node embeddings $H_b^{(l)}$ the wrapper layer aggregates output embeddings from the `HeteroGNNConv` layers associated with message types $m_2$ and $m_3$ (i.e. $H_b^{(l)[m_2]}$ and $H_b^{(l)[m_3]}$).

---

With the `HeteroGNNWrapperConv` module, we can now draw a "simplified" heterogeneous layer structure as follows:

---
**NOTE**:
As reference, it may be helpful to additionally read through PyG's introduciton to heterogeneous graph representations and buidling heterogeneous GNN models: https://pytorch-geometric.readthedocs.io/en/latest/notes/heterogeneous.html

Looking ahead, we recommend you implement the heterogeneous GNN model in following steps:

1. Implement `HeteroGNNConv`.
2. Implement **just** `mean` aggregation within `HeteroGNNWrapperConv`.
3. Implement `generate_convs`.
4. Implement the `HeteroGNN` model and the `train` function.
5. Train the model with `mean` aggregation and test your model to make sure your model has reasonable performance.
6. Once you are confident in your mean aggregation model, implement `attn` aggregation in `HeteroGNNWrapperConv`.
7. Train the model with `attn` aggregation and test your model to make sure your model has reasonable performance.

Note: The key point of advice is to work completely through implementing the mean aggregation heterogeneous GNN model before diving into the more difficult attention based model.

## Setup

```python
import copy
import torch
import deepsnap
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch_geometric.nn as pyg_nn

from sklearn.metrics import f1_score
from deepsnap.hetero_gnn import forward_op
from deepsnap.hetero_graph import HeteroGraph
from torch_sparse import SparseTensor, matmul
```

## Dataset

You need to login to your Google account and enter the verification code below.

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  from pydrive.auth import GoogleAuth
  from pydrive.drive import GoogleDrive
  from google.colab import auth
  from oauth2client.client import GoogleCredentials

  # Authenticate and create the PyDrive client
  auth.authenticate_user()
  gauth = GoogleAuth()
  gauth.credentials = GoogleCredentials.get_application_default()
  drive = GoogleDrive(gauth)
```

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  id='1ivlxd6lJMcZ9taS44TMGG72x2V1GeVvk'
  downloaded = drive.CreateFile({'id': id})
  downloaded.GetContentFile('acm.pkl')
```

## Implementing `HeteroGNNConv`

Now let's start working on our own implementation of the heterogeneous message passing layer (`HeteroGNNConv`)! Just as in Colabs 3 and 4, we will implement the layer using PyTorch Geometric.

At a high level, the `HeteroGNNConv` layer is equivalent to the homogenous GNN layers we implemented in Colab 3, but now applied to an individual heterogeous message type. Moreover, our heterogeneous GNN layer draws directly from the **GraphSAGE** message passing model ([Hamilton et al. (2017)](https://arxiv.org/abs/1706.02216)).

We begin by defining the `HeteroGNNConv` layer with respect to message type $m$:

\begin{equation}
m =(s, r, d)
\end{equation}

where each message type is a tuple containing three elements: $s$ - the source node type, $r$ - the edge (relation) type, and $d$ - the destination node type.

The message passing update rule that we implement is very similar to that of GraphSAGE, except we now need to include the node types and the edge relation type. The update rule for message type $m$ is described below:

\begin{equation}
h_v^{(l)[m]} = W^{(l)[m]} \cdot \text{CONCAT} \Big( W_d^{(l)[m]} \cdot h_v^{(l-1)}, W_s^{(l)[m]} \cdot AGG(\{h_u^{(l-1)}, \forall u \in N_{m}(v) \})\Big)
\end{equation}

where we compute $h_v^{(l)[m]}$, the node embedding representation for node $v$ after `HeteroGNNConv` layer $l$ with respect message type $m$. Further unpacking the forumla we have:
- $W_s^{(l)[m]}$ - linear transformation matrix for the messages of neighboring source nodes of type $s$ along message type $m$.
- $W_d^{(l)[m]}$ - linear transformation matrix for the message from the node $v$ itself of type $d$.
- $W^{(l)[m]}$ - linear transformation matrix for the concatenated messages from neighboring node's and the central node.
- $h_u^{(l-1)}$ - the hidden embedding representation for node $u$ after the $(l-1)$th `HeteroGNNWrapperConv` layer. Note, that this embedding is not associated with a particular message type (see layer diagrams above).
- $N_{m}(v)$ - the set of neighbor source nodes $s$ for the node v that we are embedding along message type $m = (s, r, d)$.

**NOTE**: We emphasize that each weight matrix is associated with a specific message type $[m]$ and additionally, the weight matrices applied to node messages are differentiated by node type (i.e. $W_s$ and $W_d$).

Lastly, for simplicity, we use mean aggregations for $AGG$ where:

\begin{equation}
AGG(\{h_u^{(l-1)}, \forall u \in N_{m}(v) \}) = \frac{1}{|N_{m}(v)|} \sum_{u\in N_{m}(v)} h_u^{(l-1)}
\end{equation}

```python
class HeteroGNNConv(pyg_nn.MessagePassing):
    def __init__(self, in_channels_src, in_channels_dst, out_channels):
        super(HeteroGNNConv, self).__init__(aggr="mean")

        self.in_channels_src = in_channels_src
        self.in_channels_dst = in_channels_dst
        self.out_channels = out_channels

        # To simplify implementation, please initialize both self.lin_dst
        # and self.lin_src out_features to out_channels
        self.lin_dst = None
        self.lin_src = None

        self.lin_update = None

        ############# Your code here #############
        ## (~3 lines of code)
        ## Note:
        ## 1. Initialize the 3 linear layers.
        ## 2. Think through the connection between the mathematical
        ##    definition of the update rule and torch linear layers!

        pass

        ##########################################

    def forward(
        self,
        node_feature_src,
        node_feature_dst,
        edge_index,
        size=None
    ):
        ############# Your code here #############
        ## (~1 line of code)
        ## Note:
        ## 1. Unlike Colabs 3 and 4, we just need to call self.propagate with
        ## proper/custom arguments.

        pass

        ##########################################

    def message_and_aggregate(self, edge_index, node_feature_src):

        out = None
        ############# Your code here #############
        ## (~1 line of code)
        ## Note:
        ## 1. Different from what we implemented in Colabs 3 and 4, we use message_and_aggregate
        ##    to combine the previously seperate message and aggregate functions.
        ##    The benefit is that we can avoid materializing x_i and x_j
        ##    to make the implementation more efficient.
        ## 2. To implement efficiently, refer to PyG documentation for message_and_aggregate
        ##    and sparse-matrix multiplication:
        ##    https://pytorch-geometric.readthedocs.io/en/latest/notes/sparse_tensor.html
        ## 3. Here edge_index is torch_sparse SparseTensor. Although interesting, you
        ##    do not need to deeply understand SparseTensor represenations!
        ## 4. Conceptually, think through how the message passing and aggregation
        ##    expressed mathematically can be expressed through matrix multiplication.

        pass

        ##########################################

        return out

    def update(self, aggr_out, node_feature_dst):

        ############# Your code here #############
        ## (~4 lines of code)
        ## Note:
        ## 1. The update function is called after message_and_aggregate
        ## 2. Think through the one-one connection between the mathematical update
        ##    rule and the 3 linear layers defined in the constructor.

       pass

        ##########################################

        return aggr_out
```

## Heterogeneous GNN Wrapper Layer

After implementing the `HeteroGNNConv` layer for each message type, we need to manage and aggregate the node embedding results (with respect to each message types). Here we will implement two types of message type level aggregation.

The first one is simply mean aggregation over message types:

\begin{equation}
h_v^{(l)} = \frac{1}{M}\sum_{m=1}^{M}h_v^{(l)[m]}
\end{equation}

where node $v$ has node type $d$ and we sum over the $M$ message types that have destination node type $d$. From our original example, for a node v of type $b$ we aggregate v's `HeteroGNNConv` embeddings for message types $m_2$ and $m_3$ (i.e. $h_v^{(l)[m_2]}$ and $h_v^{(l)[m_3]}$).

The second method we implement is the semantic level attention introduced in **HAN** ([Wang et al. (2019)](https://arxiv.org/abs/1903.07293)). Instead of directly averaging on the message type aggregation results, we use attention to learn which message type result is more important, then aggregate across all the message types. Below are the equations for semantic level attention:

\begin{equation}
e_{m} = \frac{1}{|V_{d}|} \sum_{v \in V_{d}} q_{attn}^T \cdot tanh \Big( W_{attn}^{(l)} \cdot h_v^{(l)[m]} + b \Big)
\end{equation}

where $m$ is the message type and $d$ refers to the destination node type for that message ($m = (s, r, d)$). Additionally, $V_{d}$ refers to the set of nodes v with type $d$. Lastly, the unormalized attention weight $e_m$ is a scaler computed for each message type $m$.

Next, we can compute the normalized attention weights and update $h_v^{(l)}$:

\begin{equation}
\alpha_{m} = \frac{\exp(e_{m})}{\sum_{m=1}^M \exp(e_{m})}
\end{equation}

\begin{equation}
h_v^{(l)} = \sum_{m=1}^{M} \alpha_{m} \cdot h_v^{(l)[m]}
\end{equation}

, where we emphasize that $M$ here is the number of message types associated with the destination node type $d$.

**Note**: The implementation of the attention aggregation is tricky and nuanced. We strongly recommend working carefully through the math equations to undersatnd exactly what each notation refers to and how all the pieces fit together. If you can, try to connect the math to our original example, focusing on node type $b$, which depends on two different message types!

**_We've implemented most of this for you but you'll need to initialize self.attn_proj in the initializer_**

```python
class HeteroGNNWrapperConv(deepsnap.hetero_gnn.HeteroConv):
    def __init__(self, convs, args, aggr="mean"):
        super(HeteroGNNWrapperConv, self).__init__(convs, None)
        self.aggr = aggr

        # Map the index and message type
        self.mapping = {}

        # A numpy array that stores the final attention probability
        self.alpha = None

        self.attn_proj = None

        if self.aggr == "attn":
            ############# Your code here #############
            ## (~1 line of code)
            ## Note:
            ## 1. Initialize self.attn_proj, where self.attn_proj should include
            ##    two linear layers. Note, make sure you understand
            ##    which part of the equation self.attn_proj captures.
            ## 2. You should use nn.Sequential for self.attn_proj
            ## 3. nn.Linear and nn.Tanh are useful.
            ## 4. You can model a weight vector (rather than matrix) by using:
            ##    nn.Linear(some_size, 1, bias=False).
            ## 5. The first linear layer should have out_features as args['attn_size']
            ## 6. You can assume we only have one "head" for the attention.
            ## 7. We recommend you to implement the mean aggregation first. After
            ##    the mean aggregation works well in the training, then you can
            ##    implement this part.

            pass

            ##########################################

    def reset_parameters(self):
        super(HeteroConvWrapper, self).reset_parameters()
        if self.aggr == "attn":
            for layer in self.attn_proj.children():
                layer.reset_parameters()

    def forward(self, node_features, edge_indices):
        message_type_emb = {}
        for message_key, message_type in edge_indices.items():
            src_type, edge_type, dst_type = message_key
            node_feature_src = node_features[src_type]
            node_feature_dst = node_features[dst_type]
            edge_index = edge_indices[message_key]
            message_type_emb[message_key] = (
                self.convs[message_key](
                    node_feature_src,
                    node_feature_dst,
                    edge_index,
                )
            )
        node_emb = {dst: [] for _, _, dst in message_type_emb.keys()}
        mapping = {}
        for (src, edge_type, dst), item in message_type_emb.items():
            mapping[len(node_emb[dst])] = (src, edge_type, dst)
            node_emb[dst].append(item)
        self.mapping = mapping
        for node_type, embs in node_emb.items():
            if len(embs) == 1:
                node_emb[node_type] = embs[0]
            else:
                node_emb[node_type] = self.aggregate(embs)
        return node_emb

    def aggregate(self, xs):
        # TODO: Implement this function that aggregates all message type results.
        # Here, xs is a list of tensors (embeddings) with respect to message
        # type aggregation results.

        if self.aggr == "mean":

            ############# Your code here #############
            ## (~2 lines of code)
            ## Note:
            ## 1. Explore the function parameter `xs`!

            pass

            ##########################################

        elif self.aggr == "attn":
            N = xs[0].shape[0] # Number of nodes for that node type
            M = len(xs) # Number of message types for that node type

            x = torch.cat(xs, dim=0).view(M, N, -1) # M * N * D
            z = self.attn_proj(x).view(M, N) # M * N * 1
            z = z.mean(1) # M * 1
            alpha = torch.softmax(z, dim=0) # M * 1

            # Store the attention result to self.alpha as np array
            self.alpha = alpha.view(-1).data.cpu().numpy()

            alpha = alpha.view(M, 1, 1)
            x = x * alpha
            return x.sum(dim=0)
```

## Initialize Heterogeneous GNN Layers

Now let's put it all together and initialize the Heterogeneous GNN Layers. Different from the homogeneous graph case, heterogeneous graphs can be a little bit complex.

In general, we need to create a dictionary of `HeteroGNNConv` layers where the keys are message types.

* To get all message types, `deepsnap.hetero_graph.HeteroGraph.message_types` is useful.
* If we are initializing the first conv layers, we need to get the feature dimension of each node type. Using `deepsnap.hetero_graph.HeteroGraph.num_node_features(node_type)` will return the node feature dimension of `node_type`. In this function, we will set each `HeteroGNNConv` `out_channels` to be `hidden_size`.
* If we are not initializing the first conv layers, all node types will have the smae embedding dimension `hidden_size` and we still set `HeteroGNNConv` `out_channels` to be `hidden_size` for simplicity.

```python
def generate_convs(hetero_graph, conv, hidden_size, first_layer=False):
    # TODO: Implement this function that returns a dictionary of `HeteroGNNConv`
    # layers where the keys are message types. `hetero_graph` is deepsnap `HeteroGraph`
    # object and the `conv` is the `HeteroGNNConv`.

    convs = {}

    ############# Your code here #############
    ## (~9 lines of code)
    ## Note:
    ## 1. See the hints above!
    ## 2. conv is of type `HeteroGNNConv`

    pass

    ##########################################

    return convs
```

## HeteroGNN

Now we will make a simple HeteroGNN model which contains only two `HeteroGNNWrapperConv` layers.

For the forward function in `HeteroGNN`, the model is going to be run as following:

$\text{self.convs1} \rightarrow \text{self.bns1} \rightarrow \text{self.relus1} \rightarrow \text{self.convs2} \rightarrow \text{self.bns2} \rightarrow \text{self.relus2} \rightarrow \text{self.post_mps}$

```python
class HeteroGNN(torch.nn.Module):
    def __init__(self, hetero_graph, args, aggr="mean"):
        super(HeteroGNN, self).__init__()

        self.aggr = aggr
        self.hidden_size = args['hidden_size']

        self.convs1 = None
        self.convs2 = None

        self.bns1 = nn.ModuleDict()
        self.bns2 = nn.ModuleDict()
        self.relus1 = nn.ModuleDict()
        self.relus2 = nn.ModuleDict()
        self.post_mps = nn.ModuleDict()

        ############# Your code here #############
        ## (~10 lines of code)
        ## Note:
        ## 1. For self.convs1 and self.convs2, call generate_convs at first and then
        ##    pass the returned dictionary of `HeteroGNNConv` to `HeteroGNNWrapperConv`.
        ## 2. For self.bns, self.relus and self.post_mps, the keys are node_types.
        ##    `deepsnap.hetero_graph.HeteroGraph.node_types` will be helpful.
        ## 3. Initialize all batchnorms to torch.nn.BatchNorm1d(hidden_size, eps=1).
        ## 4. Initialize all relus to nn.LeakyReLU().
        ## 5. For self.post_mps, each value in the ModuleDict is a linear layer
        ##    where the `out_features` is the number of classes for that node type.
        ##    `deepsnap.hetero_graph.HeteroGraph.num_node_labels(node_type)` will be
        ##    useful.

        pass

        ##########################################

    def forward(self, node_feature, edge_index):
        # TODO: Implement the forward function. Notice that `node_feature` is
        # a dictionary of tensors where keys are node types and values are
        # corresponding feature tensors. The `edge_index` is a dictionary of
        # tensors where keys are message types and values are corresponding
        # edge index tensors (with respect to each message type).

        x = node_feature

        ############# Your code here #############
        ## (~7 lines of code)
        ## Note:
        ## 1. `deepsnap.hetero_gnn.forward_op` can be helpful.

        pass

        ##########################################

        return x

    def loss(self, preds, y, indices):

        loss = 0
        loss_func = F.cross_entropy

        ############# Your code here #############
        ## (~3 lines of code)
        ## Note:
        ## 1. For each node type in preds, accumulate computed loss to `loss`
        ## 2. Loss need to be computed with respect to the given index
        ## 3. preds is a dictionary of model predictions keyed by node_type.
        ## 4. indeces is a dictionary of labeled supervision nodes keyed
        ##    by node_type

        pass

        ##########################################

        return loss
```

## Training and Testing

Here we provide you with the functions to train and test. You only need to implement one line of code here.

**Please do not modify other parts in `train` and `test` for grading purposes.**

```python
import pandas as pd

def train(model, optimizer, hetero_graph, train_idx):
    model.train()
    optimizer.zero_grad()
    preds = model(hetero_graph.node_feature, hetero_graph.edge_index)

    loss = None

    ############# Your code here #############
    ## Note:
    ## 1. Compute the loss here
    ## 2. `deepsnap.hetero_graph.HeteroGraph.node_label` is useful

    pass

    ##########################################

    loss.backward()
    optimizer.step()
    return loss.item()

def test(model, graph, indices, best_model=None, best_val=0, save_preds=False, agg_type=None):
    model.eval()
    accs = []
    for i, index in enumerate(indices):
        preds = model(graph.node_feature, graph.edge_index)
        num_node_types = 0
        micro = 0
        macro = 0
        for node_type in preds:
            idx = index[node_type]
            pred = preds[node_type][idx]
            pred = pred.max(1)[1]
            label_np = graph.node_label[node_type][idx].cpu().numpy()
            pred_np = pred.cpu().numpy()
            micro = f1_score(label_np, pred_np, average='micro')
            macro = f1_score(label_np, pred_np, average='macro')
            num_node_types += 1

        # Averaging f1 score might not make sense, but in our example we only
        # have one node type
        micro /= num_node_types
        macro /= num_node_types
        accs.append((micro, macro))

        # Only save the test set predictions and labels!
        if save_preds and i == 2:
          print ("Saving Heterogeneous Node Prediction Model Predictions with Agg:", agg_type)
          print()

          data = {}
          data['pred'] = pred_np
          data['label'] = label_np

          df = pd.DataFrame(data=data)
          # Save locally as csv
          df.to_csv('ACM-Node-' + agg_type + 'Agg.csv', sep=',', index=False)

    if accs[1][0] > best_val:
        best_val = accs[1][0]
        best_model = copy.deepcopy(model)
    return accs, best_model, best_val
```

```python
# Please do not change the following parameters
args = {
    'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
    'hidden_size': 64,
    'epochs': 100,
    'weight_decay': 1e-5,
    'lr': 0.003,
    'attn_size': 32,
}
```

## Dataset and Preprocessing

In the next, we will load the data and create a tensor backend (without a NetworkX graph) `deepsnap.hetero_graph.HeteroGraph` object.

We will use the `ACM(3025)` dataset in our node property prediction task, which is proposed in **HAN** ([Wang et al. (2019)](https://arxiv.org/abs/1903.07293)) and our dataset is extracted from [DGL](https://www.dgl.ai/)'s [ACM.mat](https://data.dgl.ai/dataset/ACM.mat).

The original ACM dataset has three node types and two edge (relation) types. For simplicity, we simplify the heterogeneous graph to one node type and two edge types (shown below). This means that in our heterogeneous graph, we have one node type (paper) and two message types (paper, author, paper) and (paper, subject, paper).

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  print("Device: {}".format(args['device']))

  # Load the data
  data = torch.load("acm.pkl")

  # Message types
  message_type_1 = ("paper", "author", "paper")
  message_type_2 = ("paper", "subject", "paper")

  # Dictionary of edge indices
  edge_index = {}
  edge_index[message_type_1] = data['pap']
  edge_index[message_type_2] = data['psp']

  # Dictionary of node features
  node_feature = {}
  node_feature["paper"] = data['feature']

  # Dictionary of node labels
  node_label = {}
  node_label["paper"] = data['label']

  # Load the train, validation and test indices
  train_idx = {"paper": data['train_idx'].to(args['device'])}
  val_idx = {"paper": data['val_idx'].to(args['device'])}
  test_idx = {"paper": data['test_idx'].to(args['device'])}

  # Construct a deepsnap tensor backend HeteroGraph
  hetero_graph = HeteroGraph(
      node_feature=node_feature,
      node_label=node_label,
      edge_index=edge_index,
      directed=True
  )

  print(f"ACM heterogeneous graph: {hetero_graph.num_nodes()} nodes, {hetero_graph.num_edges()} edges")

  # Node feature and node label to device
  for key in hetero_graph.node_feature:
      hetero_graph.node_feature[key] = hetero_graph.node_feature[key].to(args['device'])
  for key in hetero_graph.node_label:
      hetero_graph.node_label[key] = hetero_graph.node_label[key].to(args['device'])

  # Edge_index to sparse tensor and to device
  for key in hetero_graph.edge_index:
      edge_index = hetero_graph.edge_index[key]
      adj = SparseTensor(row=edge_index[0], col=edge_index[1], sparse_sizes=(hetero_graph.num_nodes('paper'), hetero_graph.num_nodes('paper')))
      hetero_graph.edge_index[key] = adj.t().to(args['device'])
  print(hetero_graph.edge_index[message_type_1])
  print(hetero_graph.edge_index[message_type_2])
```

## Start Training!

Now lets start training!

## Training the Mean Aggregation

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  best_model = None
  best_val = 0

  model = HeteroGNN(hetero_graph, args, aggr="mean").to(args['device'])
  optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay'])

  for epoch in range(args['epochs']):
      loss = train(model, optimizer, hetero_graph, train_idx)
      accs, best_model, best_val = test(model, hetero_graph, [train_idx, val_idx, test_idx], best_model, best_val)
      print(
          f"Epoch {epoch + 1}: loss {round(loss, 5)}, "
          f"train micro {round(accs[0][0] * 100, 2)}%, train macro {round(accs[0][1] * 100, 2)}%, "
          f"valid micro {round(accs[1][0] * 100, 2)}%, valid macro {round(accs[1][1] * 100, 2)}%, "
          f"test micro {round(accs[2][0] * 100, 2)}%, test macro {round(accs[2][1] * 100, 2)}%"
      )
  best_accs, _, _ = test(best_model, hetero_graph, [train_idx, val_idx, test_idx], save_preds=True, agg_type="Mean")
  print(
      f"Best model: "
      f"train micro {round(best_accs[0][0] * 100, 2)}%, train macro {round(best_accs[0][1] * 100, 2)}%, "
      f"valid micro {round(best_accs[1][0] * 100, 2)}%, valid macro {round(best_accs[1][1] * 100, 2)}%, "
      f"test micro {round(best_accs[2][0] * 100, 2)}%, test macro {round(best_accs[2][1] * 100, 2)}%"
  )
```

## Question 2.1: What is your maximum test set **micro** F1 score for the best_model when using mean aggregation? (10 points)

## Question 2.2: What is your maximum test set **macro** F1 score for the best_model when using the mean aggregation? (10 points)

## Training the Attention Aggregation

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  best_model = None
  best_val = 0

  output_size = hetero_graph.num_node_labels('paper')
  model = HeteroGNN(hetero_graph, args, aggr="attn").to(args['device'])
  optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'], weight_decay=args['weight_decay'])

  for epoch in range(args['epochs']):
      loss = train(model, optimizer, hetero_graph, train_idx)
      accs, best_model, best_val = test(model, hetero_graph, [train_idx, val_idx, test_idx], best_model, best_val)
      print(
          f"Epoch {epoch + 1}: loss {round(loss, 5)}, "
          f"train micro {round(accs[0][0] * 100, 2)}%, train macro {round(accs[0][1] * 100, 2)}%, "
          f"valid micro {round(accs[1][0] * 100, 2)}%, valid macro {round(accs[1][1] * 100, 2)}%, "
          f"test micro {round(accs[2][0] * 100, 2)}%, test macro {round(accs[2][1] * 100, 2)}%"
      )
  best_accs, _, _ = test(best_model, hetero_graph, [train_idx, val_idx, test_idx], save_preds=True, agg_type="Attention")
  print(
      f"Best model: "
      f"train micro {round(best_accs[0][0] * 100, 2)}%, train macro {round(best_accs[0][1] * 100, 2)}%, "
      f"valid micro {round(best_accs[1][0] * 100, 2)}%, valid macro {round(best_accs[1][1] * 100, 2)}%, "
      f"test micro {round(best_accs[2][0] * 100, 2)}%, test macro {round(best_accs[2][1] * 100, 2)}%"
  )
```

## Question 2.3: What is your maximum test set **micro** F1 score for the best_model when using the attention aggregation? (4 points)

## Question 2.4: What is your maximum test set **macro** F1 score for the best_model when using the attention aggregation? (4 points)

## Attention for each Message Type

Through message type level attention we can learn which message type is more important to which layer.

Here we will print out and show that each layer pay how much attention on each message type.

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
  if model.convs1.alpha is not None and model.convs2.alpha is not None:
      for idx, message_type in model.convs1.mapping.items():
          print(f"Layer 1 has attention {model.convs1.alpha[idx]} on message type {message_type}")
      for idx, message_type in model.convs2.mapping.items():
          print(f"Layer 2 has attention {model.convs2.alpha[idx]} on message type {message_type}")
```

# Submission

You will need to submit three files on Gradescope to complete this notebook.

1.   Your completed *CS224W_Colab5.ipynb*. From the "File" menu select "Download .ipynb" to save a local copy of your completed Colab.
2.  *ACM-Node-MeanAgg.csv*
3.  *ACM-Node-AttentionAgg.csv*

Download the csv files by selecting the *Folder* icon on the left panel.

To submit your work, zip the files downloaded in steps 1-3 above and submit to gradescope. **NOTE:** DO NOT rename any of the downloaded files.


```python
#remove cell visibility
from IPython.display import HTML
tag = HTML('''
Toggle cell visibility <a href="javascript:code_toggle()">here</a>.''')
display(tag)
```

Toggle cell visibility here.

```python
%matplotlib inline
import control
import numpy
import sympy as sym
from IPython.display import display, Markdown
import ipywidgets as widgets
import matplotlib.pyplot as plt

#print a matrix latex-like
def bmatrix(a):
     """Returns a LaTeX bmatrix - by Damir Arbula (ICCT project)

     :a: numpy array
     :returns: LaTeX bmatrix as a string
     """
     if len(a.shape) > 2:
         raise ValueError('bmatrix can at most display two dimensions')
     lines = str(a).replace('[', '').replace(']', '').splitlines()
     rv = [r'\begin{bmatrix}']
     rv += ['  ' + ' & '.join(l.split()) + r'\\' for l in lines]
     rv +=  [r'\end{bmatrix}']
     return '\n'.join(rv)

# Display formatted matrix:
def vmatrix(a):
    if len(a.shape) > 2:
         raise ValueError('bmatrix can at most display two dimensions')
    lines = str(a).replace('[', '').replace(']', '').splitlines()
    rv = [r'\begin{vmatrix}']
    rv += ['  ' + ' & '.join(l.split()) + r'\\' for l in lines]
    rv +=  [r'\end{vmatrix}']
    return '\n'.join(rv)

#matrixWidget is a matrix looking widget built with a VBox of HBox(es) that returns a numPy array as value !
class matrixWidget(widgets.VBox):
    def updateM(self,change):
        for irow in range(0,self.n):
            for icol in range(0,self.m):
                self.M_[irow,icol] = self.children[irow].children[icol].value
                #print(self.M_[irow,icol])
        self.value = self.M_

    def dummychangecallback(self,change):
        pass

    def __init__(self,n,m):
        self.n = n
        self.m = m
        self.M_ = numpy.matrix(numpy.zeros((self.n,self.m)))
        self.value = self.M_
        widgets.VBox.__init__(self,
                             children = [
                                 widgets.HBox(children =
                                              [widgets.FloatText(value=0.0, layout=widgets.Layout(width='90px')) for i in range(m)]
                                             )
                                 for j in range(n)
                             ])

        #fill in widgets and tell interact to call updateM each time a children changes value
        for irow in range(0,self.n):
            for icol in range(0,self.m):
                self.children[irow].children[icol].value = self.M_[irow,icol]
                self.children[irow].children[icol].observe(self.updateM, names='value')
        #value = Unicode('example@example.com', help="The email value.").tag(sync=True)
        self.observe(self.updateM, names='value', type= 'All')

    def setM(self, newM):
        #disable callbacks, change values, and reenable
        self.unobserve(self.updateM, names='value', type= 'All')
        for irow in range(0,self.n):
            for icol in range(0,self.m):
                self.children[irow].children[icol].unobserve(self.updateM, names='value')
        self.M_ = newM
        self.value = self.M_
        for irow in range(0,self.n):
            for icol in range(0,self.m):
                self.children[irow].children[icol].value = self.M_[irow,icol]
        for irow in range(0,self.n):
            for icol in range(0,self.m):
                self.children[irow].children[icol].observe(self.updateM, names='value')
        self.observe(self.updateM, names='value', type= 'All')

                #self.children[irow].children[icol].observe(self.updateM, names='value')

#overlaod class for state space systems that DO NOT remove "useless" states (what "professor" of automatic control would do this?)
class sss(control.StateSpace):
    def __init__(self,*args):
        #call base class init constructor
        control.StateSpace.__init__(self,*args)
    #disable function below in base class
    def _remove_useless_states(self):
        pass
```

## Control of robotic arm with flexible joint

A robotic arm link is driven by an electric motor via a flexible joint that behaves like a torsional spring. The system dynamics may be approximated with a third order linear time invariant system where the states are:
- $x_1$ : the difference between motor and arm angles (not zero because of the joint flexibility),
- $x_2$ : motor shaft angular speed,
- $x_3$ : link angular speed.
The input $u$ is the motor torque. The dynamic equations are:

\begin{cases}
\dot{x} = \begin{bmatrix} 0 & 1 & -1 \\ a-1 & -b_1 & b_1 \\ a & b_2 & -b_2 \end{bmatrix}x + \begin{bmatrix} 0 \\ b_3 \\ 0 \end{bmatrix}u \\
y = \begin{bmatrix} 0 & 0 & 1 \end{bmatrix}x
\end{cases}

with $a=0.1$, $b_1=0.09$, $b_2=0.01$ and $b_3=90$.

The goal of the control system design is to regulate the link angular speed in order to have dominant poles with damping equal to 0.7 and natural frequency equal to 0.5 rad/s and zero steady-state error in response to reference velocity step.

The transfer function of the system is:

```python
A = numpy.matrix('0 1 -1; -0.9 -0.09 0.09; 0.1 0.01 -0.01')
B = numpy.matrix('0; 90; 0')
C = numpy.matrix('0 0 1')
D = numpy.matrix('0')
sys_tf = control.tf(sss(A,B,C,D))
print(sys_tf)
```

      3.886e-16 s^2 + 0.9 s + 9
    -----------------------------
    s^3 + 0.1 s^2 + s - 8.674e-19

with poles

```python
import warnings
# In order to suppress the warning BadCoefficient
warnings.filterwarnings("ignore")
print(numpy.round(sys_tf.pole(),3))
```

    [-0.05+0.999j -0.05-0.999j  0.  +0.j   ]

and zeros

```python
print(numpy.round(sys_tf.zero(),3),'.')
```

    [-2.31613695e+15 -1.00000000e+01] .

First, we analyze the system in order to verify if it is controllable and observable. The controllability matrix $\mathcal{C}$ is

```python
Ctrb = control.ctrb(A,B)
display(Markdown(bmatrix(Ctrb)))
# print(numpy.linalg.matrix_rank(Ctrb))
```

\begin{bmatrix}
  0. & 90. & -9.\\
  90. & -8.1 & -80.19\\
  0. & 0.9 & 8.91\\
\end{bmatrix}

and has rank equal to 3 so the system is controllable. The observability matrix $\mathcal{O}$ is

```python
Obsv = control.obsv(A,C)
display(Markdown(bmatrix(Obsv)))
# print(numpy.linalg.matrix_rank(Obsv))
```

\begin{bmatrix}
  0. & 0. & 1.\\
  0.1 & 0.01 & -0.01\\
  -0.01 & 0.099 & -0.099\\
\end{bmatrix}

and has rank equal to 3 so the system is observable.

This could have been actually deduced from the fact that the denominator of the transfer function is of third order (equal to the state space vector dimension).

### Regulator design
#### Controller design

Due to the requirements, we know that we have to place 2 poles in $\zeta \omega_n \pm \sqrt{1-\zeta^2}\omega_n = -0.35\pm0.357i$ and place the remaining real pole at a frequency higher than that of the complex (dominant) poles. We can select to place the third pole at -3.5 rad/s. For the requirement of zero steady-state error we scale the reference signal with a  gain equal to the inverse of the closed-loop system gain.

#### Design of the observer

In order to have an observer that quickly assists the controller we simply place the poles at about -10 rad/s.

### How to use this notebook?
- Verify if the closed-loop system performs well in a case of an initial state estimation error. Try to improve its performance.
- Reduce the frequency of the real pole of the controlled closed-loop system and watch how the response differs from the reference response.

```python
# Preparatory cell

X0 = numpy.matrix('0.0; 0.0; 0.0')
K = numpy.matrix([8/15,-4.4,-4])
L = numpy.matrix([[23],[66],[107/3]])

Aw = matrixWidget(3,3)
Aw.setM(A)
Bw = matrixWidget(3,1)
Bw.setM(B)
Cw = matrixWidget(1,3)
Cw.setM(C)
X0w = matrixWidget(3,1)
X0w.setM(X0)
Kw = matrixWidget(1,3)
Kw.setM(K)
Lw = matrixWidget(3,1)
Lw.setM(L)

eig1c = matrixWidget(1,1)
eig2c = matrixWidget(2,1)
eig3c = matrixWidget(1,1)
eig1c.setM(numpy.matrix([-3.5]))
eig2c.setM(numpy.matrix([[-0.35],[-0.357]]))
eig3c.setM(numpy.matrix([-3.5]))

eig1o = matrixWidget(1,1)
eig2o = matrixWidget(2,1)
eig3o = matrixWidget(1,1)
eig1o.setM(numpy.matrix([-10.]))
eig2o.setM(numpy.matrix([[-10.],[0.]]))
eig3o.setM(numpy.matrix([-10.]))
```

```python
# Misc

#create dummy widget
DW = widgets.FloatText(layout=widgets.Layout(width='0px', height='0px'))

#create button widget
START = widgets.Button(
    description='Test',
    disabled=False,
    button_style='', # 'success', 'info', 'warning', 'danger' or ''
    tooltip='Test',
    icon='check'
)

def on_start_button_clicked(b):
    #This is a workaround to have intreactive_output call the callback:
    #   force the value of the dummy widget to change
    if DW.value> 0 :
        DW.value = -1
    else:
        DW.value = 1
    pass
START.on_click(on_start_button_clicked)

# Define type of method
selm = widgets.Dropdown(
    options= ['Set K and L', 'Set the eigenvalues'],
    value= 'Set the eigenvalues',
    description='',
    disabled=False
)

# Define the number of complex eigenvalues
sele = widgets.Dropdown(
    options= ['0 complex eigenvalues', '2 complex eigenvalues'],
    value= '2 complex eigenvalues',
    description='Complex eigenvalues:',
    style = {'description_width': 'initial'},
    disabled=False
)

#define type of ipout
selu = widgets.Dropdown(
    options=['impulse', 'step', 'sinusoid', 'square wave'],
    value='step',
    description='Type of reference:',
    style = {'description_width': 'initial'},
    disabled=False
)
# Define the values of the input
u = widgets.FloatSlider(
    value=1,
    min=0,
    max=3,
    step=0.1,
    description='Reference:',
    disabled=False,
    continuous_update=False,
    orientation='horizontal',
    readout=True,
    readout_format='.1f',
)
period = widgets.FloatSlider(
    value=0.5,
    min=0.001,
    max=10,
    step=0.001,
    description='Period: ',
    disabled=False,
    continuous_update=False,
    orientation='horizontal',
    readout=True,
    readout_format='.2f',
)

gain_w2 = widgets.FloatText(
    value=1.,
    description='',
    disabled=True
)

simTime = widgets.FloatText(
    value=20,
    description='',
    disabled=False
)
```

```python
# Support functions

def eigen_choice(sele):
    if sele == '0 complex eigenvalues':
        eig1c.children[0].children[0].disabled = False
        eig2c.children[1].children[0].disabled = True
        eig1o.children[0].children[0].disabled = False
        eig2o.children[1].children[0].disabled = True
        eig = 0
    if sele == '2 complex eigenvalues':
        eig1c.children[0].children[0].disabled = True
        eig2c.children[1].children[0].disabled = False
        eig1o.children[0].children[0].disabled = True
        eig2o.children[1].children[0].disabled = False
        eig = 2
    return eig

def method_choice(selm):
    if selm == 'Set K and L':
        method = 1
        sele.disabled = True
    if selm == 'Set the eigenvalues':
        method = 2
        sele.disabled = False
    return method
```

```python
s = control.tf('s')
Gref = (0.5)**2/(s**2 + 2*0.7*0.5*s + (0.5)**2)

def main_callback2(Aw, Bw, X0w, K, L, eig1c, eig2c, eig3c, eig1o, eig2o, eig3o, u, period, selm, sele, selu, simTime, DW):
    eige = eigen_choice(sele)
    method = method_choice(selm)

    if method == 1:
        solc = numpy.linalg.eig(A-B*K)
        solo = numpy.linalg.eig(A-L*C)
    if method == 2:
        if eige == 0:
            K = control.acker(A, B, [eig1c[0,0], eig2c[0,0], eig3c[0,0]])
            Kw.setM(K)
            L = control.acker(A.T, C.T, [eig1o[0,0], eig2o[0,0], eig3o[0,0]]).T
            Lw.setM(L)
        if eige == 2:
            K = control.acker(A, B, [eig3c[0,0],
                                     numpy.complex(eig2c[0,0],eig2c[1,0]),
                                     numpy.complex(eig2c[0,0],-eig2c[1,0])])
            Kw.setM(K)
            L = control.acker(A.T, C.T, [eig3o[0,0],
                                         numpy.complex(eig2o[0,0],eig2o[1,0]),
                                         numpy.complex(eig2o[0,0],-eig2o[1,0])]).T
            Lw.setM(L)

    sys = control.ss(A,B,numpy.vstack((C,numpy.zeros((B.shape[1],C.shape[1])))),numpy.vstack((D,numpy.eye(B.shape[1]))))
    sysC = control.ss(numpy.zeros((1,1)),
                      numpy.zeros((1,numpy.shape(A)[0])),
                      numpy.zeros((numpy.shape(B)[1],1)),
                      -K)

    sysE = control.ss(A-L*C,
                      numpy.hstack((L,B-L*D)),
                      numpy.eye(numpy.shape(A)[0]),
                      numpy.zeros((A.shape[0],C.shape[0]+B.shape[1])))

    sys_append = control.append(sys, sysE, sysC, control.ss(A,B,numpy.eye(A.shape[0]),numpy.zeros((A.shape[0],B.shape[1]))))
    Q = []
    # y in ingresso a sysE
    for i in range(C.shape[0]):
        Q.append([B.shape[1]+i+1, i+1])
    # u in ingresso a sysE
    for i in range(B.shape[1]):
        Q.append([B.shape[1]+C.shape[0]+i+1, C.shape[0]+i+1])
    # u in ingresso a sys
    for i in range(B.shape[1]):
        Q.append([i+1, C.shape[0]+B.shape[1]+A.shape[0]+i+1])
    # u in ingresso al sistema che ha come uscite gli stati reali
    for i in range(B.shape[1]):
        Q.append([2*B.shape[1]+C.shape[0]+A.shape[0]+i+1, C.shape[0]+i+1])
    # xe in ingresso a sysC
    for i in range(A.shape[0]):
        Q.append([2*B.shape[1]+C.shape[0]+i+1, C.shape[0]+B.shape[1]+i+1])

    inputv = [i+1 for i in range(B.shape[1])]
    outputv = [i+1 for i in range(numpy.shape(sys_append.C)[0])]
    sys_CL = control.connect(sys_append,
                             Q,
                             inputv,
                             outputv)

    t = numpy.linspace(0, 100000, 2)
    t, yout = control.step_response(sys_CL[0,0],T=t)
    dcgain = yout[-1]
    gain_w2.value = dcgain
    if dcgain != 0:
        u1 = u/gain_w2.value
    else:
        print('The feedforward gain setted is 0 and it is changed to 1')
        u1 = u/1
    print('The static gain of the closed-loop system (from the reference to the output) is: %.5f' %dcgain)

    X0w1 = numpy.zeros((A.shape[0],1))
    for j in range(A.shape[0]):
        X0w1 = numpy.vstack((X0w1,X0w[j]))
    X0w1 = numpy.vstack((X0w1,numpy.zeros((A.shape[0],1))))
    if simTime != 0:
        T = numpy.linspace(0, simTime, 10000)
    else:
        T = numpy.linspace(0, 1, 10000)

    if selu == 'impulse': #selu
        U = [0 for t in range(0,len(T))]
        U[0] = u
        U1 = [0 for t in range(0,len(T))]
        U1[0] = u1
        T, yout, xout = control.forced_response(sys_CL,T,U1,X0w1)
        T, yout_ref, xout_ref = control.forced_response(Gref,T,U,[0, 0])
    if selu == 'step':
        U = [u for t in range(0,len(T))]
        U1 = [u1 for t in range(0,len(T))]
        T, yout, xout = control.forced_response(sys_CL,T,U1,X0w1)
        T, yout_ref, xout_ref = control.forced_response(Gref,T,U,[0, 0])
    if selu == 'sinusoid':
        U = u*numpy.sin(2*numpy.pi/period*T)
        U1 = u1*numpy.sin(2*numpy.pi/period*T)
        T, yout, xout = control.forced_response(sys_CL,T,U1,X0w1)
        T, yout_ref, xout_ref = control.forced_response(Gref,T,U,[0, 0])
    if selu == 'square wave':
        U = u*numpy.sign(numpy.sin(2*numpy.pi/period*T))
        U1 = u1*numpy.sign(numpy.sin(2*numpy.pi/period*T))
        T, yout, xout = control.forced_response(sys_CL,T,U1,X0w1)
        T, yout_ref, xout_ref = control.forced_response(Gref,T,U,[0, 0])
    # N.B. i primi 3 stati di xout sono quelli del sistema, mentre gli ultimi 3 sono quelli dell'osservatore

    step_info_dict = control.step_info(sys_CL[0,0],SettlingTimeThreshold=0.05,T=T)
    print('Step info: \n\tRise time =',step_info_dict['RiseTime'],'\n\tSettling time (5%) =',step_info_dict['SettlingTime'],'\n\tOvershoot (%)=',step_info_dict['Overshoot'])
    # print('Max x3 value (%)=', max(abs(yout[C.shape[0]+2*B.shape[1]+A.shape[0]+2]))/(numpy.pi/180*17)*100)

    fig = plt.figure(num='Simulation1', figsize=(14,12))

    fig.add_subplot(221)
    plt.title('Output response')
    plt.ylabel('Output')
    plt.plot(T,yout[0],T,yout_ref,T,U,'r--')
    plt.xlabel('$t$ [s]')
    plt.axvline(x=0,color='black',linewidth=0.8)
    plt.axhline(y=0,color='black',linewidth=0.8)
    plt.legend(['$y$','Second-order system reference','Reference'])
    plt.grid()

    fig.add_subplot(222)
    plt.title('Input')
    plt.ylabel('$u$')
    plt.plot(T,yout[C.shape[0]])
    plt.xlabel('$t$ [s]')
    plt.axvline(x=0,color='black',linewidth=0.8)
    plt.axhline(y=0,color='black',linewidth=0.8)
    plt.grid()

    fig.add_subplot(223)
    plt.title('States response')
    plt.ylabel('States')
    plt.plot(T,yout[C.shape[0]+2*B.shape[1]+A.shape[0]],
             T,yout[C.shape[0]+2*B.shape[1]+A.shape[0]+1],
             T,yout[C.shape[0]+2*B.shape[1]+A.shape[0]+2])
    plt.xlabel('$t$ [s]')
    plt.axvline(x=0,color='black',linewidth=0.8)
    plt.axhline(y=0,color='black',linewidth=0.8)
    plt.legend(['$x_{1}$','$x_{2}$','$x_{3}$'])
    plt.grid()

    fig.add_subplot(224)
    plt.title('Estimation errors')
    plt.ylabel('Errors')
    plt.plot(T,yout[C.shape[0]+2*B.shape[1]+A.shape[0]]-yout[C.shape[0]+B.shape[1]],
             T,yout[C.shape[0]+2*B.shape[1]+A.shape[0]+1]-yout[C.shape[0]+B.shape[1]+1],
             T,yout[C.shape[0]+2*B.shape[1]+A.shape[0]+2]-yout[C.shape[0]+B.shape[1]+2])
    plt.xlabel('$t$ [s]')
    plt.axvline(x=0,color='black',linewidth=0.8)
    plt.axhline(y=0,color='black',linewidth=0.8)
    plt.legend(['$e_{1}$','$e_{2}$','$e_{3}$'])
    plt.grid()
    #plt.tight_layout()

alltogether2 = widgets.VBox([widgets.HBox([selm,
                                          sele,
                                          selu]),
                            widgets.Label(' ',border=3),
                            widgets.HBox([widgets.Label('K:',border=3), Kw,
                                          widgets.Label(' ',border=3),
                                          widgets.Label(' ',border=3),
                                          widgets.Label('Eigenvalues:',border=3),
                                          eig1c,
                                          eig2c,
                                          eig3c,
                                          widgets.Label(' ',border=3),
                                          widgets.Label(' ',border=3),
                                          widgets.Label('X0 est.:',border=3), X0w]),
                            widgets.Label(' ',border=3),
                            widgets.HBox([widgets.Label('L:',border=3), Lw,
                                          widgets.Label(' ',border=3),
                                          widgets.Label(' ',border=3),
                                          widgets.Label('Eigenvalues:',border=3),
                                          eig1o,
                                          eig2o,
                                          eig3o,
                                          widgets.Label(' ',border=3),
                                          widgets.VBox([widgets.Label('Inverse reference gain:',border=3),
                                                        widgets.Label('Simulation time [s]:',border=3)]),
                                          widgets.VBox([gain_w2,simTime])]),
                            widgets.Label(' ',border=3),
                            widgets.HBox([u,
                                          period,
                                          START])])
out2 = widgets.interactive_output(main_callback2, {'Aw':Aw, 'Bw':Bw, 'X0w':X0w, 'K':Kw, 'L':Lw,
                                                 'eig1c':eig1c, 'eig2c':eig2c, 'eig3c':eig3c, 'eig1o':eig1o, 'eig2o':eig2o, 'eig3o':eig3o,
                                                 'u':u, 'period':period, 'selm':selm, 'sele':sele, 'selu':selu, 'simTime':simTime, 'DW':DW})
out2.layout.height = '860px'
display(out2, alltogether2)
```

    Output(layout=Layout(height='860px'))

    VBox(children=(HBox(children=(Dropdown(index=1, options=('Set K and L', 'Set the eigenvalues'), value='Set the…

```python

```

```python

```