add more tests and improve code coverage (#273)

* improve tests and code coverage

* notebook

* more tests

docs/source/notebooks/0.3-Generalised_filtering.ipynb

@@ -347,30 +347,32 @@
     "tags": []
    },
    "source": [
-    "````{note} From sufficient statistics to distribution parameters and backwards\n",
+    ":::{note} From sufficient statistics to distribution parameters and backwards\n",
     ":class: dropdown\n",
     "\n",
     "When using a 1-dimensional Gaussian distribution, Setting $\\xi = [0, \\frac{1}{8}]$ is equivalent to a mean $\\mu = 0.0$ and a variance $\\sigma^2 = \\frac{1}{8}$. You can convert between distribution parameters and expected sufficient statistics using the distribution classes from PyHGF (when implemented):\n",
     "\n",
-    "```{code-cell} python\n",
+    "```python\n",
     "from pyhgf.math import Normal\n",
     "\n",
     "# from an observation to sufficient statistics\n",
     "Normal.sufficient_statistics_from_observations(x=1.5)\n",
     "```\n",
     "> Array([1.5 , 2.25], dtype=float32)\n",
-    "```{code-cell} python\n",
+    "\n",
+    "```python\n",
     "# from distribution parameters to sufficient statistics\n",
     "Normal.sufficient_statistics_from_parameters(mean=0.0, variance=4.0)\n",
     "```\n",
     "> Array([0., 4.], dtype=float32)\n",
-    "```{code-cell} python\n",
+    "\n",
+    "```python\n",
     "# from sufficient statistics to distribution parameters\n",
     "Normal.parameters_from_sufficient_statistics(xis=[0.0, 4.0])\n",
     "```\n",
     "> (0.0, 4.0)\n",
     "\n",
-    "````"
+    ":::"
    ]
   },
   {
@@ -417,7 +419,7 @@
        "</svg>\n"
       ],
       "text/plain": [
-       "<graphviz.sources.Source at 0x7f69845ec530>"
+       "<graphviz.sources.Source at 0x7fec53b15520>"
       ]
      },
      "execution_count": 7,
@@ -665,7 +667,7 @@
        "</svg>\n"
       ],
       "text/plain": [
-       "<graphviz.sources.Source at 0x7f6984591040>"
+       "<graphviz.sources.Source at 0x7fec4dea0950>"
       ]
      },
      "execution_count": 13,
@@ -1239,7 +1241,7 @@
        "</svg>\n"
       ],
       "text/plain": [
-       "<graphviz.sources.Source at 0x7f69784b16a0>"
+       "<graphviz.sources.Source at 0x7fec3c292780>"
       ]
      },
      "execution_count": 21,
@@ -1292,7 +1294,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Last updated: Tue Jan 14 2025\n",
+      "Last updated: Wed Jan 15 2025\n",
       "\n",
       "Python implementation: CPython\n",
       "Python version       : 3.12.3\n",
@@ -1302,13 +1304,13 @@
       "jax   : 0.4.31\n",
       "jaxlib: 0.4.31\n",
       "\n",
-      "numpy     : 1.26.0\n",
-      "sys       : 3.12.3 | packaged by conda-forge | (main, Apr 15 2024, 18:38:13) [GCC 12.3.0]\n",
-      "pyhgf     : 0.2.1.post4.dev0+d49aafe9\n",
+      "matplotlib: 3.10.0\n",
       "jax       : 0.4.31\n",
-      "IPython   : 8.31.0\n",
       "seaborn   : 0.13.2\n",
-      "matplotlib: 3.10.0\n",
+      "IPython   : 8.31.0\n",
+      "sys       : 3.12.3 | packaged by conda-forge | (main, Apr 15 2024, 18:38:13) [GCC 12.3.0]\n",
+      "pyhgf     : 0.2.1.post4.dev0+d49aafe9\n",
+      "numpy     : 1.26.0\n",
       "\n",
       "Watermark: 2.5.0\n",
       "\n"

pyhgf/plots.py

@@ -264,13 +264,7 @@ def plot_network(network: "Network") -> "Source":
 
         style = "filled" if idx in network.input_idxs else ""
 
-        if network.edges[idx].node_type == 0:
-            # binary state node
-            graphviz_structure.node(
-                f"x_{idx}", label=str(idx), shape="ellipse", style=style
-            )
-
-        elif network.edges[idx].node_type == 1:
+        if network.edges[idx].node_type == 1:
             # binary state node
             graphviz_structure.node(
                 f"x_{idx}", label=str(idx), shape="square", style=style

tests/test_model.py

@@ -2,6 +2,7 @@
 
 import jax.numpy as jnp
 import numpy as np
+from pytest import raises
 
 from pyhgf import load_data
 from pyhgf.model import HGF, Network
@@ -33,7 +34,11 @@ def test_network():
     custom_hgf.create_belief_propagation_fn(overwrite=False)
     custom_hgf.create_belief_propagation_fn(overwrite=True)
 
-    custom_hgf.input_data(input_data=np.ones((10, 2)))
+    custom_hgf.input_data(input_data=np.ones((10, 2)), observed=np.ones((10, 2)))
+
+    # expected error for invalid type
+    with raises(Exception):
+        custom_hgf.add_nodes(kind="error")
 
 
 def test_continuous_hgf():
@@ -155,6 +160,8 @@ def test_custom_sequence():
     update_branches = (update_sequence1, update_sequence2)
     branches_idx = np.random.binomial(n=1, p=0.5, size=len(u))
 
+    three_level_binary_hgf.inpu_dim = []
+    three_level_binary_hgf.scan_fn = None
     three_level_binary_hgf.input_custom_sequence(
         update_branches=update_branches,
         branches_idx=branches_idx,

tests/test_nodes/test_exponential_family.py

@@ -46,7 +46,10 @@ def test_multivariate_gaussian():
         + np.random.randn(N, 2) * 2
     )
 
-    # Python ---------------------------------------------------------------------------
+    # Python
+    # ----------------------------------------------------------------------------------
+
+    # generalised filtering
     bivariate_normal = (
         PyNetwork()
         .add_nodes(
@@ -64,3 +67,30 @@ def test_multivariate_gaussian():
             dtype="float32",
         ),
     ).all()
+
+    # hgf updates
+    bivariate_hgf = PyNetwork().add_nodes(
+        kind="ef-state",
+        learning="hgf-2",
+        distribution="multivariate-normal",
+        dimension=2,
+    )
+
+    # adapting prior parameter values to the sufficient statistics
+    # covariances statistics will have greater variability and amplitudes
+    for node_idx in [2, 5, 8, 11, 14]:
+        bivariate_hgf.attributes[node_idx]["tonic_volatility"] = -2.0
+    for node_idx in [1, 4, 7, 10, 13]:
+        bivariate_hgf.attributes[node_idx]["precision"] = 0.01
+    for node_idx in [9, 12, 15]:
+        bivariate_hgf.attributes[node_idx]["mean"] = 10.0
+
+    bivariate_hgf.input_data(input_data=spiral_data)
+
+    assert jnp.isclose(
+        bivariate_normal.node_trajectories[0]["xis"][-1],
+        jnp.array(
+            [3.4652710e01, -1.0609777e00, 1.2103647e03, -3.6398651e01, 3.3951855e00],
+            dtype="float32",
+        ),
+    ).all()

tests/test_plots.py

@@ -27,7 +27,7 @@ def test_plotting_functions():
     ).input_data(input_data=timeserie)
 
     # plot trajectories
-    two_level_continuous.plot_trajectories()
+    two_level_continuous.plot_trajectories(show_total_surprise=True)
 
     # plot correlations
     two_level_continuous.plot_correlations()
@@ -53,7 +53,7 @@ def test_plotting_functions():
     ).input_data(input_data=timeserie)
 
     # plot trajectories
-    three_level_continuous.plot_trajectories()
+    three_level_continuous.plot_trajectories(show_total_surprise=True)
 
     # plot correlations
     three_level_continuous.plot_correlations()
@@ -88,7 +88,7 @@ def test_plotting_functions():
     ).input_data(u)
 
     # plot trajectories
-    two_level_binary_hgf.plot_trajectories()
+    two_level_binary_hgf.plot_trajectories(show_total_surprise=True)
 
     # plot correlations
     two_level_binary_hgf.plot_correlations()
@@ -116,7 +116,7 @@ def test_plotting_functions():
     ).input_data(u)
 
     # plot trajectories
-    three_level_binary_hgf.plot_trajectories()
+    three_level_binary_hgf.plot_trajectories(show_total_surprise=True)
 
     # plot correlations
     three_level_binary_hgf.plot_correlations()

tests/test_utils.py

@@ -20,14 +20,24 @@ def test_imports():
 
 def test_add_edges():
     """Test the add_edges function."""
-    network = Network().add_nodes().add_nodes(n_nodes=3)
-    with raises(Exception):
-        network.add_edges(kind="error")
 
+    # add value coupling
+    network = Network().add_nodes(n_nodes=3)
+    network.add_edges(parent_idxs=1, children_idxs=0, coupling_strengths=1.0)
+    network.add_edges(parent_idxs=1, children_idxs=2, coupling_strengths=1.0)
+
+    # add volatility coupling
+    network = Network().add_nodes(n_nodes=3)
     network.add_edges(
-        kind="volatility", parent_idxs=2, children_idxs=0, coupling_strengths=1
+        kind="volatility", parent_idxs=1, children_idxs=0, coupling_strengths=1
     )
-    network.add_edges(parent_idxs=1, children_idxs=0, coupling_strengths=1.0)
+    network.add_edges(
+        kind="volatility", parent_idxs=1, children_idxs=2, coupling_strengths=1
+    )
+
+    # expected error for invalid type
+    with raises(Exception):
+        network.add_edges(kind="error")
 
 
 def test_find_branch():
@@ -118,13 +128,14 @@ def test_remove_node():
     """Test the remove_node function."""
     network = (
         Network()
-        .add_nodes(n_nodes=4)
+        .add_nodes(n_nodes=2)
+        .add_nodes(value_children=0, volatility_children=1)
+        .add_nodes(volatility_children=2)
         .add_nodes(value_children=2)
-        .add_nodes(value_children=3)
     )
 
     attributes, edges, _ = network.get_network()
-    new_attributes, new_edges = remove_node(attributes, edges, 1)
+    new_attributes, new_edges = remove_node(attributes, edges, 2)
 
-    assert len(new_attributes) == 6
-    assert len(new_edges) == 5
+    assert len(new_attributes) == 5
+    assert len(new_edges) == 4