Skip to content

Models

Models

Model

Base model interface.

fit(X: np.ndarray, y: np.ndarray, groups: List[np.ndarray])

Fit on grouped data.

Parameters:

  • X (ndarray) –

    Feature matrix.

  • y (ndarray) –

    Targets.

  • groups (list[ndarray]) –

    Per-group selectors (integer indices or boolean masks).

Returns:

Raises:

  • NotImplementedError –

    If not implemented by the subclass.

set_group_weights(group_weights: jnp.ndarray) -> None

Set per-group weights.

Parameters:

  • group_weights (ndarray) –

    Array of shape (n_groups,).

LogisticRegression

Bases: BaseEstimator, ClassifierMixin, Model

Group-weighted logistic regression (sklearn LogisticRegression).

Parameters:

  • group_weights (ndarray, default: np.array([]) ) –

    Group weights. If empty or invalid, falls back to uniform.

  • l2_reg (float, default: 1e-3 ) –

    L2 regularization strength (mapped to C = 1 / l2_reg).

  • fit_intercept (bool, default: True ) –

    Whether to fit an intercept.

  • random_state (int or None, default: 42 ) –

    Passed to sklearn.

fit(X: np.ndarray, y: np.ndarray, groups: List[np.ndarray])

Fit the logistic regression model on grouped data with group weights.

RidgeRegression

Bases: BaseEstimator, RegressorMixin, Model

Group-weighted ridge regression (sklearn Ridge).

Parameters:

  • group_weights (ndarray, default: np.array([]) ) –

    Group weights. If empty, uses uniform over groups.

  • l2_reg (float, default: 1e-3 ) –

    L2 regularization strength (alpha in sklearn).

  • fit_intercept (bool, default: True ) –

    Whether to fit an intercept.

  • random_state (int or None, default: 42 ) –

    Kept for API parity with other models.

score(X: np.ndarray, y: np.ndarray) -> float

Relative RMSE score: RMSE divided by mean absolute target magnitude.

SVM

Bases: BaseEstimator, ClassifierMixin, Model

Group-weighted linear SVM (sklearn LinearSVC).

Parameters:

  • group_weights (ndarray, default: np.array([]) ) –

    Group weights. If empty, uses uniform over groups.

  • l2_reg (float, default: 1e-3 ) –

    L2 regularization strength (mapped to C = 1 / l2_reg).

  • fit_intercept (bool, default: True ) –

    Whether to fit an intercept.

  • tol (float, default: 1e-9 ) –

    Solver tolerance.

  • max_iter (int, default: 1000 ) –

    Maximum iterations.

  • random_state (int, default: 42 ) –

    Passed to sklearn.