greta-dev / greta.dynamics Goto Github PK

When running iterate matrix for long enough with low enough tolerance (or large enough number of replicates), numerical underflow or overflow can lead the states to become either 0 or Inf, and therefore result in undefined growth rates and stable states (because 0/0 and Inf/Inf are both not defined, so return NaN). E.g.:

mat <- matrix(0, 2, 2)
mat[1, 2] <- 0.9
states <- iterate_matrix(mat, niter = 3, tol = -Inf)
calculate(states$lambda)

     [,1]
[1,]  NaN

calculate(states$stable_distribution)

     [,1]
[1,]  NaN
[2,]  NaN

To deal with this, we could modify the TF while loop to check at each iteration whether the new state is finite for each replicate, and only update the growth rates and states if so. The convergence trace could also be expanded to all replicates, and set based on this.

I.e. in the body:

valid <- is_valid(new_state)
# update the stored growth rates and states if the new state is finite
growth_rate <- set_defined(growth_rate, new_state / old_state, valid)
new_state <- set_defined(old_state, new_state, valid)
# if not, determine that this has converged
converged <- growth_converged(growth_rates, tf_tol)
converged <- set_defined(converged, ones, !valid)

where:

is_valid <- function (x) {
  tf$logical_and(is_non_zero(x), is_finite(x))
}
is_non_zero <- function (x) {
  zero <- tf$zeros(shape(), dtype = tf_float())
  tf$logical_not(tf$equal(x, zero))
}

This would require changing growth_converged to return a vector, passing in a vector of convergences to start with, and changing the condition function to:

tf$squeeze(tf$less(iter, maxiter)) & tf$logical_not(tf$reduce_all(converged))

I think there is a mismatch in the stable states returned by the r_iterate_matrix and iterate_matrix functions if the loop runs to the final iteration. Not sure if current tests are ending the loop before niter, hence not erroring.

The (possible) issue arises when the while loop never breaks due to diff > tol, in which case the final iteration is (i < niter => i = niter - 1). This iteration increments i by 1, then stores the output to states[[i+1]]. So the final element of states is at niter + 1.

But then the stable state is returned as states[[i]], which is states[[niter]] in this case:

Is this intentional? If not, I can prepare a PR. The same issue arises on the iterate_dynamic_matrix branch.

As mentioned in #11, we now use a TF while loop and can stop iterations early if we meet a convergence tolerance. We could also speed up convergence (and therefore reduce computation time) by warmn-starting the initial values between iterations (if the user didn't specify any). That is, the next time the iteration algorithm is run we can set the initial state to the final state from the previous run. If doing MCMC sampling, subsequent parameter values (and therefore final states) should be similar, so this could speed up sampling quite a lot.

I think we could do this with a TF variable (or constant?) tensor, setting values of the initial state tensor at the end of each run of the algorithm. Not something I've done with greta before though, so will take some thinking.

The greta.dynamics package's ode_solve() functionality also depends on contrib, using the tf$contrib$integrate$odeint and odeint_fixed methods. tfp$math$ode$Solver$solve may be a viable replacement.

Dear maintainer,

You have file 'greta.dynamics/man/greta.dynamics.Rd' with \docType{package}, likely
intended as a package overview help file, but without the appropriate
PKGNAME-package \alias as per "Documenting packages" in R-exts.

This seems to be the consequence of the breaking change

Using @doctype package no longer automatically adds a -package alias.
Instead document _PACKAGE to get all the defaults for package
documentation.

in roxygen2 7.0.0 (2019-11-12) having gone unnoticed, see
r-lib/roxygen2#1491.

As explained in the issue, to get the desired PKGNAME-package \alias
back, you should either change to the new approach and document the new
special sentinel

  "_PACKAGE"

or manually add

  @aliases greta.dynamics-package

if remaining with the old approach.

The current release of greta.dynamics does not work with TF2 because of it relies on contrib for ODE solvers (#25).

Once that is fixed, another issue will appear, which is that greta:as_tf_function() (which is used internally in greta.dynamics functions) is failing to correctly handle definition of data greta array nodes inside the function. This is causing issues in experimental branches of greta.dynamics (e.g. greta_2) which use custom functions (but not ODE solvers).

The ODE example code in the current release is one example of one which will fail, because it includes the term (1 - Prey / K), where the 1 will be turned into a data greta array inside the function. When doing MCMC, this will yield an error message returned will be related to tensors changing their shape inside the loops used to solve dynamics.

Here's a reprex using the greta_2 branch and its iterate_dynamic_function() function, using TF2:

library(greta.dynamics)
#> Loading required package: greta
#> 
#> Attaching package: 'greta'
#> The following objects are masked from 'package:stats':
#> 
#>     binomial, cov2cor, poisson
#> The following objects are masked from 'package:base':
#> 
#>     %*%, apply, backsolve, beta, chol2inv, colMeans, colSums, diag,
#>     eigen, forwardsolve, gamma, identity, rowMeans, rowSums, sweep,
#>     tapply

oneplus <- function(state, iter, one = 1) {
  state + one
}

# this fails, as 'one' is defined inside the function
x <- normal(0, 1)
#> ℹ Initialising python and checking dependencies, this may take a moment.
#> ✔ Initialising python and checking dependencies ... done!
#> 
res <- iterate_dynamic_function(transition_function = oneplus,
                                initial_state = x,
                                tol = 0, 
                                niter = 3)
m <- model(res$all_states)
draws <- mcmc(m)
#> Error in eval(expr, envir, enclos): RuntimeError: ValueError: Input tensor `chain_of_reshape_of_identity/forward_log_det_jacobian/reshape/forward/Reshape:0` enters the loop with shape (1, 1, 1), but has shape (None, 1, 1) after one iteration. To allow the shape to vary across iterations, use the `shape_invariants` argument of tf.while_loop to specify a less-specific shape.


# this works, as 'one' is defined outside the function and passed in
x <- normal(0, 1)
res <- iterate_dynamic_function(transition_function = oneplus,
                                initial_state = x,
                                tol = 0,
                                one = ones(1),
                                niter = 3)
m <- model(res$all_states)
draws <- mcmc(m)
#> running 4 chains simultaneously on up to 8 CPU cores

^{Created on 2024-03-13 with reprex v2.0.2}

Note I did a slightly weird delayed execution thing here with a default value for the one argument so I could use the same function in both examples, but it's the same thing that happens with state + 1.

This code was used in TF1 greta to overcome the issue, by ensuring that the tensors for data greta arrays were defined as constants (no batch dimensions) rather than placeholders. It seems to still be in place though, so I'll try to debug now.

Possibly (probably) a "me" issue but couldn't work it out. I've had this same issue in the past but forgot my fix. . .

edit: error is on dev branch

only on greta_2 branch for now

First release:

• usethis::use_cran_comments()
• Update (aspirational) install instructions in README
• Proofread Title: and Description:
• Check that all exported functions have @return and @examples
• Check that Authors@R: includes a copyright holder (role 'cph')
• Check licensing of included files
• Review https://github.com/DavisVaughan/extrachecks

Prepare for release:

• git pull
• Check if any deprecation processes should be advanced, as described in Gradual deprecation
• urlchecker::url_check()
• devtools::check(remote = TRUE, manual = TRUE)
• devtools::check_win_devel()
• rhub::check_for_cran()
• git push
• Draft blog post

Submit to CRAN:

• usethis::use_version('minor')
• devtools::submit_cran()
• Approve email

Wait for CRAN...

• Accepted 🎉
• git push
• usethis::use_github_release()
• usethis::use_dev_version()
• usethis::use_news_md()
• git push
• Finish blog post
• Tweet
• Add link to blog post in pkgdown news menu

Background

Gradient-based inference (like the HMC greta uses) can only operate on continuous parameter spaces. That means it cannot learn the values of parameters with discrete support (e.g. no unobserved Poisson random variables).

But demographic stochasticity due to discrete stochastic variation in population sizes between timesteps in models of populations and infections (e.g. the number of new infectees is Poisson, the number of individuals surviving is binomial) is often important, especially when populations reach low numbers and near extinction. We cannot directly model the values of these discrete random variables, but we can apply continuous relaxations to approximate them; keeping the state values continuous and replacing the discrete random variables with continuous random variables that matches the mean and variance (and ideally the full shape of the distribution) of the random variable we would like to model as stochastic.

Continuous approximations to distributions

E.g. a poisson random variable can be approximated with an appropriately-shaped gamma distribution that exactly matches the PMF at discrete values, or by some other distribution that is a a close-enough approximation:

We might write a discrete stochastic growth-rate model like this:

1. $$x_t \sim Poisson(x_{t-1} \times r)$$

where $x$ takes integer values, $Poisson(\lambda)$ is the Poisson distribution, and $r$ is a positive-valued growth rate parameter. To estimate the posterior over the values represented by $x$ in this model, but using HMC, we could instead fit the model:

2. $$y_t \sim \pi(y_{t-1} \times r)$$

where $y_t$ is a (strictly positive) real-valued parameter, and $\pi(\lambda)$ is some probability distribution with support on positive real values that has similar moments (men, variance, skewness, etc.) to $Poisson(\lambda)$.

Reparameterisation

If we structure these probability distributions such that they can be reparameterised in terms of the parameter and some latent 'innovation' or noise, with known distribution, we can significantly improve the ability to sample these models, since we can decorrelate the posterior distribution in a similar way to the reparameterisation trick for hierarchical models. This can also provide some computational advantages in greta/tensorflow by working with arrays rather than scalars.

If we know the quantile function of the continuous distribution (e.g. $q_{\pi}(p, \lambda)$ as the quantile function of $\pi(\lambda)$ , with $p$ the probability argument), we can reparameterise the innovations as $u \sim U(0, 1)$, and then plug them into the quantile function to sample poisson values. Ie. equation 2 above is equivalent to:

The vector of $u$ values can then be computed in advance, and passed into the solvers to be chopped up appropriately. The dependency structure in the model means $y_t$ depends on $y_{t-1}$, and so they are a priori (and therefore also a posteriori) correlated. But in this reparameterised formulation, HMC operates instead on $u$, and $u_t$ doesn't depend on $u_{t-1}$ so they are a priori uncorrelated, which removes a lot of correlation in the posterior and makes sampling much easier.

Note that if the quantile function is expensive to compute, other approximations and reparameterisations may be more appealing. E.g. for the Poisson, the inverse of the incomplete gamma function gives the quantile of the gamma distribution whose PDF matches the Poisson PMF at discrete values, but the function has no analytic form and is expensive to compute. A lognormal approximation (with either uniform or standard normal innovations) is imperfect but much more computationally efficient.

Currently this reparameterisation trick is only applicable in the greta_2 branch with iterate_dynamic_function(), since it requires functionality for indexing time-varying parameters.

Implementation

We just need to provide functions for the relaxations, documentation, and examples of applying this approach. I have some, I just need to add them to the greta_2 branch and work out the neatest user interface.

The examples use a matrix with columns moving to rows (matrix(A) * vec(x)) but tf_iterate_matrix calculates rows to columns (rowvec(x) * matrix(A)):

Not sure which format you prefer??