update

adata.h5ad

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1126447b46abf9c31e77a009423473988738f0295c23a702b0416dd7f56e208d
-size 32068372
+oid sha256:f4f8b3caccbb84f31fa795ad012d22c28068d2fc8a8c1a28d7b034483a168e08
+size 90959812

models/__init__.py

@@ -1,8 +0,0 @@
-# from models.schnet import SchNetModel
-# from models.dimenet import DimeNetPPModel
-# from models.spherenet import SphereNetModel
-# from models.egnn import EGNNModel
-# from models.gvpgnn import GVPGNNModel
-# from models.tfn import TFNModel
-# from models.mace import MACEModel
-from models.egnn_void_invariant import VIEGNNModel

models/cifm.py

@@ -2,10 +2,9 @@ import torch
 import torch.nn as nn
 from torch_geometric.nn import radius_graph
 import scanpy as sc
-from main import Model
 from huggingface_hub import PyTorchModelHubMixin
 from models.mlp_and_gnn import MLPBiasFree
-from models import VIEGNNModel
+from models.egnn_void_invariant import VIEGNNModel
 
 
 class CIFM(
@@ -59,10 +58,9 @@ class CIFM(
                         embs_in.append(self.gene_encoder.layers[0].weight.data[:, idx_source])
                         embs_out1.append(self.mask_cell_expression.layers[-1].weight.data[idx_source])
                         embs_out2.append(self.mask_cell_dropout.layers[-1].weight.data[idx_source])
-                    else:
-                        unmatched_channels.append(ensembl)
             
             if len(embs_in) == 0:
+                unmatched_channels += ensembls
                 continue
             
             embs_in = torch.stack(embs_in).mean(dim=0)
@@ -98,8 +96,11 @@ class CIFM(
 
         expressions_dec = self.relu(self.mask_cell_expression(embeddings_dec))
         dropouts_dec = self.sigmoid(self.mask_cell_dropout(embeddings_dec))
+
+        # import pdb ; pdb.set_trace()
+
         expressions_dec[dropouts_dec<=0.5] = 0
-        return expressions_dec[mapping]
+        return expressions_dec
     
     def embed(self, adata):
         device = next(self.parameters()).device
@@ -119,7 +120,7 @@ class CIFM(
         expressions = torch.tensor(adata.X.toarray(), dtype=torch.float32).to(device)
         expressions = torch.cat([expressions, torch.zeros(locations.shape[0], expressions.shape[1])], dim=0)
         coordinates = torch.tensor(adata.obsm['spatial'], dtype=torch.float32)
-        coordinates = torch.cat([coordinates, torch.zeros(locations.shape[0], 2)], dim=1)
+        coordinates = torch.cat([coordinates, locations], dim=0)
         coordinates = torch.cat([coordinates, torch.zeros(coordinates.shape[0], 1)], dim=1).to(device)
         edge_index = radius_graph(coordinates, r=self.radius_spatial_graph, max_num_neighbors=10000, loop=True)
         idx_cells_to_predict = torch.arange(expressions.shape[0]-locations.shape[0], expressions.shape[0]).to(device)

test.ipynb

@@ -0,0 +1,234 @@
+{
+ "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
+}