added complex functions to scalar basic

theano/scalar/basic.py

@@ -459,6 +459,32 @@ def float_out_nocomplex(*types):
         if type in complex_types:
             raise TypeError('complex argument not supported')
     return float64,
+class unary_out_lookup(gof.utils.object2):
+    """
+    get a output_types_preference object by passing a dictionary:
+
+    unary_out_lookup({int8:int32, float32:complex128})
+
+    The result is an op that maps in8 to int32 and float32 to complex128 and other input types
+    lead to a TypeError.
+    """
+    def __init__(self, type_table):
+        self.tbl = type_table
+    def __call__(self, *types):
+        if len(types) == 1:
+            types = types[0]
+        try:
+            rval = self.tbl[types]
+        except:
+            raise TypeError(types)
+        if isinstance(types, (list, tuple)):
+            return rval
+        else:
+            return [rval]
+    def __eq__(self, other):
+        return type(self) == type(other) and self.tbl == other.tbl
+    def __hash__(self):
+        return hash(type(self)) # ignore hash of table
 
 
 class ScalarOp(Op):
@@ -1478,6 +1504,84 @@ class Tanh(UnaryScalarOp):
         return "%(z)s = tanh(%(x)s);" % locals()
 tanh = Tanh(upgrade_to_float, name = 'tanh')
 
+class Real(UnaryScalarOp):
+    """Extract the real coordinate of a complex number.  """
+    def impl(self, x):
+        return numpy.real(x)
+    def grad(self, (x, ), (gz, )):
+        return [complex(gz, 0)]
+
+real = Real(unary_out_lookup({ complex64:float32, complex128:float64}), name='real')
+
+class Imag(UnaryScalarOp):
+    def impl(self, x):
+        return numpy.imag(x)
+    def grad(self, (x, ), (gz, )):
+        return [complex(0, gz)]
+imag = Imag(unary_out_lookup({ complex64:float32, complex128:float64}), name='imag')
+
+class Angle(UnaryScalarOp):
+    def impl(self, x):
+        return numpy.angle(x)
+    def grad(self, (c, ), (gtheta, )):
+        # y = x.imag
+        # r = sqrt(y**2 + x.real**2)
+        # g = y/r
+        # if x == 0 and y == 0:
+        #     theta = 0
+        # elif x >= 0:
+        #     theta = numpy.arcsin(g)
+        # else:
+        #     theta = -numpy.arcsin(g)+numpy.pi
+
+        x = real(c)
+        y = imag(c)
+        r = abs(c)
+
+        gr = -gtheta * y / (r**2 * sqrt(1 - (y/r)**2))
+
+        gx = gr * x/r
+        gy = gr * y/r
+        return [complex(gx, gy)]
+angle = Angle(unary_out_lookup( { complex64:float32, complex128:float64}), name='angle')
+
+class Complex(BinaryScalarOp):
+    @staticmethod
+    def output_types_preference(x,y):
+        if x in complex_types:
+            raise TypeError(x)
+        if y in complex_types:
+            raise TypeError(y)
+
+        up = Scalar.upcast(x, y)
+        if up in ('float64', 'int64', 'uint64', 'int32', 'uint32'):
+            return [complex128]
+        else:
+            return [complex64]
+    def impl(self, x, y):
+        return numpy.complex(x, y)
+    def grad(self, (x,y), (z,)):
+        return [real(z), imag(z)]
+complex = Complex(name='complex')
+
+class ComplexFromPolar(BinaryScalarOp):
+    @staticmethod
+    def output_types_preference(x,y):
+        return Complex.output_types_preference(x,y)
+    def impl(self, r, theta):
+        if r < 0:
+            raise ValueError('polar radius must be non-negative', r)
+        x = r*numpy.cos(theta)
+        y = r*numpy.sin(theta)
+        if x.dtype == 'float32':
+            return numpy.complex64(numpy.complex(x,y))
+        else:
+            return numpy.complex128(numpy.complex(x,y))
+    def grad(self, (r,theta), (gz,)):
+        gr = cos(theta) * real(gz) + sin(theta) * imag(gz)
+        gtheta = -real(gz) * r * sin(theta) + imag(gz) * r * cos(theta)
+        return [gr, gtheta]
+complex_from_polar = ComplexFromPolar(name='complex_from_polar')
 
 
 class Composite(ScalarOp):
@@ -1642,3 +1746,4 @@ class Composite(ScalarOp):
         #otherwise self.perform won't work.
         self.__init__(self.inputs, self.outputs)
 
+