Functors¶

@abstractmethod
def map_object(self, obj: SetObject) -> SetObject:
    """Apply the functor to an object.

    Parameters
    ----------
    obj : SetObject
        Source object.

    Returns
    -------
    SetObject
        Image object.
    """
    ...

@abstractmethod
def map_morphism(self, morph: Morphism) -> FunctorMorphism:
    """Apply the functor to a morphism.

    Returns a lazy FunctorMorphism that recomputes its tensor
    from the inner morphism, preserving gradient flow.

    Parameters
    ----------
    morph : Morphism
        Source morphism.

    Returns
    -------
    FunctorMorphism
        Image morphism.
    """
    ...

@abstractmethod
def map_tensor(self, tensor: torch.Tensor, quantale: Quantale) -> torch.Tensor:
    """Apply the functor's action on the raw tensor level.

    This is the computational core used by FunctorMorphism.tensor.

    Parameters
    ----------
    tensor : torch.Tensor
        The inner morphism's tensor.
    quantale : Quantale
        The enrichment algebra.

    Returns
    -------
    torch.Tensor
        The image tensor.
    """
    ...

def map_object(self, obj: SetObject) -> SetObject:
    """Identity on objects: A ↦ A."""
    return obj

def map_morphism(self, morph: Morphism) -> FunctorMorphism:
    """Identity on morphisms: f ↦ f (wrapped as FunctorMorphism)."""
    return FunctorMorphism(self, morph, morph.domain, morph.codomain)

def map_tensor(self, tensor: torch.Tensor, quantale: Quantale) -> torch.Tensor:
    """Identity on tensors."""
    return tensor

def map_object(self, obj: SetObject) -> SetObject:
    """F(G(A))."""
    return self._outer.map_object(self._inner.map_object(obj))

def map_morphism(self, morph: Morphism) -> FunctorMorphism:
    """F(G(f)) as a lazy FunctorMorphism."""
    domain = self.map_object(morph.domain)
    codomain = self.map_object(morph.codomain)
    return FunctorMorphism(self, morph, domain, codomain)

def map_tensor(self, tensor: torch.Tensor, quantale: Quantale) -> torch.Tensor:
    """F's tensor action applied to G's tensor action."""
    inner_result = self._inner.map_tensor(tensor, quantale)
    return self._outer.map_tensor(inner_result, quantale)

def map_object(self, obj: SetObject) -> FreeMonoid:
    """Apply functor to an object: A ↦ FreeMonoid(generators=A, max_length=max_length).

    For FinSet inputs, uses obj directly as generators. For other
    SetObject types (e.g., FreeMonoid, CoproductSet), creates a
    proxy FinSet with the same cardinality — this is needed for
    iterated application (e.g., the triangle identities of an
    adjunction require F(F(A)) = (A*)*).

    Parameters
    ----------
    obj : SetObject
        The generator set. Typically a FinSet, but any SetObject
        is accepted by treating it as a flat set.

    Returns
    -------
    FreeMonoid
        The free monoid on obj.
    """
    if not isinstance(obj, FinSet):
        # treat any SetObject as a flat set with its total cardinality
        obj = FinSet(name=getattr(obj, "name", repr(obj)), cardinality=obj.size)

    return FreeMonoid(generators=obj, max_length=self._max_length)

def map_morphism(self, morph: Morphism) -> FunctorMorphism:
    """Apply functor to a morphism: f ↦ f* (block-diagonal).

    Parameters
    ----------
    morph : Morphism
        A morphism between FinSets.

    Returns
    -------
    FunctorMorphism
        Lazy image preserving gradient flow.
    """
    domain = self.map_object(morph.domain)
    codomain = self.map_object(morph.codomain)

    return FunctorMorphism(self, morph, domain, codomain)

def map_tensor(self, tensor: torch.Tensor, quantale: Quantale) -> torch.Tensor:
    """Build the block-diagonal tensor for the free monoid action.

    For each k=0..max_length, computes componentwise_lift(f, k),
    reshapes the result to 2D (n_a^k × n_b^k), and assembles all
    blocks via torch.block_diag.

    Parameters
    ----------
    tensor : torch.Tensor
        The inner morphism's 2D tensor of shape (n_a, n_b).
    quantale : Quantale
        The enrichment algebra (passed to componentwise_lift).

    Returns
    -------
    torch.Tensor
        Block-diagonal tensor of shape (total_a, total_b).
    """
    n_a, n_b = tensor.shape
    blocks: list[torch.Tensor] = []

    for k in range(self._max_length + 1):
        lifted = componentwise_lift(tensor, k, quantale=quantale)

        # reshape from (n_a,)*k + (n_b,)*k to (n_a^k, n_b^k)
        rows = n_a**k if k > 0 else 1
        cols = n_b**k if k > 0 else 1
        blocks.append(lifted.reshape(rows, cols))

    return torch.block_diag(*blocks)