A simple little module to complete the prime number challenge. Includes 4 tests. This variant is using NumPy. To run the tests, run: `python -m unittest prime-challenge-numpy`

prime-challenge-numpy.py

#!/usr/bin/env python

import unittest
from sympy import sieve
import numpy as np


def generatePrimes(num_primes):
    """
    Returns an array of primes.
    """
    sieve._reset()
    sieve.extend_to_no(num_primes)
    return list(sieve._list)[:num_primes]


def primeMatrix(num_primes):
    """
    Generates a two dimensional array of prime products.
    """
    primes = np.array(generatePrimes(num_primes))

    # generate matrix
    matrix = np.multiply.outer(primes, primes)
    matrix = np.pad(
            matrix,
            pad_width=(1, 0),
            mode='constant',
            constant_values=0
            )

    # set first row
    matrix[0, 1:] = primes

    # set first column
    matrix[1:, 0] = primes

    return matrix


class TestPrimeTools(unittest.TestCase):

    def setUp(self):
        self.matrix = primeMatrix(10)
        self.first_10_primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]

    def testPrimeGenerator(self):
        primes = generatePrimes(10)
        self.assertEqual(
                primes,
                self.first_10_primes,
                'The first 10 primes are incorrect.'
                )

    def testFirstRow(self):
        self.assertEqual(
                list(self.matrix[0][1:]),
                self.first_10_primes,
                'The first row is incorrect.'
                )

    def testFirstColumn(self):
        self.assertEqual(
                [row[0] for row in self.matrix[1:]],
                self.first_10_primes,
                'The first column is incorrect.'
                )

    def testMatrix(self):
        expected_matrix = [
                [0,  2,  3,  5,  7, 11, 13, 17, 19, 23, 29],
                [2,  4,  6, 10, 14, 22, 26, 34, 38, 46, 58],
                [3,  6,  9, 15, 21, 33, 39, 51, 57, 69, 87],
                [5, 10, 15, 25, 35, 55, 65, 85, 95, 115, 145],
                [7, 14, 21, 35, 49, 77, 91, 119, 133, 161, 203],
                [11, 22, 33, 55, 77, 121, 143, 187, 209, 253, 319],
                [13, 26, 39, 65, 91, 143, 169, 221, 247, 299, 377],
                [17, 34, 51, 85, 119, 187, 221, 289, 323, 391, 493],
                [19, 38, 57, 95, 133, 209, 247, 323, 361, 437, 551],
                [23, 46, 69, 115, 161, 253, 299, 391, 437, 529, 667],
                [29, 58, 87, 145, 203, 319, 377, 493, 551, 667, 841]]

        np.testing.assert_array_equal(
                self.matrix,
                expected_matrix,
                'The matrix format is incorrect.'
                )


if __name__ == "__main__":
    print(primeMatrix(10))

Here is an execution profile from generating the matrix with 10,000 primes.

The numpy variant in vanilla c python is way way faster than both the stdlib python variant and the stdlib pypy variant.

Interesting to note, the NumPy variant uses way less memory, almost half as much as the standard library variant.

Here's the execution profile under PyPy. Unlike the standard library variant, this pypy numpy variant did not improve speed over the cpython numpy variant, because most of the heavy lifting is happening in NumPy, which uses C/Fortran memory structures in cpython, so marshalling has to happen in pypy between the stacks. See here for explanation: https://morepypy.blogspot.com/2018/09/inside-cpyext-why-emulating-cpython-c.html