# Source code for GPy.kern.src.eq_ode2

# Copyright (c) 2014, Cristian Guarnizo.

import numpy as np
from scipy.special import wofz
from .kern import Kern
from ...core.parameterization import Param
from paramz.transformations import Logexp
from paramz.caching import Cache_this

[docs]class EQ_ODE2(Kern):
"""
Covariance function for second order differential equation driven by an exponentiated quadratic covariance.

This outputs of this kernel have the form
.. math::
\frac{\text{d}^2y_j(t)}{\text{d}^2t} + C_j\frac{\text{d}y_j(t)}{\text{d}t} + B_jy_j(t) = \sum_{i=1}^R w_{j,i} u_i(t)

where :math:R is the rank of the system, :math:w_{j,i} is the sensitivity of the :math:jth output to the :math:ith latent function, :math:d_j is the decay rate of the :math:jth output and :math:f_i(t) and :math:g_i(t) are independent latent Gaussian processes goverened by an exponentiated quadratic covariance.

:param output_dim: number of outputs driven by latent function.
:type output_dim: int
:param W: sensitivities of each output to the latent driving function.
:type W: ndarray (output_dim x rank).
:param rank: If rank is greater than 1 then there are assumed to be a total of rank latent forces independently driving the system, each with identical covariance.
:type rank: int
:param C: damper constant for the second order system.
:type C: array of length output_dim.
:param B: spring constant for the second order system.
:type B: array of length output_dim.

"""
#This code will only work for the sparseGP model, due to limitations in models for this kernel
def __init__(self, input_dim=2, output_dim=1, rank=1, W=None, lengthscale=None, C=None, B=None, active_dims=None, name='eq_ode2'):
#input_dim should be 1, but kern._slice_X is not returning index information required to evaluate kernels
assert input_dim == 2, "only defined for 1 input dims"
super(EQ_ODE2, self).__init__(input_dim=input_dim, active_dims=active_dims, name=name)
self.rank = rank
self.output_dim = output_dim

if lengthscale is None:
lengthscale = .5+np.random.rand(self.rank)
else:
lengthscale = np.asarray(lengthscale)
assert lengthscale.size in [1, self.rank], "Bad number of lengthscales"
if lengthscale.size != self.rank:
lengthscale = np.ones(self.rank)*lengthscale

if W is None:
#W = 0.5*np.random.randn(self.output_dim, self.rank)/np.sqrt(self.rank)
W = np.ones((self.output_dim, self.rank))
else:
assert W.shape == (self.output_dim, self.rank)

if C is None:
C = np.ones(self.output_dim)

if B is None:
B = np.ones(self.output_dim)

self.C = Param('C', C, Logexp())
self.B = Param('B', B, Logexp())
self.lengthscale = Param('lengthscale', lengthscale, Logexp())
self.W = Param('W', W)

[docs]    @Cache_this(limit=3)
def K(self, X, X2=None):
#This way is not working, indexes are lost after using k._slice_X
#index = np.asarray(X, dtype=np.int)
#index = index.reshape(index.size,)
if hasattr(X, 'values'):
X = X.values
index = np.int_(np.round(X[:, 1]))
index = index.reshape(index.size,)
X_flag = index[0] >= self.output_dim
if X2 is None:
if X_flag:
#Calculate covariance function for the latent functions
index -= self.output_dim
return self._Kuu(X, index)
else: #Kff full
raise NotImplementedError
else:
#This way is not working, indexes are lost after using k._slice_X
#index2 = np.asarray(X2, dtype=np.int)
#index2 = index2.reshape(index2.size,)
if hasattr(X2, 'values'):
X2 = X2.values
index2 = np.int_(np.round(X2[:, 1]))
index2 = index2.reshape(index2.size,)
X2_flag = index2[0] >= self.output_dim
#Calculate cross-covariance function
if not X_flag and X2_flag:
index2 -= self.output_dim
return self._Kfu(X, index, X2, index2) #Kfu
elif X_flag and not X2_flag:
index -= self.output_dim
return self._Kfu(X2, index2, X, index).T #Kuf
elif X_flag and X2_flag:
index -= self.output_dim
index2 -= self.output_dim
return self._Kusu(X, index, X2, index2) #Ku_s u
else:
raise NotImplementedError #Kf_s f

#Calculate the covariance function for diag(Kff(X,X))
[docs]    def Kdiag(self, X):
if hasattr(X, 'values'):
index = np.int_(np.round(X[:, 1].values))
else:
index = np.int_(np.round(X[:, 1]))
index = index.reshape(index.size,)
X_flag = index[0] >= self.output_dim

if X_flag: #Kuudiag
return np.ones(X[:,0].shape)
else: #Kffdiag
kdiag = self._Kdiag(X)
return np.sum(kdiag, axis=1)

#Calculate the covariance function for diag(Kff(X,X))
def _Kdiag(self, X):
#This way is not working, indexes are lost after using k._slice_X
#index = np.asarray(X, dtype=np.int)
#index = index.reshape(index.size,)
if hasattr(X, 'values'):
X = X.values
index = np.int_(X[:, 1])
index = index.reshape(index.size,)

#terms that move along t
t = X[:, 0].reshape(X.shape[0], 1)
d = np.unique(index) #Output Indexes
B = self.B.values[d]
C = self.C.values[d]
S = self.W.values[d, :]
#Index transformation
indd = np.arange(self.output_dim)
indd[d] = np.arange(d.size)
index = indd[index]
#Check where wd becomes complex
wbool = C*C >= 4.*B
B = B.reshape(B.size, 1)
C = C.reshape(C.size, 1)
alpha = .5*C
C2 = C*C

wbool2 = wbool[index]
ind2t = np.where(wbool2)
ind3t = np.where(np.logical_not(wbool2))

#Terms that move along q
lq = self.lengthscale.values.reshape(1, self.lengthscale.size)
S2 = S*S
kdiag = np.empty((t.size, lq.size))

indD = np.arange(B.size)
#(1) When wd is real
if np.any(np.logical_not(wbool)):
#Indexes of index and t related to (2)
t1 = t[ind3t]
ind = index[ind3t]
d = np.asarray(np.where(np.logical_not(wbool))[0]) #Selection of outputs
indd = indD.copy()
indd[d] = np.arange(d.size)
ind = indd[ind]
#Dx1 terms
S2lq = S2[d]*(.5*lq)
c0 = S2lq*np.sqrt(np.pi)
w = .5*np.sqrt(4.*B[d] - C2[d])
w2 = w*w
c2 = .5/(gam*w2)
c = c1 - c2
#DxQ terms
nu = lq*(gam*.5)
K01 = c0*c
#Nx1 terms
gamt = -gam[ind]*t1
gamct = -gamc[ind]*t1
egamt = np.exp(gamt)
ec = egamt*c2[ind] - np.exp(gamct)*c1[ind]
#NxQ terms
t_lq = t1/lq

# Upsilon Calculations
# Using wofz
wnu = wofz(1j*nu)
lwnu = np.log(wnu)
t2_lq2 = -t_lq*t_lq
upm = wnu[ind] - np.exp(t2_lq2 + gamt + np.log(wofz(1j*(t_lq + nu[ind]))))
upm[t1[:, 0] == 0, :] = 0.

nu2 = nu*nu
z1 = nu[ind] - t_lq
indv1 = np.where(z1.real >= 0.)
indv2 = np.where(z1.real < 0.)
upv = -np.exp(lwnu[ind] + gamt)
if indv1[0].shape > 0:
upv[indv1] += np.exp(t2_lq2[indv1] + np.log(wofz(1j*z1[indv1])))
if indv2[0].shape > 0:
upv[indv2] += np.exp(nu2[ind[indv2[0]], indv2[1]] + gamt[indv2[0], 0] + np.log(2.))\
- np.exp(t2_lq2[indv2] + np.log(wofz(-1j*z1[indv2])))
upv[t1[:, 0] == 0, :] = 0.

#Covariance calculation
kdiag[ind3t] = np.real(K01[ind]*upm)
kdiag[ind3t] += np.real((c0[ind]*ec)*upv)

#(2) When w_d is complex
if np.any(wbool):
t1 = t[ind2t]
ind = index[ind2t]
#Index transformation
d = np.asarray(np.where(wbool)[0])
indd = indD.copy()
indd[d] = np.arange(d.size)
ind = indd[ind]
#Dx1 terms
S2lq = S2[d]*(lq*.25)
c0 = S2lq*np.sqrt(np.pi)
w = .5*np.sqrt(C2[d] - 4.*B[d])
w2 = -w*w
c21 = .5/(gam*w2)
c22 = .5/(gamc*w2)
c = c1 - c21
c2 = c1 - c22
#DxQ terms
K011 = c0*c
K012 = c0*c2
nu = lq*(.5*gam)
nuc = lq*(.5*gamc)
#Nx1 terms
gamt = -gam[ind]*t1
gamct = -gamc[ind]*t1
egamt = np.exp(gamt)
egamct = np.exp(gamct)
ec = egamt*c21[ind] - egamct*c1[ind]
ec2 = egamct*c22[ind] - egamt*c1[ind]
#NxQ terms
t_lq = t1/lq

#Upsilon Calculations using wofz
t2_lq2 = -t_lq*t_lq #Required when using wofz
wnu = wofz(1j*nu).real
lwnu = np.log(wnu)
upm = wnu[ind] - np.exp(t2_lq2 + gamt + np.log(wofz(1j*(t_lq + nu[ind])).real))
upm[t1[:, 0] == 0., :] = 0.

nu2 = nu*nu
z1 = nu[ind] - t_lq
indv1 = np.where(z1 >= 0.)
indv2 = np.where(z1 < 0.)
upv = -np.exp(lwnu[ind] + gamt)
if indv1[0].shape > 0:
upv[indv1] += np.exp(t2_lq2[indv1] + np.log(wofz(1j*z1[indv1]).real))
if indv2[0].shape > 0:
upv[indv2] += np.exp(nu2[ind[indv2[0]], indv2[1]] + gamt[indv2[0], 0] + np.log(2.))\
- np.exp(t2_lq2[indv2] + np.log(wofz(-1j*z1[indv2]).real))
upv[t1[:, 0] == 0, :] = 0.

wnuc = wofz(1j*nuc).real
lwnuc = np.log(wnuc)

upmc = wnuc[ind] - np.exp(t2_lq2 + gamct + np.log(wofz(1j*(t_lq + nuc[ind])).real))
upmc[t1[:, 0] == 0., :] = 0.

nuc2 = nuc*nuc
z1 = nuc[ind] - t_lq
indv1 = np.where(z1 >= 0.)
indv2 = np.where(z1 < 0.)
upvc = - np.exp(lwnuc[ind] + gamct)
if indv1[0].shape > 0:
upvc[indv1] += np.exp(t2_lq2[indv1] + np.log(wofz(1j*z1[indv1]).real))
if indv2[0].shape > 0:
upvc[indv2] += np.exp(nuc2[ind[indv2[0]], indv2[1]] + gamct[indv2[0], 0] + np.log(2.))\
- np.exp(t2_lq2[indv2] + np.log(wofz(-1j*z1[indv2]).real))
upvc[t1[:, 0] == 0, :] = 0.

#Covariance calculation
kdiag[ind2t] = K011[ind]*upm + K012[ind]*upmc + (c0[ind]*ec)*upv + (c0[ind]*ec2)*upvc
return kdiag

[docs]    def update_gradients_full(self, dL_dK, X, X2 = None):
#index = np.asarray(X, dtype=np.int)
#index = index.reshape(index.size,)
if hasattr(X, 'values'):
X = X.values
index = np.int_(X[:, 1])
index = index.reshape(index.size,)
X_flag = index[0] >= self.output_dim
if X2 is None:
if X_flag: #Kuu or Kmm
index -= self.output_dim
tmp = dL_dK*self._gkuu_lq(X, index)
for q in np.unique(index):
ind = np.where(index == q)
else:
raise NotImplementedError
else: #Kfu or Knm
#index2 = np.asarray(X2, dtype=np.int)
#index2 = index2.reshape(index2.size,)
if hasattr(X2, 'values'):
X2 = X2.values
index2 = np.int_(X2[:, 1])
index2 = index2.reshape(index2.size,)
X2_flag = index2[0] >= self.output_dim
if not X_flag and X2_flag:
index2 -= self.output_dim
else:
dL_dK = dL_dK.T #so we obtaing dL_Kfu
indtemp = index - self.output_dim
Xtemp = X
X = X2
X2 = Xtemp
index = index2
index2 = indtemp
glq, gSdq, gB, gC = self._gkfu(X, index, X2, index2)
tmp = dL_dK*glq
for q in np.unique(index2):
ind = np.where(index2 == q)
tmpB = dL_dK*gB
tmpC = dL_dK*gC
tmp = dL_dK*gSdq
for d in np.unique(index):
ind = np.where(index == d)
for q in np.unique(index2):
ind2 = np.where(index2 == q)

#index = np.asarray(X, dtype=np.int)
#index = index.reshape(index.size,)
if hasattr(X, 'values'):
X = X.values
index = np.int_(X[:, 1])
index = index.reshape(index.size,)

glq, gS, gB, gC = self._gkdiag(X, index)
if dL_dKdiag.size == X.shape[0]:
dL_dKdiag = np.reshape(dL_dKdiag, (index.size, 1))
tmp = dL_dKdiag*glq
tmpB = dL_dKdiag*gB
tmpC = dL_dKdiag*gC
tmp = dL_dKdiag*gS
for d in np.unique(index):
ind = np.where(index == d)

[docs]    def gradients_X(self, dL_dK, X, X2=None):
#index = np.asarray(X, dtype=np.int)
#index = index.reshape(index.size,)
if hasattr(X, 'values'):
X = X.values
index = np.int_(X[:, 1])
index = index.reshape(index.size,)
X_flag = index[0] >= self.output_dim
#If input_dim == 1, use this
#gX = np.zeros((X.shape[0], 1))
#Cheat to allow gradient for input_dim==2
gX = np.zeros(X.shape)
if X2 is None: #Kuu or Kmm
if X_flag:
index -= self.output_dim
gX[:, 0] = 2.*(dL_dK*self._gkuu_X(X, index)).sum(0)
return gX
else:
raise NotImplementedError
else: #Kuf or Kmn
#index2 = np.asarray(X2, dtype=np.int)
#index2 = index2.reshape(index2.size,)
if hasattr(X2, 'values'):
X2 = X2.values
index2 = np.int_(X2[:, 1])
index2 = index2.reshape(index2.size,)
X2_flag = index2[0] >= self.output_dim
if X_flag and not X2_flag: #gradient of Kuf(Z, X) wrt Z
index -= self.output_dim
gX[:, 0] = (dL_dK*self._gkfu_z(X2, index2, X, index).T).sum(1)
return gX
else:
raise NotImplementedError

#---------------------------------------#
#             Helper functions          #
#---------------------------------------#

#Evaluation of squared exponential for LFM
def _Kuu(self, X, index):
index = index.reshape(index.size,)
t = X[:, 0].reshape(X.shape[0],)
lq = self.lengthscale.values.reshape(self.rank,)
lq2 = lq*lq
#Covariance matrix initialization
kuu = np.zeros((t.size, t.size))
#Assign 1. to diagonal terms
kuu[np.diag_indices(t.size)] = 1.
#Upper triangular indices
indtri1, indtri2 = np.triu_indices(t.size, 1)
#Block Diagonal indices among Upper Triangular indices
ind = np.where(index[indtri1] == index[indtri2])
indr = indtri1[ind]
indc = indtri2[ind]
r = t[indr] - t[indc]
r2 = r*r
#Calculation of  covariance function
kuu[indr, indc] = np.exp(-r2/lq2[index[indr]])
#Completation of lower triangular part
kuu[indc, indr] = kuu[indr, indc]
return kuu

def _Kusu(self, X, index, X2, index2):
index = index.reshape(index.size,)
index2 = index2.reshape(index2.size,)
t = X[:, 0].reshape(X.shape[0],1)
t2 = X2[:, 0].reshape(1,X2.shape[0])
lq = self.lengthscale.values.reshape(self.rank,)
#Covariance matrix initialization
kuu = np.zeros((t.size, t2.size))
for q in range(self.rank):
ind1 = index == q
ind2 = index2 == q
r = t[ind1]/lq[q] - t2[0,ind2]/lq[q]
r2 = r*r
#Calculation of  covariance function
kuu[np.ix_(ind1, ind2)] = np.exp(-r2)
return kuu

#Evaluation of cross-covariance function
def _Kfu(self, X, index, X2, index2):
#terms that move along t
t = X[:, 0].reshape(X.shape[0], 1)
d = np.unique(index) #Output Indexes
B = self.B.values[d]
C = self.C.values[d]
S = self.W.values[d, :]
#Index transformation
indd = np.arange(self.output_dim)
indd[d] = np.arange(d.size)
index = indd[index]
#Check where wd becomes complex
wbool = C*C >= 4.*B
#Output related variables must be column-wise
C = C.reshape(C.size, 1)
B = B.reshape(B.size, 1)
C2 = C*C
#Input related variables must be row-wise
z = X2[:, 0].reshape(1, X2.shape[0])
lq = self.lengthscale.values.reshape((1, self.rank))
#print np.max(z), np.max(z/lq[0, index2])
alpha = .5*C

wbool2 = wbool[index]
ind2t = np.where(wbool2)
ind3t = np.where(np.logical_not(wbool2))

kfu = np.empty((t.size, z.size))

indD = np.arange(B.size)
#(1) when wd is real
if np.any(np.logical_not(wbool)):
#Indexes of index and t related to (2)
t1 = t[ind3t]
ind = index[ind3t]
#Index transformation
d = np.asarray(np.where(np.logical_not(wbool))[0])
indd = indD.copy()
indd[d] = np.arange(d.size)
ind = indd[ind]
#Dx1 terms
w = .5*np.sqrt(4.*B[d] - C2[d])

#DxQ terms
Slq = (S[d]/w)*(.5*lq)
c0 = Slq*np.sqrt(np.pi)
nu = gam*(.5*lq)
#1xM terms
z_lq = z/lq[0, index2]
#NxQ terms
t_lq = t1/lq
#NxM terms
zt_lq = z_lq - t_lq[:, index2]

# Upsilon Calculations
#Using wofz
tz = t1-z
fullind = np.ix_(ind, index2)
zt_lq2 = -zt_lq*zt_lq
z_lq2 = -z_lq*z_lq
gamt = -gam[ind]*t1

upsi = - np.exp(z_lq2 + gamt + np.log(wofz(1j*(z_lq + nu[fullind]))))
z1 = zt_lq + nu[fullind]
indv1 = np.where(z1.real >= 0.)
indv2 = np.where(z1.real < 0.)
if indv1[0].shape > 0:
upsi[indv1] += np.exp(zt_lq2[indv1] + np.log(wofz(1j*z1[indv1])))
if indv2[0].shape > 0:
nua2 = nu[ind[indv2[0]], index2[indv2[1]]]**2
upsi[indv2] += np.exp(nua2 - gam[ind[indv2[0]], 0]*tz[indv2] + np.log(2.))\
- np.exp(zt_lq2[indv2] + np.log(wofz(-1j*z1[indv2])))
upsi[t1[:, 0] == 0., :] = 0.

#Covariance calculation
kfu[ind3t] = c0[fullind]*upsi.imag

#(2) when wd is complex
if np.any(wbool):
#Indexes of index and t related to (2)
t1 = t[ind2t]
ind = index[ind2t]
#Index transformation
d = np.asarray(np.where(wbool)[0])
indd = indD.copy()
indd[d] = np.arange(d.size)
ind = indd[ind]
#Dx1 terms
w = .5*np.sqrt(C2[d] - 4.*B[d])
#DxQ terms
Slq = S[d]*(lq*.25)
c0 = -Slq*(np.sqrt(np.pi)/w)
nu = gam*(lq*.5)
nuc = gamc*(lq*.5)
#1xM terms
z_lq = z/lq[0, index2]
#NxQ terms
t_lq = t1/lq[0, index2]
#NxM terms
zt_lq = z_lq - t_lq

# Upsilon Calculations
tz = t1-z
z_lq2 = -z_lq*z_lq
zt_lq2 = -zt_lq*zt_lq
gamt = -gam[ind]*t1
gamct = -gamc[ind]*t1
fullind = np.ix_(ind, index2)
upsi = np.exp(z_lq2 + gamt + np.log(wofz(1j*(z_lq + nu[fullind])).real))\
- np.exp(z_lq2 + gamct + np.log(wofz(1j*(z_lq + nuc[fullind])).real))

z1 = zt_lq + nu[fullind]
indv1 = np.where(z1 >= 0.)
indv2 = np.where(z1 < 0.)
if indv1[0].shape > 0:
upsi[indv1] -= np.exp(zt_lq2[indv1] + np.log(wofz(1j*z1[indv1]).real))
if indv2[0].shape > 0:
nua2 = nu[ind[indv2[0]], index2[indv2[1]]]**2
upsi[indv2] -= np.exp(nua2 - gam[ind[indv2[0]], 0]*tz[indv2] + np.log(2.))\
- np.exp(zt_lq2[indv2] + np.log(wofz(-1j*z1[indv2]).real))
z1 = zt_lq + nuc[fullind]
indv1 = np.where(z1 >= 0.)
indv2 = np.where(z1 < 0.)
if indv1[0].shape > 0:
upsi[indv1] += np.exp(zt_lq2[indv1] + np.log(wofz(1j*z1[indv1]).real))
if indv2[0].shape > 0:
nuac2 = nuc[ind[indv2[0]], index2[indv2[1]]]**2
upsi[indv2] += np.exp(nuac2 - gamc[ind[indv2[0]], 0]*tz[indv2] + np.log(2.))\
- np.exp(zt_lq2[indv2] + np.log(wofz(-1j*z1[indv2]).real))
upsi[t1[:, 0] == 0., :] = 0.

kfu[ind2t] = c0[np.ix_(ind, index2)]*upsi
return kfu

def _gkuu_lq(self, X, index):
t = X[:, 0].reshape(X.shape[0],)
index = index.reshape(X.shape[0],)
lq = self.lengthscale.values.reshape(self.rank,)
lq2 = lq*lq
#Covariance matrix initialization
glq = np.zeros((t.size, t.size))
#Upper triangular indices
indtri1, indtri2 = np.triu_indices(t.size, 1)
#Block Diagonal indices among Upper Triangular indices
ind = np.where(index[indtri1] == index[indtri2])
indr = indtri1[ind]
indc = indtri2[ind]
r = t[indr] - t[indc]
r2 = r*r
r2_lq2 = r2/lq2[index[indr]]
#Calculation of  covariance function
er2_lq2 = np.exp(-r2_lq2)
c = 2.*r2_lq2/lq[index[indr]]
glq[indr, indc] = er2_lq2*c
#Complete the lower triangular
glq[indc, indr] = glq[indr, indc]
return glq

#Be careful this derivative should be transpose it
def _gkuu_X(self, X, index): #Diagonal terms are always zero
t = X[:, 0].reshape(X.shape[0],)
index = index.reshape(index.size,)
lq = self.lengthscale.values.reshape(self.rank,)
lq2 = lq*lq
#Covariance matrix initialization
gt = np.zeros((t.size, t.size))
#Upper triangular indices
indtri1, indtri2 = np.triu_indices(t.size, 1) #Offset of 1 from the diagonal
#Block Diagonal indices among Upper Triangular indices
ind = np.where(index[indtri1] == index[indtri2])
indr = indtri1[ind]
indc = indtri2[ind]
r = t[indr] - t[indc]
r2 = r*r
r2_lq2 = r2/(-lq2[index[indr]])
#Calculation of  covariance function
er2_lq2 = np.exp(r2_lq2)
c = 2.*r/lq2[index[indr]]
gt[indr, indc] = er2_lq2*c
#Complete the lower triangular
gt[indc, indr] = -gt[indr, indc]
return gt

def _gkdiag(self, X, index):
index = index.reshape(index.size,)
#terms that move along t
d = np.unique(index)
B = self.B[d].values
C = self.C[d].values
S = self.W[d, :].values
#Index transformation
indd = np.arange(self.output_dim)
indd[d] = np.arange(d.size)
index = indd[index]
#Check where wd becomes complex
wbool = C*C >= 4.*B
#Output related variables must be column-wise
t = X[:, 0].reshape(X.shape[0], 1)
B = B.reshape(B.size, 1)
C = C.reshape(C.size, 1)
alpha = .5*C
C2 = C*C
S2 = S*S

wbool2 = wbool[index]
ind2t = np.where(wbool2)
ind3t = np.where(np.logical_not(wbool2))

#Input related variables must be row-wise
lq = self.lengthscale.values.reshape(1, self.rank)
lq2 = lq*lq

gB = np.empty((t.size, lq.size))
gC = np.empty((t.size, lq.size))
glq = np.empty((t.size, lq.size))
gS = np.empty((t.size, lq.size))

indD = np.arange(B.size)
#(1) When wd is real
if np.any(np.logical_not(wbool)):
#Indexes of index and t related to (1)
t1 = t[ind3t]
ind = index[ind3t]
#Index transformation
d = np.asarray(np.where(np.logical_not(wbool))[0])
indd = indD.copy()
indd[d] = np.arange(d.size)
ind = indd[ind]
#Dx1 terms
S2lq = S2[d]*(.5*lq)
c0 = S2lq*np.sqrt(np.pi)

w = .5*np.sqrt(4.*B[d] - C2[d])
w2 = w*w
gam2 = gam*gam
c2 = 0.5/gam
c = c1 - c2

#DxQ terms
c0 = c0/w2
nu = (.5*lq)*gam
#Nx1 terms
gamt = -gam[ind]*t1
gamct = -gamc[ind]*t1
egamt = np.exp(gamt)
egamct = np.exp(gamct)
ec = egamt*c2[ind] - egamct*c1[ind]

#NxQ terms
t_lq = t1/lq
t2_lq2 = -t_lq*t_lq
t_lq2 = t_lq/lq

et2_lq2 = np.exp(t2_lq2)
etlq2gamt = np.exp(t2_lq2 + gamt)

##Upsilon calculations
#Using wofz
wnu = wofz(1j*nu)
lwnu = np.log(wnu)
t2_lq2 = -t_lq*t_lq
upm = wnu[ind] - np.exp(t2_lq2 + gamt + np.log(wofz(1j*(t_lq + nu[ind]))))
upm[t1[:, 0] == 0, :] = 0.

nu2 = nu*nu
z1 = nu[ind] - t_lq
indv1 = np.where(z1.real >= 0.)
indv2 = np.where(z1.real < 0.)
upv = -np.exp(lwnu[ind] + gamt)
if indv1[0].shape > 0:
upv[indv1] += np.exp(t2_lq2[indv1] + np.log(wofz(1j*z1[indv1])))
if indv2[0].shape > 0:
upv[indv2] += np.exp(nu2[ind[indv2[0]], indv2[1]] + gamt[indv2[0], 0] + np.log(2.))\
- np.exp(t2_lq2[indv2] + np.log(wofz(-1j*z1[indv2])))
upv[t1[:, 0] == 0, :] = 0.

Slq = S[d]*lq #For grad wrt S
c0_S = Slq*np.sqrt(np.pi)/w2
K01 = c0_S*c

gS[ind3t] = np.real(K01[ind]*upm) + np.real((c0_S[ind]*ec)*upv)

#For B and C
upmd = etlq2gamt - 1.
upvd = egamt - et2_lq2

dw_dB = 0.5/w
dgam_dB = 1j*dw_dB

Ba1 = c0*(0.5*dgam_dB/gam2 + (0.5*lq2*gam*dgam_dB - 2.*dw_dB/w)*c)
Ba2_1 = c0*(dgam_dB*(0.5/gam2 - 0.25*lq2) + dw_dB/(w*gam))
Ba2_2 = c0*dgam_dB/gam
Ba4_1 = (S2lq*lq)*dgam_dB/w2
Ba4 = Ba4_1*c

gB[ind3t] = np.real(Ba1[ind]*upm) - np.real(((Ba2_1[ind] + Ba2_2[ind]*t1)*egamt - Ba3[ind]*egamct)*upv)\
+ np.real(Ba4[ind]*upmd) + np.real((Ba4_1[ind]*ec)*upvd)

dgam_dC = 0.5 + 1j*dw_dC

Ca1 = c0*(-0.25/alpha2 + 0.5*dgam_dC/gam2 + (0.5*lq2*gam*dgam_dC - 2.*dw_dC/w)*c)
Ca2_1 = c0*(dgam_dC*(0.5/gam2 - 0.25*lq2) + dw_dC/(w*gam))
Ca2_2 = c0*dgam_dC/gam
Ca4_1 = (S2lq*lq)*dgam_dC/w2
Ca4 = Ca4_1*c

gC[ind3t] = np.real(Ca1[ind]*upm) - np.real(((Ca2_1[ind] + Ca2_2[ind]*t1)*egamt - (Ca3_1[ind] + Ca3_2[ind]*t1)*egamct)*upv)\
+ np.real(Ca4[ind]*upmd) + np.real((Ca4_1[ind]*ec)*upvd)

#DxQ terms
la = (1./lq + nu*gam)*c0
la1 = la*c

c0l = (S2[d]/w2)*lq
la3 = c0l*c
gam_2 = .5*gam
glq[ind3t] = (la1[ind]*upm).real + ((la[ind]*ec)*upv).real\
+ (la3[ind]*(-gam_2[ind] + etlq2gamt*(-t_lq2 + gam_2[ind]))).real\
+ ((c0l[ind]*ec)*(-et2_lq2*(t_lq2 + gam_2[ind]) + egamt*gam_2[ind])).real

#(2) When w_d is complex
if np.any(wbool):
t1 = t[ind2t]
ind = index[ind2t]
#Index transformation
d = np.asarray(np.where(wbool)[0])
indd = indD.copy()
indd[d] = np.arange(d.size)
ind = indd[ind]
#Dx1 terms
S2lq = S2[d]*(.25*lq)
c0 = S2lq*np.sqrt(np.pi)
w = .5*np.sqrt(C2[d]-4.*B[d])
w2 = -w*w
gam2 = gam*gam
gamc2 = gamc*gamc
c21 = .5/gam
c22 = .5/gamc
c = c1 - c21
c2 = c1 - c22
#DxQ terms
c0 = c0/w2
nu = .5*lq*gam
nuc = .5*lq*gamc

#Nx1 terms
gamt = -gam[ind]*t1
gamct = -gamc[ind]*t1
egamt = np.exp(gamt)
egamct = np.exp(gamct)
ec = egamt*c21[ind] - egamct*c1[ind]
ec2 = egamct*c22[ind] - egamt*c1[ind]
#NxQ terms
t_lq = t1/lq
t2_lq2 = -t_lq*t_lq

et2_lq2 = np.exp(t2_lq2)
etlq2gamct = np.exp(t2_lq2 + gamct)
etlq2gamt = np.exp(t2_lq2 + gamt)

#Upsilon Calculations using wofz
t2_lq2 = -t_lq*t_lq #Required when using wofz
wnu = np.real(wofz(1j*nu))
lwnu = np.log(wnu)

upm = wnu[ind] - np.exp(t2_lq2 + gamt + np.log(wofz(1j*(t_lq + nu[ind])).real))
upm[t1[:, 0] == 0., :] = 0.

nu2 = nu*nu
z1 = nu[ind] - t_lq
indv1 = np.where(z1 >= 0.)
indv2 = np.where(z1 < 0.)
upv = -np.exp(lwnu[ind] + gamt)
if indv1[0].shape > 0:
upv[indv1] += np.exp(t2_lq2[indv1] + np.log(wofz(1j*z1[indv1]).real))
if indv2[0].shape > 0:
upv[indv2] += np.exp(nu2[ind[indv2[0]], indv2[1]] + gamt[indv2[0], 0] + np.log(2.)) - np.exp(t2_lq2[indv2]\
+ np.log(wofz(-1j*z1[indv2]).real))
upv[t1[:, 0] == 0, :] = 0.

wnuc = wofz(1j*nuc).real
upmc = wnuc[ind] - np.exp(t2_lq2 + gamct + np.log(wofz(1j*(t_lq + nuc[ind])).real))
upmc[t1[:, 0] == 0., :] = 0.

lwnuc = np.log(wnuc)
nuc2 = nuc*nuc
z1 = nuc[ind] - t_lq
indv1 = np.where(z1 >= 0.)
indv2 = np.where(z1 < 0.)
upvc = -np.exp(lwnuc[ind] + gamct)
if indv1[0].shape > 0:
upvc[indv1] += np.exp(t2_lq2[indv1] + np.log(wofz(1j*z1[indv1]).real))
if indv2[0].shape > 0:
upvc[indv2] += np.exp(nuc2[ind[indv2[0]], indv2[1]] + gamct[indv2[0], 0] + np.log(2.)) - np.exp(t2_lq2[indv2]\
+ np.log(wofz(-1j*z1[indv2]).real))
upvc[t1[:, 0] == 0, :] = 0.

#NxQ terms
c0_S = (S[d]/w2)*(lq*(np.sqrt(np.pi)*.5))

K011 = c0_S*c
K012 = c0_S*c2

gS[ind2t] = K011[ind]*upm + K012[ind]*upmc + (c0_S[ind]*ec)*upv + (c0_S[ind]*ec2)*upvc

#Is required to cache this, C gradient also required them
upmd = -1. + etlq2gamt
upvd = -et2_lq2 + egamt
upmdc = -1. + etlq2gamct
upvdc = -et2_lq2 + egamct

dgam_dB = 0.5/w
dgamc_dB = -dgam_dB

Ba1 = c0*(0.5*dgam_dB/gam2 + (0.5*lq2*gam*dgam_dB - 1./w2)*c)
Ba4_1 = (S2lq*lq)*dgam_dB/w2
Ba4 = Ba4_1*c
Ba2_1 = c0*(dgam_dB*(0.5/gam2 - 0.25*lq2) + 0.5/(w2*gam))
Ba2_2 = c0*dgam_dB/gam

Ba1c = c0*(0.5*dgamc_dB/gamc2 + (0.5*lq2*gamc*dgamc_dB - 1./w2)*c2)
Ba4_1c = (S2lq*lq)*dgamc_dB/w2
Ba4c = Ba4_1c*c2
Ba2_1c = c0*(dgamc_dB*(0.5/gamc2 - 0.25*lq2) + 0.5/(w2*gamc))
Ba2_2c = c0*dgamc_dB/gamc

gB[ind2t] = Ba1[ind]*upm - ((Ba2_1[ind] + Ba2_2[ind]*t1)*egamt - Ba3[ind]*egamct)*upv\
+ Ba4[ind]*upmd + (Ba4_1[ind]*ec)*upvd\
+ Ba1c[ind]*upmc - ((Ba2_1c[ind] + Ba2_2c[ind]*t1)*egamct - Ba3c[ind]*egamt)*upvc\
+ Ba4c[ind]*upmdc + (Ba4_1c[ind]*ec2)*upvdc

dgam_dC = 0.5 - dw_dC
dgamc_dC = 0.5 + dw_dC
S2lq2 = S2lq*lq

Ca1 = c0*(-0.25/alpha2 + 0.5*dgam_dC/gam2 + (0.5*lq2*gam*dgam_dC + alphad/w2)*c)
Ca2_1 = c0*(dgam_dC*(0.5/gam2 - 0.25*lq2) - 0.5*alphad/(w2*gam))
Ca2_2 = c0*dgam_dC/gam
Ca3_1 = c0*(0.25/alpha2 - 0.25*lq2*gam*dgam_dC/alphad - 0.5/w2)
Ca4_1 = S2lq2*(dgam_dC/w2)
Ca4 = Ca4_1*c

Ca1c = c0*(-0.25/alpha2 + 0.5*dgamc_dC/gamc2 + (0.5*lq2*gamc*dgamc_dC + alphad/w2)*c2)
Ca2_1c = c0*(dgamc_dC*(0.5/gamc2 - 0.25*lq2) - 0.5*alphad/(w2*gamc))
Ca2_2c = c0*dgamc_dC/gamc
Ca3_1c = c0*(0.25/alpha2 - 0.25*lq2*gamc*dgamc_dC/alphad - 0.5/w2)
Ca4_1c = S2lq2*(dgamc_dC/w2)
Ca4c = Ca4_1c*c2

gC[ind2t] = Ca1[ind]*upm - ((Ca2_1[ind] + Ca2_2[ind]*t1)*egamt - (Ca3_1[ind] + Ca3_2[ind]*t1)*egamct)*upv\
+ Ca4[ind]*upmd + (Ca4_1[ind]*ec)*upvd\
+ Ca1c[ind]*upmc - ((Ca2_1c[ind] + Ca2_2c[ind]*t1)*egamct - (Ca3_1c[ind] + Ca3_2c[ind]*t1)*egamt)*upvc\
+ Ca4c[ind]*upmdc + (Ca4_1c[ind]*ec2)*upvdc

#DxQ terms
la = (1./lq + nu*gam)*c0
lac = (1./lq + nuc*gamc)*c0
la1 = la*c
la1c = lac*c2
t_lq2 = t_lq/lq
c0l = (S2[d]/w2)*(.5*lq)
la3 = c0l*c
la3c = c0l*c2
gam_2 = .5*gam
gamc_2 = .5*gamc
glq[ind2t] = la1c[ind]*upmc + (lac[ind]*ec2)*upvc\
+ la3c[ind]*(-gamc_2[ind] + etlq2gamct*(-t_lq2 + gamc_2[ind]))\
+ (c0l[ind]*ec2)*(-et2_lq2*(t_lq2 + gamc_2[ind]) + egamct*gamc_2[ind])\
+ la1[ind]*upm + (la[ind]*ec)*upv\
+ la3[ind]*(-gam_2[ind] + etlq2gamt*(-t_lq2 + gam_2[ind]))\
+ (c0l[ind]*ec)*(-et2_lq2*(t_lq2 + gam_2[ind]) + egamt*gam_2[ind])

return glq, gS, gB, gC

def _gkfu(self, X, index, Z, index2):
index = index.reshape(index.size,)
#TODO: reduce memory usage
#terms that move along t
d = np.unique(index)
B = self.B[d].values
C = self.C[d].values
S = self.W[d, :].values
#Index transformation
indd = np.arange(self.output_dim)
indd[d] = np.arange(d.size)
index = indd[index]
#Check where wd becomes complex
wbool = C*C >= 4.*B
#t column
t = X[:, 0].reshape(X.shape[0], 1)
C = C.reshape(C.size, 1)
B = B.reshape(B.size, 1)
C2 = C*C
#z row
z = Z[:, 0].reshape(1, Z.shape[0])
index2 = index2.reshape(index2.size,)
lq = self.lengthscale.values.reshape((1, self.rank))
lq2 = lq*lq

alpha = .5*C

wbool2 = wbool[index]
ind2t = np.where(wbool2)
ind3t = np.where(np.logical_not(wbool2))
#kfu = np.empty((t.size, z.size))
glq = np.empty((t.size, z.size))
gSdq = np.empty((t.size, z.size))
gB = np.empty((t.size, z.size))
gC = np.empty((t.size, z.size))

indD = np.arange(B.size)
#(1) when wd is real
if np.any(np.logical_not(wbool)):
#Indexes of index and t related to (2)
t1 = t[ind3t]
ind = index[ind3t]
#Index transformation
d = np.asarray(np.where(np.logical_not(wbool))[0])
indd = indD.copy()
indd[d] = np.arange(d.size)
ind = indd[ind]
#Dx1 terms
w = .5*np.sqrt(4.*B[d] - C2[d])
gam_2 = .5*gam
S_w = S[d]/w
S_wpi = S_w*(.5*np.sqrt(np.pi))
#DxQ terms
c0 = S_wpi*lq #lq*Sdq*sqrt(pi)/(2w)
nu = gam*lq
nu2 = 1.+.5*(nu*nu)
nu *= .5

#1xM terms
z_lq = z/lq[0, index2]
z_lq2 = -z_lq*z_lq
#NxQ terms
t_lq = t1/lq
#DxM terms
gamt = -gam[ind]*t1
#NxM terms
zt_lq = z_lq - t_lq[:, index2]
zt_lq2 = -zt_lq*zt_lq
ezt_lq2 = -np.exp(zt_lq2)
ezgamt = np.exp(z_lq2 + gamt)

# Upsilon calculations
fullind = np.ix_(ind, index2)
upsi = - np.exp(z_lq2 + gamt + np.log(wofz(1j*(z_lq + nu[fullind]))))
tz = t1-z
z1 = zt_lq + nu[fullind]
indv1 = np.where(z1.real >= 0.)
indv2 = np.where(z1.real < 0.)
if indv1[0].shape > 0:
upsi[indv1] += np.exp(zt_lq2[indv1] + np.log(wofz(1j*z1[indv1])))
if indv2[0].shape > 0:
nua2 = nu[ind[indv2[0]], index2[indv2[1]]]**2
upsi[indv2] += np.exp(nua2 - gam[ind[indv2[0]], 0]*tz[indv2] + np.log(2.))\
- np.exp(zt_lq2[indv2] + np.log(wofz(-1j*z1[indv2])))
upsi[t1[:, 0] == 0., :] = 0.

#DxQ term
Sa1 = lq*(.5*np.sqrt(np.pi))/w

gSdq[ind3t] = Sa1[np.ix_(ind, index2)]*upsi.imag

la1 = S_wpi*nu2
la2 = S_w*lq
uplq = ezt_lq2*(gam_2[ind])
uplq += ezgamt*(-z_lq/lq[0, index2] + gam_2[ind])

glq[ind3t] = (la1[np.ix_(ind, index2)]*upsi).imag
glq[ind3t] += la2[np.ix_(ind, index2)]*uplq.imag

#Dx1 terms
dw_dB = .5/w
dgam_dB = -1j*dw_dB
#DxQ terms
Ba1 = -c0*dw_dB/w #DXQ
Ba2 = c0*dgam_dB #DxQ
Ba3 = lq2*gam_2 #DxQ
Ba4 = (dgam_dB*S_w)*(.5*lq2) #DxQ

gB[ind3t] = ((Ba1[np.ix_(ind, index2)] + Ba2[np.ix_(ind, index2)]*(Ba3[np.ix_(ind, index2)] - (t1-z)))*upsi).imag\
+ (Ba4[np.ix_(ind, index2)]*(ezt_lq2 + ezgamt)).imag

#Gradient wrt C (it uses some calculations performed in B)
#Dx1 terms
dgam_dC = 0.5 - 1j*dw_dC
#DxQ terms
Ca1 = -c0*dw_dC/w #DXQ
Ca2 = c0*dgam_dC #DxQ
Ca4 = (dgam_dC*S_w)*(.5*lq2) #DxQ

gC[ind3t] = ((Ca1[np.ix_(ind, index2)] + Ca2[np.ix_(ind, index2)]*(Ba3[np.ix_(ind, index2)] - (t1-z)))*upsi).imag\
+ (Ca4[np.ix_(ind, index2)]*(ezt_lq2 + ezgamt)).imag

#(2) when wd is complex
if np.any(wbool):
#Indexes of index and t related to (2)
t1 = t[ind2t]
ind = index[ind2t]
#Index transformation
d = np.asarray(np.where(wbool)[0])
indd = indD.copy()
indd[d] = np.arange(d.size)
ind = indd[ind]
#Dx1 terms
w = .5*np.sqrt(C2[d] - 4.*B[d])
w2 = w*w
#DxQ terms
S_w= -S[d]/w #minus is given by j*j
S_wpi = S_w*(.25*np.sqrt(np.pi))

c0 = S_wpi*lq
gam_2 = .5*gam
gamc_2 = .5*gamc
nu = gam*lq
nuc = gamc*lq
nu2 = 1.+.5*(nu*nu)
nuc2 = 1.+.5*(nuc*nuc)
nu *= .5
nuc *= .5
#1xM terms
z_lq = z/lq[0, index2]
z_lq2 = -z_lq*z_lq
#Nx1
gamt = -gam[ind]*t1
gamct = -gamc[ind]*t1
#NxQ terms
t_lq = t1/lq[0, index2]
#NxM terms
zt_lq = z_lq - t_lq
zt_lq2 = -zt_lq*zt_lq
ezt_lq2 = -np.exp(zt_lq2)
ezgamt = np.exp(z_lq2 + gamt)
ezgamct = np.exp(z_lq2 + gamct)

# Upsilon calculations
fullind = np.ix_(ind, index2)
upsi1 = - np.exp(z_lq2 + gamct + np.log(wofz(1j*(z_lq + nuc[fullind])).real))
tz = t1-z
z1 = zt_lq + nuc[fullind]
indv1 = np.where(z1 >= 0.)
indv2 = np.where(z1 < 0.)
if indv1[0].shape > 0:
upsi1[indv1] += np.exp(zt_lq2[indv1] + np.log(wofz(1j*z1[indv1]).real))
if indv2[0].shape > 0:
nuac2 = nuc[ind[indv2[0]], index2[indv2[1]]]**2
upsi1[indv2] += np.exp(nuac2 - gamc[ind[indv2[0]], 0]*tz[indv2] + np.log(2.))\
- np.exp(zt_lq2[indv2] + np.log(wofz(-1j*z1[indv2]).real))
upsi1[t1[:, 0] == 0., :] = 0.

upsi2 = - np.exp(z_lq2 + gamt + np.log(wofz(1j*(z_lq + nu[fullind])).real))
z1 = zt_lq + nu[fullind]
indv1 = np.where(z1 >= 0.)
indv2 = np.where(z1 < 0.)
if indv1[0].shape > 0:
upsi2[indv1] += np.exp(zt_lq2[indv1] + np.log(wofz(1j*z1[indv1]).real))
if indv2[0].shape > 0:
nua2 = nu[ind[indv2[0]], index2[indv2[1]]]**2
upsi2[indv2] += np.exp(nua2 - gam[ind[indv2[0]], 0]*tz[indv2] + np.log(2.))\
- np.exp(zt_lq2[indv2] + np.log(wofz(-1j*z1[indv2]).real))
upsi2[t1[:, 0] == 0., :] = 0.

la1 = S_wpi*nu2
la1c = S_wpi*nuc2
la2 = S_w*(.5*lq)
uplq = ezt_lq2*(gamc_2[ind]) + ezgamct*(-z_lq/lq[0, index2] + gamc_2[ind])\
- ezt_lq2*(gam_2[ind]) - ezgamt*(-z_lq/lq[0, index2] + gam_2[ind])

glq[ind2t] = la1c[np.ix_(ind, index2)]*upsi1 - la1[np.ix_(ind, index2)]*upsi2\
+ la2[np.ix_(ind, index2)]*uplq

Sa1 = (lq*(-.25*np.sqrt(np.pi)))/w

gSdq[ind2t] = Sa1[np.ix_(ind, index2)]*(upsi1 - upsi2)

#Dx1 terms
dgam_dB = .5/w
dgamc_dB = -dgam_dB
#DxQ terms
Ba1 = .5*(c0/w2)
Ba2 = c0*dgam_dB
Ba3 = lq2*gam_2
Ba4 = (dgam_dB*S_w)*(.25*lq2)

Ba2c = c0*dgamc_dB
Ba3c = lq2*gamc_2
Ba4c = (dgamc_dB*S_w)*(.25*lq2)

gB[ind2t] = (Ba1[np.ix_(ind, index2)] + Ba2c[np.ix_(ind, index2)]*(Ba3c[np.ix_(ind, index2)] - (t1-z)))*upsi1\
+ Ba4c[np.ix_(ind, index2)]*(ezt_lq2 + ezgamct)\
- (Ba1[np.ix_(ind, index2)] + Ba2[np.ix_(ind, index2)]*(Ba3[np.ix_(ind, index2)] - (t1-z)))*upsi2\
- Ba4[np.ix_(ind, index2)]*(ezt_lq2 + ezgamt)

#Dx1 terms
#DxQ terms
Ca2 = c0*dgam_dC
Ca4 = (dgam_dC*S_w)*(.25*lq2)

Ca2c = c0*dgamc_dC
Ca4c = (dgamc_dC*S_w)*(.25*lq2)

gC[ind2t] = (Ca1[np.ix_(ind, index2)] + Ca2c[np.ix_(ind, index2)]*(Ba3c[np.ix_(ind, index2)] - (t1-z)))*upsi1\
+ Ca4c[np.ix_(ind, index2)]*(ezt_lq2 + ezgamct)\
- (Ca1[np.ix_(ind, index2)] + Ca2[np.ix_(ind, index2)]*(Ba3[np.ix_(ind, index2)] - (t1-z)))*upsi2\
- Ca4[np.ix_(ind, index2)]*(ezt_lq2 + ezgamt)

return glq, gSdq, gB, gC

#TODO: reduce memory usage
def _gkfu_z(self, X, index, Z, index2): #Kfu(t,z)
index = index.reshape(index.size,)
#terms that move along t
d = np.unique(index)
B = self.B[d].values
C = self.C[d].values
S = self.W[d, :].values
#Index transformation
indd = np.arange(self.output_dim)
indd[d] = np.arange(d.size)
index = indd[index]
#Check where wd becomes complex
wbool = C*C >= 4.*B
wbool2 = wbool[index]
ind2t = np.where(wbool2)
ind3t = np.where(np.logical_not(wbool2))
#t column
t = X[:, 0].reshape(X.shape[0], 1)
C = C.reshape(C.size, 1)
B = B.reshape(B.size, 1)
C2 = C*C
alpha = .5*C
#z row
z = Z[:, 0].reshape(1, Z.shape[0])
index2 = index2.reshape(index2.size,)
lq = self.lengthscale.values.reshape((1, self.rank))

#kfu = np.empty((t.size, z.size))
gz = np.empty((t.size, z.size))
indD = np.arange(B.size)
#(1) when wd is real
if np.any(np.logical_not(wbool)):
#Indexes of index and t related to (2)
t1 = t[ind3t]
ind = index[ind3t]
#TODO: Find a better way of doing this
#Index transformation
d = np.asarray(np.where(np.logical_not(wbool))[0])
indd = indD.copy()
indd[d] = np.arange(d.size)
ind = indd[ind]
#Dx1 terms
w = .5*np.sqrt(4.*B[d] - C2[d])
S_w = S[d]/w
S_wpi =S_w*(.5*np.sqrt(np.pi))
#DxQ terms
c0 = S_wpi*lq #lq*Sdq*sqrt(pi)/(2w)
nu = (.5*gam)*lq

#1xM terms
z_lq = z/lq[0, index2]
z_lq2 = -z_lq*z_lq
#NxQ terms
t_lq = t1/lq
#DxM terms
gamt = -gam[ind]*t1
#NxM terms
zt_lq = z_lq - t_lq[:, index2]
zt_lq2 = -zt_lq*zt_lq
#ezt_lq2 = -np.exp(zt_lq2)
ezgamt = np.exp(z_lq2 + gamt)

# Upsilon calculations
fullind = np.ix_(ind, index2)
upsi = - np.exp(z_lq2 + gamt + np.log(wofz(1j*(z_lq + nu[fullind]))))
tz = t1-z
z1 = zt_lq + nu[fullind]
indv1 = np.where(z1.real >= 0.)
indv2 = np.where(z1.real < 0.)
if indv1[0].shape > 0:
upsi[indv1] += np.exp(zt_lq2[indv1] + np.log(wofz(1j*z1[indv1])))
if indv2[0].shape > 0:
nua2 = nu[ind[indv2[0]], index2[indv2[1]]]**2
upsi[indv2] += np.exp(nua2 - gam[ind[indv2[0]], 0]*tz[indv2] + np.log(2.))\
- np.exp(zt_lq2[indv2] + np.log(wofz(-1j*z1[indv2])))
upsi[t1[:, 0] == 0., :] = 0.

za1 = c0*gam
#za2 = S_w
gz[ind3t] = (za1[np.ix_(ind, index2)]*upsi).imag + S_w[np.ix_(ind, index2)]*ezgamt.imag

#(2) when wd is complex
if np.any(wbool):
#Indexes of index and t related to (2)
t1 = t[ind2t]
ind = index[ind2t]
#Index transformation
d = np.asarray(np.where(wbool)[0])
indd = indD.copy()
indd[d] = np.arange(d.size)
ind = indd[ind]
#Dx1 terms
w = .5*np.sqrt(C2[d] - 4.*B[d])
#DxQ terms
S_w = -S[d]/w #minus is given by j*j
S_wpi = S_w*(.25*np.sqrt(np.pi))
c0 = S_wpi*lq
nu = .5*gam*lq
nuc = .5*gamc*lq

#1xM terms
z_lq = z/lq[0, index2]
z_lq2 = -z_lq*z_lq
#Nx1
gamt = -gam[ind]*t1
gamct = -gamc[ind]*t1
#NxQ terms
t_lq = t1/lq
#NxM terms
zt_lq = z_lq - t_lq[:, index2]
ezgamt = np.exp(z_lq2 + gamt)
ezgamct = np.exp(z_lq2 + gamct)

# Upsilon calculations
zt_lq2 = -zt_lq*zt_lq
fullind = np.ix_(ind, index2)
upsi1 = - np.exp(z_lq2 + gamct + np.log(wofz(1j*(z_lq + nuc[fullind])).real))
tz = t1-z
z1 = zt_lq + nuc[fullind]
indv1 = np.where(z1 >= 0.)
indv2 = np.where(z1 < 0.)
if indv1[0].shape > 0:
upsi1[indv1] += np.exp(zt_lq2[indv1] + np.log(wofz(1j*z1[indv1]).real))
if indv2[0].shape > 0:
nuac2 = nuc[ind[indv2[0]], index2[indv2[1]]]**2
upsi1[indv2] += np.exp(nuac2 - gamc[ind[indv2[0]], 0]*tz[indv2] + np.log(2.))\
- np.exp(zt_lq2[indv2] + np.log(wofz(-1j*z1[indv2]).real))
upsi1[t1[:, 0] == 0., :] = 0.

upsi2 = - np.exp(z_lq2 + gamt + np.log(wofz(1j*(z_lq + nu[fullind])).real))
z1 = zt_lq + nu[fullind]
indv1 = np.where(z1 >= 0.)
indv2 = np.where(z1 < 0.)
if indv1[0].shape > 0:
upsi2[indv1] += np.exp(zt_lq2[indv1] + np.log(wofz(1j*z1[indv1]).real))
if indv2[0].shape > 0:
nua2 = nu[ind[indv2[0]], index2[indv2[1]]]**2
upsi2[indv2] += np.exp(nua2 - gam[ind[indv2[0]], 0]*tz[indv2] + np.log(2.))\
- np.exp(zt_lq2[indv2] + np.log(wofz(-1j*z1[indv2]).real))
upsi2[t1[:, 0] == 0., :] = 0.