AF2 metrics prediction jax error

[mdamgaard-biolib]

As also reported in #13, running the af2_interface_metrics.py step in the SH3 tutorial results in the following Jax error:

TypeError: _forward_fn() takes 1 positional argument but 2 were given

Specifically, this is the code that causes it.

 prediction_result = jax.vmap(model_runner.apply, in_axes=(None,None,0,0))(model_runner.params,
          jax.random.PRNGKey(0), processed_feature_dict, processed_initial_guess_dict)

13 got closed by the issue author as removing the initial_guess logic allowed the code to run.

I opened this issue as it is my understanding that this initial guess logic is a big part of the method which was outlined in https://doi.org/10.1101/2022.06.15.495993 (N. Bennet et al 2022).

Do you know of a way to fix this error or is it still valuable to run it without the initial_guess? Seems to be the last step i need before being able to run through the tutorial in its entirety :)

[jongseo-park]

I totally agree with your opinion that the processed_initial_guess_dict is one of the core logic of the af2_interface_metrics.py

But intriguingly, in my case (https://github.com/RosettaCommons/RFDesign/issues/13),

the initial_guess seems working fine even after remove that option.

To validate it, first, I performed a binder design task based on a small and flexible target (receptor) protein that is too small to properly predicted its structure even in the AlphaFold2 or ColabFold.

And using the af2_interface_metrics.py, finally I could obtain a PDB file comprising an appropriately folded receptor and a designed small binder protein.

When I tried to model this structure using AlphaFold2-Multimer (on local) as well as ColabFold with or w/o MSA, It wasn't folded properly.

Based on it, I concluded that the modified script works fine, but I still doubt whether the removal of processed_initial_guess_dict is a right way or not.

[mdamgaard-biolib]

@jongseo-park Thanks for the additional info!

Interesting that it folds without the initial guess, though not completely unexpected as it still passes the template through the processed_feature_dict.

Maybe @jueseph can weigh in on this? Otherwise we might have to wait for publishing of the Bennett paper, which also is fair :)

[mdamgaard-biolib]

I have some time to dig into this and find the root cause. @jueseph Do you know if there was AF2 code changes made to enable this initial_guess logic or is it more likely that it is a dependency problem (Jax + Jaxlib version etc.)?

[jueseph]

I've asked Nate Bennett about this (who wrote af2_interface_metrics.py) and will get back to you soon.

On Fri, Sep 16, 2022 at 2:51 AM mdamgaard-biolib @.***> wrote:

I am willing to dig into this and find the root cause. @jueseph https://github.com/jueseph Do you know if there was AF2 code changes made to enable this initial_guess logic or is it more likely that it is a dependency problem (Jax + Jaxlib version etc.)?

— Reply to this email directly, view it on GitHub https://github.com/RosettaCommons/RFDesign/issues/25#issuecomment-1249164462, or unsubscribe https://github.com/notifications/unsubscribe-auth/AABCU5ZM4YKTMB3LSXDPLGLV6Q7IRANCNFSM6AAAAAAQKPMIHA . You are receiving this because you were mentioned.Message ID: @.***>

[jueseph]

Ok, looks like Nate Bennett actually customized AF2 a little in order to perform the initial guess prediction. He will upload his code for his paper soon, at which point I'll update this repo with code or instructions to make this script work.

For now, either use jongseo's solution of removing the initial guess input or just skip using this script.

Sorry for the confusion and inconvenience.

On Fri, Sep 16, 2022 at 6:38 AM Jue Wang @.***> wrote:

I've asked Nate Bennett about this (who wrote af2_interface_metrics.py) and will get back to you soon.

On Fri, Sep 16, 2022 at 2:51 AM mdamgaard-biolib @.***> wrote:

I am willing to dig into this and find the root cause. @jueseph https://github.com/jueseph Do you know if there was AF2 code changes made to enable this initial_guess logic or is it more likely that it is a dependency problem (Jax + Jaxlib version etc.)?

— Reply to this email directly, view it on GitHub https://github.com/RosettaCommons/RFDesign/issues/25#issuecomment-1249164462, or unsubscribe https://github.com/notifications/unsubscribe-auth/AABCU5ZM4YKTMB3LSXDPLGLV6Q7IRANCNFSM6AAAAAAQKPMIHA . You are receiving this because you were mentioned.Message ID: @.***>

[mdamgaard-biolib]

Okay that makes sense. Thanks for the great help as always Jue!

Closing the issue :)