{ "cells": [ { "cell_type": "markdown", "id": "ca821458", "metadata": {}, "source": [ "{download}`Download this notebook <5_quantum_internet.ipynb>`" ] }, { "cell_type": "markdown", "id": "5aef78fd", "metadata": {}, "source": [ "## Optimizing quantum internet\n", "> **Concepts demonstrated below**\n", "> - Cost matrix (edge construction costs are proportional to node separations)\n", "> - Loss landscapes riddled with local minima\n", "> - Batching\n", "> - Noise in `GradNet`\n", "\n", "### Problem setup\n", "Here, we optimize the communication capacity of a quantum internet with lossy channels. The geographic placement of quantum computers are sampled randomly, and wiring costs are assumed to be proportional to the distance between nodes. We compute the pairwise distances $D_{ij}$, and optimize the adjacency matrix $A_{ij}$ under the constrained construction budget $b=\\frac{1}{2}\\sum_{ij}D_{ij}A_{ij}$.\n", "\n", "Following the derivations in [1], we define the transmissivity of each edge $e$ as $\\eta_e$. We model the edge transmissivities from the edge weights using the mapping $\\eta_{ij} = \\tanh{A_{ij}}$. The composite transmissivity along some path $P=\\{e_1, e_2, ..., e_K\\}$ is given by the bottleneck transmissivity along the chain $\\eta_P = \\min_{e\\in P}\\eta_e$ (see [1] for details). And the capacity of this chain is given by\n", "\n", "$$\n", "\\mathcal C_P = -\\log_2(1-\\eta_P).\n", "$$\n", "\n", "The communication capacity between two nodes $(i,j)$ in a network is the capacity of the best path $P_{ij}$ connecting them\n", "\n", "$$\n", "\\mathcal C_{ij} = \\max_{P_{ij}} C_P = -\\log_2(1-\\max_{P_{ij}}\\min_{e\\in P}\\eta_e).\n", "$$\n", "\n", "The mean network capacity, and the target of our optimization, is then given by \n", "\n", "$$\n", "\\mathcal C = \\frac{1}{n(n-1)} \\sum_{i\\ne j} {\\mathcal C_{ij}} = -\\frac{1}{n(n-1)} \\sum_{i\\ne j}\\log_2(1-\\max_{P_{ij}}\\min_{e\\in P}\\eta_e).\n", "$$\n", "\n", "[1] Pirandola, S., 2019. End-to-end capacities of a quantum communication network. *Communications Physics*, 2(1), p.51.\n" ] }, { "cell_type": "markdown", "id": "ab2fee11", "metadata": {}, "source": [ "### Install\n", "install the required dependencies silently" ] }, { "cell_type": "code", "execution_count": 1, "id": "aae9a798", "metadata": {}, "outputs": [], "source": [ "%%capture\n", "!pip install 'gradnet[examples]'" ] }, { "cell_type": "markdown", "id": "7058dac0", "metadata": {}, "source": [ "### Functions for computing channel capacities and the loss function " ] }, { "cell_type": "code", "execution_count": 2, "id": "b5338994", "metadata": {}, "outputs": [], "source": [ "import torch\n", "import matplotlib.pyplot as plt\n", "from gradnet import GradNet, fit\n", "from gradnet.utils import positions_to_distance_matrix, plot_graph, random_seed\n", "import networkx as nx\n", "import numpy as np\n", "\n", "\n", "def mean_pairwise_min_cut(A: torch.Tensor) -> torch.Tensor:\n", " \"\"\"\n", " Compute the mean value of pairwise minimum s–t cuts on an undirected\n", " weighted graph with adjacency/capacity matrix A (nonnegative, symmetric).\n", " \"\"\"\n", " assert A.ndim == 2 and A.shape[0] == A.shape[1], \"A must be square\"\n", " n = A.shape[0]\n", " device = A.device\n", " dtype = A.dtype\n", "\n", " # Build NetworkX graph using *detached* capacities, so no grad flows here.\n", " A_cpu = A.detach().cpu().numpy()\n", " G = nx.Graph()\n", " G.add_nodes_from(range(n))\n", " for i in range(n):\n", " for j in range(i + 1, n):\n", " w = float(A_cpu[i, j])\n", " if w > 0.0: # treat zero as 'no edge'\n", " G.add_edge(i, j, capacity=w)\n", "\n", " # Accumulate total cut value as a Torch scalar on the original device.\n", " total = A.new_tensor(0.0)\n", "\n", " # Loop over unordered pairs (s, t)\n", " for s in range(n):\n", " for t in range(s + 1, n):\n", " # Min s–t cut: we only use the partition, *not* the numeric value\n", " cut_value, partition = nx.minimum_cut(G, s, t, capacity=\"capacity\")\n", " S, T = partition # S and T = V \\ S\n", "\n", " # Build a boolean mask for nodes in S on the Torch device\n", " s_mask = torch.zeros(n, dtype=torch.bool, device=device)\n", " s_indices = torch.tensor(list(S), dtype=torch.long, device=device)\n", " s_mask[s_indices] = True\n", "\n", " # Edges crossing the cut: exactly one endpoint in S\n", " cross = s_mask.unsqueeze(0) ^ s_mask.unsqueeze(1) # XOR\n", " # For an undirected graph with symmetric A, count each edge once:\n", " cross = torch.triu(cross, diagonal=1)\n", "\n", " # Cut value expressed as linear combination of A entries\n", " cut_val_torch = (A * cross.to(dtype)).sum()\n", " total = total + cut_val_torch\n", "\n", " num_pairs = n * (n - 1) // 2\n", " mean_cut = total / num_pairs\n", " return mean_cut\n", "\n", "\n", "def multi_path_routing_capacity(A):\n", " eta = torch.tanh(A)\n", " c = - torch.log(1-eta)/torch.log(torch.tensor(2.0))\n", " return mean_pairwise_min_cut(c)\n", "\n", "\n", "def average_channel_capacity(A):\n", " ''' computes the mean channel capacity in the given weighted network based on reference [1] '''\n", " eta = torch.tanh(A).clone() # edge transmissivities\n", " n = eta.size(0)\n", " for k in range(n): # widest-path Floyd–Warshall (max-min)\n", " via_k = torch.minimum(eta[:, k].unsqueeze(1), eta[k, :].unsqueeze(0))\n", " eta = torch.maximum(eta, via_k)\n", " eta = eta.clamp(max=1-1e-12) # path transmissivities\n", " C = -torch.log2(1 - eta) # path capacities\n", " mask = ~torch.eye(n, dtype=bool, device=A.device)\n", " return C[mask].mean()\n", "\n", "\n", "def loss_fn_constructor(batch_size, noise_amplitude):\n", " def loss_fn(gn):\n", " capacities = torch.zeros(batch_size)\n", " for i in range(batch_size):\n", " A = gn(noise_amplitude=noise_amplitude)\n", " capacities[i] = average_channel_capacity(A)\n", " # capacities[i] = multi_path_routing_capacity(A)\n", " return -capacities.mean(), {'capacity': average_channel_capacity(gn()).item()}\n", " return loss_fn" ] }, { "cell_type": "markdown", "id": "00ce85ce", "metadata": {}, "source": [ "### Optimization" ] }, { "cell_type": "code", "execution_count": 3, "id": "ea9e2d55", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "GPU available: True (mps), used: False\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n", "Missing logger folder: lightning_logs/updates-100_lr-0.005_batch100_noise1.0\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "ada7fab172e14246a05486f83053e4e3", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Updates: 0%| | 0/100 [00:00" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stderr", "output_type": "stream", "text": [ "GPU available: True (mps), used: False\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n", "Missing logger folder: lightning_logs/updates-100_lr-0.005_batch100_noise0.8333333333333334\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "8caee61a35884626a607a2bbdbd3fc29", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Updates: 0%| | 0/100 [00:00" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stderr", "output_type": "stream", "text": [ "GPU available: True (mps), used: False\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n", "Missing logger folder: lightning_logs/updates-100_lr-0.005_batch100_noise0.6666666666666667\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "32f13deb05164ab59fbbf02bcca0d566", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Updates: 0%| | 0/100 [00:00" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stderr", "output_type": "stream", "text": [ "GPU available: True (mps), used: False\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n", "Missing logger folder: lightning_logs/updates-100_lr-0.005_batch100_noise0.5\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "d3ebf41523264aa98935433792fb39ed", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Updates: 0%| | 0/100 [00:00" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stderr", "output_type": "stream", "text": [ "GPU available: True (mps), used: False\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n", "Missing logger folder: lightning_logs/updates-200_lr-0.005_batch100_noise0.33333333333333337\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "e394b9f378544a408c50e9eb6460b988", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Updates: 0%| | 0/200 [00:00" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stderr", "output_type": "stream", "text": [ "GPU available: True (mps), used: False\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n", "Missing logger folder: lightning_logs/updates-1000_lr-0.005_batch100_noise0.16666666666666674\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "2169c74b72a64897b298e7dd86e3ec4e", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Updates: 0%| | 0/1000 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "N = 50\n", "L = 10.0\n", "budget_per_node = 1\n", "\n", "seed = 0\n", "random_seed(seed)\n", "\n", "positions = torch.rand(N, 2) * L # random coordinates in an LxL square\n", "cost_matrix = positions_to_distance_matrix(positions) # all-to-all distances\n", "\n", "pos_dict = {idx: pos.tolist() for idx, pos in enumerate(positions)}\n", "draw_kwargs = {\n", " \"pos\": pos_dict,\n", " \"with_labels\": False,\n", " \"node_size\": 20\n", "}\n", "\n", "updates = [100, 100, 100, 100, 200, 1000]\n", "lrs = [0.005]*7\n", "batch_sizes = [100]*6 + [1]\n", "noises = np.linspace(1, 0, 7)\n", "\n", "gn = GradNet(\n", " num_nodes=N,\n", " budget=budget_per_node * N,\n", " cost_matrix=cost_matrix,\n", " rand_init_weights=0\n", ")\n", "\n", "for update, lr, batch_size, noise in zip(updates, lrs, batch_sizes, noises):\n", " loss_fn = loss_fn_constructor(batch_size=batch_size, noise_amplitude=noise)\n", " fit(\n", " gn=gn,\n", " loss_fn=loss_fn,\n", " num_updates=update,\n", " optim_kwargs={\"lr\": lr},\n", " enable_checkpointing=False,\n", " logger=True,\n", " log_dir=f'lightning_logs/updates-{update}_lr-{lr}_batch{batch_size}_noise{noise}',\n", " accelerator=\"cpu\",\n", " verbose=True,\n", " )\n", " plot_graph(gn, layout=\"networkx\", draw_kwargs=dict(draw_kwargs))\n", " plt.pause(0.1)\n", "\n", "# plot_graph(gn, layout=\"networkx\", draw_kwargs=dict(draw_kwargs))" ] } ], "metadata": { "kernelspec": { "display_name": "torch", "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.11.8" } }, "nbformat": 4, "nbformat_minor": 5 }