Flat Encoders

FlatRelationEncoder

Bases: EncoderBase[FlatRelationData]

Encode one planning state as packed flat relations.

The output uses one flat entity table and packed relation tensors instead of relation nodes. Optional goals, subgoals, explicit actions, and history payloads can add more relations and helper rows on that same table.

Source code in src/mifrost/encoders/flat.py

class FlatRelationEncoder(EncoderBase[FlatRelationData]):
    """Encode one planning state as packed flat relations.

    The output uses one flat entity table and packed relation tensors instead
    of relation nodes. Optional goals, subgoals, explicit actions, and history
    payloads can add more relations and helper rows on that same table.
    """

    def __init__(
        self,
        domain: DomainInput,
        *,
        max_goal_level: int = 0,
        support_literals: bool = False,
        include_static: bool = True,
        export_node_names: bool = True,
        ignore_zero_arity_relations: bool = True,
        use_predicate_virtual_nodes: bool = False,
        include_lgan_edges: bool = False,
        lgan_anchor_sources: Iterable[TargetSource | str] | None = None,
        target_sources: Iterable[TargetSource | str] | None = None,
        target_symbol_prefix: str = "target:",
        lgan_tn_edge_pos: str = DEFAULT_LGAN_TN_EDGE_POS,
        lgan_nn_edge_pos: str = DEFAULT_LGAN_NN_EDGE_POS,
        lgan_rr_edge_pos: str = DEFAULT_LGAN_RR_EDGE_POS,
        goal_derivations: Iterable[Any] | None = None,
    ) -> None:
        """Build a flat encoder for state-style workloads.

        Parameters follow the flat main state lane.

        `target_sources` answers "what should count as a selectable target?":

        - `action`: explicit grounded actions from `actions=...`
        - `goal`: literals from the root `goals=...` input
        - `subgoal`: literals from `subgoal_layers=...`
        - `history`: literals from `history_subgoals=...`
        - `state`: not used here; state targets belong to `FlatHorizonEncoder`

        `lgan_anchor_sources` is separate. It only creates extra LGAN anchor
        rows for `goal`, `subgoal`, and `history`, without turning them into
        prediction targets. `include_lgan_edges=True` emits the packed LGAN
        edge tensors.

        `state` targets are not supported on this lane. Use
        `FlatHorizonEncoder` or the flat transition encoders for state
        candidates.
        """
        normalized_target_sources = normalize_target_sources(target_sources)
        normalized_lgan_anchor_sources = normalize_target_sources(lgan_anchor_sources)

        def _validate_flat_sources(
            sources: set[TargetSource] | None,
            field_name: str,
        ) -> None:
            if sources is None:
                return
            unsupported_sources = sources.intersection({TargetSource.states})
            if unsupported_sources:
                raise ValueError(
                    "FlatRelationEncoder currently supports "
                    f"{field_name}={{'action', 'goal', 'subgoal', 'history'}} "
                    "only; 'state' is reserved for the upcoming flat "
                    "successor/horizon encoders"
                )

        _validate_flat_sources(normalized_target_sources, "target_sources")
        _validate_flat_sources(normalized_lgan_anchor_sources, "lgan_anchor_sources")
        config_kwargs: dict[str, Any] = {
            "max_goal_level": max_goal_level,
            "support_literals": support_literals,
            "include_static": include_static,
            "export_node_names": export_node_names,
            "ignore_zero_arity_relations": ignore_zero_arity_relations,
            "use_predicate_virtual_nodes": use_predicate_virtual_nodes,
            "include_lgan_edges": include_lgan_edges,
            "target_symbol_prefix": target_symbol_prefix,
            "lgan_tn_edge_pos": lgan_tn_edge_pos,
            "lgan_nn_edge_pos": lgan_nn_edge_pos,
            "lgan_rr_edge_pos": lgan_rr_edge_pos,
        }
        if normalized_lgan_anchor_sources is not None:
            config_kwargs["lgan_anchor_sources"] = normalized_lgan_anchor_sources
        if normalized_target_sources is not None:
            config_kwargs["target_sources"] = normalized_target_sources
        if goal_derivations is not None:
            config_kwargs["goal_derivations"] = goal_derivations
        config = FlatRelationEncoderConfig(**config_kwargs)
        self._engine = FlatRelationEncoderEngine(_advanced_domain(domain), config)
        self._config = config
        self.entity_node_type = "entity"
        self.use_predicate_virtual_nodes = bool(config.use_predicate_virtual_nodes)
        self.include_lgan_edges = bool(config.include_lgan_edges)
        self.target_sources = set(config.target_sources)
        self.lgan_anchor_sources = set(config.lgan_anchor_sources)
        self.target_symbol_prefix = str(config.target_symbol_prefix)
        self.lgan_tn_edge_pos = str(config.lgan_tn_edge_pos)
        self.lgan_nn_edge_pos = str(config.lgan_nn_edge_pos)
        self.lgan_rr_edge_pos = str(config.lgan_rr_edge_pos)
        self._lgan_edge_positions = {
            self.lgan_tn_edge_pos,
            self.lgan_nn_edge_pos,
            self.lgan_rr_edge_pos,
        }

    @property
    def engine(self) -> FlatRelationEncoderEngine:
        """Expose the native flat relation engine."""
        return self._engine

    @property
    def config(self) -> FlatRelationEncoderConfig:
        """Expose the resolved native config."""
        return self._config

    @property
    def relation_dict(self):
        """Expose the relation schema used by the native engine."""
        return self._engine.relation_dict

    @property
    def relation_names(self) -> tuple[str, ...]:
        """Expose the ordered flat relation names declared by the native engine."""
        return tuple(str(name) for name in self._engine.relation_names)

    @property
    def relation_arities(self) -> tuple[int, ...]:
        """Expose the ordered flat relation arities declared by the native engine."""
        return tuple(int(arity) for arity in self._engine.relation_arities)

    @property
    def relation_sources(self) -> tuple[str, ...]:
        """Expose the ordered flat relation source labels declared by the native engine."""
        return tuple(str(source) for source in self._engine.relation_sources)

    @property
    def relation_logical_arities(self) -> tuple[int, ...]:
        """Expose the logical flat relation arities declared by the native engine."""
        return tuple(int(arity) for arity in self._engine.relation_logical_arities)

    @property
    def relation_encoded_arities(self) -> tuple[int, ...]:
        """Expose the encoded flat relation arities declared by the native engine."""
        return tuple(int(arity) for arity in self._engine.relation_encoded_arities)

    @property
    def relation_slot_roles(self) -> tuple[int, ...]:
        """Expose flattened per-relation slot-role ids."""
        return tuple(int(role_id) for role_id in self._engine.relation_slot_roles)

    @property
    def relation_slot_role_offsets(self) -> tuple[int, ...]:
        """Expose offsets into `relation_slot_roles` for each relation."""
        return tuple(int(offset) for offset in self._engine.relation_slot_role_offsets)

    @property
    def slot_role_names(self) -> tuple[str, ...]:
        """Expose the ordered slot-role labels used by flat schema metadata."""
        return tuple(str(name) for name in self._engine.slot_role_names)

    def _encode_one_into_builder(
        self,
        state: StateInput,
        builder: BatchBuilder,
        *,
        goals: GoalBatchInput = None,
        actions: ActionBatchInput = None,
        subgoal_layers: SubgoalLayersInput = None,
        history_subgoals: HistorySubgoalInput | None = None,
        history_max_steps: int | None = None,
    ) -> None:
        """Append one flat encoding step into a caller-owned builder."""
        adv_state = _advanced_state(state)
        action_inputs = parse_flat_actions(actions)
        action_list = _prepare_actions(action_inputs)
        history_list = _prepare_history_subgoals(history_subgoals)
        subgoal_layers_list = None if subgoal_layers is None else list(subgoal_layers)
        _validate_subgoal_layers_state_payload(
            subgoal_layers_list,
            state_index=0,
            max_goal_level=int(self._config.max_goal_level),
        )
        if goals is None and subgoal_layers is None and not history_list:
            if action_list:
                self._engine.encode(adv_state, action_list, builder)
                return
            self._engine.encode(adv_state, builder)
            return

        goals_input = goals if goals is not None else default_goals_from_state(state)
        split_goals = _split_goals(goals_input, subgoal_layers_list)
        if history_list:
            self._engine.encode(
                adv_state,
                split_goals,
                action_list,
                history_list,
                history_max_steps,
                builder,
            )
            return
        if action_list:
            self._engine.encode(adv_state, split_goals, action_list, builder)
            return
        self._engine.encode(adv_state, split_goals, builder)

    def _accepted_kwargs(self) -> set[str]:
        return {"history_subgoals", "history_max_steps"}

    def _encode(
        self,
        state: StateInput,
        *,
        goals: GoalBatchInput = None,
        actions: ActionBatchInput = None,
        subgoal_layers: SubgoalLayersInput = None,
        history_subgoals: HistorySubgoalInput | None = None,
        history_max_steps: int | None = None,
    ) -> FlatEncoding:
        builder = BatchBuilder()
        builder.set_graph_kind("flat")
        self._encode_one_into_builder(
            state,
            builder,
            goals=goals,
            actions=actions,
            subgoal_layers=subgoal_layers,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
        )
        builder.next_graph()
        return builder.build()

    def encode(
        self,
        state: StateInput,
        *,
        goals: GoalBatchInput = None,
        actions: ActionBatchInput = None,
        subgoal_layers: SubgoalLayersInput = None,
        history_subgoals: HistorySubgoalInput | None = None,
        history_max_steps: int | None = None,
        include_metadata: bool = True,
        **kwargs,
    ) -> FlatEncoding:
        """Encode one state into the native flat carrier.

        If `goals` are omitted, the problem goals from `state` are used when
        needed. `actions`, `subgoal_layers`, and `history_subgoals` are all
        optional. Use `encode_pyg()` when you want a `FlatRelationData`
        wrapper directly.
        """
        return super().encode(
            state,
            goals=goals,
            actions=actions,
            subgoal_layers=subgoal_layers,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
            include_metadata=include_metadata,
            **kwargs,
        )

    def _encode_batch(
        self,
        states: StateBatchInput,
        *,
        goals: GoalBatchParam = None,
        actions: ActionBatchParam = None,
        subgoal_layers: SubgoalLayersBatchParam = None,
        history_subgoals: HistorySubgoalsBatchParam = None,
        history_max_steps: int | None = None,
    ) -> FlatEncoding:
        states_for_core = _convert_batch_payload(
            states,
            is_leaf=is_state_input,
            convert_leaf=to_advanced_state,
        )
        goals_for_core = _convert_batch_payload(
            goals,
            is_leaf=is_goal_literal_input,
            convert_leaf=to_advanced_literal,
        )
        actions_for_core = _convert_batch_payload(
            actions,
            is_leaf=is_action_input,
            convert_leaf=to_advanced_action,
        )
        subgoal_layers_for_core = _convert_batch_payload(
            subgoal_layers,
            is_leaf=is_goal_literal_input,
            convert_leaf=to_advanced_literal,
        )
        history_subgoals_for_core = _convert_batch_payload(
            history_subgoals,
            is_leaf=is_goal_literal_input,
            convert_leaf=to_advanced_literal,
        )
        if _is_sequence_like(states_for_core):
            state_count = len(states_for_core)
        else:
            state_count = 1
        _validate_subgoal_layers_batch_payload(
            subgoal_layers_for_core,
            state_count=state_count,
            max_goal_level=int(self._config.max_goal_level),
        )
        return self._engine.encode_batch(
            states_for_core,
            goals=goals_for_core,
            actions=actions_for_core,
            subgoal_layers=subgoal_layers_for_core,
            history_subgoals=history_subgoals_for_core,
            history_max_steps=history_max_steps,
        )

    def encode_batch(
        self,
        states: StateBatchInput,
        *,
        goals: GoalBatchParam = None,
        actions: ActionBatchParam = None,
        subgoal_layers: SubgoalLayersBatchParam = None,
        history_subgoals: HistorySubgoalsBatchParam = None,
        history_max_steps: int | None = None,
        batch_attrs: Mapping[str, Any] | None = None,
        collate_spec: CollateSpecParam = None,
        include_metadata: bool = True,
        **kwargs,
    ) -> FlatEncoding:
        """Encode many states with shared or per-state optional payloads.

        Batch kwargs follow the same rules as `encode`: each optional payload
        may be shared for all states or given separately per state.
        """
        return super().encode_batch(
            states,
            goals=goals,
            actions=actions,
            subgoal_layers=subgoal_layers,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
            batch_attrs=batch_attrs,
            collate_spec=collate_spec,
            include_metadata=include_metadata,
            **kwargs,
        )

    def stream(self) -> FlatRelationEncoderStream:
        """Return an append-only stream for flat state encodings."""
        return FlatRelationEncoderStream(self)

    def mutable_stream(self) -> FlatRelationMutableEncoderStream:
        """Return a mutable stream with `append`, `update`, and `remove`."""
        return FlatRelationMutableEncoderStream(self)

    def to_networkx(
        self,
        data: FlatRelationData,
        *,
        graph_index: int = 0,
        mode: str = "star",
    ) -> nx.MultiDiGraph:
        """Build a debug graph view for one encoded flat graph.

        The returned graph is only for inspection. It expands each relation
        instance into a synthetic node and can overlay LGAN edges when they are
        present in `data`.
        """
        if mode != "star":
            raise ValueError(f"Unsupported flat visualization mode: {mode!r}")

        graph = nx.MultiDiGraph()
        lgan_tn_edge_pos = str(getattr(data, "lgan_tn_edge_pos", self.lgan_tn_edge_pos))
        lgan_nn_edge_pos = str(getattr(data, "lgan_nn_edge_pos", self.lgan_nn_edge_pos))
        lgan_rr_edge_pos = str(getattr(data, "lgan_rr_edge_pos", self.lgan_rr_edge_pos))
        start, end = data.graph_node_range(graph_index)
        node_names = data.graph_node_names(graph_index)
        name_by_global = {start + idx: node_names[idx] for idx in range(end - start)}
        history_entity_dt_by_index = {
            int(global_idx): int(dt)
            for global_idx, dt in zip(
                data.graph_history_entity_indices(graph_index).tolist(),
                data.graph_history_entity_dt(graph_index).tolist(),
                strict=True,
            )
        }
        target_entity_indices = data.graph_target_entity_indices(graph_index)
        target_entity_group_ids = data.graph_target_entity_group_ids(graph_index)
        target_entity_groups = list(getattr(data, "target_entity_groups", ()))
        target_entity_group_by_index: dict[int, tuple[int | None, str | None]] = {}
        for global_idx, group_id in zip(
            target_entity_indices.tolist(),
            target_entity_group_ids.tolist(),
            strict=True,
        ):
            group_name = None
            if 0 <= group_id < len(target_entity_groups):
                group_name = str(target_entity_groups[group_id])
            target_entity_group_by_index[int(global_idx)] = (int(group_id), group_name)
        target_groups = list(getattr(data, "target_groups", ()))
        target_positions = data.graph_target_positions(graph_index).tolist()
        target_indices = data.graph_target_indices(graph_index).tolist()
        target_candidate_ids = data.graph_target_candidate_ids(graph_index).tolist()
        target_group_ids = data.graph_target_group_ids(graph_index).tolist()
        target_names = data.graph_target_names(graph_index)
        target_depths_tensor = getattr(data, "target_depths", None)
        target_depths = (
            data.graph_target_depths(graph_index).tolist()
            if target_depths_tensor is not None
            else [None] * len(target_positions)
        )
        target_rows_by_position: dict[int, list[dict[str, Any]]] = {}
        for (
            position,
            target_index,
            candidate_id,
            group_id,
            target_name,
            target_depth,
        ) in zip(
            target_positions,
            target_indices,
            target_candidate_ids,
            target_group_ids,
            target_names,
            target_depths,
            strict=True,
        ):
            group_name = None
            if 0 <= group_id < len(target_groups):
                group_name = str(target_groups[group_id])
            target_rows_by_position.setdefault(int(position), []).append(
                {
                    "target_index": int(target_index),
                    "target_candidate_id": int(candidate_id),
                    "target_group_id": int(group_id),
                    "target_group": group_name,
                    "target_name": str(target_name),
                    "target_depth": (
                        None if target_depth is None else int(target_depth)
                    ),
                }
            )
        entity_roles = data.graph_entity_roles(graph_index)
        entity_role_by_index = {
            start + idx: entity_roles[idx]
            for idx in range(min(len(entity_roles), end - start))
        }

        for global_idx in range(start, end):
            label = name_by_global.get(global_idx, f"entity:{global_idx}")
            target_group_id, target_group_name = target_entity_group_by_index.get(
                global_idx, (None, None)
            )
            history_dt = history_entity_dt_by_index.get(global_idx)
            target_rows = target_rows_by_position.get(global_idx, [])
            if target_group_name is not None:
                entity_kind = "target_entity"
            elif history_dt is not None:
                entity_kind = "history_entity"
            else:
                entity_kind = "object"
            entity_role = entity_role_by_index.get(global_idx)
            graph.add_node(
                label,
                type="entity",
                kind="entity",
                entity_kind=entity_kind,
                entity_role=entity_role,
                history_dt=history_dt,
                target_group_id=target_group_id,
                target_group=target_group_name,
                target_rows=target_rows,
                global_index=global_idx,
                graph_index=graph_index,
            )
            if len(target_rows) == 1:
                graph.nodes[label].update(target_rows[0])

        flattened = data.flattened_relations_view(graph_index=graph_index)
        schema = data.schema
        relation_cursor, _ = data.graph_relation_instance_range(graph_index)
        relation_node_by_global: dict[int, str] = {}
        for relation_idx, relation_name in enumerate(schema.names):
            instances = flattened[relation_name]
            source = (
                schema.sources[relation_idx]
                if relation_idx < len(schema.sources)
                else "relation"
            )
            for instance_idx, args in enumerate(instances.tolist()):
                relation_node = f"{relation_name}#{instance_idx}"
                slot_roles = (
                    schema.slot_roles[relation_idx]
                    if relation_idx < len(schema.slot_roles)
                    else ()
                )
                graph.add_node(
                    relation_node,
                    type=relation_name,
                    kind="relation",
                    source=source,
                    relation_name=relation_name,
                    relation_id=relation_idx,
                    instance_index=instance_idx,
                    global_relation_index=relation_cursor,
                )
                relation_node_by_global[relation_cursor] = relation_node
                for slot, global_idx in enumerate(args):
                    entity_node = name_by_global.get(global_idx, f"entity:{global_idx}")
                    slot_role = slot_roles[slot] if slot < len(slot_roles) else None
                    graph.add_edge(
                        relation_node,
                        entity_node,
                        position=str(slot),
                        slot=slot,
                        slot_role=slot_role,
                    )
                relation_cursor += 1

        lgan_specs = (
            ("tn", data.graph_lgan_tn_edges(graph_index), lgan_tn_edge_pos),
            ("nn", data.graph_lgan_nn_edges(graph_index), lgan_nn_edge_pos),
        )
        for lgan_kind, edges, edge_pos in lgan_specs:
            if edges.numel() == 0:
                continue
            for relation_index, entity_index in edges.t().tolist():
                relation_node = relation_node_by_global.get(int(relation_index))
                if relation_node is None:
                    continue
                entity_node = name_by_global.get(
                    int(entity_index), f"entity:{int(entity_index)}"
                )
                graph.add_edge(
                    relation_node,
                    entity_node,
                    position=edge_pos,
                    lgan_kind=lgan_kind,
                    style="dashed",
                )

        rr_edges = data.graph_lgan_rr_edges(graph_index)
        if rr_edges.numel() > 0:
            for src_relation_index, dst_relation_index in rr_edges.t().tolist():
                src_relation_node = relation_node_by_global.get(int(src_relation_index))
                dst_relation_node = relation_node_by_global.get(int(dst_relation_index))
                if src_relation_node is None or dst_relation_node is None:
                    continue
                graph.add_edge(
                    src_relation_node,
                    dst_relation_node,
                    position=lgan_rr_edge_pos,
                    lgan_kind="rr",
                    style="dashed",
                )
        return graph

    def draw(
        self,
        data: FlatRelationData | nx.Graph,
        *,
        graph_index: int = 0,
        ax=None,
        with_labels: bool = True,
        edge_labels: bool = True,
        layout: dict | None = None,
    ):
        """Draw a flat debug graph with matplotlib.

        Pass either `FlatRelationData` or a graph produced by
        :meth:`to_networkx`. LGAN edges are shown as dashed overlays when they
        exist.
        """
        try:
            import matplotlib.pyplot as plt
        except ModuleNotFoundError as exc:
            raise RuntimeError(
                "FlatRelationEncoder.draw requires matplotlib to be installed"
            ) from exc

        graph = (
            data
            if isinstance(data, nx.Graph)
            else self.to_networkx(data, graph_index=graph_index)
        )
        if ax is None:
            _, ax = plt.subplots()

        pos = layout or _default_flat_debug_layout(graph)
        entity_nodes = [
            node for node, attrs in graph.nodes(data=True) if attrs["kind"] == "entity"
        ]
        relation_nodes = [
            node
            for node, attrs in graph.nodes(data=True)
            if attrs["kind"] == "relation"
        ]

        if entity_nodes:
            entity_groups = [
                graph.nodes[node].get("target_group")
                or graph.nodes[node].get("entity_kind", "object")
                for node in entity_nodes
            ]
            entity_palette = {
                "object": "#f3efe0",
                "state": "#d6f0c0",
                "goal": "#f9d7dd",
                "subgoal": "#dbecc8",
                "action": "#d8ecff",
                "history": "#e6d7ff",
                "history_entity": "#e8def8",
                "target_entity": "#d8ecff",
            }
            nx.draw_networkx_nodes(
                graph,
                pos,
                nodelist=entity_nodes,
                node_color=[
                    entity_palette.get(group, "#d8ecff") for group in entity_groups
                ],
                edgecolors="#222222",
                linewidths=1.3,
                node_size=420,
                ax=ax,
            )
        if relation_nodes:
            nx.draw_networkx_nodes(
                graph,
                pos,
                nodelist=relation_nodes,
                node_color=[
                    _relation_name_color(
                        str(graph.nodes[node].get("relation_name", node))
                    )
                    for node in relation_nodes
                ],
                node_shape="s",
                edgecolors="#111111",
                linewidths=1.1,
                node_size=520,
                ax=ax,
            )

        normal_edges = []
        lgan_edges = []
        lgan_colors = []
        lgan_palette = {
            "tn": "#d35454",
            "nn": "#4f7cac",
            "rr": "#5a9367",
        }
        for u, v, attrs in graph.edges(data=True):
            lgan_kind = attrs.get("lgan_kind")
            if lgan_kind is None:
                normal_edges.append((u, v))
                continue
            lgan_edges.append((u, v))
            lgan_colors.append(lgan_palette.get(str(lgan_kind), "#666666"))

        if normal_edges:
            nx.draw_networkx_edges(
                graph,
                pos,
                edgelist=normal_edges,
                ax=ax,
                arrows=True,
                width=1.2,
                alpha=0.8,
            )
        if lgan_edges:
            nx.draw_networkx_edges(
                graph,
                pos,
                edgelist=lgan_edges,
                ax=ax,
                arrows=True,
                width=1.4,
                alpha=0.85,
                style="dashed",
                edge_color=lgan_colors,
            )

        if with_labels:
            nx.draw_networkx_labels(graph, pos, ax=ax, font_size=8)

        if edge_labels:
            if graph.is_multigraph():
                labels = {
                    (u, v, key): attrs.get("position")
                    for u, v, key, attrs in graph.edges(keys=True, data=True)
                    if attrs.get("position") is not None
                }
            else:
                labels = {
                    (u, v): attrs.get("position")
                    for u, v, attrs in graph.edges(data=True)
                    if attrs.get("position") is not None
                }
            if labels:
                nx.draw_networkx_edge_labels(
                    graph,
                    pos,
                    edge_labels=labels,
                    ax=ax,
                    font_size=7,
                )

        ax.set_axis_off()
        return ax

engine property

Expose the native flat relation engine.

config property

Expose the resolved native config.

relation_dict property

Expose the relation schema used by the native engine.

relation_names property

Expose the ordered flat relation names declared by the native engine.

relation_arities property

Expose the ordered flat relation arities declared by the native engine.

relation_sources property

Expose the ordered flat relation source labels declared by the native engine.

relation_logical_arities property

Expose the logical flat relation arities declared by the native engine.

relation_encoded_arities property

Expose the encoded flat relation arities declared by the native engine.

relation_slot_roles property

Expose flattened per-relation slot-role ids.

relation_slot_role_offsets property

Expose offsets into relation_slot_roles for each relation.

slot_role_names property

Expose the ordered slot-role labels used by flat schema metadata.

__init__(domain, *, max_goal_level=0, support_literals=False, include_static=True, export_node_names=True, ignore_zero_arity_relations=True, use_predicate_virtual_nodes=False, include_lgan_edges=False, lgan_anchor_sources=None, target_sources=None, target_symbol_prefix='target:', lgan_tn_edge_pos=DEFAULT_LGAN_TN_EDGE_POS, lgan_nn_edge_pos=DEFAULT_LGAN_NN_EDGE_POS, lgan_rr_edge_pos=DEFAULT_LGAN_RR_EDGE_POS, goal_derivations=None)

Build a flat encoder for state-style workloads.

Parameters follow the flat main state lane.

target_sources answers "what should count as a selectable target?":

• action: explicit grounded actions from actions=...
• goal: literals from the root goals=... input
• subgoal: literals from subgoal_layers=...
• history: literals from history_subgoals=...
• state: not used here; state targets belong to FlatHorizonEncoder

lgan_anchor_sources is separate. It only creates extra LGAN anchor rows for goal, subgoal, and history, without turning them into prediction targets. include_lgan_edges=True emits the packed LGAN edge tensors.

state targets are not supported on this lane. Use FlatHorizonEncoder or the flat transition encoders for state candidates.

Source code in src/mifrost/encoders/flat.py

def __init__(
    self,
    domain: DomainInput,
    *,
    max_goal_level: int = 0,
    support_literals: bool = False,
    include_static: bool = True,
    export_node_names: bool = True,
    ignore_zero_arity_relations: bool = True,
    use_predicate_virtual_nodes: bool = False,
    include_lgan_edges: bool = False,
    lgan_anchor_sources: Iterable[TargetSource | str] | None = None,
    target_sources: Iterable[TargetSource | str] | None = None,
    target_symbol_prefix: str = "target:",
    lgan_tn_edge_pos: str = DEFAULT_LGAN_TN_EDGE_POS,
    lgan_nn_edge_pos: str = DEFAULT_LGAN_NN_EDGE_POS,
    lgan_rr_edge_pos: str = DEFAULT_LGAN_RR_EDGE_POS,
    goal_derivations: Iterable[Any] | None = None,
) -> None:
    """Build a flat encoder for state-style workloads.

    Parameters follow the flat main state lane.

    `target_sources` answers "what should count as a selectable target?":

    - `action`: explicit grounded actions from `actions=...`
    - `goal`: literals from the root `goals=...` input
    - `subgoal`: literals from `subgoal_layers=...`
    - `history`: literals from `history_subgoals=...`
    - `state`: not used here; state targets belong to `FlatHorizonEncoder`

    `lgan_anchor_sources` is separate. It only creates extra LGAN anchor
    rows for `goal`, `subgoal`, and `history`, without turning them into
    prediction targets. `include_lgan_edges=True` emits the packed LGAN
    edge tensors.

    `state` targets are not supported on this lane. Use
    `FlatHorizonEncoder` or the flat transition encoders for state
    candidates.
    """
    normalized_target_sources = normalize_target_sources(target_sources)
    normalized_lgan_anchor_sources = normalize_target_sources(lgan_anchor_sources)

    def _validate_flat_sources(
        sources: set[TargetSource] | None,
        field_name: str,
    ) -> None:
        if sources is None:
            return
        unsupported_sources = sources.intersection({TargetSource.states})
        if unsupported_sources:
            raise ValueError(
                "FlatRelationEncoder currently supports "
                f"{field_name}={{'action', 'goal', 'subgoal', 'history'}} "
                "only; 'state' is reserved for the upcoming flat "
                "successor/horizon encoders"
            )

    _validate_flat_sources(normalized_target_sources, "target_sources")
    _validate_flat_sources(normalized_lgan_anchor_sources, "lgan_anchor_sources")
    config_kwargs: dict[str, Any] = {
        "max_goal_level": max_goal_level,
        "support_literals": support_literals,
        "include_static": include_static,
        "export_node_names": export_node_names,
        "ignore_zero_arity_relations": ignore_zero_arity_relations,
        "use_predicate_virtual_nodes": use_predicate_virtual_nodes,
        "include_lgan_edges": include_lgan_edges,
        "target_symbol_prefix": target_symbol_prefix,
        "lgan_tn_edge_pos": lgan_tn_edge_pos,
        "lgan_nn_edge_pos": lgan_nn_edge_pos,
        "lgan_rr_edge_pos": lgan_rr_edge_pos,
    }
    if normalized_lgan_anchor_sources is not None:
        config_kwargs["lgan_anchor_sources"] = normalized_lgan_anchor_sources
    if normalized_target_sources is not None:
        config_kwargs["target_sources"] = normalized_target_sources
    if goal_derivations is not None:
        config_kwargs["goal_derivations"] = goal_derivations
    config = FlatRelationEncoderConfig(**config_kwargs)
    self._engine = FlatRelationEncoderEngine(_advanced_domain(domain), config)
    self._config = config
    self.entity_node_type = "entity"
    self.use_predicate_virtual_nodes = bool(config.use_predicate_virtual_nodes)
    self.include_lgan_edges = bool(config.include_lgan_edges)
    self.target_sources = set(config.target_sources)
    self.lgan_anchor_sources = set(config.lgan_anchor_sources)
    self.target_symbol_prefix = str(config.target_symbol_prefix)
    self.lgan_tn_edge_pos = str(config.lgan_tn_edge_pos)
    self.lgan_nn_edge_pos = str(config.lgan_nn_edge_pos)
    self.lgan_rr_edge_pos = str(config.lgan_rr_edge_pos)
    self._lgan_edge_positions = {
        self.lgan_tn_edge_pos,
        self.lgan_nn_edge_pos,
        self.lgan_rr_edge_pos,
    }

encode(state, *, goals=None, actions=None, subgoal_layers=None, history_subgoals=None, history_max_steps=None, include_metadata=True, **kwargs)

Encode one state into the native flat carrier.

If goals are omitted, the problem goals from state are used when needed. actions, subgoal_layers, and history_subgoals are all optional. Use encode_pyg() when you want a FlatRelationData wrapper directly.

Source code in src/mifrost/encoders/flat.py

def encode(
    self,
    state: StateInput,
    *,
    goals: GoalBatchInput = None,
    actions: ActionBatchInput = None,
    subgoal_layers: SubgoalLayersInput = None,
    history_subgoals: HistorySubgoalInput | None = None,
    history_max_steps: int | None = None,
    include_metadata: bool = True,
    **kwargs,
) -> FlatEncoding:
    """Encode one state into the native flat carrier.

    If `goals` are omitted, the problem goals from `state` are used when
    needed. `actions`, `subgoal_layers`, and `history_subgoals` are all
    optional. Use `encode_pyg()` when you want a `FlatRelationData`
    wrapper directly.
    """
    return super().encode(
        state,
        goals=goals,
        actions=actions,
        subgoal_layers=subgoal_layers,
        history_subgoals=history_subgoals,
        history_max_steps=history_max_steps,
        include_metadata=include_metadata,
        **kwargs,
    )

encode_batch(states, *, goals=None, actions=None, subgoal_layers=None, history_subgoals=None, history_max_steps=None, batch_attrs=None, collate_spec=None, include_metadata=True, **kwargs)

Encode many states with shared or per-state optional payloads.

Batch kwargs follow the same rules as encode: each optional payload may be shared for all states or given separately per state.

Source code in src/mifrost/encoders/flat.py

def encode_batch(
    self,
    states: StateBatchInput,
    *,
    goals: GoalBatchParam = None,
    actions: ActionBatchParam = None,
    subgoal_layers: SubgoalLayersBatchParam = None,
    history_subgoals: HistorySubgoalsBatchParam = None,
    history_max_steps: int | None = None,
    batch_attrs: Mapping[str, Any] | None = None,
    collate_spec: CollateSpecParam = None,
    include_metadata: bool = True,
    **kwargs,
) -> FlatEncoding:
    """Encode many states with shared or per-state optional payloads.

    Batch kwargs follow the same rules as `encode`: each optional payload
    may be shared for all states or given separately per state.
    """
    return super().encode_batch(
        states,
        goals=goals,
        actions=actions,
        subgoal_layers=subgoal_layers,
        history_subgoals=history_subgoals,
        history_max_steps=history_max_steps,
        batch_attrs=batch_attrs,
        collate_spec=collate_spec,
        include_metadata=include_metadata,
        **kwargs,
    )

stream()

Return an append-only stream for flat state encodings.

Source code in src/mifrost/encoders/flat.py

def stream(self) -> FlatRelationEncoderStream:
    """Return an append-only stream for flat state encodings."""
    return FlatRelationEncoderStream(self)

mutable_stream()

Return a mutable stream with append, update, and remove.

Source code in src/mifrost/encoders/flat.py

def mutable_stream(self) -> FlatRelationMutableEncoderStream:
    """Return a mutable stream with `append`, `update`, and `remove`."""
    return FlatRelationMutableEncoderStream(self)

to_networkx(data, *, graph_index=0, mode='star')

Build a debug graph view for one encoded flat graph.

The returned graph is only for inspection. It expands each relation instance into a synthetic node and can overlay LGAN edges when they are present in data.

Source code in src/mifrost/encoders/flat.py

def to_networkx(
    self,
    data: FlatRelationData,
    *,
    graph_index: int = 0,
    mode: str = "star",
) -> nx.MultiDiGraph:
    """Build a debug graph view for one encoded flat graph.

    The returned graph is only for inspection. It expands each relation
    instance into a synthetic node and can overlay LGAN edges when they are
    present in `data`.
    """
    if mode != "star":
        raise ValueError(f"Unsupported flat visualization mode: {mode!r}")

    graph = nx.MultiDiGraph()
    lgan_tn_edge_pos = str(getattr(data, "lgan_tn_edge_pos", self.lgan_tn_edge_pos))
    lgan_nn_edge_pos = str(getattr(data, "lgan_nn_edge_pos", self.lgan_nn_edge_pos))
    lgan_rr_edge_pos = str(getattr(data, "lgan_rr_edge_pos", self.lgan_rr_edge_pos))
    start, end = data.graph_node_range(graph_index)
    node_names = data.graph_node_names(graph_index)
    name_by_global = {start + idx: node_names[idx] for idx in range(end - start)}
    history_entity_dt_by_index = {
        int(global_idx): int(dt)
        for global_idx, dt in zip(
            data.graph_history_entity_indices(graph_index).tolist(),
            data.graph_history_entity_dt(graph_index).tolist(),
            strict=True,
        )
    }
    target_entity_indices = data.graph_target_entity_indices(graph_index)
    target_entity_group_ids = data.graph_target_entity_group_ids(graph_index)
    target_entity_groups = list(getattr(data, "target_entity_groups", ()))
    target_entity_group_by_index: dict[int, tuple[int | None, str | None]] = {}
    for global_idx, group_id in zip(
        target_entity_indices.tolist(),
        target_entity_group_ids.tolist(),
        strict=True,
    ):
        group_name = None
        if 0 <= group_id < len(target_entity_groups):
            group_name = str(target_entity_groups[group_id])
        target_entity_group_by_index[int(global_idx)] = (int(group_id), group_name)
    target_groups = list(getattr(data, "target_groups", ()))
    target_positions = data.graph_target_positions(graph_index).tolist()
    target_indices = data.graph_target_indices(graph_index).tolist()
    target_candidate_ids = data.graph_target_candidate_ids(graph_index).tolist()
    target_group_ids = data.graph_target_group_ids(graph_index).tolist()
    target_names = data.graph_target_names(graph_index)
    target_depths_tensor = getattr(data, "target_depths", None)
    target_depths = (
        data.graph_target_depths(graph_index).tolist()
        if target_depths_tensor is not None
        else [None] * len(target_positions)
    )
    target_rows_by_position: dict[int, list[dict[str, Any]]] = {}
    for (
        position,
        target_index,
        candidate_id,
        group_id,
        target_name,
        target_depth,
    ) in zip(
        target_positions,
        target_indices,
        target_candidate_ids,
        target_group_ids,
        target_names,
        target_depths,
        strict=True,
    ):
        group_name = None
        if 0 <= group_id < len(target_groups):
            group_name = str(target_groups[group_id])
        target_rows_by_position.setdefault(int(position), []).append(
            {
                "target_index": int(target_index),
                "target_candidate_id": int(candidate_id),
                "target_group_id": int(group_id),
                "target_group": group_name,
                "target_name": str(target_name),
                "target_depth": (
                    None if target_depth is None else int(target_depth)
                ),
            }
        )
    entity_roles = data.graph_entity_roles(graph_index)
    entity_role_by_index = {
        start + idx: entity_roles[idx]
        for idx in range(min(len(entity_roles), end - start))
    }

    for global_idx in range(start, end):
        label = name_by_global.get(global_idx, f"entity:{global_idx}")
        target_group_id, target_group_name = target_entity_group_by_index.get(
            global_idx, (None, None)
        )
        history_dt = history_entity_dt_by_index.get(global_idx)
        target_rows = target_rows_by_position.get(global_idx, [])
        if target_group_name is not None:
            entity_kind = "target_entity"
        elif history_dt is not None:
            entity_kind = "history_entity"
        else:
            entity_kind = "object"
        entity_role = entity_role_by_index.get(global_idx)
        graph.add_node(
            label,
            type="entity",
            kind="entity",
            entity_kind=entity_kind,
            entity_role=entity_role,
            history_dt=history_dt,
            target_group_id=target_group_id,
            target_group=target_group_name,
            target_rows=target_rows,
            global_index=global_idx,
            graph_index=graph_index,
        )
        if len(target_rows) == 1:
            graph.nodes[label].update(target_rows[0])

    flattened = data.flattened_relations_view(graph_index=graph_index)
    schema = data.schema
    relation_cursor, _ = data.graph_relation_instance_range(graph_index)
    relation_node_by_global: dict[int, str] = {}
    for relation_idx, relation_name in enumerate(schema.names):
        instances = flattened[relation_name]
        source = (
            schema.sources[relation_idx]
            if relation_idx < len(schema.sources)
            else "relation"
        )
        for instance_idx, args in enumerate(instances.tolist()):
            relation_node = f"{relation_name}#{instance_idx}"
            slot_roles = (
                schema.slot_roles[relation_idx]
                if relation_idx < len(schema.slot_roles)
                else ()
            )
            graph.add_node(
                relation_node,
                type=relation_name,
                kind="relation",
                source=source,
                relation_name=relation_name,
                relation_id=relation_idx,
                instance_index=instance_idx,
                global_relation_index=relation_cursor,
            )
            relation_node_by_global[relation_cursor] = relation_node
            for slot, global_idx in enumerate(args):
                entity_node = name_by_global.get(global_idx, f"entity:{global_idx}")
                slot_role = slot_roles[slot] if slot < len(slot_roles) else None
                graph.add_edge(
                    relation_node,
                    entity_node,
                    position=str(slot),
                    slot=slot,
                    slot_role=slot_role,
                )
            relation_cursor += 1

    lgan_specs = (
        ("tn", data.graph_lgan_tn_edges(graph_index), lgan_tn_edge_pos),
        ("nn", data.graph_lgan_nn_edges(graph_index), lgan_nn_edge_pos),
    )
    for lgan_kind, edges, edge_pos in lgan_specs:
        if edges.numel() == 0:
            continue
        for relation_index, entity_index in edges.t().tolist():
            relation_node = relation_node_by_global.get(int(relation_index))
            if relation_node is None:
                continue
            entity_node = name_by_global.get(
                int(entity_index), f"entity:{int(entity_index)}"
            )
            graph.add_edge(
                relation_node,
                entity_node,
                position=edge_pos,
                lgan_kind=lgan_kind,
                style="dashed",
            )

    rr_edges = data.graph_lgan_rr_edges(graph_index)
    if rr_edges.numel() > 0:
        for src_relation_index, dst_relation_index in rr_edges.t().tolist():
            src_relation_node = relation_node_by_global.get(int(src_relation_index))
            dst_relation_node = relation_node_by_global.get(int(dst_relation_index))
            if src_relation_node is None or dst_relation_node is None:
                continue
            graph.add_edge(
                src_relation_node,
                dst_relation_node,
                position=lgan_rr_edge_pos,
                lgan_kind="rr",
                style="dashed",
            )
    return graph

draw(data, *, graph_index=0, ax=None, with_labels=True, edge_labels=True, layout=None)

Draw a flat debug graph with matplotlib.

Pass either FlatRelationData or a graph produced by :meth:to_networkx. LGAN edges are shown as dashed overlays when they exist.

Source code in src/mifrost/encoders/flat.py

def draw(
    self,
    data: FlatRelationData | nx.Graph,
    *,
    graph_index: int = 0,
    ax=None,
    with_labels: bool = True,
    edge_labels: bool = True,
    layout: dict | None = None,
):
    """Draw a flat debug graph with matplotlib.

    Pass either `FlatRelationData` or a graph produced by
    :meth:`to_networkx`. LGAN edges are shown as dashed overlays when they
    exist.
    """
    try:
        import matplotlib.pyplot as plt
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "FlatRelationEncoder.draw requires matplotlib to be installed"
        ) from exc

    graph = (
        data
        if isinstance(data, nx.Graph)
        else self.to_networkx(data, graph_index=graph_index)
    )
    if ax is None:
        _, ax = plt.subplots()

    pos = layout or _default_flat_debug_layout(graph)
    entity_nodes = [
        node for node, attrs in graph.nodes(data=True) if attrs["kind"] == "entity"
    ]
    relation_nodes = [
        node
        for node, attrs in graph.nodes(data=True)
        if attrs["kind"] == "relation"
    ]

    if entity_nodes:
        entity_groups = [
            graph.nodes[node].get("target_group")
            or graph.nodes[node].get("entity_kind", "object")
            for node in entity_nodes
        ]
        entity_palette = {
            "object": "#f3efe0",
            "state": "#d6f0c0",
            "goal": "#f9d7dd",
            "subgoal": "#dbecc8",
            "action": "#d8ecff",
            "history": "#e6d7ff",
            "history_entity": "#e8def8",
            "target_entity": "#d8ecff",
        }
        nx.draw_networkx_nodes(
            graph,
            pos,
            nodelist=entity_nodes,
            node_color=[
                entity_palette.get(group, "#d8ecff") for group in entity_groups
            ],
            edgecolors="#222222",
            linewidths=1.3,
            node_size=420,
            ax=ax,
        )
    if relation_nodes:
        nx.draw_networkx_nodes(
            graph,
            pos,
            nodelist=relation_nodes,
            node_color=[
                _relation_name_color(
                    str(graph.nodes[node].get("relation_name", node))
                )
                for node in relation_nodes
            ],
            node_shape="s",
            edgecolors="#111111",
            linewidths=1.1,
            node_size=520,
            ax=ax,
        )

    normal_edges = []
    lgan_edges = []
    lgan_colors = []
    lgan_palette = {
        "tn": "#d35454",
        "nn": "#4f7cac",
        "rr": "#5a9367",
    }
    for u, v, attrs in graph.edges(data=True):
        lgan_kind = attrs.get("lgan_kind")
        if lgan_kind is None:
            normal_edges.append((u, v))
            continue
        lgan_edges.append((u, v))
        lgan_colors.append(lgan_palette.get(str(lgan_kind), "#666666"))

    if normal_edges:
        nx.draw_networkx_edges(
            graph,
            pos,
            edgelist=normal_edges,
            ax=ax,
            arrows=True,
            width=1.2,
            alpha=0.8,
        )
    if lgan_edges:
        nx.draw_networkx_edges(
            graph,
            pos,
            edgelist=lgan_edges,
            ax=ax,
            arrows=True,
            width=1.4,
            alpha=0.85,
            style="dashed",
            edge_color=lgan_colors,
        )

    if with_labels:
        nx.draw_networkx_labels(graph, pos, ax=ax, font_size=8)

    if edge_labels:
        if graph.is_multigraph():
            labels = {
                (u, v, key): attrs.get("position")
                for u, v, key, attrs in graph.edges(keys=True, data=True)
                if attrs.get("position") is not None
            }
        else:
            labels = {
                (u, v): attrs.get("position")
                for u, v, attrs in graph.edges(data=True)
                if attrs.get("position") is not None
            }
        if labels:
            nx.draw_networkx_edge_labels(
                graph,
                pos,
                edge_labels=labels,
                ax=ax,
                font_size=7,
            )

    ax.set_axis_off()
    return ax

FlatHorizonEncoder

Bases: FlatRelationEncoder

Encode a root state plus lookahead candidates as packed flat relations.

Source code in src/mifrost/encoders/flat_horizon.py

class FlatHorizonEncoder(FlatRelationEncoder):
    """Encode a root state plus lookahead candidates as packed flat relations."""

    def __init__(
        self,
        domain: DomainInput,
        *,
        transition_mode: (
            FlatHorizonEncoderMode | HorizonEncoderMode | str | None
        ) = None,
        target_symbol_prefix: str = "target:",
        parent_relation: str = DEFAULT_PARENT_RELATION,
        sibling_relation: str = "_sibling_",
        cousin_relation: str = "_cousin_",
        enable_parent_relation: bool = False,
        enable_sibling_relation: bool = False,
        enable_cousin_relation: bool = False,
        root_policy: RootPolicy | str = "exclude",
        max_goal_level: int = 0,
        support_literals: bool = False,
        include_static: bool = True,
        export_node_names: bool = True,
        ignore_zero_arity_relations: bool = True,
        ignore_actions: bool = True,
        use_predicate_virtual_nodes: bool = False,
        include_lgan_edges: bool = False,
        lgan_tn_edge_pos: str = DEFAULT_LGAN_TN_EDGE_POS,
        lgan_nn_edge_pos: str = DEFAULT_LGAN_NN_EDGE_POS,
        lgan_rr_edge_pos: str = DEFAULT_LGAN_RR_EDGE_POS,
        goal_derivations: Any | None = None,
    ) -> None:
        """Create a flat horizon encoder.

        This lane reads a root state plus a `TransitionDAG` and creates
        candidate state rows.

        `root_policy` controls how the root is treated:

        - `include`: root is encoded and is also a target
        - `encode_only`: root is encoded, but not a target
        - `exclude`: root is not a target, and root facts stay base relations

        When `include_lgan_edges=True`, LGAN anchors follow the target rows.
        There is no `lgan_anchor_sources` switch here.
        """
        normalized_root_policy = normalize_root_policy(root_policy)
        config_kwargs: dict[str, Any] = {
            "max_goal_level": max_goal_level,
            "support_literals": support_literals,
            "include_static": include_static,
            "export_node_names": export_node_names,
            "ignore_zero_arity_relations": ignore_zero_arity_relations,
            "ignore_actions": ignore_actions,
            "use_predicate_virtual_nodes": use_predicate_virtual_nodes,
            "include_lgan_edges": include_lgan_edges,
            "target_symbol_prefix": target_symbol_prefix,
            "parent_relation": parent_relation,
            "sibling_relation": sibling_relation,
            "cousin_relation": cousin_relation,
            "lgan_tn_edge_pos": lgan_tn_edge_pos,
            "lgan_nn_edge_pos": lgan_nn_edge_pos,
            "lgan_rr_edge_pos": lgan_rr_edge_pos,
            "enable_parent_relation": enable_parent_relation,
            "enable_sibling_relation": enable_sibling_relation,
            "enable_cousin_relation": enable_cousin_relation,
            "root_policy": normalized_root_policy,
        }
        normalized_mode = _normalize_flat_horizon_mode(transition_mode)
        if normalized_mode is not None:
            config_kwargs["transition_mode"] = normalized_mode
        if goal_derivations is not None:
            config_kwargs["goal_derivations"] = goal_derivations
        config = FlatHorizonEncoderConfig(**config_kwargs)
        self._engine = FlatHorizonEncoderEngine(_advanced_domain(domain), config)
        self._config = config
        self.entity_node_type = "entity"
        self.use_predicate_virtual_nodes = bool(config.use_predicate_virtual_nodes)
        self.target_symbol_prefix = config.target_symbol_prefix
        self.parent_relation = config.parent_relation
        self.sibling_relation = config.sibling_relation
        self.cousin_relation = config.cousin_relation
        self.include_lgan_edges = bool(config.include_lgan_edges)
        self.lgan_tn_edge_pos = str(config.lgan_tn_edge_pos)
        self.lgan_nn_edge_pos = str(config.lgan_nn_edge_pos)
        self.lgan_rr_edge_pos = str(config.lgan_rr_edge_pos)
        self._lgan_edge_positions = {
            self.lgan_tn_edge_pos,
            self.lgan_nn_edge_pos,
            self.lgan_rr_edge_pos,
        }

    @property
    def engine(self) -> FlatHorizonEncoderEngine:
        """Expose the native flat horizon engine."""
        return self._engine

    @property
    def config(self) -> FlatHorizonEncoderConfig:
        """Expose the resolved native config."""
        return self._config

    @property
    def relation_dict(self):
        """Expose the relation schema used by the native engine."""
        return self._engine.relation_dict

    def _accepted_kwargs(self) -> set[str]:
        return {"dag", "dags", "history_subgoals", "history_max_steps"}

    def _encode(
        self,
        root: StateInput,
        *,
        goals: GoalBatchInput = None,
        actions: ActionBatchInput = None,
        subgoal_layers: SubgoalLayersInput = None,
        dag: TransitionDAG | RXStateDAG | None = None,
        history_subgoals: HistorySubgoalInput | None = None,
        history_max_steps: int | None = None,
    ) -> FlatEncoding:
        validate_single_optional_payloads(
            FLAT_HORIZON_LANE_SPEC,
            actions=actions,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
        )
        adv_root = _advanced_state(root)
        subgoal_layers_list = None if subgoal_layers is None else list(subgoal_layers)
        _validate_subgoal_layers_state_payload(
            subgoal_layers_list,
            state_index=0,
            max_goal_level=int(self._config.max_goal_level),
        )
        dag = ensure_transition_dag(root, dag)
        inputs = prepare_goal_inputs(root, goals, subgoal_layers_list)
        return self._engine.encode(adv_root, dag, inputs)

    def encode(
        self,
        root: StateInput,
        dag: TransitionDAG | RXStateDAG | None = None,
        *,
        goals: GoalBatchInput = None,
        actions: ActionBatchInput = None,
        subgoal_layers: SubgoalLayersInput = None,
        history_subgoals: HistorySubgoalInput | None = None,
        history_max_steps: int | None = None,
        include_metadata: bool = True,
        **kwargs,
    ) -> FlatEncoding:
        """Encode one root state and optional lookahead DAG.

        If `dag` is omitted, a one-node DAG for the root is used. `actions` and
        `history_subgoals` are accepted for API consistency but non-empty
        payloads are rejected on this lane.
        """
        return super().encode(
            root,
            goals=goals,
            actions=actions,
            subgoal_layers=subgoal_layers,
            dag=dag,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
            include_metadata=include_metadata,
            **kwargs,
        )

    def _encode_batch(
        self,
        roots: StateBatchInput,
        dags: DagBatchParam = None,
        *,
        goals: GoalBatchParam = None,
        subgoal_layers: SubgoalLayersBatchParam = None,
        actions: ActionBatchParam = None,
        history_subgoals: HistorySubgoalsBatchParam = None,
        history_max_steps: int | None = None,
    ) -> FlatEncoding:
        validate_batch_optional_payloads(
            FLAT_HORIZON_LANE_SPEC,
            actions=actions,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
        )
        roots_for_core = _convert_batch_payload(
            roots,
            is_leaf=is_state_input,
            convert_leaf=to_advanced_state,
        )
        goals_for_core = _convert_batch_payload(
            goals,
            is_leaf=is_goal_literal_input,
            convert_leaf=to_advanced_literal,
        )
        subgoal_layers_for_core = _convert_batch_payload(
            subgoal_layers,
            is_leaf=is_goal_literal_input,
            convert_leaf=to_advanced_literal,
        )
        dags_for_core = _normalize_dag_batch_data(dags)
        parsed_roots = parse_states_batch(roots_for_core)
        _validate_subgoal_layers_batch_payload(
            subgoal_layers_for_core,
            state_count=len(parsed_roots),
            max_goal_level=int(self._config.max_goal_level),
        )
        if dags_for_core is None:
            parsed_dags = [None] * len(parsed_roots)
        else:
            parsed_dags = parse_dags_batch_param(
                dags_for_core, state_count=len(parsed_roots)
            )
        for adv_root, dag in zip(parsed_roots, parsed_dags, strict=True):
            if dag is not None and dag.root().get_index() != adv_root.get_index():
                raise ValueError("dag root must match root state")
        return self._engine.encode_batch(
            parsed_roots,
            dags=parsed_dags,
            goals=goals_for_core,
            actions=None,
            subgoal_layers=subgoal_layers_for_core,
            history_subgoals=None,
            history_max_steps=None,
        )

    def encode_batch(
        self,
        roots: StateBatchInput,
        dags: DagBatchParam = None,
        *,
        goals: GoalBatchParam = None,
        actions: ActionBatchParam = None,
        subgoal_layers: SubgoalLayersBatchParam = None,
        history_subgoals: HistorySubgoalsBatchParam = None,
        history_max_steps: int | None = None,
        batch_attrs: Mapping[str, Any] | None = None,
        collate_spec: CollateSpecParam = None,
        include_metadata: bool = True,
        **kwargs,
    ) -> FlatEncoding:
        """Encode many root/DAG inputs into one flat batch."""
        return super().encode_batch(
            roots,
            goals=goals,
            actions=actions,
            subgoal_layers=subgoal_layers,
            dags=dags,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
            batch_attrs=batch_attrs,
            collate_spec=collate_spec,
            include_metadata=include_metadata,
            **kwargs,
        )

    def stream(self) -> FlatHorizonEncoderStream:
        """Return an append-only stream for root/DAG horizon inputs."""
        return FlatHorizonEncoderStream(self)

    def mutable_stream(self) -> FlatHorizonMutableEncoderStream:
        """Return a mutable stream with `append`, `update`, and `remove`."""
        return FlatHorizonMutableEncoderStream(self)

engine property

Expose the native flat horizon engine.

config property

Expose the resolved native config.

relation_dict property

Expose the relation schema used by the native engine.

__init__(domain, *, transition_mode=None, target_symbol_prefix='target:', parent_relation=DEFAULT_PARENT_RELATION, sibling_relation='_sibling_', cousin_relation='_cousin_', enable_parent_relation=False, enable_sibling_relation=False, enable_cousin_relation=False, root_policy='exclude', max_goal_level=0, support_literals=False, include_static=True, export_node_names=True, ignore_zero_arity_relations=True, ignore_actions=True, use_predicate_virtual_nodes=False, include_lgan_edges=False, lgan_tn_edge_pos=DEFAULT_LGAN_TN_EDGE_POS, lgan_nn_edge_pos=DEFAULT_LGAN_NN_EDGE_POS, lgan_rr_edge_pos=DEFAULT_LGAN_RR_EDGE_POS, goal_derivations=None)

Create a flat horizon encoder.

This lane reads a root state plus a TransitionDAG and creates candidate state rows.

root_policy controls how the root is treated:

• include: root is encoded and is also a target
• encode_only: root is encoded, but not a target
• exclude: root is not a target, and root facts stay base relations

When include_lgan_edges=True, LGAN anchors follow the target rows. There is no lgan_anchor_sources switch here.

Source code in src/mifrost/encoders/flat_horizon.py

def __init__(
    self,
    domain: DomainInput,
    *,
    transition_mode: (
        FlatHorizonEncoderMode | HorizonEncoderMode | str | None
    ) = None,
    target_symbol_prefix: str = "target:",
    parent_relation: str = DEFAULT_PARENT_RELATION,
    sibling_relation: str = "_sibling_",
    cousin_relation: str = "_cousin_",
    enable_parent_relation: bool = False,
    enable_sibling_relation: bool = False,
    enable_cousin_relation: bool = False,
    root_policy: RootPolicy | str = "exclude",
    max_goal_level: int = 0,
    support_literals: bool = False,
    include_static: bool = True,
    export_node_names: bool = True,
    ignore_zero_arity_relations: bool = True,
    ignore_actions: bool = True,
    use_predicate_virtual_nodes: bool = False,
    include_lgan_edges: bool = False,
    lgan_tn_edge_pos: str = DEFAULT_LGAN_TN_EDGE_POS,
    lgan_nn_edge_pos: str = DEFAULT_LGAN_NN_EDGE_POS,
    lgan_rr_edge_pos: str = DEFAULT_LGAN_RR_EDGE_POS,
    goal_derivations: Any | None = None,
) -> None:
    """Create a flat horizon encoder.

    This lane reads a root state plus a `TransitionDAG` and creates
    candidate state rows.

    `root_policy` controls how the root is treated:

    - `include`: root is encoded and is also a target
    - `encode_only`: root is encoded, but not a target
    - `exclude`: root is not a target, and root facts stay base relations

    When `include_lgan_edges=True`, LGAN anchors follow the target rows.
    There is no `lgan_anchor_sources` switch here.
    """
    normalized_root_policy = normalize_root_policy(root_policy)
    config_kwargs: dict[str, Any] = {
        "max_goal_level": max_goal_level,
        "support_literals": support_literals,
        "include_static": include_static,
        "export_node_names": export_node_names,
        "ignore_zero_arity_relations": ignore_zero_arity_relations,
        "ignore_actions": ignore_actions,
        "use_predicate_virtual_nodes": use_predicate_virtual_nodes,
        "include_lgan_edges": include_lgan_edges,
        "target_symbol_prefix": target_symbol_prefix,
        "parent_relation": parent_relation,
        "sibling_relation": sibling_relation,
        "cousin_relation": cousin_relation,
        "lgan_tn_edge_pos": lgan_tn_edge_pos,
        "lgan_nn_edge_pos": lgan_nn_edge_pos,
        "lgan_rr_edge_pos": lgan_rr_edge_pos,
        "enable_parent_relation": enable_parent_relation,
        "enable_sibling_relation": enable_sibling_relation,
        "enable_cousin_relation": enable_cousin_relation,
        "root_policy": normalized_root_policy,
    }
    normalized_mode = _normalize_flat_horizon_mode(transition_mode)
    if normalized_mode is not None:
        config_kwargs["transition_mode"] = normalized_mode
    if goal_derivations is not None:
        config_kwargs["goal_derivations"] = goal_derivations
    config = FlatHorizonEncoderConfig(**config_kwargs)
    self._engine = FlatHorizonEncoderEngine(_advanced_domain(domain), config)
    self._config = config
    self.entity_node_type = "entity"
    self.use_predicate_virtual_nodes = bool(config.use_predicate_virtual_nodes)
    self.target_symbol_prefix = config.target_symbol_prefix
    self.parent_relation = config.parent_relation
    self.sibling_relation = config.sibling_relation
    self.cousin_relation = config.cousin_relation
    self.include_lgan_edges = bool(config.include_lgan_edges)
    self.lgan_tn_edge_pos = str(config.lgan_tn_edge_pos)
    self.lgan_nn_edge_pos = str(config.lgan_nn_edge_pos)
    self.lgan_rr_edge_pos = str(config.lgan_rr_edge_pos)
    self._lgan_edge_positions = {
        self.lgan_tn_edge_pos,
        self.lgan_nn_edge_pos,
        self.lgan_rr_edge_pos,
    }

encode(root, dag=None, *, goals=None, actions=None, subgoal_layers=None, history_subgoals=None, history_max_steps=None, include_metadata=True, **kwargs)

Encode one root state and optional lookahead DAG.

If dag is omitted, a one-node DAG for the root is used. actions and history_subgoals are accepted for API consistency but non-empty payloads are rejected on this lane.

Source code in src/mifrost/encoders/flat_horizon.py

def encode(
    self,
    root: StateInput,
    dag: TransitionDAG | RXStateDAG | None = None,
    *,
    goals: GoalBatchInput = None,
    actions: ActionBatchInput = None,
    subgoal_layers: SubgoalLayersInput = None,
    history_subgoals: HistorySubgoalInput | None = None,
    history_max_steps: int | None = None,
    include_metadata: bool = True,
    **kwargs,
) -> FlatEncoding:
    """Encode one root state and optional lookahead DAG.

    If `dag` is omitted, a one-node DAG for the root is used. `actions` and
    `history_subgoals` are accepted for API consistency but non-empty
    payloads are rejected on this lane.
    """
    return super().encode(
        root,
        goals=goals,
        actions=actions,
        subgoal_layers=subgoal_layers,
        dag=dag,
        history_subgoals=history_subgoals,
        history_max_steps=history_max_steps,
        include_metadata=include_metadata,
        **kwargs,
    )

encode_batch(roots, dags=None, *, goals=None, actions=None, subgoal_layers=None, history_subgoals=None, history_max_steps=None, batch_attrs=None, collate_spec=None, include_metadata=True, **kwargs)

Encode many root/DAG inputs into one flat batch.

Source code in src/mifrost/encoders/flat_horizon.py

def encode_batch(
    self,
    roots: StateBatchInput,
    dags: DagBatchParam = None,
    *,
    goals: GoalBatchParam = None,
    actions: ActionBatchParam = None,
    subgoal_layers: SubgoalLayersBatchParam = None,
    history_subgoals: HistorySubgoalsBatchParam = None,
    history_max_steps: int | None = None,
    batch_attrs: Mapping[str, Any] | None = None,
    collate_spec: CollateSpecParam = None,
    include_metadata: bool = True,
    **kwargs,
) -> FlatEncoding:
    """Encode many root/DAG inputs into one flat batch."""
    return super().encode_batch(
        roots,
        goals=goals,
        actions=actions,
        subgoal_layers=subgoal_layers,
        dags=dags,
        history_subgoals=history_subgoals,
        history_max_steps=history_max_steps,
        batch_attrs=batch_attrs,
        collate_spec=collate_spec,
        include_metadata=include_metadata,
        **kwargs,
    )

stream()

Return an append-only stream for root/DAG horizon inputs.

Source code in src/mifrost/encoders/flat_horizon.py

def stream(self) -> FlatHorizonEncoderStream:
    """Return an append-only stream for root/DAG horizon inputs."""
    return FlatHorizonEncoderStream(self)

mutable_stream()

Return a mutable stream with append, update, and remove.

Source code in src/mifrost/encoders/flat_horizon.py

def mutable_stream(self) -> FlatHorizonMutableEncoderStream:
    """Return a mutable stream with `append`, `update`, and `remove`."""
    return FlatHorizonMutableEncoderStream(self)

FlatTransitionEncoder

Bases: _FlatTransitionEncoderBase

Encode full current-to-successor structure on the flat carrier.

Source code in src/mifrost/encoders/flat_transition.py

class FlatTransitionEncoder(_FlatTransitionEncoderBase):
    """Encode full current-to-successor structure on the flat carrier."""

    def __init__(self, domain, **kwargs) -> None:
        """Create a full flat transition encoder.

        LGAN, when enabled, uses successor-state candidate rows from the
        underlying flat horizon lane.
        """
        kwargs.setdefault("root_policy", "exclude")
        super().__init__(
            domain,
            transition_mode="full",
            enable_parent_relation=False,
            enable_sibling_relation=False,
            enable_cousin_relation=False,
            ignore_actions=True,
            **kwargs,
        )

    def stream(self) -> "FlatTransitionEncoderStream":
        """Return a mutable stream for full flat transitions."""
        return FlatTransitionEncoderStream(self)

__init__(domain, **kwargs)

Create a full flat transition encoder.

LGAN, when enabled, uses successor-state candidate rows from the underlying flat horizon lane.

Source code in src/mifrost/encoders/flat_transition.py

def __init__(self, domain, **kwargs) -> None:
    """Create a full flat transition encoder.

    LGAN, when enabled, uses successor-state candidate rows from the
    underlying flat horizon lane.
    """
    kwargs.setdefault("root_policy", "exclude")
    super().__init__(
        domain,
        transition_mode="full",
        enable_parent_relation=False,
        enable_sibling_relation=False,
        enable_cousin_relation=False,
        ignore_actions=True,
        **kwargs,
    )

stream()

Return a mutable stream for full flat transitions.

Source code in src/mifrost/encoders/flat_transition.py

def stream(self) -> "FlatTransitionEncoderStream":
    """Return a mutable stream for full flat transitions."""
    return FlatTransitionEncoderStream(self)

FlatTransitionEffectsEncoder

Bases: _FlatTransitionEncoderBase

Encode only changed successor structure on the flat carrier.

Source code in src/mifrost/encoders/flat_transition.py

class FlatTransitionEffectsEncoder(_FlatTransitionEncoderBase):
    """Encode only changed successor structure on the flat carrier."""

    def __init__(self, domain, **kwargs) -> None:
        """Create a delta/effects flat transition encoder."""
        kwargs.setdefault("root_policy", "exclude")
        super().__init__(
            domain,
            transition_mode="delta",
            enable_parent_relation=False,
            enable_sibling_relation=False,
            enable_cousin_relation=False,
            ignore_actions=True,
            **kwargs,
        )

    def stream(self) -> "FlatTransitionEffectsEncoderStream":
        """Return a mutable stream for flat delta/effects transitions."""
        return FlatTransitionEffectsEncoderStream(self)

__init__(domain, **kwargs)

Create a delta/effects flat transition encoder.

Source code in src/mifrost/encoders/flat_transition.py

def __init__(self, domain, **kwargs) -> None:
    """Create a delta/effects flat transition encoder."""
    kwargs.setdefault("root_policy", "exclude")
    super().__init__(
        domain,
        transition_mode="delta",
        enable_parent_relation=False,
        enable_sibling_relation=False,
        enable_cousin_relation=False,
        ignore_actions=True,
        **kwargs,
    )

stream()

Return a mutable stream for flat delta/effects transitions.

Source code in src/mifrost/encoders/flat_transition.py

def stream(self) -> "FlatTransitionEffectsEncoderStream":
    """Return a mutable stream for flat delta/effects transitions."""
    return FlatTransitionEffectsEncoderStream(self)

FlatRelationData

Bases: Data

Packed flat relation carrier with PyG-compatible batching behavior.

Source code in src/mifrost/encoders/flat_data.py

class FlatRelationData(Data):
    """Packed flat relation carrier with PyG-compatible batching behavior."""

    def __inc__(self, key: str, value: Any, *args, **kwargs) -> Any:
        if key in {
            "relation_args",
            "object_indices",
            "history_entity_indices",
            "target_entity_indices",
            "target_positions",
            "lgan_tn_entity_indices",
            "lgan_nn_entity_indices",
        }:
            return int(getattr(self, "num_nodes", 0))
        if key in {
            "lgan_tn_relation_indices",
            "lgan_nn_relation_indices",
            "lgan_rr_src_relation_indices",
            "lgan_rr_dst_relation_indices",
        }:
            return self._relation_instance_offset()
        return super().__inc__(key, value, *args, **kwargs)

    @cached_property
    def schema(self) -> FlatRelationSchema:
        """Return the normalized flat relation schema for this batch."""
        return FlatRelationSchema(
            names=_normalize_str_tuple(getattr(self, "relation_names", ()) or ()),
            arities=_normalize_int_tuple(getattr(self, "relation_arities", ()) or ()),
            logical_arities=_normalize_int_tuple(
                getattr(self, "relation_logical_arities", None)
                or getattr(self, "relation_arities", ())
            ),
            encoded_arities=_normalize_int_tuple(
                getattr(self, "relation_encoded_arities", None)
                or getattr(self, "relation_arities", ())
            ),
            sources=_normalize_str_tuple(getattr(self, "relation_sources", ()) or ()),
            slot_role_names=_normalize_str_tuple(
                getattr(self, "slot_role_names", ()) or ()
            ),
            slot_role_ids=_normalize_int_tuple(
                getattr(self, "relation_slot_roles", ()) or ()
            ),
            slot_role_offsets=_normalize_int_tuple(
                getattr(self, "relation_slot_role_offsets", ()) or ()
            ),
            fingerprint=_normalize_optional_int(
                getattr(self, "schema_fingerprint", None)
            ),
        )

    @property
    def flattened_relations(self) -> dict[str, torch.Tensor]:
        """Return the flat relations grouped by relation name."""
        return self.flattened_relations_view()

    def relation_instance_counts_total(self) -> torch.Tensor:
        """Return total instance counts per relation across the whole batch."""
        counts = getattr(self, "relation_counts", None)
        if counts is None:
            return torch.zeros((len(self.schema.names),), dtype=torch.long)
        counts = counts.long()
        if counts.dim() == 1:
            return counts
        return counts.sum(dim=0)

    def relation_slot_offsets(self, graph_index: int | None = None) -> torch.Tensor:
        """Return slot offsets into `relation_args` for one graph or the whole batch."""
        counts = self._relation_counts_for(graph_index)
        arities = self._relation_arities_tensor(counts.device)
        slot_counts = counts * arities
        prefix = torch.zeros(
            (slot_counts.numel() + 1,), dtype=torch.long, device=slot_counts.device
        )
        if slot_counts.numel() > 0:
            prefix[1:] = torch.cumsum(slot_counts, dim=0)
        return prefix

    def flattened_relations_view(
        self, graph_index: int | None = None
    ) -> dict[str, torch.Tensor]:
        """Slice `relation_args` into per-relation matrices.

        With `graph_index=None`, the result covers the whole batch. With a graph
        index, the result is limited to one graph.
        """
        relation_args = getattr(self, "relation_args", None)
        if relation_args is None:
            return {
                name: torch.empty((0, arity), dtype=torch.long)
                for name, arity in zip(self.schema.names, self.schema.arities)
            }

        relation_args = relation_args.long().view(-1)
        if graph_index is None and self.num_graphs > 1:
            counts = getattr(self, "relation_counts", None)
            if counts is None:
                return {
                    name: relation_args.new_empty((0, arity))
                    for name, arity in zip(self.schema.names, self.schema.arities)
                }
            counts = counts.long()
            arities = self.schema.arities
            chunks: dict[str, list[torch.Tensor]] = {
                name: [] for name in self.schema.names
            }
            start = 0
            for row in counts:
                for relation_idx, name in enumerate(self.schema.names):
                    arity = arities[relation_idx]
                    instances = int(row[relation_idx].item()) if row.numel() > 0 else 0
                    slots = instances * arity
                    chunk = relation_args[start : start + slots]
                    if arity > 0:
                        chunks[name].append(chunk.view(instances, arity))
                    else:
                        chunks[name].append(chunk.new_empty((instances, 0)))
                    start += slots
            out: dict[str, torch.Tensor] = {}
            for relation_name, arity in zip(self.schema.names, self.schema.arities):
                parts = chunks[relation_name]
                out[relation_name] = (
                    torch.cat(parts, dim=0)
                    if parts
                    else relation_args.new_empty((0, arity))
                )
            return out

        counts = self._relation_counts_for(graph_index)
        arities = self.schema.arities
        start = self._relation_arg_start(graph_index)
        out: dict[str, torch.Tensor] = {}
        for relation_idx, name in enumerate(self.schema.names):
            arity = arities[relation_idx]
            instances = int(counts[relation_idx].item()) if counts.numel() > 0 else 0
            slots = instances * arity
            chunk = relation_args[start : start + slots]
            if arity > 0:
                out[name] = chunk.view(instances, arity)
            else:
                out[name] = chunk.new_empty((instances, 0))
            start += slots
        return out

    def graph_node_names(self, graph_index: int = 0) -> list[str]:
        """Return entity-row names for one graph."""
        node_names = getattr(self, "node_names", None)
        if node_names is None:
            start, end = self.graph_node_range(graph_index)
            return [f"entity:{idx}" for idx in range(start, end)]
        if node_names and isinstance(node_names[0], (list, tuple)):
            return [str(name) for name in node_names[graph_index]]
        return [str(name) for name in node_names]

    def graph_object_names(self, graph_index: int = 0) -> list[str]:
        """Return object names for one graph."""
        object_names = getattr(self, "object_names", None)
        if object_names is None:
            return self.graph_node_names(graph_index)
        if object_names and isinstance(object_names[0], (list, tuple)):
            return [str(name) for name in object_names[graph_index]]
        return [str(name) for name in object_names]

    def graph_object_indices(self, graph_index: int = 0) -> torch.Tensor:
        """Return global entity rows that correspond to objects."""
        return self._graph_cat_field_slice(
            field_name="object_indices",
            size_field_name="object_sizes",
            graph_index=graph_index,
        )

    def graph_history_entity_indices(self, graph_index: int = 0) -> torch.Tensor:
        """Return global entity rows used as history carriers."""
        return self._graph_cat_field_slice(
            field_name="history_entity_indices",
            size_field_name="history_entity_sizes",
            graph_index=graph_index,
        )

    def graph_history_entity_dt(self, graph_index: int = 0) -> torch.Tensor:
        """Return the `dt` values for history carrier rows."""
        return self._graph_cat_field_slice(
            field_name="history_entity_dt",
            size_field_name="history_entity_sizes",
            graph_index=graph_index,
        )

    def graph_history_entity_names(self, graph_index: int = 0) -> list[str]:
        """Return names for the history carrier rows of one graph."""
        history_entity_indices = self.graph_history_entity_indices(graph_index)
        if history_entity_indices.numel() == 0:
            return []
        start, _end = self.graph_node_range(graph_index)
        local_names = self.graph_node_names(graph_index)
        return [
            str(local_names[int(global_idx.item()) - start])
            for global_idx in history_entity_indices
        ]

    def graph_target_entity_indices(
        self, graph_index: int = 0, group: str | int | None = None
    ) -> torch.Tensor:
        """Return candidate-carrier rows for one graph.

        Use `group` to limit the result to one source, such as `goal`,
        `subgoal`, `action`, or `history`.
        """
        indices = self._graph_cat_field_slice(
            field_name="target_entity_indices",
            size_field_name="target_entity_sizes",
            graph_index=graph_index,
        )
        if group is None or indices.numel() == 0:
            return indices
        group_ids = self.graph_target_entity_group_ids(graph_index)
        mask = self._group_mask(
            group_ids, group=group, attr_name="target_entity_groups"
        )
        return indices[mask]

    def graph_target_entity_group_ids(self, graph_index: int = 0) -> torch.Tensor:
        """Return source-group ids for target-entity rows."""
        return self._graph_cat_field_slice(
            field_name="target_entity_group_ids",
            size_field_name="target_entity_sizes",
            graph_index=graph_index,
        )

    def graph_target_entity_names(
        self, graph_index: int = 0, group: str | int | None = None
    ) -> list[str]:
        """Return names for target-entity rows, optionally filtered by group."""
        target_entity_indices = self.graph_target_entity_indices(
            graph_index, group=group
        )
        if target_entity_indices.numel() == 0:
            return []
        start, _end = self.graph_node_range(graph_index)
        local_names = self.graph_node_names(graph_index)
        return [
            str(local_names[int(global_idx.item()) - start])
            for global_idx in target_entity_indices
        ]

    def graph_target_positions(self, graph_index: int = 0) -> torch.Tensor:
        """Return entity-row positions for prediction targets."""
        return self._graph_cat_field_slice(
            field_name="target_positions",
            size_field_name="target_sizes",
            graph_index=graph_index,
        )

    def graph_target_indices(self, graph_index: int = 0) -> torch.Tensor:
        """Return per-graph target indices in encounter order."""
        return self._graph_cat_field_slice(
            field_name="target_indices",
            size_field_name="target_sizes",
            graph_index=graph_index,
        )

    def graph_target_candidate_ids(self, graph_index: int = 0) -> torch.Tensor:
        """Return stable candidate ids for prediction targets."""
        return self._graph_cat_field_slice(
            field_name="target_candidate_ids",
            size_field_name="target_sizes",
            graph_index=graph_index,
        )

    def graph_target_depths(self, graph_index: int = 0) -> torch.Tensor:
        """Return target depths when the encoder emitted them."""
        return self._graph_cat_field_slice(
            field_name="target_depths",
            size_field_name="target_sizes",
            graph_index=graph_index,
        )

    def graph_target_group_ids(self, graph_index: int = 0) -> torch.Tensor:
        """Return source-group ids for prediction targets."""
        return self._graph_cat_field_slice(
            field_name="target_group_ids",
            size_field_name="target_sizes",
            graph_index=graph_index,
        )

    def graph_target_names(self, graph_index: int = 0) -> list[str]:
        """Return display names for prediction targets."""
        target_names = getattr(self, "target_names", None)
        if target_names is None:
            return []
        if target_names and isinstance(target_names[0], (list, tuple)):
            return [str(name) for name in target_names[graph_index]]
        return [str(name) for name in target_names]

    def graph_relation_instance_range(self, graph_index: int = 0) -> tuple[int, int]:
        """Return the global relation-instance index range for one graph."""
        relation_instance_sizes = getattr(self, "relation_instance_sizes", None)
        if relation_instance_sizes is None:
            counts = getattr(self, "relation_counts", None)
            if counts is None:
                return (0, 0)
            counts = counts.long()
            if counts.dim() == 1:
                total = int(counts.sum().item())
                return (0, total)
            if graph_index < 0 or graph_index >= counts.size(0):
                raise IndexError(
                    f"graph_index {graph_index} out of range for {counts.size(0)} graphs"
                )
            sizes = counts.sum(dim=1)
        else:
            sizes = relation_instance_sizes.long().view(-1)
            if graph_index < 0 or graph_index >= len(sizes):
                raise IndexError(
                    f"graph_index {graph_index} out of range for {len(sizes)} graphs"
                )
        start = int(sizes[:graph_index].sum().item()) if graph_index > 0 else 0
        end = start + int(sizes[graph_index].item())
        return start, end

    def graph_lgan_tn_edges(self, graph_index: int = 0) -> torch.Tensor:
        """Return packed TN LGAN edges for one graph."""
        return self._graph_edge_slice(
            src_field_name="lgan_tn_relation_indices",
            dst_field_name="lgan_tn_entity_indices",
            size_field_name="lgan_tn_sizes",
            graph_index=graph_index,
        )

    def graph_lgan_nn_edges(self, graph_index: int = 0) -> torch.Tensor:
        """Return packed NN LGAN edges for one graph."""
        return self._graph_edge_slice(
            src_field_name="lgan_nn_relation_indices",
            dst_field_name="lgan_nn_entity_indices",
            size_field_name="lgan_nn_sizes",
            graph_index=graph_index,
        )

    def graph_lgan_rr_edges(self, graph_index: int = 0) -> torch.Tensor:
        """Return packed RR LGAN edges for one graph."""
        return self._graph_edge_slice(
            src_field_name="lgan_rr_src_relation_indices",
            dst_field_name="lgan_rr_dst_relation_indices",
            size_field_name="lgan_rr_sizes",
            graph_index=graph_index,
        )

    def graph_node_range(self, graph_index: int = 0) -> tuple[int, int]:
        """Return the global entity-row range for one graph."""
        node_sizes = getattr(self, "node_sizes", None)
        if node_sizes is None:
            return (0, int(getattr(self, "num_nodes", 0)))
        node_sizes = node_sizes.long().view(-1)
        if graph_index < 0 or graph_index >= len(node_sizes):
            raise IndexError(
                f"graph_index {graph_index} out of range for {len(node_sizes)} graphs"
            )
        start = int(node_sizes[:graph_index].sum().item()) if graph_index > 0 else 0
        end = start + int(node_sizes[graph_index].item())
        return start, end

    def graph_entity_role_ids(self, graph_index: int = 0) -> torch.Tensor:
        """Return per-entity role ids for one graph."""
        start, end = self.graph_node_range(graph_index)
        entity_role_ids = getattr(self, "entity_role_ids", None)
        if entity_role_ids is None:
            return torch.empty((0,), dtype=torch.long)
        return entity_role_ids.long().view(-1)[start:end]

    def graph_entity_roles(self, graph_index: int = 0) -> list[str]:
        """Return decoded per-entity role labels for one graph."""
        role_names = [
            str(value) for value in getattr(self, "entity_role_names", []) or []
        ]
        out: list[str] = []
        for role_id in self.graph_entity_role_ids(graph_index).tolist():
            if 0 <= role_id < len(role_names):
                out.append(role_names[role_id])
            else:
                out.append(str(role_id))
        return out

    @property
    def num_graphs(self) -> int:
        """Return how many graphs are stored in this flat carrier."""
        if hasattr(self, "_num_graphs"):
            return int(self._num_graphs)
        node_sizes = getattr(self, "node_sizes", None)
        if node_sizes is not None and torch.is_tensor(node_sizes):
            return int(node_sizes.view(-1).numel())
        return 1

    def _relation_counts_for(self, graph_index: int | None) -> torch.Tensor:
        counts = getattr(self, "relation_counts", None)
        if counts is None:
            return torch.zeros((len(self.schema.names),), dtype=torch.long)
        counts = counts.long()
        if counts.dim() == 1:
            return counts
        if graph_index is None:
            return counts.sum(dim=0)
        if graph_index < 0 or graph_index >= counts.size(0):
            raise IndexError(
                f"graph_index {graph_index} out of range for {counts.size(0)} graphs"
            )
        return counts[graph_index]

    def _graph_cat_field_slice(
        self,
        *,
        field_name: str,
        size_field_name: str,
        graph_index: int,
    ) -> torch.Tensor:
        values = getattr(self, field_name, None)
        if values is None:
            return torch.empty((0,), dtype=torch.long)
        values = values.long().view(-1)
        sizes = getattr(self, size_field_name, None)
        if sizes is None:
            return values
        sizes = sizes.long().view(-1)
        if graph_index < 0 or graph_index >= len(sizes):
            raise IndexError(
                f"graph_index {graph_index} out of range for {len(sizes)} graphs"
            )
        start = int(sizes[:graph_index].sum().item()) if graph_index > 0 else 0
        end = start + int(sizes[graph_index].item())
        return values[start:end]

    def _graph_edge_slice(
        self,
        *,
        src_field_name: str,
        dst_field_name: str,
        size_field_name: str,
        graph_index: int,
    ) -> torch.Tensor:
        src = self._graph_cat_field_slice(
            field_name=src_field_name,
            size_field_name=size_field_name,
            graph_index=graph_index,
        )
        dst = self._graph_cat_field_slice(
            field_name=dst_field_name,
            size_field_name=size_field_name,
            graph_index=graph_index,
        )
        if src.numel() == 0 or dst.numel() == 0:
            device = src.device if src.numel() else dst.device
            return torch.empty((2, 0), dtype=torch.long, device=device)
        return torch.stack((src, dst), dim=0)

    def _relation_arities_tensor(
        self, device: torch.device | None = None
    ) -> torch.Tensor:
        return torch.tensor(
            list(self.schema.arities),
            dtype=torch.long,
            device=device or torch.device("cpu"),
        )

    def _relation_arg_start(self, graph_index: int | None) -> int:
        if graph_index is None or self.num_graphs <= 1:
            return 0
        counts = getattr(self, "relation_counts", None)
        if counts is None:
            return 0
        counts = counts.long()
        if graph_index <= 0:
            return 0
        arities = self._relation_arities_tensor(counts.device)
        prior_counts = counts[:graph_index].sum(dim=0)
        return int((prior_counts * arities).sum().item())

    def _group_mask(
        self,
        group_ids: torch.Tensor,
        *,
        group: str | int,
        attr_name: str,
    ) -> torch.Tensor:
        target_group_id = self._group_id(group=group, attr_name=attr_name)
        return group_ids.long().view(-1) == target_group_id

    def _group_id(self, *, group: str | int, attr_name: str) -> int:
        if isinstance(group, int):
            return group
        group_names = getattr(self, attr_name, None)
        if group_names is None:
            raise ValueError(f"{attr_name} metadata is not available")
        names = [str(value) for value in group_names]
        if group not in names:
            raise ValueError(
                f"Unknown {attr_name} group {group!r}; expected one of {names!r}"
            )
        return names.index(group)

    def _relation_instance_offset(self) -> int:
        relation_instance_sizes = getattr(self, "relation_instance_sizes", None)
        if relation_instance_sizes is None:
            return 0
        if torch.is_tensor(relation_instance_sizes):
            return int(relation_instance_sizes.long().sum().item())
        return int(sum(int(value) for value in relation_instance_sizes))

schema cached property

Return the normalized flat relation schema for this batch.

flattened_relations property

Return the flat relations grouped by relation name.

num_graphs property

Return how many graphs are stored in this flat carrier.

relation_instance_counts_total()

Return total instance counts per relation across the whole batch.

Source code in src/mifrost/encoders/flat_data.py

def relation_instance_counts_total(self) -> torch.Tensor:
    """Return total instance counts per relation across the whole batch."""
    counts = getattr(self, "relation_counts", None)
    if counts is None:
        return torch.zeros((len(self.schema.names),), dtype=torch.long)
    counts = counts.long()
    if counts.dim() == 1:
        return counts
    return counts.sum(dim=0)

relation_slot_offsets(graph_index=None)

Return slot offsets into relation_args for one graph or the whole batch.

Source code in src/mifrost/encoders/flat_data.py

def relation_slot_offsets(self, graph_index: int | None = None) -> torch.Tensor:
    """Return slot offsets into `relation_args` for one graph or the whole batch."""
    counts = self._relation_counts_for(graph_index)
    arities = self._relation_arities_tensor(counts.device)
    slot_counts = counts * arities
    prefix = torch.zeros(
        (slot_counts.numel() + 1,), dtype=torch.long, device=slot_counts.device
    )
    if slot_counts.numel() > 0:
        prefix[1:] = torch.cumsum(slot_counts, dim=0)
    return prefix

flattened_relations_view(graph_index=None)

Slice relation_args into per-relation matrices.

With graph_index=None, the result covers the whole batch. With a graph index, the result is limited to one graph.

Source code in src/mifrost/encoders/flat_data.py

def flattened_relations_view(
    self, graph_index: int | None = None
) -> dict[str, torch.Tensor]:
    """Slice `relation_args` into per-relation matrices.

    With `graph_index=None`, the result covers the whole batch. With a graph
    index, the result is limited to one graph.
    """
    relation_args = getattr(self, "relation_args", None)
    if relation_args is None:
        return {
            name: torch.empty((0, arity), dtype=torch.long)
            for name, arity in zip(self.schema.names, self.schema.arities)
        }

    relation_args = relation_args.long().view(-1)
    if graph_index is None and self.num_graphs > 1:
        counts = getattr(self, "relation_counts", None)
        if counts is None:
            return {
                name: relation_args.new_empty((0, arity))
                for name, arity in zip(self.schema.names, self.schema.arities)
            }
        counts = counts.long()
        arities = self.schema.arities
        chunks: dict[str, list[torch.Tensor]] = {
            name: [] for name in self.schema.names
        }
        start = 0
        for row in counts:
            for relation_idx, name in enumerate(self.schema.names):
                arity = arities[relation_idx]
                instances = int(row[relation_idx].item()) if row.numel() > 0 else 0
                slots = instances * arity
                chunk = relation_args[start : start + slots]
                if arity > 0:
                    chunks[name].append(chunk.view(instances, arity))
                else:
                    chunks[name].append(chunk.new_empty((instances, 0)))
                start += slots
        out: dict[str, torch.Tensor] = {}
        for relation_name, arity in zip(self.schema.names, self.schema.arities):
            parts = chunks[relation_name]
            out[relation_name] = (
                torch.cat(parts, dim=0)
                if parts
                else relation_args.new_empty((0, arity))
            )
        return out

    counts = self._relation_counts_for(graph_index)
    arities = self.schema.arities
    start = self._relation_arg_start(graph_index)
    out: dict[str, torch.Tensor] = {}
    for relation_idx, name in enumerate(self.schema.names):
        arity = arities[relation_idx]
        instances = int(counts[relation_idx].item()) if counts.numel() > 0 else 0
        slots = instances * arity
        chunk = relation_args[start : start + slots]
        if arity > 0:
            out[name] = chunk.view(instances, arity)
        else:
            out[name] = chunk.new_empty((instances, 0))
        start += slots
    return out

graph_node_names(graph_index=0)

Return entity-row names for one graph.

Source code in src/mifrost/encoders/flat_data.py

def graph_node_names(self, graph_index: int = 0) -> list[str]:
    """Return entity-row names for one graph."""
    node_names = getattr(self, "node_names", None)
    if node_names is None:
        start, end = self.graph_node_range(graph_index)
        return [f"entity:{idx}" for idx in range(start, end)]
    if node_names and isinstance(node_names[0], (list, tuple)):
        return [str(name) for name in node_names[graph_index]]
    return [str(name) for name in node_names]

graph_object_names(graph_index=0)

Return object names for one graph.

Source code in src/mifrost/encoders/flat_data.py

def graph_object_names(self, graph_index: int = 0) -> list[str]:
    """Return object names for one graph."""
    object_names = getattr(self, "object_names", None)
    if object_names is None:
        return self.graph_node_names(graph_index)
    if object_names and isinstance(object_names[0], (list, tuple)):
        return [str(name) for name in object_names[graph_index]]
    return [str(name) for name in object_names]

graph_object_indices(graph_index=0)

Return global entity rows that correspond to objects.

Source code in src/mifrost/encoders/flat_data.py

def graph_object_indices(self, graph_index: int = 0) -> torch.Tensor:
    """Return global entity rows that correspond to objects."""
    return self._graph_cat_field_slice(
        field_name="object_indices",
        size_field_name="object_sizes",
        graph_index=graph_index,
    )

graph_history_entity_indices(graph_index=0)

Return global entity rows used as history carriers.

Source code in src/mifrost/encoders/flat_data.py

def graph_history_entity_indices(self, graph_index: int = 0) -> torch.Tensor:
    """Return global entity rows used as history carriers."""
    return self._graph_cat_field_slice(
        field_name="history_entity_indices",
        size_field_name="history_entity_sizes",
        graph_index=graph_index,
    )

graph_history_entity_dt(graph_index=0)

Return the dt values for history carrier rows.

Source code in src/mifrost/encoders/flat_data.py

def graph_history_entity_dt(self, graph_index: int = 0) -> torch.Tensor:
    """Return the `dt` values for history carrier rows."""
    return self._graph_cat_field_slice(
        field_name="history_entity_dt",
        size_field_name="history_entity_sizes",
        graph_index=graph_index,
    )

graph_history_entity_names(graph_index=0)

Return names for the history carrier rows of one graph.

Source code in src/mifrost/encoders/flat_data.py

def graph_history_entity_names(self, graph_index: int = 0) -> list[str]:
    """Return names for the history carrier rows of one graph."""
    history_entity_indices = self.graph_history_entity_indices(graph_index)
    if history_entity_indices.numel() == 0:
        return []
    start, _end = self.graph_node_range(graph_index)
    local_names = self.graph_node_names(graph_index)
    return [
        str(local_names[int(global_idx.item()) - start])
        for global_idx in history_entity_indices
    ]

graph_target_entity_indices(graph_index=0, group=None)

Return candidate-carrier rows for one graph.

Use group to limit the result to one source, such as goal, subgoal, action, or history.

Source code in src/mifrost/encoders/flat_data.py

def graph_target_entity_indices(
    self, graph_index: int = 0, group: str | int | None = None
) -> torch.Tensor:
    """Return candidate-carrier rows for one graph.

    Use `group` to limit the result to one source, such as `goal`,
    `subgoal`, `action`, or `history`.
    """
    indices = self._graph_cat_field_slice(
        field_name="target_entity_indices",
        size_field_name="target_entity_sizes",
        graph_index=graph_index,
    )
    if group is None or indices.numel() == 0:
        return indices
    group_ids = self.graph_target_entity_group_ids(graph_index)
    mask = self._group_mask(
        group_ids, group=group, attr_name="target_entity_groups"
    )
    return indices[mask]

graph_target_entity_group_ids(graph_index=0)

Return source-group ids for target-entity rows.

Source code in src/mifrost/encoders/flat_data.py

def graph_target_entity_group_ids(self, graph_index: int = 0) -> torch.Tensor:
    """Return source-group ids for target-entity rows."""
    return self._graph_cat_field_slice(
        field_name="target_entity_group_ids",
        size_field_name="target_entity_sizes",
        graph_index=graph_index,
    )

graph_target_entity_names(graph_index=0, group=None)

Return names for target-entity rows, optionally filtered by group.

Source code in src/mifrost/encoders/flat_data.py

def graph_target_entity_names(
    self, graph_index: int = 0, group: str | int | None = None
) -> list[str]:
    """Return names for target-entity rows, optionally filtered by group."""
    target_entity_indices = self.graph_target_entity_indices(
        graph_index, group=group
    )
    if target_entity_indices.numel() == 0:
        return []
    start, _end = self.graph_node_range(graph_index)
    local_names = self.graph_node_names(graph_index)
    return [
        str(local_names[int(global_idx.item()) - start])
        for global_idx in target_entity_indices
    ]

graph_target_positions(graph_index=0)

Return entity-row positions for prediction targets.

Source code in src/mifrost/encoders/flat_data.py

def graph_target_positions(self, graph_index: int = 0) -> torch.Tensor:
    """Return entity-row positions for prediction targets."""
    return self._graph_cat_field_slice(
        field_name="target_positions",
        size_field_name="target_sizes",
        graph_index=graph_index,
    )

graph_target_indices(graph_index=0)

Return per-graph target indices in encounter order.

Source code in src/mifrost/encoders/flat_data.py

def graph_target_indices(self, graph_index: int = 0) -> torch.Tensor:
    """Return per-graph target indices in encounter order."""
    return self._graph_cat_field_slice(
        field_name="target_indices",
        size_field_name="target_sizes",
        graph_index=graph_index,
    )

graph_target_candidate_ids(graph_index=0)

Return stable candidate ids for prediction targets.

Source code in src/mifrost/encoders/flat_data.py

def graph_target_candidate_ids(self, graph_index: int = 0) -> torch.Tensor:
    """Return stable candidate ids for prediction targets."""
    return self._graph_cat_field_slice(
        field_name="target_candidate_ids",
        size_field_name="target_sizes",
        graph_index=graph_index,
    )

graph_target_depths(graph_index=0)

Return target depths when the encoder emitted them.

Source code in src/mifrost/encoders/flat_data.py

def graph_target_depths(self, graph_index: int = 0) -> torch.Tensor:
    """Return target depths when the encoder emitted them."""
    return self._graph_cat_field_slice(
        field_name="target_depths",
        size_field_name="target_sizes",
        graph_index=graph_index,
    )

graph_target_group_ids(graph_index=0)

Return source-group ids for prediction targets.

Source code in src/mifrost/encoders/flat_data.py

def graph_target_group_ids(self, graph_index: int = 0) -> torch.Tensor:
    """Return source-group ids for prediction targets."""
    return self._graph_cat_field_slice(
        field_name="target_group_ids",
        size_field_name="target_sizes",
        graph_index=graph_index,
    )

graph_target_names(graph_index=0)

Return display names for prediction targets.

Source code in src/mifrost/encoders/flat_data.py

def graph_target_names(self, graph_index: int = 0) -> list[str]:
    """Return display names for prediction targets."""
    target_names = getattr(self, "target_names", None)
    if target_names is None:
        return []
    if target_names and isinstance(target_names[0], (list, tuple)):
        return [str(name) for name in target_names[graph_index]]
    return [str(name) for name in target_names]

graph_relation_instance_range(graph_index=0)

Return the global relation-instance index range for one graph.

Source code in src/mifrost/encoders/flat_data.py

def graph_relation_instance_range(self, graph_index: int = 0) -> tuple[int, int]:
    """Return the global relation-instance index range for one graph."""
    relation_instance_sizes = getattr(self, "relation_instance_sizes", None)
    if relation_instance_sizes is None:
        counts = getattr(self, "relation_counts", None)
        if counts is None:
            return (0, 0)
        counts = counts.long()
        if counts.dim() == 1:
            total = int(counts.sum().item())
            return (0, total)
        if graph_index < 0 or graph_index >= counts.size(0):
            raise IndexError(
                f"graph_index {graph_index} out of range for {counts.size(0)} graphs"
            )
        sizes = counts.sum(dim=1)
    else:
        sizes = relation_instance_sizes.long().view(-1)
        if graph_index < 0 or graph_index >= len(sizes):
            raise IndexError(
                f"graph_index {graph_index} out of range for {len(sizes)} graphs"
            )
    start = int(sizes[:graph_index].sum().item()) if graph_index > 0 else 0
    end = start + int(sizes[graph_index].item())
    return start, end

graph_lgan_tn_edges(graph_index=0)

Return packed TN LGAN edges for one graph.

Source code in src/mifrost/encoders/flat_data.py

def graph_lgan_tn_edges(self, graph_index: int = 0) -> torch.Tensor:
    """Return packed TN LGAN edges for one graph."""
    return self._graph_edge_slice(
        src_field_name="lgan_tn_relation_indices",
        dst_field_name="lgan_tn_entity_indices",
        size_field_name="lgan_tn_sizes",
        graph_index=graph_index,
    )

graph_lgan_nn_edges(graph_index=0)

Return packed NN LGAN edges for one graph.

Source code in src/mifrost/encoders/flat_data.py

def graph_lgan_nn_edges(self, graph_index: int = 0) -> torch.Tensor:
    """Return packed NN LGAN edges for one graph."""
    return self._graph_edge_slice(
        src_field_name="lgan_nn_relation_indices",
        dst_field_name="lgan_nn_entity_indices",
        size_field_name="lgan_nn_sizes",
        graph_index=graph_index,
    )

graph_lgan_rr_edges(graph_index=0)

Return packed RR LGAN edges for one graph.

Source code in src/mifrost/encoders/flat_data.py

def graph_lgan_rr_edges(self, graph_index: int = 0) -> torch.Tensor:
    """Return packed RR LGAN edges for one graph."""
    return self._graph_edge_slice(
        src_field_name="lgan_rr_src_relation_indices",
        dst_field_name="lgan_rr_dst_relation_indices",
        size_field_name="lgan_rr_sizes",
        graph_index=graph_index,
    )

graph_node_range(graph_index=0)

Return the global entity-row range for one graph.

Source code in src/mifrost/encoders/flat_data.py

def graph_node_range(self, graph_index: int = 0) -> tuple[int, int]:
    """Return the global entity-row range for one graph."""
    node_sizes = getattr(self, "node_sizes", None)
    if node_sizes is None:
        return (0, int(getattr(self, "num_nodes", 0)))
    node_sizes = node_sizes.long().view(-1)
    if graph_index < 0 or graph_index >= len(node_sizes):
        raise IndexError(
            f"graph_index {graph_index} out of range for {len(node_sizes)} graphs"
        )
    start = int(node_sizes[:graph_index].sum().item()) if graph_index > 0 else 0
    end = start + int(node_sizes[graph_index].item())
    return start, end

graph_entity_role_ids(graph_index=0)

Return per-entity role ids for one graph.

Source code in src/mifrost/encoders/flat_data.py

def graph_entity_role_ids(self, graph_index: int = 0) -> torch.Tensor:
    """Return per-entity role ids for one graph."""
    start, end = self.graph_node_range(graph_index)
    entity_role_ids = getattr(self, "entity_role_ids", None)
    if entity_role_ids is None:
        return torch.empty((0,), dtype=torch.long)
    return entity_role_ids.long().view(-1)[start:end]

graph_entity_roles(graph_index=0)

Return decoded per-entity role labels for one graph.

Source code in src/mifrost/encoders/flat_data.py

def graph_entity_roles(self, graph_index: int = 0) -> list[str]:
    """Return decoded per-entity role labels for one graph."""
    role_names = [
        str(value) for value in getattr(self, "entity_role_names", []) or []
    ]
    out: list[str] = []
    for role_id in self.graph_entity_role_ids(graph_index).tolist():
        if 0 <= role_id < len(role_names):
            out.append(role_names[role_id])
        else:
            out.append(str(role_id))
    return out