Quickstart

Let’s see how to make a sychrotron simulation quickly! All required files can be found in example/ folder in the repo

#%env XLA_PYTHON_CLIENT_PREALLOCATE=false
%env XLA_PYTHON_CLIENT_ALLOCATOR=platform

env: XLA_PYTHON_CLIENT_ALLOCATOR=platform

import os
import sys
import jax
jax.config.update("jax_enable_x64", True)
#sys.path.append('../synax/')

import synax,importlib
import jax.numpy as jnp
import interpax
import healpy as hp
import numpy as np
import matplotlib.pyplot as plt
from functools import partial
import scipy.constants as const

#for debug
def reload_package(package):
    importlib.reload(package)
    for attribute_name in dir(package):
        attribute = getattr(package, attribute_name)
        if type(attribute) == type(package):
            importlib.reload(attribute)

reload_package(synax)

Generating integration points

First thing to do is to generate the coordinates of all integration points. You can use synax.coords.get_healpix_positions() to quickly generate these coordinates for a healpix map.

poss: coordinates for integration points, e.g. poss[0] represents x coordinates.

dls: length of integration intervals.

nhats: unit vector of different line-of-sights (LoS).

You don’t need to understand dls and nhats and just need to pass them to synax.synax.Synax.sim().

nside = 64
num_int_points = 512

poss,dls,nhats = synax.coords.get_healpix_positions(nside=nside,num_int_points = num_int_points )

2024-08-01 14:35:49.395744: W external/xla/xla/service/gpu/nvptx_compiler.cc:765] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.5.40). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.

Let’s see the projection of 500-th integration points for every 10 LoS to the \(x-y\) plane.

plt.scatter(poss[0,::10,500],poss[1,::10,500],alpha=0.05)
poss.shape

(3, 49152, 512)

Generating fields

Let’s then generating different fields. First let’s generate an analytical cosmic ray electron spectrum cosntant \(C\) field:

\[N(\gamma)d\gamma = C\gamma^{-p}d\gamma\]

. The model here is the WMAP model presented in Page et al. 2007

You can define a field generator synax.cfield.C_WMAP() first. It need the integration points coordinates poss. Then C_generator.C_field() takes model parameters and calculates the \(C\) field at these points.

C_generator = synax.cfield.C_WMAP(poss)

C_field = C_generator.C_field()

Let’s generate the thermal electron density \(n_e\) model with the intepolation on 3D grids.

First let’s define the coordinates of the voxels in the 3D grid and read the 3D grids in.

nx,ny,nz = 256,256,64 # number of voxels along xyz axis

xs,step = jnp.linspace(-20,20,nx,endpoint=False,retstep=True)
xs = xs + step*0.5# get x coordinates

ys,step = jnp.linspace(-20,20,ny,endpoint=False,retstep=True)
ys = ys + step*0.5# get x coordinates

zs,step = jnp.linspace(-5,5,nz,endpoint=False,retstep=True)
zs = zs + step*0.5# get x coordinates

coords = jnp.meshgrid(xs,ys,zs,indexing='ij')
coords[0].shape

(256, 256, 64)

tereg = np.load('te.npy')

Let’s visualize the grid \(n_e\) field.

plt.imshow(tereg[:,:,32],vmax=0.2)
plt.colorbar()

<matplotlib.colorbar.Colorbar at 0x7f70f0231e10>

Similarly, define a generator first. Now despite of the integration point coordinates, it also needs the coordinates of your 3D grid voxels.

TE_generator = synax.tefield.TE_grid(poss,(xs,ys,zs))

%%time
TE_field = TE_generator.TE_field(tereg)
TE_field.shape

CPU times: user 3.49 s, sys: 7.58 s, total: 11.1 s
Wall time: 11.1 s

(49152, 512)

Last thing is repeat the same thing on \(\mathbf{B}\) field

lsa_params = {"b0":1.2,
               "psi0":27.0*np.pi/180,
               "psi1":0.9*np.pi/180,
               "chi0":25.0*np.pi/180}

B_generator = synax.bfield.B_lsa(poss)
B_field = B_generator.B_field(lsa_params)
B_field.shape

(49152, 512, 3)

Simulate!

Now we have every field we need. Last two things to specify is the frequency in GHz and spectral index. We can initialize a simer synax.synax.Synax() first.

simer = synax.synax.Synax()
freq = 2.4
spectral_index = 3.

Then give everything we have to this simer by calling simer.sim()

sync = simer.sim(freq,B_field,C_field,TE_field,nhats,dls,spectral_index)

We can visualize the results now!

sync['I'] gives I map;

sync['Q'] gives Q map;

sync['U'] gives U map;

They are all in units of K.

plt.figure(dpi=200)
np.random.seed(42)
plt.subplot(131)
hp.mollview(sync['I'],format='%.2g',norm='hist',cmap='coolwarm',hold=True,title='Synax I')

plt.subplot(132)
hp.mollview(sync['Q'],format='%.2g',norm='hist',cmap='coolwarm',hold=True,title='Synax Q')

plt.subplot(133)
hp.mollview(sync['U'],format='%.2g',norm='hist',cmap='coolwarm',hold=True,title='Synax U')

I hope follow these lines, you did generate the first simulation of your own with synax!