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"+((refSeq==null) ? "" : (" for " + refSeq.getName()));
 228   }
 229
 230   @Override
 231   public String getAnnotLabel()
 232   {
 233     StringBuilder label = new StringBuilder("PAE Matrix");
 234     //if (this.getReferenceSeq() != null)
 235     //{
 236     //  label.append(":").append(this.getReferenceSeq().getDisplayId(false));
 237     //}
 238     return label.toString();
 239   }
 240
 241   public static final String PAEMATRIX = "PAE_MATRIX";
 242
 243   @Override
 244   public String getType()
 245   {
 246     return PAEMATRIX;
 247   }
 248
 249   @Override
 250   public int getWidth()
 251   {
 252     return length;
 253   }
 254
 255   @Override
 256   public int getHeight()
 257   {
 258     return length;
 259   }
 260   List<BitSet> groups=null;
 261   @Override
 262   public boolean hasGroups()
 263   {
 264     return groups!=null;
 265   }
 266   String newick=null;
 267   @Override
 268   public String getNewick()
 269   {
 270     return newick;
 271   }
 272   @Override
 273   public boolean hasTree()
 274   {
 275     return newick!=null && newick.length()>0;
 276   }
 277   boolean abs;
 278   double thresh;
 279   String treeType=null;
 280   public void makeGroups(float thresh,boolean abs)
 281   {
 282     AverageDistanceEngine clusterer = new AverageDistanceEngine(null, null, this);
 283     double height = clusterer.findHeight(clusterer.getTopNode());
 284     newick = new jalview.io.NewickFile(clusterer.getTopNode(),false,true).print();
 285     treeType = "UPGMA";
 286     Console.trace("Newick string\n"+newick);
 287
 288     List<BinaryNode> nodegroups;
 289     if (abs ? height > thresh : 0 < thresh && thresh < 1)
 290     {
 291       float cut = abs ? (float) (thresh / height) : thresh;
 292       Console.debug("Threshold "+cut+" for height="+height);
 293
 294       nodegroups = clusterer.groupNodes(cut);
 295     }
 296     else
 297     {
 298       nodegroups = new ArrayList<BinaryNode>();
 299       nodegroups.add(clusterer.getTopNode());
 300     }
 301     this.abs=abs;
 302     this.thresh=thresh;
 303     groups = new ArrayList<>();
 304     for (BinaryNode root:nodegroups)
 305     {
 306       BitSet gpset=new BitSet();
 307       for (BinaryNode leaf:clusterer.findLeaves(root))
 308       {
 309         gpset.set((Integer)leaf.element());
 310       }
 311       groups.add(gpset);
 312     }
 313   }
 314   @Override
 315   public void updateGroups(List<BitSet> colGroups)
 316   {
 317     if (colGroups!=null)
 318     {
 319       groups=colGroups;
 320     }
 321   }
 322   @Override
 323   public BitSet getGroupsFor(int column)
 324   {
 325     for (BitSet gp:groups) {
 326       if (gp.get(column))
 327       {
 328         return gp;
 329       }
 330     }
 331     return ContactMatrixI.super.getGroupsFor(column);
 332   }
 333
 334   HashMap<BitSet,Color> colorMap = new HashMap<>();
 335   @Override
 336   public Color getColourForGroup(BitSet bs)
 337   {
 338     if (bs==null) {
 339       return Color.white;
 340     }
 341     Color groupCol=colorMap.get(bs);
 342     if (groupCol==null)
 343     {
 344       return Color.white;
 345     }
 346     return groupCol;
 347   }
 348   @Override
 349   public void setColorForGroup(BitSet bs,Color color)
 350   {
 351     colorMap.put(bs,color);
 352   }
 353   public void restoreGroups(List<BitSet> newgroups, String treeMethod,
 354           String tree, double thresh2)
 355   {
 356     treeType=treeMethod;
 357     groups = newgroups;
 358     thresh=thresh2;
 359     newick =tree;
 360
 361   }
 362   @Override
 363   public boolean hasCutHeight() {
 364     return groups!=null && thresh!=0;
 365   }
 366   @Override
 367   public double getCutHeight()
 368   {
 369     return thresh;
 370   }
 371   @Override
 372   public String getTreeMethod()
 373   {
 374     return treeType;
 375   }
 376 }