3.1-GNN-intro.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ignore scanpy warnings\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import scipy\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import squidpy as sq\n",
    "import scanpy as sc\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import Parameter\n",
    "import torch.nn.functional as F\n",
    "import torch_geometric as pyg\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.nn import MessagePassing, GCNConv, MLP\n",
    "from torch_geometric.nn.dense.linear import Linear\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = sq.datasets.merfish()\n",
    "adata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing\n",
    "Our data is stored in Anndata format, which behaves similar to pandas DataFrames. We can access the data matrix using the `.X` attribute. We can also access the cell and gene names using the `.obs_names` and `.var_names` attributes, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select one slice\n",
    "bregma = 1\n",
    "adata = adata[adata.obs[\"Bregma\"] == bregma, :].copy()\n",
    "\n",
    "# filtering\n",
    "sc.pp.filter_cells(adata, min_counts=10)\n",
    "sc.pp.filter_genes(adata, min_cells=5)\n",
    "\n",
    "# normalization\n",
    "adata.layers[\"counts\"] = adata.X.copy()\n",
    "sc.pp.normalize_total(adata, inplace=True)\n",
    "sc.pp.log1p(adata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_gene_idx = adata.X.toarray().sum(axis=0).argmax()\n",
    "top_gene = adata.var_names[top_gene_idx]\n",
    "print(f\"The gene with the highest total expression is {top_gene}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.spatial(adata, color=[\"Cell_class\", top_gene], spot_size=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construct graph by connecting each cells to its k nearest neighbors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we want to connect each node to its k nearest neighbors, while ensuring the graph is undirected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 10\n",
    "\n",
    "coords = adata.obsm[\"spatial\"]\n",
    "kdtree = scipy.spatial.KDTree(coords)\n",
    "distances, indices = kdtree.query(coords, k=k + 1, p=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Indices:\\n\", indices)\n",
    "print(f\"Indices shape: {indices.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_index = torch.cat(\n",
    "    [\n",
    "        torch.tensor(indices.flatten())[None, :],  # source\n",
    "        torch.arange(0, coords.shape[0]).repeat_interleave(k + 1)[None, :],  # target\n",
    "    ],\n",
    "    axis=0,\n",
    ")\n",
    "edge_weight = torch.tensor(distances.flatten()).unsqueeze(-1).to(torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_index, edge_weight = pyg.utils.to_undirected(edge_index, edge_weight)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The graph is undirected: {pyg.utils.is_undirected(edge_index)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets visualize the graph using networkx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = Data(edge_index=edge_index, num_nodes=coords.shape[0])\n",
    "g = pyg.utils.to_networkx(data, to_undirected=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove self-loops for better visualization\n",
    "g.remove_edges_from(nx.selfloop_edges(g))\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "nx.draw(g, pos=coords, node_size=10, width=0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 1: construct a graph by connecting cells only if they are within a certain distance of each other\n",
    "Hint: you can loop over all possible pairs of cells using a nested for loop and calculate the distance between them using the `np.linalg.norm` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radius = 0.1\n",
    "\n",
    "# define empty matrices to hold the results\n",
    "dist_mat = np.zeros((coords.shape[0], coords.shape[0]))\n",
    "adj_mat = np.zeros((coords.shape[0], coords.shape[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "insert your code here\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Scipy provides some optimized functions for this task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kdtree = scipy.spatial.KDTree(coords)\n",
    "dist_mat_hat = kdtree.sparse_distance_matrix(kdtree, radius, p=2)\n",
    "dist_mat_hat = scipy.sparse.csr_matrix(dist_mat_hat)\n",
    "adj_mat_hat = (dist_mat_hat > 0).toarray().astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    \"The results are the same:\",\n",
    "    (adj_mat_hat + np.eye(adj_mat_hat.shape[0]) == adj_mat).all(),\n",
    ")\n",
    "\n",
    "print(\n",
    "    \"The results are the same:\",\n",
    "    (dist_mat_hat.toarray() == dist_mat).all(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute the cell below if you want to run the subsequent computation on the distance graph instead of the nearest neighbor graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# edge_index, edge_weight = pyg.utils.from_scipy_sparse_matrix(dist_mat_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 2: aggregate the gene expression per neighborhood\n",
    "This is our target per cell neighborhood for the graph autoencoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In a first version we perform the aggregation with a simple matrix multiplication (for the distance-based case we could have directly used the adj mat.)  \n",
    "Hint 1: you can convert the edge_index to a sparse matrix using the `pyg.utils.to_scipy_sparse_matrix` function  \n",
    "Hint 2: the gene expression is stored in the `.X` attribute of the Anndata object  \n",
    "Hint 3: make sure to normalize the adjacency matrix by the degree of each node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "insert your code here, name the result X_agg\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can perform the same cell aggregation in the Pytorch Geometric Message Passing Framework"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphAggregation(MessagePassing):\n",
    "    \"\"\"\n",
    "    GraphAggregation class for aggregating node features in a graph.\n",
    "\n",
    "    Args:\n",
    "        aggr (str): Aggregation method to use. Default is \"mean\".\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, aggr=\"mean\"):\n",
    "        super(GraphAggregation, self).__init__(aggr=aggr)\n",
    "\n",
    "    def forward(self, x, edge_index, **kwargs):\n",
    "        \"\"\"\n",
    "        Forward pass of the GraphAggregation module.\n",
    "\n",
    "        Args:\n",
    "            x (Tensor): Node features.\n",
    "            edge_index (LongTensor): Graph edge indices.\n",
    "\n",
    "        Returns:\n",
    "            Tensor: Aggregated node features.\n",
    "\n",
    "        \"\"\"\n",
    "        return self.propagate(edge_index, x=x)\n",
    "\n",
    "    def message(self, x_j):\n",
    "        \"\"\"\n",
    "        Message function for the GraphAggregation module.\n",
    "\n",
    "        Args:\n",
    "            x_j (Tensor): Node features of neighboring nodes.\n",
    "\n",
    "        Returns:\n",
    "            Tensor: The input node features.\n",
    "\n",
    "        \"\"\"\n",
    "        return x_j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_agg = GraphAggregation(aggr=\"mean\")\n",
    "X = torch.Tensor(adata.X.toarray())\n",
    "X_agg_pyg = mean_agg(X, edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    \"Results are the same (up to numeric error):\",\n",
    "    np.allclose(X_agg.toarray(), X_agg_pyg.numpy(), atol=1e-9),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 3: define your own graph convolutional network in Pytorch Geometric"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The equation for the graph convolutional layer is given by:\n",
    "$$\n",
    "H^{i+1} = (\\hat{D}^{-\\frac{1}{2}} \\hat{A} \\hat{D}^{-\\frac{1}{2}}) H^{i} W + b\n",
    "$$\n",
    "However, you can ignore the normalization for now (already handled by `gcn_norm`) and just implement the last part of the equation:\n",
    "$$\n",
    "H^{i+1} = \\dots H^{i} W + b\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gcn_norm(edge_index, num_nodes, edge_weight=None, dtype=None):\n",
    "    \"\"\"\n",
    "    Applies graph convolutional network (GCN) normalization to the given edge index and edge weight.\n",
    "\n",
    "    Args:\n",
    "        edge_index (Tensor): The edge index tensor of shape (2, num_edges) representing the connectivity of the graph.\n",
    "        num_nodes (int): The total number of nodes in the graph.\n",
    "        edge_weight (Tensor, optional): The edge weight tensor of shape (num_edges,) representing the weight of each edge. Defaults to None.\n",
    "        dtype (torch.dtype, optional): The desired data type of the edge weight tensor. Defaults to None.\n",
    "\n",
    "    Returns:\n",
    "        Tuple[Tensor, Tensor]: A tuple containing the updated edge index tensor and the normalized edge weight tensor.\n",
    "\n",
    "    \"\"\"\n",
    "    if edge_weight is None:\n",
    "        edge_weight = torch.ones(\n",
    "            (edge_index.size(1),), device=edge_index.device, dtype=dtype\n",
    "        )\n",
    "\n",
    "    row, col = edge_index[0], edge_index[1]\n",
    "    idx = col\n",
    "    deg = pyg.utils.scatter(edge_weight, idx, dim=0, dim_size=num_nodes, reduce=\"sum\")\n",
    "\n",
    "    deg_inv_sqrt = deg.pow_(-0.5)\n",
    "    deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float(\"inf\"), 0)\n",
    "    edge_weight = deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]\n",
    "\n",
    "    return edge_index, edge_weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GCNLayer(MessagePassing):\n",
    "    \"\"\"Graph Convolutional Network (GCN) layer implementation.\n",
    "\n",
    "    This class represents a single layer of a Graph Convolutional Network (GCN).\n",
    "    It performs message passing and aggregation operations on a graph.\n",
    "\n",
    "    Args:\n",
    "        in_channels (int): Number of input channels/features.\n",
    "        out_channels (int): Number of output channels/features.\n",
    "        normalize (bool, optional): Whether to normalize the edge weights. Defaults to True.\n",
    "        aggr (str, optional): Aggregation method for message passing. Defaults to \"add\".\n",
    "\n",
    "    Attributes:\n",
    "        linear (torch.nn.Linear): Linear transformation layer.\n",
    "        bias (torch.nn.Parameter): Bias parameter.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, normalize=True, aggr=\"add\"):\n",
    "        super(GCNLayer, self).__init__(aggr=aggr)\n",
    "        self.normalize = normalize\n",
    "\n",
    "        self.linear = Linear(\n",
    "            in_channels, out_channels, bias=False, weight_initializer=\"glorot\"\n",
    "        )\n",
    "        self.bias = Parameter(torch.zeros(out_channels))\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        super().reset_parameters()\n",
    "        self.linear.reset_parameters()\n",
    "\n",
    "    def forward(self, x, edge_index, edge_weight=None):\n",
    "        if self.normalize:\n",
    "            edge_index, edge_weight = gcn_norm(\n",
    "                edge_index, x.size(0), edge_weight, dtype=x.dtype\n",
    "            )\n",
    "        \"\"\"\n",
    "        Here you need to implement the missing parts of the forward pass of the GCN layer. you need to use the self.linear and self.bias\n",
    "        \"\"\"\n",
    "\n",
    "        # here something is missing\n",
    "        x = ...\n",
    "\n",
    "        # this part is correct\n",
    "        out = self.propagate(edge_index, x=x, edge_weight=edge_weight)\n",
    "\n",
    "        # here something is missing\n",
    "        out = ...\n",
    "        return out\n",
    "\n",
    "    def message(self, x_i, x_j, edge_weight=None):\n",
    "        # the message function allows you to access the node features of the source nodes (x_i) and the target nodes (x_j)\n",
    "        # any other node level attributes that you pass to the propagate function can also be accessed here\n",
    "        return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the seed for reproducibility\n",
    "torch.manual_seed(42)\n",
    "\n",
    "gcn_layer_custom = GCNLayer(adata.X.shape[1], 32)\n",
    "h_custom = gcn_layer_custom(X_agg_pyg, edge_index)\n",
    "h_custom"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check that it gives the same result as the built-in GCN implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the seed for reproducibility\n",
    "torch.manual_seed(42)\n",
    "\n",
    "gcn_layer = GCNConv(adata.X.shape[1], 32, add_self_loops=False)\n",
    "h = gcn_layer(X_agg_pyg, edge_index)\n",
    "h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Results are the same:\", torch.allclose(h, h_custom, atol=1e-9))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus 1: formulate the same architecture in native Pytorch using matrix multiplications\n",
    "Hint: for simplicity you can reuse the gcn_norm function and convert the adjacency matrix to a tensor "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus 2: implement a more general Message Passing Layer\n",
    "Hint: adapt the Pytorch Geometric Message Passing class according to the general message passing equation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define your node-level graph autoencoder\n",
    "This model takes as input a cell graph with gene expression features and learns a latent representation of the cell neighborhood by reconstructing the neighborhood gene expression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphEncoder(nn.Module):\n",
    "    \"\"\"GraphEncoder is a class that represents a graph encoder module.\n",
    "\n",
    "    Args:\n",
    "        in_channels (int): The number of input channels.\n",
    "        hidden_channels (int): The number of hidden channels.\n",
    "        n_layers (int, optional): The number of graph convolutional layers. Defaults to 2.\n",
    "        normalize (bool, optional): Whether to apply normalization. Defaults to True.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, hidden_channels, n_layers=2, normalize=True):\n",
    "        super(GraphEncoder, self).__init__()\n",
    "        self.linear = Linear(in_channels, hidden_channels)\n",
    "        self.convs = nn.ModuleList(\n",
    "            [\n",
    "                GCNLayer(hidden_channels, hidden_channels, normalize=normalize)\n",
    "                for _ in range(n_layers)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def forward(self, x, edge_index, edge_weight=None):\n",
    "        x = self.linear(x)\n",
    "        for conv in self.convs:\n",
    "            x = conv(x, edge_index, edge_weight)\n",
    "            x = F.relu(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    \"\"\"Decoder module for the GNN model, uses the Pytorch Geometric MLP for convenience.\n",
    "\n",
    "    Args:\n",
    "        in_channels (int): Number of input channels.\n",
    "        hidden_channels (int): Number of hidden channels.\n",
    "        out_channels (int): Number of output channels.\n",
    "        n_layers (int, optional): Number of MLP layers. Defaults to 2.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self, in_channels, hidden_channels, out_channels, n_layers=2, **kwargs\n",
    "    ):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.mlp = MLP(\n",
    "            in_channels=in_channels,\n",
    "            hidden_channels=hidden_channels,\n",
    "            out_channels=out_channels,\n",
    "            num_layers=n_layers,\n",
    "            plain_last=False,\n",
    "            norm=None,\n",
    "            **kwargs,\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.mlp(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphAutoEncoder(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels,\n",
    "        hidden_channels,\n",
    "        out_channels,\n",
    "        n_layers_encoder=1,\n",
    "        n_layers_decoder=1,\n",
    "        **kwargs,\n",
    "    ):\n",
    "        \"\"\"Initializes a GraphAutoEncoder object.\n",
    "\n",
    "        Args:\n",
    "            in_channels (int): Number of input channels.\n",
    "            hidden_channels (int): Number of hidden channels.\n",
    "            out_channels (int): Number of output channels.\n",
    "            n_layers_encoder (int, optional): Number of layers in the encoder. Defaults to 1.\n",
    "            n_layers_decoder (int, optional): Number of layers in the decoder. Defaults to 1.\n",
    "            **kwargs: Additional keyword arguments.\n",
    "\n",
    "        Returns:\n",
    "            None\n",
    "        \"\"\"\n",
    "        super(GraphAutoEncoder, self).__init__()\n",
    "        self.encoder = GraphEncoder(\n",
    "            in_channels, hidden_channels, n_layers=n_layers_encoder, **kwargs\n",
    "        )\n",
    "        self.decoder = Decoder(\n",
    "            hidden_channels, hidden_channels, out_channels, n_layers=n_layers_decoder\n",
    "        )\n",
    "\n",
    "    def forward(self, x, edge_index, edge_weight=None):\n",
    "        h = self.encoder(x, edge_index, edge_weight)\n",
    "        x = self.decoder(h)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train your model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"You are using {device}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 42\n",
    "n_epochs = 400\n",
    "n_genes = adata.X.shape[1]\n",
    "n_layers_encoder = 1\n",
    "n_layers_decoder = 1\n",
    "\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed(seed)\n",
    "\n",
    "model = GraphAutoEncoder(n_genes, 32, n_genes, n_layers_encoder, n_layers_decoder)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    out = model(X_agg_pyg, edge_index, edge_weight)\n",
    "    loss = criterion(out, X_agg_pyg)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    if epoch % 10 == 0:\n",
    "        print(f\"Epoch {epoch}, Loss: {loss.item()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 4: Extract and visualize the latent embeddings\n",
    "Hint: our model has an encoder that maps the input gene expression to the latent space. We can use this encoder to extract the latent embeddings of the cells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "insert your code here\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Encoded features shape: {h.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obsm[\"X_gnn\"] = h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hint to visualize the embeddings you can use PCA or UMAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "insert your code here\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define spatial domains via Leiden clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.neighbors(adata, use_rep=\"X_gnn\")\n",
    "sc.tl.leiden(adata, resolution=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize the spatial domains"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.spatial(adata, color=\"leiden\", spot_size=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Task 5: Analyze the cell type proportions for each spatial domain\n",
    "Hint: the only data you need is given in the below DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = adata.obs[[\"leiden\", \"Cell_class\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Further reading\n",
    "1. [Geometric deep learning resources](https://geometricdeeplearning.com/)\n",
    "2. [Graph (variational) autoencoder paper](https://arxiv.org/abs/1611.07308)\n",
    "3. [Pytorch Geometric documentation](https://pytorch-geometric.readthedocs.io/en/latest/)\n",
    "4. [Google tuning playbook](https://github.com/google-research/tuning_playbook)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spacet",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}