lifejacket 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

lifejacket/after_study_analysis.py

@@ -128,6 +128,12 @@ def cli():
     default=-1000,
     help="Index of the argument holding the decision times the action probabilities correspond to in the tuple of algorithm update func args, if applicable.",
 )
+@click.option(
+    "--alg_update_func_args_previous_betas_index",
+    type=int,
+    default=-1000,
+    help="Index of the previous betas array in the tuple of algorithm update func args, if applicable. Note that these are only post-update betas. Sometimes a beta_0 may be defined pre-update; this should not be in here.",
+)
 @click.option(
     "--inference_func_filename",
     type=click.Path(exists=True),
@@ -299,6 +305,7 @@ def analyze_dataset(
     alg_update_func_args_beta_index: int,
     alg_update_func_args_action_prob_index: int,
     alg_update_func_args_action_prob_times_index: int,
+    alg_update_func_args_previous_betas_index: int,
     inference_func: Callable,
     inference_func_type: str,
     inference_func_args_theta_index: int,
@@ -421,6 +428,7 @@ def analyze_dataset(
             alg_update_func_args_beta_index,
             alg_update_func_args_action_prob_index,
             alg_update_func_args_action_prob_times_index,
+            alg_update_func_args_previous_betas_index,
             theta_est,
             beta_dim,
             suppress_interactive_data_checks,
@@ -497,6 +505,7 @@ def analyze_dataset(
         alg_update_func_args_beta_index,
         alg_update_func_args_action_prob_index,
         alg_update_func_args_action_prob_times_index,
+        alg_update_func_args_previous_betas_index,
         inference_func,
         inference_func_type,
         inference_func_args_theta_index,
@@ -1019,6 +1028,7 @@ def get_avg_weighted_estimating_function_stacks_and_aux_values(
     alg_update_func_args_beta_index: int,
     alg_update_func_args_action_prob_index: int,
     alg_update_func_args_action_prob_times_index: int,
+    alg_update_func_args_previous_betas_index: int,
     inference_func: callable,
     inference_func_type: str,
     inference_func_args_theta_index: int,
@@ -1075,6 +1085,8 @@ def get_avg_weighted_estimating_function_stacks_and_aux_values(
         alg_update_func_args_action_prob_times_index (int):
             The index in the update function arguments tuple where an array of times for which the
             given action probabilities apply is provided, if applicable. -1 otherwise.
+        alg_update_func_args_previous_betas_index (int):
+            The index in the update function arguments tuple where previous betas are provided.
         inference_func (callable):
             The inference loss or estimating function.
         inference_func_type (str):
@@ -1179,6 +1191,7 @@ def get_avg_weighted_estimating_function_stacks_and_aux_values(
         alg_update_func_args_beta_index,
         alg_update_func_args_action_prob_index,
         alg_update_func_args_action_prob_times_index,
+        alg_update_func_args_previous_betas_index,
         threaded_action_prob_func_args_by_decision_time_by_user_id,
         action_prob_func,
     )
@@ -1278,6 +1291,7 @@ def construct_classical_and_adaptive_sandwiches(
     alg_update_func_args_beta_index: int,
     alg_update_func_args_action_prob_index: int,
     alg_update_func_args_action_prob_times_index: int,
+    alg_update_func_args_previous_betas_index: int,
     inference_func: callable,
     inference_func_type: str,
     inference_func_args_theta_index: int,
@@ -1354,6 +1368,8 @@ def construct_classical_and_adaptive_sandwiches(
         alg_update_func_args_action_prob_times_index (int):
             The index in the update function arguments tuple where an array of times for which the
             given action probabilities apply is provided, if applicable. -1 otherwise.
+        alg_update_func_args_previous_betas_index (int):
+            The index in the update function arguments tuple where the previous betas are provided.
         inference_func (callable):
             The inference loss or estimating function.
         inference_func_type (str):
@@ -1463,6 +1479,7 @@ def construct_classical_and_adaptive_sandwiches(
         alg_update_func_args_beta_index,
         alg_update_func_args_action_prob_index,
         alg_update_func_args_action_prob_times_index,
+        alg_update_func_args_previous_betas_index,
         inference_func,
         inference_func_type,
         inference_func_args_theta_index,

lifejacket/arg_threading_helpers.py

@@ -63,6 +63,7 @@ def thread_action_prob_func_args(
         action_prob_func_args_beta_index (int):
             The index in the action probability function arguments tuple
             where the beta value should be inserted.
+
     Returns:
         dict[collections.abc.Hashable, dict[int, tuple[Any, ...]]]:
             A map from user ids to maps of decision times to action probability function
@@ -125,6 +126,7 @@ def thread_update_func_args(
     alg_update_func_args_beta_index: int,
     alg_update_func_args_action_prob_index: int,
     alg_update_func_args_action_prob_times_index: int,
+    alg_update_func_args_previous_betas_index: int,
     threaded_action_prob_func_args_by_decision_time_by_user_id: dict[
         collections.abc.Hashable, dict[int, tuple[Any, ...]]
     ],
@@ -165,6 +167,9 @@ def thread_update_func_args(
             to the update function, this is the index in the arguments where an array of times for
             which the given action probabilities apply is provided.
 
+        alg_update_func_args_previous_betas_index (int):
+            The index in the update function with previous beta parameters
+
         threaded_action_prob_func_args_by_decision_time_by_user_id (dict[collections.abc.Hashable, dict[int, tuple[Any, ...]]]):
             A dictionary mapping decision times to the function arguments required to compute action
             probabilities for this user, and with the shared betas thread in.
@@ -206,6 +211,24 @@ def thread_update_func_args(
                     beta_to_introduce,
                 )
             )
+            if alg_update_func_args_previous_betas_index >= 0:
+                previous_betas_to_introduce = all_post_update_betas[
+                    : len(
+                        update_func_args_by_user_id[user_id][
+                            alg_update_func_args_previous_betas_index
+                        ]
+                    )
+                ]
+                if previous_betas_to_introduce.size > 0:
+                    threaded_update_func_args_by_policy_num_by_user_id[user_id][
+                        policy_num
+                    ] = replace_tuple_index(
+                        threaded_update_func_args_by_policy_num_by_user_id[user_id][
+                            policy_num
+                        ],
+                        alg_update_func_args_previous_betas_index,
+                        previous_betas_to_introduce,
+                    )
 
             if alg_update_func_args_action_prob_index >= 0:
                 logger.debug(

lifejacket/calculate_derivatives.py

@@ -743,7 +743,7 @@ def get_loss_gradient_derivatives_wrt_pi_batched(
     *batched_arg_tensors,
 ):
     if update_func_type == FunctionTypes.LOSS:
-        return jax.jit(
+        return jax.jit(  # pylint: disable=not-callable
             jax.vmap(
                 fun=jax.jacrev(
                     jax.grad(update_func, update_func_args_beta_index),
@@ -754,7 +754,7 @@ def get_loss_gradient_derivatives_wrt_pi_batched(
             )
         )(*batched_arg_tensors)
     if update_func_type == FunctionTypes.ESTIMATING:
-        return jax.jit(
+        return jax.jit(  # pylint: disable=not-callable
             jax.vmap(
                 fun=jax.jacrev(
                     update_func,

lifejacket/form_adaptive_meat_adjustments_directly.py

@@ -62,7 +62,7 @@ def form_adaptive_meat_adjustments_directly(
     # 3. Split the effective M blocks into (theta_dim, beta_dim) blocks and the
     # estimating function stacks into (num_updates, beta_dim) stacks.
 
-    # effective_M_blocks is shape (theta_dim, num_updates * beta_dim)
+    # effective_M_blocks_together is shape (theta_dim, num_updates * beta_dim)
     # We want to split it into a list of (theta_dim, beta_dim) arrays
     M_blocks = np.split(
         effective_M_blocks_together,

lifejacket/input_checks.py

@@ -41,6 +41,7 @@ def perform_first_wave_input_checks(
     alg_update_func_args_beta_index,
     alg_update_func_args_action_prob_index,
     alg_update_func_args_action_prob_times_index,
+    alg_update_func_args_previous_betas_index,
     theta_est,
     beta_dim,
     suppress_interactive_data_checks,
@@ -56,6 +57,9 @@ def perform_first_wave_input_checks(
     require_beta_is_1D_array_in_alg_update_args(
         alg_update_func_args, alg_update_func_args_beta_index
     )
+    require_previous_betas_is_2D_array_in_alg_update_args(
+        alg_update_func_args, alg_update_func_args_previous_betas_index
+    )
     require_all_policy_numbers_in_study_df_except_possibly_initial_and_fallback_present_in_alg_update_args(
         study_df, in_study_col_name, policy_num_col_name, alg_update_func_args
     )
@@ -73,6 +77,9 @@ def perform_first_wave_input_checks(
     require_betas_match_in_alg_update_args_each_update(
         alg_update_func_args, alg_update_func_args_beta_index
     )
+    require_previous_betas_match_in_alg_update_args_each_update(
+        alg_update_func_args, alg_update_func_args_previous_betas_index
+    )
     require_action_prob_args_in_alg_update_func_correspond_to_study_df(
         study_df,
         action_prob_col_name,
@@ -675,6 +682,23 @@ def require_beta_is_1D_array_in_alg_update_args(
                 == 1
             ), "Beta is not a 1D array in the algorithm update function args."
 
+def require_previous_betas_is_2D_array_in_alg_update_args(
+    alg_update_func_args, alg_update_func_args_previous_betas_index
+):
+    if alg_update_func_args_previous_betas_index < 0:
+        return
+    
+    for policy_num in alg_update_func_args:
+        for user_id in alg_update_func_args[policy_num]:
+            if not alg_update_func_args[policy_num][user_id]:
+                continue
+            assert (
+                alg_update_func_args[policy_num][user_id][
+                    alg_update_func_args_previous_betas_index
+                ].ndim
+                == 2
+            ), "Previous betas is not a 2D array in the algorithm update function args."
+
 
 def require_beta_is_1D_array_in_action_prob_args(
     action_prob_func_args, action_prob_func_args_beta_index
@@ -827,6 +851,30 @@ def require_betas_match_in_alg_update_args_each_update(
                     beta, first_beta
                 ), f"Betas do not match across users in the algorithm update function args for policy number {policy_num}. Please see the contract for details."
 
+def require_previous_betas_match_in_alg_update_args_each_update(
+    alg_update_func_args, alg_update_func_args_previous_betas_index
+):
+    logger.info(
+        "Checking that previous betas match across users for each update in the algorithm update function args."
+    )
+    if alg_update_func_args_previous_betas_index < 0:
+        return 
+    
+    for policy_num in alg_update_func_args:
+        first_previous_betas = None
+        for user_id in alg_update_func_args[policy_num]:
+            if not alg_update_func_args[policy_num][user_id]:
+                continue
+            previous_betas = alg_update_func_args[policy_num][user_id][
+                alg_update_func_args_previous_betas_index
+            ]
+            if first_previous_betas is None:
+                first_previous_betas = previous_betas
+            else:
+                assert np.array_equal(
+                    previous_betas, first_previous_betas
+                ), f"Previous betas do not match across users in the algorithm update function args for policy number {policy_num}. Please see the contract for details."
+
 
 def require_betas_match_in_action_prob_func_args_each_decision(
     action_prob_func_args, action_prob_func_args_beta_index

{lifejacket-0.1.0.dist-info → lifejacket-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lifejacket
-Version: 0.1.0
+Version: 0.2.0
 Summary: A package for after-study analysis of adaptive experiments in which data is pooled across users.
 Author-email: Nowell Closser <nowellclosser@gmail.com>
 Requires-Python: >=3.10

lifejacket-0.2.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+lifejacket/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+lifejacket/after_study_analysis.py,sha256=_Weeca51EXWlQkx3IWd2t1jImBxODOTe9-8gtKBSlus,82168
+lifejacket/arg_threading_helpers.py,sha256=7HV7qkiJtm-E-cpAhtv4n_BpCVjz9tC0nGENbc090h8,17282
+lifejacket/calculate_derivatives.py,sha256=SceXFWtK56uCCdXGD7v8JijgYz0UCBKzcnrPH_nAqNE,37536
+lifejacket/constants.py,sha256=2L05p6NJ7l3qRZ1hD2KlrvzWF1ReSmWRUkULPIhdvJo,842
+lifejacket/form_adaptive_meat_adjustments_directly.py,sha256=bSLrVYLZR1-Qlm5yIdktzv8ZQTVhHTlhVVL2wEjLTmw,13737
+lifejacket/get_datum_for_blowup_supervised_learning.py,sha256=V8H4PE49dQwsKjj93QEu2BKbhwPr56QMtx2jhan39-c,58357
+lifejacket/helper_functions.py,sha256=xOhRG-Cm4ZdRNm-O0faHna583d74pLWY5_jfnokegWc,23295
+lifejacket/input_checks.py,sha256=KcDdfsdCVCKKcx07FfOKJb3KVX6xFuWwAufGJ3msAuc,46972
+lifejacket/small_sample_corrections.py,sha256=f8jmg9U9ZN77WadJud70tt6NMxCTsSGtlsdF_-mfws4,5543
+lifejacket/trial_conditioning_monitor.py,sha256=qNTHh0zY2P7zJxox_OwhEEK8Ed1l0TPOjGDqNxMNoIQ,42164
+lifejacket/vmap_helpers.py,sha256=pZqYN3p9Ty9DPOeeY9TKbRJXR2AV__HBwwDFOvdOQ84,2688
+lifejacket-0.2.0.dist-info/METADATA,sha256=XzZQXRKzohTg0FOv0sIh7Ii6Snj2m5HVIdeXVYUphe4,7274
+lifejacket-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lifejacket-0.2.0.dist-info/entry_points.txt,sha256=4k8ibVIUT-OHxPaaDv-QwWpC64ErzhdemHpTAXCnb8w,67
+lifejacket-0.2.0.dist-info/top_level.txt,sha256=vKl8m7jOQ4pkbzVuHCJsq-8LcXRrOAWnok3bBo9qpsE,11
+lifejacket-0.2.0.dist-info/RECORD,,

lifejacket-0.1.0.dist-info/RECORD

@@ -1,17 +0,0 @@
-lifejacket/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lifejacket/after_study_analysis.py,sha256=HtAHRYCni8Lh7HCB9A0SsRkyAol2FbczrkxFn1ww1HA,81170
-lifejacket/arg_threading_helpers.py,sha256=kxtGlN_B1C1WKcWtJMlgWJwzzI1THytflmWsL5ZML7k,16228
-lifejacket/calculate_derivatives.py,sha256=3rYukD1wbjDof7d4_3QdQ-A4GSK9H8z8HJsbNQh0DzA,37472
-lifejacket/constants.py,sha256=2L05p6NJ7l3qRZ1hD2KlrvzWF1ReSmWRUkULPIhdvJo,842
-lifejacket/form_adaptive_meat_adjustments_directly.py,sha256=_BaziGfYjEySN78nU3lCrVtf2KWIuZ8PmzfMZypAaWI,13728
-lifejacket/get_datum_for_blowup_supervised_learning.py,sha256=V8H4PE49dQwsKjj93QEu2BKbhwPr56QMtx2jhan39-c,58357
-lifejacket/helper_functions.py,sha256=xOhRG-Cm4ZdRNm-O0faHna583d74pLWY5_jfnokegWc,23295
-lifejacket/input_checks.py,sha256=A0f2owqRUjeBAh5jLULKu1nXW1SgZDR5eK7xBm1ahZw,44878
-lifejacket/small_sample_corrections.py,sha256=f8jmg9U9ZN77WadJud70tt6NMxCTsSGtlsdF_-mfws4,5543
-lifejacket/trial_conditioning_monitor.py,sha256=qNTHh0zY2P7zJxox_OwhEEK8Ed1l0TPOjGDqNxMNoIQ,42164
-lifejacket/vmap_helpers.py,sha256=pZqYN3p9Ty9DPOeeY9TKbRJXR2AV__HBwwDFOvdOQ84,2688
-lifejacket-0.1.0.dist-info/METADATA,sha256=VT6H9TNcYleRhp-wLda4HXrbI7EJYj6ZV0_7K5fraI4,7274
-lifejacket-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lifejacket-0.1.0.dist-info/entry_points.txt,sha256=4k8ibVIUT-OHxPaaDv-QwWpC64ErzhdemHpTAXCnb8w,67
-lifejacket-0.1.0.dist-info/top_level.txt,sha256=vKl8m7jOQ4pkbzVuHCJsq-8LcXRrOAWnok3bBo9qpsE,11
-lifejacket-0.1.0.dist-info/RECORD,,