[ENH] PyAFQ in the superficial white matter

[36000]

TODO:
waiting for tinygrad update

This is a few things I have been trying out to modernize pyAFQ on multi-shell images, as well as some random things.

These are used:

1. Require T1; use Brainchop to make brain mask, and to segment into CSF, GM, WM; generates WM/GM interface highly reliably using some basic scikit-image functions and the 3T segmentations.
2. I am adding an option to use IsolationForest from scikit-learn for cleaning and making that the default for VOF and callosal bundles over the old system. This makes scikit-learn a dependency. It does not assume a tube-like shape for a given bundle, instead finding outlier nodes and removing streamlines with too many outlier nodes.
3. I made it so mahalanobis cleaning by default only considers the middle 60% of the bundle. This means streamlines are allowed to deviate in the starting and ending 20% of the bundle. This is useful for allowing more diverse endpoints.
4. Added ability to export endpoint maps as distance from endpoint
5. Made some changes to VOF segmentation. There is still more work to be done in a separate PR.
6. Streamlined API for parallelization. Made parallelization with Ray turned on by default.
7. Made PFT and WMGMI seeding the default over local tracking and uniform seeding in the white matter. More work needs to be done on smarter WMGMI seeding, but this is a good start.
8. I am vendorizing here some code from scilpy for doing unified filtering to get asymmetric ODFs. I have modified this to also use numba and ray, and it is efficient enough, ~20 min. could be an option for anyone interested in endpoints.
9. Use osqpy directly to fit MSMT

[36000]

An example HBN subject results:
https://drive.google.com/file/d/1SoR7fumNvxkmY2uxDl74MAR5HAhjyhXb/view?usp=sharing

[36000]

Some interesting ORs found using seeding by endpoints on wmgmi using 2M seeds

[36000]

Added an interesting output. For each bundle, you can see the distance to the nearest endpoint for every voxel in the gray matter. The distances are in millimeters. This could be thresholded to get a boolean endpoint mask for each bundle.

endpoint_distance_visualization.mov

[Copilot]

Pull Request Overview

This PR introduces modern multi-shell DWI processing capabilities to pyAFQ, along with improvements to tractography, segmentation, and bundle recognition workflows. The primary focus is enhancing analysis of superficial white matter through better seeding, tracking, and cleaning methods.

Key Changes:

• Require T1-weighted images: All workflows now require T1w data for tissue segmentation and interface generation
• Enhanced tracking defaults: Switch to PFT tracking with WM/GM interface seeding and ACT stopping criteria by default
• Improved bundle cleaning: Add IsolationForest cleaning as default for VOF and callosal bundles, plus enhanced Mahalanobis cleaning

Reviewed Changes

Copilot reviewed 43 out of 46 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
setup.cfg	Add scikit-learn, numba, and osqp dependencies
AFQ/tasks/data.py	Add T1w tissue segmentation, DAM fitting, MSMT CSD, and asymmetric ODF processing
AFQ/tasks/tractography.py	Update default tracking to PFT with WM/GM interface seeding
AFQ/recognition/cleaning.py	Add IsolationForest cleaning method and orientation-based Mahalanobis cleaning
AFQ/api/bundle_dict.py	Update VOF and callosal bundle definitions with new cleaning methods
AFQ/models/	Add new models for MSMT, DAM fitting, asymmetric filtering, and WM/GM interface generation

Comments suppressed due to low confidence (2)

AFQ/models/msmt.py:328

• [nitpick] The parameter name 'use_osqppy' is unclear. Consider renaming to 'use_osqp' to better reflect that it uses the OSQP solver.

         use_osqppy=True,

[36000]

Now with brainchop and isolation forest cleaning, here are the callosal bundles from 1million seeds on example HBN subject

[36000]

@arokem right now I am just trying to get the documentation to finish without running out of memory, I think due to increases in parallelization. Anyways, this is otherwise ready for review!

[arokem]

I'll get right to it! Shouldn't take me more than...

... 6-8 months 😄

[arokem]

First batch of comments

[arokem]

This will be a problem in the context of a release of pyAFQ. I don't think that pypi allows specifying github installs for dependencies, see pypa/pip#6301

[arokem]

Maybe we need to wait until tinygrad has another release? Looks like they have these quite often. Otherwise, add to _fixes.py a monkey patch of the bits we need?

[arokem]

I am experiencing an interesting regression with this PR, where calling:

my_afq = GroupAFQ(
    bids_path=study_dir,
    preproc_pipeline="qsiprep",
    brain_mask_definition=brain_mask_definition,
    tracking_params={"n_seeds": 4,
                     "directions": "prob",
                     "odf_model": "CSD",
                     "seed_mask": RoiImage()},
    bundle_info=bundles)

proceeds to use only 4 seeds for tracking (as though I am passing "random_seeds": True as part of the tracking_params dictionary. Is that an intentional consequence of these changes?

[36000]

This is not a regression, its a change in defaults. The old defaults were:

"random_seeds": False
"n_seeds": 1

New defaults are:

"random_seeds": True
"n_seeds": 2000000

I think this is more intuitive but old users might get confused initially.

[arokem]

This old user certainly did 😆

[arokem]

Hey @36000 : when you get a chance, could you please rebase this on main? Thanks!

[arokem]

Interestingly, with this branch, we found that installing pyAFQ into a clean environment doesn't work (under some circumstances?). Instead, both @johndromero and I see this:

Obtaining file:///Users/arokem/tmp/afq_dam_test/pyAFQ
  Installing build dependencies ... done
  Checking if build backend supports build_editable ... done
  Getting requirements to build editable ... done
  Preparing editable metadata (pyproject.toml) ... error
  error: subprocess-exited-with-error
  
  × Preparing editable metadata (pyproject.toml) did not run successfully.
  │ exit code: 1
  ╰─> [71 lines of output]
      /private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools_scm/_integration/version_inference.py:51: UserWarning: version of pyAFQ already set
        warnings.warn(self.message)
      Traceback (most recent call last):
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/expand.py", line 71, in __getattr__
          return next(
                 ^^^^^
      StopIteration
      
      The above exception was the direct cause of the following exception:
      
      Traceback (most recent call last):
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/expand.py", line 185, in read_attr
          value = getattr(StaticModule(module_name, spec), attr_name)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/expand.py", line 77, in __getattr__
          raise AttributeError(f"{self.name} has no attribute {attr}") from e
      AttributeError: AFQ has no attribute __version__
      
      During handling of the above exception, another exception occurred:
      
      Traceback (most recent call last):
        File "/Users/arokem/miniforge3/envs/afq_dam_test/lib/python3.11/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 389, in <module>
          main()
        File "/Users/arokem/miniforge3/envs/afq_dam_test/lib/python3.11/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 373, in main
          json_out["return_val"] = hook(**hook_input["kwargs"])
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/Users/arokem/miniforge3/envs/afq_dam_test/lib/python3.11/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 209, in prepare_metadata_for_build_editable
          return hook(metadata_directory, config_settings)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/build_meta.py", line 478, in prepare_metadata_for_build_editable
          return self.prepare_metadata_for_build_wheel(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/build_meta.py", line 374, in prepare_metadata_for_build_wheel
          self.run_setup()
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/build_meta.py", line 317, in run_setup
          exec(code, locals())
        File "<string>", line 33, in <module>
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/__init__.py", line 115, in setup
          return distutils.core.setup(**attrs)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/_distutils/core.py", line 160, in setup
          dist.parse_config_files()
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/dist.py", line 752, in parse_config_files
          setupcfg.parse_configuration(
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/setupcfg.py", line 188, in parse_configuration
          meta.parse()
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/setupcfg.py", line 502, in parse
          section_parser_method(section_options)
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/setupcfg.py", line 477, in parse_section
          self[name] = value
          ~~~~^^^^^^
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/setupcfg.py", line 294, in __setitem__
          parsed = self.parsers.get(option_name, lambda x: x)(value)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/setupcfg.py", line 600, in _parse_version
          return expand.version(self._parse_attr(value, self.package_dir, self.root_dir))
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/setupcfg.py", line 419, in _parse_attr
          return expand.read_attr(attr_desc, package_dir, root_dir)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/expand.py", line 190, in read_attr
          module = _load_spec(spec, module_name)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/private/var/folders/px/t5ym05pn5k16nz97mzcpqqkr0000gn/T/pip-build-env-2rgpqld6/overlay/lib/python3.11/site-packages/setuptools/config/expand.py", line 211, in _load_spec
          spec.loader.exec_module(module)
        File "<frozen importlib._bootstrap_external>", line 940, in exec_module
        File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
        File "/Users/arokem/tmp/afq_dam_test/pyAFQ/AFQ/__init__.py", line 1, in <module>
          from .version import version as __version__  # noqa
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      ModuleNotFoundError: No module named 'AFQ.version'
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.

I think that it's related to changes in setup/pyproject, but not sure.

If I try to install this on top of an already-existing installation (i.e., in an environment where pyAFQ was previously installed) then no such issue. Not sure why this works in the CI and not locally, but thought I'd raise this here so we can debug.

[36000]

Just updated this to use immlib, rebased on master

[arokem]

A thought I had while thinking about this error: https://neurostars.org/t/some-qsirecon-runs-encountering-csdnanresponseerror-in-pyafq/34222, is that we have been using APM as a registration target for the MNI T1-weighted template. If we are going to require a T1-weighted image, maybe we can also use that T1-weighted image as a target for template registration? Should generally give better registrations to the template and we assume that the T1-weighted is registered to the DWI anyway.

[36000]

In that case we are using a T1 that was registered to the DWI. Is that better? That's chaining two registrations. Not sure if that would be better or worse. I guess it's fine?

Also, should we support registering the T1 to the DWI in pyAFQ or assume the preprocessing pipeline does that? Right now we assume the preprocessing pipeline puts them all in the same space.

[arokem]

I think that we should assume that the T1 and DWI are well-registered to each other and use the T1 as an intermediate for registration to the template. I think that it's a bit tricky to reason about the cost/benefit of each approach (i.e., two registrations vs. registration of APM to the template), but ultimately this means that if the registration DWI <=> T1w improves over time, we will be here to benefit from that, and it seems that it's a more general problem than the ones we would like to solve here. Does that make sense?

[36000]

Yes, that makes sense!