diff --git a/src/dist/number/fix.jl b/src/dist/number/fix.jl index c1394916..c25df082 100644 --- a/src/dist/number/fix.jl +++ b/src/dist/number/fix.jl @@ -3,7 +3,7 @@ using SymPy @vars varint @vars v2 -export DistFix, bitblast, bitblast_linear, bitblast_exponential, bitblast_exact, unit_exponential, exponential, laplace, unit_gamma, shift_point_gamma, n_unit_exponentials +export DistFix, bitblast, bitblast_linear, bitblast_exponential, bitblast_exact, unit_exponential, exponential, laplace, unit_gamma, shift_point_gamma, n_unit_exponentials, geometric ################################## # types, structs, and constructors @@ -53,7 +53,7 @@ function Base.convert(::Type{DistFix{W2, F2}}, x::DistFix{W1, F1}) where {W1,W2, mantissa = convert(DistInt{W1+(F2-F1)}, x.mantissa) mantissa <<= (F2-F1) else #F2 < F1 - mantissa = drop_bits(DistInt{W1+(F2-F1)}, x.mantissa; last=true) + mantissa = drop_bits(DistInt{W1+(F2-F1)}, x.mantissa; last=false) end convert(DistFix{W2, F2}, DistFix{W1+(F2-F1), F2}(mantissa)) end @@ -646,3 +646,16 @@ end # Y # end +###################################################################################################### +# bit blasting geometric distribution: https://en.wikipedia.org/wiki/Geometric_distribution +###################################################################################################### + +# Creates a geometric distribution over the integers in the range [0, stop-1] +function geometric(::Type{DistFix{W, F}}, success::Float64, stop::Int) where {W, F} + #TODO: use convert + @assert ispow2(stop) + bits = Int(log2(stop)) + @assert W - F > bits + + convert(DistFix{W, F}, DistFix{W, 0}(unit_exponential(DistFix{bits+1, bits}, log(1 - success)*2^bits).mantissa)) +end diff --git a/test/dist/number/fix_test.jl b/test/dist/number/fix_test.jl index 9c0293b7..a2c80f1c 100644 --- a/test/dist/number/fix_test.jl +++ b/test/dist/number/fix_test.jl @@ -388,5 +388,11 @@ end #TODO test for beta = 0 #TODO test for positive beta + # Truncated Geometric + x = Truncated(Geometric(0.8), 0, 32) + @test_throws Exception a = geometric(DistFix{6, 2}, 0.8, 32) + a = geometric(DistFix{8, 2}, 0.8, 32) + @test pr(a)[0.0] ≈ pdf(x, 0.0) + @test pr(a)[19] ≈ pdf(x, 19.0) end