ColorEncoder

ColorEncoder

Bases: EncoderBase[Data]

Homogeneous color encoder backed by ColorEncoderEngine.

Produces compact Data/Batch outputs with integer-like node/edge attributes suitable for color-based graph models.

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

class ColorEncoder(EncoderBase[Data]):
    """
    Homogeneous color encoder backed by ``ColorEncoderEngine``.

    Produces compact ``Data``/``Batch`` outputs with integer-like node/edge
    attributes suitable for color-based graph models.
    """

    def __init__(
        self,
        domain: DomainInput,
        *,
        edge_features: bool = False,
        enable_global_predicate_nodes: bool = False,
    ) -> None:
        """Create a color encoder for one domain."""
        config = ColorEncoderConfig()
        config.edge_features = edge_features
        config.enable_global_predicate_nodes = enable_global_predicate_nodes
        self._engine = ColorEncoderEngine(_advanced_domain(domain), config)
        self.edge_features = edge_features
        self.predicate_nodes_enabled = enable_global_predicate_nodes

    @property
    def engine(self) -> ColorEncoderEngine:
        """Expose the underlying C++ color engine."""
        return self._engine

    def _encode(
        self,
        state: StateInput,
        *,
        goals: GoalBatchInput = None,
        actions: ActionBatchInput = None,
        subgoal_layers: SubgoalLayersInput = None,
    ) -> HomoEncoding:
        """Encode one state into homogeneous encoding dictionary."""
        adv_state = _advanced_state(state)
        action_inputs = parse_flat_actions(actions)
        action_list = _prepare_actions(action_inputs)
        if goals is None and subgoal_layers is None and not action_list:
            return self._engine.encode(adv_state)
        if goals is None:
            goals = default_goals_from_state(state)
        inputs = _split_goals(goals, subgoal_layers)
        return self._engine.encode(adv_state, inputs, action_list)

    def encode(
        self,
        state: StateInput,
        *,
        goals: GoalBatchInput = None,
        actions: ActionBatchInput = None,
        subgoal_layers: SubgoalLayersInput = None,
        include_metadata: bool = True,
        **kwargs,
    ) -> HomoEncoding:
        """Encode one state into native ``BatchEncoding``."""
        return super().encode(
            state,
            goals=goals,
            actions=actions,
            subgoal_layers=subgoal_layers,
            include_metadata=include_metadata,
            **kwargs,
        )

    def _encode_batch(
        self,
        states: StateBatchInput,
        *,
        goals: GoalBatchParam = None,
        actions: ActionBatchParam = None,
        subgoal_layers: SubgoalLayersBatchParam = None,
    ) -> HomoEncoding:
        """Encode one or many states into homogeneous batch encoding_dict."""
        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,
        )
        return self._engine.encode_batch(
            states_for_core,
            goals=goals_for_core,
            actions=actions_for_core,
            subgoal_layers=subgoal_layers_for_core,
        )

    def encode_batch(
        self,
        states: StateBatchInput,
        *,
        goals: GoalBatchParam = None,
        actions: ActionBatchParam = None,
        subgoal_layers: SubgoalLayersBatchParam = None,
        batch_attrs: Mapping[str, Any] | None = None,
        collate_spec: CollateSpecParam = None,
        include_metadata: bool = True,
        **kwargs,
    ) -> HomoEncoding:
        """Encode one or many states into native ``BatchEncoding``."""
        return super().encode_batch(
            states,
            goals=goals,
            actions=actions,
            subgoal_layers=subgoal_layers,
            batch_attrs=batch_attrs,
            collate_spec=collate_spec,
            include_metadata=include_metadata,
            **kwargs,
        )

    def stream(self) -> ColorEncoderStream:
        """Create a streaming encoder sharing this encoder's C++ engine."""
        return ColorEncoderStream(self)

    def to_networkx(self, data: Data) -> nx.Graph:
        """Convert a color-encoded PyG graph into a NetworkX graph."""
        graph = nx.Graph()
        node_names = getattr(data, "node_names", None)
        if not node_names:
            if hasattr(data, "x") and data.x is not None:
                count = data.x.shape[0]
            else:
                count = data.num_nodes
            node_names = [str(i) for i in range(count)]

        has_scalar_x = (
            hasattr(data, "x")
            and data.x is not None
            and torch.is_tensor(data.x)
            and data.x.dim() == 2
            and data.x.size(1) > 0
        )
        for i, name in enumerate(node_names):
            val = data.x[i] if has_scalar_x else 0
            attrs = {"type": val.item() if torch.is_tensor(val) else val}
            if hasattr(data, "goal_level") and data.goal_level is not None:
                gval = data.goal_level[i]
                attrs["goal_level"] = gval.item() if torch.is_tensor(gval) else gval
            graph.add_node(name, **attrs)

        for i in range(data.edge_index.shape[1]):
            u_idx = data.edge_index[0, i].item()
            v_idx = data.edge_index[1, i].item()
            u_name = node_names[u_idx]
            v_name = node_names[v_idx]
            attrs = {}
            if hasattr(data, "edge_attr") and data.edge_attr is not None:
                val = data.edge_attr[i]
                attrs["type"] = int(val.item()) if torch.is_tensor(val) else int(val)
            graph.add_edge(u_name, v_name, **attrs)

        graph.graph["encoder_hash"] = getattr(data, "encoder_hash", None)
        return graph

    def draw(
        self,
        data: Data,
        *,
        with_labels: bool = True,
        edge_labels: bool = False,
        ax: Any | None = None,
        node_size: int = 300,
        font_size: int = 8,
    ) -> Any:
        """Render a color-encoded graph with matplotlib and return the axis."""
        try:
            import matplotlib.pyplot as plt
        except ModuleNotFoundError as exc:
            raise RuntimeError(
                "ColorEncoder.draw requires matplotlib to be installed"
            ) from exc

        graph = self.to_networkx(data)
        if ax is None:
            _, ax = plt.subplots()

        positions = nx.spring_layout(graph, seed=0)
        nx.draw_networkx(
            graph,
            pos=positions,
            ax=ax,
            with_labels=with_labels,
            node_size=node_size,
            font_size=font_size,
        )

        if edge_labels:
            labels = {}
            for src, dst, attrs in graph.edges(data=True):
                if "type" in attrs:
                    labels[(src, dst)] = attrs["type"]
            if labels:
                nx.draw_networkx_edge_labels(
                    graph,
                    pos=positions,
                    edge_labels=labels,
                    ax=ax,
                    font_size=max(6, font_size - 1),
                )

        return ax

engine property

Expose the underlying C++ color engine.

__init__(domain, *, edge_features=False, enable_global_predicate_nodes=False)

Create a color encoder for one domain.

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

def __init__(
    self,
    domain: DomainInput,
    *,
    edge_features: bool = False,
    enable_global_predicate_nodes: bool = False,
) -> None:
    """Create a color encoder for one domain."""
    config = ColorEncoderConfig()
    config.edge_features = edge_features
    config.enable_global_predicate_nodes = enable_global_predicate_nodes
    self._engine = ColorEncoderEngine(_advanced_domain(domain), config)
    self.edge_features = edge_features
    self.predicate_nodes_enabled = enable_global_predicate_nodes

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

Encode one state into native BatchEncoding.

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

def encode(
    self,
    state: StateInput,
    *,
    goals: GoalBatchInput = None,
    actions: ActionBatchInput = None,
    subgoal_layers: SubgoalLayersInput = None,
    include_metadata: bool = True,
    **kwargs,
) -> HomoEncoding:
    """Encode one state into native ``BatchEncoding``."""
    return super().encode(
        state,
        goals=goals,
        actions=actions,
        subgoal_layers=subgoal_layers,
        include_metadata=include_metadata,
        **kwargs,
    )

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

Encode one or many states into native BatchEncoding.

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

def encode_batch(
    self,
    states: StateBatchInput,
    *,
    goals: GoalBatchParam = None,
    actions: ActionBatchParam = None,
    subgoal_layers: SubgoalLayersBatchParam = None,
    batch_attrs: Mapping[str, Any] | None = None,
    collate_spec: CollateSpecParam = None,
    include_metadata: bool = True,
    **kwargs,
) -> HomoEncoding:
    """Encode one or many states into native ``BatchEncoding``."""
    return super().encode_batch(
        states,
        goals=goals,
        actions=actions,
        subgoal_layers=subgoal_layers,
        batch_attrs=batch_attrs,
        collate_spec=collate_spec,
        include_metadata=include_metadata,
        **kwargs,
    )

stream()

Create a streaming encoder sharing this encoder's C++ engine.

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

def stream(self) -> ColorEncoderStream:
    """Create a streaming encoder sharing this encoder's C++ engine."""
    return ColorEncoderStream(self)

to_networkx(data)

Convert a color-encoded PyG graph into a NetworkX graph.

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

def to_networkx(self, data: Data) -> nx.Graph:
    """Convert a color-encoded PyG graph into a NetworkX graph."""
    graph = nx.Graph()
    node_names = getattr(data, "node_names", None)
    if not node_names:
        if hasattr(data, "x") and data.x is not None:
            count = data.x.shape[0]
        else:
            count = data.num_nodes
        node_names = [str(i) for i in range(count)]

    has_scalar_x = (
        hasattr(data, "x")
        and data.x is not None
        and torch.is_tensor(data.x)
        and data.x.dim() == 2
        and data.x.size(1) > 0
    )
    for i, name in enumerate(node_names):
        val = data.x[i] if has_scalar_x else 0
        attrs = {"type": val.item() if torch.is_tensor(val) else val}
        if hasattr(data, "goal_level") and data.goal_level is not None:
            gval = data.goal_level[i]
            attrs["goal_level"] = gval.item() if torch.is_tensor(gval) else gval
        graph.add_node(name, **attrs)

    for i in range(data.edge_index.shape[1]):
        u_idx = data.edge_index[0, i].item()
        v_idx = data.edge_index[1, i].item()
        u_name = node_names[u_idx]
        v_name = node_names[v_idx]
        attrs = {}
        if hasattr(data, "edge_attr") and data.edge_attr is not None:
            val = data.edge_attr[i]
            attrs["type"] = int(val.item()) if torch.is_tensor(val) else int(val)
        graph.add_edge(u_name, v_name, **attrs)

    graph.graph["encoder_hash"] = getattr(data, "encoder_hash", None)
    return graph

draw(data, *, with_labels=True, edge_labels=False, ax=None, node_size=300, font_size=8)

Render a color-encoded graph with matplotlib and return the axis.

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

def draw(
    self,
    data: Data,
    *,
    with_labels: bool = True,
    edge_labels: bool = False,
    ax: Any | None = None,
    node_size: int = 300,
    font_size: int = 8,
) -> Any:
    """Render a color-encoded graph with matplotlib and return the axis."""
    try:
        import matplotlib.pyplot as plt
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "ColorEncoder.draw requires matplotlib to be installed"
        ) from exc

    graph = self.to_networkx(data)
    if ax is None:
        _, ax = plt.subplots()

    positions = nx.spring_layout(graph, seed=0)
    nx.draw_networkx(
        graph,
        pos=positions,
        ax=ax,
        with_labels=with_labels,
        node_size=node_size,
        font_size=font_size,
    )

    if edge_labels:
        labels = {}
        for src, dst, attrs in graph.edges(data=True):
            if "type" in attrs:
                labels[(src, dst)] = attrs["type"]
        if labels:
            nx.draw_networkx_edge_labels(
                graph,
                pos=positions,
                edge_labels=labels,
                ax=ax,
                font_size=max(6, font_size - 1),
            )

    return ax