prisma-oed 0.0.1__py3-none-any.whl

prisma/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.0.1'

prisma/apf/api/__init__.py

@@ -0,0 +1,7 @@
+from .active_parameter_fit import ActiveParameterFitSigmaSqPred, \
+    ActiveParameterFitRP, ActiveParameterFitEI, \
+    AbstractActiveParameterFit
+
+__all__ = ['ActiveParameterFitSigmaSqPred', 'ActiveParameterFitRP',
+           'ActiveParameterFitEI',
+           'AbstractActiveParameterFit']

prisma/apf/api/active_parameter_fit.py

@@ -0,0 +1,182 @@
+from abc import ABC, abstractmethod
+import jax
+import jax.numpy as jnp
+from jax.stages import Wrapped
+import numpy as np
+from numpy.typing import ArrayLike
+
+from ..util.active_parameter_fit_jax_lib import get_nonlin_f_x_j, \
+    acq_f_sigma_sq_predictive_j, acq_f_rp_j, \
+    acq_f_ei_j, \
+    init_arr_j, init_arr_rp_j, init_arr_ei_j, \
+    posterior_update_j
+from prisma.file import save_np_val
+from prisma.plot import plot_predictive_train_iter, \
+    plot_acq_train_iter, plot_predictive_acq_train_iter, plot_w_rmse, \
+    plot_w_i_err, plot_predictive_rmse
+from prisma.samplers.abstract import AbstractDomainSampler
+
+
+class AbstractActiveParameterFit(ABC):
+
+    def __init__(self, nonlin_f_j: Wrapped, acq_domain: ArrayLike,
+                 t_sampler: AbstractDomainSampler,
+                 beta: float, m_0: ArrayLike, S_0: ArrayLike,
+                 sample_ct: int):
+        # TODO: initialize all jit argument vars as jnp
+        self.nonlin_f_j = nonlin_f_j
+        self.acq_domain = acq_domain
+        self.t_sampler = t_sampler
+        self.beta = beta
+        self.param_ct = len(m_0)
+        self.sample_ct = sample_ct
+        self.acq_domain_len = len(acq_domain)
+        self.nonlin_f_j_N = get_nonlin_f_x_j(nonlin_f_j, acq_domain)
+        self.x_n_jnp, self.t_n_jnp, self.phi_n_jnp, \
+            self.m_prior_jnp, self.var_prior_jnp, self.acq_val_jnp \
+            = init_arr_j(m_0, S_0, self.param_ct, sample_ct,
+                         self.acq_domain_len)
+        self.x_n_np = self.t_n_np = self.phi_n_np \
+            = self.m_prior_np = self.var_prior_np = self.acq_val_np = None
+
+    def fit(self):
+        for i in range(self.sample_ct):
+            x_next, acq_val_next = self._acquisition_f(i)
+            t_next = self.t_sampler.sample(x_next)
+            self.x_n_jnp, self.t_n_jnp, self.acq_val_jnp, \
+                self.phi_n_jnp, self.m_prior_jnp, self.var_prior_jnp \
+                = posterior_update_j(x_next, acq_val_next, t_next,
+                                     self.nonlin_f_j, self.beta,
+                                     i, self.x_n_jnp, self.t_n_jnp,
+                                     self.phi_n_jnp, self.m_prior_jnp,
+                                     self.var_prior_jnp, self.acq_val_jnp)
+
+    def fit_sampler_reset(self):
+        self.t_sampler.reset()
+        self.fit()
+
+    def init_np_arrs(self):
+        self.x_n_np, self.t_n_np, self.phi_n_np, \
+            self.m_prior_np, self.var_prior_np, self.acq_val_np \
+            = self._jax_arrs_to_np_arrs(self.x_n_jnp, self.t_n_jnp,
+                                        self.phi_n_jnp, self.m_prior_jnp,
+                                        self.var_prior_jnp, self.acq_val_jnp)
+
+    def save_np_arrs(self, save_dir):
+        save_np_val(save_dir, "x_n.npy", self.x_n_np)
+        save_np_val(save_dir, "t_n.npy", self.t_n_np)
+        save_np_val(save_dir, "phi_n.npy", self.phi_n_np)
+        save_np_val(save_dir, "m_prior.npy", self.m_prior_np)
+        save_np_val(save_dir, "var_prior.npy", self.var_prior_np)
+        save_np_val(save_dir, "acq_n.npy", self.acq_val_np)
+
+    def plot_predictive_train_iter(self, y=None, title=True,
+                                   title_samples=False):
+        plot_predictive_train_iter(self.nonlin_f_j, self.acq_domain,
+                                   self.beta, self.m_prior_jnp,
+                                   self.var_prior_jnp,
+                                   y=y,
+                                   x_n=self.x_n_jnp, t_n=self.t_n_jnp,
+                                   title=title, title_samples=title_samples)
+
+    def plot_acq_train_iter(self):
+        plot_acq_train_iter(self.acq_domain, self.acq_val_jnp)
+
+    def plot_predictive_acq_train_iter(self, y=None, title=True,
+                                       title_samples=False):
+        plot_predictive_acq_train_iter(self.nonlin_f_j, self.acq_domain,
+                                       self.beta, self.m_prior_jnp,
+                                       self.var_prior_jnp,
+                                       self.acq_val_jnp,
+                                       y=y,
+                                       x_n=self.x_n_jnp, t_n=self.t_n_jnp,
+                                       title=title,
+                                       title_samples=title_samples)
+
+    def plot_w_rmse(self, w_act):
+        plot_w_rmse(w_act, m_N_l=self.m_prior_jnp)
+
+    def plot_w_i_err(self, w_act, i, w_i_name=None):
+        plot_w_i_err(w_act, i, m_N_l=self.m_prior_jnp, w_i_name=w_i_name)
+
+    def plot_predictive_rmse(self, t, y):
+        plot_predictive_rmse(self.nonlin_f_j, self.acq_domain, t, y,
+                             m_N_l=self.m_prior_jnp)
+
+    @abstractmethod
+    def _acquisition_f(self, i):
+        ...
+
+    def _jax_arrs_to_np_arrs(self, *jax_arrs):
+        return tuple(map(np.array, jax_arrs))
+
+
+class ActiveParameterFitSigmaSqPred(AbstractActiveParameterFit):
+    '''
+    Predictive variance
+    '''
+
+    def _acquisition_f(self, i):
+        x_next, acq = acq_f_sigma_sq_predictive_j(self.nonlin_f_j_N,
+                                                  self.acq_domain,
+                                                  self.beta, i,
+                                                  self.m_prior_jnp,
+                                                  self.var_prior_jnp)
+        return x_next, acq
+
+
+class ActiveParameterFitRP(AbstractActiveParameterFit):
+    '''
+    Normal repulsion
+    '''
+
+    def __init__(self, nonlin_f_j: Wrapped, acq_domain: ArrayLike,
+                 t_sampler: AbstractDomainSampler,
+                 beta: float, m_0: ArrayLike, S_0: ArrayLike,
+                 sample_ct: int,
+                 magnitude_rp: float, scale_rp: float):
+        super().__init__(nonlin_f_j, acq_domain, t_sampler, beta, m_0, S_0,
+                         sample_ct)
+        self.mag_rp = magnitude_rp
+        self.scl_rp = scale_rp
+        self.rp_jnp = init_arr_rp_j(self.acq_domain_len)
+
+    def _acquisition_f(self, i):
+        x_next, acq, rp_next = acq_f_rp_j(self.nonlin_f_j_N, self.acq_domain,
+                                          self.beta, i,
+                                          self.m_prior_jnp, self.var_prior_jnp,
+                                          self.mag_rp, self.scl_rp,
+                                          self.rp_jnp)
+        self.rp_jnp = rp_next
+        return x_next, acq
+
+
+class ActiveParameterFitEI(AbstractActiveParameterFit):
+    '''
+    Expected Improvement
+    '''
+
+    def __init__(self, nonlin_f_j: Wrapped, acq_domain: ArrayLike,
+                 t_sampler: AbstractDomainSampler,
+                 beta: float, m_0: ArrayLike, S_0: ArrayLike,
+                 sample_ct: int,
+                 range_delta: ArrayLike,
+                 m_mask: ArrayLike | None = None):
+        super().__init__(nonlin_f_j, acq_domain, t_sampler, beta, m_0, S_0,
+                         sample_ct)
+        if m_mask is None:
+            m_mask = np.ones(self.param_ct)
+        if jax.config.read('jax_enable_x64'):
+            self.eps_j = jnp.finfo(jnp.float64).eps
+        else:
+            self.eps_j = jnp.finfo(jnp.float32).eps
+        self.m_mask_jnp, self.range_delta_jnp = init_arr_ei_j(m_mask,
+                                                              range_delta)
+
+    def _acquisition_f(self, i):
+        x_next, acq, = acq_f_ei_j(self.nonlin_f_j_N, self.nonlin_f_j,
+                                  self.acq_domain, self.beta, i,
+                                  self.m_prior_jnp, self.var_prior_jnp,
+                                  self.range_delta_jnp, self.m_mask_jnp,
+                                  self.eps_j)
+        return x_next, acq

prisma/apf/util/active_parameter_fit_jax_lib.py

@@ -0,0 +1,262 @@
+from functools import partial
+import jax
+from jax import Array
+from jax import jit
+import jax.numpy as jnp
+from jax.tree_util import Partial as Partialj
+
+from prisma.bonlr import get_S_N_MAP_j, get_m_N_MAP_j, \
+    get_S_N_MAP_i_j, get_m_N_MAP_i_j, get_sigma_sq_predictive_j
+
+
+@jit
+def train_posterior_j(Phi, t, beta, m_prior, var_prior):
+    S_N = get_S_N_MAP_j(Phi, beta, var_prior)
+    m_N = get_m_N_MAP_j(Phi, t, beta, m_prior, var_prior, S_N)
+    return m_N, S_N
+
+
+@jit
+def train_posterior_i_j(phi, t, beta, m_prior, var_prior):
+    S_N = get_S_N_MAP_i_j(phi, beta, var_prior)
+    m_N = get_m_N_MAP_i_j(phi, t, beta, m_prior, var_prior, S_N)
+    return m_N, S_N
+
+
+@partial(jit, static_argnames='nonlin_f_j')
+def get_nonlin_f_x_j(nonlin_f_j, x):
+    return Partialj(nonlin_f_j, x)
+
+
+@partial(jit, static_argnames='nonlin_f_j')
+def get_nonlin_f_x_i_j(nonlin_f_j, x_i):
+    x_i_jnp = jnp.array([x_i])
+    nonlin_f_j_x_i = get_nonlin_f_x_j(nonlin_f_j, x_i_jnp)
+    return nonlin_f_j_x_i
+
+
+# @jit
+# def get_phi_ts_j(nonlin_f, m_prior):
+#     return jax.grad(nonlin_f)(m_prior)
+
+
+@jit
+def get_Phi_ts_j(nonlin_f, m_prior):
+    return jax.jacfwd(nonlin_f)(m_prior)
+
+
+@jit
+def get_Phi_vec_ts_j(nonlin_f, m_prior):
+    Phi = get_Phi_ts_j(nonlin_f, m_prior)
+    phi = jnp.squeeze(Phi)
+    return phi
+
+
+@jit
+def get_t_lin_ts_j(nonlin_f_x_n, t_n, Phi_n, m_prior):
+    return t_n - nonlin_f_x_n(m_prior) + Phi_n @ m_prior
+
+
+@jit
+def get_sigma_sq_predictive_N_j(nonlin_f_N, m_n, S_n, beta):
+    Phi = get_Phi_ts_j(nonlin_f_N, m_n)
+    sigma_sq_predictive = get_sigma_sq_predictive_j(Phi, S_n, beta)
+    return sigma_sq_predictive
+
+
+@jit
+def get_sigma_sq_predictive_N_acq_j(nonlin_f_N, m_prior_jnp, var_prior_jnp,
+                                    beta, i):
+    m_n = m_prior_jnp[i]
+    S_n = var_prior_jnp[i]
+    sigma_sq_predictive = get_sigma_sq_predictive_N_j(nonlin_f_N, m_n, S_n,
+                                                      beta)
+    return sigma_sq_predictive
+
+
+@jit
+def get_x_next_j(acq, x):
+    n_next = jnp.argmax(acq)
+    x_next = x[n_next]
+    return x_next
+
+
+@jit
+def acq_f_sigma_sq_predictive_j(nonlin_f_N, acq_domain, beta, i,
+                                m_prior_jnp, var_prior_jnp):
+    sigma_sq_predictive \
+        = get_sigma_sq_predictive_N_acq_j(nonlin_f_N, m_prior_jnp,
+                                          var_prior_jnp, beta, i)
+    x_next = get_x_next_j(sigma_sq_predictive, acq_domain)
+    return x_next, sigma_sq_predictive
+
+
+@jit
+def get_rp_es_j(x_next, x, mag_rp, scale_rp):
+    return mag_rp * jax.scipy.stats.norm.pdf(x, x_next, scale_rp)
+
+
+@jit
+def acq_f_rp_j(nonlin_f_N, acq_domain, beta, i,
+               m_prior_jnp, var_prior_jnp,
+               mag_rp, scale_rp, rp_prev):
+    sigma_sq_predictive \
+            = get_sigma_sq_predictive_N_acq_j(nonlin_f_N, m_prior_jnp,
+                                              var_prior_jnp, beta, i)
+    acq = sigma_sq_predictive - rp_prev
+    x_next = get_x_next_j(acq, acq_domain)
+    rp_next = rp_prev + get_rp_es_j(x_next, acq_domain, mag_rp, scale_rp)
+    return x_next, acq, rp_next
+
+
+@jit
+def posterior_improvement_j(t, nonlin_f_j_x, beta, phi_x,
+                            m_n, S_n,
+                            m_mask, eps_j):
+    m_hyp, S_hyp = get_next_posterior_x_j(t, nonlin_f_j_x, beta,
+                                          phi_x, m_n, S_n)
+    m_improvement_vec = m_hyp - m_n
+    m_n_eps = m_n + eps_j
+    m_improvement_prop_vec = m_improvement_vec / m_n_eps
+    m_improvement_prop_vec = jnp.where(m_mask, m_improvement_prop_vec, 0)
+    norm_improvement = jnp.linalg.norm(m_improvement_prop_vec)
+    return norm_improvement
+
+
+@jit
+def predictive_probability(t, predictive_loc, predictive_scale):
+    pred_prob = jax.scipy.stats.norm.pdf(t, predictive_loc, predictive_scale)
+    return pred_prob
+
+
+@jit
+def probability_weighted_improvement_j(nonlin_f_j_x, beta, phi_x,
+                                       m_n, S_n,
+                                       predictive_loc, predictive_scale,
+                                       m_mask, eps_j,
+                                       t):
+    norm_improvement = posterior_improvement_j(t, nonlin_f_j_x, beta,
+                                               phi_x, m_n, S_n, m_mask, eps_j)
+    pred_prob = predictive_probability(t, predictive_loc, predictive_scale)
+    prob_w_imp = pred_prob * norm_improvement
+    return prob_w_imp
+
+
+@partial(jit, static_argnames='nonlin_f_j')
+def expected_improvement_j(nonlin_f_j, beta,
+                           m_n, S_n,
+                           range_delta, m_mask, eps_j, x_i,
+                           predictive_loc, predictive_scale):
+    nonlin_f_j_x_i = get_nonlin_f_x_i_j(nonlin_f_j, x_i)
+    phi_i = get_Phi_vec_ts_j(nonlin_f_j_x_i, m_n)
+    range_domain = predictive_loc + range_delta
+    probability_weighted_improvement_j_p \
+        = Partialj(probability_weighted_improvement_j,
+                   nonlin_f_j_x_i, beta, phi_i, m_n, S_n,
+                   predictive_loc, predictive_scale,
+                   m_mask, eps_j)
+    probability_weighted_improvement_j_p_v \
+        = jax.vmap(probability_weighted_improvement_j_p,
+                   in_axes=0)
+    y_trapezoid \
+        = probability_weighted_improvement_j_p_v(range_domain)
+    ei = jax.scipy.integrate.trapezoid(y_trapezoid, range_domain)
+    return ei
+
+
+@partial(jit, static_argnames='nonlin_f_j')
+def acq_f_ei_j(nonlin_f_j_N, nonlin_f_j,
+               acq_domain,
+               beta, i,
+               m_prior_jnp, var_prior_jnp,
+               range_delta, m_mask, eps_j):
+    m_n = m_prior_jnp[i]
+    S_n = var_prior_jnp[i]
+    Phi = get_Phi_ts_j(nonlin_f_j_N, m_n)
+    mu_predictive = Phi @ m_n
+    sigma_sq_predictive = get_sigma_sq_predictive_j(Phi, S_n, beta)
+    scale_predictive = jnp.sqrt(sigma_sq_predictive)
+    expected_improvement_j_p = Partialj(expected_improvement_j,
+                                        nonlin_f_j, beta, m_n, S_n,
+                                        range_delta, m_mask, eps_j)
+    expected_improvement_j_p_v = jax.vmap(expected_improvement_j_p,
+                                          in_axes=(0, 0, 0))
+    acq = expected_improvement_j_p_v(acq_domain,
+                                     mu_predictive, scale_predictive)
+    x_next = get_x_next_j(acq, acq_domain)
+    return x_next, acq
+
+
+@partial(jit, static_argnames=('param_ct', 'sample_ct', 'acq_domain_len'))
+def init_arr_j(m_prior, var_prior, param_ct, sample_ct, acq_domain_len):
+    prior_ct = sample_ct + 1
+    m_prior_jnp = jnp.zeros((prior_ct, param_ct))
+    m_prior_jnp = m_prior_jnp.at[0].set(m_prior)
+    var_prior_jnp = jnp.zeros((prior_ct, param_ct, param_ct))
+    var_prior_jnp = var_prior_jnp.at[0].set(var_prior)
+    x_n_jnp = jnp.zeros(sample_ct)
+    acq_val_jnp = jnp.zeros((sample_ct, acq_domain_len))
+    t_n_jnp = jnp.zeros(sample_ct)
+    phi_n_jnp = jnp.zeros((sample_ct, param_ct))
+    return x_n_jnp, t_n_jnp, phi_n_jnp, m_prior_jnp, var_prior_jnp, acq_val_jnp
+
+
+@partial(jit, static_argnames='acq_domain_len')
+def init_arr_rp_j(acq_domain_len):
+    rp_jnp = jnp.zeros(acq_domain_len)
+    return rp_jnp
+
+
+@jit
+def init_arr_ei_j(m_mask, range_delta):
+    m_mask_jnp = jnp.array(m_mask, dtype=bool)
+    range_delta_jnp = jnp.array(range_delta)
+    return m_mask_jnp, range_delta_jnp
+
+
+@jit
+def get_rp_next_es_j(x_next, x, rp_prev, mag_rp, scale_rp):
+    return rp_prev + get_rp_es_j(x_next, x, mag_rp, scale_rp)
+
+
+def get_next_posterior_x_j(t, nonlin_f_j_x, beta,
+                           phi, m_prior, var_prior):
+    t_lin = get_t_lin_ts_j(nonlin_f_j_x, t,
+                           phi, m_prior)
+    m_n, S_n = train_posterior_i_j(phi, t_lin, beta,
+                                   m_prior, var_prior)
+    return m_n, S_n
+
+
+@partial(jit, static_argnames='nonlin_f_j')
+def get_next_posterior_j(nonlin_f_j, beta, i,
+                         x_n_jnp, t_n_jnp,
+                         m_prior_jnp, var_prior_jnp):
+    x_i = x_n_jnp[i]
+    t_i = t_n_jnp[i]
+    m_prior = m_prior_jnp[i]
+    var_prior = var_prior_jnp[i]
+    nonlin_f_j_x_i = get_nonlin_f_x_i_j(nonlin_f_j, x_i)
+    phi_i = get_Phi_vec_ts_j(nonlin_f_j_x_i, m_prior)
+    m_n, S_n = get_next_posterior_x_j(t_i, nonlin_f_j_x_i, beta,
+                                      phi_i, m_prior, var_prior)
+    return phi_i, m_n, S_n
+
+
+@partial(jit, static_argnames='nonlin_f_j')
+def posterior_update_j(x_next, acq_val_next, t_next,
+                       nonlin_f_j, beta, i,
+                       x_n_jnp: Array, t_n_jnp: Array,
+                       phi_n_jnp: Array, m_prior_jnp: Array,
+                       var_prior_jnp: Array, acq_val_jnp: Array):
+    x_n_jnp = x_n_jnp.at[i].set(x_next)
+    acq_val_jnp = acq_val_jnp.at[i].set(acq_val_next)
+    t_n_jnp = t_n_jnp.at[i].set(t_next)
+    phi_i, m_n, S_n = get_next_posterior_j(nonlin_f_j, beta, i,
+                                           x_n_jnp, t_n_jnp,
+                                           m_prior_jnp, var_prior_jnp)
+    phi_n_jnp = phi_n_jnp.at[i].set(phi_i)
+    i_posterior = i + 1
+    m_prior_jnp = m_prior_jnp.at[i_posterior].set(m_n)
+    var_prior_jnp = var_prior_jnp.at[i_posterior].set(S_n)
+    return x_n_jnp, t_n_jnp, acq_val_jnp, phi_n_jnp, m_prior_jnp, var_prior_jnp

prisma/bnlr/__init__.py

@@ -0,0 +1,3 @@
+from .posterior_optim import optimize_posterior
+
+__all__ = ['optimize_posterior']

prisma/bnlr/posterior_optim.py

@@ -0,0 +1,12 @@
+from .posterior_optim_iter_jax_lib import optimize_posterior_j
+
+
+def optimize_posterior(beem_integral_f_vector_j,
+                       beta, x_N, t_N,
+                       m_prior, var_prior,
+                       eps=1e-3, i_max=100):
+    Phi, m_N, S_N = optimize_posterior_j(beem_integral_f_vector_j,
+                                         beta, x_N, t_N,
+                                         m_prior, var_prior,
+                                         eps=eps, i_max=i_max)
+    return Phi, m_N, S_N

prisma/bnlr/posterior_optim_iter_jax_lib.py

@@ -0,0 +1,57 @@
+from functools import partial
+from jax import jit, lax
+import jax.numpy as jnp
+from jax.tree_util import Partial as Partialj
+
+from prisma.apf.util.active_parameter_fit_jax_lib \
+    import get_Phi_ts_j, get_t_lin_ts_j, train_posterior_j, \
+    get_nonlin_f_x_j
+
+
+@jit
+def gn_cond_posterior_j(Phi_mprior_mN_SN_i, eps=1e-3, i_max=100):
+    _, m_prior, m_N, _, i = Phi_mprior_mN_SN_i
+    m_delta = m_N - m_prior
+    # jax.debug.print("m_delta: {m_delta}", m_delta=m_delta)
+    # jax.debug.print("eps: {eps}", eps=jnp.any(jnp.abs(m_delta) > eps))
+    # jax.debug.print("i cond: {i}", i=i)
+    return jnp.logical_and(jnp.any(jnp.abs(m_delta) > eps), i < i_max)
+
+
+@jit
+def gn_step_posterior_j(Phi_mprior_mN_SN_i, nonlin_f_N, beta, t_N, var_prior):
+    _, _, m_prior, _, i = Phi_mprior_mN_SN_i
+    Phi = get_Phi_ts_j(nonlin_f_N, m_prior)
+    t_lin = get_t_lin_ts_j(nonlin_f_N, t_N, Phi, m_prior)
+    m_N, S_N = train_posterior_j(Phi, t_lin, beta, m_prior, var_prior)
+    # jax.debug.print("i: {i}", i=i)
+    i += 1
+    return Phi, m_prior, m_N, S_N, i
+
+
+@partial(jit, static_argnames=['param_ct', 'sample_ct'])
+def gn_opt_posterior_j(gn_cond_pj, gn_step_pj, param_ct, sample_ct, m_prior):
+    Phi, _, m_N, S_N, _ \
+        = lax.while_loop(gn_cond_pj, gn_step_pj,
+                         (jnp.full((sample_ct, param_ct), jnp.nan),
+                          jnp.full(param_ct, -jnp.inf),
+                          m_prior,
+                          jnp.full((param_ct, param_ct), jnp.nan),
+                          0))  # TODO: pass rather than generate?
+    return Phi, m_N, S_N
+
+
+@partial(jit, static_argnames='nonlin_f')
+def optimize_posterior_j(nonlin_f, beta, x_N, t_N, m_prior, var_prior,
+                         eps=1e-3, i_max=100):
+    param_ct = len(m_prior)
+    sample_ct = len(x_N)
+    nonlin_f_N = get_nonlin_f_x_j(nonlin_f, x_N)
+    gn_cond = Partialj(gn_cond_posterior_j,
+                       **{'eps': eps, 'i_max': i_max})
+    gn_step = Partialj(gn_step_posterior_j,
+                       **{'nonlin_f_N': nonlin_f_N, 'beta': beta, 't_N': t_N,
+                          'var_prior': var_prior})
+    Phi, m_N, S_N = gn_opt_posterior_j(gn_cond, gn_step, param_ct, sample_ct,
+                                       m_prior)
+    return Phi, m_N, S_N

prisma/bonlr/__init__.py

@@ -0,0 +1,7 @@
+from .bonlr_jax_lib import get_S_N_MAP_j, get_m_N_MAP_j, \
+    get_S_N_MAP_i_j, get_m_N_MAP_i_j, get_sigma_sq_predictive_n_j, \
+    get_sigma_sq_predictive_j
+
+__all__ = ['get_S_N_MAP_j', 'get_m_N_MAP_j', 'get_S_N_MAP_i_j',
+           'get_m_N_MAP_i_j', 'get_sigma_sq_predictive_n_j',
+           'get_sigma_sq_predictive_j',]

prisma/bonlr/bonlr_jax_lib.py

@@ -0,0 +1,152 @@
+import jax
+from jax import Array
+from jax import jit
+import jax.numpy as jnp
+from numbers import Real
+
+
+@jit
+def get_S_N_MAP_j(Phi: Array, beta: Real, S_0: Array) -> Array:
+    '''
+    Maximum a posteriori closed from for posterior covariance matrix.
+
+    Parameters
+    ----------
+    Phi
+        Design matrix of linear model
+    beta
+        Precision (inverse variance) of measurement noise
+    S_0
+        Prior covariance matrix
+
+    Returns
+    -------
+        Posterior covariance matrix
+    '''
+    return jnp.linalg.pinv(jnp.linalg.pinv(S_0) + beta * Phi.T @ Phi)
+
+
+@jit
+def get_m_N_MAP_j(Phi: Array, t: Array, beta: Real, m_0: Array,
+                  S_0: Array, S_N: Array) -> Array:
+    '''
+    Maximum a posteriori closed form for posterior mean.
+
+    Parameters
+    ----------
+    Phi
+        Design matrix of linear model
+    t
+        Measured outputs of system
+    beta
+        Precision (inverse variance) of measurement noise
+    m_0
+        Prior mean
+    S_0
+        Prior covariance matrix
+    S_N
+        Posterior covariance matrix
+
+    Returns
+    -------
+        Posterior mean
+    '''
+    return S_N @ (jnp.linalg.pinv(S_0) @ m_0 + beta * Phi.T @ t)
+
+
+@jit
+def get_S_N_MAP_i_j(phi: Array, beta: Real, S_0: Array) -> Array:
+    '''
+    Online maximum a posteriori closed from for posterior covariance matrix.
+
+    Parameters
+    ----------
+    phi
+        Next sample vector for design matrix of linear model
+    beta
+        Precision (inverse variance) of measurement noise
+    S_0
+        Prior covariance matrix
+
+    Returns
+    -------
+        Posterior covariance matrix
+    '''
+    S_N = jnp.linalg.pinv(jnp.linalg.pinv(S_0) + beta * jnp.outer(phi, phi))
+    return S_N
+
+
+@jit
+def get_m_N_MAP_i_j(phi: Array, t: Array, beta: Real, m_0: Array,
+                    S_0: Array, S_N: Array) -> Array:
+    '''
+    Online maximum a posteriori closed form for posterior mean.
+
+    Parameters
+    ----------
+    phi
+        Next sample vector for design matrix of linear model
+    t
+        Measured outputs of system
+    beta
+        Precision (inverse variance) of measurement noise
+    m_0
+        Prior mean
+    S_0
+        Prior covariance matrix
+    S_N
+        Posterior covariance matrix
+
+    Returns
+    -------
+        Posterior mean
+    '''
+    return S_N @ (jnp.linalg.pinv(S_0) @ m_0 + beta * phi * t)
+
+
+@jit
+def get_sigma_sq_predictive_n_j(phi_n: Array, S_N: Array,
+                                beta: Real) -> Real:
+    '''
+    Bayes' rule for Gaussian distributions predictive distribution variance
+    for a given domain value.
+
+    Parameters
+    ----------
+    phi_n
+        Sample vector for design matrix
+    S_N
+        Posterior covariance matrix
+    beta
+        Precision (inverse variance) of measurement noise
+
+    Returns
+    -------
+        Posterior predictive variance
+    '''
+    return 1. / beta + phi_n @ S_N @ phi_n
+
+
+@jit
+def get_sigma_sq_predictive_j(Phi_test: Array, S_N: Array,
+                              beta: Real) -> Array:
+    '''
+    Get predictive distribution variance over domain.
+
+    Parameters
+    ----------
+    Phi_test
+        Design matrix over domain
+    S_N
+        Posterior covariance matrix
+    beta
+        Precision (inverse variance) of measurement noise
+
+    Returns
+    -------
+        Posterior predictive variance over domain
+    '''
+    get_sigma_sq_predictive_n_j_v = jax.vmap(get_sigma_sq_predictive_n_j,
+                                             in_axes=(0, None, None))
+    sigma_sq_N = get_sigma_sq_predictive_n_j_v(Phi_test, S_N, beta)
+    return sigma_sq_N