API Reference

API Reference#

This page provides detailed information about the main functions in lightweight-mcnnm.

Main Functions#

estimate#

mcnnm.wrappers.estimate(Y, Mask, X=None, Z=None, V=None, Omega=None, use_unit_fe=True, use_time_fe=True, lambda_L=None, lambda_H=None, n_lambda=10, max_iter=10000.0, tol=0.0001, validation_method='cv', K=5, initial_window=None, step_size=None, horizon=None, max_window_size=None)[source]#

Estimate the Matrix Completion with Nuclear Norm Minimization (MC-NNM) model.

This function performs the complete estimation process for the MC-NNM model, including parameter selection, model fitting, and treatment effect computation. It handles various input configurations and validation methods to provide a comprehensive analysis of the given data.

Detailed Process:

  1. Input Validation and Initialization: - Checks the dimensions of input matrices Y and W. - Initializes the inverse of the Omega matrix (Omega_inv) if provided, otherwise uses an identity matrix. - Validates that Omega_inv is positive definite. - Initializes matrices L, X_tilde, Z_tilde, and V using the initialize_matrices function. - Initializes fixed effects and H matrix using the initialize_fixed_effects_and_H function.

  2. Lambda Selection: - If lambda_L or lambda_H are not provided, performs validation to select optimal values:

    1. Cross-validation (CV): - Uses the cross_validate function to find optimal lambda values. - Performs K-fold cross-validation on the data.

    2. Holdout validation: - Uses the holdout_validate function to find optimal lambda values. - Requires initial_window, step_size, and horizon parameters. - Generates default values for these parameters if not provided. - Validates the holdout configuration using validate_holdout_config function.

    • If lambda_L and lambda_H are provided, uses these values directly.

  3. Model Fitting: - Calls the final_fit function to fit the MC-NNM model using the selected lambda values. - Uses a warm-start approach, fitting the model along a path of lambda values.

  4. Results Computation: - Computes the completed outcome matrix using the compute_Y_hat function. - Calculates the average treatment effect using the compute_treatment_effect function.

  5. Return Results: - Returns an MCNNMResults object containing the estimated parameters, completed matrix, and treatment effect.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T), where N is the number of units and T is the number of time periods.

  • Mask (Array) – The treatment assignment matrix of shape (N, T). Binary values where 1 indicates treatment and 0 indicates control.

  • X (Array, optional) – Unit-specific covariates of shape (N, P), where P is the number of unit-specific covariates.

  • Z (Array, optional) – Time-specific covariates of shape (T, Q), where Q is the number of time-specific covariates.

  • V (Array, optional) – Unit-time specific covariates of shape (N, T, J), where J is the number of unit-time specific covariates.

  • Omega (Array, optional) – The covariance matrix of shape (T, T) representing the time series correlation structure. If not provided, an identity matrix is used.

  • use_unit_fe (bool) – Whether to include unit fixed effects in the model. Default is True.

  • use_time_fe (bool) – Whether to include time fixed effects in the model. Default is True.

  • lambda_L (Scalar, optional) – The regularization parameter for the nuclear norm of L. If not provided, it will be selected via validation.

  • lambda_H (Scalar, optional) – The regularization parameter for the nuclear norm of H. If not provided, it will be selected via validation.

  • n_lambda (int) – The number of lambda values to consider in the validation grid. Default is 10.

  • max_iter (Scalar, optional) – Maximum number of iterations for the optimization algorithm. Default is 10,000.

  • tol (Scalar, optional) – Tolerance for convergence of the optimization algorithm. Default is 1e-4.

  • validation_method (str) – The method to use for selecting lambda values. Must be either ‘cv’ for cross-validation or ‘holdout’ for holdout validation. Default is ‘cv’.

  • K (int) – The number of folds to use in cross-validation. Default is 5.

  • initial_window (int, optional) – The size of the initial window for holdout validation. Required if validation_method is ‘holdout’.

  • step_size (int, optional) – The step size for moving the window in holdout validation. Required if validation_method is ‘holdout’.

  • horizon (int, optional) – The size of the holdout horizon in holdout validation. Required if validation_method is ‘holdout’.

  • max_window_size (int, optional) – The maximum window size for holdout validation. If not provided, no maximum is imposed.

Returns:

An object containing the following attributes:

  • tau: The estimated average treatment effect.

  • lambda_L: The final lambda_L value used in the model.

  • lambda_H: The final lambda_H value used in the model.

  • L: The estimated low-rank matrix.

  • Y_completed: The completed outcome matrix.

  • gamma: The estimated unit fixed effects (if use_unit_fe is True).

  • delta: The estimated time fixed effects (if use_time_fe is True).

  • beta: The estimated coefficients for covariates.

  • H: The estimated interactive fixed effects matrix.

Return type:

MCNNMResults

Raises:

ValueError – If Omega_inv is not positive definite, if an invalid validation method is specified, or if required parameters for holdout validation are missing.

Notes

  • The function uses JAX for numerical computations, which allows for automatic differentiation and potential GPU acceleration.

  • The estimation process involves several steps of initialization, validation, and optimization, which can be computationally intensive for large datasets.

  • The choice of validation method and associated parameters can significantly affect the final estimates and computational time.

  • The function handles missing data in the outcome matrix Y through the treatment assignment matrix W.

  • The implementation supports various types of fixed effects and covariates, allowing for flexible model specifications.

Example

>>> from mcnnm.wrappers import estimate
>>> from mcnnm.utils import generate_data
>>> N, T = 100, 50  # Number of units and time periods
>>> noise_scale = 1.0 # Scale of the noise in the data
>>> Y, W, X, Z, V, true_params = generate_data(
...     nobs=N,
...     nperiods=T,
...     unit_fe=True,
...     time_fe=True,
...     X_cov=True,
...     Z_cov=True,
...     V_cov=True,
...     seed=2024,
...     noise_scale=noise_scale,
... )
>>> results = estimate(Y, W, X, Z, V, validation_method='cv', K=5)
>>> print(results.tau)  # Print the estimated average treatment effect
>>> print(f"True ATE: {true_params['treatment_effect']}")  # Compare with the true ATE

See also

  • validation.cross_validate(): Function used for cross-validation.

  • validation.holdout_validate(): Function used for holdout validation.

  • validation.final_fit(): Function used for the final model fitting.

  • core.compute_Y_hat(): Function used to compute the completed outcome matrix.

  • core.compute_treatment_effect(): Function used to compute the average treatment effect.

complete_matrix#

mcnnm.wrappers.complete_matrix(Y, Mask, X=None, Z=None, V=None, Omega=None, use_unit_fe=True, use_time_fe=True, lambda_L=None, lambda_H=None, n_lambda=10, max_iter=10000, tol=0.0001, validation_method='cv', K=5, initial_window=None, step_size=None, horizon=None, max_window_size=None)[source]#

Complete a matrix using the Matrix Completion with Nuclear Norm Minimization (MC-NNM) model.

This function is a thin wrapper around the estimate function, focusing on matrix completion. It performs the estimation process for the MC-NNM model and returns the completed matrix along with the regularization parameters used.

For a detailed description of the estimation process, input parameters, and their meanings, please refer to the documentation of the estimate function.

Returns:

A tuple containing:

  • Y_completed (Array): The completed outcome matrix.

  • lambda_L (Scalar): The final lambda_L value used in the model.

  • lambda_H (Scalar): The final lambda_H value used in the model.

Return type:

tuple

Note

This function uses the same parameters as the estimate function. For a comprehensive explanation of each parameter, including optional covariates, fixed effects, validation methods, and other configuration options, please consult the estimate function’s documentation.

Example

>>> Y_completed, lambda_L, lambda_H = complete_matrix(Y, W, X=X, Z=Z, V=V)
>>> print(Y_completed.shape)  # Print the shape of the completed matrix
>>> print(f"Lambda L: {lambda_L}, Lambda H: {lambda_H}")  # Print the regularization parameters

See also

estimate: The main function performing the complete MC-NNM estimation process.

Parameters:

Classes#

MCNNMResults#

class mcnnm.wrappers.MCNNMResults(tau: float | int | Array | None = None, lambda_L: float | int | Array | None = None, lambda_H: float | int | Array | None = None, L: Array | None = None, Y_completed: Array | None = None, gamma: Array | None = None, delta: Array | None = None, beta: Array | None = None, H: Array | None = None)[source]

Bases: NamedTuple

A named tuple containing the results of the MC-NNM (Matrix Completion with Nuclear Norm Minimization) estimation.

This class encapsulates all the key outputs from the MC-NNM model, including the estimated treatment effect, selected regularization parameters, and various estimated matrices and vectors.

tau

The estimated average treatment effect.

Type:

Optional[Scalar]

lambda_L

The selected regularization parameter for the low-rank matrix L.

Type:

Optional[Scalar]

lambda_H

The selected regularization parameter for the covariate coefficient matrix H.

Type:

Optional[Scalar]

L

The estimated low-rank matrix.

Type:

Optional[Array]

Y_completed

The completed outcome matrix (including counterfactuals).

Type:

Optional[Array]

gamma

The estimated unit fixed effects.

Type:

Optional[Array]

delta

The estimated time fixed effects.

Type:

Optional[Array]

beta

The estimated unit-time specific covariate coefficients.

Type:

Optional[Array]

H

The estimated covariate coefficient matrix.

Type:

Optional[Array]

All attributes are optional and initialized to None by default.

H: Array | None

Alias for field number 8

L: Array | None

Alias for field number 3

Y_completed: Array | None

Alias for field number 4

beta: Array | None

Alias for field number 7

delta: Array | None

Alias for field number 6

gamma: Array | None

Alias for field number 5

lambda_H: float | int | Array | None

Alias for field number 2

lambda_L: float | int | Array | None

Alias for field number 1

tau: float | int | Array | None

Alias for field number 0

Internal Functions#

Functions in wrappers.py#

compute_treatment_effect#

mcnnm.wrappers.compute_treatment_effect(Y, W, L, X_tilde, Z_tilde, V, H_tilde, gamma, delta, beta, use_unit_fe, use_time_fe)[source]#

Compute the average treatment effect using the MC-NNM model estimates. Thin wrapper around the compute_Y_hat function.

This function calculates the difference between the observed outcomes and the completed (counterfactual) outcomes for treated units, then averages this difference to estimate the average treatment effect.

Parameters:

  • Y (Array) – The observed outcome matrix.

  • W (Array) – The binary treatment matrix.

  • L (Array) – The estimated low-rank matrix.

  • X_tilde (Array) – The augmented unit-specific covariates matrix.

  • Z_tilde (Array) – The augmented time-specific covariates matrix.

  • V (Array) – The unit-time specific covariates tensor.

  • H_tilde (Array) – The augmented covariate coefficient matrix.

  • gamma (Array) – The estimated unit fixed effects.

  • delta (Array) – The estimated time fixed effects.

  • beta (Array) – The estimated unit-time specific covariate coefficients.

  • use_unit_fe (bool) – Whether to use unit fixed effects.

  • use_time_fe (bool) – Whether to use time fixed effects.

Returns:

The estimated average treatment effect.

Return type:

Scalar

Functions in validation.py#

cross_validate#

mcnnm.validation.cross_validate(Y, X, Z, V, W, Omega_inv, use_unit_fe, use_time_fe, num_lam, max_iter=1000, tol=1e-05, cv_ratio=0.8, K=5)[source]#

Perform K-fold cross-validation to select the best regularization parameters for the model.

This function splits the data into K folds, trains the model on K-1 folds, and evaluates it on the remaining fold. The process is repeated for each fold and for different combinations of regularization parameters (lambda_L and lambda_H) specified in the lambda grid. The best lambda values are selected based on the minimum average root mean squared error (RMSE) across all folds.

Steps:

  1. Create K-fold masks using the create_folds function, which randomly assigns observations to folds based on the cv_ratio.

  2. Initialize the low-rank matrix L and the augmented covariate matrices X_tilde, Z_tilde, and V using the initialize_matrices function.

  3. Initialize the model parameters (gamma, delta, beta, H_tilde) and compute the maximum lambda values for each fold using the initialize_fold function and jax.vmap.

  4. Determine the overall maximum lambda_L and lambda_H values across all folds.

  5. Generate lambda_L and lambda_H value ranges using the propose_lambda_values function.

  6. Create a lambda grid by combining the lambda_L and lambda_H value ranges using the generate_lambda_grid function.

  7. Define the fold_loss function that computes the validation RMSE for each lambda combination within a fold: - Split the data into training and validation sets based on the fold mask. - Use jax.lax.scan to iterate over the lambda grid and compute the RMSE for each combination. - Train the model using the fit function on the training set for each lambda combination. - Compute the validation RMSE using the compute_objective_value function on the validation set. - Return the validation RMSE for each lambda combination.

  8. Apply the fold_loss function to each fold using jax.vmap to compute the validation RMSE for each lambda combination across all folds.

  9. Compute the average validation RMSE for each lambda combination across all folds.

  10. Select the best lambda_L and lambda_H values based on the minimum average RMSE.

  11. Determine the optimal lambda_L and lambda_H ranges by slicing the corresponding value ranges based on the best lambda values.

  12. Return the best lambda_L and lambda_H values along with their optimal ranges.

Parameters:

  • Y (Array) – The target variable matrix of shape (N, T).

  • X (Array) – The feature matrix for unit-specific covariates of shape (N, P).

  • Z (Array) – The feature matrix for time-specific covariates of shape (T, Q).

  • V (Array) – The feature matrix for unit-time covariates of shape (N, T, R).

  • W (Array) – The binary matrix indicating the presence of observations of shape (N, T).

  • Omega_inv (Array, optional) – The inverse of the covariance matrix of shape (T, T). If not provided, the identity matrix is used.

  • use_unit_fe (bool) – Whether to include unit fixed effects in the model.

  • use_time_fe (bool) – Whether to include time fixed effects in the model.

  • num_lam (int) – The number of lambda values to include in the lambda grid.

  • max_iter (int, optional) – The maximum number of iterations for model fitting. Default is 1000.

  • tol (float, optional) – The tolerance for convergence in model fitting. Default is 1e-5.

  • cv_ratio (float, optional) – The ratio of data to use for training in each fold. Default is 0.8.

  • K (int, optional) – The number of folds for cross-validation. Default is 5.

Returns:

A tuple containing the following elements:

  • best_lambda_L (Array): The best lambda_L value based on the minimum average RMSE.

  • best_lambda_H (Array): The best lambda_H value based on the minimum average RMSE.

  • lambda_L_opt_range (Array): The optimal lambda_L range.

  • lambda_H_opt_range (Array): The optimal lambda_H range.

Return type:

Tuple[Array, Array, Array, Array]

Raises:

ValueError – If the input arrays have inconsistent shapes.

Note

  • The function assumes that the input arrays are of type jax.numpy.ndarray.

  • The function uses jax.vmap and jax.lax.scan for parallelization and efficient computation.

  • The function initializes the model parameters using the initialize_matrices and initialize_fixed_effects_and_H functions.

  • The function generates a lambda grid using the propose_lambda_values and generate_lambda_grid functions.

  • The function computes the RMSE for each fold and lambda combination using the fit function.

  • The function selects the best lambda values based on the minimum average RMSE across all folds.

holdout_validate#

mcnnm.validation.holdout_validate(Y, X, Z, V, W, Omega_inv, use_unit_fe, use_time_fe, num_lam, initial_window, step_size, horizon, K, max_window_size=None, max_iter=1000, tol=1e-05)[source]#

Perform holdout validation to select the optimal regularization parameters for the MC-NNM model.

This function splits the data into K holdout folds along the time dimension, initializes the model configurations for each fold based on the observed data within the specified time window, and computes the holdout loss and RMSE for each fold and lambda pair. The lambda pair that yields the lowest average RMSE across all folds is selected as the optimal regularization parameters.

Steps:

  1. Create K holdout masks using the create_holdout_masks function, which generates masks based on the specified time windows determined by initial_window, step_size, and horizon.

  2. Initialize the low-rank matrix L and the augmented covariate matrices X_tilde, Z_tilde, and V using the initialize_matrices function.

  3. Initialize the model configurations (gamma, delta, beta, H_tilde) and compute the maximum lambda values for each holdout fold using the initialize_holdout function and jax.vmap.

  4. Determine the overall maximum lambda_L and lambda_H values across all holdout folds.

  5. Generate lambda_L and lambda_H value ranges using the propose_lambda_values function.

  6. Create a lambda grid by combining the lambda_L and lambda_H value ranges using the generate_lambda_grid function.

  7. Define the holdout_fold_loss function that computes the holdout RMSE for each lambda combination within a fold: - Split the data into training and validation sets based on the holdout mask. - Use jax.lax.scan to iterate over the lambda grid and compute the RMSE for each combination. - Train the model using the fit function on the training set for each lambda combination. - Compute the holdout RMSE using the compute_objective_value function on the validation set. - Return the holdout RMSE for each lambda combination.

  8. Apply the holdout_fold_loss function to each holdout fold using jax.vmap to compute the holdout RMSE for each lambda combination across all folds.

  9. Compute the average holdout RMSE for each lambda combination across all folds.

  10. Select the best lambda_L and lambda_H values based on the minimum average RMSE.

  11. Return the best lambda_L and lambda_H values along with the maximum lambda_L and lambda_H values.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • X (Array) – The unit-specific covariates matrix of shape (N, P).

  • Z (Array) – The time-specific covariates matrix of shape (T, Q).

  • V (Array) – The unit-time specific covariates tensor of shape (N, T, J).

  • W (Array) – The binary matrix indicating observed (0) and missing (1) entries in Y, shape (N, T).

  • Omega_inv (Array, optional) – The autocorrelation matrix of shape (T, T). If not provided, an identity matrix is used.

  • use_unit_fe (bool) – Whether to use unit fixed effects.

  • use_time_fe (bool) – Whether to use time fixed effects.

  • num_lam (int) – The number of lambda values to generate in the grid.

  • initial_window (int) – The size of the initial time window for holdout validation. It determines the number of time steps used to initialize the model configurations for each holdout fold.

  • step_size (int) – The step size for moving the time window in each holdout fold. It determines the number of time steps to shift the window for each subsequent fold.

  • horizon (int) – The size of the holdout horizon (number of time steps to predict). It determines the number of time steps used for evaluating the model’s performance in each holdout fold.

  • K (int) – The number of holdout folds.

  • max_window_size (int, optional) – The maximum size of the time window. If specified, it limits the size of the time window used for initializing the model configurations in each holdout fold. The window size will not exceed max_window_size even if initial_window + horizon is larger. Defaults to None, meaning no limit on the window size.

  • max_iter (int, optional) – Maximum number of iterations for fitting the model. Defaults to 1000.

  • tol (float, optional) – Convergence tolerance for fitting the model. Defaults to 1e-5.

Returns:

A tuple containing the following elements:

  • best_lambda_L (Array): The best lambda_L value based on the minimum average RMSE.

  • best_lambda_H (Array): The best lambda_H value based on the minimum average RMSE.

  • max_lambda_L (Array): The maximum lambda_L value across all holdout folds.

  • max_lambda_H (Array): The maximum lambda_H value across all holdout folds.

Return type:

Tuple[Array, Array, Array, Array]

Raises:

ValueError – If the input arrays have inconsistent shapes or if the time window parameters are invalid.

Note

  • The binary matrix W indicates observed (0) and missing (1) entries in Y. The missing entries (1) are relevant for the loss computation.

  • The function uses JAX’s vmap and scan operations to efficiently compute the holdout losses and RMSEs for multiple holdout folds and lambda pairs in parallel.

  • The function initializes the model configurations using the initialize_matrices and initialize_fixed_effects_and_H functions.

  • The function generates a lambda grid using the propose_lambda_values and generate_lambda_grid functions.

  • The function computes the RMSE for each holdout fold and lambda combination using the fit function.

  • The function selects the best lambda values based on the minimum average RMSE across all holdout folds.

final_fit#

mcnnm.validation.final_fit(Y, X, Z, V, W, Omega_inv, use_unit_fe, use_time_fe, best_lambda_L, best_lambda_H, lambda_L_opt_range, lambda_H_opt_range, max_iter=1000, tol=1e-05)[source]#

Perform the final fit of the MC-NNM model using the optimal regularization parameters.

This function fits the model using a sequence of lambda values, starting from the largest values in the optimal ranges and moving towards the best lambda values. This approach, known as warm-starting, helps to improve convergence and stability of the final fit.

This function is similar to NNM_H, with one key difference. Instead of fitting all models on the grid described by lambda_Ls and lambda_Hs, it only considers the shortest path from the point on the grid with highest lambda_L and lambda_H to the point on the grid with smallest values of lambda_L and lambda_H. The key benefit of using this function is that, for chosen values of lambda_L and lambda_H, training can be much faster as the number of trained models is M+N-1 compared to M*N, where M is the length of lambda_L and N is the length of lambda_H.

Steps:

  1. Initialize matrices and model parameters using the initialize_matrices and initialize_fixed_effects_and_H functions.

  2. Generate a lambda grid using the optimal lambda ranges.

  3. Extract the shortest path of lambda values from the largest to the best values.

  4. Iteratively fit the model for each lambda pair in the shortest path, using the results from the previous fit as initial values for the next fit.

  5. Return the final model parameters and loss.

Parameters:

  • Y (Array) – The target variable matrix of shape (N, T).

  • X (Array) – The feature matrix for unit-specific covariates of shape (N, P).

  • Z (Array) – The feature matrix for time-specific covariates of shape (T, Q).

  • V (Array) – The feature matrix for unit-time covariates of shape (N, T, R).

  • W (Array) – The binary matrix indicating the presence of observations of shape (N, T).

  • Omega_inv (Array, optional) – The inverse of the covariance matrix of shape (T, T). If not provided, an identity matrix is used.

  • use_unit_fe (bool) – Whether to include unit fixed effects in the model.

  • use_time_fe (bool) – Whether to include time fixed effects in the model.

  • best_lambda_L – The best lambda_L value determined by cross-validation or holdout validation.

  • best_lambda_H – The best lambda_H value determined by cross-validation or holdout validation.

  • lambda_L_opt_range (Array) – The optimal range of lambda_L values.

  • lambda_H_opt_range (Array) – The optimal range of lambda_H values.

  • max_iter (int, optional) – The maximum number of iterations for model fitting. Default is 1000.

  • tol (float, optional) – The tolerance for convergence in model fitting. Default is 1e-5.

Returns:

A tuple containing the following elements:

  • L_final (Array): The final low-rank matrix.

  • H_final (Array): The final interactive fixed effects matrix.

  • in_prod_final (Array): The final in-product matrix.

  • gamma_final (Array): The final unit fixed effects.

  • delta_final (Array): The final time fixed effects.

  • beta_final (Array): The final coefficients for covariates.

  • loss_final (float): The final loss value.

Return type:

Tuple[Array, Array, Array, Array, Array, Array, float]

Note

  • The function uses the fit function to perform the model fitting for each lambda pair.

  • The function uses jax.lax.scan for efficient iteration over the lambda pairs.

  • The warm-starting approach helps to improve the stability and convergence of the final fit.

Functions in core.py#

initialize_coefficients#

mcnnm.core.initialize_coefficients(Y, X_tilde, Z_tilde, V)[source]#

Initialize covariate and fixed effects coefficients for the MC-NNM model.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • X_tilde (Array) – The augmented unit-specific covariates matrix of shape (N, P+N).

  • Z_tilde (Array) – The augmented time-specific covariates matrix of shape (T, Q+T).

  • V (Array) – The unit-time specific covariates tensor of shape (N, T, J).

Returns:

A tuple containing initial values for:

  • H_tilde: Initial value for the interactive fixed effects matrix.

  • gamma: Initial value for unit fixed effects.

  • delta: Initial value for time fixed effects.

  • beta: Initial value for covariate coefficients.

Return type:

Tuple[Array, Array, Array, Array]

initialize_matrices#

mcnnm.core.initialize_matrices(Y, X, Z, V)[source]#

Initialize the matrices for the MC-NNM model.

This function initializes the low-rank matrix L and the covariate matrices X_tilde and Z_tilde. If the covariate matrices X, Z, or V are not provided, they are initialized to zero matrices/tensors.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • X (Optional[Array]) – The unit-specific covariates matrix of shape (N, P). If None, initialized to zeros.

  • Z (Optional[Array]) – The time-specific covariates matrix of shape (T, Q). If None, initialized to zeros.

  • V (Optional[Array]) – The unit-time-specific covariates tensor of shape (N, T, J). If None, initialized to zeros.

Returns:

A tuple containing:

  • L (Array): The low-rank matrix of shape (N, T).

  • X_tilde (Array): The augmented unit-specific covariates matrix of shape (N, P+N).

  • Z_tilde (Array): The augmented time-specific covariates matrix of shape (T, Q+T).

  • V (Array): The unit-time-specific covariates tensor of shape (N, T, J).

Return type:

Tuple[Array, Array, Array, Array]

compute_svd#

mcnnm.core.compute_svd(M)[source]#

Compute the Singular Value Decomposition (SVD) of the input matrix M. Return the left singular vectors (U), right singular vectors (V), and singular values (Sigma).

Parameters:

M (Array)

Return type:

tuple[Array, Array, Array]

svt#

mcnnm.core.svt(U, V, sing_vals, threshold)[source]#

Perform soft singular value thresholding (SVT) on the given singular value decomposition.

Parameters:

  • U (Array) – The left singular vectors matrix.

  • V (Array) – The right singular vectors matrix.

  • sing_vals (Array) – The singular values array.

  • threshold (Scalar) – The thresholding value.

Returns:

The thresholded low-rank matrix.

Return type:

Array

update_unit_fe#

mcnnm.core.update_unit_fe(Y, X_tilde, Z_tilde, H_tilde, W, L, time_fe, use_unit_fe)[source]#

Update the unit fixed effects in the coordinate descent algorithm when covariates are available.

This function calculates the unit fixed effects by minimizing the squared error of the model predictions with respect to the observed outcomes. It accounts for the covariate effects, low-rank matrix, and time fixed effects in the calculation.

The update process involves the following steps: 1. Compute the covariate contribution to the predicted outcomes. 2. Calculate the total predicted outcomes (including L and time fixed effects). 3. Compute the residuals (observed - predicted). 4. Apply the observation mask to the residuals. 5. Calculate the average residual for each unit, accounting for missing observations. 6. Return the negative of these averages as the updated unit fixed effects.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • X_tilde (Array) – The augmented unit-specific covariates matrix of shape (N, P+N).

  • Z_tilde (Array) – The augmented time-specific covariates matrix of shape (T, Q+T).

  • H_tilde (Array) – The augmented covariate coefficients matrix of shape (P+N, Q+T).

  • W (Array) – The mask matrix indicating observed entries of shape (N, T).

  • L (Array) – The low-rank matrix of shape (N, T).

  • time_fe (Array) – The time fixed effects vector of shape (T,).

  • use_unit_fe (bool) – Whether to estimate unit fixed effects.

Returns:

The updated unit fixed effects vector of shape (N,) if use_unit_fe is True, else a zero vector.

Return type:

Array

update_time_fe#

mcnnm.core.update_time_fe(Y, X_tilde, Z_tilde, H_tilde, W, L, unit_fe, use_time_fe)[source]#

Update the time fixed effects in the coordinate descent algorithm when covariates are available.

This function calculates the time fixed effects by minimizing the squared error of the model predictions with respect to the observed outcomes. It accounts for the covariate effects, low-rank matrix, and unit fixed effects in the calculation.

The update process involves the following steps: 1. Compute the covariate contribution to the predicted outcomes. 2. Calculate the total predicted outcomes (including L and unit fixed effects). 3. Compute the residuals (observed - predicted). 4. Apply the observation mask to the residuals. 5. Calculate the average residual for each time period, accounting for missing observations. 6. Return the negative of these averages as the updated time fixed effects.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • X_tilde (Array) – The augmented unit-specific covariates matrix of shape (N, P+N).

  • Z_tilde (Array) – The augmented time-specific covariates matrix of shape (T, Q+T).

  • H_tilde (Array) – The augmented covariate coefficients matrix of shape (P+N, Q+T).

  • W (Array) – The mask matrix indicating observed entries of shape (N, T).

  • L (Array) – The low-rank matrix of shape (N, T).

  • unit_fe (Array) – The unit fixed effects vector of shape (N,).

  • use_time_fe (bool) – Whether to estimate time fixed effects.

Returns:

The updated time fixed effects vector of shape (T,) if use_time_fe is True, else a zero vector.

Return type:

Array

update_beta#

mcnnm.core.update_beta(Y, X_tilde, Z_tilde, V, H_tilde, W, L, unit_fe, time_fe)[source]#

Update the unit-time-specific covariate coefficients (beta) in the coordinate descent algorithm.

This function calculates the coefficients for unit-time-specific covariates by minimizing the squared error of the model predictions with respect to the observed outcomes. It accounts for the other components of the model (covariate effects, low-rank matrix, unit and time fixed effects) in the calculation.

The update process involves the following steps: 1. Compute the covariate contribution to the predicted outcomes. 2. Calculate the total predicted outcomes (including all components except V*beta). 3. Compute the residuals (observed - predicted). 4. Apply the observation mask to the residuals and unit-time-specific covariates. 5. Calculate the updated beta coefficients using the masked residuals and covariates.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • X_tilde (Array) – The augmented unit-specific covariates matrix of shape (N, P+N).

  • Z_tilde (Array) – The augmented time-specific covariates matrix of shape (T, Q+T).

  • V (Array) – The unit-time-specific covariates tensor of shape (N, T, J).

  • H_tilde (Array) – The augmented covariate coefficients matrix of shape (P+N, Q+T).

  • W (Array) – The mask matrix indicating observed entries of shape (N, T).

  • L (Array) – The low-rank matrix of shape (N, T).

  • unit_fe (Array) – The unit fixed effects vector of shape (N,).

  • time_fe (Array) – The time fixed effects vector of shape (T,).

Returns:

The updated unit-time-specific covariate coefficients vector of shape (J,).

Return type:

Array

compute_Y_hat#

mcnnm.core.compute_Y_hat(L, X_tilde, Z_tilde, V, H_tilde, gamma, delta, beta, use_unit_fe, use_time_fe)[source]#

Compute the decomposition of the observed outcome matrix Y.

This function computes the decomposition of the observed outcome matrix Y into its low-rank component L, covariate effects, and fixed effects (unit and time). The decomposition is given by:

\[\hat{Y}_{it} = \hat{L} + \tilde{X}\hat{\tilde{H}}\tilde{Z}^T + \hat{\Gamma}1_T^T + 1_N\hat{\Delta}^T + [V_{i,t}^T \hat{\beta}]_{i,t}\]

where:

  • \(\hat{L}\) is the estimated low-rank matrix of shape (N, T)

  • \(\tilde{X}\) is the augmented unit-specific covariates matrix of shape (N, P+N)

  • \(\tilde{Z}\) is the augmented time-specific covariates matrix of shape (T, Q+T)

  • \(V\) is the unit-time-specific covariates tensor of shape (N, T, J)

  • \(\hat{\tilde{H}}\) is the estimated augmented covariate coefficients matrix of shape (P+N, Q+T)

  • \(\hat{\Gamma}\) is the estimated unit fixed effects vector of shape (N,)

  • \(\hat{\Delta}\) is the estimated time fixed effects vector of shape (T,)

  • \(\hat{\beta}\) is the estimated unit-time-specific covariate coefficients vector of shape (J,)

Parameters:

  • L (Array) – The low-rank matrix of shape (N, T).

  • X_tilde (Array) – The unit-specific covariates matrix of shape (N, P).

  • Z_tilde (Array) – The time-specific covariates matrix of shape (T, Q).

  • V (Array) – The unit-time-specific covariates tensor of shape (N, T, J).

  • H_tilde (Array) – The covariate coefficients matrix of shape (P + N, Q + T).

  • gamma (Array) – The unit fixed effects vector of shape (N,).

  • delta (Array) – The time fixed effects vector of shape (T,).

  • beta (Array) – The unit-time-specific covariate coefficients vector of shape (J,).

  • use_unit_fe (bool) – Whether to include unit fixed effects in the decomposition.

  • use_time_fe (bool) – Whether to include time fixed effects in the decomposition.

Returns:

The estimated matrix \(\hat{Y}\) of shape (N, T).

Return type:

Array

compute_objective_value#

mcnnm.core.compute_objective_value(Y, X_tilde, Z_tilde, V, H_tilde, W, L, gamma, delta, beta, sum_sing_vals, lambda_L, lambda_H, use_unit_fe, use_time_fe, inv_omega=None)[source]#

Compute the objective value for the MC-NNM model with covariates, fixed effects, and time series correlation.

The objective function is defined as:

\[\frac{1}{|\mathcal{O}|} \sum_{(i,t) \in \mathcal{O}} \sum_{(i,s) \in \mathcal{O}} (Y_{it} - \hat{Y}_{it}) [\Omega^{-1}]_{ts} (Y_{is} - \hat{Y}_{is}) + \lambda_L \|L^*\|_* + \lambda_H \|H^*\|_1\]

where:

  • \(Y_{it}\) is the observed outcome for unit \(i\) at time \(t\)

  • \(\widehat{Y}_{it}\) is the estimated outcome for unit \(i\) at time \(t\),

    computed by compute_Y_hat(), and given by:

    \[ \hat{Y}_{it} =\hat{L} + \tilde{X}\hat{\tilde{H}}\tilde{Z}^T + \hat{\Gamma}1_T^T - 1_N\hat{\Delta}^T - [V_{i,t}^T \hat{\beta}]_{i,t}\]

and:

  • \(\mathcal{O}\) is the set of observed entries in the outcome matrix

  • \(\Omega^{-1}\) is the inverse of the omega matrix, capturing the time series correlation

  • \(\lambda_L\) is the regularization parameter for the nuclear norm of \(L\)

  • \(\lambda_H\) is the regularization parameter for the element-wise L1 norm of \(\tilde{H}\)

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • X_tilde (Array) – The augmented unit-specific covariates matrix \(\tilde{X}\) of shape (N, P+N).

  • Z_tilde (Array) – The augmented time-specific covariates matrix \(\tilde{Z}\) of shape (T, Q+T).

  • V (Array) – The unit-time-specific covariates tensor \(V\) of shape (N, T, J).

  • H_tilde (Array) – The augmented covariate coefficients matrix \(\tilde{H}\) of shape (P+N, Q+T).

  • W (Array) – The mask matrix indicating observed entries of shape (N, T).

  • L (Array) – The low-rank matrix \(L\) of shape (N, T).

  • gamma (Array) – The unit fixed effects vector \(\Gamma\) of shape (N,).

  • delta (Array) – The time fixed effects vector \(\Delta\) of shape (T,).

  • beta (Array) – The unit-time-specific covariate coefficients vector \(\beta\) of shape (J,).

  • sum_sing_vals (Scalar) – The sum of singular values of L.

  • lambda_L (Scalar) – The regularization parameter for the nuclear norm of L.

  • lambda_H (Scalar) – The regularization parameter for the element-wise L1 norm of H.

  • use_unit_fe (bool) – Whether to include unit fixed effects in the decomposition.

  • use_time_fe (bool) – Whether to include time fixed effects in the decomposition.

  • inv_omega (Optional[Array]) – The inverse of the omega matrix of shape (T, T). If None, the identity matrix is used.

Returns:

The computed objective value.

Return type:

Scalar

initialize_fixed_effects_and_H#

mcnnm.core.initialize_fixed_effects_and_H(Y, L, X_tilde, Z_tilde, V, W, use_unit_fe, use_time_fe, niter=1000, rel_tol=1e-05, verbose=False)[source]#

Initialize fixed effects and the matrix H_tilde for the MC-NNM model.

This function initializes the fixed effects (unit and time) and the matrix H_tilde using an iterative coordinate descent algorithm. It also computes the maximum regularization parameters for the low-rank matrix L and the covariate coefficients matrix H_tilde.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • L (Array) – The zero-initialised low-rank matrix of shape (N, T).

  • X_tilde (Array) – The unit-specific covariates matrix of shape (N, P+N).

  • Z_tilde (Array) – The time-specific covariates matrix of shape (T, Q+T).

  • V (Array) – The unit-time-specific covariates tensor of shape (N, T, J).

  • W (Array) – The mask matrix indicating observed entries of shape (N, T).

  • use_unit_fe (bool) – Whether to include unit fixed effects in the decomposition.

  • use_time_fe (bool) – Whether to include time fixed effects in the decomposition.

  • niter (int, optional) – The maximum number of iterations for the coordinate descent algorithm. Default is 1000.

  • rel_tol (float, optional) – The relative tolerance for convergence. Default is 1e-5.

  • verbose (bool, optional) – Whether to print the objective value after initialization. Default is False.

Returns:

A tuple containing:

  • gamma (Array): The unit fixed effects vector of shape (N,).

  • delta (Array): The time fixed effects vector of shape (T,).

  • beta (Array): The unit-time-specific covariate coefficients vector of shape (J,).

  • H_tilde (Array): The covariate coefficients matrix of shape (P+N, Q+T).

  • T_mat (Array): The matrix T used for computing the regularization parameter lambda_H_max.

  • in_prod_T (Array): The inner product of T_mat used for computing lambda_H_max.

  • in_prod (Array): The inner product vector used for updating H_tilde. Initialized as zeros.

  • lambda_L_max (Scalar): The maximum regularization parameter for the nuclear norm of L.

  • lambda_H_max (Scalar): The maximum regularization parameter for the element-wise L1 norm of H_tilde.

Return type:

Tuple[Array, Array, Scalar, Scalar, Array, Array]

update_H#

mcnnm.core.update_H(Y, X_tilde, Z_tilde, V, H_tilde, T_mat, in_prod, in_prod_T, W, L, unit_fe, time_fe, beta, lambda_H, use_unit_fe, use_time_fe)[source]#

Update the covariate coefficients matrix H_tilde in the coordinate descent algorithm.

This function implements the update step for the covariate coefficients matrix (H_tilde) in the matrix completion algorithm. It uses a soft-thresholding operator to update each element of the coefficient matrix, taking into account the regularization parameter and the current residuals.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • X_tilde (Array) – The augmented unit-specific covariates matrix of shape (N, P+N).

  • Z_tilde (Array) – The augmented time-specific covariates matrix of shape (T, Q+T).

  • V (Array) – The unit-time-specific covariates tensor of shape (N, T, J).

  • H_tilde (Array) – The covariate coefficients matrix of shape (P+N, Q+T).

  • T_mat (Array) – The precomputed matrix T of shape (N * T, (P+N) * (Q+T)).

  • in_prod (Array) – The inner product vector of shape (N * T,).

  • in_prod_T (Array) – The inner product vector of T of shape ((P+N) * (Q+T),).

  • W (Array) – The mask matrix indicating observed entries of shape (N, T).

  • L (Array) – The low-rank matrix of shape (N, T).

  • unit_fe (Array) – The unit fixed effects vector of shape (N,).

  • time_fe (Array) – The time fixed effects vector of shape (T,).

  • beta (Array) – The unit-time-specific covariate coefficients vector of shape (J,).

  • lambda_H (Scalar) – The regularization parameter for the element-wise L1 norm of H_tilde.

  • use_unit_fe (bool) – Whether to include unit fixed effects in the decomposition.

  • use_time_fe (bool) – Whether to include time fixed effects in the decomposition.

Returns:

A tuple containing the updated covariate coefficients matrix H_tilde and the updated inner product vector in_prod.

Return type:

Tuple[Array, Array]

update_L#

mcnnm.core.update_L(Y, X_tilde, Z_tilde, V, H_tilde, W, L, unit_fe, time_fe, beta, lambda_L, use_unit_fe, use_time_fe)[source]#

Update the low-rank matrix L in the coordinate descent algorithm.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • X_tilde (Array) – The augmented unit-specific covariates matrix of shape (N, P+N).

  • Z_tilde (Array) – The augmented time-specific covariates matrix of shape (T, Q+T).

  • V (Array) – The unit-time-specific covariates tensor of shape (N, T, J).

  • H_tilde (Array) – The covariate coefficients matrix of shape (P+N, Q+T).

  • W (Array) – The mask matrix indicating observed entries of shape (N, T).

  • L (Array) – The low-rank matrix of shape (N, T).

  • unit_fe (Array) – The unit fixed effects vector of shape (N,).

  • time_fe (Array) – The time fixed effects vector of shape (T,).

  • beta (Array) – The unit-time-specific covariate coefficients vector of shape (J,).

  • lambda_L (Scalar) – The regularization parameter for the nuclear norm of L.

  • use_unit_fe (bool) – Whether to include unit fixed effects in the decomposition.

  • use_time_fe (bool) – Whether to include time fixed effects in the decomposition.

Returns:

A tuple containing the updated low-rank matrix L and the singular values.

Return type:

Tuple[Array, Array]

fit#

mcnnm.core.fit(Y, X_tilde, Z_tilde, V, H_tilde, T_mat, in_prod, in_prod_T, W, L, gamma, delta, beta, lambda_L, lambda_H, use_unit_fe, use_time_fe, Omega_inv=None, niter=1000, rel_tol=1e-05, verbose=False, print_iters=False)[source]#

Perform cyclic coordinate descent to estimate the MC-NNM model parameters.

This function implements the core estimation algorithm for the Matrix Completion with Nuclear Norm Minimization (MC-NNM) model. It uses a cyclic coordinate descent approach to iteratively update the model parameters, including the low-rank matrix L, the covariate coefficients matrix H_tilde, fixed effects vectors (gamma and delta), and unit-time-specific covariate coefficients (beta).

The algorithm minimizes an objective function that balances the fit to the observed data with regularization terms. It uses nuclear norm regularization for the low-rank matrix L and element-wise L1 regularization for the covariate coefficients matrix H_tilde. The function handles both unit and time fixed effects, as well as various types of covariates (unit-specific, time-specific, and unit-time-specific).

The estimation process continues until convergence is reached (determined by the relative change in the objective function) or the maximum number of iterations is exceeded. The function can optionally provide verbose output to track the optimization progress.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • X_tilde (Array) – The augmented unit-specific covariates matrix of shape (N, P+N).

  • Z_tilde (Array) – The augmented time-specific covariates matrix of shape (T, Q+T).

  • V (Array) – The unit-time-specific covariates tensor of shape (N, T, J).

  • H_tilde (Array) – The initial covariate coefficients matrix of shape (P+N, Q+T).

  • T_mat (Array) – The precomputed matrix T of shape (N * T, (P+N) * (Q+T)).

  • in_prod (Array) – The inner product vector of shape (N * T,).

  • in_prod_T (Array) – The inner product vector of T of shape ((P+N) * (Q+T),).

  • W (Array) – The mask matrix indicating observed entries of shape (N, T).

  • L (Array) – The initial low-rank matrix of shape (N, T).

  • gamma (Array) – The initial unit fixed effects vector of shape (N,).

  • delta (Array) – The initial time fixed effects vector of shape (T,).

  • beta (Array) – The initial unit-time-specific covariate coefficients vector of shape (J,).

  • lambda_L (Scalar) – The regularization parameter for the nuclear norm of L.

  • lambda_H (Scalar) – The regularization parameter for the element-wise L1 norm of H_tilde.

  • use_unit_fe (bool) – Whether to include unit fixed effects in the decomposition.

  • use_time_fe (bool) – Whether to include time fixed effects in the decomposition.

  • Omega_inv (Optional[Array]) – The inverse of the omega matrix of shape (T, T). If None, the identity matrix is

  • niter (int, optional) – The maximum number of iterations for the coordinate descent algorithm. Default is 1000.

  • rel_tol (float, optional) – The relative tolerance for convergence. Default is 1e-5.

  • verbose (bool, optional) – Whether to print the objective value at each iteration. Default is False.

  • print_iters (bool, optional) – Whether to print in each iteration. Default is False.

Returns:

A tuple containing:

  • The updated covariate coefficient matrix H_tilde

  • The updated low-rank matrix L

  • The updated unit fixed effects vector

  • The updated time fixed effects vector

  • The updated unit-time-specific covariate vector

  • The updated in_prod vector

  • The final objective value

Return type:

Tuple[Array, Array, Array, Array, Array, Array]

Note: This function is computationally intensive and forms the core of the MC-NNM estimation procedure. It is designed to handle large matrices efficiently, leveraging JAX for potential GPU acceleration. The convergence speed and final estimates can be sensitive to the choice of regularization parameters (lambda_L and lambda_H) and the initial values of the matrices and vectors.

Functions in core_utils.py#

is_positive_definite#

mcnnm.core_utils.is_positive_definite(mat)[source]#

Check if a matrix is positive definite.

Parameters:

mat (jnp.ndarray) – The input matrix to check.

Returns:

True if the matrix is positive definite, False otherwise.

Return type:

bool

This function checks if a matrix is positive definite by computing its eigenvalues and checking if they are all positive.

mask_observed#

mcnnm.core_utils.mask_observed(A, mask)[source]#

Projects the matrix A onto the observed entries specified by the binary mask. Corresponds to \(P_{\mathcal{O}}\) in the paper.

Parameters:

  • A (Array) – The input matrix.

  • mask (Array) – The binary mask matrix, where 1 indicates an observed entry and 0 indicates an unobserved entry.

Returns:

The projected matrix.

Return type:

Array

Raises:

ValueError – If the shapes of A and mask do not match.

\[P_{\mathcal{O}}(A) = A \odot \text{mask}\]

where \(\odot\) denotes the element-wise product.

mask_unobserved#

mcnnm.core_utils.mask_unobserved(A, mask)[source]#

Projects the matrix A onto the unobserved entries specified by the binary mask. Corresponds to \(P_{\mathcal{O}}^\perp\) in the paper.

Parameters:

  • A (Array) – The input matrix.

  • mask (Array) – The binary mask matrix, where 1 indicates an observed entry and 0 indicates an unobserved entry.

Returns:

The projected matrix.

Return type:

Array

Raises:

ValueError – If the shapes of A and mask do not match.

\[P_{\mathcal{O}}^\perp(A) = A \odot (\mathbf{1} - \text{mask})\]

where \(\odot\) denotes the element-wise product and \(\mathbf{1}\) is a matrix of 1s.

frobenius_norm#

mcnnm.core_utils.frobenius_norm(A)[source]#

Computes the Frobenius norm of a matrix A.

Parameters:

A (Array) – The input matrix.

Returns:

The Frobenius norm of the matrix A.

Return type:

Scalar

Raises:

ValueError – If the input is not a 2D array.

nuclear_norm#

mcnnm.core_utils.nuclear_norm(A)[source]#

Computes the nuclear norm (sum of singular values) of a matrix A.

Parameters:

A (Array) – The input matrix.

Returns:

The nuclear norm of the matrix A.

Return type:

Scalar

Raises:

ValueError – If the input is not a 2D array.

element_wise_l1_norm#

mcnnm.core_utils.element_wise_l1_norm(A)[source]#

Computes the element-wise L1 norm of a matrix A.

Parameters:

A (Array) – The input matrix.

Returns:

The element-wise L1 norm of the matrix A.

Return type:

Scalar

Raises:

ValueError – If the input is not a 2D array.

normalize#

mcnnm.core_utils.normalize(mat)[source]#

Normalize the columns of the input matrix. Return the normalized matrix and the column norms.

Parameters:

mat (Array)

Return type:

tuple[Array, Array]

normalize_back#

mcnnm.core_utils.normalize_back(mat, row_scales, col_scales)[source]#

Rescale the rows and columns of the matrix H using the provided scales.

Parameters:
Return type:

Array

Functions in utils.py#

convert_inputs#

mcnnm.utils.convert_inputs(Y, W, X=None, Z=None, V=None, Omega=None)[source]#

Convert input DataFrames to JAX arrays for the MC-NNM model.

Parameters:

  • Y (DataFrame)

  • W (DataFrame)

  • X (DataFrame | None)

  • Z (DataFrame | None)

  • V (list[DataFrame] | None)

  • Omega (DataFrame | None)

Return type:

tuple[Array, Array, Array | None, Array | None, Array | None, Array | None]

check_inputs#

mcnnm.utils.check_inputs(Y, W, X=None, Z=None, V=None, Omega=None)[source]#

Check and preprocess input arrays for the MC-NNM model.

Parameters:

  • Y (Array) – The observed outcome matrix of shape (N, T).

  • W (Array) – The binary treatment matrix of shape (N, T).

  • X (Optional[Array]) – The unit-specific covariates matrix of shape (N, P). Default is None.

  • Z (Optional[Array]) – The time-specific covariates matrix of shape (T, Q). Default is None.

  • V (Optional[Array]) – The unit-time specific covariates tensor of shape (N, T, J). Default is None.

  • Omega (Optional[Array]) – The autocorrelation matrix of shape (T, T). Default is None.

Returns:

A tuple containing preprocessed X, Z, V, and Omega arrays.

Return type:

Tuple

Raises:

ValueError – If the input array dimensions are mismatched or invalid.

generate_data#

mcnnm.utils.generate_data(nobs=500, nperiods=100, Y_mean=10.0, treatment_probability=0.5, rank=5, treatment_effect=5.0, unit_fe=True, time_fe=True, X_cov=True, Z_cov=True, V_cov=True, fixed_effects_scale=0.5, covariates_scale=0.5, noise_scale=1, assignment_mechanism='last_periods', treated_fraction=0.2, last_treated_periods=2, autocorrelation=0.0, seed=None)[source]#
Parameters:

  • nobs (int)

  • nperiods (int)

  • Y_mean (float)

  • treatment_probability (float)

  • rank (int)

  • treatment_effect (float)

  • unit_fe (bool)

  • time_fe (bool)

  • X_cov (bool)

  • Z_cov (bool)

  • V_cov (bool)

  • fixed_effects_scale (float)

  • covariates_scale (float)

  • noise_scale (float)

  • assignment_mechanism (Literal['staggered', 'block', 'single_treated_period', 'single_treated_unit', 'last_periods'])

  • treated_fraction (float)

  • last_treated_periods (int)

  • autocorrelation (float)

  • seed (int | None)

Return type:

tuple[Array, Array, Array | None, Array | None, Array | None, dict]

propose_lambda_values#

mcnnm.utils.propose_lambda_values(max_lambda, min_lambda=None, n_lambdas=6)[source]#

Creates a decreasing, log-spaced list of proposed lambda values between max_lambda and min_lambda, and appends 0 to the end of the list.

Parameters:

  • max_lambda (float | int | Array) – The maximum lambda value.

  • min_lambda (float | int | Array | None) – The minimum lambda value. If None, it is set to max_lambda - 3 (in log10 scale).

  • n_lambdas (int) – The number of lambda values to generate (excluding the appended 0).

Returns:

The decreasing sequence of proposed lambda values, including 0 at the end.

Return type:

Array

Raises:

ValueError – If max_lambda is smaller than the default minimum lambda value (1e-10).

generate_lambda_grid#

mcnnm.utils.generate_lambda_grid(lambda_L_values, lambda_H_values)[source]#

Generates a grid of lambda values for the MC-NNM model.

This function creates a 2D grid of lambda values by forming a meshgrid from the provided sequences for both lambda_L and lambda_H.

Parameters:

  • lambda_L_values (Array) – The decreasing sequence of lambda values for the L dimension.

  • lambda_H_values (Array) – The decreasing sequence of lambda values for the H dimension.

Returns:

A 2D array where each row represents a pair of lambda values (lambda_L, lambda_H).

Return type:

jnp.ndarray

extract_shortest_path#

mcnnm.utils.extract_shortest_path(lambda_grid)[source]#

Extracts the shortest path along the edges of a 2D lambda grid.

This function traverses the edges of the given lambda grid starting from the top-right corner (highest lambda_L and lambda_H), moving left to the top-left corner (lowest lambda_L, highest lambda_H), and then down to the bottom-left corner (lowest lambda_L and lambda_H).

Parameters:

lambda_grid (jnp.ndarray) – A 2D array representing the flattened lambda grid.

Returns:

A 2D array containing the lambda pairs along the shortest path.

Return type:

jnp.ndarray

generate_holdout_val_defaults#

mcnnm.utils.generate_holdout_val_defaults(Y)[source]#

Generates default parameters for time-based validation.

This function calculates default values for various parameters used in time-based validation, including the initial window size, step size, horizon, and lambda grid.

Parameters:

Y (Array) – The observed outcome matrix of shape (N, T).

Returns:

A dictionary containing the default parameters for time-based validation.

  • initial_window (int): Number of initial time periods to use for the first training set.

  • step_size (int): Number of time periods to move forward for each split.

  • horizon (int): Number of future time periods to predict (forecast horizon).

Return type:

dict

validate_holdout_config#

mcnnm.utils.validate_holdout_config(initial_window, step_size, horizon, K, max_window_size, T)[source]#

Validate the configuration of initial_window, step_size, horizon, K, and max_window_size for holdout validation.

Parameters:

  • initial_window (int) – The size of the initial time window for holdout validation.

  • step_size (int) – The step size for moving the time window in each holdout fold.

  • horizon (int) – The size of the holdout horizon (number of time steps to predict).

  • K (int) – The number of holdout folds.

  • max_window_size (int, optional) – The maximum size of the time window. If specified, it limits the size of the time window used for initializing the model configurations in each holdout fold.

  • T (int) – The total number of time steps in the data.

Returns:

A tuple containing the validated or default values for

initial_window, step_size, horizon, K, and max_window_size.

Return type:

Tuple[int, int, int, int, Optional[int]]

Raises:

ValueError – If the configuration is invalid or inconsistent and cannot be adjusted to sensible defaults.

Contents