<jupyter_start><jupyter_text>[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/jfcrenshaw/pzflow/blob/main/examples/marginalization.ipynb)If running in Colab, to switch to GPU, go to the menu and select Runtime -> Change runtime type -> Hardware accelerator -> GPU.In addition, uncomment and run the following code:<jupyter_code># !pip install pzflow<jupyter_output><empty_output><jupyter_text>------------------- Marginalization during posterior calculationThis example notebook demonstrates how to marginalize over missing variables during posterior calculation.We will use the Flow trained in the [redshift example](https://github.com/jfcrenshaw/pzflow/blob/main/examples/redshift_example.ipynb).<jupyter_code>import jax.numpy as np
import matplotlib.pyplot as plt
from pzflow.examples import get_example_flow<jupyter_output><empty_output><jupyter_text>First let's load the pre-trained flow, and use it to generate some samples:<jupyter_code>flow = get_example_flow()
samples = flow.sample(2, seed=123)
samples<jupyter_output><empty_output><jupyter_text>Remember that we can calculate posteriors for the data in samples. For example, let's plot redshift posteriors:<jupyter_code>grid = np.linspace(0.25, 1.45, 100)
pdfs = flow.posterior(samples, column="redshift", grid=grid)
fig, axes = plt.subplots(1, 2, figsize=(5.5, 2), dpi=120, constrained_layout=True)

for i, ax in enumerate(axes.flatten()):
    
    ax.plot(grid, pdfs[i], label="Redshift posterior")
    
    ztrue = samples["redshift"][i]
    ax.axvline(ztrue, c="C3", label="True redshift")
    ax.set(
        xlabel="redshift", 
        xlim=(ztrue - 0.25, ztrue + 0.25), 
        yticks=[]
    )
    
axes[0].legend(
    bbox_to_anchor=(0.55, 1.05, 1, 0.2), 
    loc="lower left",
    mode="expand", 
    borderaxespad=0, 
    ncol=2,
    fontsize=8,
)

plt.show()<jupyter_output><empty_output><jupyter_text>But what if we have missing values? E.g. let's imagine that galaxy 1 wasn't observed in the u band, while galaxy 2 wasn't observed in the u or y bands. We will mark these non-observations with the value 99:<jupyter_code># make a new copy of the samples
samples2 = samples.copy()
# make the non-observations
samples2.iloc[0, 1] = 99
samples2.iloc[1, 1] = 99
samples2.iloc[1, -1] = 99
# print the new samples
samples2<jupyter_output><empty_output><jupyter_text>Now if we want to calculate posteriors, we can't simply call `flow.posterior()` as before because the flow will think that 99 is the actual value for those bands, rather than just a flag for a missing value. What we can do, however, is pass `marg_rules`, which is a dictionary of rules that tells the Flow how to marginalize over missing variables.`marg_rules` must include:- "flag": 99, which tells the posterior method that 99 is the flag for missing values- "u": callable, which returns an array of values for the u band over which to marginalize- "y": callable, which returns an array of values for the y band over which to marginalize"u" and "y" both map to callable, because you can use a function of the other values to decide what values of u and y to marginalize over. For example, maybe you expect the value of u to be close to the value of g. In which case you might use:```"u": lambda row: np.linspace(row["g"] - 1, row["g"] + 1, 100)```The only constraint is that regardless of the values of the other variables, the callable must *always* return an array of the same length.For this example, we won't make the marginalization rules a function of the other variables, but will instead return a fixed array.<jupyter_code>marg_rules = {
    "flag": 99, # tells the posterior method that 99 means missing value
    "u": lambda row: np.linspace(26, 28, 40), # the array of u values to marginalize over
    "y": lambda row: np.linspace(24, 26, 40), # the array of y values to marginalize over
}

pdfs2 = flow.posterior(samples2, column="redshift", grid=grid, marg_rules=marg_rules)
fig, axes = plt.subplots(1, 2, figsize=(5.5, 2), dpi=120, constrained_layout=True)

for i, ax in enumerate(axes.flatten()):
    
    ax.plot(grid, pdfs[i], label="Posterior w/ all bands")
    ax.plot(grid, pdfs2[i], label="Posterior w/ missing bands marginalized")
    
    ztrue = samples["redshift"][i]
    ax.axvline(ztrue, c="C3", label="True redshift")
    ax.set(
        xlabel="redshift", 
        xlim=(ztrue - 0.25, ztrue + 0.25), 
        yticks=[]
    )
    
axes[0].legend(
    bbox_to_anchor=(0, 1.05, 2, 0.2), 
    loc="lower left",
    mode="expand", 
    borderaxespad=0, 
    ncol=3,
    fontsize=7.5,
)

plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11533505154639176, 'fraction_numerical': 0.025558419243986254, 'mean_word_length': 4.32837528604119, 'pattern_counts': {'":': 7, '<': 25, '<?xml version=': 0, '>': 28, 'https://': 3, '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': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '5815', 'license': 'MIT', 'n_tokens_mistral': 1621, 'n_tokens_neox': 1497, 'n_words': 520, 'path': 'examples/marginalization.ipynb', 'repo_name': 'jfcrenshaw/pzflow', 'size': '4656'}