Expose activity energy gradient for potential advanced use

docs/api/Gradients.md

@@ -1,3 +1,7 @@
 # Gradients
 
+::: jpc.neg_activity_grad
+
+---
+
 ::: jpc.compute_pc_param_grads

jpc/__init__.py

@@ -5,6 +5,7 @@
     init_activities_from_gaussian as init_activities_from_gaussian,
     init_activities_with_amort as init_activities_with_amort,
     pc_energy_fn as pc_energy_fn,
+    neg_activity_grad as neg_activity_grad,
     solve_pc_activities as solve_pc_activities,
     compute_pc_param_grads as compute_pc_param_grads,
     linear_equilib_energy as linear_equilib_energy

jpc/_core/__init__.py

@@ -4,6 +4,9 @@
     init_activities_with_amort as init_activities_with_amort
 )
 from ._energies import pc_energy_fn as pc_energy_fn
+from ._grads import (
+    neg_activity_grad as neg_activity_grad,
+    compute_pc_param_grads as compute_pc_param_grads
+)
 from ._infer import solve_pc_activities as solve_pc_activities
-from ._grads import compute_pc_param_grads as compute_pc_param_grads
 from ._analytical import linear_equilib_energy_batch as linear_equilib_energy

jpc/_core/_grads.py

@@ -8,7 +8,7 @@
 from ._energies import pc_energy_fn
 
 
-def _neg_activity_grad(
+def neg_activity_grad(
         t: float | int,
         activities: PyTree[ArrayLike],
         args: Tuple[PyTree[Callable], ArrayLike, Optional[ArrayLike]],
@@ -54,7 +54,7 @@ def compute_pc_param_grads(
         y: ArrayLike,
         x: Optional[ArrayLike] = None
 ) -> PyTree[Array]:
-    """Computes the gradient of the energy with respect to network parameters $\partial \mathcal{F} / \partial θ$.
+    """Computes the gradient of the energy with respect to model parameters $\partial \mathcal{F} / \partial θ$.
 
     **Main arguments:**
 

jpc/_core/_infer.py

@@ -2,7 +2,7 @@
 
 from jaxtyping import PyTree, ArrayLike, Array
 from typing import Callable, Optional
-from ._grads import _neg_activity_grad
+from ._grads import neg_activity_grad
 from diffrax import (
     AbstractSolver,
     AbstractStepSizeController,
@@ -65,7 +65,7 @@ def solve_pc_activities(
 
     """
     sol = diffeqsolve(
-        terms=ODETerm(_neg_activity_grad),
+        terms=ODETerm(neg_activity_grad),
         solver=solver,
         t0=0,
         t1=t1,