Метод класса Numba-оптимизации без изменения класса ⇐ Python

[Anonymous]

Предположим, у меня есть класс ForceGenerator, определенный в следующем коде:

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

import numba
import numpy as np

class ForceGenerator:
def __init__(self, a, b, c):
self.a = a
self.b = b
self.c = c
def force(self, x):
if x > 0:
return self.a*x**2 + self.b*x + self.c
else:
return self.b*x
def simple_force(self, x):
return self.a*x**2 + self.

@numba.jit
def euler_integration(force, x_0=0, v_0=0, m=1, dt=1e-3, t=10):
steps = int(t//dt)
X, V = np.zeros(steps), np.zeros(steps)
X[0], V[0] = x_0, v_0
for i in range(1,steps):
V[i] = force(X[i-1])*dt/m
X[i] = V[i]*dt
return X, V

example_simple_force = numba.jit(lambda x: x**2 + 2*x + 1)
example_force = numba.jit(lambda x: (x**2 + 2*x + 1)*(x>0) + (2*x)*(~(x>0)))

example_generator = ForceGenerator(1, 2, 1)

X, V = euler_integration(example_force)
print(X)
print(V)

Этот код выполняется и печатает два массива. Однако, если вместо этого я установлю example_force = numba.jit(lambda x: example_generator.force(x)), код завершится с ошибкой, поскольку numba не знает, как обращаться с классами. Я хочу знать, есть ли какой-нибудь способ скопировать example_generator.force как чистую функцию, которая не вызывает никаких методов класса, чтобы это могло работать - своего рода «плоская копия», эквивалентная лямбда-функции, которую я определил в этом коде. Один из хакерских способов сделать это для менее сложной функции simple_force:

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

import sympy as sym

x = sym.symbols('x')
simple_force = sym.lambdify(x, example_generator.simple_force(x))

X, V = euler_integration(numba.jit(simple_force))

Это работает, но не будет работать для example_generator.force, например, из-за условия. (Даже если бы я мог найти обходной путь, мне было интересно, есть ли более элегантный способ, чем этот.)
Использование jitclass из numba.experimental для меня не вариант, поскольку мой реальный вариант использования содержит наследование и множество других вещей, с которыми numba не может справиться, и добавление @numba.jit над определением силы (как в этой статье geeksforgeeks, которую я сильно подозрительно, что это сгенерировано ИИ, предполагает) также терпит неудачу. Я нашел эту ветку stackoverflow, в которой есть два ответа, ни один из которых мне не подходит. (Первый ответ предлагает использовать статический метод , что лишает смысла использование класса; второй ответ предлагает изменить метод класса, чтобы он вызывал функцию, оптимизированную для числовых значений, что для меня неосуществимо, поскольку на практике форма силы вычисляется внутри класса с использованием набора библиотек, таких как Sympy.)

Подробнее здесь: https://stackoverflow.com/questions/798 ... ying-class