Implementation of external Zernike3D distance

[DavidHerreros]

Implementation of the Zernike3D on the map_to_map pipeline.

The implementation includes an initial approximation on how future external distances could be implemented. Overall, the proposed solution involved:

• Custom fields in the config file for specific external distances to provide specific parameters
• Some checks on specific software availability (conda availability, is software installed...)
• New docs folder including setup instructions for a given external package and tips on how to configure and run the new distance

Regarding tutorials, I don't think it make sense to add test for external distances, as it would involve installing and running different packages on GitHub side. In some cases, installations might not be possible without user input, and in some others the external software might need resources not available on GitHub (such as GPUs). I would rather specify on the documentation that any erros arising when running external distances should be reported to the external package developers.

[geoffwoollard]

write a validator for a config that has the zernike parameters

[DSilva27]

Feel free to use Pydantic if you are familiar with it, the idea is to switch all validation to it in the future, I've just been busy with other things. Check the end of this file

https://github.com/flatironinstitute/Cryo-EM-Heterogeneity-Challenge-1/blob/60-improve-svd-pipeline-with-randomized-svd/src/cryo_challenge/data/_validation/config_validators.py

if you want to get a better idea of how it works.

[geoffwoollard]

@DavidHerreros I am getting the following error. What is some sample data that this should run on? Perhaps it's just some shape / pixel size differences in the test data...

(flexutils-tensorflow) (base) [gwoollard@workergpu004 Cryo-EM-Heterogeneity-Challenge-1]$ compute_distance_matrix_zernike3deep.py --references_file tmp_zernike/reference_maps.npy  --targets_file tmp_zernike/target_maps.npy  --out_path tmp_zernike/ --gpu 0 --thr 20
2024-11-14 14:32:32.847240: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-11-14 14:32:32.886429: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-11-14 14:32:32.886456: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-11-14 14:32:32.887769: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-11-14 14:32:32.894664: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-11-14 14:32:33.663224: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Projecting volumes:   0%|                                                                                                                                          | 0/10 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/mnt/home/gwoollard/software/mambaforge/envs/flexutils-tensorflow/bin/compute_distance_matrix_zernike3deep.py", line 8, in <module>
    sys.exit(main())
  File "/mnt/ceph/users/gwoollard/repos/Flexutils-Toolkit/tensorflow_toolkit/scripts/compute_distance_matrix_zernike3deep.py", line 164, in main
    compute_distance_matrix(**inputs)
  File "/mnt/ceph/users/gwoollard/repos/Flexutils-Toolkit/tensorflow_toolkit/scripts/compute_distance_matrix_zernike3deep.py", line 60, in compute_distance_matrix
    volume = np.reshape(volume, (target_bx, target_bx, target_bx))
  File "<__array_function__ internals>", line 180, in reshape
  File "/mnt/home/gwoollard/software/mambaforge/envs/flexutils-tensorflow/lib/python3.9/site-packages/numpy/core/fromnumeric.py", line 298, in reshape
    return _wrapfunc(a, 'reshape', newshape, order=order)
  File "/mnt/home/gwoollard/software/mambaforge/envs/flexutils-tensorflow/lib/python3.9/site-packages/numpy/core/fromnumeric.py", line 57, in _wrapfunc
    return bound(*args, **kwds)
ValueError: cannot reshape array of size 4096 into shape (3,3,3)

[geoffwoollard]

and where the original config (that generates the same error) is

data:
  n_pix: 16
  psize: 30.044
  submission:
    fname: tests/data/dataset_2_submissions/submission_1000.pt
    volume_key: volumes
    metadata_key: populations
    label_key: id
  ground_truth:
    volumes: tests/data/Ground_truth/test_maps_gt_flat_10.pt
    metadata: tests/data/Ground_truth/test_metadata_10.csv
  mask:
    do: true
    volume: tests/data/Ground_truth/test_mask_dilated_wide.mrc
analysis:
  zernike3d_extra_params:
    gpuID: 0
    tmpDir: tmp_zernike
    thr: 20
  metrics:
    - l2
    - zernike3d
  chunk_size_submission: 4
  chunk_size_gt: 5
  low_memory:
    do: false
    chunk_size_low_memory: null
  normalize:
    do: true
    method: median_zscore
output: tests/results/test_map_to_map_distance_matrix_submission_0.pkl

[geoffwoollard]

In [15]: d = np.load('tmp_zernike/reference_maps.npy')

In [16]: d.shape
Out[16]: (8, 4096)

In [17]: d = np.load('tmp_zernike/target_maps.npy')

In [18]: d.shape
Out[18]: (10, 16, 16, 16)

[DavidHerreros]

In [15]: d = np.load('tmp_zernike/reference_maps.npy')

In [16]: d.shape
Out[16]: (8, 4096)

In [17]: d = np.load('tmp_zernike/target_maps.npy')

In [18]: d.shape
Out[18]: (10, 16, 16, 16)

Hi @geoffwoollard

Thanks for the report! I did not consider in the code the possibility of having volumes with variable shapes in the input. Currently, the program assumes that the spatial dimensions are collapsed instead of having a 4D array.

I have added a fix to solve this issue so that the program is no longer restricted by the number of dimensions in the input volumes. Here you have the PR with the fixes.

Could you please change your branch to the one in the PR to check if everything is now working?

[geoffwoollard]

Yes @DavidHerreros , that fix seems to have addressed the tensor shape issue. However, I'm getting another error. I isolated the issue to Flexutils-Toolkit, apart from how it is being externally called here, and made an issue on your repo if you want to discuss there: I2PC/Flexutils-Toolkit#18

[DavidHerreros]

Yes @DavidHerreros , that fix seems to have addressed the tensor shape issue. However, I'm getting another error. I isolated the issue to Flexutils-Toolkit, apart from how it is being externally called here, and made an issue on your repo if you want to discuss there: I2PC/Flexutils-Toolkit#18

Hi @geoffwoollard,

Sorry, I have missed that one...

We can continue discussing the other repo to better isolate the issues as you propose. I have already added some fixes and replied to your comments there.

Thanks for the report!