From 6bf638e4cf4ca80747b1c42e495802561ce92375 Mon Sep 17 00:00:00 2001 From: Henry Tregidgo Date: Thu, 14 May 2026 14:56:10 +0100 Subject: [PATCH 1/4] Add local subregions test config profiles Refactor the shared subregions smoke-test harness so Henry and Jackson can run it without editing hard-coded paths in the script. The script now looks for repo-relative JSON profile pairs and uses the first one whose required FreeSurfer, subject, segmentation, and DTI paths exist. Add Henry-specific DTI and multichannel JSON files while leaving Jackson's existing JSON files as the fallback profile. Keep temp and output locations separate for DTI, standard thalamus, and multichannel runs so results remain inspectable and avoid filename collisions. Replace the hard-coded means_groupings.json path in thalamusDTI.py with a path derived from the module location, so the support file is loaded from the repo on either machine. --- .../for_testing/dti_args_FA_henry.json | 16 ++++ .../for_testing/dti_args_henry.json | 16 ++++ samseg/subregions/for_testing/shell_test.py | 80 +++++++++++++++---- samseg/subregions/thalamusDTI.py | 6 +- 4 files changed, 100 insertions(+), 18 deletions(-) create mode 100644 samseg/subregions/for_testing/dti_args_FA_henry.json create mode 100644 samseg/subregions/for_testing/dti_args_henry.json 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/shell_test.py b/samseg/subregions/for_testing/shell_test.py index 89f47fc..3386618 100644 --- a/samseg/subregions/for_testing/shell_test.py +++ b/samseg/subregions/for_testing/shell_test.py @@ -12,35 +12,81 @@ ### 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}") ### 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 @@ -94,4 +140,4 @@ # not really needed for tests, should just remove temp files dti.cleanup() thal.cleanup() -multi.cleanup() \ No newline at end of file +multi.cleanup() diff --git a/samseg/subregions/thalamusDTI.py b/samseg/subregions/thalamusDTI.py index 9c12cc3..16b77da 100644 --- a/samseg/subregions/thalamusDTI.py +++ b/samseg/subregions/thalamusDTI.py @@ -323,7 +323,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", ) ) From 155d45a200c1b5c4389ac953ad4c72467bfcff70 Mon Sep 17 00:00:00 2001 From: Henry Tregidgo Date: Thu, 14 May 2026 15:20:22 +0100 Subject: [PATCH 2/4] Document subregions smoke-test outputs Add notes to the local subregions test harness explaining the three objects it drives: the DTI-aware run, the standard thalamus run, and the multi-channel DTI/FA run. Annotate each process.py-style command group with the temp files or final outputs it writes. This keeps the dirty shared harness easier to inspect after long local smoke tests without changing runtime behavior. --- samseg/subregions/for_testing/shell_test.py | 45 +++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/samseg/subregions/for_testing/shell_test.py b/samseg/subregions/for_testing/shell_test.py index 3386618..9be28b7 100644 --- a/samseg/subregions/for_testing/shell_test.py +++ b/samseg/subregions/for_testing/shell_test.py @@ -72,6 +72,19 @@ def select_args_profile(): 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) @@ -90,54 +103,86 @@ def select_args_profile(): ### 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() From 7b8430943f7745f3d0658d8dd5a480b71a4bdd92 Mon Sep 17 00:00:00 2001 From: Henry Tregidgo Date: Wed, 27 May 2026 11:53:27 +0100 Subject: [PATCH 3/4] Document multichannel Gaussian plan --- .../for_testing/multichannel_gaussian_plan.md | 184 ++++++++++++++++++ 1 file changed, 184 insertions(+) create mode 100644 samseg/subregions/for_testing/multichannel_gaussian_plan.md 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. From 3be107e344e46ed8e4d4dd579f316d8379734ca8 Mon Sep 17 00:00:00 2001 From: Henry Tregidgo Date: Wed, 27 May 2026 16:38:50 +0100 Subject: [PATCH 4/4] Add subregions covariance modes Add shared Gaussian helpers for diagonal and full-covariance subregions fitting, including likelihoods, posterior parameter updates, covariance conversion for GEMS, and validation for covariance-mode inputs. Wire covariance_mode through MeshModel and ThalamicNucleiDTI while preserving diagonal mode as the default. Refactor image fitting and final posterior extraction to use the shared helpers, and pass full covariance matrices to GEMS in both covariance modes. Update DTI hyperparameter estimation so multi-channel runs produce channel-aware meanHyper arrays while keeping nHyper one-dimensional. Add focused pytest coverage for the helper math, covariance shape handling, low-mass fallbacks, singular covariance rejection, and the full mean-prior cost. --- samseg/subregions/core.py | 98 +++++-- samseg/subregions/gaussian.py | 310 +++++++++++++++++++++++ samseg/subregions/thalamusDTI.py | 60 ++++- samseg/tests/test_subregions_gaussian.py | 125 +++++++++ 4 files changed, 563 insertions(+), 30 deletions(-) create mode 100644 samseg/subregions/gaussian.py create mode 100644 samseg/tests/test_subregions_gaussian.py 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/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 16b77da..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, @@ -882,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] @@ -952,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: @@ -965,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: @@ -981,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: @@ -1070,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]) @@ -1107,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")