A naive quanutm register

Register.jl

#################
# Utils
#################

"""
    log2i(x)

logrithm for integer pow of 2
"""
function log2i(x::T)::T where T
    local n::T = 0
    while x&0x1!=1
        n += 1
        x >>= 1
    end
    return n
end

"""
    batch_normalize!(matrix)

normalize a batch of vector.
"""
function batch_normalize!(s::AbstractMatrix)
    B = size(s, 2)
    for i = 1:B
        normalize!(view(s, :, i))
    end
    s
end

"""
    batch_normalize

normalize a batch of vector.
"""
function batch_normalize(s::AbstractMatrix)
    ts = copy(s)
    batch_normalize!(ts)
end

##############################################################


using Compat
import Base: show


# TODO: move this to document

"""
    AbstractRegister{B, T}

abstract type that registers will subtype from. `B` is the batch size, `T` is the
data type.

## Required Properties

|    Property    |                                                     Description                                                      |     default      |
|:--------------:|:--------------------------------------------------------------------------------------------------------------------:|:----------------:|
| `nqubit(reg)`  | get the total number of qubits.                                                                                      |                  |
| `nactive(reg)` | get the number of active qubits.                                                                                     |                  |
| `nremain(reg)` | get the number of remained qubits.                                                                                   | nqubit - nactive |
| `nbatch(reg)`  | get the number of batch.                                                                                             | `B`              |
| `address(reg)` | get the address of this register.                                                                                    |                  |
| `state(reg)`   | get the state of this register. It always return the matrix stored inside.                                           |                  |
| `eltype(reg)`  | get the element type stored by this register on classical memory. (the type Julia should use to represent amplitude) | `T`              |
| `copy(reg)`    | copy this register.                                                                                                  |                  |
| `similar(reg)` | construct a new register with similar configuration.                                                                 |                  |

## Required Methods

### Multiply

    *(op, reg)

define how operator `op` act on this register. This is quite useful when
there is a special approach to apply an operator on this register. (e.g
a register with no batch, or a register with a MPS state, etc.)

!!! note

    be careful, generally, operators can only be applied to a register, thus
    we should only overload this operation and do not overload `*(reg, op)`.

### Pack Address

    pack_address!(reg, addrs)

pack `addrs` together to the first k-dimensions.

#### Example

Given a register with dimension `[2, 3, 1, 5, 4]`, we pack `[5, 4]`
to the first 2 dimensions. We will get `[5, 4, 2, 3, 1]`.

### Focus Address

    focus!(reg, range)

merge address in `range` together as one dimension (the active space).

#### Example

Given a register with dimension `(2^4)x3` and address [1, 2, 3, 4], we focus
address `[3, 4]`, will pack `[3, 4]` together and merge them as the active
space. Then we will have a register with size `2^2x(2^2x3)`, and address
`[3, 4, 1, 2]`.

## Initializers

Initializers are functions that provide specific quantum states, e.g zero states,
random states, GHZ states and etc.

    register(::Type{RT}, raw, nbatch)

an general initializer for input raw state array.

    register(::Type{InitMethod}, ::Type{RT}, ::Type{T}, n, nbatch)

init register type `RT` with `InitMethod` type (e.g `InitMethod{:zero}`) with
element type `T` and total number qubits `n` with `nbatch`. This will be
auto-binded to some shortcuts like `zero_state`, `rand_state`, `randn_state`.
"""
abstract type AbstractRegister{B, T} end

"""
    nqubit(reg)->Int

get the total number of qubits.
"""
function nqubit end

"""
    nactive(reg)->Int

get the number of active qubits.
"""
function nactive end

"""
nremain(reg)->Int

get the number of remained qubits.
"""
function nremain end

"""
    nbatch(reg)->Int

get the number of batch.
"""
function nbatch end

"""
    address(reg)->Int

get the address of this register.
"""
function address end

"""
    state(reg)

get the state of this register. It always return
the matrix stored inside.
"""
function state end

"""
    register(::Type{RT}, raw, nbatch)

an general initializer for input raw state array.

    register(::Type{InitMethod}, ::Type{RT}, ::Type{T}, n, nbatch)

init register type `RT` with `InitMethod` type (e.g `InitMethod{:zero}`) with
element type `T` and total number qubits `n` with `nbatch`. This will be
auto-binded to some shortcuts like `zero_state`, `rand_state`, `randn_state`.
"""
function register end

"""
    zero_state(n, nbatch)

construct a zero state ``|00\\cdots 00\\rangle``.
"""
function zero_state end

"""
    rand_state(n, nbatch)

construct a normalized random state with uniform distributed
``\\theta`` and ``r`` with amplitude ``r\\cdot e^{i\\theta}``.
"""
function rand_state end

"""

    randn_state(n, nbatch)

construct normalized a random state with normal distributed
``\\theta`` and ``r`` with amplitude ``r\\cdot e^{i\\theta}``.
"""
function randn_state end

##############
## Interfaces
##############

# nqubit
# nactive
nremain(r::AbstractRegister) = nqubit(r) - nactive(r)
nbatch(r::AbstractRegister{B}) where B = Int(B)
eltype(r::AbstractRegister{B, T}) where {B, T} = T

# Factory Methods

# set unsigned conversion rules for nbatch
function register(::Type{RT}, raw::AbstractArray, nbatch::Int) where RT
    register(RT, raw, unsigned(nbatch))
end

# set default register
function register(raw::AbstractArray, nbatch::Int)
    register(Register, raw, nbatch)
end

## Config Initializers

abstract type InitMethod{T} end

# define type conversion
function register(::Type{InitMethod{IM}}, ::Type{RT}, ::Type{T}, n::Int, nbatch::Int) where {IM, RT, T}
    register(InitMethod{IM}, RT, T, n, unsigned(nbatch))
end

# enable multiple dispatch for different initializers
function register(::Type{RT}, ::Type{T}, n::Int, nbatch::Int, method::Symbol=:rand) where {RT, T}
    register(InitMethod{method}, RT, T, n, nbatch)
end

# config default register type
function register(::Type{T}, n::Int, nbatch::Int, method::Symbol=:rand) where T
    register(Register, T, n, nbatch, method)
end

# config default eltype
register(n::Int, nbatch::Int, method::Symbol=:rand) = register(Compat.ComplexF64, n, nbatch, method)

# shortcuts
zero_state(n::Int, nbatch::Int) = register(n, nbatch, :zero)
rand_state(n::Int, nbatch::Int) = register(n, nbatch, :rand)
randn_state(n::Int, nbatch::Int) = register(n, nbatch, :randn)

import Base: *

#############################################################################

###############################################
# Naive Implementation (single process CPU)
###############################################

@inline function _len_active_remain(raw::Matrix, nbatch)
    active_len, nbatch_and_remain = size(raw)
    remain_len = nbatch_and_remain ÷ nbatch
    active_len, remain_len
end

mutable struct Register{B, T} <: AbstractRegister{B, T}
    state::Matrix{T} # this stores a batched state

    nactive::Int # this is the total number of active qubits
    # NOTE: we should replace this with a static mutable vector in the future
    address::Vector{UInt} # this indicates the absolute address of each qubit

    function Register(raw::Matrix{T}, address::Vector{UInt}, nactive::UInt, nbatch::UInt) where T
        active_len, remain_len = _len_active_remain(raw, nbatch)

        ispow2(active_len) && ispow2(remain_len) ||
            throw(Compat.InexactError(:Register, Register, raw))

        new{nbatch, T}(raw, nactive, address)
    end

    # copy method
    function Register(r::Register{B, T}) where {B, T}
        new{B, T}(copy(r.state), r.nactive, copy(r.address))
    end
end

function Register(raw::Matrix, nbatch::UInt)
    active_len, remain_len = _len_active_remain(raw, nbatch)
    N = unsigned(log2i(active_len * remain_len))
    Register(raw, collect(0x1:N), N, nbatch)
end

function Register(raw::Vector, nbatch::UInt)
    Register(reshape(raw, length(raw), 1), nbatch)
end

# Required Properties

nqubit(r::Register) = length(r.address)
nactive(r::Register) = r.nactive
address(r::Register) = r.address
state(r::Register) = r.state
copy(r::Register) = Register(r)

function similar(r::Register{B, T}) where {B, T}
    Register(similar(r.state), copy(r.address), r.nactive, B)
end

# factory methods
register(::Type{Register}, raw, nbatch::UInt) = Register(raw, nbatch)

function register(::Type{InitMethod{:zero}}, ::Type{Register}, ::Type{T}, n::Int, nbatch::UInt) where T
    raw = zeros(T, 1 << n, nbatch)
    raw[1, :] = 1
    Register(raw, nbatch)
end

function register(::Type{InitMethod{:rand}}, ::Type{Register}, ::Type{T}, n::Int, nbatch::UInt) where T
    theta = rand(real(T), 1 << n, nbatch)
    radius = rand(real(T), 1 << n, nbatch)
    raw = @. radius * exp(im * theta)
    Register(batch_normalize!(raw), nbatch)
end

function register(::Type{InitMethod{:randn}}, ::Type{Register}, ::Type{T}, n::Int, nbatch::UInt) where T
    theta = randn(real(T), 1 << n, nbatch)
    radius = randn(real(T), 1 << n, nbatch)
    raw = @. radius * exp(im * theta)
    Register(batch_normalize!(raw), nbatch)
end

function swap_first!(addr::Vector, index)
    temp = addr[1]
    addr[1] = addr[index]
    addr[index] = temp
    addr
end

# NOTE: we use relative address here
# the input are desired orders of current
# address, not the desired address.
# orders here can be any iterable
function pack_address!(tensor::Array{T, N}, address, orders) where {T, N}
    curr_orders = collect(1:(N-1))

    for each in reverse(orders)
        swap_first!(curr_orders, each)
        swap_first!(address, each)
    end

    # we preserve last dim
    permutedims!(tensor, tensor, (curr_orders..., N))
end

function focus!(r::Register, range)
end

nexposed(orders) = 1 << length(orders)
function total_exposed(orders...)
    total = 0
    for each in orders
        total = length(each)
    end
    1 << total
end

function focus!(r::Register{B}, range...) where B
    tensor = reshape(state(r), ntuple(x->2, nqubit(r))..., B)
    for each in reverse(range)
        pack_address!(tensor, r.address, each)
    end

    r.state = reshape(r.state, (nexposed(range...), :))
    r
end

# we convert state to a vector to use
# intrincs like gemv, when nremain is
# 0 and the state is actually a vector
function *(op, r::Register{1})
    if nremain(r) == 0
        return op * vec(r.state)
    end
    op * r.state
end

function *(op, r::Register)
    op * r.state
end

function show(io::IO, r::Register{B, T}) where {B, T}
    println(io, "Default Register (CPU, $T):")
    println(io, "    total: ", nqubit(r))
    println(io, "    batch: ", B)
    print(io, "    active: ", nactive(r))
end