Fix random number draws in simulation (#186)

Co-authored-by: Sebastian Gsell <sebastian.gsell93@gmail.com>

Author: MaxBlesch

.github/workflows/main.yml

@@ -21,9 +21,9 @@ jobs:
         os:
           - ubuntu-latest
         python-version:
-          - '3.10'
           - '3.11'
           - '3.12'
+          - '3.13'
     steps:
       - uses: actions/checkout@v3
       - name: create build environment

.pre-commit-config.yaml

@@ -10,11 +10,11 @@ repos:
     hooks:
       - id: yamllint
   - repo: https://github.com/lyz-code/yamlfix
-    rev: 1.17.0
+    rev: 1.19.0
     hooks:
       - id: yamlfix
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: check-added-large-files
         args:
@@ -48,7 +48,7 @@ repos:
   #       args:
   #         - --py37-plus
   - repo: https://github.com/pycqa/isort
-    rev: 6.0.1
+    rev: 7.0.0
     hooks:
       - id: isort
         name: isort
@@ -58,8 +58,8 @@ repos:
 #    rev: v2.8.0
 #    hooks:
 #      - id: setup-cfg-fmt
-  - repo: https://github.com/psf/black
-    rev: 25.1.0
+  - repo: https://github.com/psf/black-pre-commit-mirror
+    rev: 25.9.0
     hooks:
       - id: black
         language_version: python3.12
@@ -90,7 +90,7 @@ repos:
       - id: nbqa-ruff
         exclude: tests/sandbox/
   - repo: https://github.com/executablebooks/mdformat
-    rev: 0.7.22
+    rev: 1.0.0
     hooks:
       - id: mdformat
         additional_dependencies:

src/dcegm/simulation/sim_utils.py

@@ -279,11 +279,12 @@ def vectorized_utility(consumption_period, state, choice, params, compute_utilit
 def realize_stochastic_states(state, choice, key, params, processed_stochastic_funcs):
     stochastic_states_next_period = {}
     for state_name in processed_stochastic_funcs.keys():
+        key, subkey = jax.random.split(key)
         state_vec = processed_stochastic_funcs[state_name](
             params=params, **state, choice=choice
         )
         stochastic_states_next_period[state_name] = jax.random.choice(
-            key=key, a=state_vec.shape[0], p=state_vec
+            key=subkey, a=state_vec.shape[0], p=state_vec
         )
     return stochastic_states_next_period