diff --git a/pyproject.toml b/pyproject.toml index 0c065e0..35f2220 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "respondpy" -version = "0.3.0" +version = "0.3.1" description = "The Syndemic Lab's RESPOND Simulation Python Extension Module." readme = "README.md" requires-python = ">=3.11" diff --git a/src/respondpy/data/database_helpers.py b/src/respondpy/data/database_helpers.py index 3e08014..af35cc4 100644 --- a/src/respondpy/data/database_helpers.py +++ b/src/respondpy/data/database_helpers.py @@ -4,7 +4,7 @@ # Created Date: 2026-06-05 # # Author: Matthew Carroll # # ----- # -# Last Modified: 2026-06-17 # +# Last Modified: 2026-08-03 # # Modified By: Matthew Carroll # # ----- # # Copyright (c) 2026 Syndemics Lab at Boston Medical Center # @@ -13,6 +13,30 @@ import polars as pl +def _normalize_state_pairs( + values: list[tuple[int, str]] | list[list] +) -> list[tuple[int, str]]: + """Normalize state id/name mappings to ``[(id, name), ...]`` format.""" + if not values: + return [] + + first = values[0] + if isinstance(first, tuple) and len(first) == 2: + return values # type: ignore[return-value] + + # Legacy shape: [[id1, id2, ...], [name1, name2, ...]] + if ( + len(values) == 2 + and isinstance(values[0], list) + and isinstance(values[1], list) + and len(values[0]) == len(values[1]) + ): + return list(zip(values[0], values[1], strict=True)) + + raise ValueError( + "Invalid state mapping format. Expected list of (id, name) pairs.") + + def _sort_state_vector( lf: pl.LazyFrame, behaviors: list[tuple[int, str]], @@ -43,19 +67,25 @@ def _sort_state_vector( if 'intervention' not in s or 'behavior' not in s: raise ValueError( f"Invalid columns provided when attempting to sort state vector: {s}") + behavior_pairs = _normalize_state_pairs(behaviors) + intervention_pairs = _normalize_state_pairs(interventions) + behav = pl.LazyFrame( - behaviors, schema=["b_id", "b_name"], orient='row') + behavior_pairs, schema=["b_id", "b_name"], orient='row') inter = pl.LazyFrame( - interventions, schema=["i_id", "i_name"], orient='row') + intervention_pairs, schema=["i_id", "i_name"], orient='row') - # Sort first by behaviors, then by interventions - return lf.join( + # Sort by intervention id, then behavior id to match Input.get_state_names. + sorted_lf = lf.join( behav, left_on="behavior", right_on="b_name", how="inner" ).join( inter, left_on="intervention", right_on="i_name", how="inner" ).sort(["i_id", "b_id"]).drop(["i_id", "b_id"]) + # Preserve the original column order to avoid accidental schema drift. + return sorted_lf.select(s) + def _sort_transition_matrix( lf: pl.LazyFrame, @@ -83,27 +113,61 @@ def _sort_transition_matrix( ValueError If required transition columns are missing. """ - if 'intervention' not in lf.columns or 'behavior' not in lf.columns or 'next_intervention' not in lf.columns or 'next_behavior' not in lf.columns: + s = lf.collect_schema().names() + cols = set(s) + canonical_required = { + "initial_intervention", + "initial_behavior", + "new_intervention", + "new_behavior", + } + + if canonical_required.issubset(cols): + working = lf + else: raise ValueError( - f"Invalid columns provided when attempting to sort transition matrix: {lf.columns}") - behav = pl.LazyFrame(behaviors, schema=["b_id", "b_name"]) - inter = pl.LazyFrame(interventions, schema=["i_id", "i_name"]) - - # Sort order: - # 1. Next Behavior - # 2. Next Behavior - # 3. Initial Behavior - # 4. Initial Intervention - # e.g. [active_injection, no_treatment, active_injection, no_treatment], [active_injection, no_treatment, active_injection, buprenorphine], [active_injection, no_treatment, active_injection, methadone], etc. - return lf.drop("i_id").join( - behav, left_on="next_behavior", right_on="b_name", how="inner" - ).sort(pl.col("b_id")).drop("b_id").join( - inter, left_on="next_intervention", right_on="i_name", how="inner" - ).sort(pl.col("i_id")).drop("i_id").join( - behav, left_on="behavior", right_on="b_name", how="inner" - ).sort(pl.col("b_id")).drop("b_id").join( - inter, left_on="intervention", right_on="i_name", how="inner" - ).sort(pl.col("i_id")) + f"Invalid columns provided when attempting to sort transition matrix: {s}") + + behavior_pairs = _normalize_state_pairs(behaviors) + intervention_pairs = _normalize_state_pairs(interventions) + + behav = pl.LazyFrame( + behavior_pairs, schema=["b_id", "b_name"], orient='row' + ) + inter = pl.LazyFrame( + intervention_pairs, schema=["i_id", "i_name"], orient='row' + ) + + # Destination-major ordering guarantees column-stochastic matrices after + # reshape for y = Mx when source-state labels define columns. + sorted_lf = working.join( + inter.rename({"i_id": "new_i_id", "i_name": "new_i_name"}), + left_on="new_intervention", + right_on="new_i_name", + how="inner", + ).join( + behav.rename({"b_id": "new_b_id", "b_name": "new_b_name"}), + left_on="new_behavior", + right_on="new_b_name", + how="inner", + ).join( + inter.rename({"i_id": "initial_i_id", "i_name": "initial_i_name"}), + left_on="initial_intervention", + right_on="initial_i_name", + how="inner", + ).join( + behav.rename({"b_id": "initial_b_id", "b_name": "initial_b_name"}), + left_on="initial_behavior", + right_on="initial_b_name", + how="inner", + ).sort(["new_i_id", "new_b_id", "initial_i_id", "initial_b_id"]).drop([ + "new_i_id", + "new_b_id", + "initial_i_id", + "initial_b_id", + ]) + + return sorted_lf.select(s) def sort_dataframes( @@ -130,10 +194,19 @@ def sort_dataframes( polars.LazyFrame Sorted LazyFrame when shape is recognized, otherwise input. """ - if len(lf.collect_schema().names()) == 3: - return _sort_state_vector(lf, behaviors, interventions) - if len(lf.collect_schema().names()) == 4: + cols = set(lf.collect_schema().names()) + + if { + "initial_intervention", + "initial_behavior", + "new_intervention", + "new_behavior", + }.issubset(cols): return _sort_transition_matrix(lf, behaviors, interventions) + + if {"intervention", "behavior"}.issubset(cols): + return _sort_state_vector(lf, behaviors, interventions) + return lf diff --git a/src/respondpy/data/input.py b/src/respondpy/data/input.py index d325968..acaa36e 100644 --- a/src/respondpy/data/input.py +++ b/src/respondpy/data/input.py @@ -360,8 +360,19 @@ def _get_parameter_filled( ParameterType.INTERVENTION_TRANSITION_PROBABILITY, ParameterType.BEHAVIOR_TRANSITION_PROBABILITY] ) - if complete_state_vector or complete_transition: - return lf + if complete_state_vector: + return sort_dataframes( + lf, + self._get_single_state_table("behavior"), + self._get_single_state_table("intervention") + ) + + if complete_transition: + return sort_dataframes( + lf, + self._get_single_state_table("behavior"), + self._get_single_state_table("intervention") + ) if param.is_state_vector_operation(): value_col = param.get_value_column_name() @@ -401,8 +412,8 @@ def _get_parameter_filled( return sort_dataframes( res.lazy(), - self._get_single_state_table("intervention"), - self._get_single_state_table("behavior") + self._get_single_state_table("behavior"), + self._get_single_state_table("intervention") ) # External Functions diff --git a/tests/test_data_database_helpers.py b/tests/test_data_database_helpers.py index cb671d6..8983bab 100644 --- a/tests/test_data_database_helpers.py +++ b/tests/test_data_database_helpers.py @@ -4,7 +4,7 @@ # Created Date: 2026-06-10 # # Author: Matthew Carroll # # ----- # -# Last Modified: 2026-06-17 # +# Last Modified: 2026-08-03 # # Modified By: Matthew Carroll # # ----- # # Copyright (c) 2026 Syndemics Lab at Boston Medical Center # @@ -99,33 +99,6 @@ def test_sort_state_vector_happy_path() -> None: assert out.select("count").to_series().to_list() == [10, 20, 30, 40] -@pytest.mark.unit -@pytest.mark.filterwarnings( - "ignore:Determining the column names of a LazyFrame requires resolving its schema.*:polars.exceptions.PerformanceWarning" -) -def test_sort_transition_matrix_happy_path() -> None: - lf = pl.LazyFrame( - { - "intervention": ["i2", "i1", "i2", "i1"], - "behavior": ["b2", "b1", "b1", "b2"], - "next_intervention": ["i1", "i2", "i2", "i1"], - "next_behavior": ["b1", "b2", "b1", "b2"], - "probability": [0.1, 0.2, 0.3, 0.4], - # Current implementation drops this column before re-adding it. - "i_id": [0, 0, 0, 0], - } - ) - - out = _sort_transition_matrix( - lf, - behaviors=[[1, 2], ["b1", "b2"]], - interventions=[[1, 2], ["i1", "i2"]], - ).collect() - - assert out.height == 4 - assert "probability" in out.columns - - @pytest.mark.unit def test_sort_dataframes_calls_state_vector_sort(monkeypatch: pytest.MonkeyPatch) -> None: lf = pl.LazyFrame( @@ -150,30 +123,3 @@ def fake_sort_state_vector( out = sort_dataframes(lf, [(1, "b1")], [(1, "i1")]).collect() assert out.columns == ["sentinel"] - - -@pytest.mark.unit -def test_sort_dataframes_calls_transition_matrix_sort(monkeypatch: pytest.MonkeyPatch) -> None: - lf = pl.LazyFrame( - { - "intervention": ["i1"], - "behavior": ["b1"], - "next_intervention": ["i1"], - "next_behavior": ["b1"], - } - ) - - def fake_sort_transition_matrix( - _lf: pl.LazyFrame, - _behaviors: list[tuple[int, str]], - _interventions: list[tuple[int, str]], - ) -> pl.LazyFrame: - return pl.LazyFrame({"sentinel": [2]}) - - monkeypatch.setattr( - "respondpy.data.database_helpers._sort_transition_matrix", - fake_sort_transition_matrix, - ) - - out = sort_dataframes(lf, [(1, "b1")], [(1, "i1")]).collect() - assert out.columns == ["sentinel"] diff --git a/tests/test_integration.py b/tests/test_integration.py index 76a6397..fa9ed21 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -427,7 +427,9 @@ def test_single_timestep_numerical_state_matches_expected( sim.run() final_state = sim.get_model(0).get_state() - expected_state = np.array([216.2933685, 0.0, 0.0, 210.24611685000002]) + expected_state = np.array( + [105.36565049999999, 71.49283020000001, 61.394525775000005, 38.11650975] + ) np.testing.assert_allclose( final_state, expected_state, rtol=1e-10, atol=1e-10) @@ -448,8 +450,104 @@ def test_fifty_two_timestep_numerical_state_matches_expected( final_state = sim.get_model(0).get_state() expected_state = np.array( - [5.764271919791796e26, 0.0, 0.0, 5.6072656482224475e26] + [63.99771165143241, 28.195217933852412, + 39.53089805842006, 19.00629313025105] ) np.testing.assert_allclose( final_state, expected_state, rtol=1e-9, atol=1e-9) + + +@pytest.mark.integration +def test_behavior_and_intervention_matrices_are_column_stochastic(setup_data): + """Behavior and intervention matrices should be column-stochastic for y=Mx.""" + db_path, config_path = setup_data + inp = rpy.data.Input(db_path=db_path, conf_path=config_path) + sim = rpy.build_simulation(inp) + + timestep = sim[0].get_timestep_at_index(0) + transition_names = timestep.get_transition_names() + + behavior_matrix = timestep[ + transition_names.index("behavior") + ].get_matrices()[0] + intervention_matrix = timestep[ + transition_names.index("intervention") + ].get_matrices()[0] + + np.testing.assert_allclose( + behavior_matrix.sum(axis=0), + np.ones(behavior_matrix.shape[1]), + atol=1e-12, + rtol=0.0, + ) + np.testing.assert_allclose( + intervention_matrix.sum(axis=0), + np.ones(intervention_matrix.shape[1]), + atol=1e-12, + rtol=0.0, + ) + + +@pytest.mark.integration +def test_behavior_then_intervention_preserves_mass_one_step(setup_data): + """Applying behavior then intervention should preserve total mass.""" + db_path, config_path = setup_data + inp = rpy.data.Input(db_path=db_path, conf_path=config_path) + sim = rpy.build_simulation(inp) + + model = sim[0] + x0 = model.get_state() + + timestep = model.get_timestep_at_index(0) + transition_names = timestep.get_transition_names() + behavior_matrix = timestep[ + transition_names.index("behavior") + ].get_matrices()[0] + intervention_matrix = timestep[ + transition_names.index("intervention") + ].get_matrices()[0] + + x1 = behavior_matrix @ x0 + x2 = intervention_matrix @ x1 + + np.testing.assert_allclose( + float(np.sum(x1)), float(np.sum(x0)), atol=1e-12, rtol=0.0 + ) + np.testing.assert_allclose( + float(np.sum(x2)), float(np.sum(x1)), atol=1e-12, rtol=0.0 + ) + + +@pytest.mark.integration +def test_transition_matrix_indices_match_state_name_order(setup_data): + """Matrix indices should map as M[destination_state, source_state].""" + db_path, config_path = setup_data + inp = rpy.data.Input(db_path=db_path, conf_path=config_path) + + state_names = inp.get_state_names() + state_to_idx = {state: idx for idx, state in enumerate(state_names)} + + for param_type in ( + rpy.data.ParameterType.BEHAVIOR_TRANSITION_PROBABILITY, + rpy.data.ParameterType.INTERVENTION_TRANSITION_PROBABILITY, + ): + param = rpy.data.Parameter(param_type) + matrix = inp.select_parameter(param, cohort_id=1, time=1) + + sample_id = inp._get_sample_id_for_parameter(param, cohort_id=1) + long_df = inp._get_parameter_filled( + param, sample_id=sample_id, time=1).collect() + + for row in long_df.iter_rows(named=True): + src = (row["initial_intervention"], row["initial_behavior"]) + dst = (row["new_intervention"], row["new_behavior"]) + src_idx = state_to_idx[src] + dst_idx = state_to_idx[dst] + + np.testing.assert_allclose( + matrix[dst_idx, src_idx], + row["probability"], + atol=1e-12, + rtol=0.0, + )