Fast Graph-based Higher-Order Clustering Statistics on the GPU

Resumo Técnico: Estatísticas de Agrupamento de Ordem Superior Baseadas em Grafos Rápidas na GPU

Enunciado do Problema
A distribuição espacial das galáxias codifica informações críticas sobre o campo de densidade primordial e a evolução gravitacional. Embora a função de correlação de dois pontos (2pCF) seja há muito tempo a ferramenta padrão para analisar essa distribuição, ela é insuficiente para capturar a estrutura não-Gaussiana gerada pela evolução gravitacional não-linear tardia e pelo viés (bias) das galáxias. Estatísticas de ordem superior, especificamente as funções de correlação de três pontos (3pCF) e quatro pontos (4pCF), fornecem informações complementares, como a quebra de degenerescências nos parâmetros de viés e a sondagem da não-Gaussianidade primordial e da violação de paridade. No entanto, computar essas estatísticas para catálogos cosmológicos de grande escala é computacionalmente proibitivo. Uma abordagem de força bruta escala como $O(N^k)$, e embora métodos baseados em árvores reduzam isso para $O(N m^{k-1})$ (onde $m$ é a contagem média de vizinhos), o custo permanece alto para $k=3$ e $k=4$ em escalas de levantamentos modernos (por exemplo, DESI). Além disso, os códigos existentes carecem de implementações diretas no espaço de configuração de 4pCFs conectados com decomposição de paridade ou de aceleração eficiente em GPU capaz de lidar com grafos que excedem a memória do dispositivo.

Metodologia
Este artigo apresenta a Versão 2 do GRAMSCI (GRAph Made Statistics for Cosmological Information), um algoritmo que utiliza uma abordagem de banco de dados de grafos construída sobre uma árvore $k$-d. A metodologia central envolve a construção de uma estrutura de adjacência compacta e ordenada a partir das coordenadas iniciais, que é então consultada para contar $N$-uplas sem revisitar os dados brutos.

A atualização introduz cinco avanços metodológicos principais:

1. Algoritmo de Merge-Walk: O loop interno original de busca binária, que incorria em um custo de $O(m \log m)$ por par hub-spoke, é substituado por um algoritmo de merge-walk. Ao explorar a ordem ordenada das listas de adjacência, o algoritmo realiza uma passagem linear para encontrar interseções de conjuntos, reduzindo a complexidade para $O(m)$.
2. 4pCF com Decomposição de Paridade: Um novo estimador separa o sinal da 4pCF em canais de paridade par ($\zeta_+$) e ímpar ($\zeta_-$). Isso é alcançado calculando o volume assinado dos tetraedros de galáxias. Para evitar a contaminação por ruído de ponto flutuante no canal de paridade ímpar, tetraedros com volumes degenerados (efetivamente zero) são deterministicamente atribuídos como peso zero no canal ímpar.
3. Estimativa de 4pCF Conectada: O código agora computa internamente a contribuição desconectada (produtos de 2pCFs) no mesmo grafo e agrupamento (binning). Isso permite a subtração direta do termo desconectado da 4pCF total para isolar o sinal conectado (trispectro), o que é essencial para estudar a não-Gaussianidade primordial.
4. Interface Python: Um wrapper leve permite que o motor Fortran seja chamado diretamente de arrays NumPy, lidando com E/S de arquivos e retornando objetos estruturados de resultados.
5. Porte para GPU com Suporte Out-of-Core: Todo o motor de consulta (3pCF, 4pCF e 4pCF de paridade) foi portado para GPUs NVIDIA usando OpenACC. A implementação utiliza um layout de matriz esparsa compacta (CSR) e emprega uma estratégia de mapeamento "gang-and-vector", onde um bloco de threads processa um hub e as vias vetoriais (vector lanes) lidam com os loops internos de vizinhos. Para lidar com grafos maiores que a memória do dispositivo (por exemplo, levantamentos de escala BAO), um esquema de tiling out-of-core divide os arrays de arestas em janelas, processando o grafo em blocos enquanto mantém a contagem exata.

Resultos Principais
Os autores testaram o novo código contra a implementação original em CPU e injeções analíticas:

• Aceleração Algorítmica: O kernel de merge-walk fornece um aumento de velocidade de $1,5\times$ a $2,0\times$ para a 3pCF e de $2,8\times$ a $3,8\times$ para a 4pCF em relação à abordagem de busca binária em um nó de CPU de 64 threads, com ganhos que aumentam conforme a escala máxima ($r_{max}$) e a contagem de vizinhos crescem.
• Desempenho em GPU: Em uma única GPU de consumo (NVIDIA RTX 3090 Ti), o código atinge acelerações de $2,6\times$ (3pCF) a $9\times$ (4pCF) comparado a um nó de CPU de 64 threads. O desempenho da GPU escala com a largura de banda da memória, em vez do rendimento de ponto flutuante.
• Escalabilidade: O esquema de tiling out-of-core permite o processamento de grafos que excedem em muito a memória do dispositivo. Um teste específico mediu uma 3pCF de $9 \times 10^9$ arestas em uma placa de 24 GB com apenas $\sim20\%$ de overhead em comparação com uma execução hipotética totalmente residente.
• Corretude: Os kernels de GPU reproduzem exatamente os resultados de referência da CPU em contagens agrupadas. A validação contra testes de injeção de sinal analítico para violação de paridade e testes nulos em simulações simétricas de paridade (Quijote) confirma que o estimador é não-viesado e lida corretamente com geometrias degeneradas.
• Aplicação: O código foi aplicado à amostra LRG do DESI DR1 (3,0 milhões de galáxias). A 3pCF e a 4pCF conectada medidas seguem as previsões do conjunto EZmock, com a 3pCF mostrando a característica BAO esperada e a 4pCF conectada isolando o sinal não-Gaussiano.

Significância
O artigo afirma que estas atualizações expandem significativamente a utilidade prática das estatísticas de ordem superior na cosmologia. Ao reduzir o custo computacional da 4pCF em quase uma ordem de magnitude em hardware de consumo, o código torna viável a avaliação de 4pCFs conectadas em toda a faixa BAO para grandes catálogos como o DESI, uma tarefa anteriormente limitada pela necessidade de milhares de realizações de mocks para estimativa de covariância. A inclusão de um estimador com decomposição de paridade fornece uma verificação cruzada direta no espaço de configuração para buscas de violação de paridade, independente de análises de harmônicos esféricos. Finalmente, a capacidade de computar 4pCFs conectadas internamente sem entradas externas garante a autossuficiência no agrupamento e nos conjuntos de pares, o que é crítico para isolar o sinal do trispectro não-Gaussiano. Os autores concluem que o código está disponível publicamente e pronto para aplicação em levantamentos atuais e futuros de grande escala da estrutura do universo.