<commit_before>"""Classes and algorithms related to 1D tensor networks.
"""

import re
import functools
from math import log2
from numbers import Integral

import numpy as np
import opt_einsum as oe
from autoray import do, dag

from ..utils import check_opt, three_line_multi_print, pairwise
import quimb as qu
from .tensor_core import (
    Tensor,
    TensorNetwork,
    rand_uuid,
    bonds,
    tags2set,
    get_tags,
    _asarray,
    _ndim,
)


def align_TN_1D(*tns, ind_ids=None, inplace=False):
    r"""Align an arbitrary number of 1D tensor networks in a stack-like
    geometry::

        a-a-a-a-a-a-a-a-a-a-a-a-a-a-a-a-a-a
        | | | | | | | | | | | | | | | | | | <- ind_ids[0] (defaults to 1st id)
        b-b-b-b-b-b-b-b-b-b-b-b-b-b-b-b-b-b
        | | | | | | | | | | | | | | | | | | <- ind_ids[1]
                       ...
        | | | | | | | | | | | | | | | | | | <- ind_ids[-2]
        y-y-y-y-y-y-y-y-y-y-y-y-y-y-y-y-y-y
        | | | | | | | | | | | | | | | | | | <- ind_ids[-1]
        z-z-z-z-z-z-z-z-z-z-z-z-z-z-z-z-z-z

    Parameters
    ----------
    tns : sequence of TensorNetwork1D
        The 1D TNs to align.
    ind_ids : None, or sequence of str
        String with format specifiers to id each level of sites with. Will be
        automatically generated like ``(tns[0].site_ind_id, "__ind_a{}__",
        "__ind_b{}__", ...)`` if not given.
    """
    if not inplace:
        tns = [tn.copy() for tn in tns]

    if ind_ids is None:
        ind_ids = ([tns[0].site_ind_id] +
                   ["__ind_{}".format(oe.get_symbol(i)) + "{}__"
                    for i in range(len(tns) - 2)])
    else:
        ind_ids = tuple(ind_ids)

    for i, tn in enumerate(tns):
        if isinstance(tn, TensorNetwork1DVector):
            if i == 0:
                tn.site_ind_id = ind_ids[i]
            elif i == len(tns) - 1:
                tn.site_ind_id = ind_ids[i - 1]
            else:
                raise ValueError("An 1D TN vector can only be aligned as the "
                                 "first or last TN in a sequence.")

        elif isinstance(tn, MatrixProductOperator):
            tn.upper_ind_id = ind_ids[i - 1]
            tn.lower_ind_id = ind_ids[i]

        else:
            raise ValueError("Can only align MPS and MPOs currently.")

    return tns


def expec_TN_1D(*tns, compress=None, eps=1e-15):
    """Compute the expectation of several 1D TNs, using transfer matrix
    compression if any are periodic.

    Parameters
    ----------
    tns : sequence of TensorNetwork1D
        The MPS and MPO to find expectation of. Should start and begin with
        an MPS e.g. ``(MPS, MPO, ...,  MPS)``.
    compress : {None, False, True}, optional
        Whether to perform transfer matrix compression on cyclic systems. If
        set to ``None`` (the default), decide heuristically.
    eps : float, optional
        The accuracy of the transfer matrix compression.

    Returns
    -------
    x : float
        The expectation value.
    """
    expec_tn = TensorNetwork(align_TN_1D(*tns))

    # if OBC or <= 0.0 specified use exact contraction
    cyclic = any(tn.cyclic for tn in tns)<commit_msg>TN align_TN_1D: allow MPO at beginning and end<commit_after>"""Classes and algorithms related to 1D tensor networks.
"""

import re
import functools
from math import log2
from numbers import Integral

import numpy as np
import opt_einsum as oe
from autoray import do, dag

from ..utils import check_opt, three_line_multi_print, pairwise
import quimb as qu
from .tensor_core import (
    Tensor,
    TensorNetwork,
    rand_uuid,
    bonds,
    tags2set,
    get_tags,
    _asarray,
    _ndim,
)


def align_TN_1D(*tns, ind_ids=None, inplace=False):
    r"""Align an arbitrary number of 1D tensor networks in a stack-like
    geometry::

        a-a-a-a-a-a-a-a-a-a-a-a-a-a-a-a-a-a
        | | | | | | | | | | | | | | | | | | <- ind_ids[0] (defaults to 1st id)
        b-b-b-b-b-b-b-b-b-b-b-b-b-b-b-b-b-b
        | | | | | | | | | | | | | | | | | | <- ind_ids[1]
                       ...
        | | | | | | | | | | | | | | | | | | <- ind_ids[-2]
        y-y-y-y-y-y-y-y-y-y-y-y-y-y-y-y-y-y
        | | | | | | | | | | | | | | | | | | <- ind_ids[-1]
        z-z-z-z-z-z-z-z-z-z-z-z-z-z-z-z-z-z

    Parameters
    ----------
    tns : sequence of TensorNetwork1D
        The 1D TNs to align.
    ind_ids : None, or sequence of str
        String with format specifiers to id each level of sites with. Will be
        automatically generated like ``(tns[0].site_ind_id, "__ind_a{}__",
        "__ind_b{}__", ...)`` if not given.
    """
    if not inplace:
        tns = [tn.copy() for tn in tns]

    n = len(tns)

    if ind_ids is None:
        ind_ids = ["__ind_{}".format(oe.get_symbol(i)) + "{}__"
                   for i in range(n - 1)]
    else:
        ind_ids = tuple(ind_ids)

    for i, tn in enumerate(tns):
        if isinstance(tn, TensorNetwork1DVector):
            if i == 0:
                tn.site_ind_id = ind_ids[i]
            elif i == n - 1:
                tn.site_ind_id = ind_ids[i - 1]
            else:
                raise ValueError("An 1D TN vector can only be aligned as the "
                                 "first or last TN in a sequence.")

        elif isinstance(tn, MatrixProductOperator):
            if i != 0:
                tn.upper_ind_id = ind_ids[i - 1]
            if i != n - 1:
                tn.lower_ind_id = ind_ids[i]

        else:
            raise ValueError("Can only align MPS and MPOs currently.")

    return tns


def expec_TN_1D(*tns, compress=None, eps=1e-15):
    """Compute the expectation of several 1D TNs, using transfer matrix
    compression if any are periodic.

    Parameters
    ----------
    tns : sequence of TensorNetwork1D
        The MPS and MPO to find expectation of. Should start and begin with
        an MPS e.g. ``(MPS, MPO, ...,  MPS)``.
    compress : {None, False, True}, optional
        Whether to perform transfer matrix compression on cyclic systems. If
        set to ``None`` (the default), decide heuristically.
    eps : float, optional
        The accuracy of the transfer matrix compression.

    Returns
    -------
    x : float
        The expectation value.
    """
    expec_tn = TensorNetwork(align_TN_1D(*tns))

    # if OBC or <= 0.0 specified use exact contraction
    cyclic = any(tn.cyclic for tn in tns)
{'fraction_non_alphanumeric': 0.13309352517985612, 'fraction_numerical': 0.009696590553644042, 'mean_word_length': 1.814700704225352, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 26, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7627408', 'n_tokens_mistral': 2295, 'n_tokens_neox': 2196, 'n_words': 612}