I/O Operations & Data Formats
==============================

Overview
--------
Pyralysis provides comprehensive I/O capabilities for various data formats commonly used in radio astronomy and scientific computing. The I/O system is designed to be efficient, scalable, and automatically handle metadata parsing.

.. note::
   All I/O classes inherit from the base `Io` class and provide consistent interfaces for reading and writing operations.

.. figure:: diagrams/png/fig07_io-1.png
   :alt: IO hierarchy diagram
   :width: 100%

   **I/O hierarchy** (``Io`` and concrete backends such as ``DaskMS``, ``Zarr``, and FITS-related classes).

Available I/O Formats
---------------------

FITS (Flexible Image Transport System)
--------------------------------------
The standard format for astronomical images with comprehensive header support.

**Key Features:**

- **Automatic Header Parsing**: FITS headers are automatically parsed for coordinate information
- **Dask Integration**: Support for out-of-core processing with chunking
- **Memory Mapping**: Efficient memory usage for large files
- **Header Preservation**: Maintains all FITS header information

**Usage Examples:**

.. code-block:: python

   from pyralysis.io import FITS

   # Read FITS image with automatic header parsing
   fits_io = FITS(
       input_name="image.fits",
       hdu=0,                    # HDU index (default: 0)
       use_dask=True,            # Enable Dask for large files
       chunks=(512, 512),        # Custom chunking
       memmap=True,              # Memory mapping (default: True)
       auto_parse_header=True    # Parse headers automatically
   )
   image = fits_io.read()

   # Write image to FITS with header preservation
   fits_io = FITS(output_name="output.fits")
   fits_io.write(
       image,
       overwrite=True,           # Overwrite existing file
       preserve_attributes=True  # Keep all attributes
   )

   # Read large FITS file efficiently
   fits_io = FITS(
       input_name="huge_image.fits",
       use_dask=True,
       chunks=(256, 256),        # Optimal chunking for large files
       lazy_load_hdus=True       # Load HDUs only when needed
   )
   image = fits_io.read()

.. note::

   ``write`` treats disk output as a **host boundary**. Pass the live ``Image``
   (NumPy or CuPy Dask) — do **not** call ``to_host()`` first. The handler
   converts CuPy-backed Dask to a host NumPy-meta collection lazily via
   :func:`~pyralysis.runtime.array_backend.convert_dask_array_to_host` before
   ``dafits.writeto``; eager CuPy ndarrays are copied with
   :func:`~pyralysis.runtime.array_backend.to_host_ndarray`. ``use_dask`` selects
   the write path (streaming ``dafits`` vs eager ``astropy.io.fits``), not
   whether the host sync happens.

.. _partial-parameter-io:

Partial parameter snapshots (optimization)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

:class:`~pyralysis.optimization.optimizer.Optimizer.save_partial_parameter`
passes the **live** iterate to ``io_handler.write``:

- ``Image`` → full ``Image`` object (so FITS can call ``create_header()`` when
  line search stripped ``xarray`` attrs)
- other ``Parameter`` → ``parameter.data``

Device-to-host transfer and Dask materialization are **I/O responsibility**,
not the optimizer's. After a partial save, the in-memory ``parameter`` keeps
its active backend (for example CuPy Dask on GPU sandboxes).

Control cadence with ``partial_parameter_interval`` on ``optimize()`` (default
``1`` = every iteration; CLI ``--partial-parameter-interval``). Snapshots use a
1-based iteration index; the final iteration is always written when partial
saves are enabled.

See :ref:`io-host-boundary-contract` for the backend matrix.

**FITS-Specific Features:**

- **HDU Selection**: Choose specific HDU (Header Data Unit)
- **Memory Mapping**: Efficient memory usage for large files
- **Header Validation**: Automatic validation of FITS headers
- **Coordinate System**: Automatic extraction of WCS information

Zarr Arrays
-----------
Modern, cloud-native format for large datasets with excellent compression and parallel I/O.

**Key Features:**

- **Native Chunking**: Built-in support for chunked arrays
- **Compression**: Multiple compression algorithms (blosc, lz4, zstd)
- **Parallel I/O**: Excellent performance for distributed computing
- **Cloud Storage**: Native support for cloud storage systems

**Usage Examples:**

.. code-block:: python

   from pyralysis.io import ZarrArray

   # Read Zarr image with Dask support
   zarr_io = ZarrArray(
       input_name="image.zarr",
       use_dask=True,            # Enable Dask arrays
       chunks="auto",            # Automatic chunking
       preserve_attributes=True, # Keep Zarr attributes
       auto_parse_attributes=True # Parse attributes automatically
   )
   image = zarr_io.read()

   # Write image to Zarr with compression
   zarr_io = ZarrArray(output_name="output.zarr")
   zarr_io.write(
       image,
       overwrite=True,
       compressor="blosc",       # Use blosc compression
       chunks=(256, 256)         # Set chunk size
   )

   # Read with custom chunking strategy
   zarr_io = ZarrArray(
       input_name="large_image.zarr",
       chunks=(128, 128),        # Small chunks for memory efficiency
       use_dask=True
   )
   image = zarr_io.read()

.. note::

   Like FITS, ``ZarrArray.write`` is a **host boundary**. Dask CuPy payloads
   are converted lazily per chunk in ``_fill_zarr_array``; eager CuPy ndarrays
   are copied with ``to_host_ndarray``. String chunk specs such as ``"auto"``
   and ``"100 MiB"`` (valid for Dask reads) are translated to Zarr-2-compatible
   values for ``zarr.create`` when the payload is eager NumPy. Dask payloads use
   ``data.chunksize`` directly.

**Zarr-Specific Features:**

- **Compression Options**: blosc, lz4, zstd, gzip
- **Chunking Strategies**: Flexible chunking for optimal performance
- **Attribute Preservation**: Maintains all Zarr metadata
- **Cloud Integration**: Works with S3, GCS, Azure

Zarr dataset (pyralysis ``Dataset``)
------------------------------------
For a **full** Pyralysis :class:`~pyralysis.base.dataset.Dataset` (antenna, field, ``ms_list`` / visibilities in a tree of groups), use :class:`~pyralysis.io.zarr.zarr_dataset.ZarrDataset`, not to be confused with :class:`~pyralysis.io.zarr.zarr_array.ZarrArray` (single image / ``Parameter``).

**Typical usage:**

.. code-block:: python

   from pyralysis.io import ZarrDataset

   zio = ZarrDataset(chunks={"ms_list": {"row": 10000, "chan": 32}})
   dataset = zio.read("my_dataset.zarr")

   # Same optional sanitization as Measurement Set reads (non-finite data / NaN)
   dataset = zio.read("my_dataset.zarr", sanitize=True)

   zio.write(dataset, "out.zarr")

**Simulation → Zarr → imager:** :class:`~pyralysis.pipelines.simulation.SimulationPipeline`
can export with ``export_format="zarr"``; :class:`~pyralysis.pipelines.imager.ImagerPipeline`
``LoadData`` accepts ``data_format="zarr"`` so visibilities stay as lazy Dask arrays on disk.
See ``tests/integration/pipelines/test_zarr_simulation_e2e_int.py``.

See :ref:`visibility-sanitization` for what ``sanitize`` does and how to customize it.


Measurement Sets (MS)
---------------------
Radio interferometry data format for visibility data and antenna information.

**Key Features:**

- **Visibility Data**: Handle complex visibility measurements
- **Antenna Information**: Access antenna positions and properties
- **Spectral Windows**: Manage frequency channels and bandwidth
- **Advanced Chunking**: Sophisticated chunking strategies for large datasets

**Usage Examples:**

.. code-block:: python

   from pyralysis.io import DaskMS

   # Read Measurement Set with default row chunking (100_000 rows per chunk)
   ms_io = DaskMS(input_name="data.ms")
   dataset = ms_io.read()

   # Read with advanced chunking strategy
   ms_io = DaskMS(
       input_name="data.ms",
       chunks={
           'row': (40000, 60000, 40000, 60000),  # Custom row chunks
           'chan': (16, 16, 16, 16),             # Channel chunks
           'corr': (1, 2, 1)                     # Correlation chunks
       }
   )
   dataset = ms_io.read()

   # Read with specific options
   dataset = ms_io.read(
       read_flagged_data=False,   # Skip flagged data
       filter_flag_column=True,   # Apply flag filtering
       calculate_psf=True,        # Calculate PSF
       taql_query="ANTENNA1 == 1" # Apply TAQL filtering
   )

   # Optional sanitization for NaN/Inf (visibilities and weights); see :ref:`visibility-sanitization`
   dataset = ms_io.read(sanitize=True)

   # Optional CuPy-backed Dask arrays (Milestone 1: read-time placement only)
   dataset = ms_io.read(array_backend="cupy", calculate_psf=False)

See :doc:`array_backend` for install requirements, GPU imaging/optimization status, and alternatives (``ArrayBackendContext``, ``DASK_ARRAY__BACKEND``, CLI ``--array-backend``).

.. code-block:: python

   # Write Measurement Set from simulated dataset
   ms_io = DaskMS(output_name="simulated.ms")
   ms_io.write(
       dataset=simulated_dataset,
       tables="ALL",  # Write all tables (MAIN + subtables)
   )

   # Write specific tables
   ms_io.write(
       dataset=dataset,
       tables=["MAIN", "ANTENNA", "FIELD"],  # Write only specified tables
   )

   # Write with lazy computation
   writes = ms_io.write(
       dataset=dataset,
       tables="ALL",
       dask_compute=False,  # Return lazy operations
   )
   dask.compute(writes)  # Compute manually

When visibilities were processed on GPU (CuPy-backed Dask arrays), ``write`` and ``write_xarray_ds`` automatically materialize numeric variables to host NumPy collections before calling dask-ms (see :func:`~pyralysis.runtime.array_backend.materialize_pyralysis_dataset_for_write`).

**MS-Specific Features:**

- **TAQL Queries**: Filter data using Table Query Language
- **Flag Management**: Handle data quality flags
- **PSF Calculation**: Automatic Point Spread Function computation
- **Multi-dimensional Chunking**: Optimize for different data dimensions
- **Optional Sanitization**: Read-time sanitization of invalid visibility/weight/sigma values
- **Complete MS Creation**: Write all tables (MAIN + subtables) from scratch
- **ROWID Handling**: Automatic append/update behavior based on ROWID presence
- **Batch Writing**: Efficient single compute call for all tables

.. _visibility-sanitization:

Visibility sanitization (NaN / Inf)
-----------------------------------
Some MS or archived Zarr stores contain **non-finite** values (NaN, Inf) in ``DATA``-like columns or in weights / sigmas. That can break gridding or optimization. Pyralysis can clean these **at read time** or on an **existing** :class:`~pyralysis.base.dataset.Dataset`.

**I/O hooks (same semantics):**

- :meth:`pyralysis.io.daskms.DaskMS.read` — optional ``sanitize=True`` and optional ``sanitizer=``.
- :meth:`pyralysis.io.zarr.zarr_dataset.ZarrDataset.read` — same ``sanitize`` / ``sanitizer`` arguments after opening each visibility group from Zarr.

Default behavior uses :class:`~pyralysis.sanitizers.composite_sanitizer.CompositeSanitizer` with :class:`~pyralysis.sanitizers.visibility_sanitizer.VisibilitySanitizer` and :class:`~pyralysis.sanitizers.weight_sanitizer.WeightSanitizer`. What counts as invalid is controlled by ``invalid_kind`` (e.g. all non-finite vs NaN only). How values and ``FLAG`` are updated is controlled per sanitizer (e.g. ``action``, ``action_3d``, ``action_2d`` on weights).

**CompositeSanitizer and ``FLAG``:** Each child sanitizer applies its own ``FLAG`` updates to the visibility ``xarray.Dataset`` as it runs. The composite also builds a **merged** ``flag_delta`` on the returned :class:`~pyralysis.sanitizers.sanitizer.SanitizationResult` for reporting or introspection only; it is **not** written again to ``FLAG`` (that would duplicate bitwise-OR work in the Dask graph).

**Important: ``FLAG`` shape vs compact weights**

- **3D** spectra-shaped arrays (row, chan, corr), e.g. ``DATA``, ``WEIGHT_SPECTRUM``, ``IMAGING_WEIGHT_SPECTRUM``: any column whose name contains ``WEIGHT`` or ``SIGMA`` and has three dimensions is treated like a per-channel weight; invalid samples can be **flagged** at the same indices as ``FLAG``.
- **2D** row–correlation weights (``WEIGHT``, ``SIGMA``): values can be **repaired** (e.g. zeroed), but the sanitizer **does not** write those into the 3D ``FLAG`` grid (which would wrongly imply a per-channel choice).

**Standalone (no ``read()``):** You do not need I/O to run sanitizers.

- **Entire** :class:`~pyralysis.base.dataset.Dataset`: call :meth:`~pyralysis.sanitizers.sanitizer.Sanitizer.apply_dataset` on a composite or single sanitizer; it walks ``dataset.ms_list`` and sanitizes each visibility dataset in place.
- **Single** partition: call :meth:`~pyralysis.sanitizers.sanitizer.Sanitizer.sanitize_visibility_dataset` on that SubMS’s ``ms.visibilities.dataset`` (an ``xarray.Dataset``).

**Examples**

The snippets below are **separate recipes**. Use **either** read-time ``sanitize=True`` **or** an in-memory pass (``apply_dataset`` / per-SubMS loop)—not both on the same load, or you would sanitize twice.

**Measurement Set — default pipeline**

.. code-block:: python

   from pyralysis.io import DaskMS

   dataset = DaskMS(input_name="data.ms").read(sanitize=True)

**Measurement Set — custom sanitizer**

.. code-block:: python

   from pyralysis.io import DaskMS
   from pyralysis.sanitizers import CompositeSanitizer, VisibilitySanitizer, WeightSanitizer

   sanitizer = CompositeSanitizer(
       sanitizers=[
           VisibilitySanitizer(action="flag_and_zero"),
           WeightSanitizer(action_3d="flag_and_zero", action_2d="zero"),
       ]
   )
   dataset = DaskMS(input_name="data.ms").read(sanitize=True, sanitizer=sanitizer)

**Zarr full ``Dataset`` — same read-time hook**

.. code-block:: python

   from pyralysis.io import ZarrDataset
   from pyralysis.sanitizers import CompositeSanitizer, VisibilitySanitizer, WeightSanitizer

   sanitizer = CompositeSanitizer(
       sanitizers=[
           VisibilitySanitizer(action="flag_and_zero"),
           WeightSanitizer(action_3d="flag_and_zero", action_2d="zero"),
       ]
   )
   dataset = ZarrDataset().read("data.zarr", sanitize=True, sanitizer=sanitizer)

**Already-loaded ``Dataset`` (no I/O sanitizer)**

Equivalent to what ``read(..., sanitize=True)`` does per partition, but for data you already hold (e.g. ``sanitize=False`` earlier, simulation, or another source):

.. code-block:: python

   from pyralysis.sanitizers import CompositeSanitizer

   CompositeSanitizer().apply_dataset(dataset)

**One SubMS visibility ``xarray.Dataset`` only**

Same per-partition call as inside I/O; use when you only need to clean one block:

.. code-block:: python

   from pyralysis.sanitizers import CompositeSanitizer

   sanitizer = CompositeSanitizer()
   # ms: one SubMS from dataset.ms_list
   result = sanitizer.sanitize_visibility_dataset(ms.visibilities.dataset)
   # Optional: result.stats, result.flag_delta for introspection

API reference: :mod:`pyralysis.sanitizers`.

.. _io-host-boundary-contract:

Host boundary contract (FITS, Zarr, MS)
---------------------------------------

Disk I/O always produces **host** files. Conversion from device or lazy Dask
graphs happens inside each handler at write time — callers pass live objects.

.. list-table::
   :header-rows: 1
   :widths: 22 18 18 42

   * - Payload
     - FITS ``use_dask=True``
     - Zarr ``use_dask=True``
     - Host transfer
   * - Dask, NumPy chunks
     - ``dafits.writeto``
     - ``da.store``
     - No device copy; lazy compute at write
   * - Dask, CuPy chunks
     - ``convert_dask_array_to_host`` → ``dafits``
     - ``convert_dask_array_to_host`` → ``da.store``
     - Lazy chunk GPU→host during write
   * - Eager NumPy
     - ``astropy.io.fits.writeto``
     - ``zarr_array[...] = data``
     - Already on host
   * - Eager CuPy
     - ``to_host_ndarray`` → ``writeto``
     - ``to_host_ndarray`` → assign
     - Eager copy before write

**Optimizer partial saves** (:ref:`partial-parameter-io`) follow the same rule:
the iterate stays on GPU/CPU in memory; only the snapshot file is host-backed.

**Zarr 2 vs GPU-native Zarr:** Pyralysis targets ``zarr>=2.12,<3`` (via
``dask-ms[complete]``). ``zarr.config.enable_gpu()`` (Zarr-Python ≥ 3) is not
available in this stack; see :ref:`zarr-gpu-policy` in :doc:`array_backend`.

**Understanding ROWID Behavior:**

When writing to Measurement Sets, the presence or absence of the `ROWID` coordinate determines the write behavior:

.. code-block:: python

   # Simulated datasets (no ROWID) → Always APPEND rows
   ms_io.write(dataset=simulated_dataset, tables="ALL")

   # Datasets with ROWID → UPDATE existing rows
   # (ROWID coordinates are automatically handled by dask-ms)

**Tables Parameter Options:**

.. code-block:: python

   # Write only MAIN table (default, backward compatible)
   ms_io.write(dataset=dataset, tables=None)

   # Write all tables
   ms_io.write(dataset=dataset, tables="ALL")

   # Write specific tables
   ms_io.write(dataset=dataset, tables=["MAIN", "ANTENNA", "POLARIZATION"])

Performance Optimization
------------------------
Optimize I/O performance for different use cases:

**Chunking Strategies:**

.. code-block:: python

   # Small images (< 1024x1024): No chunking needed
   fits_io = FITS(input_name="small_image.fits", use_dask=False)

   # Medium images (1024x1024 - 4096x4096): Moderate chunking
   fits_io = FITS(
       input_name="medium_image.fits",
       use_dask=True,
       chunks=(512, 512)
   )

   # Large images (> 4096x4096): Aggressive chunking
   fits_io = FITS(
       input_name="large_image.fits",
       use_dask=True,
       chunks=(256, 256)
   )

   # Very large images: Very small chunks
   fits_io = FITS(
       input_name="huge_image.fits",
       use_dask=True,
       chunks=(128, 128)
   )

**Memory Management:**

.. code-block:: python

   # Enable memory mapping for FITS files
   fits_io = FITS(
       input_name="large_image.fits",
       memmap=True,              # Memory mapping
       use_dask=True,            # Dask for out-of-core
       chunks=(256, 256)         # Optimal chunking
   )

   # Use lazy loading for multi-HDU files
   fits_io = FITS(
       input_name="multi_hdu.fits",
       lazy_load_hdus=True,      # Load HDUs only when needed
       use_dask=True
   )

**Parallel I/O:**

.. code-block:: python

   # Configure for distributed computing
   import dask.config
   dask.config.set(scheduler='distributed')

   # Use Zarr for best parallel performance
   zarr_io = ZarrArray(
       input_name="distributed_image.zarr",
       use_dask=True,
       chunks=(256, 256)
   )

Format Selection Guide
----------------------
Choose the right format for your use case:

**For Astronomical Images:**

- **FITS**: Standard format, excellent header support, wide compatibility
- **Zarr**: Large datasets, cloud storage, parallel processing

**For Scientific Data:**

- **Zarr**: Large datasets, compression, cloud-native

**For Radio Astronomy:**

- **Measurement Set**: Visibility data, antenna information, spectral data
- **FITS**: Image products, calibration solutions

**For Large Datasets:**

- **Zarr**: Best performance, compression, parallel I/O
- **FITS + Dask**: Good performance with wide compatibility

**For Cloud Computing:**

- **Zarr**: Native cloud storage support
- **FITS**: Standard format, good cloud compatibility

Best Practices
--------------

**1. Format Selection:**

- Use FITS for standard astronomical images
- Use Zarr for large datasets requiring compression
- Use Measurement Sets for radio interferometry data

**2. Performance Optimization:**

- Enable Dask for files larger than 1GB
- Use appropriate chunking for your data size
- Enable memory mapping for FITS files
- Use compression for Zarr files

**3. Metadata Preservation:**

- Always enable `preserve_attributes`
- Use `auto_parse_header` for automatic coordinate extraction
- Check that metadata is properly preserved after I/O operations

**4. Error Handling:**

- Check file existence before reading
- Handle large files with appropriate chunking
- Validate data after reading operations
- Use try-catch blocks for robust I/O operations

**Example: Robust I/O Workflow**

.. code-block:: python

   from pyralysis.io import FITS, ZarrArray
   import os

   def robust_read_image(filename):
       """Robustly read an image file with automatic format detection."""

       # Check file existence
       if not os.path.exists(filename):
           raise FileNotFoundError(f"File not found: {filename}")

       # Determine format and create appropriate I/O handler
       if filename.endswith('.fits'):
           io_handler = FITS(
               input_name=filename,
               use_dask=True,
               chunks=(256, 256),
               auto_parse_header=True
           )
       elif filename.endswith('.zarr'):
           io_handler = ZarrArray(
               input_name=filename,
               use_dask=True,
               chunks=(256, 256),
               auto_parse_attributes=True
           )
       elif filename.endswith('.ms'):
           io_handler = DaskMS(
               input_name=filename,
               chunks={'row': 100_000}
           )
       else:
           raise ValueError(f"Unsupported file format: {filename}. Supported formats: .fits, .zarr, .ms")

       try:
           # Read the image
           image = io_handler.read()

           # Validate the result
           if image is None:
               raise ValueError("Failed to read image data")

           return image

       except Exception as e:
           raise RuntimeError(f"Error reading {filename}: {e}")

Troubleshooting
---------------
Common I/O issues and solutions:

**1. Memory Issues:**

.. code-block:: python

   # Problem: Out of memory when reading large file
   # Solution: Enable Dask and use chunking
   fits_io = FITS(
       input_name="large_file.fits",
       use_dask=True,
       chunks=(256, 256)
   )

   # Problem: Memory mapping fails
   # Solution: Disable memory mapping
   fits_io = FITS(
       input_name="problematic_file.fits",
       memmap=False
   )

**2. Performance Issues:**

.. code-block:: python

   # Problem: Slow reading
   # Solution: Optimize chunking
   zarr_io = ZarrArray(
       input_name="slow_file.zarr",
       chunks=(512, 512)  # Larger chunks for better performance
   )

   # Problem: No parallel processing
   # Solution: Check Dask configuration
   import dask.config
   print(f"Number of workers: {dask.config.get('num_workers')}")

**3. Metadata Issues:**

.. code-block:: python

   # Problem: Headers not parsed
   # Solution: Check auto_parse_header setting
   fits_io = FITS(
       input_name="file.fits",
       auto_parse_header=True
   )

   # Problem: Attributes lost
   # Solution: Enable preserve_attributes
   zarr_io = ZarrArray(
       input_name="file.zarr",
       preserve_attributes=True
   )

----

:doc:`image_processing` | :doc:`performance`