Defining custom embedders

In this notebook, you will learn how to:

• understand how the embedding module is structured,
• define a custom embedder from a function/algorithm and a config dataclass,

The available embedders page described some pre-defined embedders. The embedding module in QoolQit is designed to be flexible and extendable to various forms of embedding algorithms developed for the Rydberg analog model, with potentially different inputs and outputs, and different configuration parameters. It is structured in three levels:

Level 1: Concretizing the interface

The BaseEmbedder is the abstract base class for all embedders, but it is defined on generic input and output data types. It enforces the interface for all embedders by defining the info and config properties, as well as the embed method. It also defines abstract validate_input and validate_output methods that must be overwritten by subclasses.

Level 2: Concretizing input/output data

The next level is to define the concrete data types involved in the mapping, thus defining a family of embedders. Currently, there are two families of embedders defined in QoolQit:

• GraphToGraphEmbedder which concretizes the BaseEmbedder with a DataGraph input type and a DataGraph output type.
• MatrixToGraphEmbedder which concretizes the BaseEmbedder with a np.ndarray input type and a DataGraph output type.

In both cases, the validate_input and validate_output are overridden to check the input and output are of the correct type. In the case of the MatrixToGraphEmbedder, conditions on the input matrix are also checked such as if the array has the right dimensions and is symmetric. Still, at this level, no specific embedding algorithm is defined.

In the future, more families of embedders can be defined that may require different input and output data types.

Level 2: Concretizing the algorithm and config

The final level is defining concrete embedders, such as the ones we have used in the available embedders page. Here the requirement is to define a concrete function that maps the input to the output, along with any parameters required, and a config dataclass inheriting from EmbedderConfig holding all the configuration parameters.

Let's exemplify the case of defining a custom embedder in the family of graph to graph embedders.

In [1]:

from dataclasses import dataclass

from qoolqit import DataGraph
from qoolqit.embedding import EmbedderConfig, GraphToGraphEmbedder


def my_embedding_function(graph: DataGraph, param1: float) -> DataGraph:
    """Some embedding function that manipulates the input graph.

    This docstring should be clear on the embedding logic, because it will be
    directly accessed by the embedder.info property.

    Arguments:
        param1: a useless parameter...
    """
    return graph

@dataclass
class MyEmbedderConfig(EmbedderConfig):
    param1: float = 1.0

embedder = GraphToGraphEmbedder(my_embedding_function, MyEmbedderConfig())
print(embedder)

GraphToGraphEmbedder:
| Algorithm: my_embedding_function
| Config: MyEmbedderConfig(param1=1.0)

It can now be used like any other embedder.

In [2]:

print(embedder.info)

embedder.config.param1 = 2.0

graph = DataGraph.random_er(5, 0.5)

embedded_graph = embedder.embed(graph)

-- Embedding algorithm docstring:

Some embedding function that manipulates the input graph.

This docstring should be clear on the embedding logic, because it will be
directly accessed by the embedder.info property.

Arguments:
    param1: a useless parameter...

To share this embedder or potentially add it to the QoolQit codebase, we might want to create a concrete embedder wrapper that users can easily import.

In [3]:

class MyNewEmbedder(GraphToGraphEmbedder):
    def __init__(self):
        super().__init__(my_embedding_function, MyEmbedderConfig())

Automatic validation

To define a custom embedder, the extra arguments in the embedding function (besides the data) must match the fields in the configuration dataclass, otherwise an error will be raised.

In [4]:

def my_embedding_function(graph: DataGraph, param1: float) -> DataGraph:
    return graph

@dataclass
class MyWrongConfig(EmbedderConfig):
    some_other_param: float = 1.0

try:
    wrong_embedder = GraphToGraphEmbedder(my_embedding_function, MyWrongConfig())
except TypeError as error:
    print(error)

Config MyWrongConfig is not compatible with the algorithm my_embedding_function, as not all configuration fields correspond to keyword arguments in the algorithm function.

Config MyWrongConfig keys:
	- some_other_param

Algorithm signature parameters:
	- param1
	- graph

Furthermore, because we are defining an embedder in the GraphToGraphEmbedder, the input must be an instance of a DataGraph:

In [5]:

embedder = GraphToGraphEmbedder(my_embedding_function, MyEmbedderConfig())

try:
    data = 1.0 # Not a DataGraph
    embedded_data = embedder.embed(data)
except TypeError as error:
    print(error)

Embedding data of type <class 'float'> not supported by this embedder. GraphToGraphEmbedder requires data of type DataGraph.

The output of the embedding function must also be a DataGraph:

In [6]:

def my_wrong_embedding_function(graph: DataGraph, param1: float) -> DataGraph:
    return param1 # Not a DataGraph

embedder = GraphToGraphEmbedder(my_wrong_embedding_function, MyEmbedderConfig())

try:
    graph = DataGraph.random_er(5, 0.5)
    embedded_graph = embedder.embed(graph)
except TypeError as error:
    print(error)

Expected embedding result to be of type DataGraph, algorithm returned <class 'float'> instead.