@ тел сделал большую часть работы.Вот как вы можете получить еще 2-кратное ускорение поверх их 20-кратного.
Делайте линейную алгебру вручную.Когда я попробовал это, я был шокирован тем, насколько расточительна крошка на маленьких матрицах:
>>> from timeit import timeit
# using eigvalsh
>>> print(timeit("test(False, 0.1)", setup="from __main__ import test", number=3))
-2301206.495955009
-2301206.495955009
-2301206.495955009
55.794611917983275
>>> print(timeit("test(False, 0.3)", setup="from __main__ import test", number=5))
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
-85240.46154498367
3.400342195003759
# by hand
>>> print(timeit("test(True, 0.1)", setup="from __main__ import test", number=3))
-2301206.495955076
-2301206.495955076
-2301206.495955076
26.67294767702697
>>> print(timeit("test(True, 0.3)", setup="from __main__ import test", number=5))
-85240.46154498379
-85240.46154498379
-85240.46154498379
-85240.46154498379
-85240.46154498379
1.5047460949863307
Обратите внимание, что часть ускорения, вероятно, замаскирована общим кодом, на одной только линейной алгебре, кажется, больше, хотя я нене проверяйте слишком остро.
Одно предостережение: я использую дополнение Шура на разбиении матриц 2 на 2 для вычисления диагональных элементов обратного.Это потерпит неудачу, если дополнения Шура не существуют, т. Е. Если верхняя левая или нижняя правая субматрица необратимы.
Вот модифицированный код:
import itertools as it
import numpy as np
import sympy as sp
from sympy.abc import x, y, z
from sympy.core.numbers import Number
from sympy.utilities.lambdify import implemented_function
xones = implemented_function('xones', lambda x: np.ones(len(x)))
lfuncs = {'xones': xones}
def vectorizemat(mat):
ret = mat.copy()
for x in mat.free_symbols: break
for i,j in it.product(*(range(s) for s in mat.shape)):
if isinstance(mat[i,j], Number):
ret[i,j] = xones(x) * mat[i,j]
return ret
class Solver:
def __init__(self, vmat):
vmat = vectorizemat(vmat)
self._vfunc = sp.lambdify((x, y, z),
expr=vmat,
modules=[lfuncs, 'numpy'])
self._q_count, self._qs = None, [] # these depend on ksep!
def eval_s_vectorized_completely(self, stiff):
assert len(self._qs) == self._q_count, "Run 'populate_qs' first!"
mats = self._vfunc(*self._qs.T).T
evs = np.linalg.eigvalsh(mats)
result = np.sum(np.divide(1., (stiff + evs)))
return result.real - 4 * self._q_count
def eval_s_pp(self, stiff):
assert len(self._qs) == self._q_count, "Run 'populate_qs' first!"
mats = self._vfunc(*self._qs.T).T
np.einsum('...ii->...i', mats)[...] += stiff
(A, B), (C, D) = mats.reshape(-1, 2, 2, 2, 2).transpose(1, 3, 0, 2, 4)
res = 0
for AA, BB, CC, DD in ((A, B, C, D), (D, C, B, A)):
(a, b), (c, d) = DD.transpose(1, 2, 0)
rdet = 1 / (a*d - b*b)[:, None]
iD = DD[..., ::-1, ::-1].copy()
iD.reshape(-1, 4)[..., 1:3] *= -rdet
np.einsum('...ii->...i', iD)[...] *= rdet
(Aa, Ab), (Ac, Ad) = AA.transpose(1, 2, 0)
(Ba, Bb), (Bc, Bd) = BB.transpose(1, 2, 0)
(Da, Db), (Dc, Dd) = iD.transpose(1, 2, 0)
a = Aa - Ba*Ba*Da - 2*Bb*Ba*Db - Bb*Bb*Dd
d = Ad - Bd*Bd*Dd - 2*Bc*Bd*Db - Bc*Bc*Da
b = Ab - Ba*Bc*Da - Ba*Bd*Db - Bb*Bd*Dd - Bb*Bc*Dc
res += ((a + d) / (a*d - b*b)).sum()
return res - 4 * self._q_count
def populate_qs(self, ksep: float = 1.7):
self._qs = np.array([(kx, ky, kz) for kx, ky, kz
in it.product(np.arange(-3*np.pi, 3.01*np.pi, ksep),
np.arange(-3*np.pi, 3.01*np.pi, ksep),
np.arange(-3*np.pi, 3.01*np.pi, ksep))])
self._q_count = len(self._qs)
def test(manual=False, ksep=0.3):
vmat = sp.Matrix([[1, sp.cos(x/4+y/4), sp.cos(x/4+z/4), sp.cos(y/4+z/4)],
[sp.cos(x/4+y/4), 1, sp.cos(y/4-z/4), sp.cos(x/4 - z/4)],
[sp.cos(x/4+z/4), sp.cos(y/4-z/4), 1, sp.cos(x/4-y/4)],
[sp.cos(y/4+z/4), sp.cos(x/4-z/4), sp.cos(x/4-y/4), 1]]) * 2
solver = Solver(vmat)
solver.populate_qs(ksep=ksep) # <---- Performance starts to worsen (in eval_s) when ksep is reduced!
if manual:
print(solver.eval_s_pp(0.65))
else:
print(solver.eval_s_vectorized_completely(0.65))