create_adios_table.py

# This script requires the changes in https://github.com/casacore/casacore/pull/1116,
# which as of 14/Jul/2021 haven't been merged upstream yet

import functools
import operator
import os
import sys
import time

import casacore.tables as tables
import numpy as np

def get_size(msdir):
    size = 0
    for root, dirs, files in os.walk(msdir):
        size += sum(os.path.getsize(os.path.join(root, name))
                    for name in files + dirs)
    return size

def write_ms(filename, compressor, nrows, cell_shape, accuracy):

    size = functools.reduce(operator.mul, cell_shape, nrows * 8)
    print(f"Will write {size / 1024 / 1024:.2f} MB of data into {filename}")

    # Prepare stman information
    tabdesc = tables.maketabdesc(
        tables.makearrcoldesc('DATA', '',
            valuetype='double', shape=cell_shape,
            datamanagergroup='group0', datamanagertype='Adios2StMan'
        )
    )
    dminfo = tables.makedminfo(
        tabdesc,
        {
            'group0': {
                'OPERATORPARAMS': {
                    'DATA': {
                        'Operator': compressor,
                        'Accuracy': str(accuracy),
                    }
                }
            }
        }
    )

    # The fake data, only ones
    ones = np.ones((nrows,) + cell_shape)

    # Create table, write some data
    t0 = time.time()
    t = tables.table(filename, tabledesc=tabdesc, dminfo=dminfo, ack=False)
    t.addrows(nrows)
    data = t.col('DATA')
    data.putcol(ones)
    t.summary()
    print(t.showstructure())
    t.close()
    t1 = time.time()

    on_disk_size = get_size(filename)

    print(f'Table created and written in {(t1 - t0):.2f} s')
    print(f'Table size on disk: {on_disk_size / 1024 / 1024:.2f} MB')
    print(f'Compression ratio: {on_disk_size / size:.2f}')

def main():
    filename = sys.argv[1]
    compressor = sys.argv[2]
    nrows = int(sys.argv[3])
    cell_shape = tuple(map(int, sys.argv[4:]))
    write_ms(filename, compressor, nrows, cell_shape, 0.1)

if __name__ == '__main__':
    main()