From e6036ee73f225e4fe32b1cdd5defc54bceb89a31 Mon Sep 17 00:00:00 2001 From: Steven Murray Date: Wed, 8 Jun 2022 07:59:28 -0700 Subject: [PATCH 1/6] maint: don't point to git repos for deps --- .vscode/configurationCache.log | 1 + .vscode/dryrun.log | 6 ++++++ .vscode/settings.json | 3 +++ setup.py | 4 ++-- 4 files changed, 12 insertions(+), 2 deletions(-) create mode 100644 .vscode/configurationCache.log create mode 100644 .vscode/dryrun.log create mode 100644 .vscode/settings.json diff --git a/.vscode/configurationCache.log b/.vscode/configurationCache.log new file mode 100644 index 00000000..618ed811 --- /dev/null +++ b/.vscode/configurationCache.log @@ -0,0 +1 @@ +{"buildTargets":[],"launchTargets":[],"customConfigurationProvider":{"workspaceBrowse":{"browsePath":[]},"fileIndex":[]}} \ No newline at end of file diff --git a/.vscode/dryrun.log b/.vscode/dryrun.log new file mode 100644 index 00000000..cdd55217 --- /dev/null +++ b/.vscode/dryrun.log @@ -0,0 +1,6 @@ +make --dry-run --always-make --keep-going --print-directory +make: Entering directory '/home/steven/Documents/Projects/radio/EOR/HERA/hera_pspec' +make: Leaving directory '/home/steven/Documents/Projects/radio/EOR/HERA/hera_pspec' + +make: *** No targets specified and no makefile found. Stop. + diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 00000000..65e1ec07 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "makefile.extensionOutputFolder": "./.vscode" +} \ No newline at end of file diff --git a/setup.py b/setup.py index 73b29e0c..6067f846 100644 --- a/setup.py +++ b/setup.py @@ -44,8 +44,8 @@ def package_files(package_dir, subdirectory): 'astropy>=2.0', 'pyyaml', 'h5py', - 'uvtools @ git+https://github.com/HERA-Team/uvtools', - 'hera_cal @ git+https://github.com/HERA-Team/hera_cal' + 'uvtools', + 'hera-calibration' ], 'include_package_data': True, 'scripts': ['scripts/pspec_run.py', 'scripts/pspec_red.py', From e8b21766fbe63af38ec997e03bbdfb7a043ef86d Mon Sep 17 00:00:00 2001 From: Steven Murray Date: Wed, 8 Jun 2022 08:41:46 -0700 Subject: [PATCH 2/6] maint: add pyproject.toml and setuptools_scm --- .gitignore | 1 + MANIFEST.in | 6 -- hera_pspec/VERSION | 1 - hera_pspec/__init__.py | 17 +++- hera_pspec/tests/test_version.py | 73 ---------------- hera_pspec/utils.py | 20 +++++ hera_pspec/version.py | 137 ------------------------------- pyproject.toml | 8 ++ setup.cfg | 101 +++++++++++++++++++++++ setup.py | 58 +------------ 10 files changed, 147 insertions(+), 275 deletions(-) delete mode 100644 MANIFEST.in delete mode 100644 hera_pspec/VERSION delete mode 100644 hera_pspec/tests/test_version.py delete mode 100644 hera_pspec/version.py create mode 100644 pyproject.toml create mode 100644 setup.cfg diff --git a/.gitignore b/.gitignore index 6ac93e2b..5fa4ebb2 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,4 @@ eggs/ *egg-info *GIT_INFO *ipynb_checkpoints +hera_pspec/_version.py \ No newline at end of file diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index 7c110647..00000000 --- a/MANIFEST.in +++ /dev/null @@ -1,6 +0,0 @@ -include *.md -include LICENSE -include hera_pspec/VERSION -include hera_pspec/GIT_INFO -include pipelines/*/*.yaml - diff --git a/hera_pspec/VERSION b/hera_pspec/VERSION deleted file mode 100644 index 0d91a54c..00000000 --- a/hera_pspec/VERSION +++ /dev/null @@ -1 +0,0 @@ -0.3.0 diff --git a/hera_pspec/__init__.py b/hera_pspec/__init__.py index b99bf143..44f206e5 100644 --- a/hera_pspec/__init__.py +++ b/hera_pspec/__init__.py @@ -1,6 +1,11 @@ """ __init__.py file for hera_pspec """ +try: + from importlib.metadata import version, PackageNotFoundError +except ImportError: + from importlib_metadata import version, PackageNotFoundError + from . import version, conversions, grouping, pspecbeam, plot, pstokes, testing, utils from . import uvpspec_utils as uvputils @@ -11,4 +16,14 @@ from .parameter import PSpecParam from .pspecbeam import PSpecBeamUV, PSpecBeamGauss, PSpecBeamFromArray -__version__ = version.version +try: + from ._version import version as __version__ +except ModuleNotFoundError: # pragma: no cover + try: + __version__ = version("hera_pspec") + except PackageNotFoundError: + # package is not installed + __version__ = "unknown" + +del version +del PackageNotFoundError \ No newline at end of file diff --git a/hera_pspec/tests/test_version.py b/hera_pspec/tests/test_version.py deleted file mode 100644 index 0d581ab3..00000000 --- a/hera_pspec/tests/test_version.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Tests for version.py.""" -import os -import sys -import pytest - -try: - # Python 2 - from cStringIO import StringIO -except: - # Python 3 - from io import StringIO -import hera_pspec -from .. import version -import json - - -def test_main(): - version_info = version.construct_version_info() - - saved_stdout = sys.stdout - try: - out = StringIO() - sys.stdout = out - hera_pspec.version.main() - output = out.getvalue() - assert (output == 'Version = {v}\ngit origin = {o}\n' - 'git branch = {b}\ngit description = {d}\n' - .format(v=version_info['version'], - o=version_info['git_origin'], - b=version_info['git_branch'], - d=version_info['git_description'])) - - - finally: - sys.stdout = saved_stdout - - # Test history string function - history = hera_pspec.version.history_string() - - -def test_get_gitinfo_file(): - dir = version.hera_pspec_dir - - git_file = os.path.join(dir, 'GIT_INFO') - if not os.path.exists(git_file): - # write a file to read in - temp_git_file = os.path.join(dir, 'GIT_INFO') - version_info = version.construct_version_info() - data = [version_info['git_origin'], version_info['git_origin'], - version_info['git_origin'], version_info['git_origin']] - with open(temp_git_file, 'w') as outfile: - json.dump(data, outfile) - git_file = temp_git_file - - with open(git_file) as data_file: - data = [version._unicode_to_str(x) for x in json.loads(data_file.read().strip())] - git_origin = data[0] - git_hash = data[1] - git_description = data[2] - git_branch = data[3] - - test_file_info = { - 'git_origin': git_origin, 'git_hash': git_hash, - 'git_description': git_description, 'git_branch': git_branch - } - - if 'temp_git_file' in locals(): - file_info = version._get_gitinfo_file(git_file=temp_git_file) - os.remove(temp_git_file) - else: - file_info = version._get_gitinfo_file() - - assert file_info == test_file_info diff --git a/hera_pspec/utils.py b/hera_pspec/utils.py index 98a9c20e..ef99bcef 100644 --- a/hera_pspec/utils.py +++ b/hera_pspec/utils.py @@ -11,6 +11,8 @@ import uvtools as uvt import argparse from .conversions import Cosmo_Conversions +import inspect +from . import __version__ def cov(d1, w1, d2=None, w2=None, conj_1=False, conj_2=True): @@ -1517,3 +1519,21 @@ def apply_P_SN_correction(uvp, P_SN='P_SN', P_N='P_N'): corr[np.isnan(corr)] = np.inf # apply correction uvp.stats_array[P_SN][spw] *= corr + + +def history_string(notes=''): + """ + Creates a standardized history string that all functions that write to + disk can use. Optionally add notes. + """ + notes = f"""\n\nNotes:\n{notes}""" if notes else "" + + stack = inspect.stack()[1] + history = f""" + ------------ + This file was produced by the function {stack[3]}() in {stack[1]} using version {__version__} + {notes} + ------------ + """ + return history + diff --git a/hera_pspec/version.py b/hera_pspec/version.py deleted file mode 100644 index ab58aa57..00000000 --- a/hera_pspec/version.py +++ /dev/null @@ -1,137 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2019 the HERA Project -# Licensed under the MIT License - -from __future__ import print_function, division, absolute_import - -import inspect -import json -import os -import subprocess -import sys - -hera_pspec_dir = os.path.dirname(os.path.realpath(__file__)) - - - -def _get_git_output(args, capture_stderr=False): - """ - Get output from git, ensuring that it is of the ``str`` type, not bytes. - """ - - argv = ['git', '-C', hera_pspec_dir] + args - - if capture_stderr: - data = subprocess.check_output(argv, stderr=subprocess.STDOUT) - else: - data = subprocess.check_output(argv) - - data = data.strip() - - return data.decode('utf8') - - -def _get_gitinfo_file(git_file=None): - """ - Get saved info from GIT_INFO file that was created when installing package - """ - if git_file is None: - git_file = os.path.join(hera_pspec_dir, 'GIT_INFO') - - with open(git_file) as data_file: - data = [_unicode_to_str(x) for x in json.loads(data_file.read().strip())] - git_origin = data[0] - git_hash = data[1] - git_description = data[2] - git_branch = data[3] - - return {'git_origin': git_origin, 'git_hash': git_hash, - 'git_description': git_description, 'git_branch': git_branch} - - -def _unicode_to_str(u): - return u - - -def construct_version_info(): - version_file = os.path.join(hera_pspec_dir, 'VERSION') - with open(version_file) as f: - version = f.read().strip() - - git_origin = '' - git_hash = '' - git_description = '' - git_branch = '' - - version_info = {'version': version, 'git_origin': '', 'git_hash': '', - 'git_description': '', 'git_branch': ''} - - try: - version_info['git_origin'] = _get_git_output( - ['config', '--get', 'remote.origin.url'], - capture_stderr=True) - version_info['git_hash'] = _get_git_output( - ['rev-parse', 'HEAD'], - capture_stderr=True) - version_info['git_description'] = _get_git_output( - ['describe', '--dirty', '--tag', '--always']) - version_info['git_branch'] = _get_git_output( - ['rev-parse', '--abbrev-ref', 'HEAD'], - capture_stderr=True) - except subprocess.CalledProcessError: # pragma: no cover - try: - # Check if a GIT_INFO file was created when installing package - version_info.update(_get_gitinfo_file()) - except (IOError, OSError): - pass - - return version_info - - -def history_string(notes=''): - """ - Creates a standardized history string that all functions that write to - disk can use. Optionally add notes. - """ - history = '\n------------\nThis file was produced by the function ' \ - + str(inspect.stack()[1][3]) + '()' - # inspect.stack()[1][3] is the name of the function that called this fn - - history += ' in ' + os.path.basename(inspect.stack()[1][1]) + ' using: ' - # inspect.stack()[1][1] is path to the file that contains the function - # that called this function - version_info = construct_version_info() - - for v in sorted(version_info.keys()): - history += '\n ' + v + ': ' + version_info[v] - - if (notes is not None) and (notes != ''): - history += '\n\nNotes:\n' - history += notes - return history + '\n------------\n' - - -def print_version_info(): - """ - Print git/version info. - """ - print('Version = {0}'.format(version)) - print('git origin = {0}'.format(git_origin)) - print('git branch = {0}'.format(git_branch)) - print('git description = {0}'.format(git_description)) - - -version_info = construct_version_info() -version = version_info['version'] -git_origin = version_info['git_origin'] -git_hash = version_info['git_hash'] -git_description = version_info['git_description'] -git_branch = version_info['git_branch'] - - -def main(): - print_version_info() - - -if __name__ == '__main__': - main() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..6d675668 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,8 @@ +[build-system] +requires = ["setuptools>=45", "wheel", "setuptools_scm[toml]>=6.2"] +build-backend = "setuptools.build_meta" + +[tool.setuptools_scm] +write_to = "hera_pspec/_version.py" +parentdir_prefix_version = "hera_pspec-" +fallback_version = "0.0.0" \ No newline at end of file diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 00000000..655d0033 --- /dev/null +++ b/setup.cfg @@ -0,0 +1,101 @@ +# This file is used to configure your project. +# Read more about the various options under: +# http://setuptools.readthedocs.io/en/latest/setuptools.html#configuring-setup-using-setup-cfg-files + +[metadata] +name = hera_pspec +description = HERA Power Spectrum Estimator Code +author = HERA Team +license = BSD +long_description = file: README.md +long_description_content_type = text/x-rst; charset=UTF-8 +url = https://github.com/HERA-Team/hera_pspec +project_urls = + Documentation = https://hera-pspec.readthedocs.io/en/latest/ +# Change if running only on Windows, Mac or Linux (comma-separated) +platforms = OSX,Linux +# Add here all kinds of additional classifiers as defined under +# https://pypi.python.org/pypi?%3Aaction=list_classifiers +classifiers = + Development Status :: 5 - Production/Stable + Programming Language :: Python :: 3.8 + Programming Language :: Python :: 3.9 + Intended Audience :: Science/Research + License :: OSI Approved + Natural Language :: English + Topic :: Scientific/Engineering :: Physics + Topic :: Scientific/Engineering :: Astronomy + +[options] +zip_safe = False +packages = find: +include_package_data = True +scripts = + scripts/pspec_run.py + scripts/pspec_red.py + scripts/bootstrap_run.py + scripts/generate_pstokes_run.py + scripts/auto_noise_run.py + +install_requires = + numpy>=1.14 + scipy + matplotlib>=2.2 + pyuvdata + astropy>=2.0 + pyyaml + h5py + uvtools + hera-calibration + +[options.packages.find] +exclude = + tests + +[options.extras_require] +docs = + sphinx>=1.8 + nbsphinx + ipython + sphinx_autorun + numpydoc>=0.8 + nbsphinx +tests = + coverage>=4.5.1 + pytest>=3.5.1 + pytest-cov>=2.5.1 +dev = + sphinx>=1.8 + numpydoc>=0.8.0 + nbsphinx + ipython + coverage>=4.5.1 + pytest>=3.5.1 + pytest-cov>=2.5.1 + +[tool:pytest] +# Options for py.test: +# Specify command line options as you would do when invoking py.test directly. +# e.g. --cov-report html (or xml) for html/xml output or --junitxml junit.xml +# in order to write a coverage file that can be read by Jenkins. +addopts = + --cov hera_pspec + --cov-config=.coveragerc + --cov-report xml:./coverage.xml + --verbose +norecursedirs = + dist + build + .tox +testpaths = hera_sim/tests + +[aliases] +dists = bdist_wheel + +[bdist_wheel] +# Use this option if your package is pure-python +universal = 1 + +[build_sphinx] +source_dir = docs +build_dir = build/sphinx diff --git a/setup.py b/setup.py index 6067f846..60684932 100644 --- a/setup.py +++ b/setup.py @@ -1,59 +1,3 @@ -import json -import os -import sys - from setuptools import setup -sys.path.append("hera_pspec") -import version - -data = [version.git_origin, version.git_hash, version.git_description, version.git_branch] -with open(os.path.join('hera_pspec', 'GIT_INFO'), 'w') as outfile: - json.dump(data, outfile, default=str) - - -def package_files(package_dir, subdirectory): - # walk the input package_dir/subdirectory - # return a package_data list - paths = [] - directory = os.path.join(package_dir, subdirectory) - for (path, directories, filenames) in os.walk(directory): - for filename in filenames: - path = path.replace(package_dir + '/', '') - paths.append(os.path.join(path, filename)) - return paths - - -data_files = package_files('hera_pspec', 'data') + package_files('hera_pspec', '../pipelines') - -setup_args = { - 'name': 'hera_pspec', - 'author': 'HERA Team', - 'url': 'https://github.com/HERA-Team/hera_pspec', - 'license': 'BSD', - 'version': version.version, - 'description': 'HERA Power Spectrum Estimator Code.', - 'packages': ['hera_pspec'], - 'package_dir': {'hera_pspec': 'hera_pspec'}, - 'package_data': {'hera_pspec': data_files}, - 'install_requires': [ - 'numpy>=1.15', - 'scipy', - 'matplotlib>=2.2' - 'pyuvdata', - 'astropy>=2.0', - 'pyyaml', - 'h5py', - 'uvtools', - 'hera-calibration' - ], - 'include_package_data': True, - 'scripts': ['scripts/pspec_run.py', 'scripts/pspec_red.py', - 'scripts/bootstrap_run.py', - 'scripts/generate_pstokes_run.py', - 'scripts/auto_noise_run.py'], - 'zip_safe': False, -} - -if __name__ == '__main__': - setup(**setup_args) +setup() From e101be1e8ef07fae4aa2ae72a6267babf8564be0 Mon Sep 17 00:00:00 2001 From: Steven Murray Date: Wed, 8 Jun 2022 14:24:34 -0700 Subject: [PATCH 3/6] many linting updates --- .coveragerc | 1 - .flake8 | 38 + .gitignore | 2 +- .pre-commit-config.yaml | 45 + .vscode/configurationCache.log | 2 +- .vscode/dryrun.log | 3 +- .vscode/settings.json | 5 +- CHANGELOG.md | 8 +- README.md | 26 +- ci/hera_pspec_tests.yml | 1 - docs/Makefile | 2 +- docs/conf.py | 117 +- docs/index.rst | 8 +- docs/plot.rst | 18 +- docs/pspec.rst | 60 +- docs/pspecbeam.rst | 18 +- .../error_bar_choices/bibtex.bib | 4 +- .../hera_pspec_error_bar_choice.tex | 76 +- hera_pspec/__init__.py | 2 +- hera_pspec/container.py | 264 +- hera_pspec/conversions.py | 103 +- hera_pspec/grouping.py | 1063 ++++--- hera_pspec/noise.py | 138 +- hera_pspec/parameter.py | 9 +- hera_pspec/plot.py | 697 +++-- hera_pspec/pspecbeam.py | 277 +- hera_pspec/pspecdata.py | 2655 ++++++++++++----- hera_pspec/pstokes.py | 206 +- hera_pspec/testing.py | 8 +- hera_pspec/tests/compare_legacy_oqe.py | 100 +- hera_pspec/tests/test_container.py | 243 +- hera_pspec/tests/test_conversions.py | 64 +- hera_pspec/tests/test_grouping.py | 741 +++-- hera_pspec/tests/test_noise.py | 163 +- hera_pspec/tests/test_plot.py | 606 +++- hera_pspec/tests/test_pspecbeam.py | 418 ++- hera_pspec/tests/test_pspecdata.py | 2380 ++++++++++----- hera_pspec/tests/test_pstokes.py | 51 +- hera_pspec/tests/test_testing.py | 191 +- hera_pspec/tests/test_utils.py | 334 ++- hera_pspec/tests/test_uvpspec.py | 798 +++-- hera_pspec/tests/test_uvpspec_utils.py | 241 +- hera_pspec/tests/test_uvwindow.py | 748 +++-- hera_pspec/utils.py | 574 ++-- hera_pspec/uvpspec.py | 1665 +++++++---- hera_pspec/uvpspec_utils.py | 531 ++-- hera_pspec/uvwindow.py | 751 +++-- pipelines/web_ui/data.js | 10 +- pipelines/web_ui/hera_webui.js | 178 +- pipelines/web_ui/report.html | 48 +- pipelines/web_ui/w3.css | 6 +- pyproject.toml | 21 +- scripts/auto_noise_run.py | 5 +- scripts/bootstrap_run.py | 5 +- scripts/generate_pstokes_run.py | 1 + scripts/psc_merge_spectra.py | 8 +- scripts/pspec_red.py | 104 +- scripts/pspec_run.py | 8 +- setup.cfg | 4 +- 59 files changed, 11139 insertions(+), 5714 deletions(-) create mode 100644 .flake8 create mode 100644 .pre-commit-config.yaml diff --git a/.coveragerc b/.coveragerc index 98f5d134..392a36c6 100644 --- a/.coveragerc +++ b/.coveragerc @@ -3,4 +3,3 @@ omit = */tests/* [report] omit = */tests/* - diff --git a/.flake8 b/.flake8 new file mode 100644 index 00000000..5ff6686d --- /dev/null +++ b/.flake8 @@ -0,0 +1,38 @@ +[flake8] +ignore = + E203 # No space before colon + W503 + C408 # Ignore using dict() function. + D107 # Missing docstring in __init__ (we do it in the class) + D401 # "First line should be in imperative mood" -- this doesn't work for properties, see https://github.com/PyCQA/pydocstyle/issues/301 + A003 # allow method names to be the same as python builtins + RST210 # inline strong start-string without end-string. This is OK in the case of **kwargs in parameters. + D # TODO: add doc checks back in by removing this line! + E501 # TODO: add this back in to ensure lines aren't too long! + T001 # Allow print statements (eventually should use logging instead...) + B008 # allow function calls in argument defaults. + A002 # allow arguments to shadow python builtins +max-line-length = 88 +# TODO: max-complexity should be 18! +max-complexity = 60 +per-file-ignores = + docs/conf.py:D,A + */__init__.py:F401 + scripts/*:T001,T201 +rst-roles = + class + func + mod + data + const + meth + attr + exc + obj +rst-directives = + note + warning + versionadded + versionchanged + deprecated + seealso diff --git a/.gitignore b/.gitignore index 5fa4ebb2..b717abd6 100644 --- a/.gitignore +++ b/.gitignore @@ -11,4 +11,4 @@ eggs/ *egg-info *GIT_INFO *ipynb_checkpoints -hera_pspec/_version.py \ No newline at end of file +hera_pspec/_version.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..437a4116 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,45 @@ +exclude: 'setup.py|hera_pspec/data/' + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.2.0 + hooks: + - id: trailing-whitespace + - id: check-added-large-files + - id: check-ast + - id: check-json + # - id: check-merge-conflict + - id: check-xml + - id: check-yaml + args: ['--unsafe'] # Only check syntax rather than try to load. + - id: debug-statements + - id: end-of-file-fixer + - id: requirements-txt-fixer + - id: mixed-line-ending + args: ['--fix=no'] + + - repo: https://github.com/PyCQA/flake8 + rev: 4.0.1 # pick a git hash / tag to point to + hooks: + - id: flake8 + additional_dependencies: + - flake8-rst-docstrings + - flake8-docstrings + - flake8-builtins + - flake8-logging-format + - flake8-rst-docstrings + - flake8-rst + - flake8-markdown + - flake8-bugbear + - flake8-comprehensions + - flake8-print + + - repo: https://github.com/psf/black + rev: 22.3.0 + hooks: + - id: black + + - repo: https://github.com/pre-commit/pygrep-hooks + rev: v1.9.0 + hooks: + - id: rst-backticks diff --git a/.vscode/configurationCache.log b/.vscode/configurationCache.log index 618ed811..c5d5faf5 100644 --- a/.vscode/configurationCache.log +++ b/.vscode/configurationCache.log @@ -1 +1 @@ -{"buildTargets":[],"launchTargets":[],"customConfigurationProvider":{"workspaceBrowse":{"browsePath":[]},"fileIndex":[]}} \ No newline at end of file +{"buildTargets":[],"launchTargets":[],"customConfigurationProvider":{"workspaceBrowse":{"browsePath":[]},"fileIndex":[]}} diff --git a/.vscode/dryrun.log b/.vscode/dryrun.log index cdd55217..5abf6dee 100644 --- a/.vscode/dryrun.log +++ b/.vscode/dryrun.log @@ -1,6 +1,5 @@ make --dry-run --always-make --keep-going --print-directory make: Entering directory '/home/steven/Documents/Projects/radio/EOR/HERA/hera_pspec' make: Leaving directory '/home/steven/Documents/Projects/radio/EOR/HERA/hera_pspec' - -make: *** No targets specified and no makefile found. Stop. +make: *** No targets specified and no makefile found. Stop. diff --git a/.vscode/settings.json b/.vscode/settings.json index 65e1ec07..5a635a5e 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,3 +1,4 @@ { - "makefile.extensionOutputFolder": "./.vscode" -} \ No newline at end of file + "makefile.extensionOutputFolder": "./.vscode", + "esbonio.sphinx.confDir": "" +} diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c81146c..550e388a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,12 +1,12 @@ v0.2.0 (2019-08-09) * Python 3 compatibility. - * Allow cross-polarization spectra to be calculated as long as they + * Allow cross-polarization spectra to be calculated as long as they aren't beam-normalized; polarizations now specified as polpairs. * Time-dependent noise power spectra capability in generate_noise. * Methods to fetch redundant baselines/blpairs from UVPSpec objects. - * Exact normalization mode (Saurabh Singh) with optimization (Ronan + * Exact normalization mode (Saurabh Singh) with optimization (Ronan Legin). - * Updated covariance handling, incl. averaging and analytic variance + * Updated covariance handling, incl. averaging and analytic variance (Jianrong Tan). * New 'lazy covariance' weighting mode (Aaron Ewall-Wice). * Fix bug where little_h was passed incorrectly (Zachary Martinot). @@ -15,6 +15,6 @@ v0.2.0 (2019-08-09) * Various minor interface changes and improved tests. v0.1.0 (2018-07-18) - * Initial release; implements core functionality, + * Initial release; implements core functionality, documentation, and tests. * Includes basic pipeline run scripts. diff --git a/README.md b/README.md index 1c328873..dde8bd48 100644 --- a/README.md +++ b/README.md @@ -4,23 +4,23 @@ [![codecov](https://codecov.io/gh/HERA-Team/hera_pspec/branch/master/graph/badge.svg)](https://codecov.io/gh/HERA-Team/hera_pspec) [![Documentation](https://readthedocs.org/projects/hera-pspec/badge/?version=latest)](https://readthedocs.org/projects/hera-pspec/badge/?version=latest) -The ``hera_pspec`` library provides all of the tools and data structures needed to perform a delay -spectrum analysis on interferometric data. The input data can be in any format supported by ``pyuvdata``, +The ``hera_pspec`` library provides all of the tools and data structures needed to perform a delay +spectrum analysis on interferometric data. The input data can be in any format supported by ``pyuvdata``, and the output data are stored in HDF5 containers. For usage examples and documentation, see http://hera-pspec.readthedocs.io/en/latest/. ## Installation Preferred method of installation for users is simply `pip install .` -(or `pip install git+https://github.com/HERA-Team/hera_pspec`). This will install +(or `pip install git+https://github.com/HERA-Team/hera_pspec`). This will install required dependencies. See below for manual dependency management. - + ### Dependencies If you are using `conda`, you may wish to install the following dependencies manually to avoid them being installed automatically by `pip`:: $ conda install -c conda-forge "numpy>=1.15" "astropy>=2.0" h5py pyuvdata scipy matplotlib pyyaml - + ### Developing If you are developing `hera_pspec`, it is preferred that you do so in a fresh `conda` environment. The following commands will install all relevant development packages:: @@ -30,9 +30,9 @@ environment. The following commands will install all relevant development packag $ conda create -n hera_pspec python=3 $ conda activate hera_pspec $ conda env update -n hera_pspec -f ci/hera_pspec_tests.yml - $ pip install -e . + $ pip install -e . -This will install extra dependencies required for testing/development as well as the +This will install extra dependencies required for testing/development as well as the standard ones. ### Running Tests @@ -42,10 +42,10 @@ From the source `hera_pspec` directory run: `pytest`. ## Running `hera_pspec` -See the documentation for an -[overview and examples](http://hera-pspec.readthedocs.io/en/latest/pspec.html) -of how to run `hera_pspec`. There are also some example Jupyter notebooks, -including [`examples/PS_estimation_examples.ipynb`](examples/PS_estimation_example.ipynb) -(a brief tutorial on how to create delay spectra), and -[`examples/PSpecBeam_tutorial.ipynb`](examples/PSpecBeam_tutorial.ipynb) (a brief +See the documentation for an +[overview and examples](http://hera-pspec.readthedocs.io/en/latest/pspec.html) +of how to run `hera_pspec`. There are also some example Jupyter notebooks, +including [`examples/PS_estimation_examples.ipynb`](examples/PS_estimation_example.ipynb) +(a brief tutorial on how to create delay spectra), and +[`examples/PSpecBeam_tutorial.ipynb`](examples/PSpecBeam_tutorial.ipynb) (a brief tutorial on handling beam objects). diff --git a/ci/hera_pspec_tests.yml b/ci/hera_pspec_tests.yml index eee21afb..353fba09 100644 --- a/ci/hera_pspec_tests.yml +++ b/ci/hera_pspec_tests.yml @@ -22,4 +22,3 @@ dependencies: - git+https://github.com/RadioAstronomySoftwareGroup/pyuvdata - git+https://github.com/HERA-Team/uvtools - git+https://github.com/HERA-Team/hera_cal - diff --git a/docs/Makefile b/docs/Makefile index 12c13560..ba5f41ed 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -17,4 +17,4 @@ help: # Catch-all target: route all unknown targets to Sphinx using the new # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). %: Makefile - @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) \ No newline at end of file + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/conf.py b/docs/conf.py index dd50f68a..66062156 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -25,19 +25,39 @@ # Mock-import modules to allow build to complete without throwing errors import mock -MOCK_MODULES = ['numpy', 'scipy', 'scipy.interpolate', 'scipy.integrate', - 'pyuvdata', 'h5py', 'omnical', 'linsolve', - 'hera_qm', 'uvtools', 'uvtools.dspec', 'hera_cal', - 'hera_cal.utils', 'healpy', 'astropy', - 'astropy.cosmology', 'astropy.units', 'astropy.constants', - 'matplotlib', 'matplotlib.pyplot', 'pylab', 'yaml', - 'pyuvdata.utils', 'hera_sim'] + +MOCK_MODULES = [ + "numpy", + "scipy", + "scipy.interpolate", + "scipy.integrate", + "pyuvdata", + "h5py", + "omnical", + "linsolve", + "hera_qm", + "uvtools", + "uvtools.dspec", + "hera_cal", + "hera_cal.utils", + "healpy", + "astropy", + "astropy.cosmology", + "astropy.units", + "astropy.constants", + "matplotlib", + "matplotlib.pyplot", + "pylab", + "yaml", + "pyuvdata.utils", + "hera_sim", +] for mod_name in MOCK_MODULES: sys.modules[mod_name] = mock.Mock() -sys.path.insert(0, os.path.abspath('../')) -#import hera_pspec +sys.path.insert(0, os.path.abspath("../")) +# import hera_pspec # -- General configuration ------------------------------------------------ @@ -49,38 +69,40 @@ # 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.napoleon', - 'sphinx.ext.coverage', - 'sphinx.ext.mathjax', - 'sphinx.ext.viewcode', - 'sphinx.ext.githubpages'] +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.napoleon", + "sphinx.ext.coverage", + "sphinx.ext.mathjax", + "sphinx.ext.viewcode", + "sphinx.ext.githubpages", +] # Add any paths that contain templates here, relative to this directory. -templates_path = ['_templates'] +templates_path = ["_templates"] # The suffix(es) of source filenames. # You can specify multiple suffix as a list of string: # # source_suffix = ['.rst', '.md'] -source_suffix = '.rst' +source_suffix = ".rst" # The master toctree document. -master_doc = 'index' +master_doc = "index" # General information about the project. -project = u'hera_pspec' -copyright = u'2018, HERA Collaboration' -author = u'HERA Collaboration' +project = "hera_pspec" +copyright = "2018, HERA Collaboration" +author = "HERA Collaboration" # The version info for the project you're documenting, acts as replacement for # |version| and |release|, also used in various other places throughout the # built documents. # # The short X.Y version. -version = u'0.1' +version = "0.1" # The full version, including alpha/beta/rc tags. -release = u'0.1' +release = "0.1" # The language for content autogenerated by Sphinx. Refer to documentation # for a list of supported languages. @@ -92,10 +114,10 @@ # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This patterns also effect to html_static_path and html_extra_path -exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] # The name of the Pygments (syntax highlighting) style to use. -pygments_style = 'sphinx' +pygments_style = "sphinx" # If true, `todo` and `todoList` produce output, else they produce nothing. todo_include_todos = False @@ -106,12 +128,9 @@ # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. # -#html_theme = 'alabaster' +# html_theme = 'alabaster' html_theme = "default" -html_theme_options = { - "rightsidebar": "false", - "relbarbgcolor": "black" -} +html_theme_options = {"rightsidebar": "false", "relbarbgcolor": "black"} # Theme options are theme-specific and customize the look and feel of a theme # further. For a list of options available for each theme, see the @@ -122,18 +141,18 @@ # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". -#html_static_path = ['_static'] +# html_static_path = ['_static'] # -- Options for HTMLHelp output ------------------------------------------ # Output file base name for HTML help builder. -htmlhelp_basename = 'hera_pspecdoc' +htmlhelp_basename = "hera_pspecdoc" # -- Options for numpydoc ------------------------------------------ -#numpydoc_show_class_members = False +# numpydoc_show_class_members = False napoleon_google_docstring = False napoleon_use_param = False napoleon_use_ivar = True @@ -145,15 +164,12 @@ # The paper size ('letterpaper' or 'a4paper'). # # 'papersize': 'letterpaper', - # The font size ('10pt', '11pt' or '12pt'). # # 'pointsize': '10pt', - # Additional stuff for the LaTeX preamble. # # 'preamble': '', - # Latex figure (float) alignment # # 'figure_align': 'htbp', @@ -163,8 +179,13 @@ # (source start file, target name, title, # author, documentclass [howto, manual, or own class]). latex_documents = [ - (master_doc, 'hera_pspec.tex', u'hera\\_pspec Documentation', - u'HERA Collaboration', 'manual'), + ( + master_doc, + "hera_pspec.tex", + "hera\\_pspec Documentation", + "HERA Collaboration", + "manual", + ), ] @@ -172,10 +193,7 @@ # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). -man_pages = [ - (master_doc, 'hera_pspec', u'hera_pspec Documentation', - [author], 1) -] +man_pages = [(master_doc, "hera_pspec", "hera_pspec Documentation", [author], 1)] # -- Options for Texinfo output ------------------------------------------- @@ -184,10 +202,13 @@ # (source start file, target name, title, author, # dir menu entry, description, category) texinfo_documents = [ - (master_doc, 'hera_pspec', u'hera_pspec Documentation', - author, 'hera_pspec', 'HERA delay power spectrum estimation.', - 'Miscellaneous'), + ( + master_doc, + "hera_pspec", + "hera_pspec Documentation", + author, + "hera_pspec", + "HERA delay power spectrum estimation.", + "Miscellaneous", + ), ] - - - diff --git a/docs/index.rst b/docs/index.rst index 4246e426..fdf4e618 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -2,9 +2,9 @@ HERA delay power spectrum estimation ==================================== -The ``hera_pspec`` library provides all of the tools and data structures needed -to perform a delay spectrum analysis on interferometric data. The input data -can be in any format supported by ``pyuvdata``, and the output data are stored in +The ``hera_pspec`` library provides all of the tools and data structures needed +to perform a delay spectrum analysis on interferometric data. The input data +can be in any format supported by ``pyuvdata``, and the output data are stored in HDF5 containers. You can find the code in the ``hera_pspec`` `GitHub repository `_. A set of `example Jupyter notebooks `_ are also available on GitHub. @@ -13,7 +13,7 @@ You can find the code in the ``hera_pspec`` `GitHub repository `_. -The only plotting function currently available in the `hera_pspec.plot` module is `delay_spectrum()`. +The only plotting function currently available in the ``hera_pspec.plot`` module is ``delay_spectrum()``. .. autofunction:: hera_pspec.plot.delay_spectrum - - diff --git a/docs/pspec.rst b/docs/pspec.rst index 49229e36..af8f42ad 100644 --- a/docs/pspec.rst +++ b/docs/pspec.rst @@ -13,46 +13,46 @@ Example delay power spectrum calculation The following example shows how to load ``UVData`` objects into a ``PSpecData`` object, specify a set of baselines and datasets that should be cross-multiplied, specify a set of spectral windows, weights, and tapers, and output a set of delay power spectra into a ``UVPSpec`` object. .. code-block:: python - + # Load into UVData objects uvd1 = UVData(); uvd2 = UVData() uvd1.read_miriad(datafile1) uvd2.read_miriad(datafile2) - + # Create a new PSpecData object ds = hp.PSpecData(dsets=[uvd1, uvd2], beam=beam) - + # bls1 and bls2 are lists of tuples specifying baselines (antenna pairs) - # Here, we specify three baseline-pairs, i.e. bls1[0] x bls2[0], + # Here, we specify three baseline-pairs, i.e. bls1[0] x bls2[0], # bls1[1] x bls2[1], and bls1[2] x bls2[2]. bls1 = [(24,25), (37,38), (38,39)] bls2 = [(37,38), (38,39), (24,25)] - - # Calculate cross-spectra of visibilities for baselines bls1[i] x bls2[i], - # where bls1 are the baselines to take from dataset 0 and bls2 are taken from - # dataset 1 (and i goes from 0..2). This is done for two spectral windows - # (freq. channel indexes between 300-400 and 600-721), with unity weights + + # Calculate cross-spectra of visibilities for baselines bls1[i] x bls2[i], + # where bls1 are the baselines to take from dataset 0 and bls2 are taken from + # dataset 1 (and i goes from 0..2). This is done for two spectral windows + # (freq. channel indexes between 300-400 and 600-721), with unity weights # and a Blackman-Harris taper in each spectral window - uvp = ds.pspec(bls1, bls2, dsets=(0, 1), spw_ranges=[(300, 400), (600, 721)], - input_data_weight='identity', norm='I', taper='blackman-harris', + uvp = ds.pspec(bls1, bls2, dsets=(0, 1), spw_ranges=[(300, 400), (600, 721)], + input_data_weight='identity', norm='I', taper='blackman-harris', verbose=True) -``uvp`` is now a ``UVPSpec`` object containing 2 x 3 x Ntimes delay spectra, where -3 is the number of baseline-pairs (i.e. ``len(bls1) == len(bls2) == 3``), 2 is -the number of spectral windows, and Ntimes is the number of LST bins in the -input ``UVData`` objects. Each delay spectrum has length ``Nfreqs``, i.e. the +``uvp`` is now a ``UVPSpec`` object containing 2 x 3 x Ntimes delay spectra, where +3 is the number of baseline-pairs (i.e. ``len(bls1) == len(bls2) == 3``), 2 is +the number of spectral windows, and Ntimes is the number of LST bins in the +input ``UVData`` objects. Each delay spectrum has length ``Nfreqs``, i.e. the number of frequency channels in each spectral window. To get power spectra from the ``UVPSpec`` object that was returned by ``pspec``: .. code-block:: python - + # Key specifying desired spectral window, baseline-pair, and polarization pair spw = 1 polpair = ('xx', 'xx') blpair =((24, 25), (37, 38)) key = (spw, blpair, polpair) - + # Get delay bins and power spectrum dlys = uvp.get_dlys(spw) ps = uvp.get_data(key) @@ -78,10 +78,10 @@ Each call to :meth:`~hera_pspec.PSpecData.pspec` must specify a set of baseline- * **Baseline pairs** are specified as two lists: ``bls1`` is the list of baselines from the first dataset in the pair specified by the ``dsets`` argument, and ``bls2`` is the list from the second. The baseline pairs are formed by matching each element from the first list with the corresponding element from the second, e.g. ``blpair[i] = bls1[i] x bls2[i]``. A couple of helper functions are provided to construct appropriately paired lists to calculate all of the cross-spectra within a redundant baseline group, for example; see :func:`~hera_pspec.pspecdata.construct_blpairs` and :func:`~hera_pspec.pspecdata.validate_bls`. - * **Spectral windows** are specified as a list of tuples using the ``spw_ranges`` argument, with each tuple specifying the beginning and end frequency channel of the spectral window. For example, ``spw_ranges=[(220, 320)]`` defines a spectral window from channel 220 to 320, as indexed by the ``UVData`` objects. The larger the spectral window, the longer it will take to calculate the power spectra. Note that - + * **Spectral windows** are specified as a list of tuples using the ``spw_ranges`` argument, with each tuple specifying the beginning and end frequency channel of the spectral window. For example, ``spw_ranges=[(220, 320)]`` defines a spectral window from channel 220 to 320, as indexed by the ``UVData`` objects. The larger the spectral window, the longer it will take to calculate the power spectra. Note that + * **Polarizations** are looped over by default. At the moment, ``pspec()`` will only compute power spectra for matching polarizations from datasets 1 and 2. If the ``UVData`` objects stored inside the ``PSpecData`` object have incompatible polarizations, :meth:`~hera_pspec.PSpecData.validate_datasets` will raise an exception. - + .. note:: If the input datasets are phased slightly differently (e.g. due to offsets in LST bins), you can rephase (fringe-stop) them to help reduce decoherence by using the :meth:`~hera_pspec.PSpecData.rephase_to_dset` method. Note that the :meth:`~hera_pspec.PSpecData.validate_datasets` method automatically checks for discrepancies in how the ``UVData`` objects are phased, and will raise warnings or errors if any problems are found. @@ -104,13 +104,13 @@ The :meth:`~hera_pspec.PSpecData.pspec` method outputs power spectra as a single To access the power spectra, use the :meth:`~hera_pspec.UVPSpec.get_data` method, which takes a key of the form: ``(spw, blpair, polpair)``. For example: .. code-block:: python - + # Key specifying desired spectral window, baseline-pair, and polarization spw = 1 polpair = ('xx', 'xx') blpair =((24, 25), (37, 38)) key = (spw, blpair, polpair) - + # Get delay bins and power spectrum dlys = uvp.get_dlys(spw) ps = uvp.get_data(key) @@ -118,9 +118,9 @@ To access the power spectra, use the :meth:`~hera_pspec.UVPSpec.get_data` method A number of methods are provided for returning useful metadata: * :meth:`~hera_pspec.UVPSpec.get_integrations`: Get the average integration time (in seconds) for a given delay spectrum. - + * :meth:`~hera_pspec.UVPSpec.get_nsamples`: If the power spectra have been incoherently averaged (i.e. averaged after squaring), this is the effective number of samples in that average. - + * :meth:`~hera_pspec.UVPSpec.get_dlys`: Get the delay for each bin of the delay power spectra (in seconds). * :meth:`~hera_pspec.UVPSpec.get_blpair_seps`: Get the average baseline separation for a baseline pair, in the ENU frame, in meters. @@ -133,15 +133,15 @@ The power spectra are stored internally in ``UVPSpec.data_array``, which is a li * ``Npols`` is the number of polarizations. Available polarizations can be retrieved from the ``UVPSpec.pol_array`` attribute. This dimension can be indexed using the :meth:`~hera_pspec.UVPSpec.pol_to_indices` method. * ``Ndlys`` is the number of delays, which is equal to the number of frequency channels within the spectral window. The available frequencies/delays can be retrievd from the ``UVPSpec.freq_array`` and ``UVPSpec.dly_array`` attributes. Alternatively, use the :meth:`~hera_pspec.UVPSpec.get_dlys` method to get the delay values. - + * ``Nblpairts`` is the number of unique combinations of baseline-pairs and times (or equivalently LSTs), i.e. the total number of delay spectra that were calculated for a given polarization and spectral window. Baseline-pairs and times have been collapsed into a single dimension because each baseline-pair can have a different number of time samples. - + You can access slices of the baseline-pair/time dimension using the :meth:`~hera_pspec.UVPSpec.blpair_to_indices` and :meth:`~hera_pspec.UVPSpec.time_to_indices` methods. The baseline-pairs and times contained in the object can be retrieved from the ``UVPSpec.blpair_array`` and ``UVPSpec.time_avg_array`` (or ``UVPSpec.lst_avg_array``) attributes. .. note:: The ``UVPSpec.time_avg_array`` attribute is just the average of the times of the input datasets. To access the original times from each dataset, see the ``UVPSpec.time_1_array`` and ``UVPSpec.time_2_array`` attributes (or equivalently ``UVPSpec.lst_1_array`` and ``UVPSpec.lst_2_array``). - + Averaging and folding spectra ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -149,10 +149,10 @@ Averaging and folding spectra By default, separate delay spectra are produced for every LST bin and polarization available in the input datasets, and for every baseline-pair passed to :meth:`~hera_pspec.PSpecData.pspec`. To (incoherently) average these down into a single delay spectrum, use the :meth:`~hera_pspec.UVPSpec.average_spectra` method. For example, to average over all times and baseline-pairs (i.e. assuming that the ``UVPSpec`` contains spectra for a single redundant baseline group): .. code-block:: python - + # Build a list of all baseline-pairs in the UVPSpec object blpair_group = [sorted(np.unique(uvp.blpair_array))] - + # Average over all baseline pairs and times, and return in a new ``UVPSpec`` object uvp_avg = uvp.average_spectra(blpair_groups=blpair_group, time_avg=True, inplace=False) @@ -169,5 +169,3 @@ The most relevant methods from ``UVPSpec`` are listed below. See :any:`uvpspec` .. autoclass:: hera_pspec.UVPSpec :members: __init__, get_data, get_wgts, get_integrations, get_nsamples, get_dlys, get_kperps, get_kparas, get_blpair_seps, select, read_hdf5, write_hdf5, generate_noise_spectra, average_spectra, fold_spectra, get_blpair_groups_from_bl_groups :noindex: - - diff --git a/docs/pspecbeam.rst b/docs/pspecbeam.rst index f0da63eb..c963b8e2 100644 --- a/docs/pspecbeam.rst +++ b/docs/pspecbeam.rst @@ -6,10 +6,10 @@ The main purpose of ``PSpecBeam`` objects is to provide the :class:`~hera_pspec.PSpecData` class with a way of normalizing the power spectra that it produces, using the :meth:`~hera_pspec.pspecbeam.PSpecBeamBase.compute_pspec_scalar` method. To attach a ``PSpecBeam`` object to a ``PSpecData`` object, pass one in when you instantiate the class, e.g. .. code-block:: python - + # Create PSpecBeamUV from a beamfits file beam = hp.PSpecBeamUV('HERA_Beams.beamfits') - + # Attach beam to PSpecData object psd = hp.PSpecData(dsets=[], wgts=[], beam=beam) @@ -23,7 +23,9 @@ Another purpose of ``PSpecBeam`` objects is to convert flux densities to tempera uvd.data_array *= beam.Jy_to_mK(np.unique(uvd.freq_array))[None, None, :, None] # (The brackets [] are needed to match the shape of uvd.data_array) -Note that ``PSpecBeam`` objects have a cosmology attached to them. If you don't specify a cosmology (with the ``cosmo`` keyword argument), they will be instantiated with the default cosmology from `hera_pspec.conversions`. +Note that ``PSpecBeam`` objects have a cosmology attached to them. If you don't specify +a cosmology (with the ``cosmo`` keyword argument), they will be instantiated with the +default cosmology from ``hera_pspec.conversions``. There are several different types of ``PSpecBeam`` object: @@ -48,14 +50,14 @@ This class allows you to use any beam that is supported by the ``UVBeam`` class To create a beam that uses this format, simply create a new ``PSpecBeamUV`` instance with the name of a ``beamfits`` file that is supported by ``UVBeam``, e.g. .. code-block:: python - + beam = hp.PSpecBeamUV('HERA_Beams.beamfits') .. autoclass:: hera_pspec.pspecbeam.PSpecBeamUV :members: __init__, power_beam_int, power_beam_sq_int :inherited-members: :member-order: bysource - + ``PSpecBeamGauss``: Simple Gaussian beam model ---------------------------------------------- @@ -65,14 +67,12 @@ A Gaussian beam type is provided for simple testing purposes. You can specify a For example, to specify a Gaussian beam with a constant FWHM of 0.1 radians, defined over a frequency interval of [100, 200] MHz: .. code-block:: python - + # Each beam is defined over a frequency interval: beam_freqs = np.linspace(100e6, 200e6, 200) # in Hz - + # Create a new Gaussian beam object with full-width at half-max. of 0.1 radians beam_gauss = hp.PSpecBeamGauss(fwhm=0.1, beam_freqs=beam_freqs) .. autoclass:: hera_pspec.pspecbeam.PSpecBeamGauss :members: __init__, power_beam_int, power_beam_sq_int - - diff --git a/examples/internal_memos/error_bar_choices/bibtex.bib b/examples/internal_memos/error_bar_choices/bibtex.bib index 150f836e..8af690c2 100644 --- a/examples/internal_memos/error_bar_choices/bibtex.bib +++ b/examples/internal_memos/error_bar_choices/bibtex.bib @@ -706,7 +706,7 @@ @article{deboer2017hydrogen pages={045001}, year={2017}, publisher={IOP Publishing} -} +} @article{aghanim2018planck, title={Planck 2018 results. VI. Cosmological parameters}, @@ -812,4 +812,4 @@ @article{dillon2014overcoming pages={023002}, year={2014}, publisher={APS} -} \ No newline at end of file +} diff --git a/examples/internal_memos/error_bar_choices/hera_pspec_error_bar_choice.tex b/examples/internal_memos/error_bar_choices/hera_pspec_error_bar_choice.tex index d9b0b837..4c1664b0 100644 --- a/examples/internal_memos/error_bar_choices/hera_pspec_error_bar_choice.tex +++ b/examples/internal_memos/error_bar_choices/hera_pspec_error_bar_choice.tex @@ -5,7 +5,7 @@ \usepackage{natbib} \bibliographystyle{abbrvnat} \setcitestyle{authoryear,open={(},close={)}} -\newcommand{\JRT}[1]{\textcolor{red}{JRT: #1}} +\newcommand{\JRT}[1]{\textcolor{red}{JRT: #1}} \title{Error Bar Choices in HERA PSPEC} \author{pspec team} @@ -13,21 +13,21 @@ \maketitle{} \begin{abstract} - In this memo we summarize the math behind several ways to derive error bars on power spectra which are available in the current pipeline of HERA PSPEC. + In this memo we summarize the math behind several ways to derive error bars on power spectra which are available in the current pipeline of HERA PSPEC. \end{abstract} \section{Foreground/systematics dependent variance} \label{ap:fg_dependent_var} -We begin with a general expression of the variance on power spectra with the existence of foregrounds or systematics, which includes both the noise variance and the signal-noise coupling term. Given two delay spectra $\tilde{x}_1 = \tilde{s} + \tilde{n}_1$ and $\tilde{x}_2 = \tilde{s} + \tilde{n}_2$, and we express $\tilde{s} = a + b i$, $\tilde{n}_1 = c_1 + d_1 i$ and $\tilde{n}_2 = c_2 + d_2 i$, the power spectra formed from $\tilde{x}_1^* \tilde{x}_2$ is +We begin with a general expression of the variance on power spectra with the existence of foregrounds or systematics, which includes both the noise variance and the signal-noise coupling term. Given two delay spectra $\tilde{x}_1 = \tilde{s} + \tilde{n}_1$ and $\tilde{x}_2 = \tilde{s} + \tilde{n}_2$, and we express $\tilde{s} = a + b i$, $\tilde{n}_1 = c_1 + d_1 i$ and $\tilde{n}_2 = c_2 + d_2 i$, the power spectra formed from $\tilde{x}_1^* \tilde{x}_2$ is \begin{align*} -P_{\tilde{x}_1\tilde{x}_2} & = \tilde{s}^*\tilde{s} + \tilde{s}^*\tilde{n}_2 + \tilde{n}_1^*\tilde{s} + +P_{\tilde{x}_1\tilde{x}_2} & = \tilde{s}^*\tilde{s} + \tilde{s}^*\tilde{n}_2 + \tilde{n}_1^*\tilde{s} + \tilde{n}_1^*\tilde{n}_2 \notag \\ & = \{a^2+b^2 + a(c_1+c_2) + b(d_1+d_2) + c_1c_2 + d_1d_2\} \notag \\ & \phantom{=} + \{a(d_2-d_1)+b(c_1-c_2)+d_2 c_1-d_1 c_2\}i \,. \end{align*} -Here we consider $\langle s \rangle = s$, which means $a$ and $b$ are not random variables, but related to the true signal power spectrum by $P_{\tilde{s}\tilde{s}} = a^2 +b^2$, and $c_1$, $d_1$, $c_2$ and $d_2$ are i.i.d random normal variables. We then have +Here we consider $\langle s \rangle = s$, which means $a$ and $b$ are not random variables, but related to the true signal power spectrum by $P_{\tilde{s}\tilde{s}} = a^2 +b^2$, and $c_1$, $d_1$, $c_2$ and $d_2$ are i.i.d random normal variables. We then have \begin{align} \label{eq:var_real_ps} \text{var} \left[\text{Re} (P_{\tilde{x}_1\tilde{x}_2}) \right] &= \text{var} \left[a^2+b^2 + a(c_1+c_2) + b(d_1+d_2) + c_1c_2 + d_1d_2 \right] \notag \\ @@ -35,20 +35,20 @@ \section{Foreground/systematics dependent variance} & = \sqrt{2}P_{\tilde{s}\tilde{s}}P_\text{N} + P_\text{N}^2 \notag \\ & = \sqrt{2}\langle \text{Re} (P_{\tilde{x}_1\tilde{x}_2})\rangle P_\text{N} + P_\text{N}^2 \,. \end{align} -In the equations above we have used the relation $\text{var} ( c_1c_2 + d_1d_2) = 2\langle c_1^2\rangle^2 = P_\text{N}^2$, where $P_\text{N}$ is the analytic noise power spectrum we will refer to again later. We have also used the fact $\langle \text{Re} (P_{\tilde{x}_1\tilde{x}_2})\rangle = P_{\tilde{s}\tilde{s}}$, therefore we can choose $\sqrt{2}\text{Re} (P_{\tilde{x}_1\tilde{x}_2}) P_\text{N} + P_\text{N}^2$ as a general form of error bars with the existence of foregrounds or systematics. Also, since $\text{Re} (P_{\tilde{x}_1\tilde{x}_2})$ could be negative due to noise randomness, we explicitly exert a zero clipping for negative values of $\text{Re} (P_{\tilde{x}_1\tilde{x}_2})$. This will of course come with excess variance, which is not ideal, but may still be a good first-order approximation in the limit that we don't have good signal or residual systematic models. +In the equations above we have used the relation $\text{var} ( c_1c_2 + d_1d_2) = 2\langle c_1^2\rangle^2 = P_\text{N}^2$, where $P_\text{N}$ is the analytic noise power spectrum we will refer to again later. We have also used the fact $\langle \text{Re} (P_{\tilde{x}_1\tilde{x}_2})\rangle = P_{\tilde{s}\tilde{s}}$, therefore we can choose $\sqrt{2}\text{Re} (P_{\tilde{x}_1\tilde{x}_2}) P_\text{N} + P_\text{N}^2$ as a general form of error bars with the existence of foregrounds or systematics. Also, since $\text{Re} (P_{\tilde{x}_1\tilde{x}_2})$ could be negative due to noise randomness, we explicitly exert a zero clipping for negative values of $\text{Re} (P_{\tilde{x}_1\tilde{x}_2})$. This will of course come with excess variance, which is not ideal, but may still be a good first-order approximation in the limit that we don't have good signal or residual systematic models. -If we consider the variance on the whole complex values of power spectrum, then +If we consider the variance on the whole complex values of power spectrum, then \begin{align} \label{eq:var_ps} \text{var} \left[P_{\tilde{x}_1\tilde{x}_2} \right] &= \text{var} \left[a^2+b^2 + a(c_1+c_2) + b(d_1+d_2) + c_1c_2 + d_1d_2 \right] \notag \\ &\phantom{=} + \text{var} \left[a(d_2-d_1)+b(c_1-c_2)+d_2 c_1-d_1 c_2 \right] \notag \\ & = 4(a^2+b^2)\langle c_1^2\rangle + 4\langle c_1^2\rangle^2 \,. \end{align} -It is just the form in \citet{kolopanis2019simplified}, while they used the notation $P_\text{N} = 2\langle c_1^2\rangle$, thus $\text{var} \left[P_{\tilde{x}_1\tilde{x}_2} \right] = 2 P_{\tilde{s}\tilde{s}} P_\text{N} + P_\text{N}^2$ there. +It is just the form in \citet{kolopanis2019simplified}, while they used the notation $P_\text{N} = 2\langle c_1^2\rangle$, thus $\text{var} \left[P_{\tilde{x}_1\tilde{x}_2} \right] = 2 P_{\tilde{s}\tilde{s}} P_\text{N} + P_\text{N}^2$ there. \section{Analytic method from QE formalism} \label{ap:analytic} -The QE formalism used in HERA PSPEC power spectrum estimation \footnote{\url{http://reionization.org/wp-content/uploads/2020/04/HERA044_power_spectrum_normalization_v2.pdf}} naturally leads to an analytic expression of the output covariance between bandpowers. +The QE formalism used in HERA PSPEC power spectrum estimation \footnote{\url{http://reionization.org/wp-content/uploads/2020/04/HERA044_power_spectrum_normalization_v2.pdf}} naturally leads to an analytic expression of the output covariance between bandpowers. In HERA PSPEC, an unnormalized estimator to $\alpha$th bandpowers $\hat{q}_\alpha$ is defined as $\hat{q}_\alpha = \bm{x}_1^\dagger \bm{Q}^{12,\alpha} \bm{x}_2 = \sum_{ij} \bm{x}_{1,i}^*\bm{Q}^{12,\alpha}_{ij}\bm{x}_{2,j}$, where $\bm{x}_{1}$ and $\bm{x}_{2}$ are visibilities across frequencies. The key idea here is to propagate the input covariance on visibilities between frequencies into the output covariance on bandpowers between delays. We continue to define three sets of input covariance matrices $\bm{C}^{12}$, $\bm{U}^{12}$ and $\bm{S}^{12}$ \begin{align} @@ -58,7 +58,7 @@ \section{Analytic method from QE formalism} \end{align} and we have \begin{align} -\langle \hat{q}_\alpha \hat{q}_\beta \rangle - \langle \hat{q}_\alpha \rangle\langle \hat{q}_\beta\rangle =& +\langle \hat{q}_\alpha \hat{q}_\beta \rangle - \langle \hat{q}_\alpha \rangle\langle \hat{q}_\beta\rangle =& \sum_{ijkl}\langle \bm{x}_{1,i}^*\bm{Q}^{12,\alpha}_{ij}\bm{x}_{2,j}\bm{x}_{1,k}^*\bm{Q}^{12,\beta}_{kl}\bm{x}_{2,l}\rangle - \langle \bm{x}_{1,i}^*\bm{Q}^{12,\alpha}_{ij}\bm{x}_{2,j}\rangle\langle \bm{x}_{1,k}^*\bm{Q}^{12,\beta}_{kl}\bm{x}_{2,l}\rangle\notag\\ =&\sum_{ijkl}\bm{Q}^{12,\alpha}_{ij}\bm{Q}^{12,\beta}_{kl}(\langle \bm{x}_{1,i}^*\bm{x}_{2,j}\bm{x}_{1,k}^*\bm{x}_{2,l}\rangle - \langle \bm{x}_{1,i}^*\bm{x}_{2,j}\rangle\langle \bm{x}_{1,k}^*\bm{x}_{2,l}\rangle)\notag\\ =&\sum_{ijkl}\bm{Q}^{12,\alpha}_{ij}\bm{Q}^{12,\beta}_{kl}(\langle \bm{x}_{1,i}^*\bm{x}_{1,k}^*\rangle\langle \bm{x}_{2,j}\bm{x}_{2,l}\rangle + \langle \bm{x}_{1,i}^*\bm{x}_{2,l}\rangle \langle \bm{x}_{1,k}^*\bm{x}_{2,j}\rangle)\notag\\ @@ -88,43 +88,43 @@ \section{Analytic method from QE formalism} \end{align} where $\bm{Q}^{12,\alpha*}_{ij}= \bm{Q}^{21,\alpha}_{ji}$. -Therefore the covariance between the real part of $\hat{q}_\alpha$ and the real part of $\hat{q}_\beta$ is +Therefore the covariance between the real part of $\hat{q}_\alpha$ and the real part of $\hat{q}_\beta$ is \begin{equation} \frac{1}{4}\left\{ (\langle \hat{q}_\alpha \hat{q}_\beta \rangle - \langle \hat{q}_\alpha \rangle\langle \hat{q}_\beta \rangle) + (\langle \hat{q}_\alpha \hat{q}_\beta^* \rangle - \langle \hat{q}_\alpha \rangle\langle \hat{q}_\beta ^*\rangle) + (\langle \hat{q}_\alpha^* \hat{q}_\beta \rangle - \langle \hat{q}_\alpha^* \rangle\langle \hat{q}_\beta \rangle) + (\langle \hat{q}_\alpha ^* \hat{q}_\beta ^*\rangle - \langle \hat{q}_\alpha^*\rangle\langle \hat{q}_\beta^*\rangle) \right\}\,, \end{equation} -and the covariance between the imaginary part of $\hat{q}_\alpha$ and the imaginary part of $\hat{q}_\beta$ is +and the covariance between the imaginary part of $\hat{q}_\alpha$ and the imaginary part of $\hat{q}_\beta$ is \begin{equation} \frac{1}{4}\left\{ (\langle \hat{q}_\alpha \hat{q}_\beta \rangle - \langle \hat{q}_\alpha \rangle\langle \hat{q}_\beta \rangle) - (\langle \hat{q}_\alpha \hat{q}_\beta^* \rangle - \langle \hat{q}_\alpha \rangle\langle \hat{q}_\beta ^*\rangle) - (\langle \hat{q}_\alpha^* \hat{q}_\beta \rangle - \langle \hat{q}_\alpha^* \rangle\langle \hat{q}_\beta \rangle) + (\langle \hat{q}_\alpha ^* \hat{q}_\beta ^*\rangle - \langle \hat{q}_\alpha^*\rangle\langle \hat{q}_\beta^*\rangle) \right\}\,. \end{equation} -$\hat{q}_\alpha$ should be normalized via multiplying a proper matrix $\bm{M}$ as +$\hat{q}_\alpha$ should be normalized via multiplying a proper matrix $\bm{M}$ as \begin{equation} \label{eq:palpha} \hat{P}_\alpha = \sum_{\beta} \bm{M}_{\alpha\beta} \hat{q}_\beta \,. \end{equation} -We then update the results above for $\hat{P}_\alpha$. The covariance between the real part of $\hat{P}_\alpha$ and the real part of $\hat{P}_\beta$ is +We then update the results above for $\hat{P}_\alpha$. The covariance between the real part of $\hat{P}_\alpha$ and the real part of $\hat{P}_\beta$ is \begin{align} &\frac{1}{4} \sum_{\gamma\delta} \Big\{ \bm{M}_{\alpha\gamma} \bm{M}_{\beta\delta} (\langle \hat{q}_\gamma q_\delta \rangle - \langle \hat{q}_\gamma \rangle\langle q_\delta\rangle) + \bm{M}_{\alpha\gamma} \bm{M}_{\beta\delta}^* (\langle \hat{q}_\gamma q_\delta^*\rangle - \langle \hat{q}_\gamma \rangle\langle q_\delta^*\rangle) + \notag\\& - \bm{M}_{\alpha\gamma}^* \bm{M}_{\beta\delta} (\langle \hat{q}_\gamma^* q_\delta \rangle - \langle \hat{q}_\gamma^* \rangle\langle q_\delta \rangle) + + \bm{M}_{\alpha\gamma}^* \bm{M}_{\beta\delta} (\langle \hat{q}_\gamma^* q_\delta \rangle - \langle \hat{q}_\gamma^* \rangle\langle q_\delta \rangle) + \bm{M}_{\alpha\gamma}^* \bm{M}_{\beta\delta}^* (\langle \hat{q}_\gamma^* q_\delta^*\rangle - \langle \hat{q}_\gamma^*\rangle\langle q_\delta^*\rangle) \Big\}\,, \end{align} and the covariance in the imaginary part of $\hat{P}_\alpha$ and the imaginary part of $\hat{P}_\beta$ is \begin{align} &\frac{1}{4} \sum_{\gamma\delta} \Big\{ \bm{M}_{\alpha\gamma} \bm{M}_{\beta\delta} (\langle \hat{q}_\gamma q_\delta \rangle - \langle \hat{q}_\gamma \rangle\langle q_\delta\rangle) - \bm{M}_{\alpha\gamma} \bm{M}_{\beta\delta}^* (\langle \hat{q}_\gamma q_\delta^*\rangle - \langle \hat{q}_\gamma \rangle\langle q_\delta^*\rangle) - \notag\\& - \bm{M}_{\alpha\gamma}^* \bm{M}_{\beta\delta} (\langle \hat{q}_\gamma^* q_\delta \rangle - \langle \hat{q}_\gamma^* \rangle\langle q_\delta \rangle) + + \bm{M}_{\alpha\gamma}^* \bm{M}_{\beta\delta} (\langle \hat{q}_\gamma^* q_\delta \rangle - \langle \hat{q}_\gamma^* \rangle\langle q_\delta \rangle) + \bm{M}_{\alpha\gamma}^* \bm{M}_{\beta\delta}^* (\langle \hat{q}_\gamma^* q_\delta^*\rangle - \langle \hat{q}_\gamma^*\rangle\langle q_\delta^*\rangle) \Big\}\,. \end{align} -Remarkably, the variance of the real part of $\hat{P}_\alpha$ is +Remarkably, the variance of the real part of $\hat{P}_\alpha$ is \begin{align} \label{eq:var_in_ps_real} - &\frac{1}{4} \sum_{\beta\gamma} + &\frac{1}{4} \sum_{\beta\gamma} \Big\{\bm{M}_{\alpha\beta} \bm{M}_{\alpha\gamma} \big[ \text{tr}(\bm{Q}^{12,\beta} \bm{U}^{22} \bm{Q}^{21,\gamma*} \bm{S}^{11}) + \text{tr}(\bm{Q}^{12,\beta} \bm{C}^{21}\notag \\ - & \phantom{=} \bm{Q}^{12,\gamma} \bm{C}^{21}) \big] + & \phantom{=} \bm{Q}^{12,\gamma} \bm{C}^{21}) \big] \, + 2\times \bm{M}_{\alpha\beta} \bm{M}_{\alpha\gamma}^* \big[ \text{tr}(\bm{Q}^{12,\beta} \bm{U}^{21} \bm{Q}^{12,\gamma *} \bm{S}^{21}) \, + \notag\\ - & \phantom{=} \text{tr}(\bm{Q}^{12,\beta} \bm{C}^{22} \bm{Q}^{21,\gamma} \bm{C}^{11}) \big] + & \phantom{=} \text{tr}(\bm{Q}^{12,\beta} \bm{C}^{22} \bm{Q}^{21,\gamma} \bm{C}^{11}) \big] +\bm{M}_{\alpha\beta}^* \bm{M}_{\alpha\gamma}^* \big[ \text{tr}(\bm{Q}^{21,\beta} \bm{U}^{11} \bm{Q}^{12,\gamma*} \notag\\ & \phantom{=} \bm{S}^{22})+ \text{tr}( \bm{Q}^{21,\beta} \bm{C}^{12} \bm{Q}^{21,\gamma} \bm{C}^{12}) \big] \Big\}\,, \end{align} @@ -132,7 +132,7 @@ \section{Analytic method from QE formalism} \begin{align} \label{eq:var_in_ps_imag} &\frac{-1}{4} \sum_{\beta\gamma} \Big\{ \bm{M}_{\alpha\beta} \bm{M}_{\alpha\gamma} \big[ \text{tr}(\bm{Q}^{12,\beta} \bm{U}^{22} \bm{Q}^{21,\gamma*} \bm{S}^{11}) + \text{tr}(\bm{Q}^{12,\beta} \bm{C}^{21}\notag \\ - & \phantom{=} \bm{Q}^{12,\gamma} \bm{C}^{21}) \big] + & \phantom{=} \bm{Q}^{12,\gamma} \bm{C}^{21}) \big] \, - 2 \times \bm{M}_{\alpha\beta} \bm{M}_{\alpha\gamma}^* \big[ \text{tr}(\bm{Q}^{12,\beta} \bm{U}^{21} \bm{Q}^{12,\gamma *} \bm{S}^{21}) \, + \notag\\ & \phantom{=} \text{tr}(\bm{Q}^{12,\beta} \bm{C}^{22} \bm{Q}^{21,\gamma} \bm{C}^{11}) \big] +\bm{M}_{\alpha\beta}^* \bm{M}_{\alpha\gamma}^* \big[ \text{tr}(\bm{Q}^{21,\beta} \bm{U}^{11} \bm{Q}^{12,\gamma*} \notag \\ @@ -146,29 +146,29 @@ \section{Analytic method from QE formalism} & - \langle V_\text{n}(\{a,b\},\nu_i,t) \rangle \langle V_\text{n}^*(\{a,b\},\nu_i,t) \rangle \notag \\ \approx & \phantom{-} \left|\frac{V(\{a,a\}, \nu_i,t)V(\{b,b\}, \nu_i,t)}{N_\text{nights} B \Delta t}\right| \,, \end{align} -where $B\Delta t$ is the product of the channel bandwidth and the integration time. Non-zero parts in Equation \ref{eq:var_in_ps_real} or \ref{eq:var_in_ps_imag} give us the noise variance on either real or imaginary parts of power spectra as +where $B\Delta t$ is the product of the channel bandwidth and the integration time. Non-zero parts in Equation \ref{eq:var_in_ps_real} or \ref{eq:var_in_ps_imag} give us the noise variance on either real or imaginary parts of power spectra as \begin{equation} \label{eq:analytic_noise_variance} -\frac{1}{2} \sum_{\beta\gamma} - \Big\{\bm{M}_{\alpha\beta} \bm{M}_{\alpha\gamma}^* \big[\text{tr}(\bm{Q}^{12,\beta} \bm{C}_\text{n}^{22} \bm{Q}^{21,\gamma} \bm{C}_\text{n}^{11}) \Big\}\,. +\frac{1}{2} \sum_{\beta\gamma} + \Big\{\bm{M}_{\alpha\beta} \bm{M}_{\alpha\gamma}^* \big[\text{tr}(\bm{Q}^{12,\beta} \bm{C}_\text{n}^{22} \bm{Q}^{21,\gamma} \bm{C}_\text{n}^{11}) \Big\}\,. \end{equation} -In the following we will show it is an equivalent form of $P_\text{N}^2$ where $P_\text{N}$ is what we call `Analytic Noise Power Spectrum' estimated in another parallel way given a system temperature input. +In the following we will show it is an equivalent form of $P_\text{N}^2$ where $P_\text{N}$ is what we call `Analytic Noise Power Spectrum' estimated in another parallel way given a system temperature input. -If we also consider a `Foreground/systematics dependent variance', Equation \ref{eq:var_in_ps_real} reduces to +If we also consider a `Foreground/systematics dependent variance', Equation \ref{eq:var_in_ps_real} reduces to \begin{align} \label{eq:reduced_var_in_ps_real} - & \frac{1}{2} \sum_{\beta\gamma} - \Big\{\bm{M}_{\alpha\beta} \bm{M}_{\alpha\gamma}^* \big[\text{tr}(\bm{Q}^{12,\beta} \bm{C}_\text{n}^{22} \bm{Q}^{21,\gamma} \bm{C}_\text{n}^{11}) + & \frac{1}{2} \sum_{\beta\gamma} + \Big\{\bm{M}_{\alpha\beta} \bm{M}_{\alpha\gamma}^* \big[\text{tr}(\bm{Q}^{12,\beta} \bm{C}_\text{n}^{22} \bm{Q}^{21,\gamma} \bm{C}_\text{n}^{11}) + \notag \\ + \phantom{=} & + \text{tr}(\bm{Q}^{12,\beta} \bm{C}_\text{signal}^{22} \bm{Q}^{21,\gamma} \bm{C}_\text{n}^{11}) \notag \\ - \phantom{=} & + \text{tr}(\bm{Q}^{12,\beta} \bm{C}_\text{signal}^{22} \bm{Q}^{21,\gamma} \bm{C}_\text{n}^{11}) - \notag \\ \phantom{=} & + \text{tr}(\bm{Q}^{12,\beta} \bm{C}_\text{n}^{22} \bm{Q}^{21,\gamma} \bm{C}_\text{signal}^{11}) \big] \Big\}\,, \end{align} where \begin{align} \bm{C}_{\text{signal}, ij}^{11} = \bm{C}_{\text{signal}, ij}^{22} = \frac{1}{2}\left[\bm{x}_{1,i} \bm{x}_{2,j}^* + \bm{x}_{2,i} \bm{x}_{1,j}^*\right]\,. \end{align} -Equation \ref{eq:reduced_var_in_ps_real} is an equivalent form to $\sqrt{2}\text{Re} (P_{\tilde{x}_1\tilde{x}_2}) P_\text{N} + P_\text{N}^2$. We also apply a similar zero clipping on $\bm{C}_{\text{signal}, ij}^{11}$, where rows and columns containing negative diagonal elements are set to be zero. +Equation \ref{eq:reduced_var_in_ps_real} is an equivalent form to $\sqrt{2}\text{Re} (P_{\tilde{x}_1\tilde{x}_2}) P_\text{N} + P_\text{N}^2$. We also apply a similar zero clipping on $\bm{C}_{\text{signal}, ij}^{11}$, where rows and columns containing negative diagonal elements are set to be zero. \subsection{Direct Noise Estimation By Differencing Visibility} \label{subsubsec:diff} @@ -183,15 +183,15 @@ \subsection{Direct Noise Estimation By Differencing Visibility} \bm{P}_\text{diff} & =\frac{(\tilde{n}_{1,t2} - \tilde{n}_{1,t_1})^*}{\sqrt{2}}\frac{(\tilde{n}_{2,t2} - \tilde{n}_{2,t_1})}{\sqrt{2}} \notag \\ & = \{\frac{(c_{1,t2}-c_{1,t1})}{\sqrt{2}}\frac{(c_{2,t2}-c_{2,t1})}{\sqrt{2}} + \frac{(d_{1,t2}-d_{1,t1})}{\sqrt{2}}\frac{(d_{2,t2}-d_{2,t1})}{\sqrt{2}}\} \notag \\ & \phantom{==} + \{\frac{(c_{1,t2}-c_{1,t1})}{\sqrt{2}}\frac{(d_{2,t2}-d_{2,t1})}{\sqrt{2}}-\frac{(c_{2,t2}-c_{2,t1})}{\sqrt{2}}\frac{(d_{1,t2}-d_{1,t1})}{\sqrt{2}}\}i \,, -\end{align} -where we see $\langle\{\text{Re}(\bm{P}_\text{diff})\}^2 \rangle \equiv \langle \{\frac{(c_{1,t2}-c_{1,t1})}{\sqrt{2}}\frac{(c_{2,t2}-c_{2,t1})}{\sqrt{2}} + \frac{(d_{1,t2}-d_{1,t1})}{\sqrt{2}}\frac{(d_{2,t2}-d_{2,t1})}{\sqrt{2}}\}^2 \rangle = \langle c_1^2\rangle\langle c_2^2\rangle + \langle d_1^2\rangle\langle d_2^2\rangle = P_\text{N}^2$. Therefore we could use $|\text{Re}(\bm{P}_\text{diff})|$ as a realization of error bars of $\bm{P}_{\tilde{x}_1\tilde{x}_2}$ in the noise-dominated region. +\end{align} +where we see $\langle\{\text{Re}(\bm{P}_\text{diff})\}^2 \rangle \equiv \langle \{\frac{(c_{1,t2}-c_{1,t1})}{\sqrt{2}}\frac{(c_{2,t2}-c_{2,t1})}{\sqrt{2}} + \frac{(d_{1,t2}-d_{1,t1})}{\sqrt{2}}\frac{(d_{2,t2}-d_{2,t1})}{\sqrt{2}}\}^2 \rangle = \langle c_1^2\rangle\langle c_2^2\rangle + \langle d_1^2\rangle\langle d_2^2\rangle = P_\text{N}^2$. Therefore we could use $|\text{Re}(\bm{P}_\text{diff})|$ as a realization of error bars of $\bm{P}_{\tilde{x}_1\tilde{x}_2}$ in the noise-dominated region. Intuitively, $\bm{P}_\text{diff}$ can be computed from time-differenced or frequency differenced visibility. However, by differencing the neighbouring points in frequency, we actually apply a high-pass filter in the delay space which means we suppress the power at low delay modes. To illustrate it, we replace the original data vector $\bm{x}_i$ with the difference data vector $\bm{x}'_i \equiv V'(\bm{b},\nu_i) = \left[V(\bm{b},\nu_{i+1})-V(\bm{b},\nu_i)\right] / \sqrt{2} \equiv \left(\bm{x}_{i+1} - \bm{x}_i\right)/\sqrt{2}\, (i=1,\cdots,N-1)$ and $\bm{x}'_N = \bm{x}_N$, and the new estimation of the same bandpower is %%\acl{We will probably want to use a different notation than tilde to avoid confusion with the delay transform} \begin{eqnarray} \hat{q}'_\alpha & \equiv & \sum_{ij}\frac{1}{2}e^{i2\pi\eta_\alpha(\nu_i - \nu_j)}\bm{R}_{1,i}\bm{R}_{2,j} \bm{x}'^*_{1,i} \bm{x}'_{2,j} \notag \\ & = &\sum_{i=1,\cdots, N-1;j}\frac{1}{2}e^{i2\pi\eta_\alpha(\nu_i - \nu_j)} \bm{R}_{1,i}\bm{R}_{2,j} \frac{(\bm{x}_{1,i+1}-\bm{x}_{1,i})^*}{\sqrt{2}} \bm{x}'_{2,j} \notag \\ - & \phantom{=} & + \sum_{j} \frac{1}{2} e^{i2\pi\eta_\alpha(\nu_N - \nu_j)} \bm{R}_{1,N}\bm{R}_{2,j} \bm{x}_{1,N}^* \bm{x}'_{2,j} \notag \\ - & = & \sum_{i=1,\cdots, N-1;j}\frac{1}{2} e^{i2\pi\eta_\alpha(\nu_i - \nu_j)} \bm{R}_{1,i}\bm{R}_{2,j} \frac{\bm{x}^*_{1,i+1}}{\sqrt{2}} \bm{x}'_{2,j}\notag \\ + & \phantom{=} & + \sum_{j} \frac{1}{2} e^{i2\pi\eta_\alpha(\nu_N - \nu_j)} \bm{R}_{1,N}\bm{R}_{2,j} \bm{x}_{1,N}^* \bm{x}'_{2,j} \notag \\ + & = & \sum_{i=1,\cdots, N-1;j}\frac{1}{2} e^{i2\pi\eta_\alpha(\nu_i - \nu_j)} \bm{R}_{1,i}\bm{R}_{2,j} \frac{\bm{x}^*_{1,i+1}}{\sqrt{2}} \bm{x}'_{2,j}\notag \\ & \phantom{=} & - \sum_{i=1,\cdots, N-1;j}\frac{1}{2} e^{i2\pi\eta_\alpha(\nu_i - \nu_j)} \bm{R}_{1,i}\bm{R}_{2,j} \frac{\bm{x}^*_{1,i}}{\sqrt{2}} \bm{x}'_{2,j} \notag \\ & \phantom{=} & + \sum_{j} \frac{1}{2} e^{i2\pi\eta_\alpha(\nu_N - \nu_j)} \bm{R}_{1,N}\bm{R}_{2,j} \bm{x}_{1,N}^* \bm{x}'_{2,j} \notag \\ & \approx & (e^{-i2\pi\eta_\alpha \Delta\nu} -1 )\sum_{ij} \frac{1}{2} e^{i2\pi\eta_\alpha(\nu_i-\nu_j)} \bm{R}_{1,i}\bm{R}_{2,j} \frac{\bm{x}^*_{1,i}}{\sqrt{2}} \bm{x}'_{2,j} \notag \\ @@ -210,7 +210,7 @@ \subsection{Noise Power Spectrum} \end{equation} where $X^2Y$ are conversion factors from signal angles and frequencies to cosmological coordinates, $\Omega_\text{eff}$ is the effective beam area, $t_\text{int}$ is the integration time, $N_\text{coherent}$ is the number of samples averaged at the level of visibility while $N_\text{incoherent}$ is the numbers of samples averaged at the level of power spectrum. -Generally, $T_\text{sys} = T_\text{signal} + T_\text{rcvr}$. It can be estimated via the RMS of the differenced visibilities over samples, where we take the differences of raw visibilities in adjacent time and frequency channels to obtain the differenced visibilities first. By the relation +Generally, $T_\text{sys} = T_\text{signal} + T_\text{rcvr}$. It can be estimated via the RMS of the differenced visibilities over samples, where we take the differences of raw visibilities in adjacent time and frequency channels to obtain the differenced visibilities first. By the relation \begin{equation} \label{eq:RMS_Tsys} V_\text{RMS} = \frac{2k_b \nu^2 \Omega_p}{c^2}\frac{T_\text{sys}}{B \Delta t} \,, @@ -225,7 +225,7 @@ \subsection{Noise Power Spectrum} \label{eq:auto2RMS} V^2_\text{RMS,\{a,b\}} = \frac{V(\{a,a\}) V(\{b,b\})}{B \Delta t}\,, \end{equation} -which is equivalent to Equation \ref{eq:auto_vis_noise} for the input noise covariance matrix. Thus the noise power spectrum estimated in this way essentially reduces to a special case of the analytic method we introduced earlier. While the analytic method is more preferred since in Equation \ref{eq:analytic_P_N} we actually use a spectral-window-averaged system temperature which might lose some information on frequency spectra during the averaging process. +which is equivalent to Equation \ref{eq:auto_vis_noise} for the input noise covariance matrix. Thus the noise power spectrum estimated in this way essentially reduces to a special case of the analytic method we introduced earlier. While the analytic method is more preferred since in Equation \ref{eq:analytic_P_N} we actually use a spectral-window-averaged system temperature which might lose some information on frequency spectra during the averaging process. -\bibliography{bibtex} +\bibliography{bibtex} \end{document} diff --git a/hera_pspec/__init__.py b/hera_pspec/__init__.py index 44f206e5..a2e8e98b 100644 --- a/hera_pspec/__init__.py +++ b/hera_pspec/__init__.py @@ -26,4 +26,4 @@ __version__ = "unknown" del version -del PackageNotFoundError \ No newline at end of file +del PackageNotFoundError diff --git a/hera_pspec/container.py b/hera_pspec/container.py index f015150d..51037d3a 100644 --- a/hera_pspec/container.py +++ b/hera_pspec/container.py @@ -9,8 +9,8 @@ def transactional(fn): """ - Handle 'transactional' operations on PSpecContainer, where the HDF5 file is - opened and then closed again for every operation. This is done when + Handle 'transactional' operations on PSpecContainer, where the HDF5 file is + opened and then closed again for every operation. This is done when keep_open = False. For calling a @transactional function within another @transactional function, @@ -19,16 +19,17 @@ def transactional(fn): @transactional function exits, in which case it will close the container if keep_open = False. """ + @wraps(fn) def wrapper(*args, **kwargs): - psc = args[0] # self object - + psc = args[0] # self object + # Open HDF5 file if needed if psc.data is None: psc._open() # if passed 'nested' kwarg get it - nested = kwargs.pop('nested', False) + nested = kwargs.pop("nested", False) # Run function try: @@ -42,10 +43,10 @@ def wrapper(*args, **kwargs): # Close HDF5 file if necessary if not psc.keep_open and not nested: psc._close() - + # Return function result return f - + return wrapper @@ -53,8 +54,10 @@ class PSpecContainer(object): """ Container class for managing multiple UVPSpec objects. """ - def __init__(self, filename, mode='r', keep_open=True, swmr=False, - tsleep=0.1, maxiter=1): + + def __init__( + self, filename, mode="r", keep_open=True, swmr=False, tsleep=0.1, maxiter=1 + ): """ Manage a collection of UVPSpec objects that are stored in a structured HDF5 file. @@ -72,16 +75,16 @@ def __init__(self, filename, mode='r', keep_open=True, swmr=False, Whether to load the HDF5 file as read/write ('rw') or read-only ('r'). If 'rw' is specified and the file doesn't exist, an empty one will be created. - + keep_open : bool, optional - Whether the HDF5 file should be kept open, or opened and then - closed again each time an operation is performed. Setting + Whether the HDF5 file should be kept open, or opened and then + closed again each time an operation is performed. Setting `keep_open=False` is helpful for multi-process access patterns. Default: True (keep file open). swmr : bool, optional Enable Single-Writer Multiple-Reader (SWMR) feature of HDF5. - Note that SWMR can only be used on POSIX-compliant + Note that SWMR can only be used on POSIX-compliant filesystems, and so may not work on some network filesystems. Default: False (do not use SWMR) @@ -89,7 +92,7 @@ def __init__(self, filename, mode='r', keep_open=True, swmr=False, Time to wait in seconds after each attempt at opening the file. maxiter : int, optional - Maximum number of attempts to open file (useful for concurrent + Maximum number of attempts to open file (useful for concurrent access when file may be locked temporarily by other processes). """ self.filename = filename @@ -98,7 +101,7 @@ def __init__(self, filename, mode='r', keep_open=True, swmr=False, self.tsleep = tsleep self.maxiter = maxiter self.swmr = swmr - if mode not in ['r', 'rw']: + if mode not in ["r", "rw"]: raise ValueError("Must set mode to either 'r' or 'rw'.") # Open file ready for reading and/or writing (if not in transactional mode) @@ -108,76 +111,78 @@ def __init__(self, filename, mode='r', keep_open=True, swmr=False, def _open(self): """ - Open HDF5 file ready for reading/writing. Does nothing if the file is + Open HDF5 file ready for reading/writing. Does nothing if the file is already open. - - This method uses HDF5's single writer, multiple reader (swmr) mode, - which allows multiple handles to exist for the same file at the same - time, as long as only one is in rw mode. The rw instance should be the - *first* one that is created; if a read-only instance is already open - when a rw instance is created, an error will be raised by h5py. + + This method uses HDF5's single writer, multiple reader (swmr) mode, + which allows multiple handles to exist for the same file at the same + time, as long as only one is in rw mode. The rw instance should be the + *first* one that is created; if a read-only instance is already open + when a rw instance is created, an error will be raised by h5py. """ - if self.data is not None: return + if self.data is not None: + return # Convert user-specified mode to a mode that HDF5 recognizes. We only # allow non-destructive operations! - if self.mode == 'rw': - mode = 'a' + if self.mode == "rw": + mode = "a" else: - mode = 'r' - if self.swmr and self.mode == 'r': + mode = "r" + if self.swmr and self.mode == "r": swmr = True else: swmr = False # check HDF5 version if swmr if swmr: - hdf5_v = float('.'.join(h5py.version.hdf5_version.split('.')[:2])) + hdf5_v = float(".".join(h5py.version.hdf5_version.split(".")[:2])) if hdf5_v < 1.1: - raise NotImplementedError("HDF5 version is {}: must " - "be >= 1.10 for SWMR".format(hdf5_v)) + raise NotImplementedError( + "HDF5 version is {}: must " "be >= 1.10 for SWMR".format(hdf5_v) + ) # Try to open the file Ncount = 0 while True: try: - self.data = h5py.File(self.filename, mode, libver='latest', - swmr=swmr) - if self.mode == 'rw': + self.data = h5py.File(self.filename, mode, libver="latest", swmr=swmr) + if self.mode == "rw": try: - # Enable single writer, multiple reader mode on HDF5 file. - # This allows multiple handles to exist for the same file + # Enable single writer, multiple reader mode on HDF5 file. + # This allows multiple handles to exist for the same file # at the same time, as long as only one is in rw mode if self.swmr: self.data.swmr_mode = True except ValueError: pass break - except (IOError, OSError): + except OSError: # raise Exception if exceeded maxiter if Ncount >= self.maxiter: - if self.mode == 'rw': + if self.mode == "rw": raise OSError( "Failed to open HDF5 file. Another process may " "be holding it open; use \nkeep_open=False to " "help prevent this from happening (single " "process), or use the\nlock kwarg (multiple " - "processes).") + "processes)." + ) else: raise - + # sleep and try again Ncount += 1 time.sleep(self.tsleep) # Update header info - if self.mode == 'rw': + if self.mode == "rw": # Update header self._update_header() # Denote as Container - if 'pspec_type' not in list(self.data.attrs.keys()): - self.data.attrs['pspec_type'] = self.__class__.__name__ + if "pspec_type" not in list(self.data.attrs.keys()): + self.data.attrs["pspec_type"] = self.__class__.__name__ def _close(self): """ @@ -187,7 +192,7 @@ def _close(self): return self.data.close() self.data = None - + def _store_pspec(self, pspec_group, uvp): """ Store a UVPSpec object as group of datasets within the HDF5 file. @@ -200,7 +205,7 @@ def _store_pspec(self, pspec_group, uvp): uvp : UVPSpec Object containing power spectrum and related data. """ - if self.mode == 'r': + if self.mode == "r": raise IOError("HDF5 file was opened read-only; cannot write to file.") # Get data and attributes from UVPSpec object (stored in dicts) @@ -218,7 +223,7 @@ def _load_pspec(self, pspec_group, **kwargs): pspec_group : HDF5 group Group containing datasets that contain power spectrum and supporting information, in a standard format expected by UVPSpec. - + kwargs : dict, optional UVPSpec.read_from_group partial IO keyword arguments @@ -228,11 +233,12 @@ def _load_pspec(self, pspec_group, **kwargs): Returns a UVPSpec object constructed from the input HDF5 group. """ # Check that group is tagged as containing UVPSpec (pspec_type attribute) - if 'pspec_type' in list(pspec_group.attrs.keys()): - + if "pspec_type" in list(pspec_group.attrs.keys()): + # Convert bytes -> str if needed - pspec_type = pspec_group.attrs['pspec_type'] - if isinstance(pspec_type, bytes): pspec_type = pspec_type.decode() + pspec_type = pspec_group.attrs["pspec_type"] + if isinstance(pspec_type, bytes): + pspec_type = pspec_type.decode() if pspec_type != uvpspec.UVPSpec.__name__: raise TypeError("HDF5 group is not tagged as a UVPSpec object.") else: @@ -248,22 +254,24 @@ def _update_header(self): Update the header in the HDF5 file with useful metadata, including the git version of hera_pspec. """ - if 'header' not in list(self.data.keys()): + if "header" not in list(self.data.keys()): if not self.swmr: - hdr = self.data.create_group('header') + hdr = self.data.create_group("header") else: raise ValueError("Cannot create a header group with SWMR") else: - hdr = self.data['header'] + hdr = self.data["header"] # Check if versions of hera_pspec are the same - if 'hera_pspec.git_hash' in list(hdr.attrs.keys()): - if hdr.attrs['hera_pspec.git_hash'] != version.git_hash: - print("WARNING: HDF5 file was created by a different version " - "of hera_pspec.") + if "hera_pspec.git_hash" in list(hdr.attrs.keys()): + if hdr.attrs["hera_pspec.git_hash"] != version.git_hash: + print( + "WARNING: HDF5 file was created by a different version " + "of hera_pspec." + ) elif not self.swmr: - hdr.attrs['hera_pspec.git_hash'] = version.git_hash - + hdr.attrs["hera_pspec.git_hash"] = version.git_hash + @transactional def set_pspec(self, group, psname, pspec, overwrite=False): """ @@ -284,7 +292,7 @@ def set_pspec(self, group, psname, pspec, overwrite=False): If the power spectrum already exists in the file, whether it should overwrite it or raise an error. Default: False (does not overwrite). """ - if self.mode == 'r': + if self.mode == "r": raise IOError("HDF5 file was opened read-only; cannot write to file.") if isinstance(group, (tuple, list, dict)): @@ -296,14 +304,16 @@ def set_pspec(self, group, psname, pspec, overwrite=False): # Recursively call set_pspec() on each item of the list for _psname, _pspec in zip(psname, pspec): if not isinstance(_pspec, uvpspec.UVPSpec): - raise TypeError("pspec lists must only contain UVPSpec " - "objects.") + raise TypeError( + "pspec lists must only contain UVPSpec " "objects." + ) self.set_pspec(group, _psname, _pspec, overwrite=overwrite) return else: # Raise exception if psname is a list, but pspec is not - raise ValueError("If psname is a list, pspec must be a list of " - "the same length.") + raise ValueError( + "If psname is a list, pspec must be a list of " "the same length." + ) if isinstance(pspec, list) and not isinstance(psname, list): raise ValueError("If pspec is a list, psname must also be a list.") # No lists should pass beyond this point @@ -339,14 +349,15 @@ def set_pspec(self, group, psname, pspec, overwrite=False): psgrp = grp.create_group(key2) else: raise AttributeError( - "Power spectrum %s/%s already exists and overwrite=False." \ - % (key1, key2) ) + "Power spectrum %s/%s already exists and overwrite=False." + % (key1, key2) + ) # Add power spectrum to this group self._store_pspec(psgrp, pspec) # Store info about what kind of power spectra are in the group - psgrp.attrs['pspec_type'] = pspec.__class__.__name__ + psgrp.attrs["pspec_type"] = pspec.__class__.__name__ @transactional def get_pspec(self, group, psname=None, **kwargs): @@ -361,7 +372,7 @@ def get_pspec(self, group, psname=None, **kwargs): psname : str, optional The name of the power spectrum to return. If None, extract all available power spectra. - + kwargs : dict UVPSpec.read_from_group partial IO keyword arguments @@ -390,14 +401,15 @@ def get_pspec(self, group, psname=None, **kwargs): # Otherwise, extract all available power spectra uvp = [] + def pspec_filter(n, obj): - if u'pspec_type' in list(obj.attrs.keys()): + if "pspec_type" in list(obj.attrs.keys()): uvp.append(self._load_pspec(obj, **kwargs)) # Traverse the entire set of groups/datasets looking for pspecs - grp.visititems(pspec_filter) # This adds power spectra to the uvp list + grp.visititems(pspec_filter) # This adds power spectra to the uvp list return uvp - + @transactional def spectra(self, group): """ @@ -422,8 +434,9 @@ def spectra(self, group): # Filter to look for pspec objects ps_list = [] + def pspec_filter(n, obj): - if u'pspec_type' in list(obj.attrs.keys()): + if "pspec_type" in list(obj.attrs.keys()): ps_list.append(n) # Traverse the entire set of groups/datasets looking for pspecs @@ -441,8 +454,8 @@ def groups(self): List of group names. """ groups = list(self.data.keys()) - if u'header' in groups: - groups.remove(u'header') + if "header" in groups: + groups.remove("header") return groups @transactional @@ -475,14 +488,14 @@ def tree(self, return_str=True): for pspec in tree[grp]: s += " |--%s\n" % pspec return s - + @transactional def save(self): """ Force HDF5 file to flush to disk. """ self.data.flush() - + def __del__(self): """ Make sure that HDF5 file is closed on destruct. @@ -490,14 +503,21 @@ def __del__(self): # Uses try-except construct just as a safeguard try: self.data.close() - except: + except Exception: pass -def combine_psc_spectra(psc, groups=None, dset_split_str='_x_', ext_split_str='_', - merge_history=True, verbose=True, overwrite=False): +def combine_psc_spectra( + psc, + groups=None, + dset_split_str="_x_", + ext_split_str="_", + merge_history=True, + verbose=True, + overwrite=False, +): """ - Iterate through a PSpecContainer and, within each specified group, combine + Iterate through a PSpecContainer and, within each specified group, combine UVPSpec (i.e. spectra) of similar name but varying psname extension. Power spectra to-be-merged are assumed to follow the naming convention @@ -505,15 +525,15 @@ def combine_psc_spectra(psc, groups=None, dset_split_str='_x_', ext_split_str='_ dset1_x_dset2_ext1, dset1_x_dset2_ext2, ... where _x_ is the default dset_split_str, and _ is the default ext_split_str. - The spectra names are first split by dset_split_str, and then by - ext_split_str. In this particular case, all instances of dset1_x_dset2* + The spectra names are first split by dset_split_str, and then by + ext_split_str. In this particular case, all instances of dset1_x_dset2* will be merged together. - In order to merge spectra names with no dset distinction and only an - extension, feed dset_split_str as '' or None. Example, to merge together: - uvp_1, uvp_2, uvp_3, feed dset_split_str=None and ext_split_str='_'. + In order to merge spectra names with no dset distinction and only an + extension, feed dset_split_str as '' or None. Example, to merge together: + ``uvp_1, uvp_2, uvp_3, feed dset_split_str=None and ext_split_str='_'``. - Note this is a destructive and inplace operation, all of the *_ext1 objects + Note this is a destructive and inplace operation, all of the *_ext1 objects are removed after merge. Parameters @@ -541,10 +561,10 @@ def combine_psc_spectra(psc, groups=None, dset_split_str='_x_', ext_split_str='_ """ # Load container if isinstance(psc, (str, np.str)): - psc = PSpecContainer(psc, mode='rw') + psc = PSpecContainer(psc, mode="rw") else: assert isinstance(psc, PSpecContainer) - + # Get groups _groups = psc.groups() if groups is None: @@ -561,11 +581,12 @@ def combine_psc_spectra(psc, groups=None, dset_split_str='_x_', ext_split_str='_ # Get unique spectra by splitting and then re-joining unique_spectra = [] for spc in spectra: - if dset_split_str == '' or dset_split_str is None: + if dset_split_str == "" or dset_split_str is None: sp = spc.split(ext_split_str)[0] else: - sp = utils.flatten([s.split(ext_split_str) - for s in spc.split(dset_split_str)])[:2] + sp = utils.flatten( + [s.split(ext_split_str) for s in spc.split(dset_split_str)] + )[:2] sp = dset_split_str.join(sp) if sp not in unique_spectra: unique_spectra.append(sp) @@ -573,20 +594,25 @@ def combine_psc_spectra(psc, groups=None, dset_split_str='_x_', ext_split_str='_ # Iterate over each unique spectra, and merge all spectra extensions for spc in unique_spectra: # check for overwrite - if spc in spectra and overwrite == False: + if spc in spectra and not overwrite: if verbose: - print("spectra {}/{} already exists and overwrite == False, " - "skipping...".format(grp, spc)) + print( + "spectra {}/{} already exists and overwrite == False, " + "skipping...".format(grp, spc) + ) continue # get merge list - to_merge = [spectra[i] for i in \ - np.where([spc in _sp for _sp in spectra])[0]] + to_merge = [ + spectra[i] for i in np.where([spc in _sp for _sp in spectra])[0] + ] try: # merge uvps = [psc.get_pspec(grp, uvp) for uvp in to_merge] if len(uvps) > 1: - merged_uvp = uvpspec.combine_uvpspec(uvps, merge_history=merge_history, verbose=verbose) + merged_uvp = uvpspec.combine_uvpspec( + uvps, merge_history=merge_history, verbose=verbose + ) else: merged_uvp = uvps[0] # write to file @@ -598,25 +624,41 @@ def combine_psc_spectra(psc, groups=None, dset_split_str='_x_', ext_split_str='_ except Exception as exc: # merge failed, so continue if verbose: - print("uvp merge failed for spectra {}/{}, exception: " \ - "{}".format(grp, spc, exc)) - + print( + "uvp merge failed for spectra {}/{}, exception: " + "{}".format(grp, spc, exc) + ) def get_combine_psc_spectra_argparser(): a = argparse.ArgumentParser( - description="argument parser for hera_pspec.container.combine_psc_spectra") + description="argument parser for hera_pspec.container.combine_psc_spectra" + ) # Add list of arguments - a.add_argument("filename", type=str, - help="Filename of HDF5 container (PSpecContainer) containing " - "groups / input power spectra.") - a.add_argument("--dset_split_str", default='_x_', type=str, - help='The pattern used to split dset1 from dset2 in the ' - 'psname.') - a.add_argument("--ext_split_str", default='_', type=str, - help='The pattern used to split the dset names from their ' - 'extension in the psname (if it exists).') - a.add_argument("--verbose", default=False, action='store_true', - help='Report feedback to stdout.') + a.add_argument( + "filename", + type=str, + help="Filename of HDF5 container (PSpecContainer) containing " + "groups / input power spectra.", + ) + a.add_argument( + "--dset_split_str", + default="_x_", + type=str, + help="The pattern used to split dset1 from dset2 in the " "psname.", + ) + a.add_argument( + "--ext_split_str", + default="_", + type=str, + help="The pattern used to split the dset names from their " + "extension in the psname (if it exists).", + ) + a.add_argument( + "--verbose", + default=False, + action="store_true", + help="Report feedback to stdout.", + ) return a diff --git a/hera_pspec/conversions.py b/hera_pspec/conversions.py index 6d2eeabd..8e3e2a0e 100644 --- a/hera_pspec/conversions.py +++ b/hera_pspec/conversions.py @@ -10,6 +10,7 @@ try: from astropy.cosmology import LambdaCDM + _astropy = True except: print("Could not import astropy") @@ -29,6 +30,7 @@ class units: f21 : 21cm frequency in Hz w21 : 21cm wavelength in meters """ + c = 2.99792458e8 # speed of light in m s-1 ckm = c / 1e3 # speed of light in km s-1 G = 6.67408e-11 # Newton constant m3 kg-1 s-2 @@ -37,13 +39,14 @@ class units: sb = 5.670367e-8 # stefan boltzmann constant W m-2 K-4 f21 = 1.420405751e9 # frequency of 21cm transition in Hz w21 = 0.211061140542 # 21cm wavelength in meters - H0_to_SI = 3.24078e-20 # km s-1 Mpc-1 to s-1 + H0_to_SI = 3.24078e-20 # km s-1 Mpc-1 to s-1 class cgs_units: """ fundamental constants in ** CGS ** units """ + c = 2.99792458e10 # cm s-1 kb = 1.38064852e-16 # erg K-1 @@ -62,11 +65,20 @@ class Cosmo_Conversions(object): Furlanetto 2006 (2006PhR...433..181F) """ + # astropy load attribute _astropy = _astropy - def __init__(self, Om_L=0.68440, Om_b=0.04911, Om_c=0.26442, H0=67.27, - Om_M=None, Om_k=None, **kwargs): + def __init__( + self, + Om_L=0.68440, + Om_b=0.04911, + Om_c=0.26442, + H0=67.27, + Om_M=None, + Om_k=None, + **kwargs + ): """ Default parameter values are Planck 2015 TT,TE,EE+lowP. (Table 4 of https://doi.org/10.1051/0004-6361/201525830) @@ -142,8 +154,8 @@ def f2z(self, freq, ghz=False): if ghz: freq = freq * 1e9 - return (units.f21 / freq - 1) - + return units.f21 / freq - 1 + @staticmethod def z2f(z, ghz=False): """ @@ -163,7 +175,8 @@ def z2f(z, ghz=False): Frequency in Hz (or GHz if ghz=True) """ freq = units.f21 / (z + 1) - if ghz: freq /= 1e9 + if ghz: + freq /= 1e9 return freq def E(self, z): @@ -176,7 +189,7 @@ def E(self, z): z : float Redshift. """ - return np.sqrt(self.Om_M*(1+z)**3 + self.Om_k*(1+z)**2 + self.Om_L) + return np.sqrt(self.Om_M * (1 + z) ** 3 + self.Om_k * (1 + z) ** 2 + self.Om_L) def DC(self, z, little_h=True): """ @@ -189,14 +202,14 @@ def DC(self, z, little_h=True): Redshift. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc (True) or + Whether to have cosmological length units be h^-1 Mpc (True) or Mpc (False). Default: True (h^-1 Mpc) """ - d = integrate.quad(lambda z: 1/self.E(z), 0, z)[0] + d = integrate.quad(lambda z: 1 / self.E(z), 0, z)[0] if little_h: - return d * units.ckm / 100. + return d * units.ckm / 100.0 else: - return d * units.ckm / self.H0 + return d * units.ckm / self.H0 def DM(self, z, little_h=True): """ @@ -209,22 +222,28 @@ def DM(self, z, little_h=True): Redshift. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc (True) or + Whether to have cosmological length units be h^-1 Mpc (True) or Mpc (False). Default: True (h^-1 Mpc) """ if little_h: - DH = units.ckm / 100. + DH = units.ckm / 100.0 else: DH = units.ckm / self.H0 if self.Om_k > 0: - DM = DH * np.sinh(np.sqrt(self.Om_k) \ - * self.DC(z, little_h=little_h) / DH) \ - / np.sqrt(self.Om_k) + DM = ( + DH + * np.sinh(np.sqrt(self.Om_k) * self.DC(z, little_h=little_h) / DH) + / np.sqrt(self.Om_k) + ) elif self.Om_k < 0: - DM = DH * np.sin(np.sqrt(np.abs(self.Om_k)) \ - * self.DC(z, little_h=little_h) / DH) \ - / np.sqrt(np.abs(self.Om_k)) + DM = ( + DH + * np.sin( + np.sqrt(np.abs(self.Om_k)) * self.DC(z, little_h=little_h) / DH + ) + / np.sqrt(np.abs(self.Om_k)) + ) else: DM = self.DC(z, little_h=little_h) @@ -241,7 +260,7 @@ def DA(self, z, little_h=True): Redshift. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc (True) or + Whether to have cosmological length units be h^-1 Mpc (True) or Mpc (False). Default: True (h^-1 Mpc) """ return self.DM(z, little_h=little_h) / (1 + z) @@ -257,32 +276,32 @@ def dRperp_dtheta(self, z, little_h=True): Redshift. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc (True) or + Whether to have cosmological length units be h^-1 Mpc (True) or Mpc (False). Default: True (h^-1 Mpc) """ - return self.DM(z, little_h=little_h) + return self.DM(z, little_h=little_h) def dRpara_df(self, z, ghz=False, little_h=True): """ - Conversion from frequency bandwidth to radial comoving distance at a + Conversion from frequency bandwidth to radial comoving distance at a specific redshift: [Mpc / Hz] Parameters: ----------- z : float Redshift. - + ghz : bool, optional Whether to convert output to [Mpc / GHz] (if True). Default: False. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc (True) or + Whether to have cosmological length units be h^-1 Mpc (True) or Mpc (False). Default: True (h^-1 Mpc) """ if little_h: - y = (1 + z)**2.0 / self.E(z) * units.ckm / 100. / units.f21 + y = (1 + z) ** 2.0 / self.E(z) * units.ckm / 100.0 / units.f21 else: - y = (1 + z)**2.0 / self.E(z) * units.ckm / self.H0 / units.f21 + y = (1 + z) ** 2.0 / self.E(z) * units.ckm / self.H0 / units.f21 if ghz: return y * 1e9 else: @@ -298,22 +317,23 @@ def X2Y(self, z, little_h=True): Redshift. little_h : boolean, optional - Whether to have cosmological length units be h^-1 Mpc (True) or + Whether to have cosmological length units be h^-1 Mpc (True) or Mpc (False). Default: True (h^-1 Mpc) Notes: ------ Calls Cosmo_Conversions.dRperp_dtheta() and Cosmo_Conversions.dRpara_df(). """ - return self.dRperp_dtheta(z, little_h=little_h)**2 \ - * self.dRpara_df(z, little_h=little_h) + return self.dRperp_dtheta(z, little_h=little_h) ** 2 * self.dRpara_df( + z, little_h=little_h + ) def bl_to_kperp(self, z, little_h=True): """ - Produce the conversion factor from baseline length [meters] to - k_perpendicular mode [h Mpc-1] at a specified redshift. + Produce the conversion factor from baseline length [meters] to + k_perpendicular mode [h Mpc-1] at a specified redshift. - Multiply this conversion factor by a baseline-separation length in + Multiply this conversion factor by a baseline-separation length in [meters] to get its corresponding k_perp mode in [h Mpc-1]. Parameters @@ -331,13 +351,16 @@ def bl_to_kperp(self, z, little_h=True): Conversion factor in units [h Mpc-1 / meters] """ # Parsons 2012, Pober 2014, Kohn 2018 - bl2kpara = 2*np.pi / (self.dRperp_dtheta(z, little_h=little_h) \ - * (units.c / self.z2f(z))) + bl2kpara = ( + 2 + * np.pi + / (self.dRperp_dtheta(z, little_h=little_h) * (units.c / self.z2f(z))) + ) return bl2kpara def tau_to_kpara(self, z, little_h=True): """ - Produce the conversion factor from delay [seconds] to k_parallel mode + Produce the conversion factor from delay [seconds] to k_parallel mode [h Mpc-1] at a specified redshift. Multiply this conversion factor by a delay mode in [seconds] @@ -358,14 +381,15 @@ def tau_to_kpara(self, z, little_h=True): Conversion factor in units [h Mpc-1 / seconds] """ # Parsons 2012, Pober 2014, Kohn 2018 - tau2kpara = 2*np.pi / self.dRpara_df(z, little_h=little_h, ghz=False) + tau2kpara = 2 * np.pi / self.dRpara_df(z, little_h=little_h, ghz=False) return tau2kpara def __str__(self): message = "Cosmo_Conversions object at <{}>\n".format(hex(id(self))) - message += "; ".join( ["{:s} : {:0.4f}".format(p, getattr(self, p)) - for p in self.params] ) + message += "; ".join( + ["{:s} : {:0.4f}".format(p, getattr(self, p)) for p in self.params] + ) return message def __eq__(self, other): @@ -373,4 +397,3 @@ def __eq__(self, other): Check two Cosmo_Conversion objects are equivalent """ return self.get_params() == other.get_params() - diff --git a/hera_pspec/grouping.py b/hera_pspec/grouping.py index e79d9daa..356895fb 100644 --- a/hera_pspec/grouping.py +++ b/hera_pspec/grouping.py @@ -12,8 +12,7 @@ from .uvwindow import UVWindow -def group_baselines(bls, Ngroups, keep_remainder=False, randomize=False, - seed=None): +def group_baselines(bls, Ngroups, keep_remainder=False, randomize=False, seed=None): """ Group baselines together into equal-sized sets. @@ -49,25 +48,29 @@ def group_baselines(bls, Ngroups, keep_remainder=False, randomize=False, grouped_bls : list of lists of tuples List of grouped baselines. """ - Nbls = len(bls) # Total baselines - n = Nbls // Ngroups # Baselines per group - rem = Nbls - n*Ngroups + Nbls = len(bls) # Total baselines + n = Nbls // Ngroups # Baselines per group + rem = Nbls - n * Ngroups # Sanity check on number of groups - if Nbls < Ngroups: raise ValueError("Can't have more groups than baselines.") + if Nbls < Ngroups: + raise ValueError("Can't have more groups than baselines.") # Make sure only tuples were provided (can't have groups of groups) - for bl in bls: assert isinstance(bl, tuple) + for bl in bls: + assert isinstance(bl, tuple) # Randomize baseline order if requested if randomize: - if seed is not None: random.seed(seed) + if seed is not None: + random.seed(seed) bls = copy.deepcopy(bls) random.shuffle(bls) # Assign to groups sequentially - grouped_bls = [bls[i*n:(i+1)*n] for i in range(Ngroups)] - if keep_remainder and rem > 0: grouped_bls[-1] += bls[-rem:] + grouped_bls = [bls[i * n : (i + 1) * n] for i in range(Ngroups)] + if keep_remainder and rem > 0: + grouped_bls[-1] += bls[-rem:] return grouped_bls @@ -94,16 +97,24 @@ def sample_baselines(bls, seed=None): Bootstrap-sampled set of baselines (will include multiple instances of some baselines). """ - if seed is not None: random.seed(seed) + if seed is not None: + random.seed(seed) # Sample with replacement; return as many baselines/groups as were input return [random.choice(bls) for i in range(len(bls))] -def average_spectra(uvp_in, blpair_groups=None, time_avg=False, - blpair_weights=None, error_field=None, - error_weights=None, normalize_weights=True, - inplace=True, add_to_history=''): +def average_spectra( + uvp_in, + blpair_groups=None, + time_avg=False, + blpair_weights=None, + error_field=None, + error_weights=None, + normalize_weights=True, + inplace=True, + add_to_history="", +): """ Average power spectra across the baseline-pair-time axis, weighted by each spectrum's integration time or a specified kind of error bars. @@ -199,14 +210,16 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, if blpair_groups is not None: # Enforce shape of blpair_groups - assert isinstance(blpair_groups[0], (list, np.ndarray)), \ - "blpair_groups must be fed as a list of baseline-pair lists. " \ - "See docstring." + assert isinstance(blpair_groups[0], (list, np.ndarray)), ( + "blpair_groups must be fed as a list of baseline-pair lists. " + "See docstring." + ) # Convert blpair_groups to list of blpair group integers if isinstance(blpair_groups[0][0], tuple): - new_blpair_grps = [[uvp.antnums_to_blpair(blp) for blp in blpg] - for blpg in blpair_groups] + new_blpair_grps = [ + [uvp.antnums_to_blpair(blp) for blp in blpg] for blpg in blpair_groups + ] blpair_groups = new_blpair_grps # Get all baseline pairs in uvp object (in integer form) @@ -216,7 +229,7 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, blvecs_groups.append(uvp.get_blpair_blvecs()[uvp_blpairs.index(group[0])]) # get baseline length for each group of baseline pairs # assuming only redundant baselines are paired together - blpair_lens, _ = utils.get_bl_lens_angs(blvecs_groups, bl_error_tol=1.) + blpair_lens, _ = utils.get_bl_lens_angs(blvecs_groups, bl_error_tol=1.0) else: # If not, each baseline pair is its own group @@ -225,27 +238,31 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, # get baseline length for each group of baseline pairs # assuming only redundant baselines are paired together blpair_lens = [blv for blv in uvp.get_blpair_seps()[np.sort(idx)]] - assert blpair_weights is None, "Cannot specify blpair_weights if "\ - "blpair_groups is None." + assert blpair_weights is None, ( + "Cannot specify blpair_weights if " "blpair_groups is None." + ) # Print warning if a blpair appears more than once in all of blpair_groups all_blpairs = [item for sublist in blpair_groups for item in sublist] if len(set(all_blpairs)) < len(all_blpairs): - print("Warning: some baseline-pairs are repeated between blpair "\ - "averaging groups.") + print( + "Warning: some baseline-pairs are repeated between blpair " + "averaging groups." + ) # Create baseline-pair weights list if not specified if blpair_weights is None: # Assign unity weights to baseline-pair groups that were specified - blpair_weights = [[1. for item in grp] for grp in blpair_groups] + blpair_weights = [[1.0 for item in grp] for grp in blpair_groups] else: # Check that blpair_weights has the same shape as blpair_groups for i, grp in enumerate(blpair_groups): try: len(blpair_weights[i]) == len(grp) except: - raise IndexError("blpair_weights must have the same shape as " - "blpair_groups") + raise IndexError( + "blpair_weights must have the same shape as " "blpair_groups" + ) # pre-check for error_weights if error_weights is None: @@ -253,7 +270,9 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, else: if hasattr(uvp, "stats_array"): if error_weights not in uvp.stats_array.keys(): - raise KeyError("error_field \"%s\" not found in stats_array keys." % error_weights) + raise KeyError( + 'error_field "%s" not found in stats_array keys.' % error_weights + ) use_error_weights = True # stat_l is a list of supplied error_fields, to sum over. @@ -269,13 +288,18 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, for stat in stat_l: if hasattr(uvp, "stats_array"): if stat not in uvp.stats_array.keys(): - raise KeyError("error_field \"%s\" not found in stats_array keys." % stat) + raise KeyError('error_field "%s" not found in stats_array keys.' % stat) if not uvp.exact_windows: # For baseline pairs not in blpair_groups, add them as their own group extra_blpairs = set(uvp.blpair_array) - set(all_blpairs) blpair_groups += [[blp] for blp in extra_blpairs] - blpair_weights += [[1.,] for blp in extra_blpairs] + blpair_weights += [ + [ + 1.0, + ] + for blp in extra_blpairs + ] # Create new data arrays data_array, wgts_array = odict(), odict() @@ -288,7 +312,7 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, cov_array_imag = odict() # same for window function - store_window = hasattr(uvp, 'window_function_array') + store_window = hasattr(uvp, "window_function_array") if store_window: window_function_array = odict() window_function_kperp, window_function_kpara = odict(), odict() @@ -329,12 +353,15 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, # calculated so that Sum (blpair wgts) = no. baselines. if blpair_weights is not None: blpg_wgts = np.array(blpair_weights[j]) - norm = np.sum(blpg_wgts) if normalize_weights else 1. - - if norm <= 0.: - raise ValueError("Sum of baseline-pair weights in " - "group %d is <= 0." % j) - blpg_wgts = blpg_wgts * float(blpg_wgts.size) / norm # Apply normalization + norm = np.sum(blpg_wgts) if normalize_weights else 1.0 + + if norm <= 0.0: + raise ValueError( + "Sum of baseline-pair weights in " "group %d is <= 0." % j + ) + blpg_wgts = ( + blpg_wgts * float(blpg_wgts.size) / norm + ) # Apply normalization else: blpg_wgts = np.ones(len(blpg)) @@ -361,7 +388,9 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, errws = {} for stat in stat_l: errws[stat] = uvp.get_stats(stat, (spw, blp, p)).copy() - np.square(errws[stat], out=errws[stat], where=np.isfinite(errws[stat])) + np.square( + errws[stat], out=errws[stat], where=np.isfinite(errws[stat]) + ) # shape of errs: (Ntimes, Ndlys) if use_error_weights: @@ -373,51 +402,67 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, # epsilon_avg = \sum{ (epsilon_i / (sigma_i)^4 } / ( \sum{ 1 / (sigma_i)^2 } )^2 # For reference: M. Tegmark 1997, The Astrophysical Journal Letters, 480, L87, Table 1, #3 # or J. Dillon 2014, Physical Review D, 89, 023002 , Equation 34. - stat_val = uvp.get_stats(error_weights, (spw, blp, p)).copy().real #shape (Ntimes, Ndlys) + stat_val = ( + uvp.get_stats(error_weights, (spw, blp, p)).copy().real + ) # shape (Ntimes, Ndlys) np.square(stat_val, out=stat_val, where=np.isfinite(stat_val)) - #corrects for potential nan values - stat_val = np.nan_to_num(stat_val, copy=False, nan=np.inf, posinf=np.inf) - w = np.real(1. / stat_val.clip(1e-40, np.inf)) + # corrects for potential nan values + stat_val = np.nan_to_num( + stat_val, copy=False, nan=np.inf, posinf=np.inf + ) + w = np.real(1.0 / stat_val.clip(1e-40, np.inf)) # shape of w: (Ntimes, Ndlys) else: # Otherwise all arrays are averaged in a way weighted by the integration time, # including the error_filed in stats_array and cov_array. # Since P_N ~ Tsys^2 / sqrt{N_incoherent} t_int (see N. Kern, The Astrophysical Journal 888.2 (2020): 70, Equation 7), # we choose w ~ P_N^{-2} ~ (ints * sqrt{nsmp})^2 - w = (ints * np.sqrt(nsmp))**2 + w = (ints * np.sqrt(nsmp)) ** 2 # shape of w: (Ntimes, 1) # Take time average if desired if time_avg: wsum = np.sum(w, axis=0).clip(1e-40, np.inf) - data = (np.sum(data * w, axis=0) \ - / wsum)[None] - wgts = (np.sum(wgts * w[:, :1, None], axis=0) \ - / wsum[:1, None])[None] + data = (np.sum(data * w, axis=0) / wsum)[None] + wgts = (np.sum(wgts * w[:, :1, None], axis=0) / wsum[:1, None])[ + None + ] # wgts has a shape of (Ntimes, Nfreqs, 2), while # w has a shape of (Ntimes, Ndlys) or (Ntimes, 1) # To handle with the case when Nfreqs != Ntimes, # we choose to multiply wgts with w[:,:1,None]. - ints = (np.sum(ints * w, axis=0) \ - / wsum)[None] + ints = (np.sum(ints * w, axis=0) / wsum)[None] nsmp = np.sum(nsmp, axis=0)[None] if store_window: if uvp.exact_windows: - window_function = (np.sum(window_function * w[:, :, None, None], axis=0)\ - / (wsum)[:, None, None])[None] + window_function = ( + np.sum( + window_function * w[:, :, None, None], axis=0 + ) + / (wsum)[:, None, None] + )[None] if not uvp.exact_windows: - window_function = (np.sum(window_function * w[:, :, None], axis=0) \ - / (wsum)[:, None])[None] + window_function = ( + np.sum(window_function * w[:, :, None], axis=0) + / (wsum)[:, None] + )[None] if store_cov: - cov_real = (np.sum(cov_real * w[:, :, None] * w[:, None, :], axis=0) \ - / wsum[:, None] / wsum[None, :])[None] - cov_imag = (np.sum(cov_imag * w[:, :, None] * w[:, None, :], axis=0) \ - / wsum[:, None] / wsum[None, :])[None] + cov_real = ( + np.sum(cov_real * w[:, :, None] * w[:, None, :], axis=0) + / wsum[:, None] + / wsum[None, :] + )[None] + cov_imag = ( + np.sum(cov_imag * w[:, :, None] * w[:, None, :], axis=0) + / wsum[:, None] + / wsum[None, :] + )[None] for stat in stat_l: # clip errws to eliminate nan: inf * 0 yields nans weighted_errws = errws[stat].clip(0, 1e40) * w**2 - errws[stat] = (np.sum(weighted_errws, axis=0) \ - / wsum**2)[None] + errws[stat] = (np.sum(weighted_errws, axis=0) / wsum**2)[ + None + ] # set near-zero errws to inf, as they should be errws[stat][np.isclose(errws[stat], 0)] = np.inf w = np.sum(w, axis=0)[None] @@ -437,23 +482,39 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, bpg_stats[stat].append(errws[stat].clip(0, 1e40) * w**2) if store_window: if uvp.exact_windows: - bpg_window_function.append(window_function * w[:, :, None, None]) + bpg_window_function.append( + window_function * w[:, :, None, None] + ) else: - bpg_window_function.append(window_function * w[:, :, None]) + bpg_window_function.append( + window_function * w[:, :, None] + ) if store_cov: - bpg_cov_real.append(cov_real * w[:, :, None] * w[:, None, :]) - bpg_cov_imag.append(cov_imag * w[:, :, None] * w[:, None, :]) + bpg_cov_real.append( + cov_real * w[:, :, None] * w[:, None, :] + ) + bpg_cov_imag.append( + cov_imag * w[:, :, None] * w[:, None, :] + ) w_list.append(w) # normalize sum: clip to deal with w_list_sum == 0 w_list_sum = np.sum(w_list, axis=0).clip(1e-40, np.inf) bpg_data = np.sum(bpg_data, axis=0) / w_list_sum - bpg_wgts = np.sum(bpg_wgts, axis=0) / w_list_sum[:,:1, None] + bpg_wgts = np.sum(bpg_wgts, axis=0) / w_list_sum[:, :1, None] bpg_nsmp = np.sum(bpg_nsmp, axis=0) bpg_ints = np.sum(bpg_ints, axis=0) / w_list_sum if store_cov: - bpg_cov_real = np.sum(bpg_cov_real, axis=0) / w_list_sum[:, :, None] / w_list_sum[:, None, :] - bpg_cov_imag = np.sum(bpg_cov_imag, axis=0) / w_list_sum[:, :, None] / w_list_sum[:, None, :] + bpg_cov_real = ( + np.sum(bpg_cov_real, axis=0) + / w_list_sum[:, :, None] + / w_list_sum[:, None, :] + ) + bpg_cov_imag = ( + np.sum(bpg_cov_imag, axis=0) + / w_list_sum[:, :, None] + / w_list_sum[:, None, :] + ) for stat in stat_l: stat_avg = np.sum(bpg_stats[stat], axis=0) / w_list_sum**2 # set near-zero stats to inf, as they should be @@ -462,12 +523,18 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, bpg_stats[stat] = np.sqrt(stat_avg) if store_window: if uvp.exact_windows: - bpg_window_function = np.sum(bpg_window_function, axis=0) # / w_list_sum[:, :, None, None] + bpg_window_function = np.sum( + bpg_window_function, axis=0 + ) # / w_list_sum[:, :, None, None] else: - bpg_window_function = np.sum(bpg_window_function, axis=0) / w_list_sum[:, :, None] + bpg_window_function = ( + np.sum(bpg_window_function, axis=0) / w_list_sum[:, :, None] + ) # Append to lists (polarization) - pol_data.extend(bpg_data); pol_wgts.extend(bpg_wgts) - pol_ints.extend(bpg_ints); pol_nsmp.extend(bpg_nsmp) + pol_data.extend(bpg_data) + pol_wgts.extend(bpg_wgts) + pol_ints.extend(bpg_ints) + pol_nsmp.extend(bpg_nsmp) for stat in stat_l: pol_stats[stat].extend(bpg_stats[stat]) if store_window: @@ -476,8 +543,10 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, pol_cov_real.extend(bpg_cov_real) pol_cov_imag.extend(bpg_cov_imag) # Append to lists (spectral window) - spw_data.append(pol_data); spw_wgts.append(pol_wgts) - spw_ints.append(pol_ints); spw_nsmp.append(pol_nsmp) + spw_data.append(pol_data) + spw_wgts.append(pol_wgts) + spw_ints.append(pol_ints) + spw_nsmp.append(pol_nsmp) for stat in stat_l: spw_stats[stat].append(pol_stats[stat]) if store_window: @@ -506,7 +575,7 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, cov_array_imag[spw] = np.moveaxis(np.array(spw_cov_imag), 0, -1) # Iterate over blpair groups one more time to assign metadata - time_1, time_2, time_avg_arr = [], [], [] + time_1, time_2, time_avg_arr = [], [], [] lst_1, lst_2, lst_avg_arr = [], [], [] blpair_arr, bl_arr = [], [] @@ -522,9 +591,11 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, time_1.extend([np.mean(uvp.time_1_array[blpairts])]) time_2.extend([np.mean(uvp.time_2_array[blpairts])]) time_avg_arr.extend([np.mean(uvp.time_avg_array[blpairts])]) - lst_1.extend([np.mean(np.unwrap(uvp.lst_1_array[blpairts]))%(2*np.pi)]) - lst_2.extend([np.mean(np.unwrap(uvp.lst_2_array[blpairts]))%(2*np.pi)]) - lst_avg_arr.extend([np.mean(np.unwrap(uvp.lst_avg_array[blpairts]))%(2*np.pi)]) + lst_1.extend([np.mean(np.unwrap(uvp.lst_1_array[blpairts])) % (2 * np.pi)]) + lst_2.extend([np.mean(np.unwrap(uvp.lst_2_array[blpairts])) % (2 * np.pi)]) + lst_avg_arr.extend( + [np.mean(np.unwrap(uvp.lst_avg_array[blpairts])) % (2 * np.pi)] + ) else: blpair_arr.extend(np.ones_like(blpairts, np.int) * blpg[0]) time_1.extend(uvp.time_1_array[blpairts]) @@ -536,8 +607,7 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, # Update arrays bl_arr = np.array(sorted(set(bl_arr))) - bl_vecs = np.array([uvp.bl_vecs[uvp.bl_array.tolist().index(bl)] - for bl in bl_arr]) + bl_vecs = np.array([uvp.bl_vecs[uvp.bl_array.tolist().index(bl)] for bl in bl_arr]) # Assign arrays and metadata to UVPSpec object uvp.Ntimes = len(np.unique(time_avg_arr)) @@ -573,13 +643,15 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, if store_cov: uvp.cov_array_real = cov_array_real uvp.cov_array_imag = cov_array_imag - if len(stat_l) >=1 : + if len(stat_l) >= 1: uvp.stats_array = stats_array elif hasattr(uvp, "stats_array"): delattr(uvp, "stats_array") # Add to history - uvp.history = "Spectra averaged with hera_pspec [{}]\n{}\n{}\n{}".format(version.git_hash[:15], add_to_history, '-'*40, uvp.history) + uvp.history = "Spectra averaged with hera_pspec [{}]\n{}\n{}\n{}".format( + version.git_hash[:15], add_to_history, "-" * 40, uvp.history + ) # Validity check uvp.check() @@ -588,9 +660,20 @@ def average_spectra(uvp_in, blpair_groups=None, time_avg=False, return uvp -def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=False, blpair_weights=None, - weight_by_cov=False, error_weights=None, - add_to_history='', little_h=True, A={}, run_check=True): +def spherical_average( + uvp_in, + kbins, + bin_widths, + blpair_groups=None, + time_avg=False, + blpair_weights=None, + weight_by_cov=False, + error_weights=None, + add_to_history="", + little_h=True, + A={}, + run_check=True, +): """ Perform a spherical average of a UVPSpec, mapping k_perp & k_para onto a |k| grid. Use UVPSpec.set_stats_slice to downweight regions of k_perp and k_para grid before averaging. @@ -659,7 +742,9 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa """ # input checks if weight_by_cov: - assert hasattr(uvp_in, 'cov_array_real'), "cannot weight by cov with no cov_array_real" + assert hasattr( + uvp_in, "cov_array_real" + ), "cannot weight by cov with no cov_array_real" if isinstance(bin_widths, (float, int)): bin_widths = np.ones_like(kbins) * bin_widths @@ -671,21 +756,30 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa assert np.all(kbin_left[1:] >= kbin_right[:-1] - 1e-6), "kbins must not overlap" # copy input - uvp = copy.deepcopy(uvp_in) + uvp = copy.deepcopy(uvp_in) # perform time and cylindrical averaging upfront if requested if not uvp.exact_windows and (blpair_groups is not None or time_avg): - uvp.average_spectra(blpair_groups=blpair_groups, time_avg=time_avg, - blpair_weights=blpair_weights, error_weights=error_weights, - inplace=True) + uvp.average_spectra( + blpair_groups=blpair_groups, + time_avg=time_avg, + blpair_weights=blpair_weights, + error_weights=error_weights, + inplace=True, + ) # initialize blank arrays and dicts Nk = len(kbins) dlys_array, spw_dlys_array = [], [] - data_array, wgt_array, integration_array, nsample_array = odict(), odict(), odict(), odict() - store_stats = hasattr(uvp, 'stats_array') + data_array, wgt_array, integration_array, nsample_array = ( + odict(), + odict(), + odict(), + odict(), + ) + store_stats = hasattr(uvp, "stats_array") store_cov = hasattr(uvp, "cov_array_real") - store_window = hasattr(uvp, 'window_function_array') or uvp.exact_windows + store_window = hasattr(uvp, "window_function_array") or uvp.exact_windows if store_cov: cov_array_real = odict() cov_array_imag = odict() @@ -705,25 +799,39 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa for spw in uvp.spw_array: # setup non delay-based arrays for this spw spw_range = spw_ranges[spw] - wgt_array[spw] = np.zeros((uvp.Ntimes, spw_range[2], 2, uvp.Npols), dtype=np.float64) + wgt_array[spw] = np.zeros( + (uvp.Ntimes, spw_range[2], 2, uvp.Npols), dtype=np.float64 + ) integration_array[spw] = np.zeros((uvp.Ntimes, uvp.Npols), dtype=np.float64) nsample_array[spw] = np.zeros((uvp.Ntimes, uvp.Npols), dtype=np.float64) # setup arrays with delays in them Ndlys = spw_range[3] Ndlyblps = Ndlys * uvp.Nblpairs - data_array[spw] = np.zeros((uvp.Ntimes, Ndlyblps, uvp.Npols), dtype=np.complex128) + data_array[spw] = np.zeros( + (uvp.Ntimes, Ndlyblps, uvp.Npols), dtype=np.complex128 + ) if store_cov: - cov_array_real[spw] = np.zeros((uvp.Ntimes, Ndlyblps, Ndlyblps, uvp.Npols), dtype=np.float64) - cov_array_imag[spw] = np.zeros((uvp.Ntimes, Ndlyblps, Ndlyblps, uvp.Npols), dtype=np.float64) + cov_array_real[spw] = np.zeros( + (uvp.Ntimes, Ndlyblps, Ndlyblps, uvp.Npols), dtype=np.float64 + ) + cov_array_imag[spw] = np.zeros( + (uvp.Ntimes, Ndlyblps, Ndlyblps, uvp.Npols), dtype=np.float64 + ) if store_stats: for stat in uvp.stats_array.keys(): - stats_array[stat][spw] = np.zeros((uvp.Ntimes, Ndlyblps, uvp.Npols), dtype=np.complex128) + stats_array[stat][spw] = np.zeros( + (uvp.Ntimes, Ndlyblps, uvp.Npols), dtype=np.complex128 + ) if store_window: if uvp.exact_windows: - window_function_array[spw] = np.zeros((uvp.Ntimes, Nk, Nk, uvp.Npols), dtype=np.float64) + window_function_array[spw] = np.zeros( + (uvp.Ntimes, Nk, Nk, uvp.Npols), dtype=np.float64 + ) else: - window_function_array[spw] = np.zeros((uvp.Ntimes, Ndlyblps, Ndlyblps, uvp.Npols), dtype=np.float64) + window_function_array[spw] = np.zeros( + (uvp.Ntimes, Ndlyblps, Ndlyblps, uvp.Npols), dtype=np.float64 + ) # setup the design matrix: P_cyl = A P_sph A[spw] = np.zeros((uvp.Ntimes, Ndlyblps, Nk, uvp.Npols), dtype=np.float64) @@ -732,7 +840,6 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa # we can represent the Ndlyblps x Ndlyblps block diagonal matrix as Ndlyblps x Ndlys E = np.zeros((uvp.Ntimes, Ndlyblps, Ndlys, uvp.Npols), dtype=np.float64) - # get kperps for this spw: shape (Nblpairts,) kperps = uvp.get_kperps(spw, little_h=True) @@ -754,7 +861,7 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa blpt_inds = uvp.blpair_to_indices(blp) # get k magnitude of data: (Ndlys,) - kmags = np.sqrt(kperps[blpt_inds][0]**2 + kparas**2) + kmags = np.sqrt(kperps[blpt_inds][0] ** 2 + kparas**2) # shape of nsmap: (Ntimes, Npols) nsmp = uvp.nsample_array[spw][blpt_inds] @@ -773,35 +880,51 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa data_array[spw][:, dslice] = uvp.data_array[spw][blpt_inds] if store_window and not uvp.exact_windows: - window_function_array[spw][:, dslice, dslice] = uvp.window_function_array[spw][blpt_inds] + window_function_array[spw][ + :, dslice, dslice + ] = uvp.window_function_array[spw][blpt_inds] if store_stats: for stat in stats_array: - stats_array[stat][spw][:, dslice] = uvp.stats_array[stat][spw][blpt_inds] + stats_array[stat][spw][:, dslice] = uvp.stats_array[stat][spw][ + blpt_inds + ] if store_cov: - cov_array_real[spw][:, dslice, dslice] = uvp.cov_array_real[spw][blpt_inds] + cov_array_real[spw][:, dslice, dslice] = uvp.cov_array_real[spw][ + blpt_inds + ] # fill weighting matrix E if weight_by_cov: # weight by inverse (real) covariance for p in range(uvp.Npols): # the covariance is block diagonal assuming no correlations between baseline-pairs - E[:, dslice, :, p] = np.linalg.pinv(cov_array_real[spw][:, dslice, dslice, p].real) + E[:, dslice, :, p] = np.linalg.pinv( + cov_array_real[spw][:, dslice, dslice, p].real + ) elif error_weights is not None: # fill diagonal with by 1/stats_array^2 as weight stat_weight = stats_array[error_weights][spw][:, dslice].real.copy() np.square(stat_weight, out=stat_weight, where=np.isfinite(stat_weight)) - E[:, range(dstart, dstop), range(0, Ndlys)] = 1 / stat_weight.clip(1e-40, np.inf) + E[:, range(dstart, dstop), range(0, Ndlys)] = 1 / stat_weight.clip( + 1e-40, np.inf + ) else: E[:, range(dstart, dstop), range(0, Ndlys)] = 1.0 - f = np.isclose(uvp.integration_array[spw][blpt_inds] * uvp.nsample_array[spw][blpt_inds], 0) - E[:, range(dstart, dstop), range(0, Ndlys)] *= (~f[:, None, :]) + f = np.isclose( + uvp.integration_array[spw][blpt_inds] + * uvp.nsample_array[spw][blpt_inds], + 0, + ) + E[:, range(dstart, dstop), range(0, Ndlys)] *= ~f[:, None, :] # append to non-dly arrays - Emean = np.trace(E[:, dslice, :], axis1=1, axis2=2) # use sum of E across delay as weight + Emean = np.trace( + E[:, dslice, :], axis1=1, axis2=2 + ) # use sum of E across delay as weight wgt_array[spw] += wgts * Emean[:, None, None, :] integration_array[spw] += ints * Emean nsample_array[spw] += nsmp @@ -821,7 +944,9 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa A[spw][:, i + Ndlys * b, kind, :] = 1.0 # normalize metadata sums - wgt_array[spw] /= np.max(wgt_array[spw], axis=(2, 3), keepdims=True).clip(1e-40, np.inf) + wgt_array[spw] /= np.max(wgt_array[spw], axis=(2, 3), keepdims=True).clip( + 1e-40, np.inf + ) integration_array[spw] /= np.trace(E.real, axis1=1, axis2=2).clip(1e-40, np.inf) # project onto spherically binned space @@ -866,7 +991,9 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa sq_stat = stats_array[stat][spw].copy() np.square(sq_stat, out=sq_stat, where=np.isfinite(sq_stat)) # einsum is fast enough for this, and is more succinct than matmul - avg_stat = np.sqrt(np.einsum("ptik,tip,ptik->tkp", H, sq_stat.clip(0, 1e40), H)) + avg_stat = np.sqrt( + np.einsum("ptik,tip,ptik->tkp", H, sq_stat.clip(0, 1e40), H) + ) # set zeroed stats to large number avg_stat[np.isclose(avg_stat, 0)] = 1e40 # update stats_array @@ -881,14 +1008,18 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa cov_array_imag[spw] = np.zeros_like(cov_array_real[spw]) if uvp.exact_windows: - window_function_array[spw] = spherical_wf_from_uvp(uvp, kbins, bin_widths, - blpair_groups=blpair_groups, - blpair_weights=blpair_weights, - time_avg=time_avg, - error_weights=error_weights, - spw_array=spw, - little_h=little_h, - verbose=True)[spw] + window_function_array[spw] = spherical_wf_from_uvp( + uvp, + kbins, + bin_widths, + blpair_groups=blpair_groups, + blpair_weights=blpair_weights, + time_avg=time_avg, + error_weights=error_weights, + spw_array=spw, + little_h=little_h, + verbose=True, + )[spw] # handle data arrays uvp.data_array = data_array @@ -916,7 +1047,9 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa uvp.Nblpairts = uvp.Ntimes uvp.Nblpairs = 1 bl_array = np.unique([uvp.antnums_to_bl(an) for an in uvp.blpair_to_antnums(blp)]) - uvp.bl_vecs = np.asarray([uvp.bl_vecs[np.argmin(uvp.bl_array - bl)] for bl in bl_array]) + uvp.bl_vecs = np.asarray( + [uvp.bl_vecs[np.argmin(uvp.bl_array - bl)] for bl in bl_array] + ) uvp.bl_array = bl_array uvp.Nbls = len(bl_array) uvp.label_1_array = uvp.label_1_array[:, blp_inds] @@ -935,7 +1068,9 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa uvp.lst_2_array = np.unique(uvp_in.lst_2_array) # Add to history - uvp.history = "Spherically averaged with hera_pspec [{}]\n{}\n{}\n{}".format(version.git_hash[:15], add_to_history, '-'*40, uvp.history) + uvp.history = "Spherically averaged with hera_pspec [{}]\n{}\n{}\n{}".format( + version.git_hash[:15], add_to_history, "-" * 40, uvp.history + ) # validity check if run_check: @@ -943,14 +1078,24 @@ def spherical_average(uvp_in, kbins, bin_widths, blpair_groups=None, time_avg=Fa return uvp -def spherical_wf_from_uvp(uvp_in, kbins, bin_widths, - blpair_groups=None, blpair_lens=None, blpair_weights=None, - error_weights=None, time_avg=False, spw_array=None, - little_h=True, verbose=False): - + +def spherical_wf_from_uvp( + uvp_in, + kbins, + bin_widths, + blpair_groups=None, + blpair_lens=None, + blpair_weights=None, + error_weights=None, + time_avg=False, + spw_array=None, + little_h=True, + verbose=False, +): + """ Obtains exact spherical window functions from an UVPspec object, - given a set of baseline-pair groups, their associated lengths, and + given a set of baseline-pair groups, their associated lengths, and a set of spherical k-bins. Parameters @@ -970,8 +1115,8 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths, blpair_weights : list relative weights of blpairs in blpair averaging (used for bootstrapping) - - blpair_lens : list + + blpair_lens : list lengths of blpairs in blpair_groups error_weights : str, optional @@ -1013,30 +1158,36 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths, Nk = len(kbins) # copy input - uvp = copy.deepcopy(uvp_in) + uvp = copy.deepcopy(uvp_in) if blpair_groups is None: if blpair_lens is not None: - warnings.warn('blpair_lens given but blpair_groups is None... overriding blpair_lens.') + warnings.warn( + "blpair_lens given but blpair_groups is None... overriding blpair_lens." + ) blpair_groups, blpair_lens, _ = uvp.get_red_blpairs() else: # Enforce shape of blpair_groups - assert isinstance(blpair_groups[0], (list, np.ndarray)), \ - "blpair_groups must be fed as a list of baseline-pair lists. " \ - "See docstring." + assert isinstance(blpair_groups[0], (list, np.ndarray)), ( + "blpair_groups must be fed as a list of baseline-pair lists. " + "See docstring." + ) if blpair_lens is None: # Get all baseline pairs in uvp object (in integer form) uvp_blpairs = [uvp.antnums_to_blpair(blp) for blp in uvp.get_blpairs()] blvecs_groups = [] for group in blpair_groups: - blvecs_groups.append(uvp.get_blpair_blvecs()[uvp_blpairs.index(group[0])]) + blvecs_groups.append( + uvp.get_blpair_blvecs()[uvp_blpairs.index(group[0])] + ) # get baseline length for each group of baseline pairs # assuming only redundant baselines are paired together - blpair_lens, _ = utils.get_bl_lens_angs(blvecs_groups, bl_error_tol=1.) - else: + blpair_lens, _ = utils.get_bl_lens_angs(blvecs_groups, bl_error_tol=1.0) + else: # ensure consistency between inputs - assert len(blpair_groups)==len(blpair_lens), "Baseline-pair groups" \ - " are inconsistent with baseline lengths" + assert len(blpair_groups) == len(blpair_lens), ( + "Baseline-pair groups" " are inconsistent with baseline lengths" + ) blpair_lens = np.array(blpair_lens) # check spw input and create array of spws to loop over @@ -1044,24 +1195,31 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths, # if no spw specified, use attribute spw_array = uvp.spw_array else: - spw_array = spw_array if isinstance(spw_array, (list, tuple, np.ndarray)) else [int(spw_array)] + spw_array = ( + spw_array + if isinstance(spw_array, (list, tuple, np.ndarray)) + else [int(spw_array)] + ) # check if spw given is in uvp - assert np.all([spw in uvp.spw_array for spw in spw_array]), \ - "input spw is not in UVPSpec.spw_array." + assert np.all( + [spw in uvp.spw_array for spw in spw_array] + ), "input spw is not in UVPSpec.spw_array." assert uvp.exact_windows, "Need to compute exact window functions first." if blpair_weights is None: # assign weight of one to each baseline length - blpair_weights = [[1. for item in grp] for grp in blpair_groups] + blpair_weights = [[1.0 for item in grp] for grp in blpair_groups] # perform redundant cylindrical averaging upfront # and apply weights to window functions - uvp.average_spectra(blpair_groups=blpair_groups, - blpair_weights=blpair_weights, - error_weights=error_weights, - time_avg=time_avg, - inplace=True) + uvp.average_spectra( + blpair_groups=blpair_groups, + blpair_weights=blpair_weights, + error_weights=error_weights, + time_avg=time_avg, + inplace=True, + ) # transform kgrid to little_h units if not little_h: @@ -1074,15 +1232,22 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths, # iterate over spectral windows for spw in spw_array: - avg_nu = (uvp.get_spw_ranges(spw)[0][1]+uvp.get_spw_ranges(spw)[0][0])/2 + avg_nu = (uvp.get_spw_ranges(spw)[0][1] + uvp.get_spw_ranges(spw)[0][0]) / 2 # construct array giving the k probed by each baseline-tau pair - kperps = uvp.cosmo.bl_to_kperp(uvp.cosmo.f2z(avg_nu), little_h=little_h) * blpair_lens - kparas = uvp.cosmo.tau_to_kpara(uvp.cosmo.f2z(avg_nu), little_h=little_h) * uvp.get_dlys(spw) - kmags = np.sqrt(kperps[:, None]**2+kparas**2) - - # setup arrays - window_function_array[spw] = np.zeros((uvp.Ntimes, Nk, Nk, uvp.Npols), dtype=np.float64) + kperps = ( + uvp.cosmo.bl_to_kperp(uvp.cosmo.f2z(avg_nu), little_h=little_h) + * blpair_lens + ) + kparas = uvp.cosmo.tau_to_kpara( + uvp.cosmo.f2z(avg_nu), little_h=little_h + ) * uvp.get_dlys(spw) + kmags = np.sqrt(kperps[:, None] ** 2 + kparas**2) + + # setup arrays + window_function_array[spw] = np.zeros( + (uvp.Ntimes, Nk, Nk, uvp.Npols), dtype=np.float64 + ) # iterate over polarisation spw_window_function = [] @@ -1091,11 +1256,11 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths, # grids used to compute the window functions kperp_bins = uvp.window_function_kperp[spw][:, ip] kpara_bins = uvp.window_function_kpara[spw][:, ip] - ktot = np.sqrt(kperp_bins[:, None]**2 + kpara_bins**2) + ktot = np.sqrt(kperp_bins[:, None] ** 2 + kpara_bins**2) cyl_wf = uvp.window_function_array[spw][:, :, :, :, ip] # separate baseline-time axis to iterate over times - cyl_wf = cyl_wf.reshape((uvp.Ntimes, uvp.Nblpairs, *cyl_wf.shape[1:] )) + cyl_wf = cyl_wf.reshape((uvp.Ntimes, uvp.Nblpairs, *cyl_wf.shape[1:])) # take average for each time for it in range(uvp.Ntimes): @@ -1103,15 +1268,26 @@ def spherical_wf_from_uvp(uvp_in, kbins, bin_widths, for m1 in range(Nk): mask1 = (kbin_left[m1] <= kmags) & (kmags < kbin_right[m1]) if np.any(mask1): - wf_temp = np.sum(cyl_wf[it, :, :, :, :]*mask1[:, :, None, None].astype(int), axis=(0, 1))/np.sum(mask1) - if np.sum(wf_temp) > 0.: + wf_temp = np.sum( + cyl_wf[it, :, :, :, :] + * mask1[:, :, None, None].astype(int), + axis=(0, 1), + ) / np.sum(mask1) + if np.sum(wf_temp) > 0.0: for m2 in range(Nk): - mask2 = (kbin_left[m2] <= ktot) & (ktot < kbin_right[m2]) - if np.any(mask2): #cannot compute mean if zero elements + mask2 = (kbin_left[m2] <= ktot) & ( + ktot < kbin_right[m2] + ) + if np.any( + mask2 + ): # cannot compute mean if zero elements wf_spherical[m1, m2] = np.mean(wf_temp[mask2]) # normalisation - wf_spherical[m1,:] = np.divide(wf_spherical[m1, :], np.sum(wf_spherical[m1, :]), - where = np.sum(wf_spherical[m1,:]) != 0) + wf_spherical[m1, :] = np.divide( + wf_spherical[m1, :], + np.sum(wf_spherical[m1, :]), + where=np.sum(wf_spherical[m1, :]) != 0, + ) spw_window_function.append(wf_spherical) window_function_array[spw][:, :, :, ip] = np.copy(spw_window_function) @@ -1150,120 +1326,173 @@ def fold_spectra(uvp): if Ndlys % 2 == 0: # even number of dlys - left = uvp.data_array[spw][:, 1:Ndlys//2, :][:, ::-1, :] - right = uvp.data_array[spw][:, Ndlys//2+1:, :] - uvp.data_array[spw][:, Ndlys//2+1:, :] = np.mean([left, right], axis=0) - uvp.data_array[spw][:, :Ndlys//2, :] = 0.0 + left = uvp.data_array[spw][:, 1 : Ndlys // 2, :][:, ::-1, :] + right = uvp.data_array[spw][:, Ndlys // 2 + 1 :, :] + uvp.data_array[spw][:, Ndlys // 2 + 1 :, :] = np.mean([left, right], axis=0) + uvp.data_array[spw][:, : Ndlys // 2, :] = 0.0 uvp.nsample_array[spw] *= 2.0 - if hasattr(uvp, 'window_function_array'): + if hasattr(uvp, "window_function_array"): if uvp.exact_windows: - left = uvp.window_function_array[spw][:, 1:Ndlys//2, ...][:, ::-1, ...] - right = uvp.window_function_array[spw][:, Ndlys//2+1: , ...] - uvp.window_function_array[spw][:, Ndlys//2+1:, ...] = .50*(left+right) + left = uvp.window_function_array[spw][:, 1 : Ndlys // 2, ...][ + :, ::-1, ... + ] + right = uvp.window_function_array[spw][:, Ndlys // 2 + 1 :, ...] + uvp.window_function_array[spw][:, Ndlys // 2 + 1 :, ...] = 0.50 * ( + left + right + ) else: - leftleft = uvp.window_function_array[spw][:, 1:Ndlys//2, 1:Ndlys//2, :][:, ::-1, ::-1, :] - leftright = uvp.window_function_array[spw][:, 1:Ndlys//2, Ndlys//2+1:, :][:, ::-1, :, :] - rightleft = uvp.window_function_array[spw][:, Ndlys//2+1: , 1:Ndlys//2, :][:, :, ::-1, :] - rightright = uvp.window_function_array[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] - uvp.window_function_array[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] = .25*(leftleft\ - +leftright\ - +rightleft\ - +rightright) - uvp.window_function_array[spw][:, :, :Ndlys//2, :] = 0.0 - uvp.window_function_array[spw][:, :Ndlys//2, ...] = 0.0 + leftleft = uvp.window_function_array[spw][ + :, 1 : Ndlys // 2, 1 : Ndlys // 2, : + ][:, ::-1, ::-1, :] + leftright = uvp.window_function_array[spw][ + :, 1 : Ndlys // 2, Ndlys // 2 + 1 :, : + ][:, ::-1, :, :] + rightleft = uvp.window_function_array[spw][ + :, Ndlys // 2 + 1 :, 1 : Ndlys // 2, : + ][:, :, ::-1, :] + rightright = uvp.window_function_array[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] + uvp.window_function_array[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] = 0.25 * (leftleft + leftright + rightleft + rightright) + uvp.window_function_array[spw][:, :, : Ndlys // 2, :] = 0.0 + uvp.window_function_array[spw][:, : Ndlys // 2, ...] = 0.0 # fold covariance array if it exists. - if hasattr(uvp,'cov_array_real'): - leftleft = uvp.cov_array_real[spw][:, 1:Ndlys//2, 1:Ndlys//2, :][:, ::-1, ::-1, :] - leftright = uvp.cov_array_real[spw][:, 1:Ndlys//2, Ndlys//2+1:, :][:, ::-1, :, :] - rightleft = uvp.cov_array_real[spw][:, Ndlys//2+1: , 1:Ndlys//2, :][:, :, ::-1, :] - rightright = uvp.cov_array_real[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] - uvp.cov_array_real[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] = .25*(leftleft\ - +leftright\ - +rightleft\ - +rightright) - uvp.cov_array_real[spw][:, :Ndlys//2, :, :] = 0.0 - uvp.cov_array_real[spw][:, :, :Ndlys//2, : :] = 0.0 - - leftleft = uvp.cov_array_imag[spw][:, 1:Ndlys//2, 1:Ndlys//2, :][:, ::-1, ::-1, :] - leftright = uvp.cov_array_imag[spw][:, 1:Ndlys//2, Ndlys//2+1:, :][:, ::-1, :, :] - rightleft = uvp.cov_array_imag[spw][:, Ndlys//2+1: , 1:Ndlys//2, :][:, :, ::-1, :] - rightright = uvp.cov_array_imag[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] - uvp.cov_array_imag[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] = .25*(leftleft\ - +leftright\ - +rightleft\ - +rightright) - uvp.cov_array_imag[spw][:, :Ndlys//2, :, :] = 0.0 - uvp.cov_array_imag[spw][:, :, :Ndlys//2, : :] = 0.0 + if hasattr(uvp, "cov_array_real"): + leftleft = uvp.cov_array_real[spw][ + :, 1 : Ndlys // 2, 1 : Ndlys // 2, : + ][:, ::-1, ::-1, :] + leftright = uvp.cov_array_real[spw][ + :, 1 : Ndlys // 2, Ndlys // 2 + 1 :, : + ][:, ::-1, :, :] + rightleft = uvp.cov_array_real[spw][ + :, Ndlys // 2 + 1 :, 1 : Ndlys // 2, : + ][:, :, ::-1, :] + rightright = uvp.cov_array_real[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] + uvp.cov_array_real[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] = 0.25 * (leftleft + leftright + rightleft + rightright) + uvp.cov_array_real[spw][:, : Ndlys // 2, :, :] = 0.0 + uvp.cov_array_real[spw][:, :, : Ndlys // 2, ::] = 0.0 + + leftleft = uvp.cov_array_imag[spw][ + :, 1 : Ndlys // 2, 1 : Ndlys // 2, : + ][:, ::-1, ::-1, :] + leftright = uvp.cov_array_imag[spw][ + :, 1 : Ndlys // 2, Ndlys // 2 + 1 :, : + ][:, ::-1, :, :] + rightleft = uvp.cov_array_imag[spw][ + :, Ndlys // 2 + 1 :, 1 : Ndlys // 2, : + ][:, :, ::-1, :] + rightright = uvp.cov_array_imag[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] + uvp.cov_array_imag[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] = 0.25 * (leftleft + leftright + rightleft + rightright) + uvp.cov_array_imag[spw][:, : Ndlys // 2, :, :] = 0.0 + uvp.cov_array_imag[spw][:, :, : Ndlys // 2, ::] = 0.0 # fold stats array if it exists: sum in inverse quadrature - if hasattr(uvp, 'stats_array'): + if hasattr(uvp, "stats_array"): for stat in uvp.stats_array.keys(): - left = uvp.stats_array[stat][spw][:, 1:Ndlys//2, :][:, ::-1, :] - right = uvp.stats_array[stat][spw][:, Ndlys//2+1:, :] - uvp.stats_array[stat][spw][:, Ndlys//2+1:, :] = (np.sum([1./left**2.0, 1./right**2.0], axis=0))**(-0.5) - uvp.data_array[spw][:, :Ndlys//2, :] = np.nan + left = uvp.stats_array[stat][spw][:, 1 : Ndlys // 2, :][:, ::-1, :] + right = uvp.stats_array[stat][spw][:, Ndlys // 2 + 1 :, :] + uvp.stats_array[stat][spw][:, Ndlys // 2 + 1 :, :] = ( + np.sum([1.0 / left**2.0, 1.0 / right**2.0], axis=0) + ) ** (-0.5) + uvp.data_array[spw][:, : Ndlys // 2, :] = np.nan else: # odd number of dlys - left = uvp.data_array[spw][:, :Ndlys//2, :][:, ::-1, :] - right = uvp.data_array[spw][:, Ndlys//2+1:, :] - uvp.data_array[spw][:, Ndlys//2+1:, :] = np.mean([left, right], axis=0) - uvp.data_array[spw][:, :Ndlys//2, :] = 0.0 + left = uvp.data_array[spw][:, : Ndlys // 2, :][:, ::-1, :] + right = uvp.data_array[spw][:, Ndlys // 2 + 1 :, :] + uvp.data_array[spw][:, Ndlys // 2 + 1 :, :] = np.mean([left, right], axis=0) + uvp.data_array[spw][:, : Ndlys // 2, :] = 0.0 uvp.nsample_array[spw] *= 2.0 - if hasattr(uvp, 'window_function_array'): + if hasattr(uvp, "window_function_array"): if uvp.exact_windows: - left = uvp.window_function_array[spw][:, :Ndlys//2, ...][:, ::-1, ...] - right = uvp.window_function_array[spw][:, Ndlys//2+1: , ...] - uvp.window_function_array[spw][:, Ndlys//2+1:, ...] = .50*(left+right) + left = uvp.window_function_array[spw][:, : Ndlys // 2, ...][ + :, ::-1, ... + ] + right = uvp.window_function_array[spw][:, Ndlys // 2 + 1 :, ...] + uvp.window_function_array[spw][:, Ndlys // 2 + 1 :, ...] = 0.50 * ( + left + right + ) else: - leftleft = uvp.window_function_array[spw][:, :Ndlys//2, :Ndlys//2, :][:, ::-1, ::-1, :] - leftright = uvp.window_function_array[spw][:, :Ndlys//2, Ndlys//2+1:, :][:, ::-1, :, :] - rightleft = uvp.window_function_array[spw][:, Ndlys//2+1: , :Ndlys//2, :][:, :, ::-1, :] - rightright = uvp.window_function_array[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] - uvp.window_function_array[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] = .25*(leftleft\ - +leftright\ - +rightleft\ - +rightright) - uvp.window_function_array[spw][:, :, :Ndlys//2, :] = 0.0 - uvp.window_function_array[spw][:, :Ndlys//2, ...] = 0.0 + leftleft = uvp.window_function_array[spw][ + :, : Ndlys // 2, : Ndlys // 2, : + ][:, ::-1, ::-1, :] + leftright = uvp.window_function_array[spw][ + :, : Ndlys // 2, Ndlys // 2 + 1 :, : + ][:, ::-1, :, :] + rightleft = uvp.window_function_array[spw][ + :, Ndlys // 2 + 1 :, : Ndlys // 2, : + ][:, :, ::-1, :] + rightright = uvp.window_function_array[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] + uvp.window_function_array[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] = 0.25 * (leftleft + leftright + rightleft + rightright) + uvp.window_function_array[spw][:, :, : Ndlys // 2, :] = 0.0 + uvp.window_function_array[spw][:, : Ndlys // 2, ...] = 0.0 # fold covariance array if it exists. - if hasattr(uvp,'cov_array_real'): - leftleft = uvp.cov_array_real[spw][:, :Ndlys//2, :Ndlys//2, :][:, ::-1, ::-1, :] - leftright = uvp.cov_array_real[spw][:, :Ndlys//2, Ndlys//2+1:, :][:, ::-1, :, :] - rightleft = uvp.cov_array_real[spw][:, Ndlys//2+1: , :Ndlys//2, :][:, :, ::-1, :] - rightright = uvp.cov_array_real[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] - uvp.cov_array_real[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] = .25*(leftleft\ - +leftright\ - +rightleft\ - +rightright) - uvp.cov_array_real[spw][:, :Ndlys//2, :, :] = 0.0 - uvp.cov_array_real[spw][:, :, :Ndlys//2, : :] = 0.0 - - leftleft = uvp.cov_array_imag[spw][:, :Ndlys//2, :Ndlys//2, :][:, ::-1, ::-1, :] - leftright = uvp.cov_array_imag[spw][:, :Ndlys//2, Ndlys//2+1:, :][:, ::-1, :, :] - rightleft = uvp.cov_array_imag[spw][:, Ndlys//2+1: , :Ndlys//2, :][:, :, ::-1, :] - rightright = uvp.cov_array_imag[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] - uvp.cov_array_imag[spw][:, Ndlys//2+1:, Ndlys//2+1:, :] = .25*(leftleft\ - +leftright\ - +rightleft\ - +rightright) - uvp.cov_array_imag[spw][:, :Ndlys//2, :, :] = 0.0 - uvp.cov_array_imag[spw][:, :, :Ndlys//2, : :] = 0.0 + if hasattr(uvp, "cov_array_real"): + leftleft = uvp.cov_array_real[spw][:, : Ndlys // 2, : Ndlys // 2, :][ + :, ::-1, ::-1, : + ] + leftright = uvp.cov_array_real[spw][ + :, : Ndlys // 2, Ndlys // 2 + 1 :, : + ][:, ::-1, :, :] + rightleft = uvp.cov_array_real[spw][ + :, Ndlys // 2 + 1 :, : Ndlys // 2, : + ][:, :, ::-1, :] + rightright = uvp.cov_array_real[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] + uvp.cov_array_real[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] = 0.25 * (leftleft + leftright + rightleft + rightright) + uvp.cov_array_real[spw][:, : Ndlys // 2, :, :] = 0.0 + uvp.cov_array_real[spw][:, :, : Ndlys // 2, ::] = 0.0 + + leftleft = uvp.cov_array_imag[spw][:, : Ndlys // 2, : Ndlys // 2, :][ + :, ::-1, ::-1, : + ] + leftright = uvp.cov_array_imag[spw][ + :, : Ndlys // 2, Ndlys // 2 + 1 :, : + ][:, ::-1, :, :] + rightleft = uvp.cov_array_imag[spw][ + :, Ndlys // 2 + 1 :, : Ndlys // 2, : + ][:, :, ::-1, :] + rightright = uvp.cov_array_imag[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] + uvp.cov_array_imag[spw][ + :, Ndlys // 2 + 1 :, Ndlys // 2 + 1 :, : + ] = 0.25 * (leftleft + leftright + rightleft + rightright) + uvp.cov_array_imag[spw][:, : Ndlys // 2, :, :] = 0.0 + uvp.cov_array_imag[spw][:, :, : Ndlys // 2, ::] = 0.0 # fold stats array if it exists: sum in inverse quadrature - if hasattr(uvp, 'stats_array'): + if hasattr(uvp, "stats_array"): for stat in uvp.stats_array.keys(): - left = uvp.stats_array[stat][spw][:, :Ndlys//2, :][:, ::-1, :] - right = uvp.stats_array[stat][spw][:, Ndlys//2+1:, :] - uvp.stats_array[stat][spw][:, Ndlys//2+1:, :] = (np.sum([1./left**2.0, 1./right**2.0], axis=0))**(-0.5) - uvp.data_array[spw][:, :Ndlys//2, :] = np.nan + left = uvp.stats_array[stat][spw][:, : Ndlys // 2, :][:, ::-1, :] + right = uvp.stats_array[stat][spw][:, Ndlys // 2 + 1 :, :] + uvp.stats_array[stat][spw][:, Ndlys // 2 + 1 :, :] = ( + np.sum([1.0 / left**2.0, 1.0 / right**2.0], axis=0) + ) ** (-0.5) + uvp.data_array[spw][:, : Ndlys // 2, :] = np.nan uvp.folded = True -def bootstrap_average_blpairs(uvp_list, blpair_groups, time_avg=False, - seed=None): +def bootstrap_average_blpairs(uvp_list, blpair_groups, time_avg=False, seed=None): """ Generate a bootstrap-sampled average over a set of power spectra. The sampling is done over user-defined groups of baseline pairs (with @@ -1318,59 +1547,72 @@ def bootstrap_average_blpairs(uvp_list, blpair_groups, time_avg=False, """ # Validate input data from hera_pspec import UVPSpec + single_uvp = False if isinstance(uvp_list, UVPSpec): single_uvp = True - uvp_list = [uvp_list,] - assert isinstance(uvp_list, list), \ - "uvp_list must be a list of UVPSpec objects" + uvp_list = [ + uvp_list, + ] + assert isinstance(uvp_list, list), "uvp_list must be a list of UVPSpec objects" # Check that uvp_list contains UVPSpec objects with the correct dimensions - assert np.all([isinstance(uvp, UVPSpec) for uvp in uvp_list]), \ - "uvp_list must be a list of UVPSpec objects" + assert np.all( + [isinstance(uvp, UVPSpec) for uvp in uvp_list] + ), "uvp_list must be a list of UVPSpec objects" # Check for same length of time axis if no time averaging will be done if not time_avg: if np.unique([uvp.Ntimes for uvp in uvp_list]).size > 1: - raise IndexError("Input UVPSpec objects must have the same number " - "of time samples if time_avg is False.") + raise IndexError( + "Input UVPSpec objects must have the same number " + "of time samples if time_avg is False." + ) # Check that blpair_groups is a list of lists of groups - assert isinstance(blpair_groups, list), \ - "blpair_groups must be a list of lists of baseline-pair tuples/integers" + assert isinstance( + blpair_groups, list + ), "blpair_groups must be a list of lists of baseline-pair tuples/integers" for grp in blpair_groups: - assert isinstance(grp, list), \ - "blpair_groups must be a list of lists of baseline-pair tuples/integers" + assert isinstance( + grp, list + ), "blpair_groups must be a list of lists of baseline-pair tuples/integers" # Convert blpair tuples into blpair integers if necessary if isinstance(blpair_groups[0][0], tuple): - new_blp_grps = [[uvputils._antnums_to_blpair(blp) for blp in blpg] - for blpg in blpair_groups] + new_blp_grps = [ + [uvputils._antnums_to_blpair(blp) for blp in blpg] for blpg in blpair_groups + ] blpair_groups = new_blp_grps # Homogenise input UVPSpec objects in terms of available polarizations # and spectral windows if len(uvp_list) > 1: - uvp_list = uvpspec_utils.select_common(uvp_list, spws=True, - pols=True, inplace=False) + uvp_list = uvpspec_utils.select_common( + uvp_list, spws=True, pols=True, inplace=False + ) # Loop over UVPSpec objects, looking for available blpairs in each avail_blpairs = [_ordered_unique(uvp.blpair_array) for uvp in uvp_list] - all_avail_blpairs = _ordered_unique([blp for blplist in avail_blpairs - for blp in blplist]) + all_avail_blpairs = _ordered_unique( + [blp for blplist in avail_blpairs for blp in blplist] + ) # Check that all requested blpairs exist in at least one UVPSpec object all_req_blpairs = _ordered_unique([blp for grp in blpair_groups for blp in grp]) missing = [] for blp in all_req_blpairs: - if blp not in all_avail_blpairs: missing.append(blp) + if blp not in all_avail_blpairs: + missing.append(blp) if len(missing) > 0: - raise KeyError("The following baseline-pairs were specified, but do " - "not exist in any of the input UVPSpec objects: %s" \ - % str(missing)) + raise KeyError( + "The following baseline-pairs were specified, but do " + "not exist in any of the input UVPSpec objects: %s" % str(missing) + ) # Set random seed if specified - if seed is not None: np.random.seed(seed) + if seed is not None: + np.random.seed(seed) # Sample with replacement from the full list of available baseline-pairs # in each group, and create a blpair_group and blpair_weights entry for @@ -1401,18 +1643,22 @@ def bootstrap_average_blpairs(uvp_list, blpair_groups, time_avg=False, j = 0 for i in range(len(uvp_list)): n_blps = len(avail[i]) - blpair_grps_list[i].append( list(avail[i]) ) - _wgts = wgt[np.arange(j, j+n_blps)].astype(float) #+ 1e-4 - blpair_wgts_list[i].append( list(_wgts) ) + blpair_grps_list[i].append(list(avail[i])) + _wgts = wgt[np.arange(j, j + n_blps)].astype(float) # + 1e-4 + blpair_wgts_list[i].append(list(_wgts)) j += n_blps # Loop over UVPSpec objects and calculate averages in each blpair group, # using the bootstrap-sampled blpair weights uvp_avg = [] for i, uvp in enumerate(uvp_list): - _uvp = average_spectra(uvp, blpair_groups=blpair_grps_list[i], - blpair_weights=blpair_wgts_list[i], - time_avg=time_avg, inplace=False) + _uvp = average_spectra( + uvp, + blpair_groups=blpair_grps_list[i], + blpair_weights=blpair_wgts_list[i], + time_avg=time_avg, + inplace=False, + ) uvp_avg.append(_uvp) # Return list of averaged spectra for now @@ -1422,9 +1668,19 @@ def bootstrap_average_blpairs(uvp_list, blpair_groups, time_avg=False, return uvp_avg, blpair_wgts_list -def bootstrap_resampled_error(uvp, blpair_groups=None, time_avg=False, Nsamples=1000, seed=None, - normal_std=True, robust_std=False, cintervals=None, bl_error_tol=1.0, - add_to_history='', verbose=False): +def bootstrap_resampled_error( + uvp, + blpair_groups=None, + time_avg=False, + Nsamples=1000, + seed=None, + normal_std=True, + robust_std=False, + cintervals=None, + bl_error_tol=1.0, + add_to_history="", + verbose=False, +): """ Given a UVPSpec object, generate bootstrap resamples of its average and calculate their spread as an estimate of the errorbar on the @@ -1477,8 +1733,11 @@ def bootstrap_resampled_error(uvp, blpair_groups=None, time_avg=False, Nsamples= estimates calculated from the bootstrap resampling. """ from hera_pspec import UVPSpec + # type check - assert isinstance(uvp, (UVPSpec, str, np.str)), "uvp must be fed as a UVPSpec object or filepath" + assert isinstance( + uvp, (UVPSpec, str, np.str) + ), "uvp must be fed as a UVPSpec object or filepath" if isinstance(uvp, (str, np.str)): _uvp = UVPSpec() _uvp.read_hdf5(uvp) @@ -1489,57 +1748,88 @@ def bootstrap_resampled_error(uvp, blpair_groups=None, time_avg=False, Nsamples= blpair_groups, _, _, _ = utils.get_blvec_reds(uvp, bl_error_tol=bl_error_tol) # Uniform average - uvp_avg = average_spectra(uvp, blpair_groups=blpair_groups, time_avg=time_avg, - inplace=False) + uvp_avg = average_spectra( + uvp, blpair_groups=blpair_groups, time_avg=time_avg, inplace=False + ) # initialize a seed - if seed is not None: np.random.seed(seed) + if seed is not None: + np.random.seed(seed) # Iterate over Nsamples and create bootstrap resamples uvp_boots = [] uvp_wgts = [] for i in range(Nsamples): # resample - boot, wgt = bootstrap_average_blpairs(uvp, blpair_groups=blpair_groups, time_avg=time_avg, seed=None) + boot, wgt = bootstrap_average_blpairs( + uvp, blpair_groups=blpair_groups, time_avg=time_avg, seed=None + ) uvp_boots.append(boot) uvp_wgts.append(wgt) # get all keys in uvp_avg and get data from each uvp_boot keys = uvp_avg.get_all_keys() - uvp_boot_data = odict([(k, np.array([u.get_data(k) for u in uvp_boots])) - for k in keys]) + uvp_boot_data = odict( + [(k, np.array([u.get_data(k) for u in uvp_boots])) for k in keys] + ) # calculate various error estimates if normal_std: for k in keys: - nstd = np.std(uvp_boot_data[k].real, axis=0) \ - + 1j*np.std(uvp_boot_data[k].imag, axis=0) + nstd = np.std(uvp_boot_data[k].real, axis=0) + 1j * np.std( + uvp_boot_data[k].imag, axis=0 + ) uvp_avg.set_stats("bs_std", k, nstd) if robust_std: for k in keys: - rstd = np.sqrt(astats.biweight_midvariance(uvp_boot_data[k].real, axis=0)) \ - + 1j*np.sqrt(astats.biweight_midvariance(uvp_boot_data[k].imag, axis=0)) + rstd = np.sqrt( + astats.biweight_midvariance(uvp_boot_data[k].real, axis=0) + ) + 1j * np.sqrt(astats.biweight_midvariance(uvp_boot_data[k].imag, axis=0)) uvp_avg.set_stats("bs_robust_std", k, rstd) if cintervals is not None: for ci in cintervals: ci_tag = "bs_cinterval_{:05.2f}".format(ci) for k in keys: - cint = np.percentile(uvp_boot_data[k].real, ci, axis=0) \ - + 1j*np.percentile(uvp_boot_data[k].imag, ci, axis=0) + cint = np.percentile( + uvp_boot_data[k].real, ci, axis=0 + ) + 1j * np.percentile(uvp_boot_data[k].imag, ci, axis=0) uvp_avg.set_stats(ci_tag, k, cint) # Update history - uvp_avg.history = "Bootstrap errors estimated w/ hera_pspec [{}], {} samples, {} seed\n{}\n{}\n{}" \ - "".format(version.git_hash[:15], Nsamples, seed, add_to_history, '-'*40, uvp_avg.history) + uvp_avg.history = ( + "Bootstrap errors estimated w/ hera_pspec [{}], {} samples, {} seed\n{}\n{}\n{}" + "".format( + version.git_hash[:15], + Nsamples, + seed, + add_to_history, + "-" * 40, + uvp_avg.history, + ) + ) return uvp_avg, uvp_boots, uvp_wgts -def bootstrap_run(filename, spectra=None, blpair_groups=None, time_avg=False, Nsamples=1000, seed=0, - normal_std=True, robust_std=True, cintervals=None, keep_samples=False, - bl_error_tol=1.0, overwrite=False, add_to_history='', verbose=True, maxiter=1): +def bootstrap_run( + filename, + spectra=None, + blpair_groups=None, + time_avg=False, + Nsamples=1000, + seed=0, + normal_std=True, + robust_std=True, + cintervals=None, + keep_samples=False, + bl_error_tol=1.0, + overwrite=False, + add_to_history="", + verbose=True, + maxiter=1, +): """ Run bootstrap resampling on a PSpecContainer object to estimate errorbars. For each group/spectrum specified in the PSpecContainer, this function produces @@ -1612,10 +1902,18 @@ def bootstrap_run(filename, spectra=None, blpair_groups=None, time_avg=False, Ns """ from hera_pspec import uvpspec from hera_pspec import PSpecContainer + # type check if isinstance(filename, (str, np.str)): # open in transactional mode - psc = PSpecContainer(filename, mode='rw', keep_open=False, swmr=False, tsleep=0.5, maxiter=maxiter) + psc = PSpecContainer( + filename, + mode="rw", + keep_open=False, + swmr=False, + tsleep=0.5, + maxiter=maxiter, + ) elif isinstance(filename, PSpecContainer): psc = filename assert not psc.swmr, "PSpecContainer should not be in SWMR mode" @@ -1627,9 +1925,9 @@ def bootstrap_run(filename, spectra=None, blpair_groups=None, time_avg=False, Ns assert len(groups) > 0, "No groups exist in PSpecContainer" # get spectra if not fed - all_spectra = utils.flatten([ [os.path.join(grp, s) - for s in psc.spectra(grp)] - for grp in groups]) + all_spectra = utils.flatten( + [[os.path.join(grp, s) for s in psc.spectra(grp)] for grp in groups] + ) if spectra is None: spectra = all_spectra else: @@ -1639,64 +1937,131 @@ def bootstrap_run(filename, spectra=None, blpair_groups=None, time_avg=False, Ns # iterate over spectra for spc_name in spectra: # split group and spectra names - grp, spc = spc_name.split('/') + grp, spc = spc_name.split("/") # run boostrap_resampled_error uvp = psc.get_pspec(grp, spc) - (uvp_avg, uvp_boots, - uvp_wgts) = bootstrap_resampled_error(uvp, blpair_groups=blpair_groups, - time_avg=time_avg, - Nsamples=Nsamples, seed=seed, - normal_std=normal_std, - robust_std=robust_std, - cintervals=cintervals, - bl_error_tol=bl_error_tol, - add_to_history=add_to_history, - verbose=verbose) + (uvp_avg, uvp_boots, uvp_wgts) = bootstrap_resampled_error( + uvp, + blpair_groups=blpair_groups, + time_avg=time_avg, + Nsamples=Nsamples, + seed=seed, + normal_std=normal_std, + robust_std=robust_std, + cintervals=cintervals, + bl_error_tol=bl_error_tol, + add_to_history=add_to_history, + verbose=verbose, + ) # set averaged uvp - psc.set_pspec(grp, spc+"_avg", uvp_avg, overwrite=overwrite) + psc.set_pspec(grp, spc + "_avg", uvp_avg, overwrite=overwrite) # if keep_samples write uvp_boots if keep_samples: for i, uvpb in enumerate(uvp_boots): - psc.set_pspec(grp, spc+"_bs{}".format(i), uvpb, overwrite=overwrite) + psc.set_pspec(grp, spc + "_bs{}".format(i), uvpb, overwrite=overwrite) def get_bootstrap_run_argparser(): a = argparse.ArgumentParser( - description="argument parser for grouping.bootstrap_run()") + description="argument parser for grouping.bootstrap_run()" + ) def list_of_lists_of_tuples(s): - s = [[int(_x) for _x in x.split()] for x in s.split(',')] + s = [[int(_x) for _x in x.split()] for x in s.split(",")] return s # Add list of arguments - a.add_argument("filename", type=str, - help="Filename of HDF5 container (PSpecContainer) containing " - "input power spectra.") - a.add_argument("--spectra", default=None, type=str, nargs='+', - help="List of power spectra names (with group prefix) to bootstrap over.") - a.add_argument("--blpair_groups", default=None, type=list_of_lists_of_tuples, - help="List of baseline-pair groups (must be space-delimited blpair integers) " - "wrapped in quotes to use in resampling. Default is to solve for and use redundant groups (recommended)." - "Ex: --blpair_groups '101102103104 102103014015, 101013102104' --> " - "[ [((1, 2), (3, 4)), ((2, 3), (4, 5))], [((1, 3), (2, 4))], ...]") - a.add_argument("--time_avg", default=False, type=bool, help="Perform time-average in averaging step.") - a.add_argument("--Nsamples", default=100, type=int, help="Number of bootstrap resamples to generate.") - a.add_argument("--seed", default=0, type=int, help="random seed to initialize bootstrap resampling with.") - a.add_argument("--normal_std", default=True, type=bool, - help="Whether to calculate a 'normal' standard deviation (np.std).") - a.add_argument("--robust_std", default=False, type=bool, - help="Whether to calculate a 'robust' standard deviation (astropy.stats.biweight_midvariance).") - a.add_argument("--cintervals", default=None, type=float, nargs='+', - help="Confidence intervals (precentage from 0 < ci < 100) to calculate.") - a.add_argument("--keep_samples", default=False, action='store_true', - help="If True, store bootstrap resamples in PSpecContainer object with *_bs# extension.") - a.add_argument("--bl_error_tol", type=float, default=1.0, - help="Baseline redudancy tolerance if calculating redundant groups.") - a.add_argument("--overwrite", default=False, action='store_true', help="overwrite outputs if they exist.") - a.add_argument("--add_to_history", default='', type=str, help="String to add to history of power spectra.") - a.add_argument("--verbose", default=False, action='store_true', help="report feedback to stdout.") + a.add_argument( + "filename", + type=str, + help="Filename of HDF5 container (PSpecContainer) containing " + "input power spectra.", + ) + a.add_argument( + "--spectra", + default=None, + type=str, + nargs="+", + help="List of power spectra names (with group prefix) to bootstrap over.", + ) + a.add_argument( + "--blpair_groups", + default=None, + type=list_of_lists_of_tuples, + help="List of baseline-pair groups (must be space-delimited blpair integers) " + "wrapped in quotes to use in resampling. Default is to solve for and use redundant groups (recommended)." + "Ex: --blpair_groups '101102103104 102103014015, 101013102104' --> " + "[ [((1, 2), (3, 4)), ((2, 3), (4, 5))], [((1, 3), (2, 4))], ...]", + ) + a.add_argument( + "--time_avg", + default=False, + type=bool, + help="Perform time-average in averaging step.", + ) + a.add_argument( + "--Nsamples", + default=100, + type=int, + help="Number of bootstrap resamples to generate.", + ) + a.add_argument( + "--seed", + default=0, + type=int, + help="random seed to initialize bootstrap resampling with.", + ) + a.add_argument( + "--normal_std", + default=True, + type=bool, + help="Whether to calculate a 'normal' standard deviation (np.std).", + ) + a.add_argument( + "--robust_std", + default=False, + type=bool, + help="Whether to calculate a 'robust' standard deviation (astropy.stats.biweight_midvariance).", + ) + a.add_argument( + "--cintervals", + default=None, + type=float, + nargs="+", + help="Confidence intervals (precentage from 0 < ci < 100) to calculate.", + ) + a.add_argument( + "--keep_samples", + default=False, + action="store_true", + help="If True, store bootstrap resamples in PSpecContainer object with *_bs# extension.", + ) + a.add_argument( + "--bl_error_tol", + type=float, + default=1.0, + help="Baseline redudancy tolerance if calculating redundant groups.", + ) + a.add_argument( + "--overwrite", + default=False, + action="store_true", + help="overwrite outputs if they exist.", + ) + a.add_argument( + "--add_to_history", + default="", + type=str, + help="String to add to history of power spectra.", + ) + a.add_argument( + "--verbose", + default=False, + action="store_true", + help="report feedback to stdout.", + ) return a diff --git a/hera_pspec/noise.py b/hera_pspec/noise.py index c77a47b5..9710473d 100644 --- a/hera_pspec/noise.py +++ b/hera_pspec/noise.py @@ -1,24 +1,31 @@ import numpy as np -import os -import copy import ast from collections import OrderedDict as odict -from . import conversions, pspecbeam +from . import conversions -def calc_P_N(scalar, Tsys, t_int, Ncoherent=1, Nincoherent=None, form='Pk', k=None, component='real'): +def calc_P_N( + scalar, + Tsys, + t_int, + Ncoherent=1, + Nincoherent=None, + form="Pk", + k=None, + component="real", +): """ Calculate the noise power spectrum via Eqn. (22) of Cheng et al. 2018 for a specified component of the power spectrum. - The noise power spectrum is written as + The noise power spectrum is written as P_N = scalar * (Tsys * 1e3)^2 / (t_int * Ncoherent) / sqrt(Nincoherent) where scalar is a nomalization given by the cosmological model and beam response, i.e. X2Y * Omega_eff - Tsys is the system temp in Kelvin, t_int is the integration time of the underlying data [sec], - Ncoherent is the number of coherent averages before forming power spectra, and Nincoherent is the + Tsys is the system temp in Kelvin, t_int is the integration time of the underlying data [sec], + Ncoherent is the number of coherent averages before forming power spectra, and Nincoherent is the number of incoherent averages after squaring. If component is 'real' or 'imag' an additional factor of 1/sqrt(2) is multiplied. For an ensemble of power spectra with the same Tsys, this estimate should equal their RMS value. @@ -42,33 +49,40 @@ def calc_P_N(scalar, Tsys, t_int, Ncoherent=1, Nincoherent=None, form='Pk', k=No if form == 'DelSq', then units include a factor of h^3 k^3 / (2pi^2) """ # assert form - assert form in ('Pk', 'DelSq'), "form must be either 'Pk' or 'DelSq' for P(k) or Delta^2(k) respectively" - assert component in ['abs', 'real', 'imag'], "component must be one of 'real', 'imag', 'abs'" + assert form in ( + "Pk", + "DelSq", + ), "form must be either 'Pk' or 'DelSq' for P(k) or Delta^2(k) respectively" + assert component in [ + "abs", + "real", + "imag", + ], "component must be one of 'real', 'imag', 'abs'" # construct prefactor in mK^2 - P_N = scalar * (Tsys * 1e3)**2 + P_N = scalar * (Tsys * 1e3) ** 2 # Multiply in effective integration time - P_N /= (t_int * Ncoherent) + P_N /= t_int * Ncoherent # Mulitply in incoherent averaging if Nincoherent is not None: P_N /= np.sqrt(Nincoherent) # parse component - if component in ['real', 'imag']: + if component in ["real", "imag"]: P_N /= np.sqrt(2) # Convert to Delta Sq - if form == 'DelSq': + if form == "DelSq": assert k is not None, "if form == 'DelSq' then k must be fed" - P_N = P_N * k**3 / (2*np.pi**2) + P_N = P_N * k**3 / (2 * np.pi**2) return P_N class Sensitivity(object): - """ Power spectrum thermal sensitivity calculator """ + """Power spectrum thermal sensitivity calculator""" def __init__(self, cosmo=None, beam=None): """ @@ -110,28 +124,30 @@ def set_beam(self, beam): beam : pspecbeam.PSpecBeam instance """ # ensure self.cosmo and beam.cosmo are consistent, if they both exist - if hasattr(beam, 'cosmo'): - if hasattr(self, 'cosmo'): + if hasattr(beam, "cosmo"): + if hasattr(self, "cosmo"): # attach self.cosmo to beam if they both exist beam.cosmo = self.cosmo else: # attach beam.cosmo to self if self.cosmo doesn't exist self.cosmo = beam.cosmo else: - if hasattr(self, 'cosmo'): + if hasattr(self, "cosmo"): # attach self.cosmo to beam if beam.cosmo doesn't exist. beam.cosmo = self.cosmo else: # neither beam nor self have cosmo, raise AssertionError - raise AssertionError("neither self nor beam have a Cosmo_Conversions instance attached. "\ - "See self.set_cosmology().") + raise AssertionError( + "neither self nor beam have a Cosmo_Conversions instance attached. " + "See self.set_cosmology()." + ) self.beam = beam def calc_scalar(self, freqs, pol, num_steps=5000, little_h=True): """ Calculate noise power spectrum prefactor from Eqn. (1) of Pober et al. 2014, ApJ 782, 66, - equal to + equal to scalar = X2Y(z) * Omega_P^2 / Omega_PP @@ -157,55 +173,89 @@ def calc_scalar(self, freqs, pol, num_steps=5000, little_h=True): self.pol : str, polarization used to calculate self.scalar """ # compute scalar - self.scalar = self.beam.compute_pspec_scalar(freqs.min(), freqs.max(), len(freqs), num_steps=num_steps, - pol=pol, little_h=little_h, noise_scalar=True) + self.scalar = self.beam.compute_pspec_scalar( + freqs.min(), + freqs.max(), + len(freqs), + num_steps=num_steps, + pol=pol, + little_h=little_h, + noise_scalar=True, + ) self.subband = freqs self.pol = pol - def calc_P_N(self, Tsys, t_int, Ncoherent=1, Nincoherent=None, form='Pk', k=None, component='real'): + def calc_P_N( + self, + Tsys, + t_int, + Ncoherent=1, + Nincoherent=None, + form="Pk", + k=None, + component="real", + ): """ Calculate the noise power spectrum via Eqn. (22) of Cheng et al. 2018 for a specified component of the power spectrum. - The noise power spectrum is written as + The noise power spectrum is written as P_N = scalar * (Tsys * 1e3)^2 / (t_int * Ncoherent) / sqrt(Nincoherent) where scalar is a nomalization given by the cosmological model and beam response, i.e. X2Y * Omega_eff - Tsys is the system temp in Kelvin, t_int is the integration time of the underlying data [sec], - Ncoherent is the number of coherent averages before forming power spectra, and Nincoherent is the + Tsys is the system temp in Kelvin, t_int is the integration time of the underlying data [sec], + Ncoherent is the number of coherent averages before forming power spectra, and Nincoherent is the number of incoherent averages after squaring. If component is 'real' or 'imag' a factor of 1/sqrt(2) is multiplied. Parameters ---------- - scalar : float, Power spectrum normalization factor: X2Y(z) * Omega_P^2 / Omega_PP - Tsys : float, System temperature in Kelvin - t_int : float, integration time of power spectra in seconds - Ncoherent : int, number of coherent averages of visibility data with integration time t_int + scalar : float + Power spectrum normalization factor: X2Y(z) * Omega_P^2 / Omega_PP + Tsys : float + System temperature in Kelvin + t_int : float + integration time of power spectra in seconds + Ncoherent : int + number of coherent averages of visibility data with integration time t_int Total integration time is t_int * Ncoherent - Nincoherent : int, number of incoherent averages of pspectra (i.e. after squaring). - form : str, power spectra form 'Pk' for P(k) and 'DelSq' for Delta^2(k) - k : float ndarray, cosmological wave-vectors in h Mpc^-1, only needed if form == 'DelSq' + Nincoherent : int + number of incoherent averages of pspectra (i.e. after squaring). + form : str + power spectra form 'Pk' for P(k) and 'DelSq' for Delta^2(k) + k : float ndarray + cosmological wave-vectors in h Mpc^-1, only needed if form == 'DelSq' component : str, options=['real', 'imag', 'abs'] If component is real or imag, divide by an extra factor of sqrt(2) - Returns (P_N) + Returns ------- - P_N : estimated noise power spectrum in units of mK^2 * h^-3 Mpc^3 - if form == 'DelSq', then units include a factor of h^3 k^3 / (2pi^2) + P_N + estimated noise power spectrum in units of mK^2 * h^-3 Mpc^3 + if form == 'DelSq', then units include a factor of h^3 k^3 / (2pi^2) """ # assert scalar exists - assert hasattr(self, 'scalar'), "self.scalar must exist before one can calculate a noise spectrum, see self.calc_scalar()" + assert hasattr( + self, "scalar" + ), "self.scalar must exist before one can calculate a noise spectrum, see self.calc_scalar()" # assert form - assert form in ('Pk', 'DelSq'), "form must be either 'Pk' or 'DelSq' for P(k) or Delta^2(k) respectively" + assert form in ( + "Pk", + "DelSq", + ), "form must be either 'Pk' or 'DelSq' for P(k) or Delta^2(k) respectively" # calculate P_N - P_N = calc_P_N(self.scalar, Tsys, t_int, Ncoherent=Ncoherent, Nincoherent=Nincoherent, form=form, - k=k, component=component) + P_N = calc_P_N( + self.scalar, + Tsys, + t_int, + Ncoherent=Ncoherent, + Nincoherent=Nincoherent, + form=form, + k=k, + component=component, + ) return P_N - - - diff --git a/hera_pspec/parameter.py b/hera_pspec/parameter.py index 08cd4904..9e6a6706 100644 --- a/hera_pspec/parameter.py +++ b/hera_pspec/parameter.py @@ -5,12 +5,17 @@ by hera_pspec. """ + class PSpecParam(object): - def __init__(self, name, description=None, value=None, expected_type=None, form=None): + def __init__( + self, name, description=None, value=None, expected_type=None, form=None + ): """PSpecParam init""" self.name = name self.description = description self.value = value self.expected_type = expected_type self.form = form - self.__doc__ = "name : {}, form : {}, expected_type : {} \n {}".format(name, form, expected_type, description) \ No newline at end of file + self.__doc__ = "name : {}, form : {}, expected_type : {} \n {}".format( + name, form, expected_type, description + ) diff --git a/hera_pspec/plot.py b/hera_pspec/plot.py index 8c0c69d9..ff64917f 100644 --- a/hera_pspec/plot.py +++ b/hera_pspec/plot.py @@ -1,69 +1,84 @@ import numpy as np -import pyuvdata import copy from collections import OrderedDict as odict -import astropy.units as u -import astropy.constants as c from pyuvdata import UVData import uvtools from . import conversions, uvpspec, utils -def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False, - average_times=False, fold=False, plot_noise=False, - delay=True, deltasq=False, legend=False, ax=None, - component='real', lines=True, markers=False, error=None, - times=None, logscale=True, force_plot=False, - label_type='blpairt', plot_stats=None, **kwargs): +def delay_spectrum( + uvp, + blpairs, + spw, + pol, + average_blpairs=False, + average_times=False, + fold=False, + plot_noise=False, + delay=True, + deltasq=False, + legend=False, + ax=None, + component="real", + lines=True, + markers=False, + error=None, + times=None, + logscale=True, + force_plot=False, + label_type="blpairt", + plot_stats=None, + **kwargs +): """ Plot a 1D delay spectrum (or spectra) for a group of baselines. - + Parameters ---------- uvp : UVPspec - UVPSpec object, containing delay spectra for a set of baseline-pairs, + UVPSpec object, containing delay spectra for a set of baseline-pairs, times, polarizations, and spectral windows. - + blpairs : list of tuples or lists of tuples List of baseline-pair tuples, or groups of baseline-pair tuples. - + spw, pol : int or str Which spectral window and polarization to plot. - + average_blpairs : bool, optional If True, average over the baseline pairs within each group. - + average_times : bool, optional If True, average spectra over the time axis. Default: False. - + fold : bool, optional - Whether to fold the power spectrum in :math:`|k_\parallel|`. + Whether to fold the power spectrum in :math:`|k_\parallel|`. Default: False. - + plot_noise : bool, optional Whether to plot noise power spectrum curves or not. Default: False. - + delay : bool, optional - Whether to plot the power spectrum in delay units (ns) or cosmological + Whether to plot the power spectrum in delay units (ns) or cosmological units (h/Mpc). Default: True. - + deltasq : bool, optional - If True, plot dimensionless power spectra, Delta^2. This is ignored if + If True, plot dimensionless power spectra, Delta^2. This is ignored if delay=True. Default: False. - + legend : bool, optional Whether to switch on the plot legend. Default: False. - + ax : matplotlib.axes, optional - Use this to pass in an existing Axes object, which the power spectra - will be added to. (Warning: Labels and legends will not be altered in - this case, even if the existing plot has completely different axis + Use this to pass in an existing Axes object, which the power spectra + will be added to. (Warning: Labels and legends will not be altered in + this case, even if the existing plot has completely different axis labels etc.) If None, a new Axes object will be created. Default: None. component : str, optional Component of complex spectra to plot, options=['abs', 'real', 'imag']. - Default: 'real'. + Default: 'real'. lines : bool, optional If True, plot lines between bandpowers for a given pspectrum. @@ -122,47 +137,56 @@ def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False, # Add ungrouped baseline-pairs into a group of their own (expected by the # averaging routines) blpairs_in = blpairs - blpairs = [] # Must be a list, not an array + blpairs = [] # Must be a list, not an array for i, blpgrp in enumerate(blpairs_in): if not isinstance(blpgrp, list): - blpairs.append([blpairs_in[i],]) + blpairs.append( + [ + blpairs_in[i], + ] + ) else: blpairs.append(blpairs_in[i]) # Average over blpairs or times if requested - blpairs_in = copy.deepcopy(blpairs) # Save input blpair list + blpairs_in = copy.deepcopy(blpairs) # Save input blpair list if average_blpairs: - uvp_plt = uvp.average_spectra(blpair_groups=blpairs, - time_avg=average_times, inplace=False) + uvp_plt = uvp.average_spectra( + blpair_groups=blpairs, time_avg=average_times, inplace=False + ) else: uvp_plt = copy.deepcopy(uvp) if average_times: # Average over times, but not baseline-pairs uvp_plt.average_spectra(time_avg=True, inplace=True) - + # Check plot size if uvp_plt.Ntimes * len(blpairs) > 100 and force_plot == False: - raise ValueError("Trying to plot > 100 spectra... Set force_plot=True to continue.") + raise ValueError( + "Trying to plot > 100 spectra... Set force_plot=True to continue." + ) # Fold the power spectra if requested if fold: uvp_plt.fold_spectra() - + # Convert to Delta^2 units if requested if deltasq and not delay: uvp_plt.convert_to_deltasq() - + # Get x-axis units (delays in ns, or k_parallel in Mpc^-1 or h Mpc^-1) if delay: - dlys = uvp_plt.get_dlys(spw) * 1e9 # ns + dlys = uvp_plt.get_dlys(spw) * 1e9 # ns x = dlys else: k_para = uvp_plt.get_kparas(spw) x = k_para - + # Check plot_stats if plot_stats is not None: - assert plot_stats in uvp_plt.stats_array, "specified key {} not found in stats_array".format(plot_stats) + assert ( + plot_stats in uvp_plt.stats_array + ), "specified key {} not found in stats_array".format(plot_stats) # Plot power spectra for blgrp in blpairs: @@ -170,11 +194,11 @@ def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False, for blp in blgrp: # setup key and casting function key = (spw, blp, pol) - if component == 'real': + if component == "real": cast = np.real - elif component == 'imag': + elif component == "imag": cast = np.imag - elif component == 'abs': + elif component == "abs": cast = np.abs # get data array and repeat x array @@ -188,7 +212,7 @@ def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False, data[flags] = np.nan # get errs if requessted - if error is not None and hasattr(uvp_plt, 'stats_array'): + if error is not None and hasattr(uvp_plt, "stats_array"): if error in uvp_plt.stats_array: errs = uvp_plt.get_stats(error, key) errs[flags] = np.nan @@ -203,14 +227,16 @@ def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False, t = blptimes[i] # form label - if label_type == 'key': + if label_type == "key": label = "{}".format(key) - elif label_type == 'blpair': + elif label_type == "blpair": label = "{}".format(blp) - elif label_type == 'blpairt': + elif label_type == "blpairt": label = "{}, {:0.5f}".format(blp, t) else: - raise ValueError("Couldn't undestand label_type {}".format(label_type)) + raise ValueError( + "Couldn't undestand label_type {}".format(label_type) + ) # plot elements cax = None @@ -220,7 +246,7 @@ def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False, else: _y = y - cax, = ax.plot(x, _y, marker='None', label=label, **kwargs) + (cax,) = ax.plot(x, _y, marker="None", label=label, **kwargs) if markers: if cax is None: @@ -233,17 +259,43 @@ def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False, # plot markers if logscale: # plot positive w/ filled circles - cax, = ax.plot(x[y >= 0], np.abs(y[y >= 0]), c=c, ls='None', marker='o', - markerfacecolor=c, markeredgecolor=c, label=label, **kwargs) + (cax,) = ax.plot( + x[y >= 0], + np.abs(y[y >= 0]), + c=c, + ls="None", + marker="o", + markerfacecolor=c, + markeredgecolor=c, + label=label, + **kwargs + ) # plot negative w/ unfilled circles c = cax.get_color() - cax, = ax.plot(x[y < 0], np.abs(y[y < 0]), c=c, ls='None', marker='o', - markerfacecolor='None', markeredgecolor=c, **kwargs) + (cax,) = ax.plot( + x[y < 0], + np.abs(y[y < 0]), + c=c, + ls="None", + marker="o", + markerfacecolor="None", + markeredgecolor=c, + **kwargs + ) else: - cax, = ax.plot(x, y, c=c, ls='None', marker='o', - markerfacecolor=c, markeredgecolor=c, label=label, **kwargs) - - if error is not None and hasattr(uvp_plt, 'stats_array'): + (cax,) = ax.plot( + x, + y, + c=c, + ls="None", + marker="o", + markerfacecolor=c, + markeredgecolor=c, + label=label, + **kwargs + ) + + if error is not None and hasattr(uvp_plt, "stats_array"): if error in uvp_plt.stats_array: if cax is None: c = None @@ -253,109 +305,134 @@ def delay_spectrum(uvp, blpairs, spw, pol, average_blpairs=False, _y = np.abs(y) else: _y = y - cax = ax.errorbar(x, _y, fmt='none', ecolor=c, yerr=cast(errs[i]), **kwargs) + cax = ax.errorbar( + x, _y, fmt="none", ecolor=c, yerr=cast(errs[i]), **kwargs + ) else: - raise KeyError("Error variable '%s' not found in stats_array of UVPSpec object." % error) + raise KeyError( + "Error variable '%s' not found in stats_array of UVPSpec object." + % error + ) - # If blpairs were averaged, only the first blpair in the group + # If blpairs were averaged, only the first blpair in the group # exists any more (so skip the rest) - if average_blpairs: break - + if average_blpairs: + break + # Set log scale if logscale: - ax.set_yscale('log') - + ax.set_yscale("log") + # Add legend if legend: - ax.legend(loc='upper left') - + ax.legend(loc="upper left") + # Add labels with units if ax.get_xlabel() == "": if delay: ax.set_xlabel(r"$\tau$ $[{\rm ns}]$", fontsize=16) else: - ax.set_xlabel("$k_{\parallel}\ h\ Mpc^{-1}$", fontsize=16) + ax.set_xlabel(r"$k_{\parallel}\ h\ Mpc^{-1}$", fontsize=16) if ax.get_ylabel() == "" and plot_stats is None: - # Sanitize power spectrum units + # Sanitize power spectrum units psunits = uvp_plt.units - if "h^-1" in psunits: psunits = psunits.replace("h^-1", "h^{-1}") - if "h^-3" in psunits: psunits = psunits.replace("h^-3", "h^{-3}") - if "Mpc" in psunits and "\\rm" not in psunits: + if "h^-1" in psunits: + psunits = psunits.replace("h^-1", "h^{-1}") + if "h^-3" in psunits: + psunits = psunits.replace("h^-3", "h^{-3}") + if "Mpc" in psunits and "\\rm" not in psunits: psunits = psunits.replace("Mpc", r"{\rm Mpc}") - if "pi" in psunits and "\\pi" not in psunits: + if "pi" in psunits and "\\pi" not in psunits: psunits = psunits.replace("pi", r"\pi") - + # Power spectrum type if deltasq: - ax.set_ylabel("$\Delta^2$ $[%s]$" % psunits, fontsize=16) + ax.set_ylabel(r"$\Delta^2$ $[%s]$" % psunits, fontsize=16) else: - ax.set_ylabel("$P(k_\parallel)$ $[%s]$" % psunits, fontsize=16) - + ax.set_ylabel(r"$P(k_\parallel)$ $[%s]$" % psunits, fontsize=16) + # Return Figure: the axis is an attribute of figure if new_plot: return fig -def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', - average_blpairs=False, fold=False, delay=True, - deltasq=False, log=True, lst_in_hrs=True, - vmin=None, vmax=None, cmap='YlGnBu', axes=None, - figsize=(14, 6), force_plot=False, times=None, - title_type='blpair', colorbar=True, **kwargs): - """ +def delay_waterfall( + uvp, + blpairs, + spw, + pol, + component="abs-real", + average_blpairs=False, + fold=False, + delay=True, + deltasq=False, + log=True, + lst_in_hrs=True, + vmin=None, + vmax=None, + cmap="YlGnBu", + axes=None, + figsize=(14, 6), + force_plot=False, + times=None, + title_type="blpair", + colorbar=True, + **kwargs +): + r""" Plot a 1D delay spectrum waterfall (or spectra) for a group of baselines. - + Parameters ---------- uvp : UVPspec - UVPSpec object, containing delay spectra for a set of baseline-pairs, + UVPSpec object, containing delay spectra for a set of baseline-pairs, times, polarizations, and spectral windows. - + blpairs : list of tuples or lists of tuples List of baseline-pair tuples, or groups of baseline-pair tuples. - + spw, pol : int or str Which spectral window and polarization to plot. - + component : str Component of complex spectra to plot, options=['abs', 'real', 'imag', 'abs-real', 'abs-imag']. abs-real is abs(real(data)), whereas 'real' is real(data) - Default: 'abs-real'. + Default: 'abs-real'. average_blpairs : bool, optional If True, average over the baseline pairs within each group. - + fold : bool, optional - Whether to fold the power spectrum in :math:`|k_\parallel|`. + Whether to fold the power spectrum in :math:`|k_\parallel|`. Default: False. - + delay : bool, optional - Whether to plot the power spectrum in delay units (ns) or cosmological + Whether to plot the power spectrum in delay units (ns) or cosmological units (h/Mpc). Default: True. - + deltasq : bool, optional - If True, plot dimensionless power spectra, Delta^2. This is ignored if + If True, plot dimensionless power spectra, Delta^2. This is ignored if delay=True. Default: False. - + log : bool, optional Whether to plot the log10 of the data. Default: True. - + lst_in_hrs : bool, optional If True, LST is plotted in hours, otherwise its plotted in radians. vmin, vmax : float, optional - Clip the color scale of the delay spectrum to these min./max. values. + Clip the color scale of the delay spectrum to these min./max. values. If None, use the natural range of the data. Default: None. - + cmap : str, optional Matplotlib colormap to use. Default: 'YlGnBu'. - + axes : array of matplotlib.axes, optional Use this to pass in an existing Axes object or array of axes, which the power spectra will be added to. (Warning: Labels and legends will - not be altered in this case, even if the existing plot has completely different axis + not be altered in this case, even if the existing plot has completely different axis labels etc.) If None, a new Axes object will be created. Default: None. - + figsize : tuple len-2 integer tuple specifying figure size if axes is None @@ -369,7 +446,7 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', title_type : str, optional Type of title to put above plot(s). Options = ['blpair', 'blvec'] blpair : "bls: {bl1} x {bl2}" - blvec : "bl len {len} m & ang {ang} deg" + blvec : "bl len {len} m & ang {ang} deg" colorbar : bool, optional Whether to make a colorbar. Default: True @@ -386,16 +463,26 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', import matplotlib.pyplot as plt # assert component - assert component in ['real', 'abs', 'imag', 'abs-real', 'abs-imag'], "Can't parse specified component {}".format(component) - fix_negval = component in ['real', 'imag'] and log + assert component in [ + "real", + "abs", + "imag", + "abs-real", + "abs-imag", + ], "Can't parse specified component {}".format(component) + fix_negval = component in ["real", "imag"] and log # Add ungrouped baseline-pairs into a group of their own (expected by the # averaging routines) blpairs_in = blpairs - blpairs = [] # Must be a list, not an array + blpairs = [] # Must be a list, not an array for i, blpgrp in enumerate(blpairs_in): if not isinstance(blpgrp, list): - blpairs.append([blpairs_in[i],]) + blpairs.append( + [ + blpairs_in[i], + ] + ) else: blpairs.append(blpairs_in[i]) @@ -417,13 +504,14 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', uvp = uvp.select(times=times, inplace=False) # Average over blpairs or times if requested - blpairs_in = copy.deepcopy(blpairs) # Save input blpair list + blpairs_in = copy.deepcopy(blpairs) # Save input blpair list if average_blpairs: - uvp_plt = uvp.average_spectra(blpair_groups=blpairs, - time_avg=False, inplace=False) + uvp_plt = uvp.average_spectra( + blpair_groups=blpairs, time_avg=False, inplace=False + ) else: uvp_plt = copy.deepcopy(uvp) - + # Fold the power spectra if requested if fold: uvp_plt.fold_spectra() @@ -431,15 +519,15 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', # Convert to Delta^2 units if requested if deltasq and not delay: uvp_plt.convert_to_deltasq() - + # Get x-axis units (delays in ns, or k_parallel in Mpc^-1 or h Mpc^-1) if delay: - dlys = uvp_plt.get_dlys(spw) * 1e9 # ns + dlys = uvp_plt.get_dlys(spw) * 1e9 # ns x = dlys else: k_para = uvp_plt.get_kparas(spw) x = k_para - + # Extract power spectra into array waterfall = odict() for blgrp in blpairs: @@ -454,15 +542,15 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', power[flags, :] = np.nan # get component - if component == 'abs': + if component == "abs": power = np.abs(power) - elif component == 'real': + elif component == "real": power = np.real(power) - elif component == 'abs-real': + elif component == "abs-real": power = np.abs(np.real(power)) - elif component == 'imag': + elif component == "imag": power = np.imag(power) - elif component == 'abs-imag': + elif component == "abs-imag": power = np.abs(np.real(power)) # if real or imag and log is True, set negative values to near zero @@ -473,13 +561,14 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', # assign to waterfall waterfall[key] = power - # If blpairs were averaged, only the first blpair in the group + # If blpairs were averaged, only the first blpair in the group # exists any more (so skip the rest) - if average_blpairs: break - + if average_blpairs: + break + # check for reasonable number of blpairs to plot... Nkeys = len(waterfall) - if Nkeys > 20 and force_plot == False: + if Nkeys > 20 and not force_plot: raise ValueError("Nblps > 20 and force_plot == False, quitting...") # Take logarithm of data if requested @@ -489,7 +578,7 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', logunits = "\log_{10}" else: logunits = "" - + # Create new Axes if none specified new_plot = False if axes is None: @@ -498,13 +587,13 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', Nkeys = len(waterfall) Nside = int(np.ceil(np.sqrt(Nkeys))) fig, axes = plt.subplots(Nside, Nside, figsize=figsize) - + # Ensure axes is an ndarray if isinstance(axes, matplotlib.axes._subplots.Axes): axes = np.array([[axes]]) if isinstance(axes, list): axes = np.array(axes) - + # Ensure its 2D and get side lengths if axes.ndim == 1: axes = axes[:, None] @@ -512,8 +601,7 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', Nvert, Nhorz = axes.shape # Get LST range: setting y-ticks is tricky due to LST wrapping... - y = uvp_plt.lst_avg_array[ - uvp_plt.key_to_indices(list(waterfall.keys())[0])[1] ] + y = uvp_plt.lst_avg_array[uvp_plt.key_to_indices(list(waterfall.keys())[0])[1]] y = np.unwrap(y) if y[0] > np.pi: # if start is closer to 2pi than 0, lower axis by an octave @@ -525,21 +613,29 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', lst_units = "rad" # get baseline vectors - blvecs = dict(zip([uvp_plt.bl_to_antnums(bl) for bl in uvp_plt.bl_array], uvp_plt.get_ENU_bl_vecs())) - - # Sanitize power spectrum units + blvecs = dict( + zip( + [uvp_plt.bl_to_antnums(bl) for bl in uvp_plt.bl_array], + uvp_plt.get_ENU_bl_vecs(), + ) + ) + + # Sanitize power spectrum units psunits = uvp_plt.units - if "h^-1" in psunits: psunits = psunits.replace("h^-1", "h^{-1}") - if "h^-3" in psunits: psunits = psunits.replace("h^-3", "h^{-3}") - if "Hz" in psunits: psunits = psunits.replace("Hz", r"{\rm Hz}") - if "str" in psunits: psunits = psunits.replace("str", r"\,{\rm str}") - if "Mpc" in psunits and "\\rm" not in psunits: + if "h^-1" in psunits: + psunits = psunits.replace("h^-1", "h^{-1}") + if "h^-3" in psunits: + psunits = psunits.replace("h^-3", "h^{-3}") + if "Hz" in psunits: + psunits = psunits.replace("Hz", r"{\rm Hz}") + if "str" in psunits: + psunits = psunits.replace("str", r"\,{\rm str}") + if "Mpc" in psunits and "\\rm" not in psunits: psunits = psunits.replace("Mpc", r"{\rm Mpc}") - if "pi" in psunits and "\\pi" not in psunits: + if "pi" in psunits and "\\pi" not in psunits: psunits = psunits.replace("pi", r"\pi") if "beam normalization not specified" in psunits: - psunits = psunits.replace("beam normalization not specified", - r"{\rm unnormed}") + psunits = psunits.replace("beam normalization not specified", r"{\rm unnormed}") # Iterate over waterfall keys keys = list(waterfall.keys()) @@ -551,7 +647,7 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', # turn off subplot if no more plots to make if k >= Nkeys: - ax.axis('off') + ax.axis("off") continue # get blpair key for this subplot @@ -559,31 +655,42 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', blp = uvp_plt.blpair_to_antnums(key[1]) # plot waterfall - cax = ax.matshow(waterfall[key], cmap=cmap, aspect='auto', - vmin=vmin, vmax=vmax, - extent=[x[0], x[-1], y[-1], y[0]], **kwargs) + cax = ax.matshow( + waterfall[key], + cmap=cmap, + aspect="auto", + vmin=vmin, + vmax=vmax, + extent=[x[0], x[-1], y[-1], y[0]], + **kwargs + ) # ax config - ax.xaxis.set_ticks_position('bottom') + ax.xaxis.set_ticks_position("bottom") ax.tick_params(labelsize=12) - if ax.get_title() == '': - if title_type == 'blpair': + if ax.get_title() == "": + if title_type == "blpair": ax.set_title("bls: {} x {}".format(*blp), y=1) - elif title_type == 'blvec': + elif title_type == "blvec": blv = 0.5 * (blvecs[blp[0]] + blvecs[blp[1]]) lens, angs = utils.get_bl_lens_angs([blv], bl_error_tol=1.0) - ax.set_title("bl len {len:0.2f} m & {ang:0.0f} deg".format(len=lens[0], ang=angs[0]), y=1) + ax.set_title( + "bl len {len:0.2f} m & {ang:0.0f} deg".format( + len=lens[0], ang=angs[0] + ), + y=1, + ) # set colorbar if colorbar: if fix_negval: - cb_extend = 'min' + cb_extend = "min" else: - cb_extend = 'neither' + cb_extend = "neither" cbar = ax.get_figure().colorbar(cax, ax=ax, extend=cb_extend) cbar.ax.tick_params(labelsize=14) if fix_negval: - cbar.ax.set_title("$< 0$",y=-0.05, fontsize=16) + cbar.ax.set_title("$< 0$", y=-0.05, fontsize=16) # configure left-column plots if j == 0: @@ -602,40 +709,71 @@ def delay_waterfall(uvp, blpairs, spw, pol, component='abs-real', else: ax.set_xticklabels([]) - k += 1 + k += 1 # make suptitle if axes[0][0].get_figure()._suptitle is None: if deltasq: - units = "$%s\Delta^2$ $[%s]$" % (logunits, psunits) + units = r"$%s\Delta^2$ $[%s]$" % (logunits, psunits) else: - units = "$%sP(k_\parallel)$ $[%s]$" % (logunits, psunits) + units = r"$%sP(k_\parallel)$ $[%s]$" % (logunits, psunits) spwrange = np.around(np.array(uvp_plt.get_spw_ranges()[spw][:2]) / 1e6, 2) - axes[0][0].get_figure().suptitle("{}\n{} polarization | {} -- {} MHz".format(units, pol, *spwrange), - y=1.03, fontsize=14) + axes[0][0].get_figure().suptitle( + "{}\n{} polarization | {} -- {} MHz".format(units, pol, *spwrange), + y=1.03, + fontsize=14, + ) # Return Axes if new_plot: return fig -def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fold=False, delay=True, - rotate=False, component='abs-real', log10=True, loglog=False, - red_tol=1.0, center_line=False, horizon_lines=False, - title=None, ax=None, cmap='viridis', figsize=(8, 6), - deltasq=False, colorbar=False, cbax=None, vmin=None, vmax=None, - edgecolor='none', flip_xax=False, flip_yax=False, lw=2, - set_bl_tick_major=False, set_bl_tick_minor=False, - xtick_size=10, xtick_rot=0, ytick_size=10, ytick_rot=0, - **kwargs): +def delay_wedge( + uvp, + spw, + pol, + blpairs=None, + times=None, + error_weights=None, + fold=False, + delay=True, + rotate=False, + component="abs-real", + log10=True, + loglog=False, + red_tol=1.0, + center_line=False, + horizon_lines=False, + title=None, + ax=None, + cmap="viridis", + figsize=(8, 6), + deltasq=False, + colorbar=False, + cbax=None, + vmin=None, + vmax=None, + edgecolor="none", + flip_xax=False, + flip_yax=False, + lw=2, + set_bl_tick_major=False, + set_bl_tick_minor=False, + xtick_size=10, + xtick_rot=0, + ytick_size=10, + ytick_rot=0, + **kwargs +): """ Plot a 2D delay spectrum (or spectra) from a UVPSpec object. Note that - all integrations and redundant baselines are averaged (unless specifying + all integrations and redundant baselines are averaged (unless specifying times) before plotting. Note: this deepcopies input uvp before averaging. - + Parameters ---------- uvp : UVPSpec @@ -655,11 +793,11 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fol select on before plotting. Default: None. error_weights : string, optional - error_weights specify which kind of errors we use for weights + error_weights specify which kind of errors we use for weights during averaging power spectra. fold : bool, optional - Whether to fold the power spectrum in k_parallel. + Whether to fold the power spectrum in k_parallel. Default: False. delay : bool, optional @@ -736,7 +874,7 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fol Line-width of horizon and center lines if plotted. Default: 2. set_bl_tick_major : bool, optional - If True, use the baseline lengths as major ticks, rather than default + If True, use the baseline lengths as major ticks, rather than default uniform grid. set_bl_tick_minor : bool, optional @@ -758,10 +896,10 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fol assert isinstance(uvp, uvpspec.UVPSpec), "input uvp must be a UVPSpec object" assert isinstance(spw, (int, np.integer)) assert isinstance(pol, (int, np.integer, tuple)) - fix_negval = component in ['real', 'imag'] and log10 + fix_negval = component in ["real", "imag"] and log10 # check pspec units for little h - little_h = 'h^-3' in uvp.norm_units + little_h = "h^-3" in uvp.norm_units # Create new ax if none specified new_plot = False @@ -777,9 +915,12 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fol # Average across redundant groups and time # this also ensures blpairs are ordered from short_bl --> long_bl - blp_grps, lens, angs, tags = utils.get_blvec_reds(uvp, bl_error_tol=red_tol, - match_bl_lens=True) - uvp.average_spectra(blpair_groups=blp_grps, time_avg=True, error_weights=error_weights, inplace=True) + blp_grps, lens, angs, tags = utils.get_blvec_reds( + uvp, bl_error_tol=red_tol, match_bl_lens=True + ) + uvp.average_spectra( + blpair_groups=blp_grps, time_avg=True, error_weights=error_weights, inplace=True + ) # get blpairs and order by len and enforce bl len ordering anyways blpairs, blpair_seps = uvp.get_blpairs(), uvp.get_blpair_seps() @@ -808,31 +949,34 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fol # Conigure Units psunits = "({})^2\ {}".format(uvp.vis_units, uvp.norm_units) - if "h^-1" in psunits: psunits = psunits.replace("h^-1", "h^{-1}\ ") - if "h^-3" in psunits: psunits = psunits.replace("h^-3", "h^{-3}\ ") - if "Hz" in psunits: psunits = psunits.replace("Hz", r"{\rm Hz}\ ") - if "str" in psunits: psunits = psunits.replace("str", r"\,{\rm str}\ ") - if "Mpc" in psunits and "\\rm" not in psunits: + if "h^-1" in psunits: + psunits = psunits.replace(r"h^-1", "h^{-1}\ ") + if "h^-3" in psunits: + psunits = psunits.replace(r"h^-3", "h^{-3}\ ") + if "Hz" in psunits: + psunits = psunits.replace(r"Hz", r"{\rm Hz}\ ") + if "str" in psunits: + psunits = psunits.replace("str", r"\,{\rm str}\ ") + if "Mpc" in psunits and "\\rm" not in psunits: psunits = psunits.replace("Mpc", r"{\rm Mpc}") - if "pi" in psunits and "\\pi" not in psunits: + if "pi" in psunits and "\\pi" not in psunits: psunits = psunits.replace("pi", r"\pi") if "beam normalization not specified" in psunits: - psunits = psunits.replace("beam normalization not specified", - r"{\rm unnormed}") + psunits = psunits.replace("beam normalization not specified", r"{\rm unnormed}") # get data with shape (Nblpairs, Ndlys) data = [uvp.get_data((spw, blp, pol)).squeeze() for blp in blpairs] # get component - if component == 'real': + if component == "real": data = np.real(data) - elif component == 'abs-real': + elif component == "abs-real": data = np.abs(np.real(data)) - elif component == 'imag': + elif component == "imag": data = np.imag(data) - elif component == 'abs-imag': + elif component == "abs-imag": data = np.abs(np.imag(data)) - elif component == 'abs': + elif component == "abs": data = np.abs(data) else: raise ValueError("Did not understand component {}".format(component)) @@ -848,8 +992,8 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fol # loglog if loglog: - ax.set_xscale('log') - ax.set_yscale('log') + ax.set_xscale("log") + ax.set_yscale("log") ax.yaxis.set_major_formatter(matplotlib.ticker.LogFormatterSciNotation()) ax.yaxis.set_minor_formatter(matplotlib.ticker.NullFormatter()) ax.xaxis.set_major_formatter(matplotlib.ticker.LogFormatterSciNotation()) @@ -861,53 +1005,72 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fol # Get bin edges xdiff = np.diff(x_axis) - x_edges = np.array([x_axis[0]-xdiff[0]/2.0] + list(x_axis[:-1]+xdiff/2.0) + [x_axis[-1]+xdiff[-1]/2.0]) + x_edges = np.array( + [x_axis[0] - xdiff[0] / 2.0] + + list(x_axis[:-1] + xdiff / 2.0) + + [x_axis[-1] + xdiff[-1] / 2.0] + ) ydiff = np.diff(y_axis) - y_edges = np.array([y_axis[0]-ydiff[0]/2.0] + list(y_axis[:-1]+ydiff/2.0) + [y_axis[-1]+ydiff[-1]/2.0]) + y_edges = np.array( + [y_axis[0] - ydiff[0] / 2.0] + + list(y_axis[:-1] + ydiff / 2.0) + + [y_axis[-1] + ydiff[-1] / 2.0] + ) X, Y = np.meshgrid(x_edges, y_edges) - # plot - cax = ax.pcolormesh(X, Y, data, cmap=cmap, edgecolor=edgecolor, lw=0.01, - vmin=vmin, vmax=vmax, **kwargs) + # plot + cax = ax.pcolormesh( + X, + Y, + data, + cmap=cmap, + edgecolor=edgecolor, + lw=0.01, + vmin=vmin, + vmax=vmax, + **kwargs + ) # Configure ticks if set_bl_tick_major: if rotate: - ax.set_xticks([np.around(x, _get_sigfig(x)+2) for x in x_axis]) + ax.set_xticks([np.around(x, _get_sigfig(x) + 2) for x in x_axis]) else: - ax.set_yticks([np.around(x, _get_sigfig(x)+2) for x in y_axis]) + ax.set_yticks([np.around(x, _get_sigfig(x) + 2) for x in y_axis]) if set_bl_tick_minor: if rotate: - ax.set_xticks([np.around(x, _get_sigfig(x)+2) for x in x_axis], - minor=True) + ax.set_xticks( + [np.around(x, _get_sigfig(x) + 2) for x in x_axis], minor=True + ) else: - ax.set_yticks([np.around(x, _get_sigfig(x)+2) for x in y_axis], - minor=True) + ax.set_yticks( + [np.around(x, _get_sigfig(x) + 2) for x in y_axis], minor=True + ) # Add colorbar if colorbar: if fix_negval: - cb_extend = 'min' + cb_extend = "min" else: - cb_extend = 'neither' + cb_extend = "neither" if cbax is None: cbax = ax cbar = fig.colorbar(cax, ax=cbax, extend=cb_extend) if deltasq: - p = "\Delta^2" + p = r"\Delta^2" else: p = "P" if delay: - p = "{}({},\ {})".format(p, r'\tau', r'|\vec{b}|') + p = r"{}({},\ {})".format(p, r"\tau", r"|\vec{b}|") else: - p = "{}({},\ {})".format(p, r'k_\parallel', r'k_\perp') + p = r"{}({},\ {})".format(p, r"k_\parallel", r"k_\perp") if log10: psunits = r"$\log_{{10}}\ {}\ [{}]$".format(p, psunits) else: psunits = r"${}\ [{}]$".format(p, psunits) cbar.set_label(psunits, fontsize=14) if fix_negval: - cbar.ax.set_title("$< 0$",y=-0.05, fontsize=16) + cbar.ax.set_title("$< 0$", y=-0.05, fontsize=16) # Configure tick labels if delay: @@ -920,17 +1083,17 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fol _xlabel = ylabel ylabel = xlabel xlabel = _xlabel - if ax.get_xlabel() == '': + if ax.get_xlabel() == "": ax.set_xlabel(xlabel, fontsize=16) - if ax.get_ylabel() == '': + if ax.get_ylabel() == "": ax.set_ylabel(ylabel, fontsize=16) # Configure center line if center_line: if rotate: - ax.axhline(y=0, color='#000000', ls='--', lw=lw) + ax.axhline(y=0, color="#000000", ls="--", lw=lw) else: - ax.axvline(x=0, color='#000000', ls='--', lw=lw) + ax.axvline(x=0, color="#000000", ls="--", lw=lw) # Plot horizons if horizon_lines: @@ -950,13 +1113,29 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fol bin_edges = y_edges for i, hor in enumerate(horizons): if rotate: - ax.plot(bin_edges[i:i+2], [hor, hor], color='#ffffff', ls='--', lw=lw) + ax.plot( + bin_edges[i : i + 2], [hor, hor], color="#ffffff", ls="--", lw=lw + ) if not uvp.folded: - ax.plot(bin_edges[i:i+2], [-hor, -hor], color='#ffffff', ls='--', lw=lw) + ax.plot( + bin_edges[i : i + 2], + [-hor, -hor], + color="#ffffff", + ls="--", + lw=lw, + ) else: - ax.plot([hor, hor], bin_edges[i:i+2], color='#ffffff', ls='--', lw=lw) + ax.plot( + [hor, hor], bin_edges[i : i + 2], color="#ffffff", ls="--", lw=lw + ) if not uvp.folded: - ax.plot([-hor, -hor], bin_edges[i:i+2], color='#ffffff', ls='--', lw=lw) + ax.plot( + [-hor, -hor], + bin_edges[i : i + 2], + color="#ffffff", + ls="--", + lw=lw, + ) # flip axes if flip_xax: @@ -981,85 +1160,93 @@ def delay_wedge(uvp, spw, pol, blpairs=None, times=None, error_weights=None, fol return fig -def plot_uvdata_waterfalls(uvd, basename, data='data', plot_mode='log', - vmin=None, vmax=None, recenter=False, format='png', - **kwargs): +def plot_uvdata_waterfalls( + uvd, + basename, + data="data", + plot_mode="log", + vmin=None, + vmax=None, + recenter=False, + format="png", + **kwargs +): """ - Plot waterfalls for all baselines and polarizations within a UVData object, + Plot waterfalls for all baselines and polarizations within a UVData object, and save to individual files. - + Parameters ---------- uvd : UVData object - Input data object. Waterfalls will be stored for all baselines and - polarizations within the object; use uvd.select() to remove unwanted + Input data object. Waterfalls will be stored for all baselines and + polarizations within the object; use uvd.select() to remove unwanted information. basename : str - Base filename for the output plots. This must have two placeholders - for baseline ID ('bl') and polarization ('pol'), + Base filename for the output plots. This must have two placeholders + for baseline ID ('bl') and polarization ('pol'), e.g. basename='plots/uvdata.{pol}.{bl}'. - + data : str, optional - Which data array to plot from the UVData object. Options are: + Which data array to plot from the UVData object. Options are: 'data', 'flags', 'nsamples'. Default: 'data'. - + plot_mode : str, optional Plot mode, passed to uvtools.plot.waterfall. Default: 'log'. - + vmin, vmax : float, optional - Min./max. values of the plot colorscale. Default: None (uses the + Min./max. values of the plot colorscale. Default: None (uses the min./max. of the data). - + recenter : bool, optional - Whether to apply recentering (see uvtools.plot.waterfall). + Whether to apply recentering (see uvtools.plot.waterfall). Default: False. - + format : str, optional - The file format of the output images. If None, the image format will be + The file format of the output images. If None, the image format will be deduced from the basename string. Default: 'png'. - + **kwargs : dict - Keyword arguments passed to uvtools.plot.waterfall, which passes them + Keyword arguments passed to uvtools.plot.waterfall, which passes them on to matplotlib.imshow. """ - import matplotlib import matplotlib.pyplot as plt assert isinstance(uvd, UVData), "'uvd' must be a UVData object." - assert data in ['data', 'flags', 'nsamples'], \ - "'%s' not a valid data array; use 'data', 'flags', or 'nsamples'" \ - % data - + assert data in ["data", "flags", "nsamples"], ( + "'%s' not a valid data array; use 'data', 'flags', or 'nsamples'" % data + ) + # Set plot colorscale max/min if specified drng = None - if vmin is not None: + if vmin is not None: assert vmax is not None, "Must also specify vmax if vmin is specified." drng = vmax - vmin - + # Empty figure fig, ax = plt.subplots(1, 1) - + # Loop over antenna pairs and pols for (ant1, ant2, pol), d in uvd.antpairpol_iter(): - + # Get chosen data array - if data == 'data': + if data == "data": pass - elif data == 'flags': + elif data == "flags": d = uvd.get_flags((ant1, ant2, pol)) - elif data == 'nsamples': + elif data == "nsamples": d = uvd.get_nsamples((ant1, ant2, pol)) else: raise KeyError("Invalid data array type '%s'" % data) - + # Make plot - img = uvtools.plot.waterfall(d, mode=plot_mode, mx=vmax, drng=drng, - recenter=recenter, **kwargs) + img = uvtools.plot.waterfall( + d, mode=plot_mode, mx=vmax, drng=drng, recenter=recenter, **kwargs + ) fig.colorbar(img) - + # Save to file - outfile = basename.format(bl="%d.%d"%(ant1, ant2), pol=pol) + outfile = basename.format(bl="%d.%d" % (ant1, ant2), pol=pol) if format is not None: # Make sure format extension is given if outfile[-len(format)].lower() != format.lower(): @@ -1068,12 +1255,14 @@ def plot_uvdata_waterfalls(uvd, basename, data='data', plot_mode='log', fig.savefig(outfile, format=format) fig.clf() + def _get_sigfig(x): return -int(np.floor(np.log10(np.abs(x)))) + def _round_sigfig(x, up=True): - sigfigs = get_sigfig(x) + sigfigs = _get_sigfig(x) if up: return np.ceil(10**sigfigs * x) / 10**sigfigs else: - return np.floor(10**sigfigs * x) / 10**sigfigs \ No newline at end of file + return np.floor(10**sigfigs * x) / 10**sigfigs diff --git a/hera_pspec/pspecbeam.py b/hera_pspec/pspecbeam.py index ed05d413..78e816fa 100644 --- a/hera_pspec/pspecbeam.py +++ b/hera_pspec/pspecbeam.py @@ -10,9 +10,17 @@ from . import conversions as conversions, uvpspec_utils as uvputils -def _compute_pspec_scalar(cosmo, beam_freqs, omega_ratio, pspec_freqs, - num_steps=5000, taper='none', little_h=True, - noise_scalar=False, exact_norm=False): +def _compute_pspec_scalar( + cosmo, + beam_freqs, + omega_ratio, + pspec_freqs, + num_steps=5000, + taper="none", + little_h=True, + noise_scalar=False, + exact_norm=False, +): """ This is not to be used by the novice user to calculate a pspec scalar. Instead, look at the PSpecBeamUV and PSpecBeamGauss classes. @@ -67,9 +75,13 @@ def _compute_pspec_scalar(cosmo, beam_freqs, omega_ratio, pspec_freqs, """ # Get integration freqs df = np.median(np.diff(pspec_freqs)) - integration_freqs = np.linspace(pspec_freqs.min(), - pspec_freqs.min() + df*len(pspec_freqs), - num_steps, endpoint=True, dtype=np.float) + integration_freqs = np.linspace( + pspec_freqs.min(), + pspec_freqs.min() + df * len(pspec_freqs), + num_steps, + endpoint=True, + dtype=np.float, + ) # The interpolations are generally more stable in MHz integration_freqs_MHz = integration_freqs / 1e6 @@ -78,39 +90,44 @@ def _compute_pspec_scalar(cosmo, beam_freqs, omega_ratio, pspec_freqs, redshifts = cosmo.f2z(integration_freqs).flatten() X2Y = np.array([cosmo.X2Y(z, little_h=little_h) for z in redshifts]) - if exact_norm: #Beam and spectral tapering are already taken into account in normalization. We only use averaged X2Y - scalar = integrate.trapz(X2Y, x=integration_freqs)/(np.abs(integration_freqs[-1]-integration_freqs[0])) + if ( + exact_norm + ): # Beam and spectral tapering are already taken into account in normalization. We only use averaged X2Y + scalar = integrate.trapz(X2Y, x=integration_freqs) / ( + np.abs(integration_freqs[-1] - integration_freqs[0]) + ) return scalar # Use linear interpolation to interpolate the frequency-dependent # quantities derived from the beam model to the same frequency grid as the # power spectrum estimation beam_model_freqs_MHz = beam_freqs / 1e6 - dOpp_over_Op2_fit = interp1d(beam_model_freqs_MHz, omega_ratio, - kind='quadratic', fill_value='extrapolate') + dOpp_over_Op2_fit = interp1d( + beam_model_freqs_MHz, omega_ratio, kind="quadratic", fill_value="extrapolate" + ) dOpp_over_Op2 = dOpp_over_Op2_fit(integration_freqs_MHz) # Get B_pp = \int dnu taper^2 and Bp = \int dnu - if taper == 'none': + if taper == "none": dBpp_over_BpSq = np.ones_like(integration_freqs, np.float) else: - dBpp_over_BpSq = dspec.gen_window(taper, len(pspec_freqs))**2. - dBpp_over_BpSq = interp1d(pspec_freqs, dBpp_over_BpSq, kind='nearest', - fill_value='extrapolate')(integration_freqs) - dBpp_over_BpSq /= (integration_freqs[-1] - integration_freqs[0])**2. + dBpp_over_BpSq = dspec.gen_window(taper, len(pspec_freqs)) ** 2.0 + dBpp_over_BpSq = interp1d( + pspec_freqs, dBpp_over_BpSq, kind="nearest", fill_value="extrapolate" + )(integration_freqs) + dBpp_over_BpSq /= (integration_freqs[-1] - integration_freqs[0]) ** 2.0 # Keep dBpp_over_BpSq term or not if noise_scalar: - dBpp_over_BpSq = 1. / (integration_freqs[-1] - integration_freqs[0]) + dBpp_over_BpSq = 1.0 / (integration_freqs[-1] - integration_freqs[0]) # Integrate to get scalar d_inv_scalar = dBpp_over_BpSq * dOpp_over_Op2 / X2Y - scalar = 1. / integrate.trapz(d_inv_scalar, x=integration_freqs) + scalar = 1.0 / integrate.trapz(d_inv_scalar, x=integration_freqs) return scalar class PSpecBeamBase(object): - def __init__(self, cosmo=None): """ Base class for PSpecBeam objects. Provides compute_pspec_scalar() @@ -128,9 +145,18 @@ def __init__(self, cosmo=None): else: self.cosmo = conversions.Cosmo_Conversions() - def compute_pspec_scalar(self, lower_freq, upper_freq, num_freqs, - num_steps=5000, pol='pI', taper='none', - little_h=True, noise_scalar=False, exact_norm=False): + def compute_pspec_scalar( + self, + lower_freq, + upper_freq, + num_freqs, + num_steps=5000, + pol="pI", + taper="none", + little_h=True, + noise_scalar=False, + exact_norm=False, + ): """ Computes the scalar function to convert a power spectrum estimate in "telescope units" to cosmological units @@ -188,22 +214,26 @@ def compute_pspec_scalar(self, lower_freq, upper_freq, num_freqs, Units: h^-3 Mpc^3 or Mpc^3. """ # Get pspec_freqs - pspec_freqs = np.linspace(lower_freq, upper_freq, num_freqs, - endpoint=False) + pspec_freqs = np.linspace(lower_freq, upper_freq, num_freqs, endpoint=False) # Get omega_ratio - omega_ratio = self.power_beam_sq_int(pol) \ - / self.power_beam_int(pol)**2 + omega_ratio = self.power_beam_sq_int(pol) / self.power_beam_int(pol) ** 2 # Get scalar - scalar = _compute_pspec_scalar(self.cosmo, self.beam_freqs, - omega_ratio, pspec_freqs, - num_steps=num_steps, taper=taper, - little_h=little_h, - noise_scalar=noise_scalar, exact_norm=exact_norm) + scalar = _compute_pspec_scalar( + self.cosmo, + self.beam_freqs, + omega_ratio, + pspec_freqs, + num_steps=num_steps, + taper=taper, + little_h=little_h, + noise_scalar=noise_scalar, + exact_norm=exact_norm, + ) return scalar - def Jy_to_mK(self, freqs, pol='pI'): + def Jy_to_mK(self, freqs, pol="pI"): """ Return the multiplicative factor [mK / Jy], to convert a visibility from Jy -> mK, @@ -235,21 +265,39 @@ def Jy_to_mK(self, freqs, pol='pI'): freqs = np.array([freqs]) elif not isinstance(freqs, np.ndarray): raise TypeError("freqs must be fed as a float ndarray") - elif isinstance(freqs, np.ndarray) \ - and freqs.dtype not in (float, np.float, np.float64): + elif isinstance(freqs, np.ndarray) and freqs.dtype not in ( + float, + np.float, + np.float64, + ): raise TypeError("freqs must be fed as a float ndarray") # Check frequency bounds if np.min(freqs) < self.beam_freqs.min(): - print("Warning: min freq {} < self.beam_freqs.min(), extrapolating...".format(np.min(freqs))) + print( + "Warning: min freq {} < self.beam_freqs.min(), extrapolating...".format( + np.min(freqs) + ) + ) if np.max(freqs) > self.beam_freqs.max(): - print("Warning: max freq {} > self.beam_freqs.max(), extrapolating...".format(np.max(freqs))) - - Op = interp1d(self.beam_freqs/1e6, self.power_beam_int(pol=pol), - kind='quadratic', fill_value='extrapolate')(freqs/1e6) - - return 1e-20 * conversions.cgs_units.c**2 \ - / (2 * conversions.cgs_units.kb * freqs**2 * Op) + print( + "Warning: max freq {} > self.beam_freqs.max(), extrapolating...".format( + np.max(freqs) + ) + ) + + Op = interp1d( + self.beam_freqs / 1e6, + self.power_beam_int(pol=pol), + kind="quadratic", + fill_value="extrapolate", + )(freqs / 1e6) + + return ( + 1e-20 + * conversions.cgs_units.c**2 + / (2 * conversions.cgs_units.kb * freqs**2 * Op) + ) def get_Omegas(self, polpairs): """ @@ -272,14 +320,19 @@ def get_Omegas(self, polpairs): # Unpack polpairs into tuples if not isinstance(polpairs, (list, np.ndarray)): if isinstance(polpairs, (tuple, int, np.integer)): - polpairs = [polpairs,] + polpairs = [ + polpairs, + ] else: raise TypeError("polpairs is not a list of integers or tuples") # Convert integers to tuples - polpairs = [uvputils.polpair_int2tuple(p) - if isinstance(p, (int, np.integer, np.int32)) else p - for p in polpairs] + polpairs = [ + uvputils.polpair_int2tuple(p) + if isinstance(p, (int, np.integer, np.int32)) + else p + for p in polpairs + ] # Calculate Omegas for each pol pair OmegaP, OmegaPP = [], [] @@ -292,9 +345,10 @@ def get_Omegas(self, polpairs): # Check for cross-pol; only same-pol calculation currently supported if pol1 != pol2: raise NotImplementedError( - "get_Omegas does not support cross-correlation between " - "two different visibility polarizations yet. " - "Could not calculate Omegas for (%s, %s)" % (pol1, pol2)) + "get_Omegas does not support cross-correlation between " + "two different visibility polarizations yet. " + "Could not calculate Omegas for (%s, %s)" % (pol1, pol2) + ) # Calculate Omegas OmegaP.append(self.power_beam_int(pol=pol1)) @@ -306,7 +360,6 @@ def get_Omegas(self, polpairs): class PSpecBeamGauss(PSpecBeamBase): - def __init__(self, fwhm, beam_freqs, cosmo=None): """ Object to store a simple (frequency independent) Gaussian beam in a @@ -332,7 +385,7 @@ def __init__(self, fwhm, beam_freqs, cosmo=None): else: self.cosmo = conversions.Cosmo_Conversions() - def power_beam_int(self, pol='pI'): + def power_beam_int(self, pol="pI"): """ Computes the integral of the beam over solid angle to give a beam area (in sr). Uses analytic formula that the answer @@ -356,10 +409,15 @@ def power_beam_int(self, pol='pI'): primary_beam_area: float, array-like Primary beam area. """ - return np.ones_like(self.beam_freqs) * 2. * np.pi * self.fwhm**2 \ - / (8. * np.log(2.)) + return ( + np.ones_like(self.beam_freqs) + * 2.0 + * np.pi + * self.fwhm**2 + / (8.0 * np.log(2.0)) + ) - def power_beam_sq_int(self, pol='pI'): + def power_beam_sq_int(self, pol="pI"): """ Computes the integral of the beam**2 over solid angle to give a beam area (in str). Uses analytic formula that the answer @@ -383,12 +441,12 @@ def power_beam_sq_int(self, pol='pI'): primary_beam_area: float, array-like Primary beam area. """ - return np.ones_like(self.beam_freqs) * np.pi * self.fwhm**2 \ - / (8. * np.log(2.)) + return ( + np.ones_like(self.beam_freqs) * np.pi * self.fwhm**2 / (8.0 * np.log(2.0)) + ) class PSpecBeamUV(PSpecBeamBase): - def __init__(self, uvbeam, cosmo=None): """ Object to store the primary beam for a pspec observation. @@ -428,11 +486,11 @@ def __init__(self, uvbeam, cosmo=None): # setup primary power beam self.primary_beam = uvb - if uvb.beam_type == 'efield': + if uvb.beam_type == "efield": self.primary_beam.efield_to_power(inplace=True) self.primary_beam.peak_normalize() - def beam_normalized_response(self, pol='pI', freq=None, x_orientation=None): + def beam_normalized_response(self, pol="pI", freq=None, x_orientation=None): """ Outputs beam response for given polarization as a function of pixels on the sky and input frequencies. @@ -464,17 +522,19 @@ def beam_normalized_response(self, pol='pI', freq=None, x_orientation=None): used to compute resolution """ - if self.primary_beam.beam_type != 'power': - raise ValueError('beam_type must be power') + if self.primary_beam.beam_type != "power": + raise ValueError("beam_type must be power") if self.primary_beam.Naxes_vec > 1: - raise ValueError('Expect scalar for power beam, found vector') - if self.primary_beam._data_normalization.value != 'peak': - raise ValueError('beam must be peak normalized') - if self.primary_beam.pixel_coordinate_system != 'healpix': - raise ValueError('Currently only healpix format supported') + raise ValueError("Expect scalar for power beam, found vector") + if self.primary_beam._data_normalization.value != "peak": + raise ValueError("beam must be peak normalized") + if self.primary_beam.pixel_coordinate_system != "healpix": + raise ValueError("Currently only healpix format supported") nside = self.primary_beam.nside - beam_res = self.primary_beam._interp_freq(freq) # interpolate beam in frequency, based on the data frequencies + beam_res = self.primary_beam._interp_freq( + freq + ) # interpolate beam in frequency, based on the data frequencies beam_res = beam_res[0] if isinstance(pol, (str, np.str)): @@ -484,15 +544,19 @@ def beam_normalized_response(self, pol='pI', freq=None, x_orientation=None): if pol in pol_array: stokes_p_ind = np.where(np.isin(pol_array, pol))[0][0] - beam_res = beam_res[0, 0, stokes_p_ind] # extract the beam with the correct polarization, dim (nfreq X npix) + beam_res = beam_res[ + 0, 0, stokes_p_ind + ] # extract the beam with the correct polarization, dim (nfreq X npix) else: - raise ValueError('Do not have the right polarization information') + raise ValueError("Do not have the right polarization information") - omega = np.sum(beam_res, axis=-1) * np.pi / (3. * nside**2) #compute beam solid angle as a function of frequency + omega = ( + np.sum(beam_res, axis=-1) * np.pi / (3.0 * nside**2) + ) # compute beam solid angle as a function of frequency return beam_res, omega, nside - def power_beam_int(self, pol='pI'): + def power_beam_int(self, pol="pI"): """ Computes the integral of the beam over solid angle to give a beam area (in str) as a function of frequency. Uses function @@ -513,12 +577,12 @@ def power_beam_int(self, pol='pI'): primary_beam_area: float, array-like Scalar integral over beam solid angle. """ - if hasattr(self.primary_beam, 'get_beam_area'): + if hasattr(self.primary_beam, "get_beam_area"): return np.real(self.primary_beam.get_beam_area(pol)) else: raise NotImplementedError("Outdated version of pyuvdata.") - def power_beam_sq_int(self, pol='pI'): + def power_beam_sq_int(self, pol="pI"): """ Computes the integral of the beam**2 over solid angle to give a beam**2 area (in str) as a function of frequency. Uses function @@ -538,14 +602,13 @@ def power_beam_sq_int(self, pol='pI'): ------- primary_beam_area: float, array-like """ - if hasattr(self.primary_beam, 'get_beam_area'): + if hasattr(self.primary_beam, "get_beam_area"): return np.real(self.primary_beam.get_beam_sq_area(pol)) else: raise NotImplementedError("Outdated version of pyuvdata.") class PSpecBeamFromArray(PSpecBeamBase): - def __init__(self, OmegaP, OmegaPP, beam_freqs, cosmo=None, x_orientation=None): """ Primary beam model built from user-defined arrays for the integrals @@ -586,7 +649,8 @@ def __init__(self, OmegaP, OmegaPP, beam_freqs, cosmo=None, x_orientation=None): Orientation in cardinal direction east or north of X dipole. Default keeps polarization in X and Y basis. """ - self.OmegaP = {}; self.OmegaPP = {} + self.OmegaP = {} + self.OmegaPP = {} self.x_orientation = x_orientation # these are allowed pols in AIPS polarization integer convention # see pyuvdata.utils.polstr2num() for details @@ -602,16 +666,21 @@ def __init__(self, OmegaP, OmegaPP, beam_freqs, cosmo=None, x_orientation=None): elif isinstance(OmegaP, np.ndarray) or isinstance(OmegaPP, np.ndarray): # Mixed dict and array types are not allowed - raise TypeError("OmegaP and OmegaPP must both be either dicts " - "or arrays. Mixing dicts and arrays is not " - "allowed.") + raise TypeError( + "OmegaP and OmegaPP must both be either dicts " + "or arrays. Mixing dicts and arrays is not " + "allowed." + ) else: pass # Should now have two dicts if everything is OK - if not isinstance(OmegaP, (odict, dict)) or not isinstance(OmegaPP, (odict, dict)): - raise TypeError("OmegaP and OmegaPP must both be either dicts or " - "arrays.") + if not isinstance(OmegaP, (odict, dict)) or not isinstance( + OmegaPP, (odict, dict) + ): + raise TypeError( + "OmegaP and OmegaPP must both be either dicts or " "arrays." + ) # Check for disallowed polarizations for key in list(OmegaP.keys()): @@ -622,7 +691,7 @@ def __init__(self, OmegaP, OmegaPP, beam_freqs, cosmo=None, x_orientation=None): key = new_key # check its an allowed pol if key not in self.allowed_pols: - raise KeyError("Unrecognized polarization '%s' in OmegaP." % key) + raise KeyError("Unrecognized polarization '%s' in OmegaP." % key) for key in list(OmegaPP.keys()): # turn into pol integer if a pol string if isinstance(key, str): @@ -631,14 +700,16 @@ def __init__(self, OmegaP, OmegaPP, beam_freqs, cosmo=None, x_orientation=None): key = new_key # check its an allowed pol if key not in self.allowed_pols: - raise KeyError("Unrecognized polarization '%s' in OmegaPP." % key) + raise KeyError("Unrecognized polarization '%s' in OmegaPP." % key) # Check for available polarizations for pol in self.allowed_pols: if pol in OmegaP.keys() or pol in OmegaPP.keys(): if pol not in OmegaP.keys() or pol not in OmegaPP.keys(): - raise KeyError("Polarization '%s' must be specified for" - " both OmegaP and OmegaPP." % pol) + raise KeyError( + "Polarization '%s' must be specified for" + " both OmegaP and OmegaPP." % pol + ) # Add arrays for this polarization self.add_pol(pol, OmegaP[pol], OmegaPP[pol]) @@ -649,7 +720,6 @@ def __init__(self, OmegaP, OmegaPP, beam_freqs, cosmo=None, x_orientation=None): else: self.cosmo = cosmo - def add_pol(self, pol, OmegaP, OmegaPP): """ Add OmegaP and OmegaPP for a new polarization. @@ -690,10 +760,13 @@ def add_pol(self, pol, OmegaP, OmegaPP): raise TypeError("OmegaP and OmegaPP must both be array_like.") # Check that array dimensions are consistent - if OmegaP.shape != self.beam_freqs.shape \ - or OmegaPP.shape != self.beam_freqs.shape: - raise ValueError("OmegaP and OmegaPP should both " - "have the same shape as beam_freqs.") + if ( + OmegaP.shape != self.beam_freqs.shape + or OmegaPP.shape != self.beam_freqs.shape + ): + raise ValueError( + "OmegaP and OmegaPP should both " "have the same shape as beam_freqs." + ) # Store arrays self.OmegaP[pol] = OmegaP self.OmegaPP[pol] = OmegaPP @@ -701,7 +774,7 @@ def add_pol(self, pol, OmegaP, OmegaPP): # get available pols self.available_pols = ", ".join(map(uvutils.polnum2str, self.OmegaP.keys())) - def power_beam_int(self, pol='pI'): + def power_beam_int(self, pol="pI"): """ Computes the integral of the beam over solid angle to give a beam area (in str) as a function of frequency. @@ -726,11 +799,12 @@ def power_beam_int(self, pol='pI'): if pol in self.OmegaP.keys(): return self.OmegaP[pol] else: - raise KeyError("OmegaP not specified for polarization '%s'. " - "Available polarizations are: %s" \ - % (pol, self.available_pols)) + raise KeyError( + "OmegaP not specified for polarization '%s'. " + "Available polarizations are: %s" % (pol, self.available_pols) + ) - def power_beam_sq_int(self, pol='pI'): + def power_beam_sq_int(self, pol="pI"): """ Computes the integral of the beam**2 over solid angle to give a beam**2 area (in str) as a function of frequency. @@ -755,16 +829,19 @@ def power_beam_sq_int(self, pol='pI'): if pol in self.OmegaPP.keys(): return self.OmegaPP[pol] else: - raise KeyError("OmegaPP not specified for polarization '%s'. " - "Available polarizations are: %s" \ - % (pol, self.available_pols)) + raise KeyError( + "OmegaPP not specified for polarization '%s'. " + "Available polarizations are: %s" % (pol, self.available_pols) + ) def __str__(self): """ Return a string with useful information about this object. """ s = "PSpecBeamFromArray object\n" - s += "\tFrequency range: Min. %4.4e Hz, Max. %4.4e Hz\n" \ - % (np.min(self.beam_freqs), np.max(self.beam_freqs)) + s += "\tFrequency range: Min. %4.4e Hz, Max. %4.4e Hz\n" % ( + np.min(self.beam_freqs), + np.max(self.beam_freqs), + ) s += "\tAvailable pols: %s" % (self.available_pols) return s diff --git a/hera_pspec/pspecdata.py b/hera_pspec/pspecdata.py index d5496f6c..8fafdf8c 100644 --- a/hera_pspec/pspecdata.py +++ b/hera_pspec/pspecdata.py @@ -17,9 +17,16 @@ class PSpecData(object): - - def __init__(self, dsets=[], wgts=None, dsets_std=None, labels=None, - beam=None, cals=None, cal_flag=True): + def __init__( + self, + dsets=[], + wgts=None, + dsets_std=None, + labels=None, + beam=None, + cals=None, + cal_flag=True, + ): """ Object to store multiple sets of UVData visibilities and perform operations such as power spectrum estimation on them. @@ -58,7 +65,9 @@ def __init__(self, dsets=[], wgts=None, dsets_std=None, labels=None, If True, propagate flags from calibration into data """ self.clear_cache() # clear matrix cache - self.dsets = []; self.wgts = []; self.labels = [] + self.dsets = [] + self.wgts = [] + self.labels = [] self.dsets_std = [] self.Nfreqs = None self.spw_range = None @@ -68,11 +77,11 @@ def __init__(self, dsets=[], wgts=None, dsets_std=None, labels=None, # parametric R matrices. self.r_params = {} self.filter_extension = (0, 0) - self.cov_regularization = 0. + self.cov_regularization = 0.0 # set data weighting to identity by default # and taper to none by default - self.data_weighting = 'identity' - self.taper = 'none' + self.data_weighting = "identity" + self.taper = "none" self.symmetric_taper = True # Set all weights to None if wgts=None if wgts is None: @@ -84,7 +93,14 @@ def __init__(self, dsets=[], wgts=None, dsets_std=None, labels=None, # Store the input UVData objects if specified if len(dsets) > 0: - self.add(dsets, wgts, dsets_std=dsets_std, labels=labels, cals=cals, cal_flag=cal_flag) + self.add( + dsets, + wgts, + dsets_std=dsets_std, + labels=labels, + cals=cals, + cal_flag=cal_flag, + ) # Store a primary beam self.primary_beam = beam @@ -125,27 +141,27 @@ def add(self, dsets, wgts, labels=None, dsets_std=None, cals=None, cal_flag=True if isinstance(dsets, dict): # Disallow labels kwarg if a dict was passed if labels is not None: - raise ValueError("If 'dsets' is a dict, 'labels' cannot be " - "specified.") + raise ValueError( + "If 'dsets' is a dict, 'labels' cannot be " "specified." + ) labels = list(dsets.keys()) if wgts is None: wgts = dict([(l, None) for l in labels]) elif not isinstance(wgts, dict): - raise TypeError("If 'dsets' is a dict, 'wgts' must also be " - "a dict") + raise TypeError("If 'dsets' is a dict, 'wgts' must also be " "a dict") if dsets_std is None: dsets_std = dict([(l, None) for l in labels]) elif not isinstance(dsets_std, dict): - raise TypeError("If 'dsets' is a dict, 'dsets_std' must also be " - "a dict") + raise TypeError( + "If 'dsets' is a dict, 'dsets_std' must also be " "a dict" + ) if cals is None: cals = dict([(l, None) for l in labels]) elif not isinstance(cals, dict): - raise TypeError("If 'cals' is a dict, 'cals' must also be " - "a dict") + raise TypeError("If 'cals' is a dict, 'cals' must also be " "a dict") # Unpack dsets and wgts dicts dsets = [dsets[key] for key in labels] @@ -154,47 +170,83 @@ def add(self, dsets, wgts, labels=None, dsets_std=None, cals=None, cal_flag=True cals = [cals[key] for key in labels] # Convert input args to lists if possible - if isinstance(dsets, UVData): dsets = [dsets,] - if isinstance(wgts, UVData): wgts = [wgts,] - if isinstance(labels, str): labels = [labels,] - if isinstance(dsets_std, UVData): dsets_std = [dsets_std,] - if isinstance(cals, UVCal): cals = [cals,] - if wgts is None: wgts = [wgts,] - if dsets_std is None: dsets_std = [dsets_std for m in range(len(dsets))] - if cals is None: cals = [cals for m in range(len(dsets))] - if isinstance(dsets, tuple): dsets = list(dsets) - if isinstance(wgts, tuple): wgts = list(wgts) - if isinstance(dsets_std, tuple): dsets_std = list(dsets_std) - if isinstance(cals, tuple): cals = list(cals) + if isinstance(dsets, UVData): + dsets = [ + dsets, + ] + if isinstance(wgts, UVData): + wgts = [ + wgts, + ] + if isinstance(labels, str): + labels = [ + labels, + ] + if isinstance(dsets_std, UVData): + dsets_std = [ + dsets_std, + ] + if isinstance(cals, UVCal): + cals = [ + cals, + ] + if wgts is None: + wgts = [ + wgts, + ] + if dsets_std is None: + dsets_std = [dsets_std for m in range(len(dsets))] + if cals is None: + cals = [cals for m in range(len(dsets))] + if isinstance(dsets, tuple): + dsets = list(dsets) + if isinstance(wgts, tuple): + wgts = list(wgts) + if isinstance(dsets_std, tuple): + dsets_std = list(dsets_std) + if isinstance(cals, tuple): + cals = list(cals) # Only allow UVData or lists - if not isinstance(dsets, list) or not isinstance(wgts, list)\ - or not isinstance(dsets_std, list) or not isinstance(cals, list): - raise TypeError("dsets, dsets_std, wgts and cals must be UVData" - "UVCal, or lists of UVData or UVCal") + if ( + not isinstance(dsets, list) + or not isinstance(wgts, list) + or not isinstance(dsets_std, list) + or not isinstance(cals, list) + ): + raise TypeError( + "dsets, dsets_std, wgts and cals must be UVData" + "UVCal, or lists of UVData or UVCal" + ) # Make sure enough weights were specified - assert len(dsets) == len(wgts), \ - "The dsets and wgts lists must have equal length" - assert len(dsets_std) == len(dsets), \ - "The dsets and dsets_std lists must have equal length" - assert len(cals) == len(dsets), \ - "The dsets and cals lists must have equal length" + assert len(dsets) == len( + wgts + ), "The dsets and wgts lists must have equal length" + assert len(dsets_std) == len( + dsets + ), "The dsets and dsets_std lists must have equal length" + assert len(cals) == len( + dsets + ), "The dsets and cals lists must have equal length" if labels is not None: - assert len(dsets) == len(labels), \ - "If labels are specified, the dsets and labels lists " \ + assert len(dsets) == len(labels), ( + "If labels are specified, the dsets and labels lists " "must have equal length" + ) # Check that everything is a UVData object for d, w, s in zip(dsets, wgts, dsets_std): if not isinstance(d, UVData): raise TypeError("Only UVData objects can be used as datasets.") if not isinstance(w, UVData) and w is not None: - raise TypeError("Only UVData objects (or None) can be used as " - "weights.") + raise TypeError( + "Only UVData objects (or None) can be used as " "weights." + ) if not isinstance(s, UVData) and s is not None: - raise TypeError("Only UVData objects (or None) can be used as " - "error sets") + raise TypeError( + "Only UVData objects (or None) can be used as " "error sets" + ) for c in cals: if not isinstance(c, UVCal) and c is not None: raise TypeError("Only UVCal objects can be used for calibration.") @@ -203,18 +255,26 @@ def add(self, dsets, wgts, labels=None, dsets_std=None, cals=None, cal_flag=True if self.labels is None: self.labels = [] if labels is None: - labels = ["dset{:d}".format(i) - for i in range(len(self.dsets), len(dsets) + len(self.dsets))] + labels = [ + "dset{:d}".format(i) + for i in range(len(self.dsets), len(dsets) + len(self.dsets)) + ] # Apply calibration if provided for dset, dset_std, cal in zip(dsets, dsets_std, cals): if cal is not None: if dset is not None: uvutils.uvcalibrate(dset, cal, inplace=True, prop_flags=cal_flag) - dset.extra_keywords['calibration'] = cal.extra_keywords.get('filename', '""') + dset.extra_keywords["calibration"] = cal.extra_keywords.get( + "filename", '""' + ) if dset_std is not None: - uvutils.uvcalibrate(dset_std, cal, inplace=True, prop_flags=cal_flag) - dset_std.extra_keywords['calibration'] = cal.extra_keywords.get('filename', '""') + uvutils.uvcalibrate( + dset_std, cal, inplace=True, prop_flags=cal_flag + ) + dset_std.extra_keywords["calibration"] = cal.extra_keywords.get( + "filename", '""' + ) # Append to list self.dsets += dsets @@ -250,16 +310,28 @@ def __str__(self): # Basic info s = "PSpecData object\n" s += " %d datasets" % len(self.dsets) - if len(self.dsets) == 0: return s + if len(self.dsets) == 0: + return s # Dataset summary for i, d in enumerate(self.dsets): if self.labels[i] is None: - s += " dset (%d): %d bls (freqs=%d, times=%d, pols=%d)\n" \ - % (i, d.Nbls, d.Nfreqs, d.Ntimes, d.Npols) + s += " dset (%d): %d bls (freqs=%d, times=%d, pols=%d)\n" % ( + i, + d.Nbls, + d.Nfreqs, + d.Ntimes, + d.Npols, + ) else: - s += " dset '%s' (%d): %d bls (freqs=%d, times=%d, pols=%d)\n" \ - % (self.labels[i], i, d.Nbls, d.Nfreqs, d.Ntimes, d.Npols) + s += " dset '%s' (%d): %d bls (freqs=%d, times=%d, pols=%d)\n" % ( + self.labels[i], + i, + d.Nbls, + d.Nfreqs, + d.Ntimes, + d.Npols, + ) return s def validate_datasets(self, verbose=True): @@ -272,8 +344,9 @@ def validate_datasets(self, verbose=True): raise ValueError("self.wgts does not have same len as self.dsets") if len(self.dsets_std) != len(self.dsets): - raise ValueError("self.dsets_std does not have the same len as " - "self.dsets") + raise ValueError( + "self.dsets_std does not have the same len as " "self.dsets" + ) if len(self.labels) != len(self.dsets): raise ValueError("self.labels does not have same len as self.dsets") @@ -292,42 +365,59 @@ def validate_datasets(self, verbose=True): # raise warnings if times don't match if len(self.dsets) > 1: - lst_diffs = np.array( [ np.unique(self.dsets[0].lst_array) - - np.unique(dset.lst_array) - for dset in self.dsets[1:]] ) + lst_diffs = np.array( + [ + np.unique(self.dsets[0].lst_array) - np.unique(dset.lst_array) + for dset in self.dsets[1:] + ] + ) if np.max(np.abs(lst_diffs)) > 0.001: - raise_warning("Warning: LST bins in dsets misaligned by more than 15 seconds", - verbose=verbose) + raise_warning( + "Warning: LST bins in dsets misaligned by more than 15 seconds", + verbose=verbose, + ) # raise warning if frequencies don't match - freq_diffs = np.array( [ np.unique(self.dsets[0].freq_array) - - np.unique(dset.freq_array) - for dset in self.dsets[1:]] ) + freq_diffs = np.array( + [ + np.unique(self.dsets[0].freq_array) - np.unique(dset.freq_array) + for dset in self.dsets[1:] + ] + ) if np.max(np.abs(freq_diffs)) > 0.001e6: - raise_warning("Warning: frequency bins in dsets misaligned by more than 0.001 MHz", - verbose=verbose) + raise_warning( + "Warning: frequency bins in dsets misaligned by more than 0.001 MHz", + verbose=verbose, + ) # Check phase type phase_types = [] - for d in self.dsets: phase_types.append(d.phase_type) + for d in self.dsets: + phase_types.append(d.phase_type) if np.unique(phase_types).size > 1: - raise ValueError("all datasets must have the same phase type " - "(i.e. 'drift', 'phased', ...)\ncurrent phase " - "types are {}".format(phase_types)) + raise ValueError( + "all datasets must have the same phase type " + "(i.e. 'drift', 'phased', ...)\ncurrent phase " + "types are {}".format(phase_types) + ) # Check phase centers if phase type is phased - if 'phased' in set(phase_types): + if "phased" in set(phase_types): phase_ra = [d.phase_center_ra_degrees for d in self.dsets] phase_dec = [d.phase_center_dec_degrees for d in self.dsets] - max_diff_ra = np.max( [np.diff(d) - for d in itertools.combinations(phase_ra, 2)]) - max_diff_dec = np.max([np.diff(d) - for d in itertools.combinations(phase_dec, 2)]) + max_diff_ra = np.max( + [np.diff(d) for d in itertools.combinations(phase_ra, 2)] + ) + max_diff_dec = np.max( + [np.diff(d) for d in itertools.combinations(phase_dec, 2)] + ) max_diff = np.sqrt(max_diff_ra**2 + max_diff_dec**2) if max_diff > 0.15: - raise_warning("Warning: maximum phase-center difference " - "between datasets is > 10 arcmin", - verbose=verbose) + raise_warning( + "Warning: maximum phase-center difference " + "between datasets is > 10 arcmin", + verbose=verbose, + ) def check_key_in_dset(self, key, dset_ind): """ @@ -347,7 +437,7 @@ def check_key_in_dset(self, key, dset_ind): exists : bool True if the key exists, False otherwise """ - #FIXME: Fix this to enable label keys + # FIXME: Fix this to enable label keys # get iterable key = uvutils._get_iterable(key) if isinstance(key, str): @@ -378,18 +468,30 @@ def clear_cache(self, keys=None): self._identity_G, self._identity_H, self._identity_Y = {}, {}, {} else: for k in keys: - try: del(self._C[k]) - except(KeyError): pass - try: del(self._I[k]) - except(KeyError): pass - try: del(self._iC[k]) - except(KeyError): pass - try: del(self.r_params[k]) - except(KeyError): pass - try: del(self._Y[k]) - except(KeyError): pass - try: del(self._R[k]) - except(KeyError): pass + try: + del self._C[k] + except (KeyError): + pass + try: + del self._I[k] + except (KeyError): + pass + try: + del self._iC[k] + except (KeyError): + pass + try: + del self.r_params[k] + except (KeyError): + pass + try: + del self._Y[k] + except (KeyError): + pass + try: + del self._R[k] + except (KeyError): + pass def dset_idx(self, dset): """ @@ -455,8 +557,9 @@ def parse_blkey(self, key): # put pol into bl key if it exists if len(key) > 0: pol = key[0] - assert isinstance(pol, (str, int, np.integer)), \ - "pol must be fed as a str or int" + assert isinstance( + pol, (str, int, np.integer) + ), "pol must be fed as a str or int" bl += (key[0],) return dset_idx, bl @@ -508,7 +611,7 @@ def dx(self, key, include_extension=False): Array of std data from the requested UVData dataset and baseline. """ assert isinstance(key, tuple) - dset,bl = self.parse_blkey(key) + dset, bl = self.parse_blkey(key) spw = slice(*self.get_spw(include_extension=include_extension)) return self.dsets_std[dset].get_data(bl).T[spw] @@ -557,7 +660,8 @@ def set_C(self, cov): while the ndarrays should have shape (spw_Nfreqs, spw_Nfreqs) """ self.clear_cache(cov.keys()) - for key in cov: self._C[key] = cov[key] + for key in cov: + self._C[key] = cov[key] def get_spw(self, include_extension=False): """ @@ -577,11 +681,21 @@ def get_spw(self, include_extension=False): include_extension = True # if there is non-zero self.filter_extension, include_extension is automatically set to be True if include_extension: - return (self.spw_range[0] - self.filter_extension[0], self.spw_range[1] + self.filter_extension[1]) + return ( + self.spw_range[0] - self.filter_extension[0], + self.spw_range[1] + self.filter_extension[1], + ) else: return self.spw_range - def C_model(self, key, model='empirical', time_index=None, known_cov=None, include_extension=False): + def C_model( + self, + key, + model="empirical", + time_index=None, + known_cov=None, + include_extension=False, + ): """ Return a covariance model having specified a key and model type. Note: Time-dependent flags that differ from frequency channel-to-channel @@ -633,13 +747,25 @@ def C_model(self, key, model='empirical', time_index=None, known_cov=None, inclu # parse key dset, bl = self.parse_blkey(key) - if model == 'empirical': + if model == "empirical": # add model to key - Ckey = ((dset, dset), (bl,bl), ) + (model, None, False, True,) + Ckey = ((dset, dset), (bl, bl),) + ( + model, + None, + False, + True, + ) else: - assert isinstance(time_index, int), "time_index must be integer if cov-model=={}".format(model) + assert isinstance( + time_index, int + ), "time_index must be integer if cov-model=={}".format(model) # add model to key - Ckey = ((dset, dset), (bl,bl), ) + (model, time_index, False, True,) + Ckey = ((dset, dset), (bl, bl),) + ( + model, + time_index, + False, + True, + ) # Check if Ckey exists in known_cov. If so, just update self._C[Ckey] with known_cov. if known_cov is not None: @@ -651,20 +777,58 @@ def C_model(self, key, model='empirical', time_index=None, known_cov=None, inclu # check cache if Ckey not in self._C: # calculate covariance model - if model == 'empirical': - self.set_C({Ckey: utils.cov(self.x(key, include_extension=include_extension), self.w(key, include_extension=include_extension))}) - elif model == 'dsets': - self.set_C({Ckey: np.diag( np.abs(self.w(key, include_extension=include_extension)[:,time_index] * self.dx(key, include_extension=include_extension)[:,time_index]) ** 2. )}) - elif model == 'autos': + if model == "empirical": + self.set_C( + { + Ckey: utils.cov( + self.x(key, include_extension=include_extension), + self.w(key, include_extension=include_extension), + ) + } + ) + elif model == "dsets": + self.set_C( + { + Ckey: np.diag( + np.abs( + self.w(key, include_extension=include_extension)[ + :, time_index + ] + * self.dx(key, include_extension=include_extension)[ + :, time_index + ] + ) + ** 2.0 + ) + } + ) + elif model == "autos": spw_range = self.get_spw(include_extension=include_extension) - self.set_C({Ckey: np.diag(utils.variance_from_auto_correlations(self.dsets[dset], bl, spw_range, time_index))}) + self.set_C( + { + Ckey: np.diag( + utils.variance_from_auto_correlations( + self.dsets[dset], bl, spw_range, time_index + ) + ) + } + ) else: raise ValueError("didn't recognize Ckey {}".format(Ckey)) return self._C[Ckey] - def cross_covar_model(self, key1, key2, model='empirical', - time_index=None, conj_1=False, conj_2=True, known_cov=None, include_extension=False): + def cross_covar_model( + self, + key1, + key2, + model="empirical", + time_index=None, + conj_1=False, + conj_2=True, + known_cov=None, + include_extension=False, + ): """ Return a covariance model having specified a key and model type. Note: Time-dependent flags that differ from frequency channel-to-channel @@ -727,16 +891,31 @@ def cross_covar_model(self, key1, key2, model='empirical', dset2, bl2 = self.parse_blkey(key2) covar = None - if model == 'empirical': - covar = utils.cov(self.x(key1, include_extension=include_extension), self.w(key1, include_extension=include_extension), - self.x(key2, include_extension=include_extension), self.w(key2, include_extension=include_extension), - conj_1=conj_1, conj_2=conj_2) - if model in ['dsets','autos']: - covar = np.zeros((np.diff(self.get_spw(include_extension=include_extension))[0], - np.diff(self.get_spw(include_extension=include_extension))[0]), dtype=np.float64) + if model == "empirical": + covar = utils.cov( + self.x(key1, include_extension=include_extension), + self.w(key1, include_extension=include_extension), + self.x(key2, include_extension=include_extension), + self.w(key2, include_extension=include_extension), + conj_1=conj_1, + conj_2=conj_2, + ) + if model in ["dsets", "autos"]: + covar = np.zeros( + ( + np.diff(self.get_spw(include_extension=include_extension))[0], + np.diff(self.get_spw(include_extension=include_extension))[0], + ), + dtype=np.float64, + ) # Check if model exists in known_cov. If so, just overwrite covar with known_cov. if known_cov is not None: - Ckey = ((dset1, dset2), (bl1,bl2), ) + (model, time_index, conj_1, conj_2,) + Ckey = ((dset1, dset2), (bl1, bl2),) + ( + model, + time_index, + conj_1, + conj_2, + ) if Ckey in known_cov.keys(): spw = slice(*self.get_spw(include_extension=include_extension)) covar = known_cov[Ckey][spw, spw] @@ -771,7 +950,7 @@ def I(self, key): self._I[key] = np.identity(self.spw_Nfreqs + np.sum(self.filter_extension)) return self._I[key] - def iC(self, key, model='empirical', time_index=None): + def iC(self, key, model="empirical", time_index=None): """ Return the inverse covariance matrix, C^-1. @@ -807,23 +986,28 @@ def iC(self, key, model='empirical', time_index=None): dset, bl = self.parse_blkey(key) key = (dset,) + (bl,) - Ckey = ((dset, dset), (bl,bl), ) + (model, time_index, False, True,) + Ckey = ((dset, dset), (bl, bl),) + ( + model, + time_index, + False, + True, + ) # Calculate inverse covariance if not in cache if Ckey not in self._iC: C = self.C_model(key, model=model, time_index=time_index) - #U,S,V = np.linalg.svd(C.conj()) # conj in advance of next step + # U,S,V = np.linalg.svd(C.conj()) # conj in advance of next step if np.linalg.cond(C) >= 1e9: warnings.warn("Poorly conditioned covariance. Computing Pseudo-Inverse") ic = np.linalg.pinv(C) else: ic = np.linalg.inv(C) # FIXME: Not sure what these are supposed to do - #if self.lmin is not None: S += self.lmin # ensure invertibility - #if self.lmode is not None: S += S[self.lmode-1] + # if self.lmin is not None: S += self.lmin # ensure invertibility + # if self.lmode is not None: S += S[self.lmode-1] # FIXME: Is series of dot products quicker? - self.set_iC({Ckey:ic}) + self.set_iC({Ckey: ic}) return self._iC[Ckey] def Y(self, key): @@ -861,8 +1045,9 @@ def Y(self, key): if key not in self._Y: self._Y[key] = np.diag(np.max(self.w(key), axis=1)) - if not np.all(np.isclose(self._Y[key], 0.0) \ - + np.isclose(self._Y[key], 1.0)): + if not np.all( + np.isclose(self._Y[key], 0.0) + np.isclose(self._Y[key], 1.0) + ): raise NotImplementedError("Non-binary weights not currently implmented") return self._Y[key] @@ -933,17 +1118,25 @@ def R(self, key): # Only add to Rkey if a particular mode is enabled # If you do add to this, you need to specify this in self.set_R docstring! Rkey = key + (self.data_weighting,) + (self.taper,) - if self.data_weighting == 'dayenu': + if self.data_weighting == "dayenu": # add extra dayenu params - Rkey = Rkey + tuple(self.filter_extension,) + (self.spw_Nfreqs,) \ - + (self.symmetric_taper,) + Rkey = ( + Rkey + + tuple( + self.filter_extension, + ) + + (self.spw_Nfreqs,) + + (self.symmetric_taper,) + ) if Rkey not in self._R: # form sqrt(taper) matrix - if self.taper == 'none': + if self.taper == "none": sqrtT = np.ones(self.spw_Nfreqs).reshape(1, -1) else: - sqrtT = np.sqrt(dspec.gen_window(self.taper, self.spw_Nfreqs)).reshape(1, -1) + sqrtT = np.sqrt(dspec.gen_window(self.taper, self.spw_Nfreqs)).reshape( + 1, -1 + ) # get flag weight vector: straight multiplication of vectors # mimics matrix multiplication @@ -954,48 +1147,82 @@ def R(self, key): sqrtT[np.isnan(sqrtT)] = 0.0 sqrtY[np.isnan(sqrtY)] = 0.0 fext = self.filter_extension - #if we want to use a full-band filter, set the R-matrix to filter and then truncate. - tmat = np.zeros((self.spw_Nfreqs, - self.spw_Nfreqs+np.sum(fext)),dtype=complex) - tmat[:,fext[0]:fext[0] + self.spw_Nfreqs] = np.identity(self.spw_Nfreqs,dtype=complex) + # if we want to use a full-band filter, set the R-matrix to filter and then truncate. + tmat = np.zeros( + (self.spw_Nfreqs, self.spw_Nfreqs + np.sum(fext)), dtype=complex + ) + tmat[:, fext[0] : fext[0] + self.spw_Nfreqs] = np.identity( + self.spw_Nfreqs, dtype=complex + ) # form R matrix - if self.data_weighting == 'identity': + if self.data_weighting == "identity": if self.symmetric_taper: - self._R[Rkey] = sqrtT.T * sqrtY.T * self.I(key) * sqrtY * sqrtT + self._R[Rkey] = sqrtT.T * sqrtY.T * self.I(key) * sqrtY * sqrtT else: - self._R[Rkey] = sqrtT.T ** 2. * np.dot(tmat, sqrtY.T * self.I(key) * sqrtY) + self._R[Rkey] = sqrtT.T**2.0 * np.dot( + tmat, sqrtY.T * self.I(key) * sqrtY + ) - elif self.data_weighting == 'iC': + elif self.data_weighting == "iC": if self.symmetric_taper: self._R[Rkey] = sqrtT.T * sqrtY.T * self.iC(key) * sqrtY * sqrtT else: - self._R[Rkey] = sqrtT.T ** 2. * np.dot(tmat, sqrtY.T * self.iC(key) * sqrtY ) + self._R[Rkey] = sqrtT.T**2.0 * np.dot( + tmat, sqrtY.T * self.iC(key) * sqrtY + ) - elif self.data_weighting == 'dayenu': + elif self.data_weighting == "dayenu": r_param_key = (self.data_weighting,) + key if not r_param_key in self.r_params: - raise ValueError("r_param not set for %s!"%str(r_param_key)) + raise ValueError("r_param not set for %s!" % str(r_param_key)) r_params = self.r_params[r_param_key] - if not 'filter_centers' in r_params or\ - not 'filter_half_widths' in r_params or\ - not 'filter_factors' in r_params: - raise ValueError("filtering parameters not specified!") - #This line retrieves a the psuedo-inverse of a lazy covariance - #matrix given by dspec.dayenu_mat_inv. + if ( + not "filter_centers" in r_params + or not "filter_half_widths" in r_params + or not "filter_factors" in r_params + ): + raise ValueError("filtering parameters not specified!") + # This line retrieves a the psuedo-inverse of a lazy covariance + # matrix given by dspec.dayenu_mat_inv. # Note that we multiply sqrtY inside of the pinv - #to apply flagging weights before taking psuedo inverse. + # to apply flagging weights before taking psuedo inverse. if self.symmetric_taper: - self._R[Rkey] = sqrtT.T * np.linalg.pinv(sqrtY.T * \ - dspec.dayenu_mat_inv(x=self.freqs[self.spw_range[0]-fext[0]:self.spw_range[1]+fext[1]], - filter_centers=r_params['filter_centers'], - filter_half_widths=r_params['filter_half_widths'], - filter_factors=r_params['filter_factors']) * sqrtY) * sqrtT + self._R[Rkey] = ( + sqrtT.T + * np.linalg.pinv( + sqrtY.T + * dspec.dayenu_mat_inv( + x=self.freqs[ + self.spw_range[0] + - fext[0] : self.spw_range[1] + + fext[1] + ], + filter_centers=r_params["filter_centers"], + filter_half_widths=r_params["filter_half_widths"], + filter_factors=r_params["filter_factors"], + ) + * sqrtY + ) + * sqrtT + ) else: - self._R[Rkey] = sqrtT.T ** 2. * np.dot(tmat, np.linalg.pinv(sqrtY.T * \ - dspec.dayenu_mat_inv(x=self.freqs[self.spw_range[0]-fext[0]:self.spw_range[1]+fext[1]], - filter_centers=r_params['filter_centers'], - filter_half_widths=r_params['filter_half_widths'], - filter_factors=r_params['filter_factors']) * sqrtY)) + self._R[Rkey] = sqrtT.T**2.0 * np.dot( + tmat, + np.linalg.pinv( + sqrtY.T + * dspec.dayenu_mat_inv( + x=self.freqs[ + self.spw_range[0] + - fext[0] : self.spw_range[1] + + fext[1] + ], + filter_centers=r_params["filter_centers"], + filter_half_widths=r_params["filter_half_widths"], + filter_factors=r_params["filter_factors"], + ) + * sqrtY + ), + ) return self._R[Rkey] @@ -1014,13 +1241,15 @@ def set_symmetric_taper(self, use_symmetric_taper): use_taper : bool, do you want to use a symmetric taper? True or False? """ - if use_symmetric_taper and (self.filter_extension[0] > 0 or self.filter_extension[1] > 0): - raise ValueError("You cannot use a symmetric taper when there are nonzero filter extensions.") + if use_symmetric_taper and ( + self.filter_extension[0] > 0 or self.filter_extension[1] > 0 + ): + raise ValueError( + "You cannot use a symmetric taper when there are nonzero filter extensions." + ) else: self.symmetric_taper = use_symmetric_taper - - def set_filter_extension(self, filter_extension): """ Set extensions to filtering matrix @@ -1032,22 +1261,40 @@ def set_filter_extension(self, filter_extension): filter will be applied to data. filter_extensions will be clipped to not extend beyond data range. """ - if self.symmetric_taper and not filter_extension[0] == 0 and not filter_extension[1]==0: - raise_warning("You cannot set filter extensions greater then zero when symmetric_taper==True! Setting symmetric_taper==False!") + if ( + self.symmetric_taper + and not filter_extension[0] == 0 + and not filter_extension[1] == 0 + ): + raise_warning( + "You cannot set filter extensions greater then zero when symmetric_taper==True! Setting symmetric_taper==False!" + ) self.symmetric_taper = False - assert isinstance(filter_extension, (list, tuple)), "filter_extension must a tuple or list" + assert isinstance( + filter_extension, (list, tuple) + ), "filter_extension must a tuple or list" assert len(filter_extension) == 2, "filter extension must be length 2" - assert isinstance(filter_extension[0], int) and\ - isinstance(filter_extension[1], int) and \ - filter_extension[0] >= 0 and\ - filter_extension[1] >=0, "filter extension must contain only positive integers" - filter_extension=list(filter_extension) + assert ( + isinstance(filter_extension[0], int) + and isinstance(filter_extension[1], int) + and filter_extension[0] >= 0 + and filter_extension[1] >= 0 + ), "filter extension must contain only positive integers" + filter_extension = list(filter_extension) if filter_extension[0] > self.spw_range[0]: - warnings.warn("filter_extension[0] exceeds data spw_range. Defaulting to spw_range[0]!") + warnings.warn( + "filter_extension[0] exceeds data spw_range. Defaulting to spw_range[0]!" + ) if filter_extension[1] > self.Nfreqs - self.spw_range[1]: - warnings.warn("filter_extension[1] exceeds channels between spw_range[1] and Nfreqs. Defaulting to Nfreqs-spw_range[1]!") - filter_extension[0] = np.min([self.spw_range[0], filter_extension[0]])#clip extension to not extend beyond data range - filter_extension[1] = np.min([self.Nfreqs - self.spw_range[1], filter_extension[1]])#clip extension to not extend beyond data range + warnings.warn( + "filter_extension[1] exceeds channels between spw_range[1] and Nfreqs. Defaulting to Nfreqs-spw_range[1]!" + ) + filter_extension[0] = np.min( + [self.spw_range[0], filter_extension[0]] + ) # clip extension to not extend beyond data range + filter_extension[1] = np.min( + [self.Nfreqs - self.spw_range[1], filter_extension[1]] + ) # clip extension to not extend beyond data range self.filter_extension = tuple(filter_extension) def set_weighting(self, data_weighting): @@ -1119,10 +1366,12 @@ def set_spw(self, spw_range, ndlys=None): Number of delay bins. Default: None, sets number of delay bins equal to the number of frequency channels in the spw. """ - assert isinstance(spw_range, tuple), \ - "spw_range must be fed as a len-2 integer tuple" - assert isinstance(spw_range[0], (int, np.integer)), \ - "spw_range must be fed as len-2 integer tuple" + assert isinstance( + spw_range, tuple + ), "spw_range must be fed as a len-2 integer tuple" + assert isinstance( + spw_range[0], (int, np.integer) + ), "spw_range must be fed as len-2 integer tuple" self.spw_range = spw_range self.spw_Nfreqs = spw_range[1] - spw_range[0] self.set_Ndlys(ndlys=ndlys) @@ -1143,11 +1392,20 @@ def set_Ndlys(self, ndlys=None): else: # Check that one is not trying to estimate more delay channels than there are frequencies if self.spw_Nfreqs < ndlys: - raise ValueError("Cannot estimate more delays than there are frequency channels") + raise ValueError( + "Cannot estimate more delays than there are frequency channels" + ) self.spw_Ndlys = ndlys - def cov_q_hat(self, key1, key2, model='empirical', exact_norm=False, pol=False, - time_indices=None): + def cov_q_hat( + self, + key1, + key2, + model="empirical", + exact_norm=False, + pol=False, + time_indices=None, + ): """ Compute the un-normalized covariance matrix for q_hat for a given pair of visibility vectors. Returns the following matrix: @@ -1204,32 +1462,45 @@ def cov_q_hat(self, key1, key2, model='empirical', exact_norm=False, pol=False, time_indices = [time_indices] if not isinstance(time_indices, list): raise ValueError("time_indices must be an integer or list of integers.") - if isinstance(key1,list): + if isinstance(key1, list): assert isinstance(key2, list), "key1 is a list, key2 must be a list" assert len(key2) == len(key1), "key1 length must equal key2 length" - if isinstance(key2,list): + if isinstance(key2, list): assert isinstance(key1, list), "key2 is a list, key1 must be a list" - #check time_indices + # check time_indices for tind in time_indices: if not (tind >= 0 and tind <= self.Ntimes): raise ValueError("Invalid time index provided.") - if not isinstance(key1,list): + if not isinstance(key1, list): key1 = [key1] - if not isinstance(key2,list): + if not isinstance(key2, list): key2 = [key2] - output = np.zeros((len(time_indices), self.spw_Ndlys, self.spw_Ndlys), dtype=complex) + output = np.zeros( + (len(time_indices), self.spw_Ndlys, self.spw_Ndlys), dtype=complex + ) for k1, k2 in zip(key1, key2): - if model == 'dsets': - output+=1./np.asarray([self.get_unnormed_V(k1, k2, model=model, - exact_norm=exact_norm, pol=pol, time_index=t)\ - for t in time_indices]) - - elif model == 'empirical': - cm = self.get_unnormed_V(k1, k2, model=model, - exact_norm=exact_norm, pol=pol) - output+=1./np.asarray([cm for m in range(len(time_indices))]) + if model == "dsets": + output += 1.0 / np.asarray( + [ + self.get_unnormed_V( + k1, + k2, + model=model, + exact_norm=exact_norm, + pol=pol, + time_index=t, + ) + for t in time_indices + ] + ) + + elif model == "empirical": + cm = self.get_unnormed_V( + k1, k2, model=model, exact_norm=exact_norm, pol=pol + ) + output += 1.0 / np.asarray([cm for m in range(len(time_indices))]) return float(len(key1)) / output @@ -1299,7 +1570,7 @@ def q_hat(self, key1, key2, allow_fft=False, exact_norm=False, pol=False): R1 += self.R(_key) else: Rx1 = np.dot(self.R(key1), self.x(key1)) - R1 = self.R(key1) + R1 = self.R(key1) # Calculate R x_2 if isinstance(key2, list): @@ -1308,21 +1579,29 @@ def q_hat(self, key1, key2, allow_fft=False, exact_norm=False, pol=False): R2 += self.R(_key) else: Rx2 = np.dot(self.R(key2), self.x(key2)) - R2 = self.R(key2) + R2 = self.R(key2) # The set of operations for exact_norm == True are drawn from Equations # 11(a) and 11(b) from HERA memo #44. We are incorporating the # multiplicatives to the exponentials, and sticking to quantities in # their physical units. - if exact_norm and allow_fft: #exact_norm approach is meant to enable non-uniform binnning as well, where FFT is not - #applicable. As of now, we are using uniform binning. - raise NotImplementedError("Exact normalization does not support FFT approach at present") + if ( + exact_norm and allow_fft + ): # exact_norm approach is meant to enable non-uniform binnning as well, where FFT is not + # applicable. As of now, we are using uniform binning. + raise NotImplementedError( + "Exact normalization does not support FFT approach at present" + ) - elif exact_norm and not(allow_fft): - q = [] - del_tau = np.median(np.diff(self.delays()))*1e-9 #Get del_eta in Eq.11(a) (HERA memo #44) (seconds) - integral_beam = self.get_integral_beam(pol) #Integral of beam in Eq.11(a) (HERA memo #44) + elif exact_norm and not (allow_fft): + q = [] + del_tau = ( + np.median(np.diff(self.delays())) * 1e-9 + ) # Get del_eta in Eq.11(a) (HERA memo #44) (seconds) + integral_beam = self.get_integral_beam( + pol + ) # Integral of beam in Eq.11(a) (HERA memo #44) for i in range(self.spw_Ndlys): # Ideally, del_tau and integral_beam should be part of get_Q. We use them here to @@ -1331,25 +1610,28 @@ def q_hat(self, key1, key2, allow_fft=False, exact_norm=False, pol=False): QRx2 = np.dot(Q, Rx2) # Square and sum over columns - qi = 0.5 * np.einsum('i...,i...->...', Rx1.conj(), QRx2) + qi = 0.5 * np.einsum("i...,i...->...", Rx1.conj(), QRx2) q.append(qi) - q = np.asarray(q) #(Ndlys X Ntime) + q = np.asarray(q) # (Ndlys X Ntime) return q # use FFT if possible and allowed elif allow_fft and (self.spw_Nfreqs == self.spw_Ndlys): _Rx1 = np.fft.fft(Rx1, axis=0) _Rx2 = np.fft.fft(Rx2, axis=0) - return 0.5 * np.fft.fftshift(_Rx1, axes=0).conj() \ - * np.fft.fftshift(_Rx2, axes=0) + return ( + 0.5 + * np.fft.fftshift(_Rx1, axes=0).conj() + * np.fft.fftshift(_Rx2, axes=0) + ) else: q = [] for i in range(self.spw_Ndlys): Q = self.get_Q_alt(i) QRx2 = np.dot(Q, Rx2) - qi = np.einsum('i...,i...->...', Rx1.conj(), QRx2) + qi = np.einsum("i...,i...->...", Rx1.conj(), QRx2) q.append(qi) return 0.5 * np.array(q) @@ -1387,45 +1669,47 @@ def get_G(self, key1, key2, exact_norm=False, pol=False): Fisher matrix, with dimensions (Nfreqs, Nfreqs). """ if self.spw_Ndlys == None: - raise ValueError("Number of delay bins should have been set" - "by now! Cannot be equal to None") + raise ValueError( + "Number of delay bins should have been set" + "by now! Cannot be equal to None" + ) G = np.zeros((self.spw_Ndlys, self.spw_Ndlys), dtype=np.complex) R1 = self.R(key1) R2 = self.R(key2) iR1Q1, iR2Q2 = {}, {} - if (exact_norm): + if exact_norm: integral_beam = self.get_integral_beam(pol) - del_tau = np.median(np.diff(self.delays()))*1e-9 + del_tau = np.median(np.diff(self.delays())) * 1e-9 if exact_norm: - qnorm = del_tau * integral_beam + qnorm = del_tau * integral_beam else: - qnorm = 1. + qnorm = 1.0 for ch in range(self.spw_Ndlys): - #G is given by Tr[E^\alpha C,\beta] - #where E^\alpha = R_1^\dagger Q^\apha R_2 - #C,\beta = Q2 and Q^\alpha = Q1 - #Note that we conjugate transpose R - #because we want to E^\alpha to - #give the absolute value squared of z = m_\alpha \dot R @ x - #where m_alpha takes the FT from frequency to the \alpha fourier mode. - #Q is essentially m_\alpha^\dagger m + # G is given by Tr[E^\alpha C,\beta] + # where E^\alpha = R_1^\dagger Q^\apha R_2 + # C,\beta = Q2 and Q^\alpha = Q1 + # Note that we conjugate transpose R + # because we want to E^\alpha to + # give the absolute value squared of z = m_\alpha \dot R @ x + # where m_alpha takes the FT from frequency to the \alpha fourier mode. + # Q is essentially m_\alpha^\dagger m # so we need to sandwhich it between R_1^\dagger and R_2 Q1 = self.get_Q_alt(ch) * qnorm Q2 = self.get_Q_alt(ch, include_extension=True) * qnorm - iR1Q1[ch] = np.dot(np.conj(R1).T, Q1) # R_1 Q - iR2Q2[ch] = np.dot(R2, Q2) # R_2 Q + iR1Q1[ch] = np.dot(np.conj(R1).T, Q1) # R_1 Q + iR2Q2[ch] = np.dot(R2, Q2) # R_2 Q for i in range(self.spw_Ndlys): for j in range(self.spw_Ndlys): # tr(R_2 Q_i R_1 Q_j) - G[i,j] = np.einsum('ab,ba', iR1Q1[i], iR2Q2[j]) + G[i, j] = np.einsum("ab,ba", iR1Q1[i], iR2Q2[j]) # check if all zeros, in which case turn into identity if np.count_nonzero(G) == 0: G = np.eye(self.spw_Ndlys) - return G / 2. + return G / 2.0 def get_H(self, key1, key2, sampling=False, exact_norm=False, pol=False): """ @@ -1486,55 +1770,57 @@ def get_H(self, key1, key2, sampling=False, exact_norm=False, pol=False): Dimensions (Nfreqs, Nfreqs). """ if self.spw_Ndlys == None: - raise ValueError("Number of delay bins should have been set" - "by now! Cannot be equal to None.") + raise ValueError( + "Number of delay bins should have been set" + "by now! Cannot be equal to None." + ) H = np.zeros((self.spw_Ndlys, self.spw_Ndlys), dtype=np.complex) R1 = self.R(key1) R2 = self.R(key2) if not sampling: - nfreq=np.sum(self.filter_extension) + self.spw_Nfreqs + nfreq = np.sum(self.filter_extension) + self.spw_Nfreqs sinc_matrix = np.zeros((nfreq, nfreq)) for i in range(nfreq): for j in range(nfreq): - sinc_matrix[i,j] = np.float(i - j) + sinc_matrix[i, j] = np.float(i - j) sinc_matrix = np.sinc(sinc_matrix / np.float(nfreq)) iR1Q1, iR2Q2 = {}, {} - if (exact_norm): + if exact_norm: integral_beam = self.get_integral_beam(pol) - del_tau = np.median(np.diff(self.delays()))*1e-9 + del_tau = np.median(np.diff(self.delays())) * 1e-9 if exact_norm: qnorm = del_tau * integral_beam else: - qnorm = 1. + qnorm = 1.0 for ch in range(self.spw_Ndlys): Q1 = self.get_Q_alt(ch) * qnorm Q2 = self.get_Q_alt(ch, include_extension=True) * qnorm if not sampling: Q2 *= sinc_matrix - #H is given by Tr([E^\alpha C,\beta]) - #where E^\alpha = R_1^\dagger Q^\apha R_2 - #C,\beta = Q2 and Q^\alpha = Q1 - #Note that we conjugate transpose R - #because we want to E^\alpha to - #give the absolute value squared of z = m_\alpha \dot R @ x - #where m_alpha takes the FT from frequency to the \alpha fourier mode. - #Q is essentially m_\alpha^\dagger m + # H is given by Tr([E^\alpha C,\beta]) + # where E^\alpha = R_1^\dagger Q^\apha R_2 + # C,\beta = Q2 and Q^\alpha = Q1 + # Note that we conjugate transpose R + # because we want to E^\alpha to + # give the absolute value squared of z = m_\alpha \dot R @ x + # where m_alpha takes the FT from frequency to the \alpha fourier mode. + # Q is essentially m_\alpha^\dagger m # so we need to sandwhich it between R_1^\dagger and R_2 - iR1Q1[ch] = np.dot(np.conj(R1).T, Q1) # R_1 Q_alt - iR2Q2[ch] = np.dot(R2, Q2) # R_2 Q + iR1Q1[ch] = np.dot(np.conj(R1).T, Q1) # R_1 Q_alt + iR2Q2[ch] = np.dot(R2, Q2) # R_2 Q - for i in range(self.spw_Ndlys): # this loop goes as nchan^4 + for i in range(self.spw_Ndlys): # this loop goes as nchan^4 for j in range(self.spw_Ndlys): # tr(R_2 Q_i R_1 Q_j) - H[i,j] = np.einsum('ab,ba', iR1Q1[i], iR2Q2[j]) + H[i, j] = np.einsum("ab,ba", iR1Q1[i], iR2Q2[j]) # check if all zeros, in which case turn into identity if np.count_nonzero(H) == 0: H = np.eye(self.spw_Ndlys) - return H / 2. + return H / 2.0 def get_unnormed_E(self, key1, key2, exact_norm=False, pol=False): """ @@ -1578,26 +1864,35 @@ def get_unnormed_E(self, key1, key2, exact_norm=False, pol=False): """ if self.spw_Ndlys == None: - raise ValueError("Number of delay bins should have been set" - "by now! Cannot be equal to None") + raise ValueError( + "Number of delay bins should have been set" + "by now! Cannot be equal to None" + ) nfreq = self.spw_Nfreqs + np.sum(self.filter_extension) - E_matrices = np.zeros((self.spw_Ndlys, nfreq, nfreq), - dtype=np.complex) + E_matrices = np.zeros((self.spw_Ndlys, nfreq, nfreq), dtype=np.complex) R1 = self.R(key1) R2 = self.R(key2) - if (exact_norm): + if exact_norm: integral_beam = self.get_integral_beam(pol) - del_tau = np.median(np.diff(self.delays()))*1e-9 + del_tau = np.median(np.diff(self.delays())) * 1e-9 for dly_idx in range(self.spw_Ndlys): - if exact_norm: QR2 = del_tau * integral_beam * np.dot(self.get_Q_alt(dly_idx), R2) - else: QR2 = np.dot(self.get_Q_alt(dly_idx), R2) + if exact_norm: + QR2 = del_tau * integral_beam * np.dot(self.get_Q_alt(dly_idx), R2) + else: + QR2 = np.dot(self.get_Q_alt(dly_idx), R2) E_matrices[dly_idx] = np.dot(np.conj(R1).T, QR2) return 0.5 * E_matrices - - def get_unnormed_V(self, key1, key2, model='empirical', exact_norm=False, - pol=False, time_index=None): + def get_unnormed_V( + self, + key1, + key2, + model="empirical", + exact_norm=False, + pol=False, + time_index=None, + ): """ Calculates the covariance matrix for unnormed bandpowers (i.e., the q vectors). If the data were real and x_1 = x_2, the expression would be @@ -1690,22 +1985,32 @@ def get_unnormed_V(self, key1, key2, model='empirical', exact_norm=False, E_matrices = self.get_unnormed_E(key1, key2, exact_norm=exact_norm, pol=pol) C1 = self.C_model(key1, model=model, time_index=time_index) C2 = self.C_model(key2, model=model, time_index=time_index) - P21 = self.cross_covar_model(key2, key1, model=model, conj_1=False, - conj_2=False, time_index=time_index) - S21 = self.cross_covar_model(key2, key1, model=model, conj_1=True, - conj_2=True, time_index=time_index) - - E21C1 = np.dot(np.transpose(E_matrices.conj(), (0,2,1)), C1) + P21 = self.cross_covar_model( + key2, key1, model=model, conj_1=False, conj_2=False, time_index=time_index + ) + S21 = self.cross_covar_model( + key2, key1, model=model, conj_1=True, conj_2=True, time_index=time_index + ) + + E21C1 = np.dot(np.transpose(E_matrices.conj(), (0, 2, 1)), C1) E12C2 = np.dot(E_matrices, C2) - auto_term = np.einsum('aij,bji', E12C2, E21C1) + auto_term = np.einsum("aij,bji", E12C2, E21C1) E12starS21 = np.dot(E_matrices.conj(), S21) E12P21 = np.dot(E_matrices, P21) - cross_term = np.einsum('aij,bji', E12P21, E12starS21) + cross_term = np.einsum("aij,bji", E12P21, E12starS21) return auto_term + cross_term - def get_analytic_covariance(self, key1, key2, M=None, exact_norm=False, - pol=False, model='empirical', known_cov=None): + def get_analytic_covariance( + self, + key1, + key2, + M=None, + exact_norm=False, + pol=False, + model="empirical", + known_cov=None, + ): """ Calculates the auto-covariance matrix for both the real and imaginary parts of bandpowers (i.e., the q vectors and the p vectors). @@ -1869,35 +2174,60 @@ def get_analytic_covariance(self, key1, key2, M=None, exact_norm=False, # E_matrices has a shape of (spw_Ndlys, spw_Nfreqs, spw_Nfreqs) # using numpy.einsum_path to speed up the array products with numpy.einsum - einstein_path_0 = np.einsum_path('bij, cji->bc', E_matrices, E_matrices, optimize='optimal')[0] - einstein_path_1 = np.einsum_path('bi, ci,i->bc', E_matrices[:,:,0], E_matrices[:,:,0],E_matrices[0,:,0], optimize='optimal')[0] - einstein_path_2 = np.einsum_path('ab,cd,bd->ac', M[0], M[0], M[0], optimize='optimal')[0] + einstein_path_0 = np.einsum_path( + "bij, cji->bc", E_matrices, E_matrices, optimize="optimal" + )[0] + einstein_path_1 = np.einsum_path( + "bi, ci,i->bc", + E_matrices[:, :, 0], + E_matrices[:, :, 0], + E_matrices[0, :, 0], + optimize="optimal", + )[0] + einstein_path_2 = np.einsum_path( + "ab,cd,bd->ac", M[0], M[0], M[0], optimize="optimal" + )[0] # check if the covariance matrix is uniform along the time axis. If so, we just calculate the result for one timestamp and duplicate its copies # along the time axis. check_uniform_input = False - if model != 'foreground_dependent': - # When model is 'foreground_dependent', since we are processing the outer products of visibilities from different times, - # we are expected to have time-dependent inputs, thus check_uniform_input is always set to be False here. - C11_first = self.C_model(key1, model=model, known_cov=known_cov, time_index=0) - C11_last = self.C_model(key1, model=model, known_cov=known_cov, time_index=self.dsets[0].Ntimes-1) + if model != "foreground_dependent": + # When model is 'foreground_dependent', since we are processing the outer products of visibilities from different times, + # we are expected to have time-dependent inputs, thus check_uniform_input is always set to be False here. + C11_first = self.C_model( + key1, model=model, known_cov=known_cov, time_index=0 + ) + C11_last = self.C_model( + key1, + model=model, + known_cov=known_cov, + time_index=self.dsets[0].Ntimes - 1, + ) if np.isclose(C11_first, C11_last).all(): check_uniform_input = True cov_q_real, cov_q_imag, cov_p_real, cov_p_imag = [], [], [], [] for time_index in range(self.dsets[0].Ntimes): - if model in ['dsets','autos']: + if model in ["dsets", "autos"]: # calculate - = tr[ E^{12,a} C^{22} E^{21,b} C^{11} ] # We have used tr[A D_1 B D_2] = \sum_{ijkm} A_{ij} d_{1j} \delta_{jk} B_{km} d_{2m} \delta_{mi} = \sum_{ik} [A_{ik}*d_{1k}] * [B_{ki}*d_{2i}] # to simplify the computation. - C11 = self.C_model(key1, model=model, known_cov=known_cov, time_index=time_index) - C22 = self.C_model(key2, model=model, known_cov=known_cov, time_index=time_index) - E21C11 = np.multiply(np.transpose(E_matrices.conj(), (0,2,1)), np.diag(C11)) + C11 = self.C_model( + key1, model=model, known_cov=known_cov, time_index=time_index + ) + C22 = self.C_model( + key2, model=model, known_cov=known_cov, time_index=time_index + ) + E21C11 = np.multiply( + np.transpose(E_matrices.conj(), (0, 2, 1)), np.diag(C11) + ) E12C22 = np.multiply(E_matrices, np.diag(C22)) # Get q_q, q_qdagger, qdagger_qdagger - q_q, qdagger_qdagger = 0.+1.j*0, 0.+1.j*0 - q_qdagger = np.einsum('bij, cji->bc', E12C22, E21C11, optimize=einstein_path_0) - elif model == 'foreground_dependent': + q_q, qdagger_qdagger = 0.0 + 1.0j * 0, 0.0 + 1.0j * 0 + q_qdagger = np.einsum( + "bij, cji->bc", E12C22, E21C11, optimize=einstein_path_0 + ) + elif model == "foreground_dependent": # calculate tr[ E^{12,b} Cautos^{22} E^{21,c} Cautos^{11} + # E^{12,b} Cs E^{21,c} Cautos^{11} + # E^{12,b} Cautos^{22} E^{21,c} Cs ], @@ -1906,112 +2236,266 @@ def get_analytic_covariance(self, key1, key2, M=None, exact_norm=False, # we have used tr[A u u*^t B D_2] = \sum_{ijkm} A_{ij} u_j u*_k B_{km} D_{2mi} \\ # = \sum_{i} [ \sum_j A_{ij} u_j ] * [\sum_k u*_k B_{ki} ] * d_{2i} # to simplify the computation. - C11_autos = self.C_model(key1, model='autos', known_cov=known_cov, time_index=time_index) - C22_autos = self.C_model(key2, model='autos', known_cov=known_cov, time_index=time_index) - E21C11_autos = np.multiply(np.transpose(E_matrices.conj(), (0,2,1)), np.diag(C11_autos)) + C11_autos = self.C_model( + key1, model="autos", known_cov=known_cov, time_index=time_index + ) + C22_autos = self.C_model( + key2, model="autos", known_cov=known_cov, time_index=time_index + ) + E21C11_autos = np.multiply( + np.transpose(E_matrices.conj(), (0, 2, 1)), np.diag(C11_autos) + ) E12C22_autos = np.multiply(E_matrices, np.diag(C22_autos)) # Get q_q, q_qdagger, qdagger_qdagger - q_q, qdagger_qdagger = 0.+1.j*0, 0.+1.j*0 - q_qdagger = np.einsum('bij, cji->bc', E12C22_autos, E21C11_autos, optimize=einstein_path_0) - x1 = self.w(key1)[:,time_index] * self.x(key1)[:,time_index] - x2 = self.w(key2)[:,time_index] * self.x(key2)[:,time_index] + q_q, qdagger_qdagger = 0.0 + 1.0j * 0, 0.0 + 1.0j * 0 + q_qdagger = np.einsum( + "bij, cji->bc", E12C22_autos, E21C11_autos, optimize=einstein_path_0 + ) + x1 = self.w(key1)[:, time_index] * self.x(key1)[:, time_index] + x2 = self.w(key2)[:, time_index] * self.x(key2)[:, time_index] E12_x1 = np.dot(E_matrices, x1) E12_x2 = np.dot(E_matrices, x2) x2star_E21 = E12_x2.conj() x1star_E21 = E12_x1.conj() - x1star_E12 = np.dot(np.transpose(E_matrices,(0,2,1)), x1.conj()) - x2star_E12 = np.dot(np.transpose(E_matrices,(0,2,1)), x2.conj()) + x1star_E12 = np.dot(np.transpose(E_matrices, (0, 2, 1)), x1.conj()) + x2star_E12 = np.dot(np.transpose(E_matrices, (0, 2, 1)), x2.conj()) E21_x1 = x1star_E12.conj() E21_x2 = x2star_E12.conj() - SN_cov = np.einsum('bi,ci,i->bc', E12_x1, x2star_E21, np.diag(C11_autos), optimize=einstein_path_1)/2. + np.einsum('bi,ci,i->bc', E12_x2, x1star_E21, np.diag(C11_autos), optimize=einstein_path_1)/2.\ - + np.einsum('bi,ci,i->bc', x2star_E12, E21_x1, np.diag(C22_autos), optimize=einstein_path_1)/2. + np.einsum('bi,ci,i->bc', x1star_E12, E21_x2, np.diag(C22_autos), optimize=einstein_path_1)/2. + SN_cov = ( + np.einsum( + "bi,ci,i->bc", + E12_x1, + x2star_E21, + np.diag(C11_autos), + optimize=einstein_path_1, + ) + / 2.0 + + np.einsum( + "bi,ci,i->bc", + E12_x2, + x1star_E21, + np.diag(C11_autos), + optimize=einstein_path_1, + ) + / 2.0 + + np.einsum( + "bi,ci,i->bc", + x2star_E12, + E21_x1, + np.diag(C22_autos), + optimize=einstein_path_1, + ) + / 2.0 + + np.einsum( + "bi,ci,i->bc", + x1star_E12, + E21_x2, + np.diag(C22_autos), + optimize=einstein_path_1, + ) + / 2.0 + ) # Apply zero clipping on the columns and rows containing negative diagonal elements - SN_cov[np.real(np.diag(SN_cov))<=0., :] = 0. + 1.j*0 - SN_cov[:, np.real(np.diag(SN_cov))<=0.,] = 0. + 1.j*0 + SN_cov[np.real(np.diag(SN_cov)) <= 0.0, :] = 0.0 + 1.0j * 0 + SN_cov[:, np.real(np.diag(SN_cov)) <= 0.0,] = ( + 0.0 + 1.0j * 0 + ) q_qdagger += SN_cov else: # for general case (which is the slowest without simplification) - C11 = self.C_model(key1, model=model, known_cov=known_cov, time_index=time_index) - C22 = self.C_model(key2, model=model, known_cov=known_cov, time_index=time_index) - C21 = self.cross_covar_model(key2, key1, model=model, conj_1=False, conj_2=True, known_cov=known_cov, time_index=time_index) - C12 = self.cross_covar_model(key1, key2, model=model, conj_1=False, conj_2=True, known_cov=known_cov, time_index=time_index) - P11 = self.cross_covar_model(key1, key1, model=model, conj_1=False, conj_2=False, known_cov=known_cov, time_index=time_index) - S11 = self.cross_covar_model(key1, key1, model=model, conj_1=True, conj_2=True, known_cov=known_cov, time_index=time_index) - P22 = self.cross_covar_model(key2, key2, model=model, conj_1=False, conj_2=False, known_cov=known_cov, time_index=time_index) - S22 = self.cross_covar_model(key2, key2, model=model, conj_1=True, conj_2=True, known_cov=known_cov, time_index=time_index) - P21 = self.cross_covar_model(key2, key1, model=model, conj_1=False, conj_2=False, known_cov=known_cov, time_index=time_index) - S21 = self.cross_covar_model(key2, key1, model=model, conj_1=True, conj_2=True, known_cov=known_cov, time_index=time_index) + C11 = self.C_model( + key1, model=model, known_cov=known_cov, time_index=time_index + ) + C22 = self.C_model( + key2, model=model, known_cov=known_cov, time_index=time_index + ) + C21 = self.cross_covar_model( + key2, + key1, + model=model, + conj_1=False, + conj_2=True, + known_cov=known_cov, + time_index=time_index, + ) + C12 = self.cross_covar_model( + key1, + key2, + model=model, + conj_1=False, + conj_2=True, + known_cov=known_cov, + time_index=time_index, + ) + P11 = self.cross_covar_model( + key1, + key1, + model=model, + conj_1=False, + conj_2=False, + known_cov=known_cov, + time_index=time_index, + ) + S11 = self.cross_covar_model( + key1, + key1, + model=model, + conj_1=True, + conj_2=True, + known_cov=known_cov, + time_index=time_index, + ) + P22 = self.cross_covar_model( + key2, + key2, + model=model, + conj_1=False, + conj_2=False, + known_cov=known_cov, + time_index=time_index, + ) + S22 = self.cross_covar_model( + key2, + key2, + model=model, + conj_1=True, + conj_2=True, + known_cov=known_cov, + time_index=time_index, + ) + P21 = self.cross_covar_model( + key2, + key1, + model=model, + conj_1=False, + conj_2=False, + known_cov=known_cov, + time_index=time_index, + ) + S21 = self.cross_covar_model( + key2, + key1, + model=model, + conj_1=True, + conj_2=True, + known_cov=known_cov, + time_index=time_index, + ) # Get q_q, q_qdagger, qdagger_qdagger - if np.isclose(P22, 0).all() or np.isclose(S11,0).all(): - q_q = 0.+1.j*0 + if np.isclose(P22, 0).all() or np.isclose(S11, 0).all(): + q_q = 0.0 + 1.0j * 0 else: E12P22 = np.matmul(E_matrices, P22) - E21starS11 = np.matmul(np.transpose(E_matrices, (0,2,1)), S11) - q_q = np.einsum('bij, cji->bc', E12P22, E21starS11, optimize=einstein_path_0) + E21starS11 = np.matmul(np.transpose(E_matrices, (0, 2, 1)), S11) + q_q = np.einsum( + "bij, cji->bc", E12P22, E21starS11, optimize=einstein_path_0 + ) if np.isclose(C21, 0).all(): - q_q += 0.+1.j*0 + q_q += 0.0 + 1.0j * 0 else: E12C21 = np.matmul(E_matrices, C21) - q_q += np.einsum('bij, cji->bc', E12C21, E12C21, optimize=einstein_path_0) - E21C11 = np.matmul(np.transpose(E_matrices.conj(), (0,2,1)), C11) + q_q += np.einsum( + "bij, cji->bc", E12C21, E12C21, optimize=einstein_path_0 + ) + E21C11 = np.matmul(np.transpose(E_matrices.conj(), (0, 2, 1)), C11) E12C22 = np.matmul(E_matrices, C22) - q_qdagger = np.einsum('bij, cji->bc', E12C22, E21C11, optimize=einstein_path_0) - if np.isclose(P21, 0).all() or np.isclose(S21,0).all(): - q_qdagger += 0.+1.j*0 + q_qdagger = np.einsum( + "bij, cji->bc", E12C22, E21C11, optimize=einstein_path_0 + ) + if np.isclose(P21, 0).all() or np.isclose(S21, 0).all(): + q_qdagger += 0.0 + 1.0j * 0 else: E12P21 = np.matmul(E_matrices, P21) E12starS21 = np.matmul(E_matrices.conj(), S21) - q_qdagger += np.einsum('bij, cji->bc', E12P21, E12starS21, optimize=einstein_path_0) + q_qdagger += np.einsum( + "bij, cji->bc", E12P21, E12starS21, optimize=einstein_path_0 + ) if np.isclose(C12, 0).all(): - qdagger_qdagger = 0.+1.j*0 + qdagger_qdagger = 0.0 + 1.0j * 0 else: - E21C12 = np.matmul(np.transpose(E_matrices.conj(), (0,2,1)), C12) - qdagger_qdagger = np.einsum('bij, cji->bc', E21C12, E21C12, optimize=einstein_path_0) - if np.isclose(P11, 0).all() or np.isclose(S22,0).all(): - qdagger_qdagger += 0.+1.j*0 + E21C12 = np.matmul(np.transpose(E_matrices.conj(), (0, 2, 1)), C12) + qdagger_qdagger = np.einsum( + "bij, cji->bc", E21C12, E21C12, optimize=einstein_path_0 + ) + if np.isclose(P11, 0).all() or np.isclose(S22, 0).all(): + qdagger_qdagger += 0.0 + 1.0j * 0 else: - E21P11 = np.matmul(np.transpose(E_matrices.conj(), (0,2,1)), P11) + E21P11 = np.matmul(np.transpose(E_matrices.conj(), (0, 2, 1)), P11) E12starS22 = np.matmul(E_matrices.conj(), S22) - qdagger_qdagger += np.einsum('bij, cji->bc', E21P11, E12starS22, optimize=einstein_path_0) + qdagger_qdagger += np.einsum( + "bij, cji->bc", E21P11, E12starS22, optimize=einstein_path_0 + ) - cov_q_real_temp = (q_q + qdagger_qdagger + q_qdagger + q_qdagger.conj() ) / 4. - cov_q_imag_temp = -(q_q + qdagger_qdagger - q_qdagger - q_qdagger.conj() ) / 4. + cov_q_real_temp = ( + q_q + qdagger_qdagger + q_qdagger + q_qdagger.conj() + ) / 4.0 + cov_q_imag_temp = ( + -(q_q + qdagger_qdagger - q_qdagger - q_qdagger.conj()) / 4.0 + ) m = M[time_index] # calculate \sum_{bd} [ M_{ab} M_{cd} ( - ) ] if np.isclose([q_q], 0).all(): - MMq_q = np.zeros((E_matrices.shape[0],E_matrices.shape[0])).astype(np.complex128) + MMq_q = np.zeros((E_matrices.shape[0], E_matrices.shape[0])).astype( + np.complex128 + ) else: - assert np.shape(q_q) == np.shape(m), "covariance matrix and normalization matrix has different shapes." - MMq_q = np.einsum('ab,cd,bd->ac', m, m, q_q, optimize=einstein_path_2) + assert np.shape(q_q) == np.shape( + m + ), "covariance matrix and normalization matrix has different shapes." + MMq_q = np.einsum("ab,cd,bd->ac", m, m, q_q, optimize=einstein_path_2) # calculate \sum_{bd} [ M_{ab} M_{cd}^* ( - ) ] # and \sum_{bd} [ M_{ab}^* M_{cd} ( - ) ] if np.isclose([q_qdagger], 0).all(): - MM_q_qdagger = 0.+1.j*0 - M_Mq_qdagger_ = 0.+1.j*0 + MM_q_qdagger = 0.0 + 1.0j * 0 + M_Mq_qdagger_ = 0.0 + 1.0j * 0 else: - assert np.shape(q_qdagger) == np.shape(m), "covariance matrix and normalization matrix has different shapes." - MM_q_qdagger = np.einsum('ab,cd,bd->ac', m, m.conj(), q_qdagger, optimize=einstein_path_2) - M_Mq_qdagger_ = np.einsum('ab,cd,bd->ac', m.conj(), m, q_qdagger.conj(), optimize=einstein_path_2) + assert np.shape(q_qdagger) == np.shape( + m + ), "covariance matrix and normalization matrix has different shapes." + MM_q_qdagger = np.einsum( + "ab,cd,bd->ac", m, m.conj(), q_qdagger, optimize=einstein_path_2 + ) + M_Mq_qdagger_ = np.einsum( + "ab,cd,bd->ac", + m.conj(), + m, + q_qdagger.conj(), + optimize=einstein_path_2, + ) # calculate \sum_{bd} [ M_{ab}^* M_{cd}^* ( - ) ] if np.isclose([qdagger_qdagger], 0).all(): - M_M_qdagger_qdagger = 0.+1.j*0 + M_M_qdagger_qdagger = 0.0 + 1.0j * 0 else: - assert np.shape(qdagger_qdagger) == np.shape(m), "covariance matrix and normalization matrix has different shapes." - M_M_qdagger_qdagger = np.einsum('ab,cd,bd->ac', m.conj(), m.conj(), qdagger_qdagger, optimize=einstein_path_2) - - cov_p_real_temp = ( MMq_q + MM_q_qdagger + M_Mq_qdagger_ + M_M_qdagger_qdagger)/ 4. - cov_p_imag_temp = -( MMq_q - MM_q_qdagger - M_Mq_qdagger_ + M_M_qdagger_qdagger)/ 4. + assert np.shape(qdagger_qdagger) == np.shape( + m + ), "covariance matrix and normalization matrix has different shapes." + M_M_qdagger_qdagger = np.einsum( + "ab,cd,bd->ac", + m.conj(), + m.conj(), + qdagger_qdagger, + optimize=einstein_path_2, + ) + + cov_p_real_temp = ( + MMq_q + MM_q_qdagger + M_Mq_qdagger_ + M_M_qdagger_qdagger + ) / 4.0 + cov_p_imag_temp = ( + -(MMq_q - MM_q_qdagger - M_Mq_qdagger_ + M_M_qdagger_qdagger) / 4.0 + ) # cov_p_real_temp has a shaoe of (spw_Ndlys, spw_Ndlys) if check_uniform_input: - # if the covariance matrix is uniform along the time axis, we just calculate the result for one timestamp and duplicate its copies - # along the time axis. - cov_q_real.extend([cov_q_real_temp]*self.dsets[0].Ntimes) - cov_q_imag.extend([cov_q_imag_temp]*self.dsets[0].Ntimes) - cov_p_real.extend([cov_p_real_temp]*self.dsets[0].Ntimes) - cov_p_imag.extend([cov_p_imag_temp]*self.dsets[0].Ntimes) - warnings.warn("Producing time-uniform covariance matrices between bandpowers.") + # if the covariance matrix is uniform along the time axis, we just calculate the result for one timestamp and duplicate its copies + # along the time axis. + cov_q_real.extend([cov_q_real_temp] * self.dsets[0].Ntimes) + cov_q_imag.extend([cov_q_imag_temp] * self.dsets[0].Ntimes) + cov_p_real.extend([cov_p_real_temp] * self.dsets[0].Ntimes) + cov_p_imag.extend([cov_p_imag_temp] * self.dsets[0].Ntimes) + warnings.warn( + "Producing time-uniform covariance matrices between bandpowers." + ) break else: cov_q_real.append(cov_q_real_temp) @@ -2027,8 +2511,7 @@ def get_analytic_covariance(self, key1, key2, M=None, exact_norm=False, return cov_q_real, cov_q_imag, cov_p_real, cov_p_imag - - def get_MW(self, G, H, mode='I', band_covar=None, exact_norm=False, rcond=1e-15): + def get_MW(self, G, H, mode="I", band_covar=None, exact_norm=False, rcond=1e-15): """ Construct the normalization matrix M and window function matrix W for the power spectrum estimator. These are defined through Eqs. 14-16 of @@ -2098,33 +2581,36 @@ def get_MW(self, G, H, mode='I', band_covar=None, exact_norm=False, rcond=1e-15) # return M, W # Check that mode is supported - modes = ['H^-1', 'V^-1/2', 'I', 'L^-1'] - assert (mode in modes) + modes = ["H^-1", "V^-1/2", "I", "L^-1"] + assert mode in modes - if mode != 'I' and exact_norm is True: + if mode != "I" and exact_norm is True: raise NotImplementedError("Exact norm is not supported for non-I modes") # Build M matrix according to specified mode - if mode == 'H^-1': + if mode == "H^-1": try: M = np.linalg.inv(H) except np.linalg.LinAlgError as err: - if 'Singular matrix' in str(err): + if "Singular matrix" in str(err): M = np.linalg.pinv(H, rcond=rcond) - raise_warning("Warning: Window function matrix is singular " - "and cannot be inverted, so using " - " pseudoinverse instead.") + raise_warning( + "Warning: Window function matrix is singular " + "and cannot be inverted, so using " + " pseudoinverse instead." + ) else: - raise np.linalg.LinAlgError("Linear algebra error with H matrix " - "during MW computation.") + raise np.linalg.LinAlgError( + "Linear algebra error with H matrix " "during MW computation." + ) W = np.dot(M, H) W_norm = np.sum(W, axis=1) W = (W.T / W_norm).T - elif mode == 'V^-1/2': + elif mode == "V^-1/2": if np.sum(band_covar) == None: raise ValueError("Covariance not supplied for V^-1/2 normalization") # First find the eigenvectors and eigenvalues of the unnormalizd covariance @@ -2132,27 +2618,33 @@ def get_MW(self, G, H, mode='I', band_covar=None, exact_norm=False, rcond=1e-15) eigvals, eigvects = np.linalg.eigh(band_covar) nonpos_eigvals = eigvals <= 1e-20 if (nonpos_eigvals).any(): - raise_warning("At least one non-positive eigenvalue for the " - "unnormed bandpower covariance matrix.") + raise_warning( + "At least one non-positive eigenvalue for the " + "unnormed bandpower covariance matrix." + ) # truncate them eigvals = eigvals[~nonpos_eigvals] eigvects = eigvects[:, ~nonpos_eigvals] - V_minus_half = np.dot(eigvects, np.dot(np.diag(1./np.sqrt(eigvals)), eigvects.T)) + V_minus_half = np.dot( + eigvects, np.dot(np.diag(1.0 / np.sqrt(eigvals)), eigvects.T) + ) - W_norm = np.diag(1. / np.sum(np.dot(V_minus_half, H), axis=1)) + W_norm = np.diag(1.0 / np.sum(np.dot(V_minus_half, H), axis=1)) M = np.dot(W_norm, V_minus_half) W = np.dot(M, H) - elif mode == 'I': + elif mode == "I": # This is not the M matrix as is rigorously defined in the # OQE formalism, because the power spectrum scalar is excluded # in this matrix normalization (i.e., M doesn't do the full # normalization) - M = np.diag(1. / np.sum(G, axis=1)) - W_norm = np.diag(1. / np.sum(H, axis=1)) + M = np.diag(1.0 / np.sum(G, axis=1)) + W_norm = np.diag(1.0 / np.sum(H, axis=1)) W = np.dot(W_norm, H) else: - raise NotImplementedError("Cholesky decomposition mode not currently supported.") + raise NotImplementedError( + "Cholesky decomposition mode not currently supported." + ) # # Cholesky decomposition # order = np.arange(G.shape[0]) - np.ceil((G.shape[0]-1.)/2.) # order[order < 0] = order[order < 0] - 0.1 @@ -2218,28 +2710,30 @@ def get_Q_alt(self, mode, allow_fft=True, include_extension=False): self.set_Ndlys() if mode >= self.spw_Ndlys: - raise IndexError("Cannot compute Q matrix for a mode outside" - "of allowed range of delay modes.") + raise IndexError( + "Cannot compute Q matrix for a mode outside" + "of allowed range of delay modes." + ) nfreq = self.spw_Nfreqs if include_extension: nfreq = nfreq + np.sum(self.filter_extension) phase_correction = self.filter_extension[0] else: - phase_correction = 0. + phase_correction = 0.0 if (self.spw_Ndlys == nfreq) and (allow_fft == True): _m = np.zeros((nfreq,), dtype=np.complex) - _m[mode] = 1. # delta function at specific delay mode + _m[mode] = 1.0 # delta function at specific delay mode # FFT to transform to frequency space m = np.fft.fft(np.fft.ifftshift(_m)) else: if self.spw_Ndlys % 2 == 0: - start_idx = -self.spw_Ndlys/2 + start_idx = -self.spw_Ndlys / 2 else: - start_idx = -(self.spw_Ndlys - 1)/2 + start_idx = -(self.spw_Ndlys - 1) / 2 m = (start_idx + mode) * (np.arange(nfreq) - phase_correction) m = np.exp(-2j * np.pi * m / self.spw_Ndlys) - Q_alt = np.einsum('i,j', m.conj(), m) # dot it with its conjugate + Q_alt = np.einsum("i,j", m.conj(), m) # dot it with its conjugate return Q_alt def get_integral_beam(self, pol=False): @@ -2260,24 +2754,27 @@ def get_integral_beam(self, pol=False): integral_beam : array_like integral containing the spectral beam and tapering. """ - nu = self.freqs[self.spw_range[0]:self.spw_range[1]] # in Hz + nu = self.freqs[self.spw_range[0] : self.spw_range[1]] # in Hz try: # Get beam response in (frequency, pixel), beam area(freq) and # Nside, used in computing dtheta - beam_res, beam_omega, N = \ - self.primary_beam.beam_normalized_response(pol, nu) - prod = 1. / beam_omega + beam_res, beam_omega, N = self.primary_beam.beam_normalized_response( + pol, nu + ) + prod = 1.0 / beam_omega beam_prod = beam_res * prod[:, np.newaxis] # beam_prod has omega subsumed, but taper is still part of R matrix # The nside term is dtheta^2, where dtheta is the resolution in # healpix map - integral_beam = np.pi/(3.*N*N) * np.dot(beam_prod, beam_prod.T) + integral_beam = np.pi / (3.0 * N * N) * np.dot(beam_prod, beam_prod.T) - except(AttributeError): - warnings.warn("The beam response could not be calculated. " - "PS will not be normalized!") + except (AttributeError): + warnings.warn( + "The beam response could not be calculated. " + "PS will not be normalized!" + ) integral_beam = np.ones((len(nu), len(nu))) return integral_beam @@ -2306,14 +2803,16 @@ def get_Q(self, mode): if self.spw_Ndlys == None: self.set_Ndlys() if mode >= self.spw_Ndlys: - raise IndexError("Cannot compute Q matrix for a mode outside" - "of allowed range of delay modes.") + raise IndexError( + "Cannot compute Q matrix for a mode outside" + "of allowed range of delay modes." + ) - tau = self.delays()[int(mode)] * 1.0e-9 # delay in seconds - nu = self.freqs[self.spw_range[0]:self.spw_range[1]] # in Hz + tau = self.delays()[int(mode)] * 1.0e-9 # delay in seconds + nu = self.freqs[self.spw_range[0] : self.spw_range[1]] # in Hz - eta_int = np.exp(-2j * np.pi * tau * nu) # exponential part - Q_alt = np.einsum('i,j', eta_int.conj(), eta_int) # dot with conjugate + eta_int = np.exp(-2j * np.pi * tau * nu) # exponential part + Q_alt = np.einsum("i,j", eta_int.conj(), eta_int) # dot with conjugate return Q_alt def p_hat(self, M, q): @@ -2351,11 +2850,10 @@ def cov_p_hat(self, M, q_cov): """ p_cov = np.zeros_like(q_cov) for tnum in range(len(p_cov)): - p_cov[tnum] = np.einsum('ab,cd,bd->ac', M, M, q_cov[tnum]) + p_cov[tnum] = np.einsum("ab,cd,bd->ac", M, M, q_cov[tnum]) return p_cov - def broadcast_dset_flags(self, spw_ranges=None, time_thresh=0.2, - unflag=False): + def broadcast_dset_flags(self, spw_ranges=None, time_thresh=0.2, unflag=False): """ For each dataset in self.dset, update the flag_array such that the flagging patterns are time-independent for each baseline given @@ -2399,8 +2897,9 @@ def broadcast_dset_flags(self, spw_ranges=None, time_thresh=0.2, # spw type check if spw_ranges is None: spw_ranges = [(0, self.Nfreqs)] - assert isinstance(spw_ranges, list), \ - "spw_ranges must be fed as a list of tuples" + assert isinstance( + spw_ranges, list + ), "spw_ranges must be fed as a list of tuples" # iterate over datasets for dset in self.dsets: @@ -2410,7 +2909,9 @@ def broadcast_dset_flags(self, spw_ranges=None, time_thresh=0.2, # unflag if unflag: # unflag for all times - dset.flag_array[:,:,self.spw_range[0]:self.spw_range[1],:] = False + dset.flag_array[ + :, :, self.spw_range[0] : self.spw_range[1], : + ] = False continue # enact time threshold on flag waterfalls # iterate over polarizations @@ -2428,14 +2929,27 @@ def broadcast_dset_flags(self, spw_ranges=None, time_thresh=0.2, freq_contig_flgs = np.sum(flags, axis=1) / Nfreqs > 0.999999 Ntimes_noncontig = np.sum(~freq_contig_flgs, dtype=np.float) # get freq channels where non-contiguous flags exceed threshold - exceeds_thresh = np.sum(flags[~freq_contig_flgs], axis=0, dtype=np.float) / Ntimes_noncontig > time_thresh + exceeds_thresh = ( + np.sum(flags[~freq_contig_flgs], axis=0, dtype=np.float) + / Ntimes_noncontig + > time_thresh + ) # flag channels for all times that exceed time_thresh - dset.flag_array[bl_inds, :, np.where(exceeds_thresh)[0][:, None], i] = True + dset.flag_array[ + bl_inds, :, np.where(exceeds_thresh)[0][:, None], i + ] = True # for pixels that have flags but didn't meet broadcasting limit # flag the integration within the spw flags[:, np.where(exceeds_thresh)[0]] = False - flag_ints = np.max(flags[:, self.spw_range[0]:self.spw_range[1]], axis=1) - dset.flag_array[bl_inds[flag_ints], :, self.spw_range[0]:self.spw_range[1], i] = True + flag_ints = np.max( + flags[:, self.spw_range[0] : self.spw_range[1]], axis=1 + ) + dset.flag_array[ + bl_inds[flag_ints], + :, + self.spw_range[0] : self.spw_range[1], + i, + ] = True def units(self, little_h=True): """ @@ -2455,8 +2969,10 @@ def units(self, little_h=True): """ # Work out the power spectrum units if len(self.dsets) == 0: - raise IndexError("No datasets have been added yet; cannot " - "calculate power spectrum units.") + raise IndexError( + "No datasets have been added yet; cannot " + "calculate power spectrum units." + ) # get visibility units vis_units = self.dsets[0].vis_units @@ -2486,14 +3002,27 @@ def delays(self): """ # Calculate the delays if len(self.dsets) == 0: - raise IndexError("No datasets have been added yet; cannot " - "calculate delays.") + raise IndexError( + "No datasets have been added yet; cannot " "calculate delays." + ) else: - return utils.get_delays(self.freqs[self.spw_range[0]:self.spw_range[1]], - n_dlys=self.spw_Ndlys) * 1e9 # convert to ns - - def scalar(self, polpair, little_h=True, num_steps=2000, beam=None, - taper_override='no_override', exact_norm=False): + return ( + utils.get_delays( + self.freqs[self.spw_range[0] : self.spw_range[1]], + n_dlys=self.spw_Ndlys, + ) + * 1e9 + ) # convert to ns + + def scalar( + self, + polpair, + little_h=True, + num_steps=2000, + beam=None, + taper_override="no_override", + exact_norm=False, + ): """ Computes the scalar function to convert a power spectrum estimate in "telescope units" to cosmological units, using self.spw_range to set @@ -2548,17 +3077,18 @@ def scalar(self, polpair, little_h=True, num_steps=2000, beam=None, polpair = (polpair, polpair) if polpair[0] != polpair[1]: raise NotImplementedError( - "Polarizations don't match. Beam scalar can only be " - "calculated for auto-polarization pairs at the moment.") + "Polarizations don't match. Beam scalar can only be " + "calculated for auto-polarization pairs at the moment." + ) pol = polpair[0] # set spw_range and get freqs - freqs = self.freqs[self.spw_range[0]:self.spw_range[1]] + freqs = self.freqs[self.spw_range[0] : self.spw_range[1]] start = freqs[0] end = freqs[0] + np.median(np.diff(freqs)) * len(freqs) # Override the taper if desired - if taper_override == 'no_override': + if taper_override == "no_override": taper = self.taper else: taper = taper_override @@ -2566,18 +3096,31 @@ def scalar(self, polpair, little_h=True, num_steps=2000, beam=None, # calculate scalar if beam is None: scalar = self.primary_beam.compute_pspec_scalar( - start, end, len(freqs), pol=pol, - taper=self.taper, little_h=little_h, - num_steps=num_steps, exact_norm=exact_norm) + start, + end, + len(freqs), + pol=pol, + taper=self.taper, + little_h=little_h, + num_steps=num_steps, + exact_norm=exact_norm, + ) else: - scalar = beam.compute_pspec_scalar(start, end, len(freqs), - pol=pol, taper=self.taper, - little_h=little_h, - num_steps=num_steps, exact_norm=exact_norm) + scalar = beam.compute_pspec_scalar( + start, + end, + len(freqs), + pol=pol, + taper=self.taper, + little_h=little_h, + num_steps=num_steps, + exact_norm=exact_norm, + ) return scalar - def scalar_delay_adjustment(self, key1=None, key2=None, sampling=False, - Gv=None, Hv=None): + def scalar_delay_adjustment( + self, key1=None, key2=None, sampling=False, Gv=None, Hv=None + ): """ Computes an adjustment factor for the pspec scalar. There are two reasons why this might be needed: @@ -2630,8 +3173,10 @@ def scalar_delay_adjustment(self, key1=None, key2=None, sampling=False, adjustment : float if the data_weighting is 'identity' 1d array of floats with length spw_Ndlys otherwise. """ - if Gv is None: Gv = self.get_G(key1, key2) - if Hv is None: Hv = self.get_H(key1, key2, sampling) + if Gv is None: + Gv = self.get_G(key1, key2) + if Hv is None: + Hv = self.get_H(key1, key2, sampling) # get ratio summed_G = np.sum(Gv, axis=1) @@ -2645,21 +3190,21 @@ def scalar_delay_adjustment(self, key1=None, key2=None, sampling=False, ## XXX: Adjustments like this are hacky and wouldn't be necessary ## if we deprecate the incorrectly normalized ## Q and M matrix definitions. - #In the future, we need to do our normalizations properly and - #stop introducing arbitrary normalization factors. - #if the input identity weighting is diagonal, then the - #adjustment factor is independent of alpha. + # In the future, we need to do our normalizations properly and + # stop introducing arbitrary normalization factors. + # if the input identity weighting is diagonal, then the + # adjustment factor is independent of alpha. # get mean ratio. - if self.data_weighting == 'identity': + if self.data_weighting == "identity": mean_ratio = np.mean(ratio) scatter = np.abs(ratio - mean_ratio) if (scatter > 10**-4 * mean_ratio).any(): raise ValueError("The normalization scalar is band-dependent!") adjustment = self.spw_Ndlys / (self.spw_Nfreqs * mean_ratio) - #otherwise, the adjustment factor is dependent on alpha. + # otherwise, the adjustment factor is dependent on alpha. else: adjustment = self.spw_Ndlys / (self.spw_Nfreqs * ratio) - if self.taper != 'none': + if self.taper != "none": tapering_fct = dspec.gen_window(self.taper, self.spw_Nfreqs) adjustment *= np.mean(tapering_fct**2) @@ -2696,14 +3241,20 @@ def validate_pol(self, dsets, pol_pair): # convert elements to integers if fed as strings if isinstance(pol_pair[0], (str, np.str)): - pol_pair = (uvutils.polstr2num(pol_pair[0], x_orientation=x_orientation), pol_pair[1]) + pol_pair = ( + uvutils.polstr2num(pol_pair[0], x_orientation=x_orientation), + pol_pair[1], + ) if isinstance(pol_pair[1], (str, np.str)): - pol_pair = (pol_pair[0], uvutils.polstr2num(pol_pair[1], x_orientation=x_orientation)) + pol_pair = ( + pol_pair[0], + uvutils.polstr2num(pol_pair[1], x_orientation=x_orientation), + ) assert isinstance(pol_pair[0], (int, np.integer)), err_msg assert isinstance(pol_pair[1], (int, np.integer)), err_msg - #if pol_pair[0] != pol_pair[1]: + # if pol_pair[0] != pol_pair[1]: # raise NotImplementedError("Only auto/equal polarizations are implement at the moment.") dset_ind1 = self.dset_idx(dsets[0]) @@ -2713,23 +3264,53 @@ def validate_pol(self, dsets, pol_pair): valid = True if pol_pair[0] not in dset1.polarization_array: - print("dset {} does not contain data for polarization {}".format(dset_ind1, pol_pair[0])) + print( + "dset {} does not contain data for polarization {}".format( + dset_ind1, pol_pair[0] + ) + ) valid = False if pol_pair[1] not in dset2.polarization_array: - print("dset {} does not contain data for polarization {}".format(dset_ind2, pol_pair[1])) + print( + "dset {} does not contain data for polarization {}".format( + dset_ind2, pol_pair[1] + ) + ) valid = False return valid - def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, - input_data_weight='identity', norm='I', taper='none', - sampling=False, little_h=True, spw_ranges=None, symmetric_taper=True, - baseline_tol=1.0, store_cov=False, store_cov_diag=False, - return_q=False, store_window=True, exact_windows=False, - ftbeam_file=None, verbose=True, filter_extensions=None, - exact_norm=False, history='', r_params=None, - cov_model='empirical', known_cov=None, allow_fft=False): + def pspec( + self, + bls1, + bls2, + dsets, + pols, + n_dlys=None, + input_data_weight="identity", + norm="I", + taper="none", + sampling=False, + little_h=True, + spw_ranges=None, + symmetric_taper=True, + baseline_tol=1.0, + store_cov=False, + store_cov_diag=False, + return_q=False, + store_window=True, + exact_windows=False, + ftbeam_file=None, + verbose=True, + filter_extensions=None, + exact_norm=False, + history="", + r_params=None, + cov_model="empirical", + known_cov=None, + allow_fft=False, + ): """ Estimate the delay power spectrum from a pair of datasets contained in this object, using the optimal quadratic estimator of arXiv:1502.06016. @@ -2979,49 +3560,56 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, # Currently the "pspec normalization scalar" doesn't work if a # non-identity data weighting AND a non-trivial taper are used - if taper != 'none' and input_data_weight != 'identity': - raise_warning("Warning: Scalar power spectrum normalization " - "doesn't work with current implementation " - "if the tapering AND non-identity " - "weighting matrices are both used.", - verbose=verbose) + if taper != "none" and input_data_weight != "identity": + raise_warning( + "Warning: Scalar power spectrum normalization " + "doesn't work with current implementation " + "if the tapering AND non-identity " + "weighting matrices are both used.", + verbose=verbose, + ) # get datasets - assert isinstance(dsets, (list, tuple)), \ - "dsets must be fed as length-2 tuple of integers" + assert isinstance( + dsets, (list, tuple) + ), "dsets must be fed as length-2 tuple of integers" assert len(dsets) == 2, "len(dsets) must be 2" - assert isinstance(dsets[0], (int, np.integer)) \ - and isinstance(dsets[1], (int, np.integer)), \ - "dsets must contain integer indices" + assert isinstance(dsets[0], (int, np.integer)) and isinstance( + dsets[1], (int, np.integer) + ), "dsets must contain integer indices" dset1 = self.dsets[self.dset_idx(dsets[0])] dset2 = self.dsets[self.dset_idx(dsets[1])] # assert form of bls1 and bls2 - assert isinstance(bls1, list), \ - "bls1 and bls2 must be fed as a list of antpair tuples" - assert isinstance(bls2, list), \ - "bls1 and bls2 must be fed as a list of antpair tuples" - assert len(bls1) == len(bls2) and len(bls1) > 0, \ - "length of bls1 must equal length of bls2 and be > 0" + assert isinstance( + bls1, list + ), "bls1 and bls2 must be fed as a list of antpair tuples" + assert isinstance( + bls2, list + ), "bls1 and bls2 must be fed as a list of antpair tuples" + assert ( + len(bls1) == len(bls2) and len(bls1) > 0 + ), "length of bls1 must equal length of bls2 and be > 0" for i in range(len(bls1)): if isinstance(bls1[i], tuple): - assert isinstance(bls2[i], tuple), \ - "bls1[{}] type must match bls2[{}] type".format(i, i) + assert isinstance( + bls2[i], tuple + ), "bls1[{}] type must match bls2[{}] type".format(i, i) else: - assert len(bls1[i]) == len(bls2[i]), \ - "len(bls1[{}]) must match len(bls2[{}])".format(i, i) + assert len(bls1[i]) == len( + bls2[i] + ), "len(bls1[{}]) must match len(bls2[{}])".format(i, i) # construct list of baseline pairs bl_pairs = [] for i in range(len(bls1)): if isinstance(bls1[i], tuple): - bl_pairs.append( (bls1[i], bls2[i]) ) + bl_pairs.append((bls1[i], bls2[i])) elif isinstance(bls1[i], list) and len(bls1[i]) == 1: - bl_pairs.append( (bls1[i][0], bls2[i][0]) ) + bl_pairs.append((bls1[i][0], bls2[i][0])) else: - bl_pairs.append( - [ (bls1[i][j], bls2[i][j]) for j in range(len(bls1[i])) ] ) + bl_pairs.append([(bls1[i][j], bls2[i][j]) for j in range(len(bls1[i]))]) # validate bl-pair redundancy validate_blpairs(bl_pairs, dset1, dset2, baseline_tol=baseline_tol) @@ -3030,19 +3618,26 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, if spw_ranges is None: spw_ranges = [(0, self.Nfreqs)] if isinstance(spw_ranges, tuple): - spw_ranges = [spw_ranges,] + spw_ranges = [ + spw_ranges, + ] if filter_extensions is None: filter_extensions = [(0, 0) for m in range(len(spw_ranges))] # convert to list if only a tuple was given if isinstance(filter_extensions, tuple): - filter_extensions = [filter_extensions,] + filter_extensions = [ + filter_extensions, + ] - assert len(spw_ranges) == len(filter_extensions), "must provide same number of spw_ranges as filter_extensions" + assert len(spw_ranges) == len( + filter_extensions + ), "must provide same number of spw_ranges as filter_extensions" # Check that spw_ranges is list of len-2 tuples - assert np.isclose([len(t) for t in spw_ranges], 2).all(), \ - "spw_ranges must be fed as a list of length-2 tuples" + assert np.isclose( + [len(t) for t in spw_ranges], 2 + ).all(), "spw_ranges must be fed as a list of length-2 tuples" # if using default setting of number of delay bins equal to number # of frequency channels @@ -3054,13 +3649,15 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, # if using the whole band in the dataset, then there should just be # one n_dly parameter specified if spw_ranges is None and n_dlys != None: - assert len(n_dlys) == 1, \ - "Only one spw, so cannot specify more than one n_dly value" + assert ( + len(n_dlys) == 1 + ), "Only one spw, so cannot specify more than one n_dly value" # assert that the same number of ndlys has been specified as the # number of spws - assert len(spw_ranges) == len(n_dlys), \ - "Need to specify number of delay bins for each spw" + assert len(spw_ranges) == len( + n_dlys + ), "Need to specify number of delay bins for each spw" if store_cov_diag and store_cov: store_cov = False @@ -3068,23 +3665,32 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, # no matter what the initial choice for store_cov. if exact_windows and not store_window: - warnings.warn('exact_windows is True... setting store_window to True.') + warnings.warn("exact_windows is True... setting store_window to True.") store_window = True # setup polarization selection - if isinstance(pols, (tuple, str)): pols = [pols] + if isinstance(pols, (tuple, str)): + pols = [pols] # convert all polarizations to integers if fed as strings _pols = [] for p in pols: if isinstance(p, str): # Convert string to pol-integer pair - p = (uvutils.polstr2num(p, x_orientation=self.dsets[0].x_orientation), - uvutils.polstr2num(p, x_orientation=self.dsets[0].x_orientation)) + p = ( + uvutils.polstr2num(p, x_orientation=self.dsets[0].x_orientation), + uvutils.polstr2num(p, x_orientation=self.dsets[0].x_orientation), + ) if isinstance(p[0], (str, np.str)): - p = (uvutils.polstr2num(p[0], x_orientation=self.dsets[0].x_orientation), p[1]) + p = ( + uvutils.polstr2num(p[0], x_orientation=self.dsets[0].x_orientation), + p[1], + ) if isinstance(p[1], (str, np.str)): - p = (p[0], uvutils.polstr2num(p[1], x_orientation=self.dsets[0].x_orientation)) + p = ( + p[0], + uvutils.polstr2num(p[1], x_orientation=self.dsets[0].x_orientation), + ) _pols.append(p) pols = _pols @@ -3111,7 +3717,7 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, for i in range(len(spw_ranges)): # set spectral range if verbose: - print( "\nSetting spectral range: {}".format(spw_ranges[i])) + print("\nSetting spectral range: {}".format(spw_ranges[i])) self.set_spw(spw_ranges[i], ndlys=n_dlys[i]) self.set_filter_extension(filter_extensions[i]) @@ -3130,7 +3736,7 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, spw_window_function = [] d = self.delays() * 1e-9 - f = dset1.freq_array.flatten()[spw_ranges[i][0]:spw_ranges[i][1]] + f = dset1.freq_array.flatten()[spw_ranges[i][0] : spw_ranges[i][1]] dlys.extend(d) dly_spws.extend(np.ones_like(d, np.int16) * i) freq_spws.extend(np.ones_like(f, np.int16) * i) @@ -3139,17 +3745,20 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, # Loop over polarizations for j, p in enumerate(pols): p_str = tuple([uvutils.polnum2str(_p) for _p in p]) - if verbose: print( "\nUsing polarization pair: {}".format(p_str)) + if verbose: + print("\nUsing polarization pair: {}".format(p_str)) # validating polarization pair on UVData objects valid = self.validate_pol(dsets, tuple(p)) if not valid: - # Polarization pair is invalid; skip - print("Polarization pair: {} failed the validation test, " - "continuing...".format(p_str)) - continue + # Polarization pair is invalid; skip + print( + "Polarization pair: {} failed the validation test, " + "continuing...".format(p_str) + ) + continue - spw_polpair.append( uvputils.polpair_tuple2int(p) ) + spw_polpair.append(uvputils.polpair_tuple2int(p)) pol_data = [] pol_wgts = [] pol_ints = [] @@ -3162,31 +3771,40 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, if self.primary_beam is not None: # Raise error if cross-pol is requested - if (p[0] != p[1]): + if p[0] != p[1]: raise NotImplementedError( "Visibilities with different polarizations can only " "be cross-correlated if primary_beam = None. Cannot " - "compute beam scalar for mixed polarizations.") + "compute beam scalar for mixed polarizations." + ) # using zero'th indexed polarization, as cross-polarized # beams are not yet implemented - if norm == 'H^-1': + if norm == "H^-1": # If using decorrelation, the H^-1 normalization # already deals with the taper, so we need to override # the taper when computing the scalar - scalar = self.scalar(p, little_h=little_h, - taper_override='none', - exact_norm=exact_norm) + scalar = self.scalar( + p, + little_h=little_h, + taper_override="none", + exact_norm=exact_norm, + ) else: - scalar = self.scalar(p, little_h=little_h, - exact_norm=exact_norm) + scalar = self.scalar( + p, little_h=little_h, exact_norm=exact_norm + ) else: - raise_warning("Warning: self.primary_beam is not defined, " - "so pspectra are not properly normalized", - verbose=verbose) + raise_warning( + "Warning: self.primary_beam is not defined, " + "so pspectra are not properly normalized", + verbose=verbose, + ) scalar = 1.0 - pol = (p[0]) # used in get_integral_beam function to specify the correct polarization for the beam + pol = p[ + 0 + ] # used in get_integral_beam function to specify the correct polarization for the beam spw_scalar.append(scalar) # Loop over baseline pairs @@ -3194,14 +3812,16 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, # assign keys if isinstance(blp, list): # interpet blp as group of baseline-pairs - raise NotImplementedError("Baseline lists bls1 and bls2" - " must be lists of tuples (not lists of lists" - " of tuples).\n" - "Use hera_pspec.pspecdata.construct_blpairs()" - " to construct appropriately grouped baseline" - " lists.") - #key1 = [(dsets[0],) + _blp[0] + (p[0],) for _blp in blp] - #key2 = [(dsets[1],) + _blp[1] + (p[1],) for _blp in blp] + raise NotImplementedError( + "Baseline lists bls1 and bls2" + " must be lists of tuples (not lists of lists" + " of tuples).\n" + "Use hera_pspec.pspecdata.construct_blpairs()" + " to construct appropriately grouped baseline" + " lists." + ) + # key1 = [(dsets[0],) + _blp[0] + (p[0],) for _blp in blp] + # key2 = [(dsets[1],) + _blp[1] + (p[1],) for _blp in blp] elif isinstance(blp, tuple): # interpret blp as baseline-pair key1 = (dsets[0],) + blp[0] + (p_str[0],) @@ -3213,36 +3833,46 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, # Check that number of non-zero weight chans >= n_dlys key1_dof = np.sum(~np.isclose(self.Y(key1).diagonal(), 0.0)) key2_dof = np.sum(~np.isclose(self.Y(key2).diagonal(), 0.0)) - if key1_dof - np.sum(self.filter_extension) < self.spw_Ndlys\ - or key2_dof - np.sum(self.filter_extension) < self.spw_Ndlys: + if ( + key1_dof - np.sum(self.filter_extension) < self.spw_Ndlys + or key2_dof - np.sum(self.filter_extension) < self.spw_Ndlys + ): if verbose: - print("WARNING: Number of unflagged chans for key1 " - "and/or key2 < n_dlys\n which may lead to " - "normalization instabilities.") - #if using inverse sinc weighting, set r_params - if input_data_weight == 'dayenu': + print( + "WARNING: Number of unflagged chans for key1 " + "and/or key2 < n_dlys\n which may lead to " + "normalization instabilities." + ) + # if using inverse sinc weighting, set r_params + if input_data_weight == "dayenu": key1 = (dsets[0],) + blp[0] + (p_str[0],) key2 = (dsets[1],) + blp[1] + (p_str[1],) if not key1 in r_params: - raise ValueError("No r_param dictionary supplied" - " for baseline %s"%(str(key1))) + raise ValueError( + "No r_param dictionary supplied" + " for baseline %s" % (str(key1)) + ) if not key2 in r_params: - raise ValueError("No r_param dictionary supplied" - " for baseline %s"%(str(key2))) + raise ValueError( + "No r_param dictionary supplied" + " for baseline %s" % (str(key2)) + ) self.set_r_param(key1, r_params[key1]) self.set_r_param(key2, r_params[key2]) # Build Fisher matrix - if input_data_weight == 'identity': + if input_data_weight == "identity": # in this case, all Gv and Hv differ only by flagging pattern # so check if we've already computed this # First: get flag weighting matrices given key1 & key2 - Y = np.vstack([self.Y(key1).diagonal(), - self.Y(key2).diagonal()]) + Y = np.vstack( + [self.Y(key1).diagonal(), self.Y(key2).diagonal()] + ) # Second: check cache for Y - matches = [np.isclose(Y, y).all() - for y in self._identity_Y.values()] + matches = [ + np.isclose(Y, y).all() for y in self._identity_Y.values() + ] if True in matches: # This Y exists, so pick appropriate G and H and continue match = list(self._identity_Y.keys())[matches.index(True)] @@ -3250,9 +3880,16 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, Hv = self._identity_H[match] else: # This Y doesn't exist, so compute it - if verbose: print(" Building G...") - Gv = self.get_G(key1, key2, exact_norm=exact_norm, pol = pol) - Hv = self.get_H(key1, key2, sampling=sampling, exact_norm=exact_norm, pol = pol) + if verbose: + print(" Building G...") + Gv = self.get_G(key1, key2, exact_norm=exact_norm, pol=pol) + Hv = self.get_H( + key1, + key2, + sampling=sampling, + exact_norm=exact_norm, + pol=pol, + ) # cache it self._identity_Y[(key1, key2)] = Y self._identity_G[(key1, key2)] = Gv @@ -3260,76 +3897,121 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, else: # for non identity weighting (i.e. iC weighting) # Gv and Hv are always different, so compute them - if verbose: print(" Building G...") - Gv = self.get_G(key1, key2, exact_norm=exact_norm, pol = pol) - Hv = self.get_H(key1, key2, sampling=sampling, exact_norm=exact_norm, pol = pol) + if verbose: + print(" Building G...") + Gv = self.get_G(key1, key2, exact_norm=exact_norm, pol=pol) + Hv = self.get_H( + key1, + key2, + sampling=sampling, + exact_norm=exact_norm, + pol=pol, + ) # Calculate unnormalized bandpowers - if verbose: print(" Building q_hat...") - qv = self.q_hat(key1, key2, exact_norm=exact_norm, pol=pol, allow_fft=allow_fft) + if verbose: + print(" Building q_hat...") + qv = self.q_hat( + key1, key2, exact_norm=exact_norm, pol=pol, allow_fft=allow_fft + ) - if verbose: print(" Normalizing power spectrum...") - if norm == 'V^-1/2': - V_mat = self.get_unnormed_V(key1, key2, exact_norm=exact_norm, pol = pol) - Mv, Wv = self.get_MW(Gv, Hv, mode=norm, band_covar=V_mat, exact_norm=exact_norm) + if verbose: + print(" Normalizing power spectrum...") + if norm == "V^-1/2": + V_mat = self.get_unnormed_V( + key1, key2, exact_norm=exact_norm, pol=pol + ) + Mv, Wv = self.get_MW( + Gv, Hv, mode=norm, band_covar=V_mat, exact_norm=exact_norm + ) else: Mv, Wv = self.get_MW(Gv, Hv, mode=norm, exact_norm=exact_norm) pv = self.p_hat(Mv, qv) # Multiply by scalar if self.primary_beam != None: - if verbose: print(" Computing and multiplying scalar...") + if verbose: + print(" Computing and multiplying scalar...") pv *= scalar # Wide bin adjustment of scalar, which is only needed for # the diagonal norm matrix mode (i.e., norm = 'I') - if norm == 'I' and not(exact_norm): + if norm == "I" and not (exact_norm): sa = self.scalar_delay_adjustment(Gv=Gv, Hv=Hv) if isinstance(sa, (np.float, float)): pv *= sa else: pv = np.atleast_2d(sa).T * pv - #Generate the covariance matrix if error bars provided + # Generate the covariance matrix if error bars provided if store_cov or store_cov_diag: - if verbose: print(" Building q_hat covariance...") - cov_q_real, cov_q_imag, cov_real, cov_imag \ - = self.get_analytic_covariance(key1, key2, Mv, - exact_norm=exact_norm, - pol=pol, - model=cov_model, - known_cov=known_cov, ) + if verbose: + print(" Building q_hat covariance...") + ( + cov_q_real, + cov_q_imag, + cov_real, + cov_imag, + ) = self.get_analytic_covariance( + key1, + key2, + Mv, + exact_norm=exact_norm, + pol=pol, + model=cov_model, + known_cov=known_cov, + ) if self.primary_beam != None: - cov_real = cov_real * (scalar)**2. - cov_imag = cov_imag * (scalar)**2. + cov_real = cov_real * (scalar) ** 2.0 + cov_imag = cov_imag * (scalar) ** 2.0 - if norm == 'I' and not(exact_norm): + if norm == "I" and not (exact_norm): if isinstance(sa, (np.float, float)): - cov_real = cov_real * (sa)**2. - cov_imag = cov_imag * (sa)**2. + cov_real = cov_real * (sa) ** 2.0 + cov_imag = cov_imag * (sa) ** 2.0 else: cov_real = cov_real * np.outer(sa, sa)[None] cov_imag = cov_imag * np.outer(sa, sa)[None] if not return_q: if store_cov: - pol_cov_real.extend(np.real(cov_real).astype(np.float64)) - pol_cov_imag.extend(np.real(cov_imag).astype(np.float64)) + pol_cov_real.extend( + np.real(cov_real).astype(np.float64) + ) + pol_cov_imag.extend( + np.real(cov_imag).astype(np.float64) + ) if store_cov_diag: - stats = np.sqrt(np.diagonal(np.real(cov_real), axis1=1, axis2=2)) + 1.j*np.sqrt(np.diagonal(np.real(cov_imag), axis1=1, axis2=2)) + stats = np.sqrt( + np.diagonal(np.real(cov_real), axis1=1, axis2=2) + ) + 1.0j * np.sqrt( + np.diagonal(np.real(cov_imag), axis1=1, axis2=2) + ) pol_stats_array_cov_model.extend(stats) else: if store_cov: - pol_cov_real.extend(np.real(cov_q_real).astype(np.float64)) - pol_cov_imag.extend(np.real(cov_q_imag).astype(np.float64)) + pol_cov_real.extend( + np.real(cov_q_real).astype(np.float64) + ) + pol_cov_imag.extend( + np.real(cov_q_imag).astype(np.float64) + ) if store_cov_diag: - stats = np.sqrt(np.diagonal(np.real(cov_q_real), axis1=1, axis2=2)) + 1.j*np.sqrt(np.diagonal(np.real(cov_q_imag), axis1=1, axis2=2)) + stats = np.sqrt( + np.diagonal(np.real(cov_q_real), axis1=1, axis2=2) + ) + 1.0j * np.sqrt( + np.diagonal(np.real(cov_q_imag), axis1=1, axis2=2) + ) pol_stats_array_cov_model.extend(stats) # store the window_function if store_window: - pol_window_function.extend(np.repeat(Wv[np.newaxis,:,:], qv.shape[1], axis=0).astype(np.float64)) + pol_window_function.extend( + np.repeat(Wv[np.newaxis, :, :], qv.shape[1], axis=0).astype( + np.float64 + ) + ) # Get baseline keys if isinstance(blp, list): @@ -3353,10 +4035,16 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, wgts2 = self.w(key2).T # get avg of nsample across frequency axis, weighted by wgts - nsamp1 = np.sum(dset1.get_nsamples(bl1 + (p[0],))[:, slice(*self.get_spw())] * wgts1, axis=1) \ - / np.sum(wgts1, axis=1).clip(1, np.inf) - nsamp2 = np.sum(dset2.get_nsamples(bl2 + (p[1],))[:, slice(*self.get_spw())] * wgts2, axis=1) \ - / np.sum(wgts2, axis=1).clip(1, np.inf) + nsamp1 = np.sum( + dset1.get_nsamples(bl1 + (p[0],))[:, slice(*self.get_spw())] + * wgts1, + axis=1, + ) / np.sum(wgts1, axis=1).clip(1, np.inf) + nsamp2 = np.sum( + dset2.get_nsamples(bl2 + (p[1],))[:, slice(*self.get_spw())] + * wgts2, + axis=1, + ) / np.sum(wgts2, axis=1).clip(1, np.inf) # get integ1 blts1 = dset1.antpair2ind(bl1, ordered=False) @@ -3369,11 +4057,12 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, # take inverse avg of integ1 and integ2 to get total integ # inverse avg is done b/c integ ~ 1/noise_var # and due to non-linear operation of V_1 * V_2 - pol_ints.extend(1./np.mean([1./integ1, 1./integ2], axis=0)) + pol_ints.extend(1.0 / np.mean([1.0 / integ1, 1.0 / integ2], axis=0)) # combined weight is geometric mean - pol_wgts.extend(np.concatenate([wgts1[:, :, None], - wgts2[:, :, None]], axis=2)) + pol_wgts.extend( + np.concatenate([wgts1[:, :, None], wgts2[:, :, None]], axis=2) + ) # insert time and blpair info only once per blpair if i < 1 and j < 1: @@ -3386,8 +4075,10 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, lst2.extend(dset2.lst_array[inds2]) # insert blpair info - blp_arr.extend(np.ones_like(inds1, np.int) \ - * uvputils._antnums_to_blpair(blp)) + blp_arr.extend( + np.ones_like(inds1, np.int) + * uvputils._antnums_to_blpair(blp) + ) # insert into data and wgts integrations dictionaries spw_data.append(pol_data) @@ -3402,7 +4093,9 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, spw_data = np.moveaxis(np.array(spw_data), 0, -1) spw_wgts = np.moveaxis(np.array(spw_wgts), 0, -1) spw_ints = np.moveaxis(np.array(spw_ints), 0, -1) - spw_stats_array_cov_model = np.moveaxis(np.array(spw_stats_array_cov_model), 0, -1) + spw_stats_array_cov_model = np.moveaxis( + np.array(spw_stats_array_cov_model), 0, -1 + ) if store_cov: spw_cov_real = np.moveaxis(np.array(spw_cov_real), 0, -1) spw_cov_imag = np.moveaxis(np.array(spw_cov_imag), 0, -1) @@ -3422,23 +4115,25 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, # raise error if none of pols are consistent with the UVData objects if len(spw_polpair) == 0: - raise ValueError("None of the specified polarization pairs " - "match that of the UVData objects") + raise ValueError( + "None of the specified polarization pairs " + "match that of the UVData objects" + ) self.set_filter_extension((0, 0)) # set filter_extension to be zero when ending the loop # fill uvp object uvp = uvpspec.UVPSpec() - uvp.symmetric_taper=symmetric_taper + uvp.symmetric_taper = symmetric_taper # fill meta-data uvp.time_1_array = np.array(time1) uvp.time_2_array = np.array(time2) uvp.time_avg_array = np.mean([uvp.time_1_array, uvp.time_2_array], axis=0) uvp.lst_1_array = np.array(lst1) uvp.lst_2_array = np.array(lst2) - uvp.lst_avg_array = np.mean([np.unwrap(uvp.lst_1_array), - np.unwrap(uvp.lst_2_array)], axis=0) \ - % (2*np.pi) + uvp.lst_avg_array = np.mean( + [np.unwrap(uvp.lst_1_array), np.unwrap(uvp.lst_2_array)], axis=0 + ) % (2 * np.pi) uvp.blpair_array = np.array(blp_arr) uvp.Nblpairs = len(np.unique(blp_arr)) uvp.Ntimes = len(np.unique(time1)) @@ -3466,39 +4161,54 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, uvp.weighting = input_data_weight uvp.vis_units, uvp.norm_units = self.units(little_h=little_h) uvp.telescope_location = dset1.telescope_location - filename1 = json.loads(dset1.extra_keywords.get('filename', '""')) - cal1 = json.loads(dset1.extra_keywords.get('calibration', '""')) - filename2 = json.loads(dset2.extra_keywords.get('filename', '""')) - cal2 = json.loads(dset2.extra_keywords.get('calibration', '""')) + filename1 = json.loads(dset1.extra_keywords.get("filename", '""')) + cal1 = json.loads(dset1.extra_keywords.get("calibration", '""')) + filename2 = json.loads(dset2.extra_keywords.get("filename", '""')) + cal2 = json.loads(dset2.extra_keywords.get("calibration", '""')) label1 = self.labels[self.dset_idx(dsets[0])] label2 = self.labels[self.dset_idx(dsets[1])] uvp.labels = sorted(set([label1, label2])) - uvp.label_1_array = np.ones((uvp.Nspws, uvp.Nblpairts, uvp.Npols), np.int) \ - * uvp.labels.index(label1) - uvp.label_2_array = np.ones((uvp.Nspws, uvp.Nblpairts, uvp.Npols), np.int) \ - * uvp.labels.index(label2) + uvp.label_1_array = np.ones( + (uvp.Nspws, uvp.Nblpairts, uvp.Npols), np.int + ) * uvp.labels.index(label1) + uvp.label_2_array = np.ones( + (uvp.Nspws, uvp.Nblpairts, uvp.Npols), np.int + ) * uvp.labels.index(label2) uvp.labels = np.array(uvp.labels, np.str) - uvp.history = "UVPSpec written on {} with hera_pspec git hash {}\n{}\n" \ - "dataset1: filename: {}, label: {}, cal: {}, history:\n{}\n{}\n" \ - "dataset2: filename: {}, label: {}, cal: {}, history:\n{}\n{}\n" \ - "".format(datetime.datetime.utcnow(), version.git_hash, '-'*20, - filename1, label1, cal1, dset1.history, '-'*20, - filename2, label2, cal2, dset2.history, '-'*20) + uvp.history = ( + "UVPSpec written on {} with hera_pspec git hash {}\n{}\n" + "dataset1: filename: {}, label: {}, cal: {}, history:\n{}\n{}\n" + "dataset2: filename: {}, label: {}, cal: {}, history:\n{}\n{}\n" + "".format( + datetime.datetime.utcnow(), + version.git_hash, + "-" * 20, + filename1, + label1, + cal1, + dset1.history, + "-" * 20, + filename2, + label2, + cal2, + dset2.history, + "-" * 20, + ) + ) uvp.r_params = uvputils.compress_r_params(r_params) uvp.taper = taper if not return_q: uvp.norm = norm else: - uvp.norm = 'Unnormalized' + uvp.norm = "Unnormalized" if self.primary_beam is not None: # attach cosmology uvp.cosmo = self.primary_beam.cosmo # attach beam info uvp.beam_freqs = self.primary_beam.beam_freqs - uvp.OmegaP, uvp.OmegaPP = \ - self.primary_beam.get_Omegas(uvp.polpair_array) - if hasattr(self.primary_beam, 'filename'): + uvp.OmegaP, uvp.OmegaPP = self.primary_beam.get_Omegas(uvp.polpair_array) + if hasattr(self.primary_beam, "filename"): uvp.beamfile = self.primary_beam.filename # fill data arrays @@ -3506,8 +4216,11 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, uvp.integration_array = integration_array uvp.wgt_array = wgt_array uvp.nsample_array = dict( - [ (k, np.ones_like(uvp.integration_array[k], np.float)) - for k in uvp.integration_array.keys() ] ) + [ + (k, np.ones_like(uvp.integration_array[k], np.float)) + for k in uvp.integration_array.keys() + ] + ) # covariance if store_cov: @@ -3516,15 +4229,18 @@ def pspec(self, bls1, bls2, dsets, pols, n_dlys=None, uvp.cov_model = cov_model if store_cov_diag: uvp.stats_array = odict() - uvp.stats_array[cov_model+"_diag"] = stats_array_cov_model + uvp.stats_array[cov_model + "_diag"] = stats_array_cov_model # window functions if store_window: if exact_windows: # compute and store exact window functions - uvp.get_exact_window_functions(ftbeam_file=ftbeam_file, verbose=verbose, - x_orientation=self.dsets[0].x_orientation, - inplace=True) + uvp.get_exact_window_functions( + ftbeam_file=ftbeam_file, + verbose=verbose, + x_orientation=self.dsets[0].x_orientation, + inplace=True, + ) else: uvp.window_function_array = window_function_array @@ -3578,9 +4294,9 @@ def rephase_to_dset(self, dset_index=0, inplace=True): # get LST grid we are phasing to lst_grid = [] lst_array = dsets[dset_index].lst_array.ravel() - for l in lst_array: - if l not in lst_grid: - lst_grid.append(l) + for lst in lst_array: + if lst not in lst_grid: + lst_grid.append(lst) lst_grid = np.array(lst_grid) # get polarization list @@ -3592,20 +4308,21 @@ def rephase_to_dset(self, dset_index=0, inplace=True): if i == dset_index: # even though not phasing this dset, must set to match all other # dsets due to phasing-check validation - dset.phase_type = 'unknown' + dset.phase_type = "unknown" continue # skip if dataset is not drift phased - if dset.phase_type != 'drift': + if dset.phase_type != "drift": print("skipping dataset {} b/c it isn't drift phased".format(i)) # convert UVData to DataContainers. Note this doesn't make # a copy of the data - (data, flgs, antpos, ants, freqs, times, lsts, - pols) = hc.io.load_vis(dset, return_meta=True) + (data, flgs, antpos, ants, freqs, times, lsts, pols) = hc.io.load_vis( + dset, return_meta=True + ) # make bls dictionary - bls = dict([(k, antpos[k[0]] - antpos[k[1]]) for k in data.keys()]) + bls = {k: antpos[k[0]] - antpos[k[1]] for k in data.keys()} # Get dlst array dlst = lst_grid - lsts @@ -3617,19 +4334,21 @@ def rephase_to_dset(self, dset_index=0, inplace=True): hc.utils.lst_rephase(data, bls, freqs, dlst, lat=lat) # re-insert into dataset - for j, k in enumerate(data.keys()): + for k in data.keys(): # get blts indices of basline indices = dset.antpair2ind(k[:2], ordered=False) # get index in polarization_array for this polarization - polind = pol_list.index(uvutils.polstr2num(k[-1], x_orientation=self.dsets[0].x_orientation)) + polind = pol_list.index( + uvutils.polstr2num(k[-1], x_orientation=self.dsets[0].x_orientation) + ) # insert into dset dset.data_array[indices, 0, :, polind] = data[k] # set phasing in UVData object to unknown b/c there isn't a single # consistent phasing for the entire data set. - dset.phase_type = 'unknown' + dset.phase_type = "unknown" if inplace is False: return dsets @@ -3653,16 +4372,20 @@ def Jy_to_mK(self, beam=None): beam = self.primary_beam else: if self.primary_beam is not None: - print("Warning: feeding a beam model when self.primary_beam " - "already exists...") + print( + "Warning: feeding a beam model when self.primary_beam " + "already exists..." + ) # Check beam is not None - assert beam is not None, \ - "Cannot convert Jy --> mK b/c beam object is not defined..." + assert ( + beam is not None + ), "Cannot convert Jy --> mK b/c beam object is not defined..." # assert type of beam - assert isinstance(beam, pspecbeam.PSpecBeamBase), \ - "beam model must be a subclass of pspecbeam.PSpecBeamBase" + assert isinstance( + beam, pspecbeam.PSpecBeamBase + ), "beam model must be a subclass of pspecbeam.PSpecBeamBase" # iterate over all pols and get conversion factors factors = {} @@ -3672,12 +4395,16 @@ def Jy_to_mK(self, beam=None): # iterate over datasets and apply factor for i, dset in enumerate(self.dsets): # check dset vis units - if dset.vis_units.upper() != 'JY': - print("Cannot convert dset {} Jy -> mK because vis_units = {}".format(i, dset.vis_units)) + if dset.vis_units.upper() != "JY": + print( + "Cannot convert dset {} Jy -> mK because vis_units = {}".format( + i, dset.vis_units + ) + ) continue for j, p in enumerate(dset.polarization_array): dset.data_array[:, :, :, j] *= factors[p][None, None, :] - dset.vis_units = 'mK' + dset.vis_units = "mK" def trim_dset_lsts(self, lst_tol=6): """ @@ -3697,16 +4424,17 @@ def trim_dset_lsts(self, lst_tol=6): dlst = np.median(np.diff(np.unique(self.dsets[0].lst_array))) for dset in self.dsets: _dlst = np.median(np.diff(np.unique(dset.lst_array))) - if not np.isclose(dlst, _dlst, atol=10**(-lst_tol) / dset.Ntimes): - raise ValueError("Not all datasets in self.dsets are on the same LST " - "grid, cannot LST trim.") + if not np.isclose(dlst, _dlst, atol=10 ** (-lst_tol) / dset.Ntimes): + raise ValueError( + "Not all datasets in self.dsets are on the same LST " + "grid, cannot LST trim." + ) # get lst array of each dataset, turn into string and add to common_lsts lst_arrs = [] common_lsts = set() for i, dset in enumerate(self.dsets): - lsts = ["{lst:0.{tol}f}".format(lst=l, tol=lst_tol) - for l in dset.lst_array] + lsts = [f"{lst:0.{lst_tol}f}" for lst in dset.lst_array] lst_arrs.append(lsts) if i == 0: common_lsts = common_lsts.union(set(lsts)) @@ -3721,20 +4449,62 @@ def trim_dset_lsts(self, lst_tol=6): self.dsets[i].select(times=dset.time_array[~trim_inds]) -def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, - groupname=None, dset_labels=None, dset_pairs=None, psname_ext=None, - spw_ranges=None, n_dlys=None, pol_pairs=None, blpairs=None, - input_data_weight='identity', norm='I', taper='none', sampling=False, - exclude_auto_bls=False, exclude_cross_bls=False, exclude_permutations=True, - Nblps_per_group=None, bl_len_range=(0, 1e10), - bl_deg_range=(0, 180), bl_error_tol=1.0, - store_window=True, exact_windows=False, ftbeam_file=None, - beam=None, cosmo=None, interleave_times=False, rephase_to_dset=None, - trim_dset_lsts=False, broadcast_dset_flags=True, - time_thresh=0.2, Jy2mK=False, overwrite=True, symmetric_taper=True, - file_type='miriad', verbose=True, exact_norm=False, store_cov=False, store_cov_diag=False, filter_extensions=None, - history='', r_params=None, tsleep=0.1, maxiter=1, return_q=False, known_cov=None, cov_model='empirical', - include_autocorrs=False, include_crosscorrs=True, xant_flag_thresh=0.95, allow_fft=False): +def pspec_run( + dsets, + filename, + dsets_std=None, + cals=None, + cal_flag=True, + groupname=None, + dset_labels=None, + dset_pairs=None, + psname_ext=None, + spw_ranges=None, + n_dlys=None, + pol_pairs=None, + blpairs=None, + input_data_weight="identity", + norm="I", + taper="none", + sampling=False, + exclude_auto_bls=False, + exclude_cross_bls=False, + exclude_permutations=True, + Nblps_per_group=None, + bl_len_range=(0, 1e10), + bl_deg_range=(0, 180), + bl_error_tol=1.0, + store_window=True, + exact_windows=False, + ftbeam_file=None, + beam=None, + cosmo=None, + interleave_times=False, + rephase_to_dset=None, + trim_dset_lsts=False, + broadcast_dset_flags=True, + time_thresh=0.2, + Jy2mK=False, + overwrite=True, + symmetric_taper=True, + file_type="miriad", + verbose=True, + exact_norm=False, + store_cov=False, + store_cov_diag=False, + filter_extensions=None, + history="", + r_params=None, + tsleep=0.1, + maxiter=1, + return_q=False, + known_cov=None, + cov_model="empirical", + include_autocorrs=False, + include_crosscorrs=True, + xant_flag_thresh=0.95, + allow_fft=False, +): """ Create a PSpecData object, run OQE delay spectrum estimation and write results to a PSpecContainer object. @@ -3864,7 +4634,7 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, Definition of the beam Fourier transform to be used. Options include; - Root name of the file to use, without the polarisation - Ex : FT_beam_HERA_dipole (+ path) + Ex : FT_beam_HERA_dipole (+ path) - '' for computation from beam simulations (slow) beam : PSpecBeam object, UVBeam object or string @@ -4024,14 +4794,15 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, weighting matrices. """ # type check - assert isinstance(dsets, (list, tuple, np.ndarray)), \ - "dsets must be fed as a list of dataset string paths or UVData objects." + assert isinstance( + dsets, (list, tuple, np.ndarray) + ), "dsets must be fed as a list of dataset string paths or UVData objects." # parse psname if psname_ext is not None: assert isinstance(psname_ext, (str, np.str)) else: - psname_ext = '' + psname_ext = "" # polarizations check if pol_pairs is not None: @@ -4067,31 +4838,41 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, if dset_labels is None: dset_labels = ["dset{}".format(i) for i in range(Ndsets)] else: - assert not np.any(['_' in dl for dl in dset_labels]), \ - "cannot accept underscores in input dset_labels: {}".format(dset_labels) + assert not np.any( + ["_" in dl for dl in dset_labels] + ), "cannot accept underscores in input dset_labels: {}".format(dset_labels) # if dsets are not UVData, assume they are filepaths or list of filepaths if not isinstance(dsets[0], UVData): try: # load data into UVData objects if fed as list of strings t0 = time.time() - dsets = _load_dsets(dsets, bls=bls, pols=pols, file_type=file_type, verbose=verbose) - utils.log("Loaded data in %1.1f sec." % (time.time() - t0), - lvl=1, verbose=verbose) + dsets = _load_dsets( + dsets, bls=bls, pols=pols, file_type=file_type, verbose=verbose + ) + utils.log( + "Loaded data in %1.1f sec." % (time.time() - t0), lvl=1, verbose=verbose + ) except ValueError: # at least one of the dset loads failed due to no data being present - utils.log("One of the dset loads failed due to no data overlap given " - "the bls and pols selection", verbose=verbose) + utils.log( + "One of the dset loads failed due to no data overlap given " + "the bls and pols selection", + verbose=verbose, + ) return None - assert np.all([isinstance(d, UVData) for d in dsets]), \ - "dsets must be fed as a list of dataset string paths or UVData objects." + assert np.all( + [isinstance(d, UVData) for d in dsets] + ), "dsets must be fed as a list of dataset string paths or UVData objects." # check dsets_std input if dsets_std is not None: - err_msg = "input dsets_std must be a list of UVData objects or " \ - "filepaths to miriad files" - assert isinstance(dsets_std,(list, tuple, np.ndarray)), err_msg + err_msg = ( + "input dsets_std must be a list of UVData objects or " + "filepaths to miriad files" + ) + assert isinstance(dsets_std, (list, tuple, np.ndarray)), err_msg assert len(dsets_std) == Ndsets, "len(dsets_std) must equal len(dsets)" # load data if not UVData @@ -4099,14 +4880,22 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, try: # load data into UVData objects if fed as list of strings t0 = time.time() - dsets_std = _load_dsets(dsets_std, bls=bls, pols=pols, file_type=file_type, verbose=verbose) - utils.log("Loaded data in %1.1f sec." % (time.time() - t0), - lvl=1, verbose=verbose) + dsets_std = _load_dsets( + dsets_std, bls=bls, pols=pols, file_type=file_type, verbose=verbose + ) + utils.log( + "Loaded data in %1.1f sec." % (time.time() - t0), + lvl=1, + verbose=verbose, + ) except ValueError: # at least one of the dsets_std loads failed due to no data # being present - utils.log("One of the dsets_std loads failed due to no data overlap given " - "the bls and pols selection", verbose=verbose) + utils.log( + "One of the dsets_std loads failed due to no data overlap given " + "the bls and pols selection", + verbose=verbose, + ) return None assert np.all([isinstance(d, UVData) for d in dsets_std]), err_msg @@ -4118,8 +4907,11 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, if not isinstance(cals[0], UVCal): t0 = time.time() cals = _load_cals(cals, verbose=verbose) - utils.log("Loaded calibration in %1.1f sec." % (time.time() - t0), - lvl=1, verbose=verbose) + utils.log( + "Loaded calibration in %1.1f sec." % (time.time() - t0), + lvl=1, + verbose=verbose, + ) err_msg = "cals must be a list of UVCal, filepaths, or list of filepaths" assert np.all([isinstance(c, UVCal) for c in cals]), err_msg @@ -4141,8 +4933,15 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, beam.cosmo = cosmo # package into PSpecData - ds = PSpecData(dsets=dsets, wgts=[None for d in dsets], labels=dset_labels, - dsets_std=dsets_std, beam=beam, cals=cals, cal_flag=cal_flag) + ds = PSpecData( + dsets=dsets, + wgts=[None for d in dsets], + labels=dset_labels, + dsets_std=dsets_std, + beam=beam, + cals=cals, + cal_flag=cal_flag, + ) # erase calibration as they are no longer needed del cals @@ -4155,19 +4954,32 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, if interleave_times: if len(ds.dsets) != 1: raise ValueError("interleave_times only applicable for Ndsets == 1") - Ntimes = ds.dsets[0].Ntimes # get smallest Ntimes + Ntimes = ds.dsets[0].Ntimes # get smallest Ntimes Ntimes -= Ntimes % 2 # make it an even number # update dsets - ds.dsets.append(ds.dsets[0].select(times=np.unique(ds.dsets[0].time_array)[1:Ntimes:2], inplace=False)) - ds.dsets[0].select(times=np.unique(ds.dsets[0].time_array)[0:Ntimes:2], inplace=True) + ds.dsets.append( + ds.dsets[0].select( + times=np.unique(ds.dsets[0].time_array)[1:Ntimes:2], inplace=False + ) + ) + ds.dsets[0].select( + times=np.unique(ds.dsets[0].time_array)[0:Ntimes:2], inplace=True + ) ds.labels.append("dset1") # update dsets_std if ds.dsets_std[0] is None: ds.dsets_std.append(None) else: - ds.dsets_std.append(ds.dsets_std[0].select(times=np.unique(ds.dsets_std[0].time_array)[1:Ntimes:2], inplace=False)) - ds.dsets_std[0].select(times=np.unique(ds.dsets_std[0].time_array)[0:Ntimes:2], inplace=True) + ds.dsets_std.append( + ds.dsets_std[0].select( + times=np.unique(ds.dsets_std[0].time_array)[1:Ntimes:2], + inplace=False, + ) + ) + ds.dsets_std[0].select( + times=np.unique(ds.dsets_std[0].time_array)[0:Ntimes:2], inplace=True + ) # wgts is currently always None ds.wgts.append(None) @@ -4175,7 +4987,6 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, dset_pairs = [(0, 1)] dsets = ds.dsets dsets_std = ds.dsets_std - wgts = ds.wgts dset_labels = ds.labels # rephase if desired @@ -4195,13 +5006,15 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, # is already cross-correlating different times to avoid noise bias. # See issue #160 on hera_pspec repo if exclude_auto_bls: - raise_warning("Skipping the cross-multiplications of a baseline " - "with itself may cause a bias if one is already " - "cross-correlating different times to avoid the " - "noise bias. Please see hera_pspec github issue 160 " - "to make sure you know what you are doing! " - "https://github.com/HERA-Team/hera_pspec/issues/160", - verbose=verbose) + raise_warning( + "Skipping the cross-multiplications of a baseline " + "with itself may cause a bias if one is already " + "cross-correlating different times to avoid the " + "noise bias. Please see hera_pspec github issue 160 " + "to make sure you know what you are doing! " + "https://github.com/HERA-Team/hera_pspec/issues/160", + verbose=verbose, + ) # check dset pair type err_msg = "dset_pairs must be fed as a list of len-2 integer tuples" @@ -4210,23 +5023,24 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, # Get baseline-pairs to use for each dataset pair bls1_list, bls2_list = [], [] - for i, dsetp in enumerate(dset_pairs): + for dsetp in dset_pairs: # get bls if blpairs not fed if blpairs is None: - (bls1, bls2, blps, xants1, - xants2) = utils.calc_blpair_reds( - dsets[dsetp[0]], dsets[dsetp[1]], - filter_blpairs=True, - exclude_auto_bls=exclude_auto_bls, - exclude_cross_bls=exclude_cross_bls, - exclude_permutations=exclude_permutations, - Nblps_per_group=Nblps_per_group, - bl_len_range=bl_len_range, - bl_deg_range=bl_deg_range, - include_autocorrs=include_autocorrs, - include_crosscorrs=include_crosscorrs, - bl_tol=bl_error_tol, - xant_flag_thresh=xant_flag_thresh) + (bls1, bls2, blps, xants1, xants2) = utils.calc_blpair_reds( + dsets[dsetp[0]], + dsets[dsetp[1]], + filter_blpairs=True, + exclude_auto_bls=exclude_auto_bls, + exclude_cross_bls=exclude_cross_bls, + exclude_permutations=exclude_permutations, + Nblps_per_group=Nblps_per_group, + bl_len_range=bl_len_range, + bl_deg_range=bl_deg_range, + include_autocorrs=include_autocorrs, + include_crosscorrs=include_crosscorrs, + bl_tol=bl_error_tol, + xant_flag_thresh=xant_flag_thresh, + ) bls1_list.append(bls1) bls2_list.append(bls2) @@ -4237,8 +5051,9 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, _bls1 = [] _bls2 = [] for _bl1, _bl2 in zip(bls1, bls2): - if (_bl1 in dset1_bls or _bl1[::-1] in dset1_bls) \ - and (_bl2 in dset2_bls or _bl2[::-1] in dset2_bls): + if (_bl1 in dset1_bls or _bl1[::-1] in dset1_bls) and ( + _bl2 in dset2_bls or _bl2[::-1] in dset2_bls + ): _bls1.append(_bl1) _bls2.append(_bl2) @@ -4246,12 +5061,15 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, bls2_list.append(_bls2) # Open PSpecContainer to store all output in - if verbose: print("Opening {} in transactional mode".format(filename)) - psc = container.PSpecContainer(filename, mode='rw', keep_open=False, tsleep=tsleep, maxiter=maxiter) + if verbose: + print("Opening {} in transactional mode".format(filename)) + psc = container.PSpecContainer( + filename, mode="rw", keep_open=False, tsleep=tsleep, maxiter=maxiter + ) # assign group name if groupname is None: - groupname = '_'.join(dset_labels) + groupname = "_".join(dset_labels) # Loop over dataset combinations for i, dset_idxs in enumerate(dset_pairs): @@ -4259,23 +5077,42 @@ def pspec_run(dsets, filename, dsets_std=None, cals=None, cal_flag=True, if len(bls1_list[i]) == 0 or len(bls2_list[i]) == 0: continue # Run OQE - uvp = ds.pspec(bls1_list[i], bls2_list[i], dset_idxs, pol_pairs, symmetric_taper=symmetric_taper, - spw_ranges=spw_ranges, n_dlys=n_dlys, r_params=r_params, - store_cov=store_cov, store_cov_diag=store_cov_diag, input_data_weight=input_data_weight, - exact_norm=exact_norm, sampling=sampling, - return_q=return_q, cov_model=cov_model, known_cov=known_cov, - norm=norm, taper=taper, history=history, verbose=verbose, - filter_extensions=filter_extensions, store_window=store_window, - exact_windows=exact_windows, ftbeam_file=ftbeam_file) + uvp = ds.pspec( + bls1_list[i], + bls2_list[i], + dset_idxs, + pol_pairs, + symmetric_taper=symmetric_taper, + spw_ranges=spw_ranges, + n_dlys=n_dlys, + r_params=r_params, + store_cov=store_cov, + store_cov_diag=store_cov_diag, + input_data_weight=input_data_weight, + exact_norm=exact_norm, + sampling=sampling, + return_q=return_q, + cov_model=cov_model, + known_cov=known_cov, + norm=norm, + taper=taper, + history=history, + verbose=verbose, + filter_extensions=filter_extensions, + store_window=store_window, + exact_windows=exact_windows, + ftbeam_file=ftbeam_file, + ) # Store output - psname = '{}_x_{}{}'.format(dset_labels[dset_idxs[0]], - dset_labels[dset_idxs[1]], psname_ext) + psname = "{}_x_{}{}".format( + dset_labels[dset_idxs[0]], dset_labels[dset_idxs[1]], psname_ext + ) # write in transactional mode - if verbose: print("Storing {}".format(psname)) - psc.set_pspec(group=groupname, psname=psname, pspec=uvp, - overwrite=overwrite) + if verbose: + print("Storing {}".format(psname)) + psc.set_pspec(group=groupname, psname=psname, pspec=uvp, overwrite=overwrite) return ds @@ -4285,88 +5122,300 @@ def get_pspec_run_argparser(): def list_of_int_tuples(v): """Format for parsing lists of integer pairs for different OQE args. - Two acceptable formats are - Ex1: '0~0,1~1' --> [(0, 0), (1, 1), ...] and - Ex2: '0 0, 1 1' --> [(0, 0), (1, 1), ...]""" - if '~' in v: - v = [tuple([int(_x) for _x in x.split('~')]) for x in v.split(",")] + Two acceptable formats are + Ex1: '0~0,1~1' --> [(0, 0), (1, 1), ...] and + Ex2: '0 0, 1 1' --> [(0, 0), (1, 1), ...]""" + if "~" in v: + v = [tuple([int(_x) for _x in x.split("~")]) for x in v.split(",")] else: v = [tuple([int(_x) for _x in x.split()]) for x in v.split(",")] return v def list_of_str_tuples(v): """Lists of string 2-tuples for various OQE args (ex. Polarization pairs). - Two acceptable formats are - Ex1: 'xx~xx,yy~yy' --> [('xx', 'xx'), ('yy', 'yy'), ...] and - Ex2: 'xx xx, yy yy' --> [('xx', 'xx'), ('yy', 'yy'), ...]""" - if '~' in v: - v = [tuple([str(_x) for _x in x.split('~')]) for x in v.split(",")] + Two acceptable formats are + Ex1: 'xx~xx,yy~yy' --> [('xx', 'xx'), ('yy', 'yy'), ...] and + Ex2: 'xx xx, yy yy' --> [('xx', 'xx'), ('yy', 'yy'), ...]""" + if "~" in v: + v = [tuple([str(_x) for _x in x.split("~")]) for x in v.split(",")] else: v = [tuple([str(_x) for _x in x.split()]) for x in v.split(",")] return v def list_of_tuple_tuples(v): """List of tuple tuples for various OQE args (ex. baseline pair lists). Two acceptable formats are - Ex1: '1~2~3~4,5~6~7~8' --> [((1 2), (3, 4)), ((5, 6), (7, 8)), ...] and - Ex2: '1 2 3 4, 5 6 7 8' --> [((1 2), (3, 4)), ((5, 6), (7, 8)), ...])""" - if '~' in v: - v = [tuple([int(_x) for _x in x.split('~')]) for x in v.split(",")] + Ex1: '1~2~3~4,5~6~7~8' --> [((1 2), (3, 4)), ((5, 6), (7, 8)), ...] and + Ex2: '1 2 3 4, 5 6 7 8' --> [((1 2), (3, 4)), ((5, 6), (7, 8)), ...])""" + if "~" in v: + v = [tuple([int(_x) for _x in x.split("~")]) for x in v.split(",")] else: v = [tuple([int(_x) for _x in x.split()]) for x in v.split(",")] v = [(x[:2], x[2:]) for x in v] return v - a.add_argument("dsets", nargs='*', help="List of UVData objects or miriad filepaths.") + a.add_argument( + "dsets", nargs="*", help="List of UVData objects or miriad filepaths." + ) a.add_argument("filename", type=str, help="Output filename of HDF5 container.") - a.add_argument("--dsets_std", nargs='*', default=None, type=str, help="List of miriad filepaths to visibility standard deviations.") - a.add_argument("--groupname", default=None, type=str, help="Groupname for the UVPSpec objects in the HDF5 container.") - a.add_argument("--dset_pairs", default=None, type=list_of_int_tuples, help="List of dset pairings for OQE. Two acceptable formats are " - "Ex1: '0~0,1~1' --> [(0, 0), (1, 1), ...] and " - "Ex2: '0 0, 1 1' --> [(0, 0), (1, 1), ...]") - a.add_argument("--dset_labels", default=None, type=str, nargs='*', help="List of string labels for each input dataset.") - a.add_argument("--spw_ranges", default=None, type=list_of_int_tuples, help="List of spw channel selections. Two acceptable formats are " - "Ex1: '200~300,500~650' --> [(200, 300), (500, 650), ...] and " - "Ex2: '200 300, 500 650' --> [(200, 300), (500, 650), ...]") - a.add_argument("--n_dlys", default=None, type=int, nargs='+', help="List of integers specifying number of delays to use per spectral window selection.") - a.add_argument("--pol_pairs", default=None, type=list_of_str_tuples, help="List of pol-string pairs to use in OQE. Two acceptable formats are " - "Ex1: 'xx~xx,yy~yy' --> [('xx', 'xx'), ('yy', 'yy'), ...] and " - "Ex2: 'xx xx, yy yy' --> [('xx', 'xx'), ('yy', 'yy'), ...]") - a.add_argument("--blpairs", default=None, type=list_of_tuple_tuples, help="List of baseline-pair antenna integers to run OQE on. Two acceptable formats are " - "Ex1: '1~2~3~4,5~6~7~8' --> [((1 2), (3, 4)), ((5, 6), (7, 8)), ...] and " - "Ex2: '1 2 3 4, 5 6 7 8' --> [((1 2), (3, 4)), ((5, 6), (7, 8)), ...]") - a.add_argument("--input_data_weight", default='identity', type=str, help="Data weighting for OQE. See PSpecData.pspec for details.") - a.add_argument("--norm", default='I', type=str, help='M-matrix normalization type for OQE. See PSpecData.pspec for details.') - a.add_argument("--taper", default='none', type=str, help="Taper function to use in OQE delay transform. See PSpecData.pspec for details.") - a.add_argument("--beam", default=None, type=str, help="Filepath to UVBeam healpix map of antenna beam.") - a.add_argument("--cosmo", default=None, nargs='+', type=float, help="List of float values for [Om_L, Om_b, Om_c, H0, Om_M, Om_k].") - a.add_argument("--rephase_to_dset", default=None, type=int, help="dset integer index to phase all other dsets to. Default is no rephasing.") - a.add_argument("--trim_dset_lsts", default=False, action='store_true', help="Trim non-overlapping dset LSTs.") - a.add_argument("--broadcast_dset_flags", default=False, action='store_true', help="Broadcast dataset flags across time according to time_thresh.") - a.add_argument("--time_thresh", default=0.2, type=float, help="Fractional flagging threshold across time to trigger flag broadcast if broadcast_dset_flags is True") - a.add_argument("--Jy2mK", default=False, action='store_true', help="Convert datasets from Jy to mK if a beam model is provided.") - a.add_argument("--exclude_auto_bls", default=False, action='store_true', help='If blpairs is not provided, exclude all baselines paired with itself.') - a.add_argument("--exclude_cross_bls", default=False, action='store_true', help='If blpairs is not provided, exclude all baselines paired with a different baseline.') - a.add_argument("--exclude_permutations", default=False, action='store_true', help='If blpairs is not provided, exclude a basline-pair permutations. Ex: if (A, B) exists, exclude (B, A).') - a.add_argument("--Nblps_per_group", default=None, type=int, help="If blpairs is not provided and group == True, set the number of blpairs in each group.") - a.add_argument("--bl_len_range", default=(0, 1e10), nargs='+', type=float, help="If blpairs is not provided, limit the baselines used based on their minimum and maximum length in meters.") - a.add_argument("--bl_deg_range", default=(0, 180), nargs='+', type=float, help="If blpairs is not provided, limit the baseline used based on a min and max angle cut in ENU frame in degrees.") - a.add_argument("--bl_error_tol", default=1.0, type=float, help="If blpairs is not provided, this is the error tolerance in forming redundant baseline groups in meters.") - a.add_argument("--store_cov", default=False, action='store_true', help="Compute and store covariance of bandpowers given dsets_std files or empirical covariance.") - a.add_argument("--store_cov_diag", default=False, action='store_true', help="Compute and store the error bars calculated by QE formalism.") - a.add_argument("--return_q", default=False, action='store_true', help="Return unnormalized bandpowers given dsets files.") - a.add_argument("--overwrite", default=False, action='store_true', help="Overwrite output if it exists.") - a.add_argument("--cov_model", default='empirical', type=str, help="Model for computing covariance, currently supports empirical or dsets") - a.add_argument("--psname_ext", default='', type=str, help="Extension for pspectra name in PSpecContainer.") - a.add_argument("--verbose", default=False, action='store_true', help="Report feedback to standard output.") - a.add_argument("--file_type", default="uvh5", help="filetypes of input UVData. Default is 'uvh5'") - a.add_argument("--filter_extensions", default=None, type=list_of_int_tuples, help="List of spw filter extensions wrapped in quotes. Ex:20~20,40~40' ->> [(20, 20), (40, 40), ...]") - a.add_argument("--symmetric_taper", default=True, type=bool, help="If True, apply sqrt of taper before foreground filtering and then another sqrt after. If False, apply full taper after foreground Filter. ") - a.add_argument("--include_autocorrs", default=False, action="store_true", help="Include power spectra of autocorr visibilities.") - a.add_argument("--exclude_crosscorrs", default=False, action="store_true", help="If True, exclude cross-correlations from power spectra (autocorr power spectra only).") - a.add_argument("--interleave_times", default=False, action="store_true", help="Cross multiply even/odd time intervals.") - a.add_argument("--xant_flag_thresh", default=0.95, type=float, help="fraction of baseline waterfall that needs to be flagged for entire baseline to be flagged (and excluded from pspec)") - a.add_argument("--store_window", default=False, action="store_true", help="store window function array.") - a.add_argument("--allow_fft", default=False, action="store_true", help="use an FFT to comptue q-hat.") + a.add_argument( + "--dsets_std", + nargs="*", + default=None, + type=str, + help="List of miriad filepaths to visibility standard deviations.", + ) + a.add_argument( + "--groupname", + default=None, + type=str, + help="Groupname for the UVPSpec objects in the HDF5 container.", + ) + a.add_argument( + "--dset_pairs", + default=None, + type=list_of_int_tuples, + help="List of dset pairings for OQE. Two acceptable formats are " + "Ex1: '0~0,1~1' --> [(0, 0), (1, 1), ...] and " + "Ex2: '0 0, 1 1' --> [(0, 0), (1, 1), ...]", + ) + a.add_argument( + "--dset_labels", + default=None, + type=str, + nargs="*", + help="List of string labels for each input dataset.", + ) + a.add_argument( + "--spw_ranges", + default=None, + type=list_of_int_tuples, + help="List of spw channel selections. Two acceptable formats are " + "Ex1: '200~300,500~650' --> [(200, 300), (500, 650), ...] and " + "Ex2: '200 300, 500 650' --> [(200, 300), (500, 650), ...]", + ) + a.add_argument( + "--n_dlys", + default=None, + type=int, + nargs="+", + help="List of integers specifying number of delays to use per spectral window selection.", + ) + a.add_argument( + "--pol_pairs", + default=None, + type=list_of_str_tuples, + help="List of pol-string pairs to use in OQE. Two acceptable formats are " + "Ex1: 'xx~xx,yy~yy' --> [('xx', 'xx'), ('yy', 'yy'), ...] and " + "Ex2: 'xx xx, yy yy' --> [('xx', 'xx'), ('yy', 'yy'), ...]", + ) + a.add_argument( + "--blpairs", + default=None, + type=list_of_tuple_tuples, + help="List of baseline-pair antenna integers to run OQE on. Two acceptable formats are " + "Ex1: '1~2~3~4,5~6~7~8' --> [((1 2), (3, 4)), ((5, 6), (7, 8)), ...] and " + "Ex2: '1 2 3 4, 5 6 7 8' --> [((1 2), (3, 4)), ((5, 6), (7, 8)), ...]", + ) + a.add_argument( + "--input_data_weight", + default="identity", + type=str, + help="Data weighting for OQE. See PSpecData.pspec for details.", + ) + a.add_argument( + "--norm", + default="I", + type=str, + help="M-matrix normalization type for OQE. See PSpecData.pspec for details.", + ) + a.add_argument( + "--taper", + default="none", + type=str, + help="Taper function to use in OQE delay transform. See PSpecData.pspec for details.", + ) + a.add_argument( + "--beam", + default=None, + type=str, + help="Filepath to UVBeam healpix map of antenna beam.", + ) + a.add_argument( + "--cosmo", + default=None, + nargs="+", + type=float, + help="List of float values for [Om_L, Om_b, Om_c, H0, Om_M, Om_k].", + ) + a.add_argument( + "--rephase_to_dset", + default=None, + type=int, + help="dset integer index to phase all other dsets to. Default is no rephasing.", + ) + a.add_argument( + "--trim_dset_lsts", + default=False, + action="store_true", + help="Trim non-overlapping dset LSTs.", + ) + a.add_argument( + "--broadcast_dset_flags", + default=False, + action="store_true", + help="Broadcast dataset flags across time according to time_thresh.", + ) + a.add_argument( + "--time_thresh", + default=0.2, + type=float, + help="Fractional flagging threshold across time to trigger flag broadcast if broadcast_dset_flags is True", + ) + a.add_argument( + "--Jy2mK", + default=False, + action="store_true", + help="Convert datasets from Jy to mK if a beam model is provided.", + ) + a.add_argument( + "--exclude_auto_bls", + default=False, + action="store_true", + help="If blpairs is not provided, exclude all baselines paired with itself.", + ) + a.add_argument( + "--exclude_cross_bls", + default=False, + action="store_true", + help="If blpairs is not provided, exclude all baselines paired with a different baseline.", + ) + a.add_argument( + "--exclude_permutations", + default=False, + action="store_true", + help="If blpairs is not provided, exclude a basline-pair permutations. Ex: if (A, B) exists, exclude (B, A).", + ) + a.add_argument( + "--Nblps_per_group", + default=None, + type=int, + help="If blpairs is not provided and group == True, set the number of blpairs in each group.", + ) + a.add_argument( + "--bl_len_range", + default=(0, 1e10), + nargs="+", + type=float, + help="If blpairs is not provided, limit the baselines used based on their minimum and maximum length in meters.", + ) + a.add_argument( + "--bl_deg_range", + default=(0, 180), + nargs="+", + type=float, + help="If blpairs is not provided, limit the baseline used based on a min and max angle cut in ENU frame in degrees.", + ) + a.add_argument( + "--bl_error_tol", + default=1.0, + type=float, + help="If blpairs is not provided, this is the error tolerance in forming redundant baseline groups in meters.", + ) + a.add_argument( + "--store_cov", + default=False, + action="store_true", + help="Compute and store covariance of bandpowers given dsets_std files or empirical covariance.", + ) + a.add_argument( + "--store_cov_diag", + default=False, + action="store_true", + help="Compute and store the error bars calculated by QE formalism.", + ) + a.add_argument( + "--return_q", + default=False, + action="store_true", + help="Return unnormalized bandpowers given dsets files.", + ) + a.add_argument( + "--overwrite", + default=False, + action="store_true", + help="Overwrite output if it exists.", + ) + a.add_argument( + "--cov_model", + default="empirical", + type=str, + help="Model for computing covariance, currently supports empirical or dsets", + ) + a.add_argument( + "--psname_ext", + default="", + type=str, + help="Extension for pspectra name in PSpecContainer.", + ) + a.add_argument( + "--verbose", + default=False, + action="store_true", + help="Report feedback to standard output.", + ) + a.add_argument( + "--file_type", + default="uvh5", + help="filetypes of input UVData. Default is 'uvh5'", + ) + a.add_argument( + "--filter_extensions", + default=None, + type=list_of_int_tuples, + help="List of spw filter extensions wrapped in quotes. Ex:20~20,40~40' ->> [(20, 20), (40, 40), ...]", + ) + a.add_argument( + "--symmetric_taper", + default=True, + type=bool, + help="If True, apply sqrt of taper before foreground filtering and then another sqrt after. If False, apply full taper after foreground Filter. ", + ) + a.add_argument( + "--include_autocorrs", + default=False, + action="store_true", + help="Include power spectra of autocorr visibilities.", + ) + a.add_argument( + "--exclude_crosscorrs", + default=False, + action="store_true", + help="If True, exclude cross-correlations from power spectra (autocorr power spectra only).", + ) + a.add_argument( + "--interleave_times", + default=False, + action="store_true", + help="Cross multiply even/odd time intervals.", + ) + a.add_argument( + "--xant_flag_thresh", + default=0.95, + type=float, + help="fraction of baseline waterfall that needs to be flagged for entire baseline to be flagged (and excluded from pspec)", + ) + a.add_argument( + "--store_window", + default=False, + action="store_true", + help="store window function array.", + ) + a.add_argument( + "--allow_fft", + default=False, + action="store_true", + help="use an FFT to comptue q-hat.", + ) return a @@ -4393,9 +5442,9 @@ def validate_blpairs(blpairs, uvd1, uvd2, baseline_tol=1.0, verbose=True): If True report feedback to stdout. Default: True. """ # ensure uvd1 and uvd2 are UVData objects - if isinstance(uvd1, UVData) == False: + if not isinstance(uvd1, UVData): raise TypeError("uvd1 must be a UVData instance") - if isinstance(uvd2, UVData) == False: + if not isinstance(uvd2, UVData): raise TypeError("uvd2 must be a UVData instance") # get antenna position dictionary @@ -4407,23 +5456,27 @@ def validate_blpairs(blpairs, uvd1, uvd2, baseline_tol=1.0, verbose=True): # ensure shared antenna keys match within tolerance shared = sorted(set(ap1.keys()) & set(ap2.keys())) for k in shared: - assert np.linalg.norm(ap1[k] - ap2[k]) <= baseline_tol, \ - "uvd1 and uvd2 don't agree on antenna positions within " \ + assert np.linalg.norm(ap1[k] - ap2[k]) <= baseline_tol, ( + "uvd1 and uvd2 don't agree on antenna positions within " "tolerance of {} m".format(baseline_tol) + ) ap = ap1 ap.update(ap2) # iterate through baselines and check baselines crossed with each other # are within tolerance - for i, blg in enumerate(blpairs): + for blg in blpairs: if isinstance(blg, tuple): blg = [blg] for blp in blg: bl1_vec = ap[blp[0][0]] - ap[blp[0][1]] bl2_vec = ap[blp[1][0]] - ap[blp[1][1]] if np.linalg.norm(bl1_vec - bl2_vec) >= baseline_tol: - raise_warning("blpair {} exceeds redundancy tolerance of " - "{} m".format(blp, baseline_tol), verbose=verbose) + raise_warning( + "blpair {} exceeds redundancy tolerance of " + "{} m".format(blp, baseline_tol), + verbose=verbose, + ) def raise_warning(warning, verbose=True): @@ -4434,8 +5487,16 @@ def raise_warning(warning, verbose=True): print(warning) -def _load_dsets(fnames, bls=None, pols=None, logf=None, verbose=True, - file_type='miriad', cals=None, cal_flag=True): +def _load_dsets( + fnames, + bls=None, + pols=None, + logf=None, + verbose=True, + file_type="miriad", + cals=None, + cal_flag=True, +): """ Helper function for loading UVData-compatible datasets in pspec_run. @@ -4466,8 +5527,12 @@ def _load_dsets(fnames, bls=None, pols=None, logf=None, verbose=True, dsets = [] Ndsets = len(fnames) for i, dset in enumerate(fnames): - utils.log("Reading {} / {} datasets...".format(i+1, Ndsets), - f=logf, lvl=1, verbose=verbose) + utils.log( + "Reading {} / {} datasets...".format(i + 1, Ndsets), + f=logf, + lvl=1, + verbose=verbose, + ) # read data uvd = UVData() @@ -4475,13 +5540,13 @@ def _load_dsets(fnames, bls=None, pols=None, logf=None, verbose=True, dfiles = glob.glob(dset) else: dfiles = dset - uvd.read(dfiles, bls=bls, polarizations=pols, - file_type=file_type) - uvd.extra_keywords['filename'] = json.dumps(dfiles) + uvd.read(dfiles, bls=bls, polarizations=pols, file_type=file_type) + uvd.extra_keywords["filename"] = json.dumps(dfiles) dsets.append(uvd) return dsets + def _load_cals(cnames, logf=None, verbose=True): """ Helper function for loading calibration files. @@ -4504,8 +5569,12 @@ def _load_cals(cnames, logf=None, verbose=True): cals = [] Ncals = len(cnames) for i, cfile in enumerate(cnames): - utils.log("Reading {} / {} calibrations...".format(i+1, Ncals), - f=logf, lvl=1, verbose=verbose) + utils.log( + "Reading {} / {} calibrations...".format(i + 1, Ncals), + f=logf, + lvl=1, + verbose=verbose, + ) # read data uvc = UVCal() @@ -4513,7 +5582,7 @@ def _load_cals(cnames, logf=None, verbose=True): uvc.read_calfits(glob.glob(cfile)) else: uvc.read_calfits(cfile) - uvc.extra_keywords['filename'] = json.dumps(cfile) + uvc.extra_keywords["filename"] = json.dumps(cfile) cals.append(uvc) return cals diff --git a/hera_pspec/pstokes.py b/hera_pspec/pstokes.py index 96001341..2d238061 100644 --- a/hera_pspec/pstokes.py +++ b/hera_pspec/pstokes.py @@ -13,14 +13,16 @@ # See pyuvdata.utils.polstr2num for conversion between polarization string # and polarization integer. Ex. {'XX': -5, ...} pol_weights = { - 1: odict([(-5, 1.), (-6, 1.)]), - 2: odict([(-5, 1.), (-6, -1.)]), - 3: odict([(-7, 1.), (-8, 1.)]), - 4: odict([(-7, -1.j), (-8, 1.j)]) + 1: odict([(-5, 1.0), (-6, 1.0)]), + 2: odict([(-5, 1.0), (-6, -1.0)]), + 3: odict([(-7, 1.0), (-8, 1.0)]), + 4: odict([(-7, -1.0j), (-8, 1.0j)]), } -def miriad2pyuvdata(dset, antenna_nums=None, bls=None, polarizations=None, - ant_str=None, time_range=None): + +def miriad2pyuvdata( + dset, antenna_nums=None, bls=None, polarizations=None, ant_str=None, time_range=None +): """ Reads-in a Miriad filepath to a UVData object @@ -57,13 +59,18 @@ def miriad2pyuvdata(dset, antenna_nums=None, bls=None, polarizations=None, uvd : pyuvdata.UVData object """ uvd = pyuvdata.UVData() - uvd.read_miriad(dset, antenna_nums=antenna_nums, bls=bls, - polarizations=polarizations, ant_str=ant_str, - time_range=time_range) + uvd.read_miriad( + dset, + antenna_nums=antenna_nums, + bls=bls, + polarizations=polarizations, + ant_str=ant_str, + time_range=time_range, + ) return uvd -def _combine_pol(uvd1, uvd2, pol1, pol2, pstokes='pI', x_orientation=None): +def _combine_pol(uvd1, uvd2, pol1, pol2, pstokes="pI", x_orientation=None): """ Combines UVData visibilities to form the desired pseudo-stokes visibilities. It returns UVData object containing the pseudo-stokes visibilities @@ -98,10 +105,8 @@ def _combine_pol(uvd1, uvd2, pol1, pol2, pstokes='pI', x_orientation=None): ------- uvdS : UVData object """ - assert isinstance(uvd1, pyuvdata.UVData), \ - "uvd1 must be a pyuvdata.UVData instance" - assert isinstance(uvd2, pyuvdata.UVData), \ - "uvd2 must be a pyuvdata.UVData instance" + assert isinstance(uvd1, pyuvdata.UVData), "uvd1 must be a pyuvdata.UVData instance" + assert isinstance(uvd2, pyuvdata.UVData), "uvd2 must be a pyuvdata.UVData instance" # convert pol1 and/or pol2 to integer if fed as a string if isinstance(pol1, (str, np.str)): @@ -130,33 +135,55 @@ def _combine_pol(uvd1, uvd2, pol1, pol2, pstokes='pI', x_orientation=None): pstokes_str = pyuvdata.utils.polnum2str(pstokes) # assert pstokes in pol_weights, and pol1 and pol2 in pol_weights[pstokes] - assert pstokes in pol_weights, \ - "unrecognized pstokes parameter {}".format(pstokes_str) - assert pol1 in pol_weights[pstokes], \ - "pol1 {} not used in constructing pstokes {}".format(pol1_str, pstokes_str) - assert pol2 in pol_weights[pstokes], \ - "pol2 {} not used in constructing pstokes {}".format(pol2_str, pstokes_str) + assert pstokes in pol_weights, "unrecognized pstokes parameter {}".format( + pstokes_str + ) + assert ( + pol1 in pol_weights[pstokes] + ), "pol1 {} not used in constructing pstokes {}".format(pol1_str, pstokes_str) + assert ( + pol2 in pol_weights[pstokes] + ), "pol2 {} not used in constructing pstokes {}".format(pol2_str, pstokes_str) # constructing Stokes visibilities - stdata = 0.5 * (pol_weights[pstokes][pol1]*data1 + pol_weights[pstokes][pol2]*data2) + stdata = 0.5 * ( + pol_weights[pstokes][pol1] * data1 + pol_weights[pstokes][pol2] * data2 + ) # assigning and writing data, flags and metadata to UVData object uvdS = copy.deepcopy(uvd1) uvdS.data_array = stdata # pseudo-stokes data uvdS.flag_array = flag # flag array - uvdS.polarization_array = np.array([pstokes], dtype=np.int) # polarization number - uvdS.nsample_array = uvd1.nsample_array + uvd2.nsample_array # nsamples - uvdS.history = "Merged into pseudo-stokes vis with hera_pspec version {} Git hash {}\n{}" \ - "{}{}{}{}\n".format(version.version, version.git_hash, "-"*20+'\n', - 'dset1 history:\n', uvd1.history, '\n'+'-'*20+'\ndset2 history:\n', - uvd2.history) + uvdS.polarization_array = np.array([pstokes], dtype=np.int) # polarization number + uvdS.nsample_array = uvd1.nsample_array + uvd2.nsample_array # nsamples + uvdS.history = ( + "Merged into pseudo-stokes vis with hera_pspec version {} Git hash {}\n{}" + "{}{}{}{}\n".format( + version.version, + version.git_hash, + "-" * 20 + "\n", + "dset1 history:\n", + uvd1.history, + "\n" + "-" * 20 + "\ndset2 history:\n", + uvd2.history, + ) + ) return uvdS -def construct_pstokes(dset1, dset2, pstokes='pI', run_check=True, antenna_nums=None, - bls=None, polarizations=None, ant_str=None, time_range=None, - history=''): +def construct_pstokes( + dset1, + dset2, + pstokes="pI", + run_check=True, + antenna_nums=None, + bls=None, + polarizations=None, + ant_str=None, + time_range=None, + history="", +): """ Validates datasets required to construct desired visibilities and constructs desired pseudo-Stokes visibilities. These are formed @@ -219,19 +246,31 @@ def construct_pstokes(dset1, dset2, pstokes='pI', run_check=True, antenna_nums=N """ # convert dset1 and dset2 to UVData objects if they are miriad files if isinstance(dset1, pyuvdata.UVData) == False: - assert isinstance(dset1, (str, np.str)), \ - "dset1 must be fed as a string or UVData object" - uvd1 = miriad2pyuvdata(dset1, antenna_nums=antenna_nums, bls=bls, - polarizations=polarizations, ant_str=ant_str, - time_range=time_range) + assert isinstance( + dset1, (str, np.str) + ), "dset1 must be fed as a string or UVData object" + uvd1 = miriad2pyuvdata( + dset1, + antenna_nums=antenna_nums, + bls=bls, + polarizations=polarizations, + ant_str=ant_str, + time_range=time_range, + ) else: uvd1 = dset1 if isinstance(dset2, pyuvdata.UVData) == False: - assert isinstance(dset2, (str, np.str)), \ - "dset2 must be fed as a string or UVData object" - uvd2 = miriad2pyuvdata(dset2, antenna_nums=antenna_nums, bls=bls, - polarizations=polarizations, ant_str=ant_str, - time_range=time_range) + assert isinstance( + dset2, (str, np.str) + ), "dset2 must be fed as a string or UVData object" + uvd2 = miriad2pyuvdata( + dset2, + antenna_nums=antenna_nums, + bls=bls, + polarizations=polarizations, + ant_str=ant_str, + time_range=time_range, + ) else: uvd2 = dset2 @@ -242,7 +281,7 @@ def construct_pstokes(dset1, dset2, pstokes='pI', run_check=True, antenna_nums=N # check if dset1 and dset2 habe the same spectral window spw1 = uvd1.spw_array spw2 = uvd2.spw_array - assert (spw1 == spw2), "dset1 and dset2 must have the same spectral windows." + assert spw1 == spw2, "dset1 and dset2 must have the same spectral windows." # check if dset1 and dset2 have the same frequencies freqs1 = uvd1.freq_array @@ -271,19 +310,22 @@ def construct_pstokes(dset1, dset2, pstokes='pI', run_check=True, antenna_nums=N # check polarizations of UVData objects are consistent with the required # polarization to form the desired pseudo Stokes visibilities. If multiple # exist, downselect on polarization. - assert req_pol1 in uvd1.polarization_array, \ - "Polarization {} not found in dset1 object".format(req_pol1) + assert ( + req_pol1 in uvd1.polarization_array + ), "Polarization {} not found in dset1 object".format(req_pol1) if uvd1.Npols > 1: uvd1 = uvd1.select(polarizations=req_pol1, inplace=False) - assert req_pol2 in uvd2.polarization_array, \ - "Polarization {} not found in dset2 object".format(req_pol2) + assert ( + req_pol2 in uvd2.polarization_array + ), "Polarization {} not found in dset2 object".format(req_pol2) if uvd2.Npols > 1: uvd2 = uvd2.select(polarizations=req_pol2, inplace=False) # combining visibilities to form the desired Stokes visibilties - uvdS = _combine_pol(uvd1=uvd1, uvd2=uvd2, pol1=req_pol1, pol2=req_pol2, - pstokes=pstokes) + uvdS = _combine_pol( + uvd1=uvd1, uvd2=uvd2, pol1=req_pol1, pol2=req_pol2, pstokes=pstokes + ) uvdS.history += history if run_check: @@ -316,30 +358,32 @@ def filter_dset_on_stokes_pol(dsets, pstokes): to construct_pstokes to make the desired pseudo-Stokes visibility. """ # type check - assert isinstance(dsets, list), \ - "dsets must be fed as a list of UVData objects" - assert np.all(isinstance(d, UVData) for d in dsets), \ - "dsets must be fed as a list of UVData objects" + assert isinstance(dsets, list), "dsets must be fed as a list of UVData objects" + assert np.all( + isinstance(d, UVData) for d in dsets + ), "dsets must be fed as a list of UVData objects" # get polarization of each dset pols = [d.polarization_array[0] for d in dsets] # convert pstokes to integer if a string if isinstance(pstokes, (str, np.str)): - pstokes = pyuvdata.utils.polstr2num(pstokes, x_orientation=dsets[0].x_orientation) - assert pstokes in [1, 2, 3, 4], \ - "pstokes must be fed as a pseudo-Stokes parameter" + pstokes = pyuvdata.utils.polstr2num( + pstokes, x_orientation=dsets[0].x_orientation + ) + assert pstokes in [1, 2, 3, 4], "pstokes must be fed as a pseudo-Stokes parameter" # get two necessary dipole pols given pstokes desired_pols = list(pol_weights[pstokes].keys()) - assert desired_pols[0] in pols and desired_pols[1] in pols, \ - "necessary input pols {} and {} not found in dsets".format(*desired_pols) + assert ( + desired_pols[0] in pols and desired_pols[1] in pols + ), "necessary input pols {} and {} not found in dsets".format(*desired_pols) - inp_dsets = [dsets[pols.index(desired_pols[0])], - dsets[pols.index(desired_pols[1])]] + inp_dsets = [dsets[pols.index(desired_pols[0])], dsets[pols.index(desired_pols[1])]] return inp_dsets + def generate_pstokes_argparser(): """ Get argparser to generate pstokes from linpol files. @@ -349,15 +393,39 @@ def generate_pstokes_argparser(): Returns: a: argparser object with arguments used in generate_pstokes_run.py """ - a = argparse.ArgumentParser(description="argument parser for computing " - "pstokes from linpol files.") - a.add_argument("inputdata", type=str, help="Filename of UVData object with" - "linearly polarized data to add pstokes to.") - a.add_argument("--pstokes", type=str, help="list of pStokes you wish to calculate. Default is ['pI']", - nargs="+", default="pI") - a.add_argument("--outputdata", type=str, help="Filename to write out data. Output includes original linear pols." - "if no outputdata is provided, will use inputdata, appending" - "pstokes to original linear pols.") - a.add_argument("--clobber", action="store_true", default=False, help="Overwrite outputdata or original linpol only file.") - a.add_argument("--keep_vispols", action="store_true", default=False, help="If inplace, keep the original linear polarizations in the input file. Default is False.") + a = argparse.ArgumentParser( + description="argument parser for computing " "pstokes from linpol files." + ) + a.add_argument( + "inputdata", + type=str, + help="Filename of UVData object with" + "linearly polarized data to add pstokes to.", + ) + a.add_argument( + "--pstokes", + type=str, + help="list of pStokes you wish to calculate. Default is ['pI']", + nargs="+", + default="pI", + ) + a.add_argument( + "--outputdata", + type=str, + help="Filename to write out data. Output includes original linear pols." + "if no outputdata is provided, will use inputdata, appending" + "pstokes to original linear pols.", + ) + a.add_argument( + "--clobber", + action="store_true", + default=False, + help="Overwrite outputdata or original linpol only file.", + ) + a.add_argument( + "--keep_vispols", + action="store_true", + default=False, + help="If inplace, keep the original linear polarizations in the input file. Default is False.", + ) return a diff --git a/hera_pspec/testing.py b/hera_pspec/testing.py index 6b3d1370..c6c42149 100644 --- a/hera_pspec/testing.py +++ b/hera_pspec/testing.py @@ -547,7 +547,9 @@ def gauss_cov_fg(cov_amp, cov_length_scale, freqs, Ntimes=100, constant_in_time= return s -def sky_noise_jy_autos(lsts, freqs, autovis, omega_p, integration_time, channel_width=None, Trx=0.0): +def sky_noise_jy_autos( + lsts, freqs, autovis, omega_p, integration_time, channel_width=None, Trx=0.0 +): """Make a noise realization for a given auto-visibility level and beam. This is a simple replacement for ``hera_sim.noise.sky_noise_jy``. @@ -586,7 +588,7 @@ def sky_noise_jy_autos(lsts, freqs, autovis, omega_p, integration_time, channel_ # Calculate Jansky to Kelvin conversion factor # The factor of 1e-26 converts from Jy to W/m^2/Hz. wavelengths = conversions.units.c / freqs # meters - Jy2K = 1e-26 * wavelengths ** 2 / (2 * conversions.units.kb * omega_p) + Jy2K = 1e-26 * wavelengths**2 / (2 * conversions.units.kb * omega_p) # Use autocorrelation vsibility to set noise scale Tsky = autovis * Jy2K.reshape(1, -1) @@ -621,7 +623,9 @@ def sky_noise_sim( Parameters ---------- data : str or UVData object + The data beam : str or PSpecBeam object + The beam cov_amp : float Covariance amplitude. See gauss_cov_fg() cov_length_scale : float diff --git a/hera_pspec/tests/compare_legacy_oqe.py b/hera_pspec/tests/compare_legacy_oqe.py index 76a198e3..ddd6bf51 100755 --- a/hera_pspec/tests/compare_legacy_oqe.py +++ b/hera_pspec/tests/compare_legacy_oqe.py @@ -3,54 +3,57 @@ import numpy as np import matplotlib.pyplot as plt + def run_old_oqe(fname, key1, key2, freqrange): """ Run old OQE algorithm using capo. """ # Legacy functions to load data and estimate power spectrum import capo + legacy_read_files = capo.miriad.read_files legacy_group_redundant_bls = capo.red.group_redundant_bls legacy_oqe = capo.oqe # (1) Read data from file - s,d,f = legacy_read_files([fname], antstr='all', polstr='xx') + s, d, f = legacy_read_files([fname], antstr="all", polstr="xx") bls = d.keys() - #print("Baseline keys:", bls) + # print("Baseline keys:", bls) # (1a) Legacy setting to specify whether data are conjugated or not # PB: Replace this with made up conj array (all set to False) """ aa = aipy.cal.get_aa('psa6240_v003', np.array([.15])) - _, conj = legacy_group_redundant_bls(aa.ant_layout) + _, conj = legacy_group_redundant_bls(aa.ant_layout) # conj is a dictionary containing whether bl's are conjugated or not """ - conj = {bl:False for bl in bls} + conj = {bl: False for bl in bls} # (1b) Build data and flagging dictionaries - data_dict = {}; flg_dict = {} + data_dict = {} + flg_dict = {} fmin, fmax = freqrange for key in bls: # Use only a restricted band of frequencies (e.g. to avoid RFI) - data_dict[key] = d[key]['xx'][:,fmin:fmax] - flg_dict[key] = np.logical_not(f[key]['xx'][:,fmin:fmax]) + data_dict[key] = d[key]["xx"][:, fmin:fmax] + flg_dict[key] = np.logical_not(f[key]["xx"][:, fmin:fmax]) # (2) Make dataset object ds = legacy_oqe.DataSet() ds.set_data(dsets=data_dict, conj=conj, wgts=flg_dict) # (3) Calculate unweighted power spectrum - q_I = ds.q_hat(key1, key2, use_cov=False, cov_flagging=False) # unweighted + q_I = ds.q_hat(key1, key2, use_cov=False, cov_flagging=False) # unweighted F_I = ds.get_F(key1, key2, use_cov=False, cov_flagging=False) - M_I, W_I = ds.get_MW(F_I, mode='I') + M_I, W_I = ds.get_MW(F_I, mode="I") p_I = ds.p_hat(M_I, q_I) # (4) Calculate inverse covariance-weighted power spectrum - q = ds.q_hat(key1, key2, use_cov=True, cov_flagging=False) # weighted + q = ds.q_hat(key1, key2, use_cov=True, cov_flagging=False) # weighted F = ds.get_F(key1, key2, use_cov=True, cov_flagging=False) - M, W = ds.get_MW(F, mode='I') + M, W = ds.get_MW(F, mode="I") p = ds.p_hat(M, q) - + return p_I, p @@ -64,84 +67,85 @@ def run_new_oqe(fname, key1, key2, freqrange): # (1) Read data from file d1 = UVData() d1.read_miriad(fname) - + # (1a) Use only a restricted band of frequencies (e.g. to avoid RFI) fmin, fmax = freqrange d1.select(freq_chans=np.arange(fmin, fmax)) - + # (1b) Build data and flagging lists - d = [d1,] - w = [None for _d in d] # Set weights (None => flags from UVData will be used) - #print("Baseline keys:", d[0].get_antpairs()) + d = [ + d1, + ] + w = [None for _d in d] # Set weights (None => flags from UVData will be used) + # print("Baseline keys:", d[0].get_antpairs()) # (2) Make PSpecData object ds = pspec.PSpecData(dsets=d, wgts=w) # (3) Calculate unweighted power spectrum - ds.set_R('identity') + ds.set_R("identity") q_I = ds.q_hat(key1, key2) F_I = ds.get_G(key1, key2) - M_I, W_I = ds.get_MW(F_I, mode='I') + M_I, W_I = ds.get_MW(F_I, mode="I") p_I = ds.p_hat(M_I, q_I) - + # (4) Calculate inverse covariance-weighted power spectrum - ds.set_R('iC') + ds.set_R("iC") q = ds.q_hat(key1, key2) F = ds.get_G(key1, key2) - M, W = ds.get_MW(F, mode='I') + M, W = ds.get_MW(F, mode="I") p = ds.p_hat(M, q) - - #pspec, pairs = ds.pspec(bls, input_data_weight='I', norm='I', verbose=True) + + # pspec, pairs = ds.pspec(bls, input_data_weight='I', norm='I', verbose=True) return p_I, p -if __name__ == '__main__': - +if __name__ == "__main__": + # Path to datafile - fname = '../data/zen.2458042.12552.xx.HH.uvXAA' - + fname = "../data/zen.2458042.12552.xx.HH.uvXAA" + # Baselines to use key1 = (24, 25) key2 = (24, 38) - + # Frequency channels to include freqrange = (28, 52) - + # Run old OQE pI_old, p_old = run_old_oqe(fname, key1, key2, freqrange) print("Old:", p_old.shape) - + # Run new OQE _key1 = (0,) + key1 _key2 = (0,) + key2 pI_new, p_new = run_new_oqe(fname, _key1, _key2, freqrange) print("New:", p_new.shape) - + # Calculate fractional difference of means (averaged over LST) - frac_I = np.mean(pI_new, axis=1).real / np.mean(pI_old, axis=1).real - 1. - frac_iC = np.mean(p_new, axis=1).real / np.mean(p_old, axis=1).real - 1. - + frac_I = np.mean(pI_new, axis=1).real / np.mean(pI_old, axis=1).real - 1.0 + frac_iC = np.mean(p_new, axis=1).real / np.mean(p_old, axis=1).real - 1.0 + # Plot results (averaging over LST bins) plt.subplot(221) - plt.plot(np.mean(pI_old, axis=1).real, 'k-', lw=1.8, label="capo (I)") - plt.plot(np.mean(pI_new, axis=1).real, 'r--', lw=1.8, label="hera_pspec (I)") - plt.legend(loc='lower right') + plt.plot(np.mean(pI_old, axis=1).real, "k-", lw=1.8, label="capo (I)") + plt.plot(np.mean(pI_new, axis=1).real, "r--", lw=1.8, label="hera_pspec (I)") + plt.legend(loc="lower right") plt.ylabel(r"$\hat{p}$", fontsize=18) - + plt.subplot(222) - plt.plot(frac_I, 'k-', lw=1.8) + plt.plot(frac_I, "k-", lw=1.8) plt.ylabel(r"$P_{\rm pspec}/P_{\rm capo} - 1$ $(I)$", fontsize=18) - + plt.subplot(223) - plt.plot(np.mean(p_old, axis=1).real, 'k-', lw=1.8, label="capo (iC)") - plt.plot(np.mean(p_new, axis=1).real, 'r--', lw=1.8, label="hera_pspec (iC)") - plt.legend(loc='lower right') + plt.plot(np.mean(p_old, axis=1).real, "k-", lw=1.8, label="capo (iC)") + plt.plot(np.mean(p_new, axis=1).real, "r--", lw=1.8, label="hera_pspec (iC)") + plt.legend(loc="lower right") plt.ylabel(r"$\hat{p}$", fontsize=18) - + plt.subplot(224) - plt.plot(frac_iC, 'k-', lw=1.8) + plt.plot(frac_iC, "k-", lw=1.8) plt.ylabel(r"$P_{\rm pspec}/P_{\rm capo} - 1$ $(C^{-1})$", fontsize=18) - + plt.tight_layout() plt.show() - diff --git a/hera_pspec/tests/test_container.py b/hera_pspec/tests/test_container.py index 3219f794..7a33cfa2 100644 --- a/hera_pspec/tests/test_container.py +++ b/hera_pspec/tests/test_container.py @@ -1,44 +1,44 @@ import unittest import pytest import numpy as np -import os, sys, copy +import os from hera_pspec.data import DATA_PATH from .. import container, PSpecContainer, UVPSpec, testing -class Test_PSpecContainer(unittest.TestCase): +class Test_PSpecContainer(unittest.TestCase): def setUp(self, fill_pspec=True): # create container: first without SWMR to add header - self.fname = os.path.join(DATA_PATH, '_test_container.hdf5') + self.fname = os.path.join(DATA_PATH, "_test_container.hdf5") self.uvp, self.cosmo = testing.build_vanilla_uvpspec() # Create empty container without SWMR to fill header if os.path.exists(self.fname): os.remove(self.fname) - ps_store = PSpecContainer(self.fname, mode='rw', swmr=False) + ps_store = PSpecContainer(self.fname, mode="rw", swmr=False) del ps_store def tearDown(self): # Remove HDF5 file try: os.remove(self.fname) - except: + except FileNotFoundError: print("No HDF5 file to remove.") def runTest(self): pass def fill_container(self): - ps_store = PSpecContainer(self.fname, mode='rw', swmr=False) - self.group_names = ['group1', 'group2', 'group3'] - self.pspec_names = ['pspec_dset(0,1)', 'pspec_dset(1,0)', 'pspec_dset(1,1)'] + ps_store = PSpecContainer(self.fname, mode="rw", swmr=False) + self.group_names = ["group1", "group2", "group3"] + self.pspec_names = ["pspec_dset(0,1)", "pspec_dset(1,0)", "pspec_dset(1,1)"] for grp in self.group_names: for psname in self.pspec_names: - ps_store.set_pspec(group=grp, psname=psname, - pspec=self.uvp, overwrite=False) + ps_store.set_pspec( + group=grp, psname=psname, pspec=self.uvp, overwrite=False + ) del ps_store - def test_PSpecContainer(self, keep_open=True, swmr=False): """ Test that PSpecContainer works properly. @@ -47,65 +47,98 @@ def test_PSpecContainer(self, keep_open=True, swmr=False): self.fill_container() fname = self.fname group_names, pspec_names = self.group_names, self.pspec_names - ps_store = PSpecContainer(self.fname, mode='rw', keep_open=keep_open, swmr=swmr) + ps_store = PSpecContainer(self.fname, mode="rw", keep_open=keep_open, swmr=swmr) # Make sure that invalid mode arguments are caught - pytest.raises(ValueError, PSpecContainer, fname, mode='x') + pytest.raises(ValueError, PSpecContainer, fname, mode="x") # Check that power spectra can be overwritten for psname in pspec_names: - ps_store.set_pspec(group=group_names[2], psname=psname, - pspec=self.uvp, overwrite=True) + ps_store.set_pspec( + group=group_names[2], psname=psname, pspec=self.uvp, overwrite=True + ) # Check that overwriting fails if overwrite=False - pytest.raises(AttributeError, ps_store.set_pspec, group=group_names[2], - psname=psname, pspec=self.uvp, overwrite=False) + pytest.raises( + AttributeError, + ps_store.set_pspec, + group=group_names[2], + psname=psname, + pspec=self.uvp, + overwrite=False, + ) # Check that wrong pspec types are rejected by the set() method - pytest.raises(TypeError, ps_store.set_pspec, group=group_names[2], - psname=psname, pspec=np.arange(11), overwrite=True) - pytest.raises(TypeError, ps_store.set_pspec, group=group_names[2], - psname=psname, pspec=1, overwrite=True) - pytest.raises(TypeError, ps_store.set_pspec, group=group_names[2], - psname=psname, pspec="abc", overwrite=True) + pytest.raises( + TypeError, + ps_store.set_pspec, + group=group_names[2], + psname=psname, + pspec=np.arange(11), + overwrite=True, + ) + pytest.raises( + TypeError, + ps_store.set_pspec, + group=group_names[2], + psname=psname, + pspec=1, + overwrite=True, + ) + pytest.raises( + TypeError, + ps_store.set_pspec, + group=group_names[2], + psname=psname, + pspec="abc", + overwrite=True, + ) # Check that power spectra can be retrieved one by one for i in range(len(group_names)): # Get one pspec from each group ps = ps_store.get_pspec(group_names[i], psname=pspec_names[i]) - assert(isinstance(ps, UVPSpec)) + assert isinstance(ps, UVPSpec) # Check that power spectra can be retrieved from group as a list ps_list = ps_store.get_pspec(group_names[0]) - assert(len(ps_list) == len(pspec_names)) - for p in ps_list: assert(isinstance(p, UVPSpec)) + assert len(ps_list) == len(pspec_names) + for p in ps_list: + assert isinstance(p, UVPSpec) # Check that asking for an invalid group or pspec raises an error - pytest.raises(KeyError, ps_store.get_pspec, 'x', pspec_names[0]) + pytest.raises(KeyError, ps_store.get_pspec, "x", pspec_names[0]) pytest.raises(KeyError, ps_store.get_pspec, 1, pspec_names[0]) - pytest.raises(KeyError, ps_store.get_pspec, ['x', 'y'], pspec_names[0]) - pytest.raises(KeyError, ps_store.get_pspec, group_names[0], 'x') + pytest.raises(KeyError, ps_store.get_pspec, ["x", "y"], pspec_names[0]) + pytest.raises(KeyError, ps_store.get_pspec, group_names[0], "x") pytest.raises(KeyError, ps_store.get_pspec, group_names[0], 1) - pytest.raises(KeyError, ps_store.get_pspec, group_names[0], ['x','y']) + pytest.raises(KeyError, ps_store.get_pspec, group_names[0], ["x", "y"]) # Check that printing functions work print(ps_store) - assert(len(ps_store.tree()) > 0) + assert len(ps_store.tree()) > 0 # Check that read-only mode is respected - ps_readonly = PSpecContainer(fname, mode='r', keep_open=keep_open, swmr=swmr) + ps_readonly = PSpecContainer(fname, mode="r", keep_open=keep_open, swmr=swmr) ps = ps_readonly.get_pspec(group_names[0], pspec_names[0]) - assert(isinstance(ps, UVPSpec)) - pytest.raises(IOError, ps_readonly.set_pspec, group=group_names[2], - psname=pspec_names[2], pspec=self.uvp, overwrite=True) + assert isinstance(ps, UVPSpec) + pytest.raises( + IOError, + ps_readonly.set_pspec, + group=group_names[2], + psname=pspec_names[2], + pspec=self.uvp, + overwrite=True, + ) # Check that spectra() and groups() methods return the things we put in print("ps_store:", ps_store.data) grplist = ps_store.groups() pslist = ps_store.spectra(group=group_names[0]) - assert( len(grplist) == len(group_names) ) - assert( len(pslist) == len(pspec_names) ) - for g in grplist: assert(g in group_names) + assert len(grplist) == len(group_names) + assert len(pslist) == len(pspec_names) + for g in grplist: + assert g in group_names # Check that spectra() raises an error if group doesn't exist pytest.raises(KeyError, ps_store.spectra, "x") @@ -114,31 +147,59 @@ def test_PSpecContainer(self, keep_open=True, swmr=False): for g in ps_store.groups(): for psname in ps_store.spectra(group=g): ps = ps_store.get_pspec(g, psname=psname) - assert(isinstance(ps, UVPSpec)) + assert isinstance(ps, UVPSpec) # check partial IO in get_pspec ps = ps_store.get_pspec(group_names[0], pspec_names[0], just_meta=True) - assert not hasattr(ps, 'data_array') - assert hasattr(ps, 'time_avg_array') - ps = ps_store.get_pspec(group_names[0], pspec_names[0], blpairs=[((1, 2), (1, 2))]) - assert hasattr(ps, 'data_array') + assert not hasattr(ps, "data_array") + assert hasattr(ps, "time_avg_array") + ps = ps_store.get_pspec( + group_names[0], pspec_names[0], blpairs=[((1, 2), (1, 2))] + ) + assert hasattr(ps, "data_array") assert np.all(np.isclose(ps.blpair_array, 101102101102)) # Check that invalid list arguments raise errors in set_pspec() - pytest.raises(ValueError, ps_store.set_pspec, group=group_names[:2], - psname=pspec_names[0], pspec=self.uvp, overwrite=True) - pytest.raises(ValueError, ps_store.set_pspec, group=group_names[0], - psname=pspec_names, pspec=self.uvp, overwrite=True) - pytest.raises(ValueError, ps_store.set_pspec, group=group_names[0], - psname=pspec_names[0], pspec=[self.uvp, self.uvp, self.uvp], - overwrite=True) - pytest.raises(TypeError, ps_store.set_pspec, group=group_names[0], - psname=pspec_names, pspec=[self.uvp, None, self.uvp], - overwrite=True) + pytest.raises( + ValueError, + ps_store.set_pspec, + group=group_names[:2], + psname=pspec_names[0], + pspec=self.uvp, + overwrite=True, + ) + pytest.raises( + ValueError, + ps_store.set_pspec, + group=group_names[0], + psname=pspec_names, + pspec=self.uvp, + overwrite=True, + ) + pytest.raises( + ValueError, + ps_store.set_pspec, + group=group_names[0], + psname=pspec_names[0], + pspec=[self.uvp, self.uvp, self.uvp], + overwrite=True, + ) + pytest.raises( + TypeError, + ps_store.set_pspec, + group=group_names[0], + psname=pspec_names, + pspec=[self.uvp, None, self.uvp], + overwrite=True, + ) # Check that lists can be used to set pspec - ps_store.set_pspec(group=group_names[0], psname=pspec_names, - pspec=[self.uvp, self.uvp, self.uvp], overwrite=True) + ps_store.set_pspec( + group=group_names[0], + psname=pspec_names, + pspec=[self.uvp, self.uvp, self.uvp], + overwrite=True, + ) # Check that save() can be called ps_store.save() @@ -159,14 +220,14 @@ def test_container_transactional_mode(self): group_names, pspec_names = self.group_names, self.pspec_names # Test to see whether concurrent read/write works - psc_rw = PSpecContainer(fname, mode='rw', keep_open=False, swmr=True) + psc_rw = PSpecContainer(fname, mode="rw", keep_open=False, swmr=True) # Try to open container read-only (transactional) - psc_ro = PSpecContainer(fname, mode='r', keep_open=False, swmr=True) + psc_ro = PSpecContainer(fname, mode="r", keep_open=False, swmr=True) assert len(psc_ro.groups()) == len(group_names) # Open container read-only in non-transactional mode - psc_ro_noatom = PSpecContainer(fname, mode='r', keep_open=True, swmr=True) + psc_ro_noatom = PSpecContainer(fname, mode="r", keep_open=True, swmr=True) assert len(psc_ro_noatom.groups()) == len(group_names) # Original RO handle should work fine; RW handle will throw an error @@ -176,17 +237,30 @@ def test_container_transactional_mode(self): # Close the non-transactional file; the RW file should now work psc_ro_noatom._close() - psc_rw.set_pspec(group=group_names[0], psname=pspec_names[0], - pspec=self.uvp, overwrite=True) + psc_rw.set_pspec( + group=group_names[0], psname=pspec_names[0], pspec=self.uvp, overwrite=True + ) # test that write of new group or dataset with SWMR is blocked - pytest.raises(ValueError, psc_rw.set_pspec, group='new_group', psname=pspec_names[0], - pspec=self.uvp, overwrite=True) - pytest.raises(ValueError, psc_rw.set_pspec, group=group_names[0], psname='new_psname', - pspec=self.uvp, overwrite=True) + pytest.raises( + ValueError, + psc_rw.set_pspec, + group="new_group", + psname=pspec_names[0], + pspec=self.uvp, + overwrite=True, + ) + pytest.raises( + ValueError, + psc_rw.set_pspec, + group=group_names[0], + psname="new_psname", + pspec=self.uvp, + overwrite=True, + ) # ensure SWMR attr is propagated - for m in ['r', 'rw']: + for m in ["r", "rw"]: psc = PSpecContainer(fname, mode=m, keep_open=True, swmr=True) assert psc.swmr assert psc.data.swmr_mode @@ -196,46 +270,45 @@ def test_container_transactional_mode(self): assert not psc.data.swmr_mode psc._close() + def test_combine_psc_spectra(): fname = os.path.join(DATA_PATH, "zen.2458042.17772.xx.HH.uvXA") - uvp1 = testing.uvpspec_from_data(fname, [(24, 25), (37, 38)], - spw_ranges=[(10, 40)]) - uvp2 = testing.uvpspec_from_data(fname, [(38, 39), (52, 53)], - spw_ranges=[(10, 40)]) + uvp1 = testing.uvpspec_from_data(fname, [(24, 25), (37, 38)], spw_ranges=[(10, 40)]) + uvp2 = testing.uvpspec_from_data(fname, [(38, 39), (52, 53)], spw_ranges=[(10, 40)]) # test basic execution - if os.path.exists('ex.h5'): - os.remove('ex.h5') - psc = PSpecContainer("ex.h5", mode='rw') + if os.path.exists("ex.h5"): + os.remove("ex.h5") + psc = PSpecContainer("ex.h5", mode="rw") psc.set_pspec("grp1", "uvp_a", uvp1, overwrite=True) psc.set_pspec("grp1", "uvp_b", uvp2, overwrite=True) - container.combine_psc_spectra(psc, dset_split_str=None, ext_split_str='_') - assert psc.spectra('grp1') == [u'uvp'] + container.combine_psc_spectra(psc, dset_split_str=None, ext_split_str="_") + assert psc.spectra("grp1") == ["uvp"] # test dset name handling - if os.path.exists('ex.h5'): - os.remove('ex.h5') - psc = PSpecContainer("ex.h5", mode='rw') + if os.path.exists("ex.h5"): + os.remove("ex.h5") + psc = PSpecContainer("ex.h5", mode="rw") psc.set_pspec("grp1", "d1_x_d2_a", uvp1, overwrite=True) psc.set_pspec("grp1", "d1_x_d2_b", uvp2, overwrite=True) psc.set_pspec("grp1", "d2_x_d3_a", uvp1, overwrite=True) psc.set_pspec("grp1", "d2_x_d3_b", uvp2, overwrite=True) - container.combine_psc_spectra('ex.h5', dset_split_str='_x_', ext_split_str='_') - spec_list = psc.spectra('grp1') + container.combine_psc_spectra("ex.h5", dset_split_str="_x_", ext_split_str="_") + spec_list = psc.spectra("grp1") spec_list.sort() - assert spec_list == [u'd1_x_d2', u'd2_x_d3'] + assert spec_list == ["d1_x_d2", "d2_x_d3"] # test exceptions - if os.path.exists('ex.h5'): - os.remove('ex.h5') - psc = PSpecContainer("ex.h5", mode='rw') + if os.path.exists("ex.h5"): + os.remove("ex.h5") + psc = PSpecContainer("ex.h5", mode="rw") # test no group exception pytest.raises(AssertionError, container.combine_psc_spectra, psc) # test failed combine_uvpspec psc.set_pspec("grp1", "d1_x_d2_a", uvp1, overwrite=True) psc.set_pspec("grp1", "d1_x_d2_b", uvp1, overwrite=True) - container.combine_psc_spectra(psc, dset_split_str='_x_', ext_split_str='_') - assert psc.spectra('grp1') == [u'd1_x_d2_a', u'd1_x_d2_b'] + container.combine_psc_spectra(psc, dset_split_str="_x_", ext_split_str="_") + assert psc.spectra("grp1") == ["d1_x_d2_a", "d1_x_d2_b"] if os.path.exists("ex.h5"): os.remove("ex.h5") diff --git a/hera_pspec/tests/test_conversions.py b/hera_pspec/tests/test_conversions.py index 703e999b..268c22b8 100644 --- a/hera_pspec/tests/test_conversions.py +++ b/hera_pspec/tests/test_conversions.py @@ -8,11 +8,11 @@ class Test_Cosmo(unittest.TestCase): - def setUp(self): - self.C = conversions.Cosmo_Conversions(Om_L=0.68440, Om_b=0.04911, - Om_c=0.26442, H0=100.0, - Om_M=None, Om_k=None) + self.C = conversions.Cosmo_Conversions( + Om_L=0.68440, Om_b=0.04911, Om_c=0.26442, H0=100.0, Om_M=None, Om_k=None + ) + def tearDown(self): pass @@ -31,7 +31,6 @@ def test_init(self): C = conversions.Cosmo_Conversions(H0=25.5) np.testing.assert_almost_equal(C.H0, 25.5) - def test_units(self): si = conversions.units() cgs = conversions.cgs_units() @@ -45,41 +44,55 @@ def test_distances(self): np.testing.assert_almost_equal(self.C.z2f(10.0, ghz=True), 0.12912779554545455) np.testing.assert_almost_equal(self.C.E(10.0), 20.450997530682947) np.testing.assert_almost_equal(self.C.DC(10.0), 6499.708111027144) - np.testing.assert_almost_equal(self.C.DC(10.0, little_h=False), 6499.708111027144) + np.testing.assert_almost_equal( + self.C.DC(10.0, little_h=False), 6499.708111027144 + ) np.testing.assert_almost_equal(self.C.DM(10.0), 6510.2536925709637) np.testing.assert_almost_equal(self.C.DA(10.0), 591.84124477917851) np.testing.assert_almost_equal(self.C.dRperp_dtheta(10.0), 6510.2536925709637) np.testing.assert_almost_equal(self.C.dRpara_df(10.0), 1.2487605057418872e-05) - np.testing.assert_almost_equal(self.C.dRpara_df(10.0, ghz=True), 12487.605057418872) + np.testing.assert_almost_equal( + self.C.dRpara_df(10.0, ghz=True), 12487.605057418872 + ) np.testing.assert_almost_equal(self.C.X2Y(10.0), 529.26719942209002) - def test_little_h(self): # Test that putting in a really low value of H0 into the conversions object # has no effect on the result if little_h=True units are used - self.C = conversions.Cosmo_Conversions(Om_L=0.68440, Om_b=0.04911, - Om_c=0.26442, H0=25.12, - Om_M=None, Om_k=None) + self.C = conversions.Cosmo_Conversions( + Om_L=0.68440, Om_b=0.04911, Om_c=0.26442, H0=25.12, Om_M=None, Om_k=None + ) np.testing.assert_almost_equal(self.C.f2z(100e6), 13.20405751) np.testing.assert_almost_equal(self.C.z2f(10.0), 129127795.54545455) np.testing.assert_almost_equal(self.C.E(10.0), 20.450997530682947) - np.testing.assert_almost_equal(self.C.DC(10.0, little_h=True), 6499.708111027144) - np.testing.assert_almost_equal(self.C.DC(10.0, little_h=False), - 6499.708111027144*100./25.12) - np.testing.assert_almost_equal(self.C.DM(10.0, little_h=True), 6510.2536925709637) - np.testing.assert_almost_equal(self.C.DA(10.0, little_h=True), 591.84124477917851) - np.testing.assert_almost_equal(self.C.dRperp_dtheta(10.0, little_h=True), - 6510.2536925709637) - np.testing.assert_almost_equal(self.C.dRpara_df(10.0, little_h=True), - 1.2487605057418872e-05) - np.testing.assert_almost_equal(self.C.dRpara_df(10.0, ghz=True, little_h=True), - 12487.605057418872) - np.testing.assert_almost_equal(self.C.X2Y(10.0, little_h=True), - 529.26719942209002) + np.testing.assert_almost_equal( + self.C.DC(10.0, little_h=True), 6499.708111027144 + ) + np.testing.assert_almost_equal( + self.C.DC(10.0, little_h=False), 6499.708111027144 * 100.0 / 25.12 + ) + np.testing.assert_almost_equal( + self.C.DM(10.0, little_h=True), 6510.2536925709637 + ) + np.testing.assert_almost_equal( + self.C.DA(10.0, little_h=True), 591.84124477917851 + ) + np.testing.assert_almost_equal( + self.C.dRperp_dtheta(10.0, little_h=True), 6510.2536925709637 + ) + np.testing.assert_almost_equal( + self.C.dRpara_df(10.0, little_h=True), 1.2487605057418872e-05 + ) + np.testing.assert_almost_equal( + self.C.dRpara_df(10.0, ghz=True, little_h=True), 12487.605057418872 + ) + np.testing.assert_almost_equal( + self.C.X2Y(10.0, little_h=True), 529.26719942209002 + ) def test_params(self): params = self.C.get_params() - np.testing.assert_almost_equal(params['Om_L'], self.C.Om_L) + np.testing.assert_almost_equal(params["Om_L"], self.C.Om_L) def test_kpara_kperp(self): bl2kperp = self.C.bl_to_kperp(10.0, little_h=True) @@ -87,5 +100,6 @@ def test_kpara_kperp(self): np.testing.assert_almost_equal(bl2kperp, 0.00041570092391078579) np.testing.assert_almost_equal(tau2kpara, 503153.74952115043) + if __name__ == "__main__": unittest.main() diff --git a/hera_pspec/tests/test_grouping.py b/hera_pspec/tests/test_grouping.py index 06af83fe..eaa668d2 100644 --- a/hera_pspec/tests/test_grouping.py +++ b/hera_pspec/tests/test_grouping.py @@ -3,17 +3,16 @@ import numpy as np import os from hera_pspec.data import DATA_PATH -from .. import uvpspec, conversions, parameter, pspecbeam, pspecdata, testing, utils -from .. import uvpspec_utils as uvputils +from .. import uvpspec, conversions, pspecbeam, pspecdata, testing, utils from .. import grouping, container from pyuvdata import UVData from hera_cal import redcal import copy -class Test_grouping(unittest.TestCase): +class Test_grouping(unittest.TestCase): def setUp(self): - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") self.beam = pspecbeam.PSpecBeamUV(beamfile) uvp, cosmo = testing.build_vanilla_uvpspec(beam=self.beam) uvp.check() @@ -30,12 +29,12 @@ def test_group_baselines(self): Test baseline grouping behavior. """ # Generate example lists of baselines - bls1 = [(0,i) for i in range(1)] - bls2 = [(0,i) for i in range(2)] - bls3 = [(0,i) for i in range(4)] - bls4 = [(0,i) for i in range(5)] - bls5 = [(0,i) for i in range(13)] - bls6 = [(0,i) for i in range(521)] + bls1 = [(0, i) for i in range(1)] + bls2 = [(0, i) for i in range(2)] + bls3 = [(0, i) for i in range(4)] + bls4 = [(0, i) for i in range(5)] + bls5 = [(0, i) for i in range(13)] + bls6 = [(0, i) for i in range(521)] # Check that error is raised when more groups requested than baselines pytest.raises(ValueError, grouping.group_baselines, bls1, 2) @@ -53,24 +52,23 @@ def test_group_baselines(self): # Loop over groups and check that blocks are equal in size gs = [g1a, g1b, g1c, g2a, g2b, g2c] for g in gs: - assert (np.unique([len(grp) for grp in g]).size == 1) + assert np.unique([len(grp) for grp in g]).size == 1 # Check that total no. baselines is preserved with keep_remainder=False for bls in [bls1, bls2, bls3, bls4, bls5, bls6]: for ngrp in [1, 2, 5, 10, 45]: for rand in [True, False]: try: - g = grouping.group_baselines(bls, ngrp, - keep_remainder=True, - randomize=rand) - except: + g = grouping.group_baselines( + bls, ngrp, keep_remainder=True, randomize=rand + ) + except Exception: continue count = np.sum([len(_g) for _g in g]) assert count == len(bls) # Check that random seed works g1 = grouping.group_baselines(bls5, 3, randomize=True, seed=10) - g2 = grouping.group_baselines(bls5, 3, randomize=True, seed=11) g3 = grouping.group_baselines(bls5, 3, randomize=True, seed=10) for i in range(len(g1)): for j in range(len(g1[i])): @@ -80,16 +78,16 @@ def test_average_spectra(self): """ Test average spectra behavior. """ - ## Start file prep ## - dfile = os.path.join(DATA_PATH, 'zen.all.xx.LST.1.06964.uvA') + # === Start file prep + dfile = os.path.join(DATA_PATH, "zen.all.xx.LST.1.06964.uvA") # Load into UVData objects uvd = UVData() uvd.read_miriad(dfile) cosmo = conversions.Cosmo_Conversions() - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") uvb = pspecbeam.PSpecBeamUV(beamfile, cosmo=cosmo) # find conversion factor from Jy to mK - Jy_to_mK = uvb.Jy_to_mK(np.unique(uvd.freq_array), pol='XX') + Jy_to_mK = uvb.Jy_to_mK(np.unique(uvd.freq_array), pol="XX") # reshape to appropriately match a UVData.data_array object and multiply in! uvd.data_array *= Jy_to_mK[None, None, :, None] @@ -101,29 +99,32 @@ def test_average_spectra(self): ds = pspecdata.PSpecData(dsets=[uvd1, uvd2], wgts=[None, None], beam=uvb) ds.rephase_to_dset(0) # change units of UVData objects - ds.dsets[0].vis_units = 'mK' - ds.dsets[1].vis_units = 'mK' + ds.dsets[0].vis_units = "mK" + ds.dsets[1].vis_units = "mK" - baselines = [(24,25), (37,38), (38,39)] + baselines = [(24, 25), (37, 38), (38, 39)] # calculate all baseline pairs from group - baselines1, baselines2, blpairs \ - = utils.construct_blpairs(baselines, - exclude_auto_bls=True, - exclude_permutations=True) - - uvp = ds.pspec(baselines1, baselines2, (0, 1), - [('xx', 'xx')], - spw_ranges=[(300, 350)], - input_data_weight='identity', - norm='I', - taper='blackman-harris', - store_cov=True, - cov_model="autos", - verbose=False) + baselines1, baselines2, blpairs = utils.construct_blpairs( + baselines, exclude_auto_bls=True, exclude_permutations=True + ) + + uvp = ds.pspec( + baselines1, + baselines2, + (0, 1), + [("xx", "xx")], + spw_ranges=[(300, 350)], + input_data_weight="identity", + norm="I", + taper="blackman-harris", + store_cov=True, + cov_model="autos", + verbose=False, + ) keys = uvp.get_all_keys() # Add the analytic noise to stat_array - Pn = uvp.generate_noise_spectra(0, 'xx', 220) + Pn = uvp.generate_noise_spectra(0, "xx", 220) for key in keys: blp = uvp.antnums_to_blpair(key[1]) error = Pn[blp] @@ -134,138 +135,169 @@ def test_average_spectra(self): for key in keys: uvp.set_stats("simple", key, errs) blpair_groups = [blpairs] - ## End file prep ## + # End file prep # begin tests - uvp_avg_ints_wgts = grouping.average_spectra(uvp, - blpair_groups=blpair_groups, - error_field="noise", - time_avg=True, - inplace=False) - - uvp_avg_noise_wgts = grouping.average_spectra(uvp, - time_avg=True, - blpair_groups=blpair_groups, - error_weights="noise", - inplace=False) - uvp_avg_simple_wgts = grouping.average_spectra(uvp, - blpair_groups=blpair_groups, - time_avg=True, - error_weights="simple", - inplace=False) + uvp_avg_ints_wgts = grouping.average_spectra( + uvp, + blpair_groups=blpair_groups, + error_field="noise", + time_avg=True, + inplace=False, + ) + + uvp_avg_noise_wgts = grouping.average_spectra( + uvp, + time_avg=True, + blpair_groups=blpair_groups, + error_weights="noise", + inplace=False, + ) + uvp_avg_simple_wgts = grouping.average_spectra( + uvp, + blpair_groups=blpair_groups, + time_avg=True, + error_weights="simple", + inplace=False, + ) # For using uniform error bars as weights, the error bar on the average # is 1/sqrt{N} times the error bar on one single sample. averaged_stat = uvp_avg_simple_wgts.stats_array["simple"][0][0, 0, 0] - initial_stat = uvp.stats_array["simple"][0][0, 0, 0] \ - / np.sqrt(uvp.Ntimes) / np.sqrt(len(blpairs)) + initial_stat = ( + uvp.stats_array["simple"][0][0, 0, 0] + / np.sqrt(uvp.Ntimes) + / np.sqrt(len(blpairs)) + ) assert np.all(np.isclose(initial_stat, averaged_stat)) # For non-uniform weights, we test the error bar on the average power # spectra should be smaller than one on single sample. - assert( abs(uvp_avg_ints_wgts.stats_array["noise"][0][0, 0, 0]) \ - < abs(uvp.stats_array["noise"][0][0, 0, 0])) - assert( abs(uvp_avg_noise_wgts.stats_array["noise"][0][0, 0, 0]) \ - < abs(uvp.stats_array["noise"][0][0, 0, 0])) + assert abs(uvp_avg_ints_wgts.stats_array["noise"][0][0, 0, 0]) < abs( + uvp.stats_array["noise"][0][0, 0, 0] + ) + assert abs(uvp_avg_noise_wgts.stats_array["noise"][0][0, 0, 0]) < abs( + uvp.stats_array["noise"][0][0, 0, 0] + ) # Test stats inf variance for all times, single blpair doesn't result # in nans and that the avg effectively ignores its presence: e.g. check # it matches initial over sqrt(Nblpairs - 1) uvp_inf_var = copy.deepcopy(uvp) - initial_stat = uvp.get_stats('simple', (0, blpairs[0], 'xx')) + initial_stat = uvp.get_stats("simple", (0, blpairs[0], "xx")) inf_var_stat = np.ones((uvp_inf_var.Ntimes, uvp_inf_var.Ndlys)) * np.inf - uvp_inf_var.set_stats('simple', (0, blpairs[1], 'xx'), inf_var_stat) - uvp_inf_var_avg = uvp_inf_var.average_spectra(blpair_groups=blpair_groups, - error_weights='simple', - inplace=False) - final_stat = uvp_inf_var_avg.get_stats('simple', (0, blpairs[0], 'xx')) + uvp_inf_var.set_stats("simple", (0, blpairs[1], "xx"), inf_var_stat) + uvp_inf_var_avg = uvp_inf_var.average_spectra( + blpair_groups=blpair_groups, error_weights="simple", inplace=False + ) + final_stat = uvp_inf_var_avg.get_stats("simple", (0, blpairs[0], "xx")) assert np.isclose(final_stat, initial_stat / np.sqrt(len(blpairs) - 1)).all() # Test infinite variance for single time, all blpairs doesn't result in nans # and check that averaged stat for that time is inf (not zero) uvp_inf_var = copy.deepcopy(uvp) - initial_stat = uvp.get_stats('simple', (0, blpairs[0], 'xx')) + initial_stat = uvp.get_stats("simple", (0, blpairs[0], "xx")) inf_var_stat = np.ones((uvp_inf_var.Ntimes, uvp_inf_var.Ndlys)) inf_var_stat[0] = np.inf for blp in blpairs: - uvp_inf_var.set_stats('simple', (0, blp, 'xx'), inf_var_stat) - uvp_inf_var_avg = uvp_inf_var.average_spectra(blpair_groups=blpair_groups, error_weights='simple', inplace=False) - final_stat = uvp_inf_var_avg.get_stats('simple', (0, blpairs[0], 'xx')) - assert np.isclose(final_stat[1:], initial_stat[1:] / np.sqrt(len(blpairs))).all() + uvp_inf_var.set_stats("simple", (0, blp, "xx"), inf_var_stat) + uvp_inf_var_avg = uvp_inf_var.average_spectra( + blpair_groups=blpair_groups, error_weights="simple", inplace=False + ) + final_stat = uvp_inf_var_avg.get_stats("simple", (0, blpairs[0], "xx")) + assert np.isclose( + final_stat[1:], initial_stat[1:] / np.sqrt(len(blpairs)) + ).all() assert np.all(~np.isfinite(final_stat[0])) # Tests related to exact_windows # prep objects uvd = UVData() - uvd.read_uvh5(os.path.join(DATA_PATH, 'zen.2458116.31939.HH.uvh5')) + uvd.read_uvh5(os.path.join(DATA_PATH, "zen.2458116.31939.HH.uvh5")) ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=uvb) - baselines1, baselines2, blpairs = utils.construct_blpairs(uvd.get_antpairs()[1:], - exclude_permutations=False, - exclude_auto_bls=True) + baselines1, baselines2, blpairs = utils.construct_blpairs( + uvd.get_antpairs()[1:], exclude_permutations=False, exclude_auto_bls=True + ) # compute ps - uvp = ds.pspec(baselines1, baselines2, dsets=(0, 1), pols=[('xx', 'xx')], - spw_ranges=(175,195), taper='bh',verbose=False) + uvp = ds.pspec( + baselines1, + baselines2, + dsets=(0, 1), + pols=[("xx", "xx")], + spw_ranges=(175, 195), + taper="bh", + verbose=False, + ) # get exact window functions - uvp.get_exact_window_functions(ftbeam_file = os.path.join(DATA_PATH, 'FT_beam_HERA_dipole_test'), - spw_array=None, inplace=True, verbose=False) + uvp.get_exact_window_functions( + ftbeam_file=os.path.join(DATA_PATH, "FT_beam_HERA_dipole_test"), + spw_array=None, + inplace=True, + verbose=False, + ) # time average - uvp_time_avg = grouping.average_spectra(uvp, - blpair_groups=None, - time_avg=True, - blpair_weights=None, - error_field=None, - error_weights=None, - normalize_weights=True, - inplace=False, - add_to_history='') + uvp_time_avg = grouping.average_spectra( + uvp, + blpair_groups=None, + time_avg=True, + blpair_weights=None, + error_field=None, + error_weights=None, + normalize_weights=True, + inplace=False, + add_to_history="", + ) assert uvp_time_avg.Nblpairts == uvp_time_avg.Nblpairs assert uvp_time_avg.window_function_array[0].shape[0] == uvp_time_avg.Nblpairts blpair_groups, blpair_lens, _ = uvp.get_red_blpairs() # redundant average - uvp_red_avg = grouping.average_spectra(uvp, - blpair_groups=blpair_groups, - time_avg=False, - blpair_weights=None, - error_field=None, - error_weights=None, - normalize_weights=True, - inplace=False, - add_to_history='') + uvp_red_avg = grouping.average_spectra( + uvp, + blpair_groups=blpair_groups, + time_avg=False, + blpair_weights=None, + error_field=None, + error_weights=None, + normalize_weights=True, + inplace=False, + add_to_history="", + ) assert uvp_red_avg.Nblpairts == uvp_red_avg.Ntimes # both + error_weights keys = uvp.get_all_keys() # Add the analytic noise to stat_array - Pn = uvp.generate_noise_spectra(0, 'xx', 220) + Pn = uvp.generate_noise_spectra(0, "xx", 220) for key in keys: blp = uvp.antnums_to_blpair(key[1]) error = Pn[blp] uvp.set_stats("noise", key, error) - uvp_avg = grouping.average_spectra(uvp, - blpair_groups=blpair_groups, - time_avg=True, - blpair_weights=None, - error_field='noise', - error_weights=None, - normalize_weights=True, - inplace=False, - add_to_history='') - + uvp_avg = grouping.average_spectra( + uvp, + blpair_groups=blpair_groups, + time_avg=True, + blpair_weights=None, + error_field="noise", + error_weights=None, + normalize_weights=True, + inplace=False, + add_to_history="", + ) def test_sample_baselines(self): """ Test baseline sampling (with replacement) behavior. """ # Generate example lists of baselines - bls1 = [(0,i) for i in range(1)] - bls2 = [(0,i) for i in range(2)] - bls3 = [(0,i) for i in range(4)] - bls4 = [(0,i) for i in range(5)] - bls5 = [(0,i) for i in range(13)] - bls6 = [(0,i) for i in range(521)] + bls1 = [(0, i) for i in range(1)] + bls2 = [(0, i) for i in range(2)] + bls3 = [(0, i) for i in range(4)] + bls4 = [(0, i) for i in range(5)] + bls5 = [(0, i) for i in range(13)] + bls6 = [(0, i) for i in range(521)] # Example grouped list g1 = grouping.group_baselines(bls5, 3, randomize=False) @@ -284,13 +316,23 @@ def test_bootstrap_average_blpairs(self): Test bootstrap averaging over power spectra. """ # Check that basic bootstrap averaging works - blpair_groups = [list(np.unique(self.uvp.blpair_array)),] - uvp1, wgts = grouping.bootstrap_average_blpairs([self.uvp,], - blpair_groups, - time_avg=False) - uvp2, wgts = grouping.bootstrap_average_blpairs([self.uvp,], - blpair_groups, - time_avg=True) + blpair_groups = [ + list(np.unique(self.uvp.blpair_array)), + ] + uvp1, wgts = grouping.bootstrap_average_blpairs( + [ + self.uvp, + ], + blpair_groups, + time_avg=False, + ) + uvp2, wgts = grouping.bootstrap_average_blpairs( + [ + self.uvp, + ], + blpair_groups, + time_avg=True, + ) assert uvp1[0].Nblpairs == 1 assert uvp1[0].Ntimes == self.uvp.Ntimes assert uvp2[0].Ntimes == 1 @@ -299,43 +341,73 @@ def test_bootstrap_average_blpairs(self): assert np.sum(wgts) == np.array(blpair_groups).size # Check that exceptions are raised when inputs are invalid - pytest.raises(AssertionError, grouping.bootstrap_average_blpairs, - [np.arange(5),], blpair_groups, time_avg=False) - pytest.raises(KeyError, grouping.bootstrap_average_blpairs, - [self.uvp,], [[200200200200,],], time_avg=False) + pytest.raises( + AssertionError, + grouping.bootstrap_average_blpairs, + [ + np.arange(5), + ], + blpair_groups, + time_avg=False, + ) + pytest.raises( + KeyError, + grouping.bootstrap_average_blpairs, + [ + self.uvp, + ], + [ + [ + 200200200200, + ], + ], + time_avg=False, + ) # Reduce UVPSpec to only 3 blpairs and set them all to the same values _blpairs = list(np.unique(self.uvp.blpair_array)[:3]) uvp3 = self.uvp.select(spws=0, inplace=False, blpairs=_blpairs) Nt = uvp3.Ntimes - uvp3.data_array[0][Nt:2*Nt] = uvp3.data_array[0][:Nt] - uvp3.data_array[0][2*Nt:] = uvp3.data_array[0][:Nt] - uvp3.integration_array[0][Nt:2*Nt] = uvp3.integration_array[0][:Nt] - uvp3.integration_array[0][2*Nt:] = uvp3.integration_array[0][:Nt] + uvp3.data_array[0][Nt : 2 * Nt] = uvp3.data_array[0][:Nt] + uvp3.data_array[0][2 * Nt :] = uvp3.data_array[0][:Nt] + uvp3.integration_array[0][Nt : 2 * Nt] = uvp3.integration_array[0][:Nt] + uvp3.integration_array[0][2 * Nt :] = uvp3.integration_array[0][:Nt] # Test that different bootstrap-sampled averages have the same value as # the normal average (since the data for all blpairs has been set to # the same values for uvp3) np.random.seed(10) - uvp_avg = uvp3.average_spectra(blpair_groups=[_blpairs,], - time_avg=True, inplace=False) + uvp_avg = uvp3.average_spectra( + blpair_groups=[ + _blpairs, + ], + time_avg=True, + inplace=False, + ) blpair = uvp_avg.blpair_array[0] for i in range(5): # Generate multiple samples and make sure that they are all equal # to the regular average (for the cloned data in uvp3) uvp4, wgts = grouping.bootstrap_average_blpairs( - [uvp3,], - blpair_groups=[_blpairs,], - time_avg=True) + [ + uvp3, + ], + blpair_groups=[ + _blpairs, + ], + time_avg=True, + ) try: - ps_avg = uvp_avg.get_data((0, blpair, ('xx','xx'))) - except: + ps_avg = uvp_avg.get_data((0, blpair, ("xx", "xx"))) + except Exception as e: print(uvp_avg.polpair_array) - raise - ps_boot = uvp4[0].get_data((0, blpair, ('xx','xx'))) + raise e + + ps_boot = uvp4[0].get_data((0, blpair, ("xx", "xx"))) np.testing.assert_array_almost_equal(ps_avg, ps_boot) + def test_bootstrap_resampled_error(): # generate a UVPSpec visfile = os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA") @@ -346,14 +418,18 @@ def test_bootstrap_resampled_error(): uvd.read_miriad(visfile) ap, a = uvd.get_ENU_antpos(pick_data_ants=True) reds = redcal.get_pos_reds(dict(zip(a, ap)), bl_error_tol=1.0)[:3] - uvp = testing.uvpspec_from_data(uvd, reds, spw_ranges=[(50, 100)], beam=beam, cosmo=cosmo) + uvp = testing.uvpspec_from_data( + uvd, reds, spw_ranges=[(50, 100)], beam=beam, cosmo=cosmo + ) # Lots of this function is already tested by bootstrap_run # so only test the stuff not already tested if os.path.exists("uvp.h5"): os.remove("uvp.h5") uvp.write_hdf5("uvp.h5", overwrite=True) - ua, ub, uw = grouping.bootstrap_resampled_error("uvp.h5", blpair_groups=None, Nsamples=10, seed=0, verbose=False) + ua, ub, uw = grouping.bootstrap_resampled_error( + "uvp.h5", blpair_groups=None, Nsamples=10, seed=0, verbose=False + ) # check number of boots assert len(ub) == 10 @@ -381,31 +457,58 @@ def test_validate_bootstrap_errorbar(): # generate complex gaussian noise seed = 4 - uvd1 = testing.noise_sim(uvfile, Tsys, seed=seed, whiten=True, inplace=False, Nextend=0) + uvd1 = testing.noise_sim( + uvfile, Tsys, seed=seed, whiten=True, inplace=False, Nextend=0 + ) seed = 5 - uvd2 = testing.noise_sim(uvfile, Tsys, seed=seed, whiten=True, inplace=False, Nextend=0) + uvd2 = testing.noise_sim( + uvfile, Tsys, seed=seed, whiten=True, inplace=False, Nextend=0 + ) # form (auto) baseline-pairs from only 14.6m bls - reds, lens, angs = utils.get_reds(uvd1, pick_data_ants=True, bl_len_range=(10, 50), - bl_deg_range=(0, 180)) + reds, lens, angs = utils.get_reds( + uvd1, pick_data_ants=True, bl_len_range=(10, 50), bl_deg_range=(0, 180) + ) bls1, bls2 = utils.flatten(reds), utils.flatten(reds) # setup PSpecData and form power psectra - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd2)], wgts=[None, None]) - uvp = ds.pspec(bls1, bls2, (0, 1), [('xx', 'xx')], input_data_weight='identity', norm='I', - taper='none', sampling=False, little_h=False, spw_ranges=[(0, 50)], verbose=False) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd2)], wgts=[None, None] + ) + uvp = ds.pspec( + bls1, + bls2, + (0, 1), + [("xx", "xx")], + input_data_weight="identity", + norm="I", + taper="none", + sampling=False, + little_h=False, + spw_ranges=[(0, 50)], + verbose=False, + ) # bootstrap resample Nsamples = 100 seed = 0 - uvp_avg, uvp_boots, uvp_wgts = grouping.bootstrap_resampled_error(uvp, time_avg=False, Nsamples=Nsamples, - seed=seed, normal_std=True, - blpair_groups=[uvp.get_blpairs()]) + uvp_avg, uvp_boots, uvp_wgts = grouping.bootstrap_resampled_error( + uvp, + time_avg=False, + Nsamples=Nsamples, + seed=seed, + normal_std=True, + blpair_groups=[uvp.get_blpairs()], + ) # assert z-score has std of ~1.0 along time ax to within 1/sqrt(Nsamples) - zscr_real = np.std(uvp_avg.data_array[0].real / uvp_avg.stats_array['bs_std'][0].real) - zscr_imag = np.std(uvp_avg.data_array[0].imag / uvp_avg.stats_array['bs_std'][0].imag) - assert np.abs(1.0 - zscr_real) < 1/np.sqrt(Nsamples) - assert np.abs(1.0 - zscr_imag) < 1/np.sqrt(Nsamples) + zscr_real = np.std( + uvp_avg.data_array[0].real / uvp_avg.stats_array["bs_std"][0].real + ) + zscr_imag = np.std( + uvp_avg.data_array[0].imag / uvp_avg.stats_array["bs_std"][0].imag + ) + assert np.abs(1.0 - zscr_real) < 1 / np.sqrt(Nsamples) + assert np.abs(1.0 - zscr_imag) < 1 / np.sqrt(Nsamples) def test_bootstrap_run(): @@ -418,23 +521,36 @@ def test_bootstrap_run(): uvd.read_miriad(visfile) ap, a = uvd.get_ENU_antpos(pick_data_ants=True) reds = redcal.get_pos_reds(dict(zip(a, ap)), bl_error_tol=1.0)[:3] - uvp = testing.uvpspec_from_data(uvd, reds, spw_ranges=[(50, 100)], beam=beam, cosmo=cosmo) + uvp = testing.uvpspec_from_data( + uvd, reds, spw_ranges=[(50, 100)], beam=beam, cosmo=cosmo + ) if os.path.exists("ex.h5"): os.remove("ex.h5") - psc = container.PSpecContainer("ex.h5", mode='rw', keep_open=False, swmr=False) + psc = container.PSpecContainer("ex.h5", mode="rw", keep_open=False, swmr=False) psc.set_pspec("grp1", "uvp", uvp) # Test basic bootstrap run - grouping.bootstrap_run(psc, time_avg=True, Nsamples=100, seed=0, - normal_std=True, robust_std=True, cintervals=[16, 84], keep_samples=True, - bl_error_tol=1.0, overwrite=True, add_to_history='hello!', verbose=False) + grouping.bootstrap_run( + psc, + time_avg=True, + Nsamples=100, + seed=0, + normal_std=True, + robust_std=True, + cintervals=[16, 84], + keep_samples=True, + bl_error_tol=1.0, + overwrite=True, + add_to_history="hello!", + verbose=False, + ) spcs = psc.spectra("grp1") # assert all bs samples were written - assert (np.all(["uvp_bs{}".format(i) in spcs for i in range(100)])) + assert np.all(["uvp_bs{}".format(i) in spcs for i in range(100)]) # assert average was written - assert ("uvp_avg" in spcs and "uvp" in spcs) + assert "uvp_avg" in spcs and "uvp" in spcs # assert average only has one time and 3 blpairs uvp_avg = psc.get_pspec("grp1", "uvp_avg") @@ -442,24 +558,38 @@ def test_bootstrap_run(): assert uvp_avg.Nblpairs == 3 # check avg file history - assert ("hello!" in uvp_avg.history) + assert "hello!" in uvp_avg.history # assert original uvp is unchanged - assert uvp == psc.get_pspec("grp1", 'uvp') + assert uvp == psc.get_pspec("grp1", "uvp") # check stats array - np.testing.assert_array_equal([u'bs_cinterval_16.00', u'bs_cinterval_84.00', u'bs_robust_std', u'bs_std'], list(uvp_avg.stats_array.keys())) - - for stat in [u'bs_cinterval_16.00', u'bs_cinterval_84.00', u'bs_robust_std', u'bs_std']: - assert uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), ('xx','xx'))).shape == (1, 50) - assert (np.any(np.isnan(uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), ('xx','xx')))))) == False - assert uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), ('xx','xx'))).dtype == np.complex128 + np.testing.assert_array_equal( + ["bs_cinterval_16.00", "bs_cinterval_84.00", "bs_robust_std", "bs_std"], + list(uvp_avg.stats_array.keys()), + ) + + for stat in ["bs_cinterval_16.00", "bs_cinterval_84.00", "bs_robust_std", "bs_std"]: + assert uvp_avg.get_stats( + stat, (0, ((37, 38), (38, 39)), ("xx", "xx")) + ).shape == (1, 50) + assert not ( + np.any( + np.isnan( + uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), ("xx", "xx"))) + ) + ) + ) + assert ( + uvp_avg.get_stats(stat, (0, ((37, 38), (38, 39)), ("xx", "xx"))).dtype + == np.complex128 + ) # test exceptions del psc if os.path.exists("ex.h5"): os.remove("ex.h5") - psc = container.PSpecContainer("ex.h5", mode='rw', keep_open=False, swmr=False) + psc = container.PSpecContainer("ex.h5", mode="rw", keep_open=False, swmr=False) # test empty groups pytest.raises(AssertionError, grouping.bootstrap_run, "ex.h5") @@ -469,11 +599,11 @@ def test_bootstrap_run(): # test fed spectra doesn't exist psc.set_pspec("grp1", "uvp", uvp) - pytest.raises(AssertionError, grouping.bootstrap_run, psc, spectra=['grp1/foo']) + pytest.raises(AssertionError, grouping.bootstrap_run, psc, spectra=["grp1/foo"]) # test assertionerror if SWMR - psc = container.PSpecContainer("ex.h5", mode='rw', keep_open=False, swmr=True) - pytest.raises(AssertionError, grouping.bootstrap_run, psc, spectra=['grp1/foo']) + psc = container.PSpecContainer("ex.h5", mode="rw", keep_open=False, swmr=True) + pytest.raises(AssertionError, grouping.bootstrap_run, psc, spectra=["grp1/foo"]) if os.path.exists("ex.h5"): os.remove("ex.h5") @@ -481,10 +611,25 @@ def test_bootstrap_run(): def test_get_bootstrap_run_argparser(): args = grouping.get_bootstrap_run_argparser() - a = args.parse_args(['fname', '--spectra', 'grp1/uvp1', 'grp1/uvp2', 'grp2/uvp1', - '--blpair_groups', '101102103104 101102102103, 102103104105', - '--time_avg', 'True', '--Nsamples', '100', '--cintervals', '16', '84']) - assert a.spectra == ['grp1/uvp1', 'grp1/uvp2', 'grp2/uvp1'] + a = args.parse_args( + [ + "fname", + "--spectra", + "grp1/uvp1", + "grp1/uvp2", + "grp2/uvp1", + "--blpair_groups", + "101102103104 101102102103, 102103104105", + "--time_avg", + "True", + "--Nsamples", + "100", + "--cintervals", + "16", + "84", + ] + ) + assert a.spectra == ["grp1/uvp1", "grp1/uvp2", "grp2/uvp1"] assert a.blpair_groups == [[101102103104, 101102102103], [102103104105]] assert a.cintervals == [16.0, 84.0] @@ -492,7 +637,7 @@ def test_get_bootstrap_run_argparser(): def test_spherical_average(): # create two polarization data uvd = UVData() - uvd.read(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) + uvd.read(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) # load other data, get reds and make UVPSpec beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") @@ -501,43 +646,80 @@ def test_spherical_average(): ap, a = uvd.get_ENU_antpos(pick_data_ants=True) reds = redcal.get_pos_reds(dict(zip(a, ap)), bl_error_tol=1.0) reds = [r[:2] for r in reds] - uvp = testing.uvpspec_from_data(uvd, reds, spw_ranges=[(50, 75), (100, 125)], beam=beam, cosmo=cosmo) + uvp = testing.uvpspec_from_data( + uvd, reds, spw_ranges=[(50, 75), (100, 125)], beam=beam, cosmo=cosmo + ) uvd.polarization_array[0] = -6 - uvp += testing.uvpspec_from_data(uvd, reds, spw_ranges=[(50, 75), (100, 125)], beam=beam, cosmo=cosmo) + uvp += testing.uvpspec_from_data( + uvd, reds, spw_ranges=[(50, 75), (100, 125)], beam=beam, cosmo=cosmo + ) # insert cov_array and stats_array - uvp.cov_model = 'empirical' - uvp.cov_array_real = {s: np.repeat(np.repeat(np.eye(uvp.Ndlys, dtype=np.float64)[None, : , :, None], uvp.Nblpairts, 0), uvp.Npols, -1) - for s in range(uvp.Nspws)} - uvp.cov_array_imag = {s: np.repeat(np.repeat(np.eye(uvp.Ndlys, dtype=np.float64)[None, : , :, None], uvp.Nblpairts, 0), uvp.Npols, -1) - for s in range(uvp.Nspws)} - uvp.stats_array = {'err': {s: np.ones((uvp.Nblpairts, uvp.Ndlys, uvp.Npols), dtype=np.complex128) - for s in range(uvp.Nspws)}} + uvp.cov_model = "empirical" + uvp.cov_array_real = { + s: np.repeat( + np.repeat( + np.eye(uvp.Ndlys, dtype=np.float64)[None, :, :, None], uvp.Nblpairts, 0 + ), + uvp.Npols, + -1, + ) + for s in range(uvp.Nspws) + } + uvp.cov_array_imag = { + s: np.repeat( + np.repeat( + np.eye(uvp.Ndlys, dtype=np.float64)[None, :, :, None], uvp.Nblpairts, 0 + ), + uvp.Npols, + -1, + ) + for s in range(uvp.Nspws) + } + uvp.stats_array = { + "err": { + s: np.ones((uvp.Nblpairts, uvp.Ndlys, uvp.Npols), dtype=np.complex128) + for s in range(uvp.Nspws) + } + } # try a spherical average kbins = np.arange(0, 2.9, 0.25) Nk = len(kbins) bin_widths = 0.25 A = {} - sph = grouping.spherical_average(uvp, kbins, bin_widths, add_to_history='checking 1 2 3', A=A) + sph = grouping.spherical_average( + uvp, kbins, bin_widths, add_to_history="checking 1 2 3", A=A + ) # metadata assert sph.Nblpairs == 1 - assert 'checking 1 2 3' in sph.history + assert "checking 1 2 3" in sph.history assert np.isclose(sph.get_blpair_seps(), 0).all() # assert kperp has no magnitude - assert 'err' in sph.stats_array + assert "err" in sph.stats_array for spw in sph.spw_array: # binning and normalization - assert np.isclose(sph.get_kparas(spw), kbins).all() # assert kbins are input kbins - assert np.isclose(sph.window_function_array[spw].sum(axis=2), 1).all() # assert window func is normalized + assert np.isclose( + sph.get_kparas(spw), kbins + ).all() # assert kbins are input kbins + assert np.isclose( + sph.window_function_array[spw].sum(axis=2), 1 + ).all() # assert window func is normalized # check low k modes are greater than high k modes # this is a basic "averaged data smell test" in lieu of a known pspec to compare to - assert np.all(sph.data_array[spw][:, 0, :].real / sph.data_array[spw][:, 10, :] > 1e3) + assert np.all( + sph.data_array[spw][:, 0, :].real / sph.data_array[spw][:, 10, :] > 1e3 + ) # assert errorbars are 1/sqrt(N) what they used to be - assert np.isclose(np.sqrt(sph.cov_array_real[spw])[:, range(Nk), range(Nk)], 1/np.sqrt(A[spw].sum(axis=1))).all() - assert np.isclose(sph.stats_array['err'][spw], 1/np.sqrt(A[spw].sum(axis=1))).all() + assert np.isclose( + np.sqrt(sph.cov_array_real[spw])[:, range(Nk), range(Nk)], + 1 / np.sqrt(A[spw].sum(axis=1)), + ).all() + assert np.isclose( + sph.stats_array["err"][spw], 1 / np.sqrt(A[spw].sum(axis=1)) + ).all() # bug check: time_avg_array was not down-selected to new Nblpairts assert sph.time_avg_array.size == sph.Nblpairts @@ -546,7 +728,9 @@ def test_spherical_average(): assert sph.cov_array_real[0].shape == sph.cov_array_imag[0].shape # try without little h - sph2 = grouping.spherical_average(uvp, kbins * cosmo.h, bin_widths * cosmo.h, little_h=False) + sph2 = grouping.spherical_average( + uvp, kbins * cosmo.h, bin_widths * cosmo.h, little_h=False + ) for spw in sph.spw_array: assert np.isclose(sph.get_kparas(spw), sph2.get_kparas(spw)).all() @@ -555,7 +739,7 @@ def test_spherical_average(): assert sph.Ntimes == 1 # try weighting by stats_array - sph = grouping.spherical_average(uvp, kbins, bin_widths, error_weights='err') + sph = grouping.spherical_average(uvp, kbins, bin_widths, error_weights="err") for spw in sph.spw_array: assert np.isclose(sph.window_function_array[spw].sum(axis=2), 1).all() @@ -567,21 +751,25 @@ def test_spherical_average(): # slice into stats array and set region of k_perp k_para to infinte variance uvp2 = copy.deepcopy(uvp) - uvp2.set_stats_slice('err', 0, 1000, above=False, val=np.inf) - sph2 = grouping.spherical_average(uvp2, kbins, bin_widths, error_weights='err') + uvp2.set_stats_slice("err", 0, 1000, above=False, val=np.inf) + sph2 = grouping.spherical_average(uvp2, kbins, bin_widths, error_weights="err") # assert low k modes are zeroed! assert np.isclose(sph2.data_array[0][:, :3, :], 0).all() # assert bins that weren't nulled still have proper window normalization for spw in sph2.spw_array: - assert np.isclose(sph2.window_function_array[spw].sum(axis=2)[:, 3:, :], 1).all() + assert np.isclose( + sph2.window_function_array[spw].sum(axis=2)[:, 3:, :], 1 + ).all() # assert resultant stats are not nan - assert (~np.isnan(sph2.stats_array['err'][0])).all() + assert (~np.isnan(sph2.stats_array["err"][0])).all() # try combine after spherical and assert it is equivalent - sph_a, sph_b = sph.select(spws=[0], inplace=False), sph.select(spws=[1], inplace=False) + sph_a, sph_b = sph.select(spws=[0], inplace=False), sph.select( + spws=[1], inplace=False + ) sph_c = uvpspec.combine_uvpspec([sph_a, sph_b], merge_history=False, verbose=False) # bug check: in the past, combine after spherical average erroneously changed dly_array assert sph == sph_c @@ -589,7 +777,7 @@ def test_spherical_average(): # insert an inf into the arrays as a test uvp2 = copy.deepcopy(uvp) uvp2.cov_array_real[0][0], uvp2.cov_array_imag[0][0] = np.inf, np.inf - uvp2.stats_array['err'][0][0] = np.inf + uvp2.stats_array["err"][0][0] = np.inf sph = grouping.spherical_average(uvp, kbins, bin_widths) assert np.isfinite(sph.cov_array_real[0]).all() @@ -598,93 +786,138 @@ def test_spherical_average(): # tests related to exact_windows uvd = UVData() - uvd.read_uvh5(os.path.join(DATA_PATH, 'zen.2458116.31939.HH.uvh5')) + uvd.read_uvh5(os.path.join(DATA_PATH, "zen.2458116.31939.HH.uvh5")) ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=beam) - baselines1, baselines2, blpairs = utils.construct_blpairs(uvd.get_antpairs()[1:], - exclude_permutations=False, - exclude_auto_bls=True) + baselines1, baselines2, blpairs = utils.construct_blpairs( + uvd.get_antpairs()[1:], exclude_permutations=False, exclude_auto_bls=True + ) # compute ps - uvp = ds.pspec(baselines1, baselines2, dsets=(0, 1), pols=[('xx', 'xx')], - spw_ranges=(175,195), taper='bh',verbose=False) + uvp = ds.pspec( + baselines1, + baselines2, + dsets=(0, 1), + pols=[("xx", "xx")], + spw_ranges=(175, 195), + taper="bh", + verbose=False, + ) # get exact window functions - uvp.get_exact_window_functions(ftbeam_file = os.path.join(DATA_PATH, 'FT_beam_HERA_dipole_test'), - spw_array=None, inplace=True, verbose=False) + uvp.get_exact_window_functions( + ftbeam_file=os.path.join(DATA_PATH, "FT_beam_HERA_dipole_test"), + spw_array=None, + inplace=True, + verbose=False, + ) # spherical window functions for redundant groups sph = grouping.spherical_average(uvp, kbins, bin_widths) # spherical average for input blpair groups blpair_groups, blpair_lens, _ = uvp.get_red_blpairs() - sph2 = grouping.spherical_average(uvp, kbins, bin_widths, - blpair_groups=blpair_groups) + sph2 = grouping.spherical_average( + uvp, kbins, bin_widths, blpair_groups=blpair_groups + ) + def test_spherical_wf_from_uvp(): # parameters kbins = np.arange(0, 2.9, 0.25) - Nk = len(kbins) bin_widths = 0.25 - basename = 'FT_beam_HERA_dipole_test' + basename = "FT_beam_HERA_dipole_test" # obtain uvp object uvd = UVData() - uvd.read_uvh5(os.path.join(DATA_PATH, 'zen.2458116.31939.HH.uvh5')) + uvd.read_uvh5(os.path.join(DATA_PATH, "zen.2458116.31939.HH.uvh5")) # Create a new PSpecData objec ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None]) # choose baselines - baselines1, baselines2, blpairs = utils.construct_blpairs(uvd.get_antpairs()[1:], - exclude_permutations=False, - exclude_auto_bls=True) + baselines1, baselines2, blpairs = utils.construct_blpairs( + uvd.get_antpairs()[1:], exclude_permutations=False, exclude_auto_bls=True + ) # compute ps - uvp = ds.pspec(baselines1, baselines2, dsets=(0, 1), pols=[('xx','xx')], - spw_ranges=(175,195), taper='bh',verbose=False) - uvp.cosmo = conversions.Cosmo_Conversions() #uvp.set_cosmology not overwriting + uvp = ds.pspec( + baselines1, + baselines2, + dsets=(0, 1), + pols=[("xx", "xx")], + spw_ranges=(175, 195), + taper="bh", + verbose=False, + ) + uvp.cosmo = conversions.Cosmo_Conversions() # uvp.set_cosmology not overwriting # obtain exact_windows (fiducial usage) - uvp.get_exact_window_functions(ftbeam_file = os.path.join(DATA_PATH, basename), - inplace=True) + uvp.get_exact_window_functions( + ftbeam_file=os.path.join(DATA_PATH, basename), inplace=True + ) wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths) - # blpair_groups=blpair_groups, - # blpair_lens=blpair_lens, - # blpair_weights=blpair_weights, - # time_avg=time_avg, - # error_weights=error_weights, - # spw_array=spw, - # little_h=little_h, - # verbose=True)[spw] + # blpair_groups=blpair_groups, + # blpair_lens=blpair_lens, + # blpair_weights=blpair_weights, + # time_avg=time_avg, + # error_weights=error_weights, + # spw_array=spw, + # little_h=little_h, + # verbose=True)[spw] # test different options # little_h wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths, little_h=False) # spw_array wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths, spw_array=0) - pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, bin_widths, - spw_array=2) - # blpair_groups + pytest.raises( + AssertionError, + grouping.spherical_wf_from_uvp, + uvp, + kbins, + bin_widths, + spw_array=2, + ) + # blpair_groups blpair_groups, blpair_lens, _ = uvp.get_red_blpairs() - wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths, - blpair_groups=blpair_groups, - blpair_lens=blpair_lens) - pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, bin_widths, - blpair_groups=blpair_groups[0]) + wf_array = grouping.spherical_wf_from_uvp( + uvp, kbins, bin_widths, blpair_groups=blpair_groups, blpair_lens=blpair_lens + ) + pytest.raises( + AssertionError, + grouping.spherical_wf_from_uvp, + uvp, + kbins, + bin_widths, + blpair_groups=blpair_groups[0], + ) # raise warning or error if blpair_groups inconsistent with blpair_lens - wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths, - blpair_groups=None, - blpair_lens=blpair_lens) - wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths, - blpair_groups=blpair_groups, - blpair_lens=None) - pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, bin_widths, - blpair_groups=blpair_groups, blpair_lens=[blpair_lens[0],blpair_lens[0]]) + wf_array = grouping.spherical_wf_from_uvp( + uvp, kbins, bin_widths, blpair_groups=None, blpair_lens=blpair_lens + ) + wf_array = grouping.spherical_wf_from_uvp( + uvp, kbins, bin_widths, blpair_groups=blpair_groups, blpair_lens=None + ) + pytest.raises( + AssertionError, + grouping.spherical_wf_from_uvp, + uvp, + kbins, + bin_widths, + blpair_groups=blpair_groups, + blpair_lens=[blpair_lens[0], blpair_lens[0]], + ) # error if overlapping bins pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, 1.0) # blpair_weights - wf_array = grouping.spherical_wf_from_uvp(uvp, kbins, bin_widths, - blpair_groups=blpair_groups, - blpair_lens=blpair_lens, - blpair_weights=[[1. for item in grp] for grp in blpair_groups]) + grouping.spherical_wf_from_uvp( + uvp, + kbins, + bin_widths, + blpair_groups=blpair_groups, + blpair_lens=blpair_lens, + blpair_weights=[[1.0 for item in grp] for grp in blpair_groups], + ) # raise error if uvp.exact_windows is False uvp.exact_windows = False - pytest.raises(AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, bin_widths) + pytest.raises( + AssertionError, grouping.spherical_wf_from_uvp, uvp, kbins, bin_widths + ) if __name__ == "__main__": diff --git a/hera_pspec/tests/test_noise.py b/hera_pspec/tests/test_noise.py index 0f5b4ad0..a39903ab 100644 --- a/hera_pspec/tests/test_noise.py +++ b/hera_pspec/tests/test_noise.py @@ -16,10 +16,12 @@ class Test_Sensitivity(unittest.TestCase): """ Test noise.Sensitivity object """ + def setUp(self): self.cosmo = conversions.Cosmo_Conversions() - self.beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, - 'HERA_NF_pstokes_power.beamfits')) + self.beam = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_pstokes_power.beamfits") + ) self.sense = noise.Sensitivity(beam=self.beam, cosmo=self.cosmo) def tearDown(self): @@ -45,32 +47,32 @@ def test_set(self): assert sense.cosmo.get_params() == sense.beam.cosmo.get_params() bm = copy.deepcopy(self.beam) - delattr(bm, 'cosmo') + delattr(bm, "cosmo") sense.set_beam(bm) def test_scalar(self): freqs = np.linspace(150e6, 160e6, 100, endpoint=False) - self.sense.calc_scalar(freqs, 'pI', num_steps=5000, little_h=True) + self.sense.calc_scalar(freqs, "pI", num_steps=5000, little_h=True) assert np.isclose(freqs, self.sense.subband).all() - assert self.sense.pol == 'pI' + assert self.sense.pol == "pI" def test_calc_P_N(self): # calculate scalar freqs = np.linspace(150e6, 160e6, 100, endpoint=False) - self.sense.calc_scalar(freqs, 'pI', num_steps=5000, little_h=True) + self.sense.calc_scalar(freqs, "pI", num_steps=5000, little_h=True) # basic execution k = np.linspace(0, 3, 10) Tsys = 500.0 t_int = 10.7 - P_N = self.sense.calc_P_N(Tsys, t_int, Ncoherent=1, Nincoherent=1, - form='Pk') + P_N = self.sense.calc_P_N(Tsys, t_int, Ncoherent=1, Nincoherent=1, form="Pk") assert isinstance(P_N, (float, np.float)) assert np.isclose(P_N, 642386932892.2921) # calculate DelSq - Dsq = self.sense.calc_P_N(Tsys, t_int, k=k, Ncoherent=1, - Nincoherent=1, form='DelSq') + Dsq = self.sense.calc_P_N( + Tsys, t_int, k=k, Ncoherent=1, Nincoherent=1, form="DelSq" + ) assert Dsq.shape == (10,) assert Dsq[1] < P_N @@ -82,40 +84,63 @@ def test_noise_validation(): noise simulation. """ # get simulated noise in Jy - bfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + bfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") beam = pspecbeam.PSpecBeamUV(bfile) uvfile = os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA") Tsys = 300.0 # Kelvin # generate noise seed = 0 - uvd = testing.noise_sim(uvfile, Tsys, beam, seed=seed, whiten=True, - inplace=False, Nextend=9) + uvd = testing.noise_sim( + uvfile, Tsys, beam, seed=seed, whiten=True, inplace=False, Nextend=9 + ) # get redundant baseline group - reds, lens, angs = utils.get_reds(uvd, pick_data_ants=True, - bl_len_range=(10, 20), - bl_deg_range=(0, 1)) - bls1, bls2, blps = utils.construct_blpairs(reds[0], exclude_auto_bls=True, - exclude_permutations=True) + reds, lens, angs = utils.get_reds( + uvd, pick_data_ants=True, bl_len_range=(10, 20), bl_deg_range=(0, 1) + ) + bls1, bls2, blps = utils.construct_blpairs( + reds[0], exclude_auto_bls=True, exclude_permutations=True + ) # setup PSpecData - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], - wgts=[None, None], beam=beam) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None], beam=beam + ) ds.Jy_to_mK() # get pspec - uvp = ds.pspec(bls1, bls2, (0, 1), [('xx', 'xx')], input_data_weight='identity', norm='I', - taper='none', sampling=False, little_h=True, spw_ranges=[(0, 50)], verbose=False) + uvp = ds.pspec( + bls1, + bls2, + (0, 1), + [("xx", "xx")], + input_data_weight="identity", + norm="I", + taper="none", + sampling=False, + little_h=True, + spw_ranges=[(0, 50)], + verbose=False, + ) # get noise spectra from one of the blpairs - P_N = list(uvp.generate_noise_spectra(0, ('xx','xx'), Tsys, - blpairs=uvp.get_blpairs()[:1], num_steps=2000, - component='real').values())[0][0, 0] + P_N = list( + uvp.generate_noise_spectra( + 0, + ("xx", "xx"), + Tsys, + blpairs=uvp.get_blpairs()[:1], + num_steps=2000, + component="real", + ).values() + )[0][0, 0] # get P_rms of real spectra for each baseline across time axis - Pspec = np.array([uvp.get_data((0, bl, ('xx', 'xx'))).real for bl in uvp.get_blpairs()]) - P_rms = np.sqrt(np.mean(np.abs(Pspec)**2)) + Pspec = np.array( + [uvp.get_data((0, bl, ("xx", "xx"))).real for bl in uvp.get_blpairs()] + ) + P_rms = np.sqrt(np.mean(np.abs(Pspec) ** 2)) # assert close to P_N: 2% # This should be updated to be within standard error on P_rms @@ -129,53 +154,87 @@ def test_analytic_noise(): one using QE propagated from auto-correlation second using P_N from Tsys estimate """ - bfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + bfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") beam = pspecbeam.PSpecBeamUV(bfile) uvfile = os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA") uvd = UVData() uvd.read(uvfile) # setup PSpecData - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], - wgts=[None, None], beam=beam, - dsets_std=[copy.deepcopy(uvd), copy.deepcopy(uvd)]) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], + wgts=[None, None], + beam=beam, + dsets_std=[copy.deepcopy(uvd), copy.deepcopy(uvd)], + ) ds.Jy_to_mK() # get pspec - reds, lens, angs = utils.get_reds(uvd, pick_data_ants=True, - bl_len_range=(10, 20), - bl_deg_range=(0, 1)) - bls1, bls2, blps = utils.construct_blpairs(reds[0], exclude_auto_bls=True, - exclude_permutations=True) - taper = 'bh' + reds, lens, angs = utils.get_reds( + uvd, pick_data_ants=True, bl_len_range=(10, 20), bl_deg_range=(0, 1) + ) + bls1, bls2, blps = utils.construct_blpairs( + reds[0], exclude_auto_bls=True, exclude_permutations=True + ) + taper = "bh" Nchan = 20 - uvp = ds.pspec(bls1, bls2, (0, 1), [('xx', 'xx')], input_data_weight='identity', norm='I', - taper=taper, sampling=False, little_h=True, spw_ranges=[(0, Nchan)], verbose=False, - cov_model='autos', store_cov=True) - uvp_fg = ds.pspec(bls1, bls2, (0, 1), [('xx', 'xx')], input_data_weight='identity', norm='I', - taper=taper, sampling=False, little_h=True, spw_ranges=[(0, Nchan)], verbose=False, - cov_model='foreground_dependent', store_cov=True) + uvp = ds.pspec( + bls1, + bls2, + (0, 1), + [("xx", "xx")], + input_data_weight="identity", + norm="I", + taper=taper, + sampling=False, + little_h=True, + spw_ranges=[(0, Nchan)], + verbose=False, + cov_model="autos", + store_cov=True, + ) + uvp_fg = ds.pspec( + bls1, + bls2, + (0, 1), + [("xx", "xx")], + input_data_weight="identity", + norm="I", + taper=taper, + sampling=False, + little_h=True, + spw_ranges=[(0, Nchan)], + verbose=False, + cov_model="foreground_dependent", + store_cov=True, + ) # get P_N estimate auto_Tsys = utils.uvd_to_Tsys(uvd, beam, os.path.join(DATA_PATH, "test_uvd.uvh5")) assert os.path.exists(os.path.join(DATA_PATH, "test_uvd.uvh5")) - utils.uvp_noise_error(uvp, auto_Tsys, err_type=['P_N','P_SN'], P_SN_correction=False) + utils.uvp_noise_error( + uvp, auto_Tsys, err_type=["P_N", "P_SN"], P_SN_correction=False + ) # check consistency of 1-sigma standard dev. to 1% cov_diag = uvp.cov_array_real[0][:, range(Nchan), range(Nchan)] - stats_diag = uvp.stats_array['P_N'][0] + stats_diag = uvp.stats_array["P_N"][0] frac_ratio = (cov_diag**0.5 - stats_diag) / stats_diag assert np.abs(frac_ratio).mean() < 0.01 ## check P_SN consistency of 1-sigma standard dev. to 1% cov_diag = uvp_fg.cov_array_real[0][:, range(Nchan), range(Nchan)] - stats_diag = uvp.stats_array['P_SN'][0] + stats_diag = uvp.stats_array["P_SN"][0] frac_ratio = (cov_diag**0.5 - stats_diag) / stats_diag assert np.abs(frac_ratio).mean() < 0.01 # now compute unbiased P_SN and check that it matches P_N at high-k - utils.uvp_noise_error(uvp, auto_Tsys, err_type=['P_N','P_SN'], P_SN_correction=True) - frac_ratio = (uvp.stats_array["P_SN"][0] - uvp.stats_array["P_N"][0]) / uvp.stats_array["P_N"][0] + utils.uvp_noise_error( + uvp, auto_Tsys, err_type=["P_N", "P_SN"], P_SN_correction=True + ) + frac_ratio = ( + uvp.stats_array["P_SN"][0] - uvp.stats_array["P_N"][0] + ) / uvp.stats_array["P_N"][0] dlys = uvp.get_dlys(0) * 1e9 select = np.abs(dlys) > 3000 assert np.abs(frac_ratio[:, select].mean()) < 1 / np.sqrt(uvp.Nblpairts) @@ -183,8 +242,12 @@ def test_analytic_noise(): # test single time uvp.select(times=uvp.time_avg_array[:1], inplace=True) auto_Tsys.select(times=auto_Tsys.time_array[:1], inplace=True) - utils.uvp_noise_error(uvp, auto_Tsys, err_type=['P_N','P_SN'], P_SN_correction=True) - frac_ratio = (uvp.stats_array["P_SN"][0] - uvp.stats_array["P_N"][0]) / uvp.stats_array["P_N"][0] + utils.uvp_noise_error( + uvp, auto_Tsys, err_type=["P_N", "P_SN"], P_SN_correction=True + ) + frac_ratio = ( + uvp.stats_array["P_SN"][0] - uvp.stats_array["P_N"][0] + ) / uvp.stats_array["P_N"][0] dlys = uvp.get_dlys(0) * 1e9 select = np.abs(dlys) > 3000 assert np.abs(frac_ratio[:, select].mean()) < 1 / np.sqrt(uvp.Nblpairts) diff --git a/hera_pspec/tests/test_plot.py b/hera_pspec/tests/test_plot.py index 9e90a6f6..e424a830 100644 --- a/hera_pspec/tests/test_plot.py +++ b/hera_pspec/tests/test_plot.py @@ -11,9 +11,10 @@ # Data files to use in tests dfiles = [ - 'zen.all.xx.LST.1.06964.uvA', + "zen.all.xx.LST.1.06964.uvA", ] + def axes_contains(ax, obj_list): """ Check that a matplotlib.Axes instance contains certain elements. @@ -37,15 +38,16 @@ def axes_contains(ax, obj_list): objtype, num_expected = obj num = 0 for elem in elems: - if isinstance(elem, objtype): num += 1 + if isinstance(elem, objtype): + num += 1 if num != num_expected: return False # Return True if no problems found return True -class Test_Plot(unittest.TestCase): +class Test_Plot(unittest.TestCase): def setUp(self): """ Load data and calculate power spectra. @@ -59,9 +61,9 @@ def setUp(self): self.uvd = uvd # Load beam file - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") self.bm = pspecbeam.PSpecBeamUV(beamfile) - self.bm.filename = 'HERA_NF_dipole_power.beamfits' + self.bm.filename = "HERA_NF_dipole_power.beamfits" # We only actually have 1 data file here, so slide the time axis by one # integration to avoid noise bias @@ -69,20 +71,33 @@ def setUp(self): uvd2 = uvd.select(times=np.unique(uvd.time_array)[1::2], inplace=False) # Create a new PSpecData object - self.ds = pspecdata.PSpecData(dsets=[uvd1, uvd2], wgts=[None, None], - beam=self.bm) - self.ds.rephase_to_dset(0) # Phase to the zeroth dataset + self.ds = pspecdata.PSpecData( + dsets=[uvd1, uvd2], wgts=[None, None], beam=self.bm + ) + self.ds.rephase_to_dset(0) # Phase to the zeroth dataset # Construct list of baseline pairs to calculate power spectra for - bls = [(24,25), (37,38), (38,39),] + bls = [ + (24, 25), + (37, 38), + (38, 39), + ] self.bls1, self.bls2, blp = utils.construct_blpairs( - bls, exclude_permutations=False, exclude_auto_bls=True) + bls, exclude_permutations=False, exclude_auto_bls=True + ) # Calculate the power spectrum - self.uvp = self.ds.pspec(self.bls1, self.bls2, (0, 1), - ('xx','xx'), spw_ranges=[(300, 400), (600,721)], - input_data_weight='identity', norm='I', - taper='blackman-harris', verbose=False) + self.uvp = self.ds.pspec( + self.bls1, + self.bls2, + (0, 1), + ("xx", "xx"), + spw_ranges=[(300, 400), (600, 721)], + input_data_weight="identity", + norm="I", + taper="blackman-harris", + verbose=False, + ) def tearDown(self): pass @@ -99,75 +114,152 @@ def test_plot_average(self): blps = [blp for blp in blpairs] # Plot the spectra averaged over baseline-pairs and times - f1 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), - average_blpairs=True, average_times=True) - elements = [(matplotlib.lines.Line2D, 1),] + f1 = plot.delay_spectrum( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=True, + average_times=True, + ) + elements = [ + (matplotlib.lines.Line2D, 1), + ] assert axes_contains(f1.axes[0], elements) plt.close(f1) # Average over baseline-pairs but keep the time bins intact - f2 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), - average_blpairs=True, average_times=False) - elements = [(matplotlib.lines.Line2D, self.uvp.Ntimes),] + f2 = plot.delay_spectrum( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=True, + average_times=False, + ) + elements = [ + (matplotlib.lines.Line2D, self.uvp.Ntimes), + ] assert axes_contains(f2.axes[0], elements) plt.close(f2) # Average over times, but keep the baseline-pairs separate - f3 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), - average_blpairs=False, average_times=True) - elements = [(matplotlib.lines.Line2D, self.uvp.Nblpairs),] + f3 = plot.delay_spectrum( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=False, + average_times=True, + ) + elements = [ + (matplotlib.lines.Line2D, self.uvp.Nblpairs), + ] assert axes_contains(f3.axes[0], elements) plt.close(f3) # Plot the spectra averaged over baseline-pairs and times, but also # fold the delay axis - f4 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), - average_blpairs=True, average_times=True, - fold=True) - elements = [(matplotlib.lines.Line2D, 1),] + f4 = plot.delay_spectrum( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=True, + average_times=True, + fold=True, + ) + elements = [ + (matplotlib.lines.Line2D, 1), + ] assert axes_contains(f4.axes[0], elements) plt.close(f4) # Plot imaginary part - f4 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), - average_blpairs=False, average_times=True, - component='imag') - elements = [(matplotlib.lines.Line2D, self.uvp.Nblpairs),] + f4 = plot.delay_spectrum( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=False, + average_times=True, + component="imag", + ) + elements = [ + (matplotlib.lines.Line2D, self.uvp.Nblpairs), + ] assert axes_contains(f4.axes[0], elements) plt.close(f4) # Plot abs - f5 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), - average_blpairs=False, average_times=True, - component='abs') - elements = [(matplotlib.lines.Line2D, self.uvp.Nblpairs),] + f5 = plot.delay_spectrum( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=False, + average_times=True, + component="abs", + ) + elements = [ + (matplotlib.lines.Line2D, self.uvp.Nblpairs), + ] assert axes_contains(f4.axes[0], elements) plt.close(f5) # test errorbar plotting w/ markers # bootstrap resample - (uvp_avg, _, - _) = grouping.bootstrap_resampled_error(self.uvp, time_avg=True, - Nsamples=100, normal_std=True, - robust_std=False, verbose=False) - - f6 = plot.delay_spectrum(uvp_avg, uvp_avg.get_blpairs(), spw=0, - pol=('xx','xx'), average_blpairs=False, - average_times=False, - component='real', error='bs_std', lines=False, - markers=True) + (uvp_avg, _, _) = grouping.bootstrap_resampled_error( + self.uvp, + time_avg=True, + Nsamples=100, + normal_std=True, + robust_std=False, + verbose=False, + ) + + f6 = plot.delay_spectrum( + uvp_avg, + uvp_avg.get_blpairs(), + spw=0, + pol=("xx", "xx"), + average_blpairs=False, + average_times=False, + component="real", + error="bs_std", + lines=False, + markers=True, + ) plt.close(f6) # plot errorbar instead of pspec - f7 = plot.delay_spectrum(uvp_avg, uvp_avg.get_blpairs(), spw=0, - pol=('xx','xx'), average_blpairs=False, - average_times=False, - component='real', lines=False, - markers=True, plot_stats='bs_std') + f7 = plot.delay_spectrum( + uvp_avg, + uvp_avg.get_blpairs(), + spw=0, + pol=("xx", "xx"), + average_blpairs=False, + average_times=False, + component="real", + lines=False, + markers=True, + plot_stats="bs_std", + ) plt.close(f7) - def test_plot_cosmo(self): """ Test that cosmological units can be used on plots. @@ -178,24 +270,41 @@ def test_plot_cosmo(self): # Set cosmology and plot in non-delay (i.e. cosmological) units self.uvp.set_cosmology(conversions.Cosmo_Conversions()) - f1 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), - average_blpairs=True, average_times=True, - delay=False) + f1 = plot.delay_spectrum( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=True, + average_times=True, + delay=False, + ) elements = [(matplotlib.lines.Line2D, 1), (matplotlib.legend.Legend, 0)] - self.assertTrue( axes_contains(f1.axes[0], elements) ) + self.assertTrue(axes_contains(f1.axes[0], elements)) plt.close(f1) # Plot in Delta^2 units - f2 = plot.delay_spectrum(self.uvp, [blps,], spw=0, pol=('xx','xx'), - average_blpairs=True, average_times=True, - delay=False, deltasq=True, legend=True, - label_type='blpair') + f2 = plot.delay_spectrum( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=True, + average_times=True, + delay=False, + deltasq=True, + legend=True, + label_type="blpair", + ) # Should contain 1 line and 1 legend elements = [(matplotlib.lines.Line2D, 1), (matplotlib.legend.Legend, 1)] - self.assertTrue( axes_contains(f2.axes[0], elements) ) + self.assertTrue(axes_contains(f2.axes[0], elements)) plt.close(f2) - def test_delay_spectrum_misc(self): # various other tests for plot.delay_spectrum @@ -204,10 +313,18 @@ def test_delay_spectrum_misc(self): blps = [blp for blp in blpairs] # times selection, label_type - f1 = plot.delay_spectrum(self.uvp, blpairs[:1], spw=0, pol=('xx','xx'), - times=self.uvp.time_avg_array[:1], lines=False, - markers=True, logscale=False, label_type='key', - force_plot=False) + f1 = plot.delay_spectrum( + self.uvp, + blpairs[:1], + spw=0, + pol=("xx", "xx"), + times=self.uvp.time_avg_array[:1], + lines=False, + markers=True, + logscale=False, + label_type="key", + force_plot=False, + ) plt.close(f1) # test force plot exception @@ -215,20 +332,33 @@ def test_delay_spectrum_misc(self): for i in range(3): # build-up a large uvpspec object _uvp = copy.deepcopy(uvp) - _uvp.time_avg_array += (i+1)**2 + _uvp.time_avg_array += (i + 1) ** 2 uvp = uvp + _uvp - pytest.raises(ValueError, plot.delay_spectrum, uvp, uvp.get_blpairs(), 0, 'xx') - - f2 = plot.delay_spectrum(uvp, uvp.get_blpairs(), 0, ('xx','xx'), - force_plot=True, label_type='blpairt', - logscale=False, lines=True, markers=True) + pytest.raises(ValueError, plot.delay_spectrum, uvp, uvp.get_blpairs(), 0, "xx") + + f2 = plot.delay_spectrum( + uvp, + uvp.get_blpairs(), + 0, + ("xx", "xx"), + force_plot=True, + label_type="blpairt", + logscale=False, + lines=True, + markers=True, + ) plt.close(f2) # exceptions - pytest.raises(ValueError, plot.delay_spectrum, uvp, - uvp.get_blpairs()[:3], 0, ('xx','xx'), - label_type='foo') - + pytest.raises( + ValueError, + plot.delay_spectrum, + uvp, + uvp.get_blpairs()[:3], + 0, + ("xx", "xx"), + label_type="foo", + ) def test_plot_waterfall(self): """ @@ -240,36 +370,85 @@ def test_plot_waterfall(self): # Set cosmology and plot in non-delay (i.e. cosmological) units self.uvp.set_cosmology(conversions.Cosmo_Conversions(), overwrite=True) - f1 = plot.delay_waterfall(self.uvp, [blps,], spw=0, pol=('xx','xx'), - average_blpairs=True, delay=False) + f1 = plot.delay_waterfall( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=True, + delay=False, + ) plt.close(f1) # Plot in Delta^2 units - f2 = plot.delay_waterfall(self.uvp, [blps,], spw=0, pol=('xx','xx'), - average_blpairs=True, delay=False, - deltasq=True) + f2 = plot.delay_waterfall( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=True, + delay=False, + deltasq=True, + ) plt.close(f2) # Try some other arguments - f3 = plot.delay_waterfall(self.uvp, [blpairs,], spw=0, pol=('xx','xx'), - average_blpairs=False, delay=True, - log=False, vmin=-1., vmax=3., - cmap='RdBu', fold=True, component='abs') + f3 = plot.delay_waterfall( + self.uvp, + [ + blpairs, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=False, + delay=True, + log=False, + vmin=-1.0, + vmax=3.0, + cmap="RdBu", + fold=True, + component="abs", + ) plt.close(f3) # Try with imaginary component - f4 = plot.delay_waterfall(self.uvp, [blpairs,], spw=0, pol=('xx','xx'), - average_blpairs=False, delay=True, - log=False, vmin=-1., vmax=3., - cmap='RdBu', fold=True, component='imag') + f4 = plot.delay_waterfall( + self.uvp, + [ + blpairs, + ], + spw=0, + pol=("xx", "xx"), + average_blpairs=False, + delay=True, + log=False, + vmin=-1.0, + vmax=3.0, + cmap="RdBu", + fold=True, + component="imag", + ) plt.close(f4) # Try some more arguments fig, axes = plt.subplots(1, len(blps)) - plot.delay_waterfall(self.uvp, [blps,], spw=0, pol=('xx','xx'), - lst_in_hrs=False, - times=np.unique(self.uvp.time_avg_array)[:10], - axes=axes, component='abs', title_type='blvec') + plot.delay_waterfall( + self.uvp, + [ + blps, + ], + spw=0, + pol=("xx", "xx"), + lst_in_hrs=False, + times=np.unique(self.uvp.time_avg_array)[:10], + axes=axes, + component="abs", + title_type="blvec", + ) plt.close() # exceptions @@ -278,10 +457,12 @@ def test_plot_waterfall(self): _uvp = copy.deepcopy(self.uvp) _uvp.blpair_array += i * 20 uvp += _uvp - pytest.raises(ValueError, plot.delay_waterfall, uvp, - uvp.get_blpairs(), 0, ('xx','xx')) - fig = plot.delay_waterfall(uvp, uvp.get_blpairs(), 0, ('xx','xx'), - force_plot=True) + pytest.raises( + ValueError, plot.delay_waterfall, uvp, uvp.get_blpairs(), 0, ("xx", "xx") + ) + fig = plot.delay_waterfall( + uvp, uvp.get_blpairs(), 0, ("xx", "xx"), force_plot=True + ) plt.close() def test_uvdata_waterfalls(self): @@ -292,10 +473,11 @@ def test_uvdata_waterfalls(self): basename = "test_waterfall_plots_3423523923_{bl}_{pol}" - for d in ['data', 'flags', 'nsamples']: + for d in ["data", "flags", "nsamples"]: print("running on {}".format(d)) - plot.plot_uvdata_waterfalls(uvd, basename, vmin=0, vmax=100, - data=d, plot_mode='real') + plot.plot_uvdata_waterfalls( + uvd, basename, vmin=0, vmax=100, data=d, plot_mode="real" + ) figfiles = glob.glob("test_waterfall_plots_3423523923_*_*.png") assert len(figfiles) == 15 @@ -303,85 +485,197 @@ def test_uvdata_waterfalls(self): os.remove(f) def test_delay_wedge(self): - """ Tests for plot.delay_wedge """ + """Tests for plot.delay_wedge""" # construct new uvp reds, lens, angs = utils.get_reds(self.ds.dsets[0], pick_data_ants=True) - bls1, bls2, blps, _, _ = utils.calc_blpair_reds(self.ds.dsets[0], - self.ds.dsets[1], - exclude_auto_bls=False, - exclude_permutations=True) - uvp = self.ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), - spw_ranges=[(300, 350)], - input_data_weight='identity', norm='I', - taper='blackman-harris', verbose=False) + bls1, bls2, blps, _, _ = utils.calc_blpair_reds( + self.ds.dsets[0], + self.ds.dsets[1], + exclude_auto_bls=False, + exclude_permutations=True, + ) + uvp = self.ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + spw_ranges=[(300, 350)], + input_data_weight="identity", + norm="I", + taper="blackman-harris", + verbose=False, + ) # test basic delay_wedge call - f1 = plot.delay_wedge(uvp, 0, ('xx','xx'), blpairs=None, times=None, - fold=False, delay=True, rotate=False, - component='real', log10=False, loglog=False, - red_tol=1.0, center_line=False, - horizon_lines=False, title=None, ax=None, - cmap='viridis', figsize=(8, 6), deltasq=False, - colorbar=False, cbax=None, vmin=None, vmax=None, - edgecolor='none', flip_xax=False, flip_yax=False, - lw=2) + f1 = plot.delay_wedge( + uvp, + 0, + ("xx", "xx"), + blpairs=None, + times=None, + fold=False, + delay=True, + rotate=False, + component="real", + log10=False, + loglog=False, + red_tol=1.0, + center_line=False, + horizon_lines=False, + title=None, + ax=None, + cmap="viridis", + figsize=(8, 6), + deltasq=False, + colorbar=False, + cbax=None, + vmin=None, + vmax=None, + edgecolor="none", + flip_xax=False, + flip_yax=False, + lw=2, + ) plt.close() # specify blpairs and times - f2 = plot.delay_wedge(uvp, 0, ('xx','xx'), - blpairs=uvp.get_blpairs()[-5:], - times=uvp.time_avg_array[:1], - fold=False, delay=True, component='imag', - rotate=False, log10=False, loglog=False, red_tol=1.0, - center_line=False, horizon_lines=False, title=None, - ax=None, cmap='viridis', - figsize=(8, 6), deltasq=False, colorbar=False, - cbax=None, vmin=None, vmax=None, - edgecolor='none', flip_xax=False, flip_yax=False, - lw=2) + f2 = plot.delay_wedge( + uvp, + 0, + ("xx", "xx"), + blpairs=uvp.get_blpairs()[-5:], + times=uvp.time_avg_array[:1], + fold=False, + delay=True, + component="imag", + rotate=False, + log10=False, + loglog=False, + red_tol=1.0, + center_line=False, + horizon_lines=False, + title=None, + ax=None, + cmap="viridis", + figsize=(8, 6), + deltasq=False, + colorbar=False, + cbax=None, + vmin=None, + vmax=None, + edgecolor="none", + flip_xax=False, + flip_yax=False, + lw=2, + ) plt.close() # fold, deltasq, cosmo and log10, loglog - f3 = plot.delay_wedge(uvp, 0, ('xx','xx'), blpairs=None, times=None, - fold=True, delay=False, component='abs', - rotate=False, log10=True, loglog=True, red_tol=1.0, - center_line=False, horizon_lines=False, - title='hello', ax=None, cmap='viridis', - figsize=(8, 6), deltasq=True, colorbar=False, - cbax=None, vmin=None, vmax=None, - edgecolor='none', flip_xax=False, flip_yax=False, - lw=2) + f3 = plot.delay_wedge( + uvp, + 0, + ("xx", "xx"), + blpairs=None, + times=None, + fold=True, + delay=False, + component="abs", + rotate=False, + log10=True, + loglog=True, + red_tol=1.0, + center_line=False, + horizon_lines=False, + title="hello", + ax=None, + cmap="viridis", + figsize=(8, 6), + deltasq=True, + colorbar=False, + cbax=None, + vmin=None, + vmax=None, + edgecolor="none", + flip_xax=False, + flip_yax=False, + lw=2, + ) plt.close() # colorbar, vranges, flip_axes, edgecolors, lines - f4 = plot.delay_wedge(uvp, 0, ('xx','xx'), blpairs=None, times=None, - fold=False, delay=False, component='abs', - rotate=False, log10=True, loglog=False, red_tol=1.0, - center_line=True, horizon_lines=True, title='hello', - ax=None, cmap='viridis', figsize=(8, 6), - deltasq=True, colorbar=True, cbax=None, vmin=6, - vmax=15, edgecolor='grey', flip_xax=True, - flip_yax=True, lw=2, set_bl_tick_minor=True) + f4 = plot.delay_wedge( + uvp, + 0, + ("xx", "xx"), + blpairs=None, + times=None, + fold=False, + delay=False, + component="abs", + rotate=False, + log10=True, + loglog=False, + red_tol=1.0, + center_line=True, + horizon_lines=True, + title="hello", + ax=None, + cmap="viridis", + figsize=(8, 6), + deltasq=True, + colorbar=True, + cbax=None, + vmin=6, + vmax=15, + edgecolor="grey", + flip_xax=True, + flip_yax=True, + lw=2, + set_bl_tick_minor=True, + ) plt.close() # feed axes, red_tol fig, ax = plt.subplots() cbax = fig.add_axes([0.85, 0.1, 0.05, 0.9]) - cbax.axis('off') - plot.delay_wedge(uvp, 0, ('xx','xx'), blpairs=None, times=None, - fold=False, delay=True, component='abs', - rotate=True, log10=True, loglog=False, red_tol=10.0, - center_line=False, horizon_lines=False, ax=ax, - cmap='viridis', figsize=(8, 6), deltasq=False, - colorbar=True, cbax=cbax, vmin=None, vmax=None, - edgecolor='none', flip_xax=False, flip_yax=False, - lw=2, set_bl_tick_major=True) + cbax.axis("off") + plot.delay_wedge( + uvp, + 0, + ("xx", "xx"), + blpairs=None, + times=None, + fold=False, + delay=True, + component="abs", + rotate=True, + log10=True, + loglog=False, + red_tol=10.0, + center_line=False, + horizon_lines=False, + ax=ax, + cmap="viridis", + figsize=(8, 6), + deltasq=False, + colorbar=True, + cbax=cbax, + vmin=None, + vmax=None, + edgecolor="none", + flip_xax=False, + flip_yax=False, + lw=2, + set_bl_tick_major=True, + ) plt.close() # test exceptions - pytest.raises(ValueError, plot.delay_wedge, uvp, 0, ('xx','xx'), - component='foo') + pytest.raises( + ValueError, plot.delay_wedge, uvp, 0, ("xx", "xx"), component="foo" + ) plt.close() + if __name__ == "__main__": unittest.main() diff --git a/hera_pspec/tests/test_pspecbeam.py b/hera_pspec/tests/test_pspecbeam.py index c9234d54..ab23fabf 100644 --- a/hera_pspec/tests/test_pspecbeam.py +++ b/hera_pspec/tests/test_pspecbeam.py @@ -7,7 +7,6 @@ class Test_DataSet(unittest.TestCase): - def setUp(self): pass @@ -18,8 +17,8 @@ def runTest(self): pass def test_init(self): - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') - bm = pspecbeam.PSpecBeamUV(beamfile) + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + pspecbeam.PSpecBeamUV(beamfile) def test_UVbeam(self): # Precomputed results in the following tests were done "by hand" using @@ -28,12 +27,14 @@ def test_UVbeam(self): beam = pspecbeam.PSpecBeamUV(pstokes_beamfile) Om_p = beam.power_beam_int() Om_pp = beam.power_beam_sq_int() - lower_freq = 120.*10**6 - upper_freq = 128.*10**6 + lower_freq = 120.0 * 10**6 + upper_freq = 128.0 * 10**6 num_freqs = 20 # check pI polarization - scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=2000) + scalar = beam.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, pol="pI", num_steps=2000 + ) np.testing.assert_almost_equal(Om_p[0], 0.080082680885906782) np.testing.assert_almost_equal(Om_p[18], 0.031990943334017245) @@ -43,26 +44,37 @@ def test_UVbeam(self): np.testing.assert_almost_equal(Om_pp[15], 0.018159280192894631) np.testing.assert_almost_equal(Om_pp[-1], 0.014528100116719534) - assert abs(scalar/568847837.72586381 - 1.0) <= 1e-4 + assert abs(scalar / 568847837.72586381 - 1.0) <= 1e-4 # Check array dimensionality assert Om_p.ndim == 1 assert Om_pp.ndim == 1 # convergence of integral - scalar_large_Nsteps = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=10000) + scalar_large_Nsteps = beam.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, pol="pI", num_steps=10000 + ) assert abs(scalar / scalar_large_Nsteps - 1.0) <= 1e-5 # Check that user-defined cosmology can be specified - beam = pspecbeam.PSpecBeamUV(pstokes_beamfile, - cosmo=conversions.Cosmo_Conversions()) + beam = pspecbeam.PSpecBeamUV( + pstokes_beamfile, cosmo=conversions.Cosmo_Conversions() + ) # Check that errors are not raised for other Stokes parameters - for pol in ['pQ', 'pU', 'pV',]: - scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol=pol, num_steps=2000) + for pol in [ + "pQ", + "pU", + "pV", + ]: + scalar = beam.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, pol=pol, num_steps=2000 + ) # test taper execution - scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, num_steps=5000, taper='blackman') + scalar = beam.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, num_steps=5000, taper="blackman" + ) assert abs(scalar / 1989353792.1765163 - 1.0) <= 1e-8 # test Jy_to_mK @@ -78,45 +90,82 @@ def test_UVbeam(self): pytest.raises(TypeError, beam.Jy_to_mK, np.array([1])) # test noise scalar - sclr = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=2000, noise_scalar=True) + sclr = beam.compute_pspec_scalar( + lower_freq, + upper_freq, + num_freqs, + pol="pI", + num_steps=2000, + noise_scalar=True, + ) np.testing.assert_almost_equal(sclr, 71.105979715733) # Check that invalid polarizations raise an error - pol = 'pZ' - pytest.raises(KeyError, beam.compute_pspec_scalar, - lower_freq, upper_freq, num_freqs, pol=pol) - pol = 'XX' - pytest.raises(ValueError, beam.compute_pspec_scalar, - lower_freq, upper_freq, num_freqs, pol=pol) + pol = "pZ" + pytest.raises( + KeyError, + beam.compute_pspec_scalar, + lower_freq, + upper_freq, + num_freqs, + pol=pol, + ) + pol = "XX" + pytest.raises( + ValueError, + beam.compute_pspec_scalar, + lower_freq, + upper_freq, + num_freqs, + pol=pol, + ) # check dipole beams work dipole_beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") beam = pspecbeam.PSpecBeamUV(dipole_beamfile) - scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='XX') - pytest.raises(ValueError, beam.compute_pspec_scalar, - lower_freq, upper_freq, num_freqs, pol='pI') + scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol="XX") + pytest.raises( + ValueError, + beam.compute_pspec_scalar, + lower_freq, + upper_freq, + num_freqs, + pol="pI", + ) # check efield beams work efield_beamfile = os.path.join(DATA_PATH, "HERA_NF_efield.beamfits") beam = pspecbeam.PSpecBeamUV(efield_beamfile) - scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='XX') - pytest.raises(ValueError, beam.compute_pspec_scalar, - lower_freq, upper_freq, num_freqs, pol='pI') + scalar = beam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol="XX") + pytest.raises( + ValueError, + beam.compute_pspec_scalar, + lower_freq, + upper_freq, + num_freqs, + pol="pI", + ) def test_Gaussbeam(self): - gauss = pspecbeam.PSpecBeamGauss(0.8, np.linspace(115e6, 130e6, 50, endpoint=False)) + gauss = pspecbeam.PSpecBeamGauss( + 0.8, np.linspace(115e6, 130e6, 50, endpoint=False) + ) Om_p = gauss.power_beam_int() Om_pp = gauss.power_beam_sq_int() - lower_freq = 120.*10**6 - upper_freq = 128.*10**6 + lower_freq = 120.0 * 10**6 + upper_freq = 128.0 * 10**6 num_freqs = 20 - scalar = gauss.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, pol='pI', num_steps=2000) + scalar = gauss.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, pol="pI", num_steps=2000 + ) # Check that user-defined cosmology can be specified - bgauss = pspecbeam.PSpecBeamGauss(0.8, - np.linspace(115e6, 130e6, 50, endpoint=False), - cosmo=conversions.Cosmo_Conversions()) + pspecbeam.PSpecBeamGauss( + 0.8, + np.linspace(115e6, 130e6, 50, endpoint=False), + cosmo=conversions.Cosmo_Conversions(), + ) # Check array dimensionality assert Om_p.ndim == 1 @@ -128,14 +177,18 @@ def test_Gaussbeam(self): assert Om_pp[4] == 0.36258881134617554 assert Om_pp[4] == Om_pp[0] - assert abs(scalar/6392750120.8657961 - 1.0) <= 1e-4 + assert abs(scalar / 6392750120.8657961 - 1.0) <= 1e-4 # convergence of integral - scalar_large_Nsteps = gauss.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, num_steps=5000) + scalar_large_Nsteps = gauss.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, num_steps=5000 + ) assert abs(scalar / scalar_large_Nsteps - 1.0) <= 1e-5 # test taper execution - scalar = gauss.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, num_steps=5000, taper='blackman') + scalar = gauss.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, num_steps=5000, taper="blackman" + ) assert abs(scalar / 22123832163.072491 - 1.0) <= 1e-8 def test_BeamFromArray(self): @@ -143,95 +196,145 @@ def test_BeamFromArray(self): Test PSpecBeamFromArray """ # Get Gaussian beam to use as a reference - gauss = pspecbeam.PSpecBeamGauss(0.8, np.linspace(115e6, 130e6, 50, endpoint=False)) + gauss = pspecbeam.PSpecBeamGauss( + 0.8, np.linspace(115e6, 130e6, 50, endpoint=False) + ) Om_P = gauss.power_beam_int() Om_PP = gauss.power_beam_sq_int() beam_freqs = gauss.beam_freqs # Array specs for tests - lower_freq = 120.*10**6 - upper_freq = 128.*10**6 + lower_freq = 120.0 * 10**6 + upper_freq = 128.0 * 10**6 num_freqs = 20 # Check that PSpecBeamFromArray can be instantiated - psbeam = pspecbeam.PSpecBeamFromArray(OmegaP=Om_P, OmegaPP=Om_PP, - beam_freqs=beam_freqs) + psbeam = pspecbeam.PSpecBeamFromArray( + OmegaP=Om_P, OmegaPP=Om_PP, beam_freqs=beam_freqs + ) psbeampol = pspecbeam.PSpecBeamFromArray( - OmegaP={'pI': Om_P, 'pQ': Om_P}, - OmegaPP={'pI': Om_PP, 'pQ': Om_PP}, - beam_freqs=beam_freqs) + OmegaP={"pI": Om_P, "pQ": Om_P}, + OmegaPP={"pI": Om_PP, "pQ": Om_PP}, + beam_freqs=beam_freqs, + ) # Check that user-defined cosmology can be specified - bm2 = pspecbeam.PSpecBeamFromArray(OmegaP=Om_P, OmegaPP=Om_PP, - beam_freqs=beam_freqs, - cosmo=conversions.Cosmo_Conversions()) + pspecbeam.PSpecBeamFromArray( + OmegaP=Om_P, + OmegaPP=Om_PP, + beam_freqs=beam_freqs, + cosmo=conversions.Cosmo_Conversions(), + ) # Compare scalar calculation with Gaussian case - scalar = psbeam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, - pol='pI', num_steps=2000) - g_scalar = gauss.compute_pspec_scalar(lower_freq, upper_freq, - num_freqs, pol='pI', - num_steps=2000) + scalar = psbeam.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, pol="pI", num_steps=2000 + ) + g_scalar = gauss.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, pol="pI", num_steps=2000 + ) np.testing.assert_array_almost_equal(scalar, g_scalar) # Check that polarizations are recognized and invalid ones rejected - scalarp = psbeampol.compute_pspec_scalar(lower_freq, upper_freq, - num_freqs, pol='pQ', - num_steps=2000) + psbeampol.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, pol="pQ", num_steps=2000 + ) # Test taper execution (same as Gaussian case) - scalar = psbeam.compute_pspec_scalar(lower_freq, upper_freq, num_freqs, - num_steps=5000, taper='blackman') + scalar = psbeam.compute_pspec_scalar( + lower_freq, upper_freq, num_freqs, num_steps=5000, taper="blackman" + ) assert abs(scalar / 22123832163.072491 - 1.0) <= 1e-8 # Check that invalid init args raise errors - pytest.raises(TypeError, pspecbeam.PSpecBeamFromArray, OmegaP=Om_P, - OmegaPP={'pI': Om_PP}, beam_freqs=beam_freqs) - pytest.raises(KeyError, pspecbeam.PSpecBeamFromArray, - OmegaP={'pI': Om_P, 'pQ': Om_P}, - OmegaPP={'pI': Om_PP,}, - beam_freqs=beam_freqs) - - pytest.raises(KeyError, pspecbeam.PSpecBeamFromArray, - OmegaP={'A': Om_P}, - OmegaPP={'A': Om_PP,}, - beam_freqs=beam_freqs) - - pytest.raises(TypeError, pspecbeam.PSpecBeamFromArray, - OmegaP={'pI': Om_P,}, - OmegaPP={'pI': 'string',}, - beam_freqs=beam_freqs) - - pytest.raises(ValueError, pspecbeam.PSpecBeamFromArray, - OmegaP={'pI': Om_P}, - OmegaPP={'pI': Om_PP[:-2],}, - beam_freqs=beam_freqs) - - pytest.raises(TypeError, pspecbeam.PSpecBeamFromArray, - OmegaP=Om_P, - OmegaPP={'pI': Om_PP[:-2],}, - beam_freqs=beam_freqs) - - pytest.raises(KeyError, pspecbeam.PSpecBeamFromArray, - OmegaP={'foo': Om_P}, - OmegaPP={'pI': Om_PP,}, - beam_freqs=beam_freqs) - - pytest.raises(KeyError, pspecbeam.PSpecBeamFromArray, - OmegaP={'pI': Om_P}, - OmegaPP={'foo': Om_PP,}, - beam_freqs=beam_freqs) - - pytest.raises(KeyError, psbeam.add_pol, 'foo', Om_P, Om_PP) - pytest.raises(KeyError, psbeam.power_beam_int, 'foo') - pytest.raises(KeyError, psbeam.power_beam_sq_int, 'foo') + pytest.raises( + TypeError, + pspecbeam.PSpecBeamFromArray, + OmegaP=Om_P, + OmegaPP={"pI": Om_PP}, + beam_freqs=beam_freqs, + ) + pytest.raises( + KeyError, + pspecbeam.PSpecBeamFromArray, + OmegaP={"pI": Om_P, "pQ": Om_P}, + OmegaPP={ + "pI": Om_PP, + }, + beam_freqs=beam_freqs, + ) + + pytest.raises( + KeyError, + pspecbeam.PSpecBeamFromArray, + OmegaP={"A": Om_P}, + OmegaPP={ + "A": Om_PP, + }, + beam_freqs=beam_freqs, + ) + + pytest.raises( + TypeError, + pspecbeam.PSpecBeamFromArray, + OmegaP={ + "pI": Om_P, + }, + OmegaPP={ + "pI": "string", + }, + beam_freqs=beam_freqs, + ) + + pytest.raises( + ValueError, + pspecbeam.PSpecBeamFromArray, + OmegaP={"pI": Om_P}, + OmegaPP={ + "pI": Om_PP[:-2], + }, + beam_freqs=beam_freqs, + ) + + pytest.raises( + TypeError, + pspecbeam.PSpecBeamFromArray, + OmegaP=Om_P, + OmegaPP={ + "pI": Om_PP[:-2], + }, + beam_freqs=beam_freqs, + ) + + pytest.raises( + KeyError, + pspecbeam.PSpecBeamFromArray, + OmegaP={"foo": Om_P}, + OmegaPP={ + "pI": Om_PP, + }, + beam_freqs=beam_freqs, + ) + + pytest.raises( + KeyError, + pspecbeam.PSpecBeamFromArray, + OmegaP={"pI": Om_P}, + OmegaPP={ + "foo": Om_PP, + }, + beam_freqs=beam_freqs, + ) + + pytest.raises(KeyError, psbeam.add_pol, "foo", Om_P, Om_PP) + pytest.raises(KeyError, psbeam.power_beam_int, "foo") + pytest.raises(KeyError, psbeam.power_beam_sq_int, "foo") # Check that invalid method args raise errors - pytest.raises(KeyError, psbeam.power_beam_int, pol='blah') - pytest.raises(KeyError, psbeam.power_beam_sq_int, pol='blah') - pytest.raises(KeyError, psbeam.add_pol, pol='A', - OmegaP=Om_P, OmegaPP=Om_PP) + pytest.raises(KeyError, psbeam.power_beam_int, pol="blah") + pytest.raises(KeyError, psbeam.power_beam_sq_int, pol="blah") + pytest.raises(KeyError, psbeam.add_pol, pol="A", OmegaP=Om_P, OmegaPP=Om_PP) # Check that string works assert len(str(psbeam)) > 0 @@ -240,57 +343,94 @@ def test_PSpecBeamBase(self): """ Test that base class can be instantiated. """ - bm1 = pspecbeam.PSpecBeamBase() + pspecbeam.PSpecBeamBase() # Check that user-defined cosmology can be specified - bm2 = pspecbeam.PSpecBeamBase(cosmo=conversions.Cosmo_Conversions()) + pspecbeam.PSpecBeamBase(cosmo=conversions.Cosmo_Conversions()) def test_get_Omegas(self): - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") beam = pspecbeam.PSpecBeamUV(beamfile) - OP, OPP = beam.get_Omegas(('xx','xx')) + OP, OPP = beam.get_Omegas(("xx", "xx")) assert OP.shape == (26, 1) assert OPP.shape == (26, 1) - OP, OPP = beam.get_Omegas([(-5,-5), (-6,-6)]) + OP, OPP = beam.get_Omegas([(-5, -5), (-6, -6)]) assert OP.shape == (26, 2) assert OPP.shape == (26, 2) - pytest.raises(TypeError, beam.get_Omegas, 'xx') - pytest.raises(NotImplementedError, beam.get_Omegas, [('pI','pQ'),]) - + pytest.raises(TypeError, beam.get_Omegas, "xx") + pytest.raises( + NotImplementedError, + beam.get_Omegas, + [ + ("pI", "pQ"), + ], + ) def test_beam_normalized_response(self): - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') - beam = pspecbeam.PSpecBeamUV(beamfile) - freq = np.linspace(130.0*1e6, 140.0*1e6, 10) - nside = beam.primary_beam.nside #uvbeam object - beam_res = pspecbeam.PSpecBeamUV.beam_normalized_response(beam, pol='xx', freq=freq) + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + beam = pspecbeam.PSpecBeamUV(beamfile) + freq = np.linspace(130.0 * 1e6, 140.0 * 1e6, 10) + nside = beam.primary_beam.nside # uvbeam object + beam_res = pspecbeam.PSpecBeamUV.beam_normalized_response( + beam, pol="xx", freq=freq + ) - #tests for dimensions + # tests for dimensions assert len(beam_res[1]) == len(freq) assert beam_res[0].ndim == 2 - assert np.shape(beam_res[0]) == (len(freq), (12*nside**2)) - - #tests for polarization - pytest.raises(ValueError, pspecbeam.PSpecBeamUV.beam_normalized_response, beam, pol='ll', freq=freq) - - #test if it is a power beam + assert np.shape(beam_res[0]) == (len(freq), (12 * nside**2)) + + # tests for polarization + pytest.raises( + ValueError, + pspecbeam.PSpecBeamUV.beam_normalized_response, + beam, + pol="ll", + freq=freq, + ) + + # test if it is a power beam efield_beamfile = os.path.join(DATA_PATH, "HERA_NF_efield.beamfits") - beam_efield = pspecbeam.PSpecBeamUV(efield_beamfile) - beam_efield.primary_beam.beam_type='voltage' - pytest.raises(ValueError, pspecbeam.PSpecBeamUV.beam_normalized_response, beam_efield, pol='xx', freq=freq) - - #test for right axes - beam_efield.primary_beam.beam_type='power' - beam_efield.primary_beam.Naxes_vec=2 - pytest.raises(ValueError, pspecbeam.PSpecBeamUV.beam_normalized_response, beam_efield, pol='xx', freq=freq) - - #test for peak normalization - beam_efield.primary_beam.Naxes_vec=1 - beam_efield.primary_beam._data_normalization.value = 'area' - pytest.raises(ValueError, pspecbeam.PSpecBeamUV.beam_normalized_response, beam_efield, pol='xx', freq=freq) - - #test for the coordinate system - beam_efield.primary_beam._data_normalization.value = 'peak' - beam_efield.primary_beam.pixel_coordinate_system = 'cartesian' - pytest.raises(ValueError, pspecbeam.PSpecBeamUV.beam_normalized_response, beam_efield, pol='xx', freq=freq) + beam_efield = pspecbeam.PSpecBeamUV(efield_beamfile) + beam_efield.primary_beam.beam_type = "voltage" + pytest.raises( + ValueError, + pspecbeam.PSpecBeamUV.beam_normalized_response, + beam_efield, + pol="xx", + freq=freq, + ) + + # test for right axes + beam_efield.primary_beam.beam_type = "power" + beam_efield.primary_beam.Naxes_vec = 2 + pytest.raises( + ValueError, + pspecbeam.PSpecBeamUV.beam_normalized_response, + beam_efield, + pol="xx", + freq=freq, + ) + + # test for peak normalization + beam_efield.primary_beam.Naxes_vec = 1 + beam_efield.primary_beam._data_normalization.value = "area" + pytest.raises( + ValueError, + pspecbeam.PSpecBeamUV.beam_normalized_response, + beam_efield, + pol="xx", + freq=freq, + ) + + # test for the coordinate system + beam_efield.primary_beam._data_normalization.value = "peak" + beam_efield.primary_beam.pixel_coordinate_system = "cartesian" + pytest.raises( + ValueError, + pspecbeam.PSpecBeamUV.beam_normalized_response, + beam_efield, + pol="xx", + freq=freq, + ) diff --git a/hera_pspec/tests/test_pspecdata.py b/hera_pspec/tests/test_pspecdata.py index 80c1f4d5..c8de090d 100644 --- a/hera_pspec/tests/test_pspecdata.py +++ b/hera_pspec/tests/test_pspecdata.py @@ -2,8 +2,8 @@ import pytest import numpy as np import pyuvdata as uv -import os, copy, sys -from scipy.integrate import simps, trapz +import os, copy +from scipy.integrate import trapz from .. import pspecdata, pspecbeam, conversions, container, utils, testing from hera_pspec.data import DATA_PATH from pyuvdata import UVData, UVCal, utils as uvutils @@ -14,21 +14,20 @@ import warnings import glob from uvtools import dspec + # Data files to use in tests -dfiles = [ - 'zen.2458042.12552.xx.HH.uvXAA', - 'zen.2458042.12552.xx.HH.uvXAA' -] -dfiles_std = [ - 'zen.2458042.12552.std.xx.HH.uvXAA', - 'zen.2458042.12552.std.xx.HH.uvXAA' -] +dfiles = ["zen.2458042.12552.xx.HH.uvXAA", "zen.2458042.12552.xx.HH.uvXAA"] +dfiles_std = ["zen.2458042.12552.std.xx.HH.uvXAA", "zen.2458042.12552.std.xx.HH.uvXAA"] # List of tapering function to use in tests -taper_selection = ['none', 'bh7',] -#taper_selection = ['blackman', 'blackman-harris', 'gaussian0.4', 'kaiser2', +taper_selection = [ + "none", + "bh7", +] +# taper_selection = ['blackman', 'blackman-harris', 'gaussian0.4', 'kaiser2', # 'kaiser3', 'hamming', 'hanning', 'parzen'] + def generate_pos_def(n): """ Generate a random positive definite Hermitian matrix. @@ -43,12 +42,13 @@ def generate_pos_def(n): A : array_like Positive definite matrix """ - A = np.random.normal(size=(n,n)) + 1j * np.random.normal(size=(n,n)) + A = np.random.normal(size=(n, n)) + 1j * np.random.normal(size=(n, n)) A += np.conjugate(A).T # Add just enough of an identity matrix to make all eigenvalues positive - A += -1.01*np.min(np.linalg.eigvalsh(A))*np.identity(n) + A += -1.01 * np.min(np.linalg.eigvalsh(A)) * np.identity(n) return A + def generate_pos_def_all_pos(n): """ Generate a random positive definite symmetric matrix, with all entries @@ -64,12 +64,13 @@ def generate_pos_def_all_pos(n): A : array_like Positive definite matrix """ - A = np.random.uniform(size=(n,n)) + A = np.random.uniform(size=(n, n)) A += A.T # Add just enough of an identity matrix to make all eigenvalues positive - A += -1.01*np.min(np.linalg.eigvalsh(A))*np.identity(n) + A += -1.01 * np.min(np.linalg.eigvalsh(A)) * np.identity(n) return A + def diagonal_or_not(mat, places=7): """ Tests whether a matrix is diagonal or not. @@ -86,11 +87,11 @@ def diagonal_or_not(mat, places=7): """ mat_norm = np.linalg.norm(mat) diag_mat_norm = np.linalg.norm(np.diag(np.diag(mat))) - diag = (round(mat_norm-diag_mat_norm, places) == 0) + diag = round(mat_norm - diag_mat_norm, places) == 0 return diag -class Test_PSpecData(unittest.TestCase): +class Test_PSpecData(unittest.TestCase): def setUp(self): # Instantiate empty PSpecData @@ -114,23 +115,23 @@ def setUp(self): self.w = [None for _d in dfiles] # Load beam file - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") self.bm = pspecbeam.PSpecBeamUV(beamfile) - self.bm.filename = 'HERA_NF_dipole_power.beamfits' + self.bm.filename = "HERA_NF_dipole_power.beamfits" - #Load isotropic beam file - beamfile_Q = os.path.join(DATA_PATH, 'isotropic_beam.beamfits') - self.bm_Q = pspecbeam.PSpecBeamUV(beamfile_Q) - self.bm_Q.filename = 'isotropic_beam.beamfits' + # Load isotropic beam file + beamfile_Q = os.path.join(DATA_PATH, "isotropic_beam.beamfits") + self.bm_Q = pspecbeam.PSpecBeamUV(beamfile_Q) + self.bm_Q.filename = "isotropic_beam.beamfits" # load another data file self.uvd = uv.UVData() - self.uvd.read_miriad(os.path.join(DATA_PATH, - "zen.2458042.17772.xx.HH.uvXA")) + self.uvd.read_miriad(os.path.join(DATA_PATH, "zen.2458042.17772.xx.HH.uvXA")) self.uvd_std = uv.UVData() - self.uvd_std.read_miriad(os.path.join(DATA_PATH, - "zen.2458042.17772.std.xx.HH.uvXA")) + self.uvd_std.read_miriad( + os.path.join(DATA_PATH, "zen.2458042.17772.std.xx.HH.uvXA") + ) def tearDown(self): pass @@ -143,16 +144,18 @@ def test_init(self): ds = pspecdata.PSpecData() # Test whether unequal no. of weights is picked up - self.assertRaises( AssertionError, - pspecdata.PSpecData, - [uv.UVData(), uv.UVData(), uv.UVData()], - [uv.UVData(), uv.UVData()] ) + self.assertRaises( + AssertionError, + pspecdata.PSpecData, + [uv.UVData(), uv.UVData(), uv.UVData()], + [uv.UVData(), uv.UVData()], + ) # Test passing data and weights of the wrong type d_arr = np.ones((6, 8)) - d_lst = [[0,1,2] for i in range(5)] - d_float = 12. - d_dict = {'(0,1)':np.arange(5), '(0,2)':np.arange(5)} + d_lst = [[0, 1, 2] for i in range(5)] + d_float = 12.0 + d_dict = {"(0,1)": np.arange(5), "(0,2)": np.arange(5)} self.assertRaises(TypeError, pspecdata.PSpecData, d_arr, d_arr) self.assertRaises(TypeError, pspecdata.PSpecData, d_lst, d_lst) @@ -164,7 +167,7 @@ def test_init(self): # Test get weights when fed a UVData for weights ds = pspecdata.PSpecData(dsets=[self.uvd, self.uvd], wgts=[self.uvd, self.uvd]) - key = (0, (24, 25), 'xx') + key = (0, (24, 25), "xx") assert np.all(np.isclose(ds.x(key), ds.w(key))) # Test labels when adding dsets @@ -173,11 +176,11 @@ def test_init(self): assert len(ds.labels) == 0 ds.add([uvd, uvd], [None, None]) assert len(ds.labels) == 2 - ds.add(uvd, None, labels='foo') + ds.add(uvd, None, labels="foo") assert len(ds.dsets) == len(ds.labels) == 3 - assert ds.labels == ['dset0', 'dset1', 'foo'] + assert ds.labels == ["dset0", "dset1", "foo"] ds.add(uvd, None) - assert ds.labels == ['dset0', 'dset1', 'foo', 'dset3'] + assert ds.labels == ["dset0", "dset1", "foo", "dset3"] # Test some exceptions ds = pspecdata.PSpecData() @@ -198,17 +201,16 @@ def test_add_data(self): # test adding non UVCal for cals pytest.raises(TypeError, self.ds.add, [uv], None, cals=[1]) # test TypeError if dsets is dict but other inputs are not - pytest.raises(TypeError, self.ds.add, {'d':uv}, [0]) - pytest.raises(TypeError, self.ds.add, {'d':uv}, {'d':uv}, dsets_std=[0]) - pytest.raises(TypeError, self.ds.add, {'d':uv}, {'d':uv}, cals=[0]) + pytest.raises(TypeError, self.ds.add, {"d": uv}, [0]) + pytest.raises(TypeError, self.ds.add, {"d": uv}, {"d": uv}, dsets_std=[0]) + pytest.raises(TypeError, self.ds.add, {"d": uv}, {"d": uv}, cals=[0]) # specifying labels when dsets is a dict is a ValueError - pytest.raises(ValueError, self.ds.add, {'d':uv}, None, labels=['d']) + pytest.raises(ValueError, self.ds.add, {"d": uv}, None, labels=["d"]) # use lists, but not appropriate lengths pytest.raises(AssertionError, self.ds.add, [uv], [uv, uv]) pytest.raises(AssertionError, self.ds.add, [uv], None, dsets_std=[uv, uv]) pytest.raises(AssertionError, self.ds.add, [uv], None, cals=[None, None]) - pytest.raises(AssertionError, self.ds.add, [uv], None, labels=['foo', 'bar']) - + pytest.raises(AssertionError, self.ds.add, [uv], None, labels=["foo", "bar"]) def test_set_symmetric_taper(self): """ @@ -220,83 +222,110 @@ def test_set_symmetric_taper(self): Ndlys = Nfreq - 3 self.ds.spw_Ndlys = Ndlys - # Set baselines to use for tests key1 = (0, 24, 38) key2 = (1, 25, 38) - key3 = [(0, 24, 38), (0, 24, 38)] - key4 = [(1, 25, 38), (1, 25, 38)] - - rpk1 = {'filter_centers':[0.],'filter_half_widths':[100e-9],'filter_factors':[1e-9]} - rpk2 = {'filter_centers':[0.],'filter_half_widths':[100e-9],'filter_factors':[1e-9]} - self.ds.set_weighting('dayenu') - self.ds.set_r_param(key1,rpk1) - self.ds.set_r_param(key2,rpk2) + + rpk1 = { + "filter_centers": [0.0], + "filter_half_widths": [100e-9], + "filter_factors": [1e-9], + } + rpk2 = { + "filter_centers": [0.0], + "filter_half_widths": [100e-9], + "filter_factors": [1e-9], + } + self.ds.set_weighting("dayenu") + self.ds.set_r_param(key1, rpk1) + self.ds.set_r_param(key2, rpk2) ds1 = copy.deepcopy(self.ds) - ds1.set_spw((10,Nfreq-10)) + ds1.set_spw((10, Nfreq - 10)) ds1.set_symmetric_taper(False) - ds1.set_filter_extension([10,10]) - ds1.set_filter_extension((10,10)) - rm1 = self.ds.R(key1) + ds1.set_filter_extension([10, 10]) + ds1.set_filter_extension((10, 10)) + self.ds.R(key1) self.ds.set_symmetric_taper(True) pytest.raises(ValueError, ds1.set_symmetric_taper, True) - #now make sure warnings are raised when we extend filter with - #symmetric tapering and that symmetric taper is set to false. + # now make sure warnings are raised when we extend filter with + # symmetric tapering and that symmetric taper is set to false. with warnings.catch_warnings(record=True) as w: - self.ds.set_filter_extension((10,10)) + self.ds.set_filter_extension((10, 10)) assert len(w) > 0 - self.assertTrue(not(self.ds.symmetric_taper)) + self.assertTrue(not (self.ds.symmetric_taper)) """ Now directly compare results to expectations. """ self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w) Nfreq = self.ds.spw_Nfreqs - Ntime = self.ds.Ntimes Ndlys = Nfreq - 3 self.ds.spw_Ndlys = Ndlys key1 = (0, 24, 38) - key2 = (1,25, 38) + key2 = (1, 25, 38) - rpk1 = {'filter_centers':[0.],'filter_half_widths':[100e-9],'filter_factors':[1e-9]} - self.ds.set_weighting('dayenu') - self.ds.set_taper('bh7') - self.ds.set_r_param(key1,rpk1) - #get the symmetric tapering + rpk1 = { + "filter_centers": [0.0], + "filter_half_widths": [100e-9], + "filter_factors": [1e-9], + } + self.ds.set_weighting("dayenu") + self.ds.set_taper("bh7") + self.ds.set_r_param(key1, rpk1) + # get the symmetric tapering rmat_symmetric = self.ds.R(key1) - #now set taper to be asymmetric + # now set taper to be asymmetric self.ds.set_symmetric_taper(False) rmat_a = self.ds.R(key1) - #check against independent solution - bh_taper = np.sqrt(dspec.gen_window('bh7', Nfreq).reshape(1,-1)) - rmat = dspec.dayenu_mat_inv(x=self.ds.freqs[self.ds.spw_range[0]:self.ds.spw_range[1]], - filter_centers=[0.], filter_half_widths=[100e-9], filter_factors=[1e-9]) - wmat = np.outer(np.diag(np.sqrt(self.ds.Y(key1))), np.diag(np.sqrt(self.ds.Y(key1)))) + # check against independent solution + bh_taper = np.sqrt(dspec.gen_window("bh7", Nfreq).reshape(1, -1)) + rmat = dspec.dayenu_mat_inv( + x=self.ds.freqs[self.ds.spw_range[0] : self.ds.spw_range[1]], + filter_centers=[0.0], + filter_half_widths=[100e-9], + filter_factors=[1e-9], + ) + wmat = np.outer( + np.diag(np.sqrt(self.ds.Y(key1))), np.diag(np.sqrt(self.ds.Y(key1))) + ) rmat = np.linalg.pinv(wmat * rmat) - self.assertTrue(np.all(np.isclose(rmat_symmetric, bh_taper.T * rmat * bh_taper,atol=1e-6))) - self.assertTrue(np.all(np.isclose(rmat_a, bh_taper.T ** 2. * rmat,atol=1e-6))) - self.assertTrue(not np.all(np.isclose(rmat_symmetric, rmat_a,atol=1e-6))) - + self.assertTrue( + np.all(np.isclose(rmat_symmetric, bh_taper.T * rmat * bh_taper, atol=1e-6)) + ) + self.assertTrue(np.all(np.isclose(rmat_a, bh_taper.T**2.0 * rmat, atol=1e-6))) + self.assertTrue(not np.all(np.isclose(rmat_symmetric, rmat_a, atol=1e-6))) def test_labels(self): """ Test that dataset labels work. """ # Check that specifying labels does work - psd = pspecdata.PSpecData( dsets=[self.d[0], self.d[1],], - wgts=[self.w[0], self.w[1], ], - labels=['red', 'blue']) - np.testing.assert_array_equal( psd.x(('red', 24, 38)), - psd.x((0, 24, 38)) ) + psd = pspecdata.PSpecData( + dsets=[ + self.d[0], + self.d[1], + ], + wgts=[ + self.w[0], + self.w[1], + ], + labels=["red", "blue"], + ) + np.testing.assert_array_equal(psd.x(("red", 24, 38)), psd.x((0, 24, 38))) # Check specifying labels using dicts - dsdict = {'a':self.d[0], 'b':self.d[1]} + dsdict = {"a": self.d[0], "b": self.d[1]} psd = pspecdata.PSpecData(dsets=dsdict, wgts=dsdict) - pytest.raises(ValueError, pspecdata.PSpecData, dsets=dsdict, - wgts=dsdict, labels=['a', 'b']) + pytest.raises( + ValueError, + pspecdata.PSpecData, + dsets=dsdict, + wgts=dsdict, + labels=["a", "b"], + ) # Check that invalid labels raise errors - pytest.raises(KeyError, psd.x, ('green', 24, 38)) + pytest.raises(KeyError, psd.x, ("green", 24, 38)) def test_parse_blkey(self): # make a double-pol UVData @@ -304,37 +333,41 @@ def test_parse_blkey(self): uvd.polarization_array[0] = -7 uvd = uvd + self.uvd # check parse_blkey - ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], labels=['red', 'blue']) + ds = pspecdata.PSpecData( + dsets=[uvd, uvd], wgts=[None, None], labels=["red", "blue"] + ) dset, bl = ds.parse_blkey((0, (24, 25))) assert dset == 0 assert bl == (24, 25) - dset, bl = ds.parse_blkey(('red', (24, 25), 'xx')) + dset, bl = ds.parse_blkey(("red", (24, 25), "xx")) assert dset == 0 - assert bl == (24, 25, 'xx') + assert bl == (24, 25, "xx") # check PSpecData.x works - assert ds.x(('red', (24, 25))).shape == (2, 64, 60) - assert ds.x(('red', (24, 25), 'xx')).shape == (64, 60) - assert ds.w(('red', (24, 25))).shape == (2, 64, 60) - assert ds.w(('red', (24, 25), 'xx')).shape == (64, 60) + assert ds.x(("red", (24, 25))).shape == (2, 64, 60) + assert ds.x(("red", (24, 25), "xx")).shape == (64, 60) + assert ds.w(("red", (24, 25))).shape == (2, 64, 60) + assert ds.w(("red", (24, 25), "xx")).shape == (64, 60) def test_str(self): """ Check that strings can be output. """ ds = pspecdata.PSpecData() - print(ds) # print empty psd + print(ds) # print empty psd ds.add(self.uvd, None) - print(ds) # print populated psd + print(ds) # print populated psd def test_get_Q_alt(self): """ Test the Q = dC/dp function. """ vect_length = 50 - x_vect = np.random.normal(size=vect_length) \ - + 1.j * np.random.normal(size=vect_length) - y_vect = np.random.normal(size=vect_length) \ - + 1.j * np.random.normal(size=vect_length) + x_vect = np.random.normal(size=vect_length) + 1.0j * np.random.normal( + size=vect_length + ) + y_vect = np.random.normal(size=vect_length) + 1.0j * np.random.normal( + size=vect_length + ) self.ds.spw_Nfreqs = vect_length @@ -355,18 +388,18 @@ def test_get_Q_alt(self): self.assertAlmostEqual(xQy, np.conjugate(yQx)) # x^t Q x should be real - self.assertAlmostEqual(np.imag(xQx), 0.) + self.assertAlmostEqual(np.imag(xQx), 0.0) x_vect = np.ones(vect_length) - Q_matrix = self.ds.get_Q_alt(vect_length//2) + Q_matrix = self.ds.get_Q_alt(vect_length // 2) xQx = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, x_vect)) - self.assertAlmostEqual(xQx, np.abs(vect_length**2.)) + self.assertAlmostEqual(xQx, np.abs(vect_length**2.0)) # Sending in sinusoids for x and y should give delta functions # Now do all the same tests from above but for a different number # of delay channels - self.ds.set_Ndlys(vect_length-3) - for i in range(vect_length-3): + self.ds.set_Ndlys(vect_length - 3) + for i in range(vect_length - 3): Q_matrix = self.ds.get_Q_alt(i) xQy = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, y_vect)) yQx = np.dot(np.conjugate(y_vect), np.dot(Q_matrix, x_vect)) @@ -379,17 +412,17 @@ def test_get_Q_alt(self): self.assertAlmostEqual(xQy, np.conjugate(yQx)) # x^t Q x should be real - self.assertAlmostEqual(np.imag(xQx), 0.) + self.assertAlmostEqual(np.imag(xQx), 0.0) x_vect = np.ones(vect_length) - Q_matrix = self.ds.get_Q_alt((vect_length-2)//2-1) + Q_matrix = self.ds.get_Q_alt((vect_length - 2) // 2 - 1) xQx = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, x_vect)) - self.assertAlmostEqual(xQx, np.abs(vect_length**2.)) + self.assertAlmostEqual(xQx, np.abs(vect_length**2.0)) # Sending in sinusoids for x and y should give delta functions # Make sure that error is raised when asking for a delay mode outside # of the range of delay bins - pytest.raises(IndexError, self.ds.get_Q_alt, vect_length-1) + pytest.raises(IndexError, self.ds.get_Q_alt, vect_length - 1) # Ensure that in the special case where the number of channels equals # the number of delay bins, the FFT method gives the same answer as @@ -403,8 +436,7 @@ def test_get_Q_alt(self): self.assertLessEqual(Q_diff_norm, multiplicative_tolerance) # Check for error handling - pytest.raises(ValueError, self.ds.set_Ndlys, vect_length+100) - + pytest.raises(ValueError, self.ds.set_Ndlys, vect_length + 100) def test_get_Q(self): """ @@ -419,17 +451,17 @@ def test_get_Q(self): """ vect_length = 50 - x_vect = np.random.normal(size=vect_length) \ - + 1.j * np.random.normal(size=vect_length) - y_vect = np.random.normal(size=vect_length) \ - + 1.j * np.random.normal(size=vect_length) + x_vect = np.random.normal(size=vect_length) + 1.0j * np.random.normal( + size=vect_length + ) + y_vect = np.random.normal(size=vect_length) + 1.0j * np.random.normal( + size=vect_length + ) self.ds.spw_Nfreqs = vect_length - #Test if there is a warning if user does not pass the beam - key1 = (0, 24, 38) - key2 = (1, 24, 38) + # Test if there is a warning if user does not pass the beam uvd = copy.deepcopy(self.uvd) - ds_t = pspecdata.PSpecData(dsets=[uvd, uvd]) + pspecdata.PSpecData(dsets=[uvd, uvd]) for i in range(vect_length): try: @@ -451,24 +483,24 @@ def test_get_Q(self): self.assertAlmostEqual(xQy, np.conjugate(yQx)) # x^t Q x should be real - self.assertAlmostEqual(np.imag(xQx), 0.) + self.assertAlmostEqual(np.imag(xQx), 0.0) x_vect = np.ones(vect_length) try: - Q_matrix = self.ds.get_Q(vect_length/2) + Q_matrix = self.ds.get_Q(vect_length / 2) except IndexError: Q_matrix = np.ones((vect_length, vect_length)) xQx = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, x_vect)) - self.assertAlmostEqual(xQx, np.abs(vect_length**2.)) + self.assertAlmostEqual(xQx, np.abs(vect_length**2.0)) # Now do all the same tests from above but for a different number # of delay channels - self.ds.set_Ndlys(vect_length-3) - for i in range(vect_length-3): + self.ds.set_Ndlys(vect_length - 3) + for i in range(vect_length - 3): try: Q_matrix = self.ds.get_Q(i) except IndexError: - Q_matrix = np.ones((vect_length,vect_length)) + Q_matrix = np.ones((vect_length, vect_length)) xQy = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, y_vect)) yQx = np.dot(np.conjugate(y_vect), np.dot(Q_matrix, x_vect)) xQx = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, x_vect)) @@ -480,26 +512,26 @@ def test_get_Q(self): self.assertAlmostEqual(xQy, np.conjugate(yQx)) # x^t Q x should be real - self.assertAlmostEqual(np.imag(xQx), 0.) + self.assertAlmostEqual(np.imag(xQx), 0.0) x_vect = np.ones(vect_length) try: - Q_matrix = self.ds.get_Q((vect_length-2)/2-1) + Q_matrix = self.ds.get_Q((vect_length - 2) / 2 - 1) except IndexError: - Q_matrix = np.ones((vect_length,vect_length)) + Q_matrix = np.ones((vect_length, vect_length)) xQx = np.dot(np.conjugate(x_vect), np.dot(Q_matrix, x_vect)) - self.assertAlmostEqual(xQx, np.abs(vect_length**2.)) + self.assertAlmostEqual(xQx, np.abs(vect_length**2.0)) # Make sure that error is raised when asking for a delay mode outside # of the range of delay bins - pytest.raises(IndexError, self.ds.get_Q, vect_length-1) + pytest.raises(IndexError, self.ds.get_Q, vect_length - 1) def test_get_integral_beam(self): """ Test the integral of the beam and tapering function in Q. """ - pol = 'xx' - #Test if there is a warning if user does not pass the beam + pol = "xx" + # Test if there is a warning if user does not pass the beam uvd = copy.deepcopy(self.uvd) ds_t = pspecdata.PSpecData(dsets=[uvd, uvd]) ds = pspecdata.PSpecData(dsets=[uvd, uvd], beam=self.bm) @@ -525,44 +557,54 @@ def test_get_unnormed_E(self): """ # Test that error is raised if spw_Ndlys is not set uvd = copy.deepcopy(self.uvd) - ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], labels=['red', 'blue']) + ds = pspecdata.PSpecData( + dsets=[uvd, uvd], wgts=[None, None], labels=["red", "blue"] + ) ds.spw_Ndlys = None - pytest.raises(ValueError, ds.get_unnormed_E, 'placeholder', 'placeholder') + pytest.raises(ValueError, ds.get_unnormed_E, "placeholder", "placeholder") # Test that if R1 = R2, then the result is Hermitian ds.spw_Ndlys = 7 random_R = generate_pos_def_all_pos(ds.spw_Nfreqs) - wgt_matrix_dict = {} # The keys here have no significance except they are formatted right - wgt_matrix_dict[('red', (24, 25))] = random_R - wgt_matrix_dict[('blue', (24, 25))] = random_R + wgt_matrix_dict = ( + {} + ) # The keys here have no significance except they are formatted right + wgt_matrix_dict[("red", (24, 25))] = random_R + wgt_matrix_dict[("blue", (24, 25))] = random_R ds.set_R(wgt_matrix_dict) - E_matrices = ds.get_unnormed_E(('red', (24, 25)), ('blue', (24, 25))) + E_matrices = ds.get_unnormed_E(("red", (24, 25)), ("blue", (24, 25))) multiplicative_tolerance = 0.0000001 for matrix in E_matrices: diff_norm = np.linalg.norm(matrix.T.conj() - matrix) self.assertLessEqual(diff_norm, multiplicative_tolerance) - #Test for the correct shape when exact_norm is True - ds_c = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], labels=['red', 'blue'], beam=self.bm) + # Test for the correct shape when exact_norm is True + ds_c = pspecdata.PSpecData( + dsets=[uvd, uvd], wgts=[None, None], labels=["red", "blue"], beam=self.bm + ) ds_c.spw_Ndlys = 10 random_R = generate_pos_def_all_pos(ds_c.spw_Nfreqs) wgt_matrix_dict = {} - wgt_matrix_dict[('red', (24, 25))] = random_R - wgt_matrix_dict[('blue', (24, 25))] = random_R + wgt_matrix_dict[("red", (24, 25))] = random_R + wgt_matrix_dict[("blue", (24, 25))] = random_R ds_c.set_R(wgt_matrix_dict) - E_matrices = ds_c.get_unnormed_E(('red', (24, 25)), ('blue', (24, 25)), exact_norm=True, pol='xx') - self.assertEqual(E_matrices.shape, (ds_c.spw_Ndlys, ds_c.spw_Nfreqs, ds_c.spw_Nfreqs)) + E_matrices = ds_c.get_unnormed_E( + ("red", (24, 25)), ("blue", (24, 25)), exact_norm=True, pol="xx" + ) + self.assertEqual( + E_matrices.shape, (ds_c.spw_Ndlys, ds_c.spw_Nfreqs, ds_c.spw_Nfreqs) + ) # Test that if R1 != R2, then i) E^{12,dagger} = E^{21} random_R2 = generate_pos_def_all_pos(ds.spw_Nfreqs) wgt_matrix_dict = {} - wgt_matrix_dict[('red', (24, 25))] = random_R - wgt_matrix_dict[('blue', (24, 25))] = random_R2 + wgt_matrix_dict[("red", (24, 25))] = random_R + wgt_matrix_dict[("blue", (24, 25))] = random_R2 ds.set_R(wgt_matrix_dict) - E12_matrices = ds.get_unnormed_E(('red', (24, 25)), ('blue', (24, 25))) - E21_matrices = ds.get_unnormed_E(('blue', (24, 25)), ('red', (24, 25))) + E12_matrices = ds.get_unnormed_E(("red", (24, 25)), ("blue", (24, 25))) + E21_matrices = ds.get_unnormed_E(("blue", (24, 25)), ("red", (24, 25))) multiplicative_tolerance = 0.0000001 - for mat12,mat21 in zip(E12_matrices,E21_matrices): + for mat12, mat21 in zip(E12_matrices, E21_matrices): diff_norm = np.linalg.norm(mat12.T.conj() - mat21) self.assertLessEqual(diff_norm, multiplicative_tolerance) @@ -570,29 +612,33 @@ def test_get_unnormed_E(self): # the E matrices are all 0.5s exept in flagged channels. ds.spw_Ndlys = 1 wgt_matrix_dict = {} - wgt_matrix_dict[('red', (24, 25))] = np.eye(ds.spw_Nfreqs) - wgt_matrix_dict[('blue', (24, 25))] = np.eye(ds.spw_Nfreqs) - flags1 = np.diag(ds.Y(('red', (24, 25)))) - flags2 = np.diag(ds.Y(('blue', (24, 25)))) + wgt_matrix_dict[("red", (24, 25))] = np.eye(ds.spw_Nfreqs) + wgt_matrix_dict[("blue", (24, 25))] = np.eye(ds.spw_Nfreqs) + flags1 = np.diag(ds.Y(("red", (24, 25)))) + flags2 = np.diag(ds.Y(("blue", (24, 25)))) ds.set_R(wgt_matrix_dict) - E_matrices = ds.get_unnormed_E(('red', (24, 25)), ('blue', (24, 25))) + E_matrices = ds.get_unnormed_E(("red", (24, 25)), ("blue", (24, 25))) multiplicative_tolerance = 0.0000001 for matrix in E_matrices: for i in range(ds.spw_Nfreqs): for j in range(ds.spw_Nfreqs): - if flags1[i] * flags2[j] == 0: # either channel flagged - self.assertAlmostEqual(matrix[i,j], 0.) + if flags1[i] * flags2[j] == 0: # either channel flagged + self.assertAlmostEqual(matrix[i, j], 0.0) else: - self.assertAlmostEqual(matrix[i,j], 0.5) + self.assertAlmostEqual(matrix[i, j], 0.5) def test_cross_covar_model(self): uvd = copy.deepcopy(self.uvd) - ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], labels=['red', 'blue']) - key1 = ('red', (24, 25), 'xx') - key2 = ('blue', (25, 38), 'xx') + ds = pspecdata.PSpecData( + dsets=[uvd, uvd], wgts=[None, None], labels=["red", "blue"] + ) + key1 = ("red", (24, 25), "xx") + key2 = ("blue", (25, 38), "xx") - pytest.raises(ValueError, ds.cross_covar_model, key1, key2, model='other_string') - pytest.raises(AssertionError, ds.cross_covar_model, key1, 'a_string') + pytest.raises( + ValueError, ds.cross_covar_model, key1, key2, model="other_string" + ) + pytest.raises(AssertionError, ds.cross_covar_model, key1, "a_string") conj1_conj1 = ds.cross_covar_model(key1, key1, conj_1=True, conj_2=True) conj1_real1 = ds.cross_covar_model(key1, key1, conj_1=True, conj_2=False) @@ -605,14 +651,13 @@ def test_cross_covar_model(self): for j in range(ds.spw_Nfreqs): for k in range(ds.spw_Nfreqs): # Check that the matrices that ought to be Hermitian are indeed Hermitian - self.assertAlmostEqual(conj1_real1.conj()[k,j], conj1_real1[j,k]) - self.assertAlmostEqual(real1_conj1.conj()[k,j], real1_conj1[j,k]) + self.assertAlmostEqual(conj1_real1.conj()[k, j], conj1_real1[j, k]) + self.assertAlmostEqual(real1_conj1.conj()[k, j], real1_conj1[j, k]) # Check that real_real and conj_conj are complex conjugates of each other # Also check that they are symmetric - self.assertAlmostEqual(real1_real1.conj()[j,k], conj1_conj1[j,k]) - self.assertAlmostEqual(real1_real1[k,j], real1_real1[j,k]) - self.assertAlmostEqual(conj1_conj1[k,j], conj1_conj1[j,k]) - + self.assertAlmostEqual(real1_real1.conj()[j, k], conj1_conj1[j, k]) + self.assertAlmostEqual(real1_real1[k, j], real1_real1[j, k]) + self.assertAlmostEqual(conj1_conj1[k, j], conj1_conj1[j, k]) real1_real2 = ds.cross_covar_model(key1, key2, conj_1=False, conj_2=False) real2_real1 = ds.cross_covar_model(key2, key1, conj_1=False, conj_2=False) @@ -626,27 +671,27 @@ def test_cross_covar_model(self): # And some similar tests for cross covariances for j in range(ds.spw_Nfreqs): for k in range(ds.spw_Nfreqs): - self.assertAlmostEqual(real1_real2[k,j], real2_real1[j,k]) - self.assertAlmostEqual(conj1_conj2[k,j], conj2_conj1[j,k]) - self.assertAlmostEqual(conj1_real2.conj()[k,j], conj2_real1[j,k]) - self.assertAlmostEqual(real1_conj2.conj()[k,j], real2_conj1[j,k]) + self.assertAlmostEqual(real1_real2[k, j], real2_real1[j, k]) + self.assertAlmostEqual(conj1_conj2[k, j], conj2_conj1[j, k]) + self.assertAlmostEqual(conj1_real2.conj()[k, j], conj2_real1[j, k]) + self.assertAlmostEqual(real1_conj2.conj()[k, j], real2_conj1[j, k]) def test_get_unnormed_V(self): - self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w, labels=['red', 'blue']) - key1 = ('red', (24, 25), 'xx') - key2 = ('blue', (25, 38), 'xx') + self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w, labels=["red", "blue"]) + key1 = ("red", (24, 25), "xx") + key2 = ("blue", (25, 38), "xx") self.ds.spw_Ndlys = 5 V = self.ds.get_unnormed_V(key1, key2) # Check size - self.assertEqual(V.shape, (self.ds.spw_Ndlys,self.ds.spw_Ndlys)) + self.assertEqual(V.shape, (self.ds.spw_Ndlys, self.ds.spw_Ndlys)) # Test hermiticity. Generally this is only good to about 1 part in 10^15. # If this is an issue downstream, should investigate more in the future. tol = 1e-10 - frac_non_herm = abs(V.conj().T - V)/abs(V) + frac_non_herm = abs(V.conj().T - V) / abs(V) for i in range(self.ds.spw_Ndlys): for j in range(self.ds.spw_Ndlys): - self.assertLessEqual(frac_non_herm[i,j], tol) + self.assertLessEqual(frac_non_herm[i, j], tol) def test_get_MW(self): n = 17 @@ -654,62 +699,72 @@ def test_get_MW(self): random_H = generate_pos_def_all_pos(n) random_V = generate_pos_def_all_pos(n) - pytest.raises(AssertionError, self.ds.get_MW, random_G, random_H, mode='L^3') - pytest.raises(NotImplementedError, self.ds.get_MW, random_G, random_H, mode='H^-1', exact_norm=True) - - for mode in ['H^-1', 'V^-1/2', 'I', 'L^-1']: - if mode == 'H^-1': + pytest.raises(AssertionError, self.ds.get_MW, random_G, random_H, mode="L^3") + pytest.raises( + NotImplementedError, + self.ds.get_MW, + random_G, + random_H, + mode="H^-1", + exact_norm=True, + ) + + for mode in ["H^-1", "V^-1/2", "I", "L^-1"]: + if mode == "H^-1": # Test that if we have full-rank matrices, the resulting window functions # are indeed delta functions M, W = self.ds.get_MW(random_G, random_H, mode=mode) Hinv = np.linalg.inv(random_H) for i in range(n): - self.assertAlmostEqual(W[i,i], 1.) + self.assertAlmostEqual(W[i, i], 1.0) for j in range(n): - self.assertAlmostEqual(M[i,j], Hinv[i,j]) + self.assertAlmostEqual(M[i, j], Hinv[i, j]) # When the matrices are not full rank, test that the window functions # are at least properly normalized. - deficient_H = np.ones((3,3)) + deficient_H = np.ones((3, 3)) M, W = self.ds.get_MW(deficient_H, deficient_H, mode=mode) norm = np.sum(W, axis=1) for i in range(3): - self.assertAlmostEqual(norm[i], 1.) + self.assertAlmostEqual(norm[i], 1.0) # Check that the method ignores G M, W = self.ds.get_MW(random_G, random_H, mode=mode) M_other, W_other = self.ds.get_MW(random_H, random_H, mode=mode) for i in range(n): for j in range(n): - self.assertAlmostEqual(M[i,j], M_other[i,j]) - self.assertAlmostEqual(W[i,j], W_other[i,j]) + self.assertAlmostEqual(M[i, j], M_other[i, j]) + self.assertAlmostEqual(W[i, j], W_other[i, j]) - elif mode == 'V^-1/2': + elif mode == "V^-1/2": # Test that we are checking for the presence of a covariance matrix pytest.raises(ValueError, self.ds.get_MW, random_G, random_H, mode=mode) # Test that the error covariance is diagonal - M, W = self.ds.get_MW(random_G, random_H, mode=mode, band_covar=random_V) + M, W = self.ds.get_MW( + random_G, random_H, mode=mode, band_covar=random_V + ) band_covar = np.dot(M, np.dot(random_V, M.T)) self.assertEqual(diagonal_or_not(band_covar), True) - elif mode == 'I': + elif mode == "I": # Test that the norm matrix is diagonal M, W = self.ds.get_MW(random_G, random_H, mode=mode) self.assertEqual(diagonal_or_not(M), True) - elif mode == 'L^-1': + elif mode == "L^-1": # Test that Cholesky mode is disabled - pytest.raises(NotImplementedError, - self.ds.get_MW, random_G, random_H, mode=mode) + pytest.raises( + NotImplementedError, self.ds.get_MW, random_G, random_H, mode=mode + ) # Test sizes for everyone - self.assertEqual(M.shape, (n,n)) - self.assertEqual(W.shape, (n,n)) + self.assertEqual(M.shape, (n, n)) + self.assertEqual(W.shape, (n, n)) # Window function matrices should have each row sum to unity # regardless of the mode chosen test_norm = np.sum(W, axis=1) for norm in test_norm: - self.assertAlmostEqual(norm, 1.) + self.assertAlmostEqual(norm, 1.0) def test_cov_q(self, ndlys=13): """ @@ -717,72 +772,110 @@ def test_cov_q(self, ndlys=13): format. Also validate with arbitrary number of delays. """ for d in self.d: - d.flag_array[:] = False #ensure that there are no flags! - d.select(times=np.unique(d.time_array)[:10], frequencies=d.freq_array[0, :16]) + d.flag_array[:] = False # ensure that there are no flags! + d.select( + times=np.unique(d.time_array)[:10], frequencies=d.freq_array[0, :16] + ) for d_std in self.d_std: d_std.flag_array[:] = False - d_std.select(times=np.unique(d_std.time_array)[:10], frequencies=d_std.freq_array[0, :16]) + d_std.select( + times=np.unique(d_std.time_array)[:10], + frequencies=d_std.freq_array[0, :16], + ) self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w, dsets_std=self.d_std) self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w, dsets_std=self.d_std) - Ntime = self.ds.Ntimes self.ds.set_Ndlys(ndlys) # Here is the analytic covariance matrix... chan_x, chan_y = np.meshgrid(range(self.ds.Nfreqs), range(self.ds.Nfreqs)) - cov_analytic = np.zeros((self.ds.spw_Ndlys, self.ds.spw_Ndlys), dtype=np.complex128) + cov_analytic = np.zeros( + (self.ds.spw_Ndlys, self.ds.spw_Ndlys), dtype=np.complex128 + ) for alpha in range(self.ds.spw_Ndlys): for beta in range(self.ds.spw_Ndlys): - cov_analytic[alpha, beta] = np.exp(-2j*np.pi*(alpha-beta)*(chan_x-chan_y)/self.ds.spw_Ndlys).sum() + cov_analytic[alpha, beta] = np.exp( + -2j * np.pi * (alpha - beta) * (chan_x - chan_y) / self.ds.spw_Ndlys + ).sum() key1 = (0, 24, 38) key2 = (1, 25, 38) - #print(cov_analytic) + # print(cov_analytic) - for input_data_weight in ['identity','iC', 'dayenu']: + for input_data_weight in ["identity", "iC", "dayenu"]: self.ds.set_weighting(input_data_weight) - #check error raised - if input_data_weight == 'dayenu': - pytest.raises(ValueError,self.ds.R, key1) - rpk = {'filter_centers':[0.],'filter_half_widths':[0.],'filter_factors':[0.]} - self.ds.set_r_param(key1,rpk) - self.ds.set_r_param(key2,rpk) - for taper in taper_selection: - qc = self.ds.cov_q_hat(key1,key2,model='dsets') - self.assertTrue(np.allclose(np.array(list(qc.shape)), - np.array([self.ds.Ntimes, self.ds.spw_Ndlys, self.ds.spw_Ndlys]), atol=1e-6)) - qc = self.ds.cov_q_hat(key1,key2,model='empirical') - self.assertTrue(np.allclose(np.array(list(qc.shape)), - np.array([self.ds.Ntimes, self.ds.spw_Ndlys, self.ds.spw_Ndlys]), atol=1e-6)) + # check error raised + if input_data_weight == "dayenu": + pytest.raises(ValueError, self.ds.R, key1) + rpk = { + "filter_centers": [0.0], + "filter_half_widths": [0.0], + "filter_factors": [0.0], + } + self.ds.set_r_param(key1, rpk) + self.ds.set_r_param(key2, rpk) + for _ in taper_selection: + qc = self.ds.cov_q_hat(key1, key2, model="dsets") + self.assertTrue( + np.allclose( + np.array(list(qc.shape)), + np.array( + [self.ds.Ntimes, self.ds.spw_Ndlys, self.ds.spw_Ndlys] + ), + atol=1e-6, + ) + ) + qc = self.ds.cov_q_hat(key1, key2, model="empirical") + self.assertTrue( + np.allclose( + np.array(list(qc.shape)), + np.array( + [self.ds.Ntimes, self.ds.spw_Ndlys, self.ds.spw_Ndlys] + ), + atol=1e-6, + ) + ) """ Now test that analytic Error calculation gives Nchan^2 """ - self.ds.set_weighting('identity') - qc = self.ds.cov_q_hat(key1, key2, model='dsets') - self.assertTrue(np.allclose(qc, - np.repeat(cov_analytic[np.newaxis, :, :], self.ds.Ntimes, axis=0), atol=1e-6)) + self.ds.set_weighting("identity") + qc = self.ds.cov_q_hat(key1, key2, model="dsets") + self.assertTrue( + np.allclose( + qc, + np.repeat(cov_analytic[np.newaxis, :, :], self.ds.Ntimes, axis=0), + atol=1e-6, + ) + ) """ Test lists of keys """ - self.ds.set_weighting('identity') - qc=self.ds.cov_q_hat([key1], [key2], time_indices=[0], model='dsets') - self.assertTrue(np.allclose(qc, - np.repeat(cov_analytic[np.newaxis, :, :], self.ds.Ntimes, axis=0), atol=1e-6)) + self.ds.set_weighting("identity") + qc = self.ds.cov_q_hat([key1], [key2], time_indices=[0], model="dsets") + self.assertTrue( + np.allclose( + qc, + np.repeat(cov_analytic[np.newaxis, :, :], self.ds.Ntimes, axis=0), + atol=1e-6, + ) + ) self.assertRaises(ValueError, self.ds.cov_q_hat, key1, key2, time_indices=200) - self.assertRaises(ValueError, self.ds.cov_q_hat, key1, key2, time_indices="watch out!") - + self.assertRaises( + ValueError, self.ds.cov_q_hat, key1, key2, time_indices="watch out!" + ) def test_cov_p_hat(self): """ Test cov_p_hat, verify on identity. """ self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w, dsets_std=self.d_std) - cov_p = self.ds.cov_p_hat(np.sqrt(6.)*np.identity(10),np.array([5.*np.identity(10)])) + cov_p = self.ds.cov_p_hat( + np.sqrt(6.0) * np.identity(10), np.array([5.0 * np.identity(10)]) + ) for p in range(10): for q in range(10): if p == q: - self.assertTrue(np.isclose(30., cov_p[0, p, q], atol=1e-6)) + self.assertTrue(np.isclose(30.0, cov_p[0, p, q], atol=1e-6)) else: - self.assertTrue(np.isclose(0., cov_p[0, p, q], atol=1e-6)) - + self.assertTrue(np.isclose(0.0, cov_p[0, p, q], atol=1e-6)) def test_R_truncation(self): """ @@ -791,49 +884,51 @@ def test_R_truncation(self): """ self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w) Nfreq = self.ds.spw_Nfreqs - Ntime = self.ds.Ntimes Ndlys = Nfreq - 3 self.ds.spw_Ndlys = Ndlys - # Set baselines to use for tests key1 = (0, 24, 38) key2 = (1, 25, 38) - key3 = [(0, 24, 38), (0, 24, 38)] - key4 = [(1, 25, 38), (1, 25, 38)] - - rpk1 = {'filter_centers':[0.],'filter_half_widths':[100e-9],'filter_factors':[1e-9]} - rpk2 = {'filter_centers':[0.],'filter_half_widths':[100e-9],'filter_factors':[1e-9]} - self.ds.set_weighting('dayenu') - self.ds.set_r_param(key1,rpk1) - self.ds.set_r_param(key2,rpk2) + + rpk1 = { + "filter_centers": [0.0], + "filter_half_widths": [100e-9], + "filter_factors": [1e-9], + } + rpk2 = { + "filter_centers": [0.0], + "filter_half_widths": [100e-9], + "filter_factors": [1e-9], + } + self.ds.set_weighting("dayenu") + self.ds.set_r_param(key1, rpk1) + self.ds.set_r_param(key2, rpk2) ds1 = copy.deepcopy(self.ds) - ds1.set_spw((10,Nfreq-10)) + ds1.set_spw((10, Nfreq - 10)) ds1.set_symmetric_taper(False) - ds1.set_filter_extension([10,10]) - ds1.set_filter_extension((10,10)) + ds1.set_filter_extension([10, 10]) + ds1.set_filter_extension((10, 10)) rm1 = self.ds.R(key1) rm2 = ds1.R(key2) - rm3 = ds1.R(key1) self.assertTrue(np.shape(rm2) == (ds1.spw_Nfreqs, self.ds.spw_Nfreqs)) - #check that all values that are not truncated match values of untrancated matrix. + # check that all values that are not truncated match values of untrancated matrix. self.assertTrue(np.all(np.isclose(rm1[10:-10], rm2, atol=1e-6))) - #make sure no errors are thrown by get_V, get_E, etc... + # make sure no errors are thrown by get_V, get_E, etc... ds1.get_unnormed_E(key1, key2) ds1.get_unnormed_V(key1, key2) - h=ds1.get_H(key1, key2) - g=ds1.get_G(key1, key2) + h = ds1.get_H(key1, key2) + g = ds1.get_G(key1, key2) ds1.get_MW(g, h) - #make sure identity weighting isn't broken. + # make sure identity weighting isn't broken. self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w) ds1 = copy.deepcopy(self.ds) - ds1.set_spw((10,Nfreq-10)) - ds1.set_weighting('identity') + ds1.set_spw((10, Nfreq - 10)) + ds1.set_weighting("identity") ds1.set_symmetric_taper(False) - ds1.set_filter_extension([10,10]) + ds1.set_filter_extension([10, 10]) rm1 = ds1.R(key1) - def test_q_hat(self): """ Test that q_hat has right shape and accepts keys in the right format. @@ -845,20 +940,23 @@ def test_q_hat(self): Ndlys = Nfreq - 3 self.ds.spw_Ndlys = Ndlys - # Set baselines to use for tests key1 = (0, 24, 38) key2 = (1, 25, 38) key3 = [(0, 24, 38), (0, 24, 38)] key4 = [(1, 25, 38), (1, 25, 38)] - for input_data_weight in ['identity', 'iC', 'dayenu']: + for input_data_weight in ["identity", "iC", "dayenu"]: self.ds.set_weighting(input_data_weight) - if input_data_weight == 'dayenu': - pytest.raises(ValueError,self.ds.R, key1) - rpk = {'filter_centers':[0.],'filter_half_widths':[0.],'filter_factors':[0.]} - self.ds.set_r_param(key1,rpk) - self.ds.set_r_param(key2,rpk) + if input_data_weight == "dayenu": + pytest.raises(ValueError, self.ds.R, key1) + rpk = { + "filter_centers": [0.0], + "filter_half_widths": [0.0], + "filter_factors": [0.0], + } + self.ds.set_r_param(key1, rpk) + self.ds.set_r_param(key2, rpk) # Loop over list of taper functions for taper in taper_selection: self.ds.set_taper(taper) @@ -867,35 +965,37 @@ def test_q_hat(self): q_hat_a = self.ds.q_hat(key1, key2) self.assertEqual(q_hat_a.shape, (Ndlys, Ntime)) - # Check that swapping x_1 <-> x_2 results in complex conj. only q_hat_b = self.ds.q_hat(key2, key1) q_hat_diff = np.conjugate(q_hat_a) - q_hat_b for i in range(Ndlys): for j in range(Ntime): - self.assertAlmostEqual(q_hat_diff[i,j].real, - q_hat_diff[i,j].real) - self.assertAlmostEqual(q_hat_diff[i,j].imag, - q_hat_diff[i,j].imag) + self.assertAlmostEqual( + q_hat_diff[i, j].real, q_hat_diff[i, j].real + ) + self.assertAlmostEqual( + q_hat_diff[i, j].imag, q_hat_diff[i, j].imag + ) # Check that lists of keys are handled properly - q_hat_aa = self.ds.q_hat(key1, key4) # q_hat(x1, x2+x2) - q_hat_bb = self.ds.q_hat(key4, key1) # q_hat(x2+x2, x1) - q_hat_cc = self.ds.q_hat(key3, key4) # q_hat(x1+x1, x2+x2) + self.ds.q_hat(key1, key4) # q_hat(x1, x2+x2) + self.ds.q_hat(key4, key1) # q_hat(x2+x2, x1) + q_hat_cc = self.ds.q_hat(key3, key4) # q_hat(x1+x1, x2+x2) # Effectively checks that q_hat(2*x1, 2*x2) = 4*q_hat(x1, x2) for i in range(Ndlys): for j in range(Ntime): - self.assertAlmostEqual(q_hat_a[i,j].real, - 0.25 * q_hat_cc[i,j].real) - self.assertAlmostEqual(q_hat_a[i,j].imag, - 0.25 * q_hat_cc[i,j].imag) - + self.assertAlmostEqual( + q_hat_a[i, j].real, 0.25 * q_hat_cc[i, j].real + ) + self.assertAlmostEqual( + q_hat_a[i, j].imag, 0.25 * q_hat_cc[i, j].imag + ) self.ds.spw_Ndlys = Nfreq # Check that the slow method is the same as the FFT method - for input_data_weight in ['identity', 'iC', 'dayenu']: + for input_data_weight in ["identity", "iC", "dayenu"]: self.ds.set_weighting(input_data_weight) # Loop over list of taper functions for taper in taper_selection: @@ -903,11 +1003,20 @@ def test_q_hat(self): self.ds.set_taper(taper) q_hat_a_slow = self.ds.q_hat(key1, key2, allow_fft=False) q_hat_a = self.ds.q_hat(key1, key2, allow_fft=True) - self.assertTrue(np.isclose(np.real(q_hat_a/q_hat_a_slow), 1).all()) - self.assertTrue(np.isclose(np.imag(q_hat_a/q_hat_a_slow), 0, atol=1e-6).all()) - - #Test if error is raised when one tried FFT approach on exact_norm - pytest.raises(NotImplementedError, self.ds.q_hat, key1, key2, exact_norm=True, allow_fft = True) + self.assertTrue(np.isclose(np.real(q_hat_a / q_hat_a_slow), 1).all()) + self.assertTrue( + np.isclose(np.imag(q_hat_a / q_hat_a_slow), 0, atol=1e-6).all() + ) + + # Test if error is raised when one tried FFT approach on exact_norm + pytest.raises( + NotImplementedError, + self.ds.q_hat, + key1, + key2, + exact_norm=True, + allow_fft=True, + ) def test_get_H(self): """ @@ -915,27 +1024,30 @@ def test_get_H(self): """ self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w) Nfreq = self.ds.Nfreqs - multiplicative_tolerance = 1. key1 = (0, 24, 38) key2 = (1, 25, 38) - for input_data_weight in ['identity','iC', 'dayenu']: + for input_data_weight in ["identity", "iC", "dayenu"]: self.ds.set_weighting(input_data_weight) - if input_data_weight == 'dayenu': - pytest.raises(ValueError,self.ds.R, key1) - rpk = {'filter_centers':[0.],'filter_half_widths':[0.],'filter_factors':[0.]} - self.ds.set_r_param(key1,rpk) - self.ds.set_r_param(key2,rpk) + if input_data_weight == "dayenu": + pytest.raises(ValueError, self.ds.R, key1) + rpk = { + "filter_centers": [0.0], + "filter_half_widths": [0.0], + "filter_factors": [0.0], + } + self.ds.set_r_param(key1, rpk) + self.ds.set_r_param(key2, rpk) for taper in taper_selection: self.ds.set_taper(taper) - self.ds.set_Ndlys(Nfreq//3) + self.ds.set_Ndlys(Nfreq // 3) H = self.ds.get_H(key1, key2) - self.assertEqual(H.shape, (Nfreq//3, Nfreq//3)) # Test shape + self.assertEqual(H.shape, (Nfreq // 3, Nfreq // 3)) # Test shape self.ds.set_Ndlys() H = self.ds.get_H(key1, key2) - self.assertEqual(H.shape, (Nfreq, Nfreq)) # Test shape + self.assertEqual(H.shape, (Nfreq, Nfreq)) # Test shape def test_get_G(self): """ @@ -943,28 +1055,32 @@ def test_get_G(self): """ self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w) Nfreq = self.ds.Nfreqs - multiplicative_tolerance = 1. + multiplicative_tolerance = 1.0 key1 = (0, 24, 38) key2 = (1, 25, 38) - for input_data_weight in ['identity','iC', 'dayenu']: + for input_data_weight in ["identity", "iC", "dayenu"]: self.ds.set_weighting(input_data_weight) - if input_data_weight == 'dayenu': - pytest.raises(ValueError,self.ds.R, key1) - rpk = {'filter_centers':[0.],'filter_half_widths':[0.],'filter_factors':[0.]} - self.ds.set_r_param(key1,rpk) - self.ds.set_r_param(key2,rpk) + if input_data_weight == "dayenu": + pytest.raises(ValueError, self.ds.R, key1) + rpk = { + "filter_centers": [0.0], + "filter_half_widths": [0.0], + "filter_factors": [0.0], + } + self.ds.set_r_param(key1, rpk) + self.ds.set_r_param(key2, rpk) for taper in taper_selection: self.ds.clear_cache() self.ds.set_taper(taper) - #print 'input_data_weight', input_data_weight - self.ds.set_Ndlys(Nfreq-2) + # print 'input_data_weight', input_data_weight + self.ds.set_Ndlys(Nfreq - 2) G = self.ds.get_G(key1, key2) - self.assertEqual(G.shape, (Nfreq-2, Nfreq-2)) # Test shape - #print np.min(np.abs(G)), np.min(np.abs(np.linalg.eigvalsh(G))) + self.assertEqual(G.shape, (Nfreq - 2, Nfreq - 2)) # Test shape + # print np.min(np.abs(G)), np.min(np.abs(np.linalg.eigvalsh(G))) matrix_scale = np.min(np.abs(np.linalg.eigvalsh(G))) - if input_data_weight == 'identity': + if input_data_weight == "identity": # In the identity case, there are three special properties # that are respected: # i) Symmetry: G_ab = G_ba @@ -974,8 +1090,9 @@ def test_get_G(self): # Test symmetry anti_sym_norm = np.linalg.norm(G - G.T) - self.assertLessEqual(anti_sym_norm, - matrix_scale * multiplicative_tolerance) + self.assertLessEqual( + anti_sym_norm, matrix_scale * multiplicative_tolerance + ) # Test cyclic property of trace, where key1 and key2 can be # swapped without changing the matrix. This is secretly the @@ -984,32 +1101,33 @@ def test_get_G(self): # the other. G_swapped = self.ds.get_G(key2, key1) G_diff_norm = np.linalg.norm(G - G_swapped) - self.assertLessEqual(G_diff_norm, - matrix_scale * multiplicative_tolerance) + self.assertLessEqual( + G_diff_norm, matrix_scale * multiplicative_tolerance + ) min_diagonal = np.min(np.diagonal(G)) # Test that all elements of G are positive up to numerical # noise with the threshold set to 10 orders of magnitude # down from the smallest value on the diagonal - for i in range(Nfreq-2): - for j in range(Nfreq-2): - self.assertGreaterEqual(G[i,j], - - -min_diagonal * multiplicative_tolerance) + for i in range(Nfreq - 2): + for j in range(Nfreq - 2): + self.assertGreaterEqual( + G[i, j], -min_diagonal * multiplicative_tolerance + ) else: # In general, when R_1 != R_2, there is a more restricted # symmetry where swapping R_1 and R_2 *and* taking the # transpose gives the same result - #UPDATE: Taper now occurs after filter so this - #symmetry only holds when taper = 'none'. - if taper_selection == 'none': + # UPDATE: Taper now occurs after filter so this + # symmetry only holds when taper = 'none'. + if taper_selection == "none": G_swapped = self.ds.get_G(key2, key1) G_diff_norm = np.linalg.norm(G - G_swapped.T) - self.assertLessEqual(G_diff_norm, - matrix_scale * multiplicative_tolerance) + self.assertLessEqual( + G_diff_norm, matrix_scale * multiplicative_tolerance + ) - - """ + r""" Under Construction def test_parseval(self): # Test that output power spectrum respects Parseval's theorem. @@ -1086,45 +1204,45 @@ def test_scalar_delay_adjustment(self): # Test that when: # i) Nfreqs = Ndlys, ii) Sampling, iii) No tapering, iv) R is identity # are all satisfied, the scalar adjustment factor is unity - self.ds.set_weighting('identity') + self.ds.set_weighting("identity") self.ds.spw_Ndlys = self.ds.spw_Nfreqs adjustment = self.ds.scalar_delay_adjustment(key1, key2, sampling=True) self.assertAlmostEqual(adjustment, 1.0) - self.ds.set_weighting('iC') - #if weighting is not identity, then the adjustment should be a vector. + self.ds.set_weighting("iC") + # if weighting is not identity, then the adjustment should be a vector. adjustment = self.ds.scalar_delay_adjustment(key1, key2, sampling=True) self.assertTrue(len(adjustment) == self.ds.spw_Ndlys) def test_scalar(self): self.ds = pspecdata.PSpecData(dsets=self.d, wgts=self.w, beam=self.bm) - gauss = pspecbeam.PSpecBeamGauss(0.8, - np.linspace(115e6, 130e6, 50, endpoint=False)) - ds2 = pspecdata.PSpecData(dsets=self.d, wgts=self.w, beam=gauss) + gauss = pspecbeam.PSpecBeamGauss( + 0.8, np.linspace(115e6, 130e6, 50, endpoint=False) + ) + pspecdata.PSpecData(dsets=self.d, wgts=self.w, beam=gauss) # Check normal execution - scalar = self.ds.scalar(('xx','xx')) - scalar_xx = self.ds.scalar('xx') # Can use single pol string as shorthand + scalar = self.ds.scalar(("xx", "xx")) + scalar_xx = self.ds.scalar("xx") # Can use single pol string as shorthand assert scalar == scalar_xx - scalar = self.ds.scalar(1515) # polpair-integer = ('xx', 'xx') - scalar = self.ds.scalar(('xx','xx'), taper_override='none') - scalar = self.ds.scalar(('xx','xx'), beam=gauss) - pytest.raises(NotImplementedError, self.ds.scalar, ('xx','yy')) + scalar = self.ds.scalar(1515) # polpair-integer = ('xx', 'xx') + scalar = self.ds.scalar(("xx", "xx"), taper_override="none") + scalar = self.ds.scalar(("xx", "xx"), beam=gauss) + pytest.raises(NotImplementedError, self.ds.scalar, ("xx", "yy")) # Precomputed results in the following test were done "by hand" # using iPython notebook "Scalar_dev2.ipynb" in the tests/ directory # FIXME: Uncomment when pyuvdata support for this is ready - #scalar = self.ds.scalar() - #self.assertAlmostEqual(scalar, 3732415176.85 / 10.**9) + # scalar = self.ds.scalar() + # self.assertAlmostEqual(scalar, 3732415176.85 / 10.**9) # FIXME: Remove this when pyuvdata support for the above is ready - #self.assertRaises(NotImplementedError, self.ds.scalar) + # self.assertRaises(NotImplementedError, self.ds.scalar) def test_validate_datasets(self): # test freq exception uvd = copy.deepcopy(self.d[0]) - uvd2 = uvd.select(frequencies=np.unique(uvd.freq_array)[:10], - inplace=False) + uvd2 = uvd.select(frequencies=np.unique(uvd.freq_array)[:10], inplace=False) ds = pspecdata.PSpecData(dsets=[uvd, uvd2], wgts=[None, None]) pytest.raises(ValueError, ds.validate_datasets) @@ -1154,27 +1272,31 @@ def test_validate_datasets(self): # test warnings uvd = copy.deepcopy(self.d[0]) uvd2 = copy.deepcopy(self.d[0]) - uvd.select(frequencies=np.unique(uvd.freq_array)[:10], - times=np.unique(uvd.time_array)[:10]) - uvd2.select(frequencies=np.unique(uvd2.freq_array)[10:20], - times=np.unique(uvd2.time_array)[10:20]) + uvd.select( + frequencies=np.unique(uvd.freq_array)[:10], + times=np.unique(uvd.time_array)[:10], + ) + uvd2.select( + frequencies=np.unique(uvd2.freq_array)[10:20], + times=np.unique(uvd2.time_array)[10:20], + ) ds = pspecdata.PSpecData(dsets=[uvd, uvd2], wgts=[None, None]) ds.validate_datasets() # test phasing uvd = copy.deepcopy(self.d[0]) uvd2 = copy.deepcopy(self.d[0]) - uvd.phase_to_time(Time(2458042, format='jd')) + uvd.phase_to_time(Time(2458042, format="jd")) ds = pspecdata.PSpecData(dsets=[uvd, uvd2], wgts=[None, None]) pytest.raises(ValueError, ds.validate_datasets) - uvd2.phase_to_time(Time(2458042.5, format='jd')) + uvd2.phase_to_time(Time(2458042.5, format="jd")) ds.validate_datasets() # test polarization - ds.validate_pol((0,1), ('xx', 'xx')) + ds.validate_pol((0, 1), ("xx", "xx")) # test channel widths - uvd2.channel_width *= 2. + uvd2.channel_width *= 2.0 ds2 = pspecdata.PSpecData(dsets=[uvd, uvd2], wgts=[None, None]) pytest.raises(ValueError, ds2.validate_datasets) @@ -1183,66 +1305,91 @@ def test_rephase_to_dset(self): uvd1 = copy.deepcopy(self.uvd) # give the uvd an x_orientation to test x_orientation propagation - uvd1.x_orienation = 'east' + uvd1.x_orienation = "east" # null test: check nothing changes when dsets contain same UVData object - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd1)], wgts=[None, None]) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd1)], wgts=[None, None] + ) # get normal pspec bls = [(37, 39)] - uvp1 = ds.pspec(bls, bls, (0, 1), pols=('xx','xx'), verbose=False) + uvp1 = ds.pspec(bls, bls, (0, 1), pols=("xx", "xx"), verbose=False) # rephase and get pspec ds.rephase_to_dset(0) - uvp2 = ds.pspec(bls, bls, (0, 1), pols=('xx','xx'), verbose=False) - blp = (0, ((37,39),(37,39)), ('xx','xx')) - assert np.isclose(np.abs(uvp2.get_data(blp)/uvp1.get_data(blp)), 1.0).min() + uvp2 = ds.pspec(bls, bls, (0, 1), pols=("xx", "xx"), verbose=False) + blp = (0, ((37, 39), (37, 39)), ("xx", "xx")) + assert np.isclose(np.abs(uvp2.get_data(blp) / uvp1.get_data(blp)), 1.0).min() def test_Jy_to_mK(self): # test basic execution uvd = self.uvd - uvd.vis_units = 'Jy' - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], - wgts=[None, None], beam=self.bm) + uvd.vis_units = "Jy" + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], + wgts=[None, None], + beam=self.bm, + ) ds.Jy_to_mK() - assert ds.dsets[0].vis_units == 'mK' - assert ds.dsets[1].vis_units == 'mK' - assert uvd.get_data(24, 25, 'xx')[30, 30] \ - / ds.dsets[0].get_data(24, 25, 'xx')[30, 30] < 1.0 + assert ds.dsets[0].vis_units == "mK" + assert ds.dsets[1].vis_units == "mK" + assert ( + uvd.get_data(24, 25, "xx")[30, 30] + / ds.dsets[0].get_data(24, 25, "xx")[30, 30] + < 1.0 + ) # test feeding beam - ds2 = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], - wgts=[None, None], beam=self.bm) + ds2 = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], + wgts=[None, None], + beam=self.bm, + ) ds2.Jy_to_mK(beam=self.bm) assert ds.dsets[0] == ds2.dsets[0] # test vis_units no Jansky uvd2 = copy.deepcopy(uvd) uvd2.polarization_array[0] = -6 - uvd2.vis_units = 'UNCALIB' - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd2)], - wgts=[None, None], beam=self.bm) + uvd2.vis_units = "UNCALIB" + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd2)], + wgts=[None, None], + beam=self.bm, + ) ds.Jy_to_mK() assert ds.dsets[0].vis_units == "mK" assert ds.dsets[1].vis_units == "UNCALIB" - assert ds.dsets[0].get_data(24, 25, 'xx')[30, 30] != ds.dsets[1].get_data(24, 25, 'yy')[30, 30] + assert ( + ds.dsets[0].get_data(24, 25, "xx")[30, 30] + != ds.dsets[1].get_data(24, 25, "yy")[30, 30] + ) def test_trim_dset_lsts(self): fname = os.path.join(DATA_PATH, "zen.2458042.17772.xx.HH.uvXA") uvd1 = UVData() uvd1.read_miriad(fname) uvd2 = copy.deepcopy(uvd1) - uvd2.lst_array = (uvd2.lst_array + 10. * np.median(np.diff(np.unique(uvd2.lst_array)))) % (2.*np.pi) + uvd2.lst_array = ( + uvd2.lst_array + 10.0 * np.median(np.diff(np.unique(uvd2.lst_array))) + ) % (2.0 * np.pi) # test basic execution - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd2)], wgts=[None, None]) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd2)], wgts=[None, None] + ) ds.trim_dset_lsts() assert ds.dsets[0].Ntimes == 50 assert ds.dsets[1].Ntimes == 50 - assert np.all( (2458042.178948477 < ds.dsets[0].time_array) \ - + (ds.dsets[0].time_array < 2458042.1843023109)) + assert np.all( + (2458042.178948477 < ds.dsets[0].time_array) + + (ds.dsets[0].time_array < 2458042.1843023109) + ) # test exception uvd2.lst_array += np.linspace(0, 1e-3, uvd2.Nblts) - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd2)], wgts=[None, None]) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd1), copy.deepcopy(uvd2)], wgts=[None, None] + ) pytest.raises(ValueError, ds.trim_dset_lsts) assert ds.dsets[0].Ntimes == 60 assert ds.dsets[1].Ntimes == 60 @@ -1257,8 +1404,9 @@ def test_units(self): vis_u, norm_u = ds.units() assert vis_u == "UNCALIB" assert norm_u == "Hz str [beam normalization not specified]" - ds_b = pspecdata.PSpecData(dsets=[self.uvd, self.uvd], - wgts=[None, None], beam=self.bm) + ds_b = pspecdata.PSpecData( + dsets=[self.uvd, self.uvd], wgts=[None, None], beam=self.bm + ) vis_u, norm_u = ds_b.units(little_h=False) assert norm_u == "Mpc^3" @@ -1276,67 +1424,100 @@ def test_check_in_dset(self): ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None]) # check for existing key - assert ds.check_key_in_dset(('xx'), 0) + assert ds.check_key_in_dset(("xx"), 0) assert ds.check_key_in_dset((24, 25), 0) - assert ds.check_key_in_dset((24, 25, 'xx'), 0) + assert ds.check_key_in_dset((24, 25, "xx"), 0) # check for non-existing key - assert ds.check_key_in_dset('yy', 0) == False - assert ds.check_key_in_dset((24, 26), 0) == False - assert ds.check_key_in_dset((24, 26, 'yy'), 0) == False + assert not ds.check_key_in_dset("yy", 0) + assert not ds.check_key_in_dset((24, 26), 0) + assert not ds.check_key_in_dset((24, 26, "yy"), 0) # check exception - pytest.raises(KeyError, ds.check_key_in_dset, (1,2,3,4,5), 0) + pytest.raises(KeyError, ds.check_key_in_dset, (1, 2, 3, 4, 5), 0) # test dset_idx - pytest.raises(TypeError, ds.dset_idx, (1,2)) + pytest.raises(TypeError, ds.dset_idx, (1, 2)) def test_C_model(self): # test the key format in ds._C and the shape of stored covariance uvd = UVData() - uvd.read(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) + uvd.read(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) cosmo = conversions.Cosmo_Conversions() - uvb = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits'), cosmo=cosmo) + uvb = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits"), cosmo=cosmo + ) ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=uvb) - spws = utils.spw_range_from_freqs(uvd, freq_range=[(160e6, 165e6), (160e6, 165e6)], bounds_error=True) + spws = utils.spw_range_from_freqs( + uvd, freq_range=[(160e6, 165e6), (160e6, 165e6)], bounds_error=True + ) antpos, ants = uvd.get_ENU_antpos(pick_data_ants=True) antpos = dict(zip(ants, antpos)) red_bls = redcal.get_pos_reds(antpos, bl_error_tol=1.0) - bls1, bls2, blpairs = utils.construct_blpairs(red_bls[3], exclude_auto_bls=True, exclude_permutations=True) + bls1, bls2, blpairs = utils.construct_blpairs( + red_bls[3], exclude_auto_bls=True, exclude_permutations=True + ) ds.set_spw(spws[0]) - key = (0,bls1[0],"xx") - ds.C_model(key, model='empirical', time_index=0) - assert( ((0, 0), ((bls1[0][0],bls1[0][1] ,"xx"),(bls1[0][0],bls1[0][1] ,"xx")), 'empirical', None, False, True,) in ds._C.keys()) - ds.C_model(key, model='autos', time_index=0) - assert( ((0, 0), ((bls1[0][0],bls1[0][1] ,"xx"), (bls1[0][0],bls1[0][1] ,"xx")), 'autos', 0, False, True,) in ds._C.keys()) + key = (0, bls1[0], "xx") + ds.C_model(key, model="empirical", time_index=0) + assert ( + (0, 0), + ((bls1[0][0], bls1[0][1], "xx"), (bls1[0][0], bls1[0][1], "xx")), + "empirical", + None, + False, + True, + ) in ds._C.keys() + ds.C_model(key, model="autos", time_index=0) + assert ( + (0, 0), + ((bls1[0][0], bls1[0][1], "xx"), (bls1[0][0], bls1[0][1], "xx")), + "autos", + 0, + False, + True, + ) in ds._C.keys() for Ckey in ds._C.keys(): - assert ds._C[Ckey].shape == (spws[0][1]-spws[0][0], spws[0][1]-spws[0][0]) + assert ds._C[Ckey].shape == ( + spws[0][1] - spws[0][0], + spws[0][1] - spws[0][0], + ) ds.set_spw(spws[1]) - key = (0,bls1[0],"xx") + key = (0, bls1[0], "xx") known_cov = {} - model = 'known' - Ckey = ((0, 0), ((bls1[0][0],bls1[0][1] ,"xx"),(bls1[0][0],bls1[0][1] ,"xx")), 'known', 0, False, True,) + Ckey = ( + (0, 0), + ((bls1[0][0], bls1[0][1], "xx"), (bls1[0][0], bls1[0][1], "xx")), + "known", + 0, + False, + True, + ) known_cov[Ckey] = np.diag(np.ones(uvd.Nfreqs)) - ds.C_model(key, model='known', time_index=0, known_cov=known_cov) - assert ( Ckey in ds._C.keys()) - assert ds._C[Ckey].shape == (spws[1][1]-spws[1][0], spws[1][1]-spws[1][0]) + ds.C_model(key, model="known", time_index=0, known_cov=known_cov) + assert Ckey in ds._C.keys() + assert ds._C[Ckey].shape == (spws[1][1] - spws[1][0], spws[1][1] - spws[1][0]) def test_get_analytic_covariance(self): uvd = UVData() - uvd.read(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) + uvd.read(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) uvd.nsample_array[:] = 1.0 uvd.flag_array[:] = False cosmo = conversions.Cosmo_Conversions() - uvb = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits'), cosmo=cosmo) + uvb = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits"), cosmo=cosmo + ) # extend time axis by factor of 4 - for i in range(2): + for _ in range(2): new = copy.deepcopy(uvd) new.time_array += new.Ntimes * np.diff(np.unique(new.time_array))[0] - new.lst_array = uvutils.get_lst_for_time(new.time_array, *new.telescope_location_lat_lon_alt_degrees) + new.lst_array = uvutils.get_lst_for_time( + new.time_array, *new.telescope_location_lat_lon_alt_degrees + ) uvd += new # get redundant baselines @@ -1345,21 +1526,41 @@ def test_get_analytic_covariance(self): # append roughly 20 blpairs to a list bls1, bls2 = [], [] for red in reds[:3]: - _bls1, _bls2, _ = utils.construct_blpairs(red, exclude_auto_bls=False, exclude_cross_bls=False, exclude_permutations=False) + _bls1, _bls2, _ = utils.construct_blpairs( + red, + exclude_auto_bls=False, + exclude_cross_bls=False, + exclude_permutations=False, + ) bls1.extend(_bls1) bls2.extend(_bls2) # keep only 20 blpairs for speed (each with 40 independent time samples) bls1, bls2 = bls1[:20], bls2[:20] - Nblpairs = len(bls1) - + # generate a sky and noise simulation: each bl has the same FG signal, constant in time # but has a different noise realization np.random.seed(0) - sim1 = testing.sky_noise_sim(uvd, uvb, cov_amp=1000, cov_length_scale=10, constant_per_bl=True, - constant_in_time=True, bl_loop_seed=0, divide_by_nsamp=False) + sim1 = testing.sky_noise_sim( + uvd, + uvb, + cov_amp=1000, + cov_length_scale=10, + constant_per_bl=True, + constant_in_time=True, + bl_loop_seed=0, + divide_by_nsamp=False, + ) np.random.seed(0) - sim2 = testing.sky_noise_sim(uvd, uvb, cov_amp=1000, cov_length_scale=10, constant_per_bl=True, - constant_in_time=True, bl_loop_seed=1, divide_by_nsamp=False) + sim2 = testing.sky_noise_sim( + uvd, + uvb, + cov_amp=1000, + cov_length_scale=10, + constant_per_bl=True, + constant_in_time=True, + bl_loop_seed=1, + divide_by_nsamp=False, + ) # setup ds ds = pspecdata.PSpecData(dsets=[sim1, sim2], wgts=[None, None], beam=uvb) @@ -1370,10 +1571,21 @@ def test_get_analytic_covariance(self): key2 = (1, bls2[0], "xx") ds.set_spw((60, 90)) M_ = np.diag(np.ones(ds.spw_Ndlys)) - for model in ['autos', 'empirical']: - (cov_q_real, cov_q_imag, cov_p_real, - cov_p_imag) = ds.get_analytic_covariance(key1, key2, M=M_, exact_norm=False, pol=False, - model=model, known_cov=None) + for model in ["autos", "empirical"]: + ( + cov_q_real, + cov_q_imag, + cov_p_real, + cov_p_imag, + ) = ds.get_analytic_covariance( + key1, + key2, + M=M_, + exact_norm=False, + pol=False, + model=model, + known_cov=None, + ) # assert these arrays are effectively real-valued, even though they are complex type. # some numerical noise can leak-in, so check to within a dynamic range of peak real power. for cov in [cov_q_real, cov_q_imag, cov_p_real, cov_p_imag]: @@ -1384,69 +1596,186 @@ def test_get_analytic_covariance(self): # path where we use some optimization for diagonal matrices, while 'fiducial' mode will follow the longer path # where there is no such optimization. This test should show the results from two paths are equivalent. known_cov_test = dict() - C_n_11 = np.diag([2.]*ds.Nfreqs) - P_n_11, S_n_11, C_n_12, P_n_12, S_n_12 = np.zeros_like(C_n_11), np.zeros_like(C_n_11), np.zeros_like(C_n_11), np.zeros_like(C_n_11), np.zeros_like(C_n_11) - models = ['dsets','fiducial'] + C_n_11 = np.diag([2.0] * ds.Nfreqs) + P_n_11, S_n_11, C_n_12, P_n_12, S_n_12 = ( + np.zeros_like(C_n_11), + np.zeros_like(C_n_11), + np.zeros_like(C_n_11), + np.zeros_like(C_n_11), + np.zeros_like(C_n_11), + ) + models = ["dsets", "fiducial"] for model in models: for blpair in list(zip(bls1, bls2)): for time_index in range(ds.Ntimes): - key1 = (0,blpair[0],'xx') + key1 = (0, blpair[0], "xx") dset1, bl1 = ds.parse_blkey(key1) - key2 = (1,blpair[1],'xx') + key2 = (1, blpair[1], "xx") dset2, bl2 = ds.parse_blkey(key2) - Ckey = ((dset1, dset1), (bl1,bl1), ) + (model, time_index, False, True,) + Ckey = ((dset1, dset1), (bl1, bl1),) + ( + model, + time_index, + False, + True, + ) known_cov_test[Ckey] = C_n_11 - Ckey = ((dset1, dset1), (bl1,bl1), ) + (model, time_index, False, False,) + Ckey = ((dset1, dset1), (bl1, bl1),) + ( + model, + time_index, + False, + False, + ) known_cov_test[Ckey] = P_n_11 - Ckey = ((dset1, dset1), (bl1,bl1), ) + (model, time_index, True, True,) + Ckey = ((dset1, dset1), (bl1, bl1),) + ( + model, + time_index, + True, + True, + ) known_cov_test[Ckey] = S_n_11 - Ckey = ((dset2, dset2), (bl2,bl2), ) + (model, time_index, False, True,) + Ckey = ((dset2, dset2), (bl2, bl2),) + ( + model, + time_index, + False, + True, + ) known_cov_test[Ckey] = C_n_11 - Ckey = ((dset2, dset2), (bl2,bl2), ) + (model, time_index, False, False,) + Ckey = ((dset2, dset2), (bl2, bl2),) + ( + model, + time_index, + False, + False, + ) known_cov_test[Ckey] = P_n_11 - Ckey = ((dset2, dset2), (bl2,bl2), ) + (model, time_index, True, True,) + Ckey = ((dset2, dset2), (bl2, bl2),) + ( + model, + time_index, + True, + True, + ) known_cov_test[Ckey] = S_n_11 - Ckey = ((dset1, dset2), (bl1,bl2), ) + (model, time_index, False, True,) + Ckey = ((dset1, dset2), (bl1, bl2),) + ( + model, + time_index, + False, + True, + ) known_cov_test[Ckey] = C_n_12 - Ckey = ((dset2, dset1), (bl2,bl1), ) + (model, time_index, False, True,) + Ckey = ((dset2, dset1), (bl2, bl1),) + ( + model, + time_index, + False, + True, + ) known_cov_test[Ckey] = C_n_12 - Ckey = ((dset2, dset1), (bl2,bl1), ) + (model, time_index, False, False,) + Ckey = ((dset2, dset1), (bl2, bl1),) + ( + model, + time_index, + False, + False, + ) known_cov_test[Ckey] = P_n_12 - Ckey = ((dset2, dset1), (bl2,bl1), ) + (model, time_index, True, True,) + Ckey = ((dset2, dset1), (bl2, bl1),) + ( + model, + time_index, + True, + True, + ) known_cov_test[Ckey] = S_n_12 - uvp_dsets_cov = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), spw_ranges=(60, 90), store_cov=True, - cov_model='dsets', known_cov=known_cov_test, verbose=False, taper='bh') - uvp_fiducial_cov = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), spw_ranges=(60, 90), store_cov=True, - cov_model='fiducial', known_cov=known_cov_test, verbose=False, taper='bh') + uvp_dsets_cov = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + spw_ranges=(60, 90), + store_cov=True, + cov_model="dsets", + known_cov=known_cov_test, + verbose=False, + taper="bh", + ) + uvp_fiducial_cov = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + spw_ranges=(60, 90), + store_cov=True, + cov_model="fiducial", + known_cov=known_cov_test, + verbose=False, + taper="bh", + ) # check their cov_array are equal - assert np.allclose(uvp_dsets_cov.cov_array_real[0], uvp_fiducial_cov.cov_array_real[0], rtol=1e-05) + assert np.allclose( + uvp_dsets_cov.cov_array_real[0], + uvp_fiducial_cov.cov_array_real[0], + rtol=1e-05, + ) # check noise floor computation from auto correlations - uvp_auto_cov = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), spw_ranges=(60, 90), store_cov=True, - cov_model='autos', verbose=False, taper='bh') + uvp_auto_cov = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + spw_ranges=(60, 90), + store_cov=True, + cov_model="autos", + verbose=False, + taper="bh", + ) # get RMS of noise-dominated bandpowers for uvp_auto_cov noise_dlys = np.abs(uvp_auto_cov.get_dlys(0) * 1e9) > 1000 rms = [] for key in uvp_auto_cov.get_all_keys(): - rms.append(np.std(uvp_auto_cov.get_data(key).real \ - / np.sqrt(np.diagonal(uvp_auto_cov.get_cov(key).real, axis1=1, axis2=2)), axis=0)) + rms.append( + np.std( + uvp_auto_cov.get_data(key).real + / np.sqrt( + np.diagonal(uvp_auto_cov.get_cov(key).real, axis1=1, axis2=2) + ), + axis=0, + ) + ) rms = np.mean(rms, axis=0) # assert this is close to 1.0 assert np.isclose(np.mean(rms[noise_dlys]), 1.0, atol=0.1) # check signal + noise floor computation - uvp_fgdep_cov = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), spw_ranges=(60, 90), store_cov=True, - cov_model='foreground_dependent', verbose=False, taper='bh') + uvp_fgdep_cov = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + spw_ranges=(60, 90), + store_cov=True, + cov_model="foreground_dependent", + verbose=False, + taper="bh", + ) # get RMS of data: divisor is foreground_dependent covariance this time # b/c noise in empirically estimated fg-dep cov yields biased errorbar (tavg is not unbiased, but less-biased) rms = [] for key in uvp_fgdep_cov.get_all_keys(): - rms.append(np.std(uvp_fgdep_cov.get_data(key)[:,~noise_dlys].real \ - / np.sqrt(np.mean(np.diagonal(uvp_fgdep_cov.get_cov(key).real, axis1=1, axis2=2)[:,~noise_dlys], axis=0)), axis=0)) + rms.append( + np.std( + uvp_fgdep_cov.get_data(key)[:, ~noise_dlys].real + / np.sqrt( + np.mean( + np.diagonal( + uvp_fgdep_cov.get_cov(key).real, axis1=1, axis2=2 + )[:, ~noise_dlys], + axis=0, + ) + ), + axis=0, + ) + ) rms = np.mean(rms, axis=0) # assert this is close to 1.0 assert np.isclose(np.mean(rms), 1.0, atol=0.1) @@ -1454,82 +1783,153 @@ def test_get_analytic_covariance(self): def test_pspec(self): # generate ds uvd = copy.deepcopy(self.uvd) - ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm, labels=['red', 'blue']) + ds = pspecdata.PSpecData( + dsets=[uvd, uvd], wgts=[None, None], beam=self.bm, labels=["red", "blue"] + ) # check basic execution with baseline list bls = [(24, 25), (37, 38), (38, 39), (52, 53)] - uvp = ds.pspec(bls, bls, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=False) + uvp = ds.pspec( + bls, + bls, + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=False, + ) assert len(uvp.bl_array) == len(bls) - assert (uvp.antnums_to_blpair(((24, 25), (24, 25))) in uvp.blpair_array) + assert uvp.antnums_to_blpair(((24, 25), (24, 25))) in uvp.blpair_array assert uvp.data_array[0].dtype == np.complex128 assert uvp.data_array[0].shape == (240, 64, 1) assert not uvp.exact_windows - #test for different forms of input parameters - ds.pspec(bls, bls, (0, 1), ('xx','xx'), spw_ranges=(10,20)) - ds.pspec(bls, bls, (0, 1), ('xx','xx'), n_dlys=10, spw_ranges=[(10,20)]) - ds.pspec(bls, bls, (0, 1), ('xx','xx'), n_dlys=1) + # test for different forms of input parameters + ds.pspec(bls, bls, (0, 1), ("xx", "xx"), spw_ranges=(10, 20)) + ds.pspec(bls, bls, (0, 1), ("xx", "xx"), n_dlys=10, spw_ranges=[(10, 20)]) + ds.pspec(bls, bls, (0, 1), ("xx", "xx"), n_dlys=1) my_r_params = {} my_r_params_dset0_only = {} - rp = {'filter_centers':[0.], - 'filter_half_widths':[250e-9], - 'filter_factors':[1e-9]} + rp = { + "filter_centers": [0.0], + "filter_half_widths": [250e-9], + "filter_factors": [1e-9], + } for bl in bls: - key1 = (0,) + bl + ('xx',) - key2 = (1,) + bl + ('xx',) + key1 = (0,) + bl + ("xx",) + key2 = (1,) + bl + ("xx",) my_r_params[key1] = rp my_r_params_dset0_only[key1] = rp my_r_params[key2] = rp - #test inverse sinc weighting. - ds.pspec(bls,bls,(0, 1), ('xx','xx'), - spw_ranges = (10,20), input_data_weight = 'dayenu', - r_params = my_r_params) - - #test value error - pytest.raises(ValueError, ds.pspec, bls, bls, (0, 1), ('xx','xx'), - spw_ranges = (10,20), input_data_weight = 'dayenu', r_params = {}) - - #test value error no dset1 keys - pytest.raises(ValueError, ds.pspec, bls, bls, (0, 1), ('xx','xx'), - spw_ranges = (10,20), input_data_weight = 'dayenu', - r_params = my_r_params_dset0_only) - - #assert error if baselines are not provided in the right format - pytest.raises(NotImplementedError, ds.pspec, [[(24,25),(38,39)]],[[(24,25),(38,39)]], - (0,1),[('xx','xx')]) + # test inverse sinc weighting. + ds.pspec( + bls, + bls, + (0, 1), + ("xx", "xx"), + spw_ranges=(10, 20), + input_data_weight="dayenu", + r_params=my_r_params, + ) + + # test value error + pytest.raises( + ValueError, + ds.pspec, + bls, + bls, + (0, 1), + ("xx", "xx"), + spw_ranges=(10, 20), + input_data_weight="dayenu", + r_params={}, + ) + + # test value error no dset1 keys + pytest.raises( + ValueError, + ds.pspec, + bls, + bls, + (0, 1), + ("xx", "xx"), + spw_ranges=(10, 20), + input_data_weight="dayenu", + r_params=my_r_params_dset0_only, + ) + + # assert error if baselines are not provided in the right format + pytest.raises( + NotImplementedError, + ds.pspec, + [[(24, 25), (38, 39)]], + [[(24, 25), (38, 39)]], + (0, 1), + [("xx", "xx")], + ) # compare the output of get_Q function with analytical estimates - ds_Q = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None],beam=self.bm_Q) - bls_Q = [(24, 25)] - uvp = ds_Q.pspec(bls_Q, bls_Q, (0, 1), [('xx', 'xx')], input_data_weight='identity', - norm='I', taper='none', verbose=True, exact_norm=False) - Q_sample = ds_Q.get_integral_beam('xx') #Get integral beam for pol 'xx' - - assert np.shape(Q_sample) == (ds_Q.spw_range[1] - ds_Q.spw_range[0],\ - ds_Q.spw_range[1] - ds_Q.spw_range[0]) #Check for the right shape - - estimated_Q = (1.0/(4*np.pi)) * np.ones_like(Q_sample) + ds_Q = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm_Q) + bls_Q = [(24, 25)] + uvp = ds_Q.pspec( + bls_Q, + bls_Q, + (0, 1), + [("xx", "xx")], + input_data_weight="identity", + norm="I", + taper="none", + verbose=True, + exact_norm=False, + ) + Q_sample = ds_Q.get_integral_beam("xx") # Get integral beam for pol 'xx' + + assert np.shape(Q_sample) == ( + ds_Q.spw_range[1] - ds_Q.spw_range[0], + ds_Q.spw_range[1] - ds_Q.spw_range[0], + ) # Check for the right shape + + estimated_Q = (1.0 / (4 * np.pi)) * np.ones_like(Q_sample) assert np.allclose(np.real(estimated_Q), np.real(Q_sample), rtol=1e-05) - #Test if the two pipelines match - - ds_t = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm_Q) - uvp_new = ds_t.pspec(bls_Q, bls_Q, (0, 1), [('xx', 'xx')], input_data_weight='identity', - norm='I', taper='none', verbose=True, exact_norm=True) - uvp_ext = ds_t.pspec(bls_Q, bls_Q, (0, 1), [('xx', 'xx')], input_data_weight='identity', - norm='I', taper='none', verbose=True, exact_norm=False) - spw = 0 - blp = (bls_Q[0], bls_Q[0]) - key = (spw, blp, 'xx') - power_real_new = (np.real(uvp_new.get_data(key))) - power_real_ext = (np.real(uvp_ext.get_data(key))) - - diff = np.median((power_real_new-power_real_ext)/power_real_ext) + # Test if the two pipelines match + + ds_t = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm_Q) + uvp_new = ds_t.pspec( + bls_Q, + bls_Q, + (0, 1), + [("xx", "xx")], + input_data_weight="identity", + norm="I", + taper="none", + verbose=True, + exact_norm=True, + ) + uvp_ext = ds_t.pspec( + bls_Q, + bls_Q, + (0, 1), + [("xx", "xx")], + input_data_weight="identity", + norm="I", + taper="none", + verbose=True, + exact_norm=False, + ) + spw = 0 + blp = (bls_Q[0], bls_Q[0]) + key = (spw, blp, "xx") + power_real_new = np.real(uvp_new.get_data(key)) + power_real_ext = np.real(uvp_ext.get_data(key)) + + diff = np.median((power_real_new - power_real_ext) / power_real_ext) assert diff <= 0.05 # check with redundant baseline group list @@ -1537,12 +1937,30 @@ def test_pspec(self): antpos = dict(zip(ants, antpos)) red_bls = [sorted(blg) for blg in redcal.get_pos_reds(antpos)][2] bls1, bls2, blps = utils.construct_blpairs(red_bls, exclude_permutations=True) - uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=False) + uvp = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=False, + ) assert uvp.antnums_to_blpair(((24, 25), (37, 38))) in uvp.blpair_array assert uvp.Nblpairs == 10 - uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=False) + uvp = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=False, + ) assert uvp.antnums_to_blpair(((24, 25), (52, 53))) in uvp.blpair_array assert uvp.antnums_to_blpair(((52, 53), (24, 25))) not in uvp.blpair_array assert uvp.Nblpairs == 10 @@ -1550,17 +1968,37 @@ def test_pspec(self): # test mixed bl group and non blgroup, currently bl grouping of more than 1 blpair doesn't work bls1 = [[(24, 25)], (52, 53)] bls2 = [[(24, 25)], (52, 53)] - uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=False) + uvp = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=False, + ) # test select red_bls = [(24, 25), (37, 38), (38, 39), (52, 53)] - bls1, bls2, blp = utils.construct_blpairs(red_bls, exclude_permutations=False, exclude_auto_bls=False) + bls1, bls2, blp = utils.construct_blpairs( + red_bls, exclude_permutations=False, exclude_auto_bls=False + ) uvd = copy.deepcopy(self.uvd) ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm) - uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), spw_ranges=[(20,30), (30,40)], verbose=False) + uvp = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + spw_ranges=[(20, 30), (30, 40)], + verbose=False, + ) assert uvp.Nblpairs == 16 assert uvp.Nspws == 2 - uvp2 = uvp.select(spws=0, bls=[(24, 25)], only_pairs_in_bls=False, inplace=False) + uvp2 = uvp.select( + spws=0, bls=[(24, 25)], only_pairs_in_bls=False, inplace=False + ) assert uvp2.Nspws == 1 assert uvp2.Nblpairs == 7 uvp.select(spws=0, bls=(24, 25), only_pairs_in_bls=True, inplace=True) @@ -1570,11 +2008,18 @@ def test_pspec(self): # check w/ multiple spectral ranges uvd = copy.deepcopy(self.uvd) ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm) - uvp = ds.pspec(bls, bls, (0, 1), ('xx','xx'), spw_ranges=[(10, 24), (30, 40), (45, 64)], verbose=False) + uvp = ds.pspec( + bls, + bls, + (0, 1), + ("xx", "xx"), + spw_ranges=[(10, 24), (30, 40), (45, 64)], + verbose=False, + ) assert uvp.Nspws == 3 assert uvp.Nspwdlys == 43 assert uvp.data_array[0].shape == (240, 14, 1) - assert uvp.get_data((0, 124125124125, ('xx','xx'))).shape == (60, 14) + assert uvp.get_data((0, 124125124125, ("xx", "xx"))).shape == (60, 14) uvp.select(spws=[1]) assert uvp.Nspws == 1 @@ -1584,116 +2029,233 @@ def test_pspec(self): # test polarization pairs uvd = copy.deepcopy(self.uvd) ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm) - uvp = ds.pspec(bls, bls, (0, 1), ('xx','xx'), spw_ranges=[(10, 24)], verbose=False) - #pytest.raises(NotImplementedError, ds.pspec, bls, bls, (0, 1), pols=[('xx','yy')]) + uvp = ds.pspec( + bls, bls, (0, 1), ("xx", "xx"), spw_ranges=[(10, 24)], verbose=False + ) + # pytest.raises(NotImplementedError, ds.pspec, bls, bls, (0, 1), pols=[('xx','yy')]) uvd = copy.deepcopy(self.uvd) ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm) - uvp = ds.pspec(bls, bls, (0, 1), [('xx','xx'), ('yy','yy')], spw_ranges=[(10, 24)], verbose=False) + uvp = ds.pspec( + bls, + bls, + (0, 1), + [("xx", "xx"), ("yy", "yy")], + spw_ranges=[(10, 24)], + verbose=False, + ) uvd = copy.deepcopy(self.uvd) ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm) uvp = ds.pspec(bls, bls, (0, 1), (-5, -5), spw_ranges=[(10, 24)], verbose=False) # test exceptions - pytest.raises(AssertionError, ds.pspec, bls1[:1], bls2, (0, 1), ('xx','xx')) - pytest.raises(ValueError, ds.pspec, bls, bls, (0, 1), pols=('yy','yy')) + pytest.raises(AssertionError, ds.pspec, bls1[:1], bls2, (0, 1), ("xx", "xx")) + pytest.raises(ValueError, ds.pspec, bls, bls, (0, 1), pols=("yy", "yy")) uvd1 = copy.deepcopy(self.uvd) uvd1.polarization_array = np.array([-6]) ds = pspecdata.PSpecData(dsets=[uvd, uvd1], wgts=[None, None], beam=self.bm) - pytest.raises(ValueError, ds.pspec, bls, bls, (0, 1), ('xx','xx')) + pytest.raises(ValueError, ds.pspec, bls, bls, (0, 1), ("xx", "xx")) # test files with more than one polarizations uvd1 = copy.deepcopy(self.uvd) uvd1.polarization_array = np.array([-6]) uvd2 = self.uvd + uvd1 ds = pspecdata.PSpecData(dsets=[uvd2, uvd2], wgts=[None, None], beam=self.bm) - uvp = ds.pspec(bls, bls, (0, 1), [('xx','xx'), ('yy','yy')], spw_ranges=[(10, 24)], verbose=False) + uvp = ds.pspec( + bls, + bls, + (0, 1), + [("xx", "xx"), ("yy", "yy")], + spw_ranges=[(10, 24)], + verbose=False, + ) uvd1 = copy.deepcopy(self.uvd) uvd1.polarization_array = np.array([-6]) uvd2 = self.uvd + uvd1 ds = pspecdata.PSpecData(dsets=[uvd2, uvd2], wgts=[None, None], beam=self.bm) - uvp = ds.pspec(bls, bls, (0, 1), [('xx','xx'), ('xy','xy')], spw_ranges=[(10, 24)], verbose=False) + uvp = ds.pspec( + bls, + bls, + (0, 1), + [("xx", "xx"), ("xy", "xy")], + spw_ranges=[(10, 24)], + verbose=False, + ) # test with nsamp set to zero uvd = copy.deepcopy(self.uvd) uvd.nsample_array[uvd.antpair2ind(24, 25, ordered=False)] = 0.0 ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm) - uvp = ds.pspec([(24, 25)], [(37, 38)], (0, 1), [('xx', 'xx')]) + uvp = ds.pspec([(24, 25)], [(37, 38)], (0, 1), [("xx", "xx")]) assert np.all(np.isclose(uvp.integration_array[0], 0.0)) # test covariance calculation runs with small number of delays uvd = copy.deepcopy(self.uvd) uvd_std = copy.deepcopy(self.uvd_std) - ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], - dsets_std=[uvd_std, uvd_std], beam=self.bm) + ds = pspecdata.PSpecData( + dsets=[uvd, uvd], + wgts=[None, None], + dsets_std=[uvd_std, uvd_std], + beam=self.bm, + ) # test covariance methods with non-zero filter_extension - uvp = ds.pspec(bls1[:1], bls2[:1], (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=True, spw_ranges=[(10,20)], filter_extensions=[(2,2)], symmetric_taper=False, store_cov=True, cov_model='empirical') - assert hasattr(uvp, 'cov_array_real') - key = (0, (bls1[0],bls2[0]), "xx") + uvp = ds.pspec( + bls1[:1], + bls2[:1], + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=True, + spw_ranges=[(10, 20)], + filter_extensions=[(2, 2)], + symmetric_taper=False, + store_cov=True, + cov_model="empirical", + ) + assert hasattr(uvp, "cov_array_real") + key = (0, (bls1[0], bls2[0]), "xx") # also check the output covariance is uniform along time axis when cov_model='empirical' assert np.allclose(uvp.get_cov(key)[0], uvp.get_cov(key)[-1]) - uvp = ds.pspec(bls1[:1], bls2[:1], (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=True, spw_ranges=[(10,20)], exact_norm=True, store_cov=True, cov_model='dsets') - assert hasattr(uvp, 'cov_array_real') + uvp = ds.pspec( + bls1[:1], + bls2[:1], + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=True, + spw_ranges=[(10, 20)], + exact_norm=True, + store_cov=True, + cov_model="dsets", + ) + assert hasattr(uvp, "cov_array_real") # test the results of stats_array[cov_model] - uvp_cov = ds.pspec(bls1[:1], bls2[:1], (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=True, spw_ranges=[(10,20)], exact_norm=True, store_cov=True, cov_model='foreground_dependent') - uvp_cov_diag = ds.pspec(bls1[:1], bls2[:1], (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=True, spw_ranges=[(10,20)], exact_norm=True, store_cov_diag=True, cov_model='foreground_dependent') - - key = (0, (bls1[0],bls2[0]), "xx") - assert np.isclose(np.diagonal(uvp_cov.get_cov(key), axis1=1, axis2=2), (np.real(uvp_cov_diag.get_stats('foreground_dependent_diag', key)))**2).all() + uvp_cov = ds.pspec( + bls1[:1], + bls2[:1], + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=True, + spw_ranges=[(10, 20)], + exact_norm=True, + store_cov=True, + cov_model="foreground_dependent", + ) + uvp_cov_diag = ds.pspec( + bls1[:1], + bls2[:1], + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=True, + spw_ranges=[(10, 20)], + exact_norm=True, + store_cov_diag=True, + cov_model="foreground_dependent", + ) + + key = (0, (bls1[0], bls2[0]), "xx") + assert np.isclose( + np.diagonal(uvp_cov.get_cov(key), axis1=1, axis2=2), + (np.real(uvp_cov_diag.get_stats("foreground_dependent_diag", key))) ** 2, + ).all() # test identity_Y caching works - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(self.uvd), copy.deepcopy(self.uvd)], wgts=[None, None], - beam=self.bm) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(self.uvd), copy.deepcopy(self.uvd)], + wgts=[None, None], + beam=self.bm, + ) # assert caching is used when appropriate - uvp = ds.pspec([(24, 25), (24, 25)], [(24, 25), (24, 25)], (0, 1), ('xx', 'xx'), - input_data_weight='identity', norm='I', taper='none', verbose=False, - spw_ranges=[(20, 30)]) + uvp = ds.pspec( + [(24, 25), (24, 25)], + [(24, 25), (24, 25)], + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + verbose=False, + spw_ranges=[(20, 30)], + ) assert len(ds._identity_Y) == len(ds._identity_G) == len(ds._identity_H) assert len(ds._identity_Y) == 1 - assert list(ds._identity_Y.keys())[0] == ((0, 24, 25, 'xx'), (1, 24, 25, 'xx')) + assert list(ds._identity_Y.keys())[0] == ((0, 24, 25, "xx"), (1, 24, 25, "xx")) # assert caching is not used when inappropriate - ds.dsets[0].flag_array[ds.dsets[0].antpair2ind(37, 38, ordered=False), :, 25, :] = True - uvp = ds.pspec([(24, 25), (37, 38)], [(24, 25), (37, 38)], (0, 1), ('xx', 'xx'), - input_data_weight='identity', norm='I', taper='none', verbose=False, - spw_ranges=[(20, 30)]) + ds.dsets[0].flag_array[ + ds.dsets[0].antpair2ind(37, 38, ordered=False), :, 25, : + ] = True + uvp = ds.pspec( + [(24, 25), (37, 38)], + [(24, 25), (37, 38)], + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + verbose=False, + spw_ranges=[(20, 30)], + ) assert len(ds._identity_Y) == len(ds._identity_G) == len(ds._identity_H) assert len(ds._identity_Y) == 2 - assert ((0, 24, 25, 'xx'), (1, 24, 25, 'xx')) in ds._identity_Y.keys() - assert ((0, 37, 38, 'xx'), (1, 37, 38, 'xx')) in ds._identity_Y.keys() + assert ((0, 24, 25, "xx"), (1, 24, 25, "xx")) in ds._identity_Y.keys() + assert ((0, 37, 38, "xx"), (1, 37, 38, "xx")) in ds._identity_Y.keys() # test for exact windows - print('window functions') - basename = 'FT_beam_HERA_dipole_test' + print("window functions") + basename = "FT_beam_HERA_dipole_test" # Instantiate UVWindow() # obtain uvp object - datafile = os.path.join(DATA_PATH, 'zen.2458116.31939.HH.uvh5') + datafile = os.path.join(DATA_PATH, "zen.2458116.31939.HH.uvh5") # read datafile uvd = UVData() uvd.read_uvh5(datafile) # Create a new PSpecData objec ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None]) # choose baselines - baselines1, baselines2, blpairs = utils.construct_blpairs(uvd.get_antpairs()[1:], - exclude_permutations=False, - exclude_auto_bls=True) - uvp_w = ds.pspec(baselines1, baselines2, (0, 1), ('xx','xx'), spw_ranges=(175,195), - exact_windows=True, ftbeam_file=os.path.join(DATA_PATH, basename)) + baselines1, baselines2, blpairs = utils.construct_blpairs( + uvd.get_antpairs()[1:], exclude_permutations=False, exclude_auto_bls=True + ) + uvp_w = ds.pspec( + baselines1, + baselines2, + (0, 1), + ("xx", "xx"), + spw_ranges=(175, 195), + exact_windows=True, + ftbeam_file=os.path.join(DATA_PATH, basename), + ) assert uvp_w.exact_windows def test_normalization(self): # Test Normalization of pspec() compared to PAPER legacy techniques - d1 = self.uvd.select(times=np.unique(self.uvd.time_array)[:-1:2], - frequencies=np.unique(self.uvd.freq_array)[40:51], inplace=False) - d2 = self.uvd.select(times=np.unique(self.uvd.time_array)[1::2], - frequencies=np.unique(self.uvd.freq_array)[40:51], inplace=False) + d1 = self.uvd.select( + times=np.unique(self.uvd.time_array)[:-1:2], + frequencies=np.unique(self.uvd.freq_array)[40:51], + inplace=False, + ) + d2 = self.uvd.select( + times=np.unique(self.uvd.time_array)[1::2], + frequencies=np.unique(self.uvd.freq_array)[40:51], + inplace=False, + ) freqs = np.unique(d1.freq_array) # Setup baselines @@ -1705,91 +2267,167 @@ def test_normalization(self): cosmo = conversions.Cosmo_Conversions() # Set to mK scale - d1.data_array *= beam.Jy_to_mK(freqs, pol='XX')[None, None, :, None] - d2.data_array *= beam.Jy_to_mK(freqs, pol='XX')[None, None, :, None] + d1.data_array *= beam.Jy_to_mK(freqs, pol="XX")[None, None, :, None] + d2.data_array *= beam.Jy_to_mK(freqs, pol="XX")[None, None, :, None] # Compare using no taper - OmegaP = beam.power_beam_int(pol='XX') - OmegaPP = beam.power_beam_sq_int(pol='XX') - OmegaP = interp1d(beam.beam_freqs/1e6, OmegaP)(freqs/1e6) - OmegaPP = interp1d(beam.beam_freqs/1e6, OmegaPP)(freqs/1e6) + OmegaP = beam.power_beam_int(pol="XX") + OmegaPP = beam.power_beam_sq_int(pol="XX") + OmegaP = interp1d(beam.beam_freqs / 1e6, OmegaP)(freqs / 1e6) + OmegaPP = interp1d(beam.beam_freqs / 1e6, OmegaPP)(freqs / 1e6) NEB = 1.0 Bp = np.median(np.diff(freqs)) * len(freqs) - scalar = cosmo.X2Y(np.mean(cosmo.f2z(freqs))) * np.mean(OmegaP**2/OmegaPP) * Bp * NEB + scalar = ( + cosmo.X2Y(np.mean(cosmo.f2z(freqs))) + * np.mean(OmegaP**2 / OmegaPP) + * Bp + * NEB + ) data1 = d1.get_data(bls1[0]) data2 = d2.get_data(bls2[0]) - legacy = np.fft.fftshift(np.conj(np.fft.fft(data1, axis=1)) * np.fft.fft(data2, axis=1) * scalar / len(freqs)**2, axes=1)[0] + legacy = np.fft.fftshift( + np.conj(np.fft.fft(data1, axis=1)) + * np.fft.fft(data2, axis=1) + * scalar + / len(freqs) ** 2, + axes=1, + )[0] # hera_pspec OQE ds = pspecdata.PSpecData(dsets=[d1, d2], wgts=[None, None], beam=beam) - uvp = ds.pspec(bls1, bls2, (0, 1), pols=('xx','xx'), taper='none', input_data_weight='identity', norm='I', sampling=True) - oqe = uvp.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] + uvp = ds.pspec( + bls1, + bls2, + (0, 1), + pols=("xx", "xx"), + taper="none", + input_data_weight="identity", + norm="I", + sampling=True, + ) + oqe = uvp.get_data((0, ((24, 25), (37, 38)), ("xx", "xx")))[0] # assert answers are same to within 3% - assert np.isclose(np.real(oqe)/np.real(legacy), 1, atol=0.03, rtol=0.03).all() + assert np.isclose(np.real(oqe) / np.real(legacy), 1, atol=0.03, rtol=0.03).all() # taper window = windows.blackmanharris(len(freqs)) NEB = Bp / trapz(window**2, x=freqs) - scalar = cosmo.X2Y(np.mean(cosmo.f2z(freqs))) * np.mean(OmegaP**2/OmegaPP) * Bp * NEB + scalar = ( + cosmo.X2Y(np.mean(cosmo.f2z(freqs))) + * np.mean(OmegaP**2 / OmegaPP) + * Bp + * NEB + ) data1 = d1.get_data(bls1[0]) data2 = d2.get_data(bls2[0]) - legacy = np.fft.fftshift(np.conj(np.fft.fft(data1*window[None, :], axis=1)) * np.fft.fft(data2*window[None, :], axis=1) * scalar / len(freqs)**2, axes=1)[0] + legacy = np.fft.fftshift( + np.conj(np.fft.fft(data1 * window[None, :], axis=1)) + * np.fft.fft(data2 * window[None, :], axis=1) + * scalar + / len(freqs) ** 2, + axes=1, + )[0] # hera_pspec OQE ds = pspecdata.PSpecData(dsets=[d1, d2], wgts=[None, None], beam=beam) - uvp = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), taper='blackman-harris', input_data_weight='identity', norm='I') - oqe = uvp.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] + uvp = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + taper="blackman-harris", + input_data_weight="identity", + norm="I", + ) + oqe = uvp.get_data((0, ((24, 25), (37, 38)), ("xx", "xx")))[0] # assert answers are same to within 3% - assert np.isclose(np.real(oqe)/np.real(legacy), 1, atol=0.03, rtol=0.03).all() + assert np.isclose(np.real(oqe) / np.real(legacy), 1, atol=0.03, rtol=0.03).all() def test_broadcast_dset_flags(self): # setup fname = os.path.join(DATA_PATH, "zen.all.xx.LST.1.06964.uvA") uvd = UVData() uvd.read_miriad(fname) - Nfreq = uvd.data_array.shape[2] # test basic execution w/ a spw selection - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None]) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None] + ) ds.broadcast_dset_flags(spw_ranges=[(400, 800)], time_thresh=0.2) - assert ds.dsets[0].get_flags(24, 25)[:, 550:650].any() == False + assert not ds.dsets[0].get_flags(24, 25)[:, 550:650].any() # test w/ no spw selection - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None]) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None] + ) ds.broadcast_dset_flags(spw_ranges=None, time_thresh=0.2) assert ds.dsets[0].get_flags(24, 25)[:, 550:650].any() # test unflagging - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None]) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None] + ) ds.broadcast_dset_flags(spw_ranges=None, time_thresh=0.2, unflag=True) - assert ds.dsets[0].get_flags(24, 25)[:, :].any() == False + assert not ds.dsets[0].get_flags(24, 25)[:, :].any() # test single integration being flagged within spw - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None]) - ds.dsets[0].flag_array[ds.dsets[0].antpair2ind(24, 25, ordered=False)[3], 0, 600, 0] = True + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None] + ) + ds.dsets[0].flag_array[ + ds.dsets[0].antpair2ind(24, 25, ordered=False)[3], 0, 600, 0 + ] = True ds.broadcast_dset_flags(spw_ranges=[(400, 800)], time_thresh=0.25, unflag=False) assert ds.dsets[0].get_flags(24, 25)[3, 400:800].all() - assert ds.dsets[0].get_flags(24, 25)[3, :].all() == False + assert not ds.dsets[0].get_flags(24, 25)[3, :].all() # test pspec run sets flagged integration to have zero weight uvd.flag_array[uvd.antpair2ind(24, 25, ordered=False)[3], 0, 400, :] = True - ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None]) + ds = pspecdata.PSpecData( + dsets=[copy.deepcopy(uvd), copy.deepcopy(uvd)], wgts=[None, None] + ) ds.broadcast_dset_flags(spw_ranges=[(400, 450)], time_thresh=0.25) - uvp = ds.pspec([(24, 25), (37, 38), (38, 39)], [(24, 25), (37, 38), (38, 39)], (0, 1), ('xx', 'xx'), - spw_ranges=[(400, 450)], verbose=False) + uvp = ds.pspec( + [(24, 25), (37, 38), (38, 39)], + [(24, 25), (37, 38), (38, 39)], + (0, 1), + ("xx", "xx"), + spw_ranges=[(400, 450)], + verbose=False, + ) # assert flag broadcast above hits weight arrays in uvp - assert np.all(np.isclose(uvp.get_wgts((0, ((24, 25), (24, 25)), ('xx','xx')))[3], 0.0)) + assert np.all( + np.isclose(uvp.get_wgts((0, ((24, 25), (24, 25)), ("xx", "xx")))[3], 0.0) + ) # assert flag broadcast above hits integration arrays - assert np.isclose(uvp.get_integrations((0, ((24, 25), (24, 25)), ('xx','xx')))[3], 0.0) + assert np.isclose( + uvp.get_integrations((0, ((24, 25), (24, 25)), ("xx", "xx")))[3], 0.0 + ) # average spectra - avg_uvp = uvp.average_spectra(blpair_groups=[sorted(np.unique(uvp.blpair_array))], time_avg=True, inplace=False) + avg_uvp = uvp.average_spectra( + blpair_groups=[sorted(np.unique(uvp.blpair_array))], + time_avg=True, + inplace=False, + ) # repeat but change data in flagged portion - ds.dsets[0].data_array[uvd.antpair2ind(24, 25, ordered=False)[3], 0, 400:450, :] *= 100 - uvp2 = ds.pspec([(24, 25), (37, 38), (38, 39)], [(24, 25), (37, 38), (38, 39)], (0, 1), ('xx', 'xx'), - spw_ranges=[(400, 450)], verbose=False) - avg_uvp2 = uvp.average_spectra(blpair_groups=[sorted(np.unique(uvp.blpair_array))], time_avg=True, inplace=False) + ds.dsets[0].data_array[ + uvd.antpair2ind(24, 25, ordered=False)[3], 0, 400:450, : + ] *= 100 + ds.pspec( + [(24, 25), (37, 38), (38, 39)], + [(24, 25), (37, 38), (38, 39)], + (0, 1), + ("xx", "xx"), + spw_ranges=[(400, 450)], + verbose=False, + ) + avg_uvp2 = uvp.average_spectra( + blpair_groups=[sorted(np.unique(uvp.blpair_array))], + time_avg=True, + inplace=False, + ) # assert average before and after are the same! assert avg_uvp == avg_uvp2 @@ -1800,42 +2438,88 @@ def test_RFI_flag_propagation(self): Nfreq = uvd.data_array.shape[2] # Basic test of shape ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm) - test_R = ds.R((1, 37, 38, 'XX')) + test_R = ds.R((1, 37, 38, "XX")) assert test_R.shape == (Nfreq, Nfreq) # First test that turning-off flagging does nothing if there are no flags in the data bls1 = [(24, 25)] bls2 = [(37, 38)] - ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=self.bm, labels=['red', 'blue']) - uvp_flagged = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=False) + ds = pspecdata.PSpecData( + dsets=[uvd, uvd], wgts=[None, None], beam=self.bm, labels=["red", "blue"] + ) + uvp_flagged = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=False, + ) ds.broadcast_dset_flags(unflag=True) - uvp_unflagged = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=False) - - qe_unflagged = uvp_unflagged.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] - qe_flagged = uvp_flagged.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] + uvp_unflagged = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=False, + ) + + qe_unflagged = uvp_unflagged.get_data((0, ((24, 25), (37, 38)), ("xx", "xx")))[ + 0 + ] + qe_flagged = uvp_flagged.get_data((0, ((24, 25), (37, 38)), ("xx", "xx")))[0] # assert answers are same to within 0.1% - assert np.isclose(np.real(qe_unflagged)/np.real(qe_flagged), 1, atol=0.001, rtol=0.001).all() + assert np.isclose( + np.real(qe_unflagged) / np.real(qe_flagged), 1, atol=0.001, rtol=0.001 + ).all() # Test that when flagged, the data within a channel really don't have any effect on the final result uvd2 = copy.deepcopy(uvd) uvd2.flag_array[uvd.antpair2ind(24, 25, ordered=False)] = True ds = pspecdata.PSpecData(dsets=[uvd2, uvd2], wgts=[None, None], beam=self.bm) - uvp_flagged = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=False) + uvp_flagged = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=False, + ) uvd2.data_array[uvd.antpair2ind(24, 25, ordered=False)] *= 9234.913 ds = pspecdata.PSpecData(dsets=[uvd2, uvd2], wgts=[None, None], beam=self.bm) - uvp_flagged_mod = ds.pspec(bls1, bls2, (0, 1), ('xx','xx'), input_data_weight='identity', norm='I', taper='none', - little_h=True, verbose=False) - - qe_flagged_mod = uvp_flagged_mod.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] - qe_flagged = uvp_flagged.get_data((0, ((24, 25), (37, 38)), ('xx','xx')))[0] + uvp_flagged_mod = ds.pspec( + bls1, + bls2, + (0, 1), + ("xx", "xx"), + input_data_weight="identity", + norm="I", + taper="none", + little_h=True, + verbose=False, + ) + + qe_flagged_mod = uvp_flagged_mod.get_data( + (0, ((24, 25), (37, 38)), ("xx", "xx")) + )[0] + qe_flagged = uvp_flagged.get_data((0, ((24, 25), (37, 38)), ("xx", "xx")))[0] # assert answers are same to within 0.1% - assert np.isclose(np.real(qe_flagged_mod), np.real(qe_flagged), atol=0.001, rtol=0.001).all() + assert np.isclose( + np.real(qe_flagged_mod), np.real(qe_flagged), atol=0.001, rtol=0.001 + ).all() # Test below commented out because this sort of aggressive symmetrization is not yet implemented. # # Test that flagging a channel for one dataset (e.g. just left hand dataset x2) @@ -1859,19 +2543,26 @@ def test_RFI_flag_propagation(self): # # assert answers are same to within 3% # assert np.isclose(np.real(qe_flagged_asymm)/np.real(qe_flagged), 1, atol=0.03, rtol=0.03).all() - #print(uvd.data_array.shape) + # print(uvd.data_array.shape) def test_validate_blpairs(self): # test exceptions uvd = copy.deepcopy(self.uvd) - pytest.raises(TypeError, pspecdata.validate_blpairs, [((1, 2), (2, 3))], None, uvd) - pytest.raises(TypeError, pspecdata.validate_blpairs, [((1, 2), (2, 3))], uvd, None) - - bls = [(24,25),(37,38)] - bls1, bls2, blpairs = utils.construct_blpairs(bls, exclude_permutations=False, exclude_auto_bls=True) + pytest.raises( + TypeError, pspecdata.validate_blpairs, [((1, 2), (2, 3))], None, uvd + ) + pytest.raises( + TypeError, pspecdata.validate_blpairs, [((1, 2), (2, 3))], uvd, None + ) + + bls = [(24, 25), (37, 38)] + bls1, bls2, blpairs = utils.construct_blpairs( + bls, exclude_permutations=False, exclude_auto_bls=True + ) pspecdata.validate_blpairs(blpairs, uvd, uvd) - bls1, bls2, blpairs = utils.construct_blpairs(bls, exclude_permutations=False, exclude_auto_bls=True, - group=True) + bls1, bls2, blpairs = utils.construct_blpairs( + bls, exclude_permutations=False, exclude_auto_bls=True, group=True + ) pspecdata.validate_blpairs(blpairs, uvd, uvd) @@ -1881,58 +2572,72 @@ def test_validate_blpairs(self): def test_pspec_run(): - fnames = [os.path.join(DATA_PATH, d) - for d in ['zen.even.xx.LST.1.28828.uvOCRSA', - 'zen.odd.xx.LST.1.28828.uvOCRSA']] + fnames = [ + os.path.join(DATA_PATH, d) + for d in ["zen.even.xx.LST.1.28828.uvOCRSA", "zen.odd.xx.LST.1.28828.uvOCRSA"] + ] beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") - fnames_std = [os.path.join(DATA_PATH,d) - for d in ['zen.even.std.xx.LST.1.28828.uvOCRSA', - 'zen.odd.std.xx.LST.1.28828.uvOCRSA']] + fnames_std = [ + os.path.join(DATA_PATH, d) + for d in [ + "zen.even.std.xx.LST.1.28828.uvOCRSA", + "zen.odd.std.xx.LST.1.28828.uvOCRSA", + ] + ] # test basic execution if os.path.exists("./out.h5"): os.remove("./out.h5") - ds = pspecdata.pspec_run(fnames, "./out.h5", Jy2mK=False, - verbose=False, overwrite=True, dset_pairs=[(0, 1)], - bl_len_range=(14, 15), bl_deg_range=(50, 70), - psname_ext='_0', spw_ranges=[(0, 25)]) - psc = container.PSpecContainer('./out.h5') + ds = pspecdata.pspec_run( + fnames, + "./out.h5", + Jy2mK=False, + verbose=False, + overwrite=True, + dset_pairs=[(0, 1)], + bl_len_range=(14, 15), + bl_deg_range=(50, 70), + psname_ext="_0", + spw_ranges=[(0, 25)], + ) + psc = container.PSpecContainer("./out.h5") assert isinstance(psc, container.PSpecContainer) - assert psc.groups() == ['dset0_dset1'] - assert psc.spectra(psc.groups()[0]) == ['dset0_x_dset1_0'] + assert psc.groups() == ["dset0_dset1"] + assert psc.spectra(psc.groups()[0]) == ["dset0_x_dset1_0"] assert os.path.exists("./out.h5") # test Jy2mK, blpairs, cosmo, cov_array, spw_ranges, dset labeling cosmo = conversions.Cosmo_Conversions(Om_L=0.0) if os.path.exists("./out.h5"): os.remove("./out.h5") - ds = pspecdata.pspec_run(fnames, "./out.h5", - dsets_std=fnames_std, - Jy2mK=True, - beam=beamfile, - blpairs=[((37, 38), (37, 38)), - ((37, 38), (52, 53))], - verbose=False, - overwrite=True, - pol_pairs=[('xx', 'xx'), ('xx', 'xx')], - dset_labels=["foo", "bar"], - dset_pairs=[(0, 0), (0, 1)], - spw_ranges=[(50, 75), (120, 140)], - n_dlys=[20, 20], - cosmo=cosmo, - trim_dset_lsts=False, - broadcast_dset_flags=False, - cov_model='empirical', - store_cov=True) + ds = pspecdata.pspec_run( + fnames, + "./out.h5", + dsets_std=fnames_std, + Jy2mK=True, + beam=beamfile, + blpairs=[((37, 38), (37, 38)), ((37, 38), (52, 53))], + verbose=False, + overwrite=True, + pol_pairs=[("xx", "xx"), ("xx", "xx")], + dset_labels=["foo", "bar"], + dset_pairs=[(0, 0), (0, 1)], + spw_ranges=[(50, 75), (120, 140)], + n_dlys=[20, 20], + cosmo=cosmo, + trim_dset_lsts=False, + broadcast_dset_flags=False, + cov_model="empirical", + store_cov=True, + ) # assert groupname is dset1_dset2 - psc = container.PSpecContainer('./out.h5') - assert ("foo_bar" in psc.groups()) + psc = container.PSpecContainer("./out.h5") + assert "foo_bar" in psc.groups() # assert uvp names are labeled by dset_pairs - assert (sorted(psc.spectra('foo_bar')) \ - == sorted([u'foo_x_bar', u'foo_x_foo'])) + assert sorted(psc.spectra("foo_bar")) == sorted(["foo_x_bar", "foo_x_foo"]) # get UVPSpec for further inspection uvp = psc.get_pspec("foo_bar", "foo_x_bar") @@ -1948,12 +2653,15 @@ def test_pspec_run(): assert uvp.cosmo == cosmo # assert cov_array was calculated b/c std files were passed and store_cov - assert hasattr(uvp, 'cov_array_real') + assert hasattr(uvp, "cov_array_real") # assert dset labeling propagated - assert set(uvp.labels) == set(['bar', 'foo']) + assert set(uvp.labels) == {"bar", "foo"} # assert spw_ranges and n_dlys specification worked - np.testing.assert_array_equal(uvp.get_spw_ranges(), [(163476562.5, 165917968.75, 25, 20), (170312500.0, 172265625.0, 20, 20)]) + np.testing.assert_array_equal( + uvp.get_spw_ranges(), + [(163476562.5, 165917968.75, 25, 20), (170312500.0, 172265625.0, 20, 20)], + ) # test single_dset, time_interleaving, rephasing, flag broadcasting uvd = UVData() @@ -1966,43 +2674,67 @@ def test_pspec_run(): uvd2 = uvd.select(times=np.unique(uvd.time_array)[1::2], inplace=False) if os.path.exists("./out2.h5"): os.remove("./out2.h5") - ds = pspecdata.pspec_run([copy.deepcopy(uvd)], "./out2.h5", - blpairs=[((37, 38), (37, 38)), ((37, 38), (52, 53))], interleave_times=True, - verbose=False, overwrite=True, spw_ranges=[(0, 25)], rephase_to_dset=0, - broadcast_dset_flags=True, time_thresh=0.3) - psc = container.PSpecContainer('./out2.h5') + ds = pspecdata.pspec_run( + [copy.deepcopy(uvd)], + "./out2.h5", + blpairs=[((37, 38), (37, 38)), ((37, 38), (52, 53))], + interleave_times=True, + verbose=False, + overwrite=True, + spw_ranges=[(0, 25)], + rephase_to_dset=0, + broadcast_dset_flags=True, + time_thresh=0.3, + ) + psc = container.PSpecContainer("./out2.h5") assert isinstance(psc, container.PSpecContainer) - assert psc.groups() == ['dset0_dset1'] - assert psc.spectra(psc.groups()[0]) == ['dset0_x_dset1'] + assert psc.groups() == ["dset0_dset1"] + assert psc.spectra(psc.groups()[0]) == ["dset0_x_dset1"] assert os.path.exists("./out2.h5") # assert dsets are properly interleaved - assert np.isclose((np.unique(ds.dsets[0].time_array) - np.unique(ds.dsets[1].time_array))[0], - -np.diff(np.unique(uvd.time_array))[0]) + assert np.isclose( + (np.unique(ds.dsets[0].time_array) - np.unique(ds.dsets[1].time_array))[0], + -np.diff(np.unique(uvd.time_array))[0], + ) # assert first integration flagged across entire spw assert ds.dsets[0].get_flags(37, 38)[0, 0:25].all() # assert first integration flagged *ONLY* across spw - assert ds.dsets[0].get_flags(37, 38)[0, :0].any() + ds.dsets[0].get_flags(37, 38)[0, 25:].any() == False + assert not ( + ds.dsets[0].get_flags(37, 38)[0, :0].any() + or ds.dsets[0].get_flags(37, 38)[0, 25:].any() + ) # assert channel 15 flagged for all ints assert ds.dsets[0].get_flags(37, 38)[:, 15].all() # assert phase errors decreased after re-phasing phserr_before = np.mean(np.abs(np.angle(uvd1.data_array / uvd2.data_array))) - phserr_after = np.mean(np.abs(np.angle(ds.dsets[0].data_array / ds.dsets[1].data_array))) + phserr_after = np.mean( + np.abs(np.angle(ds.dsets[0].data_array / ds.dsets[1].data_array)) + ) assert phserr_after < phserr_before # repeat feeding dsets_std and wgts if os.path.exists("./out2.h5"): os.remove("./out2.h5") - ds = pspecdata.pspec_run([copy.deepcopy(uvd)], "./out2.h5", dsets_std=[copy.deepcopy(uvd)], - blpairs=[((37, 38), (37, 38)), ((37, 38), (52, 53))], interleave_times=True, - verbose=False, overwrite=True, spw_ranges=[(0, 25)], rephase_to_dset=0, - broadcast_dset_flags=True, time_thresh=0.3) + ds = pspecdata.pspec_run( + [copy.deepcopy(uvd)], + "./out2.h5", + dsets_std=[copy.deepcopy(uvd)], + blpairs=[((37, 38), (37, 38)), ((37, 38), (52, 53))], + interleave_times=True, + verbose=False, + overwrite=True, + spw_ranges=[(0, 25)], + rephase_to_dset=0, + broadcast_dset_flags=True, + time_thresh=0.3, + ) # assert ds passes validation - psc = container.PSpecContainer('./out2.h5') + psc = container.PSpecContainer("./out2.h5") assert ds.dsets_std[0] is not None ds.validate_datasets() assert os.path.exists("./out2.h5") @@ -2013,12 +2745,17 @@ def test_pspec_run(): os.remove("./out2.h5") uvd1 = copy.deepcopy(uvd) uvd2 = uvd.select(times=np.unique(uvd.time_array)[2:], inplace=False) - ds = pspecdata.pspec_run([copy.deepcopy(uvd1), copy.deepcopy(uvd2)], "./out2.h5", - blpairs=[((37, 38), (37, 38)), ((37, 38), (52, 53))], - verbose=False, overwrite=True, spw_ranges=[(50, 100)], - trim_dset_lsts=True) + ds = pspecdata.pspec_run( + [copy.deepcopy(uvd1), copy.deepcopy(uvd2)], + "./out2.h5", + blpairs=[((37, 38), (37, 38)), ((37, 38), (52, 53))], + verbose=False, + overwrite=True, + spw_ranges=[(50, 100)], + trim_dset_lsts=True, + ) # assert first uvd1 lst_array got trimmed by 2 integrations - psc = container.PSpecContainer('./out2.h5') + psc = container.PSpecContainer("./out2.h5") assert ds.dsets[0].Ntimes == 8 assert np.isclose(np.unique(ds.dsets[0].lst_array), np.unique(uvd2.lst_array)).all() @@ -2028,10 +2765,16 @@ def test_pspec_run(): # test when no data is loaded in dset if os.path.exists("./out.h5"): os.remove("./out.h5") - ds = pspecdata.pspec_run(fnames, "./out.h5", Jy2mK=False, verbose=False, overwrite=True, - blpairs=[((500, 501), (600, 601))]) # blpairs that don't exist - assert ds == None - assert os.path.exists("./out.h5") == False + ds = pspecdata.pspec_run( + fnames, + "./out.h5", + Jy2mK=False, + verbose=False, + overwrite=True, + blpairs=[((500, 501), (600, 601))], + ) # blpairs that don't exist + assert ds is None + assert not os.path.exists("./out.h5") # same test but with pre-loaded UVDatas uvds = [] @@ -2039,27 +2782,48 @@ def test_pspec_run(): uvd = UVData() uvd.read_miriad(f) uvds.append(uvd) - ds = pspecdata.pspec_run(uvds, "./out.h5", dsets_std=fnames_std, Jy2mK=False, verbose=False, overwrite=True, - blpairs=[((500, 501), (600, 601))]) - assert ds == None - assert os.path.exists("./out.h5") == False + ds = pspecdata.pspec_run( + uvds, + "./out.h5", + dsets_std=fnames_std, + Jy2mK=False, + verbose=False, + overwrite=True, + blpairs=[((500, 501), (600, 601))], + ) + assert ds is None + assert not os.path.exists("./out.h5") # test when data is loaded, but no blpairs match if os.path.exists("./out.h5"): os.remove("./out.h5") - ds = pspecdata.pspec_run(fnames, "./out.h5", Jy2mK=False, verbose=False, overwrite=True, - blpairs=[((37, 38), (600, 601))]) + ds = pspecdata.pspec_run( + fnames, + "./out.h5", + Jy2mK=False, + verbose=False, + overwrite=True, + blpairs=[((37, 38), (600, 601))], + ) assert isinstance(ds, pspecdata.PSpecData) - assert os.path.exists("./out.h5") == False + assert not os.path.exists("./out.h5") # test glob-parseable input dataset - dsets = [os.path.join(DATA_PATH, "zen.2458042.?????.xx.HH.uvXA"), - os.path.join(DATA_PATH, "zen.2458042.?????.xx.HH.uvXA")] + dsets = [ + os.path.join(DATA_PATH, "zen.2458042.?????.xx.HH.uvXA"), + os.path.join(DATA_PATH, "zen.2458042.?????.xx.HH.uvXA"), + ] if os.path.exists("./out.h5"): os.remove("./out.h5") - ds = pspecdata.pspec_run(dsets, "./out.h5", Jy2mK=False, verbose=True, overwrite=True, - blpairs=[((24, 25), (37, 38))]) - psc = container.PSpecContainer('./out.h5', 'rw') + ds = pspecdata.pspec_run( + dsets, + "./out.h5", + Jy2mK=False, + verbose=True, + overwrite=True, + blpairs=[((24, 25), (37, 38))], + ) + psc = container.PSpecContainer("./out.h5", "rw") uvp = psc.get_pspec("dset0_dset1", "dset0_x_dset1") assert uvp.Ntimes == 120 if os.path.exists("./out.h5"): @@ -2068,50 +2832,95 @@ def test_pspec_run(): # test input calibration dfile = os.path.join(DATA_PATH, "zen.2458116.30448.HH.uvh5") cfile = os.path.join(DATA_PATH, "zen.2458116.30448.HH.flagged_abs.calfits") - ds = pspecdata.pspec_run([dfile, dfile], "./out.h5", cals=cfile, dsets_std=[dfile, dfile], - verbose=False, overwrite=True, blpairs=[((23, 24), (24, 25))], - pol_pairs=[('xx', 'xx')], interleave_times=False, - file_type='uvh5', spw_ranges=[(100, 150)], cal_flag=True) - psc = container.PSpecContainer('./out.h5', 'rw') - uvp = psc.get_pspec('dset0_dset1', 'dset0_x_dset1') + ds = pspecdata.pspec_run( + [dfile, dfile], + "./out.h5", + cals=cfile, + dsets_std=[dfile, dfile], + verbose=False, + overwrite=True, + blpairs=[((23, 24), (24, 25))], + pol_pairs=[("xx", "xx")], + interleave_times=False, + file_type="uvh5", + spw_ranges=[(100, 150)], + cal_flag=True, + ) + psc = container.PSpecContainer("./out.h5", "rw") + uvp = psc.get_pspec("dset0_dset1", "dset0_x_dset1") # test calibration flags were propagated to test that cal was applied assert ds.dsets[0].flag_array.any() assert ds.dsets[1].flag_array.any() assert ds.dsets_std[0].flag_array.any() assert ds.dsets_std[1].flag_array.any() - assert ds.dsets[0].extra_keywords['filename'] != '""' - assert ds.dsets[0].extra_keywords['calibration'] != '""' - assert 'cal: /' in uvp.history + assert ds.dsets[0].extra_keywords["filename"] != '""' + assert ds.dsets[0].extra_keywords["calibration"] != '""' + assert "cal: /" in uvp.history # test w/ conjugated blpairs dfile = os.path.join(DATA_PATH, "zen.2458116.30448.HH.uvh5") - ds = pspecdata.pspec_run([dfile, dfile], "./out.h5", cals=cfile, dsets_std=[dfile, dfile], - verbose=False, overwrite=True, blpairs=[((24, 23), (25, 24))], - pol_pairs=[('xx', 'xx')], interleave_times=False, - file_type='uvh5', spw_ranges=[(100, 150)], cal_flag=True) - psc = container.PSpecContainer('./out.h5', 'rw') - uvp = psc.get_pspec('dset0_dset1', 'dset0_x_dset1') + ds = pspecdata.pspec_run( + [dfile, dfile], + "./out.h5", + cals=cfile, + dsets_std=[dfile, dfile], + verbose=False, + overwrite=True, + blpairs=[((24, 23), (25, 24))], + pol_pairs=[("xx", "xx")], + interleave_times=False, + file_type="uvh5", + spw_ranges=[(100, 150)], + cal_flag=True, + ) + psc = container.PSpecContainer("./out.h5", "rw") + uvp = psc.get_pspec("dset0_dset1", "dset0_x_dset1") assert uvp.Nblpairs == 1 # test exceptions - pytest.raises(AssertionError, pspecdata.pspec_run, 'foo', "./out.h5") - pytest.raises(AssertionError, pspecdata.pspec_run, fnames, "./out.h5", blpairs=(1, 2), verbose=False) - pytest.raises(AssertionError, pspecdata.pspec_run, fnames, "./out.h5", blpairs=[1, 2], verbose=False) - pytest.raises(AssertionError, pspecdata.pspec_run, fnames, "./out.h5", beam=1, verbose=False) + pytest.raises(AssertionError, pspecdata.pspec_run, "foo", "./out.h5") + pytest.raises( + AssertionError, + pspecdata.pspec_run, + fnames, + "./out.h5", + blpairs=(1, 2), + verbose=False, + ) + pytest.raises( + AssertionError, + pspecdata.pspec_run, + fnames, + "./out.h5", + blpairs=[1, 2], + verbose=False, + ) + pytest.raises( + AssertionError, pspecdata.pspec_run, fnames, "./out.h5", beam=1, verbose=False + ) # test execution with list of files for each dataset and list of cals if os.path.exists("./out.h5"): os.remove("./out.h5") fnames = glob.glob(os.path.join(DATA_PATH, "zen.2458116.*.HH.uvh5")) cals = glob.glob(os.path.join(DATA_PATH, "zen.2458116.*.HH.flagged_abs.calfits")) - ds = pspecdata.pspec_run([fnames, fnames], "./out.h5", Jy2mK=False, - verbose=False, overwrite=True, file_type='uvh5', - bl_len_range=(14, 15), bl_deg_range=(0, 1), - psname_ext='_0', spw_ranges=[(0, 25)], cals=[cals, cals]) - psc = container.PSpecContainer('./out.h5', 'rw') + ds = pspecdata.pspec_run( + [fnames, fnames], + "./out.h5", + Jy2mK=False, + verbose=False, + overwrite=True, + file_type="uvh5", + bl_len_range=(14, 15), + bl_deg_range=(0, 1), + psname_ext="_0", + spw_ranges=[(0, 25)], + cals=[cals, cals], + ) + psc = container.PSpecContainer("./out.h5", "rw") assert isinstance(psc, container.PSpecContainer) - assert psc.groups() == ['dset0_dset1'] - assert psc.spectra(psc.groups()[0]) == ['dset0_x_dset1_0'] + assert psc.groups() == ["dset0_dset1"] + assert psc.spectra(psc.groups()[0]) == ["dset0_x_dset1_0"] assert os.path.exists("./out.h5") if os.path.exists("./out.h5"): @@ -2119,17 +2928,30 @@ def test_pspec_run(): # test with cov_model that requires autos w/ fname as filepath fnames = glob.glob(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) - pspecdata.pspec_run([fnames], "./out.h5", spw_ranges=[(50, 70)], dset_pairs=[(0, 0)], - verbose=False, overwrite=True, file_type='miriad', pol_pairs=[('xx', 'xx')], - blpairs=[((37, 38), (37, 38))], cov_model='foreground_dependent', store_cov=True) + pspecdata.pspec_run( + [fnames], + "./out.h5", + spw_ranges=[(50, 70)], + dset_pairs=[(0, 0)], + verbose=False, + overwrite=True, + file_type="miriad", + pol_pairs=[("xx", "xx")], + blpairs=[((37, 38), (37, 38))], + cov_model="foreground_dependent", + store_cov=True, + ) psc = container.PSpecContainer("out.h5", keep_open=False) - uvp = psc.get_pspec('dset0', 'dset0_x_dset0') - assert hasattr(uvp, 'cov_array_real') - os.remove('out.h5') + uvp = psc.get_pspec("dset0", "dset0_x_dset0") + assert hasattr(uvp, "cov_array_real") + os.remove("out.h5") + def test_input_calibration(): dfiles = sorted(glob.glob(os.path.join(DATA_PATH, "zen.2458116.30*.HH.uvh5"))) - cfiles = sorted(glob.glob(os.path.join(DATA_PATH, "zen.2458116.30*.HH.flagged_abs.calfits"))) + cfiles = sorted( + glob.glob(os.path.join(DATA_PATH, "zen.2458116.30*.HH.flagged_abs.calfits")) + ) for i, f in enumerate(zip(dfiles, cfiles)): uvd = UVData() uvd.read(f[0]) @@ -2141,28 +2963,49 @@ def test_input_calibration(): # test add pd = pspecdata.PSpecData() pd.add(dfiles, None) # w/o cal - pd.add([copy.deepcopy(uv) for uv in dfiles], None, cals=cfiles, cal_flag=False) # with cal - g = (cfiles[0].get_gains(23, 'x') * np.conj(cfiles[0].get_gains(24, 'x'))).T - np.testing.assert_array_almost_equal(pd.dsets[0].get_data(23, 24, 'xx') / g, - pd.dsets[1].get_data(23, 24, 'xx')) + pd.add( + [copy.deepcopy(uv) for uv in dfiles], None, cals=cfiles, cal_flag=False + ) # with cal + g = (cfiles[0].get_gains(23, "x") * np.conj(cfiles[0].get_gains(24, "x"))).T + np.testing.assert_array_almost_equal( + pd.dsets[0].get_data(23, 24, "xx") / g, pd.dsets[1].get_data(23, 24, "xx") + ) # test add with dictionaries - pd.add({'one': copy.deepcopy(dfiles[0])}, {'one': None}, cals={'one':cfiles[0]}, cal_flag=False) - np.testing.assert_array_almost_equal(pd.dsets[0].get_data(23, 24, 'xx') / g, - pd.dsets[2].get_data(23, 24, 'xx')) + pd.add( + {"one": copy.deepcopy(dfiles[0])}, + {"one": None}, + cals={"one": cfiles[0]}, + cal_flag=False, + ) + np.testing.assert_array_almost_equal( + pd.dsets[0].get_data(23, 24, "xx") / g, pd.dsets[2].get_data(23, 24, "xx") + ) # test dset_std calibration - pd.add([copy.deepcopy(uv) for uv in dfiles], None, dsets_std=[copy.deepcopy(uv) for uv in dfiles], - cals=cfiles, cal_flag=False) - np.testing.assert_array_almost_equal(pd.dsets[0].get_data(23, 24, 'xx') / g, - pd.dsets_std[3].get_data(23, 24, 'xx')) + pd.add( + [copy.deepcopy(uv) for uv in dfiles], + None, + dsets_std=[copy.deepcopy(uv) for uv in dfiles], + cals=cfiles, + cal_flag=False, + ) + np.testing.assert_array_almost_equal( + pd.dsets[0].get_data(23, 24, "xx") / g, pd.dsets_std[3].get_data(23, 24, "xx") + ) # test exceptions pd = pspecdata.PSpecData() - pytest.raises(TypeError, pd.add, {'one': copy.deepcopy(dfiles[0])}, {'one': None}, - cals='foo', cal_flag=False) + pytest.raises( + TypeError, + pd.add, + {"one": copy.deepcopy(dfiles[0])}, + {"one": None}, + cals="foo", + cal_flag=False, + ) pytest.raises(AssertionError, pd.add, dfiles, [None], cals=[None, None]) - pytest.raises(TypeError, pd.add, dfiles, [None], cals=['foo']) + pytest.raises(TypeError, pd.add, dfiles, [None], cals=["foo"]) def test_window_funcs(): @@ -2172,34 +3015,37 @@ def test_window_funcs(): """ # get a PSpecData uvd = UVData() - uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits")) + uvd.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) + beam = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + ) ds = pspecdata.PSpecData(dsets=[copy.deepcopy(uvd)], beam=beam) ds.set_spw((0, 20)) - ds.set_taper('bh') + ds.set_taper("bh") bl = (37, 38) - key = (0, bl, 'xx') + key = (0, bl, "xx") d = uvd.get_data(bl) C = np.cov(d[:, :20].T).real iC = np.linalg.pinv(C) # iterate over various R and M matrices and ensure # normalization and dtype is consistent - for data_weight in ['identity', 'iC']: + for data_weight in ["identity", "iC"]: ds.set_weighting(data_weight) - for norm in ['H^-1', 'I', 'V^-1/2']: + for norm in ["H^-1", "I", "V^-1/2"]: for exact_norm in [True, False]: - if exact_norm and norm != 'I': + if exact_norm and norm != "I": # exact_norm only supported for norm == 'I' continue ds.clear_cache() - if data_weight == 'iC': + if data_weight == "iC": # fill R with iC - ds._R[(0, (37, 38, 'xx'), 'iC', 'bh')] = iC + ds._R[(0, (37, 38, "xx"), "iC", "bh")] = iC # compute G and H - Gv = ds.get_G(key, key, exact_norm=exact_norm, pol='xx') - Hv = ds.get_H(key, key, exact_norm=exact_norm, pol='xx') - Mv, Wv = ds.get_MW(Gv, Hv, mode=norm, exact_norm=exact_norm, - band_covar=C) + Gv = ds.get_G(key, key, exact_norm=exact_norm, pol="xx") + Hv = ds.get_H(key, key, exact_norm=exact_norm, pol="xx") + Mv, Wv = ds.get_MW( + Gv, Hv, mode=norm, exact_norm=exact_norm, band_covar=C + ) # assert row-sum is normalized to 1 assert np.isclose(Wv.sum(axis=1).real, 1).all() # assert this is a real matrix, even though imag is populated @@ -2208,22 +3054,48 @@ def test_window_funcs(): def test_get_argparser(): args = pspecdata.get_pspec_run_argparser() - a = args.parse_args([['foo'], 'bar', '--dset_pairs', '0~0,1~1', '--pol_pairs', 'xx~xx,yy~yy', - '--spw_ranges', '300~400, 600~800', '--blpairs', '24~25~24~25, 37~38~37~38']) - assert a.pol_pairs == [('xx', 'xx'), ('yy', 'yy')] + a = args.parse_args( + [ + ["foo"], + "bar", + "--dset_pairs", + "0~0,1~1", + "--pol_pairs", + "xx~xx,yy~yy", + "--spw_ranges", + "300~400, 600~800", + "--blpairs", + "24~25~24~25, 37~38~37~38", + ] + ) + assert a.pol_pairs == [("xx", "xx"), ("yy", "yy")] assert a.dset_pairs == [(0, 0), (1, 1)] assert a.spw_ranges == [(300, 400), (600, 800)] assert a.blpairs == [((24, 25), (24, 25)), ((37, 38), (37, 38))] + def test_get_argparser_backwards_compatibility(): args = pspecdata.get_pspec_run_argparser() - a = args.parse_args([['foo'], 'bar', '--dset_pairs', '0 0, 1 1', '--pol_pairs', 'xx xx, yy yy', - '--spw_ranges', '300 400, 600 800', '--blpairs', '24 25 24 25, 37 38 37 38']) - assert a.pol_pairs == [('xx', 'xx'), ('yy', 'yy')] + a = args.parse_args( + [ + ["foo"], + "bar", + "--dset_pairs", + "0 0, 1 1", + "--pol_pairs", + "xx xx, yy yy", + "--spw_ranges", + "300 400, 600 800", + "--blpairs", + "24 25 24 25, 37 38 37 38", + ] + ) + assert a.pol_pairs == [("xx", "xx"), ("yy", "yy")] assert a.dset_pairs == [(0, 0), (1, 1)] assert a.spw_ranges == [(300, 400), (600, 800)] assert a.blpairs == [((24, 25), (24, 25)), ((37, 38), (37, 38))] + """ # LEGACY MONTE CARLO TESTS def validate_get_G(self,tolerance=0.2,NDRAWS=100,NCHAN=8): diff --git a/hera_pspec/tests/test_pstokes.py b/hera_pspec/tests/test_pstokes.py index ff146507..5b05e86f 100644 --- a/hera_pspec/tests/test_pstokes.py +++ b/hera_pspec/tests/test_pstokes.py @@ -1,21 +1,19 @@ import unittest import pytest -import os, sys +import os from hera_pspec.data import DATA_PATH from .. import pstokes import pyuvdata import pyuvdata.utils as uvutils -import copy import numpy as np -dset1 = os.path.join(DATA_PATH, 'zen.all.xx.LST.1.06964.uvA') -dset2 = os.path.join(DATA_PATH, 'zen.all.yy.LST.1.06964.uvA') -multipol_dset = os.path.join(DATA_PATH, 'zen.2458116.31193.HH.uvh5') -multipol_dset_cal = os.path.join(DATA_PATH, 'zen.2458116.31193.HH.flagged_abs.calfits') +dset1 = os.path.join(DATA_PATH, "zen.all.xx.LST.1.06964.uvA") +dset2 = os.path.join(DATA_PATH, "zen.all.yy.LST.1.06964.uvA") +multipol_dset = os.path.join(DATA_PATH, "zen.2458116.31193.HH.uvh5") +multipol_dset_cal = os.path.join(DATA_PATH, "zen.2458116.31193.HH.flagged_abs.calfits") class Test_pstokes(unittest.TestCase): - def setUp(self): # Loading pyuvdata objects self.uvd1 = pyuvdata.UVData() @@ -34,30 +32,30 @@ def test_combine_pol(self): uvd2 = self.uvd2 # basic execution on pol strings - out1 = pstokes._combine_pol(uvd1, uvd2, 'XX', 'YY') + out1 = pstokes._combine_pol(uvd1, uvd2, "XX", "YY") # again w/ pol ints out2 = pstokes._combine_pol(uvd1, uvd2, -5, -6) # assert equivalence assert out1 == out2 # check exceptions - pytest.raises(AssertionError, pstokes._combine_pol, dset1, dset2, 'XX', 'YY' ) - pytest.raises(AssertionError, pstokes._combine_pol, uvd1, uvd2, 'XX', 1) + pytest.raises(AssertionError, pstokes._combine_pol, dset1, dset2, "XX", "YY") + pytest.raises(AssertionError, pstokes._combine_pol, uvd1, uvd2, "XX", 1) def test_construct_pstokes(self): uvd1 = self.uvd1 uvd2 = self.uvd2 # test to form I and Q from single polarized UVData objects - uvdI = pstokes.construct_pstokes(dset1=uvd1, dset2=uvd2, pstokes='pI') - uvdQ = pstokes.construct_pstokes(dset1=uvd1, dset2=uvd2, pstokes='pQ') + pstokes.construct_pstokes(dset1=uvd1, dset2=uvd2, pstokes="pI") + pstokes.construct_pstokes(dset1=uvd1, dset2=uvd2, pstokes="pQ") # check exceptions pytest.raises(AssertionError, pstokes.construct_pstokes, uvd1, 1) # check baselines - uvd3 = uvd2.select(ant_str='auto', inplace=False) - pytest.raises(ValueError, pstokes.construct_pstokes, dset1=uvd1, dset2=uvd3 ) + uvd3 = uvd2.select(ant_str="auto", inplace=False) + pytest.raises(ValueError, pstokes.construct_pstokes, dset1=uvd1, dset2=uvd3) # check frequencies uvd3 = uvd2.select(frequencies=np.unique(uvd2.freq_array)[:10], inplace=False) @@ -79,10 +77,16 @@ def test_construct_pstokes(self): uvd3 = uvd1 + uvd2 # test to form I and Q from dual polarized UVData objects - uvdI = pstokes.construct_pstokes(dset1=uvd3, dset2=uvd3, pstokes='pI') + pstokes.construct_pstokes(dset1=uvd3, dset2=uvd3, pstokes="pI") # check except for same polarizations - pytest.raises(AssertionError, pstokes.construct_pstokes, dset1=uvd1, dset2=uvd1, pstokes='pI') + pytest.raises( + AssertionError, + pstokes.construct_pstokes, + dset1=uvd1, + dset2=uvd1, + pstokes="pI", + ) def test_construct_pstokes_multipol(self): """test construct_pstokes on multi-polarization files""" @@ -93,28 +97,31 @@ def test_construct_pstokes_multipol(self): uvutils.uvcalibrate(uvd, uvc) wgts = [(0.5, 0.5), (0.5, -0.5)] - for i, ps in enumerate(['pI', 'pQ']): + for i, ps in enumerate(["pI", "pQ"]): uvp = pstokes.construct_pstokes(dset1=uvd, dset2=uvd, pstokes=ps) # assert polarization array is correct assert uvp.polarization_array == np.array([i + 1]) # assert data are properly summmed - pstokes_vis = uvd.get_data(23, 24, 'xx') * wgts[i][0] + uvd.get_data(23, 24, 'yy') * wgts[i][1] + pstokes_vis = ( + uvd.get_data(23, 24, "xx") * wgts[i][0] + + uvd.get_data(23, 24, "yy") * wgts[i][1] + ) assert np.isclose(pstokes_vis, uvp.get_data(23, 24, ps)).all() def test_filter_dset_on_stokes_pol(self): dsets = [self.uvd1, self.uvd2] - out = pstokes.filter_dset_on_stokes_pol(dsets, 'pI') + out = pstokes.filter_dset_on_stokes_pol(dsets, "pI") assert out[0].polarization_array[0] == -5 assert out[1].polarization_array[0] == -6 - pytest.raises(AssertionError, pstokes.filter_dset_on_stokes_pol, dsets, 'pV') + pytest.raises(AssertionError, pstokes.filter_dset_on_stokes_pol, dsets, "pV") dsets = [self.uvd2, self.uvd1] - out2 = pstokes.filter_dset_on_stokes_pol(dsets, 'pI') + out2 = pstokes.filter_dset_on_stokes_pol(dsets, "pI") assert out == out2 def test_generate_pstokes_argparser(self): # test argparser for noise error bars. ap = pstokes.generate_pstokes_argparser() - args=ap.parse_args(["input.uvh5", "--pstokes", "pI", "pQ", "--clobber"]) + args = ap.parse_args(["input.uvh5", "--pstokes", "pI", "pQ", "--clobber"]) assert args.inputdata == "input.uvh5" assert args.outputdata is None assert args.clobber diff --git a/hera_pspec/tests/test_testing.py b/hera_pspec/tests/test_testing.py index fddaac77..f4a7f677 100644 --- a/hera_pspec/tests/test_testing.py +++ b/hera_pspec/tests/test_testing.py @@ -14,46 +14,102 @@ def test_build_vanilla_uvpspec(): assert isinstance(cosmo, conversions.Cosmo_Conversions) assert uvp.cosmo == cosmo - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits')) + beam = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + ) uvp, cosmo = testing.build_vanilla_uvpspec(beam=beam) beam_OP = beam.get_Omegas(uvp.polpair_array[0])[0] assert beam_OP.tolist() == uvp.OmegaP.tolist() + def test_uvpspec_from_data(): # get data fname = os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA") fname_std = os.path.join(DATA_PATH, "zen.even.std.xx.LST.1.28828.uvOCRSA") uvd = UVData() uvd.read_miriad(fname) - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") beam = pspecbeam.PSpecBeamUV(beamfile) # test basic execution - uvp = testing.uvpspec_from_data(fname, [(37, 38), (38, 39), (52, 53), (53, 54)], beam=beam, spw_ranges=[(50, 100)]) + uvp = testing.uvpspec_from_data( + fname, + [(37, 38), (38, 39), (52, 53), (53, 54)], + beam=beam, + spw_ranges=[(50, 100)], + ) assert uvp.Nfreqs == 50 - assert np.unique(uvp.blpair_array).tolist() == [137138138139, 137138152153, 137138153154, 138139152153, 138139153154, 152153153154] - uvp2 = testing.uvpspec_from_data(uvd, [(37, 38), (38, 39), (52, 53), (53, 54)], beam=beamfile, spw_ranges=[(50, 100)]) - uvp.history = '' - uvp2.history = '' + assert np.unique(uvp.blpair_array).tolist() == [ + 137138138139, + 137138152153, + 137138153154, + 138139152153, + 138139153154, + 152153153154, + ] + uvp2 = testing.uvpspec_from_data( + uvd, + [(37, 38), (38, 39), (52, 53), (53, 54)], + beam=beamfile, + spw_ranges=[(50, 100)], + ) + uvp.history = "" + uvp2.history = "" assert uvp == uvp2 # test multiple bl groups antpos, ants = uvd.get_ENU_antpos(pick_data_ants=True) reds = redcal.get_pos_reds(dict(zip(ants, antpos))) uvp = testing.uvpspec_from_data(fname, reds[:3], beam=beam, spw_ranges=[(50, 100)]) - assert len(set(uvp.bl_array) - set([137138, 137151, 137152, 138139, 138152, 138153, 139153, 139154, - 151152, 151167, 152153, 152167, 152168, 153154, 153168, 153169, - 154169, 167168, 168169])) == 0 + assert ( + len( + set(uvp.bl_array) + - set( + [ + 137138, + 137151, + 137152, + 138139, + 138152, + 138153, + 139153, + 139154, + 151152, + 151167, + 152153, + 152167, + 152168, + 153154, + 153168, + 153169, + 154169, + 167168, + 168169, + ] + ) + ) + == 0 + ) assert uvp.Nblpairs == 51 # test exceptions pytest.raises(AssertionError, testing.uvpspec_from_data, fname, (37, 38)) - pytest.raises(AssertionError, testing.uvpspec_from_data, fname, [([37, 38], [38, 39])]) - pytest.raises(AssertionError, testing.uvpspec_from_data, fname, [[[37, 38], [38, 39]]]) + pytest.raises( + AssertionError, testing.uvpspec_from_data, fname, [([37, 38], [38, 39])] + ) + pytest.raises( + AssertionError, testing.uvpspec_from_data, fname, [[[37, 38], [38, 39]]] + ) # test std - uvp = testing.uvpspec_from_data(fname, [(37, 38), (38, 39), (52, 53), (53, 54)], - data_std=fname_std, beam=beam, spw_ranges=[(20,28)]) + uvp = testing.uvpspec_from_data( + fname, + [(37, 38), (38, 39), (52, 53), (53, 54)], + data_std=fname_std, + beam=beam, + spw_ranges=[(20, 28)], + ) + def test_noise_sim(): uvd = UVData() @@ -69,10 +125,18 @@ def test_noise_sim(): assert uvn.Nfreqs == uvd2.Nfreqs assert uvn.Nbls == uvd2.Nbls assert uvn.Npols == uvd2.Npols - np.testing.assert_almost_equal(np.std(uvn.data_array[:, :, :, 1].real), 0.20655731998619664) - np.testing.assert_almost_equal(np.std(uvn.data_array[:, :, :, 1].imag), 0.20728471891024444) - np.testing.assert_almost_equal(np.std(uvn.data_array[:, :, :, 0].real) / np.std(uvn.data_array[:, :, :, 1].real), - 1/np.sqrt(2), decimal=2) + np.testing.assert_almost_equal( + np.std(uvn.data_array[:, :, :, 1].real), 0.20655731998619664 + ) + np.testing.assert_almost_equal( + np.std(uvn.data_array[:, :, :, 1].imag), 0.20728471891024444 + ) + np.testing.assert_almost_equal( + np.std(uvn.data_array[:, :, :, 0].real) + / np.std(uvn.data_array[:, :, :, 1].real), + 1 / np.sqrt(2), + decimal=2, + ) # test seed and inplace np.random.seed(0) @@ -81,49 +145,55 @@ def test_noise_sim(): assert uvn == uvn2 # Test with a beam! - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') - uvn = testing.noise_sim(copy.deepcopy(uvd), 300.0, beamfile, seed=0, whiten=True, inplace=False) - assert uvn.vis_units == 'Jy' + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + uvn = testing.noise_sim( + copy.deepcopy(uvd), 300.0, beamfile, seed=0, whiten=True, inplace=False + ) + assert uvn.vis_units == "Jy" # test Tsys scaling - uvn3 = testing.noise_sim(uvd2, 2*300.0, seed=0, whiten=True, inplace=False) - np.testing.assert_almost_equal(np.std(uvn3.data_array[:, :, :, 1].real), 2*0.20655731998619664) - np.testing.assert_almost_equal(np.std(uvn3.data_array[:, :, :, 1].imag), 2*0.20728471891024444) + uvn3 = testing.noise_sim(uvd2, 2 * 300.0, seed=0, whiten=True, inplace=False) + np.testing.assert_almost_equal( + np.std(uvn3.data_array[:, :, :, 1].real), 2 * 0.20655731998619664 + ) + np.testing.assert_almost_equal( + np.std(uvn3.data_array[:, :, :, 1].imag), 2 * 0.20728471891024444 + ) # test Nextend uvn = testing.noise_sim(uvd, 300.0, seed=0, whiten=True, inplace=False, Nextend=4) - assert uvn.Ntimes == uvd.Ntimes*5 + assert uvn.Ntimes == uvd.Ntimes * 5 assert uvn.Nfreqs == uvd.Nfreqs assert uvn.Nbls == uvd.Nbls assert uvn.Npols == uvd.Npols def test_sky_noise_jy_autos(): - + # Load data uvd = UVData() uvfile = os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA") uvd.read_miriad(uvfile) - + # Get input arrays lsts = np.unique(uvd.lst_array) freqs = np.unique(uvd.freq_array) int_time = np.median(uvd.integration_time) channel_width = np.mean(np.diff(freqs)) - + # Callable beam function - omega_p = lambda freq: 0.05 * (freq / 100.e6)**-1. - + omega_p = lambda freq: 0.05 * (freq / 100.0e6) ** -1.0 + # Call function n = testing.sky_noise_jy_autos( - lsts, - freqs, - autovis=150., - omega_p=omega_p, - integration_time=int_time, - channel_width=channel_width - ) - + lsts, + freqs, + autovis=150.0, + omega_p=omega_p, + integration_time=int_time, + channel_width=channel_width, + ) + # Check that results are finite assert np.all(~np.isnan(n)) assert np.all(~np.isinf(n)) @@ -138,8 +208,14 @@ def test_sky_noise_sim(): # basic test np.random.seed(0) - sim = testing.sky_noise_sim(uvd, beam, cov_amp=1000, cov_length_scale=10, constant_in_time=True, - divide_by_nsamp=False) + sim = testing.sky_noise_sim( + uvd, + beam, + cov_amp=1000, + cov_length_scale=10, + constant_in_time=True, + divide_by_nsamp=False, + ) # assert something was inserted for bl in sim.get_antpairpols(): if bl[0] != bl[1]: @@ -151,25 +227,46 @@ def test_sky_noise_sim(): uvd2.polarization_array[0] = 1 uvd2b.polarization_array[0] = 2 uvd2 += uvd2b - sim2 = testing.sky_noise_sim(uvd2, beam_ps, cov_amp=1000, cov_length_scale=10, constant_in_time=True, - divide_by_nsamp=False) + sim2 = testing.sky_noise_sim( + uvd2, + beam_ps, + cov_amp=1000, + cov_length_scale=10, + constant_in_time=True, + divide_by_nsamp=False, + ) # assert something was inserted for bl in sim2.get_antpairpols(): if bl[0] != bl[1]: assert np.all(~np.isclose(sim2.get_data(bl), uvd2.get_data(bl))) # try divide by nsamp : set cov_amp=0 so we are only probing noise - sim3 = testing.sky_noise_sim(uvd, beam, cov_amp=0, cov_length_scale=10, constant_in_time=True, - divide_by_nsamp=True) + sim3 = testing.sky_noise_sim( + uvd, + beam, + cov_amp=0, + cov_length_scale=10, + constant_in_time=True, + divide_by_nsamp=True, + ) # assert noise in channel 104 is zero (because nsample = 0) assert np.isclose(sim3.get_data(53, 69)[:, 104], 0).all() # test constant across time and bl sim3 = copy.deepcopy(uvd) - sim3.integration_time[:] = np.inf # set int_time to zero so we are only probing fg signal - sim3 = testing.sky_noise_sim(sim3, beam, cov_amp=1000, cov_length_scale=10, constant_in_time=True, - constant_per_bl=True, divide_by_nsamp=True) + sim3.integration_time[ + : + ] = np.inf # set int_time to zero so we are only probing fg signal + sim3 = testing.sky_noise_sim( + sim3, + beam, + cov_amp=1000, + cov_length_scale=10, + constant_in_time=True, + constant_per_bl=True, + divide_by_nsamp=True, + ) # assert constant in time and across baseline d = sim3.get_data(52, 53) diff --git a/hera_pspec/tests/test_utils.py b/hera_pspec/tests/test_utils.py index 06268373..a703a559 100644 --- a/hera_pspec/tests/test_utils.py +++ b/hera_pspec/tests/test_utils.py @@ -4,7 +4,6 @@ import os, sys, copy from hera_pspec.data import DATA_PATH from .. import utils, testing -from collections import OrderedDict as odict from pyuvdata import UVData from hera_cal import redcal @@ -27,46 +26,45 @@ def test_cov(): assert cov.dtype == np.complex # test exception - pytest.raises(TypeError, utils.cov, d1, w1*1j) - pytest.raises(TypeError, utils.cov, d1, w1, d2=d2, w2=w2*1j) + pytest.raises(TypeError, utils.cov, d1, w1 * 1j) + pytest.raises(TypeError, utils.cov, d1, w1, d2=d2, w2=w2 * 1j) w1 *= -1.0 pytest.raises(ValueError, utils.cov, d1, w1) + def test_load_config(): """ Check YAML config file handling. """ - fname = os.path.join(DATA_PATH, '_test_utils.yaml') + fname = os.path.join(DATA_PATH, "_test_utils.yaml") cfg = utils.load_config(fname) # Check that expected keys exist - assert('data' in cfg.keys()) - assert('pspec' in cfg.keys()) + assert "data" in cfg.keys() + assert "pspec" in cfg.keys() # Check that boolean values are read in correctly - assert(cfg['pspec']['overwrite'] == True) + assert cfg["pspec"]["overwrite"] # Check that lists are read in as lists - assert(len(cfg['data']['subdirs']) == 1) + assert len(cfg["data"]["subdirs"]) == 1 # Check that missing files cause an error pytest.raises(IOError, utils.load_config, "file_that_doesnt_exist") # Check 'None' and list of lists become Nones and list of tuples - assert cfg['data']['pairs'] == [('xx', 'xx'), ('yy', 'yy')] - assert cfg['pspec']['taper'] == 'none' - assert cfg['pspec']['groupname'] == None - assert cfg['pspec']['options']['bar'] == [('foo', 'bar')] - assert cfg['pspec']['options']['foo'] == None + assert cfg["data"]["pairs"] == [("xx", "xx"), ("yy", "yy")] + assert cfg["pspec"]["taper"] == "none" + assert cfg["pspec"]["groupname"] is None + assert cfg["pspec"]["options"]["bar"] == [("foo", "bar")] + assert cfg["pspec"]["options"]["foo"] is None class Test_Utils(unittest.TestCase): - def setUp(self): # Load data into UVData object self.uvd = UVData() - self.uvd.read_miriad(os.path.join(DATA_PATH, - "zen.2458042.17772.xx.HH.uvXA")) + self.uvd.read_miriad(os.path.join(DATA_PATH, "zen.2458042.17772.xx.HH.uvXA")) # Create UVPSpec object self.uvp, cosmo = testing.build_vanilla_uvpspec() @@ -83,36 +81,44 @@ def test_spw_range_from_freqs(self): UVPSpec files. """ # Check that type errors and bounds errors are raised - pytest.raises(AttributeError, utils.spw_range_from_freqs, np.arange(3), - freq_range=(100e6, 110e6)) + pytest.raises( + AttributeError, + utils.spw_range_from_freqs, + np.arange(3), + freq_range=(100e6, 110e6), + ) for obj in [self.uvd, self.uvp]: - pytest.raises(ValueError, utils.spw_range_from_freqs, obj, - freq_range=(98e6, 110e6)) # lower bound - pytest.raises(ValueError, utils.spw_range_from_freqs, obj, - freq_range=(190e6, 202e6)) # upper bound - pytest.raises(ValueError, utils.spw_range_from_freqs, obj, - freq_range=(190e6, 180e6)) # wrong order + pytest.raises( + ValueError, utils.spw_range_from_freqs, obj, freq_range=(98e6, 110e6) + ) # lower bound + pytest.raises( + ValueError, utils.spw_range_from_freqs, obj, freq_range=(190e6, 202e6) + ) # upper bound + pytest.raises( + ValueError, utils.spw_range_from_freqs, obj, freq_range=(190e6, 180e6) + ) # wrong order # Check that valid frequency ranges are returned freq_list = [(100e6, 120e6), (120e6, 140e6), (140e6, 160e6)] spw1 = utils.spw_range_from_freqs(self.uvd, freq_range=(110e6, 130e6)) spw2 = utils.spw_range_from_freqs(self.uvd, freq_range=freq_list) - spw3 = utils.spw_range_from_freqs(self.uvd, freq_range=(98e6, 120e6), - bounds_error=False) + spw3 = utils.spw_range_from_freqs( + self.uvd, freq_range=(98e6, 120e6), bounds_error=False + ) spw4 = utils.spw_range_from_freqs(self.uvd, freq_range=(100e6, 120e6)) # Make sure tuple vs. list arguments were handled correctly - assert( isinstance(spw1, tuple) ) - assert( isinstance(spw2, list) ) - assert( len(spw2) == len(freq_list) ) + assert isinstance(spw1, tuple) + assert isinstance(spw2, list) + assert len(spw2) == len(freq_list) # Make sure that bounds_error=False works - assert( spw3 == spw4 ) + assert spw3 == spw4 # Make sure that this also works for UVPSpec objects spw5 = utils.spw_range_from_freqs(self.uvp, freq_range=(100e6, 104e6)) - assert( isinstance(spw5, tuple) ) - assert( spw5[0] is not None ) + assert isinstance(spw5, tuple) + assert spw5[0] is not None def test_spw_range_from_redshifts(self): """ @@ -120,43 +126,56 @@ def test_spw_range_from_redshifts(self): UVPSpec files (when redshift range is specified). """ # Check that type errors and bounds errors are raised - pytest.raises(AttributeError, utils.spw_range_from_redshifts, - np.arange(3), z_range=(9.7, 12.1)) + pytest.raises( + AttributeError, + utils.spw_range_from_redshifts, + np.arange(3), + z_range=(9.7, 12.1), + ) for obj in [self.uvd, self.uvp]: - pytest.raises(ValueError, utils.spw_range_from_redshifts, obj, - z_range=(5., 8.)) # lower bound - pytest.raises(ValueError, utils.spw_range_from_redshifts, obj, - z_range=(10., 20.)) # upper bound - pytest.raises(ValueError, utils.spw_range_from_redshifts, obj, - z_range=(11., 10.)) # wrong order + pytest.raises( + ValueError, utils.spw_range_from_redshifts, obj, z_range=(5.0, 8.0) + ) # lower bound + pytest.raises( + ValueError, utils.spw_range_from_redshifts, obj, z_range=(10.0, 20.0) + ) # upper bound + pytest.raises( + ValueError, utils.spw_range_from_redshifts, obj, z_range=(11.0, 10.0) + ) # wrong order # Check that valid frequency ranges are returned z_list = [(6.5, 7.5), (7.5, 8.5), (8.5, 9.5)] - spw1 = utils.spw_range_from_redshifts(self.uvd, z_range=(7., 8.)) + spw1 = utils.spw_range_from_redshifts(self.uvd, z_range=(7.0, 8.0)) spw2 = utils.spw_range_from_redshifts(self.uvd, z_range=z_list) - spw3 = utils.spw_range_from_redshifts(self.uvd, z_range=(12., 14.), - bounds_error=False) - spw4 = utils.spw_range_from_redshifts(self.uvd, z_range=(6.2, 7.2)) + utils.spw_range_from_redshifts( + self.uvd, z_range=(12.0, 14.0), bounds_error=False + ) + utils.spw_range_from_redshifts(self.uvd, z_range=(6.2, 7.2)) # Make sure tuple vs. list arguments were handled correctly - assert( isinstance(spw1, tuple) ) - assert( isinstance(spw2, list) ) - assert( len(spw2) == len(z_list) ) + assert isinstance(spw1, tuple) + assert isinstance(spw2, list) + assert len(spw2) == len(z_list) # Make sure that this also works for UVPSpec objects spw5 = utils.spw_range_from_redshifts(self.uvp, z_range=(13.1, 13.2)) - assert( isinstance(spw5, tuple) ) - assert( spw5[0] is not None ) + assert isinstance(spw5, tuple) + assert spw5[0] is not None def test_calc_blpair_reds(self): - fname = os.path.join(DATA_PATH, 'zen.all.xx.LST.1.06964.uvA') + fname = os.path.join(DATA_PATH, "zen.all.xx.LST.1.06964.uvA") uvd = UVData() uvd.read_miriad(fname) # basic execution - (bls1, bls2, blps, xants1, xants2, rgrps, lens, - angs) = utils.calc_blpair_reds(uvd, uvd, filter_blpairs=True, extra_info=True, - exclude_auto_bls=False, exclude_permutations=True) + (bls1, bls2, blps, xants1, xants2, rgrps, lens, angs) = utils.calc_blpair_reds( + uvd, + uvd, + filter_blpairs=True, + extra_info=True, + exclude_auto_bls=False, + exclude_permutations=True, + ) assert len(bls1) == len(bls2) == 15 assert blps == list(zip(bls1, bls2)) assert xants1 == xants2 @@ -165,26 +184,46 @@ def test_calc_blpair_reds(self): assert np.max(rgrps) == len(lens) - 1 # assert rgrps indexes lens / angs # test xant_flag_thresh - (bls1, bls2, blps, xants1, - xants2) = utils.calc_blpair_reds(uvd, uvd, filter_blpairs=True, exclude_auto_bls=True, exclude_permutations=True, - xant_flag_thresh=0.0) + (bls1, bls2, blps, xants1, xants2) = utils.calc_blpair_reds( + uvd, + uvd, + filter_blpairs=True, + exclude_auto_bls=True, + exclude_permutations=True, + xant_flag_thresh=0.0, + ) assert len(bls1) == len(bls2) == 0 # test bl_len_range - (bls1, bls2, blps, xants1, - xants2) = utils.calc_blpair_reds(uvd, uvd, filter_blpairs=True, exclude_auto_bls=False, exclude_permutations=True, - bl_len_range=(0, 15.0)) + (bls1, bls2, blps, xants1, xants2) = utils.calc_blpair_reds( + uvd, + uvd, + filter_blpairs=True, + exclude_auto_bls=False, + exclude_permutations=True, + bl_len_range=(0, 15.0), + ) assert len(bls1) == len(bls2) == 12 - (bls1, bls2, blps, xants1, - xants2) = utils.calc_blpair_reds(uvd, uvd, filter_blpairs=True, exclude_auto_bls=True, exclude_permutations=True, - bl_len_range=(0, 15.0)) + (bls1, bls2, blps, xants1, xants2) = utils.calc_blpair_reds( + uvd, + uvd, + filter_blpairs=True, + exclude_auto_bls=True, + exclude_permutations=True, + bl_len_range=(0, 15.0), + ) assert len(bls1) == len(bls2) == 5 assert np.all([bls1[i] != bls2[i] for i in range(len(blps))]) # test grouping - (bls1, bls2, blps, xants1, - xants2) = utils.calc_blpair_reds(uvd, uvd, filter_blpairs=True, exclude_auto_bls=False, exclude_permutations=True, - Nblps_per_group=2) + (bls1, bls2, blps, xants1, xants2) = utils.calc_blpair_reds( + uvd, + uvd, + filter_blpairs=True, + exclude_auto_bls=False, + exclude_permutations=True, + Nblps_per_group=2, + ) assert len(blps) == 10 assert isinstance(blps[0], list) assert blps[0] == [((24, 37), (25, 38)), ((24, 37), (24, 37))] @@ -192,14 +231,20 @@ def test_calc_blpair_reds(self): # test baseline select on uvd uvd2 = copy.deepcopy(uvd) uvd2.select(bls=[(24, 25), (37, 38), (24, 39)]) - (bls1, bls2, blps, xants1, - xants2) = utils.calc_blpair_reds(uvd2, uvd2, filter_blpairs=True, exclude_auto_bls=True, exclude_permutations=True, - bl_len_range=(10.0, 20.0)) + (bls1, bls2, blps, xants1, xants2) = utils.calc_blpair_reds( + uvd2, + uvd2, + filter_blpairs=True, + exclude_auto_bls=True, + exclude_permutations=True, + bl_len_range=(10.0, 20.0), + ) assert blps == [((24, 25), (37, 38))] # test exclude_cross_bls - (bls1, bls2, blps, xants1, - xants2) = utils.calc_blpair_reds(uvd, uvd, filter_blpairs=True, exclude_cross_bls=True) + (bls1, bls2, blps, xants1, xants2) = utils.calc_blpair_reds( + uvd, uvd, filter_blpairs=True, exclude_cross_bls=True + ) for bl1, bl2 in blps: assert bl1 == bl2 @@ -207,29 +252,41 @@ def test_calc_blpair_reds(self): uvd2 = copy.deepcopy(uvd) uvd2.antenna_positions[0] += 2 pytest.raises(AssertionError, utils.calc_blpair_reds, uvd, uvd2) - pytest.raises(AssertionError, utils.calc_blpair_reds, uvd, uvd, exclude_auto_bls=True, exclude_cross_bls=True) + pytest.raises( + AssertionError, + utils.calc_blpair_reds, + uvd, + uvd, + exclude_auto_bls=True, + exclude_cross_bls=True, + ) def test_calc_blpair_reds_autos_only(self): # test include_crosscorrs selection option being set to false. - fname = os.path.join(DATA_PATH, 'zen.all.xx.LST.1.06964.uvA') + fname = os.path.join(DATA_PATH, "zen.all.xx.LST.1.06964.uvA") uvd = UVData() uvd.read_miriad(fname) # basic execution - (bls1, bls2, blps, xants1, xants2, rgrps, lens, - angs) = utils.calc_blpair_reds(uvd, uvd, filter_blpairs=True, extra_info=True, - exclude_auto_bls=False, exclude_permutations=True, include_crosscorrs=False, - include_autocorrs=True) + (bls1, bls2, blps, xants1, xants2, rgrps, lens, angs) = utils.calc_blpair_reds( + uvd, + uvd, + filter_blpairs=True, + extra_info=True, + exclude_auto_bls=False, + exclude_permutations=True, + include_crosscorrs=False, + include_autocorrs=True, + ) assert len(bls1) > 0 for bl1, bl2 in zip(bls1, bls2): assert bl1[0] == bl1[1] assert bl2[0] == bl2[1] - def test_get_delays(self): - utils.get_delays(np.linspace(100., 200., 50)*1e6) + utils.get_delays(np.linspace(100.0, 200.0, 50) * 1e6) def test_get_reds(self): - fname = os.path.join(DATA_PATH, 'zen.all.xx.LST.1.06964.uvA') + fname = os.path.join(DATA_PATH, "zen.all.xx.LST.1.06964.uvA") uvd = UVData() uvd.read_miriad(fname, read_data=False) antpos, ants = uvd.get_ENU_antpos() @@ -238,7 +295,12 @@ def test_get_reds(self): # test basic execution xants = [0, 1, 2] r, l, a = utils.get_reds(fname, xants=xants) - assert np.all([np.all([bl[0] not in xants and bl[1] not in xants for bl in _r]) for _r in r]) + assert np.all( + [ + np.all([bl[0] not in xants and bl[1] not in xants for bl in _r]) + for _r in r + ] + ) assert len(r) == len(a) == len(l) assert len(r) == 104 @@ -251,9 +313,11 @@ def test_get_reds(self): # restrict bl_len_range = (14, 16) bl_deg_range = (55, 65) - r, l, a = utils.get_reds(uvd, bl_len_range=bl_len_range, bl_deg_range=bl_deg_range) - assert (np.all([_l > bl_len_range[0] and _l < bl_len_range[1] for _l in l])) - assert (np.all([_a > bl_deg_range[0] and _a < bl_deg_range[1] for _a in a])) + r, l, a = utils.get_reds( + uvd, bl_len_range=bl_len_range, bl_deg_range=bl_deg_range + ) + assert np.all([_l > bl_len_range[0] and _l < bl_len_range[1] for _l in l]) + assert np.all([_a > bl_deg_range[0] and _a < bl_deg_range[1] for _a in a]) # min EW cut r, l, a = utils.get_reds(uvd, bl_len_range=(14, 16), min_EW_cut=14) @@ -267,14 +331,13 @@ def test_get_reds(self): assert len(r) == 105 # Check errors when wrong types input - pytest.raises(TypeError, utils.get_reds, [1., 2.]) + pytest.raises(TypeError, utils.get_reds, [1.0, 2.0]) def test_get_reds_autos_only(self): - fname = os.path.join(DATA_PATH, 'zen.all.xx.LST.1.06964.uvA') + fname = os.path.join(DATA_PATH, "zen.all.xx.LST.1.06964.uvA") uvd = UVData() uvd.read_miriad(fname, read_data=False) antpos, ants = uvd.get_ENU_antpos() - antpos_d = dict(list(zip(ants, antpos))) xants = [0, 1, 2] r, l, a = utils.get_reds(fname, xants=xants, autos_only=True, add_autos=True) assert len(r) == 1 @@ -284,30 +347,59 @@ def test_get_reds_autos_only(self): def test_config_pspec_blpairs(self): # test basic execution uv_template = os.path.join(DATA_PATH, "zen.{group}.{pol}.LST.1.28828.uvOCRSA") - groupings = utils.config_pspec_blpairs(uv_template, [('xx', 'xx')], [('even', 'odd')], verbose=False, exclude_auto_bls=True) + groupings = utils.config_pspec_blpairs( + uv_template, + [("xx", "xx")], + [("even", "odd")], + verbose=False, + exclude_auto_bls=True, + ) assert len(groupings) == 1 - assert list(groupings.keys())[0] == (('even', 'odd'), ('xx', 'xx')) + assert list(groupings.keys())[0] == (("even", "odd"), ("xx", "xx")) assert len(list(groupings.values())[0]) == 11833 # test multiple, some non-existant pairs - groupings = utils.config_pspec_blpairs(uv_template, [('xx', 'xx'), ('yy', 'yy')], [('even', 'odd'), ('even', 'odd')], verbose=False, exclude_auto_bls=True) + groupings = utils.config_pspec_blpairs( + uv_template, + [("xx", "xx"), ("yy", "yy")], + [("even", "odd"), ("even", "odd")], + verbose=False, + exclude_auto_bls=True, + ) assert len(groupings) == 1 - assert list(groupings.keys())[0] == (('even', 'odd'), ('xx', 'xx')) + assert list(groupings.keys())[0] == (("even", "odd"), ("xx", "xx")) # test xants - groupings = utils.config_pspec_blpairs(uv_template, [('xx', 'xx')], [('even', 'odd')], xants=[0, 1, 2], verbose=False, exclude_auto_bls=True) + groupings = utils.config_pspec_blpairs( + uv_template, + [("xx", "xx")], + [("even", "odd")], + xants=[0, 1, 2], + verbose=False, + exclude_auto_bls=True, + ) assert len(list(groupings.values())[0]) == 9735 # test exclude_patterns - groupings = utils.config_pspec_blpairs(uv_template, - [('xx', 'xx'), ('yy', 'yy')], - [('even', 'odd'), ('even', 'odd')], - exclude_patterns=['1.288'], - verbose=False, exclude_auto_bls=True) + groupings = utils.config_pspec_blpairs( + uv_template, + [("xx", "xx"), ("yy", "yy")], + [("even", "odd"), ("even", "odd")], + exclude_patterns=["1.288"], + verbose=False, + exclude_auto_bls=True, + ) assert len(groupings) == 0 # test exceptions - pytest.raises(AssertionError, utils.config_pspec_blpairs, uv_template, [('xx', 'xx'), ('xx', 'xx')], [('even', 'odd')], verbose=False) + pytest.raises( + AssertionError, + utils.config_pspec_blpairs, + uv_template, + [("xx", "xx"), ("xx", "xx")], + [("even", "odd")], + verbose=False, + ) def test_log(): @@ -333,8 +425,8 @@ def test_log(): utils.log("raised an exception", f=logf, tb=sys.exc_info(), verbose=False) logf.close() with open("logf.log", "r") as f: - log = ''.join(f.readlines()) - assert ("NameError" in log and "raised an exception" in log) + log = "".join(f.readlines()) + assert "NameError" in log and "raised an exception" in log os.remove("logf.log") @@ -348,43 +440,51 @@ def test_get_blvec_reds(): # test execution w/ dictionary blvecs = dict(list(zip(uvp.bl_array, uvp.get_ENU_bl_vecs()))) - (red_bl_grp, red_bl_len, red_bl_ang, - red_bl_tag) = utils.get_blvec_reds(blvecs, bl_error_tol=1.0) + (red_bl_grp, red_bl_len, red_bl_ang, red_bl_tag) = utils.get_blvec_reds( + blvecs, bl_error_tol=1.0 + ) assert len(red_bl_grp) == 2 - assert red_bl_tag == ['015_060', '015_120'] + assert red_bl_tag == ["015_060", "015_120"] # test w/ a UVPSpec - (red_bl_grp, red_bl_len, red_bl_ang, - red_bl_tag) = utils.get_blvec_reds(uvp, bl_error_tol=1.0) + (red_bl_grp, red_bl_len, red_bl_ang, red_bl_tag) = utils.get_blvec_reds( + uvp, bl_error_tol=1.0 + ) assert len(red_bl_grp) == 2 - assert red_bl_tag == ['015_060', '015_120'] + assert red_bl_tag == ["015_060", "015_120"] # test w/ zero tolerance: each blpair is its own group - (red_bl_grp, red_bl_len, red_bl_ang, - red_bl_tag) = utils.get_blvec_reds(uvp, bl_error_tol=0.0) + (red_bl_grp, red_bl_len, red_bl_ang, red_bl_tag) = utils.get_blvec_reds( + uvp, bl_error_tol=0.0 + ) assert len(red_bl_grp) == uvp.Nblpairs # test combine angles uvp = testing.uvpspec_from_data(fname, reds[:3], spw_ranges=[(10, 40)]) - (red_bl_grp, red_bl_len, red_bl_ang, - red_bl_tag) = utils.get_blvec_reds(uvp, bl_error_tol=1.0, match_bl_lens=True) + (red_bl_grp, red_bl_len, red_bl_ang, red_bl_tag) = utils.get_blvec_reds( + uvp, bl_error_tol=1.0, match_bl_lens=True + ) assert len(red_bl_grp) == 1 + def test_uvp_noise_error_arser(): # test argparser for noise error bars. ap = utils.uvp_noise_error_parser() - args=ap.parse_args(["container.hdf5", "autos.uvh5", "beam.uvbeam", "--groups", "dset0_dset1"]) + args = ap.parse_args( + ["container.hdf5", "autos.uvh5", "beam.uvbeam", "--groups", "dset0_dset1"] + ) assert args.pspec_container == "container.hdf5" assert args.auto_file == "autos.uvh5" assert args.beam == "beam.uvbeam" assert args.groups == ["dset0_dset1"] assert args.spectra is None + def test_job_monitor(): # open empty files - datafiles = ["./{}".format(i) for i in ['a', 'b', 'c', 'd']] + datafiles = ["./{}".format(i) for i in ["a", "b", "c", "d"]] for df in datafiles: - with open(df, 'w') as f: + with open(df, "w"): pass def run_func(i, datafiles=datafiles): @@ -395,7 +495,7 @@ def run_func(i, datafiles=datafiles): if rand_num > 0.7: raise ValueError df = datafiles[i] - with open(df, 'a') as f: + with open(df, "a") as f: f.write("Hello World") except: return 1 @@ -405,12 +505,16 @@ def run_func(i, datafiles=datafiles): # set seed np.random.seed(0) # run over datafiles - failures = utils.job_monitor(run_func, range(len(datafiles)), "test", maxiter=1, verbose=False) + failures = utils.job_monitor( + run_func, range(len(datafiles)), "test", maxiter=1, verbose=False + ) # assert job 1 failed np.testing.assert_array_equal(failures, np.array([1])) # try with reruns np.random.seed(0) - failures = utils.job_monitor(run_func, range(len(datafiles)), "test", maxiter=10, verbose=False) + failures = utils.job_monitor( + run_func, range(len(datafiles)), "test", maxiter=10, verbose=False + ) # assert no failures now assert len(failures) == 0 diff --git a/hera_pspec/tests/test_uvpspec.py b/hera_pspec/tests/test_uvpspec.py index c7780469..8a5097e4 100644 --- a/hera_pspec/tests/test_uvpspec.py +++ b/hera_pspec/tests/test_uvpspec.py @@ -2,42 +2,46 @@ import pytest import numpy as np import os -import sys from hera_pspec.data import DATA_PATH from .. import uvpspec, conversions, parameter, pspecbeam, pspecdata, testing, utils from .. import uvpspec_utils as uvputils import copy -import h5py from collections import OrderedDict as odict from pyuvdata import UVData from hera_cal import redcal -import json -class Test_UVPSpec(unittest.TestCase): +class Test_UVPSpec(unittest.TestCase): def setUp(self): - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") self.beam = pspecbeam.PSpecBeamUV(beamfile) uvp, cosmo = testing.build_vanilla_uvpspec(beam=self.beam) self.uvp = uvp # setup for exact windows related tests - self.ft_file = os.path.join(DATA_PATH, 'FT_beam_HERA_dipole_test') + self.ft_file = os.path.join(DATA_PATH, "FT_beam_HERA_dipole_test") # Instantiate UVWindow() # obtain uvp object - datafile = os.path.join(DATA_PATH, 'zen.2458116.31939.HH.uvh5') + datafile = os.path.join(DATA_PATH, "zen.2458116.31939.HH.uvh5") # read datafile uvd = UVData() uvd.read_uvh5(datafile) # Create a new PSpecData objec ds = pspecdata.PSpecData(dsets=[uvd, uvd], wgts=[None, None]) # choose baselines - baselines1, baselines2, blpairs = utils.construct_blpairs(uvd.get_antpairs()[1:], - exclude_permutations=False, - exclude_auto_bls=True) + baselines1, baselines2, blpairs = utils.construct_blpairs( + uvd.get_antpairs()[1:], exclude_permutations=False, exclude_auto_bls=True + ) # compute ps - self.uvp_wf = ds.pspec(baselines1, baselines2, dsets=(0, 1), pols=[('xx','xx')], - spw_ranges=(175,195), taper='bh',verbose=False) + self.uvp_wf = ds.pspec( + baselines1, + baselines2, + dsets=(0, 1), + pols=[("xx", "xx")], + spw_ranges=(175, 195), + taper="bh", + verbose=False, + ) def tearDown(self): pass @@ -49,18 +53,24 @@ def _add_optionals(self, uvp): """add dummy optional cov_array and stats_array to uvp""" uvp.cov_array_real = odict() uvp.cov_array_imag = odict() - uvp.cov_model = 'empirical' - stat = 'noise_err' + uvp.cov_model = "empirical" + stat = "noise_err" uvp.stats_array = odict({stat: odict()}) for spw in uvp.spw_array: ndlys = uvp.get_spw_ranges(spw)[0][-1] - uvp.cov_array_real[spw] = np.empty((uvp.Nblpairts, ndlys, ndlys, uvp.Npols), np.float64) - uvp.cov_array_imag[spw] = np.empty((uvp.Nblpairts, ndlys, ndlys, uvp.Npols), np.float64) - uvp.stats_array[stat][spw] = np.empty((uvp.Nblpairts, ndlys, uvp.Npols), np.complex128) + uvp.cov_array_real[spw] = np.empty( + (uvp.Nblpairts, ndlys, ndlys, uvp.Npols), np.float64 + ) + uvp.cov_array_imag[spw] = np.empty( + (uvp.Nblpairts, ndlys, ndlys, uvp.Npols), np.float64 + ) + uvp.stats_array[stat][spw] = np.empty( + (uvp.Nblpairts, ndlys, uvp.Npols), np.complex128 + ) return uvp def test_param(self): - a = parameter.PSpecParam("example", description="example", expected_type=int) + parameter.PSpecParam("example", description="example", expected_type=int) def test_eq(self): # test equivalence @@ -68,29 +78,29 @@ def test_eq(self): def test_get_funcs(self): # get_data - d = self.uvp.get_data((0, ((1, 2), (1, 2)), ('xx','xx'))) + d = self.uvp.get_data((0, ((1, 2), (1, 2)), ("xx", "xx"))) assert d.shape == (10, 30) - assert(d.dtype == np.complex) - np.testing.assert_almost_equal(d[0,0], (101.1021011020000001+0j)) + assert d.dtype == np.complex + np.testing.assert_almost_equal(d[0, 0], (101.1021011020000001 + 0j)) d = self.uvp.get_data((0, ((1, 2), (1, 2)), 1515)) - np.testing.assert_almost_equal(d[0,0], (101.1021011020000001+0j)) + np.testing.assert_almost_equal(d[0, 0], (101.1021011020000001 + 0j)) d = self.uvp.get_data((0, 101102101102, 1515)) - np.testing.assert_almost_equal(d[0,0], (101.1021011020000001+0j)) + np.testing.assert_almost_equal(d[0, 0], (101.1021011020000001 + 0j)) # get_wgts - w = self.uvp.get_wgts((0, ((1, 2), (1, 2)), ('xx','xx'))) - assert w.shape == (10, 50, 2) # should have Nfreq dim, not Ndlys + w = self.uvp.get_wgts((0, ((1, 2), (1, 2)), ("xx", "xx"))) + assert w.shape == (10, 50, 2) # should have Nfreq dim, not Ndlys assert w.dtype == np.float - assert w[0,0,0] == 1.0 + assert w[0, 0, 0] == 1.0 # get_integrations - i = self.uvp.get_integrations((0, ((1, 2), (1, 2)), ('xx','xx'))) + i = self.uvp.get_integrations((0, ((1, 2), (1, 2)), ("xx", "xx"))) assert i.shape == (10,) assert i.dtype == np.float np.testing.assert_almost_equal(i[0], 1.0) # get nsample - n = self.uvp.get_nsamples((0, ((1, 2), (1, 2)), ('xx', 'xx'))) + n = self.uvp.get_nsamples((0, ((1, 2), (1, 2)), ("xx", "xx"))) assert n.shape == (10,) assert n.dtype == np.float np.testing.assert_almost_equal(n[0], 1.0) @@ -102,7 +112,7 @@ def test_get_funcs(self): # get blpair seps blp = self.uvp.get_blpair_seps() assert len(blp) == 30 - assert(np.isclose(blp, 14.60, rtol=1e-1, atol=1e-1).all()) + assert np.isclose(blp, 14.60, rtol=1e-1, atol=1e-1).all() # get kvecs k_perp, k_para = self.uvp.get_kperps(0), self.uvp.get_kparas(0) @@ -110,12 +120,12 @@ def test_get_funcs(self): assert len(k_para) == 30 # test key expansion - key = (0, ((1, 2), (1, 2)), ('xx','xx')) + key = (0, ((1, 2), (1, 2)), ("xx", "xx")) d = self.uvp.get_data(key) assert d.shape == (10, 30) # test key as dictionary - key = {'spw':0, 'blpair':((1, 2), (1, 2)), 'polpair': ('xx','xx')} + key = {"spw": 0, "blpair": ((1, 2), (1, 2)), "polpair": ("xx", "xx")} d = self.uvp.get_data(key) assert d.shape == (10, 30) @@ -131,89 +141,137 @@ def test_get_funcs(self): # test get_polpairs polpairs = self.uvp.get_polpairs() - assert polpairs == [('xx', 'xx')] + assert polpairs == [("xx", "xx")] # test get all keys keys = self.uvp.get_all_keys() - assert keys == [(0, ((1, 2), (1, 2)), ('xx', 'xx')), - (0, ((2, 3), (2, 3)), ('xx', 'xx')), - (0, ((1, 3), (1, 3)), ('xx', 'xx'))] + assert keys == [ + (0, ((1, 2), (1, 2)), ("xx", "xx")), + (0, ((2, 3), (2, 3)), ("xx", "xx")), + (0, ((1, 3), (1, 3)), ("xx", "xx")), + ] # test omit_flags - self.uvp.integration_array[0][self.uvp.blpair_to_indices(((1, 2), (1, 2)))[:2]] = 0.0 - assert self.uvp.get_integrations((0, ((1, 2), (1, 2)), ('xx','xx')), omit_flags=True).shape == (8,) + self.uvp.integration_array[0][ + self.uvp.blpair_to_indices(((1, 2), (1, 2)))[:2] + ] = 0.0 + assert self.uvp.get_integrations( + (0, ((1, 2), (1, 2)), ("xx", "xx")), omit_flags=True + ).shape == (8,) def test_get_covariance(self): - dfile = os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA') + dfile = os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA") uvd = UVData() uvd.read(dfile) cosmo = conversions.Cosmo_Conversions() - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") uvb = pspecbeam.PSpecBeamUV(beamfile, cosmo=cosmo) - Jy_to_mK = uvb.Jy_to_mK(np.unique(uvd.freq_array), pol='XX') + Jy_to_mK = uvb.Jy_to_mK(np.unique(uvd.freq_array), pol="XX") uvd.data_array *= Jy_to_mK[None, None, :, None] - uvd1 = uvd.select(times=np.unique(uvd.time_array)[:(uvd.Ntimes//2):1], inplace=False) - uvd2 = uvd.select(times=np.unique(uvd.time_array)[(uvd.Ntimes//2):(uvd.Ntimes//2 + uvd.Ntimes//2):1], inplace=False) + uvd1 = uvd.select( + times=np.unique(uvd.time_array)[: (uvd.Ntimes // 2) : 1], inplace=False + ) + uvd2 = uvd.select( + times=np.unique(uvd.time_array)[ + (uvd.Ntimes // 2) : (uvd.Ntimes // 2 + uvd.Ntimes // 2) : 1 + ], + inplace=False, + ) ds = pspecdata.PSpecData(dsets=[uvd1, uvd2], wgts=[None, None], beam=uvb) ds.rephase_to_dset(0) - spws = utils.spw_range_from_freqs(uvd, freq_range=[(160e6, 165e6), (160e6, 165e6)], bounds_error=True) + spws = utils.spw_range_from_freqs( + uvd, freq_range=[(160e6, 165e6), (160e6, 165e6)], bounds_error=True + ) antpos, ants = uvd.get_ENU_antpos(pick_data_ants=True) antpos = dict(zip(ants, antpos)) red_bls = redcal.get_pos_reds(antpos, bl_error_tol=1.0) - bls1, bls2, blpairs = utils.construct_blpairs(red_bls[3], exclude_auto_bls=True, exclude_permutations=True) - - uvp = ds.pspec( bls1, bls2, (0, 1), [('xx', 'xx')], spw_ranges=spws, input_data_weight='identity', - norm='I', taper='blackman-harris', store_cov = True, cov_model='autos', verbose=False) - - key = (0,blpairs[0],"xx") - - cov_real = uvp.get_cov(key, component='real') + bls1, bls2, blpairs = utils.construct_blpairs( + red_bls[3], exclude_auto_bls=True, exclude_permutations=True + ) + + uvp = ds.pspec( + bls1, + bls2, + (0, 1), + [("xx", "xx")], + spw_ranges=spws, + input_data_weight="identity", + norm="I", + taper="blackman-harris", + store_cov=True, + cov_model="autos", + verbose=False, + ) + + key = (0, blpairs[0], "xx") + + cov_real = uvp.get_cov(key, component="real") assert cov_real[0].shape == (50, 50) - cov_imag = uvp.get_cov(key, component='imag') + cov_imag = uvp.get_cov(key, component="imag") assert cov_imag[0].shape == (50, 50) uvp.fold_spectra() - cov_real = uvp.get_cov(key, component='real') + cov_real = uvp.get_cov(key, component="real") assert cov_real[0].shape == (24, 24) - cov_imag = uvp.get_cov(key, component='imag') + cov_imag = uvp.get_cov(key, component="imag") assert cov_imag[0].shape == (24, 24) - def test_stats_array(self): # test get_data and set_data uvp = copy.deepcopy(self.uvp) keys = uvp.get_all_keys() - pytest.raises(ValueError, uvp.set_stats, "errors", keys[0], np.linspace(0, 1, 2)) + pytest.raises( + ValueError, uvp.set_stats, "errors", keys[0], np.linspace(0, 1, 2) + ) pytest.raises(AttributeError, uvp.get_stats, "__", keys[0]) errs = np.ones((uvp.Ntimes, uvp.Ndlys)) for key in keys: uvp.set_stats("errors", key, errs) - e = uvp.get_stats("errors", keys[0]) - assert(np.all(uvp.get_stats("errors", keys[0]) == errs)) + uvp.get_stats("errors", keys[0]) + assert np.all(uvp.get_stats("errors", keys[0]) == errs) # self.uvp.set_stats("errors", keys[0], -99.) blpairs = uvp.get_blpairs() - u = uvp.average_spectra([blpairs], time_avg=False, error_weights="errors", inplace=False) - assert(np.all(np.isclose(u.get_stats("errors", keys[0])[0], np.ones(u.Ndlys)/np.sqrt(len(blpairs))))) + u = uvp.average_spectra( + [blpairs], time_avg=False, error_weights="errors", inplace=False + ) + assert np.all( + np.isclose( + u.get_stats("errors", keys[0])[0], + np.ones(u.Ndlys) / np.sqrt(len(blpairs)), + ) + ) for key in keys: uvp.set_stats("who?", key, errs) - u = uvp.average_spectra([blpairs], time_avg=False, error_field=["errors", "who?"], inplace=False) - u2 = uvp.average_spectra([blpairs], time_avg=True, error_field=["errors", "who?"], inplace=False) - assert(np.all( u.get_stats("errors", keys[0]) == u.get_stats("who?", keys[0]))) + u = uvp.average_spectra( + [blpairs], time_avg=False, error_field=["errors", "who?"], inplace=False + ) + uvp.average_spectra( + [blpairs], time_avg=True, error_field=["errors", "who?"], inplace=False + ) + assert np.all(u.get_stats("errors", keys[0]) == u.get_stats("who?", keys[0])) u.select(times=np.unique(u.time_avg_array)[:20]) u3 = uvp.average_spectra([blpairs], time_avg=True, inplace=False) - pytest.raises(KeyError, uvp.average_spectra, [blpairs], time_avg=True, inplace=False, error_field=["..............."]) - assert hasattr(u3, "stats_array") == False - if os.path.exists('./ex.hdf5'): os.remove('./ex.hdf5') - u.write_hdf5('./ex.hdf5') - u.read_hdf5('./ex.hdf5') - os.remove('./ex.hdf5') + pytest.raises( + KeyError, + uvp.average_spectra, + [blpairs], + time_avg=True, + inplace=False, + error_field=["..............."], + ) + assert not hasattr(u3, "stats_array") + if os.path.exists("./ex.hdf5"): + os.remove("./ex.hdf5") + u.write_hdf5("./ex.hdf5") + u.read_hdf5("./ex.hdf5") + os.remove("./ex.hdf5") # test folding uvp = copy.deepcopy(self.uvp) @@ -221,30 +279,39 @@ def test_stats_array(self): uvp.set_stats("test", keys[0], errs) uvp.fold_spectra() # fold by summing in inverse quadrature - folded_errs = np.sum([1/errs[:, 1:15][:, ::-1]**2.0, 1/errs[:, 16:]**2.0], axis=0)**(-0.5) - np.testing.assert_array_almost_equal(uvp.get_stats("test", keys[0]), folded_errs) + folded_errs = np.sum( + [1 / errs[:, 1:15][:, ::-1] ** 2.0, 1 / errs[:, 16:] ** 2.0], axis=0 + ) ** (-0.5) + np.testing.assert_array_almost_equal( + uvp.get_stats("test", keys[0]), folded_errs + ) # test set_stats_slice uvp = copy.deepcopy(self.uvp) - key = (0, ((1, 2), (1, 2)), ('xx', 'xx')) - uvp.set_stats('err', key, np.ones((uvp.Ntimes, uvp.Ndlys))) - uvp.set_stats_slice('err', 50, 0, above=True, val=10) + key = (0, ((1, 2), (1, 2)), ("xx", "xx")) + uvp.set_stats("err", key, np.ones((uvp.Ntimes, uvp.Ndlys))) + uvp.set_stats_slice("err", 50, 0, above=True, val=10) # ensure all dlys above 50 * 15 ns are set to 10 and all others set to 1 - assert np.isclose(uvp.get_stats('err', key)[:, np.abs(uvp.get_dlys(0)*1e9) > 15 * 50], 10).all() - assert np.isclose(uvp.get_stats('err', key)[:, np.abs(uvp.get_dlys(0)*1e9) < 15 * 50], 1).all() + assert np.isclose( + uvp.get_stats("err", key)[:, np.abs(uvp.get_dlys(0) * 1e9) > 15 * 50], 10 + ).all() + assert np.isclose( + uvp.get_stats("err", key)[:, np.abs(uvp.get_dlys(0) * 1e9) < 15 * 50], 1 + ).all() def test_convert_deltasq(self): # setup uvp build uvd = UVData() - uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, - "HERA_NF_dipole_power.beamfits")) + uvd.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) + beam = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + ) uvd_std = copy.deepcopy(uvd) # dummy uvd_std uvd_std.data_array[:] = 1.0 bls = [(37, 38), (38, 39), (52, 53)] - uvp = testing.uvpspec_from_data(uvd, bls, data_std=uvd_std, - spw_ranges=[(20, 30), (60, 90)], - beam=beam) + uvp = testing.uvpspec_from_data( + uvd, bls, data_std=uvd_std, spw_ranges=[(20, 30), (60, 90)], beam=beam + ) # dummy stats_array build Tsys = utils.uvd_to_Tsys(uvd, beam) utils.uvp_noise_error(uvp, Tsys) @@ -253,17 +320,23 @@ def test_convert_deltasq(self): dsq = uvp.convert_to_deltasq(inplace=False) for spw in uvp.spw_array: k_perp, k_para = uvp.get_kperps(spw), uvp.get_kparas(spw) - k_mag = np.sqrt(k_perp[:, None, None]**2 + k_para[None, :, None]**2) + k_mag = np.sqrt(k_perp[:, None, None] ** 2 + k_para[None, :, None] ** 2) coeff = k_mag**3 / (2 * np.pi**2) # check data - assert np.isclose(dsq.data_array[spw][0, :, 0], (uvp.data_array[spw]*coeff)[0, :, 0]).all() + assert np.isclose( + dsq.data_array[spw][0, :, 0], (uvp.data_array[spw] * coeff)[0, :, 0] + ).all() # check stats - assert np.isclose(dsq.stats_array['P_N'][spw][0, :, 0], - (uvp.stats_array['P_N'][spw] * coeff)[0, :, 0]).all() + assert np.isclose( + dsq.stats_array["P_N"][spw][0, :, 0], + (uvp.stats_array["P_N"][spw] * coeff)[0, :, 0], + ).all() # check cov - assert np.isclose(dsq.cov_array_real[spw][0, :, :, 0].diagonal(), - uvp.cov_array_real[spw][0, :, :, 0].diagonal()*coeff[0, :, 0]**2).all() - assert dsq.norm_units == uvp.norm_units + ' k^3 / (2pi^2)' + assert np.isclose( + dsq.cov_array_real[spw][0, :, :, 0].diagonal(), + uvp.cov_array_real[spw][0, :, :, 0].diagonal() * coeff[0, :, 0] ** 2, + ).all() + assert dsq.norm_units == uvp.norm_units + " k^3 / (2pi^2)" def test_blpair_conversions(self): # test blpair -> antnums @@ -281,21 +354,23 @@ def test_blpair_conversions(self): def test_indices_funcs(self): # key to indices - spw, blpairts, pol = self.uvp.key_to_indices( (0, ((1,2),(1,2)), 1515) ) + spw, blpairts, pol = self.uvp.key_to_indices((0, ((1, 2), (1, 2)), 1515)) assert spw == 0 assert pol == 0 - assert(np.isclose(blpairts, - np.array([0,3,6,9,12,15,18,21,24,27])).min()) - spw, blpairts, pol = self.uvp.key_to_indices( (0, 101102101102, ('xx','xx')) ) + assert np.isclose( + blpairts, np.array([0, 3, 6, 9, 12, 15, 18, 21, 24, 27]) + ).min() + spw, blpairts, pol = self.uvp.key_to_indices((0, 101102101102, ("xx", "xx"))) assert spw == 0 assert pol == 0 - assert(np.isclose(blpairts, - np.array([0,3,6,9,12,15,18,21,24,27])).min()) + assert np.isclose( + blpairts, np.array([0, 3, 6, 9, 12, 15, 18, 21, 24, 27]) + ).min() # Check different polpair specification methods give the same results - s1, b1, p1 = self.uvp.key_to_indices( (0, ((1,2),(1,2)), 1515) ) - s2, b2, p2 = self.uvp.key_to_indices( (0, ((1,2),(1,2)), ('xx','xx')) ) - s3, b3, p3 = self.uvp.key_to_indices( (0, ((1,2),(1,2)), 'xx') ) + s1, b1, p1 = self.uvp.key_to_indices((0, ((1, 2), (1, 2)), 1515)) + s2, b2, p2 = self.uvp.key_to_indices((0, ((1, 2), (1, 2)), ("xx", "xx"))) + s3, b3, p3 = self.uvp.key_to_indices((0, ((1, 2), (1, 2)), "xx")) assert p1 == p2 == p3 # spw to indices @@ -315,28 +390,28 @@ def test_indices_funcs(self): np.testing.assert_array_equal(spw3, spw3b) # pol to indices - pol = self.uvp.polpair_to_indices(('xx','xx')) + pol = self.uvp.polpair_to_indices(("xx", "xx")) assert len(pol) == 1 pol = self.uvp.polpair_to_indices(1515) assert len(pol) == 1 - pol = self.uvp.polpair_to_indices([('xx','xx'), ('xx','xx')]) + pol = self.uvp.polpair_to_indices([("xx", "xx"), ("xx", "xx")]) assert len(pol) == 1 pytest.raises(TypeError, self.uvp.polpair_to_indices, 3.14) # test blpair to indices inds = self.uvp.blpair_to_indices(101102101102) - assert(np.isclose(inds, np.array([0,3,6,9,12,15,18,21,24,27])).min()) - inds = self.uvp.blpair_to_indices(((1,2),(1,2))) - assert(np.isclose(inds, np.array([0,3,6,9,12,15,18,21,24,27])).min()) + assert np.isclose(inds, np.array([0, 3, 6, 9, 12, 15, 18, 21, 24, 27])).min() + inds = self.uvp.blpair_to_indices(((1, 2), (1, 2))) + assert np.isclose(inds, np.array([0, 3, 6, 9, 12, 15, 18, 21, 24, 27])).min() inds = self.uvp.blpair_to_indices([101102101102, 101102101102]) - inds = self.uvp.blpair_to_indices([((1,2),(1,2)), ((1,2),(1,2))]) + inds = self.uvp.blpair_to_indices([((1, 2), (1, 2)), ((1, 2), (1, 2))]) # test time to indices time = self.uvp.time_avg_array[5] blpair = 101102101102 inds = self.uvp.time_to_indices(time=time) assert len(inds) == 3 - assert(np.isclose(self.uvp.time_avg_array[inds], time, rtol=1e-10).all()) + assert np.isclose(self.uvp.time_avg_array[inds], time, rtol=1e-10).all() inds = self.uvp.time_to_indices(time=time, blpairs=[blpair]) assert len(inds) == 1 assert self.uvp.blpair_array[inds] == blpair @@ -348,32 +423,41 @@ def test_select(self): uvp.select(bls=[(1, 2)], inplace=True) assert uvp.Nblpairs == 1 assert uvp.data_array[0].shape == (10, 30, 1) - np.testing.assert_almost_equal(uvp.data_array[0][0,0,0], (101.1021011020000001+0j)) + np.testing.assert_almost_equal( + uvp.data_array[0][0, 0, 0], (101.1021011020000001 + 0j) + ) # inplace vs not inplace, spw selection uvd = UVData() - uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits")) + uvd.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) + beam = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + ) bls = [(37, 38), (38, 39), (52, 53)] - rp = {'filter_centers':[0.], - 'filter_half_widths':[250e-9], - 'filter_factors':[1e-9]} + rp = { + "filter_centers": [0.0], + "filter_half_widths": [250e-9], + "filter_factors": [1e-9], + } r_params = {} for bl in bls: - key1 = bl + ('xx',) + key1 = bl + ("xx",) r_params[key1] = rp - uvp1 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(20, 30), (60, 90)], beam=beam, - r_params = r_params) + uvp1 = testing.uvpspec_from_data( + uvd, bls, spw_ranges=[(20, 30), (60, 90)], beam=beam, r_params=r_params + ) uvp2 = uvp1.select(spws=0, inplace=False) assert uvp2.Nspws == 1 uvp2 = uvp2.select(bls=[(37, 38), (38, 39)], inplace=False) assert uvp2.Nblpairs == 1 assert uvp2.data_array[0].shape == (10, 10, 1) - np.testing.assert_almost_equal(uvp2.data_array[0][0,0,0], (-3831605.3903496987+8103523.9604128916j)) + np.testing.assert_almost_equal( + uvp2.data_array[0][0, 0, 0], (-3831605.3903496987 + 8103523.9604128916j) + ) assert len(uvp2.get_r_params().keys()) == 2 for rpkey in uvp2.get_r_params(): - assert(rpkey == (37, 38, 'xx') or rpkey == (38, 39, 'xx')) + assert rpkey == (37, 38, "xx") or rpkey == (38, 39, "xx") # blpair select uvp = copy.deepcopy(self.uvp) @@ -381,7 +465,12 @@ def test_select(self): assert uvp2.Nblpairs == 2 # pol select - uvp2 = uvp.select(polpairs=[1515,], inplace=False) + uvp2 = uvp.select( + polpairs=[ + 1515, + ], + inplace=False, + ) assert uvp2.polpair_array[0] == 1515 # time select @@ -390,11 +479,12 @@ def test_select(self): # test pol and blpair select, and check dimensionality of output uvp = copy.deepcopy(self.uvp) - uvp.set_stats('hi', uvp.get_all_keys()[0], np.ones(300).reshape(10, 30)) - uvp2 = uvp.select(blpairs=uvp.get_blpairs(), polpairs=uvp.polpair_array, - inplace=False) + uvp.set_stats("hi", uvp.get_all_keys()[0], np.ones(300).reshape(10, 30)) + uvp2 = uvp.select( + blpairs=uvp.get_blpairs(), polpairs=uvp.polpair_array, inplace=False + ) assert uvp2.data_array[0].shape == (30, 30, 1) - assert uvp2.stats_array['hi'][0].shape == (30, 30, 1) + assert uvp2.stats_array["hi"][0].shape == (30, 30, 1) # test when both blp and pol array are non-sliceable uvp2, uvp3, uvp4 = copy.deepcopy(uvp), copy.deepcopy(uvp), copy.deepcopy(uvp) @@ -402,8 +492,9 @@ def test_select(self): uvp3.polpair_array[0] = 1313 uvp4.polpair_array[0] = 1212 uvp = uvp + uvp2 + uvp3 + uvp4 - uvp5 = uvp.select(blpairs=[101102101102], polpairs=[1515, 1414, 1313], - inplace=False) + uvp5 = uvp.select( + blpairs=[101102101102], polpairs=[1515, 1414, 1313], inplace=False + ) assert uvp5.data_array[0].shape == (10, 30, 3) # select only on lst @@ -426,7 +517,7 @@ def test_select(self): def test_get_ENU_bl_vecs(self): bl_vecs = self.uvp.get_ENU_bl_vecs() - assert(np.isclose(bl_vecs[0], np.array([-14.6, 0.0, 0.0]), atol=1e-6).min()) + assert np.isclose(bl_vecs[0], np.array([-14.6, 0.0, 0.0]), atol=1e-6).min() def test_check(self): uvp = copy.deepcopy(self.uvp) @@ -442,48 +533,54 @@ def test_check(self): def test_clear(self): uvp = copy.deepcopy(self.uvp) uvp._clear() - assert hasattr(uvp, 'Ntimes') == False - assert hasattr(uvp, 'data_array') == False + assert hasattr(uvp, "Ntimes") == False + assert not hasattr(uvp, "data_array") def test_get_r_params(self): # inplace vs not inplace, spw selection uvd = UVData() - uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits")) + uvd.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) + beam = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + ) bls = [(37, 38), (38, 39), (52, 53)] - rp = {'filter_centers':[0.], - 'filter_half_widths':[250e-9], - 'filter_factors':[1e-9]} + rp = { + "filter_centers": [0.0], + "filter_half_widths": [250e-9], + "filter_factors": [1e-9], + } r_params = {} for bl in bls: - key1 = bl + ('xx',) + key1 = bl + ("xx",) r_params[key1] = rp - uvp = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(20, 30), (60, 90)], beam=beam, - r_params = r_params) + uvp = testing.uvpspec_from_data( + uvd, bls, spw_ranges=[(20, 30), (60, 90)], beam=beam, r_params=r_params + ) assert r_params == uvp.get_r_params() def test_write_read_hdf5(self): # test basic write execution uvp = copy.deepcopy(self.uvp) - if os.path.exists('./ex.hdf5'): os.remove('./ex.hdf5') - uvp.write_hdf5('./ex.hdf5', overwrite=True) - assert os.path.exists('./ex.hdf5') + if os.path.exists("./ex.hdf5"): + os.remove("./ex.hdf5") + uvp.write_hdf5("./ex.hdf5", overwrite=True) + assert os.path.exists("./ex.hdf5") # test basic read uvp2 = uvpspec.UVPSpec() - uvp2.read_hdf5('./ex.hdf5') + uvp2.read_hdf5("./ex.hdf5") assert uvp == uvp2 # test just meta uvp2 = uvpspec.UVPSpec() - uvp2.read_hdf5('./ex.hdf5', just_meta=True) - assert hasattr(uvp2, 'Ntimes') - assert hasattr(uvp2, 'data_array') == False + uvp2.read_hdf5("./ex.hdf5", just_meta=True) + assert hasattr(uvp2, "Ntimes") + assert not hasattr(uvp2, "data_array") # test exception - pytest.raises(IOError, uvp.write_hdf5, './ex.hdf5', overwrite=False) + pytest.raises(IOError, uvp.write_hdf5, "./ex.hdf5", overwrite=False) # test partial I/O uvp.read_hdf5("./ex.hdf5", bls=[(1, 2)]) @@ -493,94 +590,125 @@ def test_write_read_hdf5(self): # test just meta uvp.read_hdf5("./ex.hdf5", just_meta=True) assert uvp.Nblpairs == 3 - assert hasattr(uvp, 'data_array') == False - if os.path.exists('./ex.hdf5'): os.remove('./ex.hdf5') + assert not hasattr(uvp, "data_array") + if os.path.exists("./ex.hdf5"): + os.remove("./ex.hdf5") # tests with exact windows # test basic write execution uvp = copy.deepcopy(self.uvp_wf) - uvp.get_exact_window_functions(ftbeam_file = self.ft_file, - inplace=True) - if os.path.exists('./ex.hdf5'): - os.remove('./ex.hdf5') - uvp.write_hdf5('./ex.hdf5', overwrite=True) - assert os.path.exists('./ex.hdf5') + uvp.get_exact_window_functions(ftbeam_file=self.ft_file, inplace=True) + if os.path.exists("./ex.hdf5"): + os.remove("./ex.hdf5") + uvp.write_hdf5("./ex.hdf5", overwrite=True) + assert os.path.exists("./ex.hdf5") # test basic read uvp2 = uvpspec.UVPSpec() - uvp2.read_hdf5('./ex.hdf5') + uvp2.read_hdf5("./ex.hdf5") assert uvp == uvp2 def test_sense(self): uvp = copy.deepcopy(self.uvp) # test generate noise spectra - polpair = ('xx', 'xx') - P_N = uvp.generate_noise_spectra(0, polpair, 500, form='Pk', component='real') + polpair = ("xx", "xx") + P_N = uvp.generate_noise_spectra(0, polpair, 500, form="Pk", component="real") assert P_N[101102101102].shape == (10, 30) # test smaller system temp - P_N2 = uvp.generate_noise_spectra(0, polpair, 400, form='Pk', component='real') - assert((P_N[101102101102] > P_N2[101102101102]).all()) + P_N2 = uvp.generate_noise_spectra(0, polpair, 400, form="Pk", component="real") + assert (P_N[101102101102] > P_N2[101102101102]).all() # test complex - P_N2 = uvp.generate_noise_spectra(0, polpair, 500, form='Pk', component='abs') - assert((P_N[101102101102] < P_N2[101102101102]).all()) + P_N2 = uvp.generate_noise_spectra(0, polpair, 500, form="Pk", component="abs") + assert (P_N[101102101102] < P_N2[101102101102]).all() # test Dsq - Dsq = uvp.generate_noise_spectra(0, polpair, 500, form='DelSq', component='real') + Dsq = uvp.generate_noise_spectra( + 0, polpair, 500, form="DelSq", component="real" + ) assert Dsq[101102101102].shape == (10, 30) - assert(Dsq[101102101102][0, 1] < P_N[101102101102][0, 1]) + assert Dsq[101102101102][0, 1] < P_N[101102101102][0, 1] # test a blpair selection and int polpair blpairs = uvp.get_blpairs()[:1] - P_N = uvp.generate_noise_spectra(0, 1515, 500, form='Pk', blpairs=blpairs, component='real') + P_N = uvp.generate_noise_spectra( + 0, 1515, 500, form="Pk", blpairs=blpairs, component="real" + ) assert P_N[101102101102].shape == (10, 30) # test as a dictionary of arrays - Tsys = dict([(uvp.antnums_to_blpair(k), 500 * np.ones((uvp.Ntimes, uvp.Ndlys)) * np.linspace(1, 2, uvp.Ntimes)[:, None]) for k in uvp.get_blpairs()]) - P_N = uvp.generate_noise_spectra(0, 1515, Tsys, form='Pk', blpairs=blpairs, component='real') + Tsys = { + uvp.antnums_to_blpair(k): ( + 500 + * np.ones((uvp.Ntimes, uvp.Ndlys)) + * np.linspace(1, 2, uvp.Ntimes)[:, None] + ) for k in uvp.get_blpairs() + } + + P_N = uvp.generate_noise_spectra( + 0, 1515, Tsys, form="Pk", blpairs=blpairs, component="real" + ) # assert time gradient is captured: 2 * Tsys results in 4 * P_N - assert(np.isclose(P_N[101102101102][0, 0] * 4, P_N[101102101102][-1, 0])) + assert np.isclose(P_N[101102101102][0, 0] * 4, P_N[101102101102][-1, 0]) def test_average_spectra(self): uvp = copy.deepcopy(self.uvp) # test blpair averaging - blpairs = uvp.get_blpair_groups_from_bl_groups([[101102, 102103, 101103]], - only_pairs_in_bls=False) - uvp2 = uvp.average_spectra(blpair_groups=blpairs, time_avg=False, - inplace=False) + blpairs = uvp.get_blpair_groups_from_bl_groups( + [[101102, 102103, 101103]], only_pairs_in_bls=False + ) + uvp2 = uvp.average_spectra(blpair_groups=blpairs, time_avg=False, inplace=False) assert uvp2.Nblpairs == 1 - assert(np.isclose(uvp2.get_nsamples((0, 101102101102, ('xx','xx'))), 3.0).all()) - assert uvp2.get_data((0, 101102101102, ('xx','xx'))).shape == (10, 30) + assert np.isclose(uvp2.get_nsamples((0, 101102101102, ("xx", "xx"))), 3.0).all() + assert uvp2.get_data((0, 101102101102, ("xx", "xx"))).shape == (10, 30) # Test blpair averaging (with baseline-pair weights) # Results should be identical with different weights here, as the data # are all the same) blpairs = [[101102101102, 101102101102]] - blpair_wgts = [[2., 0.,]] - uvp3a = uvp.average_spectra(blpair_groups=blpairs, time_avg=False, - blpair_weights=None, - inplace=False) - uvp3b = uvp.average_spectra(blpair_groups=blpairs, time_avg=False, - blpair_weights=blpair_wgts, - inplace=False) - #assert uvp2.Nblpairs == 1 - assert(np.isclose( - uvp3a.get_data((0, 101102101102, ('xx','xx'))), - uvp3b.get_data((0, 101102101102, ('xx','xx')))).all()) - #assert uvp2.get_data((0, 101102101102, 'xx')).shape == (10, 30) + blpair_wgts = [ + [ + 2.0, + 0.0, + ] + ] + uvp3a = uvp.average_spectra( + blpair_groups=blpairs, time_avg=False, blpair_weights=None, inplace=False + ) + uvp3b = uvp.average_spectra( + blpair_groups=blpairs, + time_avg=False, + blpair_weights=blpair_wgts, + inplace=False, + ) + # assert uvp2.Nblpairs == 1 + assert np.isclose( + uvp3a.get_data((0, 101102101102, ("xx", "xx"))), + uvp3b.get_data((0, 101102101102, ("xx", "xx"))), + ).all() + # assert uvp2.get_data((0, 101102101102, 'xx')).shape == (10, 30) # test time averaging uvp2 = uvp.average_spectra(time_avg=True, inplace=False) assert uvp2.Ntimes == 1 - assert(np.isclose( - uvp2.get_nsamples((0, 101102101102, ('xx','xx'))), 10.0).all()) - assert uvp2.get_data((0, 101102101102, ('xx','xx'))).shape == (1, 30) + assert np.isclose( + uvp2.get_nsamples((0, 101102101102, ("xx", "xx"))), 10.0 + ).all() + assert uvp2.get_data((0, 101102101102, ("xx", "xx"))).shape == (1, 30) # ensure averaging works when multiple repeated baselines are present, but only # if time_avg = True uvp.blpair_array[uvp.blpair_to_indices(102103102103)] = 101102101102 - pytest.raises(ValueError, uvp.average_spectra, blpair_groups=[list(np.unique(uvp.blpair_array))], time_avg=False, inplace=False) - uvp.average_spectra(blpair_groups=[list(np.unique(uvp.blpair_array))], time_avg=True) + pytest.raises( + ValueError, + uvp.average_spectra, + blpair_groups=[list(np.unique(uvp.blpair_array))], + time_avg=False, + inplace=False, + ) + uvp.average_spectra( + blpair_groups=[list(np.unique(uvp.blpair_array))], time_avg=True + ) assert uvp.Ntimes == 1 assert uvp.Nblpairs == 1 @@ -589,54 +717,60 @@ def test_get_exact_window_functions(self): uvp = copy.deepcopy(self.uvp_wf) # obtain exact_windows (fiducial usage) - uvp.get_exact_window_functions(ftbeam_file = self.ft_file, - inplace=True) + uvp.get_exact_window_functions(ftbeam_file=self.ft_file, inplace=True) assert uvp.exact_windows assert uvp.window_function_array[0].shape[0] == uvp.Nblpairts # if not exact window function, array dim is 4 assert uvp.window_function_array[0].ndim == 5 - ## tests + # tests # obtain exact window functions for one spw only - uvp.get_exact_window_functions(ftbeam_file = self.ft_file, - spw_array=0, inplace=True, verbose=True) + uvp.get_exact_window_functions( + ftbeam_file=self.ft_file, spw_array=0, inplace=True, verbose=True + ) # raise error if spw not in UVPSpec object - pytest.raises(AssertionError, uvp.get_exact_window_functions, - ftbeam_file = self.ft_file, - spw_array=2, inplace=True) + pytest.raises( + AssertionError, + uvp.get_exact_window_functions, + ftbeam_file=self.ft_file, + spw_array=2, + inplace=True, + ) # output exact_window functions but does not make them attributes - kperp_bins, kpara_bins, wf_array = uvp.get_exact_window_functions(ftbeam_file = self.ft_file, - inplace=False) + kperp_bins, kpara_bins, wf_array = uvp.get_exact_window_functions( + ftbeam_file=self.ft_file, inplace=False + ) # check if result is the same with and without inplace - assert np.all(wf_array[0]==uvp.window_function_array[0]) + assert np.all(wf_array[0] == uvp.window_function_array[0]) def test_fold_spectra(self): uvp = copy.deepcopy(self.uvp) uvp.fold_spectra() - assert(uvp.folded) + assert uvp.folded pytest.raises(AssertionError, uvp.fold_spectra) assert len(uvp.get_dlys(0)) == 14 - assert(np.isclose(uvp.nsample_array[0], 2.0).all()) + assert np.isclose(uvp.nsample_array[0], 2.0).all() # also run the odd case uvd = UVData() uvd_std = UVData() - uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) - uvd_std.read_miriad(os.path.join(DATA_PATH,'zen.even.xx.LST.1.28828.uvOCRSA')) - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, - "HERA_NF_dipole_power.beamfits")) + uvd.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) + uvd_std.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) + beam = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + ) bls = [(37, 38), (38, 39), (52, 53)] - uvp1 = testing.uvpspec_from_data(uvd, bls, data_std=uvd_std, - spw_ranges=[(0,17)], beam=beam) + uvp1 = testing.uvpspec_from_data( + uvd, bls, data_std=uvd_std, spw_ranges=[(0, 17)], beam=beam + ) uvp1.fold_spectra() - cov_folded = uvp1.get_cov((0, ((37, 38), (38, 39)), ('xx','xx'))) - data_folded = uvp1.get_data((0, ((37,38), (38, 39)), ('xx','xx'))) - + # Test fold_spectra method is consistent with average_spectra() - uvp = testing.uvpspec_from_data(uvd, bls, data_std=uvd_std, - spw_ranges=[(0,17)], beam=beam) + uvp = testing.uvpspec_from_data( + uvd, bls, data_std=uvd_std, spw_ranges=[(0, 17)], beam=beam + ) # Average then fold uvp_avg = uvp.average_spectra(time_avg=True, inplace=False) @@ -650,23 +784,26 @@ def test_fold_spectra(self): # Average folded spectra uvp_folded_avg = uvp_folded.average_spectra(time_avg=True, inplace=False) - assert(np.allclose(uvp_avg_folded.get_data((0, ((37, 38), (38, 39)), 'xx')), uvp_folded_avg.get_data((0, ((37, 38), (38, 39)), 'xx')), rtol=1e-5)) + assert np.allclose( + uvp_avg_folded.get_data((0, ((37, 38), (38, 39)), "xx")), + uvp_folded_avg.get_data((0, ((37, 38), (38, 39)), "xx")), + rtol=1e-5, + ) uvp = copy.deepcopy(self.uvp_wf) # obtain exact_windows (fiducial usage) - uvp.get_exact_window_functions(ftbeam_file = self.ft_file, - inplace=True) + uvp.get_exact_window_functions(ftbeam_file=self.ft_file, inplace=True) uvp.fold_spectra() def test_str(self): a = str(self.uvp) - assert(len(a) > 0) + assert len(a) > 0 def test_compute_scalar(self): uvp = copy.deepcopy(self.uvp) # test basic execution - s = uvp.compute_scalar(0, ('xx','xx'), num_steps=1000, noise_scalar=False) - np.testing.assert_almost_equal(s/553995277.90425551, 1.0, decimal=5) + s = uvp.compute_scalar(0, ("xx", "xx"), num_steps=1000, noise_scalar=False) + np.testing.assert_almost_equal(s / 553995277.90425551, 1.0, decimal=5) # test execptions del uvp.OmegaP pytest.raises(AssertionError, uvp.compute_scalar, 0, -5) @@ -682,8 +819,8 @@ def test_set_cosmology(self): # test setting cosmology uvp.set_cosmology(new_cosmo, overwrite=True) assert uvp.cosmo == new_cosmo - assert uvp.norm_units == 'h^-3 Mpc^3' - assert (uvp.scalar_array>1.0).all() + assert uvp.norm_units == "h^-3 Mpc^3" + assert (uvp.scalar_array > 1.0).all() assert (uvp.data_array[0] > 1e5).all() # test exception @@ -695,18 +832,19 @@ def test_set_cosmology(self): # try with new beam uvp.set_cosmology(new_cosmo2, overwrite=True, new_beam=self.beam) assert uvp.cosmo == new_cosmo2 - assert hasattr(uvp, 'OmegaP') + assert hasattr(uvp, "OmegaP") def test_combine_uvpspec(self): # setup uvp build uvd = UVData() - uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, - "HERA_NF_dipole_power.beamfits")) + uvd.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) + beam = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + ) bls = [(37, 38), (38, 39), (52, 53)] - uvp1 = testing.uvpspec_from_data(uvd, bls, - spw_ranges=[(20, 30), (60, 90)], - beam=beam) + uvp1 = testing.uvpspec_from_data( + uvd, bls, spw_ranges=[(20, 30), (60, 90)], beam=beam + ) uvp1 = self._add_optionals(uvp1) # test concat across pol @@ -714,23 +852,27 @@ def test_combine_uvpspec(self): uvp2.polpair_array[0] = 1414 out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) assert out.Npols == 2 - assert(len(set(out.polpair_array) ^ set([1515, 1414])) == 0) - key = (0, ((37, 38), (38, 39)), ('xx','xx')) - assert(np.all(np.isclose(out.get_nsamples(key), - np.ones(10, dtype=np.float64)))) - assert(np.all(np.isclose(out.get_integrations(key), - 190 * np.ones(10, dtype=np.float64), atol=5, rtol=2))) + assert len(set(out.polpair_array) ^ {1515, 1414}) == 0 + key = (0, ((37, 38), (38, 39)), ("xx", "xx")) + assert np.all(np.isclose(out.get_nsamples(key), np.ones(10, dtype=np.float64))) + assert np.all( + np.isclose( + out.get_integrations(key), + 190 * np.ones(10, dtype=np.float64), + atol=5, + rtol=2, + ) + ) # optionals for spw in out.spw_array: ndlys = out.get_spw_ranges(spw)[0][-1] assert out.cov_array_real[spw].shape == (30, ndlys, ndlys, 2) - assert out.stats_array['noise_err'][spw].shape == (30, ndlys, 2) + assert out.stats_array["noise_err"][spw].shape == (30, ndlys, 2) assert out.window_function_array[spw].shape == (30, ndlys, ndlys, 2) - assert out.cov_model == 'empirical' + assert out.cov_model == "empirical" # test concat across spw - uvp2 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(85, 101)], - beam=beam) + uvp2 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(85, 101)], beam=beam) uvp2 = self._add_optionals(uvp2) out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) @@ -739,14 +881,14 @@ def test_combine_uvpspec(self): assert out.Nspwdlys == 56 # optionals - assert len(out.stats_array['noise_err']) == 3 + assert len(out.stats_array["noise_err"]) == 3 assert len(out.window_function_array) == 3 assert len(out.cov_array_real) == 3 # test concat across blpairts - uvp2 = testing.uvpspec_from_data(uvd, [(53, 54), (67, 68)], - spw_ranges=[(20, 30), (60, 90)], - beam=beam) + uvp2 = testing.uvpspec_from_data( + uvd, [(53, 54), (67, 68)], spw_ranges=[(20, 30), (60, 90)], beam=beam + ) uvp2 = self._add_optionals(uvp2) out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) assert out.Nblpairs == 4 @@ -756,18 +898,18 @@ def test_combine_uvpspec(self): for spw in out.spw_array: ndlys = out.get_spw_ranges(spw)[0][-1] assert out.cov_array_real[spw].shape == (40, ndlys, ndlys, 1) - assert out.stats_array['noise_err'][spw].shape == (40, ndlys, 1) + assert out.stats_array["noise_err"][spw].shape == (40, ndlys, 1) assert out.window_function_array[spw].shape == (40, ndlys, ndlys, 1) # test feed as strings uvp1 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(20, 30)], beam=beam) uvp2 = copy.deepcopy(uvp1) uvp2.polpair_array[0] = 1414 - uvp1.write_hdf5('uvp1.hdf5', overwrite=True) - uvp2.write_hdf5('uvp2.hdf5', overwrite=True) - out = uvpspec.combine_uvpspec(['uvp1.hdf5', 'uvp2.hdf5'], verbose=False) + uvp1.write_hdf5("uvp1.hdf5", overwrite=True) + uvp2.write_hdf5("uvp2.hdf5", overwrite=True) + out = uvpspec.combine_uvpspec(["uvp1.hdf5", "uvp2.hdf5"], verbose=False) assert out.Npols == 2 - for ff in ['uvp1.hdf5', 'uvp2.hdf5']: + for ff in ["uvp1.hdf5", "uvp2.hdf5"]: if os.path.exists(ff): os.remove(ff) @@ -780,9 +922,9 @@ def test_combine_uvpspec(self): assert out.Npols == 3 # Test whether n_dlys != Nfreqs works - uvp4 = testing.uvpspec_from_data(uvd, bls, beam=beam, - spw_ranges=[(20, 30), (60, 90)], - n_dlys=[5, 15]) + uvp4 = testing.uvpspec_from_data( + uvd, bls, beam=beam, spw_ranges=[(20, 30), (60, 90)], n_dlys=[5, 15] + ) uvp4b = copy.deepcopy(uvp4) uvp4b.polpair_array[0] = 1414 out = uvpspec.combine_uvpspec([uvp4, uvp4b], verbose=False) @@ -791,29 +933,31 @@ def test_combine_uvpspec(self): uvp_a = copy.deepcopy(uvp1) uvp_b = copy.deepcopy(uvp1) uvp_b.polpair_array[0] = 1414 - uvp_a.history = 'batwing' - uvp_b.history = 'foobar' + uvp_a.history = "batwing" + uvp_b.history = "foobar" # w/ merge out = uvpspec.combine_uvpspec([uvp_a, uvp_b], merge_history=True, verbose=False) - assert 'batwing' in out.history and 'foobar' in out.history + assert "batwing" in out.history and "foobar" in out.history # w/o merge - out = uvpspec.combine_uvpspec([uvp_a, uvp_b], merge_history=False, verbose=False) - assert 'batwing' in out.history and not 'foobar' in out.history + out = uvpspec.combine_uvpspec( + [uvp_a, uvp_b], merge_history=False, verbose=False + ) + assert "batwing" in out.history and "foobar" not in out.history # test no cov_array if cov_model is not consistent uvp_a = copy.deepcopy(uvp1) uvp_b = copy.deepcopy(uvp1) - uvp_b.cov_model = 'foo' + uvp_b.cov_model = "foo" uvp_b.polpair_array = np.array([1414]) out = uvpspec.combine_uvpspec([uvp_a, uvp_b], verbose=False) - assert hasattr(out, 'cov_array_real') is False + assert hasattr(out, "cov_array_real") is False # for exact windows # test basic write execution uvp1 = copy.deepcopy(self.uvp_wf) - uvp1.get_exact_window_functions(ftbeam_file = self.ft_file, inplace=True) + uvp1.get_exact_window_functions(ftbeam_file=self.ft_file, inplace=True) uvp2 = copy.deepcopy(uvp1) uvp2.polpair_array[0] = 1414 out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) @@ -821,13 +965,14 @@ def test_combine_uvpspec(self): def test_combine_uvpspec_errors(self): # setup uvp build uvd = UVData() - uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, - "HERA_NF_dipole_power.beamfits")) + uvd.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) + beam = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + ) bls = [(37, 38), (38, 39), (52, 53)] - uvp1 = testing.uvpspec_from_data(uvd, bls, - spw_ranges=[(20, 30), (60, 90)], - beam=beam) + uvp1 = testing.uvpspec_from_data( + uvd, bls, spw_ranges=[(20, 30), (60, 90)], beam=beam + ) # test failure due to overlapping data uvp2 = copy.deepcopy(uvp1) @@ -840,13 +985,16 @@ def test_combine_uvpspec_errors(self): pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) # test partial data overlap failure - uvp2 = testing.uvpspec_from_data(uvd, [(37, 38), (38, 39), (53, 54)], - spw_ranges=[(20, 30), (60, 90)], - beam=beam) + uvp2 = testing.uvpspec_from_data( + uvd, + [(37, 38), (38, 39), (53, 54)], + spw_ranges=[(20, 30), (60, 90)], + beam=beam, + ) pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) - uvp2 = testing.uvpspec_from_data(uvd, bls, - spw_ranges=[(20, 30), (60, 105)], - beam=beam) + uvp2 = testing.uvpspec_from_data( + uvd, bls, spw_ranges=[(20, 30), (60, 105)], beam=beam + ) pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) uvp2 = copy.deepcopy(uvp1) uvp2.polpair_array[0] = 1414 @@ -854,40 +1002,42 @@ def test_combine_uvpspec_errors(self): pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) # test failure due to variable static metadata - uvp2.weighting = 'foo' + uvp2.weighting = "foo" pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) - uvp2.weighting = 'identity' + uvp2.weighting = "identity" del uvp2.OmegaP del uvp2.OmegaPP pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) - def test_combine_uvpspec_r_params(self): # setup uvp build uvd = UVData() - uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) - beam = pspecbeam.PSpecBeamUV(os.path.join(DATA_PATH, - "HERA_NF_dipole_power.beamfits")) + uvd.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) + beam = pspecbeam.PSpecBeamUV( + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + ) bls = [(37, 38), (38, 39), (52, 53)] - rp = {'filter_centers':[0.], - 'filter_half_widths':[250e-9], - 'filter_factors':[1e-9]} + rp = { + "filter_centers": [0.0], + "filter_half_widths": [250e-9], + "filter_factors": [1e-9], + } r_params = {} for bl in bls: - key1 = bl + ('xx',) + key1 = bl + ("xx",) r_params[key1] = rp # create an r_params copy with inconsistent weighting to test # error case r_params_inconsistent = copy.deepcopy(r_params) - r_params[key1]['filter_half_widths'] = [100e-9] + r_params[key1]["filter_half_widths"] = [100e-9] - uvp1 = testing.uvpspec_from_data(uvd, bls, - spw_ranges=[(20, 30), (60, 90)], - beam=beam, r_params=r_params) + uvp1 = testing.uvpspec_from_data( + uvd, bls, spw_ranges=[(20, 30), (60, 90)], beam=beam, r_params=r_params + ) # test failure due to overlapping data uvp2 = copy.deepcopy(uvp1) @@ -898,12 +1048,12 @@ def test_combine_uvpspec_r_params(self): # test errors when combining with pspecs without r_params uvp3 = copy.deepcopy(uvp2) - uvp3.r_params = '' + uvp3.r_params = "" pytest.raises(ValueError, uvpspec.combine_uvpspec, [uvp1, uvp3]) # combining multiple uvp objects without r_params should run fine uvp4 = copy.deepcopy(uvp1) - uvp4.r_params = '' + uvp4.r_params = "" uvpspec.combine_uvpspec([uvp3, uvp4]) # now test error case with inconsistent weightings. @@ -915,15 +1065,15 @@ def test_combine_uvpspec_std(self): # setup uvp build uvd = UVData() uvd_std = UVData() - uvd.read_miriad(os.path.join(DATA_PATH, 'zen.even.xx.LST.1.28828.uvOCRSA')) - uvd_std.read_miriad( - os.path.join(DATA_PATH,'zen.even.xx.LST.1.28828.uvOCRSA')) + uvd.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) + uvd_std.read_miriad(os.path.join(DATA_PATH, "zen.even.xx.LST.1.28828.uvOCRSA")) beam = pspecbeam.PSpecBeamUV( - os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits")) + os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") + ) bls = [(37, 38), (38, 39), (52, 53)] - uvp1 = testing.uvpspec_from_data(uvd, bls, data_std=uvd_std, - spw_ranges=[(20, 24), (64, 68)], - beam=beam) + uvp1 = testing.uvpspec_from_data( + uvd, bls, data_std=uvd_std, spw_ranges=[(20, 24), (64, 68)], beam=beam + ) # test failure due to overlapping data uvp2 = copy.deepcopy(uvp1) pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) @@ -932,21 +1082,24 @@ def test_combine_uvpspec_std(self): uvp2.polpair_array[0] = 1414 out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) assert out.Npols == 2 - assert len(set(out.polpair_array) ^ set([1515, 1414])) == 0 + assert len(set(out.polpair_array) ^ {1515, 1414}) == 0 # test multiple non-overlapping data axes uvp2.freq_array[0] = 0.0 pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) # test partial data overlap failure - uvp2 = testing.uvpspec_from_data(uvd, [(37, 38), (38, 39), (53, 54)], - data_std=uvd_std, - spw_ranges=[(20, 24), (64, 68)], - beam=beam) + uvp2 = testing.uvpspec_from_data( + uvd, + [(37, 38), (38, 39), (53, 54)], + data_std=uvd_std, + spw_ranges=[(20, 24), (64, 68)], + beam=beam, + ) pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) - uvp2 = testing.uvpspec_from_data(uvd, bls, - spw_ranges=[(20, 24), (64, 68)], - data_std=uvd_std, beam=beam) + uvp2 = testing.uvpspec_from_data( + uvd, bls, spw_ranges=[(20, 24), (64, 68)], data_std=uvd_std, beam=beam + ) pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) uvp2 = copy.deepcopy(uvp1) uvp2.polpair_array[0] = 1414 @@ -954,40 +1107,47 @@ def test_combine_uvpspec_std(self): pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) # test concat across spw - uvp2 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(85, 91)], - data_std=uvd_std, beam=beam) + uvp2 = testing.uvpspec_from_data( + uvd, bls, spw_ranges=[(85, 91)], data_std=uvd_std, beam=beam + ) out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) assert out.Nspws == 3 assert out.Nfreqs == 14 assert out.Nspwdlys == 14 # test concat across blpairts - uvp2 = testing.uvpspec_from_data(uvd, [(53, 54), (67, 68)], - spw_ranges=[(20, 24), (64, 68)], - data_std=uvd_std, beam=beam) + uvp2 = testing.uvpspec_from_data( + uvd, + [(53, 54), (67, 68)], + spw_ranges=[(20, 24), (64, 68)], + data_std=uvd_std, + beam=beam, + ) out = uvpspec.combine_uvpspec([uvp1, uvp2], verbose=False) assert out.Nblpairs == 4 assert out.Nbls == 5 # test failure due to variable static metadata - uvp2.weighting = 'foo' + uvp2.weighting = "foo" pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) - uvp2.weighting = 'identity' + uvp2.weighting = "identity" del uvp2.OmegaP del uvp2.OmegaPP pytest.raises(AssertionError, uvpspec.combine_uvpspec, [uvp1, uvp2]) # test feed as strings - uvp1 = testing.uvpspec_from_data(uvd, bls, spw_ranges=[(20, 30)], - data_std=uvd_std, beam=beam) + uvp1 = testing.uvpspec_from_data( + uvd, bls, spw_ranges=[(20, 30)], data_std=uvd_std, beam=beam + ) uvp2 = copy.deepcopy(uvp1) uvp2.polpair_array[0] = 1414 - uvp1.write_hdf5('uvp1.hdf5', overwrite=True) - uvp2.write_hdf5('uvp2.hdf5', overwrite=True) - out = uvpspec.combine_uvpspec(['uvp1.hdf5', 'uvp2.hdf5'], verbose=False) + uvp1.write_hdf5("uvp1.hdf5", overwrite=True) + uvp2.write_hdf5("uvp2.hdf5", overwrite=True) + out = uvpspec.combine_uvpspec(["uvp1.hdf5", "uvp2.hdf5"], verbose=False) assert out.Npols == 2 - for ff in ['uvp1.hdf5', 'uvp2.hdf5']: - if os.path.exists(ff): os.remove(ff) + for ff in ["uvp1.hdf5", "uvp2.hdf5"]: + if os.path.exists(ff): + os.remove(ff) # test UVPSpec __add__ uvp3 = copy.deepcopy(uvp1) @@ -995,22 +1155,26 @@ def test_combine_uvpspec_std(self): out = uvp1 + uvp2 + uvp3 assert out.Npols == 3 + def test_conj_blpair_int(): conj_blpair = uvputils._conj_blpair_int(101102103104) assert conj_blpair == 103104101102 + def test_conj_bl_int(): conj_bl = uvputils._conj_bl_int(101102) assert conj_bl == 102101 + def test_conj_blpair(): - blpair = uvputils._conj_blpair(101102103104, which='first') + blpair = uvputils._conj_blpair(101102103104, which="first") assert blpair == 102101103104 - blpair = uvputils._conj_blpair(101102103104, which='second') + blpair = uvputils._conj_blpair(101102103104, which="second") assert blpair == 101102104103 - blpair = uvputils._conj_blpair(101102103104, which='both') + blpair = uvputils._conj_blpair(101102103104, which="both") assert blpair == 102101104103 - pytest.raises(ValueError, uvputils._conj_blpair, 102101103104, which='foo') + pytest.raises(ValueError, uvputils._conj_blpair, 102101103104, which="foo") + def test_backwards_compatibility_read(): """This is a backwards compatibility test. @@ -1021,6 +1185,6 @@ def test_backwards_compatibility_read(): """ # test read in of a static test file dated 8/2019 uvp = uvpspec.UVPSpec() - uvp.read_hdf5(os.path.join(DATA_PATH, 'test_uvp.h5')) + uvp.read_hdf5(os.path.join(DATA_PATH, "test_uvp.h5")) # assert check does not fail uvp.check() diff --git a/hera_pspec/tests/test_uvpspec_utils.py b/hera_pspec/tests/test_uvpspec_utils.py index d264ffcc..83e7c7a7 100644 --- a/hera_pspec/tests/test_uvpspec_utils.py +++ b/hera_pspec/tests/test_uvpspec_utils.py @@ -15,66 +15,77 @@ def test_select_common(): Test selecting power spectra that two UVPSpec objects have in common. """ # setup uvp - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") beam = pspecbeam.PSpecBeamUV(beamfile) uvp, cosmo = testing.build_vanilla_uvpspec(beam=beam) # Carve up some example UVPSpec objects - uvp1 = uvp.select(times=np.unique(uvp.time_avg_array)[:-1], - inplace=False) - uvp2 = uvp.select(times=np.unique(uvp.time_avg_array)[1:], - inplace=False) - uvp3 = uvp.select(blpairs=np.unique(uvp.blpair_array)[1:], - inplace=False) - uvp4 = uvp.select(blpairs=np.unique(uvp.blpair_array)[:2], - inplace=False) - uvp5 = uvp.select(blpairs=np.unique(uvp.blpair_array)[:1], - inplace=False) - uvp6 = uvp.select(times=np.unique(uvp.time_avg_array)[:1], - inplace=False) + uvp1 = uvp.select(times=np.unique(uvp.time_avg_array)[:-1], inplace=False) + uvp2 = uvp.select(times=np.unique(uvp.time_avg_array)[1:], inplace=False) + uvp3 = uvp.select(blpairs=np.unique(uvp.blpair_array)[1:], inplace=False) + uvp4 = uvp.select(blpairs=np.unique(uvp.blpair_array)[:2], inplace=False) + uvp5 = uvp.select(blpairs=np.unique(uvp.blpair_array)[:1], inplace=False) + uvp6 = uvp.select(times=np.unique(uvp.time_avg_array)[:1], inplace=False) # Check that selecting on common times works uvp_list = [uvp1, uvp2] - uvp_new = uvputils.select_common(uvp_list, spws=True, blpairs=True, - times=True, polpairs=True, inplace=False) + uvp_new = uvputils.select_common( + uvp_list, spws=True, blpairs=True, times=True, polpairs=True, inplace=False + ) assert uvp_new[0] == uvp_new[1] - np.testing.assert_array_equal(uvp_new[0].time_avg_array, - uvp_new[1].time_avg_array) + np.testing.assert_array_equal(uvp_new[0].time_avg_array, uvp_new[1].time_avg_array) # Check that selecting on common baseline-pairs works uvp_list_2 = [uvp1, uvp2, uvp3] - uvp_new_2 = uvputils.select_common(uvp_list_2, spws=True, blpairs=True, - times=True, polpairs=True, inplace=False) + uvp_new_2 = uvputils.select_common( + uvp_list_2, spws=True, blpairs=True, times=True, polpairs=True, inplace=False + ) assert uvp_new_2[0] == uvp_new_2[1] assert uvp_new_2[0] == uvp_new_2[2] - np.testing.assert_array_equal(uvp_new_2[0].time_avg_array, - uvp_new_2[1].time_avg_array) + np.testing.assert_array_equal( + uvp_new_2[0].time_avg_array, uvp_new_2[1].time_avg_array + ) # Check that zero overlap in times raises a ValueError - pytest.raises(ValueError, uvputils.select_common, [uvp2, uvp6], - spws=True, blpairs=True, times=True, - polpairs=True, inplace=False) + pytest.raises( + ValueError, + uvputils.select_common, + [uvp2, uvp6], + spws=True, + blpairs=True, + times=True, + polpairs=True, + inplace=False, + ) # Check that zero overlap in times does *not* raise a ValueError if # not selecting on times - uvp_new_3 = uvputils.select_common([uvp2, uvp6], spws=True, - blpairs=True, times=False, - polpairs=True, inplace=False) + uvp_new_3 = uvputils.select_common( + [uvp2, uvp6], spws=True, blpairs=True, times=False, polpairs=True, inplace=False + ) # Check that zero overlap in baselines raises a ValueError - pytest.raises(ValueError, uvputils.select_common, [uvp3, uvp5], - spws=True, blpairs=True, times=True, - polpairs=True, inplace=False) + pytest.raises( + ValueError, + uvputils.select_common, + [uvp3, uvp5], + spws=True, + blpairs=True, + times=True, + polpairs=True, + inplace=False, + ) # Check that matching times are ignored when set to False - uvp_new = uvputils.select_common(uvp_list, spws=True, blpairs=True, - times=False, polpairs=True, inplace=False) - assert np.sum(uvp_new[0].time_avg_array - - uvp_new[1].time_avg_array) != 0. + uvp_new = uvputils.select_common( + uvp_list, spws=True, blpairs=True, times=False, polpairs=True, inplace=False + ) + assert np.sum(uvp_new[0].time_avg_array - uvp_new[1].time_avg_array) != 0.0 assert len(uvp_new) == len(uvp_list) # Check that in-place selection works - uvputils.select_common(uvp_list, spws=True, blpairs=True, - times=True, polpairs=True, inplace=True) + uvputils.select_common( + uvp_list, spws=True, blpairs=True, times=True, polpairs=True, inplace=True + ) assert uvp1 == uvp2 # check uvplist > 2 @@ -92,22 +103,22 @@ def test_select_common(): # check pol overlap uvp7 = copy.deepcopy(uvp1) - uvp7.polpair_array[0] = 1212 # = (-8,-8) - pytest.raises(ValueError, uvputils.select_common, [uvp1, uvp7], - polpairs=True) + uvp7.polpair_array[0] = 1212 # = (-8,-8) + pytest.raises(ValueError, uvputils.select_common, [uvp1, uvp7], polpairs=True) + def test_get_blpairs_from_bls(): """ Test conversion of bls to set of blpairs. """ # setup uvp - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") beam = pspecbeam.PSpecBeamUV(beamfile) uvp, cosmo = testing.build_vanilla_uvpspec(beam=beam) # Check that bls can be specified in several different ways blps = uvputils._get_blpairs_from_bls(uvp, bls=101102) - blps = uvputils._get_blpairs_from_bls(uvp, bls=(101,102)) + blps = uvputils._get_blpairs_from_bls(uvp, bls=(101, 102)) blps = uvputils._get_blpairs_from_bls(uvp, bls=[101102, 101103]) @@ -116,16 +127,16 @@ def test_get_red_bls(): Test retrieval of redundant baseline groups. """ # Setup uvp - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") beam = pspecbeam.PSpecBeamUV(beamfile) uvp, cosmo = testing.build_vanilla_uvpspec(beam=beam) # Get redundant baseline groups bls, lens, angs = uvp.get_red_bls() - assert len(bls) == 3 # three red grps in this file - assert len(bls) == len(lens) # Should be one length for each group - assert len(bls) == len(angs) # Ditto, for angles + assert len(bls) == 3 # three red grps in this file + assert len(bls) == len(lens) # Should be one length for each group + assert len(bls) == len(angs) # Ditto, for angles # Check that number of grouped baselines = total no. of baselines num_bls = 0 @@ -140,16 +151,16 @@ def test_get_red_blpairs(): Test retrieval of redundant baseline groups for baseline-pairs. """ # Setup uvp - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") beam = pspecbeam.PSpecBeamUV(beamfile) uvp, cosmo = testing.build_vanilla_uvpspec(beam=beam) # Get redundant baseline groups blps, lens, angs = uvp.get_red_blpairs() - assert len(blps) == 3 # three red grps in this file - assert len(blps) == len(lens) # Should be one length for each group - assert len(blps) == len(angs) # Ditto, for angles + assert len(blps) == 3 # three red grps in this file + assert len(blps) == len(lens) # Should be one length for each group + assert len(blps) == len(angs) # Ditto, for angles # Check output type assert isinstance(blps[0][0], (np.int, int)) @@ -167,16 +178,23 @@ def test_polpair_int2tuple(): Test conversion of polpair ints to tuples. """ # List of polpairs to test - polpairs = [('xx','xx'), ('xx','yy'), ('xy', 'yx'), - ('pI','pI'), ('pI','pQ'), ('pQ','pQ'), ('pU','pU'), - ('pV','pV') ] + polpairs = [ + ("xx", "xx"), + ("xx", "yy"), + ("xy", "yx"), + ("pI", "pI"), + ("pI", "pQ"), + ("pQ", "pQ"), + ("pU", "pU"), + ("pV", "pV"), + ] # Check that lists and single items work pol_ints = uvputils.polpair_tuple2int(polpairs) uvputils.polpair_tuple2int(polpairs[0]) uvputils.polpair_int2tuple(1515) - uvputils.polpair_int2tuple([1515,1414]) - uvputils.polpair_int2tuple(np.array([1515,1414])) + uvputils.polpair_int2tuple([1515, 1414]) + uvputils.polpair_int2tuple(np.array([1515, 1414])) # Test converting to int and then back again pol_pairs_returned = uvputils.polpair_int2tuple(pol_ints, pol_strings=True) @@ -184,11 +202,17 @@ def test_polpair_int2tuple(): assert polpairs[i] == pol_pairs_returned[i] # Check that errors are raised appropriately - pytest.raises(AssertionError, uvputils.polpair_int2tuple, ('xx','xx')) - pytest.raises(AssertionError, uvputils.polpair_int2tuple, 'xx') - pytest.raises(AssertionError, uvputils.polpair_int2tuple, 'pI') + pytest.raises(AssertionError, uvputils.polpair_int2tuple, ("xx", "xx")) + pytest.raises(AssertionError, uvputils.polpair_int2tuple, "xx") + pytest.raises(AssertionError, uvputils.polpair_int2tuple, "pI") pytest.raises(ValueError, uvputils.polpair_int2tuple, 999) - pytest.raises(ValueError, uvputils.polpair_int2tuple, [999,]) + pytest.raises( + ValueError, + uvputils.polpair_int2tuple, + [ + 999, + ], + ) def test_subtract_uvp(): @@ -196,13 +220,13 @@ def test_subtract_uvp(): Test subtraction of two UVPSpec objects """ # setup uvp - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") beam = pspecbeam.PSpecBeamUV(beamfile) uvp, cosmo = testing.build_vanilla_uvpspec(beam=beam) # add a dummy stats_array for k in uvp.get_all_keys(): - uvp.set_stats('mystat', k, np.ones((10, 30), dtype=np.complex)) + uvp.set_stats("mystat", k, np.ones((10, 30), dtype=np.complex)) # test execution uvs = uvputils.subtract_uvp(uvp, uvp, run_check=True) @@ -214,7 +238,7 @@ def test_subtract_uvp(): assert np.isclose(uvs.data_array[0], 0.0).all() # check stats_array is np.sqrt(2) - assert np.isclose(uvs.stats_array['mystat'][0], np.sqrt(2)).all() + assert np.isclose(uvs.stats_array["mystat"][0], np.sqrt(2)).all() def test_conj_blpair_int(): @@ -228,22 +252,44 @@ def test_conj_bl_int(): def test_conj_blpair(): - blpair = uvputils._conj_blpair(101102103104, which='first') + blpair = uvputils._conj_blpair(101102103104, which="first") assert blpair == 102101103104 - blpair = uvputils._conj_blpair(101102103104, which='second') + blpair = uvputils._conj_blpair(101102103104, which="second") assert blpair == 101102104103 - blpair = uvputils._conj_blpair(101102103104, which='both') + blpair = uvputils._conj_blpair(101102103104, which="both") assert blpair == 102101104103 - pytest.raises(ValueError, uvputils._conj_blpair, 102101103104, which='foo') + pytest.raises(ValueError, uvputils._conj_blpair, 102101103104, which="foo") def test_fast_is_in(): - blps = [ 102101103104, 102101103104, 102101103104, 102101103104, - 101102104103, 101102104103, 101102104103, 101102104103, - 102101104103, 102101104103, 102101104103, 102101104103 ] - times = [ 0.1, 0.15, 0.2, 0.25, - 0.1, 0.15, 0.2, 0.25, - 0.1, 0.15, 0.3, 0.3, ] + blps = [ + 102101103104, + 102101103104, + 102101103104, + 102101103104, + 101102104103, + 101102104103, + 101102104103, + 101102104103, + 102101104103, + 102101104103, + 102101104103, + 102101104103, + ] + times = [ + 0.1, + 0.15, + 0.2, + 0.25, + 0.1, + 0.15, + 0.2, + 0.25, + 0.1, + 0.15, + 0.3, + 0.3, + ] src_blpts = np.array(list(zip(blps, times))) assert uvputils._fast_is_in(src_blpts, [(101102104103, 0.2)])[0] @@ -251,12 +297,34 @@ def test_fast_is_in(): def test_fast_lookup_blpairts(): # Construct array of blpair-time tuples (including some out of order) - blps = [ 102101103104, 102101103104, 102101103104, 102101103104, - 101102104103, 101102104103, 101102104103, 101102104103, - 102101104103, 102101104103, 102101104103, 102101104103 ] - times = [ 0.1, 0.15, 0.2, 0.25, - 0.1, 0.15, 0.2, 0.25, - 0.1, 0.15, 0.3, 0.3, ] + blps = [ + 102101103104, + 102101103104, + 102101103104, + 102101103104, + 101102104103, + 101102104103, + 101102104103, + 101102104103, + 102101104103, + 102101104103, + 102101104103, + 102101104103, + ] + times = [ + 0.1, + 0.15, + 0.2, + 0.25, + 0.1, + 0.15, + 0.2, + 0.25, + 0.1, + 0.15, + 0.3, + 0.3, + ] src_blpts = np.array(list(zip(blps, times))) # List of blpair-times to look up @@ -266,18 +334,21 @@ def test_fast_lookup_blpairts(): idxs = uvputils._fast_lookup_blpairts(src_blpts, np.array(query_blpts)) np.testing.assert_array_equal(idxs, np.array([0, 4, 7])) + def test_r_param_compression(): - baselines = [(24,25), (37,38), (38,39)] + baselines = [(24, 25), (37, 38), (38, 39)] - rp = {'filter_centers':[0.], - 'filter_half_widths':[250e-9], - 'filter_factors':[1e-9]} + rp = { + "filter_centers": [0.0], + "filter_half_widths": [250e-9], + "filter_factors": [1e-9], + } r_params = {} for bl in baselines: - key1 = bl + ('xx',) - key2 = bl + ('xx',) + key1 = bl + ("xx",) + key2 = bl + ("xx",) r_params[key1] = rp r_params[key2] = rp @@ -293,5 +364,5 @@ def test_r_param_compression(): rp_str_2 = uvputils.compress_r_params(rp) assert json.loads(rp_str_1) == json.loads(rp_str_2) - assert uvputils.compress_r_params({}) == '' - assert uvputils.decompress_r_params('') == {} + assert uvputils.compress_r_params({}) == "" + assert uvputils.decompress_r_params("") == {} diff --git a/hera_pspec/tests/test_uvwindow.py b/hera_pspec/tests/test_uvwindow.py index f0d68d4e..348681b6 100644 --- a/hera_pspec/tests/test_uvwindow.py +++ b/hera_pspec/tests/test_uvwindow.py @@ -1,43 +1,35 @@ import unittest import pytest -from pyuvdata import utils as uvutils -import uvtools.dspec as dspec -import h5py -import warnings import numpy as np -import sys import os -import time import copy from astropy import units -from scipy.interpolate import interp2d -from pyuvdata import UVBeam, UVData +from pyuvdata import UVData from hera_pspec.data import DATA_PATH -from .. import conversions, noise, version, pspecbeam, grouping, utils +from .. import conversions, pspecbeam, utils from .. import uvwindow, pspecbeam, PSpecData -from .. import uvpspec_utils as uvputils # Data files to use in tests -dfile = 'zen.2458116.31939.HH.uvh5' -ftfile = 'FT_beam_HERA_dipole_test_xx.hdf5' -basename = 'FT_beam_HERA_dipole_test' -outfile = 'test.hdf5' +dfile = "zen.2458116.31939.HH.uvh5" +ftfile = "FT_beam_HERA_dipole_test_xx.hdf5" +basename = "FT_beam_HERA_dipole_test" +outfile = "test.hdf5" class test_FTBeam(unittest.TestCase): - def setUp(self): # parameters self.ft_file = os.path.join(DATA_PATH, ftfile) - self.pol = 'xx' + self.pol = "xx" self.spw_range = (5, 25) self.verbose = False - self.x_orientation = 'east' + self.x_orientation = "east" - self.ftbeam_obj = uvwindow.FTBeam.from_file(ftfile=self.ft_file, - spw_range=self.spw_range) + self.ftbeam_obj = uvwindow.FTBeam.from_file( + ftfile=self.ft_file, spw_range=self.spw_range + ) self.data = self.ftbeam_obj.ft_beam self.mapsize = self.ftbeam_obj.mapsize self.freq_array = self.ftbeam_obj.freq_array @@ -52,144 +44,185 @@ def runTest(self): def test_init(self): # initialise directly with array - test = uvwindow.FTBeam(data=self.data, - pol=self.pol, - freq_array=self.freq_array, - mapsize=self.mapsize, - verbose=self.verbose, - x_orientation=self.x_orientation) + test = uvwindow.FTBeam( + data=self.data, + pol=self.pol, + freq_array=self.freq_array, + mapsize=self.mapsize, + verbose=self.verbose, + x_orientation=self.x_orientation, + ) assert self.pol == test.pol assert np.all(self.data == test.ft_beam) # raise assertion error if data is not dim 3 - pytest.raises(AssertionError, uvwindow.FTBeam, - data=self.data[:, :, 0], pol=self.pol, - freq_array=self.freq_array, mapsize=self.mapsize) + pytest.raises( + AssertionError, + uvwindow.FTBeam, + data=self.data[:, :, 0], + pol=self.pol, + freq_array=self.freq_array, + mapsize=self.mapsize, + ) # raise assertion error if data has wrong shape - pytest.raises(AssertionError, uvwindow.FTBeam, - data=self.data[:, :, :-1], pol=self.pol, - freq_array=self.freq_array, mapsize=self.mapsize) + pytest.raises( + AssertionError, + uvwindow.FTBeam, + data=self.data[:, :, :-1], + pol=self.pol, + freq_array=self.freq_array, + mapsize=self.mapsize, + ) # raise assertion error if freq_array and data not compatible - pytest.raises(AssertionError, uvwindow.FTBeam, - data=self.data[:12, :, :], pol=self.pol, - freq_array=self.freq_array, mapsize=self.mapsize) + pytest.raises( + AssertionError, + uvwindow.FTBeam, + data=self.data[:12, :, :], + pol=self.pol, + freq_array=self.freq_array, + mapsize=self.mapsize, + ) # tests related to pol - test = uvwindow.FTBeam(data=self.data, - pol=-5, - freq_array=self.freq_array, - mapsize=self.mapsize) - pytest.raises(AssertionError, uvwindow.FTBeam, pol='test', - data=self.data, freq_array=self.freq_array, - mapsize=self.mapsize) - pytest.raises(AssertionError, uvwindow.FTBeam, pol=12, - data=self.data, freq_array=self.freq_array, - mapsize=self.mapsize) - pytest.raises(TypeError, uvwindow.FTBeam, pol=3.4, - data=self.data, freq_array=self.freq_array, - mapsize=self.mapsize) + test = uvwindow.FTBeam( + data=self.data, pol=-5, freq_array=self.freq_array, mapsize=self.mapsize + ) + pytest.raises( + AssertionError, + uvwindow.FTBeam, + pol="test", + data=self.data, + freq_array=self.freq_array, + mapsize=self.mapsize, + ) + pytest.raises( + AssertionError, + uvwindow.FTBeam, + pol=12, + data=self.data, + freq_array=self.freq_array, + mapsize=self.mapsize, + ) + pytest.raises( + TypeError, + uvwindow.FTBeam, + pol=3.4, + data=self.data, + freq_array=self.freq_array, + mapsize=self.mapsize, + ) def test_from_beam(self): - pytest.raises(NotImplementedError, uvwindow.FTBeam.from_beam, - beamfile='test') + pytest.raises(NotImplementedError, uvwindow.FTBeam.from_beam, beamfile="test") def test_from_file(self): - test = uvwindow.FTBeam.from_file(ftfile=self.ft_file, - spw_range=self.spw_range, - verbose=self.verbose, - x_orientation=self.x_orientation) + test = uvwindow.FTBeam.from_file( + ftfile=self.ft_file, + spw_range=self.spw_range, + verbose=self.verbose, + x_orientation=self.x_orientation, + ) assert test.pol == self.pol # tests related to ftfile - pytest.raises(TypeError, uvwindow.FTBeam.from_file, ftfile=12.) + pytest.raises(TypeError, uvwindow.FTBeam.from_file, ftfile=12.0) # if ft file does not exist, raise assertion error - pytest.raises(ValueError, uvwindow.FTBeam.from_file, ftfile='whatever') + pytest.raises(ValueError, uvwindow.FTBeam.from_file, ftfile="whatever") # tests related to spw_range - test1 = uvwindow.FTBeam.from_file(ftfile=self.ft_file, - spw_range=self.spw_range) + test1 = uvwindow.FTBeam.from_file(ftfile=self.ft_file, spw_range=self.spw_range) assert np.all(test1.freq_array == self.freq_array) - test2 = uvwindow.FTBeam.from_file(ftfile=self.ft_file, - spw_range=None) + test2 = uvwindow.FTBeam.from_file(ftfile=self.ft_file, spw_range=None) assert np.all(test2.freq_array == self.bandwidth) - pytest.raises(AssertionError, uvwindow.FTBeam.from_file, spw_range=(13), - ftfile=self.ft_file) - pytest.raises(AssertionError, uvwindow.FTBeam.from_file, spw_range=(20, 10), - ftfile=self.ft_file) - pytest.raises(AssertionError, uvwindow.FTBeam.from_file, spw_range=(1001, 1022), - ftfile=self.ft_file) + pytest.raises( + AssertionError, + uvwindow.FTBeam.from_file, + spw_range=(13), + ftfile=self.ft_file, + ) + pytest.raises( + AssertionError, + uvwindow.FTBeam.from_file, + spw_range=(20, 10), + ftfile=self.ft_file, + ) + pytest.raises( + AssertionError, + uvwindow.FTBeam.from_file, + spw_range=(1001, 1022), + ftfile=self.ft_file, + ) def test_get_bandwidth(self): test_bandwidth = uvwindow.FTBeam.get_bandwidth(self.ft_file) assert np.all(test_bandwidth == self.bandwidth) # raise error is ft_file does not exist - pytest.raises(ValueError, uvwindow.FTBeam.get_bandwidth, - ftfile='whatever') + pytest.raises(ValueError, uvwindow.FTBeam.get_bandwidth, ftfile="whatever") def test_update_spw(self): # proper usage - test = uvwindow.FTBeam.from_file(ftfile=self.ft_file, - spw_range=None) + test = uvwindow.FTBeam.from_file(ftfile=self.ft_file, spw_range=None) test.update_spw(self.spw_range) - # tests related to spw_range - test = uvwindow.FTBeam.from_file(ftfile=self.ft_file, - spw_range=None) + # tests related to spw_range + test = uvwindow.FTBeam.from_file(ftfile=self.ft_file, spw_range=None) pytest.raises(AssertionError, test.update_spw, spw_range=(13)) pytest.raises(AssertionError, test.update_spw, spw_range=(20, 10)) pytest.raises(AssertionError, test.update_spw, spw_range=(1001, 1022)) class Test_UVWindow(unittest.TestCase): - def setUp(self): # Instantiate UVWindow() self.ft_file = os.path.join(DATA_PATH, ftfile) - self.pol = 'xx' + self.pol = "xx" self.polpair = (self.pol, self.pol) self.spw_range = (5, 25) - self.taper = 'blackman-harris' + self.taper = "blackman-harris" self.verbose = False self.little_h = True self.cosmo = conversions.Cosmo_Conversions() - self.ft_beam_obj_spw = uvwindow.FTBeam.from_file(ftfile=self.ft_file, - spw_range=self.spw_range) + self.ft_beam_obj_spw = uvwindow.FTBeam.from_file( + ftfile=self.ft_file, spw_range=self.spw_range + ) self.ft_beam_obj = uvwindow.FTBeam.from_file(ftfile=self.ft_file) - self.uvw = uvwindow.UVWindow(ftbeam_obj = self.ft_beam_obj_spw, - taper=self.taper, - cosmo=self.cosmo, - little_h=self.little_h, - verbose=self.verbose) + self.uvw = uvwindow.UVWindow( + ftbeam_obj=self.ft_beam_obj_spw, + taper=self.taper, + cosmo=self.cosmo, + little_h=self.little_h, + verbose=self.verbose, + ) # set parameters self.freq_array = self.uvw.freq_array self.ngrid = self.ft_beam_obj.ft_beam.shape[-1] # HERA bandwidth - self.HERA_bw = np.linspace(1, 2, 1024, endpoint=False)*1e8 + self.HERA_bw = np.linspace(1, 2, 1024, endpoint=False) * 1e8 # define spherical kbins - kmax, dk = 1., 0.128/2 - krange = np.arange(dk*1.5, kmax, step=dk) - nbinsk = krange.size - 1 - kbins = (krange[1:]+krange[:-1])/2 + kmax, dk = 1.0, 0.128 / 2 + krange = np.arange(dk * 1.5, kmax, step=dk) + kbins = (krange[1:] + krange[:-1]) / 2 self.kbins = kbins * units.h / units.Mpc # Load datafile uvd = UVData() uvd.read(os.path.join(DATA_PATH, dfile), read_data=False) - self.reds, self.lens, _ = utils.get_reds(uvd, bl_error_tol=1.0, - pick_data_ants=False) + self.reds, self.lens, _ = utils.get_reds( + uvd, bl_error_tol=1.0, pick_data_ants=False + ) def tearDown(self): pass @@ -203,43 +236,49 @@ def test_init(self): test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw) # raise error if two ftbeam_obj are not consistent - pytest.raises(AssertionError, uvwindow.UVWindow, - ftbeam_obj=(self.ft_beam_obj_spw, self.ft_beam_obj)) + pytest.raises( + AssertionError, + uvwindow.UVWindow, + ftbeam_obj=(self.ft_beam_obj_spw, self.ft_beam_obj), + ) ftbeam_test = copy.deepcopy(self.ft_beam_obj_spw) - ftbeam_test.mapsize = 2. - pytest.raises(AssertionError, uvwindow.UVWindow, - ftbeam_obj=(self.ft_beam_obj_spw, ftbeam_test)) + ftbeam_test.mapsize = 2.0 + pytest.raises( + AssertionError, + uvwindow.UVWindow, + ftbeam_obj=(self.ft_beam_obj_spw, ftbeam_test), + ) # raise error if ftbeam_obj is wrong input - pytest.raises(AssertionError, uvwindow.UVWindow, ftbeam_obj='test') + pytest.raises(AssertionError, uvwindow.UVWindow, ftbeam_obj="test") # test taper options - test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, - taper=self.taper) + test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, taper=self.taper) assert test.taper == self.taper - test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, - taper=None) + test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, taper=None) assert test.taper is None - pytest.raises(ValueError, uvwindow.UVWindow, taper='test', - ftbeam_obj=self.ft_beam_obj_spw) + pytest.raises( + ValueError, uvwindow.UVWindow, taper="test", ftbeam_obj=self.ft_beam_obj_spw + ) # test on cosmo cosmo = conversions.Cosmo_Conversions() - test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, - cosmo=cosmo) - pytest.raises(AssertionError, uvwindow.UVWindow, cosmo=None, - ftbeam_obj=self.ft_beam_obj_spw) + test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, cosmo=cosmo) + pytest.raises( + AssertionError, + uvwindow.UVWindow, + cosmo=None, + ftbeam_obj=self.ft_beam_obj_spw, + ) # test on verbose - test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, - verbose=True) + test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, verbose=True) assert test.verbose # test on little_h - test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, - verbose=True) + test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, verbose=True) assert test.kunits.is_equivalent(units.h / units.Mpc) test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw, little_h=False) - assert test.kunits.is_equivalent(units.Mpc**(-1)) + assert test.kunits.is_equivalent(units.Mpc ** (-1)) def test_from_uvpspec(self): @@ -247,8 +286,8 @@ def test_from_uvpspec(self): datafile = os.path.join(DATA_PATH, dfile) uvd = UVData() uvd.read_uvh5(datafile) - # beam - beamfile = os.path.join(DATA_PATH, 'HERA_NF_dipole_power.beamfits') + # beam + beamfile = os.path.join(DATA_PATH, "HERA_NF_dipole_power.beamfits") uvb = pspecbeam.PSpecBeamUV(beamfile, cosmo=None) Jy_to_mK = uvb.Jy_to_mK(np.unique(uvd.freq_array), pol=self.pol) # reshape to appropriately match a UVData.data_array object and multiply in! @@ -257,34 +296,75 @@ def test_from_uvpspec(self): ds = PSpecData(dsets=[uvd, uvd], wgts=[None, None], beam=uvb) ds_nocosmo = PSpecData(dsets=[uvd, uvd], wgts=[None, None]) # choose baselines - baselines1, baselines2, blpairs = utils.construct_blpairs(uvd.get_antpairs()[1:], - exclude_permutations=False, - exclude_auto_bls=True) + baselines1, baselines2, blpairs = utils.construct_blpairs( + uvd.get_antpairs()[1:], exclude_permutations=False, exclude_auto_bls=True + ) # compute ps - uvp = ds.pspec(baselines1, baselines2, dsets=(0, 1), pols=[self.polpair], - spw_ranges=(175,195), taper=self.taper,verbose=self.verbose) - uvp_nocosmo = ds_nocosmo.pspec(baselines1, baselines2, dsets=(0, 1), pols=[self.polpair], - spw_ranges=self.spw_range, taper=self.taper,verbose=self.verbose) - uvp_crosspol = ds.pspec(baselines1, baselines2, dsets=(0, 1), pols=['xx', 'yy'], - spw_ranges=(175,195), taper=self.taper,verbose=self.verbose) + uvp = ds.pspec( + baselines1, + baselines2, + dsets=(0, 1), + pols=[self.polpair], + spw_ranges=(175, 195), + taper=self.taper, + verbose=self.verbose, + ) + uvp_nocosmo = ds_nocosmo.pspec( + baselines1, + baselines2, + dsets=(0, 1), + pols=[self.polpair], + spw_ranges=self.spw_range, + taper=self.taper, + verbose=self.verbose, + ) + uvp_crosspol = ds.pspec( + baselines1, + baselines2, + dsets=(0, 1), + pols=["xx", "yy"], + spw_ranges=(175, 195), + taper=self.taper, + verbose=self.verbose, + ) # proper usage # initialise with ftfile (read pre-computed FT of the beam in file) - uvw_ps = uvwindow.UVWindow.from_uvpspec(uvp, ipol=0, spw=0, verbose=True, - ftfile=os.path.join(DATA_PATH, basename)) + uvw_ps = uvwindow.UVWindow.from_uvpspec( + uvp, ipol=0, spw=0, verbose=True, ftfile=os.path.join(DATA_PATH, basename) + ) # if cross polarisation - test = uvwindow.UVWindow.from_uvpspec(uvp_crosspol, ipol=0, spw=0, - ftfile=os.path.join(DATA_PATH, basename)) + uvwindow.UVWindow.from_uvpspec( + uvp_crosspol, ipol=0, spw=0, ftfile=os.path.join(DATA_PATH, basename) + ) # if no cosmo, use default - uvw_ps = uvwindow.UVWindow.from_uvpspec(uvp_nocosmo, ipol=0, spw=0, verbose=True, - ftfile=os.path.join(DATA_PATH, basename)) + uvwindow.UVWindow.from_uvpspec( + uvp_nocosmo, + ipol=0, + spw=0, + verbose=True, + ftfile=os.path.join(DATA_PATH, basename), + ) # raise error if no ftfile as option is not implemented yet - pytest.raises(NotImplementedError, uvwindow.UVWindow.from_uvpspec, uvp=uvp, - ipol=0, spw=0, ftfile=None, verbose=False) + pytest.raises( + NotImplementedError, + uvwindow.UVWindow.from_uvpspec, + uvp=uvp, + ipol=0, + spw=0, + ftfile=None, + verbose=False, + ) # raise error if spw not within uvp.Nspws - pytest.raises(AssertionError, uvwindow.UVWindow.from_uvpspec, uvp=uvp_nocosmo, - ipol=0, spw=2, ftfile=os.path.join(DATA_PATH, basename)) + pytest.raises( + AssertionError, + uvwindow.UVWindow.from_uvpspec, + uvp=uvp_nocosmo, + ipol=0, + spw=2, + ftfile=os.path.join(DATA_PATH, basename), + ) def test_get_kgrid(self): @@ -293,8 +373,7 @@ def test_get_kgrid(self): # initialise object test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw) kgrid, kperp_norm = test._get_kgrid(bl_len) - pytest.raises(AssertionError, test._get_kgrid, - bl_len=bl_len, width=0.0004) + pytest.raises(AssertionError, test._get_kgrid, bl_len=bl_len, width=0.0004) def test_kperp4bl_freq(self): @@ -303,17 +382,25 @@ def test_kperp4bl_freq(self): # initialise object test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw) # test for correct input parameters - k = test._kperp4bl_freq(freq=test.freq_array[12], - bl_len=bl_len, - ngrid=self.ngrid) + test._kperp4bl_freq( + freq=test.freq_array[12], bl_len=bl_len, ngrid=self.ngrid + ) # test for frequency outside of spectral window - pytest.raises(AssertionError, test._kperp4bl_freq, - freq=1.35*1e8, bl_len=bl_len, ngrid=self.ngrid) + pytest.raises( + AssertionError, + test._kperp4bl_freq, + freq=1.35 * 1e8, + bl_len=bl_len, + ngrid=self.ngrid, + ) # test for frequency in Hhz - pytest.raises(AssertionError, test._kperp4bl_freq, - freq=test.freq_array[12]/1e6, - bl_len=bl_len, - ngrid=self.ngrid) + pytest.raises( + AssertionError, + test._kperp4bl_freq, + freq=test.freq_array[12] / 1e6, + bl_len=bl_len, + ngrid=self.ngrid, + ) def test_interpolate_ft_beam(self): @@ -325,12 +412,24 @@ def test_interpolate_ft_beam(self): interp_ft_beam, kperp_norm = test._interpolate_ft_beam(bl_len, ft_beam) # test for ft_beam of wrong dimensions - pytest.raises(AssertionError, test._interpolate_ft_beam, - bl_len=bl_len, ft_beam=ft_beam[0, :, :]) - pytest.raises(AssertionError, test._interpolate_ft_beam, - bl_len=bl_len, ft_beam=ft_beam[0:10, :, :]) - pytest.raises(AssertionError, test._interpolate_ft_beam, - bl_len=bl_len, ft_beam=ft_beam[:, :, :].T) + pytest.raises( + AssertionError, + test._interpolate_ft_beam, + bl_len=bl_len, + ft_beam=ft_beam[0, :, :], + ) + pytest.raises( + AssertionError, + test._interpolate_ft_beam, + bl_len=bl_len, + ft_beam=ft_beam[0:10, :, :], + ) + pytest.raises( + AssertionError, + test._interpolate_ft_beam, + bl_len=bl_len, + ft_beam=ft_beam[:, :, :].T, + ) def test_take_freq_FT(self): @@ -341,13 +440,15 @@ def test_take_freq_FT(self): ft_beam = np.copy(test.ftbeam_obj_pol[0].ft_beam) interp_ft_beam, kperp_norm = test._interpolate_ft_beam(bl_len, ft_beam) # frequency resolution - delta_nu = abs(test.freq_array[-1]-test.freq_array[0])/test.Nfreqs - fnu = test._take_freq_FT(interp_ft_beam, delta_nu) + delta_nu = abs(test.freq_array[-1] - test.freq_array[0]) / test.Nfreqs + test._take_freq_FT(interp_ft_beam, delta_nu) # test for ft_beam of wrong dimensions - pytest.raises(AssertionError, test._take_freq_FT, - interp_ft_beam[0, :, :], delta_nu) - pytest.raises(AssertionError, test._take_freq_FT, - interp_ft_beam[:, :, :].T, delta_nu) + pytest.raises( + AssertionError, test._take_freq_FT, interp_ft_beam[0, :, :], delta_nu + ) + pytest.raises( + AssertionError, test._take_freq_FT, interp_ft_beam[:, :, :].T, delta_nu + ) def test_get_wf_for_tau(self): @@ -361,8 +462,7 @@ def test_get_wf_for_tau(self): kpara_bins = np.array(kpara_bins.value) wf_array1 = np.zeros((kperp_bins.size, test.Nfreqs)) - kpara, cyl_wf = test._get_wf_for_tau(tau, wf_array1, - kperp_bins, kpara_bins) + kpara, cyl_wf = test._get_wf_for_tau(tau, wf_array1, kperp_bins, kpara_bins) def test_get_kperp_bins(self): @@ -376,15 +476,16 @@ def test_get_kperp_bins(self): # test for array of baseline lengths _ = test.get_kperp_bins(self.lens) # test for warning if large number of bins (> 200) - _ = test.get_kperp_bins(np.r_[1., self.lens]) + _ = test.get_kperp_bins(np.r_[1.0, self.lens]) def test_get_kpara_bins(self): # initialise object test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw) # raise error if empty freq array or length 1 - pytest.raises(AssertionError, test.get_kpara_bins, - freq_array=self.freq_array[2]) + pytest.raises( + AssertionError, test.get_kpara_bins, freq_array=self.freq_array[2] + ) # test for correct input _ = test.get_kpara_bins(self.freq_array) # test for warning if large number of bins (> 200) @@ -401,20 +502,17 @@ def test_get_cylindrical_wf(self): # initialise object test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw) - _, _, cyl_wf = test.get_cylindrical_wf(bl_len, - kperp_bins=None, - kpara_bins=None, - return_bins='weighted') - cyl_wf = test.get_cylindrical_wf(bl_len, - kperp_bins=None, - kpara_bins=None, - return_bins=None) - kperp, kpara, cyl_wf = test.get_cylindrical_wf(bl_len, - kperp_bins=None, - kpara_bins=None, - return_bins='unweighted') + _, _, cyl_wf = test.get_cylindrical_wf( + bl_len, kperp_bins=None, kpara_bins=None, return_bins="weighted" + ) + cyl_wf = test.get_cylindrical_wf( + bl_len, kperp_bins=None, kpara_bins=None, return_bins=None + ) + kperp, kpara, cyl_wf = test.get_cylindrical_wf( + bl_len, kperp_bins=None, kpara_bins=None, return_bins="unweighted" + ) # check normalisation - assert np.all(np.isclose(np.sum(cyl_wf, axis=(1, 2)), 1., atol=1e-3)) + assert np.all(np.isclose(np.sum(cyl_wf, axis=(1, 2)), 1.0, atol=1e-3)) assert kperp.size == cyl_wf.shape[1] assert kpara.size == cyl_wf.shape[2] assert test.Nfreqs == cyl_wf.shape[0] @@ -425,40 +523,47 @@ def test_get_cylindrical_wf(self): # test different key words # kperp bins - kperp2, _, cyl_wf2 = test.get_cylindrical_wf(bl_len, - kperp_bins=kperp*test.kunits, - kpara_bins=None, - return_bins='unweighted') + kperp2, _, cyl_wf2 = test.get_cylindrical_wf( + bl_len, + kperp_bins=kperp * test.kunits, + kpara_bins=None, + return_bins="unweighted", + ) assert np.all(cyl_wf2 == cyl_wf) assert np.all(kperp2 == kperp) # unweighted option to return_bins # warning raised if kperp_bins not linearly spaced kperp_log = np.logspace(-2, 0, 100) - _, _, _ = test.get_cylindrical_wf(bl_len, - kperp_bins=kperp_log*test.kunits, - kpara_bins=None, - return_bins='unweighted') + _, _, _ = test.get_cylindrical_wf( + bl_len, + kperp_bins=kperp_log * test.kunits, + kpara_bins=None, + return_bins="unweighted", + ) # kpara bins - _, kpara3, cyl_wf3 = test.get_cylindrical_wf(bl_len, - kperp_bins=None, - kpara_bins=kpara*test.kunits, - return_bins='unweighted') + _, kpara3, cyl_wf3 = test.get_cylindrical_wf( + bl_len, + kperp_bins=None, + kpara_bins=kpara * test.kunits, + return_bins="unweighted", + ) assert np.all(cyl_wf3 == cyl_wf) assert np.all(kpara == kpara3) # warning raised if kpara_bins not linearly spaced kpara_log = np.logspace(-1, 1, 100) - _, _, _ = test.get_cylindrical_wf(bl_len, - kperp_bins=None, - kpara_bins=kpara_log*test.kunits, - return_bins='unweighted') + _, _, _ = test.get_cylindrical_wf( + bl_len, + kperp_bins=None, + kpara_bins=kpara_log * test.kunits, + return_bins="unweighted", + ) # test filling array by delay symmetry for odd number of delays - ft_beam_test = uvwindow.FTBeam.from_file(ftfile=self.ft_file, - spw_range=(self.spw_range[0], - self.spw_range[1]-1)) + ft_beam_test = uvwindow.FTBeam.from_file( + ftfile=self.ft_file, spw_range=(self.spw_range[0], self.spw_range[1] - 1) + ) test = uvwindow.UVWindow(ftbeam_obj=ft_beam_test) - kperp, kpara, cyl_wf = test.get_cylindrical_wf(bl_len, - return_bins='unweighted') + kperp, kpara, cyl_wf = test.get_cylindrical_wf(bl_len, return_bins="unweighted") def test_cylindrical_to_spherical(self): @@ -466,112 +571,147 @@ def test_cylindrical_to_spherical(self): # initialise object from keywords test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw) - kperp, kpara, cyl_wf = test.get_cylindrical_wf(bl_len, - kperp_bins=None, - kpara_bins=None, - return_bins='unweighted') - ktot = np.sqrt(kperp[:, None]**2+kpara**2) - + kperp, kpara, cyl_wf = test.get_cylindrical_wf( + bl_len, kperp_bins=None, kpara_bins=None, return_bins="unweighted" + ) + ktot = np.sqrt(kperp[:, None] ** 2 + kpara**2) + # proper usage - sph_wf, weighted_k = test.cylindrical_to_spherical(cyl_wf=cyl_wf, - kbins=self.kbins, - ktot=ktot, - bl_lens=bl_len, - bl_weights=[2.]) - sph_wf, weighted_k = test.cylindrical_to_spherical(cyl_wf=cyl_wf[None], - kbins=self.kbins, - ktot=ktot, - bl_lens=bl_len, - bl_weights=None) + sph_wf, weighted_k = test.cylindrical_to_spherical( + cyl_wf=cyl_wf, kbins=self.kbins, ktot=ktot, bl_lens=bl_len, bl_weights=[2.0] + ) + sph_wf, weighted_k = test.cylindrical_to_spherical( + cyl_wf=cyl_wf[None], + kbins=self.kbins, + ktot=ktot, + bl_lens=bl_len, + bl_weights=None, + ) # ktot has shape different from cyl_wf - pytest.raises(AssertionError, test.cylindrical_to_spherical, cyl_wf=cyl_wf, - kbins=self.kbins, ktot=np.sqrt(kperp[:-2, None]**2+kpara**2), - bl_lens=bl_len) + pytest.raises( + AssertionError, + test.cylindrical_to_spherical, + cyl_wf=cyl_wf, + kbins=self.kbins, + ktot=np.sqrt(kperp[:-2, None] ** 2 + kpara**2), + bl_lens=bl_len, + ) # only one k-bin - pytest.raises(AssertionError, test.cylindrical_to_spherical, cyl_wf=cyl_wf, - kbins=self.kbins[:1], ktot=ktot, bl_lens=bl_len) + pytest.raises( + AssertionError, + test.cylindrical_to_spherical, + cyl_wf=cyl_wf, + kbins=self.kbins[:1], + ktot=ktot, + bl_lens=bl_len, + ) # weights have shape different from bl_lens - pytest.raises(AssertionError, test.cylindrical_to_spherical, cyl_wf=cyl_wf, - kbins=self.kbins, ktot=ktot, bl_lens=bl_len, - bl_weights=[1.,2.]) + pytest.raises( + AssertionError, + test.cylindrical_to_spherical, + cyl_wf=cyl_wf, + kbins=self.kbins, + ktot=ktot, + bl_lens=bl_len, + bl_weights=[1.0, 2.0], + ) # bl_lens has different size to cyl_wf.shape[0] - pytest.raises(AssertionError, test.cylindrical_to_spherical, cyl_wf=cyl_wf[None], - kbins=self.kbins, ktot=ktot, bl_lens=self.lens[:2], - bl_weights=[1.,2.]) + pytest.raises( + AssertionError, + test.cylindrical_to_spherical, + cyl_wf=cyl_wf[None], + kbins=self.kbins, + ktot=ktot, + bl_lens=self.lens[:2], + bl_weights=[1.0, 2.0], + ) # raise warning if empty bins - kbins_test = np.arange(2,5,step=.5)*test.kunits + kbins_test = np.arange(2, 5, step=0.5) * test.kunits test.verbose = True - sph_wf, weighted_k = test.cylindrical_to_spherical(cyl_wf=cyl_wf, - kbins=kbins_test, - ktot=ktot, - bl_lens=bl_len) + sph_wf, weighted_k = test.cylindrical_to_spherical( + cyl_wf=cyl_wf, kbins=kbins_test, ktot=ktot, bl_lens=bl_len + ) # raise warning if kbins not linearly spaced kbins_log = np.logspace(-2, 2, 20) - _, _ = test.cylindrical_to_spherical(cyl_wf=cyl_wf, - kbins=kbins_log*test.kunits, - ktot=ktot, - bl_lens=bl_len) + _, _ = test.cylindrical_to_spherical( + cyl_wf=cyl_wf, kbins=kbins_log * test.kunits, ktot=ktot, bl_lens=bl_len + ) def test_get_spherical_wf(self): - - bl_len = self.lens[12] - # initialise object from keywords test = uvwindow.UVWindow(ftbeam_obj=self.ft_beam_obj_spw) - WF, weighted_k = test.get_spherical_wf(kbins=self.kbins, - bl_lens=self.lens[:1], - bl_weights=[1], - kperp_bins=None, - kpara_bins=None, - return_weighted_k=True, - verbose=True) + WF, weighted_k = test.get_spherical_wf( + kbins=self.kbins, + bl_lens=self.lens[:1], + bl_weights=[1], + kperp_bins=None, + kpara_bins=None, + return_weighted_k=True, + verbose=True, + ) kperp_bins = test.get_kperp_bins(self.lens[:1]) kpara_bins = test.get_kpara_bins(test.freq_array) - WF = test.get_spherical_wf(kbins=self.kbins, - kperp_bins=kperp_bins, - kpara_bins=kpara_bins, - bl_lens=self.lens[:1], - bl_weights=None, - return_weighted_k=False, - verbose=None) + WF = test.get_spherical_wf( + kbins=self.kbins, + kperp_bins=kperp_bins, + kpara_bins=kpara_bins, + bl_lens=self.lens[:1], + bl_weights=None, + return_weighted_k=False, + verbose=None, + ) # check inputs - pytest.raises(AttributeError, test.get_spherical_wf, kbins=self.kbins.value, - bl_lens=self.lens[:2]) - pytest.raises(AssertionError, test.get_spherical_wf, kbins=self.kbins, - bl_lens=self.lens[:2], bl_weights=[1.]) - pytest.raises(AssertionError, test.get_spherical_wf, - kbins=self.kbins.value[2]*test.kunits, - bl_lens=self.lens[:1]) + pytest.raises( + AttributeError, + test.get_spherical_wf, + kbins=self.kbins.value, + bl_lens=self.lens[:2], + ) + pytest.raises( + AssertionError, + test.get_spherical_wf, + kbins=self.kbins, + bl_lens=self.lens[:2], + bl_weights=[1.0], + ) + pytest.raises( + AssertionError, + test.get_spherical_wf, + kbins=self.kbins.value[2] * test.kunits, + bl_lens=self.lens[:1], + ) # test kpara bins not outside of spectral window # will print warning - kpara_centre = test.cosmo.tau_to_kpara(test.avg_z, - little_h=test.little_h)\ + kpara_centre = ( + test.cosmo.tau_to_kpara(test.avg_z, little_h=test.little_h) * abs(test.dly_array).max() - WF = test.get_spherical_wf(kbins=self.kbins, kperp_bins=kperp_bins, - kpara_bins=np.arange(2.*kpara_centre, - 10*kpara_centre, - step=kpara_centre) - * test.kunits, - bl_lens=self.lens[:1]) + ) + test.get_spherical_wf( + kbins=self.kbins, + kperp_bins=kperp_bins, + kpara_bins=np.arange( + 2.0 * kpara_centre, 10 * kpara_centre, step=kpara_centre + ) + * test.kunits, + bl_lens=self.lens[:1], + ) # warning raised if kbins not linearly spaced kperp_log = np.logspace(-2, 0, 100) - _ = test.get_spherical_wf(kbins=self.kbins, - kperp_bins=kperp_log*test.kunits, - bl_lens=self.lens[:1]) + _ = test.get_spherical_wf( + kbins=self.kbins, kperp_bins=kperp_log * test.kunits, bl_lens=self.lens[:1] + ) kpara_log = np.logspace(-1, 1, 100) - _ = test.get_spherical_wf(kbins=self.kbins, - kpara_bins=kpara_log*test.kunits, - bl_lens=self.lens[:1]) + _ = test.get_spherical_wf( + kbins=self.kbins, kpara_bins=kpara_log * test.kunits, bl_lens=self.lens[:1] + ) kbins_log = np.logspace(-2, 2, 20) - _ = test.get_spherical_wf(kbins=kbins_log*test.kunits, - bl_lens=self.lens[:1]) - + _ = test.get_spherical_wf(kbins=kbins_log * test.kunits, bl_lens=self.lens[:1]) def test_check_kunits(self): @@ -590,34 +730,60 @@ def test_run_and_write(self): kperp_bins = test.get_kperp_bins(self.lens[:1]) kpara_bins = test.get_kpara_bins(test.freq_array) # proper usage - test.run_and_write(filepath=filepath, - bl_lens=self.lens[:1], - kperp_bins=kperp_bins, - kpara_bins=kpara_bins, - clobber=False) + test.run_and_write( + filepath=filepath, + bl_lens=self.lens[:1], + kperp_bins=kperp_bins, + kpara_bins=kpara_bins, + clobber=False, + ) # raise error if file already exists and clobber is False - pytest.raises(IOError, test.run_and_write, filepath=filepath, - bl_lens=self.lens[:1], - clobber=False) + pytest.raises( + IOError, + test.run_and_write, + filepath=filepath, + bl_lens=self.lens[:1], + clobber=False, + ) # does not if clobber is True - test.run_and_write(filepath=filepath, - bl_lens=[self.lens[:1]], - kperp_bins=None, kpara_bins=None, - clobber=True) + test.run_and_write( + filepath=filepath, + bl_lens=[self.lens[:1]], + kperp_bins=None, + kpara_bins=None, + clobber=True, + ) # check inputs - pytest.raises(AssertionError, test.run_and_write, filepath=filepath, - bl_lens=self.lens[:1], bl_weights=[1.,1.], clobber=True) - pytest.raises(AttributeError, test.run_and_write, filepath=filepath, - bl_lens=self.lens[:1], - kperp_bins=kperp_bins.value, clobber=True) - pytest.raises(AttributeError, test.run_and_write, filepath=filepath, - bl_lens=self.lens[:1], - kpara_bins=kpara_bins.value, clobber=True) + pytest.raises( + AssertionError, + test.run_and_write, + filepath=filepath, + bl_lens=self.lens[:1], + bl_weights=[1.0, 1.0], + clobber=True, + ) + pytest.raises( + AttributeError, + test.run_and_write, + filepath=filepath, + bl_lens=self.lens[:1], + kperp_bins=kperp_bins.value, + clobber=True, + ) + pytest.raises( + AttributeError, + test.run_and_write, + filepath=filepath, + bl_lens=self.lens[:1], + kpara_bins=kpara_bins.value, + clobber=True, + ) if os.path.exists(filepath): os.remove(filepath) - + + if __name__ == "__main__": unittest.main() diff --git a/hera_pspec/utils.py b/hera_pspec/utils.py index ef99bcef..54e478cb 100644 --- a/hera_pspec/utils.py +++ b/hera_pspec/utils.py @@ -50,19 +50,23 @@ def cov(d1, w1, d2=None, w2=None, conj_1=False, conj_2=True): cov : array_like Covariance (or cross-variance) matrix of size (M,M) """ - if d2 is None: d2,w2 = d1,w1 - if not np.isreal(w1).all(): raise TypeError("Weight matrices must be real") - if not np.isreal(w2).all(): raise TypeError("Weight matrices must be real") - if np.less(w1, 0.).any() or np.less(w2, 0.).any(): + if d2 is None: + d2, w2 = d1, w1 + if not np.isreal(w1).all(): + raise TypeError("Weight matrices must be real") + if not np.isreal(w2).all(): + raise TypeError("Weight matrices must be real") + if np.less(w1, 0.0).any() or np.less(w2, 0.0).any(): raise ValueError("Weight matrices must be positive") - d1sum,d1wgt = (w1*d1).sum(axis=1), w1.sum(axis=1) - d2sum,d2wgt = (w2*d2).sum(axis=1), w2.sum(axis=1) + d1sum, d1wgt = (w1 * d1).sum(axis=1), w1.sum(axis=1) + d2sum, d2wgt = (w2 * d2).sum(axis=1), w2.sum(axis=1) x1 = d1sum / np.where(d1wgt > 0, d1wgt, 1) x2 = d2sum / np.where(d2wgt > 0, d2wgt, 1) - x1.shape = (-1,1); x2.shape = (-1,1) + x1.shape = (-1, 1) + x2.shape = (-1, 1) - z1 = w1*d1 - z2 = w2*d2 + z1 = w1 * d1 + z2 = w2 * d2 if conj_1: z1 = z1.conj() @@ -77,6 +81,7 @@ def cov(d1, w1, d2=None, w2=None, conj_1=False, conj_2=True): C -= np.outer(x1, x2) return C + def variance_from_auto_correlations(uvd, bl, spw_range, time_index): """ Predict noise variance on a baseline from autocorrelation amplitudes on antennas. @@ -103,13 +108,15 @@ def variance_from_auto_correlations(uvd, bl, spw_range, time_index): var : ndarray, (spw_Nfreqs,) """ - assert isinstance(bl, tuple) and len(bl)==3, "bl must be fed as Length-3 tuple" - assert isinstance(spw_range, tuple) and len(spw_range)==2, "spw_range must be fed as Length-2 tuple" + assert isinstance(bl, tuple) and len(bl) == 3, "bl must be fed as Length-3 tuple" + assert ( + isinstance(spw_range, tuple) and len(spw_range) == 2 + ), "spw_range must be fed as Length-2 tuple" dt = np.median(uvd.integration_time) # Delta_t df = uvd.channel_width # B - bl1 = (bl[0],bl[0], bl[2]) + bl1 = (bl[0], bl[0], bl[2]) # baseline b_alpha bl2 = (bl[1], bl[1], bl[2]) # baseline b_beta @@ -117,14 +124,23 @@ def variance_from_auto_correlations(uvd, bl, spw_range, time_index): x_bl1 = uvd.get_data(bl1)[time_index, spw] x_bl2 = uvd.get_data(bl2)[time_index, spw] nsample_bl = uvd.get_nsamples(bl)[time_index, spw] - nsample_bl = np.where(nsample_bl>0, nsample_bl, np.median(uvd.nsample_array[:,:,spw,:])) + nsample_bl = np.where( + nsample_bl > 0, nsample_bl, np.median(uvd.nsample_array[:, :, spw, :]) + ) # some impainted data have zero nsample while is not flagged, and they will be assigned the median nsample within the spectral window. - var = np.abs(x_bl1*x_bl2.conj()) / dt / df / nsample_bl + var = np.abs(x_bl1 * x_bl2.conj()) / dt / df / nsample_bl return var -def construct_blpairs(bls, exclude_auto_bls=False, exclude_cross_bls=False, - exclude_permutations=False, group=False, Nblps_per_group=1): + +def construct_blpairs( + bls, + exclude_auto_bls=False, + exclude_cross_bls=False, + exclude_permutations=False, + group=False, + Nblps_per_group=1, +): """ Construct a list of baseline-pairs from a baseline-group. This function can be used to easily convert a single list of baselines into the input @@ -176,10 +192,13 @@ def construct_blpairs(bls, exclude_auto_bls=False, exclude_cross_bls=False, List of blpair tuples. """ # assert form - assert isinstance(bls, (list, np.ndarray)) and isinstance(bls[0], tuple), \ - "bls must be fed as list or ndarray of baseline antnum tuples. Use " \ + assert isinstance(bls, (list, np.ndarray)) and isinstance(bls[0], tuple), ( + "bls must be fed as list or ndarray of baseline antnum tuples. Use " "UVData.baseline_to_antnums() to convert baseline integers to tuples." - assert (not exclude_auto_bls) or (not exclude_cross_bls), "Can't exclude both auto and cross blpairs" + ) + assert (not exclude_auto_bls) or ( + not exclude_cross_bls + ), "Can't exclude both auto and cross blpairs" # form blpairs w/o explicitly forming auto blpairs # however, if there are repeated bl in bls, there will be auto bls in blpairs @@ -219,9 +238,9 @@ def construct_blpairs(bls, exclude_auto_bls=False, exclude_cross_bls=False, new_bls1 = [] new_bls2 = [] for i in range(Ngrps): - new_blps.append(blpairs[i*Nblps_per_group:(i+1)*Nblps_per_group]) - new_bls1.append(bls1[i*Nblps_per_group:(i+1)*Nblps_per_group]) - new_bls2.append(bls2[i*Nblps_per_group:(i+1)*Nblps_per_group]) + new_blps.append(blpairs[i * Nblps_per_group : (i + 1) * Nblps_per_group]) + new_bls1.append(bls1[i * Nblps_per_group : (i + 1) * Nblps_per_group]) + new_bls2.append(bls2[i * Nblps_per_group : (i + 1) * Nblps_per_group]) bls1 = new_bls1 bls2 = new_bls2 @@ -230,13 +249,23 @@ def construct_blpairs(bls, exclude_auto_bls=False, exclude_cross_bls=False, return bls1, bls2, blpairs -def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, - xant_flag_thresh=0.95, exclude_auto_bls=False, - exclude_cross_bls=False, - exclude_permutations=True, Nblps_per_group=None, - bl_len_range=(0, 1e10), bl_deg_range=(0, 180), - xants=None, include_autocorrs=False, - include_crosscorrs=True, extra_info=False): +def calc_blpair_reds( + uvd1, + uvd2, + bl_tol=1.0, + filter_blpairs=True, + xant_flag_thresh=0.95, + exclude_auto_bls=False, + exclude_cross_bls=False, + exclude_permutations=True, + Nblps_per_group=None, + bl_len_range=(0, 1e10), + bl_deg_range=(0, 180), + xants=None, + include_autocorrs=False, + include_crosscorrs=True, + extra_info=False, +): """ Use hera_cal.redcal to get matching, redundant baseline-pair groups from uvd1 and uvd2 within the specified baseline tolerance, not including @@ -335,8 +364,10 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, # assert antenna positions match for a in set(antpos1).union(set(antpos2)): if a in antpos1 and a in antpos2: - msg = "antenna positions from uvd1 and uvd2 do not agree to within " \ - "tolerance of {} m".format(bl_tol) + msg = ( + "antenna positions from uvd1 and uvd2 do not agree to within " + "tolerance of {} m".format(bl_tol) + ) assert np.linalg.norm(antpos1[a] - antpos2[a]) < bl_tol, msg # calculate xants via flags if asked @@ -350,10 +381,12 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, # continue if autocorr and we dont want to include them if not include_autocorrs: - if antnums[0] == antnums[1]: continue + if antnums[0] == antnums[1]: + continue if not include_crosscorrs: - if antnums[0] != antnums[1]: continue + if antnums[0] != antnums[1]: + continue # work on xants1 if bl in uvd1.baseline_array: @@ -386,16 +419,25 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, xants2 += xants # construct redundant groups - reds, lens, angs = get_reds(antpos, bl_error_tol=bl_tol, xants=xants1+xants2, - add_autos=include_autocorrs, autos_only=not(include_crosscorrs), - bl_deg_range=bl_deg_range, bl_len_range=bl_len_range) + reds, lens, angs = get_reds( + antpos, + bl_error_tol=bl_tol, + xants=xants1 + xants2, + add_autos=include_autocorrs, + autos_only=not (include_crosscorrs), + bl_deg_range=bl_deg_range, + bl_len_range=bl_len_range, + ) # construct baseline pairs baselines1, baselines2, blpairs, red_groups = [], [], [], [] for j, r in enumerate(reds): - (bls1, bls2, - blps) = construct_blpairs(r, exclude_auto_bls=exclude_auto_bls, - exclude_cross_bls=exclude_cross_bls, group=False, - exclude_permutations=exclude_permutations) + (bls1, bls2, blps) = construct_blpairs( + r, + exclude_auto_bls=exclude_auto_bls, + exclude_cross_bls=exclude_cross_bls, + group=False, + exclude_permutations=exclude_permutations, + ) if len(bls1) < 1: continue @@ -407,8 +449,9 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, for blp in blps: bl1 = blp[0] bl2 = blp[1] - if ((bl1 in uvd1_bls) or (bl1[::-1] in uvd1_bls)) \ - and ((bl2 in uvd2_bls) or (bl2[::-1] in uvd2_bls)): + if ((bl1 in uvd1_bls) or (bl1[::-1] in uvd1_bls)) and ( + (bl2 in uvd2_bls) or (bl2[::-1] in uvd2_bls) + ): _bls1.append(bl1) _bls2.append(bl2) bls1, bls2 = _bls1, _bls2 @@ -420,14 +463,22 @@ def calc_blpair_reds(uvd1, uvd2, bl_tol=1.0, filter_blpairs=True, # group if desired if Nblps_per_group is not None: Ngrps = int(np.ceil(float(len(blps)) / Nblps_per_group)) - bls1 = [bls1[Nblps_per_group*i:Nblps_per_group*(i+1)] - for i in range(Ngrps)] - bls2 = [bls2[Nblps_per_group*i:Nblps_per_group*(i+1)] - for i in range(Ngrps)] - blps = [blps[Nblps_per_group*i:Nblps_per_group*(i+1)] - for i in range(Ngrps)] - rinds = [rinds[Nblps_per_group*i:Nblps_per_group*(i+1)] - for i in range(Ngrps)] + bls1 = [ + bls1[Nblps_per_group * i : Nblps_per_group * (i + 1)] + for i in range(Ngrps) + ] + bls2 = [ + bls2[Nblps_per_group * i : Nblps_per_group * (i + 1)] + for i in range(Ngrps) + ] + blps = [ + blps[Nblps_per_group * i : Nblps_per_group * (i + 1)] + for i in range(Ngrps) + ] + rinds = [ + rinds[Nblps_per_group * i : Nblps_per_group * (i + 1)] + for i in range(Ngrps) + ] baselines1.extend(bls1) baselines2.extend(bls2) @@ -461,7 +512,7 @@ def get_delays(freqs, n_dlys=None): Delta_nu = np.median(np.diff(freqs)) n_freqs = freqs.size - if n_dlys == None: # assume that n_dlys = n_freqs if not specified + if n_dlys == None: # assume that n_dlys = n_freqs if not specified n_dlys = n_freqs # Calculate the delays @@ -506,10 +557,11 @@ def spw_range_from_freqs(data, freq_range, bounds_error=True): try: freqs = data.freq_array if len(freqs.shape) == 2 and freqs.shape[0] == 1: - freqs = freqs.flatten() # Support UVData 2D freq_array + freqs = freqs.flatten() # Support UVData 2D freq_array elif len(freqs.shape) > 2: - raise ValueError("data.freq_array has unsupported shape: %s" \ - % str(freqs.shape)) + raise ValueError( + "data.freq_array has unsupported shape: %s" % str(freqs.shape) + ) except: raise AttributeError("Object 'data' does not have a freq_array attribute.") @@ -517,7 +569,9 @@ def spw_range_from_freqs(data, freq_range, bounds_error=True): is_tuple = False if isinstance(freq_range, tuple): is_tuple = True - freq_range = [freq_range,] + freq_range = [ + freq_range, + ] # Make sure freq_range is now a list (of tuples) if not isinstance(freq_range, list): @@ -528,18 +582,23 @@ def spw_range_from_freqs(data, freq_range, bounds_error=True): for frange in freq_range: fmin, fmax = frange if fmin > fmax: - raise ValueError("Upper bound of spectral window is less than " - "the lower bound.") + raise ValueError( + "Upper bound of spectral window is less than " "the lower bound." + ) # Check that this doesn't go beyond the available range of freqs if fmin < np.min(freqs) and bounds_error: - raise ValueError("Lower bound of spectral window is below the " - "available frequency range. (Note: freqs should " - "be in Hz)") + raise ValueError( + "Lower bound of spectral window is below the " + "available frequency range. (Note: freqs should " + "be in Hz)" + ) if fmax > np.max(freqs) and bounds_error: - raise ValueError("Upper bound of spectral window is above the " - "available frequency range. (Note: freqs should " - "be in Hz)") + raise ValueError( + "Upper bound of spectral window is above the " + "available frequency range. (Note: freqs should " + "be in Hz)" + ) # Get indices within this range idxs = np.where(np.logical_and(freqs >= fmin, freqs < fmax))[0] @@ -547,7 +606,8 @@ def spw_range_from_freqs(data, freq_range, bounds_error=True): spw_range.append(spw) # Unpack from list if only a single tuple was specified originally - if is_tuple: return spw_range[0] + if is_tuple: + return spw_range[0] return spw_range @@ -586,21 +646,24 @@ def spw_range_from_redshifts(data, z_range, bounds_error=True): is_tuple = False if isinstance(z_range, tuple): is_tuple = True - z_range = [z_range,] + z_range = [ + z_range, + ] # Convert redshifts to frequencies (in Hz) freq_range = [] for zrange in z_range: zmin, zmax = zrange - freq_range.append( (Cosmo_Conversions.z2f(zmax), - Cosmo_Conversions.z2f(zmin)) ) + freq_range.append((Cosmo_Conversions.z2f(zmax), Cosmo_Conversions.z2f(zmin))) # Use freq. function to get spectral window - spw_range = spw_range_from_freqs(data=data, freq_range=freq_range, - bounds_error=bounds_error) + spw_range = spw_range_from_freqs( + data=data, freq_range=freq_range, bounds_error=bounds_error + ) # Unpack from list if only a single tuple was specified originally - if is_tuple: return spw_range[0] + if is_tuple: + return spw_range[0] return spw_range @@ -629,10 +692,10 @@ def log(msg, f=None, lvl=0, tb=None, verbose=True): """ # catch for traceback if provided if tb is not None: - msg += "\n{}".format('\n'.join(traceback.format_exception(*tb))) + msg += "\n{}".format("\n".join(traceback.format_exception(*tb))) # print - output = "%s%s" % (" "*lvl, msg) + output = "%s%s" % (" " * lvl, msg) if verbose: print(output) @@ -653,19 +716,22 @@ def replace(d): if isinstance(d, (dict, odict)): for k in d.keys(): # 'None' and '' turn into None - if d[k] == 'None': d[k] = None + if d[k] == "None": + d[k] = None # list of lists turn into lists of tuples - if isinstance(d[k], list) \ - and np.all([isinstance(i, list) for i in d[k]]): + if isinstance(d[k], list) and np.all( + [isinstance(i, list) for i in d[k]] + ): d[k] = [tuple(i) for i in d[k]] - elif isinstance(d[k], (dict, odict)): replace(d[k]) + elif isinstance(d[k], (dict, odict)): + replace(d[k]) # Open and read config file - with open(config_file, 'r') as cfile: + with open(config_file, "r") as cfile: try: cfg = yaml.load(cfile, Loader=yaml.FullLoader) except yaml.YAMLError as exc: - raise(exc) + raise (exc) # Replace entries replace(cfg) @@ -680,11 +746,20 @@ def flatten(nested_list): return [item for sublist in nested_list for item in sublist] -def config_pspec_blpairs(uv_templates, pol_pairs, group_pairs, exclude_auto_bls=False, - exclude_permutations=True, bl_len_range=(0, 1e10), - bl_deg_range=(0, 180), xants=None, exclude_patterns=None, - include_autocorrs=False, - file_type='miriad', verbose=True): +def config_pspec_blpairs( + uv_templates, + pol_pairs, + group_pairs, + exclude_auto_bls=False, + exclude_permutations=True, + bl_len_range=(0, 1e10), + bl_deg_range=(0, 180), + xants=None, + exclude_patterns=None, + include_autocorrs=False, + file_type="miriad", + verbose=True, +): """ Given a list of glob-parseable file templates and selections for polarization and group labels, construct a master list of @@ -763,8 +838,9 @@ def config_pspec_blpairs(uv_templates, pol_pairs, group_pairs, exclude_auto_bls= # type check if isinstance(uv_templates, (str, np.str)): uv_templates = [uv_templates] - assert len(pol_pairs) == len(group_pairs), "len(pol_pairs) must equal "\ - "len(group_pairs)" + assert len(pol_pairs) == len(group_pairs), ( + "len(pol_pairs) must equal " "len(group_pairs)" + ) # get unique pols and groups pols = sorted(set([item for sublist in pol_pairs for item in sublist])) @@ -783,8 +859,9 @@ def config_pspec_blpairs(uv_templates, pol_pairs, group_pairs, exclude_auto_bls= pol_grps.append((pol, group)) # insert into unique_files with {pol} and {group} re-inserted for _file in files: - _unique_file = _file.replace(".{pol}.".format(pol=pol), - ".{pol}.").replace(".{group}.".format(group=group), ".{group}.") + _unique_file = _file.replace( + ".{pol}.".format(pol=pol), ".{pol}." + ).replace(".{group}.".format(group=group), ".{group}.") if _unique_file not in unique_files: unique_files.append(_unique_file) unique_files = sorted(unique_files) @@ -800,8 +877,10 @@ def config_pspec_blpairs(uv_templates, pol_pairs, group_pairs, exclude_auto_bls= # Add to list of files to be excluded if pattern in f: if verbose: - print("File matches pattern '%s' and will be excluded: %s" \ - % (pattern, f)) + print( + "File matches pattern '%s' and will be excluded: %s" + % (pattern, f) + ) to_exclude.append(f) continue @@ -824,19 +903,27 @@ def config_pspec_blpairs(uv_templates, pol_pairs, group_pairs, exclude_auto_bls= uvd.read(_file, read_data=False, file_type=file_type) # get baseline pairs - (_bls1, _bls2, _, _, - _) = calc_blpair_reds(uvd, uvd, filter_blpairs=False, exclude_auto_bls=exclude_auto_bls, - exclude_permutations=exclude_permutations, bl_len_range=bl_len_range, - include_autocorrs=include_autocorrs, bl_deg_range=bl_deg_range) + (_bls1, _bls2, _, _, _) = calc_blpair_reds( + uvd, + uvd, + filter_blpairs=False, + exclude_auto_bls=exclude_auto_bls, + exclude_permutations=exclude_permutations, + bl_len_range=bl_len_range, + include_autocorrs=include_autocorrs, + bl_deg_range=bl_deg_range, + ) # take out xants if fed if xants is not None: bls1, bls2 = [], [] for bl1, bl2 in zip(_bls1, _bls2): - if bl1[0] not in xants \ - and bl1[1] not in xants \ - and bl2[0] not in xants \ - and bl2[1] not in xants: + if ( + bl1[0] not in xants + and bl1[1] not in xants + and bl2[0] not in xants + and bl2[1] not in xants + ): bls1.append(bl1) bls2.append(bl2) else: @@ -848,7 +935,11 @@ def config_pspec_blpairs(uv_templates, pol_pairs, group_pairs, exclude_auto_bls= for pp, gp in zip(pol_pairs, group_pairs): if (pp[0], gp[0]) not in pol_grps or (pp[1], gp[1]) not in pol_grps: if verbose: - print("pol_pair {} and group_pair {} not found in data files".format(pp, gp)) + print( + "pol_pair {} and group_pair {} not found in data files".format( + pp, gp + ) + ) continue groupings[(tuple(gp), tuple(pp))] = blps @@ -889,16 +980,17 @@ def get_blvec_reds(blvecs, bl_error_tol=1.0, match_bl_lens=False): A list of baseline string tags denoting bl length and angle """ from hera_pspec import UVPSpec + # type check - assert isinstance(blvecs, (dict, odict, UVPSpec)), \ - "blpairs must be fed as a dict or UVPSpec" + assert isinstance( + blvecs, (dict, odict, UVPSpec) + ), "blpairs must be fed as a dict or UVPSpec" if isinstance(blvecs, UVPSpec): # get baseline vectors uvp = blvecs bls = uvp.bl_array bl_vecs = uvp.get_ENU_bl_vecs()[:, :2] - blvecs = dict(list(zip( [uvp.bl_to_antnums(_bls) for _bls in bls], - bl_vecs ))) + blvecs = dict(list(zip([uvp.bl_to_antnums(_bls) for _bls in bls], bl_vecs))) # get baseline-pairs blpairs = uvp.get_blpairs() # form dictionary @@ -906,7 +998,7 @@ def get_blvec_reds(blvecs, bl_error_tol=1.0, match_bl_lens=False): for blp in blpairs: bl1 = blp[0] bl2 = blp[1] - _blvecs[blp] = (blvecs[bl1] + blvecs[bl2]) / 2. + _blvecs[blp] = (blvecs[bl1] + blvecs[bl2]) / 2.0 blvecs = _blvecs # create empty lists @@ -922,16 +1014,23 @@ def get_blvec_reds(blvecs, bl_error_tol=1.0, match_bl_lens=False): bl_vec = blvecs[bl][:2] bl_len = np.linalg.norm(bl_vec) bl_ang = np.arctan2(*bl_vec[::-1]) * 180 / np.pi - if bl_ang < 0: bl_ang = (bl_ang + 180) % 360 + if bl_ang < 0: + bl_ang = (bl_ang + 180) % 360 bl_tag = "{:03.0f}_{:03.0f}".format(bl_len, bl_ang) # append to list if unique within tolerance if match_bl_lens: # match only on bl length - match = [np.all(np.isclose(bll, bl_len, rtol=0.0, atol=bl_error_tol)) for bll in red_bl_len] + match = [ + np.all(np.isclose(bll, bl_len, rtol=0.0, atol=bl_error_tol)) + for bll in red_bl_len + ] else: # match on full bl vector - match = [np.all(np.isclose(blv, bl_vec, rtol=0.0, atol=bl_error_tol)) for blv in red_bl_vec] + match = [ + np.all(np.isclose(blv, bl_vec, rtol=0.0, atol=bl_error_tol)) + for blv in red_bl_vec + ] if np.any(match): match_id = np.where(match)[0][0] red_bl_grp[match_id].append(bl) @@ -954,8 +1053,9 @@ def get_blvec_reds(blvecs, bl_error_tol=1.0, match_bl_lens=False): return red_bl_grp, red_bl_len, red_bl_ang, red_bl_tag -def job_monitor(run_func, iterator, action_name, M=map, lf=None, maxiter=1, - verbose=True): +def job_monitor( + run_func, iterator, action_name, M=map, lf=None, maxiter=1, verbose=True +): """ Job monitoring function, used to send elements of iterator through calls of run_func. Can be parallelized if the input M function is from the @@ -1002,14 +1102,20 @@ def job_monitor(run_func, iterator, action_name, M=map, lf=None, maxiter=1, # check for len-0 if len(exit_codes) == 0: - raise ValueError("No output generated from run_func over iterator {}".format(iterator)) + raise ValueError( + "No output generated from run_func over iterator {}".format(iterator) + ) # inspect for failures if np.all(exit_codes != 0): # everything failed, raise error - log("\n{}\nAll {} jobs failed w/ exit codes\n {}: {}\n".format("-"*60, - action_name, exit_codes, tnow), - f=lf, verbose=verbose) + log( + "\n{}\nAll {} jobs failed w/ exit codes\n {}: {}\n".format( + "-" * 60, action_name, exit_codes, tnow + ), + f=lf, + verbose=verbose, + ) raise ValueError("All {} jobs failed".format(action_name)) # if not all failed, try re-run @@ -1032,14 +1138,20 @@ def job_monitor(run_func, iterator, action_name, M=map, lf=None, maxiter=1, # print failures if they exist if len(failures) > 0: - log("\nSome {} jobs failed after {} tries:\n{}".format(action_name, - maxiter, - failures), - f=lf, verbose=verbose) + log( + "\nSome {} jobs failed after {} tries:\n{}".format( + action_name, maxiter, failures + ), + f=lf, + verbose=verbose, + ) else: t_run = time.time() - t_start - log("\nAll {} jobs ran through ({:1.1f} sec)".format(action_name, t_run), - f=lf, verbose=verbose) + log( + "\nAll {} jobs ran through ({:1.1f} sec)".format(action_name, t_run), + f=lf, + verbose=verbose, + ) return failures @@ -1082,10 +1194,18 @@ def get_bl_lens_angs(blvecs, bl_error_tol=1.0): return lens, angs -def get_reds(uvd, bl_error_tol=1.0, pick_data_ants=False, bl_len_range=(0, 1e4), - bl_deg_range=(0, 180), xants=None, add_autos=False, - autos_only=False, min_EW_cut=0, - file_type='miriad'): +def get_reds( + uvd, + bl_error_tol=1.0, + pick_data_ants=False, + bl_len_range=(0, 1e4), + bl_deg_range=(0, 180), + xants=None, + add_autos=False, + autos_only=False, + min_EW_cut=0, + file_type="miriad", +): """ Given a UVData object, a Miriad filepath or antenna position dictionary, calculate redundant baseline groups using hera_cal.redcal and optionally @@ -1155,8 +1275,10 @@ def get_reds(uvd, bl_error_tol=1.0, pick_data_ants=False, bl_len_range=(0, 1e4), # use antenna position dictionary antpos_dict = uvd else: - raise TypeError("uvd must be a UVData object, filename string, or dict " - "of antenna positions.") + raise TypeError( + "uvd must be a UVData object, filename string, or dict " + "of antenna positions." + ) # get redundant baselines reds = redcal.get_pos_reds(antpos_dict, bl_error_tol=bl_error_tol) @@ -1167,9 +1289,12 @@ def get_reds(uvd, bl_error_tol=1.0, pick_data_ants=False, bl_len_range=(0, 1e4), # restrict baselines _reds, _lens, _angs = [], [], [] for i, (l, a) in enumerate(zip(lens, angs)): - if l < bl_len_range[0] or l > bl_len_range[1]: continue - if a < bl_deg_range[0] or a > bl_deg_range[1]: continue - if np.abs(l * np.cos(a * np.pi / 180)) < min_EW_cut: continue + if l < bl_len_range[0] or l > bl_len_range[1]: + continue + if a < bl_deg_range[0] or a > bl_deg_range[1]: + continue + if np.abs(l * np.cos(a * np.pi / 180)) < min_EW_cut: + continue _reds.append(reds[i]) _lens.append(lens[i]) _angs.append(angs[i]) @@ -1186,7 +1311,6 @@ def get_reds(uvd, bl_error_tol=1.0, pick_data_ants=False, bl_len_range=(0, 1e4), lens = lens[:1] angs = angs[:1] - # filter based on xants if xants is not None: _reds, _lens, _angs = [], [], [] @@ -1221,6 +1345,7 @@ def pspecdata_time_difference(ds, time_diff): ds_td : PSpecData object """ from hera_pspec.pspecdata import PSpecData + uvd1 = ds.dsets[0] uvd2 = ds.dsets[1] uvd10 = uvd_time_difference(uvd1, time_diff) @@ -1245,10 +1370,12 @@ def uvd_time_difference(uvd, time_diff): ------- uvd_td : UVData object """ - min_time_diff = np.mean(np.unique(uvd.time_array)[1:]-np.unique(uvd.time_array)[0:-1]) + min_time_diff = np.mean( + np.unique(uvd.time_array)[1:] - np.unique(uvd.time_array)[0:-1] + ) index_diff = int(time_diff / min_time_diff) + 1 - if index_diff > len(np.unique(uvd.time_array))-2: - index_diff = len(np.unique(uvd.time_array))-2 + if index_diff > len(np.unique(uvd.time_array)) - 2: + index_diff = len(np.unique(uvd.time_array)) - 2 uvd0 = uvd.select(times=np.unique(uvd.time_array)[0:-1:index_diff], inplace=False) uvd1 = uvd.select(times=np.unique(uvd.time_array)[1::index_diff], inplace=False) @@ -1260,8 +1387,8 @@ def uvd_time_difference(uvd, time_diff): return uvd0 -AUTOVISPOLS = ['XX', 'YY', 'EE', 'NN'] -STOKPOLS = ['PI', 'PQ', 'PU', 'PV'] +AUTOVISPOLS = ["XX", "YY", "EE", "NN"] +STOKPOLS = ["PI", "PQ", "PU", "PV"] AUTOPOLS = AUTOVISPOLS + STOKPOLS @@ -1294,37 +1421,45 @@ def uvd_to_Tsys(uvd, beam, Tsys_outfile=None): # get uvd metadata pols = [pol for pol in uvd.get_pols() if pol.upper() in AUTOPOLS] # if pseudo Stokes pol in pols, substitute for pI - pols = sorted(set([pol if pol.upper() in AUTOVISPOLS else 'pI' for pol in pols])) + pols = sorted(set([pol if pol.upper() in AUTOVISPOLS else "pI" for pol in pols])) autobls = [bl for bl in uvd.get_antpairs() if bl[0] == bl[1]] uvd.select(bls=autobls, polarizations=pols) # construct beam from hera_pspec import pspecbeam from hera_pspec import uvpspec + if isinstance(beam, str): beam = pspecbeam.PSpecBeamUV(beam) elif isinstance(beam, pspecbeam.PSpecBeamBase): pass elif isinstance(beam, uvpspec.UVPSpec): uvp = beam - if hasattr(uvp, 'OmegaP'): + if hasattr(uvp, "OmegaP"): # use first pol in each polpair - uvp_pols = [pp[0] if pp[0].upper() not in STOKPOLS else 'pI' for pp in uvp.get_polpairs()] + uvp_pols = [ + pp[0] if pp[0].upper() not in STOKPOLS else "pI" + for pp in uvp.get_polpairs() + ] Op = {uvp_pol: uvp.OmegaP[:, ii] for ii, uvp_pol in enumerate(uvp_pols)} Opp = {uvp_pol: uvp.OmegaPP[:, ii] for ii, uvp_pol in enumerate(uvp_pols)} - beam = pspecbeam.PSpecBeamFromArray(Op, Opp, uvp.beam_freqs, cosmo=uvp.cosmo) + beam = pspecbeam.PSpecBeamFromArray( + Op, Opp, uvp.beam_freqs, cosmo=uvp.cosmo + ) else: raise ValueError("UVPSpec must have OmegaP and OmegaPP to make a beam") else: - raise ValueError("beam must be a string, PSpecBeamBase subclass or UVPSpec object") + raise ValueError( + "beam must be a string, PSpecBeamBase subclass or UVPSpec object" + ) # convert autos in Jy to Tsys in Kelvin - J2K = {pol: beam.Jy_to_mK(uvd.freq_array[0], pol=pol)/1e3 for pol in pols} + J2K = {pol: beam.Jy_to_mK(uvd.freq_array[0], pol=pol) / 1e3 for pol in pols} for blpol in uvd.get_antpairpols(): bl, pol = blpol[:2], blpol[2] tinds = uvd.antpair2ind(bl) if pol.upper() in STOKPOLS: - pol = 'pI' + pol = "pI" pind = pols.index(pol) uvd.data_array[tinds, 0, :, pind] *= J2K[pol] @@ -1333,7 +1468,16 @@ def uvd_to_Tsys(uvd, beam, Tsys_outfile=None): return uvd -def uvp_noise_error(uvp, auto_Tsys=None, err_type='P_N', precomp_P_N=None, P_SN_correction=True, num_steps_scalar=2000, little_h=True): + +def uvp_noise_error( + uvp, + auto_Tsys=None, + err_type="P_N", + precomp_P_N=None, + P_SN_correction=True, + num_steps_scalar=2000, + little_h=True, +): """ Calculate analytic thermal noise error for a UVPSpec object. Adds to uvp.stats_array inplace. @@ -1391,8 +1535,13 @@ def uvp_noise_error(uvp, auto_Tsys=None, err_type='P_N', precomp_P_N=None, P_SN_ scalar = {} for spw in uvp.spw_array: for polpair in uvp.polpair_array: - scalar[(spw, polpair)] = uvp.compute_scalar(spw, polpair, num_steps=num_steps_scalar, - little_h=little_h, noise_scalar=True) + scalar[(spw, polpair)] = uvp.compute_scalar( + spw, + polpair, + num_steps=num_steps_scalar, + little_h=little_h, + noise_scalar=True, + ) # iterate over spectral window for spw in uvp.spw_array: # get spw properties @@ -1407,46 +1556,84 @@ def uvp_noise_error(uvp, auto_Tsys=None, err_type='P_N', precomp_P_N=None, P_SN_ # iterate over polarization for polpair in uvp.polpair_array: pol = uvpspec_utils.polpair_int2tuple(polpair)[0] # integer - polstr = uvutils.polnum2str(pol) # TODO: use uvp.x_orientation when attr is added + polstr = uvutils.polnum2str( + pol + ) # TODO: use uvp.x_orientation when attr is added if polstr.upper() in STOKPOLS: - pol = 'pI' + pol = "pI" key = (spw, blp, polpair) if precomp_P_N is None: # take geometric mean of four antenna autocorrs and get OR'd flags - Tsys = (auto_Tsys.get_data(blp[0][0], blp[0][0], pol)[:, spw_start:spw_stop].real * \ - auto_Tsys.get_data(blp[0][1], blp[0][1], pol)[:, spw_start:spw_stop].real * \ - auto_Tsys.get_data(blp[1][0], blp[1][0], pol)[:, spw_start:spw_stop].real * \ - auto_Tsys.get_data(blp[1][1], blp[1][1], pol)[:, spw_start:spw_stop].real)**(1./4) - Tflag = auto_Tsys.get_flags(blp[0][0], blp[0][0], pol)[:, spw_start:spw_stop] + \ - auto_Tsys.get_flags(blp[0][1], blp[0][1], pol)[:, spw_start:spw_stop] + \ - auto_Tsys.get_flags(blp[1][0], blp[1][0], pol)[:, spw_start:spw_stop] + \ - auto_Tsys.get_flags(blp[1][1], blp[1][1], pol)[:, spw_start:spw_stop] + Tsys = ( + auto_Tsys.get_data(blp[0][0], blp[0][0], pol)[ + :, spw_start:spw_stop + ].real + * auto_Tsys.get_data(blp[0][1], blp[0][1], pol)[ + :, spw_start:spw_stop + ].real + * auto_Tsys.get_data(blp[1][0], blp[1][0], pol)[ + :, spw_start:spw_stop + ].real + * auto_Tsys.get_data(blp[1][1], blp[1][1], pol)[ + :, spw_start:spw_stop + ].real + ) ** (1.0 / 4) + Tflag = ( + auto_Tsys.get_flags(blp[0][0], blp[0][0], pol)[ + :, spw_start:spw_stop + ] + + auto_Tsys.get_flags(blp[0][1], blp[0][1], pol)[ + :, spw_start:spw_stop + ] + + auto_Tsys.get_flags(blp[1][0], blp[1][0], pol)[ + :, spw_start:spw_stop + ] + + auto_Tsys.get_flags(blp[1][1], blp[1][1], pol)[ + :, spw_start:spw_stop + ] + ) # average over frequency if np.all(Tflag): # fully flagged Tsys = np.inf else: # get weights - Tsys = np.sum(Tsys * ~Tflag * taper, axis=-1) / np.sum(~Tflag * taper, axis=-1).clip(1e-20, np.inf) + Tsys = np.sum(Tsys * ~Tflag * taper, axis=-1) / np.sum( + ~Tflag * taper, axis=-1 + ).clip(1e-20, np.inf) Tflag = np.all(Tflag, axis=-1) # interpolate to appropriate LST grid if np.count_nonzero(~Tflag) > 1: - Tsys = interp1d(lsts[~Tflag], Tsys[~Tflag], kind='nearest', bounds_error=False, fill_value='extrapolate')(lst_avg) + Tsys = interp1d( + lsts[~Tflag], + Tsys[~Tflag], + kind="nearest", + bounds_error=False, + fill_value="extrapolate", + )(lst_avg) else: Tsys = Tsys[0] # calculate P_N - P_N = uvp.generate_noise_spectra(spw, polpair, Tsys, blpairs=[blp], form='Pk', component='real', scalar=scalar[(spw, polpair)])[blp_int] + P_N = uvp.generate_noise_spectra( + spw, + polpair, + Tsys, + blpairs=[blp], + form="Pk", + component="real", + scalar=scalar[(spw, polpair)], + )[blp_int] else: P_N = uvp.get_stats(precomp_P_N, key) - if 'P_N' in err_type: + if "P_N" in err_type: # set stats - uvp.set_stats('P_N', key, P_N) + uvp.set_stats("P_N", key, P_N) - if 'P_SN' in err_type: + if "P_SN" in err_type: # calculate P_SN: see Tan+2020 and # H1C_IDR2/notebooks/validation/errorbars_with_systematics_and_noise.ipynb # get signal proxy @@ -1457,13 +1644,14 @@ def uvp_noise_error(uvp, auto_Tsys=None, err_type='P_N', precomp_P_N=None, P_SN_ # catch nans, set to inf P_SN[np.isnan(P_SN)] = np.inf # set stats - uvp.set_stats('P_SN', key, P_SN) + uvp.set_stats("P_SN", key, P_SN) # P_SN correction if P_SN_correction and "P_SN" in err_type: if precomp_P_N is None: - precomp_P_N = 'P_N' - apply_P_SN_correction(uvp, P_SN='P_SN', P_N=precomp_P_N) + precomp_P_N = "P_N" + apply_P_SN_correction(uvp, P_SN="P_SN", P_N=precomp_P_N) + def uvp_noise_error_parser(): """ @@ -1474,24 +1662,49 @@ def uvp_noise_error_parser(): Returns: a: argparser object with arguments used in auto_noise_run.py. """ - a = argparse.ArgumentParser(description="argument parser for computing " - "thermal noise error bars from " - "autocorrelations") - a.add_argument("pspec_container", type=str, - help="Filename of HDF5 container (PSpecContainer) containing " - "input power spectra.") - a.add_argument("auto_file", type=str, help="Filename of UVData object containing only autocorr baselines to use" - "in thermal noise error bar estimation.") + a = argparse.ArgumentParser( + description="argument parser for computing " + "thermal noise error bars from " + "autocorrelations" + ) + a.add_argument( + "pspec_container", + type=str, + help="Filename of HDF5 container (PSpecContainer) containing " + "input power spectra.", + ) + a.add_argument( + "auto_file", + type=str, + help="Filename of UVData object containing only autocorr baselines to use" + "in thermal noise error bar estimation.", + ) a.add_argument("beam", type=str, help="Filename for UVBeam storing primary beam.") - a.add_argument("--groups", type=str, help="Name of power-spectrum group to compute noise for.", default=None, nargs="+") - a.add_argument("--spectra", default=None, type=str, nargs='+', - help="List of power spectra names (with group prefix) to calculate noise for.") - a.add_argument("--err_type", default="P_N", type=str, - nargs="+", help="Which components of noise error" - "to compute, 'P_N' or 'P_SN'") + a.add_argument( + "--groups", + type=str, + help="Name of power-spectrum group to compute noise for.", + default=None, + nargs="+", + ) + a.add_argument( + "--spectra", + default=None, + type=str, + nargs="+", + help="List of power spectra names (with group prefix) to calculate noise for.", + ) + a.add_argument( + "--err_type", + default="P_N", + type=str, + nargs="+", + help="Which components of noise error" "to compute, 'P_N' or 'P_SN'", + ) return a -def apply_P_SN_correction(uvp, P_SN='P_SN', P_N='P_N'): + +def apply_P_SN_correction(uvp, P_SN="P_SN", P_N="P_N"): """ Apply correction factor to P_SN errorbar in stats_array to account for double counting of noise by using data as proxy for signal. @@ -1513,7 +1726,9 @@ def apply_P_SN_correction(uvp, P_SN='P_SN', P_N='P_N'): p_n = uvp.stats_array[P_N][spw] p_sn = uvp.stats_array[P_SN][spw] # derive correction - corr = 1 - (np.sqrt(1 / np.sqrt(np.pi) + 1) - 1) * p_n.real / p_sn.real.clip(1e-40, np.inf) + corr = 1 - (np.sqrt(1 / np.sqrt(np.pi) + 1) - 1) * p_n.real / p_sn.real.clip( + 1e-40, np.inf + ) corr[np.isclose(corr, 0)] = np.inf corr[corr < 0] = np.inf corr[np.isnan(corr)] = np.inf @@ -1521,7 +1736,7 @@ def apply_P_SN_correction(uvp, P_SN='P_SN', P_N='P_N'): uvp.stats_array[P_SN][spw] *= corr -def history_string(notes=''): +def history_string(notes=""): """ Creates a standardized history string that all functions that write to disk can use. Optionally add notes. @@ -1536,4 +1751,3 @@ def history_string(notes=''): ------------ """ return history - diff --git a/hera_pspec/uvpspec.py b/hera_pspec/uvpspec.py index a640f260..fdbeabf5 100644 --- a/hera_pspec/uvpspec.py +++ b/hera_pspec/uvpspec.py @@ -6,10 +6,19 @@ import warnings import json -from . import conversions, noise, version, pspecbeam, grouping, utils, uvpspec_utils as uvputils +from . import ( + conversions, + noise, + version, + pspecbeam, + grouping, + utils, + uvpspec_utils as uvputils, +) from .parameter import PSpecParam from .uvwindow import UVWindow + class UVPSpec(object): """ An object for storing power spectra generated by hera_pspec and a file-format @@ -22,147 +31,487 @@ def __init__(self): by hera_pspec. """ # Summary attributes - self._Ntimes = PSpecParam("Ntimes", description="Number of unique times.", expected_type=int) - self._Nblpairts = PSpecParam("Nblpairts", description="Total number of baseline-pair times.", expected_type=int) - self._Nblpairs = PSpecParam("Nblpairs", description='Total number of baseline-pairs.', expected_type=int) - self._Nspwdlys = PSpecParam("Nspwdlys", description="Total number of delay bins across all spectral windows.", expected_type=int) - self._Nspwfreqs = PSpecParam("Nspwfreqs", description="Total number of frequency bins across all spectral windows.", expected_type=int) - self._Nspws = PSpecParam("Nspws", description="Number of unique spectral windows.", expected_type=int) - self._Ndlys = PSpecParam("Ndlys", description="Number of unique delay bins.", expected_type=int) - self._Nfreqs = PSpecParam("Nfreqs", description="Number of unique frequency bins in the data.", expected_type=int) - self._Npols = PSpecParam("Npols", description="Number of polarizations in the data.", expected_type=int) - self._history = PSpecParam("history", description='The file history.', expected_type=str) - desc = ('A compressed json string version of the r_params dictionary of pspecdata object.' - 'uvpspec only stores each unique weighting along with a list of unique baselines' - 'corresponding to that weighting.') - self._r_params = PSpecParam("r_params", description = desc, expected_type = str) + self._Ntimes = PSpecParam( + "Ntimes", description="Number of unique times.", expected_type=int + ) + self._Nblpairts = PSpecParam( + "Nblpairts", + description="Total number of baseline-pair times.", + expected_type=int, + ) + self._Nblpairs = PSpecParam( + "Nblpairs", description="Total number of baseline-pairs.", expected_type=int + ) + self._Nspwdlys = PSpecParam( + "Nspwdlys", + description="Total number of delay bins across all spectral windows.", + expected_type=int, + ) + self._Nspwfreqs = PSpecParam( + "Nspwfreqs", + description="Total number of frequency bins across all spectral windows.", + expected_type=int, + ) + self._Nspws = PSpecParam( + "Nspws", description="Number of unique spectral windows.", expected_type=int + ) + self._Ndlys = PSpecParam( + "Ndlys", description="Number of unique delay bins.", expected_type=int + ) + self._Nfreqs = PSpecParam( + "Nfreqs", + description="Number of unique frequency bins in the data.", + expected_type=int, + ) + self._Npols = PSpecParam( + "Npols", + description="Number of polarizations in the data.", + expected_type=int, + ) + self._history = PSpecParam( + "history", description="The file history.", expected_type=str + ) + desc = ( + "A compressed json string version of the r_params dictionary of pspecdata object." + "uvpspec only stores each unique weighting along with a list of unique baselines" + "corresponding to that weighting." + ) + self._r_params = PSpecParam("r_params", description=desc, expected_type=str) # Data attributes desc = "A string indicating the covariance model of cov_array. Options are ['dsets', 'empirical']. See PSpecData.pspec() for details." self._cov_model = PSpecParam("cov_model", description=desc, expected_type=str) desc = "A boolean indicating if the window functions stored in window_function_array are exact or not." - self._exact_windows = PSpecParam("exact_windows", description=desc, expected_type=bool) + self._exact_windows = PSpecParam( + "exact_windows", description=desc, expected_type=bool + ) desc = "Power spectrum data dictionary with spw integer as keys and values as complex ndarrays." - self._data_array = PSpecParam("data_array", description=desc, expected_type=np.complex128, form="(Nblpairts, spw_Ndlys, Npols)") + self._data_array = PSpecParam( + "data_array", + description=desc, + expected_type=np.complex128, + form="(Nblpairts, spw_Ndlys, Npols)", + ) desc = "Power spectrum covariance dictionary with spw integer as keys and values as float ndarrays, stored separately for real and imaginary parts." - self._cov_array_real = PSpecParam("cov_array_real", description=desc, expected_type=np.float64, form="(Nblpairts, spw_Ndlys, spw_Ndlys, Npols)") - self._cov_array_imag = PSpecParam("cov_array_imag", description=desc, expected_type=np.float64, form="(Nblpairts, spw_Ndlys, spw_Ndlys, Npols)") + self._cov_array_real = PSpecParam( + "cov_array_real", + description=desc, + expected_type=np.float64, + form="(Nblpairts, spw_Ndlys, spw_Ndlys, Npols)", + ) + self._cov_array_imag = PSpecParam( + "cov_array_imag", + description=desc, + expected_type=np.float64, + form="(Nblpairts, spw_Ndlys, spw_Ndlys, Npols)", + ) desc = "Window function dictionary of bandpowers." - self._window_function_array = PSpecParam("window_function_array", description=desc, expected_type=np.float64, form="(Nblpairts, spw_Ndlys, spw_Ndlys, Npols)") + self._window_function_array = PSpecParam( + "window_function_array", + description=desc, + expected_type=np.float64, + form="(Nblpairts, spw_Ndlys, spw_Ndlys, Npols)", + ) desc = "Dictionary of bandpowers given the kperp grid used to compute the window functions." - self._window_function_kperp = PSpecParam("window_function_kperp", description=desc, expected_type=np.float64, form="(Nblpairts, spw_Ndlys, spw_Ndlys, Npols)") + self._window_function_kperp = PSpecParam( + "window_function_kperp", + description=desc, + expected_type=np.float64, + form="(Nblpairts, spw_Ndlys, spw_Ndlys, Npols)", + ) desc = "Dictionary of bandpowers given the kparallel grid used to compute the window functions." - self._window_function_kpara = PSpecParam("window_function_kpara", description=desc, expected_type=np.float64, form="(Nblpairts, spw_Ndlys, spw_Ndlys, Npols)") + self._window_function_kpara = PSpecParam( + "window_function_kpara", + description=desc, + expected_type=np.float64, + form="(Nblpairts, spw_Ndlys, spw_Ndlys, Npols)", + ) desc = "Weight dictionary for original two datasets. The second axis holds [dset1_wgts, dset2_wgts] in that order." - self._wgt_array = PSpecParam("wgt_array", description=desc, expected_type=np.float64, form="(Nblpairts, spw_Nfreqs, 2, Npols)") - desc = "Integration time dictionary. This holds the average integration time [seconds] of each delay spectrum in the data. " \ - "This is not necessarily equal to the integration time of the visibility data: If data have been coherently averaged " \ - "(i.e. averaged before squaring), than this is the sum of each spectrum's integration time." - self._integration_array = PSpecParam("integration_array", description=desc, expected_type=np.float64, form="(Nblpairts, Npols)") - desc = "Nsample dictionary, if the pspectra have been incoherently averaged (i.e. averaged after squaring), this is " \ - "the effective number of samples in that average (float type). This is not the same as the pyuvdata.UVData nsample_array." - self._nsample_array = PSpecParam("nsample_array", description=desc, expected_type=np.float64, form="(Nblpairts, Npols)") - desc = ("Power spectrum stats array with stats type and spw integer as keys and values as complex ndarrays with same shape" - " as data_array") - self._stats_array = PSpecParam("stats_array", description=desc, expected_type=np.complex128, form="(Nblpairts, Ndlys, Npols)") - self._spw_array = PSpecParam("spw_array", description="Integer array holding unique spectral windows.", form="(Nspws,)", expected_type=np.uint16) - self._spw_dly_array = PSpecParam("spw_dly_array", description="Spw integer array for the dly_array.", form="(Nspwdlys,)", expected_type=np.uint16) - self._spw_freq_array = PSpecParam("spw_freq_array", description="Spw integer array for the freq_array.", form="(Nspwfreqs,)", expected_type=np.uint16) - self._freq_array = PSpecParam("freq_array", description="Frequency array of the original data in Hz.", form="(Nfreqs,)", expected_type=np.float64) - self._dly_array = PSpecParam("dly_array", description="Delay array in seconds.", form="(Nspwdlys,)", expected_type=np.float64) + self._wgt_array = PSpecParam( + "wgt_array", + description=desc, + expected_type=np.float64, + form="(Nblpairts, spw_Nfreqs, 2, Npols)", + ) + desc = ( + "Integration time dictionary. This holds the average integration time [seconds] of each delay spectrum in the data. " + "This is not necessarily equal to the integration time of the visibility data: If data have been coherently averaged " + "(i.e. averaged before squaring), than this is the sum of each spectrum's integration time." + ) + self._integration_array = PSpecParam( + "integration_array", + description=desc, + expected_type=np.float64, + form="(Nblpairts, Npols)", + ) + desc = ( + "Nsample dictionary, if the pspectra have been incoherently averaged (i.e. averaged after squaring), this is " + "the effective number of samples in that average (float type). This is not the same as the pyuvdata.UVData nsample_array." + ) + self._nsample_array = PSpecParam( + "nsample_array", + description=desc, + expected_type=np.float64, + form="(Nblpairts, Npols)", + ) + desc = ( + "Power spectrum stats array with stats type and spw integer as keys and values as complex ndarrays with same shape" + " as data_array" + ) + self._stats_array = PSpecParam( + "stats_array", + description=desc, + expected_type=np.complex128, + form="(Nblpairts, Ndlys, Npols)", + ) + self._spw_array = PSpecParam( + "spw_array", + description="Integer array holding unique spectral windows.", + form="(Nspws,)", + expected_type=np.uint16, + ) + self._spw_dly_array = PSpecParam( + "spw_dly_array", + description="Spw integer array for the dly_array.", + form="(Nspwdlys,)", + expected_type=np.uint16, + ) + self._spw_freq_array = PSpecParam( + "spw_freq_array", + description="Spw integer array for the freq_array.", + form="(Nspwfreqs,)", + expected_type=np.uint16, + ) + self._freq_array = PSpecParam( + "freq_array", + description="Frequency array of the original data in Hz.", + form="(Nfreqs,)", + expected_type=np.float64, + ) + self._dly_array = PSpecParam( + "dly_array", + description="Delay array in seconds.", + form="(Nspwdlys,)", + expected_type=np.float64, + ) desc = "Polarization pair integer, made up of two polarization integers concatenated in a standardized way." - self._polpair_array = PSpecParam("polpair_array", description=desc, form="(Npols,)", expected_type=np.int32) - self._lst_1_array = PSpecParam("lst_1_array", description="LST array of the first bl in the bl-pair [radians].", form="(Nblpairts,)", expected_type=np.float64) - self._lst_2_array = PSpecParam("lst_2_array", description="LST array of the second bl in the bl-pair [radians].", form="(Nblpairts,)", expected_type=np.float64) - self._lst_avg_array = PSpecParam("lst_avg_array", description="Average of the lst_1_array and lst_2_array [radians].", form="(Nblpairts,)", expected_type=np.float64) - self._time_1_array = PSpecParam("time_1_array", description="Time array of the first bl in the bl-pair [Julian Date].", form="(Nblpairts,)", expected_type=np.float64) - self._time_1_array = PSpecParam("time_2_array", description="Time array of the second bl in the bl-pair [Julian Date].", form="(Nblpairts,)", expected_type=np.float64) - self._time_avg_array = PSpecParam("time_avg_array", description="Average of the time_1_array and time_2_array [Julian Date].", form='(Nblpairts,)', expected_type=np.float64) - self._blpair_array = PSpecParam("blpair_array", description="Baseline-pair integer for all baseline-pair times.", form="(Nblpairts,)", expected_type=np.int64) - self._scalar_array = PSpecParam("scalar_array", description="Power spectrum normalization scalar from pspecbeam module.", expected_type=np.float64, form="(Nspws, Npols)") + self._polpair_array = PSpecParam( + "polpair_array", description=desc, form="(Npols,)", expected_type=np.int32 + ) + self._lst_1_array = PSpecParam( + "lst_1_array", + description="LST array of the first bl in the bl-pair [radians].", + form="(Nblpairts,)", + expected_type=np.float64, + ) + self._lst_2_array = PSpecParam( + "lst_2_array", + description="LST array of the second bl in the bl-pair [radians].", + form="(Nblpairts,)", + expected_type=np.float64, + ) + self._lst_avg_array = PSpecParam( + "lst_avg_array", + description="Average of the lst_1_array and lst_2_array [radians].", + form="(Nblpairts,)", + expected_type=np.float64, + ) + self._time_1_array = PSpecParam( + "time_1_array", + description="Time array of the first bl in the bl-pair [Julian Date].", + form="(Nblpairts,)", + expected_type=np.float64, + ) + self._time_1_array = PSpecParam( + "time_2_array", + description="Time array of the second bl in the bl-pair [Julian Date].", + form="(Nblpairts,)", + expected_type=np.float64, + ) + self._time_avg_array = PSpecParam( + "time_avg_array", + description="Average of the time_1_array and time_2_array [Julian Date].", + form="(Nblpairts,)", + expected_type=np.float64, + ) + self._blpair_array = PSpecParam( + "blpair_array", + description="Baseline-pair integer for all baseline-pair times.", + form="(Nblpairts,)", + expected_type=np.int64, + ) + self._scalar_array = PSpecParam( + "scalar_array", + description="Power spectrum normalization scalar from pspecbeam module.", + expected_type=np.float64, + form="(Nspws, Npols)", + ) # Baseline attributes - self._Nbls = PSpecParam("Nbls", description="Number of unique baseline integers.", expected_type=int) - self._bl_vecs = PSpecParam("bl_vecs", description="ndarray of baseline separation vectors in the ITRF frame [meters]. To get it in ENU frame see self.get_ENU_bl_vecs().", expected_type=np.float64, form="(Nbls,)") - self._bl_array = PSpecParam("bl_array", description="All unique baseline (antenna-pair) integers.", expected_type=np.int32, form="(Nbls,)") + self._Nbls = PSpecParam( + "Nbls", description="Number of unique baseline integers.", expected_type=int + ) + self._bl_vecs = PSpecParam( + "bl_vecs", + description="ndarray of baseline separation vectors in the ITRF frame [meters]. To get it in ENU frame see self.get_ENU_bl_vecs().", + expected_type=np.float64, + form="(Nbls,)", + ) + self._bl_array = PSpecParam( + "bl_array", + description="All unique baseline (antenna-pair) integers.", + expected_type=np.int32, + form="(Nbls,)", + ) # Misc Attributes - self._channel_width = PSpecParam("channel_width", description="width of visibility frequency channels in Hz.", expected_type=float) - self._telescope_location = PSpecParam("telescope_location", description="telescope location in ECEF frame [meters]. To get it in Lat/Lon/Alt see pyuvdata.utils.LatLonAlt_from_XYZ().", expected_type=np.float64) - self._weighting = PSpecParam("weighting", description="Form of data weighting used when forming power spectra.", expected_type=str) - self.set_symmetric_taper = PSpecParam("symmetric_taper", description="Specify whether Taper was applied symmetrically (True) or to the left(False).", expected_type=str) - self._norm = PSpecParam("norm", description="Normalization method adopted in OQE (M matrix).", expected_type=str) - self._taper = PSpecParam("taper", description='Taper function applied to visibility data before FT. See uvtools.dspec.gen_window for options."', expected_type=str) - self._vis_units = PSpecParam("vis_units", description="Units of the original visibility data used to form the power spectra.", expected_type=str) - self._norm_units = PSpecParam("norm_units", description="Power spectra normalization units, i.e. telescope units [Hz str] or cosmological [(h^-3) Mpc^3].", expected_type=str) - self._labels = PSpecParam("labels", description="Array of dataset string labels.", expected_type=np.str) - self._label_1_array = PSpecParam("label_1_array", description="Integer array w/ shape of data that indexes labels and gives label of dset1.", form="(Nspws, Nblpairts, Npols)", expected_type=np.int32) - self._label_2_array = PSpecParam("label_2_array", description="Integer array w/ shape of data that indexes labels and gives label of dset2.", form="(Nspws, Nblpairts, Npols)", expected_type=np.int32) - self._folded = PSpecParam("folded", description="if power spectra are folded (i.e. averaged) onto purely positive delay axis. Default is False", expected_type=bool) - self._beamfile = PSpecParam("beamfile", description="filename of beam-model used to normalized pspectra.", expected_type=str) - self._OmegaP = PSpecParam("OmegaP", description="Integral of unitless beam power over the sky [steradians].", form="(Nbeam_freqs, Npols)", expected_type=np.float64) - self._OmegaPP = PSpecParam("OmegaP", description="Integral of unitless beam power squared over the sky [steradians].", form="(Nbeam_freqs, Npols)", expected_type=np.float64) - self._beam_freqs = PSpecParam("beam_freqs", description="Frequency bins of the OmegaP and OmegaPP beam-integral arrays [Hz].", form="(Nbeam_freqs,)", expected_type=np.float64) - self._cosmo = PSpecParam("cosmo", description="Instance of conversion.Cosmo_Conversions class.", expected_type=conversions.Cosmo_Conversions) + self._channel_width = PSpecParam( + "channel_width", + description="width of visibility frequency channels in Hz.", + expected_type=float, + ) + self._telescope_location = PSpecParam( + "telescope_location", + description="telescope location in ECEF frame [meters]. To get it in Lat/Lon/Alt see pyuvdata.utils.LatLonAlt_from_XYZ().", + expected_type=np.float64, + ) + self._weighting = PSpecParam( + "weighting", + description="Form of data weighting used when forming power spectra.", + expected_type=str, + ) + self.set_symmetric_taper = PSpecParam( + "symmetric_taper", + description="Specify whether Taper was applied symmetrically (True) or to the left(False).", + expected_type=str, + ) + self._norm = PSpecParam( + "norm", + description="Normalization method adopted in OQE (M matrix).", + expected_type=str, + ) + self._taper = PSpecParam( + "taper", + description='Taper function applied to visibility data before FT. See uvtools.dspec.gen_window for options."', + expected_type=str, + ) + self._vis_units = PSpecParam( + "vis_units", + description="Units of the original visibility data used to form the power spectra.", + expected_type=str, + ) + self._norm_units = PSpecParam( + "norm_units", + description="Power spectra normalization units, i.e. telescope units [Hz str] or cosmological [(h^-3) Mpc^3].", + expected_type=str, + ) + self._labels = PSpecParam( + "labels", + description="Array of dataset string labels.", + expected_type=np.str, + ) + self._label_1_array = PSpecParam( + "label_1_array", + description="Integer array w/ shape of data that indexes labels and gives label of dset1.", + form="(Nspws, Nblpairts, Npols)", + expected_type=np.int32, + ) + self._label_2_array = PSpecParam( + "label_2_array", + description="Integer array w/ shape of data that indexes labels and gives label of dset2.", + form="(Nspws, Nblpairts, Npols)", + expected_type=np.int32, + ) + self._folded = PSpecParam( + "folded", + description="if power spectra are folded (i.e. averaged) onto purely positive delay axis. Default is False", + expected_type=bool, + ) + self._beamfile = PSpecParam( + "beamfile", + description="filename of beam-model used to normalized pspectra.", + expected_type=str, + ) + self._OmegaP = PSpecParam( + "OmegaP", + description="Integral of unitless beam power over the sky [steradians].", + form="(Nbeam_freqs, Npols)", + expected_type=np.float64, + ) + self._OmegaPP = PSpecParam( + "OmegaP", + description="Integral of unitless beam power squared over the sky [steradians].", + form="(Nbeam_freqs, Npols)", + expected_type=np.float64, + ) + self._beam_freqs = PSpecParam( + "beam_freqs", + description="Frequency bins of the OmegaP and OmegaPP beam-integral arrays [Hz].", + form="(Nbeam_freqs,)", + expected_type=np.float64, + ) + self._cosmo = PSpecParam( + "cosmo", + description="Instance of conversion.Cosmo_Conversions class.", + expected_type=conversions.Cosmo_Conversions, + ) # Collect all parameters: required and non-required - self._all_params = sorted( [ p[1:] for p in - fnmatch.filter(self.__dict__.keys(), '_*')]) + self._all_params = sorted( + [p[1:] for p in fnmatch.filter(self.__dict__.keys(), "_*")] + ) # Specify required params: these are required for read / write and # self.check() - self._req_params = ["Ntimes", "Nblpairts", "Nblpairs", - "Nspws", "Ndlys", "Npols", "Nfreqs", "history", - "Nspwdlys", "Nspwfreqs", "r_params", - "data_array", "wgt_array", "integration_array", - "spw_array", "freq_array", "dly_array", - "polpair_array", - "lst_1_array", "lst_2_array", "time_1_array", - "time_2_array", "blpair_array", "Nbls", - "bl_vecs", "bl_array", "channel_width", - "telescope_location", "weighting", "vis_units", - "norm_units", "taper", "norm", "nsample_array", - "lst_avg_array", "time_avg_array", "folded", - "scalar_array", "labels", "label_1_array", - "label_2_array", "spw_dly_array", "spw_freq_array", - "symmetric_taper"] + self._req_params = [ + "Ntimes", + "Nblpairts", + "Nblpairs", + "Nspws", + "Ndlys", + "Npols", + "Nfreqs", + "history", + "Nspwdlys", + "Nspwfreqs", + "r_params", + "data_array", + "wgt_array", + "integration_array", + "spw_array", + "freq_array", + "dly_array", + "polpair_array", + "lst_1_array", + "lst_2_array", + "time_1_array", + "time_2_array", + "blpair_array", + "Nbls", + "bl_vecs", + "bl_array", + "channel_width", + "telescope_location", + "weighting", + "vis_units", + "norm_units", + "taper", + "norm", + "nsample_array", + "lst_avg_array", + "time_avg_array", + "folded", + "scalar_array", + "labels", + "label_1_array", + "label_2_array", + "spw_dly_array", + "spw_freq_array", + "symmetric_taper", + ] # All parameters must fall into one and only one of the following # groups, which are used in __eq__ - self._immutables = ["Ntimes", "Nblpairts", "Nblpairs", "Nspwdlys", - "Nspwfreqs", "Nspws", "Ndlys", "Npols", "Nfreqs", - "history", "r_params", "cov_model", - "Nbls", "channel_width", "weighting", "vis_units", - "norm", "norm_units", "taper", "cosmo", "beamfile", - 'folded', 'exact_windows'] - self._ndarrays = ["spw_array", "freq_array", "dly_array", - "polpair_array", "lst_1_array", "lst_avg_array", - "time_avg_array", - "lst_2_array", "time_1_array", "time_2_array", - "blpair_array", "OmegaP", "OmegaPP", "beam_freqs", - "bl_vecs", "bl_array", "telescope_location", - "scalar_array", "labels", "label_1_array", - "label_2_array", "spw_freq_array", "spw_dly_array"] - self._dicts = ["data_array", "wgt_array", "integration_array", "window_function_array", - "window_function_kperp", "window_function_kpara", - "nsample_array", "cov_array_real", "cov_array_imag"] + self._immutables = [ + "Ntimes", + "Nblpairts", + "Nblpairs", + "Nspwdlys", + "Nspwfreqs", + "Nspws", + "Ndlys", + "Npols", + "Nfreqs", + "history", + "r_params", + "cov_model", + "Nbls", + "channel_width", + "weighting", + "vis_units", + "norm", + "norm_units", + "taper", + "cosmo", + "beamfile", + "folded", + "exact_windows", + ] + self._ndarrays = [ + "spw_array", + "freq_array", + "dly_array", + "polpair_array", + "lst_1_array", + "lst_avg_array", + "time_avg_array", + "lst_2_array", + "time_1_array", + "time_2_array", + "blpair_array", + "OmegaP", + "OmegaPP", + "beam_freqs", + "bl_vecs", + "bl_array", + "telescope_location", + "scalar_array", + "labels", + "label_1_array", + "label_2_array", + "spw_freq_array", + "spw_dly_array", + ] + self._dicts = [ + "data_array", + "wgt_array", + "integration_array", + "window_function_array", + "window_function_kperp", + "window_function_kpara", + "nsample_array", + "cov_array_real", + "cov_array_imag", + ] self._dicts_of_dicts = ["stats_array"] # define which attributes are considered meta data. Large attrs should # be constructed as datasets - self._meta_dsets = ["lst_1_array", "lst_2_array", "time_1_array", - "time_2_array", "blpair_array", "bl_vecs", - "bl_array", "lst_avg_array", "time_avg_array", - "OmegaP", "OmegaPP", "label_1_array", - "label_2_array", "spw_dly_array", "spw_freq_array"] + self._meta_dsets = [ + "lst_1_array", + "lst_2_array", + "time_1_array", + "time_2_array", + "blpair_array", + "bl_vecs", + "bl_array", + "lst_avg_array", + "time_avg_array", + "OmegaP", + "OmegaPP", + "label_1_array", + "label_2_array", + "spw_dly_array", + "spw_freq_array", + ] self._meta_attrs = sorted( - set(self._all_params) - set(self._dicts) \ - - set(self._meta_dsets) - set(self._dicts_of_dicts)) + set(self._all_params) + - set(self._dicts) + - set(self._meta_dsets) + - set(self._dicts_of_dicts) + ) self._meta = sorted(set(self._meta_dsets).union(set(self._meta_attrs))) # check all params are covered - assert len( set(self._all_params) - set(self._dicts) \ - - set(self._immutables) - set(self._ndarrays) \ - - set(self._dicts_of_dicts) ) == 0 + assert ( + len( + set(self._all_params) + - set(self._dicts) + - set(self._immutables) + - set(self._ndarrays) + - set(self._dicts_of_dicts) + ) + == 0 + ) # Default parameter values self.folded = False @@ -171,7 +520,7 @@ def __init__(self): self.get_red_bls.__func__.__doc__ = uvputils._get_red_bls.__doc__ self.get_red_blpairs.__func__.__doc__ = uvputils._get_red_blpairs.__doc__ - def get_cov(self, key, component='real', omit_flags=False): + def get_cov(self, key, component="real", omit_flags=False): """ Slice into covariance array with a specified data key in the format (spw, ((ant1, ant2),(ant3, ant4)), (pol1, pol2)) @@ -204,20 +553,30 @@ def get_cov(self, key, component='real', omit_flags=False): # Need to deal with folded data! # if data has been folded, return only positive delays - if component == 'real': - if hasattr(self,'cov_array_real'): + if component == "real": + if hasattr(self, "cov_array_real"): if self.folded: Ndlys = np.count_nonzero(self.spw_dly_array == spw) - return self.cov_array_real[spw][blpairts, -(Ndlys-Ndlys//2-1):, -(Ndlys-Ndlys//2-1):, polpair] + return self.cov_array_real[spw][ + blpairts, + -(Ndlys - Ndlys // 2 - 1) :, + -(Ndlys - Ndlys // 2 - 1) :, + polpair, + ] else: return self.cov_array_real[spw][blpairts, :, :, polpair] else: raise AttributeError("No covariance array has been calculated.") - elif component == 'imag': - if hasattr(self,'cov_array_imag'): + elif component == "imag": + if hasattr(self, "cov_array_imag"): if self.folded: Ndlys = np.count_nonzero(self.spw_dly_array == spw) - return self.cov_array_imag[spw][blpairts, -(Ndlys-Ndlys//2-1):, -(Ndlys-Ndlys//2-1):, polpair] + return self.cov_array_imag[spw][ + blpairts, + -(Ndlys - Ndlys // 2 - 1) :, + -(Ndlys - Ndlys // 2 - 1) :, + polpair, + ] else: return self.cov_array_imag[spw][blpairts, :, :, polpair] else: @@ -257,13 +616,18 @@ def get_window_function(self, key, omit_flags=False): spw, blpairts, polpair = self.key_to_indices(key, omit_flags=omit_flags) if self.exact_windows: - return self.window_function_array[spw][blpairts, :, :, :, polpair] + return self.window_function_array[spw][blpairts, :, :, :, polpair] # Need to deal with folded data! # if data has been folded, return only positive delays if self.folded: Ndlys = np.count_nonzero(self.spw_dly_array == spw) - return self.window_function_array[spw][blpairts, -(Ndlys-Ndlys//2-1):, -(Ndlys-Ndlys//2-1):, polpair] + return self.window_function_array[spw][ + blpairts, + -(Ndlys - Ndlys // 2 - 1) :, + -(Ndlys - Ndlys // 2 - 1) :, + polpair, + ] else: return self.window_function_array[spw][blpairts, :, :, polpair] @@ -301,7 +665,7 @@ def get_data(self, key, omit_flags=False): # if data has been folded, return only positive delays if self.folded: Ndlys = np.count_nonzero(self.spw_dly_array == spw) - return self.data_array[spw][blpairts, -(Ndlys-Ndlys//2-1):, polpair] + return self.data_array[spw][blpairts, -(Ndlys - Ndlys // 2 - 1) :, polpair] # else return all delays else: @@ -393,7 +757,6 @@ def get_r_params(self): """ return uvputils.decompress_r_params(self.r_params) - def get_nsamples(self, key, omit_flags=False): """ Slice into nsample_array with a specified data key in the format @@ -455,8 +818,9 @@ def get_blpair_seps(self): # get list of bl separations bl_vecs = self.get_ENU_bl_vecs() bls = self.bl_array.tolist() - blseps = np.array([np.linalg.norm(bl_vecs[bls.index(bl)]) - for bl in self.bl_array]) + blseps = np.array( + [np.linalg.norm(bl_vecs[bls.index(bl)]) for bl in self.bl_array] + ) # construct empty blp_avg_sep array blp_avg_sep = np.empty(self.Nblpairts, np.float) @@ -464,7 +828,7 @@ def get_blpair_seps(self): # construct blpair_bls blpairs = _ordered_unique(self.blpair_array) bl1 = np.floor(blpairs / 1e6) - blpair_bls = np.vstack([bl1, blpairs - bl1*1e6]).astype(np.int32).T + blpair_bls = np.vstack([bl1, blpairs - bl1 * 1e6]).astype(np.int32).T # iterate over blpairs for blp, bl in zip(blpairs, blpair_bls): @@ -524,13 +888,13 @@ def get_kperps(self, spw, little_h=True): k_perp : float ndarray, transverse wave-vectors, shape=(Nblpairs,) """ # assert cosmo - assert hasattr(self, 'cosmo'), \ - "self.cosmo must exist to form cosmological " \ + assert hasattr(self, "cosmo"), ( + "self.cosmo must exist to form cosmological " "wave-vectors. See self.set_cosmology()" + ) # calculate mean redshift of band - avg_z = self.cosmo.f2z( - np.mean(self.freq_array[self.spw_to_freq_indices(spw)]) ) + avg_z = self.cosmo.f2z(np.mean(self.freq_array[self.spw_to_freq_indices(spw)])) # get kperps blpair_seps = self.get_blpair_seps() @@ -558,32 +922,34 @@ def get_kparas(self, spw, little_h=True): Radial wave-vectors, shape=(Ndlys given spw) """ # assert cosmo - assert hasattr(self, 'cosmo'), \ - "self.cosmo must exist to form cosmological " \ + assert hasattr(self, "cosmo"), ( + "self.cosmo must exist to form cosmological " "wave-vectors. See self.set_cosmology()" + ) # calculate mean redshift of band - avg_z = self.cosmo.f2z( - np.mean(self.freq_array[self.spw_to_freq_indices(spw)]) ) + avg_z = self.cosmo.f2z(np.mean(self.freq_array[self.spw_to_freq_indices(spw)])) # get kparas - k_para = self.get_dlys(spw) \ - * self.cosmo.tau_to_kpara(avg_z, little_h=little_h) + k_para = self.get_dlys(spw) * self.cosmo.tau_to_kpara(avg_z, little_h=little_h) return k_para def get_blpairs(self): """ Returns a list of all blpair tuples in the data_array. """ - return [self.blpair_to_antnums(blp) - for blp in _ordered_unique(self.blpair_array)] + return [ + self.blpair_to_antnums(blp) for blp in _ordered_unique(self.blpair_array) + ] def get_polpairs(self): """ Returns a list of all unique polarization pairs in the data_array. """ - polpairs = [uvputils.polpair_int2tuple(pp, pol_strings=True) - for pp in self.polpair_array] + polpairs = [ + uvputils.polpair_int2tuple(pp, pol_strings=True) + for pp in self.polpair_array + ] return sorted(set(polpairs)) def get_all_keys(self): @@ -637,7 +1003,9 @@ def get_spw_ranges(self, spws=None): Nfreqs = len(spw_freqs) spw_df = np.median(np.diff(spw_freqs)) Ndlys = len(self.spw_to_dly_indices(spw)) - spw_ranges.append( (spw_freqs.min(), spw_freqs.min() + Nfreqs * spw_df, Nfreqs, Ndlys) ) + spw_ranges.append( + (spw_freqs.min(), spw_freqs.min() + Nfreqs * spw_df, Nfreqs, Ndlys) + ) return spw_ranges @@ -667,7 +1035,9 @@ def get_stats(self, stat, key, omit_flags=False): if not hasattr(self, "stats_array"): raise AttributeError("No stats have been entered to this UVPSpec object") - assert stat in self.stats_array.keys(), "Statistic name {} not found in stat keys.".format(stat) + assert ( + stat in self.stats_array.keys() + ), "Statistic name {} not found in stat keys.".format(stat) spw, blpairts, polpair = self.key_to_indices(key, omit_flags=omit_flags) statistic = self.stats_array[stat] @@ -675,7 +1045,7 @@ def get_stats(self, stat, key, omit_flags=False): # if bandpowers have been folded, return only positive delays if self.folded: Ndlys = np.count_nonzero(self.spw_dly_array == spw) - return statistic[spw][blpairts, -(Ndlys-Ndlys//2-1):, polpair] + return statistic[spw][blpairts, -(Ndlys - Ndlys // 2 - 1) :, polpair] else: return statistic[spw][blpairts, :, polpair] @@ -702,8 +1072,9 @@ def set_stats(self, stat, key, statistic): if data_shape != statistic.shape: print(data_shape, statistic.shape) - errmsg = "Input array shape {} must match " \ - "data_array shape {}.".format(statistic.shape, data_shape) + errmsg = "Input array shape {} must match " "data_array shape {}.".format( + statistic.shape, data_shape + ) raise ValueError(errmsg) if not hasattr(self, "stats_array"): @@ -711,10 +1082,12 @@ def set_stats(self, stat, key, statistic): dtype = statistic.dtype if stat not in self.stats_array.keys(): - self.stats_array[stat] = \ - odict([[i, - np.nan * np.ones(self.data_array[i].shape, dtype=dtype)] - for i in range(self.Nspws)]) + self.stats_array[stat] = odict( + [ + [i, np.nan * np.ones(self.data_array[i].shape, dtype=dtype)] + for i in range(self.Nspws) + ] + ) self.stats_array[stat][spw][blpairts, :, polpair] = statistic @@ -756,7 +1129,6 @@ def set_stats_slice(self, stat, m, b, above=False, val=1e20): for tind in tinds: self.stats_array[stat][i][tind, dinds, :] = val - def convert_to_deltasq(self, inplace=True): """ Convert from P(k) to Delta^2(k) by multiplying by k^3 / (2pi^2). @@ -782,14 +1154,14 @@ def convert_to_deltasq(self, inplace=True): uvp = copy.deepcopy(self) # determine if little_h is required - little_h = 'h^-3' in uvp.norm_units + little_h = "h^-3" in uvp.norm_units # loop over spectral windows for spw in range(uvp.Nspws): # get k vectors k_perp = uvp.get_kperps(spw, little_h=little_h) k_para = uvp.get_kparas(spw, little_h=little_h) - k_mag = np.sqrt(k_perp[:, None, None]**2 + k_para[None, :, None]**2) + k_mag = np.sqrt(k_perp[:, None, None] ** 2 + k_para[None, :, None] ** 2) # shape of (Nblpairts, spw_Ndlys, Npols) coeff = k_mag**3 / (2 * np.pi**2) @@ -797,16 +1169,20 @@ def convert_to_deltasq(self, inplace=True): uvp.data_array[spw] *= coeff # update stats array - if hasattr(uvp, 'stats_array'): + if hasattr(uvp, "stats_array"): for stat in uvp.stats_array: uvp.stats_array[stat][spw] *= coeff # update cov array - if hasattr(uvp, 'cov_array_real'): + if hasattr(uvp, "cov_array_real"): cov = uvp.cov_array_real[spw] - uvp.cov_array_real[spw] = np.einsum("tip,tijp,tjp->tijp", coeff, cov, coeff) + uvp.cov_array_real[spw] = np.einsum( + "tip,tijp,tjp->tijp", coeff, cov, coeff + ) cov = uvp.cov_array_imag[spw] - uvp.cov_array_imag[spw] = np.einsum("tip,tijp,tjp->tijp", coeff, cov, coeff) + uvp.cov_array_imag[spw] = np.einsum( + "tip,tijp,tjp->tijp", coeff, cov, coeff + ) # edit units uvp.norm_units += " k^3 / (2pi^2)" @@ -814,7 +1190,6 @@ def convert_to_deltasq(self, inplace=True): if inplace == False: return uvp - def blpair_to_antnums(self, blpair): """ Convert baseline-pair integer to nested tuple of antenna numbers. @@ -832,7 +1207,6 @@ def blpair_to_antnums(self, blpair): """ return uvputils._blpair_to_antnums(blpair) - def antnums_to_blpair(self, antnums): """ Convert nested tuple of antenna numbers to baseline-pair integer. @@ -850,7 +1224,6 @@ def antnums_to_blpair(self, antnums): """ return uvputils._antnums_to_blpair(antnums) - def bl_to_antnums(self, bl): """ Convert baseline (antenna-pair) integer to nested tuple of antenna @@ -868,7 +1241,6 @@ def bl_to_antnums(self, bl): """ return uvputils._bl_to_antnums(bl) - def antnums_to_bl(self, antnums): """ Convert tuple of antenna numbers to baseline integer. @@ -886,7 +1258,6 @@ def antnums_to_bl(self, antnums): """ return uvputils._antnums_to_bl(antnums) - def blpair_to_indices(self, blpair): """ Convert a baseline-pair nested tuple ((ant1, ant2), (ant3, ant4)) or @@ -908,11 +1279,12 @@ def blpair_to_indices(self, blpair): if isinstance(blpair[0], tuple): blpair = [self.antnums_to_blpair(blp) for blp in blpair] # assert exists in data - assert np.array([b in self.blpair_array for b in blpair]).all(), \ - "blpairs {} not all found in data".format(blpair) + assert np.array( + [b in self.blpair_array for b in blpair] + ).all(), "blpairs {} not all found in data".format(blpair) return np.arange(self.Nblpairts)[ - np.logical_or.reduce([self.blpair_array == b for b in blpair])] - + np.logical_or.reduce([self.blpair_array == b for b in blpair]) + ] def spw_to_freq_indices(self, spw): """ @@ -940,8 +1312,9 @@ def spw_to_freq_indices(self, spw): spw = [spw_ranges.index(s) for s in spw] # assert exists in data - assert np.array([s in self.spw_freq_array for s in spw]).all(), \ - "spws {} not all found in data".format(spw) + assert np.array( + [s in self.spw_freq_array for s in spw] + ).all(), "spws {} not all found in data".format(spw) # get select boolean array select = np.logical_or.reduce([self.spw_freq_array == s for s in spw]) @@ -950,7 +1323,6 @@ def spw_to_freq_indices(self, spw): freq_indices = np.arange(self.Nspwfreqs)[select] return freq_indices - def spw_to_dly_indices(self, spw): """ Convert a spectral window integer into a list of indices to index @@ -977,8 +1349,9 @@ def spw_to_dly_indices(self, spw): spw = [spw_ranges.index(s) for s in spw] # assert exists in data - assert np.array([s in self.spw_dly_array for s in spw]).all(), \ - "spws {} not all found in data".format(spw) + assert np.array( + [s in self.spw_dly_array for s in spw] + ).all(), "spws {} not all found in data".format(spw) # get select boolean array select = np.logical_or.reduce([self.spw_dly_array == s for s in spw]) @@ -1017,11 +1390,12 @@ def spw_indices(self, spw): spw = [spw_ranges.index(s) for s in spw] # assert exists in data - assert np.array([s in self.spw_array for s in spw]).all(), \ - "spws {} not all found in data".format(spw) + assert np.array( + [s in self.spw_array for s in spw] + ).all(), "spws {} not all found in data".format(spw) # get select boolean array - #select = reduce(operator.add, [self.spw_array == s for s in spw]) + # select = reduce(operator.add, [self.spw_array == s for s in spw]) select = np.logical_or.reduce([self.spw_array == s for s in spw]) # get array @@ -1029,7 +1403,6 @@ def spw_indices(self, spw): return spw_indices - def polpair_to_indices(self, polpair): """ Map a polarization-pair integer or tuple to its index in polpair_array. @@ -1050,27 +1423,33 @@ def polpair_to_indices(self, polpair): """ # Convert to list if not already a list if isinstance(polpair, (tuple, int, np.integer, str)): - polpair = [polpair,] + polpair = [ + polpair, + ] elif not isinstance(polpair, (list, np.ndarray)): raise TypeError("polpair must be list of tuple or int or str") # Convert strings to ints - polpair = [uvputils.polpair_tuple2int((p,p)) if isinstance(p, str) else p - for p in polpair] + polpair = [ + uvputils.polpair_tuple2int((p, p)) if isinstance(p, str) else p + for p in polpair + ] # Convert list items to standard format (polpair integers) - polpair = [uvputils.polpair_tuple2int(p) if isinstance(p, tuple) else p - for p in polpair] + polpair = [ + uvputils.polpair_tuple2int(p) if isinstance(p, tuple) else p + for p in polpair + ] # Ensure all pols exist in data - assert np.array([p in self.polpair_array for p in polpair]).all(), \ - "pols {} not all found in data".format(polpair) + assert np.array( + [p in self.polpair_array for p in polpair] + ).all(), "pols {} not all found in data".format(polpair) idxs = np.logical_or.reduce([self.polpair_array == pp for pp in polpair]) indices = np.arange(self.polpair_array.size)[idxs] return indices - def time_to_indices(self, time, blpairs=None): """ Convert a time [Julian Date] from self.time_avg_array to an array that @@ -1152,20 +1531,21 @@ def key_to_indices(self, key, omit_flags=False): # assert key length assert len(key) == 3, "length of key must be 3: (spw, blpair, polpair)" if isinstance(key, (odict, dict)): - key = (key['spw'], key['blpair'], key['polpair']) + key = (key["spw"], key["blpair"], key["polpair"]) # assign key elements spw_ind = key[0] blpair = key[1] polpair = key[2] - # assert types assert isinstance(spw_ind, (int, np.integer)), "spw must be an integer" - assert isinstance(blpair, (int, np.integer, tuple)), \ - "blpair must be an integer or nested tuple" - assert isinstance(polpair, (tuple, np.integer, int, str)), \ - "polpair must be a tuple, integer, or str: %s / %s" % (polpair, key) + assert isinstance( + blpair, (int, np.integer, tuple) + ), "blpair must be an integer or nested tuple" + assert isinstance( + polpair, (tuple, np.integer, int, str) + ), "polpair must be a tuple, integer, or str: %s / %s" % (polpair, key) # convert polpair string to tuple if isinstance(polpair, str): @@ -1180,18 +1560,20 @@ def key_to_indices(self, key, omit_flags=False): polpair = uvputils.polpair_tuple2int(polpair) # check attributes exist in data - assert spw_ind in self.spw_freq_array and spw_ind in self.spw_dly_array, \ - "spw {} not found in data".format(spw_ind) - assert blpair in self.blpair_array, \ - "blpair {} not found in data".format(blpair) - assert polpair in self.polpair_array, \ - "polpair {} not found in data".format(polpair) + assert ( + spw_ind in self.spw_freq_array and spw_ind in self.spw_dly_array + ), "spw {} not found in data".format(spw_ind) + assert blpair in self.blpair_array, "blpair {} not found in data".format(blpair) + assert polpair in self.polpair_array, "polpair {} not found in data".format( + polpair + ) # index polarization array polpair_ind = self.polpair_to_indices(polpair) if isinstance(polpair_ind, (list, np.ndarray)): - assert len(polpair_ind) == 1, \ - "only one polpair can be specified in key_to_indices" + assert ( + len(polpair_ind) == 1 + ), "only one polpair can be specified in key_to_indices" polpair_ind = polpair_ind[0] # index blpairts @@ -1205,10 +1587,18 @@ def key_to_indices(self, key, omit_flags=False): return spw_ind, blpairts_inds, polpair_ind - - def select(self, spws=None, bls=None, only_pairs_in_bls=True, blpairs=None, - times=None, lsts=None, polpairs=None, inplace=True, - run_check=True): + def select( + self, + spws=None, + bls=None, + only_pairs_in_bls=True, + blpairs=None, + times=None, + lsts=None, + polpairs=None, + inplace=True, + run_check=True, + ): """ Select function for selecting out certain slices of the data. @@ -1267,14 +1657,21 @@ def select(self, spws=None, bls=None, only_pairs_in_bls=True, blpairs=None, else: uvp = copy.deepcopy(self) - uvputils._select(uvp, spws=spws, bls=bls, - only_pairs_in_bls=only_pairs_in_bls, - blpairs=blpairs, times=times, lsts=lsts, - polpairs=polpairs) - - if run_check: uvp.check() - if inplace == False: return uvp - + uvputils._select( + uvp, + spws=spws, + bls=bls, + only_pairs_in_bls=only_pairs_in_bls, + blpairs=blpairs, + times=times, + lsts=lsts, + polpairs=polpairs, + ) + + if run_check: + uvp.check() + if inplace == False: + return uvp def get_ENU_bl_vecs(self): """ @@ -1287,13 +1684,22 @@ def get_ENU_bl_vecs(self): Array of baseline vectors. """ return uvutils.ENU_from_ECEF( - (self.bl_vecs + self.telescope_location), \ - *uvutils.LatLonAlt_from_XYZ(self.telescope_location[None])) - - - def read_from_group(self, grp, just_meta=False, spws=None, bls=None, - blpairs=None, times=None, lsts=None, polpairs=None, - only_pairs_in_bls=False): + (self.bl_vecs + self.telescope_location), + *uvutils.LatLonAlt_from_XYZ(self.telescope_location[None]) + ) + + def read_from_group( + self, + grp, + just_meta=False, + spws=None, + bls=None, + blpairs=None, + times=None, + lsts=None, + polpairs=None, + only_pairs_in_bls=False, + ): """ Clear current UVPSpec object and load in data from specified HDF5 group. @@ -1336,10 +1742,10 @@ def read_from_group(self, grp, just_meta=False, spws=None, bls=None, are found in the 'bls' list. Default: False. """ # Make sure the group is a UVPSpec object - assert 'pspec_type' in grp.attrs, "This object is not a UVPSpec object" - pstype = grp.attrs['pspec_type'] + assert "pspec_type" in grp.attrs, "This object is not a UVPSpec object" + pstype = grp.attrs["pspec_type"] pstype = pstype.decode() if isinstance(pstype, bytes) else pstype - assert pstype == 'UVPSpec', "This object is not a UVPSpec object" + assert pstype == "UVPSpec", "This object is not a UVPSpec object" # Clear all data in the current object self._clear() @@ -1348,48 +1754,74 @@ def read_from_group(self, grp, just_meta=False, spws=None, bls=None, for k in grp.attrs: if k in self._meta_attrs: val = grp.attrs[k] - if isinstance(val, bytes): val = val.decode() # bytes -> str + if isinstance(val, bytes): + val = val.decode() # bytes -> str setattr(self, k, val) for k in grp: if k in self._meta_dsets: setattr(self, k, grp[k][:]) # Backwards compatibility: pol_array exists (not polpair_array) - if 'pol_array' in grp.attrs: - warnings.warn("Stored UVPSpec contains pol_array attr, which has " - "been superseded by polpair_array. Converting " - "automatically.", UserWarning) - pol_arr = grp.attrs['pol_array'] + if "pol_array" in grp.attrs: + warnings.warn( + "Stored UVPSpec contains pol_array attr, which has " + "been superseded by polpair_array. Converting " + "automatically.", + UserWarning, + ) + pol_arr = grp.attrs["pol_array"] # Convert to polpair array - polpair_arr = [uvputils.polpair_tuple2int((p,p)) for p in pol_arr] - setattr(self, 'polpair_array', np.array(polpair_arr)) + polpair_arr = [uvputils.polpair_tuple2int((p, p)) for p in pol_arr] + setattr(self, "polpair_array", np.array(polpair_arr)) # Backwards compatibility: exact_windows - if 'exact_windows' not in grp.attrs: - setattr(self, 'exact_windows', False) + if "exact_windows" not in grp.attrs: + setattr(self, "exact_windows", False) # Use _select() to pick out only the requested baselines/spws if just_meta: - uvputils._select(self, spws=spws, bls=bls, lsts=lsts, - only_pairs_in_bls=only_pairs_in_bls, - blpairs=blpairs, times=times, polpairs=polpairs) + uvputils._select( + self, + spws=spws, + bls=bls, + lsts=lsts, + only_pairs_in_bls=only_pairs_in_bls, + blpairs=blpairs, + times=times, + polpairs=polpairs, + ) else: - uvputils._select(self, spws=spws, bls=bls, lsts=lsts, - only_pairs_in_bls=only_pairs_in_bls, - blpairs=blpairs, times=times, polpairs=polpairs, - h5file=grp) + uvputils._select( + self, + spws=spws, + bls=bls, + lsts=lsts, + only_pairs_in_bls=only_pairs_in_bls, + blpairs=blpairs, + times=times, + polpairs=polpairs, + h5file=grp, + ) # handle cosmo - if hasattr(self, 'cosmo'): - self.cosmo = conversions.Cosmo_Conversions( - **ast.literal_eval(self.cosmo) ) + if hasattr(self, "cosmo"): + self.cosmo = conversions.Cosmo_Conversions(**ast.literal_eval(self.cosmo)) self.check(just_meta=just_meta) - def read_hdf5(self, filepath, just_meta=False, spws=None, bls=None, - blpairs=None, times=None, lsts=None, polpairs=None, - only_pairs_in_bls=False): + def read_hdf5( + self, + filepath, + just_meta=False, + spws=None, + bls=None, + blpairs=None, + times=None, + lsts=None, + polpairs=None, + only_pairs_in_bls=False, + ): """ Clear current UVPSpec object and load in data from an HDF5 file. @@ -1433,11 +1865,17 @@ def read_hdf5(self, filepath, just_meta=False, spws=None, bls=None, are found in the 'bls' list. Default: False. """ # Open file descriptor and read data - with h5py.File(filepath, 'r') as f: - self.read_from_group(f, just_meta=just_meta, spws=spws, bls=bls, - times=times, lsts=lsts, polpairs=polpairs, - only_pairs_in_bls=only_pairs_in_bls) - + with h5py.File(filepath, "r") as f: + self.read_from_group( + f, + just_meta=just_meta, + spws=spws, + bls=bls, + times=times, + lsts=lsts, + polpairs=polpairs, + only_pairs_in_bls=only_pairs_in_bls, + ) def write_to_group(self, group, run_check=True): """ @@ -1455,7 +1893,8 @@ def write_to_group(self, group, run_check=True): """ # Run check - if run_check: self.check() + if run_check: + self.check() # Check whether the group already contains info # TODO @@ -1464,13 +1903,13 @@ def write_to_group(self, group, run_check=True): for k in self._meta_attrs: if hasattr(self, k): # handle cosmo - if k == 'cosmo': - group.attrs['cosmo'] = str(getattr(self, k).get_params()) + if k == "cosmo": + group.attrs["cosmo"] = str(getattr(self, k).get_params()) continue this_attr = getattr(self, k) # Do Unicode/string conversions, as HDF5 struggles with them - if k == 'labels': + if k == "labels": this_attr = [np.string_(lbl) for lbl in this_attr] # Store attribute in group @@ -1482,48 +1921,60 @@ def write_to_group(self, group, run_check=True): # Iterate over spectral windows and create datasets for i in np.unique(self.spw_array): - group.create_dataset("data_spw{}".format(i), - data=self.data_array[i], - dtype=np.complex128) - group.create_dataset("wgt_spw{}".format(i), - data=self.wgt_array[i], - dtype=np.float64) - group.create_dataset("integration_spw{}".format(i), - data=self.integration_array[i], - dtype=np.float64) - group.create_dataset("nsample_spw{}".format(i), - data=self.nsample_array[i], - dtype=np.float) + group.create_dataset( + "data_spw{}".format(i), data=self.data_array[i], dtype=np.complex128 + ) + group.create_dataset( + "wgt_spw{}".format(i), data=self.wgt_array[i], dtype=np.float64 + ) + group.create_dataset( + "integration_spw{}".format(i), + data=self.integration_array[i], + dtype=np.float64, + ) + group.create_dataset( + "nsample_spw{}".format(i), data=self.nsample_array[i], dtype=np.float + ) if hasattr(self, "window_function_array"): - group.create_dataset("window_function_spw{}".format(i), - data=self.window_function_array[i], - dtype=np.float64) + group.create_dataset( + "window_function_spw{}".format(i), + data=self.window_function_array[i], + dtype=np.float64, + ) if self.exact_windows: - group.create_dataset("window_function_kperp_spw{}".format(i), - data=self.window_function_kperp[i], - dtype=np.float64) - group.create_dataset("window_function_kpara_spw{}".format(i), - data=self.window_function_kpara[i], - dtype=np.float64) + group.create_dataset( + "window_function_kperp_spw{}".format(i), + data=self.window_function_kperp[i], + dtype=np.float64, + ) + group.create_dataset( + "window_function_kpara_spw{}".format(i), + data=self.window_function_kpara[i], + dtype=np.float64, + ) if hasattr(self, "cov_array_real"): - group.create_dataset("cov_real_spw{}".format(i), - data=self.cov_array_real[i], - dtype=np.float64) - group.create_dataset("cov_imag_spw{}".format(i), - data=self.cov_array_imag[i], - dtype=np.float64) + group.create_dataset( + "cov_real_spw{}".format(i), + data=self.cov_array_real[i], + dtype=np.float64, + ) + group.create_dataset( + "cov_imag_spw{}".format(i), + data=self.cov_array_imag[i], + dtype=np.float64, + ) # Store any statistics arrays if hasattr(self, "stats_array"): for s in self.stats_array.keys(): data = self.stats_array[s] for i in np.unique(self.spw_array): - group.create_dataset("stats_{}_{}".format(s, i), - data=data[i], dtype=data[i].dtype) + group.create_dataset( + "stats_{}_{}".format(s, i), data=data[i], dtype=data[i].dtype + ) # denote as a uvpspec object - group.attrs['pspec_type'] = self.__class__.__name__ - + group.attrs["pspec_type"] = self.__class__.__name__ def write_hdf5(self, filepath, overwrite=False, run_check=True): """ @@ -1548,12 +1999,10 @@ def write_hdf5(self, filepath, overwrite=False, run_check=True): os.remove(filepath) # Write file - with h5py.File(filepath, 'w') as f: + with h5py.File(filepath, "w") as f: self.write_to_group(f, run_check=run_check) - - def set_cosmology(self, new_cosmo, overwrite=False, new_beam=None, - verbose=True): + def set_cosmology(self, new_cosmo, overwrite=False, new_beam=None, verbose=True): """ Set the cosmology for this UVPSpec object by passing an instance of conversions.Cosmo_Conversions and assigning it to self.cosmo, in @@ -1591,23 +2040,27 @@ def set_cosmology(self, new_cosmo, overwrite=False, new_beam=None, verbose : bool, optional If True, report feedback to stdout. Default: True. """ - if hasattr(self, 'cosmo') and overwrite == False: + if hasattr(self, "cosmo") and overwrite == False: print("self.cosmo exists and overwrite == False, not overwriting...") return else: - if (not hasattr(self, 'OmegaP') \ - or not hasattr(self, 'OmegaPP') \ - or not hasattr(self, 'beam_freqs')) \ - and new_beam is None: - print("In order to set the cosmology, self.OmegaP, " \ - "self.OmegaPP and self.beam_freqs must exist, or you " \ - "need to pass a PSpecBeamUV object or a path to a beam " \ - "file as new_beam.") + if ( + not hasattr(self, "OmegaP") + or not hasattr(self, "OmegaPP") + or not hasattr(self, "beam_freqs") + ) and new_beam is None: + print( + "In order to set the cosmology, self.OmegaP, " + "self.OmegaPP and self.beam_freqs must exist, or you " + "need to pass a PSpecBeamUV object or a path to a beam " + "file as new_beam." + ) return # overwrite beam quantities if new_beam is not None: - if verbose: print("Updating beam data with {}".format(new_beam)) + if verbose: + print("Updating beam data with {}".format(new_beam)) if isinstance(new_beam, (str, np.str)): # PSpecBeamUV will adopt a default cosmology upon instantiation, # but this doesn't matterfor what we need from it @@ -1618,7 +2071,8 @@ def set_cosmology(self, new_cosmo, overwrite=False, new_beam=None, # Update cosmo if isinstance(new_cosmo, (dict, odict)): new_cosmo = conversions.Cosmo_Conversions(**new_cosmo) - if verbose: print("setting cosmology: \n{}".format(new_cosmo)) + if verbose: + print("setting cosmology: \n{}".format(new_cosmo)) self.cosmo = new_cosmo # Update scalar_array @@ -1626,11 +2080,12 @@ def set_cosmology(self, new_cosmo, overwrite=False, new_beam=None, print("Updating scalar array and re-normalizing power spectra") for spw in range(self.Nspws): for j, polpair in enumerate(self.polpair_array): - scalar = self.compute_scalar(spw, polpair, num_steps=1000, - little_h=True, noise_scalar=False) + scalar = self.compute_scalar( + spw, polpair, num_steps=1000, little_h=True, noise_scalar=False + ) # renormalize power spectra with new scalar - self.data_array[spw][:, :, j] *= scalar/self.scalar_array[spw, j] + self.data_array[spw][:, :, j] *= scalar / self.scalar_array[spw, j] # update self.scalar_array element self.scalar_array[spw, j] = scalar @@ -1639,9 +2094,15 @@ def set_cosmology(self, new_cosmo, overwrite=False, new_beam=None, if "Mpc" not in self.norm_units: self.norm_units = "h^-3 Mpc^3" - def get_exact_window_functions(self, ftbeam_file=None, spw_array=None, - verbose=False, inplace=True, add_to_history='', - x_orientation=None): + def get_exact_window_functions( + self, + ftbeam_file=None, + spw_array=None, + verbose=False, + inplace=True, + add_to_history="", + x_orientation=None, + ): """ Obtain the exact window functions corresponding to UVPSpec object. @@ -1661,7 +2122,7 @@ def get_exact_window_functions(self, ftbeam_file=None, spw_array=None, - '' for computation from beam simulations (slow) spw_array : list of ints, optional - Spectral window indices. If None, the window functions are computed on + Spectral window indices. If None, the window functions are computed on all the uvp.spw_ranges, successively. Default: None. verbose : bool, optional @@ -1694,15 +2155,22 @@ def get_exact_window_functions(self, ftbeam_file=None, spw_array=None, # if no spw specified, use attribute spw_array = self.spw_array else: - spw_array = spw_array if isinstance(spw_array, (list, tuple, np.ndarray)) else [int(spw_array)] + spw_array = ( + spw_array + if isinstance(spw_array, (list, tuple, np.ndarray)) + else [int(spw_array)] + ) # check if spw given is in uvp - assert np.all([spw in self.spw_array for spw in spw_array]), \ - "input spw is not in UVPSpec.spw_array." + assert np.all( + [spw in self.spw_array for spw in spw_array] + ), "input spw is not in UVPSpec.spw_array." if inplace: # set inplace to False inplace = False - warnings.warn('inplace set to False because you are not considering' \ - ' all spectral windows in object.') + warnings.warn( + "inplace set to False because you are not considering" + " all spectral windows in object." + ) # Create new window function array window_function_array = odict() @@ -1711,7 +2179,8 @@ def get_exact_window_functions(self, ftbeam_file=None, spw_array=None, # Iterate over spectral windows for spw in spw_array: - if verbose: print('spw = {}'.format(spw)) + if verbose: + print("spw = {}".format(spw)) spw_window_function = [] spw_wf_kperp_bins, spw_wf_kpara_bins = [], [] @@ -1719,36 +2188,63 @@ def get_exact_window_functions(self, ftbeam_file=None, spw_array=None, for i, p in enumerate(self.polpair_array): # initialise UVWindow object - uvw = UVWindow.from_uvpspec(self, ipol=i, spw=spw, ftfile=ftbeam_file, - x_orientation=x_orientation, verbose=verbose) - - # extract kperp bins the window functions corresponding to the baseline + uvw = UVWindow.from_uvpspec( + self, + ipol=i, + spw=spw, + ftfile=ftbeam_file, + x_orientation=x_orientation, + verbose=verbose, + ) + + # extract kperp bins the window functions corresponding to the baseline # lengths given as input kperp_bins = uvw.get_kperp_bins(blpair_lens) kpara_bins = uvw.get_kpara_bins(uvw.freq_array) - pol_window_function = np.zeros((self.Nblpairts, self.get_dlys(spw).size, kperp_bins.size, kpara_bins.size)) + pol_window_function = np.zeros( + ( + self.Nblpairts, + self.get_dlys(spw).size, + kperp_bins.size, + kpara_bins.size, + ) + ) # Iterate over baseline-pair groups for j, blpg in enumerate(blpair_groups): - if verbose: - sys.stdout.write('\rComputing for bl group {} of {}...'.format(j+1,len(blpair_groups))) + if verbose: + sys.stdout.write( + "\rComputing for bl group {} of {}...".format( + j + 1, len(blpair_groups) + ) + ) # window functions identical for all times - window_function_blg = uvw.get_cylindrical_wf(blpair_lens[j], - kperp_bins = kperp_bins, - kpara_bins = kpara_bins) + window_function_blg = uvw.get_cylindrical_wf( + blpair_lens[j], kperp_bins=kperp_bins, kpara_bins=kpara_bins + ) if np.isnan(window_function_blg).any(): - warnings.warn('nan values in the window functions at spw={}, blpair_lens={:.2f}'.format(spw, blpair_lens[j])) + warnings.warn( + "nan values in the window functions at spw={}, blpair_lens={:.2f}".format( + spw, blpair_lens[j] + ) + ) # shape of window_function: (Ndlys, Nkperp, Nkpara) - # Iterate within a baseline-pair group + # Iterate within a baseline-pair group for k, blp in enumerate(blpg): # iterate over baselines within group, which all have the same window function for iblts in self.blpair_to_indices(blp): - pol_window_function[iblts, :, :, :] = np.copy(window_function_blg) + pol_window_function[iblts, :, :, :] = np.copy( + window_function_blg + ) - if verbose: - sys.stdout.write('\rComputed wf for baseline-pair groups {} of {}.\n'.format(len(blpair_groups),len(blpair_groups))) + if verbose: + sys.stdout.write( + "\rComputed wf for baseline-pair groups {} of {}.\n".format( + len(blpair_groups), len(blpair_groups) + ) + ) # Append to lists (spectral window) spw_window_function.append(pol_window_function) @@ -1764,16 +2260,17 @@ def get_exact_window_functions(self, ftbeam_file=None, spw_array=None, self.window_function_array = window_function_array self.window_function_kperp = window_function_kperp self.window_function_kpara = window_function_kpara - if np.all(spw_array==self.spw_array): + if np.all(spw_array == self.spw_array): self.exact_windows = True # Add to history - self.history = "Computed exact window functions [{}]\n{}\n{}\n{}".format(version.git_hash[:15], add_to_history, '-'*40, self.history) + self.history = "Computed exact window functions [{}]\n{}\n{}\n{}".format( + version.git_hash[:15], add_to_history, "-" * 40, self.history + ) # Validity check self.check() else: return window_function_kperp, window_function_kpara, window_function_array - def check(self, just_meta=False): """ Run checks to make sure metadata and data arrays are properly defined @@ -1790,38 +2287,48 @@ def check(self, just_meta=False): # only enforce existance if not just_meta if not just_meta: if p in self._req_params: - assert hasattr(self, p), \ - "required parameter {} doesn't exist".format(p) + assert hasattr( + self, p + ), "required parameter {} doesn't exist".format(p) # if attribute exists, check its type if hasattr(self, p): - a = getattr(self, '_' + p) - if hasattr(a, 'expected_type'): - err_msg = "attribute {} does not have expected data type {}".format(p, a.expected_type) + a = getattr(self, "_" + p) + if hasattr(a, "expected_type"): + err_msg = "attribute {} does not have expected data type {}".format( + p, a.expected_type + ) # ndarrays if p in self._ndarrays: - assert isinstance(getattr(self, p), np.ndarray), \ - "attribute {} needs to be an ndarray, {}".format(p, getattr(self, p).tostring()) - if issubclass(getattr(self, p).dtype.type, - a.expected_type): + assert isinstance( + getattr(self, p), np.ndarray + ), "attribute {} needs to be an ndarray, {}".format( + p, getattr(self, p).tostring() + ) + if issubclass(getattr(self, p).dtype.type, a.expected_type): pass else: # try to cast into its dtype try: - setattr(self, p, - getattr(self, p).astype(a.expected_type)) + setattr( + self, p, getattr(self, p).astype(a.expected_type) + ) except: raise ValueError(err_msg) # dicts elif p in self._dicts: - assert isinstance(getattr(self, p), (dict, odict)), \ - "attribute {} needs to be a dictionary".format(p) + assert isinstance( + getattr(self, p), (dict, odict) + ), "attribute {} needs to be a dictionary".format(p) # iterate over keys for k in getattr(self, p).keys(): - assert isinstance(getattr(self, p)[k], np.ndarray), \ - "values of attribute {} need to be ndarrays".format(p) - assert issubclass(getattr(self, p)[k].dtype.type, a.expected_type), err_msg + assert isinstance( + getattr(self, p)[k], np.ndarray + ), "values of attribute {} need to be ndarrays".format(p) + assert issubclass( + getattr(self, p)[k].dtype.type, a.expected_type + ), err_msg # immutables elif p in self._immutables: if not isinstance(getattr(self, p), a.expected_type): @@ -1838,17 +2345,21 @@ def check(self, just_meta=False): assert isinstance(getattr(self, p)[k][j], np.ndarray) try: - getattr(self, p)[k][j] = a.expected_type(getattr(self, p)[k][j]) + getattr(self, p)[k][j] = a.expected_type( + getattr(self, p)[k][j] + ) except: raise AssertionError(err_msg) # check spw convention - assert set(self.spw_array) == set(np.arange(self.Nspws)), "spw_array must be np.arange(Nspws)" + assert set(self.spw_array) == set( + np.arange(self.Nspws) + ), "spw_array must be np.arange(Nspws)" # check choice of window functions if self.exact_windows: - assert hasattr(self, 'window_function_array') and hasattr(self, 'window_function_kperp'), \ - "Error with window functions: object has exact_windows=True but no related arrays stored." - + assert hasattr(self, "window_function_array") and hasattr( + self, "window_function_kperp" + ), "Error with window functions: object has exact_windows=True but no related arrays stored." def _clear(self): """ @@ -1864,9 +2375,9 @@ def __str__(self): """ s = "" s += " ATTRIBUTES\n" - s += "-"*12 + "\n" + s += "-" * 12 + "\n" for k in self._meta_attrs: - if hasattr(self, k) and 'history' not in k: + if hasattr(self, k) and "history" not in k: y = getattr(self, k) if isinstance(y, np.ndarray): s += "%18s: shape %s\n" % (k, y.shape) @@ -1874,20 +2385,21 @@ def __str__(self): s += "%18s: %s\n" % (k, y) s += "\n DATASETS\n" - s += "-"*12 + "\n" + s += "-" * 12 + "\n" for k in self._meta_dsets: if hasattr(self, k): s += "%18s: shape %s\n" % (k, getattr(self, k).shape) return s def __eq__(self, other, params=None, verbose=False): - """ Check equivalence between attributes of two UVPSpec objects """ + """Check equivalence between attributes of two UVPSpec objects""" if params is None: params = self._all_params try: for p in params: - if p not in self._req_params \ - and (not hasattr(self, p) and not hasattr(other, p)): + if p not in self._req_params and ( + not hasattr(self, p) and not hasattr(other, p) + ): continue if p in self._immutables: assert getattr(self, p) == getattr(other, p) @@ -1898,13 +2410,16 @@ def __eq__(self, other, params=None, verbose=False): assert np.isclose(getattr(self, p), getattr(other, p)).all() elif p in self._dicts: for i in getattr(self, p): - assert np.isclose(getattr(self, p)[i],\ - getattr(other, p)[i]).all() + assert np.isclose( + getattr(self, p)[i], getattr(other, p)[i] + ).all() except AssertionError: if verbose: - print("UVPSpec parameter '{}' not equivalent between {} and {}" \ - "".format(p, self.__repr__(), other.__repr__())) + print( + "UVPSpec parameter '{}' not equivalent between {} and {}" + "".format(p, self.__repr__(), other.__repr__()) + ) return False return True @@ -1922,9 +2437,18 @@ def units(self): """ return "({})^2 {}".format(self.vis_units, self.norm_units) - def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None, - little_h=True, form='Pk', num_steps=2000, - component='real', scalar=None): + def generate_noise_spectra( + self, + spw, + polpair, + Tsys, + blpairs=None, + little_h=True, + form="Pk", + num_steps=2000, + component="real", + scalar=None, + ): """ Generate the expected RMS noise power spectrum given a selection of spectral window, system temp. [K], and polarization. This estimate is @@ -2013,14 +2537,15 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None, # Calculate scalar if scalar is None: - scalar = self.compute_scalar(spw, polpair, num_steps=num_steps, - little_h=little_h, noise_scalar=True) + scalar = self.compute_scalar( + spw, polpair, num_steps=num_steps, little_h=little_h, noise_scalar=True + ) # Get k vectors - if form == 'DelSq': + if form == "DelSq": k_perp = self.get_kperps(spw, little_h=little_h) k_para = self.get_kparas(spw, little_h=little_h) - k_mag = np.sqrt(k_perp[:, None]**2 + k_para[None, :]**2) + k_mag = np.sqrt(k_perp[:, None] ** 2 + k_para[None, :] ** 2) # get blpairs if blpairs is None: @@ -2042,9 +2567,10 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None, for i, blp in enumerate(blpairs): # get indices inds = self.blpair_to_indices(blp) - assert isinstance(Tsys[blp], (float, np.float, int, np.int)) \ - or Tsys[blp].shape[0] == self.Ntimes, \ - "Tsys must be a float or an ndarray with shape[0] == Ntimes" + assert ( + isinstance(Tsys[blp], (float, np.float, int, np.int)) + or Tsys[blp].shape[0] == self.Ntimes + ), "Tsys must be a float or an ndarray with shape[0] == Ntimes" P_blp = [] # iterate over time axis @@ -2054,18 +2580,24 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None, n_samp = self.nsample_array[spw][ind, polpair_ind] # get kvecs - if form == 'DelSq': + if form == "DelSq": k = k_mag[ind] else: k = None # Get noise power spectrum - pn = noise.calc_P_N(scalar, Tsys[blp][j], t_int, k=k, - Nincoherent=n_samp, form=form, - component=component) + pn = noise.calc_P_N( + scalar, + Tsys[blp][j], + t_int, + k=k, + Nincoherent=n_samp, + form=form, + component=component, + ) # Put into appropriate form - if form == 'Pk': + if form == "Pk": pn = np.ones(len(dlys), np.float) * pn # append to P_blp @@ -2075,10 +2607,16 @@ def generate_noise_spectra(self, spw, polpair, Tsys, blpairs=None, return P_N - - def average_spectra(self, blpair_groups=None, time_avg=False, - blpair_weights=None, error_field=None, error_weights=None, - inplace=True, add_to_history=''): + def average_spectra( + self, + blpair_groups=None, + time_avg=False, + blpair_weights=None, + error_field=None, + error_weights=None, + inplace=True, + add_to_history="", + ): """ Average power spectra across the baseline-pair-time axis, weighted by each spectrum's integration time. @@ -2161,23 +2699,27 @@ def average_spectra(self, blpair_groups=None, time_avg=False, return multiple copies of their time_array. """ if inplace: - grouping.average_spectra(self, - blpair_groups=blpair_groups, - time_avg=time_avg, - error_field=error_field, - error_weights=error_weights, - blpair_weights=blpair_weights, - inplace=True, - add_to_history=add_to_history) + grouping.average_spectra( + self, + blpair_groups=blpair_groups, + time_avg=time_avg, + error_field=error_field, + error_weights=error_weights, + blpair_weights=blpair_weights, + inplace=True, + add_to_history=add_to_history, + ) else: - return grouping.average_spectra(self, - blpair_groups=blpair_groups, - time_avg=time_avg, - error_field=error_field, - error_weights=error_weights, - blpair_weights=blpair_weights, - inplace=False, - add_to_history=add_to_history) + return grouping.average_spectra( + self, + blpair_groups=blpair_groups, + time_avg=time_avg, + error_field=error_field, + error_weights=error_weights, + blpair_weights=blpair_weights, + inplace=False, + add_to_history=add_to_history, + ) def fold_spectra(self): """ @@ -2193,7 +2735,6 @@ def fold_spectra(self): """ grouping.fold_spectra(self) - def get_blpair_groups_from_bl_groups(self, blgroups, only_pairs_in_bls=False): """ Get baseline pair matches from a list of baseline groups. @@ -2217,24 +2758,26 @@ def get_blpair_groups_from_bl_groups(self, blgroups, only_pairs_in_bls=False): """ blpair_groups = [] for blg in blgroups: - blp_select = uvputils._get_blpairs_from_bls(self, blg, - only_pairs_in_bls=only_pairs_in_bls) + blp_select = uvputils._get_blpairs_from_bls( + self, blg, only_pairs_in_bls=only_pairs_in_bls + ) blp = sorted(set(self.blpair_array[blp_select])) if len(blp) > 0: blpair_groups.append(blp) return blpair_groups - def get_red_bls(self, bl_len_tol=1., bl_ang_tol=1.): - return uvputils._get_red_bls(self, bl_len_tol=bl_len_tol, - bl_ang_tol=bl_ang_tol) + def get_red_bls(self, bl_len_tol=1.0, bl_ang_tol=1.0): + return uvputils._get_red_bls(self, bl_len_tol=bl_len_tol, bl_ang_tol=bl_ang_tol) - def get_red_blpairs(self, bl_len_tol=1., bl_ang_tol=1.): - return uvputils._get_red_blpairs(self, bl_len_tol=bl_len_tol, - bl_ang_tol=bl_ang_tol) + def get_red_blpairs(self, bl_len_tol=1.0, bl_ang_tol=1.0): + return uvputils._get_red_blpairs( + self, bl_len_tol=bl_len_tol, bl_ang_tol=bl_ang_tol + ) - def compute_scalar(self, spw, polpair, num_steps=1000, little_h=True, - noise_scalar=False): + def compute_scalar( + self, spw, polpair, num_steps=1000, little_h=True, noise_scalar=False + ): """ Compute power spectrum normalization scalar given an adopted cosmology and a beam model. See pspecbeam.PSpecBeamBase.compute_pspec_scalar for @@ -2266,10 +2809,17 @@ def compute_scalar(self, spw, polpair, num_steps=1000, little_h=True, Power spectrum normalization scalar. """ # make assertions - assert hasattr(self, 'cosmo'), \ - "self.cosmo object must exist to compute scalar. See self.set_cosmology()" - assert hasattr(self, 'OmegaP') and hasattr(self, "OmegaPP") and hasattr(self, "beam_freqs"), "self.OmegaP, "\ + assert hasattr( + self, "cosmo" + ), "self.cosmo object must exist to compute scalar. See self.set_cosmology()" + assert ( + hasattr(self, "OmegaP") + and hasattr(self, "OmegaPP") + and hasattr(self, "beam_freqs") + ), ( + "self.OmegaP, " "self.OmegaPP and self.beam_freqs must exist to compute scalar." + ) # get freq array of selected spw spw_freqs = self.freq_array[self.spw_to_freq_indices(spw)] @@ -2277,12 +2827,16 @@ def compute_scalar(self, spw, polpair, num_steps=1000, little_h=True, # compute scalar OP = self.OmegaP[:, self.polpair_to_indices(polpair)].squeeze() OPP = self.OmegaPP[:, self.polpair_to_indices(polpair)].squeeze() - scalar = pspecbeam._compute_pspec_scalar(self.cosmo, self.beam_freqs, - OPP / OP**2, spw_freqs, - num_steps=num_steps, - taper=self.taper, - little_h=little_h, - noise_scalar=noise_scalar) + scalar = pspecbeam._compute_pspec_scalar( + self.cosmo, + self.beam_freqs, + OPP / OP**2, + spw_freqs, + num_steps=num_steps, + taper=self.taper, + little_h=little_h, + noise_scalar=noise_scalar, + ) return scalar @@ -2309,8 +2863,9 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): A UVPSpec object with the data of all the inputs combined. """ # Perform type checks and get concatenation axis - (uvps, concat_ax, new_spws, new_blpts, new_polpairs, - static_meta) = get_uvp_overlap(uvps, just_meta=False, verbose=verbose) + (uvps, concat_ax, new_spws, new_blpts, new_polpairs, static_meta) = get_uvp_overlap( + uvps, just_meta=False, verbose=verbose + ) Nuvps = len(uvps) # Create a new uvp @@ -2325,10 +2880,10 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): Npols = len(new_polpairs) # Store optional attrs only if all uvps have them - store_cov = np.all([hasattr(uvp, 'cov_array_real') for uvp in uvps]) - store_window = np.all([hasattr(uvp, 'window_function_array') for uvp in uvps]) + store_cov = np.all([hasattr(uvp, "cov_array_real") for uvp in uvps]) + store_window = np.all([hasattr(uvp, "window_function_array") for uvp in uvps]) exact_windows = np.all([uvp.exact_windows for uvp in uvps]) - store_stats = np.all([hasattr(uvp, 'stats_array') for uvp in uvps]) + store_stats = np.all([hasattr(uvp, "stats_array") for uvp in uvps]) # Create new empty data arrays and fill spw arrays u.data_array = odict() u.integration_array = odict() @@ -2376,15 +2931,25 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): if exact_windows: Nkperp = uvps[0].window_function_kperp[i][:, 0].size Nkpara = uvps[0].window_function_kpara[i][:, 0].size - u.window_function_array[i] = np.empty((Nblpairts, spw[3], Nkperp, Nkpara, Npols), np.float64) + u.window_function_array[i] = np.empty( + (Nblpairts, spw[3], Nkperp, Nkpara, Npols), np.float64 + ) else: - u.window_function_array[i] = np.empty((Nblpairts, spw[3], spw[3], Npols), np.float64) + u.window_function_array[i] = np.empty( + (Nblpairts, spw[3], spw[3], Npols), np.float64 + ) if store_cov: - u.cov_array_real[i] = np.empty((Nblpairts, spw[3], spw[3], Npols), np.float64) - u.cov_array_imag[i] = np.empty((Nblpairts, spw[3], spw[3], Npols), np.float64) + u.cov_array_real[i] = np.empty( + (Nblpairts, spw[3], spw[3], Npols), np.float64 + ) + u.cov_array_imag[i] = np.empty( + (Nblpairts, spw[3], spw[3], Npols), np.float64 + ) if store_stats: for stat in stored_stats: - u.stats_array[stat][i] = np.empty((Nblpairts, spw[3], Npols), np.complex128) + u.stats_array[stat][i] = np.empty( + (Nblpairts, spw[3], Npols), np.complex128 + ) # Set frequencies and delays: if concat_ax == 'spw' this is changed below # assumes spw metadata are the same for all uvps @@ -2407,12 +2972,15 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): u.Npols = Npols # Prepare time and label arrays - u.time_1_array, u.time_2_array = np.empty(Nblpairts, np.float64), \ - np.empty(Nblpairts, np.float64) - u.time_avg_array, u.lst_avg_array = np.empty(Nblpairts, np.float64), \ - np.empty(Nblpairts, np.float64) - u.lst_1_array, u.lst_2_array = np.empty(Nblpairts, np.float64), \ - np.empty(Nblpairts, np.float64) + u.time_1_array, u.time_2_array = np.empty(Nblpairts, np.float64), np.empty( + Nblpairts, np.float64 + ) + u.time_avg_array, u.lst_avg_array = np.empty(Nblpairts, np.float64), np.empty( + Nblpairts, np.float64 + ) + u.lst_1_array, u.lst_2_array = np.empty(Nblpairts, np.float64), np.empty( + Nblpairts, np.float64 + ) u.blpair_array = np.empty(Nblpairts, np.int64) u.labels = sorted(set(np.concatenate([uvp.labels for uvp in uvps]))) u.label_1_array = np.empty((Nspws, Nblpairts, Npols), np.int32) @@ -2420,37 +2988,39 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): # get each uvp's data axes uvp_spws = [_uvp.get_spw_ranges() for _uvp in uvps] - uvp_blpts = [list(zip(_uvp.blpair_array, _uvp.time_avg_array)) - for _uvp in uvps] + uvp_blpts = [list(zip(_uvp.blpair_array, _uvp.time_avg_array)) for _uvp in uvps] uvp_polpairs = [_uvp.polpair_array.tolist() for _uvp in uvps] # Construct dict of label indices, to be used for re-ordering later u_lbls = {lbl: ll for ll, lbl in enumerate(u.labels)} # Concatenate r_params arrays - #check that r_params are either all empty or all full. + # check that r_params are either all empty or all full. _r_param_strs = [_uvp.r_params for _uvp in uvps] - if '' in _r_param_strs: - if not np.all(np.asarray([rp == '' for rp in _r_param_strs])): - raise ValueError("All r_params must be set or empty." - "Combining empty with non-empty r_params" - "is not yet supported.") + if "" in _r_param_strs: + if not np.all(np.asarray([rp == "" for rp in _r_param_strs])): + raise ValueError( + "All r_params must be set or empty." + "Combining empty with non-empty r_params" + "is not yet supported." + ) _r_params = [_uvp.get_r_params() for _uvp in uvps] - #check for conflicts by iterating through each key in the first _uvp, store - #in list of new keys. + # check for conflicts by iterating through each key in the first _uvp, store + # in list of new keys. r_params = {} for _r_param in _r_params: for rkey in _r_param: if not rkey in r_params: r_params[rkey] = _r_param[rkey] elif r_params[rkey] != _r_param[rkey]: - #For now, we won't support inconsistent weightings. - raise ValueError("Conflict between weightings" - "Only consistent weightings are supported!") - + # For now, we won't support inconsistent weightings. + raise ValueError( + "Conflict between weightings" + "Only consistent weightings are supported!" + ) # fill in data arrays depending on concat ax - if concat_ax == 'spw': + if concat_ax == "spw": u.spw_array = np.arange(Nspws, dtype=np.int32) freq_array, dly_array = [], [] @@ -2474,7 +3044,8 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): # Lookup indices of new_blpts in the uvp_blpts[l] array blpts_idxs = uvputils._fast_lookup_blpairts(uvp_blpts[l], new_blpts) - if i == 0: blpts_idxs0 = blpts_idxs.copy() + if i == 0: + blpts_idxs0 = blpts_idxs.copy() # Loop over polarizations for k, p in enumerate(new_polpairs): @@ -2497,16 +3068,26 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): u.label_1_array[i, j, k] = u_lbls[uvps[l].labels[lbl1]] u.label_2_array[i, j, k] = u_lbls[uvps[l].labels[lbl2]] if store_cov: - u.cov_array_real[i][j, :, :, k] = uvps[l].cov_array_real[m][n, :, :, q] - u.cov_array_imag[i][j, :, :, k] = uvps[l].cov_array_imag[m][n, :, :, q] + u.cov_array_real[i][j, :, :, k] = uvps[l].cov_array_real[m][ + n, :, :, q + ] + u.cov_array_imag[i][j, :, :, k] = uvps[l].cov_array_imag[m][ + n, :, :, q + ] if store_window: if exact_windows: - u.window_function_array[i][j,:, :, :,k] = uvps[l].window_function_array[m][n, :, :, :, q] + u.window_function_array[i][j, :, :, :, k] = uvps[ + l + ].window_function_array[m][n, :, :, :, q] else: - u.window_function_array[i][j,:,:,k] = uvps[l].window_function_array[m][n, :, :, q] + u.window_function_array[i][j, :, :, k] = uvps[ + l + ].window_function_array[m][n, :, :, q] if store_stats: for stat in stored_stats: - u.stats_array[stat][i][j, :, k] = uvps[l].stats_array[stat][m][n, :, q] + u.stats_array[stat][i][j, :, k] = uvps[l].stats_array[stat][ + m + ][n, :, q] u.freq_array = np.array(freq_array) u.dly_array = np.array(dly_array) @@ -2528,10 +3109,9 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): u.lst_avg_array[j] = uvps[0].lst_avg_array[n] u.blpair_array[j] = uvps[0].blpair_array[n] - elif concat_ax == 'blpairts': + elif concat_ax == "blpairts": - is_in = [uvputils._fast_is_in(_blpts, new_blpts) - for _blpts in uvp_blpts] + is_in = [uvputils._fast_is_in(_blpts, new_blpts) for _blpts in uvp_blpts] # Concatenate blpair-times for j, blpt in enumerate(new_blpts): @@ -2552,15 +3132,25 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): u.nsample_array[i][j, k] = uvps[l].nsample_array[m][n, q] if store_window: if exact_windows: - u.window_function_array[i][j, :, :, :, k] = uvps[l].window_function_array[m][n, :, :, :, q] + u.window_function_array[i][j, :, :, :, k] = uvps[ + l + ].window_function_array[m][n, :, :, :, q] else: - u.window_function_array[i][j, :, :, k] = uvps[l].window_function_array[m][n, :, :, q] + u.window_function_array[i][j, :, :, k] = uvps[ + l + ].window_function_array[m][n, :, :, q] if store_cov: - u.cov_array_real[i][j, :, :, k] = uvps[l].cov_array_real[m][n, :, :, q] - u.cov_array_imag[i][j, :, :, k] = uvps[l].cov_array_imag[m][n, :, :, q] + u.cov_array_real[i][j, :, :, k] = uvps[l].cov_array_real[m][ + n, :, :, q + ] + u.cov_array_imag[i][j, :, :, k] = uvps[l].cov_array_imag[m][ + n, :, :, q + ] if store_stats: for stat in stored_stats: - u.stats_array[stat][i][j, :, k] = uvps[l].stats_array[stat][m][n, :, q] + u.stats_array[stat][i][j, :, k] = uvps[l].stats_array[stat][ + m + ][n, :, q] # Labels lbl1 = uvps[l].label_1_array[m, n, q] lbl2 = uvps[l].label_2_array[m, n, q] @@ -2580,7 +3170,7 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): u.lst_avg_array[j] = uvps[l].lst_avg_array[n] u.blpair_array[j] = uvps[l].blpair_array[n] - elif concat_ax == 'polpairs': + elif concat_ax == "polpairs": # Concatenate polarizations for k, p in enumerate(new_polpairs): @@ -2590,12 +3180,13 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): # Get mapping of blpair-time indices between old UVPSpec objects # and the new one blpts_idxs = uvputils._fast_lookup_blpairts(uvp_blpts[l], new_blpts) - if k == 0: blpts_idxs0 = blpts_idxs.copy() + if k == 0: + blpts_idxs0 = blpts_idxs.copy() # Loop over spectral windows for i, spw in enumerate(new_spws): m = [spw == _spw for _spw in uvp_spws[l]].index(True) - u.scalar_array[i,k] = uvps[l].scalar_array[m,q] + u.scalar_array[i, k] = uvps[l].scalar_array[m, q] # Loop over blpair-times for j, blpt in enumerate(new_blpts): @@ -2603,15 +3194,25 @@ def combine_uvpspec(uvps, merge_history=True, verbose=True): u.data_array[i][j, :, k] = uvps[l].data_array[m][n, :, q] if store_window: if exact_windows: - u.window_function_array[i][j, :, :, :, k] = uvps[l].window_function_array[m][n, :, :, :, q] + u.window_function_array[i][j, :, :, :, k] = uvps[ + l + ].window_function_array[m][n, :, :, :, q] else: - u.window_function_array[i][j, :, :, k] = uvps[l].window_function_array[m][n, :, :, q] + u.window_function_array[i][j, :, :, k] = uvps[ + l + ].window_function_array[m][n, :, :, q] if store_cov: - u.cov_array_real[i][j, :, :, k] = uvps[l].cov_array_real[m][n, :, :, q] - u.cov_array_imag[i][j, :, :, k] = uvps[l].cov_array_imag[m][n, :, :, q] + u.cov_array_real[i][j, :, :, k] = uvps[l].cov_array_real[m][ + n, :, :, q + ] + u.cov_array_imag[i][j, :, :, k] = uvps[l].cov_array_imag[m][ + n, :, :, q + ] if store_stats: for stat in stored_stats: - u.stats_array[stat][i][j, :, k] = uvps[l].stats_array[stat][m][n, :, q] + u.stats_array[stat][i][j, :, k] = uvps[l].stats_array[stat][ + m + ][n, :, q] u.wgt_array[i][j, :, :, k] = uvps[l].wgt_array[m][n, :, :, q] u.integration_array[i][j, k] = uvps[l].integration_array[m][n, q] u.nsample_array[i][j, k] = uvps[l].nsample_array[m][n, q] @@ -2715,10 +3316,10 @@ def get_uvp_overlap(uvps, just_meta=True, verbose=True): List of unique polarization-pair integers across all input uvps """ # type check - assert isinstance(uvps, (list, tuple, np.ndarray)), \ - "uvps must be fed as a list" - assert isinstance(uvps[0], (UVPSpec, str, np.str)), \ - "uvps must be fed as a list of UVPSpec objects or strings" + assert isinstance(uvps, (list, tuple, np.ndarray)), "uvps must be fed as a list" + assert isinstance( + uvps[0], (UVPSpec, str, np.str) + ), "uvps must be fed as a list of UVPSpec objects or strings" Nuvps = len(uvps) # load uvps if fed as strings @@ -2731,38 +3332,64 @@ def get_uvp_overlap(uvps, just_meta=True, verbose=True): uvps = _uvps # ensure static metadata agree between all objects - static_meta = ['channel_width', 'telescope_location', 'weighting', - 'OmegaP', 'beam_freqs', 'OmegaPP', 'beamfile', 'norm', - 'taper', 'vis_units', 'norm_units', 'folded', 'cosmo', 'exact_windows'] + static_meta = [ + "channel_width", + "telescope_location", + "weighting", + "OmegaP", + "beam_freqs", + "OmegaPP", + "beamfile", + "norm", + "taper", + "vis_units", + "norm_units", + "folded", + "cosmo", + "exact_windows", + ] for m in static_meta: for u in uvps[1:]: if hasattr(uvps[0], m) and hasattr(u, m): - assert uvps[0].__eq__(u, params=[m]), \ - "Cannot concatenate UVPSpec objs: not all agree on '{}' attribute".format(m) + assert uvps[0].__eq__( + u, params=[m] + ), "Cannot concatenate UVPSpec objs: not all agree on '{}' attribute".format( + m + ) else: - assert not hasattr(uvps[0], m) and not hasattr(u, m), \ - "Cannot concatenate UVPSpec objs: not all agree on '{}' attribute".format(m) + assert not hasattr(uvps[0], m) and not hasattr( + u, m + ), "Cannot concatenate UVPSpec objs: not all agree on '{}' attribute".format( + m + ) # get static metadata values - static_meta = odict([(k, getattr(uvps[0], k, None)) for k in static_meta - if getattr(uvps[0], k, None) is not None]) + static_meta = odict( + [ + (k, getattr(uvps[0], k, None)) + for k in static_meta + if getattr(uvps[0], k, None) is not None + ] + ) # create unique data axis lists unique_spws = [] unique_blpts = [] unique_polpairs = [] data_concat_axes = odict() - blpts_comb = [] # Combined blpair + time + blpts_comb = [] # Combined blpair + time # find unique items for uvp1 in uvps: for s in uvp1.get_spw_ranges(): - if s not in unique_spws: unique_spws.append(s) + if s not in unique_spws: + unique_spws.append(s) for p in uvp1.polpair_array: - if p not in unique_polpairs: unique_polpairs.append(p) + if p not in unique_polpairs: + unique_polpairs.append(p) uvp1_blpts = zip(uvp1.blpair_array, uvp1.time_avg_array) - uvp1_blpts_comb = [bl[0] + 1.j*bl[1] for bl in uvp1_blpts] + uvp1_blpts_comb = [bl[0] + 1.0j * bl[1] for bl in uvp1_blpts] blpts_comb.extend(uvp1_blpts_comb) unique_blpts_comb = np.unique(blpts_comb) @@ -2777,7 +3404,8 @@ def get_uvp_overlap(uvps, just_meta=True, verbose=True): # iterate over uvps for j, uvp2 in enumerate(uvps): - if j <= i: continue + if j <= i: + continue # get uvp2 sets uvp2_spws = uvp2.get_spw_ranges() uvp2_blpts = list(zip(uvp2.blpair_array, uvp2.time_avg_array)) @@ -2790,48 +3418,73 @@ def get_uvp_overlap(uvps, just_meta=True, verbose=True): # ensure no partial-overlaps if not spw_match: - assert len(set(uvp1_spws) & set(uvp2_spws)) == 0, \ - "uvp {} and {} have partial overlap across spw, cannot combine".format(i, j) + assert ( + len(set(uvp1_spws) & set(uvp2_spws)) == 0 + ), "uvp {} and {} have partial overlap across spw, cannot combine".format( + i, j + ) if not blpts_match: - assert len(set(uvp1_blpts) & set(uvp2_blpts)) == 0, \ - "uvp {} and {} have partial overlap across blpairts, cannot combine".format(i, j) + assert ( + len(set(uvp1_blpts) & set(uvp2_blpts)) == 0 + ), "uvp {} and {} have partial overlap across blpairts, cannot combine".format( + i, j + ) if not polpair_match: - assert len(set(uvp1_polpairs) & set(uvp2_polpairs)) == 0, \ - "uvp {} and {} have partial overlap across pol-pair, cannot combine".format(i, j) + assert ( + len(set(uvp1_polpairs) & set(uvp2_polpairs)) == 0 + ), "uvp {} and {} have partial overlap across pol-pair, cannot combine".format( + i, j + ) # assert all except 1 axis overlaps matches = [spw_match, blpts_match, polpair_match] - assert sum(matches) != 3, \ - "uvp {} and {} have completely overlapping data, cannot combine".format(i, j) - assert sum(matches) > 1, \ - "uvp {} and {} are non-overlapping across multiple data axes, cannot combine".format(i, j) - concat_ax = ['spw', 'blpairts', 'polpairs'][matches.index(False)] + assert ( + sum(matches) != 3 + ), "uvp {} and {} have completely overlapping data, cannot combine".format( + i, j + ) + assert ( + sum(matches) > 1 + ), "uvp {} and {} are non-overlapping across multiple data axes, cannot combine".format( + i, j + ) + concat_ax = ["spw", "blpairts", "polpairs"][matches.index(False)] data_concat_axes[(i, j)] = concat_ax if verbose: - print("uvp {} and {} are concatable across {} axis".format(i, j, concat_ax)) + print( + "uvp {} and {} are concatable across {} axis".format( + i, j, concat_ax + ) + ) # assert all uvp pairs have the same (single) non-overlap (concat) axis - err_msg = "Non-overlapping data in uvps span multiple data axes:\n{}" \ - "".format("\n".join( ["{} & {}: {}".format(i[0][0],i[0][1],i[1]) - for i in data_concat_axes.items()] )) + err_msg = "Non-overlapping data in uvps span multiple data axes:\n{}" "".format( + "\n".join( + [ + "{} & {}: {}".format(i[0][0], i[0][1], i[1]) + for i in data_concat_axes.items() + ] + ) + ) assert len(set(data_concat_axes.values())) == 1, err_msg # perform additional checks given concat ax - if concat_ax == 'blpairts': + if concat_ax == "blpairts": # check scalar_array - assert np.all([np.isclose(uvps[0].scalar_array, u.scalar_array) - for u in uvps[1:]]), \ - "scalar_array must be the same for all uvps given " \ - "concatenation along blpairts." + assert np.all( + [np.isclose(uvps[0].scalar_array, u.scalar_array) for u in uvps[1:]] + ), ( + "scalar_array must be the same for all uvps given " + "concatenation along blpairts." + ) - return uvps, concat_ax, unique_spws, unique_blpts, \ - unique_polpairs, static_meta + return uvps, concat_ax, unique_spws, unique_blpts, unique_polpairs, static_meta def _ordered_unique(arr): """ Get the unique elements of an array while preserving order. - + """ arr = np.asarray(arr) _, idx = np.unique(arr, return_index=True) diff --git a/hera_pspec/uvpspec_utils.py b/hera_pspec/uvpspec_utils.py index 3ddd15a9..a61c080e 100644 --- a/hera_pspec/uvpspec_utils.py +++ b/hera_pspec/uvpspec_utils.py @@ -1,5 +1,5 @@ import numpy as np -import copy, operator +import copy from collections import OrderedDict as odict from pyuvdata.utils import polstr2num, polnum2str import json @@ -37,15 +37,22 @@ def subtract_uvp(uvp1, uvp2, run_check=True, verbose=False): A copy of uvp1 with uvp2.data_array subtracted. """ # select out common parts - uvp1, uvp2 = select_common([uvp1, uvp2], spws=True, blpairs=True, lsts=True, - polpairs=True, times=False, inplace=False, - verbose=verbose) + uvp1, uvp2 = select_common( + [uvp1, uvp2], + spws=True, + blpairs=True, + lsts=True, + polpairs=True, + times=False, + inplace=False, + verbose=verbose, + ) # get metadata - spws1 = [spw for spw in uvp1.get_spw_ranges()] + spws1 = list(spw for spw in uvp1.get_spw_ranges()) polpairs1 = uvp1.polpair_array.tolist() blps1 = sorted(set(uvp1.blpair_array)) - spws2 = [spw for spw in uvp2.get_spw_ranges()] + spws2 = list(uvp2.get_spw_ranges()) # iterate over spws for i, spw in enumerate(spws1): @@ -56,7 +63,7 @@ def subtract_uvp(uvp1, uvp2, run_check=True, verbose=False): for j, polpair in enumerate(polpairs1): # iterate over blp - for k, blp in enumerate(blps1): + for blp in blps1: # form keys key1 = (i, blp, polpair) @@ -67,55 +74,65 @@ def subtract_uvp(uvp1, uvp2, run_check=True, verbose=False): uvp1.data_array[i][blp1_inds, :, j] -= uvp2.get_data(key2) # add nsample inversely in quadrature - uvp1.nsample_array[i][blp1_inds, j] \ - = np.sqrt( 1. / (1./uvp1.get_nsamples(key1)**2 - + 1. / uvp2.get_nsamples(key2)**2) ) + uvp1.nsample_array[i][blp1_inds, j] = np.sqrt( + 1.0 + / ( + 1.0 / uvp1.get_nsamples(key1) ** 2 + + 1.0 / uvp2.get_nsamples(key2) ** 2 + ) + ) # add integration inversely in quadrature - uvp1.integration_array[i][blp1_inds, j] \ - = np.sqrt(1. / (1./uvp1.get_integrations(key1)**2 - + 1. / uvp2.get_integrations(key2)**2)) + uvp1.integration_array[i][blp1_inds, j] = np.sqrt( + 1.0 + / ( + 1.0 / uvp1.get_integrations(key1) ** 2 + + 1.0 / uvp2.get_integrations(key2) ** 2 + ) + ) # add wgts inversely in quadrature - uvp1.wgt_array[i][blp1_inds, :, :, j] \ - = np.sqrt(1. / (1./uvp1.get_wgts(key1)**2 - + 1. / uvp2.get_wgts(key2)**2)) - uvp1.wgt_array[i][blp1_inds, :, :, j] /= \ - uvp1.wgt_array[i][blp1_inds, :, :, j].max() + uvp1.wgt_array[i][blp1_inds, :, :, j] = np.sqrt( + 1.0 + / (1.0 / uvp1.get_wgts(key1) ** 2 + 1.0 / uvp2.get_wgts(key2) ** 2) + ) + uvp1.wgt_array[i][blp1_inds, :, :, j] /= uvp1.wgt_array[i][ + blp1_inds, :, :, j + ].max() # add stats in quadrature: real imag separately if hasattr(uvp1, "stats_array") and hasattr(uvp2, "stats_array"): for s in uvp1.stats_array.keys(): stat1 = uvp1.get_stats(s, key1) stat2 = uvp2.get_stats(s, key2) - uvp1.stats_array[s][i][blp1_inds, :, j] \ - = np.sqrt(stat1.real**2 + stat2.real**2) \ - + 1j*np.sqrt(stat1.imag**2 + stat2.imag**2) + uvp1.stats_array[s][i][blp1_inds, :, j] = np.sqrt( + stat1.real**2 + stat2.real**2 + ) + 1j * np.sqrt(stat1.imag**2 + stat2.imag**2) - # add cov in quadrature: real and imag separately - if hasattr(uvp1, "cov_array_real") \ - and hasattr(uvp2, "cov_array_real"): + # add cov in quadrature: real and imag separately + if hasattr(uvp1, "cov_array_real") and hasattr(uvp2, "cov_array_real"): if uvp1.cov_model == uvp2.cov_model: - cov1r = uvp1.get_cov(key1, component='real') - cov2r = uvp2.get_cov(key2, component='real') - uvp1.cov_array_real[i][blp1_inds, :, :, j] \ - = np.sqrt(cov1r.real**2 + cov2r.real**2) \ - + 1j*np.sqrt(cov1r.imag**2 + cov2r.imag**2) - - cov1i = uvp1.get_cov(key1, component='imag') - cov2i = uvp2.get_cov(key2, component='imag') - uvp1.cov_array_imag[i][blp1_inds, :, :, j] \ - = np.sqrt(cov1i.real**2 + cov2i.real**2) \ - + 1j*np.sqrt(cov1i.imag**2 + cov2i.imag**2) + cov1r = uvp1.get_cov(key1, component="real") + cov2r = uvp2.get_cov(key2, component="real") + uvp1.cov_array_real[i][blp1_inds, :, :, j] = np.sqrt( + cov1r.real**2 + cov2r.real**2 + ) + 1j * np.sqrt(cov1r.imag**2 + cov2r.imag**2) + + cov1i = uvp1.get_cov(key1, component="imag") + cov2i = uvp2.get_cov(key2, component="imag") + uvp1.cov_array_imag[i][blp1_inds, :, :, j] = np.sqrt( + cov1i.real**2 + cov2i.real**2 + ) + 1j * np.sqrt(cov1i.imag**2 + cov2i.imag**2) # same for window function - if (hasattr(uvp1, 'window_function_array') - and hasattr(uvp2, 'window_function_array')): + if hasattr(uvp1, "window_function_array") and hasattr( + uvp2, "window_function_array" + ): window1 = uvp1.get_window_function(key1) window2 = uvp2.get_window_function(key2) - uvp1.window_function_array[i][blp1_inds, :, :, j] \ - = np.sqrt(window1.real**2 + window2.real**2) \ - + 1j*np.sqrt(window1.imag**2 + window2.imag**2) + uvp1.window_function_array[i][blp1_inds, :, :, j] = np.sqrt( + window1.real**2 + window2.real**2 + ) + 1j * np.sqrt(window1.imag**2 + window2.imag**2) # run check if run_check: @@ -130,43 +147,46 @@ def compress_r_params(r_params_dict): Parameters ---------- - r_params_dict: Dictionary - dictionary with parameters for weighting matrix. Proper fields - and formats depend on the mode of data_weighting. - data_weighting == 'dayenu': - dictionary with fields - 'filter_centers', list of floats (or float) specifying the (delay) channel numbers - at which to center filtering windows. Can specify fractional channel number. - 'filter_half_widths', list of floats (or float) specifying the width of each - filter window in (delay) channel numbers. Can specify fractional channel number. - 'filter_factors', list of floats (or float) specifying how much power within each filter window - is to be suppressed. + r_params_dict : dict + dictionary with parameters for weighting matrix. Proper fields + and formats depend on the mode of data_weighting. + data_weighting == 'dayenu': + * 'filter_centers', list of floats (or float) specifying the (delay) channel numbers + at which to center filtering windows. Can specify fractional channel number. + * 'filter_half_widths', list of floats (or float) specifying the width of each + filter window in (delay) channel numbers. Can specify fractional channel number. + * 'filter_factors', list of floats (or float) specifying how much power within each filter window + is to be suppressed. + Returns ------- string containing r_params dictionary in json format and only containing one copy of each unique dictionary with a list of associated baselines. """ if r_params_dict == {} or r_params_dict is None: - return '' + return "" else: r_params_unique = {} r_params_unique_bls = {} r_params_index = -1 for rp in r_params_dict: - #do not include data set in tuple key + # do not include data set in tuple key already_in = False for rpu in r_params_unique: if r_params_unique[rpu] == r_params_dict[rp]: - r_params_unique_bls[rpu] += [rp,] + r_params_unique_bls[rpu] += [ + rp, + ] already_in = True if not already_in: r_params_index += 1 r_params_unique[r_params_index] = copy.copy(r_params_dict[rp]) - r_params_unique_bls[r_params_index] = [rp,] - + r_params_unique_bls[r_params_index] = [ + rp, + ] for rpi in r_params_unique: - r_params_unique[rpi]['baselines'] = r_params_unique_bls[rpi] + r_params_unique[rpi]["baselines"] = r_params_unique_bls[rpi] r_params_str = json.dumps(r_params_unique) return r_params_str @@ -186,31 +206,38 @@ def decompress_r_params(r_params_str): Dictionary with parameters for weighting matrix. Proper fields and formats depend on the mode of data_weighting. data_weighting == 'dayenu': - dictionary with fields - 'filter_centers', list of floats (or float) specifying the (delay) channel numbers - at which to center filtering windows. Can specify fractional channel number. - 'filter_half_widths', list of floats (or float) specifying the width of each - filter window in (delay) channel numbers. Can specify fractional channel number. - 'filter_factors', list of floats (or float) specifying how much power within each filter window - is to be suppressed. + * 'filter_centers', list of floats (or float) specifying the (delay) channel numbers + at which to center filtering windows. Can specify fractional channel number. + * 'filter_half_widths', list of floats (or float) specifying the width of each + filter window in (delay) channel numbers. Can specify fractional channel number. + * 'filter_factors', list of floats (or float) specifying how much power within each filter window + is to be suppressed. """ decompressed_r_params = {} - if r_params_str != '' and not r_params_str is None: + if r_params_str != "" and r_params_str is not None: r_params = json.loads(r_params_str) for rpi in r_params: rp_dict = {} for r_field in r_params[rpi]: - if not r_field == 'baselines': + if not r_field == "baselines": rp_dict[r_field] = r_params[rpi][r_field] - for blkey in r_params[rpi]['baselines']: + for blkey in r_params[rpi]["baselines"]: decompressed_r_params[tuple(blkey)] = rp_dict else: decompressed_r_params = {} return decompressed_r_params -def select_common(uvp_list, spws=True, blpairs=True, times=True, polpairs=True, - lsts=False, inplace=False, verbose=False): +def select_common( + uvp_list, + spws=True, + blpairs=True, + times=True, + polpairs=True, + lsts=False, + inplace=False, + verbose=False, +): """ Find spectral windows, baseline-pairs, times, and/or polarization-pairs that a set of UVPSpec objects have in common and return new UVPSpec objects @@ -267,48 +294,54 @@ def select_common(uvp_list, spws=True, blpairs=True, times=True, polpairs=True, # Get times that are common to all UVPSpec objects in the list if times: common_times = np.unique(uvp_list[0].time_avg_array) - has_times = [np.isin(common_times, uvp.time_avg_array) - for uvp in uvp_list] + has_times = [np.isin(common_times, uvp.time_avg_array) for uvp in uvp_list] common_times = common_times[np.all(has_times, axis=0)] - if verbose: print("common_times:", common_times) + if verbose: + print("common_times:", common_times) # Get lsts that are common to all UVPSpec objects in the list if lsts: common_lsts = np.unique(uvp_list[0].lst_avg_array) - has_lsts = [np.isin(common_lsts, uvp.lst_avg_array) - for uvp in uvp_list] + has_lsts = [np.isin(common_lsts, uvp.lst_avg_array) for uvp in uvp_list] common_lsts = common_lsts[np.all(has_lsts, axis=0)] - if verbose: print("common_lsts:", common_lsts) + if verbose: + print("common_lsts:", common_lsts) # Get baseline-pairs that are common to all if blpairs: common_blpairs = np.unique(uvp_list[0].blpair_array) - has_blpairs = [np.isin(common_blpairs, uvp.blpair_array) - for uvp in uvp_list] + has_blpairs = [np.isin(common_blpairs, uvp.blpair_array) for uvp in uvp_list] common_blpairs = common_blpairs[np.all(has_blpairs, axis=0)] - if verbose: print("common_blpairs:", common_blpairs) + if verbose: + print("common_blpairs:", common_blpairs) # Get polarization-pairs that are common to all if polpairs: common_polpairs = np.unique(uvp_list[0].polpair_array) - has_polpairs = [np.isin(common_polpairs, uvp.polpair_array) - for uvp in uvp_list] + has_polpairs = [np.isin(common_polpairs, uvp.polpair_array) for uvp in uvp_list] common_polpairs = common_polpairs[np.all(has_polpairs, axis=0)] - if verbose: print("common_polpairs:", common_polpairs) + if verbose: + print("common_polpairs:", common_polpairs) # Get common spectral windows (the entire window must match) # Each row of common_spws is a list of that spw's index in each UVPSpec if spws: common_spws = uvp_list[0].get_spw_ranges() - has_spws = [[x in uvp.get_spw_ranges() for x in common_spws] - for uvp in uvp_list] - common_spws = [common_spws[i] for i, f in enumerate(np.all(has_spws, axis=0)) if f] - if verbose: print("common_spws:", common_spws) + has_spws = [ + [x in uvp.get_spw_ranges() for x in common_spws] for uvp in uvp_list + ] + common_spws = [ + common_spws[i] for i, f in enumerate(np.all(has_spws, axis=0)) if f + ] + if verbose: + print("common_spws:", common_spws) # Check that this won't be an empty selection if spws and len(common_spws) == 0: - raise ValueError("No spectral windows were found that exist in all " - "spectra (the entire spectral window must match).") + raise ValueError( + "No spectral windows were found that exist in all " + "spectra (the entire spectral window must match)." + ) if blpairs and len(common_blpairs) == 0: raise ValueError("No baseline-pairs were found that exist in all spectra.") @@ -324,22 +357,33 @@ def select_common(uvp_list, spws=True, blpairs=True, times=True, polpairs=True, # Apply selections out_list = [] - for i, uvp in enumerate(uvp_list): + for uvp in uvp_list: _spws, _blpairs, _times, _lsts, _polpairs = None, None, None, None, None # Set indices of blpairs, times, and pols to keep - if blpairs: _blpairs = common_blpairs - if times: _times = common_times - if lsts: _lsts = common_lsts - if polpairs: _pols = common_polpairs - if spws: _spws = [uvp.get_spw_ranges().index(j) for j in common_spws] - - _uvp = uvp.select(spws=_spws, blpairs=_blpairs, times=_times, - polpairs=_polpairs, lsts=_lsts, inplace=inplace) - if not inplace: out_list.append(_uvp) + if blpairs: + _blpairs = common_blpairs + if times: + _times = common_times + if lsts: + _lsts = common_lsts + if spws: + _spws = [uvp.get_spw_ranges().index(j) for j in common_spws] + + _uvp = uvp.select( + spws=_spws, + blpairs=_blpairs, + times=_times, + polpairs=_polpairs, + lsts=_lsts, + inplace=inplace, + ) + if not inplace: + out_list.append(_uvp) # Return if not inplace - if not inplace: return out_list + if not inplace: + return out_list def polpair_int2tuple(polpair, pol_strings=False): @@ -367,8 +411,9 @@ def polpair_int2tuple(polpair, pol_strings=False): return [polpair_int2tuple(p, pol_strings=pol_strings) for p in polpair] # Check for integer type - assert isinstance(polpair, (int, np.integer)), \ - "polpair must be integer: %s" % type(polpair) + assert isinstance(polpair, (int, np.integer)), "polpair must be integer: %s" % type( + polpair + ) # Split into pol1 and pol2 integers pol1 = int(str(polpair)[:-2]) - 20 @@ -376,9 +421,10 @@ def polpair_int2tuple(polpair, pol_strings=False): # Check that pol1 and pol2 are in the allowed range (-8, 4) if (pol1 < -8 or pol1 > 4) or (pol2 < -8 or pol2 > 4): - raise ValueError("polpair integer evaluates to an invalid " - "polarization pair: (%d, %d)" - % (pol1, pol2)) + raise ValueError( + "polpair integer evaluates to an invalid " + "polarization pair: (%d, %d)" % (pol1, pol2) + ) # Convert to strings if requested if pol_strings: return (polnum2str(pol1), polnum2str(pol2)) @@ -421,11 +467,13 @@ def polpair_tuple2int(polpair, x_orientation=None): # Convert strings to ints if necessary pol1, pol2 = polpair - if type(pol1) in (str, np.str): pol1 = polstr2num(pol1, x_orientation=x_orientation) - if type(pol2) in (str, np.str): pol2 = polstr2num(pol2, x_orientation=x_orientation) + if type(pol1) in (str, np.str): + pol1 = polstr2num(pol1, x_orientation=x_orientation) + if type(pol2) in (str, np.str): + pol2 = polstr2num(pol2, x_orientation=x_orientation) # Convert to polpair integer - ppint = (20 + pol1)*100 + (20 + pol2) + ppint = (20 + pol1) * 100 + (20 + pol2) return ppint @@ -455,12 +503,13 @@ def _get_blpairs_from_bls(uvp, bls, only_pairs_in_bls=False): """ # get blpair baselines in integer form bl1 = np.floor(uvp.blpair_array / 1e6) - blpair_bls = np.vstack([bl1, uvp.blpair_array - bl1*1e6]).astype(np.int32).T + blpair_bls = np.vstack([bl1, uvp.blpair_array - bl1 * 1e6]).astype(np.int32).T # ensure bls is in integer form if isinstance(bls, tuple): - assert isinstance(bls[0], (int, np.integer)), \ - "bls must be fed as a list of baseline tuples Ex: [(1, 2), ...]" + assert isinstance( + bls[0], (int, np.integer) + ), "bls must be fed as a list of baseline tuples Ex: [(1, 2), ...]" bls = [uvp.antnums_to_bl(bls)] elif isinstance(bls, list): if isinstance(bls[0], tuple): @@ -470,17 +519,28 @@ def _get_blpairs_from_bls(uvp, bls, only_pairs_in_bls=False): # get indices if only_pairs_in_bls: - blp_select = np.array( [np.bool((blp[0] in bls) * (blp[1] in bls)) - for blp in blpair_bls] ) + blp_select = np.array( + [np.bool((blp[0] in bls) * (blp[1] in bls)) for blp in blpair_bls] + ) else: - blp_select = np.array( [np.bool((blp[0] in bls) + (blp[1] in bls)) - for blp in blpair_bls] ) + blp_select = np.array( + [np.bool((blp[0] in bls) + (blp[1] in bls)) for blp in blpair_bls] + ) return blp_select -def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, - times=None, lsts=None, polpairs=None, h5file=None): +def _select( + uvp, + spws=None, + bls=None, + only_pairs_in_bls=False, + blpairs=None, + times=None, + lsts=None, + polpairs=None, + h5file=None, +): """ Select function for selecting out certain slices of the data, as well as loading in data from HDF5 file. @@ -542,7 +602,7 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, uvp.Nspws = len(np.unique(uvp.spw_array)) uvp.Nspwdlys = len(uvp.spw_dly_array) uvp.Nspwfreqs = len(uvp.spw_freq_array) - if hasattr(uvp, 'scalar_array'): + if hasattr(uvp, "scalar_array"): uvp.scalar_array = uvp.scalar_array[spw_select, :] # Down-convert spw indices such that spw_array == np.arange(Nspws) @@ -560,30 +620,33 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, if bls is not None: # get blpair baselines in integer form bl1 = np.floor(uvp.blpair_array / 1e6) - blpair_bls = np.vstack([bl1, uvp.blpair_array - bl1*1e6]).astype(np.int32).T - blp_select = _get_blpairs_from_bls(uvp, bls, only_pairs_in_bls=only_pairs_in_bls) + blp_select = _get_blpairs_from_bls( + uvp, bls, only_pairs_in_bls=only_pairs_in_bls + ) if blpairs is not None: if bls is None: blp_select = np.zeros(uvp.Nblpairts, np.bool) # assert form - assert isinstance(blpairs[0], (tuple, int, np.integer)), \ - "blpairs must be fed as a list of baseline-pair tuples or baseline-pair integers" + assert isinstance( + blpairs[0], (tuple, int, np.integer) + ), "blpairs must be fed as a list of baseline-pair tuples or baseline-pair integers" # if fed as list of tuples, convert to integers if isinstance(blpairs[0], tuple): blpairs = [uvp.antnums_to_blpair(blp) for blp in blpairs] blpair_select = np.logical_or.reduce( - [uvp.blpair_array == blp for blp in blpairs]) + [uvp.blpair_array == blp for blp in blpairs] + ) blp_select += blpair_select if times is not None: if bls is None and blpairs is None: blp_select = np.ones(uvp.Nblpairts, np.bool) time_select = np.logical_or.reduce( - [np.isclose(uvp.time_avg_array, t, rtol=1e-16) - for t in times]) + [np.isclose(uvp.time_avg_array, t, rtol=1e-16) for t in times] + ) blp_select *= time_select if lsts is not None: @@ -591,25 +654,26 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, if bls is None and blpairs is None: blp_select = np.ones(uvp.Nblpairts, np.bool) lst_select = np.logical_or.reduce( - [ np.isclose(uvp.lst_avg_array, t, rtol=1e-16) - for t in lsts] ) + [np.isclose(uvp.lst_avg_array, t, rtol=1e-16) for t in lsts] + ) blp_select *= lst_select if bls is None and blpairs is None and times is None and lsts is None: blp_select = slice(None) else: # assert something was selected - assert blp_select.any(), \ - "no selections provided matched any of the data... " + assert blp_select.any(), "no selections provided matched any of the data... " # turn blp_select into slice if possible blp_select = np.where(blp_select)[0] if len(set(np.diff(blp_select))) == 0: # its sliceable, turn into slice object - blp_select = slice(blp_select[0], blp_select[-1]+1) + blp_select = slice(blp_select[0], blp_select[-1] + 1) elif len(set(np.diff(blp_select))) == 1: # its sliceable, turn into slice object - blp_select = slice(blp_select[0], blp_select[-1]+1, np.diff(blp_select)[0]) + blp_select = slice( + blp_select[0], blp_select[-1] + 1, np.diff(blp_select)[0] + ) # index arrays uvp.blpair_array = uvp.blpair_array[blp_select] @@ -626,7 +690,7 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, # Calculate unique baselines from new blpair_array new_blpairs = np.unique(uvp.blpair_array) bl1 = np.floor(new_blpairs / 1e6) - new_bls = np.unique([bl1, new_blpairs - bl1*1e6]).astype(np.int32) + new_bls = np.unique([bl1, new_blpairs - bl1 * 1e6]).astype(np.int32) # Set baseline attributes bl_select = [bl in new_bls for bl in uvp.bl_array] @@ -636,22 +700,27 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, if polpairs is not None: # assert form - assert isinstance(polpairs, (list, np.ndarray)), \ - "polpairs must be passed as a list or ndarray" - assert isinstance(polpairs[0], (tuple, int, np.integer, str)), \ - "polpairs must be fed as a list of tuples or pol integers/strings" + assert isinstance( + polpairs, (list, np.ndarray) + ), "polpairs must be passed as a list or ndarray" + assert isinstance( + polpairs[0], (tuple, int, np.integer, str) + ), "polpairs must be fed as a list of tuples or pol integers/strings" # convert str to polpair integers - polpairs = [polpair_tuple2int((p,p)) if isinstance(p, str) - else p for p in polpairs] + polpairs = [ + polpair_tuple2int((p, p)) if isinstance(p, str) else p for p in polpairs + ] # convert tuples to polpair integers - polpairs = [polpair_tuple2int(p) if isinstance(p, tuple) - else p for p in polpairs] + polpairs = [ + polpair_tuple2int(p) if isinstance(p, tuple) else p for p in polpairs + ] # create selection - polpair_select = np.logical_or.reduce( [uvp.polpair_array == p - for p in polpairs] ) + polpair_select = np.logical_or.reduce( + [uvp.polpair_array == p for p in polpairs] + ) # turn into slice object if possible polpair_select = np.where(polpair_select)[0] @@ -660,14 +729,14 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, polpair_select = slice(polpair_select[0], polpair_select[-1] + 1) elif len(set(np.diff(polpair_select))) == 1: # sliceable - polpair_select = slice(polpair_select[0], - polpair_select[-1] + 1, - np.diff(polpair_select)[0]) + polpair_select = slice( + polpair_select[0], polpair_select[-1] + 1, np.diff(polpair_select)[0] + ) # edit metadata uvp.polpair_array = uvp.polpair_array[polpair_select] uvp.Npols = len(uvp.polpair_array) - if hasattr(uvp, 'scalar_array'): + if hasattr(uvp, "scalar_array"): uvp.scalar_array = uvp.scalar_array[:, polpair_select] else: polpair_select = slice(None) @@ -679,7 +748,7 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, sliceable = False # only load / select heavy data if data_array exists _or_ if h5file is passed - if h5file is not None or hasattr(uvp, 'data_array'): + if h5file is not None or hasattr(uvp, "data_array"): # select data arrays data = odict() wgts = odict() @@ -694,22 +763,28 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, # determine if certain arrays are stored if h5file is not None: - store_cov = 'cov_real_spw0' in h5file - if 'cov_array_spw0' in h5file: + store_cov = "cov_real_spw0" in h5file + if "cov_array_spw0" in h5file: store_cov = False - warnings.warn("uvp.cov_array is no longer supported and will not be loaded. Please update this to be uvp.cov_array_real and uvp.cov_array_imag. See hera_pspec PR #181 for details.") - store_window = 'window_function_spw0' in h5file - exact_windows = 'window_function_kperp_spw0' in h5file + warnings.warn( + "uvp.cov_array is no longer supported and will not be loaded. Please update this to be uvp.cov_array_real and uvp.cov_array_imag. See hera_pspec PR #181 for details." + ) + store_window = "window_function_spw0" in h5file + exact_windows = "window_function_kperp_spw0" in h5file else: - store_cov = hasattr(uvp, 'cov_array_real') - store_window = hasattr(uvp, 'window_function_array') - exact_windows = hasattr(uvp, 'window_function_kperp') + store_cov = hasattr(uvp, "cov_array_real") + store_window = hasattr(uvp, "window_function_array") + exact_windows = hasattr(uvp, "window_function_kperp") # get stats_array keys if h5file if h5file is not None: - statnames = np.unique([f[f.find("_")+1: f.rfind("_")] - for f in h5file.keys() - if f.startswith("stats")]) + statnames = np.unique( + [ + f[f.find("_") + 1 : f.rfind("_")] + for f in h5file.keys() + if f.startswith("stats") + ] + ) else: if hasattr(uvp, "stats_array"): statnames = uvp.stats_array.keys() @@ -717,24 +792,29 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, statnames = [] # iterate over spws - if spw_mapping is None: spw_mapping = uvp.spw_array + if spw_mapping is None: + spw_mapping = uvp.spw_array for s, s_old in zip(uvp.spw_array, spw_mapping): # if h5file is passed, default to loading in data if h5file is not None: # assign data arrays - _data = h5file['data_spw{}'.format(s_old)] - _wgts = h5file['wgt_spw{}'.format(s_old)] - _ints = h5file['integration_spw{}'.format(s_old)] - _nsmp = h5file['nsample_spw{}'.format(s_old)] + _data = h5file["data_spw{}".format(s_old)] + _wgts = h5file["wgt_spw{}".format(s_old)] + _ints = h5file["integration_spw{}".format(s_old)] + _nsmp = h5file["nsample_spw{}".format(s_old)] # assign non-required arrays if store_window: - _window_function = h5file['window_function_spw{}'.format(s_old)] + _window_function = h5file["window_function_spw{}".format(s_old)] if exact_windows: - _window_function_kperp = h5file['window_function_kperp_spw{}'.format(s_old)] - _window_function_kpara = h5file['window_function_kpara_spw{}'.format(s_old)] + _window_function_kperp = h5file[ + "window_function_kperp_spw{}".format(s_old) + ] + _window_function_kpara = h5file[ + "window_function_kpara_spw{}".format(s_old) + ] if store_cov: - _cov_real = h5file["cov_real_spw{}".format(s_old)] - _cov_imag = h5file["cov_imag_spw{}".format(s_old)] + _cov_real = h5file["cov_real_spw{}".format(s_old)] + _cov_imag = h5file["cov_imag_spw{}".format(s_old)] _stat = odict() for statname in statnames: if statname not in stats: @@ -772,10 +852,16 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, ints[s] = _ints[blp_select, polpair_select] nsmp[s] = _nsmp[blp_select, polpair_select] if store_window: - window_function[s] = _window_function[blp_select, ..., polpair_select] + window_function[s] = _window_function[ + blp_select, ..., polpair_select + ] if exact_windows: - window_function_kperp[s] = _window_function_kperp[:, polpair_select] - window_function_kpara[s] = _window_function_kpara[:, polpair_select] + window_function_kperp[s] = _window_function_kperp[ + :, polpair_select + ] + window_function_kpara[s] = _window_function_kpara[ + :, polpair_select + ] if store_cov: cov_real[s] = _cov_real[blp_select, :, :, polpair_select] cov_imag[s] = _cov_imag[blp_select, :, :, polpair_select] @@ -789,16 +875,30 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, nsmp[s] = _nsmp[blp_select, :][:, polpair_select] if store_window: if exact_windows: - window_function[s] = _window_function[blp_select, :, :, :, :][:, :, :, :, polpair_select] - window_function_kperp[s] = _window_function_kperp[:, polpair_select] - window_function_kpara[s] = _window_function_kpara[:, polpair_select] + window_function[s] = _window_function[blp_select, :, :, :, :][ + :, :, :, :, polpair_select + ] + window_function_kperp[s] = _window_function_kperp[ + :, polpair_select + ] + window_function_kpara[s] = _window_function_kpara[ + :, polpair_select + ] else: - window_function[s] = _window_function[blp_select, :, :, :][:, :, :, polpair_select] + window_function[s] = _window_function[blp_select, :, :, :][ + :, :, :, polpair_select + ] if store_cov: - cov_real[s] = _cov_real[blp_select, :, :, :][:, :, :, polpair_select] - cov_imag[s] = _cov_imag[blp_select, :, :, :][:, :, :, polpair_select] + cov_real[s] = _cov_real[blp_select, :, :, :][ + :, :, :, polpair_select + ] + cov_imag[s] = _cov_imag[blp_select, :, :, :][ + :, :, :, polpair_select + ] for statname in statnames: - stats[statname][s] = _stat[statname][blp_select, :, :][:, :, polpair_select] + stats[statname][s] = _stat[statname][blp_select, :, :][ + :, :, polpair_select + ] # assign arrays to uvp uvp.data_array = data @@ -818,7 +918,7 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, uvp.cov_array_imag = cov_imag # downselect on other attrs - if hasattr(uvp, 'OmegaP'): + if hasattr(uvp, "OmegaP"): uvp.OmegaP = uvp.OmegaP[:, polpair_select] uvp.OmegaPP = uvp.OmegaPP[:, polpair_select] @@ -828,12 +928,16 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, for blpkey in blp_keys: key1 = blpkey[1][0] + (blpkey[2][0],) key2 = blpkey[1][1] + (blpkey[2][1],) - if not key1 in blkeys: - blkeys += [key1,] - if not key2 in blkeys: - blkeys += [key2,] + if key1 not in blkeys: + blkeys += [ + key1, + ] + if key2 not in blkeys: + blkeys += [ + key2, + ] new_r_params = {} - if hasattr(uvp, 'r_params') and uvp.r_params != '': + if hasattr(uvp, "r_params") and uvp.r_params != "": r_params = uvp.get_r_params() for rpkey in r_params: if rpkey in blkeys: @@ -842,6 +946,7 @@ def _select(uvp, spws=None, bls=None, only_pairs_in_bls=False, blpairs=None, new_r_params = {} uvp.r_params = compress_r_params(new_r_params) + def _blpair_to_antnums(blpair): """ Convert baseline-pair integer to nested tuple of antenna numbers. @@ -859,9 +964,9 @@ def _blpair_to_antnums(blpair): """ # get antennas ant1 = int(np.floor(blpair / 1e9)) - ant2 = int(np.floor(blpair / 1e6 - ant1*1e3)) - ant3 = int(np.floor(blpair / 1e3 - ant1*1e6 - ant2*1e3)) - ant4 = int(np.floor(blpair - ant1*1e9 - ant2*1e6 - ant3*1e3)) + ant2 = int(np.floor(blpair / 1e6 - ant1 * 1e3)) + ant3 = int(np.floor(blpair / 1e3 - ant1 * 1e6 - ant2 * 1e3)) + ant4 = int(np.floor(blpair - ant1 * 1e9 - ant2 * 1e6 - ant3 * 1e3)) ant1 -= 100 ant2 -= 100 ant3 -= 100 @@ -872,6 +977,7 @@ def _blpair_to_antnums(blpair): return antnums + def _antnums_to_blpair(antnums): """ Convert nested tuple of antenna numbers to baseline-pair integer. @@ -897,7 +1003,7 @@ def _antnums_to_blpair(antnums): ant4 = antnums[1][1] + 100 # form blpair - blpair = int(ant1*1e9 + ant2*1e6 + ant3*1e3 + ant4) + blpair = int(ant1 * 1e9 + ant2 * 1e6 + ant3 * 1e3 + ant4) return blpair @@ -918,7 +1024,7 @@ def _bl_to_antnums(bl): """ # get antennas ant1 = int(np.floor(bl / 1e3)) - ant2 = int(np.floor(bl - ant1*1e3)) + ant2 = int(np.floor(bl - ant1 * 1e3)) ant1 -= 100 ant2 -= 100 @@ -951,7 +1057,7 @@ def _antnums_to_bl(antnums): ant2 = antnums[1] + 100 # form bl - bl = int(ant1*1e3 + ant2) + bl = int(ant1 * 1e3 + ant2) return bl @@ -1016,7 +1122,7 @@ def _conj_bl_int(bl): return conj_bl -def _conj_blpair(blpair, which='both'): +def _conj_blpair(blpair, which="both"): """ Conjugate one or both baseline(s) in a baseline-pair Ex. ((ant1, ant2), (ant3, ant4)) --> ((ant2, ant1), (ant4, ant3)) @@ -1035,11 +1141,11 @@ def _conj_blpair(blpair, which='both'): blpair with one or both baselines conjugated """ antnums = _blpair_to_antnums(blpair) - if which == 'first': + if which == "first": conj_blpair = _antnums_to_blpair((antnums[0][::-1], antnums[1])) - elif which == 'second': + elif which == "second": conj_blpair = _antnums_to_blpair((antnums[0], antnums[1][::-1])) - elif which == 'both': + elif which == "both": conj_blpair = _antnums_to_blpair((antnums[0][::-1], antnums[1][::-1])) else: raise ValueError("didn't recognize {}".format(which)) @@ -1078,8 +1184,8 @@ def _fast_is_in(src_blpts, query_blpts, time_prec=8): query_blpts = np.asarray(query_blpts) # Slice to create complex array - src_blpts = src_blpts[:,0] + 1.j*np.around(src_blpts[:,1], time_prec) - query_blpts = query_blpts[:,0] + 1.j*np.around(query_blpts[:,1], time_prec) + src_blpts = src_blpts[:, 0] + 1.0j * np.around(src_blpts[:, 1], time_prec) + query_blpts = query_blpts[:, 0] + 1.0j * np.around(query_blpts[:, 1], time_prec) # see if q complex number is in src_blpts return [q in src_blpts for q in query_blpts] @@ -1116,18 +1222,18 @@ def _fast_lookup_blpairts(src_blpts, query_blpts, time_prec=8): src_blpts = np.asarray(src_blpts) query_blpts = np.asarray(query_blpts) - src_blpts = src_blpts[:,0] + 1.j*np.around(src_blpts[:,1], time_prec) - query_blpts = query_blpts[:,0] + 1.j*np.around(query_blpts[:,1], time_prec) + src_blpts = src_blpts[:, 0] + 1.0j * np.around(src_blpts[:, 1], time_prec) + query_blpts = query_blpts[:, 0] + 1.0j * np.around(query_blpts[:, 1], time_prec) # Applies rounding to time values to ensure reliable float comparisons # Do np.where comparison for all new_blpts # (indices stored in second array returned by np.where) - blpts_idxs = np.where(src_blpts == query_blpts[:,np.newaxis])[1] + blpts_idxs = np.where(src_blpts == query_blpts[:, np.newaxis])[1] return blpts_idxs -def _get_red_bls(uvp, bl_len_tol=1., bl_ang_tol=1.): +def _get_red_bls(uvp, bl_len_tol=1.0, bl_ang_tol=1.0): """ Get redundant baseline groups that are present in a UVPSpec object. @@ -1168,7 +1274,9 @@ def _get_red_bls(uvp, bl_len_tol=1., bl_ang_tol=1.): # Baseline indices idxs = np.arange(len(lens)).astype(np.int) - grp_bls = []; grp_len = []; grp_ang = [] + grp_bls = [] + grp_len = [] + grp_ang = [] # Group baselines by length and angle max_loops = idxs.size @@ -1177,9 +1285,11 @@ def _get_red_bls(uvp, bl_len_tol=1., bl_ang_tol=1.): nloops += 1 # Match bls within some tolerance in length and angle - matches = np.where(np.logical_and( - np.abs(lens - lens[0]) < bl_len_tol, - np.abs(angs - angs[0]) < bl_ang_tol) ) + matches = np.where( + np.logical_and( + np.abs(lens - lens[0]) < bl_len_tol, np.abs(angs - angs[0]) < bl_ang_tol + ) + ) # Save info about this group grp_bls.append(uvp.bl_array[idxs[matches]]) @@ -1194,7 +1304,7 @@ def _get_red_bls(uvp, bl_len_tol=1., bl_ang_tol=1.): return grp_bls, grp_len, grp_ang -def _get_red_blpairs(uvp, bl_len_tol=1., bl_ang_tol=1.): +def _get_red_blpairs(uvp, bl_len_tol=1.0, bl_ang_tol=1.0): """ Group baseline-pairs from a UVPSpec object according to the redundant groups that their constituent baselines belong to. @@ -1233,13 +1343,15 @@ def _get_red_blpairs(uvp, bl_len_tol=1., bl_ang_tol=1.): Average angle of the baselines in each group. """ # Get redundant baseline groups - red_bls, red_lens, red_angs = _get_red_bls(uvp=uvp, - bl_len_tol=bl_len_tol, - bl_ang_tol=bl_ang_tol) + red_bls, red_lens, red_angs = _get_red_bls( + uvp=uvp, bl_len_tol=bl_len_tol, bl_ang_tol=bl_ang_tol + ) # Get all available blpairs and convert to pairs of integers - blps = [(uvp.antnums_to_bl(blp[0]), uvp.antnums_to_bl(blp[1])) - for blp in uvp.get_blpairs()] + blps = [ + (uvp.antnums_to_bl(blp[0]), uvp.antnums_to_bl(blp[1])) + for blp in uvp.get_blpairs() + ] bl1, bl2 = zip(*blps) # Build bl -> group index dict @@ -1264,8 +1376,7 @@ def _get_red_blpairs(uvp, bl_len_tol=1., bl_ang_tol=1.): matches = np.where(np.logical_and(bl1_grp == i, bl2_grp == i)) # Unpack into list of blpair integers - blpair_ints = [int("%d%d" % _blp) - for _blp in zip(bl1[matches], bl2[matches])] + blpair_ints = [int("%d%d" % _blp) for _blp in zip(bl1[matches], bl2[matches])] red_grps.append(blpair_ints) return red_grps, red_lens, red_angs diff --git a/hera_pspec/uvwindow.py b/hera_pspec/uvwindow.py index 92f42784..e68dac2f 100644 --- a/hera_pspec/uvwindow.py +++ b/hera_pspec/uvwindow.py @@ -5,43 +5,35 @@ import numpy as np import sys import os -import copy -import time from scipy.interpolate import interp2d -from pyuvdata import UVBeam, UVData from astropy import units from pathlib import Path -from . import conversions, noise, version, pspecbeam, grouping, utils +from . import conversions, utils from . import uvpspec_utils as uvputils class FTBeam: """Class for :class:`FTBeam` objects.""" - def __init__(self, data, pol, freq_array, mapsize, - verbose=False, x_orientation=None): + def __init__( + self, data, pol, freq_array, mapsize, verbose=False, x_orientation=None + ): """ Obtain Fourier transform of beam in sky plane for given frenquencies. Initialise an object containing all the information related to the Fourier transform of the beam. - - data contains the array of the Fourier transform, with - dimensions (Nfreqs, N, N). - - pol contains the polarisation of the beam used. - - freq_array gives the frequency coordinates of data. - - mapsize is the size of the flat map the beam was projected - onto (in deg). Parameters ---------- data : 3D array of floats - Array containing the real part of the Fourier transform of the - beam in the sky plane (flat-sky approximation), for each of + Array containing the real part of the Fourier transform of the + beam in the sky plane (flat-sky approximation), for each of the frequencies in freq_array. Has dimensions (Nfreqs,N,N). pol : str or int - Can be pseudo-Stokes or power: + Can be pseudo-Stokes or power: in str form: 'pI', 'pQ', 'pV', 'pU', 'xx', 'yy', 'xy', 'yx' in number form: 1, 2, 4, 3, -5, -6, -7, -8 freq_array : 1D array or list of floats @@ -63,14 +55,21 @@ def __init__(self, data, pol, freq_array, mapsize, # checks on data input data = np.array(data) assert data.ndim == 3, "Wrong dimensions for data input" - assert data.shape[1] == data.shape[2], \ - "Wrong dimensions for data input" + assert data.shape[1] == data.shape[2], "Wrong dimensions for data input" self.ft_beam = data # polarisation if isinstance(pol, str): - assert pol in ['pI', 'pQ', 'pV', 'pU', 'xx', 'yy', 'xy', 'yx'], \ - "Wrong polarisation" + assert pol in [ + "pI", + "pQ", + "pV", + "pU", + "xx", + "yy", + "xy", + "yx", + ], "Wrong polarisation" elif isinstance(pol, int): assert pol in [1, 2, 4, 3, -5, -6, -7, -8], "Wrong polarisation" # convert pol number to str according to AIPS Memo 117. @@ -79,24 +78,24 @@ def __init__(self, data, pol, freq_array, mapsize, raise TypeError("Must feed pol as str or int.") self.pol = pol - assert len(freq_array) == data.shape[0], \ - "data must have shape (len(freq_array), N, N)" + assert ( + len(freq_array) == data.shape[0] + ), "data must have shape (len(freq_array), N, N)" self.freq_array = np.array(freq_array) self.mapsize = float(mapsize) - @classmethod def from_beam(cls, beamfile, verbose=False, x_orientation=None): """ Compute Fourier transform of beam in sky plane given a file containing beam simulations. - Given the path to a beam simulation, obtain the Fourier transform - of the instrument beam in the sky plane for all the frequencies - in the spectral window. + Given the path to a beam simulation, obtain the Fourier transform + of the instrument beam in the sky plane for all the frequencies + in the spectral window. Output is an array of dimensions (kperpx, kperpy, freq). - Computations correspond to the Fourier transform performed in equation + Computations correspond to the Fourier transform performed in equation 10 of memo. Parameters @@ -110,21 +109,21 @@ def from_beam(cls, beamfile, verbose=False, x_orientation=None): Default keeps polarization in X and Y basis. Used to convert polstr to polnum and conversely. """ - raise NotImplementedError('Coming soon...') + raise NotImplementedError("Coming soon...") @classmethod def from_file(cls, ftfile, spw_range=None, verbose=False, x_orientation=None): """ Read Fourier transform of beam in sky plane from file. - Initialise FTBeam object by reading file containing the - Fourier transform of the instrument beam in the sky plane - for given frequencies. + Initialise FTBeam object by reading file containing the + Fourier transform of the instrument beam in the sky plane + for given frequencies. Parameters ---------- ftfile : str - Path to file constraining to the Fourier transform of the + Path to file constraining to the Fourier transform of the beam on the sky plane (Eq. 10 in Memo), including the polarisation Ex : path/to/file/ft_beam_HERA_dipole_pI.hdf5 File must be h5 format. @@ -141,33 +140,40 @@ def from_file(cls, ftfile, spw_range=None, verbose=False, x_orientation=None): """ # check file path input ftfile = Path(ftfile) - if not (ftfile.exists() and ftfile.is_file() and h5py.is_hdf5(ftfile)) : - raise ValueError('Wrong ftfile input. See docstring.') + if not (ftfile.exists() and ftfile.is_file() and h5py.is_hdf5(ftfile)): + raise ValueError("Wrong ftfile input. See docstring.") # extract polarisation channel from filename - pol = str(ftfile).split('_')[-1].split('.')[0] + pol = str(ftfile).split("_")[-1].split(".")[0] # obtain bandwidth in file to define spectral window with h5py.File(ftfile, "r") as f: - mapsize = f['mapsize'][0] - bandwidth = np.array(f['freq'][...]) - ft_beam = np.array(f['FT_beam']) + mapsize = f["mapsize"][0] + bandwidth = np.array(f["freq"][...]) + ft_beam = np.array(f["FT_beam"]) # spectral window if spw_range is not None: # check format - assert check_spw_range(spw_range, bandwidth), \ - "Wrong spw range format, see dosctring." - freq_array = bandwidth[spw_range[0]:spw_range[-1]] - ft_beam = ft_beam[spw_range[0]:spw_range[1], :, :] + assert check_spw_range( + spw_range, bandwidth + ), "Wrong spw range format, see dosctring." + freq_array = bandwidth[spw_range[0] : spw_range[-1]] + ft_beam = ft_beam[spw_range[0] : spw_range[1], :, :] else: # if spectral range is not specified, use whole bandwidth spw_range = (0, bandwidth.size) freq_array = bandwidth - return cls(data=ft_beam, pol=pol, freq_array=freq_array, - mapsize=mapsize, verbose=verbose, x_orientation=x_orientation) - + return cls( + data=ft_beam, + pol=pol, + freq_array=freq_array, + mapsize=mapsize, + verbose=verbose, + x_orientation=x_orientation, + ) + @classmethod def get_bandwidth(cls, ftfile): """ @@ -176,8 +182,8 @@ def get_bandwidth(cls, ftfile): Parameters ---------- ftfile : str - Path to file constraining to the Fourier transform of the - beam on the sky plane (Eq. 10 in Memo). The input is the + Path to file constraining to the Fourier transform of the + beam on the sky plane (Eq. 10 in Memo). The input is the root name of the file to use, including the polarisation Ex : path/to/file/ft_beam_HERA_dipole_pI.hdf5 File must be h5 format. @@ -188,11 +194,11 @@ def get_bandwidth(cls, ftfile): List of frequencies covered by the instrument, in Hz. """ ftfile = Path(ftfile) - if not (ftfile.exists() and ftfile.is_file() and h5py.is_hdf5(ftfile)) : - raise ValueError('Wrong ftfile input. See docstring.') + if not (ftfile.exists() and ftfile.is_file() and h5py.is_hdf5(ftfile)): + raise ValueError("Wrong ftfile input. See docstring.") with h5py.File(ftfile, "r") as f: - bandwidth = np.array(f['freq'][...]) + bandwidth = np.array(f["freq"][...]) assert bandwidth.size > 1, "Error reading file, empty bandwidth." @@ -212,21 +218,27 @@ def update_spw(self, spw_range): bandwidth). """ # checks on inputs - assert check_spw_range(spw_range, self.freq_array), \ - "Wrong spw range format, see dosctring." + assert check_spw_range( + spw_range, self.freq_array + ), "Wrong spw range format, see dosctring." # assign new attributes - self.spw_range = tuple((int(spw_range[0]), int(spw_range[1]))) - self.freq_array = self.freq_array[self.spw_range[0]:self.spw_range[-1]] - self.ft_beam = self.ft_beam[self.spw_range[0]:self.spw_range[-1], :, :] + self.spw_range = (int(spw_range[0]), int(spw_range[1])) + self.freq_array = self.freq_array[self.spw_range[0] : self.spw_range[-1]] + self.ft_beam = self.ft_beam[self.spw_range[0] : self.spw_range[-1], :, :] class UVWindow: """Class for :class:`UVWindow` objects.""" - def __init__(self, ftbeam_obj, taper=None, little_h=True, - cosmo=conversions.Cosmo_Conversions(), - verbose=False): + def __init__( + self, + ftbeam_obj, + taper=None, + little_h=True, + cosmo=conversions.Cosmo_Conversions(), + verbose=False, + ): """ Class for :class:`UVWindow` objects. @@ -237,9 +249,9 @@ def __init__(self, ftbeam_obj, taper=None, little_h=True, Parameters ---------- ftbeam_obj : (list of) FTBeam object(s) - List of FTBeam objects. + List of FTBeam objects. If a unique object is given, it is expanded in a matching pair of - FTBeam objects. + FTBeam objects. Its bandwidth and polarisation attributes will define the attributes of the UVWindow object. taper : str @@ -250,7 +262,7 @@ def __init__(self, ftbeam_obj, taper=None, little_h=True, Default: True (h^-1 Mpc). cosmo : conversions.Cosmo_Conversions object, optional Cosmology object. Uses the default cosmology object if not - specified. + specified. verbose : bool, optional If True, print progress, warnings and debugging info to stdout. @@ -258,8 +270,9 @@ def __init__(self, ftbeam_obj, taper=None, little_h=True, # Summary attributes # cosmology - assert cosmo is not None, \ - "If no preferred cosmology, do not call input parameter" + assert ( + cosmo is not None + ), "If no preferred cosmology, do not call input parameter" self.cosmo = cosmo # units @@ -267,7 +280,7 @@ def __init__(self, ftbeam_obj, taper=None, little_h=True, if self.little_h: self.kunits = units.h / units.Mpc else: - self.kunits = units.Mpc**(-1) + self.kunits = units.Mpc ** (-1) # verbose self.verbose = bool(verbose) @@ -276,33 +289,38 @@ def __init__(self, ftbeam_obj, taper=None, little_h=True, try: dspec.gen_window(taper, 1) except ValueError: - raise ValueError("Wrong taper. See uvtools.dspec.gen_window" - " for options.") + raise ValueError( + "Wrong taper. See uvtools.dspec.gen_window" " for options." + ) self.taper = taper - # create list of FTBeam objects for each polarisation channel - self.ftbeam_obj_pol = list(ftbeam_obj) if np.size(ftbeam_obj) > 1 else [ftbeam_obj, ftbeam_obj] - assert isinstance(self.ftbeam_obj_pol[0], FTBeam) and isinstance(self.ftbeam_obj_pol[1], FTBeam), \ - "Wrong input given in ftbeam_obj: must be (a list of) FTBeam object(s)" + self.ftbeam_obj_pol = ( + list(ftbeam_obj) if np.size(ftbeam_obj) > 1 else [ftbeam_obj, ftbeam_obj] + ) + assert isinstance(self.ftbeam_obj_pol[0], FTBeam) and isinstance( + self.ftbeam_obj_pol[1], FTBeam + ), "Wrong input given in ftbeam_obj: must be (a list of) FTBeam object(s)" # check if elements in list have same properties - assert np.all(self.ftbeam_obj_pol[0].freq_array == self.ftbeam_obj_pol[1].freq_array), \ - 'Spectral ranges of the two FTBeam objects do not match' - assert self.ftbeam_obj_pol[0].mapsize == self.ftbeam_obj_pol[1].mapsize, \ - 'Physical properties of the two FTBeam objects do not match' + assert np.all( + self.ftbeam_obj_pol[0].freq_array == self.ftbeam_obj_pol[1].freq_array + ), "Spectral ranges of the two FTBeam objects do not match" + assert ( + self.ftbeam_obj_pol[0].mapsize == self.ftbeam_obj_pol[1].mapsize + ), "Physical properties of the two FTBeam objects do not match" # extract attributes from FTBeam objects self.pols = (self.ftbeam_obj_pol[0].pol, self.ftbeam_obj_pol[1].pol) self.freq_array = np.copy(self.ftbeam_obj_pol[0].freq_array) self.Nfreqs = len(self.freq_array) - self.dly_array = utils.get_delays(self.freq_array, - n_dlys=len(self.freq_array)) + self.dly_array = utils.get_delays(self.freq_array, n_dlys=len(self.freq_array)) self.avg_nu = np.mean(self.freq_array) self.avg_z = self.cosmo.f2z(self.avg_nu) @classmethod - def from_uvpspec(cls, uvp, ipol, spw, ftfile=None, - x_orientation=None, verbose=False): + def from_uvpspec( + cls, uvp, ipol, spw, ftfile=None, x_orientation=None, verbose=False + ): """ Method for :class:`UVWindow` objects. @@ -319,12 +337,12 @@ def from_uvpspec(cls, uvp, ipol, spw, ftfile=None, ftfile : str Access to the Fourier transform of the beam on the sky plane (Eq. 10 in Memo) - Options are; + Options are: - Load from file. Then input is the root name of the file - to use, without the polarisation - Ex : ft_beam_HERA_dipole (+ path) + to use, without the polarisation + Ex : ft_beam_HERA_dipole (+ path) - None (default). Computation from beam simulations (slow). - Not yet implemented.. + Not yet implemented.. x_orientation: str, optional Orientation in cardinal direction east or north of X dipole. Default keeps polarization in X and Y basis. @@ -335,18 +353,17 @@ def from_uvpspec(cls, uvp, ipol, spw, ftfile=None, # Summary attributes: initialise attributes from UVPSpec object # cosmology - if not hasattr(uvp, 'cosmo'): + if not hasattr(uvp, "cosmo"): cosmo = conversions.Cosmo_Conversions() - warnings.warn('uvp has no cosmo attribute. Using fiducial cosmology.') + warnings.warn("uvp has no cosmo attribute. Using fiducial cosmology.") else: cosmo = uvp.cosmo # units - little_h = 'h^-3' in uvp.norm_units + little_h = "h^-3" in uvp.norm_units # spectral window - assert spw < uvp.Nspws,\ - "Input spw must be smaller or equal to uvp.Nspws" + assert spw < uvp.Nspws, "Input spw must be smaller or equal to uvp.Nspws" freq_array = uvp.freq_array[uvp.spw_to_freq_indices(spw)] # polarisation pair @@ -360,27 +377,39 @@ def from_uvpspec(cls, uvp, ipol, spw, ftfile=None, ftbeam_obj_pol.append(ftbeam_obj_pol[0]) else: if ftfile is None: - ftbeam_obj_pol.append(FTBeam.from_beam(beamfile='tbd', - verbose=verbose, - x_orientation=x_orientation)) + ftbeam_obj_pol.append( + FTBeam.from_beam( + beamfile="tbd", verbose=verbose, x_orientation=x_orientation + ) + ) else: - ftbeam_obj_pol.append(FTBeam.from_file('{}_{}.hdf5'.format(ftfile, pol), - spw_range=None, - verbose=verbose, - x_orientation=x_orientation)) + ftbeam_obj_pol.append( + FTBeam.from_file( + "{}_{}.hdf5".format(ftfile, pol), + spw_range=None, + verbose=verbose, + x_orientation=x_orientation, + ) + ) # limit spectral window of FTBeam object to the one of the UVPSpec object bandwidth = ftbeam_obj_pol[0].freq_array - spw_range = (np.argmin(abs(bandwidth-np.min(freq_array))), - np.argmin(abs(bandwidth-np.max(freq_array)))+1) + spw_range = ( + np.argmin(abs(bandwidth - np.min(freq_array))), + np.argmin(abs(bandwidth - np.max(freq_array))) + 1, + ) for ip in range(2): if (ip > 0) and (pol[ip] == pol[0]): continue ftbeam_obj_pol[ip].update_spw(spw_range=spw_range) - return cls(ftbeam_obj=ftbeam_obj_pol, - taper=uvp.taper, cosmo=cosmo, - little_h=little_h, verbose=verbose) + return cls( + ftbeam_obj=ftbeam_obj_pol, + taper=uvp.taper, + cosmo=cosmo, + little_h=little_h, + verbose=verbose, + ) def _get_kgrid(self, bl_len, width=0.020): """ @@ -410,16 +439,21 @@ def _get_kgrid(self, bl_len, width=0.020): """ # central kperp for given baseline (i.e. 2pi*b*nu/c/R) - kp_centre = self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h)\ - * bl_len + kp_centre = self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h) * bl_len # spacing of the numerical Fourier grid, in cosmological units - dk = 2.*np.pi/(2.*self.ftbeam_obj_pol[0].mapsize)\ - / self.cosmo.dRperp_dtheta(self.cosmo.f2z(self.freq_array.max()), - little_h=self.little_h) - assert width > dk, 'Change width to resolve full window function '\ - '(dk={:.2e}).'.format(dk) + dk = ( + 2.0 + * np.pi + / (2.0 * self.ftbeam_obj_pol[0].mapsize) + / self.cosmo.dRperp_dtheta( + self.cosmo.f2z(self.freq_array.max()), little_h=self.little_h + ) + ) + assert ( + width > dk + ), "Change width to resolve full window function " "(dk={:.2e}).".format(dk) # defines kgrid (kperp_x). - kgrid = np.arange(kp_centre-width, kp_centre+width, step=dk) + kgrid = np.arange(kp_centre - width, kp_centre + width, step=dk) # array of kperp norms. kperp_norm = np.sqrt((kgrid**2)[:, None] + kgrid**2) @@ -447,15 +481,20 @@ def _kperp4bl_freq(self, freq, bl_len, ngrid): Array of k_perp values to match to the FT of the beam. """ - assert (freq <= self.freq_array.max()) and (freq >= self.freq_array.min()),\ - "Choose frequency within spectral window." - assert (freq/1e6 >= 1.), "Frequency must be given in Hz." + assert (freq <= self.freq_array.max()) and ( + freq >= self.freq_array.min() + ), "Choose frequency within spectral window." + assert freq / 1e6 >= 1.0, "Frequency must be given in Hz." z = self.cosmo.f2z(freq) R = self.cosmo.DM(z, little_h=self.little_h) # Mpc # Fourier dual of sky angle theta - q = np.fft.fftshift(np.fft.fftfreq(ngrid))*ngrid/(2.*self.ftbeam_obj_pol[0].mapsize) - k = 2.*np.pi/R*(freq*bl_len/conversions.units.c/np.sqrt(2.)-q) + q = ( + np.fft.fftshift(np.fft.fftfreq(ngrid)) + * ngrid + / (2.0 * self.ftbeam_obj_pol[0].mapsize) + ) + k = 2.0 * np.pi / R * (freq * bl_len / conversions.units.c / np.sqrt(2.0) - q) k = np.flip(k) return k @@ -484,10 +523,10 @@ def _interpolate_ft_beam(self, bl_len, ft_beam): """ ft_beam = np.array(ft_beam) assert ft_beam.ndim == 3, "ft_beam must be dimension 3." - assert ft_beam.shape[0] == self.Nfreqs,\ - "ft_beam must have shape (Nfreqs,N,N)" - assert ft_beam.shape[2] == ft_beam.shape[1],\ - "ft_beam must be square in sky plane" + assert ft_beam.shape[0] == self.Nfreqs, "ft_beam must have shape (Nfreqs,N,N)" + assert ( + ft_beam.shape[2] == ft_beam.shape[1] + ), "ft_beam must be square in sky plane" # regular kperp_x grid the FT of the beam will be interpolated over. # kperp_norm is the corresponding total kperp: @@ -502,8 +541,9 @@ def _interpolate_ft_beam(self, bl_len, ft_beam): # kperp values over frequency array for bl_len k = self._kperp4bl_freq(self.freq_array[i], bl_len, ngrid=ngrid) # interpolate the FT beam over values onto regular kgrid - A_real = interp2d(k, k, ft_beam[i, :, :], bounds_error=False, - fill_value=0.) + A_real = interp2d( + k, k, ft_beam[i, :, :], bounds_error=False, fill_value=0.0 + ) interp_ft_beam[:, :, i] = A_real(kgrid, kgrid) return interp_ft_beam, kperp_norm @@ -529,19 +569,22 @@ def _take_freq_FT(self, interp_ft_beam, delta_nu): """ interp_ft_beam = np.array(interp_ft_beam) assert interp_ft_beam.ndim == 3, "interp_ft_beam must be dimension 3." - assert interp_ft_beam.shape[-1] == self.Nfreqs,\ - "interp_ft_beam must have shape (N,N,Nfreqs)" + assert ( + interp_ft_beam.shape[-1] == self.Nfreqs + ), "interp_ft_beam must have shape (N,N,Nfreqs)" # apply taper along frequency direction if self.taper is not None: tf = dspec.gen_window(self.taper, self.Nfreqs) - interp_ft_beam = interp_ft_beam*tf[None, None, :] + interp_ft_beam = interp_ft_beam * tf[None, None, :] # take numerical FT along frequency axis # normalise to appropriate units and recentre - fnu = np.fft.fftshift(np.fft.fft(np.fft.fftshift(interp_ft_beam, - axes=-1), axis=-1, norm='ortho')*delta_nu**0.5, - axes=-1) + fnu = np.fft.fftshift( + np.fft.fft(np.fft.fftshift(interp_ft_beam, axes=-1), axis=-1, norm="ortho") + * delta_nu**0.5, + axes=-1, + ) return fnu @@ -581,35 +624,32 @@ def _get_wf_for_tau(self, tau, wf_array1, kperp_bins, kpara_bins): kperp_bins = np.array(kperp_bins) nbins_kperp = kperp_bins.size dk_perp = np.diff(kperp_bins).mean() - kperp_bin_edges = np.arange(kperp_bins.min()-dk_perp/2, - kperp_bins.max() + dk_perp, - step=dk_perp) # read kpara bins kpara_bins = np.array(kpara_bins) nbins_kpara = kpara_bins.size dk_para = np.diff(kpara_bins).mean() - kpara_bin_edges = np.arange(kpara_bins.min()-dk_para/2, - kpara_bins.max() + dk_para, - step=dk_para) + kpara_bin_edges = np.arange( + kpara_bins.min() - dk_para / 2, kpara_bins.max() + dk_para, step=dk_para + ) # get kparallel grid # conversion factor to cosmological units - alpha = self.cosmo.dRpara_df(self.avg_z, little_h=self.little_h, - ghz=False) + alpha = self.cosmo.dRpara_df(self.avg_z, little_h=self.little_h, ghz=False) # frequency resolution - delta_nu = abs(self.freq_array[-1]-self.freq_array[0])/self.Nfreqs + delta_nu = abs(self.freq_array[-1] - self.freq_array[0]) / self.Nfreqs # Fourier dual of frequency (unit 1: FT along theta) - eta = np.fft.fftshift(np.fft.fftfreq(self.Nfreqs), axes=-1)/delta_nu + eta = np.fft.fftshift(np.fft.fftfreq(self.Nfreqs), axes=-1) / delta_nu # construct array of |kpara| values for given delay tau - kpar_norm = np.abs(2.*np.pi/alpha*(eta+tau)) + kpar_norm = np.abs(2.0 * np.pi / alpha * (eta + tau)) # perform binning along k_parallel cyl_wf = np.zeros((nbins_kperp, nbins_kpara)) kpara = np.zeros(nbins_kpara) for j in range(nbins_kperp): for m in range(nbins_kpara): - mask = (kpara_bin_edges[m] <= kpar_norm) &\ - (kpar_norm < kpara_bin_edges[m+1]) + mask = (kpara_bin_edges[m] <= kpar_norm) & ( + kpar_norm < kpara_bin_edges[m + 1] + ) if np.any(mask): # cannot compute mean if zero elements cyl_wf[j, m] = np.mean(wf_array1[j, mask]) kpara[m] = np.mean(kpar_norm[mask]) @@ -635,20 +675,23 @@ def get_kperp_bins(self, bl_lens): Array of kperp bins to use. """ bl_lens = np.array(bl_lens) - assert bl_lens.size > 0,\ - "get_kperp_bins() requires array of baseline lengths." + assert bl_lens.size > 0, "get_kperp_bins() requires array of baseline lengths." - dk_perp = np.diff(self._get_kgrid(np.min(bl_lens))[1]).mean()*5 - kperp_max = self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h)\ - * np.max(bl_lens) + 10.*dk_perp + dk_perp = np.diff(self._get_kgrid(np.min(bl_lens))[1]).mean() * 5 + kperp_max = ( + self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h) * np.max(bl_lens) + + 10.0 * dk_perp + ) kperp_bin_edges = np.arange(dk_perp, kperp_max, step=dk_perp) - kperp_bins = (kperp_bin_edges[1:]+kperp_bin_edges[:-1])/2 + kperp_bins = (kperp_bin_edges[1:] + kperp_bin_edges[:-1]) / 2 nbins_kperp = kperp_bins.size - if (nbins_kperp > 200): - warnings.warn('get_kperp_bins: Large number of kperp/kpara bins. ' - 'Risk of overresolving and slow computing.') + if nbins_kperp > 200: + warnings.warn( + "get_kperp_bins: Large number of kperp/kpara bins. " + "Risk of overresolving and slow computing." + ) - return kperp_bins*self.kunits + return kperp_bins * self.kunits def get_kpara_bins(self, freq_array): """ @@ -676,26 +719,32 @@ def get_kpara_bins(self, freq_array): assert freq_array.size > 1, "Must feed list of frequencies." dly_array = utils.get_delays(freq_array, n_dlys=len(freq_array)) - avg_z = self.cosmo.f2z(np.mean(freq_array)) + avg_z = self.cosmo.f2z(np.mean(freq_array)) # define default kperp bins, - dk_para = self.cosmo.tau_to_kpara(avg_z, little_h=self.little_h)\ - / (abs(freq_array[-1]-freq_array[0])) - kpara_max = self.cosmo.tau_to_kpara(avg_z, little_h=self.little_h)\ - * abs(dly_array).max()+10.*dk_para + dk_para = self.cosmo.tau_to_kpara(avg_z, little_h=self.little_h) / ( + abs(freq_array[-1] - freq_array[0]) + ) + kpara_max = ( + self.cosmo.tau_to_kpara(avg_z, little_h=self.little_h) + * abs(dly_array).max() + + 10.0 * dk_para + ) kpara_bin_edges = np.arange(dk_para, kpara_max, step=dk_para) - kpara_bins = (kpara_bin_edges[1:]+kpara_bin_edges[:-1])/2 + kpara_bins = (kpara_bin_edges[1:] + kpara_bin_edges[:-1]) / 2 nbins_kpara = kpara_bins.size - if (nbins_kpara > 200): - warnings.warn('get_kpara_bins: Large number of kperp/kpara bins. ' - 'Risk of overresolving and slow computing.') + if nbins_kpara > 200: + warnings.warn( + "get_kpara_bins: Large number of kperp/kpara bins. " + "Risk of overresolving and slow computing." + ) - return kpara_bins*self.kunits + return kpara_bins * self.kunits - def get_cylindrical_wf(self, bl_len, - kperp_bins=None, kpara_bins=None, - return_bins=None, verbose=None): + def get_cylindrical_wf( + self, bl_len, kperp_bins=None, kpara_bins=None, return_bins=None, verbose=None + ): """ Get the cylindrical window function for a baseline length. @@ -754,37 +803,47 @@ def get_cylindrical_wf(self, bl_len, else: self.check_kunits(kperp_bins) kperp_bins = np.array(kperp_bins.value) - if not np.isclose(np.diff(kperp_bins),np.diff(kperp_bins)[0]).all(): - warnings.warn('kperp_bins must be linearly spaced.') + if not np.isclose(np.diff(kperp_bins), np.diff(kperp_bins)[0]).all(): + warnings.warn("kperp_bins must be linearly spaced.") nbins_kperp = kperp_bins.size dk_perp = np.diff(kperp_bins).mean() - kperp_bin_edges = np.arange(kperp_bins.min()-dk_perp/2, - kperp_bins.max() + dk_perp, - step=dk_perp) - kperp_centre = self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h) * bl_len - if (kperp_bin_edges.max() < kperp_centre+3*dk_perp) or\ - (kperp_bin_edges.min() > kperp_centre-3*dk_perp): - warnings.warn('get_cylindrical_wf: The bin centre is not included ' - 'in the array of kperp bins given as input.') + kperp_bin_edges = np.arange( + kperp_bins.min() - dk_perp / 2, kperp_bins.max() + dk_perp, step=dk_perp + ) + kperp_centre = ( + self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h) * bl_len + ) + if (kperp_bin_edges.max() < kperp_centre + 3 * dk_perp) or ( + kperp_bin_edges.min() > kperp_centre - 3 * dk_perp + ): + warnings.warn( + "get_cylindrical_wf: The bin centre is not included " + "in the array of kperp bins given as input." + ) if kpara_bins is None: kpara_bins = self.get_kpara_bins(self.freq_array) else: self.check_kunits(kpara_bins) kpara_bins = np.array(kpara_bins.value) - if not np.isclose(np.diff(kpara_bins),np.diff(kpara_bins)[0]).all(): - warnings.warn('kpara_bins must be linearly spaced.') + if not np.isclose(np.diff(kpara_bins), np.diff(kpara_bins)[0]).all(): + warnings.warn("kpara_bins must be linearly spaced.") nbins_kpara = kpara_bins.size dk_para = np.diff(kpara_bins).mean() - kpara_bin_edges = np.arange(kpara_bins.min()-dk_para/2, - kpara_bins.max()+dk_para, - step=dk_para) - kpara_centre = self.cosmo.tau_to_kpara(self.avg_z, little_h=self.little_h)\ + kpara_bin_edges = np.arange( + kpara_bins.min() - dk_para / 2, kpara_bins.max() + dk_para, step=dk_para + ) + kpara_centre = ( + self.cosmo.tau_to_kpara(self.avg_z, little_h=self.little_h) * abs(self.dly_array).max() - if (kpara_bin_edges.max() < kpara_centre+3*dk_para) or \ - (kpara_bin_edges.min() > kpara_centre-3*dk_para): - warnings.warn('get_cylindrical_wf: The bin centre is not included ' - 'in the array of kpara bins given as input.') + ) + if (kpara_bin_edges.max() < kpara_centre + 3 * dk_para) or ( + kpara_bin_edges.min() > kpara_centre - 3 * dk_para + ): + warnings.warn( + "get_cylindrical_wf: The bin centre is not included " + "in the array of kpara bins given as input." + ) # COMPUTE CYLINDRICAL WINDOW FUNCTIONS @@ -795,7 +854,7 @@ def get_cylindrical_wf(self, bl_len, # interpolate FT of beam onto regular grid of (kperp_x,kperp_y) interp_ft_beam, kperp_norm = self._interpolate_ft_beam(bl_len, ft_beam) # frequency resolution - delta_nu = abs(self.freq_array[-1]-self.freq_array[0])/self.Nfreqs + delta_nu = abs(self.freq_array[-1] - self.freq_array[0]) / self.Nfreqs # obtain FT along frequency fnu.append(self._take_freq_FT(interp_ft_beam, delta_nu)) @@ -806,37 +865,43 @@ def get_cylindrical_wf(self, bl_len, kperp = np.zeros(nbins_kperp) for i in range(self.Nfreqs): for m in range(nbins_kperp): - mask = (kperp_bin_edges[m] <= kperp_norm)\ - & (kperp_norm < kperp_bin_edges[m+1]) + mask = (kperp_bin_edges[m] <= kperp_norm) & ( + kperp_norm < kperp_bin_edges[m + 1] + ) if np.any(mask): # cannot compute mean if zero elements - wf_array1[m, i] = np.mean(np.conj(fnu[1][mask, i])*fnu[0][mask, i]).real + wf_array1[m, i] = np.mean( + np.conj(fnu[1][mask, i]) * fnu[0][mask, i] + ).real kperp[m] = np.mean(kperp_norm[mask]) # in frequency direction cyl_wf = np.zeros((self.Nfreqs, nbins_kperp, nbins_kpara)) - for it, tau in enumerate(self.dly_array[:self.Nfreqs//2+1]): - kpara, cyl_wf[it, :, :] = self._get_wf_for_tau(tau, wf_array1, - kperp_bins, - kpara_bins) + for it, tau in enumerate(self.dly_array[: self.Nfreqs // 2 + 1]): + kpara, cyl_wf[it, :, :] = self._get_wf_for_tau( + tau, wf_array1, kperp_bins, kpara_bins + ) # fill by symmetry for tau = -tau - if (self.Nfreqs % 2 == 0): - cyl_wf[self.Nfreqs//2+1:, :, :] = np.flip(cyl_wf, axis=0)[self.Nfreqs//2:-1] + if self.Nfreqs % 2 == 0: + cyl_wf[self.Nfreqs // 2 + 1 :, :, :] = np.flip(cyl_wf, axis=0)[ + self.Nfreqs // 2 : -1 + ] else: - cyl_wf[self.Nfreqs//2+1:, :, :] = np.flip(cyl_wf, axis=0)[self.Nfreqs//2+1:] + cyl_wf[self.Nfreqs // 2 + 1 :, :, :] = np.flip(cyl_wf, axis=0)[ + self.Nfreqs // 2 + 1 : + ] # normalisation of window functions sum_per_bin = np.sum(cyl_wf, axis=(1, 2))[:, None, None] cyl_wf = np.divide(cyl_wf, sum_per_bin, where=sum_per_bin != 0) - if (return_bins == 'unweighted') or return_bins: + if (return_bins == "unweighted") or return_bins: return kperp_bins, kpara_bins, cyl_wf - elif (return_bins == 'weighted'): + elif return_bins == "weighted": return kperp, kpara, cyl_wf else: return cyl_wf - def cylindrical_to_spherical(self, cyl_wf, kbins, ktot, bl_lens, - bl_weights=None): + def cylindrical_to_spherical(self, cyl_wf, kbins, ktot, bl_lens, bl_weights=None): """ Take spherical average of cylindrical window functions. @@ -850,7 +915,7 @@ def cylindrical_to_spherical(self, cyl_wf, kbins, ktot, bl_lens, Axis 3 is kparallel. If only one bl_lens is given, then axis 0 can be omitted. kbins : array-like astropy.quantity with units - 1D float array of ascending |k| bin centers in [h] Mpc^-1 units. + 1D float array of ascending ``|k|`` bin centers in [h] Mpc^-1 units. Used for spherical binning. Must be linearly spaced. ktot : array_like 2-dimensional array giving the magnitude of k corresponding to @@ -859,7 +924,7 @@ def cylindrical_to_spherical(self, cyl_wf, kbins, ktot, bl_lens, List of baseline lengths used to compute cyl_wf. Must have same length as same size as cyl_wf.shape[0]. bl_weights : list of weights (float or int), optional - Relative weight of each baseline-length when performing + Relative weight of each baseline-length when performing the average. This should have the same shape as bl_lens. Default: None (all baseline pairs have unity weights). @@ -872,73 +937,93 @@ def cylindrical_to_spherical(self, cyl_wf, kbins, ktot, bl_lens, Returns weighted k-modes. """ # check bl_lens and bl_weights are consistent - bl_lens = bl_lens if isinstance(bl_lens, (list, tuple, np.ndarray)) else [bl_lens] + bl_lens = ( + bl_lens if isinstance(bl_lens, (list, tuple, np.ndarray)) else [bl_lens] + ) bl_lens = np.array(bl_lens) if bl_weights is None: # assign weight of one to each baseline length bl_weights = np.ones(bl_lens.size) else: bl_weights = np.array(bl_weights) - assert bl_weights.size == bl_lens.size, \ - "Blpair weights and lengths do not match" + assert ( + bl_weights.size == bl_lens.size + ), "Blpair weights and lengths do not match" # if cyl_wf were computed only for one baseline length if cyl_wf.ndim == 3: - assert bl_lens.size == 1, "If only one bl_lens is given,"\ + assert bl_lens.size == 1, ( + "If only one bl_lens is given," "cyl_wf must be of dimensions (ndlys,nkperp,nkpara)" + ) cyl_wf = cyl_wf[None] # check shapes are consistent assert bl_lens.size == cyl_wf.shape[0] - assert (ktot.shape == cyl_wf.shape[2:]), \ - "k magnitude grid does not match (kperp,kpara) grid in cyl_wf" + assert ( + ktot.shape == cyl_wf.shape[2:] + ), "k magnitude grid does not match (kperp,kpara) grid in cyl_wf" # k-bins for spherical binning - assert len(kbins) > 1, \ - "must feed array of k bins for spherical average" + assert len(kbins) > 1, "must feed array of k bins for spherical average" self.check_kunits(kbins) # check k units kbins = np.array(kbins.value) - if not np.isclose(np.diff(kbins),np.diff(kbins)[0]).all(): - warnings.warn('kbins must be linearly spaced.') + if not np.isclose(np.diff(kbins), np.diff(kbins)[0]).all(): + warnings.warn("kbins must be linearly spaced.") nbinsk = kbins.size dk = np.diff(kbins).mean() - kbin_edges = np.arange(kbins.min()-dk/2, kbins.max()+dk, step=dk) + kbin_edges = np.arange(kbins.min() - dk / 2, kbins.max() + dk, step=dk) # construct array giving the k probed by each baseline-tau pair - kperps = self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h) \ - * bl_lens - kparas = self.cosmo.tau_to_kpara(self.avg_z, little_h=self.little_h) \ - * self.dly_array - kmags = np.sqrt(kperps[:, None]**2+kparas**2) + kperps = self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h) * bl_lens + kparas = ( + self.cosmo.tau_to_kpara(self.avg_z, little_h=self.little_h) * self.dly_array + ) + kmags = np.sqrt(kperps[:, None] ** 2 + kparas**2) # take average wf_spherical = np.zeros((nbinsk, nbinsk)) kweights, weighted_k = np.zeros(nbinsk, dtype=int), np.zeros(nbinsk) for m1 in range(nbinsk): - mask2 = (kbin_edges[m1] <= kmags) & (kmags < kbin_edges[m1+1]) + mask2 = (kbin_edges[m1] <= kmags) & (kmags < kbin_edges[m1 + 1]) if np.any(mask2): weighted_k[m1] = np.mean(kmags[mask2]) - mask2 = mask2.astype(int)*bl_weights[:, None] #add weights for redundancy - kweights[m1] = np.sum(mask2) - wf_temp = np.sum(cyl_wf*mask2[:,:,None,None], axis=(0, 1))/np.sum(mask2) - if np.sum(wf_temp) > 0.: + mask2 = ( + mask2.astype(int) * bl_weights[:, None] + ) # add weights for redundancy + kweights[m1] = np.sum(mask2) + wf_temp = np.sum( + cyl_wf * mask2[:, :, None, None], axis=(0, 1) + ) / np.sum(mask2) + if np.sum(wf_temp) > 0.0: for m in range(nbinsk): - mask = (kbin_edges[m] <= ktot) & (ktot < kbin_edges[m+1]) - if np.any(mask): #cannot compute mean if zero elements - wf_spherical[m1,m]=np.mean(wf_temp[mask]) + mask = (kbin_edges[m] <= ktot) & (ktot < kbin_edges[m + 1]) + if np.any(mask): # cannot compute mean if zero elements + wf_spherical[m1, m] = np.mean(wf_temp[mask]) # normalisation - wf_spherical[m1,:] = np.divide(wf_spherical[m1, :], np.sum(wf_spherical[m1, :]), - where = np.sum(wf_spherical[m1,:]) != 0) - - if np.any(kweights == 0.) and self.verbose: - warnings.warn('Some spherical bins are empty. ' - 'Add baselines or expand spectral window.') + wf_spherical[m1, :] = np.divide( + wf_spherical[m1, :], + np.sum(wf_spherical[m1, :]), + where=np.sum(wf_spherical[m1, :]) != 0, + ) + + if np.any(kweights == 0.0) and self.verbose: + warnings.warn( + "Some spherical bins are empty. " + "Add baselines or expand spectral window." + ) return wf_spherical, weighted_k - def get_spherical_wf(self, kbins, bl_lens, bl_weights=None, - kperp_bins=None, kpara_bins=None, - return_weighted_k=False, - verbose=None): + def get_spherical_wf( + self, + kbins, + bl_lens, + bl_weights=None, + kperp_bins=None, + kpara_bins=None, + return_weighted_k=False, + verbose=None, + ): """ Get spherical window functions for a UVWindow object. @@ -948,7 +1033,7 @@ def get_spherical_wf(self, kbins, bl_lens, bl_weights=None, Parameters ---------- kbins : array-like astropy.quantity with units - 1D float array of ascending |k| bin centers in [h] Mpc^-1 units. + 1D float array of ascending ``|k|`` bin centers in [h] Mpc^-1 units. Using for spherical binning. Make sure the values are consistent with :attr:`little_h`. bl_lens : list. @@ -957,7 +1042,7 @@ def get_spherical_wf(self, kbins, bl_lens, bl_weights=None, Must have same length as bl_weights. bl_weights : list, optional. List baselines weights. Must have same length as bl_lens. - If None, a weight of 1 is attributed to + If None, a weight of 1 is attributed to each bl_len. kperp_bins : array-like astropy.quantity with units, optional. 1D float array of ascending k_perp bin centers in [h] Mpc^-1 units. @@ -993,8 +1078,9 @@ def get_spherical_wf(self, kbins, bl_lens, bl_weights=None, bl_lens = np.array(bl_lens) if bl_weights is not None: # check consistency of baseline-related inputs - assert len(bl_weights) == nbls, "bl_weights and bl_lens "\ - "must have same length" + assert len(bl_weights) == nbls, ( + "bl_weights and bl_lens " "must have same length" + ) bl_weights = np.array(bl_weights) else: # each baseline length has weight one @@ -1008,25 +1094,31 @@ def get_spherical_wf(self, kbins, bl_lens, bl_weights=None, else: self.check_kunits(kperp_bins) kperp_bins = np.array(kperp_bins.value) - if not np.isclose(np.diff(kperp_bins),np.diff(kperp_bins)[0]).all(): - warnings.warn('kperp_bins must be linearly spaced.') + if not np.isclose(np.diff(kperp_bins), np.diff(kperp_bins)[0]).all(): + warnings.warn("kperp_bins must be linearly spaced.") nbins_kperp = kperp_bins.size dk_perp = np.diff(kperp_bins).mean() - kperp_bin_edges = np.arange(kperp_bins.min()-dk_perp/2, - kperp_bins.max()+dk_perp, - step=dk_perp) + kperp_bin_edges = np.arange( + kperp_bins.min() - dk_perp / 2, kperp_bins.max() + dk_perp, step=dk_perp + ) # make sure proper kperp values are included in given bins # raise warning otherwise - kperp_max = self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h)\ - * np.max(bl_lens) #+ 10.*dk_perp - kperp_min = self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h)\ - * np.min(bl_lens) #+ 10.*dk_perp - if (kperp_bin_edges.max() <= kperp_max): - warnings.warn('get_spherical_wf: Max kperp bin centre not ' - 'included in binning array') - if (kperp_bin_edges.min() >= kperp_min): - warnings.warn('get_spherical_wf: Min kperp bin centre not ' - 'included in binning array') + kperp_max = self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h) * np.max( + bl_lens + ) # + 10.*dk_perp + kperp_min = self.cosmo.bl_to_kperp(self.avg_z, little_h=self.little_h) * np.min( + bl_lens + ) # + 10.*dk_perp + if kperp_bin_edges.max() <= kperp_max: + warnings.warn( + "get_spherical_wf: Max kperp bin centre not " + "included in binning array" + ) + if kperp_bin_edges.min() >= kperp_min: + warnings.warn( + "get_spherical_wf: Min kperp bin centre not " + "included in binning array" + ) if kpara_bins is None or len(kpara_bins) == 0: # define default kperp bins, making sure all values probed by freq @@ -1035,43 +1127,45 @@ def get_spherical_wf(self, kbins, bl_lens, bl_weights=None, else: self.check_kunits(kpara_bins) kpara_bins = np.array(kpara_bins.value) - if not np.isclose(np.diff(kpara_bins),np.diff(kpara_bins)[0]).all(): - warnings.warn('kpara_bins must be linearly spaced.') + if not np.isclose(np.diff(kpara_bins), np.diff(kpara_bins)[0]).all(): + warnings.warn("kpara_bins must be linearly spaced.") nbins_kpara = kpara_bins.size dk_para = np.diff(kpara_bins).mean() - kpara_bin_edges = np.arange(kpara_bins.min() - dk_para/2, - kpara_bins.max() + dk_para, - step=dk_para) - kpara_centre = self.cosmo.tau_to_kpara(self.avg_z, - little_h=self.little_h)\ + kpara_bin_edges = np.arange( + kpara_bins.min() - dk_para / 2, kpara_bins.max() + dk_para, step=dk_para + ) + kpara_centre = ( + self.cosmo.tau_to_kpara(self.avg_z, little_h=self.little_h) * abs(self.dly_array).max() + ) # make sure proper kpara values are included in given bins # raise warning otherwise - if (kpara_bin_edges.max() <= kpara_centre+3*dk_para) or\ - (kpara_bin_edges.min() >= kpara_centre-3*dk_para): - warnings.warn('get_spherical_wf: The bin centre is not included ' - 'in the array of kpara bins given as input.') + if (kpara_bin_edges.max() <= kpara_centre + 3 * dk_para) or ( + kpara_bin_edges.min() >= kpara_centre - 3 * dk_para + ): + warnings.warn( + "get_spherical_wf: The bin centre is not included " + "in the array of kpara bins given as input." + ) # array of |k|=sqrt(kperp**2+kpara**2) - ktot = np.sqrt(kperp_bins[:, None]**2+kpara_bins**2) + ktot = np.sqrt(kperp_bins[:, None] ** 2 + kpara_bins**2) # k-bins for spherical binning self.check_kunits(kbins) # check k units - assert kbins.value.size > 1, \ - "must feed array of k bins for spherical average" - nbinsk = kbins.value.size - if not np.isclose(np.diff(kbins),np.diff(kbins)[0]).all(): - warnings.warn('kbins must be linearly spaced.') + assert kbins.value.size > 1, "must feed array of k bins for spherical average" + if not np.isclose(np.diff(kbins), np.diff(kbins)[0]).all(): + warnings.warn("kbins must be linearly spaced.") dk = np.diff(kbins.value).mean() - kbin_edges = np.arange(kbins.value.min()-dk/2, - kbins.value.max()+dk, - step=dk) + kbin_edges = np.arange( + kbins.value.min() - dk / 2, kbins.value.max() + dk, step=dk + ) # make sure proper ktot values are included in given bins # raise warning otherwise - if (kbin_edges.max() <= ktot.max()): - warnings.warn('Max spherical k probed is not included in bins.') - if (kbin_edges.min() >= ktot.min()): - warnings.warn('Min spherical k probed is not included in bins.') + if kbin_edges.max() <= ktot.max(): + warnings.warn("Max spherical k probed is not included in bins.") + if kbin_edges.min() >= ktot.min(): + warnings.warn("Min spherical k probed is not included in bins.") # COMPUTE THE WINDOW FUNCTIONS # get cylindrical window functions for each baseline length considered @@ -1079,27 +1173,39 @@ def get_spherical_wf(self, kbins, bl_lens, bl_weights=None, cyl_wf = np.zeros((nbls, self.Nfreqs, nbins_kperp, nbins_kpara)) for ib in range(nbls): if verbose: - sys.stdout.write('\rComputing for blg {:d} of {:d}...'.format(ib + 1, nbls)) - cyl_wf[ib, :, :, :] = self.get_cylindrical_wf(bl_len=bl_lens[ib], - kperp_bins=kperp_bins*self.kunits, - kpara_bins=kpara_bins*self.kunits, - verbose=verbose) + sys.stdout.write( + "\rComputing for blg {:d} of {:d}...".format(ib + 1, nbls) + ) + cyl_wf[ib, :, :, :] = self.get_cylindrical_wf( + bl_len=bl_lens[ib], + kperp_bins=kperp_bins * self.kunits, + kpara_bins=kpara_bins * self.kunits, + verbose=verbose, + ) if verbose: - sys.stdout.write('\rComputing for blg {:d} of {:d}... \n'.format(nbls, nbls)) + sys.stdout.write( + "\rComputing for blg {:d} of {:d}... \n".format(nbls, nbls) + ) # perform spherical binning - wf_spherical, weighted_k = self.cylindrical_to_spherical(cyl_wf, - kbins, ktot, - bl_lens, bl_weights) + wf_spherical, weighted_k = self.cylindrical_to_spherical( + cyl_wf, kbins, ktot, bl_lens, bl_weights + ) if return_weighted_k: return wf_spherical, weighted_k else: return wf_spherical - def run_and_write(self, filepath, bl_lens, bl_weights=None, - kperp_bins=None, kpara_bins=None, - clobber=False): + def run_and_write( + self, + filepath, + bl_lens, + bl_weights=None, + kperp_bins=None, + kpara_bins=None, + clobber=False, + ): """ Run cylindrical wf and write result to HDF5 file. @@ -1113,7 +1219,7 @@ def run_and_write(self, filepath, bl_lens, bl_weights=None, Must have same length as bl_weights. bl_weights : list, optional. List baselines weights. Must have same length as bl_lens. - If None, a weight of 1 is attributed to + If None, a weight of 1 is attributed to each bl_len. kperp_bins : array-like astropy.quantity with unit, optional. 1D float array of ascending k_perp bin centers in [h] Mpc^-1 units. @@ -1138,8 +1244,9 @@ def run_and_write(self, filepath, bl_lens, bl_weights=None, bl_lens = np.array(bl_lens) if bl_weights is not None: # check consistency of baseline-related inputs - assert len(bl_weights) == nbls, "bl_weights and bl_lens "\ - "must have same length" + assert len(bl_weights) == nbls, ( + "bl_weights and bl_lens " "must have same length" + ) bl_weights = np.array(bl_weights) else: # each baseline length has weight one @@ -1170,31 +1277,45 @@ def run_and_write(self, filepath, bl_lens, bl_weights=None, # as a function of (kperp, kpara) cyl_wf = np.zeros((nbls, self.Nfreqs, nbins_kperp, nbins_kpara)) for ib in range(nbls): - cyl_wf[ib, :, :, :] = self.get_cylindrical_wf(bl_lens[ib], - kperp_bins*self.kunits, - kpara_bins*self.kunits)*bl_weights[ib] + cyl_wf[ib, :, :, :] = ( + self.get_cylindrical_wf( + bl_lens[ib], kperp_bins * self.kunits, kpara_bins * self.kunits + ) + * bl_weights[ib] + ) # Write file - with h5py.File(filepath, 'w') as f: + with h5py.File(filepath, "w") as f: # parameters - f.attrs['avg_nu'] = self.avg_nu - f.attrs['avg_z'] = self.avg_z - f.attrs['nfreqs'] = self.Nfreqs - f.attrs['little_h'] = self.little_h - f.attrs['polpair'] = uvputils.polpair_tuple2int(self.pols) + f.attrs["avg_nu"] = self.avg_nu + f.attrs["avg_z"] = self.avg_z + f.attrs["nfreqs"] = self.Nfreqs + f.attrs["little_h"] = self.little_h + f.attrs["polpair"] = uvputils.polpair_tuple2int(self.pols) # cannot write None to file if self.taper is None: - f.attrs['taper'] = 'none' + f.attrs["taper"] = "none" else: - f.attrs['taper'] = self.taper + f.attrs["taper"] = self.taper # arrays - f.create_dataset('kperp_bins',shape=(nbins_kperp,),data=kperp_bins,dtype=float) - f.create_dataset('kpara_bins',shape=(nbins_kpara,),data=kpara_bins,dtype=float) - f.create_dataset('bl_lens',shape=(nbls,),data=bl_lens,dtype=float) - f.create_dataset('bl_weights',shape=(nbls,),data=bl_weights,dtype=float) - f.create_dataset('dly_array',shape=(self.Nfreqs,),data=self.dly_array,dtype=float) - f.create_dataset('cyl_wf',shape=(nbls,self.Nfreqs,nbins_kperp,nbins_kpara),data=cyl_wf,dtype=float) - + f.create_dataset( + "kperp_bins", shape=(nbins_kperp,), data=kperp_bins, dtype=float + ) + f.create_dataset( + "kpara_bins", shape=(nbins_kpara,), data=kpara_bins, dtype=float + ) + f.create_dataset("bl_lens", shape=(nbls,), data=bl_lens, dtype=float) + f.create_dataset("bl_weights", shape=(nbls,), data=bl_weights, dtype=float) + f.create_dataset( + "dly_array", shape=(self.Nfreqs,), data=self.dly_array, dtype=float + ) + f.create_dataset( + "cyl_wf", + shape=(nbls, self.Nfreqs, nbins_kperp, nbins_kpara), + data=cyl_wf, + dtype=float, + ) + def check_kunits(self, karray): """ Check unit consistency between k's throughout code. @@ -1207,9 +1328,11 @@ def check_kunits(self, karray): try: karray.unit except AttributeError: - raise AttributeError('Feed k array with units (astropy.units).') - assert self.kunits.is_equivalent(karray.unit),\ - "k array units not consistent with little_h" + raise AttributeError("Feed k array with units (astropy.units).") + assert self.kunits.is_equivalent( + karray.unit + ), "k array units not consistent with little_h" + def check_spw_range(spw_range, bandwidth=None): """ @@ -1236,18 +1359,14 @@ def check_spw_range(spw_range, bandwidth=None): # check format if np.size(spw_range) != 2: return False - if spw_range[1]-spw_range[0] <= 0: + if spw_range[1] - spw_range[0] <= 0: return False if min(spw_range) < 0: - return False - # bandwidth-related checks + return False + # bandwidth-related checks if bandwidth is not None: bandwidth = np.array(bandwidth) if max(spw_range) > bandwidth.size: return False return True - - - - diff --git a/pipelines/web_ui/data.js b/pipelines/web_ui/data.js index af201374..f4b8a1a4 100644 --- a/pipelines/web_ui/data.js +++ b/pipelines/web_ui/data.js @@ -1,16 +1,16 @@ function Data() { - + this.root = 'zen.grp1.of1.', this.prefixes = ['uvOCRSLT', 'uvOCRSLTF'], this.prefix_names = ['Time-averaged', 'Time-avg. + delay filtered'], this.pols = ['XX', 'YY', 'pI', 'pQ'], this.red_bls = [ - [ [11,12], [12,13] ], - [ [1,2], [2,3], [3,4] ], + [ [11,12], [12,13] ], + [ [1,2], [2,3], [3,4] ], [ [100,101], [101,102] ], ], - this.lsts = [ '1.00641', '1.10036', '1.19432', '1.28828', '1.38224', '1.47620', + this.lsts = [ '1.00641', '1.10036', '1.19432', '1.28828', '1.38224', '1.47620', '1.57015', '1.66411', '1.75807', '1.85203', '1.94599' ]; - //this.lsts = [ '1.01267', '1.15361', '1.29454', '1.43548', '1.57642', + //this.lsts = [ '1.01267', '1.15361', '1.29454', '1.43548', '1.57642', // '1.71736', '1.85829', '1.99923' ]; }; diff --git a/pipelines/web_ui/hera_webui.js b/pipelines/web_ui/hera_webui.js index 9e38198f..e8f67074 100644 --- a/pipelines/web_ui/hera_webui.js +++ b/pipelines/web_ui/hera_webui.js @@ -8,7 +8,7 @@ function clear_div(div_name){ function load_image(basename, type, ant1, ant2){ // Create image object for a given baseline and type - + //////////////////////////////////////////// // FIXME: Testing ant1 = 11; @@ -16,11 +16,11 @@ function load_image(basename, type, ant1, ant2){ test_dir = "/home/phil/Desktop/plotty/"; // END FIXME //////////////////////////////////////////// - + function raise_image_error(){ console.log("Image load error."); } // FIXME - + var img = document.createElement("img"); img.src = test_dir + basename + "." + type + "." + ant1 + "." + ant2 + ".png"; img.id = img.src; @@ -31,7 +31,7 @@ function load_image(basename, type, ant1, ant2){ //function add_select(opts, opt_names, callback){ // // Add a dropdown menu to the document // var sel = document.createElement("select"); -// +// // // Add options from list // for (i=0; i < opts.length; i++){ // var opt = document.createElement('option'); @@ -39,13 +39,13 @@ function load_image(basename, type, ant1, ant2){ // opt.innerHTML = opt_names[i]; // sel.add(opt); // } -// +// // // Function to return selected value to a user-defined callback function // function select_callback(){ // var val = drop.value; // callback(val); // } -// +// // // Attach callback // sel.addEventListener("change", select_callback); // return sel; @@ -53,32 +53,32 @@ function load_image(basename, type, ant1, ant2){ function select_callback(dropdown, item){ mark_dropdown_selected(dropdown, item); - + // Display a grid of plots dropdown_grid(); } function populate_dropdown(dropdown, opts, opt_names, callback_fn){ // Add elements to dropdown list - + // Get the dropdown and clear it, and mark it as enabled sel = clear_dropdown(dropdown, true); - + // Add default item var item = document.createElement('a'); item.innerHTML = "—"; item.className = "w3-bar-item w3-button"; - item.setAttribute("href", "javascript:" + callback_fn + "('" + item.setAttribute("href", "javascript:" + callback_fn + "('" + dropdown + "', 'none');"); item.id = dropdown + "-none"; sel.appendChild(item); - + // Add options from list for (i=0; i < opts.length; i++){ var item = document.createElement('a'); item.innerHTML = opt_names[i]; item.className = "w3-bar-item w3-button"; - item.setAttribute("href", "javascript:" + callback_fn + "('" + item.setAttribute("href", "javascript:" + callback_fn + "('" + dropdown + "', '" + opts[i] + "');"); item.id = dropdown + "-" + opts[i]; sel.appendChild(item); @@ -87,11 +87,11 @@ function populate_dropdown(dropdown, opts, opt_names, callback_fn){ function clear_dropdown(dropdown, enabled){ // Disable a dropdown menu - + // Get select and clear it var sel = document.getElementById(dropdown + "-select"); while(sel.firstChild) { sel.removeChild(sel.firstChild); } - + // Grey-out the dropdown button and clear the info box if disabled var btn = document.getElementById(dropdown + "-button"); var info = document.getElementById(dropdown + "-status"); @@ -103,24 +103,24 @@ function clear_dropdown(dropdown, enabled){ info.style.opacity = 0.5; } info.innerHTML = "—"; - + // Return the dropdown return sel; } function mark_dropdown_selected(dropdown, option){ // Mark a given argument of a dropdown as selected - + // Loop over items in dropdown and make sure they're not highlighted var list_elements = document.getElementById(dropdown + "-select").children; for(i=0; i < list_elements.length; i++){ list_elements[i].style.fontWeight = 'normal'; } - + // Highlight selected item var item = document.getElementById(dropdown + "-" + option); item.style.fontWeight = 'bold'; - + // Change status indicator var status = document.getElementById(dropdown + "-status"); status.innerHTML = item.innerHTML; @@ -136,44 +136,44 @@ function move_next_item(evt, dropdown, array){ // Catch keypress events and go back or forward in some variable var e = evt || window.event; var c = e.keyCode; - + // Left or right key if ((c == 39) || (c == 37)){ - + // Get currently selected item var cur = get_selected_item(dropdown); if (cur[0] == -1){ return; } var i = cur[0] - 1; - + // Increment/decrement counter if (c == 39){ if (i >= array.length - 1){ i = 0; }else{ i++; } // Right key }else{ if (i == 0){ i = array.length - 1; }else{ i--; } // Left key } - + // Select new item and return mark_dropdown_selected(dropdown, array[i]); dropdown_grid(); return; - + } // end left/right check - + /* if ((c == 38) || (c == 40)){ - + // Get currently selected blgroup var redgrp = get_selected_item("redgrp"); if (redgrp[0] == -1){ return; } var i = redgrp[0] - 1; - + // Increment/decrement counter if (c == 38){ if (i >= dat.red_bls.length - 1){ i = 0; }else{ i++; } // Up key }else{ if (i == 0){ i = dat.lsts.length - 1; }else{ i--; } // Down key } - + // Select new red. group and return mark_dropdown_selected("redgrp", dat.red_bls[i]); dropdown_grid(); @@ -183,9 +183,9 @@ function move_next_item(evt, dropdown, array){ } /* -LST-binned data -RFI flagging Time averaging @@ -198,35 +198,35 @@ function move_next_item(evt, dropdown, array){ function list_baselines(){ // List all available baselines clear_div("flow"); - + var divList = document.createElement("div"); divList.className = "bl-list"; - + // List title var hTitle = document.createElement("h2"); hTitle.appendChild(document.createTextNode("Baselines")); - + // Get data structure var dat = new Data(); - + // Build dropdown select list populate_dropdown("pipeline", dat.prefixes, dat.prefix_names, "select_callback"); populate_dropdown("lst", dat.lsts, dat.lsts, "select_callback"); populate_dropdown("pol", dat.pols, dat.pols, "select_callback"); populate_dropdown("redgrp", dat.red_bls, dat.red_bls, "select_callback"); - + // Create list of redundant groups lstRedGroups = document.createElement("ul"); - + // Populate list for(i=0; i < dat.red_bls.length; i++){ - + // Create sub-list for this redundant group liRedGroups = document.createElement("li"); var hGroup = document.createElement("h3"); hGroup.appendChild(document.createTextNode("Redundant group " + i)); lstBls = document.createElement("ul"); - + // Populate sub-list with redundant baselines for(j=0; j < dat.red_bls[i].length; j++){ var li = document.createElement("li"); @@ -236,12 +236,12 @@ function list_baselines(){ document.createTextNode( "Baseline (" + ant1 + ", " + ant2 + ")" ) ); lstBls.appendChild(li); } // end loop over bls - + liRedGroups.appendChild(hGroup); liRedGroups.appendChild(lstBls); lstRedGroups.appendChild(liRedGroups); } // end loop over redundant groups - + // Attach to document divList.appendChild(hTitle); //divList.appendChild(sel); @@ -250,9 +250,9 @@ function list_baselines(){ } function get_selected_item(dropdown){ - // Get the item selected in a drop-down menu. Returns 2-element array, + // Get the item selected in a drop-down menu. Returns 2-element array, // [index, item]. - + var items = document.getElementById(dropdown + "-select").children; var selected = dropdown + "-none"; var idx = -1; @@ -262,7 +262,7 @@ function get_selected_item(dropdown){ idx = i; } } - + // Parse the selected value selected = selected.slice(dropdown.length + 1); if (selected == "none"){ idx = -1; } @@ -273,27 +273,27 @@ function get_selected_item(dropdown){ function dropdown_grid(){ // Show a grid of plots according to which drop-down options are selected // grid_red_bls("zen.grp1.of1.xx.LST.1.71736", "uvOCRSLTF.XX"); - + // Get selected items (idx, name) var pipe = get_selected_item("pipeline"); var lst = get_selected_item("lst"); var redgrp = get_selected_item("redgrp"); var pol = get_selected_item("pol"); - + // FIXME: Should be different cases depending on what's selected if ((pipe[0] != -1) && (lst[0] != -1) && (pol[0] != -1)){ var dat = new Data(); - + // Special case for xx and yy polarization strings var pol1 = pol[1], pol2 = pol[1]; if ((pol1.toLowerCase() == 'xx') || (pol1.toLowerCase() == 'yy')){ pol1 = pol1.toLowerCase(); pol2 = pol2.toUpperCase(); } - + // Construct basename and type var basename = dat.root + pol1 + ".LST." + lst[1]; var type = pipe[1] + "." + pol2; - + // Show grid of redundant baselines (or just one, if requested) if (redgrp[0] == -1){ grid_red_bls(basename, type, "all"); // show all @@ -302,7 +302,7 @@ function dropdown_grid(){ } return; } - + // By default, show nothing except a hint clear_div("flow"); var msg_box = message_box("Select some options", "Please select some combination of LST, polarization, baseline group, and pipeline stage."); @@ -311,75 +311,75 @@ function dropdown_grid(){ function message_box(title, text){ // Show a message box - + // Container var box = document.createElement("div"); box.className = "w3-panel w3-padding-16 w3-blue-gray"; box.style.width = "70%"; - + // Title var hTitle = document.createElement("h3"); hTitle.appendChild( document.createTextNode(title) ); box.appendChild(hTitle); - + // Text var ptext = document.createElement("p"); ptext.appendChild( document.createTextNode(text) ); box.appendChild(ptext); - + return box; } function grid_red_bls(basename, type, idx){ // List all available baselines (or only a particular index if specified) - + clear_div("flow"); - + var divRedGroups = document.createElement("div"); divRedGroups.className = "red-groups"; - + // Grid title var hTitle = document.createElement("h2"); hTitle.appendChild( document.createTextNode("Grid:" + basename + " " + type) ); - + // Get data structure var dat = new Data(); - + // Populate grid for(i=0; i < dat.red_bls.length; i++){ if ((idx != "all") && (i != idx)){ continue; } - + // Create div for this redundant group var divGroup = document.createElement("div"); divGroup.className = "red-group"; divGroup.id = i; var hGroup = document.createElement("h3"); hGroup.appendChild(document.createTextNode("Redundant group " + i)); - + // Populate div with plot for each baseline var divGroupImgs = document.createElement("div"); divGroupImgs.className = "red-group-imgs"; divGroupImgs.id = i; for(j=0; j < dat.red_bls[i].length; j++){ - + // Load image - img = load_image(basename, type, - dat.red_bls[i][j][0], + img = load_image(basename, type, + dat.red_bls[i][j][0], dat.red_bls[i][j][1]); - + img.href = "#img" + j; img.style.width = "20%"; - + divGroupImgs.appendChild(img); } // end loop over bls - + divGroup.appendChild(hGroup); divGroup.appendChild(divGroupImgs); divRedGroups.appendChild(divGroup); } // end loop over redundant groups - + // Attach to document document.getElementById("flow").appendChild(divRedGroups); } @@ -387,35 +387,35 @@ function grid_red_bls(basename, type, idx){ function add_plot(basename, type, ant1, ant2){ // Create a new div with a plot and basic info - + var divPlot = document.createElement("div"); divPlot.className = "plot"; - + // Plot title var hTitle = document.createElement("h2"); var txtTitle = document.createTextNode(type); hTitle.appendChild(txtTitle); - + // Load plot image and put in container var divInner = document.createElement("div"); divInner.setAttribute("class", "img-zoom-container"); divInner.style.position = "relative"; - + var imgPlot = load_image(basename, type, ant1, ant2); //var imgPlot = document.createElement("img"); //imgPlot.src = basename + "." + type + "." + ant1 + "." + ant2 + ".png"; //imgPlot.id = imgPlot.src; imgPlot.style.width = "40%"; - + // Don't attach the zoom capability until the image has loaded function attach_zoom(){ imageZoom(imgPlot.id, "zoom-window"); } imgPlot.onmouseover = attach_zoom; - - // Build div + + // Build div divPlot.appendChild(hTitle); divInner.appendChild(imgPlot); divPlot.appendChild(divInner); - + // Append div to main document document.getElementById("flow").appendChild(divPlot); } // end function @@ -423,19 +423,19 @@ function add_plot(basename, type, ant1, ant2){ /* function load_baseline(ant1, ant2){ // Load summary data and plots for a given baseline - + clear_div("flow"); flow = document.getElementById("flow"); var bl_title = document.createElement("h2"); bl_title.id = "blname"; bl_title.innerHTML = "Baseline (" + ant1 + ", " + ant2 + ")"; flow.appendChild(bl_title); - + add_plot("zen.grp1.of1.xx.LST.1.71736", "uvOCRSLTF.XX", ant1, ant2); add_plot("zen.grp1.of1.yy.LST.1.71736", "uvOCRSLTF.YY", ant1, ant2); add_plot("zen.grp1.of1.pI.LST.1.71736", "uvOCRSLTF.pI", ant1, ant2); add_plot("zen.grp1.of1.xx.tavg", "uvOCRSL.XX", ant1, ant2); - + } // end load baseline */ @@ -443,11 +443,11 @@ function load_baseline(ant1, ant2){ // Code adapted from: https://www.w3schools.com/howto/howto_js_image_zoom.asp function imageZoom(imgID, resultID) { - + var img, lens, result, cx, cy; img = document.getElementById(imgID); result = document.getElementById(resultID); - + // Check if this image has already been loaded into the zoom window if (result.style.backgroundImage == "url(\"" + img.src + "\")"){ return; @@ -477,46 +477,46 @@ function imageZoom(imgID, resultID) { img.addEventListener("mousemove", moveLens); lens.addEventListener("touchmove", moveLens); img.addEventListener("touchmove", moveLens); - + function moveLens(e) { var pos, x, y; // Prevent any other actions that may occur when moving over the image e.preventDefault(); - + // Get the cursor's x and y positions pos = getCursorPos(e); - + // Make sure lens is visible lens.style.visible = true; - + // Calculate the position of the lens x = pos.x - (lens.offsetWidth / 2); y = pos.y - (lens.offsetHeight / 2); - + // Prevent the lens from being positioned outside the image if (x > img.width - lens.offsetWidth) {x = img.width - lens.offsetWidth;} if (x < 0) {x = 0;} if (y > img.height - lens.offsetHeight) {y = img.height - lens.offsetHeight;} if (y < 0) {y = 0;} - + // Set the position of the lens lens.style.left = x + "px"; lens.style.top = y + "px"; - + // Display what the lens "sees" result.style.backgroundPosition = "-" + (x * cx) + "px -" + (y * cy) + "px"; } - + function getCursorPos(e) { var a, x = 0, y = 0; e = e || window.event; // Get the x and y positions of the image: a = img.getBoundingClientRect(); - + // Calculate the cursor's x and y coordinates, relative to the image: x = e.pageX - a.left; y = e.pageY - a.top; - + // Consider any page scrolling: x = x - window.pageXOffset; y = y - window.pageYOffset; diff --git a/pipelines/web_ui/report.html b/pipelines/web_ui/report.html index 59748ba9..0a4750e8 100644 --- a/pipelines/web_ui/report.html +++ b/pipelines/web_ui/report.html @@ -2,7 +2,7 @@ HERA - + @@ -10,14 +10,14 @@ - + @@ -95,11 +95,11 @@

GRID VIEW

- + - +
- +
- +
- +
- +
- +
- + - +

Start here.

- +
- +