Complex Numbers in Z3

A user recently asked me if Z3 had support for Complex numbers. Unfortunately, it does not. This is on my TODO list, but it is has very low priority at this moment. Internally, Z3 already has a lot of machinery for supporting Complex numbers, but there is a lot of "plumbing" missing, such as extending the parsers, simplifiers, models, and root isolation procedures. Most of this work is not really fun.

In the meantime, we can encode Complex numbers on top of the Real numbers provided by Z3. The basic idea is to represent a Complex number as a pair of Real numbers. This is not the most efficient way of implementing a decision procedure for the theory of Complex numbers, but it is simple and allows us to handle problems that mix Real and Complex number reasoning. Moreover, we can decide the resulting problem using the new nonlinear arithmetic solver nlsat that is available in Z3.

We can represent a complex number

\[ a + b i \]

as a pair of real numbers

\[ (a, b) \]

Then the operations on Complex numbers can be easily encode on top of the operations on Real numbers provided by Z3. Here is the encoding of the main operations.

\begin{eqnarray*} (a_1, b_1) + (a_2, b_2) & = & (a_1 + a_2,\, b_1 + b_2) \\ (a_1, b_1) - (a_2, b_2) & = & (a_1 - a_2,\, b_1 - b_2) \\ (a_1, b_1) \times (a_2, b_2) & = & (a_1 a_2 - b_1 b_2,\, a_1 b_2 + b_1 a_2) \\ (a, b)^{-1} & = & (\frac{a}{a^2 + b^2}, -\frac{b}{a^2 + b^2}) \mbox { if } a \neq 0 \vee b \neq 0 \\ (a_1, b_1) \div (a_2, b_2) & = & (a_1, b_1) \times (a_2, b_2)^{-1} \mbox { if } a \neq 0 \vee b \neq 0 \\ \end{eqnarray*}

The definition of the multiplicative inverse follows from basic algebraic manipulation

\[ \frac{1}{a + b i} = \frac{a - b i}{(a + b i)(a - b i)} = \frac{a - b i}{a^2 + b^2} = \frac{a}{a^2 + b^2} - \frac{b}{a^2 + b^2} i \]

It is straightforward to implement the operations above using the Z3 Python API. We can even create a class ComplexExpr to wrap the pair of Z3 real expressions and overload the usual operators +, -, *, /. This class has two fields .r (the real part) and .i (the imaginary part). The auxiliary function _to_complex is used to convert a into a Z3 "Complex" number if it is not already one.

def _to_complex(a):
    if isinstance(a, ComplexExpr):
        return a
    else:
        return ComplexExpr(a, RealVal(0))

def _is_zero(a):
    return (isinstance(a, int) and a == 0) or (is_rational_value(a) and a.numerator_as_long() == 0)

class ComplexExpr:
    def __init__(self, r, i):
        self.r = r
        self.i = i

    def __add__(self, other):
        other = _to_complex(other)
        return ComplexExpr(self.r + other.r, self.i + other.i)

    def __radd__(self, other):
        other = _to_complex(other)
        return ComplexExpr(other.r + self.r, other.i + self.i)

    def __sub__(self, other):
        other = _to_complex(other)
        return ComplexExpr(self.r - other.r, self.i - other.i)

    def __rsub__(self, other):
        other = _to_complex(other)
        return ComplexExpr(other.r - self.r, other.i - self.i)

    def __mul__(self, other):
        other = _to_complex(other)
        return ComplexExpr(self.r*other.r - self.i*other.i, self.r*other.i + self.i*other.r)

    def __mul__(self, other):
        other = _to_complex(other)
        return ComplexExpr(other.r*self.r - other.i*self.i, other.i*self.r + other.r*self.i)

    def inv(self):
        den = self.r*self.r + self.i*self.i
        return ComplexExpr(self.r/den, -self.i/den)

    def __div__(self, other):
        inv_other = _to_complex(other).inv()
        return self.__mul__(inv_other)

    def __rdiv__(self, other):
        other = _to_complex(other)
        return self.inv().__mul__(other)

    def __eq__(self, other):
        other = _to_complex(other)
        return And(self.r == other.r, self.i == other.i)

    def __neq__(self, other):
        return Not(self.__eq__(other))

    def simplify(self):
        return ComplexExpr(simplify(self.r), simplify(self.i))

    def repr_i(self):
        if is_rational_value(self.i):
            return "%s*I" % self.i
        else:
            return "(%s)*I" % str(self.i)

    def __repr__(self):
        if _is_zero(self.i):
            return str(self.r)
        elif _is_zero(self.r):
            return self.repr_i()
        else:
            return "%s + %s" % (self.r, self.repr_i())

The function Complex(a) creates a new "Complex" variable. Actually, it creates two real variables named a.r and a.i. The constant I is just an alias for the pair (0, 1).

def Complex(a):
    return ComplexExpr(Real('%s.r' % a), Real('%s.i' % a))
I = ComplexExpr(RealVal(0), RealVal(1))

The function evaluate_cexpr "retrieves" the value of a Complex expression e in the model m.

def evaluate_cexpr(m, e):
    return ComplexExpr(m[e.r], m[e.i])

Now, we have all the pieces we need to solve our first problem. Find a root of x*x + 2.

x = Complex("x")
s = Tactic('qfnra-nlsat').solver()
s.add(x*x + 2 == 0)
print(s.check())
m = s.model()
print('x = %s' % evaluate_cexpr(m, x))

In the example above, I used Tactic('qfnra-nlsat').solver() to create a solver based on nlsat. Note that, the problem above is unsatisfiable on the Reals. We can try this example online at rise4fun.