Skip to content

[numba.md] Update np.random → Generator API #550

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Chihiro2000GitHub wants to merge 2 commits into
base: main
Choose a base branch
from
+32 −15
Open

[numba.md] Update np.random → Generator API #550

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
397a891
Update rng usage in numba.md
Chihiro2000GitHub May 11, 2026
aef82bd
Address reviewer feedback on rng usage in numba.md
Chihiro2000GitHub May 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 32 additions & 15 deletions lectures/numba.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -454,11 +454,17 @@ Compare speed with and without Numba when the sample size is large.
Here is one solution:

```{code-cell} ipython3
n = 1_000_000
rng = np.random.default_rng()
u_draws = rng.uniform(size=n)
v_draws = rng.uniform(size=n)

@jit
def calculate_pi(n=1_000_000):
def calculate_pi(u_draws, v_draws):
n = len(u_draws)
count = 0
for i in range(n):
u, v = np.random.uniform(0, 1), np.random.uniform(0, 1)
u, v = u_draws[i], v_draws[i]
d = np.sqrt((u - 0.5)**2 + (v - 0.5)**2)
if d < 0.5:
count += 1
Expand All @@ -471,12 +477,12 @@ Now let's see how fast it runs:

```{code-cell} ipython3
with qe.Timer():
calculate_pi()
calculate_pi(u_draws, v_draws)
```

```{code-cell} ipython3
with qe.Timer():
calculate_pi()
calculate_pi(u_draws, v_draws)
```

If we switch off JIT compilation by removing `@jit`, the code takes around
Expand Down Expand Up @@ -550,10 +556,13 @@ p, q = 0.1, 0.2 # Prob of leaving low and high state respectively
Here's a pure Python version of the function

```{code-cell} ipython3
def compute_series(n):
n = 1_000_000
rng = np.random.default_rng()
U = rng.uniform(0, 1, size=n)

def compute_series(n, U):
x = np.empty(n, dtype=np.int64)
x[0] = 1 # Start in state 1
U = np.random.uniform(0, 1, size=n)
for t in range(1, n):
current_x = x[t-1]
if current_x == 0:
Copy link
Copy Markdown
Member

@HumphreyYang HumphreyYang May 21, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Many thanks @Chihiro2000GitHub,

I think we do not need to move the sampling outside here.

We can use RNG inside @jit-decorated functions as long as we are not using @jit(parallel=True).

(See https://manual.quantecon.org/styleguide/code.html#numba-njit-jit)

So the old commit, which passes rng as an argument, looks good to me for this case which we did not use @jit(parallel=True).

We only need to make Exercises 1 and 3 consistent in how they handle sampling so that the timing comparisons are fair.

Expand All @@ -567,8 +576,7 @@ Let's run this code and check that the fraction of time spent in the low
state is about 0.666

```{code-cell} ipython3
n = 1_000_000
x = compute_series(n)
x = compute_series(n, U)
print(np.mean(x == 0)) # Fraction of time x is in state 0
```

Expand All @@ -578,7 +586,7 @@ Now let's time it:

```{code-cell} ipython3
with qe.Timer():
compute_series(n)
compute_series(n, U)
```

Next let's implement a Numba version, which is easy
Expand All @@ -590,15 +598,15 @@ compute_series_numba = jit(compute_series)
Let's check we still get the right numbers

```{code-cell} ipython3
x = compute_series_numba(n)
x = compute_series_numba(n, U)
print(np.mean(x == 0))
```

Let's see the time

```{code-cell} ipython3
with qe.Timer():
compute_series_numba(n)
compute_series_numba(n, U)
```

This is a nice speed improvement for one line of code!
Expand Down Expand Up @@ -636,11 +644,17 @@ For the size of the Monte Carlo simulation, use something substantial, such as
Here is one solution:

```{code-cell} ipython3
n = 1_000_000
rng = np.random.default_rng()
u_draws = rng.uniform(size=n)
v_draws = rng.uniform(size=n)

@jit(parallel=True)
def calculate_pi(n=1_000_000):
def calculate_pi(u_draws, v_draws):
n = len(u_draws)
count = 0
for i in prange(n):
u, v = np.random.uniform(0, 1), np.random.uniform(0, 1)
u, v = u_draws[i], v_draws[i]
d = np.sqrt((u - 0.5)**2 + (v - 0.5)**2)
if d < 0.5:
count += 1
Comment on lines +647 to 660
Copy link
Copy Markdown
Member

@HumphreyYang HumphreyYang May 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this avoids the thread-safety issue, but the timing comparison is no longer quite apples-to-apples because the random draws have been moved outside the timed function.

Personally, I think we can generate the draws in the first solution by moving

n = 1_000_000
rng = np.random.default_rng()
u_draws = rng.uniform(size=n)
v_draws = rng.uniform(size=n)

to line 457, and modifying the function in the code block starting at line 457 (solution to exercise 1) to take u_draws and v_draws, like the function here.

In this case, we would have a fair timing comparison.

Please let me know what you think!

Copy link
Copy Markdown
Member

@HumphreyYang HumphreyYang May 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @mmcky, I think this is an example where we need to make a tough choice between keeping the legacy API for simplicity and using the new rng API with some trade-offs.

Copy link
Copy Markdown
Contributor

@mmcky mmcky May 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @HumphreyYang I agree with you. I have added some additional notes below regarding double checking nopython mode as well.

Expand All @@ -653,12 +667,12 @@ Now let's see how fast it runs:

```{code-cell} ipython3
with qe.Timer():
calculate_pi()
calculate_pi(u_draws, v_draws)
```

```{code-cell} ipython3
with qe.Timer():
calculate_pi()
calculate_pi(u_draws, v_draws)
```

By switching parallelization on and off (selecting `True` or
Expand Down Expand Up @@ -750,6 +764,9 @@ $$

Using this fact, the solution can be written as follows.

Note that random draws are kept inside the inner loop rather than pre-allocated,
to avoid creating large shock arrays of size `M * n`.

Comment on lines +767 to +769
Copy link
Copy Markdown
Member

@HumphreyYang HumphreyYang May 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a nice clarification! It would be even better to explicitly state that we are using the legacy API.

```{note}
Note that random draws are kept inside the inner loop and do not use `rng` as before.

Instead, we use the legacy `np.random.randn()` API.

This is because Numba support for `Generator` objects is not [thread-safe](https://numba.readthedocs.io/en/stable/reference/numpysupported.html#generator-objects) under parallel execution.
```


```{code-cell} ipython3
M = 10_000_000
Expand Down
Loading