<jupyter_start><jupyter_text>Fitting Mass Functions with `emcee` In this demo, we provide a very simple example of how to fit mass functions with `emcee`. This example does not do anything with the complexity required for fitting real-world data. It does not account for measurement error or Eddington bias. However, it gives a taste of what is achievable.<jupyter_code>import emcee
import hmf
import numpy as np

import matplotlib.pyplot as plt
from scipy.stats import poisson
from multiprocessing import Pool



from IPython.display import display, Markdown

%matplotlib inline
emcee.__version__
hmf.__version__<jupyter_output><empty_output><jupyter_text>Create Some Mock Data First, let's create a default model mass function:<jupyter_code>model = hmf.MassFunction(z=1.5, transfer_model='EH')<jupyter_output><empty_output><jupyter_text>Now, let's create some mock data, based on the true mass function, but with some randomness assuming the mass function was measured (perfectly) in some volume:<jupyter_code>volume = 100 ** 3 # Mpc^3 / h^3
counts_per_bin = np.random.poisson(model.dndlog10m * volume * model.dlog10m)
plt.plot(model.m, model.dndlog10m * volume * model.dlog10m)
plt.plot(model.m, counts_per_bin)
plt.xscale('log')
plt.yscale('log')<jupyter_output><empty_output><jupyter_text>Define a likelihood Now let's define a likelihood function, based on some input model for `dndlog10m`. The model is simply a Poisson model. We take the model into the space in which it describes an actual number of observations, and get the log pdf there:<jupyter_code>def log_likelihood(model, data, volume, dlog10m):
    return np.sum(poisson.logpmf(data, mu=model*volume*dlog10m))<jupyter_output><empty_output><jupyter_text>Define an `emcee`-compatible likelihood function Now we define a likelihood function for `emcee`. There is a bit more flexibility here, as this function needs to calculate priors on all input parameters, and handle exceptions as well. This means this function is rather specific to the problem at hand. We will define a very simple function, but one that is fairly general.<jupyter_code>fiducial_model = model.clone()<jupyter_output><empty_output><jupyter_text>First, we define a small utility function that will take a dictionary in which keys may be dot-paths, and converts it to a nested dictionary:<jupyter_code>def flat_to_nested_dict(dct: dict) -> dict:
    """Convert a dct of key: value pairs into a nested dict.
    
    Keys that have dots in them indicate nested structure.
    """    
    def key_to_dct(key, val, dct):
        if '.' in key:
            key, parts = key.split('.', maxsplit=1)
            
            if key not in dct:
                dct[key] = {}
                
            key_to_dct(parts, val, dct[key])
        else:
            dct[key] = val
            
    out = {} 
    for k, v in dct.items():
        key_to_dct(k, v, out)

    return out<jupyter_output><empty_output><jupyter_text>So, this will do the following:<jupyter_code>flat_to_nested_dict(
    {
        'nested.key': 1,
        'nested.key2': 2,
        'non_nested': 3
        
    }
)<jupyter_output><empty_output><jupyter_text>This will enable us to pass a list of parameter names that we want updated, which could be parameters of nested models. This means our posterior function is fairly general, and can accept any model parameters to be updated:<jupyter_code>def log_prob(param_values, param_names, data, model, volume, derived=None):
    # Update the base model with all the parameters that are being constrained.
    params = flat_to_nested_dict(dict(zip(param_names, param_values)))

    model.update(**params)
    
    if derived is not None:
        derived = [getattr(model, d) for d in derived]
    else:
        derived = []
        
    return log_likelihood(model.dndlog10m, data, volume, model.dlog10m), derived<jupyter_output><empty_output><jupyter_text>We can test that the `log_prob` function works:<jupyter_code>log_prob([0.9], ['sigma_8'], counts_per_bin, model, volume, derived=['mass_nonlinear'])<jupyter_output><empty_output><jupyter_text>Notice the derived parameters: we can pass any `quantity` of the `MassFunction`, and it will be stored on every iteration. Nice!  Run `emcee` Let's first run a simple model in which we just want to fit $\sigma_8$ and the spectral index $n_s$. We use the popular `emcee` package, and pass in our `log_prob` model:<jupyter_code>sampler = emcee.EnsembleSampler(
    nwalkers = 100,
    ndim = 2,
    log_prob_fn = log_prob,
    kwargs = {
        'param_names': ['sigma_8', 'n'], 
        'data': counts_per_bin, 
        'model': model, 
        'volume': volume, 
        'derived': ['mass_nonlinear']
    },
    pool = Pool(4),
)<jupyter_output>/home/steven/Documents/Projects/halos/HMF/hmf/src/hmf/mass_function/hmf.py:377: UserWarning: Nonlinear mass outside mass range
  warnings.warn("Nonlinear mass outside mass range")
/home/steven/Documents/Projects/halos/HMF/hmf/src/hmf/mass_function/hmf.py:377: UserWarning: Nonlinear mass outside mass range
  warnings.warn("Nonlinear mass outside mass range")
/home/steven/Documents/Projects/halos/HMF/hmf/src/hmf/mass_function/hmf.py:377: UserWarning: Nonlinear mass outside mass range
  warnings.warn("Nonlinear mass outside mass range")<jupyter_text>On the advice of the `emcee` documentation, we set up some initial positions of the walkers around the solution.<jupyter_code>initialpos = np.array([fiducial_model.sigma_8, fiducial_model.n]) + 1e-4 * np.random.normal(size=(sampler.nwalkers, sampler.ndim))
sampler.run_mcmc(initialpos, nsteps=1000)<jupyter_output><empty_output><jupyter_text>Now we can plot the posterior as a heatmap, with the true input values shown as white lines:<jupyter_code>plt.hexbin(sampler.chain[:,500:, 0].flatten(), sampler.chain[:,500:, 1].flatten())
plt.axvline(fiducial_model.sigma_8, color='white')
plt.axhline(fiducial_model.n, color='white')
plt.colorbar()<jupyter_output><empty_output><jupyter_text>And we're done! The posterior contains the truth to within 1-sigma, and we have ~1% error-bars on $\sigma_8$ and ~0.2% error-bars on $n$.  An example of fitting a new fitting function Perhaps you are a simulator and want to deliver a new fitting function, or at least new parameters for a particular fitting function. Let's look at an example of this. We choose the Sheth-Tormen function:<jupyter_code>fiducial_model = hmf.MassFunction(z=0, hmf_model='SMT', dlog10m=0.2, Mmin=10, Mmax=16)<jupyter_output><empty_output><jupyter_text>The functional form of $f(\nu)$ is:<jupyter_code>display(Markdown(f"${fiducial_model.hmf._eq}$"))<jupyter_output><empty_output><jupyter_text>The fitted parameters are:<jupyter_code>fiducial_model.hmf.params<jupyter_output><empty_output><jupyter_text>Again, we use a Poisson model to create mock data, this time in a Gpc volume:<jupyter_code>volume = 500 ** 3 # Mpc^3 / h^3

counts_per_bin = np.random.poisson(fiducial_model.dndlog10m * volume * fiducial_model.dlog10m)

plt.plot(fiducial_model.m, fiducial_model.dndlog10m * volume * fiducial_model.dlog10m)
plt.plot(fiducial_model.m, counts_per_bin)
plt.xscale('log')
plt.yscale('log')<jupyter_output><empty_output><jupyter_text>We've already set up our probability function, which was already general enough to be able to fit the hmf parameters, so off we go!<jupyter_code>smt_model = fiducial_model.clone()
sampler_hmf = emcee.EnsembleSampler(
    nwalkers = 50,
    ndim = 3,
    log_prob_fn = log_prob,
    kwargs = {
        'param_names': ['hmf_params.a', 'hmf_params.p', 'hmf_params.A'], 
        'data': counts_per_bin, 
        'model': smt_model, 
        'volume': volume, 
        'derived': ['dndlog10m']
    },
    pool = Pool(4),
)
init_pos = (
    np.array([
        smt_model.hmf.params['a'], 
        smt_model.hmf.params['p'], 
        smt_model.hmf.params['A']]) + 
    1e-4 * np.random.normal(size=(sampler_hmf.nwalkers, sampler_hmf.ndim))
)
sampler_hmf.run_mcmc(init_pos, nsteps=300, progress='notebook')<jupyter_output><empty_output><jupyter_text>Again, we can plot the posterior triangle plot:<jupyter_code>fig, ax = plt.subplots(2, 2, sharex='col', sharey='row', figsize=(10, 10), gridspec_kw={"hspace":0, 'wspace': 0})

flatchain = sampler_hmf.chain[:, 100:, :].reshape((-1, 3))

ax[0, 0].hexbin(flatchain[:, 0], flatchain[:, 1])
ax[1, 0].hexbin(flatchain[:, 0], flatchain[:, 2])
ax[1, 1].hexbin(flatchain[:, 1], flatchain[:, 2])

ax[1, 0].set_xlabel("a", fontsize=15)
ax[1, 0].set_ylabel("A", fontsize=15)
ax[1, 1].set_xlabel("p", fontsize=15)
ax[0, 0].set_ylabel("p", fontsize=15)<jupyter_output><empty_output><jupyter_text>Since we collected $dn/d\log_{10}m$ as a derived parameter, we can also plot it:<jupyter_code>dndlog10m_samples = sampler_hmf.blobs[100:, :, :].reshape((-1, 40))<jupyter_output><empty_output><jupyter_text>We'll make a plot of the posterior distribution of the mass function, as a fraction of the median of the distribution:<jupyter_code>percentiles = np.percentile(dndlog10m_samples, q=(16, 50, 84), axis=0)
plt.fill_between(fiducial_model.m, percentiles[0]/percentiles[1], percentiles[2]/percentiles[1])
plt.xscale('log')
plt.ylabel("Fractional difference to Median")
plt.xlabel("Mass")
#plt.yscale('log')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11496863508544235, 'fraction_numerical': 0.01827817434566299, 'mean_word_length': 4.878575969485061, 'pattern_counts': {'":': 1, '<': 80, '<?xml version=': 0, '>': 81, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '20', 'dir': 'jupyter-structured-clean-dedup', 'id': '230095', 'license': 'MIT', 'n_tokens_mistral': 3314, 'n_tokens_neox': 3147, 'n_words': 860, 'path': 'docs/examples/fitting.ipynb', 'repo_name': 'IvanNikolic21/hmf', 'size': '9246'}