Remove new-data reconstruction evaluation utility and README references

Delete evaluate_new_data.py and its command reference from BioMod/eval_new_data. Remove the corresponding workflow, checkpoint-interface notes, and directory-tree entries from the root and source READMEs. Model weights, paired configuration, datasets, notebooks, and other source files are unchanged.

BioMod/README.md

@@ -41,9 +41,6 @@ BioMod/                                                   # Task code: training,
 ├── drug_response/                                        # Single-drug response prediction and training
 │   ├── cmd.txt
 │   └── train_drug_performer_module.py
-├── eval_new_data/                                        # Reconstruction error and generalization on new data
-│   ├── cmd.txtx
-│   └── evaluate_new_data.py
 ├── gene_perturbation/                                    # Perturbation response prediction: processing, training and inference
 │   ├── data_utils_local.py
 │   ├── gears_reduce_local.py

BioMod/eval_new_data/cmd.txtx

@@ -1,19 +0,0 @@
-# masked 模式（与预训练 valid loss 直接可比）
-python evaluate_new_data.py \
---ckpt_dir ../mae_tost_48M_cross_attn_only_pure_memory/ \
---model_name mae_tost_48M_cross_attn_only_pure_memory \
---data_dir ./npz \
---module_data_dir ../module_data_all/ \
---eval_mode masked \
---batch_size 320 \
---output_csv eval_masked.csv 
-
-# full_recon 模式（不 mask，全 module 重建 loss） 
-python evaluate_new_data.py \
---ckpt_dir /path/to/mae_tost_48M_cross_attn_only_pure_memory/ \
---model_name mae_tost_48M_cross_attn_only_pure_memory \
---data_dir /path/to/20260423测试集/npz/npz \ 
---module_data_dir ./module_data_all/ \ 
---eval_mode full_recon \
---batch_size 32 \
---output_csv eval_full_recon.csv 

BioMod/eval_new_data/evaluate_new_data.py

@@ -1,375 +0,0 @@
-# -*- coding: utf-8 -*-
-"""
-Evaluate the reconstruction loss of a pre-trained model on a new dataset (DDP multi-card version).
-Used to verify the generalization ability of the model on unseen donor data, compared with pre-training valid loss.
-
-Two evaluation modes are supported:
-  masked — consistent with pre-training validation, using the same mask_ratio, for direct comparison
-  full_recon — without masking, calculate the reconstruction loss of all modules
-
-Usage example (single card):
-  torchrun --nproc_per_node=1 evaluate_new_data.py \
-      --ckpt_dir ./ckpts/ \
-      --model_name mae_tost_48M_cross_attn_only_pure_memory \
-      --data_dir /path/to/20260423test set/npz/npz \
-      --eval_mode masked --batch_size 32
-
-Usage example (4 card DDP):
-  torchrun --nproc_per_node=4 evaluate_new_data.py \
-      --ckpt_dir ./ckpts/ \
-      --model_name mae_tost_48M_cross_attn_only_pure_memory \
-      --data_dir /path/to/20260423test set/npz/npz \
-      --eval_mode masked --batch_size 32
-"""
-
-import os
-import re
-import csv
-import sys
-import glob
-import time
-import pickle
-import inspect
-import argparse
-from collections import defaultdict
-
-import numpy as np
-import torch
-import torch.distributed as dist
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.utils.data import Dataset, DataLoader
-from torch.utils.data.distributed import DistributedSampler
-from scipy import sparse
-
-sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
-
-from models_performer_module import PerformerEncoderOnly_Module
-from module_dataset import load_module_data
-
-
-# ---------------------------------------------------------------------------
-# Dataset
-# ---------------------------------------------------------------------------
-
-class DonorSCDataset(Dataset):
-    """Dataset of a single donor, returning single-cell expression vectors row-wise from dense numpy."""
-
-    def __init__(self, csr_matrix):
-        self.data = csr_matrix.toarray()
-
-    def __len__(self):
-        return self.data.shape[0]
-
-    def __getitem__(self, idx):
-        return torch.from_numpy(self.data[idx].astype(np.float32))
-
-
-# ---------------------------------------------------------------------------
-# Checkpoint discovery
-# ---------------------------------------------------------------------------
-
-def discover_checkpoints(ckpt_dir, model_name, epoch_filter=None):
-    pattern = os.path.join(ckpt_dir, f"{model_name}_epoch*.pth")
-    files = glob.glob(pattern)
-    results = []
-    for f in files:
-        match = re.search(r'epoch(\d+)', os.path.basename(f))
-        if match:
-            epoch_num = int(match.group(1))
-            if epoch_filter is None or epoch_num in epoch_filter:
-                results.append((epoch_num, f))
-    results.sort(key=lambda x: x[0])
-    return results
-
-
-def load_model_config(ckpt_dir, model_name):
-    pkl_path = os.path.join(ckpt_dir, f"{model_name}.pkl")
-    if not os.path.isfile(pkl_path):
-        raise FileNotFoundError(f"Model configuration file not found: {pkl_path}")
-    with open(pkl_path, 'rb') as f:
-        return pickle.load(f)
-
-
-def filter_model_config(config):
-    """Filter out obsolete parameters in pkl that are not accepted by the model __init__."""
-    valid = set(inspect.signature(PerformerEncoderOnly_Module.__init__).parameters.keys())
-    removed = {k: config.pop(k) for k in list(config) if k not in valid and k != "self"}
-    return removed
-
-
-# ---------------------------------------------------------------------------
-# Evaluation
-# ---------------------------------------------------------------------------
-
-def evaluate_donor_ddp(model, dataloader, device, eval_mode, mask_ratio, ptdtype, seed=None):
-    """
-    Calculate the loss for this rank shard of a single donor and return (loss_sum, batch_count).
-    Aggregation is done by the caller via all_reduce.
-    """
-    model.eval()
-    loss_sum = 0.0
-    batch_count = 0
-
-    raw_model = model.module if isinstance(model, DDP) else model
-
-    with torch.no_grad():
-        for batch_idx, data in enumerate(dataloader):
-            data = data.to(device, non_blocking=True)
-            if torch.isnan(data).any() or torch.isinf(data).any():
-                continue
-
-            if eval_mode == "masked" and seed is not None:
-                torch.manual_seed(seed + batch_idx)
-                torch.cuda.manual_seed(seed + batch_idx)
-
-            with torch.autocast(device_type='cuda', dtype=ptdtype):
-                if eval_mode == "masked":
-                    loss, _, _ = raw_model(data, mask_ratio=mask_ratio)
-                else:
-                    latent, mask_zero, ids_restore, module_targets = raw_model.forward_encoder(
-                        data, mask_ratio=0.0
-                    )
-                    dt = raw_model.decoder_type
-                    if dt == "performer":
-                        preds = raw_model.forward_decoder_performer(latent, ids_restore)
-                    elif dt == "cross_attn":
-                        preds = raw_model.forward_decoder_cross_attn(latent, mask_zero, ids_restore)
-                    elif dt == "simple":
-                        preds = raw_model.forward_decoder_simple(latent, ids_restore)
-                    elif dt == "none":
-                        preds = raw_model.forward_predictor_from_encoder(latent)
-                    else:
-                        raise RuntimeError(f"Unsupported decoder_type: {dt}")
-                    mask_all = torch.ones_like(mask_zero)
-                    loss = raw_model.compute_loss(preds, module_targets, mask_all)
-
-            if torch.isnan(loss) or torch.isinf(loss):
-                continue
-            loss_sum += loss.item()
-            batch_count += 1
-
-    return loss_sum, batch_count
-
-
-def reduce_scalar(value, device):
-    """all_reduce is a scalar and returns the global summation result."""
-    t = torch.tensor(value, dtype=torch.float64, device=device)
-    dist.all_reduce(t, op=dist.ReduceOp.SUM)
-    return t.item()
-
-
-def discover_datasets(data_dir):
-    datasets = []
-    for name in sorted(os.listdir(data_dir)):
-        subdir = os.path.join(data_dir, name)
-        if os.path.isdir(subdir):
-            npz_files = sorted([f for f in os.listdir(subdir) if f.endswith('.npz')])
-            if npz_files:
-                datasets.append((name, subdir, npz_files))
-    return datasets
-
-
-# ---------------------------------------------------------------------------
-# Main
-# ---------------------------------------------------------------------------
-
-def main():
-    parser = argparse.ArgumentParser(description="Evaluate pre-trained model reconstruction loss (DDP) on new datasets")
-    parser.add_argument("--ckpt_dir", type=str, required=True)
-    parser.add_argument("--model_name", type=str, required=True)
-    parser.add_argument("--data_dir", type=str, required=True)
-    parser.add_argument("--gene_list_path", type=str, default="../../selected_genes_27k.txt")
-    parser.add_argument("--gene_mapping_path", type=str, default="../../gene2ens.csv")
-    parser.add_argument("--module_data_dir", type=str, default=None)
-    parser.add_argument("--gene2vec_data_path", type=str, default="../../selected_gene2vec_27k.npy")
-    parser.add_argument("--batch_size", type=int, default=32)
-    parser.add_argument("--output_csv", type=str, default="eval_new_data_results.csv")
-    parser.add_argument("--epochs", type=str, default=None,
-                        help="Comma separated list of epochs (e.g. '1,5,10'), default is all")
-    parser.add_argument("--eval_mode", type=str, default="masked",
-                        choices=["masked", "full_recon"])
-    parser.add_argument("--mask_ratio", type=float, default=0.75)
-    parser.add_argument("--seed", type=int, default=2021)
-    parser.add_argument("--num_workers", type=int, default=4)
-    parser.add_argument("--amp_dtype", type=str, default="bfloat16",
-                        choices=["float32", "bfloat16", "float16"])
-    args = parser.parse_args()
-
-    # ---- DDP initialization ----
-    dist.init_process_group(backend='nccl')
-    local_rank = int(os.environ.get('LOCAL_RANK', 0))
-    rank = dist.get_rank()
-    world_size = dist.get_world_size()
-    device = torch.device(f'cuda:{local_rank}')
-    torch.cuda.set_device(device)
-    is_master = (rank == 0)
-
-    ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[args.amp_dtype]
-
-    if is_master:
-        print(f"World size: {world_size}", flush=True)
-        print(f"Eval mode: {args.eval_mode}", flush=True)
-        if args.eval_mode == "masked":
-            print(f"Mask ratio: {args.mask_ratio}", flush=True)
-
-    # ---- Load model configuration ----
-    if is_master:
-        print(f"\nLoad model configuration: {args.ckpt_dir}/{args.model_name}.pkl", flush=True)
-    model_config = load_model_config(args.ckpt_dir, args.model_name)
-    removed = filter_model_config(model_config)
-    if is_master and removed:
-        for k, v in removed.items():
-            print(f"  Skip obsolete parameters: {k}={v!r}", flush=True)
-
-    if args.gene2vec_data_path:
-        model_config["gene2vec_file"] = args.gene2vec_data_path
-
-    if "module_sets" in model_config and model_config["module_sets"]:
-        if is_master:
-            print(f"Loaded from pkl config {len(model_config['module_sets'])} modules", flush=True)
-    else:
-        if args.module_data_dir is None:
-            raise ValueError("There are no module_sets in the model configuration, please specify via --module_data_dir")
-        model_config["module_sets"] = load_module_data(
-            module_data_dir=args.module_data_dir,
-            gene_list_path=args.gene_list_path,
-            gene_mapping_path=args.gene_mapping_path,
-            verbose=is_master,
-        )
-
-    # ---- Discover checkpoints & data sets ----
-    epoch_filter = set(int(e) for e in args.epochs.split(",")) if args.epochs else None
-    checkpoints = discover_checkpoints(args.ckpt_dir, args.model_name, epoch_filter)
-    datasets = discover_datasets(args.data_dir)
-
-    if is_master:
-        if not checkpoints:
-            print(f"checkpoint not found: {args.ckpt_dir}/{args.model_name}_epoch*.pth", flush=True)
-        else:
-            print(f"turn up {len(checkpoints)} checkpoints: {[e for e, _ in checkpoints]}", flush=True)
-        for ds_name, _, npz_files in datasets:
-            print(f"Dataset {ds_name}: {len(npz_files)} donors", flush=True)
-        if not datasets:
-            print(f"Warning: in {args.data_dir} The subdirectory containing the .npz file was not found!", flush=True)
-
-    if not checkpoints or not datasets:
-        dist.destroy_process_group()
-        return
-
-    # ----CSV (master write only)----
-    csv_file, writer = None, None
-    if is_master:
-        csv_file = open(args.output_csv, 'w', newline='', encoding='utf-8')
-        writer = csv.writer(csv_file)
-        writer.writerow(["epoch", "dataset", "donor", "num_cells", "mean_loss", "eval_mode"])
-
-    # ----Build the model (create it only once, only reload the weights every epoch)----
-    model = PerformerEncoderOnly_Module(**model_config)
-    model.to(device)
-
-    # ---- Main loop ----
-    summary = defaultdict(float)
-
-    for epoch_num, ckpt_path in checkpoints:
-        if is_master:
-            print(f"\n{'='*60}", flush=True)
-            print(f"Epoch {epoch_num}: {os.path.basename(ckpt_path)}", flush=True)
-            print(f"{'='*60}", flush=True)
-
-        ckpt = torch.load(ckpt_path, map_location=device)
-        model.load_state_dict(ckpt['model_state_dict'])
-        del ckpt
-        torch.cuda.empty_cache()
-        model.eval()
-
-        for ds_name, ds_dir, npz_files in datasets:
-            ds_losses = []
-            ds_cells_total = 0
-            t_ds_start = time.time()
-
-            for donor_idx, npz_file in enumerate(npz_files):
-                donor_id = npz_file.replace('.npz', '')
-                npz_path = os.path.join(ds_dir, npz_file)
-
-                data_npz = np.load(npz_path, allow_pickle=True)
-                csr = sparse.csr_matrix(
-                    (data_npz['data'], data_npz['indices'], data_npz['indptr']),
-                    shape=tuple(data_npz['shape'])
-                )
-                num_cells = csr.shape[0]
-
-                dataset_obj = DonorSCDataset(csr)
-                sampler = DistributedSampler(dataset_obj, num_replicas=world_size,
-                                             rank=rank, shuffle=False, drop_last=False)
-                loader = DataLoader(
-                    dataset_obj, batch_size=args.batch_size, sampler=sampler,
-                    num_workers=args.num_workers, pin_memory=True, drop_last=False,
-                )
-
-                donor_seed = args.seed + donor_idx if args.eval_mode == "masked" else None
-                local_loss_sum, local_batch_count = evaluate_donor_ddp(
-                    model, loader, device, args.eval_mode,
-                    args.mask_ratio, ptdtype, seed=donor_seed,
-                )
-
-                global_loss_sum = reduce_scalar(local_loss_sum, device)
-                global_batch_count = reduce_scalar(local_batch_count, device)
-
-                if global_batch_count > 0:
-                    mean_loss = global_loss_sum / global_batch_count
-                else:
-                    mean_loss = 0.0
-
-                if is_master:
-                    writer.writerow([epoch_num, ds_name, donor_id, num_cells,
-                                     f"{mean_loss:.6f}", args.eval_mode])
-                    csv_file.flush()
-                    ds_losses.append(mean_loss)
-                    ds_cells_total += num_cells
-
-                    if (donor_idx + 1) % 10 == 0 or (donor_idx + 1) == len(npz_files):
-                        print(f"  {ds_name}: {donor_idx+1}/{len(npz_files)} donors done", flush=True)
-
-                dist.barrier()
-
-            if is_master:
-                ds_mean = np.mean(ds_losses) if ds_losses else 0.0
-                print(f"  {ds_name}: mean_loss={ds_mean:.6f}, "
-                      f"total_cells={ds_cells_total}, time={time.time()-t_ds_start:.1f}s", flush=True)
-                summary[(epoch_num, ds_name)] = ds_mean
-
-        torch.cuda.empty_cache()
-        dist.barrier()
-
-    del model
-    torch.cuda.empty_cache()
-
-    # ---- Summary table ----
-    if is_master:
-        csv_file.close()
-        print(f"\nDetailed results have been saved to: {args.output_csv}")
-        print(f"\n{'='*60}")
-        print("summary table")
-        print(f"{'='*60}")
-        ds_names = sorted(set(ds for _, ds in summary.keys()))
-        header = f"{'Epoch':>8}" + "".join(f"{ds:>16}" for ds in ds_names)
-        print(header)
-        print("-" * len(header))
-        for epoch_num, _ in checkpoints:
-            row = f"{epoch_num:>8}"
-            for ds in ds_names:
-                row += f"{summary.get((epoch_num, ds), float('nan')):>16.6f}"
-            print(row)
-
-    dist.destroy_process_group()
-
-
-if __name__ == "__main__":
-    import traceback
-    try:
-        main()
-    except Exception:
-        traceback.print_exc()
-        sys.stdout.flush()
-        sys.stderr.flush()
-        sys.exit(1)

README.md

@@ -50,7 +50,7 @@ Geometric 2.6. See [environment.yml](BioMod/environment.yml) and
 [requirements.txt](BioMod/requirements.txt) for the complete dependency lists.
 Match the GPU packages to your system when adapting the environment.
 
-Pretraining and new-data reconstruction evaluation use CUDA/NCCL. The zero-shot
+Pretraining uses CUDA/NCCL. The zero-shot
 runner selects CUDA when available and otherwise uses CPU. Native Windows
 execution has not been validated for the supplied environment.
 
@@ -88,14 +88,13 @@ to end with the bundled checkpoint.
 | Drug-pair synergy classification | [Training script](BioMod/drug_combination/train_drug_combination_performer_module.py) | [Notebook](BioMod/toturials/drug_combination.ipynb) |
 | Gene perturbation response prediction | [Training script](BioMod/gene_perturbation/train_module.py) | [Notebook](BioMod/toturials/gene_perturbation.ipynb) |
 
-### Pretraining, evaluation, and analysis scripts
+### Pretraining and analysis scripts
 
 The repository also includes the following supporting workflows.
 
 | Workflow | Implementation | Documentation |
 | --- | --- | --- |
 | Pretraining and continued pretraining | [Pretraining script](BioMod/performer_module.py) | [Source guide](BioMod/README.md), [launcher reference](BioMod/command.txt) |
-| Expression reconstruction evaluation on new data | [Evaluation script](BioMod/eval_new_data/evaluate_new_data.py) | [Launcher reference](BioMod/eval_new_data/cmd.txtx) |
 | Module interpretability analysis | [Module geometry](BioMod/interpretability/module_geometry.py), [module importance](BioMod/interpretability/module_importance/module_importance.py), [text alignment](BioMod/interpretability/module_text_alignment/module_text_alignment.py) | [Overview](BioMod/interpretability/README.md), [attention](BioMod/interpretability/attention/README.md), [module importance guide](BioMod/interpretability/module_importance/README.md) |
 
 Tutorials and launcher references retain some historical checkpoint paths,
@@ -138,9 +137,6 @@ Updating environment variables alone does not replace hard-coded notebook paths.
 | Cell annotation / backbone interpretation | `--performer_checkpoint` | `--config_pkl` |
 | Continued pretraining | `--model_path` | `--model_config_path` or paired-file discovery |
 
-New-data reconstruction evaluation has a separate checkpoint-directory interface;
-refer to its script for the corresponding arguments.
-
 ### Data and evaluation
 
 Datasets are stored under `data/`. Input formats, metadata fields, preprocessing,
@@ -186,9 +182,6 @@ BioMod/                                                   # Project root: source
 │   ├── drug_response/                                    # Single-drug response prediction and training
 │   │   ├── cmd.txt
 │   │   └── train_drug_performer_module.py
-│   ├── eval_new_data/                                    # Reconstruction error and generalization on new data
-│   │   ├── cmd.txtx
-│   │   └── evaluate_new_data.py
 │   ├── gene_perturbation/                                # Perturbation response prediction: processing, training and inference
 │   │   ├── data_utils_local.py
 │   │   ├── gears_reduce_local.py