[StackRnn] - add recurrent RBM stack with character-level LM notebook

Implements StackRnn: a cascaded/recurrent RBM where the visible layer at
each time step is the concatenation of the previous hidden state (context)
and the current sensory input — V[t] = [context | x_t].  Weights are
shared across time steps (RTRBM-style concatenation variant).

- stack_rnn.py: StackRnn with step(), reconstruct(), reset(), train(),
  make_layer() factory; supports greedy layer-wise training over sequences
  of shape (num_seq, T, sensory_size)
- test_rnn.py: single-layer, two-layer, and save/load tests
- moby_rnn.ipynb: character-level language model on Moby Dick; one-hot
  encoding, clamped-Gibbs next-char prediction, free text generation,
  hidden-state trace and character-distribution visualisations

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

src/rbm/stack_rnn.py

@@ -1,5 +1,184 @@
+import numpy as _np_cpu
+from .status import Status
+from .train import _to_gpu, cd_binary_binary, cd_gaussian_binary, cd_gaussian_gaussian, cd_binary_gaussian
 from .stack import Stack, StackType
+from .matrix import Mat, np, rms_error_accu, convert
+from .entity import Entity
+from .layer import Layer
+from .entity import EntityParams, TrainingParams
+
+_CD_FUNC = {
+    Entity.Type.BB_RBM: cd_binary_binary,
+    Entity.Type.GB_RBM: cd_gaussian_binary,
+    Entity.Type.GG_RBM: cd_gaussian_gaussian,
+    Entity.Type.BG_RBM: cd_binary_gaussian,
+}
+
 
 class StackRnn(Stack):
-	def __init__(self, name: str, work_dir: str = '.'):
-		Stack.__init__(self, StackType.Rnn, name, work_dir)
+    """Recurrent RBM stack (RTRBM-style, concatenation variant).
+
+    At each time step t and layer i:
+        visible_i[t] = concat(h_i[t-1], input_i[t])
+        h_i[t]       = entity_i.forward(visible_i[t])
+
+    where input_0[t] = x_t (sensory input) and input_{i+1}[t] = h_i[t].
+    Weights are shared across time — the same Entity processes every time step.
+    """
+
+    def __init__(self, name: str, work_dir: str = '.'):
+        Stack.__init__(self, StackType.Rnn, name, work_dir)
+        self._h: list[Mat] | None = None
+
+    # ── Convenience factory ───────────────────────────────────────────────────
+
+    @staticmethod
+    def make_layer(name: str, sensory_size: int, h_size: int,
+                   entity_params: EntityParams, training_params: TrainingParams) -> Layer:
+        """Create a Layer suitable for StackRnn.
+
+        The entity's visible size = sensory_size + h_size.
+        sensory_size is the raw input width; h_size is the recurrent state size.
+        """
+        return Layer(name, (1, sensory_size, h_size, h_size), entity_params, training_params)
+
+    # ── Derived sizes ─────────────────────────────────────────────────────────
+
+    def h_size(self, layer_idx: int = 0) -> int:
+        return self.from_index(layer_idx).entity.shape[1]
+
+    def sensory_size(self, layer_idx: int = 0) -> int:
+        e = self.from_index(layer_idx).entity
+        return e.shape[0] - e.shape[1]
+
+    # ── Hidden-state management ───────────────────────────────────────────────
+
+    def reset(self, batch_size: int = 1):
+        """Zero all hidden states (call before processing a new sequence)."""
+        self._h = [np.zeros((batch_size, self.h_size(i))) for i in range(self.num_layers())]
+
+    # ── Inference ─────────────────────────────────────────────────────────────
+
+    def step(self, x: Mat) -> Mat:
+        """One time step forward through all layers.
+
+        x: (batch_size, sensory_size) or (sensory_size,)
+        Returns the top-layer hidden state h.
+        """
+        if x.ndim == 1:
+            x = x[None, :]
+        batch_size = x.shape[0]
+        if self._h is None or self._h[0].shape[0] != batch_size:
+            self.reset(batch_size)
+
+        x_in = _to_gpu(x)
+        for i, layer in enumerate(self.layers):
+            visible = np.concatenate([self._h[i], x_in], axis=1)
+            h_new = layer.entity.forward(visible)
+            self._h[i] = h_new
+            x_in = h_new
+
+        return self._h[-1]
+
+    def reconstruct(self, h: Mat, layer_idx: int = -1) -> Mat:
+        """Decode h → visible, returning only the sensory portion.
+
+        The visible layer is [context | sensory]; this method strips context,
+        returning only the sensory reconstruction.
+        h: (batch_size, h_size)
+        """
+        if layer_idx < 0:
+            layer_idx = self.num_layers() + layer_idx
+        entity = self.from_index(layer_idx).entity
+        visible = entity.reconstruct(h)
+        return visible[:, entity.shape[1]:]
+
+    # ── Training ──────────────────────────────────────────────────────────────
+
+    def train(self, sequences: Mat, status: Status = None):
+        """Greedy layer-wise CD training over sequences.
+
+        sequences: (T, sensory_size)           — single sequence
+                   (num_seq, T, sensory_size)  — batch of sequences
+        """
+        if status is None:
+            status = Status()
+
+        seqs = sequences if sequences.ndim == 3 else sequences[None, :]
+        num_seq, T, _ = seqs.shape
+
+        for layer_idx, layer in enumerate(self.layers):
+            entity = layer.entity
+            print(f"Train layer {layer_idx} ({entity.name}) for {entity.training_params.num_epochs} epochs")
+
+            if entity.enable_training and entity.training_params is not None:
+                self._train_layer(entity, seqs, num_seq, T, status)
+
+            if layer_idx < self.num_layers() - 1:
+                seqs = self._pass_through(entity, seqs, num_seq, T)
+
+    def _train_layer(self, entity: Entity, seqs: Mat,
+                     num_seq: int, T: int, status: Status):
+        cd_func = _CD_FUNC[entity.type]
+        params = entity.training_params
+        h_sz = entity.shape[1]
+
+        d_progress = 100.0 / params.num_epochs
+        progress = 0.0
+        keep_running = True
+
+        entity.grad_zero()
+        status.on_change(entity)
+
+        for epoch in range(params.num_epochs):
+            if not keep_running:
+                break
+
+            h = np.zeros((num_seq, h_sz))
+            err_total = 0.0
+
+            for t in range(T):
+                x_t = _to_gpu(seqs[:, t, :])
+                visible = np.concatenate([h, x_t], axis=1)
+
+                dwhv, dbv, dbh = cd_func(entity, visible)
+                grad = entity.grad_compute(dbv, dbh, dwhv)
+                entity.state_adjust(grad, 1.0 / num_seq)
+
+                h = entity.forward(visible)
+                err_total += rms_error_accu(visible - entity.reconstruct(h))
+
+            progress += d_progress
+            if status.want_report(round(progress)):
+                if not status.on_change(entity, {
+                    "progress": {"value": round(progress), "unit": "%"},
+                    "err_rms":  {"value": err_total / T,   "unit": ""},
+                }):
+                    keep_running = False
+                    break
+
+        # Final report using one clean forward pass
+        h = np.zeros((num_seq, h_sz))
+        err_total = 0.0
+        for t in range(T):
+            x_t = _to_gpu(seqs[:, t, :])
+            visible = np.concatenate([h, x_t], axis=1)
+            h = entity.forward(visible)
+            err_total += rms_error_accu(visible - entity.reconstruct(h))
+        status.on_change(entity, {
+            "progress":      {"value": 100,           "unit": "%"},
+            "err_rms_total": {"value": err_total / T, "unit": ""},
+        })
+
+    def _pass_through(self, entity: Entity, seqs: Mat,
+                      num_seq: int, T: int) -> _np_cpu.ndarray:
+        """Run sequences through entity; return hidden outputs on CPU."""
+        h_sz = entity.shape[1]
+        outputs = _np_cpu.zeros((num_seq, T, h_sz), dtype=_np_cpu.float64)
+        h = np.zeros((num_seq, h_sz))
+        for t in range(T):
+            x_t = _to_gpu(seqs[:, t, :])
+            visible = np.concatenate([h, x_t], axis=1)
+            h = entity.forward(visible)
+            outputs[:, t, :] = convert(h)
+        return outputs

src/tests/test_rnn.py

@@ -0,0 +1,112 @@
+"""Test for StackRnn: recurrent RBM with concatenated [h_{t-1} | x_t] visible layer.
+
+Uses a repeating binary pattern as a minimal synthetic sequence so the model
+has something learnable to compress and predict.
+"""
+import numpy as _np_cpu
+from rbm.stack_rnn import StackRnn
+from rbm.matrix import np, convert
+from rbm.entity import EntityParams, TrainingParams
+from rbm.status import Status
+
+SENSORY_SIZE = 8
+H_SIZE = 4
+T = 16           # sequence length
+NUM_SEQ = 10     # sequences in the batch
+WORK_DIR = "../../results"
+
+
+def _make_sequences() -> _np_cpu.ndarray:
+    """Binary sequences: each row is one sequence of T frames."""
+    rng = _np_cpu.random.RandomState(0)
+    base = (rng.rand(NUM_SEQ, SENSORY_SIZE) > 0.5).astype(_np_cpu.float64)
+    seqs = _np_cpu.stack([base] * T, axis=1)   # (NUM_SEQ, T, SENSORY_SIZE)
+    return seqs
+
+
+def test_rnn_single_layer():
+    seqs = _make_sequences()
+
+    rnn = StackRnn("test_rnn", WORK_DIR)
+    layer = StackRnn.make_layer(
+        "layer0", SENSORY_SIZE, H_SIZE,
+        EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False),
+        TrainingParams(learning_rate=0.05, momentum=0.5, num_epochs=20,
+                       mini_batch_size=0, do_rao_blackwell=True),
+    )
+    rnn.append(layer)
+    rnn.state_init(0.01)
+
+    # Confirm entity shape
+    assert rnn.sensory_size() == SENSORY_SIZE, "sensory_size mismatch"
+    assert rnn.h_size() == H_SIZE, "h_size mismatch"
+
+    rnn.train(seqs)
+
+    # Inference: step through one sequence
+    rnn.reset(batch_size=1)
+    seq0 = seqs[0]                             # (T, SENSORY_SIZE) numpy
+    for t in range(T):
+        x_t = np.array(seq0[t][None, :])      # (1, SENSORY_SIZE) on device
+        h = rnn.step(x_t)
+        assert h.shape == (1, H_SIZE), f"step output shape wrong at t={t}"
+
+    # Reconstruction from final hidden state
+    recon = rnn.reconstruct(h)
+    assert recon.shape == (1, SENSORY_SIZE), "reconstruct shape wrong"
+
+    print(f"Original x_T : {convert(np.array(seq0[-1][None, :]))}")
+    print(f"Reconstructed: {convert(recon)}")
+    print("test_rnn_single_layer: [passed]")
+
+
+def test_rnn_two_layers():
+    """Two-layer recurrent stack: layer 1 receives h_0 as its sensory input."""
+    H_SIZE_0, H_SIZE_1 = 6, 3
+    seqs = _make_sequences()
+
+    rnn = StackRnn("test_rnn2", WORK_DIR)
+    rnn.append(StackRnn.make_layer(
+        "layer0", SENSORY_SIZE, H_SIZE_0,
+        EntityParams(),
+        TrainingParams(learning_rate=0.05, num_epochs=10),
+    ))
+    rnn.append(StackRnn.make_layer(
+        "layer1", H_SIZE_0, H_SIZE_1,
+        EntityParams(),
+        TrainingParams(learning_rate=0.05, num_epochs=10),
+    ))
+    rnn.state_init(0.01)
+
+    rnn.train(seqs)
+
+    rnn.reset(batch_size=1)
+    for t in range(T):
+        x_t = np.array(seqs[0, t][None, :])
+        h = rnn.step(x_t)
+
+    assert h.shape == (1, H_SIZE_1)
+    print("test_rnn_two_layers: [passed]")
+
+
+def test_rnn_save_load():
+    seqs = _make_sequences()
+
+    rnn = StackRnn("test_rnn_sl", WORK_DIR)
+    rnn.append(StackRnn.make_layer(
+        "layer0", SENSORY_SIZE, H_SIZE,
+        EntityParams(),
+        TrainingParams(num_epochs=5),
+    ))
+    rnn.state_init(0.01)
+    rnn.train(seqs)
+    rnn.state_save()
+    rnn.state_load()
+    print("test_rnn_save_load: [passed]")
+
+
+if __name__ == "__main__":
+    test_rnn_single_layer()
+    test_rnn_two_layers()
+    test_rnn_save_load()
+    print("All RNN tests passed.")