csiborgtools/scripts_plots/borg_acl.ipynb
Richard Stiskalek 9e4b34f579
Overlap fixing and more (#107)
* Update README

* Update density field reader

* Update name of SDSSxALFAFA

* Fix quick bug

* Add little fixes

* Update README

* Put back fit_init

* Add paths to initial snapshots

* Add export

* Remove some choices

* Edit README

* Add Jens' comments

* Organize imports

* Rename snapshot

* Add additional print statement

* Add paths to initial snapshots

* Add masses to the initial files

* Add normalization

* Edit README

* Update README

* Fix bug in CSiBORG1 so that does not read fof_00001

* Edit README

* Edit README

* Overwrite comments

* Add paths to init lag

* Fix Quijote path

* Add lagpatch

* Edit submits

* Update README

* Fix numpy int problem

* Update README

* Add a flag to keep the snapshots open when fitting

* Add a flag to keep snapshots open

* Comment out some path issue

* Keep snapshots open

* Access directly snasphot

* Add lagpatch for CSiBORG2

* Add treatment of x-z coordinates flipping

* Add radial velocity field loader

* Update README

* Add lagpatch to Quijote

* Fix typo

* Add setter

* Fix typo

* Update README

* Add output halo cat as ASCII

* Add import

* Add halo plot

* Update README

* Add evaluating field at radial distanfe

* Add field shell evaluation

* Add enclosed mass computation

* Add BORG2 import

* Add BORG boxsize

* Add BORG paths

* Edit run

* Add BORG2 overdensity field

* Add bulk flow clauclation

* Update README

* Add new plots

* Add nbs

* Edit paper

* Update plotting

* Fix overlap paths to contain simname

* Add normalization of positions

* Add default paths to CSiBORG1

* Add overlap path simname

* Fix little things

* Add CSiBORG2 catalogue

* Update README

* Add import

* Add TNG density field constructor

* Add TNG density

* Add draft of calculating BORG ACL

* Fix bug

* Add ACL of enclosed density

* Add nmean acl

* Add galaxy bias calculation

* Add BORG acl notebook

* Add enclosed mass calculation

* Add TNG300-1 dir

* Add TNG300 and BORG1 dir

* Update nb
2024-01-30 16:14:07 +00:00

50 KiB

Open In Colab

In [8]:
!pip install --quiet git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git
In [1]:
%pylab inline

import jax
import jax.numpy as jnp
import jax_cosmo as jc

from jax.experimental.ode import odeint

from jaxpm.painting import cic_paint
from jaxpm.pm import linear_field, lpt, make_ode_fn
Populating the interactive namespace from numpy and matplotlib
In [7]:
mesh_shape= [256, 256, 256]
box_size  = [256.,256.,256.]
snapshots = jnp.linspace(0.1,1.,2)

@jax.jit
def run_simulation(omega_c, sigma8):
    # Create a small function to generate the matter power spectrum
    k = jnp.logspace(-4, 1, 128)
    pk = jc.power.linear_matter_power(jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)
    pk_fn = lambda x: jc.scipy.interpolate.interp(x.reshape([-1]), k, pk).reshape(x.shape)

    # Create initial conditions
    initial_conditions = linear_field(mesh_shape, box_size, pk_fn, seed=jax.random.PRNGKey(0))

    # Create particles
    particles = jnp.stack(jnp.meshgrid(*[jnp.arange(s) for s in mesh_shape]),axis=-1).reshape([-1,3])

    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)
    
    # Initial displacement
    dx, p, f = lpt(cosmo, initial_conditions, particles, 0.1)
    
    # Evolve the simulation forward
    res = odeint(make_ode_fn(mesh_shape), [particles+dx, p], snapshots, cosmo, rtol=1e-5, atol=1e-5)
    
    # Return the simulation volume at requested 
    return res[0]
In [10]:
res = run_simulation(0.25, 0.8)
%timeit res = run_simulation(0.25, 0.8)
1 loop, best of 5: 29.2 s per loop
In [11]:
figure(figsize=[10,5])
subplot(121)
imshow(cic_paint(jnp.zeros(mesh_shape), res[0]).sum(axis=0),cmap='gist_stern')
subplot(122)
imshow(cic_paint(jnp.zeros(mesh_shape), res[1]).sum(axis=0),cmap='gist_stern')
Out[11]:
<matplotlib.image.AxesImage at 0x7f3502382c90>
No description has been provided for this image
In [ ]: