Comprehensions & Broadcasting

Comprehensions & Broadcasting#

This is my favourite chapter to teach, because it’s where Julia stops feeling like a general-purpose programming language and starts feeling like notation.

Everything here could be written as an explicit loop. None of it is faster than an explicit loop — Julia’s loops are already fast, which is exactly the point. What you get instead is code that says what you mean in about a fifth of the characters, and reads back the way you’d say it out loud.

Comprehensions#

A comprehension builds a collection by describing it:

squares = [x^2 for x in 1:5]              # [1, 4, 9, 16, 25]

Read it left to right: “the value x^2, for each x in 1:5.” Add a condition with if:

evens = [x for x in 1:20 if x % 2 == 0]   # [2, 4, ..., 20]

And use multiple iterators to get a grid:

grid = [(i, j) for i in 1:2, j in 1:3]    # a 2×3 Matrix of tuples

Note that with two iterators separated by a comma, you get a matrix whose shape matches the iterators. If you want a flat vector instead, use for i in 1:2 for j in 1:3 (no comma).

A realistic one — computing tax on a set of line items:

amounts = [1200.0, 800.0, 1500.0]
const TAX_RATE = 0.0825

with_tax = [round(a * (1 + TAX_RATE); digits = 2) for a in amounts]
# [1299.0, 866.0, 1623.75]

Generators: comprehensions without the array#

Swap the square brackets for parentheses and you get a generator, which produces values one at a time instead of building the whole collection in memory:

sum(x^2 for x in 1:1_000_000)

That never allocates a million-element array. It computes each square and adds it. For a large enough dataset this is the difference between running and running out of RAM. Note that when a generator is the only argument to a function, you can drop the extra parentheses, as above.

You can also build dictionaries and sets this way:

Dict(name => length(name) for name in ["Amit", "Priya", "Sam"])
Set(x % 5 for x in 1:20)

Broadcasting: the dot#

Here is the single most distinctive piece of Julia syntax, and once it clicks you’ll want it in every other language you use.

Suppose you have a function that works on one number:

celsius_to_fahrenheit(c) = c * 9/5 + 32
celsius_to_fahrenheit(20)        # 68.0

Now you have a thousand of them. In most languages you’d write a loop, or reach for a library function like map, or find a special vectorized version of your function. In Julia you add a dot:

temps = [0, 20, 37, 100]
celsius_to_fahrenheit.(temps)    # [32.0, 68.0, 98.6, 212.0]

That’s it. Any function can be broadcast over any collection by putting a . before the parentheses. You don’t need a vectorized version of your function, because there’s no such thing — there’s just your function, applied elementwise.

The same dot works on operators, going before the operator:

xs = [1, 2, 3, 4]
ys = [10, 20, 30, 40]

xs .+ ys          # [11, 22, 33, 44]
xs .* 2           # [2, 4, 6, 8]
xs .^ 2           # [1, 4, 9, 16]
xs .> 2           # [false, false, true, true]
sqrt.(xs)         # [1.0, 1.414..., 1.732..., 2.0]

Notice xs .* 2: one side is an array and the other is a scalar. Broadcasting expands the scalar to match, which is why you rarely need to write out the loop.

Note

Why does Julia make you type the dot? Newcomers from NumPy or MATLAB, where xs * 2 just works elementwise, often find the explicit dot annoying. Give it a week.

In NumPy, a * b is elementwise but a @ b is matrix multiplication, and you have to remember which is which. In Julia, * always means the mathematical product — matrix multiplication for matrices — and .* always means elementwise. There’s no ambiguity to memorize. When you read A * B you know it’s linear algebra; when you read A .* B you know it’s elementwise. The dot isn’t noise, it’s the information.

Fusion: why this is fast#

Here’s the part that makes broadcasting more than syntactic sugar. Consider:

result = sqrt.(xs .^ 2 .+ ys .^ 2)

A naive implementation would allocate a temporary array for xs .^ 2, another for ys .^ 2, another for the sum, and a fourth for the square root. Four allocations, four passes over memory.

Julia does not do that. It fuses the entire expression into a single loop that computes sqrt(x^2 + y^2) for each element and writes it once. One allocation, one pass.

This happens automatically for any chain of dotted operations, and it’s a genuine performance advantage over the vectorized-library approach in other languages, where each operation is a separate pre-compiled kernel that must materialize its output.

You can go further and eliminate the last allocation too, with .=:

result = similar(xs, Float64)
result .= sqrt.(xs .^ 2 .+ ys .^ 2)     # writes in place, zero allocations

The @. macro#

When an expression has a lot of dots, it gets noisy:

z = sqrt.(a .^ 2 .+ b .^ 2) ./ 2

@. puts a dot on everything in the expression for you:

z = @. sqrt(a^2 + b^2) / 2

Same meaning, much easier to read. This is idiomatic and you should use it.

Broadcasting across dimensions#

Broadcasting also expands shapes, not just scalars:

A = [1 2 3
     4 5 6]           # 2×3

col = [10, 20]        # 2-element vector
row = [100, 200, 300] # 3-element vector

A .+ col              # adds col to every column
A .+ row'             # adds row to every row (note the ' — transpose)

A dimension of length 1 (or a vector matched against a matrix) gets stretched to fit. This is how you normalize every column of a matrix by its mean, or subtract a bias vector from a batch of activations, in one line:

using Statistics
centered = A .- mean(A; dims = 2)     # subtract each row's mean

A realistic example#

Let’s do something that looks like actual work. Take a matrix where each row is a customer and each column is a month of spend, and produce a normalized version plus a per-customer summary.

using Statistics

spend = [1200.0  800.0  1500.0
          300.0  950.0   400.0
         2000.0 2100.0  1900.0]

# Total and average per customer (row-wise)
totals = sum(spend; dims = 2)          # 3×1 matrix
averages = mean(spend; dims = 2)

# Z-score each customer's months against their own mean and spread
stds = std(spend; dims = 2)
normalized = @. (spend - averages) / stds

# Flag any month more than 1 standard deviation above that customer's norm
flags = normalized .> 1.0

Six lines, no loops, no temporaries beyond what’s needed, and every line reads as the sentence you’d use to describe it. That’s the payoff.

Warning

Broadcasting is not always the answer Because the dot is so pleasant, there’s a temptation to broadcast everything. Two situations where you shouldn’t:

When the loop is clearer. If your operation has several steps, conditions, and intermediate names, a for loop is more readable, and in Julia it is not slower. Unlike in Python or R, you never have to contort your code into vectorized form for performance. Write the loop.

When you’re broadcasting a very expensive function over a huge array once. Fusion helps chains of cheap operations; it doesn’t make one expensive call cheaper. Profile before assuming.

Try it yourself#

  1. Using a comprehension, build a vector of the first 10 triangular numbers (n(n+1)/2).

  2. Given prices = [9.99, 19.50, 4.25, 30.00] and a tax rate of 8.25%, produce a vector of tax-inclusive prices rounded to 2 decimals — first with a comprehension, then with @..

  3. Given readings = [20.5, -99.0, 22.3, -99.0, 19.8], use broadcasting and a boolean mask to compute the mean of only the valid readings.

  4. Given a 3×4 matrix, subtract the mean of each column from that column, in one line.

Solutions#

# 1.
[n * (n + 1) ÷ 2 for n in 1:10]        # [1, 3, 6, 10, 15, 21, 28, 36, 45, 55]

# 2.
prices = [9.99, 19.50, 4.25, 30.00]
[round(p * 1.0825; digits = 2) for p in prices]
@. round(prices * 1.0825; digits = 2)

# 3.
using Statistics
readings = [20.5, -99.0, 22.3, -99.0, 19.8]
mean(readings[readings .!= -99.0])      # 20.866...

# 4.
using Statistics
M = rand(3, 4)
centered = M .- mean(M; dims = 1)

You can now express transformations over whole collections concisely. Next we look at the machinery underneath all of this — Julia’s type system, and how to build types of your own.