mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-27 17:23:19 +00:00
7ae2942c34
Removed the `sat_utils.` that was left there from the original zebra package. This allows the zebra generator to be used with the full set of available clues.
269 lines
8.1 KiB
Python
269 lines
8.1 KiB
Python
"""
|
|
clues.py
|
|
|
|
These are all the clue types that a puzzle can have. Things like "the tea drinker lives in the
|
|
green house" and "the cat owner lives left of the person who likes grilled cheese."
|
|
|
|
There's a Clue ABC that requires you implement an `as_cnf` method, to convert the clue to an
|
|
and-of-ors (probably using things defined in `sat_utils`), and a human-readable __repr__ that
|
|
can be used in a puzzle description.
|
|
|
|
"""
|
|
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass, field
|
|
from functools import wraps
|
|
from itertools import product
|
|
from typing import Iterable
|
|
|
|
from reasoning_gym.logic.contrib.logic_puzzle.literals import Literal
|
|
from reasoning_gym.logic.contrib.logic_puzzle.sat_utils import from_dnf, neg
|
|
|
|
|
|
def _capitalize_first(repr_func):
|
|
"""
|
|
Decorator for a __repr__ function that capitalizes the first letter without chagning the rest
|
|
|
|
(in contrast to str.capitalize(), which capitalizes the first letter and makes the rest lower)
|
|
"""
|
|
|
|
@wraps(repr_func)
|
|
def wrapper(*args, **kwargs):
|
|
output = repr_func(*args, **kwargs)
|
|
return output[0].upper() + output[1:]
|
|
|
|
return wrapper
|
|
|
|
|
|
class Clue(ABC):
|
|
"""Base class for the types of clues that we allow."""
|
|
|
|
@abstractmethod
|
|
def as_cnf(self) -> Iterable[tuple[str]]: ...
|
|
|
|
@abstractmethod
|
|
def __repr__(self) -> str: ...
|
|
|
|
def __lt__(self, other) -> bool:
|
|
return str(self) < str(other)
|
|
|
|
|
|
def comb(value: Literal, house: int) -> str:
|
|
"""Format how a value is shown at a given house"""
|
|
|
|
return f"{value} {house}"
|
|
|
|
|
|
@dataclass(eq=True, frozen=True)
|
|
class found_at(Clue):
|
|
"""
|
|
A literal is known to be at a specific house
|
|
|
|
Examples:
|
|
- the tea drinker lives in the middle house
|
|
- the fourth house is red
|
|
"""
|
|
|
|
value: Literal
|
|
house: int
|
|
|
|
def as_cnf(self) -> list[tuple[str]]:
|
|
return [(comb(self.value, self.house),)]
|
|
|
|
@_capitalize_first
|
|
def __repr__(self) -> str:
|
|
houses = [None, "first", "second", "third", "fourth", "fifth", "sixth", "seventh"]
|
|
return f"{self.value.value} is in the {houses[self.house]} house."
|
|
|
|
|
|
@dataclass(eq=True, frozen=True)
|
|
class not_at(Clue):
|
|
"""
|
|
Two values are known *not* to be at the same house
|
|
|
|
Examples:
|
|
- the musician does not drink tea
|
|
- the red house does not contain a cat
|
|
"""
|
|
|
|
value: Literal
|
|
house: int
|
|
|
|
def as_cnf(self) -> list[tuple[str]]:
|
|
return [(neg(comb(self.value, self.house)),)]
|
|
|
|
@_capitalize_first
|
|
def __repr__(self) -> str:
|
|
houses = [None, "first", "second", "third", "fourth", "fifth", "sixth", "seventh"]
|
|
return f"{self.value.value} is not in the {houses[self.house]} house."
|
|
|
|
|
|
@dataclass(eq=True, frozen=True)
|
|
class same_house(Clue):
|
|
"""
|
|
Two values are known to be at the same house
|
|
|
|
Examples:
|
|
- the musician drinks tea
|
|
- the red house contains a cat
|
|
"""
|
|
|
|
value1: Literal
|
|
value2: Literal
|
|
houses: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
|
|
|
|
def as_cnf(self) -> list[tuple[str]]:
|
|
return from_dnf((comb(self.value1, i), comb(self.value2, i)) for i in self.houses)
|
|
|
|
@_capitalize_first
|
|
def __repr__(self) -> str:
|
|
return f"{self.value1.value} is {self.value2.value}."
|
|
|
|
|
|
@dataclass(eq=True, frozen=True)
|
|
class consecutive(Clue):
|
|
"""
|
|
The first value is directly to the left of the second value
|
|
|
|
Examples:
|
|
- the green house is directly to the left of the white house
|
|
(green in 1, white in 2 OR green in 2, white in 3 OR etc.)
|
|
- the house with the kittens is directly to the right of the tea drinker's home
|
|
(kittens in 2, tea in 1 OR kittens in 3, tea in 2 OR etc.)
|
|
"""
|
|
|
|
value1: Literal
|
|
value2: Literal
|
|
houses: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
|
|
|
|
def as_cnf(self) -> list[tuple[str]]:
|
|
return from_dnf((comb(self.value1, i), comb(self.value2, j)) for i, j in zip(self.houses, self.houses[1:]))
|
|
|
|
@_capitalize_first
|
|
def __repr__(self) -> str:
|
|
return f"{self.value1.value} is directly left of {self.value2.value}."
|
|
|
|
|
|
@dataclass(eq=True, frozen=True)
|
|
class beside(Clue):
|
|
"""
|
|
The two values occur side-by-side (either left or right)
|
|
|
|
Examples:
|
|
- the coffee drinker is (left or right) of the tea drinker
|
|
- the cat owner is (left or right) of the green house
|
|
"""
|
|
|
|
value1: Literal
|
|
value2: Literal
|
|
houses: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
|
|
|
|
def as_cnf(self) -> list[tuple[str]]:
|
|
return from_dnf(
|
|
[(comb(self.value1, i), comb(self.value2, j)) for i, j in zip(self.houses, self.houses[1:])]
|
|
+ [(comb(self.value2, i), comb(self.value1, j)) for i, j in zip(self.houses, self.houses[1:])]
|
|
)
|
|
|
|
@_capitalize_first
|
|
def __repr__(self) -> str:
|
|
return f"{self.value1.value} and {self.value2.value} are next to each other."
|
|
|
|
|
|
@dataclass(eq=True, frozen=True)
|
|
class left_of(Clue):
|
|
"""
|
|
The first value is somewhere to the left of the second value
|
|
|
|
Examples:
|
|
- the tea drinker is in house 1 and the musician in 2, 3, 4, or 5;
|
|
OR the tea drinker in 2, and musician in 3, 4, or 5;
|
|
OR the tea drinker in 3, musician in 4, 5; OR tea 4, musician 5.
|
|
"""
|
|
|
|
value1: Literal
|
|
value2: Literal
|
|
houses: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
|
|
|
|
def as_cnf(self) -> list[tuple[str]]:
|
|
return from_dnf(
|
|
(comb(self.value1, i), comb(self.value2, j)) for i, j in product(self.houses, self.houses) if i < j
|
|
)
|
|
|
|
@_capitalize_first
|
|
def __repr__(self) -> str:
|
|
return f"{self.value1.value} is somewhere to the left of {self.value2.value}."
|
|
|
|
|
|
@dataclass(eq=True, frozen=True)
|
|
class right_of(Clue):
|
|
"""
|
|
The first value is somewhere to the right of the second value.
|
|
|
|
Examples:
|
|
- the coffee drinker is in house 5 and the artist in 1, 2, 3, 4;
|
|
OR the coffee drinker in 4, and artist in 1, 2, or 3;
|
|
OR the coffee drinker in 3, artist in 1, 2; OR coffee 2, artist 1.
|
|
"""
|
|
|
|
value1: Literal
|
|
value2: Literal
|
|
houses: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
|
|
|
|
def as_cnf(self) -> list[tuple[str]]:
|
|
return from_dnf(
|
|
(comb(self.value1, i), comb(self.value2, j)) for i, j in product(self.houses, self.houses) if i > j
|
|
)
|
|
|
|
@_capitalize_first
|
|
def __repr__(self) -> str:
|
|
return f"{self.value1.value} is somewhere to the right of {self.value2.value}."
|
|
|
|
|
|
@dataclass(eq=True, frozen=True)
|
|
class one_between(Clue):
|
|
"""
|
|
The values are separated by one house
|
|
|
|
Examples (if 5 houses):
|
|
- the cat is in house 1 and tea drinker in house 3; OR cat 2, tea 4;
|
|
OR cat 4 house 5
|
|
- the green house is #1 and the musician in house 3; or green house 2, musician 4;
|
|
OR green house 3, musician 5.
|
|
"""
|
|
|
|
value1: Literal
|
|
value2: Literal
|
|
houses: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
|
|
|
|
def as_cnf(self) -> list[tuple[str]]:
|
|
return from_dnf(
|
|
[(comb(self.value1, i), comb(self.value2, j)) for i, j in zip(self.houses, self.houses[2:])]
|
|
+ [(comb(self.value2, i), comb(self.value1, j)) for i, j in zip(self.houses, self.houses[2:])]
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"There is one house between {self.value1.value} and {self.value2.value}."
|
|
|
|
|
|
@dataclass(eq=True, frozen=True)
|
|
class two_between(Clue):
|
|
"""
|
|
The values are separated by two houses
|
|
|
|
Examples (if 5 houses):
|
|
- the cat is in house 1 and artist in house 4; or cat 2, artist 5
|
|
- the dog is in house 1 and red house is #4; or dog 2, red house 5
|
|
"""
|
|
|
|
value1: Literal
|
|
value2: Literal
|
|
houses: tuple[int, ...] = field(default_factory=lambda: (1, 2, 3, 4, 5))
|
|
|
|
def as_cnf(self) -> list[tuple[str]]:
|
|
return from_dnf(
|
|
[(comb(self.value1, i), comb(self.value2, j)) for i, j in zip(self.houses, self.houses[3:])]
|
|
+ [(comb(self.value2, i), comb(self.value1, j)) for i, j in zip(self.houses, self.houses[3:])]
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"There are two houses between {self.value1.value} and {self.value2.value}."
|