{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "40e2ecf6-a1ee-4d82-924a-e2f763915652",
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra, Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "89746093-dc10-4bb2-9646-c84d5db0d8f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "householder_vector (generic function with 2 methods)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function householder_vector(x::Vector{<:AbstractFloat})::Tuple{Vector, AbstractFloat}\n",
    "    # returns the normalized vector u such that H*x is a multiple of e_1\n",
    "\n",
    "    s = norm(x)\n",
    "    if x[1] ≥ 0\n",
    "        s = -s\n",
    "    end\n",
    "    u = copy(x)\n",
    "    u[1] -= s\n",
    "    u ./= norm(u)\n",
    "    return u, s\n",
    "end\n",
    "\n",
    "function householder_vector(x::Matrix{<:AbstractFloat})::Tuple{Matrix, AbstractFloat}\n",
    "    # returns the normalized vector u such that H*x is a multiple of e_1\n",
    "\n",
    "    s = norm(x)\n",
    "    if x[1] ≥ 0\n",
    "        s = -s\n",
    "    end\n",
    "    u = copy(x)\n",
    "    u[1] -= s\n",
    "    u ./= norm(u)\n",
    "    return u, s\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "262a769a-aa42-4929-bbf4-7f5a97783810",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([0.7960091839647405, 0.3357514552548967, 0.503627182882345], -3.7416573867739413)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "3-element Vector{Float64}:\n",
       " 1.0\n",
       " 2.0\n",
       " 3.0"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = [1., 2, 3]\n",
    "householder_vector(x) |> display\n",
    "x |> display\n",
    "\n",
    "# better with copy and division in place\n",
    "# @benchmark householder_vector(randn(100_000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "457b3bcf-a077-42cb-9a6c-f6d1d6e00504",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5-element Vector{Float64}:\n",
       "  3.983324158315414\n",
       " -5.457417586244235e-16\n",
       " -4.598436127585974e-17\n",
       " -1.4573782140540267e-16\n",
       "  3.2867357162787514e-16"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "A = randn(5, 4)\n",
    "\n",
    "# first step of QR factorization\n",
    "R1 = A\n",
    "(u1, s1) = householder_vector(R1[1:end,1])\n",
    "\n",
    "H1 = I - 2 * u1 * u1'\n",
    "\n",
    "Q1 = H1\n",
    "\n",
    "Q1 * R1[1:end, 1] |> display # what we expect -> a multiple of e_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fd98a89e-01b9-4403-8bca-b9e1443b2eea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×4 Matrix{Float64}:\n",
       "  3.98332      -1.23134       1.8182     0.590776\n",
       "  3.80466e-16   1.52712      -1.85383   -0.709258\n",
       " -3.09958e-16  -1.11852e-17  -0.563749  -1.85235\n",
       "  4.05338e-16  -4.61156e-17  -0.174643   1.27511\n",
       "  1.55202e-16   1.4856e-17    1.73802    0.266999"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# second step\n",
    "R2 = Q1 * R1\n",
    "\n",
    "(u2, s2) = householder_vector(R2[2:end, 2])\n",
    "H2 = I - 2 * u2 * u2'\n",
    "\n",
    "# there is no blkdiag method in julia\n",
    "# (maybe look into https://github.com/JuliaArrays/BlockDiagonals.jl)\n",
    "# there are 2 methods (blocks is an array of blocks):\n",
    "### METHOD 1:\n",
    "# cat(blocks..., dims=(1,2))\n",
    "### METHOD 2:\n",
    "# using SparseArrays\n",
    "# blockdiag(SparseMatrixCSC.(blocks)...)\n",
    "## method 2 is slightly faster with subsequent matrix multiplication\n",
    "# performance is ignored in this step\n",
    "Q2 = cat(1, H2, dims=(1, 2))\n",
    "\n",
    "Q2 * R2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a563f47e-7eda-46c7-9804-080335bcb8a3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×4 Matrix{Float64}:\n",
       "  3.98332      -1.23134       1.8182        0.590776\n",
       "  3.80466e-16   1.52712      -1.85383      -0.709258\n",
       "  2.03593e-16   2.18903e-17   1.83549       0.700423\n",
       "  4.4272e-16   -4.37079e-17   6.42625e-17   1.46093\n",
       " -2.16817e-16  -9.104e-18    -4.23002e-16  -1.58224"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# third step\n",
    "\n",
    "R3 = Q2 * R2\n",
    "\n",
    "(u3, s3) = householder_vector(R3[3:end, 3])\n",
    "H3 = I - 2 * u3 * u3'\n",
    "\n",
    "Q3 = cat(Diagonal(ones(2)), H3, dims=(1,2))\n",
    "Q3 * R3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "85f5eb54-f3fe-40c8-86d0-90e80895b753",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×4 Matrix{Float64}:\n",
       "  3.98332      -1.23134       1.8182        0.590776\n",
       "  3.80466e-16   1.52712      -1.85383      -0.709258\n",
       "  2.03593e-16   2.18903e-17   1.83549       0.700423\n",
       " -4.5963e-16    2.29618e-17  -3.54379e-16  -2.15356\n",
       "  1.78186e-16  -3.82887e-17  -2.39742e-16   1.50034e-16"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# fourth step\n",
    "\n",
    "R4 = Q3 * R3\n",
    "\n",
    "(u4, s4) = householder_vector(R4[4:end, 4])\n",
    "H4 = I - 2 * u4 * u4'\n",
    "\n",
    "Q4 = cat(Diagonal(ones(3)), H4, dims=(1,2))\n",
    "Q4 * R4\n",
    "\n",
    "# done because we arrived at the second dimension of A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eff90e91-7856-4fd7-b2a5-79c0f681147a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "qrfactorization (generic function with 1 method)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function qrfactorization(A::Matrix{<:AbstractFloat})::Tuple{Matrix{<:AbstractFloat}, Matrix{<:AbstractFloat}}\n",
    "    (m, n) = size(A)\n",
    "    R = copy(A)\n",
    "    Q = Diagonal(ones(eltype(A), m))\n",
    "\n",
    "    for k ∈ 1:n\n",
    "        (u, s) = householder_vector(R[k:end, k])\n",
    "        # construct R\n",
    "        R[k, k] = s\n",
    "        R[k+1:end, k] .= 0\n",
    "        R[k:end, k+1:end] -= 2 * u * (u' * R[k:end, k+1:end])\n",
    "        # contruct the new H\n",
    "        H = I - 2 * u * u'\n",
    "        # contruct the Q\n",
    "        Q = Q * cat(Diagonal(ones(eltype(A), k-1)), H, dims=(1,2)) # very inefficient (maybe simply send back the list of u_i)\n",
    "    end\n",
    "    return (Q, R)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4dbe13ff-44f5-4bd2-8452-a9e1477c80ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "A = randn(Float32, 1000, 20)\n",
    "(Q, R) = qrfactorization(A)\n",
    "(norm(A - Q*R) ≤ size(A)[1] * 2^-23 * norm(A)) |> display\n",
    "(norm(I - Q*Q') ≤ size(A)[1] * 2^-23) |> display"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.3",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}