ortk95 / planetmapper

PlanetMapper: An open source Python package for visualising, navigating and mapping Solar System observations
https://planetmapper.readthedocs.io
MIT License
10 stars 1 forks source link

Add convenience functions for converting arrays of coordinates #358

Closed ortk95 closed 2 months ago

ortk95 commented 3 months ago

For example, converting arrays of lons/lats into arrays of ras/decs is currently very cumbersome, and would be much more useful if there was some built-in functionality to do this.

ra, dec = zip(
    *(
        body.lonlat2radec(lon, lat)
        for lon, lat in zip(oval.longitudes, oval.latitudes)
        if body.test_if_lonlat_visible(lon, lat)
    )
)

or

ra, dec = zip(
    *(
        (
            body.lonlat2radec(lon, lat)
            if body.test_if_lonlat_visible(lon, lat)
            else (np.nan, np.nan)
        )
        for lon, lat in zip(oval.longitudes, oval.latitudes)
    )
)

Perhaps could even fold this directly into the lonlat2radec type functions? Might be doable with a decorator wrapping all the public transforms which adds array functionality. May also be worth adding a not_visible_nan type argument for the lonlat2... methods while doing this.

ortk95 commented 3 months ago

This would also set up the API nicely for future vectorisation of the internal calculations in #322

ortk95 commented 2 months ago

Options are probably either (a) do this with a decorator, or (b), do this by calling some transform within each function. E.g. as a very rough implementation:


FloatOrArray = TypeVar('FloatOrArray', float, np.ndarray)
T = TypeVar('T')
P = ParamSpec('P')

def maybe_transform_arrays(
    fn: Callable[Concatenate[float, float, P], tuple[float, float]],
    a: np.ndarray | float,
    b: np.ndarray | float,
    *args: P.args,
    **kwargs: P.kwargs,
) -> tuple[float, float] | tuple[np.ndarray, np.ndarray]:
    if isinstance(a, numbers.Number) and isinstance(b, numbers.Number):
        return fn(a, b, *args, **kwargs)
    else:
        with np.nditer([a, b, None, None]) as it:
            for x, y, u, v in it:
                u[...], v[...] = fn(x, y, *args, **kwargs)
            return it.operands[2], it.operands[3]

def array_decorator(
    fn: Callable[Concatenate[float, float, P], tuple[float, float]]
) -> Callable[
    Concatenate[FloatOrArray, FloatOrArray, P], tuple[FloatOrArray, FloatOrArray]
]:
    @functools.wraps(fn)
    def wrapper(
        a: FloatOrArray, b: FloatOrArray, *args: P.args, **kwargs: P.kwargs
    ) -> tuple[FloatOrArray, FloatOrArray]:
        return maybe_transform_arrays(fn, a, b, *args, **kwargs)

    return wrapper

Pros for the decorator:

Cons for the decorator:

For both, it's probably hard to get the type hints completely right without extensive overloads for each and every transform method - the mixed func(float, array) -> (array, array) signature is valid, but hard (impossible?) to hint properly with TypeVars.

On balance, I think it's probably best not to use a decorator, as we really want the parameter names to be visible.