ml-struct-bio/cryodrgn

Filter by indices gives an error

Closed this issue · 4 comments

Hi,

I am trying to use a pkl file that I made in jupyter to train a new model on a subset of particles. The pkl file works when I try filtering ctf.pkl or the starfile but when I use it for train_vae I gewt this error:

export CUDA_VISIBLE_DEVICES="1"; cryodrgn train_vae M09_b4_19tilt.star  --ctf ctf.pkl --poses pose.pkl --datadir particleseries  --zdim 8 -n 100 --beta .025 -o output_19tilt_50S_lower --encode-mode tilt --ntilt 19 --ind output_19tilt/epoch16_50S_lower_ind_selected.pkl
(INFO) (train_vae.py) (16-Apr-24 14:22:56) /home/janlevinson/micromamba/envs/cryodrgn/bin/cryodrgn train_vae M09_b4_19tilt.star --ctf ctf.pkl --poses pose.pkl --datadir particleseries --zdim 8 -n 100 --beta .025 -o output_19tilt_50S_lower --encode-mode tilt --ntilt 19 --ind output_19tilt/epoch16_50S_lower_ind_selected.pkl
(INFO) (train_vae.py) (16-Apr-24 14:22:56) cryoDRGN 3.0.0b1.dev63+g4b7703d
(INFO) (train_vae.py) (16-Apr-24 14:22:56) Namespace(particles='/scratch1/users/hamid/b4_crydrgn_new/06_joined_redo/M09_b4_19tilt.star', outdir='/scratch1/users/hamid/b4_crydrgn_new/06_joined_redo/output_19tilt_50S_lower', zdim=8, poses='/scratch1/
users/hamid/b4_crydrgn_new/06_joined_redo/pose.pkl', ctf='/scratch1/users/hamid/b4_crydrgn_new/06_joined_redo/ctf.pkl', load=None, checkpoint=1, log_interval=1000, verbose=False, seed=14278, ind='/scratch1/users/hamid/b4_crydrgn_new/06_joined_redo/output_19tilt/epoch16_50S_lower_ind_selected.pkl', invert_data=True, window=True, window_r=0.85, datadir='/scratch1/users/hamid/b4_crydrgn_new/06_joined_redo/particleseries', lazy=False, shuffler_size=0, num_workers=0, max_threads=16, ntilts=19, random_tilts=False, t_emb_dim=64, tlayers=3, tdim=1024, dose_per_tilt=None, angle_per_tilt=3, num_epochs=100, batch_size=8, wd=0, lr=0.0001, beta='.025', beta_control=None, norm=None, amp=True, multigpu=False, do_pose_sgd=False, pretrain=1, emb_type='quat', pose_lr=0.0003, qlayers=3, qdim=1024, encode_mode='tilt', enc_mask=None, use_real=False, players=3, pdim=1024, pe_type='gaussian', feat_sigma=0.5, pe_dim=None, domain='fourier', activation='relu', func=<function main at 0x787af41f0dc0>)
(INFO) (train_vae.py) (16-Apr-24 14:22:56) Use cuda True
(INFO) (train_vae.py) (16-Apr-24 14:22:56) Filtering image dataset with /scratch1/users/hamid/b4_crydrgn_new/06_joined_redo/output_19tilt/epoch16_50S_lower_ind_selected.pkl
(INFO) (train_vae.py) (16-Apr-24 14:23:50) Loading dataset from /scratch1/users/hamid/b4_crydrgn_new/06_joined_redo/M09_b4_19tilt.star
/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3483.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
(INFO) (dataset.py) (16-Apr-24 14:24:55) Normalizing HT by 0 +/- 64.13177490234375
(INFO) (dataset.py) (16-Apr-24 14:25:47) Loaded 1192896 tilts for 62784 particles
(INFO) (dataset.py) (16-Apr-24 14:25:47) {19} tilts per particle
/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/cryodrgn/dataset.py:170: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  self.tilt_angles = torch.tensor(self.tilt_angles).to(self.device)
Traceback (most recent call last):
  File "/home/janlevinson/micromamba/envs/cryodrgn/bin/cryodrgn", line 8, in <module>
    sys.exit(main_commands())
  File "/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/cryodrgn/command_line.py", line 65, in main_commands
    _get_commands(
  File "/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/cryodrgn/command_line.py", line 60, in _get_commands
    args.func(args)
  File "/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/cryodrgn/commands/train_vae.py", line 711, in main
    posetracker = PoseTracker.load(
  File "/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/cryodrgn/pose.py", line 91, in load
    rots = rots[ind]
IndexError: arrays used as indices must be of integer (or boolean) type

I dont understand how this can happen since the command

cryodrgn_utils filter_pkl pose.pkl --ind output_19tilt/epoch16_50S_all_ind_selected.pkl -o pose.epoch16_50S_all.pkl

works without errors.

Could the problem come from that I have one stack per tomogram coming from M and not a single stack with all my particles?

Thanks in advance,
Hamid

Hi Hamid, thanks for bringing this to our attention!

This seems to be specifically caused by the way the interactive Jupyter widget saves the chosen indices: as an array of numpy.int objects, as opposed to native Python integers. The latter can be individually used as Python list indices, but the former cannot, and thus rots[ind] fails in your case. I can replicate this error on my end by manually creating an array of numpy integers and passing it to train_vae — the same arguments but with an array of native Python integers runs error-free.

We are working on a fix to this and are hoping to include it in an upcoming release! In the meantime, you can try making sure that your indices are being saved as an array of Python integers as opposed to numpy integers:

import pickle
with open("output_19tilt/epoch16_50S_lower_ind_selected.pkl", 'rb') as f:
    indices = pickle.load(f)

print(type(indices[0])) # confirm that indices are of type numpy.int32
indices = [int(x) for x in indices]

with open("output_19tilt/epoch16_50S_lower_ind_selected.pkl", 'wb') as f:
    pickle.dump(indices, f)

Hope this helps!
-Mike

Thank you, that makes sense and I made this little code to fix the pkl file:

#/usr/bin/env python3
import pickle

#using click to make this code run from he commandline
import click

@click.command()
@click.option('--input', help='input file')
@click.option('--output', help='output file')
def main(input, output):
    with open(input, 'rb') as f:
        indices = pickle.load(f)

    print("Input type:" ,type(indices[0]))
    indices = [int(x) for x in indices]
    print("Output type:" ,type(indices[0]))
    
    with open(output, 'wb') as f:
        pickle.dump(indices, f)

if __name__ == '__main__':
    main()
    

and when I run it I get:
python3 ./fixind.py --input output_19tilt_z4/ep30_center_ind_selected.pkl --output output_19tilt_z4/ep30_center_ind_selected_fixed.pkl

Input type: <class 'numpy.int64'> Output type: <class 'int'>

But weirdly I am still getting the same error!

export CUDA_VISIBLE_DEVICES="3"; cryodrgn train_vae M09_b4_19tilt_stack_b4.star  --ctf ctf.pkl --poses pose.pkl --datadir ./  --zdim 4 -n 100 --beta .025 -o step2_center_output_19tilt_z4 --encode-mode tilt --ntilt 19 --ind output_19tilt_z4/ep30_center_ind_selected_fixed.pkl
(INFO) (train_vae.py) (29-Apr-24 11:15:29) /home/janlevinson/micromamba/envs/cryodrgn/bin/cryodrgn train_vae M09_b4_19tilt_stack_b4.star --ctf ctf.pkl --poses pose.pkl --datadir ./ --zdim 4 -n 100 --beta .025 -o step2_center_output_19tilt_z4 --encode-mode tilt --ntilt 19 --ind output_19tilt_z4/ep30_center_ind_selected_fixed.pkl
(INFO) (train_vae.py) (29-Apr-24 11:15:29) cryoDRGN 3.0.0b1.dev63+g4b7703d
(INFO) (train_vae.py) (29-Apr-24 11:15:29) Namespace(particles='/scratch1/users/hamid/b4_crydrgn_new/07_redo_stack/M09_b4_19tilt_stack_b4.star', outdir='/scratch1/users/hamid/b4_crydrgn_new/07_redo_stack/step2_center_output_19tilt_z4', zdim=4, poses='/scratch1/users/hamid/b4_crydrgn_new/07_redo_stack/pose.pkl', ctf='/scratch1/users/hamid/b4_crydrgn_new/07_redo_stack/ctf.pkl', load=None, checkpoint=1, log_interval=1000, verbose=False, seed=56442, ind='/scratch1/users/hamid/b4_crydrgn_new/07_redo_stack/output_19tilt_z4/ep30_center_ind_selected_fixed.pkl', invert_data=True, window=True, window_r=0.85, datadir='/scratch1/users/hamid/b4_crydrgn_new/07_redo_stack', lazy=False, shuffler_size=0, num_workers=0, max_threads=16, ntilts=19, random_tilts=False, t_emb_dim=64, tlayers=3, tdim=1024, dose_per_tilt=None, angle_per_tilt=3, num_epochs=100, batch_size=8, wd=0, lr=0.0001, beta='.025', beta_control=None, norm=None, amp=True, multigpu=False, do_pose_sgd=False, pretrain=1, emb_type='quat', pose_lr=0.0003, qlayers=3, qdim=1024, encode_mode='tilt', enc_mask=None, use_real=False, players=3, pdim=1024, pe_type='gaussian', feat_sigma=0.5, pe_dim=None, domain='fourier', activation='relu', func=<function main at 0x7a1f58e9cee0>)
(INFO) (train_vae.py) (29-Apr-24 11:15:29) Use cuda True
(INFO) (train_vae.py) (29-Apr-24 11:15:29) Filtering image dataset with /scratch1/users/hamid/b4_crydrgn_new/07_redo_stack/output_19tilt_z4/ep30_center_ind_selected_fixed.pkl
^[[A^[[A^[[A^[[B^[[B^[[B^[[B^[[B^[[B(INFO) (train_vae.py) (29-Apr-24 11:16:08) Loading dataset from /scratch1/users/hamid/b4_crydrgn_new/07_redo_stack/M09_b4_19tilt_stack_b4.star
/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3483.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
(INFO) (dataset.py) (29-Apr-24 11:16:57) Normalizing HT by 0 +/- 64.12897491455078
(INFO) (dataset.py) (29-Apr-24 11:17:36) Loaded 2575526 tilts for 135554 particles
(INFO) (dataset.py) (29-Apr-24 11:17:36) {19} tilts per particle
/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/cryodrgn/dataset.py:170: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  self.tilt_angles = torch.tensor(self.tilt_angles).to(self.device)
Traceback (most recent call last):
  File "/home/janlevinson/micromamba/envs/cryodrgn/bin/cryodrgn", line 8, in <module>
    sys.exit(main_commands())
  File "/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/cryodrgn/command_line.py", line 65, in main_commands
    _get_commands(
  File "/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/cryodrgn/command_line.py", line 60, in _get_commands
    args.func(args)
  File "/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/cryodrgn/commands/train_vae.py", line 711, in main
    posetracker = PoseTracker.load(
  File "/home/janlevinson/micromamba/envs/cryodrgn/lib/python3.10/site-packages/cryodrgn/pose.py", line 91, in load
    rots = rots[ind]
IndexError: arrays used as indices must be of integer (or boolean) type

Hey, I took another look at this and found that the bug was actually caused by unwanted behaviour introduced by an earlier patch to the method that parses indices passed to tilt series reconstruction (6890525). This removed an unwanted dtype=object warning, but also caused the elements of the np.array of particle indices produced by TiltSeriesData.parse_particle_tilt() to be of type object instead of int but only when all particle tilt indices were the same length (e.g. when all chosen particles had the same number of tilts).

I fixed this by making the array of particle indices a list instead, which is compatible with all downstream uses of parse_particle_tilt outputs, and also added tests to specifically cover the case of running train_vae with tilt series inputs with or without different particle tilt counts after subsetting with --ind — see tests/test_integration.TestStarFixedHetero.

This fix will be available in v3.3.1 and an early version can be found on our beta channel:
pip install -i https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ 'cryodrgn>3.3.0' --pre

Let us know if you run into any other problems!

Thank you very much that solved the problem (with no need for the little python code).