data.py

"""Loading of pre-processed breath segments (output of the preprocess_*.py scripts).

A processed data directory looks like:
    <dir>/waveforms/patient_<k>   parquet, n_breaths * SEG_LEN rows, columns paw, flow, vol, breath
    <dir>/labels/label_<k>        parquet, one column per breath, one row per label field
Files are matched on the index <k>.
"""
import re
from pathlib import Path

import numpy as np
import pandas as pd
import tensorflow as tf

from config import SEG_LEN, BATCH_SIZE

META_COLUMNS = ['disease', 'patient_nr']   # only present for simulated data


def set_gpu(gpu_nr=None):
    """Restrict TensorFlow to one GPU (if given and available) and enable memory growth."""
    gpus = tf.config.list_physical_devices('GPU')
    if not gpus:
        print("No GPU found, running on CPU.")
        return
    if gpu_nr is not None:
        if gpu_nr >= len(gpus):
            raise ValueError(f"GPU {gpu_nr} requested but only {len(gpus)} GPU(s) available.")
        gpus = [gpus[gpu_nr]]
        tf.config.set_visible_devices(gpus, 'GPU')
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)


def _index(path):
    m = re.search(r'(\d+)$', Path(path).stem)
    return int(m.group(1)) if m else -1


def list_pairs(data_dir, max_files=None):
    """Return matched (waveform file, label file) pairs, sorted by index."""
    data_dir = Path(data_dir)
    waves = {_index(p): p for p in (data_dir / 'waveforms').iterdir() if p.is_file()}
    labels = {_index(p): p for p in (data_dir / 'labels').iterdir() if p.is_file()}
    keys = sorted(set(waves) & set(labels))
    missing = set(waves) ^ set(labels)
    if missing:
        print(f"Warning: {len(missing)} file(s) in {data_dir} without a matching waveform/label file.")
    if max_files is not None:
        keys = keys[:max_files]
    if not keys:
        raise FileNotFoundError(f"No matched waveform/label files found in {data_dir}")
    return [(waves[k], labels[k]) for k in keys]


def _read(path):
    return pd.read_parquet(path, engine='fastparquet')


def _segments(wave_dfs):
    """Cut each file into SEG_LEN-long breaths and keep pressure and flow."""
    return np.array([df[['paw', 'flow']].to_numpy()[i:i + SEG_LEN]
                     for df in wave_dfs for i in range(0, len(df), SEG_LEN)], dtype=np.float32)


def _label_matrix(label_df):
    """Label parquet (fields x breaths) -> array (breaths x fields), without metadata columns."""
    return np.transpose(label_df.drop(columns=META_COLUMNS, errors='ignore').to_numpy(dtype=np.float32))


def _to_dataset(x, y, batch_size, shuffle, drop_remainder=False, repeat=1):
    ds = tf.data.Dataset.zip((tf.data.Dataset.from_tensor_slices(x), tf.data.Dataset.from_tensor_slices(y)))
    if shuffle:
        ds = ds.shuffle(len(ds), reshuffle_each_iteration=True)
    ds = ds.batch(batch_size, drop_remainder=drop_remainder)
    if repeat > 1:
        ds = ds.repeat(repeat)
    if shuffle:
        ds = ds.shuffle(1000)
        ds = ds.shuffle(len(ds), reshuffle_each_iteration=True)
    return ds


def load_unlabeled(data_dir, max_files=None, to_dataset=True):
    """Unlabeled clinical data. Labels are only the mechanical cycle: [tim, tem] (scaled)."""
    pairs = list_pairs(data_dir, max_files)
    waves = [_read(w) for w, _ in pairs]
    labels = [_label_matrix(_read(l)) for _, l in pairs]
    with tf.device('/CPU:0'):
        x = tf.constant(_segments(waves))
        y = tf.constant(np.vstack(labels))
        if not to_dataset:
            return x, y
        return _to_dataset(x, y, BATCH_SIZE, shuffle=True, drop_remainder=True)


def load_labeled(data_dir, to_dataset=True, batch_size=256):
    """Labeled data for validation/testing. Labels have 13 columns (see docs/code_overview.md)."""
    pairs = list_pairs(data_dir)
    waves = [_read(w) for w, _ in pairs]
    labels = [_label_matrix(_read(l)) for _, l in pairs]
    with tf.device('/CPU:0'):
        x = tf.constant(_segments(waves))
        y = tf.constant(np.vstack(labels))
        if not to_dataset:
            return x, y
        return _to_dataset(x, y, batch_size, shuffle=False)


def load_simulated(data_dir, repeat=1, sim_val=False, n_val_per_disease=2, seed=42):
    """Labeled simulated training data.

    Training labels have 15 columns: the 13 label fields + disease code + patient number
    (the last two are dropped in SemiSup.train_step).
    If sim_val is True, n_val_per_disease simulated patients per disease are held out and
    returned as a validation set: (train_dataset, valid_dataset). Otherwise returns train_dataset.
    """
    pairs = list_pairs(data_dir)
    waves = [_read(w) for w, _ in pairs]
    label_dfs = [_read(l) for _, l in pairs]

    meta = pd.concat([df[META_COLUMNS].iloc[[0]] for df in label_dfs], ignore_index=True)
    disease_codes = {d: i for i, d in enumerate(sorted(meta['disease'].unique()))}
    labels = []
    for df, (_, row) in zip(label_dfs, meta.iterrows()):
        lab = _label_matrix(df)
        extra = np.tile([disease_codes[row['disease']], float(row['patient_nr'])], (lab.shape[0], 1))
        labels.append(np.hstack([lab, extra]).astype(np.float32))

    if sim_val:
        held_out = (meta.drop_duplicates().groupby('disease')
                    .sample(n=n_val_per_disease, random_state=seed))
        is_val = meta.set_index(META_COLUMNS).index.isin(held_out.set_index(META_COLUMNS).index)
    else:
        is_val = np.zeros(len(meta), dtype=bool)
    train_idx = np.where(~is_val)[0]
    val_idx = np.where(is_val)[0]

    with tf.device('/CPU:0'):
        x = tf.constant(_segments([waves[i] for i in train_idx]))
        y = tf.constant(np.vstack([labels[i] for i in train_idx]))
        train_ds = _to_dataset(x, y, BATCH_SIZE, shuffle=True, drop_remainder=True, repeat=repeat)
        if not sim_val:
            return train_ds
        xv = tf.constant(_segments([waves[i] for i in val_idx]))
        yv = tf.constant(np.vstack([labels[i] for i in val_idx])[:, :-2])
        return train_ds, _to_dataset(xv, yv, 256, shuffle=False)