[1]:
%pip install -q -U circuitree[distributed]==0.11.1
Note: you may need to restart the kernel to use updated packages.

Defining custom Grammars using CircuitGrammar

For search spaces that can’t expressed using a grammar in circuitree.models, you can make a custom grammar. A grammar will have rules for how to represent each distinct state in the design space and how different actions affect the state.

You can define a grammar by subclassing CircuitGrammar. CircuitGrammar is an abstract class, which means that we must define certain methods in order to use it. Here’s a description of each method’s call signature and what it should do.

is_terminal(state) -> bool       # Return whether or not this state is terminal
get_actions(state) -> list[str]  # Return a list of actions that can be taken from this state
do_action(state, action) -> str  # Return a new state as a result of making this move
get_unique_state(state) -> str   # Return a unique representation of this state

Together, these functions allow us to explore the space of possible circuits on-the-fly, without needing to enumerate every possibility.

As an example, let’s consider an existing design space, explored by Angela Chau and colleagues in this seminal paper in their study of single-cell polarization circuits. The authors were studying all possible network topologies that could be made with two membrane-bound enzymes A and B that have active and inactive forms. Each species can catalyze either the forward or reverse reaction of the other species and of itself (autocatalysis). In other words, any of the four pairwise interactions can be present, and they can be either activating or inactivating (inhibiting).

First, let’s decide how we could represent a circuit assembly as a string of characters that we will call the state string. (Any Hashable representation can be used, but strings are convenient.) Let’s use the following convention, which is the same one used by SimpleNetworkGrammar: * Each component is an uppercase letter. A refers to enzyme A and B refers to enzyme B. * Each pairwise interaction is represented by three characters, two uppercase letters for the components and one lowercase letter referring to the type of interaction. - For instance, ABa means “A activates B”, and BBi means “B inhibits B”. * The state string consists of two parts. separated by ::. - The left side says which components are present (AB). - The right side says which interactions are present, separated by underscores (ABa_BBi). - A terminal state starts with *, denoting that assembly is complete and the game is over. * The allowed actions are - Adding a new interaction - Terminating * When adding multiple interactions, the order does not matter

For instance, a circuit where A and B both inhibit each other can be written as the state string *AB::ABi_BAi - a fully assembled circuit (*) with components A and B (AB) that inhibit each other (ABi_BAi). Since order does not matter, we want get_unique_state("*AB::ABi_BAi") and get_unique_state("*AB::BAi_ABi") to return the same string. We can achieve this by sorting the interactions alphabetically.

[2]:
from circuitree import CircuitGrammar

class PolarizationGrammar(CircuitGrammar):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def is_terminal(self, state) -> bool:
        return state.startswith("*")

    def get_actions(self, state: str) -> list[str]:
        # If termination hasn't happened yet, it's always a possibiilty
        if self.is_terminal(state):
            return []
        actions = ["*terminate*"]

        # Get the part of the string that contains the interactions
        interactions = state.split("::")[1]

        # We can add an interaction between any unused pair
        for pair in ("AA", "AB", "BA", "BB"):
            if pair not in interactions:
                actions.append(pair + "a")
                actions.append(pair + "i")
        return actions

    def do_action(self, state, action):
        if action == "*terminate*":
            return "*" + state

        prefix, interactions = state.split("::")
        if len(interactions) == 0:
            return f"{prefix}::{action}"
        else:
            return f"{prefix}::{interactions}_{action}"

    def get_unique_state(self, state):
        prefix, interactions = state.split("::")
        if len(interactions) == 0:
            return state # No interactions to sort
        else:
            interactions_list = interactions.split("_")
            interactions = "_".join(sorted(interactions_list))
            return f"{prefix}::{interactions}"

Here are some example inputs to each function and the expected outputs.

>>> grammar = PolarizationGrammar()
>>> grammar.is_terminal("*AB::AAa")              # True
>>> grammar.get_actions("AB::ABa_BAi")           # ['*terminate*', 'AAa', 'AAi', 'BBa', 'BBi']
>>> grammar.do_action("AB::ABa_BAi", "AAa")      # 'AB::ABa_BAi_AAa'
>>> grammar.get_unique_state("*AB::ABa_BAi_AAa") # '*AB::AAa_ABa_BAi'

The original paper enumerated all the possible topologies in this space and found 81. We can check our work by expanding all the possible state strings and printing the terminal ones. There should be 81. To do this, we can use the function CircuiTree.grow_tree(), which recursively builds the entire decision tree of possible states.

[3]:
from circuitree import CircuiTree

class PolarizationTree(CircuiTree):
    def __init__(self, *args, **kwargs):

        # Specify the required keyword arguments
        kwargs = {"grammar": PolarizationGrammar(), "root": "AB::"} | kwargs
        super().__init__(*args, **kwargs)

    def get_reward(self, state):
        raise NotImplementedError # no need to implement this


from pprint import pprint

tree = PolarizationTree()
tree.grow_tree()

terminal_states = list(tree.terminal_states)
print(f"# terminal states: {len(terminal_states)}")
print()
print("All terminal topologies:")
pprint(terminal_states)
# terminal states: 81

All terminal topologies:
['*AB::AAi_ABi_BAi_BBi',
 '*AB::AAa_ABi_BAi_BBi',
 '*AB::ABi_BAi_BBi',
 '*AB::AAi_ABa_BAi_BBi',
 '*AB::AAa_ABa_BAi_BBi',
 '*AB::ABa_BAi_BBi',
 '*AB::AAi_BAi_BBi',
 '*AB::AAa_BAi_BBi',
 '*AB::BAi_BBi',
 '*AB::AAi_ABi_BAa_BBi',
 '*AB::AAa_ABi_BAa_BBi',
 '*AB::ABi_BAa_BBi',
 '*AB::AAi_ABa_BAa_BBi',
 '*AB::AAa_ABa_BAa_BBi',
 '*AB::ABa_BAa_BBi',
 '*AB::AAi_BAa_BBi',
 '*AB::AAa_BAa_BBi',
 '*AB::BAa_BBi',
 '*AB::AAi_ABi_BBi',
 '*AB::AAa_ABi_BBi',
 '*AB::ABi_BBi',
 '*AB::AAi_ABa_BBi',
 '*AB::AAa_ABa_BBi',
 '*AB::ABa_BBi',
 '*AB::AAi_BBi',
 '*AB::AAa_BBi',
 '*AB::BBi',
 '*AB::AAi_ABi_BAi_BBa',
 '*AB::AAa_ABi_BAi_BBa',
 '*AB::ABi_BAi_BBa',
 '*AB::AAi_ABa_BAi_BBa',
 '*AB::AAa_ABa_BAi_BBa',
 '*AB::ABa_BAi_BBa',
 '*AB::AAi_BAi_BBa',
 '*AB::AAa_BAi_BBa',
 '*AB::BAi_BBa',
 '*AB::AAi_ABi_BAa_BBa',
 '*AB::AAa_ABi_BAa_BBa',
 '*AB::ABi_BAa_BBa',
 '*AB::AAi_ABa_BAa_BBa',
 '*AB::AAa_ABa_BAa_BBa',
 '*AB::ABa_BAa_BBa',
 '*AB::AAi_BAa_BBa',
 '*AB::AAa_BAa_BBa',
 '*AB::BAa_BBa',
 '*AB::AAi_ABi_BBa',
 '*AB::AAa_ABi_BBa',
 '*AB::ABi_BBa',
 '*AB::AAi_ABa_BBa',
 '*AB::AAa_ABa_BBa',
 '*AB::ABa_BBa',
 '*AB::AAi_BBa',
 '*AB::AAa_BBa',
 '*AB::BBa',
 '*AB::AAi_ABi_BAi',
 '*AB::AAa_ABi_BAi',
 '*AB::ABi_BAi',
 '*AB::AAi_ABa_BAi',
 '*AB::AAa_ABa_BAi',
 '*AB::ABa_BAi',
 '*AB::AAi_BAi',
 '*AB::AAa_BAi',
 '*AB::BAi',
 '*AB::AAi_ABi_BAa',
 '*AB::AAa_ABi_BAa',
 '*AB::ABi_BAa',
 '*AB::AAi_ABa_BAa',
 '*AB::AAa_ABa_BAa',
 '*AB::ABa_BAa',
 '*AB::AAi_BAa',
 '*AB::AAa_BAa',
 '*AB::BAa',
 '*AB::AAi_ABi',
 '*AB::AAa_ABi',
 '*AB::ABi',
 '*AB::AAi_ABa',
 '*AB::AAa_ABa',
 '*AB::ABa',
 '*AB::AAi',
 '*AB::AAa',
 '*AB::']

Note that if your design space is large, grow_tree() can take an extremely long time and/or cause an out-of-memory error.


[4]:
%load_ext watermark
%watermark -v -p circuitree,jupyterlab,watermark
Python implementation: CPython
Python version       : 3.10.8
IPython version      : 8.24.0

circuitree: 0.11.1
jupyterlab: 4.1.8
watermark : 2.4.3