Coverage for model_workflow/utils/subsets.py: 73%

1from inspect import getfullargspec 

2 

3from model_workflow.utils.formats import get_format_set_suitable_function 

4from model_workflow.utils.pyt_spells import get_frames_count 

5from model_workflow.utils.gmx_spells import get_trajectory_subset as gmx_get_trajectory_subset 

6from model_workflow.utils.mdt_spells import get_trajectory_subset as mdt_get_trajectory_subset 

7from model_workflow.utils.type_hints import * 

8 

9# Set the functions to perform single frame gettering 

10subset_functions = [ gmx_get_trajectory_subset, mdt_get_trajectory_subset ] 

11 

12# Get specific frames from a trajectory 

13def get_trajectory_subset ( 

14 input_structure_file : 'File', 

15 input_trajectory_file : 'File', 

16 output_trajectory_file : 'File', 

17 start : int = 0, 

18 end : int = None, 

19 step : int = 1, 

20 frames : List[int] = [], 

21 skip : List[int] = [], 

22): 

23 # If there is no end set then the end as the last frame of the simulation 

24 if end == None: 

25 end = get_frames_count(input_structure_file, input_trajectory_file) 

26 

27 # How step and skip are combined is not intuitive and it could be missleading 

28 # For this reason both arguments are not allowed together 

29 if step and step != 1 and skip and len(skip) > 0: 

30 raise SystemExit("Arguments 'step' and 'skip' are not allowed together. Please do it in 2 separated calls.") 

31 

32 # End must be grater than start 

33 if end != None and end < start: 

34 raise SystemExit('End frame must be posterior to start frame') 

35 

36 # We need an output trajectory filename 

37 if not output_trajectory_file: 

38 raise SystemExit('Missing output trajectory filename') 

39 

40 # In case the frames argument is passed 

41 if frames and len(frames) > 0: 

42 print('Specific frames were passed. Other arguments will be ignored (start, end, step and skip)') 

43 if frames != sorted(frames): 

44 print('WARNING: Note that the reduced trajectories will keep frames in their original order, not the input order') 

45 

46 # Get the input formats 

47 input_structure_format = input_structure_file.format 

48 input_trajectory_format = input_trajectory_file.format 

49 output_trajectory_format = output_trajectory_file.format 

50 format_set = { 

51 'inputs': { 

52 'input_structure_filename': None if input_structure_format == None else { input_structure_format }, 

53 'input_trajectory_filename': { input_trajectory_format } 

54 }, 

55 'outputs': { 

56 'output_trajectory_filename': { output_trajectory_format } 

57 } 

58 } 

59 # Get a suitable function to do the job 

60 suitables = next(get_format_set_suitable_function( 

61 available_functions=subset_functions, 

62 available_request_format_sets=[format_set], 

63 ), None) 

64 if not suitables: 

65 raise SystemExit('There is no subset function which supports the requested formats') 

66 suitable_function, formats = suitables 

67 # Call the subset function 

68 # Check if the subset function requires the input structure argument before 

69 suitable_function_keywords = getfullargspec(suitable_function)[0] 

70 required_structure = 'input_structure_filename' in suitable_function_keywords 

71 if required_structure: 

72 suitable_function(input_structure_file.path, input_trajectory_file.path, output_trajectory_file.path, start, end, step, frames, skip) 

73 else: 

74 suitable_function(input_trajectory_file.path, output_trajectory_file.path, start, end, step, frames, skip)