Unit Handling Update

docs/source/units.rst

@@ -1,9 +1,13 @@
 Working with Units
 ==================
 
-In general, the raw output of cosmological simulations is *scale free*: coordinates and velocities use `comoving coordinates <https://en.wikipedia.org/wiki/Comoving_and_proper_distances#Comoving_distance_and_proper_distance>`_ and values are given in terms of the reduced Hubble constant *h*.
 
-OpenCosmo allows you to convert between these conventions as needed based on what is most convinient for your analysis. The default convention is "comoving." This leaves coordinates and velocities in comoving units, but absorbs the *h* factor into the values of the data. The other available conventions are "physical", "scalefree" and "unitless." Note that "scalefree" and "unitless" will have the same numerical values.
+The raw output of cosmological simulations produced with HACC (and many other cosmology codes) is in *scale free* units: coordinates and velocities use `comoving coordinates <https://en.wikipedia.org/wiki/Comoving_and_proper_distances#Comoving_distance_and_proper_distance>`_ and values such as mass include terms of the reduced Hubble constant *h*. Some downstream data products, such as synthetic galaxy catalogs, may be in other conventions.
+
+OpenCosmo allows you to convert between these conventions as needed based on what is most convinient for your analysis. The default convention is "comoving," which leaves coordinates and velocities in comoving units but absorbs any factors of h into their values. The other available conventions are "physical", "scalefree" and "unitless." 
+
+The "comoving", "physical" and "unitless" conventions will always be available, while the "scalefree" convention will only be available if the raw data is already stored in this convention. The "unitless" convention will always return the raw data in the file, regardless of the convention the data is stored in.
+
 
 You can readily convert between these conventions with the :py:meth:`opencosmo.Dataset.with_units` method:
 

opencosmo/collection/io.py

@@ -43,6 +43,7 @@ def open_collection(
     """
     Open a file with multiple datasets.
     """
+    print("opening collection")
     CollectionType = get_collection_type(handles)
     return CollectionType.open(handles, load_kwargs)
 
@@ -71,6 +72,10 @@ def get_collection_type(handles: list[h5py.File | h5py.Group]) -> Type[Collectio
         return SimulationCollection
     elif len(list(filter(lambda x: x.endswith("properties"), datasets))) >= 1:
         return StructureCollection
+
+    elif handle["header/file"].attrs["is_lightcone"]:
+        return Lightcone
+
     else:
         raise ValueError(
             "Unknown file type. "

opencosmo/collection/lightcone/io.py

@@ -1,4 +1,3 @@
-from functools import reduce
 from pathlib import Path
 from typing import cast
 
@@ -39,9 +38,14 @@ def open_lightcone(files: list[Path], **load_kwargs):
     if len(steps) != len(files):
         raise ValueError("Each file must contain only a single lightcone step!")
 
-    datasets: dict[str, Dataset] = reduce(
-        lambda left, right: left | open_lightcone_file(right), files, {}
-    )
+    datasets = {}
+    for file in files:
+        new_ds = oc.open(file)
+        if not isinstance(new_ds, Lightcone):
+            raise ValueError("Didn't find a lightcone in a lightcone file!")
+        for key, ds in new_ds.items():
+            key = "_".join([ds.dtype, str(ds.header.file.step)])
+            datasets[key] = ds
 
     z_range = headers[0].lightcone.z_range
     return Lightcone(datasets, z_range)

opencosmo/collection/lightcone/lightcone.py

@@ -11,8 +11,9 @@
 import opencosmo as oc
 from opencosmo.dataset import Dataset
 from opencosmo.dataset.col import Mask
-from opencosmo.header import OpenCosmoHeader
+from opencosmo.header import OpenCosmoHeader, read_header
 from opencosmo.index import SimpleIndex
+from opencosmo.io.io import evaluate_load_conditions, open_single_dataset
 from opencosmo.io.schemas import LightconeSchema
 from opencosmo.parameters.hacc import HaccSimulationParameters
 from opencosmo.spatial import Region
@@ -66,7 +67,7 @@ def __init__(
         self,
         datasets: dict[str, Dataset],
         z_range: Optional[tuple[float, float]] = None,
-        hide_redshift: bool = True,
+        hide_redshift: bool = False,
     ):
         datasets = {k: with_redshift_column(ds) for k, ds in datasets.items()}
         self.update(datasets)
@@ -239,13 +240,24 @@ def open(cls, handles: list[h5py.File | h5py.Group], load_kwargs):
             raise NotImplementedError
         handle = handles[0]
         datasets: dict[str, Dataset] = {}
-        for key, group in handle.items():
-            ds = oc.open(group)
-            if not isinstance(ds, Dataset):
+        try:
+            header = read_header(handle)
+            headers = {key: header for key in handle.keys()}
+        except KeyError:
+            headers = {key: read_header(group) for key, group in handle.items()}
+
+        groups = {k: v for k, v in handle.items() if k != "header"}
+        groups = evaluate_load_conditions(groups, load_kwargs)
+
+        for key, group in groups.items():
+            if key == "header":
+                continue
+            ds = open_single_dataset(group, headers[key])
+            if not isinstance(ds, Lightcone) or len(ds.keys()) != 1:
                 raise ValueError(
                     "Lightcones can only contain datasets (not collections)"
                 )
-            datasets[key] = ds
+            datasets[key] = next(iter(ds.values()))
 
         z_range = next(iter(datasets.values())).header.lightcone.z_range
 
@@ -279,7 +291,8 @@ def with_redshift_range(self, z_low: float, z_high: float):
             new_dataset = dataset.filter(
                 oc.col("redshift") > z_low, oc.col("redshift") < z_high
             )
-            new_datasets[key] = new_dataset
+            if len(new_dataset) > 0:
+                new_datasets[key] = new_dataset
         return Lightcone(new_datasets, (z_low, z_high))
 
     def __map(self, method, *args, hide_redshift: bool = False, **kwargs):
@@ -295,6 +308,10 @@ def __map_attribute(self, attribute):
         return {k: getattr(v, attribute) for k, v in self.items()}
 
     def make_schema(self) -> LightconeSchema:
+        if len(self.keys()) == 1:
+            schema = next(iter(self.values())).make_schema()
+            schema.header = self.__header
+            return schema
         schema = LightconeSchema()
         for name, dataset in self.items():
             ds_schema = dataset.make_schema()

opencosmo/collection/simulation/simulation.py

@@ -276,7 +276,12 @@ def read_single_dataset(
     )
     index = ChunkedIndex.from_size(len(handler))
     state = DatasetState(
-        base_unit_transformations, builders, index, u.UnitConvention.COMOVING, sim_box
+        base_unit_transformations,
+        builders,
+        index,
+        u.UnitConvention.COMOVING,
+        sim_box,
+        header,
     )
 
     return Dataset(handler, header, state, tree)

opencosmo/dataset/build.py

@@ -43,6 +43,7 @@ def build_dataset(
         index,
         u.UnitConvention.SCALEFREE,
         sim_box,
+        header,
     )
 
     dataset = Dataset(

opencosmo/dataset/state.py

@@ -29,13 +29,15 @@ def __init__(
         index: DataIndex,
         convention: u.UnitConvention,
         region: Region,
+        header: OpenCosmoHeader,
         hidden: set[str] = set(),
         derived: dict[str, DerivedColumn] = {},
     ):
         self.__base_unit_transformations = base_unit_transformations
         self.__builder = builder
         self.__index = index
         self.__convention = convention
+        self.__header = header
         self.__region = region
         self.__hidden = hidden
         self.__derived: dict[str, DerivedColumn] = derived
@@ -56,6 +58,10 @@ def convention(self):
     def region(self):
         return self.__region
 
+    @property
+    def header(self):
+        return self.__header
+
     @property
     def columns(self) -> list[str]:
         columns = set(self.__builder.columns) | set(self.__derived.keys())
@@ -81,6 +87,7 @@ def with_index(self, index: DataIndex):
             index,
             self.__convention,
             self.__region,
+            self.__header,
             self.__hidden,
             self.__derived,
         )
@@ -133,6 +140,7 @@ def with_derived_columns(self, **new_columns: DerivedColumn):
             self.__index,
             self.__convention,
             self.__region,
+            self.__header,
             self.__hidden,
             new_derived,
         )
@@ -156,6 +164,7 @@ def with_region(self, region: Region):
             self.__index,
             self.__convention,
             region,
+            self.__header,
             self.__hidden,
             self.__derived,
         )
@@ -216,6 +225,7 @@ def select(self, columns: str | Iterable[str]):
             self.__index,
             self.__convention,
             self.__region,
+            self.__header,
             new_hidden,
             new_derived,
         )
@@ -245,16 +255,17 @@ def take_range(self, start: int, end: int):
         return self.with_index(new_index)
 
     def with_units(
-        self,
-        convention: str,
-        cosmology: Cosmology,
-        redshift: float | tuple[float, float],
+        self, convention: str, cosmology: Cosmology, redshift: float | table.Column
     ):
         """
         Change the unit convention
         """
         new_transformations = u.get_unit_transition_transformations(
-            convention, self.__base_unit_transformations, cosmology, redshift
+            self.__header.file.unit_convention,
+            convention,
+            self.__base_unit_transformations,
+            cosmology,
+            redshift,
         )
         convention_ = u.UnitConvention(convention)
         new_builder = get_table_builder(new_transformations, self.__builder.columns)
@@ -264,6 +275,7 @@ def with_units(
             self.__index,
             convention_,
             self.__region,
+            self.__header,
             self.__hidden,
             self.__derived,
         )

opencosmo/io/io.py

@@ -12,7 +12,7 @@
 from opencosmo.dataset import state as dss
 from opencosmo.dataset.handler import DatasetHandler
 from opencosmo.file import FileExistance, file_reader, file_writer, resolve_path
-from opencosmo.header import read_header
+from opencosmo.header import OpenCosmoHeader, read_header
 from opencosmo.index import ChunkedIndex
 from opencosmo.mpi import get_comm_world
 from opencosmo.spatial.builders import from_model
@@ -116,8 +116,13 @@ def verify_files(handles: list[h5py.File | h5py.Group]):
             raise ValueError("All files should have the same set of data types!")
 
 
-def open_single_dataset(handle: h5py.File | h5py.Group):
-    header = read_header(handle)
+def open_single_dataset(
+    handle: h5py.File | h5py.Group, header: Optional[OpenCosmoHeader] = None
+):
+    if header is None:
+        header = read_header(handle)
+    assert header is not None
+
     try:
         tree = open_tree(handle, header.simulation.box_size, header.file.is_lightcone)
     except ValueError:
@@ -152,6 +157,7 @@ def open_single_dataset(handle: h5py.File | h5py.Group):
         index,
         u.UnitConvention.COMOVING,
         sim_region,
+        header,
     )
 
     dataset = oc.Dataset(
@@ -160,6 +166,9 @@ def open_single_dataset(handle: h5py.File | h5py.Group):
         state,
         tree=tree,
     )
+
+    if header.file.is_lightcone:
+        return collection.Lightcone({"data": dataset}, header.lightcone.z_range)
     return dataset
 
 
@@ -269,10 +278,19 @@ def read(
         group, header
     )
     state = dss.DatasetState(
-        base_unit_transformations, builders, index, u.UnitConvention.COMOVING, sim_box
+        base_unit_transformations,
+        builders,
+        index,
+        u.UnitConvention.COMOVING,
+        sim_box,
+        header,
     )
 
-    return oc.Dataset(handler, header, state, tree)
+    ds = oc.Dataset(handler, header, state, tree)
+
+    if header.file.is_lightcone:
+        return collection.Lightcone({"data": ds})
+    return ds
 
 
 @file_writer

opencosmo/parameters/diffsky.py

@@ -0,0 +1,14 @@
+from typing import ClassVar
+
+from pydantic import BaseModel
+
+
+class DiffskyVersionInfo(BaseModel):
+    ACCESS_PATH: ClassVar[str] = "diffsky_versions"
+    diffmah: str
+    diffsky: str
+    diffstar: str
+    diffstarpop: str
+    dsps: str
+    jax: str
+    numpy: str

opencosmo/parameters/file.py

@@ -37,6 +37,7 @@ class FileType(Enum):
     halo_properties = "halo_properties"
     halo_profiles = "halo_profiles"
     halo_particles = "halo_particles"
+    diffsky_fits = "diffsky_fits"
 
 
 class FileParameters(BaseModel):
@@ -47,6 +48,7 @@ class FileParameters(BaseModel):
     redshift: float | tuple[float, float]
     step: int
     region: Optional[BoxRegionModel | ConeRegionModel | HealPixRegionModel] = None
+    unit_convention: str = "scalefree"
 
     @model_validator(mode="before")
     @classmethod

opencosmo/parameters/hacc.py

@@ -14,6 +14,8 @@
 
 from opencosmo.parameters import CosmologyParameters
 
+from .diffsky import DiffskyVersionInfo
+
 
 def empty_string_to_none(v):
     if isinstance(v, str) and v == "":
@@ -204,4 +206,5 @@ def empty_string_to_none(cls, data):
     "halo_particles": {},
     "galaxy_particles": {},
     "halo_profiles": {},
+    "diffsky_fits": {"diffsky_versions": DiffskyVersionInfo},
 }

opencosmo/parameters/parameters.py

@@ -11,8 +11,9 @@ def read_header_attributes(
     parameter_model: Type[BaseModel],
     **kwargs,
 ):
+    header = file["header"]
     try:
-        header_data = file["header"][header_path].attrs
+        header_data = header[header_path].attrs
     except KeyError:
         return parameter_model()  # Defaults are possible
     try: