Merge branch 'main' into release

Author: AstroPatty

README.md

@@ -7,10 +7,9 @@
 </h1><br>
 
 [![CI](https://github.com/ArgonneCPAC/OpenCosmo/actions/workflows/merge.yaml/badge.svg)](https://github.com/ArgonneCPAC/OpenCosmo/actions/workflows/merge.yaml)
-![PyPI - Version](https://img.shields.io/pypi/v/opencosmo)
-![Conda Version](https://img.shields.io/conda/vn/conda-forge/opencosmo)
-![PyPI - Python Version](https://img.shields.io/pypi/pyversions/opencosmo)
-![GitHub License](https://img.shields.io/github/license/ArgonneCPAC/opencosmo)
+[![PyPI - Version](https://img.shields.io/pypi/v/opencosmo)](https://pypi.org/project/opencosmo/)
+[![Conda Version](https://img.shields.io/conda/vn/conda-forge/opencosmo)](https://anaconda.org/conda-forge/opencosmo)
+[![GitHub License](https://img.shields.io/github/license/ArgonneCPAC/opencosmo)](https://github.com/ArgonneCPAC/OpenCosmo/blob/main/LICENSE.md)
 
 
 The OpenCosmo Python Toolkit provides utilities for reading, writing and manipulating data from cosmological simulations produced by the Cosmolgical Physics and Advanced Computing (CPAC) group at Argonne National Laboratory. It can be used to work with smaller quantities data retrieved with the CosmoExplorer, as well as the much larget datasets these queries draw from. The OpenCosmo toolkit integrates with standard tools such as AstroPy, and allows you to manipulate data in a fully-consistent cosmological context.

changes/+f722df63.misc.rst

@@ -0,0 +1 @@
+Unit handling has been fully rewritten, making it much more flexible for backend work.

changes/+fa256c74.feature.rst

@@ -0,0 +1 @@
+:code:`with_new_columns` now accepts a :code:`descriptions` argument for providing column descriptions

changes/122.feature.rst

@@ -0,0 +1 @@
+Descriptions of columns can now be accessed with :py:class:`Dataset.descriptions <opencosmo.Dataset.descriptions>`

changes/43.feature.rst

@@ -0,0 +1 @@
+:code:`with_units` can now be used to provide unit conversions, in addition to changing conventions

docs/source/units.rst

@@ -57,4 +57,132 @@ If you change unit conventions after performing a filter, the filter will still
 
 should output a value of ~1.5 x 10^10 Msun, which is about 1e10/0.67.
 
+Converting Columns to an Equivalent Unit
+----------------------------------------
+
+You can convert all columns with a given unit into a different unit with the :code:`conversions` argument:
+
+
+.. code-block:: python
+
+   import opencosmo as oc
    
+   ds = oc.read("haloproperties.hdf5")
+   conversions = {u.Mpc: u.lyr}
+   ds = ds.with_units(conversions=conversions)
+
+In the new dataset, all columns that originally had units of megaparsecs will be converted to lightyears. Composite units including megaparsec (e.g. km / s / Mpc, Mpc^2) will *not* be converted. All-column conversions are always peformed after a change of unit conventions. Changing units *after* doing a conversion always clears the conversions.
+
+.. code-block:: python
+
+   import opencosmo as oc
+   
+   ds = oc.read("haloproperties.hdf5")
+   conversions = {u.Mpc: u.lyr}
+   ds = ds.with_units(conversions=conversions)
+
+
+
+Single-Column Conversions
+-------------------------
+
+You can also use :code:`with_units` to convert the values in individual columns to their values in an equivalent unit:
+
+.. code-block:: python
+
+   import astropy.units as u
+
+   dataset = oc.read("haloproperties.hdf5").with_units(
+        fof_halo_center_x = u.lyr,
+        fof_halo_center_y = u.lyr,
+        fof_halo_center_z = u.lyr,
+   )
+
+Unit conversions like these are always performed *after* any change in unit convention, and changing unit conventions clears any existing unit conversions:
+
+.. code-block:: python
+
+    # this works
+    dataset = dataset.with_units(fof_halo_mass=u.kg)
+
+    # this clears the previous conversion,
+    # the masses are now in Msun / h
+    dataset = dataset.with_units("scalefree")
+
+    # This now fails, because the units of masses
+    # are Msun / h, which cannot be converted to kg
+    dataset = dataset.with_units(fof_halo_mass=u.kg)
+
+    # this will work, the units of halo mass in the "physical"
+    # convention are Msun (no h), and the change of convention
+    # happens before the conversions
+    dataset = dataset.with_units("physical", fof_halo_mass=u.kg, fof_halo_center_x=u.lyr)
+
+    # reset all units
+    dataset = dataset.with_units("physical")
+
+
+Unit conversions on :py:class:`Lightcones <opencosmo.Lightcone>` and :py:class:`SimulationCollections <opencosmo.SimulationCollection>` behave identically to single datasets. In :py:class:`StructureCollections <opencosmo.StructureCollections>`, unit conversions must be passed on a per-dataset basis:
+
+.. code-block:: python
+
+   import astropy.units as u
+
+   structures = oc.open("haloproperties.hdf5", "haloparticles.hdf5")
+   structures = structures.with_units(
+        halo_properties={"fof_halo_mass": u.kg},
+        dm_particles={"mass": u.kg}
+   )
+
+As with all-column conversions, composite units that include the target unit will not be converted. If you want to convert a composite unit, the conversion must be stated seperately.
+
+Conversion Precedence
+---------------------
+
+In cases where a blanket conversion is provided alongside a conversion for a specific column, the specific conversion always take precedence:
+
+.. code-block:: python
+
+   import astropy.units as u
+
+   conversions = {u.Mpc: u.lyr}
+   ds = ds.with_units(conversions=conversions, fof_halo_center_x=u.km)
+
+All columns with units of megaparsecs will be converted to lightyears, except for the :code:`fof_halo_center_x` column which will be converted to kilometers.
+
+Structure Collection Conversions
+--------------------------------
+
+When working with a structure collection, you can provide conversions that apply to the entire collection, as single dataset inside the collection, or individual columns within a given dataset. As you might expect, conversions on an individual dataset takes precedence over those that apply to all datasets.
+
+.. code-block:: python
+
+            import astropy.units as u
+
+            conversions = {u.Mpc: u.lyr}
+            structures = structures.with_units(
+                conversions=conversions
+                halo_properties = {
+                    "conversions": {u.Mpc: u.km},
+                    "fof_halo_center_x": u.m
+                }
+            )
+
+In this example, all values in Mpc will be converted to lightyears, except in the "halo_properties" dataset, where they will be converted to kilometers. The column "fof_halo_center_x" in "halo_properties" will be converted to meters instead.
+
+Clearing Conversions
+--------------------
+
+Conversions are always cleared when changing unit conventions, or you can also clear them by calling :code:`with_units` with no arguments.
+
+.. code-block:: python
+
+   dataset = oc.read("haloproperties.hdf5").with_units(
+        conversions={u.Mpc: u.lyr},
+        fof_halo_center_x = u.lyr,
+   )
+
+   dataset = dataset.with_units()
+   # all unit conversion reset
+
+

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "opencosmo"
-version = "0.9.4"
+version = "0.10.0a1"
 description = "OpenCosmo Python Toolkit"
 authors = [
   { name = "Patrick Wells", email = "pwells@anl.gov" },

src/opencosmo/collection/lightcone/lightcone.py

@@ -1,4 +1,4 @@
-from functools import reduce
+from functools import cached_property, reduce
 from itertools import chain
 from typing import Any, Callable, Generator, Iterable, Optional, Self
 
@@ -203,6 +203,25 @@ def columns(self) -> list[str]:
         cols = list(filter(lambda col: col not in self.__hidden, cols))
         return cols
 
+    @cached_property
+    def descriptions(self) -> dict[str, Optional[str]]:
+        """
+        Return the descriptions (if any) of the columns in this lightcone as a dictonary.
+        Columns without a description will be included in the dictionary with a value
+        of None
+
+        Returns
+        -------
+
+        descriptions : dict[str, str | None]
+            The column descriptions
+        """
+        descriptions = next(iter(self.values())).descriptions
+        descriptions = dict(
+            filter(lambda kv: kv[0] not in self.__hidden, descriptions.items())
+        )
+        return descriptions
+
     @property
     def cosmology(self) -> Cosmology:
         """
@@ -747,7 +766,11 @@ def take(self, n: int, at: str = "random") -> "Lightcone":
             output = {k: v for k, v in reversed(output.items())}
         return Lightcone(output, self.z_range, self.__hidden, self.__ordered_by)
 
-    def with_new_columns(self, **columns: DerivedColumn | np.ndarray | u.Quantity):
+    def with_new_columns(
+        self,
+        descriptions: str | dict[str, str] = {},
+        **columns: DerivedColumn | np.ndarray | u.Quantity,
+    ):
         """
         Create a new dataset with additional columns. These new columns can be derived
         from columns already in the dataset, a numpy array, or an Astropy quantity
@@ -758,7 +781,13 @@ def with_new_columns(self, **columns: DerivedColumn | np.ndarray | u.Quantity):
 
         Parameters
         ----------
-        ** columns : opencosmo.DerivedColumn
+        descriptions : str | dict[str, str], optional
+            A description for the new columns. These descriptions will be accessible through
+            :py:attr:`Lightcone.descriptions <opencosmo.Lighcone.descriptions>`. If a dictionary,
+            should have keys matching the column names.
+
+        ** columns : opencosmo.DerivedColumn | np.ndarray | u.quantity
+            The new columns
 
         Returns
         -------
@@ -786,7 +815,7 @@ def with_new_columns(self, **columns: DerivedColumn | np.ndarray | u.Quantity):
         for i, (ds_name, ds) in enumerate(self.items()):
             raw_columns = {name: arrs[i] for name, arrs in raw_split.items()}
             columns_input = raw_columns | derived
-            new_dataset = ds.with_new_columns(**columns_input)
+            new_dataset = ds.with_new_columns(descriptions, **columns_input)
             new_datasets[ds_name] = new_dataset
         return Lightcone(new_datasets, self.z_range, self.__hidden, self.__ordered_by)
 
@@ -828,23 +857,69 @@ def sort_by(self, column: str, invert: bool = False):
             raise ValueError(f"Column {column} does not exist in this dataset!")
         return Lightcone(dict(self), self.z_range, self.__hidden, (column, invert))
 
-    def with_units(self, convention: str) -> Self:
-        """
-        Create a new dataset from this one with a different unit convention.
+    def with_units(
+        self,
+        convention: Optional[str] = None,
+        conversions: dict[u.Unit, u.Unit] = {},
+        **columns: u.Unit,
+    ) -> Self:
+        r"""
+        Create a new lightcone from this one with a different unit convention or
+        with certain columns converted to a different compatible unit.
+
+        Unit conversions are always performed after a change of convention, and
+        changing conventions clears any existing unit conversions.
+
+        For more, see :doc:`units`.
+
+        .. code-block:: python
+
+            import astropy.units as u
+
+            # this works
+            lc = lc.with_units(fof_halo_mass=u.kg)
+
+            # this clears the previous conversion
+            lc = lc.with_units("scalefree")
+
+            # This now fails, because the units of masses
+            # are Msun / h, which cannot be converted to kg
+            lc = lc.with_units(fof_halo_mass=u.kg)
+
+            # this will now work, wince the units of halo mass in the "physical"
+            # convention are Msun (no h).
+            lc = lc.with_units("physical", fof_halo_mass=u.kg, fof_halo_center_x=u.lyr)
+
+            # Suppose you want your distances in lightyears, but the x coordinate of your
+            # halo center in kilometers, for some reason ¯\_(ツ)_/¯
+            blanket_conversions = {u.Mpc: u.lyr}
+            lc = lc.with_units(conversions = blanket_conversions, fof_halo_center_x = u.km)
 
         Parameters
         ----------
-        convention : str
+        convention : str, optional
             The unit convention to use. One of "physical", "comoving",
             "scalefree", or "unitless".
 
+        conversions: dict[astropy.units.Unit, astropy.units.Unit]
+            Conversions that apply to all columns in the lightcone with the
+            unit given by the key.
+
+        **column_conversions: astropy.units.Unit
+            Custom unit conversions for specific columns
+            in this dataset.
+
         Returns
         -------
-        dataset : Dataset
-            The new dataset with the requested unit convention.
-
-        """
-        return self.__map("with_units", convention)
+        lightcone : Lightcone
+            The new lightcone with the requested unit convention and/or conversions.
+        """
+        return self.__map(
+            "with_units",
+            convention=convention,
+            conversions=conversions,
+            **columns,
+        )
 
     def collect(self) -> "Lightcone":
         """

src/opencosmo/collection/simulation/simulation.py

@@ -1,5 +1,6 @@
 from typing import Callable, Iterable, Mapping, Optional, Self
 
+import astropy.units as u
 import h5py
 from astropy.cosmology import Cosmology  # type: ignore
 
@@ -276,7 +277,11 @@ def take(self, n: int, at: str = "random") -> Self:
         return self.__map("take", n, at)
 
     def with_new_columns(
-        self, *args, datasets: Optional[str | Iterable[str]] = None, **kwargs
+        self,
+        *args,
+        datasets: Optional[str | Iterable[str]] = None,
+        descriptions: str | dict[str, str] = {},
+        **kwargs,
     ):
         """
         Update the datasets within this collection with a set of new columns.
@@ -294,6 +299,11 @@ def with_new_columns(
         datasets: str | list[str], optional
             The datasets to add the columns to.
 
+        descriptions : str | dict[str, str], optional
+            A description for the new columns. These descriptions will be accessible through
+            :py:attr:`SimulationCollection(datasets).descriptions <opencosmo.SimulationCollection.descriptions>`.
+            If a dictionary, should have keys matching the column names.
+
         ** columns : opencosmo.DerivedColumn | np.ndarray | units.Quantity
             The new columns
         """
@@ -305,10 +315,14 @@ def with_new_columns(
 
             output = {name: ds for name, ds in self.items()}
             for ds_name in datasets:
-                output[ds_name] = output[ds_name].with_new_columns(*args, **kwargs)
+                output[ds_name] = output[ds_name].with_new_columns(
+                    *args, descriptions=descriptions, **kwargs
+                )
             return SimulationCollection(output)
 
-        return self.__map("with_new_columns", *args, **kwargs)
+        return self.__map(
+            "with_new_columns", *args, descriptions=descriptions, **kwargs
+        )
 
     def evaluate(
         self,
@@ -389,16 +403,31 @@ def sort_by(self, column: str, invert: bool = False):
         """
         return self.__map("sort_by", column=column, invert=invert)
 
-    def with_units(self, convention: str) -> Self:
+    def with_units(
+        self,
+        convention: Optional[str] = None,
+        conversions: dict[u.Unit, u.Unit] = {},
+        **columns: u.Unit,
+    ) -> Self:
         """
-        Transform all datasets or collections to use the given unit convention. This
-        method behaves exactly like :meth:`opencosmo.Dataset.with_units`.
+        Transform all datasets or collections to use the given unit convention, convert
+        all columns with a given unit into a different unit, and/or convert specific column(s)
+        to a compatible unit. This method behaves exactly like :meth:`opencosmo.Dataset.with_units`.
 
         Parameters
         ----------
         convention: str
             The unit convention to use. One of "unitless",
             "scalefree", "comoving", or "physical".
 
+        conversions: dict[astropy.units.Unit, astropy.units.Unit]
+            Conversions that apply to all columns in the collection with the
+            unit given by the key.
+
+        **column_conversions: astropy.units.Unit
+            Custom unit conversions for any column with a specific
+            name in the datasets in this collection.
+
+
         """
-        return self.__map("with_units", convention)
+        return self.__map("with_units", convention, conversions=conversions, **columns)