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

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