Skip to content

[Hotfix] Column Stochasticity #34

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MJC598 merged 3 commits into from
Aug 3, 2026
Merged

[Hotfix] Column Stochasticity #34

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
eb79b8b
addressing column stochastic issues
MJC598 Aug 3, 2026
3f70fbe
addressing broken pytests
MJC598 Aug 3, 2026
7dfb4bf
version bump
MJC598 Aug 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
129 changes: 101 additions & 28 deletions src/respondpy/data/database_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 #
Expand All @@ -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]],
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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


Expand Down
19 changes: 15 additions & 4 deletions src/respondpy/data/input.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down
56 changes: 1 addition & 55 deletions tests/test_data_database_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 #
Expand Down Expand Up @@ -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(
Expand All @@ -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"]
Loading
Loading