rusandris / StateTransitionNetworks.jl

Toolkit for dynamics on state transition networks
2 stars 0 forks source link

Huge compilation times for `renyi_entropy_spectrum` #9

Closed rusandris closed 5 months ago

rusandris commented 9 months ago

Compilation times for renyi_entropy_spectrum are huge:

 @time renyi_entropy_spectrum(P,[1.1,1.2,1.3];x=x,sparse=true)
q = 1.1
q = 1.2
q = 1.3
 22.886130 seconds (36.66 M allocations: 2.201 GiB, 5.52% gc time, 99.99% compilation time)
3-element Vector{Float64}:
 0.9418582839966015
 0.9340210407926257
 0.9267495136063403

Actual runtime is much faster:

@time renyi_entropy_spectrum(P,[1.1,1.2,1.3];x=x,sparse=true)
q = 1.1
q = 1.2
q = 1.3
  0.000846 seconds (348 allocations: 171.844 KiB)
3-element Vector{Float64}:
 0.9418582839966089
 0.9340210407926305
 0.9267495136063442

TODO: check if it's caused by KrylovKit dependencies, or too general types somewhere. Also possibly take out sparse -> non-parse switching, allowing only sparse calculation

rusandris commented 7 months ago

The problem is caused by renyi_entropy. Calling @code_warntype shows that the compiler cannot infer the return type of the function:

 MethodInstance for Core.kwcall(::@NamedTuple{x::Vector{Float64}, n::Float64}, ::typeof(renyi_entropy_old), ::SparseMatrixCSC{Float64, Int64}, ::Float64)
  from kwcall(::NamedTuple, ::typeof(renyi_entropy_old), P::SparseMatrixCSC{Float64, Int64}, q::Float64) @ Main ~/Documents/stn_research/code_warntype.jl:58
Arguments
  _::Core.Const(Core.kwcall)
  @_2::@NamedTuple{x::Vector{Float64}, n::Float64}
  @_3::Core.Const(renyi_entropy_old)
  P::SparseMatrixCSC{Float64, Int64}
  q::Float64
Locals
  x::Union{}
  n::Union{}
  tol::Union{}
  maxiter::Union{}
  verbosity::Union{}
  @_11::Union{Float64, Int64, Vector{Float64}}
Body::Any

Maybe the if-else statements which decide which formula is used (thus it can only be known at runtime, based on the values of the arguments) make it difficult to fix the return type ( becomes Any).

Simplifying renyi_entropy by dividing it into smaller methods seems to solve the problem. This technique is also mentioned in performance tips.

Now the compile time is much lower (basically the compile time of eigsolve): 6.913546 seconds (9.11 M allocations: 629.936 MiB, 18.80% gc time, 97.31% compilation time)

See the fix here f8c2f3f

rusandris commented 5 months ago

Passes tests of e317379