Fix CI/Lint issues

Author: TheSkyentist

tests/test_config.py

@@ -2,15 +2,14 @@
 
 import warnings
 
-import pytest
 from astropy import units as u
 
 from unite.config import Configuration
 from unite.continuum import ContinuumConfiguration, Linear
 from unite.disperser.base import FluxScale, RScale
 from unite.disperser.config import DispersersConfiguration
 from unite.instruments.nirspec import G235H, G395H
-from unite.line import FWHM, Flux, LineConfiguration, Redshift
+from unite.line import FWHM, LineConfiguration, Redshift
 from unite.prior import TruncatedNormal, Uniform
 
 

tests/test_continuum_functions.py

@@ -2,7 +2,6 @@
 
 import jax.numpy as jnp
 import numpy as np
-import pytest
 from scipy.special import comb as scipy_comb
 
 from unite.continuum.functions import (
@@ -13,7 +12,6 @@
     planck_function,
 )
 
-
 # ---------------------------------------------------------------------------
 # Planck function
 # ---------------------------------------------------------------------------
@@ -109,11 +107,7 @@ def _make_clamped_knots(self, n_basis, degree, low=0.0, high=1.0):
         """Create a clamped knot vector."""
         n_internal = n_basis - degree + 1
         internal = np.linspace(low, high, n_internal)
-        knots = np.concatenate([
-            np.full(degree, low),
-            internal,
-            np.full(degree, high),
-        ])
+        knots = np.concatenate([np.full(degree, low), internal, np.full(degree, high)])
         return jnp.asarray(knots)
 
     def test_basis_partition_of_unity(self):

tests/test_dependent_priors.py

@@ -11,7 +11,6 @@
 - Complex redshift hierarchies (systemic → NLR → BLR → outflow)
 """
 
-import warnings
 
 import astropy.units as u
 import jax.numpy as jnp
@@ -20,20 +19,12 @@
 from jax import random
 from numpyro.infer import Predictive
 
-from unite import line, model, prior
+from unite import model
 from unite.disperser.generic import SimpleDisperser
-from unite.line.config import FWHM, Flux, LineConfiguration, Param, Redshift
-from unite.prior import (
-    Fixed,
-    Parameter,
-    ParameterRef,
-    TruncatedNormal,
-    Uniform,
-    topological_sort,
-)
+from unite.line.config import FWHM, Flux, LineConfiguration, Redshift
+from unite.prior import ParameterRef, TruncatedNormal, Uniform, topological_sort
 from unite.spectrum import Spectra, Spectrum
 
-
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -42,29 +33,24 @@
 def _make_spectrum(wl_range=(6400, 6700), npix=200, name='test'):
     """Create a test spectrum covering the given range."""
     wl = np.linspace(*wl_range, npix) * u.AA
-    disperser = SimpleDisperser(
-        wavelength=wl.value, unit=u.AA, R=3000.0, name=name,
-    )
+    disperser = SimpleDisperser(wavelength=wl.value, unit=u.AA, R=3000.0, name=name)
     low = wl - 0.5 * np.gradient(wl)
     high = wl + 0.5 * np.gradient(wl)
     flux_unit = u.Unit('1e-17 erg / (s cm2 AA)')
     rng = np.random.default_rng(42)
     flux = (10.0 + rng.normal(0, 1, npix)) * flux_unit
     error = np.full(npix, 1.0) * flux_unit
     return Spectrum(
-        low=low, high=high, flux=flux, error=error,
-        disperser=disperser, name=name,
+        low=low, high=high, flux=flux, error=error, disperser=disperser, name=name
     )
 
 
 def _prepare_and_build(line_config, spectra, cont_config=None):
     """Prepare spectra and build model."""
     spectra.prepare(line_config, cont_config)
-    spectra.compute_scales(
-        spectra.prepared_line_config, spectra.prepared_cont_config,
-    )
+    spectra.compute_scales(spectra.prepared_line_config, spectra.prepared_cont_config)
     return model.ModelBuilder(
-        spectra.prepared_line_config, spectra.prepared_cont_config, spectra,
+        spectra.prepared_line_config, spectra.prepared_cont_config, spectra
     ).build()
 
 
@@ -100,9 +86,7 @@ def test_four_level_chain(self):
         c = FWHM('c', prior=Uniform(low=b * 1.5, high=1000))
         d = FWHM('d', prior=Uniform(low=c + 100, high=5000))
 
-        named_priors = {
-            'a': a.prior, 'b': b.prior, 'c': c.prior, 'd': d.prior,
-        }
+        named_priors = {'a': a.prior, 'b': b.prior, 'c': c.prior, 'd': d.prior}
         param_to_name = {a: 'a', b: 'b', c: 'c', d: 'd'}
         order = topological_sort(named_priors, param_to_name)
 
@@ -114,16 +98,10 @@ def test_three_level_redshift_chain(self):
         z_nlr = Redshift(
             'z_nlr',
             prior=TruncatedNormal(
-                loc=z_sys,
-                scale=0.001,
-                low=z_sys - 0.005,
-                high=z_sys + 0.005,
+                loc=z_sys, scale=0.001, low=z_sys - 0.005, high=z_sys + 0.005
             ),
         )
-        z_out = Redshift(
-            'z_out',
-            prior=Uniform(low=z_nlr - 0.01, high=z_nlr),
-        )
+        z_out = Redshift('z_out', prior=Uniform(low=z_nlr - 0.01, high=z_nlr))
 
         named_priors = {
             'z_sys': z_sys.prior,
@@ -171,8 +149,7 @@ def test_redshift_token_as_loc(self):
         """Redshift token passed directly as TruncatedNormal loc."""
         z_sys = Redshift('z_sys', prior=Uniform(-0.01, 0.01))
         z_nlr = Redshift(
-            'z_nlr',
-            prior=TruncatedNormal(loc=z_sys, scale=0.001, low=-0.02, high=0.02),
+            'z_nlr', prior=TruncatedNormal(loc=z_sys, scale=0.001, low=-0.02, high=0.02)
         )
         assert z_sys in z_nlr.prior.dependencies()
 
@@ -193,12 +170,7 @@ def test_all_three_bounds_depend_on_different_params(self):
 
         constrained = FWHM(
             'constrained',
-            prior=TruncatedNormal(
-                loc=center,
-                scale=50.0,
-                low=lower,
-                high=upper,
-            ),
+            prior=TruncatedNormal(loc=center, scale=50.0, low=lower, high=upper),
         )
 
         deps = constrained.prior.dependencies()
@@ -211,12 +183,7 @@ def test_loc_and_low_same_token(self):
         base = FWHM('base', prior=Uniform(100, 500))
         derived = FWHM(
             'derived',
-            prior=TruncatedNormal(
-                loc=base + 100,
-                scale=30.0,
-                low=base,
-                high=2000,
-            ),
+            prior=TruncatedNormal(loc=base + 100, scale=30.0, low=base, high=2000),
         )
 
         deps = derived.prior.dependencies()
@@ -244,10 +211,7 @@ def test_flux_ratio_via_parameter_ref(self):
         f_weak = Flux(
             'NII_6549',
             prior=TruncatedNormal(
-                loc=f_strong / 2.95,
-                scale=0.1,
-                low=f_strong / 4.0,
-                high=f_strong / 2.0,
+                loc=f_strong / 2.95, scale=0.1, low=f_strong / 4.0, high=f_strong / 2.0
             ),
         )
 
@@ -263,17 +227,11 @@ def test_flux_ratio_via_parameter_ref(self):
     def test_flux_chain(self):
         """Three-line flux chain: Ha → [NII]6585 → [NII]6549."""
         f_ha = Flux('Ha', prior=Uniform(0, 20))
-        f_nii_s = Flux(
-            'NII_6585',
-            prior=Uniform(low=0, high=f_ha * 2),
-        )
+        f_nii_s = Flux('NII_6585', prior=Uniform(low=0, high=f_ha * 2))
         f_nii_w = Flux(
             'NII_6549',
             prior=TruncatedNormal(
-                loc=f_nii_s / 2.95,
-                scale=0.05,
-                low=f_nii_s / 4.0,
-                high=f_nii_s / 2.0,
+                loc=f_nii_s / 2.95, scale=0.05, low=f_nii_s / 4.0, high=f_nii_s / 2.0
             ),
         )
 
@@ -316,27 +274,15 @@ def _make_three_component_config(self):
         z_blr = Redshift(
             'z_blr',
             prior=TruncatedNormal(
-                loc=z_nlr,
-                scale=0.002,
-                low=z_nlr - 0.01,
-                high=z_nlr + 0.01,
+                loc=z_nlr, scale=0.002, low=z_nlr - 0.01, high=z_nlr + 0.01
             ),
         )
-        z_out = Redshift(
-            'z_out',
-            prior=Uniform(low=z_nlr - 0.02, high=z_nlr),
-        )
+        z_out = Redshift('z_out', prior=Uniform(low=z_nlr - 0.02, high=z_nlr))
 
         # -- FWHM hierarchy --
         fwhm_narrow = FWHM('fwhm_narrow', prior=Uniform(50, 500))
-        fwhm_broad = FWHM(
-            'fwhm_broad',
-            prior=Uniform(low=fwhm_narrow + 200, high=5000),
-        )
-        fwhm_out = FWHM(
-            'fwhm_out',
-            prior=Uniform(low=fwhm_broad, high=8000),
-        )
+        fwhm_broad = FWHM('fwhm_broad', prior=Uniform(low=fwhm_narrow + 200, high=5000))
+        fwhm_out = FWHM('fwhm_out', prior=Uniform(low=fwhm_broad, high=8000))
 
         # -- Flux with doublet ratio --
         f_ha_n = Flux('Ha_n', prior=Uniform(0, 10))
@@ -346,25 +292,40 @@ def _make_three_component_config(self):
         f_nii_w = Flux(
             'NII_w',
             prior=TruncatedNormal(
-                loc=f_nii_s / 2.95,
-                scale=0.1,
-                low=f_nii_s / 4.0,
-                high=f_nii_s / 2.0,
+                loc=f_nii_s / 2.95, scale=0.1, low=f_nii_s / 4.0, high=f_nii_s / 2.0
             ),
         )
 
         lc = LineConfiguration()
 
         # Narrow lines
-        lc.add_line('Ha', 6564.61 * u.AA, redshift=z_nlr, fwhm_gauss=fwhm_narrow, flux=f_ha_n)
-        lc.add_line('NII_6585', 6585.27 * u.AA, redshift=z_nlr, fwhm_gauss=fwhm_narrow, flux=f_nii_s)
-        lc.add_line('NII_6549', 6549.86 * u.AA, redshift=z_nlr, fwhm_gauss=fwhm_narrow, flux=f_nii_w)
+        lc.add_line(
+            'Ha', 6564.61 * u.AA, redshift=z_nlr, fwhm_gauss=fwhm_narrow, flux=f_ha_n
+        )
+        lc.add_line(
+            'NII_6585',
+            6585.27 * u.AA,
+            redshift=z_nlr,
+            fwhm_gauss=fwhm_narrow,
+            flux=f_nii_s,
+        )
+        lc.add_line(
+            'NII_6549',
+            6549.86 * u.AA,
+            redshift=z_nlr,
+            fwhm_gauss=fwhm_narrow,
+            flux=f_nii_w,
+        )
 
         # Broad lines
-        lc.add_line('Ha', 6564.61 * u.AA, redshift=z_blr, fwhm_gauss=fwhm_broad, flux=f_ha_b)
+        lc.add_line(
+            'Ha', 6564.61 * u.AA, redshift=z_blr, fwhm_gauss=fwhm_broad, flux=f_ha_b
+        )
 
         # Outflow lines
-        lc.add_line('Ha', 6564.61 * u.AA, redshift=z_out, fwhm_gauss=fwhm_out, flux=f_ha_out)
+        lc.add_line(
+            'Ha', 6564.61 * u.AA, redshift=z_out, fwhm_gauss=fwhm_out, flux=f_ha_out
+        )
 
         return lc
 
@@ -463,20 +424,13 @@ def test_fwhm_depends_on_two_parents(self):
         """A FWHM with low from one token and high from another."""
         lower = FWHM('lower', prior=Uniform(50, 200))
         upper = FWHM('upper', prior=Uniform(800, 2000))
-        mid = FWHM(
-            'mid',
-            prior=Uniform(low=lower + 50, high=upper - 50),
-        )
+        mid = FWHM('mid', prior=Uniform(low=lower + 50, high=upper - 50))
 
         deps = mid.prior.dependencies()
         assert lower in deps
         assert upper in deps
 
-        named_priors = {
-            'lower': lower.prior,
-            'upper': upper.prior,
-            'mid': mid.prior,
-        }
+        named_priors = {'lower': lower.prior, 'upper': upper.prior, 'mid': mid.prior}
         param_to_name = {lower: 'lower', upper: 'upper', mid: 'mid'}
         order = topological_sort(named_priors, param_to_name)
         assert order.index('lower') < order.index('mid')
@@ -489,9 +443,7 @@ def test_diamond_with_convergent_child(self):
         c = FWHM('c', prior=Uniform(low=a, high=b))
         d = FWHM('d', prior=Uniform(low=a + 50, high=b - 50))
 
-        named_priors = {
-            'a': a.prior, 'b': b.prior, 'c': c.prior, 'd': d.prior,
-        }
+        named_priors = {'a': a.prior, 'b': b.prior, 'c': c.prior, 'd': d.prior}
         param_to_name = {a: 'a', b: 'b', c: 'c', d: 'd'}
         order = topological_sort(named_priors, param_to_name)
 
@@ -543,10 +495,7 @@ def test_nested_ref_in_truncated_normal(self):
         derived = FWHM(
             'derived',
             prior=TruncatedNormal(
-                loc=base * 1.5 + 50,
-                scale=30.0,
-                low=base + 20,
-                high=base * 3,
+                loc=base * 1.5 + 50, scale=30.0, low=base + 20, high=base * 3
             ),
         )
         context = {base: 200.0}
@@ -595,10 +544,7 @@ def test_flux_ratio_serialization(self):
         f_weak = Flux(
             'f_w',
             prior=TruncatedNormal(
-                loc=f_strong / 2.95,
-                scale=0.1,
-                low=f_strong / 4.0,
-                high=f_strong / 2.0,
+                loc=f_strong / 2.95, scale=0.1, low=f_strong / 4.0, high=f_strong / 2.0
             ),
         )
 
@@ -651,10 +597,7 @@ class TestEndToEndDeepDependencies:
     def test_narrow_broad_model_respects_ordering(self):
         """Verify sampled broad FWHM > sampled narrow FWHM + offset."""
         fwhm_narrow = FWHM('fwhm_narrow', prior=Uniform(50, 300))
-        fwhm_broad = FWHM(
-            'fwhm_broad',
-            prior=Uniform(low=fwhm_narrow + 200, high=3000),
-        )
+        fwhm_broad = FWHM('fwhm_broad', prior=Uniform(low=fwhm_narrow + 200, high=3000))
 
         lc = LineConfiguration()
         z = Redshift('z', prior=Uniform(-0.005, 0.005))
@@ -709,10 +652,7 @@ def test_redshift_hierarchy_model(self):
         z_nlr = Redshift(
             'z_nlr',
             prior=TruncatedNormal(
-                loc=z_sys,
-                scale=0.001,
-                low=z_sys - 0.003,
-                high=z_sys + 0.003,
+                loc=z_sys, scale=0.001, low=z_sys - 0.003, high=z_sys + 0.003
             ),
         )
 

tests/test_disperser_base.py

@@ -6,7 +6,6 @@
 from unite.disperser.base import Disperser, FluxScale, PixOffset, RScale
 from unite.prior import Fixed, TruncatedNormal, Uniform
 
-
 # ---------------------------------------------------------------------------
 # Calibration token construction
 # ---------------------------------------------------------------------------
@@ -120,9 +119,7 @@ def test_has_calibration_params_false(self):
         from unite.disperser.generic import GenericDisperser
 
         d = GenericDisperser(
-            R_func=lambda w: w * 0 + 1000,
-            dlam_dpix_func=lambda w: w / 1000,
-            unit=u.AA,
+            R_func=lambda w: w * 0 + 1000, dlam_dpix_func=lambda w: w / 1000, unit=u.AA
         )
         assert not d.has_calibration_params