Made all simulation submodules public and added some docstrings.

.pre-commit-config.yaml

@@ -9,9 +9,6 @@ repos:
     hooks:
     -   id: black
         language_version: python3
-
-    # Issue: pydocstyle (D) is somehow not checked in pre-commit.
-    # When running "ruff ." it is included though.
 -   repo: https://github.com/charliermarsh/ruff-pre-commit
     rev: 'v0.1.0'
     hooks:

src/lorenzpy/simulations/__init__.py

@@ -47,10 +47,10 @@
       https://github.com/mkdocstrings/mkdocstrings/issues/78
 """
 
-from ._autonomous_flows import Lorenz63, Lorenz96
-from ._discrete_maps import Logistic
-from ._driven_systems import SimplestDrivenChaotic
-from ._others import KuramotoSivashinsky, MackeyGlass
+from .autonomous_flows import Lorenz63, Lorenz96
+from .discrete_maps import Logistic
+from .driven_systems import SimplestDrivenChaotic
+from .others import KuramotoSivashinsky, MackeyGlass
 
 __all__ = [
     "Lorenz63",

src/lorenzpy/simulations/_autonomous_flows.py → src/lorenzpy/simulations/autonomous_flows.py

@@ -1,7 +1,7 @@
 """Autonomous flows."""
 import numpy as np
 
-from ._base import _BaseSimFlow
+from .base import _BaseSimFlow
 
 
 class Lorenz63(_BaseSimFlow):
@@ -16,21 +16,22 @@ class Lorenz63(_BaseSimFlow):
     """
 
     def __init__(self, sigma=10.0, rho=28.0, beta=8 / 3, dt=0.1, solver="rk4"):
-        """Initialize the Lorenz63.
+        """Initialize the Lorenz63 simulation object.
 
         Args:
-            sigma: a.
-            rho: b.
-            beta: c.
-            dt: d.
-            solver: e.
+            sigma: Sigma parameter of Lorenz63 equation.
+            rho: Rho parameter of Lorenz63 equation.
+            beta: beta parameter of Lorenz63 equation.
+            dt: Time step to simulate.
+            solver: The solver.
         """
         super().__init__(dt, solver)
         self.sigma = sigma
         self.rho = rho
         self.beta = beta
 
     def flow(self, x: np.ndarray) -> np.ndarray:
+        """Return the flow of Lorenz63 equation."""
         return np.array(
             [
                 self.sigma * (x[1] - x[0]),
@@ -40,6 +41,7 @@ def flow(self, x: np.ndarray) -> np.ndarray:
         )
 
     def get_default_starting_pnt(self) -> np.ndarray:
+        """Return default starting point of Lorenz63 system."""
         return np.array([0.0, -0.01, 9.0])
 
 
@@ -49,14 +51,20 @@ class Lorenz96(_BaseSimFlow):
     def __init__(
         self, sys_dim: int = 30, force: float = 8.0, dt: float = 0.05, solver="rk4"
     ) -> None:
+        """Initialize the Lorenz96 simulation object.
+
+        Args:
+            sys_dim: The dimension of the Lorenz96 system.
+            force: The force value.
+            dt: Time step to simulate.
+            solver: The solver.
+        """
         super().__init__(dt=dt, solver=solver)
         self.sys_dim = sys_dim
         self.force = force
 
-    def get_default_starting_pnt(self) -> np.ndarray:
-        return np.sin(np.arange(self.sys_dim))
-
     def flow(self, x: np.ndarray) -> np.ndarray:
+        """Return the flow of Lorenz96 equation."""
         system_dimension = x.shape[0]
         derivative = np.zeros(system_dimension)
         # Periodic Boundary Conditions for the 3 edge cases i=1,2,system_dimension
@@ -74,3 +82,7 @@ def flow(self, x: np.ndarray) -> np.ndarray:
 
         derivative = derivative + self.force
         return derivative
+
+    def get_default_starting_pnt(self) -> np.ndarray:
+        """Return default starting point of Lorenz96 system."""
+        return np.sin(np.arange(self.sys_dim))

src/lorenzpy/simulations/_base.py → src/lorenzpy/simulations/base.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 
-def _forward_euler(
+def forward_euler(
     f: Callable[[np.ndarray], np.ndarray], dt: float, x: np.ndarray
 ) -> np.ndarray:
     """Simulate one step for ODEs of the form dx/dt = f(x(t)) using the forward euler.
@@ -26,7 +26,7 @@ def _forward_euler(
     return x + dt * f(x)
 
 
-def _runge_kutta(
+def runge_kutta(
     f: Callable[[np.ndarray], np.ndarray], dt: float, x: np.ndarray
 ) -> np.ndarray:
     """Simulate one step for ODEs of the form dx/dt = f(x(t)) using Runge-Kutta.
@@ -48,7 +48,7 @@ def _runge_kutta(
     return next_step
 
 
-def _timestep_iterator(
+def timestep_iterator(
     f: Callable[[np.ndarray], np.ndarray], time_steps: int, starting_point: np.ndarray
 ) -> np.ndarray:
     """Iterate an iterator-function f: x(i+1) = f(x(i)) multiple times."""
@@ -142,7 +142,7 @@ def simulate(
             time_steps: Number of time steps t to simulate.
             starting_point: Starting point of the trajectory shape (sys_dim,).
                             If None, take the default starting point.
-            transient: TBD
+            transient: Discard the first <transient> steps in the output.
 
         Returns:
             Trajectory of shape (t, sys_dim).
@@ -151,7 +151,7 @@ def simulate(
         if starting_point is None:
             starting_point = self.get_default_starting_pnt()
 
-        return _timestep_iterator(self.iterate, time_steps + transient, starting_point)[
+        return timestep_iterator(self.iterate, time_steps + transient, starting_point)[
             transient:, :
         ]
 
@@ -185,9 +185,9 @@ def iterate(self, x: np.ndarray) -> np.ndarray:
         """
         if isinstance(self.solver, str):
             if self.solver == "rk4":
-                x_next = _runge_kutta(self.flow, self.dt, x)
+                x_next = runge_kutta(self.flow, self.dt, x)
             elif self.solver == "forward_euler":
-                x_next = _forward_euler(self.flow, self.dt, x)
+                x_next = forward_euler(self.flow, self.dt, x)
             else:
                 raise ValueError(f"Unsupported solver: {self.solver}")
         else:

src/lorenzpy/simulations/_discrete_maps.py → src/lorenzpy/simulations/discrete_maps.py

@@ -1,19 +1,28 @@
 """Discrete maps."""
 import numpy as np
 
-from ._base import _BaseSimIterate
+from .base import _BaseSimIterate
 
 
 class Logistic(_BaseSimIterate):
+    """Simulation class for the Logistic map."""
+
     def __init__(self, r: float = 4.0) -> None:
+        """Initialize the Logistic Map simulation object.
+
+        Args:
+            r: r parameter of the logistic map.
+        """
         self.r = r
 
     def iterate(self, x: np.ndarray) -> np.ndarray:
+        """Iterate the logistic map one step."""
         return np.array(
             [
                 self.r * x[0] * (1 - x[0]),
             ]
         )
 
     def get_default_starting_pnt(self) -> np.ndarray:
+        """Return default starting point of the Logistic map."""
         return np.array([0.1])

src/lorenzpy/simulations/_driven_systems.py → src/lorenzpy/simulations/driven_systems.py

@@ -1,16 +1,25 @@
 """Driven Systems."""
 import numpy as np
 
-from ._base import _BaseSimFlowDriven
+from .base import _BaseSimFlowDriven
 
 
 class SimplestDrivenChaotic(_BaseSimFlowDriven):
+    """Simulate the Simplest Driven Chaotic system from Sprott.
+
+    Taken from (Sprott, Julien Clinton, and Julien C. Sprott. Chaos and time-series
+    analysis. Vol. 69. Oxford: Oxford university press, 2003.)
+    """
+
     def __init__(self, omega: float = 1.88, dt: float = 0.1, solver="rk4") -> None:
+        """Initialize the SimplestDrivenChaotic simulation object."""
         super().__init__(dt, solver)
         self.omega = omega
 
     def flow(self, x: np.ndarray) -> np.ndarray:
+        """Return the flow ."""
         return np.array([x[1], -(x[0] ** 3) + np.sin(self.omega * x[2]), 1])
 
     def get_default_starting_pnt(self) -> np.ndarray:
+        """Return default starting point."""
         return np.array([0.0, 0.0, 0.0])

src/lorenzpy/simulations/_others.py → src/lorenzpy/simulations/others.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 
-from ._base import _BaseSim, _BaseSimIterate, _forward_euler, _runge_kutta
+from .base import _BaseSim, _BaseSimIterate, forward_euler, runge_kutta
 
 
 class KuramotoSivashinsky(_BaseSimIterate):
@@ -34,24 +34,21 @@ def __init__(
         eps: float = 0.0,
         dt: float = 0.1,
     ) -> None:
+        """Initialize the Kuramoto-Sivashinsky simulation object.
+
+        Args:
+            sys_dim: The dimension of the system.
+            sys_length: The physical length of the system.
+            eps: A parameter in front of the y_xx term.
+            dt: Time step to simulate.
+        """
         self.sys_dim = sys_dim
         self.sys_length = sys_length
         self.eps = eps
         self.dt = dt
 
         self._prepare()
 
-    def get_default_starting_pnt(self) -> np.ndarray:
-        x = (
-            self.sys_length
-            * np.transpose(np.conj(np.arange(1, self.sys_dim + 1)))
-            / self.sys_dim
-        )
-        return np.array(
-            np.cos(2 * np.pi * x / self.sys_length)
-            * (1 + np.sin(2 * np.pi * x / self.sys_length))
-        )
-
     def _prepare(self) -> None:
         """Calculate internal attributes.
 
@@ -117,6 +114,18 @@ def iterate(self, x: np.ndarray) -> np.ndarray:
         v = self.E * v + Nv * self.f1 + 2 * (Na + Nb) * self.f2 + Nc * self.f3
         return np.real(np.fft.ifft(v))
 
+    def get_default_starting_pnt(self) -> np.ndarray:
+        """Return default starting point of KS system."""
+        x = (
+            self.sys_length
+            * np.transpose(np.conj(np.arange(1, self.sys_dim + 1)))
+            / self.sys_dim
+        )
+        return np.array(
+            np.cos(2 * np.pi * x / self.sys_length)
+            * (1 + np.sin(2 * np.pi * x / self.sys_length))
+        )
+
 
 class MackeyGlass(_BaseSim):
     """Simulate the Mackey-Glass delay differential system.
@@ -140,6 +149,7 @@ def __init__(
         dt: float = 0.1,
         solver: str | Callable = "rk4",
     ) -> None:
+        """Initialize the Mackey-Glass simulation object."""
         self.a = a
         self.b = b
         self.c = c
@@ -151,6 +161,7 @@ def __init__(
         self.history_steps = int(self.tau / self.dt)
 
     def get_default_starting_pnt(self) -> np.ndarray:
+        """Return default starting point of MG system."""
         return np.array([0.9])
 
     def flow_mg(self, x: np.ndarray, x_past: np.ndarray) -> np.ndarray:
@@ -183,9 +194,9 @@ def flow_like(x_use: np.ndarray) -> np.ndarray:
 
         if isinstance(self.solver, str):
             if self.solver == "rk4":
-                x_next = _runge_kutta(flow_like, self.dt, x)
+                x_next = runge_kutta(flow_like, self.dt, x)
             elif self.solver == "forward_euler":
-                x_next = _forward_euler(flow_like, self.dt, x)
+                x_next = forward_euler(flow_like, self.dt, x)
             else:
                 raise ValueError(f"Unsupported solver: {self.solver}")
         else:

tests/test_simulations.py

@@ -3,7 +3,7 @@
 import pytest
 
 from lorenzpy import simulations  # type: ignore
-from lorenzpy.simulations._base import _BaseSimFlow, _BaseSimIterate  # type: ignore
+from lorenzpy.simulations.base import _BaseSimFlow, _BaseSimIterate  # type: ignore
 
 
 @pytest.fixture