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

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