Mixed-state presentation: index beliefs by a generic projection

CHANGELOG.md

@@ -29,6 +29,14 @@ brute-force reference or a closed-form value and has a regression test.
   arithmetic. A 4-state survey support with 6 maximal rectangles went from hours (all
   `C(132, 5)` row subsets in `Fraction` arithmetic) to 0.3 s; 7 rectangles take about 7 s.
 
+### Mixed-state presentation speed
+
+- `build_mixed_state_presentation` indexes numeric beliefs by a projection onto fixed positive weights
+  instead of their first coordinate. Beliefs equal to within the matching tolerance still land in the
+  same window, but many beliefs sharing a first coordinate (often 0) no longer make every lookup scan
+  all states. Reverse inference on a 4-state process reaches a 5,000-state cap in 2.5 s instead of
+  135 s; reverse machines of the 100 finite 3-state survey processes are unchanged.
+
 ### Anytime generator search
 
 - `generative_model`, `generative_model_from_bidirectional` and `minimal_wyner_model` take

sofic/generators/mixed_state_construction.py

@@ -97,12 +97,21 @@ def build_mixed_state_presentation(
     symbolic_unique: list[MixedState] = []
     numeric_keys: list[float] = []
     numeric_unique: list[MixedState] = []
+    # Index numeric beliefs by a projection onto fixed positive weights summing to one. Beliefs equal
+    # coordinatewise to within _NUMERIC_WINDOW project to within it too, and unlike a single
+    # coordinate the projection rarely ties, so a lookup scans few candidates.
+    weights = np.array([0.5 + (index * 0.6180339887498949) % 1.0 for index in range(len(basis))])
+    weights = weights / weights.sum()
+
+    def key(state: MixedState) -> float:
+        return float(np.dot(weights, np.asarray(state.belief, dtype=float)))
+
     queue: deque[MixedState] = deque()
 
     def candidates(state: MixedState) -> Sequence[MixedState]:
-        if is_symbolic(state.belief[0]):
+        if any(is_symbolic(b) for b in state.belief):
             return unique
-        lead = float(state.belief[0])
+        lead = key(state)
         window = _NUMERIC_WINDOW * (1.0 + abs(lead))
         low = bisect.bisect_left(numeric_keys, lead - window)
         high = bisect.bisect_right(numeric_keys, lead + window)
@@ -122,11 +131,12 @@ def register(state: MixedState) -> MixedState:
             )
         discovered[state] = state
         unique.append(state)
-        if is_symbolic(state.belief[0]):
+        if any(is_symbolic(b) for b in state.belief):
             symbolic_unique.append(state)
         else:
-            position = bisect.bisect_right(numeric_keys, float(state.belief[0]))
-            numeric_keys.insert(position, float(state.belief[0]))
+            lead = key(state)
+            position = bisect.bisect_right(numeric_keys, lead)
+            numeric_keys.insert(position, lead)
             numeric_unique.insert(position, state)
         graph.add_state(state)
         queue.append(state)

tests/test_mixed_state_construction.py

@@ -26,3 +26,29 @@ def test_max_states_counts_distinct_beliefs():
     with pytest.raises(MixedStateExplosionError, match="max_states=50"):
         build_mixed_state_presentation(sns(), max_states=50)
     assert len(list(build_mixed_state_presentation(golden_mean(), max_states=3).states())) == 3
+
+
+def test_infinite_reverse_presentation_reaches_its_cap_quickly() -> None:
+    """Regression: indexing beliefs by their first coordinate made every lookup scan all states when
+    many beliefs share it, so this 4-state reverse inference ran for many minutes before its cap."""
+    import numpy as np
+
+    from sofic.examples.epsilon_machines import from_symbol_matrices
+    from sofic.generators.bidirectional_construction import infer_reverse_epsilon_machine
+
+    mats = {x: np.zeros((4, 4)) for x in (0, 1)}
+    for x, i, j, p in (
+        (0, 1, 2, 0.903832833130349),
+        (0, 2, 3, 0.9529570928839382),
+        (0, 3, 0, 0.0202046632250718),
+        (1, 0, 1, 1.0),
+        (1, 1, 0, 0.0961671668696509),
+        (1, 2, 2, 0.0470429071160618),
+        (1, 3, 1, 0.9797953367749282),
+    ):
+        mats[x][i, j] = p
+    forward = from_symbol_matrices(tuple(range(4)), (0, 1), mats)
+    start = time.monotonic()
+    with pytest.raises(MixedStateExplosionError):
+        infer_reverse_epsilon_machine(forward, max_states=5_000)
+    assert time.monotonic() - start < 60