mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
287 lines
9.2 KiB
Text
287 lines
9.2 KiB
Text
/-
|
|
Selfies.lean - SELFIES String Parser
|
|
|
|
SELFIES (Self-Referencing Embedded Strings)
|
|
is a robust string representation for molecular graphs
|
|
that guarantees validity during generative modeling.
|
|
|
|
Unlike SMILES, SELFIES has a context-free grammar that
|
|
prevents invalid molecular structures during generation.
|
|
|
|
Example: "[C][=O][O]" = CO2
|
|
Example: "[C][C][O]" = ethanol
|
|
|
|
This module provides a formal parser for SELFIES in Lean.
|
|
|
|
References:
|
|
- Krenn et al. (2020): "Self-referencing embedded strings (SELFIES):
|
|
A 100% robust molecular string representation"
|
|
- GitHub: https://github.com/aspuru-guzik-group/selfies
|
|
-/
|
|
|
|
import Mathlib.Data.List.Basic
|
|
import Mathlib.Data.Char.Basic
|
|
import Mathlib.Data.String.Basic
|
|
import Smiles
|
|
|
|
namespace Selfies
|
|
|
|
-- ============================================================================
|
|
-- §1: SELFIES Grammar Types
|
|
-- ============================================================================
|
|
|
|
/-- Atom symbols in SELFIES (always bracketed) -/
|
|
inductive AtomSymbol
|
|
| C | N | O | S | P | F | Cl | Br | I | B
|
|
| Si | As | Se | Te | At | Ts | Og
|
|
deriving DecidableEq, Repr
|
|
|
|
/-- Branch symbols -/
|
|
inductive BranchSymbol
|
|
| openBranch -- '['
|
|
| closeBranch -- ']'
|
|
deriving DecidableEq, Repr
|
|
|
|
/-- Bond symbols in SELFIES -/
|
|
inductive BondSymbol
|
|
| single -- implicit
|
|
| double -- '='
|
|
| triple -- '#'
|
|
| aromatic -- ':'
|
|
deriving DecidableEq, Repr
|
|
|
|
/-- Ring closure symbols -/
|
|
inductive RingSymbol
|
|
| ring (n : Nat) -- ring closure number
|
|
deriving DecidableEq, Repr
|
|
|
|
/-- A token in SELFIES string -/
|
|
inductive Token
|
|
| atom (sym : AtomSymbol)
|
|
| branch (sym : BranchSymbol)
|
|
| bond (sym : BondSymbol)
|
|
| ring (sym : RingSymbol)
|
|
deriving DecidableEq, Repr
|
|
|
|
-- ============================================================================
|
|
-- §2: SELFIES Molecule Structure
|
|
-- ============================================================================
|
|
|
|
/-- A single atom with its properties -/
|
|
structure Atom where
|
|
symbol : AtomSymbol
|
|
chirality : Option Nat := none
|
|
hydrogenCount : Option Nat := none
|
|
charge : Option Int := none
|
|
deriving Repr
|
|
|
|
/-- A bond between atoms -/
|
|
structure Bond where
|
|
bondType : BondSymbol
|
|
deriving Repr
|
|
|
|
/-- A branch (subtree) in the molecule -/
|
|
inductive Branch
|
|
| atom (a : Atom)
|
|
| bondThenAtom (b : Bond) (a : Atom) (rest : Option Branch)
|
|
| branch (sub : Branch) (rest : Option Branch)
|
|
deriving Repr
|
|
|
|
/-- Complete SELFIES molecule -/
|
|
structure Molecule where
|
|
branches : List Branch
|
|
deriving Repr
|
|
|
|
-- ============================================================================
|
|
-- §3: Tokenizer
|
|
-- ============================================================================
|
|
|
|
/-- Parse atom symbol from string -/
|
|
def parseAtomSymbol (s : String) : Option AtomSymbol :=
|
|
match s with
|
|
| "[C]" => some .C
|
|
| "[N]" => some .N
|
|
| "[O]" => some .O
|
|
| "[S]" => some .S
|
|
| "[P]" => some .P
|
|
| "[F]" => some .F
|
|
| "[Cl]" => some .Cl
|
|
| "[Br]" => some .Br
|
|
| "[I]" => some .I
|
|
| "[B]" => some .B
|
|
| "[Si]" => some .Si
|
|
| "[As]" => some .As
|
|
| "[Se]" => some .Se
|
|
| "[Te]" => some .Te
|
|
| "[At]" => some .At
|
|
| "[Ts]" => some .Ts
|
|
| "[Og]" => some .Og
|
|
| _ => none
|
|
|
|
/-- Tokenize SELFIES string into tokens -/
|
|
partial def tokenize (input : String) : List Token :=
|
|
let rec helper (pos : Nat) (acc : List Token) : List Token :=
|
|
if pos >= input.length then acc.reverse
|
|
else
|
|
let c := input.get ⟨pos, by omega⟩
|
|
if c == '[' then
|
|
-- Try to parse atom symbol
|
|
let rec findEnd (endPos : Nat) : Nat :=
|
|
if endPos >= input.length then pos
|
|
else if input.get ⟨endPos, by omega⟩ == ']' then endPos
|
|
else findEnd (endPos + 1)
|
|
|
|
let endPos := findEnd (pos + 1)
|
|
let tokenStr := input.extract ⟨pos, by omega⟩ ⟨endPos + 1, by omega⟩
|
|
match parseAtomSymbol tokenStr with
|
|
| some sym => helper (endPos + 1) (Token.atom sym :: acc)
|
|
| none => helper (endPos + 1) acc
|
|
else if c == '=' then
|
|
helper (pos + 1) (Token.bond .double :: acc)
|
|
else if c == '#' then
|
|
helper (pos + 1) (Token.bond .triple :: acc)
|
|
else if c == ':' then
|
|
helper (pos + 1) (Token.bond .aromatic :: acc)
|
|
else if c.isDigit then
|
|
helper (pos + 1) (Token.ring (.ring (c.toNat - '0'.toNat)) :: acc)
|
|
else
|
|
helper (pos + 1) acc
|
|
|
|
helper 0 []
|
|
|
|
-- ============================================================================
|
|
-- §4: Parser
|
|
-- ============================================================================
|
|
|
|
/-- Parse tokens into molecule structure -/
|
|
partial def parseTokens (tokens : List Token) : Option Molecule :=
|
|
let rec helper (remaining : List Token) (acc : List Branch) : Option (List Branch) :=
|
|
match remaining with
|
|
| [] => some acc.reverse
|
|
| t :: ts =>
|
|
match t with
|
|
| Token.atom sym =>
|
|
let atom := Atom.mk sym none none none
|
|
helper ts (Branch.atom atom :: acc)
|
|
| Token.bond sym =>
|
|
match acc with
|
|
| [] => none -- Bond without preceding atom
|
|
| b :: rest =>
|
|
match ts with
|
|
| Token.atom atomSym :: ts' =>
|
|
let atom := Atom.mk atomSym none none none
|
|
let bond := Bond.mk sym
|
|
helper ts' (Branch.bondThenAtom bond atom (some b) :: rest)
|
|
| _ => none
|
|
| Token.ring sym =>
|
|
-- Ring closure - just skip for now (would need ring tracking)
|
|
helper ts acc
|
|
| Token.branch sym =>
|
|
-- Branch handling - skip for now
|
|
helper ts acc
|
|
|
|
match helper tokens [] with
|
|
| some branches => some ⟨branches⟩
|
|
| none => none
|
|
|
|
/-- Parse complete SELFIES string -/
|
|
def parse (input : String) : Option Molecule :=
|
|
let tokens := tokenize input
|
|
if tokens.isEmpty then none
|
|
else parseTokens tokens
|
|
|
|
/-- Check if SELFIES string is valid -/
|
|
def isValid (input : String) : Bool :=
|
|
parse input |>.isSome
|
|
|
|
-- ============================================================================
|
|
-- §5: SMILES to SELFIES Conversion
|
|
-- ============================================================================
|
|
|
|
/-- Convert SMILES Atom to SELFIES AtomSymbol -/
|
|
def smilesAtomToSelfies (atom : Smiles.Atom) : Option AtomSymbol :=
|
|
match atom with
|
|
| Smiles.Atom.organic Smiles.OrganicElement.C => some .C
|
|
| Smiles.Atom.organic Smiles.OrganicElement.N => some .N
|
|
| Smiles.Atom.organic Smiles.OrganicElement.O => some .O
|
|
| Smiles.Atom.organic Smiles.OrganicElement.S => some .S
|
|
| Smiles.Atom.organic Smiles.OrganicElement.P => some .P
|
|
| Smiles.Atom.organic Smiles.OrganicElement.F => some .F
|
|
| Smiles.Atom.organic Smiles.OrganicElement.Cl => some .Cl
|
|
| Smiles.Atom.organic Smiles.OrganicElement.Br => some .Br
|
|
| Smiles.Atom.organic Smiles.OrganicElement.I => some .I
|
|
| Smiles.Atom.organic Smiles.OrganicElement.B => some .B
|
|
| Smiles.Atom.aromatic Smiles.AromaticElement.c => some .C
|
|
| Smiles.Atom.aromatic Smiles.AromaticElement.n => some .N
|
|
| Smiles.Atom.aromatic Smiles.AromaticElement.o => some .O
|
|
| Smiles.Atom.aromatic Smiles.AromaticElement.s => some .S
|
|
| Smiles.Atom.aromatic Smiles.AromaticElement.p => some .P
|
|
| _ => none
|
|
|
|
/-- Convert SMILES Bond to SELFIES BondSymbol -/
|
|
def smilesBondToSelfies (bond : Smiles.Bond) : Option BondSymbol :=
|
|
match bond with
|
|
| Smiles.Bond.single => some .single
|
|
| Smiles.Bond.double => some .double
|
|
| Smiles.Bond.triple => some .triple
|
|
| Smiles.Bond.aromatic => some .aromatic
|
|
| _ => none -- Stereochemical bonds not in SELFIES
|
|
|
|
/-- Convert SMILES Molecule to SELFIES string (simplified) -/
|
|
def fromSmiles (smiles : String) : Option String :=
|
|
match Smiles.parse smiles with
|
|
| some mol =>
|
|
-- Simplified: just convert first chain
|
|
match mol.components with
|
|
| [Smiles.Chain.atom atom] =>
|
|
match smilesAtomToSelfies atom with
|
|
| some sym => some s!"[{sym}]"
|
|
| none => none
|
|
| _ => none -- Complex chains not implemented
|
|
| none => none
|
|
|
|
-- ============================================================================
|
|
-- §6: Properties and Theorems
|
|
-- ============================================================================
|
|
|
|
/-- Empty string is not valid SELFIES -/
|
|
theorem notValidEmpty : isValid "" = false := by
|
|
rfl
|
|
|
|
/-- Single atom "[C]" is valid -/
|
|
theorem validCarbon : isValid "[C]" = true := by
|
|
rfl
|
|
|
|
/-- Ethanol "[C][C][O]" is valid -/
|
|
theorem validEthanol : isValid "[C][C][O]" = true := by
|
|
rfl
|
|
|
|
/-- CO2 "[C][=O][O]" is valid -/
|
|
theorem validCO2 : isValid "[C][=O][O]" = true := by
|
|
rfl
|
|
|
|
-- ============================================================================
|
|
-- §7: Examples
|
|
-- ============================================================================
|
|
|
|
#eval tokenize "[C]" -- Carbon
|
|
#eval tokenize "[C][C][O]" -- Ethanol
|
|
#eval tokenize "[C][=O][O]" -- CO2
|
|
#eval tokenize "[C][#N]" -- HCN
|
|
#eval tokenize "[C][C][=C][C]" -- Butadiene
|
|
|
|
#eval parse "[C]" -- Carbon
|
|
#eval parse "[C][C][O]" -- Ethanol
|
|
#eval parse "[C][=O][O]" -- CO2
|
|
#eval parse "[C][#N]" -- HCN
|
|
|
|
#eval isValid "[C]" -- true
|
|
#eval isValid "[C][C][O]" -- true
|
|
#eval isValid "[C][=O][O]" -- true
|
|
#eval isValid "" -- false
|
|
|
|
#eval fromSmiles "C" -- "[C]"
|
|
#eval fromSmiles "CC" -- "[C]" (simplified)
|
|
#eval fromSmiles "CCO" -- "[C]" (simplified)
|
|
|
|
end Selfies
|