Rule Systems¶

def binary_tensors(
    self,
    device: torch.device | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Return binary rules as (results, lefts, rights) index tensors."""
    if not self.binary_rules:
        empty = torch.zeros(0, dtype=torch.long, device=device)
        return empty, empty, empty

    results, lefts, rights = zip(*self.binary_rules)

    return (
        torch.tensor(results, dtype=torch.long, device=device),
        torch.tensor(lefts, dtype=torch.long, device=device),
        torch.tensor(rights, dtype=torch.long, device=device),
    )

def binary_weight_tensor(
    self,
    device: torch.device | None = None,
) -> torch.Tensor:
    """Return initial binary rule weights as a float tensor.

    Returns zeros when no explicit weights are set.
    """
    if self.binary_weights is not None:
        return torch.tensor(
            self.binary_weights,
            dtype=torch.float,
            device=device,
        )

    return torch.zeros(self.n_binary, dtype=torch.float, device=device)

def unary_tensors(
    self,
    device: torch.device | None = None,
) -> tuple[torch.Tensor, torch.Tensor] | None:
    """Return unary rules as (results, inputs) index tensors, or None."""
    if not self.unary_rules:
        return None

    results, inputs = zip(*self.unary_rules)

    return (
        torch.tensor(results, dtype=torch.long, device=device),
        torch.tensor(inputs, dtype=torch.long, device=device),
    )

def unary_weight_tensor(
    self,
    device: torch.device | None = None,
) -> torch.Tensor:
    """Return initial unary rule weights as a float tensor.

    Returns zeros when no explicit weights are set.
    """
    if self.unary_weights is not None:
        return torch.tensor(
            self.unary_weights,
            dtype=torch.float,
            device=device,
        )

    return torch.zeros(self.n_unary, dtype=torch.float, device=device)

def __add__(self, other: "RuleSystem") -> "RuleSystem":
    """Merge two rule systems (deduplicated union)."""
    if self.n_categories != other.n_categories:
        raise ValueError(
            f"cannot merge rule systems with different category "
            f"counts: {self.n_categories} vs {other.n_categories}"
        )

    seen_binary: dict[tuple[int, ...], float] = {}
    for i, rule in enumerate(self.binary_rules):
        w = self.binary_weights[i] if self.binary_weights else 0.0
        seen_binary[rule] = w
    for i, rule in enumerate(other.binary_rules):
        if rule not in seen_binary:
            w = other.binary_weights[i] if other.binary_weights else 0.0
            seen_binary[rule] = w

    seen_unary: dict[tuple[int, ...], float] = {}
    for i, rule in enumerate(self.unary_rules):
        w = self.unary_weights[i] if self.unary_weights else 0.0
        seen_unary[rule] = w
    for i, rule in enumerate(other.unary_rules):
        if rule not in seen_unary:
            w = other.unary_weights[i] if other.unary_weights else 0.0
            seen_unary[rule] = w

    binary_rules = tuple(seen_binary.keys())
    unary_rules = tuple(seen_unary.keys())

    has_binary_w = (
        self.binary_weights is not None or other.binary_weights is not None
    )
    has_unary_w = self.unary_weights is not None or other.unary_weights is not None

    binary_weights = tuple(seen_binary.values()) if has_binary_w else None
    unary_weights = tuple(seen_unary.values()) if has_unary_w else None

    desc_parts = []
    if self.description:
        desc_parts.append(self.description)
    if other.description:
        desc_parts.append(other.description)

    return RuleSystem(
        binary_rules=binary_rules,
        unary_rules=unary_rules,
        n_categories=self.n_categories,
        description=" + ".join(desc_parts) if desc_parts else "",
        binary_weights=binary_weights,
        unary_weights=unary_weights,
    )