Merge pull request #2157 from Ericgig/feature.data.dispatch

Add dispatcher capacity to support any type specialisation

doc/changes/2157.feature

@@ -0,0 +1 @@
+Add capacity to dispatch on ``Data``

qutip/core/_brtools.pyx

@@ -250,10 +250,10 @@ cdef class _EigenBasisTransform:
         return self._evecs_inv
 
     cdef Data _S_converter(self, double t):
-        return _data.kron(self.evecs(t).transpose(), self._inv(t))
+        return _data.kron_transpose(self.evecs(t), self._inv(t))
 
     cdef Data _S_converter_inverse(self, double t):
-        return _data.kron(self._inv(t).transpose(), self.evecs(t))
+        return _data.kron_transpose(self._inv(t), self.evecs(t))
 
     cpdef Data to_eigbasis(self, double t, Data fock):
         """

qutip/core/data/constant.py

@@ -10,7 +10,9 @@
 from .dispatch import Dispatcher as _Dispatcher
 import inspect as _inspect
 
-__all__ = ['zeros', 'identity', 'zeros_like', 'identity_like']
+__all__ = ['zeros', 'identity', 'zeros_like', 'identity_like',
+           'zeros_like_dense', 'identity_like_dense',
+           'zeros_like_data', 'identity_like_data']
 
 zeros = _Dispatcher(
     _inspect.Signature([
@@ -72,24 +74,56 @@
 ], _defer=True)
 
 
-del _Dispatcher, _inspect
+def zeros_like_data(data, /):
+    """
+    Create an zeros matrix of the same type and shape.
+    """
+    return zeros[type(data)](*data.shape)
 
 
-def zeros_like(data, /):
+def zeros_like_dense(data, /):
     """
     Create an zeros matrix of the same type and shape.
     """
-    if type(data) is Dense:
-        return dense.zeros(*data.shape, fortran=data.fortran)
-    return zeros[type(data)](*data.shape)
+    return dense.zeros(*data.shape, fortran=data.fortran)
 
 
-def identity_like(data, /):
+def identity_like_data(data, /):
     """
     Create an identity matrix of the same type and shape.
     """
     if not data.shape[0] == data.shape[1]:
         raise ValueError("Can't create and identity like a non square matrix.")
-    if type(data) is Dense:
-        return dense.identity(data.shape[0], fortran=data.fortran)
     return identity[type(data)](data.shape[0])
+
+
+def identity_like_dense(data, /):
+    """
+    Create an identity matrix of the same type and shape.
+    """
+    if not data.shape[0] == data.shape[1]:
+        raise ValueError("Can't create and identity like a non square matrix.")
+    return dense.identity(data.shape[0], fortran=data.fortran)
+
+
+identity_like = _Dispatcher(
+    identity_like_data, name='identity_like',
+    module=__name__, inputs=("data",), out=False,
+)
+identity_like.add_specialisations([
+    (Data, identity_like_data),
+    (Dense, identity_like_dense),
+], _defer=True)
+
+
+zeros_like = _Dispatcher(
+    zeros_like_data, name='zeros_like',
+    module=__name__, inputs=("data",), out=False,
+)
+zeros_like.add_specialisations([
+    (Data, zeros_like_data),
+    (Dense, zeros_like_dense),
+], _defer=True)
+
+
+del _Dispatcher, _inspect

qutip/core/data/convert.pyx

@@ -19,10 +19,49 @@ import numbers
 import numpy as np
 from scipy.sparse import dok_matrix, csgraph
 cimport cython
+from qutip.core.data.base cimport Data
 
 __all__ = ['to', 'create']
 
 
+class _Epsilon:
+    """
+    Constant for an small weight non-null weight.
+    Use to set `Data` specialisation just over direct specialisation.
+    """
+    def __repr__(self):
+        return "EPSILON"
+
+    def __eq__(self, other):
+        if isinstance(other, _Epsilon):
+            return True
+        return NotImplemented
+
+    def __add__(self, other):
+        if isinstance(other, _Epsilon):
+            return self
+        return other
+
+    def __radd__(self, other):
+        if isinstance(other, _Epsilon):
+            return self
+        return other
+
+    def __lt__(self, other):
+        """ positive number > _Epsilon > 0 """
+        if isinstance(other, _Epsilon):
+            return False
+        return other > 0.
+
+    def __gt__(self, other):
+        if isinstance(other, _Epsilon):
+            return False
+        return other <= 0.
+
+
+EPSILON = _Epsilon()
+
+
 def _raise_if_unconnected(dtype_list, weights):
     unconnected = {}
     for i, type_ in enumerate(dtype_list):
@@ -69,19 +108,23 @@ cdef class _converter:
                 + ">")
 
 
+def identity_converter(arg):
+    return arg
+
+
 cdef class _partial_converter:
     """Convert from any known data-layer type into the type `x.to`."""
 
-    cdef dict converters
+    cdef object converter
     cdef readonly type to
 
-    def __init__(self, converters, to_type):
-        self.converters = dict(converters)
+    def __init__(self, converter, to_type):
+        self.converter = converter
         self.to = to_type
 
     def __call__(self, arg):
         try:
-            return self.converters[type(arg)](arg)
+            return self.converter[self.to, type(arg)](arg)
         except KeyError:
             raise TypeError("unknown type of input: " + str(arg)) from None
 
@@ -200,6 +243,8 @@ cdef class _to:
                 safe just to leave this blank; it is always at best an
                 approximation.  The currently defined weights are accessible in
                 the `weights` attribute of this object.
+                Weight of ~0.001 are should be used in case when no conversion
+                is needed or ``converter = lambda mat : mat``.
         """
         for arg in converters:
             if len(arg) == 3:
@@ -258,6 +303,11 @@ cdef class _to:
                 self.weight[(to_t, from_t)] = weight
                 self._convert[(to_t, from_t)] =\
                     _converter(convert[::-1], to_t, from_t)
+        for dtype in self.dtypes:
+            self.weight[(dtype, Data)] = 1.
+            self.weight[(Data, dtype)] = EPSILON
+            self._convert[(dtype, Data)] = _partial_converter(self, dtype)
+            self._convert[(Data, dtype)] = identity_converter
         for dispatcher in self.dispatchers:
             dispatcher.rebuild_lookup()
 
@@ -309,7 +359,7 @@ cdef class _to:
             type.
         """
         if type(dtype) is type:
-            if dtype not in self.dtypes:
+            if dtype not in self.dtypes and dtype is not Data:
                 raise ValueError(
                     "Type is not a data-layer type: " + repr(dtype))
             return dtype
@@ -333,10 +383,7 @@ cdef class _to:
             raise KeyError(arg)
         to_t = self.parse(arg[0])
         if len(arg) == 1:
-            converters = {
-                from_t: self._convert[to_t, from_t] for from_t in self.dtypes
-            }
-            return _partial_converter(converters, to_t)
+            return _partial_converter(self, to_t)
         from_t = self.parse(arg[1])
         return self._convert[to_t, from_t]
 

qutip/core/data/dispatch.pyx

@@ -7,10 +7,12 @@ import itertools
 import warnings
 
 from .convert import to as _to
+from .convert import EPSILON
 
 cimport cython
 from libc cimport math
 from libcpp cimport bool
+from qutip.core.data.base cimport Data
 
 __all__ = ['Dispatcher']
 
@@ -21,6 +23,8 @@ cdef double _conversion_weight(tuple froms, tuple tos, dict weight_map, bint out
     element-wise to the types in `tos`.  `weight_map` is a mapping of
     `(to_type, from_type): real`; it should almost certainly be
     `data.to.weight`.
+
+    Specialisations that support any types input should use ``Data``.
     """
     cdef double weight = 0.0
     cdef Py_ssize_t i, n=len(froms)
@@ -30,9 +34,9 @@ cdef double _conversion_weight(tuple froms, tuple tos, dict weight_map, bint out
         )
     if out:
         n = n - 1
-        weight += weight_map[froms[n], tos[n]]
+        weight = weight + weight_map[froms[n], tos[n]]
     for i in range(n):
-        weight += weight_map[tos[i], froms[i]]
+        weight = weight + weight_map[tos[i], froms[i]]
     return weight
 
 
@@ -269,14 +273,58 @@ cdef class Dispatcher:
                     + " and a callable"
                 )
             for i in range(self._n_dispatch):
-                if (not _defer) and arg[i] not in _to.dtypes:
+                if (
+                    not _defer
+                    and arg[i] not in _to.dtypes
+                    and arg[i] is not Data
+                ):
                     raise ValueError(str(arg[i]) + " is not a known data type")
             if not callable(arg[self._n_dispatch]):
                 raise TypeError(str(arg[-1]) + " is not callable")
             self._specialisations[arg[:-1]] = arg[-1]
         if not _defer:
             self.rebuild_lookup()
 
+    cdef object _find_specialization(self, tuple in_types, bint output):
+        # The complexity of building the table here is very poor, but it's a
+        # cost we pay very infrequently, and until it's proved to be a
+        # bottle-neck in real code, we stick with the simple algorithm.
+        cdef double weight, cur
+        cdef tuple types, out_types, displayed_type
+        cdef object function
+        cdef int n_dispatch
+        weight = math.INFINITY
+        types = None
+        function = None
+        n_dispatch = len(in_types)
+        for out_types, out_function in self._specialisations.items():
+            cur = _conversion_weight(
+                in_types, out_types[:n_dispatch], _to.weight, out=output)
+            if cur < weight:
+                weight = cur
+                types = out_types
+                function = out_function
+
+        if cur == math.INFINITY:
+            raise ValueError("No valid specialisations found")
+
+        if weight in [EPSILON, 0.] and not (output and types[-1] is Data):
+            self._lookup[in_types] = function
+        else:
+            if output:
+                converters = tuple(
+                    [_to[pair] for pair in zip(types[:-1], in_types[:-1])]
+                    + [_to[in_types[-1], types[-1]]]
+                )
+            else:
+                converters = tuple(_to[pair] for pair in zip(types, in_types))
+            displayed_type = in_types
+            if len(in_types) < len(types):
+                displayed_type = displayed_type + (types[-1],)
+            self._lookup[in_types] =\
+                _constructed_specialisation(function, self, displayed_type,
+                                            converters, output)
+
     def rebuild_lookup(self):
         """
         Manually trigger a rebuild of the lookup table for this dispatcher.
@@ -286,67 +334,17 @@ cdef class Dispatcher:
 
         You most likely do not need to call this function yourself.
         """
-        cdef double weight, cur
-        cdef tuple types, out_types
-        cdef object function
         if not self._specialisations:
             return
         self._dtypes = _to.dtypes.copy()
-        # The complexity of building the table here is very poor, but it's a
-        # cost we pay very infrequently, and until it's proved to be a
-        # bottle-neck in real code, we stick with the simple algorithm.
         for in_types in itertools.product(self._dtypes, repeat=self._n_dispatch):
-            weight = math.INFINITY
-            types = None
-            function = None
-            for out_types, out_function in self._specialisations.items():
-                cur = _conversion_weight(in_types, out_types, _to.weight,
-                                         out=self.output)
-                if cur < weight:
-                    weight = cur
-                    types = out_types
-                    function = out_function
-
-            if weight == 0:
-                self._lookup[in_types] = function
-            else:
-                if self.output:
-                    converters = tuple(
-                        [_to[pair] for pair in zip(types[:-1], in_types[:-1])]
-                        + [_to[in_types[-1], types[-1]]]
-                    )
-                else:
-                    converters = tuple(_to[pair] for pair in zip(types, in_types))
-                self._lookup[in_types] =\
-                    _constructed_specialisation(function, self, in_types,
-                                                converters, self.output)
+            self._find_specialization(in_types, self.output)
         # Now build the lookup table in the case that we dispatch on the output
         # type as well, but the user has called us without specifying it.
         # TODO: option to control default output type choice if unspecified?
         if self.output:
             for in_types in itertools.product(self._dtypes, repeat=self._n_dispatch-1):
-                weight = math.INFINITY
-                types = None
-                function = None
-                for out_types, out_function in self._specialisations.items():
-                    cur = _conversion_weight(in_types, out_types[:-1],
-                                             _to.weight, out=False)
-                    if cur < weight:
-                        weight = cur
-                        types = out_types
-                        function = out_function
-
-                if weight == 0:
-                    self._lookup[in_types] = function
-                else:
-                    converters = tuple(
-                        _to[pair]
-                        for pair in zip(types, in_types)
-                    )
-                    self._lookup[in_types] =\
-                        _constructed_specialisation(function, self,
-                                                    in_types + (types[-1],),
-                                                    converters, False)
+                self._find_specialization(in_types, False)
 
     def __getitem__(self, types):
         """