{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":106680,"databundleVersionId":13374319,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Need for a uniform interface of running models\n\nAs described in the official competition page, to win the prize money, a prerequisite is that the code has to be made open-source. In addition, the top 10 submissions/teams will be invited to become co-authors in a scientific paper that involves further stress-testing of their models in a subsequent phase with many other datasets outside Kaggle platform. **To enable such further analyses and re-use of the models by the community, we strongly encourage** the participants to adhere to a code template that we provide through this repository that enables a uniform interface of running models: [https://github.com/uio-bmi/predict-airr](https://github.com/uio-bmi/predict-airr)\n\n\nIdeally, all the methods can be run in a unified way, e.g.,\n\n`python3 -m submission.main --train_dir /path/to/train_dir --test_dirs /path/to/test_dir_1 /path/to/test_dir_2 --out_dir /path/to/output_dir --n_jobs 4 --device cpu`\n\n## Adhering to code template on Kaggle Notebooks\n\nThose participants who make use of Kaggle resources and Kaggle notebooks to develop and run their code are also strongly encouraged to copy the code template, particularly the `ImmuneStatePredictor` class and any utility functions from the provided code template repository and adhere to the code template to enable a unified way of running different methods at a later stage. In this notebook, we copied the code template below for participants to paste into their respective Kaggle notebooks and edit as needed.","metadata":{}},{"cell_type":"code","source":"## imports required for the basic code template below.\n\nimport os\nfrom tqdm import tqdm\nimport pandas as pd\nimport numpy as np\nimport torch\nimport glob\nimport sys\nimport argparse\nfrom collections import defaultdict\nfrom typing import Iterator, Tuple, Union, List","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"## some utility functions such as data loaders, etc.\n\ndef load_data_generator(data_dir: str, metadata_filename='metadata.csv') -> Iterator[\n    Union[Tuple[str, pd.DataFrame, bool], Tuple[str, pd.DataFrame]]]:\n    \"\"\"\n    A generator to load immune repertoire data.\n\n    This function operates in two modes:\n    1.  If metadata is found, it yields data based on the metadata file.\n    2.  If metadata is NOT found, it uses glob to find and yield all '.tsv'\n        files in the directory.\n\n    Args:\n        data_dir (str): The path to the directory containing the data.\n\n    Yields:\n        An iterator of tuples. The format depends on the mode:\n        - With metadata: (repertoire_id, pd.DataFrame, label_positive)\n        - Without metadata: (filename, pd.DataFrame)\n    \"\"\"\n    metadata_path = os.path.join(data_dir, metadata_filename)\n\n    if os.path.exists(metadata_path):\n        metadata_df = pd.read_csv(metadata_path)\n        for row in metadata_df.itertuples(index=False):\n            file_path = os.path.join(data_dir, row.filename)\n            try:\n                repertoire_df = pd.read_csv(file_path, sep='\\t')\n                yield row.repertoire_id, repertoire_df, row.label_positive\n            except FileNotFoundError:\n                print(f\"Warning: File '{row.filename}' listed in metadata not found. Skipping.\")\n                continue\n    else:\n        search_pattern = os.path.join(data_dir, '*.tsv')\n        tsv_files = glob.glob(search_pattern)\n        for file_path in sorted(tsv_files):\n            try:\n                filename = os.path.basename(file_path)\n                repertoire_df = pd.read_csv(file_path, sep='\\t')\n                yield filename, repertoire_df\n            except Exception as e:\n                print(f\"Warning: Could not read file '{file_path}'. Error: {e}. Skipping.\")\n                continue\n\n\ndef load_full_dataset(data_dir: str) -> pd.DataFrame:\n    \"\"\"\n    Loads all TSV files from a directory and concatenates them into a single DataFrame.\n\n    This function handles two scenarios:\n    1. If metadata.csv exists, it loads data based on the metadata and adds\n       'repertoire_id' and 'label_positive' columns.\n    2. If metadata.csv does not exist, it loads all .tsv files and adds\n       a 'filename' column as an identifier.\n\n    Args:\n        data_dir (str): The path to the data directory.\n\n    Returns:\n        pd.DataFrame: A single, concatenated DataFrame containing all the data.\n    \"\"\"\n    metadata_path = os.path.join(data_dir, 'metadata.csv')\n    df_list = []\n    data_loader = load_data_generator(data_dir=data_dir)\n\n    if os.path.exists(metadata_path):\n        metadata_df = pd.read_csv(metadata_path)\n        total_files = len(metadata_df)\n        for rep_id, data_df, label in tqdm(data_loader, total=total_files, desc=\"Loading files\"):\n            data_df['ID'] = rep_id\n            data_df['label_positive'] = label\n            df_list.append(data_df)\n    else:\n        search_pattern = os.path.join(data_dir, '*.tsv')\n        total_files = len(glob.glob(search_pattern))\n        for filename, data_df in tqdm(data_loader, total=total_files, desc=\"Loading files\"):\n            data_df['ID'] = os.path.basename(filename).replace(\".tsv\", \"\")\n            df_list.append(data_df)\n\n    if not df_list:\n        print(\"Warning: No data files were loaded.\")\n        return pd.DataFrame()\n\n    full_dataset_df = pd.concat(df_list, ignore_index=True)\n    return full_dataset_df\n\n\ndef load_and_encode_kmers(data_dir: str, k: int = 3) -> Tuple[pd.DataFrame, pd.DataFrame]:\n    \"\"\"\n    Loading and k-mer encoding of repertoire data.\n\n    Args:\n        data_dir: Path to data directory\n        k: K-mer length\n\n    Returns:\n        Tuple of (encoded_features_df, metadata_df)\n        metadata_df always contains 'ID', and 'label_positive' if available\n    \"\"\"\n    from collections import Counter\n\n    metadata_path = os.path.join(data_dir, 'metadata.csv')\n    data_loader = load_data_generator(data_dir=data_dir)\n\n    repertoire_features = []\n    metadata_records = []\n\n    search_pattern = os.path.join(data_dir, '*.tsv')\n    total_files = len(glob.glob(search_pattern))\n\n    for item in tqdm(data_loader, total=total_files, desc=f\"Encoding {k}-mers\"):\n        if os.path.exists(metadata_path):\n            rep_id, data_df, label = item\n        else:\n            filename, data_df = item\n            rep_id = os.path.basename(filename).replace(\".tsv\", \"\")\n            label = None\n\n        kmer_counts = Counter()\n        for seq in data_df['junction_aa'].dropna():\n            for i in range(len(seq) - k + 1):\n                kmer_counts[seq[i:i + k]] += 1\n\n        repertoire_features.append({\n            'ID': rep_id,\n            **kmer_counts\n        })\n\n        metadata_record = {'ID': rep_id}\n        if label is not None:\n            metadata_record['label_positive'] = label\n        metadata_records.append(metadata_record)\n\n        del data_df, kmer_counts\n\n    features_df = pd.DataFrame(repertoire_features).fillna(0).set_index('ID')\n    features_df.fillna(0)\n    metadata_df = pd.DataFrame(metadata_records)\n\n    return features_df, metadata_df\n\n\ndef save_tsv(df: pd.DataFrame, path: str):\n    os.makedirs(os.path.dirname(path), exist_ok=True)\n    df.to_csv(path, sep='\\t', index=False)\n\n\ndef get_repertoire_ids(data_dir: str) -> list:\n    \"\"\"\n    Retrieves repertoire IDs from the metadata file or filenames in the directory.\n\n    Args:\n        data_dir (str): The path to the data directory.\n\n    Returns:\n        list: A list of repertoire IDs.\n    \"\"\"\n    metadata_path = os.path.join(data_dir, 'metadata.csv')\n\n    if os.path.exists(metadata_path):\n        metadata_df = pd.read_csv(metadata_path)\n        repertoire_ids = metadata_df['repertoire_id'].tolist()\n    else:\n        search_pattern = os.path.join(data_dir, '*.tsv')\n        tsv_files = glob.glob(search_pattern)\n        repertoire_ids = [os.path.basename(f).replace('.tsv', '') for f in sorted(tsv_files)]\n\n    return repertoire_ids\n\n\ndef generate_random_top_sequences_df(n_seq: int = 50000) -> pd.DataFrame:\n    \"\"\"\n    Generates a random DataFrame simulating top important sequences.\n\n    Args:\n        n_seq (int): Number of sequences to generate.\n\n    Returns:\n        pd.DataFrame: A DataFrame with columns 'ID', 'dataset', 'junction_aa', 'v_call', 'j_call'.\n    \"\"\"\n    seqs = set()\n    while len(seqs) < n_seq:\n        seq = ''.join(np.random.choice(list('ACDEFGHIKLMNPQRSTVWY'), size=15))\n        seqs.add(seq)\n    data = {\n        'junction_aa': list(seqs),\n        'v_call': ['TRBV20-1'] * n_seq,\n        'j_call': ['TRBJ2-7'] * n_seq,\n        'importance_score': np.random.rand(n_seq)\n    }\n    return pd.DataFrame(data)\n\n\ndef validate_dirs_and_files(train_dir: str, test_dirs: List[str], out_dir: str) -> None:\n    assert os.path.isdir(train_dir), f\"Train directory `{train_dir}` does not exist.\"\n    train_tsvs = glob.glob(os.path.join(train_dir, \"*.tsv\"))\n    assert train_tsvs, f\"No .tsv files found in train directory `{train_dir}`.\"\n    metadata_path = os.path.join(train_dir, \"metadata.csv\")\n    assert os.path.isfile(metadata_path), f\"`metadata.csv` not found in train directory `{train_dir}`.\"\n\n    for test_dir in test_dirs:\n        assert os.path.isdir(test_dir), f\"Test directory `{test_dir}` does not exist.\"\n        test_tsvs = glob.glob(os.path.join(test_dir, \"*.tsv\"))\n        assert test_tsvs, f\"No .tsv files found in test directory `{test_dir}`.\"\n\n    try:\n        os.makedirs(out_dir, exist_ok=True)\n        test_file = os.path.join(out_dir, \"test_write_permission.tmp\")\n        with open(test_file, \"w\") as f:\n            f.write(\"test\")\n        os.remove(test_file)\n    except Exception as e:\n        print(f\"Failed to create or write to output directory `{out_dir}`: {e}\")\n        sys.exit(1)\n\n\ndef concatenate_output_files(out_dir: str) -> None:\n    \"\"\"\n    Concatenates all test predictions and important sequences TSV files from the output directory.\n\n    This function finds all files matching the patterns:\n    - *_test_predictions.tsv\n    - *_important_sequences.tsv\n\n    and concatenates them to match the expected output format of submissions.csv.\n\n    Args:\n        out_dir (str): Path to the output directory containing the TSV files.\n\n    Returns:\n        pd.DataFrame: Concatenated DataFrame with predictions followed by important sequences.\n                     Columns: ['ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call']\n    \"\"\"\n    predictions_pattern = os.path.join(out_dir, '*_test_predictions.tsv')\n    sequences_pattern = os.path.join(out_dir, '*_important_sequences.tsv')\n\n    predictions_files = sorted(glob.glob(predictions_pattern))\n    sequences_files = sorted(glob.glob(sequences_pattern))\n\n    df_list = []\n\n    for pred_file in predictions_files:\n        try:\n            df = pd.read_csv(pred_file, sep='\\t')\n            df_list.append(df)\n        except Exception as e:\n            print(f\"Warning: Could not read predictions file '{pred_file}'. Error: {e}. Skipping.\")\n            continue\n\n    for seq_file in sequences_files:\n        try:\n            df = pd.read_csv(seq_file, sep='\\t')\n            df_list.append(df)\n        except Exception as e:\n            print(f\"Warning: Could not read sequences file '{seq_file}'. Error: {e}. Skipping.\")\n            continue\n\n    if not df_list:\n        print(\"Warning: No output files were found to concatenate.\")\n        concatenated_df = pd.DataFrame(\n            columns=['ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call'])\n    else:\n        concatenated_df = pd.concat(df_list, ignore_index=True)\n    submissions_file = os.path.join(out_dir, 'submissions.csv')\n    concatenated_df.to_csv(submissions_file, index=False)\n    print(f\"Concatenated output written to `{submissions_file}`.\")\n\n\ndef get_dataset_pairs(train_dir: str, test_dir: str) -> List[Tuple[str, List[str]]]:\n    \"\"\"Returns list of (train_path, [test_paths]) tuples for dataset pairs.\"\"\"\n    test_groups = defaultdict(list)\n    for test_name in sorted(os.listdir(test_dir)):\n        if test_name.startswith(\"test_dataset_\"):\n            base_id = test_name.replace(\"test_dataset_\", \"\").split(\"_\")[0]\n            test_groups[base_id].append(os.path.join(test_dir, test_name))\n\n    pairs = []\n    for train_name in sorted(os.listdir(train_dir)):\n        if train_name.startswith(\"train_dataset_\"):\n            train_id = train_name.replace(\"train_dataset_\", \"\")\n            train_path = os.path.join(train_dir, train_name)\n            pairs.append((train_path, test_groups.get(train_id, [])))\n\n    return pairs","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"## Main ImmuneStatePredictor class, where participants will fill in their implementations within the placeholders \n## and replace any example code lines with actual code that makes sense\n\n\nclass ImmuneStatePredictor:\n    \"\"\"\n    A template for predicting immune states from TCR repertoire data.\n\n    Participants should implement the logic for training, prediction, and\n    sequence identification within this class.\n    \"\"\"\n\n    def __init__(self, n_jobs: int = 1, device: str = 'cpu', **kwargs):\n        \"\"\"\n        Initializes the predictor.\n\n        Args:\n            n_jobs (int): Number of CPU cores to use for parallel processing.\n            device (str): The device to use for computation (e.g., 'cpu', 'cuda').\n            **kwargs: Additional hyperparameters for the model.\n        \"\"\"\n        total_cores = os.cpu_count()\n        if n_jobs == -1:\n            self.n_jobs = total_cores\n        else:\n            self.n_jobs = min(n_jobs, total_cores)\n        self.device = device\n        if device == 'cuda' and not torch.cuda.is_available():\n            print(\"Warning: 'cuda' was requested but is not available. Falling back to 'cpu'.\")\n            self.device = 'cpu'\n        else:\n            self.device = device\n        # --- your code starts here ---\n        # Example: Store hyperparameters, the actual model, identified important sequences, etc.\n\n        # NOTE: we encourage you to use self.n_jobs and self.device if appropriate in\n        # your implementation instead of hardcoding these values because your code may later be run in an\n        # environment with different hardware resources.\n\n        self.model = None\n        self.important_sequences_ = None\n        # --- your code ends here ---\n\n    def fit(self, train_dir_path: str):\n        \"\"\"\n        Trains the model on the provided training data.\n\n        Args:\n            train_dir_path (str): Path to the directory with training TSV files.\n\n        Returns:\n            self: The fitted predictor instance.\n        \"\"\"\n\n        # --- your code starts here ---\n        # Load the data, prepare suited representations as needed, train your model,\n        # and find the top k important sequences that best explain the labels.\n        # Example: Load the data. One possibility could be to use the provided utility function as shown below.\n\n        # full_train_dataset_df = load_full_dataset(train_dir_path)\n\n        #   Model Training\n        #    Example: self.model = SomeClassifier().fit(X_train, y_train)\n        self.model = \"some trained model\"  # Replace with your actual learnt model\n\n        #   Identify important sequences (can be done here or in the dedicated method)\n        #    Example:\n        self.important_sequences_ = self.identify_associated_sequences(top_k=50000, dataset_name=os.path.basename(train_dir_path))\n\n        # --- your code ends here ---\n        print(\"Training complete.\")\n        return self\n\n    def predict_proba(self, test_dir_path: str) -> pd.DataFrame:\n        \"\"\"\n        Predicts probabilities for examples in the provided path.\n\n        Args:\n            test_dir_path (str): Path to the directory with test TSV files.\n\n        Returns:\n            pd.DataFrame: A DataFrame with 'ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call' columns.\n        \"\"\"\n        print(f\"Making predictions for data in {test_dir_path}...\")\n        if self.model is None:\n            raise RuntimeError(\"The model has not been fitted yet. Please call `fit` first.\")\n\n        # --- your code starts here ---\n\n        # Example: Load the data. One possibility could be to use the provided utility function as shown below.\n\n        # full_test_dataset_df = load_full_dataset(test_dir_path)\n        repertoire_ids = get_repertoire_ids(test_dir_path)  # Replace with actual repertoire IDs from the test data\n\n        # Prediction\n        #    Example:\n        # draw random probabilities for demonstration purposes\n\n        probabilities = np.random.rand(len(repertoire_ids)) # Replace with true predicted probabilities from your model\n\n        # --- your code ends here ---\n\n        predictions_df = pd.DataFrame({\n            'ID': repertoire_ids,\n            'dataset': [os.path.basename(test_dir_path)] * len(repertoire_ids),\n            'label_positive_probability': probabilities\n        })\n\n        # to enable compatibility with the expected output format that includes junction_aa, v_call, j_call columns\n        predictions_df['junction_aa'] = -999.0\n        predictions_df['v_call'] = -999.0\n        predictions_df['j_call'] = -999.0\n\n        predictions_df = predictions_df[['ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call']]\n\n        print(f\"Prediction complete on {len(repertoire_ids)} examples in {test_dir_path}.\")\n        return predictions_df\n\n    def identify_associated_sequences(self, dataset_name: str, top_k: int = 50000) -> pd.DataFrame:\n        \"\"\"\n        Identifies the top \"k\" important sequences (rows) from the training data that best explain the labels.\n\n        Args:\n            top_k (int): The number of top sequences to return (based on some scoring mechanism).\n\n        Returns:\n            pd.DataFrame: A DataFrame with 'ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call' columns.\n        \"\"\"\n\n        # --- your code starts here ---\n        \n        # Return the top k sequences, sorted based on some form of importance score.\n        # Example:\n        # all_sequences_scored = self._score_all_sequences()\n        \n        all_sequences_scored = generate_random_top_sequences_df(n_seq=top_k)  # Replace with your way of identifying top k sequences\n\n        # note that all_sequences_scored should contain a 'importance_score' column that will be used further below\n        \n        # --- your code ends here ---\n\n        top_sequences_df = all_sequences_scored.nlargest(top_k, 'importance_score')\n        top_sequences_df = top_sequences_df[['junction_aa', 'v_call', 'j_call']]\n        top_sequences_df['dataset'] = dataset_name\n        top_sequences_df['ID'] = range(1, len(top_sequences_df)+1)\n        top_sequences_df['ID'] = top_sequences_df['dataset'] + '_seq_top_' + top_sequences_df['ID'].astype(str)\n        top_sequences_df['label_positive_probability'] = -999.0 # to enable compatibility with the expected output format\n        top_sequences_df = top_sequences_df[['ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call']]\n\n        return top_sequences_df","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"## The `main` workflow that uses your implementation of the ImmuneStatePredictor class to train, identify important sequences and predict test labels\n\n\ndef _train_predictor(predictor: ImmuneStatePredictor, train_dir: str):\n    \"\"\"Trains the predictor on the training data.\"\"\"\n    print(f\"Fitting model on examples in ` {train_dir} `...\")\n    predictor.fit(train_dir)\n\n\ndef _generate_predictions(predictor: ImmuneStatePredictor, test_dirs: List[str]) -> pd.DataFrame:\n    \"\"\"Generates predictions for all test directories and concatenates them.\"\"\"\n    all_preds = []\n    for test_dir in test_dirs:\n        print(f\"Predicting on examples in ` {test_dir} `...\")\n        preds = predictor.predict_proba(test_dir)\n        if preds is not None and not preds.empty:\n            all_preds.append(preds)\n        else:\n            print(f\"Warning: No predictions returned for {test_dir}\")\n    if all_preds:\n        return pd.concat(all_preds, ignore_index=True)\n    return pd.DataFrame()\n\n\ndef _save_predictions(predictions: pd.DataFrame, out_dir: str, train_dir: str) -> None:\n    \"\"\"Saves predictions to a TSV file.\"\"\"\n    if predictions.empty:\n        raise ValueError(\"No predictions to save - predictions DataFrame is empty\")\n\n    preds_path = os.path.join(out_dir, f\"{os.path.basename(train_dir)}_test_predictions.tsv\")\n    save_tsv(predictions, preds_path)\n    print(f\"Predictions written to `{preds_path}`.\")\n\n\ndef _save_important_sequences(predictor: ImmuneStatePredictor, out_dir: str, train_dir: str) -> None:\n    \"\"\"Saves important sequences to a TSV file.\"\"\"\n    seqs = predictor.important_sequences_\n    if seqs is None or seqs.empty:\n        raise ValueError(\"No important sequences available to save\")\n\n    seqs_path = os.path.join(out_dir, f\"{os.path.basename(train_dir)}_important_sequences.tsv\")\n    save_tsv(seqs, seqs_path)\n    print(f\"Important sequences written to `{seqs_path}`.\")\n\n\ndef main(train_dir: str, test_dirs: List[str], out_dir: str, n_jobs: int, device: str) -> None:\n    validate_dirs_and_files(train_dir, test_dirs, out_dir)\n    predictor = ImmuneStatePredictor(n_jobs=n_jobs,\n                                     device=device)  # instantiate with any other parameters as defined by you in the class\n    _train_predictor(predictor, train_dir)\n    predictions = _generate_predictions(predictor, test_dirs)\n    _save_predictions(predictions, out_dir, train_dir)\n    _save_important_sequences(predictor, out_dir, train_dir)\n\n\ndef run():\n    parser = argparse.ArgumentParser(description=\"Immune State Predictor CLI\")\n    parser.add_argument(\"--train_dir\", required=True, help=\"Path to training data directory\")\n    parser.add_argument(\"--test_dirs\", required=True, nargs=\"+\", help=\"Path(s) to test data director(ies)\")\n    parser.add_argument(\"--out_dir\", required=True, help=\"Path to output directory\")\n    parser.add_argument(\"--n_jobs\", type=int, default=1,\n                        help=\"Number of CPU cores to use. Use -1 for all available cores.\")\n    parser.add_argument(\"--device\", type=str, default='cpu', choices=['cpu', 'cuda'],\n                        help=\"Device to use for computation ('cpu' or 'cuda').\")\n    args = parser.parse_args()\n    main(args.train_dir, args.test_dirs, args.out_dir, args.n_jobs, args.device)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_datasets_dir = \"/kaggle/input/adaptive-immune-profiling-challenge-2025/train_datasets/train_datasets\"\ntest_datasets_dir = \"/kaggle/input/adaptive-immune-profiling-challenge-2025/test_datasets/test_datasets\"\nresults_dir = \"/kaggle/working/results\"\n\ntrain_test_dataset_pairs = get_dataset_pairs(train_datasets_dir, test_datasets_dir)\n\nfor train_dir, test_dirs in train_test_dataset_pairs:\n    main(train_dir=train_dir, test_dirs=test_dirs, out_dir=results_dir, n_jobs=4, device=\"cpu\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"concatenate_output_files(out_dir=results_dir)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}