Skip to content

Projection

projection

VCubeProjection_v2

Projection of velocity cube with emission weighting in binned map

Source code in src/projection.py 
 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
 69
 70
 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
class VCubeProjection_v2 :
    """
    Projection of velocity cube with emission weighting in binned map
    """
    def __init__(self, binning, em_cube, PSF_kernel):
        """
        Initialize

        Parameters:
            binning (jnp.Module): Binning
            em_cube (jnp.array): 3D emissivity cube
            PSF_kernel (jnp.array): PSF discretized on the pixel grid
        """

        self.binning = binning
        self.em = em_cube
        self.PSF_kernel = PSF_kernel

    def __call__(self,v):

        """
        Uses counts to weight and project the velocity cube
        into a centroid shift map and broadening map.
        Uses a convolution with product in Fourier space.

        Parameters:
            v (jnp.array): Velocity cube

        Returns:
            v_map (jnp.array): Binned image of the emission weighted centroid shift
            std_map (jnp.array): Binned image of the emission weighted broadening
            binned_v_weighted (jnp.array): Vector of emission weighted centroid shift in each bin
            binned_std_weighted (jnp.array): Vector of emission weighted broadening in each bin

        """

        # This is a fully jaxified version

        # Summing emission weighted speed in each pixel        
        v_vec = jnp.sum(v*self.em, axis = -1)
        std_vec = jnp.sum(v**2*self.em, axis = -1)
        count_vec = jnp.sum(self.em, axis = -1)

        # Convolution by PSF
        v_vec_conv = my_convolution(v_vec, self.PSF_kernel)
        std_vec_conv = my_convolution(std_vec, self.PSF_kernel)
        count_vec_conv = my_convolution(count_vec, self.PSF_kernel)

        # Extracting values within pixels from matrix
        v_vec_conv_pix = v_vec_conv[self.binning.X_pixels,self.binning.Y_pixels]
        std_vec_conv_pix = std_vec_conv[self.binning.X_pixels,self.binning.Y_pixels]
        count_vec_conv_pix = count_vec_conv[self.binning.X_pixels,self.binning.Y_pixels]


        # Indices of which bin each pixel goes in
        bins_unique, inverse_indices = jnp.unique(self.binning.bin_num_pix, 
                                                  return_inverse=True, 
                                                  size = self.binning.nb_bins)

        # Binned vectors
        bin_v = jnp.zeros(len(bins_unique))
        bin_counts = jnp.zeros(len(bins_unique))
        bin_std = jnp.zeros(len(bins_unique))

        # We add to each bin the weighted sum of all pixels belonging to it
        bin_v = bin_v.at[inverse_indices].add(v_vec_conv_pix)
        bin_counts = bin_counts.at[inverse_indices].add(count_vec_conv_pix)
        bin_std = bin_std.at[inverse_indices].add(std_vec_conv_pix)


        # Divide by weighing (i.e. summed emission)
        binned_v_weighted = bin_v / bin_counts
        binned_std_weighted = jnp.sqrt((bin_std/bin_counts - binned_v_weighted **2))

        # Create maps
        v_map = jnp.zeros(self.binning.shape)
        v_map = v_map.at[self.binning.X_pixels, self.binning.Y_pixels].set(binned_v_weighted[inverse_indices])

        std_map = jnp.zeros(self.binning.shape)
        std_map = std_map.at[self.binning.X_pixels, self.binning.Y_pixels].set(binned_std_weighted[inverse_indices])

        # Return weighted and binned values
        return v_map, std_map, binned_v_weighted, binned_std_weighted

__call__(v)

Uses counts to weight and project the velocity cube into a centroid shift map and broadening map. Uses a convolution with product in Fourier space.

Parameters:

Name Type Description Default
v array

Velocity cube

 required

Returns:

Name Type Description
v_map array

Binned image of the emission weighted centroid shift
std_map array

Binned image of the emission weighted broadening
binned_v_weighted array

Vector of emission weighted centroid shift in each bin
binned_std_weighted array

Vector of emission weighted broadening in each bin
Source code in src/projection.py 
 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
 69
 70
 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
def __call__(self,v):

    """
    Uses counts to weight and project the velocity cube
    into a centroid shift map and broadening map.
    Uses a convolution with product in Fourier space.

    Parameters:
        v (jnp.array): Velocity cube

    Returns:
        v_map (jnp.array): Binned image of the emission weighted centroid shift
        std_map (jnp.array): Binned image of the emission weighted broadening
        binned_v_weighted (jnp.array): Vector of emission weighted centroid shift in each bin
        binned_std_weighted (jnp.array): Vector of emission weighted broadening in each bin

    """

    # This is a fully jaxified version

    # Summing emission weighted speed in each pixel        
    v_vec = jnp.sum(v*self.em, axis = -1)
    std_vec = jnp.sum(v**2*self.em, axis = -1)
    count_vec = jnp.sum(self.em, axis = -1)

    # Convolution by PSF
    v_vec_conv = my_convolution(v_vec, self.PSF_kernel)
    std_vec_conv = my_convolution(std_vec, self.PSF_kernel)
    count_vec_conv = my_convolution(count_vec, self.PSF_kernel)

    # Extracting values within pixels from matrix
    v_vec_conv_pix = v_vec_conv[self.binning.X_pixels,self.binning.Y_pixels]
    std_vec_conv_pix = std_vec_conv[self.binning.X_pixels,self.binning.Y_pixels]
    count_vec_conv_pix = count_vec_conv[self.binning.X_pixels,self.binning.Y_pixels]


    # Indices of which bin each pixel goes in
    bins_unique, inverse_indices = jnp.unique(self.binning.bin_num_pix, 
                                              return_inverse=True, 
                                              size = self.binning.nb_bins)

    # Binned vectors
    bin_v = jnp.zeros(len(bins_unique))
    bin_counts = jnp.zeros(len(bins_unique))
    bin_std = jnp.zeros(len(bins_unique))

    # We add to each bin the weighted sum of all pixels belonging to it
    bin_v = bin_v.at[inverse_indices].add(v_vec_conv_pix)
    bin_counts = bin_counts.at[inverse_indices].add(count_vec_conv_pix)
    bin_std = bin_std.at[inverse_indices].add(std_vec_conv_pix)


    # Divide by weighing (i.e. summed emission)
    binned_v_weighted = bin_v / bin_counts
    binned_std_weighted = jnp.sqrt((bin_std/bin_counts - binned_v_weighted **2))

    # Create maps
    v_map = jnp.zeros(self.binning.shape)
    v_map = v_map.at[self.binning.X_pixels, self.binning.Y_pixels].set(binned_v_weighted[inverse_indices])

    std_map = jnp.zeros(self.binning.shape)
    std_map = std_map.at[self.binning.X_pixels, self.binning.Y_pixels].set(binned_std_weighted[inverse_indices])

    # Return weighted and binned values
    return v_map, std_map, binned_v_weighted, binned_std_weighted

__init__(binning, em_cube, PSF_kernel)

Initialize

Parameters:

Name Type Description Default
binning Module

Binning

 required
em_cube array

3D emissivity cube

 required
PSF_kernel array

PSF discretized on the pixel grid

 required
Source code in src/projection.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
def __init__(self, binning, em_cube, PSF_kernel):
    """
    Initialize

    Parameters:
        binning (jnp.Module): Binning
        em_cube (jnp.array): 3D emissivity cube
        PSF_kernel (jnp.array): PSF discretized on the pixel grid
    """

    self.binning = binning
    self.em = em_cube
    self.PSF_kernel = PSF_kernel