prisma-oed 0.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (58) hide show
  1. prisma_oed-0.0.1/LICENSE.txt +30 -0
  2. prisma_oed-0.0.1/LICENSES_bundled.txt +8 -0
  3. prisma_oed-0.0.1/PKG-INFO +97 -0
  4. prisma_oed-0.0.1/README.md +69 -0
  5. prisma_oed-0.0.1/pyproject.toml +63 -0
  6. prisma_oed-0.0.1/setup.cfg +4 -0
  7. prisma_oed-0.0.1/src/prisma/__init__.py +1 -0
  8. prisma_oed-0.0.1/src/prisma/apf/__init__.py +0 -0
  9. prisma_oed-0.0.1/src/prisma/apf/api/__init__.py +7 -0
  10. prisma_oed-0.0.1/src/prisma/apf/api/active_parameter_fit.py +182 -0
  11. prisma_oed-0.0.1/src/prisma/apf/util/__init__.py +0 -0
  12. prisma_oed-0.0.1/src/prisma/apf/util/active_parameter_fit_jax_lib.py +262 -0
  13. prisma_oed-0.0.1/src/prisma/bnlr/__init__.py +3 -0
  14. prisma_oed-0.0.1/src/prisma/bnlr/posterior_optim.py +12 -0
  15. prisma_oed-0.0.1/src/prisma/bnlr/posterior_optim_iter_jax_lib.py +57 -0
  16. prisma_oed-0.0.1/src/prisma/bonlr/__init__.py +7 -0
  17. prisma_oed-0.0.1/src/prisma/bonlr/bonlr_jax_lib.py +152 -0
  18. prisma_oed-0.0.1/src/prisma/calc/__init__.py +0 -0
  19. prisma_oed-0.0.1/src/prisma/calc/dim_analysis/__init__.py +3 -0
  20. prisma_oed-0.0.1/src/prisma/calc/dim_analysis/dimensional_analysis_lib.py +34 -0
  21. prisma_oed-0.0.1/src/prisma/calc/metric/__init__.py +3 -0
  22. prisma_oed-0.0.1/src/prisma/calc/metric/metric_lib.py +21 -0
  23. prisma_oed-0.0.1/src/prisma/experiment/__init__.py +9 -0
  24. prisma_oed-0.0.1/src/prisma/experiment/performance_survey.py +95 -0
  25. prisma_oed-0.0.1/src/prisma/experiment/spectroscopy.py +35 -0
  26. prisma_oed-0.0.1/src/prisma/experiment/spectroscopy_series.py +234 -0
  27. prisma_oed-0.0.1/src/prisma/file/__init__.py +3 -0
  28. prisma_oed-0.0.1/src/prisma/file/io_lib.py +12 -0
  29. prisma_oed-0.0.1/src/prisma/format/__init__.py +3 -0
  30. prisma_oed-0.0.1/src/prisma/format/format_lib.py +6 -0
  31. prisma_oed-0.0.1/src/prisma/nlr/__init__.py +5 -0
  32. prisma_oed-0.0.1/src/prisma/nlr/nlr_parameter_fit_lib.py +59 -0
  33. prisma_oed-0.0.1/src/prisma/physics_models/__init__.py +0 -0
  34. prisma_oed-0.0.1/src/prisma/physics_models/jax/__init__.py +3 -0
  35. prisma_oed-0.0.1/src/prisma/physics_models/jax/phys_mod_jax.py +51 -0
  36. prisma_oed-0.0.1/src/prisma/physics_models/jax/phys_mod_jax_lib.py +211 -0
  37. prisma_oed-0.0.1/src/prisma/physics_models/python/__init__.py +4 -0
  38. prisma_oed-0.0.1/src/prisma/physics_models/python/phys_mod.py +128 -0
  39. prisma_oed-0.0.1/src/prisma/physics_models/python/phys_mod_lib.py +79 -0
  40. prisma_oed-0.0.1/src/prisma/plot/__init__.py +19 -0
  41. prisma_oed-0.0.1/src/prisma/plot/matplotlibrc.py +54 -0
  42. prisma_oed-0.0.1/src/prisma/plot/plot_lib.py +441 -0
  43. prisma_oed-0.0.1/src/prisma/samplers/__init__.py +0 -0
  44. prisma_oed-0.0.1/src/prisma/samplers/abstract/__init__.py +3 -0
  45. prisma_oed-0.0.1/src/prisma/samplers/abstract/abstract_samplers.py +226 -0
  46. prisma_oed-0.0.1/src/prisma/samplers/instruments/__init__.py +0 -0
  47. prisma_oed-0.0.1/src/prisma/samplers/instruments/nanonis/__init__.py +4 -0
  48. prisma_oed-0.0.1/src/prisma/samplers/instruments/nanonis/nanonis_sampler.py +223 -0
  49. prisma_oed-0.0.1/src/prisma/samplers/pseudorandom/__init__.py +0 -0
  50. prisma_oed-0.0.1/src/prisma/samplers/pseudorandom/function/__init__.py +6 -0
  51. prisma_oed-0.0.1/src/prisma/samplers/pseudorandom/function/function_samplers.py +118 -0
  52. prisma_oed-0.0.1/src/prisma/samplers/pseudorandom/noise/__init__.py +3 -0
  53. prisma_oed-0.0.1/src/prisma/samplers/pseudorandom/noise/noise_samplers.py +136 -0
  54. prisma_oed-0.0.1/src/prisma_oed.egg-info/PKG-INFO +97 -0
  55. prisma_oed-0.0.1/src/prisma_oed.egg-info/SOURCES.txt +56 -0
  56. prisma_oed-0.0.1/src/prisma_oed.egg-info/dependency_links.txt +1 -0
  57. prisma_oed-0.0.1/src/prisma_oed.egg-info/requires.txt +7 -0
  58. prisma_oed-0.0.1/src/prisma_oed.egg-info/top_level.txt +1 -0
@@ -0,0 +1,30 @@
1
+ BSD 3-Clause License
2
+
3
+ Copyright 2026 National Technology & Engineering Solutions of Sandia, LLC
4
+ (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S.
5
+ Government retains certain rights in this software.
6
+
7
+ Redistribution and use in source and binary forms, with or without
8
+ modification, are permitted provided that the following conditions are met:
9
+
10
+ 1. Redistributions of source code must retain the above copyright notice, this
11
+ list of conditions and the following disclaimer.
12
+
13
+ 2. Redistributions in binary form must reproduce the above copyright notice,
14
+ this list of conditions and the following disclaimer in the documentation
15
+ and/or other materials provided with the distribution.
16
+
17
+ 3. Neither the name of the copyright holder nor the names of its
18
+ contributors may be used to endorse or promote products derived from
19
+ this software without specific prior written permission.
20
+
21
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
22
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
23
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
24
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
25
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
26
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
27
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
28
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
29
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
30
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
@@ -0,0 +1,8 @@
1
+ This software uses methods derived from the SciPy utility
2
+ scipy.optimize.curve_fit to compute the Fisher Information Matrix (FIM) of
3
+ parameter fits. SciPy is compatibly licensed.
4
+
5
+ Name: scipy.optimize.curve_fit
6
+ Files: src/prisma/nlr/nlr_parameter_fit_lib.py
7
+ License: BSD-3-Clause
8
+ For details, see src/prisma/nlr/SciPy/LICENSE.txt
@@ -0,0 +1,97 @@
1
+ Metadata-Version: 2.4
2
+ Name: prisma-oed
3
+ Version: 0.0.1
4
+ Summary: PRISMA (Physics-informed Regression and Inference for Statistics-driven Measurement Acquisition) is a Bayesian online nonlinear regression algorithm for active learning and SOED in spectroscopy experiments.
5
+ Author: Xavier Tricoche, Andy Mounce, Sam Lilak, Devika Mehta, Leo Diaz, Luca Basso, Ezra Bussmann
6
+ Author-email: Dante Basile <dbasile@purdue.edu>
7
+ License-Expression: BSD-3-Clause
8
+ Project-URL: Repository, https://github.com/sandialabs/PRISMA.git
9
+ Keywords: machine learning,ML,Bayesian,nonlinear regression,spectroscopy,curve fitting,online learning,active learning,optimal experimental design,OED,sequential optimal experimental design,SOED,instrument control,real-time control,JAX
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Operating System :: OS Independent
12
+ Classifier: Development Status :: 4 - Beta
13
+ Classifier: Framework :: Jupyter
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
16
+ Requires-Python: >=3.13
17
+ Description-Content-Type: text/markdown
18
+ License-File: LICENSE.txt
19
+ License-File: LICENSES_bundled.txt
20
+ Requires-Dist: ipykernel>=6.29
21
+ Requires-Dist: jax>=0.6
22
+ Requires-Dist: line_profiler>=4.2
23
+ Requires-Dist: matplotlib>=3.10
24
+ Requires-Dist: nanonis_spm>=1.0
25
+ Requires-Dist: numpy>=2.3
26
+ Requires-Dist: scipy>=1.15
27
+ Dynamic: license-file
28
+
29
+ # PRISMA
30
+
31
+ PRISMA: Physics-informed Regression and Inference for Statistics-driven Measurement Acquisition
32
+
33
+ ## Description
34
+
35
+ PRISMA is a Bayesian active learning algorithm that automates and optimizes
36
+ spectroscopy measurements. PRISMA uses prior parameter distributions and known
37
+ physics model equations to optimize measurement acquisition and fit model
38
+ parameters in real-time. Active learning increases the speed and accuracy of
39
+ spectroscopy measurements while maintaining confidence and rigor.
40
+
41
+ ## Authors and acknowledgment
42
+
43
+ Developed by Dante Basile, Xavier Tricoche, Sam Lilak, Devika Mehta, Leo Diaz, Luca Basso, Andy Mounce, and Ezra Bussmann.
44
+
45
+ This software was developed as part of work at the CINT (Center for Integrated Nanotechnologies), an Office of Science User Facility operated for the U.S. Department of Energy.
46
+
47
+ ## Installation
48
+
49
+ Install in editable mode for development:
50
+
51
+ ```bash
52
+ python3 -m pip install -e .
53
+ ```
54
+
55
+ Install from `requirements.txt` for reproducible environment:
56
+
57
+ ```bash
58
+ python3 -m pip install -r requirements.txt
59
+ ```
60
+
61
+ ## Usage
62
+
63
+ Examples of PRISMA-directed spectroscopy experiments can be found in the
64
+ `experiments` folder. These Jupyter notebooks show how to use
65
+ PRISMA with both pseudorandom simulated data and instrument measurements.
66
+
67
+ As a general overview, PRISMA defines the Bayesian online nonlinear regression
68
+ algorithm in the `ActiveParameterFit` abstract base class in
69
+ `src/prisma/apf/active_parameter_fit.py`. Subclasses implement this interface
70
+ by implementing the `acquisition_f` abstract method with their acquisition
71
+ function of choice to define their sampling strategy. For example,
72
+ `ActiveParameterFitSigmaSqPred` samples at the maximum of the predictive
73
+ variance and can be imported in the API at
74
+ `prisma.apf.ActiveParameterFitSigmaSqPred`.
75
+
76
+ Sampler classes in `src/prisma/samplers` provide drivers that allow PRISMA to acquire measurements
77
+ at the domain values it specifies.
78
+
79
+ Call the `fit` method of an initialized `ActiveParameterFit` instance to
80
+ run PRISMA's active learning algorithm. This will optimize measurement
81
+ acquisitions and fit model parameters.
82
+
83
+ ## License
84
+
85
+ [BSD-3-Clause](https://opensource.org/license/bsd-3-clause)
86
+
87
+ ## Citation
88
+
89
+ ```bibtex
90
+ @software{prisma2026github,
91
+ author = {Dante Basile and Xavier Tricoche and Sam Lilak and Devika Mehta and Leo Diaz and Luca Basso and Andy Mounce and Ezra Bussmann},
92
+ title = {{PRISMA}: {P}hysics-informed {R}egression and {I}nference for {S}tatistics-driven {M}easurement {A}cquisition},
93
+ url = {https://github.com/sandialabs/PRISMA},
94
+ version = {0.0.1},
95
+ year = {2026},
96
+ }
97
+ ```
@@ -0,0 +1,69 @@
1
+ # PRISMA
2
+
3
+ PRISMA: Physics-informed Regression and Inference for Statistics-driven Measurement Acquisition
4
+
5
+ ## Description
6
+
7
+ PRISMA is a Bayesian active learning algorithm that automates and optimizes
8
+ spectroscopy measurements. PRISMA uses prior parameter distributions and known
9
+ physics model equations to optimize measurement acquisition and fit model
10
+ parameters in real-time. Active learning increases the speed and accuracy of
11
+ spectroscopy measurements while maintaining confidence and rigor.
12
+
13
+ ## Authors and acknowledgment
14
+
15
+ Developed by Dante Basile, Xavier Tricoche, Sam Lilak, Devika Mehta, Leo Diaz, Luca Basso, Andy Mounce, and Ezra Bussmann.
16
+
17
+ This software was developed as part of work at the CINT (Center for Integrated Nanotechnologies), an Office of Science User Facility operated for the U.S. Department of Energy.
18
+
19
+ ## Installation
20
+
21
+ Install in editable mode for development:
22
+
23
+ ```bash
24
+ python3 -m pip install -e .
25
+ ```
26
+
27
+ Install from `requirements.txt` for reproducible environment:
28
+
29
+ ```bash
30
+ python3 -m pip install -r requirements.txt
31
+ ```
32
+
33
+ ## Usage
34
+
35
+ Examples of PRISMA-directed spectroscopy experiments can be found in the
36
+ `experiments` folder. These Jupyter notebooks show how to use
37
+ PRISMA with both pseudorandom simulated data and instrument measurements.
38
+
39
+ As a general overview, PRISMA defines the Bayesian online nonlinear regression
40
+ algorithm in the `ActiveParameterFit` abstract base class in
41
+ `src/prisma/apf/active_parameter_fit.py`. Subclasses implement this interface
42
+ by implementing the `acquisition_f` abstract method with their acquisition
43
+ function of choice to define their sampling strategy. For example,
44
+ `ActiveParameterFitSigmaSqPred` samples at the maximum of the predictive
45
+ variance and can be imported in the API at
46
+ `prisma.apf.ActiveParameterFitSigmaSqPred`.
47
+
48
+ Sampler classes in `src/prisma/samplers` provide drivers that allow PRISMA to acquire measurements
49
+ at the domain values it specifies.
50
+
51
+ Call the `fit` method of an initialized `ActiveParameterFit` instance to
52
+ run PRISMA's active learning algorithm. This will optimize measurement
53
+ acquisitions and fit model parameters.
54
+
55
+ ## License
56
+
57
+ [BSD-3-Clause](https://opensource.org/license/bsd-3-clause)
58
+
59
+ ## Citation
60
+
61
+ ```bibtex
62
+ @software{prisma2026github,
63
+ author = {Dante Basile and Xavier Tricoche and Sam Lilak and Devika Mehta and Leo Diaz and Luca Basso and Andy Mounce and Ezra Bussmann},
64
+ title = {{PRISMA}: {P}hysics-informed {R}egression and {I}nference for {S}tatistics-driven {M}easurement {A}cquisition},
65
+ url = {https://github.com/sandialabs/PRISMA},
66
+ version = {0.0.1},
67
+ year = {2026},
68
+ }
69
+ ```
@@ -0,0 +1,63 @@
1
+ [build-system]
2
+ requires = ["setuptools >= 77.0.3"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "prisma-oed"
7
+ version = "0.0.1"
8
+ dependencies = [
9
+ "ipykernel >= 6.29",
10
+ "jax >= 0.6",
11
+ "line_profiler >= 4.2",
12
+ "matplotlib >= 3.10",
13
+ "nanonis_spm >= 1.0",
14
+ "numpy >= 2.3",
15
+ "scipy >= 1.15",
16
+ ]
17
+ requires-python = ">= 3.13"
18
+ authors = [
19
+ { name="Dante Basile", email="dbasile@purdue.edu" },
20
+ { name="Xavier Tricoche" },
21
+ { name="Andy Mounce" },
22
+ { name="Sam Lilak" },
23
+ { name="Devika Mehta" },
24
+ { name="Leo Diaz" },
25
+ { name="Luca Basso" },
26
+ { name="Ezra Bussmann" },
27
+ ]
28
+ description = """\
29
+ PRISMA (Physics-informed Regression and Inference for Statistics-driven \
30
+ Measurement Acquisition) is a Bayesian online nonlinear regression algorithm \
31
+ for active learning and SOED in spectroscopy experiments.\
32
+ """
33
+ readme = "README.md"
34
+ license = "BSD-3-Clause"
35
+ license-files = ["LICEN[CS]E*"]
36
+ keywords = [
37
+ "machine learning",
38
+ "ML",
39
+ "Bayesian",
40
+ "nonlinear regression",
41
+ "spectroscopy",
42
+ "curve fitting",
43
+ "online learning",
44
+ "active learning",
45
+ "optimal experimental design",
46
+ "OED",
47
+ "sequential optimal experimental design",
48
+ "SOED",
49
+ "instrument control",
50
+ "real-time control",
51
+ "JAX",
52
+ ]
53
+ classifiers = [
54
+ "Programming Language :: Python :: 3",
55
+ "Operating System :: OS Independent",
56
+ "Development Status :: 4 - Beta",
57
+ "Framework :: Jupyter",
58
+ "Intended Audience :: Science/Research",
59
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
60
+ ]
61
+
62
+ [project.urls]
63
+ Repository = "https://github.com/sandialabs/PRISMA.git"
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1 @@
1
+ __version__ = '0.0.1'
File without changes
@@ -0,0 +1,7 @@
1
+ from .active_parameter_fit import ActiveParameterFitSigmaSqPred, \
2
+ ActiveParameterFitRP, ActiveParameterFitEI, \
3
+ AbstractActiveParameterFit
4
+
5
+ __all__ = ['ActiveParameterFitSigmaSqPred', 'ActiveParameterFitRP',
6
+ 'ActiveParameterFitEI',
7
+ 'AbstractActiveParameterFit']
@@ -0,0 +1,182 @@
1
+ from abc import ABC, abstractmethod
2
+ import jax
3
+ import jax.numpy as jnp
4
+ from jax.stages import Wrapped
5
+ import numpy as np
6
+ from numpy.typing import ArrayLike
7
+
8
+ from ..util.active_parameter_fit_jax_lib import get_nonlin_f_x_j, \
9
+ acq_f_sigma_sq_predictive_j, acq_f_rp_j, \
10
+ acq_f_ei_j, \
11
+ init_arr_j, init_arr_rp_j, init_arr_ei_j, \
12
+ posterior_update_j
13
+ from prisma.file import save_np_val
14
+ from prisma.plot import plot_predictive_train_iter, \
15
+ plot_acq_train_iter, plot_predictive_acq_train_iter, plot_w_rmse, \
16
+ plot_w_i_err, plot_predictive_rmse
17
+ from prisma.samplers.abstract import AbstractDomainSampler
18
+
19
+
20
+ class AbstractActiveParameterFit(ABC):
21
+
22
+ def __init__(self, nonlin_f_j: Wrapped, acq_domain: ArrayLike,
23
+ t_sampler: AbstractDomainSampler,
24
+ beta: float, m_0: ArrayLike, S_0: ArrayLike,
25
+ sample_ct: int):
26
+ # TODO: initialize all jit argument vars as jnp
27
+ self.nonlin_f_j = nonlin_f_j
28
+ self.acq_domain = acq_domain
29
+ self.t_sampler = t_sampler
30
+ self.beta = beta
31
+ self.param_ct = len(m_0)
32
+ self.sample_ct = sample_ct
33
+ self.acq_domain_len = len(acq_domain)
34
+ self.nonlin_f_j_N = get_nonlin_f_x_j(nonlin_f_j, acq_domain)
35
+ self.x_n_jnp, self.t_n_jnp, self.phi_n_jnp, \
36
+ self.m_prior_jnp, self.var_prior_jnp, self.acq_val_jnp \
37
+ = init_arr_j(m_0, S_0, self.param_ct, sample_ct,
38
+ self.acq_domain_len)
39
+ self.x_n_np = self.t_n_np = self.phi_n_np \
40
+ = self.m_prior_np = self.var_prior_np = self.acq_val_np = None
41
+
42
+ def fit(self):
43
+ for i in range(self.sample_ct):
44
+ x_next, acq_val_next = self._acquisition_f(i)
45
+ t_next = self.t_sampler.sample(x_next)
46
+ self.x_n_jnp, self.t_n_jnp, self.acq_val_jnp, \
47
+ self.phi_n_jnp, self.m_prior_jnp, self.var_prior_jnp \
48
+ = posterior_update_j(x_next, acq_val_next, t_next,
49
+ self.nonlin_f_j, self.beta,
50
+ i, self.x_n_jnp, self.t_n_jnp,
51
+ self.phi_n_jnp, self.m_prior_jnp,
52
+ self.var_prior_jnp, self.acq_val_jnp)
53
+
54
+ def fit_sampler_reset(self):
55
+ self.t_sampler.reset()
56
+ self.fit()
57
+
58
+ def init_np_arrs(self):
59
+ self.x_n_np, self.t_n_np, self.phi_n_np, \
60
+ self.m_prior_np, self.var_prior_np, self.acq_val_np \
61
+ = self._jax_arrs_to_np_arrs(self.x_n_jnp, self.t_n_jnp,
62
+ self.phi_n_jnp, self.m_prior_jnp,
63
+ self.var_prior_jnp, self.acq_val_jnp)
64
+
65
+ def save_np_arrs(self, save_dir):
66
+ save_np_val(save_dir, "x_n.npy", self.x_n_np)
67
+ save_np_val(save_dir, "t_n.npy", self.t_n_np)
68
+ save_np_val(save_dir, "phi_n.npy", self.phi_n_np)
69
+ save_np_val(save_dir, "m_prior.npy", self.m_prior_np)
70
+ save_np_val(save_dir, "var_prior.npy", self.var_prior_np)
71
+ save_np_val(save_dir, "acq_n.npy", self.acq_val_np)
72
+
73
+ def plot_predictive_train_iter(self, y=None, title=True,
74
+ title_samples=False):
75
+ plot_predictive_train_iter(self.nonlin_f_j, self.acq_domain,
76
+ self.beta, self.m_prior_jnp,
77
+ self.var_prior_jnp,
78
+ y=y,
79
+ x_n=self.x_n_jnp, t_n=self.t_n_jnp,
80
+ title=title, title_samples=title_samples)
81
+
82
+ def plot_acq_train_iter(self):
83
+ plot_acq_train_iter(self.acq_domain, self.acq_val_jnp)
84
+
85
+ def plot_predictive_acq_train_iter(self, y=None, title=True,
86
+ title_samples=False):
87
+ plot_predictive_acq_train_iter(self.nonlin_f_j, self.acq_domain,
88
+ self.beta, self.m_prior_jnp,
89
+ self.var_prior_jnp,
90
+ self.acq_val_jnp,
91
+ y=y,
92
+ x_n=self.x_n_jnp, t_n=self.t_n_jnp,
93
+ title=title,
94
+ title_samples=title_samples)
95
+
96
+ def plot_w_rmse(self, w_act):
97
+ plot_w_rmse(w_act, m_N_l=self.m_prior_jnp)
98
+
99
+ def plot_w_i_err(self, w_act, i, w_i_name=None):
100
+ plot_w_i_err(w_act, i, m_N_l=self.m_prior_jnp, w_i_name=w_i_name)
101
+
102
+ def plot_predictive_rmse(self, t, y):
103
+ plot_predictive_rmse(self.nonlin_f_j, self.acq_domain, t, y,
104
+ m_N_l=self.m_prior_jnp)
105
+
106
+ @abstractmethod
107
+ def _acquisition_f(self, i):
108
+ ...
109
+
110
+ def _jax_arrs_to_np_arrs(self, *jax_arrs):
111
+ return tuple(map(np.array, jax_arrs))
112
+
113
+
114
+ class ActiveParameterFitSigmaSqPred(AbstractActiveParameterFit):
115
+ '''
116
+ Predictive variance
117
+ '''
118
+
119
+ def _acquisition_f(self, i):
120
+ x_next, acq = acq_f_sigma_sq_predictive_j(self.nonlin_f_j_N,
121
+ self.acq_domain,
122
+ self.beta, i,
123
+ self.m_prior_jnp,
124
+ self.var_prior_jnp)
125
+ return x_next, acq
126
+
127
+
128
+ class ActiveParameterFitRP(AbstractActiveParameterFit):
129
+ '''
130
+ Normal repulsion
131
+ '''
132
+
133
+ def __init__(self, nonlin_f_j: Wrapped, acq_domain: ArrayLike,
134
+ t_sampler: AbstractDomainSampler,
135
+ beta: float, m_0: ArrayLike, S_0: ArrayLike,
136
+ sample_ct: int,
137
+ magnitude_rp: float, scale_rp: float):
138
+ super().__init__(nonlin_f_j, acq_domain, t_sampler, beta, m_0, S_0,
139
+ sample_ct)
140
+ self.mag_rp = magnitude_rp
141
+ self.scl_rp = scale_rp
142
+ self.rp_jnp = init_arr_rp_j(self.acq_domain_len)
143
+
144
+ def _acquisition_f(self, i):
145
+ x_next, acq, rp_next = acq_f_rp_j(self.nonlin_f_j_N, self.acq_domain,
146
+ self.beta, i,
147
+ self.m_prior_jnp, self.var_prior_jnp,
148
+ self.mag_rp, self.scl_rp,
149
+ self.rp_jnp)
150
+ self.rp_jnp = rp_next
151
+ return x_next, acq
152
+
153
+
154
+ class ActiveParameterFitEI(AbstractActiveParameterFit):
155
+ '''
156
+ Expected Improvement
157
+ '''
158
+
159
+ def __init__(self, nonlin_f_j: Wrapped, acq_domain: ArrayLike,
160
+ t_sampler: AbstractDomainSampler,
161
+ beta: float, m_0: ArrayLike, S_0: ArrayLike,
162
+ sample_ct: int,
163
+ range_delta: ArrayLike,
164
+ m_mask: ArrayLike | None = None):
165
+ super().__init__(nonlin_f_j, acq_domain, t_sampler, beta, m_0, S_0,
166
+ sample_ct)
167
+ if m_mask is None:
168
+ m_mask = np.ones(self.param_ct)
169
+ if jax.config.read('jax_enable_x64'):
170
+ self.eps_j = jnp.finfo(jnp.float64).eps
171
+ else:
172
+ self.eps_j = jnp.finfo(jnp.float32).eps
173
+ self.m_mask_jnp, self.range_delta_jnp = init_arr_ei_j(m_mask,
174
+ range_delta)
175
+
176
+ def _acquisition_f(self, i):
177
+ x_next, acq, = acq_f_ei_j(self.nonlin_f_j_N, self.nonlin_f_j,
178
+ self.acq_domain, self.beta, i,
179
+ self.m_prior_jnp, self.var_prior_jnp,
180
+ self.range_delta_jnp, self.m_mask_jnp,
181
+ self.eps_j)
182
+ return x_next, acq
File without changes