From 61d147ac5edf21b8b1ba789693183dacd2585bd1 Mon Sep 17 00:00:00 2001
From: Alberto Mercurio <alberto.mercurio96@gmail.com>
Date: Sun, 14 Apr 2024 15:03:22 +0200
Subject: [PATCH 1/2] Improve WignerClenshaw

---
 src/wigner.jl | 46 +++++++++++++++++++++++++---------------------
 1 file changed, 25 insertions(+), 21 deletions(-)

diff --git a/src/wigner.jl b/src/wigner.jl
index 67964e21..6e3de7b4 100644
--- a/src/wigner.jl
+++ b/src/wigner.jl
@@ -48,8 +48,8 @@ function _wigner(ρ::AbstractArray, xvec::AbstractVector{T}, yvec::AbstractVecto
     return _wigner_laguerre(ρ, A, W, g, solver)
 end
 
-function _wigner(ρ::AbstractArray, xvec::AbstractVector{T}, yvec::AbstractVector{T},
-    g::Real, solver::WignerClenshaw) where {T <: BlasFloat}
+function _wigner(ρ::AbstractArray{T1}, xvec::AbstractVector{T}, yvec::AbstractVector{T},
+    g::Real, solver::WignerClenshaw) where {T1 <: BlasFloat, T <: BlasFloat}
     
     g = convert(T, g)
     M = size(ρ, 1)
@@ -62,9 +62,14 @@ function _wigner(ρ::AbstractArray, xvec::AbstractVector{T}, yvec::AbstractVecto
     W .= 2 * ρ[1, end]
     L = M - 1
 
+    y0 = similar(B, T1)
+    y1 = similar(B, T1)
+    y0_old = copy(y0)
+    res = similar(y0)
+
     while L > 0
         L -= 1
-        ρdiag = _wig_laguerre_clenshaw(L, B, (1 + Int(L!=0))*diag(ρ, L))
+        ρdiag = _wig_laguerre_clenshaw!(res, L, B, lmul!(1 + Int(L!=0), diag(ρ, L)), y0, y1, y0_old)
         @. W = ρdiag + W * A / √(L + 1)
     end
 
@@ -160,23 +165,22 @@ function check_inf(x::T) where T
 end
 
 
-function _wig_laguerre_clenshaw(L::Int, x::AbstractArray{T1}, c::AbstractVector{T2}) where {T1<:Real, T2<:BlasFloat}
-    if length(c) == 1
-        y0 = c[1]
-        y1 = 0
-    elseif length(c) == 2
-        y0 = c[1]
-        y1 = c[2]
-    else
-        k = length(c)
-        y0 = similar(x, T2); y0 .= c[end-1]
-        y1 = similar(x, T2); y1 .= c[end]
-        for i in range(3, length(c), step=1)
-            k -= 1
-            y0_old = copy(y0)
-            @. y0 = c[end+1-i] - y1 * sqrt((k - 1) * (L + k - 1) / ((L + k) * k))
-            @. y1 = y0_old - y1 * ((L + 2 * k - 1) - x) / sqrt((L + k) * k)
-        end
+function _wig_laguerre_clenshaw!(res, L::Int, x::AbstractArray{T1}, c::AbstractVector{T2}, y0, y1, y0_old) where {T1<:Real, T2<:BlasFloat}
+    length(c) == 1 && return c[1]
+    length(c) == 2 && return @. c[1] - c[2] * (L + 1 - x) / sqrt(L + 1)
+
+    y0 .= c[end-1]
+    y1 .= c[end]
+
+    k = length(c)
+    for i in range(3, length(c), step=1)
+        k -= 1
+        copyto!(y0_old, y0)
+        @. y0 = c[end+1-i] - y1 * sqrt((k - 1) * (L + k - 1) / ((L + k) * k))
+        @. y1 = y0_old - y1 * ((L + 2 * k - 1) - x) / sqrt((L + k) * k)
     end
-    return @. y0 - y1 * (L + 1 - x) / sqrt(L + 1)
+
+    @. res = y0 - y1 * (L + 1 - x) / sqrt(L + 1)
+
+    return res
 end
\ No newline at end of file

From 71a58cd5874d74aeba3b56b61de1e73b8d76f882 Mon Sep 17 00:00:00 2001
From: Alberto Mercurio <alberto.mercurio96@gmail.com>
Date: Sun, 14 Apr 2024 16:30:37 +0200
Subject: [PATCH 2/2] Minor changes

---
 src/wigner.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/wigner.jl b/src/wigner.jl
index 6e3de7b4..06d35553 100644
--- a/src/wigner.jl
+++ b/src/wigner.jl
@@ -165,7 +165,7 @@ function check_inf(x::T) where T
 end
 
 
-function _wig_laguerre_clenshaw!(res, L::Int, x::AbstractArray{T1}, c::AbstractVector{T2}, y0, y1, y0_old) where {T1<:Real, T2<:BlasFloat}
+function _wig_laguerre_clenshaw!(res, L::Int, x, c, y0, y1, y0_old)
     length(c) == 1 && return c[1]
     length(c) == 2 && return @. c[1] - c[2] * (L + 1 - x) / sqrt(L + 1)