Source code for qugradlab.pulses.composition

 1"""A collection of methods for composing functions and tensors."""
 2
 3import typing
 4from typing import Callable
 5
 6import numpy as np
 7import tensorflow as tf
 8

[docs]
 9def pad(x: np.ndarray[complex],
10        value: complex = 0,
11        left: bool = True,
12        right: bool = True
13       ) -> typing.Any:
14    """Adds padding to a tensor.
15
16    Parameters
17    ----------
18    x: NDArray[Shape[s := Any_Shape]]
19        The tensor to pad
20    value: number
21        The value to pad the tensor with, by default ``0``
22    left: bool
23        Whether to prepend padding, by default ``True``
24    right: bool
25        Whether to append padding, by default ``True``
26
27    Returns
28    -------
29    NDArray[Shape[``s[0]+left+right``, ``s[1:]``]]
30       Padded tensor
31    """
32    if not left and not right:
33        return x
34    padding = value*tf.ones_like(x[0:1])
35    if left and right:
36        return tf.concat([padding, x, padding], axis=0)
37    elif left and not right:
38        return tf.concat([padding, x], axis=0)
39    return tf.concat([x, padding], axis=0)

40

[docs]
41def concatenate_functions(functions: list[Callable]) -> Callable:
42    """Generates a function that executes the listed functions and concatenates
43    the outputs.
44
45    Parameters
46    ----------
47    functions : list[Callable]
48        A list of Callables that return concatenatable outputs.
49
50    Returns
51    -------
52    Callable
53        A function that executes the listed functions and concatenates the
54        outputs.
55
56    Example
57    -------
58    >>> def f1(x):
59    ...     return [x + 1]
60    >>> def f2(x):
61    ...     return [x - 1]
62    >>> a = f1(1)
63    >>> b = f2(2)
64    >>> c = tf.concat([a, b], axis=-1)
65    >>> f3 = concatenate_functions([f1, f2])
66    >>> d = f3([[1], [2]])
67    >>> assert all(c == d)
68    """
69    def execute_and_concatenate(function_arguments):
70        return tf.concat([d(*a) for d, a in zip(functions,
71                                                function_arguments)],
72                         axis=-1)
73    return execute_and_concatenate