From 337b571f9d499e06554798302854828cf8e18536 Mon Sep 17 00:00:00 2001
From: Emanuele Bezzi <ebezzi@chanzuckerberg.com>
Date: Fri, 1 Mar 2024 13:58:33 -0800
Subject: [PATCH] Checkpoint

---
 .../src/cellxgene_census/_get_anndata.py           | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/api/python/cellxgene_census/src/cellxgene_census/_get_anndata.py b/api/python/cellxgene_census/src/cellxgene_census/_get_anndata.py
index d8a2d7515..e3056d0e8 100644
--- a/api/python/cellxgene_census/src/cellxgene_census/_get_anndata.py
+++ b/api/python/cellxgene_census/src/cellxgene_census/_get_anndata.py
@@ -12,10 +12,12 @@
 import tiledbsoma as soma
 from somacore.options import SparseDFCoord
 
-from util import _extract_census_version
+from ._util import _extract_census_version
 
 from ._experiment import _get_experiment
 
+CENSUS_EMBEDDINGS_LOCATION_BASE_URI = "s3://cellxgene-contrib-public/contrib/cell-census/soma"
+
 
 def get_anndata(
     census: soma.Collection,
@@ -106,7 +108,15 @@ def get_anndata(
 
         # If add_obs_embeddings or add_var_embeddings are defined, inject them in the appropriate slot
         if add_obs_embeddings or add_var_embeddings:
+            obs_soma_joinids = query.obs_joinids()
             from cellxgene_census.experimental import get_embedding, get_embedding_metadata_by_name
             census_version = _extract_census_version(census)
-            get_embedding_metadata_by_name()
+            for emb in add_obs_embeddings:
+                emb_metadata = get_embedding_metadata_by_name(emb, organism, census_version, "obs_embedding")
+                uri = f"{CENSUS_EMBEDDINGS_LOCATION_BASE_URI}/{census_version}/{emb_metadata['id']}"
+                embedding = get_embedding(census_version, uri, obs_soma_joinids)
+                adata.obsm[emb] = embedding
+
+        return adata
+