Mạng Đồ Thị 3D (3D Graph Networks)¶

Mạng đồ thị 3D mở rộng GNN cho dữ liệu có hình học không gian, nơi các phép quay và tịnh tiến phải được xử lý chính xác. File này bao gồm đồ thị hình học (geometric graphs), tính tương đẳng SE(3)/E(n), SchNet, DimeNet, EGNN, tensor field networks, và các ứng dụng trong dự đoán tính chất phân tử, cấu trúc protein, khoa học vật liệu, và khám phá thuốc — những kiến trúc học từ thế giới vật lý 3D.

Các GNN trong file 3 và 4 hoạt động trên đồ thị trừu tượng: nút có đặc trưng, cạnh mã hóa kết nối, nhưng không có khái niệm không gian 3D. Một đồ thị mạng xã hội không có hình học. Nhưng nhiều ứng dụng có tác động nhất của GNN liên quan đến dữ liệu sống trong không gian vật lý 3D: phân tử, protein, tinh thể, đám mây điểm. Với những dữ liệu này, vị trí không gian của các nút mang thông tin quan trọng mà GNN trừu tượng bỏ qua.
Thách thức là dữ liệu 3D có các đối xứng hình học (geometric symmetries) (file 1): xoay một phân tử không thay đổi các tính chất của nó, và tịnh tiến nó cũng vậy. Một GNN 3D phải tôn trọng các đối xứng này. Một dự đoán năng lượng thay đổi khi bạn xoay phân tử là sai về mặt vật lý.

Đồ Thị Hình Học (Geometric Graphs)¶

Một đồ thị hình học (geometric graph) là một đồ thị được nhúng trong không gian 3D. Mỗi nút \(i\) có một vị trí \(\mathbf{r}_i \in \mathbb{R}^3\) ngoài vector đặc trưng \(\mathbf{h}_i\) của nó. Các cạnh có thể được định nghĩa bằng sự gần nhau về không gian (kết nối các nút trong khoảng cách \(r_{\text{cut}}\)) thay vì bằng các liên kết rõ ràng.
Với phân tử, đồ thị hình học có các nguyên tử là nút (với các đặc trưng: loại nguyên tố, điện tích, v.v.) và các liên kết hóa học là các cạnh. Các vị trí 3D \(\mathbf{r}_i\) là tọa độ nguyên tử, được xác định bởi cơ học lượng tử hoặc đo lường thực nghiệm (tinh thể học tia X, cryo-EM).
Với đám mây điểm (từ LiDAR hoặc máy quét 3D, chương 8 và chương 11), mỗi điểm là một nút với vị trí và các đặc trưng tùy chọn (màu sắc, cường độ). Các cạnh kết nối các điểm gần nhau, tạo thành một đồ thị k-láng giềng gần nhất (kNN graph) hoặc đồ thị bán kính.
Các đại lượng hình học chính cho truyền thông điệp là:
- Khoảng cách giữa các nguyên tử (Interatomic distances): \(d_{ij} = \|\mathbf{r}_i - \mathbf{r}_j\|\). Khoảng cách bất biến với phép quay và tịnh tiến. Hai phân tử có khoảng cách giữa các nguyên tử giống hệt nhau là cùng hình dạng, bất kể hướng.
- Góc liên kết (Bond angles): góc \(\theta_{ijk}\) giữa các vector \(\mathbf{r}_j - \mathbf{r}_i\) và \(\mathbf{r}_k - \mathbf{r}_i\) tại nút \(i\). Góc nắm bắt hình học cục bộ ngoài khoảng cách từng cặp.
- Góc xoắn (Dihedral/torsion angles): góc \(\phi_{ijkl}\) giữa các mặt phẳng xác định bởi \((i, j, k)\) và \((j, k, l)\). Góc xoắn nắm bắt cách cấu trúc xoắn trong 3D, cần thiết cho hình học khung xương protein.
- Các vector vị trí tương đối (Relative position vectors): \(\mathbf{r}_{ij} = \mathbf{r}_j - \mathbf{r}_i\). Các vector này bất biến với tịnh tiến nhưng KHÔNG bất biến với phép quay. Sử dụng chúng đòi hỏi các kiến trúc tương đẳng (equivariant), không chỉ bất biến.

Tính Tương Đẳng SE(3) và E(n)¶

Nhóm đối xứng cho dữ liệu vật lý 3D là nhóm Euclid \(E(3)\), bao gồm tất cả các phép quay, phản chiếu, và tịnh tiến. Nhóm con \(SE(3)\) (Euclid đặc biệt) bao gồm các phép quay và tịnh tiến nhưng loại trừ các phản chiếu.
Một GNN 3D nên:
- Bất biến với tịnh tiến (Translation-invariant) cho đầu ra vô hướng (năng lượng, ái lực liên kết): dịch chuyển tất cả các nguyên tử bởi cùng một vector không nên thay đổi dự đoán.
- Bất biến với quay (Rotation-invariant) cho đầu ra vô hướng: xoay phân tử không nên thay đổi năng lượng của nó.
- Tương đẳng với quay (Rotation-equivariant) cho đầu ra vector/tensor (lực, mô-men lưỡng cực): xoay phân tử nên xoay các vector lực dự đoán tương ứng.

Tương đẳng SE(3): xoay một phân tử không thay đổi các dự đoán vô hướng (năng lượng) nhưng xoay các dự đoán vector (lực) tương ứng

Một cách hình thức, với dự đoán vô hướng \(f\) và phép quay \(R \in SO(3)\):

\[f(R\mathbf{r}_1, R\mathbf{r}_2, \ldots) = f(\mathbf{r}_1, \mathbf{r}_2, \ldots) \quad \text{(bất biến)}\]

Với dự đoán vector \(\mathbf{F}\):

\[\mathbf{F}(R\mathbf{r}_1, R\mathbf{r}_2, \ldots) = R \cdot \mathbf{F}(\mathbf{r}_1, \mathbf{r}_2, \ldots) \quad \text{(tương đẳng)}\]

Các ràng buộc này phản ánh trực tiếp khuôn khổ bất biến/tương đẳng từ file 1, bây giờ được áp dụng cụ thể cho các nhóm quay và tịnh tiến 3D.
Có hai cách tiếp cận thiết kế:
1. Kiến trúc bất biến (Invariant architectures): chỉ dùng các đặc trưng hình học bất biến (khoảng cách, góc) làm đầu vào cho truyền thông điệp. Các biểu diễn bên trong là vô hướng (bất biến). Đơn giản và hiệu quả nhưng không thể tạo ra đầu ra vector mà không phá vỡ đối xứng.
2. Kiến trúc tương đẳng (Equivariant architectures): duy trì các biểu diễn vector (và tensor bậc cao hơn) trong suốt mạng, đảm bảo mỗi tầng là tương đẳng. Biểu cảm hơn, có thể dự đoán vector và tensor một cách tự nhiên, nhưng phức tạp hơn.

SchNet: Truyền Thông Điệp Dựa Trên Khoảng Cách¶

SchNet (Schütt et al., 2017) là GNN 3D bất biến nền tảng. Cải tiến chính của nó là tích chập bộ lọc liên tục (continuous filter convolution): thay vì dùng một tập cố định các loại cạnh (như loại liên kết trong GNN phân tử), SchNet tạo bộ lọc thông điệp trực tiếp từ khoảng cách giữa các nguyên tử.
Khoảng cách \(d_{ij}\) trước tiên được mở rộng thành một vector đặc trưng dùng hàm cơ sở xuyên tâm (radial basis functions — RBFs):

\[\text{RBF}(d_{ij}) = \left[\exp\left(-\gamma_1 (d_{ij} - \mu_1)^2\right), \ldots, \exp\left(-\gamma_K (d_{ij} - \mu_K)^2\right)\right]\]

Mỗi hàm cơ sở là một Gauss tập trung tại \(\mu_k\) với độ rộng \(\gamma_k\). Điều này tương tự như một mã hóa vị trí học được cho khoảng cách: khoảng cách liên tục được ánh xạ đến một không gian đặc trưng chiều cao nơi mạng có thể học các tương tác phụ thuộc khoảng cách. Các tâm \(\mu_k\) thường được đặt cách đều từ 0 đến bán kính cắt.
Thông điệp SchNet từ nút \(j\) đến nút \(i\) là:

\[\mathbf{m}_{j \to i} = \mathbf{h}_j \odot W_{\text{filter}}(\text{RBF}(d_{ij}))\]

trong đó \(W_{\text{filter}}\) là một MLP ánh xạ khai triển RBF thành một vector bộ lọc, và \(\odot\) là phép nhân theo từng phần tử (Hadamard product, chương 2). Bộ lọc phụ thuộc vào khoảng cách, vì vậy các nguyên tử gần nhau tương tác khác với các nguyên tử xa. Phép nhân theo từng phần tử giống như một cơ chế gating (chương 6): bộ lọc phụ thuộc khoảng cách kiểm soát mức độ đặc trưng của mỗi chiều đi qua.
Vì SchNet chỉ dùng các khoảng cách (bất biến), toàn bộ mô hình tự động bất biến với phép quay và tịnh tiến. Không cần xử lý đặc biệt cho đối xứng ngoài lựa chọn thiết kế này.

DimeNet và SphereNet: Góc và Góc Xoắn¶

Chỉ riêng khoảng cách không thể xác định đầy đủ cấu trúc 3D. Hai cấu hình phân tử khác nhau có thể có các khoảng cách từng cặp giống hệt nhau nhưng góc liên kết khác nhau (đây là vấn đề "mơ hồ hình học khoảng cách"). DimeNet (Gasteiger et al., 2020) đưa góc liên kết (bond angles) vào truyền thông điệp.
DimeNet dùng truyền thông điệp có hướng (directional message passing): các thông điệp chảy dọc theo các cạnh có hướng, và thông điệp trên cạnh \((j \to i)\) bị ảnh hưởng bởi góc giữa các cạnh \((k \to j)\) và \((j \to i)\):

\[\mathbf{m}_{kj \to ji} = f\left(\mathbf{m}_{kj}, d_{ji}, \theta_{kji}\right)\]

Góc \(\theta_{kji}\) được khai triển dùng các hàm cầu Bessel và hàm điều hòa cầu (spherical harmonics) — cơ sở tự nhiên cho thông tin góc trên một mặt cầu, tương tự như RBF cho khoảng cách. Điều này cung cấp cho mô hình quyền truy cập vào thông tin hướng trong khi vẫn duy trì tính bất biến.
SphereNet (Liu et al., 2022) đi xa hơn bằng cách bao gồm góc xoắn (dihedral angles) \(\phi_{lkji}\), nắm bắt cấu trúc xoắn 3D đầy đủ. Hệ thống phân cấp là:
- Khoảng cách → nắm bắt sự gần nhau giữa các cặp
- Góc → nắm bắt hình học cục bộ (gấp khúc so với thẳng hàng)
- Góc xoắn → nắm bắt sự xoắn 3D (quan trọng cho khung xương protein, liên kết thuốc)
Mỗi cấp thêm độ phân giải hình học với chi phí độ phức tạp tính toán (khoảng cách là \(O(|E|)\), góc là \(O(|E| \cdot k)\), góc xoắn là \(O(|E| \cdot k^2)\) trong đó \(k\) là bậc trung bình).

GNN Tương Đẳng E(n) (EGNN — E(n) Equivariant GNNs)¶

EGNN (Satorras et al., 2021) theo cách tiếp cận tương đẳng: thay vì chỉ dùng các đặc trưng bất biến, nó cập nhật cả đặc trưng nút và vị trí nút tại mỗi tầng, duy trì tính tương đẳng xuyên suốt.
Phép cập nhật EGNN cho nút \(i\):

\[\mathbf{m}_{ij} = \phi_e\left(\mathbf{h}_i, \mathbf{h}_j, d_{ij}^2, a_{ij}\right)\]

\[\mathbf{r}_i' = \mathbf{r}_i + C \sum_{j \neq i} (\mathbf{r}_i - \mathbf{r}_j) \cdot \phi_r(\mathbf{m}_{ij})\]

\[\mathbf{h}_i' = \phi_h\left(\mathbf{h}_i, \sum_j \mathbf{m}_{ij}\right)\]

Chìa khóa là cập nhật vị trí: các vị trí nút được điều chỉnh bởi một tổng có trọng số của các vector vị trí tương đối \((\mathbf{r}_i - \mathbf{r}_j)\). Các trọng số đến từ hàm thông điệp \(\phi_r\), chỉ phụ thuộc vào các đại lượng bất biến (đặc trưng và khoảng cách). Cấu trúc này có thể chứng minh là tương đẳng: nếu tất cả các vị trí đầu vào được quay bởi \(R\), tất cả các vị trí đầu ra cũng được quay bởi cùng \(R\) đó.
EGNN thanh lịch vì nó đạt được tính tương đẳng mà không cần dùng rõ ràng các hàm điều hòa cầu hay các biểu diễn bất khả quy. Các vector vị trí tương đối mang thông tin hướng, và hàm thông điệp bất biến kiểm soát cách thông tin hướng đó được sử dụng.
Sự đơn giản đi kèm với một sự đánh đổi: EGNN chỉ dùng các biểu diễn vector (bậc 1). Nó không thể biểu diễn các tensor bậc cao hơn như mô-men tứ cực hay tensor ứng suất mà không mở rộng.

Tensor Field Networks và Các Biểu Diễn Bậc Cao Hơn¶

Tensor Field Networks (Thomas et al., 2018) và các phiên bản kế tiếp (SE(3)-Transformers, MACE, Equiformer) dùng bộ máy đầy đủ của các biểu diễn bất khả quy (irreducible representations) của nhóm quay để xây dựng các tầng tương đẳng.
Trong lý thuyết biểu diễn (kết nối với đại số tuyến tính của chương 2), các phép quay trong 3D có thể được phân tích thành các thành phần bất khả quy đặc trưng bởi một bậc nguyên \(\ell\):
- \(\ell = 0\): vô hướng (1 thành phần, bất biến). Năng lượng, điện tích.
- \(\ell = 1\): vector (3 thành phần, quay như các vector vị trí). Lực, mô-men lưỡng cực.
- \(\ell = 2\): tensor đối xứng bậc 2 không vết (5 thành phần). Mô-men tứ cực, tensor ứng suất.
- \(\ell\) cao hơn: nắm bắt cấu trúc góc ngày càng phức tạp.
Đây được gọi là tensor cầu (spherical tensors), và chúng biến đổi dưới phép quay \(R\) qua các ma trận Wigner-D \(D^\ell(R)\): một vô hướng không thay đổi, một vector quay bởi \(R\), một tensor bậc 2 quay bởi một ma trận phức tạp hơn.
Truyền thông điệp tương đẳng (Equivariant message passing) với tensor cầu dùng tích tensor Clebsch-Gordan để kết hợp các đặc trưng có bậc khác nhau:

\[(\mathbf{f}^{\ell_1} \otimes \mathbf{f}^{\ell_2})^{\ell_{\text{out}}} = \sum_{m_1, m_2} C^{\ell_{\text{out}}, m_{\text{out}}}_{\ell_1, m_1, \ell_2, m_2} \cdot f^{\ell_1}_{m_1} \cdot f^{\ell_2}_{m_2}\]

Các hệ số Clebsch-Gordan \(C\) là các hằng số toán học cố định đảm bảo tích tensor là tương đẳng. Đây là tương tự SO(3)-tương đẳng của phép nhân ma trận.
MACE (Batatia et al., 2022) dùng các thông điệp bậc cao (tích của nhiều đặc trưng hàng xóm) để đạt được độ chính xác cao với ít tầng truyền thông điệp hơn. Bằng cách xây dựng các tương tác có bậc cơ thể (body-ordered interactions) (2-body từ khoảng cách, 3-body từ góc, nhiều-body từ tích tensor), MACE nắm bắt các tương tác nguyên tử phức tạp một cách hiệu quả.
Equiformer (Liao & Smidt, 2023) kết hợp các đặc trưng tensor cầu tương đẳng với cơ chế chú ý transformer (file 4), tạo ra một Graph Transformer tương đẳng SE(3). Điểm chú ý được tính từ các đặc trưng bất biến, trong khi việc tập hợp giá trị hoạt động trên các đặc trưng tensor tương đẳng.

Ứng Dụng (Applications)¶

Dự đoán tính chất phân tử (Molecular property prediction): với cấu trúc 3D của một phân tử, dự đoán các tính chất như năng lượng, lực, mô-men lưỡng cực, khoảng cách HOMO-LUMO, độc tính, độ hòa tan. Đây là ứng dụng trưởng thành nhất của GNN 3D. Các mô hình được huấn luyện trên các bộ dữ liệu hóa lượng tử (QM9, OC20) đạt độ chính xác hóa học trên nhiều tính chất, cho phép sàng lọc ảo hàng triệu phân tử ứng viên.
Tăng tốc mô phỏng phân tử (Molecular dynamics acceleration): tính toán lực giữa các nguyên tử dùng cơ học lượng tử (lý thuyết hàm mật độ — DFT) cực kỳ đắt (\(O(n^3)\) cho \(n\) điện tử). Một GNN 3D được huấn luyện để dự đoán lực có thể thay thế DFT trong các mô phỏng động lực học phân tử, đạt gia tốc \(10^3\)–\(10^6\) trong khi duy trì độ chính xác gần với DFT. Điều này cho phép mô phỏng các hệ thống lớn hơn và thang thời gian dài hơn, tiết lộ các hiện tượng vô hình với các phương pháp truyền thống.
Cấu trúc protein (Protein structure): protein là các chuỗi axit amin gấp lại thành các cấu trúc 3D phức tạp. Khung xương protein là một đồ thị hình học nơi các nút là các gốc (residues) và các cạnh kết nối các gốc gần nhau về không gian. GNN 3D được dùng cho dự đoán chức năng protein, xác định vị trí liên kết, và thiết kế protein (inverse folding: với một cấu trúc mong muốn, dự đoán dãy axit amin). AlphaFold dùng suy luận hình học và dựa trên đồ thị để dự đoán cấu trúc protein từ dãy.
Khoa học vật liệu và xúc tác (Materials science and catalysis): vật liệu tinh thể có cấu trúc 3D tuần hoàn. GNN mô hình hóa ô đơn vị lặp lại và dự đoán các tính chất vật liệu: độ rộng vùng cấm (band gap), năng lượng hình thành, độ bền cơ học. Dự án Open Catalyst (OC20/OC22) chuẩn hóa GNN cho dự đoán năng lượng hấp phụ trên các bề mặt xúc tác, tăng tốc việc tìm kiếm các chất xúc tác mới cho năng lượng tái tạo.
Khám phá thuốc (Drug discovery): GNN 3D dự đoán cách một phân tử thuốc sẽ liên kết với một protein mục tiêu. Ái lực liên kết phụ thuộc vào sự bổ sung hình dạng 3D và các tương tác hóa học giữa thuốc và túi liên kết của protein. Các mô hình như DiffDock dùng GNN tương đẳng với mô hình khuếch tán (chương 8) để dự đoán tư thế liên kết (hướng 3D của thuốc trong túi protein).

Sinh Đồ Thị (Graph Generation)¶

Tất cả các kiến trúc trên phân tích các đồ thị hiện có. Sinh đồ thị (Graph generation) tạo ra các đồ thị mới: thiết kế một phân tử với các tính chất mong muốn, tạo một mạng xã hội tổng hợp để kiểm thử, hoặc đề xuất một cấu trúc protein mới. Đây là đối tác sinh (generative counterpart) của dự đoán cấp đồ thị.
Thách thức là đồ thị rời rạc, có kích thước thay đổi, và mang tính tổ hợp. Sinh một đồ thị nghĩa là quyết định bao nhiêu nút để tạo, chúng có đặc trưng gì, và cặp nào để kết nối. Không gian các đồ thị khả dĩ tăng siêu mũ (super-exponential) với số lượng nút.
Sinh tự hồi quy (Autoregressive generation) xây dựng đồ thị từng nút (hoặc từng cạnh) một. GraphRNN (You et al., 2018) sinh đồ thị một cách tuần tự: một RNN duy trì một trạng thái, tạo một nút mới ở mỗi bước, và quyết định kết nối nó với những nút nào đã tồn tại. Thứ tự sinh áp đặt một dãy nhân tạo lên đồ thị vốn không có thứ tự, nhưng các thứ tự BFS giúp bằng cách giữ các nút được tạo gần đây có liên quan.
Sinh dựa trên VAE mã hóa đồ thị vào một không gian tiềm ẩn liên tục (dùng bộ mã hóa GNN), sau đó giải mã các đồ thị mới từ các vector tiềm ẩn được lấy mẫu. GraphVAE sinh một ma trận kề xác suất \(\hat{A} \in [0, 1]^{n \times n}\) trong một lần, nhưng điều này tỷ lệ \(O(n^2)\) và tạo ra đầu ra dày đặc phải được ngưỡng hóa. Không gian tiềm ẩn cho phép nội suy mượt: di chuyển giữa hai embedding phân tử tạo ra các cấu trúc trung gian có tính hợp lệ hóa học.
Sinh dựa trên khuếch tán (Diffusion-based generation) áp dụng khuôn khổ khuếch tán (chương 8) lên đồ thị. Quá trình thuận dần dần thêm nhiễu vào các đặc trưng nút và cấu trúc cạnh. Quá trình nghịch học để khử nhiễu, tạo ra các đồ thị hợp lệ từ nhiễu. DiGress (Vignac et al., 2023) áp dụng khuếch tán rời rạc cho cả loại nút và loại cạnh, xử lý một cách tự nhiên bản chất phân loại của dữ liệu đồ thị.
Với sinh phân tử, ràng buộc chính là tính hợp lệ hóa học (chemical validity): các phân tử được sinh ra phải tuân theo quy tắc hóa trị (carbon tạo 4 liên kết, oxy tạo 2, v.v.). Các phương pháp như Junction Tree VAE (JT-VAE) phân tích các phân tử thành các cấu trúc con hợp lệ (vòng, chuỗi, nhóm chức năng) và sinh bằng cách lắp ráp các khối xây dựng này, đảm bảo tính hợp lệ ngay từ thiết kế.
Sinh có định hướng mục tiêu (Goal-directed generation) tối ưu hóa cho các tính chất cụ thể: sinh một phân tử với ái lực liên kết cao với một protein mục tiêu, độc tính thấp, và độ hòa tan tốt. Điều này kết hợp sinh đồ thị với dự đoán tính chất (dùng GNN 3D làm bộ đánh giá tính chất) trong một vòng lặp: sinh → đánh giá → tinh chỉnh. Học tăng cường (chương 6) hoặc tối ưu hóa Bayes hướng dẫn việc tìm kiếm trong không gian hóa học.
DiffDock (Corso et al., 2023) dùng khuếch tán tương đẳng SE(3) để dự đoán cách một phân tử thuốc gắn (dock) vào túi liên kết protein. Mô hình sinh ra tư thế liên kết 3D (vị trí và hướng của thuốc so với protein) bằng cách khử nhiễu từ các vị trí ngẫu nhiên, kết hợp các mạng tương đẳng 3D từ file này với khuôn khổ khuếch tán từ chương 8.

Bài tập lập trình (dùng CoLab hoặc notebook)¶

Xây dựng một tầng truyền thông điệp 3D bất biến đơn dùng khoảng cách giữa các nguyên tử. Áp dụng nó lên một phân tử nhỏ (nước: H-O-H) và kiểm tra rằng đầu ra bất biến với phép quay.

import jax
import jax.numpy as jnp

# Water molecule: O at origin, two H atoms
positions = jnp.array([[0.0, 0.0, 0.0],     # O
                        [0.96, 0.0, 0.0],    # H1
                        [-0.24, 0.93, 0.0]])  # H2

# Node features: [atomic number]
features = jnp.array([[8.0], [1.0], [1.0]])

# Compute pairwise distances (invariant)
def pairwise_distances(pos):
    diff = pos[:, None, :] - pos[None, :, :]
    return jnp.sqrt(jnp.sum(diff**2, axis=-1) + 1e-8)

# Simple distance-based message passing
def invariant_message_pass(features, positions):
    dists = pairwise_distances(positions)
    # RBF expansion with 4 centres
    centres = jnp.array([0.5, 1.0, 1.5, 2.0])
    rbf = jnp.exp(-5.0 * (dists[:, :, None] - centres[None, None, :]) ** 2)

    # Message: features weighted by distance-dependent filter
    messages = jnp.einsum("ij,jd->id", rbf.sum(axis=-1), features)
    return messages

output1 = invariant_message_pass(features, positions)

# Rotate the molecule by 90 degrees around z-axis
R = jnp.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=float)
rotated_positions = (R @ positions.T).T

output2 = invariant_message_pass(features, rotated_positions)

print(f"Original output:\n{output1}")
print(f"\nRotated output:\n{output2}")
print(f"\nInvariant: {jnp.allclose(output1, output2, atol=1e-5)}")

Tính góc liên kết giữa ba nguyên tử và kiểm tra nó bất biến với phép quay.

import jax.numpy as jnp

def bond_angle(r_i, r_j, r_k):
    """Angle at node j between edges j->i and j->k."""
    v1 = r_i - r_j
    v2 = r_k - r_j
    cos_angle = jnp.dot(v1, v2) / (jnp.linalg.norm(v1) * jnp.linalg.norm(v2))
    return jnp.arccos(jnp.clip(cos_angle, -1, 1))

# Three atoms
r1 = jnp.array([1.0, 0.0, 0.0])
r2 = jnp.array([0.0, 0.0, 0.0])
r3 = jnp.array([0.0, 1.0, 0.0])

angle_original = bond_angle(r1, r2, r3)
print(f"Original angle: {jnp.degrees(angle_original):.1f}°")

# Apply random rotation
R = jnp.array([[0.36, 0.48, -0.80],
               [-0.80, 0.60, 0.00],
               [0.48, 0.64, 0.60]])
r1_rot, r2_rot, r3_rot = R @ r1, R @ r2, R @ r3

angle_rotated = bond_angle(r1_rot, r2_rot, r3_rot)
print(f"Rotated angle:  {jnp.degrees(angle_rotated):.1f}°")
print(f"Invariant: {jnp.allclose(angle_original, angle_rotated, atol=1e-4)}")

Minh họa cập nhật vị trí tương đẳng (kiểu EGNN). Cập nhật các vị trí nút dùng vector tương đối có trọng số theo khoảng cách và kiểm tra tính tương đẳng.

import jax
import jax.numpy as jnp

def egnn_position_update(positions, features):
    """Simple EGNN-style equivariant position update."""
    n = positions.shape[0]
    new_positions = jnp.zeros_like(positions)

    for i in range(n):
        shift = jnp.zeros(3)
        for j in range(n):
            if i != j:
                r_ij = positions[i] - positions[j]
                d_ij = jnp.linalg.norm(r_ij)
                # Weight based on distance (simple: inverse distance)
                weight = 1.0 / (d_ij + 1.0)
                # Scale by feature similarity
                feat_sim = jnp.dot(features[i], features[j])
                shift = shift + weight * feat_sim * r_ij
        new_positions = new_positions.at[i].set(positions[i] + 0.1 * shift)

    return new_positions

# 3 atoms
pos = jnp.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
feat = jnp.array([[1.0, 0.5], [0.5, 1.0], [0.8, 0.3]])

# Update positions
pos_new = egnn_position_update(pos, feat)

# Now rotate input, update, and check if output is rotated consistently
R = jnp.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
pos_rot = (R @ pos.T).T
pos_new_from_rot = egnn_position_update(pos_rot, feat)

# Should be the same as rotating the original output
pos_new_then_rot = (R @ pos_new.T).T

print(f"Update then rotate:\n{jnp.round(pos_new_then_rot, 4)}")
print(f"\nRotate then update:\n{jnp.round(pos_new_from_rot, 4)}")
print(f"\nEquivariant: {jnp.allclose(pos_new_then_rot, pos_new_from_rot, atol=1e-4)}")