Adds unit tests

pgmax/bp/bp_utils.py

@@ -18,8 +18,8 @@ def segment_max_opt(
 
     Args:
         data: Array of shape (a_len,) where a_len is an arbitrary integer.
-        segments_lengths: Array of shape (num_segments,), where num_segments <= a_len.
-            segments_lengths.sum() should yield a_len.
+        segments_lengths: Array of shape (num_segments,), where 0 < num_segments <= a_len.
+            segments_lengths.sum() should yield a_len, and all elements must be > 0.
         max_segment_length: The maximum value in segments_lengths.
 
     Returns:

pgmax/fg/graph.py

@@ -306,8 +306,7 @@ def message_passing_step(msgs, _):
         return msgs_after_bp
 
     def decode_map_states(self, msgs: jnp.ndarray) -> Dict[nodes.Variable, int]:
-        """Function to computes the output of MAP inference on input messages.
-
+        """Function to compute the output of MAP inference on input messages.
         The final states are computed based on evidence obtained from the self.get_evidence
         method as well as the internal wiring.
 

pyproject.toml

@@ -28,6 +28,7 @@ Sphinx = "^3.0.0"
 pre-commit = "^2.13.0"
 jupytext = "^1.11.3"
 sphinx-rtd-theme = "^0.5.2"
+coverage = "^5.5"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+filterwarnings =
+    ignore::DeprecationWarning

tests/bp/test_bputils.py

@@ -0,0 +1,32 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import pgmax.bp.bp_utils as bp_utils
+
+
+def test_with_zero():
+    assert jnp.allclose(
+        bp_utils.segment_max_opt(jnp.zeros(1, dtype=int), jnp.ones(1, dtype=int), 1),
+        jnp.zeros(1),
+    )
+
+
+def test_with_ones():
+    assert jnp.allclose(
+        bp_utils.segment_max_opt(
+            jnp.ones(10, dtype=int), jnp.ones(1, dtype=int) * 10, 10
+        ),
+        jnp.ones(1),
+    )
+
+
+def test_common_case():
+    assert jnp.allclose(
+        bp_utils.segment_max_opt(
+            jax.device_put(np.array([1.7, 2.3, 4.3, 3.8, 9.2, 111.3])),
+            jax.device_put(np.array([2, 3, 1])),
+            3,
+        ),
+        jax.device_put(np.array([2.3, 9.2, 111.3])),
+    )

tests/fg/test_nodes.py

@@ -0,0 +1,14 @@
+import dataclasses
+
+import pytest
+
+import pgmax.fg.nodes as nodes
+
+
+def test_variable_frozen():
+    v = nodes.Variable(3)
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        v.num_states = 4
+
+
+# TODO: Write more tests for this file

tests/test_pgmax.py

@@ -6,9 +6,6 @@
 from numpy.random import default_rng
 from scipy.ndimage import gaussian_filter
 
-from pgmax.fg import graph, groups
-
-
 def test_e2e_sanity_check():
     # Helper function to easily generate a list of valid configurations for a given suppression diameter
     def create_valid_suppression_config_arr(suppression_diameter):
@@ -348,4 +345,4 @@ def create_valid_suppression_config_arr(suppression_diameter):
     final_msgs = fg.run_bp(1000, 0.5)
 
     # Test that the output messages are close to the true messages
-    assert jnp.allclose(final_msgs, true_final_msgs_output, atol=1e-06)
+    assert jnp.allclose(final_msgs, true_final_msgs_output, atol=1e-06)