Repository: gwkokab/gwkokab

Issue:
Debug flag to flush information from jitable and non-jitable code
### Check for existing issues

- [X] Completed

### Description

Expose a debug flag which prints values. Relevant resources,

* [The checkify transformation](https://jax.readthedocs.io/en/latest/debugging/checkify_guide.html).
* [Debugging runtime values](https://jax.readthedocs.io/en/latest/debugging/index.html)

Discussion:

Solution patch:
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
index 1efc0fa8..e72ceb4b 100644
--- a/.github/dependabot.yml
+++ b/.github/dependabot.yml
@@ -8,6 +8,10 @@ updates:
   - package-ecosystem: "pip" # See documentation for possible values
     directory: "/" # Location of package manifests
     schedule:
-      interval: "weekly"
+      interval: "monthly"
     ignore:
       - dependency-name: "numpy"
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "monthly"
diff --git a/.github/workflows/auto-assign.yml b/.github/workflows/auto-assign.yml
index 09411e65..8f8e9d94 100644
--- a/.github/workflows/auto-assign.yml
+++ b/.github/workflows/auto-assign.yml
@@ -12,7 +12,7 @@ jobs:
       pull-requests: write
     steps:
       - name: "Auto-assign issue"
-        uses: pozil/auto-assign-issue@v1
+        uses: pozil/auto-assign-issue@v2
         with:
           repo-token: ${{ secrets.GITHUB_TOKEN }}
           assignees: Qazalbash
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 8cd28148..e0e8bc68 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -24,9 +24,9 @@ jobs:
         python-version: [ 3.11 ]
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
@@ -48,9 +48,9 @@ jobs:
         python-version: [ 3.11 ]
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index 979eaa03..472630ca 100644
--- a/.github/workflows/python-publish.yml
+++ b/.github/workflows/python-publish.yml
@@ -20,9 +20,9 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Python
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v5
         with:
           python-version: "3.x"
       - name: Install dependencies
@@ -32,7 +32,7 @@ jobs:
       - name: Build package
         run: python -m build
       - name: Publish package
-        uses: pypa/gh-action-pypi-publish@v1.8.11
+        uses: pypa/gh-action-pypi-publish@v1.10.0
         with:
           user: __token__
           password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/.github/workflows/update_all_top_ranking_issues.yml b/.github/workflows/update_all_top_ranking_issues.yml
index c5b4f76a..a40a66d3 100644
--- a/.github/workflows/update_all_top_ranking_issues.yml
+++ b/.github/workflows/update_all_top_ranking_issues.yml
@@ -13,8 +13,8 @@ jobs:
     runs-on: ubuntu-latest
     if: github.repository_owner == 'gwkokab'
     steps:
-      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4
-      - uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # v5
+      - uses: actions/checkout@ee0669bd1cc54295c223e0bb666b733df41de1c5 # v4
+      - uses: actions/setup-python@e9aba2c848f5ebd159c070c61ea2c4e2b122355e # v5
         with:
           python-version: "3.11"
           architecture: "x64"
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6f62d84c..7d221137 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,7 @@ Best view [here](https://gwkokab.readthedocs.io/en/latest/changelog.html).
   * `ComponentMassesToTotalMassAndMassRatio`
   * `SourceMassAndRedshiftToDetectedMassAndRedshift`
 * scaling factor in vt, during inference and synthetic data generation.
+* Debug mode for inference.
 
 **New Design**
 
diff --git a/pyproject.toml b/pyproject.toml
index 02120416..5279bbfb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,14 +39,14 @@ dependencies = [
     "h5py==3.11.0",
     "jax==0.4.31",
     "jaxlib==0.4.31",
-    "jaxtyping==0.2.33",
+    "jaxtyping==0.2.34",
     "matplotlib==3.9.2",
     "mplcursors==0.5.3",
     "numpy<2",
     "numpyro==0.15.2",
-    "plotly==5.23.0",
+    "plotly==5.24.0",
     "polars==1.5.0",
-    "rich==13.7.1",
+    "rich==13.8.0",
     "RIFT==0.0.15.11",
     "seaborn==0.13.2",
     "typing_extensions==4.12.2",
@@ -64,7 +64,7 @@ dev = [
     "pre-commit==3.8.0",
     "pytest==8.3.2",
     "ruff==0.6.2",
-    "setuptools==73.0.1",
+    "setuptools==74.0.0",
 ]
 
 docs = [
diff --git a/scripts/requirements.txt b/scripts/requirements.txt
index 69246c02..20063650 100644
--- a/scripts/requirements.txt
+++ b/scripts/requirements.txt
@@ -1,5 +1,5 @@
-mypy==1.11.1
-PyGithub==2.3.0
+mypy==1.11.2
+PyGithub==2.4.0
 pytz==2024.1
 ruff==0.6.2
 typer==0.12.5
diff --git a/src/gwkokab/__init__.py b/src/gwkokab/__init__.py
index 399d8b04..12428298 100644
--- a/src/gwkokab/__init__.py
+++ b/src/gwkokab/__init__.py
@@ -21,6 +21,7 @@
 
 
 from . import (
+    debug as debug,
     errors as errors,
     inference as inference,
     models as models,
diff --git a/src/gwkokab/_src/debug/__init__.py b/src/gwkokab/_src/debug/__init__.py
new file mode 100644
index 00000000..10a5904a
--- /dev/null
+++ b/src/gwkokab/_src/debug/__init__.py
@@ -0,0 +1,21 @@
+# Copyright 2023 The GWKokab Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from ._debug import (
+    debug as debug,
+    debug_flush as debug_flush,
+    debug_mode as debug_mode,
+    enable_debugging as enable_debugging,
+)
diff --git a/src/gwkokab/_src/debug/_debug.py b/src/gwkokab/_src/debug/_debug.py
new file mode 100644
index 00000000..8869e631
--- /dev/null
+++ b/src/gwkokab/_src/debug/_debug.py
@@ -0,0 +1,72 @@
+# Copyright 2023 The GWKokab Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing_extensions import Callable
+
+import jax
+
+
+DEBUG: bool = False
+
+
+def enable_debugging() -> None:
+    r"""Enables debugging mode."""
+    global DEBUG
+    DEBUG = True
+
+
+def debug_mode() -> bool:
+    r"""Returns the current debugging mode."""
+    return DEBUG
+
+
+def debug_flush(fmt_str: str, **kwargs) -> None:
+    """Flushes the debug message to the console.
+
+    .. note::
+        This function will only print message in debug mode.
+
+    .. code:: python
+
+        >>> debug_flush("Hello, {name}!", name="world")
+        Hello, world!
+
+    :param fmt_str: The format string for the debug message.
+    """
+    if debug_mode():
+        jax.debug.print("\033[32mGWKokab " + fmt_str, **kwargs)
+
+
+def debug(func) -> Callable:
+    r"""A decorator that prints the arguments and return value of a function when
+    :code:`DEBUG` is :code:`True`.
+
+    :param func: The function to be decorated.
+    :return: A wrapper function that prints the arguments and return value of the
+    decorated function.
+    """
+
+    def wrapper(*args, **kwargs):
+        result = func(*args, **kwargs)
+        if debug_mode():
+            jax.debug.print(
+                "\033[32mGWKokab Calling {func_name} with arguments {args_fmt} "
+                "and {kwargs_fmt} returned {result_fmt}",
+                func_name=func.__name__,
+                result_fmt=result,
+            )
+        return result
+
+    return wrapper
diff --git a/src/gwkokab/_src/inference/poissonlikelihood.py b/src/gwkokab/_src/inference/poissonlikelihood.py
index b5964fa1..ec898780 100644
--- a/src/gwkokab/_src/inference/poissonlikelihood.py
+++ b/src/gwkokab/_src/inference/poissonlikelihood.py
@@ -27,6 +27,7 @@
     Uniform,
 )
 
+from ..debug import debug_flush, debug_mode
 from ..models.utils import JointDistribution
 from ..parameters import DEFAULT_PRIORS, Parameter
 from .bake import Bake
@@ -251,6 +252,7 @@ def exp_rate_integral(self, model: Distribution) -> Array:
         else:
             raise ValueError("Invalid VT method.")
         vt_value *= self.time * self.scale_factor
+        debug_flush("exp_rate_integral: {vt_value}", vt_value=vt_value)
         return vt_value
 
     def log_likelihood_single_rate(self, x: Array, data: dict) -> Array:
@@ -272,11 +274,13 @@ def log_likelihood_single_rate(self, x: Array, data: dict) -> Array:
 
         log_likelihood += data["N"] * log_rate
 
+        debug_flush("model_log_likelihood: {mll}", mll=log_likelihood)
+
         rate = jnp.exp(log_rate)  # rate = e^log_rate
 
         return log_likelihood - rate * lax.cond(
             jnp.isinf(log_likelihood),
-            lambda m: 0.0,
+            lambda _: 0.0,
             lambda m: self.exp_rate_integral(m),
             model,
         )
@@ -305,6 +309,8 @@ def log_likelihood_multi_rate(self, x: Array, data: dict) -> Array:
 
         log_likelihood += data["N"] * log_sum_of_rates
 
+        debug_flush("model_log_likelihood: {mll}", mll=log_likelihood)
+
         rates = list(jnp.exp(log_rates))  # rates = e^log_rates
 
         expected_rates = jtr.reduce(
@@ -324,8 +330,19 @@ def log_likelihood(self, x: Array, data: dict) -> Array:
         :return: Log likelihood value for the given parameters.
         """
         if self.is_multi_rate_model:
-            return self.log_likelihood_multi_rate(x, data)
-        return self.log_likelihood_single_rate(x, data)
+            log_likelihood = self.log_likelihood_multi_rate(x, data)
+            if debug_mode():
+                pdict = {name: x[..., i] for name, i in self.variables_index.items()}
+                pdict["log_rates"] = x[..., 0 : self.total_pop]
+                debug_flush("parameters: {pdict}", pdict=pdict)
+        else:
+            log_likelihood = self.log_likelihood_single_rate(x, data)
+            if debug_mode():
+                pdict = {name: x[..., i] for name, i in self.variables_index.items()}
+                pdict["log_rate"] = x[..., 0]
+                debug_flush("parameters: {pdict}", pdict=pdict)
+        debug_flush("log_likelihood: {ll}", ll=log_likelihood)
+        return log_likelihood
 
     def log_posterior(self, x: Array, data: Optional[dict] = None) -> Array:
         r"""The likelihood function for the inhomogeneous Poisson process.
@@ -338,6 +355,7 @@ def log_posterior(self, x: Array, data: Optional[dict] = None) -> Array:
         :return: Log likelihood value for the given parameters.
         """
         log_prior = self.priors.log_prob(x)
+        debug_flush("log_prior: {lp}", lp=log_prior)
         return log_prior + lax.cond(
             jnp.isinf(log_prior),
             lambda x_, d_: 0.0,
diff --git a/src/gwkokab/debug.py b/src/gwkokab/debug.py
new file mode 100644
index 00000000..096f6ccd
--- /dev/null
+++ b/src/gwkokab/debug.py
@@ -0,0 +1,20 @@
+# Copyright 2023 The GWKokab Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from gwkokab._src.debug import (
+    debug as debug,
+    debug_mode as debug_mode,
+    enable_debugging as enable_debugging,
+)


Test patch:
diff --git a/tests/test_debug.py b/tests/test_debug.py
new file mode 100644
index 00000000..72c741fd
--- /dev/null
+++ b/tests/test_debug.py
@@ -0,0 +1,22 @@
+# Copyright 2023 The GWKokab Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from gwkokab.debug import debug_mode, enable_debugging
+
+
+def test_enable_debugging():
+    assert debug_mode() is False
+    enable_debugging()
+    assert debug_mode() is True
diff --git a/tests/test_math.py b/tests/test_math.py
index 8bce102d..4b9a0db0 100644
--- a/tests/test_math.py
+++ b/tests/test_math.py
@@ -1,3 +1,18 @@
+# Copyright 2023 The GWKokab Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
 import numpy as np
 import pytest
 from gwkokab.utils.math import (