Как написать общую функцию Python, которая работает с аргументами Python, Numpy или Pandas и возвращает тот же тип ⇐ Python

[Anonymous]

Как лучше всего написать функцию Python, которая может использоваться с типами данных float, Numpy или Pandas и всегда возвращает тот же тип данных, что и переданные ей аргументы. Загвоздка в том, что расчет включает одно или несколько значений с плавающей запятой.
Например. пример игрушки:

Код: Выделить всё

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

(Здесь я значительно упростил проблему, поскольку в идеале мне хотелось бы иметь более одного входного аргумента, например x, но вы можете предположить, что функция векторизована в в том смысле, что он работает с аргументами массива Numpy и сериями Pandas).
Для массивов Numpy и серий Pandas это работает нормально, поскольку тип dtype определяется входными аргументами.

Код: Выделить всё

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

Но при использовании numpy с плавающей запятой более низкой точности dtype «поднимается» до float64, предположительно из-за аргументов с плавающей запятой в формуле:

Код: Выделить всё

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

В идеале я бы хотел, чтобы функция работала с числами с плавающей запятой Python, скалярами numpy, массивами numpy, сериями Pandas, массивами Jax и даже символическими переменными Sympy, если это возможно.
/>Но я не хочу загромождать функцию слишком большим количеством дополнительных операторов для обработки каждого случая.
Я попробовал это, которое работает со скалярами Numpy, но ломается, когда вы предоставлять массивы или сериал:

Код: Выделить всё

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'

Кроме того, здесь есть ответ на аналогичный вопрос, который использует функцию декоратора для создания копий входного аргумента, что является хорошей идеей, но это было только для расширения функции из Массивы Numpy преобразуются в серии Pandas и не работают с числами с плавающей точкой Python или скалярами Numpy.

Код: Выделить всё

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'

Обновление
Как отметил @Dunes в комментариях ниже, это больше не является проблемой в Numpy версии 2.x, как описано здесь, в Руководстве по миграции Numpy 2.0.
В новой версии (np.float32(1.0) + 1).dtype == "float32". Поэтому исходная функция выше возвращает результат того же типа d, что и входной x.

Подробнее здесь: https://stackoverflow.com/questions/792 ... r-pandas-a