Rename package to ESMFold

Author: murrellb

.github/workflows/CI.yml

@@ -61,6 +61,6 @@ jobs:
         shell: julia --project=docs --color=yes {0}
         run: |
           using Documenter: DocMeta, doctest
-          using ESMEmbed
-          DocMeta.setdocmeta!(ESMEmbed, :DocTestSetup, :(using ESMEmbed); recursive=true)
-          doctest(ESMEmbed)
+          using ESMFold
+          DocMeta.setdocmeta!(ESMFold, :DocTestSetup, :(using ESMFold); recursive=true)
+          doctest(ESMFold)

Project.toml

@@ -1,5 +1,5 @@
-name = "ESMEmbed"
-uuid = "92fcdd77-1b6e-4cb4-a5ce-af12488bb353"
+name = "ESMFold"
+uuid = "992ade7d-a23b-44e1-bfa9-f18d3a6e7567"
 authors = ["Ben Murrell"]
 version = "1.0.0-DEV"
 

README.md

@@ -4,12 +4,10 @@ A Julia port of the **full ESMFold model**: ESM2 embeddings + folding trunk + st
 This repo runs end‑to‑end folding on CPU, and will run on GPU when you move the model/tensors
 to the GPU.
 
-Note: the module name is currently `ESMEmbed` for compatibility.
-
 ## Quickstart (single sequence)
 
 ```julia
-using ESMEmbed
+using ESMFold
 
 # Download weights from Hugging Face and build the full folding model
 model = load_ESMFold()
@@ -25,7 +23,7 @@ println(pdb)
 ## Batch Folding
 
 ```julia
-using ESMEmbed
+using ESMFold
 
 model = load_ESMFold()
 seqs = ["ELLKKLLEELKG", "ACDEFGHIKLMNPQRSTVWY"]

docs/Project.toml

@@ -1,3 +1,3 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
-ESMEmbed = "92fcdd77-1b6e-4cb4-a5ce-af12488bb353"
+ESMFold = "992ade7d-a23b-44e1-bfa9-f18d3a6e7567"

docs/make.jl

@@ -1,14 +1,14 @@
-using ESMEmbed
+using ESMFold
 using Documenter
 
-DocMeta.setdocmeta!(ESMEmbed, :DocTestSetup, :(using ESMEmbed); recursive=true)
+DocMeta.setdocmeta!(ESMFold, :DocTestSetup, :(using ESMFold); recursive=true)
 
 makedocs(;
-    modules=[ESMEmbed],
+    modules=[ESMFold],
     authors="Ben Murrell",
-    sitename="ESMEmbed.jl",
+    sitename="ESMFold.jl",
     format=Documenter.HTML(;
-        canonical="https://MurrellGroup.github.io/ESMEmbed.jl",
+        canonical="https://MurrellGroup.github.io/ESMFold.jl",
         edit_link="main",
         assets=String[],
     ),
@@ -18,6 +18,6 @@ makedocs(;
 )
 
 deploydocs(;
-    repo="github.com/MurrellGroup/ESMEmbed.jl",
+    repo="github.com/MurrellGroup/ESMFold.jl",
     devbranch="main",
 )

docs/src/index.md

@@ -1,42 +1,41 @@
 ```@meta
-CurrentModule = ESMEmbed
+CurrentModule = ESMFold
 ```
 
-# ESMEmbed
+# ESMFold
 
-A lightweight Julia port of the ESMFold sequence embedding stack.
+A Julia port of the **full ESMFold model**: ESM2 embeddings + folding trunk + structure module.
 
 ## Quickstart
 
 ```julia
-using ESMEmbed
+using ESMFold
 
-model = load_ESM()
-emb = model("ACDEFGHIK")
+model = load_ESMFold()
+output = infer(model, "ACDEFGHIK")
+pdb = output_to_pdb(output)[1]
 ```
 
 ## Outputs
 
-The embedding output is returned in **C × L × B** order (Julia‑native layout):
+`infer` returns a dictionary with structure + confidence outputs, including:
 
-- `C` = embedding width (`c_s`, typically 384)
-- `L` = sequence length (after padding)
-- `B` = batch size
+- `positions`, `frames`, `angles`, `states`
+- `plddt`, `mean_plddt`, `ptm`, and `predicted_aligned_error`
 
-If you request pair features with `return_pair=true` and `use_esm_attn_map=true`, the
-pair tensor is returned as **C_z × L × L × B**. Otherwise `pair = nothing`.
+Use `output_to_pdb` to export PDBs.
 
 ## Input Modes
 
 - `AbstractMatrix{Int}` shaped `(B, L)`
 - `Vector{Vector{Int}}` (auto‑padded)
 - `Vector{String}` or a single `String`
 
-See the README for more usage examples.
+See the README for more usage examples and batch folding.
 
 ```@index
 ```
 
 ```@autodocs
-Modules = [ESMEmbed]
+Modules = [ESMFold]
 ```

scripts/test.jl

@@ -1,7 +1,7 @@
 using Pkg
 Pkg.activate(joinpath(@__DIR__, ".."); io=devnull)
 
-using ESMEmbed
+using ESMFold
 
 seq = "ELLKKLLEELKG"
 expected_path = joinpath(@__DIR__, "output_ELLKKLLEELKG.pdb")
@@ -10,7 +10,7 @@ generated_path = joinpath(@__DIR__, "output_ELLKKLLEELKG.generated.pdb")
 function _load_model()
     weights_path = joinpath(@__DIR__, "..", "weights", "esm.safetensors")
     if isfile(weights_path)
-        reader = ESMEmbed.SafeTensors.Reader(weights_path)
+        reader = ESMFold.SafeTensors.Reader(weights_path)
         (
             num_layers,
             embed_dim,
@@ -22,18 +22,18 @@ function _load_model()
             position_bins,
             num_blocks,
             lddt_head_hid_dim,
-        ) = ESMEmbed._infer_esmfold_full_config(reader)
+        ) = ESMFold._infer_esmfold_full_config(reader)
 
-        alphabet = ESMEmbed.Alphabet_from_architecture("ESM-1b")
-        esm = ESMEmbed.ESM2(
+        alphabet = ESMFold.Alphabet_from_architecture("ESM-1b")
+        esm = ESMFold.ESM2(
             num_layers,
             embed_dim,
             attention_heads;
             alphabet = alphabet,
             token_dropout = true,
         )
 
-        trunk_cfg = ESMEmbed.FoldingTrunkConfig(
+        trunk_cfg = ESMFold.FoldingTrunkConfig(
             num_blocks,
             c_s,
             c_z,
@@ -45,21 +45,21 @@ function _load_model()
             false,
             4,
             nothing,
-            ESMEmbed.StructureModuleConfig(),
+            ESMFold.StructureModuleConfig(),
         )
 
-        cfg = ESMEmbed.ESMFoldConfig(; trunk=trunk_cfg, lddt_head_hid_dim=lddt_head_hid_dim, use_esm_attn_map=false)
-        model = ESMEmbed.ESMFold(esm; cfg=cfg)
-        ESMEmbed.load_esmfold_safetensors!(model, reader)
+        cfg = ESMFold.ESMFoldConfig(; trunk=trunk_cfg, lddt_head_hid_dim=lddt_head_hid_dim, use_esm_attn_map=false)
+        model = ESMFold.ESMFold(esm; cfg=cfg)
+        ESMFold.load_esmfold_safetensors!(model, reader)
         return model
     end
 
-    return ESMEmbed.load_ESMFold()
+    return ESMFold.load_ESMFold()
 end
 
 model = _load_model()
-output = ESMEmbed.infer(model, seq)
-pdb = ESMEmbed.output_to_pdb(output)[1]
+output = ESMFold.infer(model, seq)
+pdb = ESMFold.output_to_pdb(output)[1]
 open(generated_path, "w") do io
     write(io, pdb)
 end
@@ -81,7 +81,7 @@ if expected != pdb
     error("PDB does not match expected output.")
 end
 
-metrics = ESMEmbed.confidence_metrics(output)
+metrics = ESMFold.confidence_metrics(output)
 println("PDB matches: ", expected_path)
 println("mean_plddt: ", metrics.mean_plddt)
 println("ptm: ", metrics.ptm)

src/ESMEmbed.jl src/ESMFold.jlsrc/ESMEmbed.jl renamed to src/ESMFold.jl

@@ -1,4 +1,4 @@
-module ESMEmbed
+module ESMFold
 
 using LinearAlgebra
 using Statistics

test/runtests.jl

@@ -1,5 +1,5 @@
 using Test
-using ESMEmbed
+using ESMFold
 
 @testset "Alphabet" begin
     alphabet = Alphabet_from_architecture("ESM-1b")
@@ -15,7 +15,7 @@ end
 
 @testset "Padding helper" begin
     seqs = [[0, 1, 2], [3]]
-    aa, mask = ESMEmbed._pad_batch(seqs; pad_value = 0)
+    aa, mask = ESMFold._pad_batch(seqs; pad_value = 0)
     @test size(aa) == (2, 3)
     @test size(mask) == (2, 3)
     @test aa[2, 2] == 0