Type Parameters for ObservationProcess

[THargreaves]

It looks like at some point we've diverged on the interpretation of the type parameter for ObservationProcess.

As a recap, we have just the one type parameter

abstract type ObservationProcess{T} end

This is used in two places with contradictory meanings. First it is used to determine the output container type for forward simulation:

SSMProblems.jl/src/utils/forward_simulation.jl

Lines 9 to 13 in 672da2f

	T_dyn = eltype(LD)
	T_obs = eltype(OP)
	xs = Vector{T_dyn}(undef, T)
	ys = Vector{T_obs}(undef, T)

It is also used to define the main type parameter of the SSM which I believe corresponds to the type of the log-likelihoods generated by the SSM:

SSMProblems.jl/src/SSMProblems.jl

Lines 258 to 259 in 672da2f

	T = promote_type(eltype(OPT), eltype(LDT))
	return new{T,typeof(dyn),typeof(obs)}(dyn, obs)

I feel like both of these use-cases have a place. Is it worth having two type parameters?

One for the element type (which could be some weird abstract thing, e.g. a collection of Lévy jumps, rather than just a Vector{T}). One for the type returned by the logdensity function?

If we are doing this, the same should probably hold for the LatentDynamics. One for the state type, one for the type returned by logdensity.

Alternatively, it might be overkill to store the type of the state/observation. Are there any serious drawbacks that come from sampling x1, y1 and then constructing a vector from their empirical types?

Would especially appreciate thoughts from @charlesknipp as I believe it was your desire to store the log-likelihood type.

[charlesknipp]

I think you raise a very valid criticism of my idea. First of all, I agree that it doesn't make a whole lot of sense in terms of type parameterization. They both have a place, but I feel like we more commonly use the log-likelihood type for inference algorithms, so I focused on that element. Honestly, there is a world in which they coexist, but I think adding an additional parameter would just be extra clutter.

In Distributions.jl, the return type of logpdf is a translation of a state $x$ from some typed distribution. In essence, we can think of the pdf as $x \mapsto [0,1]$ where the result is a floating point representation; if things are type consistent in Distributions.jl, the return type of logpdf/pdf should be a floating point whose number of bits is equivalent to the state type XT.

Most of this is to say that we can instead of using the type parameter to determine the logpdf, we can use some type promotion given XT, since this shouldn't be a specification from the user. I will clear this up with a nicer explanation later, but for now, those are my very disorganized thoughts.

[THargreaves]

I'm not sure I'm quite following what you're saying but I agree with the general idea that it is clear cluttered to add loads of type parameters to handle both the log-likelihood type and state types.

My latest thoughts on this that the type parameters for the state/observation types are perhaps the overkill bit. Yes it's a bit more messy in the code, but the following is (I think—would appreciate confirmation) completely valid and type stable

x0 = initialise(rng model)
xs = Vector{typeof(x0)}(undef, T + 1)
xs[1] = x0

or something like that.

There would therefore only be one type parameter which would correspond to the log-likelihood type but more generally the floating point type that all fields in the model use (e.g. Q::Matrix{T}).

If you're happy with this in terms of autodiff, I think that's likely sufficient and avoids over-complication.

[charlesknipp]

the above is not necessarily type stable. Consider the scanario in which the initialization is Normal(0,1) and the state transition is Normal(A*x, 1) where we autodiff with respect to A. ForwardDiff.jl would fail here since the initialization doesn't account for the type of A

[THargreaves]

Ah I think I see what you mean. I guess I was suggesting this is the users' responsibility.

The would have to define

struct MyDynamics{T}
    mu_init::Vector{T}
    A::Matrix{T}
end

so that the type of A is known and is guaranteed to match the initialisation.

Is this what you mean?

[charlesknipp]

yeah, although it would have to be the other way around. The type of A is known and initialization is guaranteed to match it. In that case, we're kosher.

[THargreaves]

Exactly, though on reflection maybe that is a tough ask in general. I think I've shared with you before that Distributions isn't always type consistent. E.g. Uniform{T} always returns Float64 when sampling regardless of T.

[charlesknipp]

Uniform{T} always returns Float64 when sampling regardless of T.

Well that sucks. I honestly might raise an issue.

[charlesknipp]

I checked just to confirm and oh my goodness 🤮

julia> test = Uniform(0f0,1f0)
Uniform{Float32}(a=0.0f0, b=1.0f0)

julia> rand(test)
0.07834174250700543

julia> test = Uniform(BigFloat(0),BigFloat(1))
Uniform{BigFloat}(a=0.0, b=1.0)

julia> rand(test)
0.68051738282173801497521026249160058796405792236328125

julia> rand(test, 2)
2-element Vector{Float64}:
 0.8438465943304864
 0.821132972814556

[THargreaves]

Ikr. There's actually a long history of issues on this topic. Think the latest is here:
JuliaStats/Distributions.jl#1905

Think there has been some resistance to it for historical reasons, though some progress is being made. Normal used to have this behaviour but is now type-preserving.

But most other distributions can't be relied on right now.

[THargreaves]

I've been thinking of a possible third approach which whilst even more complicated, provides all the expressivity we need without the awkward T = promote_type(eltype(OPT), eltype(LDT)) step (which doesn't generalise to different state types).

The proposal would be to have two types for each of dyn and obs. The first would be the "arithmetic type" and be something like Float32, Float64, Dual etc. The second type would be the state/observation type—and it would be up to the user to make sure this matches the arithmetic type.

We then obtain the SSM type by promoting the two arithmetic types (or forcing them to be the same—I think this makes more sense).

With this approach, we can still preallocate memory for states/observations.

Maybe overkill, but let me know what you think.

[FredericWantiez]

Why do we need to parametrize the SSM with the loglikelihood type ? The states type is useful to dispatch / allocate but if we need the loglikelihood type, and it's not straightforward to derive from the state types (with the caveats around Distributions.jl instability ...), we could just have a helper function ?

[THargreaves]

@charlesknipp might be better suited to answer this as recall it being his addition

Happy to keep this discussion going, but I think until we have a good use-case for why these types are needed and a good general philosophy, we can strip things back to the simple case of just having one high-level arithmetic type.

I started the above change and found it led to a load of issues that would be faffy to fix. For example, when creating the container for the dense particle storage, you need to know the element type of the latent dynamics. You could find this by simulating x0 and then looking at its type but this seems like a unintuitive and error-prone workflow.

I'm going to return the above suggestion of having one type for the arithmetic type and one for the state/observation element type. I will make this change now.

Maybe the arithmetic type can be dropped and computed using a helper as Frederic suggests, but I will leave that for after the LAFI workshop.

[THargreaves]

On the point of creating a helper to infer these types, would this depend on something like Base.return_types? I'm surprised by its behaviour when promoting types.

julia> struct LinearTransformation{T}
           A::Matrix{T}
           b::Vector{T}
       end

julia> transform(t::LinearTransformation, x) = t.A * x + b
transform (generic function with 1 method)

julia> t = LinearTransformation(rand(2, 2), rand(2))
LinearTransformation{Float64}([0.25737465916134483 0.17756543003788494; 0.07369227300887338 0.597478835391764], [0.7104961946389236, 0.11151281420460368])

julia> x = rand(Float32, 2)
2-element Vector{Float32}:
 0.79233974
 0.34546995

julia> t.A * x + t.b
2-element Vector{Float64}:
 0.9757678862497214
 0.37631311516390975

julia> Base.return_types(transform, (LinearTransformation{Float64}, Vector{Float32}))
1-element Vector{Any}:
 Any

I'm not sure where Any is coming from.

[FredericWantiez]

Is it just a typo on transform(t::LinearTransformation, x) = t.A * x + t.b ? I get Vector{Float64} on julia 1.10

[charlesknipp]

Why do we need to parametrize the SSM with the loglikelihood type ?

The motivation for this was to easily be able to preallocate probability weights and initialize log-likelihoods without additional function calls. For a single model, this is all information we know at creation (ideally) and shouldn't require additional function calls. This is also useful for something like ForwardDiff, where the log-likelihoods should be initialized with Dual types; something that the user doesn't have to consider under the current structure.

The states type is useful to dispatch / allocate but if we need the loglikelihood type, and it's not straightforward to derive from the state types (with the caveats around Distributions.jl instability ...), we could just have a helper function ?

This is what we should do now, but I also worry that there's no easy way to determine what the likelihood type should be without running a computation every time we call this helper. I'm not sure how other modules handle cases like this, since I think Turing just operates with Float64s by default.