Files

demian3b 9e45a090af Add initial implementation of Riemannian Flow GNN with essential data structures, simulation logic, and visualization utilities

2026-04-15 23:19:53 +09:00

15 KiB

Raw Blame History

Burn 기반 Pose Manifold Graph 구현 AI 전달 문서

0. 문서 목적

이 문서는 단일 ligand pose overfitting 실험을 위해, AI 에이전트가 구현해야 할 범위를 명확히 정의한 handoff 문서이다.

구현 대상은 다음 세 가지다.

Burn 규격에 맞는 manifold tensor 표현 및 graph 구현체
graph 입력을 받아 시각화 가능한 구조 표현과 렌더링 파이프라인
time parameter t ∈ [0, 1]에 따라 translation / rotation / torsion이 변하는 simulator

이번 단계의 목표는 연구용 완성품이 아니라, 발표 및 초기 검증을 위한 작동 가능한 prototype이다.

1. 문제 설정

1.1 상태 공간

Ligand pose state는 아래 product space로 둔다.

translation: R^3
rotation: SO(3)
torsion: T^m where m = number of rotatable bonds

즉 하나의 pose state는 개념적으로 아래와 같다.

x = (p, R, theta)

p     : global translation, shape [3]
R     : global rotation
theta : torsion angles for rotatable bonds, shape [m]

1.2 그래프 입력

입력 ligand는 graph로 표현한다.

node = atom
edge = bond
torsion-available edge = 회전 가능한 bond

이번 단계에서는 protein / pocket context는 제외한다. 오직 ligand 단일 graph와 target pose만 다룬다.

1.3 실험 목적

주어진 ligand graph 하나에 대해
임의의 초기 pose에서 출발해
하나의 fixed target pose로 수렴하는 trajectory를 생성하고
이를 시각화 가능하게 만든다.

학습 로직 전체를 이 문서에서 구현하라는 뜻은 아니다. 이번 구현의 중심은 data representation + simulator + visualization이다.

2. 구현 범위 요약

AI 에이전트는 아래 항목을 구현한다.

필수 구현

Burn-friendly tensor container
ligand graph data structure
torsion metadata extraction/representation
pose state container
explicit simulator for t=0..1
3D/2D visualization export utility

이번 단계에서 제외

full training loop
optimizer / scheduler
dataset loader for many molecules
protein-ligand interaction scoring
equivariant GNN model 자체

3. 설계 원칙

3.1 Burn 친화적 설계

구현체는 Burn 모델과 쉽게 연결될 수 있어야 한다. 따라서 다음 원칙을 지킨다.

tensor payload는 Burn tensor로 쉽게 변환 가능해야 한다.
구조체는 geometry / graph / rendering 책임을 분리한다.
manifold-specific logic는 순수 수학 유틸로 분리한다.
시각화는 모델과 독립적인 모듈로 둔다.

3.2 API 변동에 덜 민감한 설계

Burn 버전 변화에 덜 흔들리도록, 내부 표현과 backend 종속성을 너무 깊게 섞지 않는다.

권장 방식:

domain object는 plain Rust struct로 유지
필요 시 to_tensor::<B: Backend>(&device) 같은 adapter 제공
학습용 Autodiff backend와 추론용 backend를 쉽게 교체할 수 있게 설계

3.3 발표용 prototype 우선

이번 구현은 성능 최적화보다 아래가 중요하다.

구조가 명확할 것
simulator가 deterministic하게 잘 동작할 것
geometry가 눈으로 드러날 것

4. 구현해야 할 핵심 데이터 구조

4.1 Atom / Bond / LigandGraph

Atom

최소 필드:

pub struct Atom {
    pub index: usize,
    pub atomic_number: u8,
    pub formal_charge: i8,
    pub is_aromatic: bool,
}

Bond

최소 필드:

pub struct Bond {
    pub index: usize,
    pub src: usize,
    pub dst: usize,
    pub bond_order: u8,
    pub is_aromatic: bool,
    pub is_rotatable: bool,
}

TorsionEdge

회전 가능한 bond에 대해, 어떤 원자 집합이 어느 쪽 fragment로 회전하는지 알아야 한다. 이 metadata가 simulator에서 중요하다.

pub struct TorsionEdge {
    pub bond_index: usize,
    pub atom_left: usize,
    pub atom_right: usize,
    pub rotating_side: Vec<usize>,
}

설명:

atom_left, atom_right는 torsion axis를 이루는 bond의 두 원자
rotating_side는 회전 적용 대상 atom index 목록

LigandGraph

pub struct LigandGraph {
    pub atoms: Vec<Atom>,
    pub bonds: Vec<Bond>,
    pub edge_index: Vec<(usize, usize)>,
    pub torsion_edges: Vec<TorsionEdge>,
}

요구사항:

undirected bond graph를 기본으로 하되
GNN 입력용으로는 directed edge pair로 변환 가능해야 함
torsion edge list를 별도로 유지해야 함

4.2 PoseState

pose state는 ligand의 geometry 상태를 담는다.

권장 표현

pub struct PoseState {
    pub translation: [f32; 3],
    pub rotation_quat_xyzw: [f32; 4],
    pub torsions: Vec<f32>,
}

주의:

quaternion은 normalize 보장 필요
torsion은 radian 단위
torsion 범위는 기본적으로 [-pi, pi) 로 유지

rotation은 내부적으로 quaternion으로 저장하되, 필요 시 axis-angle / rotation matrix 변환 유틸을 제공한다.

4.3 LigandConformer

실제 원자 좌표를 가진 구조 표현.

pub struct LigandConformer {
    pub graph: LigandGraph,
    pub coords: Vec<[f32; 3]>,
}

요구사항:

coords.len() == graph.atoms.len() 이어야 함
reference conformer를 기준으로 pose transform 적용 가능해야 함

5. Burn에 연결 가능한 tensor adapter

AI 에이전트는 아래와 같은 방향의 adapter를 제공한다.

pub struct GraphTensors<B: burn::tensor::backend::Backend> {
    pub node_features: burn::tensor::Tensor<B, 2>,
    pub edge_index: burn::tensor::Tensor<B, 2, burn::tensor::Int>,
    pub edge_features: Option<burn::tensor::Tensor<B, 2>>,
    pub torsion_edge_index: burn::tensor::Tensor<B, 2, burn::tensor::Int>,
}

목표는 정확한 최종 타입을 박제하는 것이 아니라, 아래 contract를 만족하는 것이다.

contract

node features: [num_nodes, node_feat_dim]
edge index: [2, num_edges]
torsion edge index: [2, num_rotatable_bonds] 또는 equivalent representation
pose tensor export: translation / quaternion / torsion을 batchable tensor로 변환 가능

node feature 최소 구성

최초 버전은 아래 정도면 충분하다.

atomic number
formal charge
aromatic flag
degree(optional)

edge feature 최소 구성

bond order
aromatic flag
rotatable flag

pose tensor export

pub struct PoseTensors<B: burn::tensor::backend::Backend> {
    pub translation: burn::tensor::Tensor<B, 2>, // [batch, 3]
    pub rotation: burn::tensor::Tensor<B, 2>,    // [batch, 4] quaternion
    pub torsions: burn::tensor::Tensor<B, 2>,    // [batch, m]
}

6. 구현해야 할 manifold math 유틸

이 모듈은 모델과 완전히 분리된 순수 수학 유틸이어야 한다.

파일 예시:

src/geometry/
  euclidean.rs
  so3.rs
  torus.rs
  pose.rs

6.1 Translation (`R^3`)

필수 함수:

lerp_translation(p0, p1, t)
translation_delta(p0, p1)

6.2 Rotation (`SO(3)`)

필수 함수:

quaternion normalize
quaternion multiply / inverse
quaternion to rotation matrix
rotation matrix to quaternion (optional)
slerp(q0, q1, t)
relative_rotation(q_from, q_to)
so3_log(q_rel) -> [f32; 3]
so3_exp(omega) -> quat

설명:

simulator에서는 slerp 기반 interpolation만 먼저 구현해도 충분
이후 tangent vector supervision을 위해 log/exp utility 유지 권장

6.3 Torsion (`T^m`)

필수 함수:

wrap_pi(x)
wrap_vec(theta)
shortest_angle_delta(a, b)
interpolate_torsion(theta0, theta1, t)

정의:

shortest_angle_delta(a, b) = wrap_pi(b - a)
interpolate_torsion(theta0, theta1, t) = wrap(theta0 + t * shortest_delta)

6.4 Product Pose Geometry

필수 함수:

interpolate_pose(pose0, pose1, t) -> PoseState
pose_delta(pose0, pose1) -> PoseDelta

예상 구조:

pub struct PoseDelta {
    pub translation: [f32; 3],
    pub rotation_tangent: [f32; 3],
    pub torsions: Vec<f32>,
}

7. 좌표 복원 / simulator 요구사항

핵심 요구는 다음과 같다.

reference conformer + pose state 를 받아, 시각화 가능한 현재 좌표를 만든다.

7.1 적용 순서

권장 적용 순서:

reference conformer 좌표 복사
rotatable bond별 torsion 회전 적용
global rotation 적용
global translation 적용

7.2 torsion 회전 적용

각 TorsionEdge에 대해:

axis = bond direction
pivot = bond axis 상의 원자 위치
rotating_side에 포함된 atom들만 회전
angle = 해당 torsion value

즉, torsion은 local internal transform이다.

7.3 global transform 적용

rotation quaternion으로 전체 원자 좌표 회전
translation vector 더하기

7.4 simulator

입력:

reference conformer
pose_start
pose_target
t in [0,1]

출력:

PoseState at t
LigandConformer at t

즉,

pub fn simulate_at_t(
    reference: &LigandConformer,
    pose_start: &PoseState,
    pose_target: &PoseState,
    t: f32,
) -> LigandConformer

형태의 API가 있어야 한다.

7.5 trajectory

추가로 아래 helper도 구현한다.

pub fn simulate_trajectory(
    reference: &LigandConformer,
    pose_start: &PoseState,
    pose_target: &PoseState,
    num_steps: usize,
) -> Vec<LigandConformer>

t = 0..1 균등 분할
endpoint 포함
결과는 animation 또는 프레임 저장에 활용 가능

8. Visualization 요구사항

이번 단계에서 visualization은 매우 중요하다. 발표용 데모에 직접 쓰일 수 있어야 한다.

필수 출력

ligand graph connectivity 확인용 2D/3D static view
current pose 단일 프레임 렌더링
trajectory animation용 frame sequence export

최소 기능

A. static plot

입력:

atom coordinates
bond list

출력:

atoms = points
bonds = lines
rotatable bond는 다른 색상 또는 강조

B. trajectory frames

입력:

Vec<LigandConformer>

출력:

png sequence 또는 gif용 intermediate frames

C. optional: HTML viewer

가능하면 아래 중 하나 추가 검토:

plotly 기반 HTML
three.js/json export용 포맷
simple .xyz / .sdf frame export

시각화 우선순위

정적 PNG 출력
trajectory frame sequence 출력
선택적으로 gif/mp4/HTML

주의:

visualization 구현은 model과 완전히 분리
geometry correctness 검증이 목적이지 예쁜 UI가 목적은 아님

9. 권장 디렉터리 구조

src/
  graph/
    atom.rs
    bond.rs
    torsion.rs
    ligand_graph.rs
    featurize.rs

  geometry/
    euclidean.rs
    so3.rs
    torus.rs
    pose.rs

  conformer/
    conformer.rs
    apply_torsion.rs
    apply_pose.rs

  burn_adapter/
    graph_tensors.rs
    pose_tensors.rs

  simulator/
    interpolate.rs
    trajectory.rs

  viz/
    plot_static.rs
    plot_trajectory.rs
    export.rs

  lib.rs

10. 구현 상세 계약

10.1 반드시 보장해야 할 것

quaternion normalization이 항상 유지될 것
torsion angle은 항상 wrapped range에 있을 것
rotating_side 계산이 일관될 것
bond axis 기준 회전이 올바를 것
simulate_at_t(..., 0.0) 는 시작 pose와 일치할 것
simulate_at_t(..., 1.0) 는 타깃 pose와 일치할 것

10.2 numerical sanity checks

AI 에이전트는 아래 테스트를 함께 작성한다.

Test 1: wrap correctness

pi - eps 와 -pi + eps 의 차이가 작은 값으로 처리되는지 검증

Test 2: quaternion norm

interpolation 전후 quaternion norm ≈ 1

Test 3: endpoint correctness

trajectory 첫 프레임과 마지막 프레임이 입력 pose와 일치

Test 4: rigid transform invariance

torsion 없는 molecule에서는 internal geometry가 rigid하게 유지

Test 5: torsion rotation locality

특정 torsion 회전 시 rotating_side atom만 변하는지 확인

11. 입력 데이터 가정

이번 구현은 범용 loader까지 강제하지 않는다. 아래 두 입력 경로 중 하나만 우선 지원하면 된다.

경로 A: 수동 구성

atom list
bond list
reference coordinates
rotatable bond metadata
target pose

경로 B: RDKit 등 외부 전처리 결과 import

JSON 또는 serde-friendly 포맷으로 graph + conformer metadata load

권장:

Rust 내부 구현은 외부 화학 toolkit에 강하게 묶지 말 것
외부 전처리 결과를 import하는 구조가 더 안전함

12. 구현 우선순위

Phase 1

Atom, Bond, TorsionEdge, LigandGraph
PoseState
wrap_pi, shortest_angle_delta
quaternion normalize + slerp
interpolate_pose

Phase 2

LigandConformer
torsion application
global rotation/translation application
simulate_at_t

Phase 3

Burn tensor adapter
graph/node/edge feature export
pose tensor export

Phase 4

visualization static export
trajectory frame export
tests / sanity checks

13. 기대 산출물

AI 에이전트는 최종적으로 아래를 제공해야 한다.

코드 산출물

Rust crate source
기본 unit tests
example 실행 코드

example 요구

example은 아래를 보여줘야 한다.

ligand graph 생성 또는 로드
reference conformer 준비
start pose / target pose 설정
t=0.0, 0.25, 0.5, 0.75, 1.0 프레임 생성
static image 또는 trajectory frame export
graph tensor / pose tensor 변환 예시

14. AI 에이전트에게 주는 구현 지침

반드시 지킬 것

너무 많은 abstraction을 한 번에 넣지 말 것
우선 작동하는 명시적 구현을 만들 것
geometry correctness를 최우선으로 둘 것
Burn 모델 학습 코드까지 과도하게 확장하지 말 것

선호하는 구현 스타일

plain Rust struct 중심
명시적 타입
unit test 풍부하게
실패 시 panic보다 Result 기반 오류 처리 선호
visualization/output path는 예제에서 바로 실행 가능하게 구성

금지 사항

불필요한 거대 프레임워크 도입
protein context까지 무리하게 확장
학습 코드와 시각화 코드를 강하게 결합
화학 toolkit 종속성을 core domain object 안에 깊게 삽입

15. 구현 완료 판정 기준

아래 조건을 만족하면 이번 단계 구현 완료로 본다.

ligand graph와 torsion metadata를 표현할 수 있음
pose state를 translation / rotation / torsion으로 표현할 수 있음
t=0..1 explicit interpolation simulator가 동작함
conformer 좌표를 올바르게 복원함
static frame 및 trajectory frame을 export할 수 있음
Burn 입력용 tensor adapter가 존재함
최소한의 numerical sanity tests가 모두 통과함

16. 선택적 확장 아이디어

이번 단계 이후 확장 후보:

batch support
pocket context 추가
learned vector field model 연결
explicit path supervision dataset 생성기
SE(3)-equivariant / graph neural model 연결
target pose 단일 개체가 아니라 multiple conformer target 지원

17. 마지막 메모

이 구현은 최종 모델이 아니라, 아래 질문에 답하기 위한 prototype이다.

“ligand pose manifold (R^3 × SO(3) × T^m) 위의 상태를 Rust/Burn 친화적으로 표현하고, graph 입력과 연결하며, explicit path simulator와 visualization까지 한 번에 돌릴 수 있는가?”

이 질문에 yes를 만들 수 있으면 이번 단계는 충분히 성공이다.

15 KiB Raw Blame History Unescape Escape

Burn 기반 Pose Manifold Graph 구현 AI 전달 문서

0. 문서 목적

1. 문제 설정

1.1 상태 공간

1.2 그래프 입력

1.3 실험 목적

2. 구현 범위 요약

필수 구현

이번 단계에서 제외

3. 설계 원칙

3.1 Burn 친화적 설계

3.2 API 변동에 덜 민감한 설계

3.3 발표용 prototype 우선

4. 구현해야 할 핵심 데이터 구조

4.1 Atom / Bond / LigandGraph

Atom

Bond

TorsionEdge

LigandGraph

4.2 PoseState

권장 표현

4.3 LigandConformer

5. Burn에 연결 가능한 tensor adapter

contract

node feature 최소 구성

edge feature 최소 구성

pose tensor export

6. 구현해야 할 manifold math 유틸

6.1 Translation (R^3)

6.2 Rotation (SO(3))

6.3 Torsion (T^m)

6.4 Product Pose Geometry

7. 좌표 복원 / simulator 요구사항

7.1 적용 순서

7.2 torsion 회전 적용

7.3 global transform 적용

7.4 simulator

7.5 trajectory

8. Visualization 요구사항

필수 출력

최소 기능

A. static plot

B. trajectory frames

C. optional: HTML viewer

시각화 우선순위

9. 권장 디렉터리 구조

10. 구현 상세 계약

10.1 반드시 보장해야 할 것

10.2 numerical sanity checks

Test 1: wrap correctness

Test 2: quaternion norm

Test 3: endpoint correctness

Test 4: rigid transform invariance

Test 5: torsion rotation locality

11. 입력 데이터 가정

경로 A: 수동 구성

경로 B: RDKit 등 외부 전처리 결과 import

12. 구현 우선순위

Phase 1

Phase 2

Phase 3

Phase 4

13. 기대 산출물

코드 산출물

example 요구

14. AI 에이전트에게 주는 구현 지침

반드시 지킬 것

선호하는 구현 스타일

금지 사항

15. 구현 완료 판정 기준

16. 선택적 확장 아이디어

17. 마지막 메모

15 KiB

Raw Blame History

6.1 Translation (`R^3`)

6.2 Rotation (`SO(3)`)

6.3 Torsion (`T^m`)