Fix sample and log_prob docstring

thermox/prob.py

@@ -28,10 +28,10 @@ def log_prob_identity_diffusion(
     Preprocessing (diagonalisation) costs O(d^3) and evaluation then costs O(T * d^2).
 
     Args:
-        ts: array-like, times at which samples are collected. Includes time for x0.
-        xs: initial state of the process.
-        A: drift matrix (Array or thermox.ProcessedDriftMatrix).
-        b: drift displacement vector.
+        ts: Times at which samples are collected. Includes time for x0.
+        xs: Initial state of the process.
+        A: Drift matrix (Array or thermox.ProcessedDriftMatrix).
+        b: Drift displacement vector.
     Returns:
         Scalar log probability of given xs.
     """
@@ -94,17 +94,22 @@ def log_prob(
 
     Assumes x(t_0) is given deterministically.
 
-    Preprocessing (diagonalisation) costs O(d^3) and evaluation then costs O(T * d^2).
+    Preprocessing (diagonalisation) costs O(d^3) and evaluation then costs O(T * d^2),
+    where T=len(ts).
+
+    By default, this function does the preprocessing on A and D before the evaluation.
+    However, the preprocessing can be done externally using thermox.preprocess
+    the output of which can be used as A and D here, this will skip the preprocessing.
 
     Args:
-        ts: array-like, times at which samples are collected. Includes time for x0.
-        xs: initial state of the process.
-        A: drift matrix (Array or thermox.ProcessedDriftMatrix).
+        ts: Times at which samples are collected. Includes time for x0.
+        xs: Initial state of the process.
+        A: Drift matrix (Array or thermox.ProcessedDriftMatrix).
             Note : If a thermox.ProcessedDriftMatrix instance is used as input,
-            must be transformed drift matrix, A_y, given by thermox.preprocess,
-            not thermox.preprocess_drift_matrix.
-        b: drift displacement vector.
-        D: diffusion matrix (Array or thermox.ProcessedDiffusionMatrix).
+            must be the transformed drift matrix, A_y, given by thermox.preprocess,
+            not thermox.utils._preprocess_drift_matrix.
+        b: Drift displacement vector.
+        D: Diffusion matrix (Array or thermox.ProcessedDiffusionMatrix).
 
     Returns:
         Scalar log probability of given xs.

thermox/sampler.py

@@ -88,17 +88,21 @@ def sample(
 
     by using exact diagonalization.
 
-    Preprocessing (diagonalisation) costs O(d^3) and sampling costs O(T * d^2)
+    Preprocessing (diagonalisation) costs O(d^3) and sampling costs O(T * d^2),
     where T=len(ts).
 
+    By default, this function does the preprocessing on A and D before the evaluation.
+    However, the preprocessing can be done externally using thermox.preprocess
+    the output of which can be used as A and D here, this will skip the preprocessing.
+
     Args:
         key: Jax PRNGKey.
         ts: Times at which samples are collected. Includes time for x0.
         x0: Initial state of the process.
         A: Drift matrix (Array or thermox.ProcessedDriftMatrix).
             Note : If a thermox.ProcessedDriftMatrix instance is used as input,
-            must be transformed drift matrix, A_y, given by thermox.preprocess,
-            not thermox.preprocess_drift_matrix.
+            must be the transformed drift matrix, A_y, given by thermox.preprocess,
+            not thermox.utils._preprocess_drift_matrix.
         b: Drift displacement vector.
         D: Diffusion matrix (Array or thermox.ProcessedDiffusionMatrix).