FITSWriter#

class pyvisgen.io.datawriters.FITSWriter(output_path: Path, dataset_type: str, **kwargs)[source]#

Bases: DataWriter

FITS file writer for pyvisgen visibility datasets.

This writer saves visibility data and observation information to FITS (Flexible Image Transport System) files. Each sample is written to a separate .fits file.

Parameters:

output_pathstr or Path: Directory path where FITS files will be written.
dataset_typestr: Type of dataset being written (e.g., ‘train’, ‘test’, ‘validation’). This is used in the file names.

Examples

>>> writer = FITSWriter(output_path="./data")
>>> writer.write(vis_data, obs, index=0)

Or as a context manager:

>>> with FITSWriter(output_path="./data") as writer:
...     writer.write(vis_data, obs, index=0)

Methods Summary

`get_half_image`(x, y[, overlap])	Extract half height of every image with a small overlap.
`test_shapes`(array, name)	Validate the shape of input arrays.
`write`(vis_data, obs, index[, overwrite])	Write visibility data and observation metadata to a FITS file.

Methods Documentation

get_half_image(x: ndarray, y: ndarray, overlap: int = 5) → tuple[ndarray]#

Extract half height of every image with a small overlap.

Parameters:

xnp.ndarray: Simulated data array with shape (B, C, H, W).
ynp.ndarray: Ground truth array with shape (B, C, H, W).

Returns:

tuple[np.ndarray, np.ndarray]: Tuple containing the cropped x and y arrays.

test_shapes(array: ndarray, name: str) → None#

Validate the shape of input arrays.

Arrays should have the shape (B, C, H, W), where B is the batch size, C the number of channels (2), and W and H the width and height of the images.

Parameters:

arraynp.ndarray: Array to validate.
namestr: Name of the array for error reporting.

Raises:

ValueError: If array axis 1 is not size 2.
ValueError: If array does not have exactly 4 dimensions.

write(vis_data, obs, index, overwrite=True, **kwargs) → None[source]#

Write visibility data and observation metadata to a FITS file.

Creates a new FITS file for each sample with pattern vis_{dataset_type}_{index}.fits.

Parameters:

vis_dataarray-like: Visibility data to be written to the FITS file.
obsobject: Observation metadata object from Observation.
indexint: Sample index used in the output filename.
overwritebool, optional: If True, overwrite the output file if it already exists, otherwise an error is raised. Default: True.