Specialize sqrt and cbrt (#379)

* Sqrt for real matrices

* Add cbrt

* Tests for negative and complex values

* Tests for zeros

* Compare with dense

* Test cbrt only on recent julia versions

Author: jishnub

src/FillArrays.jl

@@ -374,6 +374,32 @@ fillsimilar(a::Ones{T}, axes...) where T = Ones{T}(axes...)
 fillsimilar(a::Zeros{T}, axes...) where T = Zeros{T}(axes...)
 fillsimilar(a::AbstractFill, axes...) = Fill(getindex_value(a), axes...)
 
+# functions
+function Base.sqrt(a::AbstractFillMatrix{<:Union{Real, Complex}})
+    Base.require_one_based_indexing(a)
+    size(a,1) == size(a,2) || throw(DimensionMismatch("matrix is not square: dimensions are $(size(a))"))
+    _sqrt(a)
+end
+_sqrt(a::AbstractZerosMatrix) = float(a)
+function _sqrt(a::AbstractFillMatrix)
+    n = size(a,1)
+    n == 0 && return float(a)
+    v = getindex_value(a)
+    Fill(√(v/n), axes(a))
+end
+function Base.cbrt(a::AbstractFillMatrix{<:Real})
+    Base.require_one_based_indexing(a)
+    size(a,1) == size(a,2) || throw(DimensionMismatch("matrix is not square: dimensions are $(size(a))"))
+    _cbrt(a)
+end
+_cbrt(a::AbstractZerosMatrix) = float(a)
+function _cbrt(a::AbstractFillMatrix)
+    n = size(a,1)
+    n == 0 && return float(a)
+    v = getindex_value(a)
+    Fill(cbrt(v)/cbrt(n)^2, axes(a))
+end
+
 struct RectDiagonal{T,V<:AbstractVector{T},Axes<:Tuple{Vararg{AbstractUnitRange,2}}} <: AbstractMatrix{T}
     diag::V
     axes::Axes

test/runtests.jl

@@ -2959,3 +2959,32 @@ end
     @test triu(Z, 2) === Z
     @test tril(Z, 2) === Z
 end
+
+@testset "sqrt/cbrt" begin
+    F = Fill(4, 4, 4)
+    A = Array(F)
+    @test sqrt(F) ≈ sqrt(A) rtol=3e-8
+    @test sqrt(F)^2 ≈ F
+    F = Fill(4+4im, 4, 4)
+    A = Array(F)
+    @test sqrt(F) ≈ sqrt(A) rtol=1e-8
+    @test sqrt(F)^2 ≈ F
+    F = Fill(-4, 4, 4)
+    A = Array(F)
+    if VERSION >= v"1.11.0-rc3"
+        @test cbrt(F) ≈ cbrt(A) rtol=1e-5
+    end
+    @test cbrt(F)^3 ≈ F
+
+    # avoid overflow
+    F = Fill(4, typemax(Int), typemax(Int))
+    @test sqrt(F)^2 ≈ F
+    @test cbrt(F)^3 ≈ F
+
+    # zeros
+    F = Zeros(4, 4)
+    A = Array(F)
+    @test sqrt(F) ≈ sqrt(A) atol=1e-14
+    @test sqrt(F)^2 == F
+    @test cbrt(F)^3 == F
+end