<filename>tlquantum/tt_precontraction.py
import tensorly as tl
from tensorly import tensordot, transpose
from itertools import chain


# Author: <NAME> <<EMAIL>>

# License: BSD 3 clause


def layers_contract(layer_list, ncontrl):
    """Contracts sublists of a list of layers vertically (merging cores of multiple layers for a single qubit) up to
    some maximum contraction depth.

    Parameters
    ----------
    layer_list : List of tt-tensor layers
    ncontrl : Maximum number of layers to contract into each composite layer.

    Returns
    -------
    List of layers resulting from vertically contracting the layers of layer_list
    """
    nlayers = len(layer_list)
    contrsets, merge_layers = list(range(nlayers)), []
    contrsets = [contrsets[i:i+ncontrl] for i in range(0, nlayers, ncontrl)]
    for contrset in contrsets:
        merge_layers.append(_vertical_contract([layer_list[i] for i in contrset]))
    merge_layers = list(chain.from_iterable(merge_layers)) #merge list of core lists (layers) into single list
    return merge_layers


def _vertical_contract(layer_list):
    """Contracts list of layers vertically (merging cores of multiple layers for a single qubit).

    Parameters
    ----------
    layer_list : List of tt-tensor layers

    Returns
    -------
    Single layer (list of cores) resulting from vertically contracting the layers of layer_list
    """
    nlayers = len(layer_list)
    nqubits, merge_layers = len(layer_list[0]), []
    for q in range(nqubits):
        cont = layer_list[0][q] #first layer of each qubit to start contraction
        for l in range(1, nlayers):
            dim1, dim2 = cont.shape, layer_list[l][q].shape
            cont = tensordot(layer_list[l][q], cont, axes=([2],[1])) #abdefg --> aedbfg --> aebdfg --> 
            #cont = transpose(transpose(transpose(cont, 1, 3), 2, 3), 3, 4) #pytorch
            cont = transpose(transpose(transpose(cont, (0,3,2,1,4,5)), (0,1,3,2,4,5)), (0,1,2,4,3,5))  #numpy and interface
            cont = tl.reshape(cont, (dim1[0]*dim2[0], dim1[1], dim1[2], dim1[3]*dim2[3]))
        merge_layers.append(cont)
    return merge_layers


def qubits_contract(layer, ncontraq, contrsets=None):
    """Contracts lists (layers) of tt-tensor cores horizontally (merging multiple cores in a single layer) up to
    some maximum number of qubits.

    Parameters
    ----------
    layer : List of tt-tensor cores
    ncontrq : Maximum number of qubits to contract into each composite core.

    Returns
    -------
    Layer (list of cores) resulting from horizontally contracting the cores of layer
    """
    merge_layer, nqubits = [], len(layer)
    if contrsets == None:
        contrsets = _get_contrsets(nqubits, ncontraq)
    if len(layer[0].shape) == 3: #if cores have 3 indices, state tensor
        for contrset in contrsets:
            merge_layer.append(_horizontal_state_contract([layer[i] for i in contrset]))
    else: #if cores do not have 3 indices (here, they would have 4 indices), operator tensor
        for contrset in contrsets:
            merge_layer.append(_horizontal_contract([layer[i] for i in contrset]))
    return merge_layer


def _horizontal_contract(qubits):
    """Contracts list of qubits horizontally (merging multiple cores in a single layer).

    Parameters
    ----------
    qubits : List of tt-tensor cores
    state : tt-tensor, input state to be evolved by unitary

    Returns
    -------
    Single core resulting from horizontally contracting the cores of qubits
    """
    cont = qubits[0]
    for i in range(1, len(qubits)):
        cont = tensordot(cont, qubits[i], axes=([3],[0]))
        #cont = transpose(cont, 2, 3) # pytorch
        cont = transpose(cont, (0,1,3,2,4,5)) # numpy and cupy interface
        cont = tl.reshape(cont, (cont.shape[0], cont.shape[1]*cont.shape[2], cont.shape[3]*cont.shape[4], cont.shape[-1]))
    return cont


def _horizontal_state_contract(qubits):
    """Contracts list of qubits horizontally (merging multiple cores in a single layer).

    Parameters
    ----------
    qubits : List of tt-tensor cores
    state : tt-tensor, input state to be evolved by unitary
    
    Returns
    -------
    Single core resulting from horizontally contracting the cores of qubits
    """
    cont = qubits[0]
    for i in range(1, len(qubits)):
        cont = tensordot(cont, qubits[i], axes=([2],[0]))
        cont = tl.reshape(cont, (cont.shape[0], cont.shape[1]*cont.shape[2], cont.shape[3]))
    return cont


def _get_contrsets(nqubits, ncontraq):
    contrsets = list(range(nqubits))
    return [contrsets[i:i+ncontraq] for i in range(0, nqubits, ncontraq)]

<filename>tlquantum/tt_precontraction.py
import tensorly as tl
from tensorly import tensordot, transpose
from itertools import chain


# Author: <NAME> <<EMAIL>>

# License: BSD 3 clause


def layers_contract(layer_list, ncontrl):
    """Contracts sublists of a list of layers vertically (merging cores of multiple layers for a single qubit) up to
    some maximum contraction depth.

    Parameters
    ----------
    layer_list : List of tt-tensor layers
    ncontrl : Maximum number of layers to contract into each composite layer.

    Returns
    -------
    List of layers resulting from vertically contracting the layers of layer_list
    """
    nlayers = len(layer_list)
    contrsets, merge_layers = list(range(nlayers)), []
    contrsets = [contrsets[i:i+ncontrl] for i in range(0, nlayers, ncontrl)]
    for contrset in contrsets:
        merge_layers.append(_vertical_contract([layer_list[i] for i in contrset]))
    merge_layers = list(chain.from_iterable(merge_layers)) #merge list of core lists (layers) into single list
    return merge_layers


def _vertical_contract(layer_list):
    """Contracts list of layers vertically (merging cores of multiple layers for a single qubit).

    Parameters
    ----------
    layer_list : List of tt-tensor layers

    Returns
    -------
    Single layer (list of cores) resulting from vertically contracting the layers of layer_list
    """
    nlayers = len(layer_list)
    nqubits, merge_layers = len(layer_list[0]), []
    for q in range(nqubits):
        cont = layer_list[0][q] #first layer of each qubit to start contraction
        for l in range(1, nlayers):
            dim1, dim2 = cont.shape, layer_list[l][q].shape
            cont = tensordot(layer_list[l][q], cont, axes=([2],[1])) #abdefg --> aedbfg --> aebdfg --> 
            #cont = transpose(transpose(transpose(cont, 1, 3), 2, 3), 3, 4) #pytorch
            cont = transpose(transpose(transpose(cont, (0,3,2,1,4,5)), (0,1,3,2,4,5)), (0,1,2,4,3,5))  #numpy and interface
            cont = tl.reshape(cont, (dim1[0]*dim2[0], dim1[1], dim1[2], dim1[3]*dim2[3]))
        merge_layers.append(cont)
    return merge_layers


def qubits_contract(layer, ncontraq, contrsets=None):
    """Contracts lists (layers) of tt-tensor cores horizontally (merging multiple cores in a single layer) up to
    some maximum number of qubits.

    Parameters
    ----------
    layer : List of tt-tensor cores
    ncontrq : Maximum number of qubits to contract into each composite core.

    Returns
    -------
    Layer (list of cores) resulting from horizontally contracting the cores of layer
    """
    merge_layer, nqubits = [], len(layer)
    if contrsets == None:
        contrsets = _get_contrsets(nqubits, ncontraq)
    if len(layer[0].shape) == 3: #if cores have 3 indices, state tensor
        for contrset in contrsets:
            merge_layer.append(_horizontal_state_contract([layer[i] for i in contrset]))
    else: #if cores do not have 3 indices (here, they would have 4 indices), operator tensor
        for contrset in contrsets:
            merge_layer.append(_horizontal_contract([layer[i] for i in contrset]))
    return merge_layer


def _horizontal_contract(qubits):
    """Contracts list of qubits horizontally (merging multiple cores in a single layer).

    Parameters
    ----------
    qubits : List of tt-tensor cores
    state : tt-tensor, input state to be evolved by unitary

    Returns
    -------
    Single core resulting from horizontally contracting the cores of qubits
    """
    cont = qubits[0]
    for i in range(1, len(qubits)):
        cont = tensordot(cont, qubits[i], axes=([3],[0]))
        #cont = transpose(cont, 2, 3) # pytorch
        cont = transpose(cont, (0,1,3,2,4,5)) # numpy and cupy interface
        cont = tl.reshape(cont, (cont.shape[0], cont.shape[1]*cont.shape[2], cont.shape[3]*cont.shape[4], cont.shape[-1]))
    return cont


def _horizontal_state_contract(qubits):
    """Contracts list of qubits horizontally (merging multiple cores in a single layer).

    Parameters
    ----------
    qubits : List of tt-tensor cores
    state : tt-tensor, input state to be evolved by unitary
    
    Returns
    -------
    Single core resulting from horizontally contracting the cores of qubits
    """
    cont = qubits[0]
    for i in range(1, len(qubits)):
        cont = tensordot(cont, qubits[i], axes=([2],[0]))
        cont = tl.reshape(cont, (cont.shape[0], cont.shape[1]*cont.shape[2], cont.shape[3]))
    return cont


def _get_contrsets(nqubits, ncontraq):
    contrsets = list(range(nqubits))
    return [contrsets[i:i+ncontraq] for i in range(0, nqubits, ncontraq)]

# Author: <NAME> <<EMAIL>> # License: BSD 3 clause Contracts sublists of a list of layers vertically (merging cores of multiple layers for a single qubit) up to     some maximum contraction depth.      Parameters     ----------     layer_list : List of tt-tensor layers     ncontrl : Maximum number of layers to contract into each composite layer.      Returns     -------     List of layers resulting from vertically contracting the layers of layer_list #merge list of core lists (layers) into single list Contracts list of layers vertically (merging cores of multiple layers for a single qubit).      Parameters     ----------     layer_list : List of tt-tensor layers      Returns     -------     Single layer (list of cores) resulting from vertically contracting the layers of layer_list #first layer of each qubit to start contraction #abdefg --> aedbfg --> aebdfg --> #cont = transpose(transpose(transpose(cont, 1, 3), 2, 3), 3, 4) #pytorch #numpy and interface Contracts lists (layers) of tt-tensor cores horizontally (merging multiple cores in a single layer) up to     some maximum number of qubits.      Parameters     ----------     layer : List of tt-tensor cores     ncontrq : Maximum number of qubits to contract into each composite core.      Returns     -------     Layer (list of cores) resulting from horizontally contracting the cores of layer #if cores have 3 indices, state tensor #if cores do not have 3 indices (here, they would have 4 indices), operator tensor Contracts list of qubits horizontally (merging multiple cores in a single layer).      Parameters     ----------     qubits : List of tt-tensor cores     state : tt-tensor, input state to be evolved by unitary      Returns     -------     Single core resulting from horizontally contracting the cores of qubits #cont = transpose(cont, 2, 3) # pytorch # numpy and cupy interface Contracts list of qubits horizontally (merging multiple cores in a single layer).      Parameters     ----------     qubits : List of tt-tensor cores     state : tt-tensor, input state to be evolved by unitary          Returns     -------     Single core resulting from horizontally contracting the cores of qubits