From b0a0ca0485403e603630be97e37f76f7ce4f332f Mon Sep 17 00:00:00 2001
From: michaeldeistler <michael.deistler95@gmail.com>
Date: Sat, 3 Feb 2024 17:19:01 +0100
Subject: [PATCH] A bit of documentation

---
 jaxley/modules/base.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/jaxley/modules/base.py b/jaxley/modules/base.py
index ce1d0258c..90d432d64 100644
--- a/jaxley/modules/base.py
+++ b/jaxley/modules/base.py
@@ -534,7 +534,7 @@ def _step_channels(
         channel_nodes: pd.DataFrame,
         params: Dict[str, jnp.ndarray],
     ):
-        """One step of integration of the channels."""
+        """One step of integration of the channels and of computing their current."""
         states = self._step_channels_state(
             states, delta_t, channels, channel_nodes, params
         )
@@ -551,6 +551,7 @@ def _step_channels_state(
         channel_nodes: pd.DataFrame,
         params: Dict[str, jnp.ndarray],
     ):
+        """One integration step of the channels."""
         voltages = states["voltages"]
 
         # Update states of the channels.
@@ -586,6 +587,10 @@ def _channel_currents(
         channel_nodes: pd.DataFrame,
         params: Dict[str, jnp.ndarray],
     ):
+        """Return the current through each channel.
+        
+        This is also updates `state` because the `state` also contains the current.
+        """
         voltages = states["voltages"]
 
         # Compute current through channels.