FilesExpand file tree

main

/

run_cellpin.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

150 lines (121 loc) · 5.95 KB

main

/

run_cellpin.py

Copy path

Top

File metadata and controls

• Code
• Blame

150 lines (121 loc) · 5.95 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

#!/usr/bin/env python3

"""Simple CellPin training and spatial imputation.

Returns an anndata object containing both the original spatial data and the

imputed values + CellPin embeddings.

Usage

-----

uv run python run_cellpin.py \

--adata_path /path/to/scRNA.h5ad \

--spatial_path /path/to/spatial.h5ad \

--output_dir ./experiments/cellpin_run

"""

from __future__ import annotations

import argparse

from pathlib import Path

import anndata as ad

import yaml

from torch.utils.data import DataLoader

import cellpin

import cellpin.pp

from cellpin.models import CellPin

# Model config (YAML)

CONFIG_PATH_DEFAULT = Path(__file__).resolve().parents[1] / "configs" / "cellpin_config.yaml"

def load_model_config(config_path: Path) -> dict:

"""Load a YAML model configuration file."""

with config_path.open("r", encoding="utf-8") as f:

return yaml.safe_load(f)

def run(args: argparse.Namespace) -> None:

"""Run two-stage training and spatial imputation."""

output_dir = Path(args.output_dir)

output_dir.mkdir(parents=True, exist_ok=True)

config_path = Path(args.config).expanduser() if args.config else CONFIG_PATH_DEFAULT

if not config_path.is_absolute():

config_path = (Path.cwd() / config_path).resolve()

if not config_path.exists():

raise FileNotFoundError(f"Config file not found: {config_path}")

model_config = load_model_config(config_path)

print(f"Using model config: {config_path}")

layer = args.layer or None

# ── Load data ──────────────────────────────────────────────────────────

print(f"Loading scRNA : {args.adata_path}")

adata_sc = ad.read_h5ad(args.adata_path)

print(f"Loading spatial: {args.spatial_path}")

adata_sp = ad.read_h5ad(args.spatial_path)

print(f"sc: {adata_sc.shape} | spatial: {adata_sp.shape}")

# ── Setup datasets for training ────────────────────────────────────────

sc_dataset, st_dataset = cellpin.pp.setup_data(

sc_adata=adata_sc,

st_adata=adata_sp,

gene_symbols=None,

layer=layer,

)

# ── Create model ───────────────────────────────────────────────────────

model = CellPin(sc_dataset=sc_dataset, config=model_config)

shared = {

"precision": args.precision,

"accelerator": "auto",

"devices": args.devices,

"strategy": args.strategy,

"seed": args.seed,

"batch_size": args.batch_size,

"num_workers": args.num_workers,

"gradient_clip_val": 0.5,

"early_stopping_patience": args.early_stopping_patience,

}

# ── Stage 1: pretrain ──────────────────────────────────────────────────

print("\nStage 1: pretraining...")

model.pretrain_model(

dataset=sc_dataset,

max_epochs=args.pretrain_epochs,

output_dir=str(output_dir / "pretrain"),

**shared,

)

# ── Stage 2: main training ─────────────────────────────────────────────

print("\nStage 2: main training...")

model.train_model(

dataset=sc_dataset,

freeze_pretrained=args.freeze_pretrained,

decoder_warm_unfreeze_epoch=args.decoder_warm_unfreeze_epoch,

max_epochs=args.train_epochs,

output_dir=str(output_dir / "train"),

**shared,

)

# ── Inference on spatial data ──────────────────────────────────────────

print("\nRunning imputation on spatial data...")

dl = DataLoader(

st_dataset,

batch_size=args.batch_size,

shuffle=False,

num_workers=args.num_workers,

)

adata_imputed = model.impute(dl, obs_adata=adata_sp, return_norm=args.return_norm, nb_count_samples=20, return_int=args.return_int)

out_path = output_dir / "cellpin_imputed.h5ad"

adata_imputed.write_h5ad(str(out_path))

print(f"Imputed AnnData saved to {out_path}")

def parse_args() -> argparse.Namespace:

"""Parse CLI arguments for the training script."""

p = argparse.ArgumentParser(

description="CellPin: two-stage training + spatial imputation",

formatter_class=argparse.ArgumentDefaultsHelpFormatter,

)

p.add_argument("--adata_path", required=True, help="scRNA AnnData (.h5ad)")

p.add_argument("--spatial_path", required=True, help="Spatial AnnData (.h5ad)")

p.add_argument("--output_dir", default="./experiments/cellpin_run")

p.add_argument("--config", default=None, help=f"Path to model YAML config (default: {CONFIG_PATH_DEFAULT})")

p.add_argument("--layer", default="counts", help="Expression layer (empty → .X)")

p.add_argument("--pretrain_epochs", type=int, default=50)

p.add_argument("--train_epochs", type=int, default=60)

p.add_argument("--batch_size", type=int, default=256)

p.add_argument("--num_workers", type=int, default=8)

p.add_argument("--early_stopping_patience", type=int, default=12)

p.add_argument("--freeze_pretrained", action=argparse.BooleanOptionalAction, default=False)

p.add_argument("--decoder_warm_unfreeze_epoch", type=int, default=25)

p.add_argument("--return_int", action=argparse.BooleanOptionalAction, default=True)

p.add_argument("--return_norm", action=argparse.BooleanOptionalAction, default=True)

p.add_argument("--precision", default="16-mixed")

p.add_argument("--devices", type=int, nargs="+", default=[0])

p.add_argument("--strategy", default="auto")

p.add_argument("--seed", type=int, default=42)

return p.parse_args()

if __name__ == "__main__":

run(parse_args())