add mean and trig methods on backend

Author: refraction-ray

CHANGELOG.md

@@ -2,6 +2,14 @@
 
 ## Unreleased
 
+### Added
+
+- add `utils.append` to build function pipeline
+
+- add `mean` method on backends
+
+- add trigonometric methods on backends
+
 ## 0.0.220328
 
 ### Added

tensorcircuit/backends/abstract_backend.py

@@ -97,6 +97,138 @@ def cos(self: Any, a: Tensor) -> Tensor:
             "Backend '{}' has not implemented `cos`.".format(self.name)
         )
 
+    def acos(self: Any, a: Tensor) -> Tensor:
+        """
+        Return the acos of a tensor ``a``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: acos of ``a``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `acos`.".format(self.name)
+        )
+
+    def acosh(self: Any, a: Tensor) -> Tensor:
+        """
+        Return the acosh of a tensor ``a``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: acosh of ``a``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `acosh`.".format(self.name)
+        )
+
+    def asin(self: Any, a: Tensor) -> Tensor:
+        """
+        Return the acos of a tensor ``a``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: asin of ``a``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `asin`.".format(self.name)
+        )
+
+    def asinh(self: Any, a: Tensor) -> Tensor:
+        """
+        Return the asinh of a tensor ``a``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: asinh of ``a``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `asinh`.".format(self.name)
+        )
+
+    def atan(self: Any, a: Tensor) -> Tensor:
+        """
+        Return the atan of a tensor ``a``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: atan of ``a``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `atan`.".format(self.name)
+        )
+
+    def atan2(self: Any, y: Tensor, x: Tensor) -> Tensor:
+        """
+        Return the atan of a tensor ``y``/``x``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: atan2 of ``a``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `atan2`.".format(self.name)
+        )
+
+    def cosh(self: Any, a: Tensor) -> Tensor:
+        """
+        Return the cosh of a tensor ``a``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: cosh of ``a``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `cosh`.".format(self.name)
+        )
+
+    def tan(self: Any, a: Tensor) -> Tensor:
+        """
+        Return the tan of a tensor ``a``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: tan of ``a``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `tan`.".format(self.name)
+        )
+
+    def tanh(self: Any, a: Tensor) -> Tensor:
+        """
+        Return the tanh of a tensor ``a``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: tanh of ``a``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `tanh`.".format(self.name)
+        )
+
+    def sinh(self: Any, a: Tensor) -> Tensor:
+        """
+        Return the sinh of a tensor ``a``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: sinh of ``a``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `sinh`.".format(self.name)
+        )
+
+    # acos acosh asin asinh atan atan2 atanh cosh (cos) tan tanh (sin) sinh
+
     def abs(self: Any, a: Tensor) -> Tensor:
         """
         Return the elementwise abs value of a matrix ``a``.
@@ -283,6 +415,28 @@ def tile(self: Any, a: Tensor, rep: Tensor) -> Tensor:
             "Backend '{}' has not implemented `tile`.".format(self.name)
         )
 
+    def mean(
+        self: Any,
+        a: Tensor,
+        axis: Optional[Sequence[int]] = None,
+        keepdims: bool = False,
+    ) -> Tensor:
+        """
+        Compute the arithmetic mean for ``a`` along the specified ``axis``.
+
+        :param a: tensor to take average
+        :type a: Tensor
+        :param axis: the axis to take mean, defaults to None indicating sum over flatten array
+        :type axis: Optional[Sequence[int]], optional
+        :param keepdims: _description_, defaults to False
+        :type keepdims: bool, optional
+        :return: _description_
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `mean`.".format(self.name)
+        )
+
     def min(self: Any, a: Tensor, axis: Optional[int] = None) -> Tensor:
         """
         Return the minimum of an array or minimum along an axis.

tensorcircuit/backends/jax_backend.py

@@ -247,6 +247,39 @@ def sin(self, a: Tensor) -> Tensor:
     def cos(self, a: Tensor) -> Tensor:
         return jnp.cos(a)
 
+    def acos(self, a: Tensor) -> Tensor:
+        return jnp.arccos(a)
+
+    def acosh(self, a: Tensor) -> Tensor:
+        return jnp.arccosh(a)
+
+    def asin(self, a: Tensor) -> Tensor:
+        return jnp.arcsin(a)
+
+    def asinh(self, a: Tensor) -> Tensor:
+        return jnp.arcsinh(a)
+
+    def atan(self, a: Tensor) -> Tensor:
+        return jnp.arctan(a)
+
+    def atan2(self, y: Tensor, x: Tensor) -> Tensor:
+        return jnp.arctan2(y, x)
+
+    def atanh(self, a: Tensor) -> Tensor:
+        return jnp.arctanh(a)
+
+    def cosh(self, a: Tensor) -> Tensor:
+        return jnp.cosh(a)
+
+    def tan(self, a: Tensor) -> Tensor:
+        return jnp.tan(a)
+
+    def tanh(self, a: Tensor) -> Tensor:
+        return jnp.tanh(a)
+
+    def sinh(self, a: Tensor) -> Tensor:
+        return jnp.sinh(a)
+
     def size(self, a: Tensor) -> Tensor:
         return jnp.size(a)
 
@@ -288,6 +321,14 @@ def concat(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
     def tile(self, a: Tensor, rep: Tensor) -> Tensor:
         return jnp.tile(a, rep)
 
+    def mean(
+        self,
+        a: Tensor,
+        axis: Optional[Sequence[int]] = None,
+        keepdims: bool = False,
+    ) -> Tensor:
+        return jnp.mean(a, axis=axis, keepdims=keepdims)
+
     def min(self, a: Tensor, axis: Optional[int] = None) -> Tensor:
         return jnp.min(a, axis=axis)
 

tensorcircuit/backends/numpy_backend.py

@@ -91,6 +91,40 @@ def sin(self, a: Tensor) -> Tensor:
     def cos(self, a: Tensor) -> Tensor:
         return np.cos(a)
 
+    # acos acosh asin asinh atan atan2 atanh cosh (cos) tan tanh sinh (sin)
+    def acos(self, a: Tensor) -> Tensor:
+        return np.arccos(a)
+
+    def acosh(self, a: Tensor) -> Tensor:
+        return np.arccosh(a)
+
+    def asin(self, a: Tensor) -> Tensor:
+        return np.arcsin(a)
+
+    def asinh(self, a: Tensor) -> Tensor:
+        return np.arcsinh(a)
+
+    def atan(self, a: Tensor) -> Tensor:
+        return np.arctan(a)
+
+    def atan2(self, y: Tensor, x: Tensor) -> Tensor:
+        return np.arctan2(y, x)
+
+    def atanh(self, a: Tensor) -> Tensor:
+        return np.arctanh(a)
+
+    def cosh(self, a: Tensor) -> Tensor:
+        return np.cosh(a)
+
+    def tan(self, a: Tensor) -> Tensor:
+        return np.tan(a)
+
+    def tanh(self, a: Tensor) -> Tensor:
+        return np.tanh(a)
+
+    def sinh(self, a: Tensor) -> Tensor:
+        return np.sinh(a)
+
     def size(self, a: Tensor) -> Tensor:
         return a.size
 
@@ -116,6 +150,14 @@ def concat(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
     def tile(self, a: Tensor, rep: Tensor) -> Tensor:
         return np.tile(a, rep)
 
+    def mean(
+        self,
+        a: Tensor,
+        axis: Optional[Sequence[int]] = None,
+        keepdims: bool = False,
+    ) -> Tensor:
+        return np.mean(a, axis=axis, keepdims=keepdims)
+
     def unique_with_counts(self, a: Tensor) -> Tuple[Tensor, Tensor]:
         return np.unique(a, return_counts=True)  # type: ignore
 

tensorcircuit/backends/pytorch_backend.py

@@ -217,6 +217,39 @@ def sin(self, a: Tensor) -> Tensor:
     def cos(self, a: Tensor) -> Tensor:
         return torchlib.cos(a)
 
+    def acos(self, a: Tensor) -> Tensor:
+        return torchlib.acos(a)
+
+    def acosh(self, a: Tensor) -> Tensor:
+        return torchlib.acosh(a)
+
+    def asin(self, a: Tensor) -> Tensor:
+        return torchlib.asin(a)
+
+    def asinh(self, a: Tensor) -> Tensor:
+        return torchlib.asinh(a)
+
+    def atan(self, a: Tensor) -> Tensor:
+        return torchlib.atan(a)
+
+    def atan2(self, y: Tensor, x: Tensor) -> Tensor:
+        return torchlib.atan2(y, x)
+
+    def atanh(self, a: Tensor) -> Tensor:
+        return torchlib.atanh(a)
+
+    def cosh(self, a: Tensor) -> Tensor:
+        return torchlib.cosh(a)
+
+    def tan(self, a: Tensor) -> Tensor:
+        return torchlib.tan(a)
+
+    def tanh(self, a: Tensor) -> Tensor:
+        return torchlib.tanh(a)
+
+    def sinh(self, a: Tensor) -> Tensor:
+        return torchlib.sinh(a)
+
     def size(self, a: Tensor) -> Tensor:
         return a.size()
 
@@ -258,6 +291,16 @@ def concat(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
     def tile(self, a: Tensor, rep: Tensor) -> Tensor:
         return torchlib.tile(a, rep)
 
+    def mean(
+        self,
+        a: Tensor,
+        axis: Optional[Sequence[int]] = None,
+        keepdims: bool = False,
+    ) -> Tensor:
+        if axis is None:
+            axis = tuple([i for i in range(len(a.shape))])
+        return torchlib.mean(a, dim=axis, keepdim=keepdims)
+
     def min(self, a: Tensor, axis: Optional[int] = None) -> Tensor:
         if axis is None:
             return torchlib.min(a)

tensorcircuit/backends/tensorflow_backend.py

@@ -278,6 +278,39 @@ def sin(self, a: Tensor) -> Tensor:
     def cos(self, a: Tensor) -> Tensor:
         return tf.math.cos(a)
 
+    def acos(self, a: Tensor) -> Tensor:
+        return tf.math.acos(a)
+
+    def acosh(self, a: Tensor) -> Tensor:
+        return tf.math.acosh(a)
+
+    def asin(self, a: Tensor) -> Tensor:
+        return tf.math.asin(a)
+
+    def asinh(self, a: Tensor) -> Tensor:
+        return tf.math.asinh(a)
+
+    def atan(self, a: Tensor) -> Tensor:
+        return tf.math.atan(a)
+
+    def atan2(self, y: Tensor, x: Tensor) -> Tensor:
+        return tf.math.atan2(y, x)
+
+    def atanh(self, a: Tensor) -> Tensor:
+        return tf.math.atanh(a)
+
+    def cosh(self, a: Tensor) -> Tensor:
+        return tf.math.cosh(a)
+
+    def tan(self, a: Tensor) -> Tensor:
+        return tf.math.tan(a)
+
+    def tanh(self, a: Tensor) -> Tensor:
+        return tf.math.tanh(a)
+
+    def sinh(self, a: Tensor) -> Tensor:
+        return tf.math.sinh(a)
+
     def size(self, a: Tensor) -> Tensor:
         return tf.size(a)
 
@@ -325,6 +358,14 @@ def concat(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
     def tile(self, a: Tensor, rep: Tensor) -> Tensor:
         return tf.tile(a, rep)
 
+    def mean(
+        self,
+        a: Tensor,
+        axis: Optional[Sequence[int]] = None,
+        keepdims: bool = False,
+    ) -> Tensor:
+        return tf.math.reduce_mean(a, axis=axis, keepdims=keepdims)
+
     def sigmoid(self, a: Tensor) -> Tensor:
         return tf.nn.sigmoid(a)
 

tensorcircuit/circuit.py

@@ -761,7 +761,7 @@ def depolarizing(
     ) -> float:
         # px/y/z here not support differentiation for now
         # jit compatible for now
-        assert px + py + pz < 1 and px >= 0 and py >= 0 and pz >= 0
+        # assert px + py + pz < 1 and px >= 0 and py >= 0 and pz >= 0
 
         def step_function(x: Tensor) -> Tensor:
             r = (

tensorcircuit/utils.py

@@ -53,3 +53,31 @@ def wrapper(*args: Any, **kws: Any) -> Any:
         return tuple(nr)
 
     return wrapper
+
+
+def append(f: Callable[..., Any], *op: Callable[..., Any]) -> Any:
+    """
+    Functional programming paradigm to build function pipeline
+
+    :Example:
+
+    >>> f = tc.utils.append(lambda x: x**2, lambda x: x+1, tc.backend.mean)
+    >>> f(tc.backend.ones(2))
+    (2+0j)
+
+    :param f: The function which are attached with other functions
+    :type f: Callable[..., Any]
+    :param op: Function to be attached
+    :type op: Callable[..., Any]
+    :return: The final results after function pipeline
+    :rtype: Any
+    """
+
+    @wraps(f)
+    def wrapper(*args: Any, **kws: Any) -> Any:
+        rs = f(*args, **kws)
+        for opi in op:
+            rs = opi(rs)
+        return rs
+
+    return wrapper

tests/test_backends.py

@@ -179,6 +179,71 @@ def test_backend_methods(backend):
     )
 
 
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
+def test_backend_methods_2(backend):
+    np.testing.assert_allclose(tc.backend.mean(tc.backend.ones([10])), 1.0, atol=1e-5)
+    # acos acosh asin asinh atan atan2 atanh cosh (cos) tan tanh sinh (sin)
+    np.testing.assert_allclose(
+        tc.backend.acos(tc.backend.ones([2], dtype="float32")),
+        np.arccos(tc.backend.ones([2])),
+        atol=1e-5,
+    )
+    np.testing.assert_allclose(
+        tc.backend.acosh(tc.backend.ones([2], dtype="float32")),
+        np.arccosh(tc.backend.ones([2])),
+        atol=1e-5,
+    )
+    np.testing.assert_allclose(
+        tc.backend.asin(tc.backend.ones([2], dtype="float32")),
+        np.arcsin(tc.backend.ones([2])),
+        atol=1e-5,
+    )
+    np.testing.assert_allclose(
+        tc.backend.asinh(tc.backend.ones([2], dtype="float32")),
+        np.arcsinh(tc.backend.ones([2])),
+        atol=1e-5,
+    )
+    np.testing.assert_allclose(
+        tc.backend.atan(0.5 * tc.backend.ones([2], dtype="float32")),
+        np.arctan(0.5 * tc.backend.ones([2])),
+        atol=1e-5,
+    )
+    np.testing.assert_allclose(
+        tc.backend.atan2(
+            tc.backend.ones([1], dtype="float32"), tc.backend.ones([1], dtype="float32")
+        ),
+        np.arctan2(
+            tc.backend.ones([1], dtype="float32"), tc.backend.ones([1], dtype="float32")
+        ),
+        atol=1e-5,
+    )
+    np.testing.assert_allclose(
+        tc.backend.atanh(0.5 * tc.backend.ones([2], dtype="float32")),
+        np.arctanh(0.5 * tc.backend.ones([2])),
+        atol=1e-5,
+    )
+    np.testing.assert_allclose(
+        tc.backend.cosh(tc.backend.ones([2], dtype="float32")),
+        np.cosh(tc.backend.ones([2])),
+        atol=1e-5,
+    )
+    np.testing.assert_allclose(
+        tc.backend.tan(tc.backend.ones([2], dtype="float32")),
+        np.tan(tc.backend.ones([2])),
+        atol=1e-5,
+    )
+    np.testing.assert_allclose(
+        tc.backend.tanh(tc.backend.ones([2], dtype="float32")),
+        np.tanh(tc.backend.ones([2])),
+        atol=1e-5,
+    )
+    np.testing.assert_allclose(
+        tc.backend.sinh(0.5 * tc.backend.ones([2], dtype="float32")),
+        np.sinh(0.5 * tc.backend.ones([2])),
+        atol=1e-5,
+    )
+
+
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
 def test_arg_cmp(backend):
     np.testing.assert_allclose(tc.backend.argmax(tc.backend.ones([3], "float64")), 0)