diff --git a/samseg/subregions/core.py b/samseg/subregions/core.py index 2003b26..e81eca6 100644 --- a/samseg/subregions/core.py +++ b/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) diff --git a/samseg/subregions/for_testing/dti_args_FA_henry.json b/samseg/subregions/for_testing/dti_args_FA_henry.json new file mode 100644 index 0000000..2b95f95 --- /dev/null +++ b/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 +} diff --git a/samseg/subregions/for_testing/dti_args_henry.json b/samseg/subregions/for_testing/dti_args_henry.json new file mode 100644 index 0000000..97b6bb0 --- /dev/null +++ b/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 +} diff --git a/samseg/subregions/for_testing/multichannel_gaussian_plan.md b/samseg/subregions/for_testing/multichannel_gaussian_plan.md new file mode 100644 index 0000000..2a7f1b3 --- /dev/null +++ b/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. diff --git a/samseg/subregions/for_testing/shell_test.py b/samseg/subregions/for_testing/shell_test.py index 89f47fc..9be28b7 100644 --- a/samseg/subregions/for_testing/shell_test.py +++ b/samseg/subregions/for_testing/shell_test.py @@ -12,86 +12,177 @@ ### import required modules import json +from pathlib import Path from samseg.subregions import thalamus as thalamus from samseg.subregions import thalamusDTI as DTI -### load the json file with the MeshModel args -f = '/autofs/space/anubis_001/users/jackson/samsegDTI/port/samseg/samseg/subregions/dti_args.json' # update this path -f = open(f,'r') -args =json.load(f) -f.close() +TESTING_DIR = Path(__file__).resolve().parent + +ARG_PROFILES = [ + { + "name": "Henry", + "dti": TESTING_DIR / "dti_args_henry.json", + "multi": TESTING_DIR / "dti_args_FA_henry.json", + }, + { + "name": "Jackson", + "dti": TESTING_DIR / "dti_args.json", + "multi": TESTING_DIR / "dti_args_FA.json", + }, +] + + +def load_args(path): + with open(path, "r") as f: + return json.load(f) + + +def required_paths(args): + paths = { + "atlasDir": Path(args["atlasDir"]), + "inputSegFileName": Path(args["inputSegFileName"]), + "inputDTIDirName": Path(args["inputDTIDirName"]), + "inputImageFileNames[0]": Path(args["inputImageFileNames"][0]), + } + return paths + + +def select_args_profile(): + missing_by_profile = {} + for profile in ARG_PROFILES: + dti_args = load_args(profile["dti"]) + multi_args = load_args(profile["multi"]) + missing = [ + f"{args_name}.{path_name}: {path}" + for args_name, args in (("dti", dti_args), ("multi", multi_args)) + for path_name, path in required_paths(args).items() + if not path.exists() + ] + if not missing: + return profile["name"], dti_args, multi_args + missing_by_profile[profile["name"]] = missing + + lines = ["No valid subregions testing profile found.", "Missing paths:"] + for profile_name, missing in missing_by_profile.items(): + lines.append(f"- {profile_name}") + lines.extend(f" - {path}" for path in missing) + raise SystemExit("\n".join(lines)) + + +profile_name, args, multi_args = select_args_profile() +print(f"Using subregions testing profile: {profile_name}") + +### Object roles: +### - dti: DTI-aware thalamus run with one structural image plus the DTI dir. +### - thal: standard thalamus run with the structural image only. +### - multi: DTI-aware thalamus run with structural and FA image channels. +### +### With the Henry profile, debug/intermediate files are written under: +### - dti: tmp_dti +### - thal: tmp_thal +### - multi: tmp_mul_ch +### +### Final outputs are written under each object's outDir. For DTI objects, +### initialize() appends results/EM/ under outDir. ### init the DTI class dti = DTI.ThalamicNucleiDTI(**args) ### pop DTI specific args, update temp dir name for standard subregions init -args.pop('atlasDir') -args.pop('inputDTIDirName') -args.pop('dtiLikelihood') -args['tempDir'] = 'tmp_thal' +args.pop("atlasDir") +args.pop("inputDTIDirName") +args.pop("dtiLikelihood") +args["tempDir"] = str(Path(args["tempDir"]).parent / "tmp_thal") +args["outDir"] = str(Path(args["outDir"]).parent / "thal") ### init standard thalamus subregions class thal = thalamus.ThalamicNuclei(**args) -### load the json for the multi channel MeshModel class -f = '/autofs/space/anubis_001/users/jackson/samsegDTI/port/samseg/samseg/subregions/dti_args_FA.json' # NOTE: this file differs from previous json, also includes path to FA image -f = open(f,'r') -args =json.load(f) -f.close() - ### init the multi channel MeshModel -multi = DTI.ThalamicNucleiDTI(**args) +multi = DTI.ThalamicNucleiDTI(**multi_args) ### BEGIN PROCESS.PY CALLS ## initialize +## Writes/creates: +## - dti/tempDir, thal/tempDir, multi/tempDir +## - dti/outDir/results/EM/DSWbeta_thalamus_joint +## - thal/outDir +## - multi/outDir/results/EM/DSWbeta_thalamus_joint dti.initialize() thal.initialize() multi.initialize() ## align atlas to input seg +## Writes in each tempDir: +## - targetMask.mgz +## - alignedAtlasImage.mgz +## - trash.lta dti.align_atlas_to_seg() thal.align_atlas_to_seg() multi.align_atlas_to_seg() ## prep for seg fitting +## Writes in each tempDir when debug=True: +## - synthImage.mgz +## - synthImageMasked.mgz dti.prepare_for_seg_fitting() thal.prepare_for_seg_fitting() multi.prepare_for_seg_fitting() ## fit mesh to seg +## Writes in each tempDir: +## - warpedOriginalMesh.txt dti.fit_mesh_to_seg() thal.fit_mesh_to_seg() multi.fit_mesh_to_seg() ## additional k-means clustering step for DTI # Not performed on standard thalamus pipeline +## Writes in dti and multi tempDirs: +## - initialSegFromPriors.mgz +## - boxedASEGTHDE.mgz dti.synthseg_kmeans() multi.synthseg_kmeans() ## prepare for image fitting +## Writes in each tempDir when debug=True: +## - processedImage.mgz +## - processedImageMasked.mgz +## - processedImageMask.mgz dti.prepare_for_image_fitting() thal.prepare_for_image_fitting() multi.prepare_for_image_fitting() ## fit mesh to image +## Optimizes the in-memory mesh and Gaussian parameters. +## No direct file writes are expected from this step. dti.fit_mesh_to_image() thal.fit_mesh_to_image() multi.fit_mesh_to_image() ## extract segmentation # This will need some work for the DTI classes +## Writes in each tempDir when debug=True: +## - finalWarpedMesh.txt +## - finalWarpedMeshNoAffine.txt +## - discreteLabelsAll.mgz dti.extract_segmentation() thal.extract_segmentation() multi.extract_segmentation() ## postprocess segmentation # This will also need a bit of work for the DTI calsses +## Writes in each object's output tree: +## - ThalamicNuclei_thalamus_joint.mgz +## - ThalamicNuclei_thalamus_joint.FSvoxelSpace.mgz +## - ThalamicNuclei_thalamus_joint.volumes.txt dti.postprocess_segmentation() thal.postprocess_segmentation() multi.postprocess_segmentation() ## cleanup # not really needed for tests, should just remove temp files +## With debug=True, tempDirs are left in place for inspection. dti.cleanup() thal.cleanup() -multi.cleanup() \ No newline at end of file +multi.cleanup() diff --git a/samseg/subregions/gaussian.py b/samseg/subregions/gaussian.py new file mode 100644 index 0000000..2947b4e --- /dev/null +++ b/samseg/subregions/gaussian.py @@ -0,0 +1,310 @@ +import numpy as np + + +def validate_covariance_mode(covariance_mode): + """ + Validate the covariance parameterization used by the subregions model. + + Parameters + ---------- + covariance_mode : str + Requested covariance representation. Must be either ``"diagonal"`` + for per-channel variances or ``"full"`` for multivariate covariance + matrices. + + Returns + ------- + str + The validated covariance mode string. + + Raises + ------ + ValueError + If ``covariance_mode`` is not one of the supported values. + + Subfields Useage + ---------------- + This keeps the covariance-mode choice consistent before the EM loop, + posterior extraction, and GEMS bridge all consume it. + """ + if covariance_mode not in ("diagonal", "full"): + raise ValueError( + "covariance_mode must be one of: 'diagonal', 'full'" + ) + return covariance_mode + + +def repair_covariance_eigh(covariance, min_eigenvalue=1e-6): + """ + Repair a covariance matrix with an eigenvalue floor. + + Parameters + ---------- + covariance : array_like + Input covariance matrix to repair. + min_eigenvalue : float, optional + Absolute lower bound applied to each eigenvalue after symmetrization. + + Returns + ------- + ndarray + Symmetric positive-definite covariance matrix. + + Subfields Useage + ---------------- + This is an expensive last-resort repair for external or diagnostic + covariance inputs. The normal subfields EM path should create SPD + covariances by construction and should not call this in the fitting loop. + """ + covariance = np.asarray(covariance, dtype=float) + covariance = 0.5 * (covariance + covariance.T) + eigvals, eigvecs = np.linalg.eigh(covariance) + scale = np.max(np.abs(eigvals)) if eigvals.size else 1.0 + floor = max(min_eigenvalue, np.finfo(float).eps * scale) + eigvals = np.maximum(eigvals, floor) + repaired = eigvecs @ np.diag(eigvals) @ eigvecs.T + return 0.5 * (repaired + repaired.T) + + +def covariance_for_gems(covariance, covariance_mode): + """ + Convert a covariance representation into the matrix form expected by GEMS. + + Parameters + ---------- + covariance : array_like + Diagonal vector or full covariance matrix for a single class. + covariance_mode : {"diagonal", "full"} + Covariance representation used by the Python EM path. + + Returns + ------- + ndarray + A covariance matrix suitable for the GEMS cost/gradient calculator. + + Raises + ------ + ValueError + If the input shape is incompatible with the selected mode. + + Subfields Useage + ---------------- + The Python fitting loop stores class covariances in the mode-specific + internal form, but the mesh deformation bridge always needs matrices. + """ + validate_covariance_mode(covariance_mode) + covariance = np.asarray(covariance, dtype=float) + if covariance_mode == "diagonal": + if covariance.ndim == 1: + return np.diag(covariance) + if covariance.ndim == 2: + if covariance.shape[0] != covariance.shape[1]: + raise ValueError("Diagonal covariance matrix must be square") + return np.diag(np.diag(covariance)) + raise ValueError("Diagonal covariance expects a vector or matrix") + if covariance.ndim == 1: + return np.diag(covariance) + if covariance.ndim != 2: + raise ValueError("Full covariance expects a vector or matrix") + if covariance.shape[0] != covariance.shape[1]: + raise ValueError("Full covariance matrix must be square") + return covariance + + +def diagonal_gaussian_log_likelihood(data, mean, variances): + """ + Evaluate a diagonal multivariate Gaussian log likelihood. + + Parameters + ---------- + data : array_like + Sample matrix with shape ``(n_samples, n_channels)``. + mean : array_like + Mean vector with shape ``(n_channels,)``. + variances : array_like + Per-channel variances with shape ``(n_channels,)``. + + Returns + ------- + ndarray + Log likelihood for each sample. + + Subfields Useage + ---------------- + This is the current subregions likelihood used by both the EM update and + final posterior reconstruction when the model stays in diagonal mode. + """ + data = np.asarray(data, dtype=float) + mean = np.asarray(mean, dtype=float) + variances = np.asarray(variances, dtype=float) + return -0.5 * np.sum( + ((data - mean) ** 2) / variances + np.log(2 * np.pi * variances), + axis=-1, + ) + + +def full_gaussian_log_likelihood(data, mean, covariance): + """ + Evaluate a full-covariance multivariate Gaussian log likelihood. + + Parameters + ---------- + data : array_like + Sample matrix with shape ``(n_samples, n_channels)``. + mean : array_like + Mean vector with shape ``(n_channels,)``. + covariance : array_like + Symmetric covariance matrix with shape ``(n_channels, n_channels)``. + + Returns + ------- + ndarray + Log likelihood for each sample. + + Subfields Useage + ---------------- + Full covariance mode uses this in the same places as the diagonal helper + so the EM path and the segmentation posterior use one likelihood formula. + """ + data = np.asarray(data, dtype=float) + mean = np.asarray(mean, dtype=float) + covariance = np.asarray(covariance, dtype=float) + delta = data - mean + chol = np.linalg.cholesky(covariance) + logdet = 2 * np.sum(np.log(np.diag(chol))) + solved = np.linalg.solve(chol, delta.T).T + quadratic = np.sum(solved * solved, axis=-1) + dim = covariance.shape[0] + return -0.5 * (quadratic + logdet + dim * np.log(2 * np.pi)) + + +def diagonal_posterior_update(data, posterior, mean_hyper, n_hyper, thresh=1e-2): + """ + Update diagonal Gaussian parameters from weighted posteriors. + + Parameters + ---------- + data : array_like + Sample matrix with shape ``(n_samples, n_channels)``. + posterior : array_like + Responsibility weights for the class. + mean_hyper : array_like + Prior mean for the class. + n_hyper : float + Prior strength associated with the mean hyperparameter. + thresh : float, optional + Minimum posterior mass before falling back to a default covariance. + + Returns + ------- + mean : ndarray + Updated mean vector. + variances : ndarray + Updated per-channel variances. + + Subfields Useage + ---------------- + The diagonal M-step in `core.py` uses this to preserve the existing + per-channel update rule while keeping the code in one shared helper. + """ + data = np.asarray(data, dtype=float) + posterior = np.asarray(posterior, dtype=float) + mean_hyper = np.asarray(mean_hyper, dtype=float) + n_hyper = float(n_hyper) + total = float(np.sum(posterior)) + if total <= thresh: + return mean_hyper.copy(), np.full(data.shape[1], 100.0) + mu = (mean_hyper * n_hyper + data.T @ posterior) / (n_hyper + total + thresh) + variance = (((data - mu) ** 2).T @ posterior + n_hyper * (mu - mean_hyper) ** 2) / (total + thresh) + variance = np.maximum(variance + thresh, thresh) + return mu, variance + + +def full_mean_prior_cost(mean, covariance, mean_hyper, n_hyper): + """ + Evaluate the full-covariance mean-prior cost term. + + Parameters + ---------- + mean : array_like + Current class mean vector. + covariance : array_like + Current class covariance matrix. + mean_hyper : array_like + Prior mean vector. + n_hyper : float + Prior strength associated with the mean. + + Returns + ------- + float + Scalar negative log-prior contribution for the current class mean. + + Subfields Useage + ---------------- + This provides the full-covariance analogue of the existing diagonal + mean-prior cost used in the subregions EM objective. + """ + mean = np.asarray(mean, dtype=float) + covariance = np.asarray(covariance, dtype=float) + mean_hyper = np.asarray(mean_hyper, dtype=float) + n_hyper = float(n_hyper) + dim = mean.shape[0] + chol = np.linalg.cholesky(covariance) + logdet = 2 * np.sum(np.log(np.diag(chol))) + delta = mean - mean_hyper + solved = np.linalg.solve(chol, delta) + return 0.5 * ( + n_hyper * solved @ solved + + logdet + + dim * np.log(2 * np.pi) + - dim * np.log(n_hyper) + ) + + +def full_covariance_posterior_update(data, posterior, mean_hyper, n_hyper, thresh=1e-2): + """ + Update full covariance Gaussian parameters with a weighted ridge step. + + Parameters + ---------- + data : array_like + Sample matrix with shape ``(n_samples, n_channels)``. + posterior : array_like + Responsibility weights for the class. + mean_hyper : array_like + Prior mean vector. + n_hyper : float + Prior strength associated with the mean hyperparameter. + thresh : float, optional + Minimum posterior mass before falling back to a broad covariance. + + Returns + ------- + mean : ndarray + Updated mean vector. + covariance : ndarray + Updated full covariance matrix. + + Subfields Useage + ---------------- + Full covariance mode uses this M-step helper so `core.py` can estimate + multichannel class covariances without duplicating the weighted scatter + and mean-prior terms. + """ + data = np.asarray(data, dtype=float) + posterior = np.asarray(posterior, dtype=float) + mean_hyper = np.asarray(mean_hyper, dtype=float) + n_hyper = float(n_hyper) + total = float(np.sum(posterior)) + dim = data.shape[1] + if total <= thresh: + return mean_hyper.copy(), np.eye(dim) * 100.0 + mu = (mean_hyper * n_hyper + data.T @ posterior) / (n_hyper + total + thresh) + centered = data - mu + covariance = (centered.T * posterior) @ centered + covariance += n_hyper * np.outer(mu - mean_hyper, mu - mean_hyper) + covariance /= (total + thresh) + covariance = covariance + thresh * np.eye(dim) + covariance = 0.5 * (covariance + covariance.T) + return mu, covariance diff --git a/samseg/subregions/thalamusDTI.py b/samseg/subregions/thalamusDTI.py index 9c12cc3..c755481 100644 --- a/samseg/subregions/thalamusDTI.py +++ b/samseg/subregions/thalamusDTI.py @@ -32,6 +32,7 @@ def __init__( bbregisterMode=None, resolution=0.5, useTwoComponents=True, # maybe we just hard code this? nothing depends on this in the super.__init__, so we could always hard code post call + covariance_mode='diagonal', tempDir=None, fileSuffix="_thalamus_joint", debug=True, @@ -47,6 +48,7 @@ def __init__( bbregisterMode=bbregisterMode, resolution=resolution, useTwoComponents=useTwoComponents, + covariance_mode=covariance_mode, tempDir=tempDir, fileSuffix=fileSuffix, debug=debug, @@ -323,7 +325,11 @@ def initialize(self): print("LOADING JSON GROUPINGS") self.grouping_dict = json.load( open( - "/autofs/space/anubis_001/users/jackson/samsegDTI/port/tmp/means_groupings.json", + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "for_testing", + "means_groupings.json", + ), "r", ) ) @@ -878,7 +884,11 @@ def get_gaussian_hyps(self, sameGaussianParameters, mesh): """ # should have already called get_label_groups, so sgp should be the sharedGMM version nHyper = np.zeros(len(sameGaussianParameters)) - meanHyper = np.zeros(len(sameGaussianParameters)) + if hasattr(self, "processedImage") and self.processedImage.data.ndim > 3: + n_channels = self.processedImage.data.shape[-1] + else: + n_channels = len(self.inputImages) + meanHyper = np.zeros(len(sameGaussianParameters)) if n_channels == 1 else np.zeros((len(sameGaussianParameters), n_channels)) # TODO this needs to be adapted for multi-image cases (with masking) DATA = self.inputImages[0] @@ -948,7 +958,18 @@ def get_gaussian_hyps(self, sameGaussianParameters, mesh): ) total_mask = MASK & (DATA > 0) data = DATA[total_mask] - meanHyper[g] = np.median(data) + if n_channels == 1: + meanHyper[g] = np.median(data) + else: + channel_means = np.zeros(n_channels) + for channel in range(n_channels): + image = self.inputImages[channel] if channel < len(self.inputImages) else DATA + if image.shape[:3] == DATA.shape[:3]: + channel_mask = MASK & (image > 0) + channel_means[channel] = np.median(image[channel_mask]) + else: + channel_means[channel] = np.median(data) + meanHyper[g] = channel_means """WE NEED TO DECIDE HOW TO STORE THE LAMBDAS""" # PESUDO CODE: # if post_em_update is not None: @@ -961,7 +982,11 @@ def get_gaussian_hyps(self, sameGaussianParameters, mesh): M, H = post_em_update(self) # optionally update the meanHyper and nHyper if new values returned if M is not None: - meanHyper[g] = M + if n_channels == 1: + meanHyper[g] = M + else: + M = np.asarray(M) + meanHyper[g] = M if M.shape == (n_channels,) else np.full(n_channels, M) if H is not None: nHyper[g] = H if self.bp: @@ -977,8 +1002,11 @@ def get_gaussian_hyps(self, sameGaussianParameters, mesh): # If any NaN, replace by background # ATH: I don't there would ever be NaNs here? - nans = np.isnan(meanHyper) - meanHyper[nans] = 55 + nans = np.isnan(meanHyper) if n_channels == 1 else np.any(np.isnan(meanHyper), axis=1) + if n_channels == 1: + meanHyper[nans] = 55 + else: + meanHyper[nans] = 55 nHyper[nans] = 10 print("get_g_hyps end") if self.bp: @@ -1066,17 +1094,29 @@ def get_second_gaussian_hyps(self, sameGaussianParameters, meanHyper, nHyper): if True: # Lateral, brighter nHyper[-1] = 25 - meanHyper[-1] = ThInt + 5 + if np.ndim(meanHyper) == 1: + meanHyper[-1] = ThInt + 5 + else: + meanHyper[-1] = ThInt + 5 # Medial, darker nHyper = np.append(nHyper, 25) - meanHyper = np.append(meanHyper, ThInt - 5) + if np.ndim(meanHyper) == 1: + meanHyper = np.append(meanHyper, ThInt - 5) + else: + meanHyper = np.concatenate([meanHyper, (ThInt - 5)[None, :]], axis=0) else: nHyper[-1] = 25 nHyper = np.append(nHyper, 25) # Lateral, more WM-ish (e.g., darker, in FGATIR) - meanHyper[-1] = ThInt * (0.95 + 0.1 * (meanHyper[WMind] >= meanHyper[GMind])) + if np.ndim(meanHyper) == 1: + meanHyper[-1] = ThInt * (0.95 + 0.1 * (meanHyper[WMind] >= meanHyper[GMind])) + else: + meanHyper[-1] = ThInt * (0.95 + 0.1 * (meanHyper[WMind] >= meanHyper[GMind])) # Medial, more GM-ish (e.g., brighter, in FGATIR) - meanHyper = np.append(meanHyper, ThInt * (0.95 + 0.1 * (meanHyper[WMind] < meanHyper[GMind]))) + if np.ndim(meanHyper) == 1: + meanHyper = np.append(meanHyper, ThInt * (0.95 + 0.1 * (meanHyper[WMind] < meanHyper[GMind]))) + else: + meanHyper = np.concatenate([meanHyper, (ThInt * (0.95 + 0.1 * (meanHyper[WMind] < meanHyper[GMind])))[None, :]], axis=0) """ for g in range(len(sameGaussianParameters)): labels = np.array(sameGaussianParameters[g]) @@ -1103,7 +1143,11 @@ def get_second_gaussian_hyps(self, sameGaussianParameters, meanHyper, nHyper): M, H = post_em_update(self) # optionally update the meanHyper and nHyper if new values returned if M is not None: - meanHyper[g] = M + if np.ndim(meanHyper) == 1: + meanHyper[g] = M + else: + M = np.asarray(M) + meanHyper[g] = M if M.shape == (meanHyper.shape[1],) else np.full(meanHyper.shape[1], M) if H is not None: nHyper[g] = H diff --git a/samseg/tests/test_subregions_gaussian.py b/samseg/tests/test_subregions_gaussian.py new file mode 100644 index 0000000..7f93a31 --- /dev/null +++ b/samseg/tests/test_subregions_gaussian.py @@ -0,0 +1,125 @@ +import importlib.util +from pathlib import Path + +import numpy as np +import pytest + + +GAUSSIAN_PATH = Path(__file__).resolve().parents[1] / "subregions" / "gaussian.py" +SPEC = importlib.util.spec_from_file_location("subregions_gaussian", GAUSSIAN_PATH) +gaussian = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(gaussian) + + +def test_validate_covariance_mode(): + assert gaussian.validate_covariance_mode("diagonal") == "diagonal" + assert gaussian.validate_covariance_mode("full") == "full" + with pytest.raises(ValueError): + gaussian.validate_covariance_mode("bogus") + + +def test_diagonal_log_likelihood_matches_manual(): + data = np.array([[0.0, 1.0], [1.0, 2.0]]) + mean = np.array([0.0, 1.0]) + variances = np.array([1.0, 4.0]) + got = gaussian.diagonal_gaussian_log_likelihood(data, mean, variances) + expected = np.array([ + -0.5 * (0.0 + np.log(2 * np.pi * 1.0) + 0.0 + np.log(2 * np.pi * 4.0)), + -0.5 * (1.0 + np.log(2 * np.pi * 1.0) + 0.25 + np.log(2 * np.pi * 4.0)), + ]) + np.testing.assert_allclose(got, expected) + + +def test_full_log_likelihood_matches_manual(): + data = np.array([[1.0, 2.0]]) + mean = np.array([0.5, 1.5]) + covariance = np.array([[2.0, 0.25], [0.25, 1.5]]) + got = gaussian.full_gaussian_log_likelihood(data, mean, covariance) + delta = data - mean + expected = -0.5 * ( + delta @ np.linalg.solve(covariance, delta.T) + + np.linalg.slogdet(covariance)[1] + + 2 * np.log(2 * np.pi) + ) + np.testing.assert_allclose(got, expected.ravel()) + + +def test_diagonal_posterior_update_matches_current_formula(): + data = np.array([[0.0, 1.0], [2.0, 3.0]]) + posterior = np.array([0.25, 0.75]) + mean_hyper = np.array([1.0, 2.0]) + mu, variance = gaussian.diagonal_posterior_update(data, posterior, mean_hyper, 10.0) + np.testing.assert_allclose(mu, np.array([1.044505, 2.04359673]), rtol=1e-6) + assert np.all(variance > 0) + + +def test_full_covariance_posterior_update_and_mean_prior_cost_are_pd_safe(): + data = np.array([[0.0, 0.0], [2.0, 1.0], [1.0, 3.0]]) + posterior = np.array([0.2, 0.3, 0.5]) + mean_hyper = np.array([0.5, 1.0]) + mu, covariance = gaussian.full_covariance_posterior_update(data, posterior, mean_hyper, 10.0) + assert covariance.shape == (2, 2) + np.testing.assert_allclose(covariance, covariance.T) + np.linalg.cholesky(covariance) + prior_cost = gaussian.full_mean_prior_cost(mu, covariance, mean_hyper, 10.0) + assert np.isfinite(prior_cost) + + +def test_full_covariance_posterior_update_low_mass_returns_broad_identity(): + data = np.array([[0.0, 0.0], [2.0, 1.0]]) + posterior = np.array([0.0, 0.0]) + mean_hyper = np.array([0.5, 1.0]) + mu, covariance = gaussian.full_covariance_posterior_update(data, posterior, mean_hyper, 10.0) + np.testing.assert_allclose(mu, mean_hyper) + np.testing.assert_allclose(covariance, np.eye(2) * 100.0) + + +def test_full_mean_prior_cost_matches_diagonal_one_dimensional_formula(): + mean = np.array([2.0]) + mean_hyper = np.array([1.5]) + covariance = np.array([[4.0]]) + n_hyper = 10.0 + got = gaussian.full_mean_prior_cost(mean, covariance, mean_hyper, n_hyper) + expected = ( + 0.5 * np.log(2 * np.pi * covariance[0, 0]) + - 0.5 * np.log(n_hyper) + + 0.5 * (n_hyper / covariance[0, 0]) * (mean[0] - mean_hyper[0]) ** 2 + ) + np.testing.assert_allclose(got, expected) + + +def test_repair_covariance_eigh_repairs_singular_matrix(): + covariance = np.array([[1.0, 1.0], [1.0, 1.0]]) + repaired = gaussian.repair_covariance_eigh(covariance) + np.linalg.cholesky(repaired) + + +def test_full_log_likelihood_rejects_singular_covariance(): + data = np.array([[1.0, 2.0]]) + mean = np.array([0.5, 1.5]) + covariance = np.array([[1.0, 1.0], [1.0, 1.0]]) + with pytest.raises(np.linalg.LinAlgError): + gaussian.full_gaussian_log_likelihood(data, mean, covariance) + + +def test_covariance_for_gems_shapes(): + diag = np.array([1.0, 2.0]) + full = np.array([[1.0, 0.1], [0.1, 2.0]]) + np.testing.assert_allclose(gaussian.covariance_for_gems(diag, "diagonal"), np.diag(diag)) + np.testing.assert_allclose( + gaussian.covariance_for_gems(diag, "full"), + np.diag(diag), + ) + np.testing.assert_allclose( + gaussian.covariance_for_gems(full, "full"), + full, + ) + + +def test_covariance_for_gems_rejects_invalid_shapes(): + with pytest.raises(ValueError): + gaussian.covariance_for_gems(np.zeros((2, 3)), "full") + with pytest.raises(ValueError): + gaussian.covariance_for_gems(np.zeros((2, 2, 2)), "full") + with pytest.raises(ValueError): + gaussian.covariance_for_gems(np.zeros((2, 3)), "diagonal")