# Optimize function with complex square root

**URL:** <https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058>\
**Category:** Performance\
**Created:** [February 26, 2021, 12:41am UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058 "2021-02-26T00:41:49Z")\
**Posts on this page:** 13\
**Page:** 1

<div class="post-metadata">

**Author:** ![mleprovost](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mleprovost/32/7166_2.png) [@mleprovost](https://discourse.julialang.org/u/mleprovost)\
**Post date:** [February 26, 2021, 12:41am UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/1 "2021-02-26T00:41:49Z")

</div>

Hello,

I am writing a routine to store a matrix that contains the values of f(x, z) = \frac{1}{\sqrt{1-x^2}} \times \text{real}\left(\frac{\sqrt{z-1}\sqrt{z+1}}{z-x}\right) for \sim 1000 real values x \in [-1, 1] and \sim 50 complex numbers z.

This computation is a critical part of my code, do you have ideas of how I can improve it?

We cannot simplify \sqrt{z-1}\sqrt{z+1} since z is complex. Otherwise, the branch cut of the square root is no longer on the real axis between -1 and 1.

```julia
using LinearAlgebra
using BenchmarkTools

function test(x::Array{Float64,1}, z::Array{ComplexF64,1})
   A = zeros(size(x,1), size(z, 1))   
   @inbounds for (j, zj) in enumerate(z) 
        bj = √(zj-1.0)*√(zj+1.0)
        for (i, xi) in enumerate(x)
            ci = (1.0 - xi^2)^(-0.5)
            A[i,j] = ci*real(bj/(zj-xi))
        end
    end
    return A
end

x = 2*rand(1000) .- 1
z = randn(50) + im*randn(50)

@btime test($x, $z)
989.748 μs (2 allocations: 390.70 KiB)

```

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [February 26, 2021, 1:21am UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/2 "2021-02-26T01:21:57Z")

</div>

Although I can see several possible ways to further optimize this code, some context would be helpful.

What is this for? Is it performance-critical because the matrix is huge or because you are computing it over and over for different x and/or z?

---

<div class="post-metadata">

**Author:** ![mleprovost](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mleprovost/32/7166_2.png) [@mleprovost](https://discourse.julialang.org/u/mleprovost)\
**Post date:** [February 26, 2021, 1:48am UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/3 "2021-02-26T01:48:21Z")

</div>

The computation arises in a potential flow problem in fluid mechanics. I am interested in the induced effect of a collection of singularities at the locations {z\_j} on an infinitely-thin plate materizalized by the interval [-1,1]. Here f(x\_i,z\_j) computes the induce effect of a point vortex at location z\_j at a point x\_i along the plate. This computation is used in the forecast step of an ensemble Kalman filter. Therefore, I need to compute this matrix for about 300 (number of ensemble members) different sets of {z\_j} values. The function `test` performs this computation for a distributed set of {x\_i} values for only one set of {z\_j} values. The set of {x\_i} values is always the same. I hope this helps

---

<div class="post-metadata">

**Author:** ![PeterSimon](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/petersimon/32/25193_2.png) [@PeterSimon](https://discourse.julialang.org/u/PeterSimon)\
**Post date:** [February 26, 2021, 4:06am UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/4 "2021-02-26T04:06:55Z")

</div>

[LoopVectorization](https://github.com/chriselrod/LoopVectorization.jl) produces a huge speedup on my machine. Here is the contents of the file `testit.jl`:

```nohighlight
using LinearAlgebra
using BenchmarkTools
using LoopVectorization

function test(x::Array{Float64,1}, z::Array{ComplexF64,1})
   A = zeros(size(x,1), size(z, 1))   
   @inbounds for (j, zj) in enumerate(z) 
        bj = √(zj-1.0)*√(zj+1.0)
        for (i, xi) in enumerate(x)
            ci = (1.0 - xi^2)^(-0.5)
            A[i,j] = ci*real(bj/(zj-xi))
        end
    end
    return A
end

function test2(x, z)
    a = zeros(length(x),length(z))
    b = [√(t-1)*√(t+1) for t in z]
    c = [1/sqrt(1 - t^2) for t in x]
    @avx for i in axes(a,1), j in axes(a,2)
        a[i,j] = c[i] * real(b[j]/(z[j]-x[i])) 
    end
    a
end

x = 2*rand(1000) .- 1
z = randn(50) + im*randn(50)

versioninfo()
a = @btime test($x,$z)
a2 = @btime test2($x,$z)
@show norm(a-a2)

```

This produces the following output:

```julia
julia> include("testit.jl")
Julia Version 1.5.3
Commit 788b2c77c1 (2020-11-09 13:37 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: Intel(R) Core(TM) i7-9700 CPU @ 3.00GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-9.0.1 (ORCJIT, skylake)
Environment:
  JULIA_EDITOR = runemacs.exe
  JULIA_NUM_THREADS = 8

  947.700 μs (2 allocations: 390.70 KiB)
  66.500 μs (4 allocations: 399.52 KiB)
norm(a - a2) = 9.676938524225372e-14

```

---

<div class="post-metadata">

**Author:** ![moeddel](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/moeddel/32/18641_2.png) [@moeddel](https://discourse.julialang.org/u/moeddel)\
**Post date:** [February 26, 2021, 9:01am UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/5 "2021-02-26T09:01:25Z")

</div>

> [@PeterSimon](#):
>
> [LoopVectorization](https://github.com/chriselrod/LoopVectorization.jl) produces a huge speedup on my machine. Here is the contents of the file `testit.jl` :

I can confirm this. You even get a little better if you do calculations in single precicion.

```julia
using BenchmarkTools
using LoopVectorization

function test(x::Vector{T}, z::Vector{Complex{T}}) where T<:Real
    A = zeros(T,length(x), length(z))
    b = map(z->√(z - one(T))*√(z + one(T)),z)
	c = map(x->(one(T) - x^2)^(-0.5),x)
    @avx for j in eachindex(z)
        for i in eachindex(x)
            A[i,j] = c[i]*real(b[j]/(z[j]-x[i]))
        end
     end
     return A
end

T = Float32
x = 2*rand(T,1000) .- 1
z = randn(Complex{T},50)

```

It boosts from 70 μs to 62 μs on my machine, which is not the fastest. That is if you do not need the numerical accuracy of double precision.

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [February 26, 2021, 9:19am UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/6 "2021-02-26T09:19:28Z")

</div>

Tullio.jl can help you further, if you have a few cores:

```julia
function test2(x, z)
    b = sqrt.(z .- 1) .* sqrt.(z .+ 1)
    c = inv.(sqrt.(1 .- x.^2))
    @tullio A[i, j] := c[i] * real(b[j]/(z[j]-x[i]))
    return A
end

```

---

<div class="post-metadata">

**Author:** ![rafael.guerra](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rafael.guerra/32/216610_2.png) [@rafael.guerra](https://discourse.julialang.org/u/rafael.guerra)\
**Post date:** [February 26, 2021, 12:40pm UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/7 "2021-02-26T12:40:03Z")

</div>

For whatever it is worth, it is interesting that this function can be written without loops with adjoints. The `@fastmath` macro speeds up significantely but not as much as Tullio does:

```julia
function test4(x, z)
    b = @. √(z-1)*√(z+1)
    c = @. 1.0/sqrt(1 - x^2)
    @fastmath a = @. c*real(b'/(z' - x)) # using adjoints
    return a
end

```

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [February 26, 2021, 1:37pm UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/8 "2021-02-26T13:37:16Z")

</div>

> [@mleprovost](#):
>
> Therefore, I need to compute this matrix for about  
> 300 (number of ensemble members) different sets of \{ z\_j \} values.The set of \{ x\_i \} values is always the same. I hope this helps

In that case, you could pre-compute the 1/\sqrt{1-x\_i^2} values, and pre-allocate the A matrix.

As others have pointed out, `@fastmath` helps a lot — the reason is that complex-number division normally requires a bunch of costly checks to eliminate spurious overflows and similar conditions that can arise for extreme inputs (e.g. `x[j] = 1e300`), which you can safely eliminate if you know that you always have moderate-sized inputs and never divide by zero.

One can also use threads, but probably you are better off using threads at a higher level (e.g. to process different ensemble instances in parallel), so I’m only looking at the single-threaded performance here:

```julia
using BenchmarkTools

function test!(A, x, c, z)
    @fastmath @inbounds for j in eachindex(z)
        zj = z[j]
        bj = √(zj-1)*√(zj+1)
        for i in eachindex(x)
            A[i,j] = c[i] * real(bj / (zj - x[i]))
        end
    end
    return A
end

x = 2*rand(1000) .- 1
z = randn(50) + im*randn(50)
A = zeros(size(x,1), size(z, 1))
c = @. inv(sqrt(1 - x^2))

```

which gives:

```julia
julia> @btime test($x, $z); # your original code
  1.335 ms (2 allocations: 390.70 KiB)

julia> @btime test!($A, $x, $c, $z); # new code
  34.047 μs (0 allocations: 0 bytes)

```

for about 40× speedup.

(I kept the explicit loops in order to continue to compute `bj = -√(zj-1)*√(zj+1)` outside the inner loop without storing it in an additional array.)

I tried `@avx` from LoopVectorization.jl and it unfortunately gives me an error… it works if I write out `real(bj / (zj - x[i]))` explicitly in terms of real and imaginary parts, but in that case I get negligible speedup from `@avx`.

---

<div class="post-metadata">

**Author:** ![mleprovost](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mleprovost/32/7166_2.png) [@mleprovost](https://discourse.julialang.org/u/mleprovost)\
**Post date:** [February 26, 2021, 10:34pm UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/9 "2021-02-26T22:34:55Z")

</div>

Thank you very much for your help. This is definitly much more efficient.  
Is `eachindex` preferable to `enumerate` for the performance?  
Is there something that prevent me from using @fastmath in my future scripts everty time I am not dealing with very large entries?

---

<div class="post-metadata">

**Author:** ![moeddel](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/moeddel/32/18641_2.png) [@moeddel](https://discourse.julialang.org/u/moeddel)\
**Post date:** [February 27, 2021, 8:47am UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/10 "2021-02-27T08:47:46Z")

</div>

> [@mleprovost](#):
>
> Is there something that prevent me from using @fastmath in my future scripts everty time I am not dealing with very large entries?

If you want to be sure that `@fastmath` does what you want, please check out [this page](https://llvm.org/docs/LangRef.html#fast-math-flags), where it is documented, which kind of optimizations `@fastmath` makes and what kind of assumptions.

---

<div class="post-metadata">

**Author:** ![Henrique\_Becker](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/henrique_becker/32/15443_2.png) [@Henrique\_Becker](https://discourse.julialang.org/u/Henrique_Becker)\
**Post date:** [February 27, 2021, 1:57pm UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/11 "2021-02-27T13:57:12Z")

</div>

`eachindex` is more generic, `enumerate` does not give indexes, it gives `1` with the first element, `2` with the second, and so on; if your Array has more than one dimension, or custom indexes (start at 0 for example) then your code will break. `eachindex` will always give valid indexes, and for N-dimensional arrays, it will give the indexes in the most efficient order for linear access. If you want something that both give you valid indexes and the value then look at `pairs` documentation and use it instead of `enumerate`.

I talk a little more about this [in my answer of mine](https://discourse.julialang.org/t/any-shortcut-to-1-length-myvector/52816/27).

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [February 27, 2021, 3:39pm UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/12 "2021-02-27T15:39:32Z")

</div>

> [@Henrique\_Becker](#):
>
> `eachindex` is more generic

I wouldn’t say that. `enumerate` is generic, works also for non-standard indexed arrays, and (without testing) possibly for non-indexed collections (I should check this).

It just doesn’t do what many people think it does. It doesn’t return indices. It’s unfortunate that `enumerate` is so widely known, and `pairs` isn’t, because I suspect it’s very often misused.

I wonder if a linter would be able to catch this. If output from `enumerate` is used for indexing, it might suggest `pairs` instead.

---

<div class="post-metadata">

**Author:** ![Henrique\_Becker](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/henrique_becker/32/15443_2.png) [@Henrique\_Becker](https://discourse.julialang.org/u/Henrique_Becker)\
**Post date:** [February 27, 2021, 5:24pm UTC](https://discourse.julialang.org/t/optimize-function-with-complex-square-root/56058/13 "2021-02-27T17:24:32Z")

</div>

> [@DNF](#):
>
> I wouldn’t say that. `enumerate` is generic, works also for non-standard indexed arrays, and (without testing) possibly for non-indexed collections (I should check this).
> 
> It just doesn’t do what many people think it does. It doesn’t return indices

Fair enough. I was saying that `eachindex` is more generic for the purpose of having both index and element, but maybe it would be best to say that `eachindex` is the correct choice for that purpose and `enumerate` is designed for another purpose (easily adding an iteration counter on a loop).