math icon indicating copy to clipboard operation
math copied to clipboard

Add vectorized variant for functions of complex variable, especially exp

Open jachymb opened this issue 10 months ago • 4 comments

Currently, according to the guide https://mc-stan.org/docs/functions-reference/complex-valued_basic_functions.html it seems that basic arithmetic operators and pow are the only functions that accept complex vectors.

I am modelling some periodicities in data using Fourier series, and I find it the most compact (and computationally efficient) to express what I do in terms of complex numbers. However, it's not efficient to run a function over a matrix using a loop when there's autodiff.

The most striking case is the exp function, which is ubiquitous in complex analysis and I want to be able to run it elementwise over a vector or a matrix.

I figured out a possible workaround is to use pow(exp(1), z) which does work on vectorized z, but still I believe that exp is more fundamental than pow and it should be implemented, because the current state is confusing at best.

Of course then there are all the other functions which accept a real or complex scalar and accept vectorized real, but don't accept vectorized complex, which is inconsistent. This includes log and the trigonometric functions.

jachymb avatar Jan 09 '25 08:01 jachymb

I agree we should support these!

But, I also wanted to clear up a possible misconception:

However, it's not efficient to run a function over a matrix using a loop when there's autodiff.

In many cases, this is exactly what the elementwise versions of the functions for reals are doing, e.g.

https://github.com/stan-dev/math/blob/42d94c4840f681806ae0e0134120a4077a29e46c/stan/math/prim/fun/exp.hpp#L70-L73

(where apply_scalar_unary is basically some templated magic for, essentially, a for loop)

This is different from e.g. the distribution functions which are vectorized, which are written to share intermediate results and autodiff memory between different elements. There are also a few exceptions to this (usually around the "struct of arrays" optimization), but in general you shouldn't fear writing your own for-loop-based helper functions for these until they are supported directly in the language.

WardBrian avatar Jan 09 '25 14:01 WardBrian

On the C++ side, I don't believe there is anything specific about complex types when it comes to apply_scalar_unary and apply_scalar_binary, so if there is an existing scalar definition for complex types then it should be automatically vectorised.

Are there complex functions that are missing from the Math library, or is this just an issue for stanc3 to expose the signatures?

andrjohns avatar Jan 10 '25 04:01 andrjohns

Are there complex functions that are missing from the Math library

Generally yes, though sometimes all that's needed is loosening an existing template bound

WardBrian avatar Jan 10 '25 14:01 WardBrian

Another complex-variable function that apparently isn't vectorized but definitely should be is arg. This is in strange contrast to its sister abs which works just fine.

jachymb avatar Jan 17 '25 22:01 jachymb