@@ -195,10 +195,11 @@ If this is not the case, returns `nothing`.
195195Note: The diffusion matrix ``Σ`` is the square root of the noise covariance matrix ``Q`` (see
196196[`covariance_matrix`](@ref)), defined via the Cholesky decomposition ``Q = Σ Σ^\\ top``.
197197"""
198- function diffusion_matrix (ds:: CoupledSDEs{IIP,D} ) where {IIP,D}
198+ function diffusion_matrix (ds:: CoupledSDEs{IIP,D} ):: AbstractMatrix where {IIP,D}
199199 if ds. noise_type[:invertible ]
200200 diffusion = diffusion_function (ds)
201201 A = diffusion (zeros (D), current_parameters (ds), 0.0 )
202+ A = A isa AbstractMatrix ? A : Diagonal (A)
202203 else
203204 @warn """
204205 The diffusion function of the `CoupledSDEs` cannot be expressed as a constant
@@ -218,7 +219,7 @@ If this is not the case, returns `nothing`.
218219
219220See also [`diffusion_matrix`](@ref).
220221"""
221- function covariance_matrix (ds:: CoupledSDEs )
222+ function covariance_matrix (ds:: CoupledSDEs ):: AbstractMatrix
222223 A = diffusion_matrix (ds)
223224 (A == nothing ) ? nothing : A * A'
224225end
@@ -241,10 +242,10 @@ Returns `g, noise_prototype`.
241242"""
242243function construct_diffusion_function (
243244 g, covariance, noise_prototype, noise_strength, D, IIP
244- )
245+ )
245246 if isnothing (g) # diagonal additive noise
246247 cov = isnothing (covariance) ? LinearAlgebra. I (D) : covariance
247- size (cov) != (D,D) &&
248+ size (cov) != (D, D) &&
248249 throw (ArgumentError (" Covariance matrix must be of size $((D, D)) " ))
249250 A = sqrt (cov)
250251 if IIP
@@ -269,4 +270,4 @@ function construct_diffusion_function(
269270 end
270271 end
271272 return g, noise_prototype
272- end
273+ end
0 commit comments