Omega categories as a data structure (Rough draft)

OmegaCats.py

"""
OmegaComputadsLibrary

Collection of data structures for omega computads and omega categories.
"""

from collections import deque

class Cell:
    """
    Represents a generic cell of a given dimension.
    Source and target are only defined for dimension > 0.
    """
    def __init__(self, dimension):
        self.dimension = dimension

    def source(self):
        if self.dimension == 0:
            raise ValueError("0-cells have no source.")
        return NotImplemented

    def target(self):
        if self.dimension == 0:
            raise ValueError("0-cells have no target.")
        return NotImplemented

    def is_parallel_to(self, other: "Cell") -> bool:
        """
        Checks if self and other are parallel:
          1) Same dimension
          2) If dimension=0, trivially parallel
          3) If dimension>0, then check that the sources and targets match.
        """
        if self.dimension != other.dimension:
            return False
        if self.dimension == 0:
            return True
        return (
            self.source() == other.source() and
            self.target() == other.target()
        )


class PastingScheme:
    """
    Stores the integer-list representation of a globular pasting scheme.
    Each entry of structure[d] is a list of cell counts at dimension d.

    Requirements:
      - structure must be a list of lists.
      - structure[0] must be length 1 (exactly one 0-cell count).
      - all counts >= 1.
      - for each d, len(structure[d+1]) = sum(structure[d]) - len(structure[d]).
      - the final dimension list must be all 1s.
    
    The interpretation is that each list entry represents a block of parallel cells in
    that dimension. If there are n parallel arrows, they must be filled by n-1
    cells in the next dimension up, telling us how many block there are in that
    next dimension.
    """

    def __init__(self, structure):
        self.structure = structure
        self.check_well_formed()

    def check_well_formed(self):
        if not isinstance(self.structure, list):
            raise ValueError("Structure must be a list.")
        if len(self.structure) == 0:
            raise ValueError("Must have at least one dimension.")

        # Dimension 0 must be a list of length 1.
        if not isinstance(self.structure[0], list) or len(self.structure[0]) != 1:
            raise ValueError("Dimension 0 must be a list of length 1.")

        for d, layer in enumerate(self.structure):
            if not isinstance(layer, list):
                raise ValueError(f"Dimension {d} is not a list.")
            if not all(isinstance(x, int) and x >= 1 for x in layer):
                raise ValueError(f"Dimension {d} has invalid cell counts.")

        for d in range(len(self.structure) - 1):
            s = sum(self.structure[d])
            n = len(self.structure[d])
            expected = s - n
            if len(self.structure[d + 1]) != expected:
                raise ValueError(
                    f"Dimension {d+1} list length must be sum(dim {d}) - len(dim {d})."
                )

        # Check that the final dimension is all 1s.
        final_dim = self.structure[-1]
        if not all(x == 1 for x in final_dim):
            raise ValueError("The final dimension must consist entirely of 1s.")

    @property
    def max_dimension(self):
        return len(self.structure)

    def num_cells_at_dimension(self, d):
        if d < 0 or d >= self.max_dimension:
            raise ValueError("Dimension out of range.")
        return sum(self.structure[d])

    def all_positions(self):
        positions = set()
        for d in range(self.max_dimension):
            total = self.num_cells_at_dimension(d)
            for i in range(total):
                positions.add(PastingSchemePosition(self, d, i))
        return positions

    def source_boundary(self):
        """
        Returns a set of positions that constitute the 'source boundary'
        for dimension (max_dimension - 2). Raises ValueError if (max_dimension - 2) < 0.

        The boundary includes:
          1) All positions in dimensions < (max_dimension - 2).
          2) The first cell in each block of dimension (max_dimension - 2).
             If structure[d] = [2, 1, 3], then blocks are of sizes 2, 1, 3 and
             the first cells are at indices 0, 2, and 3 (respectively).
        """
        d = self.max_dimension - 2
        if d < 0:
            raise ValueError("Not enough dimensionality for a source boundary.")

        boundary = set()
        # Include all cells in lower dimensions
        for lower_dim in range(d):
            total = self.num_cells_at_dimension(lower_dim)
            for i in range(total):
                boundary.add(PastingSchemePosition(self, lower_dim, i))

        # Now add the first cell in each block of dimension d
        index_start = 0
        for block_size in self.structure[d]:
            boundary.add(PastingSchemePosition(self, d, index_start))
            index_start += block_size

        return boundary

    def target_boundary(self):
        """
        Returns a set of positions that constitute the 'target boundary'
        for dimension (max_dimension - 2). Raises ValueError if (max_dimension - 2) < 0.

        The boundary includes:
          1) All positions in dimensions < (max_dimension - 2).
          2) The last cell in each block of dimension (max_dimension - 2).
             If structure[d] = [2, 1, 3], then blocks are of sizes 2, 1, 3 and
             the last cells are at indices 1, 2, and 5 (respectively).
        """
        d = self.max_dimension - 2
        if d < 0:
            raise ValueError("Not enough dimensionality for a target boundary.")

        boundary = set()
        # Include all cells in lower dimensions
        for lower_dim in range(d):
            total = self.num_cells_at_dimension(lower_dim)
            for i in range(total):
                boundary.add(PastingSchemePosition(self, lower_dim, i))

        # Now add the last cell in each block of dimension d
        index_start = 0
        for block_size in self.structure[d]:
            last_cell_index = index_start + block_size - 1
            boundary.add(PastingSchemePosition(self, d, last_cell_index))
            index_start += block_size

        return boundary

    def __eq__(self, other):
        if not isinstance(other, PastingScheme):
            return False
        return self.structure == other.structure

    def maximal_cells(self):
        """
        Calculates all cells which do not appear as the boundary of a higher
        dimensional cell.
        
        Returns a set of PastingSchemePosition objects corresponding to
        any block in the structure that has exactly '1' as its entry,
        across all dimensions. In other words, for each dimension d,
        each block_size == 1 identifies exactly one cell at that dimension.
        """
        results = set()
        for d in range(self.max_dimension):
            start_index = 0
            for block_size in self.structure[d]:
                if block_size == 1:
                    results.add(PastingSchemePosition(self, d, start_index))
                start_index += block_size
        return results


class PastingSchemePosition(Cell):
    """
    Identifies a specific cell within the PastingScheme by (dimension, index).
    """

    def __init__(self, scheme, dimension, index):
        super().__init__(dimension)
        self.scheme = scheme
        self.index = index
        if not self.is_well_formed():
            raise ValueError("Invalid position.")

    def is_well_formed(self):
        if self.dimension < 0 or self.dimension >= self.scheme.max_dimension:
            return False
        if self.index < 0 or self.index >= self.scheme.num_cells_at_dimension(self.dimension):
            return False
        return True

    def __hash__(self):
        return hash((id(self.scheme), self.dimension, self.index))

    def __eq__(self, other):
        if not isinstance(other, PastingSchemePosition):
            return False
        return (
            self.scheme == other.scheme and
            self.dimension == other.dimension and
            self.index == other.index
        )

    def _source_segment_index(self):
        """
        Returns just the source (d-1)-cell index bounding this dimension-d cell.
        
        Source/target are computed by iterating over structure[d-1] to find
        which boundary segment the cell index falls into. Each lumps_count in
        structure[d-1] gives lumps_count parallel (d-1)-cells, so that yields
        (lumps_count - 1) boundary segments. The dimension-d array structure[d]
        tells how many dimension-d cells fill each boundary segment in order.
        """
        if self.dimension == 0:
            raise ValueError("0-cells have no source.")

        lumps = self.scheme.structure[self.dimension - 1]
        fills = self.scheme.structure[self.dimension]

        i = self.index
        dim_minus_1_start = 0
        seg_index = 0
        for lumps_count in lumps:
            boundary_count = lumps_count - 1
            for seg in range(boundary_count):
                c = fills[seg_index]
                if i < c:
                    return dim_minus_1_start + seg
                i -= c
                seg_index += 1
            dim_minus_1_start += lumps_count

        raise ValueError("Index out of range when searching boundary segment.")

    def source(self):
        if self.dimension == 0:
            raise ValueError("0-cells have no source.")
        src_idx = self._source_segment_index()
        return PastingSchemePosition(self.scheme, self.dimension - 1, src_idx)

    def target(self):
        if self.dimension == 0:
            raise ValueError("0-cells have no target.")
        src_idx = self._source_segment_index()
        return PastingSchemePosition(self.scheme, self.dimension - 1, src_idx + 1)


class Substitution:
    """
    Represents a substitution on a PastingScheme, mapping every Position
    to a Cell. The user supplies a map from the scheme's maximal positions
    to Cell objects, and the constructor expands this map to all positions
    recursively and matches sources and targets to verify consistency.

    If any inconsistency is found (e.g., inferred source of a position
    disagrees with the already-mapped source), a ValueError is raised.
    """
    def __init__(self, scheme, assignment):
        """
        scheme: a PastingScheme instance
        assignment: dict mapping PastingSchemePosition -> Cell
                    for exactly the scheme's maximal positions
        """
        self.scheme = scheme

        # We assume the user supplies exactly the set of maximal positions.
        # (We do not forcibly error if that fails, but the BFS approach can break if mismatched.)
        maximal_keys = scheme.maximal_cells()
        if set(assignment.keys()) != maximal_keys:
            raise ValueError("Assignment keys must match the scheme's maximal positions.")

        self._map = dict(assignment)

        queue = deque(self._map.keys())

        while queue:
            p = queue.popleft()
            cell_p = self._map[p]

            if p.dimension > 0:
                s = p.source()
                t = p.target()

                # Source
                if s in self._map:
                    if self._map[s] != cell_p.source():
                        raise ValueError(
                            f"Inconsistent source at {s}: "
                            f"expected {cell_p.source()}, found {self._map[s]}."
                        )
                else:
                    self._map[s] = cell_p.source()
                    if s.dimension > 0:
                        queue.append(s)

                # Target
                if t in self._map:
                    if self._map[t] != cell_p.target():
                        raise ValueError(
                            f"Inconsistent target at {t}: "
                            f"expected {cell_p.target()}, found {self._map[t]}."
                        )
                else:
                    self._map[t] = cell_p.target()
                    if t.dimension > 0:
                        queue.append(t)

        self.full_map = self._map

    def maximal_map(self):
        """
        Return a dictionary from the scheme's maximal positions to their assigned Cell.
        We assume BFS expansions have populated all positions in self.full_map.
        """
        result = {}
        for mp in self.scheme.maximal_cells():
            if mp not in self.full_map:
                raise ValueError(f"Maximal position {mp} missing in BFS expansions.")
            result[mp] = self.full_map[mp]
        return result


    def __getitem__(self, position):
        """Returns the Cell mapped to 'position'."""
        return self.full_map[position]

    def __eq__(self, other):
        if not isinstance(other, Substitution):
            return False
        if self.scheme != other.scheme:
            return False
        if set(self.full_map.keys()) != set(other.full_map.keys()):
            return False
        for k, v in self.full_map.items():
            if v != other.full_map[k]:
                return False
        return True


class Coherence:
    """
    A Coherence has:
      - a PastingScheme
      - a source, which is either a PastingSchemePosition or an InstantiatedCoherence
      - a target, similarly either a PastingSchemePosition or an InstantiatedCoherence

    The dimension is (scheme.max_dimension - 1).
    """

    def __init__(self, scheme, source, target):
        self.scheme = scheme
        self.source = source
        self.target = target
        self._is_operation = None

    @property
    def dimension(self):
        return self.scheme.max_dimension - 1

    @property
    def is_operation(self):
        return self._is_operation

    def __eq__(self, other):
        if not isinstance(other, Coherence):
            return False
        return (self.scheme == other.scheme
                and self.source == other.source
                and self.target == other.target)

    def check_coherence(self):
        """
        Checks that a coherence is actually coherent
        
        1) Ensure that the source and target have the correct scheme.
        2) Gather all positions from source and target (recursively).
        3) If source.dimension != target.dimension, raise an error.
           If dimension>0, check source.source()==target.source() and
           source.target()==target.target() to check that they are parallel.
        4) Compare the sets of positions from source and target to
           self.scheme.source_boundary(), self.scheme.target_boundary(),
           and self.scheme.all_positions() to decide whether the coherence is
           operational (3a) or equational (3b). Sets _is_operation accordingly.
        """
        # 1) Ensure that source and target have the correct scheme.
        if self.source.scheme != self.scheme or self.target.scheme != self.scheme:
            raise ValueError("Source or target belongs to a different scheme.")

        # 2) Gather positions.
        source_positions = set()
        target_positions = set()
        self._gather_positions(self.source, source_positions)
        self._gather_positions(self.target, target_positions)

        # 3) Check that source and target are parellel.
        if not self.source.is_parallel_to(self.target):
            raise ValueError("Source and target are not parallel.")

        # 4) Check operational vs. equational.
        sb = self.scheme.source_boundary()
        tb = self.scheme.target_boundary()
        allpos = self.scheme.all_positions()

        if source_positions == sb and target_positions == tb:
            self._is_operation = True
        elif source_positions == target_positions == allpos:
            self._is_operation = False
        else:
            raise ValueError("Coherence is neither operationally nor equationally coherent.")

    def _gather_positions(self, obj, out_set):
        if isinstance(obj, PastingSchemePosition):
            if obj.scheme != self.scheme:
                raise ValueError("Position {obj} belongs to a different scheme.")
            out_set.add(obj)
        elif isinstance(obj, InstantiatedCoherence):
            if obj.coherence.scheme != self.scheme:
                raise ValueError("InstantiatedCoherence has a different scheme.")
            for v in obj.substitution.full_map.values():
                self._gather_positions(v, out_set)
        else:
            raise ValueError("Coherence source/target must be Position or InstantiatedCoherence.")

class InstantiatedCoherence(Cell):
    """
    An InstantiatedCoherence is a Cell whose dimension =
    (coherence.scheme.max_dimension - 1). It has:
      - coherence: the underlying Coherence
      - substitution: a Substitution on that coherence's scheme.
      - scheme: a pointer to that same scheme.

    The source, target are computed by applying the substitution
    to the coherence's source/target.
    """

    def __init__(self, coherence, substitution):
        super().__init__(dimension=coherence.scheme.max_dimension - 1)
        if substitution.scheme is not coherence.scheme:
            raise ValueError("Substitution scheme does not match coherence scheme.")
        if not isinstance(coherence, Coherence):
            raise ValueError("coherence must be a Coherence.")
        if not isinstance(substitution, Substitution):
            raise ValueError("substitution must be a Substitution.")
        self.coherence = coherence
        self.substitution = substitution
        self.scheme = coherence.scheme
        
        # Check that all elements of the substitution are PastingSchemePosition or
        # InstantiatedCoherence with the same scheme.
        for k, v in substitution.full_map.items():
            if not isinstance(v, PastingSchemePosition) or not isinstance(v, InstantiatedCoherence):
                raise ValueError("Substitution value at key {k} not PastingSchemePosition or InstantiatedCoherence.")
            if v.scheme != self.scheme:
                raise ValueError("Substitution value for key {k} from a different scheme.")

    def __eq__(self, other):
        if not isinstance(other, InstantiatedCoherence):
            return False
        return (self.coherence == other.coherence
                and self.substitution == other.substitution)

    def source(self):
        return self._resolve(self.coherence.source)

    def target(self):
        return self._resolve(self.coherence.target)

    def _resolve(self, obj):
        # If it's a PastingSchemePosition, we require it to be in our map.
        # If it's an InstantiatedCoherence, build a new substitution from its map.
        # Otherwise, error.
        if isinstance(obj, PastingSchemePosition):
            if obj not in self.substitution:
                raise ValueError(f"Position {obj} missing from the substitution.")
            val = self.substitution[obj]
            return self._resolve(val)
        elif isinstance(obj, InstantiatedCoherence):
            new_map = {}
            for k, v in obj.substitution.maximal_map().items():
                new_v = self._resolve(v)
                new_map[k] = new_v
            new_sub = Substitution(obj.coherence.scheme, new_map)
            return InstantiatedCoherence(obj.coherence, new_sub)
        else:
            raise ValueError("Coherence source/target must be positions or InstantiatedCoherence.")

class OmegaComputad:
    """
    Represents an omega-computad.
      1) zero_cell_count: a non-negative integer or the string 'infinity'.
      2) generator_count_fn: a callable that takes (source, target), which may be AtomicCell
         or OmegaCategoryCell from this computad, and returns a non-negative integer or 'infinity'.

    The correctness is not fully checked; we only verify the basic form of zero_cell_count.
    """
    def __init__(self, zero_cell_count, generator_count_fn):
        if not ( (isinstance(zero_cell_count, int) and zero_cell_count >= 0) or zero_cell_count == "infinity" ):
            raise ValueError("zero_cell_count must be a non-negative int or 'infinity'.")
        self.zero_cell_count = zero_cell_count
        self._generator_count_fn = generator_count_fn

    def number_of_generators(self, source, target):
        # We do not do thorough checking; we rely on user correctness.
        return self._generator_count_fn(source, target)


class AtomicCell(Cell):
    """
    An atomic cell in an omega-computad:
      - dimension
      - index
      - source, target (either None if dim=0, or must be dimension-1 cells)
      - must live in a given OmegaComputad
    """
    def __init__(self, computad, dimension, index, source=None, target=None):
        super().__init__(dimension)
        self.computad = computad
        self._index = index
        self._source = source
        self._target = target
        self._check_correctness()

    def _check_correctness(self):
        if not isinstance(self.computad, OmegaComputad):
            raise ValueError("computad must be an OmegaComputad.")
        if not (isinstance(self._index, int) and self._index >= 0):
            raise ValueError("index must be a non-negative integer.")

        # If dimension=0:
        if self.dimension == 0:
            if self._source is not None or self._target is not None:
                raise ValueError("0-dimensional atomic cell must have source=target=None.")
            # check index < zero_cell_count if finite
            zcc = self.computad.zero_cell_count
            if zcc != "infinity" and self._index >= zcc:
                raise ValueError("AtomicCell 0-d index too large.")
        else:
            # dimension>0: source and target must be dimension-1 cells.
            if not isinstance(self._source, Cell) or not isinstance(self._target, Cell):
                raise ValueError("source and target must be Cell instances.")
            if self._source.dimension != self.dimension - 1:
                raise ValueError("source dimension must be one less.")
            if self._target.dimension != self.dimension - 1:
                raise ValueError("target dimension must be one less.")
            if (hasattr(self._source, 'computad') and self._source.computad is not self.computad) or \
               (hasattr(self._target, 'computad') and self._target.computad is not self.computad):
                raise ValueError("All cells must belong to the same computad.")
            if not self._source.is_parallel_to(self._target):
                raise ValueError("Source and target must be parallel.")

            # check index < computad.number_of_generators(source, target) if not infinite
            gen_count = self.computad.number_of_generators(self._source, self._target)
            if gen_count != "infinity" and self._index >= gen_count:
                raise ValueError("AtomicCell index too large for given source/target.")

    def source(self):
        # Overriding Cell.source.
        if self.dimension == 0:
            return None
        return self._source

    def target(self):
        # Overriding Cell.target.
        if self.dimension == 0:
            return None
        return self._target

    def __eq__(self, other):
        if not isinstance(other, AtomicCell):
            return False
        return (
            self.dimension == other.dimension and
            self.computad is other.computad and
            self._index == other._index and
            self._source == other._source and
            self._target == other._target
        )


class OmegaCategoryCell(Cell):
    """
    An omega-category cell in an omega-computad:
      - dimension = coherence.scheme.max_dimension - 1
      - computad
      - coherence (must be a Coherence)
      - substitution (must be a Substitution whose scheme matches the coherence's scheme)

    The correctness checks are similar to InstantiatedCoherence, plus we ensure
    that each substituted cell belongs to the same computad. We do not check
    parallelism here, since the Coherence itself enforces that.
    """
    def __init__(self, computad, coherence, substitution):
        if not isinstance(coherence, Coherence):
            raise ValueError("coherence must be a Coherence.")
        if not isinstance(substitution, Substitution):
            raise ValueError("substitution must be a Substitution.")
        if substitution.scheme is not coherence.scheme:
            raise ValueError("Substitution scheme does not match coherence's scheme.")
        super().__init__(dimension=coherence.scheme.max_dimension - 1)
        self.computad = computad
        self.coherence = coherence
        self.substitution = substitution
        self._check_correctness()

    def _check_correctness(self):
        if not isinstance(self.computad, OmegaComputad):
            raise ValueError("computad must be an OmegaComputad.")
        if not isinstance(self.coherence, Coherence):
            raise ValueError("coherence must be a Coherence.")

        for k, v in self.substitution.full_map.items():
            if not isinstance(v, AtomicCell) or not isinstance(v, OmegaCategoryCell):
                raise ValueError("Substitution value at key {k} not AtomicCell or OmegaCategoryCell.")
            if v.computad != self.computad:
                raise ValueError("Substitution value for key {k} from a different computad.")

    def source(self):
        return self._resolve(self.coherence.source)

    def target(self):
        return self._resolve(self.coherence.target)

    def _resolve(self, obj):
        # If it's a PastingSchemePosition, we require it to be in our map.
        # If it's an InstantiatedCoherence, build a new substitution from its map.
        # and turn it into an OmegaCategoryCell.
        # Otherwise, error.
        if isinstance(obj, PastingSchemePosition):
            if obj not in self.substitution:
                raise ValueError(f"Position {obj} missing from the substitution.")
            val = self.substitution[obj]
            return self._resolve(val)
        elif isinstance(obj, InstantiatedCoherence):
            new_map = {}
            for k, v in obj.substitution.maximal_map().items():
                new_v = self._resolve(v)
                new_map[k] = new_v
            new_sub = Substitution(obj.coherence.scheme, new_map)
            return OmegaCategoryCell(self.computad, obj.coherence, new_sub)
        else:
            raise ValueError("Coherence source/target must be positions or InstantiatedCoherence.")

    def __eq__(self, other):
        if not isinstance(other, OmegaCategoryCell):
            return False
        return (
            self.dimension == other.dimension and
            self.computad is other.computad and
            self.coherence == other.coherence and
            self.substitution == other.substitution
        )