Merge branch 'main' into feature/lm/dataset-as-config

docs/tutorials/brownian_motion.py

@@ -56,7 +56,7 @@
 # Call the model to generate 1000 steps.
 
 # %%
-df_1d = bm_1d(n_steps=1000)
+df_1d = bm_1d.generate_from(n_steps=1000)
 
 # %%
 px.line(df_1d, x="t", y="x_0")
@@ -79,7 +79,7 @@
 # We call the model to generate 1000 steps.
 
 # %%
-df_2d = bm_2d(n_steps=500)
+df_2d = bm_2d.generate_from(n_steps=500)
 
 # %%
 (

docs/tutorials/harmonic_oscillator_chain.py

@@ -113,22 +113,22 @@
 #
 # This system can be solved in terms of _travelling waves_, obtained by discrete Fourier transform.
 #
-
+#
 # We can complexify the system
 # $$S_L[x_i] = S_L[x_i, \phi_j] \equiv S_L[X^\ast_i, X_j] = \int_{t_0}^{t_1}\mathbb{d} t \left\\{ \frac{1}{2}m \dot X^\ast_i \delta_{ij} \dot X_j - \frac{1}{2}m X^\ast_i A_{ij} X_j\right\\}\\,,$$
 # where $A_{ij} / \omega^2$ is equal to $(-2)$ if $i=j$, $1$ if $|i-j|=1$ or $|i-j|=N$, and $0$ otherwise;
 # $X_i \coloneqq x_i \mathbb{e}^{-\phi_i}$, $X^\ast_i \coloneqq x_i \mathbb{e}^{+\phi_i}$.
-
+#
 # $A_{ij}$ can be diagonalised by the inverse discrete Fourier transform
 # $$X_i = (F^{-1})_{ik} Y_k = \frac{1}{\sqrt{N}}\sum_k \mathbb{e}^{i \frac{2\mathbb{\pi}}{N} k\mathbb{i}} Y_k\\,.$$
-
+#
 # Calculating gives
 # $$S_L[X^\ast_i, X_j] = S_L[Y^\ast_i, Y_j] = \sum_{k=0}^{N-1} \int_{t_0}^{t_1}\mathbb{d} t \left\\{ \frac{1}{2}m \dot Y^\ast_k \dot Y_k - \frac{1}{2}m \omega^2\cdot4\sin^2\frac{2\mathbb{\pi}k}{N} Y^\ast_k Y_k\right\\}\\,.$$
 # We can arrive at an action for the Fourier modes
 # $$S_L[Y, Y^*] = \sum_{k=0}^{N-1} \int_{t_0}^{t_1}\mathbb{d} t \left\\{ \frac{1}{2}m \dot Y_k^* Y_k - \frac{1}{2}m \omega^2\cdot4\sin^2\frac{2\mathbb{\pi}k}{N} Y_k^* Y_k\right\\}\\,.$$
-
+#
 # The origional system can then be solved by $N$ complex oscillators.
-
+#
 # Since the original degrees of freedom are real, the initial conditions of the propagating waves need to satisfy
 # $Y_k = Y^*_{-k \mod N}$, see [Wikipedia](https://en.wikipedia.org/wiki/Discrete_Fourier_transform#DFT_of_real_and_purely_imaginary_signals).
 

docs/tutorials/pendulum.py

@@ -46,7 +46,9 @@
 # ## Data
 
 # %%
-df_pen = pen(n_periods=n_periods, n_samples_per_period=n_samples_per_period)
+df_pen = pen.generate_from(
+    n_periods=n_periods, n_samples_per_period=n_samples_per_period
+)
 df_pen.head()
 
 # %%

hamilflow/models/brownian_motion.py

@@ -6,6 +6,8 @@
 import scipy as sp
 from pydantic import BaseModel, Field, computed_field, field_validator
 
+from hamilflow.models.utils.typing import TypeTime
+
 
 class BrownianMotionSystem(BaseModel):
     r"""Definition of the Brownian Motion system
@@ -183,17 +185,28 @@ def _trajectory(self, n_new_steps: int, seed: int) -> np.ndarray:
 
         return trajectory
 
-    def __call__(self, n_steps: int, seed: int = 42) -> pd.DataFrame:
-        """Simulate the coordinates of the particle
+    def generate_from(self, n_steps: int, seed: int = 42) -> pd.DataFrame:
+        """generate data from a set of interpretable params for this model
 
         :param n_steps: total number of steps to be simulated, including the inital step.
         :param seed: random generator seed for the stochastic process.
             Use it to reproduce results.
         """
+        time_steps = np.arange(0, n_steps) * self.system.delta_t
+
+        return self(t=time_steps, seed=seed)
+
+    def __call__(self, t: TypeTime, seed: int = 42) -> pd.DataFrame:
+        """Simulate the coordinates of the particle
+
+        :param t: the time sequence to be used to generate data, 1-D array like
+        :param seed: random generator seed for the stochastic process.
+            Use it to reproduce results.
+        """
+        n_steps = np.array(t).size
         trajectory = self._trajectory(n_new_steps=n_steps - 1, seed=seed)
 
         df = pd.DataFrame(trajectory, columns=self._axis_names)
-
-        df["t"] = np.arange(0, n_steps) * self.system.delta_t
+        df["t"] = t
 
         return df

hamilflow/models/pendulum.py

@@ -8,6 +8,8 @@
 from pydantic import BaseModel, Field, computed_field
 from scipy.special import ellipj, ellipk
 
+from hamilflow.models.utils.typing import TypeTime
+
 
 class PendulumSystem(BaseModel):
     r"""The params for the pendulum.
@@ -130,11 +132,26 @@ def theta(self, t: "Sequence[float] | ArrayLike[float]") -> np.ndarray[float]:
 
         return 2 * np.arcsin(cn / dn * self._k)
 
-    def __call__(self, n_periods: int, n_samples_per_period: int) -> pd.DataFrame:
+    def generate_from(self, n_periods: int, n_samples_per_period: int) -> pd.DataFrame:
+        """generate the time sequence from more interpretable params
+
+        :param n_periods: number of periods to include
+        :param n_samples_per_period: number of samples in each period
+        :return: an array that contains all the timesteps
+        """
         time_delta = self.period / n_samples_per_period
         time_steps = np.arange(0, n_periods * n_samples_per_period) * time_delta
 
-        thetas = self.theta(time_steps)
-        us = self.u(time_steps)
+        return self(time_steps)
+
+    def __call__(self, t: TypeTime) -> pd.DataFrame:
+        """generate the variables of the pendulum in time together with the time steps
+
+        :param t: time steps
+        :return: values of the variables
+            angle `x`, generalized coordinates `u`, and time `t`.
+        """
+        thetas = self.theta(t)
+        us = self.u(t)
 
-        return pd.DataFrame(dict(t=time_steps, x=thetas, u=us))
+        return pd.DataFrame(dict(t=t, x=thetas, u=us))

hamilflow/models/utils/typing.py

@@ -0,0 +1,5 @@
+from typing import Sequence, TypeVar
+
+from numpy.typing import ArrayLike
+
+TypeTime = TypeVar("TypeTime", Sequence[float], Sequence[int], ArrayLike)

pyproject.toml

@@ -29,6 +29,7 @@ mkdocs-material = "^9.5.30"
 mkdocs-autorefs = "^1.0.1"
 mkdocstrings = {version = "^0.25.2", extras = ["python"]}
 mkdocs-jupyter = "^0.24.8"
+jupytext = "^1.16.4"
 
 
 [tool.poetry.group.tutorial.dependencies]