Writing out entity info in CIF

josemduarte · josemduarte · commit 615a57045c25 · 2023-05-10T16:57:21.000-07:00
diff --git a/biojava-integrationtest/src/test/java/org/biojava/nbio/structure/test/io/cif/CifFileSupplierIntegrationTest.java b/biojava-integrationtest/src/test/java/org/biojava/nbio/structure/test/io/cif/CifFileSupplierIntegrationTest.java
@@ -59,6 +59,12 @@ public void test6ELW() throws IOException {
     	// a structure with insertion codes
     	testRoundTrip("6ELW");
     }
+
+    @Test
+    public void test4HHB() throws IOException {
+        // a structure with multiple poly entities
+        testRoundTrip("4HHB");
+    }
     
     private static void testRoundTrip(String pdbId) throws IOException {
         URL url = new URL("https://files.rcsb.org/download/" + pdbId + ".cif");
@@ -121,6 +127,14 @@ private static void testRoundTrip(String pdbId) throws IOException {
         // Test cell and symmetry
         assertEquals(originalStruct.getCrystallographicInfo().getSpaceGroup(),
                 readStruct.getCrystallographicInfo().getSpaceGroup());
+
+        // entity
+        assertEquals(originalStruct.getEntityInfos().size(), readStruct.getEntityInfos().size());
+        for (int i=0; i<originalStruct.getEntityInfos().size(); i++) {
+            assertEquals(originalStruct.getEntityInfos().get(i).getMolId(), readStruct.getEntityInfos().get(i).getMolId());
+            assertEquals(originalStruct.getEntityInfos().get(i).getType(), readStruct.getEntityInfos().get(i).getType());
+        }
+
     }
 
     /**
diff --git a/biojava-structure/src/main/java/org/biojava/nbio/structure/io/cif/AbstractCifFileSupplier.java b/biojava-structure/src/main/java/org/biojava/nbio/structure/io/cif/AbstractCifFileSupplier.java
@@ -1,12 +1,6 @@
 package org.biojava.nbio.structure.io.cif;
 
-import org.biojava.nbio.structure.Atom;
-import org.biojava.nbio.structure.Chain;
-import org.biojava.nbio.structure.Element;
-import org.biojava.nbio.structure.EntityType;
-import org.biojava.nbio.structure.Group;
-import org.biojava.nbio.structure.GroupType;
-import org.biojava.nbio.structure.Structure;
+import org.biojava.nbio.structure.*;
 import org.biojava.nbio.structure.xtal.CrystalCell;
 import org.biojava.nbio.structure.xtal.SpaceGroup;
 import org.rcsb.cif.CifBuilder;
@@ -23,8 +17,10 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.function.Consumer;
 import java.util.stream.Collector;
+import java.util.stream.Collectors;
 
 /**
  * Convert a BioJava object to a CifFile.
@@ -41,9 +37,11 @@ protected CifFile getInternal(Structure structure, List<WrappedAtom> wrappedAtom
         SpaceGroup spaceGroup = structure.getPDBHeader().getCrystallographicInfo().getSpaceGroup();
         // atom_site
         Category atomSite = wrappedAtoms.stream().collect(toAtomSite());
+        // entity information
+        List<EntityInfo> entityInfos = structure.getEntityInfos();
 
         MmCifBlockBuilder blockBuilder = CifBuilder.enterFile(StandardSchemata.MMCIF)
-                .enterBlock(structure.getPDBCode());
+                .enterBlock(structure.getPdbId() == null? "" : structure.getPdbId().getId());
 
         blockBuilder.enterStructKeywords().enterText()
         .add(String.join(", ", structure.getPDBHeader().getKeywords()))
@@ -92,6 +90,49 @@ protected CifFile getInternal(Structure structure, List<WrappedAtom> wrappedAtom
                     .leaveCategory();
         }
 
+        if (entityInfos != null) {
+
+            String[] entityIds = new String[entityInfos.size()];
+            String[] entityTypes = new String[entityInfos.size()];
+            String[] entityDescriptions = new String[entityInfos.size()];
+
+            for (int i=0; i<entityInfos.size(); i++) {
+                EntityInfo e = entityInfos.get(i);
+                entityIds[i] = Integer.toString(e.getMolId());
+                entityTypes[i] = e.getType().getEntityType();
+                entityDescriptions[i] = e.getDescription();
+            }
+
+            String[] polyEntityIds = entityInfos.stream().filter(e -> e.getType() == EntityType.POLYMER).map(e -> Integer.toString(e.getMolId())).collect(Collectors.toList()).toArray(new String[]{});
+            String[] entitySeqs = entityInfos.stream().filter(e -> e.getType() == EntityType.POLYMER).map(e -> e.getChains().get(0).getSeqResSequence()).collect(Collectors.toList()).toArray(new String[]{});
+
+            blockBuilder.enterEntity()
+                    .enterId()
+                    .add(entityIds)
+                    .leaveColumn()
+
+                    .enterType()
+                    .add(entityTypes)
+                    .leaveColumn()
+
+                    .enterPdbxDescription()
+                    .add(entityDescriptions)
+                    .leaveColumn()
+
+                    .leaveCategory();
+
+            blockBuilder.enterEntityPoly()
+                    .enterEntityId()
+                    .add(polyEntityIds)
+                    .leaveColumn()
+
+                    .enterPdbxSeqOneLetterCodeCan()
+                    .add(entitySeqs)
+                    .leaveColumn()
+
+                    .leaveCategory();
+        }
+
         return blockBuilder.leaveBlock().leaveFile();
     }
 
