Introduction - MPI Example

This lesson is a brief introduction to TOAST and its data representations. This next cell is just initializing some things for the notebook.

In [ ]:
# Are you using a special reservation for a workshop?
# If so, set it here:
nersc_reservation = None

# Load common tools for all lessons
import sys
sys.path.insert(0, "..")
from lesson_tools import (
    check_nersc
)
nersc_host, nersc_repo = check_nersc(reservation=nersc_reservation)
if nersc_host is not None:
    %reload_ext slurm_magic

# Capture C++ output in the jupyter cells
%load_ext wurlitzer
In [ ]:
%%writefile intro_mpi.py

import toast
from toast.mpi import MPI

# Load common tools for all lessons
import sys
sys.path.insert(0, "..")
from lesson_tools import (
    fake_focalplane
)

import numpy as np
import matplotlib.pyplot as plt

import toast.qarray as qa

# Runtime Environment

env = toast.Environment.get()
if MPI.COMM_WORLD.rank == 0:
    print(env, flush=True)

# Create a fake focalplane

fp = fake_focalplane()

detnames = list(sorted(fp.keys()))
detquat = {x: fp[x]["quat"] for x in detnames}
detfwhm = {x: fp[x]["fwhm_arcmin"] for x in detnames}
detlabels = {x: x for x in detnames}
detpolcol = {x: "red" if i % 2 == 0 else "blue" for i, x in enumerate(detnames)}

if MPI.COMM_WORLD.rank == 0:
    outfile = "intro_focalplane_mpi.pdf"
    toast.tod.plot_focalplane(
        detquat, 4.0, 4.0, outfile, fwhm=detfwhm, polcolor=detpolcol, labels=detlabels
    )

# Different communicators

# Default communicator (one group using COMM_WORLD)

comm = toast.Comm()
if comm.world_rank == 0:
    print("----- Default Comm -----")
for p in range(MPI.COMM_WORLD.size):
    if p == comm.world_rank:
        print(comm, flush=True)
    comm.comm_world.barrier()
    
# Communicator with group size of one.

comm = toast.Comm(world=MPI.COMM_WORLD, groupsize=1)
if comm.world_rank == 0:
    print("----- Comm: groupsize = 1 -----")
for p in range(MPI.COMM_WORLD.size):
    if p == comm.world_rank:
        print(comm, flush=True)
    comm.comm_world.barrier()
    
# Communicator with 2 groups

gsize = MPI.COMM_WORLD.size // 2

comm = toast.Comm(world=MPI.COMM_WORLD, groupsize=gsize)
if comm.world_rank == 0:
    print("----- Comm: groupsize = {} -----".format(gsize))
for p in range(MPI.COMM_WORLD.size):
    if p == comm.world_rank:
        print(comm, flush=True)
    comm.comm_world.barrier()
In [ ]:
if nersc_host is not None:
    %srun -N 1 -C haswell -n 32 -c 2 --cpu_bind=cores -t 00:03:00 python intro_mpi.py >intro_mpi.log 2>&1
else:
    # Just use mpirun
    import subprocess as sp
    sp.check_call("mpirun -np 4 python intro_mpi.py >intro_mpi.log 2>&1", shell=True)
In [ ]: