lint

tk2lab · Feb 5, 2023 · af1320a · af1320a
1 parent 3055474
commit af1320a
Show file tree

Hide file tree

Showing 22 changed files with 766 additions and 179 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -18,9 +18,12 @@ python = ">=3.8,<3.12"
 
 [tool.isort]
 profile = "black"
+multi_line_output = 3
+force_grid_wrap = 1
 
-[tool.black]
-line-length = 79
+[tool.flake8]
+max-line-length = 88
+extend-ignore = ["E203"]
 
 [tool.pytest.ini_options]
 markers = [
@@ -51,11 +54,12 @@ skip_install = True
 deps =
     isort
     black
-    autoflake
+    flake8
+    pyproject-flake8
 commands =
-    isort --sl src
-    autoflake -ri --remove-all-unused-imports --ignore-init-module-imports src
+    isort src
     black src
+    pflake8 src
 
 [testenv:eval_jax]
 deps =
@@ -96,4 +100,15 @@ deps =
     pandas
 commands =
     {envpython} -m pytest tests/test_tensorflow.py {posargs}
+
+[testenv:notebook]
+basepython = python3.11
+deps =
+    jax[cuda]
+    jupyterlab
+    jupyterlab_nvdashboard
+install_command =
+    pip install --upgrade -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html {opts} {packages}
+commands =
+    jupyter lab --no-browser
 """
diff --git a/src/logbesselk/jax/__init__.py b/src/logbesselk/jax/__init__.py
@@ -1,5 +1,17 @@
-from .integral import bessel_ke
-from .integral import bessel_kratio
-from .integral import log_abs_deriv_bessel_k
-from .integral import log_bessel_k
-from .misc import sign_deriv_bessel_k
+from .integral import (
+    bessel_ke,
+    bessel_kratio,
+    log_abs_deriv_bessel_k,
+    log_bessel_k,
+)
+from .misc import (
+    sign_deriv_bessel_k,
+)
+
+__all__ = [
+    "bessel_ke",
+    "bessel_kratio",
+    "log_abs_deriv_bessel_k",
+    "log_bessel_k",
+    "sign_deriv_bessel_k",
+]
diff --git a/src/logbesselk/jax/asymptotic.py b/src/logbesselk/jax/asymptotic.py
@@ -3,13 +3,19 @@
 import jax.lax as lax
 import jax.numpy as jnp
 
-from .math import fabs
-from .math import log
-from .math import sqrt
-from .math import square
-from .utils import epsilon
-from .utils import result_type
-from .wrap import wrap_log_bessel_k
+from .math import (
+    fabs,
+    log,
+    sqrt,
+    square,
+)
+from .utils import (
+    epsilon,
+    result_type,
+)
+from .wrap import (
+    wrap_log_bessel_k,
+)
 
 __all__ = [
     "log_bessel_k",

diff --git a/src/logbesselk/jax/cfraction.py b/src/logbesselk/jax/cfraction.py
@@ -2,14 +2,22 @@
 
 import jax.lax as lax
 
-from .math import fabs
-from .math import fround
-from .math import log
-from .math import square
-from .misc import log_bessel_recurrence
-from .utils import epsilon
-from .utils import result_type
-from .wrap import wrap_log_bessel_k
+from .math import (
+    fabs,
+    fround,
+    log,
+    square,
+)
+from .misc import (
+    log_bessel_recurrence,
+)
+from .utils import (
+    epsilon,
+    result_type,
+)
+from .wrap import (
+    wrap_log_bessel_k,
+)
 
 __all__ = [
     "log_bessel_k",

diff --git a/src/logbesselk/jax/integral.py b/src/logbesselk/jax/integral.py
@@ -1,22 +1,30 @@
 import jax.lax as lax
 import jax.numpy as jnp
-from jax import grad
-
-from .math import cosh
-from .math import is_finite
-from .math import log
-from .math import log_cosh
-from .math import log_sinh
-from .math import maximum
-from .math import square
-from .utils import epsilon
-from .utils import extend
-from .utils import find_zero
-from .utils import log_integrate
-from .utils import result_type
-from .wrap import wrap_bessel_ke
-from .wrap import wrap_bessel_kratio
-from .wrap import wrap_log_abs_deriv_bessel_k
+from jax import (
+    grad,
+)
+
+from .math import (
+    cosh,
+    is_finite,
+    log,
+    log_cosh,
+    log_sinh,
+    maximum,
+    square,
+)
+from .utils import (
+    epsilon,
+    extend,
+    find_zero,
+    log_integrate,
+    result_type,
+)
+from .wrap import (
+    wrap_bessel_ke,
+    wrap_bessel_kratio,
+    wrap_log_abs_deriv_bessel_k,
+)
 
 __all__ = [
     "log_bessel_k",
@@ -58,6 +66,9 @@ def func(t):
             out += n * log_cosh(t)
         return out
 
+    def mfunc(t):
+        return func(t) - th
+
     scale = 0.1
     tol = 1.0
     max_iter = 10
@@ -85,7 +96,6 @@ def func(t):
     tp = find_zero(deriv, start, delta, tol, max_iter)
 
     th = func(tp) + log(eps) - tol
-    mfunc = lambda t: func(t) - th
     tpl = maximum(tp - bins * eps, zero)
     tpr = maximum(tp + bins * eps, tp * (1 + bins * eps))
     mfunc_at_zero_is_negative = out_is_finite & (mfunc(zero) < 0)

diff --git a/src/logbesselk/jax/math.py b/src/logbesselk/jax/math.py
@@ -2,24 +2,32 @@
 
 import jax
 from jax.numpy import abs as fabs
-from jax.numpy import cosh
-from jax.numpy import exp
-from jax.numpy import expm1
-from jax.numpy import inf
+from jax.numpy import (
+    cosh,
+    exp,
+    expm1,
+    inf,
+)
 from jax.numpy import isfinite as is_finite
-from jax.numpy import log
-from jax.numpy import log1p
+from jax.numpy import (
+    log,
+    log1p,
+)
 from jax.numpy import logaddexp as log_add_exp
-from jax.numpy import maximum
-from jax.numpy import nan
+from jax.numpy import (
+    maximum,
+    nan,
+)
 from jax.numpy import round as fround
-from jax.numpy import sign
-from jax.numpy import sinc
-from jax.numpy import sinh
-from jax.numpy import sqrt
-from jax.numpy import square
-from jax.numpy import tanh
-from jax.numpy import where
+from jax.numpy import (
+    sign,
+    sinc,
+    sinh,
+    sqrt,
+    square,
+    tanh,
+    where,
+)
 
 __all__ = [
     "fabs",

diff --git a/src/logbesselk/jax/misc.py b/src/logbesselk/jax/misc.py
@@ -1,9 +1,13 @@
 import jax.lax as lax
 
-from .math import log
-from .math import log_add_exp
-from .math import sign
-from .utils import result_type
+from .math import (
+    log,
+    log_add_exp,
+    sign,
+)
+from .utils import (
+    result_type,
+)
 
 __all__ = [
     "sign_deriv_bessel_k",

diff --git a/src/logbesselk/jax/sca.py b/src/logbesselk/jax/sca.py
@@ -2,15 +2,23 @@
 
 from .asymptotic import log_bessel_k_naive as log_k_large_v
 from .cfraction import log_bessel_ku as log_ku_large_x
-from .math import fround
-from .math import is_finite
-from .math import log
-from .misc import log_bessel_recurrence
+from .math import (
+    fround,
+    is_finite,
+    log,
+)
+from .misc import (
+    log_bessel_recurrence,
+)
 from .series import log_bessel_ku as log_ku_small_x
-from .utils import result_type
-from .wrap import wrap_bessel_ke
-from .wrap import wrap_bessel_kratio
-from .wrap import wrap_log_bessel_k
+from .utils import (
+    result_type,
+)
+from .wrap import (
+    wrap_bessel_ke,
+    wrap_bessel_kratio,
+    wrap_log_bessel_k,
+)
 
 __all__ = [
     "log_bessel_k",

diff --git a/src/logbesselk/jax/series.py b/src/logbesselk/jax/series.py
@@ -1,17 +1,25 @@
 import jax.lax as lax
 
-from .math import cosh
-from .math import exp
-from .math import fabs
-from .math import fround
-from .math import log
-from .math import sinc
-from .math import sinhc
-from .math import square
-from .misc import log_bessel_recurrence
-from .utils import epsilon
-from .utils import result_type
-from .wrap import wrap_log_bessel_k
+from .math import (
+    cosh,
+    exp,
+    fabs,
+    fround,
+    log,
+    sinc,
+    sinhc,
+    square,
+)
+from .misc import (
+    log_bessel_recurrence,
+)
+from .utils import (
+    epsilon,
+    result_type,
+)
+from .wrap import (
+    wrap_log_bessel_k,
+)
 
 __all__ = [
     "log_bessel_k",
@@ -33,7 +41,6 @@ def log_bessel_k(v, x):
 
 
 def log_bessel_ku(u, x):
-
     def cond(args):
         ku, kn, i, p, q, r, s = args
         update = fabs(r * s) > eps * fabs(ku)

diff --git a/src/logbesselk/jax/utils.py b/src/logbesselk/jax/utils.py
@@ -2,9 +2,11 @@
 import jax.lax as lax
 import jax.numpy as jnp
 
-from .math import exp
-from .math import fabs
-from .math import log
+from .math import (
+    exp,
+    fabs,
+    log,
+)
 
 __all__ = [
     "result_type",

diff --git a/src/logbesselk/jax/wrap.py b/src/logbesselk/jax/wrap.py
@@ -3,12 +3,16 @@
 import jax
 import jax.lax as lax
 
-from .math import exp
-from .math import fabs
-from .math import inf
-from .math import nan
-from .math import sign
-from .utils import result_type
+from .math import (
+    exp,
+    fabs,
+    inf,
+    nan,
+    sign,
+)
+from .utils import (
+    result_type,
+)
 
 __all__ = [
     "wrap_log_bessel_k",

diff --git a/src/logbesselk/tensorflow/__init__.py b/src/logbesselk/tensorflow/__init__.py
@@ -1,5 +1,17 @@
-from .integral import bessel_ke
-from .integral import bessel_kratio
-from .integral import log_abs_deriv_bessel_k
-from .integral import log_bessel_k
-from .misc import sign_deriv_bessel_k
+from .integral import (
+    bessel_ke,
+    bessel_kratio,
+    log_abs_deriv_bessel_k,
+    log_bessel_k,
+)
+from .misc import (
+    sign_deriv_bessel_k,
+)
+
+__all__ = [
+    "bessel_ke",
+    "bessel_kratio",
+    "log_abs_deriv_bessel_k",
+    "log_bessel_k",
+    "sign_deriv_bessel_k",
+]