From 00993e29f1a2e0a5f0b008c490c6a643507f4d42 Mon Sep 17 00:00:00 2001 From: Jens Ahrensfeld Date: Mon, 1 Jun 2026 16:43:33 +0200 Subject: [PATCH] [rnn_helper] add shift functions, ch2idx/idx2ch, and tests Co-Authored-By: Claude Sonnet 4.6 --- src/rbm/rnn_helper.py | 41 ++++++++++++++++ src/tests/test_rnn_helper.py | 91 ++++++++++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+) create mode 100644 src/rbm/rnn_helper.py create mode 100644 src/tests/test_rnn_helper.py diff --git a/src/rbm/rnn_helper.py b/src/rbm/rnn_helper.py new file mode 100644 index 0000000..e1dd0fc --- /dev/null +++ b/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 + diff --git a/src/tests/test_rnn_helper.py b/src/tests/test_rnn_helper.py new file mode 100644 index 0000000..69bc347 --- /dev/null +++ b/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.") \ No newline at end of file