HT local subregions testing

samseg/subregions/core.py

@@ -7,6 +7,15 @@
 from samseg import gems
 from samseg.utilities import requireNumpyArray
 from samseg.subregions import utils
+from samseg.subregions.gaussian import (
+    covariance_for_gems,
+    diagonal_gaussian_log_likelihood,
+    diagonal_posterior_update,
+    full_covariance_posterior_update,
+    full_gaussian_log_likelihood,
+    full_mean_prior_cost,
+    validate_covariance_mode,
+)
 
 
 class MeshModel:
@@ -22,6 +31,7 @@ def __init__(
         bbregisterMode=None,
         resolution=0.5,
         useTwoComponents=False,
+        covariance_mode='diagonal',
         tempDir=None,
         fileSuffix='',
         debug=False,
@@ -58,6 +68,7 @@ def __init__(
         self.bbregisterMode = bbregisterMode
         self.resolution = resolution
         self.useTwoComponents = useTwoComponents
+        self.covariance_mode = validate_covariance_mode(covariance_mode)
         self.tempDir = tempDir
         self.fileSuffix = fileSuffix
         self.debug = debug
@@ -579,23 +590,45 @@ def fit_mesh_to_image(self):
 
                     n_channels = 1 if imageBuffer.ndim == 3 else imageBuffer.shape[-1]
                     self.means = np.zeros((numberOfClasses, n_channels))
-                    self.variances = np.zeros((numberOfClasses, n_channels))
+                    if self.covariance_mode == 'diagonal':
+                        self.variances = np.zeros((numberOfClasses, n_channels))
+                    else:
+                        self.variances = np.zeros((numberOfClasses, n_channels, n_channels))
 
                     thresh = 1e-2
                     for classNumber in range(numberOfClasses):
                         posterior = posteriors[:, classNumber]
                         if np.sum(posterior) > thresh:
-
-                            mu = (self.meanHyper[classNumber] * self.nHyper[classNumber] + data.T @ posterior) / (self.nHyper[classNumber] + np.sum(posterior) + thresh)
-                            variance = (((data - mu) ** 2).T @ posterior + self.nHyper[classNumber] * (mu - self.meanHyper[classNumber]) ** 2) / (np.sum(posterior) + thresh)
-                            self.means[classNumber] = mu
-                            self.variances[classNumber] = variance + thresh
+                            if self.covariance_mode == 'diagonal':
+                                mu, variance = diagonal_posterior_update(
+                                    data,
+                                    posterior,
+                                    self.meanHyper[classNumber],
+                                    self.nHyper[classNumber],
+                                    thresh=thresh,
+                                )
+                                self.means[classNumber] = mu
+                                self.variances[classNumber] = variance
+                            else:
+                                mu, covariance = full_covariance_posterior_update(
+                                    data,
+                                    posterior,
+                                    self.meanHyper[classNumber],
+                                    self.nHyper[classNumber],
+                                    thresh=thresh,
+                                )
+                                self.means[classNumber] = mu
+                                self.variances[classNumber] = covariance
                         else:
                             self.means[classNumber] = self.meanHyper[classNumber]
-                            self.variances[classNumber] = 100
+                            if self.covariance_mode == 'diagonal':
+                                self.variances[classNumber] = 100
+                            else:
+                                self.variances[classNumber] = np.eye(n_channels) * 100
 
                     # Prevents NaNs during the optimization
-                    self.variances[self.variances == 0] = 100
+                    if self.covariance_mode == 'diagonal':
+                        self.variances[self.variances == 0] = 100
 
                 stopCriterionEM = 1e-5
                 historyOfEMCost = []
@@ -611,11 +644,14 @@ def fit_mesh_to_image(self):
                         variance = self.variances[classNumber]
                         prior = priors[:, classNumber] / 65535
 
-                        log_likelihood = -0.5 * np.sum(((data - mu) ** 2) / variance + np.log(2 * np.pi * variance), axis=1)
+                        if self.covariance_mode == 'diagonal':
+                            log_likelihood = diagonal_gaussian_log_likelihood(data, mu, variance)
+                            minLogLikelihood = minLogLikelihood + 0.5 * np.sum(np.log(2 * np.pi * variance)) - 0.5 * np.log(self.nHyper[classNumber]) + 0.5 * np.sum((self.nHyper[classNumber] / variance) * (mu - self.meanHyper[classNumber]) ** 2)
+                        else:
+                            log_likelihood = full_gaussian_log_likelihood(data, mu, variance)
+                            minLogLikelihood = minLogLikelihood + full_mean_prior_cost(mu, variance, self.meanHyper[classNumber], self.nHyper[classNumber])
 
                         posteriors[:, classNumber] = np.exp(log_likelihood) * prior
-
-                        minLogLikelihood = minLogLikelihood + 0.5 * np.sum(np.log(2 * np.pi * variance)) - 0.5 * np.log(self.nHyper[classNumber]) + 0.5 * np.sum((self.nHyper[classNumber] / variance) * (mu - self.meanHyper[classNumber]) ** 2)
 
                     normalizer = np.sum(posteriors, -1) + np.finfo(np.float32).eps
                     posteriors /= normalizer[..., np.newaxis]
@@ -648,16 +684,36 @@ def fit_mesh_to_image(self):
                     for classNumber in range(numberOfClasses):
                         posterior = posteriors[:, classNumber]
                         if np.sum(posterior) > thresh:
-                            mu = (self.meanHyper[classNumber] * self.nHyper[classNumber] + data.T @ posterior) / (self.nHyper[classNumber] + np.sum(posterior) + thresh)
-                            variance = (((data - mu) ** 2).T @ posterior + self.nHyper[classNumber] * (mu - self.meanHyper[classNumber]) ** 2) / (np.sum(posterior) + thresh)
-                            self.means[classNumber] = mu
-                            self.variances[classNumber] = variance + thresh
+                            if self.covariance_mode == 'diagonal':
+                                mu, variance = diagonal_posterior_update(
+                                    data,
+                                    posterior,
+                                    self.meanHyper[classNumber],
+                                    self.nHyper[classNumber],
+                                    thresh=thresh,
+                                )
+                                self.means[classNumber] = mu
+                                self.variances[classNumber] = variance
+                            else:
+                                mu, covariance = full_covariance_posterior_update(
+                                    data,
+                                    posterior,
+                                    self.meanHyper[classNumber],
+                                    self.nHyper[classNumber],
+                                    thresh=thresh,
+                                )
+                                self.means[classNumber] = mu
+                                self.variances[classNumber] = covariance
                         else:
                             self.means[classNumber] = self.meanHyper[classNumber]
-                            self.variances[classNumber] = 100
+                            if self.covariance_mode == 'diagonal':
+                                self.variances[classNumber] = 100
+                            else:
+                                self.variances[classNumber] = np.eye(n_channels) * 100
 
                     # Prevents NaNs during the optimization
-                    self.variances[self.variances == 0] = 100
+                    if self.covariance_mode == 'diagonal':
+                        self.variances[self.variances == 0] = 100
 
                 # Part II: update the position of the mesh nodes for the current set of Gaussian parameters
 
@@ -670,7 +726,7 @@ def fit_mesh_to_image(self):
                 ##$ reshape the variances
                 full_variance = np.zeros((self.means.shape[0], self.means.shape[1], self.means.shape[1]))
                 for i in range(self.means.shape[0]):
-                    full_variance[i] = np.diag(self.variances[i])
+                    full_variance[i] = covariance_for_gems(self.variances[i], self.covariance_mode)
 
                 ##$ handle building the image list for single and multi channel
                 if imageBuffer.ndim == 3:
@@ -759,8 +815,10 @@ def extract_segmentation(self):
             mu = self.means[self.reducingLookupTable[classNumber]]
             variance = self.variances[self.reducingLookupTable[classNumber]]
 
-            # changed to handle multiple channels
-            log_likelihood = -0.5 * np.sum(((imgdata - mu) ** 2) / variance + np.log(2 * np.pi * variance), axis=-1)
+            if self.covariance_mode == 'diagonal':
+                log_likelihood = diagonal_gaussian_log_likelihood(imgdata, mu, variance)
+            else:
+                log_likelihood = full_gaussian_log_likelihood(imgdata, mu, variance)
             posteriors[:, classNumber] = np.exp(log_likelihood) * (prior[self.maskIndices] / 65535)
             #posteriors[:, classNumber] = (np.exp(-(imgdata - mu) ** 2 / 2 / variance) * (prior[self.maskIndices[:3]] / 65535)) / np.sqrt(2 * np.pi * variance)
 

samseg/subregions/for_testing/dti_args_FA_henry.json

@@ -0,0 +1,16 @@
+{
+    "atlasDir":"/Applications/freesurfer/8.2.0/average/ThalamicNuclei/atlas_DTI",
+    "outDir":"/Users/henrytregidgo/PycharmProjects/Samseg/samseg/for_testing_outputs/mul_ch",
+    "inputImageFileNames":["/Users/henrytregidgo/Documents/testVolumes/ThalamusTemplate/subject_100206/mri/norm.mgz", "/Users/henrytregidgo/Documents/testVolumes/ThalamusTemplate/subject_100206/dmri/dtifit.1+2+3K_FA.nii.gz"],
+    "inputSegFileName":"/Users/henrytregidgo/Documents/testVolumes/ThalamusTemplate/subject_100206/mri/aseg.mgz",
+    "inputDTIDirName":"/Users/henrytregidgo/Documents/testVolumes/ThalamusTemplate/subject_100206/dmri",
+    "dtiLikelihood":"DSWbeta",
+    "meshStiffness":0.05,
+    "optimizerType":"L-BFGS",
+    "bbregisterMode":null,
+    "resolution":0.5,
+    "useTwoComponents":true,
+    "tempDir":"/Users/henrytregidgo/PycharmProjects/Samseg/samseg/tmp_mul_ch",
+    "fileSuffix":"_thalamus_joint",
+    "debug":true
+}

samseg/subregions/for_testing/dti_args_henry.json

@@ -0,0 +1,16 @@
+{
+    "atlasDir":"/Applications/freesurfer/8.2.0/average/ThalamicNuclei/atlas_DTI",
+    "outDir":"/Users/henrytregidgo/PycharmProjects/Samseg/samseg/for_testing_outputs/dti",
+    "inputImageFileNames":["/Users/henrytregidgo/Documents/testVolumes/ThalamusTemplate/subject_100206/mri/norm.mgz"],
+    "inputSegFileName":"/Users/henrytregidgo/Documents/testVolumes/ThalamusTemplate/subject_100206/mri/aseg.mgz",
+    "inputDTIDirName":"/Users/henrytregidgo/Documents/testVolumes/ThalamusTemplate/subject_100206/dmri",
+    "dtiLikelihood":"DSWbeta",
+    "meshStiffness":0.05,
+    "optimizerType":"L-BFGS",
+    "bbregisterMode":null,
+    "resolution":0.5,
+    "useTwoComponents":true,
+    "tempDir":"/Users/henrytregidgo/PycharmProjects/Samseg/samseg/tmp_dti",
+    "fileSuffix":"_thalamus_joint",
+    "debug":true
+}

samseg/subregions/for_testing/multichannel_gaussian_plan.md

@@ -0,0 +1,184 @@
+# Subregions Multi-Channel Gaussian Plan
+
+This note records the current state of the DTI/multi-channel Gaussian work and
+the expected shape of the fix. It is intentionally a planning document, not a
+production design.
+
+## Current state
+
+The current Python subregions path supports multi-channel image data, but it
+models the channels as conditionally independent given the class label.
+
+In `samseg/subregions/core.py`, the fitted Gaussian parameters are stored as:
+
+- `self.means`: one mean vector per class, shape `(n_classes, n_channels)`.
+- `self.variances`: one variance vector per class, shape
+  `(n_classes, n_channels)`.
+
+The E-step likelihood is currently computed as a sum of independent univariate
+Gaussian log likelihoods:
+
+```text
+log p(x_i | c) =
+  -0.5 * sum_d [ (x_id - mu_cd)^2 / sigma_cd^2
+                 + log(2*pi*sigma_cd^2) ]
+```
+
+That happens in both the fitting loop and final segmentation extraction:
+
+- `core.py`: EM likelihood inside `fit_mesh_to_image`.
+- `core.py`: posterior reconstruction inside `extract_segmentation`.
+
+The M-step also updates only per-channel variances:
+
+```text
+mu_c =
+  (n_c0 * mu_c0 + sum_i gamma_ic * x_i)
+  / (n_c0 + sum_i gamma_ic + eps)
+
+sigma_cd^2 =
+  (sum_i gamma_ic * (x_id - mu_cd)^2
+   + n_c0 * (mu_cd - mu_c0d)^2)
+  / (sum_i gamma_ic + eps)
+```
+
+For mesh deformation, the Python path expands the per-channel variance vector
+into a diagonal covariance matrix before calling the C++ calculator:
+
+```python
+full_variance[i] = np.diag(self.variances[i])
+```
+
+That means the current implementation is mathematically consistent only for a
+diagonal covariance model. It is not a full multivariate Gaussian model.
+
+## Existing support to keep
+
+The diagonal/independent-channel behavior is still useful and should remain
+available behind an explicit option, for example:
+
+```text
+covariance_mode = "diagonal"
+```
+
+or:
+
+```text
+independent_channels = True
+```
+
+This mode gives a simpler model, easier diagnostics, and a useful regression
+target while implementing full covariance.
+
+## Required target
+
+The full multi-channel model should allow one covariance matrix per class:
+
+```text
+Sigma_c in R^(D x D)
+```
+
+where `D` is the number of channels. The log likelihood should become:
+
+```text
+log p(x_i | c) =
+  -0.5 * [ (x_i - mu_c)^T Sigma_c^-1 (x_i - mu_c)
+           + log det(Sigma_c)
+           + D * log(2*pi) ]
+```
+
+The M-step needs to estimate the full weighted covariance:
+
+```text
+Sigma_c =
+  [ sum_i gamma_ic * (x_i - mu_c)(x_i - mu_c)^T + prior_terms ]
+  / [ sum_i gamma_ic + prior_weight ]
+```
+
+The existing `meanHyper`/`nHyper` prior logic is vector/diagonal-oriented. A
+full covariance implementation therefore needs an explicit decision about the
+covariance prior or shrinkage strategy. A normal-inverse-Wishart style prior is
+the conventional full-covariance equivalent, but a simpler shrinkage-to-diagonal
+regularizer may be enough for the local DTI migration if it is documented and
+tested.
+
+## Implementation shape
+
+1. Add an explicit covariance-mode configuration path.
+
+   The default should preserve current behavior until full covariance has been
+   tested. A likely interface is:
+
+   ```text
+   covariance_mode = "diagonal" | "full"
+   ```
+
+2. Introduce shared Gaussian helpers in the Python subregions code.
+
+   The fitting loop and `extract_segmentation` should call the same likelihood
+   helper instead of duplicating the formula. The helper should support both
+   diagonal vectors and full covariance matrices.
+
+3. Normalize parameter storage.
+
+   Keep diagonal mode simple, but define clear shapes:
+
+   ```text
+   diagonal: self.variances shape = (classes, channels)
+   full:     self.variances shape = (classes, channels, channels)
+   ```
+
+   Longer term, renaming `variances` to `covariances` would be clearer, but that
+   should be handled carefully because the C++/GEMS interfaces also use the name
+   `variances` for covariance matrices.
+
+4. Update initialization and the M-step.
+
+   Diagonal mode should reproduce the current formulas. Full mode should compute
+   weighted covariance matrices and apply a positive-definite regularizer. The
+   implementation must handle disappearing classes without producing singular
+   matrices.
+
+5. Pass full covariance matrices into mesh deformation.
+
+   The C++ GEMS likelihood/filter path already accepts covariance matrices in
+   several interfaces, and `kvlGMMLikelihoodImageFilter` computes inverses and
+   determinants from those matrices. The Python binding and the active
+   calculator path still need to be verified with a small full-covariance test.
+
+6. Extend validation.
+
+   Minimum useful checks:
+
+   - Single-channel behavior is unchanged.
+   - Diagonal mode reproduces the current local smoke-test outputs within a
+     small tolerance.
+   - Full mode matches a manual or SciPy multivariate Gaussian likelihood on a
+     small synthetic two-channel example with non-zero off-diagonal covariance.
+   - Full mode rejects or regularizes singular covariance estimates.
+   - A local DTI smoke test completes with `covariance_mode = "full"`.
+
+## Open questions
+
+- What should the full covariance prior be?
+- Should full covariance be enabled only for DTI/multi-channel runs, or exposed
+  generally through the subregions model configuration?
+- How much output parity with MATLAB is needed before merging?
+- Does the MATLAB DTI path use a full covariance prior, a fixed covariance, or
+  an empirical covariance update that should be copied directly?
+- Should covariance estimation happen for all classes, or only for classes with
+  DTI-specific groupings?
+
+## Related audit material
+
+The broader MATLAB-to-Python migration audit is currently on the
+`docs/dti-migration-checklists` branch, not this branch. The most relevant
+starting points there are:
+
+- `docs/dev/migration/thalamus-gap-analysis.md`
+- `docs/dev/migration/mapping_v2/00-index.md`
+- `docs/dev/migration/mapping_v2/compact/20-gap-action-summary.md`
+
+Those documents were created against older branch tips, so they should be
+refreshed against the current `dti_integration` tip before being treated as the
+final migration state.