Repository: j-towns/craystack

Issue:
Rare bug in non-dyadic push and pop
We were a bit careless in our implementation of the non-dyadic push and pop functions, and there's a bug, which is very rare at 64 bit precision. A reproduction, with explanation:
```python
from craystack import rans

import numpy as np

# The renormalization in rans.push_with_finer_prec will keep moving
# bits into the tail until the head is less than
# (rans_l // 3) << r_t      == 3074457342754947072

# However, the true upper bound on the head after rans.pop_with_finer_prec
# is less strict, it's equal to
# ((1 << r_s) - 1) // 3 + 1 == 3074457345618258603

# We can break the inverse property by constructing a message and
# pop such that the result is between the two values above. That means
# that push thinks it needs to undo a renorm when actually no renorm
# took place in the pop:

# Set s to the largest possible value
m = np.uint64([(1 << 63) - 1]), ()

# Use N = 3 as the denominator
N = 3

u, popfun = rans.pop_with_finer_prec(m, N)
m_ = popfun(u, 1)

# u  == array([1], dtype=uint64)
# m_ == (array([3074457345618258602], dtype=uint64), ())
# s is between the two values above, and that causes a problem...

m__ = rans.push_with_finer_prec(m_, u, 1, N)

assert rans.message_equal(m, m__), 'm = ' + str(m) + '\n m__ = ' + str(m__)
```
This isn't too difficult to fix. We need to alter what we call 'renormalization', to ensure that before we perform a pop the message head h is in an interval `[a * N, a * b * N)` for some natural numbers `a`, `b`. For us it is sensible to use `b = (1 << 32)`. The value of `a` should be as large as possible, such that we don't get overflow, that is, such that `a * b * N <= (1 << 64)`. If we're fixing `b = (1 << 32)` then that condition on `a` becomes `a <= (1 << 32) / N`. I think setting `a = (1 << 32) // N` is the sensible thing to do.

We also need to 'undo' the new renormalization step at the beginning of pop, with a corresponding step at the end of push.

Hopefully I get round to fixing this as soon as possible.

Credit to Fabian's blog post https://fgiesen.wordpress.com/2014/02/02/rans-notes/ for helping me understand this.

Discussion:

Patch:
diff --git a/craystack/rans.py b/craystack/rans.py
index 0d1ed66..e37ced1 100644
--- a/craystack/rans.py
+++ b/craystack/rans.py
@@ -90,25 +90,48 @@ def unflatten(arr, shape):
 def message_equal(message1, message2):
     return np.all(flatten(message1) == flatten(message2))
 
-
 def push_with_finer_prec(ans_state, starts, freqs, precisions):
     head, tail = ans_state
+    # head in [2 ^ 32, 2 ^ 64)
     starts, freqs, precisions = map(atleast_1d, (starts, freqs, precisions))
     idxs = head >= ((rans_l // precisions) << 32) * freqs
     if np.any(idxs):
         tail = stack_extend(tail, np.uint32(head[idxs]))
         head = np.copy(head)  # Ensure no side-effects
         head[idxs] >>= 32
+    # head in [((rans_l // precisions)) * freqs), ((rans_l // precisions)) * freqs) << 32)
 
     # calculate next state s' = 2^r * (s // p) + (s % p) + c
     head_div_freqs, head_mod_freqs = np.divmod(head, freqs)
-    return head_div_freqs*precisions + head_mod_freqs + starts, tail
+    head = head_div_freqs * precisions + head_mod_freqs + starts
+
+    # head in [precisions * (2 ^ 32 // precisions), precisions * ((2 ^ 32 // precisions) << 32))
+    # check which entries need renormalizing
+    idxs = head < rans_l
 
+    # how many 32*n bits do we need from the tail?
+    n = np.sum(idxs)
+    if n > 0:
+        # new_head = 32*n bits from the tail
+        # tail = previous tail, with 32*n less bits
+        tail, new_head = stack_slice(tail, n)
+
+        # update LSBs of head, where needed
+        head[idxs] = (head[idxs] << 32) | new_head
+    # head in [2 ^ 32, 2 ^ 64)
+    return head, tail
 
 def pop_with_finer_prec(ans_state, precisions):
     precisions = atleast_1d(precisions)
     head_, tail_ = ans_state
+    # head_ in [2 ^ 32, 2 ^ 64)
+    idxs = head_ >= precisions * ((rans_l // precisions) << 32)
+    if np.any(idxs):
+        tail_ = stack_extend(tail_, np.uint32(head_[idxs]))
+        head_ = np.copy(head_)  # Ensure no side-effects
+        head_[idxs] >>= 32
 
+    # head in [precisions * (2 ^ 32 // precisions), precisions * ((2 ^ 32 // precisions) << 32))
     # s' mod 2^r
     cfs = head_ % precisions
 


Test patch:
diff --git a/craystack/rans_test.py b/craystack/rans_test.py
index 2a74552..a43e3b7 100644
--- a/craystack/rans_test.py
+++ b/craystack/rans_test.py
@@ -64,3 +64,35 @@ def test_base_message():
     num_bits = num_lower_bits*100_000
     assert (head_rnd >> (num_lower_bits - 1) == 1).all()
     assert  0.48 < popcount(head_rnd)/num_bits < 0.52
+
+def test_finer_prec():
+    # Regression test against https://github.com/j-towns/craystack/issues/14.
+    # The renormalization in rans.push_with_finer_prec will keep moving
+    # bits into the tail until the head is less than
+    # (rans_l // 3) << r_t      == 3074457342754947072
+
+    # However, the true upper bound on the head after rans.pop_with_finer_prec
+    # is less strict, it's equal to
+    # ((1 << r_s) - 1) // 3 + 1 == 3074457345618258603
+
+    # We can break the inverse property by constructing a message and
+    # pop such that the result is between the two values above. That means
+    # that push thinks it needs to undo a renorm when actually no renorm
+    # took place in the pop:
+
+    # Set s to the largest possible value
+    m = np.uint64([(1 << 63) - 1]), ()
+
+    # Use N = 3 as the denominator
+    N = 3
+
+    u, popfun = rans.pop_with_finer_prec(m, N)
+    m_ = popfun(u, 1)
+
+    # u  == array([1], dtype=uint64)
+    # m_ == (array([3074457345618258602], dtype=uint64), ())
+    # s is between the two values above, and that causes a problem...
+
+    m__ = rans.push_with_finer_prec(m_, u, 1, N)
+
+    assert rans.message_equal(m, m__), 'm = ' + str(m) + '\n m__ = ' + str(m__)