pangolin 0.0.1__py3-none-any.whl

pangolin/THOUGHTS.md

@@ -0,0 +1,41 @@
+How should indexing work?
+
+I think the best conceptual rule in terms of INTERFACE would be that:
+
+1. Individual indices can be declared "orthogonal" or "vectorized"
+2. All vectorized indices are broadcast against each other
+3. The vectorized dimensions (if any) go at the beginning
+
+The best notation I can think for this is "orthogonal by default"
+
+```python
+x = [5,7]
+y = [2,3]
+z = [6,7]
+
+x = [5,7]; y = [2,3]; z = [6,7]
+
+# 1. All Orthogonal (The Grid) -> 3D Result
+B = A[x,y,z]            # B[i,j,k] = A[x[i], y[j], z[k]]
+
+# 2. All Coupled (The Path) -> 1D Result
+B = A[v[x], v[y], v[z]] # B[i]     = A[x[i], y[i], z[i]]
+
+# 3. Mixed (Path + Grid) -> 2D Result
+# x and z move together (dim 0, 2); y is independent (dim 1)
+B = A[v[x], y, v[z]]    # B[i,j]   = A[x[i], y[j], z[i]]
+
+# 4. Mixed (Grid + Path) -> 2D Result
+# x and y move together (dim 0, 1); z is independent (dim 2)
+B = A[v[x], v[y], z]    # B[i,j]   = A[x[i], y[i], z[j]]
+
+# 5. Standard Slicing -> 3D Result
+B = A[:, ::2, 3::]      # B[i,j,k] = A[i, 2*j, 3+k]
+
+# 6. Coupled Slices (Main Diagonal) -> 1D Result
+B = A[v[:], v[:], v[:]] # B[i]     = A[i, i, i]
+
+# 7. Partial Diagonal (Trace) -> 2D Result
+# Dims 0 and 2 are diagonal; Dim 1 is full/independent
+B = A[v[:], :, v[:]]    # B[i,j]   = A[i, j, i]
+```

pangolin/__init__.py

@@ -0,0 +1,37 @@
+"""
+Pangolin's goal is to be **the world's friendliest probabilistic programming language**.
+
+Quickstart:
+
+- Use :mod:`pangolin.interface` to define probabilistic models in a cheerful interface.
+- Use :mod:`pangolin.jax_backend` to compile probabilistic models into plain JAX functions (to be used independently of the rest of Pangolin).
+- Use :mod:`pangolin.torch_backend` to compile probabilistic models into plain pytorch functions (to be used independently of the rest of Pangolin).
+- Use :mod:`pangolin.blackjax` to easily call `Blackjax <https://blackjax-devs.github.io/blackjax/>`_ to do inference on probabilistic models.
+
+In addition, there are three "private" submodules that end-users would not typically interact with:
+
+- :mod:`pangolin.util` - Various internal utility functions.
+- :mod:`pangolin.dag` - Utilities for interacting with directed acyclic graphs (DAGs).
+- :mod:`pangolin.ir` - The internal representation (IR) for probabilistic models in Pangolin.
+"""
+
+import warnings
+
+with warnings.catch_warnings():
+    # Only ignore DeprecationWarning
+    warnings.simplefilter("ignore", category=DeprecationWarning)
+    import numpyro
+
+
+from pangolin import dag, util, ir, interface, jax_backend, blackjax
+
+__all__ = ["dag", "util", "ir", "interface", "jax_backend"]
+
+try:
+    from pangolin import torch_backend
+
+    __all__.append("torch_backend")
+except ImportError:
+    pass
+
+__all__.append("blackjax")

pangolin/blackjax.py

@@ -0,0 +1,404 @@
+"""
+This module defines a convenient interface to call
+`Blackjax <https://blackjax-devs.github.io/blackjax/>`_
+to do inference. You could of course just call `pangolin.jax_backend.ancestor_log_prob`
+to get a plain jax function and then call Blackjax yourself. But this module abstracts
+away all the details.
+"""
+
+from __future__ import annotations
+from jax import numpy as jnp
+import jax.tree_util
+import numpy as np
+from typing import (
+    Any,
+    Callable,
+    TypeAlias,
+    TYPE_CHECKING,
+    Type,
+    Sequence,
+    List,
+    Optional,
+)
+from pangolin import ir
+from pangolin.ir import Op, RV, ArrayLike
+from numpy.typing import ArrayLike
+from numpyro import distributions as dist
+from jax.scipy import special as jspecial
+from jax import nn as jnn
+from pangolin import dag, util
+from pangolin import jax_backend
+import blackjax
+from jaxtyping import PyTree
+
+__all__ = ["sample", "E", "var", "std", "Calculate", "inf_until_match"]
+
+from pangolin.jax_backend.bijectors import default_bijector_dict
+
+
+def inference_loop(rng_key, kernel, initial_states, num_samples):
+    @jax.jit
+    def one_step(states, rng_key):
+        states, infos = kernel(rng_key, states)
+        return states, (states, infos)
+
+    keys = jax.random.split(rng_key, num_samples)
+    _, (states, infos) = jax.lax.scan(one_step, initial_states, keys)
+
+    return (states, infos)
+
+
+def run_nuts(log_prob, key, initial_state, num_samples):
+    # to do hmc instead:
+    # adapt = blackjax.window_adaptation(blackjax.hmc, log_prob, num_integration_steps=60)
+    # kernel = blackjax.hmc(log_prob, **parameters).step
+
+    sample_key, warmup_key = jax.random.split(key)
+    adapt = blackjax.window_adaptation(blackjax.nuts, log_prob)
+
+    (last_state, parameters), _ = adapt.run(warmup_key, initial_state, num_samples)  # type: ignore
+
+    kernel = blackjax.nuts(log_prob, **parameters).step
+    states, infos = inference_loop(sample_key, kernel, last_state, num_samples)
+    return states.position
+
+
+# TODO: Raise error if given discrete latent variable
+
+
+def sample_flat(
+    vars: list[RV],
+    given_vars: list[RV],
+    given_vals: list,
+    *,
+    niter: int,
+    bijector_dict: Optional[dict] = default_bijector_dict,
+) -> list[jnp.ndarray]:
+    """
+    Given a "flat" specification of an inference problem, do inference using Numpyro. The basic algorithm is:
+
+    1. Use `get_model_flat` to create a Numpyro model and a mapping from RVs to Numpyro variable names.
+    2. Use standard Numpyro routines to do inference using MCMC.
+    3. Use the name mapping to extract samples for all variables listed in `vars`.
+
+    If given is empty, then example samples will (automatically) be drawn using `ancestor_sample_flat`.
+
+    Parameters
+    ----------
+    vars
+        The RVs you want to sample from
+    given
+        The RVs you want to condition on
+    vals
+        The values for the conditioned RVs
+    niter: int, optional
+        The number of iterations / samples to draw
+
+    Returns
+    -------
+    samples: list[jnp.ndarray]
+        Samples for each variable in `vars`.
+
+    Examples
+    --------
+    >>> x = ir.RV(ir.Constant(0.5))
+    >>> y = ir.RV(ir.Normal(), x, x)
+    >>> z = ir.RV(ir.Normal(), y, x)
+    >>> [samps_x, samps_y] = sample_flat([x, y], [z], [3.0], niter=30)
+    >>> samps_x.shape
+    (30,)
+    >>> samps_y.shape
+    (30,)
+    >>> np.allclose(samps_x, 0.5)
+    True
+    >>> np.allclose(samps_y, 0.5)
+    False
+
+    """
+
+    if len(given_vars) != len(given_vals):
+        raise ValueError("length of given_vars not equal to length of given_vals")
+
+    if any(not v.op.random for v in given_vars):
+        nonrandom_ops = [v.op for v in given_vars if not v.op.random]
+        raise ValueError(f"Cannot condition on RV with non-random op(s) {nonrandom_ops}")
+
+    # if no given variables, just do ancestor sampling (works but disabled for simplicity)
+    # if len(given_vars) == 0:
+    #     key = jax.random.PRNGKey(0)
+    #     keys = jax.random.split(key, niter)
+    #     mysample = lambda key: backend.ancestor_sample_flat(vars, key)
+    #     return jax.vmap(mysample)(keys)
+
+    given_vals = [jnp.array(val) for val in given_vals]
+    all_vars = dag.upstream_nodes(tuple(vars) + tuple(given_vars))
+    latent_vars = [var for var in all_vars if var.op.random and var not in given_vars]
+
+    for v in latent_vars:
+        if v.op.discrete:
+            raise ValueError(f"Blackjax backend does not support discrete latent/unobserved RV (saw op {v.op})")
+
+    @jax.jit
+    def log_prob(latent_vals):
+        return jax_backend.ancestor_log_prob_flat(latent_vars + given_vars, latent_vals + given_vals, bijector_dict)
+
+    # key = jax.random.PRNGKey(0)
+    seed = np.random.randint(0, 2**32 - 1)
+    key = jax.random.PRNGKey(seed)
+
+    latent_vals = jax_backend.ancestor_sample_flat(latent_vars, key, bijector_dict=bijector_dict)
+    latent_samps = run_nuts(log_prob, key, latent_vals, niter)
+
+    if bijector_dict is not None and len(latent_samps) > 0:
+
+        def constrain(latent_vals):
+            return jax_backend.ancestor_constrain(latent_vars, latent_vals, bijector_dict)
+
+        latent_samps = jax.vmap(constrain)(latent_samps)
+
+    def fill(latent_vals):
+        return jax_backend.fill_in(latent_vars + given_vars, latent_vals + given_vals, vars)
+
+    # include niter in case latent_samps is empty
+    return jax.vmap(fill, axis_size=niter)(latent_samps)
+
+
+class Calculate:
+    """
+    A `Calculate` object just remembers a set of options and then offers inference
+    methods.
+
+    Parameters
+    ----------
+    default
+        represents a set of options for the inference engine that can be overriden later
+    frozen
+        represents a set of options for the inference engine that cannot be overridden
+
+    """
+
+    def __init__(self, default: Optional[dict] = None, frozen: Optional[dict] = None):
+
+        if default is None:
+            default = {}
+        if frozen is None:
+            frozen = {}
+
+        if util.intersects(default, frozen):
+            raise ValueError(f"default {default} intersects with frozen {frozen}")
+
+        self.default = default
+        self.frozen = frozen  # options for that engine
+
+    def sample(
+        self,
+        vars: PyTree[RV],
+        given_vars: PyTree[RV] = None,
+        given_vals: PyTree[ArrayLike] = None,
+        reduce_fn: Optional[Callable] = None,
+        **options,
+    ):
+        """
+        Draw samples!
+
+        Args:
+            vars: A `RV` or list/tuple of `RV` or pytree of `RV` to sample.
+            given_vars: A `RV` or list/tuple of `RV` or pytree of `RV` to condition on.
+                ``None`` indicates no conditioning variables.
+            given_vals: An ``ArrayLike`` or list/tuple of ``ArrayLike`` or pytree of
+                ``ArrayLike`` representing observed values. Must match the structure and
+                shape of ``given_vars``.
+            reduce_fn:  Function to apply to each leaf node in samples before returning.
+                This is used to create `E`, `var`, etc. (If ``None``, does nothing.)
+            options: extra options to pass to sampler
+
+        Returns:
+            Pytree of JAX arrays matching structure and shape of ``vars`` but with one
+            extra dimension at the start, containing the samples.
+
+        Examples
+        --------
+        >>> zero    = ir.RV(ir.Constant(0))
+        >>> one     = ir.RV(ir.Constant(1))
+        >>> x       = ir.RV(ir.Normal(), zero, one)
+        >>> y       = ir.RV(ir.Normal(), x, one)
+        >>> calc    = Calculate({'niter': 529})
+        >>> x_samps = calc.sample(x,y,2)
+        >>> x_samps.shape
+        (529,)
+        >>> np.mean(x_samps) # something close to 1.0
+        Array(...)
+        """
+
+        if util.intersects(options, self.frozen):
+            raise ValueError(f"options intersects frozen")
+
+        options = self.default | options  # overrides defaults
+
+        (
+            flat_vars,
+            flat_given_vars,
+            flat_given_vals,
+            unflatten,
+            unflatten_given,
+        ) = util.flatten_args(vars, given_vars, given_vals)
+
+        flat_samps = sample_flat(
+            flat_vars,
+            flat_given_vars,
+            flat_given_vals,
+            **options,
+            **self.frozen,
+        )
+
+        if reduce_fn is not None:
+            flat_samps = map(reduce_fn, flat_samps)
+
+        return unflatten(flat_samps)
+
+    def E(
+        self,
+        vars: PyTree[RV],
+        given_vars: PyTree[RV] = None,
+        given_vals: PyTree[ArrayLike] = None,
+        **options,
+    ):
+        """
+        Compute (conditional) expected values. This is just a thin wrapper that calls
+        `sample` and then reduces by taking the mean.
+
+        Args:
+            vars: A `RV` or list/tuple of `RV` or pytree of `RV` to sample.
+            given_vars:  A `RV` or list/tuple of `RV` or pytree of `RV` to condition on.
+                ``None`` indicates no conditioning variables.
+            given_vals: An ``ArrayLike`` or list/tuple of ``ArrayLike`` or pytree of
+                ``ArrayLike`` representing observed values. Must match the structure and
+                shape of ``given_vars``.
+            reduce_fn:  Function to apply to each leaf node in samples before returning.
+                This is used to create `E`, `var`, etc. (If ``None``, does nothing.)
+            options: extra options to pass to sampler
+
+        Returns:
+            Pytree of JAX arrays matching structure and shape of ``vars``, containing
+                the expectations.
+
+
+        Examples
+        --------
+        >>> zero    = ir.RV(ir.Constant(0))
+        >>> one     = ir.RV(ir.Constant(1))
+        >>> x       = ir.RV(ir.Normal(), zero, one)
+        >>> y       = ir.RV(ir.Normal(), x, one)
+        >>> calc    = Calculate({'niter': 529})
+        >>> calc.E(x,y,2) # something close to 1.0
+        Array(...)
+        """
+
+        return self.sample(vars, given_vars, given_vals, lambda x: np.mean(x, axis=0), **options)
+
+    def var(
+        self,
+        vars: PyTree[RV],
+        given_vars: PyTree[RV] = None,
+        given_vals: PyTree[ArrayLike] = None,
+        **options,
+    ):
+        return self.sample(vars, given_vars, given_vals, lambda x: np.var(x, axis=0), **options)
+
+    def std(
+        self,
+        vars: PyTree[RV],
+        given_vars: PyTree[RV] = None,
+        given_vals: PyTree[ArrayLike] = None,
+        **options,
+    ):
+        return self.sample(vars, given_vars, given_vals, lambda x: np.std(x, axis=0), **options)
+
+    def sample_arviz(
+        self,
+        vars: dict[str, RV],
+        given_vars: PyTree[RV] = None,
+        given_vals: PyTree[ArrayLike] = None,
+        **options,
+    ):
+        """This is an **experimental** function to draw samples in
+        `ArviZ <https://www.arviz.org/en/latest/>`__ format.
+
+        Note: ArviZ is not installed with pangolin by default: You must install it
+        manually.
+
+        Args:
+            vars: dictionary mapping names to individual random variables
+                given_vars: A `RV` or list/tuple of `RV` or pytree of `RV` to condition on.
+                ``None`` indicates no conditioning variables.
+            given_vars: A `RV` or list/tuple of `RV` or pytree of `RV` to condition on.
+                given_vals: An ``ArrayLike`` or list/tuple of ``ArrayLike`` or pytree of
+                ``ArrayLike`` representing observed values. Must match the structure and
+                shape of ``given_vars``.
+            reduce_fn:  Function to apply to each leaf node in samples before returning.
+                This is used to create `E`, `var`, etc. (If ``None``, does nothing.)
+            options: extra options to pass to sampler
+        """
+
+        try:
+            from arviz import convert_to_inference_data
+        except ImportError:
+            raise ImportError("To use this method you must install arviz manually")
+
+        samps = self.sample(vars, given_vars, given_vals, **options)
+        samps_with_none = {key: samps[key][None, ...] for key in samps}
+        dataset = convert_to_inference_data(samps_with_none)
+        return dataset
+
+
+default = {"niter": 1000}
+
+calc = Calculate(default)
+sample = calc.sample
+"""
+Default version of `Calculate.sample` that draws 1000 samples.
+"""
+E = calc.E
+"""
+Default version of `Calculate.E` that uses 1000 samples.
+"""
+var = calc.var
+"""
+Default version of `Calculate.var` that uses 1000 samples.
+"""
+std = calc.std
+"""
+Default version of `Calculate.std` that uses 1000 samples.
+"""
+
+sample_arviz = calc.sample_arviz
+"""
+Default version of `Calculate.sample_arviz` that uses 1000 samples.
+"""
+
+
+def inf_until_match(inf, vars, given, vals, testfun, niter_start=1000, niter_max=100000):
+    from time import time
+
+    niter = niter_start
+    while niter <= niter_max:
+        t0 = time()
+        out = inf(vars, given, vals, niter=niter)
+        t1 = time()
+        print(f"{niter=} {t1 - t0}")
+        if testfun(out):
+            assert True
+            return
+        else:
+            niter *= 2
+    assert False
+
+
+import functools
+
+sample_until_match = functools.partial(inf_until_match, sample)
+
+
+def sample_flat_until_match(vars, given, vals, testfun, niter_start=1000, niter_max=100000):
+    new_testfun = lambda stuff: testfun(stuff[0])
+    return inf_until_match(sample_flat, vars, given, vals, new_testfun, niter_start, niter_max)

pangolin/dag.py

@@ -0,0 +1,219 @@
+"""
+Basic code for operating on directed acyclic graphs (DAGs).
+None of these functions import or use any other parts of Pangolin.
+**End-users of Pangolin are not expected to use these functions directly**.
+"""
+
+from __future__ import annotations
+import jax.tree_util
+from typing import Sequence, Callable, Any
+from jaxtyping import PyTree
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+
+class Node:
+    """
+    The basic `Node` class. This is just an object that remembers a set of parents.
+
+    Args:
+        parents: parents of this node
+    """
+
+    def __init__(self, *parents: Self):
+        self._parents: tuple[Self, ...] = parents
+
+    @property
+    def parents(self) -> tuple[Self, ...]:
+        "The parents of this node"
+        return self._parents
+
+
+def upstream_nodes_flat[N: Node](
+    nodes_flat: Sequence[N],
+    node_block: Callable[[N], bool] | None,
+    edge_block: Callable[[N, N], bool] | None,
+    upstream: list[N],
+):
+    """
+    Do a DFS starting at all the nodes in ``nodes_flat``. But never visit nodes if
+    ``node_block(n)`` and never follow an edge from ``n`` to ``p`` if
+    ``link_block(n,p)``.
+
+    Args:
+        nodes_flat: starting nodes
+        node_block: should DFS be blocked from visting a node?
+        edge_block: should DFS be blocked from following a link?
+        upstream: list of nodes, destructively updated
+    """
+    # someday should have dual set / list for faster checking
+    for node in nodes_flat:
+        if node in upstream or (node_block and node_block(node)):
+            continue
+        for p in node.parents:
+            if not (edge_block and edge_block(node, p)) and not (node_block and node_block(p)):
+                upstream_nodes_flat([p], node_block, edge_block, upstream)
+        upstream.append(node)
+
+
+def upstream_nodes[N: Node](
+    nodes: PyTree[N],
+    node_block: Callable[[N], bool] | None = None,
+    edge_block: Callable[[N, N], bool] | None = None,
+) -> list[N]:
+    """
+    Do a DFS of all ancestors starting at all the nodes in ``nodes``.
+
+    Args:
+        nodes
+            Single node or list/pytree of node
+        node_block
+            Function defining which nodes to block from inclusion. If
+            ``node_block(node)`` is True, then ``node`` is not visited. (Defalt: Don't block.)
+        edge_block
+            Function defining which edges should not be followed.  If
+            ``edge_block(node, parent)`` is true, then don't follow edge from ``node`` to
+            ``parent``. (Defalt: Don't block)
+
+    Returns:
+        List of upstream nodes in topological order
+    """
+
+    # Args:
+    #     nodes (`PyTree[Node]`): single node or list of nodes or pytree of starting nodes
+    #     node_block: should DFS be blocked from visting a node? If `None`, then all nodes
+    #     edge_block: should DFS be blocked from following an edge? If `None`, then all
+
+    # Returns:
+    #     List of all nodes found, with a partial order so that parents always come
+    #     before children
+
+    nodes_flat, _ = jax.tree_util.tree_flatten(nodes)
+    upstream = []
+    upstream_nodes_flat(nodes_flat, node_block, edge_block, upstream)
+    return upstream
+
+
+def upstream_with_descendent_old[N: Node](requested_nodes: list[N], given_nodes: list[N]) -> list[N]:
+    """
+    First, find all nodes that are upstream (inclusive) of ``requested_nodes``
+    Then, find all the nodes that are *downstream* (inclusive) of that set
+    that have a descendant in ``given_nodes``
+
+    args:
+        requested_nodes: nodes to search above
+        given_nodes: must have a descendant in this set
+    """
+
+    has_descendent = upstream_nodes(given_nodes)
+    children = get_children(requested_nodes + given_nodes)
+
+    nodes = []
+    processed_nodes = []
+    queue = requested_nodes.copy()
+
+    def unseen(node):
+        return node not in queue and node not in processed_nodes
+
+    while queue:
+        # print(f"{queue=}")
+        node = queue.pop()
+        # if node.cond_dist.is_random and node in has_descendent:
+        if node in has_descendent:
+            nodes.append(node)
+        processed_nodes.append(node)
+
+        for p in node.parents:
+            if unseen(p) and p not in given_nodes:
+                queue.append(p)
+        for c in children[node]:
+            if unseen(c) and c in has_descendent:
+                queue.append(c)
+    return nodes
+
+
+def upstream_with_descendent[N: Node](requested_nodes: list[N], given_nodes: list[N]) -> list[N]:
+    """
+    First, find all nodes that are upstream (inclusive) of ``requested_nodes``
+    Then, find all the nodes that are *downstream* (inclusive) of that set
+    that have a descendant in ``given_nodes``
+    """
+
+    all_nodes = upstream_nodes(requested_nodes + given_nodes)
+    has_descendent = upstream_nodes(given_nodes)
+    return [n for n in all_nodes if n in has_descendent]
+
+
+def get_children[N: Node](nodes: PyTree[N], node_block: Callable[[N], bool] | None = None) -> dict[N, list[N]]:
+    """
+    Get all children for all upstream nodes
+
+    Args:
+        nodes: Starting nodes (as in ``get_upstream``)
+        node_block: Block condition (as in ``get_upstream``)
+
+    Return:
+        Dictionary mapping each upstream node to a list of children of that node
+
+    """
+
+    all_nodes = upstream_nodes(nodes, node_block)
+    children = {}
+    for n in all_nodes:
+        children[n] = []
+    for n in all_nodes:
+        for p in n.parents:
+            # print(f"{hash(n)=} {hash(p)=}")
+            if n not in children[p]:
+                children[p].append(n)
+    return children
+
+
+def is_in[N: Node](node: N, nodes: Sequence[N]) -> bool:
+    for p in nodes:
+        # if id(node) == id(p) or node.id == p.id:
+        if id(node) == id(p):
+            return True
+    return False
+
+
+def has_second_path(node: Node, par_i: int) -> bool:
+    assert par_i < len(node.parents)
+    other_parents = node.parents[:par_i] + node.parents[par_i + 1 :]
+    return node.parents[par_i] in upstream_nodes(other_parents)
+
+
+# def downstream_nodes(roots,nodes,always_include=[]):
+#     downstream = []
+#     for node in nodes:
+#         if is_in(node,roots):
+#             continue
+#         if is_in(node,always_include):
+#             downstream.append(node)
+#             continue
+#         for p in node.parents:
+#             if is_in(p,downstream+roots):
+#                 downstream.append(node)
+#                 break
+#     return downstream
+#
+# def middle_nodes(roots,leafs):
+#     "all nodes that are both upstream of a leaf and downstream of a root"
+#     # TODO: this goes through the whole graph and then filters, could be inefficient but whatever
+#     nodes = upstream_nodes(leafs)
+#     middle = {}
+#     for node in nodes:
+#         if is_in(node,leafs):
+#             continue
+#         for p in node.parents:
+#             if p in middle or is_in(p,roots):
+#                 middle[node] = None # unordered set
+#     return list(middle.keys())
+
+
+def get_graph[N: Node](nodes: PyTree[N], get_content, block_condition=None) -> dict[N, Any]:
+    upstream = upstream_nodes(nodes, block_condition)
+    return dict(zip(upstream, [get_content(n) for n in upstream]))