Deep Chemistry: Protein Folding and Molecular Modeling in Wolfram Language

Siria Sadeddin and Jason Biggs 
Wolfram Research

#WolframTechConf

Outline

From Sequence to Structure									Why protein folding matters and what it tell us about life.

The Story of Protein Structure									First discoveries: myoglobin and the importance of structure.

The AlphaFold Revolution									How AI transformed biology.

From AlphaFold to ESMFold									Faster, simpler, sequence-only AI folding.

Bringing ESMFold into Wolfram via ONNX							Importing the model into Wolfram Language

Discovering Patterns Hidden in Protein Sequences						Embeddings and clustering biomolecules.

The ESMFold Pipeline in Wolfram Language							From sequence to 3D structure in a few lines of code

ESMFold with BioMoleculePlot3D								Predict the biomolecule structures with our built-in function

From Sequence to Structure

The Journey of Protein Folding and AI

From Sequence to Structure

Proteins fold because parts of the chain attract or repel each other.

The Story of Protein Structure: from crystals to AI

Understanding life’s architecture, one fold at a time.

First protein structure: Myoglobin (Kendrew)

Decades of discovery and simulation.

Today AI predicts structure in seconds.

The AlphaFold Revolution

Average GDT tests of the contest winners. (image source: Nature news, https://www.nature.com/articles/d41586-020-03348-4)

The AlphaFold Revolution

AI cracked the folding code
What decades of physics couldn’t solve, AI achieved in months!

From AlphaFold to ESMFold

ESMFold proved folding can be learned purely from sequence — no alignments needed.
What if you could predict structure… from sequence alone?

ESMFold combines language modeling (like GPT for proteins) with a structure prediction head.

Pure transformer architecture — no Evoformer or external libraries

No evolutionary search → no MSA generation (~90% of AlphaFold’s cost)

Runs in seconds, not minutes/hours

Bringing ESMFold into Wolfram via ONNX

Making ESMFold part of the Wolfram ecosystem

esmLanguage = Import["esm_lm.onnx", "NetExternalObject"]

foldStruct = Import["structure_module.onnx", "NetExternalObject"];

process = Import["post_processing.onnx", "NetExternalObject"];

Discovering Patterns Hidden in Protein Sequences

Proteins with similar functions or structures group naturally in the model’s latent space

enzymes=Map[Style[#,Red]&,{"ACGYLKTPKLADPPVLRGDSSVTKAICKPDPVLEK","GVALDECKALDYLPGKPLPMDGKVCQCGSKTPLRP","VLPGYTCGELDCKPGKPLPKCGADKTQVATPFLRG","TCGALVQYPSCADPPVLRGSDSSVKACKKLDPQDK","GALCEECKLCPGADYKPMDGDRLPAAATSKTRPVG","PAVDCKKALVYLPKPLPMDGKVCRGSKTPKTRPYG","VLGYTCGALDCKPGKPLPKCGADKTQVATPFLRGA","CGALVQYPSCADPPVLRGSDSSVKACKKLDPQDKT","ALCEECKLCPGADYKPMDGDRLPAAATSKTRPVGK","AVDCKKALVYLPKPLPMDGKVCRGSKTPKTRPYGR"}];
structuralProteins=Map[Style[#,Green]&,{"VGKGFRYGSSQKRYLHCQKSALPPSCRRGKGQGSAT","KDPTVMTVGTYSCQCPKQDSRGSVQPTSRVKTSRSK","PLVGKACGRSSDYKCPGQMVSGGSKQTPASQRPSYD","CGKKLVGYPSSKADVPLQGRSSFSPKACKKDPQMTS","RKGVASLYCSSKLSCKAQYSKGMSDGRSPKASSTTS","RPKSAASCEQAKSYRSLSLPSMKGKVPSKCSRSKRP","RSDVSYTSCSQSKDCKPSKPPKMSGSKDSSTVATPS","LSTCSKKVAYPSSKADPPSSGRSSFSMKACKKQDPPV","RVGSASSEPKSSCSVQSYSKPSMSGDSSPKASSTSK","QPSASNCEKMSSYRPSLPSMSKGVPSSRSKSSPPYQ"}];

sequences=Join[enzymes,structuralProteins];

extractor=Map[Max,Transpose@esmLanguage[encodeSequence[#][[1]]]["s_s_0"][[1]]]&;

FeatureSpacePlot[sequences,FeatureExtractor->extractor,LabelingSize->90,LabelingFunction->Callout]

The ESMFold Pipeline in Wolfram Language

testSequence={"MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHF","VHLTPEEKSAVTALWGKVNVDEVGGEALGRLLVVYPWTQRFFESFG"};

{aatype,residx,linkerMask,chainIndex}=encodeSequence[testSequence];

esm=esmLanguage[aatype];

(*Initialize the mask and residx*)
{sS0,sZ0}={esm@"s_s_0",esm@"s_z_0"};
{b,l}=Dimensions@aatype;
{mask,residx}={ConstantArray[1,Dimensions@aatype],{Range[l]-1}};
(*Initialize recycling tensors*)
{recycleS,recycleZ,recycleBins}={ConstantArray[0,Dimensions@sS0],ConstantArray[0,Dimensions@sZ0],ConstantArray[0,{1,l,l}]};

structure=foldStruct[];

output=process[];

pdbStr=outputToPDB[output];

StringTake[pdbStr,960]

bioMolecule=ImportString[pdbStr,{"PDB","BioMolecule"},"PredictedStructure"->True]

BioMoleculePlot3D@bioMolecule

ESMFold with BioMoleculePlot3D

https://reference.wolfram.com/language/ref/BioMoleculePlot3D.html

https://resources.wolframcloud.com/NeuralNetRepository/resources/ESMFold-V1-Trained-on-UR50-Data/

To see how we can fold sequences in top-level Wolfram Language, let’s begin with a protein of known structure. The N-terminal domain of the amyloid precursor protein (APP), which is implicated in Alzheimer’s disease, is a good example.  We can retrieve the citation title for the published data using an external identifier:

biomol=BioMolecule[(ExternalIdentifierURL[ExternalIdentifier["PDBStructureID","1MWP"], "PDB structure", "1MWP", https://rdf.wwpdb.org/pdb/1MWP])]

We can fold the structure using the local model:

folded=BioMolecule[biomol["BioSequences"],BioMoleculeFoldingMethod->"Local"]

Then, after aligning the predicted structure to the experimental structure, we compare them in 3D:

BioMoleculePlot3D[{biomol,BioMoleculeAlign[biomol,folded]},PlotLegends->{"experimental","predicted"}]

Thank  you

Utilities

encodeSequence[bs_BioSequence,rest___]:=encodeSequence[bs["SequenceString"],rest];
encodeSequence[bs_BioMolecule,rest___]:=encodeSequence[Flatten@bs["BioSequence"]["SequenceString"],rest];
encodeSequence[seq_,residueIndexOffset_:512,chainLinker_:StringRepeat["G",25]]:=Module[{chains,joined,encoded,residx,linkerMask,chainIndex,offset=0,restypesWithX={"A","R","N","D","C","Q","E","G","H","I","L","K","M","F","P","S","T","W","Y","V","X"},restypeOrderWithX,chainCount,linkerLen,end},restypeOrderWithX=AssociationThread[restypesWithX->Range[1,Length[restypesWithX]]];
(*Setup*)chains=If[ListQ@seq,seq,{seq}];
chainCount=Length[chains];
linkerLen=StringLength[chainLinker];
joined=StringJoin@Riffle[chains,If[chainCount>1,chainLinker,""]];
(*Encode amino acids*)encoded=Lookup[restypeOrderWithX,Characters[joined],restypeOrderWithX["X"]]-1;
residx=Range[Length[encoded]]-1;
(*Apply offsets only if more than 1 chain*)If[residueIndexOffset>0&&chainCount>1,offset=0;
Do[Module[{start=offset+1,len=StringLength[chains[[i]]]},end=start+len-1;
residx[[start;;end]]+=(i-1)*residueIndexOffset;
offset=end;
If[i<chainCount,(*Advance offset over linker*)offset+=linkerLen;];],{i,chainCount}];];
(*Linker mask:0 where linker,1 elsewhere*)linkerMask=ConstantArray[1.0,Length[encoded]];
If[chainCount>1,offset=0;
Do[offset+=StringLength[chains[[i]]];
If[i<chainCount,linkerMask[[offset+1;;offset+linkerLen]]=0;
offset+=linkerLen;];,{i,chainCount}];];
(*Chain index*)chainIndex=Flatten[Table[Join[ConstantArray[i-1,StringLength[chains[[i]]]],If[i<chainCount,ConstantArray[i-1,linkerLen],{}]],{i,chainCount}]];
Transpose[{{encoded,residx,linkerMask,chainIndex}}]];

(*Convert atom14 positions to atom37 using indices and mask*)
atom14ToAtom37[atom14_,residxAtom37ToAtom14_,atom37Exists_]:=Module[{residues,indices,atom37Data},
indices=residxAtom37ToAtom14+1;
residues=Length@atom14;
atom37Data=ConstantArray[0.,{residues,37,3}];
(*Iterate over all batches,residues,and atom37 positions*)
Do[Do[
If[Greater[indices[[r,a]],0],atom37Data[[r,a]]=atom14[[r,indices[[r,a]]]]];,{a,1,37}],{r,1,residues}
];
(*Apply existence mask*)
atom37Data
];
formatAtomRecord[atomNumber_,atomName_,residueName_,chainID_,residueNumber_,{x_,y_,z_},occupancy_,bFactor_,element_]:=StringTake[
StringJoin[
"ATOM  ",StringPadLeft[ToString@atomNumber,5]," ",StringPadRight[atomName,4]," ",residueName," ",chainID,StringPadLeft[ToString@residueNumber,4],"    ",StringPadLeft[ToString@Round[x,0.001],8],StringPadLeft[ToString@Round[y,0.001],8],StringPadLeft[ToString@Round[z,0.001],8],StringPadLeft[ToString@Round[occupancy,0.01],6],StringPadLeft[ToString@Round[bFactor,0.01],6],"          ",StringPadRight[element,2]],UpTo@80
]
outputToPDB[output_]:=Module[{aatype,residueIndex,bFactors,chainIndices,sampleAatype,sampleResIndex,pdbs,samplePositions,sampleResidxAtom37To14,sampleAtom37Exists,atom37Positions,proteinData},aatype=output@"aatype";
residueIndex=output@"residue_index";
bFactors=output@"plddt";
chainIndices=If[KeyExistsQ[output,"chain_index"],output@"chain_index",None];
pdbs=Table[sampleAatype=aatype[[i]];
sampleResIndex=residueIndex[[i]]+1;
samplePositions=Part[Part[output@"positions",-1],i];
sampleResidxAtom37To14=Part[output@"residx_atom37_to_atom14",i];
sampleAtom37Exists=Part[output@"atom37_atom_exists",i];
atom37Positions=atom14ToAtom37[samplePositions,sampleResidxAtom37To14,sampleAtom37Exists];
proteinData=<|"aatype"->sampleAatype,"atomPositions"->atom37Positions,"atomMask"->sampleAtom37Exists,"residueIndex"->sampleResIndex,"bFactors"->bFactors[[i]],"chainIndex"->If[UnsameQ[chainIndices,None],chainIndices[[i]],None]|>;
toPDB@proteinData,{i,Length@aatype}];
Replace[pdbs,{s_String}:>s]
];
toPDB[proteinData_]:=Module[{pos,bFactor,element,atomName,atomLine,chainIndices,aatype,atomPositions,atomMask,residueIndex,bFactors,chainIndex,pdbLines,chainMapping,currentChain,currentResidue,atomIndex=1,atomTypes=,PDBCHAINIDS=,restype1to3=},
{aatype,atomPositions,atomMask,residueIndex,bFactors,chainIndex}=proteinData/@{"aatype","atomPositions","atomMask","residueIndex","bFactors","chainIndex"};
chainIndices=If[UnsameQ[chainIndex,None],chainIndex,Table[0,{Length@aatype}]];
chainMapping=AssociationThread[Union@chainIndices,PDBCHAINIDS[[Span[1,Length@Union@chainIndices]]]];
pdbLines={"MODEL     1"};
Do[currentChain=chainMapping[chainIndices[[i]]];
currentResidue=restype1to3[Part[Keys@restype1to3,aatype[[i]]+1]];
Do[atomName=atomTypes[[j]];
If[Less[atomMask[[i,j]],0.5],Continue[]];
(*Skip masked atoms*)
pos=atomPositions[[i,j]];
bFactor=bFactors[[i,j]];
element=First@Characters@atomName;
atomLine=formatAtomRecord[atomIndex,atomName,currentResidue,currentChain,residueIndex[[i]],pos,1.,bFactor,element];
AppendTo[pdbLines,atomLine];
Increment@atomIndex;,{j,Length@atomTypes}];,{i,Length@aatype}];
(*Add TER record after each residue*)AppendTo[pdbLines,StringJoin["TER   ",StringPadLeft[ToString@atomIndex,5," "],"      ",currentResidue," ",currentChain,"  ",ToString[residueIndex[[-1]]]]];
AppendTo[pdbLines,"END"];
StringRiffle[pdbLines,"\n"]]