add HETATM support to CreateEnsemble (#135)

haddocking · Dec 16, 2024 · c5879be · c5879be
1 parent 7b4f3b6
commit c5879be
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/utils/utils.go b/utils/utils.go
@@ -112,7 +112,7 @@ func CreateEnsemble(p string, out string) error {
 		modelStr := ""
 		for modelScanner.Scan() {
 			line := modelScanner.Text()
-			if strings.HasPrefix(line, "ATOM") {
+			if strings.HasPrefix(line, "ATOM") || strings.HasPrefix(line, "HETATM") {
 				modelStr += line + "\n"
 			}
 		}

diff --git a/utils/utils_test.go b/utils/utils_test.go
@@ -108,9 +108,14 @@ func TestCreateEnsemble(t *testing.T) {
 		t.Errorf("Failed to write file: %s", err)
 	}
 	defer os.Remove("dummy.pdb")
+	err = os.WriteFile("hetatm.pdb", []byte("HETATM    1  N   ALA A   1      10.000  10.000  10.000  1.00  0.00           N\n"), 0644)
+	if err != nil {
+		t.Errorf("Failed to write file: %s", err)
+	}
+	defer os.Remove("hetatm.pdb")
 
 	// Make a list of PDB files and save it to a file
-	err = os.WriteFile("pdb-files.txt", []byte("\"dummy.pdb\"\n\"dummy.pdb\"\n\"dummy.pdb\"\n"), 0644)
+	err = os.WriteFile("pdb-files.txt", []byte("\"dummy.pdb\"\n\"dummy.pdb\"\n\"hetatm.pdb\"\n"), 0644)
 	if err != nil {
 		t.Errorf("Failed to write file: %s", err)
 	}