Traced Monoidal¶

@abstractmethod
def trace(
    self,
    morph: Morphism,
    feedback: SetObject,
    domain: SetObject,
    codomain: SetObject,
) -> Morphism:
    """Apply the trace operator.

    Given f: A ⊗ U → B ⊗ U, compute Tr^U(f): A → B.

    Parameters
    ----------
    morph : Morphism
        The morphism f: A ⊗ U → B ⊗ U.
    feedback : SetObject
        The feedback object U.
    domain : SetObject
        The external domain A.
    codomain : SetObject
        The external codomain B.

    Returns
    -------
    Morphism
        The traced morphism Tr^U(f): A → B.
    """
    ...

def verify_yanking(
    self,
    feedback: SetObject,
    atol: float = 1e-5,
) -> bool:
    """Verify the yanking axiom: Tr^U(σ_{U,U}) = id_U.

    The trace of the braiding/swap should be the identity.

    Parameters
    ----------
    feedback : SetObject
        The object U to test.
    atol : float
        Absolute tolerance.

    Returns
    -------
    bool
        True if yanking holds within tolerance.
    """
    # σ_{U,U}: U ⊗ U → U ⊗ U (swap)
    if isinstance(self._monoidal, CartesianMonoidal):
        swap = self._monoidal.braiding(feedback, feedback)
        traced = self.trace(swap, feedback, feedback, feedback)
        expected = identity(feedback, quantale=self._quantale).tensor

        return torch.allclose(traced.tensor, expected, atol=atol)

    return True  # skip for non-cartesian

def trace(
    self,
    morph: Morphism,
    feedback: SetObject,
    domain: SetObject,
    codomain: SetObject,
) -> Morphism:
    """Compute the cartesian trace via diagonal extraction and join.

    Parameters
    ----------
    morph : Morphism
        f: A × U → B × U.
    feedback : SetObject
        The feedback object U.
    domain : SetObject
        The external domain A.
    codomain : SetObject
        The external codomain B.

    Returns
    -------
    ObservedMorphism
        Tr^U(f): A → B.
    """
    q = self._quantale
    t = morph.tensor

    # tensor shape: (*a.shape, *u.shape, *b.shape, *u.shape)
    # we need to extract the diagonal over the two U copies
    # and join over U

    # the U dimensions in domain are at positions [n_a, n_a + n_u)
    # the U dimensions in codomain are at positions [n_a + n_u + n_b, n_a + n_u + n_b + n_u)
    result_shape = (*domain.shape, *codomain.shape)
    result = torch.full(result_shape, q.zero)

    # extract diagonal and join over U
    for a_idx in itertools.product(*(range(s) for s in domain.shape)):
        for b_idx in itertools.product(*(range(s) for s in codomain.shape)):
            # collect f(a, u, b, u) for all u
            vals: list[torch.Tensor] = []

            for u_idx in itertools.product(*(range(s) for s in feedback.shape)):
                src_idx = a_idx + u_idx + b_idx + u_idx
                vals.append(t[src_idx].unsqueeze(0))

            if vals:
                stacked = torch.cat(vals)
                result[a_idx + b_idx] = q.join(stacked, dim=0)

    return observed(domain, codomain, result, quantale=q)

def trace(
    self,
    morph: Morphism,
    feedback: SetObject,
    domain: SetObject,
    codomain: SetObject,
) -> Morphism:
    """Compute the trace via fixpoint iteration.

    Parameters
    ----------
    morph : Morphism
        f: A ⊗ U → B ⊗ U.
    feedback : SetObject
        The feedback object U.
    domain : SetObject
        The external domain A.
    codomain : SetObject
        The external codomain B.

    Returns
    -------
    ObservedMorphism
        Tr^U(f): A → B.
    """
    q = self._quantale
    t = morph.tensor

    result_shape = (*domain.shape, *codomain.shape)

    # initialize with zero (bottom)
    prev = torch.full(result_shape, q.zero)
    current = prev

    for _ in range(self._max_iter):
        current = torch.full(result_shape, q.zero)

        for a_idx in itertools.product(*(range(s) for s in domain.shape)):
            for b_idx in itertools.product(*(range(s) for s in codomain.shape)):
                vals: list[torch.Tensor] = []

                for u_idx in itertools.product(*(range(s) for s in feedback.shape)):
                    src_idx = a_idx + u_idx + b_idx + u_idx
                    vals.append(t[src_idx].unsqueeze(0))

                if vals:
                    stacked = torch.cat(vals)
                    current[a_idx + b_idx] = q.join(stacked, dim=0)

        if torch.allclose(current, prev, atol=self._atol):
            break

        prev = current

    return observed(domain, codomain, current, quantale=q)