Fix average_size calculation

[Zac-HD]

Closes #3143.

[jebob]

I suspect a gradient descent method might run faster and be more reliable, this approach can choke for relatively small ints (if target=10000, max_size=10001 you can get float overflows when you compute the average for large inc.). I can't test the validity of this Python locally currently but testing nasty edge cases in excel shows that this can handle up to 2000 in a few iterations.

# Gradient derivation
average_size = (1.0 / (1 - p_continue) - 1) * (1 - p_continue ** max_size)
daverage/dp = (1 - p_continue ** max_size) * d/dp((1 - p_continue)**-1 - 1) + (1.0 / (1 - p_continue) - 1) * d/dp(1 - p_continue ** max_size)
daverage/dp = (1 - p_continue ** max_size) * (1 - p_continue)**-2 - (1.0 / (1 - p_continue) - 1) * max_size * p_continue ** (max_size-1)

	@staticmethod
	def _calc_p_continue(desired_avg, max_size):
		p_continue = 1 - 1.0 / (1 + desired_avg)
		if p_continue == 0 or max_size == float("inf"):
			return p_continue
		# For small max_size, the infinite-series p_continue is a poor approximation,
		# and while we can't solve the polynomial a few rounds of iteration quickly
		# gets us a good approximate solution in almost all cases (sometimes exact!).
		err = desired_avg - _p_continue_to_avg(p_continue, max_size)
		for _ in range(10):
			# Should converge in <5 iterations for nearly all cases
			if abs(err) < desired_avg * 0.001:
				# Good enough
				break
			gradient = _p_continue_to_gradient(p_continue, max_size)
			p_continue += err / gradient
			err = desired_avg - _p_continue_to_avg(p_continue, max_size)
		else:
			assert False, "Infinite loop error"
		assert 0 < p_continue < 1, p_continue
		return p_continue

def _p_continue_to_avg(p_continue, max_size):
    """Return the average_size generated by this p_continue and max_size."""
    return (1.0 / (1 - p_continue) - 1) * (1 - p_continue ** max_size)
	
def _p_continue_to_gradient(p_continue, max_size):
    """Return the gradient (with respect to p_continue) for p_continue and max_size."""
	# While the true gradient is well behaved around 1, this function is not.
	# Approximating to near-one is good enough in most cases
	# Todo: if the true value of p_continue is > this number we won't converge
	p_continue = min(p_continue, 0.9999)
	
    return (
	    (1 - p_continue ** max_size) / (1 - p_continue) ** 2 
	    - (1.0 / (1 - p_continue) - 1) * max_size * p_continue ** (max_size-1)
	)

[Zac-HD]

I've replaced my previous heuristic approach with a binary search, which is simple enough that it's clear-on-inspection that it converges and maintains the bounds we want at each step. The constant factors might be a bit larger, but it's still fast in practice and @lru_cache makes it faster again.

@jebob, I tried out your gradient-descent code, but couldn't get it numerically-stable enough to respect strict bounds when testing. Regardless, I feel that we've co-written this fix and have given you commit and coauthor credit accordingly.

[jebob]

@Zac-HD Thanks for the co-authorship!

Stability is essential so I agree binary search is the way to go.

[jebob]

LGTM other than the two comments

[jebob]

In this case I think we can just return p_continue (or do nothing, thus skipping the zeroth iteration of the while loop)? If 1 - 1.0 / (1 + desired_avg) is close enough to the correct answer that floating-point error causes p_continue to be an overestimate then we're close enough for practical purposes.

[Zac-HD]

I agree that the estimate is close enough that we could just use the initial p_continue.

However, it's cheap to make this adjustment, and that allows us to test that our binary search always maintains a strict upper bound.

[jebob]

Can we replace this special case with an @example on test_p_continue_to_average?

[Zac-HD]

We do actually have such an example, but also need this test to reliably get full coverage from our conjecture-coverage task (the other is in nocover).

[Zac-HD]

Thanks again for the review, and all your help with this!