update to boltz-2, MolecularInput -> BoltzInput

Author: AntonOresten

.github/workflows/CI.yml

@@ -37,6 +37,12 @@ jobs:
       - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v5
+        with:
+          files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: false
   docs:
     name: Documentation
     runs-on: ubuntu-latest

CondaPkg.toml

@@ -1,3 +1,3 @@
 
 [pip.deps]
-boltz = "~=2.0.0"
+boltz = "~=2.1.1"

Project.toml

@@ -19,24 +19,11 @@ BioStructuresExt = "BioStructures"
 ProteinChainsExt = "ProteinChains"
 
 [compat]
-BioAlignments = "3"
-BioSequences = "3"
 BioStructures = "4"
 Compat = "4.16"
 CondaPkg = "0.2"
 ProteinChains = "0.7"
 PythonCall = "0.9"
 Scratch = "1"
-TMscore = "0.1"
 YAML = "0.4"
 julia = "1.10"
-
-[extras]
-BioAlignments = "00701ae9-d1dc-5365-b64a-a3a3ebf5695e"
-BioSequences = "7e6ae17a-c86d-528c-b3b9-7f778a29fe59"
-BioStructures = "de9282ab-8554-53be-b2d6-f6c222edabfc"
-TMscore = "6e960f36-814a-4ceb-8c87-6ab5da0e9bff"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["Test", "BioAlignments", "BioSequences", "BioStructures", "TMscore"]

README.md

@@ -20,7 +20,7 @@ pkg"add PyBoltz"
 ```julia
 using PyBoltz, PyBoltz.Schema
 
-input = MolecularInput(
+input = BoltzInput(
     sequences = [
         protein(id="A", sequence="TTCCPSIVARSNFNVCRLPGTPEAICATYTGCIIIPGATCPGDYAN", msa="empty"),
     ]

docs/src/api.md

@@ -9,12 +9,12 @@ predict
 ## `Schema` submodule
 
 ```@docs
-PyBoltz.Schema.MolecularInput
+PyBoltz.Schema.BoltzInput
 ```
 
 ### Sequences
 
-The following sequence types go into the `sequences` vector keyword argument of `MolecularInput`.
+The following sequence types go into the `sequences` vector keyword argument of `BoltzInput`.
 
 ```@docs
 PyBoltz.Schema.protein
@@ -25,7 +25,7 @@ PyBoltz.Schema.ligand
 
 ### Constraints
 
-The following constraint types go into the `constraints` vector keyword argument of `MolecularInput`.
+The following constraint types go into the `constraints` vector keyword argument of `BoltzInput`.
 
 ```@docs
 PyBoltz.Schema.bond

ext/BioStructuresExt.jl

@@ -13,17 +13,16 @@ end
 function PyBoltz.predict(input, ::Type{MolecularStructure}; options...)
     mktempdir() do out_dir
         predict(input;
-            _prefix_index=(input isa AbstractVector{PyBoltz.Schema.MolecularInput}),
             out_dir,
             output_format="mmcif",
             options...)
         prediction_paths = readdir(joinpath(out_dir, only(readdir(out_dir)), "predictions"); join=true)
         prediction_names = basename.(prediction_paths)
 
         local results
-        if all(name -> startswith(name, PyBoltz.PYBOLTZ_INPUT_INDEX_PREFIX), prediction_names)
+        if get(task_local_storage(), "pyboltz_remember_ordering", false)
             # output vector needs to match input vector (with possible missing values)
-            @assert input isa AbstractVector{PyBoltz.Schema.MolecularInput}
+            @assert input isa AbstractVector{PyBoltz.Schema.BoltzInput}
             results = Union{MolecularStructure,Missing}[fill(missing, length(input))...]
             for prediction_name in prediction_names
                 index, name = split(split(prediction_name, PyBoltz.PYBOLTZ_INPUT_INDEX_PREFIX, limit=2)[2], "_", limit=2)
@@ -33,7 +32,7 @@ function PyBoltz.predict(input, ::Type{MolecularStructure}; options...)
                 try
                     results[idx] = read_boltz_cif(cif_path, name)
                 catch e
-                    @warn e
+                    @error name
                     results[idx] = missing
                 end
             end
@@ -44,16 +43,25 @@ function PyBoltz.predict(input, ::Type{MolecularStructure}; options...)
                 try
                     push!(results, read_boltz_cif(cif_path, basename(prediction_path)))
                 catch e
-                    @warn e
+                    @error name
                     push!(results, missing)
                 end
             end
         end
+        if any(ismissing, results)
+            @error "$(count(ismissing, results)) out of $(length(results)) boltz predictions errored."
+        end
         return results
     end
 end
 
-function PyBoltz.predict(input::PyBoltz.Schema.MolecularInput, ::Type{MolecularStructure}; options...)
+function PyBoltz.predict(input::AbstractVector{PyBoltz.Schema.BoltzInput}, ::Type{MolecularStructure}; options...)
+    task_local_storage("pyboltz_remember_ordering", true) do
+        @invoke predict(input::Any, MolecularStructure; options...)
+    end
+end
+
+function PyBoltz.predict(input::PyBoltz.Schema.BoltzInput, ::Type{MolecularStructure}; options...)
     return PyBoltz.predict([input], MolecularStructure; options...) |> only
 end
 

ext/ProteinChainsExt.jl

@@ -11,7 +11,7 @@ function PyBoltz.predict(input, ::Type{ProteinStructure}; options...)
     return PyBoltz.predict(input, MolecularStructure; options...) .|> to_structure
 end
 
-function PyBoltz.predict(input::PyBoltz.Schema.MolecularInput, ::Type{ProteinStructure}; options...)
+function PyBoltz.predict(input::PyBoltz.Schema.BoltzInput, ::Type{ProteinStructure}; options...)
     return PyBoltz.predict(input, MolecularStructure; options...) |> to_structure
 end
 

src/Schema.jl

@@ -1,18 +1,18 @@
 module Schema
 
-export MolecularInput
+export BoltzInput
 export protein, dna, rna, ligand
 export bond, pocket
 
 const BOLTZ_SCHEMA_VERSION = 1
 
 """
-    MolecularInput
+    BoltzInput
 
 A dictionary object that can be written to a YAML file.
 
 Implemented according to the schema definition in the
-[boltz documentation](https://github.com/jwohlwend/boltz/blob/a9b3abc2c1f90f26b373dd1bcb7afb5a3cb40293/docs/prediction.md),
+[boltz documentation](https://github.com/jwohlwend/boltz/blob/744b4aecb6b5e847a25692ced07c328e7995ee33/docs/prediction.md),
 allowing for easy in-memory construction of the schema.
 
 # Additions
@@ -28,7 +28,7 @@ allowing for easy in-memory construction of the schema.
 ```julia
 using PyBoltz.Schema
 
-input1 = MolecularInput(
+input1 = BoltzInput(
     name = "example1", # optional name YAML file (and thus output pdb/cif file)
     sequences = [
         protein(
@@ -47,7 +47,7 @@ input1 = MolecularInput(
     ]
 )
 
-input2 = MolecularInput(
+input2 = BoltzInput(
     sequences = [
         protein(
             id = ["A1"],
@@ -67,25 +67,33 @@ input2 = MolecularInput(
 )
 ```
 """
-struct MolecularInput <: AbstractDict{String,Any}
+struct BoltzInput <: AbstractDict{String,Any}
     dict::Dict{String,Any}
 end
 
-function MolecularInput(;
+function BoltzInput(;
     sequences,
     constraints = nothing,
+    templates = nothing,
+    properties = nothing,
     name = nothing
 )
     dict = Dict{String,Any}("version" => BOLTZ_SCHEMA_VERSION, "sequences" => sequences)
     !isnothing(constraints) && (dict["constraints"] = constraints)
+    !isnothing(templates) && (dict["templates"] = templates)
+    !isnothing(properties) && (dict["properties"] = properties)
     !isnothing(name) && (dict["name"] = name)
-    return MolecularInput(dict)
+    return BoltzInput(dict)
 end
 
-Base.length(input::MolecularInput) = length(input.dict)
-Base.iterate(input::MolecularInput, args...) = iterate(input.dict, args...)
-Base.getindex(input::MolecularInput, key::AbstractString) = input.dict[key]
-Base.get(input::MolecularInput, key::AbstractString, default) = get(input.dict, key, default)
+Base.length(input::BoltzInput) = length(input.dict)
+Base.iterate(input::BoltzInput, args...) = iterate(input.dict, args...)
+Base.getindex(input::BoltzInput, key::AbstractString) = input.dict[key]
+Base.get(input::BoltzInput, key::AbstractString, default) = get(input.dict, key, default)
+
+# deprecated
+const MolecularInput = BoltzInput
+export MolecularInput
 
 ## sequences
 
@@ -202,20 +210,66 @@ function bond(;
 end
 
 """
-    pocket(; binder, contacts)
+    pocket(; binder, contacts, max_distance=nothing)
 
 ```julia
 using PyBoltz.Schema: pocket
 # binder is a chain_id
 # contacts is a vector of vectors of (chain_id, residue_index)
-pocket(binder="A", contacts=[["B", 1], ["C", 2]])
+pocket(binder="A", contacts=[("B", 1), ("C", 2)])
 ```
 """
 function pocket(;
     binder::AbstractString,
-    contacts::AbstractVector{<:Tuple{AbstractString,Integer}}
+    contacts::AbstractVector{<:Tuple{AbstractString,Any}},
+    max_distance::Union{Real,Nothing} = nothing,
+)
+    dict = Dict("pocket" => Dict{String,Any}(
+        "binder" => binder,
+        "contacts" => [[c...] for c in contacts]))
+    !isnothing(max_distance) && (dict["max_distance"] = max_distance)
+    return dict
+end
+
+"""
+    contact(; token1, token2, max_distance=nothing)
+"""
+function contact(;
+    token1::Tuple{AbstractString,Any},
+    token2::Tuple{AbstractString,Any},
+    max_distance::Union{Real,Nothing} = nothing,
+)
+    dict = Dict("contact" => Dict(
+        "token1" => token1,
+        "token2" => token2))
+    !isnothing(max_distance) && (dict["max_distance"] = max_distance)
+    return dict
+end
+
+## templates
+
+"""
+    template(; cif, chain_id=nothing, template_id=nothing)
+"""
+function template(;
+    cif::AbstractString,
+    chain_id::Union{AbstractString,AbstractVector{<:AbstractString},Nothing} = nothing,
+    template_id::Union{AbstractString,AbstractVector{<:AbstractString},Nothing} = nothing,
+)
+    dict = Dict{String,Any}("cif" => cif)
+    !isnothing(chain_id) && (dict["chain_id"] = chain_id)
+    !isnothing(template_id) && (dict["template_id"] = template_id)
+end
+
+## properties
+
+"""
+    affinity(; binder)
+"""
+function affinity(;
+    binder::AbstractString
 )
-    return Dict("pocket" => Dict{String,Any}("binder" => binder, "contacts" => [[c...] for c in contacts]))
+    return Dict("affinity" => Dict("binder" => binder))
 end
 
 end

src/predict.jl

@@ -6,8 +6,8 @@ Run Boltz-1 prediction with the given input, output type, and options.
 # Input types
 
 - `AbstractString`: Path to a FASTA/YAML file or directory (for batching).
-- `MolecularInput`: A single [`PyBoltz.Schema.MolecularInput`](@ref) object.
-- `Vector{MolecularInput}`: A vector of [`PyBoltz.Schema.MolecularInput`](@ref) objects for batching.
+- `BoltzInput`: A single [`PyBoltz.Schema.BoltzInput`](@ref) object.
+- `Vector{BoltzInput}`: A vector of [`PyBoltz.Schema.BoltzInput`](@ref) objects for batching.
 
 # Output types
 
@@ -16,8 +16,8 @@ By default, raw results will be written to disk in the `out_dir` directory (see
 For convenience, `output_type` can be provided as a second argument to reduce manual file I/O.
 
 If `output_type` is provided, the function will return a
-single object if a `MolecularInput` was provided as input,
-otherwise a vector if an `AbstractString` or `Vector{MolecularInput}` was provided.
+single object if a `BoltzInput` was provided as input,
+otherwise a vector if an `AbstractString` or `Vector{BoltzInput}` was provided.
 
 The following output types are supported:
 
@@ -46,35 +46,28 @@ Defaults to a Scratch.jl-backed directory created at module init; call `clear_ca
 - `msa_pairing_strategy::String`: 'greedy' or 'complete'; requires `use_msa_server=true`.
 
 ## Boolean Flags
-- `verbose::Bool`: Whether to print boltz logs to stdout. Default: true.
 - `write_full_pae::Bool`: Dump PAE to a npz file. Default: true.
 - `write_full_pde::Bool`: Dump PDE to a npz file. Default: false.
 - `override::Bool`: Override existing predictions. Default: false.
 - `use_msa_server::Bool`: Use MMSeqs2 server for MSA generation. Default: false.
 """
-function predict(input_path::AbstractString; verbose=true, _prefix_index=false, options...)
-    @assert !_prefix_index "`_prefix_index` is reserved for internal use"
-    cmd = predict_cmd(input_path; options...)
-    if verbose
-        run(cmd)
-    else
-        read(cmd, String)
-    end
+function predict(input_path::AbstractString; options...)
+    run(predict_cmd(input_path; options...))
     return nothing
 end
 
 
 const PYBOLTZ_INPUT_INDEX_PREFIX = "__pyboltz_index_"
 
-function predict(inputs::AbstractVector{Schema.MolecularInput}; _prefix_index=false, options...)
+function predict(inputs::AbstractVector{Schema.BoltzInput}; options...)
     mktempdir() do dir
         input_dir = joinpath(dir, "inputs")
         mkdir(input_dir)
         msa_dir = joinpath(dir, "msas")
         mkdir(msa_dir)
         for (i, input) in enumerate(inputs)
-            name = get(input, "name", "noname")
-            prefix = _prefix_index ? "$(PYBOLTZ_INPUT_INDEX_PREFIX)$(i)_" : ""
+            name = get(input, "name", "pyboltz_structure_$i")
+            prefix = get(task_local_storage(), "pyboltz_remember_ordering", false) ? "$(PYBOLTZ_INPUT_INDEX_PREFIX)$(i)_" : ""
             path = joinpath(input_dir, "$prefix$name.yaml")
             YAML.write_file(path, MSAs_to_files!(deepcopy(input), msa_dir; prefix=prefix))
         end
@@ -83,4 +76,4 @@ function predict(inputs::AbstractVector{Schema.MolecularInput}; _prefix_index=fa
     return nothing
 end
 
-predict(input::Schema.MolecularInput; options...) = predict([input]; options...)
+predict(input::Schema.BoltzInput; options...) = predict([input]; options...)

src/utils.jl

@@ -25,7 +25,7 @@ function write_alignment(path::AbstractString, alignment::Vector{<:AbstractStrin
     end
 end
 
-function MSAs_to_files!(input::Schema.MolecularInput, dir::AbstractString; prefix="")
+function MSAs_to_files!(input::Schema.BoltzInput, dir::AbstractString; prefix="")
     for (i, type_dict) in enumerate(input["sequences"])
         type, seq_dict = only(type_dict)
         type == "protein" || continue # only proteins have msas

test/Project.toml

@@ -0,0 +1,7 @@
+[deps]
+BioAlignments = "00701ae9-d1dc-5365-b64a-a3a3ebf5695e"
+BioSequences = "7e6ae17a-c86d-528c-b3b9-7f778a29fe59"
+BioStructures = "de9282ab-8554-53be-b2d6-f6c222edabfc"
+PyBoltz = "0bca4874-a2b8-4dca-bc03-d2d86e344f1f"
+TMscore = "6e960f36-814a-4ceb-8c87-6ab5da0e9bff"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"