Trajectree 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

trajectree/quimb/quimb/tensor/fitting.py

@@ -1,385 +0,0 @@
-"""Tools for computing distances between and fitting tensor networks."""
-from autoray import dag, do
-
-from .contraction import contract_strategy
-from ..utils import check_opt
-
-
-def tensor_network_distance(
-    tnA,
-    tnB,
-    xAA=None,
-    xAB=None,
-    xBB=None,
-    method="auto",
-    normalized=False,
-    **contract_opts,
-):
-    r"""Compute the Frobenius norm distance between two tensor networks:
-
-    .. math::
-
-            D(A, B)
-            = | A - B |_{\mathrm{fro}}
-            = \mathrm{Tr} [(A - B)^{\dagger}(A - B)]^{1/2}
-            = ( \langle A | A \rangle - 2 \mathrm{Re} \langle A | B \rangle|
-            + \langle B | B \rangle ) ^{1/2}
-
-    which should have matching outer indices. Note the default approach to
-    computing the norm is precision limited to about ``eps**0.5`` where ``eps``
-    is the precision of the data type, e.g. ``1e-8`` for float64. This is due
-    to the subtraction in the above expression.
-
-    Parameters
-    ----------
-    tnA : TensorNetwork or Tensor
-        The first tensor network operator.
-    tnB : TensorNetwork or Tensor
-        The second tensor network operator.
-    xAA : None or scalar
-        The value of ``A.H @ A`` if you already know it (or it doesn't matter).
-    xAB : None or scalar
-        The value of ``A.H @ B`` if you already know it (or it doesn't matter).
-    xBB : None or scalar
-        The value of ``B.H @ B`` if you already know it (or it doesn't matter).
-    method : {'auto', 'overlap', 'dense'}, optional
-        How to compute the distance. If ``'overlap'``, the default, the
-        distance will be computed as the sum of overlaps, without explicitly
-        forming the dense operators. If ``'dense'``, the operators will be
-        directly formed and the norm computed, which can be quicker when the
-        exterior dimensions are small. If ``'auto'``, the dense method will
-        be used if the total operator (outer) size is ``<= 2**16``.
-    normalized : bool, optional
-        If ``True``, then normalize the distance by the norm of the two
-        operators, i.e. ``2 * D(A, B) / (|A| + |B|)``. The resulting distance
-        lies between 0 and 2 and is more useful for assessing convergence.
-    contract_opts
-        Supplied to :meth:`~quimb.tensor.tensor_core.TensorNetwork.contract`.
-
-    Returns
-    -------
-    D : float
-    """
-    check_opt("method", method, ("auto", "dense", "overlap"))
-
-    tnA = tnA.as_network()
-    tnB = tnB.as_network()
-
-    oix = tnA.outer_inds()
-    if set(oix) != set(tnB.outer_inds()):
-        raise ValueError(
-            "Can only compute distance between tensor "
-            "networks with matching outer indices."
-        )
-
-    if method == "auto":
-        d = tnA.inds_size(oix)
-        if d <= 1 << 16:
-            method = "dense"
-        else:
-            method = "overlap"
-
-    # directly from vectorizations of both
-    if method == "dense":
-        tnA = tnA.contract(..., output_inds=oix, preserve_tensor=True)
-        tnB = tnB.contract(..., output_inds=oix, preserve_tensor=True)
-
-    # overlap method
-    if xAA is None:
-        xAA = (tnA | tnA.H).contract(..., **contract_opts)
-    if xAB is None:
-        xAB = (tnA | tnB.H).contract(..., **contract_opts)
-    if xBB is None:
-        xBB = (tnB | tnB.H).contract(..., **contract_opts)
-
-    dAB = do("abs", xAA - 2 * do("real", xAB) + xBB) ** 0.5
-
-    if normalized:
-        dAB *= 2 / (do("abs", xAA)**0.5 + do("abs", xBB)**0.5)
-
-    return dAB
-
-
-
-
-def tensor_network_fit_autodiff(
-    tn,
-    tn_target,
-    steps=1000,
-    tol=1e-9,
-    autodiff_backend="autograd",
-    contract_optimize="auto-hq",
-    distance_method="auto",
-    inplace=False,
-    progbar=False,
-    **kwargs,
-):
-    """Optimize the fit of ``tn`` with respect to ``tn_target`` using
-    automatic differentation. This minimizes the norm of the difference
-    between the two tensor networks, which must have matching outer indices,
-    using overlaps.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to fit.
-    tn_target : TensorNetwork
-        The target tensor network to fit ``tn`` to.
-    steps : int, optional
-        The maximum number of autodiff steps.
-    tol : float, optional
-        The target norm distance.
-    autodiff_backend : str, optional
-        Which backend library to use to perform the gradient computation.
-    contract_optimize : str, optional
-        The contraction path optimized used to contract the overlaps.
-    distance_method : {'auto', 'dense', 'overlap'}, optional
-        Supplied to :func:`~quimb.tensor.tensor_core.tensor_network_distance`,
-        controls how the distance is computed.
-    inplace : bool, optional
-        Update ``tn`` in place.
-    progbar : bool, optional
-        Show a live progress bar of the fitting process.
-    kwargs
-        Passed to :class:`~quimb.tensor.tensor_core.optimize.TNOptimizer`.
-
-    See Also
-    --------
-    tensor_network_distance, tensor_network_fit_als
-    """
-    from .optimize import TNOptimizer
-    from .tensor_core import tensor_network_distance
-
-    xBB = (tn_target | tn_target.H).contract(
-        ...,
-        output_inds=(),
-        optimize=contract_optimize,
-    )
-
-    tnopt = TNOptimizer(
-        tn=tn,
-        loss_fn=tensor_network_distance,
-        loss_constants={"tnB": tn_target, "xBB": xBB},
-        loss_kwargs={"method": distance_method, "optimize": contract_optimize},
-        autodiff_backend=autodiff_backend,
-        progbar=progbar,
-        **kwargs,
-    )
-
-    tn_fit = tnopt.optimize(steps, tol=tol)
-
-    if not inplace:
-        return tn_fit
-
-    for t1, t2 in zip(tn, tn_fit):
-        t1.modify(data=t2.data)
-
-    return tn
-
-
-def _tn_fit_als_core(
-    var_tags,
-    tnAA,
-    tnAB,
-    xBB,
-    tol,
-    contract_optimize,
-    steps,
-    enforce_pos,
-    pos_smudge,
-    solver="solve",
-    progbar=False,
-):
-    from .tensor_core import group_inds
-
-    # shared intermediates + greedy = good reuse of contractions
-    with contract_strategy(contract_optimize):
-        # prepare each of the contractions we are going to repeat
-        env_contractions = []
-        for tg in var_tags:
-            # varying tensor and conjugate in norm <A|A>
-            tk = tnAA["__KET__", tg]
-            tb = tnAA["__BRA__", tg]
-
-            # get inds, and ensure any bonds come last, for linalg.solve
-            lix, bix, rix = group_inds(tb, tk)
-            tk.transpose_(*rix, *bix)
-            tb.transpose_(*lix, *bix)
-
-            # form TNs with 'holes', i.e. environment tensors networks
-            A_tn = tnAA.select((tg,), "!all")
-            y_tn = tnAB.select((tg,), "!all")
-
-            env_contractions.append((tk, tb, lix, bix, rix, A_tn, y_tn))
-
-        if tol != 0.0:
-            old_d = float("inf")
-
-        if progbar:
-            import tqdm
-
-            pbar = tqdm.trange(steps)
-        else:
-            pbar = range(steps)
-
-        # the main iterative sweep on each tensor, locally optimizing
-        for _ in pbar:
-            for tk, tb, lix, bix, rix, A_tn, y_tn in env_contractions:
-                Ni = A_tn.to_dense(lix, rix)
-                Wi = y_tn.to_dense(rix, bix)
-
-                if enforce_pos:
-                    el, ev = do("linalg.eigh", Ni)
-                    el = do("clip", el, el[-1] * pos_smudge, None)
-                    Ni_p = ev * do("reshape", el, (1, -1)) @ dag(ev)
-                else:
-                    Ni_p = Ni
-
-                if solver == "solve":
-                    x = do("linalg.solve", Ni_p, Wi)
-                elif solver == "lstsq":
-                    x = do("linalg.lstsq", Ni_p, Wi, rcond=pos_smudge)[0]
-
-                x_r = do("reshape", x, tk.shape)
-                # n.b. because we are using virtual TNs -> updates propagate
-                tk.modify(data=x_r)
-                tb.modify(data=do("conj", x_r))
-
-            # assess | A - B | for convergence or printing
-            if (tol != 0.0) or progbar:
-                xAA = do("trace", dag(x) @ (Ni @ x))  # <A|A>
-                xAB = do("trace", do("real", dag(x) @ Wi))  # <A|B>
-                d = do("abs", (xAA - 2 * xAB + xBB)) ** 0.5
-                if abs(d - old_d) < tol:
-                    break
-                old_d = d
-
-            if progbar:
-                pbar.set_description(str(d))
-
-
-def tensor_network_fit_als(
-    tn,
-    tn_target,
-    tags=None,
-    steps=100,
-    tol=1e-9,
-    solver="solve",
-    enforce_pos=False,
-    pos_smudge=None,
-    tnAA=None,
-    tnAB=None,
-    xBB=None,
-    contract_optimize="greedy",
-    inplace=False,
-    progbar=False,
-):
-    """Optimize the fit of ``tn`` with respect to ``tn_target`` using
-    alternating least squares (ALS). This minimizes the norm of the difference
-    between the two tensor networks, which must have matching outer indices,
-    using overlaps.
-
-    Parameters
-    ----------
-    tn : TensorNetwork
-        The tensor network to fit.
-    tn_target : TensorNetwork
-        The target tensor network to fit ``tn`` to.
-    tags : sequence of str, optional
-        If supplied, only optimize tensors matching any of given tags.
-    steps : int, optional
-        The maximum number of ALS steps.
-    tol : float, optional
-        The target norm distance.
-    solver : {'solve', 'lstsq', ...}, optional
-        The underlying driver function used to solve the local minimization,
-        e.g. ``numpy.linalg.solve`` for ``'solve'`` with ``numpy`` backend.
-    enforce_pos : bool, optional
-        Whether to enforce positivity of the locally formed environments,
-        which can be more stable.
-    pos_smudge : float, optional
-        If enforcing positivity, the level below which to clip eigenvalues
-        for make the local environment positive definite.
-    tnAA : TensorNetwork, optional
-        If you have already formed the overlap ``tn.H & tn``, maybe
-        approximately, you can supply it here. The unconjugated layer should
-        have tag ``'__KET__'`` and the conjugated layer ``'__BRA__'``. Each
-        tensor being optimized should have tag ``'__VAR{i}__'``.
-    tnAB : TensorNetwork, optional
-        If you have already formed the overlap ``tn_target.H & tn``, maybe
-        approximately, you can supply it here. Each tensor being optimized
-        should have tag ``'__VAR{i}__'``.
-    xBB : float, optional
-        If you have already know, have computed ``tn_target.H @ tn_target``,
-        or it doesn't matter, you can supply the value here.
-    contract_optimize : str, optional
-        The contraction path optimized used to contract the local environments.
-        Note ``'greedy'`` is the default in order to maximize shared work.
-    inplace : bool, optional
-        Update ``tn`` in place.
-    progbar : bool, optional
-        Show a live progress bar of the fitting process.
-
-    Returns
-    -------
-    TensorNetwork
-
-    See Also
-    --------
-    tensor_network_fit_autodiff, tensor_network_distance
-    """
-    # mark the tensors we are going to optimize
-    tna = tn.copy()
-    tna.add_tag("__KET__")
-
-    if tags is None:
-        to_tag = tna
-    else:
-        to_tag = tna.select_tensors(tags, "any")
-
-    var_tags = []
-    for i, t in enumerate(to_tag):
-        var_tag = f"__VAR{i}__"
-        t.add_tag(var_tag)
-        var_tags.append(var_tag)
-
-    # form the norm of the varying TN (A) and its overlap with the target (B)
-    if tnAA is None:
-        tnAA = tna | tna.H.retag_({"__KET__": "__BRA__"})
-    if tnAB is None:
-        tnAB = tna | tn_target.H
-
-    if (tol != 0.0) and (xBB is None):
-        # <B|B>
-        xBB = (tn_target | tn_target.H).contract(
-            ...,
-            optimize=contract_optimize,
-            output_inds=(),
-        )
-
-    if pos_smudge is None:
-        pos_smudge = max(tol, 1e-15)
-
-    _tn_fit_als_core(
-        var_tags=var_tags,
-        tnAA=tnAA,
-        tnAB=tnAB,
-        xBB=xBB,
-        tol=tol,
-        contract_optimize=contract_optimize,
-        steps=steps,
-        enforce_pos=enforce_pos,
-        pos_smudge=pos_smudge,
-        solver=solver,
-        progbar=progbar,
-    )
-
-    if not inplace:
-        tn = tn.copy()
-
-    for t1, t2 in zip(tn, tna):
-        # transpose so only thing changed in original TN is data
-        t2.transpose_like_(t1)
-        t1.modify(data=t2.data)
-
-    return tn