AlphaFold vs. JAX

[davis797@…]

The following bug report has been submitted:
Platform:        macOS-10.16-x86_64-i386-64bit
ChimeraX Version: 1.3 (2021-12-08 23:08:33 UTC)
Description
Hi ChimeraX Team, 

I noticed yesterday on my mid-2015 MBP’s v1.3 (2021-12-08) production release of ChimeraX that the linked alphafold_predict_colab.ipynb under the Structure Prediction > Alphafold interface fails to compute any models (link attached). The model error tracebacks I am receiving (file attached) point to an error on Alphafold’s side of things: 

  File "/usr/local/lib/python3.7/dist-packages/alphafold/model/modules.py", line 1733, in __call__
    msa_activations = jax.ops.index_add(msa_activations, 0,
AttributeError: module 'jax.ops' has no attribute ‘index_add’

I believe this is an issue with the JAX library, where the jax.ops.index_add function was removed on March 16th, 2022 with with JAX’s v0.3.2 update (https://jax.readthedocs.io/en/latest/changelog.html). It seems this has been temporarily fixed elsewhere by using v<0.3.0 (https://github.com/sokrypton/ColabFold/blob/28e9d5812d54cdb0f2261a6326dee2f3020c767b/AlphaFold2.ipynb). 

You may likely already be aware of the thes issue, but wanted to chime in just in case. I have not tried recreating the issue on the daily build of ChimeraX but am downloading it now. Thanks for all your hard work!

Best,
Barrett

Copy of my colab instance: https://colab.research.google.com/drive/1uZlvU2wVl0F0zTjxul1t6VM0WcPmDdEL?usp=sharing

Log:
UCSF ChimeraX version: 1.3 (2021-12-08)  
© 2016-2021 Regents of the University of California. All rights reserved.  
How to cite UCSF ChimeraX  

> close #

Expected a models specifier or a keyword  

> ui tool show AlphaFold

> alphafold predict
> MKDNTVPLKLIALLANGEFHSGEQLGETLGMSRAAINKHIQTLRDWGVDVFTVPGKGYSLPEPIPLLNAKQILGQLDGGSVAVLPVVDSTNQYLLDRIGELKSGDACIAEYQQAGRGSRGRKWFSPFGANLYLSMFWRLKRGPAAAIGLGPVIGIVMAEALRKLGADKVRVKWPNDLYLQDRKLAGILVELAGITGDAAQIVIGAGINVAMRRVEESVVNQGWITLQEAGINLDRNTLAATLIRELRAALELFEQEGLAPYLPRWEKLDNFINRPVKLIIGDKEIFGISRGIDKQGALLLEQDGVIKPWMGGEISLRSAEK

Running AlphaFold prediction  

Downloaded prediction file not found:
/Users/barrettdavis/Downloads/ChimeraX/AlphaFold/prediction_5/best_model.pdb  

AlphaFold prediction finished  
Results in /Users/barrettdavis/Downloads/ChimeraX/AlphaFold/prediction_5  




OpenGL version: 4.1 INTEL-18.4.6
OpenGL renderer: Intel Iris Pro OpenGL Engine
OpenGL vendor: Intel Inc.Hardware:

    Hardware Overview:

      Model Name: MacBook Pro
      Model Identifier: MacBookPro11,4
      Processor Name: Quad-Core Intel Core i7
      Processor Speed: 2.8 GHz
      Number of Processors: 1
      Total Number of Cores: 4
      L2 Cache (per Core): 256 KB
      L3 Cache: 6 MB
      Hyper-Threading Technology: Enabled
      Memory: 16 GB
      System Firmware Version: 428.60.3.0.0
      OS Loader Version: 540.80.2~11
      SMC Version (system): 2.29f24

Software:

    System Software Overview:

      System Version: macOS 12.2.1 (21D62)
      Kernel Version: Darwin 21.3.0
      Time since boot: 2 days 23:40

Graphics/Displays:

    Intel Iris Pro:

      Chipset Model: Intel Iris Pro
      Type: GPU
      Bus: Built-In
      VRAM (Dynamic, Max): 1536 MB
      Vendor: Intel
      Device ID: 0x0d26
      Revision ID: 0x0008
      Metal Family: Supported, Metal GPUFamily macOS 1
      Displays:
        Color LCD:
          Display Type: Built-In Retina LCD
          Resolution: 2880 x 1800 Retina
          Framebuffer Depth: 24-Bit Color (ARGB8888)
          Main Display: Yes
          Mirror: Off
          Online: Yes
          Automatically Adjust Brightness: No
          Connection Type: Internal
        E2411:
          Resolution: 1920 x 1080 (1080p FHD - Full High Definition)
          UI Looks like: 1920 x 1080 @ 60.00Hz
          Framebuffer Depth: 24-Bit Color (ARGB8888)
          Display Serial Number: 107NDEZ58868
          Mirror: Off
          Online: Yes
          Rotation: Supported
          Automatically Adjust Brightness: Yes
        E2411:
          Resolution: 1920 x 1080 (1080p FHD - Full High Definition)
          UI Looks like: 1920 x 1080 @ 60.00Hz
          Framebuffer Depth: 24-Bit Color (ARGB8888)
          Display Serial Number: 107NDQA58884
          Mirror: Off
          Online: Yes
          Rotation: Supported
          Automatically Adjust Brightness: Yes

Locale: (None, 'UTF-8')
PyQt5 5.15.2, Qt 5.15.2
Installed Packages:
    alabaster: 0.7.12
    appdirs: 1.4.4
    appnope: 0.1.2
    Babel: 2.9.1
    backcall: 0.2.0
    blockdiag: 2.0.1
    certifi: 2021.5.30
    cftime: 1.5.1.1
    charset-normalizer: 2.0.9
    ChimeraX-AddCharge: 1.2.2
    ChimeraX-AddH: 2.1.11
    ChimeraX-AlignmentAlgorithms: 2.0
    ChimeraX-AlignmentHdrs: 3.2
    ChimeraX-AlignmentMatrices: 2.0
    ChimeraX-Alignments: 2.2.3
    ChimeraX-AlphaFold: 1.0
    ChimeraX-AltlocExplorer: 1.0.1
    ChimeraX-AmberInfo: 1.0
    ChimeraX-Arrays: 1.0
    ChimeraX-Atomic: 1.31
    ChimeraX-AtomicLibrary: 4.2
    ChimeraX-AtomSearch: 2.0
    ChimeraX-AtomSearchLibrary: 1.0
    ChimeraX-AxesPlanes: 2.0
    ChimeraX-BasicActions: 1.1
    ChimeraX-BILD: 1.0
    ChimeraX-BlastProtein: 2.0
    ChimeraX-BondRot: 2.0
    ChimeraX-BugReporter: 1.0
    ChimeraX-BuildStructure: 2.6.1
    ChimeraX-Bumps: 1.0
    ChimeraX-BundleBuilder: 1.1
    ChimeraX-ButtonPanel: 1.0
    ChimeraX-CageBuilder: 1.0
    ChimeraX-CellPack: 1.0
    ChimeraX-Centroids: 1.2
    ChimeraX-ChemGroup: 2.0
    ChimeraX-Clashes: 2.2.2
    ChimeraX-Clipper: 0.17.0
    ChimeraX-ColorActions: 1.0
    ChimeraX-ColorGlobe: 1.0
    ChimeraX-ColorKey: 1.5
    ChimeraX-CommandLine: 1.1.5
    ChimeraX-ConnectStructure: 2.0
    ChimeraX-Contacts: 1.0
    ChimeraX-Core: 1.3
    ChimeraX-CoreFormats: 1.1
    ChimeraX-coulombic: 1.3.2
    ChimeraX-Crosslinks: 1.0
    ChimeraX-Crystal: 1.0
    ChimeraX-CrystalContacts: 1.0
    ChimeraX-DataFormats: 1.2.2
    ChimeraX-Dicom: 1.0
    ChimeraX-DistMonitor: 1.1.5
    ChimeraX-DistUI: 1.0
    ChimeraX-Dssp: 2.0
    ChimeraX-EMDB-SFF: 1.0
    ChimeraX-ExperimentalCommands: 1.0
    ChimeraX-FileHistory: 1.0
    ChimeraX-FunctionKey: 1.0
    ChimeraX-Geometry: 1.1
    ChimeraX-gltf: 1.0
    ChimeraX-Graphics: 1.1
    ChimeraX-Hbonds: 2.1.2
    ChimeraX-Help: 1.2
    ChimeraX-HKCage: 1.3
    ChimeraX-IHM: 1.1
    ChimeraX-ImageFormats: 1.2
    ChimeraX-IMOD: 1.0
    ChimeraX-IO: 1.0.1
    ChimeraX-ISOLDE: 1.3
    ChimeraX-ItemsInspection: 1.0
    ChimeraX-Label: 1.1
    ChimeraX-ListInfo: 1.1.1
    ChimeraX-Log: 1.1.4
    ChimeraX-LookingGlass: 1.1
    ChimeraX-Maestro: 1.8.1
    ChimeraX-Map: 1.1
    ChimeraX-MapData: 2.0
    ChimeraX-MapEraser: 1.0
    ChimeraX-MapFilter: 2.0
    ChimeraX-MapFit: 2.0
    ChimeraX-MapSeries: 2.1
    ChimeraX-Markers: 1.0
    ChimeraX-Mask: 1.0
    ChimeraX-MatchMaker: 2.0.4
    ChimeraX-MDcrds: 2.6
    ChimeraX-MedicalToolbar: 1.0.1
    ChimeraX-Meeting: 1.0
    ChimeraX-MLP: 1.1
    ChimeraX-mmCIF: 2.4
    ChimeraX-MMTF: 2.1
    ChimeraX-Modeller: 1.2.6
    ChimeraX-ModelPanel: 1.2.1
    ChimeraX-ModelSeries: 1.0
    ChimeraX-Mol2: 2.0
    ChimeraX-Morph: 1.0
    ChimeraX-MouseModes: 1.1
    ChimeraX-Movie: 1.0
    ChimeraX-Neuron: 1.0
    ChimeraX-Nucleotides: 2.0.2
    ChimeraX-OpenCommand: 1.7
    ChimeraX-PDB: 2.6.5
    ChimeraX-PDBBio: 1.0
    ChimeraX-PDBLibrary: 1.0.2
    ChimeraX-PDBMatrices: 1.0
    ChimeraX-PickBlobs: 1.0
    ChimeraX-Positions: 1.0
    ChimeraX-PresetMgr: 1.0.1
    ChimeraX-PubChem: 2.1
    ChimeraX-ReadPbonds: 1.0.1
    ChimeraX-Registration: 1.1
    ChimeraX-RemoteControl: 1.0
    ChimeraX-ResidueFit: 1.0
    ChimeraX-RestServer: 1.1
    ChimeraX-RNALayout: 1.0
    ChimeraX-RotamerLibMgr: 2.0.1
    ChimeraX-RotamerLibsDunbrack: 2.0
    ChimeraX-RotamerLibsDynameomics: 2.0
    ChimeraX-RotamerLibsRichardson: 2.0
    ChimeraX-SaveCommand: 1.5
    ChimeraX-SchemeMgr: 1.0
    ChimeraX-SDF: 2.0
    ChimeraX-Segger: 1.0
    ChimeraX-Segment: 1.0
    ChimeraX-SelInspector: 1.0
    ChimeraX-SeqView: 2.4.6
    ChimeraX-Shape: 1.0.1
    ChimeraX-Shell: 1.0
    ChimeraX-Shortcuts: 1.1
    ChimeraX-ShowAttr: 1.0
    ChimeraX-ShowSequences: 1.0
    ChimeraX-SideView: 1.0
    ChimeraX-Smiles: 2.1
    ChimeraX-SmoothLines: 1.0
    ChimeraX-SpaceNavigator: 1.0
    ChimeraX-StdCommands: 1.6.1
    ChimeraX-STL: 1.0
    ChimeraX-Storm: 1.0
    ChimeraX-Struts: 1.0
    ChimeraX-Surface: 1.0
    ChimeraX-SwapAA: 2.0
    ChimeraX-SwapRes: 2.1
    ChimeraX-TapeMeasure: 1.0
    ChimeraX-Test: 1.0
    ChimeraX-Toolbar: 1.1
    ChimeraX-ToolshedUtils: 1.2
    ChimeraX-Tug: 1.0
    ChimeraX-UI: 1.13.7
    ChimeraX-uniprot: 2.2
    ChimeraX-UnitCell: 1.0
    ChimeraX-ViewDockX: 1.0.1
    ChimeraX-VIPERdb: 1.0
    ChimeraX-Vive: 1.1
    ChimeraX-VolumeMenu: 1.0
    ChimeraX-VTK: 1.0
    ChimeraX-WavefrontOBJ: 1.0
    ChimeraX-WebCam: 1.0
    ChimeraX-WebServices: 1.0
    ChimeraX-Zone: 1.0
    colorama: 0.4.4
    cxservices: 1.1
    cycler: 0.11.0
    Cython: 0.29.24
    decorator: 5.1.0
    docutils: 0.17.1
    filelock: 3.0.12
    funcparserlib: 0.3.6
    grako: 3.16.5
    h5py: 3.6.0
    html2text: 2020.1.16
    idna: 3.3
    ihm: 0.21
    imagecodecs: 2021.4.28
    imagesize: 1.3.0
    ipykernel: 5.5.5
    ipython: 7.23.1
    ipython-genutils: 0.2.0
    jedi: 0.18.0
    Jinja2: 3.0.1
    jupyter-client: 6.1.12
    jupyter-core: 4.9.1
    kiwisolver: 1.3.2
    lxml: 4.6.3
    lz4: 3.1.3
    MarkupSafe: 2.0.1
    matplotlib: 3.4.3
    matplotlib-inline: 0.1.3
    msgpack: 1.0.2
    netCDF4: 1.5.7
    networkx: 2.6.3
    numexpr: 2.8.0
    numpy: 1.21.2
    openvr: 1.16.801
    packaging: 21.0
    ParmEd: 3.2.0
    parso: 0.8.3
    pexpect: 4.8.0
    pickleshare: 0.7.5
    Pillow: 8.3.2
    pip: 21.2.4
    pkginfo: 1.7.1
    prompt-toolkit: 3.0.23
    psutil: 5.8.0
    ptyprocess: 0.7.0
    pycollada: 0.7.1
    pydicom: 2.1.2
    Pygments: 2.10.0
    PyOpenGL: 3.1.5
    PyOpenGL-accelerate: 3.1.5
    pyparsing: 3.0.6
    PyQt5-commercial: 5.15.2
    PyQt5-sip: 12.8.1
    PyQtWebEngine-commercial: 5.15.2
    python-dateutil: 2.8.2
    pytz: 2021.3
    pyzmq: 22.3.0
    qtconsole: 5.1.1
    QtPy: 1.11.3
    RandomWords: 0.3.0
    requests: 2.26.0
    scipy: 1.7.1
    setuptools: 57.5.0
    sfftk-rw: 0.7.1
    six: 1.16.0
    snowballstemmer: 2.2.0
    sortedcontainers: 2.4.0
    Sphinx: 4.2.0
    sphinx-autodoc-typehints: 1.12.0
    sphinxcontrib-applehelp: 1.0.2
    sphinxcontrib-blockdiag: 2.0.0
    sphinxcontrib-devhelp: 1.0.2
    sphinxcontrib-htmlhelp: 2.0.0
    sphinxcontrib-jsmath: 1.0.1
    sphinxcontrib-qthelp: 1.0.3
    sphinxcontrib-serializinghtml: 1.1.5
    suds-jurko: 0.6
    tifffile: 2021.4.8
    tinyarray: 1.2.3
    tornado: 6.1
    traitlets: 5.1.1
    urllib3: 1.26.7
    wcwidth: 0.2.5
    webcolors: 1.11.1
    wheel: 0.37.0
    wheel-filename: 1.3.0
File attachment: model_1_error

model_1_error​

[davis797@…]

Added by email2trac

[pett]

Reported by Barrett Davis

[Tom Goddard]

Thanks! AlphaFold is really fragile. I will look into it hopefully today (preparing a talk to give on AlphaFold for tomorrow morning).

[Tom Goddard]

Fixed, I think. Still needs one more test.

A change in the default jax and jaxlib provided by Google Colab broke the ChimeraX 1.3 AlphaFold Colab script. The ChimeraX script installed a known to work jax and jaxlib version to avoid exactly this kind of problem, but AlphaFold did not use the installed jax because before the pinned jax was installed the default jax was imported to check if a GPU was available. I changed the ChimeraX colab script to not import jax to check for a GPU. After that change it used the pinned jax and jaxlib but would not use the GPU because the PyPi jaxlib does not have CUDA support. I changed the script to install jaxlib with CUDA from Google, and that got it working again.

ChimeraX daily builds use AlphaFold 2.1 which was not broken by the new default jax and jaxlib in Colab. ChimeraX 1.3 uses the older AlphaFold 2.0 which the new jax broke. But the daily builds have the same bug that the colab script imports jax to check for a gpu before installing its pinned jax and jaxlib versions. I will fix that, probably tomorrow.

[Tom Goddard]

ChiimeraX 1.3 AlphaFold prediction is fixed.

ChimeraX 1.4 daily builds continue to work but need a similar fix to pin the jax and jaxlib versions to avoid breaking during future jax and jaxlib updates.

[Tom Goddard]

I also updated the ChimeraX 1.4 daily build so it will not be broken in the future by JAX library updates.

And I updated ChimeraX 1.4 daily builds to use AlphaFold 2.2.0.