Merge pull request #5 from MurrellGroup/consistent-order

fix: consistent ordering

Author: AntonOresten

Project.toml

@@ -1,7 +1,7 @@
 name = "PyBoltz"
 uuid = "0bca4874-a2b8-4dca-bc03-d2d86e344f1f"
 authors = ["Anton Oresten and Aron Stålmarck"]
-version = "0.1.1"
+version = "0.1.2"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -27,16 +27,16 @@ CondaPkg = "0.2"
 ProteinChains = "0.7"
 PythonCall = "0.9"
 Scratch = "1"
-TMscore_jll = "1.0.0"
+TMscore = "0.1.0"
 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_jll = "1d8dad79-2f3d-554c-bf22-543753cb6ff4"
+TMscore = "6e960f36-814a-4ceb-8c87-6ab5da0e9bff"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "BioAlignments", "BioSequences", "BioStructures", "TMscore_jll"]
+test = ["Test", "BioAlignments", "BioSequences", "BioStructures", "TMscore"]

ext/BioStructuresExt.jl

@@ -19,24 +19,37 @@ function PyBoltz.predict(input, ::Type{MolecularStructure}; options...)
             options...)
         prediction_paths = readdir(joinpath(out_dir, only(readdir(out_dir)), "predictions"); join=true)
         prediction_names = basename.(prediction_paths)
-        structures = MolecularStructure[]
-        perm, structure_names = if all(name -> startswith(name, PyBoltz.PYBOLTZ_INPUT_INDEX_PREFIX), prediction_names)
-            indices = Int[]
-            structure_names = String[]
+
+        local results
+        if all(name -> startswith(name, PyBoltz.PYBOLTZ_INPUT_INDEX_PREFIX), prediction_names)
+            # output vector needs to match input vector (with possible missing values)
+            @assert input isa AbstractVector{PyBoltz.Schema.MolecularInput}
+            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)
-                push!(indices, parse(Int, index))
-                push!(structure_names, name)
+                idx = parse(Int, index)
+                prediction_path = joinpath(joinpath(out_dir, only(readdir(out_dir)), "predictions"), prediction_name)
+                cif_path = joinpath(prediction_path, basename(prediction_path)*"_model_0.cif")
+                try
+                    results[idx] = read_boltz_cif(cif_path, name)
+                catch e
+                    @warn e
+                    results[idx] = missing
+                end
             end
-            sortperm(indices), structure_names
         else
-            collect(1:length(prediction_names)), prediction_names
-        end
-        for (structure_name, prediction_path) in zip(structure_names, prediction_paths)
-            cif_path = joinpath(prediction_path, basename(prediction_path)*"_model_0.cif")
-            push!(structures, read_boltz_cif(cif_path, structure_name))
+            results = Union{MolecularStructure,Missing}[]
+            for prediction_path in prediction_paths
+                cif_path = joinpath(prediction_path, basename(prediction_path)*"_model_0.cif")
+                try
+                    push!(results, read_boltz_cif(cif_path, basename(prediction_path)))
+                catch e
+                    @warn e
+                    push!(results, missing)
+                end
+            end
         end
-        return structures[perm]
+        return results
     end
 end
 

test/TMscore.jl

@@ -5,8 +5,7 @@ using BioAlignments
 using BioSequences
 using BioStructures
 
-include("TMscore.jl") # remove this once TMscore.jl is registered: https://github.com/JuliaRegistries/General/pull/130867
-using .TMscore
+using TMscore
 
 # ENV["PyBoltz_TEST_ACCELERATOR"] = "gpu"