Setup stacktrace truncation on common types

SciML · Feb 26, 2023 · 9b9f4ae · 9b9f4ae
1 parent b9eee8a
commit 9b9f4ae
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 0 deletions.
diff --git a/Project.toml b/Project.toml
@@ -37,6 +37,7 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
@@ -71,6 +72,7 @@ SimpleNonlinearSolve = "0.1.8"
 StaticArraysCore = "1.4"
 StochasticDiffEq = "6.20"
 Tracker = "0.2"
+TruncatedStacktraces = "1"
 Zygote = "0.6"
 ZygoteRules = "0.2"
 julia = "1.6"

diff --git a/src/adjoint_common.jl b/src/adjoint_common.jl
@@ -21,6 +21,19 @@ struct AdjointDiffCache{UF, PF, G, TJ, PJT, uType, JC, GC, PJC, JNC, PJNC, rateT
  issemiexplicitdae::Bool
 end
 
+function Base.show(io::IO,
+ t::Type{AdjointDiffCache{UF, PF, G, TJ, PJT, uType, JC, GC, PJC, JNC, PJNC, rateType, DG1,
+ DG2, DI,
+ AI, FM}}) where {UF, PF, G, TJ, PJT, uType, JC, GC, PJC, JNC, PJNC, rateType, DG1,
+ DG2, DI,
+ AI, FM}
+ if TruncatedStacktraces.VERBOSE[]
+ print(io, "AdjointDiffCache{$UF,$PF,$G,$TJ,$PJT,$uType,$JC,$GC,$PJC,$JNC,$PJNC,$rateType,$DG1,$DG2,$DI,$AI,$FM}")
+ else
+ print(io, "AdjointDiffCache{…}")
+ end
+end
+
 """
  adjointdiffcache(g,sensealg,discrete,sol,dg,alg;quad=false)
 

diff --git a/src/backsolve_adjoint.jl b/src/backsolve_adjoint.jl
@@ -11,6 +11,19 @@ struct ODEBacksolveSensitivityFunction{C <: AdjointDiffCache, Alg <: BacksolveAd
  noiseterm::Bool
 end
 
+function Base.show(io::IO,
+ t::Type{ODEBacksolveSensitivityFunction{C, Alg,
+ uType, pType,
+ fType}}) where {C, Alg,
+ uType, pType,
+ fType}
+ if TruncatedStacktraces.VERBOSE[]
+ print(io, "ODEBacksolveSensitivityFunction{$C,$Alg,$uType,$pType,$fType}")
+ else
+ print(io, "ODEBacksolveSensitivityFunction{…}")
+ end
+end
+
 function ODEBacksolveSensitivityFunction(g, sensealg, discrete, sol, dgdu, dgdp, f, alg;
  noiseterm = false)
  diffcache, y = adjointdiffcache(g, sensealg, discrete, sol, dgdu, dgdp, f, alg;

diff --git a/src/forward_sensitivity.jl b/src/forward_sensitivity.jl
@@ -35,6 +35,18 @@ struct ODEForwardSensitivityFunction{iip, F, A, Tt, OJ, J, JP, S, PJ, TW, TWt, U
  isautojacmat::Bool
  colorvec::CV
 end
+
+function Base.show(io::IO,
+ t::Type{ODEForwardSensitivityFunction{iip, F, A, Tt, OJ, J, JP, S, PJ, TW, TWt, UF, PF, JC,
+ PJC, Alg, fc, JM, pJM, MM, CV}}) where {iip, F, A, Tt, OJ, J, JP, S, PJ, TW, TWt, UF, PF, JC,
+ PJC, Alg, fc, JM, pJM, MM, CV}
+ if TruncatedStacktraces.VERBOSE[]
+ print(io, "ODEForwardSensitivityFunction{$iip,$F,$A,$Tt,$OJ,$J,$JP,$S,$PJ,$TW,$TWt,$UF,$PF,$JC,$PJC,$Alg,$fc,$JM,$pJM,$MM,$CV}")
+ else
+ print(io, "ODEForwardSensitivityFunction{…}")
+ end
+end
+
 has_original_jac(S) = isdefined(S, :original_jac) && S.jac !== nothing
 
 struct NILSSForwardSensitivityFunction{iip, sensefunType, senseType, MM} <:

diff --git a/src/interpolating_adjoint.jl b/src/interpolating_adjoint.jl
@@ -15,6 +15,19 @@ struct ODEInterpolatingAdjointSensitivityFunction{C <: AdjointDiffCache,
  noiseterm::Bool
 end
 
+function Base.show(io::IO,
+ t::Type{ODEInterpolatingAdjointSensitivityFunction{C,
+ Alg,uType, SType, CPS, pType,
+ fType}}) where {C,
+ Alg,uType, SType, CPS, pType,
+ fType}
+ if TruncatedStacktraces.VERBOSE[]
+ print(io, "ODEInterpolatingAdjointSensitivityFunction{$C,$Alg,$uType,$SType,$CPS,$pType,$fType}")
+ else
+ print(io, "ODEInterpolatingAdjointSensitivityFunction{…}")
+ end
+end
+
 mutable struct CheckpointSolution{S, I, T, T2}
  cpsol::S # solution in a checkpoint interval
  intervals::I # checkpoint intervals

diff --git a/src/quadrature_adjoint.jl b/src/quadrature_adjoint.jl
@@ -11,6 +11,21 @@ struct ODEQuadratureAdjointSensitivityFunction{C <: AdjointDiffCache,
  f::fType
 end
 
+function Base.show(io::IO,
+ t::Type{ODEQuadratureAdjointSensitivityFunction{C,
+ Alg,
+ uType, SType,
+ fType}}) where {C,
+ Alg,
+ uType, SType,
+ fType}
+ if TruncatedStacktraces.VERBOSE[]
+ print(io, "ODEQuadratureAdjointSensitivityFunction{$C,$Alg,$uType,$SType,$fType}")
+ else
+ print(io, "ODEQuadratureAdjointSensitivityFunction{…}")
+ end
+end
+
 function ODEQuadratureAdjointSensitivityFunction(g, sensealg, discrete, sol, dgdu, dgdp,
  alg)
  diffcache, y = adjointdiffcache(g, sensealg, discrete, sol, dgdu, dgdp, sol.prob.f, alg;