gradient_descent_on_complex_data_least_squares.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4ff5888c",
   "metadata": {},
   "source": [
    "Sascha Spors,\n",
    "Professorship Signal Theory and Digital Signal Processing,\n",
    "Institute of Communications Engineering (INT),\n",
    "Faculty of Computer Science and Electrical Engineering (IEF),\n",
    "University of Rostock,\n",
    "Germany\n",
    "\n",
    "# Data Driven Audio Signal Processing - A Tutorial with Computational Examples\n",
    "\n",
    "Winter Semester 2023/24 (Master Course #24512)\n",
    "\n",
    "- lecture: https://github.com/spatialaudio/data-driven-audio-signal-processing-lecture\n",
    "- tutorial: https://github.com/spatialaudio/data-driven-audio-signal-processing-exercise\n",
    "\n",
    "Feel free to contact lecturer frank.schultz@uni-rostock.de"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a4b3cf8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Linear problem y = X w\n",
    "# - with complex-valued data\n",
    "# - with full column rank F, tall/thin, pure column space matrix X\n",
    "# - given feature matrix X and ground truth outcome y\n",
    "# - unknown weights w\n",
    "# 1. solve for w with left inverse of X (complex valued closed form solution)\n",
    "# 2. iteratively solve with complex, linear layer (without bias) and\n",
    "# ADAM stochastic gradient descent\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "from complextorch.nn.modules.linear import CVLinear\n",
    "from complextorch.nn.modules.loss import CVQuadError\n",
    "\n",
    "torch.manual_seed(1)\n",
    "rng = np.random.default_rng(1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96c70af9",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 2**10  # number of rows for tall/thin X  = number of  data samples\n",
    "F = 3  # number of columns for tall/thin X = number of features\n",
    "w_ground_truth = (np.arange(F)+1) - 1j*((np.arange(F)-F))  # nice numbers\n",
    "X_train = rng.normal(size=(N, F)) + 1j * rng.normal(size=(N, F))\n",
    "print('\\nmatrix rank == F ? ', np.allclose(np.linalg.matrix_rank(X_train), F))\n",
    "U, _, _ = np.linalg.svd(X_train)\n",
    "X_train = U[:, 0:F]  # X is now pure column space\n",
    "y_pure_column_space = X_train @ w_ground_truth  # linear combination of pure column space\n",
    "y_train = y_pure_column_space + np.sqrt(N)*U[:, F+1]  # add 'noise' from left null space,\n",
    "# such that we precisely know the residual\n",
    "\n",
    "residual = y_train - y_pure_column_space\n",
    "theoretical_empirical_risk = np.inner(np.conj(residual), residual) / N\n",
    "print('\\ntheoretical_empirical_risk', theoretical_empirical_risk)\n",
    "# theoretical empirical risk -> any optimisation can never get it lower than that,\n",
    "# because linear algebra fundamentals cannot be beaten -> keep this in mind if\n",
    "# desperately trying to reduce this loss further\n",
    "# if the values above are unchanged, theoretical_empirical_risk = 1\n",
    "\n",
    "# note: CVQuadError loss used below is normalised by 1/2\n",
    "# and not averaged by batch_size\n",
    "# so only the empirical risk of the finally trained model should be\n",
    "# directly compared with the above theoretical_empirical_risk\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57f40877",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prep for torch / complextorch\n",
    "X_train = torch.from_numpy(X_train.astype('complex64'))\n",
    "y_train = torch.from_numpy(y_train.astype('complex64'))\n",
    "print()\n",
    "print(X_train.shape, X_train.dtype)\n",
    "print(y_train.shape, y_train.dtype)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7462572a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# closed form solution w = (X^H X)^-1 X^H y with torch tensor handling\n",
    "w_left_inverse = torch.matmul(torch.matmul(torch.inverse(torch.matmul(X_train.H, X_train)), X_train.H), y_train)\n",
    "print('\\nweights true vs. weights from left inverse',\n",
    "      '\\n', w_ground_truth,\n",
    "      '\\n', w_left_inverse.detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a0ccf52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ML parameters\n",
    "B = N // 64\n",
    "batch_size = N // B\n",
    "learning_rate = 1e-3\n",
    "num_epochs = 1000 + 1\n",
    "log_epoch = 100\n",
    "\n",
    "\n",
    "class Model(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # simple layer: 1 perceptron with F inputs\n",
    "        self.layer1 = CVLinear(F, 1, bias=False)\n",
    "\n",
    "    def predict_train(self, x):\n",
    "        # backprop/autograd is by default enabled\n",
    "        # so x.real and x.imag have grad_fn pointers\n",
    "        # for training\n",
    "        x = self.layer1(x)\n",
    "        return x\n",
    "\n",
    "    def predict_test(self, x):\n",
    "        # we don't need all the backprop stuff in test prediction\n",
    "        # so x has no grad_fn object assigned\n",
    "        with torch.no_grad():\n",
    "            return self.predict_train(x)\n",
    "\n",
    "\n",
    "# data handling, we do no split into train/test\n",
    "train_ds = TensorDataset(X_train, y_train)\n",
    "train_dl = DataLoader(train_ds, batch_size, shuffle=True)\n",
    "# prep model\n",
    "model = Model()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "loss_fn = CVQuadError()\n",
    "\n",
    "print(model)\n",
    "print('batch_size', batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad98512c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('\\nlearn / train ...')\n",
    "for epoch in range(num_epochs):\n",
    "    for X_batch, y_batch in train_dl:\n",
    "        y_pred = model.predict_train(X_batch)\n",
    "        loss = loss_fn(y_pred, y_batch[:, None])\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "    if epoch % log_epoch == 0:\n",
    "        print(f'epoch {epoch} last batch loss {loss.item():.4e}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92d234eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('\\npredict...')\n",
    "# model is trained, check it\n",
    "residual = (model.predict_test(X_train)[:, 0] - y_train).detach().numpy()\n",
    "empirical_risk = np.inner(np.conj(residual), residual) / N\n",
    "print('\\nempirical_risk', empirical_risk)\n",
    "print('\\ntheoretical_empirical_risk', theoretical_empirical_risk)\n",
    "# check the learned weights\n",
    "layer = model.layer1.state_dict()\n",
    "print('\\nweights true vs. from trained model')\n",
    "print('real part')\n",
    "print(w_ground_truth.real)\n",
    "print(layer['linear_r.weight'].detach().numpy())\n",
    "print('imag part')\n",
    "print(w_ground_truth.imag)\n",
    "print(layer['linear_i.weight'].detach().numpy())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "651b1eff",
   "metadata": {},
   "source": [
    "## Copyright\n",
    "\n",
    "- the notebooks are provided as [Open Educational Resources](https://en.wikipedia.org/wiki/Open_educational_resources)\n",
    "- feel free to use the notebooks for your own purposes\n",
    "- the text is licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/)\n",
    "- the code of the IPython examples is licensed under the [MIT license](https://opensource.org/licenses/MIT)\n",
    "- please attribute the work as follows: *Frank Schultz, Data Driven Audio Signal Processing - A Tutorial Featuring Computational Examples, University of Rostock* ideally with relevant file(s), github URL https://github.com/spatialaudio/data-driven-audio-signal-processing-exercise, commit number and/or version tag, year."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myddasp",
   "language": "python",
   "name": "myddasp"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}