class TestFeaturizer(unittest.TestCase): @classmethod def setUpClass(cls): import tempfile cls.asn_leu_pdbfile = tempfile.mkstemp(suffix=".pdb")[1] with open(cls.asn_leu_pdbfile, 'w') as fh: fh.write(asn_leu_pdb) cls.asn_leu_traj = tempfile.mktemp(suffix='.xtc') cls.bogus_geom_pdbfile = tempfile.mkstemp(suffix=".pdb")[1] with open(cls.bogus_geom_pdbfile, 'w') as fh: fh.write(bogus_geom_pdbfile) # create traj for asn_leu n_frames = 4001 traj = mdtraj.load(cls.asn_leu_pdbfile) ref = traj.xyz new_xyz = np.empty((n_frames, ref.shape[1], 3)) noise = np.random.random(new_xyz.shape) new_xyz[:, :, :] = noise + ref traj.xyz = new_xyz traj.time = np.arange(n_frames) traj.save(cls.asn_leu_traj) @classmethod def tearDownClass(cls): try: os.unlink(cls.asn_leu_pdbfile) except EnvironmentError: pass try: os.unlink(cls.bogus_geom_pdbfile) except EnvironmentError: pass def setUp(self): self.pdbfile = pdbfile self.traj = mdtraj.load(xtcfile, top=self.pdbfile) self.feat = MDFeaturizer(self.pdbfile) self.atol = 1e-5 self.ref_frame = 0 self.atom_indices = np.arange(0, self.traj.n_atoms / 2) def test_select_backbone(self): inds = self.feat.select_Backbone() def test_select_non_symmetry_heavy_atoms(self): try: inds = self.feat.select_Heavy(exclude_symmetry_related=True) except RuntimeError as e: if "recursion depth" in e.args: import sys raise Exception( "recursion limit reached. Interpreter limit: {}".format( sys.getrecursionlimit())) def test_select_all(self): self.feat.add_all() assert (self.feat.dimension() == self.traj.n_atoms * 3) refmap = np.reshape(self.traj.xyz, (len(self.traj), self.traj.n_atoms * 3)) assert (np.all(refmap == self.feat.transform(self.traj))) def test_select(self): sel = np.array([1, 2, 5, 20], dtype=int) self.feat.add_selection(sel) assert (self.feat.dimension() == sel.shape[0] * 3) refmap = np.reshape(self.traj.xyz[:, sel, :], (len(self.traj), sel.shape[0] * 3)) assert (np.all(refmap == self.feat.transform(self.traj))) def test_distances(self): sel = np.array([1, 2, 5, 20], dtype=int) pairs_expected = np.array([[1, 5], [1, 20], [2, 5], [2, 20], [5, 20]]) pairs = self.feat.pairs(sel, excluded_neighbors=2) assert (pairs.shape == pairs_expected.shape) assert (np.all(pairs == pairs_expected)) self.feat.add_distances( pairs, periodic=False) # unperiodic distances such that we can compare assert (self.feat.dimension() == pairs_expected.shape[0]) X = self.traj.xyz[:, pairs_expected[:, 0], :] Y = self.traj.xyz[:, pairs_expected[:, 1], :] D = np.sqrt(np.sum((X - Y)**2, axis=2)) assert (np.allclose(D, self.feat.transform(self.traj))) def test_inverse_distances(self): sel = np.array([1, 2, 5, 20], dtype=int) pairs_expected = np.array([[1, 5], [1, 20], [2, 5], [2, 20], [5, 20]]) pairs = self.feat.pairs(sel, excluded_neighbors=2) assert (pairs.shape == pairs_expected.shape) assert (np.all(pairs == pairs_expected)) self.feat.add_inverse_distances( pairs, periodic=False) # unperiodic distances such that we can compare assert (self.feat.dimension() == pairs_expected.shape[0]) X = self.traj.xyz[:, pairs_expected[:, 0], :] Y = self.traj.xyz[:, pairs_expected[:, 1], :] Dinv = 1.0 / np.sqrt(np.sum((X - Y)**2, axis=2)) assert (np.allclose(Dinv, self.feat.transform(self.traj))) def test_ca_distances(self): sel = self.feat.select_Ca() assert (np.all(sel == list(range(self.traj.n_atoms))) ) # should be all for this Ca-traj pairs = self.feat.pairs(sel, excluded_neighbors=0) self.feat.add_distances_ca( periodic=False, excluded_neighbors=0 ) # unperiodic distances such that we can compare assert (self.feat.dimension() == pairs.shape[0]) X = self.traj.xyz[:, pairs[:, 0], :] Y = self.traj.xyz[:, pairs[:, 1], :] D = np.sqrt(np.sum((X - Y)**2, axis=2)) assert (np.allclose(D, self.feat.transform(self.traj))) def test_ca_distances_with_all_atom_geometries(self): feat = MDFeaturizer(pdbfile_ops_aa) feat.add_distances_ca(excluded_neighbors=0) D_aa = feat.transform(mdtraj.load(pdbfile_ops_aa)) # Create a reference feat_just_ca = MDFeaturizer(pdbfile_ops_Ca) feat_just_ca.add_distances(np.arange(feat_just_ca.topology.n_atoms)) D_ca = feat_just_ca.transform(mdtraj.load(pdbfile_ops_Ca)) assert (np.allclose(D_aa, D_ca)) def test_ca_distances_with_all_atom_geometries_and_exclusions(self): feat = MDFeaturizer(pdbfile_ops_aa) feat.add_distances_ca(excluded_neighbors=2) D_aa = feat.transform(mdtraj.load(pdbfile_ops_aa)) # Create a reference feat_just_ca = MDFeaturizer(pdbfile_ops_Ca) ca_pairs = feat.pairs(feat_just_ca.select_Ca(), excluded_neighbors=2) feat_just_ca.add_distances(ca_pairs) D_ca = feat_just_ca.transform(mdtraj.load(pdbfile_ops_Ca)) assert (np.allclose(D_aa, D_ca)) def test_ca_distances_with_residues_not_containing_cas_no_exclusions(self): # Load test geom geom = mdtraj.load(self.pdbfile) # No exclusions feat_EN0 = MDFeaturizer(self.bogus_geom_pdbfile) feat_EN0.add_distances_ca(excluded_neighbors=0) ENO_pairs = [[1, 3], [1, 5], [1, 7], [3, 5], [3, 7], [5, 7]] # Check indices assert (np.allclose(ENO_pairs, feat_EN0.active_features[0].distance_indexes)) # Check distances D = mdtraj.compute_distances(geom, ENO_pairs) assert (np.allclose(D, feat_EN0.transform(geom))) # excluded_neighbors=1 ## will yield the same as before, because the first neighbor # doesn't conting CA's anyway feat_EN1 = MDFeaturizer(self.bogus_geom_pdbfile) feat_EN1.add_distances_ca(excluded_neighbors=1) EN1_pairs = [[1, 3], [1, 5], [1, 7], [3, 5], [3, 7], [5, 7]] assert (np.allclose(EN1_pairs, feat_EN1.active_features[0].distance_indexes)) D = mdtraj.compute_distances(geom, EN1_pairs) assert (np.allclose(D, feat_EN1.transform(geom))) def test_ca_distances_with_residues_not_containing_cas_with_exclusions( self): # Load test geom geom = mdtraj.load(self.pdbfile) # No exclusions feat_EN2 = MDFeaturizer(self.bogus_geom_pdbfile) feat_EN2.add_distances_ca(excluded_neighbors=2) EN2_pairs = [ [1, 5], [1, 7], [3, 7], ] # Check indices assert (np.allclose(EN2_pairs, feat_EN2.active_features[0].distance_indexes)) # Check distances D = mdtraj.compute_distances(geom, EN2_pairs) assert (np.allclose(D, feat_EN2.transform(geom))) # excluded_neighbors=1 ## will yield the same as before, because the first neighbor # doesn't conting CA's anyway feat_EN1 = MDFeaturizer(self.bogus_geom_pdbfile) feat_EN1.add_distances_ca(excluded_neighbors=1) EN1_pairs = [[1, 3], [1, 5], [1, 7], [3, 5], [3, 7], [5, 7]] assert (np.allclose(EN1_pairs, feat_EN1.active_features[0].distance_indexes)) D = mdtraj.compute_distances(geom, EN1_pairs) assert (np.allclose(D, feat_EN1.transform(geom))) def test_contacts(self): sel = np.array([1, 2, 5, 20], dtype=int) pairs_expected = np.array([[1, 5], [1, 20], [2, 5], [2, 20], [5, 20]]) pairs = self.feat.pairs(sel, excluded_neighbors=2) assert (pairs.shape == pairs_expected.shape) assert (np.all(pairs == pairs_expected)) self.feat.add_contacts( pairs, threshold=0.5, periodic=False) # unperiodic distances such that we can compare assert (self.feat.dimension() == pairs_expected.shape[0]) X = self.traj.xyz[:, pairs_expected[:, 0], :] Y = self.traj.xyz[:, pairs_expected[:, 1], :] D = np.sqrt(np.sum((X - Y)**2, axis=2)) C = np.zeros(D.shape) I = np.argwhere(D <= 0.5) C[I[:, 0], I[:, 1]] = 1.0 assert (np.allclose(C, self.feat.transform(self.traj))) def test_contacts_count_contacts(self): sel = np.array([1, 2, 5, 20], dtype=int) pairs_expected = np.array([[1, 5], [1, 20], [2, 5], [2, 20], [5, 20]]) pairs = self.feat.pairs(sel, excluded_neighbors=2) assert (pairs.shape == pairs_expected.shape) assert (np.all(pairs == pairs_expected)) self.feat.add_contacts( pairs, threshold=0.5, periodic=False, count_contacts=True ) # unperiodic distances such that we can compare # The dimensionality of the feature is now one assert (self.feat.dimension() == 1) X = self.traj.xyz[:, pairs_expected[:, 0], :] Y = self.traj.xyz[:, pairs_expected[:, 1], :] D = np.sqrt(np.sum((X - Y)**2, axis=2)) C = np.zeros(D.shape) I = np.argwhere(D <= 0.5) C[I[:, 0], I[:, 1]] = 1.0 # Count the contacts C = C.sum(1, keepdims=True) assert (np.allclose(C, self.feat.transform(self.traj))) def test_angles(self): sel = np.array([[1, 2, 5], [1, 3, 8], [2, 9, 10]], dtype=int) self.feat.add_angles(sel) assert (self.feat.dimension() == sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) self.assertEqual(len(self.feat.describe()), self.feat.dimension()) def test_angles_deg(self): sel = np.array([[1, 2, 5], [1, 3, 8], [2, 9, 10]], dtype=int) self.feat.add_angles(sel, deg=True) assert (self.feat.dimension() == sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -180.0)) assert (np.alltrue(Y <= 180.0)) def test_angles_cossin(self): sel = np.array([[1, 2, 5], [1, 3, 8], [2, 9, 10]], dtype=int) self.feat.add_angles(sel, cossin=True) assert (self.feat.dimension() == 2 * sel.shape[0]) Y = self.feat.transform(self.traj) self.assertEqual(Y.shape, (self.traj.n_frames, 2 * sel.shape[0])) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension()) def test_dihedrals(self): sel = np.array([[1, 2, 5, 6], [1, 3, 8, 9], [2, 9, 10, 12]], dtype=int) self.feat.add_dihedrals(sel) assert (self.feat.dimension() == sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) self.assertEqual(len(self.feat.describe()), self.feat.dimension()) def test_dihedrals_deg(self): sel = np.array([[1, 2, 5, 6], [1, 3, 8, 9], [2, 9, 10, 12]], dtype=int) self.feat.add_dihedrals(sel, deg=True) assert (self.feat.dimension() == sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -180.0)) assert (np.alltrue(Y <= 180.0)) self.assertEqual(len(self.feat.describe()), self.feat.dimension()) def test_dihedrials_cossin(self): sel = np.array([[1, 2, 5, 6], [1, 3, 8, 9], [2, 9, 10, 12]], dtype=int) self.feat.add_dihedrals(sel, cossin=True) assert (self.feat.dimension() == 2 * sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension()) def test_backbone_dihedrals(self): self.feat = MDFeaturizer(topfile=self.asn_leu_pdbfile) self.feat.add_backbone_torsions() traj = mdtraj.load(self.asn_leu_pdbfile) Y = self.feat.transform(traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension()) # test ordering of indices backbone_feature = self.feat.active_features[0] angle_indices = backbone_feature.angle_indexes np.testing.assert_equal(angle_indices[0], backbone_feature._phi_inds[0]) np.testing.assert_equal(angle_indices[1], backbone_feature._psi_inds[0]) np.testing.assert_equal(angle_indices[2], backbone_feature._phi_inds[1]) np.testing.assert_equal(angle_indices[3], backbone_feature._psi_inds[1]) def test_backbone_dihedrals_deg(self): self.feat = MDFeaturizer(topfile=self.asn_leu_pdbfile) self.feat.add_backbone_torsions(deg=True) traj = mdtraj.load(self.asn_leu_pdbfile) Y = self.feat.transform(traj) assert (np.alltrue(Y >= -180.0)) assert (np.alltrue(Y <= 180.0)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension()) def test_backbone_dihedrals_cossin(self): self.feat = MDFeaturizer(topfile=self.asn_leu_pdbfile) self.feat.add_backbone_torsions(cossin=True) traj = mdtraj.load(self.asn_leu_traj, top=self.asn_leu_pdbfile) Y = self.feat.transform(traj) self.assertEqual(Y.shape, (len(traj), 3 * 4)) # (3 phi + 3 psi)*2 [cos, sin] assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension(), msg=desc) self.assertIn("COS", desc[0]) self.assertIn("SIN", desc[1]) def test_backbone_dihedrials_chi(self): self.feat = MDFeaturizer(topfile=self.asn_leu_pdbfile) self.feat.add_chi1_torsions() traj = mdtraj.load(self.asn_leu_pdbfile) Y = self.feat.transform(traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension()) def test_backbone_dihedrials_chi_cossin(self): self.feat = MDFeaturizer(topfile=self.asn_leu_pdbfile) self.feat.add_chi1_torsions(cossin=True) traj = mdtraj.load(self.asn_leu_pdbfile) Y = self.feat.transform(traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() assert "COS" in desc[0] assert "SIN" in desc[1] self.assertEqual(len(desc), self.feat.dimension()) def test_custom_feature(self): # TODO: test me pass def test_MinRmsd(self): # Test the Trajectory-input variant self.feat.add_minrmsd_to_ref(self.traj[self.ref_frame]) # and the file-input variant self.feat.add_minrmsd_to_ref(xtcfile, ref_frame=self.ref_frame) test_Y = self.feat.transform(self.traj).squeeze() # now the reference ref_Y = mdtraj.rmsd(self.traj, self.traj[self.ref_frame]) verbose_assertion_minrmsd(ref_Y, test_Y, self) assert self.feat.dimension() == 2 assert len(self.feat.describe()) == 2 def test_MinRmsd_with_atom_indices(self): # Test the Trajectory-input variant self.feat.add_minrmsd_to_ref(self.traj[self.ref_frame], atom_indices=self.atom_indices) # and the file-input variant self.feat.add_minrmsd_to_ref(xtcfile, ref_frame=self.ref_frame, atom_indices=self.atom_indices) test_Y = self.feat.transform(self.traj).squeeze() # now the reference ref_Y = mdtraj.rmsd(self.traj, self.traj[self.ref_frame], atom_indices=self.atom_indices) verbose_assertion_minrmsd(ref_Y, test_Y, self) assert self.feat.dimension() == 2 assert len(self.feat.describe()) == 2 def test_MinRmsd_with_atom_indices_precentered(self): # Test the Trajectory-input variant self.feat.add_minrmsd_to_ref(self.traj[self.ref_frame], atom_indices=self.atom_indices, precentered=True) # and the file-input variant self.feat.add_minrmsd_to_ref(xtcfile, ref_frame=self.ref_frame, atom_indices=self.atom_indices, precentered=True) test_Y = self.feat.transform(self.traj).squeeze() # now the reference ref_Y = mdtraj.rmsd(self.traj, self.traj[self.ref_frame], atom_indices=self.atom_indices, precentered=True) verbose_assertion_minrmsd(ref_Y, test_Y, self) assert self.feat.dimension() == 2 assert len(self.feat.describe()) == 2 def test_Residue_Mindist_Ca_all(self): n_ca = self.feat.topology.n_atoms self.feat.add_residue_mindist(scheme='ca') D = self.feat.transform(self.traj) Dref = mdtraj.compute_contacts(self.traj, scheme='ca')[0] assert np.allclose(D, Dref) assert len(self.feat.describe()) == self.feat.dimension() def test_Residue_Mindist_Ca_all_threshold(self): threshold = .7 self.feat.add_residue_mindist(scheme='ca', threshold=threshold) D = self.feat.transform(self.traj) Dref = mdtraj.compute_contacts(self.traj, scheme='ca')[0] Dbinary = np.zeros_like(Dref) I = np.argwhere(Dref <= threshold) Dbinary[I[:, 0], I[:, 1]] = 1 assert np.allclose(D, Dbinary) assert len(self.feat.describe()) == self.feat.dimension() def test_Residue_Mindist_Ca_array(self): contacts = np.array([[ 20, 10, ], [10, 0]]) self.feat.add_residue_mindist(scheme='ca', residue_pairs=contacts) D = self.feat.transform(self.traj) Dref = mdtraj.compute_contacts(self.traj, scheme='ca', contacts=contacts)[0] assert np.allclose(D, Dref) assert len(self.feat.describe()) == self.feat.dimension() def test_Residue_Mindist_Ca_array_periodic(self): traj = mdtraj.load(pdbfile) # Atoms most far appart in Z atom_minz = traj.xyz.argmin(1).squeeze()[-1] atom_maxz = traj.xyz.argmax(1).squeeze()[-1] # Residues with the atoms most far appart in Z res_minz = traj.topology.atom(atom_minz).residue.index res_maxz = traj.topology.atom(atom_maxz).residue.index contacts = np.array([[res_minz, res_maxz]]) # Tweak the trajectory so that a (bogus) PBC exists (otherwise traj._have_unitcell is False) traj.unitcell_angles = [90, 90, 90] traj.unitcell_lengths = [1, 1, 1] self.feat.add_residue_mindist(scheme='ca', residue_pairs=contacts, periodic=False) D = self.feat.transform(traj) Dperiodic_true = mdtraj.compute_contacts(traj, scheme='ca', contacts=contacts, periodic=True)[0] Dperiodic_false = mdtraj.compute_contacts(traj, scheme='ca', contacts=contacts, periodic=False)[0] # This asserts that the periodic option is having an effect at all assert not np.allclose( Dperiodic_false, Dperiodic_true, ) # This asserts that the periodic option is being handled correctly by pyemma assert np.allclose(D, Dperiodic_false) assert len(self.feat.describe()) == self.feat.dimension() def test_Group_Mindist_One_Group(self): group0 = [0, 20, 30, 0] self.feat.add_group_mindist( group_definitions=[group0]) # Even with duplicates D = self.feat.transform(self.traj) dist_list = list(combinations(np.unique(group0), 2)) Dref = mdtraj.compute_distances(self.traj, dist_list) assert np.allclose(D.squeeze(), Dref.min(1)) assert len(self.feat.describe()) == self.feat.dimension() def test_Group_Mindist_All_Three_Groups(self): group0 = [0, 20, 30, 0] group1 = [1, 21, 31, 1] group2 = [2, 22, 32, 2] self.feat.add_group_mindist(group_definitions=[group0, group1, group2]) D = self.feat.transform(self.traj) # Now the references, computed separately for each combination of groups dist_list_01 = np.array( list(product(np.unique(group0), np.unique(group1)))) dist_list_02 = np.array( list(product(np.unique(group0), np.unique(group2)))) dist_list_12 = np.array( list(product(np.unique(group1), np.unique(group2)))) Dref_01 = mdtraj.compute_distances(self.traj, dist_list_01).min(1) Dref_02 = mdtraj.compute_distances(self.traj, dist_list_02).min(1) Dref_12 = mdtraj.compute_distances(self.traj, dist_list_12).min(1) Dref = np.vstack((Dref_01, Dref_02, Dref_12)).T assert np.allclose(D.squeeze(), Dref) assert len(self.feat.describe()) == self.feat.dimension() def test_Group_Mindist_All_Three_Groups_threshold(self): threshold = .7 group0 = [0, 20, 30, 0] group1 = [1, 21, 31, 1] group2 = [2, 22, 32, 2] self.feat.add_group_mindist(group_definitions=[group0, group1, group2], threshold=threshold) D = self.feat.transform(self.traj) # Now the references, computed separately for each combination of groups dist_list_01 = np.array( list(product(np.unique(group0), np.unique(group1)))) dist_list_02 = np.array( list(product(np.unique(group0), np.unique(group2)))) dist_list_12 = np.array( list(product(np.unique(group1), np.unique(group2)))) Dref_01 = mdtraj.compute_distances(self.traj, dist_list_01).min(1) Dref_02 = mdtraj.compute_distances(self.traj, dist_list_02).min(1) Dref_12 = mdtraj.compute_distances(self.traj, dist_list_12).min(1) Dref = np.vstack((Dref_01, Dref_02, Dref_12)).T Dbinary = np.zeros_like(Dref) I = np.argwhere(Dref <= threshold) Dbinary[I[:, 0], I[:, 1]] = 1 assert np.allclose(D, Dbinary) assert len(self.feat.describe()) == self.feat.dimension() def test_Group_Mindist_Some_Three_Groups(self): group0 = [0, 20, 30, 0] group1 = [1, 21, 31, 1] group2 = [2, 22, 32, 2] group_pairs = np.array([[0, 1], [2, 2], [0, 2]]) self.feat.add_group_mindist(group_definitions=[group0, group1, group2], group_pairs=group_pairs) D = self.feat.transform(self.traj) # Now the references, computed separately for each combination of groups dist_list_01 = np.array( list(product(np.unique(group0), np.unique(group1)))) dist_list_02 = np.array( list(product(np.unique(group0), np.unique(group2)))) dist_list_22 = np.array(list(combinations(np.unique(group2), 2))) Dref_01 = mdtraj.compute_distances(self.traj, dist_list_01).min(1) Dref_02 = mdtraj.compute_distances(self.traj, dist_list_02).min(1) Dref_22 = mdtraj.compute_distances(self.traj, dist_list_22).min(1) Dref = np.vstack((Dref_01, Dref_22, Dref_02)).T assert np.allclose(D.squeeze(), Dref) assert len(self.feat.describe()) == self.feat.dimension()
class TestFeaturizer(unittest.TestCase): @classmethod def setUpClass(cls): import tempfile cls.asn_leu_pdbfile = tempfile.mkstemp(suffix=".pdb")[1] with open(cls.asn_leu_pdbfile, 'w') as fh: fh.write(asn_leu_pdb) cls.asn_leu_traj = tempfile.mktemp(suffix='.xtc') # create traj for asn_leu n_frames = 4001 traj = mdtraj.load(cls.asn_leu_pdbfile) ref = traj.xyz new_xyz = np.empty((n_frames, ref.shape[1], 3)) noise = np.random.random(new_xyz.shape) new_xyz[:, :, :] = noise + ref traj.xyz = new_xyz traj.time = np.arange(n_frames) traj.save(cls.asn_leu_traj) super(TestFeaturizer, cls).setUpClass() @classmethod def tearDownClass(cls): try: os.unlink(cls.asn_leu_pdbfile) except EnvironmentError: pass super(TestFeaturizer, cls).tearDownClass() def setUp(self): self.pdbfile = pdbfile self.traj = mdtraj.load(xtcfile, top=self.pdbfile) self.feat = MDFeaturizer(self.pdbfile) self.atol = 1e-5 self.ref_frame = 0 self.atom_indices = np.arange(0, self.traj.n_atoms / 2) def test_select_backbone(self): inds = self.feat.select_Backbone() def test_select_all(self): self.feat.add_all() assert (self.feat.dimension() == self.traj.n_atoms * 3) refmap = np.reshape(self.traj.xyz, (len(self.traj), self.traj.n_atoms * 3)) assert (np.all(refmap == self.feat.transform(self.traj))) def test_select(self): sel = np.array([1, 2, 5, 20], dtype=int) self.feat.add_selection(sel) assert (self.feat.dimension() == sel.shape[0] * 3) refmap = np.reshape(self.traj.xyz[:, sel, :], (len(self.traj), sel.shape[0] * 3)) assert (np.all(refmap == self.feat.transform(self.traj))) def test_distances(self): sel = np.array([1, 2, 5, 20], dtype=int) pairs_expected = np.array([[1, 5], [1, 20], [2, 5], [2, 20], [5, 20]]) pairs = self.feat.pairs(sel, excluded_neighbors=2) assert (pairs.shape == pairs_expected.shape) assert (np.all(pairs == pairs_expected)) self.feat.add_distances( pairs, periodic=False) # unperiodic distances such that we can compare assert (self.feat.dimension() == pairs_expected.shape[0]) X = self.traj.xyz[:, pairs_expected[:, 0], :] Y = self.traj.xyz[:, pairs_expected[:, 1], :] D = np.sqrt(np.sum((X - Y)**2, axis=2)) assert (np.allclose(D, self.feat.transform(self.traj))) def test_inverse_distances(self): sel = np.array([1, 2, 5, 20], dtype=int) pairs_expected = np.array([[1, 5], [1, 20], [2, 5], [2, 20], [5, 20]]) pairs = self.feat.pairs(sel, excluded_neighbors=2) assert (pairs.shape == pairs_expected.shape) assert (np.all(pairs == pairs_expected)) self.feat.add_inverse_distances( pairs, periodic=False) # unperiodic distances such that we can compare assert (self.feat.dimension() == pairs_expected.shape[0]) X = self.traj.xyz[:, pairs_expected[:, 0], :] Y = self.traj.xyz[:, pairs_expected[:, 1], :] Dinv = 1.0 / np.sqrt(np.sum((X - Y)**2, axis=2)) assert (np.allclose(Dinv, self.feat.transform(self.traj))) def test_ca_distances(self): sel = self.feat.select_Ca() assert (np.all(sel == list(range(self.traj.n_atoms))) ) # should be all for this Ca-traj pairs = self.feat.pairs(sel, excluded_neighbors=0) self.feat.add_distances_ca( periodic=False, excluded_neighbors=0 ) # unperiodic distances such that we can compare assert (self.feat.dimension() == pairs.shape[0]) X = self.traj.xyz[:, pairs[:, 0], :] Y = self.traj.xyz[:, pairs[:, 1], :] D = np.sqrt(np.sum((X - Y)**2, axis=2)) assert (np.allclose(D, self.feat.transform(self.traj))) def test_ca_distances_with_all_atom_geometries(self): feat = MDFeaturizer(pdbfile_ops_aa) feat.add_distances_ca(excluded_neighbors=0) D_aa = feat.transform(mdtraj.load(pdbfile_ops_aa)) # Create a reference feat_just_ca = MDFeaturizer(pdbfile_ops_Ca) feat_just_ca.add_distances(np.arange(feat_just_ca.topology.n_atoms)) D_ca = feat_just_ca.transform(mdtraj.load(pdbfile_ops_Ca)) assert (np.allclose(D_aa, D_ca)) def test_ca_distances_with_all_atom_geometries_and_exclusions(self): feat = MDFeaturizer(pdbfile_ops_aa) feat.add_distances_ca(excluded_neighbors=2) D_aa = feat.transform(mdtraj.load(pdbfile_ops_aa)) # Create a reference feat_just_ca = MDFeaturizer(pdbfile_ops_Ca) ca_pairs = feat.pairs(feat_just_ca.select_Ca(), excluded_neighbors=2) feat_just_ca.add_distances(ca_pairs) D_ca = feat_just_ca.transform(mdtraj.load(pdbfile_ops_Ca)) assert (np.allclose(D_aa, D_ca)) def test_contacts(self): sel = np.array([1, 2, 5, 20], dtype=int) pairs_expected = np.array([[1, 5], [1, 20], [2, 5], [2, 20], [5, 20]]) pairs = self.feat.pairs(sel, excluded_neighbors=2) assert (pairs.shape == pairs_expected.shape) assert (np.all(pairs == pairs_expected)) self.feat.add_contacts( pairs, threshold=0.5, periodic=False) # unperiodic distances such that we can compare assert (self.feat.dimension() == pairs_expected.shape[0]) X = self.traj.xyz[:, pairs_expected[:, 0], :] Y = self.traj.xyz[:, pairs_expected[:, 1], :] D = np.sqrt(np.sum((X - Y)**2, axis=2)) C = np.zeros(D.shape) I = np.argwhere(D <= 0.5) C[I[:, 0], I[:, 1]] = 1.0 assert (np.allclose(C, self.feat.transform(self.traj))) def test_contacts_count_contacts(self): sel = np.array([1, 2, 5, 20], dtype=int) pairs_expected = np.array([[1, 5], [1, 20], [2, 5], [2, 20], [5, 20]]) pairs = self.feat.pairs(sel, excluded_neighbors=2) assert (pairs.shape == pairs_expected.shape) assert (np.all(pairs == pairs_expected)) self.feat.add_contacts( pairs, threshold=0.5, periodic=False, count_contacts=True ) # unperiodic distances such that we can compare # The dimensionality of the feature is now one assert (self.feat.dimension() == 1) X = self.traj.xyz[:, pairs_expected[:, 0], :] Y = self.traj.xyz[:, pairs_expected[:, 1], :] D = np.sqrt(np.sum((X - Y)**2, axis=2)) C = np.zeros(D.shape) I = np.argwhere(D <= 0.5) C[I[:, 0], I[:, 1]] = 1.0 # Count the contacts C = C.sum(1, keepdims=True) assert (np.allclose(C, self.feat.transform(self.traj))) def test_angles(self): sel = np.array([[1, 2, 5], [1, 3, 8], [2, 9, 10]], dtype=int) self.feat.add_angles(sel) assert (self.feat.dimension() == sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) self.assertEqual(len(self.feat.describe()), self.feat.dimension()) def test_angles_deg(self): sel = np.array([[1, 2, 5], [1, 3, 8], [2, 9, 10]], dtype=int) self.feat.add_angles(sel, deg=True) assert (self.feat.dimension() == sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -180.0)) assert (np.alltrue(Y <= 180.0)) def test_angles_cossin(self): sel = np.array([[1, 2, 5], [1, 3, 8], [2, 9, 10]], dtype=int) self.feat.add_angles(sel, cossin=True) assert (self.feat.dimension() == 2 * sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension()) def test_dihedrals(self): sel = np.array([[1, 2, 5, 6], [1, 3, 8, 9], [2, 9, 10, 12]], dtype=int) self.feat.add_dihedrals(sel) assert (self.feat.dimension() == sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) self.assertEqual(len(self.feat.describe()), self.feat.dimension()) def test_dihedrals_deg(self): sel = np.array([[1, 2, 5, 6], [1, 3, 8, 9], [2, 9, 10, 12]], dtype=int) self.feat.add_dihedrals(sel, deg=True) assert (self.feat.dimension() == sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -180.0)) assert (np.alltrue(Y <= 180.0)) self.assertEqual(len(self.feat.describe()), self.feat.dimension()) def test_dihedrials_cossin(self): sel = np.array([[1, 2, 5, 6], [1, 3, 8, 9], [2, 9, 10, 12]], dtype=int) self.feat.add_dihedrals(sel, cossin=True) assert (self.feat.dimension() == 2 * sel.shape[0]) Y = self.feat.transform(self.traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension()) def test_backbone_dihedrals(self): self.feat = MDFeaturizer(topfile=self.asn_leu_pdbfile) self.feat.add_backbone_torsions() traj = mdtraj.load(self.asn_leu_pdbfile) Y = self.feat.transform(traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension()) def test_backbone_dihedrals_deg(self): self.feat = MDFeaturizer(topfile=self.asn_leu_pdbfile) self.feat.add_backbone_torsions(deg=True) traj = mdtraj.load(self.asn_leu_pdbfile) Y = self.feat.transform(traj) assert (np.alltrue(Y >= -180.0)) assert (np.alltrue(Y <= 180.0)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension()) def test_backbone_dihedrals_cossin(self): self.feat = MDFeaturizer(topfile=self.asn_leu_pdbfile) self.feat.add_backbone_torsions(cossin=True) traj = mdtraj.load(self.asn_leu_traj, top=self.asn_leu_pdbfile) Y = self.feat.transform(traj) self.assertEqual(Y.shape, (len(traj), 3 * 4)) # (3 phi + 3 psi)*2 [cos, sin] assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() assert "COS" in desc[0] assert "SIN" in desc[1] self.assertEqual(len(desc), self.feat.dimension()) def test_backbone_dihedrials_chi(self): self.feat = MDFeaturizer(topfile=self.asn_leu_pdbfile) self.feat.add_chi1_torsions() traj = mdtraj.load(self.asn_leu_pdbfile) Y = self.feat.transform(traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() self.assertEqual(len(desc), self.feat.dimension()) def test_backbone_dihedrials_chi_cossin(self): self.feat = MDFeaturizer(topfile=self.asn_leu_pdbfile) self.feat.add_chi1_torsions(cossin=True) traj = mdtraj.load(self.asn_leu_pdbfile) Y = self.feat.transform(traj) assert (np.alltrue(Y >= -np.pi)) assert (np.alltrue(Y <= np.pi)) desc = self.feat.describe() assert "COS" in desc[0] assert "SIN" in desc[1] self.assertEqual(len(desc), self.feat.dimension()) def test_custom_feature(self): # TODO: test me pass def test_MinRmsd(self): # Test the Trajectory-input variant self.feat.add_minrmsd_to_ref(self.traj[self.ref_frame]) # and the file-input variant self.feat.add_minrmsd_to_ref(xtcfile, ref_frame=self.ref_frame) test_Y = self.feat.transform(self.traj).squeeze() # now the reference ref_Y = mdtraj.rmsd(self.traj, self.traj[self.ref_frame]) verbose_assertion_minrmsd(ref_Y, test_Y, self) assert self.feat.dimension() == 2 assert len(self.feat.describe()) == 2 def test_MinRmsd_with_atom_indices(self): # Test the Trajectory-input variant self.feat.add_minrmsd_to_ref(self.traj[self.ref_frame], atom_indices=self.atom_indices) # and the file-input variant self.feat.add_minrmsd_to_ref(xtcfile, ref_frame=self.ref_frame, atom_indices=self.atom_indices) test_Y = self.feat.transform(self.traj).squeeze() # now the reference ref_Y = mdtraj.rmsd(self.traj, self.traj[self.ref_frame], atom_indices=self.atom_indices) verbose_assertion_minrmsd(ref_Y, test_Y, self) assert self.feat.dimension() == 2 assert len(self.feat.describe()) == 2 def test_MinRmsd_with_atom_indices_precentered(self): # Test the Trajectory-input variant self.feat.add_minrmsd_to_ref(self.traj[self.ref_frame], atom_indices=self.atom_indices, precentered=True) # and the file-input variant self.feat.add_minrmsd_to_ref(xtcfile, ref_frame=self.ref_frame, atom_indices=self.atom_indices, precentered=True) test_Y = self.feat.transform(self.traj).squeeze() # now the reference ref_Y = mdtraj.rmsd(self.traj, self.traj[self.ref_frame], atom_indices=self.atom_indices, precentered=True) verbose_assertion_minrmsd(ref_Y, test_Y, self) assert self.feat.dimension() == 2 assert len(self.feat.describe()) == 2 def test_Residue_Mindist_Ca_all(self): n_ca = self.feat.topology.n_atoms self.feat.add_residue_mindist(scheme='ca') D = self.feat.transform(self.traj) Dref = mdtraj.compute_contacts(self.traj, scheme='ca')[0] assert np.allclose(D, Dref) assert len(self.feat.describe()) == self.feat.dimension() def test_Residue_Mindist_Ca_all_threshold(self): threshold = .7 self.feat.add_residue_mindist(scheme='ca', threshold=threshold) D = self.feat.transform(self.traj) Dref = mdtraj.compute_contacts(self.traj, scheme='ca')[0] Dbinary = np.zeros_like(Dref) I = np.argwhere(Dref <= threshold) Dbinary[I[:, 0], I[:, 1]] = 1 assert np.allclose(D, Dbinary) assert len(self.feat.describe()) == self.feat.dimension() def test_Residue_Mindist_Ca_array(self): contacts = np.array([[ 20, 10, ], [10, 0]]) self.feat.add_residue_mindist(scheme='ca', residue_pairs=contacts) D = self.feat.transform(self.traj) Dref = mdtraj.compute_contacts(self.traj, scheme='ca', contacts=contacts)[0] assert np.allclose(D, Dref) assert len(self.feat.describe()) == self.feat.dimension() def test_Group_Mindist_One_Group(self): group0 = [0, 20, 30, 0] self.feat.add_group_mindist( group_definitions=[group0]) # Even with duplicates D = self.feat.transform(self.traj) dist_list = list(combinations(np.unique(group0), 2)) Dref = mdtraj.compute_distances(self.traj, dist_list) assert np.allclose(D.squeeze(), Dref.min(1)) assert len(self.feat.describe()) == self.feat.dimension() def test_Group_Mindist_All_Three_Groups(self): group0 = [0, 20, 30, 0] group1 = [1, 21, 31, 1] group2 = [2, 22, 32, 2] self.feat.add_group_mindist(group_definitions=[group0, group1, group2]) D = self.feat.transform(self.traj) # Now the references, computed separately for each combination of groups dist_list_01 = np.array( list(product(np.unique(group0), np.unique(group1)))) dist_list_02 = np.array( list(product(np.unique(group0), np.unique(group2)))) dist_list_12 = np.array( list(product(np.unique(group1), np.unique(group2)))) Dref_01 = mdtraj.compute_distances(self.traj, dist_list_01).min(1) Dref_02 = mdtraj.compute_distances(self.traj, dist_list_02).min(1) Dref_12 = mdtraj.compute_distances(self.traj, dist_list_12).min(1) Dref = np.vstack((Dref_01, Dref_02, Dref_12)).T assert np.allclose(D.squeeze(), Dref) assert len(self.feat.describe()) == self.feat.dimension() def test_Group_Mindist_All_Three_Groups_threshold(self): threshold = .7 group0 = [0, 20, 30, 0] group1 = [1, 21, 31, 1] group2 = [2, 22, 32, 2] self.feat.add_group_mindist(group_definitions=[group0, group1, group2], threshold=threshold) D = self.feat.transform(self.traj) # Now the references, computed separately for each combination of groups dist_list_01 = np.array( list(product(np.unique(group0), np.unique(group1)))) dist_list_02 = np.array( list(product(np.unique(group0), np.unique(group2)))) dist_list_12 = np.array( list(product(np.unique(group1), np.unique(group2)))) Dref_01 = mdtraj.compute_distances(self.traj, dist_list_01).min(1) Dref_02 = mdtraj.compute_distances(self.traj, dist_list_02).min(1) Dref_12 = mdtraj.compute_distances(self.traj, dist_list_12).min(1) Dref = np.vstack((Dref_01, Dref_02, Dref_12)).T Dbinary = np.zeros_like(Dref) I = np.argwhere(Dref <= threshold) Dbinary[I[:, 0], I[:, 1]] = 1 assert np.allclose(D, Dbinary) assert len(self.feat.describe()) == self.feat.dimension() def test_Group_Mindist_Some_Three_Groups(self): group0 = [0, 20, 30, 0] group1 = [1, 21, 31, 1] group2 = [2, 22, 32, 2] group_pairs = np.array([[0, 1], [2, 2], [0, 2]]) self.feat.add_group_mindist(group_definitions=[group0, group1, group2], group_pairs=group_pairs) D = self.feat.transform(self.traj) # Now the references, computed separately for each combination of groups dist_list_01 = np.array( list(product(np.unique(group0), np.unique(group1)))) dist_list_02 = np.array( list(product(np.unique(group0), np.unique(group2)))) dist_list_22 = np.array(list(combinations(np.unique(group2), 2))) Dref_01 = mdtraj.compute_distances(self.traj, dist_list_01).min(1) Dref_02 = mdtraj.compute_distances(self.traj, dist_list_02).min(1) Dref_22 = mdtraj.compute_distances(self.traj, dist_list_22).min(1) Dref = np.vstack((Dref_01, Dref_22, Dref_02)).T assert np.allclose(D.squeeze(), Dref) assert len(self.feat.describe()) == self.feat.dimension()