Skip to content

Refactor BaseTrajSeq key_attribute #242

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

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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