JAL-629 improvements to argparser toString. Improvements to cli paeFile structure...
[jalview.git] / src / jalview / ws / dbsources / EBIAlfaFold.java
index f98ae6d..5165d04 100644 (file)
@@ -23,11 +23,19 @@ package jalview.ws.dbsources;
 
 import java.io.File;
 import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
+import java.util.Date;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import org.json.simple.JSONArray;
+import org.json.simple.JSONObject;
+import org.json.simple.parser.ParseException;
+
 import com.stevesoft.pat.Regex;
 
 import jalview.api.FeatureSettingsModelI;
@@ -39,11 +47,15 @@ import jalview.datamodel.DBRefEntry;
 import jalview.datamodel.PDBEntry;
 import jalview.datamodel.SequenceFeature;
 import jalview.datamodel.SequenceI;
+import jalview.gui.Desktop;
 import jalview.io.DataSourceType;
 import jalview.io.FileFormat;
 import jalview.io.FileFormatI;
 import jalview.io.FormatAdapter;
 import jalview.io.PDBFeatureSettings;
+import jalview.structure.StructureImportSettings.TFType;
+import jalview.structure.StructureMapping;
+import jalview.structure.StructureSelectionManager;
 import jalview.util.MessageManager;
 import jalview.util.Platform;
 import jalview.ws.datamodel.alphafold.PAEContactMatrix;
@@ -187,6 +199,7 @@ public class EBIAlfaFold extends EbiFileRetrievedProxy
       {
         return null;
       }
+      // TODO Get the PAE file somewhere around here and remove from JmolParser
 
       pdbAlignment = importDownloadedStructureFromUrl(alphaFoldCif, tmpFile,
               id, chain, getDbSource(), getDbVersion());
@@ -209,22 +222,22 @@ public class EBIAlfaFold extends EbiFileRetrievedProxy
   }
 
   /**
-   * get an alphafold pAE for the given id, and add it to sequence 0 in
-   * pdbAlignment (assuming it came from structurefile parser).
+   * get an alphafold pAE for the given id and return the File object of the
+   * downloaded (temp) file
    * 
    * @param id
    * @param pdbAlignment
    * @param retrievalUrl
    *          - URL of .mmcif from EBI-AlphaFold - will be used to generate the
    *          pAE URL automatically
+   * @throws IOException
    * @throws Exception
    */
-  public static void retrieve_AlphaFold_pAE(String id,
-          AlignmentI pdbAlignment, String retrievalUrl) throws Exception
+  public static File fetchAlphaFoldPAE(String id, String retrievalUrl)
+          throws IOException
   {
     // import PAE as contact matrix - assume this will work if there was a
     // model
-    File pae = File.createTempFile(id, "pae_json");
     String paeURL = getAlphaFoldPaeDownloadUrl(id, AF_VERSION);
 
     if (retrievalUrl != null)
@@ -233,24 +246,147 @@ public class EBIAlfaFold extends EbiFileRetrievedProxy
       paeURL = retrievalUrl.replace("model", "predicted_aligned_error")
               .replace(".cif", ".json");
     }
+
+    // check the cache
+    File pae = paeDownloadCache.get(paeURL);
+    if (pae != null && pae.exists() && (new Date().getTime()
+            - pae.lastModified()) < PAE_CACHE_STALE_TIME)
+    {
+      Console.debug(
+              "Using existing file in PAE cache for '" + paeURL + "'");
+      return pae;
+    }
+
+    try
+    {
+      pae = File.createTempFile(id == null ? "af_pae" : id, "pae_json");
+    } catch (IOException e)
+    {
+      e.printStackTrace();
+    }
     Console.debug("Downloading pae from " + paeURL + " to " + pae.toString()
             + "");
-
     try
     {
       UrlDownloadClient.download(paeURL, pae);
-      FileInputStream pae_input = new FileInputStream(pae);
+    } catch (IOException e)
+    {
+      throw e;
+    }
+    // cache and it if successful
+    paeDownloadCache.put(paeURL, pae);
+    return pae;
+  }
 
-      if (!importPaeJSONAsContactMatrix(pdbAlignment, pae_input))
+  /**
+   * get an alphafold pAE for the given id, and add it to sequence 0 in
+   * pdbAlignment (assuming it came from structurefile parser).
+   * 
+   * @param id
+   * @param pdbAlignment
+   * @param retrievalUrl
+   *          - URL of .mmcif from EBI-AlphaFold - will be used to generate the
+   *          pAE URL automatically
+   * @throws IOException
+   * @throws Exception
+   */
+  public static void retrieve_AlphaFold_pAE(String id,
+          AlignmentI pdbAlignment, String retrievalUrl) throws IOException
+  {
+    File pae = fetchAlphaFoldPAE(id, retrievalUrl);
+    addAlphaFoldPAE(pdbAlignment, pae, 0, null, false, false, null);
+  }
+
+  public static void addAlphaFoldPAE(AlignmentI pdbAlignment, File pae,
+          int index, String id, boolean isStruct, boolean isStructId,
+          String label)
+  {
+    FileInputStream paeInput = null;
+    try
+    {
+      paeInput = new FileInputStream(pae);
+    } catch (FileNotFoundException e)
+    {
+      Console.error(
+              "Could not find pAE file '" + pae.getAbsolutePath() + "'", e);
+      return;
+    }
+
+    if (isStruct)
+    {
+      StructureSelectionManager ssm = StructureSelectionManager
+              .getStructureSelectionManager(Desktop.instance);
+      if (ssm != null)
       {
-        Console.warn("Couln't import contact matrix from " + paeURL
-                + " (stored in " + pae.toString() + ")");
+        String structFilename = isStructId ? ssm.findFileForPDBId(id) : id;
+        addPAEToStructure(ssm, structFilename, pae, label);
       }
-    } catch (Exception pae_ex)
+
+    }
+    else
+    {
+      // attach to sequence?!
+      try
+      {
+        if (!importPaeJSONAsContactMatrixToSequence(pdbAlignment, paeInput,
+                index, id, label))
+        {
+          Console.warn("Could not import contact matrix from '"
+                  + pae.getAbsolutePath() + "' to sequence.");
+        }
+      } catch (IOException e1)
+      {
+        Console.error("Error when importing pAE file '"
+                + pae.getAbsolutePath() + "'", e1);
+      } catch (ParseException e2)
+      {
+        Console.error("Error when parsing pAE file '"
+                + pae.getAbsolutePath() + "'", e2);
+      }
+    }
+
+  }
+
+  public static void addPAEToStructure(StructureSelectionManager ssm,
+          String structFilename, File pae, String label)
+  {
+    FileInputStream paeInput = null;
+    try
+    {
+      paeInput = new FileInputStream(pae);
+    } catch (FileNotFoundException e)
     {
-      Console.error("Couldn't download PAE", pae_ex);
+      Console.error(
+              "Could not find pAE file '" + pae.getAbsolutePath() + "'", e);
+      return;
     }
+    if (ssm == null)
+    {
+      ssm = StructureSelectionManager
+              .getStructureSelectionManager(Desktop.instance);
+    }
+    if (ssm != null)
+    {
+      StructureMapping[] smArray = ssm.getMapping(structFilename);
 
+      try
+      {
+        if (!importPaeJSONAsContactMatrixToStructure(smArray, paeInput,
+                label))
+        {
+          Console.warn("Could not import contact matrix from '"
+                  + pae.getAbsolutePath() + "' to structure.");
+        }
+      } catch (IOException e1)
+      {
+        Console.error("Error when importing pAE file '"
+                + pae.getAbsolutePath() + "'", e1);
+      } catch (ParseException e2)
+      {
+        Console.error("Error when parsing pAE file '"
+                + pae.getAbsolutePath() + "'", e2);
+      }
+    }
   }
 
   /**
@@ -260,24 +396,119 @@ public class EBIAlfaFold extends EbiFileRetrievedProxy
    * @param pdbAlignment
    * @param pae_input
    * @return true if there was a pAE matrix added
+   * @throws ParseException
+   * @throws IOException
    * @throws Exception
    */
-  public static boolean importPaeJSONAsContactMatrix(
-          AlignmentI pdbAlignment, InputStream pae_input) throws Exception
+  public static boolean importPaeJSONAsContactMatrixToSequence(
+          AlignmentI pdbAlignment, InputStream pae_input, int index,
+          String seqId, String label) throws IOException, ParseException
   {
+    SequenceI sequence = null;
+    if (seqId == null)
+    {
+      int seqToGet = index > 0 ? index : 0;
+      sequence = pdbAlignment.getSequenceAt(seqToGet);
+    }
+    if (sequence == null)
+    {
+      SequenceI[] sequences = pdbAlignment.findSequenceMatch(seqId);
+      if (sequences == null || sequences.length < 1)
+      {
+        Console.warn("Could not find sequence with id '" + seqId
+                + "' to attach pAE matrix to. Ignoring matrix.");
+        return false;
+      }
+      else
+      {
+        sequence = sequences[0]; // just use the first sequence with this seqId
+      }
+    }
+    if (sequence == null)
+    {
+      return false;
+    }
+    return importPaeJSONAsContactMatrixToSequence(pdbAlignment, pae_input,
+            sequence, label);
+  }
 
-    List<Object> pae_obj = (List<Object>) Platform.parseJSON(pae_input);
-    if (pae_obj == null)
+  public static boolean importPaeJSONAsContactMatrixToSequence(
+          AlignmentI pdbAlignment, InputStream pae_input,
+          SequenceI sequence, String label)
+          throws IOException, ParseException
+  {
+    JSONObject paeDict = parseJSONtoPAEContactMatrix(pae_input);
+    if (paeDict == null)
     {
+      Console.debug("JSON file did not parse properly.");
       return false;
     }
-    ContactMatrixI matrix = new PAEContactMatrix(
-            pdbAlignment.getSequenceAt(0),
-            (Map<String, Object>) pae_obj.get(0));
+    ContactMatrixI matrix = new PAEContactMatrix(sequence,
+            (Map<String, Object>) paeDict);
 
-    AlignmentAnnotation cmannot = pdbAlignment.getSequenceAt(0)
-            .addContactList(matrix);
+    AlignmentAnnotation cmannot = sequence.addContactList(matrix);
+    if (label != null)
+      cmannot.label = label;
     pdbAlignment.addAnnotation(cmannot);
+
+    return true;
+  }
+
+  public static JSONObject parseJSONtoPAEContactMatrix(
+          InputStream pae_input) throws IOException, ParseException
+  {
+    Object paeJson = Platform.parseJSON(pae_input);
+    JSONObject paeDict = null;
+    if (paeJson instanceof JSONObject)
+    {
+      paeDict = (JSONObject) paeJson;
+    }
+    else if (paeJson instanceof JSONArray)
+    {
+      JSONArray jsonArray = (JSONArray) paeJson;
+      if (jsonArray.size() > 0)
+        paeDict = (JSONObject) jsonArray.get(0);
+    }
+
+    return paeDict;
+  }
+
+  public static boolean importPaeJSONAsContactMatrixToStructure(
+          StructureMapping[] smArray, InputStream paeInput, String label)
+          throws IOException, ParseException
+  {
+    boolean someDone = false;
+    for (StructureMapping sm : smArray)
+    {
+      boolean thisDone = importPaeJSONAsContactMatrixToStructure(sm,
+              paeInput, label);
+      someDone |= thisDone;
+    }
+    return someDone;
+  }
+
+  public static boolean importPaeJSONAsContactMatrixToStructure(
+          StructureMapping sm, InputStream paeInput, String label)
+          throws IOException, ParseException
+  {
+    JSONObject pae_obj = parseJSONtoPAEContactMatrix(paeInput);
+    if (pae_obj == null)
+    {
+      Console.debug("JSON file did not parse properly.");
+      return false;
+    }
+
+    ContactMatrixI matrix = new PAEContactMatrix(sm.getSequence(),
+            (Map<String, Object>) pae_obj);
+
+    AlignmentAnnotation cmannot = sm.getSequence().addContactList(matrix);
+    if (label != null)
+    {
+      Console.debug("Setting annotation label to '" + label + "'");
+      cmannot.label = label;
+    }
+    sm.getSequence().addAlignmentAnnotation(cmannot);
+
     return true;
   }
 
@@ -301,8 +532,10 @@ public class EBIAlfaFold extends EbiFileRetrievedProxy
     String file = tmpFile.getAbsolutePath();
     // todo get rid of Type and use FileFormatI instead?
     FileFormatI fileFormat = FileFormat.MMCif;
-    AlignmentI pdbAlignment = new FormatAdapter().readFile(tmpFile,
-            DataSourceType.FILE, fileFormat);
+    TFType tempfacType = TFType.PLDDT;
+    AlignmentI pdbAlignment = new FormatAdapter().readFile(tmpFile, file,
+            DataSourceType.FILE, fileFormat, tempfacType);
+
     if (pdbAlignment != null)
     {
       List<SequenceI> toremove = new ArrayList<SequenceI>();
@@ -315,7 +548,6 @@ public class EBIAlfaFold extends EbiFileRetrievedProxy
           if (pid.getFile() == file)
           {
             chid = pid.getChainCode();
-
           }
         }
         if (chain == null || (chid != null && (chid.equals(chain)
@@ -438,4 +670,9 @@ public class EBIAlfaFold extends EbiFileRetrievedProxy
     return new PDBFeatureSettings();
   }
 
+  // days * 86400000
+  private static final long PAE_CACHE_STALE_TIME = 1 * 86400000;
+
+  private static Map<String, File> paeDownloadCache = new HashMap<>();
+
 }