implement single step retrieval

main.py

@@ -2,10 +2,15 @@ import argparse
 import ringfit.data as data
 import click
 
+import click
+
+
+@click.command()
+@click.option("--count", default=1, help="Number of greetings.")
+@click.option("--name", prompt="Your name", help="The person to greet.")
+def main():
+    pass
+
+
 if __name__ == "main":
-    parser = argparse.ArgumentParser(
-        description="Fitting the transient response of the fibre loops."
-    )
-    parser.add_argument(
-        "file_path", type=str, help="A required integer positional argument"
-    )
+    main()

pyproject.toml

@@ -14,6 +14,7 @@ matplotlib = "^3.8.4"
 
 [tool.poetry.group.dev.dependencies]
 ipython = "^8.24.0"
+jupyter = "^1.0.0"
 
 [build-system]
 requires = ["poetry-core"]

ringfit/data.py

@@ -80,7 +80,7 @@ class ScanData:
         return self._time
 
     @functools.cache
-    def laser_steps(self, *args, **kwargs):
+    def laser_steps(self, *args, **kwargs) -> np.ndarray:
         """Find the indices of the laser signal ``laser`` where the
         frequency of the laser changes.  For the parameters, see
         :func:`utils.find_frequency_steps`.
@@ -100,7 +100,9 @@ class ScanData:
         return self._time[self.laser_steps()]
 
     @functools.cache
-    def output_end_averages(self, end_fraction: float = 0.1, *args, **kwargs):
+    def output_end_averages(
+        self, end_fraction: float = 0.1, steps: np.ndarray | None = None
+    ) -> np.ndarray:
         """
         The average output intensity at the end of each laser
         frequency modulation step.  **If** the system is in its
@@ -108,7 +110,9 @@ class ScanData:
         transmission at that frequency.
         """
 
+        if steps is None:
             steps = self.laser_steps()
+
         step_size = np.mean(np.diff(steps))
 
         if end_fraction > 1:
@@ -117,3 +121,20 @@ class ScanData:
         window = int(step_size * end_fraction)
 
         return np.array([self._output[step - window : step].mean() for step in steps])
+
+    def for_step(self, step: int, steps: np.ndarray | None = None):
+        """Return time and output for the ``step``.  If ``steps`` is
+        not provided, then they are retrieved using the default
+        settings.  See :any:`laser_steps`.
+        """
+        time_steps: np.ndarray = self.laser_steps() if steps is None else steps
+
+        if step < 0 or step >= len(time_steps):
+            raise ValueError("The step must be between 0 and the number of steps.")
+
+        padded_steps = [0, *time_steps, len(self._output) - 1]
+        return (
+            self._time[padded_steps[step] : padded_steps[step + 1]],
+            self._output[padded_steps[step] : padded_steps[step + 1]],
+            self._laser[padded_steps[step] : padded_steps[step + 1]],
+        )

ringfit/plotting.py

@@ -61,37 +61,45 @@ def plot_scan(
     data: data.ScanData,
     laser=False,
     output=True,
-    steps=False,
+    steps: bool | int = False,
     normalize=False,
-    smoothe_output=False,
+    smoothe_output: bool | int = False,
     ax=None,
     **kwargs,
 ):
     if not (laser or output):
         raise ValueError("At least one of 'laser' or 'output' must be True.")
 
-    if laser:
-        laser = data.laser
-        if normalize:
-            laser = (laser - laser.min()) / (laser.max() - laser.min())
+    time, output_data, laser_data = (
+        (data.time, data.output, data.laser)
+        if isinstance(steps, bool)
+        else data.for_step(steps)
+    )
 
-        lines = ax.plot(data.time, laser, **kwargs)
+    if laser:
+        if normalize:
+            laser_data = (laser_data - laser_data.min()) / (
+                laser_data.max() - laser_data.min()
+            )
+
+        lines = ax.plot(time, laser_data, **kwargs)
 
     if output:
-        output = data.output
         if normalize:
-            output = (output - output.min()) / (output.max() - output.min())
+            output_data = (output_data - output_data.min()) / (
+                output_data.max() - output_data.min()
+            )
 
         if smoothe_output:
             if not isinstance(smoothe_output, int):
-                smoothe_output = 60
+                smoothe_output_data = 60
 
-            window = len(output) // smoothe_output
-            output = uniform_filter1d(output, window)
+            window = len(output_data) // smoothe_output
+            output_data = uniform_filter1d(output_data, window)
 
-        lines = ax.plot(data.time, output, **kwargs)
+        lines = ax.plot(time, output_data, **kwargs)
 
-    if steps:
+    if isinstance(steps, bool) and steps:
         peaks = data.laser_steps()
         for peak in peaks:
             ax.axvline(

scripts/parse_and_plot_scan.py

@@ -11,6 +11,4 @@ path = (
 )
 scan = ScanData.from_dir(path, truncation=[0, 50])
 
-fig, ax = plot_scan(scan, steps=True, smoothe_output=1000)
-
-plt.savefig("../figures/non_steady.png")
+fig, ax = plot_scan(scan, smoothe_output=10, steps=1)

scripts/startup.ipy

@@ -1,2 +1,9 @@
 %matplotlib tk
 plt.ion()
+
+import sys
+sys.path.append("../")
+import matplotlib.pyplot as plt
+from ringfit.data import *
+from ringfit.utils import *
+from ringfit.plotting import *