computational performance about shenfun HOT 9 OPEN

def compute_rhs(duvh, uvh):
    global count
    count += 1
    duvh.fill(0)
    u, v, h = uvh[:]
    du, dv, dh = duvh[:]
    #
    du[:] = -g*D(hf,dv0,0) # -g*dhdx
    du += -u*D(uf,dv0,0)   # -u*dudx
    du += -v*D(uf,dv0,1)   # -v*dudy
    du += -Cd/H*u*np.sqrt(u**2+v**2)
    du += f*v

# init variables for derivatives
uf = Function(T, False, buffer=u)
vf = Function(T, False, buffer=v)
hf = Function(T, False, buffer=h)
dv0 = np.zeros_like(u)
def D(varf,dvar,dim, padded=False):
    ''' Wrapper around Dx
    '''
    if not padded:
        dvar[:] = T.backward(project(Dx(varf, dim, 1), T))
    else:
        dvar[:] = Tp.backward(project(Dx(varf, dim, 1), T))
    return dvar

work = Array(T, True)
K = T.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
def D(varf,dvar,dim):
    ''' Wrapper around Dx
    '''
    global work
    work = T.forward(varf, work)
    dvar = T.backward((1j*K[dim])*work, dvar)
    return dvar

Wrote profile results to swater_1L_physical.py.lprof
Timer unit: 1e-06 s

Total time: 11.2486 s
File: swater_1L_physical.py
Function: compute_rhs at line 102

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
   102                                           @profile
   103                                           def compute_rhs(duvh, uvh):
   104                                               global count
   105       576       1163.0      2.0      0.0      count += 1
   106       576      16547.0     28.7      0.1      duvh.fill(0)
   107       576      11322.0     19.7      0.1      u, v, h = uvh[:]
   108       576       5874.0     10.2      0.1      du, dv, dh = duvh[:]
   109                                               #
   110       576    1152978.0   2001.7     10.2      du[:] = -g*D(hf,dv0,0) # -g*dhdx
   111       576    1095948.0   1902.7      9.7      du += -u*D(uf,dv0,0)   # -u*dudx
   112       576    1091663.0   1895.2      9.7      du += -v*D(uf,dv0,1)   # -v*dudy
   113       576     100679.0    174.8      0.9      du += -Cd/H*u*np.sqrt(u**2+v**2)
   114       576      20936.0     36.3      0.2      du += f*v
   115                                               #

(shenfun) br222-003:swater_1L aponte$ python -m line_profiler swater_1L_physical.py_opt.lprof
Timer unit: 1e-06 s

Total time: 3.60774 s
File: swater_1L_physical.py
Function: compute_rhs at line 100

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
   100                                           @profile
   101                                           def compute_rhs(duvh, uvh):
   102                                               global count
   103       576       1484.0      2.6      0.0      count += 1
   104       576      18370.0     31.9      0.5      duvh.fill(0)
   105       576      14141.0     24.6      0.4      u, v, h = uvh[:]
   106       576       6692.0     11.6      0.2      du, dv, dh = duvh[:]
   107                                               #
   108       576     349008.0    605.9      9.7      du[:] = -g*D(h,dv0,0) # -g*dhdx
   109       576     326703.0    567.2      9.1      du += -u*D(u,dv0,0)   # -u*dudx
   110       576     321089.0    557.4      8.9      du += -v*D(u,dv0,1)   # -v*dudy
   111       576     117991.0    204.8      3.3      du += -Cd/H*u*np.sqrt(u**2+v**2)
   112       576      24379.0     42.3      0.7      du += f*v
   113                                               #

from numba import jit, complex128, float64, int64
@jit((complex128[:, ::1], complex128[:, ::1], float64[::1], float64[::1], int64), nopython=True)
def deriv(work, wx, kx, ky, dim):
    if dim == 0:
        for i in range(wx.shape[0]):
            k0 = kx[i]*1j
            for j in range(wx.shape[1]):
                wx[i, j] = work[i, j]*k0

    else:
        for i in range(wx.shape[0]):
            for j in range(wx.shape[1]):
                k1 = ky[j]*1j
                wx[i, j] = work[i, j]*k1

work2 = Array(T, True)
def D(varf, dvar, dim):
    ''' Wrapper around Dx
    '''
    global K, work, work2
    work = T.forward(varf, work)
    deriv(work, work2, np.squeeze(K[0]), np.squeeze(K[1]), dim)
    dvar = T.backward(work2, dvar)
    return dvar