KNNGroupedFixedEffects

Guide & worked example · Source

KNNGroupedFixedEffects(
    n_groups,
    n_neighbors=10,
    knn_block_size=2048,
    kmeans_niter=25,
    seed=0,
    use_flash_kmeans=True,
    standardize_embeddings=True,
    include_outcome=True,
    device=None,
)

Experimental grouped fixed-effects estimator with KNN-smoothed embeddings.

Parameters

Name	Type	Description	Default
n_groups	int	Number of latent groups used in the second-step grouped FE regression.	required
n_neighbors	int	Number of neighbors for embedding smoothing.	`10`
knn_block_size	int	Query block size for exact KNN evaluation.	`2048`
kmeans_niter	int	Maximum K-Means iterations.	`25`
seed	int	Random seed used by clustering.	`0`
use_flash_kmeans	bool	Whether to use flash-kmeans when importable.	`True`
standardize_embeddings	bool	Whether to z-score panel embeddings before KNN and K-Means.	`True`
include_outcome	bool	Whether to include the outcome in the default panel embeddings.	`True`

Attributes

Name	Description
`device`
`embeddings_`
`group_ids_`
`group_time_ids_`
`include_outcome`
`kmeans_niter`
`knn_block_size`
`linear_model_`
`n_groups`
`n_neighbors`
`neighbor_indices_`
`params`
`seed`
`smoothed_embeddings_`
`standardize_embeddings`
`time_levels_`
`unit_levels_`
`use_flash_kmeans`

Methods

Name	Description
fit	Estimate grouped fixed effects through KNN smoothing and clustered FE OLS.
predict	Delegate prediction to the second-step linear model.
summary	Print the second-step linear regression summary.
to	Move all model parameters to specified device.

fit

KNNGroupedFixedEffects.fit(
    X,
    y,
    unit_ids,
    time_ids,
    classification_features=None,
    classification_unit_ids=None,
    se=None,
)

Estimate grouped fixed effects through KNN smoothing and clustered FE OLS.

If classification_features has one row per unit, rows are assumed to be ordered by sorted(unique(unit_ids)) unless classification_unit_ids is provided. Passing classification_unit_ids is safer for shuffled or non-contiguous unit identifiers; features are then aligned internally to the fitted unit ordering.

predict

KNNGroupedFixedEffects.predict(X)

Delegate prediction to the second-step linear model.

summary

KNNGroupedFixedEffects.summary()

Print the second-step linear regression summary.

to

KNNGroupedFixedEffects.to(device)

Move all model parameters to specified device.

Parameters

Name	Type	Description	Default
device	torch.device or str	Target device (‘cuda’, ‘cpu’, etc.)	required

Returns

Name	Type	Description
	BaseEstimator	Self for method chaining