Create triage.py

core/triage.py

@@ -0,0 +1,179 @@
+import torch
+import numpy as np
+from dataclasses import dataclass
+from typing import List, Dict, Tuple
+
+# ==========================================
+# 1. CONSTANTS (Aligned with V7/V5 Specs)
+# ==========================================
+PHONEMES = [
+    'p', 'pʰ', 'b', 'bʰ', 't', 'tʰ', 'd', 'dʰ', 't̪', 't̪ʰ', 'd̪', 'd̪ʱ',
+    'ʈ', 'ɖ', 'k', 'kʰ', 'g', 'ɡʰ', 'q', 'ʔ', 'f', 'v', 's', 'z',
+    'ʃ', 'ʒ', 'ʂ', 'ʐ', 'x', 'ɣ', 'h', 'θ', 'ð', 'tʃ', 'dʒ', 'm',
+    'mʰ', 'n', 'nʰ', 'ŋ', 'ɳ', 'l', 'lʰ', 'r', 'rʰ', 'ɽ', 'ɽʰ', 'w',
+    'j', 'a', 'aː', 'i', 'iː', 'u', 'uː', 'e', 'eː', 'o', 'oː',
+    'ə', 'æ', 'ɔ', 'ɪ', 'ʊ'
+]
+P2I = {p: i for i, p in enumerate(PHONEMES)}
+
+ERRORS = ["Sub", "Front", "Stop", "Glide", "Clust", "Epen", "Block", "Prolong", "Repeat"]
+
+# Mapping Error Index (0-8) to Major Type Category
+# 0-4 = Articulation, 5 = Motor, 6-8 = Fluency
+MAJOR_MAP = {
+    0: "Artic", 1: "Artic", 2: "Artic", 3: "Artic", 4: "Artic",
+    5: "Motor", 6: "Fluency", 7: "Fluency", 8: "Fluency"
+}
+
+# Clinical Priority: Treat Motor issues first, then Fluency, then Artic
+PRIORITY = {"Motor": 3, "Fluency": 2, "Artic": 1}
+
+@dataclass
+class ClinicalDiagnosis:
+    id: int
+    error_idx: int
+    error_name: str
+    major_type: str
+    severity: float
+    active_indices: List[int]
+    active_phonemes: List[str]
+    position_idx: int
+
+class TriageEngineV2:
+    def __init__(self, device='cpu'):
+        self.device = device
+
+    def scan_tensor(self, tensor_3d: np.ndarray) -> List[ClinicalDiagnosis]:
+        """
+        Scans the 64x9x3 AMLM output map for clinical syndromes.
+        Returns a sorted list of diagnoses.
+        """
+        diagnoses = []
+        uid_counter = 0
+
+        # tensor_3d shape: (3, 64, 9) -> (Channels, Phonemes, Errors)
+        
+        for err_idx in range(9):
+            # Check maximum probability for this error type across all phonemes/positions
+            slice_max = np.max(tensor_3d[:, :, err_idx])
+
+            # STRICT FILTERING: Threshold > 0.6 means the AMLM is confident
+            if slice_max > 0.6:
+                # Find which phonemes are affected
+                # We collapse the channel dim (positions) to see if the phoneme is affected anywhere
+                hits = np.where(tensor_3d[:, :, err_idx].max(axis=0) > 0.6)[0]
+                
+                # Filter specific weird edge cases (e.g. Gliding on vowels) if necessary
+                # Logic from Integration.ipynb: valid_hits = [int(h) for h in hits if not (err_idx == 3 and h >= 48)]
+                valid_hits = [int(h) for h in hits]
+
+                # NOISE REDUCTION: Cap at Top 10 worst errors per category
+                if len(valid_hits) > 10:
+                    severity_scores = []
+                    for h in valid_hits:
+                        score = tensor_3d[:, h, err_idx].max()
+                        severity_scores.append((score, h))
+                    severity_scores.sort(key=lambda x: x[0], reverse=True)
+                    valid_hits = [x[1] for x in severity_scores[:10]]
+
+                valid_hits.sort()
+
+                if not valid_hits: continue
+
+                # Determine dominant position (Initial, Medial, or Final)
+                pos_sums = tensor_3d[:, valid_hits, err_idx].sum(axis=1)
+                dom_pos = np.argmax(pos_sums)
+
+                diag = ClinicalDiagnosis(
+                    id=uid_counter,
+                    error_idx=err_idx,
+                    error_name=ERRORS[err_idx],
+                    major_type=MAJOR_MAP[err_idx],
+                    severity=float(slice_max),
+                    active_indices=valid_hits,
+                    active_phonemes=[PHONEMES[i] for i in valid_hits],
+                    position_idx=int(dom_pos)
+                )
+                diagnoses.append(diag)
+                uid_counter += 1
+
+        # Sort by Priority (Motor > Fluency > Artic) then Severity
+        return sorted(diagnoses, key=lambda x: (PRIORITY[x.major_type], x.severity), reverse=True)
+
+    def prepare_aeg_input(self, focus_diag: ClinicalDiagnosis, all_diagnoses: List[ClinicalDiagnosis], user_stats: List[float]) -> Dict[str, torch.Tensor]:
+        """
+        Generates the 162-Dim Input Vector for the AEG Neural Network.
+        Adheres strictly to 'Data Generator.py' logic.
+        """
+        # 1. Cluster Mask (Targets)
+        # Which sounds should we target in therapy? (The ones in the focus diagnosis)
+        cluster_mask = [0.0] * 64
+        target_set = set(focus_diag.active_indices)
+        for idx in focus_diag.active_indices:
+            cluster_mask[idx] = 1.0
+
+        # 2. Trap Mask (Forbidden Sounds)
+        # Sounds we must NOT use to avoid negative reinforcement.
+        trap_mask = [0.0] * 64
+
+        def attempt_trap(idx):
+            # Logic: A sound cannot be a Trap if it is currently a Target.
+            if idx not in target_set:
+                trap_mask[idx] = 1.0
+
+        # A. Standard Traps (Rule-based from Data Generator)
+        if focus_diag.error_name == "Front": # If Fronting K->T
+            # Trap T/D so we don't ask user to say them while trying to learn K
+            if any(p in ['k', 'g', 'kʰ', 'ɡʰ'] for p in focus_diag.active_phonemes):
+                for x in ['t', 'd', 'tʰ', 'dʰ']:
+                    if x in P2I: attempt_trap(P2I[x])
+
+        # B. Comorbidity Traps (Dynamic)
+        # If user also has a Lisp (Sub S), Trap 'S' even if we are treating Fronting.
+        for other in all_diagnoses:
+            if other.id != focus_diag.id:
+                for idx in other.active_indices:
+                    attempt_trap(idx)
+
+        # 3. Phonetic Features (8-Dim Vector)
+        # Describes the target sound (e.g., +Stop, +Velar)
+        feat = [0.0] * 8
+        t_str = "".join(focus_diag.active_phonemes)
+
+        if any(x in t_str for x in "pbtdkgʈɖʔ"): feat[0] = 1.0 # Stop
+        if any(x in t_str for x in "fvszʃʒxhθð"): feat[1] = 1.0 # Fricative
+        if any(x in t_str for x in "lrw"): feat[2] = 1.0 # Liquid
+        if any(x in t_str for x in "mnŋ"): feat[3] = 1.0 # Nasal
+        if any(x in t_str for x in "bdgvzʒðmnlrw"): feat[4] = 1.0 # Voiced
+        if any(x in t_str for x in "pbfvmθð"): feat[5] = 1.0 # Front
+        if any(x in t_str for x in "kgŋxɣhʔ"): feat[6] = 1.0 # Back
+        if len(focus_diag.active_indices) > 1 or focus_diag.error_name == "Clust": feat[7] = 1.0 # Cluster
+
+        # 4. Flags & Stats
+        trap_count = sum(trap_mask)
+        # user_stats expects: [Severity, Fatigue, Age, Delta]
+        sev, fat, age, delta = user_stats
+
+        # Binary Flags for the Network
+        is_systemic = 1.0 if len(focus_diag.active_indices) > 1 else 0.0
+        is_trap = 1.0 if trap_count > 0 else 0.0
+        is_comorb = 1.0 if len(all_diagnoses) > 1 else 0.0
+        is_child = 1.0 if age < 0.3 else 0.0 # Normalized age (0.3 ~ 7 years old)
+        is_maint = 1.0 if sev <= 0.2 else 0.0
+        is_fatigued = 1.0 if fat > 0.8 else 0.0
+
+        flags = [is_systemic, is_trap, is_comorb, is_child, is_maint, is_fatigued]
+
+        maj_int = {"Artic": 0, "Fluency": 1, "Motor": 2}[focus_diag.major_type]
+
+        # Convert to PyTorch Tensors on correct Device
+        return {
+            "error_idx": torch.tensor([focus_diag.error_idx]).to(self.device),
+            "major_idx": torch.tensor([maj_int]).to(self.device),
+            "pos_idx": torch.tensor([focus_diag.position_idx]).to(self.device),
+            "cluster": torch.tensor([cluster_mask], dtype=torch.float32).to(self.device),
+            "trap": torch.tensor([trap_mask], dtype=torch.float32).to(self.device),
+            "feat": torch.tensor([feat], dtype=torch.float32).to(self.device),
+            "flags": torch.tensor([flags], dtype=torch.float32).to(self.device),
+            "stats": torch.tensor([user_stats], dtype=torch.float32).to(self.device)
+        }