# Geometric deep learning for galaxy-halo connection: a case study for galaxy intrinsic alignments

Yesukhei Jagvaral, François Lanusse, Rachel Mandelbaum

► To cite this version:

Yesukhei Jagvaral, François Lanusse, Rachel Mandelbaum. Geometric deep learning for galaxy-halo connection: a case study for galaxy intrinsic alignments. Monthly Notices of the Royal Astronomical Society, 2025, 542 (3), pp.2560-2571. 10.1093/mnras/staf592. hal-05412124

HAL Id: hal-05412124

https://hal.science/hal-05412124v1

Submitted on 11 Dec 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>
# Geometric deep learning for galaxy-halo connection: a case study for galaxy intrinsic alignments

Yesukhei Jagvaral<sup>1,2,4*</sup>①, François Lanusse<sup>3,4</sup>①, Rachel Mandelbaum<sup>1,2</sup>①

<sup>1</sup>McWilliams Center for Cosmology and Astrophysics, Department of Physics, Carnegie Mellon University, Pittsburgh, PA 15213, USA

<sup>2</sup>NSF AI Planning Institute for Data-Driven Discovery in Physics, Carnegie Mellon University, Pittsburgh, PA 15213, USA

<sup>3</sup>AIM, CEA, CNRS, Université Paris-Saclay, Université Paris Diderot, Sorbonne Paris Cité, F-91191 Gif-sur-Yvette, France

<sup>4</sup>Simons Foundation, Flatiron Institute, New York, NY 10010, USA

Accepted XXX. Received YYY; in original form ZZZ

## ABSTRACT

Forthcoming cosmological imaging surveys, such as the Rubin Observatory LSST, require large-scale simulations encompassing realistic galaxy populations for a variety of scientific applications. Of particular concern is the phenomenon of intrinsic alignments (IA), whereby galaxies orient themselves towards overdensities, potentially introducing significant systematic biases in weak gravitational lensing analyses if they are not properly modeled. Due to computational constraints, simulating the intricate details of galaxy formation and evolution relevant to IA across vast volumes is impractical. As an alternative, we propose a Deep Generative Model trained on the IllustrisTNG-100 simulation to sample 3D galaxy shapes and orientations to accurately reproduce intrinsic alignments along with correlated scalar features. We model the cosmic web as a set of graphs, each graph representing a halo with nodes representing the subhalos/galaxies. The architecture consists of a SO(3) × R<sup>n</sup> diffusion generative model, for galaxy orientations and *n* scalars, implemented with E(3) equivariant Graph Neural Networks that explicitly respect the Euclidean symmetries of our Universe. The model is able to learn and predict features such as galaxy orientations that are statistically consistent with the reference simulation. Notably, our model demonstrates the ability to jointly model Euclidean-valued scalars (galaxy sizes, shapes, and colors) along with non-Euclidean valued SO(3) quantities (galaxy orientations) that are governed by highly complex galactic physics at non-linear scales.

**Key words:** methods: numerical – cosmology: theory – galaxies: statistics – galaxies: structure – gravitational lensing: weak

## 1 INTRODUCTION

Ongoing and future cosmological weak lensing surveys, including the European Space Agency's Euclid Mission<fnref n="1" />, the Vera C. Rubin Observatory Legacy Survey of Space and Time<fnref n="2" />, and the Nancy Grace Roman Space Telescope<fnref n="3" /> High Latitude Imaging Survey, promise a wealth of data for precise cosmological parameter measurements (Kilbinger 2015). However, extracting the full cosmological information from these datasets requires robust and scalable analysis methods that correct for systematic biases and model systematic uncertainties. An important systematic uncertainty, sourced by coherent galaxy alignments with their neighbors and the large scale structure, is the intrinsic alignment (IA) of galaxies (e.g., for a review see Troxel & Ishak 2015). This effect, if not accounted for, can falsely imitate weak gravitational lensing signals and bias cosmological analyses (Mandelbaum 2018). Cosmological simulations are often employed to derive synthetic galaxy catalogs for testing, improving, and validating analysis pipelines and their robustness to systematic uncertainties (DeRose et al 2019, 2021; Euclid Collaboration et al. 2024); and in future analyses, they may be used to constrain cosmology via Simulation-Based Inference (SBI; Jeffrey et al. 2021; Lu et al. 2023). As future surveys will cover large areas of the sky and produce high-quality data, cosmological simulations of increasingly larger volumes and higher resolutions, along with realistic systematics, like intrinsic alignments, are needed (Vogelsberger et al. 2020).

Cosmological simulations with galaxy evolution are computationally intractable at the large scales and high resolutions demanded by future surveys (Vogelsberger et al. 2020). One common approach for producing synthetic galaxy catalogs is to run N-body simulations and then “paint” in the galaxies using semi-analytic models (Benson 2012; Lagos et al. 2018; De Lucia 2019). Though these methods are less computationally costly than running a full hydrodynamical simulation, they are still costly and introduce tunable parameters. The multivariate distribution describing various galaxy and dark matter (sub)halo properties is called the galaxy-halo connection. Even though halo occupation distribution (HOD) and subhalo abundance matching (SHAM) models have been shown to predict galaxy mass, abundance and clustering with some success (Somerville & Davé 2015; Guo et al. 2016), recent studies show that the full galaxy-halo connection is highly non-linear and high dimensional (see, e.g., Wechsler & Tinker 2018, for a comprehensive review).

Recent advances in machine learning have sparked a revolution in

<footnote>* E-mail: yjagvara@andrew.cmu.edu</footnote>

<footnote n="1">¹ https://www.euclid-ec.org/</footnote>

<footnote n="2">² https://www.lsst.org/</footnote>

<footnote n="3">³ https://roman.gsfc.nasa.gov/</footnote>
astrophysics and cosmology, offering promising avenues for tackling complex challenges in modeling the universe (for a recent review, see Huertas-Company & Lanusse 2023). Relevant previous attempts to harness machine learning for the problem of intrinsic alignments considered in this work include simple and limited implementations with generative adversarial networks (GANs) (Jagvaral et al. 2022a); novel approaches with geometric deep learning methods are emerging to overcome the hurdles identified in that work (Jagvaral et al. 2024). Geometric deep learning methods have garnered a lot of attention with their ability to handle complex datasets that do not fit a grid-like structure (e.g., images) or datasets that do not live in a Euclidean space (Bronstein et al. 2021). Geometric deep learning is an area of research that seeks to exploit symmetries and structures that are inherent in the data using mathematical graphs, group theory and other geometric concepts. As such, these methods allow us to model galaxy intrinsic alignments and orientations, while respecting the inherent symmetry of the problem.

At small scales, non-linear gravitational evolution and baryonic physics drive the physics of galaxies, making it difficult to model and capture the dependencies and correlations among galaxy and dark matter subhalo properties (Weinberg et al. 2015; Del Popolo & Le Delliou 2017). Fortunately, graph neural networks (GNNs), based on mathematical graphs, are a natural way to model galaxies, given their ability to model sparsely distributed objects and capture correlations among neighbors (Duval et al. 2023). Additionally, galaxy orientations belong to the Lie group SO(3), requiring these data to be modeled on the corresponding manifold.

Recently, diffusion-based generative models have become the state-of-the-art for various generative tasks, such as images, audio, video and molecules (Yang et al. 2023). Modeling distributions on the manifold of 3D rotations is, however, a non-trivial task, and to address this problem we developed a new type of score-based diffusion model specifically for the SO(3) manifold in Jagvaral et al. (2024), by extending the Euclidean framework introduced in Song et al. (2021). Diffusion models are flexible in their ability to model data that live in various different spaces (e.g., scalars and manifold-valued data) compared to normalizing flows (Bond-Taylor et al. 2021). Additionally, diffusion models are more stable compared to Generative Adversarial Networks (Bond-Taylor et al. 2021).

Based on these developments, we propose to build a conditional generative model for joint galaxy properties, using diffusion-based models in Euclidean space (Song & Ermon 2020) for scalars; and diffusion-based models in SO(3) (Jagvaral et al. 2024) for galaxy orientations. Moreover, the diffusion model is implemented with a GNN architecture that is used to jointly model the correlated galaxy properties within halos with E(3) equivariance (Satorras et al. 2021) to exploit global symmetries. In this regard, generative machine learning approaches combined with state-of-the-art geometric deep learning methods could serve as robust emulators for synthetic galaxy catalog production. In particular, they could be used to populate realistic galaxies in large volume dark matter-only simulations.

The paper is organized as follows: in § 2 we briefly introduce weak lensing and IA. In § 3 we describe the problem specifics. In § 5, we describe the cosmological simulation and the astrophysical observables. We describe our methodology in § 4. In § 6 we present the results, and finally in § 7 we conclude our paper and propose future work.

## 2 GRAVITATIONAL WEAK LENSING AND INTRINSIC ALIGNMENTS

Coherent galaxy shape distortions caused by gravitational lensing due to large-scale structure are referred to as cosmic shear (Kilbinger 2015). In gravitational weak lensing studies, the observed galaxy shape distortions $γ = γ^I + γ^G$ are assumed to come from coherent intrinsic alignments due to localized physics such as tidal alignments (first term) and the true weak lensing signal caused by structure along the line of sight (second term). These shape distortions are directly and linearly related to the observed galaxy ellipticities, with the exact relation depending on the ellipticity definition (Mandelbaum et al. 2014). Assuming that the tidal fields at the formation of the galaxy set the intrinsic shapes of galaxies through a linear response (Catelan et al. 2001; Hirata & Seljak 2004), the linear alignment model can be expressed as

$$γ^I = (γ^I_+, γ^I_\times) \propto (\partial_x^2 - \partial_y^2, \partial_x \partial_y) \Phi \quad (1)$$

where $Φ$ is the gravitational potential at the time of galaxy formation and $(x, y)$ is an arbitrary coordinate system on the plane of the sky (Kiessling et al. 2015). However, this assumption is expected to break down on small scales (below $2-6 h^{-1} Mpc$) where non-linear gravitational effects take over, along with other small-scale effects such as baryonic physics (for a recent study, see Samuroff et al. 2023) and necessitates specialized models with high capacity and expressivity, such as the halo model and GNNs (Van Alfen et al. 2023; Jagvaral et al. 2022a).

The $γ^G$ term takes a similar form that differs in its geometric configuration and involves the matter overdensity along the line of sight; given the scope of our paper we focus only on the $γ^I$ term. For the interested reader, we refer to the reviews by Troxel & Ishak (2015); Joachimi et al. (2015); Kiessling et al. (2015); Kirk et al. (2015). In the literature, the IA formalism is mostly expressed in 2D, given the observational nature of cosmological measurements. For a 3D simulation box, we make similar assumption relating tidal fields and intrinsic galaxy orientations, and condition the large scale alignments on tidal fields.

## 3 PROBLEM STATEMENT

In general, the problem of emulating galaxies in cosmological simulations can be formulated as a joint conditional probability density estimation problem over correlated properties of different galaxies $x_{\text{galaxy}}$, that may be conditioned on some information about their local environment $Q$:

$$p(x_{\text{galaxy}}|Q) \quad (2)$$

In practice, addressing this joint conditional density estimation problem with analytic methods has been shown to be challenging given the highly complicated physics that drive galaxy evolution (Wechsler & Tinker 2018). Recently, there have been studies focused on tackling this problem with deep learning, particularly for scalar quantities (Stiskalek et al. 2022; Dai et al. 2024). In addition, some non-scalar galaxy properties, such as orientation and position, obey symmetries, adding additional constraints and complexity to the problem. Graph neural networks are specifically designed to capture these joint correlated densities and can be made to respect the symmetries of the problem (Bronstein et al. 2021).

The other aspect of the problem is learning how to model various regimes in cosmological distance scales (linear vs. non-linear
scales). Here we try to capture the complex interconnections and dependencies in the non-linear scales of the universe using graph neural networks (GNNs), namely E(3) equivariant graph neural networks (E3GNN). In order to exploit the global Euclidean symmetry, we chose the analytically imposed E3GNN from Satorras et al. (2021), as opposed to implementing the symmetry solely through data augmentation.

Galaxies are not randomly distributed, but rather are organized into clusters, filaments, and voids, forming a rich network of connections. GNNs can leverage this network structure to learn representations of galaxies that incorporate information from their neighboring environments, though the exact specificity of the connections may be designed to fit the problem and the computational resources. By iteratively aggregating information from neighboring nodes through message passing, GNNs can capture the complex spatial correlations and hierarchical relationships inherent in the cosmic web (Bronstein et al. 2017; Jagvaral et al. 2022a). GNNs provide a powerful framework for representing and analyzing these intricate relationships by treating the cosmic large-scale structure as a set of graphs, where nodes represent galaxies, and edges encode interactions between them.

As for the linear alignments on large scales, we leverage the gravitational tidal forces experienced by galaxies to condition their spatial orientations. By incorporating this information into the input features of a neural network architecture, it becomes possible to predict the statistical properties of intrinsic alignments of galaxies with much less complexity and computational cost, as opposed to building one single graph for the whole galaxy population in the simulation volume (Jagvaral et al. 2022a).

In this paper, we focus on the problem of intrinsic alignments of galaxies, where we are learning and predicting galaxy orientations in 3D jointly with their sizes. To do so, we model DM subhalos and galaxies as triaxial ellipsoids. In subsequent sections we will enumerate the specific quantities $x_{\text{galaxy}}$ that we predict, and the set of Q values on which they are conditioned.

# 4 METHODOLOGY

## 4.1 Score-based Diffusion for SO(3) and Euclidean data

Diffusion generative models have emerged as a powerful class of machine learning techniques for generating complex data distributions. These models operate by simulating a diffusion process that progressively adds noise to the data, transforming it into a simple, often Gaussian distribution. The generative process is then learned by training the model to reverse this diffusion, effectively denoising the data step-by-step to reconstruct the original distribution. Diffusion generative models are particularly notable for their stability and ability to model diverse types of data, including images, audio, and 3D molecular structures (Yang et al. 2023). These models offer a robust framework for generating high-fidelity synthetic data, making them very promising for applications in astrophysics, where they could potentially be used to simulate galaxy distributions, orientations, and other complex galactic physics.

We start with a generic stochastic differential equation (SDE) with a deterministic and a stochastic term:

$$dx = f(x, t) dt + g(t) dw \quad (3)$$

where $x$ is a data vector, $t$ is time and $dw$ is an infinitesimal Brownian motion. As noted in Anderson (1982); De Bortoli et al. (2022), under mild regularity conditions this SDE has a reverse diffusion process,

defined by the reverse-time SDE:

$$dx = [f(x, t) - g(t)^2 \nabla \log p_t(x)] dt + g(t) d\bar{w}, \quad (4)$$

where $\bar{w}$ is a reversed-time Brownian motion and the score function<fnref n="4" /> $\nabla \log p_t(x)$ is the derivative of the logarithmic marginal density of the forward process at time $t$. If we denote the data distribution at $t=0$ as $x(0) \sim p_{\text{data}}$ and $p_t$ the marginal distribution of $x(t)$ at time $t>0$, $p_0 = p_{\text{data}}$ then at final time $T$, $p_T$ will typically tend to a known target distribution.

Following Song & Ermon (2019), we choose $f(x, t) = 0$ and $g(t) = \sqrt{\frac{d\epsilon(t)}{dt}}$ for a given choice of noise schedule<fnref n="5" /> $\epsilon(t)$, which reduces Eq. (3) to:

$$dx = \sqrt{\frac{d\epsilon(t)}{dt}} dw. \quad (5)$$

Next we adopt the variance exploding noise schedule and set the noise schedules for the SDE:

$$\sigma_i^2 = 2\epsilon(t_i), \quad \sigma_{\min} = \sigma_1 < \sigma_i \cdots < \sigma_N = \sigma_{\max} \quad (6)$$

which is usually chosen to be a geometric sequence and $i$ denote different timesteps. Since we are dealing with data of multiple types (scalar vs. SO(3)), here our formalism slightly diverges from Song & Ermon (2019). Given a dataset $x = \{x^S, x^M\}$, where $x^S$ denotes the scalar-valued data and $x^M$ denotes SO(3) manifold-valued data, we can assign the prior distributions for each:

$$\mathcal{U}_{\text{SO}(3)} \text{ Uniform noise for Manifold of SO(3)} \quad (7)$$

$$\mathcal{N}(0, 2\epsilon_{\max} \mathbf{I}) \text{ Gaussian noise for Euclidean scalars} \quad (8)$$

where $2\epsilon_{\max} = 2\epsilon(t_N) = \sigma_{\max}^2$. Next, for scalar values we define the standard noise kernels with a Gaussian distribution, while the $\mathcal{IG}_{\text{SO}(3)}$ from Nikolayev & Savyolov (1970); Matthies et al. (1988) is used for the SO(3) data:

$$p_{\epsilon}(\tilde{x}^M | x^M) = \mathcal{IG}_{\text{SO}(3)}(\tilde{x}^M; x^M, \epsilon) \quad \text{for } x^M, \tilde{x}^M \in \text{SO}(3) \quad (9)$$

$$p_{\epsilon}(\tilde{x}^S | x^S) = \mathcal{N}(\tilde{x}^S; x^S, 2\epsilon\mathbf{I}) \quad \text{for } x^S, \tilde{x}^S \in \mathbb{R}^n \quad (10)$$

where $n$ is the dimension of the scalars and $\tilde{x}(t)$ denote noised data at time $t = t_i^6$. In short, the Isotropic Gaussian Distribution on SO(3): $\mathcal{IG}_{\text{SO}(3)}$ is the heat kernel (solution to the diffusion equation on the SO(3) manifold) that has properties that are well suited for our method, such as closure under convolution (for the noising process) and computational tractability (often heat kernels on non-Euclidean manifolds are infinite series). For the algebraic expressions and quick discussion about $\mathcal{IG}_{\text{SO}(3)}$ distribution please refer to §2.1 of Jagvaral et al. (2024).

These definitions allow us to write down the noised data distribution for each data type:

$$p_{\epsilon}^{M}(\tilde{x}^{M}) = \int_{\text{SO}(3)} p_{\text{data}}^{M}(x^{M}) p_{\epsilon}^{M}(\tilde{x}^{M} | x^{M}) dx^{M}, \quad (11)$$

<footnote n="4">⁴ In the literature there are two different definitions of this score function: when the derivative is taken with respect to the model parameters $\theta$: Fisher score – $\nabla_\theta \log p(x)$; and when the derivative is taken with respect to the data $x$: Stein score – $\nabla_x \log p(x)$. In diffusion models, the Stein score function is almost always used, and we adopt it here as well.</footnote>

<footnote n="5">⁵ The noise schedule is a sequence that controls how much noise is added to data during the diffusion process.</footnote>

<footnote n="6">⁶ For notational brevity, we omit the time dependence on $\tilde{x}(t)$ and just write $\tilde{x}$.</footnote>
$$p_{\epsilon}^{S}(\tilde{\mathbf{x}}^{S}) = \int_{\mathbb{R}^{n}} p_{\text{data}}^{S}(\mathbf{x}^{S}) p_{\epsilon}^{S}(\tilde{\mathbf{x}}^{S} | \mathbf{x}^{S}) d\mathbf{x}^{S}, \qquad (12)$$

and correspond to $p_t$, the marginal distribution of the diffusion process at time $t$: $p_{\epsilon}(t) = p_t$.

This process is entirely defined as soon as the score function of the marginal distribution at any intermediate time $t$, $\nabla \log p_{\epsilon}(t)$, is known.

Samples from $p_0$ can be obtained first by sampling:

$$\mathbf{x} = \begin{cases} \mathbf{x}_T^M, & \text{at } t = T \sim \mathcal{U}_{\text{SO}(3)}. \\ \mathbf{x}_T^S, & \text{at } t = T \sim \mathcal{N}(0, 2\epsilon_{\text{max}}\mathbf{I}). \end{cases} \qquad (13)$$

then evolving them with Eq. (4) down to $t = 0$. However, we still need the score function at intermediate times $t$, $\nabla \log p_{\epsilon}(t)$.

### 4.1.1 The (Stein) score function

We need to approximate the score function of the form:

$$\nabla_x \log p_{\epsilon}(\tilde{\mathbf{x}}|\mathbf{x}) \qquad (14)$$

First, let us consider $\{X_i\}_{i=0}^3$, an orthonormal basis of the tangent space $T_{\epsilon}SO(3)$. The directional derivative of the log density of the noise kernel $p_{\epsilon}(\tilde{\mathbf{x}}^M|\mathbf{x}^M)$ can be computed as:

$$\nabla_{X_i} \log p_{\epsilon}(\tilde{\mathbf{x}}^M | \mathbf{x}^M) = \frac{d}{ds} \log p_{\epsilon}(\tilde{\mathbf{x}}^M \exp(sX_i)|\mathbf{x}^M) \Big|_{s=0}, \qquad (15)$$

which can be computed in practice by automatic differentiation given the explicit approximation formulae for the $\mathcal{I}\mathcal{G}_{\text{SO}(3)}$ distribution from Jagvaral et al. (2024). Thus, the score computations for the SO(3) manifold valued data are tractable, compared to the general Riemannian manifolds.

Now for the scalar valued data, since we chose the noising kernel to be $\mathcal{N}(\tilde{\mathbf{x}}^S; \mathbf{x}^S, 2\epsilon\mathbf{I})$, and noting Eq. (6), we obtain the explicit form

$$\nabla_x \log p_{\epsilon}(\tilde{\mathbf{x}}^S | \mathbf{x}^S) = - \frac{\tilde{\mathbf{x}}^S - \mathbf{x}^S}{\sigma^2} = - \frac{\tilde{\mathbf{x}}^S - \mathbf{x}^S}{2\epsilon}. \qquad (16)$$

Therefore, the score function from Eq. (14) can be decomposed:

$$\nabla_x \log p_{\epsilon}(\tilde{\mathbf{x}}|\mathbf{x}) = \begin{cases} \nabla_{X_i} \log p_{\epsilon}(\tilde{\mathbf{x}}^M|\mathbf{x}) \\ \nabla_x \log p_{\epsilon}(\tilde{\mathbf{x}}^S|\mathbf{x}) \end{cases} \qquad (17)$$

To match this derivative, we introduce a neural score estimator $s_{\theta}$:

$$\nabla_x \log p_{\epsilon}(\tilde{\mathbf{x}}|\mathbf{x}) \approx s_{\theta}(\mathbf{x}, \epsilon) : (\text{SO}(3) \times \mathbb{R}^n) \times \mathbb{R}^{+*} \rightarrow \mathbb{R}^3 \times \mathbb{R}^n \qquad (18)$$

where $\mathbb{R}^{+*}$ denotes strictly positive numbers. This can be trained directly under a conventional denoising score matching loss:

$$\mathcal{L}_{\text{DSM}} = \mathbb{E}_{p_{\text{data}}(\mathbf{x})} \mathbb{E}_{\epsilon \sim \mathcal{N}(0, \sigma_{\epsilon}^2)} \mathbb{E}_{p_{|\epsilon|}(\tilde{\mathbf{x}}|\mathbf{x})} \left[ |\epsilon| \| s_{\theta}(\tilde{\mathbf{x}}, \epsilon) - \nabla_{\mathbf{x}} \log p_{|\epsilon|}(\tilde{\mathbf{x}}|\mathbf{x}) \|_{2}^{2} \right]$$

where we sample at training time random noise scales $\epsilon \sim \mathcal{N}(0, \sigma_{\epsilon}^2/2)$ similarly to Song & Ermon (2020), and we write $|\epsilon|$ to ensure the variance is positive. The minimum of this loss will be achieved for $s_{\theta}(\mathbf{x}, \epsilon) = \nabla \log p_{\epsilon}$.

We then run it backwards starting from:

$$\mathbf{x}_{t=0} = \begin{cases} \mathbf{x}_T^M \sim \mathcal{U}_{\text{SO}(3)}. \\ \mathbf{x}_T^S \sim \mathcal{N}(0, 2\epsilon_{\text{max}}\mathbf{I}). \end{cases} \qquad (19)$$

then using the learned score and the corresponding differential equation

$$d\mathbf{x}_t = -\frac{1}{2} \frac{d\epsilon_t}{dt} s_{\theta}(\mathbf{x}_t, \epsilon(t)) dt \qquad (20)$$

or in the decomposed form:

$$d\mathbf{x}_t = -\frac{1}{2} \frac{d\epsilon_t}{dt} \left\{ s_{\theta}^{M}(\mathbf{x}_t, \epsilon_t) \begin{array}{l} s_{\theta}^{S}(\mathbf{x}_t, \epsilon_t) \end{array} \right\} dt \qquad (21)$$

The network specifications for $s_{\theta}(\mathbf{x}_t, \epsilon(t))$ can be found in Appendix A.

We note that integrating these equations require special care for quantities on the SO(3) manifold. Indeed, in addition to minimizing the integration error, the iterative solver must also ensure that the variable remains on the manifold at each integration step. We use in practice Heun's method to integrate this system, taking advantage of a special geometric flavor of the method for the SO(3) part of the signal. We direct the interested reader to Appendix B for the algorithm and to (Iserles et al. 2000) for more information on geometric ODE solvers.

## 4.2 E(3) GNN: 3D Euclidean group equivariant graph neural networks

Graphs $\mathcal{G} = (\mathcal{V}, \mathcal{E})$ are a mathematical structure composed of set of objects: nodes $v_i \in \mathcal{V}$, and their pair-wise connections – edges $e_{ij} \in \mathcal{E}$. In graph-based data, each node is associated with features that can be broadly categorized into scalar node features $\mathbf{h}_i \in \mathbb{R}^n$, which represent various properties such as mass or color; and geometric node features $v_i \in \mathbb{R}^3$, which may represent spatial information like position, orientation, velocity. We will closely follow the formalism and notation of Satorras et al. (2021) and Duval et al. (2023) – thus $\mathbf{x}^S = \mathbf{h}$ and $\mathbf{x}^M = \mathbf{v}$ to connect with notation from the previous subsection.

Graph Neural Networks (GNNs) are a class of neural networks specifically designed to operate on graph structures, making them particularly adept at capturing the relationships and interactions between nodes. A GNN can be made geometric, dubbed Geometric Graph Neural Networks, where each node $i$ possesses scalar and geometric features/representations $\mathbf{h}_i^l$ and $\mathbf{v}_i^l$, respectively, at layer $l$. Nodes exchange information through message passing, where messages are computed using a learnable function Msg. These messages are then aggregated from the neighbors $N_i$ of each node $i$ using a permutation-invariant operator $\oplus$, ensuring that the order of nodes does not affect the result. Finally, the features of each node are updated using its current representation and the aggregated messages via a learnable function Upd. The general form of the message passing between nodes $i$ and $j$ at layer $l$: $\mathbf{m}_{ij}^l$, and the subsequent feature update mechanism on node $i$, are mathematically described by the following equations:

$$\mathbf{m}_{ij}^l = \text{Msg}(\mathbf{h}_i^l, \mathbf{h}_j^l, \mathbf{v}_i^l, \mathbf{v}_j^l, d_{ij}) \qquad (22)$$

$$\mathbf{h}_i^{l+1}, \mathbf{v}_i^{l+1} = \text{Upd}(\mathbf{h}_i^l, \mathbf{v}_i^l, \bigoplus_{j \in N_i} (\mathbf{m}_{ij}^l)) \qquad (23)$$

with $d_{ij} = |\mathbf{v}_i^l - \mathbf{v}_j^l|_2$ denoting the Euclidean distance between nodes $i$ and $j$. Further, in order to enforce E(3) equivariance which include symmetries of rotations, translations and reflections, we follow Satorras et al. (2021); Hoogeboom et al. (2022) and chose the Msg and Upd function as follows:

$$\mathbf{m}_{ij}^l = \phi_e(\mathbf{h}_i^l, \mathbf{h}_j^l, d_{ij}^2) \qquad (24)$$
$$ \mathbf{h}_i^{l+1} = \phi_h(\mathbf{h}_i^l; \sum_{j \neq i} \tilde{e}_{ij} \mathbf{m}_{ij}^l) \qquad (25) $$

$$ \mathbf{v}_i^{l+1} = \mathbf{v}_i^l + \sum_{j \neq i} \frac{\mathbf{v}_i^l - \mathbf{v}_j^l}{d_{ij} + 1} \phi_v(\mathbf{h}_i^l, \mathbf{h}_j^l, d_{ij}^2) \qquad (26) $$

with $\tilde{e}_{ij} = \phi_{\text{inf}}(\mathbf{m}_{ij}^l) \in (0, 1)$. The various $\phi$ functions are implementations of the Msg and Upd mentioned above, which are small Multilayer Perceptrons whose specifications can be found in the Appendix B of Hoogeboom et al. (2022).

Eqs. (24)-(26) need to be equivariant under the group transformations. Trivially, $\mathbf{m}_{ij}$ and $\mathbf{h}_i^{l+1}$ are invariant, since their inputs and outputs are scalars. For Eq. (26), the second term has $(\mathbf{v}_i^l - \mathbf{v}_j^l)$ in which the two translation transformation will cancel each other. Next, since (roto-)reflections are linear operators, we can factor out the (roto-)reflections outside the parenthesis of the second term. Finally, when we add the first term $\mathbf{v}_i^l$, we get a translation term and we can again factor out the (roto-)reflections. Thus, the whole equation is equivariant under E(3) transformations; for details see Appendix A of Hoogeboom et al. (2022).

We can stack $L$ number of GNN layers to form the E3GNN architecture block, which essentially transforms (non-linearly) the initial features, $\mathbf{h}^l=0$ and $\mathbf{v}^l=0$, to obtain abstracted representations $\mathbf{h}^l=1..L$ and $\mathbf{v}^l=1..L$. Typically, for molecular modeling, $L$ is on the order of 2-8; we chose $L=5$ for our study. Empirically, we found that $L < 4$ were not expressive enough and $L > 6$ were too slow to use in practice.

## 4.3 Graph construction

To construct the graph for the cosmic web, we follow Jagvaral et al. (2022a). First, we grouped all of the subhalos and galaxies based on their parent halo using their group membership ID from the halo finder. There exist a few choices for modeling proximity graph relations between the members in a group, such as the Gabriel graph (which is a subset of the Delaunay triangulation) and the radius nearest neighbor graph (r-NNG). These different graph types differ by their connectivity, i.e., they have different adjacency matrices (Mathieson & Moscato 2019). In this study we employ the r-NNG to model the connectivity of our graphs with a radius of $1 h^{-1}$Mpc.

Given a galaxy catalog, an undirected graph based on the 3D positions is built by placing each galaxy on a graph node. Then, each node will have a list of features such as halo mass, subhalo mass, central vs. satellite identification (binary column) and tidal fields smoothed on several scales. Then for a given group (i.e., within a halo) the graphs are connected using r-NNG with a radius of 1 Mpc/h. To build the graph connection, the nearest neighbors within a specified radius for a given node are connected via the undirected edges with signals on the graphs representing the alignments and the scalar features.

# 5 ASTROPHYSICAL DATA PRODUCTS AND OBSERVABLES

## 5.1 The Cosmological Simulation

The IllustrisTNG suite of cosmological hydrodynamical simulations was run using the moving mesh code AREPO (Springel 2010) in

periodic boxes of ~ 50, 100, 300 Mpc, each with three different varying resolutions (Nelson et al. 2018; Pillepich et al. 2018b; Springel et al. 2018; Naiman et al. 2018; Marinacci et al. 2018; Nelson et al. 2019). The simulations were initialized using cosmological parameters from Planck CMB measurements, assuming a flat $\Lambda$CDM model (Planck Collaboration et al. 2016). Here, we chose the IllustrisTNG100 dataset due to its balance between large volume and high resolution. Alongside magneto-hydrodynamic prescriptions of galaxy evolution, the simulations also include radiative gas dynamics; star formation and evolution; various forms of feedback from supernova and AGN (for more details, see Pillepich et al. 2018a; Weinberger et al. 2017). The friends-of-friends (FoF; Davis et al. 1985) and SUBFIND (Springel et al. 2001) algorithms were employed to identify halos and subhalos in the simulations, respectively.

We used the latest snapshot at $z=0$ for our analysis. A mass cut of $\log_{10}(M_*/M_{\odot}) \ge 9$ was enforced in order to obtain reliable galaxy shape measurements and to avoid resolution limitations (Jagvaral et al. 2021). The galaxy morphologies were determined based on a dynamical decomposition of the star particles into angular momentum-dominated and dispersion-dominated components; for more details, see Jagvaral et al. (2022b).

## 5.2 Tidal field

In order to predict galaxy and DM halo shapes and orientations, we use the tidal field as a conditioning feature. The tidal field is defined as the Hessian of the gravitational potential $\Phi$:

$$ T_{ij}(\mathbf{r}) = \frac{\partial^2 \Phi(\mathbf{r})}{\partial r_i \partial r_j}, \qquad (27) $$

Using Poisson's equation, the Fourier-space smoothed tidal field, expressed in terms of the overdensity field $\delta$, with a smoothing scale $\eta$ is:

$$ \hat{T}_{ij}(\mathbf{k}) = 4\pi G \bar{\rho} \frac{k_i k_j}{k^2} \delta(\mathbf{k}) e^{-k^2 \eta^2 / 2}. \qquad (28) $$

The overdensity field was obtained by downsampling 1 million DM particles to reduce computational cost. The tidal field was evaluated at the position of each galaxy, using a cloud-in-cell window kernel to interpolate between the centers of the grid points, with $\eta:= 1 h^{-1}$Mpc on a mesh of size $1024^3$, with cell sizes given as $L_{\text{box}}/1024 = 0.073 h^{-1}$Mpc.

## 5.3 Shapes and Orientation of Halos and Galaxies

We use the simple (as opposed to the reduced or the iterative) mass quadrupole moments $I_{ij}$ to measure the shapes of galaxies (DM (sub)halos) using star particles (DM particles):

$$ I_{ij} = \frac{\sum_n m_n r_{ni} r_{nj}}{\sum_n m_n}. \qquad (29) $$

the summation index $n$ runs over all particles of a given type in a given galaxy, where $m_n$ is the mass of the $n^{th}$ particle and $r$ is the distance between the galaxy/subhalo centre of mass and the $n^{th}$ particle, with $i$ and $j$ indexing the three spatial directions, the $x, y, z$ Cartesian coordinates of the simulation box.

The three eigenvectors of $I_{ij}$, defined as $s_\mu = \{s_{x,\mu}, s_{y,\mu}, s_{z,\mu}\}^\tau$ and $\mu \in \{a, b, c\}$ with eigenvalues $\omega_a, \omega_b, \omega_c$, are related to the half-lengths of the principal axes of the ellipsoid by: $a \propto \sqrt{\omega_a}$, $b \propto \sqrt{\omega_b}$, and $c \propto \sqrt{\omega_c}$, such that $a \ge b \ge c$.

To compute the projected alignment signals, we need to use the 3D mass quadrupole moments to define 2D projected shapes. Following
Joachimi et al. (2013); Jagvaral et al. (2022a), we can obtain the projected 2D ellipse as such:

$$ (e_1, e_2) = \frac{(B_{xx} - B_{yy}, 2B_{xy})}{B_{xx} + B_{yy} + 2\sqrt{\det\mathbf{B}}}. \qquad (30) $$

where

$$ \mathbf{B}^{-1} = \sum_{\mu=1}^{3} \frac{s_{\perp, \mu} \mathbf{b} s_{\perp, \mu}^{\tau}}{\omega_{\mu}^{2}} - \frac{\mathbf{k} \mathbf{k}^{\tau}}{\alpha^{2}}, \qquad (31) $$

and

$$ \mathbf{k} = \sum_{\mu=1}^{3} \frac{s_{\parallel, \mu} s_{\perp, \mu}}{\omega_{\mu}^{2}} \quad \text{and} \quad \alpha^{2} = \sum_{\mu=1}^{3} \left( \frac{s_{\parallel, \mu}}{\omega_{\mu}} \right)^{2}. \qquad (32) $$

Here, $s_{\perp, \mu} = \{s_{x, \mu}, s_{y, \mu}\}^{\tau}$ are the eigenvectors projected along the projection axis (for which we arbitrarily choose the z-axis of the 3D simulation box). In terms of the projected simulation box, the x, y directions correspond to the positive and negative direction of $e_1$.

## 5.4 Two-point estimators

Here we describe the two-point correlation functions that are used to quantify IA. The ellipticity-direction (ED) correlation captures the position and the orientation correlation angles in 3D, whereas the projected density-shape correlation function ($w_{g+}$) captures the correlation between overdensity and projected intrinsic ellipticity (which includes information on both shape and orientation in 2D). All of the two-point statistic were measured using the HALOTOOLS package v0.7 (Hearin et al. 2017) and the supporting halotools_ia package.

### 5.4.1 Density-Orientation Correlation Functions in 3D

The ellipticity-direction (ED) correlation function is defined for each orientation axis $\mu \in \{a, b, c\}$ as follows (Lee et al. 2008):

$$ \omega_{\mu}(r) = \langle |\mathbf{s}_{\mu}(\mathbf{x}_{\text{pos}}) \cdot \mathbf{\vec{r}}(\mathbf{x}_{\text{pos}})|^2 \rangle - \frac{1}{3} = \omega_{\mu}^{1h}(r) + \omega_{\mu}^{2h}(r) \qquad (33) $$

for a subhalo/galaxy at position $\mathbf{x}_{\text{pos}}$ with axes direction $s_{\mu}$ and the unit vector $\mathbf{\vec{r}}$ denoting the direction of a density tracer at a distance $r$. Additionally, this correlation function may be decomposed into the sum of 1-halo and 2-halo terms, where the 1-halo term captures the correlation among subhalos within the same halo, and the 2-halo term captures contributions from pairs of subhalos belonging to different halos. The estimator we use to compute it in the simulations is as follows:

$$ \omega_{\mu}(r) = \sum_{\alpha \neq \beta} |s_{\mu}^{\alpha} \cdot \mathbf{\vec{r}}^{\alpha \to \beta}|^2 - \frac{1}{3} \qquad (34) $$

where the orientation vector $s_{\mu}^{\alpha}$ of galaxy $\alpha$ is dotted with the unit vector pointing from the position of galaxy $\alpha$ to the position of galaxy $\beta$, $\mathbf{\vec{r}}^{\alpha \to \beta}$.

### 5.4.2 Density-Shape Correlation Functions in 2D

The cross correlation function of galaxy positions (or density tracers) and intrinsic ellipticities is defined as:

$$ \xi_{g+}(\mathbf{r}) = \langle \delta_g(\mathbf{x}) \delta_+(\mathbf{x} + \mathbf{r}) \rangle \qquad (35) $$

where $\delta_g(\mathbf{r})$ and $\delta_+(\mathbf{r})$ represent the galaxy overdensity field and the intrinsic shape field, respectively. It can be estimated using the

method described in Mandelbaum et al. (2011) as a function of $r_p$ (plane-of-sky distance) and $\Pi$ (line-of-sight distance):

$$ \xi_{g+}(r_p, \Pi) = \frac{S_+ D - S_+ R}{RR}. \qquad (36) $$

Here, $RR$ are counts of random-random pairs binned based on their perpendicular and line-of-sight separation;

$$ S_+ D \equiv \frac{1}{2} \sum_{\alpha \neq \beta} e_+(\beta | \alpha), \qquad (37) $$

represent the shape correlations, where $e_+(\beta | \alpha)$ is the + component of the ellipticity of galaxy $\beta$ (from the shape sample) measured relative to the direction of galaxy $\alpha$ (from the density tracer sample). $S_+R$ is defined in an equivalent way, but instead of using galaxy positions we use randoms.

The two-point correlation functions can be projected onto 2D by integrating over the third dimension, with the integral approximated as sums over the line-of-sight separation ($\Pi$) bins:

$$ w_{g+}(r_p) = \sum_{\Pi_{\text{max}}} \Delta\Pi \xi_{g+}(r_p, \Pi), \qquad (38) $$

where we chose a $\Pi_{\text{max}}$ value of $20 h^{-1} \text{Mpc}$, following Jagvaral et al. (2022b).

## 6 RESULTS

In this section, we present and investigate the outputs from the neural network. The results presented here are from a single random seed, since stochastic variation due to different random seeds has been shown to be subdominant to other sources of uncertainty (Jagvaral et al. 2022a; Jordan 2023). The train-test splits were roughly 50/50, whilst preserving the halo membership of galaxies. For orientation/alignment statistics and subsamples, we present the results on the whole dataset in order to increase signal-to-noise.

First, for the scalar outputs, we will present the corner plot to illustrate the pairwise correlations. Second, in order to quantify the correlations of geometric quantities we measure and present the 3D Ellipticity-Direction (ED) correlation functions. Third, we present 2D projected quantities that are of great interest for weak gravitational lensing. We remind the reader of Eq. (2), and we list the inputs and outputs of the model in Table 1.

The scalar features of interest for weak lensing science are the $a, b, c$ values that characterize galaxy/subhalo size, defined in §5.3; or equivalently defined as $q = b/a, s = c/a$ on the interval (0, 1]. In addition, we predicted relevant scalars such as galaxy mass, DM subhalo mass and galaxy color. In addition, the output geometric features are galaxy orientations and DM-subhalo orientations.

In principle, one can add as many scalar quantities as one wishes, however due to the neural scaling laws the network capacity needs to be increased, in turn leading to a need for increased computer resources (Bahri et al. 2021).

### 6.1 Investigations into generated galaxy features

In Figure 1, we visualize and compare the joint distribution of scalar quantities using a corner plot of the generated and the TNG samples. Below the diagonal cells of the plot grid, the scatter plot along with contour lines helps us visualize the correlations across scalar quantities. On the primary diagonal, the marginal distributions of
<figure><img src="image_2.png" /><figcaption>**Figure 1.** Corner plot of correlations and 1D histograms between scalar quantities in the true TNG training sample versus the generated testing sample. The two joint distributions visually match well, and quantitative metrics are provided in Table 2.</figcaption></figure>

each scalar feature are shown. As shown, both the marginal and scatter/contours show good agreement between the generated and the TNG sample. Interestingly, the model has captured the bimodal distribution of galaxy colors as well as the unimodal distributions, such as the galaxy sizes/shapes. Further, when broken down by mass, morphology and central/satellite type, the model has learned the trends with these quantities to a good degree, showing similar agreement between the two distributions (not shown here for brevity, but Figure 4 implicitly shows this, since it uses three scalars to compute the projected 2D galaxy shapes).

In order to quantify the agreement, we calculate Kolmogorov-Smirnov test statistics (KS-test); the ratio of the means; the ratio of

the variances; and Wasserstein-1 distances on these quantities and present them in Table 2. For all quantities, the KS-test results are all below $10^{-1}$; for $\{a^G, s^G, a^{DM}, M^{DM}/M^G, g-i\}$, the p-values support the null hypothesis that the two samples came from the same distribution. However, the corresponding p-values for the constrained quantities such as q, s and mass are below the standard threshold of 0.05. This may be due to the “edge” of the distribution shifting due to the cutoffs imposed on these parameters, which are challenging for the neural network to pin down. However, in our intended astrophysical application, these low p-values may be negligible given that the shape of the distributions match very well, as seen by the ratios ~ 1 of the mean and variance in Table 2. Lastly, we have computed the
<table><caption>**Table 1.** Summary of input/output features with their names, data-type, constraints and description.</caption><thead><tr><th>Q (Conditional Input)</th><th>Data-Type</th><th>Constraints</th><th>Description</th></tr></thead><tbody><tr><td>λ&#773;<sub>1</sub>, λ&#773;<sub>2</sub>, λ&#773;<sub>3</sub></td><td>R<sup>3</sup></td><td>||λ&#773;<sub>i</sub>||=1</td><td>eigenvectors of tidal field, defined in §5.2 (at subhalo positions)</td></tr><tr><td>λ<sub>1</sub>, λ<sub>2</sub>, λ<sub>3</sub></td><td>Scalar</td><td>–</td><td>eigenvalues of tidal field, defined in §5.2 (at subhalo positions)</td></tr><tr><td>M<sub>total</sub></td><td>Scalar</td><td>–</td><td>subhalo total mass: stellar+DM+gas</td></tr><tr><td>s<sup>c</sup></td><td>Binary</td><td>{0,1}</td><td>binary column indicating centrals vs. satellites</td></tr><tr><td>x (Modeled Quantities)</td><td></td><td></td><td></td></tr><tr><td>log[a<sup>G</sup>]</td><td>Scalar</td><td>(0, ∞)</td><td>galaxy size: semi-major axis length, defined in §5.3</td></tr><tr><td>q<sup>G</sup></td><td>Scalar</td><td>(0, 1]</td><td>galaxy shape: intermediate-to-major axis ratio</td></tr><tr><td>s<sup>G</sup></td><td>Scalar</td><td>(0, 1]</td><td>galaxy shape: minor-to-major axis ratio</td></tr><tr><td>log[a<sup>DM</sup>]</td><td>Scalar</td><td>(0, ∞)</td><td>DM-subhalo size: semi-major axis length, defined in §5.3</td></tr><tr><td>q<sup>DM</sup></td><td>Scalar</td><td>(0, 1]</td><td>DM-subhalo shape: intermediate-to-major axis ratio</td></tr><tr><td>s<sup>DM</sup></td><td>Scalar</td><td>(0, 1]</td><td>DM-subhalo shape: minor-to-major axis ratio</td></tr><tr><td>log[M<sup>G</sup>/M<sub>☉</sub>]</td><td>Scalar</td><td>(9, ∞)</td><td>log(galaxy stellar mass)</td></tr><tr><td>log[M<sup>DM</sup>/M<sup>G</sup>]</td><td>Scalar</td><td>–</td><td>log of ratio of (DM-subhalo mass) to the (galaxy stellar mass)</td></tr><tr><td>g − i</td><td>Scalar</td><td>–</td><td>galaxy color</td></tr><tr><td>x<sub>t=0</sub><sup>M,G</sup></td><td>SO(3)</td><td>S<sup>3</sup>: unit 3-sphere</td><td>3D orientation of galaxy</td></tr><tr><td>x<sub>t=0</sub><sup>M,DM</sup></td><td>SO(3)</td><td>S<sup>3</sup>: unit 3-sphere</td><td>3D orientation of DM-subhalo</td></tr></tbody></table>

<table><caption>**Table 2.** Kolmogorov-Smirnov test statistics, the corresponding *p*-values; ratio of the means; ratio of the variances; and the Wasserstein-1 distance between the true TNG and generated scalar galaxy properties included within *x*<sub>galaxy</sub>, as defined in Eq. (2). The KS test *p*-values are consistent with the null hypothesis that the generated samples and the TNG samples are drawn from the same underlying distribution for unconstrained quantities, while constrained quantities such as *q*, *s* and galaxy mass have *p*-values ~ 0. The ratios of the first 2 moments are all ~ 1. The Wasserstein-1 distances of ≤ 10<sup>-2</sup> also indicate the similarity of the distributions.</caption><thead><tr><th>x<sup>S</sup> (Scalars)</th><th>KS</th><th>p-value</th><th>Mean(Gen) Mean(TNG)</th><th>Var(Gen) Var(TNG)</th><th>W<sub>1</sub></th></tr></thead><tbody><tr><td>log[a<sup>G</sup>]</td><td>0.0077</td><td>0.68</td><td>1.00</td><td>1.01</td><td>0.0028</td></tr><tr><td>q<sup>G</sup></td><td>0.014</td><td>0.07</td><td>0.99</td><td>1.04</td><td>0.0032</td></tr><tr><td>s<sup>G</sup></td><td>0.0056</td><td>0.94</td><td>1.00</td><td>1.01</td><td>0.0014</td></tr><tr><td>log[a<sup>DM</sup>]</td><td>0.0062</td><td>0.88</td><td>1.00</td><td>1.00</td><td>0.0034</td></tr><tr><td>q<sup>DM</sup></td><td>0.028</td><td>0</td><td>1.00</td><td>1.04</td><td>0.0035</td></tr><tr><td>s<sup>DM</sup></td><td>0.018</td><td>0.05</td><td>0.99</td><td>1.02</td><td>0.0033</td></tr><tr><td>log[M<sup>G</sup>/M<sub>☉</sub>]</td><td>0.052</td><td>0</td><td>0.99</td><td>0.99</td><td>0.0049</td></tr><tr><td>log[M<sup>DM</sup>/M<sup>G</sup>]</td><td>0.0045</td><td>0.99</td><td>1.00</td><td>1.00</td><td>0.0022</td></tr><tr><td>g − i</td><td>0.0075</td><td>0.71</td><td>1.00</td><td>1.00</td><td>0.0018</td></tr></tbody></table>

Wasserstein-1 ($W_1$) metric for these scalars and it shows distances below 10<sup>-2</sup>. The $W_1$ distance measures the minimum cost required to move one probability distribution to another, where the cost is the distance moved times the amount of “probability distribution”. Compared to the KS-test, the $W_1$ metric is more sensitive to the overall shape of the distributions, whereas for large samples, even small differences can lead to a rejection of the null hypothesis for a KS-test (Corder & Foreman 2009). Based on these findings, we conclude that

the generated scalars are robust. Additionally, when broken down by mass, morphology and central/satellite type (not shown on Fig. 1), the scalar quantities exhibited similar agreement between TNG-100 and the generated samples, both visually and quantitatively.

Moving on to the geometric quantities, in Figure 2 we present the ED correlations as a function of scale, as well as the 1- and 2-halo decomposed curves. The errorbars were obtained using jackknife estimates. The first (second) row contains the ED correlation function calculated using each of the 3 axes orientation of galaxies (DM subhalos) correlated with the density tracers. The generated samples reproduce the same correlated statistics as the TNG simulation across all scales within the given statistical precision. We remind the reader that the 1-halo term is dominated by highly complicated physical interactions that are hard to analytically model and the GNN has successfully handled this regime. While the 2-halo term is often assumed to be dominated by (linear) gravity, it still encompasses highly nonlinear scales, and was conditioned on the smoothed tidal field interpolated to the galaxy positions. This was done to reduce the computational cost of the method. By not connecting the halos, we save in the number of graph edges – which dominates the computational expense.

When the shape sample is broken down by mass, morphology and the central-satellite classification, the results for the generated samples and TNG still agree well, as shown in Figure 3 for representative subsamples. Here we show the ratio of the ED curves of generated galaxy orientations against the true TNG galaxy orientations for the whole sample (black solid line) and for various subsamples (colored dashed/dotted lines). Given the very low signal of the intermediate axis and its non-prevalent use in literature, we omitted it for brevity, focusing only on the longest and shortest axis. The ratio curves show trends consistent with 1 given the errorbars. Of these curves, the centrals show trends furthest away from 1, probably due to our choice to build the graphs separately for each halo (thus centrals do not pass
<figure><img src="image_3.png" /><figcaption>**Figure 2.** Ellipticity-Direction (ED) correlation function of galaxies (top row) and DM subhalos (bottom) for the whole sample. The ratio of the Generated ED function to the TNG ED function is shown in Figure 3 along with the same ratio for subsamples. The errorbars are shown for only 1 curve for visual clarity. We see an agreement between the true and generated values across all panels and on all scales.</figcaption></figure>

messages to each other during optimization via graphs). The alignment of centrals is therefore mostly inferred from the tidal fields.

Therefore, the learned model reproduces the target joint distributions to a good statistical agreement within our error bars. These results show us that geometric deep learning methods are a powerful way to model the highly complex joint probability density of the galaxy-halo connection. Compared with the GAN result from Jagvaral et al. (2022a), this model robustly and jointly captures all the features concurrently and in full 3D.

## 6.2 Projected data products for weak gravitational lensing

We project the 3D triaxial ellipsoidal model for the galaxy onto 2D in order to produce and compare the quantity that is typically measured to characterize the IA contamination to the weak gravitational lensing observables, namely the $w_{g+}$ correlation function. We follow the procedure described in § 5.3 to project the 3D inertia tensor onto 2D to obtain the complex ellipticities $e_1$ and $e_2$. The measured $w_{g+}$ functions are presented in Figure 4; in all of these curves the density tracers are the same – all galaxies above our adopted mass threshold. The curves from generated samples follow the TNG curves within the statistical errors. The sub-populations also exhibit good agreement with the truth. Notably, morphological information was not given to the model, yet the model learned the trends with morphology. Hence, the method we propose in this paper can potentially be used in current and future weak lensing surveys to generate realistically complex mock galaxy catalogs that can be used to test the analysis pipelines.

In the application of machine learning to cosmology, one critical step involves avoiding overfitting, which is commonly done by dividing the data into training and testing sets. However, simulated data are limited by computational cost. In order to decrease the shot noise, we used the full TNG sample as both the training set and testing set in Figures 2, 3, and 4. To test the impact of this choice and ascertain whether our results may be subject to overfitting, in Figure 5 we show the train and test split samples for the $w_{g+}$ correlation function. The train-test splits were done roughly 50/50, preserving the group membership of galaxies by sampling the halos. The two TNG curves presented in Figure 5 are more noisy than the curve presented in Figure 4 due to the decreased sample size, but they do agree within the errors. As for the $w_{g+}$ measurements from the generated samples, they follow very closely the true curve and are statistically consistent, if we assume the same covariance estimates. Therefore, we conclude that our model has not been subject to significant overfitting.

To summarize, our results demonstrate that the integration of machine learning techniques with geometric deep learning methodologies provides a robust framework for modeling complex joint probability distributions of galaxy properties that are governed by highly non-linear galactic physics. By effectively capturing both scalars and orientations, the model shows strong agreement with the true TNG samples across various scales, sub-populations and parameters. This advance underscores the potential of geometric deep learning to tackle complex problems in astrophysics, paving the way for more precise and tractable mock galaxy catalogs for testing and improving future cosmological analyses pipelines. One critical step that was not demonstrated in this paper is the performance of the model when var-
<figure><img src="image_4.png" /><figcaption>**Figure 3.** Ratio of the Ellipticity-Direction (ED) correlation function between the entire generated and TNG samples (black) and representative subsamples (colors indicated in legend). The top (bottom) panel shows results for the longest (shortest) axis; we do not show results for the intermediate axis due to its very low signal. The errorbars on the subsample curves were horizontally shifted by 3% for visual clarity. These ratio curves exhibit consistency with 1 given the statistical error of the measurements.</figcaption></figure>

ious simulation suites are considered; we will explore this in future work.

## 7 CONCLUSIONS

We have presented a machine learning based method for modeling the galaxy-halo connection, specifically focusing on the problem of intrinsic alignments of galaxies and their dependence on properties such as mass. Leveraging recent advances in geometric deep learning methods, we have modeled the galaxy distributions within dark matter halos as a set of graphs, where the nodes represent galaxies and edges representing their interaction, whilst incorporating the global E(3) symmetry. Additionally, the 3D intrinsic orientations of galaxies were incorporated into a novel diffusion-based model for SO(3). We have demonstrated the efficacy of our method on the TNG100-1 hydrodynamical simulation from the IllustrisTNG suite with ~18k galaxies. The generated samples from the trained model agree well with the true data, showing consistent correlations for both geometric and scalar quantities.

This successful demonstration of our methodology motivates its extensions towards practical applications in future work. One issue to be addressed in future is that this model has custom-tuned, possibly overtuned, hyperparameters that may be specific to TNG. We do not

know if these hyperparameters will give us similar quality results when training or testing on other simulations. Extending this work to other simulations is a crucial step for future work. Also, in future work, the model is planned to be deployed on a much lower resolution high-volume simulation than the one considered here. Additional tests must be done on lower resolution simulations within the TNG suite to understand resolution effects on these models.

Moving forward, it is essential to continue refining and expanding upon these machine learning methodologies to enhance the efficiency and robustness of mock galaxy catalog production, as well as to deploy these methods to test the data analysis pipelines for future cosmological surveys, such as Rubin LSST. Future studies could explore additional ML architectures and datasets, whilst incorporating more sophisticated machine learning techniques such as transfer learning (Iman et al. 2022; Vilalta 2018) on various cosmological simulations. In an upcoming paper, we will investigate how to fully generate galaxy samples given a low resolution halo catalog, where we anticipate the outcome to be of comparable quality and realism as hydrodynamical simulations at a far lower computational cost. This in turn will allow an alternative and economical way of producing mock galaxy catalogs.

## ACKNOWLEDGEMENTS

We thank Chad Schafer and Rupert Croft for useful feedback. This work was supported in part by a grant from the Simons Foundation (Simons Investigator in Astrophysics, Award ID 620789). This work is supported by the NSF AI Institute: Physics of the Future, NSF PHY-2020295. The computations in this work were, in part, run at facilities supported by the Scientific Computing Core at the Flatiron Institute, a division of the Simons Foundation. YJ is thankful to Francisco Villaescusa-Navarro for hosting YJ at the Simons Foundation.

## DATA AVAILABILITY

The IllustrisTNG data can be obtained through the website at https://www.tng-project.org/data/. The catalog data with morphological decompositions of galaxies is available at https://github.com/McWilliamsCenter/gal_decomp_paper. The software developed as part of this work will be made available after peer-review.

## REFERENCES

* Anderson B. D., 1982, Stochastic Processes and their Applications, 12, 313

* Bahri Y., Dyer E., Kaplan J., Lee J., Sharma U., 2021, arXiv e-prints, p. arXiv:2102.06701

* Benson A. J., 2012, New Astron., 17, 175

* Bond-Taylor S., Leach A., Long Y., Willcocks C. G., 2021, arXiv e-prints, p. arXiv:2103.04922

* Bronstein M. M., Bruna J., LeCun Y., Szlam A., Vandergheynst P., 2017, IEEE Signal Processing Magazine, 34, 18

* Bronstein M. M., Bruna J., Cohen T., Veličković P., 2021, arXiv e-prints, p. arXiv:2104.13478

* Catelan P., Kamionkowski M., Blandford R. D., 2001, MNRAS, 320, L7

* Corder G. W., Foreman D. I., 2009, Nonparametric Statistics for Non-Statisticians: A Step-by-Step Approach, 1 edn. John Wiley & Sons

* Dai Z., Moews B., Vilalta R., Davé R., 2024, MNRAS, 527, 3381

* Davis M., Efstathiou G., Frenk C. S., White S. D. M., 1985, ApJ, 292, 371
<figure><img src="image_5.png" /><figcaption>**Figure 4.** Projected two-point correlation functions $w_{g+}$ of all galaxy positions and projected 2D ellipticities of (sub)samples. For visual clarity we only show the errorbars on a single measurement. The left panel contains measurements for all galaxies and for central and satellite subsamples. The center (right) panel contains measurements for samples split by mass (morphology). The measurements with generated quantities show consistency across all scales with the true TNG measurements.</figcaption></figure>

<figure><img src="image_6.png" /><figcaption>**Figure 5.** Projected two-point correlation functions $w_{g+}$ of galaxy positions and the projected 2D ellipticities of train-test split samples. The generated curves follow the TNG measurements closely, showing no signs of overfitting.</figcaption></figure>

De Bortoli V., Mathieu E., Hutchinson M., Thornton J., Teh Y. W., Doucet A., 2022, Riemannian Score-Based Generative Modeling, doi:10.48550/arXiv.2202.02763, http://arxiv.org/abs/2202.02763

De Lucia G., 2019, Galaxies, 7

DeRose et al J., 2019, arXiv e-prints, p. arXiv:1901.02401

DeRose et al J., 2021, arXiv e-prints, p. arXiv:2105.13547

Del Popolo A., Le Delliou M., 2017, Galaxies, 5, 17

Duval A., et al., 2023, arXiv e-prints, p. arXiv:2312.07511

Euclid Collaboration et al., 2024, arXiv e-prints, p. arXiv:2405.13495

Guo H., et al., 2016, MNRAS, 459, 3040

Hearin A. P., et al., 2017, AJ, 154, 190

Hirata C. M., Seljak U., 2004, Phys. Rev. D, 70, 063526

Hoogeboom E., Satorras V. G., Vignac C., Welling M., 2022, in Chaudhuri K., Jegelka S., Song L., Szepesvari C., Niu G., Sabato S., eds, Proceedings of Machine Learning Research Vol. 162, Proceedings of the 39th International Conference on Machine Learning. PMLR, pp 8867–8887, https://proceedings.mlr.press/v162/hoogeboom22a.html

Huertas-Company M., Lanusse F., 2023, Publ. Astron. Soc. Australia, 40, e001

Iman M., Rasheed K., Arabnia H. R., 2022, arXiv e-prints, p. arXiv:2201.09679

Iserles A., Munthe-Kaas H. Z., Nørsett S. P., Zanna A., 2000, Acta Numerica, 9, 215–365

Jagvaral Y., Campbell D., Mandelbaum R., Rau M. M., 2021, arXiv e-prints, p. arXiv:2105.02237

Jagvaral Y., Lanusse F., Singh S., Mandelbaum R., Ravanbakhsh S., Campbell D., 2022a, Monthly Notices of the Royal Astronomical Society

Jagvaral Y., Singh S., Mandelbaum R., 2022b, arXiv e-prints, p. arXiv:2202.08849

Jagvaral Y., Lanusse F., Mandelbaum R., 2024, Proceedings of the AAAI Conference on Artificial Intelligence, 38, 12754

Jeffrey N., Alsing J., Lanusse F., 2021, Monthly Notices of the Royal Astronomical Society, 501, 954

Joachimi B., Semboloni E., Bett P. E., Hartlap J., Hilbert S., Hoekstra H., Schneider P., Schrabback T., 2013, MNRAS, 431, 477

Joachimi B., et al., 2015, Space Science Reviews, 193, 1

Jordan K., 2023, arXiv e-prints, p. arXiv:2304.01910

Kiessling A., et al., 2015, Space Sci. Rev., 193, 67

Kilbinger M., 2015, Reports on Progress in Physics, 78, 086901

Kirk D., et al., 2015, Space Sci. Rev., 193, 139

Lagos C. d. P., Tobar R. J., Robotham A. S. G., Obreschkow D., Mitchell P. D., Power C., Elahi P. J., 2018, MNRAS, 481, 3573

Lee J., Springel V., Pen U.-L., Lemson G., 2008, MNRAS, 389, 1266

Lu T., Haiman Z., Li X., 2023, MNRAS, 521, 2050

Mandelbaum R., 2018, ARA&A, 56, 393

Mandelbaum R., et al., 2011, MNRAS, 410, 844

Mandelbaum R., et al., 2014, ApJS, 212, 5

Marinacci F., et al., 2018, Mon. Not. Roy. Astron. Soc., 480, 5113

Mathieson L., Moscato P., 2019, An Introduction to Proximity Graphs. Springer International Publishing, Cham, pp 213–233, doi:10.1007/978-3-030-06222-4_4, https://doi.org/10.1007/978-3-030-06222-4_4

Matthies S., Muller J., Vinel G. W., 1988, Textures and Microstructures, 10, 77

Naiman J. P., et al., 2018, MNRAS, 477, 1206

Nelson D., et al., 2018, Mon. Not. Roy. Astron. Soc., 475, 624

Nelson D., et al., 2019, Computational Astrophysics and Cosmology, 6, 2

Nikolayev D. I., Savyolov T. I., 1970, Textures and Microstructures, 29

Pillepich A., et al., 2018a, MNRAS, 473, 4077

Pillepich A., et al., 2018b, MNRAS, 475, 648

Planck Collaboration et al., 2016, A&A, 594, A13
Samuroff S., et al., 2023, MNRAS, 524, 2195

Satorras V. G., Hoogeboom E., Welling M., 2021, in Meila M., Zhang T., eds, Proceedings of Machine Learning Research Vol. 139, Proceedings of the 38th International Conference on Machine Learning. PMLR, pp 9323–9332, https://proceedings.mlr.press/v139/satorras21a.html

Somerville R. S., Davé R., 2015, ARA&A, 53, 51

Song Y., Ermon S., 2019, in Advances in Neural Information Processing Systems. Curran Associates, Inc., https://proceedings.neurips.cc/paper/2019/file/3001ef257407d5a371a96dcd947c7d93-Paper.pdf

Song Y., Ermon S., 2020, Improved Techniques for Training Score-Based Generative Models, doi:10.48550/arXiv.2006.09011, http://arxiv.org/abs/2006.09011

Song Y., Sohl-Dickstein J., Kingma D. P., Kumar A., Ermon S., Poole B., 2021, in International Conference on Learning Representations. https://openreview.net/forum?id=PxTIG12RRHS

Springel V., 2010, MNRAS, 401, 791

Springel V., White S. D. M., Tormen G., Kauffmann G., 2001, MNRAS, 328, 726

Springel V., et al., 2018, Mon. Not. Roy. Astron. Soc., 475, 676

Stiskalek R., Bartlett D. J., Desmond H., Anbajagane D., 2022, MNRAS, 514, 4026

Troxel M. A., Ishak M., 2015, Phys. Rep., 558, 1

Van Alfen N., Campbell D., Blazek J., Leonard C. D., Lanusse F., Hearin A., Mandelbaum R., The LSST Dark Energy Science Collaboration 2023, arXiv e-prints, p. arXiv:2311.07374

Vilalta R., 2018, in Journal of Physics Conference Series. IOP, p. 052014 (arXiv:1812.10403), doi:10.1088/1742-6596/1085/5/052014

Vogelsberger M., Marinacci F., Torrey P., Puchwein E., 2020, Nature Reviews Physics, 2, 42

Wechsler R. H., Tinker J. L., 2018, ARA&A, 56, 435

Weinberg D. H., Bullock J. S., Governato F., Kuzio de Naray R., Peter A. H. G., 2015, Proceedings of the National Academy of Science, 112, 12249

Weinberger R., et al., 2017, MNRAS, 465, 3291

Yang L., et al., 2023, ACM Comput. Surv., 56

$$P_1 = h f(t_i, y_i) \qquad (B2)$$

$$P_2 = h f(t_i + \frac{1}{2}h, \frac{1}{2}P_1 + y_i) \qquad (B3)$$

$$y_{i+1} = P_2 + y_i \qquad (B4)$$

## B2 SO(3) manifold case

Let us define $\phi$ to be the map from Lie algebra $\mathcal{G}$ to the Lie group $G$, for example the exponential map ($\phi = \exp m$), then we can write:

$$F_1 = h f(t_i, y_i) \qquad (B5)$$

$$F_2 = h f\left(t_i + \frac{1}{2}h, \phi\left(\frac{1}{2}F_1\right)y_i\right) \qquad (B6)$$

$$y_{i+1} = \phi(F_2)y_i \qquad (B7)$$

These solutions can be obtained numerically for Eq. (20) for a given set of time steps $t_i = 0...T$ (Iserles et al. 2000). We use equal timesteps with $T = 1024$.

This paper has been typeset from a TeX/LA T E X file prepared by the author.

# APPENDIX A: DETAILS OF THE NEURAL NETWORK
ARCHITECTURE

The diffusion neural network has a size of {4096, 2048, 1024, 512, 256, 128} neurons each with GELU activation along with 3 layers of E(3) GNNs. We trained our models using the Adam optimizer with a learning rate of $10^{-4}$, exponential decay rates of $\beta_1 = 0.90$ and $\beta_2 = 0.95$, 1 million iterations, and then we decrease the learning rate to $10^{-5}$ for 0.5 million steps with a batch size of 128. NVIDIA Tesla A100 GPU was used as the hardware, with JAX, Jraph and DeepMind-Haiku Python libraries as the software.

# APPENDIX B: SOLVING DIFFERENTIAL EQUATIONS:
HEUN'S METHOD

Heun's method, also known as the explicit trapezoidal rule, is a numerical method of solving initial value ODEs.

## B1 Euclidean manifold case

Given an ODE of the form:

$$y'(t) = f(t, y(t)), \qquad y(t_0) = y_0, \qquad (B1)$$

with timestep $h$ and $t_{i+1} = t_i + h$, the approximate final value at the next integration point is given by: