Closed ortk95 closed 2 months ago
This would also set up the API nicely for future vectorisation of the internal calculations in #322
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 TypeVar
s.
On balance, I think it's probably best not to use a decorator, as we really want the parameter names to be visible.
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.
or
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 anot_visible_nan
type argument for thelonlat2...
methods while doing this.