├── africanus ├── model │ ├── __init__.py │ ├── coherency │ │ ├── tests │ │ │ └── __init__.py │ │ ├── cuda │ │ │ ├── tests │ │ │ │ ├── __init__.py │ │ │ │ └── test_convert.py │ │ │ ├── __init__.py │ │ │ └── conversion.cu.j2 │ │ ├── __init__.py │ │ └── dask.py │ ├── shape │ │ ├── tests │ │ │ ├── __init__.py │ │ │ └── test_gaussian_shape.py │ │ ├── __init__.py │ │ ├── dask.py │ │ └── gaussian_shape.py │ ├── spectral │ │ ├── tests │ │ │ ├── __init__.py │ │ │ └── test_spectral_model.py │ │ ├── __init__.py │ │ └── dask.py │ ├── wsclean │ │ ├── tests │ │ │ ├── __init__.py │ │ │ ├── conftest.py │ │ │ ├── test_spectral_model.py │ │ │ └── test_wsclean_model_file.py │ │ ├── __init__.py │ │ └── dask.py │ └── spi │ │ ├── examples │ │ ├── __init__.py │ │ └── README.rst │ │ ├── __init__.py │ │ ├── dask.py │ │ └── tests │ │ └── test_component_spi.py ├── util │ ├── __init__.py │ ├── tests │ │ ├── test_cub_download.py │ │ ├── test_trove_download.py │ │ ├── test_type_inference.py │ │ ├── test_progress_bar.py │ │ ├── test_nvcc_compiler.py │ │ ├── test_util_shape.py │ │ ├── test_requirements.py │ │ └── test_patterns.py │ ├── jax_init.py │ ├── files.py │ ├── appdirs.py │ ├── type_inference.py │ ├── testing.py │ ├── casa_types.py │ ├── numba.py │ ├── docs.py │ ├── cuda.py │ ├── shapes.py │ ├── code.py │ └── jinja2.py ├── rime │ ├── jax │ │ ├── __init__.py │ │ ├── tests │ │ │ ├── __init__.py │ │ │ └── test_jax_phase_delay.py │ │ ├── README.rst │ │ └── phase.py │ ├── tests │ │ ├── __init__.py │ │ ├── conftest.py │ │ └── test_rime.py │ ├── examples │ │ ├── __init__.py │ │ ├── README.rst │ │ ├── tests │ │ │ ├── test_examples.py │ │ │ └── meqtrees │ │ │ │ └── tdlconf.profiles │ │ └── N6251-sky-model.txt │ ├── cuda │ │ ├── __init__.py │ │ ├── tests │ │ │ ├── test_cuda_feed_rotation.py │ │ │ ├── test_cuda_phase_delay.py │ │ │ ├── test_warp_transpose.cu.j2 │ │ │ ├── test_macros.py │ │ │ ├── test_cuda_beam.py │ │ │ └── test_cuda_predict.py │ │ ├── feeds.cu.j2 │ │ ├── phase.cu.j2 │ │ ├── beam_freq_interp.cu.j2 │ │ ├── feeds.py │ │ └── phase.py │ ├── __init__.py │ ├── parangles_astropy.py │ ├── parangles_casa.py │ ├── parangles.py │ ├── transform.py │ └── feeds.py ├── testing │ ├── __init__.py │ └── tests │ │ ├── __init__.py │ │ └── test_beam_factory.py ├── averaging │ ├── tests │ │ ├── __init__.py │ │ ├── test_splines.py │ │ └── test_support.py │ ├── __init__.py │ └── support.py ├── experimental │ ├── __init__.py │ └── rime │ │ ├── __init__.py │ │ └── fused │ │ ├── terms │ │ ├── __init__.py │ │ ├── feed_rotation.py │ │ ├── gaussian.py │ │ └── phase.py │ │ ├── tests │ │ ├── __init__.py │ │ ├── test_transformer.py │ │ └── test_structref_setter.py │ │ ├── transformers │ │ ├── __init__.py │ │ └── lm.py │ │ ├── error.py │ │ ├── __init__.py │ │ └── dask.py ├── gridding │ ├── nifty │ │ ├── __init__.py │ │ └── tests │ │ │ └── __init__.py │ ├── perleypolyhedron │ │ ├── __init__.py │ │ ├── tests │ │ │ └── __init__.py │ │ └── policies │ │ │ ├── __init__.py │ │ │ └── phase_transform_policies.py │ ├── __init__.py │ ├── wgridder │ │ ├── tests │ │ │ └── __init__.py │ │ └── __init__.py │ └── util.py ├── calibration │ ├── tests │ │ └── __init__.py │ ├── __init__.py │ ├── utils │ │ ├── tests │ │ │ └── __init__.py │ │ ├── __init__.py │ │ └── utils.py │ └── phase_only │ │ ├── tests │ │ └── __init__.py │ │ ├── __init__.py │ │ └── dask.py ├── deconv │ ├── __init__.py │ └── hogbom │ │ └── __init__.py ├── constants │ ├── __init__.py │ └── consts.py ├── dft │ ├── __init__.py │ └── dask.py ├── linalg │ └── __init__.py ├── gps │ ├── __init__.py │ ├── utils.py │ ├── kernels.py │ └── examples │ │ └── generate_phase_only_gains.py ├── install │ └── __init__.py ├── coordinates │ ├── __init__.py │ └── dask.py ├── __init__.py └── conftest.py ├── docs ├── authors.rst ├── history.rst ├── readme.rst ├── contributing.rst ├── usage.rst ├── deconv-api.rst ├── api.rst ├── cmdline-utils.rst ├── gps-api.rst ├── index.rst ├── Makefile ├── coordinates-api.rst ├── make.bat ├── linalg-api.rst ├── installation.rst ├── gridding-api.rst ├── util-api.rst ├── dft-api.rst ├── rime-api.rst ├── calibration-api.rst └── model-api.rst ├── MANIFEST.in ├── .github ├── workflows │ ├── pre-commit.yml │ └── readthedocs.yml ├── ISSUE_TEMPLATE.md ├── dependabot.yml └── PULL_REQUEST_TEMPLATE.md ├── .editorconfig ├── AUTHORS.rst ├── .pre-commit-config.yaml ├── .readthedocs.yaml ├── README.rst ├── CITATION.cff ├── LICENSE ├── .gitignore ├── Makefile └── pyproject.toml /africanus/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/util/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/rime/jax/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/rime/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/testing/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/averaging/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/experimental/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/gridding/nifty/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/rime/examples/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/rime/jax/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/testing/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/calibration/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/experimental/rime/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/gridding/nifty/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/model/coherency/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/model/shape/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/model/spectral/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/model/wsclean/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/gridding/perleypolyhedron/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/model/coherency/cuda/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /docs/authors.rst: -------------------------------------------------------------------------------- 1 | .. include:: ../AUTHORS.rst 2 | -------------------------------------------------------------------------------- /docs/history.rst: -------------------------------------------------------------------------------- 1 | .. include:: ../HISTORY.rst 2 | -------------------------------------------------------------------------------- /docs/readme.rst: -------------------------------------------------------------------------------- 1 | .. include:: ../README.rst 2 | -------------------------------------------------------------------------------- /africanus/calibration/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/terms/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/gridding/perleypolyhedron/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/deconv/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/transformers/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /africanus/gridding/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | -------------------------------------------------------------------------------- /africanus/gridding/perleypolyhedron/policies/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /docs/contributing.rst: -------------------------------------------------------------------------------- 1 | .. include:: ../CONTRIBUTING.rst 2 | -------------------------------------------------------------------------------- /africanus/model/spi/examples/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | -------------------------------------------------------------------------------- /africanus/calibration/utils/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | -------------------------------------------------------------------------------- /africanus/constants/__init__.py: -------------------------------------------------------------------------------- 1 | from .consts import * # noqa 2 | -------------------------------------------------------------------------------- /africanus/gridding/wgridder/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | -------------------------------------------------------------------------------- /africanus/calibration/phase_only/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | -------------------------------------------------------------------------------- /africanus/dft/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from .kernels import im_to_vis, vis_to_im 4 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/error.py: -------------------------------------------------------------------------------- 1 | class InvalidSignature(ValueError): 2 | pass 3 | -------------------------------------------------------------------------------- /africanus/model/coherency/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from .conversion import convert 4 | -------------------------------------------------------------------------------- /docs/usage.rst: -------------------------------------------------------------------------------- 1 | ===== 2 | Usage 3 | ===== 4 | 5 | To use Codex Africanus in a project:: 6 | 7 | import africanus 8 | -------------------------------------------------------------------------------- /africanus/linalg/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from africanus.linalg.kronecker_tools import kron_matvec, kron_cholesky 4 | -------------------------------------------------------------------------------- /africanus/model/spi/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from africanus.model.spi.component_spi import fit_spi_components 4 | -------------------------------------------------------------------------------- /africanus/model/coherency/cuda/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from africanus.model.coherency.cuda.conversion import convert 4 | -------------------------------------------------------------------------------- /africanus/deconv/hogbom/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | __all__ = ["hogbom_clean"] 4 | 5 | from .clean import hogbom_clean 6 | -------------------------------------------------------------------------------- /africanus/model/spectral/__init__.py: -------------------------------------------------------------------------------- 1 | __all__ = ["spectral_model"] 2 | 3 | from africanus.model.spectral.spec_model import spectral_model 4 | -------------------------------------------------------------------------------- /africanus/gps/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from africanus.gps.kernels import exponential_squared 4 | from africanus.gps.utils import abs_diff 5 | -------------------------------------------------------------------------------- /africanus/util/tests/test_cub_download.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | def test_cub_download(): 5 | from africanus.util.cub import cub_dir 6 | 7 | cub_dir() 8 | -------------------------------------------------------------------------------- /africanus/util/tests/test_trove_download.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | def test_trove_download(): 5 | from africanus.util.trove import trove_dir 6 | 7 | trove_dir() 8 | -------------------------------------------------------------------------------- /africanus/install/__init__.py: -------------------------------------------------------------------------------- 1 | # NOTE(sjperkins) 2 | # Imports at this level should be avoided, 3 | # or should fail gracefully as functionality 4 | # in these modules is called by setup.py 5 | -------------------------------------------------------------------------------- /africanus/model/wsclean/__init__.py: -------------------------------------------------------------------------------- 1 | __all__ = ["load", "spectra"] 2 | 3 | from africanus.model.wsclean.file_model import load 4 | from africanus.model.wsclean.spec_model import spectra 5 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/__init__.py: -------------------------------------------------------------------------------- 1 | from africanus.experimental.rime.fused.specification import RimeSpecification # noqa 2 | from africanus.experimental.rime.fused.core import rime # noqa 3 | -------------------------------------------------------------------------------- /docs/deconv-api.rst: -------------------------------------------------------------------------------- 1 | ------------------------ 2 | Deconvolution Algorithms 3 | ------------------------ 4 | 5 | .. currentmodule:: africanus.deconv.hogbom 6 | 7 | 8 | .. autofunction:: hogbom_clean 9 | -------------------------------------------------------------------------------- /africanus/model/shape/__init__.py: -------------------------------------------------------------------------------- 1 | __all__ = ["gaussian", "shapelet", "shapelet_1d"] 2 | 3 | from africanus.model.shape.gaussian_shape import gaussian 4 | from africanus.model.shape.shapelets import shapelet, shapelet_1d 5 | -------------------------------------------------------------------------------- /africanus/averaging/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | __all__ = ["bda", "time_and_channel"] 4 | 5 | from africanus.averaging.bda_avg import bda 6 | from africanus.averaging.time_and_channel_avg import time_and_channel 7 | -------------------------------------------------------------------------------- /africanus/rime/cuda/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from africanus.rime.cuda.beam import beam_cube_dde 4 | from africanus.rime.cuda.feeds import feed_rotation 5 | from africanus.rime.cuda.phase import phase_delay 6 | from africanus.rime.cuda.predict import predict_vis 7 | -------------------------------------------------------------------------------- /africanus/rime/jax/README.rst: -------------------------------------------------------------------------------- 1 | Experimental RIME with Google JAX 2 | ================================= 3 | 4 | Google `jax `_ supports 5 | 6 | - jit compilation of numpy code to both CPU and GPU targets 7 | - autodiff gradients, jacobians and hessians 8 | -------------------------------------------------------------------------------- /africanus/gridding/wgridder/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from africanus.gridding.wgridder.im2vis import model 4 | from africanus.gridding.wgridder.vis2im import dirty 5 | from africanus.gridding.wgridder.im2residim import residual 6 | from africanus.gridding.wgridder.hessian import hessian 7 | -------------------------------------------------------------------------------- /africanus/coordinates/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from africanus.coordinates.coordinates import radec_to_lmn 4 | from africanus.coordinates.coordinates import radec_to_lm 5 | from africanus.coordinates.coordinates import lmn_to_radec 6 | from africanus.coordinates.coordinates import lm_to_radec 7 | -------------------------------------------------------------------------------- /africanus/constants/consts.py: -------------------------------------------------------------------------------- 1 | __all__ = ["c", "minus_two_pi_over_c", "two_pi_over_c"] 2 | 3 | import math 4 | 5 | # Lightspeed 6 | c = 2.99792458e8 7 | 8 | two_pi_over_c = 2 * math.pi / c 9 | minus_two_pi_over_c = -two_pi_over_c 10 | 11 | DEG2RAD = math.pi / 180.0 12 | ARCSEC2RAD = math.pi / (180 * 3600) 13 | -------------------------------------------------------------------------------- /africanus/util/jax_init.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | """ 4 | Configure jax to use default 64 bit precision 5 | https://github.com/google/jax/issues/216 6 | """ 7 | 8 | 9 | try: 10 | import jax 11 | except ImportError: 12 | pass 13 | else: 14 | jax.config.update("jax_enable_x64", True) 15 | -------------------------------------------------------------------------------- /africanus/rime/examples/README.rst: -------------------------------------------------------------------------------- 1 | Predict Script 2 | ============== 3 | 4 | Predicts from sources.txt. 5 | 6 | .. code-block:: bash 7 | 8 | $ python predict.py -sm sky-model.txt OBSERVATION.MS 9 | $ wsclean -weight uniform -name obs -size 512 512 -scale 1.5asec -make-psf -data-column MODEL_DATA OBSERVATION.MS 10 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/tests/test_transformer.py: -------------------------------------------------------------------------------- 1 | from africanus.experimental.rime.fused.transformers.lm import LMTransformer 2 | 3 | 4 | def test_transformers(): 5 | T = LMTransformer() # noqa 6 | assert T.ARGS == ("radec", "phase_dir") 7 | assert T.KWARGS == {} 8 | assert T.ALL_ARGS == ("radec", "phase_dir") 9 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include AUTHORS.rst 2 | include CONTRIBUTING.rst 3 | include HISTORY.rst 4 | include LICENSE 5 | include README.rst 6 | 7 | recursive-include tests * 8 | recursive-include * *.j2 9 | recursive-exclude * __pycache__ 10 | recursive-exclude * *.py[co] 11 | 12 | recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif 13 | -------------------------------------------------------------------------------- /africanus/calibration/phase_only/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from africanus.calibration.phase_only.phase_only import gauss_newton 4 | from africanus.calibration.phase_only.phase_only import compute_jhj 5 | from africanus.calibration.phase_only.phase_only import compute_jhr 6 | from africanus.calibration.phase_only.phase_only import compute_jhj_and_jhr 7 | -------------------------------------------------------------------------------- /.github/workflows/pre-commit.yml: -------------------------------------------------------------------------------- 1 | name: pre-commit 2 | 3 | on: [push] 4 | 5 | jobs: 6 | pre-commit: 7 | runs-on: ubuntu-latest 8 | steps: 9 | - uses: actions/checkout@v6 10 | with: 11 | fetch-depth: 1 12 | - uses: actions/setup-python@v6.1.0 13 | with: 14 | python-version: 3.11 15 | - uses: pre-commit/action@v3.0.1 16 | -------------------------------------------------------------------------------- /africanus/calibration/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from .utils import check_type, chunkify_rows 4 | from africanus.calibration.utils.corrupt_vis import corrupt_vis 5 | from africanus.calibration.utils.correct_vis import correct_vis 6 | from africanus.calibration.utils.residual_vis import residual_vis 7 | from .compute_and_corrupt_vis import compute_and_corrupt_vis 8 | -------------------------------------------------------------------------------- /docs/api.rst: -------------------------------------------------------------------------------- 1 | API 2 | === 3 | 4 | .. toctree:: 5 | :maxdepth: 1 6 | 7 | rime-api.rst 8 | dft-api.rst 9 | gridding-api.rst 10 | deconv-api.rst 11 | coordinates-api.rst 12 | model-api.rst 13 | averaging-api.rst 14 | util-api.rst 15 | calibration-api.rst 16 | linalg-api.rst 17 | gps-api.rst 18 | 19 | experimental.rst 20 | -------------------------------------------------------------------------------- /africanus/rime/examples/tests/test_examples.py: -------------------------------------------------------------------------------- 1 | def test_examples(): 2 | """ 3 | Test that import works at least as this should flush out 4 | other import issues caused by renames etc. 5 | """ 6 | 7 | from africanus.rime.examples.predict import predict # noqa 8 | 9 | # TODO(sjperkins) 10 | # Call with a fake MS once pytest-ms is available 11 | -------------------------------------------------------------------------------- /.editorconfig: -------------------------------------------------------------------------------- 1 | # http://editorconfig.org 2 | 3 | root = true 4 | 5 | [*] 6 | indent_style = space 7 | indent_size = 4 8 | trim_trailing_whitespace = true 9 | insert_final_newline = true 10 | charset = utf-8 11 | end_of_line = lf 12 | 13 | [*.bat] 14 | indent_style = tab 15 | end_of_line = crlf 16 | 17 | [LICENSE] 18 | insert_final_newline = false 19 | 20 | [Makefile] 21 | indent_style = tab 22 | -------------------------------------------------------------------------------- /africanus/util/files.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from hashlib import sha1 5 | 6 | 7 | def sha_hash_file(filename): 8 | """Compute the SHA1 hash of filename""" 9 | hash_sha = sha1() 10 | 11 | with open(filename, "rb") as f: 12 | for chunk in iter(lambda: f.read(1024 * 1024), b""): 13 | hash_sha.update(chunk) 14 | 15 | return hash_sha.hexdigest() 16 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | * Codex Africanus version: 2 | * Python version: 3 | * Operating System: 4 | 5 | ### Description 6 | 7 | Describe what you were trying to get done. 8 | Tell us what happened, what went wrong, and what you expected to happen. 9 | 10 | ### What I Did 11 | 12 | ``` 13 | Paste the command(s) you ran and the output. 14 | If there was a crash, please include the traceback here. 15 | ``` 16 | -------------------------------------------------------------------------------- /docs/cmdline-utils.rst: -------------------------------------------------------------------------------- 1 | Command Line Utilities 2 | ---------------------- 3 | 4 | The following command line utilities are installed. 5 | Run each utility's help for further information. 6 | 7 | .. code-block:: console 8 | 9 | $ utility --help 10 | 11 | plot-filter 12 | ~~~~~~~~~~~ 13 | 14 | Plots convolution filters. 15 | 16 | plot-taper 17 | ~~~~~~~~~~ 18 | 19 | Plots tapers associated with convolution filters. 20 | -------------------------------------------------------------------------------- /docs/gps-api.rst: -------------------------------------------------------------------------------- 1 | ------------------ 2 | Gaussian processes 3 | ------------------ 4 | 5 | This module provides a collection of tools 6 | that are useful when performing Gaussian 7 | process regression. 8 | 9 | Numpy 10 | ~~~~~ 11 | 12 | .. currentmodule:: africanus.gps 13 | 14 | .. autosummary:: 15 | abs_diff 16 | exponential_squared 17 | 18 | .. autofunction:: abs_diff 19 | .. autofunction:: exponential_squared 20 | -------------------------------------------------------------------------------- /AUTHORS.rst: -------------------------------------------------------------------------------- 1 | ======= 2 | Credits 3 | ======= 4 | 5 | Development Lead 6 | ---------------- 7 | 8 | * Simon Perkins 9 | 10 | Contributors 11 | ------------ 12 | 13 | * Landman Bester 14 | * Benjamin Hugo 15 | * Jonathan Kenyon 16 | * Gijs Molenaar 17 | * Joshua van Staden 18 | * Oleg Smirnov 19 | -------------------------------------------------------------------------------- /africanus/util/appdirs.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from os.path import join as pjoin 5 | 6 | from appdirs import AppDirs 7 | 8 | from africanus import __version__ 9 | 10 | _dirs = AppDirs("codex-africanus", "radio-astronomer", __version__) 11 | 12 | user_data_dir = _dirs.user_data_dir 13 | downloads_dir = pjoin(user_data_dir, "downloads") 14 | include_dir = pjoin(user_data_dir, "include") 15 | 16 | del __version__ 17 | del _dirs 18 | -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- 1 | =========================================== 2 | Welcome to Codex Africanus's documentation! 3 | =========================================== 4 | 5 | .. toctree:: 6 | :maxdepth: 2 7 | :caption: Contents: 8 | 9 | installation 10 | usage 11 | cmdline-utils 12 | api 13 | contributing 14 | authors 15 | history 16 | 17 | Indices and tables 18 | ================== 19 | * :ref:`genindex` 20 | * :ref:`modindex` 21 | * :ref:`search` 22 | -------------------------------------------------------------------------------- /africanus/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | """Top-level package for Codex Africanus.""" 4 | 5 | 6 | # NOTE(sjperkins) 7 | # Imports at this level within this module should be avoided, 8 | # or should fail gracefully as this is the base africanus module. 9 | # The setup.py file accesses the ``africanus.install`` modules 10 | import africanus.util.jax_init # noqa 11 | 12 | __author__ = """Simon Perkins""" 13 | __email__ = "sperkins@ska.ac.za" 14 | __version__ = "0.4.3" 15 | -------------------------------------------------------------------------------- /africanus/rime/__init__.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | 3 | from africanus.rime.phase import phase_delay 4 | from africanus.rime.feeds import feed_rotation 5 | from africanus.rime.fast_beam_cubes import beam_cube_dde 6 | from africanus.rime.parangles import parallactic_angles 7 | from africanus.rime.transform import transform_sources 8 | from africanus.rime.zernike import zernike_dde 9 | from africanus.rime.predict import predict_vis, apply_gains 10 | from africanus.rime.wsclean_predict import wsclean_predict 11 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | # See https://pre-commit.com for more information 2 | # See https://pre-commit.com/hooks.html for more hooks 3 | repos: 4 | - repo: https://github.com/pre-commit/pre-commit-hooks 5 | rev: v4.5.0 6 | hooks: 7 | - id: trailing-whitespace 8 | - id: end-of-file-fixer 9 | - id: check-yaml 10 | - id: check-added-large-files 11 | - repo: https://github.com/astral-sh/ruff-pre-commit 12 | rev: v0.1.3 13 | hooks: 14 | - id: ruff-format 15 | name: ruff format 16 | -------------------------------------------------------------------------------- /.readthedocs.yaml: -------------------------------------------------------------------------------- 1 | # .readthedocs.yaml 2 | # Read the Docs configuration file 3 | # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details 4 | 5 | # Required 6 | version: 2 7 | python: 8 | install: 9 | - method: pip 10 | path: . 11 | extra_requirements: 12 | - doc 13 | - testing 14 | 15 | build: 16 | os: ubuntu-24.04 17 | tools: 18 | python: "3.12" 19 | 20 | # Build documentation in the "docs/" directory with Sphinx 21 | sphinx: 22 | configuration: docs/conf.py 23 | -------------------------------------------------------------------------------- /.github/workflows/readthedocs.yml: -------------------------------------------------------------------------------- 1 | # .github/workflows/documentation-links.yml 2 | 3 | name: readthedocs/actions 4 | on: 5 | pull_request_target: 6 | types: 7 | - opened 8 | # Execute this action only on PRs that touch 9 | # documentation files. 10 | # paths: 11 | # - "doc/**" 12 | 13 | permissions: 14 | pull-requests: write 15 | 16 | jobs: 17 | documentation-links: 18 | runs-on: ubuntu-latest 19 | steps: 20 | - uses: readthedocs/actions/preview@v1 21 | with: 22 | project-slug: "codex-africanus" 23 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | # To get started with Dependabot version updates, you'll need to specify which 2 | # package ecosystems to update and where the package manifests are located. 3 | # Please see the documentation for all configuration options: 4 | # https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file 5 | 6 | version: 2 7 | updates: 8 | - package-ecosystem: "github-actions" # See documentation for possible values 9 | directory: "/" # Location of package manifests 10 | schedule: 11 | interval: "weekly" 12 | -------------------------------------------------------------------------------- /africanus/testing/tests/test_beam_factory.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from africanus.testing.beam_factory import beam_factory 5 | 6 | import pytest 7 | 8 | 9 | @pytest.mark.parametrize("pol_type", ["linear", "circular"]) 10 | def test_beam_factory(tmp_path, pol_type): 11 | fits = pytest.importorskip("astropy.io.fits") 12 | schema = tmp_path / "test_beam_$(corr)_$(reim).fits" 13 | 14 | filenames = beam_factory(schema=schema, npix=15, polarisation_type=pol_type) 15 | 16 | for corr, (re_file, im_file) in filenames.items(): 17 | with fits.open(re_file), fits.open(im_file): 18 | pass 19 | -------------------------------------------------------------------------------- /.github/PULL_REQUEST_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | - [ ] Tests added / passed 2 | 3 | ```bash 4 | $ py.test -v -s africanus 5 | ``` 6 | 7 | If the pre-commit tests fail, install and 8 | run the pre-commit hooks in your development 9 | virtuale environment: 10 | 11 | ```bash 12 | $ pip install pre-commit 13 | $ pre-commit install 14 | $ pre-commit run -a 15 | ``` 16 | 17 | - [ ] Fully documented, including `HISTORY.rst` for all changes 18 | and one of the `docs/*-api.rst` files for new API 19 | 20 | To build the docs locally: 21 | 22 | ```bash 23 | pip install -r requirements.readthedocs.txt 24 | cd docs 25 | READTHEDOCS=True make html 26 | ``` 27 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line. 5 | SPHINXOPTS = 6 | SPHINXBUILD = python -msphinx 7 | SPHINXPROJ = africanus 8 | SOURCEDIR = . 9 | BUILDDIR = _build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | .PHONY: help Makefile 16 | 17 | # Catch-all target: route all unknown targets to Sphinx using the new 18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 19 | %: Makefile 20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 21 | -------------------------------------------------------------------------------- /africanus/util/tests/test_type_inference.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | """Tests for `codex-africanus` package.""" 5 | 6 | 7 | import numpy as np 8 | 9 | from africanus.util.type_inference import infer_complex_dtype 10 | 11 | 12 | def test_type_inference(): 13 | i32 = np.empty(64, dtype=np.int32) 14 | i64 = np.empty(64, dtype=np.int64) 15 | f32 = np.empty(64, dtype=np.float32) 16 | f64 = np.empty(64, dtype=np.float64) 17 | 18 | assert infer_complex_dtype(f32, f64) == np.complex128 19 | assert infer_complex_dtype(f32, f32) == np.complex64 20 | assert infer_complex_dtype(i32, f32) == np.complex128 21 | assert infer_complex_dtype(i64, f32) == np.complex128 22 | -------------------------------------------------------------------------------- /africanus/conftest.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import numba 4 | from numba.core.runtime import rtsys 5 | import pytest 6 | 7 | from africanus.util.testing import mark_in_pytest 8 | 9 | 10 | @pytest.fixture(scope="function", autouse=bool(numba.config.NRT_STATS)) 11 | def check_allocations(): 12 | """Check allocations match frees""" 13 | try: 14 | yield 15 | start = rtsys.get_allocation_stats() 16 | finally: 17 | end = rtsys.get_allocation_stats() 18 | assert start.alloc - end.alloc == start.free - end.free 19 | 20 | 21 | # content of conftest.py 22 | def pytest_configure(config): 23 | mark_in_pytest(True) 24 | 25 | 26 | def pytest_unconfigure(config): 27 | mark_in_pytest(False) 28 | -------------------------------------------------------------------------------- /africanus/util/type_inference.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | try: 5 | from dask.array import Array as dask_array 6 | except ImportError: 7 | dask_array = object() 8 | 9 | import numba 10 | import numpy as np 11 | 12 | 13 | def _numpy_dtype(arg): 14 | if isinstance(arg, np.ndarray): 15 | return arg.dtype 16 | elif isinstance(arg, numba.types.npytypes.Array): 17 | return np.dtype(arg.dtype.name) 18 | elif isinstance(arg, dask_array): 19 | return arg.dtype 20 | else: 21 | raise ValueError("Unhandled type %s" % type(arg)) 22 | 23 | 24 | def infer_complex_dtype(*args): 25 | """Infer complex datatype from arg inputs""" 26 | return np.result_type(np.complex64, *(_numpy_dtype(a) for a in args)) 27 | -------------------------------------------------------------------------------- /docs/coordinates-api.rst: -------------------------------------------------------------------------------- 1 | --------------------- 2 | Coordinate Transforms 3 | --------------------- 4 | 5 | Numpy 6 | ~~~~~ 7 | 8 | .. currentmodule:: africanus.coordinates 9 | 10 | .. autosummary:: 11 | radec_to_lm 12 | radec_to_lmn 13 | lm_to_radec 14 | lmn_to_radec 15 | 16 | .. autofunction:: radec_to_lm 17 | .. autofunction:: radec_to_lmn 18 | .. autofunction:: lm_to_radec 19 | .. autofunction:: lmn_to_radec 20 | 21 | Dask 22 | ~~~~ 23 | 24 | .. currentmodule:: africanus.coordinates.dask 25 | 26 | .. autosummary:: 27 | radec_to_lm 28 | radec_to_lmn 29 | lm_to_radec 30 | lmn_to_radec 31 | 32 | .. autofunction:: radec_to_lm 33 | .. autofunction:: radec_to_lmn 34 | .. autofunction:: lm_to_radec 35 | .. autofunction:: lmn_to_radec 36 | -------------------------------------------------------------------------------- /africanus/rime/cuda/tests/test_cuda_feed_rotation.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | import pytest 6 | 7 | from africanus.rime import feed_rotation as np_feed_rotation 8 | from africanus.rime.cuda.feeds import feed_rotation as cp_feed_rotation 9 | 10 | 11 | @pytest.mark.parametrize("feed_type", ["linear", "circular"]) 12 | @pytest.mark.parametrize("shape", [(10, 7), (8,)]) 13 | @pytest.mark.parametrize("dtype", [np.float32, np.float64]) 14 | def test_cuda_feed_rotation(feed_type, shape, dtype): 15 | cp = pytest.importorskip("cupy") 16 | 17 | pa = np.random.random(shape).astype(dtype) 18 | 19 | cp_feed_rot = cp_feed_rotation(cp.asarray(pa), feed_type=feed_type) 20 | np_feed_rot = np_feed_rotation(pa, feed_type=feed_type) 21 | 22 | np.testing.assert_array_almost_equal(cp.asnumpy(cp_feed_rot), np_feed_rot) 23 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | =============== 2 | Codex Africanus 3 | =============== 4 | 5 | 6 | .. image:: https://img.shields.io/pypi/v/codex-africanus.svg 7 | :target: https://pypi.python.org/pypi/codex-africanus 8 | 9 | .. image:: https://img.shields.io/travis/ska-sa/codex-africanus.svg 10 | :target: https://travis-ci.org/ska-sa/codex-africanus 11 | 12 | .. image:: https://readthedocs.org/projects/codex-africanus/badge/?version=latest 13 | :target: https://codex-africanus.readthedocs.io/en/latest/?badge=latest 14 | :alt: Documentation Status 15 | 16 | 17 | .. image:: https://pyup.io/repos/github/ska-sa/codex-africanus/shield.svg 18 | :target: https://pyup.io/repos/github/ska-sa/codex-africanus/ 19 | :alt: Updates 20 | 21 | 22 | 23 | Radio Astronomy Building Blocks 24 | 25 | 26 | Documentation 27 | ------------- 28 | 29 | https://codex-africanus.readthedocs.io. 30 | -------------------------------------------------------------------------------- /africanus/rime/cuda/tests/test_cuda_phase_delay.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | import pytest 6 | 7 | from africanus.rime import phase_delay as np_phase_delay 8 | from africanus.rime.cuda.phase import phase_delay as cp_phase_delay 9 | 10 | 11 | @pytest.mark.parametrize("dtype, decimal", [(np.float32, 5), (np.float64, 6)]) 12 | def test_cuda_phase_delay(dtype, decimal): 13 | cp = pytest.importorskip("cupy") 14 | 15 | lm = 0.01 * np.random.random((10, 2)).astype(dtype) 16 | uvw = np.random.random((100, 3)).astype(dtype) 17 | freq = np.linspace(0.856e9, 2 * 0.856e9, 70, dtype=dtype) 18 | 19 | cp_cplx_phase = cp_phase_delay(cp.asarray(lm), cp.asarray(uvw), cp.asarray(freq)) 20 | np_cplx_phase = np_phase_delay(lm, uvw, freq) 21 | 22 | np.testing.assert_array_almost_equal( 23 | cp.asnumpy(cp_cplx_phase), np_cplx_phase, decimal=decimal 24 | ) 25 | -------------------------------------------------------------------------------- /africanus/rime/jax/phase.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | try: 5 | import jax.numpy as jnp 6 | except ImportError as e: 7 | opt_import_error = e 8 | else: 9 | opt_import_error = None 10 | 11 | from africanus.constants import minus_two_pi_over_c 12 | from africanus.util.requirements import requires_optional 13 | 14 | 15 | @requires_optional("jax", opt_import_error) 16 | def phase_delay(lm, uvw, frequency): 17 | one = lm.dtype.type(1.0) 18 | neg_two_pi_over_c = lm.dtype.type(minus_two_pi_over_c) 19 | 20 | l = lm[:, 0, None, None] # noqa 21 | m = lm[:, 1, None, None] 22 | 23 | u = uvw[None, :, 0, None] 24 | v = uvw[None, :, 1, None] 25 | w = uvw[None, :, 2, None] 26 | 27 | n = jnp.sqrt(one - l**2 - m**2) - one 28 | 29 | real_phase = neg_two_pi_over_c * (l * u + m * v + n * w) * frequency[None, None, :] 30 | 31 | return jnp.exp(jnp.complex64(1j) * real_phase) 32 | -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=python -msphinx 9 | ) 10 | set SOURCEDIR=. 11 | set BUILDDIR=_build 12 | set SPHINXPROJ=africanus 13 | 14 | if "%1" == "" goto help 15 | 16 | %SPHINXBUILD% >NUL 2>NUL 17 | if errorlevel 9009 ( 18 | echo. 19 | echo.The Sphinx module was not found. Make sure you have Sphinx installed, 20 | echo.then set the SPHINXBUILD environment variable to point to the full 21 | echo.path of the 'sphinx-build' executable. Alternatively you may add the 22 | echo.Sphinx directory to PATH. 23 | echo. 24 | echo.If you don't have Sphinx installed, grab it from 25 | echo.http://sphinx-doc.org/ 26 | exit /b 1 27 | ) 28 | 29 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% 30 | goto end 31 | 32 | :help 33 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% 34 | 35 | :end 36 | popd 37 | -------------------------------------------------------------------------------- /africanus/util/testing.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | try: 5 | from dask.utils import SerializableLock as Lock 6 | except ImportError: 7 | from threading import Lock 8 | 9 | 10 | __run_marker = {"in_pytest": False} 11 | __run_marker_lock = Lock() 12 | 13 | 14 | # Tag indicating that missing packages should generate an 15 | # exception, regardless of the 'in_pytest' marker 16 | # Used for testing exception raising behaviour 17 | force_missing_pkg_exception = object() 18 | 19 | 20 | def in_pytest(): 21 | """Return True if we're marked as executing inside pytest""" 22 | with __run_marker_lock: 23 | return __run_marker["in_pytest"] 24 | 25 | 26 | def mark_in_pytest(in_pytest=True): 27 | """Mark if we're in a pytest run""" 28 | if type(in_pytest) is not bool: 29 | raise TypeError("in_pytest %s is not a boolean" % in_pytest) 30 | 31 | with __run_marker_lock: 32 | __run_marker["in_pytest"] = in_pytest 33 | -------------------------------------------------------------------------------- /africanus/util/tests/test_progress_bar.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from africanus.util.dask_util import EstimatingProgressBar, format_time 4 | 5 | 6 | def test_progress_bar(): 7 | da = pytest.importorskip("dask.array") 8 | 9 | A = da.zeros((10, 10, 2)) 10 | B = A.map_blocks(lambda a: a + 1.0) 11 | 12 | with EstimatingProgressBar(out=None): 13 | B.compute() 14 | 15 | assert " 0s" == format_time(0) 16 | assert " 59s" == format_time(59) 17 | assert " 1m 0s" == format_time(60) 18 | assert " 1m 1s" == format_time(61) 19 | 20 | assert " 2h 6m" == format_time(2 * 60 * 60 + 6 * 60) 21 | assert " 2h 6m" == format_time(2 * 60 * 60 + 6 * 60 + 59) 22 | assert " 2h 7m" == format_time(2 * 60 * 60 + 7 * 60) 23 | assert " 2h 7m" == format_time(2 * 60 * 60 + 7 * 60 + 1) 24 | 25 | assert " 5d 2h" == format_time(5 * 60 * 60 * 24 + 2 * 60 * 60 + 500) 26 | 27 | assert " 5w 2d" == format_time(5 * 60 * 60 * 24 * 7 + 2 * 60 * 60 * 24 + 500) 28 | -------------------------------------------------------------------------------- /docs/linalg-api.rst: -------------------------------------------------------------------------------- 1 | -------------- 2 | Linear Algebra 3 | -------------- 4 | 5 | This module contains specialised linear algebra 6 | tools that are not currently available in the 7 | :code:`python` standard scientific libraries. 8 | 9 | Kronecker tools 10 | --------------- 11 | 12 | A kronecker matrix is matrix that can be written 13 | as a kronecker matrix of the individual matrices i.e. 14 | 15 | .. math:: 16 | K = K_0 \\otimes K_1 \\otimes K_2 \\otimes \\cdots 17 | 18 | Matrices which exhibit this structure can exploit 19 | properties of the kronecker product to avoid 20 | explicitly expanding the matrix :math:`K`. This 21 | module implements some common linear algebra 22 | operations which leverages this property for 23 | computational gains and a reduced memory footprint. 24 | 25 | Numpy 26 | ~~~~~ 27 | 28 | .. currentmodule:: africanus.linalg 29 | 30 | .. autosummary:: 31 | kron_matvec 32 | kron_cholesky 33 | 34 | .. autofunction:: kron_matvec 35 | .. autofunction:: kron_cholesky 36 | -------------------------------------------------------------------------------- /africanus/rime/cuda/tests/test_warp_transpose.cu.j2: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | {%- from "rime/cuda/macros.j2" import warp_transpose %} 4 | 5 | extern "C" __global__ void kernel( 6 | const CArray<{{type}}, 2> input, 7 | CArray<{{type}}, 2> output) 8 | { 9 | const ptrdiff_t & nvis = input.shape()[0]; 10 | int v = blockIdx.x*blockDim.x + threadIdx.x; 11 | 12 | 13 | // Array to hold our variables 14 | {{type}} values[{{corrs}}]; 15 | 16 | if(v < nvis) 17 | { 18 | {% for corr in range(corrs) %} 19 | values[{{corr}}] = input[v + {{corr}}*nvis]; 20 | {%- endfor %} 21 | } 22 | 23 | 24 | if(v < {{corrs}}*((nvis + {{corrs - 1}}) / {{corrs}})) 25 | { 26 | {{ warp_transpose("values", type, corrs) }} 27 | {{ warp_transpose("values", type, corrs) }} 28 | } 29 | 30 | if(v < nvis) 31 | { 32 | {% for corr in range(corrs) %} 33 | output[v + {{corr}}*nvis] = values[{{corr}}]; 34 | {%- endfor %} 35 | } 36 | } 37 | -------------------------------------------------------------------------------- /africanus/rime/jax/tests/test_jax_phase_delay.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | import pytest 6 | 7 | from africanus.rime.phase import phase_delay as np_phase_delay 8 | from africanus.rime.jax.phase import phase_delay 9 | 10 | 11 | @pytest.mark.parametrize("dtype", [np.float32, np.float64]) 12 | def test_jax_phase_delay(dtype): 13 | jax = pytest.importorskip("jax") 14 | 15 | np.random.seed(0) 16 | 17 | uvw = np.random.random(size=(100, 3)).astype(dtype) 18 | lm = np.random.random(size=(10, 2)).astype(dtype) * 0.001 19 | frequency = np.linspace(0.856e9, 0.856e9 * 2, 64).astype(dtype) 20 | 21 | # Compute complex phase 22 | np_complex_phase = np_phase_delay(lm, uvw, frequency) 23 | complex_phase = jax.jit(phase_delay)(lm, uvw, frequency) 24 | 25 | np.testing.assert_array_almost_equal(complex_phase, np_complex_phase) 26 | expected_ctype = np.result_type(dtype, np.complex64) 27 | assert np_complex_phase.dtype == expected_ctype 28 | assert complex_phase.dtype == expected_ctype 29 | -------------------------------------------------------------------------------- /africanus/util/tests/test_nvcc_compiler.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from africanus.util.cub import cub_dir 5 | 6 | import pytest 7 | 8 | 9 | def test_nvcc_compiler(tmpdir): 10 | cp = pytest.importorskip("cupy") 11 | from africanus.util.nvcc import compile_using_nvcc 12 | 13 | code = """ 14 | #include 15 | #include 16 | 17 | extern "C" __global__ void kernel(const CArray in, 18 | CArray out) 19 | { 20 | int x = blockDim.x*blockIdx.x + threadIdx.x; 21 | int n = int(in.shape()[0]); 22 | int value = in[x]; 23 | // printf("[%d, %d] value = %d\\n", x, n, value); 24 | out[x] = value; 25 | } 26 | 27 | """ 28 | mod = compile_using_nvcc(code, options=["-I " + cub_dir()]) 29 | kernel = mod.get_function("kernel") 30 | inputs = cp.arange(1024, dtype=cp.int32) 31 | outputs = cp.empty_like(inputs) 32 | 33 | kernel((1, 1, 1), (1024, 1, 1), (inputs, outputs)) 34 | cp.testing.assert_array_almost_equal(inputs, outputs) 35 | -------------------------------------------------------------------------------- /africanus/averaging/tests/test_splines.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import numpy as np 4 | from numpy.testing import assert_almost_equal 5 | import pytest 6 | 7 | from africanus.averaging.splines import fit_cubic_spline, evaluate_spline 8 | 9 | 10 | # Generate y,z coords from given x coords 11 | def generate_y_coords(x): 12 | y = -0.5 * x**2 - 0.3 * x + 5.0 13 | # z = 0.1 * x**3 + 5 14 | return y 15 | 16 | 17 | @pytest.mark.flaky(min_passes=1, max_runs=3) 18 | @pytest.mark.parametrize("order", [0]) 19 | def test_fit_cubic_spline(order): 20 | # Generate function y for x 21 | x = np.linspace(-2.0, 2.0, 16) 22 | y = generate_y_coords(x) 23 | 24 | spline = fit_cubic_spline(x, y) 25 | 26 | # Evaluation of spline at knot points is exact 27 | sy = evaluate_spline(spline, x, order=order) 28 | assert_almost_equal(sy, y) 29 | 30 | # Evaluate spline at points between knots is pretty inaccurate 31 | dx = np.diff(x) / 2 32 | dx += x[:-1] 33 | dy = generate_y_coords(dx) 34 | sdy = evaluate_spline(spline, dx, order=order) 35 | assert_almost_equal(sdy, dy, decimal=2) 36 | -------------------------------------------------------------------------------- /africanus/gps/utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | 6 | 7 | def abs_diff(x, xp): 8 | """ 9 | Gets matrix of differences between 10 | :math:`D`-dimensional vectors x and xp 11 | i.e. 12 | 13 | .. math:: 14 | X_{ij} = |x_i - x_j| 15 | 16 | Parameters 17 | ---------- 18 | x : :class:`numpy.ndarray` 19 | Array of inputs of shape :code:`(N, D)`. 20 | xp : :class:`numpy.ndarray` 21 | Array of inputs of shape :code:`(Np, D)`. 22 | 23 | Returns 24 | ------- 25 | XX : :class:`numpy.ndarray` 26 | Array of differences of shape :code:`(N, Np)`. 27 | 28 | """ 29 | try: 30 | N, D = x.shape 31 | Np, D = xp.shape 32 | except Exception: 33 | N = x.size 34 | D = 1 35 | Np = xp.size 36 | x = np.reshape(x, (N, D)) 37 | xp = np.reshape(xp, (Np, D)) 38 | xD = np.zeros([D, N, Np]) 39 | xpD = np.zeros([D, N, Np]) 40 | for i in range(D): 41 | xD[i] = np.tile(x[:, i], (Np, 1)).T 42 | xpD[i] = np.tile(xp[:, i], (N, 1)) 43 | return np.linalg.norm(xD - xpD, axis=0) 44 | -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- 1 | cff-version: 1.2.0 2 | message: If you use this software, please cite both the article from preferred-citation and the software itself. 3 | authors: 4 | - family-names: S.J. Perkins 5 | given-names: J.S. Kenyon 6 | - family-names: L.A.L. Andati 7 | given-names: H.L. Bester 8 | - family-names: O.M. Smirnov 9 | given-names: B.V. Hugo 10 | title: Africanus I. Scalable, distributed and efficient radio data processing with Dask-MS and Codex Africanus 11 | version: 1.0.0 12 | url: https://www.sciencedirect.com/science/article/pii/S2213133725000319 13 | doi: https://doi.org/10.1016/j.ascom.2025.100958 14 | date-released: '2025-06-05' 15 | preferred-citation: 16 | authors: 17 | - family-names: S.J. Perkins 18 | given-names: J.S. Kenyon 19 | - family-names: L.A.L. Andati 20 | given-names: H.L. Bester 21 | - family-names: O.M. Smirnov 22 | given-names: B.V. Hugo 23 | title: Africanus I. Scalable, distributed and efficient radio data processing with Dask-MS and Codex Africanus 24 | doi: https://doi.org/10.1016/j.ascom.2025.100958 25 | url: https://www.sciencedirect.com/science/article/pii/S2213133725000319 26 | type: article-journal 27 | pages: '100958' 28 | year: '2025' 29 | conference: {} 30 | publisher: {} 31 | -------------------------------------------------------------------------------- /africanus/util/tests/test_util_shape.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | """Tests for `codex-africanus` package.""" 5 | 6 | 7 | import pytest 8 | 9 | 10 | def test_corr_shape(): 11 | from africanus.util.shapes import corr_shape 12 | 13 | for i in range(10): 14 | assert corr_shape(i, "flat") == (i,) 15 | 16 | assert corr_shape(1, "matrix") == (1,) 17 | assert corr_shape(2, "matrix") == (2,) 18 | assert corr_shape(4, "matrix") == ( 19 | 2, 20 | 2, 21 | ) 22 | 23 | with pytest.raises(ValueError, match=r"ncorr not in \(1, 2, 4\)"): 24 | corr_shape(3, "matrix") 25 | 26 | 27 | def test_aggregate_chunks(): 28 | from africanus.util.shapes import aggregate_chunks 29 | 30 | chunks, max_c = (3, 4, 6, 3, 6, 7), 10 31 | expected = (7, 9, 6, 7) 32 | assert aggregate_chunks(chunks, max_c) == expected 33 | 34 | chunks, max_c = ((3, 4, 6, 3, 6, 7), (1, 1, 1, 1, 1, 1)), (10, 3) 35 | expected = ((7, 9, 6, 7), (2, 2, 1, 1)) 36 | assert aggregate_chunks(chunks, max_c) == expected 37 | 38 | chunks, max_c = ((3, 4, 6, 3, 6, 7), (1, 1, 1, 1, 1, 1)), (10, 1) 39 | assert aggregate_chunks(chunks, max_c) == chunks 40 | 41 | chunks, max_c = (5, 5, 5), 5 42 | assert aggregate_chunks(chunks, max_c) == chunks 43 | -------------------------------------------------------------------------------- /africanus/rime/tests/conftest.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | """Tests for `codex-africanus` package.""" 5 | 6 | 7 | import numpy as np 8 | import pytest 9 | 10 | 11 | @pytest.fixture 12 | def wsrt_ants(): 13 | """Westerbork antenna positions""" 14 | return np.array( 15 | [ 16 | [3828763.10544699, 442449.10566454, 5064923.00777], 17 | [3828746.54957258, 442592.13950824, 5064923.00792], 18 | [3828729.99081359, 442735.17696417, 5064923.00829], 19 | [3828713.43109885, 442878.2118934, 5064923.00436], 20 | [3828696.86994428, 443021.24917264, 5064923.00397], 21 | [3828680.31391933, 443164.28596862, 5064923.00035], 22 | [3828663.75159173, 443307.32138056, 5064923.00204], 23 | [3828647.19342757, 443450.35604638, 5064923.0023], 24 | [3828630.63486201, 443593.39226634, 5064922.99755], 25 | [3828614.07606798, 443736.42941621, 5064923.0], 26 | [3828609.94224429, 443772.19450029, 5064922.99868], 27 | [3828601.66208572, 443843.71178407, 5064922.99963], 28 | [3828460.92418735, 445059.52053929, 5064922.99071], 29 | [3828452.64716351, 445131.03744105, 5064922.98793], 30 | ], 31 | dtype=np.float64, 32 | ) 33 | -------------------------------------------------------------------------------- /africanus/util/casa_types.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | STOKES_TYPES = [ 5 | "Undefined", 6 | "I", 7 | "Q", 8 | "U", 9 | "V", 10 | "RR", 11 | "RL", 12 | "LR", 13 | "LL", 14 | "XX", 15 | "XY", 16 | "YX", 17 | "YY", 18 | "RX", 19 | "RY", 20 | "LX", 21 | "LY", 22 | "XR", 23 | "XL", 24 | "YR", 25 | "YL", 26 | "PP", 27 | "PQ", 28 | "QP", 29 | "QQ", 30 | "RCircular", 31 | "LCircular", 32 | "Linear", 33 | "Ptotal", 34 | "Plinear", 35 | "PFtotal", 36 | "PFlinear", 37 | "Pangle", 38 | ] 39 | """ 40 | List of stokes types as defined in 41 | Measurement Set 2.0 and Stokes.h in casacore: 42 | https://casacore.github.io/casacore/classcasacore_1_1Stokes.html 43 | """ 44 | 45 | 46 | STOKES_TYPE_MAP = {k: i for i, k in enumerate(STOKES_TYPES)} 47 | """ 48 | Map of stokes type to enumeration as defined in 49 | Measurement Set 2.0 and Stokes.h in casacore: 50 | https://casacore.github.io/casacore/classcasacore_1_1Stokes.html 51 | """ 52 | 53 | STOKES_ID_MAP = {v: k for k, v in STOKES_TYPE_MAP.items()} 54 | """ 55 | Map of stokes ID to stokes type string as defined in 56 | Measurement Set 2.0 and Stokes.h in casacore: 57 | https://casacore.github.io/casacore/classcasacore_1_1Stokes.html 58 | """ 59 | -------------------------------------------------------------------------------- /africanus/model/wsclean/dask.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from africanus.model.wsclean.spec_model import spectra as np_spectra, SPECTRA_DOCS 5 | from africanus.util.requirements import requires_optional 6 | 7 | try: 8 | import dask.array as da 9 | except ImportError as e: 10 | opt_import_error = e 11 | else: 12 | opt_import_error = None 13 | 14 | 15 | def spectra_wrapper(stokes, spi, log_si, ref_freq, frequency): 16 | return np_spectra(stokes, spi[0], log_si, ref_freq, frequency) 17 | 18 | 19 | @requires_optional("dask.array", opt_import_error) 20 | def spectra(stokes, spi, log_si, ref_freq, frequency): 21 | corrs = tuple("corr-%d" % i for i in range(len(stokes.shape[1:]))) 22 | log_si_schema = None if isinstance(log_si, bool) else ("source",) 23 | 24 | return da.blockwise( 25 | spectra_wrapper, 26 | ("source", "chan") + corrs, 27 | stokes, 28 | ("source",) + corrs, 29 | spi, 30 | ("source", "spi") + corrs, 31 | log_si, 32 | log_si_schema, 33 | ref_freq, 34 | ("source",), 35 | frequency, 36 | ("chan",), 37 | dtype=stokes.dtype, 38 | ) 39 | 40 | 41 | try: 42 | spectra.__doc__ = SPECTRA_DOCS.substitute(array_type=":class:`dask.array.Array`") 43 | except AttributeError: 44 | pass 45 | -------------------------------------------------------------------------------- /africanus/util/tests/test_requirements.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | 5 | import sys 6 | 7 | import pytest 8 | 9 | from africanus.util.requirements import requires_optional, MissingPackageException 10 | from africanus.util.testing import force_missing_pkg_exception as force_tag 11 | 12 | 13 | def test_requires_optional_missing_import(): 14 | @requires_optional("sys", "bob", force_tag) 15 | def f(*args, **kwargs): 16 | pass 17 | 18 | with pytest.raises(MissingPackageException) as e: 19 | f(1, a=2) 20 | 21 | assert "f requires installation of the following packages: ('bob',)." in str( 22 | e.value 23 | ) 24 | 25 | 26 | def test_requires_optional_pass_import_error(): 27 | assert "clearly_missing_and_nonexistent_package" not in sys.modules 28 | 29 | try: 30 | import clearly_missing_and_nonexistent_package # noqa 31 | except ImportError as e: 32 | me = e 33 | else: 34 | me = None 35 | 36 | with pytest.raises(ImportError) as e: 37 | 38 | @requires_optional("sys", "os", me, force_tag) 39 | def f(*args, **kwargs): 40 | pass 41 | 42 | msg = str(e.value) 43 | assert "Successfully imported ['sys', 'os']" in msg 44 | assert "No module named" in msg 45 | assert "clearly_missing_and_nonexistent_package" in msg 46 | -------------------------------------------------------------------------------- /africanus/model/shape/tests/test_gaussian_shape.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | from numpy.testing import assert_array_almost_equal 6 | import pytest 7 | 8 | from africanus.model.shape import gaussian as np_gaussian 9 | 10 | 11 | def test_gauss_shape(): 12 | row = 10 13 | chan = 16 14 | 15 | shape_params = np.array([[0.4, 0.3, 0.2], [0.4, 0.3, 0.2]]) 16 | uvw = np.random.random((row, 3)) 17 | freq = np.linspace(0.856e9, 2 * 0.856e9, chan) 18 | 19 | gauss_shape = np_gaussian(uvw, freq, shape_params) 20 | 21 | assert gauss_shape.shape == (shape_params.shape[0], row, chan) 22 | 23 | 24 | def test_dask_gauss_shape(): 25 | da = pytest.importorskip("dask.array") 26 | from africanus.model.shape.dask import gaussian as da_gaussian 27 | 28 | row_chunks = (5, 5) 29 | chan_chunks = (4, 4) 30 | 31 | row = sum(row_chunks) 32 | chan = sum(chan_chunks) 33 | 34 | shape_params = da.asarray([[0.4, 0.3, 0.2], [0.4, 0.3, 0.2]]) 35 | uvw = da.random.random((row, 3), chunks=(row_chunks, 3)) 36 | freq = da.linspace(0.856e9, 2 * 0.856e9, chan, chunks=chan_chunks) 37 | da_gauss_shape = da_gaussian(uvw, freq, shape_params).compute() 38 | np_gauss_shape = np_gaussian(uvw.compute(), freq.compute(), shape_params.compute()) 39 | 40 | assert_array_almost_equal(da_gauss_shape, np_gauss_shape) 41 | -------------------------------------------------------------------------------- /docs/installation.rst: -------------------------------------------------------------------------------- 1 | .. highlight:: shell 2 | 3 | ============ 4 | Installation 5 | ============ 6 | 7 | 8 | Stable release 9 | -------------- 10 | 11 | To install Codex Africanus, run this command in your terminal: 12 | 13 | .. code-block:: console 14 | 15 | $ pip install codex-africanus 16 | 17 | This is the preferred method to install Codex Africanus, 18 | as it will always install the most recent stable release. 19 | 20 | If you don't have `pip`_ installed, this `Python installation guide`_ 21 | can guide you through the process. 22 | 23 | By default, Codex Africanus will install with a minimal set of 24 | dependencies, numpy and numba. 25 | 26 | Further functionality can be enabled by installing extra requirements 27 | as follows: 28 | 29 | .. code-block:: console 30 | 31 | $ pip install codex-africanus[dask] 32 | $ pip install codex-africanus[scipy] 33 | $ pip install codex-africanus[astropy] 34 | $ pip install codex-africanus[python-casacore] 35 | 36 | 37 | To install the complete set of dependencies for the CPU: 38 | 39 | .. code-block:: console 40 | 41 | $ pip install codex-africanus[complete] 42 | 43 | To install the complete set of dependencies including CUDA: 44 | 45 | .. code-block:: console 46 | 47 | $ pip install codex-africanus[complete-cuda] 48 | 49 | 50 | .. _pip: https://pip.pypa.io 51 | .. _Python installation guide: http://docs.python-guide.org/en/latest/starting/installation/ 52 | -------------------------------------------------------------------------------- /africanus/model/coherency/cuda/tests/test_convert.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | import pytest 6 | 7 | from africanus.model.coherency import convert as np_convert 8 | from africanus.model.coherency.cuda import convert 9 | from africanus.model.coherency.tests.test_convert import ( 10 | stokes_corr_cases, 11 | stokes_corr_int_cases, 12 | visibility_factory, 13 | vis_shape, 14 | ) 15 | 16 | 17 | @pytest.mark.skip 18 | def test_stokes_schemas(input_schema, output_schema, vis_shape): 19 | input_shape = np.asarray(input_schema).shape 20 | output_shape = np.asarray(output_schema).shape 21 | 22 | vis = visibility_factory(vis_shape, input_shape) 23 | xformed_vis = np_convert(vis, input_schema, output_schema) 24 | assert xformed_vis.shape == vis_shape + output_shape 25 | 26 | 27 | @pytest.mark.parametrize( 28 | "input_schema, output_schema", 29 | stokes_corr_cases + stokes_corr_int_cases, 30 | ) 31 | @pytest.mark.parametrize("vis_shape", vis_shape) 32 | def test_cuda_convert(input_schema, output_schema, vis_shape): 33 | cp = pytest.importorskip("cupy") 34 | 35 | input_shape = np.asarray(input_schema).shape 36 | vis = visibility_factory(vis_shape, input_shape) 37 | 38 | cp_out = convert(cp.asarray(vis), input_schema, output_schema) 39 | np_out = np_convert(vis, input_schema, output_schema) 40 | 41 | np.testing.assert_array_almost_equal(cp.asnumpy(cp_out), np_out) 42 | -------------------------------------------------------------------------------- /africanus/util/numba.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from decorator import decorate 5 | from numba import types 6 | 7 | from africanus.util.docs import on_rtd 8 | 9 | JIT_OPTIONS = { 10 | "cache": True, 11 | "nogil": True, 12 | } 13 | 14 | if on_rtd(): 15 | # Fake decorators when on readthedocs 16 | def _fake_decorator(*args, **kwargs): 17 | def decorator(fn): 18 | def wrapper(*args, **kwargs): 19 | return fn(*args, **kwargs) 20 | 21 | return decorate(fn, wrapper) 22 | 23 | return decorator 24 | 25 | cfunc = _fake_decorator 26 | jit = _fake_decorator 27 | njit = _fake_decorator 28 | stencil = _fake_decorator 29 | 30 | overload = _fake_decorator 31 | register_jitable = _fake_decorator 32 | intrinsic = _fake_decorator 33 | else: 34 | from numba import cfunc, jit, njit, stencil # noqa 35 | from numba.extending import overload, register_jitable, intrinsic # noqa 36 | 37 | 38 | def is_numba_type_none(arg): 39 | """ 40 | Returns True if the numba type represents None 41 | 42 | 43 | Parameters 44 | ---------- 45 | arg : :class:`numba.Type` 46 | The numba type 47 | 48 | Returns 49 | ------- 50 | boolean 51 | True if the type represents None 52 | """ 53 | return isinstance(arg, types.misc.NoneType) or ( 54 | isinstance(arg, types.misc.Omitted) and arg.value is None 55 | ) 56 | -------------------------------------------------------------------------------- /africanus/model/spectral/dask.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from africanus.model.spectral.spec_model import ( 5 | spectral_model as np_spectral_model, 6 | SPECTRAL_MODEL_DOC, 7 | ) 8 | from africanus.util.requirements import requires_optional 9 | 10 | try: 11 | import dask.array as da 12 | except ImportError as e: 13 | opt_import_error = e 14 | else: 15 | opt_import_error = None 16 | 17 | 18 | def spectral_model_wrapper(stokes, spi, ref_freq, frequencies, base=None): 19 | return np_spectral_model(stokes, spi[0], ref_freq, frequencies, base=base) 20 | 21 | 22 | @requires_optional("dask.array", opt_import_error) 23 | def spectral_model(stokes, spi, ref_freq, frequencies, base=0): 24 | if len(spi.chunks[1]) != 1: 25 | raise ValueError("Chunking along the spi dimension unsupported") 26 | 27 | pol_dim = () if stokes.ndim == 1 else ("pol",) 28 | 29 | return da.blockwise( 30 | spectral_model_wrapper, 31 | ( 32 | "source", 33 | "chan", 34 | ) 35 | + pol_dim, 36 | stokes, 37 | ("source",) + pol_dim, 38 | spi, 39 | ("source", "spi") + pol_dim, 40 | ref_freq, 41 | ("source",), 42 | frequencies, 43 | ("chan",), 44 | base=base, 45 | dtype=stokes.dtype, 46 | ) 47 | 48 | 49 | try: 50 | spectral_model.__doc__ = SPECTRAL_MODEL_DOC.substitute( 51 | array_type=":class:`dask.array.Array`" 52 | ) 53 | except AttributeError: 54 | pass 55 | -------------------------------------------------------------------------------- /africanus/rime/parangles_astropy.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from africanus.util.requirements import requires_optional 5 | 6 | try: 7 | from astropy.coordinates import EarthLocation, SkyCoord, AltAz, CIRS 8 | from astropy.time import Time 9 | from astropy import units 10 | except ImportError as e: 11 | astropy_import_error = e 12 | have_astropy_parangles = False 13 | else: 14 | astropy_import_error = None 15 | have_astropy_parangles = True 16 | 17 | 18 | @requires_optional("astropy", astropy_import_error) 19 | def astropy_parallactic_angles(times, antenna_positions, field_centre): 20 | """ 21 | Computes parallactic angles per timestep for the given 22 | reference antenna position and field centre. 23 | """ 24 | ap = antenna_positions 25 | fc = field_centre 26 | 27 | # Convert from MJD second to MJD 28 | times = Time(times / 86400.00, format="mjd", scale="utc") 29 | 30 | ap = EarthLocation.from_geocentric(ap[:, 0], ap[:, 1], ap[:, 2], unit="m") 31 | fc = SkyCoord(ra=fc[0], dec=fc[1], unit=units.rad, frame="fk5") 32 | pole = SkyCoord(ra=0, dec=90, unit=units.deg, frame="fk5") 33 | 34 | cirs_frame = CIRS(obstime=times) 35 | pole_cirs = pole.transform_to(cirs_frame) 36 | fc_cirs = fc.transform_to(cirs_frame) 37 | 38 | altaz_frame = AltAz(location=ap[None, :], obstime=times[:, None]) 39 | pole_altaz = pole_cirs[:, None].transform_to(altaz_frame) 40 | fc_altaz = fc_cirs[:, None].transform_to(altaz_frame) 41 | return fc_altaz.position_angle(pole_altaz) 42 | -------------------------------------------------------------------------------- /africanus/gps/kernels.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | from africanus.gps.utils import abs_diff 6 | 7 | 8 | def exponential_squared(x, xp, sigmaf, l, pspec=False): # noqa: E741 9 | """ 10 | Create exponential squared covariance 11 | function between :math:`D` dimensional 12 | vectors :math:`x` and :math:`x_p` i.e. 13 | 14 | .. math:: 15 | k(x, x_p) = \\sigma_f^2 \\exp\\left(-\\frac{(x-x_p)^2}{2l^2}\\right) 16 | 17 | Parameters 18 | ---------- 19 | x : :class:`numpy.ndarray` 20 | Array of shape :code:`(N, D)`. 21 | xp : :class:`numpy.ndarray` 22 | Array of shape :code:`(Np, D)`. 23 | sigmaf : float 24 | The signal variance hyper-parameter 25 | l : float 26 | The length scale hyper-parameter 27 | 28 | Returns 29 | ------- 30 | K : :class:`numpy.ndarray` 31 | Array of shape :code:`(N, Np)` 32 | """ 33 | if pspec: 34 | N, D = x.shape 35 | if D != 1: 36 | raise NotImplementedError("Only 1D pspecs supported") 37 | if (x != xp).any(): 38 | raise ValueError("pspec only defined if x = xp") 39 | x = x.squeeze() 40 | delx = x[1] - x[0] 41 | if (x[1::] - x[0:-1] != delx).any(): 42 | raise ValueError("pspec only defined on regular grid") 43 | s = np.fft.fftshift(np.fft.fftfreq(N, d=delx)) 44 | return np.sqrt(2 * np.pi * l) * sigmaf**2.0 * np.exp(-(l**2) * s**2 / 2.0) 45 | else: 46 | xxp = abs_diff(x, xp) 47 | return sigmaf**2 * np.exp(-(xxp**2) / (2 * l**2)) 48 | -------------------------------------------------------------------------------- /docs/gridding-api.rst: -------------------------------------------------------------------------------- 1 | ----------------------- 2 | Gridding and Degridding 3 | ----------------------- 4 | 5 | This section contains routines for 6 | 7 | 1. Gridding complex visibilities onto an image. 8 | 2. Degridding complex visibilities from an image. 9 | 10 | Nifty 11 | ~~~~~ 12 | 13 | Dask wrappers around 14 | `Nifty's Gridder `_. 15 | 16 | Dask 17 | ++++ 18 | 19 | .. currentmodule:: africanus.gridding.nifty.dask 20 | 21 | .. autosummary:: 22 | grid_config 23 | grid 24 | dirty 25 | degrid 26 | model 27 | 28 | .. autofunction:: grid_config 29 | .. autofunction:: grid 30 | .. autofunction:: dirty 31 | .. autofunction:: degrid 32 | .. autofunction:: model 33 | 34 | wgridder 35 | ~~~~~~~~ 36 | 37 | Wrappers around 'ducc.wgridder `_. 38 | 39 | 40 | Numpy 41 | +++++ 42 | 43 | .. currentmodule:: africanus.gridding.wgridder 44 | 45 | .. autosummary:: 46 | dirty 47 | model 48 | residual 49 | hessian 50 | 51 | .. autofunction:: dirty 52 | .. autofunction:: model 53 | .. autofunction:: residual 54 | .. autofunction:: hessian 55 | 56 | Dask 57 | ++++ 58 | 59 | .. currentmodule:: africanus.gridding.wgridder.dask 60 | 61 | .. autosummary:: 62 | dirty 63 | model 64 | residual 65 | hessian 66 | 67 | .. autofunction:: dirty 68 | .. autofunction:: model 69 | .. autofunction:: residual 70 | .. autofunction:: hessian 71 | 72 | Utilities 73 | ~~~~~~~~~ 74 | 75 | .. currentmodule:: africanus.gridding.util 76 | 77 | .. autosummary:: 78 | estimate_cell_size 79 | 80 | .. autofunction:: estimate_cell_size 81 | -------------------------------------------------------------------------------- /africanus/model/wsclean/tests/conftest.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from os.path import join as pjoin 5 | 6 | import pytest 7 | 8 | 9 | _WSCLEAN_MODEL_FILE = """ 10 | Format = Name, Type, Ra, Dec, I, SpectralIndex, LogarithmicSI, ReferenceFrequency='125584411.621094', MajorAxis, MinorAxis, Orientation 11 | s0c0,POINT,-08:28:05.152,39.35.08.511,0.000748810650400475,[-0.00695379313004673,-0.0849693907803257],true,125584411.621094,,, 12 | s0c1,POINT,08:22:27.658,39.37.38.353,-0.000154968071120503,[-0.000898135869319762,0.0183710297781511],false,125584411.621094,,, 13 | s0c2,POINT,08:18:44,39.38.37,0.000233552686127518,[-0.000869089801859608,0.0828587947079702],false,125584411.621094,,, 14 | s0c3,POINT,08:03:07.538,39.37.02.717,0.000919058240247659,[0.001264109956439,0.0201438425344451],false,125584411.621094,,, 15 | s1c0,GAUSSIAN,08:31:10.37,+41.47.17.131,0.000723326710524984,[0.00344317919656096,-0.115990377833407],true,125584411.621094,83.6144111272856,83.6144111272856,0 16 | s1c1,GAUSSIAN,07:51:09.24,42.32.46.177,0.000660490865128381,[0.00404869217508666,-0.011844732049232],false,125584411.621094,83.6144111272856,83.6144111272856,0 17 | s1c2,GAUSSIAN,07:51:09.24,42.32.46.177,0.000660490865128381,[0.00404869217508666,-0.011844732049232],false,,83.6144111272856,83.6144111272856,45.0 18 | s1c3,GAUSSIAN,07:51:09.24,42.32.46.177,nan,[nan,inf],false,,83.6144111272856,83.6144111272856,45.0 19 | """ # noqa 20 | 21 | 22 | @pytest.fixture 23 | def wsclean_model_file(tmpdir): 24 | filename = pjoin(str(tmpdir), "model.txt") 25 | 26 | with open(filename, "w") as f: 27 | f.write(_WSCLEAN_MODEL_FILE) 28 | 29 | return filename 30 | -------------------------------------------------------------------------------- /africanus/rime/cuda/tests/test_macros.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from os.path import join as pjoin 5 | 6 | import numpy as np 7 | import pytest 8 | 9 | from africanus.util.code import format_code 10 | from africanus.util.jinja2 import jinja_env 11 | 12 | 13 | @pytest.mark.parametrize("ncorrs", [1, 2, 4, 8]) 14 | @pytest.mark.parametrize( 15 | "dtype", [np.int32, np.float64, np.float32, np.complex64, np.complex128] 16 | ) 17 | @pytest.mark.parametrize("nvis", [9, 10, 11, 32, 1025]) 18 | @pytest.mark.parametrize("debug", ["false"]) 19 | def test_cuda_inplace_warp_transpose(ncorrs, dtype, nvis, debug): 20 | cp = pytest.importorskip("cupy") 21 | 22 | path = pjoin("rime", "cuda", "tests", "test_warp_transpose.cu.j2") 23 | render = jinja_env.get_template(path).render 24 | 25 | dtypes = { 26 | np.float32: "float", 27 | np.float64: "double", 28 | np.int32: "int", 29 | np.complex64: "float2", 30 | np.complex128: "double2", 31 | } 32 | 33 | code = render(type=dtypes[dtype], warp_size=32, corrs=ncorrs, debug=debug) 34 | kernel = cp.RawKernel(code, "kernel") 35 | 36 | inputs = cp.arange(nvis * ncorrs, dtype=dtype).reshape(nvis, ncorrs) 37 | outputs = cp.empty_like(inputs) 38 | args = (inputs, outputs) 39 | block = (256, 1, 1) 40 | grid = tuple((d + b - 1) // b for d, b in zip((nvis, 1, 1), block)) 41 | 42 | try: 43 | kernel(grid, block, args) 44 | except cp.cuda.compiler.CompileException: 45 | print(format_code(kernel.code)) 46 | raise 47 | 48 | np.testing.assert_array_almost_equal(cp.asnumpy(inputs), cp.asnumpy(outputs)) 49 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2019, National Research Foundation 2 | (South African Radio Astronomy Observatory) 3 | All rights reserved. 4 | 5 | Redistribution and use in source and binary forms, with or without modification, 6 | are permitted provided that the following conditions are met: 7 | 8 | 1. Redistributions of source code must retain the above copyright notice, 9 | this list of conditions and the following disclaimer. 10 | 11 | 2. Redistributions in binary form must reproduce the above copyright notice, 12 | this list of conditions and the following disclaimer in the documentation 13 | and/or other materials provided with the distribution. 14 | 15 | 3. Neither the name of the copyright holder nor the names of its contributors 16 | may be used to endorse or promote products derived from this software 17 | without specific prior written permission. 18 | 19 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 20 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, 21 | THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 22 | ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 23 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 24 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS 25 | OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 26 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 27 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 28 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 29 | -------------------------------------------------------------------------------- /africanus/util/docs.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import os 5 | import re 6 | from string import Template 7 | 8 | _on_rtd = bool(os.environ.get("READTHEDOCS")) 9 | 10 | 11 | def on_rtd(): 12 | return _on_rtd 13 | 14 | 15 | def mod_docs(docstring, replacements): 16 | for search, replace in replacements: 17 | docstring = docstring.replace(search, replace) 18 | 19 | return docstring 20 | 21 | 22 | def doc_tuple_to_str(doc_tuple, replacements=None): 23 | fields = getattr(doc_tuple, "_fields", None) 24 | 25 | if fields is not None: 26 | fields = (getattr(doc_tuple, f) for f in doc_tuple._fields) 27 | elif isinstance(doc_tuple, dict): 28 | fields = doc_tuple.values() 29 | 30 | if replacements is not None: 31 | fields = (mod_docs(f, replacements) for f in fields) 32 | 33 | return "".join(fields) 34 | 35 | 36 | class DefaultOut(object): 37 | def __init__(self, arg): 38 | self.arg = arg 39 | 40 | def __repr__(self): 41 | return self.arg 42 | 43 | __str__ = __repr__ 44 | 45 | 46 | class DocstringTemplate(Template): 47 | """ 48 | Overrides the ${identifer} braced pattern in the string Template 49 | with a $(identifier) braced pattern 50 | """ 51 | 52 | pattern = r""" 53 | %(delim)s(?: 54 | (?P%(delim)s) | # Escape sequence of two delimiters 55 | (?P%(id)s) | # delimiter and a Python identifier 56 | \((?P%(id)s)\) | # delimiter and a braced identifier 57 | (?P) # Other ill-formed delimiter exprs 58 | ) 59 | """ % {"delim": re.escape(Template.delimiter), "id": Template.idpattern} 60 | -------------------------------------------------------------------------------- /africanus/model/coherency/cuda/conversion.cu.j2: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | {%- from "rime/cuda/macros.j2" import warp_transpose %} 4 | 5 | extern "C" __global__ void {{kernel_name}}( 6 | const CArray<{{input_type}}, 2> inputs, 7 | CArray<{{output_type}}, 2> outputs) 8 | { 9 | // Number of sources 10 | const ptrdiff_t & nsrc = inputs.shape()[0]; 11 | // Number of sources rounded up to a multiple of the number of elements 12 | int nsrcmulup = {{elements}}*((nsrc + {{elements - 1}}) / {{elements}}); 13 | // Current source 14 | int source = blockIdx.x*blockDim.x + threadIdx.x; 15 | 16 | // Guard 17 | if(source >= nsrcmulup) 18 | { return; } 19 | 20 | // Array to hold our variables 21 | {{input_type}} in[{{elements}}]; 22 | {{output_type}} out[{{elements}}]; 23 | 24 | // Read in data 25 | { 26 | int idx = source; 27 | int upper = {{elements}}*nsrc; 28 | {%- for elem in range(elements) %} 29 | if(idx < upper) 30 | { 31 | in[{{elem}}] = inputs[idx]; 32 | idx += nsrcmulup; 33 | } 34 | {%- endfor %} 35 | } 36 | 37 | // Transpose forward 38 | {{ warp_transpose("in", input_type, elements) }} 39 | 40 | {%- for expr in assign_exprs %} 41 | {{expr}} 42 | {%- endfor %} 43 | 44 | // Transpose backward 45 | {{ warp_transpose("out", output_type, elements) }} 46 | 47 | // Write out data 48 | { 49 | int idx = source; 50 | int upper = {{elements}}*nsrc; 51 | {%- for elem in range(elements) %} 52 | if(idx < upper) 53 | { 54 | outputs[idx] = out[{{elem}}]; 55 | idx += nsrcmulup; 56 | } 57 | 58 | {%- endfor %} 59 | } 60 | } 61 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/transformers/lm.py: -------------------------------------------------------------------------------- 1 | from numba.core import types 2 | from numba.core.errors import TypingError 3 | import numpy as np 4 | 5 | from africanus.experimental.rime.fused.transformers.core import Transformer 6 | 7 | 8 | class LMTransformer(Transformer): 9 | OUTPUTS = ["lm"] 10 | 11 | def init_fields(self, typingctx, init_state, radec, phase_dir): 12 | dt = typingctx.unify_types(radec.dtype, phase_dir.dtype) 13 | fields = [("lm", dt[:, :])] 14 | 15 | if not isinstance(phase_dir, types.Array) or phase_dir.ndim != 1: 16 | raise TypingError(f"phase_dir {phase_dir} is not a (2,) shaped array") 17 | 18 | assert radec.ndim == 2 19 | 20 | def lm(init_state, radec, phase_dir): 21 | lm = np.empty_like(radec) 22 | pc_ra = phase_dir[0] 23 | pc_dec = phase_dir[1] 24 | 25 | sin_pc_dec = np.sin(pc_dec) 26 | cos_pc_dec = np.cos(pc_dec) 27 | 28 | for s in range(radec.shape[0]): 29 | da = radec[s, 0] - pc_ra 30 | sin_ra_delta = np.sin(da) 31 | cos_ra_delta = np.cos(da) 32 | 33 | sin_dec = np.sin(radec[s, 1]) 34 | cos_dec = np.cos(radec[s, 1]) 35 | 36 | lm[s, 0] = cos_dec * sin_ra_delta 37 | lm[s, 1] = sin_dec * cos_pc_dec - cos_dec * sin_pc_dec * cos_ra_delta 38 | 39 | return lm 40 | 41 | return fields, lm 42 | 43 | def dask_schema(self, radec, phase_dir): 44 | assert radec.ndim == 2 45 | assert phase_dir.ndim == 1 46 | 47 | inputs = {"radec": ("source", "radec"), "phase_dir": ("radec",)} 48 | 49 | outputs = {"lm": np.empty((0, 0), dtype=radec.dtype)} 50 | 51 | return inputs, outputs 52 | -------------------------------------------------------------------------------- /africanus/rime/cuda/tests/test_cuda_beam.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | from numpy.testing import assert_array_almost_equal 6 | import pytest 7 | 8 | from africanus.rime import beam_cube_dde as np_beam_cube_dde 9 | from africanus.rime.cuda.beam import beam_cube_dde as cp_beam_cude_dde 10 | 11 | cp = pytest.importorskip("cupy") 12 | 13 | 14 | @pytest.mark.parametrize("corrs", [(2, 2), (4,), (2,), (1,)]) 15 | def test_cuda_beam(corrs): 16 | rs = np.random.RandomState(42) 17 | 18 | src, time, ant, chan = 20, 29, 14, 64 19 | beam_lw = beam_mh = beam_nud = 50 20 | 21 | beam = ( 22 | rs.normal(size=(beam_lw, beam_mh, beam_nud) + corrs) 23 | + rs.normal(size=(beam_lw, beam_mh, beam_nud) + corrs) * 1j 24 | ) 25 | 26 | beam_lm_ext = np.array(([[-0.5, 0.5], [-0.5, 0.5]])) 27 | 28 | lm = rs.normal(size=(src, 2)) - 0.5 29 | 30 | if chan == 1: 31 | freqs = np.array([0.856e9 * 3 / 2]) 32 | else: 33 | freqs = np.linspace(0.856e9, 2 * 0.856e9, chan) 34 | 35 | beam_freq_map = np.linspace(0.856e9, 2 * 0.856e9, beam_nud) 36 | 37 | parangles = rs.normal(size=(time, ant)) 38 | point_errors = rs.normal(size=(time, ant, chan, 2)) 39 | ant_scales = rs.normal(size=(ant, chan, 2)) 40 | 41 | np_ddes = np_beam_cube_dde( 42 | beam, beam_lm_ext, beam_freq_map, lm, parangles, point_errors, ant_scales, freqs 43 | ) 44 | 45 | cp_ddes = cp_beam_cude_dde( 46 | cp.asarray(beam), 47 | cp.asarray(beam_lm_ext), 48 | cp.asarray(beam_freq_map), 49 | cp.asarray(lm), 50 | cp.asarray(parangles), 51 | cp.asarray(point_errors), 52 | cp.asarray(ant_scales), 53 | cp.asarray(freqs), 54 | ) 55 | 56 | assert_array_almost_equal(np_ddes, cp.asnumpy(cp_ddes)) 57 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Generated by setup.py 2 | africanus/install/extras_require.py 3 | 4 | # Byte-compiled / optimized / DLL files 5 | __pycache__/ 6 | *.py[cod] 7 | *$py.class 8 | 9 | # C extensions 10 | *.so 11 | 12 | # Distribution / packaging 13 | .Python 14 | env/ 15 | build/ 16 | develop-eggs/ 17 | dist/ 18 | downloads/ 19 | eggs/ 20 | .eggs/ 21 | lib/ 22 | lib64/ 23 | parts/ 24 | sdist/ 25 | var/ 26 | wheels/ 27 | *.egg-info/ 28 | .installed.cfg 29 | *.egg 30 | 31 | # PyInstaller 32 | # Usually these files are written by a python script from a template 33 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 34 | *.manifest 35 | *.spec 36 | 37 | # Installer logs 38 | pip-log.txt 39 | pip-delete-this-directory.txt 40 | 41 | # Unit test / coverage reports 42 | htmlcov/ 43 | .tox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | 53 | # Translations 54 | *.mo 55 | *.pot 56 | 57 | # Django stuff: 58 | *.log 59 | local_settings.py 60 | 61 | # Flask stuff: 62 | instance/ 63 | .webassets-cache 64 | 65 | # Scrapy stuff: 66 | .scrapy 67 | 68 | # Sphinx documentation 69 | docs/_build/ 70 | 71 | # PyBuilder 72 | target/ 73 | 74 | # Jupyter Notebook 75 | .ipynb_checkpoints 76 | 77 | # pyenv 78 | .python-version 79 | 80 | # celery beat schedule file 81 | celerybeat-schedule 82 | 83 | # SageMath parsed files 84 | *.sage.py 85 | 86 | # dotenv 87 | .env 88 | 89 | # virtualenv 90 | .venv 91 | venv/ 92 | ENV/ 93 | 94 | # Spyder project settings 95 | .spyderproject 96 | .spyproject 97 | 98 | # Rope project settings 99 | .ropeproject 100 | 101 | # mkdocs documentation 102 | /site 103 | 104 | # mypy 105 | .mypy_cache/ 106 | -------------------------------------------------------------------------------- /africanus/averaging/tests/test_support.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | from numpy.testing import assert_array_equal 6 | import pytest 7 | 8 | from africanus.averaging.support import unique_baselines, unique_time 9 | 10 | 11 | @pytest.fixture 12 | def time(): 13 | return np.asarray([1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0]) # noqa 14 | 15 | 16 | @pytest.fixture 17 | def ant1(): 18 | return np.asarray( 19 | [0, 0, 1, 0, 0, 1, 2, 0, 0, 1], # noqa 20 | dtype=np.int32, 21 | ) 22 | 23 | 24 | @pytest.fixture 25 | def ant2(): 26 | return np.asarray( 27 | [1, 2, 2, 0, 1, 2, 3, 0, 1, 2], # noqa 28 | dtype=np.int32, 29 | ) 30 | 31 | 32 | @pytest.fixture 33 | def vis(): 34 | def _vis(row, chan, fcorrs): 35 | return ( 36 | np.arange(row * chan * fcorrs, dtype=np.float32) 37 | + np.arange(1, row * chan * fcorrs + 1, dtype=np.float32) * 1j 38 | ) 39 | 40 | return _vis 41 | 42 | 43 | def test_unique_time(time): 44 | # Reverse time to test that sort works 45 | time = np.flipud(time) 46 | 47 | utime, idx, inv, counts = unique_time(time) 48 | assert_array_equal(utime, [1.0, 2.0, 3.0]) 49 | assert_array_equal(utime[inv], time) 50 | assert_array_equal(time[idx], utime) 51 | assert_array_equal(counts, [3, 4, 3]) 52 | 53 | 54 | def test_unique_baselines(ant1, ant2): 55 | # Reverse ant1, ant2 to test that sort works 56 | ant1 = np.flipud(ant1) 57 | ant2 = np.flipud(ant2) 58 | 59 | test_bl = np.stack([ant1, ant2], axis=1) 60 | 61 | bl, idx, inv, counts = unique_baselines(ant1, ant2) 62 | assert_array_equal(bl, [[0, 0], [0, 1], [0, 2], [1, 2], [2, 3]]) 63 | assert_array_equal(bl[inv], test_bl) 64 | assert_array_equal(test_bl[idx], bl) 65 | assert_array_equal(counts, [2, 3, 1, 3, 1]) 66 | -------------------------------------------------------------------------------- /africanus/model/spi/examples/README.rst: -------------------------------------------------------------------------------- 1 | Simple SPI Fitter 2 | ================= 3 | 4 | Fits a simple spectral index model to image cubes. Usage is as follows 5 | 6 | .. code-block:: bash 7 | 8 | $ ./simple_spi_fitter.py --fitsmodel=/path/to/model.fits 9 | 10 | Run 11 | 12 | .. code-block:: bash 13 | 14 | $ ./simple_spi_fitter.py -h 15 | 16 | for documentation of the various options. In principle the model file is the 17 | only compulsary input if the beam parameters are specified. 18 | If they are not supplied the residual image cube needs to be provided as input 19 | so that these can be taken from the header. This means you either have to 20 | specify the beam parameters manually or pass in a residual cube with a header 21 | which contains beam parameters. 22 | 23 | The residual is also used to determine the weights in the different imaging 24 | bands. The weights will be set as 1/rms**2 in each imaging band, given that 25 | the rms is not 0 in which case the weight is also set to zero for that band. 26 | 27 | The threshold above which to fit components is set as a multiple of the rms 28 | in the residual, where the multiple is specified by the --threshold parameter. 29 | If the residual is not provided, then this threshold can be specified through 30 | the --maxDR parameter. In this case only components above model.max()/maxDR 31 | will be fit. 32 | 33 | It is also possible to perform an image space correction for the primary beam 34 | pattern. Currently only real and imaginary fits beams are supported. 35 | Please see the documentation for the --beammodel parameter for the required 36 | format. 37 | 38 | For parallel FFT's you will also need to install pypocketfft which can be 39 | found at git+https://gitlab.mpcdf.mpg.de/mtr/pypocketfft. Install via 40 | 41 | .. code-block:: bash 42 | 43 | $ pip3 install git+https://gitlab.mpcdf.mpg.de/mtr/pypocketfft 44 | -------------------------------------------------------------------------------- /africanus/model/spi/dask.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from africanus.model.spi.component_spi import SPI_DOCSTRING 5 | from africanus.model.spi.component_spi import ( 6 | fit_spi_components as np_fit_spi_components, 7 | ) 8 | 9 | from africanus.util.requirements import requires_optional 10 | 11 | try: 12 | from dask.array.core import blockwise 13 | except ImportError as e: 14 | opt_import_error = e 15 | else: 16 | opt_import_error = None 17 | 18 | 19 | def _fit_spi_components_wrapper( 20 | data, weights, freqs, freq0, alphai, I0i, beam, tol, maxiter 21 | ): 22 | return np_fit_spi_components( 23 | data[0], 24 | weights[0], 25 | freqs[0], 26 | freq0, 27 | alphai, 28 | I0i, 29 | beam[0] if beam is not None else beam, 30 | tol=tol, 31 | maxiter=maxiter, 32 | ) 33 | 34 | 35 | @requires_optional("dask.array", opt_import_error) 36 | def fit_spi_components( 37 | data, weights, freqs, freq0, alphai=None, I0i=None, beam=None, tol=1e-5, maxiter=100 38 | ): 39 | """Dask wrapper fit_spi_components function""" 40 | return blockwise( 41 | _fit_spi_components_wrapper, 42 | ("vars", "comps"), 43 | data, 44 | ("comps", "chan"), 45 | weights, 46 | ("chan",), 47 | freqs, 48 | ("chan",), 49 | freq0, 50 | None, 51 | alphai, 52 | ("comps",) if alphai is not None else None, 53 | I0i, 54 | ("comps",) if I0i is not None else None, 55 | beam, 56 | ("comps", "chan") if beam is not None else None, 57 | tol, 58 | None, 59 | maxiter, 60 | None, 61 | new_axes={"vars": 4}, 62 | dtype=data.dtype, 63 | ) 64 | 65 | 66 | fit_spi_components.__doc__ = SPI_DOCSTRING.substitute( 67 | array_type=":class:`dask.array.Array`" 68 | ) 69 | -------------------------------------------------------------------------------- /africanus/gridding/perleypolyhedron/policies/phase_transform_policies.py: -------------------------------------------------------------------------------- 1 | from africanus.util.numba import overload 2 | from numpy import pi, cos, sin, sqrt 3 | 4 | 5 | def phase_norotate(vis, uvw, lambdas, ra0, dec0, ra, dec, policy_type, phasesign=1.0): 6 | pass 7 | 8 | 9 | def phase_rotate(vis, uvw, lambdas, ra0, dec0, ra, dec, policy_type, phasesign=1.0): 10 | """ 11 | Convert ra,dec to l,m,n based on Synthesis Imaging II, Pg. 388 12 | The phase term (as documented in Perley & Cornwell (1992)) 13 | calculation requires the delta l,m,n coordinates. 14 | Through simplification l0,m0,n0 = (0,0,1) (assume dec == dec0 and 15 | ra == ra0, and the simplification follows) 16 | l,m,n is then calculated using the new and original phase centres 17 | as per the relation on Pg. 388 18 | lambdas has the same shape as vis 19 | """ 20 | d_ra = ra - ra0 21 | d_dec = dec 22 | d_decp = dec0 23 | c_d_dec = cos(d_dec) 24 | s_d_dec = sin(d_dec) 25 | s_d_ra = sin(d_ra) 26 | c_d_ra = cos(d_ra) 27 | c_d_decp = cos(d_decp) 28 | s_d_decp = sin(d_decp) 29 | ll = c_d_dec * s_d_ra 30 | mm = s_d_dec * c_d_decp - c_d_dec * s_d_decp * c_d_ra 31 | nn = -(1 - sqrt(1 - ll * ll - mm * mm)) 32 | for c in range(lambdas.size): 33 | x = phasesign * 2 * pi * (uvw[0] * ll + uvw[1] * mm + uvw[2] * nn) / lambdas[c] 34 | vis[c, :] *= cos(x) + 1.0j * sin(x) 35 | 36 | 37 | def policy(vis, uvw, lambdas, ra0, dec0, ra, dec, policy_type, phasesign=1.0): 38 | pass 39 | 40 | 41 | @overload(policy, inline="always", prefer_literal=True) 42 | def policy_impl(vis, uvw, lambdas, ra0, dec0, ra, dec, policy_type, phasesign=1.0): 43 | if policy_type.literal_value == "None" or policy_type.literal_value is None: 44 | return phase_norotate 45 | elif policy_type.literal_value == "phase_rotate": 46 | return phase_rotate 47 | else: 48 | raise ValueError("Invalid baseline transform policy type") 49 | -------------------------------------------------------------------------------- /docs/util-api.rst: -------------------------------------------------------------------------------- 1 | --------- 2 | Utilities 3 | --------- 4 | 5 | Command Line 6 | ~~~~~~~~~~~~ 7 | 8 | .. currentmodule:: africanus.util.cmdline 9 | 10 | .. autosummary:: 11 | parse_python_assigns 12 | 13 | .. autofunction:: parse_python_assigns 14 | 15 | 16 | Requirements Handling 17 | ~~~~~~~~~~~~~~~~~~~~~ 18 | 19 | .. currentmodule:: africanus.util.requirements 20 | 21 | .. autosummary:: 22 | requires_optional 23 | 24 | .. autofunction:: requires_optional 25 | 26 | Shapes 27 | ~~~~~~ 28 | 29 | .. currentmodule:: africanus.util.shapes 30 | 31 | .. autosummary:: 32 | aggregate_chunks 33 | corr_shape 34 | 35 | .. autofunction:: aggregate_chunks 36 | .. autofunction:: corr_shape 37 | 38 | 39 | Beams 40 | ~~~~~ 41 | 42 | .. currentmodule:: africanus.util.beams 43 | 44 | .. autosummary:: 45 | beam_filenames 46 | beam_grids 47 | 48 | 49 | .. autofunction:: beam_filenames 50 | .. autofunction:: beam_grids 51 | 52 | Code 53 | ~~~~ 54 | 55 | .. currentmodule:: africanus.util.code 56 | 57 | .. autosummary:: 58 | format_code 59 | memoize_on_key 60 | 61 | .. autofunction:: format_code 62 | .. autoclass:: memoize_on_key 63 | 64 | 65 | dask 66 | ~~~~ 67 | 68 | .. currentmodule:: africanus.util.dask_util 69 | 70 | .. autosummary:: 71 | EstimatingProgressBar 72 | 73 | .. autoclass:: EstimatingProgressBar 74 | :members: 75 | :no-inherited-members: 76 | :exclude-members: register, unregister 77 | 78 | CUDA 79 | ~~~~ 80 | 81 | .. currentmodule:: africanus.util.cuda 82 | 83 | .. autosummary:: 84 | grids 85 | 86 | .. autofunction:: grids 87 | 88 | Patterns 89 | ~~~~~~~~ 90 | 91 | .. currentmodule:: africanus.util.patterns 92 | 93 | .. autosummary:: 94 | Multiton 95 | LazyProxy 96 | LazyProxyMultiton 97 | 98 | .. autoclass:: Multiton 99 | :exclude-members: __call__, mro 100 | .. autoclass:: LazyProxy 101 | .. autoclass:: LazyProxyMultiton 102 | -------------------------------------------------------------------------------- /africanus/rime/cuda/feeds.cu.j2: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | {%- if out_type in ("int", "unsigned int", "float") %} 4 | {% set vis_align = 8 %} 5 | {%- elif out_type == "double" %} 6 | {% set vis_align = 16 %} 7 | {%- else %} 8 | {{ throw("Unhandled alignment type %s" % out_type) }} 9 | {%- endif %} 10 | 11 | {%- if feed_type == "linear" %} 12 | typedef struct __align__({{vis_align}}) { 13 | {{out_type}}2 XX; 14 | {{out_type}}2 XY; 15 | {{out_type}}2 YX; 16 | {{out_type}}2 YY; 17 | } correlation; 18 | {%- elif feed_type == "circular" %} 19 | typedef struct __align__({{vis_align}}) { 20 | {{out_type}}2 RR; 21 | {{out_type}}2 RL; 22 | {{out_type}}2 LR; 23 | {{out_type}}2 LL; 24 | } correlation; 25 | {%- else %} 26 | {{ throw("Invalid feed_type %s" % feed_type) }} 27 | {%- endif %} 28 | 29 | extern "C" __global__ void {{kernel_name}}( 30 | const CArray<{{pa_type}}, 1> parangles, 31 | CArray<{{out_type}}2, 2> feed_rotation) 32 | { 33 | const ptrdiff_t & npa = parangles.shape()[0]; 34 | const int pa = blockIdx.x*blockDim.x + threadIdx.x; 35 | 36 | // Guard 37 | if(pa >= npa) 38 | { return; } 39 | 40 | {{pa_type}} pa_sin; 41 | {{pa_type}} pa_cos; 42 | 43 | {{sincos_fn}}(parangles[pa], &pa_sin, &pa_cos); 44 | 45 | correlation vis; 46 | {%- if feed_type == 'linear' %} 47 | vis.XX = make_{{out_type}}2(pa_cos, 0.0); 48 | vis.XY = make_{{out_type}}2(pa_sin, 0.0); 49 | vis.YX = make_{{out_type}}2(-pa_sin, 0.0); 50 | vis.YY = make_{{out_type}}2(pa_cos, 0.0); 51 | {%- elif feed_type == 'circular' %} 52 | vis.RR = make_{{out_type}}2(pa_cos, -pa_sin); 53 | vis.RL = make_{{out_type}}2(0.0, 0.0); 54 | vis.LR = make_{{out_type}}2(0.0, 0.0); 55 | vis.LL = make_{{out_type}}2(pa_cos, pa_sin); 56 | {%- else %} 57 | {{ throw("Invalid feed_type %s" % feed_type) }} 58 | {% endif %} 59 | 60 | // Recast output type 61 | correlation * base_ptr = reinterpret_cast(&feed_rotation[0]); 62 | base_ptr[pa] = vis; 63 | 64 | } 65 | -------------------------------------------------------------------------------- /africanus/rime/parangles_casa.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import threading 4 | 5 | import numpy as np 6 | 7 | from africanus.util.requirements import requires_optional 8 | 9 | try: 10 | import pyrap.measures 11 | import pyrap.quanta as pq 12 | except ImportError as e: 13 | casa_import_error = e 14 | have_casa_parangles = False 15 | else: 16 | casa_import_error = None 17 | have_casa_parangles = True 18 | 19 | # Create thread local storage for the measures server 20 | _thread_local = threading.local() 21 | 22 | 23 | @requires_optional("pyrap.measures", "pyrap.quanta", casa_import_error) 24 | def casa_parallactic_angles( 25 | times, antenna_positions, field_centre, zenith_frame="AZEL" 26 | ): 27 | """ 28 | Computes parallactic angles per timestep for the given 29 | reference antenna position and field centre. 30 | """ 31 | 32 | try: 33 | meas_serv = _thread_local.meas_serv 34 | except AttributeError: 35 | # Create a measures server 36 | _thread_local.meas_serv = meas_serv = pyrap.measures.measures() 37 | 38 | # Create direction measure for the zenith 39 | zenith = meas_serv.direction(zenith_frame, "0deg", "90deg") 40 | 41 | # Create position measures for each antenna 42 | reference_positions = [ 43 | meas_serv.position("itrf", *(pq.quantity(x, "m") for x in pos)) 44 | for pos in antenna_positions 45 | ] 46 | 47 | # Compute field centre in radians 48 | fc_rad = meas_serv.direction( 49 | "J2000", *(pq.quantity(f, "rad") for f in field_centre) 50 | ) 51 | 52 | return np.asarray( 53 | [ 54 | # Set current time as the reference frame 55 | meas_serv.do_frame(meas_serv.epoch("UTC", pq.quantity(t, "s"))) 56 | and [ # Set antenna position as the reference frame 57 | meas_serv.do_frame(rp) 58 | and meas_serv.posangle(fc_rad, zenith).get_value("rad") 59 | for rp in reference_positions 60 | ] 61 | for t in times 62 | ] 63 | ) 64 | -------------------------------------------------------------------------------- /africanus/model/coherency/dask.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from africanus.model.coherency.conversion import ( 5 | convert_setup, 6 | convert_impl, 7 | CONVERT_DOCS, 8 | ) 9 | 10 | from africanus.util.requirements import requires_optional 11 | 12 | try: 13 | import dask.array as da 14 | except ImportError as e: 15 | da_import_error = e 16 | else: 17 | da_import_error = None 18 | 19 | 20 | def convert_wrapper(np_input, mapping=None, in_shape=None, out_shape=None, dtype_=None): 21 | result = convert_impl(np_input, mapping, in_shape, out_shape, dtype_) 22 | 23 | # Introduce extra singleton dimension at the end of our shape 24 | return result.reshape(result.shape + (1,) * len(in_shape)) 25 | 26 | 27 | @requires_optional("dask.array", da_import_error) 28 | def convert(input, input_schema, output_schema, implicit_stokes=False): 29 | mapping, in_shape, out_shape, dtype = convert_setup( 30 | input, input_schema, output_schema, implicit_stokes 31 | ) 32 | 33 | n_free_dims = len(input.shape) - len(in_shape) 34 | free_dims = tuple("dim-%d" % i for i in range(n_free_dims)) 35 | in_corr_dims = tuple("icorr-%d" % i for i in range(len(in_shape))) 36 | out_corr_dims = tuple("ocorr-%d" % i for i in range(len(out_shape))) 37 | 38 | # Output dimension are new dimensions 39 | new_axes = {d: s for d, s in zip(out_corr_dims, out_shape)} 40 | 41 | # Note the dummy in_corr_dims introduced at the end of our output, 42 | # We do this to prevent a contraction over the input dimensions 43 | # (which can be arbitrary) within the wrapper class 44 | res = da.core.blockwise( 45 | convert_wrapper, 46 | free_dims + out_corr_dims + in_corr_dims, 47 | input, 48 | free_dims + in_corr_dims, 49 | mapping=mapping, 50 | in_shape=in_shape, 51 | out_shape=out_shape, 52 | new_axes=new_axes, 53 | dtype_=dtype, 54 | dtype=dtype, 55 | ) 56 | 57 | # Now contract over the dummy dimensions 58 | start = len(free_dims) + len(out_corr_dims) 59 | end = start + len(in_corr_dims) 60 | return res.sum(axis=list(range(start, end))) 61 | 62 | 63 | try: 64 | convert.__doc__ = CONVERT_DOCS.substitute(array_type=":class:`dask.array.Array`") 65 | except AttributeError: 66 | pass 67 | -------------------------------------------------------------------------------- /africanus/util/tests/test_patterns.py: -------------------------------------------------------------------------------- 1 | """Keep this file in sync with the dask-ms version""" 2 | 3 | import pickle 4 | 5 | import pytest 6 | 7 | from africanus.util.patterns import Multiton, LazyProxy, LazyProxyMultiton 8 | 9 | 10 | class DummyResource: 11 | def __init__(self, arg, tracker, kw=None): 12 | self.arg = arg 13 | self.tracker = tracker 14 | self.kw = kw 15 | self.value = None 16 | 17 | def set(self, value): 18 | self.value = value 19 | 20 | def close(self): 21 | self.tracker.closed = True 22 | 23 | 24 | class Tracker: 25 | def __init__(self): 26 | self.closed = False 27 | 28 | def __eq__(self, other): 29 | return True 30 | 31 | def __hash__(self): 32 | return 0 33 | 34 | 35 | @pytest.mark.parametrize("finalise", [True, False]) 36 | @pytest.mark.parametrize("cls", [LazyProxyMultiton, LazyProxy]) 37 | def test_lazy(cls, finalise): 38 | def _inner(tracker): 39 | if finalise: 40 | fn = (DummyResource, DummyResource.close) 41 | else: 42 | fn = DummyResource 43 | 44 | obj = cls(fn, "test.txt", tracker, kw="w") 45 | obj.set(5) 46 | 47 | assert obj.arg == "test.txt" 48 | assert obj.kw == "w" 49 | assert obj.value == 5 50 | 51 | obj2 = pickle.loads(pickle.dumps(obj)) 52 | assert obj.__lazy_eq__(obj2) 53 | assert obj.__lazy_hash__() == obj2.__lazy_hash__() 54 | 55 | if cls is LazyProxyMultiton: 56 | assert obj is obj2 57 | else: 58 | assert obj is not obj2 59 | 60 | tracker = Tracker() 61 | assert tracker.closed is False 62 | _inner(tracker) 63 | assert tracker.closed is finalise 64 | 65 | 66 | def test_multiton(): 67 | class A(metaclass=Multiton): 68 | def __init__(self, *args, **kwargs): 69 | self.args = args 70 | self.kwargs = kwargs 71 | 72 | class B(metaclass=Multiton): 73 | def __init__(self, *args, **kwargs): 74 | self.args = args 75 | self.kwargs = kwargs 76 | 77 | a1 = A(1, 2, 3) 78 | a2 = A(1, 2) 79 | a3 = A(1) 80 | b1 = B(1) 81 | 82 | assert a1 is A(1, 2, 3) 83 | assert a2 is A(1, 2) 84 | assert a3 is A(1) 85 | assert b1 is B(1) 86 | assert a1 is not a2 87 | assert a1 is not a3 88 | assert a1 is not b1 89 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/tests/test_structref_setter.py: -------------------------------------------------------------------------------- 1 | import numba 2 | from numba.experimental import structref 3 | from numba.extending import intrinsic 4 | from numba.core import types, errors 5 | 6 | import numpy as np 7 | 8 | 9 | @structref.register 10 | class StateStructRef(types.StructRef): 11 | def preprocess_fields(self, fields): 12 | """Disallow literal types in field definitions""" 13 | return tuple((n, types.unliteral(t)) for n, t in fields) 14 | 15 | 16 | def test_structref_setter(): 17 | @intrinsic 18 | def constructor(typingctx, arg_tuple): 19 | if not isinstance(arg_tuple, types.BaseTuple): 20 | raise errors.TypingError(f"{arg_tuple} must be a Tuple") 21 | 22 | # Determine the fields for the StructRef 23 | names = [f"arg_{i}" for i in range(len(arg_tuple))] 24 | fields = list((name, typ) for name, typ in zip(names, arg_tuple)) 25 | state_type = StateStructRef(fields) 26 | 27 | sig = state_type(arg_tuple) 28 | 29 | def codegen(context, builder, signature, args): 30 | def make_struct(): 31 | """Allocate the structure""" 32 | return structref.new(state_type) 33 | 34 | state = context.compile_internal(builder, make_struct, state_type(), []) 35 | 36 | # Now assign each argument 37 | U = structref._Utils(context, builder, state_type) 38 | data_struct = U.get_data_struct(state) 39 | 40 | for i, name in enumerate(names): 41 | value = builder.extract_value(args[0], i) 42 | value_type = signature.args[0][i] 43 | field_type = state_type.field_dict[name] 44 | casted = context.cast(builder, value, value_type, field_type) 45 | old_value = getattr(data_struct, name) 46 | context.nrt.incref(builder, value_type, casted) 47 | context.nrt.decref(builder, value_type, old_value) 48 | setattr(data_struct, name, casted) 49 | 50 | return state 51 | 52 | return sig, codegen 53 | 54 | @numba.njit 55 | def fn(*args): 56 | s = constructor(args) 57 | print(s.arg_0) 58 | print(s.arg_1) 59 | print(s.arg_2) 60 | return s.arg_2 61 | 62 | args = (2, "b", np.arange(10)) 63 | assert np.array_equal(fn(*args), args[2]) 64 | -------------------------------------------------------------------------------- /africanus/rime/cuda/phase.cu.j2: -------------------------------------------------------------------------------- 1 | // #include 2 | #include 3 | // #include 4 | 5 | #define blockdimx {{blockdimx}} 6 | #define blockdimy {{blockdimy}} 7 | 8 | #define minus_two_pi_over_c {{minus_two_pi_over_c}} 9 | 10 | extern "C" __global__ void {{kernel_name}}( 11 | const CArray<{{lm_type}}, 2> lm, 12 | const CArray<{{uvw_type}}, 2> uvw, 13 | const CArray<{{freq_type}}, 1> frequency, 14 | CArray<{{out_type}}2, 3> complex_phase) 15 | { 16 | int row = blockIdx.y*blockDim.y + threadIdx.y; 17 | int chan = blockIdx.x*blockDim.x + threadIdx.x; 18 | 19 | // Return if outside the grid 20 | if(row >= uvw.shape()[0] || chan >= frequency.shape()[0]) 21 | { return; } 22 | 23 | // Reinterpret inputs as vector types 24 | const {{lm_type}}2 * lm_ptr = reinterpret_cast( 25 | &lm[0]); 26 | const {{uvw_type}}3 * uvw_ptr = reinterpret_cast( 27 | &uvw[0]); 28 | {{out_type}}2 * complex_phase_ptr = reinterpret_cast<{{out_type}}2 *>( 29 | &complex_phase[0]); 30 | 31 | __shared__ struct { 32 | {{uvw_type}}3 uvw[{{blockdimy}}]; 33 | {{freq_type}} frequency[{{blockdimx}}]; 34 | } shared; 35 | 36 | // UVW coordinates vary along y dimension only 37 | if(threadIdx.x == 0) 38 | { shared.uvw[threadIdx.y] = uvw_ptr[row]; } 39 | 40 | // Frequencies vary along x dimension only 41 | if(threadIdx.y == 0) 42 | { shared.frequency[threadIdx.x] = frequency[chan]; } 43 | 44 | __syncthreads(); 45 | 46 | for(int source = 0; source < lm.shape()[0]; ++source) 47 | { 48 | {{lm_type}}2 rlm = lm_ptr[source]; 49 | {{lm_type}} n = {{sqrt_fn}}(1.0 - rlm.x*rlm.x - rlm.y*rlm.y) - 1.0; 50 | {{out_type}} real_phase = rlm.x*shared.uvw[threadIdx.y].x + 51 | rlm.y*shared.uvw[threadIdx.y].y + 52 | n*shared.uvw[threadIdx.y].z; 53 | 54 | real_phase = minus_two_pi_over_c * 55 | real_phase * 56 | shared.frequency[threadIdx.x]; 57 | 58 | {{out_type}}2 cplx_phase; 59 | {{sincos_fn}}(real_phase, &cplx_phase.y, &cplx_phase.x); 60 | 61 | 62 | ptrdiff_t idx [] = {source, row, chan}; 63 | complex_phase[idx] = cplx_phase; 64 | } 65 | } 66 | -------------------------------------------------------------------------------- /africanus/gps/examples/generate_phase_only_gains.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Simulates phase-only gains where phases are drawn from a 5 | Gaussian process with covariance function given by cov_func 6 | """ 7 | from __future__ import absolute_import 8 | from __future__ import division 9 | from __future__ import print_function 10 | 11 | import numpy as np 12 | from pyrap.tables import table 13 | from africanus.gps.kernels import exponential_squared as cov_func 14 | from africanus.linalg import kronecker_tools as kt 15 | from africanus.coordinates.coordinates import radec_to_lm 16 | import argparse 17 | 18 | 19 | def create_parser(): 20 | p = argparse.ArgumentParser() 21 | p.add_argument("--ms", type=str) 22 | p.add_argument("--lsm", type=str) 23 | p.add_argument("--gain_file", type=str) 24 | return p 25 | 26 | 27 | args = create_parser().parse_args() 28 | 29 | # get times and normalise 30 | ms = table(args.ms) 31 | time = ms.getcol("TIME") 32 | ant1 = ms.getcol("ANTENNA1") 33 | ant2 = ms.getcol("ANTENNA2") 34 | n_ant = int(np.maximum(ant1.max(), ant2.max()) + 1) 35 | ms.close() 36 | time = np.unique(time) 37 | time -= time.min() 38 | time /= time.max() 39 | time = np.ascontiguousarray(time) 40 | n_time = time.size 41 | time_cov = cov_func(time, time, 0.25, 0.2) 42 | 43 | 44 | # get freqs and normalise 45 | spw = table(args.ms + "::SPECTRAL_WINDOW") 46 | freq = spw.getcol("CHAN_FREQ").squeeze() 47 | spw.close() 48 | freq -= freq.min() 49 | freq /= freq.max() 50 | freq = np.ascontiguousarray(freq) 51 | n_freq = freq.size 52 | freq_cov = cov_func(freq, freq, 1.0, 0.25) 53 | 54 | # get directions 55 | sources = np.loadtxt(args.lsm, skiprows=1, usecols=(1, 2)) 56 | lm = radec_to_lm(sources) 57 | lm = np.ascontiguousarray(lm) 58 | n_dir = lm.shape[0] 59 | lm_cov = cov_func(lm, lm, 1.0, 0.5) 60 | 61 | # create kronecker matrix and get cholesky factor 62 | K = np.array([time_cov, freq_cov, lm_cov], dtype=object) 63 | L = kt.kron_cholesky(K) 64 | 65 | # sample phases 66 | gains = np.zeros((n_time, n_ant, n_freq, n_dir, 2)) 67 | for p in range(n_ant): 68 | xi = np.random.randn(n_time * n_freq * n_dir) 69 | samp = kt.kron_matvec(L, xi).reshape(n_time, n_freq, n_dir) 70 | gains[:, p, :, :, 0] = samp 71 | xi = np.random.randn(n_time * n_freq * n_dir) 72 | samp = kt.kron_matvec(L, xi).reshape(n_time, n_freq, n_dir) 73 | gains[:, p, :, :, 1] = samp 74 | 75 | # convert to gains 76 | gains = np.exp(1.0j * gains) 77 | 78 | # save result 79 | np.save(args.gain_file, gains) 80 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/dask.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from africanus.util.requirements import requires_optional 4 | from africanus.experimental.rime.fused.core import RimeFactory, consolidate_args 5 | 6 | try: 7 | import dask.array as da 8 | except ImportError as e: 9 | opt_import_err = e 10 | else: 11 | opt_import_err = None 12 | 13 | 14 | def rime_dask_wrapper(factory, names, nconcat_dims, *args): 15 | # Call the rime factory 16 | assert len(names) == len(args) 17 | out = factory(**dict(zip(names, args))) 18 | # (1) Reintroduce source dimension, 19 | # (2) slice the existing dimensions 20 | # (3) expand by the contraction dims which will 21 | # be removed in the later dask reduction 22 | return out[(None,) + (slice(None),) * out.ndim + (None,) * nconcat_dims] 23 | 24 | 25 | @requires_optional("dask.array", opt_import_err) 26 | def rime(rime_spec, *args, **kw): 27 | """Like :func:`~africanus.experimental.rime.fused.core.rime`, but for 28 | a dask paradigm""" 29 | factory = RimeFactory(rime_spec) 30 | names, args = factory.dask_blockwise_args(**consolidate_args(args, kw)) 31 | 32 | dims = ("source", "row", "chan", "corr") 33 | contract_dims = set(d for ds in args[1::2] if ds is not None for d in ds) 34 | contract_dims -= set(dims) 35 | out_dims = dims + tuple(contract_dims) 36 | 37 | # Source and concatenation dimension are reduced to 1 element 38 | adjust_chunks = {"source": 1, **{d: 1 for d in contract_dims}} 39 | new_axes = {"corr": len(factory.rime_spec.corrs)} 40 | 41 | # This is needed otherwise, dask will call rime_dask_wrapper 42 | # with dummy arguments to infer the output dtype. 43 | # This incurs memory allocations within numba, as well as 44 | # exceptions, leading to memory leaks as described 45 | # in https://github.com/numba/numba/issues/3263 46 | meta = np.empty((0,) * len(out_dims), dtype=np.complex128) 47 | 48 | # Construct the wrapper call from given arguments 49 | out = da.blockwise( 50 | rime_dask_wrapper, 51 | out_dims, 52 | factory, 53 | None, 54 | names, 55 | None, 56 | len(contract_dims), 57 | None, 58 | *args, 59 | concatenate=False, 60 | adjust_chunks=adjust_chunks, 61 | new_axes=new_axes, 62 | meta=meta, 63 | ) 64 | 65 | # Contract over source and concatenation dims 66 | axes = (0,) + tuple(range(len(dims), len(dims) + len(contract_dims))) 67 | return out.sum(axis=axes) 68 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/terms/feed_rotation.py: -------------------------------------------------------------------------------- 1 | from africanus.experimental.rime.fused.terms.core import Term 2 | 3 | 4 | class FeedRotation(Term): 5 | """Feed Rotation Term""" 6 | 7 | def __init__(self, configuration, feed_type, corrs): 8 | if configuration not in {"left", "right"}: 9 | raise ValueError( 10 | f"FeedRotation configuration must " 11 | f"be either 'left' or 'right'. " 12 | f"Got {configuration}" 13 | ) 14 | 15 | if feed_type not in {"linear", "circular"}: 16 | raise ValueError( 17 | f"FeedRotation feed_type must be " 18 | f"either 'linear' or 'circular'. " 19 | f"Got {feed_type}" 20 | ) 21 | 22 | if len(corrs) != 4: 23 | raise ValueError( 24 | f"Four correlations required for " 25 | f"feed rotation but {corrs} were " 26 | f"specified" 27 | ) 28 | 29 | super().__init__(configuration) 30 | self.feed_type = feed_type 31 | 32 | def init_fields(self, typingctx, init_state, feed_parangle): 33 | def dummy(init_state, feed_parangle): 34 | pass 35 | 36 | return [], dummy 37 | 38 | def dask_schema(self, feed_parangle): 39 | return {} 40 | 41 | def sampler(self): 42 | left = self.configuration == "left" 43 | linear = self.feed_type == "linear" 44 | 45 | def feed_rotation(state, s, r, t, f1, f2, a1, a2, c): 46 | a = state.antenna1_inverse[r] if left else state.antenna2_inverse[r] 47 | f = state.feed1_inverse[r] if left else state.feed2_inverse[r] 48 | sin_a = state.feed_parangle[t, f, a, 0, 0] 49 | cos_a = state.feed_parangle[t, f, a, 0, 1] 50 | sin_b = state.feed_parangle[t, f, a, 1, 0] 51 | cos_b = state.feed_parangle[t, f, a, 1, 1] 52 | 53 | # https://casa.nrao.edu/aips2_docs/notes/185/node6.html 54 | if linear: 55 | return cos_a, sin_a, -sin_b, cos_b 56 | else: 57 | # e^{ix} = cos(x) + i.sin(x) 58 | return ( 59 | 0.5 * ((cos_a + cos_b) - (sin_a + sin_b) * 1j), 60 | 0.5 * ((cos_a - cos_b) + (sin_a - sin_b) * 1j), 61 | 0.5 * ((cos_a - cos_b) - (sin_a - sin_b) * 1j), 62 | 0.5 * ((cos_a + cos_b) + (sin_a + sin_b) * 1j), 63 | ) 64 | 65 | return feed_rotation 66 | -------------------------------------------------------------------------------- /africanus/rime/cuda/beam_freq_interp.cu.j2: -------------------------------------------------------------------------------- 1 | // #include 2 | #include 3 | // #include 4 | 5 | #define blockdimx {{blockdimx}} 6 | #define beam_nud_limit {{beam_nud_limit}} 7 | 8 | extern "C" __global__ void {{kernel_name}}( 9 | const CArray<{{freq_type}}, 1> frequencies, 10 | const CArray<{{beam_freq_type}}, 1> beam_freq_map, 11 | CArray<{{freq_type}}, 2> freq_data) 12 | { 13 | const int nud = beam_freq_map.shape()[0]; 14 | const int nchan = frequencies.shape()[0]; 15 | const int chan = blockIdx.x*blockDim.x + threadIdx.x; 16 | 17 | __shared__ struct { 18 | {{beam_freq_type}} beam_freqs[beam_nud_limit]; 19 | } shared; 20 | 21 | // Load in beam frequency map across 22 | if(chan < beam_nud_limit) 23 | { shared.beam_freqs[threadIdx.x] = beam_freq_map[chan]; } 24 | 25 | __syncthreads(); 26 | 27 | if(chan >= nchan) 28 | { return; } 29 | 30 | {{freq_type}} freq = frequencies[chan]; 31 | int lower = 0; 32 | int upper = nud - 1; 33 | 34 | // Warp divergence 35 | while(lower <= upper) 36 | { 37 | int mid = lower + (upper - lower)/2; 38 | {{freq_type}} beam_freq = shared.beam_freqs[mid]; 39 | 40 | if(beam_freq < freq) 41 | { lower = mid + 1; } 42 | else if(beam_freq > freq) 43 | { upper = mid - 1; } 44 | else 45 | { 46 | lower = mid; 47 | break; 48 | } 49 | } 50 | 51 | lower = min(lower, upper); 52 | upper = lower + 1; 53 | 54 | {{freq_type}} freq_scale; 55 | {{freq_type}} freq_ratio; 56 | {{freq_type}} freq_grid; 57 | 58 | // Frequency below the beam cube 59 | if(lower == -1) 60 | { 61 | freq_scale = freq / shared.beam_freqs[0]; 62 | freq_ratio = 1.0; 63 | freq_grid = 0.0; 64 | } 65 | // Frequency above the beam cube 66 | else if(upper == nud) 67 | { 68 | freq_scale = freq / shared.beam_freqs[nud - 1]; 69 | freq_ratio = 0.0; 70 | freq_grid = nud - 2; 71 | } 72 | // Standard interpolation case 73 | else 74 | { 75 | freq_scale = 1.0; 76 | {{freq_type}} freq_low = shared.beam_freqs[lower]; 77 | {{freq_type}} freq_high = shared.beam_freqs[upper]; 78 | freq_ratio = (freq_high - freq) / (freq_high - freq_low); 79 | freq_grid = lower; 80 | } 81 | 82 | freq_data[0*nchan + chan] = freq_scale; 83 | freq_data[1*nchan + chan] = freq_ratio; 84 | freq_data[2*nchan + chan] = freq_grid; 85 | } 86 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: clean clean-test clean-pyc clean-build docs help 2 | .DEFAULT_GOAL := help 3 | 4 | define BROWSER_PYSCRIPT 5 | import os, webbrowser, sys 6 | 7 | try: 8 | from urllib import pathname2url 9 | except: 10 | from urllib.request import pathname2url 11 | 12 | webbrowser.open("file://" + pathname2url(os.path.abspath(sys.argv[1]))) 13 | endef 14 | export BROWSER_PYSCRIPT 15 | 16 | define PRINT_HELP_PYSCRIPT 17 | import re, sys 18 | 19 | for line in sys.stdin: 20 | match = re.match(r'^([a-zA-Z_-]+):.*?## (.*)$$', line) 21 | if match: 22 | target, help = match.groups() 23 | print("%-20s %s" % (target, help)) 24 | endef 25 | export PRINT_HELP_PYSCRIPT 26 | 27 | BROWSER := python -c "$$BROWSER_PYSCRIPT" 28 | 29 | help: 30 | @python -c "$$PRINT_HELP_PYSCRIPT" < $(MAKEFILE_LIST) 31 | 32 | clean: clean-build clean-pyc clean-test ## remove all build, test, coverage and Python artifacts 33 | 34 | clean-build: ## remove build artifacts 35 | rm -fr build/ 36 | rm -fr dist/ 37 | rm -fr .eggs/ 38 | find . -name '*.egg-info' -exec rm -fr {} + 39 | find . -name '*.egg' -exec rm -f {} + 40 | 41 | clean-pyc: ## remove Python file artifacts 42 | find . -name '*.pyc' -exec rm -f {} + 43 | find . -name '*.pyo' -exec rm -f {} + 44 | find . -name '*~' -exec rm -f {} + 45 | find . -name '__pycache__' -exec rm -fr {} + 46 | 47 | clean-test: ## remove test and coverage artifacts 48 | rm -fr .tox/ 49 | rm -f .coverage 50 | rm -fr htmlcov/ 51 | 52 | lint: ## check style with flake8 53 | flake8 africanus tests 54 | 55 | test: ## run tests quickly with the default Python 56 | py.test 57 | 58 | test-all: ## run tests on every Python version with tox 59 | tox 60 | 61 | coverage: ## check code coverage quickly with the default Python 62 | coverage run --source africanus -m pytest 63 | coverage report -m 64 | coverage html 65 | $(BROWSER) htmlcov/index.html 66 | 67 | docs: ## generate Sphinx HTML documentation, including API docs 68 | rm -f docs/africanus.rst 69 | rm -f docs/modules.rst 70 | sphinx-apidoc -o docs/ africanus 71 | $(MAKE) -C docs clean 72 | $(MAKE) -C docs html 73 | $(BROWSER) docs/_build/html/index.html 74 | 75 | servedocs: docs ## compile the docs watching for changes 76 | watchmedo shell-command -p '*.rst' -c '$(MAKE) -C docs html' -R -D . 77 | 78 | release: dist ## package and upload a release 79 | twine upload dist/* 80 | 81 | dist: clean ## builds source and wheel package 82 | python setup.py sdist 83 | python setup.py bdist_wheel 84 | ls -l dist 85 | 86 | install: clean ## install the package to the active Python's site-packages 87 | python setup.py install 88 | -------------------------------------------------------------------------------- /africanus/rime/parangles.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import warnings 5 | 6 | from .parangles_astropy import have_astropy_parangles, astropy_parallactic_angles 7 | from .parangles_casa import have_casa_parangles, casa_parallactic_angles 8 | 9 | _discovered_backends = [] 10 | 11 | if have_astropy_parangles: 12 | _discovered_backends.append("astropy") 13 | 14 | if have_casa_parangles: 15 | _discovered_backends.append("casa") 16 | 17 | 18 | _standard_backends = set(["casa", "astropy", "test"]) 19 | 20 | 21 | def parallactic_angles(times, antenna_positions, field_centre, backend="casa"): 22 | """ 23 | Computes parallactic angles per timestep for the given 24 | reference antenna position and field centre. 25 | 26 | Parameters 27 | ---------- 28 | times : :class:`numpy.ndarray` 29 | Array of Mean Julian Date times in 30 | *seconds* with shape :code:`(time,)`, 31 | antenna_positions : :class:`numpy.ndarray` 32 | Antenna positions of shape :code:`(ant, 3)` 33 | in *metres* in the *ITRF* frame. 34 | field_centre : :class:`numpy.ndarray` 35 | Field centre of shape :code:`(2,)` in *radians* 36 | backend : {'casa', 'test'}, optional 37 | Backend to use for calculating the parallactic angles. 38 | 39 | * ``casa`` defers to an implementation 40 | depending on ``python-casacore``. 41 | This backend should be used by default. 42 | * ``test`` creates parallactic angles 43 | by multiplying the ``times`` and ``antenna_position`` 44 | arrays. It exist solely for testing. 45 | 46 | Returns 47 | ------- 48 | parallactic_angles : :class:`numpy.ndarray` 49 | Parallactic angles of shape :code:`(time,ant)` 50 | """ 51 | if backend not in _standard_backends: 52 | raise ValueError( 53 | "'%s' is not one of the " 54 | "standard backends '%s'" % (backend, _standard_backends) 55 | ) 56 | 57 | if not field_centre.shape == (2,): 58 | raise ValueError("Invalid field_centre shape %s" % (field_centre.shape,)) 59 | 60 | if backend == "astropy": 61 | warnings.warn("astropy backend currently returns the incorrect values") 62 | return astropy_parallactic_angles(times, antenna_positions, field_centre) 63 | elif backend == "casa": 64 | return casa_parallactic_angles(times, antenna_positions, field_centre) 65 | elif backend == "test": 66 | return times[:, None] * (antenna_positions.sum(axis=1)[None, :]) 67 | else: 68 | raise ValueError("Invalid backend %s" % backend) 69 | -------------------------------------------------------------------------------- /africanus/rime/cuda/tests/test_cuda_predict.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | import pytest 6 | 7 | from africanus.rime.predict import predict_vis as np_predict_vis 8 | from africanus.rime.cuda.predict import predict_vis 9 | from africanus.rime.tests.test_predict import ( 10 | corr_shape_parametrization, 11 | die_presence_parametrization, 12 | dde_presence_parametrization, 13 | chunk_parametrization, 14 | rc, 15 | ) 16 | 17 | 18 | @corr_shape_parametrization 19 | @dde_presence_parametrization 20 | @die_presence_parametrization 21 | @chunk_parametrization 22 | def test_cuda_predict_vis( 23 | corr_shape, idm, einsum_sig1, einsum_sig2, a1j, blj, a2j, g1j, bvis, g2j, chunks 24 | ): 25 | np.random.seed(40) 26 | 27 | cp = pytest.importorskip("cupy") 28 | 29 | s = sum(chunks["source"]) 30 | t = sum(chunks["time"]) 31 | a = sum(chunks["antenna"]) 32 | c = sum(chunks["channels"]) 33 | r = sum(chunks["rows"]) 34 | 35 | a1_jones = rc((s, t, a, c) + corr_shape) 36 | bl_jones = rc((s, r, c) + corr_shape) 37 | a2_jones = rc((s, t, a, c) + corr_shape) 38 | g1_jones = rc((t, a, c) + corr_shape) 39 | base_vis = rc((r, c) + corr_shape) 40 | g2_jones = rc((t, a, c) + corr_shape) 41 | 42 | # Add 10 to the index to test time index normalisation 43 | time_idx = np.concatenate( 44 | [np.full(rows, i + 10, dtype=np.int32) for i, rows in enumerate(chunks["rows"])] 45 | ) 46 | 47 | ant1 = np.concatenate( 48 | [np.random.randint(0, a, rows, dtype=np.int32) for rows in chunks["rows"]] 49 | ) 50 | 51 | ant2 = np.concatenate( 52 | [np.random.randint(0, a, rows, dtype=np.int32) for rows in chunks["rows"]] 53 | ) 54 | 55 | assert ant1.size == r 56 | 57 | model_vis = predict_vis( 58 | cp.asarray(time_idx), 59 | cp.asarray(ant1), 60 | cp.asarray(ant2), 61 | cp.asarray(a1_jones) if a1j else None, 62 | cp.asarray(bl_jones) if blj else None, 63 | cp.asarray(a2_jones) if a2j else None, 64 | cp.asarray(g1_jones) if g1j else None, 65 | cp.asarray(base_vis) if bvis else None, 66 | cp.asarray(g2_jones) if g2j else None, 67 | ) 68 | 69 | np_model_vis = np_predict_vis( 70 | time_idx, 71 | ant1, 72 | ant2, 73 | a1_jones if a1j else None, 74 | bl_jones if blj else None, 75 | a2_jones if a2j else None, 76 | g1_jones if g1j else None, 77 | base_vis if bvis else None, 78 | g2_jones if g2j else None, 79 | ) 80 | 81 | np.testing.assert_array_almost_equal(cp.asnumpy(model_vis), np_model_vis) 82 | -------------------------------------------------------------------------------- /docs/dft-api.rst: -------------------------------------------------------------------------------- 1 | ------------------------ 2 | Direct Fourier Transform 3 | ------------------------ 4 | 5 | Functions used to compute the discretised 6 | direct Fourier transform (DFT) 7 | for an ideal interferometer. 8 | The DFT for an ideal interferometer is 9 | defined as 10 | 11 | .. math:: 12 | 13 | V(u,v,w) = \int B(l,m) e^{-2\pi i 14 | \left( ul + vm + w(n-1)\right)} 15 | \frac{dl dm}{n} 16 | 17 | where :math:`u,v,w` are data space coordinates and 18 | where visibilities :math:`V` have been obtained. 19 | The :math:`l,m,n` are signal space coordinates at which 20 | we wish to reconstruct the signal :math:`B`. Note that 21 | the signal correspondes to the brightness matrix 22 | and not the Stokes parameters. We adopt the convention 23 | where we absorb the fixed coordinate :math:`n` in the 24 | denominator into the image. 25 | Note that the data space coordinates have an implicit 26 | dependence on frequency and time and that the image 27 | has an implicit dependence on frequency. 28 | The discretised form of the DFT can be written as 29 | 30 | .. math:: 31 | 32 | V(u,v,w) = \sum_s e^{-2 \pi i 33 | (u l_s + v m_s + w (n_s - 1))} \cdot B_s 34 | 35 | where :math:`s` labels the source (or pixel) location. 36 | If only a single correlation is present :math:`B = I`, 37 | this can be cast into a matrix equation as follows 38 | 39 | .. math:: 40 | 41 | V = R I 42 | 43 | where :math:`R` is the operator that maps an 44 | image to visibility space. This mapping is 45 | implemented by the :func:`~africanus.dft.im_to_vis` 46 | function. If multiple correlations are present then 47 | each one is mapped to its corresponding visibility. 48 | An imaging algorithm also requires the adjoint 49 | denoted :math:`R^\dagger` which is simply the 50 | complex conjugate transpose of :math:`R`. 51 | The dirty image is obtained by applying the 52 | adjoint operator to the visibilities 53 | 54 | .. math:: 55 | 56 | I^D = R^\dagger V 57 | 58 | This is implemented by the 59 | :func:`~africanus.dft.vis_to_im` 60 | function. 61 | Note that an imaging algorithm using these 62 | operators will actually reconstruct 63 | :math:`\frac{I}{n}` but that it is trivial 64 | to obtain :math:`I` since :math:`n` is 65 | known at each location in the image. 66 | 67 | 68 | Numpy 69 | ~~~~~ 70 | 71 | .. currentmodule:: africanus.dft 72 | 73 | .. autosummary:: 74 | im_to_vis 75 | vis_to_im 76 | 77 | .. autofunction:: im_to_vis 78 | .. autofunction:: vis_to_im 79 | 80 | Dask 81 | ~~~~ 82 | 83 | .. currentmodule:: africanus.dft.dask 84 | 85 | .. autosummary:: 86 | im_to_vis 87 | vis_to_im 88 | 89 | .. autofunction:: im_to_vis 90 | .. autofunction:: vis_to_im 91 | -------------------------------------------------------------------------------- /africanus/rime/cuda/feeds.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import logging 5 | from os.path import join as pjoin 6 | 7 | import numpy as np 8 | 9 | from africanus.rime.feeds import FEED_ROTATION_DOCS 10 | from africanus.util.code import format_code, memoize_on_key 11 | from africanus.util.cuda import cuda_function, grids 12 | from africanus.util.jinja2 import jinja_env 13 | from africanus.util.requirements import requires_optional 14 | 15 | try: 16 | import cupy as cp 17 | from cupy._core._scalar import get_typename as _get_typename 18 | from cupy.cuda.compiler import CompileException 19 | except ImportError as e: 20 | opt_import_error = e 21 | else: 22 | opt_import_error = None 23 | 24 | 25 | log = logging.getLogger(__name__) 26 | 27 | 28 | _TEMPLATE_PATH = pjoin("rime", "cuda", "feeds.cu.j2") 29 | 30 | 31 | def _key_fn(parallactic_angles, feed_type): 32 | return (parallactic_angles.dtype, feed_type) 33 | 34 | 35 | @memoize_on_key(_key_fn) 36 | def _generate_kernel(parallactic_angles, feed_type): 37 | dtype = parallactic_angles.dtype 38 | 39 | # Block sizes 40 | if dtype == np.float32: 41 | block = (1024, 1, 1) 42 | elif dtype == np.float64: 43 | block = (512, 1, 1) 44 | else: 45 | raise TypeError("Unhandled type %s" % dtype) 46 | 47 | # Create template 48 | render = jinja_env.get_template(_TEMPLATE_PATH).render 49 | name = "feed_rotation" 50 | 51 | code = render( 52 | kernel_name=name, 53 | feed_type=feed_type, 54 | sincos_fn=cuda_function("sincos", dtype), 55 | pa_type=_get_typename(dtype), 56 | out_type=_get_typename(dtype), 57 | ) 58 | 59 | # Complex output type 60 | out_dtype = np.result_type(dtype, np.complex64) 61 | return cp.RawKernel(code, name), block, out_dtype 62 | 63 | 64 | @requires_optional("cupy", opt_import_error) 65 | def feed_rotation(parallactic_angles, feed_type="linear"): 66 | """Cupy implementation of the feed_rotation kernel.""" 67 | kernel, block, out_dtype = _generate_kernel(parallactic_angles, feed_type) 68 | in_shape = parallactic_angles.shape 69 | parallactic_angles = parallactic_angles.ravel() 70 | grid = grids((parallactic_angles.shape[0], 1, 1), block) 71 | out = cp.empty(shape=(parallactic_angles.shape[0], 4), dtype=out_dtype) 72 | 73 | try: 74 | kernel(grid, block, (parallactic_angles, out)) 75 | except CompileException: 76 | log.exception(format_code(kernel.code)) 77 | raise 78 | 79 | return out.reshape(in_shape + (2, 2)) 80 | 81 | 82 | try: 83 | feed_rotation.__doc__ = FEED_ROTATION_DOCS.substitute( 84 | array_type=":class:`cupy.ndarray`" 85 | ) 86 | except AttributeError: 87 | pass 88 | -------------------------------------------------------------------------------- /africanus/rime/cuda/phase.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import logging 5 | from os.path import join as pjoin 6 | 7 | import numpy as np 8 | 9 | from africanus.constants import minus_two_pi_over_c 10 | from africanus.util.jinja2 import jinja_env 11 | from africanus.rime.phase import PHASE_DELAY_DOCS 12 | from africanus.util.code import memoize_on_key, format_code 13 | from africanus.util.cuda import cuda_function, grids 14 | from africanus.util.requirements import requires_optional 15 | 16 | try: 17 | import cupy as cp 18 | from cupy._core._scalar import get_typename as _get_typename 19 | from cupy.cuda.compiler import CompileException 20 | except ImportError: 21 | pass 22 | 23 | log = logging.getLogger(__name__) 24 | 25 | 26 | def _key_fn(lm, uvw, frequency): 27 | return (lm.dtype, uvw.dtype, frequency.dtype) 28 | 29 | 30 | _TEMPLATE_PATH = pjoin("rime", "cuda", "phase.cu.j2") 31 | 32 | 33 | @memoize_on_key(_key_fn) 34 | def _generate_kernel(lm, uvw, frequency): 35 | # Floating point output type 36 | out_dtype = np.result_type(lm, uvw, frequency) 37 | 38 | # Block sizes 39 | blockdimx = 32 if frequency.dtype == np.float32 else 16 40 | blockdimy = 32 if uvw.dtype == np.float32 else 16 41 | block = (blockdimx, blockdimy, 1) 42 | 43 | # Create template 44 | render = jinja_env.get_template(_TEMPLATE_PATH).render 45 | name = "phase_delay" 46 | 47 | code = render( 48 | kernel_name=name, 49 | lm_type=_get_typename(lm.dtype), 50 | uvw_type=_get_typename(uvw.dtype), 51 | freq_type=_get_typename(frequency.dtype), 52 | out_type=_get_typename(out_dtype), 53 | sqrt_fn=cuda_function("sqrt", lm.dtype), 54 | sincos_fn=cuda_function("sincos", out_dtype), 55 | minus_two_pi_over_c=minus_two_pi_over_c, 56 | blockdimx=blockdimx, 57 | blockdimy=blockdimy, 58 | ) 59 | 60 | # Complex output type 61 | out_dtype = np.result_type(out_dtype, np.complex64) 62 | return cp.RawKernel(code, name), block, out_dtype 63 | 64 | 65 | @requires_optional("cupy") 66 | def phase_delay(lm, uvw, frequency): 67 | kernel, block, out_dtype = _generate_kernel(lm, uvw, frequency) 68 | grid = grids((frequency.shape[0], uvw.shape[0], 1), block) 69 | out = cp.empty( 70 | shape=(lm.shape[0], uvw.shape[0], frequency.shape[0]), dtype=out_dtype 71 | ) 72 | 73 | try: 74 | kernel(grid, block, (lm, uvw, frequency, out)) 75 | except CompileException: 76 | log.exception(format_code(kernel.code)) 77 | raise 78 | 79 | return out 80 | 81 | 82 | try: 83 | phase_delay.__doc__ = PHASE_DELAY_DOCS.substitute( 84 | array_type=":class:`cupy.ndarray`" 85 | ) 86 | except AttributeError: 87 | pass 88 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/terms/gaussian.py: -------------------------------------------------------------------------------- 1 | from africanus.constants import c as lightspeed 2 | from numba import types 3 | from numba.core.errors import TypingError 4 | import numpy as np 5 | 6 | from africanus.experimental.rime.fused.terms.core import Term 7 | 8 | 9 | class Gaussian(Term): 10 | """Gaussian Amplitude Term""" 11 | 12 | def dask_schema(self, uvw, chan_freq, gauss_shape): 13 | assert uvw.ndim == 2 14 | assert chan_freq.ndim == 1 15 | assert gauss_shape.ndim == 2 16 | 17 | return { 18 | "uvw": ("row", "uvw"), 19 | "chan_freq": ("chan",), 20 | "gauss_shape": ("source", "gauss_shape_params"), 21 | } 22 | 23 | def init_fields(self, typingctx, init_state, uvw, chan_freq, gauss_shape): 24 | if not isinstance(uvw, types.Array) or uvw.ndim != 2: 25 | raise TypingError(f"uvw {uvw} should be a (row, uvw) array") 26 | 27 | if not isinstance(chan_freq, types.Array) or chan_freq.ndim != 1: 28 | raise TypingError(f"chan_freq {chan_freq} should be a (chan,) array") 29 | 30 | if not isinstance(gauss_shape, types.Array) or gauss_shape.ndim != 2: 31 | raise TypingError( 32 | f"gauss_shape {gauss_shape} should be a (source, gauss_shape_params) array" 33 | ) 34 | 35 | guv_dtype = typingctx.unify_types(uvw.dtype, chan_freq.dtype, gauss_shape.dtype) 36 | fields = [("gauss_uv", guv_dtype[:, :, :]), ("scaled_freq", chan_freq)] 37 | 38 | fwhm = 2.0 * np.sqrt(2.0 * np.log(2.0)) 39 | fwhminv = 1.0 / fwhm 40 | gauss_scale = fwhminv * np.sqrt(2.0) * np.pi / lightspeed 41 | 42 | def gaussian_init(init_state, uvw, chan_freq, gauss_shape): 43 | nsrc, _ = gauss_shape.shape 44 | nrow, _ = uvw.shape 45 | 46 | gauss_uv = np.empty((nsrc, nrow, 2), dtype=guv_dtype) 47 | scaled_freq = chan_freq * gauss_scale 48 | 49 | for s in range(nsrc): 50 | emaj, emin, angle = gauss_shape[s] 51 | 52 | # Convert to l-projection, m-projection, ratio 53 | el = emaj * np.sin(angle) 54 | em = emaj * np.cos(angle) 55 | er = emin / (1.0 if emaj == 0.0 else emaj) 56 | 57 | for r in range(uvw.shape[0]): 58 | u = uvw[r, 0] 59 | v = uvw[r, 1] 60 | 61 | gauss_uv[s, r, 0] = (u * em - v * el) * er 62 | gauss_uv[s, r, 1] = u * el + v * em 63 | 64 | return gauss_uv, scaled_freq 65 | 66 | return fields, gaussian_init 67 | 68 | def sampler(self): 69 | def gaussian_sample(state, s, r, t, f1, f2, a1, a2, c): 70 | fu1 = state.gauss_uv[s, r, 0] * state.scaled_freq[c] 71 | fv1 = state.gauss_uv[s, r, 1] * state.scaled_freq[c] 72 | return np.exp(-(fu1 * fu1 + fv1 * fv1)) 73 | 74 | return gaussian_sample 75 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "codex-africanus" 3 | description = "Radio Astronomy Building Blocks" 4 | version = "0.4.3" 5 | readme = "README.rst" 6 | authors = [{name = "Simon Perkins", email = "simon.perkins@gmail.com"}] 7 | requires-python = '>=3.10,<4.0' 8 | dependencies = ['appdirs >=1.4.4', 'decorator >=5.1.1', 'numpy >=2.0', 'numba >=0.60'] 9 | 10 | [project.optional-dependencies] 11 | astropy = ['astropy >=6.1.4'] 12 | cuda = ['cupy >=13.3.0', 'jinja2 >=3.1.4'] 13 | dask = ['dask[array] >=2024.0'] 14 | ducc0 = ['ducc0 >=0.35.0'] 15 | jax = ['jax >=0.4.35', 'jaxlib >=0.4.35'] 16 | scipy = ['scipy >=1.14.1'] 17 | testing = ['flaky >=3.8.1', 'pytest >=8.0.0'] 18 | python-casacore = ['python-casacore >=3.6.1'] 19 | complete = ['astropy', 'dask', 'ducc0', 'jax', 'jaxlib', 'python-casacore', 'scipy'] 20 | complete-cuda = ['astropy', 'cupy', 'jinja2', 'dask', 'ducc0', 'jax', 'jaxlib', 'python-casacore', 'scipy'] 21 | dev = [ 22 | "pre-commit>=4.3.0", 23 | "ruff>=0.14.0", 24 | "tbump>=6.11.0", 25 | ] 26 | doc = [ 27 | "numpydoc>=1.9.0", 28 | "pydata-sphinx-theme>=0.16.1", 29 | "pygments>=2.19.2", 30 | "ruff>=0.14.0", 31 | "sphinx>=8.1.3", 32 | "sphinx-copybutton>=0.5.2", 33 | ] 34 | 35 | [build-system] 36 | requires = ["hatchling"] 37 | build-backend = "hatchling.build" 38 | 39 | [tool.hatch.build.targets.wheel] 40 | packages = ["africanus"] 41 | 42 | [tool.ruff] 43 | exclude = ["turbo-sim.py"] 44 | line-length = 88 45 | indent-width = 4 46 | target-version = "py311" 47 | 48 | [tool.ruff.lint] 49 | extend-select = ["I"] 50 | select = [ 51 | # flake8-builtins 52 | "A", 53 | # flake8-bugbear 54 | "B", 55 | # isort 56 | "I001", 57 | "I002", 58 | # tidy imports 59 | "TID" 60 | ] 61 | 62 | [tool.tbump.version] 63 | current = "0.4.3" 64 | 65 | # Example of a semver regexp. 66 | # Make sure this matches current_version before 67 | # using tbump 68 | regex = ''' 69 | (?P\d+) 70 | \. 71 | (?P\d+) 72 | \. 73 | (?P\d+) 74 | ''' 75 | 76 | [tool.tbump.git] 77 | message_template = "Bump to {new_version}" 78 | tag_template = "{new_version}" 79 | 80 | # For each file to patch, add a [[tool.tbump.file]] config 81 | # section containing the path of the file, relative to the 82 | # tbump.toml location. 83 | [[tool.tbump.file]] 84 | src = "pyproject.toml" 85 | 86 | [[tool.tbump.file]] 87 | src = "africanus/__init__.py" 88 | 89 | [[tool.tbump.file]] 90 | src = "docs/conf.py" 91 | 92 | # You can specify a list of commands to 93 | # run after the files have been patched 94 | # and before the git commit is made 95 | 96 | # [[tool.tbump.before_commit]] 97 | # name = "check changelog" 98 | # cmd = "grep -q {new_version} Changelog.rst" 99 | 100 | # Or run some commands after the git tag and the branch 101 | # have been pushed: 102 | # [[tool.tbump.after_push]] 103 | # name = "publish" 104 | # cmd = "./publish.sh" 105 | -------------------------------------------------------------------------------- /africanus/model/shape/dask.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from africanus.model.shape.gaussian_shape import ( 5 | gaussian as np_gaussian, 6 | GAUSSIAN_DOCS, 7 | ) 8 | from africanus.util.requirements import requires_optional 9 | 10 | from africanus.model.shape.shapelets import shapelet as nb_shapelet 11 | from africanus.model.shape.shapelets import ( 12 | shapelet_with_w_term as nb_shapelet_with_w_term, 13 | ) 14 | 15 | import numpy as np 16 | 17 | try: 18 | import dask.array as da 19 | except ImportError as e: 20 | opt_import_error = e 21 | else: 22 | opt_import_error = None 23 | 24 | 25 | def gaussian_wrapper(uvw, frequency, shape_params): 26 | return np_gaussian(uvw[0], frequency, shape_params[0]) 27 | 28 | 29 | @requires_optional("dask.array", opt_import_error) 30 | def gaussian(uvw, frequency, shape_params): 31 | dtype = np.result_type(uvw.dtype, frequency.dtype, shape_params.dtype) 32 | 33 | return da.blockwise( 34 | gaussian_wrapper, 35 | ("source", "row", "chan"), 36 | uvw, 37 | ("row", "uvw-comp"), 38 | frequency, 39 | ("chan",), 40 | shape_params, 41 | ("source", "shape-comp"), 42 | dtype=dtype, 43 | ) 44 | 45 | 46 | def _shapelet_wrapper(coords, frequency, coeffs, beta, delta_lm): 47 | return nb_shapelet(coords[0], frequency, coeffs[0][0], beta[0], delta_lm[0]) 48 | 49 | 50 | @requires_optional("dask.array", opt_import_error) 51 | def shapelet(coords, frequency, coeffs, beta, delta_lm): 52 | dtype = np.complex128 53 | return da.blockwise( 54 | _shapelet_wrapper, 55 | ("row", "chan", "source"), 56 | coords, 57 | ("row", "coord-comp"), 58 | frequency, 59 | ("chan",), 60 | coeffs, 61 | ("source", "nmax1", "nmax2"), 62 | beta, 63 | ("source", "beta-comp"), 64 | delta_lm, 65 | ("delta_lm-comp",), 66 | dtype=dtype, 67 | ) 68 | 69 | 70 | def _shapelet_with_w_term_wrapper(coords, frequency, coeffs, beta, delta_lm, lm): 71 | return nb_shapelet_with_w_term( 72 | coords[0], frequency, coeffs[0][0], beta[0], delta_lm[0], lm[0] 73 | ) 74 | 75 | 76 | @requires_optional("dask.array", opt_import_error) 77 | def shapelet_with_w_term(coords, frequency, coeffs, beta, delta_lm, lm): 78 | dtype = np.complex128 79 | return da.blockwise( 80 | _shapelet_with_w_term_wrapper, 81 | ("row", "chan", "source"), 82 | coords, 83 | ("row", "coord-comp"), 84 | frequency, 85 | ("chan",), 86 | coeffs, 87 | ("source", "nmax1", "nmax2"), 88 | beta, 89 | ("source", "beta-comp"), 90 | delta_lm, 91 | ("delta_lm-comp",), 92 | lm, 93 | ("source", "lm-comp"), 94 | dtype=dtype, 95 | ) 96 | 97 | 98 | try: 99 | gaussian.__doc__ = GAUSSIAN_DOCS.substitute(array_type=":class:`dask.array.Array`") 100 | except AttributeError: 101 | pass 102 | -------------------------------------------------------------------------------- /docs/rime-api.rst: -------------------------------------------------------------------------------- 1 | .. _rime-api-anchor: 2 | 3 | ----------------------------------------- 4 | Radio Interferometer Measurement Equation 5 | ----------------------------------------- 6 | 7 | Functions used to compute the terms of the 8 | Radio Interferometer Measurement Equation (RIME). 9 | It describes the response of an interferometer to a sky model. 10 | 11 | .. math:: 12 | 13 | V_{pq} = G_{p} \left( 14 | \sum_{s} E_{ps} L_{p} K_{ps} 15 | B_{s} 16 | K_{qs}^H L_{q}^H E_{qs}^H 17 | \right) G_{q}^H 18 | 19 | where for antenna :math:`p` and :math:`q`, and source :math:`s`: 20 | 21 | * :math:`G_{p}` represents direction-independent effects. 22 | * :math:`E_{ps}` represents direction-dependent effects. 23 | * :math:`L_{p}` represents the feed rotation. 24 | * :math:`K_{ps}` represents the phase delay term. 25 | * :math:`B_{s}` represents the brightness matrix. 26 | 27 | 28 | The RIME is more formally described in the following four papers: 29 | 30 | * `I. A full-sky Jones formalism `_ 31 | * `II. Calibration and direction-dependent effects `_ 32 | * `III. Addressing direction-dependent effects in 21cm WSRT observations of 3C147 `_ 33 | * `IV. A generalized tensor formalism `_ 34 | 35 | .. _rime_paper_i: https://arxiv.org/abs/1101.1764 36 | .. _rime_paper_ii: https://arxiv.org/abs/1101.1765 37 | .. _rime_paper_iii: https://arxiv.org/abs/1101.1768 38 | .. _rime_paper_iv: https://arxiv.org/abs/1106.0579 39 | 40 | 41 | Numpy 42 | ~~~~~ 43 | 44 | .. currentmodule:: africanus.rime 45 | 46 | .. autosummary:: 47 | predict_vis 48 | phase_delay 49 | parallactic_angles 50 | feed_rotation 51 | transform_sources 52 | beam_cube_dde 53 | zernike_dde 54 | wsclean_predict 55 | 56 | .. autofunction:: predict_vis 57 | .. autofunction:: phase_delay 58 | .. autofunction:: parallactic_angles 59 | .. autofunction:: feed_rotation 60 | .. autofunction:: transform_sources 61 | .. autofunction:: beam_cube_dde 62 | .. autofunction:: zernike_dde 63 | .. autofunction:: wsclean_predict 64 | 65 | Cuda 66 | ~~~~ 67 | 68 | .. currentmodule:: africanus.rime.cuda 69 | 70 | .. autosummary:: 71 | predict_vis 72 | phase_delay 73 | feed_rotation 74 | beam_cube_dde 75 | 76 | .. autofunction:: predict_vis 77 | .. autofunction:: phase_delay 78 | .. autofunction:: feed_rotation 79 | .. autofunction:: beam_cube_dde 80 | 81 | 82 | Dask 83 | ~~~~ 84 | 85 | .. currentmodule:: africanus.rime.dask 86 | 87 | .. autosummary:: 88 | predict_vis 89 | phase_delay 90 | parallactic_angles 91 | feed_rotation 92 | transform_sources 93 | beam_cube_dde 94 | zernike_dde 95 | wsclean_predict 96 | 97 | 98 | .. autofunction:: predict_vis 99 | .. autofunction:: phase_delay 100 | .. autofunction:: parallactic_angles 101 | .. autofunction:: feed_rotation 102 | .. autofunction:: transform_sources 103 | .. autofunction:: beam_cube_dde 104 | .. autofunction:: zernike_dde 105 | .. autofunction:: wsclean_predict 106 | -------------------------------------------------------------------------------- /africanus/rime/examples/N6251-sky-model.txt: -------------------------------------------------------------------------------- 1 | #format: name ra_d dec_d i spi freq0 emaj_s emin_s shapelet_coeffs 2 | N6251 0.00 -30.00 1.00 -0.700000 1085000000 8.33 20.0 1.9988189101699798123,-0.0271946441057400456,-0.2770276847939583398,-0.1908514566337720653,0.0724469969123310853,2.4401275492195297900,-0.0012046783205888478,0.0468895929047210280,-0.1213605238156697658,-0.3906235244057300049,0.0341431730487933641,-0.0119580607767994090,-0.2152219568731133237,0.0412995224154704826,0.8920076937275780438,-0.0037899175608049024,0.0033804923124753222,0.0082975562403374767,0.0684835444031911544,-0.1234270623207482376,-0.0758891195371745736,0.0082838993103408736,0.0026103188161130829,0.0296927393465825570,-0.0141320968864564189,-0.0514444296751662283,-0.0687210458695055038,-0.5679711359716888008,-0.0026726340887466525,0.0087095327660485413,-0.0072509561138205531,0.0032535310089492862,0.0163444683903988600,0.0341214378308733121,0.0034862872460355299,0.4964005201552225688,0.0117130367781834686,-0.0017458496564881406,-0.0025757968336746848,-0.0024951097474205236,0.0187499390817604389,0.0030249104473347374,0.0884866440683228922,-0.0631312466352219881,0.1602869523945408270,-0.0030560488631701595,0.0100273320877909307,-0.0022090140201403534,0.0101081468354888467,-0.0049381861147342523,0.0082325366434215057,-0.0081524880137551839,-0.0167729062412844167,-0.0321526535402662655,-0.0667729822368420822,0.0094561061040604744,-0.0027967370091000082,0.0019377134956919363,-0.0058944327475736052,0.0090372958167671269,-0.0038227851253948286,0.0048696065384362872,-0.0049489854658295064,0.0234292977300321163,-0.0283520548606985695,0.3959838652725587438,-0.0033680102151816636,0.0112060193741995787,-0.0033556328733418115,0.0097725199579149734,-0.0026228974137796775,0.0114149511415664560,-0.0036130424518936964,0.0203616137606379213,-0.0135385617600301544,0.0639233536320822915,-0.0867405945839860859,0.1146370756095908411,0.0077799787758118504,-0.0037090778938207811,0.0022337532882431071,-0.0060476667302850071,0.0092260915464434188,-0.0024560365638246449,0.0122368364228536350,-0.0076201981148041817,0.0172620503274020058,-0.0163421267350482881,0.0132957760927871343,-0.0709417084897278061,-0.0966586356857347118,-0.0034779403977029909,0.0107213272222486840,-0.0049423787241187275,0.0099352259245091452,-0.0035996367889153643,0.0099791065297096734,-0.0057015671002194034,0.0119779477500032427,-0.0062875740367258313,0.0170993404828291284,-0.0099559942710626872,0.0530078158186475315,-0.0452094166884028717,0.2040528594521089922,0.0067177399821413547,-0.0042047257579186014,0.0026617760731779891,-0.0055901939509616092,0.0070844273693341613,-0.0029839583017661174,0.0101018823806401196,-0.0062780890554478067,0.0169044096580723384,-0.0044054750777043870,0.0220117477191779098,-0.0135840937556933172,0.0673949945424858199,-0.0490309136417489133,0.1216595194855662893,-0.0033331110107027006,0.0089931300370690147,-0.0065534495402234990,0.0091255619444002376,-0.0035871131386896243,0.0093725652205782547,-0.0064653001557596563,0.0108673119018303753,-0.0054683529498851212,0.0140624537614863418,-0.0064475816532327257,0.0208312458615564147,-0.0212634674052380283,0.0418037655388663887,-0.0441691295990301949,-0.0540919313185133921 3 | -------------------------------------------------------------------------------- /africanus/experimental/rime/fused/terms/phase.py: -------------------------------------------------------------------------------- 1 | from africanus.constants import c as lightspeed 2 | from numba import types 3 | from numba.core.errors import TypingError, RequireLiteralValue 4 | import numpy as np 5 | 6 | from africanus.experimental.rime.fused.terms.core import Term 7 | 8 | 9 | class Phase(Term): 10 | """Phase Delay Term""" 11 | 12 | def dask_schema(self, lm, uvw, chan_freq, convention="fourier"): 13 | assert lm.ndim == 2 14 | assert uvw.ndim == 2 15 | assert chan_freq.ndim == 1 16 | assert isinstance(convention, str) 17 | 18 | return { 19 | "lm": ("source", "lm"), 20 | "uvw": ("row", "uvw"), 21 | "chan_freq": ("chan",), 22 | "convention": None, 23 | } 24 | 25 | def init_fields( 26 | self, typingctx, init_state, lm, uvw, chan_freq, convention="fourier" 27 | ): 28 | if not isinstance(lm, types.Array) or lm.ndim != 2: 29 | raise TypingError(f"lm {lm} should be a (source, lm) array") 30 | 31 | if not isinstance(uvw, types.Array) or uvw.ndim != 2: 32 | raise TypingError(f"uvw {uvw} should be a (row, uvw) array") 33 | 34 | if not isinstance(chan_freq, types.Array) or chan_freq.ndim != 1: 35 | raise TypingError(f"chan_freq {chan_freq} should be a (chan,) array") 36 | 37 | if not isinstance(convention, types.misc.UnicodeType): 38 | raise TypingError(f"convention {convention} should be a UnicodeType") 39 | 40 | assert lm.ndim == 2 41 | assert uvw.ndim == 2 42 | assert chan_freq.ndim == 1 43 | phase_dt = typingctx.unify_types(lm.dtype, uvw.dtype, chan_freq.dtype) 44 | fields = [("phase_dot", phase_dt[:, :])] 45 | 46 | def phase(init_state, lm, uvw, chan_freq, convention="fourier"): 47 | nsrc, _ = lm.shape 48 | nrow, _ = uvw.shape 49 | (nchan,) = chan_freq.shape 50 | 51 | phase_dot = np.empty((nsrc, nrow), dtype=phase_dt) 52 | 53 | zero = lm.dtype.type(0.0) 54 | one = lm.dtype.type(1.0) 55 | 56 | if convention == "fourier": 57 | C = phase_dt(-2.0 * np.pi / lightspeed) 58 | elif convention == "casa": 59 | C = phase_dt(2.0 * np.pi / lightspeed) 60 | else: 61 | raise ValueError('convention not in ("fourier", "casa")') 62 | 63 | for s in range(nsrc): 64 | l = lm[s, 0] # noqa 65 | m = lm[s, 1] 66 | n = one - l**2 - m**2 67 | n = np.sqrt(zero if n < zero else n) - one 68 | 69 | for r in range(nrow): 70 | u = uvw[r, 0] 71 | v = uvw[r, 1] 72 | w = uvw[r, 2] 73 | 74 | phase_dot[s, r] = C * (l * u + m * v + n * w) 75 | 76 | return phase_dot 77 | 78 | return fields, phase 79 | 80 | def sampler(self): 81 | def phase_sample(state, s, r, t, f1, f2, a1, a2, c): 82 | p = state.phase_dot[s, r] * state.chan_freq[c] 83 | return np.cos(p) + np.sin(p) * 1j 84 | 85 | return phase_sample 86 | -------------------------------------------------------------------------------- /africanus/calibration/phase_only/dask.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from africanus.calibration.phase_only.phase_only import COMPUTE_JHJ_DOCS 4 | from africanus.calibration.phase_only.phase_only import COMPUTE_JHR_DOCS 5 | from africanus.calibration.utils import check_type 6 | from africanus.calibration.phase_only import compute_jhj as np_compute_jhj 7 | from africanus.calibration.phase_only import compute_jhr as np_compute_jhr 8 | from africanus.util.requirements import requires_optional 9 | from africanus.calibration.utils.utils import DIAG_DIAG 10 | 11 | try: 12 | from dask.array.core import blockwise 13 | except ImportError as e: 14 | dask_import_error = e 15 | else: 16 | dask_import_error = None 17 | 18 | 19 | @requires_optional("dask.array", dask_import_error) 20 | def compute_jhj( 21 | time_bin_indices, time_bin_counts, antenna1, antenna2, jones, model, flag 22 | ): 23 | mode = check_type(jones, model, vis_type="model") 24 | 25 | if mode != DIAG_DIAG: 26 | raise NotImplementedError("Only DIAG-DIAG case has been implemented") 27 | 28 | jones_shape = ("row", "ant", "chan", "dir", "corr") 29 | vis_shape = ("row", "chan", "corr") 30 | model_shape = ("row", "chan", "dir", "corr") 31 | return blockwise( 32 | np_compute_jhj, 33 | jones_shape, 34 | time_bin_indices, 35 | ("row",), 36 | time_bin_counts, 37 | ("row",), 38 | antenna1, 39 | ("row",), 40 | antenna2, 41 | ("row",), 42 | jones, 43 | jones_shape, 44 | model, 45 | model_shape, 46 | flag, 47 | vis_shape, 48 | adjust_chunks={"row": antenna1.chunks[0]}, 49 | new_axes={"corr2": 2}, # why? 50 | dtype=model.dtype, 51 | align_arrays=False, 52 | ) 53 | 54 | 55 | @requires_optional("dask.array", dask_import_error) 56 | def compute_jhr( 57 | time_bin_indices, time_bin_counts, antenna1, antenna2, jones, residual, model, flag 58 | ): 59 | mode = check_type(jones, residual) 60 | 61 | if mode != DIAG_DIAG: 62 | raise NotImplementedError("Only DIAG-DIAG case has been implemented") 63 | 64 | jones_shape = ("row", "ant", "chan", "dir", "corr") 65 | vis_shape = ("row", "chan", "corr") 66 | model_shape = ("row", "chan", "dir", "corr") 67 | return blockwise( 68 | np_compute_jhr, 69 | jones_shape, 70 | time_bin_indices, 71 | ("row",), 72 | time_bin_counts, 73 | ("row",), 74 | antenna1, 75 | ("row",), 76 | antenna2, 77 | ("row",), 78 | jones, 79 | jones_shape, 80 | residual, 81 | vis_shape, 82 | model, 83 | model_shape, 84 | flag, 85 | vis_shape, 86 | adjust_chunks={"row": antenna1.chunks[0]}, 87 | new_axes={"corr2": 2}, # why? 88 | dtype=model.dtype, 89 | align_arrays=False, 90 | ) 91 | 92 | 93 | compute_jhj.__doc__ = COMPUTE_JHJ_DOCS.substitute( 94 | array_type=":class:`dask.array.Array`" 95 | ) 96 | 97 | compute_jhr.__doc__ = COMPUTE_JHR_DOCS.substitute( 98 | array_type=":class:`dask.array.Array`" 99 | ) 100 | -------------------------------------------------------------------------------- /africanus/util/cuda.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | 6 | _array_types = [np.ndarray] 7 | 8 | try: 9 | import dask.array as da 10 | except ImportError: 11 | pass 12 | else: 13 | _array_types.append(da.Array) 14 | 15 | try: 16 | import cupy as cp 17 | except ImportError: 18 | pass 19 | else: 20 | _array_types.append(cp.ndarray) 21 | 22 | _array_types = tuple(_array_types) 23 | 24 | cuda_fns = { 25 | np.dtype(np.float32): { 26 | "abs": "fabsf", 27 | "cos": "cosf", 28 | "floor": "floorf", 29 | "make2": "make_float2", 30 | "max": "fmaxf", 31 | "min": "fminf", 32 | "rsqrt": "rsqrtf", 33 | "sqrt": "sqrtf", 34 | "sin": "sinf", 35 | "sincos": "sincosf", 36 | "sincospi": "sincospif", 37 | }, 38 | np.dtype(np.float64): { 39 | "abs": "fabs", 40 | "cos": "cos", 41 | "floor": "floor", 42 | "make2": "make_double2", 43 | "max": "fmax", 44 | "min": "fmin", 45 | "rsqrt": "rsqrt", 46 | "sin": "sin", 47 | "sincos": "sincos", 48 | "sincospi": "sincospi", 49 | "sqrt": "sqrt", 50 | }, 51 | } 52 | 53 | 54 | numpy_to_cuda_type_map = { 55 | np.dtype("int8"): "char", 56 | np.dtype("uint8"): "unsigned char", 57 | np.dtype("int16"): "short", 58 | np.dtype("uint16"): "unsigned short", 59 | np.dtype("int32"): "int", 60 | np.dtype("uint32"): "unsigned int", 61 | np.dtype("float32"): "float", 62 | np.dtype("float64"): "double", 63 | np.dtype("complex64"): "float2", 64 | np.dtype("complex128"): "double2", 65 | } 66 | 67 | # Also map the types 68 | numpy_to_cuda_type_map.update({k.type: v for k, v in numpy_to_cuda_type_map.items()}) 69 | 70 | 71 | def grids(dims, blocks): 72 | """ 73 | Determine the grid size, given space dimensions sizes and blocks 74 | 75 | Parameters 76 | ---------- 77 | dims : tuple of ints 78 | `(x, y, z)` tuple 79 | 80 | Returns 81 | ------- 82 | tuple 83 | `(x, y, z)` grid size tuple 84 | """ 85 | if not len(dims) == 3: 86 | raise ValueError( 87 | "dims must be an (x, y, z) tuple. " 88 | "CUDA dimension ordering is inverted compared " 89 | "to NumPy" 90 | ) 91 | 92 | if not len(blocks) == 3: 93 | raise ValueError( 94 | "blocks must be an (x, y, z) tuple. " 95 | "CUDA dimension ordering is inverted compared " 96 | "to NumPy" 97 | ) 98 | 99 | return tuple((d + b - 1) // b for d, b in zip(dims, blocks)) 100 | 101 | 102 | def cuda_function(function_name, dtype): 103 | try: 104 | type_map = cuda_fns[dtype] 105 | except KeyError: 106 | raise ValueError("No registered functions for type %s" % dtype) 107 | 108 | try: 109 | return type_map[function_name] 110 | except KeyError: 111 | raise ValueError("Unknown CUDA function %s" % function_name) 112 | 113 | 114 | def cuda_type(dtype): 115 | if isinstance(dtype, _array_types): 116 | dtype = dtype.dtype 117 | 118 | try: 119 | return numpy_to_cuda_type_map[dtype] 120 | except KeyError: 121 | raise ValueError("No registered map for type %s" % dtype) 122 | -------------------------------------------------------------------------------- /africanus/gridding/util.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def estimate_cell_size(u, v, wavelength, factor=3.0, ny=None, nx=None): 5 | r""" 6 | Estimate the cell size in arcseconds given 7 | baseline ``u`` and ``v`` coordinates, as well 8 | as the ``wavelengths``, :math:`\lambda`. 9 | 10 | The cell size is computed as: 11 | 12 | .. math:: 13 | 14 | \Delta u = 1.0 / \left( 2 \times \text{ factor } 15 | \times \max (\vert u \vert) 16 | / \min( \lambda) \right) 17 | 18 | \Delta v = 1.0 / \left( 2 \times \text{ factor } 19 | \times \max (\vert v \vert) 20 | / \min( \lambda) \right) 21 | 22 | 23 | If ``ny`` and ``nx`` are provided the following checks are performed 24 | and exceptions are raised on failure: 25 | 26 | .. math:: 27 | 28 | \Delta u * \text{ ny } \leq \min (\lambda) / \min (\vert u \vert) 29 | 30 | \Delta v * \text{ nx } \leq \min (\lambda) / \min (\vert v \vert) 31 | 32 | Parameters 33 | ---------- 34 | u : :class:`numpy.ndarray` or float 35 | Maximum ``u`` coordinate in metres. 36 | v : :class:`numpy.ndarray` or float 37 | Maximum ``v`` coordinate in metres. 38 | wavelength : :class:`numpy.ndarray` or float 39 | Wavelengths, in metres. 40 | factor : float, optional 41 | Scaling factor 42 | ny : int, optional 43 | Grid y dimension 44 | nx : int, optional 45 | Grid x dimension 46 | 47 | Raises 48 | ------ 49 | ValueError 50 | If the cell size criteria are not matched. 51 | 52 | Returns 53 | ------- 54 | :class:`numpy.ndarray` 55 | Cell size of ``u`` and ``v`` in arcseconds with shape :code:`(2,)` 56 | """ 57 | if isinstance(u, np.ndarray): 58 | abs_u = np.abs(u) 59 | umax = abs_u.max() 60 | umin = abs_u.min() 61 | elif isinstance(u, float): 62 | umax = umin = abs(u) 63 | else: 64 | raise TypeError("Invalid u type %s" % type(u)) 65 | 66 | if isinstance(v, np.ndarray): 67 | abs_v = np.abs(v) 68 | vmax = abs_v.max() 69 | vmin = abs_v.min() 70 | elif isinstance(v, float): 71 | vmax = vmin = abs(v) 72 | else: 73 | raise TypeError("Invalid v type %s" % type(v)) 74 | 75 | if isinstance(wavelength, np.ndarray): 76 | wave_min = wavelength.min() 77 | elif isinstance(wavelength, float): 78 | wave_min = wavelength 79 | else: 80 | raise TypeError("Invalid wavelength type %s" % type(v)) 81 | 82 | umax /= wave_min 83 | vmax /= wave_min 84 | umin /= wave_min 85 | vmin /= wave_min 86 | 87 | u_cell_size = 1.0 / (2.0 * factor * umax) 88 | v_cell_size = 1.0 / (2.0 * factor * vmax) 89 | 90 | if ny is not None and u_cell_size * ny < (1.0 / umin): 91 | raise ValueError( 92 | "v_cell_size*ny [%f] < (1.0 / umin) [%f]" % (u_cell_size * ny, 1.0 / umin) 93 | ) 94 | 95 | if nx is not None and v_cell_size * nx < (1.0 / vmin): 96 | raise ValueError( 97 | "v_cell_size*nx [%f] < (1.0 / vmin) [%f]" % (v_cell_size * nx, 1.0 / vmin) 98 | ) 99 | 100 | # Convert radians to arcseconds 101 | return np.rad2deg([u_cell_size, v_cell_size]) * (60 * 60) 102 | -------------------------------------------------------------------------------- /africanus/util/shapes.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | def aggregate_chunks(chunks, max_chunks): 5 | """ 6 | Aggregate dask ``chunks`` together into chunks no larger than 7 | ``max_chunks``. 8 | 9 | .. code-block:: python 10 | 11 | chunks, max_c = ((3,4,6,3,6,7),(1,1,1,1,1,1)), (10,3) 12 | expected = ((7,9,6,7), (2,2,1,1)) 13 | assert aggregate_chunks(chunks, max_c) == expected 14 | 15 | 16 | Parameters 17 | ---------- 18 | chunks : sequence of tuples or tuple 19 | max_chunks : sequence of ints or int 20 | 21 | Returns 22 | ------- 23 | sequence of tuples or tuple 24 | 25 | """ 26 | 27 | if isinstance(max_chunks, int): 28 | chunks = (chunks,) 29 | max_chunks = (max_chunks,) 30 | 31 | singleton = True if len(max_chunks) == 1 else False 32 | 33 | if len(chunks) != len(max_chunks): 34 | raise ValueError("len(chunks) != len(max_chunks)") 35 | 36 | if not all(len(chunks[0]) == len(c) for c in chunks): 37 | raise ValueError("Number of chunks do not match") 38 | 39 | agg_chunks = [[] for _ in max_chunks] 40 | agg_chunk_counts = [0] * len(max_chunks) 41 | chunk_scratch = [0] * len(max_chunks) 42 | ndim = len(chunks[0]) 43 | 44 | # For each chunk dimension 45 | for di in range(ndim): 46 | # For each chunk 47 | aggregate = False 48 | 49 | for ci, chunk in enumerate(chunks): 50 | chunk_scratch[ci] = agg_chunk_counts[ci] + chunk[di] 51 | if chunk_scratch[ci] > max_chunks[ci]: 52 | aggregate = True 53 | 54 | if aggregate: 55 | for ci, chunk in enumerate(chunks): 56 | agg_chunks[ci].append(agg_chunk_counts[ci]) 57 | agg_chunk_counts[ci] = chunk[di] 58 | else: 59 | for ci, chunk in enumerate(chunks): 60 | agg_chunk_counts[ci] = chunk_scratch[ci] 61 | 62 | # Do the final aggregation 63 | for ci, chunk in enumerate(chunks): 64 | agg_chunks[ci].append(agg_chunk_counts[ci]) 65 | agg_chunk_counts[ci] = chunk[di] 66 | 67 | agg_chunks = tuple(tuple(ac) for ac in agg_chunks) 68 | 69 | return agg_chunks[0] if singleton else agg_chunks 70 | 71 | 72 | def corr_shape(ncorr, corr_shape): 73 | """ 74 | Returns the shape of the correlations, given 75 | ``ncorr`` and the type of correlation shape requested 76 | 77 | Parameters 78 | ---------- 79 | ncorr : integer 80 | Number of correlations 81 | corr_shape : {'flat', 'matrix'} 82 | Shape of output correlations 83 | 84 | 85 | Returns 86 | ------- 87 | tuple 88 | Shape tuple describing the correlation dimensions 89 | 90 | * If ``flat`` returns :code:`(ncorr,)` 91 | * If ``matrix`` returns 92 | 93 | * :code:`(1,)` if :code:`ncorr == 1` 94 | * :code:`(2,)` if :code:`ncorr == 2` 95 | * :code:`(2,2)` if :code:`ncorr == 4` 96 | 97 | 98 | """ 99 | if corr_shape == "flat": 100 | return (ncorr,) 101 | elif corr_shape == "matrix": 102 | if ncorr == 1: 103 | return (1,) 104 | elif ncorr == 2: 105 | return (2,) 106 | elif ncorr == 4: 107 | return (2, 2) 108 | else: 109 | raise ValueError("ncorr not in (1, 2, 4)") 110 | else: 111 | raise ValueError("corr_shape must be 'flat' or 'matrix'") 112 | -------------------------------------------------------------------------------- /africanus/rime/examples/tests/meqtrees/tdlconf.profiles: -------------------------------------------------------------------------------- 1 | [codex-compare-circular] 2 | me.e_advanced = 0 3 | me.e_enable = 0 4 | me.e_module = Siamese_OMS_pybeams_fits 5 | me.epe_enable = 0 6 | me.g_enable = 0 7 | me.l_advanced = 0 8 | me.l_enable = 0 9 | me.l_module = Rotation 10 | me.ncorr_enable = 0 11 | me.p_enable = 0 12 | me.sky.siamese_agw_azel_sky = 0 13 | me.sky.siamese_oms_fitsimage_sky = 0 14 | me.sky.siamese_oms_gridded_sky = 0 15 | me.sky.siamese_oms_transient_sky = 0 16 | me.sky.tiggerskymodel = 1 17 | me.use_jones_inspectors = 1 18 | me.use_skyjones_visualizers = 0 19 | me.use_smearing = 0 20 | me.z_enable = 0 21 | ms_sel.ddid_index = 0 22 | ms_sel.field_index = 0 23 | ms_sel.ms_corr_sel = 2x2 24 | ms_sel.ms_ifr_subset_str = all 25 | ms_sel.ms_polarization = LL LR RL RR 26 | ms_sel.ms_taql_str = None 27 | ms_sel.msname = data/vla_test.ms 28 | ms_sel.output_column = MODEL_DATA 29 | ms_sel.select_channels = 0 30 | ms_sel.tile_size = 100 31 | noise_from_sefd = 0 32 | noise_stddev = None 33 | pybeams_fits.ampl_interpolation = 1 34 | pybeams_fits.filename_pattern = JVLA-L-centred-$(corr)_$(reim).fits 35 | pybeams_fits.l_axis = -X 36 | pybeams_fits.l_beam_offset = 0.0 37 | pybeams_fits.m_axis = Y 38 | pybeams_fits.m_beam_offset = 0.0 39 | pybeams_fits.missing_is_null = 1 40 | pybeams_fits.normalize_gains = 0 41 | pybeams_fits.sky_rotation = 1 42 | pybeams_fits.spline_order = 1 43 | pybeams_fits.verbose_level = None 44 | random_seed = time 45 | read_ms_model = 0 46 | run_purr = 0 47 | sim_mode = sim only 48 | tensormeqmaker.psv_class = PSVTensor 49 | tiggerlsm.filename = meqtrees/sky_model.lsm.html 50 | tiggerlsm.lsm_subset = all 51 | tiggerlsm.null_subset = None 52 | tiggerlsm.solvable_sources = 0 53 | uvw_refant = default 54 | uvw_source = from MS 55 | rot_l.enable_pa = enabled 56 | 57 | 58 | [codex-compare-linear] 59 | me.e_advanced = 0 60 | me.e_enable = 0 61 | me.e_module = Siamese_OMS_pybeams_fits 62 | me.epe_enable = 0 63 | me.g_enable = 0 64 | me.l_advanced = 0 65 | me.l_enable = 0 66 | me.l_module = Rotation 67 | me.ncorr_enable = 0 68 | me.p_enable = 0 69 | me.sky.siamese_agw_azel_sky = 0 70 | me.sky.siamese_oms_fitsimage_sky = 0 71 | me.sky.siamese_oms_gridded_sky = 0 72 | me.sky.siamese_oms_transient_sky = 0 73 | me.sky.tiggerskymodel = 1 74 | me.use_jones_inspectors = 1 75 | me.use_skyjones_visualizers = 0 76 | me.use_smearing = 0 77 | me.z_enable = 0 78 | ms_sel.ddid_index = 0 79 | ms_sel.field_index = 0 80 | ms_sel.ms_corr_sel = 2x2 81 | ms_sel.ms_ifr_subset_str = all 82 | ms_sel.ms_polarization = XX XY YX YY 83 | ms_sel.ms_taql_str = None 84 | ms_sel.msname = data/WSRT.MS 85 | ms_sel.output_column = MODEL_DATA 86 | ms_sel.select_channels = 0 87 | ms_sel.tile_size = 100 88 | noise_from_sefd = 0 89 | noise_stddev = None 90 | pybeams_fits.ampl_interpolation = 1 91 | pybeams_fits.filename_pattern = beam_$(xy)_$(reim).fits 92 | pybeams_fits.l_axis = X 93 | pybeams_fits.l_beam_offset = 0.0 94 | pybeams_fits.m_axis = Y 95 | pybeams_fits.m_beam_offset = 0.0 96 | pybeams_fits.missing_is_null = 1 97 | pybeams_fits.normalize_gains = 0 98 | pybeams_fits.sky_rotation = 1 99 | pybeams_fits.spline_order = 1 100 | pybeams_fits.verbose_level = None 101 | random_seed = time 102 | read_ms_model = 0 103 | run_purr = 0 104 | sim_mode = sim only 105 | tensormeqmaker.psv_class = PSVTensor 106 | tiggerlsm.filename = meqtrees/sky_model.lsm.html 107 | tiggerlsm.lsm_subset = all 108 | tiggerlsm.null_subset = None 109 | tiggerlsm.solvable_sources = 0 110 | uvw_refant = default 111 | uvw_source = from MS 112 | rot_l.enable_pa = enabled 113 | -------------------------------------------------------------------------------- /docs/calibration-api.rst: -------------------------------------------------------------------------------- 1 | ----------- 2 | Calibration 3 | ----------- 4 | 5 | This module provides basic radio interferometry 6 | calibration utilities. Calibration is the 7 | process of estimating the :math:`2\times 2` 8 | Jones matrices which describe transformations 9 | of the signal as it propagates from source to 10 | observer. Currently, all utilities assume a 11 | discretised form of the radio interferometer 12 | measurement equation (RIME) as described in 13 | :ref:`rime-api-anchor`. 14 | 15 | Calibration is usually divided into three 16 | phases viz. 17 | 18 | * First generation calibration (1GC): using 19 | an external calibrator to infer the gains during 20 | the target observation. Sometimes also refered to 21 | as calibrator transfer 22 | 23 | * Second generation calibration (2GC): using 24 | a partially incomplete sky model to perform 25 | direction independent calibration. Also known 26 | as direction independent self-calibration. 27 | 28 | * Third generation calibration (3GC): using 29 | a partially incomplete sky model to perform 30 | direction dependent calibration. Also known 31 | as direction dependent self-calibration. 32 | 33 | On top of these three phases, there are usually 34 | three possible calibration scenarios. The first 35 | is when both the Jones terms and the visibilities 36 | are assumed to be diagonal. In this case the two 37 | correlations can be calibrated separately and it 38 | is refered to as :code:`diag-diag` calibration. 39 | The second case is when the Jones matrices are 40 | assumed to be diagonal but the visibility data 41 | are full :math:`2\times 2` matrices. This is 42 | refered to as :code:`diag` calibration. The final 43 | scenario is when both the full :math:`2\times 2` 44 | Jones matrices and the full :math:`2\times 2` 45 | visibilities are used for calibration. This is 46 | simply refered to as calibration. The specific 47 | scenario is determined from the shapes of the input 48 | gains and the input data. 49 | 50 | This module also provides a number of utilities which 51 | are useful for calibration. 52 | 53 | Utils 54 | +++++ 55 | 56 | Numpy 57 | ~~~~~ 58 | 59 | .. currentmodule:: africanus.calibration.utils 60 | 61 | .. autosummary:: 62 | corrupt_vis 63 | residual_vis 64 | correct_vis 65 | compute_and_corrupt_vis 66 | 67 | 68 | .. autofunction:: corrupt_vis 69 | .. autofunction:: residual_vis 70 | .. autofunction:: correct_vis 71 | .. autofunction:: compute_and_corrupt_vis 72 | 73 | Dask 74 | ~~~~ 75 | 76 | .. currentmodule:: africanus.calibration.utils.dask 77 | 78 | .. autosummary:: 79 | corrupt_vis 80 | residual_vis 81 | correct_vis 82 | compute_and_corrupt_vis 83 | 84 | 85 | .. autofunction:: corrupt_vis 86 | .. autofunction:: residual_vis 87 | .. autofunction:: correct_vis 88 | .. autofunction:: compute_and_corrupt_vis 89 | 90 | 91 | Phase only 92 | ++++++++++ 93 | 94 | Numpy 95 | ~~~~~ 96 | 97 | .. currentmodule:: africanus.calibration.phase_only 98 | 99 | .. autosummary:: 100 | compute_jhr 101 | compute_jhj 102 | compute_jhj_and_jhr 103 | gauss_newton 104 | 105 | 106 | .. autofunction:: compute_jhr 107 | .. autofunction:: compute_jhj 108 | .. autofunction:: compute_jhj_and_jhr 109 | .. autofunction:: gauss_newton 110 | 111 | 112 | Dask 113 | ~~~~~ 114 | 115 | .. currentmodule:: africanus.calibration.phase_only.dask 116 | 117 | .. autosummary:: 118 | compute_jhr 119 | compute_jhj 120 | 121 | 122 | .. autofunction:: compute_jhr 123 | .. autofunction:: compute_jhj 124 | -------------------------------------------------------------------------------- /africanus/dft/dask.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from africanus.dft.kernels import im_to_vis_docs, vis_to_im_docs 4 | from africanus.dft.kernels import im_to_vis as np_im_to_vis 5 | from africanus.dft.kernels import vis_to_im as np_vis_to_im 6 | 7 | from africanus.util.docs import doc_tuple_to_str 8 | from africanus.util.requirements import requires_optional 9 | 10 | import numpy as np 11 | 12 | try: 13 | import dask.array as da 14 | except ImportError as e: 15 | dask_import_error = e 16 | else: 17 | dask_import_error = None 18 | 19 | 20 | def _im_to_vis_wrapper(image, uvw, lm, frequency, convention, dtype_): 21 | return np_im_to_vis( 22 | image[0], uvw[0], lm[0][0], frequency, convention=convention, dtype=dtype_ 23 | ) 24 | 25 | 26 | @requires_optional("dask.array", dask_import_error) 27 | def im_to_vis(image, uvw, lm, frequency, convention="fourier", dtype=np.complex128): 28 | """Dask wrapper for im_to_vis function""" 29 | if lm.chunks[0][0] != lm.shape[0]: 30 | raise ValueError("lm chunks must match lm shape " "on first axis") 31 | if image.chunks[0][0] != image.shape[0]: 32 | raise ValueError("Image chunks must match image " "shape on first axis") 33 | if image.chunks[0][0] != lm.chunks[0][0]: 34 | raise ValueError("Image chunks and lm chunks must " "match on first axis") 35 | if image.chunks[1] != frequency.chunks[0]: 36 | raise ValueError("Image chunks must match frequency " "chunks on second axis") 37 | return da.core.blockwise( 38 | _im_to_vis_wrapper, 39 | ("row", "chan", "corr"), 40 | image, 41 | ("source", "chan", "corr"), 42 | uvw, 43 | ("row", "(u,v,w)"), 44 | lm, 45 | ("source", "(l,m)"), 46 | frequency, 47 | ("chan",), 48 | convention=convention, 49 | dtype=dtype, 50 | dtype_=dtype, 51 | ) 52 | 53 | 54 | def _vis_to_im_wrapper(vis, uvw, lm, frequency, flags, convention, dtype_): 55 | return np_vis_to_im( 56 | vis, uvw[0], lm[0], frequency, flags, convention=convention, dtype=dtype_ 57 | )[None, :] 58 | 59 | 60 | @requires_optional("dask.array", dask_import_error) 61 | def vis_to_im(vis, uvw, lm, frequency, flags, convention="fourier", dtype=np.float64): 62 | """Dask wrapper for vis_to_im function""" 63 | 64 | if vis.chunks[0] != uvw.chunks[0]: 65 | raise ValueError("Vis chunks and uvw chunks must " "match on first axis") 66 | if vis.chunks[1] != frequency.chunks[0]: 67 | raise ValueError("Vis chunks must match frequency " "chunks on second axis") 68 | if vis.chunks != flags.chunks: 69 | raise ValueError("Vis chunks must match flags " "chunks on all axes") 70 | 71 | ims = da.core.blockwise( 72 | _vis_to_im_wrapper, 73 | ("row", "source", "chan", "corr"), 74 | vis, 75 | ("row", "chan", "corr"), 76 | uvw, 77 | ("row", "(u,v,w)"), 78 | lm, 79 | ("source", "(l,m)"), 80 | frequency, 81 | ("chan",), 82 | flags, 83 | ("row", "chan", "corr"), 84 | adjust_chunks={"row": 1}, 85 | convention=convention, 86 | dtype=dtype, 87 | dtype_=dtype, 88 | ) 89 | 90 | return ims.sum(axis=0) 91 | 92 | 93 | im_to_vis.__doc__ = doc_tuple_to_str( 94 | im_to_vis_docs, [(":class:`numpy.ndarray`", ":class:`dask.array.Array`")] 95 | ) 96 | 97 | vis_to_im.__doc__ = doc_tuple_to_str( 98 | vis_to_im_docs, [(":class:`numpy.ndarray`", ":class:`dask.array.Array`")] 99 | ) 100 | -------------------------------------------------------------------------------- /africanus/rime/transform.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | 5 | import math 6 | 7 | import numpy as np 8 | 9 | from africanus.util.numba import jit 10 | 11 | 12 | @jit(nopython=True, nogil=True, cache=True) 13 | def _nb_transform_sources( 14 | lm, parallactic_angles, pointing_errors, antenna_scaling, frequency, coords 15 | ): 16 | """ 17 | numba implementation of 18 | :func:`~africanus.rime.transform_sources` 19 | """ 20 | _, nsrc, ntime, na, nchan = coords.shape 21 | 22 | for t in range(ntime): 23 | for a in range(na): 24 | pa_sin = math.sin(parallactic_angles[t, a]) 25 | pa_cos = math.cos(parallactic_angles[t, a]) 26 | 27 | for s in range(nsrc): 28 | l, m = lm[s] 29 | 30 | # Rotate source coordinate by parallactic angle 31 | l = l * pa_cos - m * pa_sin # noqa 32 | m = l * pa_sin + m * pa_cos 33 | 34 | # Add pointing errors 35 | l += pointing_errors[t, a, 0] # noqa 36 | m += pointing_errors[t, a, 1] 37 | 38 | # Scale by antenna scaling factors 39 | for c in range(nchan): 40 | coords[0, s, t, a, c] = l * antenna_scaling[a, c] 41 | coords[1, s, t, a, c] = m * antenna_scaling[a, c] 42 | coords[2, s, t, a, c] = frequency[c] 43 | 44 | return coords 45 | 46 | 47 | def transform_sources( 48 | lm, parallactic_angles, pointing_errors, antenna_scaling, frequency, dtype=None 49 | ): 50 | """ 51 | Creates beam sampling coordinates suitable for use 52 | in :func:`~africanus.rime.beam_cube_dde` by: 53 | 54 | 1. Rotating ``lm`` coordinates by the ``parallactic_angles`` 55 | 2. Adding ``pointing_errors`` 56 | 3. Scaling by ``antenna_scaling`` 57 | 58 | Parameters 59 | ---------- 60 | lm : :class:`numpy.ndarray` 61 | LM coordinates of shape :code:`(src,2)` in radians 62 | offset from the phase centre. 63 | parallactic_angles : :class:`numpy.ndarray` 64 | parallactic angles of shape :code:`(time, antenna)` 65 | in radians. 66 | pointing_errors : :class:`numpy.ndarray` 67 | LM pointing errors for each antenna at 68 | each timestep in radians. 69 | Has shape :code:`(time, antenna, 2)` 70 | antenna_scaling : :class:`numpy.ndarray` 71 | antenna scaling factor for each channel and 72 | each antenna. Has shape :code:`(antenna, chan)` 73 | frequency : :class:`numpy.ndarray` 74 | frequencies for each channel. Has shape :code:`(chan,)` 75 | dtype : :class:`numpy.dtype`, optional 76 | Numpy dtype of result array. Should be float32 or float64. 77 | Defaults to float64 78 | 79 | 80 | Returns 81 | ------- 82 | coords : :class:`numpy.ndarray` 83 | coordinates of shape :code:`(3, src, time, antenna, chan)` 84 | where each coordinate component represents **l**, **m** and 85 | **frequency**, respectively. 86 | """ 87 | 88 | ntime, na = parallactic_angles.shape 89 | nsrc = lm.shape[0] 90 | assert (ntime, na, 2) == pointing_errors.shape 91 | nchan = antenna_scaling.shape[1] 92 | assert nchan == frequency.shape[0] 93 | 94 | dtype = np.float64 if dtype is None else dtype 95 | coords = np.empty((3, nsrc, ntime, na, nchan), dtype=dtype) 96 | 97 | return _nb_transform_sources( 98 | lm, parallactic_angles, pointing_errors, antenna_scaling, frequency, coords 99 | ) 100 | -------------------------------------------------------------------------------- /africanus/model/wsclean/tests/test_spectral_model.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | from numpy.testing import assert_array_almost_equal 6 | import pytest 7 | 8 | from africanus.model.wsclean.file_model import load 9 | from africanus.model.wsclean.spec_model import ( 10 | ordinary_spectral_model, 11 | log_spectral_model, 12 | spectra, 13 | ) 14 | from africanus.model.wsclean.dask import spectra as dask_spectra 15 | 16 | 17 | @pytest.fixture 18 | def freq(): 19 | return np.linspace(0.856e9, 0.856e9 * 2, 16) 20 | 21 | 22 | @pytest.fixture 23 | def spectral_model_inputs(wsclean_model_file): 24 | sources = dict(load(wsclean_model_file)) 25 | 26 | I, spi, log_si, ref_freq = ( 27 | sources[n] 28 | for n in ("I", "SpectralIndex", "LogarithmicSI", "ReferenceFrequency") 29 | ) 30 | 31 | I = np.asarray(I) # noqa 32 | spi = np.asarray(spi) 33 | log_si = np.asarray(log_si) 34 | ref_freq = np.asarray(ref_freq) 35 | 36 | return I, spi, log_si, ref_freq 37 | 38 | 39 | def test_standard_spectral_model(spectral_model_inputs, freq): 40 | flux, spi, log_si, ref_freq = spectral_model_inputs 41 | 42 | mask = log_si == False # noqa 43 | assert mask.any() 44 | flux = flux[mask] 45 | spi = spi[mask] 46 | log_si = log_si[mask] 47 | ref_freq = ref_freq[mask] 48 | 49 | # Compute spectral model with numpy implementations 50 | np_model = ordinary_spectral_model(flux, spi, log_si, ref_freq, freq) 51 | model = spectra(flux, spi, log_si, ref_freq, freq) 52 | assert_array_almost_equal(np_model, model) 53 | 54 | 55 | def test_log_spectral_model(spectral_model_inputs, freq): 56 | flux, spi, log_si, ref_freq = spectral_model_inputs 57 | 58 | mask = log_si == True # noqa 59 | assert mask.any() 60 | flux = flux[mask] 61 | spi = spi[mask] 62 | log_si = log_si[mask] 63 | ref_freq = ref_freq[mask] 64 | 65 | # Compute spectral model with numpy implementations 66 | np_model = log_spectral_model(flux, spi, log_si, ref_freq, freq) 67 | model = spectra(flux, spi, log_si, ref_freq, freq) 68 | assert_array_almost_equal(np_model, model) 69 | 70 | 71 | def test_dask_spectral_model(spectral_model_inputs, freq): 72 | da = pytest.importorskip("dask.array") 73 | 74 | I, spi, log_si, ref_freq = spectral_model_inputs 75 | 76 | # Ensure positive flux for logarithmic polynomials 77 | I[log_si] = np.abs(I[log_si]) 78 | spi[log_si] = np.abs(spi[log_si]) 79 | 80 | # Compute spectral model with numpy implementations 81 | ordinary_spec_model = ordinary_spectral_model(I, spi, log_si, ref_freq, freq) 82 | log_spec_model = log_spectral_model(I, spi, log_si, ref_freq, freq) 83 | 84 | # Choose between ordinary and log spectral index 85 | # based on log_si array 86 | spec_model = np.where( 87 | log_si[:, None] == True, # noqa 88 | log_spec_model, 89 | ordinary_spec_model, 90 | ) 91 | 92 | # Create dask arrays 93 | src_chunks = (4, 4) 94 | spi_chunks = (2,) 95 | freq_chunks = (4, 4, 4, 4) 96 | 97 | I = da.from_array(I, chunks=(src_chunks,)) # noqa 98 | spi = da.from_array(spi, chunks=(src_chunks, spi_chunks)) 99 | log_si = da.from_array(log_si, chunks=(src_chunks,)) 100 | ref_freq = da.from_array(ref_freq, chunks=(src_chunks,)) 101 | freq = da.from_array(freq, chunks=(freq_chunks,)) 102 | 103 | # Compute spectra and compare 104 | model = dask_spectra(I, spi, log_si, ref_freq, freq) 105 | assert_array_almost_equal(model, spec_model) 106 | -------------------------------------------------------------------------------- /docs/model-api.rst: -------------------------------------------------------------------------------- 1 | --------- 2 | Sky Model 3 | --------- 4 | 5 | Functionality related to the Sky Model. 6 | 7 | Coherency Conversion 8 | -------------------- 9 | 10 | Utilities for converting back and forth between 11 | stokes parameters and correlations 12 | 13 | Numpy 14 | ~~~~~ 15 | 16 | .. currentmodule:: africanus.model.coherency 17 | 18 | .. autosummary:: 19 | convert 20 | 21 | .. autofunction:: convert 22 | 23 | Cuda 24 | ~~~~ 25 | 26 | .. currentmodule:: africanus.model.coherency.cuda 27 | 28 | .. autosummary:: 29 | convert 30 | 31 | .. autofunction:: convert 32 | 33 | Dask 34 | ~~~~ 35 | 36 | .. currentmodule:: africanus.model.coherency.dask 37 | 38 | .. autosummary:: 39 | convert 40 | 41 | .. autofunction:: convert 42 | 43 | 44 | Spectral Model 45 | -------------- 46 | 47 | Functionality for computing a Spectral Model. 48 | 49 | 50 | Numpy 51 | ~~~~~ 52 | 53 | .. currentmodule:: africanus.model.spectral 54 | 55 | .. autosummary:: 56 | spectral_model 57 | 58 | .. autofunction:: spectral_model 59 | 60 | Dask 61 | ~~~~ 62 | 63 | .. currentmodule:: africanus.model.spectral.dask 64 | 65 | .. autosummary:: 66 | spectral_model 67 | 68 | .. autofunction:: spectral_model 69 | 70 | 71 | Spectral Index 72 | -------------- 73 | 74 | Functionality related to the spectral index. 75 | 76 | For example, we may want 77 | to compute the spectral indices of 78 | components in a sky model defined by 79 | 80 | .. math:: 81 | 82 | I(\nu) = I(\nu_0) \left(\frac{\nu}{\nu_0}\right)^\alpha 83 | 84 | where :math:`\nu` are frequencies ay 85 | which we want to construct the intensity 86 | of a Stokes I image and the :math:`\nu_0` 87 | is the corresponding reference frequency. 88 | The spectral index :math:`\alpha` 89 | determines how quickly the intensity grows 90 | or decays as a function of frequency. 91 | Given a list of model image components 92 | (preferably with the residuals added back 93 | in) we can recover the corresponding 94 | spectral indices and reference intensities 95 | using the :func:`~africanus.model.spi.fit_spi_components` 96 | function. This will also return a lower bound 97 | on the associated uncertainties on these 98 | components. 99 | 100 | Numpy 101 | ~~~~~ 102 | 103 | .. currentmodule:: africanus.model.spi 104 | 105 | .. autosummary:: 106 | fit_spi_components 107 | 108 | .. autofunction:: fit_spi_components 109 | 110 | 111 | Dask 112 | ~~~~~ 113 | 114 | .. currentmodule:: africanus.model.spi.dask 115 | 116 | .. autosummary:: 117 | fit_spi_components 118 | 119 | .. autofunction:: fit_spi_components 120 | 121 | 122 | Source Morphology 123 | ----------------- 124 | 125 | Shape functions for different Source Morphologies 126 | 127 | Numpy 128 | ~~~~~ 129 | 130 | .. currentmodule:: africanus.model.shape 131 | 132 | .. autosummary:: 133 | gaussian 134 | 135 | .. autofunction:: gaussian 136 | 137 | Dask 138 | ~~~~ 139 | 140 | .. currentmodule:: africanus.model.shape.dask 141 | 142 | .. autosummary:: 143 | gaussian 144 | 145 | .. autofunction:: gaussian 146 | 147 | 148 | WSClean Spectral Model 149 | ---------------------- 150 | 151 | Utilities for creating a spectral model from a wsclean component file. 152 | 153 | Numpy 154 | ~~~~~ 155 | 156 | .. currentmodule:: africanus.model.wsclean 157 | 158 | .. autosummary:: 159 | load 160 | spectra 161 | 162 | .. autofunction:: load 163 | .. autofunction:: spectra 164 | 165 | Dask 166 | ~~~~ 167 | 168 | .. currentmodule:: africanus.model.wsclean.dask 169 | 170 | .. autosummary:: 171 | spectra 172 | 173 | .. autofunction:: spectra 174 | -------------------------------------------------------------------------------- /africanus/rime/feeds.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from functools import reduce 5 | from operator import mul 6 | 7 | import numpy as np 8 | 9 | from africanus.util.docs import DocstringTemplate 10 | from africanus.util.numba import jit 11 | 12 | 13 | @jit(nopython=True, nogil=True, cache=True) 14 | def _nb_feed_rotation(parallactic_angles, feed_type, feed_rotation): 15 | shape = parallactic_angles.shape 16 | parangles = parallactic_angles.flat 17 | 18 | # Linear feeds 19 | if feed_type == 0: 20 | for i, pa in enumerate(parangles): 21 | pa_cos = np.cos(pa) 22 | pa_sin = np.sin(pa) 23 | 24 | feed_rotation.real[i, 0, 0] = pa_cos 25 | feed_rotation.imag[i, 0, 0] = 0.0 26 | feed_rotation.real[i, 0, 1] = pa_sin 27 | feed_rotation.imag[i, 0, 1] = 0.0 28 | feed_rotation.real[i, 1, 0] = -pa_sin 29 | feed_rotation.imag[i, 1, 0] = 0.0 30 | feed_rotation.real[i, 1, 1] = pa_cos 31 | feed_rotation.imag[i, 1, 1] = 0.0 32 | 33 | # Circular feeds 34 | elif feed_type == 1: 35 | for i, pa in enumerate(parangles): 36 | pa_cos = np.cos(pa) 37 | pa_sin = np.sin(pa) 38 | 39 | feed_rotation.real[i, 0, 0] = pa_cos 40 | feed_rotation.imag[i, 0, 0] = -pa_sin 41 | feed_rotation[i, 0, 1] = 0.0 + 0.0 * 1j 42 | feed_rotation[i, 1, 0] = 0.0 + 0.0 * 1j 43 | feed_rotation.real[i, 1, 1] = pa_cos 44 | feed_rotation.imag[i, 1, 1] = pa_sin 45 | else: 46 | raise ValueError("Invalid feed_type") 47 | 48 | return feed_rotation.reshape(shape + (2, 2)) 49 | 50 | 51 | def feed_rotation(parallactic_angles, feed_type="linear"): 52 | if feed_type == "linear": 53 | poltype = 0 54 | elif feed_type == "circular": 55 | poltype = 1 56 | else: 57 | raise ValueError("Invalid feed_type '%s'" % feed_type) 58 | 59 | if parallactic_angles.dtype == np.float32: 60 | dtype = np.complex64 61 | elif parallactic_angles.dtype == np.float64: 62 | dtype = np.complex128 63 | else: 64 | raise ValueError( 65 | "parallactic_angles has " 66 | "none-floating point type %s" % parallactic_angles.dtype 67 | ) 68 | 69 | # Create result array with flattened parangles 70 | shape = (reduce(mul, parallactic_angles.shape),) + (2, 2) 71 | result = np.empty(shape, dtype=dtype) 72 | 73 | return _nb_feed_rotation(parallactic_angles, poltype, result) 74 | 75 | 76 | FEED_ROTATION_DOCS = DocstringTemplate( 77 | r""" 78 | Computes the 2x2 feed rotation (L) matrix 79 | from the ``parallactic_angles``. 80 | 81 | .. math:: 82 | 83 | \textrm{linear} 84 | \begin{bmatrix} 85 | cos(pa) & sin(pa) \\ 86 | -sin(pa) & cos(pa) 87 | \end{bmatrix} 88 | \qquad 89 | \textrm{circular} 90 | \begin{bmatrix} 91 | e^{-i pa} & 0 \\ 92 | 0 & e^{i pa} 93 | \end{bmatrix} 94 | 95 | Parameters 96 | ---------- 97 | parallactic_angles : $(array_type) 98 | floating point parallactic angles. Of shape 99 | :code:`(pa0, pa1, ..., pan)`. 100 | feed_type : {'linear', 'circular'} 101 | The type of feed 102 | 103 | Returns 104 | ------- 105 | feed_matrix : $(array_type) 106 | Feed rotation matrix of shape :code:`(pa0, pa1,...,pan,2,2)` 107 | """ 108 | ) 109 | 110 | try: 111 | feed_rotation.__doc__ = FEED_ROTATION_DOCS.substitute( 112 | array_type=":class:`numpy.ndarray`" 113 | ) 114 | except AttributeError: 115 | pass 116 | -------------------------------------------------------------------------------- /africanus/model/wsclean/tests/test_wsclean_model_file.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | from numpy.testing import assert_array_almost_equal, assert_array_equal 6 | 7 | from africanus.model.wsclean.file_model import load, arcsec2rad 8 | 9 | 10 | def test_wsclean_model_file(wsclean_model_file): 11 | sources = dict(load(wsclean_model_file)) 12 | 13 | (name, stype, ra, dec, I, spi, log_si, ref_freq, major, minor, orientation) = ( 14 | sources[n] 15 | for n in ( 16 | "Name", 17 | "Type", 18 | "Ra", 19 | "Dec", 20 | "I", 21 | "SpectralIndex", 22 | "LogarithmicSI", 23 | "ReferenceFrequency", 24 | "MajorAxis", 25 | "MinorAxis", 26 | "Orientation", 27 | ) 28 | ) 29 | 30 | # Seven sources 31 | assert len(I) == len(spi) == len(log_si) == len(ref_freq) == 8 32 | 33 | # Name and type read correctly 34 | assert name[0] == "s0c0" and stype[0] == "POINT" 35 | 36 | # Check ra conversion for line 0 file entry (-float, float, float) 37 | hours, mins, secs = (-8.0, 28.0, 5.152) 38 | expected_ra0 = ( 39 | -2.0 40 | * np.pi 41 | * ((-hours / 24.0) + (mins / (24.0 * 60.0)) + (secs / (24.0 * 60.0 * 60.0))) 42 | ) 43 | 44 | assert ra[0] == expected_ra0 45 | 46 | # Check dec conversion for line 0 file entry 47 | degs, mins, secs = (39.0, 35.0, 8.511) 48 | expected_dec0 = ( 49 | 2.0 50 | * np.pi 51 | * ((degs / 360.0) + (mins / (360.0 * 60.0)) + (secs / (360.0 * 60.0 * 60.0))) 52 | ) 53 | 54 | assert dec[0] == expected_dec0 55 | 56 | # SPI read correctly 57 | assert_array_equal(spi[0], [-0.00695379313004673, -0.0849693907803257]) 58 | 59 | # LogrithmicSI read correctly 60 | assert log_si[0] is True 61 | 62 | # Check ra conversion for line 2 file entry (int, not float, seconds) 63 | hours, mins, secs = (8, 18, 44) 64 | expected_ra2 = ( 65 | 2.0 66 | * np.pi 67 | * ((hours / 24.0) + (mins / (24.0 * 60.0)) + (secs / (24.0 * 60.0 * 60.0))) 68 | ) 69 | 70 | assert ra[2] == expected_ra2 71 | 72 | # Check dec conversion for line 2 file entry (int, not float, seconds) 73 | degs, mins, secs = (39, 38, 37) 74 | expected_dec2 = ( 75 | 2.0 76 | * np.pi 77 | * ((degs / 360.0) + (mins / (360.0 * 60.0)) + (secs / (360.0 * 60.0 * 60.0))) 78 | ) 79 | 80 | assert dec[2] == expected_dec2 81 | 82 | assert log_si[2] is False 83 | 84 | assert I[2] == 0.000233552686127518 85 | 86 | # Check dec conversion for line 4 file entry (+int, not float, seconds) 87 | degs, mins, secs = (+41, 47, 17.131) 88 | expected_dec4 = ( 89 | 2.0 90 | * np.pi 91 | * ((degs / 360.0) + (mins / (360.0 * 60.0)) + (secs / (360.0 * 60.0 * 60.0))) 92 | ) 93 | 94 | assert dec[4] == expected_dec4 95 | 96 | # Missing reference frequency set in the last 97 | assert ref_freq[6] == ref_freq[0] 98 | 99 | # Last name and type correct 100 | assert name[6] == "s1c2" and stype[6] == "GAUSSIAN" 101 | 102 | # https://www.convertunits.com/from/arcsecond/to/radian 103 | assert_array_almost_equal(major[6], arcsec2rad(83.6144111272856)) 104 | assert_array_almost_equal(minor[6], arcsec2rad(83.6144111272856)) 105 | assert_array_almost_equal(orientation[6], np.deg2rad(45)) 106 | 107 | assert I[6] == 0.000660490865128381 108 | 109 | assert I[7] == 0 110 | assert_array_equal(spi[7], [0.0, 0.0]) 111 | -------------------------------------------------------------------------------- /africanus/coordinates/dask.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | try: 5 | import dask.array as da 6 | except ImportError as e: 7 | dask_import_error = e 8 | else: 9 | dask_import_error = None 10 | 11 | 12 | from africanus.util.requirements import requires_optional 13 | from africanus.coordinates.coordinates import ( 14 | radec_to_lmn as np_radec_to_lmn, 15 | radec_to_lm as np_radec_to_lm, 16 | lmn_to_radec as np_lmn_to_radec, 17 | lm_to_radec as np_lm_to_radec, 18 | RADEC_TO_LMN_DOCS, 19 | LMN_TO_RADEC_DOCS, 20 | ) 21 | 22 | 23 | def _radec_to_lmn(radec, phase_centre): 24 | return np_radec_to_lmn(radec[0], phase_centre[0] if phase_centre else None) 25 | 26 | 27 | @requires_optional("dask.array", dask_import_error) 28 | def radec_to_lmn(radec, phase_centre=None): 29 | phase_centre_dims = ("radec",) if phase_centre is not None else None 30 | 31 | return da.core.blockwise( 32 | _radec_to_lmn, 33 | ("source", "lmn"), 34 | radec, 35 | ("source", "radec"), 36 | phase_centre, 37 | phase_centre_dims, 38 | new_axes={"lmn": 3}, 39 | dtype=radec.dtype, 40 | ) 41 | 42 | 43 | def _lmn_to_radec(lmn, phase_centre): 44 | return np_lmn_to_radec(lmn[0], phase_centre) 45 | 46 | 47 | @requires_optional("dask.array", dask_import_error) 48 | def lmn_to_radec(lmn, phase_centre=None): 49 | phase_centre_dims = ("radec",) if phase_centre is not None else None 50 | 51 | return da.core.blockwise( 52 | _lmn_to_radec, 53 | ("source", "radec"), 54 | lmn, 55 | ("source", "lmn"), 56 | phase_centre, 57 | phase_centre_dims, 58 | new_axes={"radec": 2}, 59 | dtype=lmn.dtype, 60 | ) 61 | 62 | 63 | def _radec_to_lm(radec, phase_centre): 64 | return np_radec_to_lm(radec[0], phase_centre[0] if phase_centre else None) 65 | 66 | 67 | @requires_optional("dask.array", dask_import_error) 68 | def radec_to_lm(radec, phase_centre=None): 69 | phase_centre_dims = ("radec",) if phase_centre is not None else None 70 | 71 | return da.core.blockwise( 72 | _radec_to_lm, 73 | ("source", "lm"), 74 | radec, 75 | ("source", "radec"), 76 | phase_centre, 77 | phase_centre_dims, 78 | new_axes={"lm": 2}, 79 | dtype=radec.dtype, 80 | ) 81 | 82 | 83 | def _lm_to_radec(lm, phase_centre): 84 | return np_lm_to_radec(lm[0], phase_centre) 85 | 86 | 87 | @requires_optional("dask.array", dask_import_error) 88 | def lm_to_radec(lm, phase_centre=None): 89 | phase_centre_dims = ("radec",) if phase_centre is not None else None 90 | 91 | return da.core.blockwise( 92 | _lm_to_radec, 93 | ("source", "radec"), 94 | lm, 95 | ("source", "lm"), 96 | phase_centre, 97 | phase_centre_dims, 98 | new_axes={"radec": 2}, 99 | dtype=lm.dtype, 100 | ) 101 | 102 | 103 | try: 104 | radec_to_lmn.__doc__ = RADEC_TO_LMN_DOCS.substitute( 105 | lm_components="3", array_type=":class:`dask.array.Array`" 106 | ) 107 | radec_to_lm.__doc__ = RADEC_TO_LMN_DOCS.substitute( 108 | lm_components="2", array_type=":class:`dask.array.Array`" 109 | ) 110 | lmn_to_radec.__doc__ = LMN_TO_RADEC_DOCS.substitute( 111 | lm_name="lmn", lm_components="3", array_type=":class:`dask.array.Array`" 112 | ) 113 | lm_to_radec.__doc__ = LMN_TO_RADEC_DOCS.substitute( 114 | lm_name="lm", lm_components="2", array_type=":class:`dask.array.Array`" 115 | ) 116 | 117 | 118 | except AttributeError: 119 | pass 120 | -------------------------------------------------------------------------------- /africanus/averaging/support.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | import numba 6 | 7 | from africanus.util.numba import JIT_OPTIONS, overload, njit 8 | 9 | 10 | @njit(nogil=True, cache=True) 11 | def _unique_internal(data): 12 | if len(data.shape) != 1: 13 | raise ValueError("_unique_internal currently " "only supports 1D arrays") 14 | 15 | # Handle the empty array case 16 | if data.shape[0] == 0: 17 | return ( 18 | data, 19 | np.empty((0,), dtype=np.intp), 20 | np.empty((0,), dtype=np.intp), 21 | np.empty((0,), dtype=np.intp), 22 | ) 23 | 24 | # See numpy's unique1d 25 | perm = np.argsort(data, kind="mergesort") 26 | 27 | # Combine these arrays to save on allocations? 28 | aux = np.empty_like(data) 29 | mask = np.empty(aux.shape, dtype=np.bool_) 30 | inv_idx = np.empty(mask.shape, dtype=np.intp) 31 | 32 | # Hard code first iteration 33 | p = perm[0] 34 | aux[0] = data[p] 35 | mask[0] = True 36 | cumsum = 1 37 | inv_idx[p] = cumsum - 1 38 | counts = [np.intp(0)] 39 | 40 | for i in range(1, aux.shape[0]): 41 | p = perm[i] 42 | aux[i] = data[p] 43 | d = aux[i] != aux[i - 1] 44 | mask[i] = d 45 | cumsum += d 46 | inv_idx[p] = cumsum - 1 47 | 48 | if d: 49 | counts.append(np.intp(i)) 50 | 51 | counts.append(aux.shape[0]) 52 | 53 | # (uniques, indices, inverse index, counts) 54 | return aux[mask], perm[mask], inv_idx, np.diff(np.array(counts)) 55 | 56 | 57 | @njit(**JIT_OPTIONS) 58 | def unique_time(time): 59 | return unique_time_impl(time) 60 | 61 | 62 | def unique_time_impl(time): 63 | return NotImplementedError 64 | 65 | 66 | @overload(unique_time_impl, jit_options=JIT_OPTIONS) 67 | def nb_unique_time(time): 68 | """Return unique time, inverse index and counts""" 69 | if time.dtype not in (numba.float32, numba.float64): 70 | raise ValueError("time must be floating point but is %s" % time.dtype) 71 | 72 | def impl(time): 73 | return _unique_internal(time) 74 | 75 | return impl 76 | 77 | 78 | @njit(**JIT_OPTIONS) 79 | def unique_baselines(ant1, ant2): 80 | return unique_baselines_impl(ant1, ant2) 81 | 82 | 83 | def unique_baselines_impl(ant1, ant2): 84 | return NotImplementedError 85 | 86 | 87 | @overload(unique_baselines_impl, jit_options=JIT_OPTIONS) 88 | def nb_unique_baselines(ant1, ant2): 89 | """Return unique baselines, inverse index and counts""" 90 | if not ant1.dtype == numba.int32 or not ant2.dtype == numba.int32: 91 | # Need these to be int32 for the bl_32bit.view(np.int64) trick 92 | raise ValueError( 93 | "ant1 and ant2 must be np.int32 " 94 | "but received %s and %s" % (ant1.dtype, ant2.dtype) 95 | ) 96 | 97 | def impl(ant1, ant2): 98 | # Trickery, stack the two int32 antenna pairs in an array 99 | # and cast to int64 100 | bl_32bit = np.empty((ant1.shape[0], 2), dtype=np.int32) 101 | 102 | # Copy data 103 | for r in range(ant1.shape[0]): 104 | bl_32bit[r, 0] = ant1[r] 105 | bl_32bit[r, 1] = ant2[r] 106 | 107 | # Cast to int64 for the unique operation 108 | bl = bl_32bit.view(np.int64).reshape(ant1.shape[0]) 109 | 110 | ret, idx, inv, counts = _unique_internal(bl) 111 | 112 | # Recast to int32 and reshape 113 | ubl = ret.view(np.int32).reshape(ret.shape[0], 2) 114 | 115 | return ubl, idx, inv, counts 116 | 117 | return impl 118 | -------------------------------------------------------------------------------- /africanus/model/shape/gaussian_shape.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | 6 | from africanus.util.docs import DocstringTemplate 7 | from africanus.util.numba import njit, overload, JIT_OPTIONS 8 | from africanus.constants import c as lightspeed 9 | 10 | 11 | @njit(**JIT_OPTIONS) 12 | def gaussian(uvw, frequency, shape_params): 13 | return gaussian_impl(uvw, frequency, shape_params) 14 | 15 | 16 | def gaussian_impl(uvw, frequency, shape_params): 17 | raise NotImplementedError 18 | 19 | 20 | @overload(gaussian_impl, jit_options=JIT_OPTIONS) 21 | def nb_gaussian(uvw, frequency, shape_params): 22 | # https://en.wikipedia.org/wiki/Full_width_at_half_maximum 23 | fwhm = 2.0 * np.sqrt(2.0 * np.log(2.0)) 24 | fwhminv = 1.0 / fwhm 25 | gauss_scale = fwhminv * np.sqrt(2.0) * np.pi / lightspeed 26 | 27 | dtype = np.result_type( 28 | *(np.dtype(a.dtype.name) for a in (uvw, frequency, shape_params)) 29 | ) 30 | 31 | def impl(uvw, frequency, shape_params): 32 | nsrc = shape_params.shape[0] 33 | nrow = uvw.shape[0] 34 | nchan = frequency.shape[0] 35 | 36 | shape = np.empty((nsrc, nrow, nchan), dtype=dtype) 37 | scaled_freq = np.empty_like(frequency) 38 | 39 | # Scale each frequency 40 | for f in range(frequency.shape[0]): 41 | scaled_freq[f] = frequency[f] * gauss_scale 42 | 43 | for s in range(shape_params.shape[0]): 44 | emaj, emin, angle = shape_params[s] 45 | 46 | # Convert to l-projection, m-projection, ratio 47 | el = emaj * np.sin(angle) 48 | em = emaj * np.cos(angle) 49 | er = emin / (1.0 if emaj == 0.0 else emaj) 50 | 51 | for r in range(uvw.shape[0]): 52 | u, v, w = uvw[r] 53 | 54 | u1 = (u * em - v * el) * er 55 | v1 = u * el + v * em 56 | 57 | for f in range(scaled_freq.shape[0]): 58 | fu1 = u1 * scaled_freq[f] 59 | fv1 = v1 * scaled_freq[f] 60 | 61 | shape[s, r, f] = np.exp(-(fu1 * fu1 + fv1 * fv1)) 62 | 63 | return shape 64 | 65 | return impl 66 | 67 | 68 | GAUSSIAN_DOCS = DocstringTemplate( 69 | r""" 70 | Computes the Gaussian Shape Function. 71 | 72 | .. math:: 73 | 74 | & \nu^\prime = \frac{\nu \pi}{2 c \sqrt{\log{2}}} \\ 75 | & r = \frac{e_{min}}{e_{maj}} \\ 76 | & u_{1} = (u \, e_{maj} \, cos(\alpha) - v \, e_{maj} \, sin(\alpha)) 77 | r \nu^\prime \\ 78 | & v_{1} = (u \, e_{maj} \, sin(\alpha) - v \, e_{maj} \, cos(\alpha)) 79 | \nu^\prime \\ 80 | & \textrm{shape} = e^{(-(u_{1}^2) - (v_{1}^2))} 81 | 82 | where: 83 | 84 | - :math:`u` and :math:`v` are the UV coordinates and 85 | :math:`\nu` the frequency. 86 | - :math:`e_{maj}` and :math:`e_{min}` are the major and minor axes 87 | and :math:`\alpha` the position angle. 88 | 89 | Parameters 90 | ---------- 91 | uvw : $(array_type) 92 | UVW coordinates of shape :code:`(row, 3)` 93 | frequency : $(array_type) 94 | frequencies of shape :code:`(chan,)` 95 | shape_param : $(array_type) 96 | Gaussian Shape Parameters of shape :code:`(source, 3)` 97 | where the second dimension contains the 98 | `(emajor, eminor, angle)` parameters describing 99 | the shape of the Gaussian 100 | 101 | Returns 102 | ------- 103 | gauss_shape : $(array_type) 104 | Shape parameters of shape :code:`(source, row, chan)` 105 | """ 106 | ) 107 | 108 | try: 109 | gaussian.__doc__ = GAUSSIAN_DOCS.substitute(array_type=":class:`numpy.ndarray`") 110 | except KeyError: 111 | pass 112 | -------------------------------------------------------------------------------- /africanus/util/code.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from functools import wraps 5 | 6 | try: 7 | from dask.utils import SerializableLock as Lock 8 | except ImportError: 9 | from threading import Lock 10 | 11 | 12 | class SingletonMixin(object): 13 | __singleton_lock = Lock() 14 | __singleton_instance = None 15 | 16 | @classmethod 17 | def instance(cls): 18 | if not cls.__singleton_instance: 19 | with cls.__singleton_lock: 20 | if not cls.__singleton_instance: 21 | cls.__singleton_instance = cls() 22 | 23 | return cls.__singleton_instance 24 | 25 | 26 | def format_code(code): 27 | """ 28 | Formats some code with line numbers 29 | 30 | Parameters 31 | ---------- 32 | code : str 33 | Code 34 | 35 | Returns 36 | ------- 37 | str 38 | Code prefixed with line numbers 39 | """ 40 | lines = [""] 41 | lines.extend(["%-5d %s" % (i, l) for i, l in enumerate(code.split("\n"), 1)]) 42 | return "\n".join(lines) 43 | 44 | 45 | class memoize_on_key(object): 46 | """ 47 | Memoize based on a key function supplied by the user. 48 | The key function should return a custom key 49 | for memoizing the decorated function, based on the arguments 50 | passed to it. 51 | 52 | In the following example, the arguments required to generate 53 | the `_generate_phase_delay_kernel` function are the types of 54 | the `lm`, `uvw` and `frequency` arrays, as well as the number 55 | of correlations, `ncorr`. 56 | 57 | The supplied ``key_fn`` produces a unique key based on these types 58 | and the number of correlations, which is used to cache the 59 | generated function. 60 | 61 | .. code-block:: python 62 | 63 | def key_fn(lm, uvw, frequency, ncorrs=4): 64 | ''' 65 | Produce a unique key for the arguments of 66 | _generate_phase_delay_kernel 67 | ''' 68 | return (lm.dtype, uvw.dtype, frequency.dtype, ncorrs) 69 | 70 | _code_template = jinja2.Template(''' 71 | #define ncorrs {{ncorrs}} 72 | 73 | __global__ void phase_delay( 74 | const {{lm_type}} * lm, 75 | const {{uvw_type}} * uvw, 76 | const {{freq_type}} * frequency, 77 | {{out_type}} * out) 78 | { 79 | ... 80 | } 81 | ''') 82 | 83 | _type_map = { 84 | np.float32: 'float', 85 | np.float64: 'double' 86 | } 87 | 88 | @memoize_on_key(key_fn) 89 | def _generate_phase_delay_kernel(lm, uvw, frequency, ncorrs=4): 90 | ''' Generate the phase delay kernel ''' 91 | out_dtype = np.result_type(lm.dtype, uvw.dtype, frequency.dtype) 92 | code = _code_template.render(lm_type=_type_map[lm.dtype], 93 | uvw_type=_type_map[uvw.dtype], 94 | freq_type=_type_map[frequency.dtype], 95 | ncorrs=ncorrs) 96 | return cp.RawKernel(code, "phase_delay") 97 | """ 98 | 99 | def __init__(self, key_fn): 100 | self._key_fn = key_fn 101 | self._lock = Lock() 102 | self._cache = {} 103 | 104 | def __call__(self, fn): 105 | @wraps(fn) 106 | def wrapper(*args, **kwargs): 107 | key = self._key_fn(*args, **kwargs) 108 | 109 | with self._lock: 110 | try: 111 | return self._cache[key] 112 | except KeyError: 113 | self._cache[key] = entry = fn(*args, **kwargs) 114 | return entry 115 | 116 | return wrapper 117 | -------------------------------------------------------------------------------- /africanus/model/spi/tests/test_component_spi.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | """Tests for `codex-africanus` package.""" 4 | 5 | import numpy as np 6 | import pytest 7 | 8 | 9 | def test_fit_spi_components_vs_scipy(): 10 | """ 11 | Here we just test the per component spi fitter against 12 | a looped version of scipy's curve_fit 13 | :return: 14 | """ 15 | from africanus.model.spi import fit_spi_components 16 | 17 | curve_fit = pytest.importorskip("scipy.optimize").curve_fit 18 | 19 | np.random.seed(123) 20 | 21 | ncomps = 250 22 | alphas = -0.7 + 0.25 * np.random.randn(ncomps, 1) 23 | i0s = 5.0 + np.random.randn(ncomps, 1) 24 | nfreqs = 100 25 | freqs = np.linspace(0.5, 1.5, nfreqs).reshape(1, nfreqs) 26 | freq0 = 0.7 27 | beams = np.zeros((ncomps, nfreqs)) 28 | for i in range(ncomps): 29 | beams[i, :] = np.sinc(freqs - freqs[0]) 30 | model = beams * i0s * (freqs / freq0) ** alphas 31 | sigma = np.abs(0.25 + 0.1 * np.random.randn(nfreqs)) 32 | data = model + sigma[None, :] * np.random.randn(ncomps, nfreqs) 33 | 34 | weights = 1.0 / sigma**2 35 | alpha1, alphavar1, I01, I0var1 = fit_spi_components( 36 | data, weights, freqs.squeeze(), freq0, tol=1e-8 37 | ) 38 | 39 | def spi_func(nu, I0, alpha, beam=1.0): 40 | return beam * I0 * nu**alpha 41 | 42 | I02 = np.zeros(ncomps) 43 | I0var2 = np.zeros(ncomps) 44 | alpha2 = np.zeros(ncomps) 45 | alphavar2 = np.zeros(ncomps) 46 | 47 | for i in range(ncomps): 48 | 49 | def fit_func(nu, I0, alpha): 50 | return spi_func(nu, I0, alpha, beam=beams[i]) 51 | 52 | popt, pcov = curve_fit( 53 | fit_func, 54 | (freqs / freq0).squeeze(), 55 | data[i, :], 56 | sigma=np.diag(sigma**2), 57 | p0=np.array([1.0, -0.7]), 58 | absolute_sigma=False, 59 | ) 60 | I02[i] = popt[0] 61 | I0var2[i] = pcov[0, 0] 62 | alpha2[i] = popt[1] 63 | alphavar2[i] = pcov[1, 1] 64 | 65 | np.testing.assert_array_almost_equal(alpha1, alpha2, decimal=5) 66 | np.testing.assert_array_almost_equal(alphavar1, alphavar2, decimal=5) 67 | np.testing.assert_array_almost_equal(I01, I02, decimal=5) 68 | np.testing.assert_array_almost_equal(I0var1, I0var2, decimal=5) 69 | 70 | 71 | def test_dask_fit_spi_components_vs_np(): 72 | from africanus.model.spi import fit_spi_components as np_fit_spi 73 | from africanus.model.spi.dask import fit_spi_components 74 | 75 | da = pytest.importorskip("dask.array") 76 | 77 | np.random.seed(123) 78 | 79 | ncomps = 800 80 | alphas = -0.7 + 0.25 * np.random.randn(ncomps, 1) 81 | i0s = 5.0 + np.random.randn(ncomps, 1) 82 | nfreqs = 1000 83 | freqs = np.linspace(0.5, 1.5, nfreqs).reshape(1, nfreqs) 84 | freq0 = 0.7 85 | model = i0s * (freqs / freq0) ** alphas 86 | sigma = np.abs(0.25 + 0.1 * np.random.randn(nfreqs)) 87 | data = model + sigma[None, :] * np.random.randn(ncomps, nfreqs) 88 | 89 | weights = 1.0 / sigma**2 90 | freqs = freqs.squeeze() 91 | alpha1, alphavar1, I01, I0var1 = np_fit_spi(data, weights, freqs, freq0) 92 | 93 | # now for the dask version 94 | data_dask = da.from_array(data, chunks=(100, nfreqs)) 95 | weights_dask = da.from_array(weights, chunks=(nfreqs)) 96 | freqs_dask = da.from_array(freqs, chunks=(nfreqs)) 97 | 98 | alpha2, alphavar2, I02, I0var2 = fit_spi_components( 99 | data_dask, weights_dask, freqs_dask, freq0 100 | ).compute() 101 | 102 | np.testing.assert_array_almost_equal(alpha1, alpha2, decimal=6) 103 | np.testing.assert_array_almost_equal(alphavar1, alphavar2, decimal=6) 104 | np.testing.assert_array_almost_equal(I01, I02, decimal=6) 105 | np.testing.assert_array_almost_equal(I0var1, I0var2, decimal=6) 106 | -------------------------------------------------------------------------------- /africanus/rime/tests/test_rime.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | 4 | """Tests for `codex-africanus` package.""" 5 | 6 | import numpy as np 7 | 8 | import pytest 9 | 10 | 11 | def rf(*a, **kw): 12 | return np.random.random(*a, **kw) 13 | 14 | 15 | def rc(*a, **kw): 16 | return rf(*a, **kw) + 1j * rf(*a, **kw) 17 | 18 | 19 | @pytest.mark.parametrize("convention, sign", [("fourier", 1), ("casa", -1)]) 20 | def test_phase_delay(convention, sign): 21 | from africanus.rime import phase_delay 22 | 23 | uvw = np.random.random(size=(100, 3)) 24 | lm = np.random.random(size=(10, 2)) 25 | frequency = np.linspace(0.856e9, 0.856e9 * 2, 64, endpoint=True) 26 | 27 | from africanus.constants import minus_two_pi_over_c 28 | 29 | # Test complex phase at a particular index in the output 30 | uvw_i, lm_i, freq_i = 2, 3, 5 31 | 32 | u, v, w = [1, 2, 3] 33 | l, m = [0.1, 0.2] 34 | freq = 0.856e9 35 | 36 | # Set up values in the input 37 | uvw[uvw_i] = [u, v, w] 38 | lm[lm_i] = [l, m] 39 | frequency[freq_i] = freq 40 | 41 | # Compute complex phase 42 | complex_phase = phase_delay(lm, uvw, frequency, convention=convention) 43 | 44 | # Test singular value vs a point in the output 45 | n = np.sqrt(1.0 - l**2 - m**2) - 1.0 46 | phase = sign * minus_two_pi_over_c * (u * l + v * m + w * n) * freq 47 | assert np.all(np.exp(1j * phase) == complex_phase[lm_i, uvw_i, freq_i]) 48 | 49 | 50 | def test_feed_rotation(): 51 | import numpy as np 52 | from africanus.rime import feed_rotation 53 | 54 | parangles = np.random.random((10, 5)) 55 | pa_sin = np.sin(parangles) 56 | pa_cos = np.cos(parangles) 57 | 58 | fr = feed_rotation(parangles, feed_type="linear") 59 | np_expr = np.stack([pa_cos, pa_sin, -pa_sin, pa_cos], axis=2) 60 | assert np.allclose(fr, np_expr.reshape(10, 5, 2, 2)) 61 | 62 | fr = feed_rotation(parangles, feed_type="circular") 63 | zeros = np.zeros_like(pa_sin) 64 | np_expr = np.stack( 65 | [pa_cos - 1j * pa_sin, zeros, zeros, pa_cos + 1j * pa_sin], axis=2 66 | ) 67 | assert np.allclose(fr, np_expr.reshape(10, 5, 2, 2)) 68 | 69 | 70 | @pytest.mark.parametrize("convention, sign", [("fourier", 1), ("casa", -1)]) 71 | def test_dask_phase_delay(convention, sign): 72 | da = pytest.importorskip("dask.array") 73 | from africanus.rime import phase_delay as np_phase_delay 74 | from africanus.rime.dask import phase_delay as dask_phase_delay 75 | 76 | # So that 1 > 1 - l**2 - m**2 >= 0 77 | lm = np.random.random(size=(10, 2)) * 0.01 78 | uvw = np.random.random(size=(100, 3)) 79 | frequency = np.linspace(0.856e9, 0.856e9 * 2, 64, endpoint=True) 80 | 81 | dask_lm = da.from_array(lm, chunks=(5, 2)) 82 | dask_uvw = da.from_array(uvw, chunks=(25, 3)) 83 | dask_frequency = da.from_array(frequency, chunks=16) 84 | 85 | dask_phase = dask_phase_delay( 86 | dask_lm, dask_uvw, dask_frequency, convention=convention 87 | ) 88 | np_phase = np_phase_delay(lm, uvw, frequency, convention=convention) 89 | 90 | # Should agree completely 91 | assert np.all(np_phase == dask_phase.compute()) 92 | 93 | 94 | def test_dask_feed_rotation(): 95 | da = pytest.importorskip("dask.array") 96 | import numpy as np 97 | from africanus.rime import feed_rotation as np_feed_rotation 98 | from africanus.rime.dask import feed_rotation 99 | 100 | parangles = np.random.random((10, 5)) 101 | dask_parangles = da.from_array(parangles, chunks=(5, (2, 3))) 102 | 103 | np_fr = np_feed_rotation(parangles, feed_type="linear") 104 | assert np.all(np_fr == feed_rotation(dask_parangles, feed_type="linear")) 105 | 106 | np_fr = np_feed_rotation(parangles, feed_type="circular") 107 | assert np.all(np_fr == feed_rotation(dask_parangles, feed_type="circular")) 108 | -------------------------------------------------------------------------------- /africanus/util/jinja2.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | from africanus.util.requirements import requires_optional 5 | 6 | 7 | def register_assign_cycles(N, case=0): 8 | """ 9 | Determine cycles that stem from performing 10 | an in-place assignment of the elements of an array. 11 | In the following. we cannot naively assign the source index 12 | to the dest index, 13 | If we assigned the value at index 0 to index 1, the assignment 14 | from index 1 to index 3 would be invalid, for example: 15 | 16 | src: [3, 0, 2, 1] 17 | dest: [0, 1, 2, 3] 18 | 19 | assignment cycles can be broken by using a 20 | temporary variable to store the contents of a source index 21 | until dependent assignments have been performed. 22 | 23 | In this way, the minimum number of registers can be used 24 | to perform the in-place assignment. 25 | 26 | Returns 27 | ------- 28 | list of lists of tuples 29 | For example, `[[(0, 2), (2, 0)], [1, 3], [3, 1]]` 30 | 31 | """ 32 | dest = range(N) 33 | src = [(N - case + n) % N for n in dest] 34 | 35 | deps = {d: s for d, s in zip(dest, src) if d != s} 36 | 37 | for di, d in enumerate(dest): 38 | si = src.index(d) 39 | if si > di: 40 | deps[si] = di 41 | 42 | cycles = [] 43 | 44 | while len(deps) > 0: 45 | k, v = deps.popitem() 46 | cycle = [(k, v)] 47 | 48 | while True: 49 | try: 50 | k = v 51 | v = deps.pop(k) 52 | except KeyError: 53 | # Check that the last key we're trying 54 | # to get is the first one in the cycle 55 | assert k == cycle[0][0] 56 | break 57 | 58 | cycle.append((k, v)) 59 | 60 | cycles.append(cycle) 61 | 62 | return cycles 63 | 64 | 65 | class TemplatingException(Exception): 66 | def __init__(self, msg): 67 | super(TemplatingException, self).__init__(msg) 68 | 69 | 70 | def throw_helper(msg): 71 | raise TemplatingException(msg) 72 | 73 | 74 | class FakeEnvironment(object): 75 | """ 76 | Fake jinja2 environment, for which attribute/dict 77 | type access will fail 78 | """ 79 | 80 | @requires_optional("jinja2") 81 | def __getitem__(self, key): 82 | raise NotImplementedError() 83 | 84 | @requires_optional("jinja2") 85 | def __setitem__(self, key, value): 86 | raise NotImplementedError() 87 | 88 | @requires_optional("jinja2") 89 | def __delitem__(self, key): 90 | raise NotImplementedError() 91 | 92 | @requires_optional("jinja2") 93 | def __getattr__(self, name): 94 | raise NotImplementedError() 95 | 96 | @requires_optional("jinja2") 97 | def __setattr__(self, name, value): 98 | raise NotImplementedError() 99 | 100 | @requires_optional("jinja2") 101 | def __delattr__(self, name): 102 | raise NotImplementedError() 103 | 104 | 105 | def _jinja2_env_factory(): 106 | try: 107 | from jinja2 import Environment, PackageLoader, select_autoescape 108 | except ImportError: 109 | return FakeEnvironment() 110 | 111 | loader = PackageLoader("africanus", ".") 112 | autoescape = select_autoescape(["j2", "cu.j2"]) 113 | env = Environment( 114 | loader=loader, autoescape=autoescape, extensions=["jinja2.ext.do"] 115 | ) 116 | 117 | # TODO(sjperkins) 118 | # Find a better way to set globals 119 | # perhaps search the package tree for e.g. 120 | # `jinja2_setup`.py files, whose contents 121 | # are inspected and assigned into the globals dict 122 | env.globals["register_assign_cycles"] = register_assign_cycles 123 | env.globals["throw"] = throw_helper 124 | 125 | return env 126 | 127 | 128 | jinja_env = _jinja2_env_factory() 129 | -------------------------------------------------------------------------------- /africanus/model/spectral/tests/test_spectral_model.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | import numpy as np 5 | from numpy.testing import assert_array_almost_equal 6 | import pytest 7 | 8 | from africanus.model.spectral.spec_model import spectral_model, numpy_spectral_model 9 | 10 | 11 | @pytest.fixture 12 | def flux(): 13 | def impl(nsrc): 14 | return np.random.normal(size=nsrc) 15 | 16 | return impl 17 | 18 | 19 | @pytest.fixture 20 | def frequency(): 21 | def impl(nchan): 22 | return np.linspace(0.856e9, 2 * 0.856e9, nchan) 23 | 24 | return impl 25 | 26 | 27 | @pytest.fixture 28 | def ref_freq(): 29 | def impl(shape): 30 | return np.full(shape, 3 * 0.856e9 / 2) 31 | 32 | return impl 33 | 34 | 35 | @pytest.fixture 36 | def spi(): 37 | def impl(shape): 38 | return 0.7 + np.random.random(shape) * 0.2 39 | 40 | return impl 41 | 42 | 43 | @pytest.mark.parametrize( 44 | "base", [0, 1, 2, "std", "log", "log10", ["log", "std", "std", "std"]] 45 | ) 46 | @pytest.mark.parametrize("npol", [0, 1, 2, 4]) 47 | def test_spectral_model_multiple_spi(flux, ref_freq, frequency, base, npol): 48 | nsrc = 10 49 | nchan = 16 50 | nspi = 6 51 | 52 | # Clip to number of polarisations 53 | if isinstance(base, list): 54 | base = base[0] if npol == 0 else base[:npol] 55 | 56 | if npol > 0: 57 | flux_slice = (slice(None), None) 58 | flux_shape = (nsrc, npol) 59 | spi_shape = (nsrc, nspi, npol) 60 | else: 61 | flux_slice = (slice(None),) 62 | flux_shape = (nsrc,) 63 | spi_shape = (nsrc, nspi) 64 | 65 | I = flux(nsrc) # noqa 66 | stokes = np.broadcast_to(I[flux_slice], flux_shape) # noqa 67 | spi = 0.7 + np.random.random(spi_shape) * 0.2 68 | ref_freq = ref_freq(nsrc) 69 | freq = frequency(nchan) 70 | 71 | model = spectral_model(stokes, spi, ref_freq, freq, base=base) 72 | np_model = numpy_spectral_model(stokes, spi, ref_freq, freq, base=base) 73 | 74 | assert_array_almost_equal(model, np_model) 75 | assert model.flags.c_contiguous is True 76 | 77 | 78 | @pytest.mark.parametrize( 79 | "base", [0, 1, 2, "std", "log", "log10", ["log", "std", "std", "std"]] 80 | ) 81 | @pytest.mark.parametrize("npol", [0, 1, 2, 4]) 82 | def test_dask_spectral_model(flux, ref_freq, frequency, base, npol): 83 | da = pytest.importorskip("dask.array") 84 | from africanus.model.spectral.spec_model import spectral_model as np_spectral_model 85 | from africanus.model.spectral.dask import spectral_model 86 | 87 | sc = (5, 5) 88 | fc = (8, 8) 89 | spc = (6,) 90 | 91 | nsrc = sum(sc) 92 | nchan = sum(fc) 93 | nspi = sum(spc) 94 | 95 | # Clip to number of polarisations 96 | if isinstance(base, list): 97 | base = base[0] if npol == 0 else base[:npol] 98 | 99 | if npol > 0: 100 | flux_slice = (slice(None), None) 101 | flux_shape = (nsrc, npol) 102 | spi_shape = (nsrc, nspi, npol) 103 | cc = (npol,) 104 | else: 105 | flux_slice = (slice(None),) 106 | flux_shape = (nsrc,) 107 | spi_shape = (nsrc, nspi) 108 | cc = () 109 | 110 | I = flux(nsrc) # noqa 111 | stokes = np.broadcast_to(I[flux_slice], flux_shape) # noqa 112 | spi = 0.7 + np.random.random(spi_shape) * 0.2 113 | ref_freq = ref_freq(nsrc) 114 | freq = frequency(nchan) 115 | 116 | da_stokes = da.from_array(stokes, chunks=(sc,) + cc) 117 | da_spi = da.from_array(spi, chunks=(sc, spc) + cc) 118 | da_ref_freq = da.from_array(ref_freq, chunks=sc) 119 | da_freq = da.from_array(freq, chunks=fc) 120 | 121 | da_model = spectral_model(da_stokes, da_spi, da_ref_freq, da_freq, base=base) 122 | 123 | np_model = np_spectral_model(stokes, spi, ref_freq, freq, base=base) 124 | assert_array_almost_equal(da_model, np_model) 125 | -------------------------------------------------------------------------------- /africanus/calibration/utils/utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import numpy as np 4 | from africanus.util.docs import DocstringTemplate 5 | 6 | DIAG_DIAG = 0 7 | DIAG = 1 8 | FULL = 2 9 | 10 | 11 | def check_type(jones, vis, vis_type="vis"): 12 | if vis_type == "vis": 13 | vis_ndim = (3, 4) 14 | elif vis_type == "model": 15 | vis_ndim = (4, 5) 16 | else: 17 | raise ValueError("Unknown vis_type") 18 | 19 | vis_axes_count = vis.ndim 20 | jones_axes_count = jones.ndim 21 | if vis_axes_count == vis_ndim[0]: 22 | mode = DIAG_DIAG 23 | if jones_axes_count != 5: 24 | raise RuntimeError( 25 | "Jones axes not compatible with \ 26 | visibility axes. Expected length \ 27 | 5 but got length %d" 28 | % jones_axes_count 29 | ) 30 | 31 | elif vis_axes_count == vis_ndim[1]: 32 | if jones_axes_count == 5: 33 | mode = DIAG 34 | 35 | elif jones_axes_count == 6: 36 | mode = FULL 37 | else: 38 | raise RuntimeError("Jones term has incorrect shape") 39 | else: 40 | raise RuntimeError("Visibility data has incorrect shape") 41 | 42 | return mode 43 | 44 | 45 | def chunkify_rows(time, utimes_per_chunk): 46 | utimes, time_bin_counts = np.unique(time, return_counts=True) 47 | n_time = len(utimes) 48 | if utimes_per_chunk <= 0: 49 | utimes_per_chunk = n_time 50 | row_chunks = [ 51 | np.sum(time_bin_counts[i : i + utimes_per_chunk]) 52 | for i in range(0, n_time, utimes_per_chunk) 53 | ] 54 | time_bin_indices = np.zeros(n_time, dtype=np.int32) 55 | time_bin_indices[1::] = np.cumsum(time_bin_counts)[0:-1] 56 | time_bin_counts = time_bin_counts.astype(np.int32) 57 | return tuple(row_chunks), time_bin_indices, time_bin_counts 58 | 59 | 60 | CHECK_TYPE_DOCS = DocstringTemplate( 61 | """ 62 | Determines which calibration scenario to apply i.e. 63 | DIAG_DIAG, DIAG or COMPLEX2x2. 64 | 65 | Parameters 66 | ---------- 67 | jones : $(array_type) 68 | Jones term of shape :code:`(time, ant, chan, dir, corr)` 69 | or :code:`(time, ant, chan, dir, corr, corr)` 70 | vis : $(array_type) 71 | Visibility data of shape :code:`(row, chan, corr)` 72 | or :code:`(row, chan, corr, corr)` 73 | vis_type : str 74 | String specifying what kind of visibility we are checking 75 | against. Options are 'vis' or 'model' 76 | 77 | Returns 78 | ------- 79 | mode : integer 80 | An integer representing the calibration mode. 81 | Options are 0 -> DIAG_DIAG, 1 -> DIAG, 2 -> FULL 82 | 83 | """ 84 | ) 85 | 86 | try: 87 | check_type.__doc__ = CHECK_TYPE_DOCS.substitute(array_type=":class:`numpy.ndarray`") 88 | except AttributeError: 89 | pass 90 | 91 | CHUNKIFY_ROWS_DOCS = DocstringTemplate( 92 | """ 93 | Divides rows into chunks containing integer 94 | numbers of times keeping track of the indices 95 | at which the unique time changes and the number 96 | of times per unique time. 97 | 98 | Parameters: 99 | ----------- 100 | 101 | time : $(array_type) 102 | TIME column of MS 103 | utimes_per_chunk : integer 104 | The number of unique times to place in each chunk 105 | 106 | Returns 107 | ------- 108 | row_chunks : tuple 109 | A tuple of row chunks that can be used to initialise 110 | an xds with chunks={'row': row_chunks} for example. 111 | time_bin_indices : $(array_type) 112 | Array containing the indices at which unique time 113 | changes 114 | times_bin_counts : $(array_type) 115 | Array containing the number of times per unique time. 116 | """ 117 | ) 118 | 119 | try: 120 | chunkify_rows.__doc__ = CHUNKIFY_ROWS_DOCS.substitute( 121 | array_type=":class:`numpy.ndarray`" 122 | ) 123 | except AttributeError: 124 | pass 125 | --------------------------------------------------------------------------------