TUNeS: Neural Emulation of Large-Scale Structure Across Redshifts

Título: TUNeS: Emulação Neural da Estrutura em Grande Escala através de Redshifts

1. Problema e Motivação

A exploração de dados de precisão de novos levantamentos cosmológicos exige grandes conjuntos de simulações numéricas para prever teorias, estimar matrizes de covariância e caracterizar incertezas sistemáticas. As simulações de corpo N (N-body) são a ferramenta padrão para modelar a evolução gravitacional não linear da estrutura em grande escala. No entanto, elas apresentam um custo computacional proibitivo, especialmente para simulações de alta resolução que exigem grandes volumes e muitas partículas. Métodos aproximados (como COLA ou FastPM) reduzem o custo, mas perdem precisão em regimes profundamente não lineares e em pequenas escalas. O objetivo deste trabalho é desenvolver uma alternativa baseada em aprendizado de máquina que acelere drasticamente essas simulações sem sacrificar a precisão estatística.

2. Metodologia: O Framework TUNeS

Os autores introduzem o TUNeS (Temporal UNet emulator for Structure formation), um framework de rede neural projetado para prever a evolução não linear do campo de densidade de matéria a partir de uma distribuição inicial de partículas. O modelo opera diretamente sobre dados de partículas, sem depender de condições iniciais baseadas em grades ou deslocamentos de perturbação linear.

A arquitetura adota uma estratégia de dois estágios:

• Estágio 1 (Inferência Baseada em Partículas):

  • Utiliza uma rede MLP (Perceptron Multicamada) residual com modulação linear de características (FiLM).
  • Entrada: Posições e velocidades iniciais das partículas, além do par de redshifts (inicial e alvo).
  • Função: Prever deslocamentos de partículas ($\Delta x$) para capturar a evolução aproximadamente linear e de grande escala.
  • Técnica: As posições são codificadas com recursos de Fourier e o tempo é incorporado via fatores de escala logarítmicos ($\ln a$), modulando a rede através de FiLM.
  • Saída: Uma distribuição de partículas "grossa" (coarse) no redshift alvo.

• Estágio 2 (Refinamento do Campo de Densidade):

  • As partículas do Estágio 1 são mapeadas para um campo de densidade 3D em uma grade cartesiana usando o esquema de atribuição de massa "Cloud-in-Cell" (CIC).
  • O campo é decomposto em janelas espaciais sobrepostas.
  • Arquitetura: Uma rede U-Net 3D (convolucional) refina o campo de densidade bruto.
  • Função: Modelar a evolução não linear e estruturas de pequena escala.
  • Condição de Redshift: A dependência do redshift é incorporada explicitamente na rede U-Net através de modulação FiLM 3D baseada nos fatores de escala logarítmicos.
  • Saída: Campos de densidade refinados que são costurados (stitched) para formar o volume completo.

Função de Perda:

• Estágio 1: Erro quadrático médio (MSE) nos deslocamentos das partículas.
• Estágio 2: MSE ponderado espacialmente nos pixels da janela, dando mais peso ao centro da janela para minimizar artefatos nas bordas.

3. Contribuições Principais

• Arquitetura Híbrida: Combina a interpretabilidade física do movimento de partículas (Estágio 1) com a capacidade representacional de redes convolucionais para estruturas densas (Estágio 2).
• Flexibilidade de Redshift: O modelo pode, em princípio, ser inicializado em redshifts intermediários arbitrários, não estando intrinsicamente ligado a condições iniciais de redshifts extremamente altos (embora o teste atual use $z=100$).
• Escalabilidade: O uso de processamento em lotes (batching) de partículas e janelas sobrepostas permite que o modelo escale para volumes maiores e geometrias irregulares sem aumentar os requisitos de memória da GPU.
• Eficiência: Treinado com apenas 8 simulações de N-body, o modelo é capaz de generalizar bem para estatísticas complexas.

4. Resultados e Avaliação

O modelo foi avaliado comparando suas previsões com simulações de referência de N-body (Gadget-4) em um volume cúbico de $500 \, h^{-1}\text{Mpc}$ com $256^3$ partículas.

• Desempenho Computacional:

  • Inferência de ponta a ponta (de $256^3$ partículas para uma grade de densidade $256^3$) leva ~25 segundos em uma única GPU (NVIDIA RTX 4090).
  • Uso de memória de pico: ~4 GB.

• Estatísticas de Ponto Único (One-point):

  • A distribuição de probabilidade (PDF) do campo de densidade sobre-densa ($\delta$) coincide bem com a referência, com erros relativos abaixo de 2% na faixa de densidade intermediária.

• Estatísticas de Dois Pontos (Two-point):

  • Espectro de Potência: O erro no espectro de potência permanece no nível de alguns por cento até $k \simeq 1 \, h \text{Mpc}^{-1}$.
  • Correlação de Fase: A "estocasticidade" ($1-r^2$) indica alinhamento de fase excelente, com erros abaixo de 5% para $z \gtrsim 1$ e até ~15% em redshifts mais baixos em pequenas escalas.

• Estatísticas de Picos (Peak Statistics):

  • A abundância de picos locais no campo de densidade é reproduzida com erro relativo abaixo de 10% para sinais de ruído ($S/N$) até 3.

• Funcionais de Minkowski 3D:

  • O modelo reproduz com sucesso descritores morfológicos e topológicos (volume, área superficial, curvatura média integrada e característica de Euler), capturando a evolução não gaussiana da estrutura. As discrepâncias são tipicamente uma ordem de magnitude menores que os valores absolutos dos funcionais.

5. Significado e Conclusão

O TUNeS representa um avanço significativo na emulação de simulações cosmológicas. Ao oferecer uma aceleração de várias ordens de magnitude em relação às simulações de N-body tradicionais, mantendo alta precisão em estatísticas gaussianas e não gaussianas, o modelo torna viável a geração de grandes ensembles de simulações para análises de lentes gravitacionais fracas e inferência de parâmetros cosmológicos.

A capacidade de operar diretamente sobre distribuições de partículas e a flexibilidade para iniciar em redshifts intermediários abrem novas possibilidades para reconstruções baseadas em levantamentos observacionais e simulações de cone de luz. O trabalho sugere que, com mais dados de treinamento e maior resolução, a precisão em escalas ainda menores pode ser melhorada, consolidando o aprendizado de máquina como uma ferramenta essencial para a cosmologia de precisão futura.