Skip to content

fitting

xsb_fluc.fitting.mean_model

MeanModelFitter

This class is meant to fit a 3D model using the X-ray surface

Source code in src/xsb_fluc/fitting/mean_model.py 
 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
 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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
class MeanModelFitter:
    """
    This class is meant to fit a 3D model using the X-ray surface
    """

    def __init__(self,
                 cluster: Cluster,
                 n_samples: int = 1000,
                 n_warmup: int = 1000,
                 n_chains: int = 1,
                 model: Literal[None, 'beta', 'circ'] = None,
                 max_tree_depth: int = 10,
                 ref_params=None):

        """
        Constructor for the MeanModelFitter which fit the mean model using MCMC

        Parameters:
            cluster (Cluster): the cluster object which will be fitted
            n_samples (int): the number of samples
            n_warmup (int): the number of warmup samples
            n_chains (int): the number of chains
            model (Literal[None, 'beta', 'circ']): the model to be fitted. Leave None to use the best.
            max_tree_depth (int): the maximum recursion depth of the MCMC. Lower to 7 or 5 if too slow.
            ref_params (dict): the parameters where the chain should be started.
        """

        self.cluster = cluster
        self.model_var = model
        self.ref_params = ref_params
        self.max_tree_depth = max_tree_depth

        if self.model_var == 'beta' or self.model_var == 'circ':

            self.counts = hk.without_apply_rng(hk.transform(lambda: MockXrayCountsBetaModel(self.cluster)()))

        else:

            self.counts = hk.without_apply_rng(hk.transform(lambda: MockXrayCounts(self.cluster)()))

        self.mcmc_config = {'num_samples': n_samples,
                            'num_warmup': n_warmup,
                            'num_chains': n_chains,
                            'progress_bar': True}

    @property
    def prior(self):
        """
        Default priors distributions.
        """

        sb = self.cluster.img/self.cluster.exp
        X = np.stack([self.cluster.x_ref, self.cluster.y_ref])
        np.average(X, weights=sb, axis=1)
        X_cov = np.cov(X, aweights=sb)

        eig_val, eig_vec = np.linalg.eig(X_cov)
        angle = np.arccos(np.dot([1., 0.], eig_vec[:, np.argmin(eig_val)]))
        (1-(min(eig_val)/max(eig_val))**2)**(1/2)

        if self.model_var is None:

            prior = {'angle': numpyro.sample('angle', dist.Uniform(-jnp.pi/2, jnp.pi/2)),
                     'eccentricity': numpyro.sample('eccentricity', dist.Uniform(0., 0.99)),
                     'x_c': numpyro.sample('x_c', dist.Normal(0., 0.5)),
                     'y_c': numpyro.sample('y_c', dist.Normal(0., 0.5)),
                     'log_bkg': numpyro.sample('log_bkg', dist.Uniform(-10., -4.)),
                     'log_ne2': numpyro.sample('log_ne2', dist.Uniform(-8., -3.)),
                     'log_r_c': numpyro.sample('log_r_c', dist.Uniform(-2., 0.)),
                     'log_r_s': numpyro.sample('log_r_s', dist.Uniform(-1., 1.)),
                     'beta': numpyro.sample('beta', dist.Uniform(0., 5.)),
                     'epsilon': numpyro.sample('epsilon', dist.Uniform(0., 5.))}

        if self.model_var == 'beta':

            prior = {'angle': numpyro.sample('angle', dist.TruncatedNormal(loc=angle,
                                                                           scale=0.2,
                                                                           low=-jnp.pi/2+angle,
                                                                           high=jnp.pi/2+angle)),

                     'eccentricity': numpyro.sample('eccentricity', dist.TruncatedNormal(
                                                                           loc=angle,
                                                                           scale=0.2,
                                                                           low=0,
                                                                           high=0.9)),
                     'x_c': numpyro.sample('x_c', dist.Normal(0., 0.2)),
                     'y_c': numpyro.sample('y_c', dist.Normal(0., 0.2)),
                     'log_bkg': numpyro.sample('log_bkg', dist.Uniform(-10., -4.)),
                     'log_e_0': numpyro.sample('log_e_0', dist.Uniform(-7., 0.)),
                     'log_r_c': numpyro.sample('log_r_c', dist.Uniform(-4., 1.)),
                     'beta': numpyro.sample('beta', dist.Uniform(0., 5.))}

        if self.model_var == 'circ':

            prior = {'angle': jnp.array(0.),
                     'eccentricity': jnp.array(0.),
                     'x_c': numpyro.sample('x_c', dist.Normal(0., 0.2)),
                     'y_c': numpyro.sample('y_c', dist.Normal(0., 0.2)),
                     'log_bkg': numpyro.sample('log_bkg', dist.Uniform(-10., -4.)),
                     'log_e_0': numpyro.sample('log_e_0', dist.Uniform(-7., 0.)),
                     'log_r_c': numpyro.sample('log_r_c', dist.Uniform(-4., 1.)),
                     'beta': numpyro.sample('beta', dist.Uniform(0., 5.))}

        return prior

    def model(self):
        """
        The numpyro model which is used in the fitting routine.
        """

        params = set_params(self.counts.init(None), self.prior)

        numpyro.sample('likelihood',
                       dist.Poisson(self.counts.apply(params)),
                       obs=jnp.array(self.cluster.img, dtype=jnp.float32))

    def fit(self) -> az.InferenceData:
        """
        Perform the fitting routine.

        Returns:
            An az.InferenceData object containing the fit results.
        """

        if self.ref_params is not None:

            kernel = BarkerMH(self.model,
                          #max_tree_depth=self.max_tree_depth,
                          init_strategy=init_to_value(values=self.ref_params),
                        )
                          #dense_mass=True,
                          #target_accept_prob=0.95)

        else :

            kernel = BarkerMH(self.model)
              #max_tree_depth=self.max_tree_depth,
              #dense_mass=True,
              #target_accept_prob=0.95)

        posterior = MCMC(kernel, **self.mcmc_config)

        key = random.split(rng_key(), 4)
        posterior.run(key[0])

        posterior_samples = posterior.get_samples()
        posterior_predictive = Predictive(self.model, posterior_samples)(key[1])
        prior = Predictive(self.model, num_samples=100000)(key[2])

        inference_data = az.from_numpyro(posterior, prior=prior, posterior_predictive=posterior_predictive)

        return inference_data
prior property

Default priors distributions.

__init__(cluster, n_samples=1000, n_warmup=1000, n_chains=1, model=None, max_tree_depth=10, ref_params=None)

Constructor for the MeanModelFitter which fit the mean model using MCMC

Parameters:

Name Type Description Default
cluster Cluster

the cluster object which will be fitted

 required
n_samples int

the number of samples

 1000
n_warmup int

the number of warmup samples

 1000
n_chains int

the number of chains

 1
model Literal[None, 'beta', 'circ']

the model to be fitted. Leave None to use the best.

 None
max_tree_depth int

the maximum recursion depth of the MCMC. Lower to 7 or 5 if too slow.

 10
ref_params dict

the parameters where the chain should be started.

 None
Source code in src/xsb_fluc/fitting/mean_model.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
def __init__(self,
             cluster: Cluster,
             n_samples: int = 1000,
             n_warmup: int = 1000,
             n_chains: int = 1,
             model: Literal[None, 'beta', 'circ'] = None,
             max_tree_depth: int = 10,
             ref_params=None):

    """
    Constructor for the MeanModelFitter which fit the mean model using MCMC

    Parameters:
        cluster (Cluster): the cluster object which will be fitted
        n_samples (int): the number of samples
        n_warmup (int): the number of warmup samples
        n_chains (int): the number of chains
        model (Literal[None, 'beta', 'circ']): the model to be fitted. Leave None to use the best.
        max_tree_depth (int): the maximum recursion depth of the MCMC. Lower to 7 or 5 if too slow.
        ref_params (dict): the parameters where the chain should be started.
    """

    self.cluster = cluster
    self.model_var = model
    self.ref_params = ref_params
    self.max_tree_depth = max_tree_depth

    if self.model_var == 'beta' or self.model_var == 'circ':

        self.counts = hk.without_apply_rng(hk.transform(lambda: MockXrayCountsBetaModel(self.cluster)()))

    else:

        self.counts = hk.without_apply_rng(hk.transform(lambda: MockXrayCounts(self.cluster)()))

    self.mcmc_config = {'num_samples': n_samples,
                        'num_warmup': n_warmup,
                        'num_chains': n_chains,
                        'progress_bar': True}
fit()

Perform the fitting routine.

Returns:

Type Description
InferenceData

An az.InferenceData object containing the fit results.
Source code in src/xsb_fluc/fitting/mean_model.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def fit(self) -> az.InferenceData:
    """
    Perform the fitting routine.

    Returns:
        An az.InferenceData object containing the fit results.
    """

    if self.ref_params is not None:

        kernel = BarkerMH(self.model,
                      #max_tree_depth=self.max_tree_depth,
                      init_strategy=init_to_value(values=self.ref_params),
                    )
                      #dense_mass=True,
                      #target_accept_prob=0.95)

    else :

        kernel = BarkerMH(self.model)
          #max_tree_depth=self.max_tree_depth,
          #dense_mass=True,
          #target_accept_prob=0.95)

    posterior = MCMC(kernel, **self.mcmc_config)

    key = random.split(rng_key(), 4)
    posterior.run(key[0])

    posterior_samples = posterior.get_samples()
    posterior_predictive = Predictive(self.model, posterior_samples)(key[1])
    prior = Predictive(self.model, num_samples=100000)(key[2])

    inference_data = az.from_numpyro(posterior, prior=prior, posterior_predictive=posterior_predictive)

    return inference_data
model()

The numpyro model which is used in the fitting routine.

Source code in src/xsb_fluc/fitting/mean_model.py 
120
121
122
123
124
125
126
127
128
129
def model(self):
    """
    The numpyro model which is used in the fitting routine.
    """

    params = set_params(self.counts.init(None), self.prior)

    numpyro.sample('likelihood',
                   dist.Poisson(self.counts.apply(params)),
                   obs=jnp.array(self.cluster.img, dtype=jnp.float32))