Remove nearly all mentions of sparse variables from state-choice-space creation

Author: timmens

src/lcm/discrete_problem.py

@@ -196,27 +196,19 @@ def _segment_logsumexp(a, segment_info):
 def _determine_dense_discrete_choice_axes(
     variable_info: pd.DataFrame,
 ) -> tuple[int, ...] | None:
-    """Get axes of a state choice space that correspond to dense discrete choices.
-
-    Note: The dense choice axes determine over which axes we reduce the conditional
-    continuation values using a non-segmented operation. The axes ordering of the
-    conditional continuation value array is given by [sparse_variable, dense_variables].
-    The dense continuous choice dimension is already reduced as we are working with
-    the conditional continuation values.
+    """Get axes of a state-choice-space that correspond to discrete choices.
 
     Args:
-        variable_info (pd.DataFrame): DataFrame with information about the variables.
+        variable_info: DataFrame with information about the variables.
 
     Returns:
-        tuple[int, ...] | None: A tuple of indices representing the axes in the value
-            function that correspond to discrete choices. Returns None if there are no
-            discrete choice axes.
+        tuple[int, ...] | None: A tuple of indices representing the axes positions in
+            the value function that correspond to discrete choices. Returns None if
+            there are no discrete choice axes.
 
     """
     # List of dense variables excluding continuous choice variables.
-    axes = variable_info.query(
-        "~(is_choice & is_continuous)",
-    ).index.tolist()
+    axes = variable_info.query("~(is_choice & is_continuous)").index.tolist()
 
     choice_vars = set(variable_info.query("is_choice").index.tolist())
 

src/lcm/entry_point.py

@@ -94,11 +94,14 @@ def get_lcm_function(
 
         # call state space creation function, append trivial items to their lists
         # ==============================================================================
-        sc_space, space_info, state_indexer, segments = create_state_choice_space(
+        sc_space, space_info = create_state_choice_space(
             model=_mod,
             is_last_period=is_last_period,
         )
 
+        state_indexer = {}
+        segments = None
+
         state_choice_spaces.append(sc_space)
         choice_segments.append(segments)
 

src/lcm/state_space.py

@@ -3,23 +3,13 @@
 from lcm.interfaces import InternalModel, Space, SpaceInfo
 
 
-def create_state_choice_space(model: InternalModel, *, is_last_period: bool):
-    """Create a state choice space for the model.
+def create_state_choice_space(
+    model: InternalModel, *, is_last_period: bool
+) -> tuple[Space, SpaceInfo]:
+    """Create a state-choice-space for the model.
 
-    A state_choice_space is a compressed representation of all feasible states and the
-    feasible discrete choices within that state. We currently use the following
-    compressions:
-
-    We distinguish between dense and sparse variables (dense_vars and sparse_vars).
-    Dense state or choice variables are those whose set of feasible values does not
-    depend on any other state or choice variables. Sparse state or choice variables are
-    all other state variables. For dense state variables it is thus enough to store the
-    grid of feasible values (value_grid), whereas for sparse variables all feasible
-    combinations (combination_grid) have to be stored.
-
-    Note:
-    -----
-    - We only use the filter mask, not the forward mask (yet).
+    A state-choice-space is a compressed representation of all feasible states and the
+    feasible discrete choices within that state.
 
     Args:
         model (Model): A processed model.
@@ -30,9 +20,6 @@ def create_state_choice_space(model: InternalModel, *, is_last_period: bool):
             to execute a function on an entire space.
         SpaceInfo: A SpaceInfo object that contains all information needed to work with
             the output of a function evaluated on the space.
-        dict: Dictionary containing state indexer arrays.
-        jnp.ndarray: Jax array containing the choice segments needed for the emax
-            calculations.
 
     """
     # ==================================================================================
@@ -54,12 +41,6 @@ def create_state_choice_space(model: InternalModel, *, is_last_period: bool):
         sparse_vars={},
         dense_vars=_value_grid,
     )
-    # ==================================================================================
-    # create indexers and segments
-    # ==================================================================================
-    choice_segments = None
-
-    state_indexers = {}  # type: ignore[var-annotated]
 
     # ==================================================================================
     # create state space info
@@ -85,7 +66,7 @@ def create_state_choice_space(model: InternalModel, *, is_last_period: bool):
         indexer_infos=indexer_infos,
     )
 
-    return state_choice_space, space_info, state_indexers, choice_segments
+    return state_choice_space, space_info
 
 
 def _create_value_grid(grids, subset):

tests/test_entry_point.py

@@ -182,7 +182,7 @@ def test_create_compute_conditional_continuation_value():
         },
     }
 
-    _, space_info, _, _ = create_state_choice_space(
+    _, space_info = create_state_choice_space(
         model=model,
         is_last_period=False,
     )
@@ -228,7 +228,7 @@ def test_create_compute_conditional_continuation_value_with_discrete_model():
         },
     }
 
-    _, space_info, _, _ = create_state_choice_space(
+    _, space_info = create_state_choice_space(
         model=model,
         is_last_period=False,
     )
@@ -279,7 +279,7 @@ def test_create_compute_conditional_continuation_policy():
         },
     }
 
-    _, space_info, _, _ = create_state_choice_space(
+    _, space_info = create_state_choice_space(
         model=model,
         is_last_period=False,
     )
@@ -326,7 +326,7 @@ def test_create_compute_conditional_continuation_policy_with_discrete_model():
         },
     }
 
-    _, space_info, _, _ = create_state_choice_space(
+    _, space_info = create_state_choice_space(
         model=model,
         is_last_period=False,
     )

tests/test_model_functions.py

@@ -56,7 +56,7 @@ def test_get_utility_and_feasibility_function():
         },
     }
 
-    _, space_info, _, _ = create_state_choice_space(
+    _, space_info = create_state_choice_space(
         model=model,
         is_last_period=False,
     )

tests/test_simulate.py

@@ -40,7 +40,7 @@ def simulate_inputs():
     model_config = get_model_config("iskhakov_et_al_2017_stripped_down", n_periods=1)
     model = process_model(model_config)
 
-    _, space_info, _, _ = create_state_choice_space(
+    _, space_info = create_state_choice_space(
         model=model,
         is_last_period=False,
     )