Schmidt decomposition of a vector.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Schmidt decomposition of a vector\n",
    "\n",
    "This problem can be solved using Singular Value Decomposition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cirq\n",
    "\n",
    "def schmidt_decompose(vec, dim1, dim2, tol=1e-9):\n",
    "    \"\"\"Schmidt decomposition of a unit vector.\n",
    "    \n",
    "    Returns a, b, k, s.t. vec = Σ_i (k_i * a_i x b_i).\n",
    "    \"\"\"\n",
    "    assert vec.shape == (dim1 * dim2, )\n",
    "    assert np.allclose(np.linalg.norm(vec), 1), \"Not unit vector.\"\n",
    "    A = np.reshape(vec, (dim1, dim2))\n",
    "    a, k, bT = np.linalg.svd(A)\n",
    "    size = sum(k>=tol)\n",
    "    a = [a[:, i] for i in range(size)]\n",
    "    b = [bT[i, :] for i in range(size)]\n",
    "    k = k[0:size]\n",
    "    return a, b, k  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OK\n"
     ]
    }
   ],
   "source": [
    "def test_schmidt_decompose(vec):\n",
    "    def is_orthonormal_basis(vecs):\n",
    "        V = np.array(vecs)\n",
    "        return np.allclose(V @ V.T.conj(), np.eye(len(vecs)))\n",
    "        \n",
    "    vec = np.array(vec)\n",
    "    dim = len(vec)\n",
    "    for dim1 in range(1, dim+1):\n",
    "        if dim % dim1 != 0: continue\n",
    "        dim2 = dim // dim1\n",
    "        a, b, k = schmidt_decompose(vec, dim1, dim2)\n",
    "        \n",
    "        assert np.allclose(np.linalg.norm(k), 1)\n",
    "        assert is_orthonormal_basis(a)\n",
    "        assert is_orthonormal_basis(b)\n",
    "        \n",
    "        vec_restored = sum([k[i] * np.kron(a[i], b[i]) for i in range(len(k))])\n",
    "        assert np.allclose(vec, vec_restored)\n",
    "\n",
    "test_schmidt_decompose([1.0])\n",
    "test_schmidt_decompose([0.5, 0.5, 0.5, 0.5])\n",
    "\n",
    "for i in range(100):\n",
    "    vec = np.random.random(size=(8,2)) \n",
    "    vec = vec[:,0] + 1j * vec[:,1]\n",
    "    vec /= np.linalg.norm(vec)\n",
    "    test_schmidt_decompose(vec)\n",
    "print(\"OK\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}