But I would use an enum for consistency with C/Cython, to keep single type for the key_kind (remove key_attribute) , and to avoid misunderstandings
@staticmethod
def interpolate(
key_attribute: str,
key_value: float,
p0: BaseTrajData,
p1: BaseTrajData,
p2: BaseTrajData,
method: InterpolationMethod = "pchip",
) -> BaseTrajData:
"""
Interpolate a BaseTrajData point using monotone PCHIP (default) or linear.
Args:
key_attribute: Can be 'time', 'mach', or a vector component like 'position.x' or 'velocity.z'.
key_value: The value to interpolate for.
p0: First bracketing point.
p1: Second (middle) bracketing point.
p2: Third bracketing point.
method: 'pchip' (default, monotone cubic Hermite) or 'linear'.
Returns:
The interpolated data point.
Raises:
AttributeError: If the key_attribute is not a member of BaseTrajData.
ZeroDivisionError: If the interpolation fails due to zero division.
(This will result if two of the points are identical).
ValueError: If method is not one of 'pchip' or 'linear'.
"""
def get_key_val(td: "BaseTrajData", path: str) -> float:
"""Helper to get the key value from a BaseTrajData point."""
if "." in path:
top, component = path.split(".", 1)
obj = getattr(td, top)
return getattr(obj, component)
return getattr(td, path)
# independent variable values
x0 = get_key_val(p0, key_attribute)
x1 = get_key_val(p1, key_attribute)
x2 = get_key_val(p2, key_attribute)
def _interp_scalar(y0, y1, y2):
if method == "pchip":
return interpolate_3_pt(key_value, x0, y0, x1, y1, x2, y2)
elif method == "linear":
pts = sorted(((x0, y0), (x1, y1), (x2, y2)), key=lambda p: p[0])
(sx0, sy0), (sx1, sy1), (sx2, sy2) = pts
if key_value <= sx1:
return interpolate_2_pt(key_value, sx0, sy0, sx1, sy1)
else:
return interpolate_2_pt(key_value, sx1, sy1, sx2, sy2)
else:
raise ValueError("method must be 'pchip' or 'linear'")
time = _interp_scalar(p0.time, p1.time, p2.time) if key_attribute != "time" else key_value
px = _interp_scalar(p0.position.x, p1.position.x, p2.position.x)
py = _interp_scalar(p0.position.y, p1.position.y, p2.position.y)
pz = _interp_scalar(p0.position.z, p1.position.z, p2.position.z)
position = Vector(px, py, pz)
vx = _interp_scalar(p0.velocity.x, p1.velocity.x, p2.velocity.x)
vy = _interp_scalar(p0.velocity.y, p1.velocity.y, p2.velocity.y)
vz = _interp_scalar(p0.velocity.z, p1.velocity.z, p2.velocity.z)
velocity = Vector(vx, vy, vz)
mach = _interp_scalar(p0.mach, p1.mach, p2.mach) if key_attribute != "mach" else key_value
return BaseTrajData(time=time, position=position, velocity=velocity, mach=mach)
TRAJECTORY_DATA_ATTRIBUTES = Literal[
"time",
"distance",
"velocity",
"mach",
"height",
"slant_height",
"drop_angle",
"windage",
"windage_angle",
"slant_distance",
"angle",
"density_ratio",
"drag",
"energy",
"ogw",
"flag",
"x",
"y",
"z",
]
TRAJECTORY_DATA_SYNONYMS: dict[TRAJECTORY_DATA_ATTRIBUTES, TRAJECTORY_DATA_ATTRIBUTES] = {
"x": "distance",
"y": "height",
"z": "windage",
}
I do not like it, and it's not good way to annotate "choices" argument
key_attribute: strThere should be at least
Literal["mach", "time", "position.x", ...]But I would use an enum for consistency with C/Cython, to keep single type for the key_kind (remove key_attribute) , and to avoid misunderstandings