Recipe: Fast Relax 2.0

created by Artoria2e5

Profile

• Name: Fast Relax 2.0
• ID: 108397
• Shared with: Public
• Parent: Fast Relax 1.0
• Children:
  • Bravo MFR v3.0
• Created on: September 07, 2023 at 07:25 AM UTC
• Updated on: May 09, 2024 at 03:40 AM UTC
• Description:

  The classic Fast Relax script, now using "MonomerRelax2019". One day I might even add all four variants...

Code

--[[
        FastRelax v2

        Based on the updated Rosetta relax recipe, "MonomerRelax2019"

            Maguire JB, Haddox HK, Strickland D, Halabiya SF, Coventry B, et al. (December 2020).
            "Perturbing the energy landscape for improved packing during computational protein design.".
            Proteins 89(4):436-449. PMID 33249652.

            https://europepmc.org/article/PMC/8299543

        Converted from v1 Lua by Rav3n_pl's v1 to v2 converter, FolditLuaConverter.

        Fast Relax 1.0 - LociOiling -- 2022/09/17
        + based on the classic PNAS paper,  "Algorithm discovery by protein folding game players"
        + converted using FolditLuaConverter
        + fixed #NEED EDIT! lines (only two in this case)
        + minimal testing

        Fast Relax 2.0 -- Artoria2e5 -- 2023/09/07
        + Use the new algo from https://www.rosettacommons.org/docs/latest/application_documentation/structure_prediction/relax#description-of-algorithm
        + Implement wiggle-with-tolerance, as opposed to fixed iterations
        + (TODO) implement all variants

]]
-- aliases
P = print
function CI(x)
    behavior.SetClashImportance(x)
    cival = x
end

S = structure.ShakeSidechainsSelected -- TODO allow mutate-repack
WA = structure.WiggleAll

-- options; TODO: make GUI maybe, now that there's no in-GUI editor to easily change these values
energy = true --set true to seek energy in exploration puzzles
minPpi = 2    --minimum points gain per loop
shakes = 1    --number of shakes
wiggles = 5   --number of wiggles

-- band stuff. the original has a coord_cst_weight, which is like a bis to starting pos.
-- we don't have that, so turn it off for now.
usebands = false --do we mess with bands
-- p("If you want us to scale band strengths, select them.")
-- bands = dialog.SelectBands()
bands = {}

-- TODO allow change of schedule
CI_SCHEDULES = {
    ["MonomerRelax2019"] = { 0.040, 0.051, 0.265, 0.280, 0.559, 0.581, 1 },
    ["MonomerDesign2019"] = { 0.059, 0.092, 0.280, 0.323, 0.568, 0.633, 1 },
    ["InterfaceRelax2019"] = { 0.069, 0.080, 0.288, 0.302, 0.573, 0.594, 1 },
    ["InterfaceDesign2019"] = { 0.079, 0.100, 0.295, 0.323, 0.577, 0.619, 1 },
    ["Legacy"] = { 0.02, 0.02, 0.25, 0.25, 0.55, 0.55, 1 },
}
CI_SCHEDULE = CI_SCHEDULES["MonomerRelax2019"]

has_density = 0
--[[
for _, part in ipairs(puzzle.GetPuzzleSubscoreNames()) do
    if part == "Density" then
        has_density = 1
        P("Found density")
    end
end
]]

function BandScale(multiplier)
    for ii = 1, #bands do
        band.SetStrength(ii, band.GetStrength(ii) * multiplier)
    end
end

function Score()
    if energy == true then
        return current.GetEnergyScore()
    else
        return current.GetScore()
    end
end

function round(x) --cut all afer 3-rd place
    return x - x % 0.001
end

-- Minimizer
-- crude approximation of rosetta minimize "relative tolerance".
-- wiggle until the (adjusted) score improves by less than
-- RelTol2AbsTol(newscore, tol).
--
-- I know the real minimizer knows better, but I can't exactly ask it, can I?
function Minimize(tol)
    -- always do something
    local oldscore = AdjScore()
    WA(wiggles)
    local newscore = AdjScore()
    local abstol = RelTol2AbsTol(newscore, tol)

    -- now loop
    local iter = 1
    while math.abs(newscore - oldscore) > abstol and iter < 20 do
        oldscore = newscore
        WA(math.floor(wiggles))
        newscore = AdjScore()
        abstol = RelTol2AbsTol(newscore, tol)
        iter = iter + 1
    end
end

-- just vibes. 
function RelTol2AbsTol(score, tol)
    -- why 30? because each residue eventually should get about -3 REU. Yes, it's vibes.
    local expectedScore =  structure.GetCount() * (30 + 20 * has_density)
    return expectedScore * tol
end

function ClashScore()
    local sum = 0
    local segs = structure.GetCount()

    for i = 1, segs do
        sum = sum + current.GetSegmentEnergySubscore(i, "Clashing")
    end

    return sum
end

-- Re-weight clashing (fa_rep) according to clashing importance (CIval), so
-- it follows the wiggle gradient.
-- This recipe only changes CI, so it only adjusts for CI. Change accordingly
-- if other scoreparts are reweighted.
function AdjScore()
    local ret = Score() - ClashScore() * (1 - cival)
    -- for debug
    P("AdjScore (CI ", cival, "): ", ret)
    return ret
end

function PrintScores()
    P("Score: ", Score())
    P("ClashScore: ", ClashScore())
    P("AdjScore (CI " .. tostring(cival) .. "): ", AdjScore())
end

function SaveBest()
    local g = Score() - bestScore
    if g > 0 then
        bestScore = Score()
        save.Quicksave(3)
    else
        save.Quickload(3)
    end
end

function FastRelax()
    bestScore = Score()
    save.Quicksave(3)
    P("Starting FastRelax, Score: ", round(bestScore))
    selection.SelectAll()
    recentbest.Save()
    repeat
        local ss = Score()
        if usebands then
            band.EnableAll()
        end
        CI(CI_SCHEDULE[1])
        S(shakes)
        CI(CI_SCHEDULE[2])
        Minimize(0.01)

        if usebands then
            BandScale(0.5) -- coord_cst_weight 0.5
        end
        CI(CI_SCHEDULE[3])
        S(shakes)
        CI(CI_SCHEDULE[4])
        Minimize(0.01)

        if usebands then
            BandScale(2)      -- restore from 0.5 for coord_cst_weight 1.0
            band.DisableAll() -- coord_cst_weight 0.0
        end
        CI(CI_SCHEDULE[5])
        S(shakes)
        CI(CI_SCHEDULE[6])
        Minimize(0.01)

        CI(CI_SCHEDULE[7])
        S(shakes)
        Minimize(0.0001)
        recentbest.Restore()
        SaveBest()
        local g = Score() - ss
        P("Loop gain: ", round(g))
    until g < minPpi
    P("End score: ", round(Score()))
end

-- main call
FastRelax()

--end of script

Comments

[Artoria2e5]

None of the TODOs are done, but it works well enough. Might as well publish.