API Reference

Documentation for the RewardEstimation public interface.

Index

General

IAI.fit_predict!Function

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:

Categorical Treatments

IAI.fit_predict!Method
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 predicted counterfactual rewards 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.

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

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


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.

IAI.scoreMethod
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.

IAI.EqualPropensityEstimatorType

Learner that estimates equal propensity for all treatments.

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

Numeric Treatments

IAI.fit_predict!Method
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 predicted counterfactual rewards for each observation under each treatment option in treatment_candidates, as well as the score of the internal outcome estimator.

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.

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

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


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.

IAI.scoreMethod
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.

IAI.get_estimation_densitiesFunction
get_estimation_densities(lnr::NumericRewardEstimator)

Return the total kernel density surrounding each treatment candidate for the propensity/outcome estimation problems in the fitted lnr.

IAI.tune_reward_kernel_bandwidthFunction
tune_reward_kernel_bandwidth(lnr::NumericRewardEstimator,
                             input_bandwidths::AbstractVector{<:Real})

Conduct the reward kernel bandwidth tuning procedure using lnr for each starting value in input_bandwidths and return a Vector containing the final tuned values.

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

IAI.set_reward_kernel_bandwidth!Function
set_reward_kernel_bandwidth!(lnr::NumericRewardEstimator,
                             bandwidth::Real)

Save the new value of bandwidth as the reward kernel bandwidth inside lnr, and return a DataFrame containing new reward predictions generated using this bandwidth for the original data used to train lnr.

IAI.all_treatment_combinationsFunction
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
IAI.convert_treatments_to_numericFunction
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:

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]