04_mpi.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MPI\n",
    "\n",
    "MPI (Message Passing Interface) is a standard for writing parallel programs that can run on a distributed memory system. It is widely used in the HPC (High Performance Computing) community. There are implementations of MPI for many programming languages, including C, C++, Fortran, and Python. Much of the functionality and many of the commands we will see here will be similar in other languages. \n",
    "\n",
    "There are multiple implementations of MPI, including [OpenMPI](https://www.open-mpi.org/), [MPICH](https://www.mpich.org/), and [Intel MPI](https://www.intel.com/content/www/us/en/developer/tools/oneapi/mpi-library.html). If you are running this notebook in a Github Codespace, MPICH will already be installed. If you want to install MPICH on your local system you can follow the instructions [here](https://www.mpich.org/downloads/).\n",
    "\n",
    "We will also be using the `mpi4py` package which is a Python wrapper for MPI. This is already installed if you'r running this notebook in a GitHUb Codespace, or you can install it locally using pip:\n",
    "\n",
    "```bash\n",
    "pip install mpi4py\n",
    "```\n",
    "\n",
    "Unlike the other methods we have seen so far, we run MPI from the terminal using the `mpiexec` command. We can run a Python script named `python_script.py` using MPI like this:\n",
    "\n",
    "```bash\n",
    "mpiexec -n 4 python python_script.py\n",
    "```\n",
    "In this command, the `-n 4` flag tells MPI to run the script using 4 processes. The text `python python_script.py` tells which command MPI should be running on each process. This will run the script 4 times in parallel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ranks\n",
    "\n",
    "Each copy of the code will be executed on its own process and will be identified by its \"rank\". The rank is a unique integer identifier for each process which can be accessed using the `mpi4py.MPI.COMM_WORLD.Get_rank()` function. We can also get the number of ranks using the method `Get_rank()`. The code below shows how to access the rank:\n",
    "\n",
    "```python\n",
    "# Run this script with the terminal command `mpiexec -n 4 python get_rank.py`\n",
    "\n",
    "import mpi4py.MPI as MPI\n",
    "\n",
    "# Get a reference to the current MPI.COMM_WORLD communicator\n",
    "comm = MPI.COMM_WORLD\n",
    "\n",
    "# Get the total number of ranks in the communicator\n",
    "n_rank = comm.Get_size()\n",
    "\n",
    "# Get the rank of the current process\n",
    "rank = comm.Get_rank()\n",
    "\n",
    "# Print the rank of the current process\n",
    "print(f'This script is being run by Rank {rank} out of {n_rank} total ranks')\n",
    "```\n",
    "\n",
    "This code can be found in the file [`04_mpi_scripts/get_rank.py`](04_mpi_scripts/get_rank.py). To run this command we will need to change directory in the terminal using the command:\n",
    "\n",
    "```bash\n",
    "cd 04_mpi_scripts\n",
    "```\n",
    "\n",
    "and run the script using MPI using the command:\n",
    "\n",
    "```bash\n",
    "mpiexec -n 4 python get_rank.py\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Communicating Between Ranks\n",
    "\n",
    "In the above example, we access the variable `MPI.COMM_WORLD` which references a communicator. We can use a communicator to send messages between the processes it contains. The communicator `MPI.COMM_WORLD` contains all the processes that are running the script. It is possible to create other communicators which contain only a subset of the processes, which can be useful for more complex parallel programs, but we won't be covering that here.\n",
    "\n",
    "We can send messages between ranks using the `send` and `recv` methods of the communicator. The code below shows how to send a message from rank 0 to rank 1:\n",
    "\n",
    "```python\n",
    "from mpi4py import MPI\n",
    "\n",
    "# Get the communicator and the rank of the process\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = comm.Get_rank()\n",
    "\n",
    "if rank == 0:\n",
    "    # If we're in rank 0, send a dictionary to rank 1\n",
    "    data = {'a': 7, 'b': 3.14}\n",
    "    comm.send(data, dest=1)\n",
    "elif rank == 1:\n",
    "    # If we're in rank 1, receive the dictionary from rank 0\n",
    "    data = comm.recv(source=0)\n",
    "    print(rank, data)\n",
    "```\n",
    "\n",
    "In the example above the `send` method is used to send a dictionary from rank 0 to rank 1. By including it in the if-block, we make sure it is only called by the rank 0. The first argument to `send` is the data to be sent. We also specify the destination rank so the message can be sent to rank 1. The `recv` method is called from rank 1 to receive the message from rank 0. The `source` argument specifies the rank of the process that sent the message. The data which is received is saved into the variable `data` and then printed. As we can see, we can send any type of data between ranks, including dictionaries, lists, and numpy arrays.\n",
    "\n",
    "Both `send` and `recv` block, meaning that the program will wait at the `send` line until the message has been received by the destination rank, and will wait at the `recv` line until the message has been sent by the source rank. This means we need to plan carefully to make sure that the program doesn't get deadlocked. For example, the code below will deadlock as both ranks are waiting for the other to receive a message:\n",
    "\n",
    "```python\n",
    "from mpi4py import MPI\n",
    "\n",
    "# Get the communicator and the rank of the process\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = comm.Get_rank()\n",
    "\n",
    "if rank == 0:\n",
    "    # Send a message to rank 1\n",
    "    comm.send(\"Hello from rank 0\", dest=1)\n",
    "    # Receive the message from rank 1\n",
    "    data = comm.recv(source=1)\n",
    "elif rank == 1:\n",
    "    # Send a message to rank 0\n",
    "    comm.send(\"Hello from rank 1\", dest=0)\n",
    "    # Receive the message from rank 0\n",
    "    data = comm.recv(source=0)\n",
    "```\n",
    "\n",
    "This code can be found in the file [`04_mpi_scripts/deadlock.py`](04_mpi_scripts/deadlock.py) and should be run with two processes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Non-Blocking Communication\n",
    "\n",
    "To avoid deadlocks, we can use non-blocking communication. This allows the program to continue running while the message is being sent or received. This can be done using the `isend` and `irecv` functions. These returned a `Request` object. We can use the `wait` method of the `Request` object to wait until the message has been sent or received. If the `Request` object was created by `irecv` the `wait` method waits until it receives a value, then returns the value received. Using non-blocking communication can free up processes to do other work while the communication is pending. The code below shows how to use non-blocking communication:\n",
    "\n",
    "```python\n",
    "from mpi4py import MPI\n",
    "\n",
    "# Get the communicator and the rank of the process\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = comm.Get_rank()\n",
    "\n",
    "if rank == 0:\n",
    "    # Send the integer 100 to rank 1\n",
    "    req = comm.isend(100,dest=1)\n",
    "    # Wait for the request to complete\n",
    "    req.wait()\n",
    "elif rank == 1:\n",
    "    # Receive the integer from rank 0\n",
    "    req = comm.irecv(source=0)\n",
    "    # Wait for the request to complete and get the data\n",
    "    data = req.wait()\n",
    "    print(data)\n",
    "```\n",
    "\n",
    "This code can be found in the file [`04_mpi_scripts/non_blocking.py`](04_mpi_scripts/non_blocking.py) and should be run with two processes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sending Numpy Arrays\n",
    "\n",
    "The methods we've seen before like `send`, `recv`, `isend`, and `irecv` send Python objects between ranks using a process called [pickling](https://docs.python.org/3/library/pickle.html). This process allows an arbitrarily complex object to be serialized so they can be sent between ranks. Whilst flexible in terms of the types of object that can be sent, the process of pickling and unpicking adds a significant performance overhead to the sending of data between ranks.\n",
    "\n",
    "However, some data types in Python, such as Numpy arrays, do not need to be pickled to be sent between ranks, which can speed up communication significantly. This can be done using functions with similar names to those we have already seen, except they begin with a capital letter, such as `Send` and `Recv`.\n",
    "\n",
    "The syntax we have to use is a little different than before. We need to prepare an object to receive the data before we call the `Recv` method. This array should be a Numpy array with the same shape and data type as the array we are sending. This prepares a section of the memory in receiving rank to receive the data and is known as a buffer. The function `numpy.empty` is an efficient way to create this buffer. It will allocate the memory for th Numpy array but will not initialize it, meaning it will contain junk values. This is faster than using `numpy.zeros` which would initialize the array to zeros. For many Numpy functions, including `empty`, the `dtype` argument is used to [specify the data type](https://numpy.org/doc/2.1/reference/arrays.dtypes.html) of the array. There are a few ways to do this, but one ay is to use the Python type names such as `int`, `float`, `complex`, etc.\n",
    "\n",
    "The code below shows how to send a Numpy array between ranks:\n",
    "\n",
    "```python\n",
    "from mpi4py import MPI\n",
    "import numpy as np\n",
    "\n",
    "# Get the communicator and the rank of the process\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = comm.Get_rank()\n",
    "\n",
    "if rank == 0:\n",
    "    # If we're in rank 0, create an array of ten integers to send\n",
    "    data = np.arange(10, dtype=float)\n",
    "    # Send the array to rank 1\n",
    "    comm.Send(data, dest=1)\n",
    "elif rank == 1:\n",
    "    # If we're in rank 1, create an array to receive the data\n",
    "    data = np.empty(10, dtype=float)\n",
    "    # data will initially contain junk values\n",
    "    print('data before: ', data)\n",
    "    # Receive the data from rank 0\n",
    "    comm.Recv(data, source=0)\n",
    "    print(rank, data)\n",
    "```\n",
    "\n",
    "This code can be found in the file [`04_mpi_scripts/numpy_send_recv.py`](04_mpi_scripts/numpy_send_recv.py) and should be run with two processes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise: Sum of Powers of a Array\n",
    "\n",
    "The function `random_float_array` in the file [`sum_of_powers.py`](04_mpi_scripts/sum_of_powers.py) generates a random array of integers. You should generate a random array of 100 floats with a minimum of zero and maximum of 10 using the function call `random_float_array(0, 10, 100)` on rank 0. Then send this array to all other ranks. Each rank, including rank 0, should calculate the value:\n",
    "\n",
    "$$\n",
    "\\sum_{i=0}^{n-1} x_i^{r+1}\n",
    "$$\n",
    "\n",
    "where $x_i$ is the $i$ th element of the array and $r$ is the rank of the process. Each rank should then send the result back to rank 0. Rank 0 should assemble the results into a list `results` whose $i$ th element is the sum of the powers of the array elements calculated by rank $i$. Finally, rank 0 should print the list of results. This code should be able to be run with any number of ranks. For example, if the code is run on 2 ranks, you might receive the result:\n",
    "\n",
    "```\n",
    "[5079.372714885385, 34822.40206813637]\n",
    "```\n",
    "\n",
    "while on four ranks you might get\n",
    "\n",
    "```\n",
    "[5022.381368813411, 33480.47733076967, 251675.5452958118, 2018665.9149521096]\n",
    "```\n",
    "\n",
    "Note that you will get slightly different results as your array will contain different random numbers. The first entry is the sum of the array and is generated on rank 0, the next value is the sum of the square of the array and is calculated on rank 1, the next value is the sum of the cube of the array and is calculated on rank 2, and so on.\n",
    "\n",
    "There is a sample solution in the file [`sample_solutions/sum_of_powers_solution.py`](sample_solutions/sum_of_powers.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Collective Communication\n",
    "\n",
    "In the last exercise, we send data from one rank to all other ranks. This sort of communication is a common thing you might want to do in parallel programs. MPI allows for collective communication which provides a convenient and efficient way to do this. The most common collective communication functions are `bcast`, `scatter`, `gather`, and `reduce`. These functions are called by all ranks in the communicator and communicates with all other ranks in the communicator. There are also equivalent functions that start with a capital letter, such as `Bcast`, `Scatter`, `Gather`, and `Reduce`, which you can use for sending Numpy arrays and other types of data that don't need to be pickled.\n",
    "\n",
    "### Broadcast\n",
    "\n",
    "The `bcast` function sends a single object from one rank to all other ranks in the communicator. The syntax is:\n",
    "\n",
    "```python\n",
    "import mpi4py.MPI as MPI\n",
    "\n",
    "# Get the communicator and the rank of the process\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = comm.Get_rank()\n",
    "\n",
    "if rank == 0:\n",
    "    # If the rank is 0, set the data to be broadcasted\n",
    "    data = ['apples', 'bananas', 'cherries', 'dates']\n",
    "else:\n",
    "    # If the rank is not 0, we still need the variable to exist\n",
    "    # Set it to None for now\n",
    "    data = None\n",
    "\n",
    "# Broadcast the data from rank 0 to all other ranks\n",
    "data = comm.bcast(data, root=0)\n",
    "\n",
    "# Each rank now has a copy of the data\n",
    "print(f'Rank {rank} has data: {data}')\n",
    "```\n",
    "\n",
    "The code above can be found in the file [`04_mpi_scripts/broadcast.py`](04_mpi_scripts/broadcast.py) and should be run with 4 processes. This code is more compact and efficient that the equivalent code using `send` and `recv`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scatter\n",
    "\n",
    "The `scatter` function distributes the elements of an array from one rank to all other ranks in the communicator. The array should have the same number of entries as the number of ranks in the communicator. The syntax is:\n",
    "\n",
    "```python\n",
    "import mpi4py.MPI as MPI\n",
    "\n",
    "# Get the communicator and the rank of the process\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = comm.Get_rank()\n",
    "n_rank = comm.Get_size()\n",
    "\n",
    "if rank == 0:\n",
    "    # If the rank is 0, set the data to be scattered\n",
    "    data = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]\n",
    "    # Divide the data into equal parts\n",
    "    n_per_rank = len(data) / n_rank\n",
    "    data_local = []\n",
    "    for i in range(n_rank):\n",
    "        data_local.append(data[int(i*n_per_rank):int((i+1)*n_per_rank)])\n",
    "else:\n",
    "    # If the rank is not 0, we still need the variable to exist\n",
    "    # Set it to None for now\n",
    "    data_local = None\n",
    "\n",
    "# Scatter the data from rank 0 to all other ranks\n",
    "data_local = comm.scatter(data_local, root=0)\n",
    "\n",
    "# Each rank now has a piece of the data in data_local\n",
    "print(f'Rank {rank} has data: {data_local}')\n",
    "```\n",
    "\n",
    "The code above can be found in the file [`04_mpi_scripts/scatter.py`](04_mpi_scripts/scatter.py) and should be run with 4 processes.\n",
    "\n",
    "One complication of the method that is reflected in the code above is that the data must be divided into as many parts as there are ranks. This can necessitate some preparation of the data as in the code above, which is a little cumbersome. However, the `scatter` function is more efficient than sending the data to each rank individually."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gather\n",
    "\n",
    "The `gather` function is the opposite of `scatter`. It collects the data from all ranks in the communicator and assembles it into an array on one rank. The syntax is:\n",
    "\n",
    "```python\n",
    "import mpi4py.MPI as MPI\n",
    "\n",
    "# Get the communicator and the rank of the process\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = comm.Get_rank()\n",
    "n_rank = comm.Get_size()\n",
    "\n",
    "# We're aiming to write an array with square of each number.\n",
    "# We'll compute part of the array on each rank and then gather the results.\n",
    "n = 10\n",
    "i_local_start = int(rank * n / n_rank)\n",
    "i_local_stop = int((rank + 1) * n / n_rank)\n",
    "data_local = [i**2 for i in range(i_local_start, i_local_stop)]\n",
    "\n",
    "# Print the local data on each rank\n",
    "print(f'Rank {rank} has data: {data_local}')\n",
    "\n",
    "# Gather the data from all ranks to rank 0\n",
    "data = comm.gather(data_local, root=0)\n",
    "\n",
    "# Rank 0 now has all the data\n",
    "if rank == 0:\n",
    "    print(f'Rank {rank} has data: {data} before flattening')\n",
    "    # Flatten the list of lists into a single list\n",
    "    data = sum(data, [])\n",
    "    print(f'Rank {rank} has data: {data} after flattening')\n",
    "```\n",
    "\n",
    "The code above can be found in the file [`04_mpi_scripts/gather.py`](04_mpi_scripts/gather.py) and should be run with 4 processes. Again, we've had to do a little work to work out which parts of the data each rank should be working on in advance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reduce\n",
    "\n",
    "The `reduce` function is used to combine data from all ranks in the communicator into a single value. The syntax is:\n",
    "\n",
    "```python\n",
    "# Run this script with the terminal command `mpiexec -n 4 python gather.py`\n",
    "\n",
    "import mpi4py.MPI as MPI\n",
    "import numpy as np\n",
    "\n",
    "# Get the communicator and the rank of the process\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = comm.Get_rank()\n",
    "n_rank = comm.Get_size()\n",
    "\n",
    "# We're going to generate 10 random numbers between 0 and 1 and count how many are less than 0.5\n",
    "n = 10\n",
    "\n",
    "# Calculate how many numbers each rank will generate\n",
    "n_local = int(n / n_rank)\n",
    "if rank < n % n_rank:\n",
    "    # If the rank is less than the remainder, it will generate one more number\n",
    "    # This is to ensure that the right number of numbers are generated\n",
    "    n_local += 1\n",
    "\n",
    "# Generate the random numbers\n",
    "numbers_local = np.random.rand(n_local)\n",
    "count_local = np.sum(numbers_local < 0.5)\n",
    "\n",
    "# Print the local data on each rank\n",
    "print(f'Rank {rank} has generated {n_local} numbers, and {count_local} are less than 0.5')\n",
    "\n",
    "# Reduce the count of numbers less than 0.5 from all ranks to rank 0\n",
    "count = comm.reduce(count_local, op=MPI.SUM, root=0)\n",
    "\n",
    "# Rank 0 now has the total count\n",
    "if rank == 0:\n",
    "    print(f'Rank {rank} has the overall result of {count} numbers less than 0.5')\n",
    "else:\n",
    "    # Reduce returns None on all ranks other than the root\n",
    "    print(f'On rank {rank}, the value of count is {count}')\n",
    "\n",
    "# If we want all ranks to have the total count, we can use allreduce\n",
    "# This does not require a root rank\n",
    "count_all = comm.allreduce(count_local, op=MPI.SUM)\n",
    "\n",
    "print(f'On rank {rank}, the value of count_all is {count_all}')\n",
    "```\n",
    "\n",
    "The code above can be found in the file [`04_mpi_scripts/reduce.py`](04_mpi_scripts/reduce.py) and should be run with 4 processes. There are a couple of things to note in the code above. The first is that the `reduce` function only returns a value on the root rank. On all other ranks, it returns `None`. If we want each rank to have a copy of the result, we can use the `allreduce` function. The second thing to note is that we have to specify an operation to the `op` argument to perform on the data to combine it into a single value. In this case, we are using `MPI.SUM` to add the counts from each rank together. There are other operations available, some common ones are:\n",
    "\n",
    "* `MPI.SUM` - Adds the values together\n",
    "* `MPI.PROD` - Multiplies the values together\n",
    "* `MPI.MAX` - Returns the maximum value\n",
    "* `MPI.MIN` - Returns the minimum value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise: Share Price Prediction\n",
    "\n",
    "Share prices change over time in an unpredictable way. However, we can simulate how the share price might change over time using a geometric random walk. We can say that the price of a share on a given day is equal to the price of the share on the previous day multiplied by a random number drawn from a normal distribution with a given mean and standard deviation. We can simulate this process by generating a random number for each day and multiplying it by the share price on the previous day. We can do this for a number of days to simulate the share price over time. In the file [`share_price.py`](04_mpi_scripts/share_price.py) there is a function `simulate_share_price` which takes the initial share price, the mean of the normal distribution, the standard deviation of the normal distribution, and the number of days to simulate. This function returns the share price at the end of the final day of the simulated period.\n",
    "\n",
    "Each call to this function will return a different final share price as a different set of random numbers will have been generated. This means we can call the function multiple times to get a distribution of final share prices. Your task is to write code which will call this code `n` times in total, split across each rank in the communicator. Once this is done, use collective communication to collect the results on rank zero and print the mean, standard deviation, minimum and maximum values of the final share prices. The code should be able to be run with any number of ranks. Your code should have values for the initial share price, mean, standard deviation, and number of days to simulate, and the number of simulations to run hard-coded into the script. Try the following values:\n",
    "\n",
    "* Initial share price: 100\n",
    "* Mean fractional daily change: 0.001\n",
    "* Standard deviation of fractional daily change: 0.02\n",
    "* Number of days to simulate: 100\n",
    "* Number of simulations: 1000\n",
    "\n",
    "As a reminder, the equations are the mean and standard deviation are:\n",
    "\n",
    "$$\n",
    "\\bar{x} = \\frac{1}{n} \\sum_{i=0}^{n-1} x_i\\\\\n",
    "\\sigma = \\sqrt{\\frac{1}{n} \\sum_{i=0}^{n-1} (x_i - \\bar{x})^2} = \\sqrt{\\frac{1}{n} \\sum_{i=0}^{n-1} x_i^2 - \\bar{x}^2}\n",
    "$$\n",
    "\n",
    "where $\\bar{x}$ is the mean, $\\sigma$ is the standard deviation and $x_i$ is the value of each piece of data (the final share price of the $i$ th simulation in this case). There are two forms of the standard deviation equation given, you may use whichever you think is more appropriate.\n",
    "\n",
    "Think carefully about how to structure your code and use the collective communication functions you've seen. There is a sample solution in the file [`sample_solutions/share_price.py`](sample_solutions/share_price.py). For reference, when I ran the sample solution, I got the following output:\n",
    "\n",
    "* Mean: 110.70897767200135\n",
    "* Standard deviation: 23.170420016985695\n",
    "* Minimum: 54.494856062641915\n",
    "* Maximum: 220.93130434066072\n",
    "\n",
    "Your code will probably not produce the same results as the random numbers generated will be different, but your results should be similar."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}