Source code for mojito.writer

"""Contain functions to create and write datasets in Mojito L1 HDF5 files.

You can use the functions in this module to create and write datasets and
attributes in a Mojito L1 HDF5 file. In particular, you can combine multiple
single-brick files into a single file containing all the data from the input
files, using the :func:`combine_bricks` function.

.. note::

    You must open the HDF5 file in write ('w') or append ('a') mode to use these
    functions, see :class:`mojito.reader.MojitoL1File`.

Combining bricks
----------------

Use the :func:`combine_bricks` function to combine multiple single-brick files
into a single file containing all the data from the input files.

The function can check that the input files are consistent with each other
before combining, and will raise an error if any of the files are not
consistent.

.. code-block:: python

    from mojito.writer import combine_bricks

    combine_bricks(
        paths=["brick1.h5", "brick2.h5", "brick3.h5"],
        output_path="combined.h5",
        mode="w",
        check_consistent=True,
        compression="lzf",
    )

.. autofunction:: combine_bricks

Writing attributes
------------------

Use the following function to write the main attributes to a Mojito L1 file.

.. autofunction:: write_attrs

Creating groups and datasets
----------------------------

You can create the main groups and datasets in a Mojito L1 file using the
following functions:

.. autofunction:: create_tdis
.. autofunction:: create_ltts
.. autofunction:: create_orbits
.. autofunction:: create_noise_estimates

The datasets will be created empty and with default compression (LZF), unless
you provide data to initialize them with, or specify a different compression
algorithm.

Note that missing datasets in the provided data will not be copied, and will be
created empty instead.

Writing sampling attributes
---------------------------

.. autofunction:: write_uniform_time_sampling
.. autofunction:: write_log_uniform_frequency_sampling

"""

import logging
from typing import Any

from h5py import Group

from .reader import LTT, TDI, FileOpenMode, MojitoL1File
from .sampling import LogUniformFrequencySampling, UniformTimeSampling

logger = logging.getLogger(__name__)

DEFAULT_COMPRESSION = "lzf"
"""Default compression algorithm for HDF5 datasets."""


[docs] def write_uniform_time_sampling(group: Group, sampling: UniformTimeSampling) -> None: """Write the sampling attributes to an HDF5 group. Parameters ---------- group HDF5 group where to write the sampling attributes. sampling Uniform time sampling instance containing the attributes to write. """ group.attrs["t0"] = sampling.t0 group.attrs["dt"] = sampling.dt group.attrs["size"] = sampling.size
[docs] def write_log_uniform_frequency_sampling( group: Group, sampling: LogUniformFrequencySampling ) -> None: """Write the log-uniform frequency sampling attributes to an HDF5 group. Parameters ---------- group HDF5 group where to write the frequency sampling attributes. sampling Log-uniform frequency sampling instance containing the attributes to write. """ group.attrs["fmin"] = sampling.fmin group.attrs["fmax"] = sampling.fmax group.attrs["size"] = sampling.size
[docs] def create_tdis( mojito_file: MojitoL1File, time_sampling: UniformTimeSampling, *, data: TDI | None = None, compression: str | None = DEFAULT_COMPRESSION, **kwargs, ) -> None: """Create TDI group and datasets in the given HDF5 file. Parameters ---------- mojito_file Single-brick MojitoL1File instance where to create the TDI group. time_sampling Uniform time sampling for the TDI datasets. data TDI data to initialize the datasets with. Missing datasets will not be copied and will be created empty. If None, datasets are created empty. compression Compression algorithm to use for the datasets. **kwargs Additional keyword arguments for dataset creation. Raises ------ ValueError If the provided MojitoL1File is a combined file (multi-brick). """ # Check that the file is single-brick if mojito_file.is_combined: raise ValueError("TDI group can only be created in a single-brick file") group = mojito_file.files[0] tdi_group = group.create_group("tdis") time_sampling_group = tdi_group.create_group("sampling") write_uniform_time_sampling(time_sampling_group, time_sampling) def create_dataset( dname: str, source_dname: str | None = None, *, dtype: str ) -> None: """Create a dataset in the TDI group. Parameters ---------- dname Name of the dataset to create. source_dname Name of the dataset in the source data to copy from. If None, use ``dname``. dtype Data type of the dataset. """ source_dname = dname if source_dname is None else source_dname data_to_copy = None try: data_to_copy = getattr(data, source_dname)[:] except (AttributeError, KeyError): logger.warning("Missing dataset '%s' will not be copied", dname) tdi_group.create_dataset( dname, dtype=dtype, data=data_to_copy, shape=(time_sampling.size,), compression=compression, **kwargs, ) create_dataset("eta_12", dtype="f8") create_dataset("eta_23", dtype="f8") create_dataset("eta_31", dtype="f8") create_dataset("eta_13", dtype="f8") create_dataset("eta_32", dtype="f8") create_dataset("eta_21", dtype="f8") create_dataset("A2", "a2", dtype="f8") create_dataset("E2", "e2", dtype="f8") create_dataset("T2", "t2", dtype="f8") create_dataset("X2", "x2", dtype="f8") create_dataset("Y2", "y2", dtype="f8") create_dataset("Z2", "z2", dtype="f8") create_dataset("eta_flags", dtype="i8") create_dataset("tdi_flags", dtype="i8")
[docs] def create_ltts( mojito_file: MojitoL1File, time_sampling: UniformTimeSampling, *, data: LTT | None = None, compression: str | None = DEFAULT_COMPRESSION, **kwargs, ) -> None: """Create LTT group and datasets in the given HDF5 file. Parameters ---------- mojito_file Single-brick MojitoL1File instance where to create the LTT group. time_sampling Uniform time sampling for the LTT datasets. data LTT data to initialize the datasets with. Missing datasets will not be copied and will be created empty. If None, datasets are created empty. compression Compression algorithm to use for the datasets. **kwargs Additional keyword arguments for dataset creation. Raises ------ ValueError If the provided MojitoL1File is a combined file (multi-brick). """ # Check that the file is single-brick if mojito_file.is_combined: raise ValueError("LTT group can only be created in a single-brick file") group = mojito_file.files[0] ltt_group = group.create_group("ltts") time_sampling_group = ltt_group.create_group("sampling") write_uniform_time_sampling(time_sampling_group, time_sampling) def create_dataset(dname: str) -> None: """Create a dataset in the LTT group. Parameters ---------- dname Name of the dataset to create. """ data_to_copy = None try: data_to_copy = getattr(data, dname)[:] except (AttributeError, KeyError): logger.warning("Missing dataset '%s' will not be copied", dname) ltt_group.create_dataset( dname, dtype="f8", data=data_to_copy, shape=(time_sampling.size,), compression=compression, **kwargs, ) create_dataset("ltt_12") create_dataset("ltt_23") create_dataset("ltt_31") create_dataset("ltt_13") create_dataset("ltt_32") create_dataset("ltt_21") create_dataset("ltt_derivative_12") create_dataset("ltt_derivative_23") create_dataset("ltt_derivative_31") create_dataset("ltt_derivative_13") create_dataset("ltt_derivative_32") create_dataset("ltt_derivative_21")
[docs] def create_orbits( mojito_file: MojitoL1File, time_sampling: UniformTimeSampling, *, data: Any | None = None, compression: str | None = DEFAULT_COMPRESSION, **kwargs, ) -> None: """Create Orbits group and datasets in the given HDF5 file. Parameters ---------- mojito_file Single-brick MojitoL1File instance where to create the Orbits group. time_sampling Uniform time sampling for the Orbits datasets. data Orbits data to initialize the datasets with. Missing datasets will not be copied and will be created empty. If None, datasets are created empty. compression Compression algorithm to use for the datasets. **kwargs Additional keyword arguments for dataset creation. Raises ------ ValueError If the provided MojitoL1File is a combined file (multi-brick). """ # Check that the file is single-brick if mojito_file.is_combined: raise ValueError("Orbits group can only be created in a single-brick file") group = mojito_file.files[0] orbits_group = group.create_group("orbits") time_sampling_group = orbits_group.create_group("sampling") write_uniform_time_sampling(time_sampling_group, time_sampling) def create_dataset(dname: str, source_dname: str | None = None) -> None: """Create a dataset in the orbits group. Parameters ---------- dname Name of the dataset to create. source_dname Name of the dataset in the source data to copy from. If None, use ``dname``. """ source_dname = dname if source_dname is None else source_dname data_to_copy = None try: data_to_copy = getattr(data, source_dname)[:] except (AttributeError, KeyError): logger.warning("Missing dataset '%s' will not be copied", dname) orbits_group.create_dataset( dname, dtype="f8", data=data_to_copy, shape=(time_sampling.size, 3), compression=compression, **kwargs, ) create_dataset("sc_position_1", "position_1") create_dataset("sc_position_2", "position_2") create_dataset("sc_position_3", "position_3") create_dataset("sc_velocity_1", "velocity_1") create_dataset("sc_velocity_2", "velocity_2") create_dataset("sc_velocity_3", "velocity_3")
[docs] def create_noise_estimates( mojito_file: MojitoL1File, time_sampling: UniformTimeSampling, freq_sampling: LogUniformFrequencySampling, *, data: Any | None = None, compression: str | None = DEFAULT_COMPRESSION, **kwargs, ) -> None: """Create Noise Estimates group and datasets in the given HDF5 file. Parameters ---------- mojito_file Single-brick MojitoL1File instance where to create the Noise Estimates group. time_sampling Uniform time sampling for the noise estimates datasets. freq_sampling Log-uniform frequency sampling for the noise estimates datasets. data Noise estimates data to initialize the datasets with. Missing datasets will not be copied and will be created empty. If None, datasets are created empty. compression Compression algorithm to use for the datasets. **kwargs Additional keyword arguments for dataset creation. Raises ------ ValueError If the provided MojitoL1File is a combined file (multi-brick). """ # Check that the file is single-brick if mojito_file.is_combined: raise ValueError( "Noise estimate group can only be created in a single-brick file" ) group = mojito_file.files[0] noise_group = group.create_group("noise_estimates") time_sampling_group = noise_group.create_group("sampling") write_uniform_time_sampling(time_sampling_group, time_sampling) freq_sampling_group = noise_group.create_group("log_frequency_sampling") write_log_uniform_frequency_sampling(freq_sampling_group, freq_sampling) def create_dataset( dname: str, source_dname: str | None = None, *, size: int ) -> None: """Create a dataset in the noise estimate group. Parameters ---------- dname Name of the dataset to create. source_dname Name of the dataset in the source data to copy from. If None, use ``dname``. size Size of the last two dimensions of the dataset. """ source_dname = dname if source_dname is None else source_dname data_to_copy = None try: data_to_copy = getattr(data, source_dname)[:] except (AttributeError, KeyError): logger.warning("Missing dataset '%s' will not be copied", dname) noise_group.create_dataset( dname, dtype="c16", data=data_to_copy, shape=(time_sampling.size, freq_sampling.size, size, size), compression=compression, **kwargs, ) create_dataset("XYZ", "xyz", size=3) create_dataset("AET", "aet", size=3) create_dataset("eta", "eta", size=6)
[docs] def write_attrs( mojito_file: MojitoL1File, *, pipeline_name: str, laser_frequency: float, lolipops_version: str, ) -> None: """Write attributes to the Mojito L1 file. Parameters ---------- mojito_file Single-brick MojitoL1File instance where to write the attributes. pipeline_name Name of the pipeline that generated the file. laser_frequency Laser frequency in Hz. lolipops_version Version of the Lolipops software used. Raises ------ ValueError If the provided MojitoL1File is a combined file (multi-brick). """ # Check that the file is single-brick if mojito_file.is_combined: raise ValueError("Attributes can only be written in a single-brick file") file = mojito_file.files[0] file.attrs["pipeline_name"] = str(pipeline_name) file.attrs["laser_frequency"] = float(laser_frequency) file.attrs["lolipops_version"] = str(lolipops_version)
[docs] def combine_bricks( paths: list, output_path: str, *, mode: FileOpenMode = "w-", check_consistent: bool = True, compression: str | None = DEFAULT_COMPRESSION, ) -> None: """Combine multiple Mojito brick files into a single file. A new file is created at ``output_path`` and filled with the combined data from the input files. Consistency between the input files is verified using :func:`mojito.MojitoL1File.check_consistent` and raises an error if any of the files are not consistent with each other. Parameters ---------- paths A list of paths to the input files to be combined. output_path The path to the output file to be created. mode The mode to open the output file in. Must be a mode that allows writing. check_consistent Whether to check that all input files are consistent with each other before combining. If True, then an error is raised if any of the files are not consistent with each other. If False, then the consistency check is skipped. compression The compression to use for the output file. Must be a valid h5py compression method or None for no compression. If not specified, then the default compression defined in :attr:`mojito.writer.DEFAULT_COMPRESSION`` is used. Raises ------ ValueError If the list of paths is empty. Exception If any of the input files are not consistent with each other and ``check_consistent`` is True. """ logger.info("Combining bricks: %s", ", ".join(paths)) # Check the list of paths is not empty if not paths: raise ValueError("No input files provided") # Open the output file using requested mode logger.debug("Opening output file %s with mode '%s'...", output_path, mode) with MojitoL1File(output_path, mode) as output_file: # Open input files with MojitoL1File(paths, "r") as input_files: # Check consistency of input files if requested if check_consistent: logger.info("Checking consistency of input files...") input_files.check_consistent() # Write the attributes logger.info("Writing attributes...") pipeline_name = f"combined-brick ({', '.join(input_files.pipeline_names)})" write_attrs( output_file, laser_frequency=input_files.laser_frequency, lolipops_version=input_files.lolipops_version, pipeline_name=pipeline_name, ) # Write each group in the output file logger.info("Writing TDI group...") create_tdis( output_file, input_files.tdis.time_sampling, data=input_files.tdis, compression=compression, ) logger.info("Writing LTT group...") create_ltts( output_file, input_files.ltts.time_sampling, data=input_files.ltts, compression=compression, ) logger.info("Writing orbits group...") create_orbits( output_file, input_files.orbits.time_sampling, data=input_files.orbits, compression=compression, ) logger.info("Writing noise_estimates group...") create_noise_estimates( output_file, input_files.noise_estimates.time_sampling, input_files.noise_estimates.freq_sampling, data=input_files.noise_estimates, compression=compression, ) logger.info("Combined bricks into %s", output_path)