greta-dev / greta.dynamics

a greta extension for modelling dynamical systems
https://greta-dev.github.io/greta.dynamics/
Other
6 stars 2 forks source link

iterate_matrix under/overflow #16

Closed goldingn closed 5 years ago

goldingn commented 5 years ago

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))