VesNet: Neural network accelerated solver for simulating Stokesian vesicle suspensions

Resumen Técnico: VesNet: Solucionador acelerado por redes neuronales para la simulación de suspensiones de vesículas estocásticas

Planteamiento del Problema
La simulación numérica de suspensiones de partículas deformables, específicamente vesículas, en el régimen de flujo de Stokes es computacionalmente prohibitiva. Estos sistemas se caracteran por interacciones fluido-estructura fuertemente no lineales, interfaces en evolución y fenómenos hidrodinámicos multiescala que involucran efectos de cuerpo múltiple de largo alcance y fuerzas de lubricación de corto alcance. Los métodos tradicionales de alta fidelidad, como el Método de la Ecuación de Integral de Frontera (BIEM), requieren resolver grandes sistemas lineales en cada paso de tiempo y gestionar integrales casi singulares costosas, lo que dificulta su escalabilidad para suspensiones densas o horizontes temporales largos. Si bien el aprendizaje automático (ML) se ha aplicado a flujos de partículas, la mayoría de los enfoques existentes se centran en partículas rígidas o actúan como modelos sustitutos puros que carecen de la fidelidad física necesaria para cuerpos deformables con un acoplamiento hidrodinámico complejo.

Metodología: Marco de Trabajo VesNet
Los autores presentan VesNet, un solucionador híbrido que integra redes neuronales (NN) con módulos algorítmicos convencionales para acelerar las simulaciones de vesículas en dos dimensiones. En lugar de reemplazar el solucionador físico completo, VesNet utiliza NN para aproximar operadores integro-diferenciales específicos y computacionalmente costosos, manteniendo algoritmos estándar para la reparametrización de la frontera e interacciones de campo lejano.

• Ecuaciones de Gobierno y División de Operadores: El solucionador aborda las ecuaciones acopladas de interacción fluido-sólido para membranas inextensibles que resisten la flexión. El esquema de pasos de tiempo emplea la división de operadores para desacoplar los pasos de advección (impulsados por el flujo de fondo y las interacciones entre vesículas) y de relajación (impulsados por la elasticidad de la membrana).
• Aproximación de Operadores mediante Redes Neuronales: La innovación central radica en la aproximación de los operadores no lineales utilizando la Red de V-forma Iterada (IVNet), una arquitectura convolucional multinivel similar a U-Net.
  • Autointeracciones: La IVNet aproxima el mapa de flujo para el paso de relajación y el cálculo de la tensión, los cuales dependen únicamente de la configuración de una sola vesícula.
  • Campo Lejano y Advección: Para asegurar la generalización a condiciones de flujo arbitrarias, el campo de velocidades se descompone en una serie de Fourier truncada. Subredes separadas se entrenan para mapear las configuraciones de las vesículas con la respuesta de cada modo de Fourier, permitiendo que el solucionador maneje flujos de fondo y conteos de partículas no vistos sin necesidad de reentrenamiento.
  • Integración de Casi-Singularidades: Un novedoso esquema basado en NN reemplaza la cuadratura singular tradicional para las interacciones de campo cercano. El método construye cinco capas concéntricas alrededor de la membrana de la vesícula (incluyendo la propia membrana y capas en $\pm h/2$ y $\pm h$). Las NN aproximan la integral de capa única en estas capas, y se utiliza la interpolación de Funciones de Base Radial (RBF) para recuperar las velocidades en puntos arbitrarios del campo cercano. Esto evita la necesidad de una evaluación directa de integrales singulares costosas durante la simulación.
• Invarianzas Físicas y Entrenamiento: Las redes se entrenan en simulaciones de horizonte temporal corto de vesículas individuales en diversas condiciones de flujo (cizalladura, Poiseuille, cilindros rotatorios). Para asegurar la generalización, los datos de entrada se estandarizan para la invarianza de traslación, rotación y escala. Las redes se entrenan para minimizar el error relativo tanto en la solución como en sus derivadas para mejorar la suavidad.
• Implementación: Todo el flujo de trabajo, incluyendo los pasos del solucionador y la inferencia de la red, está implementado en GPUs utilizando PyTorch. El código utiliza el procesamiento por lotes (batching), operaciones de tensores vectorizados y álgebra lineal por lotes (por ejemplo, descomposición de Cholesky) para maximizar el paralelismo.

Contribuciones Clave

1. Arquitectura de Solucionador Híbrido: VesNet combina la precisión del BIEM con la velocidad del ML al utilizar NN específicamente para los operadores más costosos (integrales de casi-singularidad y autointeracciones) mientras retiene algoritmos exactos para las interacciones de campo lejano y restricciones geométricas.
2. IVNet para Flujo de Stokes: La aplicación de Redes de V-forma Iteradas para aproximar operadores no lineales en flujos de partículas de Stokes, demostrando su eficacia para capturar la naturaleza multiescala y las estructuras convolucionales inherentes a la solución fundamental.
3. Integración Neuronal de Casi-Singularidades: Una nueva estrategia para manejar interacciones de campo cercano que elimina la necesidad de reglas de cuadratura singular durante la simulación, reemplazándolas con aproximaciones de NN e interpolación RBF.
4. Generalización mediante Descomposición de Fourier: Un método para desacoplar la red de condiciones de flujo específicas mediante el entrenamiento en los modos de Fourier del campo de velocidades, permitiendo que el solucionador maneje flios de fondo arbitrarios y densidades de partículas sin reentrenamiento.
5. Rendimiento Acelerado por GPU: Una implementación totalmente acelerada por GPU que logra aceleraciones significativas sobre versiones de CPU y GPU de los mejores solucionadores BIEM actuales.

Resultados y Rendimiento
Los autores evalúan VesNet frente a simulaciones de BIEM de verdad fundamental (ground-truth) en varios escenarios:

• Relajación: VesNet recupera con precisión la forma de equilibrio y la energía de flexión de una vesícula deformada en un fluido estacionario, incluso para configuraciones no presentes en el conjunto de datos de entrenamiento.
• Flujo de Poiseuille: El solucionador predice correctamente la migración lateral y las configuraciones de equilibrio de las vesículas bajo variaciones de confinamiento y números capilares. VesNet con $N=32$ puntos de discretización logra una precisión comparable o superior al BIEM con la misma resolución, y se aproxima a la precisión del BIEM con $N=128$.
• Flujo de Cizalladura (Pares de Vesículas): En simulaciones de dos vesículas interactuando, se demuestra que el esquema de integración de casi-singularidad propuesto es esencial para la estabilidad. Sin él, ocurren colisiones erróneas; con la corrección de la NN, VesNet captura con precisión la distancia de separación y evita colisiones, coincidiendo con los resultados de BIEM de alta resolución.
• Suspensiones Densas: Simulaciones de 2,000 vesículas en flujos de Taylor-Green y Poiseuille demuestran que VesNet reproduce observables clave (por ejemplo, la formación de la capa libre de células, estadísticas de migración) con alta fidelidad.
• Aceleración: VesNet ofrece una aceleración de 100$\times$ respecto a una implementación de CPU multihilo de un solucionador BIEM totalmente resuelto en MATLAB y una aceleración de 5$\times$ respecto a una implementación en GPU del mismo BIEM. Un solo paso de tiempo para 2,000 vesículas toma 1.4 segundos en una NVIDIA GH200.

Significancia y Limitaciones
El artículo sostiene que VesNet proporciona una herramienta robusta para la exploración de parámetros, optimización y cuantificación de la incertidumbre en suspensiones de vesículas, escenarios donde las simulaciones tradicionales suelen ser demasiado lentas. Al lograr una precisión de nivel BIEM con resoluciones espaciales y temporales reducidas, VesNet permite realizar simulaciones extensas con recursos computacionales relativamente modestos.

Los autores señalan explícitamente las limitaciones actuales:

• El marco está restringido a simulaciones en 2D.
• Actualmente no acomoda contrastes de viscosidad entre los fluidos interior y exterior de la vesícula.
• No puede garantizar interacciones libres de colisiones en todos los regímenes, un desafío que sigue siendo difícil incluso para el BIEM estándar.
• El solucionador asume que el fluido dentro de la vesícula es idéntico al fluido de suspensión.

Este trabajo establece un marco fundacional para extender estos enfoques híbridos de ML-física a sistemas en 3D y aplicaciones biológicas e industriales más complejas.