Plan to fix conjugate dirichlet problems

[axch]

Our dirichlet-categorical woes, namely #407, #451, and #452, are ultimately due to representing the probabilities in terms of the indexes of the options rather than in terms of their actual values. Changing the representation to values is slightly tricky, but doable.

The plan:

• [ ] Add known-failing unit tests for #407.
• [ ] Find as a library, or implement in Python, a key-value search tree that supports "accumulate over interval" queries. This would have the following spec:
  • A tree node has an element at self, a subtree of strictly smaller elements, a subtree of strictly greater elements, and a total weight. The weight of the current element is the difference between the self total and the totals of the two children.
  • A tree of size n supports insertion, removal, and "increment value" queries in O(log n).
  • A tree of size n supports inclusive and exclusive interval-sum queries in O(log n). (A interval-sum query takes an interval of elements and returns the sum of the values of all elements in the interval). Cute hack: primitively support querying interval from min (which takes only one descent into the tree), and compose double-ended interval queries as differences.
  • [Edit: Don't need split at element, in light of @luac's comment below] A tree supports a 'split at element k' operation in O(log n) that returns either two subtrees or two lists of subtrees, at the implementor's convenience. Semantics of how to indicate whether k was in tree and whether it appears in any of the returned subtrees at implementor discretion.
  • A tree supports a 'split somewhere' operation in O(1), which yields a decomposition into a subtree, an element (with value), and a subtree, with the property that all the above are in order. The implementation is to yield the left subtree, the current element, and the right subtree.
  • The 'split somewhere' operation can be used to sample an element proportional to its value in O(log n), by repeatedly splitting and choosing one of the three options proportional to its total weight.
  • Option: Do not officially store elements in the interior nodes, and have the semantics of 'split somewhere' be that the mid-point element also appears as, say, the least element of the second subtree.
  • Con: The intermediate nodes still have to have all the same information (pivot element, total weight), but now there will be twice as many nodes in the tree.
  • Pro: Runs of distinct elements that compare equal in the order can be represented reasonably; insertion and removal might be simpler too, because elements are only stored at the leaves.
• [ ] The operation we actually want for dirichlet categorical accepts two such trees of total size O(n), (the counts and the pseudo-counts) with potentially repeated elements, and samples an element proportional to its total weight in both trees, in O(log n).
  • This can be done by defining a representation of a "lazy pointwise sum" of these trees that stores both and implements 'split somewhere' (in O(log n) instead of O(1)) by invoking 'split somewhere' on the first of its trees and 'split here' on the second. I can imagine various ways to minimize the amount of traversing of the second tree, but in any case, this permits sampling in O(log n)^2, which is still much better than the O(n) that naively pointwise summing two such trees would cost.
  • [Edit: From @luac's comment below] Or we can implement this simply by picking one of the two trees proportional to its total weight in O(1) and then invoking its sample operation.
• [ ] Rewrite each dirichlet categorical in terms of the values returned rather than their indexes.
  • Redefine the SPAux for dirichlet categorical to store a map from the VentureValues that were actually returned to the counts of how many times.
  • In the made SP constructors, replace self.index with constructing a tree for the pseudo-counts of the given objects. This tree also represents the set of options that are possible. Don't forget to sum if equal objects appear multiple times (fixing #407).
  • simulate calls the "sample proportional to total weight" operation.
  • logDensity finds the value's total weight by looking the value up in the aux and the pseudo-count tree.
  • if the value is absent from the pseudo-count tree, the situation is impossible, return float('-inf') (fixing #451)
  • incorporate and unincorporate invoke "increment, inserting if absent" and "decrement, deleting if zero", respectively.
  • enumerateValues should not produce duplicates.
  • logDensityOfData should apply the same formula to traversing both trees.
  • if a value present in the counts is absent from the pseudo-counts, return float('-inf') (I think that's what the formula would do anyway if one of the alpha terms is zero).
  • if a value present in pseudo-counts is absent from the counts, treat as zero count (which should be fine).
  • doing this by value instead of index should fix #452.
  • Update gradientOfLogDensityOfData accordingly.
• [ ] This whole story relies on global equality and comparison of Venture values. Most should be fine; explicit IDs can be added to records, SPRefs, and any other values that need them in order to make this work for all types (not that people put SPs into dirichlet categorical options that often).
• [ ] Optional: There is a performance optimization possible if the options are known to be constant, for example if that argument is not supplied. In this case, by knowing the possible key set of the counts tree, one can arrange for it to always have the same structure as the pseudo-counts tree. Then all their top elements will always be the same, and the sum-sampling operation will be O(log n) instead of O(log n)^2.

Doing all this for Puma is #340.

[Edited: Replaced "multinomial" with "categorical" in names, pursuant to #450.]

[lenaqr]

The operation we actually want for dirichlet multinomial accepts two such trees of total size O(n), (the counts and the pseudo-counts) with potentially repeated elements, and samples an element proportional to its total weight in both trees, in O(log n).

Why not just sample one of the two trees proportional to their weight, then sample within the chosen tree in the straightforward way?

[axch]

Oh. Duh. That's even better.