Add network plotting functionalities (#44)

.github/workflows/linting.yml

@@ -29,7 +29,7 @@ jobs:
     - name: Install dependencies
       run: |
         pip install -r requirements.txt
-        pip install isort==5.10.1 flake8 black==22.12.0 pydocstyle==6.3.0 mypy==1.1.1
+        pip install isort==5.10.1 flake8 black==23.1.0 pydocstyle==6.3.0 mypy==1.1.1
     - name: Run linting
       run: |
         flake8 ./pyhgf/

.github/workflows/test.yml

@@ -28,6 +28,7 @@ jobs:
           ${{ runner.os }}-pip-
     - name: Install dependencies
       run: |
+        sudo apt-get install graphviz
         pip install -r requirements-tests.txt
         pip install ipykernel coverage pytest pytest-cov
         python -m ipykernel install --user --name python3

.pre-commit-config.yaml

@@ -5,7 +5,7 @@ repos:
     - id: isort
       files: ^pyhgf/
 -   repo: https://github.com/ambv/black
-    rev: 22.12.0
+    rev: 23.1.0
     hooks:
     - id: black
       language_version: python3

docs/source/api.rst

@@ -82,6 +82,7 @@ observing new data.
 
    plot_trajectories
    plot_correlations
+   plot_network
 
 Response
 --------

docs/source/notebooks/1-Binary_HGF.md

@@ -5,7 +5,7 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.14.1
+    jupytext_version: 1.14.5
 kernelspec:
   display_name: Python 3 (ipykernel)
   language: python
@@ -68,18 +68,15 @@ two_levels_hgf = HGF(
     model_type="binary",
     initial_mu={"1": .0, "2": .5},
     initial_pi={"1": .0, "2": 1e4},
-    omega={"1": None, "2": -3.0},
-    rho={"1": None, "2": 0.0},
-    kappas={"1": None},
-    eta0=0.0,
-    eta1=1.0,
-    pihat = jnp.inf,
+    omega={"2": -3.0},
 )
 ```
 
-This function create an instance of a HGF model automatically parametrized for a 2-levels binary structure, so we do not have to worry about creating the nodes structure ourself. This class also embed function to add new observations and plots results that we are going to use below.
+This function create an instance of a HGF model automatically parametrized for a 2-levels binary structure, so we do not have to worry about creating the nodes structure ourself. This class also embed function to add new observations and plots results that we are going to use below. We can have a look at the node structure itself using the {ref}`pyhgf.plots.plot_network` function. This function will automatically dray the provided node structure using [Graphviz](https://github.com/xflr6/graphviz).
 
-+++
+```{code-cell} ipython3
+two_levels_hgf.plot_network()
+```
 
 #### Add data
 
@@ -132,6 +129,12 @@ three_levels_hgf = HGF(
 )
 ```
 
+The node structure now includes a volatility parent at the third level.
+
+```{code-cell} ipython3
+three_levels_hgf.plot_network()
+```
+
 #### Add data
 
 ```{code-cell} ipython3

docs/source/notebooks/2-Continuous_HGF.md

@@ -5,7 +5,7 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.14.1
+    jupytext_version: 1.14.5
 kernelspec:
   display_name: Python 3 (ipykernel)
   language: python
@@ -67,8 +67,6 @@ The response function used is the [sum of the Guassian surprise](pyhgf.response.
 ```
 
 ```{code-cell} ipython3
-:tags: []
-
 two_levels_continuous_hgf = HGF(
     n_levels=2,
     model_type="continuous",
@@ -79,9 +77,11 @@ two_levels_continuous_hgf = HGF(
     kappas={"1": 1.0})
 ```
 
-This function create an instance of a HGF model automatically parametrized for a 2-levels continuous structure, so we do not have to worry about creating the nodes structure ourself. This class also embed function to add new observations and plots results that we are going to use below.
+This function create an instance of a HGF model automatically parametrized for a 2-levels continuous structure, so we do not have to worry about creating the nodes structure ourself. This class also embed function to add new observations and plots results that we are going to use below. We can have a look at the node structure itself using the {ref}`pyhgf.plots.plot_network` function. This function will automatically dray the provided node structure using [Graphviz](https://github.com/xflr6/graphviz).
 
-+++
+```{code-cell} ipython3
+two_levels_continuous_hgf.plot_network()
+```
 
 #### Add data
 
@@ -138,6 +138,12 @@ three_levels_continuous_hgf = HGF(
     kappas={"1": 1.0, "2": 1.0})
 ```
 
+The node structure now includes a volatility parent at the third level.
+
+```{code-cell} ipython3
+three_levels_continuous_hgf.plot_network()
+```
+
 #### Add data
 
 ```{code-cell} ipython3
@@ -388,7 +394,7 @@ hgf_mcmc = HGF(
 ```
 
 ```{code-cell} ipython3
-hgf_mcmc.plot_trajectories(ci=False)
+hgf_mcmc.plot_trajectories(ci=True)
 ```
 
 ```{code-cell} ipython3

pyhgf/continuous.py

@@ -73,12 +73,10 @@ def continuous_node_update(
     # Update the continuous value parents #
     #######################################
     if value_parents_idx is not None:
-
         vape = node_parameters["mu"] - node_parameters["muhat"]
         psis = node_parameters["psis"]
 
         for va_pa_idx, psi in zip(value_parents_idx, psis):
-
             # unpack the current parent's parameters with value and volatility parents
             va_pa_node_parameters = parameters_structure[va_pa_idx]
             va_pa_value_parents_idx = node_structure[va_pa_idx].value_parents
@@ -125,7 +123,6 @@ def continuous_node_update(
     # Update volatility parents #
     #############################
     if volatility_parents_idx is not None:
-
         nu = node_parameters["nu"]
         kappas = node_parameters["kappas"]
         vope = (
@@ -134,7 +131,6 @@ def continuous_node_update(
         ) * node_parameters["pihat"] - 1
 
         for vo_pa_idx, kappa in zip(volatility_parents_idx, kappas):
-
             # unpack the current parent's parameters with value and volatility parents
             vo_pa_node_parameters = parameters_structure[vo_pa_idx]
             vo_pa_value_parents_idx = node_structure[vo_pa_idx].value_parents
@@ -248,7 +244,6 @@ def continuous_input_update(
     # Update value parents #
     ########################
     if value_parents_idx is not None:
-
         # unpack the current parent's parameters with value and volatility parents
         va_pa_node_parameters = parameters_structure[value_parents_idx[0]]
         va_pa_value_parents_idx = node_structure[value_parents_idx[0]].value_parents
@@ -306,7 +301,6 @@ def continuous_input_update(
     # Update volatility parents #
     #############################
     if volatility_parents_idx is not None:
-
         # unpack the current parent's parameters with value and volatility parents
         vo_pa_node_parameters = parameters_structure[volatility_parents_idx[0]]
         vo_pa_value_parents_idx = node_structure[

pyhgf/distribution.py

@@ -179,7 +179,6 @@ def hgf_logp(
 
     # Fitting n HGF models to the n datasets
     for i in range(n):
-
         # Format HGF parameters
         initial_mu: Dict = {
             "1": _mu_1[i],

pyhgf/model.py

@@ -16,10 +16,10 @@
     continuous_node_update,
     gaussian_surprise,
 )
-from pyhgf.plots import plot_correlations, plot_trajectories
+from pyhgf.plots import plot_correlations, plot_network, plot_trajectories
 from pyhgf.response import total_binary_surprise, total_gaussian_surprise
 from pyhgf.structure import loop_inputs
-from pyhgf.typing import Indexes, NodeStructure
+from pyhgf.typing import Indexes, InputIndexes, NodeStructure
 
 
 class HGF(object):
@@ -30,6 +30,9 @@ class HGF(object):
 
     Attributes
     ----------
+    input_nodes_idx :
+        Indexes of the input nodes. Defaults to `(0,)` if the network only has one input
+        node.
     model_type :
         The model implemented (can be `"continuous"`, `"binary"` or `"custom"`).
     n_levels :
@@ -38,7 +41,8 @@ class HGF(object):
     node_structure :
         A tuple of :py:class:`pyhgf.typing.Indexes` representing the nodes hierarchy.
     node_trajectories :
-        The node structure updated at each new observation.
+        The parameter structure that incluse the consecutive updates at each new
+        observation.
     parameters_structure :
         The structure of nodes' parameters. Each parameter is a dictionary with the
         following parameters: `"pihat", "pi", "muhat", "mu", "nu", "psis", "omega"` for
@@ -140,6 +144,7 @@ def __init__(
         self.node_structure: NodeStructure
         self.node_trajectories: Dict
         self.parameters_structure: Dict
+        self.input_nodes_idx: Tuple[InputIndexes, ...]
 
         if model_type in ["continuous", "binary"]:
             if self.verbose:
@@ -166,8 +171,8 @@ def __init__(
                 value_coupling=1.0,
                 mu=initial_mu["1"],
                 pi=initial_pi["1"],
-                omega=omega["1"],
-                rho=rho["1"],
+                omega=omega["1"] if self.model_type != "binary" else np.nan,
+                rho=rho["1"] if self.model_type != "binary" else np.nan,
             )
 
             #########
@@ -258,6 +263,10 @@ def plot_correlations(self):
         """Plot the heatmap of cross-trajectories correlation."""
         return plot_correlations(hgf=self)
 
+    def plot_network(self):
+        """Visualization of node network using GraphViz."""
+        return plot_network(hgf=self)
+
     def surprise(
         self,
         response_function: Optional[Callable] = None,
@@ -333,12 +342,31 @@ def to_pandas(self) -> pd.DataFrame:
     def add_input_node(
         self,
         kind: str,
+        input_idx: int = 0,
         omega_input: Union[float, np.ndarray, ArrayLike] = log(1e-4),
         pihat: Union[float, np.ndarray, ArrayLike] = jnp.inf,
         eta0: Union[float, np.ndarray, ArrayLike] = 0.0,
         eta1: Union[float, np.ndarray, ArrayLike] = 1.0,
     ):
-        """Create an input node."""
+        """Create an input node.
+
+        Parameters
+        ----------
+        kind :
+            The kind of input that should be created (can be `"continuous"` or
+            `"binary"`).
+        input_idx :
+            The index of the new input (defaults to `0`).
+        omega_input :
+            The input precision (only relevant if `kind="continuous"`).
+        pihat :
+            The input precision (only relevant if `kind="binary"`).
+        eta0 :
+            The lower bound of the binary process (only relevant if `kind="binary"`).
+        eta1 :
+            The lower bound of the binary process (only relevant if `kind="binary"`).
+
+        """
         if kind == "continuous":
             input_node_parameters = {
                 "kappas": None,
@@ -356,8 +384,16 @@ def add_input_node(
                 "time_step": jnp.nan,
                 "value": jnp.nan,
             }
-        self.parameters_structure = {0: input_node_parameters}
-        self.node_structure = (Indexes(None, None),)
+        if input_idx == 0:
+            # this is the first node, create the node structure
+            self.parameters_structure = {input_idx: input_node_parameters}
+            self.node_structure = (Indexes(None, None),)
+            self.input_nodes_idx = (InputIndexes(input_idx, kind),)
+        else:
+            # update the node structure
+            self.parameters_structure[input_idx] = input_node_parameters
+            self.node_structure += (Indexes(None, None),)
+            self.input_nodes_idx += (InputIndexes(input_idx, kind),)
         return self
 
     def add_value_parent(
@@ -458,7 +494,7 @@ def add_value_parent(
     def add_volatility_parent(
         self,
         children_idxs: List,
-        volatility_coupling: Union[float, np.ndarray, ArrayLike],
+        volatility_coupling: Union[float, np.ndarray, ArrayLike] = 1.0,
         mu: Union[float, np.ndarray, ArrayLike] = 0.0,
         mu_hat: Union[float, np.ndarray, ArrayLike] = jnp.nan,
         pi: Union[float, np.ndarray, ArrayLike] = 1.0,