Skip to content

Can we get rid of the 4D arrays in XraySourceBox? #552

@jordanflitter

Description

@jordanflitter

The problem

Each XraySourceBox contains now up to three 4D arrays (filtered_xray, filtered_sfr and filtered_sfr_mini) which take a lot of memory (in the case of Ly $\alpha$ multiple scattering, we have two more of these, filtered_sfr_lw and filtered_sfr_mini_lw to distinguish the straight line propagation of LW photons compared to the scattered trajectory of Ly $\alpha$ photons). Memorywise, I think this implementation could be improved.

Proposed Solution (Ly $\alpha$ flux)

Let us look for example on how filtered_sfr contributes to sfr_term here and how sfr_term eventually is used to compute $J_\alpha$. In that line, we have

$$ \text{sfr term}=\dot\rho_*^{\mathrm{(II)},R}\left(\mathbf x,z''\right)\times\Delta z''\left|\frac{dt''}{dz''}\right| $$

while the Ly $\alpha$ flux in the code is eventually given by

$$ J_{\alpha,*}\left(\mathbf x,z'\right)=\sum_{i=\mathrm{II,III}}\int_{z'}^{z_\mathrm{max}}\dot\rho_*^{i,R}\left(\mathbf x,z''\right)\Delta z''\left|\frac{dt''}{dz''}\right|\left(1+z'\right)^2\left(1+z''\right) $$

$$ \times\sum_{n=2}^{23}f_\mathrm{recycle}\left(n\right)N_*^i\left(\nu''_n\right)I^i\left(\nu''_n\right)\Theta\left(z{max}\left(z',n\right)-z''\right)\frac{c}{4\pi} $$

Here, the integral is computed numerically via the trapezoidal integration rule, and can be thought as a weighted sum of the filtered SFRD $\dot\rho_*^{i,R}$. Instead of saving the entire 4D array of $\dot\rho_*^{i,R}\left(\mathbf x,z''\right)$, we could write $J_{\alpha,*}$ as

$$ J_{\alpha,*}\left(\mathbf x,z'\right)=\sum_{i=\mathrm{II,III}}\sum_{R=R_\mathrm{min}}^{R_\mathrm{max}}\dot\rho_*^{i,R}\left(\mathbf x,z''\right)\times w_i\left(z',R\right) = \sum_{R=R_\mathrm{min}}^{R_\mathrm{max}}\sum_{i=\mathrm{II,III}}\dot\rho_*^{i,R}\left(\mathbf x,z''\right)\times w_i\left(z',R\right) $$

where $w_i\left(z',R\right)$ is a "global" quantity (doesn't depend on the cell $x$) that depends only on the the stellar population $i$, the snapshot redshift $z'$ and the filter radius $R$ (note that the dummy redshift $z''$ is determined from $z'$ and $R$). From here, it is easy to see that we don't really need to save the entire 4D array of $\dot\rho_*^{i,R}\left(\mathbf x,z''\right)$ (filtered_sfr in the code) but only save an "evolving" 3D box that evolves when we account higher filter radii. At the final step, the evolved 3D box is $J_{\alpha,*}\left(\mathbf x,z'\right)$.

Proposed Solution (X-ray flux)

The expression for X-ray flux is more complicated as it depends also on frequency and thus the quantities of interest involve integrals over frequency. Here, filtered_xray contributes to xray_sfr here. This line reads (note that $L_X$, which is halo_xray in the code, contains both contributions from popII and popIII stars)

$$ \text{xray sfr}=\frac{10^{-38}L_X^{R}\left(\mathbf x,z''\right)}{V_\mathrm{cell}}\times\Delta z''\left|\frac{dt''}{dz''}\right|\times\left(1+z''\right)^{-\alpha_X}\times10^{38} $$

This quantity is then used to compute more physical quantities that involve integration over frequency. For example, the X-ray heating is then given by

$$ \Gamma_X\left(\mathbf x,z'\right)=\int_{z'}^{z_\mathrm{max}}\left(1+z''\right)^{-\alpha_X}\frac{L_X^{R}\left(\mathbf x,z''\right)}{V_\mathrm{cell}}\Delta z''\left|\frac{dt''}{dz''}\right| $$

$$ \times \int_{\max{[\nu_0,\nu_{\tau_X=1}}]}^{\nu_\mathrm{max}} d\nu\sum_i f_\mathrm{heat}\left(x_\mathrm{e}\left(\mathbf x,z'\right),\nu-\nu^\mathrm{th}\right)h_\mathrm{P}\left(\nu-\nu^\mathrm{th}_i\right)f_ix_i\left(\mathbf x,z'\right)\sigma_i\left(\nu\right)\left(\frac{\nu}{\nu_0}\right)^{-\alpha_X-1} $$

$$ \times \frac{\left[\int_{\nu_0}^{\nu_\mathrm{max}} d\nu\left(\frac{\nu}{\nu_0}\right)^{-\alpha_X}\right]^{-1}}{h_\mathrm{P},\nu_0}c\left(1+z'\right)^{\alpha_X+3} $$

Here, the frequency integral is computed via the trapezoidal integration rule. Similarly as above, we may write this expression as

$$ \Gamma_X\left(\mathbf x,z'\right)= \sum_{R=R_\mathrm{min}}^{R_\mathrm{max}}\sum_{\nu=\nu_\mathrm{min}}^{\nu_\mathrm{max}}L_X^{R}\left(\mathbf x,z''\right)\times w\left(\mathbf x,z',R,\nu\right)=\sum_{R=R_\mathrm{min}}^{R_\mathrm{max}}L_X^{R}\left(\mathbf x,z''\right)\times\tilde w\left(\mathbf x,z',R\right) $$

where $\tilde w\left(\mathbf x,z',R\right)$ is just as a summed version of $w\left(\mathbf x,z',R,\nu\right)$. Note that we cannot switch the order between the $\nu$ summation and the $R$ summation, since the former depends on the latter through the minimum frequency we consider: $\nu_\mathrm{min}=\max{[\nu_0,\nu_{\tau_X=1}}]$ and $\tau_X$ depends on $R$. Regardless, here we ended with a local weight (that depends on the cell $\mathbf x$). But still, we see that we don't need to save the entire 4D array $L_X^{R}\left(\mathbf x,z''\right)$ (filtered_xray in the code), but rather the "evolving" 3D box of it, just like in the case of the Ly $\alpha$ flux.

Conclusions

I think it is indeed possible to come up with an implementation that doesn't require saving 4D arrays, only 3D boxes, thus reducing the amount of used memory. I also believe that this proposed implementation will not harm the runtime of the calculation, and I also think it could be done at the python wrapper for that purpose (instead the C backend), without slowing down things too much, as all we need to do is basically sum and multiply 3D boxes. I imagine replacing the XraySourceBox with something like RadiationFields that includes 3D boxes of physical radiation fields like $J_{\alpha,*}$, $J_\mathrm{LW}$, $J_{\alpha,X}$, $\Gamma_X$ and $\Lambda_X$ (the last one is X-ray ionization rate). Then, the computation of the spin temperature becomes much simplified as it would only require promoting $T_k$ and $x_e$ to the next redshift iteration and use the promoted values together with $J_\alpha$ to compute the spin temperature. This kind of implementation could also be generalized for USE_HALO_FIELD=False (though I need to give it more thought), which would make ComputeTsBox much simplified.

This is a rather big change because it requires modifying both wrapper and backend. Before making a PR for it, I'd like to hear @daviesje and @steven-murray thoughts on it.

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions