API Reference

Abstract type encompassing all reward estimators

Fit a reward estimator to the given data and then estimate the rewards for the same data, with measures to guard against overfitting the in-sample reward estimates.

The specific method depends on the type of reward estimation:

• fit_predict! for categorical treatments
• fit_predict! for numeric treatments

The predictions are returned as a Dict containing the following entries:

• :reward: the estimated reward for each observation under each treatment
• :propensity: the estimated propensity score for each observation under each treatment
• :outcome: the estimated outcome for each observation under each treatment
• :censoring: the estimated censoring curve for each observation under each treatment (for problems with survival outcomes only)

Based on the type of the reward estimation being conducted, some entries may not be relevant, and in these cases will instead contain nothing.

Abstract type encompassing all reward estimators with categorical treatments

Learner for reward estimation with categorical treatments and classification outcomes.

The following parameters are supported (refer to the documentation for each):

• shared Learner parameters
• shared RewardEstimator parameters

Learner for reward estimation with categorical treatments and regression outcomes.

The following parameters are supported (refer to the documentation for each):

• shared Learner parameters
• shared RewardEstimator parameters

Learner for reward estimation with categorical treatments and survival outcomes.

The following parameters are supported (refer to the documentation for each):

• shared Learner parameters
• shared RewardEstimator parameters

fit_predict!(lnr::CategoricalRewardEstimator, X::FeatureInput,
             treatments::AbstractVector, outcomes::AbstractVector;
             keyword_arguments...)

Fit lnr with a reward estimation model on features X, treatments treatments, and outcomes outcomes, and return predictions for each observation along with the scores of the internal estimators during training.

Supports the same keyword_arguments as score.

fit_predict!(lnr::CategoricalSurvivalRewardEstimator, X::FeatureInput,
             treatments::AbstractVector, deaths::AbstractVector{Bool},
             times::AbstractVector; keyword_arguments...)

Alternate version of fit_predict! for survival problems, where the survival outcomes are given by deaths and times.

predict(lnr::CategoricalRewardEstimator, X::FeatureInput,
        treatments::AbstractVector, outcomes::AbstractVector)

Return predictions estimated by lnr for each observation in the data given by X, treatments and outcomes.

The returned predictions are in the same format as for fit_predict!.

predict(lnr::CategoricalSurvivalRewardEstimator, X::FeatureInput,
        treatments::AbstractVector, deaths::AbstractVector{Bool},
        times::AbstractVector)

Alternate version of predict for survival problems, where the survival outcomes are given by deaths and times.

predict(lnr::CategoricalRewardEstimator, X::FeatureInput)

Alternate version of predict that can be used when using the direct method for reward estimation.

predict_reward(lnr::CategoricalRewardEstimator, treatments::AbstractVector,
               outcomes::AbstractVector, predictions::AbstractDict)

Return reward predictions estimated by lnr for each observation for the data contained in treatments and outcomes with model estimates predictions.

predictions should be supplied in the same format as returned by fit_predict! or predict, and the predictions are returned in the same format.

predict_reward(lnr::CategoricalSurvivalRewardEstimator,
               treatments::AbstractVector, deaths::AbstractVector{Bool},
               times::AbstractVector, predictions::AbstractDict)

Alternate version of predict_reward for survival problems, where the survival outcomes are given by deaths and times.

score(lnr::CategoricalRewardEstimator, X::FeatureInput,
      treatments::AbstractVector, outcomes::AbstractVector;
      keyword_arguments...)

Calculate the scores of the internal estimators in lnr on the data given by X, treatments and outcomes.

Returns a Dict with the following entries:

• :propensity: the score for the propensity estimator
• :outcome: a Dict where the keys are the possible treatments, and the values are the scores of the outcome estimator corresponding to each treatment

A score of NaN will be returned for any internal estimator that is not present.

Keyword Arguments

• propensity_score_criterion=:default: the criterion to use for evaluating the propensity estimator
• propensity_score_kwargs::Dict: any additional keyword arguments to use when evaluating the score for the propensity estimator
• outcome_score_criterion=:default: the criterion to use for evaluating the outcome estimator
• outcome_score_kwargs::Dict: any additional keyword arguments to use when evaluating the score for the outcome estimator

score(lnr::CategoricalSurvivalRewardEstimator, X::FeatureInput,
      treatments::AbstractVector, deaths::AbstractVector{Bool},
      times::AbstractVector; keyword_arguments...)

Alternate version of score for survival problems, where the survival outcomes are given by deaths and times.

An additional entry :censoring will be present in the returned Dict, containing a Dict where the keys are the possible treatments, and the values are the scores of the censoring estimator corresponding to each treatment.

Learner that estimates equal propensity for all treatments.

For use with data from randomized experiments where treatments are known to be randomly assigned.

Abstract type encompassing all reward estimators with numeric treatments

Learner for reward estimation with numeric treatments and classification outcomes.

The following parameters are supported (refer to the documentation for each):

• shared Learner parameters
• shared RewardEstimator parameters
• shared NumericRewardEstimator parameters

Learner for reward estimation with numeric treatments and regression outcomes.

The following parameters are supported (refer to the documentation for each):

• shared Learner parameters
• shared RewardEstimator parameters
• shared NumericRewardEstimator parameters

Learner for reward estimation with numeric treatments and survival outcomes.

The following parameters are supported (refer to the documentation for each):

• shared Learner parameters
• shared RewardEstimator parameters
• shared NumericRewardEstimator parameters

fit_predict!(lnr::NumericRewardEstimator, X::FeatureInput,
             treatments::Union{FeatureInput,AbstractVector},
             outcomes::AbstractVector,
             treatment_candidates::Union{FeatureInput,AbstractVector},
             keyword_arguments...)

Fit lnr with a reward estimation model on features X, treatments treatments, and outcomes outcomes, and return predictions for each observation under each treatment option in treatment_candidates, as well as the score of the internal estimators during training.

all_treatment_combinations can be used to construct treatment_candidates in the correct format.

Supports the same keyword_arguments as score.

fit_predict!(lnr::NumericSurvivalRewardEstimator, X::FeatureInput,
             treatments::Union{FeatureInput,AbstractVector},
             deaths::AbstractVector{Bool}, times::AbstractVector,
             treatment_candidates::Union{FeatureInput,AbstractVector},
             keyword_arguments...)

Alternate version of fit_predict! for survival problems, where the survival outcomes are given by deaths and times.

predict(lnr::NumericRewardEstimator, X::FeatureInput,
        treatments::Union{FeatureInput,AbstractVector},
        outcomes::AbstractVector)

Return predictions estimated by lnr for each observation in the data given by X, treatments and outcomes.

The returned predictions are in the same format as for fit_predict!.

predict(lnr::NumericSurvivalRewardEstimator, X::FeatureInput,
        treatments::Union{FeatureInput,AbstractVector},
        deaths::AbstractVector{Bool}, times::AbstractVector)

Alternate version of predict for survival problems, where the survival outcomes are given by deaths and times.

predict(lnr::NumericRewardEstimator, X::FeatureInput)

Alternate version of predict that can be used when using the direct method for reward estimation.

predict_reward(lnr::NumericRewardEstimator,
               treatments::Union{FeatureInput,AbstractVector},
               outcomes::AbstractVector, predictions::AbstractDict)

Return reward predictions estimated by lnr for each observation for the data contained in treatments and outcomes with model estimates predictions.

predictions should be supplied in the same format as returned by fit_predict! or predict, and the predictions are returned in the same format.

predict_reward(lnr::NumericSurvivalRewardEstimator,
               treatments::Union{FeatureInput,AbstractVector},
               deaths::AbstractVector{Bool}, times::AbstractVector,
               predictions::AbstractDict)

Alternate version of predict_reward for survival problems, where the survival outcomes are given by deaths and times.

score(lnr::NumericRewardEstimator, X::FeatureInput,
      treatments::Union{FeatureInput,AbstractVector},
      outcomes::AbstractVector; keyword_arguments...)

Calculate the score of the internal estimators in lnr on the data given by X, treatments and outcomes.

Returns a Dict with the following entries:

• :propensity: a Dict where the keys are the treatment candidates, and the values are the scores of the propensity estimator corresponding to each candidate
• :outcome: a Dict where the keys are the treatment candidates, and the values are the scores of the outcome estimator corresponding to each candidate

A score of NaN will be returned for any internal estimator that is not present.

Keyword Arguments

• propensity_score_criterion=:default: the criterion to use for evaluating the propensity estimator
• propensity_score_kwargs::Dict: any additional keyword arguments to use when evaluating the score for the propensity estimator
• outcome_score_criterion=:default: the criterion to use for evaluating the outcome estimator
• outcome_score_kwargs::Dict: any additional keyword arguments to use when evaluating the score for the outcome estimator

score(lnr::NumericSurvivalRewardEstimator, X::FeatureInput,
      treatments::Union{FeatureInput,AbstractVector},
      deaths::AbstractVector{Bool}, times::AbstractVector;
      keyword_arguments...)

Alternate version of score for survival problems, where the survival outcomes are given by deaths and times.

An additional entry :censoring will be present in the returned Dict, containing a Dict where the keys are the treatment candidates, and the values are the scores of the censoring estimator corresponding to each candidate.

get_estimation_densities(lnr::NumericRewardEstimator,
                         treatments::Union{FeatureInput,AbstractVector},
                         treatment_candidates::Union{FeatureInput,AbstractVector})

Return the total kernel density surrounding each treatment option in treatment_candidates when applying the estimation kernel function and bandwidth in lnr to treatments.

tune_reward_kernel_bandwidth(lnr::NumericRewardEstimator,
                             treatments::Union{FeatureInput,AbstractVector},
                             outcomes::AbstractVector,
                             predictions::AbstractDict,
                             input_bandwidths::AbstractVector{<:Real})

Conduct the reward kernel bandwidth tuning procedure (using lnr, the data contained in treatments and outcomes, and the model estimates predictions) for each starting value in input_bandwidths and return a Vector containing the final tuned values.

predictions should be supplied in the same format as returned by fit_predict! or predict.

Please refer to the documentation on kernel bandwidth tuning for more information on the recommended tuning approach.

tune_reward_kernel_bandwidth(lnr::NumericRewardEstimator,
                             treatments::Union{FeatureInput,AbstractVector},
                             deaths::AbstractVector{Bool},
                             times::AbstractVector,
                             predictions::AbstractDict,
                             input_bandwidths::AbstractVector{<:Real})

Alternate version of tune_reward_kernel_bandwidth for survival problems, where the survival outcomes are given by deaths and times.

set_reward_kernel_bandwidth!(lnr::NumericRewardEstimator,
                             treatments::Union{FeatureInput,AbstractVector},
                             outcomes::AbstractVector,
                             predictions::AbstractDict,
                             bandwidth::Real)

Save the new value of bandwidth as the reward kernel bandwidth inside lnr, and return new reward predictions generated using this bandwidth for the data contained in treatments and outcomes with model estimates predictions.

predictions should be supplied in the same format as returned by fit_predict! or predict.

set_reward_kernel_bandwidth!(lnr::NumericRewardEstimator,
                             treatments::Union{FeatureInput,AbstractVector},
                             deaths::AbstractVector{Bool},
                             times::AbstractVector,
                             predictions::AbstractDict,
                             bandwidth::Real)

Alternate version of set_reward_kernel_bandwidth! for survival problems, where the survival outcomes are given by deaths and times.

all_treatment_combinations(args...)
all_treatment_combinations(keyword_arguments...)

Return a DataFrame containing all treatment combinations of one or more treatment vectors, ready for use as treatment_candidates in fit_predict! or predict.

Each argument should be an AbstractVector containing the possible options for each treatment, or a keyword argument with the name of the treatment and the value as the possible options.

Examples

All combinations of two treatments ranging between 2 and 4:

IAI.all_treatment_combinations(2:4, 2:4)

9×2 DataFrame
 Row │ treatment1  treatment2
     │ Int64       Int64
─────┼────────────────────────
   1 │          2           2
   2 │          3           2
   3 │          4           2
   4 │          2           3
   5 │          3           3
   6 │          4           3
   7 │          2           4
   8 │          3           4
   9 │          4           4

All combinations of treatments A (ranging between 1 and 3) and B (with three possible values):

IAI.all_treatment_combinations(A=1:3, B=[0.5, 0.75, 1.0])

9×2 DataFrame
 Row │ A     B
     │ Real  Real
─────┼────────────
   1 │    1  0.5
   2 │    2  0.5
   3 │    3  0.5
   4 │    1  0.75
   5 │    2  0.75
   6 │    3  0.75
   7 │    1  1.0
   8 │    2  1.0
   9 │    3  1.0

convert_treatments_to_numeric(treatments::Union{Symbol,AbstractString})
convert_treatments_to_numeric(treatments::AbstractArray)

Convert treatments from symbol/string format into numeric values.

When numeric treatments are used, each column in the rewards dataframe corresponds to numeric doses of one or more numeric treatments. However, dataframe column names are saved internally as strings, meaning that some functions can only return these values as strings rather than in their original numeric form, including:

• predict for numeric-dose policy learners
• predict_treatment_rank
• get_policy_treatment_rank

This function provides a convenient method to convert these strings back into numeric doses, either as a scalar value for single treatment problems, or as a vector of doses for multiple treatment problems.

Examples

Converting treatment names back to numeric values with a single numeric treatment:

predictions = ["1", "2", "3", "4", "5", "6"]
IAI.convert_treatments_to_numeric(predictions)

6-element Vector{Int64}:
 1
 2
 3
 4
 5
 6

Converting treatment names back to numeric values with multiple numeric treatments:

predictions = ["[1,1]", "[2,1]", "[3,1]", "[1,2]", "[2,2]", "[3,2]"]
IAI.convert_treatments_to_numeric(predictions)

6-element Vector{Vector{Float64}}:
 [1.0, 1.0]
 [2.0, 1.0]
 [3.0, 1.0]
 [1.0, 2.0]
 [2.0, 2.0]
 [3.0, 2.0]