Implement Series.nearest() to allow finding value by nearest key in a Series object
57 lines
1.4 KiB
Python
57 lines
1.4 KiB
Python
import enum
|
|
from typing import Self
|
|
|
|
from xmet.list import nearest
|
|
|
|
class SeriesIntersection(enum.Enum):
|
|
GREATER = -1
|
|
EQUAL = 0
|
|
LESSER = 1
|
|
|
|
@staticmethod
|
|
def cmp(a, b) -> Self:
|
|
if a > b:
|
|
return SeriesIntersection.GREATER
|
|
elif a == b:
|
|
return SeriesIntersection.EQUAL
|
|
elif a < b:
|
|
return SeriesIntersection.LESSER
|
|
|
|
def inverse(self):
|
|
if self is SeriesIntersection.GREATER:
|
|
return SeriesIntersection.LESSER
|
|
elif self is SeriesIntersection.LESSER:
|
|
return SeriesIntersection.GREATER
|
|
|
|
return self
|
|
|
|
class Series(dict):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def neighbors(self, series: Self):
|
|
return nearest(sorted(self.keys(), reverse=True),
|
|
sorted(series.keys(), reverse=True))
|
|
|
|
def intersect(self,
|
|
series: Self,
|
|
intersection: SeriesIntersection,
|
|
start: float=None) -> tuple[float]:
|
|
pairs = self.neighbors(series)
|
|
|
|
for pair in pairs:
|
|
v1, v2 = self[pair[0]], series[pair[1]]
|
|
|
|
if start is not None:
|
|
if pair[0] > start:
|
|
continue
|
|
|
|
sign = SeriesIntersection.cmp(v1, v2)
|
|
|
|
if sign is intersection:
|
|
return v1, pair[0]
|
|
|
|
def nearest(self, key):
|
|
k = sorted(self.keys(), key=lambda k: abs(k-key))[0]
|
|
|
|
return self[k]
|