PtyRANNOSAUR: Ptychography with Robust Artificial Neural Networks Optimized for Sub-Angstrom Accuracy and Ultrafast Reconstruction

Resumen Técnico de PtyRANNOSAUR

Planteamiento del Problema
La ptecografía electrónica es una técnica poderosa para recuperar imágenes de fase de objetos 2D o 3D a partir de conjuntos de datos de microscopía electrónica de barrido por transmisión (STEM) 4D, lo que permite una resolución sub-angstrom profunda y una alta eficiencia de dosis. Sin embargo, su adopción generalizada se ve actualmente obstaculizada por el inmenso costo computacional del procesamiento de datos. La reconstrucción de vanguardia se basa en resolvedores iterativos basados en la física (por ejemplo, PtyRAD) que utilizan diferenciación automática para optimizar los modelos de propagación hacia adelante de la dispersión de electrones. Aunque son precisos, estos métodos requieren cientos o miles de iteraciones, lo que a menudo toma horas en GPUs modernas (por ejemplo, NVIDIA A100). Esta latencia impide la retroalimentación de la muestra en tiempo real durante la adquisición de datos y limita el rendimiento experimental. Por el contrario, los métodos existentes de aprendizaje automático (ML) de disparo único ofrecen velocidad, pero no logran dar cuenta de la compleja física de la dispersión de electrones, las variaciones ópticas del microscopio o los errores de posición de escaneo, lo que resulta en una resolución y calidad inferiores a los métodos iterativos. Existe una necesidad insatisfecha de un método de reconstrucción que opere a velocidades cercanas al tiempo real manteniendo la precisión, la generalizabilidad y la resolución de los resolvedores iterativos a través de diversos sistemas de materiales y condiciones experimentales.

Metodología
Los autores presentan PtyRANNOSAUR, un marco de red neuronal basado en datos diseñado para reconstruir datos de ptecografía electrónica de resolución atómica en segundos. La metodología central consiste en entrenar autoencoders convolucionales especializados para mapear datos 4D-STEM (específicamente, una cuadrícula de 5×5 de patrones de difracción de haz convergente) directamente a imágenes de fase proyectadas en 2D.

• Generación de Datos de Entrenamiento: En lugar de utilizar imágenes genéricas o simulaciones simples, los autores generaron una base de datos masiva de 6.4 millones de pares entrada-objetivo. Los objetivos son imágenes de fase proyectadas en 2D (3 Å × 3 Å) derivadas de 16,000 estructuras cristalinas estables obtenidas del Materials Project. Las entradas son conjuntos correspondientes de 25 patrones CBED simulados generados mediante simulaciones de dispersión electrónica de multislice en abTEM.
• Simulación Informada por la Física: Los datos de entrenamiento incorporan explícitamente fenómenos físicos complejos, incluyendo la dispersión múltiple (vía el método multislice), la coherencia parcial espacial (modelada mediante convolución Gaussiana para simular el tamaño de la fuente), y variaciones en los parámetros experimentales como el voltaje de aceleración, el ángulo de convergencia, el desenfoque y el espesor de la muestra. Esto asegura que la red aprenda el mapeo inverso para condiciones experimentales realistas e imperfectas en lugar de escenarios idealizados.
• Especialización del Modelo: Se entrenaron tres modelos distintos para cubrir diferentes regímenes experimentales:
  • Modelo 2: Optimizado para muestras ultra delgadas (por ejemplo, materiales 2D) a 80 kV con un espesor de 2 ± 0.6 nm.
  • Modelo 10 y Modelo 20: Optimizados para láminas estándar de TEM a 300 kV con espesores nominales de 10 ± 2 nm y 20 ± 2 nm, respectivamente.
• Función de Pérdida: Las redes se entrenan utilizando el Índice de Similitud Estructural (SSIM) como función de pérdida, la cual cuantifica las similitudes de luminancia, contraste y estructura entre la salida de la red y la imagen de fase real (ground truth).
• Inferencia y Ensamblaje (Stitching): Durante la inferencia, la red procesa parches locales superpuestos de los datos 4D-STEM. Estos parches se ensamblan para formar la imagen completa. Para abordar el desplazamiento de la etapa y los errores de posición de la sonda comunes en los datos experimentales, los autores implementaron un módulo opcional de corrección de posición mediante correlación cruzada de fase, que se ejecuta en aproximadamente 20 segundos en una CPU.

Contribuciones Clave

1. Reconstrucción Ultrarrápida: PtyRANNOSAUR logra velocidades de reconstrucción de aproximadamente 5 segundos por conjunto de datos en una sola GPU NVIDIA A100, lo que representa una aceleración de 10 a 100 veces respecto a los métodos iterativos estándar (que toman ~30 minutos).
2. Robustez ante la Variabilidad Experimental: Los modelos demuestran una tolerancia significativa a las variaciones de desenfoque (±8 nm), tamaño de paso de escaneo (±0.12 Å), espesor de la muestra (±10 nm) y ángulo de convergencia (±2 mrad), eliminando la necesidad de la sintonización manual de hiperparámetros que típicamente requieren los resolvedores iterativos.
3. Manejo de la Física Compleja: A diferencia de los enfoques previos de ML de disparo único, PtyRANNOSAUR contabiliza con éxito la dispersión múltiple y la coherencia parcial espacial, permitiendo reconstrucciones precisas de muestras más gruesas (hasta ~25 nm) y geometrías complejas.
4. Accesibilidad: Los autores proporcionan modelos entrenados a través de una interfaz de Jupyter notebook amigable para el usuario, lo que permite realizar reconstrucciones sin simular datos ni entrenar nuevos modelos, ofreciendo también la capacidad de entrenar modelos personalizados para configuraciones específicas de microscopio.

Resultados
Los autores validaron PtyRANNOSAUR tanto en conjuntos de datos experimentales como simulados:

• Desempeño Experimental: El método fue probado en diversos especímenes incluyendo zafiro masivo, multicapas de Al/AlOx/Al, nanopartículas de oro sobre carbono amorfo y MoS2 de bicapa retorcida. En todos los casos, PtyRANNOSAUR produjo imágenes de resolución atómica comparables, y en algunos casos superiores, en calidad a las reconstrucciones iterativas de PtyRAD. Notablemente, para el MoS2 de bicapa retorcida, la red resolvió columnas atómicas espaciadas 0.4 Å, logrando una resolución de < 0.5 Å, la cual era apenas resuelta en la reconstrucción iterativa.
• Detección de Defectos: La red identificó con éxito defectos de átomo único, como vacantes de azufre en MoS2, sin requerir entrenamiento específico del material.
• Validación Simulada: En un conjunto de validación de 2,950 materiales no vistos, la pérdida mediana de SSIM a nivel de material cayó a 0.020 tras el ensamblaje, donde la red reprodujo fielmente las estructuras de red periódicas y las intensidades relativas.
• Robustez: Los modelos mantuvieron una alta precisión a través de los rangos de parámetros probados. Aunque el desempeño degradó fuera de estos rangos (por ejemplo, desenfoque extremo o espesor), la red aún capturó la estructura gruesa, actuando como un diagnóstico rápido para problemas de calidad de datos.

Significancia y Reivindicaciones
El artículo afirma que PtyRANNOSAUR permite "reconstrucciones de ptecografía electrónica de vanguardia casi en vivo". Al reducir el cuello de botella computacional de horas a segundos, el marco permite la retroalimentación en tiempo real durante la adquisición de datos, permitiendo a los usuarios localizar características de interés y optimizar los parámetros de adquisición sobre la marcha. Los autores posicionan la herramienta no solo como un método de reconstrucción independiente, sino también como un generador de priors para acelerar los resolvedores iterativos o como una herramienta de cribado rápido para grandes conjuntos de datos (por ejemplo, experimentos in-situ o tomografía). El trabajo demuestra que el aprendizaje profundo puede aprender eficazmente el problema inverso no lineal y complejo de la ptecografía electrónica, logrando una resolución profunda sub-angstrom (< 0.5 Å) con robustez ante imperfecciones experimentales, democratizando así el acceso a la microscopía de alta resolución.