Writing readable mathematical code

In Statnett, we solve problems that involve a significant amount of mathematics. We need code to solve these problems, and writing readable, maintainable, testable and efficient code that computes intricate mathematical equations can be challenging. This blog post will look at one way to tackle this challenge. As a practical example, we will discuss the architectural and style choices behind the (soon to be open source) linerate Python package, which implements the conductor heat model described in CIGRE TB 601.

Dynamic line rating and CIGRE TB 601

Let’s first look briefly at the usecase we consider:  CIGRE’s Technical Brochure number 601, or CIGRE601. CIGRE601 is an essential part of our dynamic line rating work, which you can read more about in one of our previous blog posts. Here, we focus on the architectural and programming style choices we made when we developed code that implements the CIGRE 601 thermal heating model.

The mathematics of CIGRE 601 is not overly complicated, but there is a lot of it, which can make it difficult to follow if you’re not careful. At its core, CIGRE601 is a heat balance equation:

$$P_J(I) + P_s(\mathbf{w}) = P_c(T, \mathbf{w}) + P_r(T, \mathbf{w}). $$

In other words, the heating from the current ($P_J$) plus the solar heating ($P_s$) is equal to the convective cooling ($P_c$) plus the radiative (black body) cooling ($P_r$). In other words: the heating from the sun and the electricity flowing through is equal to the cooling from the wind and the radiated heat waves when the line has a constant temperature. In our case, we want to provide a current, $I$, and weather conditions, $\mathbf{w}$, and use these to solve the above equation with respect to the line temperature, $T$.

To evaluate only a single term of the above equation, we must evaluate up to another 20 other “sub-equations”. In the code, we need to keep track of all these sub-equations: what should be their input and what should be their output? We also want the implementations to be readable correct, and efficient.

Architecture

To get readable correct, and efficient code, we need to consider the software architecture. We need an architecture that strives for readability and testability without sacrificing speed (too much). How can we accomplish this?

A user-friendly frontend with a functional backend

As a basic user of a toolkit, I want to run code on the form

Python
model = Cigre601(line_parameters, weather_parameters, time)
convective_cooling = model.compute_convective_cooling(temperature)

However, we also want access to the equations themselves. While the code above is excellent for computing properties for specific lines (like the total convective cooling or ampacity rating), it could be better structured for exploring how the equations behave. To facilitate such exploration, we also want access to the sub-terms themselves.

What’s more, if we use an object-oriented approach to implement these sub-terms, then it might be appealing to store their values as properties or attributes. By doing so, we can save time on repeated computations. However, we also hide what each function depends on and provides. We want an interface that makes it clear what information we need for each computation, and if the methods implementing the equations have access to self.anything_at_all, then it is unclear what they need.

Instead, I think that a functional approach, where each mathematical equation is implemented as a pure function— a deterministic function without state — is better. Therefore, linerate has a single function for (more or less) each equation in CIGRE 601. This makes is very easy to see how the code ties together with the equations. For example, the method that computes the Joule heating is defined by

Python
def compute_joule_heating(current, resistance):
    I = current  # noqa
    R = resistance
    return I * I * R

And yes, this approach will create a minor additional computational overhead, but I found it insignificant compared to the benefits such an approach gave.

Still, I’m not saying that an object-oriented approach is only a bad thing. For linerate, we have the best of both worlds: a functional “backend” that makes it easy to explore each equation in CIGRE601 and an object-oriented “frontend” that makes it easy to compute ampacity ratings for a given power line.

Performant Python code

An essential part of mathematical code is its performance. We do not want to use native Python code for any numerical operation. Instead, we want to use the vectorisation capabilities of NumPy.

Specifically, we only use basic binary operators and NumPy universal functions (ufuncs). In the cases when we cannot use builtin ufuncs or binary operators (for example table lookups), I recommend using numba.vectorize to create your own blazing fast JIT-compiled ufuncs. By structuring the code this way, we never need to write another loop!

Say, for example, that we want to compute the line rating for all integer air temperatures between $0^\circ \text{C}$ and $6^\circ \text{C}$ and five linearly spaced wind speeds between $0 \frac{\text{m}}{\text{s}}$ and $1 \frac{\text{m}}{\text{s}}$. If we make sure that all code only uses binary operators and numpy ufuncs, then we only need to write

Python
[...]

air_temperature = np.arange(0, 7).reshape(7, 1)
wind_speed = np.linspace(0, 1, 5).reshape(1, 5)
weather = Weather(
    air_temperature=air_temperature,
    wind_speed=wind_speed
    [...]
)
model = Cigre601(weather=weather, [...])
line_rating = model.compute_steady_state_ampacity(current=2000)

The magic happens in the highlighted lines. By adding singleton axes to the array, we enable NumPy array broadcasting. Broadcasting is almost magic! Typically, when we do array computations, we need array dimensions that match up. However, singleton dimensions are “wildcards” and will match any other dimension size by repeating values in the singleton dimension until they match up. I think the best way to understand broadcasting is by computing your own multiplication table:

Python
x = np.arange(0, 5).reshape(1, 5)
y = np.arange(0, 5).reshape(5, 1)

print("x:")
print(x)
print("y:")
print(y)
print("x * y:")
print(x * y)
x:
[[0 1 2 3 4]]
y:
[[0]
 [1]
 [2]
 [3]
 [4]]
x * y:
[[ 0  0  0  0  0]
 [ 0  1  2  3  4]
 [ 0  2  4  6  8]
 [ 0  3  6  9 12]
 [ 0  4  8 12 16]]

We see that with broadcasting we can very easily compute the product of all possible combinations of entries in x and y. By only using ufuncs in our numerical code, we enable similar approaches, for example, to compute the ampacity rating for all possible combinations of weather- and line parameters.

Automatic testing of mathematical code

Okay, we got the code architecture sorted and the code super fast, so now we only need tests. Fortunately, another benefit of using pure functions for the equations is that they become straightforward to test.

Here, we will discuss two types of tests that we can make for our code: acceptance tests and property-based unit tests. The acceptance test defines a minimal set of requirements for the implementation and is provided by the standard. In CIGRE 601, we can find these under Annex E. Examples of calculation.

The more interesting part is the property-based unit tests. Property-based testing is a beautiful way to write elegant tests on the form:

given input that satisfies X, output should satisfy Y.

In Python, we use the Hypothesis library for property based testing, and I think the best way to understand it is with an example.

Consider the sorted-function. Once we sort a list, its numbers should be in ascending order, which means that lst[i] <= lst[i+1], and a test for this property could look like this:

Python
from itertools import pairwise
import hypothesis.strategies as st
from hypothesis import given


@given(lst=st.lists(st.floats()))
def test_sorted_is_ascending(lst):
    lst = sorted(lst)
    for l1, l2 in pairwise(lst):
        assert l1 <= l2

Do you think this test will pass? Click here to try it out and see for yourself!

In our case, we have mathematical equations that we want to test. And the nice thing with mathematical equations is that they have a vast amount of properties we can test! For example, the equation to compute the convective cooling, $P_c$, states that

$$ P_c = \pi \lambda_f (T_s – T_a) \text{Nu}, $$

where $\lambda_f$ is the thermal conductivity of air, $\text{Nu}$ is the Nusselt number, and $T_s$ and $T_a$ is the temperature of the conductor (surface) and air, respectively. We know that this function is linear with respect to $\lambda_f$. Moreover, if we set $T_s – T_a = \text{Nu} = 1$, then $P_c = \pi \lambda_f$, which gives us the following property-based test

Python
@given(
    thermal_conductivity_of_air=st.floats(
        min_value=0,
        max_value=1e10,
        allow_nan=False,
    )
)
def test_convective_cooling_scales_linearly_with_thermal_conductivity_of_air(
    thermal_conductivity_of_air,
):
    Nu = 1
    T_s = 1
    T_a = 0
    lambda_f = thermal_conductivity_of_air
    P_c = convective_cooling.compute_convective_cooling(T_s, T_a, Nu, lambda_f)
    
    assert P_c == np.pi * approx(lambda_f)

Here, we let Hypothesis select arbitrary values for $\lambda_f$, and only test whether whether $P_c = \pi \ \lambda_f$ when $\text{Nu} = T_s – T_a = 1$. If we didn’t use Hypothesis, we would either have to hard-code the test cases (e.g. with pytest.mark.parametrize) or iterate over random values for $\lambda_f$, which would make it less obvious what the code actually tests.

Style considerations

Selecting good variable names

One of the first style considerations we encounter when implementing mathematical equations with code is what we should name our variables. Do we use mathematical symbols or semantic names for the variable names?

Recall the equation for the convective cooling:

$$ P_c = \pi \lambda_f (T_s – T_a) \text{Nu}. $$

With Python, we could write the following equivalent function:

Python
def compute_convective_cooling(
    lambda_f: float,
    T_s: float,
    T_a: float,
    Nu: float,
) -> float:
    r"""Compute the convective cooling of the conductor.

    Equation (17) on page 24 of :cite:p:`cigre601`.
    """
    return pi * lambda_f * (T_s - T_a) * Nu

However, while this makes it easy to see how the mathematics works, it could be more user-friendly. For example, do you remember what lambda_f represents?

We could, of course, switch it around and use argument names with semantic meaning:

Python
def compute_convective_cooling(
    thermal_conductivity_of_air: float,
    surface_temperature: float,
    air_temperature: float,
    nusselt_number: float,
) -> float:
    r"""Compute the convective cooling of the conductor.

    Equation (17) on page 24 of :cite:p:`cigre601`.
    """
    return pi * thermal_conductivity_of_air * (surface_temperature - air_temperature) * nusselt_number

but then it becomes challenging to connect the code and the equations.

To get code where the variables both have semantic meaning, and the equations are easy to follow, we can combine these two approaches:

Python
def compute_convective_cooling(
    thermal_conductivity_of_air: float,
    surface_temperature: float,
    air_temperature: float,
    nusselt_number: float,
) -> float:
    r"""Compute the convective cooling of the conductor.

    Equation (17) on page 24 of :cite:p:`cigre601`.
    """
    lambda_f = thermal_conductivity_of_air
    T_s = surface_temperature
    T_a = air_temperature
    Nu = nusselt_number
    return pi * lambda_f * (T_s - T_a) * Nu

This way, the API clarifies what we need to know to compute the convective cooling, while it is also easy to compare the code with the reference equations.

Selecting good function names

One of my main pet peeves when I deal with mathematical code is function names. All functions (unless you have a very good reason) should have a verb in their name. After all, a function does do something, so why not explicitly state what it does?

In mathematical code, naming functions with verbs is even more critical. Take, for example, the convective_cooling function we defined above. The function is very readable. We see what the inputs are, what it computes and how. Still, if we want to use this function to compute some convective cooling, we should also give that variable a good name. Any suggestions for a good name? I would probably want to go with convective_cooling. But wait, that’s the name of the function, so we might end up writing

Python
convective_cooling = convective_cooling([...])

This is not ideal.

What is the solution to the naming problem? Simply add a verb to the function name. For example, “compute”:

Python
def compute_convective_cooling(
    thermal_conductivity_of_air,
    surface_temperature,
    air_temperature,
    nusselt_number,
):
    r"""Compute the convective cooling of the conductor.

    Equation (17) on page 24 of :cite:p:`cigre601`.
    """
    lambda_f = thermal_conductivity_of_air
    T_s = surface_temperature
    T_a = air_temperature
    Nu = nusselt_number
    
    return lambda * (T_s - T_a) * Nu


convective_cooling = compute_convective_cooling(
    thermal_conductivity_of_air=thermal_conductivity_of_air,
    surface_temperature=surface_temperature,
    air_temperature=air_temperature,
    nusselt_number=nusselt_number,
)

Sure, the function names get longer, but the code also becomes much more pleasant to read!

Using typing.annotated to provide units

A handy new Python feature (3.9) is annotated types. Annotated types allow us to provide extra information to our type hints. This is very useful for physical equations, where we can use the annotations to represent units. In linerate, we do this in the units.py-file

Python
from typing import Union

try:
    from typing import Annotated
except ImportError:  # Python version <3.9
    from typing_extensions import Annotated

import numpy as np
import numpy.typing as npt

FloatOrFloatArray = Union[float, np.float64, npt.NDArray[np.float64]]

OhmPerMeter = Annotated[FloatOrFloatArray, "Ω/m"]
Ampere = Annotated[FloatOrFloatArray, "A"]
Radian = Annotated[FloatOrFloatArray, "rad"]

We use these annotations in the functions that implement the mathematical equations:

def compute_joule_heating(current: Ampere, resistance: OhmPerMeter) -> WattPerMeter:
    I = current  # noqa
    R = resistance
    return I * I * R

Using annotations like this makes it extraordinarily easy to know what units we use in the different functions. But unfortunately, the units are not type-checked, so we can still write the wrong units by accident. I’m hopeful that a package like Pint can provide static checking of units in the future, but that is not possible (at least for now).

Extensive use of documentation strings

I believe docstrings are particularly important when we implement mathematical equations. A bare minimum is listing the source of the equations (e.g. Equation (17) on page 24 of CIGRE TB 601.). Even better, if you’re using Sphinx for the documentation, you can include the sphinxcontrib.bibtex plugin, which enables you to link the documentation against a BibTeX database (e.g. Equation (17) on page 24 of :cite:p:`cigre601`.).

After listing which equation(s) each function represents, we may want also to describe the notation. I like using the function arguments and returns field of the docstring for this, preferably formatted with LaTeX and using either the NumPy docstring style or the Google docstring style. A nice benefit we get by including type hints is that Sphinx can automatically parse these to provide the types in the docstring, also. To enable this, set

Python
autodoc_typehints = "description"
autodoc_typehints_description_target = "documented"

in the Sphinx conf.py file. Unfortunately, I have yet to figure out how to get it to include the return types properly, but at least we don’t need to worry about the input types.

A final and somewhat optional step is to also include information about the quantities in the docstring so that the user does not need to look at the references. If we combine all these steps, we get a docstring kind of like the one below

Python
def compute_convective_cooling(
    surface_temperature: Celsius,
    air_temperature: Celsius,
    nusselt_number: Unitless,
    thermal_conductivity_of_air: WattPerMeterPerKelvin,
) -> WattPerMeter:
    r"""Compute the convective cooling of the conductor.

    Equation (17) on page 24 of :cite:p:`cigre601`, given by
    
    .. math::
    
        P_c = \lambda_f \ Nu \ (T_a - T_s).

    Parameters
    ----------
    surface_temperature:
        :math:`T_s~\left[^\circ\text{C}\right]`. The conductor surface temperature.
    air_temperature:
        :math:`T_a~\left[^\circ\text{C}\right]`. The ambient air temperature.
    nusselt_number:
        :math:`Nu`. The nusselt number.
    thermal_conductivity_of_air:
        :math:`\lambda_f~\left[\text{W}~\text{m}^{-1}~\text{K}^{-1}\right]`. The thermal
        conductivity of air at the given temperature.

    Returns
    -------
    Union[float, float64, ndarray[Any, dtype[float64]]]
        :math:`P_c~\left[\text{W}~\text{m}^{-1}\right]`. The convective cooling of the conductor.
        Either due to wind, or passive convection, whichever is largest.
    """

Closing notes

There is no single solution to writing maintainable and readable mathematical code. This blog post demonstrates the choices behind linerate, but I am sure there are ways to improve these decisions. My main pain point is units, and for a future project, I hope to check if Pint might be part of a solution to ensuring consistent type usage. Moreover, I would test the code more thoroughly if given infinite time. For example, instead of checking that a function is linear by fixing all parameters but one, we could numerically differentiate it and check that the derivative is constant almost everywhere. If you have any ideas for writing maintainable and readable mathematical code, please share them in the comments below!

%d bloggers like this: