[Python] How to write a generic Python function that works with Python, Numpy or Pandas...

What's the best way to write a python function that can be used with either float, Numpy, or Pandas data types and always returns the same data type as the arguments it was given. The catch is, the calculation includes one or more float values.

E.g. toy example:

def mycalc(x, a=1.0, b=1.0):
return a * x + b


(I've simplified the problem a lot here as I would ideally want to have more than one input argument like x, but you can assume that the function is vectorized in the sense that it works with Numpy array arguments and Pandas series).

For Numpy arrays and Pandas Series this works fine because the dtype is dictated by the input arguments.

import numpy as np
x = np.array([1, 2, 3], dtype="float32")
print(mycalc(x).dtype) # float32


import pandas as pd
x = pd.Series([1.0, 2.0, 3.0], dtype="float32")
print(mycalc(x).dtype) # float32


But when using numpy floats of lower precision, the dtype is 'lifted' to float64, presumably due to the float arguments in the formula:

x = np.float32(1.0)
print(mycalc(x).dtype) # float64


Ideally, I would like the function to work with Python floats, numpy scalars, numpy arrays, Pandas series, Jax arrays, and even Sympy symbolic variables if possible.

But I don't want to clutter up the function with too many additional statements to handle each case.

I tried this, which works with Numpy scalars but breaks when you provide arrays or series:

def mycalc(x, a=1.0, b=1.0):
a = type(x)(a)
b = type(x)(b)
return a * x + b

assert isinstance(mycalc(1.0), float)
assert isinstance(mycalc(np.float32(1.0)), np.float32)
mycalc(np.array([1, 2, 3], dtype="float32")) # raises TypeError: expected a sequence of integers or a single integer, got '1.0'


Also, there is an answer here to a similar question which uses a decorator function to make copies of the input argument, which is a nice idea, but this was only for extending the function from Numpy arrays to Pandas series and doesn't work with Python floats or Numpy scalars.

import functools

def apply_to_pandas(func):
@functools.wraps(func)
def wrapper_func(x, *args, **kwargs):
if isinstance(x, (np.ndarray, list)):
out = func(x, *args, **kwargs)
else:
out = x.copy(deep=False)
out[:] = np.apply_along_axis(func, 0, x, *args, **kwargs)
return out
return wrapper_func

@apply_to_pandas
def mycalc(x, a=1.0, b=1.0):
return a * x + b

mycalc(1.0) # TypeError: copy() got an unexpected keyword argument 'deep'


Update

As pointed out by @Dunes in the comments below, this is no longer a problem in Numpy versions 2.x as explained here in the Numpy 2.0 Migration Guide.

In the new version, (np.float32(1.0) + 1).dtype == "float32". Therefore the original function above returns a result of the same dtype as the input x.

Continue reading...