I/O Operations & Data Formats

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.

IO hierarchy diagram — **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:

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 convert_dask_array_to_host() before dafits.writeto; eager CuPy ndarrays are copied with to_host_ndarray(). use_dask selects the write path (streaming dafits vs eager astropy.io.fits), not whether the host sync happens.

Partial parameter snapshots (optimization)#

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 Host boundary contract (FITS, Zarr, MS) 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:

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 Dataset (antenna, field, ms_list / visibilities in a tree of groups), use ZarrDataset, not to be confused with ZarrArray (single image / Parameter).

Typical usage:

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: SimulationPipeline can export with export_format="zarr"; 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 Visibility sanitization (NaN / Inf) 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:

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 Array backends for install requirements, GPU imaging/optimization status, and alternatives (ArrayBackendContext, DASK_ARRAY__BACKEND, CLI --array-backend).

# 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 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 (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 Dataset.

I/O hooks (same semantics):

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

Default behavior uses CompositeSanitizer with VisibilitySanitizer and 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 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 Dataset: call apply_dataset() on a composite or single sanitizer; it walks dataset.ms_list and sanitizes each visibility dataset in place.
Single partition: call 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

from pyralysis.io import DaskMS

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

Measurement Set — custom sanitizer

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

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):

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:

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: pyralysis.sanitizers.

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.

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 (Partial parameter snapshots (optimization)) 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 zarr-gpu-policy in Array backends.

Understanding ROWID Behavior:

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

# 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:

# 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:

# 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:

# 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:

# 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

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:

# 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:

# 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:

# 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
)

Image Processing | Performance & Scalability

I/O Operations & Data Formats

Contents

I/O Operations & Data Formats#

Overview#

Available I/O Formats#

FITS (Flexible Image Transport System)#

Partial parameter snapshots (optimization)#

Zarr Arrays#

Zarr dataset (pyralysis Dataset)#

Measurement Sets (MS)#

Visibility sanitization (NaN / Inf)#

Host boundary contract (FITS, Zarr, MS)#

Performance Optimization#

Format Selection Guide#

Best Practices#

Troubleshooting#

Zarr dataset (pyralysis `Dataset`)#