src/jalview/ws/datamodel/alphafold/PAEContactMatrix.java

   1 package jalview.ws.datamodel.alphafold;
   2
   3 import java.awt.Color;
   4 import java.io.BufferedInputStream;
   5 import java.io.File;
   6 import java.io.FileInputStream;
   7 import java.io.IOException;
   8 import java.util.ArrayList;
   9 import java.util.BitSet;
  10 import java.util.HashMap;
  11 import java.util.Iterator;
  12 import java.util.List;
  13 import java.util.Map;
  14 import java.util.Map.Entry;
  15
  16 import org.json.simple.JSONObject;
  17
  18 import jalview.analysis.AverageDistanceEngine;
  19 import jalview.bin.Console;
  20 import jalview.datamodel.Annotation;
  21 import jalview.datamodel.BinaryNode;
  22 import jalview.datamodel.ContactListI;
  23 import jalview.datamodel.ContactListImpl;
  24 import jalview.datamodel.ContactListProviderI;
  25 import jalview.datamodel.ContactMatrixI;
  26 import jalview.datamodel.GroupSet;
  27 import jalview.datamodel.GroupSetI;
  28 import jalview.datamodel.Mapping;
  29 import jalview.datamodel.SequenceDummy;
  30 import jalview.datamodel.SequenceI;
  31 import jalview.io.DataSourceType;
  32 import jalview.io.FileFormatException;
  33 import jalview.io.FileParse;
  34 import jalview.util.MapList;
  35 import jalview.util.MapUtils;
  36 import jalview.ws.dbsources.EBIAlfaFold;
  37
  38 public class PAEContactMatrix extends
  39         MappableContactMatrix<PAEContactMatrix> implements ContactMatrixI
  40 {
  41   int maxrow = 0, maxcol = 0;
  42
  43   float[][] elements;
  44
  45   float maxscore;
  46
  47   @SuppressWarnings("unchecked")
  48   public PAEContactMatrix(SequenceI _refSeq, Map<String, Object> pae_obj)
  49           throws FileFormatException
  50   {
  51     setRefSeq(_refSeq);
  52     // convert the lists to primitive arrays and store
  53
  54     if (!MapUtils.containsAKey(pae_obj, "predicted_aligned_error", "pae"))
  55     {
  56       parse_version_1_pAE(pae_obj);
  57       return;
  58     }
  59     else
  60     {
  61       parse_version_2_pAE(pae_obj);
  62     }
  63   }
  64
  65   /**
  66    * construct a sequence associated PAE matrix directly from a float array
  67    *
  68    * @param _refSeq
  69    * @param matrix
  70    */
  71   public PAEContactMatrix(SequenceI _refSeq, float[][] matrix)
  72   {
  73     setRefSeq(_refSeq);
  74     maxcol = 0;
  75     for (float[] row : matrix)
  76     {
  77       if (row.length > maxcol)
  78       {
  79         maxcol = row.length;
  80       }
  81       maxscore = row[0];
  82       for (float f : row)
  83       {
  84         if (maxscore < f)
  85         {
  86           maxscore = f;
  87         }
  88       }
  89     }
  90     maxrow = matrix.length;
  91     elements = matrix;
  92
  93   }
  94
  95   /**
  96    * new matrix with specific mapping to a reference sequence
  97    *
  98    * @param newRefSeq
  99    * @param newFromMapList
 100    * @param elements2
 101    * @param grps2
 102    */
 103   public PAEContactMatrix(SequenceI newRefSeq, MapList newFromMapList,
 104           float[][] elements2, GroupSet grps2)
 105   {
 106     this(newRefSeq, elements2);
 107     toSeq = newFromMapList;
 108     grps = grps2;
 109   }
 110
 111   /**
 112    * parse a sane JSON representation of the pAE
 113    *
 114    * @param pae_obj
 115    */
 116   @SuppressWarnings("unchecked")
 117   private void parse_version_2_pAE(Map<String, Object> pae_obj)
 118   {
 119     maxscore = -1;
 120     // look for a maxscore element - if there is one...
 121     try
 122     {
 123       // this is never going to be reached by the integer rounding.. or is it ?
 124       maxscore = ((Double) MapUtils.getFirst(pae_obj,
 125               "max_predicted_aligned_error", "max_pae")).floatValue();
 126     } catch (Throwable t)
 127     {
 128       // ignore if a key is not found.
 129     }
 130     List<List<Long>> scoreRows = ((List<List<Long>>) MapUtils
 131             .getFirst(pae_obj, "predicted_aligned_error", "pae"));
 132     elements = new float[scoreRows.size()][scoreRows.size()];
 133     int row = 0, col = 0;
 134     for (List<Long> scoreRow : scoreRows)
 135     {
 136       Iterator<Long> scores = scoreRow.iterator();
 137       while (scores.hasNext())
 138       {
 139         Object d = scores.next();
 140
 141         if (d instanceof Double)
 142         {
 143           elements[row][col++] = ((Double) d).longValue();
 144         }
 145         else
 146         {
 147           elements[row][col++] = (float) ((Long) d).longValue();
 148         }
 149
 150         if (maxscore < elements[row][col - 1])
 151         {
 152           maxscore = elements[row][col - 1];
 153         }
 154       }
 155       row++;
 156       col = 0;
 157     }
 158     maxcol = length;
 159     maxrow = length;
 160   }
 161
 162   /**
 163    * v1 format got ditched 28th July 2022 see
 164    * https://alphafold.ebi.ac.uk/faq#:~:text=We%20updated%20the%20PAE%20JSON%20file%20format%20on%2028th%20July%202022
 165    *
 166    * @param pae_obj
 167    */
 168   @SuppressWarnings("unchecked")
 169   private void parse_version_1_pAE(Map<String, Object> pae_obj)
 170   {
 171     // assume indices are with respect to range defined by _refSeq on the
 172     // dataset refSeq
 173     Iterator<Long> rows = ((List<Long>) pae_obj.get("residue1")).iterator();
 174     Iterator<Long> cols = ((List<Long>) pae_obj.get("residue2")).iterator();
 175     // two pass - to allocate the elements array
 176     while (rows.hasNext())
 177     {
 178       int row = rows.next().intValue();
 179       int col = cols.next().intValue();
 180       if (maxrow < row)
 181       {
 182         maxrow = row;
 183       }
 184       if (maxcol < col)
 185       {
 186         maxcol = col;
 187       }
 188
 189     }
 190     rows = ((List<Long>) pae_obj.get("residue1")).iterator();
 191     cols = ((List<Long>) pae_obj.get("residue2")).iterator();
 192     Iterator<Double> scores = ((List<Double>) pae_obj.get("distance"))
 193             .iterator();
 194     elements = new float[maxrow][maxcol];
 195     while (scores.hasNext())
 196     {
 197       float escore = scores.next().floatValue();
 198       int row = rows.next().intValue();
 199       int col = cols.next().intValue();
 200       if (maxrow < row)
 201       {
 202         maxrow = row;
 203       }
 204       if (maxcol < col)
 205       {
 206         maxcol = col;
 207       }
 208       elements[row - 1][col - 1] = escore;
 209     }
 210
 211     maxscore = ((Double) MapUtils.getFirst(pae_obj,
 212             "max_predicted_aligned_error", "max_pae")).floatValue();
 213   }
 214
 215   @Override
 216   public ContactListI getContactList(final int column)
 217   {
 218     // final int _column;
 219     // if (toSeq != null)
 220     // {
 221     // int[] word = toSeq.locateInTo(column, column);
 222     // if (word == null)
 223     // {
 224     // return null;
 225     // }
 226     // _column = word[0];
 227     // }
 228     // else
 229     // {
 230     // _column = column;
 231     // }
 232     if (column < 0 || column >= elements.length)
 233     {
 234       return null;
 235     }
 236
 237     return new ContactListImpl(new ContactListProviderI()
 238     {
 239       @Override
 240       public int getPosition()
 241       {
 242         return column;
 243       }
 244
 245       @Override
 246       public int getContactHeight()
 247       {
 248         return maxcol - 1;
 249       }
 250
 251       @Override
 252       public double getContactAt(int mcolumn)
 253       {
 254         if (mcolumn < 0 || mcolumn >= elements[column].length)
 255         {
 256           return -1;
 257         }
 258         return elements[column][mcolumn];
 259       }
 260     });
 261   }
 262
 263   @Override
 264   protected double getElementAt(int _column, int i)
 265   {
 266     return elements[_column][i];
 267   }
 268
 269   @Override
 270   public float getMin()
 271   {
 272     return 0;
 273   }
 274
 275   @Override
 276   public float getMax()
 277   {
 278     return maxscore;
 279   }
 280
 281   @Override
 282   public String getAnnotDescr()
 283   {
 284     return "Predicted Alignment Error"
 285             + ((refSeq == null) ? "" : (" for " + refSeq.getName()));
 286   }
 287
 288   @Override
 289   public String getAnnotLabel()
 290   {
 291     StringBuilder label = new StringBuilder("PAE Matrix");
 292     // if (this.getReferenceSeq() != null)
 293     // {
 294     // label.append(":").append(this.getReferenceSeq().getDisplayId(false));
 295     // }
 296     return label.toString();
 297   }
 298
 299   public static final String PAEMATRIX = "PAE_MATRIX";
 300
 301   @Override
 302   public String getType()
 303   {
 304     return PAEMATRIX;
 305   }
 306
 307   @Override
 308   public int getWidth()
 309   {
 310     return length;
 311   }
 312
 313   @Override
 314   public int getHeight()
 315   {
 316     return length;
 317   }
 318
 319   public static void validateContactMatrixFile(String fileName)
 320           throws FileFormatException, IOException
 321   {
 322     FileInputStream infile = null;
 323     try
 324     {
 325       infile = new FileInputStream(new File(fileName));
 326     } catch (Throwable t)
 327     {
 328       new IOException("Couldn't open " + fileName, t);
 329     }
 330
 331     JSONObject paeDict = null;
 332     try
 333     {
 334       paeDict = EBIAlfaFold.parseJSONtoPAEContactMatrix(infile);
 335     } catch (Throwable t)
 336     {
 337       new FileFormatException("Couldn't parse " + fileName
 338               + " as a JSON dict or array containing a dict");
 339     }
 340
 341     PAEContactMatrix matrix = new PAEContactMatrix(
 342             new SequenceDummy("Predicted"), (Map<String, Object>) paeDict);
 343     if (matrix.getWidth() <= 0)
 344     {
 345       throw new FileFormatException(
 346               "No data in PAE matrix read from '" + fileName + "'");
 347     }
 348   }
 349
 350   @Override
 351   protected PAEContactMatrix newMappableContactMatrix(SequenceI newRefSeq,
 352           MapList newFromMapList)
 353   {
 354     PAEContactMatrix pae = new PAEContactMatrix(newRefSeq, newFromMapList,
 355             elements, new GroupSet(grps));
 356     return pae;
 357   }
 358 }