v2 API

Constant

src.inspeqtor.v2.constant

Control

src.inspeqtor.v2.control

BaseControl dataclass

Source code in src/inspeqtor/v2/control.py

@dataclass
class BaseControl(ABC):
    # duration: int

    def __post_init__(self):
        # self.t_eval = jnp.arange(0, self.duration, 1)
        self.validate()

    def validate(self):
        # Validate that all attributes are json serializable
        try:
            json.dumps(self.to_dict())
        except TypeError as e:
            raise TypeError(
                f"Cannot serialize {self.__class__.__name__} to json"
            ) from e

        lower, upper = self.get_bounds()
        # Validate that the sampling function is working
        key = jax.random.key(0)
        params = sample_params(key, lower, upper)
        # waveform = self.get_waveform(params)

        assert all([isinstance(k, str) for k in params.keys()]), (
            "All key of params dict must be string"
        )
        assert all([isinstance(v, float) for v in params.values()]) or all(
            [isinstance(v, jnp.ndarray) for v in params.values()]
        ), "All value of params dict must be float"
        # assert isinstance(waveform, jnp.ndarray), "Waveform must be jnp.ndarray"

    @abstractmethod
    def get_bounds(
        self, *arg, **kwarg
    ) -> tuple[ParametersDictType, ParametersDictType]: ...

    @abstractmethod
    def get_envelope(self, params: ParametersDictType) -> typing.Callable:
        raise NotImplementedError("get_envelopes method is not implemented")

    def to_dict(self) -> dict[str, typing.Union[int, float, str]]:
        """Convert the control configuration to dictionary

        Returns:
            dict[str, typing.Union[int, float, str]]: Configuration of the control
        """
        return asdict(self)

    def to_dict_new(self) -> dict[str, typing.Union[int, float, str]]:
        """Convert the control configuration to dictionary

        Returns:
            dict[str, typing.Union[int, float, str]]: Configuration of the control
        """
        return {**asdict(self), "classname": self.__class__.__name__}

    @classmethod
    def from_dict(cls, data):
        """Construct the control instace from the dictionary.

        Args:
            data (dict): Dictionary for construction of the control instance.

        Returns:
            The instance of the control.
        """
        return cls(**data)

to_dict

to_dict() -> dict[str, Union[int, float, str]]

Convert the control configuration to dictionary

Returns:

Type	Description
dict[str, Union[int, float, str]]	dict[str, typing.Union[int, float, str]]: Configuration of the control

Source code in src/inspeqtor/v2/control.py

def to_dict(self) -> dict[str, typing.Union[int, float, str]]:
    """Convert the control configuration to dictionary

    Returns:
        dict[str, typing.Union[int, float, str]]: Configuration of the control
    """
    return asdict(self)

to_dict_new

to_dict_new() -> dict[str, Union[int, float, str]]

Convert the control configuration to dictionary

Returns:

Type	Description
dict[str, Union[int, float, str]]	dict[str, typing.Union[int, float, str]]: Configuration of the control

Source code in src/inspeqtor/v2/control.py

def to_dict_new(self) -> dict[str, typing.Union[int, float, str]]:
    """Convert the control configuration to dictionary

    Returns:
        dict[str, typing.Union[int, float, str]]: Configuration of the control
    """
    return {**asdict(self), "classname": self.__class__.__name__}

from_dict classmethod

from_dict(data)

Construct the control instace from the dictionary.

Parameters:

Name	Type	Description	Default
data	dict	Dictionary for construction of the control instance.	required

Returns:

Type	Description
	The instance of the control.

Source code in src/inspeqtor/v2/control.py

@classmethod
def from_dict(cls, data):
    """Construct the control instace from the dictionary.

    Args:
        data (dict): Dictionary for construction of the control instance.

    Returns:
        The instance of the control.
    """
    return cls(**data)

ControlSequence dataclass

Control sequence, expect to be sum of atomic control.

Source code in src/inspeqtor/v2/control.py

@dataclass
class ControlSequence:
    """Control sequence, expect to be sum of atomic control."""

    controls: dict[str, BaseControl]
    total_dt: int
    structure: typing.Sequence[typing.Sequence[str]] | None = field(default=None)

    def __post_init__(self):
        # Cache the bounds
        self.lower, self.upper = self.get_bounds()
        # Create the order

        if self.structure is None:
            self.auto_order = []
            for ctrl_key in self.controls.keys():
                sub_control = []
                for ctrl_param_key in self.lower[ctrl_key].keys():
                    sub_control.append((ctrl_key, ctrl_param_key))

                self.auto_order += sub_control
            self.structure = self.auto_order
        else:
            self.auto_order = self.structure

    def get_structure(self) -> typing.Sequence[typing.Sequence[str]]:
        return self.auto_order

    def sample_params_v1(self, key: jnp.ndarray) -> dict[str, ParametersDictType]:
        """Sample control parameter

        Args:
            key (jnp.ndarray): Random key

        Returns:
            dict[str, ParametersDictType]: control parameters
        """
        params_dict: dict[str, ParametersDictType] = {}
        for idx, ctrl_key in enumerate(self.controls.keys()):
            subkey = jax.random.fold_in(key, idx)
            params = sample_params(subkey, self.lower[ctrl_key], self.upper[ctrl_key])
            params_dict[ctrl_key] = params

        return params_dict

    def sample_params(self, key: jnp.ndarray) -> dict[str, ParametersDictType]:
        return self.sample_params_v2(key)

    def sample_params_v2(self, key: jnp.ndarray) -> dict[str, ParametersDictType]:
        """Sample control parameter

        Args:
            key (jnp.ndarray): Random key

        Returns:
            dict[str, ParametersDictType]: control parameters
        """
        return nested_sample(key, merge_lower_upper(self.lower, self.upper))

    def get_bounds(
        self,
    ) -> tuple[dict[str, ParametersDictType], dict[str, ParametersDictType]]:
        """Get the bounds of the controls

        Returns:
            tuple[list[ParametersDictType], list[ParametersDictType]]: tuple of list of lower and upper bounds.
        """

        lower_bounds = jax.tree.map(
            lambda x: x.get_bounds()[0],
            self.controls,
            is_leaf=lambda x: isinstance(x, BaseControl),
        )
        upper_bounds = jax.tree.map(
            lambda x: x.get_bounds()[1],
            self.controls,
            is_leaf=lambda x: isinstance(x, BaseControl),
        )

        return lower_bounds, upper_bounds

    def get_envelope(
        self, params_dict: dict[str, ParametersDictType]
    ) -> typing.Callable:
        """Create envelope function with given control parameters

        Args:
            params_list (dict[str, ParametersDictType]): control parameter to be used

        Returns:
            typing.Callable: Envelope function
        """
        callables = []
        for params_key, params_val in params_dict.items():
            callables.append(self.controls[params_key].get_envelope(params_val))

        # Create a function that returns the sum of the envelopes
        def envelope(t):
            return sum([c(t) for c in callables])

        return envelope

    def to_dict_new(self) -> dict[str, str | dict[str, str | float]]:
        """Convert self to dict

        Returns:
            dict[str, str | dict[str, str | float]]: dict contain argument necessary for re-initialization.
        """
        return {
            **asdict(self),
            "classname": {k: v.__class__.__name__ for k, v in self.controls.items()},
            "controls": jax.tree.map(
                lambda x: x.to_dict(),
                self.controls,
                is_leaf=lambda x: isinstance(x, BaseControl),
            ),
        }

    def to_dict(self) -> dict[str, str | dict[str, str | float]]:
        """Convert self to dict

        Returns:
            dict[str, str | dict[str, str | float]]: dict contain argument necessary for re-initialization.
        """
        return {
            **asdict(self),
            "classname": {k: v.__class__.__name__ for k, v in self.controls.items()},
        }

    @classmethod
    def from_dict(
        cls,
        data: dict[str, str | dict[str, str | float]],
        controls: dict[str, type[BaseControl]],
    ) -> "ControlSequence":
        """Construct self with the provided dictionary

        Args:
            data (dict[str, str  |  dict[str, str  |  float]]): The dictionary contain initialization arguments
            controls (dict[str, type[BaseControl]]): The map of control name and class of the control

        Returns:
            ControlSequence: the instance of control sequence
        """
        controls_data = data["controls"]
        assert isinstance(controls_data, dict)

        instantiated_controls = {}

        for (ctrl_key, ctrl_data), (ctrl_key_match, ctrl_cls) in zip(
            controls_data.items(), controls.items()
        ):
            assert ctrl_key == ctrl_key_match
            assert isinstance(ctrl_data, dict), f"Expected dict, got {type(ctrl_data)}"
            instantiated_controls[ctrl_key] = ctrl_cls.from_dict(ctrl_data)

        total_dt = data["total_dt"]
        assert isinstance(total_dt, int)
        structure = data["structure"]
        assert isinstance(structure, list)
        # Explicitly convert each item in the structure to be tuple.
        structure = [tuple(item) for item in structure]

        return cls(
            controls=instantiated_controls, total_dt=total_dt, structure=structure
        )

    @classmethod
    def from_dict_new(
        cls,
        data: dict[str, str | dict[str, str | float]],
        controls: dict[str, type[BaseControl]],
    ) -> "ControlSequence":
        controls_data = data["controls"]
        assert isinstance(controls_data, dict)

        def check_if_control_dict(leave) -> bool:
            if isinstance(leave, dict):
                if "classname" in leave:
                    return True

            return False

        def initialize_control(control_data: dict) -> BaseControl:
            cls_name = control_data["classname"]
            clean_data = {k: v for k, v in control_data.items() if k != "classname"}
            return controls[cls_name].from_dict(clean_data)

        # Initialize contols
        initialized_controls = jax.tree.map(
            initialize_control, controls_data, is_leaf=check_if_control_dict
        )

        total_dt = data["total_dt"]
        assert isinstance(total_dt, int)
        structure = data["structure"]
        assert isinstance(structure, list)
        # Explicitly convert each item in the structure to be tuple.
        structure = [tuple(item) for item in structure]

        return cls(
            controls=initialized_controls, total_dt=total_dt, structure=structure
        )

    def to_file(self, path: typing.Union[str, pathlib.Path]):
        """Save configuration of the pulse to file given folder path.

        Args:
            path (typing.Union[str, pathlib.Path]): Path to the folder to save sequence, will be created if not existed.
        """
        if isinstance(path, str):
            path = pathlib.Path(path)

        save_pytree_to_json(self.to_dict(), path / "control_sequence.json")

    @classmethod
    def from_file(
        cls,
        path: typing.Union[str, pathlib.Path],
        controls: dict[str, type[BaseControl]],
    ):
        """Initialize itself from a file.

        Args:
            path (typing.Union[str, pathlib.Path]): Path to file.
            controls (dict[str, type[BaseControl]]): The map of control name and class of the control

        Returns:
            ControlSequence: the instance of control sequence
        """
        if isinstance(path, str):
            path = pathlib.Path(path)

        ctrl_loaded_dict = load_pytree_from_json(
            path / "control_sequence.json", lambda k, v: (True, v)
        )

        return cls.from_dict(ctrl_loaded_dict, controls=controls)

sample_params_v1

sample_params_v1(
    key: ndarray,
) -> dict[str, ParametersDictType]

Sample control parameter

Parameters:

Name	Type	Description	Default
key	ndarray	Random key	required

Returns:

Type	Description
dict[str, ParametersDictType]	dict[str, ParametersDictType]: control parameters

Source code in src/inspeqtor/v2/control.py

def sample_params_v1(self, key: jnp.ndarray) -> dict[str, ParametersDictType]:
    """Sample control parameter

    Args:
        key (jnp.ndarray): Random key

    Returns:
        dict[str, ParametersDictType]: control parameters
    """
    params_dict: dict[str, ParametersDictType] = {}
    for idx, ctrl_key in enumerate(self.controls.keys()):
        subkey = jax.random.fold_in(key, idx)
        params = sample_params(subkey, self.lower[ctrl_key], self.upper[ctrl_key])
        params_dict[ctrl_key] = params

    return params_dict

sample_params_v2

sample_params_v2(
    key: ndarray,
) -> dict[str, ParametersDictType]

Sample control parameter

Parameters:

Name	Type	Description	Default
key	ndarray	Random key	required

Returns:

Type	Description
dict[str, ParametersDictType]	dict[str, ParametersDictType]: control parameters

Source code in src/inspeqtor/v2/control.py

def sample_params_v2(self, key: jnp.ndarray) -> dict[str, ParametersDictType]:
    """Sample control parameter

    Args:
        key (jnp.ndarray): Random key

    Returns:
        dict[str, ParametersDictType]: control parameters
    """
    return nested_sample(key, merge_lower_upper(self.lower, self.upper))

get_bounds

get_bounds() -> tuple[
    dict[str, ParametersDictType],
    dict[str, ParametersDictType],
]

Get the bounds of the controls

Returns:

Type	Description
tuple[dict[str, ParametersDictType], dict[str, ParametersDictType]]	tuple[list[ParametersDictType], list[ParametersDictType]]: tuple of list of lower and upper bounds.

Source code in src/inspeqtor/v2/control.py

def get_bounds(
    self,
) -> tuple[dict[str, ParametersDictType], dict[str, ParametersDictType]]:
    """Get the bounds of the controls

    Returns:
        tuple[list[ParametersDictType], list[ParametersDictType]]: tuple of list of lower and upper bounds.
    """

    lower_bounds = jax.tree.map(
        lambda x: x.get_bounds()[0],
        self.controls,
        is_leaf=lambda x: isinstance(x, BaseControl),
    )
    upper_bounds = jax.tree.map(
        lambda x: x.get_bounds()[1],
        self.controls,
        is_leaf=lambda x: isinstance(x, BaseControl),
    )

    return lower_bounds, upper_bounds

get_envelope

get_envelope(
    params_dict: dict[str, ParametersDictType],
) -> Callable

Create envelope function with given control parameters

Parameters:

Name	Type	Description	Default
params_list	dict[str, ParametersDictType]	control parameter to be used	required

Returns:

Type	Description
Callable	typing.Callable: Envelope function

Source code in src/inspeqtor/v2/control.py

def get_envelope(
    self, params_dict: dict[str, ParametersDictType]
) -> typing.Callable:
    """Create envelope function with given control parameters

    Args:
        params_list (dict[str, ParametersDictType]): control parameter to be used

    Returns:
        typing.Callable: Envelope function
    """
    callables = []
    for params_key, params_val in params_dict.items():
        callables.append(self.controls[params_key].get_envelope(params_val))

    # Create a function that returns the sum of the envelopes
    def envelope(t):
        return sum([c(t) for c in callables])

    return envelope

to_dict_new

to_dict_new() -> dict[str, str | dict[str, str | float]]

Convert self to dict

Returns:

Type	Description
dict[str, str | dict[str, str | float]]	dict[str, str | dict[str, str | float]]: dict contain argument necessary for re-initialization.

Source code in src/inspeqtor/v2/control.py

def to_dict_new(self) -> dict[str, str | dict[str, str | float]]:
    """Convert self to dict

    Returns:
        dict[str, str | dict[str, str | float]]: dict contain argument necessary for re-initialization.
    """
    return {
        **asdict(self),
        "classname": {k: v.__class__.__name__ for k, v in self.controls.items()},
        "controls": jax.tree.map(
            lambda x: x.to_dict(),
            self.controls,
            is_leaf=lambda x: isinstance(x, BaseControl),
        ),
    }

to_dict

to_dict() -> dict[str, str | dict[str, str | float]]

Convert self to dict

Returns:

Type	Description
dict[str, str | dict[str, str | float]]	dict[str, str | dict[str, str | float]]: dict contain argument necessary for re-initialization.

Source code in src/inspeqtor/v2/control.py

def to_dict(self) -> dict[str, str | dict[str, str | float]]:
    """Convert self to dict

    Returns:
        dict[str, str | dict[str, str | float]]: dict contain argument necessary for re-initialization.
    """
    return {
        **asdict(self),
        "classname": {k: v.__class__.__name__ for k, v in self.controls.items()},
    }

from_dict classmethod

from_dict(
    data: dict[str, str | dict[str, str | float]],
    controls: dict[str, type[BaseControl]],
) -> ControlSequence

Construct self with the provided dictionary

Parameters:

Name	Type	Description	Default
data	dict[str, str | dict[str, str | float]]	The dictionary contain initialization arguments	required
controls	dict[str, type[BaseControl]]	The map of control name and class of the control	required

Returns:

Name	Type	Description
ControlSequence	ControlSequence	the instance of control sequence

Source code in src/inspeqtor/v2/control.py

@classmethod
def from_dict(
    cls,
    data: dict[str, str | dict[str, str | float]],
    controls: dict[str, type[BaseControl]],
) -> "ControlSequence":
    """Construct self with the provided dictionary

    Args:
        data (dict[str, str  |  dict[str, str  |  float]]): The dictionary contain initialization arguments
        controls (dict[str, type[BaseControl]]): The map of control name and class of the control

    Returns:
        ControlSequence: the instance of control sequence
    """
    controls_data = data["controls"]
    assert isinstance(controls_data, dict)

    instantiated_controls = {}

    for (ctrl_key, ctrl_data), (ctrl_key_match, ctrl_cls) in zip(
        controls_data.items(), controls.items()
    ):
        assert ctrl_key == ctrl_key_match
        assert isinstance(ctrl_data, dict), f"Expected dict, got {type(ctrl_data)}"
        instantiated_controls[ctrl_key] = ctrl_cls.from_dict(ctrl_data)

    total_dt = data["total_dt"]
    assert isinstance(total_dt, int)
    structure = data["structure"]
    assert isinstance(structure, list)
    # Explicitly convert each item in the structure to be tuple.
    structure = [tuple(item) for item in structure]

    return cls(
        controls=instantiated_controls, total_dt=total_dt, structure=structure
    )

to_file

to_file(path: Union[str, Path])

Save configuration of the pulse to file given folder path.

Parameters:

Name	Type	Description	Default
path	Union[str, Path]	Path to the folder to save sequence, will be created if not existed.	required

Source code in src/inspeqtor/v2/control.py

def to_file(self, path: typing.Union[str, pathlib.Path]):
    """Save configuration of the pulse to file given folder path.

    Args:
        path (typing.Union[str, pathlib.Path]): Path to the folder to save sequence, will be created if not existed.
    """
    if isinstance(path, str):
        path = pathlib.Path(path)

    save_pytree_to_json(self.to_dict(), path / "control_sequence.json")

from_file classmethod

from_file(
    path: Union[str, Path],
    controls: dict[str, type[BaseControl]],
)

Initialize itself from a file.

Parameters:

Name	Type	Description	Default
path	Union[str, Path]	Path to file.	required
controls	dict[str, type[BaseControl]]	The map of control name and class of the control	required

Returns:

Name	Type	Description
ControlSequence		the instance of control sequence

Source code in src/inspeqtor/v2/control.py

@classmethod
def from_file(
    cls,
    path: typing.Union[str, pathlib.Path],
    controls: dict[str, type[BaseControl]],
):
    """Initialize itself from a file.

    Args:
        path (typing.Union[str, pathlib.Path]): Path to file.
        controls (dict[str, type[BaseControl]]): The map of control name and class of the control

    Returns:
        ControlSequence: the instance of control sequence
    """
    if isinstance(path, str):
        path = pathlib.Path(path)

    ctrl_loaded_dict = load_pytree_from_json(
        path / "control_sequence.json", lambda k, v: (True, v)
    )

    return cls.from_dict(ctrl_loaded_dict, controls=controls)

sample_params

sample_params(
    key: ndarray,
    lower: ParametersDictType,
    upper: ParametersDictType,
) -> ParametersDictType

Sample parameters with the same shape with given lower and upper bounds

Parameters:

Name	Type	Description	Default
key	ndarray	Random key	required
lower	ParametersDictType	Lower bound	required
upper	ParametersDictType	Upper bound	required

Returns:

Name	Type	Description
ParametersDictType	ParametersDictType	Dict of the sampled parameters

Source code in src/inspeqtor/v2/control.py

def sample_params(
    key: jnp.ndarray, lower: ParametersDictType, upper: ParametersDictType
) -> ParametersDictType:
    """Sample parameters with the same shape with given lower and upper bounds

    Args:
        key (jnp.ndarray): Random key
        lower (ParametersDictType): Lower bound
        upper (ParametersDictType): Upper bound

    Returns:
        ParametersDictType: Dict of the sampled parameters
    """
    # This function is general because it is depend only on lower and upper structure
    param: ParametersDictType = {}
    param_names = lower.keys()
    for name in param_names:
        sample_key, key = jax.random.split(key)
        param[name] = jax.random.uniform(
            sample_key, shape=(), minval=lower[name], maxval=upper[name]
        )

    # return jax.tree.map(float, param)
    return param

sequence_waveform

sequence_waveform(
    params: dict[str, ParametersDictType],
    t_eval: ndarray,
    control_seqeunce: ControlSequence,
) -> ndarray

Samples the pulse sequence by generating random parameters for each pulse and computing the total waveform.

Parameters:

Name	Type	Description	Default
key	Key	The random key used for generating the parameters.	required

Returns:

Type	Description
ndarray	tuple[list[ParametersDictType], Complex[Array, "time"]]: A tuple containing a list of parameter dictionaries for each pulse and the total waveform.

Example

key = jax.random.PRNGKey(0) params, total_waveform = sample(key)

Source code in src/inspeqtor/v2/control.py

def sequence_waveform(
    params: dict[str, ParametersDictType],
    t_eval: jnp.ndarray,
    control_seqeunce: ControlSequence,
) -> jnp.ndarray:
    """
    Samples the pulse sequence by generating random parameters for each pulse and computing the total waveform.

    Parameters:
        key (Key): The random key used for generating the parameters.

    Returns:
        tuple[list[ParametersDictType], Complex[Array, "time"]]: A tuple containing a list of parameter dictionaries for each pulse and the total waveform.

    Example:
        key = jax.random.PRNGKey(0)
        params, total_waveform = sample(key)
    """
    # Create base waveform
    total_waveform = jnp.zeros_like(t_eval, dtype=jnp.complex64)

    for (param_key, param_val), (ctrl_key, control) in zip(
        params.items(), control_seqeunce.controls.items()
    ):
        waveform = control_waveform(param_val, t_eval, control)
        total_waveform += waveform

    return total_waveform

merge_lower_upper

merge_lower_upper(lower: LowerBound, upper: UpperBound)

Merge lower and upper bound into bounds

Parameters:

Name	Type	Description	Default
lower	_type_	The lower bound	required
upper	_type_	The upper bound	required

Returns:

Name	Type	Description
_type_		Bound from the lower and upper.

Source code in src/inspeqtor/v2/control.py

def merge_lower_upper(lower: LowerBound, upper: UpperBound):
    """Merge lower and upper bound into bounds

    Args:
        lower (_type_): The lower bound
        upper (_type_): The upper bound

    Returns:
        _type_: Bound from the lower and upper.
    """
    return jax.tree.map(lambda x, y: (x, y), lower, upper)

split_bounds

split_bounds(
    bounds: Bounds,
) -> tuple[LowerBound, UpperBound]

Create lower and upper bound from bounds

Parameters:

Name	Type	Description	Default
bounds	_type_	The bounds to extract the lower and upper bound	required

Returns:

Name	Type	Description
_type_	tuple[LowerBound, UpperBound]	The lower and upper bound

Source code in src/inspeqtor/v2/control.py

def split_bounds(bounds: Bounds) -> tuple[LowerBound, UpperBound]:
    """Create lower and upper bound from bounds

    Args:
        bounds (_type_): The bounds to extract the lower and upper bound

    Returns:
        _type_: The lower and upper bound
    """
    return jax.tree.map(
        lambda x: x[0], bounds, is_leaf=lambda x: isinstance(x, tuple)
    ), jax.tree.map(lambda x: x[1], bounds, is_leaf=lambda x: isinstance(x, tuple))

nested_sample

nested_sample(
    key: ndarray, bounds, sample_fn=uniform_sample
)

Sample from nested bounds with custom sampling function sample_fn

Parameters:

Name	Type	Description	Default
key	ndarray	Random key	required
bounds	_type_	Bound of the control parameter	required
sample_fn	_type_	Custom sampling function. Defaults to uniform_sample.	uniform_sample

Returns:

Name	Type	Description
_type_		Control parameter sample from bound

Source code in src/inspeqtor/v2/control.py

def nested_sample(key: jnp.ndarray, bounds, sample_fn=uniform_sample):
    """Sample from nested bounds with custom sampling function `sample_fn`

    Args:
        key (jnp.ndarray): Random key
        bounds (_type_): Bound of the control parameter
        sample_fn (_type_, optional): Custom sampling function. Defaults to uniform_sample.

    Returns:
        _type_: Control parameter sample from bound
    """
    return unflatten_dict(
        {
            k: sample_fn(jax.random.fold_in(key, idx), bound)
            for idx, (k, bound) in enumerate(flatten_dict(bounds).items())
        }
    )

check_bounds

check_bounds(
    param: ParametersDictType, bounds: Bounds
) -> bool

Check if the given control parameter violate the bound or not.

Parameters:

Name	Type	Description	Default
param	_type_	Control parameter	required
bounds	_type_	Bound of control parameter	required

Returns:

Name	Type	Description
bool	bool	True if parameter do not violate the bound, otherwise False

Source code in src/inspeqtor/v2/control.py

def check_bounds(param: ParametersDictType, bounds: Bounds) -> bool:
    """Check if the given control parameter violate the bound or not.

    Args:
        param (_type_): Control parameter
        bounds (_type_): Bound of control parameter

    Returns:
        bool: `True` if parameter do not violate the bound, otherwise `False`
    """
    valid_container = jax.tree.map(
        lambda x, bound: (bound[0] < x) & (x < bound[1]), param, bounds
    )
    return jax.tree.reduce(lambda init, x: init & x, valid_container, initializer=True)

ravel_unravel_fn

ravel_unravel_fn(structure: Iterable[Iterable[str]])

This function return the ravel and unravel functions for the provided control sequence

Parameters:

Name	Type	Description	Default
structure	Iterable[Iterable[str]]	The structure of the pytree	required

Returns:

Type	Description
	tuple[typing.Callable, typing.Callable]: The first element is the function that convert structured parameter to array, the second is a function that reverse the action of the first.

Source code in src/inspeqtor/v2/control.py

def ravel_unravel_fn(structure: typing.Iterable[typing.Iterable[str]]):
    """This function return the ravel and unravel functions for the provided control sequence

    Args:
        structure (typing.Iterable[typing.Iterable[str]]): The structure of the pytree

    Returns:
        tuple[typing.Callable, typing.Callable]: The first element is the function that convert structured parameter to array, the second is a function that reverse the action of the first.
    """

    def ravel_fn(param: ParametersDictType):
        return jnp.array(
            [get_value_by_keys(param, dict_keys) for dict_keys in structure]
        )

    def unravel_fn(param: jnp.ndarray):
        return unflatten_dict(
            {dict_keys: param[idx] for idx, dict_keys in enumerate(structure)}
        )

    return ravel_fn, unravel_fn

ravel_unravel_fn_old

ravel_unravel_fn_old(control_sequence: ControlSequence)

This function return the ravel and unravel functions for the provided control sequence

Parameters:

Name	Type	Description	Default
control_sequence	ControlSequence	The control sequence	required

Returns:

Type	Description
	tuple[typing.Callable, typing.Callable]: The first element is the function that convert structured parameter to array, the second is a function that reverse the action of the first.

Source code in src/inspeqtor/v2/control.py

@deprecated(reason="Old implementation of the ravel_unravel_fn")
def ravel_unravel_fn_old(control_sequence: ControlSequence):
    """This function return the ravel and unravel functions for the provided control sequence

    Args:
        control_sequence (ControlSequence): The control sequence

    Returns:
        tuple[typing.Callable, typing.Callable]: The first element is the function that convert structured parameter to array, the second is a function that reverse the action of the first.
    """
    structure = control_sequence.get_structure()

    def ravel_fn(param_dict: dict[str, ParametersDictType]) -> jnp.ndarray:
        tmp = []
        for sub_order in structure:
            tmp.append(param_dict[sub_order[0]][sub_order[1]])
        return jnp.array(tmp)

    def unravel_fn(param: jnp.ndarray) -> dict[str, ParametersDictType]:
        tmp = {}
        for idx, sub_order in enumerate(structure):
            if sub_order[0] not in tmp:
                tmp[sub_order[0]] = {}

            tmp[sub_order[0]][sub_order[1]] = param[idx]
        return tmp

    return ravel_fn, unravel_fn

construct_control_sequence_reader

construct_control_sequence_reader(
    controls: list[type[BaseControl]] = [],
) -> Callable[[Union[str, Path]], ControlSequence]

Construct the control sequence reader

Parameters:

Name	Type	Description	Default
controls	list[type[BasePulse]]	List of control constructor. Defaults to [].	[]

Returns:

Type	Description
Callable[[Union[str, Path]], ControlSequence]	typing.Callable[[typing.Union[str, pathlib.Path]], controlsequence]: Control sequence reader that will automatically contruct control sequence from path.

Source code in src/inspeqtor/v2/control.py

def construct_control_sequence_reader(
    controls: list[type[BaseControl]] = [],
) -> typing.Callable[[typing.Union[str, pathlib.Path]], ControlSequence]:
    """Construct the control sequence reader

    Args:
        controls (list[type[BasePulse]], optional): List of control constructor. Defaults to [].

    Returns:
        typing.Callable[[typing.Union[str, pathlib.Path]], controlsequence]: Control sequence reader that will automatically contruct control sequence from path.
    """
    default_controls: list[type[BaseControl]] = []

    # Merge the default controls with the provided controls
    controls_list = default_controls + controls
    control_dict = {ctrl.__name__: ctrl for ctrl in controls_list}

    def control_sequence_reader(
        path: typing.Union[str, pathlib.Path],
    ) -> ControlSequence:
        """Construct control sequence from path

        Args:
            path (typing.Union[str, pathlib.Path]): Path of the saved control sequence configuration.

        Returns:
            ControlSeqence: Control sequence instance.
        """
        if isinstance(path, str):
            path = pathlib.Path(path)

        control_sequence_dict = load_pytree_from_json(
            path / "control_sequence.json", lambda k, v: (True, v)
        )

        parsed_controls = {}
        assert isinstance(control_sequence_dict["classname"], dict)

        for ctrl_key, ctrl_classname in control_sequence_dict["classname"].items():
            parsed_controls[ctrl_key] = control_dict[ctrl_classname]

        return ControlSequence.from_dict(
            control_sequence_dict, controls=parsed_controls
        )

    return control_sequence_reader

get_envelope_transformer

get_envelope_transformer(control_sequence: ControlSequence)

Generate get_envelope function with control parameter array as an input instead of list form

Parameters:

Name	Type	Description	Default
control_sequence	ControlSequence	Control seqence instance	required

Returns:

Type	Description
	typing.Callable[[jnp.ndarray], typing.Any]: Transformed get envelope function

Source code in src/inspeqtor/v2/control.py

def get_envelope_transformer(control_sequence: ControlSequence):
    """Generate get_envelope function with control parameter array as an input instead of list form

    Args:
        control_sequence (ControlSequence): Control seqence instance

    Returns:
        typing.Callable[[jnp.ndarray], typing.Any]: Transformed get envelope function
    """
    _, unravel_fn = ravel_unravel_fn(control_sequence.get_structure())

    def get_envelope(params: jnp.ndarray) -> typing.Callable[..., typing.Any]:
        return control_sequence.get_envelope(unravel_fn(params))

    return get_envelope

ravel_transform

ravel_transform(
    fn: Callable,
    control_sequence: ControlSequence,
    at: int = 0,
) -> Callable

Transform the argument at index at of the function fn with unravel_fn of the control sequence

Note

signal_fn = sq.control.ravel_transform(
    sq.physics.signal_func_v5(control_sequence.get_envelope, qubit_info.frequency, dt),
    control_sequence,
)

Parameters:

Name	Type	Description	Default
fn	Callable	The function to be transformed	required
control_sequence	ControlSequence	The control sequence that will use to produce unravel_fn.	required

Returns:

Type	Description
Callable	typing.Callable: A function that its first argument is transformed by unravel_fn

Source code in src/inspeqtor/v2/control.py

def ravel_transform(
    fn: typing.Callable, control_sequence: ControlSequence, at: int = 0
) -> typing.Callable:
    """Transform the argument at index `at` of the function `fn` with `unravel_fn` of the control sequence

    Note:
        ```python
        signal_fn = sq.control.ravel_transform(
            sq.physics.signal_func_v5(control_sequence.get_envelope, qubit_info.frequency, dt),
            control_sequence,
        )
        ```

    Args:
        fn (typing.Callable): The function to be transformed
        control_sequence (ControlSequence): The control sequence that will use to produce `unravel_fn`.

    Returns:
        typing.Callable: A function that its first argument is transformed by `unravel_fn`
    """
    _, unravel_fn = ravel_unravel_fn(control_sequence.get_structure())

    def wrapper(*args, **kwargs):
        list_args = list(args)
        list_args[at] = unravel_fn(list_args[at])

        return fn(*tuple(list_args), **kwargs)

    return wrapper

get_envelope

get_envelope(
    param: ParametersDictType, seq: ControlSequence
)

Return an envelope function create from envelope of all controls in seq with control parameter param

Parameters:

Name	Type	Description	Default
param	_type_	Control parameter	required
seq	ControlSequence	Control Sequence	required

Returns:

Name	Type	Description
_type_		A function of time which is a sum of all envelope of control in seq with parameter param

Source code in src/inspeqtor/v2/control.py

def get_envelope(param: ParametersDictType, seq: ControlSequence):
    """Return an envelope function create from envelope of all controls in `seq` with control parameter `param`

    Args:
        param (_type_): Control parameter
        seq (ControlSequence): Control Sequence

    Returns:
        _type_: A function of time which is a sum of all envelope of control in `seq` with parameter `param`
    """
    tree = jax.tree.map(
        lambda ctrl, x: ctrl.get_envelope(x),
        seq.controls,
        param,
        is_leaf=lambda x: isinstance(x, BaseControl),
    )

    def envelope(t):
        return jax.tree.reduce(lambda value, fn: fn(t) + value, tree, initializer=0.0)

    return envelope

envelope_fn

envelope_fn(
    param: ParametersDictType,
    t: ndarray,
    seq: ControlSequence,
)

Return an envelope of all of the control in control sequence seq given paramter param at time t

Parameters:

Name	Type	Description	Default
param	_type_	Control parameter	required
t	ndarray	Time to evaluate the envelope	required
seq	ControlSequence	The control sequence to get the envelope	required

Returns:

Name	Type	Description
_type_		Envelope of all control in seq evaluate at time t with parameter param

Source code in src/inspeqtor/v2/control.py

def envelope_fn(param: ParametersDictType, t: jnp.ndarray, seq: ControlSequence):
    """Return an envelope of all of the control in control sequence `seq` given paramter `param` at time `t`

    Args:
        param (_type_): Control parameter
        t (jnp.ndarray): Time to evaluate the envelope
        seq (ControlSequence): The control sequence to get the envelope

    Returns:
        _type_: Envelope of all control in `seq` evaluate at time `t` with parameter `param`
    """
    tree = jax.tree.map(
        lambda ctrl, x: ctrl.get_envelope(x)(t),
        seq.controls,
        param,
        is_leaf=lambda x: isinstance(x, BaseControl),
    )

    return jax.tree.reduce(jnp.add, tree)

Data

src.inspeqtor.v2.data

ExpectationValue dataclass

Class representing a single experimental setting of state initialization and observable measurement.

Supports both single-qubit and multi-qubit configurations using string representation: - Observable: "XYZ" (instead of ["X", "Y", "Z"]) - Initial state: "+0r" (instead of ["+", "0", "r"])

Source code in src/inspeqtor/v2/data.py

@dataclass
class ExpectationValue:
    """Class representing a single experimental setting of state initialization and observable measurement.

    Supports both single-qubit and multi-qubit configurations using string representation:
    - Observable: "XYZ" (instead of ["X", "Y", "Z"])
    - Initial state: "+0r" (instead of ["+", "0", "r"])
    """

    initial_state: str
    # String where each character represents an observable for one qubit
    observable: str
    # String where each character represents an initial state for one qubit

    def __post_init__(self):
        # Ensure both strings have the same length (number of qubits)
        assert len(self.observable) == len(self.initial_state), (
            f"Observable and initial state must have same number of qubits: {len(self.observable)} != {len(self.initial_state)}"
        )

        # Validate observable characters
        for o in self.observable:
            assert o in "IXYZ", (
                f"Invalid observable '{o}'. Must be one of 'I', 'X', 'Y', or 'Z'"
            )

        # Validate initial state characters
        valid_states = "+-rl01"
        for s in self.initial_state:
            assert s in valid_states, (
                f"Invalid initial state '{s}'. Must be one of {valid_states}"
            )

    def to_dict(self):
        return {
            "initial_state": self.initial_state,
            "observable": self.observable,
        }

    def __eq__(self, __value: object) -> bool:
        if not isinstance(__value, ExpectationValue):
            return False

        return (
            self.initial_state == __value.initial_state
            and self.observable == __value.observable
        )

    @classmethod
    def from_dict(cls, data):
        return cls(**data)

    def __str__(self) -> str:
        return self.initial_state + "/" + self.observable

ExperimentConfiguration dataclass

Experiment configuration dataclass

Source code in src/inspeqtor/v2/data.py

@dataclass
class ExperimentConfiguration:
    """Experiment configuration dataclass"""

    qubits: typing.Sequence[QubitInformation]
    expectation_values_order: typing.Sequence[ExpectationValue]
    parameter_structure: typing.Sequence[
        typing.Sequence[str]
    ]  # Get from the pulse sequence .get_parameter_names()
    backend_name: str
    shots: int
    EXPERIMENT_IDENTIFIER: str
    EXPERIMENT_TAGS: typing.Sequence[str]
    description: str
    device_cycle_time_ns: float
    sequence_duration_dt: int
    sample_size: int
    date: str = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    additional_info: dict[str, str | int | float] = field(default_factory=dict)

    def to_dict(self):
        return {
            **asdict(self),
            "qubits": [qubit.to_dict() for qubit in self.qubits],
            "expectation_values_order": [
                exp.to_dict() for exp in self.expectation_values_order
            ],
        }

    @classmethod
    def from_dict(cls, dict_experiment_config):
        dict_experiment_config["qubits"] = [
            QubitInformation.from_dict(qubit)
            for qubit in dict_experiment_config["qubits"]
        ]

        dict_experiment_config["expectation_values_order"] = [
            ExpectationValue.from_dict(exp)
            for exp in dict_experiment_config["expectation_values_order"]
        ]

        dict_experiment_config["parameter_structure"] = [
            tuple(control) for control in dict_experiment_config["parameter_structure"]
        ]

        return cls(**dict_experiment_config)

    def to_file(self, path: typing.Union[Path, str]):
        if isinstance(path, str):
            path = Path(path)

        # os.makedirs(path, exist_ok=True)
        path.mkdir(parents=True, exist_ok=True)
        with open(path / "config.json", "w") as f:
            json.dump(self.to_dict(), f, indent=4)

    @classmethod
    def from_file(cls, path: typing.Union[Path, str]):
        if isinstance(path, str):
            path = Path(path)
        with open(path / "config.json", "r") as f:
            dict_experiment_config = json.load(f)

        return cls.from_dict(dict_experiment_config)

    def __str__(self):
        lines = [
            "=" * 60,
            "EXPERIMENT CONFIGURATION",
            "=" * 60,
            f"Identifier: {self.EXPERIMENT_IDENTIFIER}",
            f"Backend: {self.backend_name}",
            f"Date: {self.date}",
            f"Description: {self.description}",
            "",
            f"Shots: {self.shots:,}",
            f"Sample Size: {self.sample_size}",
            f"Device Cycle Time: {self.device_cycle_time_ns:.4f} ns",
            f"Sequence Duration: {self.sequence_duration_dt} dt",
            "",
            f"Qubits: {len(self.qubits)}",
            *[f"  - {qubit}" for qubit in self.qubits],
            "",
            f"Expectation Values: {len(self.expectation_values_order)}",
            f"  (States: {set(e.initial_state for e in self.expectation_values_order)})",
            f"  (Observables: {set(e.observable for e in self.expectation_values_order)})",
            "",
            f"Parameter Structure: {self.parameter_structure}",
            f"Tags: {', '.join(self.EXPERIMENT_TAGS)}",
            "=" * 60,
        ]
        return "\n".join(lines)

ExperimentalData dataclass

Dataclass for processing of the characterization dataset. A difference between preprocess and postprocess dataset is that postprocess group expectation values same control parameter id within single row instead of multiple rows.

Source code in src/inspeqtor/v2/data.py

@dataclass
class ExperimentalData:
    """Dataclass for processing of the characterization dataset.
    A difference between preprocess and postprocess dataset is that postprocess group
    expectation values same control parameter id within single row instead of multiple rows.
    """

    config: ExperimentConfiguration
    parameter_dataframe: pl.DataFrame
    observed_dataframe: pl.DataFrame
    mode: typing.Literal["expectation_value", "binary"] = "expectation_value"

    def __post_init__(self):
        self.validate()

    def validate(self):
        assert "parameter_id" in self.parameter_dataframe
        assert "parameter_id" in self.observed_dataframe

        assert (
            self.parameter_dataframe["parameter_id"]
            .unique()
            .sort()
            .equals(self.observed_dataframe["parameter_id"].unique().sort())
        )

    def get_parameter(self) -> jnp.ndarray:
        col_selector = ["/".join(param) for param in self.config.parameter_structure]
        return self.parameter_dataframe[col_selector].to_jax("array")

    def get_observed(self) -> jnp.ndarray:
        col_selector = [str(expval) for expval in self.config.expectation_values_order]

        if self.mode == "binary":
            return jnp.array(
                [
                    calculate_expectation_value_from_binary_dataframe(
                        str(exp), self.observed_dataframe
                    )
                    for exp in self.config.expectation_values_order
                ]
            ).transpose()

        return self.observed_dataframe[col_selector].to_jax("array")

    def save_to_folder(self, path: str | Path):
        if isinstance(path, str):
            path = Path(path)

        path.mkdir(parents=True, exist_ok=True)
        self.config.to_file(path)

        self.parameter_dataframe.write_csv(path / "parameter.csv")
        self.observed_dataframe.write_csv(path / "observed.csv")

    @classmethod
    def from_folder(cls, path: str | Path) -> "ExperimentalData":
        if isinstance(path, str):
            path = Path(path)

        config = ExperimentConfiguration.from_file(path)
        parameter_dataframe = pl.read_csv(path / "parameter.csv")
        observed_dataframe = pl.read_csv(path / "observed.csv")

        return cls(
            config=config,
            parameter_dataframe=parameter_dataframe,
            observed_dataframe=observed_dataframe,
        )

    def __eq__(self, __value: object) -> bool:
        if not isinstance(__value, ExperimentalData):
            return False

        return (
            self.config == __value.config
            and self.parameter_dataframe.equals(__value.parameter_dataframe)
            and self.observed_dataframe.equals(__value.observed_dataframe)
        )

    def __str__(self):
        lines = [
            "=" * 60,
            "EXPERIMENTAL DATA",
            str(self.config),
            "",
            "Parameter DataFrame",
            str(self.parameter_dataframe),
            "",
            "Observed DataFrame",
            str(self.observed_dataframe),
            "=" * 60,
        ]
        return "\n".join(lines)

tensor_product

tensor_product(*operators) -> ndarray

Create tensor product of multiple operators

Source code in src/inspeqtor/v2/data.py

def tensor_product(*operators) -> jnp.ndarray:
    """Create tensor product of multiple operators"""

    result = operators[0]
    for op in operators[1:]:
        result = jnp.kron(result, op)
    return result

get_observable_operator

get_observable_operator(observable: str) -> ndarray

Get the full observable operator as a tensor product

Source code in src/inspeqtor/v2/data.py

def get_observable_operator(observable: str) -> jnp.ndarray:
    """Get the full observable operator as a tensor product"""
    ops = [operator_from_label(label) for label in observable]
    if len(ops) == 1:
        return ops[0]
    return tensor_product(*ops)

get_initial_state

get_initial_state(
    initial_state: str, dm: bool = True
) -> ndarray

Get the initial state as state vector or density matrix

Source code in src/inspeqtor/v2/data.py

def get_initial_state(initial_state: str, dm: bool = True) -> jnp.ndarray:
    """Get the initial state as state vector or density matrix"""
    states = [state_from_label(label, dm=False) for label in initial_state]

    if len(states) == 1:
        state = states[0]
    else:
        # For multi-qubit state, compute the tensor product
        result = states[0]
        for s in states[1:]:
            result = jnp.kron(result, s)
        state = result

    # Convert to vector shape if needed
    if state.shape == (2, 1) or state.shape == (2 ** len(states), 1):
        # Already in correct shape
        pass
    elif state.shape == (2,) or state.shape == (2 ** len(states),):
        # Reshape to column vector
        state = state.reshape(-1, 1)

    if dm:
        return jnp.outer(state, state.conj())
    return state

get_complete_expectation_values

get_complete_expectation_values(
    num_qubits: int,
    observables: Iterable[Literal["I", "X", "Y", "Z"]] = [
        "I",
        "X",
        "Y",
        "Z",
    ],
    states: Iterable[
        Literal["+", "-", "r", "l", "0", "1"]
    ] = ["+", "-", "r", "l", "0", "1"],
    exclude_all_identities: bool = True,
) -> list[ExpectationValue]

Generate a complete set of expectation values for characterizing a multi-qubit system

Source code in src/inspeqtor/v2/data.py

def get_complete_expectation_values(
    num_qubits: int,
    observables: typing.Iterable[typing.Literal["I", "X", "Y", "Z"]] = [
        "I",
        "X",
        "Y",
        "Z",
    ],
    states: typing.Iterable[typing.Literal["+", "-", "r", "l", "0", "1"]] = [
        "+",
        "-",
        "r",
        "l",
        "0",
        "1",
    ],
    exclude_all_identities: bool = True,
) -> list[ExpectationValue]:
    """Generate a complete set of expectation values for characterizing a multi-qubit system"""

    # For n qubits, we need all combinations of observables and states
    result: typing.Iterable[ExpectationValue] = []

    # Generate all combinations of observables
    for obs_combo in itertools.product(observables, repeat=num_qubits):
        for state_combo in itertools.product(states, repeat=num_qubits):
            obs_str = "".join(obs_combo)
            state_str = "".join(state_combo)
            result.append(ExpectationValue(observable=obs_str, initial_state=state_str))

    if exclude_all_identities:
        result = [exp for exp in result if exp.observable != "I" * num_qubits]

    return result

check_parity

check_parity(n)

Determines the parity of a number using bitwise_count.

Efficiently computes parity by counting all 1 bits and taking modulo 2. This is much faster than the iterative approach as it uses hardware intrinsics for population count.

Parameters:

Name	Type	Description	Default
n		The input integer.	required

Returns:

Type	Description
	0 if the number has even parity, 1 if it has odd parity.

Example

check_parity(7) # 0b111 -> three 1s -> odd parity 1 check_parity(6) # 0b110 -> two 1s -> even parity 0

Source code in src/inspeqtor/v2/data.py

def check_parity(n):
    """
    Determines the parity of a number using bitwise_count.

    Efficiently computes parity by counting all 1 bits and taking modulo 2.
    This is much faster than the iterative approach as it uses hardware
    intrinsics for population count.

    Args:
        n: The input integer.

    Returns:
        0 if the number has even parity, 1 if it has odd parity.

    Example:
        >>> check_parity(7)  # 0b111 -> three 1s -> odd parity
        1
        >>> check_parity(6)  # 0b110 -> two 1s -> even parity
        0
    """
    return jnp.bitwise_count(n) % 2

Optimize

src.inspeqtor.v2.optimize

BayesOptState

The dataclass holding optimization state for the gaussian process.

Source code in src/inspeqtor/v2/optimize.py

@struct.dataclass
class BayesOptState:
    """The dataclass holding optimization state for the gaussian process."""

    dataset: gpx.Dataset
    control: ControlSequence

fit_gaussian_process

fit_gaussian_process(D: Dataset)

Fit the Gaussian process given an instance of Dataset

Parameters:

Name	Type	Description	Default
D	Dataset	The gpx.Dataset instance	required

Returns:

Type	Description
	tuple[]: description

Source code in src/inspeqtor/v2/optimize.py

def fit_gaussian_process(D: gpx.Dataset):
    """Fit the Gaussian process given an instance of Dataset

    Args:
        D (gpx.Dataset): The `gpx.Dataset` instance

    Returns:
        tuple[]: _description_
    """
    kernel = gpx.kernels.RBF()  # 1-dimensional input
    meanf = gpx.mean_functions.Zero()
    prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)

    likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n)

    posterior = prior * likelihood

    opt_posterior, history = gpx.fit_scipy(
        model=posterior,
        objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),  # type: ignore
        train_data=D,
        trainable=gpx.parameters.Parameter,
        verbose=True,
    )
    return opt_posterior, history

predict_mean_and_std

predict_mean_and_std(
    x: ndarray, D: Dataset
) -> tuple[ndarray, ndarray]

Predict a Gaussian distribution to the given x using the dataset D

Parameters:

Name	Type	Description	Default
x	ndarray	The array of points to evaluate the gaussian process.	required
D	Dataset	The dataset contain observation from the real process.	required

Returns:

Type	Description
tuple[ndarray, ndarray]	tuple[jnp.ndarray, jnp.ndarray]: The array of mean and standard deviation of the Gaussian process at ponits x.

Source code in src/inspeqtor/v2/optimize.py

def predict_mean_and_std(
    x: jnp.ndarray, D: gpx.Dataset
) -> tuple[jnp.ndarray, jnp.ndarray]:
    """Predict a Gaussian distribution to the given `x` using the dataset `D`

    Args:
        x (jnp.ndarray): The array of points to evaluate the gaussian process.
        D (gpx.Dataset): The dataset contain observation from the real process.

    Returns:
        tuple[jnp.ndarray, jnp.ndarray]: The array of mean and standard deviation of the Gaussian process at ponits `x`.
    """
    opt_posterior, _ = fit_gaussian_process(D)

    return predict_with_gaussian_process(x, opt_posterior, D)

expected_improvement

expected_improvement(
    y_best: ndarray,
    posterior_mean: ndarray,
    posterior_var: ndarray,
    exploration_factor: float,
) -> ndarray

The expected improvement calculated using posterior mean and variance of the gaussian process. The exploration factor can be adjust to balance between exploration and exploitation.

Parameters:

Name	Type	Description	Default
y_best	ndarray	The current maximum value of y	required
posterior_mean	ndarray	The posterior mean of the gaussian process	required
posterior_var	ndarray	The posterior variance of the gaussian process	required
exploration_factor	float	The factor that balance between exploration and exploitation. Set to 0. to maximize exploitation.	required

Returns:

Type	Description
ndarray	jnp.ndarray: The expeced improvement corresponding to the points given from array of the posterior.

Source code in src/inspeqtor/v2/optimize.py

def expected_improvement(
    y_best: jnp.ndarray,
    posterior_mean: jnp.ndarray,
    posterior_var: jnp.ndarray,
    exploration_factor: float,
) -> jnp.ndarray:
    """The expected improvement calculated using posterior mean and variance of the gaussian process.
    The exploration factor can be adjust to balance between exploration and exploitation.


    Args:
        y_best (jnp.ndarray): The current maximum value of y
        posterior_mean (jnp.ndarray): The posterior mean of the gaussian process
        posterior_var (jnp.ndarray): The posterior variance of the gaussian process
        exploration_factor (float): The factor that balance between exploration and exploitation. Set to 0. to maximize exploitation.

    Returns:
        jnp.ndarray: The expeced improvement corresponding to the points given from array of the posterior.
    """
    # https://github.com/alonfnt/bayex/blob/main/bayex/acq.py
    std = jnp.sqrt(posterior_var)
    a = posterior_mean - y_best - exploration_factor
    z = a / std

    return a * norm.cdf(z) + std * norm.pdf(z)

init_opt_state

init_opt_state(x, y, control) -> BayesOptState

Function to intialize the optimizer

Parameters:

Name	Type	Description	Default
x	ndarray	The input arguments	required
y	ndarray	The observation corresponding to the input x	required
control	_type_	The intance of control sequence.	required

Returns:

Name	Type	Description
BayesOptState	BayesOptState	The state of optimizer.

Source code in src/inspeqtor/v2/optimize.py

def init_opt_state(x, y, control) -> BayesOptState:
    """Function to intialize the optimizer

    Args:
        x (jnp.ndarray): The input arguments
        y (jnp.ndarray): The observation corresponding to the input `x`
        control (_type_): The intance of control sequence.

    Returns:
        BayesOptState: The state of optimizer.
    """
    return BayesOptState(dataset=gpx.Dataset(X=x, y=y), control=control)

suggest_next_candidates

suggest_next_candidates(
    key: ndarray,
    opt_state: BayesOptState,
    sample_size: int = 1000,
    num_suggest: int = 1,
    exploration_factor: float = 0.0,
) -> ndarray

Sample new candidates for experiment using expected improvement.

Parameters:

Name	Type	Description	Default
key	ndarray	The jax random key	required
opt_state	BayesOptState	The current optimizer state	required
sample_size	int	The internal number of sample size. Defaults to 1000.	1000
num_suggest	int	The number of suggestion for next experiment. Defaults to 1.	1
exploration_factor	float	The factor that balance between exploration and exploitation. Set to 0. to maximize exploitation. Defaults to 0.0.	0.0

Returns:

Type	Description
ndarray	jnp.ndarray: The suggest data points to evalute in the experiment.

Source code in src/inspeqtor/v2/optimize.py

def suggest_next_candidates(
    key: jnp.ndarray,
    opt_state: BayesOptState,
    sample_size: int = 1000,
    num_suggest: int = 1,
    exploration_factor: float = 0.0,
) -> jnp.ndarray:
    """Sample new candidates for experiment using expected improvement.

    Args:
        key (jnp.ndarray): The jax random key
        opt_state (BayesOptState): The current optimizer state
        sample_size (int, optional): The internal number of sample size. Defaults to 1000.
        num_suggest (int, optional): The number of suggestion for next experiment. Defaults to 1.
        exploration_factor (float, optional): The factor that balance between exploration and exploitation. Set to 0. to maximize exploitation. Defaults to 0.0.

    Returns:
        jnp.ndarray: The suggest data points to evalute in the experiment.
    """
    y = opt_state.dataset.y
    assert isinstance(y, jnp.ndarray)
    y_best = jnp.max(y)

    ravel_fn, unravel_fn = ravel_unravel_fn(opt_state.control.get_structure())
    params = jax.vmap(opt_state.control.sample_params)(
        jax.random.split(key, sample_size)
    )
    # In shape of (sample_size, ctrl_feature)
    ravel_param = jax.vmap(ravel_fn)(params)

    mean, variance = predict_mean_and_std(ravel_param, opt_state.dataset)

    ei = expected_improvement(
        y_best, mean, variance, exploration_factor=exploration_factor
    )

    selected_indice = jnp.argsort(ei, descending=True)[:num_suggest]

    return ravel_param[selected_indice]

add_observations

add_observations(
    opt_state: BayesOptState, x, y
) -> BayesOptState

Function to update the optimization state using new data points x and y

Parameters:

Name	Type	Description	Default
opt_state	BayesOptState	The current optimization state	required
x	ndarray	The input arguments	required
y	ndarray	The observation corresponding to the input x	required

Returns:

Name	Type	Description
BayesOptState	BayesOptState	The updated optimization state.

Source code in src/inspeqtor/v2/optimize.py

def add_observations(opt_state: BayesOptState, x, y) -> BayesOptState:
    """Function to update the optimization state using new data points `x` and `y`

    Args:
        opt_state (BayesOptState): The current optimization state
        x (jnp.ndarray): The input arguments
        y (jnp.ndarray): The observation corresponding to the input `x`

    Returns:
        BayesOptState: The updated optimization state.
    """
    return BayesOptState(
        dataset=opt_state.dataset + gpx.Dataset(X=x, y=y), control=opt_state.control
    )

Predefined

src.inspeqtor.v2.predefined

HamiltonianSpec dataclass

Source code in src/inspeqtor/v2/predefined.py

@dataclass
class HamiltonianSpec:
    method: WhiteboxStrategy
    hamiltonian_enum: HamiltonianEnum = HamiltonianEnum.rotating_transmon_hamiltonian
    # For Trotterization
    trotter_steps: int = 1000
    # For ODE sovler
    max_steps = int(2**16)

    def get_hamiltonian_fn(self):
        if self.hamiltonian_enum == HamiltonianEnum.rotating_transmon_hamiltonian:
            return rotating_transmon_hamiltonian
        elif self.hamiltonian_enum == HamiltonianEnum.transmon_hamiltonian:
            return transmon_hamiltonian
        else:
            raise ValueError(f"Unsupport Hamiltonian: {self.hamiltonian_enum}")

    def get_solver(
        self,
        control_sequence: ControlSequence,
        qubit_info: QubitInformation,
        dt: float,
    ):
        """Return Unitary solver from the given specification of the Hamiltonian and solver

        Args:
            control_sequence (ControlSequence): The control sequence object
            qubit_info (QubitInformation): The qubit information object
            dt (float): The time step size of the device

        Raises:
            ValueError: Unsupport Solver method

        Returns:
            typing.Any: The unitary solver
        """
        if self.method == WhiteboxStrategy.TROTTER:
            hamiltonian = partial(
                self.get_hamiltonian_fn(),
                qubit_info=qubit_info,
                signal=make_signal_fn(
                    get_envelope=control_sequence.get_envelope,
                    drive_frequency=qubit_info.frequency,
                    dt=dt,
                ),
            )

            hamiltonian = ravel_transform(hamiltonian, control_sequence)

            whitebox = make_trotterization_solver(
                hamiltonian=hamiltonian,
                total_dt=control_sequence.total_dt,
                dt=dt,
                trotter_steps=self.trotter_steps,
                y0=jnp.eye(2, dtype=jnp.complex128),
            )

        elif self.method == WhiteboxStrategy.ODE:
            t_eval = jnp.linspace(
                0, control_sequence.total_dt * dt, control_sequence.total_dt
            )

            hamiltonian = partial(
                self.get_hamiltonian_fn(),
                qubit_info=qubit_info,
                signal=make_signal_fn(
                    control_sequence.get_envelope,
                    qubit_info.frequency,
                    dt,
                ),
            )

            hamiltonian = ravel_transform(hamiltonian, control_sequence)

            whitebox = partial(
                solver,
                t_eval=t_eval,
                hamiltonian=hamiltonian,
                y0=jnp.eye(2, dtype=jnp.complex_),
                t0=0,
                t1=control_sequence.total_dt * dt,
                max_steps=self.max_steps,
            )
        else:
            raise ValueError("Unsupport method")

        return whitebox

get_solver

get_solver(
    control_sequence: ControlSequence,
    qubit_info: QubitInformation,
    dt: float,
)

Return Unitary solver from the given specification of the Hamiltonian and solver

Parameters:

Name	Type	Description	Default
control_sequence	ControlSequence	The control sequence object	required
qubit_info	QubitInformation	The qubit information object	required
dt	float	The time step size of the device	required

Raises:

Type	Description
ValueError	Unsupport Solver method

Returns:

Type	Description
	typing.Any: The unitary solver

Source code in src/inspeqtor/v2/predefined.py

def get_solver(
    self,
    control_sequence: ControlSequence,
    qubit_info: QubitInformation,
    dt: float,
):
    """Return Unitary solver from the given specification of the Hamiltonian and solver

    Args:
        control_sequence (ControlSequence): The control sequence object
        qubit_info (QubitInformation): The qubit information object
        dt (float): The time step size of the device

    Raises:
        ValueError: Unsupport Solver method

    Returns:
        typing.Any: The unitary solver
    """
    if self.method == WhiteboxStrategy.TROTTER:
        hamiltonian = partial(
            self.get_hamiltonian_fn(),
            qubit_info=qubit_info,
            signal=make_signal_fn(
                get_envelope=control_sequence.get_envelope,
                drive_frequency=qubit_info.frequency,
                dt=dt,
            ),
        )

        hamiltonian = ravel_transform(hamiltonian, control_sequence)

        whitebox = make_trotterization_solver(
            hamiltonian=hamiltonian,
            total_dt=control_sequence.total_dt,
            dt=dt,
            trotter_steps=self.trotter_steps,
            y0=jnp.eye(2, dtype=jnp.complex128),
        )

    elif self.method == WhiteboxStrategy.ODE:
        t_eval = jnp.linspace(
            0, control_sequence.total_dt * dt, control_sequence.total_dt
        )

        hamiltonian = partial(
            self.get_hamiltonian_fn(),
            qubit_info=qubit_info,
            signal=make_signal_fn(
                control_sequence.get_envelope,
                qubit_info.frequency,
                dt,
            ),
        )

        hamiltonian = ravel_transform(hamiltonian, control_sequence)

        whitebox = partial(
            solver,
            t_eval=t_eval,
            hamiltonian=hamiltonian,
            y0=jnp.eye(2, dtype=jnp.complex_),
            t0=0,
            t1=control_sequence.total_dt * dt,
            max_steps=self.max_steps,
        )
    else:
        raise ValueError("Unsupport method")

    return whitebox

get_gaussian_control_sequence

get_gaussian_control_sequence(
    qubit_info: QubitInformation, max_amp: float = 0.5
)

Get predefined Gaussian control sequence with single Gaussian pulse.

Parameters:

Name	Type	Description	Default
qubit_info	QubitInformation	Qubit information	required
max_amp	float	The maximum amplitude. Defaults to 0.5.	0.5

Returns:

Name	Type	Description
ControlSequence		Control sequence instance

Source code in src/inspeqtor/v2/predefined.py

def get_gaussian_control_sequence(
    qubit_info: QubitInformation,
    max_amp: float = 0.5,  # NOTE: Choice of maximum amplitude is arbitrary
):
    """Get predefined Gaussian control sequence with single Gaussian pulse.

    Args:
        qubit_info (QubitInformation): Qubit information
        max_amp (float, optional): The maximum amplitude. Defaults to 0.5.

    Returns:
        ControlSequence: Control sequence instance
    """
    total_length = 320
    dt = 2 / 9

    control_sequence = ControlSequence(
        controls={
            "gaussian": GaussianPulse(
                duration=total_length,
                qubit_drive_strength=qubit_info.drive_strength,
                dt=dt,
                max_amp=max_amp,
                min_theta=0.0,
                max_theta=2 * jnp.pi,
            )
        },
        total_dt=total_length,
    )

    return control_sequence

get_drag_pulse_v2_sequence

get_drag_pulse_v2_sequence(
    qubit_info_drive_strength: float,
    max_amp: float = 0.5,
    min_theta=0.0,
    max_theta=2 * pi,
    min_beta=-2.0,
    max_beta=2.0,
    dt=2 / 9,
)

Get predefined DRAG control sequence with single DRAG pulse.

Parameters:

Name	Type	Description	Default
qubit_info	QubitInformation	Qubit information	required
max_amp	float	The maximum amplitude. Defaults to 0.5.	0.5

Returns:

Name	Type	Description
ControlSequence		Control sequence instance

Source code in src/inspeqtor/v2/predefined.py

def get_drag_pulse_v2_sequence(
    qubit_info_drive_strength: float,
    max_amp: float = 0.5,  # NOTE: Choice of maximum amplitude is arbitrary
    min_theta=0.0,
    max_theta=2 * jnp.pi,
    min_beta=-2.0,
    max_beta=2.0,
    dt=2 / 9,
):
    """Get predefined DRAG control sequence with single DRAG pulse.

    Args:
        qubit_info (QubitInformation): Qubit information
        max_amp (float, optional): The maximum amplitude. Defaults to 0.5.

    Returns:
        ControlSequence: Control sequence instance
    """
    total_length = 320
    control_sequence = ControlSequence(
        controls={
            "0": DragPulseV2(
                duration=total_length,
                qubit_drive_strength=qubit_info_drive_strength,
                dt=dt,
                max_amp=max_amp,
                min_theta=min_theta,
                max_theta=max_theta,
                min_beta=min_beta,
                max_beta=max_beta,
            ),
        },
        total_dt=total_length,
    )

    return control_sequence

load_data_from_path

load_data_from_path(
    path: str | Path,
    hamiltonian_spec: HamiltonianSpec,
    control_reader=default_control_reader,
) -> LoadedData

Load and prepare the experimental data from given path and hamiltonian spec.

Parameters:

Name	Type	Description	Default
path	str | Path	The path to the folder that contain experimental data.	required
hamiltonian_spec	HamiltonianSpec	The specification of the Hamiltonian	required
control_reader	Any	description. Defaults to default_control_reader.	default_control_reader

Returns:

Name	Type	Description
LoadedData	LoadedData	The object contatin necessary information for device characterization.

Source code in src/inspeqtor/v2/predefined.py

def load_data_from_path(
    path: str | pathlib.Path,
    hamiltonian_spec: HamiltonianSpec,
    control_reader=default_control_reader,
) -> LoadedData:
    """Load and prepare the experimental data from given path and hamiltonian spec.

    Args:
        path (str | pathlib.Path): The path to the folder that contain experimental data.
        hamiltonian_spec (HamiltonianSpec): The specification of the Hamiltonian
        control_reader (typing.Any, optional): _description_. Defaults to default_control_reader.

    Returns:
        LoadedData: The object contatin necessary information for device characterization.
    """
    exp_data = ExperimentalData.from_folder(path)
    control_sequence = control_reader(path)

    assert isinstance(control_sequence, ControlSequence)

    qubit_info = exp_data.config.qubits[0]
    dt = exp_data.config.device_cycle_time_ns

    whitebox = hamiltonian_spec.get_solver(
        control_sequence,
        qubit_info,
        dt,
    )

    return prepare_data(exp_data, control_sequence, whitebox)

save_data_to_path

save_data_to_path(
    path: str | Path,
    experiment_data: ExperimentalData,
    control_sequence: ControlSequence,
)

Save the experimental data to the path

Parameters:

Name	Type	Description	Default
path	str | Path	The path to folder to save the experimental data	required
experiment_data	ExperimentData	The experimental data object	required
control_sequence	ControlSequence	The control sequence that used to create the experimental data.	required

Returns:

Name	Type	Description
None

Source code in src/inspeqtor/v2/predefined.py

def save_data_to_path(
    path: str | pathlib.Path,
    experiment_data: ExperimentalData,
    control_sequence: ControlSequence,
):
    """Save the experimental data to the path

    Args:
        path (str | pathlib.Path): The path to folder to save the experimental data
        experiment_data (ExperimentData): The experimental data object
        control_sequence (ControlSequence): The control sequence that used to create the experimental data.

    Returns:
        None:
    """
    path = pathlib.Path(path)
    path.mkdir(parents=True, exist_ok=True)
    experiment_data.save_to_folder(path)
    control_sequence.to_file(path)

generate_single_qubit_experimental_data

generate_single_qubit_experimental_data(
    key: ndarray,
    hamiltonian: Callable[..., ndarray],
    sample_size: int = 10,
    shots: int = 1000,
    strategy: SimulationStrategy = SHOT,
    qubit_inforamtion: QubitInformation = get_mock_qubit_information(),
    control_sequence: ControlSequence = get_drag_pulse_v2_sequence(
        drive_strength
    ),
    max_steps: int = int(2**16),
    method: WhiteboxStrategy = ODE,
    trotter_steps: int = 1000,
    expectation_value_receipt: list[
        ExpectationValue
    ] = get_complete_expectation_values(1),
) -> tuple[
    ExperimentalData,
    ControlSequence,
    ndarray,
    Callable[[ndarray], ndarray],
]

Generate simulated dataset

Parameters:

Name	Type	Description	Default
key	ndarray	Random key	required
hamiltonian	Callable[..., ndarray]	Total Hamiltonian of the device	required
sample_size	int	Sample size of the control parameters. Defaults to 10.	10
shots	int	Number of shots used to estimate expectation value, will be used if SimulationStrategy is SHOT, otherwise ignored. Defaults to 1000.	1000
strategy	SimulationStrategy	Simulation strategy. Defaults to SimulationStrategy.RANDOM.	SHOT
get_qubit_information_fn	Callable[[], QubitInformation]	Function that return qubit information. Defaults to get_mock_qubit_information.	required
get_control_sequence_fn	Callable[[], ControlSequence]	Function that return control sequence. Defaults to get_multi_drag_control_sequence_v3.	required
max_steps	int	Maximum step of solver. Defaults to int(2**16).	int(2 ** 16)
method	WhiteboxStrategy	Unitary solver method. Defaults to WhiteboxStrategy.ODE.	ODE
trotter_steps	int	Trotterization step. Defualts to 1000	1000

Raises:

Type	Description
NotImplementedError	Not support strategy

Returns:

Type	Description
tuple[ExperimentalData, ControlSequence, ndarray, Callable[[ndarray], ndarray]]	tuple[ExperimentData, ControlSequence, jnp.ndarray, typing.Callable[[jnp.ndarray], jnp.ndarray]]: tuple of (1) Experiment data, (2) Pulse sequence, (3) Noisy unitary, (4) Noisy solver

Source code in src/inspeqtor/v2/predefined.py

def generate_single_qubit_experimental_data(
    key: jnp.ndarray,
    hamiltonian: typing.Callable[..., jnp.ndarray],
    sample_size: int = 10,
    shots: int = 1000,
    strategy: SimulationStrategy = SimulationStrategy.SHOT,
    qubit_inforamtion: QubitInformation = get_mock_qubit_information(),
    control_sequence: ControlSequence = get_drag_pulse_v2_sequence(
        get_mock_qubit_information().drive_strength
    ),
    max_steps: int = int(2**16),
    method: WhiteboxStrategy = WhiteboxStrategy.ODE,
    trotter_steps: int = 1000,
    expectation_value_receipt: list[ExpectationValue] = get_complete_expectation_values(
        1
    ),
) -> tuple[
    ExperimentalData,
    ControlSequence,
    jnp.ndarray,
    typing.Callable[[jnp.ndarray], jnp.ndarray],
]:
    """Generate simulated dataset

    Args:
        key (jnp.ndarray): Random key
        hamiltonian (typing.Callable[..., jnp.ndarray]): Total Hamiltonian of the device
        sample_size (int, optional): Sample size of the control parameters. Defaults to 10.
        shots (int, optional): Number of shots used to estimate expectation value, will be used if `SimulationStrategy` is `SHOT`, otherwise ignored. Defaults to 1000.
        strategy (SimulationStrategy, optional): Simulation strategy. Defaults to SimulationStrategy.RANDOM.
        get_qubit_information_fn (typing.Callable[ [], QubitInformation ], optional): Function that return qubit information. Defaults to get_mock_qubit_information.
        get_control_sequence_fn (typing.Callable[ [], ControlSequence ], optional): Function that return control sequence. Defaults to get_multi_drag_control_sequence_v3.
        max_steps (int, optional): Maximum step of solver. Defaults to int(2**16).
        method (WhiteboxStrategy, optional): Unitary solver method. Defaults to WhiteboxStrategy.ODE.
        trotter_steps (int): Trotterization step. Defualts to 1000

    Raises:
        NotImplementedError: Not support strategy

    Returns:
        tuple[ExperimentData, ControlSequence, jnp.ndarray, typing.Callable[[jnp.ndarray], jnp.ndarray]]: tuple of (1) Experiment data, (2) Pulse sequence, (3) Noisy unitary, (4) Noisy solver
    """
    experiment_config = ExperimentConfiguration(
        qubits=[qubit_inforamtion],
        expectation_values_order=get_complete_expectation_values(1),
        parameter_structure=control_sequence.get_structure(),
        backend_name="stardust",
        sample_size=sample_size,
        shots=shots,
        EXPERIMENT_IDENTIFIER="0001",
        EXPERIMENT_TAGS=["test", "test2"],
        description="This is a test experiment",
        device_cycle_time_ns=2 / 9,
        sequence_duration_dt=control_sequence.total_dt,
    )

    # Generate mock expectation value
    key, exp_key = jax.random.split(key)

    dt = experiment_config.device_cycle_time_ns

    if method == WhiteboxStrategy.TROTTER:
        noisy_simulator = jax.jit(
            make_trotterization_solver(
                hamiltonian=hamiltonian,
                total_dt=control_sequence.total_dt,
                dt=dt,
                trotter_steps=trotter_steps,
                y0=jnp.eye(2, dtype=jnp.complex128),
            )
        )
    else:
        t_eval = jnp.linspace(
            0, control_sequence.total_dt * dt, control_sequence.total_dt
        )
        noisy_simulator = jax.jit(
            partial(
                solver,
                t_eval=t_eval,
                hamiltonian=hamiltonian,
                y0=jnp.eye(2, dtype=jnp.complex64),
                t0=0,
                t1=control_sequence.total_dt * dt,
                max_steps=max_steps,
            )
        )

    key, sample_key = jax.random.split(key)

    ravel_fn, _ = ravel_unravel_fn(control_sequence.get_structure())
    # Sample the parameter by vectorization.
    params_dict = jax.vmap(control_sequence.sample_params)(
        jax.random.split(sample_key, experiment_config.sample_size)
    )
    # Prepare parameter in single line
    control_params = jax.vmap(ravel_fn)(params_dict)

    unitaries = jax.vmap(noisy_simulator)(control_params)
    SHOTS = experiment_config.shots

    # Calculate the expectation values depending on the strategy
    unitaries_f = jnp.asarray(unitaries)[:, -1, :, :]

    assert unitaries_f.shape == (
        sample_size,
        2,
        2,
    ), f"Final unitaries shape is {unitaries_f.shape}"

    if strategy == SimulationStrategy.RANDOM:
        # Just random expectation values with key
        expectation_values = 2 * (
            jax.random.uniform(exp_key, shape=(experiment_config.sample_size, 18))
            - (1 / 2)
        )
    elif strategy == SimulationStrategy.IDEAL:
        expectation_values = calculate_expectation_values(unitaries_f)

    elif strategy == SimulationStrategy.SHOT:
        key, sample_key = jax.random.split(key)
        # The `shot_quantum_device` function will re-calculate the unitary
        expectation_values = single_qubit_shot_quantum_device(
            sample_key,
            control_params,
            noisy_simulator,
            SHOTS,
            expectation_value_receipt,
        )
    else:
        raise NotImplementedError

    assert expectation_values.shape == (
        sample_size,
        18,
    ), f"Expectation values shape is {expectation_values.shape}"

    param_df = pl.DataFrame(
        jax.tree.map(lambda x: np.array(x), flatten_dict(params_dict, sep="/"))
    ).with_row_index("parameter_id")

    obs_df = pl.DataFrame(
        jax.tree.map(
            lambda x: np.array(x),
            flatten_dict(
                dictorization(
                    expectation_values.T, order=get_complete_expectation_values(1)
                ),
                sep="/",
            ),
        )
    ).with_row_index("parameter_id")

    exp_data = ExperimentalData(experiment_config, param_df, obs_df)

    return (
        exp_data,
        control_sequence,
        jnp.array(unitaries),
        noisy_simulator,
    )

Utilities

src.inspeqtor.v2.utils

LoadedData dataclass

A utility dataclass holding objects necessary for device characterization.

Source code in src/inspeqtor/v2/utils.py

@dataclass
class LoadedData:
    """A utility dataclass holding objects necessary for device characterization."""

    experiment_data: ExperimentalData
    control_parameters: jnp.ndarray
    unitaries: jnp.ndarray
    observed_values: jnp.ndarray
    control_sequence: ControlSequence
    whitebox: typing.Callable
    noisy_whitebox: typing.Callable | None = None
    noisy_unitaries: jnp.ndarray | None = None

SyntheticDataModel dataclass

A utility dataclass holding objects necessary for simulating single qubit quantum device.

Source code in src/inspeqtor/v2/utils.py

@dataclass
class SyntheticDataModel:
    """A utility dataclass holding objects necessary for simulating single qubit quantum device."""

    control_sequence: ControlSequence
    qubit_information: QubitInformation
    dt: float
    ideal_hamiltonian: typing.Callable[..., jnp.ndarray]
    total_hamiltonian: typing.Callable[..., jnp.ndarray]
    solver: typing.Callable[..., jnp.ndarray]
    quantum_device: typing.Callable[..., jnp.ndarray] | None
    whitebox: typing.Callable[..., jnp.ndarray] | None

prepare_data

prepare_data(
    exp_data: ExperimentalData,
    control_sequence: ControlSequence,
    whitebox: Callable,
) -> LoadedData

Prepare the data for easy accessing from experiment data, control sequence, and Whitebox.

Parameters:

Name	Type	Description	Default
exp_data	ExperimentData	ExperimentData instance	required
control_sequence	ControlSequence	Control sequence of the experiment	required
whitebox	Callable	Ideal unitary solver.	required

Returns:

Name	Type	Description
LoadedData	LoadedData	LoadedData instance

Source code in src/inspeqtor/v2/utils.py

def prepare_data(
    exp_data: ExperimentalData,
    control_sequence: ControlSequence,
    whitebox: typing.Callable,
) -> LoadedData:
    """Prepare the data for easy accessing from experiment data, control sequence, and Whitebox.

    Args:
        exp_data (ExperimentData): `ExperimentData` instance
        control_sequence (ControlSequence): Control sequence of the experiment
        whitebox (typing.Callable): Ideal unitary solver.

    Returns:
        LoadedData: `LoadedData` instance
    """
    logging.info(f"Loaded data from {exp_data.config.EXPERIMENT_IDENTIFIER}")

    control_parameters = exp_data.get_parameter()

    expectation_values = exp_data.get_observed()
    unitaries = jax.vmap(whitebox)(control_parameters)

    logging.info(
        f"Finished preparing the data for the experiment {exp_data.config.EXPERIMENT_IDENTIFIER}"
    )

    return LoadedData(
        experiment_data=exp_data,
        control_parameters=control_parameters,
        unitaries=unitaries[:, -1, :, :],
        observed_values=expectation_values,
        control_sequence=control_sequence,
        whitebox=whitebox,
    )

single_qubit_shot_quantum_device

single_qubit_shot_quantum_device(
    key: ndarray,
    control_parameters: ndarray,
    solver: Callable[[ndarray], ndarray],
    SHOTS: int,
    expectation_value_receipt: Sequence[
        ExpectationValue
    ] = get_complete_expectation_values(1),
) -> ndarray

This is the shot estimate expectation value quantum device

Parameters:

Name	Type	Description	Default
control_parameters	ndarray	The control parameter to be feed to simlulator	required
key	ndarray	Random key	required
solver	Callable[[ndarray], ndarray]	The ODE solver for propagator	required
SHOTS	int	The number of shots used to estimate expectation values	required

Returns:

Type	Description
ndarray	jnp.ndarray: The expectation value of shape (control_parameters.shape[0], 18)

Source code in src/inspeqtor/v2/utils.py

def single_qubit_shot_quantum_device(
    key: jnp.ndarray,
    control_parameters: jnp.ndarray,
    solver: typing.Callable[[jnp.ndarray], jnp.ndarray],
    SHOTS: int,
    expectation_value_receipt: typing.Sequence[
        ExpectationValue
    ] = get_complete_expectation_values(1),
) -> jnp.ndarray:
    """This is the shot estimate expectation value quantum device

    Args:
        control_parameters (jnp.ndarray): The control parameter to be feed to simlulator
        key (jnp.ndarray): Random key
        solver (typing.Callable[[jnp.ndarray], jnp.ndarray]): The ODE solver for propagator
        SHOTS (int): The number of shots used to estimate expectation values

    Returns:
        jnp.ndarray: The expectation value of shape (control_parameters.shape[0], 18)
    """

    expectation_values = jnp.zeros(
        (control_parameters.shape[0], len(expectation_value_receipt))
    )
    unitaries = jax.vmap(solver)(control_parameters)[:, -1, :, :]

    for idx, exp in enumerate(expectation_value_receipt):
        key, sample_key = jax.random.split(key)
        sample_keys = jax.random.split(sample_key, num=unitaries.shape[0])

        expectation_value = jax.vmap(
            calculate_shots_expectation_value,
            in_axes=(0, None, 0, None, None),
        )(
            sample_keys,
            get_initial_state(exp.initial_state, dm=True),
            unitaries,
            get_observable_operator(exp.observable),
            SHOTS,
        )

        expectation_values = expectation_values.at[..., idx].set(expectation_value)

    return expectation_values

dictorization

dictorization(
    expvals: ndarray, order: list[ExpectationValue]
)

This function formats expectation values of shape (18, N) to a dictionary with the initial state as outer key and the observable as inner key.

Parameters:

Name	Type	Description	Default
expvals	ndarray	Expectation values of shape (18, N). Assumes that order is as in default_expectation_values_order.	required

Returns:

Type	Description
	dict[str, dict[str, jnp.ndarray]]: A dictionary with the initial state as outer key and the observable as inner key.

Source code in src/inspeqtor/v2/utils.py

def dictorization(expvals: jnp.ndarray, order: list[ExpectationValue]):
    """This function formats expectation values of shape (18, N) to a dictionary
    with the initial state as outer key and the observable as inner key.

    Args:
        expvals (jnp.ndarray): Expectation values of shape (18, N). Assumes that order is as in default_expectation_values_order.

    Returns:
        dict[str, dict[str, jnp.ndarray]]: A dictionary with the initial state as outer key and the observable as inner key.
    """
    expvals_dict: dict[str, dict[str, jnp.ndarray]] = {}
    for idx, exp in enumerate(order):
        if exp.initial_state not in expvals_dict:
            expvals_dict[exp.initial_state] = {}

        expvals_dict[exp.initial_state][exp.observable] = expvals[idx]

    return expvals_dict

tensor_product

tensor_product(*operators) -> ndarray

Create tensor product of multiple operators

Source code in src/inspeqtor/v2/utils.py

def tensor_product(*operators) -> jnp.ndarray:
    """Create tensor product of multiple operators"""
    return jax.tree.reduce(jnp.kron, operators)

get_measurement_probability

get_measurement_probability(
    state: ndarray, operator: str
) -> ndarray

Calculate the probability of measuring each projector of tensor product of Pauli operators

Parameters:

Name	Type	Description	Default
state	ndarray	The quantum state to measure	required
operator	str	The string representation of the measurement operator, e.g., 'XY'	required

Returns:

Type	Description
ndarray	jnp.ndarray: An array of probability where each index is a base 10 representation of base 2 measurement result.

Source code in src/inspeqtor/v2/utils.py

def get_measurement_probability(state: jnp.ndarray, operator: str) -> jnp.ndarray:
    """Calculate the probability of measuring each projector of tensor product of Pauli operators

    Args:
        state (jnp.ndarray): The quantum state to measure
        operator (str): The string representation of the measurement operator, e.g., 'XY'

    Returns:
        jnp.ndarray: An array of probability where each index is a base 10 representation of base 2 measurement result.
    """

    return jnp.array(
        [
            jnp.trace(state @ tensor_product(*g_projector))
            for g_projector in product(
                *[(projectors[op][0], projectors[op][1]) for op in operator]
            )
        ]
    ).real