Partial rever

Author: AstroPatty

src/opencosmo/collection/structure/collection.py

@@ -4,12 +4,12 @@
 from warnings import warn
 
 import astropy  # type: ignore
-import numpy as np
 
 import opencosmo as oc
 from opencosmo.collection.structure import io as sio
 from opencosmo.collection.structure import visit
 from opencosmo.dataset.column import DerivedColumn
+from opencosmo.index import DataIndex
 from opencosmo.io import io
 from opencosmo.io.schemas import StructCollectionSchema
 from opencosmo.parameters import HaccSimulationParameters
@@ -166,8 +166,8 @@ def __getitem__(self, key: str) -> oc.Dataset | oc.StructureCollection:
         elif key == self.__header.file.data_type:
             return self.__source
 
-        start, size = self.__links[key].make_indices(self.__index)
-        return self.__datasets[key].take_rows(start, size)
+        index = self.__links[key].make_index(self.__index)
+        return self.__datasets[key].with_index(index)
 
     def __enter__(self):
         return self
@@ -609,8 +609,8 @@ def with_new_columns(self, dataset: str, **new_columns: DerivedColumn):
             self.__hide_source,
         )
 
-    def take_rows(self, indices: np.ndarray, size: Optional[np.ndarray] = None):
-        new_source = self.__source.take_rows(indices, size)
+    def with_index(self, index: DataIndex):
+        new_source = self.__source.with_index(index)
         return StructureCollection(
             new_source, self.__header, self.__datasets, self.__links, self.__hide_source
         )
@@ -650,8 +650,8 @@ def objects(
         for i, row in enumerate(self.__source.rows()):
             index = self.__source.index[i]
             output = {
-                key: self.__datasets[key].take_rows(
-                    *self.__links[key].make_indices(index)
+                key: self.__datasets[key].with_index(
+                    self.__links[key].make_index(index)
                 )
                 for key in data_types
             }

src/opencosmo/collection/structure/handler.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 import h5py
 import numpy as np
 
-from opencosmo.index import DataIndex
+from opencosmo.index import ChunkedIndex, DataIndex, SimpleIndex
 from opencosmo.io import schemas as ios
 
 if TYPE_CHECKING:
@@ -35,17 +35,22 @@ def has_linked_data(self, index: DataIndex) -> np.ndarray:
             rows = index.get_data(self.link)
             return rows != -1
 
-    def make_indices(self, index: DataIndex) -> tuple[np.ndarray, Optional[np.ndarray]]:
+    def make_index(self, index: DataIndex) -> DataIndex:
         if isinstance(self.link, tuple):
             start = index.get_data(self.link[0])
             size = index.get_data(self.link[1])
             valid_rows = size > 0
-            starts = start[valid_rows]
-            sizes = size[valid_rows]
-            return starts, sizes
+            start = start[valid_rows]
+            size = size[valid_rows]
+            if not start.size:
+                return SimpleIndex(np.array([], dtype=int))
+            else:
+                return ChunkedIndex(start, size)
         else:
             indices_into_data = index.get_data(self.link)
-            return indices_into_data[indices_into_data >= 0], None
+            indices_into_data = indices_into_data[indices_into_data >= 0]
+
+            return SimpleIndex(indices_into_data)
 
     def make_schema(self, name: str, index: DataIndex) -> ios.LinkSchema:
         if isinstance(self.link, h5py.Dataset):

src/opencosmo/collection/structure/io.py

@@ -90,12 +90,12 @@ def get_linked_datasets(
 
 def make_index_with_linked_data(
     index: DataIndex, links: dict[str, LinkedDatasetHandler]
-) -> np.ndarray:
+):
     mask = np.ones(len(index), dtype=bool)
     for link in links.values():
         mask &= link.has_linked_data(index)
 
-    return np.where(mask)[0]
+    return index.mask(mask)
 
 
 def build_structure_collection(targets: list[io.io.OpenTarget], ignore_empty: bool):
@@ -139,7 +139,7 @@ def build_structure_collection(targets: list[io.io.OpenTarget], ignore_empty: bo
         source_dataset = io.io.open_single_dataset(link_sources["galaxy_properties"][0])
         if ignore_empty:
             new_index = make_index_with_linked_data(source_dataset.index, handlers)
-            source_dataset = source_dataset.take_rows(new_index)
+            source_dataset = source_dataset.with_index(new_index)
         collection = sc.StructureCollection(
             source_dataset,
             source_dataset.header,
@@ -161,7 +161,7 @@ def build_structure_collection(targets: list[io.io.OpenTarget], ignore_empty: bo
 
         if ignore_empty:
             new_index = make_index_with_linked_data(source_dataset.index, handlers)
-            source_dataset = source_dataset.take_rows(new_index)
+            source_dataset = source_dataset.with_index(new_index)
 
         return sc.StructureCollection(
             source_dataset,

src/opencosmo/dataset/dataset.py

@@ -12,7 +12,7 @@
 from opencosmo.dataset.state import DatasetState
 from opencosmo.dataset.visit import visit_dataset
 from opencosmo.header import OpenCosmoHeader
-from opencosmo.index import ChunkedIndex, DataIndex, SimpleIndex
+from opencosmo.index import ChunkedIndex, DataIndex
 from opencosmo.io.schemas import DatasetSchema
 from opencosmo.parameters import HaccSimulationParameters
 from opencosmo.spatial import check
@@ -569,44 +569,7 @@ def take_range(self, start: int, end: int) -> Table:
             self.__tree,
         )
 
-    def take_rows(self, indices: np.ndarray, size: Optional[np.ndarray] = None):
-        """
-        Create a new dataset with the rows specified by the arguments. This method allows for
-        two different types of input
-
-        1. A single array, representing the indices of the rows to take
-        2. Two arrays. The first representing the start of a set of row chunks, and the second representing
-           their size
-
-        Parameters
-        ----------
-        indices: np.ndarray[int]
-            The rows to take, or the beginnings of the row chunks
-
-        size: np.ndarray[int], optional
-            The size of the row chunks
-
-        Returns
-        -------
-        dataset : opencosmo.Dataset
-            The new dataset with only the selected rows.
-
-        """
-        if np.any(indices < 0):
-            raise ValueError("Row indices cannot be negative!")
-
-        size_ = 0 if size is None else size
-        if np.any(indices + size_) > len(self):
-            raise ValueError(
-                "Row indices cannot be larger than the length of this dataset!"
-            )
-
-        index: DataIndex
-        if size is None:
-            index = SimpleIndex(indices)
-        else:
-            index = ChunkedIndex(indices, size)
-
+    def with_index(self, index: DataIndex):
         new_state = self.__state.with_index(index)
         return Dataset(self.__handler, self.__header, new_state, self.__tree)
 

test/test_collection.py

@@ -291,21 +291,6 @@ def test_data_linking(halo_paths):
     assert n_profiles > 0
 
 
-def test_data_gets_all_particles(halo_paths):
-    collection = oc.open(*halo_paths)
-    collection = collection.filter(oc.col("sod_halo_mass") > 10**14).take(
-        10, at="random"
-    )
-    for halo in collection.halos():
-        for name, particle_species in halo.items():
-            if "particle" not in name:
-                continue
-            halo_tag = halo["halo_properties"]["fof_halo_tag"]
-            tag_filter = oc.col("fof_halo_tag") == halo_tag
-            ds = collection[name].filter(tag_filter)
-            assert len(ds) == len(particle_species)
-
-
 def test_data_linking_bound(halo_paths):
     collection = oc.open(*halo_paths)
     p1 = tuple(random.uniform(10, 20) for _ in range(3))