Function Design for Scientific Code

class DataPipeline:
    def __init__(self, data):
        self.data = data    # shared mutable state

    def remove_nans(self):
        self.data = self.data[~np.isnan(self.data)]

    def log_transform(self):
        self.data = np.log(self.data + 1e-10)

    def normalize(self):
        self.data = (self.data - self.data.mean()) / self.data.std()

# Problem: order matters, state is implicit, hard to test individually
pipe = DataPipeline(raw_data)
pipe.remove_nans()
pipe.log_transform()
pipe.normalize()
result = pipe.data

def remove_nans(x: np.ndarray) -> np.ndarray:
    return x[~np.isnan(x)]

def log_transform(x: np.ndarray, eps: float = 1e-10) -> np.ndarray:
    return np.log(x + eps)

def normalize(x: np.ndarray) -> np.ndarray:
    return (x - x.mean()) / x.std()

# Each function is independently testable
result = normalize(log_transform(remove_nans(raw_data)))

# Or compose them:
from functools import reduce
pipeline = [remove_nans, log_transform, normalize]
result = reduce(lambda data, fn: fn(data), pipeline, raw_data)

• Each function can be unit tested in isolation
• Functions can be reordered or swapped easily
• No hidden coupling through shared state
• Easy to parallelize (no race conditions on shared data)

import matplotlib.pyplot as plt
import numpy as np

def plot_results(
    *datasets: np.ndarray,
    labels: list[str] | None = None,
    title: str = "Simulation Results",
    xlabel: str = "Index",
    ylabel: str = "Value",
    **plot_kwargs,
) -> plt.Figure:
    """Plot one or more datasets on the same axes.

    Parameters
    ----------
    *datasets : np.ndarray
        One or more 1-D arrays to plot.
    labels : list[str] or None
        Legend labels for each dataset.
    title, xlabel, ylabel : str
        Axis labels and title.
    **plot_kwargs
        Forwarded to ``ax.plot()`` (e.g., linewidth, marker).

    Returns
    -------
    plt.Figure
        The created figure.
    """
    fig, ax = plt.subplots()
    for i, data in enumerate(datasets):
        label = labels[i] if labels else f"Dataset {i}"
        ax.plot(data, label=label, **plot_kwargs)
    ax.set(title=title, xlabel=xlabel, ylabel=ylabel)
    ax.legend()
    return fig

snr_values = np.linspace(0, 20, 50)
ber_zf = np.exp(-snr_values / 5)
ber_mmse = np.exp(-snr_values / 3)

fig = plot_results(
    ber_zf, ber_mmse,
    labels=["ZF", "MMSE"],
    title="BER Comparison",
    xlabel="SNR (dB)",
    ylabel="BER",
    linewidth=2,
    marker='o',
    markersize=3,
)

def resample_signal(
    signal: np.ndarray,
    original_rate: float,
    *,                           # everything after is keyword-only
    target_rate: float,
    method: str = "linear",
    preserve_energy: bool = True,
) -> np.ndarray:
    """Resample a signal to a new sample rate.

    Parameters
    ----------
    signal : np.ndarray
        Input signal.
    original_rate : float
        Original sample rate in Hz.
    target_rate : float
        Desired sample rate in Hz (keyword-only).
    method : str
        Interpolation method: 'linear', 'cubic', 'nearest'.
    preserve_energy : bool
        If True, scale amplitude to preserve signal energy.

    Returns
    -------
    np.ndarray
        Resampled signal.
    """
    ratio = target_rate / original_rate
    n_new = int(len(signal) * ratio)
    x_old = np.linspace(0, 1, len(signal))
    x_new = np.linspace(0, 1, n_new)
    resampled = np.interp(x_new, x_old, signal)
    if preserve_energy:
        resampled *= np.sqrt(1 / ratio)
    return resampled

# This works — clear and unambiguous:
resample_signal(sig, 44100, target_rate=22050, method="cubic")

# This raises TypeError — prevents positional mistakes:
resample_signal(sig, 44100, 22050, "cubic")
# TypeError: resample_signal() takes 2 positional arguments
#            but 4 were given