cTreeBalls: a fast 3-point correlation function code for clustering measurements

Resumo Técnico: cTreeBalls – Um Código Rápido para Funções de Correlação de 3 Pontos

1. O Problema
As funções de correlação são ferramentas fundamentais na cosmologia para extrair informações sobre a estrutura em grande escala do universo e discriminar entre modelos cosmológicos. Enquanto estatísticas de segunda ordem (como a função de correlação de dois pontos, 2PCF) caracterizam campos gaussianos aleatórios, a evolução não-linear da matéria e não-gaussianidades primordiais geram desvios que exigem estatísticas de ordem superior, especificamente a função de correlação de três pontos (3PCF) ou o espectro de bis.

O desafio principal reside na complexidade computacional. Um estimador ingênuo da 3PCF requer $O(N^3)$ avaliações, onde $N$ é o número de pontos de dados (galáxias ou pixels). Para levantamentos futuros como o do Observatório Vera C. Rubin (LSST), que prevê mais de 20 bilhões de objetos, esse escalonamento torna a análise proibitiva em termos de tempo de CPU. Mesmo o uso de estruturas de árvore (como kd-tree ou octree) para busca de vizinhos, que reduz a complexidade para $O(N^2 \log N)$, ainda é insuficiente para lidar com a escala de dados do LSST.

2. Metodologia
O código cTreeBalls (ou cBalls) foi desenvolvido para superar essas limitações através de uma combinação de otimizações algorítmicas e matemáticas:

• Decomposição Harmônica e Isotropia: Em vez de calcular a 3PCF diretamente, o código projeta a informação da correlação em uma base ortogonal (decomposição harmônica de um dos ângulos nas configurações triangulares). Isso permite que os coeficientes escalem como contagens de pares, $O(N \log N)$, assumindo que um número finito e pequeno de autovetores captura a maior parte da informação.
• Estruturas de Árvore Otimizadas: O código utiliza estruturas de octree (Barnes & Hut, 1986) e kd-tree (Bentley, 1975) para organizar espacialmente os dados.
  • Poda de Células (Oc-ball-tree): As células da octree são "apertadas" (podadas) para criar uma estrutura semelhante a uma "ball-tree", visitando apenas regiões preenchidas com galáxias e ignorando grandes espaços vazios (vazios).
  • Compartilhamento de Listas de Vizinhos: O algoritmo identifica que pivôs (pontos de referência) espacialmente próximos compartilham listas de vizinhos. O código define automaticamente um "raio de pivô" baseado na separação média entre galáxias para reutilizar esses cálculos.
• Avaliação Polinomial Rápida: Para evitar o custo computacional de funções trigonométricas ($\sin$ e $\cos$) durante o preenchimento dos bins de multipolos, o código substitui essas avaliações por polinômios de Chebyshev, que são significativamente mais rápidos.
• Paralelização: O motor de busca em C é paralelizado usando threads OpenMP, distribuindo os pixels pivô entre múltiplos núcleos de CPU.
• Interface e Flexibilidade: O pacote possui uma interface Python (via Cython) que permite o uso em terminais Linux ou Jupyter Notebooks. Ele suporta diversos formatos de entrada (CFITSIO, HEALPix, GADGET-2, ASCII) e pode calcular estatísticas tanto em caixas cúbicas com condições de contorno periódicas (apenas 2PCF) quanto em esferas unitárias (para campos projetados como lentes gravitacionais).

3. Principais Contribuições

• Desempenho Extremo: O código consegue calcular correlações de 3 pontos para mais de 200 milhões de pixels (simulação de céu completo com $N_{side}=4096$) em menos de 10 minutos em um único nó de computação de alto desempenho.
• Escalabilidade: Demonstra um escalonamento de tempo de execução próximo a $O(N^{1.1})$ para uma faixa de resoluções HEALPix ($N_{side}$ de 256 a 4096), tornando viável a análise dos dados futuros do LSST (bilhões de objetos).
• Complementaridade: Oferece uma abordagem complementar a outros códigos existentes (como TreeCorr e Triumvirate), utilizando uma decomposição específica que pode oferecer poder de restrição diferente sobre parâmetros cosmológicos.
• Correções de Borda: Implementa correções de borda (baseadas em Slepian & Eisenstein, 2015) para cálculos de multipolos de 3PCF em volumes que não cobrem o céu inteiro.
• Integração em Fluxos de Trabalho: Projetado para ser integrado em pipelines massivos como o TXPipe do LSST, com suporte a logs detalhados, tratamento de exceções e arquivos de parâmetros flexíveis.

4. Resultados de Desempenho

• Simulação de Céu Completo: Para uma simulação com $N_{side} = 8192$ (aproximadamente 800 milhões de pixels) e um raio de busca máximo de 200 minutos de arco, o tempo de processamento foi de cerca de 40 minutos de CPU (tempo de parede) utilizando 128 threads em um nó Perlmutter-NERSC.
• Dependência de Multipolos: O tempo de execução em função do número máximo de multipolos ($m$) segue uma relação linear $t(m) = 1.85m + 63.7$ segundos (para $N_{side}=1024$). O aumento do número de multipolos tem um impacto moderado; calcular até $m=8$ é apenas ~20% mais lento do que calcular até $m=1$.
• Custo Dominante: Para volumes massivos, o tempo de CPU é gasto principalmente na varredura da árvore e na atualização dos histogramas somando as listas de vizinhos, e não na avaliação de polinômios.

5. Significado
O cTreeBalls preenche uma lacuna crítica na análise de dados cosmológicos de próxima geração. Ao reduzir a complexidade da 3PCF de $O(N^3)$ para algo próximo a $O(N \log N)$, ele torna viável a extração de informações não-gaussianas de levantamentos de galáxias massivos como o LSST. Isso permitirá aos cosmólogos testar modelos de energia escura e matéria escura com uma precisão sem precedentes, utilizando estatísticas de ordem superior que antes eram computacionalmente inacessíveis para dados de escala de bilhões de objetos. O código é uma ferramenta essencial para a colaboração de Ciência da Energia Escura (DESC) do Observatório Rubin.