Fix diagonal estimate

Author: ppegolo

examples/07_example_negative_log_likelihood.ipynb

@@ -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.")
 
@@ -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
@@ -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))
@@ -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:

sportran/current/current.py

@@ -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.
         """
@@ -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:
@@ -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())
 
@@ -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)
@@ -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):
@@ -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