<jupyter_start><jupyter_text>Exploring PyMC3 for Parameter estimation of photobleaching<jupyter_code>"""
Attempting to support multiple BLAS
MKL = 
OPENBLAS = default pip insetlal blas?
BLIS = Maybe better on AMP

OMP = "Open Multi-Process"? library

Desire is to turn off all under-the-hood parallelization and control that ourselves
But AS thinks that there's JAX issue causing unexpectedd multi-thread. TBD.
"""

# BLAS implenentions
%env MKL_NUM_THREADS=1
%env OPENBLAS_NUM_THREADS=1
%env BLIS_NUM_THREADS=1

# Open MP
%env OMP_NUM_THREADS=1

# JAX
%env NUM_INTER_THREADS=1
%env NUM_INTRA_THREADS=1

# The xla_force_host_platform_device_count flag is the manual control of 8 threads
n_threads=8
%env XLA_FLAGS=--xla_cpu_multi_thread_eigen=false --xla_force_host_platform_device_count={n_threads}
import theano.tensor as aet
import theano as aesara
import jax
import jax.numpy as jnp
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import pymc3 as pm
from scipy import stats, optimize
assert jax.device_count() == n_threads<jupyter_output><empty_output><jupyter_text>Generate Simulated Data  TheoryPhotobleaching experiments are simplified control experiments of the Erisyon platform whereinwe add purified control peptides with some known "count" of attached dyes (fluorophores).We then illuminate the sample and record cycles with no stage movement and minimal chemistry.The result is a series of images of the single-molecule peptides, one measurement per cycle,over many indivdual peptide "peaks".In the below code we model only the post-signal processing measurement data which we callthe "radiometry". This takes the form of a "Radiometry Matrix" or "radmat" where each columnis a measurement over time and each row is an independent peak that isassumed to have come from one peptide.(The order of the rows of a radmat is arbitrary and we often sort by mean intensity merely fordisplay athesthtics.)Realistic data has a variety of problems/annoyances that are not modelled here. Examples:  - collisions between peptides  - contamination  - untethered diffusing dyes/peptides that transiently corrupt the signal of a tehtered peptide  - dye-dye interactions Naming Conventions* `n_*:` A count of something* `p_*`: A probability* `*_i`: An index* `*_iz`: A vector of indicies (_iz vs _is because _is looks too much like the verb "is")* `true_*`: True parameters values for fake data* `hyper_*`: A hyper parameter that is not fit Nomenclature* `cycle`: A measuremnt/chemistry cycle. For these experiments effectively 1 cycle = 1 unit of time = 1 measurment* `n_cycles`: The number of measurements (images) that were captured.* `survival`: The number of cycles a dye survives before bleaaching* `dyetrack` abbreviated (`dyt`): A vector of counts over cycles. Eg: "22211111000" means a given track  had 2 live dyes for 3 cycles, 1 live for 5 cycles followed by 3 cycles of darkness.* `dytmat`: The matrix of dyetracks. n_columns == n_cycles, n_rows == n_samples where each value is an integer count of live dyes.* `radmat`: The matrix of radiometry. n_columns == n_cycles, n_rows == n_samples where each value is a floating point value  that was measured as some peak.* `n_count`: The maximum number of dyes that could possibly be attached to each peptide.* `row_k`: A constant that multiples each row of the radmat that is due to illumination differences from peak-to-peak.  This term conflated with gain and is defined as having mean == 1.0. Parameters* `gain`: The illumination scale by which a single dye is measured; assumed to be linear.* `bg_sigma`: Standard deviations of background (dark) pixels* `sigma`: Standard deviations of the measurement of a dye. Variance (ie sigma^2) is assumed to grow linearly  with number of dyes. It is also assumed independent of bg_sigma. Ie, measured variance = bg_sigma^2 + sigma^2 * n_dyes* `p_bleach`: Probabilty of a dye irreversibly going dark during a measurement cycle. This is  modeled as a discrete event occurring *between* cycles (although this is not realistic because real dyes can  bleach *during* a measurment and result in a fractional measurement).* `p_attach`: Probability that a dye successfully attaches to a peptide *and* is fluorescent at the  first measurement (cycle 0).<jupyter_code>def find_gamma_params(lower, upper, alpha=0.05):
    """
    Convert upper-lower boundaries into a Gamma distibution which is strictly positive.
    Ie converts lower, upper to 5th and 95th percentile
    """
    def cdf(x, lower, upper):
        a, scale = x
        dist = stats.gamma(a=a, scale=scale)
        cdf_lower = dist.cdf(lower)
        cdf_upper = dist.cdf(upper)
        return np.array([cdf_lower - alpha, cdf_upper - (1 - alpha)])

    # There's probably a closed form solution but this is an easy approximation
    # and speed is irrelevant.
    opt = optimize.least_squares(cdf, [1, 1000], args=(lower, upper))
    if not opt.success:
        raise ValueError("Optimization of gamma parameters failed.")

    distr = stats.gamma(a=opt.x[0], scale=opt.x[1])
    
    # Convert to a 2-parameter distribtiuon that pm.Gamma(**kwargs) will accept
    return dict(
        mu=distr.mean(),
        sd=distr.std(),
    )

def find_beta_params(lower, upper, alpha=0.05):
    """
    Convert lower, upper priors into a Beta distr. that is 0-1 (ie used for probabilitic priors).
    Smilrly to find_gamma_params this is defined on 5th and 95th percentile
    """
    def cdf(x, lower, upper):
        a, b = x
        dist = stats.beta(a=a, b=b)
        cdf_lower = dist.cdf(lower)
        cdf_upper = dist.cdf(upper)
        return np.array([cdf_lower - alpha, cdf_upper - (1 - alpha)])

    opt = optimize.least_squares(cdf, [1, 1], args=(lower, upper))
    if not opt.success:
        raise ValueError("Optimization of gamma parameters failed.")

    # Convert to a 2-parameter distribtiuon that pm.Beta(**kwargs) will accept
    return dict(
        alpha=opt.x[0],
        beta=opt.x[1],
    )<jupyter_output><empty_output><jupyter_text>Simulate data<jupyter_code># True parameters of fake data
true_gain = 10000.0  # Realistic: 5000 - 20000
true_sigma = 900.0  # Realistic: 500 - 2000
true_bg_sigma = 500.0  # Realistic: 100 - 1000
true_p_bleach = 0.05   # Realistic: 0.01 - 0.20
true_p_attach = 0.9  # Realistic: 0.5 - 0.9
true_row_k_sigma = 0.15  # Realistic: 0.05 - 0.15

# Hyper parameters
hyper_n_count = 2  # Realisitc: 1-4
hyper_n_cycles = 50  # Realistic: 50-200
hyper_n_training_rows = 128 # Realistic: 10_000 - 100_000
np.random.seed(42)

cycle_iz = np.tile(np.arange(hyper_n_cycles), (hyper_n_training_rows, 1))
attached = np.random.uniform(0, 1, size=(hyper_n_training_rows, hyper_n_count)) < true_p_attach
survival = np.random.geometric(true_p_bleach, size=(hyper_n_training_rows, hyper_n_count))
dytmat = np.where(cycle_iz[:, None, :] <= survival[:, :, None], 1, 0)
dytmat = dytmat * attached[:, :, None]
dytmat = np.sum(dytmat, axis=1)
radmat = np.random.normal(dytmat * true_gain, np.sqrt(dytmat * true_sigma**2 + true_bg_sigma**2))

# We defined row_k originally as a Gaussian with mean 1 and sigma as a free parameters
# But this has the problem that it allows nonsensical negative values so aa better
# choice would be strictly positive.
# In this reformulation, we model the row_k as a lognormal with mu 0
#   row_k = exp(normal(mu=0, sigma=??))
# The sigma of that normal in log is chosen so that the stdev of row_k is equa to
# the row_k_sigma which results in the following sqrt(4 * ...)
# Another option might be to use a Gamma or define the parameters as log_row_k_sigma.
log_row_k_scale = np.sqrt(np.log((1 + np.sqrt(4 * true_row_k_sigma ** 2 + 1)) / 2))
row_k = np.exp(log_row_k_scale * np.random.randn(hyper_n_training_rows))
radmat = radmat * row_k[:, None]

# The instrument can not measure a dark cycle 0 (it can't find the LACK of data)
# TODO: Think about how to deal with these dark rows.
#mask = radmat[:, 0] > 2000.0

X_train = radmat
sns.heatmap(X_train[np.argsort(X_train.sum(1))[::-1], :], center=0)<jupyter_output><empty_output><jupyter_text>Log-scale integrationZBS Definitions:* $ \theta $ = The vector of global parameters (shared between all rows)* $ i $ = An index into the data; the "i-th" row.* $ y_i $ = The i-th row of data* $ L_i $ = The Log-liklihood of the i-th row.* $ P(y_i | \theta) $ = Probability of measuring the observed data $y_i$ at parameters $\theta$* $ P(y_i | \theta, \text{row_k}_i = x) $ = Probability of measuring the observed data $y_i$ at parameters $\theta$ and with row_k on the i-th row equal to x.* $ P(\text{row_k}_i = x | \theta) $ = Probability of $\text{row_k}_i$ being x given $\theta$.Note that a particular row's value of k is not a global parameter but its probability does depend on a global paratmer, specifically "row_k_sigma".Our goal is to marginalize row_k in the model for each row.That means we need to evaluate an integral:$$L_i(\theta) = P(y_i | \theta) = \int P(y_i | \theta, \text{row_k}_i = x) P(\text{row_k}_i = x | \theta) dx $$$$ = \int e^{f(x)} dx $$where we break the P product up into the sum of two log terms:$$ f(x) = \log P(y_i | \theta, \text{row_k}_i = x) + \log P(\text{row_k}_i | \theta) $$So... $f(x)$ is the tricky piece here...because $f(x)$ is expensive to evaluate from our HMM (it requires the forward-pass to be evaluated).Thus, we want to find a spline $\hat{f}$ that approximates our log-likelihood $L$on an interval so that we can do integration (marginalization).To find that spline, we evaluate the log-likelihood function at a number ofdfferent values of row_k as well as the derivatives at those row_k values.The spline allows us to evaluate intermediate values along row_k and thus integrate $e^{\hat{f}(x)}$ smoothly.Ideally we would use scipy for spline evaluations but scipy doesn't provide derivaties of the function $L_i{\theta}$this the need for an autodiff framework, ie JAX.When performing integration, we want to do it accurately and thus we need to focus on the regions of highest density.To do this we estimate the integration error in each sub-intervaland sub-divide the interval with the largest estimated integration error.The integration error is estimated using estimated curvature on that sectionand assume there's a small variation in spline and see how much it changesthe integrated value. (Details to follow...)We might be able to change the integration strategyto leverage Taylor expansions, which gives us bounds on error.<jupyter_code># Utility functions for modelling the integration in JAX with splines

def make_cubic_spline_basis(N):
    """Construct cubic spline basis."""
    x = np.linspace(0, 1, N)

    p00 = (2 * x - 3) * x * x + 1
    p10 = ((x - 2) * x + 1) * x
    p01 = (-2 * x + 3) * x * x
    p11 = x * x * (x - 1)

    # Next order, val and deriv zero at 0 and 1, second derivative 1 at 1/2
    p22 = x * x * (x - 1) ** 2

    A = np.c_[p00, p01, p10, p11]
    return A, p22


def insert_one(x, idx, value):
    """Helper function to insert value on JAX arrays. (They are immutable.)"""
    out = jnp.zeros(len(x) + 1)

    idxs = jnp.arange(len(x))
    zeros = jnp.zeros_like(x)
    out = out + jax.lax.dynamic_update_slice(out, jnp.select([idxs < idx], [x]), (0,))
    out = jax.lax.dynamic_update_slice(
        out, jnp.select([idxs >= idx, idxs < idx], [x, out[1:]]), (1,)
    )
    out = jax.ops.index_update(out, idx, value)

    return out

def sample_log_integral_exp(rng_seed_key, x0, x1, y0, y1, dydx0, dydx1):
    """
    Sample from the distr. that is defined by the log-density defined by a spline.
    This is used when we reconstuct the marginalized parameters.
    """
    
    # Note that JAX does not like kwargs so for now the n_samples is hard-coded.
    n_spline_interpolation_points = 1000

    A, _ = make_cubic_spline_basis(n_spline_interpolation_points)
    h = x1 - x0
    coefs = jnp.stack([y0, y1, dydx0 * h, dydx1 * h])

    x = jnp.linspace(x0, x1, n_spline_interpolation_points)[:-1]
    interpolated = A[:-1, :] @ coefs

    log_norm = jax.scipy.special.logsumexp(interpolated)

    # DRAW 1 sample from the spline-section that models our row_k distr.
    return jax.random.choice(rng_seed_key, x, p=jnp.exp(interpolated - log_norm))<jupyter_output><empty_output><jupyter_text>Functions for performing integrationAnd then now we write the functions that perform integration for us.<jupyter_code>n_spline = 100
n_refine = 5


def log_integral_exp_step_jax(x0, x1, y0, y1, dydx0, dydx1):
    """
    Computes the integral from x0 to x1 of one section of the spline given
    the spline and its boundary conditions including derivatives.
    
    JAX will compile this function.
    """
    
    # spline_mat.shape == (n_spline, n_coef)
    # Each column is the basis function of one spline segment
    spline_mat, error_basis = make_cubic_spline_basis(n_spline)
    interval_len = x1 - x0
    
    # We will later need to refine some spline sections thus we need
    # some estimate of the error in the integral of this section 
    # which we compute here and store for later decision whether or not
    # to subdivide in that we subdive the interval with the largest
    # error estimate.

    # Two different error estimation functions were explored
    # A suboptimal eror function only means slower conversion to
    # as opposed to getting the wrong answer.

    # error_coef = h * (jnp.abs(dydx0) + jnp.abs(dydx1)) / 4
    error_coef = interval_len * jnp.maximum(jnp.abs(dydx0), jnp.abs(dydx1))

    coefs = jnp.stack([y0, y1, dydx0 * interval_len, dydx1 * interval_len])

    # COMPUTE the values of the spline. The basis functino already had
    # the sampling along the independent variable of number of subdivisions of n_spline
    interpolated_without_error = spline_mat @ coefs
    
    # ADD in the error terms
    interpolated_with_error = interpolated_without_error + error_coef * error_basis

    def integrate(interpolated):
        """
        Integrate the the interpolted spline samples using the trapzoidal rule
        which is:
            (boundary_values / 2 + sum(interior_values)) * delta_x
        where delta_x = the distance between the equally spacing samples
        
        Arguments:
            interpolated: a vector of log(values) of the spline function evaluated at
            regular intervals.
        
        Note inheritance of interval_len from outer scope.
        """
        
        # Interior integral
        interior_values = interpolated[1:-1]
        interior_log_integral = jax.scipy.special.logsumexp(interior_values)
        
        # Bounday integral
        # Rescale the "/ 2" of the trapazoldal rule to log becomes "- log(2)"
        bounday_log_integral = jnp.logaddexp(interpolated[0], interpolated[-1]) - jnp.log(2)
        
        full_log_integral = jnp.logaddexp(interior_log_integral, bounday_log_integral)
        log_delta_x = jnp.log(interval_len) - jnp.log(n_spline - 1)
        return full_log_integral + log_delta_x

    # INTEGRATE spline samples with and without error
    log_integral_without_error = integrate(interpolated_without_error)
    log_integral_with_error = integrate(interpolated_with_error)

    # This is like the logexpsum trick but we want a log difference.
    #
    #   error
    #   == exp(log_error)
    #   == integral_without_error * (exp(log_integral_with_error - log_integral_without_error) - 1)
    #      ... factor out integral_without_error ...
    #   == integral_without_error * (integral_with_error / integral_with_error - 1)
    #   == exp(log_integral_with_error) - exp(log_integral_without_error)
    #
    # Note: jnp.expm1 computes a numerically stable version of (exp() - 1)

    log_error = log_integral_without_error + jnp.log(jnp.expm1(log_integral_with_error - log_integral_without_error))

    # Right now we always split at the middle of the interval if we choose to subdivide.
    # A possible improvment would be to choose the point were the residual is maximum
    # which we could calculate here.

    return log_integral_without_error, log_error


def _log_integral_exp_adapt_jax_debug(func, x):
    """
    A helper called by the integration or sampling entrypoints below so that
    we can return intermediate calculation results.
    
    Args:
        func: The function we want want to evaluate
        x: initial subdivisions of the integration domain
    
    Returns:
        total: sum(exp(log_p)) == the total integral of the function
        x: where we evalulate the spline
        y: log function values at x
        dydx: derivate of the log function values
        logp_val: integral between each x point (1 fewer sample than len(x))
        errors: error estimates of each section
        errors_trace: The history of the spline subdivisions
    """

    # We are assuming that func can be called by JAX and thus JAX can do autodiff on func

    # Forward or backward gradients? From experimentation they
    # have different compilation vs execution speeds.
    
    # Backwards...
    # func_grad = jax.value_and_grad(func)
    
    # Forwards...
    func_grad = lambda x: jax.jvp(func, (x,), (1.,))

    _, error_basis = make_cubic_spline_basis(n_spline)

    y, dydx = jax.vmap(func_grad, out_axes=(0, 0))(x)

    log_integral_exp_step_jax_vec = jax.vmap(
        log_integral_exp_step_jax, in_axes=[0, 0, 0, 0, 0, 0], out_axes=(0, 0)
    )

    logp_vals, errors = log_integral_exp_step_jax_vec(
        x[:-1], x[1:], y[:-1], y[1:], dydx[:-1], dydx[1:]
    )

    total = jax.scipy.special.logsumexp(logp_vals)

    errors_trace = []

    for _ in range(n_refine):
        max_idx = jnp.argmax(errors)

        errors_trace.append((errors, max_idx))

        old_logp = logp_vals[max_idx]

        x_middle = (x[max_idx + 1] + x[max_idx]) / 2
        y_middle, dydx_middle = func_grad(x_middle)

        left, left_error = log_integral_exp_step_jax(
            x[max_idx], x_middle, y[max_idx], y_middle, dydx[max_idx], dydx_middle
        )
        right, right_error = log_integral_exp_step_jax(
            x_middle,
            x[max_idx + 1],
            y_middle,
            y[max_idx + 1],
            dydx_middle,
            dydx[max_idx + 1],
        )

        x = insert_one(x, max_idx + 1, x_middle)
        y = insert_one(y, max_idx + 1, y_middle)
        dydx = insert_one(dydx, max_idx + 1, dydx_middle)

        logp_vals = insert_one(logp_vals, max_idx, left)
        logp_vals = jax.ops.index_update(logp_vals, max_idx + 1, right)

        errors = insert_one(errors, max_idx, left_error)
        errors = jax.ops.index_update(errors, max_idx + 1, right_error)

        # Todo update total instead of recomputing to simplify grad ops
        total = jax.scipy.special.logsumexp(logp_vals)
        #total = jnp.logaddexp(total, jnp.logaddexp(left, right))
        #total = jnp.expm1(total - scipy.)

    return total, x, y, dydx, logp_vals, errors, errors_trace


def log_integral_exp_adapt_jax(func, x):
    """
    Wrapper around log_integral_exp_adapt_jax_debug() to give
    JAX back only the [0] which is the log_total_integral.
    """
    return log_integral_exp_adapt_jax_debug(func, x)[0]


def log_integral_exp_adapt_jax_sample(key, func, x):
    """
    Choose one of the spline intervals based on its
    relative probabilty mass using jax.random.choice()
    
    Then inside of that spline interval pick a point based
    on its relative probabilty mass within that spline
    using our spline function sample_log_integral_exp()
    """
    total, x, y, dydx, logp_vals, *_ = log_integral_exp_adapt_jax_debug(func, x)

    key, key2 = jax.random.split(key)

    interval = jnp.arange(len(logp_vals), dtype=np.int32)
    interval_idx = jax.random.choice(key, interval, p=jnp.exp(logp_vals - total))

    # i, j are start and stop indices of the randomly chosen spline section
    i = interval_idx
    j = interval_idx + 1
    return sample_log_integral_exp(key2, x[i], x[j], y[i], y[j], dydx[i], dydx[j])
# # Test-drive the integration

# def example_function_to_integrate(x):
#     """A test function to integrate. Something this is not a polynomial in log-space."""
#     return 100 * (jnp.sin(x) - x ** 2 - 0.3)

# x = np.linspace(-3, 3, 200)
# y = example_function_to_integrate(x)
# fig, (ax1, ax2) = plt.subplots(1, 2)
# ax1.plot(x, y)
# ax2.plot(x, np.exp(y))

# # For a baseline, see how scipy's integration functions as a ground truth.
# # (This function doesn't appear to have a closed-form solution when integrating w.r.t. x.)
# from scipy import integrate
# quad, *out = integrate.quad(lambda x: np.exp(example_function_to_integrate(x)), -3, 3., full_output=1)

# print(
#     f"SCIPY METHOD:\n"
#     f"  integral value={quad:2.1e}\n"
#     f"  estimated_integration_error_by_scipy={out[0]:2.1e}\n"
#     f"  n_func_evals_by_scipy={out[1]['neval']:d}"
# )

# # Versus our method
# integral, x, y, dydx, logp_vals, error_estimates, err_trace = log_integral_exp_adapt_jax_debug(
#     example_function_to_integrate, np.array([-3., 3.])
# )

# print(
#     f"OUR METHOD:\n"
#     f"  integral value={integral:2.1e}\n"
#     f"  relative_error_compared_scipy={100 * (quad - np.exp(integral)) / quad:2.1e}%\n"
#     f"  n_func_evals_our_method={len(x):d}"
# )

# # VISUAIZE the splines
# fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12,6))

# A, _ = make_cubic_spline_basis(n_spline)
# h = x[1:] - x[:-1]
# c = np.stack([y[:-1], y[1:], dydx[:-1] * h, dydx[1:] * h])
# splines = A @ c
# x_splines = np.stack([np.linspace(x0, x1, len(A)) for x0, x1 in zip(x[:-1], x[1:])])
# errors = [
#     integrate.quad(lambda x: np.exp(example_function_to_integrate(x)), x0, x1)[0] - np.exp(logp_val)
#     for logp_val, (x0, x1) in zip(logp_vals, zip(x[:-1], x[1:]))
# ]
# ax1.plot(error_estimates + np.log(10), label="error estimate")
# ax1.plot(np.log10(np.abs(errors)), label="actual error")
# ax1.legend()
# ax1.set_title("error estimates log-scale")

# x_dense = np.linspace(-3, 3, 100)
# ax2.plot(x_splines.T, np.exp(splines), linewidth=4)
# ax2.plot(x_dense, np.exp(example_function_to_integrate(x_dense)), linewidth=0.8, color='red')
# ax2.set_title("non-log space")

# x_dense = np.linspace(-3, 3, 100)
# ax3.plot(x_splines.T, splines, linewidth=4)
# ax3.plot(x_dense, example_function_to_integrate(x_dense), linewidth=0.8, color='red')
# ax3.set_title("log-space")

# # As we can see in the bottom plot, our splines (thick lines)
# # diverge from the original (red thin line) very slightly. (See the blue section on the left.)
# # However if you look at the non-log-space plot, the blue region has
# # very little probability density associated with it. Therefore, the integration error
# # located there is really negligible.

# # The integral above is not one so if we normalize it then it is
# # a probabity density function we can sample from.
# # Test sampling from that derived PDF.
# vectorized_func = jax.vmap(
#     log_integral_exp_adapt_jax_sample,
#     in_axes=[0, None, None]
# )

# # Vectoring over random nuber seeds
# initial_seed = jax.random.PRNGKey(10)

# # Use jax.random to generate 10000 random numbers from the inital_seed
# # Which we will then use as the seeds for each sample
# # (which jax calls "keys")
# keys = jax.random.split(initial_seed, 10_000)
# x_samples = vectorized_func(keys, example_function_to_integrate, np.array([-3., 3.]))
# sns.displot(x_samples)<jupyter_output><empty_output><jupyter_text>On marginalization in MCMCThe key concept of marginalizing in MCMC:  - With samples (ie from MCMC) we can estimate an integral  - or, with an integral we can sample...    that is, an algorithm that integrates can typically be easily adapated to sample.The HMC algorithm that we which to use in PyMC3 (specially NUTS) use gradients and thus do not deal well with:  - discerete distribtions  - multi-modal distributionsThus, we wish to marginalize out either of those cases by usingan intergal for multi-modal or sums for discretes.The inner loop of the MCMC sampler is doing a callback which asks:"At theta, I want a logp and its gradient."To compute that logp and its gradient we need the integral of the distributionsthat we have marginalized.In our case, we get the gradient from autodiff via JAX. Ie webuild up splines in JAX's framework and it differentiates themwhich is why we couldn't use the scipy implementations of the splines.<jupyter_code>from collections import namedtuple
Theta = namedtuple('Theta', ['lp_init', 'ltrans', 'row_gain_sigma', 'mu_unscaled', 'sigma_unscaled'])
# HERE
# Dyetrack marginalization

def forward_step_log(lalpha, ltrans, lp_data):
    return jax.scipy.special.logsumexp(lalpha + ltrans.T, axis=1) + lp_data


def forward_log_scan(lp_init, ltrans, lp_data):
    lalpha = lp_init + lp_data[0, :]
    lalpha, _ = jax.lax.scan(lambda d, x: (forward_step_log(d, ltrans, x), None), lalpha, lp_data[1:])
    return lalpha


def sample_path(key, lp_init, ltrans, lp_data):
    """
    Sample the dyetracks from 
    """
    lalpha_first = lp_init + lp_data[0, :]

    def inner(lalpha_current, lp_data_item):
        lalpha_next = forward_step_log(lalpha_current, ltrans, lp_data_item)
        return (lalpha_next, lalpha_next)

    _, lalphas = jax.lax.scan(inner, lalpha_first, lp_data[1:])
    lalphas = jnp.concatenate([lalpha_first[None, :], lalphas], axis=0)

    lalphas_rev = lalphas[::-1]

    def choice(key, log_p):
        log_norm = jax.scipy.special.logsumexp(log_p)
        idxs = jnp.arange(len(log_p), dtype=np.int8)
        return jax.random.choice(key, idxs, p=jnp.exp(log_p - log_norm))

    def draw_step(state, args):
        key, lalpha = args
        log_p = lalpha + ltrans[:, state]
        state_prev = choice(key, log_p)
        return (state_prev, state_prev)

    keys = jax.random.split(key, len(lp_data))

    state = choice(keys[0], lalphas_rev[0])

    _, path = jax.lax.scan(draw_step, state, (keys[1:], lalphas_rev[1:]))

    return jnp.concatenate([state[None], path])[::-1]

def _compute_lp_data(row_gain, theta, data):
    mu = theta.mu_unscaled * row_gain
    sigma = theta.sigma_unscaled * row_gain

    lp_data = jax.scipy.stats.norm.logpdf(
        loc=mu[None, :], scale=sigma[None, :], x=data[:, None]
    )

    return lp_data


def marginal_dye_track(row_gain, theta, data):
    lp_data = _compute_lp_data(row_gain, theta, data)

    lp_marginal = forward_log_scan(theta.lp_init, theta.ltrans, lp_data)
    return jax.scipy.special.logsumexp(lp_marginal)


def _row_gain_from_raw(row_gain_raw, theta):
    scale = jnp.sqrt(jnp.log((1 + jnp.sqrt(4 * theta.row_gain_sigma ** 2 + 1)) / 2))
    return jnp.exp(scale * row_gain_raw)


def likelihood(row_gain_raw, theta, observed):
    row_gain = _row_gain_from_raw(row_gain_raw, theta)

    logp = marginal_dye_track(
        row_gain, theta, observed
    )
    return logp + jax.scipy.stats.norm.logpdf(row_gain_raw)


def marginal_likelihood(gauss_points, theta, observed):
    return log_integral_exp_adapt_jax(
        lambda x: likelihood(x, theta, observed),
        gauss_points,
    )


def marginal_likelihood_grad(gauss_points, theta, observed):
    val, grads = jax.value_and_grad(marginal_likelihood, argnums=1)(
        gauss_points,
        theta,
        observed,
    )
    return jax.lax.psum(val, "peak"), jax.lax.psum(grads, "peak")


def draw_posterior(key, gauss_points, theta, observed):
    key_curr, key = jax.random.split(key, 2)
    row_gain_raw = log_integral_exp_adapt_jax_sample(
        key_curr,
        lambda x: likelihood(x, theta, observed),
        gauss_points,
    )

    row_gain = _row_gain_from_raw(row_gain_raw, theta)
    lp_data = _compute_lp_data(row_gain, theta, observed)

    return row_gain, sample_path(key, theta.lp_init, theta.ltrans, lp_data)


# ZBS: Shall we just bump the jax version in pipfile? I get a warning abot xmap being experimenal
# Also, this needs some docs, I have no idea why this is being set up

# This requires a higher version of JAX than what is specified in environment.yml.
# We can swap out xmap for vmap if we're operating in a lower version of JAX.
with jax.experimental.maps.mesh(jax.local_devices(), ["peak"]):
    marginal_likelihood_grad_xmap = jax.experimental.maps.xmap(
        marginal_likelihood_grad,
        in_axes=[
            [...],
            [...],
            ["peak", ...],
        ],
        out_axes=[...],
        axis_resources={"peak": "peak"},
    )<jupyter_output><empty_output><jupyter_text>**ZBS: No idea why we need this hermitian.**<jupyter_code>gauss_points, _ = np.polynomial.hermite_e.hermegauss(13)<jupyter_output><empty_output><jupyter_text>**ZBS: Are the zthe HMM args?**<jupyter_code>row_k_sigma = 1.
lp_init = np.log(np.ones(5) / 5)
ltrans = np.log(np.ones((5, 5)) / 25)
mu_unscaled = np.arange(5) * 10_000.
sigma_unscaled = (np.arange(5) + 1) * 100.
lp_outlier = np.log(0.05)
outlier_mu = X_train.mean()
outlier_sigma = X_train.std()
observed = X_train
%%time
with jax.experimental.maps.mesh(jax.local_devices(), ["peak"]):
    theta = Theta(lp_init, ltrans, np.array(row_k_sigma), mu_unscaled, sigma_unscaled)
    #theta = Theta(lp_init, ltrans, np.array(row_k_sigma), mu_unscaled, sigma_unscaled, lp_outlier, outlier_mu, outlier_sigma)

    logp, grads = marginal_likelihood_grad_xmap(gauss_points, theta, observed)
    jax.device_get(grads)
    logp.block_until_ready()
    print(logp)
%%timeit
with jax.experimental.maps.mesh(jax.local_devices(), ["peak"]):
    #theta = (np.array(row_k_sigma), lp_init, ltrans, mu_unscaled, sigma_unscaled)
    theta = Theta(lp_init, ltrans, np.array(row_k_sigma), mu_unscaled, sigma_unscaled)
    logp, grads = marginal_likelihood_grad_xmap(gauss_points, theta, observed)
    jax.device_get(grads)
    logp.block_until_ready()<jupyter_output><empty_output><jupyter_text>Wrap jax code in op<jupyter_code>gauss_points, _ = np.polynomial.hermite_e.hermegauss(19)
gauss_points = gauss_points / 1.5
class MarginalHMMLh(aesara.graph.op.Op):
    itypes = [aet.dvector, aet.dmatrix, aet.dscalar, aet.dvector, aet.dvector, aet.dmatrix]
    otypes = [aet.dscalar] + itypes.copy()[:-1]

    def perform(self, node, inputs, outputs):
        use_xmap = True

        if use_xmap:
            *theta, observed = inputs
            val, grads = marginal_likelihood_grad_xmap(gauss_points, Theta(*theta), observed)
            outputs[0][0] = jax.device_get(val).astype('d')
            grads = jax.device_get(grads)
            for i in range(len(grads)):
                outputs[i + 1][0] = grads[i].astype('d')
        else:
            val, grads = marginal_likelihood_grad(gauss_points, *inputs)
            outputs[0][0] = val.to_py().astype('d')
            for i in range(len(grads)):
                outputs[i + 1][0] = grads[i].to_py().astype('d')

    def grad(self, inputs, dout):
        val, *grads = self(*inputs)

        if not all(isinstance(g.type, aet.DisconnectedType) for g in dout[1:]):
            raise ValueError("Undefined gradient.")

        return [
            grad * dout[0]
            for grad in grads
        ] + [aet.grad_not_implemented(self, 5, inputs)]

def marginal_hmm_lh(theta, observed):
    op = MarginalHMMLh()
    val, *grads = op(*theta, observed)
    return val<jupyter_output><empty_output><jupyter_text>PyMC3 modelHere's an implementation of the PyMC3 model,including marginalization over row_k as well as marginalization over dye counts.<jupyter_code>def generative_model(
    hyper_n_cycles: int,
    hyper_n_training_rows: int,
    hyper_n_count: int,
    X_train: np.ndarray,
):
    coords = {
        "peak": pd.RangeIndex(hyper_n_training_rows),
        "cycle": pd.RangeIndex(hyper_n_cycles),
        "state": pd.RangeIndex(hyper_n_count + 1),
        "state_to": pd.RangeIndex(hyper_n_count + 1),
    }
    with pm.Model(coords=coords) as model:
        gain_raw = pm.Gamma("gain_raw", **find_gamma_params(lower=1.5, upper=2.0))
        gain = pm.Deterministic("gain", 10_000 * gain_raw)

        sigma_raw = pm.Gamma("sigma_raw", **find_gamma_params(lower=5, upper=20))
        sigma = pm.Deterministic("sigma", 100 * sigma_raw)

        bg_sigma_raw = pm.Gamma("bg_sigma_raw", **find_gamma_params(lower=1, upper=10))
        bg_sigma = pm.Deterministic("bg_sigma", 100 * bg_sigma_raw)

        p_bleach = pm.Beta("p_bleach", **find_beta_params(lower=0.01, upper=0.20))
        p_attach = pm.Beta("p_attach", **find_beta_params(lower=0.5, upper=0.9))
        
        row_k_sigma = pm.Gamma(
            "row_k_sigma", **find_gamma_params(lower=0.01, upper=0.15)
        )
        
        state_base = aet.as_tensor_variable(np.arange(hyper_n_count + 1))

        mu_unscaled = state_base * gain
        sigma_unscaled = aet.sqrt(state_base * sigma ** 2 + bg_sigma ** 2)
        
        pm.Deterministic("mu_unscaled", mu_unscaled, dims="state")
        pm.Deterministic("sigma_unscaled", sigma_unscaled, dims="state")

        lp_init = aet.log(pm.Dirichlet("p_init", a=np.ones(hyper_n_count + 1, dtype='d'), dims="state"))
        #pm.Binomial.dist... TODO

        ltrans = pm.Binomial.dist(p=1 - p_bleach, n=np.arange(hyper_n_count + 1)[:, None]).logp(np.arange(hyper_n_count + 1)[None, :])

        pm.Deterministic("ltrans", ltrans, dims=("state", "state_to"))
        pm.Deterministic("lp_init", lp_init, dims="state")

        theta = Theta(lp_init, ltrans, row_k_sigma, mu_unscaled, sigma_unscaled)
        pm.Potential("logp_hmm", marginal_hmm_lh(theta, aet.as_tensor_variable(X_train)))

    return model
import arviz as az
import pandas as pd<jupyter_output><empty_output><jupyter_text>Here, we fit the model assuming we know the true dye tracks.<jupyter_code>with jax.experimental.maps.mesh(jax.local_devices(), ["peak"]):
    model = generative_model(
        hyper_n_cycles,
        hyper_n_training_rows,
        2,
        X_train,
    )
with model:
    with jax.experimental.maps.mesh(jax.local_devices(), ["peak"]):
        trace = pm.sample(return_inferencedata=False, cores=1, chains=2, discard_tuned_samples=False)
import arviz
idata_kd = arviz.from_pymc3(trace, model=model, save_warmup=True)
tr = idata_kd.posterior
arviz.plot_trace(tr);
tr_ = tr.stack(sample=('chain', 'draw'))#.sel(chain=0, draw=0)
key = jax.random.PRNGKey(10)
theta_draw = Theta(
    tr_.lp_init.values,
    tr_.ltrans.values,
    tr_.row_k_sigma.values,
    tr_.mu_unscaled.values,
    tr_.sigma_unscaled.values,
)
draw_func = jax.jit(
    jax.vmap(  # Map over data
        jax.vmap(  # Map over posterior draws
            draw_posterior,
            in_axes=[
                0,
                None,
                -1,
                None,
            ],
            out_axes=(0, 0)
        ),
        in_axes=[
            None,
            None,
            None,
            0,
        ],
        out_axes=(0, 0),
    )
)
keys = jax.random.split(key, len(tr_.sample))

draws_gain, draws_dye_track = draw_func(
    keys,
    jax.device_put(gauss_points),
    jax.device_put(theta_draw),
    jax.device_put(X_train),
)
tr_['row_gain'] = (("peak", "sample"), draws_gain)
#tr.to_zarr("/home/adr/data/erisyon-trace3.zarr")
tr_['dye_track'] = (("peak", "sample", "cycle"), draws_dye_track)
import xarray as xr
tr_.dye_track
sns.heatmap(tr_.dye_track.std('sample'))
is_correct = tr_.dye_track - xr.DataArray(dytmat, dims=("peak", "cycle")) == 0
is_correct.mean()
(~is_correct).sum("sample").values.nonzero()
sns.catplot(data=tr_.sigma_unscaled.to_dataframe().reset_index().astype(dict(state='category')), x='sigma_unscaled', y='state', kind='violin', cut=0, orient='h')
idx = 92
plt.plot(X_train[idx])
plt.plot(tr_.dye_track.isel(peak=idx).values[::1].T * tr_.gain.mean().values * tr_.row_gain.isel(peak=idx).mean().values)
plt.ylim(bottom=0)
sns.displot(tr_.row_gain.isel(peak=idx).unstack('sample').values.T)
import arviz
idata_kd.posterior.plot.scatter('sigma', 'bg_sigma')
import matplotlib.pyplot as plt
sns.histplot(idata_kd.posterior['bg_sigma'].values.ravel())
plt.axvline(true_bg_sigma)
sns.histplot(idata_kd.posterior['gain'].values.ravel())
plt.axvline(true_gain)
sns.histplot(idata_kd.posterior['sigma'].values.ravel())
plt.axvline(true_sigma)<jupyter_output><empty_output><jupyter_text>Sanity-check that we still recover parameters correctly  Row K Sigma<jupyter_code>sns.histplot(idata_kd.posterior['row_k_sigma'].values.ravel())

plt.axvline(true_row_k_sigma)
tr = idata_kd.posterior<jupyter_output><empty_output><jupyter_text>Background sigma<jupyter_code>sns.histplot(idata_kd.posterior['bg_sigma'].values.ravel())
plt.axvline(true_bg_sigma)<jupyter_output><empty_output><jupyter_text>Attachment probability<jupyter_code>sns.histplot(idata_kd.posterior['p_attach'].values.ravel())
plt.axvline(true_p_attach)<jupyter_output><empty_output><jupyter_text>Bleaching probability<jupyter_code>sns.histplot(idata_kd.posterior['p_bleach'].values.ravel())
plt.axvline(true_p_bleach)
sns.histplot(idata_kd.posterior['sigma'].values.ravel())
plt.axvline(true_sigma)
sns.heatmap(dytmat, center=0)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11072945122116139, 'fraction_numerical': 0.014902766883942336, 'mean_word_length': 3.725715746421268, 'pattern_counts': {'":': 5, '<': 65, '<?xml version=': 0, '>': 63, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '15', 'dir': 'jupyter-structured-clean-dedup', 'id': '590710', 'license': 'MIT', 'n_tokens_mistral': 13448, 'n_tokens_neox': 12635, 'n_words': 3532, 'path': 'plaster/explore/pymc3_poc1.ipynb', 'repo_name': 'pymc-labs/plaster', 'size': '36973'}