Improve static analysis: type stability and JET tests

src/DataReduction/POD.jl

@@ -1,7 +1,7 @@
 using TSVD: tsvd
 using RandomizedLinAlg: rsvd
 
-function matricize(VoV::Vector{Vector{T}}) where {T}
+function matricize(VoV::Vector{Vector{T}})::Matrix{T} where {T}
     return reduce(hcat, VoV)
 end
 
@@ -26,35 +26,51 @@ end
 
 _rsvd(data, n::Int, p::Int) = rsvd(data, n, p)
 
-mutable struct POD <: AbstractDRProblem
+mutable struct POD{S, T <: AbstractFloat} <: AbstractDRProblem
     # specified
-    snapshots::Any
-    min_renergy::Any
+    snapshots::S
+    min_renergy::T
     min_nmodes::Int
     max_nmodes::Int
     # computed
     nmodes::Int
-    rbasis::Any
-    renergy::Any
-    spectrum::Any
+    rbasis::Union{Missing, Matrix{T}}
+    renergy::T
+    spectrum::Union{Missing, Vector{T}}
     # constructors
     function POD(
-            snaps;
-            min_renergy = 1.0,
+            snaps::S;
+            min_renergy::T = 1.0,
             min_nmodes::Int = 1,
             max_nmodes::Int = length(snaps[1])
-        )
+        ) where {S <: AbstractMatrix{T}} where {T <: AbstractFloat}
         nmodes = min_nmodes
         errorhandle(snaps, nmodes, min_renergy, min_nmodes, max_nmodes)
-        return new(snaps, min_renergy, min_nmodes, max_nmodes, nmodes, missing, 1.0, missing)
+        return new{S, T}(snaps, min_renergy, min_nmodes, max_nmodes, nmodes, missing, one(T), missing)
     end
-    function POD(snaps, nmodes::Int)
-        errorhandle(snaps, nmodes, 0.0, nmodes, nmodes)
-        return new(snaps, 0.0, nmodes, nmodes, nmodes, missing, 1.0, missing)
+    function POD(
+            snaps::S;
+            min_renergy::T = 1.0,
+            min_nmodes::Int = 1,
+            max_nmodes::Int = length(snaps[1])
+        ) where {T <: AbstractFloat, S <: AbstractVector{<:AbstractVector{T}}}
+        nmodes = min_nmodes
+        errorhandle(snaps, nmodes, min_renergy, min_nmodes, max_nmodes)
+        return new{S, T}(snaps, min_renergy, min_nmodes, max_nmodes, nmodes, missing, one(T), missing)
+    end
+    function POD(snaps::S, nmodes::Int) where {S <: AbstractMatrix{T}} where {T <: AbstractFloat}
+        errorhandle(snaps, nmodes, zero(T), nmodes, nmodes)
+        return new{S, T}(snaps, zero(T), nmodes, nmodes, nmodes, missing, one(T), missing)
+    end
+    function POD(snaps::S, nmodes::Int) where {T <: AbstractFloat, S <: AbstractVector{<:AbstractVector{T}}}
+        errorhandle(snaps, nmodes, zero(T), nmodes, nmodes)
+        return new{S, T}(snaps, zero(T), nmodes, nmodes, nmodes, missing, one(T), missing)
     end
 end
 
-function determine_truncation(s, min_nmodes, min_renergy, max_nmodes)
+function determine_truncation(
+        s::AbstractVector{T}, min_nmodes::Int, max_nmodes::Int, min_renergy::T
+    )::Tuple{Int, T} where {T <: AbstractFloat}
     nmodes = min_nmodes
     overall_energy = sum(s)
     energy = sum(s[1:nmodes]) / overall_energy
@@ -65,42 +81,43 @@ function determine_truncation(s, min_nmodes, min_renergy, max_nmodes)
     return nmodes, energy
 end
 
-function reduce!(pod::POD, alg::SVD)
+function reduce!(pod::POD{S, T}, alg::SVD)::Nothing where {S, T}
     u, s, v = _svd(pod.snapshots; alg.kwargs...)
     pod.nmodes,
         pod.renergy = determine_truncation(
         s, pod.min_nmodes, pod.max_nmodes,
         pod.min_renergy
     )
-    pod.rbasis = u[:, 1:(pod.nmodes)]
-    pod.spectrum = s
+    pod.rbasis = Matrix{T}(u[:, 1:(pod.nmodes)])
+    pod.spectrum = Vector{T}(s)
     return nothing
 end
 
-function reduce!(pod::POD, alg::TSVD)
+function reduce!(pod::POD{S, T}, alg::TSVD)::Nothing where {S, T}
     u, s, v = _tsvd(pod.snapshots, pod.nmodes; alg.kwargs...)
     n_max = min(size(u, 1), size(v, 1))
-    pod.renergy = sum(s) / (sum(s) + (n_max - pod.nmodes) * s[end])
-    pod.rbasis = u
-    pod.spectrum = s
+    pod.renergy = T(sum(s) / (sum(s) + (n_max - pod.nmodes) * s[end]))
+    pod.rbasis = Matrix{T}(u)
+    pod.spectrum = Vector{T}(s)
     return nothing
 end
 
-function reduce!(pod::POD, alg::RSVD)
+function reduce!(pod::POD{S, T}, alg::RSVD)::Nothing where {S, T}
     u, s, v = _rsvd(pod.snapshots, pod.nmodes, alg.p)
     n_max = min(size(u, 1), size(v, 1))
-    pod.renergy = sum(s) / (sum(s) + (n_max - pod.nmodes) * s[end])
-    pod.rbasis = u
-    pod.spectrum = s
+    pod.renergy = T(sum(s) / (sum(s) + (n_max - pod.nmodes) * s[end]))
+    pod.rbasis = Matrix{T}(u)
+    pod.spectrum = Vector{T}(s)
     return nothing
 end
 
-function Base.show(io::IO, pod::POD)
+function Base.show(io::IO, pod::POD)::Nothing
     print(io, "POD \n")
     print(io, "Reduction Order = ", pod.nmodes, "\n")
     print(
         io, "Snapshot size = (", size(pod.snapshots, 1), ",", size(pod.snapshots[1], 2),
         ")\n"
     )
-    return print(io, "Relative Energy = ", pod.renergy, "\n")
+    print(io, "Relative Energy = ", pod.renergy, "\n")
+    return nothing
 end

src/Types.jl

@@ -4,17 +4,19 @@ abstract type AbstractDRProblem <: AbstractReductionProblem end
 
 abstract type AbstractSVD end
 
-struct SVD <: AbstractSVD
-    kwargs::Any
+struct SVD{K <: NamedTuple} <: AbstractSVD
+    kwargs::K
     function SVD(; kwargs...)
-        return new(kwargs)
+        kw = NamedTuple(kwargs)
+        return new{typeof(kw)}(kw)
     end
 end
 
-struct TSVD <: AbstractSVD
-    kwargs::Any
+struct TSVD{K <: NamedTuple} <: AbstractSVD
+    kwargs::K
     function TSVD(; kwargs...)
-        return new(kwargs)
+        kw = NamedTuple(kwargs)
+        return new{typeof(kw)}(kw)
     end
 end
 

src/utils.jl

@@ -31,7 +31,7 @@ $(SIGNATURES)
 Returns `true` if `expr` contains variables in `dvs` only and does not contain `iv`.
 
 """
-function only_dvs(expr, dvs, iv)
+function only_dvs(expr, dvs, iv)::Bool
     if isequal(expr, iv)
         return false
     elseif expr in dvs

test/Project.toml

@@ -1,6 +1,8 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MethodOfLines = "94925ecb-adb7-4558-8ed8-f975c56a0bf4"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
@@ -10,6 +12,8 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 [compat]
 Aqua = "0.8"
 ExplicitImports = "1"
+JET = "0.10"
+LinearAlgebra = "1"
 MethodOfLines = "0.11"
 ModelingToolkit = "10.10"
 OrdinaryDiffEq = "6"

test/jet.jl

@@ -0,0 +1,41 @@
+using Test, JET
+using ModelOrderReduction
+using LinearAlgebra: qr
+
+@testset "JET Static Analysis" begin
+    # Create test data
+    n = 20  # state dimension
+    m = 10  # number of snapshots
+    snapshot_matrix = Float64[sin(i * j / n) for i in 1:n, j in 1:m]
+    snapshot_vov = [Float64[sin(i * j / n) for i in 1:n] for j in 1:m]
+
+    # Create an orthonormal basis for deim_interpolation_indices
+    Q, _ = qr(snapshot_matrix)
+    deim_basis = Matrix(Q[:, 1:5])
+
+    @testset "deim_interpolation_indices type stability" begin
+        rep = JET.report_call(ModelOrderReduction.deim_interpolation_indices, (Matrix{Float64},))
+        @test length(JET.get_reports(rep)) == 0
+    end
+
+    @testset "matricize type stability" begin
+        rep = JET.report_call(ModelOrderReduction.matricize, (Vector{Vector{Float64}},))
+        @test length(JET.get_reports(rep)) == 0
+    end
+
+    @testset "POD constructor type stability" begin
+        # Matrix constructor
+        rep1 = JET.report_call(ModelOrderReduction.POD, (Matrix{Float64}, Int))
+        @test length(JET.get_reports(rep1)) == 0
+
+        # Vector{Vector} constructor
+        rep2 = JET.report_call(ModelOrderReduction.POD, (Vector{Vector{Float64}}, Int))
+        @test length(JET.get_reports(rep2)) == 0
+    end
+
+    @testset "reduce! with SVD type stability" begin
+        pod = POD(snapshot_matrix, 3)
+        rep = JET.report_call(ModelOrderReduction.reduce!, (typeof(pod), typeof(SVD())))
+        @test length(JET.get_reports(rep)) == 0
+    end
+end

test/runtests.jl

@@ -6,6 +6,9 @@ end
 @safetestset "Explicit Imports" begin
     include("explicit_imports.jl")
 end
+@safetestset "JET Static Analysis" begin
+    include("jet.jl")
+end
 @safetestset "POD" begin
     include("DataReduction.jl")
 end