pytme 0.1.9__cp311-cp311-macosx_14_0_arm64.whl → 0.2.0__cp311-cp311-macosx_14_0_arm64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (42) hide show
  1. pytme-0.2.0.data/scripts/match_template.py +1019 -0
  2. pytme-0.2.0.data/scripts/postprocess.py +570 -0
  3. {pytme-0.1.9.data → pytme-0.2.0.data}/scripts/preprocessor_gui.py +244 -60
  4. {pytme-0.1.9.dist-info → pytme-0.2.0.dist-info}/METADATA +3 -1
  5. pytme-0.2.0.dist-info/RECORD +72 -0
  6. {pytme-0.1.9.dist-info → pytme-0.2.0.dist-info}/WHEEL +1 -1
  7. scripts/extract_candidates.py +218 -0
  8. scripts/match_template.py +459 -218
  9. pytme-0.1.9.data/scripts/match_template.py → scripts/match_template_filters.py +459 -218
  10. scripts/postprocess.py +380 -435
  11. scripts/preprocessor_gui.py +244 -60
  12. scripts/refine_matches.py +218 -0
  13. tme/__init__.py +2 -1
  14. tme/__version__.py +1 -1
  15. tme/analyzer.py +533 -78
  16. tme/backends/cupy_backend.py +80 -15
  17. tme/backends/npfftw_backend.py +35 -6
  18. tme/backends/pytorch_backend.py +15 -7
  19. tme/density.py +173 -78
  20. tme/extensions.cpython-311-darwin.so +0 -0
  21. tme/matching_constrained.py +195 -0
  22. tme/matching_data.py +76 -33
  23. tme/matching_exhaustive.py +354 -225
  24. tme/matching_memory.py +1 -0
  25. tme/matching_optimization.py +753 -649
  26. tme/matching_utils.py +152 -8
  27. tme/orientations.py +561 -0
  28. tme/preprocessing/__init__.py +2 -0
  29. tme/preprocessing/_utils.py +176 -0
  30. tme/preprocessing/composable_filter.py +30 -0
  31. tme/preprocessing/compose.py +52 -0
  32. tme/preprocessing/frequency_filters.py +322 -0
  33. tme/preprocessing/tilt_series.py +967 -0
  34. tme/preprocessor.py +35 -25
  35. tme/structure.py +2 -37
  36. pytme-0.1.9.data/scripts/postprocess.py +0 -625
  37. pytme-0.1.9.dist-info/RECORD +0 -61
  38. {pytme-0.1.9.data → pytme-0.2.0.data}/scripts/estimate_ram_usage.py +0 -0
  39. {pytme-0.1.9.data → pytme-0.2.0.data}/scripts/preprocess.py +0 -0
  40. {pytme-0.1.9.dist-info → pytme-0.2.0.dist-info}/LICENSE +0 -0
  41. {pytme-0.1.9.dist-info → pytme-0.2.0.dist-info}/entry_points.txt +0 -0
  42. {pytme-0.1.9.dist-info → pytme-0.2.0.dist-info}/top_level.txt +0 -0
@@ -1,4 +1,4 @@
1
- #!python
1
+ #!python3
2
2
  """ CLI interface for basic pyTME template matching functions.
3
3
 
4
4
  Copyright (c) 2023 European Molecular Biology Laboratory
@@ -11,14 +11,16 @@ import warnings
11
11
  import importlib.util
12
12
  from sys import exit
13
13
  from time import time
14
+ from typing import Tuple
14
15
  from copy import deepcopy
15
- from os.path import abspath
16
+ from os.path import abspath, exists
16
17
 
17
18
  import numpy as np
18
19
 
19
- from tme import Density, Preprocessor, __version__
20
+ from tme import Density, __version__
20
21
  from tme.matching_utils import (
21
22
  get_rotation_matrices,
23
+ get_rotations_around_vector,
22
24
  compute_parallelization_schedule,
23
25
  euler_from_rotationmatrix,
24
26
  scramble_phases,
@@ -32,6 +34,7 @@ from tme.analyzer import (
32
34
  PeakCallerMaximumFilter,
33
35
  )
34
36
  from tme.backends import backend
37
+ from tme.preprocessing import Compose
35
38
 
36
39
 
37
40
  def get_func_fullname(func) -> str:
@@ -150,77 +153,291 @@ def crop_data(data: Density, cutoff: float, data_mask: Density = None) -> bool:
150
153
  return True
151
154
 
152
155
 
156
+ def parse_rotation_logic(args, ndim):
157
+ if args.angular_sampling is not None:
158
+ rotations = get_rotation_matrices(
159
+ angular_sampling=args.angular_sampling,
160
+ dim=ndim,
161
+ use_optimized_set=not args.no_use_optimized_set,
162
+ )
163
+ if args.angular_sampling >= 180:
164
+ rotations = np.eye(ndim).reshape(1, ndim, ndim)
165
+ return rotations
166
+
167
+ if args.axis_sampling is None:
168
+ args.axis_sampling = args.cone_sampling
169
+
170
+ rotations = get_rotations_around_vector(
171
+ cone_angle=args.cone_angle,
172
+ cone_sampling=args.cone_sampling,
173
+ axis_angle=args.axis_angle,
174
+ axis_sampling=args.axis_sampling,
175
+ n_symmetry=args.axis_symmetry,
176
+ )
177
+ return rotations
178
+
179
+
180
+ # TODO: Think about whether wedge mask should also be added to target
181
+ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Compose]:
182
+ from tme.preprocessing import LinearWhiteningFilter, BandPassFilter
183
+ from tme.preprocessing.tilt_series import (
184
+ Wedge,
185
+ WedgeReconstructed,
186
+ ReconstructFromTilt,
187
+ )
188
+
189
+ template_filter, target_filter = [], []
190
+ if args.tilt_angles is not None:
191
+ try:
192
+ wedge = Wedge.from_file(args.tilt_angles)
193
+ wedge.weight_type = args.tilt_weighting
194
+ if args.tilt_weighting in ("angle", None) and args.ctf_file is None:
195
+ wedge = WedgeReconstructed(
196
+ angles=wedge.angles, weight_wedge=args.tilt_weighting == "angle"
197
+ )
198
+ except FileNotFoundError:
199
+ tilt_step, create_continuous_wedge = None, True
200
+ tilt_start, tilt_stop = args.tilt_angles.split(",")
201
+ if ":" in tilt_stop:
202
+ create_continuous_wedge = False
203
+ tilt_stop, tilt_step = tilt_stop.split(":")
204
+ tilt_start, tilt_stop = float(tilt_start), float(tilt_stop)
205
+ tilt_angles = (tilt_start, tilt_stop)
206
+ if tilt_step is not None:
207
+ tilt_step = float(tilt_step)
208
+ tilt_angles = np.arange(
209
+ -tilt_start, tilt_stop + tilt_step, tilt_step
210
+ ).tolist()
211
+
212
+ if args.tilt_weighting is not None and tilt_step is None:
213
+ raise ValueError(
214
+ "Tilt weighting is not supported for continuous wedges."
215
+ )
216
+ if args.tilt_weighting not in ("angle", None):
217
+ raise ValueError(
218
+ "Tilt weighting schemes other than 'angle' or 'None' require "
219
+ "a specification of electron doses."
220
+ )
221
+
222
+ wedge = Wedge(
223
+ angles=tilt_angles,
224
+ opening_axis=args.wedge_axes[0],
225
+ tilt_axis=args.wedge_axes[1],
226
+ shape=None,
227
+ weight_type=None,
228
+ weights=np.ones_like(tilt_angles),
229
+ )
230
+ if args.tilt_weighting in ("angle", None) and args.ctf_file is None:
231
+ wedge = WedgeReconstructed(
232
+ angles=tilt_angles,
233
+ weight_wedge=args.tilt_weighting == "angle",
234
+ create_continuous_wedge=create_continuous_wedge,
235
+ )
236
+
237
+ wedge.opening_axis = args.wedge_axes[0]
238
+ wedge.tilt_axis = args.wedge_axes[1]
239
+ wedge.sampling_rate = template.sampling_rate
240
+ template_filter.append(wedge)
241
+ if not isinstance(wedge, WedgeReconstructed):
242
+ template_filter.append(ReconstructFromTilt(
243
+ reconstruction_filter = args.reconstruction_filter
244
+ ))
245
+
246
+ if args.ctf_file is not None:
247
+ from tme.preprocessing.tilt_series import CTF
248
+
249
+ ctf = CTF.from_file(args.ctf_file)
250
+ n_tilts_ctfs, n_tils_angles = len(ctf.defocus_x), len(wedge.angles)
251
+ if n_tilts_ctfs != n_tils_angles:
252
+ raise ValueError(
253
+ f"CTF file contains {n_tilts_ctfs} micrographs, but match_template "
254
+ f"recieved {n_tils_angles} tilt angles. Expected one angle "
255
+ "per micrograph."
256
+ )
257
+ ctf.angles = wedge.angles
258
+ ctf.opening_axis, ctf.tilt_axis = args.wedge_axes
259
+
260
+ if isinstance(template_filter[-1], ReconstructFromTilt):
261
+ template_filter.insert(-1, ctf)
262
+ else:
263
+ template_filter.insert(0, ctf)
264
+ template_filter.insert(1, ReconstructFromTilt(
265
+ reconstruction_filter = args.reconstruction_filter
266
+ ))
267
+
268
+ if args.lowpass or args.highpass is not None:
269
+ lowpass, highpass = args.lowpass, args.highpass
270
+ if args.pass_format == "voxel":
271
+ if lowpass is not None:
272
+ lowpass = np.max(np.multiply(lowpass, template.sampling_rate))
273
+ if highpass is not None:
274
+ highpass = np.max(np.multiply(highpass, template.sampling_rate))
275
+ elif args.pass_format == "frequency":
276
+ if lowpass is not None:
277
+ lowpass = np.max(np.divide(template.sampling_rate, lowpass))
278
+ if highpass is not None:
279
+ highpass = np.max(np.divide(template.sampling_rate, highpass))
280
+
281
+ bandpass = BandPassFilter(
282
+ use_gaussian=args.no_pass_smooth,
283
+ lowpass=lowpass,
284
+ highpass=highpass,
285
+ sampling_rate=template.sampling_rate,
286
+ )
287
+ template_filter.append(bandpass)
288
+ target_filter.append(bandpass)
289
+
290
+ if args.whiten_spectrum:
291
+ whitening_filter = LinearWhiteningFilter()
292
+ template_filter.append(whitening_filter)
293
+ target_filter.append(whitening_filter)
294
+
295
+ template_filter = Compose(template_filter) if len(template_filter) else None
296
+ target_filter = Compose(target_filter) if len(target_filter) else None
297
+
298
+ return template_filter, target_filter
299
+
300
+
153
301
  def parse_args():
154
302
  parser = argparse.ArgumentParser(description="Perform template matching.")
155
- parser.add_argument(
303
+
304
+ io_group = parser.add_argument_group("Input / Output")
305
+ io_group.add_argument(
156
306
  "-m",
157
307
  "--target",
158
308
  dest="target",
159
309
  type=str,
160
310
  required=True,
161
- help="Path to a target in CCP4/MRC format.",
311
+ help="Path to a target in CCP4/MRC, EM, H5 or another format supported by "
312
+ "tme.density.Density.from_file "
313
+ "https://kosinskilab.github.io/pyTME/reference/api/tme.density.Density.from_file.html",
162
314
  )
163
- parser.add_argument(
315
+ io_group.add_argument(
164
316
  "--target_mask",
165
317
  dest="target_mask",
166
318
  type=str,
167
319
  required=False,
168
- help="Path to a mask for the target target in CCP4/MRC format.",
169
- )
170
- parser.add_argument(
171
- "--cutoff_target",
172
- dest="cutoff_target",
173
- type=float,
174
- required=False,
175
- help="Target contour level (used for cropping).",
176
- default=None,
320
+ help="Path to a mask for the target in a supported format (see target).",
177
321
  )
178
- parser.add_argument(
179
- "--cutoff_template",
180
- dest="cutoff_template",
181
- type=float,
182
- required=False,
183
- help="Template contour level (used for cropping).",
184
- default=None,
185
- )
186
- parser.add_argument(
187
- "--no_centering",
188
- dest="no_centering",
189
- action="store_true",
190
- help="If set, assumes the template is centered and omits centering.",
191
- )
192
- parser.add_argument(
322
+ io_group.add_argument(
193
323
  "-i",
194
324
  "--template",
195
325
  dest="template",
196
326
  type=str,
197
327
  required=True,
198
- help="Path to a template in PDB/MMCIF or CCP4/MRC format.",
328
+ help="Path to a template in PDB/MMCIF or other supported formats (see target).",
199
329
  )
200
- parser.add_argument(
330
+ io_group.add_argument(
201
331
  "--template_mask",
202
332
  dest="template_mask",
203
333
  type=str,
204
334
  required=False,
205
- help="Path to a mask for the template in CCP4/MRC format.",
335
+ help="Path to a mask for the template in a supported format (see target).",
206
336
  )
207
- parser.add_argument(
337
+ io_group.add_argument(
208
338
  "-o",
339
+ "--output",
209
340
  dest="output",
210
341
  type=str,
211
342
  required=False,
212
343
  default="output.pickle",
213
- help="Path to output pickle file.",
344
+ help="Path to the output pickle file.",
345
+ )
346
+ io_group.add_argument(
347
+ "--invert_target_contrast",
348
+ dest="invert_target_contrast",
349
+ action="store_true",
350
+ default=False,
351
+ help="Invert the target's contrast and rescale linearly between zero and one. "
352
+ "This option is intended for targets where templates to-be-matched have "
353
+ "negative values, e.g. tomograms.",
354
+ )
355
+ io_group.add_argument(
356
+ "--scramble_phases",
357
+ dest="scramble_phases",
358
+ action="store_true",
359
+ default=False,
360
+ help="Phase scramble the template to generate a noise score background.",
214
361
  )
215
- parser.add_argument(
362
+
363
+ scoring_group = parser.add_argument_group("Scoring")
364
+ scoring_group.add_argument(
216
365
  "-s",
217
366
  dest="score",
218
367
  type=str,
219
368
  default="FLCSphericalMask",
369
+ choices=list(MATCHING_EXHAUSTIVE_REGISTER.keys()),
220
370
  help="Template matching scoring function.",
221
- choices=MATCHING_EXHAUSTIVE_REGISTER.keys(),
222
371
  )
223
- parser.add_argument(
372
+ scoring_group.add_argument(
373
+ "-p",
374
+ dest="peak_calling",
375
+ action="store_true",
376
+ default=False,
377
+ help="Perform peak calling instead of score aggregation.",
378
+ )
379
+
380
+ angular_group = parser.add_argument_group("Angular Sampling")
381
+ angular_exclusive = angular_group.add_mutually_exclusive_group(required=True)
382
+
383
+ angular_exclusive.add_argument(
384
+ "-a",
385
+ dest="angular_sampling",
386
+ type=check_positive,
387
+ default=None,
388
+ help="Angular sampling rate using optimized rotational sets."
389
+ "A lower number yields more rotations. Values >= 180 sample only the identity.",
390
+ )
391
+ angular_exclusive.add_argument(
392
+ "--cone_angle",
393
+ dest="cone_angle",
394
+ type=check_positive,
395
+ default=None,
396
+ help="Half-angle of the cone to be sampled in degrees. Allows to sample a "
397
+ "narrow interval around a known orientation, e.g. for surface oversampling.",
398
+ )
399
+ angular_group.add_argument(
400
+ "--cone_sampling",
401
+ dest="cone_sampling",
402
+ type=check_positive,
403
+ default=None,
404
+ help="Sampling rate of the cone in degrees.",
405
+ )
406
+ angular_group.add_argument(
407
+ "--axis_angle",
408
+ dest="axis_angle",
409
+ type=check_positive,
410
+ default=360.0,
411
+ required=False,
412
+ help="Sampling angle along the z-axis of the cone. Defaults to 360.",
413
+ )
414
+ angular_group.add_argument(
415
+ "--axis_sampling",
416
+ dest="axis_sampling",
417
+ type=check_positive,
418
+ default=None,
419
+ required=False,
420
+ help="Sampling rate along the z-axis of the cone. Defaults to --cone_sampling.",
421
+ )
422
+ angular_group.add_argument(
423
+ "--axis_symmetry",
424
+ dest="axis_symmetry",
425
+ type=check_positive,
426
+ default=1,
427
+ required=False,
428
+ help="N-fold symmetry around z-axis of the cone.",
429
+ )
430
+ angular_group.add_argument(
431
+ "--no_use_optimized_set",
432
+ dest="no_use_optimized_set",
433
+ action="store_true",
434
+ default=False,
435
+ required=False,
436
+ help="Whether to use random uniform instead of optimized rotation sets.",
437
+ )
438
+
439
+ computation_group = parser.add_argument_group("Computation")
440
+ computation_group.add_argument(
224
441
  "-n",
225
442
  dest="cores",
226
443
  required=False,
@@ -228,7 +445,24 @@ def parse_args():
228
445
  default=4,
229
446
  help="Number of cores used for template matching.",
230
447
  )
231
- parser.add_argument(
448
+ computation_group.add_argument(
449
+ "--use_gpu",
450
+ dest="use_gpu",
451
+ action="store_true",
452
+ default=False,
453
+ help="Whether to perform computations on the GPU.",
454
+ )
455
+ computation_group.add_argument(
456
+ "--gpu_indices",
457
+ dest="gpu_indices",
458
+ type=str,
459
+ default=None,
460
+ help="Comma-separated list of GPU indices to use. For example,"
461
+ " 0,1 for the first and second GPU. Only used if --use_gpu is set."
462
+ " If not provided but --use_gpu is set, CUDA_VISIBLE_DEVICES will"
463
+ " be respected.",
464
+ )
465
+ computation_group.add_argument(
232
466
  "-r",
233
467
  "--ram",
234
468
  dest="memory",
@@ -237,168 +471,186 @@ def parse_args():
237
471
  default=None,
238
472
  help="Amount of memory that can be used in bytes.",
239
473
  )
240
- parser.add_argument(
474
+ computation_group.add_argument(
241
475
  "--memory_scaling",
242
476
  dest="memory_scaling",
243
477
  required=False,
244
- type=check_positive,
478
+ type=float,
245
479
  default=0.85,
246
- help="Fraction of available memory that can be used."
247
- "Defaults to 0.85. Ignored if --ram is set",
480
+ help="Fraction of available memory that can be used. Defaults to 0.85 and is "
481
+ "ignored if --ram is set",
248
482
  )
249
- parser.add_argument(
250
- "-a",
251
- dest="angular_sampling",
252
- type=check_positive,
253
- default=40.0,
254
- help="Angular sampling rate for template matching. "
255
- "A lower number yields more rotations. Values >= 180 sample only the identity.",
483
+ computation_group.add_argument(
484
+ "--temp_directory",
485
+ dest="temp_directory",
486
+ default=None,
487
+ help="Directory for temporary objects. Faster I/O improves runtime.",
256
488
  )
257
- parser.add_argument(
258
- "-p",
259
- dest="peak_calling",
260
- action="store_true",
261
- default=False,
262
- help="When set perform peak calling instead of score aggregation.",
489
+
490
+ filter_group = parser.add_argument_group("Filters")
491
+ filter_group.add_argument(
492
+ "--lowpass",
493
+ dest="lowpass",
494
+ type=float,
495
+ required=False,
496
+ help="Resolution to lowpass filter template and target to in the same unit "
497
+ "as the sampling rate of template and target (typically Ångstrom).",
263
498
  )
264
- parser.add_argument(
265
- "--use_gpu",
266
- dest="use_gpu",
499
+ filter_group.add_argument(
500
+ "--highpass",
501
+ dest="highpass",
502
+ type=float,
503
+ required=False,
504
+ help="Resolution to highpass filter template and target to in the same unit "
505
+ "as the sampling rate of template and target (typically Ångstrom).",
506
+ )
507
+ filter_group.add_argument(
508
+ "--no_pass_smooth",
509
+ dest="no_pass_smooth",
510
+ action="store_false",
511
+ default=True,
512
+ help="Whether a hard edge filter should be used for --lowpass and --highpass."
513
+ )
514
+ filter_group.add_argument(
515
+ "--pass_format",
516
+ dest="pass_format",
517
+ type=str,
518
+ required=False,
519
+ choices=["sampling_rate", "voxel", "frequency"],
520
+ help="How values passed to --lowpass and --highpass should be interpreted. "
521
+ "By default, they are assumed to be in units of sampling rate, e.g. Ångstrom."
522
+ )
523
+ filter_group.add_argument(
524
+ "--whiten_spectrum",
525
+ dest="whiten_spectrum",
267
526
  action="store_true",
268
- default=False,
269
- help="Whether to perform computations on the GPU.",
527
+ default=None,
528
+ help="Apply spectral whitening to template and target based on target spectrum.",
270
529
  )
271
- parser.add_argument(
272
- "--gpu_indices",
273
- dest="gpu_indices",
530
+ filter_group.add_argument(
531
+ "--wedge_axes",
532
+ dest="wedge_axes",
274
533
  type=str,
534
+ required=False,
275
535
  default=None,
276
- help="Comma-separated list of GPU indices to use. For example,"
277
- " 0,1 for the first and second GPU. Only used if --use_gpu is set."
278
- " If not provided but --use_gpu is set, CUDA_VISIBLE_DEVICES will"
279
- " be respected.",
536
+ help="Indices of wedge opening and tilt axis, e.g. 0,2 for a wedge that is open "
537
+ "in z-direction and tilted over the x axis.",
280
538
  )
281
- parser.add_argument(
282
- "--invert_target_contrast",
283
- dest="invert_target_contrast",
539
+ filter_group.add_argument(
540
+ "--tilt_angles",
541
+ dest="tilt_angles",
542
+ type=str,
543
+ required=False,
544
+ default=None,
545
+ help="Path to a tab-separated file containing the column angles and optionally "
546
+ " weights, or comma separated start and stop stage tilt angle, e.g. 50,45, which "
547
+ " yields a continuous wedge mask. Alternatively, a tilt step size can be "
548
+ "specified like 50,45:5.0 to sample 5.0 degree tilt angle steps.",
549
+ )
550
+ filter_group.add_argument(
551
+ "--tilt_weighting",
552
+ dest="tilt_weighting",
553
+ type=str,
554
+ required=False,
555
+ choices=["angle", "relion", "grigorieff"],
556
+ default=None,
557
+ help="Weighting scheme used to reweight individual tilts. Available options: "
558
+ "angle (cosine based weighting), "
559
+ "relion (relion formalism for wedge weighting) requires,"
560
+ "grigorieff (exposure filter as defined in Grant and Grigorieff 2015)."
561
+ "relion and grigorieff require electron doses in --tilt_angles weights column.",
562
+ )
563
+ # filter_group.add_argument(
564
+ # "--ctf_file",
565
+ # dest="ctf_file",
566
+ # type=str,
567
+ # required=False,
568
+ # default=None,
569
+ # help="Path to a file with CTF parameters from CTFFIND4.",
570
+ # )
571
+ filter_group.add_argument(
572
+ "--reconstruction_filter",
573
+ dest="reconstruction_filter",
574
+ type=str,
575
+ required=False,
576
+ choices = ["ram-lak", "ramp", "shepp-logan", "cosine", "hamming"],
577
+ default=None,
578
+ help="Filter applied when reconstructing (N+1)-D from N-D filters.",
579
+ )
580
+
581
+ performance_group = parser.add_argument_group("Performance")
582
+ performance_group.add_argument(
583
+ "--cutoff_target",
584
+ dest="cutoff_target",
585
+ type=float,
586
+ required=False,
587
+ default=None,
588
+ help="Target contour level (used for cropping).",
589
+ )
590
+ performance_group.add_argument(
591
+ "--cutoff_template",
592
+ dest="cutoff_template",
593
+ type=float,
594
+ required=False,
595
+ default=None,
596
+ help="Template contour level (used for cropping).",
597
+ )
598
+ performance_group.add_argument(
599
+ "--no_centering",
600
+ dest="no_centering",
284
601
  action="store_true",
285
- default=False,
286
- help="Invert the target contrast via multiplication with negative one and"
287
- " linear rescaling between zero and one. Note that this might lead to"
288
- " different baseline scores of individual target splits when using"
289
- " unnormalized scores. This option is intended for targets, where the"
290
- " object to-be-matched has negative values, i.e. tomograms.",
602
+ help="Assumes the template is already centered and omits centering.",
291
603
  )
292
- parser.add_argument(
604
+ performance_group.add_argument(
293
605
  "--no_edge_padding",
294
606
  dest="no_edge_padding",
295
607
  action="store_true",
296
608
  default=False,
297
- help="Whether to pad the edges of the target. This is useful, if the target"
298
- " has a well defined bounding box, e.g. a density map.",
609
+ help="Whether to not pad the edges of the target. Can be set if the target"
610
+ " has a well defined bounding box, e.g. a masked reconstruction.",
299
611
  )
300
- parser.add_argument(
612
+ performance_group.add_argument(
301
613
  "--no_fourier_padding",
302
614
  dest="no_fourier_padding",
303
615
  action="store_true",
304
616
  default=False,
305
- help="Whether input arrays should be zero-padded to the full convolution shape"
306
- " for numerical stability. When working with very large targets such as"
307
- " tomograms it is safe to use this flag and benefit from the performance gain.",
617
+ help="Whether input arrays should not be zero-padded to full convolution shape "
618
+ "for numerical stability. When working with very large targets, e.g. tomograms, "
619
+ "it is safe to use this flag and benefit from the performance gain.",
308
620
  )
309
- parser.add_argument(
310
- "--scramble_phases",
311
- dest="scramble_phases",
312
- action="store_true",
313
- default=False,
314
- help="Whether to phase scramble the template for subsequent normalization.",
315
- )
316
- parser.add_argument(
621
+ performance_group.add_argument(
317
622
  "--interpolation_order",
318
623
  dest="interpolation_order",
319
624
  required=False,
320
625
  type=int,
321
626
  default=3,
322
- help="Spline interpolation used during rotations. If less than zero"
323
- " no interpolation is performed.",
627
+ help="Spline interpolation used for template rotations. If less than zero "
628
+ "no interpolation is performed.",
324
629
  )
325
- parser.add_argument(
630
+ performance_group.add_argument(
326
631
  "--use_mixed_precision",
327
632
  dest="use_mixed_precision",
328
633
  action="store_true",
329
634
  default=False,
330
635
  help="Use float16 for real values operations where possible.",
331
636
  )
332
- parser.add_argument(
637
+ performance_group.add_argument(
333
638
  "--use_memmap",
334
639
  dest="use_memmap",
335
640
  action="store_true",
336
641
  default=False,
337
- help="Use memmaps to offload large data objects to disk. This is"
338
- " particularly useful for large inputs when using --use_gpu..",
339
- )
340
- parser.add_argument(
341
- "--temp_directory",
342
- dest="temp_directory",
343
- default=None,
344
- help="Directory for temporary objects. Faster I/O typically improves runtime.",
345
- )
346
- parser.add_argument(
347
- "--gaussian_sigma",
348
- dest="gaussian_sigma",
349
- type=float,
350
- required=False,
351
- help="Sigma parameter for Gaussian filtering the template.",
352
- )
353
- parser.add_argument(
354
- "--bandpass_band",
355
- dest="bandpass_band",
356
- type=str,
357
- required=False,
358
- help="Comma separated start and stop frequency for bandpass filtering the"
359
- " template, e.g. 0.1, 0.5",
360
- )
361
- parser.add_argument(
362
- "--bandpass_smooth",
363
- dest="bandpass_smooth",
364
- type=float,
365
- required=False,
366
- default=None,
367
- help="Smooth parameter for the bandpass filter.",
368
- )
369
- parser.add_argument(
370
- "--tilt_range",
371
- dest="tilt_range",
372
- type=str,
373
- required=False,
374
- help="Comma separated start and stop stage tilt angle, e.g. '50,45'. Used"
375
- " to create a wedge mask to be applied to the template.",
642
+ help="Use memmaps to offload large data objects to disk. "
643
+ "Particularly useful for large inputs in combination with --use_gpu.",
376
644
  )
377
- parser.add_argument(
378
- "--tilt_step",
379
- dest="tilt_step",
380
- type=float,
381
- required=False,
382
- default=None,
383
- help="Step size between tilts, e.g. '5'. When set the wedge mask"
384
- " reflects the individual tilts, otherwise a continuous mask is used.",
385
- )
386
- parser.add_argument(
387
- "--wedge_axes",
388
- dest="wedge_axes",
389
- type=str,
645
+
646
+ analyzer_group = parser.add_argument_group("Analyzer")
647
+ analyzer_group.add_argument(
648
+ "--score_threshold",
649
+ dest="score_threshold",
390
650
  required=False,
391
- default="0,2",
392
- help="Axis index of wedge opening and tilt axis, e.g. 0,2 for a wedge that is open in"
393
- " z and tilted over x.",
394
- )
395
- parser.add_argument(
396
- "--wedge_smooth",
397
- dest="wedge_smooth",
398
651
  type=float,
399
- required=False,
400
- default=None,
401
- help="Gaussian sigma used to smooth the wedge mask.",
652
+ default=0,
653
+ help="Minimum template matching scores to consider for analysis.",
402
654
  )
403
655
 
404
656
  args = parser.parse_args()
@@ -406,6 +658,8 @@ def parse_args():
406
658
  if args.interpolation_order < 0:
407
659
  args.interpolation_order = None
408
660
 
661
+ args.ctf_file = None
662
+
409
663
  if args.temp_directory is None:
410
664
  default = abspath(".")
411
665
  if os.environ.get("TMPDIR", None) is not None:
@@ -438,6 +692,21 @@ def parse_args():
438
692
  int(x) for x in os.environ["CUDA_VISIBLE_DEVICES"].split(",")
439
693
  ]
440
694
 
695
+ if args.tilt_angles is not None:
696
+ if args.wedge_axes is None:
697
+ raise ValueError("Need to specify --wedge_axes when --tilt_angles is set.")
698
+ if not exists(args.tilt_angles):
699
+ try:
700
+ float(args.tilt_angles.split(",")[0])
701
+ except ValueError:
702
+ raise ValueError(f"{args.tilt_angles} is not a file nor a range.")
703
+
704
+ if args.ctf_file is not None and args.tilt_angles is None:
705
+ raise ValueError("Need to specify --tilt_angles when --ctf_file is set.")
706
+
707
+ if args.wedge_axes is not None:
708
+ args.wedge_axes = tuple(int(i) for i in args.wedge_axes.split(","))
709
+
441
710
  return args
442
711
 
443
712
 
@@ -514,51 +783,6 @@ def main():
514
783
  },
515
784
  )
516
785
 
517
- template_filter = {}
518
- if args.gaussian_sigma is not None:
519
- template.data = Preprocessor().gaussian_filter(
520
- sigma=args.gaussian_sigma, template=template.data
521
- )
522
-
523
- if args.bandpass_band is not None:
524
- bandpass_start, bandpass_stop = [
525
- float(x) for x in args.bandpass_band.split(",")
526
- ]
527
- if args.bandpass_smooth is None:
528
- args.bandpass_smooth = 0
529
-
530
- template_filter["bandpass_mask"] = {
531
- "minimum_frequency": bandpass_start,
532
- "maximum_frequency": bandpass_stop,
533
- "gaussian_sigma": args.bandpass_smooth,
534
- }
535
-
536
- if args.tilt_range is not None:
537
- args.wedge_smooth if args.wedge_smooth is not None else 0
538
- tilt_start, tilt_stop = [float(x) for x in args.tilt_range.split(",")]
539
- opening_axis, tilt_axis = [int(x) for x in args.wedge_axes.split(",")]
540
-
541
- if args.tilt_step is not None:
542
- template_filter["step_wedge_mask"] = {
543
- "start_tilt": tilt_start,
544
- "stop_tilt": tilt_stop,
545
- "tilt_step": args.tilt_step,
546
- "sigma": args.wedge_smooth,
547
- "opening_axis": opening_axis,
548
- "tilt_axis": tilt_axis,
549
- "omit_negative_frequencies": True,
550
- }
551
- else:
552
- template_filter["continuous_wedge_mask"] = {
553
- "start_tilt": tilt_start,
554
- "stop_tilt": tilt_stop,
555
- "tilt_axis": tilt_axis,
556
- "opening_axis": opening_axis,
557
- "infinite_plane": True,
558
- "sigma": args.wedge_smooth,
559
- "omit_negative_frequencies": True,
560
- }
561
-
562
786
  if template_mask is None:
563
787
  template_mask = template.empty
564
788
  if not args.no_centering:
@@ -672,21 +896,13 @@ def main():
672
896
  )
673
897
  exit(-1)
674
898
 
675
- analyzer_args = {
676
- "score_threshold": 0,
677
- "number_of_peaks": 1000,
678
- "convolution_mode": "valid",
679
- "use_memmap": args.use_memmap,
680
- }
681
-
682
899
  matching_setup, matching_score = MATCHING_EXHAUSTIVE_REGISTER[args.score]
683
900
  matching_data = MatchingData(target=target, template=template.data)
684
- matching_data.rotations = get_rotation_matrices(
685
- angular_sampling=args.angular_sampling, dim=target.data.ndim
686
- )
687
- if args.angular_sampling >= 180:
688
- ndim = target.data.ndim
689
- matching_data.rotations = np.eye(ndim).reshape(1, ndim, ndim)
901
+ matching_data.rotations = parse_rotation_logic(args=args, ndim=target.data.ndim)
902
+
903
+ template_filter, target_filter = setup_filter(args, template, target)
904
+ matching_data.template_filter = template_filter
905
+ matching_data.target_filter = target_filter
690
906
 
691
907
  matching_data.template_filter = template_filter
692
908
  matching_data._invert_target = args.invert_target_contrast
@@ -724,10 +940,35 @@ def main():
724
940
  label_width=max(len(key) for key in options.keys()) + 2,
725
941
  )
726
942
 
727
- options = {"Analyzer": callback_class, **analyzer_args}
943
+ filter_args = {
944
+ "Lowpass": args.lowpass,
945
+ "Highpass": args.highpass,
946
+ "Smooth Pass": args.no_pass_smooth,
947
+ "Pass Format" : args.pass_format,
948
+ "Spectral Whitening": args.whiten_spectrum,
949
+ "Wedge Axes": args.wedge_axes,
950
+ "Tilt Angles": args.tilt_angles,
951
+ "Tilt Weighting": args.tilt_weighting,
952
+ "CTF": args.ctf_file,
953
+ }
954
+ filter_args = {k: v for k, v in filter_args.items() if v is not None}
955
+ if len(filter_args):
956
+ print_block(
957
+ name="Filters",
958
+ data=filter_args,
959
+ label_width=max(len(key) for key in options.keys()) + 2,
960
+ )
961
+
962
+ analyzer_args = {
963
+ "score_threshold": args.score_threshold,
964
+ "number_of_peaks": 1000,
965
+ "convolution_mode": "valid",
966
+ "use_memmap": args.use_memmap,
967
+ }
968
+ analyzer_args = {"Analyzer": callback_class, **analyzer_args}
728
969
  print_block(
729
970
  name="Score Analysis Options",
730
- data=options,
971
+ data=analyzer_args,
731
972
  label_width=max(len(key) for key in options.keys()) + 2,
732
973
  )
733
974
  print("\n" + "-" * 80)