From e5c8f0e70ffe9dc2ea959ccd9ea87c5365b77779 Mon Sep 17 00:00:00 2001
From: "valentin.boettcher@mailbox.tu-dresden.de" <hiro@protagon.space>
Date: Mon, 24 Jul 2023 15:23:11 -0400
Subject: [PATCH] use builtin strobe for energy plot

---
 python/otto_motor/otto_utilities.py |  2 +-
 python/plot_utils.py                | 12 ++++++++----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/python/otto_motor/otto_utilities.py b/python/otto_motor/otto_utilities.py
index fc3823d..9f39eb1 100644
--- a/python/otto_motor/otto_utilities.py
+++ b/python/otto_motor/otto_utilities.py
@@ -348,7 +348,7 @@ def full_report(model):
 def plot_energy(model):
     f, a = pu.plot_energy_overview(
         model,
-        strobe_frequency=model.Ω,
+        strobe_data=model.strobe,
         hybrid=True,
         bath_names=["cold", "hot"],
         online=True,
diff --git a/python/plot_utils.py b/python/plot_utils.py
index be3591c..4d12c54 100644
--- a/python/plot_utils.py
+++ b/python/plot_utils.py
@@ -229,21 +229,25 @@ def plot_with_σ(
     bath=None,
     strobe_frequency=None,
     strobe_tolerance=1e-3,
+    strobe_data=None,
     hybrid=False,
     **kwargs,
 ):
     err = (y.σ[bath] if bath is not None else y.σ).real
     y_final = transform(y.value[bath] if bath is not None else y.value)
 
-    strobe_mode = strobe_frequency is not None
+    strobe_mode = strobe_frequency is not None or strobe_data
     strobe_indices = None
     strobe_times = None
     strobe_style = dict(linestyle="none", marker="o", markersize=2) | kwargs
 
     if strobe_mode:
-        strobe_times, strobe_indices = ut.strobe_times(
-            x, strobe_frequency, strobe_tolerance
-        )
+        if strobe_data:
+            strobe_times, strobe_indices = strobe_data
+        else:
+            strobe_times, strobe_indices = ut.strobe_times(
+                x, strobe_frequency, strobe_tolerance
+            )
 
         line = ax.errorbar(
             strobe_times,