Skip to content

Refactor BaseTrajSeq key_attribute #242

@o-murphy

Description

@o-murphy

I do not like it, and it's not good way to annotate "choices" argument key_attribute: str
There 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

    @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",
}

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions