Skip to content

Commit

Permalink
Fix diagonal estimate
Browse files Browse the repository at this point in the history
  • Loading branch information
ppegolo committed Aug 9, 2024
1 parent dcaa678 commit 9b09ea6
Show file tree
Hide file tree
Showing 3 changed files with 395 additions and 200 deletions.
386 changes: 225 additions & 161 deletions examples/07_example_negative_log_likelihood.ipynb

Large diffs are not rendered by default.

47 changes: 35 additions & 12 deletions sportran/current/current.py
Original file line number Diff line number Diff line change
Expand Up @@ -405,14 +405,15 @@ def maxlike_estimate(self,
likelihood = 'wishart',
solver = 'BFGS',
guess_runave_window = 50,
minimize_kwargs={}
):

if likelihood.lower() == 'wishart':
data = self.cospectrum.real # TODO figure out this business of real vs abs value
data = self.cospectrum.real * self.N_CURRENTS
elif likelihood.lower() == 'chisquare' or likelihood.lower() == 'chisquared':
data = self.psd
data = self.psd * self.N_CURRENTS
elif likelihood == 'variancegamma' or likelihood == 'variance-gamma':
data = self.cospectrum.real[0,1]
data = self.cospectrum.real[0,1] # * self.N_CURRENTS
else:
raise ValueError("Likelihood must be Wishart, Chi-square, or Variance-Gamma.")

Expand All @@ -429,7 +430,8 @@ def maxlike_estimate(self,
# Minimize the negative log-likelihood with a fixed number of parameters
self.maxlike.maxlike(mask = mask,
guess_runave_window = guess_runave_window,
n_parameters = n_parameters)
n_parameters = n_parameters,
minimize_kwargs=minimize_kwargs)

elif isinstance(n_parameters, list) or isinstance(n_parameters, np.ndarray):
do_AIC = True
Expand All @@ -452,7 +454,8 @@ def maxlike_estimate(self,
guess_runave_window = guess_runave_window,
n_parameters = int(n_par),
omega_fixed = None,
write_log = False)
write_log = False,
minimize_kwargs=minimize_kwargs)

_aic.append(self.maxlike.log_likelihood_value - n_par)
_filters.append(deepcopy(self.maxlike))
Expand All @@ -466,21 +469,41 @@ def maxlike_estimate(self,
params = self.maxlike.parameters_mean
params_std = self.maxlike.parameters_std

om = self.maxlike.omega
if likelihood == 'wishart':
self.NLL_spline = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params)(x)))
# nw = params.shape[0]//2
# re_NLL_spline = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params[:nw])(x)))
# im_NLL_spline = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params[nw:])(x)))
# self.NLL_mean = re_NLL_spline(om) + 1j*im_NLL_spline(om)
# _spl = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params)(x)))
# self.NLL_mean = _spl(om)
from sportran.md.maxlike import scale_matrix
self.NLL_mean = scale_matrix(model, params, om, omega_fixed, self.N_CURRENTS) * self.N_EQUIV_COMPONENTS / self.N_CURRENTS
try:
self.NLL_spline_upper = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params + params_std)(x)))
self.NLL_spline_lower = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params - params_std)(x)))
# _NLL_spline_upper = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params + params_std)(x)))
# _NLL_spline_lower = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params - params_std)(x)))
# re_NLL_spline_upper = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params[:nw] + params_std[:nw])(x)))
# re_NLL_spline_lower = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params[:nw] - params_std[:nw])(x)))
# im_NLL_spline_upper = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params[nw:] + params_std[nw:])(x)))
# im_NLL_spline_lower = lambda x: np.einsum('wba,wbc->wac', *(lambda y: (y, y))(model(omega_fixed, params[nw:] - params_std[nw:])(x)))
# self.NLL_upper = re_NLL_spline_upper(om) + 1j*im_NLL_spline_upper(om)
# self.NLL_lower = re_NLL_spline_lower(om) + 1j*im_NLL_spline_lower(om)
self.NLL_upper = scale_matrix(model, params+params_std, om, omega_fixed, self.N_CURRENTS) * self.N_EQUIV_COMPONENTS / self.N_CURRENTS
self.NLL_lower = scale_matrix(model, params-params_std, om, omega_fixed, self.N_CURRENTS) * self.N_EQUIV_COMPONENTS / self.N_CURRENTS
except TypeError:
pass
else:
self.NLL_spline = model(omega_fixed, params)
_NLL_spline = model(omega_fixed, params)
fact = np.sqrt(self.N_EQUIV_COMPONENTS) / self.N_CURRENTS / np.sqrt(2)
self.NLL_mean = _NLL_spline(om) * fact
try:
self.NLL_spline_upper = model(omega_fixed, params + params_std)
self.NLL_spline_lower = model(omega_fixed, params - params_std)
_NLL_spline_upper = model(omega_fixed, params + params_std)
_NLL_spline_lower = model(omega_fixed, params - params_std)
self.NLL_upper = _NLL_spline_upper(om) * fact
self.NLL_lower = _NLL_spline_lower(om) * fact
except TypeError:
pass

# self.estimate = self.maxlike.parameters_mean[0]*self.maxlike.factor

# try:
Expand Down
162 changes: 135 additions & 27 deletions sportran/md/maxlike.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,19 @@ def __repr__(self):
msg = 'MaxLikeFilter:\n'
return msg

def maxlike(self, data=None, model=None, n_parameters=None, likelihood=None, solver=None, mask=None, n_components=None, n_currents=None, guess_runave_window=50, omega_fixed=None, write_log=True):
def maxlike(self,
data=None,
model=None,
n_parameters=None,
likelihood=None,
solver=None,
mask=None,
n_components=None,
n_currents=None,
guess_runave_window=50,
omega_fixed=None,
write_log=True,
minimize_kwargs=None):
"""
Perform the maximum-likelihood estimation.
"""
Expand Down Expand Up @@ -160,10 +172,13 @@ def maxlike(self, data=None, model=None, n_parameters=None, likelihood=None, sol

# Minimize the negative log-likelihood
self._guess_data = guess_data
guess_par = guess_data #normalize_parameters(guess_data, guess_data)
# print(guess_data - denormalize_parameters(guess_par, guess_data))
res = minimize(fun = self.log_like,
x0 = guess_data,
x0 = guess_par,
args = (self.model, omega, self.omega_fixed, self.data, nu, ell),
method = self.solver)
method = self.solver,
**minimize_kwargs)

# Covariance of the parameters
try:
Expand All @@ -176,6 +191,7 @@ def maxlike(self, data=None, model=None, n_parameters=None, likelihood=None, sol
cov = None

self.parameters_mean = res.x
# self.parameters_mean = denormalize_parameters(res.x, self._guess_data)
if cov is not None:
self.parameters_std = np.sqrt(cov.diagonal())

Expand All @@ -199,56 +215,110 @@ def guess_data(self, data, omega, omega_fixed, ell, nu, window=10, loglike='wish

guess_data = np.array([guess_data[..., j] for j in [np.argmin(np.abs(omega - omega_fixed[i])) for i in range(len(omega_fixed))]])
print(guess_data.shape)

if loglike == 'wishart':
guess_data = np.array([cholesky(g, lower=False) for g in guess_data]) / np.sqrt(ell)

guess_data = np.array([cholesky(g, lower=False) for g in guess_data]) #/ np.sqrt(ell)
upper_triangle_indices = np.triu_indices(nu)
guess_data = guess_data[:, upper_triangle_indices[0], upper_triangle_indices[1]].T.reshape(-1)

# guess_data = np.array([guess_data[:, 0, 0], guess_data[:, 0, 1], guess_data[:, 1, 1]]).reshape(-1)

nw = omega_fixed.shape[0]
ie = 0
if guess_data.dtype == np.complex128:
guess_params = np.zeros((nw, nu**2))
for i, j in zip(*upper_triangle_indices):
if i == j:
guess_params[:, ie] = guess_data[:, i, j].real
ie += 1
else:
guess_params[:, ie] = guess_data[:, i, j].real
ie += 1
guess_params[:, ie] = guess_data[:, i, j].imag
ie += 1
guess_data = guess_params.flatten()
else:
guess_params = np.zeros((nw, nu*(nu+1)//2))
for i, j in zip(*upper_triangle_indices):
guess_params[:, ie] = guess_data[:, i, j]
ie += 1
guess_data = guess_params.flatten()

return guess_data
def log_likelihood_wishart(self, w, model, omega, omega_fixed, data_, nu, ell, eps = 1e-9):

def log_likelihood_wishart(self, w, model, omega, omega_fixed, data_, nu, ell, eps = 1e-3):
"""
Logarithm of the Wishart probability density function.
"""
n = ell
p = nu

# Compute scale matrix from the model (symmetrize to ensure positive definiteness)
spline = model(omega_fixed, w)
V = spline(omega)
V = opt_einsum.contract('wba,wbc->wac', V, V) / n

V = scale_matrix(model, w, omega, omega_fixed, p)
X = data_

if n < p:
# Singular Wishart
multig = multigammaln(0.5 * n, n)
S = np.linalg.svd(X, hermitian = True, compute_uv = False)
S = np.linalg.svd(X, hermitian=True, compute_uv=False)
detX = np.array([np.prod(s[abs(s) > eps]) for s in S])

else:
multig = multigammaln(0.5 * n, p)
detX = np.linalg.det(X)

invV = np.linalg.inv(V)
detV = np.linalg.det(V)

trinvV_X = opt_einsum.contract('wab,wba->w', invV, X)

# log_pdf = - (0.5 * (-n * p * LOG2 - n * np.log(detV) + (n - p - 1) * np.log(detX) - trinvV_X) - multig)
coeff_detV = -n
coeff_detX = n-p-1
log_pdf = coeff_detV * np.log(detV) + coeff_detX * np.log(detX) - trinvV_X
# print(-np.sum(log_pdf))

log_pdf = coeff_detV * np.log(detV + eps) + coeff_detX * np.log(detX + eps) - trinvV_X
tot = -np.sum(log_pdf)
return tot

def log_likelihood_complex_wishart(self, w, model, omega, omega_fixed, data_, nu, ell, eps = 1e-9):
"""
Logarithm of the Complex Wishart probability density function.
"""
n = ell
p = nu

# Compute scale matrix from the model (symmetrize to ensure positive definiteness)
S = scale_matrix(model, w, omega, omega_fixed, p)
# nw = w.shape[0]//2
# real_part = model(omega_fixed, w[:nw])
# imag_part = model(omega_fixed, w[nw:])

# # Lower Cholesky factor of S
# L = (real_part(omega) + 1j*imag_part(omega))
# S = opt_einsum.contract('wba,wbc->wac', L.conj(), L)

# The distribution refers to the sample covariance matrix of the (complex) multinormal vectors, not their average
X = data_

if n < p:
raise ValueError('n must be greater or equal than p')
# Singular Wishart
multig = multigammaln(0.5 * n, n)
S = np.linalg.svd(X, hermitian = True, compute_uv = False)
detX = np.array([np.prod(s[abs(s) > eps]) for s in S])
else:
# multig = multigammaln(0.5 * n, p)
logdetX = np.log(np.abs(np.linalg.det(X).real))

invS = np.linalg.inv(S)
logdetS = np.log(np.abs(np.linalg.det(S).real))

trinvS_X = opt_einsum.contract('wab,wba->w', invS, X).real

log_pdf = (n-p)*logdetX - trinvS_X - n*logdetS
# coeff_detV = -n
# coeff_detX = n-p-1
# log_pdf = coeff_detV * np.log(detV) + coeff_detX * np.log(detX) - trinvV_X
# # print(-np.sum(log_pdf))
return -np.sum(log_pdf)

def log_likelihood_offdiag(self, w, model, omega, omega_fixed, data_, nu, ell):
"""
Negative of the logarithm of the Variance-Gamma probability density function.
"""
print(w)
spline = model(omega_fixed, w)
rho = np.clip(spline(omega), -0.98, 0.98)
_alpha = 1 / (1 - rho**2)
Expand All @@ -257,12 +327,18 @@ def log_likelihood_offdiag(self, w, model, omega, omega_fixed, data_, nu, ell):
_gamma2 = _alpha**2 - _beta**2
_lambda_minus_half = _lambda - 0.5

# Non sono più sicuro sia sensata questa definizione di z. Non è semplicemtente data_? AH! Forse è la stessa cosa che succede al Chi2, va moltiplicato per il numero di dof. Capire meglio e fare prove.
z = data_ * ell * nu
# Non sono più sicuro sia sensata questa definizione di z.
# Non è semplicemtente data_? AH! Forse è la stessa cosa che succede al Chi2, va moltiplicato per il numero di dof. Capire meglio e fare prove.
z = data_ * nu * ell
absz = np.abs(z)
log_pdf = _lambda * np.log(_gamma2) + _lambda_minus_half * np.log(absz) + np.log(sp.kv(_lambda_minus_half, _alpha * absz)) + \
_beta * z - 0.5 * np.log(np.pi) - np.log(sp.gamma(_lambda)) - _lambda_minus_half * np.log(2 * _alpha)

term1 = _lambda * np.log(_gamma2)
term2 = _lambda_minus_half * np.log(absz)
term3 = np.log(sp.kv(_lambda_minus_half, _alpha * absz))
term4 = _beta * z
term5 = - _lambda_minus_half * np.log(2 * _alpha)
# log_pdf = _lambda * np.log(_gamma2) + _lambda_minus_half * np.log(absz) + np.log(sp.kv(_lambda_minus_half, _alpha * absz)) + \
# _beta * z - _lambda_minus_half * np.log(2 * _alpha) # + const
log_pdf = term1 + term2 + term3 + term4 + term5
return -np.sum(log_pdf)

def log_likelihood_diag(self, w, model, omega, omega_fixed, data_, nu, ell):
Expand Down Expand Up @@ -327,6 +403,38 @@ def log_posterior_normal(self, w, omega, omega_fixed, data, nu=6, ell=3):
return self.log_prior_offdiag(w) + self.log_likelihood_normal(w, omega, omega_fixed, data, nu, ell)


def scale_matrix(model, w, omega, omega_fixed, n):
'''
'''
elements = model(omega_fixed, w)(omega)
ie = 0
if elements.dtype == np.complex128:
L = np.zeros((n, n, omega.shape[0]), dtype = np.complex128)
for i, j in zip(*np.triu_indices(n)):
if i == j:
L[i, j] = elements[:, ie]
ie += 1
else:
L[i, j] = elements[:, ie] + 1j*elements[:, ie+1]
ie += 2

S = np.einsum('jiw,jkw->wik', L.conj(), L)
else:
L = np.zeros((n, n, omega.shape[0]))
for i, j in zip(*np.triu_indices(n)):
L[i, j] = elements[:, ie]
ie += 1

S = np.einsum('jiw,jkw->wik', L, L)
return S

def normalize_parameters(p, guess):
mini, maxi = 1.2*np.abs(guess), -1.2*np.abs(guess)
return (p-mini) / (maxi-mini)

def denormalize_parameters(p, guess):
mini, maxi = 1.2*np.abs(guess), -1.2*np.abs(guess)
return p * (maxi-mini) + mini

# # Methods to perform a bayesian estimation of the transport coefficients

Expand Down

0 comments on commit 9b09ea6

Please sign in to comment.