from __future__ import annotations import math import numpy as np import pandas as pd import pytest from sctm_next_round.tensor_builder_v2 import ( AUDIENCE_NA_ID, AUDIENCE_OTHER_ID, CADENCE_TEETH, NUMERIC_DROP_FIELDS_V2, STATIC_NUMERIC_DATE_COLS, STATIC_NUMERIC_DURATION_COLS, build_static_numeric, compute_observation_features, compute_patient_cadence, encode_audience, encode_drug_flags, encode_v1_event_values, enumerate_windows, fit_cat_vocab_v2, fit_dose_winsor_caps, split_static_columns, ) # --------------------------------------------------------------------------- # window enumeration parity with the v1 inline algorithm # --------------------------------------------------------------------------- def v1_reference_windows(n_visits: int, max_seq_len: int, stride: int) -> list[tuple[int, int]]: """Literal transcription of tensor_builder.build_split_arrays enumeration.""" starts = list(range(0, n_visits, max(1, stride))) if starts and starts[-1] + max_seq_len < n_visits: starts.append(max(0, n_visits - max_seq_len)) out = [] for start in starts: length = int(min(max_seq_len, max(0, n_visits - start))) if length < 1: continue if n_visits > 1 and length < 2: continue out.append((start, length)) return out class TestWindowEnumeration: @pytest.mark.parametrize("n", [1, 2, 63, 64, 65, 127, 128, 129, 200, 475]) def test_matches_v1_reference(self, n): assert enumerate_windows(n, 128, 64) == v1_reference_windows(n, 128, 64) def test_singleton_patient_gets_one_window(self): assert enumerate_windows(1, 128, 64) == [(0, 1)] def test_tail_window_appended(self): # 200 visits, stride 64: starts 0,64,128,192 -> the 192 window has length 8 (>=2, kept) wins = enumerate_windows(200, 128, 64) assert (192, 8) in wins # --------------------------------------------------------------------------- # cadence block # --------------------------------------------------------------------------- class TestCadence: def test_trailing_median_and_deviation(self): # gaps: -, 14, 0(same-day), 28, 90 time_days = np.array([0.0, 14.0, 14.0, 42.0, 132.0], dtype=np.float32) cad = compute_patient_cadence(time_days) # first visit undefined assert cad["cadence_trailing_median"][0] == 0.0 assert cad["cadence_deviation_log"][0] == 0.0 # t=1: first positive gap 14 -> trailing median 14, no prior rhythm -> deviation 0 assert cad["cadence_trailing_median"][1] == 14.0 assert cad["cadence_deviation_log"][1] == 0.0 # t=2: same-day 0-gap excluded from the rhythm (red line 3) assert cad["cadence_trailing_median"][2] == 14.0 assert cad["cadence_deviation_log"][2] == 0.0 assert cad["cadence_tooth_distance"][2] == 0.0 # t=3: gap 28 vs prior rhythm 14 assert cad["cadence_trailing_median"][3] == pytest.approx(21.0) assert cad["cadence_deviation_log"][3] == pytest.approx(math.log(29.0 / 15.0), abs=1e-6) # t=4: gap 90 vs prior rhythm median([14, 28]) = 21 assert cad["cadence_deviation_log"][4] == pytest.approx(math.log(91.0 / 22.0), abs=1e-6) assert cad["cadence_trailing_median"][4] == pytest.approx(28.0) def test_trailing_window_is_five(self): gaps = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0] time_days = np.cumsum([0.0] + gaps).astype(np.float32) cad = compute_patient_cadence(time_days) # at the last visit the trailing window is the last 5 gaps [30,40,50,60,70] assert cad["cadence_trailing_median"][-1] == pytest.approx(50.0) def test_tooth_distance(self): time_days = np.array([0.0, 14.0, 44.0, 134.0, 154.0], dtype=np.float32) cad = compute_patient_cadence(time_days) # gap 14 -> exactly on the 14 tooth assert cad["cadence_tooth_distance"][1] == 0.0 # gap 30 -> nearest tooth 28: |30-28|/28 assert cad["cadence_tooth_distance"][2] == pytest.approx(2.0 / 28.0) # gap 90 -> on the 90 tooth assert cad["cadence_tooth_distance"][3] == 0.0 # gap 20 -> min(|20-14|/14, |20-28|/28) = 6/14 vs 8/28 -> 8/28 assert cad["cadence_tooth_distance"][4] == pytest.approx(min(6.0 / 14.0, 8.0 / 28.0, 70.0 / 90.0)) assert set(CADENCE_TEETH) == {14.0, 28.0, 90.0} # --------------------------------------------------------------------------- # observation-feature precompute == formal_model torch loop # --------------------------------------------------------------------------- class TestObservationFeatures: def test_matches_formal_model_torch_loop(self): torch = pytest.importorskip("torch") from sctm_next_round.formal_model import SCTMFormalModel rng = np.random.default_rng(20260611) bsz, seq_len, n_fields = 7, 19, 11 observed = rng.random((bsz, seq_len, n_fields)) < 0.35 gaps = rng.integers(0, 60, size=(bsz, seq_len)).astype(np.float64) gaps[:, 0] = 0.0 times = np.cumsum(gaps, axis=1).astype(np.float32) batch = {"time_since_start_days": torch.from_numpy(times)} # the method touches no instance state — call it unbound on the literal implementation torch_since, torch_carry = SCTMFormalModel._time_since_last_observed( None, batch, torch.from_numpy(observed) ) np_since, np_carry = compute_observation_features(observed, times) np.testing.assert_allclose(np_since, torch_since.numpy(), rtol=0, atol=1e-4) np.testing.assert_array_equal(np_carry, torch_carry.numpy()) def test_window_reset_semantics(self): # state resets at the window start: nothing seen before position 0 observed = np.zeros((1, 4, 1), dtype=bool) observed[0, 1, 0] = True times = np.array([[0.0, 10.0, 25.0, 60.0]], dtype=np.float32) since, carry = compute_observation_features(observed, times) np.testing.assert_allclose(since[0, :, 0], [0.0, 0.0, 15.0, 50.0]) np.testing.assert_array_equal(carry[0, :, 0], [False, True, True, True]) def test_visit_completeness_formula(self): # completeness = mean over (cat observed, numeric mask, ordinal mask) missing_ids = np.array([[[0, 3], [3, 3]]]) # cat observed: [T0: 1/2, T1: 0/2] numeric_mask = np.array([[[True], [False]]]) ordinal_mask = np.array([[[True], [True]]]) observed_all = np.concatenate( [missing_ids == 0, numeric_mask, ordinal_mask], axis=2 ) completeness = observed_all.mean(axis=2) np.testing.assert_allclose(completeness[0], [0.75, 0.25]) # --------------------------------------------------------------------------- # encoders # --------------------------------------------------------------------------- class TestEncoders: def test_audience_vocab(self): series = pd.Series(["本人", "家属", "居委干部", "社区民警", "民警", "邻居", "其他", "亲戚", None, ""]) out = encode_audience(series) assert out.tolist() == [0, 1, 2, 3, 3, 4, 5, AUDIENCE_OTHER_ID, AUDIENCE_NA_ID, AUDIENCE_NA_ID] assert out.dtype == np.int8 def test_drug_flags(self): df = pd.DataFrame( { "随访药品1名称": [ "棕榈酸帕利哌酮注射液", # LAI by name regex "氯氮平片|阿立哌唑", # clozapine "盐酸苯海索片", # anticholinergic "利培酮", # none None, ], "药品1剂量长效药": [None, None, None, "25", None], } ) flags = encode_drug_flags(df) assert flags.shape == (5, 3) assert flags[0].tolist() == [True, False, False] assert flags[1].tolist() == [False, True, False] assert flags[2].tolist() == [False, False, True] # 剂量长效药 nonnull -> LAI even without a regex hit (codex red line 10) assert flags[3].tolist() == [True, False, False] assert flags[4].tolist() == [False, False, False] def test_v1_event_values_vectorized(self): df = pd.DataFrame( { "是否转诊": ["是", "否", ""], "是否应急处置": ["否", "是", ""], "治疗方式": ["住院", "门诊", ""], "随访类型": ["", "住院随访", ""], "危险行为等级": ["0级", "1级 ", ""], "滋事次数": [0, 2, None], "肇事次数": [0, 0, 0], "肇祸次数": [0, 0, 0], "自杀未遂次数": [0, 0, 0], "自伤次数": [0, 0, 0], "其他危害他人行为": [0, 0, 0], } ) ev = encode_v1_event_values(df) np.testing.assert_allclose(ev[0], [1, 0, 1, 0, 0]) # trailing-space danger token still counts as high danger (matches v1 value_text strip) np.testing.assert_allclose(ev[1], [0, 1, 1, 1, 1]) np.testing.assert_allclose(ev[2], [0, 0, 0, 0, 0]) def test_cat_vocab_whitespace_normalized(self): df = pd.DataFrame( { "split": ["train"] * 4, "危险行为等级": ["1级", "1级 ", "1级　", "2级"], } ) vocab = fit_cat_vocab_v2(df, ["危险行为等级"], min_freq=1) keys = [k for k in vocab if k.startswith("危险行为等级=")] assert "危险行为等级=1级" in keys # whitespace variants collapse into ONE token assert len([k for k in keys if "1级" in k]) == 1 def test_dose_winsor_caps_train_only(self): train_vals = list(np.linspace(1, 100, 999)) + [100150.0] # keying error in train df = pd.DataFrame( { "split": ["train"] * 1000 + ["val"], "__dose": train_vals + [50.0], } ) caps = fit_dose_winsor_caps(df, {"药品1非长效药物晨间服用剂量": "__dose"}) cap = caps["药品1非长效药物晨间服用剂量"] assert cap < 100150.0 assert cap >= 99.0 def test_numeric_drop_list_matches_codex_c1(self): assert set(NUMERIC_DROP_FIELDS_V2) == { "公安滋事总数", "公安肇事总数", "以往住院情况", "本次住院时长天数", "易肇事肇祸次数", "其他危害他人行为", "药品1用法长效药", } # --------------------------------------------------------------------------- # static routing # --------------------------------------------------------------------------- class TestStaticRouting: def test_split_static_columns(self): cols = ( ["has_static", "A_static_性别", "A_static_性别_missing"] + ["A_static_曾肇事肇祸", "A_static_曾肇事肇祸_missing", "A_static_家族史", "A_static_关解锁情况"] + STATIC_NUMERIC_DATE_COLS + [c + "_missing" for c in STATIC_NUMERIC_DATE_COLS] + STATIC_NUMERIC_DURATION_COLS + ["A_static_recent_violence_date", "A_static_years_since_recent_violence_at_first_visit"] ) kept, numeric_routed, dropped = split_static_columns(cols) assert kept == ["has_static", "A_static_性别", "A_static_性别_missing"] assert set(numeric_routed) == set(STATIC_NUMERIC_DATE_COLS + STATIC_NUMERIC_DURATION_COLS) assert "A_static_曾肇事肇祸" in dropped assert "A_static_recent_violence_date" in dropped def test_static_numeric_sentinel_guard_and_zscore(self): static = pd.DataFrame( { "patient_id": ["a", "b", "c", "d"], "A_static_birth_date": ["1960-05-01", "1900-01-01", "1975-01-01", ""], "A_static_age_at_first_visit_years": [50.0, 40.0, 999.0, np.nan], } ) values, mask, cols, scaler = build_static_numeric(static, train_ids={"a", "b"}) bidx = cols.index("A_static_birth_date__frac_year") didx = cols.index("A_static_age_at_first_visit_years") # 1900 sentinel masked, empty masked assert mask[:, bidx].tolist() == [True, False, True, False] # out-of-range duration masked assert mask[:, didx].tolist() == [True, True, False, False] # masked entries carry 0 assert values[1, bidx] == 0.0 assert values[3, didx] == 0.0 # z-score fitted on train only (patients a,b): age mean 45 std 5 assert scaler["A_static_age_at_first_visit_years"]["mean"] == pytest.approx(45.0) assert values[0, didx] == pytest.approx((50.0 - 45.0) / 5.0) # --------------------------------------------------------------------------- # tiny end-to-end through build_split_arrays_v2 (synthetic) # --------------------------------------------------------------------------- class TestSplitAssembly: def test_label_arrays_window_aligned(self): from sctm_next_round.tensor_builder_v2 import build_split_arrays_v2, encode_rows n = 5 rows = pd.DataFrame( { "A表编号": ["p1"] * 3 + ["p2"] * 2, "visit_index": [0, 1, 2, 0, 1], "visit_date": pd.to_datetime( ["2020-01-01", "2020-01-15", "2020-02-12", "2021-01-01", "2021-04-01"] ), "__year__": [2020, 2020, 2020, 2021, 2021], "split": ["train"] * n, "terminal_or_disposition_label": ["none", "none", "none", "none", "death"], "危险行为等级": ["0级", "1级", "0级", "0级", "0级"], "本次访视对象": ["本人", "家属", "本人", "居委干部", "本人"], "随访药品1名称": ["氯氮平", None, "棕榈酸帕利哌酮", None, None], # label-v2 columns "terminal_label_v2": ["none", "none", "none", "none", "death"], "diseng_state": [0, 1, 0, 0, 0], "diseng_episode_start": [False, True, False, False, False], "time_to_return_days": [np.nan, 28.0, np.nan, np.nan, np.nan], "return_within_365": [-1, 1, -1, -1, -1], "post_discharge_gap90": [-1, -1, 0, -1, -1], "same_day_flag": [False, False, False, False, False], "n_same_day": [1, 1, 1, 1, 1], "event_hosp_newdate": [False, True, False, False, False], "event_discharge_newdate": [False, False, True, False, False], } ) cat_cols = ["危险行为等级", "本次访视对象"] ord_cols = ["危险行为等级"] num_cols: list[str] = [] cat_vocab = {"": 0, "": 1, "危险行为等级=0级": 2, "危险行为等级=1级": 3, "本次访视对象=本人": 4, "本次访视对象=家属": 5, "本次访视对象=居委干部": 6} drug_vocab = {"": 0, "": 1, "氯氮平": 2} rank_maps = {"危险行为等级": {"0级": 0, "1级": 1}} encoded = encode_rows( rows, cat_cols=cat_cols, num_cols=num_cols, ord_cols=ord_cols, cat_vocab=cat_vocab, drug_vocab=drug_vocab, numeric_scaler={}, parsed_cols={}, rank_maps=rank_maps, cbe_dim=4, max_drug_atoms=2, ) arrays = build_split_arrays_v2( rows, encoded, static_index_by_patient={"p1": 0, "p2": 1}, static_cat_ids=np.array([[7], [8]], dtype=np.int64), static_numeric=np.array([[0.5], [-0.5]], dtype=np.float32), static_numeric_mask=np.array([[True], [False]]), n_cat=len(cat_cols), n_num=0, n_ord=len(ord_cols), cbe_dim=4, max_drug_atoms=2, max_seq_len=8, stride=4, progress_every=0, ) assert arrays["valid_mask"].shape[0] == 2 # one window per patient w0 = list(arrays["patient_ids"]).index("p1") w1 = list(arrays["patient_ids"]).index("p2") # window-aligned label arrays assert arrays["diseng_state"][w0, :3].tolist() == [0, 1, 0] assert arrays["diseng_episode_start"][w0, 1] assert arrays["time_to_return_days"][w0, 1] == 28.0 assert np.isnan(arrays["time_to_return_days"][w0, 0]) assert np.isnan(arrays["time_to_return_days"][w0, 5]) # padding stays NaN assert arrays["return_within_365"][w0, :3].tolist() == [-1, 1, -1] assert arrays["post_discharge_gap90"][w0, :3].tolist() == [-1, -1, 0] assert arrays["terminal_label_v2"][w1, :2].tolist() == [0, 1] # none, death # event_labels_v2 channel order: referral, emergency, hosp_new, discharge_new, danger, harm assert arrays["event_labels_v2"][w0, 1, 2] == 1.0 assert arrays["event_labels_v2"][w0, 2, 3] == 1.0 # v2 inputs assert arrays["audience_id"][w0, :3].tolist() == [0, 1, 0] assert arrays["drug_flags"][w0, 0].tolist() == [False, True, False] # clozapine assert arrays["drug_flags"][w0, 2].tolist() == [True, False, False] # LAI by regex # cadence: p1 gaps 14, 28 assert arrays["cadence_trailing_median"][w0, :3].tolist() == pytest.approx([0.0, 14.0, 21.0]) assert arrays["cadence_deviation_log"][w0, 2] == pytest.approx(math.log(29.0 / 15.0), abs=1e-6) # static channels broadcast to windows assert arrays["static_numeric"][w0, 0] == 0.5 assert arrays["static_numeric_mask"][w1, 0] == False # noqa: E712 # obs features exist with the family-ordered F axis assert arrays["obs_time_since_last"].shape == (2, 8, len(cat_cols) + 0 + len(ord_cols)) assert arrays["visit_completeness"].shape == (2, 8) # v1-compatible arrays still present for key in ("cat_value_ids", "missing_ids", "ordinal_cbe", "service_state", "delta_t_next_log", "terminal_label"): assert key in arrays assert arrays["terminal_label"][w1, :2].tolist() == [0, 1]