Optics¶

@abstractmethod
def forward(self) -> Morphism:
    """The forward (get/match) morphism.

    Returns
    -------
    Morphism
        A morphism extracting the focus from the source.
    """
    ...

@abstractmethod
def backward(self) -> Morphism:
    """The backward (put/build) morphism.

    Returns
    -------
    Morphism
        A morphism putting the focus back into the whole.
    """
    ...

def as_profunctor(self) -> Profunctor:
    """View this optic as a profunctor S ↛ T.

    Computes the profunctor representation by composing
    forward and backward through the focus.

    Returns
    -------
    Profunctor
        The profunctor representation of this optic.
    """
    fwd = self.forward()
    bwd = self.backward()

    # the profunctor view depends on the optic type,
    # but generally it's fwd >> bwd or a similar composition
    tensor = self._quantale.compose(fwd.tensor, bwd.tensor, self._focus_source.ndim)

    return Profunctor(
        contra=self._source,
        co=self._target,
        tensor=tensor,
        quantale=self._quantale,
    )

def forward(self) -> Morphism:
    """get: S → A (projection to focus component).

    Returns
    -------
    ObservedMorphism
        The projection morphism.
    """
    q = self._quantale
    whole = self._source
    focus = self._focus_source

    data = torch.full((*whole.shape, *focus.shape), q.zero)

    # build projection tensor
    assert isinstance(whole, ProductSet)
    components = whole.components

    # compute dimension offset for focus
    offset = 0

    for i in range(self._focus_index):
        offset += components[i].ndim

    # for each element of the whole, project to the focus
    for idx in itertools.product(*(range(s) for s in whole.shape)):
        focus_idx = idx[offset : offset + focus.ndim]
        data[idx + focus_idx] = q.unit

    return observed(whole, focus, data, quantale=q)

def backward(self) -> Morphism:
    """put: A → S (embed focus back, averaging over complement).

    In the V-enriched setting, the "put" creates a morphism
    A → S where fixing the focus component and joining over
    the complement gives a fuzzy relation.

    Returns
    -------
    ObservedMorphism
        The embedding morphism.
    """
    q = self._quantale
    whole = self._source
    focus = self._focus_source

    # put: A → S where S = A × C
    # for each a, put(a, (a', c)) = δ(a, a') for the focus,
    # and the unit for the complement
    assert isinstance(whole, ProductSet)

    data = torch.full((*focus.shape, *whole.shape), q.zero)

    # compute dimension offset for focus
    components = whole.components
    offset = 0

    for i in range(self._focus_index):
        offset += components[i].ndim

    for focus_idx in itertools.product(*(range(s) for s in focus.shape)):
        # for each complement index, set the delta
        complement_shapes: list[tuple[int, ...]] = []

        for i, comp in enumerate(components):
            if i != self._focus_index:
                complement_shapes.append(tuple(s for s in comp.shape))

        # iterate over all whole indices where focus matches
        for whole_idx in itertools.product(*(range(s) for s in whole.shape)):
            w_focus = whole_idx[offset : offset + focus.ndim]

            if w_focus == focus_idx:
                data[focus_idx + whole_idx] = q.unit

    return observed(focus, whole, data, quantale=q)

def forward(self) -> Morphism:
    """match: S → A (partial extraction, zero on non-matching).

    Returns
    -------
    ObservedMorphism
        The matching morphism.
    """
    q = self._quantale
    whole = self._source
    focus = self._focus_source

    assert isinstance(whole, CoproductSet)
    start, end = whole.component_range(self._focus_index)

    data = torch.full((whole.size, focus.size), q.zero)

    for i in range(focus.size):
        data[start + i, i] = q.unit

    return observed(whole, focus, data, quantale=q)

def backward(self) -> Morphism:
    """build: A → S (inject focus into the coproduct).

    Returns
    -------
    ObservedMorphism
        The injection morphism.
    """
    q = self._quantale
    whole = self._source
    focus = self._focus_source

    assert isinstance(whole, CoproductSet)
    start, end = whole.component_range(self._focus_index)

    data = torch.full((focus.size, whole.size), q.zero)

    for i in range(focus.size):
        data[i, start + i] = q.unit

    return observed(focus, whole, data, quantale=q)

def forward(self) -> Morphism:
    """from: S → A."""
    return self._from_morph

def backward(self) -> Morphism:
    """to: B → T."""
    return self._to_morph

def verify_isomorphism(self, atol: float = 1e-5) -> bool:
    """Verify that forward and backward form an isomorphism.

    Checks from >> to ≈ id_S and to >> from ≈ id_A (when S=T, A=B).

    Parameters
    ----------
    atol : float
        Absolute tolerance.

    Returns
    -------
    bool
        True if the adapter is an isomorphism.
    """
    if self._source != self._target or self._focus_source != self._focus_target:
        return False

    fwd = self._from_morph
    bwd = self._to_morph

    # fwd >> bwd ≈ id_S
    roundtrip_s = (fwd >> bwd).tensor
    id_s = identity(self._source, quantale=self._quantale).tensor

    if not torch.allclose(roundtrip_s, id_s, atol=atol):
        return False

    # bwd >> fwd ≈ id_A
    roundtrip_a = (bwd >> fwd).tensor
    id_a = identity(self._focus_source, quantale=self._quantale).tensor

    return torch.allclose(roundtrip_a, id_a, atol=atol)

def forward(self) -> Morphism:
    """Extract focus by evaluating at each index.

    Produces S → A by marginalizing over the index.
    """
    q = self._quantale
    source = self._source
    focus = self._focus_source

    # use the cotraverse tensor transposed as the getter
    # simplified: just project
    data = q.identity_tensor(source.shape)

    # if source and focus match, return identity
    if source.shape == focus.shape:
        return observed(source, focus, data, quantale=q)

    # otherwise, use marginalization
    return observed(source, focus, self._cotraverse_tensor, quantale=q)

def backward(self) -> Morphism:
    """Rebuild whole from focus via cotraverse."""
    q = self._quantale
    focus = self._focus_target
    source = self._target

    if source.shape == focus.shape:
        data = q.identity_tensor(source.shape)
        return observed(focus, source, data, quantale=q)

    # transpose the cotraverse
    n_src = len(source.shape)
    n_foc = len(focus.shape)
    perm = list(range(n_src, n_src + n_foc)) + list(range(n_src))

    return observed(
        focus,
        source,
        self._cotraverse_tensor.permute(*perm),
        quantale=q,
    )