Wirtinger derivatives

[lucasgrjn]

Hi everyone !

I am a beginner with Jax. I understand the principles but I am trying to use it with complex calculus.
In my case, I am trying to optimize nanophotonics structures.

On a part of my code, I want to differentiate this kind of equation $\frac{\partial}{\partial z}|z|^2$ which is equal to $\overline{z}$ (with respect to Wirtinger derivatives).

And this is where problems start.. I do not understand how I can find such a result using jvp / vjp. From Autodiff Cookbook, I understand that we get a Jacobian but I am having a hard time to connect it to Wirtinger derivatives. (I also took a look to #1853 or #3402).

Can anyone give me some insight ?

Thanks in advance,
Differentially yours,
Lucas.

[YouJiacheng]

import jax
def wirtinger_derivatives(fun):
    def fun_xy(x, y):
        return fun(x + y *1j)
    def derivative(z):
        x, y = z.real, z.imag
        dx, dy = jax.grad(fun_xy, argnums=(0, 1))(x, y)
        return (dx - dy * 1j) / 2
    return derivative

[lucasgrjn]

Dear @YouJiacheng ,

Thank for your answer ! I try your code, it works for $\frac{\partial}{\partial z} |z|^2$ but if I replace with $\frac{\partial}{\partial z} z\overline{z}$, it stops working since jax.grad "requires real-valued outputs"...

I join my code below

import jax


def wirtinger_derivatives(fun):
    def fun_xy(x, y):
        return fun(x + y *1j)
    def derivative(z):
        x, y = z.real, z.imag
        dx, dy = jax.grad(fun_xy, argnums=(0, 1))(x, y)
        return (dx - dy * 1j) / 2
    return derivative


f = lambda z: jax.numpy.abs(z) ** 2
g = lambda z: z * jax.numpy.conj(z)

val_test = 3. + 1.j
assert f(val_test) == g(val_test)
print(g(val_test))

dfdz = wirtinger_derivatives(f)
assert dfdz(val_test) == jax.numpy.conj(val_test)

dgdz = wirtinger_derivatives(g)
assert dgdz(val_test) == jax.numpy.conj(val_test)

[YouJiacheng]

I think wirtinger derivatives is defined for real-valued function only. Here z * z.conj has complex dtype(though its imaginary part is 0).
It can be fixed by modifying fun_xy:

def fun_xy(x, y):
    return fun(x + y *1j).real

But I suggest modifying g:

g = lambda z: (z * jax.numpy.conj(z)).real

[lucasgrjn]

Actually Wirtinger are also defined for complex-function !

I managed to find a way to obtain the derivatives of $u(x,y)$ and $v(x,y)$ defined from Autodiff Cookbook

I use the following code :

import jax


def wirtinger_derivatives(fun):
    def derivative(z):
        z += 0.j  # To ensure z is complex

        du = jax.vjp(fun,z)[1](1. + 0.j)
        dv = jax.vjp(fun,z)[1](0. - 1.j)

        dudx = du[0].real
        dudy = -du[0].imag
        dvdx = dv[0].real
        dvdy = -dv[0].imag

        dx = dudx + 1.j * dvdx
        dy = dudy + 1.j * dvdy

        return (dx - dy * 1j) / 2
    return derivative


f = lambda z: jax.numpy.abs(z) ** 2 + 0.j
g = lambda z: z * jax.numpy.conj(z)
h = lambda z: z

val_test = 3. + 1.j
assert f(val_test) == g(val_test)
print(g(val_test))

dfdz = wirtinger_derivatives(f)
assert dfdz(val_test) == jax.numpy.conj(val_test)

dgdz = wirtinger_derivatives(g)
assert dgdz(val_test) == jax.numpy.conj(val_test)

assert dfdz(val_test) == dgdz(val_test)

dhdz = wirtinger_derivatives(h)
assert dhdz(val_test) == 1.

One should notice that if we are interested for $\frac{\partial}{\partial \overline{z}}$, only the last line need to be changed to

        return (dx + dy * 1j) / 2

since we have $\frac{\partial}{\partial z} = \frac{1}{2} \left(\frac{\partial}{\partial x} - j \frac{\partial}{\partial y} \right)$ and $\frac{\partial}{\partial \overline{z}} = \frac{1}{2} \left(\frac{\partial}{\partial x} + j \frac{\partial}{\partial y} \right)$

[YouJiacheng]

Oh I didn't know this before!
FYI, I think your code can be optimized(for faster tracing and compiling):

import jax
def wirtinger_derivatives(fun):
    def derivative(z):
        z += 0.j  # To ensure z is complex
        _, vjp_f = jax.vjp(fun, z) # only linearize once!
        du, = vjp_f(1. + 0.j)
        dv, = vjp_f(0. - 1.j)
        dudx = du.real
        dudy = -du.imag
        dvdx = dv.real
        dvdy = -dv.imag
        dx = dudx + 1.j * dvdx
        dy = dudy + 1.j * dvdy
        return (dx - dy * 1j) / 2
    return derivative

[lucasgrjn]

Thanks for the insights !!

[lucasgrjn]

@mattjj without being presumptuous, do you want I try to write a little paragraph on this and the Wirtinger derivatives to be added on the page Autodiff Cookbook ?