I already cut out any other self-defined functions and made it so that tall lists are gone so this should be the bare minimum, but i can not understand what goes wrong, all are either booleans, numpy arrays or scalars.
@njit(cache=True, parallel=True, fastmath=True)
def fFilter(mY, mX, vPsi, vClusterAssign, iG, bGroupSpecificVar=False, bGroupSpecificTheta=False, bExoRegressors=False, bIndividualFE=False, bGroupedFE=False):
breakpoint()
iN = mY.shape[0]
iP = mY.shape[1]
iT = mY.shape[2]
vFilteredTheta = np.zeros(iT)
# vTransPar = fTransPar(vPsi, mX, iG, iT, iN, iP, bGroupSpecificVar, bGroupSpecificTheta, bExoRegressors, bIndividualFE, bGroupedFE)
vTransPar = vPsi
vEta = vTransPar[:(bIndividualFE*iN + bGroupedFE*iG)*iP]
vAlpha = vTransPar[(bIndividualFE*iN + bGroupedFE*iG) * iP : (bIndividualFE*iN + bGroupedFE*iG) * iP + iP * iG * iT]
vTheta = vEta
lSigmas = vTransPar[(bIndividualFE*iN + bGroupedFE*iG) * iP + iP * iG * iT:]
for i in range(iG):
# Extract the elements for sigma_i
# mSigma_i = fUnvech(mSigma[i* (iP*(iP + 1) // 2):(i+1) * (iP*(iP + 1) // 2) ], iP)
vSigma_i = lSigmas[i* (iP*(iP + 1) // 2):(i+1) * (iP*(iP + 1) // 2) ]
mSigma_i = np.zeros((iP,iP))
mSigma_i[np.tril_indices(iP)] = vSigma_i
mSigma_i = mSigma_i @ mSigma_i.T
mAlpha = vAlpha.reshape(iG,iP,iT)
# lSigmas = vTransPar[1:]
if bIndividualFE:
mEta = vEta.reshape(iN, iP)
elif bGroupedFE:
mEta = vEta.reshape(iG, iP)
else:
mEta = np.zeros((iN, iP))
if bExoRegressors:
mTheta = vTheta.reshape((mX.shape[0]-1), iP, np.maximum(1 , (iG *bGroupSpecificTheta)))
else:
mTheta = np.zeros((mX.shape[0], iP, np.maximum(1 , (iG *bGroupSpecificTheta))))
# update with likelihood contribution
mLogLik = np.zeros((iN, iT))
for i in range(iN):
for t in range(iT):
mSigma = mSigma_i
mAlpha_g = mAlpha[vClusterAssign[i]-1]
mTheta_g = mTheta[:, :, ((not bGroupSpecificTheta)*0 + bGroupSpecificTheta*(vClusterAssign[i]-1)) * bExoRegressors]
mEta_g = mEta[(bIndividualFE or (not bIndividualFE and not bGroupedFE ))*i + bGroupedFE*(vClusterAssign[i]-1),:]
if np.isnan(mY[i,:,t]).any():
# if np.isnan(mY[i,:,t]).all():
# breakpoint()
mask = ~np.isnan(mY[i,:,t])
vY = mY[i,:,t][mask]
mAlpha_g = mAlpha_g[:, t][mask]
mSigma = mSigma[mask][:, mask]
mTheta_g = mTheta_g[:,mask]
mEta_g = mEta_g[mask]
else:
vY = mY[i,:,t]
mAlpha_g = mAlpha_g[:, t]
vLogLik = (-0.5* np.log(2*np.pi)
- 0.5* np.log(np.linalg.det(mSigma))
- 0.5 * (np.dot(np.dot((vY - mAlpha_g - mTheta_g.T @ mX[:,t] - mEta_g).T
,np.linalg.inv(mSigma) ),
(vY - mAlpha_g - mTheta_g.T @ mX[:,t] - mEta_g) ))
)
mLogLik[i,t] = vLogLik
#vFilteredTheta later maybe GAS update
# return filtered log-variances and likelihood values
return (vFilteredTheta, mLogLik)