3_AoA_Estimation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8cd4169b-1b0b-418c-b37e-1b37dbbdaea4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-02-06 09:23:03.140005: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2024-02-06 09:23:03.167713: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-02-06 09:23:03.610513: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "2024-02-06 09:23:04.182257: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-02-06 09:23:04.197271: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-02-06 09:23:04.197384: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-02-06 09:23:04.198838: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-02-06 09:23:04.198963: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-02-06 09:23:04.199035: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-02-06 09:23:04.444590: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-02-06 09:23:04.444700: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-02-06 09:23:04.444770: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2024-02-06 09:23:04.444831: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1639] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 10 MB memory:  -> device: 0, name: NVIDIA GeForce GTX 1080 Ti, pci bus id: 0000:01:00.0, compute capability: 6.1\n",
      "2024-02-06 09:23:04.455783: I tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:753] failed to allocate 10.06MiB (10551296 bytes) from device: CUDA_ERROR_OUT_OF_MEMORY: out of memory\n",
      "2024-02-06 09:23:04.456159: I tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:753] failed to allocate 9.06MiB (9496320 bytes) from device: CUDA_ERROR_OUT_OF_MEMORY: out of memory\n",
      "2024-02-06 09:23:04.456519: I tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:753] failed to allocate 8.15MiB (8546816 bytes) from device: CUDA_ERROR_OUT_OF_MEMORY: out of memory\n",
      "2024-02-06 09:23:04.843601: I tensorflow/core/grappler/optimizers/data/replicate_on_split.cc:32] Running replicate on split optimization\n",
      "2024-02-06 09:23:04.865021: I tensorflow/core/grappler/optimizers/data/replicate_on_split.cc:32] Running replicate on split optimization\n"
     ]
    }
   ],
   "source": [
    "from dichasus_cf0x import training_set\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "84658b93-77ba-435b-ab5d-5219a3858591",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Count number of datapoints in training set (for progress bar)\n",
    "TOTAL_DATAPOINTS = sum(1 for _ in training_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "71d829c0-b795-4c74-bbe3-4d223880cd9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ANTENNAS_PER_ROW = tf.shape(training_set.take(1).get_single_element()[0])[2].numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "13e0cc6f-c58f-47d8-9c88-5e00a5284899",
   "metadata": {},
   "outputs": [],
   "source": [
    "estimated_toas = np.load(\"results/estimated_toas.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "062eb5d6-12f4-40b3-a362-1bde5686f778",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_unitary_rootmusic_estimator(chunksize = 32, shed_coeff_ratio = 0):\n",
    "    I = np.eye(chunksize // 2)\n",
    "    J = np.flip(np.eye(chunksize // 2), axis = -1)\n",
    "    Q = np.asmatrix(np.block([[I, 1.0j * I], [J, -1.0j * J]]) / np.sqrt(2))\n",
    "    \n",
    "    def unitary_rootmusic(R):\n",
    "        assert(len(R) == chunksize)\n",
    "        C = np.real(Q.H @ R @ Q)\n",
    "    \n",
    "        eig_val, eig_vec = np.linalg.eigh(C)\n",
    "        eig_val = eig_val[::-1]\n",
    "        eig_vec = eig_vec[:,::-1]\n",
    "\n",
    "        source_count = 1\n",
    "        En = eig_vec[:,source_count:]\n",
    "        ENSQ = Q @ En @ En.T @ Q.H\n",
    "    \n",
    "        coeffs = np.asarray([np.trace(ENSQ, offset = diag) for diag in range(1, len(R))])\n",
    "        coeffs = coeffs[:int(len(coeffs) * (1 - shed_coeff_ratio))]\n",
    "\n",
    "        # Remove some of the smaller noise coefficients, trade accuracy for speed\n",
    "        coeffs = np.hstack((coeffs[::-1], np.trace(ENSQ), coeffs.conj()))\n",
    "        roots = np.roots(coeffs)\n",
    "        roots = roots[abs(roots) < 1.0]\n",
    "        largest_root = np.argmax(1 / (1.0 - np.abs(roots)))\n",
    "        \n",
    "        return np.angle(roots[largest_root])\n",
    "\n",
    "    return unitary_rootmusic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6a962da6-601f-4582-8970-2b4ad124ffed",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20973/20973 [01:50<00:00, 189.39it/s]\n"
     ]
    }
   ],
   "source": [
    "umusic = get_unitary_rootmusic_estimator(ANTENNAS_PER_ROW)\n",
    "\n",
    "estimated_aoas = []\n",
    "for index, data in enumerate(tqdm.tqdm(training_set, total = TOTAL_DATAPOINTS)):\n",
    "    csi, pos, time = data[0], data[1], data[2]\n",
    "\n",
    "    # Extract channel coefficients at the arrival time of the LoS component\n",
    "    # We use only the lower 512 subcarriers, since the observed phase offset for them is lower.\n",
    "    # This is a well-known side effect of the FDD-based DICHASUS synchronization method used at the time of recording the dc0x dataset.\n",
    "    transform_vectors = np.exp(1.0j * 2 * np.pi * np.outer(estimated_toas[index], np.arange(-np.shape(csi)[-1] // 2, np.shape(csi)[-1] // 2)) / np.shape(csi)[-1])\n",
    "    csi_at_los = np.einsum(\"arms,as->arm\", csi[:,:,:,:512], transform_vectors[:,:512])\n",
    "    R = np.einsum(\"arm,arn->amn\", csi_at_los, np.conj(csi_at_los))\n",
    "\n",
    "    estimated_aoas.append([np.arcsin(umusic(R[array]) / np.pi) for array in range(len(R))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7857b029-08db-47a5-89c0-9ce391c4f064",
   "metadata": {},
   "outputs": [],
   "source": [
    "estimated_aoas = np.asarray(estimated_aoas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "48d31b9a-0739-4ce2-bb93-410b2f969a5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"results/estimated_aoas.npy\", estimated_aoas)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}