lifejacket 0.2.1__py3-none-any.whl → 1.0.2__py3-none-any.whl

lifejacket/helper_functions.py

@@ -11,8 +11,6 @@ import numpy as np
 import jax.numpy as jnp
 import pandas as pd
 
-from .constants import InverseStabilizationMethods
-
 logger = logging.getLogger(__name__)
 logging.basicConfig(
     format="%(asctime)s,%(msecs)03d %(levelname)-2s [%(filename)s:%(lineno)d] %(message)s",
@@ -27,11 +25,7 @@ def conditional_x_or_one_minus_x(x, condition):
 
 def invert_matrix_and_check_conditioning(
     matrix: np.ndarray,
-    inverse_stabilization_method: str = InverseStabilizationMethods.NONE,
     condition_num_threshold: float = 10**4,
-    ridge_median_singular_value_fraction: str = 0.01,
-    beta_dim: int = None,
-    theta_dim: int = None,
 ):
     """
     Check a matrix's condition number and invert it. If the condition number is
@@ -39,139 +33,15 @@ def invert_matrix_and_check_conditioning(
     Parameters
     """
     inverse = None
-    pre_inversion_condition_number = np.linalg.cond(matrix)
-    if pre_inversion_condition_number > condition_num_threshold:
+    condition_number = np.linalg.cond(matrix)
+    if condition_number > condition_num_threshold:
         logger.warning(
-            "You are inverting a matrix with a large condition number: %s",
-            pre_inversion_condition_number,
+            "You are inverting a matrix with a potentially large condition number: %s",
+            condition_number,
         )
-        if (
-            inverse_stabilization_method
-            == InverseStabilizationMethods.TRIM_SMALL_SINGULAR_VALUES
-        ):
-            logger.info("Trimming small singular values to improve conditioning.")
-            u, s, vT = np.linalg.svd(matrix, full_matrices=False)
-            logger.info(
-                " Sorted singular values: %s",
-                s,
-            )
-            sing_values_above_threshold_cond = s > s.max() / condition_num_threshold
-            if not np.any(sing_values_above_threshold_cond):
-                raise RuntimeError(
-                    f"All singular values are below the threshold of {s.max() / condition_num_threshold}. Singular value trimming will not work.",
-                )
-            trimmed_pseudoinverse = (
-                vT.T[:, sing_values_above_threshold_cond]
-                / s[sing_values_above_threshold_cond]
-            ) @ u[:, sing_values_above_threshold_cond].T
-            inverse = trimmed_pseudoinverse
-            pre_inversion_condition_number = (
-                s[sing_values_above_threshold_cond].max()
-                / s[sing_values_above_threshold_cond].min()
-            )
-
-            logger.info(
-                "Kept %s out of %s singular values. Condition number of resulting lower-rank-approximation before inversion: %s",
-                sum(sing_values_above_threshold_cond),
-                len(s),
-                pre_inversion_condition_number,
-            )
-        elif (
-            inverse_stabilization_method
-            == InverseStabilizationMethods.ADD_RIDGE_FIXED_CONDITION_NUMBER
-        ):
-            logger.info("Adding ridge/Tikhonov regularization to improve conditioning.")
-            _, singular_values, _ = np.linalg.svd(matrix, full_matrices=False)
-            logger.info(
-                "Using fixed condition number threshold of %s to determine lambda.",
-                condition_num_threshold,
-            )
-            lambda_ = (
-                singular_values.max() / condition_num_threshold - singular_values.min()
-            )
-            logger.info("Lambda for ridge regularization: %s", lambda_)
-            new_matrix = matrix + lambda_ * np.eye(matrix.shape[0])
-            pre_inversion_condition_number = np.linalg.cond(new_matrix)
-            logger.info(
-                "Condition number of matrix after ridge regularization: %s",
-                pre_inversion_condition_number,
-            )
-            inverse = np.linalg.solve(new_matrix, np.eye(matrix.shape[0]))
-        elif (
-            inverse_stabilization_method
-            == InverseStabilizationMethods.ADD_RIDGE_MEDIAN_SINGULAR_VALUE_FRACTION
-        ):
-            logger.info("Adding ridge/Tikhonov regularization to improve conditioning.")
-            _, singular_values, _ = np.linalg.svd(matrix, full_matrices=False)
-            logger.info(
-                "Using median singular value times %s as lambda.",
-                ridge_median_singular_value_fraction,
-            )
-            lambda_ = ridge_median_singular_value_fraction * np.median(singular_values)
-            logger.info("Lambda for ridge regularization: %s", lambda_)
-            new_matrix = matrix + lambda_ * np.eye(matrix.shape[0])
-            pre_inversion_condition_number = np.linalg.cond(new_matrix)
-            logger.info(
-                "Condition number of matrix after ridge regularization: %s",
-                pre_inversion_condition_number,
-            )
-            inverse = np.linalg.solve(new_matrix, np.eye(matrix.shape[0]))
-        elif (
-            inverse_stabilization_method
-            == InverseStabilizationMethods.INVERSE_BREAD_STRUCTURE_AWARE_INVERSION
-        ):
-            if not beta_dim or not theta_dim:
-                raise ValueError(
-                    "When using structure-aware inversion, beta_dim and theta_dim must be provided."
-                )
-            logger.info(
-                "Using inverse bread's block lower triangular structure to invert only diagonal blocks."
-            )
-            pre_inversion_condition_number = np.linalg.cond(matrix)
-            inverse = invert_inverse_bread_matrix(
-                matrix,
-                beta_dim,
-                theta_dim,
-                InverseStabilizationMethods.ADD_RIDGE_FIXED_CONDITION_NUMBER,
-            )
-        elif (
-            inverse_stabilization_method
-            == InverseStabilizationMethods.ZERO_OUT_SMALL_OFF_DIAGONALS
-        ):
-            if not beta_dim or not theta_dim:
-                raise ValueError(
-                    "When zeroing out small off diagonals, beta_dim and theta_dim must be provided."
-                )
-            logger.info(
-                "Zeroing out small off-diagonal blocks to improve conditioning."
-            )
-            zeroed_matrix = zero_small_off_diagonal_blocks(
-                matrix,
-                ([beta_dim] * (matrix.shape[0] // beta_dim)) + [theta_dim],
-            )
-            pre_inversion_condition_number = np.linalg.cond(zeroed_matrix)
-            logger.info(
-                "Condition number of matrix after zeroing out small off-diagonal blocks: %s",
-                pre_inversion_condition_number,
-            )
-            inverse = np.linalg.solve(zeroed_matrix, np.eye(zeroed_matrix.shape[0]))
-        elif (
-            inverse_stabilization_method
-            == InverseStabilizationMethods.ALL_METHODS_COMPETITION
-        ):
-            # TODO: Choose right metric for competition... identity diff might not be it.
-            raise NotImplementedError(
-                "All methods competition is not implemented yet. Please choose a specific method."
-            )
-        elif inverse_stabilization_method == InverseStabilizationMethods.NONE:
-            logger.info("No inverse stabilization method applied. Inverting directly.")
-        else:
-            raise ValueError(
-                f"Unknown inverse stabilization method: {inverse_stabilization_method}"
-            )
     if inverse is None:
         inverse = np.linalg.solve(matrix, np.eye(matrix.shape[0]))
-    return inverse, pre_inversion_condition_number
+    return inverse, condition_number
 
 
 def zero_small_off_diagonal_blocks(
@@ -183,7 +53,7 @@ def zero_small_off_diagonal_blocks(
     Zero off-diagonal blocks whose Frobenius norm is < frobenius_norm_threshold_fraction x
     Frobenius norm of the diagonal block in the same ROW. One could compare to
     the same column or both the row and column, but we choose row here since
-    rows correspond to a single RL update or inference step in the adaptive bread
+    rows correspond to a single RL update or inference step in the bread
     inverse matrices this method is designed for.
 
     Args:
@@ -237,18 +107,17 @@ def zero_small_off_diagonal_blocks(
     return J_trim
 
 
-def invert_inverse_bread_matrix(
-    inverse_bread,
+def invert_bread_matrix(
+    bread,
     beta_dim,
     theta_dim,
-    diag_inverse_stabilization_method=InverseStabilizationMethods.TRIM_SMALL_SINGULAR_VALUES,
 ):
     """
-    Invert the inverse bread matrix to get the bread matrix.  This is a special
+    Invert the bread matrix to get the inverse bread matrix.  This is a special
     function in order to take advantage of the block lower triangular structure.
 
     The procedure is as follows:
-    1. Initialize the inverse matrix B = A^{-1} as a block lower triangular matrix
+    1. Initialize the matrix B = A^{-1} as a block lower triangular matrix
        with the same block structure as A.
 
     2. Compute the diagonal blocks B_{ii}:
@@ -260,24 +129,23 @@ def invert_inverse_bread_matrix(
            B_{ij} = -A_{ii}^{-1} * sum(A_{ik} * B_{kj} for k in range(j, i))
     """
     blocks = []
-    num_beta_block_rows = (inverse_bread.shape[0] - theta_dim) // beta_dim
+    num_beta_block_rows = (bread.shape[0] - theta_dim) // beta_dim
 
     # Create upper rows of block of bread (just the beta portion)
     for i in range(0, num_beta_block_rows):
         beta_block_row = []
         beta_diag_inverse = invert_matrix_and_check_conditioning(
-            inverse_bread[
+            bread[
                 beta_dim * i : beta_dim * (i + 1),
                 beta_dim * i : beta_dim * (i + 1),
             ],
-            diag_inverse_stabilization_method,
         )[0]
         for j in range(0, num_beta_block_rows):
             if i > j:
                 beta_block_row.append(
                     -beta_diag_inverse
                     @ sum(
-                        inverse_bread[
+                        bread[
                             beta_dim * i : beta_dim * (i + 1),
                             beta_dim * k : beta_dim * (k + 1),
                         ]
@@ -299,17 +167,16 @@ def invert_inverse_bread_matrix(
     # Create the bottom block row of bread (the theta portion)
     theta_block_row = []
     theta_diag_inverse = invert_matrix_and_check_conditioning(
-        inverse_bread[
+        bread[
             -theta_dim:,
             -theta_dim:,
         ],
-        diag_inverse_stabilization_method,
     )[0]
     for k in range(0, num_beta_block_rows):
         theta_block_row.append(
             -theta_diag_inverse
             @ sum(
-                inverse_bread[
+                bread[
                     -theta_dim:,
                     beta_dim * h : beta_dim * (h + 1),
                 ]
@@ -378,9 +245,9 @@ def confirm_input_check_result(message, suppress_interaction, error=None):
             print("\nPlease enter 'y' or 'n'.\n")
 
 
-def get_in_study_df_column(study_df, col_name, in_study_col_name):
+def get_active_df_column(analysis_df, col_name, active_col_name):
     return jnp.array(
-        study_df.loc[study_df[in_study_col_name] == 1, col_name]
+        analysis_df.loc[analysis_df[active_col_name] == 1, col_name]
         .to_numpy()
         .reshape(-1, 1)
     )
@@ -408,7 +275,7 @@ def get_radon_nikodym_weight(
     action_prob_func: callable,
     action_prob_func_args_beta_index: int,
     action: int,
-    *action_prob_func_args_single_user: tuple[Any, ...],
+    *action_prob_func_args_single_subject: tuple[Any, ...],
 ):
     """
     Computes a ratio of action probabilities under two sets of algorithm parameters:
@@ -426,13 +293,13 @@ def get_radon_nikodym_weight(
             The beta value to use in the denominator. NOT involved in differentation!
         action_prob_func (callable):
             The function used to compute the probability of action 1 at a given decision time for
-            a particular user given their state and the algorithm parameters.
+            a particular subject given their state and the algorithm parameters.
         action_prob_func_args_beta_index (int):
             The index of the beta argument in the action probability function's arguments.
         action (int):
             The actual taken action at the relevant decision time.
-        *action_prob_func_args_single_user (tuple[Any, ...]):
-            The arguments to the action probability function for the relevant user at this time.
+        *action_prob_func_args_single_subject (tuple[Any, ...]):
+            The arguments to the action probability function for the relevant subject at this time.
 
     Returns:
         jnp.float32: The Radon-Nikodym weight.
@@ -440,15 +307,17 @@ def get_radon_nikodym_weight(
     """
 
     # numerator
-    pi_beta = action_prob_func(*action_prob_func_args_single_user)
+    pi_beta = action_prob_func(*action_prob_func_args_single_subject)
 
     # denominator, where we thread in beta_target so that differentiation with respect to the
     # original beta in the arguments leaves this alone.
-    beta_target_action_prob_func_args_single_user = [*action_prob_func_args_single_user]
-    beta_target_action_prob_func_args_single_user[action_prob_func_args_beta_index] = (
-        beta_target
-    )
-    pi_beta_target = action_prob_func(*beta_target_action_prob_func_args_single_user)
+    beta_target_action_prob_func_args_single_subject = [
+        *action_prob_func_args_single_subject
+    ]
+    beta_target_action_prob_func_args_single_subject[
+        action_prob_func_args_beta_index
+    ] = beta_target
+    pi_beta_target = action_prob_func(*beta_target_action_prob_func_args_single_subject)
 
     return conditional_x_or_one_minus_x(pi_beta, action) / conditional_x_or_one_minus_x(
         pi_beta_target, action
@@ -456,7 +325,7 @@ def get_radon_nikodym_weight(
 
 
 def get_min_time_by_policy_num(
-    single_user_policy_num_by_decision_time, beta_index_by_policy_num
+    single_subject_policy_num_by_decision_time, beta_index_by_policy_num
 ):
     """
     Returns a dictionary mapping each policy number to the first time it was applicable,
@@ -464,12 +333,12 @@ def get_min_time_by_policy_num(
     """
     min_time_by_policy_num = {}
     first_time_after_first_update = None
-    for decision_time, policy_num in single_user_policy_num_by_decision_time.items():
+    for decision_time, policy_num in single_subject_policy_num_by_decision_time.items():
         if policy_num not in min_time_by_policy_num:
             min_time_by_policy_num[policy_num] = decision_time
 
         # Grab the first time where a non-initial, non-fallback policy is used.
-        # Assumes single_user_policy_num_by_decision_time is sorted.
+        # Assumes single_subject_policy_num_by_decision_time is sorted.
         if (
             policy_num in beta_index_by_policy_num
             and first_time_after_first_update is None
@@ -494,10 +363,10 @@ def calculate_beta_dim(
         int: The dimension of the beta vector.
     """
     for decision_time in action_prob_func_args:
-        for user_id in action_prob_func_args[decision_time]:
-            if action_prob_func_args[decision_time][user_id]:
+        for subject_id in action_prob_func_args[decision_time]:
+            if action_prob_func_args[decision_time][subject_id]:
                 return len(
-                    action_prob_func_args[decision_time][user_id][
+                    action_prob_func_args[decision_time][subject_id][
                         action_prob_func_args_beta_index
                     ]
                 )
@@ -507,7 +376,7 @@ def calculate_beta_dim(
 
 
 def construct_beta_index_by_policy_num_map(
-    study_df: pd.DataFrame, policy_num_col_name: str, in_study_col_name: str
+    analysis_df: pd.DataFrame, policy_num_col_name: str, active_col_name: str
 ) -> tuple[dict[int | float, int], int | float]:
     """
     Constructs a mapping from non-initial, non-fallback policy numbers to the index of the
@@ -524,8 +393,9 @@ def construct_beta_index_by_policy_num_map(
     """
 
     unique_sorted_non_fallback_policy_nums = sorted(
-        study_df[
-            (study_df[policy_num_col_name] >= 0) & (study_df[in_study_col_name] == 1)
+        analysis_df[
+            (analysis_df[policy_num_col_name] >= 0)
+            & (analysis_df[active_col_name] == 1)
         ][policy_num_col_name]
         .unique()
         .tolist()
@@ -550,10 +420,10 @@ def collect_all_post_update_betas(
     """
     all_post_update_betas = []
     for policy_num in sorted(beta_index_by_policy_num.keys()):
-        for user_id in alg_update_func_args[policy_num]:
-            if alg_update_func_args[policy_num][user_id]:
+        for subject_id in alg_update_func_args[policy_num]:
+            if alg_update_func_args[policy_num][subject_id]:
                 all_post_update_betas.append(
-                    alg_update_func_args[policy_num][user_id][
+                    alg_update_func_args[policy_num][subject_id][
                         alg_update_func_args_beta_index
                     ]
                 )
@@ -561,27 +431,31 @@ def collect_all_post_update_betas(
     return jnp.array(all_post_update_betas)
 
 
-def extract_action_and_policy_by_decision_time_by_user_id(
-    study_df,
-    user_id_col_name,
-    in_study_col_name,
+def extract_action_and_policy_by_decision_time_by_subject_id(
+    analysis_df,
+    subject_id_col_name,
+    active_col_name,
     calendar_t_col_name,
     action_col_name,
     policy_num_col_name,
 ):
-    action_by_decision_time_by_user_id = {}
-    policy_num_by_decision_time_by_user_id = {}
-    for user_id, user_df in study_df.groupby(user_id_col_name):
-        in_study_user_df = user_df[user_df[in_study_col_name] == 1]
-        action_by_decision_time_by_user_id[user_id] = dict(
+    action_by_decision_time_by_subject_id = {}
+    policy_num_by_decision_time_by_subject_id = {}
+    for subject_id, subject_df in analysis_df.groupby(subject_id_col_name):
+        active_subject_df = subject_df[subject_df[active_col_name] == 1]
+        action_by_decision_time_by_subject_id[subject_id] = dict(
             zip(
-                in_study_user_df[calendar_t_col_name], in_study_user_df[action_col_name]
+                active_subject_df[calendar_t_col_name],
+                active_subject_df[action_col_name],
             )
         )
-        policy_num_by_decision_time_by_user_id[user_id] = dict(
+        policy_num_by_decision_time_by_subject_id[subject_id] = dict(
             zip(
-                in_study_user_df[calendar_t_col_name],
-                in_study_user_df[policy_num_col_name],
+                active_subject_df[calendar_t_col_name],
+                active_subject_df[policy_num_col_name],
             )
         )
-    return action_by_decision_time_by_user_id, policy_num_by_decision_time_by_user_id
+    return (
+        action_by_decision_time_by_subject_id,
+        policy_num_by_decision_time_by_subject_id,
+    )