Add the end-to-end runner, chest X-ray as a third modality, and run provenance

pyhealth/datasets/configs/mimic4_cxr.yaml

@@ -5,7 +5,7 @@ tables:
     patient_id: "subject_id"
     timestamp:
       - "studydate"
-      - "studytime"
+      - "studytime_normalized"
     timestamp_format: "%Y%m%d%H%M%S"
     attributes:
       - "image_path"

pyhealth/datasets/mimic4.py

@@ -191,6 +191,15 @@ def __init__(
         log_memory_usage(f"After initializing {dataset_name}")
 
     def prepare_metadata(self, root: str) -> None:
+        prepared_path = os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth.csv")
+        # The prepared file holds absolute JPEG paths. To rewrite its 100+ MB
+        # CSV for every CXR experiment is needless shared-filesystem IO, and it
+        # races a concurrent reader, so reuse the file when it is present. An
+        # installation with raw metadata only still uses the builder below.
+        if os.path.exists(prepared_path):
+            header = pd.read_csv(prepared_path, nrows=1)
+            if "image_path" in header.columns:
+                return
         metadata = pd.read_csv(
             os.path.join(root, "mimic-cxr-2.0.0-metadata.csv.gz"), dtype=str
         )
@@ -217,9 +226,7 @@ def process_image_path(x):
 
         metadata["image_path"] = metadata.apply(process_image_path, axis=1)
 
-        metadata.to_csv(
-            os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth.csv"), index=False
-        )
+        metadata.to_csv(prepared_path, index=False)
         return
 
 
@@ -228,7 +235,7 @@ class MIMIC4CXRSunlabDataset(BaseDataset):
     Sunlab variant of the MIMIC-CXR Chest X-ray dataset.
 
     This variant uses the existing metadata CSV and derives flattened image
-    paths at ``images/{dicom_id}.jpg``.
+    paths at ``{images|resized_images}/{dicom_id}.jpg``.
     """
 
     def __init__(
@@ -245,7 +252,11 @@ def __init__(
                 os.path.dirname(__file__), "configs", "mimic4_cxr_sunlab.yaml"
             )
             logger.info(f"Using default Sunlab CXR config: {config_path}")
-        self.prepare_metadata(root)
+        metadata_csv = self.prepare_metadata(root, cache_dir=cache_dir)
+        if os.path.dirname(os.path.abspath(metadata_csv)) != os.path.abspath(root):
+            config_path = self._rewrite_sunlab_config(
+                config_path, metadata_csv, cache_dir or os.path.dirname(metadata_csv)
+            )
         log_memory_usage(f"Before initializing {dataset_name}")
         super().__init__(
             root=root,
@@ -267,20 +278,52 @@ def _resolve_column_name(columns: List[str], target: str) -> str:
             )
         return resolved
 
-    def prepare_metadata(self, root: str) -> None:
+    @staticmethod
+    def _rewrite_sunlab_config(
+        config_path: str, metadata_csv: str, dest_dir: str
+    ) -> str:
+        """Point the sunlab YAML at a metadata CSV that is not under root."""
+        with open(config_path, encoding="utf-8") as f:
+            text = f.read()
+        rewritten = text.replace(
+            "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv",
+            metadata_csv,
+        )
+        os.makedirs(dest_dir, exist_ok=True)
+        out = os.path.join(dest_dir, "mimic4_cxr_sunlab.generated.yaml")
+        with open(out, "w", encoding="utf-8") as f:
+            f.write(rewritten)
+        return out
+
+    def prepare_metadata(
+        self, root: str, cache_dir: Optional[str] = None
+    ) -> str:
         metadata_path = os.path.join(root, "mimic-cxr-2.0.0-metadata.csv")
         if not os.path.exists(metadata_path):
             raise FileNotFoundError(
                 f"Sunlab metadata file not found: {metadata_path}. "
                 "Expected existing metadata linked by dicom_id/subject_id/study_id."
             )
 
-        images_dir = os.path.join(root, "images")
-        if not os.path.isdir(images_dir):
+        # The flattened layout appears under more than one directory name
+        # depending on how the set was produced, so accept either rather than
+        # hardcoding one and failing on a complete, correct dataset.
+        candidates = ("images", "resized_images")
+        images_dir = next(
+            (
+                os.path.join(root, name)
+                for name in candidates
+                if os.path.isdir(os.path.join(root, name))
+            ),
+            None,
+        )
+        if images_dir is None:
             raise FileNotFoundError(
-                f"Sunlab images directory not found: {images_dir}. "
-                "Expected flattened image files at images/{dicom_id}.jpg."
+                f"No flattened image directory under {root}. Looked for "
+                f"{', '.join(candidates)}, each expected to hold "
+                "{dicom_id}.jpg."
             )
+        images_subdir = os.path.basename(images_dir)
 
         metadata = pd.read_csv(metadata_path, dtype=str)
 
@@ -307,15 +350,35 @@ def normalize_studytime(value: Optional[str]) -> str:
         metadata[study_time_col] = metadata[study_time_col].apply(normalize_studytime)
 
         metadata["image_path"] = metadata[dicom_col].apply(
-            lambda dicom_id: os.path.join(root, "images", f"{dicom_id}.jpg")
+            lambda dicom_id: os.path.join(root, images_subdir, f"{dicom_id}.jpg")
         )
 
         # Align with existing config conventions by using lowercase headers.
         metadata.columns = [col.lower() for col in metadata.columns]
 
-        metadata.to_csv(
-            os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"),
-            index=False,
+        filename = "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"
+        dest_dirs = []
+        if cache_dir:
+            dest_dirs.append(str(cache_dir))
+        dest_dirs.append(root)
+
+        for d in dest_dirs:
+            existing = os.path.join(d, filename)
+            if os.path.isfile(existing):
+                return existing
+
+        last_err: Optional[OSError] = None
+        for d in dest_dirs:
+            os.makedirs(d, exist_ok=True)
+            dest = os.path.join(d, filename)
+            try:
+                metadata.to_csv(dest, index=False)
+                return dest
+            except OSError as exc:
+                last_err = exc
+                continue
+        raise PermissionError(
+            f"Could not write {filename} under {dest_dirs}: {last_err}"
         )
 
 

pyhealth/models/jamba_ehr.py

@@ -73,7 +73,7 @@ class JambaLayer(nn.Module):
     Args:
         feature_size (int): Hidden dimension shared by all layers.
         num_transformer_layers (int): Number of attention layers. Default 2.
-        num_mamba_layers (int): Number of SSM layers. Default 6.
+        num_mamba_layers (int): Number of SSM layers. Default 2.
         heads (int): Attention heads for Transformer layers. Default 4.
         dropout (float): Dropout rate for Transformer layers. Default 0.3.
         state_size (int): SSM state size for Mamba layers. Default 16.
@@ -94,7 +94,7 @@ def __init__(
         self,
         feature_size: int,
         num_transformer_layers: int = 2,
-        num_mamba_layers: int = 6,
+        num_mamba_layers: int = 2,
         heads: int = 4,
         dropout: float = 0.3,
         state_size: int = 16,
@@ -187,7 +187,7 @@ class JambaEHR(BaseModel):
         dataset (SampleDataset): Dataset providing processed inputs.
         embedding_dim (int): Embedding and hidden dimension. Default 128.
         num_transformer_layers (int): Transformer layers per stream. Default 2.
-        num_mamba_layers (int): Mamba layers per stream. Default 6.
+        num_mamba_layers (int): Mamba layers per stream. Default 2.
         heads (int): Attention heads per Transformer block. Default 4.
         dropout (float): Dropout rate. Default 0.3.
         state_size (int): SSM state size in Mamba blocks. Default 16.
@@ -237,7 +237,7 @@ def __init__(
         dataset: SampleDataset,
         embedding_dim: int = 128,
         num_transformer_layers: int = 2,
-        num_mamba_layers: int = 6,
+        num_mamba_layers: int = 2,
         heads: int = 4,
         dropout: float = 0.3,
         state_size: int = 16,

pyhealth/processors/time_image_processor.py

@@ -178,6 +178,26 @@ def _zero_image_tensor(self) -> torch.Tensor:
             c = 3
         return torch.zeros(c, self.image_size, self.image_size)
 
+    @property
+    def in_channels(self) -> int:
+        """Channel count implied by ``mode``.
+
+        The unified embedding sizes its patch embedding from this. Without it
+        the model defaulted to 3 while a greyscale task produced 1, and the
+        mismatch only appeared at the first forward pass:
+
+            RuntimeError: Given groups=1, weight of size [128, 3, 16, 16],
+            expected input[16, 1, 224, 224] to have 3 channels
+
+        Deriving it here means the two cannot disagree.
+        """
+        # Must match _zero_image_tensor exactly, or a placeholder image would
+        # carry a different channel count from a real one.
+        if self.n_channels is not None:
+            return int(self.n_channels)
+        return {"1": 1, "L": 1, "LA": 2, "RGB": 3, "RGBA": 4}.get(self.mode or "RGB", 3)
+
+
     def _load_single_image(self, path: Union[str, Path]) -> torch.Tensor:
         """Load and transform a single image from disk.
 

pyhealth/tasks/__init__.py

@@ -74,5 +74,6 @@
     ClinicalNotesMIMIC4,
     ClinicalNotesICDLabsMIMIC4,
     ClinicalNotesICDLabsCXRMIMIC4,
+    CXRMultimodalMIMIC4,
 )
 from .patient_linkage_mimic3 import PatientLinkageMIMIC3Task