Use of brainmask results into dimensionality mismatch in estimator

[jhlegarreta]

What happened?

A user reported offline that when providing a brainmask to the DWI data class and attempting to estimate the head motion a dimension mismatch happens. Using a DTI model as the estimator, the data dimensions provided to the model and the data dimensions do not match and result in an error.

Commenting

nifreeze/src/nifreeze/model/dmri.py

Line 99 in 18bec0d

self._data_mask[dataset.bzero < S0_EPSILON] = False

allows the estimation to run.

What command did you use?

My current attempt looks something like this:

from nifreeze.data import dmri
from nifreeze.estimator import Estimator
dwi_data = dmri.from_nii(
    dti.fpath,
    brainmask_file=os.path.dirname(dti.fpath) + "/mask.nii.gz",
    bvec_file=dti.fpath.replace(".nii.gz", ".bvec"),
    bval_file=dti.fpath.replace(".nii.gz", ".bval"),
)
estimator = Estimator("dti")
estimated_affine = estimator.run(
    dwi_data,
    omp_nthreads=1,
    n_jobs=1,
    seed=42,
)

What version of the software are you running?

HEAD as of 08/28/2025

How are you running this software?

Local installation ("bare-metal")

Is your data BIDS valid?

Yes

Are you reusing any previously computed results?

No

Please copy and paste any relevant log output.

Traceback (most recent call last):
  File "/gpfs01/home/bbzmsc/code/fproc/fproc/pipeline.py", line 93, in run
    module.run(self)
  File "/gpfs01/home/bbzmsc/code/fproc/fproc/module.py", line 47, in run
    self.process()
  File "/gpfs01/home/bbzmsc/code/fproc/pipelines/diffad.py", line 666, in process
    estimated_affine = estimator.run(
                       ^^^^^^^^^^^^^^
  File "/gpfs01/software/imaging/miniconda3/envs/fproc/lib/python3.11/site-packages/nifreeze/estimator.py", line 187, in run
    predicted = model.fit_predict(  # type: ignore[union-attr]
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs01/software/imaging/miniconda3/envs/fproc/lib/python3.11/site-packages/nifreeze/model/dmri.py", line 197, in fit_predict
    predicted, _ = _exec_predict(
                   ^^^^^^^^^^^^^^
  File "/gpfs01/software/imaging/miniconda3/envs/fproc/lib/python3.11/site-packages/nifreeze/model/dmri.py", line 45, in _exec_predict
    return np.squeeze(model.predict(**kwargs)), chunk
                      ^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs01/software/imaging/miniconda3/envs/fproc/lib/python3.11/site-packages/dipy/testing/decorators.py", line 201, in wrapper
    return convert_positional_to_keyword(func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs01/software/imaging/miniconda3/envs/fproc/lib/python3.11/site-packages/dipy/testing/decorators.py", line 192, in convert_positional_to_keyword
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs01/software/imaging/miniconda3/envs/fproc/lib/python3.11/site-packages/dipy/reconst/dti.py", line 1255, in predict
    predict[i : i + step] = tensor_prediction(
                            ^^^^^^^^^^^^^^^^^^
  File "/gpfs01/software/imaging/miniconda3/envs/fproc/lib/python3.11/site-packages/dipy/reconst/dti.py", line 690, in tensor_prediction
    lower_tri = lower_triangular(qform, b0=S0)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs01/software/imaging/miniconda3/envs/fproc/lib/python3.11/site-packages/dipy/testing/decorators.py", line 201, in wrapper
    return convert_positional_to_keyword(func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs01/software/imaging/miniconda3/envs/fproc/lib/python3.11/site-packages/dipy/testing/decorators.py", line 192, in convert_positional_to_keyword
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs01/software/imaging/miniconda3/envs/fproc/lib/python3.11/site-packages/dipy/reconst/dti.py", line 2372, in lower_triangular
    D[..., 6] = -np.log(b0)
    ~^^^^^^^^
ValueError: could not broadcast input array from shape (6751,) into shape (89936,)

Additional information / screenshots

Data is 116x116x81x127.

May be related to #79.