Skip to content

Commit

Permalink
[python] Update add_new_collection method to support Scene collec…
Browse files Browse the repository at this point in the history
…tions (#3491) (#3499)

* Add kwargs to `add_new_collection` for adding `Scene` objects

* Use `add_new_collection` for creating scenes in tests

Co-authored-by: Julia Dark <[email protected]>
  • Loading branch information
github-actions[bot] and jp-dark authored Dec 20, 2024
1 parent 4ad92ba commit ac4e912
Show file tree
Hide file tree
Showing 3 changed files with 152 additions and 154 deletions.
7 changes: 7 additions & 0 deletions apis/python/src/tiledbsoma/_collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,7 @@ def add_new_collection(
*,
uri: str | None = ...,
platform_config: options.PlatformConfig | None = ...,
**kwargs: Any,
) -> "Collection[AnySOMAObject]": ...

@overload
Expand All @@ -158,6 +159,7 @@ def add_new_collection(
*,
uri: str | None = ...,
platform_config: options.PlatformConfig | None = ...,
**kwargs: Any,
) -> _Coll: ...

def add_new_collection(
Expand All @@ -167,6 +169,7 @@ def add_new_collection(
*,
uri: str | None = None,
platform_config: options.PlatformConfig | None = None,
**kwargs: Any,
) -> AnyTileDBCollection:
"""Adds a new sub-collection to this collection.
Expand All @@ -190,6 +193,9 @@ def add_new_collection(
Platform configuration options to use when
creating this sub-collection. This is passed directly to
``[CurrentCollectionType].create()``.
kwargs:
Other keyword arguments to pass to the ``create`` method of the
new sub-collection type.
Examples:
>>> with tiledbsoma.Collection.create("/tmp/parent") as parent_collection:
Expand Down Expand Up @@ -217,6 +223,7 @@ def add_new_collection(
platform_config=platform_config,
context=self.context,
tiledb_timestamp=self.tiledb_timestamp_ms,
**kwargs,
),
uri,
)
Expand Down
245 changes: 120 additions & 125 deletions apis/python/tests/test_basic_spatialdata_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,138 +30,133 @@ def experiment_with_single_scene(tmp_path_factory, sample_2d_data) -> soma.Exper
with soma.Experiment.create(uri) as exp:
assert exp.uri == uri
# Create spatial folder.
exp.add_new_collection("spatial")
with exp.add_new_collection("spatial") as spatial:

# Create scene 1.
scene1_uri = urljoin(exp.spatial.uri, "scene1")
exp.spatial["scene1"] = soma.Scene.create(scene1_uri)
scene1 = exp.spatial["scene1"]
assert scene1_uri == scene1.uri
scene1.coordinate_space = soma.CoordinateSpace.from_axis_names(
["x_scene1", "y_scene1"]
)
scene1.add_new_collection("obsl")
scene1.add_new_collection("varl")
scene1.varl.add_new_collection("RNA")
scene1.add_new_collection("img")
# Create scene 1.
with spatial.add_new_collection(
"scene1", soma.Scene, coordinate_space=("x_scene1", "y_scene1")
) as scene1:
scene1.add_new_collection("obsl")
scene1.add_new_collection("varl")
scene1.varl.add_new_collection("RNA")
scene1.add_new_collection("img")

# Add point cloud with shape to scene 1 obsl.
points1 = scene1.add_new_point_cloud_dataframe(
"points1",
"obsl",
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), 2.0
),
schema=pa.schema([("x", pa.float64()), ("y", pa.float64())]),
domain=[[0, 1], [0, 1]],
)
points1.write(
pa.Table.from_pydict(
{
"soma_joinid": np.arange(4),
"x": np.array([0, 0, 0.5, 0.5]),
"y": np.array([0, 0.5, 0, 0.5]),
}
)
)
points1.metadata["soma_geometry"] = 1.0
points1.metadata["soma_geometry_type"] = "radius"
# Add point cloud with shape to scene 1 obsl.
points1 = scene1.add_new_point_cloud_dataframe(
"points1",
"obsl",
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), 2.0
),
schema=pa.schema([("x", pa.float64()), ("y", pa.float64())]),
domain=[[0, 1], [0, 1]],
)
points1.write(
pa.Table.from_pydict(
{
"soma_joinid": np.arange(4),
"x": np.array([0, 0, 0.5, 0.5]),
"y": np.array([0, 0.5, 0, 0.5]),
}
)
)
points1.metadata["soma_geometry"] = 1.0
points1.metadata["soma_geometry_type"] = "radius"

# Add point cloud wihtout shape to scene 1 obsl
points3 = scene1.add_new_point_cloud_dataframe(
"points3",
"obsl",
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), 4.0
),
schema=pa.schema([("x", pa.float64()), ("y", pa.float64())]),
domain=[[-1, 0], [-1, 0]],
)
points3.write(
pa.Table.from_pydict(
{
"soma_joinid": np.arange(4),
"x": np.array([0, 0, -0.5, -0.5]),
"y": np.array([0, -0.5, 0, -0.5]),
}
)
)
# Add point cloud wihtout shape to scene 1 obsl
points3 = scene1.add_new_point_cloud_dataframe(
"points3",
"obsl",
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), 4.0
),
schema=pa.schema([("x", pa.float64()), ("y", pa.float64())]),
domain=[[-1, 0], [-1, 0]],
)
points3.write(
pa.Table.from_pydict(
{
"soma_joinid": np.arange(4),
"x": np.array([0, 0, -0.5, -0.5]),
"y": np.array([0, -0.5, 0, -0.5]),
}
)
)

# Add point cloud without shape to scene 1 varl.
points2 = scene1.add_new_point_cloud_dataframe(
"points2",
["varl", "RNA"],
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), -1.0
),
schema=pa.schema([("x", pa.float64()), ("y", pa.float64())]),
domain=[[-1, 0], [-1, 0]],
)
points2.write(
pa.Table.from_pydict(
{
"soma_joinid": np.arange(4),
"x": np.array([0, 0, -0.5, -0.5]),
"y": np.array([0, -0.5, 0, -0.5]),
}
)
)
# Add point cloud without shape to scene 1 varl.
points2 = scene1.add_new_point_cloud_dataframe(
"points2",
["varl", "RNA"],
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), -1.0
),
schema=pa.schema([("x", pa.float64()), ("y", pa.float64())]),
domain=[[-1, 0], [-1, 0]],
)
points2.write(
pa.Table.from_pydict(
{
"soma_joinid": np.arange(4),
"x": np.array([0, 0, -0.5, -0.5]),
"y": np.array([0, -0.5, 0, -0.5]),
}
)
)

# Add point cloud with shape to scene 1 varl.
points4 = scene1.add_new_point_cloud_dataframe(
"points4",
["varl", "RNA"],
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), 0.25
),
schema=pa.schema([("x", pa.float64()), ("y", pa.float64())]),
domain=[[0, 1], [0, 1]],
)
points4.write(
pa.Table.from_pydict(
{
"soma_joinid": np.arange(4),
"x": np.array([0, 0, 0.5, 0.5]),
"y": np.array([0, 0.5, 0, 0.5]),
}
)
)
points4.metadata["soma_geometry"] = 2.0
points4.metadata["soma_geometry_type"] = "radius"
# Add point cloud with shape to scene 1 varl.
points4 = scene1.add_new_point_cloud_dataframe(
"points4",
["varl", "RNA"],
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), 0.25
),
schema=pa.schema([("x", pa.float64()), ("y", pa.float64())]),
domain=[[0, 1], [0, 1]],
)
points4.write(
pa.Table.from_pydict(
{
"soma_joinid": np.arange(4),
"x": np.array([0, 0, 0.5, 0.5]),
"y": np.array([0, 0.5, 0, 0.5]),
}
)
)
points4.metadata["soma_geometry"] = 2.0
points4.metadata["soma_geometry_type"] = "radius"

# Add multiscale image with a single image.
with scene1.add_new_multiscale_image(
"image1",
"img",
type=pa.uint8(),
level_shape=(3, 64, 64),
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), 0.5
),
) as image1:
coords = (slice(None), slice(None), slice(None))
l0 = image1["level0"]
l0.write(coords, pa.Tensor.from_numpy(sample_2d_data[0]))
# Add multiscale image with a single image.
with scene1.add_new_multiscale_image(
"image1",
"img",
type=pa.uint8(),
level_shape=(3, 64, 64),
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), 0.5
),
) as image1:
coords = (slice(None), slice(None), slice(None))
l0 = image1["level0"]
l0.write(coords, pa.Tensor.from_numpy(sample_2d_data[0]))

# Add multiscale image with multiple resolutions.
with scene1.add_new_multiscale_image(
"image2",
"img",
type=pa.uint8(),
level_key="fullres",
level_shape=(3, 32, 32),
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), 0.5
),
) as image2:
coords = (slice(None), slice(None), slice(None))
fullres = image2["fullres"]
fullres.write(coords, pa.Tensor.from_numpy(sample_2d_data[1]))
hires = image2.add_new_level("hires", shape=(3, 16, 16))
hires.write(coords, pa.Tensor.from_numpy(sample_2d_data[2]))
lowres = image2.add_new_level("lowres", shape=(3, 8, 8))
lowres.write(coords, pa.Tensor.from_numpy(sample_2d_data[3]))
scene1.close()
# Add multiscale image with multiple resolutions.
with scene1.add_new_multiscale_image(
"image2",
"img",
type=pa.uint8(),
level_key="fullres",
level_shape=(3, 32, 32),
transform=soma.UniformScaleTransform(
("x_scene1", "y_scene1"), ("x", "y"), 0.5
),
) as image2:
coords = (slice(None), slice(None), slice(None))
fullres = image2["fullres"]
fullres.write(coords, pa.Tensor.from_numpy(sample_2d_data[1]))
hires = image2.add_new_level("hires", shape=(3, 16, 16))
hires.write(coords, pa.Tensor.from_numpy(sample_2d_data[2]))
lowres = image2.add_new_level("lowres", shape=(3, 8, 8))
lowres.write(coords, pa.Tensor.from_numpy(sample_2d_data[3]))

return soma.Experiment.open(uri, mode="r")

Expand Down
54 changes: 25 additions & 29 deletions apis/python/tests/test_experiment_query_spatial.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import itertools
from typing import Dict, List, Tuple, Union
from urllib.parse import urljoin

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -105,36 +104,33 @@ def add_scene(
circles: Dict[Tuple[Union[str, Tuple[str, ...]], str], Dict[str, np.ndarray]],
images: Dict[str, Tuple[Tuple[int, ...]]],
) -> None:
uri = urljoin(coll.uri, key)
coll[key] = soma.Scene.create(uri, coordinate_space=("x", "y"))
scene = coll[key]
scene.add_new_collection("obsl")
varl = scene.add_new_collection("varl")
varl.add_new_collection("RNA")
varl.add_new_collection("other")
scene.add_new_collection("img")

for (subcoll, key), data in points.items():
add_point_cloud_dataframe(
scene,
subcoll if isinstance(subcoll, str) else list(subcoll),
key,
data,
circles=False,
)
with coll.add_new_collection(key, soma.Scene, coordinate_space=("x", "y")) as scene:
scene.add_new_collection("obsl")
varl = scene.add_new_collection("varl")
varl.add_new_collection("RNA")
varl.add_new_collection("other")
scene.add_new_collection("img")

for (subcoll, key), data in points.items():
add_point_cloud_dataframe(
scene,
subcoll if isinstance(subcoll, str) else list(subcoll),
key,
data,
circles=False,
)

for (subcoll, key), data in circles.items():
add_point_cloud_dataframe(
scene,
subcoll if isinstance(subcoll, str) else list(subcoll),
key,
data,
circles=True,
)
for (subcoll, key), data in circles.items():
add_point_cloud_dataframe(
scene,
subcoll if isinstance(subcoll, str) else list(subcoll),
key,
data,
circles=True,
)

for key, shapes in images.items():
add_multiscale_image(scene, key, shapes)
scene.close()
for key, shapes in images.items():
add_multiscale_image(scene, key, shapes)


@pytest.fixture(scope="module")
Expand Down

0 comments on commit ac4e912

Please sign in to comment.