Stepping Stone model - Torus

[JaneeshBansal]

Hi,

I am trying to model a stepping stone model in a torus shape. I saw #2100 issue where a 2D stepping stone model was discussed but I am confused about how to add boundary conditions.

I tried to add in boundary conditions below but I don't think they are correct. I don't undestand parts of the script - why is the remainder being taken in this part right = (index + 1) % N?

import msprime
import numpy as np

def stepping_stone2d(initial_size, rate, boundaries=True):
    # check if 2D
    assert len(initial_size.shape) == 2
    # dim of lattice
    n, m = initial_size.shape

    N = n * m
    # linear?
    model = msprime.Demography.isolated_model(initial_size.reshape(N))
    M = model.migration_matrix
        
    for j in range(n):
        for k in range(m):
            index = j * m + k
            model.populations[index].name = f"pop_{j}_{k}"
            
            # define the migration rates for the current index
            left = index - 1
            # % returns remainder - d=9 would be 0
            right = (index + 1) % N
            up = index - m
            down = (index + m) % N

            # if bondary condition enabled
            if boundaries:
                if left < 0: # Periodic boundary condition for left edge
                    left = n
                if right >= n: # Periodic boundary condition for right edge
                    right = 0
                if up < 0: # Periodic boundary condition for top edge
                    up = m
                if down >= m: # Periodic boundary condition for bottom edge
                    down = 0

            # Update migration rates
            M[index, left] = rate
            M[index, right] = rate
            M[index, up] = rate
            M[index, down] = rate

            M[left, index] = rate
            M[right, index] = rate
            M[up, index] = rate
            M[down, index] = rate

    return model

Thank you :)
Janeesh

[jeromekelleher]

Hi @JaneeshBansal 👋

The "mod" conditions here usually used to implement the "wraparound" that you get for a torus. So, if we were at the rightmost k = m - 1, then the next deme = (k + 1) % m is 0.

Printing out the migration matrix for a small (say 3x3) example would probably help?

[JaneeshBansal]

Hi @jeromekelleher ,

Thank you for the reply :)

So the script below already deals with the edges (I added more variables to make more sense for me).

As if left = index - 1 or up = index - m turns out to be -1 then it loops around to the end. And the right = (index + 1) % N or down = (index + m) % N is 0 it would loop it to the start.

import msprime
import numpy as np

def stepping_stone2d(initial_size, rate):
    # check if 2D
    assert len(initial_size.shape) == 2
    # shape of lattice
    n, m = initial_size.shape
    # total number of demes
    N = n * m
    
    # base model on isolated pops
    model = msprime.Demography.isolated_model(initial_size.reshape(N))
    # migration matrix 
    M = model.migration_matrix
    
    # loop over 2d lattice
    for j in range(n):
        for k in range(m):
            # different states
            index = j * m + k
            model.populations[index].name = f"pop_{j}_{k}"
            
            # movements
            left = index - 1
            right = (index + 1) % N
            up = index - m
            down = (index + m) % N
            
            M[index, left] = rate
            M[index, right] = rate
            M[index, up] = rate
            M[index, down] = rate

            M[left, index] = rate
            M[right, index] = rate
            M[up, index] = rate
            M[down, index] = rate
            
    return model

Janeesh

[jeromekelleher]

Ah right, I see - why are we modding by N, rather than m? Should it be something like this?

def transform(row, col):
     return row * m + col
focal = transform(j, k)
left = transform(j, (k - 1) % m)
right = transform(j, k + 1) % m)
up = transform((j - 1) % n, k)
down = transform((j + 1) % n, k)
for coord in [left, right, up, down]:
    M[focal, coord] = rate
    M[coord, focal] = rate

I haven't tested this at all though. Maybe you could print out a sample migration matrix to so we can see what it looks like?

[JaneeshBansal]

From this script ...

import msprime
import numpy as np

def stepping_stone2d(initial_size, rate):
    # check if 2D
    assert len(initial_size.shape) == 2
    # shape of lattice
    n, m = initial_size.shape
    # total number of demes
    N = n * m
    
    # base model on isolated pops
    model = msprime.Demography.isolated_model(initial_size.reshape(N))
    # migration matrix 
    M = model.migration_matrix
    
    # loop over 2d lattice
    for j in range(n):
        for k in range(m):
            # different states
            index = j * m + k
            model.populations[index].name = f"pop_{j}_{k}"
            
            def transform(row, col):
                 return row * m + col
            
            focal = transform(j, k)
            left = transform(j, (k - 1) % m)
            right = transform(j, (k + 1) % m)
            up = transform((j - 1) % n, k)
            down = transform((j + 1) % n, k)
            
            for coord in [left, right, up, down]:
                M[focal, coord] = rate
                M[coord, focal] = rate

    return model


model = stepping_stone2d(np.zeros((3, 3)) + 5, 1)

print(model)

I got this output ...

Demography
╟  Populations
║  ┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
║  │ id │name     │description  │initial_size  │ growth_rate │  default_sampling_time│extra_metadata  │
║  ├──────────────────────────────────────────────────────────────────────────────────────────────────┤
║  │ 0  │pop_0_0  │             │5.0           │      0      │                      0│{}              │
║  │ 1  │pop_0_1  │             │5.0           │      0      │                      0│{}              │
║  │ 2  │pop_0_2  │             │5.0           │      0      │                      0│{}              │
║  │ 3  │pop_1_0  │             │5.0           │      0      │                      0│{}              │
║  │ 4  │pop_1_1  │             │5.0           │      0      │                      0│{}              │
║  │ 5  │pop_1_2  │             │5.0           │      0      │                      0│{}              │
║  │ 6  │pop_2_0  │             │5.0           │      0      │                      0│{}              │
║  │ 7  │pop_2_1  │             │5.0           │      0      │                      0│{}              │
║  │ 8  │pop_2_2  │             │5.0           │      0      │                      0│{}              │
║  └──────────────────────────────────────────────────────────────────────────────────────────────────┘
╟  Migration Matrix
║  ┌───────────────────────────────────────────────────────────────────────────────────────────────────┐
║  │         │ pop_0_0 │ pop_0_1 │ pop_0_2 │ pop_1_0 │ pop_1_1 │ pop_1_2 │ pop_2_0 │ pop_2_1 │ pop_2_2 │
║  ├───────────────────────────────────────────────────────────────────────────────────────────────────┤
║  │  pop_0_0│    0    │    1    │    1    │    1    │    0    │    0    │    1    │    0    │    0    │
║  │  pop_0_1│    1    │    0    │    1    │    0    │    1    │    0    │    0    │    1    │    0    │
║  │  pop_0_2│    1    │    1    │    0    │    0    │    0    │    1    │    0    │    0    │    1    │
║  │  pop_1_0│    1    │    0    │    0    │    0    │    1    │    1    │    1    │    0    │    0    │
║  │  pop_1_1│    0    │    1    │    0    │    1    │    0    │    1    │    0    │    1    │    0    │
║  │  pop_1_2│    0    │    0    │    1    │    1    │    1    │    0    │    0    │    0    │    1    │
║  │  pop_2_0│    1    │    0    │    0    │    1    │    0    │    0    │    0    │    1    │    1    │
║  │  pop_2_1│    0    │    1    │    0    │    0    │    1    │    0    │    1    │    0    │    1    │
║  │  pop_2_2│    0    │    0    │    1    │    0    │    0    │    1    │    1    │    1    │    0    │
║  └───────────────────────────────────────────────────────────────────────────────────────────────────┘
╟  Events
║  ┌───────────────────────────────────┐
║  │  time│type  │parameters  │effect  │
║  ├───────────────────────────────────┤
║  └───────────────────────────────────┘

So I think this works? for pop_0_0 there is a value to pop_0_2 and for pop_2_0. So there is migration from edge to edge?

Janeesh

[jeromekelleher]

I think so? I'd have to draw it out to be sure...

[JaneeshBansal]

From drawings it seems to work!

Thank you :)