mirror of
https://github.com/Richard-Sti/csiborgtools_public.git
synced 2025-07-08 15:21:11 +00:00
* 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
50 KiB
50 KiB
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
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)
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]:
In [ ]: