How to take Laplacian of R^n->C function

[B-Lorentz]

Hi!
I am new to jax and have a probably noobish question:
I have a jax function that maps a real array to a complex scalar (it will represent a wavefunction)

I was able to implement its laplacian like this:
laplace_psi = lambda x: jnp.diag( jacobian( jacobian(psi, holomorphic = True ), holomorphic = True) (x) ).sum(-1)

It does not work without the holomorphism promises because of the complex output, so I am forced to input complex numbers, and it produces the correct result only when the function is in fact holomorphic.

Is there a way to get the Laplacian of a non-holomorphic R^n->C function?

The crux of my predicament (it is not a bug, it is expected behavior explained in the documentation, just not the behavior I need) is illustrated in this code:

x = jnp.linspace(0, 1, 10)+0j
psi1 = lambda x: jnp.exp(1j*x)
psi2 = lambda x: jnp.cos(jnp.real(x)) + 1j*jnp.sin(jnp.real(x))
assert(jnp.allclose(psi1(x) ,psi2(x)))

d1 = vmap( grad(psi1, holomorphic=True),  0) (x)
d2 = vmap( grad(psi2, holomorphic=True),  0) (x)
assert(jnp.allclose(1j*psi1(x), d1))
print(d1-d2)

The two versions of psi produce the same result as long as x is on the real line. As C->C functions they are however completely different, and psi2 is not holomorphic. Therefore in jax's convention for gradients of complex functions they have different derivatives.

Is there a way to "look on them" as R->C functions and get the same grads?

[EyalRozenberg1]

Have you managed to solve this? Is there a way to bypass?
@B-Lorentz

[PhilipVinc]

Hey @EyalRozenberg1 , over at netket we faced several issues related to that.

In general, If you only need laplacian-vector products (Aka, Hessian-diagonal - vector products) you can write it only using vip and jvp products which are well defined for both holomorphic and non-holomorphic functions.

If you can't, you can make a function that splits the real and imaginary part of the output into two different functions and stacks the Jacobian of the two before doing further calculations and then recombines it. We do something similar around here, but that becomes more involved.

However, I have to say that I Fail to see, in this particular case, why that would be needed.
If you are trying to compute an object like a kinetic energy (aka, a laplacian) your function is necessarily R->C so you are not concerned with the non-holomorphicity.

[EyalRozenberg1]

Thank you for your reply, @PhilipVinc
As you said, I would like to take the Laplacian of a function psi(r, m) with respect to the first argument, r.
The function tasks real-valued arguments and returns a complex number (R->C).

I can do that by taking the diagonal values of the Hessian function as you said. But Jax will raise an error for the following formulation:
hessian(psi)(r, m)
TypeError: jacrev requires real-valued outputs...

To avoid this error, a modification is required:
hessian(psi, holomorphic=True)(r.astype(jnp.complex64), m.astype(jnp.complex64))

Looking at the results I have received, I am a bit worried about this implementation (the Laplacian was not negative definite). The problem presented above, by @B-Lorentz, reinforces my concerns.

just to be sure.
Does your last paragraph refer to the jvp, vjp realization? So, in your opinion, I can get the Laplacian as required in my problem. This is very gratifying!
Can you pls explain to me how Laplacian can be calculated in that manner?

Much appreciated,
Eyal

[pourion]

Hi Eyal, did you manage to compute laplacian with vjp/jvp? I am thinking about it in another application...

[EyalRozenberg1]

Hey, @pourion. I solved it in a different way that I think is more convenient and straight-forward.

The easiest way to do this is to simply split the function into its real and imaginary parts, and take the Hessian to each one individually. This can be done by using lambda function, for example.
Say you would like to take the Hessian of a function psi() with respect to the first argument, r.

r_hessian_real = hessian(lambda r_in: jax.numpy.real(psi(r, m)), argnums=0, holomorphic=False)(r)
r_hessian_imag = hessian(lambda r_in: jax.numpy.imag(psi(r, m)), argnums=0, holomorphic=False)(r)
You may discard the argument holomorphic=False, as it the default. I only used it here to emphasize.

To obtain the Laplacian, all you should to is to take the diagonal of the Hessian:
laplacian = -0.5 * ( jax.numpy.diag(r_hessian_real) + 1j * jnp.numpy.diag(r_hessian_imag) ).sum()

Hope I helped.
Eyal

[pourion]

Hey @EyalRozenberg1 Thank you so much for your response. For my application I meant to compute \nabla \cdot ( \nabla f / \vert \nabla f \vert) which is the mean curvature of the field f and position x. The following two methods provide identical results (v2 is based on your suggestion, and v1 is adapted from netket as mentioned by @PhilipVinc ):

def node_curvature_fn(f):
    def _lapl_over_f_v1(x):
        n = x.shape[0]
        eye = jnp.eye(n, dtype=f32)
        grad_f = jax.grad(f)
        grad_f_closure = lambda y: grad_f(y) / jnp.linalg.norm(grad_f(y))
        def _body_fun(i, val):
            primal_out, tangent = jax.jvp(grad_f_closure, (x,), (eye[i],))
            return val + tangent[i]
        return lax.fori_loop(0, n, _body_fun, 0.0)
    def _lapl_over_f_v2(x):
        r_hessian_real = jax.hessian(lambda r_in: f(r_in), argnums=0, holomorphic=False)(x)
        return ( jnp.diag(r_hessian_real)).sum()

    return _lapl_over_f_v1