lifejacket 0.1.0__py3-none-any.whl

lifejacket/vmap_helpers.py

@@ -0,0 +1,71 @@
+import collections
+
+from jax import numpy as jnp
+import numpy as np
+
+
+# TODO: Check for exactly the required types earlier
+# TODO: Try except and nice error message
+# TODO: This is complicated enough to deserve its own unit tests
+def stack_batched_arg_lists_into_tensors(batched_arg_lists):
+    """
+    Stack a simple Python list of lists of function arguments into a list of jnp arrays that can be
+    supplied to vmap as batch arguments. vmap requires all elements of such a batched array to be
+    the same shape, as do the stacking functions we use here.  Thus we require this be called on
+    batches with the same data shape. We also supply the axes one must iterate over to get
+    each users's args in a batch.
+    """
+
+    batched_arg_tensors = []
+
+    # This ends up being all zeros because of the way we are (now) doing the
+    # stacking, but better to not assume that externally and send out what
+    # we've done with this list.
+    batch_axes = []
+
+    for batched_arg_list in batched_arg_lists:
+        if (
+            isinstance(
+                batched_arg_list[0],
+                (jnp.ndarray, np.ndarray),
+            )
+            and batched_arg_list[0].ndim > 2
+        ):
+            raise TypeError("Arrays with dimension greater that 2 are not supported.")
+        if (
+            isinstance(
+                batched_arg_list[0],
+                (jnp.ndarray, np.ndarray),
+            )
+            and batched_arg_list[0].ndim == 2
+        ):
+            ########## We have a matrix (2D array) arg
+
+            batched_arg_tensors.append(jnp.stack(batched_arg_list, 0))
+            batch_axes.append(0)
+        elif isinstance(
+            batched_arg_list[0],
+            (collections.abc.Sequence, jnp.ndarray, np.ndarray),
+        ) and not isinstance(batched_arg_list[0], str):
+            ########## We have a vector (1D array) arg
+            if not isinstance(batched_arg_list[0], (jnp.ndarray, np.ndarray)):
+                try:
+                    batched_arg_list = [jnp.array(x) for x in batched_arg_list]
+                except Exception as e:
+                    raise TypeError(
+                        "Argument of sequence type that cannot be cast to JAX numpy array"
+                    ) from e
+            assert batched_arg_list[0].ndim == 1
+
+            batched_arg_tensors.append(jnp.vstack(batched_arg_list))
+            batch_axes.append(0)
+        else:
+            ########## Otherwise we should have a list of scalars.
+            # Just turn into a jnp array.
+            batched_arg_tensors.append(jnp.array(batched_arg_list))
+            batch_axes.append(0)
+
+    return (
+        batched_arg_tensors,
+        batch_axes,
+    )

lifejacket-0.1.0.dist-info/METADATA

@@ -0,0 +1,100 @@
+Metadata-Version: 2.4
+Name: lifejacket
+Version: 0.1.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
+Description-Content-Type: text/markdown
+Requires-Dist: click>=8.0
+Requires-Dist: jax>=0.4.0
+Requires-Dist: jaxlib>=0.4.0
+Requires-Dist: numpy>=1.20.0
+Requires-Dist: pandas>=1.3.0
+Requires-Dist: scipy>=1.7.0
+Requires-Dist: plotext>=5.0.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: black>=22.0; extra == "dev"
+Requires-Dist: flake8>=4.0; extra == "dev"
+
+```python
+  _ _  __     _            _        _
+ | (_)/ _|   (_)          | |      | |
+ | |_| |_ ___ _  __ _  ___| | _____| |_
+ | | |  _/ _ \ |/ _` |/ __| |/ / _ \ __|
+ | | | ||  __/ | (_| | (__|   <  __/ |_
+ |_|_|_| \___| |\__,_|\___|_|\_\___|\__|
+            _/ |
+           |__/
+```
+
+Save your standard errors from pooling in adaptive experiments.
+
+## Setup (if not using conda)
+### Create and activate a virtual environment
+- `python3 -m venv .venv; source /.venv/bin/activate`
+
+### Adding a package
+- Add to `requirements.txt` with a specific version or no version if you want the latest stable
+- Run `pip freeze > requirements.txt` to lock the versions of your package and all its subpackages
+
+## Running the code
+- `export PYTHONPATH to the absolute path of this repository on your computer
+- `./run_local_synthetic.sh`, which outputs to `simulated_data/` by default. See all the possible flags to be toggled in the script code.
+
+## Linting/Formatting
+
+## Testing
+python -m pytest
+python -m pytest tests/unit_tests
+python -m pytest tests/integration_tests
+
+
+# Talk about gitignored cluster simulation scripts
+
+
+
+
+
+
+### Important Large-Scale Simulations
+
+#### No adaptivity
+sbatch --array=[0-999] -t 0-5:00 --mem=50G run_and_analysis_parallel_synthetic --T=10 --n=50000 --recruit_n=50000 --steepness=0.0 --alg_state_feats=intercept,past_reward --action_centering_RL=0 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_no_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_no_action_centering.py"
+
+#### No adaptivity, 5 batches incremental recruitment
+sbatch --array=[0-999] -t 0-5:00 --mem=50G run_and_analysis_parallel_synthetic --T=10 --n=50000 --recruit_n=10000 --steepness=0.0 --alg_state_feats=intercept,past_reward --action_centering_RL=0 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_no_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_no_action_centering.py"
+
+#### Some adaptivity, no action_centering
+sbatch --array=[0-999] -t 0-5:00 --mem=50G run_and_analysis_parallel_synthetic --T=10 --n=50000 --recruit_n=50000 --steepness=3.0 --alg_state_feats=intercept,past_reward --action_centering_RL=0 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_no_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_no_action_centering.py"
+
+#### Some adaptivity, no action_centering, 5 batches incremental recruitment
+sbatch --array=[0-999] -t 0-5:00 --mem=50G run_and_analysis_parallel_synthetic --T=10 --n=50000 --recruit_n=10000 --steepness=3.0 --alg_state_feats=intercept,past_reward --action_centering_RL=0 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_no_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_no_action_centering.py"
+
+#### More adaptivity, no action_centering
+sbatch --array=[0-999] -t 0-5:00 --mem=50G run_and_analysis_parallel_synthetic --T=10 --n=50000 --recruit_n=50000 --steepness=5.0 --alg_state_feats=intercept,past_reward --action_centering_RL=0 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_no_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_no_action_centering.py"
+
+#### Even more adaptivity, no action_centering
+sbatch --array=[0-999] -t 0-5:00 --mem=50G run_and_analysis_parallel_synthetic --T=10 --n=50000 --recruit_n=50000 --steepness=10.0 --alg_state_feats=intercept,past_reward --action_centering_RL=0 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_no_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_no_action_centering.py"
+
+#### Some adaptivity, RL action_centering, no inference action centering
+sbatch --array=[0-999] -t 0-5:00 --mem=50G run_and_analysis_parallel_synthetic --T=10 --n=50000 --recruit_n=50000 --steepness=3.0 --alg_state_feats=intercept,past_reward --action_centering_RL=1 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_no_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_no_action_centering.py"
+
+#### Some adaptivity, inference action_centering, no RL action centering
+sbatch --array=[0-999] -t 0-5:00 --mem=50G run_and_analysis_parallel_synthetic --T=10 --n=50000 --recruit_n=50000 --steepness=3.0 --alg_state_feats=intercept,past_reward --action_centering_RL=0 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_action_centering.py"
+
+#### Some adaptivity, inference and RL action_centering
+sbatch --array=[0-999] -t 0-5:00 --mem=50G run_and_analysis_parallel_synthetic --T=10 --n=50000 --recruit_n=50000 --steepness=3.0 --alg_state_feats=intercept,past_reward --action_centering_RL=1 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_action_centering.py"
+
+#### Some adaptivity, inference and RL action_centering, even more T
+sbatch --array=[0-999] -t 1-00:00 --mem=50G run_and_analysis_parallel_synthetic --T=25 --n=50000 --recruit_n=50000 --steepness=3.0 --alg_state_feats=intercept,past_reward --action_centering_RL=1 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_no_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_no_action_centering.py"
+
+#### Some adaptivity, inference and RL action_centering, even more T, 5 batches incremental recruitment
+sbatch --array=[0-999] -t 1-00:00 --mem=50G run_and_analysis_parallel_synthetic --T=25 --n=50000 --recruit_n=10000 --steepness=3.0 --alg_state_feats=intercept,past_reward --action_centering_RL=1 --inference_loss_func_filename="functions_to_pass_to_analysis/get_least_squares_loss_inference_no_action_centering.py" --theta_calculation_func_filename="functions_to_pass_to_analysis/estimate_theta_least_squares_no_action_centering.py"
+
+
+
+## TODO
+1. Add precommit hooks (pip freeze, linting, formatting)
+2. Run tests on PRs
+

lifejacket-0.1.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+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,,

lifejacket-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

lifejacket-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+lifejacket = lifejacket.after_study_analysis:cli

lifejacket-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+lifejacket