Rewrite Open Logic

.pre-commit-config.yaml

@@ -5,7 +5,7 @@ repos:
     rev: 'v1.15.0'  # Use the sha / tag you want to point at
     hooks:
     -   id: mypy
-        additional_dependencies: ["numpy"]
+        additional_dependencies: ["numpy", "types-deprecated"]
 -   repo: https://github.com/astral-sh/ruff-pre-commit
     rev: 'v0.11.9'
     hooks:

README.md

@@ -124,7 +124,7 @@ We could then plot the data, or perform further transformations. This is cool on
 ```python
 import opencosmo as oc
 
-data = oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5")
+data = oc.open("haloproperties.hdf5", "haloparticles.hdf5")
 ```
 This will return a data *collection* that will allow you to query and transform the data as before, but will associate the halos with their particles. 
 

docs/source/analysis.rst

@@ -17,7 +17,7 @@ Here is an example for how to use `create_yt_dataset` to load a selection of dat
     from opencosmo.analysis import create_yt_dataset, ParticleProjectionPlot
 
     # select a random halo
-    with oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5").take(1, at="random") as data:
+    with oc.open("haloproperties.hdf5", "haloparticles.hdf5").take(1, at="random") as data:
         # get yt data container
         yt_ds = create_yt_dataset(data.objects())
 
@@ -82,7 +82,7 @@ We will now edit the code-block from before to compute X-ray luminosities:
     }
 
     # select a random halo
-    with oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5").take(1, at="random") as data:
+    with oc.open("haloproperties.hdf5", "haloparticles.hdf5").take(1, at="random") as data:
         # get yt data container
         ds_yt, source_model = create_yt_dataset(data.objects(), 
             compute_xray_fields = True, return_source_model = True)
@@ -121,7 +121,7 @@ This function essentially uses :func:`halo_projection_array` with pre-filled set
     import matplotlib.pyplot as plt
 
     # load one halo at random
-    with oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5").take(1, at="random") as data:
+    with oc.open("haloproperties.hdf5", "haloparticles.hdf5").take(1, at="random") as data:
         halo = next(data.halos())
         halo_id = halo['halo_properties']['unique_tag']
 
@@ -162,7 +162,7 @@ The outputted figure is an array of images, with the shape matching that of the
     import numpy as np
 
     # load 16 halos at random
-    with oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5").take(16, at="random") as data:
+    with oc.open("haloproperties.hdf5", "haloparticles.hdf5").take(16, at="random") as data:
 
         halo_ids = [halo['halo_properties']['unique_tag'] for halo in data.halos()]
 
@@ -192,7 +192,7 @@ One can also define a dictionary of plotting parameters to plot different fields
     import matplotlib.pyplot as plt
     import numpy as np
 
-    with oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5").take(2, at="random") as data:
+    with oc.open("haloproperties.hdf5", "haloparticles.hdf5").take(2, at="random") as data:
         halo_ids = [halo['halo_properties']['unique_tag'] for halo in data.halos()]
 
         # We are going to make a 2x3 panel figure, where each row is a different halo, and

docs/source/collections.rst

@@ -71,7 +71,7 @@ However the real power of working with a :py:class:`StructureCollection` is the
 .. code-block:: python
 
    import opencosmo as oc
-   data = oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5")
+   data = oc.open("haloproperties.hdf5", "haloparticles.hdf5")
    for halo in data.halos():
       print(halo)
 
@@ -95,7 +95,7 @@ It is also possible for structure collections to contain other structure collect
 .. code-block:: python
 
    import opencosmo as oc
-   ds = oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5", "galaxyproperties.hdf5", "galaxyparticles.hdf5")
+   ds = oc.open("haloproperties.hdf5", "haloparticles.hdf5", "galaxyproperties.hdf5", "galaxyparticles.hdf5")
    for structure in ds.halos():
         gals_ds = structure["galaxies"]
         for galaxy in gals_ds.galaxies():

docs/source/cols.rst

@@ -67,7 +67,7 @@ Queries can generally be performed as usual on collections. In a :py:class:`open
 .. code-block:: python
 
    import opencosmo as oc
-   ds = oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5")
+   ds = oc.open("haloproperties.hdf5", "haloparticles.hdf5")
    ds = ds.filter(oc.col("fof_halo_mass") > 1e14)
 
 The resultant :py:class:`opencosmo.StructureCollection` will contain only halos with a mass greater than 10^14, along with all their associated particles.
@@ -134,7 +134,7 @@ Calls to :py:meth:`opencosmo.StructureCollection.with_new_columns` must explicit
 .. code-block:: python
 
    import opencosmo as oc
-   ds = oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5")
+   ds = oc.open("haloproperties.hdf5", "haloparticles.hdf5")
    fof_halo_speed_sqrd = oc.col("fof_halo_com_vx") + oc.col("fof_halo_com_vy") ** 2 + oc.col("fof_halo_com_vz") ** 2
    fof_halo_ke = 0.5 * oc.col("fof_halo_mass") * fof_halo_speed_sqrd
    fof_halo_p = oc.col("fof_halo_mass") * fof_halo_speed_sqrd ** 0.5

docs/source/first_steps.rst

@@ -84,7 +84,7 @@ This is cool on its own, but the real power of :code:`opencosmo` comes from its
 .. code-block:: python
 
    import opencosmo as oc
-   data = oc.open_linked_files("haloproperties.hdf5", "haloparticles.hdf5")
+   data = oc.open("haloproperties.hdf5", "haloparticles.hdf5")
 
 This will return a data *collection* that will allow you to query and transform the data as before, but will associate the halos with their particles. 
 

docs/source/io.rst

@@ -5,7 +5,7 @@ OpenCosmo defines a data format for storing simulation data in hdf5 files. A dat
 Options for Reading Data
 ------------------------
 
-Any single opencosmo file can be open with :py:meth:`opencosmo.open`. This function will parse the structure of the file and return the appropriate object, such as an :py:class`opencosmo.Dataset` or :py:class:`opencosmo.StructureCollection`
+Any single opencosmo file can be open with :py:meth:`opencosmo.open`. This function will parse the structure of the file and return the appropriate object, such as a :py:class:`opencosmo.Dataset` or :py:class:`opencosmo.StructureCollection`
 
 .. code-block:: python
 
@@ -22,7 +22,7 @@ You can also use open as a context manager to automatically close the file when
        print(ds.data)
 
 
-When opening multiple files that are linked to each other, use :py:meth:`opencosmo.open_linked_files`
+When opening multiple files that are linked to each other, use :py:meth:`opencosmo.open`
 
 Writing Data
 ------------

opencosmo/analysis/yt_viz.py

@@ -147,7 +147,7 @@ def visualize_halo(
         Identifier of the halo to be visualized.
     data : opencosmo.StructureCollection
         OpenCosmo StructureCollection object containing both halo properties and particle data
-        (e.g. output of ``opencosmo.open_linked_files([haloproperties, sodbighaloparticles])``).
+        (e.g. output of ``opencosmo.open([haloproperties, sodbighaloparticles])``).
     length_scale : str or None, optional
         Optionally add a horizontal bar denoting length scale in Mpc.
 
@@ -276,7 +276,7 @@ def halo_projection_array(
         array of projections). If `int`, a single panel is output while preserving formatting.
     data : opencosmo.StructureCollection
         OpenCosmo StructureCollection dataset containing both halo properties and particle data
-        (e.g., output of ``opencosmo.open_linked_files([haloproperties, sodbighaloparticles])``).
+        (e.g., output of ``opencosmo.open([haloproperties, sodbighaloparticles])``).
     field : tuple of str
         Field to plot for all panels. Follows yt naming conventions (e.g., ``("dm", "particle_mass")``,
         ``("gas", "temperature")``). Overridden if ``params["fields"]`` is provided.

opencosmo/collection/__init__.py

@@ -1,8 +1,8 @@
-from .io import open_collection, open_simulation_files
-from .lightcone import Lightcone, open_lightcone
+from .io import get_collection_type, open_simulation_files
+from .lightcone import Lightcone
 from .protocols import Collection
 from .simulation import SimulationCollection
-from .structure import StructureCollection, open_linked_file, open_linked_files
+from .structure import StructureCollection, open_linked_files
 
 __all__ = [
     "Collection",
@@ -11,9 +11,7 @@
     "open_multi_dataset_file",
     "SimulationCollection",
     "open_simulation_files",
-    "open_linked_file",
     "open_linked_files",
-    "open_collection",
     "Lightcone",
-    "open_lightcone",
+    "get_collection_type",
 ]

opencosmo/collection/io.py

@@ -1,14 +1,9 @@
 from pathlib import Path
-from typing import Type
-
-import h5py
 
 import opencosmo as oc
-from opencosmo import dataset as ds
-from opencosmo.collection.lightcone import Lightcone
-from opencosmo.collection.protocols import Collection
 from opencosmo.collection.simulation import SimulationCollection
-from opencosmo.collection.structure import StructureCollection
+from opencosmo.collection.structure.io import validate_linked_groups
+from opencosmo.io.io import FILE_TYPE, OpenTarget
 
 
 def open_simulation_files(**paths: Path) -> SimulationCollection:
@@ -37,46 +32,31 @@ def open_simulation_files(**paths: Path) -> SimulationCollection:
     return SimulationCollection(datasets)
 
 
-def open_collection(
-    handles: list[h5py.File | h5py.Group], load_kwargs: dict[str, bool]
-) -> Collection | ds.Dataset:
-    """
-    Open a file with multiple datasets.
-    """
-    CollectionType = get_collection_type(handles)
-    return CollectionType.open(handles, load_kwargs)
-
-
-def get_collection_type(handles: list[h5py.File | h5py.Group]) -> Type[Collection]:
+def get_collection_type(targets: list[OpenTarget], file_types: list[FILE_TYPE]):
     """
-    Determine the type of a single file containing multiple datasets. Currently
-    we support multi_simulation, particle, and linked collections.
+    If there are multiple files, determine their collection type. There are
+    three options we support at present:
 
-    multi_simulation == multiple simulations, same data types
-    particle == single simulation, multiple particle species
-    linked == A properties dataset, linked with other particle or profile datasets
+    1. All files contain a single lightcone dataset, all of the same type
+    2. The files contain a single non-lightcone datatype.
+    3. The files are linked together into a structure collection
     """
-    if len(handles) > 1:
-        return SimulationCollection
-
-    handle = handles[0]
-
-    datasets = [k for k in handle.keys() if k != "header"]
-    if len(datasets) == 0:
-        raise ValueError("No datasets found in file.")
-
-    if "header" not in handle.keys():
-        if all(group["header/file"].attrs["is_lightcone"] for group in handle.values()):
-            return Lightcone
-        return SimulationCollection
-    elif len(list(filter(lambda x: x.endswith("properties"), datasets))) >= 1:
-        return StructureCollection
-
-    elif handle["header/file"].attrs["is_lightcone"]:
-        return Lightcone
-
+    handles_by_type = {target.data_type: target.group for target in targets}
+    is_lightcone = [target.header.file.is_lightcone for target in targets]
+    unique_data_types = set(handles_by_type.keys())
+    unique_file_types = set(file_types)
+    if len(unique_data_types) == 1 and all(is_lightcone):
+        return oc.Lightcone
+    elif len(unique_data_types) == 1 and all(not il for il in is_lightcone):
+        return oc.SimulationCollection
+    elif unique_file_types == set([FILE_TYPE.PARTICLES, FILE_TYPE.DATASET]):
+        validate_linked_groups(handles_by_type)
+        return oc.StructureCollection
+    elif (
+        len(unique_file_types) == 1
+        and unique_file_types.pop() == FILE_TYPE.STRUCTURE_COLLECTION
+    ):
+        validate_linked_groups(handles_by_type)
+        return oc.StructureCollection
     else:
-        raise ValueError(
-            "Unknown file type. "
-            "It appears to have multiple datasets, but organized incorrectly"
-        )
+        raise ValueError("Invalid combination of files")

opencosmo/collection/lightcone/__init__.py

@@ -1,4 +1,3 @@
-from .io import open_lightcone
 from .lightcone import Lightcone
 
-__all__ = ["open_lightcone", "Lightcone"]
+__all__ = ["Lightcone"]