Architecture and API Design#

Understanding the internal architecture of ModernDiD will help you add new estimators or extend existing functionality. The patterns described here ensure your code integrates seamlessly with the rest of the package.

Overview#

ModernDiD is organized around a shared core and specialized estimator modules. The core module provides infrastructure that all estimators rely on, including a preprocessing pipeline for validating and transforming input data, configuration dataclasses that define estimator parameters, shared formatting utilities for structured table output, and base utilities for result objects. Each estimator module (did, drdid, didcont, didtriple) contains the statistical implementation for its methodology, following the patterns established in core. Each module also has a format.py that registers display formatting for its result objects. The didhonest module provides sensitivity analysis that works with results from any estimator, and the plots module offers unified visualization across all result types.

This architecture means that adding a new estimator involves implementing the statistical logic while reusing the preprocessing, result handling, formatting, and plotting infrastructure.

The Preprocessing Pipeline#

All estimators in ModernDiD use a shared preprocessing pipeline built on the builder pattern. The PreprocessDataBuilder provides a single entry point for data validation and transformation, with module-specific behavior controlled by configuration classes, validators, and transformers. Each module defines its own config class (e.g., DIDConfig, TwoPeriodDIDConfig, DDDConfig, DIDInterConfig) and data container (e.g., DIDData, DDDData), but the builder itself is shared. The builder selects the appropriate validators and transformers based on the config type, so module-specific requirements like partition-based subgroup creation or time-varying treatment handling are expressed through the pipeline rather than bypassing it.

Configuration Classes#

Each estimator type has a configuration dataclass that inherits from BasePreprocessConfig. These classes serve as both input validation specifications and metadata containers. They define the expected parameters and their types, and after preprocessing completes, computed values like time_periods and treated_groups are stored on the config for use during estimation.

from dataclasses import dataclass
from moderndid.core.preprocess.config import BasePreprocessConfig

@dataclass
class DIDConfig(BasePreprocessConfig):
    yname: str
    tname: str
    gname: str
    idname: str | None = None
    xformla: str = "~1"
    panel: bool = True
    allow_unbalanced_panel: bool = True
    weightsname: str | None = None
    control_group: ControlGroup = ControlGroup.NEVER_TREATED
    est_method: EstimationMethod = EstimationMethod.DOUBLY_ROBUST
    base_period: BasePeriod = BasePeriod.VARYING
    anticipation: int = 0
    alp: float = 0.05
    boot: bool = False
    biters: int = 1000
    clustervars: list[str] = field(default_factory=list)

Data Container Classes#

Preprocessed data is stored in dataclass containers that provide a clean interface for estimators. These containers bundle all the data an estimator needs, including the processed DataFrame, precomputed tensors for efficient computation, and the configuration that produced them. Properties on the container provide convenient access to common queries about the data structure.

@dataclass
class DIDData:
    data: pl.DataFrame              # Processed data in long format
    time_invariant_data: pl.DataFrame  # Unit-level data
    weights: np.ndarray             # Observation weights
    outcomes_tensor: list[np.ndarray] | None  # For balanced panels
    covariates_matrix: np.ndarray | None
    covariates_tensor: list[np.ndarray] | None
    config: DIDConfig
    cluster: np.ndarray | None

    @property
    def is_panel(self) -> bool: ...

    @property
    def is_balanced_panel(self) -> bool: ...

    @property
    def has_covariates(self) -> bool: ...

The Builder Pattern#

The PreprocessDataBuilder class provides a fluent interface for the preprocessing pipeline. The configuration object passed to the builder determines which validators and transformers run, so the caller only needs to construct the right config and the builder handles the rest.

from moderndid.core.preprocess import PreprocessDataBuilder, DIDConfig

config = DIDConfig(
    yname="outcome",
    tname="period",
    idname="unit_id",
    gname="treatment_group",
)

preprocessed = (
    PreprocessDataBuilder()
    .with_data(raw_dataframe)
    .with_config(config)
    .validate()      # Runs all validators
    .transform()     # Applies all transformations
    .build()         # Returns DIDData
)

Constants and Enums#

Categorical parameters use string enums for type safety. Enums prevent typos and enable IDE autocompletion while still accepting string values from users. When a user passes est_method="dr", the enum automatically converts it to the proper type.

from enum import Enum

class ControlGroup(str, Enum):
    NEVER_TREATED = "nevertreated"
    NOT_YET_TREATED = "notyettreated"

class EstimationMethod(str, Enum):
    DOUBLY_ROBUST = "dr"
    IPW = "ipw"
    REGRESSION = "reg"

class BasePeriod(str, Enum):
    UNIVERSAL = "universal"
    VARYING = "varying"

Consistent Argument Naming#

All estimators follow consistent naming conventions for parameters. This makes the API predictable and reduces cognitive load for users. When you learn that yname specifies the outcome column in one estimator, you know it means the same thing in every other estimator.

Column Name Parameters#

These parameters specify column names in the input data:

  • yname — outcome/dependent variable column

  • tname — time period column

  • idname — unit/entity identifier column (required for panel data)

  • gname — treatment group column (first period treated, 0 for never-treated)

  • dname — dose/treatment intensity column (continuous treatment only)

  • pname — partition/eligibility indicator (triple differences only)

  • treatname — binary treatment indicator (two-period models only)

  • weightsname — sampling weights column

Estimation Parameters#

  • xformla — covariate formula in Wilkinson notation (e.g., "~ x1 + x2")

  • est_method — estimation method: "dr", "ipw", or "reg"

  • control_group — control group: "nevertreated" or "notyettreated"

  • base_period — base period: "varying" or "universal"

  • anticipation — number of anticipation periods (default: 0)

  • panel — whether data is panel (True) or repeated cross-section (False)

  • allow_unbalanced_panel — whether to allow unbalanced panel data

Inference Parameters#

  • alp — significance level (default: 0.05)

  • boot — whether to use bootstrap inference

  • biters — number of bootstrap iterations (default: 1000)

  • boot_type — bootstrap type: "weighted" or "multiplier"

  • cband — whether to compute simultaneous confidence bands

  • clustervars — variables for clustered standard errors

  • random_state — random seed for reproducibility

Result Object Design#

All estimators return immutable NamedTuple result objects. Immutability ensures results cannot be accidentally modified during interactive analysis. The explicit attribute definitions with type hints create a clear interface that documents exactly what each estimator returns. Because NamedTuples are also tuples, results can be unpacked for quick access when only a few values are needed.

Docstring and Field Documentation#

Result class docstrings should begin with "Container for ...". Each field must have a #: doc comment immediately above it. Sphinx autodoc uses these comments for the autosummary tables and individual attribute pages in the API docs. Without them, the generated pages show Python’s default "Alias for field number N" instead of a meaningful description.

The class docstring should also include a numpydoc Attributes section with full descriptions. The #: comments drive the short summaries shown in the API reference tables, while the Attributes section provides the detailed descriptions rendered on the class page itself.

Core Result Attributes#

All result objects include these core attributes:

from typing import NamedTuple

class MPResult(NamedTuple):
    """Container for group-time average treatment effect results.

    Attributes
    ----------
    groups : ndarray
        Which group (defined by period first treated) each ATT is for.
    times : ndarray
        Which time period each group-time ATT is for.
    ...
    """

    #: Which group (defined by period first treated) each ATT is for.
    groups: np.ndarray
    #: Which time period each group-time ATT is for.
    times: np.ndarray
    #: Group-time average treatment effects.
    att_gt: np.ndarray

    #: Standard errors for group-time ATTs.
    se_gt: np.ndarray
    #: Critical value for confidence intervals.
    critical_value: float
    #: Analytical variance-covariance matrix.
    vcov_analytical: np.ndarray | None

    #: Influence function for estimating group-time ATTs.
    influence_func: np.ndarray | None

    #: DID estimation parameters.
    estimation_params: dict
    #: Number of unique cross-sectional units.
    n_units: int

Influence Functions#

Influence functions are a key architectural element that enable several important capabilities. Analytical standard errors can be computed from the influence function without resampling, which is faster and often more stable than bootstrap methods. When bootstrap inference is needed, multiplier bootstrap uses influence functions for efficient resampling without refitting the model. Aggregation of group-time effects into summary measures uses influence functions to propagate uncertainty correctly. The honest_did sensitivity analysis function uses influence functions to construct robust confidence intervals under violations of parallel trends.

All estimators should compute and return influence functions when possible. This is not optional infrastructure but a core part of what makes the downstream analysis tools work correctly.

Metadata Preservation#

Every result object includes an estimation_params dictionary containing all parameters used in estimation. This enables reproducibility because users can see exactly what was computed. Downstream analysis functions like aggregation need to know the original setup, such as whether bootstrap was used and what the significance level was. Having all relevant information in one place also simplifies debugging when results are unexpected.

Maketables Functionality#

ModernDiD result objects expose the maketables plug-in extractor interface so that maketables.ETable can consume them directly, with no coupling between packages.

The Plug-In Interface#

maketables.ETable discovers result data through a set of standard attributes. Every ModernDiD result class implements them.

Attribute

Description

__maketables_coef_table__

Property. Pandas DataFrame indexed by coefficient name with columns b, se, t, p, and optionally ci95l/ci95u and ci90l/ci90u.

__maketables_stat__(key)

Method. Returns a model-level statistic by key ("N", "se_type", "control_group", "aggregation", etc.).

__maketables_depvar__

Property. Dependent variable label for column headers.

__maketables_fixef_string__

Property. Fixed-effects specification string, or None.

__maketables_vcov_info__

Property. Dict with "vcov_type" and "clustervar" keys.

__maketables_stat_labels__

Property (optional). Maps stat keys to display labels.

__maketables_default_stat_keys__

Property (optional). Stat keys shown when the user does not specify model_stats.

Shared Helpers#

moderndid.core.maketables provides helpers that all result classes share so that table-building logic is not duplicated across modules.

Category

Functions

Coefficient tables

build_coef_table_with_ci, build_single_coef_table

Row labels

make_effect_names, make_group_time_names

Stat labels

se_type_label, control_group_label, est_method_label, vcov_info_from_bootstrap

A typical implementation on a result class looks like this.

@property
def __maketables_coef_table__(self):
    names = ["Overall ATT"]
    estimates = [self.overall_att]
    se = [self.overall_se]

    if self.event_times is not None:
        names.extend(make_effect_names(self.event_times, prefix="Event"))
        estimates.extend(self.att_by_event.tolist())
        se.extend(self.se_by_event.tolist())

    return build_coef_table_with_ci(names, estimates, se, alpha=0.05)

Adding Maketables Support to a New Estimator#

See Step 9: Add Maketables Support in the new estimator guide for a full implementation walkthrough covering all seven interface components, the shared helpers, and label mapping conventions.

Implementation Standards#

ModernDiD uses specific libraries for data handling and performance-critical code. Following these standards ensures consistency and maintains the performance characteristics users expect.

Polars for Data Handling#

All internal data manipulation uses Polars rather than pandas. Polars provides better performance for the operations common in DiD estimation and has a more consistent API. When users pass any Arrow-compatible DataFrame (polars, pandas, pyarrow, duckdb, etc.), the preprocessing pipeline converts it to Polars at the boundary via the Arrow PyCapsule Interface using narwhals, and all subsequent operations work with Polars DataFrames.

If you are adding new preprocessing logic or data transformations within an estimator, use Polars operations. Avoid converting back to pandas for intermediate steps. The moderndid.core.dataframe module provides to_polars() for converting user input at the API boundary.

Common Polars patterns used throughout the codebase:

import polars as pl
from moderndid.core.dataframe import to_polars

# Convert data to Polars
df = to_polars(data)

# Filter rows
treated = df.filter(pl.col("group") > 0)
balanced_ids = counts.filter(pl.col("len") == n_periods)

# Add or modify columns
df = df.with_columns(pl.Series(name="_weights", values=weights))
df = df.with_columns(pl.col("group").cast(pl.Float64))

# Conditional expressions
df = df.with_columns(
    pl.when(pl.col("group") == 0)
    .then(pl.lit(float("inf")))
    .otherwise(pl.col("group"))
    .alias("group")
)

# Group operations
counts = df.group_by("unit_id").len()
complete_ids = counts.filter(pl.col("len") == n_periods)["unit_id"].to_list()

# Window operations for panel data
df = df.with_columns(
    (pl.col("outcome") - pl.col("outcome").shift(1).over("unit_id")).alias("dy")
)

Numba for Performance#

Computationally intensive operations should use Numba JIT compilation rather than pure Python loops. Numba compiles Python functions to machine code, providing performance comparable to C while keeping the code readable. This is particularly important for bootstrap procedures where the same computation runs thousands of times, and for any code that runs in tight loops over large arrays.

Consider computing bootstrap standard errors for group-time effects. Each bootstrap iteration requires resampling clusters and recomputing weighted means. In pure Python, this nested loop over bootstrap iterations and clusters is painfully slow.

import numpy as np

# Pure Python - slow
def bootstrap_means_python(data, weights, cluster_ids, n_boot):
    n_clusters = len(np.unique(cluster_ids))
    unique_clusters = np.unique(cluster_ids)
    results = np.zeros(n_boot)

    for b in range(n_boot):
        sampled = np.random.choice(unique_clusters, size=n_clusters, replace=True)
        total = 0.0
        weight_sum = 0.0
        for c in sampled:
            mask = cluster_ids == c
            total += np.sum(data[mask] * weights[mask])
            weight_sum += np.sum(weights[mask])
        results[b] = total / weight_sum

    return results

With 1000 bootstrap iterations and 500 clusters, this function spends most of its time in Python’s interpreter rather than doing actual computation. Numba eliminates this overhead by compiling the function to machine code.

import numba as nb

@nb.njit(cache=True, parallel=True)
def bootstrap_means_numba(data, weights, cluster_ids, n_boot, seed):
    n_clusters = len(np.unique(cluster_ids))
    unique_clusters = np.unique(cluster_ids)
    results = np.zeros(n_boot)

    for b in nb.prange(n_boot):
        np.random.seed(seed + b)
        sampled = np.random.choice(unique_clusters, size=n_clusters, replace=True)
        total = 0.0
        weight_sum = 0.0
        for c in sampled:
            for i in range(len(cluster_ids)):
                if cluster_ids[i] == c:
                    total += data[i] * weights[i]
                    weight_sum += weights[i]
        results[b] = total / weight_sum

    return results

The Numba version uses nb.prange instead of range for the outer loop, enabling automatic parallelization across CPU cores. The cache=True argument stores the compiled function on disk, avoiding recompilation on subsequent runs. Speedups vary depending on the workload and data size, but can be substantial for the nested loops common in bootstrap procedures.

The pattern used throughout ModernDiD defines a pure Python/NumPy fallback first, then conditionally overrides it with a Numba-compiled version. This ensures the code works even when Numba is not installed.

try:
    import numba as nb
    HAS_NUMBA = True
except ImportError:
    HAS_NUMBA = False
    nb = None


def _compute_impl(data, weights):
    # Pure NumPy fallback
    return np.sum(data * weights) / np.sum(weights)


if HAS_NUMBA:

    @nb.njit(cache=True)
    def _compute_impl(data, weights):
        total = 0.0
        weight_sum = 0.0
        for i in range(len(data)):
            total += data[i] * weights[i]
            weight_sum += weights[i]
        return total / weight_sum

For element-wise operations on arrays, guvectorize provides a cleaner interface than writing explicit loops. It defines a generalized ufunc that NumPy can broadcast automatically. The signature specifies input and output array shapes.

from numba import float64, guvectorize

# Pure Python/NumPy fallback
def _safe_divide_impl(x, y, out=None):
    if out is None:
        out = np.zeros_like(x, dtype=float)
    mask = np.abs(y) >= 1e-10
    np.divide(x, y, out=out, where=mask)
    out[~mask] = 0.0
    return out


if HAS_NUMBA:

    @guvectorize(
        [(float64[:], float64[:], float64[:])],
        "(n),(n)->(n)",
        nopython=True,
        cache=True,
    )
    def _safe_divide_impl(x, y, result):
        for i in range(x.shape[0]):
            if np.abs(y[i]) < 1e-10:
                result[i] = 0.0
            else:
                result[i] = x[i] / y[i]

The signature "(n),(n)->(n)" means the function takes two arrays of the same length and produces an output array of that length. Numba handles memory allocation and broadcasting, so the compiled function works seamlessly with NumPy’s array operations.

When writing Numba functions, avoid Python objects and stick to NumPy arrays and scalar types. Numba works best with simple numerical code. If you need complex logic, keep it in pure Python and only JIT-compile the hot loops.

Thread-Based Parallelism#

Estimators that loop over group-time cells can use the parallel_map utility in moderndid.core.parallel to distribute work across threads. Threads work well here because the per-cell computation is dominated by NumPy, SciPy, and statsmodels C extensions that release the GIL. This avoids the serialization overhead of multiprocessing while still achieving concurrency.

from moderndid.core.parallel import parallel_map

# Build a list of argument tuples, one per group-time cell
args_list = [
    (group_idx, time_idx, data)
    for group_idx in range(n_groups)
    for time_idx in range(n_times)
]

# Run sequentially or in parallel depending on n_jobs
results = parallel_map(estimate_single_cell, args_list, n_jobs=n_jobs)

The n_jobs parameter follows scikit-learn conventions: 1 runs sequentially, -1 uses all available cores, and any value > 1 uses that many worker threads. Expose n_jobs as a parameter on your estimator function with a default of 1 so that sequential execution remains the default and parallelism is opt-in.

CuPy GPU Acceleration#

On machines with NVIDIA GPUs, ModernDiD can offload regression and propensity score estimation to the GPU via CuPy. All GPU-related code lives in the moderndid/cupy/ module, which provides three files:

  • backend.pyget_backend(), set_backend(), use_backend(), to_device(), to_numpy() for switching between NumPy and CuPy array libraries. The active backend is stored in a ContextVar, so use_backend() scopes the override to a block and reverts automatically.

  • regression.pycupy_wls (weighted least squares) and cupy_logistic_irls (logistic regression via IRLS), implemented with the current backend’s array operations.

  • bootstrap.py — GPU-accelerated multiplier bootstrap and cluster aggregation helpers.

Estimators use a dispatch pattern that checks the active backend at runtime. When the backend is CuPy, arrays are moved to the GPU with to_device() and regression is performed with cupy_wls or cupy_logistic_irls. Results are moved back to NumPy with to_numpy() before being returned. When the backend is NumPy (the default), the standard statsmodels path runs with no overhead.

from moderndid.cupy.backend import get_backend, to_numpy
from moderndid.cupy.regression import cupy_wls

xp = get_backend()
if xp is not np:
    beta, fitted = cupy_wls(xp.asarray(y), xp.asarray(X), xp.asarray(w))
    params = to_numpy(beta)
else:
    result = sm.WLS(y, X, weights=w).fit()
    params = result.params

Users enable GPU acceleration by installing the gpu extra (uv pip install moderndid[gpu]) and either passing backend="cupy" to att_gt/ddd or calling set_backend("cupy") before running an estimator. The gpu extra is not included in all because it requires CUDA hardware.

The Formatting System#

Every result object in ModernDiD has a formatted __repr__ and __str__ that produces structured table output when users call print() on a result. This is implemented through a formatting layer in moderndid/core/format.py and per-module format.py files.

Format Registration#

The core mechanism is attach_format, which monkey-patches __repr__ and __str__ onto a result class so that printing it produces formatted output instead of raw NamedTuple contents.

from moderndid.core.format import attach_format

# At module level in myestimator/format.py
attach_format(MyEstimatorResult, format_my_result)

After this call, print(result) will call format_my_result(result) instead of showing the default NamedTuple representation.

Shared Format Helpers#

The moderndid.core.format module provides helper functions that produce consistent output across all estimators. Format functions should compose these helpers rather than building strings from scratch.

The helpers fall into two categories:

  • Structureformat_title, format_section_header, format_footer, format_significance_note, and adjust_separators produce the framing elements (title blocks, section dividers, footers) that wrap every printed result.

  • Tablesformat_single_result_table, format_event_table, format_group_time_table, and format_horizon_table build the data tables for different result layouts (single-row summaries, event studies, group-time matrices, horizon-indexed effects).

All table helpers use prettytable with the SINGLE_BORDER style internally.

Writing a Format Function#

A format function takes a result object and returns a string. The typical structure is title block, summary table, detail table, significance note, metadata sections, and footer.

# myestimator/format.py
from scipy import stats

from moderndid.core.format import (
    attach_format,
    format_event_table,
    format_footer,
    format_section_header,
    format_significance_note,
    format_single_result_table,
    format_title,
)

from .results import MyEstimatorResult


def format_my_result(result):
    lines = []

    lines.extend(format_title("My Estimator Results"))

    alpha = result.estimation_params.get("alp", 0.05)
    conf_level = int((1 - alpha) * 100)
    z_crit = stats.norm.ppf(1 - alpha / 2)

    lci = result.overall_att - z_crit * result.overall_se
    uci = result.overall_att + z_crit * result.overall_se

    lines.extend(
        format_single_result_table(
            "ATT", result.overall_att, result.overall_se,
            conf_level, lci, uci,
        )
    )

    lines.extend(format_significance_note())

    lines.extend(format_section_header("Data Info"))
    lines.append(f" Observations: {result.n_units}")

    lines.extend(format_footer("Reference: Author (Year)"))
    return "\n".join(lines)


attach_format(MyEstimatorResult, format_my_result)

The attach_format call at module level means the format is registered as soon as the module is imported. Each estimator module’s __init__.py should import the format module to ensure registration happens at import time.

Creating a New Estimator#

See the dedicated Creating a New Estimator guide for a complete step-by-step walkthrough covering configuration, result objects, estimation, formatting, plotting, public API export, aggregation, maketables support, distributed backends, and testing.

Plotting Architecture#

The plotting module provides a unified interface for visualizing results from all estimators. This design means users learn one set of plotting functions that work across the entire package.

Converter Functions#

Each result type has a converter in moderndid/plots/converters.py that transforms it to a long-format Polars DataFrame suitable for plotting. Converters handle computing confidence intervals from standard errors and critical values, creating indicator columns for pre/post treatment periods, and dealing with missing values and edge cases.

All converters follow the naming convention {resulttype}_to_polars and produce DataFrames with standardized column names. Rows where the standard error is NaN (reference periods) must be filtered out:

df = df.filter(~pl.col("se").is_nan())
# plots/converters.py

def mpresult_to_polars(result: MPResult) -> pl.DataFrame:
    """Convert MPResult to plotting DataFrame."""
    ...

def aggteresult_to_polars(result: AGGTEResult) -> pl.DataFrame:
    """Convert AGGTEResult to plotting DataFrame."""
    ...

Plot Functions#

Plot functions accept result objects and delegate to the appropriate converter based on the result type. Adding support for a new result type requires two changes: writing a converter function in converters.py and adding an isinstance() branch in the relevant plot function in plots.py.

# plots/plots.py

def plot_event_study(result, show_ci=True, ref_line=0, ...):
    # Detect result type and use appropriate converter
    if isinstance(result, AGGTEResult):
        df = aggteresult_to_polars(result)
    elif isinstance(result, PTEResult):
        df = pteresult_to_polars(result)
    # ... build ggplot

Integration with Sensitivity Analysis#

The honest_did function performs sensitivity analysis on event study results. To make your estimator compatible with sensitivity analysis, ensure your aggregated result object implements the EventStudyProtocol. Any result object with the required attributes can be passed to honest_did, which makes it straightforward to add sensitivity analysis support for new estimators.

from typing import Protocol, runtime_checkable

@runtime_checkable
class EventStudyProtocol(Protocol):
    aggregation_type: str
    influence_func: np.ndarray | None
    event_times: np.ndarray | None
    att_by_event: np.ndarray | None
    estimation_params: dict