Skip to content

Simulation

simulation

Simulation

Bases: Module

Source code in src/simulation.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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
class Simulation(hk.Module):

    def __init__(self, 
                spatial_grid,
                structure_function,
                binning,
                projection,
                radial_bins_mes_errors,
                censhift_offsets,
                censhift_errors,
                broad_offsets,
                broad_errors,
                ):
        """
        Initialize simulation model

        Parameters:
            spatial_grid (hk.Module): 3D spatial grid
            structure_function (hk.Module): Structure Function
            binning (hk.Module): Binning
            projection (hk.Module) : Projection
            radial_bins_mes_errors (jnp.array): Bounds of the radial bins defininf the measurement error
            censhift_offsets (jnp.array): Means of censhift measurement error
            censhift_errors (jnp.array): stds of censhift measurement error
            broad_offsets (jnp.array): Means of broadening measurement error
            broad_errors (jnp.array): std of censhift measurement error
        """

        super(Simulation, self).__init__()


        # Initialize 
        self.fluctuation_generator = FluctuationCube(spatial_grid)
        #self.pars = self.fluctuation_generator.init(rng_key())

        # Binning
        self.binning = binning

        # Structure Function function
        self.StructureFunction = structure_function

        # Projection
        self.projection = projection

        # Half size of the grid
        half_grid_size = spatial_grid.shape[0]/2

        # Distances of bins to center of grid
        rBar_bins = jnp.sqrt((binning.xBar_bins - half_grid_size)**2 + (binning.yBar_bins - half_grid_size)**2)

        # Spread in centroid shift measurement error in each bin
        self.offsets_v = censhift_offsets[jnp.searchsorted(radial_bins_mes_errors, rBar_bins)-1]
        self.errors_v = censhift_errors[jnp.searchsorted(radial_bins_mes_errors, rBar_bins)-1]

        # Spread in broadening measurement error in each bin
        self.offsets_std = broad_offsets[jnp.searchsorted(radial_bins_mes_errors, rBar_bins)-1]
        self.errors_std = broad_errors[jnp.searchsorted(radial_bins_mes_errors, rBar_bins)-1]

    def __call__(self):
        """
        Creates a realization of a GRF for the speed along the los
        Projects the speed with em-weighting in binned maps
        Adds measurement error
        Returns the structure function

        Returns:
            dist (jn.array): Vector of separations at which the SF is computed
            sf (jnp.array): SF of centroid shift
            sf_std (jnp.array): SF of broadening
            v_vec (jnp.array): Vector of centroid shifts
            std_vec (jnp.array): Vector of broadenings
        """

        key = hk.next_rng_key()

        v_cube = self.fluctuation_generator()

        _,_,v_vec, std_vec = self.projection(v_cube)

        #Add measurement error on centroid shift
        err_v = random.multivariate_normal(key = key,
                                            mean = self.offsets_v, 
                                            cov = jnp.diag(self.errors_v**2))

        v_vec += jnp.where(jnp.invert(jnp.isnan(err_v)), err_v, 0)
        #v_vec = v_vec.at[jnp.invert(jnp.isnan(err_v))].add(err_v)


        # Add measurement error on broadening 
        # with a Gamma distribution, so that for high broadening values, 
        # the errors are distributed normally around the expected vector
        # and for low broadening values, the errors are always positive

        mu = std_vec + self.offsets_std
        a = mu**2 / self.errors_std**2
        scale = self.errors_std**2 / mu
        std_vec = random.gamma(key = rng_key(),a = a) * scale


        #SF of velocity map
        dist, sf = self.StructureFunction(v_vec , 
                                          self.binning.xBar_bins, 
                                          self.binning.yBar_bins)
        dist, sf_std = self.StructureFunction(std_vec, 
                                              self.binning.xBar_bins, 
                                              self.binning.yBar_bins)

        return dist, sf, sf_std, v_vec, std_vec

__call__()

Creates a realization of a GRF for the speed along the los Projects the speed with em-weighting in binned maps Adds measurement error Returns the structure function

Returns:

Name Type Description
dist array

Vector of separations at which the SF is computed
sf array

SF of centroid shift
sf_std array

SF of broadening
v_vec array

Vector of centroid shifts
std_vec array

Vector of broadenings
Source code in src/simulation.py 
 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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def __call__(self):
    """
    Creates a realization of a GRF for the speed along the los
    Projects the speed with em-weighting in binned maps
    Adds measurement error
    Returns the structure function

    Returns:
        dist (jn.array): Vector of separations at which the SF is computed
        sf (jnp.array): SF of centroid shift
        sf_std (jnp.array): SF of broadening
        v_vec (jnp.array): Vector of centroid shifts
        std_vec (jnp.array): Vector of broadenings
    """

    key = hk.next_rng_key()

    v_cube = self.fluctuation_generator()

    _,_,v_vec, std_vec = self.projection(v_cube)

    #Add measurement error on centroid shift
    err_v = random.multivariate_normal(key = key,
                                        mean = self.offsets_v, 
                                        cov = jnp.diag(self.errors_v**2))

    v_vec += jnp.where(jnp.invert(jnp.isnan(err_v)), err_v, 0)
    #v_vec = v_vec.at[jnp.invert(jnp.isnan(err_v))].add(err_v)


    # Add measurement error on broadening 
    # with a Gamma distribution, so that for high broadening values, 
    # the errors are distributed normally around the expected vector
    # and for low broadening values, the errors are always positive

    mu = std_vec + self.offsets_std
    a = mu**2 / self.errors_std**2
    scale = self.errors_std**2 / mu
    std_vec = random.gamma(key = rng_key(),a = a) * scale


    #SF of velocity map
    dist, sf = self.StructureFunction(v_vec , 
                                      self.binning.xBar_bins, 
                                      self.binning.yBar_bins)
    dist, sf_std = self.StructureFunction(std_vec, 
                                          self.binning.xBar_bins, 
                                          self.binning.yBar_bins)

    return dist, sf, sf_std, v_vec, std_vec

__init__(spatial_grid, structure_function, binning, projection, radial_bins_mes_errors, censhift_offsets, censhift_errors, broad_offsets, broad_errors)

Initialize simulation model

Parameters:

Name Type Description Default
spatial_grid Module

3D spatial grid

 required
structure_function Module

Structure Function

 required
binning Module

Binning

 required
projection hk.Module)

Projection

 required
radial_bins_mes_errors array

Bounds of the radial bins defininf the measurement error

 required
censhift_offsets array

Means of censhift measurement error

 required
censhift_errors array

stds of censhift measurement error

 required
broad_offsets array

Means of broadening measurement error

 required
broad_errors array

std of censhift measurement error

 required
Source code in src/simulation.py 
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
def __init__(self, 
            spatial_grid,
            structure_function,
            binning,
            projection,
            radial_bins_mes_errors,
            censhift_offsets,
            censhift_errors,
            broad_offsets,
            broad_errors,
            ):
    """
    Initialize simulation model

    Parameters:
        spatial_grid (hk.Module): 3D spatial grid
        structure_function (hk.Module): Structure Function
        binning (hk.Module): Binning
        projection (hk.Module) : Projection
        radial_bins_mes_errors (jnp.array): Bounds of the radial bins defininf the measurement error
        censhift_offsets (jnp.array): Means of censhift measurement error
        censhift_errors (jnp.array): stds of censhift measurement error
        broad_offsets (jnp.array): Means of broadening measurement error
        broad_errors (jnp.array): std of censhift measurement error
    """

    super(Simulation, self).__init__()


    # Initialize 
    self.fluctuation_generator = FluctuationCube(spatial_grid)
    #self.pars = self.fluctuation_generator.init(rng_key())

    # Binning
    self.binning = binning

    # Structure Function function
    self.StructureFunction = structure_function

    # Projection
    self.projection = projection

    # Half size of the grid
    half_grid_size = spatial_grid.shape[0]/2

    # Distances of bins to center of grid
    rBar_bins = jnp.sqrt((binning.xBar_bins - half_grid_size)**2 + (binning.yBar_bins - half_grid_size)**2)

    # Spread in centroid shift measurement error in each bin
    self.offsets_v = censhift_offsets[jnp.searchsorted(radial_bins_mes_errors, rBar_bins)-1]
    self.errors_v = censhift_errors[jnp.searchsorted(radial_bins_mes_errors, rBar_bins)-1]

    # Spread in broadening measurement error in each bin
    self.offsets_std = broad_offsets[jnp.searchsorted(radial_bins_mes_errors, rBar_bins)-1]
    self.errors_std = broad_errors[jnp.searchsorted(radial_bins_mes_errors, rBar_bins)-1]