rename `LogDensityFunction` to `NamedLogDensity`

[yebai]

In a discussion with @penelopeysm, I suggested making LogDensityFunction more user-facing and using the name NamedLogDensity for better clarity.

In the longer term, we could move NamedLogDensity to AbstractMCMC, replacing LogDensityModel.

[yebai]

For clarity, the main difference between NamedLogDensity and LogDensityProblem is that

1. the former has VarNamed variables (VarNamedTuple #900) instead of a plain vector of floating numbers
2. the parameters could have mixed discrete and continuous RVs

(1) facilitates defining modelling and inference-related operations like fix and condition. Meanwhile, (2) is needed by Turing.

[penelopeysm]

I'm still waiting for the final bunch of context-related PRs to be reviewed + merged, but am starting to think about this. I think we need to move beyond the LogDensityProblems interface, or rather, we need to create a superset of it.

NamedLogDensity is a fine name for that, it would look something like this:

                                                                           
                ┌───────────────────────┐                                  
                │                       │        ┌────────────────────┐    
                │   something that      │        │ something that     │    
                │   satisfies the       │        │ satisfies the      │    
                │   NamedLogDensity     │        │ LogDensityProblems │    
                │     interface         │        │     interface      │    
                │   contains a mapping  │        └──────────┬─────────┘    
                │   of keys => indices  │                   │              
                │                       │                   │              
                └────────────┬──────────┘                   │              
                             │                              │              
  ┌───────────────────┐      │                              │              
  │  mapping of       │      ▼     ┌────────────────────┐   ▼   ┌───────┐  
  │  keys => values   │◄──────────►│  vector of values  ├──────►│ logp  │  
  │                   │            └────────────────────┘       └───────┘  
  │  e.g. NamedTuple  │                                                    
  │   or Dict         │                                                    
  └───────────────────┘                                                    
                                                                                                                                                        

The right hand side is just the ordinary LDP interface: it's some object which when passed a vector of values can calculate logp. The left hand side is where the conversion needs to take place, and right now those arrows are essentially vi[:] and unflatten(vi, vec).

It follows from this diagram that anything that satisfies the NamedLogDensity interface can also satisfy the LogDensityProblems interface.

IMO this interface should be defined in a separate package, a very lightweight one like LDP. So I will probably start playing around with defining an interface. The concrete application to DPPL can come a bit later (especially since it may shift depending on the progress of other things like of-types).

[mhauru]

Note that dynamic models complicate the LDP ideal of having a vector map to a logp, in that the dimension of the vector is variable and, worse, the meaning of the elements changes based on values. Here's a pretty pathological demo of the latter:

@model function f()
    x ~ Normal()
    if x > 0
        a = Vector{Float64)(undef, 1)
        b = Vector{Float64)(undef, 2)
    else
        a = Vector{Float64)(undef, 2)
        b = Vector{Float64)(undef, 1)
    end
    a .~ Normal(10.0)
    b .~ Normal(-10.0)
end

This model always takes 4 parameters, but the 3rd parameter has a very strange and discontinuous logp landscape (because sometimes it's b[1], sometimes a[2]). A more natural (or just plain correct?) way to consider this model would be to think of it as a model with 5 parameters, one of which does not contribute to the logp (but which one that is varies).

I'm not sure exactly how this should affect the design in this issue, but it feels relevant. Maybe NamedLogDensity allows us to do this properly, and in the above example feed 5-parameter vectors to LDP?

[penelopeysm]

I think for that model, my current understanding is that all of the arrows in this diagram should be well-defined relations, in that if you put in one of the inputs, there is a correct output that can be obtained (which may be an error). I think to first-order, that's all the interface really cares about; whether it's a valid statistical model or discontinuous or differentiable is IMO a secondary concern.

The NamedLogDensity interface demands Dict -> logp.* We have a pretty good way to do this already using init!! (which does feel a bit like a misnomer now). If you attempt to calculate logp for Dict(x => 0.5, a = [1.0, 2.0], b = [3.0]) it should error because the length doesn't match, so we do have a well-defined way to obtain a logp or error from an input dict.

Even the Vectors are, in principle, well-behaved For example, let's say we want to convert a Vector to a Dict, then we expect:

[0.5, 1.0, 2.0, 3.0] should be sent to Dict(x => 0.5, a => [1.0], b => [2.0, 3.0])

[-0.5, 1.0, 2.0, 3.0] should be sent to Dict(x => -0.5, a => [1.0, 2.0], b => [3.0])

From the diagram above, this should be accomplished by values_as ∘ unflatten.† To do this correctly you would need to evaluate the model: firstly you need to know what order the variables come in, and secondly you need to check the conditional on x, which is only found in the model. Currently, unflatten probably won't get this right, because it doesn't evaluate the model, it just uses whatever the shape of the VarInfo currently is. But in principle it should be possible to make it work correctly via an extra model evaluation.

The LogDensityProblems interface demands Vector -> logp. Right now we accomplish this with unflatten + evaluate!!, which is problematic for the same reasons as above. But again it should in principle be made to work correctly.

Putting all this together, I think the following makes sense modulo naming:

function dict_to_logp end

• The main interface function with NamedLogDensity. If you implement this, it satisfies the NamedLogDensity interface.
• We overload this with init!!. This would always be correct for the model above and if correctness is the main aim, then we should always use this function and never touch the vector stuff again. Not sure about performance.

function vector_to_dict end

• This would be optional but if this is overloaded then the resulting object would also satisfies the LDP interface via composition.
• We could overload this with values_as ∘ unflatten. It would be incorrect for this model, but no worse than where we are currently. Or we could choose to just not overload this.

LogDensityProblems.logdensity

• Since this exists already we don't need to make a new function.
• By default it would use dict_to_logp ∘ vector_to_dict. Naturally people can still overload that for efficiency purposes if desired (for example maybe we want to keep using evaluate!! ∘ unflatten).

function dict_to_vector end

• We don't need to touch this in DynamicPPL (hence this arrow doesn't exist), but for a general purpose interface what this means is that if somebody has an object which already satisfies the LogDensityProblems interface, they can implement this and then it will automatically satisfy the NamedLogDensity interface via LogDensityProblems.logdensity ∘ dict_to_vector.

---

* Maybe Dict isn't the best, maybe it's NamedTuple or the hypothetical VarNamedTuple, the point is that it's some sort of key-value map.

† the functions don't compose as written, I'm just being loose with the arguments

[penelopeysm]

The main concern I have about this, from a software perspective, is that I'm unsure how to make Julia work correctly for this sort of situation where overloading a function means it will obey a new interface. For example I don't know how to make LogDensityProblems.capability return a different thing if someone implements vector_to_dict and indeed it might be impossible.

If that's the case (I haven't tried), then that suggests that we should just completely decouple NamedLogDensity from LogDensityProblems. So if somebody implements a NLD interface, they have to implement the LDP interface separately (although they can of course make use of composition to reduce code duplication).

That might make life easier for us in some ways because it can just be a Turing-only thing and we don't need to generalise it beyond what we need out of it. I'm very uncertain about which direction to take this in but I see some merit in making NLD independent of LDP.

[mhauru]

I think for that model, my current understanding is that all of the arrows in this diagram should be well-defined relations, in that if you put in one of the inputs, there is a correct output that can be obtained (which may be an error). I think to first-order, that's all the interface really cares about; whether it's a valid statistical model or discontinuous or differentiable is IMO a secondary concern.

[...]

Even the Vectors are, in principle, well-behaved For example, let's say we want to convert a Vector to a Dict, then we expect:

[0.5, 1.0, 2.0, 3.0] should be sent to Dict(x => 0.5, a => [1.0], b => [2.0, 3.0])

[-0.5, 1.0, 2.0, 3.0] should be sent to Dict(x => -0.5, a => [1.0, 2.0], b => [3.0])

This is a well-defined interface. However, I don't think it's the maximally useful interface for us. It's similar to how LDP currently behaves with dynamic models, and I would like to use NLD as an opportunity to do something better. Namely, I would like for each position in the vector to be associated with a particular VarName, and for that position to never change. Thus the above conversions would have to be

[0.5, 1.0, x, 2.0, 3.0] sent to Dict(x => 0.5, a => [1.0], b => [2.0, 3.0])
[-0.5, 1.0, 2.0, x, 3.0] sent to Dict(x => -0.5, a => [1.0, 2.0], b => [3.0])

where I think x should probably be allowed to be any value. Actually, more principled would be to demand that x === missing, but I'm not sure how that would play with samplers.

The reason this would be better is that I suspect it would make the life of samplers easier. For instance, say some sampler tries to figure out the correct step size to take along the third dimension v[3] of the space it's exploring. The logp landscape in that direction would still have a discontinuity at v[1] = 0.0, but I would think that would be easier to deal with than a landscape that just changes wildly in shape as you cross v[1] = 0.0. It also feels to me like a more rich and faithful description of what is going on with this model.

[mhauru]

If you want to make the above situation worse, imagine in my example model that a is Float64s and b is Ints. Now the type of one element in the parameter vector v has to change based on the value of another element. Whereas in my proposal at least all the elements have a fixed type, and some of them just don't have an effect on the logp in some parts of the parameter space.

My point is that a[2] and b[1] are different parameters and conceptually have nothing to do with each other. And because samplers just see an interface that takes vectors and gives logps, if sometimes some element of that vector maps to a[2] and sometimes to b[1] this has the potential to cause all sorts of confusions and trouble. This has been a problem so far, but NLD sounds like an opportunity to fix that.