Document random sampling

Author: MilesCranmer

docs/make.jl

@@ -1,5 +1,6 @@
 using Documenter
 using DynamicExpressions
+using Random: AbstractRNG
 
 makedocs(;
     sitename="DynamicExpressions.jl",

docs/src/utils.md

@@ -20,6 +20,16 @@ convert(::Type{<:AbstractExpressionNode{T1}}, n::AbstractExpressionNode{T2}) whe
 hash(tree::AbstractExpressionNode{T}, h::UInt; break_sharing::Val=Val(false)) where {T}
 ```
 
+## Sampling
+
+There are also methods for random sampling of nodes:
+
+```@docs
+NodeSampler
+rand(rng::AbstractRNG, tree::AbstractNode; break_sharing::Val=Val(false))
+rand(rng::AbstractRNG, sampler::NodeSampler{N,F,Nothing}) where {N,F}
+```
+
 ## Internal utilities
 
 Almost all node utilities are crafted using the `tree_mapreduce` function,

src/Random.jl

@@ -2,10 +2,11 @@ module RandomModule
 
 import Compat: Returns
 import Random: AbstractRNG
+import Base: rand
 import ..EquationModule: AbstractNode, tree_mapreduce, filter_map
 
 """
-    NodeSampler(; tree, filter=Returns(true), weighting=nothing, break_sharing=Val(false))
+    NodeSampler(; tree, filter::Function=Returns(true), weighting::Union{Nothing,Function}=nothing, break_sharing::Val=Val(false))
 
 Defines a sampler of nodes in a tree. `filter` can be used to pre-filter
 nodes on which to sample.
@@ -23,51 +24,52 @@ nodes on which to sample.
   returns a weight for the node, if it passes the filter, proportional
   to the probability of sampling the node. If `nothing`, all nodes are
   sampled uniformly.
-- `break_sharing::Union{Bool,Val}`: If `true` or `Val(true)`, the
+- `break_sharing::Val`: If `Val(true)`, the
   sampler will break sharing in the tree, and sample nodes uniformly
   from the tree.
 """
 Base.@kwdef struct NodeSampler{
-    N<:AbstractNode,F<:Function,W<:Union{Nothing,Function},B<:Union{Bool,Val}
+    N<:AbstractNode,F<:Function,W<:Union{Nothing,Function},B<:Val
 }
     tree::N
     weighting::W = nothing
     filter::F = Returns(true)
     break_sharing::B = Val(false)
 end
 
-Base.rand(rng::AbstractRNG, tree::AbstractNode) = rand(rng, NodeSampler(; tree))
-function Base.rand(rng::AbstractRNG, sampler::NodeSampler{N,F,Nothing}) where {N,F}
-    break_sharing = if sampler.break_sharing isa Val
-        sampler.break_sharing
-    else
-        sampler.break_sharing ? Val(true) : Val(false)
-    end
-    n = count(sampler.filter, sampler.tree; break_sharing)
+"""
+    rand(rng::AbstractRNG, tree::AbstractNode; break_sharing::Val=Val(false))
+
+Sample a node from a tree according to the default sampler `NodeSampler(; tree, break_sharing)`.
+"""
+rand(rng::AbstractRNG, tree::AbstractNode; break_sharing::Val=Val(false)) = rand(rng, NodeSampler(; tree, break_sharing))
+
+"""
+    rand(rng::AbstractRNG, sampler::NodeSampler)
+
+Sample a node from a tree according to the sampler `sampler`.
+"""
+function rand(rng::AbstractRNG, sampler::NodeSampler{N,F,Nothing}) where {N,F}
+    n = count(sampler.filter, sampler.tree; sampler.break_sharing)
     idx = rand(rng, 1:n)
     i = Ref(0)
     out = Ref(sampler.tree)
-    foreach(sampler.tree; break_sharing) do node
+    foreach(sampler.tree; sampler.break_sharing) do node
         if @inline(sampler.filter(node)) && (i[] += 1) == idx
             out[] = node
         end
         nothing
     end
     return out[]
 end
-function Base.rand(rng::AbstractRNG, sampler::NodeSampler{N,F,W}) where {N,F,W<:Function}
-    break_sharing = if sampler.break_sharing isa Val
-        sampler.break_sharing
-    else
-        sampler.break_sharing ? Val(true) : Val(false)
-    end
+function rand(rng::AbstractRNG, sampler::NodeSampler{N,F,W}) where {N,F,W<:Function}
     weights = filter_map(
-        sampler.filter, sampler.weighting, sampler.tree, Float64; break_sharing
+        sampler.filter, sampler.weighting, sampler.tree, Float64; sampler.break_sharing
     )
     idx = sample_idx(rng, weights)
     i = Ref(0)
     out = Ref(sampler.tree)
-    foreach(sampler.tree; break_sharing) do node
+    foreach(sampler.tree; sampler.break_sharing) do node
         if @inline(sampler.filter(node)) && (i[] += 1) == idx
             out[] = node
         end