File size: 7,299 Bytes
be4e6f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from models.cifm import CIFM\n",
    "import scanpy as sc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CIFM(\n",
       "  (gene_encoder): MLPBiasFree(\n",
       "    (layers): ModuleList(\n",
       "      (0): Linear(in_features=18289, out_features=1024, bias=False)\n",
       "      (1-3): 3 x Linear(in_features=1024, out_features=1024, bias=False)\n",
       "    )\n",
       "    (layernorms): ModuleList(\n",
       "      (0-2): 3 x LayerNorm((1024,), eps=1e-05, elementwise_affine=False)\n",
       "    )\n",
       "    (activation): ReLU()\n",
       "  )\n",
       "  (model): VIEGNNModel(\n",
       "    (emb_in): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "    (convs): ModuleList(\n",
       "      (0-1): 2 x EGNNLayer(emb_dim=1024, aggr=sum)\n",
       "    )\n",
       "    (pred): MLPBiasFree(\n",
       "      (layers): ModuleList(\n",
       "        (0-3): 4 x Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      )\n",
       "      (layernorms): ModuleList(\n",
       "        (0-2): 3 x LayerNorm((1024,), eps=1e-05, elementwise_affine=False)\n",
       "      )\n",
       "      (activation): ReLU()\n",
       "    )\n",
       "  )\n",
       "  (mask_cell_decoder): VIEGNNModel(\n",
       "    (emb_in): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "    (convs): ModuleList(\n",
       "      (0-1): 2 x EGNNLayer(emb_dim=1024, aggr=sum)\n",
       "    )\n",
       "    (pred): MLPBiasFree(\n",
       "      (layers): ModuleList(\n",
       "        (0-3): 4 x Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      )\n",
       "      (layernorms): ModuleList(\n",
       "        (0-2): 3 x LayerNorm((1024,), eps=1e-05, elementwise_affine=False)\n",
       "      )\n",
       "      (activation): ReLU()\n",
       "    )\n",
       "  )\n",
       "  (mask_cell_expression): MLPBiasFree(\n",
       "    (layers): ModuleList(\n",
       "      (0-2): 3 x Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (3): Linear(in_features=1024, out_features=18289, bias=False)\n",
       "    )\n",
       "    (layernorms): ModuleList(\n",
       "      (0-2): 3 x LayerNorm((1024,), eps=1e-05, elementwise_affine=False)\n",
       "    )\n",
       "    (activation): ReLU()\n",
       "  )\n",
       "  (mask_cell_dropout): MLPBiasFree(\n",
       "    (layers): ModuleList(\n",
       "      (0-2): 3 x Linear(in_features=1024, out_features=1024, bias=False)\n",
       "      (3): Linear(in_features=1024, out_features=18289, bias=False)\n",
       "    )\n",
       "    (layernorms): ModuleList(\n",
       "      (0-2): 3 x LayerNorm((1024,), eps=1e-05, elementwise_affine=False)\n",
       "    )\n",
       "    (activation): ReLU()\n",
       "  )\n",
       "  (mask_embedding): Embedding(1, 1024)\n",
       "  (relu): ReLU()\n",
       "  (sigmoid): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args_model = torch.load('./model_files/args.pt')\n",
    "model = CIFM.from_pretrained('ynyou/CIFM', args=args_model)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 24844 × 18289\n",
       "    obs: 'in_tissue'\n",
       "    var: 'feature_types', 'genome', 'gene_names'\n",
       "    uns: 'log1p'\n",
       "    obsm: 'spatial'\n",
       "    layers: 'counts'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "channel2ensembl = torch.load('./model_files/channel2ensembl.pt')\n",
    "adata = sc.read_h5ad('./adata.h5ad')\n",
    "adata.layers['counts'] = adata.X.copy()\n",
    "sc.pp.normalize_total(adata)\n",
    "sc.pp.log1p(adata)\n",
    "adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "matching 18289 gene channels out of 18289 unmatched channels: []\n"
     ]
    }
   ],
   "source": [
    "model.channel_matching(adata, channel2ensembl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[-0.4132, -0.9847,  0.1647,  ..., -0.8351, -0.8177, -1.3235],\n",
       "         [ 0.8701,  0.0967, -0.3676,  ...,  0.2687, -1.4821,  0.1605],\n",
       "         [-0.5178, -0.4442, -0.0862,  ..., -0.7446, -0.5761, -0.5571],\n",
       "         ...,\n",
       "         [ 1.2264,  1.2326,  0.2791,  ...,  0.8018, -1.4069,  1.4567],\n",
       "         [ 0.6699, -0.6107,  0.2450,  ..., -0.1975, -0.6034, -0.6608],\n",
       "         [-1.9240, -1.8125, -0.0766,  ..., -0.2799, -0.0217, -2.2051]]),\n",
       " torch.Size([13898, 1024]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    embeddings = model.embed(adata)\n",
    "embeddings, embeddings.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[0.0000, 0.0000, 0.8603,  ..., 0.0000, 0.0000, 0.0000],\n",
       "         [0.0000, 0.0000, 0.6644,  ..., 0.0000, 0.0000, 0.0000],\n",
       "         [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "         ...,\n",
       "         [0.0000, 0.0000, 0.9809,  ..., 0.0000, 0.0000, 0.0000],\n",
       "         [0.6641, 0.0000, 0.6858,  ..., 0.0000, 0.0000, 0.0000],\n",
       "         [0.4999, 0.0000, 0.5311,  ..., 0.0000, 0.0000, 0.0000]]),\n",
       " torch.Size([10, 18289]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rand_loc = np.random.rand(10, 2)\n",
    "x_min, x_max = adata.obsm['spatial'][:, 0].min(), adata.obsm['spatial'][:, 0].max()\n",
    "y_min, y_max = adata.obsm['spatial'][:, 1].min(), adata.obsm['spatial'][:, 1].max()\n",
    "rand_loc[:, 0] = rand_loc[:, 0] * (x_max - x_min) + x_min\n",
    "rand_loc[:, 1] = rand_loc[:, 1] * (y_max - y_min) + y_min\n",
    "\n",
    "with torch.no_grad():\n",
    "    expressions = model.predict_cells_at_locations(adata, rand_loc)\n",
    "expressions, expressions.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}