aocov 0.1.2

Efficient covariance computations for Adaptive Optics in Python

Efficient covariance computations for Adaptive Optics in Python

Motivation

In Adaptive Optics control/estimation/simulations, it's common to find yourself performing an operation similar to:

import numpy as np
import aotools

# define some grid coordinates in x and y
n = 64
xx, yy = np.meshgrid(np.arange(n), np.arange(n), indexing="xy")
xx = xx.flatten()
yy = yy.flatten()

# compute the distances between every pair of coordinates in this grid
rr = ((xx[:, None]-xx[None, :])**2 + (yy[:, None]-yy[None, :])**2)**0.5

# evaluate the covariance at all points in 
r0 = 0.15
L0 = 25.0
cov = aotools.phase_covariance(rr, r0, L0)

The goal of this package is to address 2 main limitations of this simple operation:

1. The evaluation of the covariance function is highly parallelisable, though the implementation above is limited to CPU execution.
2. The distance matrix (rr in this case) is highly redundant, with many repeated values due to the regular geometry of the inputs.

These problems are addressed simultaneously by:

1. using PyTorch to execute von Karman covariance functions on (e.g.) GPU, by simply defining a device to execute on. Note that PyTorch must be installed with the appropriate configuration, though for modern CUDA GPU devices, this is done by default. Just use device="gpu" when you call aocov.phase_covariance.
2. calculating the unique elements of the distance matrix, and only evaluating the covariance function for these values, then copying the values for all other duplicated values.

Modifying the above example, we would instead use:

import numpy as np
import aocov # <-- import this package

# define some grid coordinates in x and y
n = 64
xx, yy = np.meshgrid(np.arange(n), np.arange(n), indexing="xy")
xx = xx.flatten()
yy = yy.flatten()

# compute the distances between every pair of coordinates in this grid
rr = ((xx[:, None]-xx[None, :])**2 + (yy[:, None]-yy[None, :])**2)**0.5

# evaluate a von Karman covariance on GPU
r0 = 0.15
L0 = 25.0
cov = aocov.phase_covariance(rr, r0, L0, device="cuda:0")

Furthermore, since the distance computations can become expensive, and is prone to typos, we provide a further convenience function, phase_covariance_xyxy. Finally, the example becomes:

import numpy as np
import aocov # <-- import this package

# define some grid coordinates in x and y
n = 64
xx, yy = np.meshgrid(np.arange(n), np.arange(n), indexing="xy")
xx = xx.flatten()
yy = yy.flatten()

# evaluate a von Karman covariance on GPU
r0 = 0.15
L0 = 25.0
cov = aocov.phase_covariance_xyxy(xx, yy, xx, yy, r0, L0, device="cuda:0")

Note that the last option is likely to perform the fastest, see performance comparison below. Also, for the time being, PyTorch does not implement the required bessel function (scipy.special.kv(nu,x)), so the code makes a single call to the scipy function on CPU. If someone makes a torch-friendly version of kv(nu,x), this would further improve the performance of this module - especially on GPU.

Installation

Simply clone and install this git repo using pip.

git clone [e-mail]:jcranney/aocov
cd aocov
pip install -e .

Performance

This repo comes with a performance test script perf/perf_test.py, you can tweak it to match your system dimensions. Note that $n$ is the number of samples in each dimension for the grid, so assuming a square grid, the covariance matrix will have $n^4$ elements, i.e.,:

n	elements
16	65,536
32	1,048,576
64	16,777,216
128	268,435,456

Using the default parameters on a system 40 cores and a Tesla V100 32GB GPU, I measured the following:

n	device	experiment	sec/matrix
16	cpu	no pytorch	2.880e-02
16	cpu	rr in numpy, rest in aocov	1.222e-02
16	cpu	all in aocov	1.435e-02
16	cuda:0	rr in numpy, rest in aocov	1.110e+00
16	cuda:0	all in aocov	2.640e-03
32	cpu	no pytorch	4.535e-01
32	cpu	rr in numpy, rest in aocov	8.379e-02
32	cpu	all in aocov	7.282e-02
32	cuda:0	rr in numpy, rest in aocov	8.874e-03
32	cuda:0	all in aocov	2.020e-03
64	cpu	no pytorch	7.233e+00
64	cpu	rr in numpy, rest in aocov	1.269e+00
64	cpu	all in aocov	1.159e+00
64	cuda:0	rr in numpy, rest in aocov	1.777e-01
64	cuda:0	all in aocov	2.087e-02
128	cpu	no pytorch	1.143e+02
128	cpu	rr in numpy, rest in aocov	2.218e+01
128	cpu	all in aocov	2.056e+01
128	cuda:0	rr in numpy, rest in aocov	2.560e+00
128	cuda:0	all in aocov	3.152e-01

For $n=2^7=128$, the covariance matrix has $2^{28}=268, 435, 456$ elements. In that case, with a decent GPU, you can produce a covariance matrix $362$ times faster. In the case above, we go from almost 2 minutes to less than 0.5 seconds.