Parallel Execution¶

from parallel_mpi import ParallelMPI  # import our model definition  # ty:ignore[unresolved-import]

import neworder

# neworder.verbose()
# neworder.checked(False)

# must be MPI enabled
assert neworder.mpi.SIZE > 1, "This configuration requires MPI with >1 process"

population_size = 100
p = 0.01
timeline = neworder.LinearTimeline(0, 10, 10)
model = ParallelMPI(timeline, p, population_size)
neworder.run(model)

import numpy as np
import pandas as pd

import neworder


class ParallelMPI(neworder.Model):
    def __init__(self, timeline: neworder.Timeline, p: float, n: int):
        """This model uses MPI to run parallel process that synchronise timesteps and exchange data with each other"""
        # initialise base model (essential!)
        super().__init__(timeline, neworder.MonteCarlo.nondeterministic_stream)

        # enumerate possible states
        self.s = np.arange(neworder.mpi.SIZE)

        # create transition matrix with all off-diagonal probabilities equal to p
        self.p = np.identity(neworder.mpi.SIZE) * (1 - neworder.mpi.SIZE * p) + p

        # record initial population size
        self.n = n

        # individuals get a unique id and their initial state is the MPI rank
        self.pop = pd.DataFrame({"id": neworder.df.unique_index(n), "state": np.full(n, neworder.mpi.RANK)}).set_index(
            "id"
        )

    def step(self) -> None:
        # generate some movement
        neworder.df.transition(self, self.s, self.p, self.pop, "state")

        # send emigrants to other processes
        for s in range(neworder.mpi.SIZE):
            if s != neworder.mpi.RANK:
                emigrants = self.pop[self.pop.state == s]
                neworder.log(f"sending {len(emigrants)} emigrants to {s}")
                neworder.mpi.COMM.send(emigrants, dest=s)

        # remove the emigrants
        self.pop = self.pop[self.pop.state == neworder.mpi.RANK]

        # receive immigrants
        for s in range(neworder.mpi.SIZE):
            if s != neworder.mpi.RANK:
                immigrants = neworder.mpi.COMM.recv(source=s)
                if len(immigrants):
                    neworder.log(f"received {len(immigrants)} immigrants from {s}")
                    self.pop = pd.concat((self.pop, immigrants))

    def check(self) -> bool:
        # Ensure we haven't lost (or gained) anybody
        totals = neworder.mpi.COMM.gather(len(self.pop), root=0)
        if totals and sum(totals) != self.n * neworder.mpi.SIZE:
            return False
        # And check each process only has individuals that it should have
        out_of_place = neworder.mpi.COMM.gather(len(self.pop[self.pop.state != neworder.mpi.RANK]))
        return not (out_of_place and any(out_of_place))

    def finalise(self) -> None:
        # process 0 assembles all the data and prints a summary
        pops = neworder.mpi.COMM.gather(self.pop, root=0)
        if pops:
            pop = pd.concat(pops)
            neworder.log(f"State counts (total {len(pop)}):\n{pop['state'].value_counts().to_string()}")

import time
from threading import Barrier, Lock

import pandas as pd

import neworder


class ParallelThreaded(neworder.Model):
    """
    A simple model demonstrating multithreaded rather than multiprocess execution, sharing a dataset and synchronising
    at each timestep. An expensive computation is simulated that is O(n), showing how partitioning the dataset can
    achieve significant performance gains.
    """

    def __init__(self, pop: pd.DataFrame, indices: list[int], lock: Lock, barrier: Barrier) -> None:
        """
        pop: the overall population
        indices: the chunk of the population that this thread deals with
        """
        super().__init__(neworder.LinearTimeline(0, 10, 10), lambda: indices[0])  # deterministic, independent
        self.lock = lock
        self.barrier = barrier
        self.pop = pop  # ref not copy
        self.indices = indices

    def step(self) -> None:
        # simulate some complex CPU-bound task that scales linearly
        time.sleep(len(self.indices) / len(self.pop))
        value = self.mc.ustream(1)

        # update the dataset
        with self.lock:  # throwing in a lock is bad
            self.pop.loc[self.indices] += value
        # (optional) synchronise threads at the end of each timestep
        self.barrier.wait()

"""
Run script for multithreaded parallel example
"""

import sys
import time
from concurrent.futures import ThreadPoolExecutor
from threading import Barrier, Lock

import pandas as pd
from parallel_mt import ParallelThreaded  # ty:ignore[unresolved-import]

import neworder

# neworder.verbose()


def sync() -> None:
    """Optionally add code here to run before the barrier is crossed"""
    pass


def run_thread(index: int, lock: Lock, barrier: Barrier) -> None:
    indices = list(range(index * N_PEOPLE // N_THREADS, (index + 1) * N_PEOPLE // N_THREADS))
    m = ParallelThreaded(pop, indices, lock, barrier)
    neworder.run(m)


N_THREADS = 8
N_PEOPLE = 100
pop = pd.DataFrame(index=range(N_PEOPLE), data={"state": 0.0})

# threading primitives
lock = Lock()
barrier = Barrier(N_THREADS, action=sync)


def _is_ft() -> bool:
    if sys.version_info.minor > 12:
        return not sys._is_gil_enabled()
    return False


neworder.log(
    f"neworder FT: {neworder.freethreaded()}, python FT:{_is_ft()}",
    sys.version,
)

t = time.perf_counter()
with ThreadPoolExecutor() as executor:
    futures = [executor.submit(run_thread, i, lock, barrier) for i in range(N_THREADS)]
    results = [f.result() for f in futures]
neworder.log(f"Execution time: {time.perf_counter() - t:.3f}s, (total CPU time ~ 10s)")

# This should output N_THREADS unique values, with (approximately) even counts
print(pop.value_counts())