达摩院ElementsClaw深度解析：AI智能体如何用28GPU小时发现4种全新超导材料

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import MessagePassing, global_mean_pool
from torch_geometric.data import Data, Dataset
import numpy as np
from typing import Optional, Tuple, List

# ─── 原子级消息传递层 ───
class AtomicMessagePassing(MessagePassing):
    """
    几何感知的消息传递层。
    融入原子间距离、键角等几何信息，实现三维结构感知。
    """
    def __init__(self, in_channels: int, out_channels: int, num_filters: int = 128):
        super().__init__(aggr='add')  # 聚合方式: 求和
        self.in_channels = in_channels
        self.out_channels = out_channels
        
        # 边特征网络: 编码原子间距、键角、周期性等几何信息
        self.edge_mlp = nn.Sequential(
            nn.Linear(3, num_filters),  # 3维几何特征: 距离、角度、周期性偏移
            nn.SiLU(),
            nn.Linear(num_filters, out_channels),
            nn.SiLU(),
        )
        
        # 节点特征更新网络
        self.node_mlp = nn.Sequential(
            nn.Linear(in_channels + out_channels, out_channels * 2),
            nn.SiLU(),
            nn.Linear(out_channels * 2, out_channels),
        )
        
        # 残差连接
        self.residual = nn.Linear(in_channels, out_channels) if in_channels != out_channels else nn.Identity()
    
    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, 
                edge_attr: torch.Tensor) -> torch.Tensor:
        """前向传播"""
        return self.propagate(edge_index, x=x, edge_attr=edge_attr)
    
    def message(self, x_j: torch.Tensor, edge_attr: torch.Tensor) -> torch.Tensor:
        """消息生成: 融合邻居节点特征 + 边几何特征"""
        edge_features = self.edge_mlp(edge_attr)
        return edge_features * x_j  # 门控机制: 几何特征调制消息
    
    def update(self, aggr_out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
        """节点更新: 聚合消息 + 残差连接"""
        combined = torch.cat([x, aggr_out], dim=-1)
        out = self.node_mlp(combined)
        return out + self.residual(x)


# ─── Elements基础模型 ───
class ElementsFoundationModel(nn.Module):
    """
    1B参数原子基础模型。
    架构: 嵌入层 → 24层几何消息传递 → 读出层
    支持: 超导性分类(Tc阈值)、临界温度回归、能量预测、结构生成
    """
    def __init__(
        self,
        num_atom_types: int = 118,     # 元素周期表原子种类
        hidden_dim: int = 1024,         # 隐层维度
        num_layers: int = 24,           # 消息传递层数
        num_heads: int = 16,            # 多头注意力头数
        dropout: float = 0.1,
    ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        
        # 原子类型嵌入 (118种元素 + 特殊token)
        self.atom_embedding = nn.Embedding(num_atom_types + 1, hidden_dim)
        
        # 位置编码: 原子在晶胞中的分数坐标
        self.pos_mlp = nn.Sequential(
            nn.Linear(3, hidden_dim // 4),
            nn.SiLU(),
            nn.Linear(hidden_dim // 4, hidden_dim),
        )
        
        # 多层几何消息传递
        self.layers = nn.ModuleList([
            AtomicMessagePassing(hidden_dim, hidden_dim)
            for _ in range(num_layers)
        ])
        
        # LayerNorm + Dropout (每个子层后)
        self.layer_norms = nn.ModuleList([
            nn.LayerNorm(hidden_dim) for _ in range(num_layers)
        ])
        self.dropout = nn.Dropout(dropout)
        
        # 全局读出: 从原子表示聚合到晶体表示
        self.readout = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim * 2),
            nn.SiLU(),
            nn.Linear(hidden_dim * 2, hidden_dim),
        )
        
        # Scaling Law验证: 记录参数量
        self.total_params = sum(p.numel() for p in self.parameters())
        print(f"Elements Foundation Model Parameters: {self.total_params:,}")
    
    def forward(self, data: Data) -> torch.Tensor:
        """
        Args:
            data: PyG Data对象，包含:
                - x: 原子类型索引 [num_atoms]
                - pos: 原子分数坐标 [num_atoms, 3]
                - edge_index: 边连接 [2, num_edges]
                - edge_attr: 边几何特征 [num_edges, 3]
                - batch: 批次分配 [num_atoms]
        Returns:
            晶体级表示向量 [batch_size, hidden_dim]
        """
        # 原子嵌入 + 位置编码
        atom_feat = self.atom_embedding(data.x)  # [num_atoms, hidden_dim]
        pos_feat = self.pos_mlp(data.pos)        # [num_atoms, hidden_dim]
        h = atom_feat + pos_feat                 # [num_atoms, hidden_dim]
        
        # 多层消息传递 (带残差和LayerNorm)
        for i in range(self.num_layers):
            h_new = self.layers[i](h, data.edge_index, data.edge_attr)
            h_new = self.layer_norms[i](h_new)
            h = h + self.dropout(h_new)  # 残差连接
        
        # 全局池化: 从原子级到晶体级
        h_graph = global_mean_pool(h, data.batch)  # [batch_size, hidden_dim]
        
        # 读出网络
        return self.readout(h_graph)  # [batch_size, hidden_dim]


# ─── 四分支输出头 ───
class ElementsMultiHead(nn.Module):
    """
    Elements模型的四分支输出头。
    共享基础模型表示，各分支独立预测不同属性。
    """
    def __init__(self, hidden_dim: int = 1024):
        super().__init__()
        
        # 分支1: Elements-C (超导性分类)
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, 512),
            nn.SiLU(),
            nn.Dropout(0.1),
            nn.Linear(512, 128),
            nn.SiLU(),
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )
        
        # 分支2: Elements-T (临界温度回归)
        self.tc_regressor = nn.Sequential(
            nn.Linear(hidden_dim, 512),
            nn.SiLU(),
            nn.Dropout(0.1),
            nn.Linear(512, 256),
            nn.SiLU(),
            nn.Linear(256, 1),
        )
        
        # 分支3: Elements-E (能量/稳定性预测)
        self.energy_predictor = nn.Sequential(
            nn.Linear(hidden_dim, 512),
            nn.SiLU(),
            nn.Linear(512, 256),
            nn.SiLU(),
            nn.Linear(256, 1),
        )
        
        # 分支4: Elements-G (结构生成，通过扩散头实现)
        self.structure_generator = nn.Sequential(
            nn.Linear(hidden_dim, 1024),
            nn.SiLU(),
            nn.Linear(1024, 2048),
            nn.SiLU(),
            nn.Linear(2048, 1024),
        )
    
    def forward(self, h: torch.Tensor) -> dict:
        return {
            'superconducting_prob': self.classifier(h),       # [B, 1]
            'critical_temp': self.tc_regressor(h),            # [B, 1]
            'energy_stability': self.energy_predictor(h),     # [B, 1]
            'structure_latent': self.structure_generator(h),  # [B, 1024]
        }


# ─── 完整Elements模型 ───
class ElementsModel(nn.Module):
    """1B参数完整模型: 基础模型 + 多头输出"""
    def __init__(self, hidden_dim: int = 1024, num_layers: int = 24):
        super().__init__()
        self.foundation = ElementsFoundationModel(
            hidden_dim=hidden_dim, num_layers=num_layers
        )
        self.heads = ElementsMultiHead(hidden_dim=hidden_dim)
        
        # 验证Scaling Law: 记录参数量
        self.total_params = sum(p.numel() for p in self.parameters())
        
    def forward(self, data: Data) -> dict:
        h = self.foundation(data)
        return self.heads(h)
    
    def get_param_count(self) -> dict:
        """按组件统计参数量"""
        counts = {
            'foundation': sum(p.numel() for p in self.foundation.parameters()),
            'classifier_head': sum(p.numel() for p in self.heads.classifier.parameters()),
            'tc_head': sum(p.numel() for p in self.heads.tc_regressor.parameters()),
            'energy_head': sum(p.numel() for p in self.heads.energy_predictor.parameters()),
            'generator_head': sum(p.numel() for p in self.heads.structure_generator.parameters()),
        }
        counts['total'] = sum(counts.values())
        return counts


# ─── 训练验证 ───
def validate_elements_scaling():
    """
    Scaling Law验证：Elements首次在非LLM架构上验证了Scaling Law。
    更多参数 + 更多数据 → 性能持续提升。
    """
    configs = [
        {'hidden_dim': 256, 'num_layers': 8},
        {'hidden_dim': 512, 'num_layers': 16},
        {'hidden_dim': 768, 'num_layers': 20},
        {'hidden_dim': 1024, 'num_layers': 24},  # 1B参数版本
    ]
    
    results = []
    for cfg in configs:
        model = ElementsModel(**cfg)
        params = model.get_param_count()
        results.append({
            'config': cfg,
            'total_params': params['total'],
            'foundation_params': params['foundation'],
        })
        print(f"Config {cfg}: {params['total']:,} total params")
    
    return results


if __name__ == "__main__":
    # 实例化1B参数模型
    model = ElementsModel(hidden_dim=1024, num_layers=24)
    counts = model.get_param_count()
    
    print("=" * 60)
    print("Elements 1B Atom Foundation Model - Parameter Distribution")
    print("=" * 60)
    for name, count in counts.items():
        print(f"  {name:20s}: {count:>12,} params")
    print("=" * 60)
    
    # 验证Scaling Law
    print("\nScaling Law Verification:")
    validate_elements_scaling()

"""
ElementsClaw智能体核心工作流
"""
from enum import Enum
from dataclasses import dataclass, field
from typing import Optional, List, Dict
import json


class DiscoveryStage(Enum):
    """材料发现四阶段"""
    SCREENING = "screening"           # 大规模筛选
    CROSS_VALIDATION = "cross_val"    # 文献/数据库交叉验证
    SYNTHESIS_EVAL = "synthesis"      # 可合成性评估
    EXPERIMENT_DESIGN = "experiment"  # 实验方案设计


@dataclass
class CandidateMaterial:
    """候选材料数据结构"""
    formula: str                       # 化学式
    space_group: str                   # 空间群
    predicted_tc: float                # 预测临界温度(K)
    superconducting_prob: float        # 超导概率
    energy_above_hull: float           # 能量高于凸包(meV/atom)
    source: str                        # 来源: screening/cross_val/generated
    literature_matched: bool = False   # 文献是否已有报道
    synthesis_score: float = 0.0       # 可合成性评分(0-1)
    
    def to_dict(self) -> dict:
        return {
            "formula": self.formula,
            "space_group": self.space_group,
            "predicted_tc": self.predicted_tc,
            "superconducting_prob": self.superconducting_prob,
            "synthesis_score": self.synthesis_score,
            "source": self.source,
        }


class ElementsClawAgent:
    """
    ElementsClaw智能体核心类。
    将专有模型、文献检索、数据库查询、合成评估整合为自动化流程。
    """
    def __init__(self, elements_model, llm_backend):
        self.model = elements_model            # Elements专有模型
        self.llm = llm_backend                 # LLM决策引擎
        self.candidates: List[CandidateMaterial] = []
        self.verified_materials: List[CandidateMaterial] = []
        self.literature_db = {}                # 文献数据库缓存
        self.structure_db = {}                 # 结构数据库缓存
        
    def screen_crystals(self, crystal_data) -> List[CandidateMaterial]:
        """
        第一阶段：大规模筛选。
        240万晶体 → 快速过滤 → 保留top候选
        """
        print(f"[Stage 1] Screening {len(crystal_data)} crystals...")
        
        candidates = []
        batch_size = 1024
        
        for i in range(0, len(crystal_data), batch_size):
            batch = crystal_data[i:i+batch_size]
            
            # 模型推理
            with torch.no_grad():
                predictions = self.model(batch)
            
            # 提取预测结果
            for j, pred in enumerate(zip(
                predictions['superconducting_prob'],
                predictions['critical_temp'],
                predictions['energy_stability']
            )):
                prob, tc, energy = pred
                
                # 多条件过滤
                if (prob.item() > 0.8 and            # 超导概率高
                    tc.item() > 1.0 and               # 有超导转变
                    energy.item() < 50.0):            # 结构稳定
                    
                    candidates.append(CandidateMaterial(
                        formula=batch[j].formula,
                        space_group=batch[j].space_group,
                        predicted_tc=round(tc.item(), 2),
                        superconducting_prob=round(prob.item(), 4),
                        energy_above_hull=round(energy.item(), 2),
                        source="screening",
                    ))
        
        # 按超导概率排序，保留top 6.8万
        candidates.sort(key=lambda x: x.superconducting_prob, reverse=True)
        self.candidates = candidates[:68000]
        
        print(f"[Stage 1] Complete: {len(self.candidates)} candidates retained")
        return self.candidates
    
    def cross_validate_literature(self, candidates: List[CandidateMaterial]) -> List[CandidateMaterial]:
        """
        第二阶段：文献/数据库交叉验证。
        自动检索SuperCon等数据库，排除已有报道材料。
        """
        print(f"[Stage 2] Cross-validating {len(candidates)} candidates against literature...")
        
        validated = []
        for mat in candidates:
            # 查询已有数据库和文献
            is_known = self._check_database(mat.formula)
            literature_entries = self._search_literature(mat.formula)
            
            mat.literature_matched = is_known
            
            if not is_known:
                # 全新候选，计算文献支持度
                mat.synthesis_score = self._evaluate_synthesis(mat)
                validated.append(mat)
            elif is_known and self._has_structural_error(mat):
                # 已知但数据库构型有误，标记为"沉冤得雪"候选
                mat.literature_matched = False
                mat.source = "corrected_structure"
                validated.append(mat)
        
        print(f"[Stage 2] Complete: {len(validated)} novel/revised candidates")
        return validated
    
    def _check_database(self, formula: str) -> bool:
        """检查数据库中是否已有该材料"""
        # 模拟SuperCon数据库查询
        known_materials = {
            "YBa2Cu3O7", "Bi2Sr2CaCu2O8", "HgBa2Ca2Cu3O8",
            "MgB2", "Nb3Sn", "NbTi", "FeSe", "LaFeAsO1-xFx"
        }
        return formula in known_materials
    
    def _search_literature(self, formula: str) -> List[dict]:
        """检索文献中与化学式相关的信息"""
        # 模拟文献检索
        return []
    
    def _has_structural_error(self, mat: CandidateMaterial) -> bool:
        """检查数据库中该材料的结构是否正确"""
        # Elements模型预测的结构与数据库记录对比
        return False
    
    def _evaluate_synthesis(self, mat: CandidateMaterial) -> float:
        """
        可合成性评估。
        综合评估：元素丰度、合成温度、毒性、相稳定性等。
        """
        score = 0.5  # 基准分
        
        # 基于元素丰度调整
        abundant_elements = {'Hf', 'Zr', 'Re', 'Sc', 'Ti', 'V', 'Nb', 'Mo', 'W'}
        rare_elements = {'Ru', 'Rh', 'Pd', 'Os', 'Ir', 'Pt', 'Au'}
        
        elements = self._parse_elements(mat.formula)
        for elem in elements:
            if elem in abundant_elements:
                score += 0.1
            elif elem in rare_elements:
                score -= 0.1
        
        # 基于结构复杂度
        if mat.energy_above_hull < 10:
            score += 0.2
        
        return min(max(score, 0.0), 1.0)
    
    def _parse_elements(self, formula: str) -> set:
        """从化学式解析元素种类"""
        import re
        elements = set()
        pattern = r'([A-Z][a-z]?)'
        for match in re.finditer(pattern, formula):
            elements.add(match.group(1))
        return elements
    
    def design_experiment(self, mat: CandidateMaterial) -> dict:
        """
        第三阶段：实验方案设计。
        根据材料特性自动设计合成路线和表征方案。
        """
        return {
            "material": mat.formula,
            "synthesis_method": self._recommend_synthesis(mat),
            "sintering_temp": self._estimate_sintering_temp(mat),
            "characterization": ["XRD", "SQUID", "resistivity", "heat_capacity"],
        }
    
    def _recommend_synthesis(self, mat: CandidateMaterial) -> str:
        """推荐合成方法"""
        elements = self._parse_elements(mat.formula)
        if 'Re' in elements:
            return "arc_melting_high_temp"
        return "solid_state_reaction"
    
    def _estimate_sintering_temp(self, mat: CandidateMaterial) -> int:
        """估算烧结温度"""
        return 1200  # 摄氏度
    
    def self_evolve(self, new_data: List[dict]):
        """
        自我进化机制。
        当从文献中发现新数据时，自动微调模型参数。
        """
        print("[Self-Evolution] Fine-tuning on new literature data...")
        # 实际实现中会触发模型的增量训练
        # 这里模拟微调流程
        n_samples = len(new_data)
        print(f"[Self-Evolution] Ingested {n_samples} new data points")
        print("[Self-Evolution] Model updated. Ready for next screening cycle.")
        

# ─── 4种新超导体的发现路径模拟 ───
def simulate_discovery_pathways():
    """
    模拟ElementsClaw发现4种新超导体的完整路径。
    每条路径代表一种不同的材料发现策略。
    """
    print("=" * 70)
    print("ElementsClaw Discovery: 4 Pathways to New Superconductors")
    print("=" * 70)
    
    pathways = {
        "Hf21Re25": {
            "method": "漏网之鱼 (Missed Catch)",
            "tc": 2.5,
            "description": (
                "材料早已存在于理论数据库，但从未被实验合成。"
                "AI在文献-数据库交叉比对时捞出，发现其声子谱"
                "和电子结构均满足超导条件。"
            ),
            "key_insight": "数据库挖潜：已存在但被忽略的宝藏",
        },
        "Zr4VRe7": {
            "method": "沉冤得雪 (Structural Correction)",
            "tc": 3.5,
            "description": (
                "数据库记录的结构构型有误（空间群标注错误），"
                "导致此前判断为非超导。Elements预测了不同的"
                "晶体结构，实验验证果然超导。"
            ),
            "key_insight": "结构纠偏：数据库质量决定了发现上限",
        },
        "HfZrRe4": {
            "method": "无中生有 (De Novo Design)",
            "tc": 5.9,
            "description": (
                "AI从头设计：锁定Hf-Zr-Re三元体系后，使用"
                "Elements-G生成器产生全新晶体结构，自我验证超导性。"
                "材料不存在于任何已知数据库中。"
            ),
            "key_insight": "AI创造力：超越人类知识边界的发现",
        },
        "Zr3ScRe8": {
            "method": "举一反三 (Analogous Reasoning)",
            "tc": 6.5,
            "description": (
                "基于Hf-Zr-Re体系验证的超导相，AI总结出结构模体："
                "保留P6/mmm富Re六方框架，保持Re子晶格完整。"
                "将Hf替换为同族Sc元素，生成新结构。"
            ),
            "key_insight": "结构模体：举一反三的AI推理能力",
        },
    }
    
    for formula, info in pathways.items():
        print(f"\n{'─' * 60}")
        print(f"■ {formula} | Tc = {info['tc']}K")
        print(f"  方法: {info['method']}")
        print(f"  描述: {info['description']}")
        print(f"  关键洞察: {info['key_insight']}")
    
    return pathways


if __name__ == "__main__":
    simulate_discovery_pathways()

# ─── 晶体结构模体分析 ───
def analyze_structure_motif(formula: str, tc: float, family: str):
    """
    结构模体分析函数。
    验证"保留P6/mmm富Re六方框架"这一结构模体通则。
    """
    print(f"\n{'='*50}")
    print(f"Material: {formula} (Tc={tc}K, Family: {family})")
    print(f"{'='*50}")
    
    motifs = {
        "P6/mmm": {
            "present": True,
            "description": "六方晶系，最高对称性空间群",
            "role": "提供高对称性电子结构",
        },
        "Re_sub_lattice": {
            "present": "Re" in formula,
            "description": "Re子晶格完整性",
            "role": "Re的5d电子提供高DOS",
        },
        "transition_metal_mix": {
            "present": True,
            "description": "多种过渡金属混合",
            "role": "d电子杂化增强超导耦合",
        },
    }
    
    for motif, info in motifs.items():
        print(f"  ■ {motif:20s}: {'✓' if info['present'] else '✗'}  {info['role']}")
    
    if all(m['present'] for m in motifs.values()):
        print(f"\n  → 结构模体验证通过！{formula}符合超导候选结构通式")
    else:
        print(f"\n  → 部分模体不匹配，需进一步分析")


# 执行模体分析
materials = [
    ("Hf21Re25", 2.5, "missed"),
    ("Zr4VRe7", 3.5, "corrected"),
    ("HfZrRe4", 5.9, "generated"),
    ("Zr3ScRe8", 6.5, "analogous"),
]

for formula, tc, family in materials:
    analyze_structure_motif(formula, tc, family)