[rnn_helper] add shift functions, ch2idx/idx2ch, and tests

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

src/rbm/rnn_helper.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+VOCAB = ' .!?ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'
+_CH2IDX = {ch: i for i, ch in enumerate(VOCAB)}
+
+def ch2idx(ch: str) -> int:
+	return _CH2IDX[ch]
+
+def idx2ch(idx: int) -> str:
+	return VOCAB[idx]
+
+
+def concat(v: np.ndarray, c: np.ndarray) -> np.ndarray:
+	return np.concatenate([v, c])
+
+def split(vc: np.ndarray, v_size: np.ndarray|int) -> tuple[np.ndarray, np.ndarray]:
+	n = len(v_size) if isinstance(v_size, np.ndarray) else v_size
+	return vc[:n], vc[n:]
+
+def shift_right(m: np.ndarray) -> np.ndarray:
+	return np.hstack([np.zeros((m.shape[0], 1)), m[:, :-1]])
+
+def shift_left(m: np.ndarray) -> np.ndarray:
+	return np.hstack([m[:, 1:], np.zeros((m.shape[0], 1))])
+
+def shift_up(m: np.ndarray) -> np.ndarray:
+	return np.vstack([m[1:, :], np.zeros((1, m.shape[1]))])
+
+def shift_down(m: np.ndarray) -> np.ndarray:
+	return np.vstack([np.zeros((1, m.shape[1])), m[:-1, :]])
+
+def clamp_one_hot(src_dst: np.ndarray) -> np.ndarray:
+	index = np.argmax(src_dst)
+	result = np.zeros_like(src_dst)
+	result[index] = 1
+	return result
+
+class Rnn:
+	def __init__(self):
+		pass
+

src/tests/test_rnn_helper.py

@@ -0,0 +1,91 @@
+import numpy as np
+import pytest
+from rbm.rnn_helper import shift_right, shift_left, shift_up, shift_down, ch2idx, idx2ch, VOCAB
+
+M = np.array([
+    [1, 2, 3],
+    [4, 5, 6],
+], dtype=float)
+
+
+def test_shift_right():
+    result = shift_right(M)
+    expected = np.array([
+        [0, 1, 2],
+        [0, 4, 5],
+    ], dtype=float)
+    assert np.array_equal(result, expected)
+    assert result.shape == M.shape
+
+
+def test_shift_left():
+    result = shift_left(M)
+    expected = np.array([
+        [2, 3, 0],
+        [5, 6, 0],
+    ], dtype=float)
+    assert np.array_equal(result, expected)
+    assert result.shape == M.shape
+
+
+def test_shift_up():
+    result = shift_up(M)
+    expected = np.array([
+        [4, 5, 6],
+        [0, 0, 0],
+    ], dtype=float)
+    assert np.array_equal(result, expected)
+    assert result.shape == M.shape
+
+
+def test_shift_down():
+    result = shift_down(M)
+    expected = np.array([
+        [0, 0, 0],
+        [1, 2, 3],
+    ], dtype=float)
+    assert np.array_equal(result, expected)
+    assert result.shape == M.shape
+
+
+def test_ch2idx_known():
+    assert ch2idx(' ') == 0
+    assert ch2idx('A') == 4
+    assert ch2idx('Z') == 29
+    assert ch2idx('0') == 30
+    assert ch2idx('9') == 39
+
+
+def test_idx2ch_known():
+    assert idx2ch(0) == ' '
+    assert idx2ch(4) == 'A'
+    assert idx2ch(29) == 'Z'
+    assert idx2ch(30) == '0'
+    assert idx2ch(39) == '9'
+
+
+def test_ch2idx_invalid():
+    with pytest.raises(KeyError):
+        ch2idx('a')
+
+
+def test_idx2ch_invalid():
+    with pytest.raises(IndexError):
+        idx2ch(len(VOCAB))
+
+
+def test_ch2idx_idx2ch_roundtrip():
+    for i, ch in enumerate(VOCAB):
+        assert ch2idx(ch) == i
+        assert idx2ch(i) == ch
+
+
+if __name__ == "__main__":
+    test_shift_right()
+    test_shift_left()
+    test_shift_up()
+    test_shift_down()
+    test_ch2idx_known()
+    test_idx2ch_known()
+    test_ch2idx_idx2ch_roundtrip()
+    print("All tests passed.")