diff --git a/.travis.yml b/.travis.yml index 98041fb..2393bd0 100644 --- a/.travis.yml +++ b/.travis.yml @@ -27,7 +27,7 @@ before_install: - ls - conda update --yes conda install: -- conda install --yes python=$TRAVIS_PYTHON_VERSION numpy scipy nose matplotlib +- conda install --yes python=$TRAVIS_PYTHON_VERSION numpy scipy nose matplotlib numba - pip install codecov pypandoc - python setup.py install script: diff --git a/aotools/turbulence/infinitephasescreen.py b/aotools/turbulence/infinitephasescreen.py index c565588..25e04d7 100644 --- a/aotools/turbulence/infinitephasescreen.py +++ b/aotools/turbulence/infinitephasescreen.py @@ -11,6 +11,13 @@ import numpy from numpy import pi +# Numba compiles python code to machine code for faster execution +try: + import numba +except: + numba = None + + from . import phasescreen, turb __all__ = ["PhaseScreenVonKarman", "PhaseScreenKolmogorov"] @@ -109,14 +116,19 @@ def calc_seperations(self): positions = numpy.append(self.stencil_positions, self.X_positions, axis=0) self.seperations = numpy.zeros((len(positions), len(positions))) - for i, (x1, y1) in enumerate(positions): - for j, (x2, y2) in enumerate(positions): - delta_x = x2 - x1 - delta_y = y2 - y1 + if numba: + calc_seperations_fast(positions, self.seperations) + else: + for i, (x1, y1) in enumerate(positions): + for j, (x2, y2) in enumerate(positions): + delta_x = x2 - x1 + delta_y = y2 - y1 + + delta_r = numpy.sqrt(delta_x ** 2 + delta_y ** 2) + + self.seperations[i, j] = delta_r - delta_r = numpy.sqrt(delta_x ** 2 + delta_y ** 2) - self.seperations[i, j] = delta_r def make_covmats(self): """ @@ -139,7 +151,9 @@ def makeAMatrix(self): cf = linalg.cho_factor(self.cov_mat_zz) inv_cov_zz = linalg.cho_solve(cf, numpy.identity(self.cov_mat_zz.shape[0])) except linalg.LinAlgError: - raise linalg.LinAlgError("Could not invert Covariance Matrix to for A and B Matrices. Try with a larger pixel scale") + # print("Cholesky solve failed. Performing SVD inversion...") + # inv_cov_zz = numpy.linalg.pinv(self.cov_mat_zz) + raise linalg.LinAlgError("Could not invert Covariance Matrix to for A and B Matrices. Try with a larger pixel scale or smaller L0") self.A_mat = self.cov_mat_xz.dot(inv_cov_zz) @@ -398,4 +412,19 @@ def get_new_row(self): def __repr__(self): return str(self.scrn) - \ No newline at end of file + + + +@numba.jit(nopython=True, parallel=True) +def calc_seperations_fast(positions, seperations): + + for i in numba.prange(len(positions)): + x1, y1 = positions[i] + for j in range(len(positions)): + x2, y2 = positions[j] + delta_x = x2 - x1 + delta_y = y2 - y1 + + delta_r = numpy.sqrt(delta_x ** 2 + delta_y ** 2) + + seperations[i, j] = delta_r \ No newline at end of file