Skip to content

Instantly share code, notes, and snippets.

@jpivarski
Last active May 3, 2024 21:04
Show Gist options
  • Save jpivarski/cfe39deb0118bf1f36b0987fb06b3b3c to your computer and use it in GitHub Desktop.
Save jpivarski/cfe39deb0118bf1f36b0987fb06b3b3c to your computer and use it in GitHub Desktop.
Package to demonstrate testing
import numpy as np
class System:
G = 3
min_distance = 0.1
@classmethod
def random(cls, num_particles, num_dimensions, mass_mean, mass_width, x_width, p_width):
m = np.random.gamma(mass_mean / mass_width, mass_width, num_particles)
x = np.random.normal(0, x_width, (num_particles, num_dimensions))
p = np.random.normal(0, p_width, (num_particles, num_dimensions))
return cls(m, x, p)
def __init__(self, m, x, p):
self.x, self.p = np.broadcast_arrays(x, p)
assert self.x.shape == self.p.shape
if len(self.x.shape) != 2:
raise ValueError(f"arrays of position and momentum must each have 2 components, not {len(self.x.shape)}")
if self.num_dimensions < 2:
raise ValueError("number of dimensions must be at least 1")
self.m, _ = np.broadcast_arrays(m, self.x[:, 0])
assert len(self.m) == len(self.x)
if len(self.m.shape) != 1:
raise ValueError(f"array of masses must have only 1 component, not {len(self.m.shape)}")
if issubclass(self.m.dtype.type, np.integer):
self.m = self.m.astype(np.float64)
if issubclass(self.x.dtype.type, np.integer):
self.x = self.x.astype(np.float64)
if issubclass(self.p.dtype.type, np.integer):
self.p = self.p.astype(np.float64)
if not issubclass(self.m.dtype.type, np.floating):
raise TypeError(f"masses must have floating-point type, not {self.m.dtype}")
if not issubclass(self.x.dtype.type, np.floating):
raise TypeError(f"positions must have floating-point type, not {self.m.dtype}")
if not issubclass(self.p.dtype.type, np.floating):
raise TypeError(f"momenta must have floating-point type, not {self.m.dtype}")
self.history = [self.Step(0, self.x, self.p)]
class Step:
def __init__(self, t, x, p):
self.t, self.x, self.p = t, x, p
def __repr__(self):
return f"<Step t={self.t} x={self.x.tolist()} p={self.p.tolist()}>"
def __repr__(self):
return f"<System of {self.num_particles} particles in {self.num_dimensions} dimensions>"
@property
def num_particles(self):
return self.x.shape[0]
@property
def num_dimensions(self):
return self.x.shape[1]
@property
def forces(self):
# indexes to pick out particle 1 and particle 2
p1, p2 = np.triu_indices(len(self.x), k=1)
# pairwise (pw) displacements between all particle pairs
pw_displacement = self.x[p2] - self.x[p1]
# pairwise displacement is a sum in quadrature over all dimensions
pw_distance = np.maximum(np.sqrt(np.sum(pw_displacement**2, axis=-1)), self.min_distance)
# direction is a unit vector
pw_direction = pw_displacement / pw_distance[:, np.newaxis]
m1 = self.m[p1, np.newaxis]
m2 = self.m[p2, np.newaxis]
# 1/r in 2D, 1/r**2 in 3D, 1/r**3 in 4D...
power = self.num_dimensions - 1
# law of universal gravitation
pw_force = self.G * m1 * m2 * pw_direction / pw_distance[:, np.newaxis]**power
# vector sum over pairs for each particle, np.add.at inverts p1, p2 indexing
total_force = np.zeros_like(self.x)
np.add.at(total_force, p1, pw_force)
np.add.at(total_force, p2, -pw_force)
return total_force
def step(self, dt=0.1):
half_dt = dt / 2
# half kick: update p by a half time-step using current positions
self.p = self.p + self.forces * half_dt
# full drift: update x by a full time-step using new momenta
self.x = self.x + self.p * dt / self.m[:, np.newaxis]
# half kick: update p by another half time-step using new positions
self.p = self.p + self.forces * half_dt
# save the history
self.history.append(self.Step(self.history[-1].t + dt, self.x, self.p))
def steps(self, n, dt=0.01):
for _ in range(n):
self.step(dt=dt)
def plot(self, figsize=(5, 5), method="to_html5_video", num_frames=100, frame_ms=50):
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML
fig, ax = plt.subplots(figsize=(5, 5))
x = np.empty((len(self.history), self.num_particles, 2))
for i, step in enumerate(self.history):
for j in range(self.num_particles):
x[i, j, :] = step.x[j, :2]
x0 = np.mean(x[:, :, 0])
y0 = np.mean(x[:, :, 1])
scale = np.percentile(np.max(abs(x), axis=0), 75) * 1.5
ax.set(xlim=(x0 - scale, x0 + scale), ylim=(y0 - scale, y0 + scale))
x = x[::len(x) // num_frames]
lines = []
for j in range(self.num_particles):
lines.append(ax.plot(x[:1, j, 0], x[:1, j, 1])[0])
dots = ax.scatter(x[0, :, 0], x[0, :, 1], color="black")
def update(i):
for j, line in enumerate(lines):
line.set_xdata(x[:i, j, 0])
line.set_ydata(x[:i, j, 1])
dots.set_offsets(x[i, :, :])
return lines + [dots]
ani = animation.FuncAnimation(fig=fig, func=update, frames=num_frames, interval=frame_ms, blit=True)
out = HTML(getattr(ani, method)())
plt.close()
return out
# system = System([1, 1], [[0, -1, 0], [0, 1, 0]], [[1, 0, 0], [-1, 0, 0]])
# system.steps(300, dt=0.1)
# system.plot()
# system = System([1, 1], [[0, -1], [0, 1]], [[1, 0], [-1, 0]])
# system.steps(1500, dt=0.01)
# system.plot()
# system = System([1000, 1], [[0, 0, 0], [0, 10, 0]], [[0, 0, 0], [-10, 0, 0]])
# system.steps(300)
# system.plot()
# system = System([1000, 1, 1], [[0, 0, 0], [0, 10, 0], [0, 11, 0]], [[0, 0, 0], [-16, 0, 0], [-13, 0, 0]])
# system.steps(600)
# system.plot()
# p1 = 0.347111
# p2 = 0.532728
# system = System([1, 1, 1], [[-1, 0, 0], [1, 0, 0], [0, 0, 0]], [[p1, p2, 0], [p1, p2, 0], [-2*p1, -2*p2, 0]])
# system.G = 1
# system.steps(1000, dt=0.01)
# system.plot()
# system = System.random(20, 3, 10, 1, 100, 10)
# system.steps(10000, dt=0.1)
# system.plot()
# system = System.random(50, 3, 10, 1, 100, 15)
# system.steps(10000, dt=0.1)
# system.plot()

Give me an example of a miniature codebase that would be a good example to use to teach students about how to write unit tests?


Certainly! A simple calculator program is often a good choice to teach students about writing unit tests, as it involves clear, straightforward logic that's easy to follow and test. Here’s an example of such a codebase in Python along with some basic unit tests using unittest, a built-in testing framework.

Simple Calculator Code

# calculator.py

def add(x, y):
    return x + y

def subtract(x, y):
    return x - y

def multiply(x, y):
    return x * y

def divide(x, y):
    if y == 0:
        raise ValueError("Cannot divide by zero.")
    return x / y

Unit Tests for the Calculator

# test_calculator.py

import unittest
from calculator import add, subtract, multiply, divide

class TestCalculator(unittest.TestCase):

    def test_add(self):
        self.assertEqual(add(5, 3), 8)
        self.assertEqual(add(-1, 1), 0)
        self.assertEqual(add(-1, -1), -2)

    def test_subtract(self):
        self.assertEqual(subtract(10, 5), 5)
        self.assertEqual(subtract(-1, 1), -2)
        self.assertEqual(subtract(-1, -1), 0)

    def test_multiply(self):
        self.assertEqual(multiply(3, 7), 21)
        self.assertEqual(multiply(-1, 3), -3)
        self.assertEqual(multiply(0, 3), 0)

    def test_divide(self):
        self.assertEqual(divide(10, 2), 5)
        self.assertRaises(ValueError, divide, 10, 0)

        with self.assertRaises(ValueError):
            divide(10, 0)

if __name__ == '__main__':
    unittest.main()

Running the Tests

You can run these tests by executing the test_calculator.py script. This environment setup ensures students can learn how a function’s logic is translated into a test case and how different scenarios (like edge cases) are handled in unit testing.

Learning Points

  • Understanding how to write basic functions.
  • Writing unit tests using assertions to check for expected outcomes.
  • Testing for exceptions in scenarios that could result in runtime errors.
  • Integrating tests within a development workflow to check code before deployment.

No, it needs to be more involved than a calculator because with a calculator, the implementation and the test are the same calculation. I want to show some examples of property-based testing, such as invariants that are not manifest in the implementation, but would be in the tests. Also, I want to show the importance of not over-testing, to allow the codebase to grow in some ways while being fixed in others.


For a more complex example while introducing concepts of property-based testing and maintaining a balance in testing scope, consider a "String Manipulator" codebase. This example involves more complex operations than simple arithmetic, including rotations, reversals, and transformations, providing a nice platform to demonstrate invariants and properties in tests.

String Manipulator Code

# string_manipulator.py

def reverse_string(s):
    return s[::-1]

def capitalize_string(s):
    return s.capitalize()

def add_prefix(s, prefix):
    return f"{prefix}{s}"

def rotate_string(s, n):
    n = n % len(s)  # Prevent rotating more than string length
    return s[n:] + s[:n]

Property-Based Tests

Here, we'll use unittest along with the hypothesis library, which is designed for property-based testing in Python.

# test_string_manipulator.py

import unittest
from hypothesis import given, strategies as st
from string_manipulator import reverse_string, capitalize_string, add_prefix, rotate_string

class TestStringManipulator(unittest.TestCase):

    @given(st.text())
    def test_reverse_invariant(self, s):
        # Reversing twice should give the original string
        self.assertEqual(reverse_string(reverse_string(s)), s)

    @given(st.text())
    def test_capitalize_idempotent(self, s):
        # Capitalizing twice should be the same as capitalizing once
        self.assertEqual(capitalize_string(capitalize_string(s)), capitalize_string(s))

    @given(st.text(), st.text())
    def test_prefix_addition_length(self, s, prefix):
        # The length after adding a prefix should be sum of both lengths
        result = add_prefix(s, prefix)
        self.assertEqual(len(result), len(s) + len(prefix))

    @given(st.text(min_size=1), st.integers())  # Ensuring non-empty string for rotation to be meaningful
    def test_rotation_invariant(self, s, n):
        # A full rotation (len(s) times) should yield the original string
        rotated = rotate_string(s, n)
        self.assertEqual(rotate_string(rotated, len(s) - (n % len(s))), s)

if __name__ == '__main__':
    unittest.main()

Key Learning Points

  1. Property-Based Testing: Using properties or invariants, such as idempotence ("capitalizing twice is the same as capitalizing once") and specific behaviors after transformations (like reversing or rotating), which might not be directly apparent from just looking at the implementation.

  2. Testing Invariants: Ensuring fundamental truths like "reversing a string twice returns the original string" hold, which are crucial in many practical applications beyond the original scope of the function.

  3. Scope of Testing: The tests ensure correctness of specific behaviors without over-specifying how the implementation should achieve these results, thus allowing for internal implementation to change as long as the properties hold, illustrating how to avoid over-testing.

  4. Usage of Hypothesis: Demonstrates using an external library to generate test cases based on defined strategies, which can be more effective and thorough compared to manually written static tests.

This approach not only covers more complex testing scenarios, but also introduces tools and techniques useful in modern testing practices, reflecting real-world software development complexities.


Maybe a simulation? So that we can use physics invariants, like (approximate) energy conservation?


A simulation-based example does lend itself very well to teaching testing concepts, specifically through the use of physics invariants like energy conservation. We'll consider a basic particle physics simulation where particles in a system interact with each other, and we'll test for conservation of energy and momentum throughout the simulation.

Particle System Physics Simulation

Codebase Structure

The simulation includes:

  • Particle class with properties like mass, position, velocity.
  • Simulation class that manages multiple particles and updates their state.

Particle Class (particle.py)

import numpy as np

class Particle:
    def __init__(self, mass, position, velocity):
        self.mass = mass
        self.position = np.array(position)
        self.velocity = np.array(velocity)

    def update_position(self, dt):
        self.position += self.velocity * dt

Simulation Class (simulation.py)

from particle import Particle

class Simulation:
    def __init__(self):
        self.particles = []

    def add_particle(self, particle):
        self.particles.append(particle)

    def update(self, dt):
        for particle in self.particles:
            particle.update_position(dt)

    def total_energy(self):
        # Simplified kinetic energy calculation: 0.5 * m * v^2
        return sum(0.5 * p.mass * np.dot(p.velocity, p.velocity) for p in self.particles)

Property-Based Tests (test_simulation.py)

Using property-based testing to ensure that energy and momentum are conserved during simulations without external forces. We use unittest framework and hypothesis for generating test cases:

import unittest
from hypothesis import given, strategies as st
from particle import Particle
from simulation import Simulation
import numpy as np

class TestSimulation(unittest.TestCase):
    
    @given(st.lists(st.tuples(st.floats(min_value=0.1, max_value=10), 
                              st.lists(st.floats(min_value=-100, max_value=100), min_size=3, max_size=3), 
                              st.lists(st.floats(min_value=-100, max_value=100), min_size=3, max_size=3)),
                    min_size=1, max_size=10), 
           st.floats(min_value=0.01, max_value=1))
    def test_energy_conservation(self, particles_data, dt):
        # Initialize simulation with particles
        sim = Simulation()
        for mass, position, velocity in particles_data:
            sim.add_particle(Particle(mass, position, velocity))
        
        initial_energy = sim.total_energy()

        # Small simulation steps to approximate continuity
        sim.update(dt)
        
        final_energy = sim.total_energy()
        
        # Check if the initial and final energy are approximately equal
        np.testing.assert_allclose(initial_energy, final_energy, atol=1e-5)

if __name__ == '__main__':
    unittest.main()

Explanation

  1. Creating a Physics-Based Model: The simplified model considers kinetic energy of particles using their mass and velocity, handling changes in position over time.

  2. Property-Based Testing in Physics: We harnessed physics invariants to create robust tests. The test for energy conservation generates random particles with random mass, position, and velocity. It ensures that despite these random configurations, the system adheres to the conservation of energy law after updating.

  3. Use of Libraries: Utilized NumPy for numerical operations and Hypothesis for generating complex strategy-based test inputs, particularly useful for testing simulation-type applications.

This setup provides practical insights into testing simulations, emphasizing how well-designed unit tests can validate complex dynamics and invariants in a codebase, using physics rules as a fundamental guideline.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment