1 - Julia basics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we'll work through some basic features of Julia while building up some code to generate a single transverse field Ising model Hamiltonian. Our goal is to cover:\n",
    "- Basic functions\n",
    "- Methods vs functions\n",
    "- Multiple dispatch\n",
    "- Vectorization/broadcasting\n",
    "- Structs\n",
    "- Using packages\n",
    "- Unit tests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In Julia we can call functions that are defined in \"base Julia\" itself:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "abs(-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = 10\n",
    "b = 5.5\n",
    "a * b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = rand(5, 5)\n",
    "B = rand(5, 10)\n",
    "A * B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@which A * B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C = rand(5)\n",
    "D = rand(1, 5)\n",
    "C * D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@which C * D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In fact, we can ask Julia to show us *all* the type combinations for which a method of `*` is defined:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "methods(*)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define our own functions, but sometimes it's nice to be concise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function my_mult(A, B)\n",
    "   return A*B\n",
    "end\n",
    "@assert my_mult(A, B) == A*B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_mult2(A, B) = A*B\n",
    "@assert my_mult2(A, B) == A*B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One thing that makes Julia a good choice for scientific computing is that it is often very fast. But it's easy to write slow Julia code (it's easy to write bad code in any language). Let's examine the difference between a fast Julia function and a slow one, and see how we could spot the issue and fix it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function my_adder(A, B)\n",
    "    C = 1\n",
    "    C += A + B\n",
    "    return C\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Julia allows us to examine the final instructions emitted by LLVM after the compilation process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_llvm my_mult(1, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_llvm my_mult(1.1, 3.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Julia's compiler can generate highly efficient code that is correctly specialized on the input type, but **only** if the code is **type stable**. Let's see what that means:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_warntype my_adder(1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_llvm my_adder(1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_warntype my_adder(1.2, 2.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_llvm my_adder(1.2, 2.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the case with floating point arguments, `my_adder` has extra instructions related to type promotion. How can we fix this and ensure type stability?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function my_adder2(A::T, B::T) where {T}\n",
    "    C = one(T)\n",
    "    C += A + B\n",
    "    return C\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_warntype my_adder2(1.2, 2.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_warntype my_adder2(1, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This seems to have helped. But what if `A` and `B` were of different types (e.g. were not both of type `T`)? We could write additional methods to *promote* them to a joint type (e.g. promote an `Int` to a `Float64` and then add them). `@code_warntype` and more sophisticated tools like [Cthulhu](https://github.com/JuliaDebug/Cthulhu.jl) can be very helpful ways to diagnose type instability and performance problems in your code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Julia has a lot of sophisticated functionality for working with arrays, which makes it a great language for our purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M = zeros(3, 3)\n",
    "@show typeof(M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = ones(Int8, 2,2,2)\n",
    "@show typeof(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@show size(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@show typeof(size(N))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "O = cat(N, ones(Int8, 2, 2), dims=3)\n",
    "@show size(O)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M2 = vcat(M, M)\n",
    "@show size(M2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M2 = hcat(M, M)\n",
    "@show size(M2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also apply the same functions over multiple elements in a collection:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = [rand(2, 3), rand(3, 4), rand(5, 10)] .* [rand(3, 6), rand(4, 2), rand(10, 1)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll create our own \"abstract\" type to represent the generic idea of a Hamiltonian. We can define generic methods on this type which its child types will inherit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "abstract type Hamiltonian{N} end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's extend some methods from the LinearAlgebra standard library\n",
    "using LinearAlgebra\n",
    "LinearAlgebra.ishermitian(H::Hamiltonian) = true\n",
    "Base.ndims(H::Hamiltonian{N}) where {N} = N"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's build a **mutable** struct to represent the Hamiltonians of a specific model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct TransverseFieldIsing{N} <: Hamiltonian{N}\n",
    "    dims::NTuple{N, Int}\n",
    "    h::Array{Float64, N}\n",
    "    J::Array{Float64, N}\n",
    "    row_ixs::Vector{Int}\n",
    "    col_ixs::Vector{Int}\n",
    "    nz_vals::Vector{ComplexF64}\n",
    "    # inner c-tor\n",
    "    function TransverseFieldIsing{N}(dims, h, J, row_ixs, col_ixs, nz_vals) where {N}\n",
    "        new(dims, h, J, row_ixs, col_ixs, nz_vals)\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# outer c-tor\n",
    "function TransverseFieldIsing(dims, h, J)\n",
    "    N = length(dims) # number of dimensions in lattice\n",
    "    ndims(h) != N && throw(ArgumentError(\"h must have the same number of dimensions as dims!\"))\n",
    "    ndims(J) != N && throw(ArgumentError(\"J must have the same number of dimensions as dims!\"))\n",
    "    # create indices representing (hypercubic) lattice\n",
    "    cis = CartesianIndices(dims)\n",
    "    for (ii, ci) in enumerate(cis)\n",
    "        # fill this in later...?\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# outer c-tor -- let's specialize this a bit!\n",
    "function TransverseFieldIsing(dims::Tuple{Int}, h::Vector{Float64}, J::Vector{Float64})\n",
    "    N_sites = dims[1]\n",
    "    length(h) == 1 || length(h) == N_sites || throw(ArgumentError(\"h must be either length 1 or length $N_sites\"))\n",
    "    length(J) == 1 || length(J) == N_sites || throw(ArgumentError(\"J must be either length 1 or length $N_sites\"))\n",
    "    if length(h) == 1\n",
    "        h = fill(h, N_sites)\n",
    "    end\n",
    "    if length(J) == 1\n",
    "        J = fill(J, N_sites)\n",
    "    end\n",
    "    row_ixs = Int[]\n",
    "    col_ixs = Int[]\n",
    "    nz_vals = ComplexF64[]\n",
    "    for input_ind in 0:2^N_sites-1\n",
    "        input_bits = Vector{Bool}(digits(input_ind, base=2, pad=N_sites))\n",
    "        # h-term\n",
    "        @inbounds for ii in 1:N_sites\n",
    "            output_bits = copy(input_bits)\n",
    "            output_bits[ii] ⊻= true\n",
    "            output_ind = sum([output_bits[jj]<<(jj-1) for jj in 1:N_sites])\n",
    "            push!(row_ixs, input_ind+1)\n",
    "            push!(col_ixs, output_ind+1)\n",
    "            push!(nz_vals, h[ii])\n",
    "        end\n",
    "        # J-term\n",
    "        push!(row_ixs, input_ind+1)\n",
    "        push!(col_ixs, input_ind+1)\n",
    "        nz_val = 0.0\n",
    "        @inbounds for site in 1:N_sites-1\n",
    "            next_site   = site + 1\n",
    "            flip = (input_bits[site] ^ input_bits[next_site]) == 1\n",
    "            bond_J = flip ? -J[site] : J[site]\n",
    "            nz_val += bond_J\n",
    "        end\n",
    "        push!(nz_vals, nz_val)\n",
    "    end\n",
    "    return TransverseFieldIsing{1}(dims, h, J, row_ixs, col_ixs, nz_vals)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf = TransverseFieldIsing((2,), [1.0, 1.0], [2.0, 2.0]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's extend some methods that Julia itself provides to make better use of our new type. At first we will be doing exact diagonalization, so it makes sense to extend some methods from Julia for sparse matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using SparseArrays\n",
    "Base.convert(SparseMatrixCSC, tf::TransverseFieldIsing) = sparse(tf.row_ixs, tf.col_ixs, tf.nz_vals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convert(SparseMatrixCSC, tf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ishermitian(tf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately, Julia doesn't provide a sparse eigensolver by default. For now we can get around this by implementing an ugly hack: converting to a dense matrix. Of course, this won't work for large systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function LinearAlgebra.eigen(tf::TransverseFieldIsing)\n",
    "    sp_mat = convert(SparseMatrixCSC, tf)\n",
    "    return eigen(Hermitian(Matrix(sp_mat)), 1:1)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eigen(tf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Converting to a dense matrix like this isn't very efficient -- can we use a package to do better?\n",
    "\n",
    "Yes. [Arpack.jl](https://github.com/JuliaLinearAlgebra/Arpack.jl) wraps the ARPACK sparse eigensolving/SVD library and we can use it to perform the diagonalization. To use it, we'll first engage Julia's package manager (helpfully called `Pkg`) to install the package. Note that becaus ARPACK itself is written in FORTRAN, you might get install errors if your system doesn't have a `gfortran` in its `PATH`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Pkg\n",
    "Pkg.add(\"Arpack\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Arpack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function LinearAlgebra.eigen(tf::TransverseFieldIsing)\n",
    "    sp_mat = convert(SparseMatrixCSC, tf)\n",
    "    λ, ϕ = eigs(Hermitian(sp_mat), nev=1, which=:SR)\n",
    "    return λ, ϕ\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eigen(tf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We should also write some basic unit tests for our code to try to catch bugs. Julia provides testing functionality to ensure that a function throws an expected error, or that the result is (approximately) equal to some expected value. We can also `@test_broken` to allow a test to pass which shouldn't and mark it as something we should return to in the future."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Test\n",
    "@test_throws ArgumentError TransverseFieldIsing((2,), [1.0, 1.0, 1.0], [2.0, 2.0])\n",
    "@test_throws ArgumentError TransverseFieldIsing((2,), [1.0, 1.0], [2.0, 2.0, 2.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf = TransverseFieldIsing((2,), [0.0, 0.0], [2.0, 2.0])\n",
    "@test eigen(tf)[1] ≈ [-2.0] atol=1e012"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.7.0-DEV",
   "language": "julia",
   "name": "julia-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}