Merge pull request #248 from Astroua/ellipse2Dfit_speedup

Add unbinned power spectrum fitting

.github/workflows/main.yml

@@ -14,6 +14,11 @@ jobs:
       fail-fast: false
       matrix:
         include:
+          - os: ubuntu-latest
+            python-version: '3.11'
+            name: Python 3.11 with minimal dependencies
+            toxenv: py311-test
+
           - os: ubuntu-latest
             python-version: '3.10'
             name: Python 3.10 with minimal dependencies
@@ -60,9 +65,9 @@ jobs:
             toxenv: build_docs
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install testing dependencies
@@ -71,6 +76,6 @@ jobs:
       run: tox -v -e ${{ matrix.toxenv }}
     - name: Upload coverage to codecov
       if: ${{ contains(matrix.toxenv,'-cov') }}
-      uses: codecov/codecov-action@v1.0.13
+      uses: codecov/codecov-action@v3
       with:
         file: ./coverage.xml

.github/workflows/python-publish.yml

@@ -10,22 +10,18 @@ jobs:
       matrix:
         os: [ubuntu-latest]
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions/setup-python@v2
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
         name: Install Python
         with:
-          python-version: '3.9'
-      - name: Install cibuildwheel
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install cibuildwheel
+          python-version: '3.11'
       - name: Build wheels
-        run: python -m cibuildwheel --output-dir wheelhouse
+        uses: pypa/cibuildwheel@v2.16.2
         env:
-          CIBW_BUILD: cp37-manylinux_x86_64 cp38-manylinux_x86_64 cp39-manylinux_x86_64 cp310-manylinux_x86_64
+          CIBW_BUILD: cp37-manylinux_x86_64 cp38-manylinux_x86_64 cp39-manylinux_x86_64 cp310-manylinux_x86_64 cp311-manylinux_x86_64
           CIBW_TEST_EXTRAS: test
           CIBW_TEST_COMMAND: pytest --pyargs turbustat
-      - uses: actions/upload-artifact@v2
+      - uses: actions/upload-artifact@v3
         with:
           path: ./wheelhouse/*.whl
 
@@ -37,22 +33,18 @@ jobs:
       matrix:
         os: [macos-latest]
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions/setup-python@v2
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
         name: Install Python
         with:
-          python-version: '3.9'
-      - name: Install cibuildwheel
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install cibuildwheel
+          python-version: '3.11'
       - name: Build wheels
-        run: python -m cibuildwheel --output-dir wheelhouse
+        uses: pypa/cibuildwheel@v2.16.2
         env:
-          CIBW_BUILD: cp37-* cp38-* cp39-* cp310-*
+          CIBW_BUILD: cp37-* cp38-* cp39-* cp310-* cp311-*
           CIBW_TEST_EXTRAS: test
           CIBW_TEST_COMMAND: pytest --pyargs turbustat
-      - uses: actions/upload-artifact@v2
+      - uses: actions/upload-artifact@v3
         with:
           path: ./wheelhouse/*.whl
 
@@ -64,41 +56,37 @@ jobs:
       matrix:
         os: [windows-latest]
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions/setup-python@v2
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
         name: Install Python
         with:
-          python-version: '3.9'
-      - name: Install cibuildwheel
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install cibuildwheel
+          python-version: '3.11'
       - name: Build wheels
-        run: python -m cibuildwheel --output-dir wheelhouse
+        uses: pypa/cibuildwheel@v2.16.2
         env:
-          CIBW_BUILD: cp37-* cp38-* cp39-*
+          CIBW_BUILD: cp39-*win_amd64 cp310-*win_amd64 cp311-*win_amd64
           CIBW_TEST_EXTRAS: test
           CIBW_TEST_COMMAND: pytest --pyargs turbustat
-      - uses: actions/upload-artifact@v2
+      - uses: actions/upload-artifact@v3
         with:
           path: ./wheelhouse/*.whl
 
   build_sdist:
     name: Build source distribution
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions/setup-python@v2
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
         name: Install Python
         with:
-          python-version: '3.9'
+          python-version: '3.11'
       - name: Install build
         run: |
           python -m pip install --upgrade pip
           python -m pip install build
       - name: Build sdist
         run: python -m build --sdist --outdir dist/ .
-      - uses: actions/upload-artifact@v2
+      - uses: actions/upload-artifact@v3
         with:
           path: dist/*.tar.gz
 
@@ -108,11 +96,11 @@ jobs:
     runs-on: ubuntu-latest
     if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags/v')
     steps:
-      - uses: actions/download-artifact@v2
+      - uses: actions/download-artifact@v3
         with:
           name: artifact
           path: dist
-      - uses: pypa/gh-action-pypi-publish@master
+      - uses: pypa/gh-action-pypi-publish@release/v1
         with:
           user: __token__
           password: ${{ secrets.pypi_password }}

.gitignore

@@ -64,6 +64,7 @@ docs/api
 *.pyc
 
 *.DS_Store
+*.ipynb_checkpoints
 
 turbustat/version.py
 turbustat/cython_version.py

setup.cfg

@@ -48,7 +48,7 @@ testpaths = "turbustat" "docs"
 astropy_header = true
 doctest_plus = enabled
 text_file_format = rst
-addopts = --doctest-rst
+addopts = --doctest-rst -p no:warnings
 
 [coverage:run]
 omit =

tox.ini

@@ -1,6 +1,6 @@
 [tox]
 envlist =
-    py{37,38,39,310}-test{,-all,-dev,-cov}
+    py{37,38,39,310,311}-test{,-all,-dev,-cov}
     build_docs
     codestyle
 requires =
@@ -32,8 +32,8 @@ extras =
     all: all
 commands =
     pip freeze
-    !cov: pytest --open-files --pyargs turbustat {toxinidir}/docs {posargs}
-    cov: pytest --open-files --pyargs turbustat {toxinidir}/docs --cov turbustat --cov-config={toxinidir}/setup.cfg {posargs}
+    !cov: pytest --pyargs turbustat {toxinidir}/docs {posargs}
+    cov: pytest --pyargs turbustat {toxinidir}/docs --cov turbustat --cov-config={toxinidir}/setup.cfg {posargs}
     cov: coverage xml -o {toxinidir}/coverage.xml
 
 [testenv:build_docs]

turbustat/moments/_moment_errs.py

@@ -238,7 +238,7 @@ def _slice0(cube, axis, scale):
     result = np.zeros(shp) * cube.unit ** 2
 
     view = [slice(None)] * 3
-    valid = np.zeros(shp, dtype=np.bool)
+    valid = np.zeros(shp, dtype=bool)
 
     for i in range(cube.shape[axis]):
         view[axis] = i

turbustat/statistics/base_pspec2.py

@@ -88,7 +88,7 @@ def compute_radial_pspec(self, logspacing=False, max_bin=None, **kwargs):
         '''
 
         # Check if azimuthal constraints are given
-        if kwargs.get("theta_0"):
+        if "theta_0" in kwargs:
             azim_constraint_flag = True
         else:
             azim_constraint_flag = False
@@ -106,7 +106,8 @@ def compute_radial_pspec(self, logspacing=False, max_bin=None, **kwargs):
         # Attach units to freqs
         self._freqs = self.freqs / u.pix
 
-    def fit_pspec(self, brk=None, log_break=False, low_cut=None,
+    def fit_pspec(self, fit_unbinned=False,
+                  brk=None, log_break=False, low_cut=None,
                   high_cut=None, min_fits_pts=10, weighted_fit=False,
                   bootstrap=False, bootstrap_kwargs={},
                   verbose=False):
@@ -118,6 +119,10 @@ def fit_pspec(self, brk=None, log_break=False, low_cut=None,
 
         Parameters
         ----------
+        fit_unbinned : bool, optional
+            Fits the unbinned 2D power-spectrum to the linear model. Default is True.
+            Use False to fit the binned 1D power-spectrum instead and replicate fitting
+            in earlier TurbuStat versions.
         brk : float or None, optional
             Guesses for the break points. If given as a list, the length of
             the list sets the number of break points to be fit. If a choice is
@@ -150,29 +155,62 @@ def fit_pspec(self, brk=None, log_break=False, low_cut=None,
 
         self._bootstrap_flag = bootstrap
 
-        # Make the data to fit to
-        if low_cut is None:
-            # Default to the largest frequency, since this is just 1 pixel
-            # in the 2D PSpec.
-            self.low_cut = 1. / (0.5 * float(max(self.ps2D.shape)) * u.pix)
-        else:
-            self.low_cut = self._to_pixel_freq(low_cut)
+        if fit_unbinned:
+            yy_freq, xx_freq = make_radial_freq_arrays(self.ps2D.shape)
+
+            freqs_2d = np.sqrt(yy_freq**2 + xx_freq**2) / u.pix
+
+            # Make the data to fit to
+            if low_cut is None:
+                # Default to the largest frequency, since this is just 1 pixel
+                # in the 2D PSpec.
+                self.low_cut = 1. / (0.5 * float(max(self.ps2D.shape)) * u.pix)
+            else:
+                self.low_cut = self._to_pixel_freq(low_cut)
+
+            if high_cut is None:
+                # self.high_cut = self.freqs.max().value / u.pix
+                self.high_cut = freqs_2d.value.max() / u.pix
+            else:
+                self.high_cut = self._to_pixel_freq(high_cut)
+
+
+            clip_mask = clip_func(freqs_2d.value,
+                                self.low_cut.value,
+                                self.high_cut.value)
+
+
+            x = np.log10(freqs_2d.value[clip_mask])
+            y = np.log10(self.ps2D[clip_mask])
 
-        if high_cut is None:
-            self.high_cut = self.freqs.max().value / u.pix
         else:
-            self.high_cut = self._to_pixel_freq(high_cut)
+            # Make the data to fit to
+            if low_cut is None:
+                # Default to the largest frequency, since this is just 1 pixel
+                # in the 2D PSpec.
+                self.low_cut = 1. / (0.5 * float(max(self.ps2D.shape)) * u.pix)
+            else:
+                self.low_cut = self._to_pixel_freq(low_cut)
+
+            if high_cut is None:
+                self.high_cut = self.freqs.max().value / u.pix
+            else:
+                self.high_cut = self._to_pixel_freq(high_cut)
 
-        x = np.log10(self.freqs[clip_func(self.freqs.value, self.low_cut.value,
-                                          self.high_cut.value)].value)
+            x = np.log10(self.freqs[clip_func(self.freqs.value, self.low_cut.value,
+                                            self.high_cut.value)].value)
 
-        clipped_ps1D = self.ps1D[clip_func(self.freqs.value,
-                                           self.low_cut.value,
-                                           self.high_cut.value)]
-        y = np.log10(clipped_ps1D)
+            clipped_ps1D = self.ps1D[clip_func(self.freqs.value,
+                                            self.low_cut.value,
+                                            self.high_cut.value)]
+            y = np.log10(clipped_ps1D)
 
         if weighted_fit:
+            if fit_unbinned:
+                raise NotImplementedError("Error propagation for the unbinned modeling is not "
+                                          "implemented yet.")
 
+            # Currently this will run only for the binned fitting.
             clipped_stddev = self.ps1D_stddev[clip_func(self.freqs.value,
                                                         self.low_cut.value,
                                                         self.high_cut.value)]
@@ -181,6 +219,10 @@ def fit_pspec(self, brk=None, log_break=False, low_cut=None,
 
             y_err = 0.434 * clipped_stddev / clipped_ps1D
 
+            weights = 1 / y_err**2
+        else:
+            weights = None
+
         if brk is not None:
             # Try the fit with a break in it.
             if not log_break:
@@ -192,11 +234,6 @@ def fit_pspec(self, brk=None, log_break=False, low_cut=None,
                     assert brk.unit == u.dimensionless_unscaled
                     brk = brk.value
 
-            if weighted_fit:
-                weights = 1 / y_err**2
-            else:
-                weights = None
-
             brk_fit = Lm_Seg(x, y, brk, weights=weights)
             brk_fit.fit_model(verbose=verbose, cov_type='HC3')
 
@@ -246,7 +283,7 @@ def fit_pspec(self, brk=None, log_break=False, low_cut=None,
             x = sm.add_constant(x)
 
             if weighted_fit:
-                model = sm.WLS(y, x, missing='drop', weights=1 / y_err**2)
+                model = sm.WLS(y, x, missing='drop', weights=weights)
             else:
                 model = sm.OLS(y, x, missing='drop')
 
@@ -614,8 +651,6 @@ def plot_fit(self, show_2D=False, show_residual=True,
         good_interval = clip_func(self.freqs.value, self.low_cut.value,
                                   self.high_cut.value)
 
-        y_fit = self.fit.fittedvalues
-
         if show_residual:
             if isinstance(self.slope, np.ndarray):
                 # Broken linear model
@@ -679,7 +714,11 @@ def plot_fit(self, show_2D=False, show_residual=True,
                                    fmt=symbol, markersize=5, alpha=0.5,
                                    capsize=10, elinewidth=3)
 
-        ax_1D.plot(np.log10(xvals[fit_index]), y_fit, linestyle='-',
+        xvals_plot = np.log10(xvals[fit_index]).ravel()
+        # y_fit = self.fit.fittedvalues
+        y_fit = self.fit.predict(sm.add_constant(xvals_plot))
+
+        ax_1D.plot(xvals_plot, y_fit, linestyle='-',
                    label=label, linewidth=3, color=fit_color)
 
         if show_residual: