add simple data interpolator

examples/common.py

@@ -1,4 +1,5 @@
-from typing import cast, TypedDict
+from typing import cast, TypedDict, Optional, Sequence, Union
+from datetime import datetime
 
 import requests
 
@@ -30,3 +31,77 @@ def format_coord(x, y): # type: ignore
         coords = [ax_coord, (x, y)]
         return f'{date} funds={ax_coord[1]:1.0f} profit={y:1.0f}'
     return format_coord
+
+def _to_float_time(t: Union[float, int, datetime]) -> float:
+    if isinstance(t, datetime):
+        return t.timestamp()
+    return float(t)
+
+def interpolate(
+    y: Sequence[Union[int, float]],
+    x: Optional[Sequence[Union[int, float, datetime]]] = None,
+    max_hole_size: int = 0
+) -> list[Union[float, int]]:
+    """
+    Replace data holes (zeros) between two non-zero values with linearly
+    interpolated values.
+
+    If a data hole is bigger than the input max_hole_size, leave the hole as is.
+
+    Parameters
+    ----------
+    y:
+        Input data, potentially with holes
+    x: [optional]
+        Sequence of x values (same length as y).
+        If omitted, x values are assumed equally spaced.
+    max_hole_size: [optional][default: 0]
+        Maximum number of consecutive zeros to interpolate across.
+        0 means no limit.
+    """
+    result = [float(u) for u in y]
+    n = len(result)
+
+    if n < 3:
+        return result
+
+    if x is not None:
+        if len(x) != n:
+            raise ValueError("x and y must have the same length.")
+        t = [_to_float_time(u) for u in x]
+    else:
+        t = list(range(n))
+
+    i = 0
+    while i < n:
+        if result[i] != 0:
+            i += 1
+            continue
+
+        start = i
+        while ((i < n) and (result[i] == 0)):
+            i += 1
+        end = i
+
+        if start == 0 or end >= n:
+            # can't interpolate when boundary values are missing
+            continue
+
+        gap = end - start
+        if max_hole_size > 0 and gap > max_hole_size:
+            # don't interpolate on holes bigger than max_hole_size
+            continue
+
+        v_left = result[start - 1]
+        v_right = result[end]
+        t_left = t[start - 1]
+        t_right = t[end]
+
+        if t_right == t_left:
+            raise ValueError("found duplicate x (time) values")
+
+        slope = (v_right - v_left) / (t_right - t_left)
+        for j in range(start, end):
+            result[j] = v_left + slope * (t[j] - t_left)
+
+    return result

examples/gpu_temperature.py

@@ -6,7 +6,7 @@
 import matplotlib.dates as mdates
 import requests
 
-from common import Config, generate_token, make_format
+from common import Config, generate_token, interpolate
 
 if __name__ == '__main__':
     with open(os.path.join(os.path.dirname(__file__), 'config.yaml')) as f:
@@ -18,7 +18,8 @@
         r = requests.get(f'{config['base_url']}/metrics/{cluster}/{jobid}/gpu/temperature', headers=auth_header)
         r.raise_for_status()
 
-        time = [datetime.datetime.fromtimestamp(x) for x in r.json()['time']]
+        response = r.json()
+        time = [datetime.datetime.fromtimestamp(x) for x in response['time']]
         nodeTemps = r.json()['nodes']
 
         fig, ax = plt.subplots()
@@ -27,7 +28,11 @@
 
         for nodeid, nodeTemps in nodeTemps.items():
             for gpuTemps in nodeTemps:
-                ax.plot(time, gpuTemps['temperature'], label=f'GPU {gpuTemps['gpu_id']}')
+                temps_interpolated = interpolate(
+                    y=gpuTemps['temperature'],
+                    x=time,
+                    max_hole_size=5)
+                ax.plot(time, temps_interpolated, label=f'GPU {gpuTemps['gpu_id']}')
             # only plot GPU temperatures of first node
             break