In [8]:

import astropy.units as u
from xsb_fluc.data.cluster import Cluster

cluster = Cluster(
    imglink='data/A2142/mosaic_a2142.fits.gz',
    explink='data/A2142/mosaic_a2142_expo.fits.gz',
    bkglink='data/A2142/mosaic_a2142_bkg.fits.gz',
    reglink='data/A2142/src_ps.reg',
    nhlink='data/A2142/A2142_nh.fits',
    ra=239.58615,
    dec=27.229433,
    r_500=1.403*u.Mpc, 
    redshift=0.09,
)

cluster = cluster.reduce_to_r500(1).rebin(5)
cluster.plot_cluster()

import astropy.units as u from xsb_fluc.data.cluster import Cluster cluster = Cluster( imglink='data/A2142/mosaic_a2142.fits.gz', explink='data/A2142/mosaic_a2142_expo.fits.gz', bkglink='data/A2142/mosaic_a2142_bkg.fits.gz', reglink='data/A2142/src_ps.reg', nhlink='data/A2142/A2142_nh.fits', ra=239.58615, dec=27.229433, r_500=1.403*u.Mpc, redshift=0.09, ) cluster = cluster.reduce_to_r500(1).rebin(5) cluster.plot_cluster()

WARNING: FITSFixedWarning: RADECSYS= 'FK5 ' / Stellar reference frame 
the RADECSYS keyword is deprecated, use RADESYSa. [astropy.wcs.wcs]
WARNING: FITSFixedWarning: EQUINOX = '2000.0 ' / Coordinate system equinox 
a floating-point value was expected. [astropy.wcs.wcs]

In [9]:

import json 
import haiku as hk
import jax.numpy as jnp
from xsb_fluc.simulation.mock_image import MockXrayCountsBetaModel

# We reload the parameters and parse them to the appropriate format
with open('data/A2142/posterior_parameters.json', 'r') as file:
    posterior_parameters = {key:jnp.asarray(value) for key, value in json.load(file).items()}

images_simulator = hk.without_apply_rng(hk.transform(lambda : MockXrayCountsBetaModel(cluster)()))

import json import haiku as hk import jax.numpy as jnp from xsb_fluc.simulation.mock_image import MockXrayCountsBetaModel # We reload the parameters and parse them to the appropriate format with open('data/A2142/posterior_parameters.json', 'r') as file: posterior_parameters = {key:jnp.asarray(value) for key, value in json.load(file).items()} images_simulator = hk.without_apply_rng(hk.transform(lambda : MockXrayCountsBetaModel(cluster)()))

In [11]:

class SpectraSimulator(hk.Module):
    
    pass

class SpectraSimulator(hk.Module): pass

Simulate spectra¶

In [12]:

import os
import numpy as np
import sys
import jax
import haiku as hk
from xsb_fluc.data.mean_model import MeanModel

from xsb_fluc.utils.mock_image import ImageGenerator
from xsb_fluc.data.power_spectrum import PowerSpectrum
from xsb_fluc.utils.config import config
from xsb_fluc.utils.misc import rng_key
from xsb_fluc.utils.ps import chexmate_low_scale
import jax.numpy as jnp
from jax import config as jax_config

jax_config.update("jax_enable_x64", True)




mean_model = MeanModel.from_data(data)
reduced_data = data.reduce_fov(mean_model, aperture)
reduced_mean_model = MeanModel.from_data(reduced_data)

mask_list = [(reduced_data.exp > 0.) & (reduced_mean_model.rad < 1.)]
scales_list = [jnp.geomspace(chexmate_low_scale(cluster.z), 1., 20)]
ps_func_list = [jax.pmap(lambda i: hk.without_apply_rng(
    hk.transform(lambda img: PowerSpectrum(reduced_data, mask=mask_list[0])(img, scales_list[0]))).apply(None, i))]

gen = ImageGenerator(reduced_data, reduced_mean_model)
bkg = jnp.asarray(reduced_data.bkg[None, ...])
exp = jnp.asarray(reduced_data.exp[None, ...])
C = jnp.ones_like(exp) * 10 ** float(reduced_mean_model.posterior_median['log_bkg'])
#Reference simulations here

x_regions = [[] for region in regions]
x_regions_rel = [[] for region in regions_rel]
theta = []

for _ in range(1000):

    log_sigma = -0.5 * jnp.ones((1,))
    log_inj = -1.3 * jnp.ones((1,))
    alpha = 11 / 3 * jnp.ones((1,))

    theta.append(jnp.hstack([log_sigma, log_inj, alpha]))

    counts_perturbed, counts_rest, pars_val = gen(log_sigma, log_inj, alpha)

    null = jax.random.poisson(rng_key(), counts_rest)
    img_list_abs = jnp.nan_to_num((counts_perturbed - counts_rest) / (2 * exp)) * (exp > 0.)
    img_list_abs_null = jnp.nan_to_num((null - counts_rest) / (2 * exp)) * (exp > 0.)

    img_list_rel = jnp.nan_to_num(((counts_perturbed - bkg) / exp - C) / ((counts_rest - bkg) / exp - C)) * (exp > 0.)
    img_list_rel_null = jnp.nan_to_num(((null - bkg) / exp - C) / ((counts_rest - bkg) / exp - C)) * (exp > 0.)

    for x_reg, ps in zip(x_regions, ps_func_list):
        x_reg.append(jnp.log(ps(img_list_abs)) - jnp.log(ps(img_list_abs_null)))

    for x_reg, ps in zip(x_regions_rel, ps_func_list):
        x_reg.append(jnp.log(ps(img_list_rel)) - jnp.log(ps(img_list_rel_null)))

np.savetxt(os.path.join(config.RESULTS_PATH, f'mock_observables/{name}/ref/theta.txt'), np.stack(theta))

for x_reg, region in zip(x_regions, regions):
    np.savetxt(os.path.join(config.RESULTS_PATH, f'mock_observables/{name}/ref/x_{region}.txt'), np.vstack(x_reg))

for x_reg, region in zip(x_regions_rel, regions_rel):
    np.savetxt(os.path.join(config.RESULTS_PATH, f'mock_observables/{name}/ref/x_{region}.txt'), np.vstack(x_reg))

#Simulations with prior here

x_regions = [[] for region in regions]
theta = []
first_key = rng_key()

for _ in range(500000):

    key = jax.random.split(first_key, 5)
    first_key = key[0]

    log_sigma = jax.random.uniform(key[1], minval=-2., maxval=1., shape=(1,))
    log_inj = jax.random.uniform(key[2], minval=-2., maxval=1., shape=(1,))
    alpha = jax.random.uniform(key[3], minval=0., maxval=7., shape=(1,))

    theta.append(jnp.hstack([log_sigma, log_inj, alpha]))

    counts_perturbed, counts_rest, pars_val = gen(log_sigma, log_inj, alpha)

    null = jax.random.poisson(key[4], counts_rest)
    img_list_abs = jnp.nan_to_num((counts_perturbed - counts_rest) / (2 * exp)) * (exp > 0.)
    img_list_abs_null = jnp.nan_to_num((null - counts_rest) / (2 * exp)) * (exp > 0.)

    for x_reg, ps in zip(x_regions, ps_func_list):
        x_reg.append(jnp.log(ps(img_list_abs)) - jnp.log(ps(img_list_abs_null)))
"""
np.savetxt(os.path.join(config.RESULTS_PATH, f'mock_observables/{name}/theta.txt'), np.stack(theta))

for x_reg, region in zip(x_regions, regions):
    np.savetxt(os.path.join(config.RESULTS_PATH, f'mock_observables/{name}/x_{region}.txt'), np.vstack(x_reg))
"""

import os import numpy as np import sys import jax import haiku as hk from xsb_fluc.data.mean_model import MeanModel from xsb_fluc.utils.mock_image import ImageGenerator from xsb_fluc.data.power_spectrum import PowerSpectrum from xsb_fluc.utils.config import config from xsb_fluc.utils.misc import rng_key from xsb_fluc.utils.ps import chexmate_low_scale import jax.numpy as jnp from jax import config as jax_config jax_config.update("jax_enable_x64", True) mean_model = MeanModel.from_data(data) reduced_data = data.reduce_fov(mean_model, aperture) reduced_mean_model = MeanModel.from_data(reduced_data) mask_list = [(reduced_data.exp > 0.) & (reduced_mean_model.rad < 1.)] scales_list = [jnp.geomspace(chexmate_low_scale(cluster.z), 1., 20)] ps_func_list = [jax.pmap(lambda i: hk.without_apply_rng( hk.transform(lambda img: PowerSpectrum(reduced_data, mask=mask_list[0])(img, scales_list[0]))).apply(None, i))] gen = ImageGenerator(reduced_data, reduced_mean_model) bkg = jnp.asarray(reduced_data.bkg[None, ...]) exp = jnp.asarray(reduced_data.exp[None, ...]) C = jnp.ones_like(exp) * 10 ** float(reduced_mean_model.posterior_median['log_bkg']) #Reference simulations here x_regions = [[] for region in regions] x_regions_rel = [[] for region in regions_rel] theta = [] for _ in range(1000): log_sigma = -0.5 * jnp.ones((1,)) log_inj = -1.3 * jnp.ones((1,)) alpha = 11 / 3 * jnp.ones((1,)) theta.append(jnp.hstack([log_sigma, log_inj, alpha])) counts_perturbed, counts_rest, pars_val = gen(log_sigma, log_inj, alpha) null = jax.random.poisson(rng_key(), counts_rest) img_list_abs = jnp.nan_to_num((counts_perturbed - counts_rest) / (2 * exp)) * (exp > 0.) img_list_abs_null = jnp.nan_to_num((null - counts_rest) / (2 * exp)) * (exp > 0.) img_list_rel = jnp.nan_to_num(((counts_perturbed - bkg) / exp - C) / ((counts_rest - bkg) / exp - C)) * (exp > 0.) img_list_rel_null = jnp.nan_to_num(((null - bkg) / exp - C) / ((counts_rest - bkg) / exp - C)) * (exp > 0.) for x_reg, ps in zip(x_regions, ps_func_list): x_reg.append(jnp.log(ps(img_list_abs)) - jnp.log(ps(img_list_abs_null))) for x_reg, ps in zip(x_regions_rel, ps_func_list): x_reg.append(jnp.log(ps(img_list_rel)) - jnp.log(ps(img_list_rel_null))) np.savetxt(os.path.join(config.RESULTS_PATH, f'mock_observables/{name}/ref/theta.txt'), np.stack(theta)) for x_reg, region in zip(x_regions, regions): np.savetxt(os.path.join(config.RESULTS_PATH, f'mock_observables/{name}/ref/x_{region}.txt'), np.vstack(x_reg)) for x_reg, region in zip(x_regions_rel, regions_rel): np.savetxt(os.path.join(config.RESULTS_PATH, f'mock_observables/{name}/ref/x_{region}.txt'), np.vstack(x_reg)) #Simulations with prior here x_regions = [[] for region in regions] theta = [] first_key = rng_key() for _ in range(500000): key = jax.random.split(first_key, 5) first_key = key[0] log_sigma = jax.random.uniform(key[1], minval=-2., maxval=1., shape=(1,)) log_inj = jax.random.uniform(key[2], minval=-2., maxval=1., shape=(1,)) alpha = jax.random.uniform(key[3], minval=0., maxval=7., shape=(1,)) theta.append(jnp.hstack([log_sigma, log_inj, alpha])) counts_perturbed, counts_rest, pars_val = gen(log_sigma, log_inj, alpha) null = jax.random.poisson(key[4], counts_rest) img_list_abs = jnp.nan_to_num((counts_perturbed - counts_rest) / (2 * exp)) * (exp > 0.) img_list_abs_null = jnp.nan_to_num((null - counts_rest) / (2 * exp)) * (exp > 0.) for x_reg, ps in zip(x_regions, ps_func_list): x_reg.append(jnp.log(ps(img_list_abs)) - jnp.log(ps(img_list_abs_null))) """ np.savetxt(os.path.join(config.RESULTS_PATH, f'mock_observables/{name}/theta.txt'), np.stack(theta)) for x_reg, region in zip(x_regions, regions): np.savetxt(os.path.join(config.RESULTS_PATH, f'mock_observables/{name}/x_{region}.txt'), np.vstack(x_reg)) """

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[12], line 9
      6 from xsb_fluc.data.mean_model import MeanModel
      8 from xsb_fluc.utils.mock_image import ImageGenerator
----> 9 from xsb_fluc.data.power_spectrum import PowerSpectrum
     10 from xsb_fluc.utils.config import config
     11 from xsb_fluc.utils.misc import rng_key

ModuleNotFoundError: No module named 'xsb_fluc.data.power_spectrum'