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
  15 import org.json.simple.JSONObject;
  16
  17 import jalview.analysis.AverageDistanceEngine;
  18 import jalview.bin.Console;
  19 import jalview.datamodel.BinaryNode;
  20 import jalview.datamodel.ContactListI;
  21 import jalview.datamodel.ContactListImpl;
  22 import jalview.datamodel.ContactListProviderI;
  23 import jalview.datamodel.ContactMatrixI;
  24 import jalview.datamodel.SequenceDummy;
  25 import jalview.datamodel.SequenceI;
  26 import jalview.io.DataSourceType;
  27 import jalview.io.FileFormatException;
  28 import jalview.io.FileParse;
  29 import jalview.util.MapUtils;
  30 import jalview.ws.dbsources.EBIAlfaFold;
  31
  32 public class PAEContactMatrix implements ContactMatrixI
  33 {
  34
  35   SequenceI refSeq = null;
  36
  37   /**
  38    * the length that refSeq is expected to be (excluding gaps, of course)
  39    */
  40   int length;
  41
  42   int maxrow = 0, maxcol = 0;
  43
  44   int[] indices1, indices2;
  45
  46   float[][] elements;
  47
  48   float maxscore;
  49
  50   private void setRefSeq(SequenceI _refSeq)
  51   {
  52     refSeq = _refSeq;
  53     while (refSeq.getDatasetSequence() != null)
  54     {
  55       refSeq = refSeq.getDatasetSequence();
  56     }
  57     length = _refSeq.getEnd() - _refSeq.getStart() + 1;
  58   }
  59
  60   @SuppressWarnings("unchecked")
  61   public PAEContactMatrix(SequenceI _refSeq, Map<String, Object> pae_obj) throws FileFormatException
  62   {
  63     setRefSeq(_refSeq);
  64     // convert the lists to primitive arrays and store
  65
  66     if (!MapUtils.containsAKey(pae_obj, "predicted_aligned_error", "pae"))
  67     {
  68       parse_version_1_pAE(pae_obj);
  69       return;
  70     }
  71     else
  72     {
  73       parse_version_2_pAE(pae_obj);
  74     }
  75   }
  76
  77   /**
  78    * construct a sequence associated PAE matrix directly from a float array
  79    *
  80    * @param _refSeq
  81    * @param matrix
  82    */
  83   public PAEContactMatrix(SequenceI _refSeq, float[][] matrix)
  84   {
  85     setRefSeq(_refSeq);
  86     maxcol = 0;
  87     for (float[] row : matrix)
  88     {
  89       if (row.length > maxcol)
  90       {
  91         maxcol = row.length;
  92       }
  93       maxscore = row[0];
  94       for (float f : row)
  95       {
  96         if (maxscore < f)
  97         {
  98           maxscore = f;
  99         }
 100       }
 101     }
 102     maxrow = matrix.length;
 103     elements = matrix;
 104
 105   }
 106
 107   /**
 108    * parse a sane JSON representation of the pAE
 109    *
 110    * @param pae_obj
 111    */
 112   @SuppressWarnings("unchecked")
 113   private void parse_version_2_pAE(Map<String, Object> pae_obj)
 114   {
 115     maxscore = -1;
 116     // look for a maxscore element - if there is one...
 117     try
 118     {
 119       // this is never going to be reached by the integer rounding.. or is it ?
 120       maxscore = ((Double) MapUtils.getFirst(pae_obj,
 121               "max_predicted_aligned_error", "max_pae")).floatValue();
 122     } catch (Throwable t)
 123     {
 124       // ignore if a key is not found.
 125     }
 126     List<List<Long>> scoreRows = ((List<List<Long>>) MapUtils
 127             .getFirst(pae_obj, "predicted_aligned_error", "pae"));
 128     elements = new float[scoreRows.size()][scoreRows.size()];
 129     int row = 0, col = 0;
 130     for (List<Long> scoreRow : scoreRows)
 131     {
 132       Iterator<Long> scores = scoreRow.iterator();
 133       while (scores.hasNext())
 134       {
 135         Object d = scores.next();
 136
 137         if (d instanceof Double)
 138         {
 139           elements[row][col++] = ((Double) d).longValue();
 140         }
 141         else
 142         {
 143           elements[row][col++] = (float) ((Long) d).longValue();
 144         }
 145
 146         if (maxscore < elements[row][col - 1])
 147         {
 148           maxscore = elements[row][col - 1];
 149         }
 150       }
 151       row++;
 152       col = 0;
 153     }
 154     maxcol = length;
 155     maxrow = length;
 156   }
 157
 158   /**
 159    * v1 format got ditched 28th July 2022 see
 160    * https://alphafold.ebi.ac.uk/faq#:~:text=We%20updated%20the%20PAE%20JSON%20file%20format%20on%2028th%20July%202022
 161    *
 162    * @param pae_obj
 163    */
 164   @SuppressWarnings("unchecked")
 165   private void parse_version_1_pAE(Map<String, Object> pae_obj)
 166   {
 167     // assume indices are with respect to range defined by _refSeq on the
 168     // dataset refSeq
 169     Iterator<Long> rows = ((List<Long>) pae_obj.get("residue1")).iterator();
 170     Iterator<Long> cols = ((List<Long>) pae_obj.get("residue2")).iterator();
 171     // two pass - to allocate the elements array
 172     while (rows.hasNext())
 173     {
 174       int row = rows.next().intValue();
 175       int col = cols.next().intValue();
 176       if (maxrow < row)
 177       {
 178         maxrow = row;
 179       }
 180       if (maxcol < col)
 181       {
 182         maxcol = col;
 183       }
 184
 185     }
 186     rows = ((List<Long>) pae_obj.get("residue1")).iterator();
 187     cols = ((List<Long>) pae_obj.get("residue2")).iterator();
 188     Iterator<Double> scores = ((List<Double>) pae_obj.get("distance"))
 189             .iterator();
 190     elements = new float[maxrow][maxcol];
 191     while (scores.hasNext())
 192     {
 193       float escore = scores.next().floatValue();
 194       int row = rows.next().intValue();
 195       int col = cols.next().intValue();
 196       if (maxrow < row)
 197       {
 198         maxrow = row;
 199       }
 200       if (maxcol < col)
 201       {
 202         maxcol = col;
 203       }
 204       elements[row - 1][col - 1] = escore;
 205     }
 206
 207     maxscore = ((Double) MapUtils.getFirst(pae_obj,
 208             "max_predicted_aligned_error", "max_pae")).floatValue();
 209   }
 210
 211   @Override
 212   public ContactListI getContactList(final int _column)
 213   {
 214     if (_column < 0 || _column >= elements.length)
 215     {
 216       return null;
 217     }
 218
 219     return new ContactListImpl(new ContactListProviderI()
 220     {
 221       @Override
 222       public int getPosition()
 223       {
 224         return _column;
 225       }
 226
 227       @Override
 228       public int getContactHeight()
 229       {
 230         return maxcol - 1;
 231       }
 232
 233       @Override
 234       public double getContactAt(int column)
 235       {
 236         if (column < 0 || column >= elements[_column].length)
 237         {
 238           return -1;
 239         }
 240         return elements[_column][column];
 241       }
 242     });
 243   }
 244
 245   @Override
 246   public float getMin()
 247   {
 248     return 0;
 249   }
 250
 251   @Override
 252   public float getMax()
 253   {
 254     return maxscore;
 255   }
 256
 257   @Override
 258   public boolean hasReferenceSeq()
 259   {
 260     return (refSeq != null);
 261   }
 262
 263   @Override
 264   public SequenceI getReferenceSeq()
 265   {
 266     return refSeq;
 267   }
 268
 269   @Override
 270   public String getAnnotDescr()
 271   {
 272     return "Predicted Alignment Error"+((refSeq==null) ? "" : (" for " + refSeq.getName()));
 273   }
 274
 275   @Override
 276   public String getAnnotLabel()
 277   {
 278     StringBuilder label = new StringBuilder("PAE Matrix");
 279     //if (this.getReferenceSeq() != null)
 280     //{
 281     //  label.append(":").append(this.getReferenceSeq().getDisplayId(false));
 282     //}
 283     return label.toString();
 284   }
 285
 286   public static final String PAEMATRIX = "PAE_MATRIX";
 287
 288   @Override
 289   public String getType()
 290   {
 291     return PAEMATRIX;
 292   }
 293
 294   @Override
 295   public int getWidth()
 296   {
 297     return length;
 298   }
 299
 300   @Override
 301   public int getHeight()
 302   {
 303     return length;
 304   }
 305   List<BitSet> groups=null;
 306   @Override
 307   public boolean hasGroups()
 308   {
 309     return groups!=null;
 310   }
 311   String newick=null;
 312   @Override
 313   public String getNewick()
 314   {
 315     return newick;
 316   }
 317   @Override
 318   public boolean hasTree()
 319   {
 320     return newick!=null && newick.length()>0;
 321   }
 322   boolean abs;
 323   double thresh;
 324   String treeType=null;
 325   public void makeGroups(float thresh,boolean abs)
 326   {
 327     AverageDistanceEngine clusterer = new AverageDistanceEngine(null, null, this);
 328     double height = clusterer.findHeight(clusterer.getTopNode());
 329     newick = new jalview.io.NewickFile(clusterer.getTopNode(),false,true).print();
 330     treeType = "UPGMA";
 331     Console.trace("Newick string\n"+newick);
 332
 333     List<BinaryNode> nodegroups;
 334     if (abs ? height > thresh : 0 < thresh && thresh < 1)
 335     {
 336       float cut = abs ? (float) (thresh / height) : thresh;
 337       Console.debug("Threshold "+cut+" for height="+height);
 338
 339       nodegroups = clusterer.groupNodes(cut);
 340     }
 341     else
 342     {
 343       nodegroups = new ArrayList<BinaryNode>();
 344       nodegroups.add(clusterer.getTopNode());
 345     }
 346     this.abs=abs;
 347     this.thresh=thresh;
 348     groups = new ArrayList<>();
 349     for (BinaryNode root:nodegroups)
 350     {
 351       BitSet gpset=new BitSet();
 352       for (BinaryNode leaf:clusterer.findLeaves(root))
 353       {
 354         gpset.set((Integer)leaf.element());
 355       }
 356       groups.add(gpset);
 357     }
 358   }
 359   @Override
 360   public void updateGroups(List<BitSet> colGroups)
 361   {
 362     if (colGroups!=null)
 363     {
 364       groups=colGroups;
 365     }
 366   }
 367   @Override
 368   public BitSet getGroupsFor(int column)
 369   {
 370     for (BitSet gp:groups) {
 371       if (gp.get(column))
 372       {
 373         return gp;
 374       }
 375     }
 376     return ContactMatrixI.super.getGroupsFor(column);
 377   }
 378
 379   HashMap<BitSet,Color> colorMap = new HashMap<>();
 380   @Override
 381   public Color getColourForGroup(BitSet bs)
 382   {
 383     if (bs==null) {
 384       return Color.white;
 385     }
 386     Color groupCol=colorMap.get(bs);
 387     if (groupCol==null)
 388     {
 389       return Color.white;
 390     }
 391     return groupCol;
 392   }
 393   @Override
 394   public void setColorForGroup(BitSet bs,Color color)
 395   {
 396     colorMap.put(bs,color);
 397   }
 398   public void restoreGroups(List<BitSet> newgroups, String treeMethod,
 399           String tree, double thresh2)
 400   {
 401     treeType=treeMethod;
 402     groups = newgroups;
 403     thresh=thresh2;
 404     newick =tree;
 405
 406   }
 407   @Override
 408   public boolean hasCutHeight() {
 409     return groups!=null && thresh!=0;
 410   }
 411   @Override
 412   public double getCutHeight()
 413   {
 414     return thresh;
 415   }
 416   @Override
 417   public String getTreeMethod()
 418   {
 419     return treeType;
 420   }
 421
 422   public static void validateContactMatrixFile(String fileName) throws FileFormatException,IOException
 423   {
 424     FileInputStream infile=null;
 425     try {
 426       infile = new FileInputStream(new File(fileName));
 427     } catch (Throwable t)
 428     {
 429       new IOException("Couldn't open "+fileName,t);
 430     }
 431
 432
 433     JSONObject paeDict=null;
 434     try {
 435       paeDict = EBIAlfaFold.parseJSONtoPAEContactMatrix(infile);
 436     } catch (Throwable t)
 437     {
 438       new FileFormatException("Couldn't parse "+fileName+" as a JSON dict or array containing a dict");
 439     }
 440
 441     PAEContactMatrix matrix = new PAEContactMatrix(new SequenceDummy("Predicted"), (Map<String,Object>)paeDict);
 442     if (matrix.getWidth()<=0)
 443     {
 444       throw new FileFormatException("No data in PAE matrix read from '"+fileName+"'");
 445     }
 446   }
 447 }