refactor!: merge lossless source coding encoders/decoders into code c…

…lasses

- Removed `FixedToVariableEncoder`, `FixedToVariableDecoder`, `VariableToFixedEncoder`, `VariableToFixedDecoder`. Instead of `encoder = FixedToVariableEncoder(code); output = encoder(input)`, use `output = code.encode(input)`.

site/toc.yaml

@@ -74,12 +74,8 @@ Source coding:
   Lossless coding:
     - FixedToVariableCode
     - HuffmanCode
-    - FixedToVariableEncoder
-    - FixedToVariableDecoder
     - VariableToFixedCode
     - TunstallCode
-    - VariableToFixedEncoder
-    - VariableToFixedDecoder
   Quantization:
     - ScalarQuantizer
     - LloydMaxQuantizer

src/komm/_lossless_coding/FixedToVariableCode.py

@@ -6,7 +6,7 @@
 from typing_extensions import Self
 
 from .._util.information_theory import PMF
-from .util import Word, is_prefix_free, is_uniquely_decodable
+from .util import Word, is_prefix_free, is_uniquely_decodable, parse_prefix_free
 
 
 @frozen
@@ -202,3 +202,48 @@ def rate(self, pmf: npt.ArrayLike) -> float:
         probabilities = [np.prod(ps) for ps in it.product(pmf, repeat=k)]
         lengths = [len(word) for word in self.codewords]
         return np.dot(lengths, probabilities) / k
+
+    def encode(self, source_symbols: npt.ArrayLike) -> npt.NDArray[np.integer]:
+        r"""
+        Encodes a sequence of source symbols using the code.
+
+        Parameters:
+            source_symbols: The sequence of symbols to be encoded. Must be a 1D-array with elements in $[0:S)$, where $S$ is the source cardinality of the code.
+
+        Returns:
+            The sequence of encoded symbols. It is a 1D-array with elements in $[0:T)$, where $T$ is the target cardinality of the code.
+
+        Examples:
+            >>> code = komm.FixedToVariableCode.from_codewords(3, [(0,), (1,0), (1,1)])
+            >>> code.encode([1, 0, 1, 0, 2, 0])
+            array([1, 0, 0, 1, 0, 0, 1, 1, 0])
+        """
+        source_symbols = np.asarray(source_symbols)
+        k, enc = self.source_block_size, self.enc_mapping
+        return np.concatenate([enc[tuple(s)] for s in source_symbols.reshape(-1, k)])
+
+    def decode(self, target_symbols: npt.ArrayLike) -> npt.NDArray[np.integer]:
+        r"""
+        Decodes a sequence of target symbols using the code. Only works if the code is prefix-free.
+
+        Parameters:
+            target_symbols: The sequence of symbols to be decoded. Must be a 1D-array with elements in $[0:T)$, where $T$ is the target cardinality of the code.
+
+        Returns:
+            output: The sequence of decoded symbols. It is a 1D-array with elements in $[0:S)$, where $S$ is the source cardinality of the code.
+
+        Examples:
+            >>> code = komm.FixedToVariableCode.from_codewords(3, [(0,), (1,0), (1,1)])
+            >>> code.decode([1, 0, 0, 1, 0, 0, 1, 1, 0])
+            array([1, 0, 1, 0, 2, 0])
+
+            >>> code = komm.FixedToVariableCode.from_codewords(2, [(0,), (1,0), (1,1), (1,1,0)])
+            >>> code.decode([1, 0, 0, 1, 0, 0, 1, 1, 0])
+            Traceback (most recent call last):
+            ...
+            ValueError: code is not prefix-free
+        """
+        if not self.is_prefix_free():
+            raise ValueError("code is not prefix-free")
+        target_symbols = np.asarray(target_symbols)
+        return parse_prefix_free(target_symbols, self.inv_enc_mapping)

src/komm/_lossless_coding/VariableToFixedCode.py

@@ -6,7 +6,7 @@
 from typing_extensions import Self
 
 from .._util.information_theory import PMF
-from .util import Word, is_prefix_free
+from .util import Word, is_prefix_free, parse_prefix_free
 
 
 @frozen
@@ -178,3 +178,48 @@ def rate(self, pmf: npt.ArrayLike) -> float:
         probabilities = [np.prod([pmf[x] for x in word]) for word in self.sourcewords]
         lengths = [len(word) for word in self.sourcewords]
         return self.target_block_size / np.dot(lengths, probabilities)
+
+    def encode(self, source_symbols: npt.ArrayLike) -> npt.NDArray[np.integer]:
+        r"""
+        Encodes a sequence of source symbols using the code.
+
+        Parameters:
+            source_symbols: The sequence of symbols to be encoded. Must be a 1D-array with elements in $[0:S)$, where $S$ is the source cardinality of the code.
+
+        Returns:
+            The sequence of encoded symbols. It is a 1D-array with elements in $[0:T)$, where $T$ is the target cardinality of the code.
+
+        Examples:
+            >>> code = komm.VariableToFixedCode.from_sourcewords(2, [(0,0,0), (0,0,1), (0,1), (1,)])
+            >>> code.encode([0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0])
+            array([0, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 0])
+
+            >>> code = komm.VariableToFixedCode.from_sourcewords(2, [(0,0,0), (0,0,1), (0,1), (0,)])
+            >>> code.encode([0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0])
+            Traceback (most recent call last):
+            ...
+            ValueError: code is not prefix-free
+        """
+        if not self.is_prefix_free():
+            raise ValueError("code is not prefix-free")
+        source_symbols = np.asarray(source_symbols)
+        return parse_prefix_free(source_symbols, self.inv_dec_mapping)
+
+    def decode(self, target_symbols: npt.ArrayLike) -> npt.NDArray[np.integer]:
+        r"""
+        Decodes a sequence of target symbols using the code.
+
+        Parameters:
+            target_symbols: The sequence of symbols to be decoded. Must be a 1D-array with elements in $[0:T)$, where $T$ is the target cardinality of the code.
+
+        Returns:
+            The sequence of decoded symbols. It is a 1D-array with elements in $[0:S)$, where $S$ is the source cardinality of the code.
+
+        Examples:
+            >>> code = komm.VariableToFixedCode.from_sourcewords(2, [(0,0,0), (0,0,1), (0,1), (1,)])
+            >>> code.decode([0, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 0])
+            array([0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0])
+        """
+        target_symbols = np.asarray(target_symbols)
+        n, dec = self.target_block_size, self.dec_mapping
+        return np.concatenate([dec[tuple(s)] for s in target_symbols.reshape(-1, n)])

src/komm/_lossless_coding/__init__.py

@@ -1,19 +1,11 @@
 from .FixedToVariableCode import FixedToVariableCode
-from .FixedToVariableDecoder import FixedToVariableDecoder
-from .FixedToVariableEncoder import FixedToVariableEncoder
 from .HuffmanCode import HuffmanCode
 from .TunstallCode import TunstallCode
 from .VariableToFixedCode import VariableToFixedCode
-from .VariableToFixedDecoder import VariableToFixedDecoder
-from .VariableToFixedEncoder import VariableToFixedEncoder
 
 __all__ = [
     "FixedToVariableCode",
-    "FixedToVariableDecoder",
-    "FixedToVariableEncoder",
     "HuffmanCode",
     "TunstallCode",
     "VariableToFixedCode",
-    "VariableToFixedDecoder",
-    "VariableToFixedEncoder",
 ]

tests/lossless_coding/test_fixed_to_variable_code.py

@@ -165,7 +165,7 @@ def test_non_injective_enc_mapping():
 def test_decoding_not_prefix_free(source_cardinality, codewords):
     code = komm.FixedToVariableCode.from_codewords(source_cardinality, codewords)
     with pytest.raises(ValueError):
-        komm.FixedToVariableDecoder(code)
+        code.decode([0, 0, 0, 0])
 
 
 @pytest.mark.parametrize(
@@ -228,11 +228,7 @@ def test_rate_invalid_pmf(pmf):
     ],
 )
 def test_encoding_decoding(code_parameters, x, y):
-    source_cardinality, target_cardinality, source_block_size, codewords = (
-        code_parameters.values()
-    )
+    source_cardinality, _, _, codewords = code_parameters.values()
     code = komm.FixedToVariableCode.from_codewords(source_cardinality, codewords)
-    encoder = komm.FixedToVariableEncoder(code)
-    decoder = komm.FixedToVariableDecoder(code)
-    assert np.array_equal(encoder(x), y)
-    assert np.array_equal(decoder(y), x)
+    assert np.array_equal(code.encode(x), y)
+    assert np.array_equal(code.decode(y), x)

tests/lossless_coding/test_variable_to_fixed_code.py

@@ -106,9 +106,7 @@ def test_encoding_decoding():
     code = komm.VariableToFixedCode.from_sourcewords(
         2, [(0, 0, 0), (0, 0, 1), (0, 1), (1,)]
     )
-    encoder = komm.VariableToFixedEncoder(code)
-    decoder = komm.VariableToFixedDecoder(code)
     x = [0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0]
     y = [0, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 0]
-    assert np.array_equal(encoder(x), y)
-    assert np.array_equal(decoder(y), x)
+    assert np.array_equal(code.encode(x), y)
+    assert np.array_equal(code.decode(y), x)