diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000000000000000000000000000000000000..16beae29cc7f31303c709b4981421f395512d84b --- /dev/null +++ b/.editorconfig @@ -0,0 +1,148 @@ +root = true + +[*] +charset = utf-8 +end_of_line = lf +indent_size = 4 +indent_style = space +insert_final_newline = true +trim_trailing_whitespace = true +max_line_length = 120 +tab_width = 4 +ij_continuation_indent_size = 8 +ij_formatter_off_tag = @formatter:off +ij_formatter_on_tag = @formatter:on +ij_formatter_tags_enabled = false +ij_smart_tabs = false +ij_visual_guides = none +ij_wrap_on_typing = false + +[{*.bash,*.sh,*.zsh}] +indent_size = 2 +tab_width = 2 +ij_shell_binary_ops_start_line = false +ij_shell_keep_column_alignment_padding = false +ij_shell_minify_program = false +ij_shell_redirect_followed_by_space = false +ij_shell_switch_cases_indented = false +ij_shell_use_unix_line_separator = true + +[{*.har,*.jsb2,*.jsb3,*.json,.babelrc,.eslintrc,.stylelintrc,bowerrc,jest.config}] +indent_size = 2 +ij_json_array_wrapping = split_into_lines +ij_json_keep_blank_lines_in_code = 0 +ij_json_keep_indents_on_empty_lines = false +ij_json_keep_line_breaks = true +ij_json_keep_trailing_comma = false +ij_json_object_wrapping = split_into_lines +ij_json_property_alignment = do_not_align +ij_json_space_after_colon = true +ij_json_space_after_comma = true +ij_json_space_before_colon = false +ij_json_space_before_comma = false +ij_json_spaces_within_braces = false +ij_json_spaces_within_brackets = false +ij_json_wrap_long_lines = false + +[{*.markdown,*.md}] +ij_markdown_force_one_space_after_blockquote_symbol = true +ij_markdown_force_one_space_after_header_symbol = true +ij_markdown_force_one_space_after_list_bullet = true +ij_markdown_force_one_space_between_words = true +ij_markdown_insert_quote_arrows_on_wrap = true +ij_markdown_keep_indents_on_empty_lines = false +ij_markdown_keep_line_breaks_inside_text_blocks = true +ij_markdown_max_lines_around_block_elements = 1 +ij_markdown_max_lines_around_header = 1 +ij_markdown_max_lines_between_paragraphs = 1 +ij_markdown_min_lines_around_block_elements = 1 +ij_markdown_min_lines_around_header = 1 +ij_markdown_min_lines_between_paragraphs = 1 +ij_markdown_wrap_text_if_long = true +ij_markdown_wrap_text_inside_blockquotes = true + +[{*.py,*.pyw}] +max_line_length = 88 +ij_python_align_collections_and_comprehensions = true +ij_python_align_multiline_imports = true +ij_python_align_multiline_parameters = true +ij_python_align_multiline_parameters_in_calls = false +ij_python_blank_line_at_file_end = true +ij_python_blank_lines_after_imports = 1 +ij_python_blank_lines_after_local_imports = 0 +ij_python_blank_lines_around_class = 1 +ij_python_blank_lines_around_method = 1 +ij_python_blank_lines_around_top_level_classes_functions = 2 +ij_python_blank_lines_before_first_method = 0 +ij_python_call_parameters_new_line_after_left_paren = true +ij_python_call_parameters_right_paren_on_new_line = true +ij_python_call_parameters_wrap = normal +ij_python_dict_alignment = 2 +ij_python_dict_new_line_after_left_brace = true +ij_python_dict_new_line_before_right_brace = true +ij_python_dict_wrapping = 1 +ij_python_from_import_new_line_after_left_parenthesis = true +ij_python_from_import_new_line_before_right_parenthesis = true +ij_python_from_import_parentheses_force_if_multiline = true +ij_python_from_import_trailing_comma_if_multiline = false +ij_python_from_import_wrapping = 1 +ij_python_hang_closing_brackets = false +ij_python_keep_blank_lines_in_code = 1 +ij_python_keep_blank_lines_in_declarations = 1 +ij_python_keep_indents_on_empty_lines = false +ij_python_keep_line_breaks = true +ij_python_method_parameters_new_line_after_left_paren = true +ij_python_method_parameters_right_paren_on_new_line = true +ij_python_method_parameters_wrap = normal +ij_python_new_line_after_colon = true +ij_python_new_line_after_colon_multi_clause = true +ij_python_optimize_imports_always_split_from_imports = false +ij_python_optimize_imports_case_insensitive_order = false +ij_python_optimize_imports_join_from_imports_with_same_source = false +ij_python_optimize_imports_sort_by_type_first = true +ij_python_optimize_imports_sort_imports = true +ij_python_optimize_imports_sort_names_in_from_imports = false +ij_python_space_after_comma = true +ij_python_space_after_number_sign = true +ij_python_space_after_py_colon = true +ij_python_space_before_backslash = true +ij_python_space_before_comma = false +ij_python_space_before_for_semicolon = false +ij_python_space_before_lbracket = false +ij_python_space_before_method_call_parentheses = false +ij_python_space_before_method_parentheses = false +ij_python_space_before_number_sign = true +ij_python_space_before_py_colon = false +ij_python_space_within_empty_method_call_parentheses = false +ij_python_space_within_empty_method_parentheses = false +ij_python_spaces_around_additive_operators = true +ij_python_spaces_around_assignment_operators = true +ij_python_spaces_around_bitwise_operators = true +ij_python_spaces_around_eq_in_keyword_argument = false +ij_python_spaces_around_eq_in_named_parameter = false +ij_python_spaces_around_equality_operators = true +ij_python_spaces_around_multiplicative_operators = true +ij_python_spaces_around_power_operator = false +ij_python_spaces_around_relational_operators = true +ij_python_spaces_around_shift_operators = true +ij_python_spaces_within_braces = false +ij_python_spaces_within_brackets = false +ij_python_spaces_within_method_call_parentheses = false +ij_python_spaces_within_method_parentheses = false +ij_python_use_continuation_indent_for_arguments = true +ij_python_use_continuation_indent_for_collection_and_comprehensions = false +ij_python_use_continuation_indent_for_parameters = true +ij_python_wrap_long_lines = false + +[{*.yaml,*.yml}] +indent_size = 2 +ij_yaml_align_values_properties = do_not_align +ij_yaml_autoinsert_sequence_marker = true +ij_yaml_block_mapping_on_new_line = false +ij_yaml_indent_sequence_value = true +ij_yaml_keep_indents_on_empty_lines = false +ij_yaml_keep_line_breaks = true +ij_yaml_sequence_on_new_line = false +ij_yaml_space_before_colon = false +ij_yaml_spaces_within_braces = true +ij_yaml_spaces_within_brackets = true diff --git a/.gitignore b/.gitignore index 68bc17f9ff2104a9d7b6777058bb4c343ca72609..88155d14ac9d21d9128c8c53b227ab8369c9161c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,160 +1,22 @@ -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class +dist/* +*.egg-info +*.pyc +.tox -# C extensions -*.so - -# Distribution / packaging -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -share/python-wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.nox/ .coverage -.coverage.* -.cache -nosetests.xml coverage.xml -*.cover -*.py,cover -.hypothesis/ -.pytest_cache/ -cover/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py -db.sqlite3 -db.sqlite3-journal - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -.pybuilder/ -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# IPython -profile_default/ -ipython_config.py - -# pyenv -# For a library or package, you might want to ignore these files since the code is -# intended to run in multiple environments; otherwise, check them in: -# .python-version - -# pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -#Pipfile.lock - -# poetry -# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control -#poetry.lock - -# pdm -# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. -#pdm.lock -# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it -# in version control. -# https://pdm.fming.dev/#use-with-ide -.pdm.toml - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm -__pypackages__/ - -# Celery stuff -celerybeat-schedule -celerybeat.pid - -# SageMath parsed files -*.sage.py - -# Environments -.env -.venv -env/ -venv/ -ENV/ -env.bak/ -venv.bak/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy -.mypy_cache/ -.dmypy.json -dmypy.json - -# Pyre type checker -.pyre/ +htmlcov/* +build +dist -# pytype static type analyzer -.pytype/ +# Documentation +docs/source/source_documentation +!docs/source/source_documentation/index.rst +docs/build -# Cython debug symbols -cython_debug/ +# Setuptools SCM +lofty/_version.py -# PyCharm -# JetBrains specific template is maintained in a separate JetBrains.gitignore that can -# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore -# and can be added to the global gitignore or merged into this file. For a more nuclear -# option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +# IDE configuration +.vscode +.idea diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml new file mode 100644 index 0000000000000000000000000000000000000000..2b625b44bb02e8f677552d53165ca15551fea0ce --- /dev/null +++ b/.gitlab-ci.yml @@ -0,0 +1,203 @@ +default: + image: $CI_REGISTRY_IMAGE/ci-build-runner:$CI_COMMIT_REF_SLUG + before_script: + - python --version # For debugging + cache: + paths: + - .cache/pip + # Do not cache .tox, to recreate virtualenvs for every step + +stages: + - prepare + - lint + - test + - package + - integration + - publish # publish instead of deploy + +# Caching of dependencies to speed up builds +variables: + PIP_CACHE_DIR: "$CI_PROJECT_DIR/.cache/pip" + +include: + - template: Security/SAST.gitlab-ci.yml + - template: Security/Dependency-Scanning.gitlab-ci.yml + - template: Security/Secret-Detection.gitlab-ci.yml + +# Prepare image to run ci on +trigger_prepare: + stage: prepare + trigger: + strategy: depend + include: .prepare.gitlab-ci.yml + +run_black: + stage: lint + script: + - tox -e black + allow_failure: true + +run_flake8: + stage: lint + script: + - tox -e pep8 + allow_failure: true + +run_pylint: + stage: lint + script: + - tox -e pylint + allow_failure: true + +sast: + variables: + SAST_EXCLUDED_ANALYZERS: brakeman, flawfinder, kubesec, nodejs-scan, phpcs-security-audit, + pmd-apex, security-code-scan, sobelow, spotbugs + stage: test + +dependency_scanning: + # override default before_script, job won't have Python available + before_script: + - uname + +secret_detection: + # override default before_script, job won't have Python available + before_script: + - uname + +# Basic setup for all Python versions for which we don't have a base image +.run_unit_test_version_base: + before_script: + - python --version # For debugging + - python -m pip install --upgrade pip + - pip install --upgrade tox twine + +# Run all unit tests for Python versions except the base image +run_unit_tests: + extends: .run_unit_test_version_base + stage: test + image: python:3.${PY_VERSION} + script: + - tox -e py3${PY_VERSION} + parallel: + matrix: # use the matrix for testing + - PY_VERSION: [9, 10, 11] + +# Run code coverage on the base image thus also performing unit tests +run_unit_tests_coverage: + stage: test + script: + - tox -e coverage + coverage: '/(?i)total.*? (100(?:\.0+)?\%|[1-9]?\d(?:\.\d+)?\%)$/' + artifacts: + reports: + coverage_report: + coverage_format: cobertura + path: coverage.xml + paths: + - htmlcov/* + +package_files: + stage: package + artifacts: + expire_in: 1w + paths: + - dist/* + script: + - tox -e build + +package_docs: + stage: package + artifacts: + expire_in: 1w + paths: + - docs/build/* + script: + - tox -e docs + +run_integration_tests: + stage: integration + allow_failure: true + needs: + - package_files + script: + - echo "make sure to move out of source dir" + - echo "install package from filesystem (or use the artefact)" + - echo "run against foreign systems (e.g. databases, cwl etc.)" + - exit 1 + +publish_on_gitlab: + stage: publish + environment: gitlab + needs: + - package_files + when: manual + rules: + - if: $CI_COMMIT_TAG + script: + - echo "run twine for gitlab" + - | + TWINE_PASSWORD=${CI_JOB_TOKEN} \ + TWINE_USERNAME=gitlab-ci-token \ + python -m twine upload \ + --repository-url ${CI_API_V4_URL}/projects/${CI_PROJECT_ID}/packages/pypi dist/* + +publish_on_test_pypi: + stage: publish + environment: pypi-test + needs: + - package_files + when: manual + rules: + - if: '$CI_COMMIT_TAG && $CI_COMMIT_REF_PROTECTED == "true"' + script: + - echo "run twine for test pypi" + # - | + # TWINE_PASSWORD=${PIPY_TOKEN} \ + # TWINE_USERNAME=${PIPY_USERNAME} \ + # TODO: replace URL with a pipy URL + # python -m twine upload \ + # --repository-url ${CI_API_V4_URL}/projects/${CI_PROJECT_ID}/packages/pypi dist/* + - exit 1 + +publish_on_pypi: + stage: publish + environment: pypi + needs: + - package_files + when: manual + rules: + - if: '$CI_COMMIT_TAG && $CI_COMMIT_REF_PROTECTED == "true"' + script: + - echo "run twine for pypi" + # - | + # TWINE_PASSWORD=${PIPY_TOKEN} \ + # TWINE_USERNAME=${PIPY_USERNAME} \ + # TODO: replace URL with a pipy URL + # python -m twine upload \ + # --repository-url ${CI_API_V4_URL}/projects/${CI_PROJECT_ID}/packages/pypi dist/* + - exit 1 + +publish_to_readthedocs: + stage: publish + allow_failure: true + environment: readthedocs + needs: + - package_docs + when: manual + rules: + - if: '$CI_COMMIT_TAG && $CI_COMMIT_REF_PROTECTED == "true"' + script: + - echo "scp docs/* ???" + - exit 1 + +release_job: + stage: publish + image: registry.gitlab.com/gitlab-org/release-cli:latest + rules: + - if: '$CI_COMMIT_TAG && $CI_COMMIT_REF_PROTECTED == "true"' + script: + - echo "running release_job" + release: + tag_name: '$CI_COMMIT_TAG' + description: '$CI_COMMIT_TAG - $CI_COMMIT_TAG_MESSAGE' diff --git a/.prepare.gitlab-ci.yml b/.prepare.gitlab-ci.yml new file mode 100644 index 0000000000000000000000000000000000000000..3e48a271564faec892e42aab9f947d946ecc4d7b --- /dev/null +++ b/.prepare.gitlab-ci.yml @@ -0,0 +1,23 @@ +stages: + - build + +build_ci_runner_image: + stage: build + image: docker + tags: + - dind + script: + - docker login -u $CI_REGISTRY_USER -p $CI_REGISTRY_PASSWORD $CI_REGISTRY + - | + if docker pull $CI_REGISTRY_IMAGE/ci-build-runner:$CI_COMMIT_REF_SLUG; then + docker build --cache-from $CI_REGISTRY_IMAGE/ci-build-runner:$CI_COMMIT_REF_SLUG --tag $CI_REGISTRY_IMAGE/ci-build-runner:$CI_COMMIT_REF_SLUG docker/ci-runner + else + docker pull $CI_REGISTRY_IMAGE/ci-build-runner:latest || true + docker build --cache-from $CI_REGISTRY_IMAGE/ci-build-runner:latest --tag $CI_REGISTRY_IMAGE/ci-build-runner:$CI_COMMIT_REF_SLUG docker/ci-runner + fi + - docker push $CI_REGISTRY_IMAGE/ci-build-runner:$CI_COMMIT_REF_SLUG # push the image + - | + if [[ "$CI_COMMIT_BRANCH" == "$CI_DEFAULT_BRANCH" ]]; then + docker image tag $CI_REGISTRY_IMAGE/ci-build-runner:$CI_COMMIT_REF_SLUG $CI_REGISTRY_IMAGE/ci-build-runner:latest + docker push $CI_REGISTRY_IMAGE/ci-build-runner:latest + fi diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000000000000000000000000000000000000..c4a3399ac9b36be786467e020c84f95e09db0606 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,5 @@ +include LICENSE +include README.md + +recursive-include docs * +recursive-exclude tests * diff --git a/bstplot.py b/bstplot.py deleted file mode 100644 index f879927f24ab35f21e6c95d18dc7df92afa22cf2..0000000000000000000000000000000000000000 --- a/bstplot.py +++ /dev/null @@ -1,33 +0,0 @@ -#!/usr/bin/env python3 -import os -import glob -import argparse - -import numpy as np -import matplotlib.pyplot as plt -import matplotlib.dates as mdates -from matplotlib.ticker import FormatStrFormatter, MultipleLocator - -from lofty.io import read_bst, read_minio_bst -from lofty.plotting import plot_bst_timeseries, plot_bst_dynamic_spectrum - -if __name__ == "__main__": - # Read command line arguments - parser = argparse.ArgumentParser(description="Plot BST data") - parser.add_argument("-m", "--minio", help="Minio HDF5", action="store_true") - parser.add_argument("-d", "--dynspec", help="Plot dynamic spectrum", action="store_true") - parser.add_argument("-t", "--timeseries", help="Plot timeseries", action="store_true") - parser.add_argument("-s", "--subband", help="Subband to plot timeseries [int, default: 300]", default=300, type=int) - parser.add_argument("-p", "--pol", help="Polarization to plot dynamic spectrum [stokes/X/Y, default: stokes]", default="stokes") - parser.add_argument("filenames", help="SST filenames", nargs="*", metavar="FILE") - args = parser.parse_args() - - # Read data - if args.minio: - metadata, data = read_minio_bst(args.filenames) - else: - metadata, data = read_bst(args.filenames) - - plot_bst_dynamic_spectrum(metadata, data, args.pol) - - plot_bst_timeseries(metadata, data, args.subband) diff --git a/docs/cleanup.py b/docs/cleanup.py new file mode 100644 index 0000000000000000000000000000000000000000..95e5ac8f5bf802cd428722129a958b7916459818 --- /dev/null +++ b/docs/cleanup.py @@ -0,0 +1,23 @@ +#!/usr/bin/env python3 + +# Copyright (C) 2023 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import os + +file_dir = os.path.dirname(os.path.realpath(__file__)) + +clean_dir = os.path.join(file_dir, "source", "source_documentation") +print(f"Cleaning.. {clean_dir}/*") + +if not os.path.exists(clean_dir): + exit() + +for file_name in os.listdir(clean_dir): + file = os.path.join(clean_dir, file_name) + + if file_name == "index.rst": + continue + + print(f"Removing.. {file}") + os.remove(file) diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..3c6e46c6db7ddaf65e47cfa22c9ec0b914f7fd38 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,5 @@ +sphinx!=1.6.6,!=1.6.7,>=1.6.5 # BSD +sphinx-rtd-theme>=0.4.3 #MIT +sphinxcontrib-apidoc>=0.3.0 #BSD +myst-parser>=2.0 # MIT +docutils>=0.17 # BSD diff --git a/docs/source/conf.py b/docs/source/conf.py new file mode 100644 index 0000000000000000000000000000000000000000..40c7e9d3002c5dc0093d7426d47a00ea920e8fd9 --- /dev/null +++ b/docs/source/conf.py @@ -0,0 +1,95 @@ +# Copyright (C) 2023 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import os + +from lofty import __version__ + +# -- General configuration ---------------------------------------------------- + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom ones. +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.viewcode", + "sphinxcontrib.apidoc", + "sphinx_rtd_theme", + "myst_parser" +] + + + + +project_root_directory = os.getcwd() + +apidoc_module_dir = "../../lofty" +apidoc_output_dir = "source_documentation" +apidoc_excluded_paths = [] +apidoc_separate_modules = True +apidoc_toc_file = False +# This should include private methods but does not work +# https://github.com/sphinx-contrib/apidoc/issues/14 +apidoc_extra_args = ["--private"] + +# The suffix of source filenames. +source_suffix = [".rst"] + +# The master toctree document. +master_doc = "index" + +# General information about the project. +project = "lofty" +copyright = "2023, ASTRON" + +# openstackdocstheme options +repository_name = "git.astron.nl/lukken/lofty" +bug_project = "none" +bug_tag = "" +html_last_updated_fmt = "%Y-%m-%d %H:%M" + +# If true, '()' will be appended to :func: etc. cross-reference text. +add_function_parentheses = True + +version = __version__ + +modindex_common_prefix = ["lofty."] + +# If true, the current module name will be prepended to all description +# unit titles (such as .. function::). +add_module_names = True + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = "sphinx" + +# -- Options for HTML output -------------------------------------------------- + +# The theme to use for HTML and HTML Help pages. Major themes that come with +# Sphinx are currently 'default' and 'sphinxdoc'. +# html_theme_path = ["."] +html_theme = "sphinx_rtd_theme" +html_static_path = ["static"] +html_css_files = [ + "css/custom.css", +] + +# Output file base name for HTML help builder. +htmlhelp_basename = "%sdoc" % project + +# Conf.py variables exported to sphinx rst files access using |NAME| +variables_to_export = [ + "project", + "copyright", + "version", +] + +# Write to rst_epilog to export `variables_to_export` extract using `locals()` +# https://www.sphinx-doc.org/en/master/usage/configuration.html#confval-rst_epilog +frozen_locals = dict(locals()) +rst_epilog = "\n".join( + map( + lambda x: f".. |{x}| replace:: {frozen_locals[x]}", # noqa: F821 + variables_to_export, + ) +) +# Pep is not able to determine that frozen_locals always exists so noqa +del frozen_locals diff --git a/docs/source/index.rst b/docs/source/index.rst new file mode 100644 index 0000000000000000000000000000000000000000..12bebc64dd7c40223ad1bdb81901a150d4e06c83 --- /dev/null +++ b/docs/source/index.rst @@ -0,0 +1,16 @@ +==================================================== +Welcome to the documentation of lofty +==================================================== + +.. + To define more variables see rst_epilog generation in conf.py + +Documentation for version: |version| + +Contents: + +.. toctree:: + :maxdepth: 2 + + readme + source_documentation/index diff --git a/docs/source/readme.rst b/docs/source/readme.rst new file mode 100644 index 0000000000000000000000000000000000000000..87c96deef60fc04b35a8b9b9cca2f72b7e64e70c --- /dev/null +++ b/docs/source/readme.rst @@ -0,0 +1,2 @@ +.. include:: ../../README.md + :parser: myst_parser.sphinx_ diff --git a/docs/source/static/css/custom.css b/docs/source/static/css/custom.css new file mode 100644 index 0000000000000000000000000000000000000000..3ea8a2fc0c9c1ecb2b318cbc32edf90d2940c38d --- /dev/null +++ b/docs/source/static/css/custom.css @@ -0,0 +1,14 @@ +.orange { color: #c65d09; } + +.green { color: #5dc609; } + +.yellow { color: #c6c609; } + +.bolditalic { + font-weight: bold; + font-style: italic; +} + +.rst-content code, .rst-content tt, code { + white-space: break-spaces; +} diff --git a/lofty/__init__.py b/lofty/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..7c5373409232a73802f713c3af43511467d57a7f 100644 --- a/lofty/__init__.py +++ b/lofty/__init__.py @@ -0,0 +1,8 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +""" lofty """ + +from importlib import metadata + +__version__ = metadata.version("lofty") diff --git a/lofty/animate.py b/lofty/animate.py new file mode 100644 index 0000000000000000000000000000000000000000..d74c1da651a72fc28550ef08208b53d13b8854d7 --- /dev/null +++ b/lofty/animate.py @@ -0,0 +1,114 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# Copyright (C) 2024 Corne Lukken +# SPDX-License-Identifier: GPL-3.0-or-later + +import logging +import time +from typing import Callable + +import fastplotlib as fpl +import matplotlib.pyplot as plt +from matplotlib.patches import Circle +from nptyping import NDArray, Shape, Float64 + +from lofty.sources import lmn_sources + +logger = logging.getLogger() + + +def animate( + fn: Callable, + visibility_intervals: NDArray[Shape["Interval, Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + """ + :param fn: all-sky imaging function + :param visibility_intervals: Intervals with each 2d rectangular array of visibility + :param baselines: 3d array with u, v, w per antenna baseline (N^2) + :param freq: the frequency in hertz + :param npix_l: number of pixels length + :param npix_m: number of pixels height + """ + + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1) + circle = Circle((0, 0), 1.0, edgecolor="k", fill=False, facecolor="none", alpha=0.3) + + initial_data = fn(visibility_intervals[0], baselines, freq, npix_l, npix_m) + img = plt.imshow( + initial_data, cmap="hot", animated=True, origin="lower", extent=[1, -1, -1, 1] + ) + plt.colorbar(img, ax=ax, location="right", format="%.2e") + + ax.set_axis_off() + ax.add_artist(circle) + + # img.set_label("FPS: 0") + props = dict(boxstyle="round", facecolor="wheat", alpha=0.5) + # place a text box in upper left in axes coords + text = plt.text( + 0.05, + 0.95, + f"FPS: 0\nResulution: {npix_l}x{npix_m}", + fontsize=14, + verticalalignment="top", + bbox=props, + ) + + i = 0 + frames = 0 + start_time = time.time() - 0.1 + while True: + frames += 1 + fps = frames / (time.time() - start_time) + index = i % len(visibility_intervals) + data = fn(visibility_intervals[index], baselines, freq, npix_l, npix_m) + plt.pause(0.00001) + img.set_data(data) + text.set_text(f"FPS: {fps}\nResulution: {npix_l}x{npix_m}") + plt.draw() + plt.pause(0.00001) + i += 1 + + +def animate_glfw( + fn: Callable, + visibility_intervals: NDArray[Shape["Interval, Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + """ + :param fn: all-sky imaging function + :param visibility_intervals: Intervals with each 2d rectangular array of visibility + :param baselines: 3d array with u, v, w per antenna baseline (N^2) + :param freq: the frequency in hertz + :param npix_l: number of pixels length + :param npix_m: number of pixels height + """ + + fig = fpl.Figure() + initial_data = fn(visibility_intervals[0], baselines, freq, npix_l, npix_m) + graphic = fig[0, 0].add_image(data=initial_data, name="All-Sky image") + fig.show() + + def update_data(figure_instance, frames, i): + frames += 1 + start_time = time.time() + index = i % len(visibility_intervals) + data = fn(visibility_intervals[index], baselines, freq, npix_l, npix_m) + figure_instance[0, 0]["All-Sky image"].data = data + fps = frames / (time.time() - start_time) + logger.error(f"FPS: {fps}") + i += 1 + + i = 0 + frames = 0 + fig.add_animations(lambda update_instance: update_data(update_instance, frames, i)) + fig.show() + + fpl.run() diff --git a/lofty/data/__init__.py b/lofty/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lofty/data/baselines.npy b/lofty/data/baselines.npy new file mode 100644 index 0000000000000000000000000000000000000000..6d97890489f10600944095ef96003065ecfb2cf5 Binary files /dev/null and b/lofty/data/baselines.npy differ diff --git a/lofty/data/cs001lba_baselines.npy b/lofty/data/cs001lba_baselines.npy new file mode 100644 index 0000000000000000000000000000000000000000..eb96eb3126dcfc270a62736a2f5d4dbb565f22f7 Binary files /dev/null and b/lofty/data/cs001lba_baselines.npy differ diff --git a/lofty/data/cs001lba_visibilities.npy b/lofty/data/cs001lba_visibilities.npy new file mode 100644 index 0000000000000000000000000000000000000000..61add3a419d917db0029ee7b2691b3273a8d695c Binary files /dev/null and b/lofty/data/cs001lba_visibilities.npy differ diff --git a/lofty/data/freq.npy b/lofty/data/freq.npy new file mode 100644 index 0000000000000000000000000000000000000000..1ceb7a674078453f97f1c83a18d7f43f836f8bb1 Binary files /dev/null and b/lofty/data/freq.npy differ diff --git a/lofty/data/vis.npy b/lofty/data/vis.npy new file mode 100644 index 0000000000000000000000000000000000000000..0b5451b740ac0bbecb48f0cf133e4ba94de5f797 Binary files /dev/null and b/lofty/data/vis.npy differ diff --git a/lofty/data/visies.npy b/lofty/data/visies.npy new file mode 100644 index 0000000000000000000000000000000000000000..61add3a419d917db0029ee7b2691b3273a8d695c Binary files /dev/null and b/lofty/data/visies.npy differ diff --git a/lofty/entry/__init__.py b/lofty/entry/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lofty/entry/conversion.py b/lofty/entry/conversion.py new file mode 100644 index 0000000000000000000000000000000000000000..bd8b0f1d4af82463ef594436cdf4d807c5746ede --- /dev/null +++ b/lofty/entry/conversion.py @@ -0,0 +1,37 @@ +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import argparse + +import numpy as np +import lofarantpos + +from lofty.io import ( + read_single_npz, + extract_antenna_field_from_filename, + extract_station_name_from_filename, +) +from lofty.utility import get_station_xyz + + +def main(): + # Read command line arguments + parser = argparse.ArgumentParser(description="Extract correlated visibilities") + parser.add_argument("files", help="files", nargs="*", metavar="FILE") + args = parser.parse_args() + + import pdb + + pdb.set_trace() + for file in args.files: + station_name = extract_station_name_from_filename(file) + antenna_field = extract_antenna_field_from_filename(file) + metadata, data = read_single_npz(file) + + baselines, _ = get_station_xyz( + station_name, antenna_field, lofarantpos.db.LofarAntennaDatabase() + ) + + np.save("baselines.npy", baselines) + np.save("visibilities.npy", data) diff --git a/lofty/entry/imaging/__init__.py b/lofty/entry/imaging/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lofty/entry/imaging/animate.py b/lofty/entry/imaging/animate.py new file mode 100644 index 0000000000000000000000000000000000000000..e30ecf5fed751aff50e913b7c39f5d85a6bfac5f --- /dev/null +++ b/lofty/entry/imaging/animate.py @@ -0,0 +1,32 @@ +from importlib.resources import files + +import numpy as np + +from lofty.animate import animate +from lofty.imaging import ( + sky_imager_jax_ravel_real_jit, +) + + +def fold_polarizations(input): + cube_xx = input[:, 0::2, 0::2].astype("complex64") + cube_yy = input[:, 1::2, 1::2].astype("complex64") + return cube_xx + cube_yy + + +def main(): + baselines = np.load(files("lofty.data").joinpath("baselines.npy")) + visibilities = np.load(files("lofty.data").joinpath("vis.npy")) + frequency = np.load(files("lofty.data").joinpath("freq.npy")) + + if len(baselines.shape) < 3: + baselines = baselines[:, np.newaxis, :] - baselines[np.newaxis, :, :] + + # its a reasonable assumption that if we have both polarizations for the + # baselines we also still have them for the visibilities in the input data + visibilities = fold_polarizations(visibilities) + + if len(visibilities.shape) >= 2 and len(visibilities[1]) > 96: + visibilities = fold_polarizations(visibilities) + + animate(sky_imager_jax_ravel_real_jit, visibilities, baselines, frequency, 256, 256) diff --git a/xstimager.py b/lofty/entry/imaging/xst.py similarity index 64% rename from xstimager.py rename to lofty/entry/imaging/xst.py index cf268252b4461ad2e94f4c098e0186af6d8bf074..206aa3a1c29aec3c5f56080c090aafb3412df6a1 100644 --- a/xstimager.py +++ b/lofty/entry/imaging/xst.py @@ -1,21 +1,28 @@ -#!/usr/bin/env python3 -import os -import glob -import tqdm +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + import argparse +import tqdm import numpy as np import matplotlib.pyplot as plt - from lofarantpos.db import LofarAntennaDatabase -from lofty.io import read_xst, read_minio_xst +from lofty.io import read_npz +from lofty.io import read_minio_xst +from lofty.io import read_xst from lofty.imaging import sky_imager from lofty.sources import lmn_sources from lofty.plotting import sky_plot -from lofty.utils import get_station_type, get_full_station_name, freq_from_sb, get_station_xyz +from lofty.utility import ( + get_full_station_name, + frequency_from_subband, + get_station_xyz, +) + -if __name__ == "__main__": +def main(): # Read command line arguments parser = argparse.ArgumentParser(description="Create XST images") parser.add_argument("-m", "--minio", help="Minio HDF5", action="store_true") @@ -29,7 +36,7 @@ if __name__ == "__main__": metadata, data = read_xst(sorted(args.filenames), fill_empty=True) # Get information - station_type = get_station_type(metadata["station_name"]) + # station_type = get_station_type(metadata["station_name"]) antenna_set = metadata["antennafield"] station_name = get_full_station_name(metadata["station_name"], antenna_set) @@ -46,14 +53,14 @@ if __name__ == "__main__": # Apply mask to antennas and visibilities station_xyz = np.compress(mask, station_xyz, axis=0) data = np.compress(np.repeat(mask, 2), data, axis=1) - data = np.compress(np.repeat(mask, 2), data, axis=2) - + data = np.compress(np.repeat(mask, 2), data, axis=2) + # Generate visibilities cube_xx = data[:, 0::2, 0::2].astype("complex64") - cube_xy = data[:, 0::2, 1::2].astype("complex64") - cube_yx = data[:, 1::2, 0::2].astype("complex64") + # cube_xy = data[:, 0::2, 1::2].astype("complex64") + # cube_yx = data[:, 1::2, 0::2].astype("complex64") cube_yy = data[:, 1::2, 1::2].astype("complex64") - + # Baselines baselines = station_xyz[:, np.newaxis, :] - station_xyz[np.newaxis, :, :] @@ -63,29 +70,33 @@ if __name__ == "__main__": band = metadata["frequency_bands"][0][0].replace("LBA_", "").replace("HBA_", "") for isub in range(nsub): subband = metadata["subbands"][isub] - freqs[isub] = freq_from_sb(subband, band) + freqs[isub] = frequency_from_subband(subband, band) # Create imager - npix_l, npix_m = 205, 205 - l, m = np.meshgrid(np.linspace(-1, 1, npix_l), np.linspace(1, -1, npix_m)) + x, y = 205, 205 # Compute visibilities vis_i = cube_xx + cube_yy -# vis_q = cube_xx - cube_yy -# vis_u = cube_xy + cube_yx -# vis_v = 1j * (cube_xy - cube_yy) + # vis_q = cube_xx - cube_yy + # vis_u = cube_xy + cube_yx + # vis_v = 1j * (cube_xy - cube_yy) source_names = ["Cas A", "Cyg A", "Sun", "Tau A", "Moon", "Jupiter"] - lmn = lmn_sources(metadata["timestamps"], metadata["durations"], db.phase_centres[station_name], source_names) - + lmn = lmn_sources( + metadata["timestamps"], + metadata["durations"], + db.phase_centres[station_name], + source_names, + ) + # Compute images - img_i = np.zeros((nsub, npix_l, npix_m)) + img_i = np.zeros((nsub, x, y)) for isub in tqdm.tqdm(range(nsub)): - img_i[isub] = sky_imager(vis_i[isub], baselines, freqs[isub], l, m) + img_i[isub] = sky_imager(vis_i[isub], baselines, freqs[isub], x, y) for isub in range(nsub): - subband = metadata['subbands'][isub] - tstr = metadata['timestamps'][isub][:19] + subband = metadata["subbands"][isub] + tstr = metadata["timestamps"][isub][:19] freq = freqs[isub] * 1e-6 fname = f"{station_name}_{tstr}_SB{subband:03d}.png".replace(":", "_") @@ -95,22 +106,22 @@ if __name__ == "__main__": sky_plot(img_i[isub], title, subtitle, source_names, lmn[isub], fig) plt.savefig(fname) plt.close() - + + # for isub in range(nsub): # fig, ax = plt.subplots(figsize=(10, 10)) # ax.set_title(f"{metadata['timestamps'][isub]} {freqs[isub]}") # ax.imshow(img_i[isub], origin="lower") # plt.show() - - - # Write to FITS - #hdr = fits.Header() - - #hdu = fits.PrimaryHDU(data=img_i, header=hdr) - #hdu.writeto("stokes_i.fits", overwrite=True) - #hdu = fits.PrimaryHDU(data=img_q, header=hdr) - #hdu.writeto("stokes_q.fits", overwrite=True) - #hdu = fits.PrimaryHDU(data=img_u, header=hdr) - #hdu.writeto("stokes_u.fits", overwrite=True) - #hdu = fits.PrimaryHDU(data=img_v, header=hdr) - #hdu.writeto("stokes_v.fits", overwrite=True) + +# Write to FITS +# hdr = fits.Header() + +# hdu = fits.PrimaryHDU(data=img_i, header=hdr) +# hdu.writeto("stokes_i.fits", overwrite=True) +# hdu = fits.PrimaryHDU(data=img_q, header=hdr) +# hdu.writeto("stokes_q.fits", overwrite=True) +# hdu = fits.PrimaryHDU(data=img_u, header=hdr) +# hdu.writeto("stokes_u.fits", overwrite=True) +# hdu = fits.PrimaryHDU(data=img_v, header=hdr) +# hdu.writeto("stokes_v.fits", overwrite=True) diff --git a/lofty/entry/plot/__init__.py b/lofty/entry/plot/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lofty/entry/plot/bst.py b/lofty/entry/plot/bst.py new file mode 100644 index 0000000000000000000000000000000000000000..0bd44c04b3b5c36134093601f8544767ab758a25 --- /dev/null +++ b/lofty/entry/plot/bst.py @@ -0,0 +1,45 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import argparse + +from lofty.io import read_bst, read_minio_bst +from lofty.plotting import plot_bst_timeseries, plot_bst_dynamic_spectrum + + +def main(): + # Read command line arguments + parser = argparse.ArgumentParser(description="Plot BST data") + parser.add_argument("-m", "--minio", help="Minio HDF5", action="store_true") + parser.add_argument( + "-d", "--dynspec", help="Plot dynamic spectrum", action="store_true" + ) + parser.add_argument( + "-t", "--timeseries", help="Plot timeseries", action="store_true" + ) + parser.add_argument( + "-s", + "--subband", + help="Subband to plot timeseries [int, default: 300]", + default=300, + type=int, + ) + parser.add_argument( + "-p", + "--pol", + help="Polarization to plot dynamic spectrum [stokes/X/Y, default: stokes]", + default="stokes", + ) + parser.add_argument("filenames", help="SST filenames", nargs="*", metavar="FILE") + args = parser.parse_args() + + # Read data + if args.minio: + metadata, data = read_minio_bst(args.filenames) + else: + metadata, data = read_bst(args.filenames) + + plot_bst_dynamic_spectrum(metadata, data, args.pol) + + plot_bst_timeseries(metadata, data, args.subband) diff --git a/lofty/entry/plot/sst.py b/lofty/entry/plot/sst.py new file mode 100644 index 0000000000000000000000000000000000000000..724487cd8210ac1ad6228791a6145ab8e2120f97 --- /dev/null +++ b/lofty/entry/plot/sst.py @@ -0,0 +1,63 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import argparse + +from lofty.io import read_sst +from lofty.io import read_minio_sst +from lofty.plotting import ( + plot_sst_timeseries, + plot_sst_bandpass, + plot_sst_dynamic_spectrum, +) + + +def main(): + # Read command line arguments + parser = argparse.ArgumentParser(description="Plot SST data") + parser.add_argument("-m", "--minio", help="Minio HDF5", action="store_true") + parser.add_argument( + "-d", "--dynspec", help="Plot dynamic spectrum", action="store_true" + ) + parser.add_argument( + "-t", "--timeseries", help="Plot timeseries", action="store_true" + ) + parser.add_argument("-b", "--bandpass", help="Plot bandpass", action="store_true") + parser.add_argument( + "-s", + "--subband", + help="Subband to plot timeseries [int, default: 300]", + default=300, + type=int, + ) + parser.add_argument( + "-p", + "--pol", + help="Polarization to plot dynamic spectrum [stokes/X/Y, default: stokes]", + default="stokes", + ) + parser.add_argument( + "-a", + "--ant", + help="Antenna to plot dynamic spectrum [int (-1 for mean), default: -1]", + default=-1, + type=int, + ) + parser.add_argument("filenames", help="SST filenames", nargs="*", metavar="FILE") + args = parser.parse_args() + + print(args.filenames) + print(args.pol) + + # Read data + if args.minio: + metadata, data = read_minio_sst(args.filenames) + else: + metadata, data = read_sst(args.filenames) + + plot_sst_dynamic_spectrum(metadata, data, args.ant, args.pol) + + plot_sst_bandpass(metadata, data) + + plot_sst_timeseries(metadata, data, args.subband) diff --git a/xstplot.py b/lofty/entry/plot/xst.py similarity index 75% rename from xstplot.py rename to lofty/entry/plot/xst.py index dc8d1d85d0299c72d5212ab40927f34d4eb01e23..59588d8b82dbfad60f61183c68b0127fd0ba1786 100644 --- a/xstplot.py +++ b/lofty/entry/plot/xst.py @@ -1,17 +1,14 @@ -#!/usr/bin/env python3 -import os -import glob -import argparse +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np -import matplotlib.pyplot as plt -import matplotlib.dates as mdates -from matplotlib.ticker import FormatStrFormatter, MultipleLocator +import argparse from lofty.io import read_xst, read_minio_xst from lofty.plotting import plot_xst_matrix -if __name__ == "__main__": + +def main(): # Read command line arguments parser = argparse.ArgumentParser(description="Plot XST data") parser.add_argument("-m", "--minio", help="Minio HDF5", action="store_true") @@ -19,15 +16,15 @@ if __name__ == "__main__": args = parser.parse_args() print(args.filenames) - + # Read data if args.minio: metadata, data = read_minio_xst(sorted(args.filenames), fill_empty=False) else: - metadata, data = read_xst(sorted(args.filenames), fill_empty=False) + metadata, data = read_xst(sorted(args.filenames), fill_empty=False) print(metadata["subbands"]) print(metadata["nsub"]) - + for isub in range(data.shape[0]): plot_xst_matrix(metadata, data, isub) diff --git a/lofty/imaging.py b/lofty/imaging.py deleted file mode 100644 index c33a55be1c12d1a53e2bf34f3dd74cc2c6cdfcca..0000000000000000000000000000000000000000 --- a/lofty/imaging.py +++ /dev/null @@ -1,49 +0,0 @@ -#/usr/bin/env python3 -import numpy as np - -try: - import cupy as cp - - def sky_imager(visibilities, baselines, freq, l, m): - l, m = cp.array(l).astype("float32"), cp.array(m).astype("float32") - visibilities = cp.array(visibilities) - - # Select and ravel - c = l**2 + m**2 < 1 - lt, mt = l[c].ravel(), m[c].ravel() - nt = np.sqrt(1 - lt**2 - mt**2) - - u, v, w = cp.array(baselines).astype("float32").T - prod = (u[:, :, np.newaxis] * lt + - v[:, :, np.newaxis] * mt + - w[:, :, np.newaxis] * (nt - 1)) - phase = -2 * cp.pi * freq * prod / 299792458.0 - prod = None - pr, pi = cp.cos(phase), cp.sin(phase) - phase = None - vr, vi = cp.real(visibilities), cp.imag(visibilities) - img = cp.full(np.prod(l.shape), cp.nan, dtype="float32") - img[c.ravel()] = cp.mean(vr[:, :, cp.newaxis] * pr - vi[:, :, cp.newaxis] * pi, axis=(0, 1)) - - return cp.asnumpy(img.reshape(l.shape)) -except: - def sky_imager(visibilities, baselines, freq, l, m): - l, m = l.astype("float32"), m.astype("float32") - - # Select and ravel - c = l**2 + m**2 < 1 - lt, mt = l[c].ravel(), m[c].ravel() - nt = np.sqrt(1 - lt**2 - mt**2) - - u, v, w = baselines.astype("float32").T - prod = (u[:, :, np.newaxis] * lt + - v[:, :, np.newaxis] * mt + - w[:, :, np.newaxis] * (nt - 1)) - phase = -2 * np.pi * freq * prod / 299792458.0 - pr, pi = np.cos(phase), np.sin(phase) - vr, vi = np.real(visibilities), np.imag(visibilities) - img = np.full(np.prod(l.shape), np.nan, dtype="float32") - img[c.ravel()] = np.mean(vr[:, :, np.newaxis] * pr - vi[:, :, np.newaxis] * pi, axis=(0, 1)) - - return img.reshape(l.shape) - diff --git a/lofty/imaging/__init__.py b/lofty/imaging/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a199d1281fcbf918839205ab6a842c817993ce88 --- /dev/null +++ b/lofty/imaging/__init__.py @@ -0,0 +1,56 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from ._jax import sky_imager_jax +from ._jax import sky_imager_jax_jit +from ._jax import sky_imager_jax_ravel +from ._jax import sky_imager_jax_ravel_jit +from ._jax import sky_imager_jax_real +from ._jax import sky_imager_jax_real_jit +from ._jax import sky_imager_jax_ravel_real_jit +from ._jax import sky_imager_jax_precompute_real_jit +from ._numpy import sky_imager_numpy_ravel_32 +from ._numpy import sky_imager_numpy_ravel_real +from ._numpy import sky_imager_numpy_real +from ._numpy import sky_imager_numpy_exp +from ._simple import sky_imager_simple + +try: + from ._cupy import sky_imager_cupy # noqa + + __all__ = [ + "sky_imager_cupy", + "sky_imager_jax", + "sky_imager_jax_jit", + "sky_imager_jax_ravel", + "sky_imager_jax_ravel_jit", + "sky_imager_jax_real", + "sky_imager_jax_real_jit", + "sky_imager_jax_ravel_real_jit", + "sky_imager_jax_precompute_real_jit", + "sky_imager_numpy_exp", + "sky_imager_numpy_ravel_32", + "sky_imager_numpy_real", + "sky_imager_numpy_ravel_real", + "sky_imager_simple", + ] + + sky_imager = sky_imager_cupy +except ImportError: + __all__ = [ + "sky_imager_jax", + "sky_imager_jax_jit", + "sky_imager_jax_ravel", + "sky_imager_jax_ravel_jit", + "sky_imager_jax_real", + "sky_imager_jax_real_jit", + "sky_imager_jax_ravel_real_jit", + "sky_imager_jax_precompute_real_jit", + "sky_imager_numpy_exp", + "sky_imager_numpy_ravel_32", + "sky_imager_numpy_real", + "sky_imager_numpy_ravel_real", + "sky_imager_simple", + ] + + sky_imager = sky_imager_numpy_ravel_32 diff --git a/lofty/imaging/_constants.py b/lofty/imaging/_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..6f9a159227baae0fb086a1be6cd7bbaef68512f5 --- /dev/null +++ b/lofty/imaging/_constants.py @@ -0,0 +1,4 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +SPEED_OF_LIGHT = 299792458.0 diff --git a/lofty/imaging/_cupy.py b/lofty/imaging/_cupy.py new file mode 100644 index 0000000000000000000000000000000000000000..0aec82bc5f4e5d569a2f0a386a65c2991bfe8236 --- /dev/null +++ b/lofty/imaging/_cupy.py @@ -0,0 +1,63 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later +import copy + +try: + import cupy as cp + import numpy as np + + def sky_imager_cupy(visibilities, baselines, freq, npix_l, npix_m): + """ + :param visibilities: ? + :param baselines: ? + :param freq: ? + :param npix_l: Array of 2D normalized (-1, 1) pixel coordinates + :param npix_m: Array of 2D normalized (-1, 1) pixel coordinates + """ + cp.cuda.set_allocator(None) + cp.cuda.set_pinned_memory_allocator(None) + mempool = cp.get_default_memory_pool() + pinned_mempool = cp.get_default_pinned_memory_pool() + + npix_l, npix_m = np.meshgrid( + np.linspace(-1, 1, npix_l), np.linspace(1, -1, npix_m) + ) + npix_l, npix_m = npix_l.astype("float32"), npix_m.astype("float32") + l, m = cp.array(npix_l).astype("float32"), cp.array(npix_m).astype("float32") + visibilities = cp.array(visibilities) + freq = cp.array(freq) + + # Select and ravel + c = l**2 + m**2 < 1 + lt, mt = l[c].ravel(), m[c].ravel() + nt = np.sqrt(1 - lt**2 - mt**2) + del m + + u, v, w = cp.array(baselines).astype("float32").T + prod = ( + u[:, :, np.newaxis] * lt + + v[:, :, np.newaxis] * mt + + w[:, :, np.newaxis] * (nt - 1) + ) + del u, v, w, lt, mt, nt + phase = -2 * cp.pi * freq * prod / 299792458.0 + del prod + pr, pi = cp.cos(phase), cp.sin(phase) + del phase + vr, vi = cp.real(visibilities), cp.imag(visibilities) + del visibilities + img = cp.full(np.prod(l.shape), cp.nan, dtype="float32") + # img = cp.full(1, cp.nan, dtype="float32") + img[c.ravel()] = cp.mean( + vr[:, :, cp.newaxis] * pr - vi[:, :, cp.newaxis] * pi, axis=(0, 1) + ) + + result = cp.asnumpy(img.reshape(l.shape)) + del vr, vi, pr, pi + mempool.free_all_blocks() + pinned_mempool.free_all_blocks() + return result + +except ImportError: + pass diff --git a/lofty/imaging/_jax.py b/lofty/imaging/_jax.py new file mode 100644 index 0000000000000000000000000000000000000000..453b953b7407873ed2bc5df8d103964e1b9b50a2 --- /dev/null +++ b/lofty/imaging/_jax.py @@ -0,0 +1,304 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# Copyright (C) 2024 Corne Lukken +# SPDX-License-Identifier: GPL-3.0-or-later +import logging +from functools import lru_cache + +import jax +from matplotlib import pyplot as plt + +# import numpy as np +from nptyping import NDArray, Shape, Float64, Float32, Bool +from jax import numpy as jnp + +from ._constants import SPEED_OF_LIGHT + + +def sky_imager_jax( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + """ + :param visibilities: 2d rectangular array of visibilities + :param baselines: 3d array with u, v, w per antenna baseline (N^2) + :param freq: the frequency in hertz + :param npix_l: number of pixels length + :param npix_m: number of pixels height + :return: 2d image from the imaging process + """ + + # def compute(_visibilities, _baselines, _freq, _npix_l, _npix_m): + visibilities = jnp.array(visibilities) + freq = jnp.array(freq) + l, m = jnp.meshgrid(jnp.linspace(-1, 1, npix_l), jnp.linspace(1, -1, npix_m)) + n = jnp.sqrt(1 - l**2 - m**2) - 1 + + u, v, w = jnp.array(baselines).T + prod = ( + u[:, :, jnp.newaxis, jnp.newaxis] * l + + v[:, :, jnp.newaxis, jnp.newaxis] * m + + w[:, :, jnp.newaxis, jnp.newaxis] * n + ) + phasor = jnp.exp(-2j * jnp.pi * freq * prod / SPEED_OF_LIGHT) + + img = jnp.real( + jnp.mean(visibilities[:, :, jnp.newaxis, jnp.newaxis] * phasor, axis=(0, 1)) + ) + return jnp.real(img) + + +f_imager = jax.jit(sky_imager_jax, static_argnums=(3, 4)) + + +def sky_imager_jax_jit( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + # f_imager = jax.jit(sky_imager_jax, static_argnums=(3, 4)) + return f_imager(visibilities, baselines, freq, npix_l, npix_m) + + +def sky_imager_jax_ravel( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + visibilities = jnp.array(visibilities) + freq = jnp.array(freq) + + npix_l, npix_m = jnp.meshgrid( + jnp.linspace(-1, 1, npix_l), jnp.linspace(1, -1, npix_m) + ) + # Select and ravel + c = npix_l**2 + npix_m**2 < 1 # Create unit circle 2D image + lt, mt = npix_l[c].ravel(), npix_m[c].ravel() + nt = jnp.sqrt(1 - lt**2 - mt**2) + + u, v, w = jnp.array(baselines.astype("float32")).T + prod = ( + u[:, :, jnp.newaxis] * lt + + v[:, :, jnp.newaxis] * mt + + w[:, :, jnp.newaxis] * (nt - 1) + ) + phase = -2 * jnp.pi * freq * prod / SPEED_OF_LIGHT + pr, pi = jnp.cos(phase), jnp.sin(phase) + vr, vi = jnp.real(visibilities), jnp.imag(visibilities) + img = jnp.full(jnp.prod(jnp.array(npix_l.shape)), jnp.nan, dtype="float32") + img = img.at[c.ravel()].set( + jnp.mean(vr[:, :, jnp.newaxis] * pr - vi[:, :, jnp.newaxis] * pi, axis=(0, 1)) + ) + + return img.reshape(npix_l.shape) + + +def _sky_imager_jax_ravel_jit( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + lt: NDArray[Shape["Ravel"], Float32], + mt: NDArray[Shape["Ravel"], Float32], +): + visibilities = jnp.array(visibilities) + freq = jnp.array(freq) + nt = jnp.sqrt(1 - lt**2 - mt**2) + + u, v, w = jnp.array(baselines.astype("float32")).T + prod = ( + u[:, :, jnp.newaxis] * lt + + v[:, :, jnp.newaxis] * mt + + w[:, :, jnp.newaxis] * (nt - 1) + ) + phase = -2 * jnp.pi * freq * prod / SPEED_OF_LIGHT + pr, pi = jnp.cos(phase), jnp.sin(phase) + vr, vi = jnp.real(visibilities), jnp.imag(visibilities) + return jnp.mean( + vr[:, :, jnp.newaxis] * pr - vi[:, :, jnp.newaxis] * pi, axis=(0, 1) + ) + + +f_imager_ravel = jax.jit(_sky_imager_jax_ravel_jit) # , static_argnums=(3,)) + + +def sky_imager_jax_ravel_jit( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + npix_l, npix_m = jnp.meshgrid( + jnp.linspace(-1, 1, npix_l), jnp.linspace(1, -1, npix_m) + ) + c = npix_l**2 + npix_m**2 < 1 + lt, mt = npix_l[c].ravel(), npix_m[c].ravel() + img = jnp.full(jnp.prod(jnp.array(npix_l.shape)), jnp.nan, dtype="float32") + img = img.at[c.ravel()].set(f_imager_ravel(visibilities, baselines, freq, lt, mt)) + + return img.reshape(npix_l.shape) + + +def sky_imager_jax_real( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + """ + :param visibilities: 2d rectangular array of visibilities + :param baselines: 3d array with u, v, w per antenna baseline (N^2) + :param freq: the frequency in hertz + :param npix_l: number of pixels length + :param npix_m: number of pixels height + :return: 2d image from the imaging process + """ + + # def compute(_visibilities, _baselines, _freq, _npix_l, _npix_m): + visibilities = jnp.array(visibilities) + freq = jnp.array(freq) + l, m = jnp.meshgrid(jnp.linspace(-1, 1, npix_l), jnp.linspace(1, -1, npix_m)) + n = jnp.sqrt(1 - l**2 - m**2) - 1 + + u, v, w = jnp.array(baselines).T + prod = ( + u[:, :, jnp.newaxis, jnp.newaxis] * l + + v[:, :, jnp.newaxis, jnp.newaxis] * m + + w[:, :, jnp.newaxis, jnp.newaxis] * n + ) + phase = -2 * jnp.pi * freq * prod / SPEED_OF_LIGHT + vis = visibilities[:, :, jnp.newaxis, jnp.newaxis] + + return jnp.mean(jnp.cos(phase + jnp.angle(vis)) * jnp.abs(vis), axis=(0, 1)) + + +f_imager_real = jax.jit(sky_imager_jax_real, static_argnums=(3, 4)) + + +def sky_imager_jax_real_jit( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + return f_imager_real(visibilities, baselines, freq, npix_l, npix_m) + + +def _sky_imager_jax_ravel_real_jit( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + lt: NDArray[Shape["Ravel"], Float32], + mt: NDArray[Shape["Ravel"], Float32], +): + visibilities = jnp.array(visibilities) + freq = jnp.array(freq) + nt = jnp.sqrt(1 - lt**2 - mt**2) + + u, v, w = jnp.array(baselines.astype("float32")).T + prod = ( + u[:, :, jnp.newaxis] * lt + + v[:, :, jnp.newaxis] * mt + + w[:, :, jnp.newaxis] * (nt - 1) + ) + phase = -2 * jnp.pi * freq * prod / SPEED_OF_LIGHT + vis = visibilities[:, :, jnp.newaxis] + + return jnp.mean(jnp.cos(phase + jnp.angle(vis)) * jnp.abs(vis), axis=(0, 1)) + + +f_imager_ravel_real = jax.jit(_sky_imager_jax_ravel_real_jit) # , static_argnums=(3,)) + + +def sky_imager_jax_ravel_real_jit( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + npix_l, npix_m = jnp.meshgrid( + jnp.linspace(-1, 1, npix_l), jnp.linspace(1, -1, npix_m) + ) + c = npix_l**2 + npix_m**2 < 1 + lt, mt = npix_l[c].ravel(), npix_m[c].ravel() + img = jnp.full(jnp.prod(jnp.array(npix_l.shape)), jnp.nan, dtype="float32") + img = img.at[c.ravel()].set( + f_imager_ravel_real(visibilities, baselines, freq, lt, mt) + ) + + return img.reshape(npix_l.shape) + + +def _compute_phasors( + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + l: NDArray[Shape["Width, Height"], Float32], + m: NDArray[Shape["Width, Height"], Float32], + n: NDArray[Shape["Width, Height"], Float32], +): + u, v, w = jnp.array(baselines).T + prod = ( + u[:, :, jnp.newaxis, jnp.newaxis] * l + + v[:, :, jnp.newaxis, jnp.newaxis] * m + + w[:, :, jnp.newaxis, jnp.newaxis] * n + ) + return -2 * jnp.pi * freq * prod / SPEED_OF_LIGHT + + +def _sky_imager_jax_precompute_real( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + phasors: NDArray[Shape["Dim, Dim, Width, Height"], Float64], +): + vis = visibilities[:, :, jnp.newaxis, jnp.newaxis] + + return jnp.mean(jnp.cos(phasors + jnp.angle(vis)) * jnp.abs(vis), axis=(0, 1)) + + +f_imager_precompute_real = jax.jit( + _sky_imager_jax_precompute_real +) # , static_argnums=(3,)) +f_compute_phasors = jax.jit(_compute_phasors) # , static_argnums=(3,)) + + +def sky_imager_jax_precompute_real_jit( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + l, m = jnp.meshgrid(jnp.linspace(-1, 1, npix_l), jnp.linspace(1, -1, npix_m)) + n = jnp.sqrt(1 - l**2 - m**2) - 1 + phasors = sky_imager_jax_precompute_real_jit.phasors + base_hash = sky_imager_jax_precompute_real_jit.base_hash + + if phasors is None or ( + base_hash != baselines.__hash__ + or sky_imager_jax_precompute_real_jit.freq != freq[0] + or npix_l != sky_imager_jax_precompute_real_jit.width + or npix_m != sky_imager_jax_precompute_real_jit.height + ): + phasors = f_compute_phasors(baselines, freq, l, m, n) + sky_imager_jax_precompute_real_jit.width = npix_l + sky_imager_jax_precompute_real_jit.height = npix_m + sky_imager_jax_precompute_real_jit.phasors = phasors + sky_imager_jax_precompute_real_jit.base_hash = baselines.__hash__ + sky_imager_jax_precompute_real_jit.freq = freq[0] + + return f_imager_precompute_real(visibilities, phasors) + + +sky_imager_jax_precompute_real_jit.width = None +sky_imager_jax_precompute_real_jit.height = None +sky_imager_jax_precompute_real_jit.base_hash = None +sky_imager_jax_precompute_real_jit.freq = None +sky_imager_jax_precompute_real_jit.phasors = None diff --git a/lofty/imaging/_numpy.py b/lofty/imaging/_numpy.py new file mode 100644 index 0000000000000000000000000000000000000000..8f35800eb15118e521366a4868d1fa49934147d0 --- /dev/null +++ b/lofty/imaging/_numpy.py @@ -0,0 +1,149 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from nptyping import NDArray, Shape, Float64 + +from lofty.imaging._constants import SPEED_OF_LIGHT + + +def sky_imager_numpy_exp( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + """ + :param visibilities: 2d rectangular array of visibilities + :param baselines: 3d array with u, v, w per antenna baseline (N^2) + :param freq: the frequency in hertz + :param npix_l: number of pixels length + :param npix_m: number of pixels height + :return: 2d image from the imaging process + """ + l, m = np.meshgrid(np.linspace(-1, 1, npix_l), np.linspace(1, -1, npix_m)) + n = np.sqrt(1 - l**2 - m**2) - 1 + + u, v, w = baselines.T + prod = ( + u[:, :, np.newaxis, np.newaxis] * l + + v[:, :, np.newaxis, np.newaxis] * m + + w[:, :, np.newaxis, np.newaxis] * n + ) + phasor = np.exp(-2j * np.pi * freq * prod / SPEED_OF_LIGHT) + + img = np.real( + np.mean(visibilities[:, :, np.newaxis, np.newaxis] * phasor, axis=(0, 1)) + ) + + return img + + +def sky_imager_numpy_ravel_32( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + """ + :param visibilities: 2d rectangular array of visibilities + :param baselines: 3d array with u, v, w per antenna baseline (N^2) + :param freq: the frequency in hertz + :param npix_l: number of pixels length + :param npix_m: number of pixels height + :return: 2d image from the imaging process + """ + + npix_l, npix_m = np.meshgrid(np.linspace(-1, 1, npix_l), np.linspace(1, -1, npix_m)) + npix_l, npix_m = npix_l.astype("float32"), npix_m.astype("float32") + + # Select and ravel + c = npix_l**2 + npix_m**2 < 1 # Create unit circle 2D image + lt, mt = npix_l[c].ravel(), npix_m[c].ravel() + nt = np.sqrt(1 - lt**2 - mt**2) + + u, v, w = baselines.astype("float32").T + prod = ( + u[:, :, np.newaxis] * lt + + v[:, :, np.newaxis] * mt + + w[:, :, np.newaxis] * (nt - 1) + ) + phase = -2 * np.pi * freq * prod / SPEED_OF_LIGHT + pr, pi = np.cos(phase), np.sin(phase) + vr, vi = np.real(visibilities), np.imag(visibilities) + img = np.full(np.prod(npix_l.shape), np.nan, dtype="float32") + img[c.ravel()] = np.mean( + vr[:, :, np.newaxis] * pr - vi[:, :, np.newaxis] * pi, axis=(0, 1) + ) + + return img.reshape(npix_l.shape) + + +def sky_imager_numpy_real( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + l, m = np.meshgrid(np.linspace(-1, 1, npix_l), np.linspace(1, -1, npix_m)) + n = np.sqrt(1 - l**2 - m**2) - 1 + + u, v, w = baselines.T + prod = ( + u[:, :, np.newaxis, np.newaxis] * l + + v[:, :, np.newaxis, np.newaxis] * m + + w[:, :, np.newaxis, np.newaxis] * n + ) + phase = -2 * np.pi * freq * prod / SPEED_OF_LIGHT + vf = np.angle(visibilities[:, :, np.newaxis, np.newaxis]) + vr = np.abs(visibilities[:, :, np.newaxis, np.newaxis]) + pr = np.cos(phase + vf) * vr + + img = np.mean(pr, axis=(0, 1)) + + return img + + +def sky_imager_numpy_ravel_real( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + """ + :param visibilities: 2d rectangular array of visibilities + :param baselines: 3d array with u, v, w per antenna baseline (N^2) + :param freq: the frequency in hertz + :param npix_l: number of pixels length + :param npix_m: number of pixels height + :return: 2d image from the imaging process + """ + + npix_l, npix_m = np.meshgrid(np.linspace(-1, 1, npix_l), np.linspace(1, -1, npix_m)) + npix_l, npix_m = npix_l.astype("float32"), npix_m.astype("float32") + + # Select and ravel + c = npix_l**2 + npix_m**2 < 1 # Create unit circle 2D image + lt, mt = npix_l[c].ravel(), npix_m[c].ravel() + nt = np.sqrt(1 - lt**2 - mt**2) + + u, v, w = baselines.astype("float32").T + prod = ( + u[:, :, np.newaxis] * lt + + v[:, :, np.newaxis] * mt + + w[:, :, np.newaxis] * (nt - 1) + ) + phase = -2 * np.pi * freq * prod / SPEED_OF_LIGHT + vf = np.angle(visibilities[:, :, np.newaxis]) + vr = np.abs(visibilities[:, :, np.newaxis]) + pr = np.cos(phase + vf) * vr + + img = np.full(np.prod(npix_l.shape), np.nan, dtype="float32") + img[c.ravel()] = np.mean(pr, axis=(0, 1)) + + return img.reshape(npix_l.shape) diff --git a/lofty/imaging/_simple.py b/lofty/imaging/_simple.py new file mode 100644 index 0000000000000000000000000000000000000000..4d9a562884bf2fc62c55859f2fcea1ce7c292d0d --- /dev/null +++ b/lofty/imaging/_simple.py @@ -0,0 +1,191 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import numba +import numpy as np +from nptyping import NDArray, Shape, Float64 + +from ._constants import SPEED_OF_LIGHT + + +def sky_imager_simple( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + """ + :param visibilities: 2d rectangular array of visibilities + :param baselines: 3d array with u, v, w per antenna baseline (N^2) + :param freq: the frequency in hertz + :param npix_l: number of pixels length + :param npix_m: number of pixels height + :return: 2d image from the imaging process + """ + + img = np.zeros((npix_l, npix_m), dtype=np.complex128) + l, m = np.meshgrid(np.linspace(-1, 1, npix_l), np.linspace(1, -1, npix_m)) + n = np.sqrt(1 - l**2 - m**2) - 1 + + for l_ix in range(npix_l): + for m_ix in range(npix_m): + img[npix_l - l_ix - 1, npix_m - m_ix - 1] = np.mean( + visibilities + * np.exp( + -2j + * np.pi + * freq + * ( + baselines[:, :, 0] * l[l_ix, m_ix] + + baselines[:, :, 1] * m[l_ix, m_ix] + + baselines[:, :, 2] * n[l_ix, m_ix] + ) + / SPEED_OF_LIGHT + ) + ) + return np.real(img) + + +@numba.jit(parallel=True, fastmath=True) +def sky_imager_simple_numba( + visibilities: NDArray[Shape["Dim, Dim"], Float64], + baselines: NDArray[Shape["Dim, Dim, 3"], Float64], + freq: NDArray[Shape["1"], Float64], + npix_l: int, + npix_m: int, +): + """ + :param visibilities: 2d rectangular array of visibilities + :param baselines: 3d array with u, v, w per antenna baseline (N^2) + :param freq: the frequency in hertz + :param npix_l: number of pixels length + :param npix_m: number of pixels height + :return: 2d image from the imaging process + """ + + img = np.zeros((npix_m, npix_l), dtype=np.complex128) + + for m_ix in range(npix_m): + m = -1 + m_ix * 2 / npix_m + for l_ix in range(npix_l): + l = 1 - l_ix * 2 / npix_l + n = np.sqrt(1 - l * l - m * m) - 1 + img[m_ix, l_ix] = np.mean( + visibilities + * np.exp( + -2j + * np.pi + * freq + * ( + baselines[:, :, 0] * l + + baselines[:, :, 1] * m + + baselines[:, :, 2] * n + ) + / SPEED_OF_LIGHT + ) + ) + return np.real(img) + + +# +# +# def sky_imager_cupy(visibilities, baselines, freq, npix_l, npix_m): +# l, m = cp.meshgrid( +# cp.linspace(-1, 1, npix_l, dtype="float32"), +# cp.linspace(1, -1, npix_m, dtype="float32") +# ) +# n = cp.sqrt(1 - l**2 - m**2) - 1 +# +# vis = cp.array(visibilities) +# +# u, v, w = cp.array(baselines.astype("float32")).T +# prod = (u[:, :, cp.newaxis, cp.newaxis] * l + +# v[:, :, cp.newaxis, cp.newaxis] * m + +# w[:, :, cp.newaxis, cp.newaxis] * n).astype("complex64") +# phasor = cp.exp(-2j * cp.pi * freq * prod / SPEED_OF_LIGHT) +# prod = None +# +# img = cp.real(cp.mean(vis[:, :, cp.newaxis, cp.newaxis] * phasor, axis=(0, 1))) +# +# return cp.asnumpy(img) +# +# +# def sky_imager_numpy(visibilities, baselines, freq, npix_l, npix_m): +# l, m = np.meshgrid(np.linspace(-1, 1, npix_l), np.linspace(1, -1, npix_m)) +# n = np.sqrt(1 - l**2 - m**2) - 1 +# +# u, v, w = baselines.T +# prod = (u[:, :, np.newaxis, np.newaxis] * l + +# v[:, :, np.newaxis, np.newaxis] * m + +# w[:, :, np.newaxis, np.newaxis] * n) +# phasor = np.exp(-2j * np.pi * freq * prod / SPEED_OF_LIGHT) +# +# img = np.real( +# np.mean(visibilities[:, :, np.newaxis, np.newaxis] * phasor, axis=(0, 1)) +# ) +# return img +# +# +# def sky_imager_numpy_real(visibilities, baselines, freq, npix_l, npix_m): +# l, m = np.meshgrid(np.linspace(-1, 1, npix_l), np.linspace(1, -1, npix_m)) +# n = np.sqrt(1 - l**2 - m**2) - 1 +# +# u, v, w = baselines.T +# prod = (u[:, :, np.newaxis, np.newaxis] * l + +# v[:, :, np.newaxis, np.newaxis] * m + +# w[:, :, np.newaxis, np.newaxis] * n) +# phase = -2 * np.pi * freq * prod / SPEED_OF_LIGHT +# pr, pi = np.cos(phase), np.sin(phase) +# vr, vi = np.real(visibilities), np.imag(visibilities) +# img = np.mean( +# vr[:, :, np.newaxis, np.newaxis] * pr - vi[:, :, np.newaxis, +# np.newaxis] * pi, axis=(0, 1) +# ) +# +# return img +# +# +# def sky_imager_numpy_float32(visibilities, baselines, freq, npix_l, npix_m): +# l, m = np.meshgrid( +# np.linspace(-1, 1, npix_l).astype("float32"), +# np.linspace(1, -1, npix_m).astype("float32") +# ) +# n = np.sqrt(1 - l**2 - m**2) - 1 +# +# u, v, w = baselines.astype("float32").T +# prod = (u[:, :, np.newaxis, np.newaxis] * l + +# v[:, :, np.newaxis, np.newaxis] * m + +# w[:, :, np.newaxis, np.newaxis] * n) +# phasor = np.exp(-2j * np.pi * freq * prod / SPEED_OF_LIGHT) +# +# img = np.real( +# np.mean(visibilities[:, :, np.newaxis, np.newaxis] * phasor, axis=(0, 1)) +# ) +# return img +# +# +# def sky_imager_numpy_float32_ravel(visibilities, baselines, freq, npix_l, npix_m): +# l, m = np.meshgrid( +# np.linspace(-1, 1, npix_l).astype("float32"), +# np.linspace(1, -1, npix_m).astype("float32") +# ) +# +# # Select and ravel +# c = l**2 + m**2 < 1 +# l, m = l[c].ravel(), m[c].ravel() +# n = np.sqrt(1 - l**2 - m**2) - 1 +# +# u, v, w = baselines.astype("float32").T +# prod = (u[:, :, np.newaxis] * l + +# v[:, :, np.newaxis] * m + +# w[:, :, np.newaxis] * n) +# phasor = np.exp(-2j * np.pi * freq * prod / SPEED_OF_LIGHT) +# +# img = np.full(npix_l * npix_m, np.nan) +# img[c.ravel()] = np.real( +# np.mean(visibilities[:, :, np.newaxis] * phasor, axis=(0, 1)) +# ) +# +# return img.reshape(npix_l, npix_m) diff --git a/lofty/io.py b/lofty/io.py deleted file mode 100644 index db8963581e5450ee4bf4563d42d6d89cffda46b1..0000000000000000000000000000000000000000 --- a/lofty/io.py +++ /dev/null @@ -1,618 +0,0 @@ -#!/usr/bin/env python3 -import h5py - -import numpy as np - -from datetime import datetime - -def read_single_xst(xst_filename, fill_empty=False): - # Open HDF5 file - h5 = h5py.File(xst_filename, "r") - - # Store main header meta data - metadata = {} - for key in ["antenna_names", "antenna_quality", "antenna_type", "antenna_usage_mask", "station_name"]: - metadata[key] = h5["/"].attrs[key] - - # Antennafield - metadata["antennafield"] = h5["/"].attrs["antennafield_device"].replace("STAT/AFH/", "").replace("STAT/AFL/", "") - - # Number of subintegrations - nsub = len(h5.keys()) - metadata["nsub"] = nsub - - # Loop over keys - data_dict = [] - for key in h5.keys(): - # Keys in group - attribute_keys = list(h5[key].attrs.keys()) - - # Keys to store in dict - keys_to_store = ["frequency_band", "subbands", "timestamp", "integration_interval"] - - # Create dict - d = {"data": np.array(h5[key])} - for key_to_store in keys_to_store: - if key_to_store in attribute_keys: - d[key_to_store] = h5[key].attrs[key_to_store] - else: - print(f"Failed to find {key_to_store} in {key}") - d[key_to_store] = None - - # Store tile beam for HBA - if "HBA" in metadata["antennafield"]: - d["tile_beam_pointing_direction"] = h5[key].attrs["tile_beam_pointing_direction"] - - # Store dict - data_dict.append(d) - - # Store arrays - metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) - metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) - metadata["subbands"] = [data_dict[i]["subbands"][0] for i in range(nsub)] - - # Assume first entry holds for all (ASSUMPTION!!) - metadata["frequency_bands"] = data_dict[0]["frequency_band"] - if "HBA" in metadata["antennafield"]: - metadata["tile_beam_pointing_direction"] = data_dict[0]["tile_beam_pointing_direction"][0] - else: - metadata["tile_beam_pointing_direction"] = "" - - # Data shape - nrcu, _ = data_dict[0]["data"].shape - nant = nrcu // 2 - metadata["nrcu"] = nrcu - metadata["nant"] = nant - - # Store as nsub x nrcu x nrcu - data = np.zeros((nsub, nrcu, nrcu), dtype=data_dict[0]["data"].dtype) - for isub in range(nsub): - data[isub] = data_dict[isub]["data"] - - # Fill empty parts - if fill_empty: - data += np.conj(np.transpose(data, (0, 2, 1))) * (data==0) - - # Close file - h5.close() - - return metadata, data - -def read_xst(xst_filenames, fill_empty=False): - first = True - for xst_filename in xst_filenames: - # Get data - if first: - metadata, data = read_single_xst(xst_filename, fill_empty) - first = False - else: - newmetadata, newdata = read_single_xst(xst_filename, fill_empty) - - # Concatenate timestamps - metadata["timestamps"] = np.hstack((metadata["timestamps"], - newmetadata["timestamps"])) - - # Concatenate integration_interval - metadata["durations"] = np.hstack((metadata["durations"], - newmetadata["durations"])) - - # Concatenate subbands - metadata["subbands"] = np.hstack((metadata["subbands"], - newmetadata["subbands"])) - - # Concatenate data - data = np.vstack((data, newdata)) - - # Update nsub - metadata["nsub"] = data.shape[0] - - return metadata, data - - -def read_single_minio_xst(xst_filename, fill_empty=False): - # Open HDF5 file - h5 = h5py.File(xst_filename, "r") - - # Store main header meta data - metadata = {} - for key in ["antenna_names", "antenna_status", "antenna_type", "antenna_usage_mask", "station_name"]: - metadata[key] = h5["/"].attrs[key] - - # Antennafield - metadata["antennafield"] = h5["/"].attrs["antennafield_device"].upper().replace("STAT/AFH/", "").replace("STAT/AFL/", "") - - # Number of subintegrations - nsub = len(h5.keys()) - metadata["nsub"] = nsub - - # Loop over keys - data_dict = [] - for key in h5.keys(): - # Keys in group - attribute_keys = list(h5[key].attrs.keys()) - - # Keys to store in dict - keys_to_store = [] - - # Create dict - d = {"data": np.array(h5[key]["data"]).squeeze()} - for key_to_store in keys_to_store: - if key_to_store in attribute_keys: - d[key_to_store] = h5[key].attrs[key_to_store] - else: - print(f"Failed to find {key_to_store} in {key}") - d[key_to_store] = None - - # Add fake data - if "LBA" in metadata["antennafield"]: - d["frequency_band"] = [["LBA_10_90"] * 2] * 96 - else: - d["frequency_band"] = [["HBA_110_190"] * 2] * 48 - d["subbands"] = h5[key]["subband"].squeeze() - d["integration_interval"] = 1.0 - - # Timestamp - d["timestamp"] = key.replace("XST_", "") - - # Store dict - data_dict.append(d) - - # Close file - h5.close() - - # Store arrays - metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) - metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) - metadata["subbands"] = np.array([d["subbands"] for d in data_dict]) - - # Assume first entry holds for all (ASSUMPTION!!) - metadata["frequency_bands"] = data_dict[0]["frequency_band"] - if "HBA" in metadata["antennafield"]: - metadata["tile_beam_pointing_direction"] = data_dict[0]["tile_beam_pointing_direction"][0] - else: - metadata["tile_beam_pointing_direction"] = "" - - # Data shape - nrcu, _ = data_dict[0]["data"].shape - nant = nrcu // 2 - metadata["nrcu"] = nrcu - metadata["nant"] = nant - - # Store as nsub x nrcu x nrcu - data = np.zeros((nsub, nrcu, nrcu), dtype=data_dict[0]["data"].dtype) - for isub in range(nsub): - data[isub] = data_dict[isub]["data"] - - # Fill empty parts - if fill_empty: - data += np.conj(np.transpose(data, (0, 2, 1))) * (data==0) - - return metadata, data - -def read_minio_xst(xst_filenames, fill_empty=False): - first = True - for xst_filename in xst_filenames: - # Get data - if first: - metadata, data = read_single_minio_xst(xst_filename, fill_empty) - first = False - - metadata["station_name"] = metadata["station_name"].upper() - return metadata, data - - -def read_single_sst(sst_filename): - # Open HDF5 file - h5 = h5py.File(sst_filename, "r") - - # Store main header meta data - metadata = {} - for key in ["antenna_names", "antenna_quality", "antenna_type", "antenna_usage_mask", "station_name"]: - metadata[key] = h5["/"].attrs[key] - - # Antennafield - metadata["antennafield"] = h5["/"].attrs["antennafield_device"].replace("STAT/AFH/", "").replace("STAT/AFL/", "") - - # Number of subintegrations - nsub = len(h5.keys()) - metadata["nsub"] = nsub - - # Loop over keys - data_dict = [] - for key in h5.keys(): - # Keys in group - attribute_keys = list(h5[key].attrs.keys()) - - # Keys to store in dict - keys_to_store = ["frequency_band", "subbands", "timestamp", "integration_interval"] - - # Create dict - d = {"data": np.array(h5[key])} - for key_to_store in keys_to_store: - if key_to_store in attribute_keys: - d[key_to_store] = h5[key].attrs[key_to_store] - else: - print(f"Failed to find {key_to_store} in {key}") - d[key_to_store] = None - - # Store tile beam for HBA - if "HBA" in metadata["antennafield"]: - d["tile_beam_pointing_direction"] = h5[key].attrs["tile_beam_pointing_direction"] - - # Store dict - data_dict.append(d) - - # Store arrays - metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) - metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) - - # Assume first entry holds for all (ASSUMPTION!!) - metadata["subbands"] = data_dict[0]["subbands"] - metadata["frequency_bands"] = data_dict[0]["frequency_band"] - if "HBA" in metadata["antennafield"]: - metadata["tile_beam_pointing_direction"] = data_dict[0]["tile_beam_pointing_direction"][0] - else: - metadata["tile_beam_pointing_direction"] = "" - - # Data shape - nrcu, nchan = data_dict[0]["data"].shape - nant = nrcu // 2 - metadata["nrcu"] = nrcu - metadata["nant"] = nant - metadata["nchan"] = nchan - - # Store as nsub x npol x nant x nchan - data = np.zeros((nsub, 2, nant, nchan)) - for isub in range(nsub): - data[isub, 0] = data_dict[isub]["data"][0::2, :] - data[isub, 1] = data_dict[isub]["data"][1::2, :] - - # Close file - h5.close() - - return metadata, data - -def read_sst(sst_filenames): - first = True - for sst_filename in sst_filenames: - # Get data - if first: - metadata, data = read_single_sst(sst_filename) - first = False - else: - newmetadata, newdata = read_single_sst(sst_filename) - - # Concatenate timestamps - metadata["timestamps"] = np.hstack((metadata["timestamps"], - newmetadata["timestamps"])) - - # Concatenate integration_interval - metadata["durations"] = np.hstack((metadata["durations"], - newmetadata["durations"])) - - # Concatenate data - data = np.vstack((data, newdata)) - - return metadata, data - -def read_single_minio_sst(sst_filename): - # Open HDF5 file - h5 = h5py.File(sst_filename, "r") - - # Store main header meta data - metadata = {} - for key in ["antenna_names", "antenna_status", "antenna_type", "antenna_usage_mask", "station_name"]: - metadata[key] = h5["/"].attrs[key] - - # Antennafield - metadata["antennafield"] = h5["/"].attrs["antennafield_device"].upper().replace("STAT/AFH/", "").replace("STAT/AFL/", "") - - # Number of subintegrations - nsub = len(h5.keys()) - metadata["nsub"] = nsub - - # Loop over keys - data_dict = [] - for key in h5.keys(): - # Keys in group - attribute_keys = list(h5[key].attrs.keys()) - - # Keys to store in dict - keys_to_store = [] - - # Create dict - d = {"data": np.array(h5[key])} - for key_to_store in keys_to_store: - if key_to_store in attribute_keys: - d[key_to_store] = h5[key].attrs[key_to_store] - else: - print(f"Failed to find {key_to_store} in {key}") - d[key_to_store] = None - - # Add fake data - if "LBA" in metadata["antennafield"]: - d["frequency_band"] = [["LBA_10_90"] * 2] * 96 - else: - d["frequency_band"] = [["HBA_110_190"] * 2] * 48 - d["subbands"] = list(range(0, 512)) - d["integration_interval"] = 1.0 - - # Timestamp - d["timestamp"] = key.replace("SST_", "") - - # Store dict - data_dict.append(d) - - # Close file - h5.close() - - # Store arrays - metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) - metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) - - # Assume first entry holds for all (ASSUMPTION!!) - metadata["subbands"] = data_dict[0]["subbands"] - metadata["frequency_bands"] = data_dict[0]["frequency_band"] - if "HBA" in metadata["antennafield"]: - metadata["tile_beam_pointing_direction"] = data_dict[0]["tile_beam_pointing_direction"][0] - else: - metadata["tile_beam_pointing_direction"] = "" - - # Data shape - nrcu, nchan = data_dict[0]["data"].shape - nant = nrcu // 2 - metadata["nrcu"] = nrcu - metadata["nant"] = nant - metadata["nchan"] = nchan - - # Store as nsub x npol x nant x nchan - data = np.zeros((nsub, 2, nant, nchan)) - for isub in range(nsub): - data[isub, 0] = data_dict[isub]["data"][0::2, :] - data[isub, 1] = data_dict[isub]["data"][1::2, :] - - return metadata, data - -def read_minio_sst(sst_filenames): - first = True - for sst_filename in sst_filenames: - # Get data - if first: - metadata, data = read_single_minio_sst(sst_filename) - first = False - - return metadata, data - - -def read_single_bst(bst_filename): - # Open HDF5 file - h5 = h5py.File(bst_filename, "r") - - # Store main header meta data - metadata = {} - for key in ["antenna_names", "antenna_quality", "antenna_type", "antenna_usage_mask", "station_name"]: - metadata[key] = h5["/"].attrs[key] - - # Antennafield - metadata["antennafield"] = h5["/"].attrs["antennafield_device"].replace("STAT/AFH/", "").replace("STAT/AFL/", "") - - # Number of subintegrations - nsub = len(h5.keys()) - metadata["nsub"] = nsub - - # Loop over keys - data_dict = [] - for key in h5.keys(): - # Keys in group - attribute_keys = list(h5[key].attrs.keys()) - - # Keys to store in dict - keys_to_store = ["frequency_band", "subbands", "timestamp", "integration_interval"] - - # Create dict - d = {"data": np.array(h5[key])} - for key_to_store in keys_to_store: - if key_to_store in attribute_keys: - d[key_to_store] = h5[key].attrs[key_to_store] - else: - print(f"Failed to find {key_to_store} in {key}") - d[key_to_store] = None - - # Store tile beam for HBA - if "HBA" in metadata["antennafield"]: - d["tile_beam_pointing_direction"] = h5[key].attrs["tile_beam_pointing_direction"] - - # Store dict - data_dict.append(d) - - # Store arrays - metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) - metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) - - # Assume first entry holds for all (ASSUMPTION!!) - metadata["subbands"] = data_dict[0]["subbands"] - metadata["frequency_bands"] = data_dict[0]["frequency_band"] - if "HBA" in metadata["antennafield"]: - metadata["tile_beam_pointing_direction"] = data_dict[0]["tile_beam_pointing_direction"][0] - else: - metadata["tile_beam_pointing_direction"] = "" - - # Data shape - nchan, npol = data_dict[0]["data"].shape - metadata["nchan"] = nchan - - # Store as nsub x npol x nchan - data = np.zeros((nsub, 2, nchan)) - for isub in range(nsub): - data[isub, 0] = data_dict[isub]["data"][:, 0::2].T - data[isub, 1] = data_dict[isub]["data"][:, 1::2].T - - # Close file - h5.close() - - return metadata, data - - -def read_bst(bst_filenames): - first = True - for bst_filename in bst_filenames: - # Get data - if first: - metadata, data = read_single_bst(bst_filename) - first = False - else: - newmetadata, newdata = read_single_bst(bst_filename) - - # Concatenate timestamps - metadata["timestamps"] = np.hstack((metadata["timestamps"], - newmetadata["timestamps"])) - - # Concatenate integration_interval - metadata["durations"] = np.hstack((metadata["durations"], - newmetadata["durations"])) - - # Concatenate data - data = np.vstack((data, newdata)) - - return metadata, data - -def read_single_minio_bst(bst_filename): - # Open HDF5 file - h5 = h5py.File(bst_filename, "r") - - # Store main header meta data - metadata = {} - for key in ["antenna_names", "antenna_status", "antenna_type", "antenna_usage_mask", "station_name"]: - metadata[key] = h5["/"].attrs[key] - - # Antennafield - metadata["antennafield"] = h5["/"].attrs["antennafield_device"].upper().replace("STAT/AFH/", "").replace("STAT/AFL/", "") - - # Number of subintegrations - nsub = len(h5.keys()) - metadata["nsub"] = nsub - - # Loop over keys - data_dict = [] - for key in h5.keys(): - # Keys in group - attribute_keys = list(h5[key].attrs.keys()) - - # Keys to store in dict - keys_to_store = [] - - # Create dict - d = {"data": np.array(h5[key])} - for key_to_store in keys_to_store: - if key_to_store in attribute_keys: - d[key_to_store] = h5[key].attrs[key_to_store] - else: - print(f"Failed to find {key_to_store} in {key}") - d[key_to_store] = None - - # Add fake data - if "LBA" in metadata["antennafield"]: - d["frequency_band"] = [["LBA_10_90"] * 2] * 96 - else: - d["frequency_band"] = [["HBA_110_190"] * 2] * 48 - d["subbands"] = list(range(12, 500)) - d["integration_interval"] = 1.0 - - # Timestamp - d["timestamp"] = key.replace("BST_", "") - - # Store dict - data_dict.append(d) - - # Close file - h5.close() - - # Store arrays - metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) - metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) - - # Assume first entry holds for all (ASSUMPTION!!) - metadata["subbands"] = data_dict[0]["subbands"] - metadata["frequency_bands"] = data_dict[0]["frequency_band"] - if "HBA" in metadata["antennafield"]: - metadata["tile_beam_pointing_direction"] = data_dict[0]["tile_beam_pointing_direction"][0] - else: - metadata["tile_beam_pointing_direction"] = "" - - # Data shape - nchan, npol = data_dict[0]["data"].shape - metadata["nchan"] = nchan - - # Store as nsub x npol x nchan - data = np.zeros((nsub, 2, nchan)) - for isub in range(nsub): - data[isub, 0] = data_dict[isub]["data"][:, 0::2].T - data[isub, 1] = data_dict[isub]["data"][:, 1::2].T - - return metadata, data - -def read_minio_bst(bst_filenames): - first = True - for bst_filename in bst_filenames: - # Get data - if first: - metadata, data = read_single_minio_bst(bst_filename) - first = False - - return metadata, data - - -def read_acm_cube(xst_filename, fill=True): - # Open file - h5 = h5py.File(xst_filename, "r") - - writer_version = h5.attrs["writer_version"] - - # Number of sub integrations - nsub = len(h5.keys()) - - # Loop over sub integrations - timestrings = [] - for isub, key in enumerate(h5.keys()): - # Read data - data = np.array(h5[key]) - - # Create cube - if isub == 0: - nant = data.shape[0] - cube = np.zeros((nsub, nant, nant), dtype=data.dtype) - subbands = np.zeros(nsub, dtype="int64") - durations = np.zeros(nsub, dtype="float32") - - # Store - cube[isub] = data - - # Read attributes - timestrings.append(h5[key].attrs["timestamp"]) - subbands[isub] = h5[key].attrs["data_id_subband_index"] - durations[isub] = h5[key].attrs["integration_interval"] - - # Decode timestamps - if writer_version == "0.3": - timestamps = [datetime.strptime(timestring, "%Y-%m-%dT%H:%M:%S.%f%z") for timestring in timestrings] - else: - timestamps = [datetime.strptime(timestring, "%Y-%m-%dT%H:%M:%S.%f") for timestring in timestrings] - - #if np.all(np.diagonal(cube[0, 48:96, 48:96]) == 0 + 0j): - # # Move block to right location - # cube[:, 48:96, 48:96] = cube[:, 48:96, 0:48] - # # Zero original block - # cube[:, 48:96, 0:48] = 0 + 0j - - # Redimension HBA data to 96x96 - if np.all(np.diagonal(cube[0, 96:, 96:]) == 0 + 0j): - cube = cube[:, :96, :96] - - # Fill empty parts - if fill: - cube += np.conj(np.transpose(cube, (0, 2, 1))) * (cube==0) - - # Close file - h5.close() - - return cube, timestamps, subbands, durations diff --git a/lofty/io/__init__.py b/lofty/io/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..72cc3236fec88fdfeb76a287b432c95166981e10 --- /dev/null +++ b/lofty/io/__init__.py @@ -0,0 +1,46 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from ._read_acm_cube import read_acm_cube + +from ._read_single_bst import read_single_bst +from ._read_single_npz import read_single_npz +from ._read_single_sst import read_single_sst +from ._read_single_xst import read_single_xst + +from ._read_single_minio_bst import read_single_minio_bst +from ._read_single_minio_sst import read_single_minio_sst +from ._read_single_minio_xst import read_single_minio_xst + +from ._read_bst import read_bst +from ._read_npz import read_npz +from ._read_sst import read_sst +from ._read_xst import read_xst + +from ._read_minio_bst import read_minio_bst +from ._read_minio_sst import read_minio_sst +from ._read_minio_xst import read_minio_xst + +from ._station_name import extract_station_name_from_filename +from ._antenna_field import extract_antenna_field_from_filename + + +__all__ = [ + "read_acm_cube", + "read_single_bst", + "read_single_npz", + "read_single_sst", + "read_single_xst", + "read_single_minio_bst", + "read_single_minio_sst", + "read_single_minio_xst", + "read_bst", + "read_npz", + "read_sst", + "read_xst", + "read_minio_bst", + "read_minio_sst", + "read_minio_xst", + "extract_antenna_field_from_filename", + "extract_station_name_from_filename", +] diff --git a/lofty/io/_antenna_field.py b/lofty/io/_antenna_field.py new file mode 100644 index 0000000000000000000000000000000000000000..44542872eeaf939623bc30d38a7ba26f6920958f --- /dev/null +++ b/lofty/io/_antenna_field.py @@ -0,0 +1,28 @@ +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import logging +import re +from typing import Optional + +logger = logging.getLogger() + +ANTENNA_FIELD_NAME_REGEX = "(?i)lba|hba" + + +def extract_antenna_field_from_filename(fname: str) -> Optional[str]: + """Extract the antenna field from the filename""" + + pattern = re.compile(ANTENNA_FIELD_NAME_REGEX) + fields = pattern.findall(fname) + + if len(fields) > 1: + logger.warning( + "Found multiple antenna field: %s matches in file: %s", fields, fname + ) + return None + if len(fields) > 0: + return fields[0] + + return None diff --git a/lofty/io/_read_acm_cube.py b/lofty/io/_read_acm_cube.py new file mode 100644 index 0000000000000000000000000000000000000000..1d3753abaed9a618b2ac822fd756f28df5b524c8 --- /dev/null +++ b/lofty/io/_read_acm_cube.py @@ -0,0 +1,71 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from datetime import datetime + +import numpy as np +import h5py + + +def read_acm_cube(xst_filename, fill=True): + # Open file + h5 = h5py.File(xst_filename, "r") + + writer_version = h5.attrs["writer_version"] + + # Number of sub integrations + nsub = len(h5.keys()) + + # Loop over sub integrations + timestrings = [] + for isub, key in enumerate(h5.keys()): + # Read data + data = np.array(h5[key]) + + # Create cube + if isub == 0: + nant = data.shape[0] + cube = np.zeros((nsub, nant, nant), dtype=data.dtype) + subbands = np.zeros(nsub, dtype="int64") + durations = np.zeros(nsub, dtype="float32") + + # Store + cube[isub] = data + + # Read attributes + timestrings.append(h5[key].attrs["timestamp"]) + subbands[isub] = h5[key].attrs["data_id_subband_index"] + durations[isub] = h5[key].attrs["integration_interval"] + + # Decode timestamps + if writer_version == "0.3": + timestamps = [ + datetime.strptime(timestring, "%Y-%m-%dT%H:%M:%S.%f%z") + for timestring in timestrings + ] + else: + timestamps = [ + datetime.strptime(timestring, "%Y-%m-%dT%H:%M:%S.%f") + for timestring in timestrings + ] + + # if np.all(np.diagonal(cube[0, 48:96, 48:96]) == 0 + 0j): + # # Move block to right location + # cube[:, 48:96, 48:96] = cube[:, 48:96, 0:48] + # # Zero original block + # cube[:, 48:96, 0:48] = 0 + 0j + + # Redimension HBA data to 96x96 + if np.all(np.diagonal(cube[0, 96:, 96:]) == 0 + 0j): + cube = cube[:, :96, :96] + + # Fill empty parts + if fill: + cube += np.conj(np.transpose(cube, (0, 2, 1))) * (cube == 0) + + # Close file + h5.close() + + return cube, timestamps, subbands, durations diff --git a/lofty/io/_read_bst.py b/lofty/io/_read_bst.py new file mode 100644 index 0000000000000000000000000000000000000000..441d775c1ab284b21d6df1a4d19c2c2b1fbe27ed --- /dev/null +++ b/lofty/io/_read_bst.py @@ -0,0 +1,34 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from lofty.io._read_single_bst import read_single_bst + + +def read_bst(bst_filenames): + first = True + for bst_filename in bst_filenames: + # Get data + if first: + metadata, data = read_single_bst(bst_filename) + first = False + else: + newmetadata, newdata = read_single_bst(bst_filename) + + # Concatenate timestamps + metadata["timestamps"] = np.hstack( + (metadata["timestamps"], newmetadata["timestamps"]) + ) + + # Concatenate integration_interval + metadata["durations"] = np.hstack( + (metadata["durations"], newmetadata["durations"]) + ) + + # Concatenate data + data = np.vstack((data, newdata)) + + return metadata, data diff --git a/lofty/io/_read_minio_bst.py b/lofty/io/_read_minio_bst.py new file mode 100644 index 0000000000000000000000000000000000000000..4a02f6f1e97bbdacf73b782c0c605ed618276618 --- /dev/null +++ b/lofty/io/_read_minio_bst.py @@ -0,0 +1,16 @@ +# Copyright (C) 2024 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from lofty.io._read_single_minio_bst import read_single_minio_bst + + +def read_minio_bst(bst_filenames): + first = True + for bst_filename in bst_filenames: + # Get data + if first: + metadata, data = read_single_minio_bst(bst_filename) + first = False + + return metadata, data diff --git a/lofty/io/_read_minio_sst.py b/lofty/io/_read_minio_sst.py new file mode 100644 index 0000000000000000000000000000000000000000..4bb916017431e3fb937b88d4e3ff13072a88f200 --- /dev/null +++ b/lofty/io/_read_minio_sst.py @@ -0,0 +1,16 @@ +# Copyright (C) 2024 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from lofty.io._read_single_minio_sst import read_single_minio_sst + + +def read_minio_sst(sst_filenames): + first = True + for sst_filename in sst_filenames: + # Get data + if first: + metadata, data = read_single_minio_sst(sst_filename) + first = False + + return metadata, data diff --git a/lofty/io/_read_minio_xst.py b/lofty/io/_read_minio_xst.py new file mode 100644 index 0000000000000000000000000000000000000000..0c810991c78e747eaf13099eb413cd2ac9a3400c --- /dev/null +++ b/lofty/io/_read_minio_xst.py @@ -0,0 +1,17 @@ +# Copyright (C) 2024 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from lofty.io._read_single_minio_xst import read_single_minio_xst + + +def read_minio_xst(xst_filenames, fill_empty=False): + first = True + for xst_filename in xst_filenames: + # Get data + if first: + metadata, data = read_single_minio_xst(xst_filename, fill_empty) + first = False + + metadata["station_name"] = metadata["station_name"].upper() + return metadata, data diff --git a/lofty/io/_read_npz.py b/lofty/io/_read_npz.py new file mode 100644 index 0000000000000000000000000000000000000000..068ddc86e3413f815d35bfdedf751df97a3acc0d --- /dev/null +++ b/lofty/io/_read_npz.py @@ -0,0 +1,42 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from lofty.io._read_single_npz import read_single_npz + + +def read_npz(xst_filenames): + first = True + for xst_filename in xst_filenames: + # Get data + if first: + metadata, data = read_single_npz(xst_filename) + first = False + else: + newmetadata, newdata = read_single_npz(xst_filename) + + # Concatenate timestamps + metadata["timestamps"] = np.hstack( + (metadata["timestamps"], newmetadata["timestamps"]) + ) + + # Concatenate integration_interval + metadata["durations"] = np.hstack( + (metadata["durations"], newmetadata["durations"]) + ) + + # Concatenate subbands + metadata["subbands"] = np.hstack( + (metadata["subbands"], newmetadata["subbands"]) + ) + + # Concatenate data + data = np.vstack((data, newdata)) + + # Update nsub + metadata["nsub"] = data.shape[0] + + return metadata, data diff --git a/lofty/io/_read_single_bst.py b/lofty/io/_read_single_bst.py new file mode 100644 index 0000000000000000000000000000000000000000..4febdff5afbdac9cc634ed469a9cf9ade256ed4b --- /dev/null +++ b/lofty/io/_read_single_bst.py @@ -0,0 +1,93 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +import h5py + + +def read_single_bst(bst_filename): + # Open HDF5 file + h5 = h5py.File(bst_filename, "r") + + # Store main header meta data + metadata = {} + # for key in ["antenna_names", "antenna_quality", "antenna_type", "antenna_usage_mask", "station_name"]: + # metadata[key] = h5["/"].attrs[key] + for key in ["antenna_names", "antenna_type", "station_name"]: + metadata[key] = h5["/"].attrs[key] + + # Antennafield + metadata["antennafield"] = ( + h5["/"] + .attrs["antennafield_device"] + .replace("STAT/AFH/", "") + .replace("STAT/AFL/", "") + ) + + # Number of subintegrations + nsub = len(h5.keys()) + metadata["nsub"] = nsub + + # Loop over keys + data_dict = [] + for key in h5.keys(): + # Keys in group + attribute_keys = list(h5[key].attrs.keys()) + + # Keys to store in dict + keys_to_store = [ + "frequency_band", + "subbands", + "timestamp", + "integration_interval", + ] + + # Create dict + d = {"data": np.array(h5[key])} + for key_to_store in keys_to_store: + if key_to_store in attribute_keys: + d[key_to_store] = h5[key].attrs[key_to_store] + else: + print(f"Failed to find {key_to_store} in {key}") + d[key_to_store] = None + + # Store tile beam for HBA + if "HBA" in metadata["antennafield"]: + d["tile_beam_pointing_direction"] = h5[key].attrs[ + "tile_beam_pointing_direction" + ] + + # Store dict + data_dict.append(d) + + # Store arrays + metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) + metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) + + # Assume first entry holds for all (ASSUMPTION!!) + metadata["subbands"] = data_dict[0]["subbands"] + metadata["frequency_bands"] = data_dict[0]["frequency_band"] + if "HBA" in metadata["antennafield"]: + metadata["tile_beam_pointing_direction"] = data_dict[0][ + "tile_beam_pointing_direction" + ][0] + else: + metadata["tile_beam_pointing_direction"] = "" + + # Data shape + nchan, npol = data_dict[0]["data"].shape + metadata["nchan"] = nchan + + # Store as nsub x npol x nchan + data = np.zeros((nsub, 2, nchan)) + for isub in range(nsub): + data[isub, 0] = data_dict[isub]["data"][:, 0::2].T + data[isub, 1] = data_dict[isub]["data"][:, 1::2].T + + # Close file + h5.close() + + return metadata, data diff --git a/lofty/io/_read_single_minio_bst.py b/lofty/io/_read_single_minio_bst.py new file mode 100644 index 0000000000000000000000000000000000000000..ad5367f765c78ae8538d43ce382cd88af58c19f5 --- /dev/null +++ b/lofty/io/_read_single_minio_bst.py @@ -0,0 +1,96 @@ +# Copyright (C) 2024 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import h5py +import numpy as np + + +def read_single_minio_bst(bst_filename): + # Open HDF5 file + h5 = h5py.File(bst_filename, "r") + + # Store main header meta data + metadata = {} + for key in [ + "antenna_names", + "antenna_status", + "antenna_type", + "antenna_usage_mask", + "station_name", + ]: + metadata[key] = h5["/"].attrs[key] + + # Antennafield + metadata["antennafield"] = ( + h5["/"] + .attrs["antennafield_device"] + .upper() + .replace("STAT/AFH/", "") + .replace("STAT/AFL/", "") + ) + + # Number of subintegrations + nsub = len(h5.keys()) + metadata["nsub"] = nsub + + # Loop over keys + data_dict = [] + for key in h5.keys(): + # Keys in group + attribute_keys = list(h5[key].attrs.keys()) + + # Keys to store in dict + keys_to_store = [] + + # Create dict + d = {"data": np.array(h5[key])} + for key_to_store in keys_to_store: + if key_to_store in attribute_keys: + d[key_to_store] = h5[key].attrs[key_to_store] + else: + print(f"Failed to find {key_to_store} in {key}") + d[key_to_store] = None + + # Add fake data + if "LBA" in metadata["antennafield"]: + d["frequency_band"] = [["LBA_10_90"] * 2] * 96 + else: + d["frequency_band"] = [["HBA_110_190"] * 2] * 48 + d["subbands"] = list(range(12, 500)) + d["integration_interval"] = 1.0 + + # Timestamp + d["timestamp"] = key.replace("BST_", "") + + # Store dict + data_dict.append(d) + + # Close file + h5.close() + + # Store arrays + metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) + metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) + + # Assume first entry holds for all (ASSUMPTION!!) + metadata["subbands"] = data_dict[0]["subbands"] + metadata["frequency_bands"] = data_dict[0]["frequency_band"] + if "HBA" in metadata["antennafield"]: + metadata["tile_beam_pointing_direction"] = data_dict[0][ + "tile_beam_pointing_direction" + ][0] + else: + metadata["tile_beam_pointing_direction"] = "" + + # Data shape + nchan, npol = data_dict[0]["data"].shape + metadata["nchan"] = nchan + + # Store as nsub x npol x nchan + data = np.zeros((nsub, 2, nchan)) + for isub in range(nsub): + data[isub, 0] = data_dict[isub]["data"][:, 0::2].T + data[isub, 1] = data_dict[isub]["data"][:, 1::2].T + + return metadata, data diff --git a/lofty/io/_read_single_minio_sst.py b/lofty/io/_read_single_minio_sst.py new file mode 100644 index 0000000000000000000000000000000000000000..e3c2860f61fb7671d367f24636346629af434050 --- /dev/null +++ b/lofty/io/_read_single_minio_sst.py @@ -0,0 +1,99 @@ +# Copyright (C) 2024 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import h5py +import numpy as np + + +def read_single_minio_sst(sst_filename): + # Open HDF5 file + h5 = h5py.File(sst_filename, "r") + + # Store main header meta data + metadata = {} + for key in [ + "antenna_names", + "antenna_status", + "antenna_type", + "antenna_usage_mask", + "station_name", + ]: + metadata[key] = h5["/"].attrs[key] + + # Antennafield + metadata["antennafield"] = ( + h5["/"] + .attrs["antennafield_device"] + .upper() + .replace("STAT/AFH/", "") + .replace("STAT/AFL/", "") + ) + + # Number of subintegrations + nsub = len(h5.keys()) + metadata["nsub"] = nsub + + # Loop over keys + data_dict = [] + for key in h5.keys(): + # Keys in group + attribute_keys = list(h5[key].attrs.keys()) + + # Keys to store in dict + keys_to_store = [] + + # Create dict + d = {"data": np.array(h5[key])} + for key_to_store in keys_to_store: + if key_to_store in attribute_keys: + d[key_to_store] = h5[key].attrs[key_to_store] + else: + print(f"Failed to find {key_to_store} in {key}") + d[key_to_store] = None + + # Add fake data + if "LBA" in metadata["antennafield"]: + d["frequency_band"] = [["LBA_10_90"] * 2] * 96 + else: + d["frequency_band"] = [["HBA_110_190"] * 2] * 48 + d["subbands"] = list(range(0, 512)) + d["integration_interval"] = 1.0 + + # Timestamp + d["timestamp"] = key.replace("SST_", "") + + # Store dict + data_dict.append(d) + + # Close file + h5.close() + + # Store arrays + metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) + metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) + + # Assume first entry holds for all (ASSUMPTION!!) + metadata["subbands"] = data_dict[0]["subbands"] + metadata["frequency_bands"] = data_dict[0]["frequency_band"] + if "HBA" in metadata["antennafield"]: + metadata["tile_beam_pointing_direction"] = data_dict[0][ + "tile_beam_pointing_direction" + ][0] + else: + metadata["tile_beam_pointing_direction"] = "" + + # Data shape + nrcu, nchan = data_dict[0]["data"].shape + nant = nrcu // 2 + metadata["nrcu"] = nrcu + metadata["nant"] = nant + metadata["nchan"] = nchan + + # Store as nsub x npol x nant x nchan + data = np.zeros((nsub, 2, nant, nchan)) + for isub in range(nsub): + data[isub, 0] = data_dict[isub]["data"][0::2, :] + data[isub, 1] = data_dict[isub]["data"][1::2, :] + + return metadata, data diff --git a/lofty/io/_read_single_minio_xst.py b/lofty/io/_read_single_minio_xst.py new file mode 100644 index 0000000000000000000000000000000000000000..058d438aa017d847c1dd49cb94291b317ccf62d9 --- /dev/null +++ b/lofty/io/_read_single_minio_xst.py @@ -0,0 +1,101 @@ +# Copyright (C) 2024 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import h5py +import numpy as np + + +def read_single_minio_xst(xst_filename, fill_empty=False): + # Open HDF5 file + h5 = h5py.File(xst_filename, "r") + + # Store main header meta data + metadata = {} + for key in [ + "antenna_names", + "antenna_status", + "antenna_type", + "antenna_usage_mask", + "station_name", + ]: + metadata[key] = h5["/"].attrs[key] + + # Antennafield + metadata["antennafield"] = ( + h5["/"] + .attrs["antennafield_device"] + .upper() + .replace("STAT/AFH/", "") + .replace("STAT/AFL/", "") + ) + + # Number of subintegrations + nsub = len(h5.keys()) + metadata["nsub"] = nsub + + # Loop over keys + data_dict = [] + for key in h5.keys(): + # Keys in group + attribute_keys = list(h5[key].attrs.keys()) + + # Keys to store in dict + keys_to_store = [] + + # Create dict + d = {"data": np.array(h5[key]["data"]).squeeze()} + for key_to_store in keys_to_store: + if key_to_store in attribute_keys: + d[key_to_store] = h5[key].attrs[key_to_store] + else: + print(f"Failed to find {key_to_store} in {key}") + d[key_to_store] = None + + # Add fake data + if "LBA" in metadata["antennafield"]: + d["frequency_band"] = [["LBA_10_90"] * 2] * 96 + else: + d["frequency_band"] = [["HBA_110_190"] * 2] * 48 + d["subbands"] = h5[key]["subband"].squeeze() + d["integration_interval"] = 1.0 + + # Timestamp + d["timestamp"] = key.replace("XST_", "") + + # Store dict + data_dict.append(d) + + # Close file + h5.close() + + # Store arrays + metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) + metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) + metadata["subbands"] = np.array([d["subbands"] for d in data_dict]) + + # Assume first entry holds for all (ASSUMPTION!!) + metadata["frequency_bands"] = data_dict[0]["frequency_band"] + if "HBA" in metadata["antennafield"]: + metadata["tile_beam_pointing_direction"] = data_dict[0][ + "tile_beam_pointing_direction" + ][0] + else: + metadata["tile_beam_pointing_direction"] = "" + + # Data shape + nrcu, _ = data_dict[0]["data"].shape + nant = nrcu // 2 + metadata["nrcu"] = nrcu + metadata["nant"] = nant + + # Store as nsub x nrcu x nrcu + data = np.zeros((nsub, nrcu, nrcu), dtype=data_dict[0]["data"].dtype) + for isub in range(nsub): + data[isub] = data_dict[isub]["data"] + + # Fill empty parts + if fill_empty: + data += np.conj(np.transpose(data, (0, 2, 1))) * (data == 0) + + return metadata, data diff --git a/lofty/io/_read_single_npz.py b/lofty/io/_read_single_npz.py new file mode 100644 index 0000000000000000000000000000000000000000..acec8645390218405da31028b61b327991e94199 --- /dev/null +++ b/lofty/io/_read_single_npz.py @@ -0,0 +1,40 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from datetime import datetime +from typing import Union +from typing import Any +import os + +import numpy as np + +from lofty.io._antenna_field import extract_antenna_field_from_filename +from lofty.io._station_name import extract_station_name_from_filename + + +def read_single_npz(fname: Union[str, bytes, os.PathLike]): + # Load file + d = np.load(fname) + + # Get data + data = d["data"] + nsub, nant, _ = data.shape + bands = np.array([[d["band"]]]) + timestamps = [datetime.utcfromtimestamp(t).strftime("%FT%T") for t in d["times"]] + + # Construct metadata + metadata = { + "antennafield": extract_antenna_field_from_filename(fname), + "station_name": extract_station_name_from_filename(fname), + "antenna_usage_mask": d["antenna_status"], + "nant": nant, + "nsub": nsub, + "frequency_bands": bands, + "subbands": d["subband"], + "timestamps": timestamps, + "durations": d["duration"], + } + + return metadata, data diff --git a/lofty/io/_read_single_sst.py b/lofty/io/_read_single_sst.py new file mode 100644 index 0000000000000000000000000000000000000000..ae219cf157e3c1a82a96194e1b5c42b1377c43bd --- /dev/null +++ b/lofty/io/_read_single_sst.py @@ -0,0 +1,99 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import h5py + + +def read_single_sst(sst_filename): + # Open HDF5 file + h5 = h5py.File(sst_filename, "r") + + # Store main header meta data + metadata = {} + for key in [ + "antenna_names", + "antenna_quality", + "antenna_type", + "antenna_usage_mask", + "station_name", + ]: + metadata[key] = h5["/"].attrs[key] + + # Antennafield + metadata["antennafield"] = ( + h5["/"] + .attrs["antennafield_device"] + .replace("STAT/AFH/", "") + .replace("STAT/AFL/", "") + ) + + # Number of subintegrations + nsub = len(h5.keys()) + metadata["nsub"] = nsub + + # Loop over keys + data_dict = [] + for key in h5.keys(): + # Keys in group + attribute_keys = list(h5[key].attrs.keys()) + + # Keys to store in dict + keys_to_store = [ + "frequency_band", + "subbands", + "timestamp", + "integration_interval", + ] + + # Create dict + d = {"data": np.array(h5[key])} + for key_to_store in keys_to_store: + if key_to_store in attribute_keys: + d[key_to_store] = h5[key].attrs[key_to_store] + else: + print(f"Failed to find {key_to_store} in {key}") + d[key_to_store] = None + + # Store tile beam for HBA + if "HBA" in metadata["antennafield"]: + d["tile_beam_pointing_direction"] = h5[key].attrs[ + "tile_beam_pointing_direction" + ] + + # Store dict + data_dict.append(d) + + # Store arrays + metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) + metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) + + # Assume first entry holds for all (ASSUMPTION!!) + metadata["subbands"] = data_dict[0]["subbands"] + metadata["frequency_bands"] = data_dict[0]["frequency_band"] + if "HBA" in metadata["antennafield"]: + metadata["tile_beam_pointing_direction"] = data_dict[0][ + "tile_beam_pointing_direction" + ][0] + else: + metadata["tile_beam_pointing_direction"] = "" + + # Data shape + nrcu, nchan = data_dict[0]["data"].shape + nant = nrcu // 2 + metadata["nrcu"] = nrcu + metadata["nant"] = nant + metadata["nchan"] = nchan + + # Store as nsub x npol x nant x nchan + data = np.zeros((nsub, 2, nant, nchan)) + for isub in range(nsub): + data[isub, 0] = data_dict[isub]["data"][0::2, :] + data[isub, 1] = data_dict[isub]["data"][1::2, :] + + # Close file + h5.close() + + return metadata, data diff --git a/lofty/io/_read_single_xst.py b/lofty/io/_read_single_xst.py new file mode 100644 index 0000000000000000000000000000000000000000..3a2d8f0b99a9f8bd0a0c9f30487d586cfcac1ea6 --- /dev/null +++ b/lofty/io/_read_single_xst.py @@ -0,0 +1,98 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import h5py + + +def read_single_xst(xst_filename, fill_empty=False): + # Open HDF5 file + h5 = h5py.File(xst_filename, "r") + + # Store main header meta data + metadata = {} + # for key in ["antenna_names", "antenna_quality", "antenna_type", "antenna_usage_mask", "station_name"]: + # metadata[key] = h5["/"].attrs[key] + for key in ["antenna_names", "antenna_type", "station_name"]: + metadata[key] = h5["/"].attrs[key] + metadata["antenna_usage_mask"] = np.ones(96, dtype="bool") + + # Antennafield + metadata["antennafield"] = ( + h5["/"] + .attrs["antennafield_device"] + .replace("STAT/AFH/", "") + .replace("STAT/AFL/", "") + ) + + # Number of subintegrations + nsub = len(h5.keys()) + metadata["nsub"] = nsub + + # Loop over keys + data_dict = [] + for key in h5.keys(): + # Keys in group + attribute_keys = list(h5[key].attrs.keys()) + + # Keys to store in dict + keys_to_store = [ + "frequency_band", + "subbands", + "timestamp", + "integration_interval", + ] + + # Create dict + d = {"data": np.array(h5[key])} + for key_to_store in keys_to_store: + if key_to_store in attribute_keys: + d[key_to_store] = h5[key].attrs[key_to_store] + else: + print(f"Failed to find {key_to_store} in {key}") + d[key_to_store] = None + + # Store tile beam for HBA + if "HBA" in metadata["antennafield"]: + d["tile_beam_pointing_direction"] = h5[key].attrs[ + "tile_beam_pointing_direction" + ] + + # Store dict + data_dict.append(d) + + # Store arrays + metadata["timestamps"] = np.array([d["timestamp"] for d in data_dict]) + metadata["durations"] = np.array([d["integration_interval"] for d in data_dict]) + metadata["subbands"] = [data_dict[i]["subbands"][0] for i in range(nsub)] + + # Assume first entry holds for all (ASSUMPTION!!) + metadata["frequency_bands"] = data_dict[0]["frequency_band"] + if "HBA" in metadata["antennafield"]: + metadata["tile_beam_pointing_direction"] = data_dict[0][ + "tile_beam_pointing_direction" + ][0] + else: + metadata["tile_beam_pointing_direction"] = "" + + # Data shape + nrcu, _ = data_dict[0]["data"].shape + nant = nrcu // 2 + metadata["nrcu"] = nrcu + metadata["nant"] = nant + + # Store as nsub x nrcu x nrcu + data = np.zeros((nsub, nrcu, nrcu), dtype=data_dict[0]["data"].dtype) + for isub in range(nsub): + data[isub] = data_dict[isub]["data"] + + # Fill empty parts + if fill_empty: + data += np.conj(np.transpose(data, (0, 2, 1))) * (data == 0) + + # Close file + h5.close() + + return metadata, data diff --git a/lofty/io/_read_sst.py b/lofty/io/_read_sst.py new file mode 100644 index 0000000000000000000000000000000000000000..d5a5b00455e01d4d52ec8f328d23c00b98fca1a9 --- /dev/null +++ b/lofty/io/_read_sst.py @@ -0,0 +1,34 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from lofty.io._read_single_sst import read_single_sst + + +def read_sst(sst_filenames): + first = True + for sst_filename in sst_filenames: + # Get data + if first: + metadata, data = read_single_sst(sst_filename) + first = False + else: + newmetadata, newdata = read_single_sst(sst_filename) + + # Concatenate timestamps + metadata["timestamps"] = np.hstack( + (metadata["timestamps"], newmetadata["timestamps"]) + ) + + # Concatenate integration_interval + metadata["durations"] = np.hstack( + (metadata["durations"], newmetadata["durations"]) + ) + + # Concatenate data + data = np.vstack((data, newdata)) + + return metadata, data diff --git a/lofty/io/_read_xst.py b/lofty/io/_read_xst.py new file mode 100644 index 0000000000000000000000000000000000000000..b27ca869dc6469274b3438fdecc318ad8aba4a7d --- /dev/null +++ b/lofty/io/_read_xst.py @@ -0,0 +1,42 @@ +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from lofty.io._read_single_xst import read_single_xst + + +def read_xst(xst_filenames, fill_empty=False): + first = True + for xst_filename in xst_filenames: + # Get data + if first: + metadata, data = read_single_xst(xst_filename, fill_empty) + first = False + else: + newmetadata, newdata = read_single_xst(xst_filename, fill_empty) + + # Concatenate timestamps + metadata["timestamps"] = np.hstack( + (metadata["timestamps"], newmetadata["timestamps"]) + ) + + # Concatenate integration_interval + metadata["durations"] = np.hstack( + (metadata["durations"], newmetadata["durations"]) + ) + + # Concatenate subbands + metadata["subbands"] = np.hstack( + (metadata["subbands"], newmetadata["subbands"]) + ) + + # Concatenate data + data = np.vstack((data, newdata)) + + # Update nsub + metadata["nsub"] = data.shape[0] + + return metadata, data diff --git a/lofty/io/_station_name.py b/lofty/io/_station_name.py new file mode 100644 index 0000000000000000000000000000000000000000..0ebaf325b7b7c4180c591e46aa36eec46e735cc6 --- /dev/null +++ b/lofty/io/_station_name.py @@ -0,0 +1,27 @@ +# Copyright (C) 2024 Corne Lukken +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +import logging +import re +from typing import Optional + +logger = logging.getLogger() + +STATION_NAME_REGEX = "[^\\W\\d_]{2}[0-9]{3}" + + +def extract_station_name_from_filename(fname: str) -> Optional[str]: + """Extract the name of a station from the filename""" + + pattern = re.compile(STATION_NAME_REGEX) + names = pattern.findall(fname) + + if len(names) > 1: + logger.warning( + "Found multiple station names: %s matches in file: %s", names, fname + ) + if len(names) > 0: + return names[0] + + return None diff --git a/lofty/plotting.py b/lofty/plotting.py index b876412f7bdd1eb08b298313d21e0cba19e2e9b2..3c0f7520ac859e8664e88c9b911d641747255ec6 100644 --- a/lofty/plotting.py +++ b/lofty/plotting.py @@ -1,10 +1,16 @@ -#!/usr/bin/env python3 +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + import numpy as np import matplotlib.pyplot as plt import matplotlib.dates as mdates -from matplotlib.ticker import FormatStrFormatter, MultipleLocator +from matplotlib.axes import Axes +from matplotlib.figure import Figure +from matplotlib.ticker import MultipleLocator from matplotlib.patches import Circle + def decibel(x): return 10 * np.log10(x) @@ -27,7 +33,7 @@ def plot_xst_matrix(metadata, data, isub): else: offset = 8 tick_step = nant // 6 - + # Get bad antennas mask = metadata["antenna_usage_mask"][:nant] bad_ant = np.arange(nant)[~mask] @@ -41,8 +47,8 @@ def plot_xst_matrix(metadata, data, isub): # XY in lower right, XX in upper left xst_img = np.zeros((offset + nant, offset + nant), dtype=cube_xx.dtype) - xst_img[offset:offset+nant, :nant] += cube_xx - xst_img[:nant, offset:offset+nant] += cube_xy.T + xst_img[offset : offset + nant, :nant] += cube_xx + xst_img[:nant, offset : offset + nant] += cube_xy.T # Amplitude ax1_top = ax1.secondary_xaxis("top") @@ -53,7 +59,7 @@ def plot_xst_matrix(metadata, data, isub): labels = [f"{x}" for x in ticks] ax1_top.set_xticks(ticks) ax1_top.set_xticklabels(labels) - ax1_top.xaxis.set_minor_locator(MultipleLocator(4)) + ax1_top.xaxis.set_minor_locator(MultipleLocator(4)) ax1_right.set_yticks(ticks) ax1_right.set_yticklabels(labels) ax1_right.yaxis.set_minor_locator(MultipleLocator(4)) @@ -64,13 +70,21 @@ def plot_xst_matrix(metadata, data, isub): ax1.set_yticklabels([f"{x}" for x in bad_yy]) ax1.set_ylabel("XX Antenna number") - fig.colorbar(img, ax=ax1, extend="both", shrink=0.8, location="bottom", label="Power (dB)") + fig.colorbar( + img, ax=ax1, extend="both", shrink=0.8, location="bottom", label="Power (dB)" + ) # Phase ax2_top = ax2.secondary_xaxis("top") ax2_right = ax2.secondary_yaxis("right") - - img = ax2.imshow(np.angle(xst_img) * 180 / np.pi, origin="lower", vmin=-180, vmax=180, cmap="twilight") + + img = ax2.imshow( + np.angle(xst_img) * 180 / np.pi, + origin="lower", + vmin=-180, + vmax=180, + cmap="twilight", + ) labels = [f"{x}" for x in ticks] ax2_top.set_xticks(ticks) ax2_top.set_xticklabels(labels) @@ -86,12 +100,14 @@ def plot_xst_matrix(metadata, data, isub): ax2.set_yticklabels([f"{x}" for x in bad_yy]) ax2.set_ylabel("XX Antenna number") - fig.colorbar(img, ax=ax2, extend="both", shrink=0.8, location="bottom", label="Phase (deg)") + fig.colorbar( + img, ax=ax2, extend="both", shrink=0.8, location="bottom", label="Phase (deg)" + ) # YY in lower right, YX in upper left xst_img = np.zeros((offset + nant, offset + nant), dtype=cube_xx.dtype) - xst_img[offset:offset+nant, :nant] += cube_yx - xst_img[:nant, offset:offset+nant] += cube_yy.T + xst_img[offset : offset + nant, :nant] += cube_yx + xst_img[:nant, offset : offset + nant] += cube_yy.T # Amplitude ax3_top = ax3.secondary_xaxis("top") @@ -102,7 +118,7 @@ def plot_xst_matrix(metadata, data, isub): labels = [f"{x}" for x in ticks] ax3_top.set_xticks(ticks) ax3_top.set_xticklabels(labels) - ax3_top.xaxis.set_minor_locator(MultipleLocator(4)) + ax3_top.xaxis.set_minor_locator(MultipleLocator(4)) ax3_right.set_yticks(ticks) ax3_right.set_yticklabels(labels) ax3_right.yaxis.set_minor_locator(MultipleLocator(4)) @@ -113,13 +129,21 @@ def plot_xst_matrix(metadata, data, isub): ax3.set_yticklabels([f"{x}" for x in bad_yy]) ax3.set_ylabel("YX Antenna number") - fig.colorbar(img, ax=ax3, extend="both", shrink=0.8, location="bottom", label="Power (dB)") + fig.colorbar( + img, ax=ax3, extend="both", shrink=0.8, location="bottom", label="Power (dB)" + ) # Phase ax4_top = ax4.secondary_xaxis("top") ax4_right = ax4.secondary_yaxis("right") - - img = ax4.imshow(np.angle(xst_img) * 180 / np.pi, origin="lower", vmin=-180, vmax=180, cmap="twilight") + + img = ax4.imshow( + np.angle(xst_img) * 180 / np.pi, + origin="lower", + vmin=-180, + vmax=180, + cmap="twilight", + ) labels = [f"{x}" for x in ticks] ax4_top.set_xticks(ticks) ax4_top.set_xticklabels(labels) @@ -135,10 +159,13 @@ def plot_xst_matrix(metadata, data, isub): ax4.set_yticklabels([f"{x}" for x in bad_yy]) ax4.set_ylabel("YX Antenna number") - fig.colorbar(img, ax=ax4, extend="both", shrink=0.8, location="bottom", label="Phase (deg)") - + fig.colorbar( + img, ax=ax4, extend="both", shrink=0.8, location="bottom", label="Phase (deg)" + ) + plt.show() + def plot_sst_timeseries(metadata, data, subband): # Timestamp t = mdates.date2num(metadata["timestamps"]) @@ -161,7 +188,7 @@ def plot_sst_timeseries(metadata, data, subband): color = "r" ax1.plot_date(t, px[:, iant], color, alpha=0.1) ax2.plot_date(t, py[:, iant], color, alpha=0.1) - + field = metadata["antennafield"] if "HBA" in field: tile_beam = metadata["tile_beam_pointing_direction"] @@ -174,9 +201,10 @@ def plot_sst_timeseries(metadata, data, subband): ax2.xaxis.set_major_formatter(date_format) ax2.set_xlabel("Date (UTC)") ax2.set_ylabel("Power (dB)") - + plt.show() + def plot_bst_timeseries(metadata, data, subband): # Timestamp t = mdates.date2num(metadata["timestamps"]) @@ -196,7 +224,7 @@ def plot_bst_timeseries(metadata, data, subband): ax1.plot_date(t, px, "k", alpha=1) ax2.plot_date(t, py, "k", alpha=1) - + field = metadata["antennafield"] if "HBA" in field: tile_beam = metadata["tile_beam_pointing_direction"] @@ -209,10 +237,10 @@ def plot_bst_timeseries(metadata, data, subband): ax2.xaxis.set_major_formatter(date_format) ax2.set_xlabel("Date (UTC)") ax2.set_ylabel("Power (dB)") - + plt.show() - + def freq_from_sb(sb, band): if "LBA" in band: freq = sb * 0.1953125 @@ -222,20 +250,26 @@ def freq_from_sb(sb, band): freq = 200 + sb * 0.1953125 return freq - + + def plot_sst_bandpass(metadata, data): - fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(16, 6), sharex=True, gridspec_kw={"height_ratios": [0.2, 0.8], "hspace": 0.02, "wspace": 0.15}) + fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots( + 2, + 2, + figsize=(16, 6), + sharex=True, + gridspec_kw={"height_ratios": [0.2, 0.8], "hspace": 0.02, "wspace": 0.15}, + ) # Frequencies band_x, band_y = metadata["frequency_bands"][0] freq_x = np.array([freq_from_sb(sb, band_x) for sb in metadata["subbands"]]) - freq_y = np.array([freq_from_sb(sb, band_y) for sb in metadata["subbands"]]) + freq_y = np.array([freq_from_sb(sb, band_y) for sb in metadata["subbands"]]) px = decibel(data[0, 0, :, :].squeeze()) py = decibel(data[0, 1, :, :].squeeze()) - nant = metadata["nant"] - mask = metadata["antenna_usage_mask"][:nant] # Array is longer by default + mask = metadata["antenna_usage_mask"][:nant] # Array is longer by default for iant in range(nant): if mask[iant]: color = "k" @@ -244,71 +278,93 @@ def plot_sst_bandpass(metadata, data): ax1.plot(freq_x, px[iant], color, linewidth=0.2) ax2.plot(freq_y, py[iant], color, linewidth=0.2) - img_ax3 = ax3.imshow(px, aspect="auto", origin="lower", interpolation="None", - extent=[np.min(freq_x), np.max(freq_x), 0, nant]) - cbar_ax3 = plt.colorbar(img_ax3, ax=ax3, extend="both", shrink=0.99, - location="bottom", label="Power (dB)") - img_ax4 = ax4.imshow(py, aspect="auto", origin="lower", interpolation="None", - extent=[np.min(freq_y), np.max(freq_y), 0, nant]) - cbar_ax4 = plt.colorbar(img_ax4, ax=ax4, extend="both", shrink=0.99, location="bottom", label="Power (dB)") + img_ax3 = ax3.imshow( + px, + aspect="auto", + origin="lower", + interpolation="None", + extent=[np.min(freq_x), np.max(freq_x), 0, nant], + ) + _ = plt.colorbar( # cbar_ax3 + img_ax3, + ax=ax3, + extend="both", + shrink=0.99, + location="bottom", + label="Power (dB)", + ) + img_ax4 = ax4.imshow( + py, + aspect="auto", + origin="lower", + interpolation="None", + extent=[np.min(freq_y), np.max(freq_y), 0, nant], + ) + _ = plt.colorbar( # cbar_ax4 + img_ax4, + ax=ax4, + extend="both", + shrink=0.99, + location="bottom", + label="Power (dB)", + ) freq_ticks = np.arange(np.min(freq_x), np.max(freq_x), 10) freq_labels = [f"{x:g}" for x in freq_ticks] - + if nant == 96: - ticks = np.arange(0, nant + 1, 16) + ticks = np.arange(0, nant + 1, 16) else: ticks = np.arange(0, nant + 1, 4) labels = [f"{x}" for x in ticks] ax1.set_ylabel("Power (dB)") -# ax1.set_title(title) - + # ax1.set_title(title) + ax3.set_xticks(freq_ticks) ax3.set_xticklabels(freq_labels) - ax3.xaxis.set_minor_locator(MultipleLocator(4)) + ax3.xaxis.set_minor_locator(MultipleLocator(4)) ax3.set_yticks(ticks) ax3.set_yticklabels(labels) - ax3.yaxis.set_minor_locator(MultipleLocator(4)) - + ax3.yaxis.set_minor_locator(MultipleLocator(4)) + ax3.set_xlabel("Frequency (MHz)") ax3.set_ylabel("Antenna number") - bad = np.arange(nant)[mask == False] - + bad = np.arange(nant)[not mask] + ax3_right = ax3.secondary_yaxis("right") ax3_right.set_yticks(bad + 0.5) ax3_right.set_yticklabels([f"{x}" for x in bad]) ax2.set_ylabel("Power (dB)") -# ax1.set_title(title) - + # ax1.set_title(title) + ax4.set_xticks(freq_ticks) ax4.set_xticklabels(freq_labels) - ax4.xaxis.set_minor_locator(MultipleLocator(4)) + ax4.xaxis.set_minor_locator(MultipleLocator(4)) ax4.set_yticks(ticks) ax4.set_yticklabels(labels) - ax4.yaxis.set_minor_locator(MultipleLocator(4)) - + ax4.yaxis.set_minor_locator(MultipleLocator(4)) + ax4.set_xlabel("Frequency (MHz)") ax4.set_ylabel("Antenna number") ax4_right = ax4.secondary_yaxis("right") ax4_right.set_yticks(bad + 0.5) ax4_right.set_yticklabels([f"{x}" for x in bad]) - - + plt.show() - + def plot_sst_dynamic_spectrum(metadata, data, iant, pol): # Timestamps t = mdates.date2num(metadata["timestamps"]) - + # Frequencies band_x, band_y = metadata["frequency_bands"][0] freq_x = np.array([freq_from_sb(sb, band_x) for sb in metadata["subbands"]]) - freq_y = np.array([freq_from_sb(sb, band_y) for sb in metadata["subbands"]]) + # freq_y = np.array([freq_from_sb(sb, band_y) for sb in metadata["subbands"]]) if iant >= 0: data = data[:, :, iant, :] @@ -316,7 +372,7 @@ def plot_sst_dynamic_spectrum(metadata, data, iant, pol): data = np.mean(data, axis=2) print(data.shape) - + if pol.lower() == "x": p = decibel(data[:, 0, :].squeeze()) elif pol.lower() == "y": @@ -330,8 +386,15 @@ def plot_sst_dynamic_spectrum(metadata, data, iant, pol): fig.autofmt_xdate(rotation=0, ha="center") vmin, vmax = np.percentile(p, (5, 99)) - ax.imshow(p.T, origin="lower", aspect="auto", interpolation="None", vmin=vmin, vmax=vmax, - extent=[np.min(t), np.max(t), np.min(freq_x), np.max(freq_x)]) + ax.imshow( + p.T, + origin="lower", + aspect="auto", + interpolation="None", + vmin=vmin, + vmax=vmax, + extent=[np.min(t), np.max(t), np.min(freq_x), np.max(freq_x)], + ) ax.xaxis_date() ax.xaxis.set_major_formatter(date_format) @@ -340,14 +403,15 @@ def plot_sst_dynamic_spectrum(metadata, data, iant, pol): plt.show() + def plot_bst_dynamic_spectrum(metadata, data, pol): # Timestamps t = mdates.date2num(metadata["timestamps"]) - + # Frequencies band_x, band_y = metadata["frequency_bands"][0] freq_x = np.array([freq_from_sb(sb, band_x) for sb in metadata["subbands"]]) - freq_y = np.array([freq_from_sb(sb, band_y) for sb in metadata["subbands"]]) + # freq_y = np.array([freq_from_sb(sb, band_y) for sb in metadata["subbands"]]) if pol.lower() == "x": p = decibel(data[:, 0, :].squeeze()) @@ -362,8 +426,15 @@ def plot_bst_dynamic_spectrum(metadata, data, pol): fig.autofmt_xdate(rotation=0, ha="center") vmin, vmax = np.percentile(p, (5, 99)) - ax.imshow(p.T, origin="lower", aspect="auto", interpolation="None", vmin=vmin, vmax=vmax, - extent=[np.min(t), np.max(t), np.min(freq_x), np.max(freq_x)]) + ax.imshow( + p.T, + origin="lower", + aspect="auto", + interpolation="None", + vmin=vmin, + vmax=vmax, + extent=[np.min(t), np.max(t), np.min(freq_x), np.max(freq_x)], + ) ax.xaxis_date() ax.xaxis.set_major_formatter(date_format) @@ -371,33 +442,58 @@ def plot_bst_dynamic_spectrum(metadata, data, pol): ax.set_ylabel("Frequency (MHz)") plt.show() - -def sky_plot(img, title, subtitle, source_names, source_lmn, fig): + + +def mark_sources(source_names, source_lmn, ax: Axes): + if source_lmn is not None: + for name, lmn in zip(source_names, source_lmn): + if lmn[2] >= -1: + ax.plot(lmn[0], lmn[1], marker="x", color="k") + ax.text(lmn[0], lmn[1], f" {name}") + + +def sky_plot(img, title, subtitle, source_names, source_lmn, fig: Figure): if fig is None: fig = plt.figure(figsize=(10, 10)) ax = fig.add_subplot(1, 1, 1) - circle1 = Circle((0, 0), 1.0, edgecolor='k', fill=False, facecolor='none', alpha=0.3) + circle1 = Circle( + (0, 0), 1.0, edgecolor="k", fill=False, facecolor="none", alpha=0.3 + ) # Plot image and colorbar cimg = ax.imshow(img, origin="lower", cmap="Blues_r", extent=[1, -1, -1, 1]) - cbar = fig.colorbar(cimg, ax=ax, location="right", format="%.2e") + _ = fig.colorbar(cimg, ax=ax, location="right", format="%.2e") # cbar ax.set_axis_off() ax.add_artist(circle1) # Plot sources - if source_lmn is not None: - for name, lmn in zip(source_names, source_lmn): - if lmn[2] >= -1: - ax.plot(lmn[0], lmn[1], marker="x", color="k") - ax.text(lmn[0], lmn[1], f" {name}") + mark_sources(source_names, source_lmn, ax) # Plot the compass directions - for direction, xoffset, yoffset in zip(["E", "W", "N", "S"], [1.05, -1.05, 0.0, 0.0], [0.0, 0.0, 1.05, -1.05]): - ax.text(xoffset, yoffset, direction, horizontalalignment='center', verticalalignment='center', color='k') - - ax.text(0.5, 1.1, title, fontsize=17, ha="center", va="bottom", transform=ax.transAxes) - ax.text(0.5, 1.05, subtitle, fontsize=12, ha="center", va="bottom", transform=ax.transAxes) + for direction, xoffset, yoffset in zip( + ["E", "W", "N", "S"], [1.05, -1.05, 0.0, 0.0], [0.0, 0.0, 1.05, -1.05] + ): + ax.text( + xoffset, + yoffset, + direction, + horizontalalignment="center", + verticalalignment="center", + color="k", + ) + + ax.text( + 0.5, 1.1, title, fontsize=17, ha="center", va="bottom", transform=ax.transAxes + ) + ax.text( + 0.5, + 1.05, + subtitle, + fontsize=12, + ha="center", + va="bottom", + transform=ax.transAxes, + ) return fig - diff --git a/lofty/sources.py b/lofty/sources.py index 65cf657429009129499ac1576ab9edbb114e644a..b0eb495eeb0caf34df501bc33e8c2ad426d67b36 100644 --- a/lofty/sources.py +++ b/lofty/sources.py @@ -1,4 +1,7 @@ -#!/usr/bin/env python3 +# Copyright (C) 2023 Cees Bassa +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + import numpy as np import astropy.units as u @@ -8,7 +11,6 @@ from astropy.coordinates import EarthLocation, SkyCoord, get_icrs_coordinates from astropy.coordinates import solar_system_ephemeris, get_body, GCRS from astropy.coordinates import SkyOffsetFrame, CartesianRepresentation -#from lofarimaging import skycoord_to_lmn def skycoord_to_lmn(pos: SkyCoord, phasecentre: SkyCoord): """ @@ -22,7 +24,8 @@ def skycoord_to_lmn(pos: SkyCoord, phasecentre: SkyCoord): Note that this means that l increases east-wards - This function was taken from https://github.com/SKA-ScienceDataProcessor/algorithm-reference-library + This function was taken from: + https://github.com/SKA-ScienceDataProcessor/algorithm-reference-library """ # Determine relative sky position @@ -39,7 +42,7 @@ def lmn_sources(timestamps, durations, loc, source_names): # Array lengths ntime = len(timestamps) nsrc = len(source_names) - + # Get times t = Time(timestamps, format="isot", scale="utc") + 0.5 * np.array(durations) * u.s @@ -47,17 +50,21 @@ def lmn_sources(timestamps, durations, loc, source_names): loc = EarthLocation.from_geocentric(*loc * u.m) # Get phase centre - zenith = SkyCoord(az=np.zeros(ntime), alt=90 * np.ones(ntime), unit="deg", - obstime=t, location=loc, frame="altaz").transform_to(GCRS) + zenith = SkyCoord( + az=np.zeros(ntime), + alt=90 * np.ones(ntime), + unit="deg", + obstime=t, + location=loc, + frame="altaz", + ).transform_to(GCRS) source_lmn = np.zeros((nsrc, ntime, 3)) for i, source_name in enumerate(source_names): if source_name.lower() in solar_system_ephemeris.bodies: p = get_body(source_name.lower(), t, loc) else: - p = get_icrs_coordinates(source_name.lower()) + p = get_icrs_coordinates(source_name.lower()) source_lmn[i] = np.array(skycoord_to_lmn(p, zenith)).T return np.moveaxis(source_lmn, 1, 0) - - diff --git a/lofty/utility/__init__.py b/lofty/utility/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..85d7d5b6ff5cd51bec3b068499d093bbf06e86e9 --- /dev/null +++ b/lofty/utility/__init__.py @@ -0,0 +1,16 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from lofty.utility._frequency_from_subband import frequency_from_subband +from lofty.utility._station_name import get_full_station_name +from lofty.utility._station_pqr import get_station_pqr +from lofty.utility._station_type import get_station_type +from lofty.utility._station_xyz import get_station_xyz + +__all__ = [ + "frequency_from_subband", + "get_full_station_name", + "get_station_pqr", + "get_station_type", + "get_station_xyz", +] diff --git a/lofty/utility/_frequency_from_subband.py b/lofty/utility/_frequency_from_subband.py new file mode 100644 index 0000000000000000000000000000000000000000..0c0e25e81e1acccb833c46f2eecb089d2420af99 --- /dev/null +++ b/lofty/utility/_frequency_from_subband.py @@ -0,0 +1,39 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from typing import Optional + + +def frequency_from_subband(subband: int, band: str) -> Optional[float]: + """ + Convert central frequency to subband number + + Args: + sb: subband number + band: filter band + + Returns: + float: frequency in Hz + + Example: + >>> frequency_from_subband(297, '30_90') + 58007812.5 + """ + if band not in ["10_90", "30_90", "110_190", "170_230", "210_250"]: + return None + + if band == "10_90" or band == "30_90": + clock, zone = 200e6, 1 + elif band == "110_190": + clock, zone = 200e6, 2 + elif band == "170_230": + clock, zone = 160e6, 3 + elif band == "210_250": + clock, zone = 200e6, 3 + else: + return None + + sb_bandwidth = 0.5 * clock / 512.0 + freq_offset = 0.5 * clock * (zone - 1) + freq = (subband * sb_bandwidth) + freq_offset + return freq diff --git a/lofty/utility/_station_name.py b/lofty/utility/_station_name.py new file mode 100644 index 0000000000000000000000000000000000000000..79018e97d3fc729fee49fd86f36e20f64132e1f1 --- /dev/null +++ b/lofty/utility/_station_name.py @@ -0,0 +1,34 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + + +def get_full_station_name(station_name: str, antenna_set: str) -> str: + """ + Get full station name with the field appended, e.g. DE603LBA + + Args: + station_name (str): Short station name, e.g. 'DE603' + antenna_set (str): antenna_set, e.g. LBA_OUTER + + Returns: + str: Full station name, e.g. DE603LBA + + Example: + >>> get_full_station_name("DE603", 'LBA_INNER') + 'DE603LBA' + + >>> get_full_station_name("LV614", 'HBA') + 'LV614HBA' + + >>> get_full_station_name("CS013LBA", 'LBA_OUTER') + 'CS013LBA' + + >>> get_full_station_name("CS002", 'LBA_OUTER') + 'CS002LBA' + """ + if len(station_name) > 5: + return station_name + elif antenna_set[0:3] in ["LBA", "HBA"]: + station_name += antenna_set[0:3] + + return station_name diff --git a/lofty/utility/_station_pqr.py b/lofty/utility/_station_pqr.py new file mode 100644 index 0000000000000000000000000000000000000000..944684e8d6d0e1c1d206efc3b6a36982c1005c72 --- /dev/null +++ b/lofty/utility/_station_pqr.py @@ -0,0 +1,279 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from lofarantpos.db import LofarAntennaDatabase +import numpy as np + +from lofty.utility._station_name import get_full_station_name +from lofty.utility._station_type import get_station_type + +# Configurations for HBA observations with a single dipole activated per tile. +GENERIC_INT_201512 = [ + 0, + 5, + 3, + 1, + 8, + 3, + 12, + 15, + 10, + 13, + 11, + 5, + 12, + 12, + 5, + 2, + 10, + 8, + 0, + 3, + 5, + 1, + 4, + 0, + 11, + 6, + 2, + 4, + 9, + 14, + 15, + 3, + 7, + 5, + 13, + 15, + 5, + 6, + 5, + 12, + 15, + 7, + 1, + 1, + 14, + 9, + 4, + 9, + 3, + 9, + 3, + 13, + 7, + 14, + 7, + 14, + 2, + 8, + 8, + 0, + 1, + 4, + 2, + 2, + 12, + 15, + 5, + 7, + 6, + 10, + 12, + 3, + 3, + 12, + 7, + 4, + 6, + 0, + 5, + 9, + 1, + 10, + 10, + 11, + 5, + 11, + 7, + 9, + 7, + 6, + 4, + 4, + 15, + 4, + 1, + 15, +] + +GENERIC_CORE_201512 = [ + 0, + 10, + 4, + 3, + 14, + 0, + 5, + 5, + 3, + 13, + 10, + 3, + 12, + 2, + 7, + 15, + 6, + 14, + 7, + 5, + 7, + 9, + 0, + 15, + 0, + 10, + 4, + 3, + 14, + 0, + 5, + 5, + 3, + 13, + 10, + 3, + 12, + 2, + 7, + 15, + 6, + 14, + 7, + 5, + 7, + 9, + 0, + 15, +] + +GENERIC_REMOTE_201512 = [ + 0, + 13, + 12, + 4, + 11, + 11, + 7, + 8, + 2, + 7, + 11, + 2, + 10, + 2, + 6, + 3, + 8, + 3, + 1, + 7, + 1, + 15, + 13, + 1, + 11, + 1, + 12, + 7, + 10, + 15, + 8, + 2, + 12, + 13, + 9, + 13, + 4, + 5, + 5, + 12, + 5, + 5, + 9, + 11, + 15, + 12, + 2, + 15, +] + + +def get_station_pqr(station_name: str, antenna_set: str, db: LofarAntennaDatabase): + """ + Get PQR coordinates for the relevant subset of antennas in a station. + + Args: + station_name: Station name, e.g. 'DE603LBA' or 'DE603' + antenna_set: antenna_set, e.g. LBA_INNER + db: instance of LofarAntennaDatabase from lofarantpos + + Example: + >>> from lofarantpos.db import LofarAntennaDatabase + >>> db = LofarAntennaDatabase() + >>> pqr = get_station_pqr("DE603", "LBA_OUTER", db) + >>> pqr.shape + (96, 3) + >>> pqr[0, 0] + 1.7434713 + + >>> pqr = get_station_pqr("LV614", "HBA", db) + >>> pqr.shape + (96, 3) + """ + + full_station_name = get_full_station_name(station_name, antenna_set) + station_type = get_station_type(full_station_name) + + all_pqr = db.antenna_pqr(full_station_name) + if "LBA" in antenna_set: + if antenna_set == "LBA_OUTER": + station_pqr = all_pqr[48:, :] + elif antenna_set == "LBA_INNER": + station_pqr = all_pqr[:48, :] + elif antenna_set == "LBA_SPARSE_EVEN": + station_pqr = np.ravel( + np.column_stack((all_pqr[:48:2], all_pqr[49::2])) + ).reshape(48, 3) + elif antenna_set == "LBA_SPARSE_ODD": + station_pqr = np.ravel( + np.column_stack((all_pqr[1:48:2], all_pqr[48::2])) + ).reshape(48, 3) + else: + station_pqr = all_pqr + elif "HBA" in antenna_set: + selected_dipole_config = { + "intl": GENERIC_INT_201512, + "remote": GENERIC_REMOTE_201512, + "core": GENERIC_CORE_201512, + } + selected_dipoles = ( + selected_dipole_config[station_type] + + np.arange(selected_dipole_config[station_type]) * 16 + ) + single_dipole_pqr = db.hba_dipole_pqr(full_station_name)[selected_dipoles] + if antenna_set == "HBA_SINGLE": + station_pqr = single_dipole_pqr + elif antenna_set == "HBA0_SINGLE": + station_pqr = single_dipole_pqr[:24, :] + elif antenna_set == "HBA1_SINGLE": + station_pqr = single_dipole_pqr[24:, :] + elif antenna_set == "HBA0": + station_pqr = all_pqr[:24, :] + elif antenna_set == "HBA1": + station_pqr = all_pqr[24:, :] + else: + station_pqr = all_pqr + + return station_pqr.astype("float32") diff --git a/lofty/utility/_station_type.py b/lofty/utility/_station_type.py new file mode 100644 index 0000000000000000000000000000000000000000..298953bc4b53458c657f8c4532c9e2d4cbc99f03 --- /dev/null +++ b/lofty/utility/_station_type.py @@ -0,0 +1,24 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + + +def get_station_type(station_name: str) -> str: + """ + Get the station type, one of 'intl', 'core' or 'remote' + + Args: + station_name: Station name, e.g. "DE603LBA" or just "DE603" + + Returns: + str: station type, one of 'intl', 'core' or 'remote' + + Example: + >>> get_station_type("DE603") + 'intl' + """ + if station_name[0] == "C": + return "core" + elif station_name[0] == "R" or station_name[:5] == "PL611": + return "remote" + else: + return "intl" diff --git a/lofty/utility/_station_xyz.py b/lofty/utility/_station_xyz.py new file mode 100644 index 0000000000000000000000000000000000000000..fc22dac62dbc8e62ccfcabee5acbdf424b78403d --- /dev/null +++ b/lofty/utility/_station_xyz.py @@ -0,0 +1,55 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +from lofarantpos.db import LofarAntennaDatabase +import numpy as np + +from lofty.utility._station_pqr import get_station_pqr +from lofty.utility._station_name import get_full_station_name + + +def get_station_xyz(station_name: str, antenna_set: str, db: LofarAntennaDatabase): + """ + Get XYZ coordinates for the relevant subset of antennas in a station. + The XYZ system is defined as the PQR system rotated along the R axis to make + the Q-axis point towards local north. + + Args: + station_name: Station name, e.g. 'DE603LBA' or 'DE603' + antenna_set: antenna_set, e.g. LBA_OUTER + db: instance of LofarAntennaDatabase from lofarantpos + + Returns: + np.array: Antenna xyz, shape [n_ant, 3] + np.array: rotation matrix pqr_to_xyz, shape [3, 3] + + Example: + >>> from lofarantpos.db import LofarAntennaDatabase + >>> db = LofarAntennaDatabase() + >>> xyz, _ = get_station_xyz("DE603", "LBA_OUTER", db) + >>> xyz.shape + (96, 3) + >>> f"{xyz[0, 0]:.7f}" + '2.7033776' + + >>> xyz, _ = get_station_xyz("LV614", "HBA", db) + >>> xyz.shape + (96, 3) + """ + station_pqr = get_station_pqr(station_name, antenna_set, db) + + station_name = get_full_station_name(station_name, antenna_set) + + rotation = db.rotation_from_north(station_name) + + pqr_to_xyz = np.array( + [ + [np.cos(-rotation), -np.sin(-rotation), 0], + [np.sin(-rotation), np.cos(-rotation), 0], + [0, 0, 1], + ] + ) + + station_xyz = (pqr_to_xyz @ station_pqr.T).T + + return station_xyz, pqr_to_xyz diff --git a/lofty/utils.py b/lofty/utils.py deleted file mode 100644 index bec5cd25af858bd00e026ee252439925e1325774..0000000000000000000000000000000000000000 --- a/lofty/utils.py +++ /dev/null @@ -1,198 +0,0 @@ -#!/usr/bin/env python3 -import numpy as np - -# Configurations for HBA observations with a single dipole activated per tile. -GENERIC_INT_201512 = [0, 5, 3, 1, 8, 3, 12, 15, 10, 13, 11, 5, 12, 12, 5, 2, 10, 8, 0, 3, 5, 1, 4, 0, 11, 6, 2, 4, 9, - 14, 15, 3, 7, 5, 13, 15, 5, 6, 5, 12, 15, 7, 1, 1, 14, 9, 4, 9, 3, 9, 3, 13, 7, 14, 7, 14, 2, 8, - 8, 0, 1, 4, 2, 2, 12, 15, 5, 7, 6, 10, 12, 3, 3, 12, 7, 4, 6, 0, 5, 9, 1, 10, 10, 11, 5, 11, 7, 9, - 7, 6, 4, 4, 15, 4, 1, 15] -GENERIC_CORE_201512 = [0, 10, 4, 3, 14, 0, 5, 5, 3, 13, 10, 3, 12, 2, 7, 15, 6, 14, 7, 5, 7, 9, 0, 15, 0, 10, 4, 3, 14, - 0, 5, 5, 3, 13, 10, 3, 12, 2, 7, 15, 6, 14, 7, 5, 7, 9, 0, 15] -GENERIC_REMOTE_201512 = [0, 13, 12, 4, 11, 11, 7, 8, 2, 7, 11, 2, 10, 2, 6, 3, 8, 3, 1, 7, 1, 15, 13, 1, 11, 1, 12, 7, - 10, 15, 8, 2, 12, 13, 9, 13, 4, 5, 5, 12, 5, 5, 9, 11, 15, 12, 2, 15] - -def get_station_type(station_name: str) -> str: - """ - Get the station type, one of 'intl', 'core' or 'remote' - - Args: - station_name: Station name, e.g. "DE603LBA" or just "DE603" - - Returns: - str: station type, one of 'intl', 'core' or 'remote' - - Example: - >>> get_station_type("DE603") - 'intl' - """ - if station_name[0] == "C": - return "core" - elif station_name[0] == "R" or station_name[:5] == "PL611": - return "remote" - else: - return "intl" - -def get_full_station_name(station_name: str, antenna_set: str) -> str: - """ - Get full station name with the field appended, e.g. DE603LBA - - Args: - station_name (str): Short station name, e.g. 'DE603' - antenna_set (str): antenna_set, e.g. LBA_OUTER - - Returns: - str: Full station name, e.g. DE603LBA - - Example: - >>> get_full_station_name("DE603", 'LBA_INNER') - 'DE603LBA' - - >>> get_full_station_name("LV614", 'HBA') - 'LV614HBA' - - >>> get_full_station_name("CS013LBA", 'LBA_OUTER') - 'CS013LBA' - - >>> get_full_station_name("CS002", 'LBA_OUTER') - 'CS002LBA' - """ - if len(station_name) > 5: - return station_name - elif antenna_set[0:3] in ["LBA", "HBA"]: - station_name += antenna_set[0:3] - - return station_name - -def freq_from_sb(sb: int, band: str) -> float: - """ - Convert central frequency to subband number - - Args: - sb: subband number - band: filter band - - Returns: - float: frequency in Hz - - Example: - >>> freq_from_sb(297, '30_90') - 58007812.5 - """ - if band not in ["10_90", "30_90", "110_190", "170_230", "210_250"]: - return None - - if band == "10_90" or band == "30_90": - clock, zone = 200e6, 1 - elif band == "110_190": - clock, zone = 200e6, 2 - elif band == "170_230": - clock, zone = 160e6, 3 - elif band == "210_250": - clock, zone = 200e6, 3 - - sb_bandwidth = 0.5 * clock / 512. - freq_offset = 0.5 * clock * (zone -1) - freq = (sb * sb_bandwidth) + freq_offset - return freq - -def get_station_xyz(station_name: str, antenna_set: str, db): - """ - Get XYZ coordinates for the relevant subset of antennas in a station. - The XYZ system is defined as the PQR system rotated along the R axis to make - the Q-axis point towards local north. - - Args: - station_name: Station name, e.g. 'DE603LBA' or 'DE603' - antenna_set: antenna_set, e.g. LBA_OUTER - db: instance of LofarAntennaDatabase from lofarantpos - - Returns: - np.array: Antenna xyz, shape [n_ant, 3] - np.array: rotation matrix pqr_to_xyz, shape [3, 3] - - Example: - >>> from lofarantpos.db import LofarAntennaDatabase - >>> db = LofarAntennaDatabase() - >>> xyz, _ = get_station_xyz("DE603", "LBA_OUTER", db) - >>> xyz.shape - (96, 3) - >>> f"{xyz[0, 0]:.7f}" - '2.7033776' - - >>> xyz, _ = get_station_xyz("LV614", "HBA", db) - >>> xyz.shape - (96, 3) - """ - station_pqr = get_station_pqr(station_name, antenna_set, db) - - station_name = get_full_station_name(station_name, antenna_set) - - rotation = db.rotation_from_north(station_name) - - pqr_to_xyz = np.array([[np.cos(-rotation), -np.sin(-rotation), 0], - [np.sin(-rotation), np.cos(-rotation), 0], - [0, 0, 1]]) - - station_xyz = (pqr_to_xyz @ station_pqr.T).T - - return station_xyz, pqr_to_xyz - -def get_station_pqr(station_name: str, antenna_set: str, db): - """ - Get PQR coordinates for the relevant subset of antennas in a station. - - Args: - station_name: Station name, e.g. 'DE603LBA' or 'DE603' - antenna_set: antenna_set, e.g. LBA_INNER - db: instance of LofarAntennaDatabase from lofarantpos - - Example: - >>> from lofarantpos.db import LofarAntennaDatabase - >>> db = LofarAntennaDatabase() - >>> pqr = get_station_pqr("DE603", "LBA_OUTER", db) - >>> pqr.shape - (96, 3) - >>> pqr[0, 0] - 1.7434713 - - >>> pqr = get_station_pqr("LV614", "HBA", db) - >>> pqr.shape - (96, 3) - """ - full_station_name = get_full_station_name(station_name, antenna_set) - station_type = get_station_type(full_station_name) - - all_pqr = db.antenna_pqr(full_station_name) - if "LBA" in antenna_set: - if antenna_set == "LBA_OUTER": - station_pqr = all_pqr[48:, :] - elif antenna_set == "LBA_INNER": - station_pqr = all_pqr[:48, :] - elif antenna_set == "LBA_SPARSE_EVEN": - station_pqr = np.ravel(np.column_stack((all_pqr[:48:2], all_pqr[49::2]))).reshape(48, 3) - elif antenna_set == "LBA_SPARSE_ODD": - station_pqr = np.ravel(np.column_stack((all_pqr[1:48:2], all_pqr[48::2]))).reshape(48, 3) - else: - station_pqr = all_pqr - elif "HBA" in antenna_set: - selected_dipole_config = { - 'intl': GENERIC_INT_201512, 'remote': GENERIC_REMOTE_201512, 'core': GENERIC_CORE_201512 - } - selected_dipoles = selected_dipole_config[station_type] + \ - np.arange(len(selected_dipole_config[station_type])) * 16 - single_dipole_pqr = db.hba_dipole_pqr(full_station_name)[selected_dipoles] - if antenna_set == "HBA_SINGLE": - station_pqr = single_dipole_pqr - elif antenna_set == "HBA0_SINGLE": - station_pqr = single_dipole_pqr[:24, :] - elif antenna_set == "HBA1_SINGLE": - station_pqr = single_dipole_pqr[24:, :] - elif antenna_set == "HBA0": - station_pqr = all_pqr[:24, :] - elif antenna_set == "HBA1": - station_pqr = all_pqr[24:, :] - else: - station_pqr = all_pqr - - return station_pqr.astype('float32') - diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..5e67506dcfb737ccfee89a48957f34088d44c3b9 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,13 @@ +[build-system] +requires = [ + "setuptools>=70.0", + "setuptools_scm[toml]>=8.1", + "wheel" +] +build-backend = "setuptools.build_meta" + +[tool.setuptools_scm] +version_file = "lofty/_version.py" + +[tool.pylint] +ignore = "_version.py" diff --git a/requirements-cupy.txt b/requirements-cupy.txt new file mode 100644 index 0000000000000000000000000000000000000000..4c6e89f3b2f4f48b52e3fbe5c8a4fea896023b25 --- /dev/null +++ b/requirements-cupy.txt @@ -0,0 +1,3 @@ +cupy-cuda12x +# cupy-rocm-5-0 + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..f0cac92dbaca0b5e404b833fcd909df56de5a034 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,16 @@ +numpy +dask +jax +# jax[cuda12] +jaxlib +xarray +astropy +matplotlib +fastplotlib +glfw +h5py +lofarantpos +tqdm +nptyping +numba +ddt diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000000000000000000000000000000000000..5fd0524479b8e66f63be5ef4decee3cd75bf6f07 --- /dev/null +++ b/setup.cfg @@ -0,0 +1,49 @@ +[metadata] +name = lofty +description = LOFAR2.0 statistics analysis tools +long_description = file: README.md +long_description_content_type = text/markdown +url = https://git.astron.nl/bassa/lofty +license = GNU General Public License v3.0 +classifiers = + Development Status :: 4- Beta + Intended Audience :: Developers + Intended Audience :: Science/Research + License :: OSI Approved :: GNU General Public License v3.0 + Operating System :: OS Independent + Programming Language :: Python + Programming Language :: Python :: 3 + Programming Language :: Python :: 3 :: Only + Programming Language :: Python :: 3.9 + Programming Language :: Python :: 3.10 + Programming Language :: Python :: 3.11 + Programming Language :: Python :: 3.12 + Programming Language :: Python :: 3.13 + Topic :: Scientific/Engineering + Topic :: Scientific/Engineering :: Astronomy + +[options] +include_package_data = true +packages = find: +python_requires = >=3.9 +install_requires = file: requirements.txt + +[options.package_data] +* = *.npy + +[options.extras_require] +cupy = file: requirements-cupy.txt + +[options.entry_points] +console_scripts = + bst-plot = lofty.entry.plot.bst:main + sst-plot = lofty.entry.plot.sst:main + xst-plot = lofty.entry.plot.xst:main + xst-image = lofty.entry.imaging.xst:main + convert = lofty.entry.conversion:main + animate = lofty.entry.imaging.animate:main + +[flake8] +max-line-length = 88 +extend-ignore = E203 +exclude=.venv,.git,.tox,dist,docs,*lib/python*,*egg,_version.py diff --git a/setup.py b/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..b908cbe55cb344569d32de1dfc10ca7323828dc5 --- /dev/null +++ b/setup.py @@ -0,0 +1,3 @@ +import setuptools + +setuptools.setup() diff --git a/sstplot.py b/sstplot.py deleted file mode 100644 index d72c0591af7c8ba41171a20e5e56b88fa574955a..0000000000000000000000000000000000000000 --- a/sstplot.py +++ /dev/null @@ -1,40 +0,0 @@ -#!/usr/bin/env python3 -import os -import glob -import argparse - -import numpy as np -import matplotlib.pyplot as plt -import matplotlib.dates as mdates -from matplotlib.ticker import FormatStrFormatter, MultipleLocator - -from lofty.io import read_sst, read_minio_sst -from lofty.plotting import plot_sst_timeseries, plot_sst_bandpass, plot_sst_dynamic_spectrum - -if __name__ == "__main__": - # Read command line arguments - parser = argparse.ArgumentParser(description="Plot SST data") - parser.add_argument("-m", "--minio", help="Minio HDF5", action="store_true") - parser.add_argument("-d", "--dynspec", help="Plot dynamic spectrum", action="store_true") - parser.add_argument("-t", "--timeseries", help="Plot timeseries", action="store_true") - parser.add_argument("-b", "--bandpass", help="Plot bandpass", action="store_true") - parser.add_argument("-s", "--subband", help="Subband to plot timeseries [int, default: 300]", default=300, type=int) - parser.add_argument("-p", "--pol", help="Polarization to plot dynamic spectrum [stokes/X/Y, default: stokes]", default="stokes") - parser.add_argument("-a", "--ant", help="Antenna to plot dynamic spectrum [int (-1 for mean), default: -1]", default=-1, type=int) - parser.add_argument("filenames", help="SST filenames", nargs="*", metavar="FILE") - args = parser.parse_args() - - print(args.filenames) - print(args.pol) - - # Read data - if args.minio: - metadata, data = read_minio_sst(args.filenames) - else: - metadata, data = read_sst(args.filenames) - - plot_sst_dynamic_spectrum(metadata, data, args.ant, args.pol) - - plot_sst_bandpass(metadata, data) - - plot_sst_timeseries(metadata, data, args.subband) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/base.py b/tests/base.py new file mode 100644 index 0000000000000000000000000000000000000000..7cadf1d5a623aa49235292d46b23a76c89b70a93 --- /dev/null +++ b/tests/base.py @@ -0,0 +1,30 @@ +# Copyright (C) 2024 Corne Lukken +# SPDX-License-Identifier: GPL-3.0-or-later + +import sys +import unittest +import logging + +logger = logging.getLogger() + + +def configure_logging(): + hdlr_info = logging.StreamHandler(sys.stdout) + formatter = logging.Formatter( + fmt="[%(asctime)s][PID:%(process)d, TID:%(thread)d][%(levelname)s][" + "%(module)s.%(funcName)s:%(lineno)d] %(message)s" + ) + hdlr_info.setFormatter(formatter) + logging.getLogger("matplotlib").setLevel(logging.WARN) + logger.addHandler(hdlr_info) + # logging.basicConfig(handlers=[hdlr_info], level=logging.DEBUG) + + +configure_logging() + + +class BaseTestCase(unittest.TestCase): + """Test base class.""" + + def setUp(self): + super().setUp() diff --git a/tests/image_16_16.npy b/tests/image_16_16.npy new file mode 100644 index 0000000000000000000000000000000000000000..7041a2f29a7a898991891ffa1226b465ad58a53f Binary files /dev/null and b/tests/image_16_16.npy differ diff --git a/tests/image_205_205.npy b/tests/image_205_205.npy new file mode 100644 index 0000000000000000000000000000000000000000..8861c60890817dbea2d71b697a06b85f9dc2ea89 Binary files /dev/null and b/tests/image_205_205.npy differ diff --git a/tests/image_256_256.npy b/tests/image_256_256.npy new file mode 100644 index 0000000000000000000000000000000000000000..553fa964e894a497d9687f273159ac0713db5a06 Binary files /dev/null and b/tests/image_256_256.npy differ diff --git a/tests/image_32_32.npy b/tests/image_32_32.npy new file mode 100644 index 0000000000000000000000000000000000000000..1273acb8d8c5788437c37a3e464e56d758bb540b Binary files /dev/null and b/tests/image_32_32.npy differ diff --git a/tests/image_512_512.npy b/tests/image_512_512.npy new file mode 100644 index 0000000000000000000000000000000000000000..08c6f9e61fce6d9a813c371dbe27f37f2b6e96f0 Binary files /dev/null and b/tests/image_512_512.npy differ diff --git a/tests/image_64_64.npy b/tests/image_64_64.npy new file mode 100644 index 0000000000000000000000000000000000000000..5c082093b2172f0e1ed88620d393c04baaf854ef Binary files /dev/null and b/tests/image_64_64.npy differ diff --git a/tests/image_95_95.npy b/tests/image_95_95.npy new file mode 100644 index 0000000000000000000000000000000000000000..38257fb0bca8dcef8c3b1dc211d73d1b1d62ddb2 Binary files /dev/null and b/tests/image_95_95.npy differ diff --git a/tests/measurements/__init__.py b/tests/measurements/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/measurements/data.py b/tests/measurements/data.py new file mode 100644 index 0000000000000000000000000000000000000000..e38f6341b622ca3c47acb168b1eac02dfc3fc92f --- /dev/null +++ b/tests/measurements/data.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass + + +@dataclass +class MeasureData: + seconds: float + + +@dataclass +class ResultData: + min: float = float("inf") + max: float = float("-inf") + mean: float = 0.0 + std: float = 0.0 diff --git a/tests/measurements/measure.py b/tests/measurements/measure.py new file mode 100644 index 0000000000000000000000000000000000000000..3c2d0becc0c64a16c43d6ad649806fc9568f8e17 --- /dev/null +++ b/tests/measurements/measure.py @@ -0,0 +1,46 @@ +# Copyright (C) 2024 Corne Lukken +# SPDX-License-Identifier: GPL-3.0-or-later + +import logging +import time +from typing import Callable +import statistics + +from tests.measurements.data import MeasureData, ResultData +from tests.measurements.settings import Settings + +logger = logging.getLogger() + + +class Measure: + + def __init__(self, settings: Settings = None): + self.measures: list[MeasureData] = [] + + if settings is None: + settings = Settings() + self.settings = settings + + def _add_measure(self, seconds: float): + self.measures.append(MeasureData(seconds=seconds)) + + def run(self, fn: Callable[[int, int], None]): + """Run the function under test and measure""" + + for i in range(self.settings.iterations): + start_time = time.time() + result = fn(self.settings.image_size_x, self.settings.image_size_y) + if hasattr(result, "block_until_ready"): + result.block_until_ready() + self._add_measure(time.time() - start_time) + + def compute(self) -> ResultData: + result = ResultData() + result.mean = statistics.mean([x.seconds for x in self.measures]) + result.std = statistics.stdev([x.seconds for x in self.measures]) + for measure in self.measures: + if measure.seconds < result.min: + result.min = measure.seconds + if measure.seconds > result.max: + result.max = measure.seconds + return result diff --git a/tests/measurements/settings.py b/tests/measurements/settings.py new file mode 100644 index 0000000000000000000000000000000000000000..760eeda47cad0b7c2a2c20a2e8c3ec9e862036b8 --- /dev/null +++ b/tests/measurements/settings.py @@ -0,0 +1,11 @@ +# Copyright (C) 2024 Corne Lukken +# SPDX-License-Identifier: GPL-3.0-or-later + +import dataclasses + + +@dataclasses.dataclass +class Settings: + image_size_x: int = 205 + image_size_y: int = 205 + iterations: int = 3 diff --git a/tests/requirements.txt b/tests/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..6fb8d2d7bed4b468dfd3da0b579c6386380a833a --- /dev/null +++ b/tests/requirements.txt @@ -0,0 +1,8 @@ +autopep8 >= 1.7.0 # MIT +black >= 22.0.0 # MIT +build >= 0.8.0 # MIT +flake8 >= 5.0.0 # MIT +pylint >= 2.15.0 # GPLv2 +pytest >= 7.0.0 # MIT +pytest-cov >= 3.0.0 # MIT +ddt diff --git a/tests/test_antenna_field.py b/tests/test_antenna_field.py new file mode 100644 index 0000000000000000000000000000000000000000..c9705323596c77af92e19e2328474b19b48f6387 --- /dev/null +++ b/tests/test_antenna_field.py @@ -0,0 +1,31 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +"""Testing of the station name extraction""" + +import logging + +from lofty.io import extract_antenna_field_from_filename + +from tests.base import BaseTestCase + +logger = logging.getLogger() + + +class TestAntennaField(BaseTestCase): + """Test antennafield file extraction""" + + def test_antenna_field_none(self): + name = extract_antenna_field_from_filename("1201jaka") + + self.assertEqual(None, name) + + def test_antenna_field_one(self): + name = extract_antenna_field_from_filename("LBA") + + self.assertEqual("LBA", name) + + def test_antenna_field_multiple(self): + name = extract_antenna_field_from_filename("HBAHBA") + + self.assertEqual(None, name) diff --git a/tests/test_station_name.py b/tests/test_station_name.py new file mode 100644 index 0000000000000000000000000000000000000000..730218510aea7c3c68d128a0bb19af72bb764a37 --- /dev/null +++ b/tests/test_station_name.py @@ -0,0 +1,31 @@ +# Copyright (C) 2024 ASTRON (Netherlands Institute for Radio Astronomy) +# SPDX-License-Identifier: GPL-3.0-or-later + +"""Testing of the station name extraction""" + +import logging + +from lofty.io import extract_station_name_from_filename + +from tests.base import BaseTestCase + +logger = logging.getLogger() + + +class TestStationName(BaseTestCase): + """Test Station Name file extraction""" + + def test_station_name_none(self): + name = extract_station_name_from_filename("1201jaka") + + self.assertEqual(None, name) + + def test_station_name_core(self): + name = extract_station_name_from_filename("cs001") + + self.assertEqual("cs001", name) + + def test_station_name_international(self): + name = extract_station_name_from_filename("DE603") + + self.assertEqual("DE603", name) diff --git a/tests/test_xst_imaging.py b/tests/test_xst_imaging.py new file mode 100644 index 0000000000000000000000000000000000000000..22a11bdf2d12231425f56beb764f336bfbfeb1e4 --- /dev/null +++ b/tests/test_xst_imaging.py @@ -0,0 +1,234 @@ +# Copyright (C) 2023 ASTRON (Netherlands Institute for Radio Astronomy) +# Copyright (C) 2024 Corne Lukken +# SPDX-License-Identifier: GPL-3.0-or-later + +"""Testing of the xst imager""" + +from importlib.resources import files +from typing import Callable, Optional +import logging +import os + +import numpy as np +from matplotlib import pyplot as plt +from ddt import ddt, data, unpack + +from lofty.imaging import ( + sky_imager_simple, + sky_imager_numpy_exp, + sky_imager_numpy_ravel_32, + sky_imager_jax, + sky_imager_jax_jit, + sky_imager_jax_ravel, + sky_imager_jax_ravel_jit, + sky_imager_jax_real, + sky_imager_jax_real_jit, + sky_imager_jax_ravel_real_jit, + sky_imager_jax_precompute_real_jit, +) +from lofty.imaging._numpy import sky_imager_numpy_ravel_real, sky_imager_numpy_real + +from tests.base import BaseTestCase +from tests.measurements.measure import Measure +from tests.measurements.settings import Settings + +logger = logging.getLogger() + + +@ddt +class TestXstImaging(BaseTestCase): + """Test performance and verify correctness of xst all-sky imaging""" + + @classmethod + def setUpClass(cls): + cls.measure_performance = False + if os.getenv("MEASURE_PERFORMANCE"): + cls.measure_performance = True + + if cls.measure_performance: + logger.info("Running performance analysis, this might take a while") + else: + logger.info( + "Verifying image results only, use `tox -e performance` to analyse " + "performace" + ) + + def setUp(self): + self.baselines = np.load(files("lofty.data").joinpath("baselines.npy")) + self.visibilities = np.load(files("lofty.data").joinpath("vis.npy")) + self.frequency = np.load(files("lofty.data").joinpath("freq.npy")) + + # Storing the results of an imaging computation in a file + # self.image = sky_imager_simple( + # self.visibilities[0], self.baselines, self.frequency, 512, 512 + # ) + # np.save("image", self.image) + + def _measure_imager(self, fn: Callable, name: str, x: int = 205, y: int = 205): + if not self.measure_performance: + return + + measure = Measure(Settings(iterations=3, image_size_x=x, image_size_y=y)) + for x in range(10): + measure.run( + lambda var_x, var_y: fn( + self.visibilities[x % 10], + self.baselines, + self.frequency, + var_x, + var_y, + ) + ) + result = measure.compute() + logger.info( + "[%s]: min: %.4f max: %.4f, mean: %.4f, stddev: %.4f", + name, + result.min, + result.max, + result.mean, + result.std, + ) + + def _verify_imager( + self, fn: Callable, fn2: Optional[Callable] = None, x: int = 205, y: int = 205 + ): + """Compare against precomputed stored results or two implementations + + :warning: This method relies on L, M, and N being derived in the same way across + implementations to have a meaningful comparison. The order in which + pixels are computed is also essential for the comparison, otherwise + the result will be a mirrored image. + """ + + if fn2 is None: + reference_image = np.load( + os.path.join(os.path.dirname(__file__), f"image_{x}_{y}.npy") + ) + else: + + def reference_image(var_x, var_y): + return fn2( + self.visibilities[0], self.baselines, self.frequency, var_x, var_y + ) + + reference_image = reference_image(x, y) + + def result_image(var_x, var_y): + return fn( + self.visibilities[0], self.baselines, self.frequency, var_x, var_y + ) + + result_image = result_image(x, y) + + # Can be used for diagnostic comparison of differences + # plt.imshow(result_image, cmap="hot") + # plt.colorbar() + # plt.show() + + # Create a circle as mask just below unit length. and remove those results from + # the evaluation. The (all-sky imaging) computation does not resolve beyond the + # horizon. This boundary can shift slightly due to numerical error between + # implementations. + npix_l, npix_m = np.meshgrid(np.linspace(-1, 1, x), np.linspace(1, -1, y)) + c = npix_l**2 + npix_m**2 < 0.99 + reference_image = np.where(c, reference_image, float("nan")) + result_image = np.where(c, result_image, float("nan")) + + np.testing.assert_allclose(reference_image, result_image, rtol=1e-04) + + # def test_sky_imager_simple(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of simple imager""" + # self._verify_imager(sky_imager_simple) + # self._measure_imager(sky_imager_simple, "sky imager simple") + # + # def test_sky_imager_numpy_exp(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of numpy exp imager""" + # self._verify_imager(sky_imager_numpy_exp) + # self._measure_imager(sky_imager_numpy_exp, "sky imager numpy exp") + # + # def test_sky_imager_numpy_ravel_32(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of ravelling numpy imager""" + # self._verify_imager(sky_imager_numpy_ravel_32) + # self._measure_imager(sky_imager_numpy_ravel_32, "sky imager numpy ravel") + # + # def test_sky_imager_numpy_real(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of numpy real imager""" + # self._verify_imager(sky_imager_numpy_real) + # self._measure_imager(sky_imager_numpy_real, "sky imager numpy real") + # + # def test_sky_imager_numpy_ravel_real(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of ravelling numpy imager""" + # self._verify_imager(sky_imager_numpy_ravel_real, x=x, y=y) + # self._measure_imager( + # sky_imager_numpy_ravel_real, "sky imager numpy ravel real", x=x, y=y + # ) + # + # def test_sky_imager_jax(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of basic jax imager""" + # self._verify_imager(sky_imager_jax) + # self._measure_imager(sky_imager_jax, "sky imager jax") + # + # @data((16, 16), (32, 32), (64, 64), (95, 95), (205, 205), (256, 256), (512, 512)) + # @unpack + # def test_sky_imager_jax_jit(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of jitted basic jax imager""" + # self._verify_imager(sky_imager_jax_jit, x=x, y=y) + # self._measure_imager(sky_imager_jax_jit, "sky imager jax jit", x=x, y=y) + # + # def test_sky_imager_jax_ravel(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of ravelling jax imager""" + # self._verify_imager(sky_imager_jax_ravel) + # self._measure_imager(sky_imager_jax_ravel, "sky imager jax ravel") + # + # @data((16, 16), (32, 32), (64, 64), (95, 95), (205, 205), (256, 256), (512, 512)) + # @unpack + # def test_sky_imager_jax_ravel_jit(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of jitted ravelling jax imager""" + # self._verify_imager(sky_imager_jax_ravel_jit, x=x, y=y) + # self._measure_imager( + # sky_imager_jax_ravel_jit, "sky imager jit jax ravel", x=x, y=y + # ) + # + # def test_sky_imager_jax_real(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of jax real imager""" + # self._verify_imager(sky_imager_jax_real, x=x, y=y) + # self._measure_imager(sky_imager_jax_real, "sky imager jax real") + # + # @data((16, 16), (32, 32), (64, 64), (95, 95), (205, 205), (256, 256), (512, 512)) + # @unpack + # def test_sky_imager_jax_real_jit(self, x: int = 205, y: int = 205): + # """Testing the performance and verifying of jitted jax real imager""" + # self._verify_imager(sky_imager_jax_real_jit, x=x, y=x) + # self._measure_imager( + # sky_imager_jax_real_jit, "sky imager jitted jax real", x=x, y=y + # ) + # + # @data((16, 16), (32, 32), (64, 64), (95, 95), (205, 205), (256, 256), (512, 512)) + # @unpack + # def test_sky_imager_jax_ravel_real_jit(self, x: int = 256, y: int = 256): + # """Testing the performance and verifying of jitted jax real imager""" + # self._verify_imager(sky_imager_jax_ravel_real_jit, x=x, y=x) + # self._measure_imager( + # sky_imager_jax_ravel_real_jit, "sky imager jitted jax ravel real", x=x, y=y + # ) + + # @data((16, 16), (32, 32), (64, 64), (95, 95), (205, 205), (256, 256), (512, 512)) + # @unpack + def test_sky_imager_jax_precompute_real_jit(self, x: int = 205, y: int = 205): + """""" + self._verify_imager(sky_imager_jax_precompute_real_jit, x=x, y=x) + self._measure_imager( + sky_imager_jax_precompute_real_jit, + "sky imager jitted jax precomputed real", + x=x, + y=y, + ) + + # def test_sky_imager_cupy(self, x: int = 95, y: int = 95): + # try: + # from lofty.imaging import sky_imager_cupy + # + # self._verify_imager(sky_imager_cupy, x=x, y=y) + # self._measure_imager(sky_imager_cupy, "sky imager cupy", x=95, y=95) + # except ImportError: + # logger.info("Skipping sky_imager_cupy") diff --git a/tox.ini b/tox.ini new file mode 100644 index 0000000000000000000000000000000000000000..0624c2a917a140394cd77a1110d742e47fdbcc65 --- /dev/null +++ b/tox.ini @@ -0,0 +1,77 @@ +[tox] +# Generative environment list to test all supported Python versions +envlist = py3{9,10,11,12,13},black,pep8,pylint +min_version = 4.3.3 +requires = + tox-ignore-env-name-mismatch >= 0.2.0 + +[testenv] +usedevelop = True +package = wheel +wheel_build_env = .pkg +setenv = + LANGUAGE=en_US + LC_ALL=en_US.UTF-8 + PYTHONWARNINGS=default::DeprecationWarning + JAX_TRACEBACK_FILTERING=off +deps = + -r{toxinidir}/requirements.txt + -r{toxinidir}/tests/requirements.txt +commands = + {envpython} --version + {envpython} -m pytest -v -rP --log-level=INFO {posargs} + +[testenv:performance] +setenv = + MEASURE_PERFORMANCE=1 + +[testenv:cupy] +deps = + -r{toxinidir}/requirements.txt + -r{toxinidir}/tests/requirements.txt + -r{toxinidir}/requirements-cupy.txt + +[testenv:cupy-performance] +setenv = + MEASURE_PERFORMANCE=1 +deps = + -r{toxinidir}/requirements.txt + -r{toxinidir}/tests/requirements.txt + -r{toxinidir}/requirements-cupy.txt + +[testenv:coverage] +commands = + {envpython} --version + {envpython} -m pytest --cov-report term --cov-report xml --cov-report html --cov=lofty + +# Use generative name and command prefixes to reuse the same virtualenv +# for all linting jobs. +[testenv:{pep8,black,pylint,format}] +usedevelop = False +package = editable +envdir = {toxworkdir}/linting +commands = + pep8: {envpython} -m flake8 --version + pep8: {envpython} -m flake8 lofty tests + black: {envpython} -m black --version + black: {envpython} -m black --check --diff lofty tests + pylint: {envpython} -m pylint --version + pylint: {envpython} -m pylint lofty tests + format: {envpython} -m autopep8 -v -aa --in-place --recursive lofty + format: {envpython} -m autopep8 -v -aa --in-place --recursive tests + format: {envpython} -m black -v lofty tests + +[testenv:docs] +; unset LC_ALL / LANGUAGE from testenv, would fail sphinx otherwise +setenv = +deps = + -r{toxinidir}/requirements.txt + -r{toxinidir}/docs/requirements.txt +changedir = {toxinidir} +commands = + {envpython} docs/cleanup.py + sphinx-build -b html docs/source docs/build/html + +[testenv:build] +usedevelop = False +commands = {envpython} -m build