import holoviews as hv
import numpy as np
import scistanpy as ssp
hv.extension('bokeh', inline=True)
Overview¶
This notebook contains an example workflow for modeling deep mutational scanning data (DMS). If you're familiar with DMS, continue to the next section ("Bayesian Analysis of DMS"). Otherwise, this cell provides an overview of the core concept:
DMS Background¶
When engineering a protein, how well it performs a desired task (e.g., catalysis, binding, etc.) is commonly referred to as its "fitness". Deep mutational scanning is a method for assigning fitness labels to different protein variants at scale. While the exact implementation varies, a typical DMS study follows the below workflow:
- We begin with a population of organisms (typically a model organism such as E. coli or S. cerevisiae). Within this population are subpopulations consisting of organisms carrying DNA that encode a specific protein variant--each subpopulation is defined by the variant DNA it carries.
- DNA is harvested from the full population and sequenced by next-generation sequencing. This results in a set number of reads (or "counts" in the context of DMS), with the number of reads pertaining to each subpopulation reflecting the overall frequency of that subpopulation within the full population.
- The population is then subjected to a selective pressure. Importantly, the ability of organisms to resist this selective pressure is tied to the fitness of the protein of interest. For example, if the goal were to engineer a protein critical for the production of a core element of cellular metabolism, the population would be grown in an environment absent that element, meaning that only subpopulations whose proteins were active could grow and propagate.
- Once selection is complete, next generation sequencing is performed again. Subpopulations expressing proteins with higher fitness should be enriched in the returned counts relative to their initial abundance and vice versa.
- "Fitness" for a protein is then defined as the ratio of normalized ending counts (where "normalization" is relative to the wild-type (unmutated) protein) to normalized beginning counts. The next section covers this in more detail.
Bayesian Analysis of DMS¶
Notably, DMS does not provide direct readouts of fitness. Instead, fitness values must be inferred. While, again, the exact procedure for doing this will vary with the DMS study, a typical approach is to take the relative ratio of normalized ending counts to normalized starting counts as below:
$$ \begin{align*} f_k = \frac{c_{k, f} / c_{wt, f}}{c_{k, i} / c_{wt, i}} \end{align*}, $$
where $c_{k, f}$ and $c_{k, i}$ are the final and starting counts for variant $k$, respectively, and $c_{wt, f}$ and $c_{wt, i}$ are the final and starting counts for the wild-type protein. Performing the above calculation for all variants gives us their fitnesses.
The value $f_k$ is also commonly called the "enrichment ratio" of the DMS study. Importantly, it approximates both (i) the abundance of each variant relative to other variants in the population at a given timepoint (hence division by the wild-type counts in both the numerator and denominator) and (ii) the abundance of a given variant to itself across timepoints (hence division of the end timepoint by the starting timepoint). The word "approximates" is stressed in the previous sentence, as the counts returned by NGS are an indirect correlate of true variant abundance.
As an analogy, if I had a bag full of marbles (variants) and drew 10 marbles from it, 4 being blue and 6 being green (e.g., "blue" being one variant and "green" being another), I cannot know with certainty what the true proportion of marbles is in that bag. Absent other information, a guess of 40% blue and 60% green seems most reasonable, but we intuitively know that it could be something else. 45% blue and 55% green might be the actual proportions. It's even possible, however unlikely, that there are only 4 blue marbles in the entire bag and we just so happened to pick them all, in which case, assuming the bag contains >>10 total marbles, the proportion of blue marbles is vanishingly small. The key point is that, because our marble counts are the result of a random process, there will always be some measure of uncertainty in our inference of the true proportions.
NGS can be thought of as drawing counts of specific variants from a population of unknown proportions. Just as with the marble example, absent additional information, the relative ratios of counts is the most probable reflection of true proportions, and this is, of course, what the enrichment ratio equation above is calculating. "Most probable" does not necessarily mean "probable", however, and the goal of Bayesian analysis is, in addition to identifying the most probable solution, to also quantify how probable a given solution may be.
Okay, all background and motivation out of the way, this brings us to building a Bayesian model of the above typical DMS scenario. One way to think about constructing a Bayesian model of any observable is to reconceptualize the data collecton process as a data generation process. The output of this process is our observables, which are a reflection of all underlying contributors of uncertainty (modeled as "parameters" SciStanPy) and their deterministic transformations (modeled as "transformations" in SciStanPy). For example, we have already established that one contributor of uncertainty is the sequencing process itself, which outputs discrete counts reflecting the true abundance of each variant in the selected population. A natural distribution for modeling this situation would be the Dirichlet distribution followed by the multinomial distribution
$$ \begin{align*} \mathbf{\theta_i} &\sim \text{Dirichlet}(\mathbf{\alpha}) \\ \mathbf{c}_i &\sim \text{Multinomial}(\mathbf{\theta_i}), \end{align*} $$
where $\alpha \in \mathbb{R}^K$ is a hyperparameter controlling our prior beliefs about the initial distribution of proportions, $\mathbf{c}_i \in \mathbb{W}^{K}$ defines the starting counts for $K$ total proteins and $\mathbf{\theta_i} \in \{\mathbb{R}^K | 0 \leq \theta_{i_k} \leq 1, \sum_{k=1}^K \theta_{i_k} = 1\}$ their proportions (i.e., the probability simplex).
The above is the first step of our "data generation" process. The next step is to apply the selective pressure. We will model the selection process as applying some multiplicative correction (i.e. an enrichment ratio) to the input proportions, not counts.
$$ \begin{align*} \mathbf{f} &\sim \text{Exponential}(\beta) \\ \mathbf{\theta_f} &= \frac{\mathbf{\theta_i} * \mathbf{f}}{\sum_{k=1}^K \theta_{i_k} * f_k}, \end{align*} $$
where $\mathbf{f} \in \{\mathbb{R}^K | R_k \geq 0\}$, $*$ indicates elementwise multiplication, and $\mathbf{\theta_f}$ gives the proportions after selection. Two things should be highlighted:
- We place an exponential prior on $\mathbf{f}$. This is to reflect the belief that variants are more likely to have low fitness than high fitness. The extent to which this belief is enforced is controlled by the hyperparameter $\beta$.
- $\mathbf{\theta_f}$ is not drawn from a distribution. Instead, it is modeled as the deterministic transformation of two parameters, $\mathbf{\theta_i}$ and $\mathbf{f}$. Implicitly, then, our uncertainty of the value of $\mathbf{\theta_f}$ results from our uncertainty in both the initial population proportions and the fitness of those proportions. In terms of data generation, $\mathbf{\theta_f}$ is found at a later step in the generation process than the other parameters (it is further along in the dependency graph).
From the transformed parameter $\mathbf{\theta_f}$, we can model our output counts using another multinomial distribution:
$$ \mathbf{c}_f \sim \text{Multinomial}(\mathbf{\theta_f}). $$
In all, this gives us the following model:
$$ \begin{align*} \mathbf{\theta_i} &\sim \text{Dirichlet}(\mathbf{\alpha}) \\ \mathbf{f} &\sim \text{Exponential}(\beta) \\ \mathbf{\theta_f} &= \frac{\mathbf{\theta_i} * \mathbf{f}}{\sum_{k=1}^K \theta_{i_k} * f_k} \\ \mathbf{c}_i &\sim \text{Multinomial}(\mathbf{\theta_i}) \\ \mathbf{c}_f &\sim \text{Multinomial}(\mathbf{\theta_f}) \end{align*} $$
SciStanPy is designed to allow models like the above to be encoded in a Pythonic way and, by extension, to allow us to infer values of unobservables ($\mathbf{f}$ in our example here) from observables, propagating uncertainty between different model components. In the case of DMS, then, fitting the above model in SciStanPy gives us a distribution of potential models, and so a distribution of potential fitness values, all informed by the collected data.
Fitting a DMS Model in SciStanPy¶
The above section covered the key concepts of DMS. It also covered the core differences between standard modeling of DMS data and Bayesian modeling of DMS data. This section demonstrates how to fit the above-described Bayesian model using SciStanPy:
To begin, let's just simulate some example data:
def sample_data():
"""
Generate sample data for deep mutational scanning analysis.
Returns:
INPUT_COUNTS: Array of input counts for each variant.
LOG_INPUT_FREQS: Log frequencies of input variants.
LOG_OUTPUT_FREQS: Log frequencies of output variants after selection.
"""
# Sample input counts
rng = np.random.default_rng(1025)
input_freqs = rng.dirichlet(np.ones(10))
log_input_freqs = np.log(input_freqs)
input_counts = np.stack([rng.multinomial(10000, input_freqs)
for _ in range(3)])
# Sample enrichment factors
log_enrichment_factors = np.log(rng.exponential(0.1, size=(10,)))
# Generate output counts after selection
log_output_freqs = log_input_freqs + log_enrichment_factors
log_output_freqs -= np.log(np.sum(np.exp(log_output_freqs))) # Normalize
output_counts = np.stack([rng.multinomial(10000, np.exp(log_output_freqs))
for _ in range(3)])
return {
"INPUT_COUNTS": input_counts,
"LOG_INPUT_FREQS": log_input_freqs,
"OUTPUT_COUNTS": output_counts,
"LOG_OUTPUT_FREQS": log_output_freqs,
"LOG_ENRICHMENT_FACTORS": log_enrichment_factors
}
SAMPLE_DATA = sample_data()
The above simulates a deep mutational scanning experiment of 10 variants with sequencing performed in triplicate at both the beginning and end. Note that we are working in the log space for our implementation relative to the model described in the previous section. For 10 variants, this is not likely to be necessary, but as the size of the probability simplex grows, numerical precision will become a problem on the standard scale. The log scale is used here to both (i) demonstrate a unique feature of SciStanPy (log-simplexes are not currently natively supported in Stan or PyTorch) and (ii) demonstrate a model that would be more practical for the typical DMS experiment (which tend to have >>10 variants).
We can now model the process using SciStanPy:
# All model's inherit from `ssp.Model`
class DMSModel(ssp.Model):
# Define the structure of the model in the `__init__` method
def __init__(self, input_counts, output_counts):
# We're going to register default data for the input and output counts.
# This isn't necessary, but means you won't need to pass the observables
# into later methods.
super().__init__(
default_data={"input_counts": input_counts, "output_counts": output_counts}
)
# We now define our priors. Let's assume that we expect our enrichment
# ratios to follow an exponential distribution. The log-enrichment factors
# will then follow a exponential-exponential (Gumbel) distribution.
# Note: We define 10 independent log-enrichment factors using the "shape"
# argument.
# Note: The "beta" parameter here is the inverse of the scale parameter
# in Numpy/Scipy.
self.log_enrichment = ssp.parameters.ExpExponential(beta=10.0, shape=(10,))
# We reason that the input and output counts are multinomially distributed
# with some unknown frequency, which are the values we want to infer. To
# handle potentially small values, we will use an Exp-Dirichlet prior to
# model the log-frequencies.
self.log_input_freqs = ssp.parameters.ExpDirichlet(alpha=1.0, shape=(10,))
# From the log-input frequencies and log-enrichment factors, we can define
# a transformation that takes us to the output frequencies. We're in log
# space, so this is just addition followed by normalization. Note that,
# currently, all reductions and normalizations are performed over the last
# axis (this cannot be changed yet).
self.log_output_freqs = ssp.operations.normalize_log(
self.log_input_freqs + self.log_enrichment
)
# Finally, we can model our observed counts at both the beginning and end
# as multinomially distributed. Note that the name of the observable must
# match the name we used when registering default data. If not registering
# default data, you will need to provide the observables as keyword arguments
# in the relevant functions (again, with matching names).
# Note: We are using an alternate parametrization of the multinomial distribution
# here to keep in log space.
# Note: Numpy broadcasting rules apply, so the below will use the same
# 10 log-frequencies for all 3 replicates. Note that this is why we need
# `keepdims=True` when summing the counts to get `N`: (shapes (3, 1) and
# (10,) broadcast to (3 x 10), while (3,) and (10,) do not broadcast.
self.input_counts = ssp.parameters.MultinomialLogTheta(
log_theta=self.log_input_freqs,
N=input_counts.sum(axis=-1, keepdims=True),
shape=(3, 10),
)
self.output_counts = ssp.parameters.MultinomialLogTheta(
log_theta=self.log_output_freqs,
N=output_counts.sum(axis=-1, keepdims=True),
shape=(3, 10),
)
That's it! Model is defined. You'll note that the variable names in the above model correspond neatly to variables defined in the model definition from the previous section:
| Name Previous Section | Name in Model |
|---|---|
| $\ln{\mathbf{\theta_i}}$ | log_input_freqs |
| $\ln{\mathbf{\theta_f}}$ | log_output_freqs |
| $\ln{\mathbf{f}}$ | log_enrichment |
| $\mathbf{c_i}$ | input_counts |
| $\mathbf{c_i}$ | output_counts |
That is, SciStanPy models are designed to follow syntax that closely reflects standard probabilistic model definitions, with instance variables becoming model names and transformations automatically recorded.
Another note: SciStanPy models are, of course, Python classes, and can be extended in all the usual ways. This can allow for the construction of class hierarchies that greatly reduce the need for duplicated code.
Now, what can we do with our model? Let's create an instance of it and test out some SciStanPy operations. First up, let's do a prior predictive check:
# Build an instance
EXAMPLE_MODEL = DMSModel(
input_counts=SAMPLE_DATA["INPUT_COUNTS"],
output_counts=SAMPLE_DATA["OUTPUT_COUNTS"]
)
# Run a prior predictive check
EXAMPLE_MODEL.prior_predictive()
BokehModel(combine_events=True, render_bundle={'docs_json': {'143225d7-76f4-456e-b0bc-57edc7246af7': {'version…
The prior_predictive function brings up an interactive dashboard that lets you test out the effects of different values for hyperparameters on model observables and parameters. By default, updating any parameters with the sliders (and subsequently clicking "update model") will also update their values in the model. This allows you to explore model hyperparameter values interactively before moving on to fitting a model.
Depending on the parameter selected, you may want to increase the value for "Number of Experiments"--this is the number of draws made from the model to build the figure. The default ECDF view flattens the displayed array before calculation, which obscures any relationships within or between variables. You can additionally choose values for "Group By", which will plot separate lines (or violins) for each grouping dimension; any constants with appropriate dimensionality can also be selected as "Independent Variable" to plot relationships between components.
Let's say we are happy with our hyperparameter selection. We're not yet ready to commit to full MCMC sampling, but we want to get an estimate for our parameter values. We can perform a maximum likelihood estimate using PyTorch as the backend:
MLE = EXAMPLE_MODEL.mle(lr=0.01)
Epochs: 7%|▋ | 6823/100000 [01:03<14:21, 108.17it/s, -log pdf/pmf=245.32]
By default, maximum likelihood estiation will run for 100,000 steps or until the loss (negative log likelihood) has not decreased for 10 steps, whichever comes first. The output results object exposes a lot of additional functionality which is covered in great detail in the documentation. For the purposes of this example notebook, however, we'll look at just two: Extracting maximum likelihood estimates and bootstrapping observations:
MLE_ESTIMATES = {k: v.mle for k, v in MLE.model_varname_to_mle.items()}
MLE_ESTIMATES
{'log_enrichment': array([-0.88175307, -3.0444229 , -2.52572081, -2.72127213, -2.63158634,
-3.36346962, -4.35208336, -2.55653866, -0.82543532, -3.89821476]),
'log_input_freqs': array([-2.31319812, -1.04716232, -2.52949171, -2.33462835, -2.90948524,
-2.95035977, -2.54973785, -6.09674183, -2.15798691, -2.64728865]),
'input_counts': None,
'output_counts': None}
A dictionary of results linking model varnames to MLE estimates can be accessed with the model_to_varname_property. Note that observables do not have an MLE--observables are ground truth, so there is nothing to estimate.
Let's compare the estimates above to our known values from our simulated experiment:
hv.Scatter(
data={
"x": SAMPLE_DATA["LOG_ENRICHMENT_FACTORS"],
"y": MLE_ESTIMATES["log_enrichment"],
}
)
hv.Scatter(
data={"x": SAMPLE_DATA["LOG_INPUT_FREQS"], "y": MLE_ESTIMATES["log_input_freqs"]}
)
As we'd expect considering we know the exact generative process, pretty good alignment between MLE and true values. Obviously, in practice, we do not know the generative process--our model defines what we believe it to be, then we fit the model and evaluate how well the fit model describes the data. One simple way to evaluate goodness of fit is to bootstrap samples from the MLE and perform a posterior predictive check on the results:
INFERENCE_OBJ = MLE.get_inference_obj() # Bootstrapping
INFERENCE_OBJ.run_ppc() # Checking fit
BokehModel(combine_events=True, render_bundle={'docs_json': {'315e15f2-3cee-49ae-bc1b-8ab500f48833': {'version…
The above example is contrived, so the plots look a little too pristine compared to what you'd get in a real-world setting. Here's how to interpret them, though:
- The first plot shows the true rank (x-axis) of an observation against its values bootstrapped from the model. The x-axis does not mean anything in this case--it is shown by true rank to simplify data presentation when there are thousands or more observed points. The y-axis, however, shows the distribution of bootstrapped values (grey) and associated observed values (gold). A model that is fit well will have most gold dots falling within the shaded regions.
- The next plot has the same axis, but now shows the quantile of true observations relative to the bootstrapped distribution of observations. This figure is also designed to work with thousands or more data points, so it aggregates nearby points into hexagonal bins. It is not the most useful figure for data at the scale of this example; however, for larger datasets, what you would want to see is uniform distribution of probability across all quantiles at all values with the median line (grey) running down the center.
- The final plot is a quantile-quantile plot. Effectively, it is an ECDF over the quantile values for observables relative to bootstrapped samples. Because, by definition, each quantile should be uniformly represented in a sufficiently large sample from a perfectly calibrated, a perfectly calibrated (and fit) model will have a diagonal ECDF, indicated in the figure by the dashed line. The annotation "Absolute Deviance" gives the absolute area between the observed and ideal ECDF curves and can be used to measure how well calibrated and fit a given model is relative to another.
It should also be noted that the above plot is interactive! Use the dropdown to update the plot for different observables.
Now, one final note on the INFERENCE_OBJ variable. It exposes a special property inference_obj which is an ArviZ InferenceData instance holding all bootstrapped data. Use it to plug directly into the ArviZ ecosystem:
INFERENCE_OBJ.inference_obj # ArviZ InferenceData instance
-
<xarray.Dataset> Size: 160kB Dimensions: (chain: 1, draw: 1000, a: 10) Coordinates: log_input_freqs.alpha (a) float64 80B 1.0 1.0 1.0 1.0 ... 1.0 1.0 1.0 1.0 Dimensions without coordinates: chain, draw, a Data variables: log_enrichment (chain, draw, a) float64 80kB -3.052 ... -0.9929 log_input_freqs (chain, draw, a) float64 80kB -3.075 ... -2.559 -
<xarray.Dataset> Size: 480kB Dimensions: (chain: 1, draw: 1000, b: 3, a: 10) Coordinates: input_counts.N (b) int64 24B 10000 10000 10000 log_input_freqs.alpha (a) float64 80B 1.0 1.0 1.0 1.0 ... 1.0 1.0 1.0 1.0 output_counts.N (b) int64 24B 10000 10000 10000 Dimensions without coordinates: chain, draw, b, a Data variables: input_counts (chain, draw, b, a) float64 240kB 979.0 ... 716.0 output_counts (chain, draw, b, a) float64 240kB 3.211e+03 ... 122.0 -
<xarray.Dataset> Size: 480B Dimensions: (b: 3, a: 10) Dimensions without coordinates: b, a Data variables: input_counts (b, a) int64 240B 965 3504 777 963 545 ... 755 21 1125 691 output_counts (b, a) int64 240B 3172 1273 481 471 312 ... 86 14 3936 108
Alright, we're satisfied with our maximum likelihood estimation and ready to move on to full monte carlo sampling with Stan. Just call the below and SciStanPy will:
- Convert your SciStanPy model into Stan code.
- Compile that code.
- Run that code.
- Organize and return the results.
HMC_RES = EXAMPLE_MODEL.mcmc(iter_warmup=2000, iter_sampling=1000)
13:16:34 - cmdstanpy - INFO - compiling stan file /tmp/tmpmmet_h7o/model.stan to exe file /tmp/tmpmmet_h7o/model
13:17:03 - cmdstanpy - INFO - compiled model executable: /tmp/tmpmmet_h7o/model
13:17:03 - cmdstanpy - WARNING - Stan compiler has produced 2 warnings:
13:17:03 - cmdstanpy - WARNING -
--- Translating Stan model to C++ code ---
bin/stanc --filename-in-msg=model.stan --warn-pedantic --O1 --include-paths=/home/bwittmann/micromamba/envs/ssp_test/lib/python3.12/site-packages/scistanpy/model/stan --o=/tmp/tmpmmet_h7o/model.hpp /tmp/tmpmmet_h7o/model.stan
Warning: The parameter log_input_freqs_raw has no priors. This means either
no prior is provided, or the prior(s) depend on data variables. In the
later case, this may be a false positive.
Warning: The parameter log_enrichment has no priors. This means either no
prior is provided, or the prior(s) depend on data variables. In the later
case, this may be a false positive.
--- Compiling C++ code ---
g++ -std=c++17 -pthread -D_REENTRANT -Wno-sign-compare -Wno-ignored-attributes -Wno-class-memaccess -DSTAN_THREADS -I stan/lib/stan_math/lib/tbb_2020.3/include -O3 -I src -I stan/src -I stan/lib/rapidjson_1.1.0/ -I lib/CLI11-1.9.1/ -I stan/lib/stan_math/ -I stan/lib/stan_math/lib/eigen_3.4.0 -I stan/lib/stan_math/lib/boost_1.87.0 -I stan/lib/stan_math/lib/sundials_6.1.1/include -I stan/lib/stan_math/lib/sundials_6.1.1/src/sundials -DBOOST_DISABLE_ASSERTS -c -Wno-ignored-attributes -x c++ -o /tmp/tmpmmet_h7o/model.o /tmp/tmpmmet_h7o/model.hpp
--- Linking model ---
g++ -std=c++17 -pthread -D_REENTRANT -Wno-sign-compare -Wno-ignored-attributes -Wno-class-memaccess -DSTAN_THREADS -I stan/lib/stan_math/lib/tbb_2020.3/include -O3 -I src -I stan/src -I stan/lib/rapidjson_1.1.0/ -I lib/CLI11-1.9.1/ -I stan/lib/stan_math/ -I stan/lib/stan_math/lib/eigen_3.4.0 -I stan/lib/stan_math/lib/boost_1.87.0 -I stan/lib/stan_math/lib/sundials_6.1.1/include -I stan/lib/stan_math/lib/sundials_6.1.1/src/sundials -DBOOST_DISABLE_ASSERTS -Wl,-L,"/home/bwittmann/.cmdstan/cmdstan-2.37.0/stan/lib/stan_math/lib/tbb" -Wl,-rpath,"/home/bwittmann/.cmdstan/cmdstan-2.37.0/stan/lib/stan_math/lib/tbb" /tmp/tmpmmet_h7o/model.o src/cmdstan/main_threads.o -ltbb stan/lib/stan_math/lib/sundials_6.1.1/lib/libsundials_nvecserial.a stan/lib/stan_math/lib/sundials_6.1.1/lib/libsundials_cvodes.a stan/lib/stan_math/lib/sundials_6.1.1/lib/libsundials_idas.a stan/lib/stan_math/lib/sundials_6.1.1/lib/libsundials_kinsol.a stan/lib/stan_math/lib/tbb/libtbb.so.2 -o /tmp/tmpmmet_h7o/model
rm /tmp/tmpmmet_h7o/model.hpp /tmp/tmpmmet_h7o/model.o
13:17:03 - cmdstanpy - INFO - CmdStan start processing
chain 1 | | 00:00 Status
chain 2 | | 00:00 Status
chain 3 | | 00:00 Status
chain 4 | | 00:00 Status
13:17:05 - cmdstanpy - INFO - CmdStan done processing.
Converting CSV to NetCDF: 100%|██████████| 4/4 [00:27<00:00, 6.75s/it]
You can ignore warnings about parameters not having priors. This is a result of how data is passed to the model: that is, priors depend on data variables.
Now that we have the results, lets run some diagnostics:
_ = HMC_RES.diagnose()
Sample diagnostic tests results' summaries: ------------------------------------------- 0 of 4000 (0.00%) samples had a low energy. 0 of 4000 (0.00%) samples reached the maximum tree depth. 0 of 4000 (0.00%) samples diverged. R_hat diagnostic tests results' summaries: ------------------------------------------ 0 of 10 (0.00%) r_hats tests failed for log_enrichment. 0 of 10 (0.00%) r_hats tests failed for log_input_freqs. 0 of 10 (0.00%) r_hats tests failed for log_output_freqs. Ess_bulk diagnostic tests results' summaries: --------------------------------------------- 0 of 10 (0.00%) ess_bulks tests failed for log_enrichment. 0 of 10 (0.00%) ess_bulks tests failed for log_input_freqs. 0 of 10 (0.00%) ess_bulks tests failed for log_output_freqs. Ess_tail diagnostic tests results' summaries: --------------------------------------------- 0 of 10 (0.00%) ess_tails tests failed for log_enrichment. 0 of 10 (0.00%) ess_tails tests failed for log_input_freqs. 0 of 10 (0.00%) ess_tails tests failed for log_output_freqs.
The diagnostic tests and their meanings are described in greater detail in the full documentation. For our purposes here, what matters is that they all passed!
We stored the output of the function in a throw away variable, as we don't need it. However, it contains indices of failed samples and variables, as relevant. Also note that the SampleResults object has additional functionality for helping to diagnose failed samples, when they're present, via the plot_sample_failure_quantile_traces and plot_variable_failure_quantile_traces methods. See the full documentation for details on these methods.
Finally, we're going to want to do a posterior predictive check on our Stan samples. Note that, unlike the MLE example, these samples should be representative of the full posterior, not just the MLE. The same workflow as for evaluating MLE applies here, however:
HMC_RES.run_ppc()
BokehModel(combine_events=True, render_bundle={'docs_json': {'1dffcae8-2c6d-4ab4-b57c-21e451fb8dc8': {'version…
As before, the fit looks reasonable. Also as before, note that we can access an underlying ArviZ InferenceData object for further analysis. It will have some additional fields compared to the MLE-associated one, reflecting the richer information content of HMC samples:
HMC_RES.inference_obj
-
<xarray.Dataset> Size: 480kB Dimensions: (chain: 4, draw: 1000, a: 10) Dimensions without coordinates: chain, draw, a Data variables: log_enrichment (chain, draw, a) float32 160kB -1.069 -3.236 ... -3.867 log_input_freqs (chain, draw, a) float32 160kB -2.328 -1.053 ... -2.638 log_output_freqs (chain, draw, a) float32 160kB -1.142 -2.034 ... -4.514 -
<xarray.Dataset> Size: 960kB Dimensions: (chain: 4, draw: 1000, b: 3, a: 10) Dimensions without coordinates: chain, draw, b, a Data variables: input_counts (chain, draw, b, a) int32 480kB 1020 3470 792 ... 24 1168 681 output_counts (chain, draw, b, a) int32 480kB 3149 1251 454 ... 8 3920 119 -
<xarray.Dataset> Size: 112kB Dimensions: (chain: 4, draw: 1000) Dimensions without coordinates: chain, draw Data variables: lp__ (chain, draw) float32 16kB ... accept_stat__ (chain, draw) float32 16kB ... stepsize__ (chain, draw) float32 16kB ... treedepth__ (chain, draw) int32 16kB 7 7 8 7 7 7 7 7 ... 7 7 7 7 7 7 8 7 n_leapfrog__ (chain, draw) int32 16kB ... divergent__ (chain, draw) int32 16kB 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0 energy__ (chain, draw) float32 16kB 270.7 269.4 275.2 ... 284.3 285.1 -
<xarray.Dataset> Size: 240B Dimensions: (b: 3, a: 10) Dimensions without coordinates: b, a Data variables: input_counts (b, a) int32 120B 965 3504 777 963 545 ... 755 21 1125 691 output_counts (b, a) int32 120B 3172 1273 481 471 312 ... 86 14 3936 108 -
<xarray.Dataset> Size: 1kB Dimensions: (metric: 5, a: 10) Coordinates: * metric (metric) <U9 180B 'mcse_mean' 'mcse_sd' ... 'r_hat' Dimensions without coordinates: a Data variables: log_enrichment (metric, a) float64 400B 0.01 0.01 0.01 ... 1.0 1.0 1.0 log_input_freqs (metric, a) float64 400B 0.0 0.0 0.0 0.0 ... 1.0 1.0 1.0 log_output_freqs (metric, a) float64 400B 0.0 0.0 0.0 0.0 ... 1.0 1.0 1.0 -
<xarray.Dataset> Size: 1kB Dimensions: (metric: 5, a: 10) Coordinates: * metric (metric) <U9 180B 'ess_bulk' 'ess_tail' ... 'r_hat' Dimensions without coordinates: a Data variables: log_enrichment (metric, a) float64 400B 599.8 600.1 597.2 ... 1.0 1.0 1.0 log_input_freqs (metric, a) float64 400B 4.149e+03 5.995e+03 ... 1.0 1.0 log_output_freqs (metric, a) float64 400B 3.942e+03 3.907e+03 ... 1.0 1.0 -
<xarray.Dataset> Size: 12kB Dimensions: (chain: 4, draw: 1000) Dimensions without coordinates: chain, draw Data variables: low_ebfmi (chain, draw) bool 4kB False False ... False False max_tree_depth_reached (chain, draw) bool 4kB False False ... False False diverged (chain, draw) bool 4kB False False ... False False -
<xarray.Dataset> Size: 198B Dimensions: (metric: 3, a: 10) Coordinates: * metric (metric) <U9 108B 'r_hat' 'ess_bulk' 'ess_tail' Dimensions without coordinates: a Data variables: log_enrichment (metric, a) bool 30B False False False ... False False log_input_freqs (metric, a) bool 30B False False False ... False False log_output_freqs (metric, a) bool 30B False False False ... False False