Merge pull request #965 from qiboteam/derivatives

`Parameter` class

src/qibo/__init__.py

@@ -2,7 +2,16 @@
 
 __version__ = im.version(__package__)
 
-from qibo import callbacks, gates, hamiltonians, models, optimizers, parallel, solvers
+from qibo import (
+    callbacks,
+    gates,
+    hamiltonians,
+    models,
+    optimizers,
+    parallel,
+    parameter,
+    solvers,
+)
 from qibo.backends import (
     get_backend,
     get_device,

src/qibo/gates/gates.py

@@ -5,6 +5,7 @@
 
 from qibo.config import PRECISION_TOL, raise_error
 from qibo.gates.abstract import Gate, ParametrizedGate
+from qibo.parameter import Parameter
 
 
 class H(Gate):
@@ -497,6 +498,10 @@ def __init__(self, q, theta, trainable=True):
         self.target_qubits = (q,)
         self.unitary = True
 
+        self.initparams = theta
+        if isinstance(theta, Parameter):
+            theta = theta()
+
         if isinstance(theta, (float, int)) and (theta % (np.pi / 2)).is_integer():
             self.clifford = True
 

src/qibo/parameter.py

@@ -0,0 +1,131 @@
+import numpy as np
+import sympy as sp
+
+from qibo.config import raise_error
+
+
+def calculate_derivatives(func):
+    """Calculates derivatives w.r.t to all parameters of a target function `func`."""
+    vars = []
+    for i in range(func.__code__.co_argcount):
+        vars.append(sp.Symbol(f"p{i}"))
+
+    expr = sp.sympify(func(*vars))
+
+    derivatives = []
+    for i in range(len(vars)):
+        derivative_expr = sp.diff(expr, vars[i])
+        derivatives.append(sp.lambdify(vars, derivative_expr))
+
+    return derivatives
+
+
+class Parameter:
+    """Object which allows for variational gate parameters. Several trainable parameters
+    and possibly features are linked through a lambda function which returns the
+    final gate parameter. All possible analytical derivatives of the lambda function are
+    calculated at the object initialisation using Sympy.
+
+    Example::
+
+        from qibo.parameter import Parameter
+        param = Parameter(
+                lambda x, th1, th2, th3: x**2 * th1 + th2 * th3**2,
+                features=[7.0],
+                trainable=[1.5, 2.0, 3.0],
+            )
+
+        partial_derivative = param.get_partial_derivative(3)
+
+        param.update_parameters(trainable=[15.0, 10.0, 7.0], feature=[5.0])
+        param_value = param()
+
+
+    Args:
+        func (function): lambda function which builds the gate parameter. If both features and trainable parameters
+        compose the function, it must be passed by first providing the features and then the parameters, as
+        described in the code example above.
+        features (list or np.ndarray): array containing possible input features x.
+        trainable (list or np.ndarray): array with initial trainable parameters theta.
+        nofeatures (bool): flag to explicitly ban the updating of the features. This simplifies the task of updating Parameter objects simultaneously when some have embedded features and some do not.
+    """
+
+    def __init__(self, func, trainable=None, features=None):
+        self.trainable = trainable if trainable is not None else []
+        self.features = features if features is not None else []
+
+        if self.nfeat + self.nparams != func.__code__.co_argcount:
+            raise_error(
+                TypeError,
+                f"{self.nfeat + self.nparams} parameters are provided, but {func.__code__.co_argcount} are required, please initialize features and trainable according to the defined function.",
+            )
+        # lambda function
+        self.lambdaf = func
+
+        # calculate derivatives
+        # maybe here use JAX ?
+        self.derivatives = calculate_derivatives(func=self.lambdaf)
+
+    def __call__(self, features=None, trainable=None):
+        """Return parameter value with given features and/or trainable."""
+
+        params = []
+
+        if features is None:
+            params.extend(self.features)
+        else:
+            if len(features) != self.nfeat:
+                raise_error(
+                    TypeError,
+                    f"The number of features provided is not compatible with the problem's dimensionality, which is {self.nfeat}.",
+                )
+            else:
+                params.extend(features)
+        if trainable is None:
+            params.extend(self.trainable)
+        else:
+            if len(trainable) != self.nparams:
+                raise_error(
+                    TypeError,
+                    f"The number of trainable provided is different from the number of required parameters, which is {self.nparams}.",
+                )
+            else:
+                params.extend(trainable)
+
+        return self.lambdaf(*params)
+
+    @property
+    def nparams(self):
+        """Returns the number of trainable parameters"""
+        return len(self.trainable)
+
+    @property
+    def nfeat(self):
+        """Returns the number of features"""
+        return len(self.features)
+
+    @property
+    def ncomponents(self):
+        """Returns the number of elements which compose the Parameter"""
+        return self.nparams + self.nfeat
+
+    def trainable_parameter_indices(self, start_index):
+        """Return list of respective indices of trainable parameters within
+        the larger trainable parameter list of a circuit for example"""
+        return (np.arange(self.nparams) + start_index).tolist()
+
+    def unaffected_by(self, trainable_idx):
+        """Retrieve constant term of lambda function with regard to a specific trainable parameter"""
+        params = self.trainable.copy()
+        params[trainable_idx] = 0.0
+        return self(trainable=params)
+
+    def partial_derivative(self, trainable_idx):
+        """Get derivative w.r.t a trainable parameter"""
+        deriv = self.derivatives[trainable_idx]
+
+        params = []
+        params.extend(self.features)
+        params.extend(self.trainable)
+
+        return deriv(*params)

tests/test_gates_gates.py

@@ -3,6 +3,7 @@
 import pytest
 
 from qibo import gates
+from qibo.parameter import Parameter
 from qibo.quantum_info import random_hermitian, random_statevector, random_unitary
 
 
@@ -253,6 +254,17 @@ def test_rx(backend, theta):
     else:
         assert not gates.RX(0, theta=theta).clifford
 
+    # test Parameter
+    assert (
+        gates.RX(
+            0,
+            theta=Parameter(
+                lambda x, th1: 10 * th1 + x, trainable=[0.2], features=[40]
+            ),
+        ).init_kwargs["theta"]
+        == 42
+    )
+
 
 @pytest.mark.parametrize("theta", [np.random.rand(), np.pi / 2, -np.pi / 2, np.pi])
 def test_ry(backend, theta):

tests/test_parameter.py

@@ -0,0 +1,126 @@
+import numpy as np
+import pytest
+
+from qibo.parameter import Parameter
+
+
+def test_parameter():
+    # single feature
+    param = Parameter(
+        lambda x, th1, th2, th3: x**2 * th1 + th2 * th3**2,
+        features=[7.0],
+        trainable=[1.5, 2.0, 3.0],
+    )
+
+    indices = param.trainable_parameter_indices(10)
+    assert indices == [10, 11, 12]
+
+    fixed = param.unaffected_by(1)
+    assert fixed == 73.5
+
+    factor = param.partial_derivative(3)
+    assert factor == 12.0
+
+    param.trainable = [15.0, 10.0, 7.0]
+    param.features = [5.0]
+    gate_value = param()
+    assert gate_value == 865
+
+    # single feature, no list
+    param2 = Parameter(
+        lambda x, th1, th2, th3: x**2 * th1 + th2 * th3**2,
+        features=[7.0],
+        trainable=[1.5, 2.0, 3.0],
+    )
+
+    gate_value2 = param2()
+    assert gate_value2 == 91.5
+
+    # multiple features
+    param = Parameter(
+        lambda x1, x2, th1, th2, th3: x1**2 * th1 + x2 * th2 * th3,
+        features=[7.0, 4.0],
+        trainable=[1.5, 2.0, 3.0],
+    )
+
+    fixed = param.unaffected_by(1)
+    assert fixed == 73.5
+    assert param.nparams == 3
+    assert param.nfeat == 2
+
+    factor = param.partial_derivative(4)
+    assert factor == 8.0
+
+    param.trainable = np.array([15.0, 10.0, 7.0])
+    param.features = [5.0, 3.0]
+    gate_value = param()
+    assert gate_value == 585
+
+    # testing call with new values
+    executed = param(features=[0.5, 2.0], trainable=[2.0, 0.1, 4.0])
+    assert executed == 1.3
+
+    # injecting only trainable
+    param = Parameter(lambda x: x, trainable=[0.8])
+    nparams = param.nparams
+    nfeat = param.nfeat
+    ncomponents = param.ncomponents
+    assert nparams == 1
+    assert nfeat == 0
+    assert ncomponents == 1
+
+    # injecting only features
+    param = Parameter(lambda x: x, features=[0.8])
+    nparams = param.nparams
+    nfeat = param.nfeat
+    assert nparams == 0
+    assert nfeat == 1
+
+
+def test_parameter_errors():
+    param = Parameter(
+        lambda x, th1, th2, th3: x**2 * th1 + th2 * th3**2,
+        features=[7.0],
+        trainable=[1.5, 2.0, 3.0],
+    )
+
+    param.trainable = [1, 1, 1]
+    param.features = 1
+
+    try:
+        param()
+        assert False
+    except Exception as e:
+        assert True
+
+    param.trainable = [1, 1]
+    param.features = [1]
+    try:
+        param()
+        assert False
+    except Exception as e:
+        assert True
+
+    param.trainable = [1, 1, 1]
+    param.features = [1, 1]
+    try:
+        param()
+        assert False
+    except Exception as e:
+        assert True
+
+    # test type error due to wrong initialization
+    with pytest.raises(TypeError):
+        param = Parameter(func=lambda x, y: x + y**2)
+
+    # test call function with wrong features and trainable dimensionality
+    param = Parameter(
+        func=lambda x, th1, th2: th1 * x + th2, features=[1.2], trainable=[0.2, 9.1]
+    )
+
+    # wrong features length
+    with pytest.raises(TypeError):
+        param(features=[2.3, 9.2], trainable=[0.4, 9.3])
+    # wrong trainable length
+    with pytest.raises(TypeError):
+        param(features=[0.4], trainable=[3.4, 0.1, 5.6])