{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "51c97544-aee1-493d-94c5-0db0c4db6a96",
   "metadata": {},
   "source": [
    "# Channel Models from Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c854a11-320e-4d9d-bf4b-3b60e2d408f6",
   "metadata": {},
   "source": [
    "In this notebook, you will learn how to create a channel model from a [generator](https://wiki.python.org/moin/Generators). This can be used, e.g., to import datasets of channel impulse responses."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f25c031b-7c18-4942-8301-781e9555e000",
   "metadata": {},
   "source": [
    "* [Configuration and Imports](#Configuration-and-Imports)\n",
    "* [Simulation Parameters](#Simulation-Parameters)\n",
    "* [Creating a Simple Dataset](#Creating-a-Simple-Dataset)\n",
    "* [Generators](#Generators)\n",
    "* [Use the Channel Model for OFDM Transmissions](#Use-the-Channel-Model-for-OFDM-Transmissions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0c41452-1a20-4b53-a02f-a2981ee42836",
   "metadata": {},
   "source": [
    "## Configuration and Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bac912ac-4a02-4e58-a697-e53d021852a5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:25.810030Z",
     "iopub.status.busy": "2026-02-13T14:13:25.809921Z",
     "iopub.status.idle": "2026-02-13T14:13:28.630286Z",
     "shell.execute_reply": "2026-02-13T14:13:28.629359Z"
    }
   },
   "outputs": [],
   "source": [
    "# Import Sionna\n",
    "try:\n",
    "    import sionna.phy\n",
    "except ImportError as e:\n",
    "    import os\n",
    "    import sys\n",
    "    if 'google.colab' in sys.modules:\n",
    "       # Install Sionna in Google Colab\n",
    "       print(\"Installing Sionna and restarting the runtime. Please run the cell again.\")\n",
    "       os.system(\"pip install sionna\")\n",
    "       os.kill(os.getpid(), 5)\n",
    "    else:\n",
    "       raise e\n",
    "\n",
    "import numpy as np\n",
    "import h5py\n",
    "\n",
    "config = sionna.phy.config\n",
    "config.seed = 42 # Set seed for reproducible random number generation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d02a418-ff93-4d27-a519-7c9e563ae91c",
   "metadata": {},
   "source": [
    "## Simulation Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "01c892be-713d-401a-b6de-e984fd426d6f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:28.632675Z",
     "iopub.status.busy": "2026-02-13T14:13:28.632397Z",
     "iopub.status.idle": "2026-02-13T14:13:28.635170Z",
     "shell.execute_reply": "2026-02-13T14:13:28.634481Z"
    }
   },
   "outputs": [],
   "source": [
    "num_rx = 2\n",
    "num_rx_ant = 2\n",
    "num_tx = 1\n",
    "num_tx_ant = 8\n",
    "num_time_steps = 100\n",
    "num_paths = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "632ffb75-c8d0-47e2-9861-ad12daffa20b",
   "metadata": {},
   "source": [
    "## Creating a Simple Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4bd653e-59b6-4a22-8201-3d1d7024ca6d",
   "metadata": {},
   "source": [
    "To illustrate how to load dataset, we will first create one.\n",
    "\n",
    "The next cell creates a very small HDF5 file storing Gaussian distributed i.i.d. path coefficients and uniformly distributed i.i.d. path delays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a54379aa-b75c-448a-b5f8-4e97ebd1d282",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:28.637147Z",
     "iopub.status.busy": "2026-02-13T14:13:28.637023Z",
     "iopub.status.idle": "2026-02-13T14:13:29.583881Z",
     "shell.execute_reply": "2026-02-13T14:13:29.583007Z"
    }
   },
   "outputs": [],
   "source": [
    "# Number of examples in the dataset\n",
    "dataset_size = 1000\n",
    "\n",
    "# Random path coefficients\n",
    "a_shape = [dataset_size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_paths, num_time_steps]\n",
    "a = (config.np_rng.normal(size=a_shape) + 1j*config.np_rng.normal(size=a_shape))/np.sqrt(2)\n",
    "\n",
    "# Random path delays\n",
    "tau = config.np_rng.uniform(size=[dataset_size, num_rx, num_tx, num_paths])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "206c14de-660a-4112-ad65-948595af540e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:29.585983Z",
     "iopub.status.busy": "2026-02-13T14:13:29.585831Z",
     "iopub.status.idle": "2026-02-13T14:13:29.977435Z",
     "shell.execute_reply": "2026-02-13T14:13:29.976497Z"
    }
   },
   "outputs": [],
   "source": [
    "filename = 'my_dataset.h5'\n",
    "hf = h5py.File(filename, 'w')\n",
    "hf.create_dataset('a', data=a)\n",
    "hf.create_dataset('tau', data=tau)\n",
    "hf.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e83cf683-609e-4d88-b26e-9c6df4bc2d3a",
   "metadata": {},
   "source": [
    "## Generators"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ed9063-4d5b-4097-8b57-9005dcb9b597",
   "metadata": {},
   "source": [
    "The first step to load a dataset is to create a [generator](https://wiki.python.org/moin/Generators).\n",
    "A generator is a callable object, i.e., a function or a class that implements the `__call__()` method, and that behaves like an iterator.\n",
    "\n",
    "The next cell shows how to create a generator that parses an HDF5 file storing path coefficients and delays.\n",
    "Note that how the HDF5 file is parsed depends on its structure. The following generator is specific to the dataset previously created.\n",
    "\n",
    "If you have another dataset, you will need to change the way it is parsed in the generator. The generator can also carry out any type of desired pre-processing of your data, e.g., normalization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "787f5ba7-38ae-4ecd-8dde-028e45eb6a39",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:29.979528Z",
     "iopub.status.busy": "2026-02-13T14:13:29.979396Z",
     "iopub.status.idle": "2026-02-13T14:13:29.982525Z",
     "shell.execute_reply": "2026-02-13T14:13:29.981822Z"
    }
   },
   "outputs": [],
   "source": [
    "class HD5CIRGen:\n",
    "    def __init__(self, filename):\n",
    "        self.filename = filename\n",
    "\n",
    "    def __call__(self):\n",
    "        with h5py.File(self.filename, 'r') as hf:\n",
    "            for im in zip(hf[\"a\"], hf[\"tau\"]):\n",
    "                a = im[0]\n",
    "                tau = im[1]\n",
    "                # One could do some preprocessing on the dataset here\n",
    "                # ...\n",
    "                yield im"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7584d077-2261-4dab-bf16-fd5ac3b93901",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:29.984760Z",
     "iopub.status.busy": "2026-02-13T14:13:29.984641Z",
     "iopub.status.idle": "2026-02-13T14:13:29.986943Z",
     "shell.execute_reply": "2026-02-13T14:13:29.986235Z"
    }
   },
   "outputs": [],
   "source": [
    "generator = HD5CIRGen(filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "113c36b6-dd8c-4aa0-9a9d-5dc4db7312e7",
   "metadata": {},
   "source": [
    "We can use the generator to sample the first 5 items of the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7a30d141-0f79-4bbe-9293-9d163e0e33b3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:29.988855Z",
     "iopub.status.busy": "2026-02-13T14:13:29.988728Z",
     "iopub.status.idle": "2026-02-13T14:13:29.997386Z",
     "shell.execute_reply": "2026-02-13T14:13:29.996721Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 2, 1, 8, 10, 100)\n",
      "(2, 1, 10)\n",
      "(2, 2, 1, 8, 10, 100)\n",
      "(2, 1, 10)\n",
      "(2, 2, 1, 8, 10, 100)\n",
      "(2, 1, 10)\n",
      "(2, 2, 1, 8, 10, 100)\n",
      "(2, 1, 10)\n",
      "(2, 2, 1, 8, 10, 100)\n",
      "(2, 1, 10)\n"
     ]
    }
   ],
   "source": [
    "i = 0\n",
    "for (a,tau) in generator():\n",
    "    print(a.shape)\n",
    "    print(tau.shape)\n",
    "    i = i + 1\n",
    "    if i == 5:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00146ccc-023e-4343-8a8e-75eea0faf5e9",
   "metadata": {},
   "source": [
    "Let us create a channel model from this dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0fdaa876-cefd-433a-ad74-2d944192149a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:29.999708Z",
     "iopub.status.busy": "2026-02-13T14:13:29.999581Z",
     "iopub.status.idle": "2026-02-13T14:13:30.003988Z",
     "shell.execute_reply": "2026-02-13T14:13:30.003335Z"
    }
   },
   "outputs": [],
   "source": [
    "from sionna.phy.channel import CIRDataset\n",
    "\n",
    "batch_size = 64 # The batch_size cannot be changed after the creation of the channel model\n",
    "channel_model = CIRDataset(generator,\n",
    "                           batch_size,\n",
    "                           num_rx,\n",
    "                           num_rx_ant,\n",
    "                           num_tx,\n",
    "                           num_tx_ant,\n",
    "                           num_paths,\n",
    "                           num_time_steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "234582db-aad6-4150-812f-9ae86f83d780",
   "metadata": {},
   "source": [
    "We can now sample from this dataset in the same way as we would from a stochastic channel model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "84b6373a-0c2e-43b5-a201-f999e210d549",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:30.006288Z",
     "iopub.status.busy": "2026-02-13T14:13:30.006171Z",
     "iopub.status.idle": "2026-02-13T14:13:30.197603Z",
     "shell.execute_reply": "2026-02-13T14:13:30.196810Z"
    }
   },
   "outputs": [],
   "source": [
    "# Note that the arguments batch_size, num_time_steps, and smapling_frequency\n",
    "# of the __call__ function are ignored as they are already specified by the dataset.\n",
    "a, tau = channel_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1a07eef8-540b-4434-a144-32d6da174454",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:30.199778Z",
     "iopub.status.busy": "2026-02-13T14:13:30.199643Z",
     "iopub.status.idle": "2026-02-13T14:13:30.202513Z",
     "shell.execute_reply": "2026-02-13T14:13:30.201866Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 2, 2, 1, 8, 10, 100])\n",
      "torch.complex64\n",
      "torch.Size([64, 2, 1, 10])\n",
      "torch.float32\n"
     ]
    }
   ],
   "source": [
    "print(a.shape)\n",
    "print(a.dtype)\n",
    "print(tau.shape)\n",
    "print(tau.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9de862ae-001f-4eb9-8f73-9232373a9f96",
   "metadata": {},
   "source": [
    "## Use the Channel Model for OFDM Transmissions\n",
    "\n",
    "The following code demonstrates how you can use the channel model to generate channel frequency responses that can be used for the simulation of communication system based on OFDM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6865d499-7bd9-4d0b-9ba7-677df685d3b7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:30.203975Z",
     "iopub.status.busy": "2026-02-13T14:13:30.203823Z",
     "iopub.status.idle": "2026-02-13T14:13:30.269338Z",
     "shell.execute_reply": "2026-02-13T14:13:30.268453Z"
    }
   },
   "outputs": [],
   "source": [
    "# Create an OFDM resource grid\n",
    "# Each time step is assumed to correspond to one OFDM symbol over which it is constant.\n",
    "resource_grid = sionna.phy.ofdm.ResourceGrid(\n",
    "                                num_ofdm_symbols=num_time_steps,\n",
    "                                fft_size=76,\n",
    "                                subcarrier_spacing=15e3,\n",
    "                                num_tx=num_tx,\n",
    "                                num_streams_per_tx=num_tx_ant)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cd862bcd-26e9-4774-97a2-0d32c2ba8613",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:30.271256Z",
     "iopub.status.busy": "2026-02-13T14:13:30.271119Z",
     "iopub.status.idle": "2026-02-13T14:13:30.281969Z",
     "shell.execute_reply": "2026-02-13T14:13:30.281231Z"
    }
   },
   "outputs": [],
   "source": [
    "ofdm_channel = sionna.phy.channel.GenerateOFDMChannel(channel_model, resource_grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6180d32d-c5e7-45bd-a71d-333f405a548b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:30.283967Z",
     "iopub.status.busy": "2026-02-13T14:13:30.283830Z",
     "iopub.status.idle": "2026-02-13T14:13:30.336667Z",
     "shell.execute_reply": "2026-02-13T14:13:30.335980Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 2, 2, 1, 8, 100, 76])\n"
     ]
    }
   ],
   "source": [
    "# Generate a batch of frequency responses\n",
    "# Shape: [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_ofdm_symbols, num_subcarriers]\n",
    "h_freq = ofdm_channel()\n",
    "print(h_freq.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3bcc8deb-db22-4b48-9ae0-489cb0f083df",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-02-13T14:13:30.338483Z",
     "iopub.status.busy": "2026-02-13T14:13:30.338351Z",
     "iopub.status.idle": "2026-02-13T14:13:30.508960Z",
     "shell.execute_reply": "2026-02-13T14:13:30.508100Z"
    }
   },
   "outputs": [],
   "source": [
    "# Delete dataset\n",
    "%rm my_dataset.h5"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.12"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}