Add WeightedIndex::update_weights_relative

[MichaelOwenDyer]

• Added a CHANGELOG.md entry

Summary

This PR adds a new method update_weights_relative to the WeightedIndex distribution, whose purpose is to update a subset of weights relative to the current values, following a provided mapping function.

Motivation

Sometimes it is desirable to update a weight in a distribution relative to its current value, for instance subtracting one from the weight of a sampled index in order to achieve drawing without replacement. However, the update_weights method requires you to provide the new weights directly, and there is no external way to check the current weight at an index. This means that you have to track the current weights manually:

let ordinal = thread_rng().sample(&weighted_index);
remaining[ordinal] -= 1; // Lower the weight of the drawn ordinal by one
let new_weight = remaining[ordinal];
if weighted_index.update_weights(&[(ordinal, &new_weight)]).is_err() {
    // all weights are zero
}

In this code snippet, remaining has the type [X; len], where X is the element type of the WeightedIndex and len is the number of weights in the index.
Keeping this array around externally is tedious, and wastes space, since all the information we need is technically already stored within the index.
It would be nice to instead be able to do something like this, to lower the weight at ordinal by 1:

let ordinal = thread_rng().sample(&weighted_index);
if weighted_index.update_weights_relative(&[(ordinal, &-1)]).is_err() {
    // all weights are zero
}

But this only works with signed types - unsigned types would only ever be able to be increased. Furthermore, it seems like an opportunity here to allow not just for addition and subtraction, but for arbitrary mappings via a closure:

let ordinal = thread_rng().sample(&weighted_index);
if weighted_index.update_weights_relative(&[(ordinal, &1)], |current, relative| current - relative).is_err() {
    // all weights are zero
}

Here, we provide a closure describing in what way the provided values should be combined with the current weights: the relative value should be subtracted from the existing weight to yield the new weight.

This yields another potential issue: what if the arithmetic in this closure fails (overflow, underflow, etc)? We could have the closure return an Option<X> to account for failures. Then we can do something like this:

let ordinal = thread_rng().sample(&weighted_index);
if weighted_index.update_weights_relative(&[(ordinal, &1)], |current, relative| current.checked_sub(relative)).is_err() {
    // all weights are zero
}

...and the whole method will return a WeightError if any of these subtractions underflows.

Details

The signature of the new update_weights_relative method is as follows:

pub fn update_weights_relative<R>(&mut self, relative_weights: &[(usize, &R)], op: impl Fn(X, &R) -> Option<X>) -> Result<(), WeightError>
    where
        X: for<'a> ::core::ops::AddAssign<&'a X>
            + for<'a> ::core::ops::SubAssign<&'a X>
            + Clone
            + Default
{ ...

The generic type parameter R allows the provided values to not necessarily be the same type as the values in the WeightedIndex, as long as the provided op can combine the two into an X.

In the error checking phase, update_weights_relative follows the same logic as update_weights, with a couple differences:

• the provided weights are not required to be >= zero like in update_weights (they do not even necessarily share zero's type), however op(old_weight, relative), the calculated new weight, must be >= zero. If the calculated weight is None or < zero, WeightError::InvalidWeight is returned.
• the new absolute weights are calculated in the initial check for errors. In order to avoid doing this a second time, a Vec is allocated to store these new weights until they can be applied. I would like to avoid this allocation but I don't see a way to do so -> suggestions welcome.

After performing the necessary checks (and calculating the new absolute weights), the actual application of the new weights is identical to the update_weights method. Therefore, this section of the update_weights method has been factored out, and is now called by both methods.

Positive and negative test cases have been added for the new method (and for the existing update_weights method).

This is my first contribution to this crate, I appreciate your feedback!

[dhardy]

the new absolute weights are calculated in the initial check for errors. In order to avoid doing this a second time, a Vec is allocated to store these new weights until they can be applied. I would like to avoid this allocation but I don't see a way to do so -> suggestions welcome.

I can see two possible ways around this:

• Pass in values as an array of [(usize, X)] by value or mutable reference. This is somewhat restrictive, so may be undesirable.
• Declare that in the case of error, the state of self is undefined. In case the caller still needs the WeightedIndex state after an error, they must copy the entire distribution (Vec<X> plus a couple of constant-size parameters).

It comes down to how update_weights_relative will get used most.

---

The old update_weights is now a special case of update_weights_relative. Can it be implemented in terms of the latter without significant run-time cost?

[MichaelOwenDyer]

Hey @dhardy, thanks for the input. I prefer the first of the two options you proposed; I think the second sacrifices too much in the way of error handling. Since taking ownership of the passed weights is not really necessary I would aim for borrowing mutably; it might be interesting to the caller to see what the weights have been changed to after the fact.

I think update_weights could be implemented in terms of update_weights_relative where op simply returns the new weight irrespective of the old one. I'll experiment with that and report back :)

[MichaelOwenDyer]

I think update_weights could be implemented in terms of update_weights_relative where op simply returns the new weight irrespective of the old one.

This would only work if we also changed the signature of update_weights to take ownership or mutable reference of the new weights, which would be breaking. I assume that is not acceptable.

We could also discuss alternative ways to solve this problem, as the update_weights_relative method may not be the best one. The root of the problem is really that there is no way to look into the distribution from outside and see the current weights upon which to calculate the desired new weights. If we could expose the distribution weights in an ergonomic way (I don't think the cumulative form is particularly ergonomic for the outside user) then I would consider it good enough to cover my use case.

[dhardy]

I think the second sacrifices too much in the way of error handling.

But does this really matter? In many cases, errors are not designed to be handled, but just reported.

---

Another half-baked idea would be to add the following:

• fn iter_weights(&self) -> impl Iterator<Item = &X>
• fn visit_weights_mut(&self, visitor: impl FnMut(usize, &mut X)) -> Result<(), ErrorType>

The latter would make a copy of each weight before calling visitor (either on the copy or the original), updating the total weight as it goes. But this also does not allow for recovering from errors.

[MichaelOwenDyer]

After working on an implementation of update_weights_relative which accepts a &[(usize, &mut X)] as a parameter, I realized that forcing the caller to retain ownership of all the Xs while calling this method is actually a rather large burden, because inlining the values inside the method arguments becomes impossible. So I don't think that is the way to go anymore.

I'm starting to think it would be easier to just expose a method like the following and let the caller figure out the new weights themselves:

pub fn weight_at(&self, index: usize) -> X
where
    X: for<'a> ::core::ops::SubAssign<&'a X>
        + Clone
{
    let mut weight = if index < self.cumulative_weights.len() {
        self.cumulative_weights[index].clone()
    } else {
        self.total_weight.clone()
    };
    if index > 0 {
        weight -= &self.cumulative_weights[index - 1];
    }
    weight
}

This is the same as how the old weights are currently calculated inside the update_weights method, and uses the same trait bounds for consistency. We probably should return an Option for index-safety, and a way to get all weights at once would be useful too.

While we're at it, maybe exposing more of the inner workings of the distribution (immutably) would offer additional convenience at no extra risk, just in case it is useful:

pub fn total_weight(&self) -> &X {
    &self.total_weight
}

pub fn cumulative_weights(&self) -> &Vec<X> {
    &self.cumulative_weights
}

@dhardy what do you think about this? Is there reason not to do this?

[dhardy]

Reporting more inner details makes it harder to change the implementation later, but I don't see the implementation of WegihtedIndex changing in incompatible ways anyway, so this may be a good approach.

cumulative_weights should return &[X] not &Vec<X> since the latter is more restrictive and offers no advantage.

weight_at should return Option<X> unless there is a good reason to do otherwise.

[MichaelOwenDyer]

Closing this as #1420 was merged.