pymc-extras 0.2.0__py3-none-any.whl

pymc_extras/utils/linear_cg.py

@@ -0,0 +1,290 @@
+#   Copyright 2022 The PyMC Developers
+#
+#   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
+
+EVAL_CG_TOLERANCE = 0.01
+CG_TOLERANCE = 1
+
+
+def masked_fill(vector, mask, fill_value):
+    masked_vector = np.ma.array(vector, mask=mask)
+    vector = masked_vector.filled(fill_value=fill_value)
+    return vector
+
+
+def linear_cg_updates(
+    result, alpha, residual_inner_prod, eps, beta, residual, precond_residual, curr_conjugate_vec
+):
+    # Everything inside _jit_linear_cg_updates
+    result = result + alpha * curr_conjugate_vec
+    beta = np.copy(residual_inner_prod)
+
+    residual_inner_prod = residual.T @ precond_residual
+
+    # safe division
+    is_zero = beta < eps
+    beta = masked_fill(beta, mask=is_zero, fill_value=1)
+
+    beta = residual_inner_prod / beta
+    beta = masked_fill(beta, mask=is_zero, fill_value=0)
+    curr_conjugate_vec = beta * curr_conjugate_vec + precond_residual
+    return (
+        result,
+        alpha,
+        residual_inner_prod,
+        eps,
+        beta,
+        residual,
+        precond_residual,
+        curr_conjugate_vec,
+    )
+
+
+def linear_cg(
+    mat: np.matrix,
+    rhs,
+    n_tridiag=0,
+    tolerance=None,
+    eps=1e-10,
+    stop_updating_after=1e-10,
+    max_iter=1000,
+    max_tridiag_iter=20,
+    initial_guess=None,
+    preconditioner=None,
+    terminate_cg_by_size=False,
+    use_eval_tolerance=False,
+):
+    if initial_guess is None:
+        initial_guess = np.zeros_like(rhs)
+
+    if preconditioner is None:
+
+        def preconditioner(x):
+            return x
+
+        precond = False
+    else:
+        precond = True
+
+    if tolerance is None:
+        if use_eval_tolerance:
+            tolerance = EVAL_CG_TOLERANCE
+        else:
+            tolerance = CG_TOLERANCE
+
+    # If we are running m CG iterations, we obviously can't get more than m Lanczos coefficients
+    if max_tridiag_iter > max_iter:
+        raise RuntimeError(
+            "Getting a tridiagonalization larger than the number of CG iterations run is not possible!"
+        )
+
+    is_vector = len(rhs.shape) == 1
+    if is_vector:
+        rhs = rhs[:, np.newaxis]
+
+    num_rows = rhs.size
+    n_iter = min(max_iter, num_rows) if terminate_cg_by_size else max_iter
+    n_tridiag_iter = min(max_tridiag_iter, num_rows)
+
+    # norm of rhs for convergence tests
+    rhs_norm = np.linalg.norm(rhs, 2)
+    # make almost-zero norms be 1 (so we don't get divide-by-zero errors)
+    rhs_is_zero = rhs_norm < eps
+    rhs_norm = masked_fill(rhs_norm, mask=rhs_is_zero, fill_value=1)
+
+    # lets normalize rhs
+    rhs = rhs / rhs_norm
+
+    # residuals
+    residual = rhs - mat @ initial_guess
+    batch_shape = residual.shape[:-2]
+
+    result = np.copy(initial_guess)
+
+    if not np.allclose(residual, residual):
+        raise RuntimeError("NaNs encountered when trying to perform matrix-vector multiplication")
+
+    # sometimes we are lucky and preconditioner solves the system right away
+    # check for convergence
+    residual_norm = np.linalg.norm(residual, 2)
+    has_converged = residual_norm < stop_updating_after
+
+    if has_converged.all() and not n_tridiag:
+        n_iter = 0  # skip iterations
+    else:
+        precond_residual = preconditioner(residual)
+
+        curr_conjugate_vec = precond_residual
+        residual_inner_prod = residual.T @ precond_residual
+
+        # define storage matrices
+        np.zeros_like(residual)
+        alpha = np.zeros((*batch_shape, 1, rhs.shape[-1]))
+        beta = np.zeros_like(alpha)
+        is_zero = np.zeros((*batch_shape, 1, rhs.shape[-1]))
+
+    # Define tridiagonal matrices if applicable
+    if n_tridiag:
+        t_mat = np.zeros((n_tridiag_iter, n_tridiag_iter, *batch_shape, n_tridiag))
+        alpha_tridiag_is_zero = np.zeros(*batch_shape, n_tridiag)
+        alpha_reciprocal = np.zeros(*batch_shape, n_tridiag)
+        prev_alpha_reciprocal = np.zeros_like(alpha_reciprocal)
+        prev_beta = np.zeros_like(alpha_reciprocal)
+
+    update_tridiag = True
+    last_tridiag_iter = 0
+
+    # it is possible that we don't reach tolerance even after all the iterations are over
+    tolerance_reached = False
+
+    # start iteration
+    for k in range(n_iter):
+        mvms = mat @ curr_conjugate_vec
+        if precond:
+            alpha = curr_conjugate_vec @ mvms  # scalar
+
+            # safe division
+            is_zero = alpha < eps
+            alpha = masked_fill(alpha, mask=is_zero, fill_value=1)
+            alpha = residual_inner_prod / alpha
+            alpha = masked_fill(alpha, mask=is_zero, fill_value=0)
+
+            # cancel out updates by setting directions which have converged to zero
+            alpha = masked_fill(alpha, mask=has_converged, fill_value=0)
+            residual = residual - alpha * mvms
+
+            # update precond_residual
+            precond_residual = preconditioner(residual)
+
+            # Everything inside _jit_linear_cg_updates
+            (
+                result,
+                alpha,
+                residual_inner_prod,
+                eps,
+                beta,
+                residual,
+                precond_residual,
+                curr_conjugate_vec,
+            ) = linear_cg_updates(
+                result,
+                alpha,
+                residual_inner_prod,
+                eps,
+                beta,
+                residual,
+                precond_residual,
+                curr_conjugate_vec,
+            )
+
+        else:
+            # everything inside _jit_linear_cg_updates_no_precond
+            alpha = curr_conjugate_vec.T @ mvms
+
+            # safe division
+            is_zero = alpha < eps
+            alpha = masked_fill(alpha, mask=is_zero, fill_value=1)
+            alpha = residual_inner_prod / alpha
+            alpha = masked_fill(alpha, is_zero, fill_value=0)
+
+            alpha = masked_fill(alpha, has_converged, fill_value=0)  # <- I'm here
+            residual = residual - alpha * mvms
+            precond_residual = np.copy(residual)
+
+            (
+                result,
+                alpha,
+                residual_inner_prod,
+                eps,
+                beta,
+                residual,
+                precond_residual,
+                curr_conjugate_vec,
+            ) = linear_cg_updates(
+                result,
+                alpha,
+                residual_inner_prod,
+                eps,
+                beta,
+                residual,
+                precond_residual,
+                curr_conjugate_vec,
+            )
+
+        residual_norm = np.linalg.norm(residual, 2)
+        residual_norm = masked_fill(residual_norm, mask=rhs_is_zero, fill_value=0)
+        has_converged = residual_norm < stop_updating_after
+
+        if (
+            k >= min(10, max_iter - 1)
+            and bool(residual_norm.mean() < tolerance)
+            and not (n_tridiag and k < min(n_tridiag_iter, max_iter - 1))
+        ):
+            tolerance_reached = True
+            break
+
+        # Update tridiagonal matrices, if applicable
+        if n_tridiag and k < n_tridiag_iter and update_tridiag:
+            alpha_tridiag = np.copy(alpha)
+            beta_tridiag = np.copy(beta)
+
+            alpha_tridiag_is_zero = alpha_tridiag == 0
+            alpha_tridiag = masked_fill(alpha_tridiag, mask=alpha_tridiag_is_zero, fill_value=1)
+            alpha_reciprocal = 1 / alpha_tridiag
+            alpha_tridiag = masked_fill(alpha_tridiag, mask=alpha_tridiag_is_zero, fill_value=0)
+
+            if k == 0:
+                t_mat[k, k] = alpha_reciprocal
+            else:
+                t_mat[k, k] += np.squeeze(alpha_reciprocal + prev_beta * prev_alpha_reciprocal)
+                t_mat[k, k - 1] = np.sqrt(prev_beta) * prev_alpha_reciprocal
+                t_mat[k - 1, k] = np.copy(t_mat[k, k - 1])
+
+                if t_mat[k - 1, k].max() < 1e-6:
+                    update_tridiag = False
+
+            last_tridiag_iter = k
+
+            prev_alpha_reciprocal = np.copy(alpha_reciprocal)
+            prev_beta = np.copy(beta_tridiag)
+
+    # Un-normalize
+    result = result * rhs_norm
+    if not tolerance_reached and n_iter > 0:
+        raise RuntimeError(
+            f"CG terminated in {k + 1} iterations with average residual norm {residual_norm.mean()}"
+            f" which is larger than the tolerance of {tolerance} specified by"
+            " gpytorch.settings.cg_tolerance."
+            " If performance is affected, consider raising the maximum number of CG iterations by running code in"
+            " a gpytorch.settings.max_cg_iterations(value) context."
+        )
+
+    if n_tridiag:
+        t_mat = t_mat[: last_tridiag_iter + 1, : last_tridiag_iter + 1]
+        return result, t_mat.transpose(-1, *range(2, 2 + len(batch_shape)), 0, 1)
+    else:
+        # We set the estimated Lanczos tri-diagonal matrices to be identity so that
+        # the subsequent eigen decomposition https://arxiv.org/pdf/1809.11165.pdf (eq.S7)
+        # would work fine.
+        # t_mat = np.zeros((n_tridiag_iter, n_tridiag_iter, *batch_shape, n_tridiag))
+        # Note that after transpose the last two dimensions are dimensions 0 and 1 of the matrix above
+        # Which are the same values i.e. n_tridiag_iter
+        # So we generate identity matrices of size n_tridiag_iter and repeat them [n_iter, *range(2, 2+len(batch_shape))] times
+        # TODO: for same input, n_tridiag = True and n_tridiag = False must produce t_mat with same shape (with assumed n_tridiag=1)
+        n_tridiag = 1
+        eye = np.eye(n_tridiag_iter)
+        t_mat_eye = np.tile(eye, [n_tridiag] + [1] * (len(batch_shape) + 2))
+        return result, t_mat_eye

pymc_extras/utils/pivoted_cholesky.py

@@ -0,0 +1,69 @@
+try:
+    import torch
+
+    from gpytorch.utils.permutation import apply_permutation
+except ImportError as e:
+    raise ImportError("PyTorch and GPyTorch not found.") from e
+
+import numpy as np
+
+
+def pp(x):
+    return np.array2string(x, precision=4, floatmode="fixed")
+
+
+def pivoted_cholesky(mat: np.matrix, error_tol=1e-6, max_iter=np.inf):
+    """
+    mat: numpy matrix of N x N
+
+    This is to replicate what is done in GPyTorch verbatim.
+    """
+    n = mat.shape[-1]
+    max_iter = min(int(max_iter), n)
+
+    d = np.array(np.diag(mat))
+    orig_error = np.max(d)
+    error = np.linalg.norm(d, 1) / orig_error
+    pi = np.arange(n)
+
+    L = np.zeros((max_iter, n))
+
+    m = 0
+    while m < max_iter and error > error_tol:
+        permuted_d = d[pi]
+        max_diag_idx = np.argmax(permuted_d[m:])
+        max_diag_idx = max_diag_idx + m
+        max_diag_val = permuted_d[max_diag_idx]
+        i = max_diag_idx
+
+        # swap pi_m and pi_i
+        pi[m], pi[i] = pi[i], pi[m]
+        pim = pi[m]
+
+        L[m, pim] = np.sqrt(max_diag_val)
+
+        if m + 1 < n:
+            row = apply_permutation(
+                torch.from_numpy(mat), torch.tensor(pim), right_permutation=None
+            )  # left permutation just swaps row
+            row = row.numpy().flatten()
+            pi_i = pi[m + 1 :]
+            L_m_new = row[pi_i]  # length = 9
+
+            if m > 0:
+                L_prev = L[:m, pi_i]
+                update = L[:m, pim]
+                prod = update @ L_prev
+                L_m_new = L_m_new - prod  # np.sum(prod, axis=-1)
+
+            L_m = L[m, :]
+            L_m_new = L_m_new / L_m[pim]
+            L_m[pi_i] = L_m_new
+
+            matrix_diag_current = d[pi_i]
+            d[pi_i] = matrix_diag_current - L_m_new**2
+
+            L[m, :] = L_m
+            error = np.linalg.norm(d[pi_i], 1) / orig_error
+        m = m + 1
+    return L, pi

pymc_extras/utils/prior.py

@@ -0,0 +1,200 @@
+#   Copyright 2022 The PyMC Developers
+#
+#   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 collections.abc import Sequence
+from typing import TypedDict
+
+import arviz
+import numpy as np
+import pymc as pm
+import pytensor.tensor as pt
+
+from pymc.logprob.transforms import Transform
+
+
+class ParamCfg(TypedDict):
+    name: str
+    transform: Transform | None
+    dims: str | tuple[str] | None
+
+
+class ShapeInfo(TypedDict):
+    # shape might not match slice due to a transform
+    shape: tuple[int]  # transformed shape
+    slice: slice
+
+
+class VarInfo(TypedDict):
+    sinfo: ShapeInfo
+    vinfo: ParamCfg
+
+
+class FlatInfo(TypedDict):
+    data: np.ndarray
+    info: list[VarInfo]
+
+
+def _arg_to_param_cfg(key, value: ParamCfg | Transform | str | tuple | None = None):
+    if value is None:
+        cfg = ParamCfg(name=key, transform=None, dims=None)
+    elif isinstance(value, tuple):
+        cfg = ParamCfg(name=key, transform=None, dims=value)
+    elif isinstance(value, str):
+        cfg = ParamCfg(name=value, transform=None, dims=None)
+    elif isinstance(value, Transform):
+        cfg = ParamCfg(name=key, transform=value, dims=None)
+    else:
+        cfg = value.copy()
+        cfg.setdefault("name", key)
+        cfg.setdefault("transform", None)
+        cfg.setdefault("dims", None)
+    return cfg
+
+
+def _parse_args(
+    var_names: Sequence[str], **kwargs: ParamCfg | Transform | str | tuple
+) -> dict[str, ParamCfg]:
+    results = dict()
+    for var in var_names:
+        results[var] = _arg_to_param_cfg(var)
+    for key, val in kwargs.items():
+        results[key] = _arg_to_param_cfg(key, val)
+    return results
+
+
+def _flatten(idata: arviz.InferenceData, **kwargs: ParamCfg) -> FlatInfo:
+    posterior = idata.posterior
+    vars = list()
+    info = list()
+    begin = 0
+    for key, cfg in kwargs.items():
+        data = (
+            posterior[key]
+            # combine all draws from all chains
+            .stack(__sample__=["chain", "draw"])
+            # move sample dim to the first position
+            # no matter where it was before
+            .transpose("__sample__", ...)
+            # we need numpy data for all the rest functionality
+            .values
+        )
+        # omitting __sample__
+        # we need shape in the untransformed space
+        if cfg["transform"] is not None:
+            # some transforms need original shape
+            data = cfg["transform"].forward(data).eval()
+        shape = data.shape[1:]
+        # now we can get rid of shape
+        data = data.reshape(data.shape[0], -1)
+        end = begin + data.shape[1]
+        vars.append(data)
+        sinfo = dict(shape=shape, slice=slice(begin, end))
+        info.append(dict(sinfo=sinfo, vinfo=cfg))
+        begin = end
+    return dict(data=np.concatenate(vars, axis=-1), info=info)
+
+
+def _mean_chol(flat_array: np.ndarray):
+    mean = flat_array.mean(0)
+    cov = np.cov(flat_array, rowvar=False)
+    cov = np.atleast_2d(cov)
+    chol = np.linalg.cholesky(cov)
+    return mean, chol
+
+
+def _mvn_prior_from_flat_info(name, flat_info: FlatInfo):
+    mean, chol = _mean_chol(flat_info["data"])
+    base_dist = pm.Normal(name, np.zeros_like(mean))
+    interim = mean + chol @ base_dist
+    result = dict()
+    for var_info in flat_info["info"]:
+        sinfo = var_info["sinfo"]
+        vinfo = var_info["vinfo"]
+        var = interim[sinfo["slice"]].reshape(sinfo["shape"])
+        if vinfo["transform"] is not None:
+            var = vinfo["transform"].backward(var)
+        var = pm.Deterministic(vinfo["name"], var, dims=vinfo["dims"])
+        result[vinfo["name"]] = var
+    return result
+
+
+def prior_from_idata(
+    idata: arviz.InferenceData,
+    name="trace_prior_",
+    *,
+    var_names: Sequence[str] = (),
+    **kwargs: ParamCfg | Transform | str | tuple,
+) -> dict[str, pt.TensorVariable]:
+    """
+    Create a prior from posterior using MvNormal approximation.
+
+    The approximation uses MvNormal distribution.
+    Keep in mind that this function will only work well for unimodal
+    posteriors and will fail when complicated interactions happen.
+
+    Moreover, if a retrieved variable is constrained, you
+    should specify a transform for the variable, e.g.
+    ``pymc.distributions.transforms.log`` for standard
+    deviation posterior.
+
+    Parameters
+    ----------
+    idata: arviz.InferenceData
+        Inference data with posterior group
+    var_names: Sequence[str]
+        names of variables to take as is from the posterior
+    kwargs: Union[ParamCfg, Transform, str, Tuple]
+        names of variables with additional configuration, see more in Examples
+
+    Examples
+    --------
+    >>> import pymc as pm
+    >>> import pymc.distributions.transforms as transforms
+    >>> import numpy as np
+    >>> with pm.Model(coords=dict(test=range(4), options=range(3))) as model1:
+    ...     a = pm.Normal("a")
+    ...     b = pm.Normal("b", dims="test")
+    ...     c = pm.HalfNormal("c")
+    ...     d = pm.Normal("d")
+    ...     e = pm.Normal("e")
+    ...     f = pm.Dirichlet("f", np.ones(3), dims="options")
+    ...     trace = pm.sample(progressbar=False)
+
+    You can reuse the posterior in the new model.
+
+    >>> with pm.Model(coords=dict(test=range(4), options=range(3))) as model2:
+    ...     priors = prior_from_idata(
+    ...         trace,                  # the old trace (posterior)
+    ...         var_names=["a", "d"],   # take variables as is
+    ...
+    ...         e="new_e",              # assign new name "new_e" for a variable
+    ...                                 # similar to dict(name="new_e")
+    ...
+    ...         b=("test", ),           # set a dim to "test"
+    ...                                 # similar to dict(dims=("test", ))
+    ...
+    ...         c=transforms.log,       # apply log transform to a positive variable
+    ...                                 # similar to dict(transform=transforms.log)
+    ...
+    ...                                 # set a name, assign a dim and apply simplex transform
+    ...         f=dict(name="new_f", dims="options", transform=transforms.simplex)
+    ...     )
+    ...     trace1 = pm.sample_prior_predictive(100)
+    """
+    param_cfg = _parse_args(var_names=var_names, **kwargs)
+    if not param_cfg:
+        return {}
+    flat_info = _flatten(idata, **param_cfg)
+    return _mvn_prior_from_flat_info(name, flat_info)

pymc_extras/utils/spline.py

@@ -0,0 +1,131 @@
+#   Copyright 2022 The PyMC Developers
+#
+#   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 pytensor
+import pytensor.sparse as ps
+import pytensor.tensor as pt
+import scipy.interpolate
+
+from pytensor.graph.op import Apply, Op
+
+
+def numpy_bspline_basis(eval_points: np.ndarray, k: int, degree=3):
+    k_knots = k + degree + 1
+    knots = np.linspace(0, 1, k_knots - 2 * degree)
+    knots = np.r_[[0] * degree, knots, [1] * degree]
+    basis_funcs = scipy.interpolate.BSpline(knots, np.eye(k), k=degree)
+    Bx = basis_funcs(eval_points).astype(eval_points.dtype)
+    return Bx
+
+
+class BSplineBasis(Op):
+    __props__ = ("sparse",)
+
+    def __init__(self, sparse=True) -> None:
+        super().__init__()
+        if not isinstance(sparse, bool):
+            raise TypeError("sparse should be True or False")
+        self.sparse = sparse
+
+    def make_node(self, *inputs) -> Apply:
+        eval_points, k, d = map(pt.as_tensor, inputs)
+        if not (eval_points.ndim == 1 and np.issubdtype(eval_points.dtype, np.floating)):
+            raise TypeError("eval_points should be a vector of floats")
+        if k.type not in pt.int_types:
+            raise TypeError("k should be integer")
+        if d.type not in pt.int_types:
+            raise TypeError("degree should be integer")
+        if self.sparse:
+            out_type = ps.SparseTensorType("csr", eval_points.dtype)()
+        else:
+            out_type = pt.matrix(dtype=eval_points.dtype)
+        return Apply(self, [eval_points, k, d], [out_type])
+
+    def perform(self, node, inputs, output_storage, params=None) -> None:
+        eval_points, k, d = inputs
+        Bx = numpy_bspline_basis(eval_points, int(k), int(d))
+        if self.sparse:
+            Bx = scipy.sparse.csr_matrix(Bx, dtype=eval_points.dtype)
+        output_storage[0][0] = Bx
+
+    def infer_shape(self, fgraph, node, ins_shapes):
+        return [(node.inputs[0].shape[0], node.inputs[1])]
+
+
+def bspline_basis(n, k, degree=3, dtype=None, sparse=True):
+    dtype = dtype or pytensor.config.floatX
+    eval_points = np.linspace(0, 1, n, dtype=dtype)
+    return BSplineBasis(sparse=sparse)(eval_points, k, degree)
+
+
+def bspline_interpolation(x, *, n=None, eval_points=None, degree=3, sparse=True):
+    """Interpolate sparse grid to dense grid using bsplines.
+
+    Parameters
+    ----------
+    x : Variable
+        Input Variable to interpolate.
+        0th coordinate assumed to be mapped regularly on [0, 1] interval
+    n : int (optional)
+        Resolution of interpolation
+    eval_points : vector (optional)
+        Custom eval points in [0, 1] interval (or scaled properly using min/max scaling)
+    degree : int, optional
+        BSpline degree, by default 3
+    sparse : bool, optional
+        Use sparse operation, by default True
+
+    Returns
+    -------
+    Variable
+        The interpolated variable, interpolation is across 0th axis
+
+    Examples
+    --------
+    >>> import pymc as pm
+    >>> import numpy as np
+    >>> half_months = np.linspace(0, 365, 12*2)
+    >>> with pm.Model(coords=dict(knots_time=half_months, time=np.arange(365))) as model:
+    ...     kernel = pm.gp.cov.ExpQuad(1, ls=365/12)
+    ...     # ready to define gp (a latent process over parameters)
+    ...     gp = pm.gp.gp.Latent(
+    ...         cov_func=kernel
+    ...     )
+    ...     y_knots = gp.prior("y_knots", half_months[:, None], dims="knots_time")
+    ...     y = pm.Deterministic(
+    ...         "y",
+    ...         bspline_interpolation(y_knots, n=365, degree=3),
+    ...         dims="time"
+    ...     )
+    ...     trace = pm.sample_prior_predictive(1)
+
+    Notes
+    -----
+    Adopted from `BayesAlpha <https://github.com/quantopian/bayesalpha/blob/676f4f194ad20211fd040d3b0c6e82969aafb87e/bayesalpha/dists.py#L97>`_
+    where it was written by @aseyboldt
+    """
+    x = pt.as_tensor(x)
+    if n is not None and eval_points is not None:
+        raise ValueError("Please provide one of n or eval_points")
+    elif n is not None:
+        eval_points = np.linspace(0, 1, n, dtype=x.dtype)
+    elif eval_points is None:
+        raise ValueError("Please provide one of n or eval_points")
+    basis = BSplineBasis(sparse=sparse)(eval_points, x.shape[0], degree)
+    if sparse:
+        return ps.dot(basis, x)
+    else:
+        return pt.dot(basis, x)

pymc_extras/version.py

@@ -0,0 +1,11 @@
+import os
+
+
+def get_version():
+    version_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "version.txt")
+    with open(version_file) as f:
+        version = f.read().strip()
+    return version
+
+
+__version__ = get_version()

pymc_extras/version.txt

@@ -0,0 +1 @@
+0.2.0