Coverage for mddb_workflow / utils / subsets.py: 74%

1from inspect import getfullargspec 

2 

3from mddb_workflow.utils.auxiliar import InputError 

4from mddb_workflow.utils.formats import get_format_set_suitable_function 

5from mddb_workflow.utils.pyt_spells import get_frames_count 

6from mddb_workflow.utils.gmx_spells import get_trajectory_subset as gmx_get_trajectory_subset 

7from mddb_workflow.utils.mdt_spells import get_trajectory_subset as mdt_get_trajectory_subset 

8from mddb_workflow.utils.type_hints import * 

9 

10# Set the functions to perform single frame gettering 

11subset_functions = [ gmx_get_trajectory_subset, mdt_get_trajectory_subset ] 

12 

13# Get specific frames from a trajectory 

14def get_trajectory_subset ( 

15 input_structure_file : 'File', 

16 input_trajectory_file : 'File', 

17 output_trajectory_file : 'File', 

18 start : int = 0, 

19 end : int = None, 

20 step : int = 1, 

21 frames : list[int] = [], 

22 skip : list[int] = [], 

23): 

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

25 if end == None: 

26 end = get_frames_count(input_structure_file, input_trajectory_file) 

27 

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

29 # For this reason both arguments are not allowed together 

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

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

32 

33 # End must be grater than start 

34 if end != None and end < start: 

35 raise InputError('End frame must be posterior to start frame') 

36 

37 # We need an output trajectory filename 

38 if not output_trajectory_file: 

39 raise InputError('Missing output trajectory filename') 

40 

41 # In case the frames argument is passed 

42 if frames and len(frames) > 0: 

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

44 if frames != sorted(frames): 

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

46 

47 # Get the input formats 

48 input_structure_format = input_structure_file.format 

49 input_trajectory_format = input_trajectory_file.format 

50 output_trajectory_format = output_trajectory_file.format 

51 format_set = { 

52 'inputs': { 

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

54 'input_trajectory_filename': { input_trajectory_format } 

55 }, 

56 'outputs': { 

57 'output_trajectory_filename': { output_trajectory_format } 

58 } 

59 } 

60 # Get a suitable function to do the job 

61 suitables = next(get_format_set_suitable_function( 

62 available_functions=subset_functions, 

63 available_request_format_sets=[format_set], 

64 ), None) 

65 if not suitables: 

66 raise InputError('There is no subset function which supports the requested formats') 

67 suitable_function, formats = suitables 

68 # Call the subset function 

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

70 suitable_function_keywords = getfullargspec(suitable_function)[0] 

71 required_structure = 'input_structure_filename' in suitable_function_keywords 

72 if required_structure: 

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

74 else: 

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