Remove unused global cancer type mappings from aeon.py

- Remove unused imports of INT_TO_CANCER_TYPE_MAP and CANCER_TYPE_TO_INT_MAP
- Remove unused col_indices_to_drop global variable
- Simplify CANCER_TYPES_TO_DROP to only the 5 types that are excluded
- Update tests to import mappings from data.py instead of aeon.py
- All 8 tests pass successfully

The Aeon inference now exclusively uses metadata-loaded mappings (160 classes)
instead of the old global constants (183 classes), eliminating confusion.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

src/mosaic/inference/aeon.py

@@ -16,8 +16,6 @@ from torch.utils.data import DataLoader
 from mosaic.inference.data import (
     SiteType,
     TileFeatureTensorDataset,
-    INT_TO_CANCER_TYPE_MAP,
-    CANCER_TYPE_TO_INT_MAP,
     encode_sex,
     encode_tissue_site,
 )
@@ -25,29 +23,14 @@ from mosaic.inference.data import (
 from loguru import logger
 
 # Cancer types excluded from prediction (too broad or ambiguous)
-cancer_types_to_drop = [
+# These are used to mask out predictions for overly general cancer types
+CANCER_TYPES_TO_DROP = [
     "UDMN",
     "ADNOS",
     "CUP",
     "CUPNOS",
-    "BRCNOS",
-    "GNOS",
-    "SCCNOS",
-    "PDC",
-    "NSCLC",
-    "BRCA",
-    "SARCNOS",
-    "NETNOS",
-    "MEL",
-    "RCC",
-    "BRCANOS",
-    "COADREAD",
-    "MUP",
-    "NECNOS",
-    "UCEC",
     "NOT",
 ]
-col_indices_to_drop = [CANCER_TYPE_TO_INT_MAP[x] for x in cancer_types_to_drop]
 
 
 BATCH_SIZE = 8
@@ -98,8 +81,7 @@ def run(
     CANCER_TYPE_TO_INT_MAP_LOCAL = {v: k for k, v in INT_TO_CANCER_TYPE_MAP_LOCAL.items()}
 
     # Calculate col_indices_to_drop using local mapping
-    cancer_types_to_drop = ["UDMN", "ADNOS", "CUP", "CUPNOS", "NOT"]
-    col_indices_to_drop_local = [CANCER_TYPE_TO_INT_MAP_LOCAL[x] for x in cancer_types_to_drop if x in CANCER_TYPE_TO_INT_MAP_LOCAL]
+    col_indices_to_drop_local = [CANCER_TYPE_TO_INT_MAP_LOCAL[x] for x in CANCER_TYPES_TO_DROP if x in CANCER_TYPE_TO_INT_MAP_LOCAL]
 
     site_type = SiteType.METASTASIS if metastatic else SiteType.PRIMARY
 

tests/inference/test_aeon.py

@@ -4,43 +4,43 @@ import numpy as np
 import pytest
 import torch
 
-from mosaic.inference.aeon import (
+from mosaic.inference.aeon import CANCER_TYPES_TO_DROP
+from mosaic.inference.data import (
     CANCER_TYPE_TO_INT_MAP,
     INT_TO_CANCER_TYPE_MAP,
-    col_indices_to_drop,
 )
 
 
 class TestAeonConstants:
     """Test constants defined in aeon module."""
 
-    def test_col_indices_to_drop_is_list(self):
-        """Test that col_indices_to_drop is a list."""
-        assert isinstance(col_indices_to_drop, list)
+    def test_cancer_types_to_drop_is_list(self):
+        """Test that CANCER_TYPES_TO_DROP is a list."""
+        assert isinstance(CANCER_TYPES_TO_DROP, list)
 
-    def test_col_indices_to_drop_has_entries(self):
-        """Test that col_indices_to_drop has entries."""
-        assert len(col_indices_to_drop) > 0
+    def test_cancer_types_to_drop_has_entries(self):
+        """Test that CANCER_TYPES_TO_DROP has entries."""
+        assert len(CANCER_TYPES_TO_DROP) > 0
 
-    def test_col_indices_to_drop_are_integers(self):
-        """Test that all indices are integers."""
-        for idx in col_indices_to_drop:
-            assert isinstance(idx, int)
+    def test_cancer_types_to_drop_are_strings(self):
+        """Test that all cancer types are strings."""
+        for cancer_type in CANCER_TYPES_TO_DROP:
+            assert isinstance(cancer_type, str)
 
-    def test_col_indices_to_drop_are_valid(self):
-        """Test that all indices are valid cancer type indices."""
-        max_idx = max(CANCER_TYPE_TO_INT_MAP.values())
-        for idx in col_indices_to_drop:
-            assert 0 <= idx <= max_idx
+    def test_cancer_types_to_drop_are_valid(self):
+        """Test that all cancer types to drop are valid cancer type codes."""
+        # They should all be uppercase alphanumeric codes
+        for cancer_type in CANCER_TYPES_TO_DROP:
+            assert cancer_type.isupper()
+            assert len(cancer_type) >= 2
+            assert len(cancer_type) <= 10
 
-    def test_col_indices_to_drop_contains_expected_types(self):
+    def test_cancer_types_to_drop_contains_expected_types(self):
         """Test that specific cancer types are in the drop list."""
         # Check that some known cancer types to drop are in the list
-        drop_types = ["UDMN", "CUP", "BRCA", "MEL"]
-        for cancer_type in drop_types:
-            if cancer_type in CANCER_TYPE_TO_INT_MAP:
-                idx = CANCER_TYPE_TO_INT_MAP[cancer_type]
-                assert idx in col_indices_to_drop
+        expected_types = ["UDMN", "CUP", "NOT"]
+        for cancer_type in expected_types:
+            assert cancer_type in CANCER_TYPES_TO_DROP
 
     def test_cancer_type_maps_available(self):
         """Test that cancer type maps are available."""