mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
314 lines
12 KiB
Text
314 lines
12 KiB
Text
import Mathlib.Data.Real.Basic
|
||
import Mathlib.Data.List.Basic
|
||
|
||
/-!
|
||
PeptideMoE.lean — Mixture-of-Experts for Peptide Conformational Analysis
|
||
|
||
This module formalizes a Mixture-of-Experts (MoE) system for peptide
|
||
conformational state analysis, incorporating thermodynamic parameters,
|
||
admissibility constraints, and expert coordination.
|
||
|
||
Purpose: Mathematical formalization of peptide conformational search
|
||
using expert gating, thermodynamic scoring, and constraint-based filtering.
|
||
|
||
Key structures:
|
||
- PeptideState: Conformational state (φ, ψ angles, energies)
|
||
- Expert: MoE expert with gating and advice functions
|
||
- AdmissibilityParams: Steric/bond/angle constraints
|
||
- ThermoParams: Temperature and Boltzmann constant
|
||
|
||
The module provides functions for:
|
||
- Free energy computation
|
||
- Admissibility checking
|
||
- Score filtering and penalization
|
||
- Expert usefulness evaluation
|
||
- Candidate selection and reporting
|
||
-/
|
||
|
||
namespace PeptideMoE
|
||
|
||
/-- Peptide conformational state with Ramachandran angles and energies -/
|
||
structure PeptideState where
|
||
phi : ℝ
|
||
psi : ℝ
|
||
internalEnergy : ℝ
|
||
conformationalEntropy : ℝ
|
||
structuralCoherence : ℝ
|
||
stericEnergy : ℝ
|
||
bondEnergy : ℝ
|
||
|
||
/-- MoE expert with gating function and advice for φ/ψ angles -/
|
||
structure Expert where
|
||
name : String
|
||
gate : PeptideState → ℝ
|
||
advicePhi : PeptideState → ℝ
|
||
advicePsi : PeptideState → ℝ
|
||
|
||
/-- Candidate peptide conformation with label -/
|
||
structure Candidate where
|
||
state : PeptideState
|
||
label : String
|
||
|
||
/-- Admissibility parameters for conformational constraints -/
|
||
structure AdmissibilityParams where
|
||
stericMax : ℝ
|
||
bondMax : ℝ
|
||
phiMin : ℝ
|
||
phiMax : ℝ
|
||
psiMin : ℝ
|
||
psiMax : ℝ
|
||
c0 : ℝ
|
||
|
||
/-- Thermodynamic parameters for free energy computation -/
|
||
structure ThermoParams where
|
||
kB : ℝ
|
||
temperature : ℝ
|
||
|
||
/-- Learning parameters for gate weight updates -/
|
||
structure LearningParams where
|
||
learningRate : ℝ
|
||
updateSignal : PeptideState → Expert → ℝ
|
||
previousEfficiency : ℝ -- Track previous efficiency for ΔΦ computation
|
||
|
||
/-- Free energy: E + kB·T·S -/
|
||
noncomputable def freeEnergy (tp : ThermoParams) (P : PeptideState) : ℝ :=
|
||
P.internalEnergy + tp.kB * tp.temperature * P.conformationalEntropy
|
||
|
||
/-- Cost function: C(x) - measures computational or thermodynamic cost -/
|
||
noncomputable def costFunction (tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState) : ℝ :=
|
||
freeEnergy tp P + ap.c0
|
||
|
||
/-- Utility function: U(x) - measures structural coherence or benefit -/
|
||
noncomputable def utilityFunction (P : PeptideState) : ℝ :=
|
||
P.structuralCoherence
|
||
|
||
/-- Efficiency metric: Φ(x) = C(x) / U(x) - cost/utility ratio -/
|
||
noncomputable def efficiency (tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState) : ℝ :=
|
||
costFunction tp ap P / utilityFunction P
|
||
|
||
/-- φ-peptide score: structural coherence / (free energy + c0) - legacy name for efficiency -/
|
||
noncomputable def phiPeptide (tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState) : ℝ :=
|
||
efficiency tp ap P
|
||
|
||
/-- Admissibility predicate: steric, bond, and angle constraints -/
|
||
def admissible (ap : AdmissibilityParams) (P : PeptideState) : Prop :=
|
||
ap.phiMin ≤ P.phi ∧ P.phi ≤ ap.phiMax ∧
|
||
ap.psiMin ≤ P.psi ∧ P.psi ≤ ap.psiMax ∧
|
||
P.stericEnergy < ap.stericMax ∧
|
||
P.bondEnergy < ap.bondMax
|
||
|
||
noncomputable instance decidableAdmissible (ap : AdmissibilityParams) (P : PeptideState) : Decidable (admissible ap P) :=
|
||
inferInstanceAs (Decidable (ap.phiMin ≤ P.phi ∧ P.phi ≤ ap.phiMax ∧
|
||
ap.psiMin ≤ P.psi ∧ P.psi ≤ ap.psiMax ∧
|
||
P.stericEnergy < ap.stericMax ∧
|
||
P.bondEnergy < ap.bondMax))
|
||
|
||
/-- Admissibility indicator: 1 if admissible, 0 otherwise -/
|
||
noncomputable def admissibilityIndicator (ap : AdmissibilityParams) (P : PeptideState) : ℝ :=
|
||
if admissible ap P then 1 else 0
|
||
|
||
/-- Filtered score: zero if not admissible, otherwise φ-peptide score -/
|
||
noncomputable def filteredScore (tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState) : ℝ :=
|
||
admissibilityIndicator ap P * phiPeptide tp ap P
|
||
|
||
/-- Penalized score: subtract penalty if not admissible -/
|
||
noncomputable def penalizedScore (tp : ThermoParams) (ap : AdmissibilityParams) (penalty : ℝ) (P : PeptideState) : ℝ :=
|
||
phiPeptide tp ap P - (if admissible ap P then 0 else penalty)
|
||
|
||
/-- Expert usefulness: negative of gate-weighted advice alignment with gradient -/
|
||
def expertUsefulness
|
||
(gradPhi gradPsi : PeptideState → ℝ)
|
||
(E : Expert)
|
||
(P : PeptideState) : ℝ :=
|
||
-(E.gate P) * ((E.advicePhi P) * (gradPhi P) + (E.advicePsi P) * (gradPsi P))
|
||
|
||
/-- Expert helpful predicate: usefulness is non-negative -/
|
||
def expertHelpful
|
||
(gradPhi gradPsi : PeptideState → ℝ)
|
||
(E : Expert)
|
||
(P : PeptideState) : Prop :=
|
||
0 ≤ expertUsefulness gradPhi gradPsi E P
|
||
|
||
/-- MoE drift: sum of gate-weighted advice across all experts -/
|
||
def moeDrift (experts : List Expert) (P : PeptideState) : ℝ × ℝ :=
|
||
let dphi := List.sum (experts.map fun E => E.gate P * E.advicePhi P)
|
||
let dpsi := List.sum (experts.map fun E => E.gate P * E.advicePsi P)
|
||
(dphi, dpsi)
|
||
|
||
/-- Gates normalized: sum to 1, all non-negative -/
|
||
def gatesNormalized (experts : List Expert) (P : PeptideState) : Prop :=
|
||
0 ≤ List.sum (experts.map fun E => E.gate P) ∧
|
||
List.sum (experts.map fun E => E.gate P) = 1 ∧
|
||
∀ E ∈ experts, 0 ≤ E.gate P
|
||
|
||
/-- Gate weight update: z_k' = z_k + α ΔΦ · U_k(t) where ΔΦ = Φ(x) - Φ_prev -/
|
||
noncomputable def gateUpdate
|
||
(lp : LearningParams)
|
||
(tp : ThermoParams)
|
||
(ap : AdmissibilityParams)
|
||
(E : Expert)
|
||
(P : PeptideState) : ℝ :=
|
||
let currentEfficiency := efficiency tp ap P
|
||
let deltaEfficiency := currentEfficiency - lp.previousEfficiency
|
||
E.gate P + lp.learningRate * deltaEfficiency * lp.updateSignal P E
|
||
|
||
/-- Temporal transformation T(P_t, z_t) = (∂t/∂Θ_t, z_{k(t+1)}) -/
|
||
structure TemporalTransformation where
|
||
driftPhi : ℝ
|
||
driftPsi : ℝ
|
||
updatedGates : List Expert
|
||
|
||
/-- Apply temporal transformation: compute drift and update all gate weights -/
|
||
noncomputable def applyTemporalTransformation
|
||
(lp : LearningParams)
|
||
(tp : ThermoParams)
|
||
(ap : AdmissibilityParams)
|
||
(experts : List Expert)
|
||
(P : PeptideState) : TemporalTransformation :=
|
||
let drift := moeDrift experts P
|
||
let updatedExperts := experts.map fun E =>
|
||
{ name := E.name
|
||
, gate := fun _ => gateUpdate lp tp ap E P
|
||
, advicePhi := E.advicePhi
|
||
, advicePsi := E.advicePsi }
|
||
{ driftPhi := drift.1
|
||
, driftPsi := drift.2
|
||
, updatedGates := updatedExperts }
|
||
|
||
/-- Best candidate: fold-based selection maximizing filtered score among admissible candidates -/
|
||
noncomputable def bestCandidate?
|
||
(tp : ThermoParams)
|
||
(ap : AdmissibilityParams)
|
||
(cands : List Candidate) : Option Candidate :=
|
||
cands.foldl
|
||
(fun best cand =>
|
||
match best with
|
||
| none =>
|
||
if admissible ap cand.state then some cand else none
|
||
| some b =>
|
||
if admissible ap cand.state ∧
|
||
filteredScore tp ap b.state < filteredScore tp ap cand.state
|
||
then some cand
|
||
else some b)
|
||
none
|
||
|
||
/-- Candidate report: label, free energy, φ-peptide score, filtered score -/
|
||
noncomputable def candidateReport
|
||
(tp : ThermoParams)
|
||
(ap : AdmissibilityParams)
|
||
(cands : List Candidate) : List (String × ℝ × ℝ × ℝ) :=
|
||
cands.map fun c =>
|
||
( c.label
|
||
, freeEnergy tp c.state
|
||
, phiPeptide tp ap c.state
|
||
, filteredScore tp ap c.state
|
||
)
|
||
|
||
/-- Denominator safe: free energy + c0 is positive -/
|
||
def denominatorSafe (tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState) : Prop :=
|
||
0 < freeEnergy tp P + ap.c0
|
||
|
||
/-- All denominators safe: holds for all candidates -/
|
||
def allDenominatorsSafe (tp : ThermoParams) (ap : AdmissibilityParams) (cands : List Candidate) : Prop :=
|
||
∀ c ∈ cands, denominatorSafe tp ap c.state
|
||
|
||
/-- Theorem: filtered score is zero when not admissible -/
|
||
theorem filteredScore_of_not_admissible
|
||
(tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState)
|
||
(h : ¬ admissible ap P) :
|
||
filteredScore tp ap P = 0 := by
|
||
unfold filteredScore admissibilityIndicator
|
||
split
|
||
· contradiction
|
||
· simp
|
||
|
||
/-- Theorem: filtered score equals φ-peptide score when admissible -/
|
||
theorem filteredScore_of_admissible
|
||
(tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState)
|
||
(h : admissible ap P) :
|
||
filteredScore tp ap P = phiPeptide tp ap P := by
|
||
unfold filteredScore admissibilityIndicator
|
||
split
|
||
· simp
|
||
· contradiction
|
||
|
||
/-- Theorem: expert helpful iff usefulness is non-negative (reflexive) -/
|
||
theorem expertHelpful_iff
|
||
(gradPhi gradPsi : PeptideState → ℝ)
|
||
(E : Expert) (P : PeptideState) :
|
||
expertHelpful gradPhi gradPsi E P ↔ 0 ≤ expertUsefulness gradPhi gradPsi E P := by
|
||
rfl
|
||
|
||
/-- Theorem: gate mass is one when gates are normalized -/
|
||
theorem gate_mass_one
|
||
(experts : List Expert) (P : PeptideState)
|
||
(h : gatesNormalized experts P) :
|
||
List.sum (experts.map fun E => E.gate P) = 1 := by
|
||
exact h.2.1
|
||
|
||
/-
|
||
Intended invariant:
|
||
Any candidate returned by `bestCandidate?` should be admissible.
|
||
This is an external correctness property of the fold-based selection.
|
||
-/
|
||
structure BestCandidateAdmissibleHypothesis where
|
||
property (tp : ThermoParams) (ap : AdmissibilityParams) (cands : List Candidate) (c : Candidate) :
|
||
bestCandidate? tp ap cands = some c → admissible ap c.state
|
||
|
||
/-
|
||
Transformation T(P_t) properties:
|
||
|
||
The transformation T(P_t) = (∂t/∂Θ_t, Φ_filtered[P_t]) should preserve
|
||
key invariants of the peptide-MoE system.
|
||
-/
|
||
|
||
/-- Theorem: filtered score is zero when not admissible -/
|
||
theorem filteredScore_zero_of_not_admissible
|
||
(tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState)
|
||
(h : ¬ admissible ap P) :
|
||
filteredScore tp ap P = 0 := by
|
||
unfold filteredScore admissibilityIndicator
|
||
split
|
||
· contradiction
|
||
· simp
|
||
|
||
/-- Theorem: filtered score equals φ_peptide when admissible -/
|
||
theorem filteredScore_eq_phiPeptide_of_admissible
|
||
(tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState)
|
||
(h : admissible ap P) :
|
||
filteredScore tp ap P = phiPeptide tp ap P := by
|
||
unfold filteredScore admissibilityIndicator
|
||
split
|
||
· simp
|
||
· contradiction
|
||
|
||
/-- Hypothesis: filtered score is bounded when structural coherence is bounded -/
|
||
structure FilteredScoreBoundedHypothesis where
|
||
property (tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState)
|
||
(h : 0 ≤ P.structuralCoherence) (hdenom : 0 < freeEnergy tp P + ap.c0) :
|
||
0 ≤ filteredScore tp ap P ∧ filteredScore tp ap P ≤ P.structuralCoherence
|
||
|
||
/-- Hypothesis: φ_peptide is positive when denominator is safe and structural coherence is positive -/
|
||
structure PhiPeptidePosHypothesis where
|
||
property (tp : ThermoParams) (ap : AdmissibilityParams) (P : PeptideState)
|
||
(h : 0 < P.structuralCoherence) (hdenom : 0 < freeEnergy tp P + ap.c0) :
|
||
0 < phiPeptide tp ap P
|
||
|
||
/-- Hypothesis: MoE drift is bounded when expert advice is bounded -/
|
||
structure MoEDriftBoundedHypothesis where
|
||
property (B : ℝ) (experts : List Expert) (P : PeptideState)
|
||
(hgate : gatesNormalized experts P)
|
||
(hbound : ∀ E ∈ experts, |E.advicePhi P| ≤ B ∧ |E.advicePsi P| ≤ B) :
|
||
|(moeDrift experts P).1| ≤ B ∧ |(moeDrift experts P).2| ≤ B
|
||
|
||
/-- Hypothesis: MoE drift preserves angle bounds when gates are normalized -/
|
||
structure MoEDriftPreservesBoundsHypothesis where
|
||
property (experts : List Expert) (ap : AdmissibilityParams) (P : PeptideState)
|
||
(hgate : gatesNormalized experts P) (h : admissible ap P)
|
||
(hbound : ∀ E ∈ experts, |E.advicePhi P| ≤ 1 ∧ |E.advicePsi P| ≤ 1) :
|
||
ap.phiMin ≤ P.phi + (moeDrift experts P).1 ∧
|
||
P.phi + (moeDrift experts P).1 ≤ ap.phiMax ∧
|
||
ap.psiMin ≤ P.psi + (moeDrift experts P).2 ∧
|
||
P.psi + (moeDrift experts P).2 ≤ ap.psiMax
|
||
|
||
end PeptideMoE
|