Expose simulation parameters

docs/requirements.txt

@@ -1 +1,2 @@
 sphinx_rtd_theme>=3.0.2
+autodoc_pydantic>=2.2.0

docs/source/collection_ref.rst

@@ -8,6 +8,7 @@ Collections within a single file can always be loaded with :py:func:`opencosmo.o
    :members:
    :exclude-members: open,read,close,write
    :undoc-members:
+   :member-order: bysource
 
 
 .. autoclass:: opencosmo.StructureCollection

docs/source/conf.py

@@ -14,7 +14,12 @@
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
-extensions = ["sphinx.ext.autodoc", "sphinx.ext.napoleon", "sphinx_rtd_theme"]
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.napoleon",
+    "sphinx_rtd_theme",
+    "sphinxcontrib.autodoc_pydantic",
+]
 
 templates_path = ["_templates"]
 exclude_patterns = []
@@ -28,3 +33,5 @@
 # -- Options for autodoc -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html
 autodoc_typehints = "description"
+autodoc_pydantic_model_show_validator_summary = False
+autodoc_pydantic_field_list_validators = False

docs/source/index.rst

@@ -39,5 +39,6 @@ The OpenCosmo Python Toolkit provides utilities for reading, writing and manipul
    io_ref
    dataset_ref
    collection_ref
+   parameters_ref
 
 

docs/source/mpi.rst

@@ -1,7 +1,7 @@
 Working with MPI
 ================
 
-OpenCosmo can read and write data inside an MPI environment. In general the API works exactly the same within an MPI context as it does otherwise, but there are some things to be aware of in the current version of the library (see below). More flexibility in working in an MPI context is planned for the next version of the library.
+OpenCosmo can read and write data in an MPI environment. In general the API works exactly the same within an MPI context as it does otherwise, but there are some things to be aware of in the current version of the library (see below). More flexibility in working in an MPI context is planned for the next version of the library.
 
 I/O with Parallel HDF5 and Select Operations
 --------------------------------------------

docs/source/parameters_ref.rst

@@ -0,0 +1,24 @@
+Simulation Parameters
+=====================
+
+You can access the parameters of the simulation any dataset is drawn from with :py:attr:`opencosmo.Dataset.simulation`. All datasets regardless of simulation will have the parameters in :py:class:`opencosmo.parameters.SimulationParameters`. Hydrodynamic simulations will additionally contain the parameters in :py:class:`opencosmo.parameters.SubgridParameters`
+
+
+.. autoclass:: opencosmo.parameters.SimulationParameters
+   :members:
+   :undoc-members:
+   :exclude-members: model_config,empty_string_to_none,cosmology_parameters
+   :member-order: bysource
+
+
+.. autoclass:: opencosmo.parameters.HydroSimulationParameters
+   :members:
+   :undoc-members:
+   :exclude-members: model_config
+   :member-order: bysource
+
+.. autoclass:: opencosmo.parameters.SubgridParameters
+   :members:
+   :undoc-members:
+   :exclude-members: model_config
+   :member-order: bysource

opencosmo/collection/collection.py

@@ -11,13 +11,15 @@
 
 
 import h5py
+from astropy.cosmology import Cosmology  # type: ignore
 
 import opencosmo as oc
 from opencosmo.dataset.index import ChunkedIndex
 from opencosmo.dataset.mask import Mask
 from opencosmo.handler import InMemoryHandler, OpenCosmoDataHandler, OutOfMemoryHandler
 from opencosmo.header import OpenCosmoHeader, read_header
 from opencosmo.link import StructureCollection
+from opencosmo.parameters import SimulationParameters
 from opencosmo.transformations import units as u
 
 
@@ -172,6 +174,43 @@ def __map(self, method, *args, **kwargs):
         output = {k: getattr(v, method)(*args, **kwargs) for k, v in self.items()}
         return SimulationCollection(output)
 
+    def __map_attribute(self, attribute):
+        return {k: getattr(v, attribute) for k, v in self.items()}
+
+    @property
+    def cosmology(self) -> dict[str, Cosmology]:
+        """
+        Get the cosmologies of the simulations in the collection
+
+        Returns
+        --------
+        cosmologies: dict[str, astropy.cosmology.Cosmology]
+        """
+        return self.__map_attribute("cosmology")
+
+    @property
+    def redshift(self) -> dict[str, float]:
+        """
+        Get the redshift slices for the simulations in the collection
+
+        Returns
+        --------
+        redshifts: dict[str, float]
+        """
+        return self.__map_attribute("redshift")
+
+    @property
+    def simulation(self) -> dict[str, SimulationParameters]:
+        """
+        Get the simulation parameters for the simulations in the collection
+
+        Returns
+        --------
+        simulation_parameters: dict[str, opencosmo.parameters.SimulationParameters]
+        """
+
+        return self.__map_attribute("simulation")
+
     def filter(self, *masks: Mask, **kwargs) -> SimulationCollection:
         """
         Filter the datasets in the collection. This method behaves
@@ -213,6 +252,28 @@ def select(self, *args, **kwargs) -> SimulationCollection:
         """
         return self.__map("select", *args, **kwargs)
 
+    def take(self, n: int, at: str = "random") -> SimulationCollection:
+        """
+        Take a subest of rows from all datasets or collections in this collection.
+        This method will delegate to the underlying method in
+        :class:`opencosmo.Dataset`, or :class:`opencosmo.StructureCollection` depending
+        on  the context. As such, behavior may vary depending on what this collection
+        contains. See their documentation for more info.
+
+        Parameters
+        ----------
+        n: int
+            The number of rows to take
+        at: str, default = "random"
+            The method to use to take rows. Must be one of "start", "end", "random".
+
+        """
+        if any(len(ds) < n for ds in self.values()):
+            raise ValueError(
+                f"Not all datasets in this collection have at least {n} rows!"
+            )
+        return self.__map("take", n, at)
+
     def with_units(self, convention: str) -> SimulationCollection:
         """
         Transform all datasets or collections to use the given unit convention. This
@@ -227,16 +288,6 @@ def with_units(self, convention: str) -> SimulationCollection:
         """
         return self.__map("with_units", convention)
 
-    def take(self, *args, **kwargs) -> SimulationCollection:
-        """
-        Take a subest of rows from all datasets or collections in this collection.
-        This method will delegate to the underlying method in
-        :class:`opencosmo.Dataset`, or :class:`opencosmo.Collection` depending on  the
-        context. As such, behaviormay vary depending on what this collection contains.
-        """
-
-        return self.__map("take", *args, **kwargs)
-
 
 def open_single_dataset(
     file: h5py.File, dataset_key: str, header: Optional[OpenCosmoHeader] = None

opencosmo/collection/io.py

@@ -29,7 +29,7 @@ def open_simulation_files(**paths: Path) -> SimulationCollection:
         dataset = oc.open(path)
         if not isinstance(dataset, oc.Dataset):
             raise ValueError("All datasets must be of the same type.")
-    dtypes = set(dataset.header.file.data_type for dataset in datasets.values())
+    dtypes = set(dataset for dataset in datasets.values())
     if len(dtypes) != 1:
         raise ValueError("All datasets must be of the same type.")
     return SimulationCollection(datasets)

opencosmo/dataset/dataset.py

@@ -14,6 +14,7 @@
 from opencosmo.dataset.mask import Mask, apply_masks
 from opencosmo.handler import OpenCosmoDataHandler
 from opencosmo.header import OpenCosmoHeader, write_header
+from opencosmo.parameters import SimulationParameters
 
 
 class Dataset:
@@ -67,10 +68,45 @@ def cosmology(self) -> Cosmology:
 
         Returns
         -------
-        cosmology : astropy.cosmology.Cosmology
+        cosmology: astropy.cosmology.Cosmology
         """
         return self.__header.cosmology
 
+    @property
+    def dtype(self) -> str:
+        """
+        The data type of this dataset.
+
+        Returns
+        -------
+        dtype: str
+        """
+        return self.__header.file.data_type
+
+    @property
+    def redshift(self) -> float:
+        """
+        The redshift slice this dataset was drawn from
+
+        Returns
+        -------
+        redshift: float
+
+        """
+        return self.__header.file.redshift
+
+    @property
+    def simulation(self) -> SimulationParameters:
+        """
+        The parameters of the simulation this dataset is drawn
+        from.
+
+        Returns
+        -------
+        parameters: opencosmo.parameters.SimulationParameters
+        """
+        return self.__header.simulation
+
     @property
     def data(self) -> Table | Column:
         """
@@ -85,13 +121,7 @@ def data(self) -> Table | Column:
         """
         # should rename this, dataset.data can get confusing
         # Also the point is that there's MORE data than just the table
-        return self.__handler.get_data(
-            builders=self.__builders, index=self.__index
-        )
-
-    @property
-    def header(self) -> OpenCosmoHeader:
-        return self.__header
+        return self.__handler.get_data(builders=self.__builders, index=self.__index)
 
     @property
     def index(self) -> DataIndex:
@@ -332,7 +362,7 @@ def with_units(self, convention: str) -> Dataset:
             convention,
             self.__base_unit_transformations,
             self.__header.cosmology,
-            self.__header.file.redshift,
+            self.redshift,
         )
         new_builders = get_column_builders(new_transformations, self.__builders.keys())