path: root/src/core/quantum_components.c

#include "core/quantum_components.h"

#include <assert.h>
#include <stdbool.h>
#include <stdlib.h>

struct PsiVector psi_new_state_0(void)
{
    return psi_column_vector(psi_new_complex(1.0, 0.0), psi_new_complex(0.0, 0.0));
}

struct PsiVector psi_new_state_1(void)
{
    return psi_column_vector(psi_new_complex(0.0, 0.0), psi_new_complex(1.0, 0.0));
}

struct PsiQuantumGate psi_new_quantum_gate(const char* name, struct PsiMatrix matrix,
                                           size_t num_qubits)
{
    size_t expected_dim = (size_t)1 << num_qubits;
    assert(matrix.rows == expected_dim);
    assert(matrix.cols == expected_dim);

    return (struct PsiQuantumGate){
        name,
        matrix,
        num_qubits,
    };
}

struct PsiQuantumGate psi_new_quantum_gate_from_matrix(const char* name, struct PsiMatrix matrix)
{
    assert(matrix.rows == matrix.cols);

    size_t dim = matrix.rows;
    assert(dim > 0 && (dim & (dim - 1)) == 0);

    size_t num_qubits = 0;
    while (((size_t)1 << num_qubits) < dim)
        num_qubits++;

    return (struct PsiQuantumGate){
        name,
        matrix,
        num_qubits,
    };
}

void psi_free_quantum_gate(struct PsiQuantumGate* gate)
{
    psi_free_matrix(&gate->matrix);
}

struct PsiQuantumBit psi_new_quantum_bit(const char* name, struct PsiVector state)
{
    return (struct PsiQuantumBit){
        name,
        state,
    };
}

void psi_free_quantum_bit(struct PsiQuantumBit* bit)
{
    psi_free_vector(&bit->state);
}

static void update_register(struct PsiQuantumRegister* reg)
{
    psi_free_vector(&reg->state_vector);

    if (reg->num_qubits == 0)
    {
        reg->state_vector = psi_new_vector(0, PSI_COLUMN_VECTOR);
        return;
    }

    struct PsiMatrix result = psi_matrix_from_vector(reg->qubits[0].state);
    for (size_t i = 1; i < reg->num_qubits; i++)
    {
        struct PsiMatrix part = psi_matrix_from_vector(reg->qubits[i].state);
        struct PsiMatrix next = psi_kronecker_matrix(result, part);
        psi_free_matrix(&result);
        psi_free_matrix(&part);
        result = next;
    }

    reg->state_vector = psi_vector_from_matrix(result, PSI_COLUMN_VECTOR);
    psi_free_matrix(&result);
}

struct PsiQuantumRegister psi_new_quantum_register(const char* name, const char** names,
                                                   size_t count)
{
    struct PsiQuantumBit* qubits = malloc(count * sizeof(struct PsiQuantumBit));
    assert(qubits != NULL || count == 0);

    for (size_t i = 0; i < count; i++)
        qubits[i] = psi_new_quantum_bit(names[i], psi_new_state_0());

    struct PsiQuantumRegister reg = {
        name,
        psi_new_vector(0, PSI_COLUMN_VECTOR),
        qubits,
        count,
    };

    update_register(&reg);
    return reg;
}

struct PsiQuantumRegister
psi_new_quantum_register_from(const char* name, const struct PsiQuantumBit* bits, size_t count)
{
    struct PsiQuantumBit* qubits = malloc(count * sizeof(struct PsiQuantumBit));
    assert(qubits != NULL || count == 0);

    for (size_t i = 0; i < count; i++)
        qubits[i] = psi_new_quantum_bit(bits[i].name, psi_clone_vector(bits[i].state));

    struct PsiQuantumRegister reg = {
        name,
        psi_new_vector(0, PSI_COLUMN_VECTOR),
        qubits,
        count,
    };

    update_register(&reg);
    return reg;
}

void psi_free_quantum_register(struct PsiQuantumRegister* reg)
{
    for (size_t i = 0; i < reg->num_qubits; i++)
        psi_free_quantum_bit(&reg->qubits[i]);

    free(reg->qubits);
    reg->qubits = NULL;
    reg->num_qubits = 0;
    psi_free_vector(&reg->state_vector);
}

static bool targets_contain(const size_t* targets, size_t count, size_t value)
{
    for (size_t i = 0; i < count; i++)
        if (targets[i] == value)
            return true;

    return false;
}

static struct PsiMatrix build_contiguous_operator(struct PsiQuantumRegister reg,
                                                  struct PsiQuantumGate gate, size_t start_idx)
{
    size_t n = reg.num_qubits;
    size_t g = gate.num_qubits;

    bool has_result = false;
    struct PsiMatrix result = { 0 };

    for (size_t i = 0; i < n; i++)
    {
        if (i > start_idx && i < start_idx + g)
            continue;

        struct PsiMatrix part =
                i == start_idx ? psi_clone_matrix(gate.matrix) : psi_identity_matrix(2);

        if (!has_result)
        {
            result = part;
            has_result = true;
            continue;
        }

        struct PsiMatrix next = psi_kronecker_matrix(result, part);
        psi_free_matrix(&result);
        psi_free_matrix(&part);
        result = next;
    }

    if (!has_result)
        return psi_identity_matrix((size_t)1 << n);

    return result;
}

static struct PsiMatrix build_full_operator(struct PsiQuantumRegister reg,
                                            struct PsiQuantumGate gate, const size_t* targets,
                                            size_t target_count)
{
    size_t n = reg.num_qubits;
    size_t g = gate.num_qubits;
    size_t dim = (size_t)1 << n;

    bool contiguous = true;
    for (size_t i = 1; i < target_count; i++)
        if (targets[i] != targets[i - 1] + 1)
        {
            contiguous = false;
            break;
        }

    if (contiguous && g == n)
        return psi_clone_matrix(gate.matrix);

    if (contiguous)
        return build_contiguous_operator(reg, gate, targets[0]);

    struct PsiMatrix result = psi_new_matrix(dim, dim);

    for (size_t col = 0; col < dim; col++)
        for (size_t row = 0; row < dim; row++)
        {
            size_t target_row_bits = 0;
            size_t target_col_bits = 0;

            for (size_t i = 0; i < target_count; i++)
            {
                size_t qubit_pos = n - 1 - targets[i];
                if ((row >> qubit_pos) & 1)
                    target_row_bits |= (size_t)1 << (g - 1 - i);
                if ((col >> qubit_pos) & 1)
                    target_col_bits |= (size_t)1 << (g - 1 - i);
            }

            bool non_target_match = true;
            for (size_t q = 0; q < n; q++)
            {
                if (targets_contain(targets, target_count, q))
                    continue;

                size_t qubit_pos = n - 1 - q;
                if (((row >> qubit_pos) & 1) != ((col >> qubit_pos) & 1))
                {
                    non_target_match = false;
                    break;
                }
            }

            if (non_target_match)
                result.data[row * result.cols + col] =
                        psi_get_matrix(gate.matrix, target_row_bits, target_col_bits);
        }

    return result;
}

void psi_apply_gate(struct PsiQuantumRegister* reg, struct PsiQuantumGate gate,
                    const size_t* targets, size_t target_count)
{
    size_t n = reg->num_qubits;

    assert(gate.num_qubits == target_count);
    for (size_t i = 0; i < target_count; i++)
        assert(targets[i] < n);

    for (size_t i = 0; i < target_count; i++)
        for (size_t j = i + 1; j < target_count; j++)
            assert(targets[i] != targets[j]);

    struct PsiMatrix full_operator = build_full_operator(*reg, gate, targets, target_count);
    struct PsiVector new_state = psi_mul_vector_matrix(reg->state_vector, full_operator);

    psi_free_vector(&reg->state_vector);
    psi_free_matrix(&full_operator);
    reg->state_vector = new_state;
}