From 757e10cc4f5808e1c676bc96212b7ee1e0824537 Mon Sep 17 00:00:00 2001 From: martinvoegele Date: Sat, 25 Jan 2025 17:06:11 -0500 Subject: [PATCH] make extract_coordinates work with a single string as trajectory input --- pensa/preprocessing/coordinates.py | 7 +++++-- tests/test_workflow.py | 5 ++++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/pensa/preprocessing/coordinates.py b/pensa/preprocessing/coordinates.py index 707a1c8..cf537ee 100644 --- a/pensa/preprocessing/coordinates.py +++ b/pensa/preprocessing/coordinates.py @@ -18,8 +18,8 @@ def extract_coordinates(top, pdb, trj_list, out_name, sel_string, start_frame=0, Can read all MDAnalysis-compatible topology formats. pdb : str File name for the reference PDB file. - trj_list : list of str - File names for the input trajectory. + trj_list : str or list of str + File name(s) for the input trajectory. Can read all MDAnalysis-compatible trajectory formats. out_name : str Core of the file names for the output files. @@ -29,6 +29,9 @@ def extract_coordinates(top, pdb, trj_list, out_name, sel_string, start_frame=0, First frame to read from the trajectory. """ + # Convert the trajectory input to a list if it is a single string + if type(trj_list) is str: + trj_list = [trj_list] # Read the topology+PDB files and extract selected parts. u = mda.Universe(top, pdb) u.residues.resids -= residues_offset diff --git a/tests/test_workflow.py b/tests/test_workflow.py index 1d43c23..9925336 100644 --- a/tests/test_workflow.py +++ b/tests/test_workflow.py @@ -100,14 +100,17 @@ def setUp(self): ) # Extract the coordinates of the TM region from the trajectory + # - Test using a single trajectory directly self.file_tm_single_a = extract_coordinates( - self.ref_file_a, self.pdb_file_a, [self.trj_file_a], + self.ref_file_a, self.pdb_file_a, self.trj_file_a, self.trj_name_a + "_tm", sel_string_a ) + # - Test using a single trajectory as a list self.file_tm_single_b = extract_coordinates( self.ref_file_b, self.pdb_file_b, [self.trj_file_b], self.trj_name_b + "_tm", sel_string_b ) + # - Test using multiple trajectories self.file_tm_combined = extract_coordinates_combined( [self.ref_file_a] * 3 + [self.ref_file_b] * 3, self.trj_file_a + self.trj_file_b,