update entire initial_size matrix and/or migration rate matrix at once

[connor-french]

Hello all!

I'm running large 2D stepping stone models (deme arrays ~100 x 100 or larger) that change through time and have come across a time limiting step when setting up the Demography models. After initializing the Demography object with an initial_size matrix and directly assigning a matrix of migration rates to the the model.migration_matrix slot, I need to update the values at different time points in the past.

The current way I'm doing this is to use add_population_parameters_change() to update the initial_sizes and add_migration_rate_change() for the migration rates at discrete time points. However, this takes a very long time when updating many demes and migration rates (~2 minutes for a sparse initial_size array of size 110 x 117, and I stopped tracking time after 20 minutes for a fully filled initial_size array of the same size).

Would it be possible to implement a feature where add_population_parameters_change() and add_migration_rate_change() are able to have arrays as an input, so multiple populations and multiple migration rates can be updated at a specific time point? This could drastically speed up model setup for these unreasonably large models I and possibly others are trying to run.

Oh yeah, I meant to loop @petrelharp in on this

[jeromekelleher]

I'm not sure that we're going to speed things up much by doing that @connor-french, because we're still going to be storing separate objects representing the low-level events (e.g. here)

However, I'm sure there must be some low-hanging fruit here perf-wise. Perhaps you could construct a standalone example with similar properties that takes (say) 10 seconds to run . If you post the code (and ideally a timed profile) we can see how to start speeding things up. There's no consideration at all given to performance at the moment, so I'm sure we can improve things.

[connor-french]

Ah, that makes sense.

Here is a standalone example that takes ~15 seconds to run, along with a timed profile generated by cProfile. I have a screenshot of my snakeviz visualization below for quick reference.

It looks like __getitem__ from demography.py is what is slowing the script down. Although a single run of the function isn't so slow, it's called many times (266371 times!) because I'm calling add_migration_rate_change so many times (118020) and having to check the population identifiers with _check_population_references. Maybe there's a way to optimize the __getitem__ function? My gut tells me this for loop may be slowing the function down, but I'm not 100%:

if isinstance(identifier, str):
            for population in self.populations:
                if population.name == identifier:
                    return population

[petrelharp]

To add a bit more information here: I think this is the use case where there's a fixed sparse pattern of nonzero entries in the migration matrix, and we want to update many of those nonzero entries often.

[jeromekelleher]

This is very helpful, thanks @connor-french!

I think the most useful option for you (which would allow us to optimise quite a bit) would be to make a new method add_migration_matrix_change. You could feed in the migration matrix, then the lower-level code would be able to optimise out all the population reference checking, etc. It may not be worth worrying about which particular rates change, and just feeding the whole thing through.

Ideally, this would be passed directly through to the simulated, but this would require some more low-level plumbing.

In general, there's a bunch of things we could do to make this application of a large 2D model that changes over time work better - I'm happy to work with you to get those changes into msprime, if you're interested!

[petrelharp]

It may not be worth worrying about which particular rates change, and just feeding the whole thing through.

I think that in cases where it's worth having this method (instead of just updating a bunch of rates individually), the migration matrix will be very sparse. Note that in this case the migration matrix is 1e4 x 1e4, but only ~4e4 entries are nonzero. So - you may be right that this is the way to start, but I'm trying to avoid putting in a bunch of work that doesn't solve the issue.

[connor-french]

Agreed that it might still be worth checking which particular rates change, given how sparse the matrices typically are. I can think of a simple way to obtain the indices for rates that need to be updated, and update the migration matrix based on these indices. Maybe the original donor/recipient populations could be backtracked from these indices?

# migration matrix for t1
m1 = np.random.randint(0, 100, (10, 10))

# migration matrix for t2
m2 = np.random.randint(0, 100, (10, 10))

# subtract the two matrices
mdiff = np.subtract(m2, m1)

# if the difference is not zero, these need to be updated, so get their indices
np.transpose(np.nonzero(mdiff !=0))

I'm definitely down to work on making these large 2D stepping stone models work better!