Skip to content

simulation

xsb_fluc.simulation.grid

FourierGrid3D

The equivalent of a Spatial grid in Fourier space

Source code in src/xsb_fluc/simulation/grid.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
class FourierGrid3D:
    """
    The equivalent of a Spatial grid in Fourier space
    """

    def __init__(self, spatial_grid: SpatialGrid3D):
        """
        Constructor of a FourierGrid3D object as the dual of a SpatialGrid3D
        """

        self.kx = fft.fftfreq(len(spatial_grid.x), d=spatial_grid.x[1] - spatial_grid.x[0])
        self.ky = fft.fftfreq(len(spatial_grid.y), d=spatial_grid.y[1] - spatial_grid.y[0])
        self.kz = fft.rfftfreq(len(spatial_grid.z), d=spatial_grid.z[1] - spatial_grid.z[0])
        KX, KY, KZ = jnp.meshgrid(self.kx, self.ky, self.kz, indexing='ij', sparse=True)
        self.K = jnp.sqrt(KX ** 2 + KY ** 2 + KZ ** 2)

        self.shape = self.K.shape
__init__(spatial_grid)

Constructor of a FourierGrid3D object as the dual of a SpatialGrid3D

Source code in src/xsb_fluc/simulation/grid.py 
64
65
66
67
68
69
70
71
72
73
74
75
def __init__(self, spatial_grid: SpatialGrid3D):
    """
    Constructor of a FourierGrid3D object as the dual of a SpatialGrid3D
    """

    self.kx = fft.fftfreq(len(spatial_grid.x), d=spatial_grid.x[1] - spatial_grid.x[0])
    self.ky = fft.fftfreq(len(spatial_grid.y), d=spatial_grid.y[1] - spatial_grid.y[0])
    self.kz = fft.rfftfreq(len(spatial_grid.z), d=spatial_grid.z[1] - spatial_grid.z[0])
    KX, KY, KZ = jnp.meshgrid(self.kx, self.ky, self.kz, indexing='ij', sparse=True)
    self.K = jnp.sqrt(KX ** 2 + KY ** 2 + KZ ** 2)

    self.shape = self.K.shape

SpatialGrid3D

Helper function to define a spatial grid to simulate a 3D cluster

Source code in src/xsb_fluc/simulation/grid.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class SpatialGrid3D:
    """
    Helper function to define a spatial grid to simulate a 3D cluster
    """

    def __init__(self, pixsize=2./1000., shape=(100, 100), crop_r_500=5.):
        r"""
        Constructor for SpatialGrid3D.

        Parameters:
            pixsize (float): Size of the pixel in $R_{500}$ units.
            shape (tuple): Shape of the cube on the sky.
            crop_r_500 (float): Size along the line of sight in $\pm R_{500}$.
        """

        self.pixsize = pixsize

        x_size, y_size = shape
        z_size = np.ceil(2*crop_r_500/pixsize).astype(int)

        self.x = jnp.linspace(-self.pixsize * (x_size - 1)/2, self.pixsize * (x_size - 1)/2, x_size)
        self.y = jnp.linspace(-self.pixsize * (y_size - 1)/2, self.pixsize * (y_size - 1)/2, y_size)
        self.z = jnp.linspace(-self.pixsize * (z_size - 1)/2, self.pixsize * (z_size - 1)/2, z_size)
        self.X, self.Y, self.Z = jnp.meshgrid(self.x, self.y, self.z, indexing='ij', sparse=True)
        self.y_i, self.x_i = jnp.indices(shape)
        self.R = jnp.sqrt(self.X**2 + self.Y**2 + self.Z**2)
        self.shape = (x_size, y_size, z_size)
        self.volume = np.prod(self.shape)*pixsize**3

    @classmethod
    def from_data(cls, data: Cluster, crop_r_500=5., pixsize=None):
        """
        Constructor for SpatialGrid3D using a Cluster object

        Parameters:
            crop_r_500: size along the line of sight.
            pixsize: should be None since it is read from the Cluster
        """

        if pixsize is None:

            return cls(pixsize=(data.kpc_per_pixel/data.r_500).to(1/u.pixel).value,
                       shape = data.shape, 
                       crop_r_500=crop_r_500)

        else : 

            return cls(pixsize=pixsize,
                       shape = data.shape, 
                       crop_r_500=crop_r_500)
__init__(pixsize=2.0 / 1000.0, shape=(100, 100), crop_r_500=5.0)

Constructor for SpatialGrid3D.

Parameters:

Name Type Description Default
pixsize float

Size of the pixel in \(R_{500}\) units.

 2.0 / 1000.0
shape tuple

Shape of the cube on the sky.

 (100, 100)
crop_r_500 float

Size along the line of sight in \(\pm R_{500}\).

 5.0
Source code in src/xsb_fluc/simulation/grid.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def __init__(self, pixsize=2./1000., shape=(100, 100), crop_r_500=5.):
    r"""
    Constructor for SpatialGrid3D.

    Parameters:
        pixsize (float): Size of the pixel in $R_{500}$ units.
        shape (tuple): Shape of the cube on the sky.
        crop_r_500 (float): Size along the line of sight in $\pm R_{500}$.
    """

    self.pixsize = pixsize

    x_size, y_size = shape
    z_size = np.ceil(2*crop_r_500/pixsize).astype(int)

    self.x = jnp.linspace(-self.pixsize * (x_size - 1)/2, self.pixsize * (x_size - 1)/2, x_size)
    self.y = jnp.linspace(-self.pixsize * (y_size - 1)/2, self.pixsize * (y_size - 1)/2, y_size)
    self.z = jnp.linspace(-self.pixsize * (z_size - 1)/2, self.pixsize * (z_size - 1)/2, z_size)
    self.X, self.Y, self.Z = jnp.meshgrid(self.x, self.y, self.z, indexing='ij', sparse=True)
    self.y_i, self.x_i = jnp.indices(shape)
    self.R = jnp.sqrt(self.X**2 + self.Y**2 + self.Z**2)
    self.shape = (x_size, y_size, z_size)
    self.volume = np.prod(self.shape)*pixsize**3
from_data(data, crop_r_500=5.0, pixsize=None) classmethod

Constructor for SpatialGrid3D using a Cluster object

Parameters:

Name Type Description Default
crop_r_500

size along the line of sight.

 5.0
pixsize

should be None since it is read from the Cluster

 None
Source code in src/xsb_fluc/simulation/grid.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
@classmethod
def from_data(cls, data: Cluster, crop_r_500=5., pixsize=None):
    """
    Constructor for SpatialGrid3D using a Cluster object

    Parameters:
        crop_r_500: size along the line of sight.
        pixsize: should be None since it is read from the Cluster
    """

    if pixsize is None:

        return cls(pixsize=(data.kpc_per_pixel/data.r_500).to(1/u.pixel).value,
                   shape = data.shape, 
                   crop_r_500=crop_r_500)

    else : 

        return cls(pixsize=pixsize,
                   shape = data.shape, 
                   crop_r_500=crop_r_500)

xsb_fluc.simulation.cube

EmissivityCube

Bases: Module

Compute an emissivity cube for a given cluster. Used as a part of simulations.

Source code in src/xsb_fluc/simulation/cube.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class EmissivityCube(hk.Module):
    """
    Compute an emissivity cube for a given cluster. Used as a part of simulations.
    """

    def __init__(self, data: Cluster):
        """
        Constructor for EmissivityCube using a Cluster object
        """

        super(EmissivityCube, self).__init__()
        self.emissivity = XrayEmissivity.from_data(data)
        self.radius = EllipseRadius.from_data(data)
        self.spatial_grid = SpatialGrid3D.from_data(data, crop_r_500=5)
        self.nh = data.nh
        self.x_i, self.x_j = self.spatial_grid.x_i, self.spatial_grid.y_i
        self.Z = self.spatial_grid.Z

    def __call__(self) -> jax.Array:
        """
        Return the emissivity cube.
        """

        Rho = self.radius(self.x_i, self.x_j) 
        R = jnp.sqrt(Rho[..., None]**2 + self.Z**2)
        epsilon = self.emissivity(R, self.nh[..., None])

        return epsilon
__call__()

Return the emissivity cube.

Source code in src/xsb_fluc/simulation/cube.py 
32
33
34
35
36
37
38
39
40
41
def __call__(self) -> jax.Array:
    """
    Return the emissivity cube.
    """

    Rho = self.radius(self.x_i, self.x_j) 
    R = jnp.sqrt(Rho[..., None]**2 + self.Z**2)
    epsilon = self.emissivity(R, self.nh[..., None])

    return epsilon
__init__(data)

Constructor for EmissivityCube using a Cluster object

Source code in src/xsb_fluc/simulation/cube.py 
19
20
21
22
23
24
25
26
27
28
29
30
def __init__(self, data: Cluster):
    """
    Constructor for EmissivityCube using a Cluster object
    """

    super(EmissivityCube, self).__init__()
    self.emissivity = XrayEmissivity.from_data(data)
    self.radius = EllipseRadius.from_data(data)
    self.spatial_grid = SpatialGrid3D.from_data(data, crop_r_500=5)
    self.nh = data.nh
    self.x_i, self.x_j = self.spatial_grid.x_i, self.spatial_grid.y_i
    self.Z = self.spatial_grid.Z

FluctuationCube

Bases: Module

Compute a fluctuation cube for a given cluster. The density fluctuations are computed assuming a Gaussian Random Field defined with a turbulent power spectrum.

Source code in src/xsb_fluc/simulation/cube.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
class FluctuationCube(hk.Module):
    """
    Compute a fluctuation cube for a given cluster. The density fluctuations are computed
    assuming a Gaussian Random Field defined with a turbulent power spectrum.
    """

    def __init__(self, data: Cluster):
        """
        Constructor for FluctuationCube using a Cluster object
        """

        super(FluctuationCube, self).__init__()
        self.power_spectrum = KolmogorovPowerSpectrum()
        self.spatial_grid = SpatialGrid3D.from_data(data, crop_r_500=5)
        self.fourier_grid = FourierGrid3D(self.spatial_grid)
        self.K = self.fourier_grid.K.astype(np.float32)

    def __call__(self) -> jax.Array:
        """
        Return the fluctuation cube.
        """

        key = hk.next_rng_key()

        field_spatial = random.normal(key, shape=self.spatial_grid.shape)
        field_fourier = fft.rfftn(field_spatial)*jnp.sqrt(self.power_spectrum(self.K)/self.spatial_grid.pixsize**3)
        field_spatial = fft.irfftn(field_fourier, s=self.spatial_grid.shape)

        return field_spatial
__call__()

Return the fluctuation cube.

Source code in src/xsb_fluc/simulation/cube.py 
61
62
63
64
65
66
67
68
69
70
71
72
def __call__(self) -> jax.Array:
    """
    Return the fluctuation cube.
    """

    key = hk.next_rng_key()

    field_spatial = random.normal(key, shape=self.spatial_grid.shape)
    field_fourier = fft.rfftn(field_spatial)*jnp.sqrt(self.power_spectrum(self.K)/self.spatial_grid.pixsize**3)
    field_spatial = fft.irfftn(field_fourier, s=self.spatial_grid.shape)

    return field_spatial
__init__(data)

Constructor for FluctuationCube using a Cluster object

Source code in src/xsb_fluc/simulation/cube.py 
50
51
52
53
54
55
56
57
58
59
def __init__(self, data: Cluster):
    """
    Constructor for FluctuationCube using a Cluster object
    """

    super(FluctuationCube, self).__init__()
    self.power_spectrum = KolmogorovPowerSpectrum()
    self.spatial_grid = SpatialGrid3D.from_data(data, crop_r_500=5)
    self.fourier_grid = FourierGrid3D(self.spatial_grid)
    self.K = self.fourier_grid.K.astype(np.float32)

xsb_fluc.simulation.mexican_hat

MexicanHat

Bases: Module

Mexican Hat filter for the power spectrum, using the implementation of Arévalo et al. 2012.

Source code in src/xsb_fluc/simulation/mexican_hat.py 
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
class MexicanHat(hk.Module):
    """
    Mexican Hat filter for the power spectrum, using the implementation of Arévalo et al. 2012.
    """

    def __init__(self, data, mask=None):
        """
        Constructor for the MexicanHat class.

        Parameters:
            data (Cluster): Cluster object containing the data.
            mask (np.ndarray): Mask to apply to the data.

        !!! note
            The `PowerSpectrum`class is a `Cluster` oriented wrapper around this class.
        """
        super(MexicanHat, self).__init__()
        self.mexican_hat = MexicanHatMask.from_data(data, mask=mask)
        self.pixsize = self.mexican_hat.pixsize
        self.epsilon = self.mexican_hat.epsilon
        self.ratio = 1./np.sum(mask)
        self.Y = jnp.pi 

    def __call__(self, img, scale):
        """
        Computes the Mexican Hat filter for a given scale.

        Parameters:
            img (np.ndarray): Image to filter.
            scale (float): Scale of the filter.
        """

        img_convolved = self.mexican_hat(img, scale)
        k = 1/scale

        return jnp.sum(img_convolved**2)/(self.epsilon**2*self.Y*k**2)*self.ratio#/self.pixsize**2

__call__(img, scale)

Computes the Mexican Hat filter for a given scale.

Parameters:

Name Type Description Default
img ndarray

Image to filter.

 required
scale float

Scale of the filter.

 required
Source code in src/xsb_fluc/simulation/mexican_hat.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def __call__(self, img, scale):
    """
    Computes the Mexican Hat filter for a given scale.

    Parameters:
        img (np.ndarray): Image to filter.
        scale (float): Scale of the filter.
    """

    img_convolved = self.mexican_hat(img, scale)
    k = 1/scale

    return jnp.sum(img_convolved**2)/(self.epsilon**2*self.Y*k**2)*self.ratio#/self.pixsize**2

__init__(data, mask=None)

Constructor for the MexicanHat class.

Parameters:

Name Type Description Default
data Cluster

Cluster object containing the data.

 required
mask ndarray

Mask to apply to the data.

 None

Note

The PowerSpectrumclass is a Cluster oriented wrapper around this class.

Source code in src/xsb_fluc/simulation/mexican_hat.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def __init__(self, data, mask=None):
    """
    Constructor for the MexicanHat class.

    Parameters:
        data (Cluster): Cluster object containing the data.
        mask (np.ndarray): Mask to apply to the data.

    !!! note
        The `PowerSpectrum`class is a `Cluster` oriented wrapper around this class.
    """
    super(MexicanHat, self).__init__()
    self.mexican_hat = MexicanHatMask.from_data(data, mask=mask)
    self.pixsize = self.mexican_hat.pixsize
    self.epsilon = self.mexican_hat.epsilon
    self.ratio = 1./np.sum(mask)
    self.Y = jnp.pi

MexicanHatMask

Bases: Module

Auxiliary class for the MexicanHat class. Handle the convolution in Fourier space, the padding, etc.

Source code in src/xsb_fluc/simulation/mexican_hat.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
class MexicanHatMask(hk.Module):
    """
    Auxiliary class for the MexicanHat class. Handle the convolution in Fourier space,
    the padding, etc.
    """

    def __init__(self, shape, mask=None, pixsize=0.005):
        super(MexicanHatMask, self).__init__()

        self.pixsize = pixsize 
        self.epsilon = 1e-3 
        self.pad_size = 50
        self.mask = jnp.pad(mask, self.pad_size)

        kx = fft.fftfreq(shape[0] + 2*self.pad_size, d=self.pixsize)
        ky = fft.rfftfreq(shape[1]+ 2*self.pad_size, d=self.pixsize)
        KX, KY = jnp.meshgrid(kx, ky, indexing='ij')
        self.K = (KX**2 + KY**2)**(1/2)

    @classmethod
    def from_data(cls, data: Cluster, mask=None):

        if mask is None:
            mask = data.exp > 0.

        return cls(mask.shape, mask=mask, pixsize=(data.kpc_per_pixel/data.r_500).to(1/u.pixel).value)

    def fourier_kernel(self, sigma):

        return jnp.exp(-2*(jnp.pi*self.K*sigma)**2)

    def gaussian_blur(self, to_blur, sigma):

        blurred_img = fft.irfft2(fft.rfft2(to_blur)*self.fourier_kernel(sigma), s=to_blur.shape)
        return blurred_img#[self.pad_size:-self.pad_size, self.pad_size:-self.pad_size]

    def filter(self, img, sigma):

        img_masked = jnp.pad(img, self.pad_size) * self.mask
        gsigma1 = self.gaussian_blur(img_masked, sigma*(1+self.epsilon)**(-1/2))
        gsigma2 = self.gaussian_blur(img_masked, sigma*(1+self.epsilon)**(+1/2))
        # FFT-convolve mask with the two scales
        gmask1 = self.gaussian_blur(self.mask, sigma*(1+self.epsilon)**(-1/2))
        gmask2 = self.gaussian_blur(self.mask, sigma*(1+self.epsilon)**(+1/2))
        # Eq. 6 of Arevalo et al. 2012

        r1 = jnp.where(gmask1 != 0., gsigma1/gmask1, 0.)
        r2 = jnp.where(gmask2 != 0., gsigma2/gmask2, 0.)

        return (r1 - r2)*self.mask

    def __call__(self, img, scale):

        k = 1/scale

        #Arévalo & al A5
        sigma = 1 / jnp.sqrt(2 * jnp.pi ** 2) / k
        convolved_img = self.filter(img, sigma)

        return convolved_img

xsb_fluc.simulation.power_spectrum

PowerSpectrum

Bases: Module

Compute the power spectrum using the Mexican Hat filter for Clusterdata.

Source code in src/xsb_fluc/simulation/power_spectrum.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class PowerSpectrum(hk.Module):
    """
    Compute the power spectrum using the Mexican Hat filter for `Cluster`data.
    """

    def __init__(self, data, mask=None):
        """
        Constructor for the PowerSpectrum class.

        Parameters:
            data (Cluster): Cluster object containing the data.
            mask (np.ndarray): Mask to apply to the data.
        """
        super(PowerSpectrum, self).__init__()
        self.mexican_hat = MexicanHat(data, mask=mask)

    def __call__(self, image, scales):
        """
        Computes the power spectrum for a given image and scales.

        Parameters:
            image (np.ndarray): Image to compute the power spectrum of.
            scales (np.ndarray): Scales to compute the power spectrum at.
        """

        return hk.vmap(lambda s: self.mexican_hat(image, s), split_rng=False)(scales)

__call__(image, scales)

Computes the power spectrum for a given image and scales.

Parameters:

Name Type Description Default
image ndarray

Image to compute the power spectrum of.

 required
scales ndarray

Scales to compute the power spectrum at.

 required
Source code in src/xsb_fluc/simulation/power_spectrum.py 
22
23
24
25
26
27
28
29
30
31
def __call__(self, image, scales):
    """
    Computes the power spectrum for a given image and scales.

    Parameters:
        image (np.ndarray): Image to compute the power spectrum of.
        scales (np.ndarray): Scales to compute the power spectrum at.
    """

    return hk.vmap(lambda s: self.mexican_hat(image, s), split_rng=False)(scales)

__init__(data, mask=None)

Constructor for the PowerSpectrum class.

Parameters:

Name Type Description Default
data Cluster

Cluster object containing the data.

 required
mask ndarray

Mask to apply to the data.

 None
Source code in src/xsb_fluc/simulation/power_spectrum.py 
11
12
13
14
15
16
17
18
19
20
def __init__(self, data, mask=None):
    """
    Constructor for the PowerSpectrum class.

    Parameters:
        data (Cluster): Cluster object containing the data.
        mask (np.ndarray): Mask to apply to the data.
    """
    super(PowerSpectrum, self).__init__()
    self.mexican_hat = MexicanHat(data, mask=mask)