Backport PR wtbarnes#241: Fix vectorize_where bug and RTD config

.pre-commit-config.yaml

@@ -1,30 +1,20 @@
 repos:
-  # The warnings/errors we check for here are:
-  # E101 - mix of tabs and spaces
-  # W191 - use of tabs
-  # W291 - trailing whitespace
-  # W292 - no newline at end of file
-  # W293 - trailing whitespace
-  # W391 - blank line at end of file
-  # E111 - 4 spaces per indentation level
-  # E112 - 4 spaces per indentation level
-  # E113 - 4 spaces per indentation level
-  # E303 - too many blank lines (3)
-  # E304 - blank lines found after function decorator
-  # E305 - expected 2 blank lines after class or function definition
-  # E306 - expected 1 blank line before a nested definition
-  # E502 - the backslash is redundant between brackets
-  # E722 - do not use bare except
-  # E901 - SyntaxError or IndentationError
-  # E902 - IOError
-  # F822: undefined name in __all__
-  # F823: local variable name referenced before assignment
-  - repo: https://github.com/pycqa/flake8
-    rev: 6.0.0
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: "v0.1.5"
     hooks:
-      - id: flake8
-        additional_dependencies: [flake8-use-pathlib]
-        args: ['--count', '--select', 'E101,W191,W291,W292,W293,W391,E111,E112,E113,E303,E304,E306,E502,E722,E901,E902,F822,F823,PL100,PL102,PL103,PL104,PL105,PL110,PL112,PL113,PL18,PL120,PL122,PL123']
+      - id: ruff
+        args: ["--fix"]
+
+  - repo: https://github.com/PyCQA/isort
+    rev: 5.12.0
+    hooks:
+    - id: isort
+      name: isort
+      entry: isort
+      require_serial: true
+      language: python
+      types:
+      - python
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.4.0
@@ -41,31 +31,9 @@ repos:
       - id: mixed-line-ending
         exclude: ".*(.fits|.fts|.fit|.txt|.bib|tca.*)$"
 
-  - repo: https://github.com/asottile/pyupgrade
-    rev: v3.3.1
-    hooks:
-    - id: pyupgrade
-      args: [--keep-runtime-typing, --py38-plus]
-      files: ^fiasco/
-
   - repo: https://github.com/codespell-project/codespell
-    rev: v2.2.2
+    rev: v2.2.6
     hooks:
       - id: codespell
-        args: ['--config setup.cfg']
-
-  - repo: https://github.com/MarcoGorelli/absolufy-imports
-    rev: v0.3.1
-    hooks:
-    - id: absolufy-imports
-
-  - repo: https://github.com/PyCQA/isort
-    rev: 5.12.0
-    hooks:
-    - id: isort
-      name: isort
-      entry: isort
-      require_serial: true
-      language: python
-      types:
-      - python
+        additional_dependencies:
+        - tomli

.readthedocs.yml

@@ -5,6 +5,12 @@
 # Required
 version: 2
 
+# Set the version of Python and other tools you might need
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.11"
+
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
   configuration: docs/conf.py
@@ -14,10 +20,9 @@ formats: []
 
 # Optionally set the version of Python and requirements required to build your docs
 python:
-  version: 3.8
   install:
-    - method: pip
-      path: .
-      extra_requirements:
-        - docs
-        - all
+  - method: pip
+    path: .
+    extra_requirements:
+    - docs
+    - all

docs/conf.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 #
 # Configuration file for the Sphinx documentation builder.
 #

examples/nei_populations.py

@@ -159,7 +159,7 @@
 Te_eff =  carbon.temperature[[(np.fabs(carbon.equilibrium_ionization - carbon_nei[i, :])).sum(axis=1).argmin()
                           for i in range(carbon_nei.shape[0])]]
 plt.plot(t, Te.to('MK'), label=r'$T$')
-plt.plot(t, Te_eff.to('MK'), label='$T_{\mathrm{eff}}$')
+plt.plot(t, Te_eff.to('MK'), label=r'$T_{\mathrm{eff}}$')
 plt.xlim(t[[0,-1]].value)
 plt.legend()
 plt.show()

fiasco/elements.py

@@ -35,7 +35,7 @@ class Element(fiasco.IonCollection):
 
     @u.quantity_input
     def __init__(self, element_name, temperature: u.K, **kwargs):
-        if type(element_name) is str:
+        if isinstance(element_name, str):
             element_name = element_name.capitalize()
         Z = plasmapy.particles.atomic_number(element_name)
         ion_list = []

fiasco/io/datalayer.py

@@ -106,7 +106,7 @@ def __contains__(self, key):
         return key_in_grp
 
     def __getitem__(self, key):
-        if type(key) is int:
+        if isinstance(key, int):
             raise NotImplementedError('Iteration not supported.')
         if key not in self:
             raise KeyError(f'{key} not found in {self.top_level_path}')

fiasco/io/generic.py

@@ -1,7 +1,6 @@
 """
 Base class for file parser
 """
-import astropy
 import astropy.units as u
 import h5py
 import numpy as np
@@ -94,7 +93,7 @@ def preprocessor(self, table, line, *args):
             line = self.fformat.read(line)
         else:
             line = line.strip().split()
-        line = [item.strip() if type(item) is str else item for item in line]
+        line = [item.strip() if isinstance(item, str) else item for item in line]
         table.append(line)
 
     def postprocessor(self, df):

fiasco/io/sources/continuum_sources.py

@@ -4,7 +4,6 @@
 import astropy.units as u
 import h5py
 import numpy as np
-import os
 import pathlib
 import plasmapy
 

fiasco/io/sources/non_ion_sources.py

@@ -4,7 +4,6 @@
 import astropy.units as u
 import fortranformat
 import numpy as np
-import os
 import pathlib
 import plasmapy
 

fiasco/ions.py

@@ -127,7 +127,7 @@ def __radd__(self, value):
 
     @property
     def _instance_kwargs(self):
-        # Keyword arguments used to istantiate this Ion. These are useful when
+        # Keyword arguments used to instantiate this Ion. These are useful when
         # constructing a new Ion instance that pulls from exactly the same
         # data sources.
         kwargs = {

fiasco/tests/idl/test_idl_ioneq.py

@@ -2,7 +2,6 @@
 IDL comparison tests for ionization equilibria
 """
 import astropy.units as u
-import numpy as np
 import pathlib
 import pytest
 
@@ -16,7 +15,7 @@ def ioneq_from_idl(idl_env, ascii_dbase_root):
     read_ioneq, ioneqfile, ioneq_logt, ioneq, ioneq_ref
     """
     args = {
-        'ioneq_filename': ascii_dbase_root / 'ioneq' / f'chianti.ioneq',
+        'ioneq_filename': ascii_dbase_root / 'ioneq' / 'chianti.ioneq',
     }
     res_idl = idl_env.run(script, args=args, verbose=True)
     return res_idl

fiasco/tests/test_fiasco.py

@@ -3,21 +3,20 @@
 """
 import astropy.units as u
 import numpy as np
-import pytest
 
 import fiasco
 
 
 def test_list_elements(hdf5_dbase_root):
     # FIXME: actually test the expected elements
     elements = fiasco.list_elements(hdf5_dbase_root)
-    assert type(elements) is list
+    assert isinstance(elements, list)
 
 
 def test_list_ions(hdf5_dbase_root):
     # FIXME: actually test the expected ions
     ions = fiasco.list_ions(hdf5_dbase_root)
-    assert type(ions) is list
+    assert isinstance(ions, list)
 
 
 def test_proton_electron_ratio(hdf5_dbase_root):

fiasco/tests/test_ion.py

@@ -198,8 +198,12 @@ def test_contribution_function(ion):
     # This value has not been tested for correctness
     assert u.allclose(cont_func[0, 0, 0], 2.08668713e-30 * u.cm**3 * u.erg / u.s)
 
-
-def test_emissivity_shape(c6):
+@pytest.mark.parametrize(('density', 'use_coupling'), [
+    ([1e9,] * u.cm**(-3), False),
+    ([1e8, 1e9, 1e10] * u.cm**(-3), False),
+    (None, True)
+])
+def test_emissivity_shape(c6, density, use_coupling):
     # NOTE: Explicitly testing C VI here because it has a psplups file
     # and thus will compute the proton rates as well which have a specific
     # codepath for coupled density/temperature.
@@ -208,23 +212,15 @@ def test_emissivity_shape(c6):
     # Using the emissivity quantity here because it is the highest level
     # product that needs to manipulate the density. This will implicitly test the
     # contribution function as well.
-    #
-    # Scalar, no coupling
-    density = 1e9 * u.cm**(-3)
-    emiss = c6.emissivity(density)
-    wavelength = c6.transitions.wavelength[~c6.transitions.is_twophoton]
-    assert emiss.shape == c6.temperature.shape + (1,) + wavelength.shape
-    # Array, no coupling
-    density = [1e8, 1e9, 1e10] * u.cm**(-3)
-    emiss = c6.emissivity(density)
-    wavelength = c6.transitions.wavelength[~c6.transitions.is_twophoton]
-    assert emiss.shape == c6.temperature.shape + density.shape + wavelength.shape
-    # Array, with coupling
-    pressure = 1e15 * u.K * u.cm**(-3)
-    density = pressure / c6.temperature
-    emiss = c6.emissivity(density, couple_density_to_temperature=True)
+    if use_coupling:
+        pressure = 1e15 * u.K * u.cm**(-3)
+        density = pressure / c6.temperature
+        density_shape = (1,)
+    else:
+        density_shape = density.shape
+    emiss = c6.emissivity(density, couple_density_to_temperature=use_coupling)
     wavelength = c6.transitions.wavelength[~c6.transitions.is_twophoton]
-    assert emiss.shape == c6.temperature.shape + (1,) + wavelength.shape
+    assert emiss.shape == c6.temperature.shape + density_shape + wavelength.shape
 
 
 def test_coupling_unequal_dimensions_exception(ion):

fiasco/util/tests/test_tools.py

@@ -1,7 +1,25 @@
 import numpy as np
 import pytest
 
-from fiasco.util import burgess_tully_descale
+from fiasco.util import burgess_tully_descale, vectorize_where
+
+
+@pytest.mark.parametrize(('a', 'b', 'c'), [
+    (np.arange(10), [3, 5, 9], np.array([3, 5, 9])),
+    (np.arange(10), [6], np.array([6])),
+    (np.arange(10), 6, np.array([6])),
+])
+def test_vectorize_where(a, b, c):
+    indices = vectorize_where(a, b)
+    assert (indices == c).all()
+
+
+def test_vectorize_where_result_in_array():
+    a = np.arange(20)
+    b = np.array([5, 7, 17, 19])
+    c = np.random.rand(*(15, 10,)+a.shape)
+    indices = vectorize_where(a, b)
+    assert c[:, :, indices].shape == (15, 10) + b.shape
 
 
 @pytest.mark.parametrize('x', [

fiasco/util/tests/test_util.py

@@ -32,6 +32,6 @@ def test_get_chianti_catalog(ascii_dbase_root):
             'ip_files',
             'continuum_files',
             'ion_files']
-    for k in catalog:
+    for k in keys:
         assert k in catalog
         assert isinstance(catalog[k], list)

fiasco/util/tools.py

@@ -21,7 +21,9 @@ def vectorize_where(x_1, x_2):
     x_2 : array-like
         Values to search for
     """
-    return np.vectorize(lambda a, b: np.where(a == b)[0], excluded=[0])(x_1, x_2)
+    x_1 = np.atleast_1d(x_1)
+    x_2 = np.atleast_1d(x_2)
+    return np.array([np.where(x_1==x)[0] for x in x_2]).squeeze()
 
 
 def vectorize_where_sum(x_1, x_2, y, axis=None):

pyproject.toml

@@ -59,3 +59,58 @@ include_trailing_comma = true
 force_alphabetical_sort_within_sections = true
 honor_noqa = true
 lines_between_types = 1
+
+[tool.codespell]
+skip = "*.fts,*.fits,venv,*.pro,*.bib"
+ignore-words-list = "te,emiss"
+
+[tool.ruff]
+target-version = "py39"
+line-length = 110
+exclude=[
+    ".git,",
+    "__pycache__",
+    "build",
+    "fiasco/version.py",
+]
+show-fixes = true
+show-source = true
+
+select = [
+       "E",
+       "F",
+       "W",
+       "UP",
+       "PT",
+       #"RET",
+       #"TID",
+
+]
+extend-ignore = [
+    # pycodestyle (E, W)
+    "E501",  # LineTooLong                             # TODO! fix
+    "E741",  # Ambiguous variable name
+
+    # pytest (PT)
+    "PT001",  # Always use pytest.fixture()
+    "PT004",  # Fixtures which don't return anything should have leading _
+    "PT007",  # Parametrize should be lists of tuples  # TODO! fix
+    "PT011",  # Too broad exception assert             # TODO! fix
+    "PT023",  # Always use () on pytest decorators
+]
+
+[tool.ruff.per-file-ignores]
+# Part of configuration, not a package.
+"setup.py" = ["INP001"]
+"conftest.py" = ["INP001"]
+# implicit-namespace-package. The examples are not a package.
+"docs/*.py" = ["INP001"]
+# Module level imports do not need to be at the top of a file here
+"docs/conf.py" = ["E402"]
+"__init__.py" = ["E402", "F401", "F403"]
+"test_*.py" = ["B011", "D", "E402", "PGH001", "S101"]
+# Allow star import for the factory
+"fiasco/io/factory.py" = ["F403"]
+
+[tool.ruff.pydocstyle]
+convention = "numpy"

setup.cfg

@@ -98,9 +98,3 @@ exclude_lines =
   pragma: py{ignore_python_version}
   # Don't complain about IPython completion helper
   def _ipython_key_completions_
-
-[codespell]
-skip = *.fts,*.fits,venv,*.pro
-ignore-words-list =
-    te,
-    emiss