Refactor utility to apply a PointFunc over a dataset

Author: ricardoV94

pymc/backends/arviz.py

@@ -32,12 +32,15 @@
 from arviz.data.base import CoordSpec, DimSpec, dict_to_dataset, requires
 from pytensor.graph.basic import Constant
 from pytensor.tensor.sharedvar import SharedVariable
+from rich.progress import Console, Progress
+from rich.theme import Theme
+from xarray import Dataset
 
 import pymc
 
 from pymc.model import Model, modelcontext
-from pymc.pytensorf import extract_obs_data
-from pymc.util import get_default_varnames
+from pymc.pytensorf import PointFunc, extract_obs_data
+from pymc.util import default_progress_theme, get_default_varnames
 
 if TYPE_CHECKING:
     from pymc.backends.base import MultiTrace
@@ -637,3 +640,49 @@ def dataset_to_point_list(
     ]
     # use the list of points
     return cast(list[dict[str, np.ndarray]], points), stacked_dims
+
+
+def apply_function_over_dataset(
+    fn: PointFunc,
+    dataset: Dataset,
+    *,
+    output_var_names: Sequence[str],
+    coords,
+    dims,
+    sample_dims: Sequence[str] = ("chain", "draw"),
+    progressbar: bool = True,
+    progressbar_theme: Theme | None = default_progress_theme,
+) -> Dataset:
+    posterior_pts, stacked_dims = dataset_to_point_list(dataset, sample_dims)
+
+    n_pts = len(posterior_pts)
+    out_dict = _DefaultTrace(n_pts)
+    indices = range(n_pts)
+
+    with Progress(console=Console(theme=progressbar_theme)) as progress:
+        task = progress.add_task("Computinng ...", total=n_pts, visible=progressbar)
+        for idx in indices:
+            out = fn(posterior_pts[idx])
+            fn.f.trust_input = True  # If we arrive here the dtypes are valid
+            for var_name, val in zip(output_var_names, out):
+                out_dict.insert(var_name, val, idx)
+
+            progress.advance(task)
+
+    out_trace = out_dict.trace_dict
+    for key, val in out_trace.items():
+        out_trace[key] = val.reshape(
+            (
+                *[len(coord) for coord in stacked_dims.values()],
+                *val.shape[1:],
+            )
+        )
+
+    return dict_to_dataset(
+        out_trace,
+        library=pymc,
+        dims=dims,
+        coords=coords,
+        default_dims=list(sample_dims),
+        skip_event_dims=True,
+    )

pymc/stats/log_density.py

@@ -12,22 +12,16 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 from collections.abc import Sequence
-from typing import cast
+from typing import Literal
 
-from arviz import InferenceData, dict_to_dataset
-from rich.console import Console
-from rich.progress import Progress
-
-import pymc
+from arviz import InferenceData
+from xarray import Dataset
 
 from pymc.backends.arviz import (
-    _DefaultTrace,
+    apply_function_over_dataset,
     coords_and_dims_for_inferencedata,
-    dataset_to_point_list,
 )
 from pymc.model import Model, modelcontext
-from pymc.pytensorf import PointFunc
-from pymc.util import default_progress_theme
 
 __all__ = ("compute_log_likelihood", "compute_log_prior")
 
@@ -117,10 +111,10 @@ def compute_log_density(
     var_names: Sequence[str] | None = None,
     extend_inferencedata: bool = True,
     model: Model | None = None,
-    kind="likelihood",
+    kind: Literal["likelihood", "prior"] = "likelihood",
     sample_dims: Sequence[str] = ("chain", "draw"),
     progressbar=True,
-):
+) -> InferenceData | Dataset:
     """
     Compute elemwise log_likelihood or log_prior of model given InferenceData with posterior group
     """
@@ -163,40 +157,20 @@ def compute_log_density(
             outs=model.logp(vars=vars, sum=False),
             on_unused_input="ignore",
         )
-        elemwise_logdens_fn = cast(PointFunc, elemwise_logdens_fn)
     finally:
         model.rvs_to_values = original_rvs_to_values
         model.rvs_to_transforms = original_rvs_to_transforms
 
-    # Ignore Deterministics
-    posterior_values = posterior[[rv.name for rv in model.free_RVs]]
-    posterior_pts, stacked_dims = dataset_to_point_list(posterior_values, sample_dims)
-
-    n_pts = len(posterior_pts)
-    logdens_dict = _DefaultTrace(n_pts)
-
-    with Progress(console=Console(theme=default_progress_theme)) as progress:
-        task = progress.add_task("Computing log density...", total=n_pts, visible=progressbar)
-        for idx in range(n_pts):
-            logdenss_pts = elemwise_logdens_fn(posterior_pts[idx])
-            for rv_name, rv_logdens in zip(var_names, logdenss_pts):
-                logdens_dict.insert(rv_name, rv_logdens, idx)
-            progress.update(task, advance=1)
-
-    logdens_trace = logdens_dict.trace_dict
-    for key, array in logdens_trace.items():
-        logdens_trace[key] = array.reshape(
-            (*[len(coord) for coord in stacked_dims.values()], *array.shape[1:])
-        )
-
     coords, dims = coords_and_dims_for_inferencedata(model)
-    logdens_dataset = dict_to_dataset(
-        logdens_trace,
-        library=pymc,
+
+    logdens_dataset = apply_function_over_dataset(
+        elemwise_logdens_fn,
+        posterior[[rv.name for rv in model.free_RVs]],
+        output_var_names=var_names,
+        sample_dims=sample_dims,
         dims=dims,
         coords=coords,
-        default_dims=list(sample_dims),
-        skip_event_dims=True,
+        progressbar=progressbar,
     )
 
     if extend_inferencedata: