diff --git a/foldingdiff/angles_and_coords.py b/foldingdiff/angles_and_coords.py index 7a49c9c..62e68c9 100644 --- a/foldingdiff/angles_and_coords.py +++ b/foldingdiff/angles_and_coords.py @@ -60,25 +60,19 @@ def canonical_distances_and_dihedrals( for a in non_dihedral_angles: if a == "tau" or a == "N:CA:C": # tau = N - CA - C internal angles - idx = np.array( - [list(range(i, i + 3)) for i in range(3, len(backbone_atoms), 3)] - + [(0, 0, 0)] - ) + r = np.arange(3, len(backbone_atoms), 3) + idx = np.hstack([np.vstack([r, r + 1, r + 2]), np.zeros((3, 1))]).T elif a == "CA:C:1N": # Same as C-N angle in nerf # This measures an angle between two residues. Due to the way we build # proteins out later, we do not need to meas - idx = np.array( - [(i + 1, i + 2, i + 3) for i in range(0, len(backbone_atoms) - 3, 3)] - + [(0, 0, 0)] - ) + r = np.arange(0, len(backbone_atoms) - 3, 3) + idx = np.hstack([np.vstack([r + 1, r + 2, r + 3]), np.zeros((3, 1))]).T elif a == "C:1N:1CA": - idx = np.array( - [(i + 2, i + 3, i + 4) for i in range(0, len(backbone_atoms) - 3, 3)] - + [(0, 0, 0)] - ) + r = np.arange(0, len(backbone_atoms) - 3, 3) + idx = np.hstack([np.vstack([r + 2, r + 3, r + 4]), np.zeros((3, 1))]).T else: raise ValueError(f"Unrecognized angle: {a}") - calc_angles[a] = struc.index_angle(backbone_atoms, indices=idx) + calc_angles[a] = struc.index_angle(backbone_atoms, indices=idx.astype(int)) # At this point we've only looked at dihedral and angles; check value range for k, v in calc_angles.items(): @@ -92,29 +86,25 @@ def canonical_distances_and_dihedrals( # Since this is measuring the distance between pairs of residues, there # is one fewer such measurement than the total number of residues like # for dihedrals. Therefore, we pad this with a null 0 value at the end. - idx = np.array( - [(i + 2, i + 3) for i in range(0, len(backbone_atoms) - 3, 3)] - + [(0, 0)] - ) + r = np.arange(0, len(backbone_atoms) - 3, 3) + idx = np.hstack([np.vstack([r + 2, r + 3]), np.zeros((2, 1))]).T elif d == "N:CA": # We start resconstructing with a fixed initial residue so we do not need # to predict or record the initial distance. Additionally we pad with a # null value at the end - idx = np.array( - [(i, i + 1) for i in range(3, len(backbone_atoms), 3)] + [(0, 0)] - ) + r = np.arange(3, len(backbone_atoms), 3) + idx = np.hstack([np.vstack([r, r + 1]), np.zeros((2, 1))]).T assert len(idx) == len(calc_angles["phi"]) elif d == "CA:C": # We start reconstructing with a fixed initial residue so we do not need # to predict or record the initial distance. Additionally, we pad with a # null value at the end. - idx = np.array( - [(i + 1, i + 2) for i in range(3, len(backbone_atoms), 3)] + [(0, 0)] - ) + r = np.arange(3, len(backbone_atoms), 3) + idx = np.hstack([np.vstack([r + 1, r + 2]), np.zeros((2, 1))]).T assert len(idx) == len(calc_angles["phi"]) else: raise ValueError(f"Unrecognized distance: {d}") - calc_angles[d] = struc.index_distance(backbone_atoms, indices=idx) + calc_angles[d] = struc.index_distance(backbone_atoms, indices=idx.astype(int)) return pd.DataFrame({k: calc_angles[k].squeeze() for k in distances + angles}) @@ -368,18 +358,20 @@ def build_aa_sidechain_dict( that specify how to build out that sidechain's atoms from the backbone """ if not reference_pdbs: - glob.glob( + reference_pdbs = glob.glob( os.path.join(os.path.dirname(os.path.dirname(__file__)), "data/*.pdb") ) + ref_file_counter = 0 retval = {} for pdb in reference_pdbs: try: sidechain_angles = collect_aa_sidechain_angles(pdb) retval.update(sidechain_angles) # Overwrites any existing key/value pairs + ref_file_counter += 1 except ValueError: continue - logging.info(f"Built sidechain dictionary with {len(retval)} amino acids") + logging.info(f"Built sidechain dictionary with {len(retval)} amino acids from {ref_file_counter} files") return retval diff --git a/tests/test_data.py b/tests/test_data.py index f2e9096..f3bcc4d 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -164,4 +164,5 @@ class TestNoisedDataset(unittest.TestCase): "sqrt_alphas_cumprod_t" ] recovered = utils.modulo_with_wrapped_range(recovered, -np.pi, np.pi) - self.assertTrue(torch.allclose(recovered, orig_angles, atol=1e-5)) + delta = recovered - orig_angles + self.assertTrue(torch.allclose(delta, torch.zeros_like(delta), atol=1e-4), f"Got non-zero delta on de-noise: {delta}") diff --git a/tests/test_losses.py b/tests/test_losses.py index 8ba8ffc..5f82846 100644 --- a/tests/test_losses.py +++ b/tests/test_losses.py @@ -11,7 +11,7 @@ class TestRadianSmoothL1Loss(unittest.TestCase): """ Easy test of basic wrapping functionality """ - l = losses.radian_smooth_l1_loss(torch.tensor(0.1), 2 * torch.pi, beta=1.0) + l = losses.radian_smooth_l1_loss(torch.tensor(0.1), torch.tensor(2 * torch.pi), beta=1.0) self.assertAlmostEqual(0.0050, l.item()) def test_rounding(self):