Repository: kjappelbaum/pyepal (language: python)

Issue:
Minimised objectives lead to mirrored palinstance.means
Hi! I noticed some unexpected behaviour when minimising objectives. This isn't really a bug and the workaround is quite simple.

## Steps to reproduce the problem

- Set `goals` to any combination involving min
- Perform active learning using exhaust loop
- Observe that `palinstance.means` gives expected values for maximised objectives
- Observe that `palinstance.means` gives mirrored values for minimised objectives (multiplying it by -1 gives correct predictions)
- Observe that using `.predict` with the models gives mirrored values for minimised objectives (multiplying it by -1 gives correct predictions)

[Here](https://deepnote.com/project/Pyepal-mirrored-results-Yeuc095QRE6ToMXFtwFjnA/%2Fbinh_korn_goals_issue.ipynb) is a notebook detailing the issue.

## Specifications

- pyepal version: 0.6.1
- operating system: DeepNote cloud
- Python version: 3.7

Minimised objectives lead to mirrored palinstance.means
Hi! I noticed some unexpected behaviour when minimising objectives. This isn't really a bug and the workaround is quite simple.

## Steps to reproduce the problem

- Set `goals` to any combination involving min
- Perform active learning using exhaust loop
- Observe that `palinstance.means` gives expected values for maximised objectives
- Observe that `palinstance.means` gives mirrored values for minimised objectives (multiplying it by -1 gives correct predictions)
- Observe that using `.predict` with the models gives mirrored values for minimised objectives (multiplying it by -1 gives correct predictions)

[Here](https://deepnote.com/project/Pyepal-mirrored-results-Yeuc095QRE6ToMXFtwFjnA/%2Fbinh_korn_goals_issue.ipynb) is a notebook detailing the issue.

## Specifications

- pyepal version: 0.6.1
- operating system: DeepNote cloud
- Python version: 3.7



Discussion / hints:
you're a fan of DeepNote? :D 

> you're a fan of DeepNote? :D

I'm a fan of anyone who gives me free swag and computing hours 👍 
you're a fan of DeepNote? :D 

> you're a fan of DeepNote? :D

I'm a fan of anyone who gives me free swag and computing hours 👍 

Fix patch:
diff --git a/README.md b/README.md
index 1afb066..981d2e4 100644
--- a/README.md
+++ b/README.md
@@ -125,7 +125,7 @@ class PALSklearn(PALBase):
             means.append(mean.reshape(-1, 1))
             stds.append(std.reshape(-1, 1))
 
-        self.means = np.hstack(mean)
+        self._means = np.hstack(mean)
         self.std = np.hstack(stds)
 ```
 
diff --git a/pyepal/pal/pal_base.py b/pyepal/pal/pal_base.py
index 9f39d40..198ef08 100644
--- a/pyepal/pal/pal_base.py
+++ b/pyepal/pal/pal_base.py
@@ -85,7 +85,7 @@ def __init__(  # pylint:disable=too-many-arguments
                 Defaults to 1/9.
             goals (List[str], optional): If a list, provide "min" for every objective
                 that shall be minimized and "max" for every objective
-                that shall be maximized. Defaults to None, which means
+                that shall be maximized. Defaults to None, which _means
                 that the code maximizes all objectives.
             coef_var_threshold (float, optional): Use only points with
                 a coefficient of variation below this threshold
@@ -113,8 +113,8 @@ def __init__(  # pylint:disable=too-many-arguments
         design_space_size = len(X_design)
         self.coef_var_threshold = validate_coef_var(coef_var_threshold)
         self.coef_var_mask = np.array([True] * design_space_size)
-        # means/std are the model predictions
-        self.means: np.array = None
+        # _means/std are the model predictions
+        self._means: np.array = None
         self.std: np.array = None
 
         self.design_space = X_design
@@ -127,7 +127,7 @@ def __init__(  # pylint:disable=too-many-arguments
         # self._y contains the data as provided by the user
         self.y = np.zeros((design_space_size, self.ndim))  # pylint:disable=invalid-name
         self._y = self.y
-        # measurement_uncertainity is provided in update_train_set by the user
+        # measurement_uncertainty is provided in update_train_set by the user
         self.measurement_uncertainty = np.zeros((design_space_size, self.ndim))
         self._has_train_set = False
 
@@ -159,8 +159,8 @@ def _reset(self):
 
         self.coef_var_mask = np.array([True] * self.number_design_points)
 
-        # means/std are the model predictions
-        self.means: np.array = None
+        # _means/std are the model predictions
+        self._means: np.array = None
         self.std: np.array = None
         self.beta = None
 
@@ -171,7 +171,7 @@ def _reset(self):
             (self.number_design_points, self.ndim)
         )  # pylint:disable=invalid-name
         self._y = self.y
-        # measurement_uncertainity is provided in update_train_set by the user
+        # measurement_uncertainty is provided in update_train_set by the user
         self.measurement_uncertainty = np.zeros((self.number_design_points, self.ndim))
         self._has_train_set = False
 
@@ -243,7 +243,7 @@ def number_pareto_optimal_points(self):
 
     @property
     def number_discarded_points(self):
-        """Return the nnumber of discarded points"""
+        """Return the number of discarded points"""
         return sum(self.discarded)
 
     @property
@@ -259,7 +259,17 @@ def number_sampled_points(self):
     @property
     def hyperrectangle_sizes(self):
         """Return the sizes of the hyperrectangles"""
-        return _uncertainty(self.rectangle_ups, self.rectangle_lows, self.means)
+        return _uncertainty(self.rectangle_ups, self.rectangle_lows, self._means)
+
+    @property
+    def means(self):
+        """Return the means predicted by the model"""
+        if self._means is None:
+            raise ValueError(
+                "Predicted means is None. Execute run_one_step() \
+                    to obtain predicted means for each model."
+            )
+        return self._means * self.goals
 
     def _update_beta(self):
         """Update beta according to section 7.2. of the epsilon-PAL paper"""
@@ -320,7 +330,7 @@ def _crossvalidate(self):
         for sampled_idx in sample_subset:
             # make sure that we do not run into errors due to np.nan
             if self.sampled[sampled_idx].sum() == self.ndim:
-                # copy here is important, otherewise all
+                # copy here is important, otherwise all
                 # points we set to False remain False
                 self.sampled = sampled_original.copy()
                 self.sampled[sampled_idx, :] = False
@@ -331,7 +341,7 @@ def _crossvalidate(self):
                 self._train()
                 self._predict()
                 error = mean_absolute_error(
-                    self.y[sampled_idx], self.means[sampled_idx]
+                    self.y[sampled_idx], self._means[sampled_idx]
                 )
                 errors.append(error)
 
@@ -346,7 +356,7 @@ def should_cross_validate(self):
 
     def _update_hyperrectangles(self, new_indices: np.ndarray = None):
         """Computes new hyperrectangles based on beta,
-        the means and the standard deviations.
+        the _means and the standard deviations.
         If the iteration is > 0,
         then it uses iterative intersection to ensure that the size of the
         hyperrectangles is decreasing.
@@ -357,7 +367,7 @@ def _update_hyperrectangles(self, new_indices: np.ndarray = None):
                 Instead, it will just use the scaled edges based on the
                 model's prediction. Defaults to None.
         """
-        lows, ups = _get_uncertainty_regions(self.means, self.std, np.sqrt(self.beta))
+        lows, ups = _get_uncertainty_regions(self._means, self.std, np.sqrt(self.beta))
         if self.iteration == 1:
             # initialization
             self.rectangle_lows, self.rectangle_ups = lows, ups
@@ -390,11 +400,11 @@ def _update_coef_var_mask(self):
         # this will fail if everything is zero in this case,
         # I feel we might just give up for now,
         # in the future we can think fo just applying a shift
-        if self.means.sum() != 0:
-            means_no_zero = self.means.copy()
-            means_no_zero[means_no_zero == 0] = np.median(means_no_zero)
+        if self._means.sum() != 0:
+            _means_no_zero = self._means.copy()
+            _means_no_zero[_means_no_zero == 0] = np.median(_means_no_zero)
             self.coef_var_mask = (
-                np.max(self.std / means_no_zero, axis=1) < self.coef_var_threshold
+                np.max(self.std / _means_no_zero, axis=1) < self.coef_var_threshold
             )
         else:
             mean_variation = self.std.mean()
@@ -438,7 +448,7 @@ def _replace_by_measurements(
         predictions for the sampled points we use the data that
         was actually measured and the actual uncertainty."""
         if replace_mean:
-            self.means[self.sampled] = self.y[self.sampled]
+            self._means[self.sampled] = self.y[self.sampled]
         if replace_std:
             self.std[self.sampled] = self.measurement_uncertainty[self.sampled]
 
@@ -464,7 +474,7 @@ def run_one_step(  # pylint:disable=too-many-arguments
                 and not only from the unclassified and Pareto optimal ones
             use_coef_var (bool): If True, uses the coefficient of variation instead of
                 the unscaled rectangle sizes
-            replace_mean (bool): If true uses the measured means for the sampled points
+            replace_mean (bool): If true uses the measured _means for the sampled points
             replace_std (bool): If true uses the measured standard deviation for the
                 sampled points
 
@@ -662,9 +672,9 @@ def augment_design_space(  # pylint: disable=invalid-name
             self.coef_var_mask, np.array([True] * number_new_points), 0
         )
 
-        # means/std are the model predictions
-        self.means = np.append(
-            self.means, np.full([number_new_points, self.ndim], np.nan), 0
+        # _means/std are the model predictions
+        self._means = np.append(
+            self._means, np.full([number_new_points, self.ndim], np.nan), 0
         )
         self.std = np.append(
             self.std, np.full([number_new_points, self.ndim], np.nan), 0
@@ -677,7 +687,7 @@ def augment_design_space(  # pylint: disable=invalid-name
         # self._y contains the data as provided by the user
         self._y = self.y
 
-        # measurement_uncertainity is provided in update_train_set by the user
+        # measurement_uncertainty is provided in update_train_set by the user
         self.measurement_uncertainty = np.append(
             self.measurement_uncertainty, np.zeros((number_new_points, self.ndim)), 0
         )
@@ -759,7 +769,7 @@ def sample(
             sampled_idx = _get_max_wt_all(
                 self.rectangle_lows,
                 self.rectangle_ups,
-                self.means,
+                self._means,
                 sampled_mask,
                 pooling_method,
                 use_coef_var,
@@ -768,7 +778,7 @@ def sample(
             sampled_idx = _get_max_wt(
                 self.rectangle_lows,
                 self.rectangle_ups,
-                self.means,
+                self._means,
                 self.pareto_optimal,
                 self.unclassified,
                 sampled_mask,
diff --git a/pyepal/pal/pal_coregionalized.py b/pyepal/pal/pal_coregionalized.py
index 5af259a..3f67255 100644
--- a/pyepal/pal/pal_coregionalized.py
+++ b/pyepal/pal/pal_coregionalized.py
@@ -94,7 +94,7 @@ def _predict(self):
             means.append(mean.reshape(-1, 1))
             stds.append(std.reshape(-1, 1))
 
-        self.means = np.hstack(means)
+        self._means = np.hstack(means)
         self.std = np.hstack(stds)
 
     def _set_hyperparameters(self):
diff --git a/pyepal/pal/pal_finite_ensemble.py b/pyepal/pal/pal_finite_ensemble.py
index 9128bf3..d18fae0 100644
--- a/pyepal/pal/pal_finite_ensemble.py
+++ b/pyepal/pal/pal_finite_ensemble.py
@@ -46,7 +46,7 @@ def _ensemble_train_one_finite_width(  # pylint:disable=too-many-arguments, too-
     ensemble_size: Sequence[int],
 ):
     from jax import random  # pylint:disable=import-outside-toplevel
-    from jax.api import grad, jit, vmap  # pylint:disable=import-outside-toplevel
+    from jax import grad, jit, vmap  # pylint:disable=import-outside-toplevel
 
     model = models[i]
     optimizer = optimizers[i]
@@ -158,7 +158,7 @@ def _set_data(self):
         self.design_space = self.design_space_scaler.fit_transform(self.design_space)
 
     def _train(self):
-        for i in range(len(self.models)):
+        for i, _ in enumerate(self.models):
             params, scaler = _ensemble_train_one_finite_width(
                 i,
                 self.models,
@@ -183,5 +183,5 @@ def _predict(self):
             means.append(mean.reshape(-1, 1))
             stds.append(std.reshape(-1, 1))
 
-        self.means = np.hstack(means)
+        self._means = np.hstack(means)
         self.std = np.hstack(stds)
diff --git a/pyepal/pal/pal_gbdt.py b/pyepal/pal/pal_gbdt.py
index 71c10b3..595fc9f 100644
--- a/pyepal/pal/pal_gbdt.py
+++ b/pyepal/pal/pal_gbdt.py
@@ -141,7 +141,7 @@ def _predict(self):
             means.append(mean.reshape(-1, 1))
             stds.append(std.reshape(-1, 1))
 
-        self.means = np.hstack(means)
+        self._means = np.hstack(means)
         self.std = np.hstack(stds)
 
     def _set_hyperparameters(self):
diff --git a/pyepal/pal/pal_gpflowgpr.py b/pyepal/pal/pal_gpflowgpr.py
index adeeeb5..38b9d12 100644
--- a/pyepal/pal/pal_gpflowgpr.py
+++ b/pyepal/pal/pal_gpflowgpr.py
@@ -120,7 +120,7 @@ def _predict(self):
             means.append(mean.reshape(-1, 1))
             stds.append(np.sqrt(std.reshape(-1, 1)))
 
-        self.means = np.hstack(means)
+        self._means = np.hstack(means)
         self.std = np.hstack(stds)
 
     def _set_hyperparameters(self):
diff --git a/pyepal/pal/pal_gpy.py b/pyepal/pal/pal_gpy.py
index bbb8fdf..088b4a0 100644
--- a/pyepal/pal/pal_gpy.py
+++ b/pyepal/pal/pal_gpy.py
@@ -96,7 +96,7 @@ def _predict(self):
             means.append(mean.reshape(-1, 1))
             stds.append(std.reshape(-1, 1))
 
-        self.means = np.hstack(means)
+        self._means = np.hstack(means)
         self.std = np.hstack(stds)
 
     def _set_hyperparameters(self):
diff --git a/pyepal/pal/pal_neural_tangent.py b/pyepal/pal/pal_neural_tangent.py
index 5078c24..6415ee7 100644
--- a/pyepal/pal/pal_neural_tangent.py
+++ b/pyepal/pal/pal_neural_tangent.py
@@ -126,7 +126,7 @@ def __init__(self, *args, **kwargs):
 
     def _set_data(self):
         self.design_space = self.design_space_scaler.fit_transform(self.design_space)
-        for i in range(len(self.models)):
+        for i, _ in enumerate(self.models):
 
             predict_fn, scaler = _set_one_infinite_width_model(
                 i,
@@ -144,7 +144,7 @@ def _train(self):
 
     def _predict(self):
         means, stds = [], []
-        for i in range(len(self.models)):
+        for i, _ in enumerate(self.models):
             mean, std = _predict_one_infinite_width_model(
                 i,
                 self.models,
@@ -154,5 +154,5 @@ def _predict(self):
             means.append(mean.reshape(-1, 1))
             stds.append(std.reshape(-1, 1))
 
-        self.means = np.hstack(means)
+        self._means = np.hstack(means)
         self.std = np.hstack(stds)
diff --git a/pyepal/pal/pal_sklearn.py b/pyepal/pal/pal_sklearn.py
index f6027f6..5647cd9 100644
--- a/pyepal/pal/pal_sklearn.py
+++ b/pyepal/pal/pal_sklearn.py
@@ -96,7 +96,7 @@ def _predict(self):
             means.append(mean.reshape(-1, 1))
             stds.append(std.reshape(-1, 1))
 
-        self.means = np.hstack(means)
+        self._means = np.hstack(means)
         self.std = np.hstack(stds)
 
     def _set_hyperparameters(self):
diff --git a/pyepal/pal/validate_inputs.py b/pyepal/pal/validate_inputs.py
index ed9e294..199f064 100644
--- a/pyepal/pal/validate_inputs.py
+++ b/pyepal/pal/validate_inputs.py
@@ -426,19 +426,17 @@ def validate_positive_integer_list(
 
     if not isinstance(seq, collections.Sequence):
         if (not isinstance(seq, int)) or (seq < 1):
-            raise ValueError("{} must be a positive integer".format(parameter_name))
+            raise ValueError(f"{parameter_name} must be a positive integer")
         return [seq] * ndim
 
     if not len(seq) == ndim:
         raise ValueError(
-            "If you provide a sequence for {} its length must match \
-                the number of objectives".format(
-                parameter_name
-            )
+            f"If you provide a sequence for {parameter_name} its length must match \
+                the number of objectives"
         )
     for elem in seq:
         if (not isinstance(elem, int)) or (elem < 1):
-            raise ValueError("{} must be a positive integer".format(parameter_name))
+            raise ValueError(f"{parameter_name} must be a positive integer")
 
     return seq
 
diff --git a/pyepal/plotting/__init__.py b/pyepal/plotting/__init__.py
index fa70bd8..d855604 100644
--- a/pyepal/plotting/__init__.py
+++ b/pyepal/plotting/__init__.py
@@ -280,7 +280,7 @@ def plot_residuals(  # pylint:disable=invalid-name
 
     for index in range(num_targets):
         sampled = palinstance.sampled[:, index]
-        fitted = palinstance.means[sampled, index]
+        fitted = palinstance._means[sampled, index]  # pylint:disable=protected-access
         residuals = y[sampled, index] - fitted
         ax[index].scatter(fitted, residuals)
         ax[index].spines["top"].set_color("none")
@@ -445,7 +445,7 @@ def plot_learning_curve(  # pylint:disable=dangerous-default-value, too-many-arg
                 test_errors[metric_name].append(
                     metric_function(
                         observations[test_idx],
-                        palinstance.means[test_idx],
+                        palinstance._means[test_idx],  # pylint:disable=protected-access
                     )
                 )
         for metric_name, results in test_errors.items():
diff --git a/requirements.txt b/requirements.txt
index 2f39e76..025fbf4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,3 @@
-numpy==1.*
-numba==0.*
-scikit-learn==0.*
+numpy==1.19.*
+numba==0.53.*
+scikit-learn==0.24.*
diff --git a/setup.py b/setup.py
index d1e05f4..d883511 100644
--- a/setup.py
+++ b/setup.py
@@ -18,11 +18,11 @@
 
 import versioneer
 
-with open("requirements.txt", "r") as fh:
+with open("requirements.txt", "r", encoding="utf-8") as fh:
     REQUIREMENTS = fh.readlines()
 
 
-with open("README.md", encoding="utf-8") as fh:
+with open("README.md", "r", encoding="utf-8") as fh:
     LONG_DESCRIPTION = fh.read()
 
 gpy_requirements = [
@@ -30,8 +30,16 @@
     "matplotlib==3.*",
 ]
 gbdt_requirements = ["lightgbm==3.*"]
-neural_tangents_requirements = ["neural_tangents==0.*", "jaxlib==0.*"]
-gpflow_requirements = ["gpflow"]
+neural_tangents_requirements = [
+    "jaxlib<=0.1.65",
+    "jax<=0.2.12",
+    "neural_tangents==0.3.*",
+]
+gpflow_requirements = [
+    "tensorflow==2.5.*",
+    "tensorflow-probability==0.12.*",
+    "gpflow==2.2.*",
+]
 setup(
     name="pyepal",
     version=versioneer.get_version(),


Test patch:
diff --git a/tests/test_pal_base.py b/tests/test_pal_base.py
index cdb7c47..8a254b7 100644
--- a/tests/test_pal_base.py
+++ b/tests/test_pal_base.py
@@ -18,7 +18,9 @@
 import numpy as np
 import pytest
 
+from pyepal.models.gpr import build_coregionalized_model
 from pyepal.pal.pal_base import PALBase
+from pyepal.pal.pal_coregionalized import PALCoregionalized
 
 
 def test_pal_base(make_random_dataset):
@@ -67,7 +69,7 @@ def test_reset_classification(make_random_dataset):
     highs = np.zeros((100, 3))
 
     means = np.full((100, 3), 1)
-    palinstance.means = means
+    palinstance._means = means
     palinstance.std = np.full((100, 3), 0.1)
     pareto_optimal = np.array([False] * 98 + [True, True])
     sampled = np.array([[False] * 3, [False] * 3, [False] * 3, [False] * 3])
@@ -89,6 +91,7 @@ def test_reset_classification(make_random_dataset):
 def test_update_train_set(make_random_dataset):
     """Check if the update of the training set works"""
     X, y = make_random_dataset  # pylint:disable=invalid-name
+
     palinstance = PALBase(X, ["model"], 3)
     assert not palinstance._has_train_set
     assert palinstance.sampled.sum() == 0
@@ -99,6 +102,43 @@ def test_update_train_set(make_random_dataset):
     assert (palinstance.y[0] == y[0, :]).all()
 
 
+def test_means_property(binh_korn_points):
+    """based on #185 and example in
+    https://deepnote.com/project/
+    Pyepal-mirrored-results-Yeuc095QRE6ToMXFtwFjnA/
+    %2Fbinh_korn_goals_issue.ipynb/#00012-ebf692a0-f995-4e25-98b6-8d34a7db7f54"""
+    x, points = binh_korn_points  # pylint:disable=invalid-name
+    np.random.seed(10)
+    points[:, 0] = -points[:, 0]
+    model = build_coregionalized_model(x, points)
+    palinstance = PALCoregionalized(
+        x, [model], 2, goals=["min", "max"], epsilon=[0.5, 0.5], restarts=3
+    )
+    indices = np.array([0, 1, 5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 99])
+    palinstance.update_train_set(indices, points[indices])
+    palinstance.run_one_step()
+    assert np.all(palinstance.means[:, 0] >= 0)
+
+
+def test_means_property2(binh_korn_points):
+    """based on #185 and example in
+    https://deepnote.com/project/
+    Pyepal-mirrored-results-Yeuc095QRE6ToMXFtwFjnA/
+    %2Fbinh_korn_goals_issue.ipynb/#00012-ebf692a0-f995-4e25-98b6-8d34a7db7f54"""
+    x, points = binh_korn_points  # pylint:disable=invalid-name
+    points = -points
+    np.random.seed(10)
+    model = build_coregionalized_model(x, points)
+    palinstance = PALCoregionalized(
+        x, [model], 2, goals=["min", "min"], epsilon=[0.5, 0.5], restarts=3
+    )
+    indices = np.array([0, 1, 5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 99])
+    palinstance.update_train_set(indices, points[indices])
+    palinstance.run_one_step()
+    assert np.all(palinstance.means[:, 0] >= 0)
+    assert np.all(palinstance.means[:, 1] >= 0)
+
+
 def test_augment_design_space(make_random_dataset):
     """Testing the basic functionality of the augmentation method
     Does NOT test the re-classification step, which needs a model"""
@@ -122,7 +162,7 @@ def test_augment_design_space(make_random_dataset):
     highs = np.zeros((100, 3))
 
     means = np.full((100, 3), 1)
-    palinstance.means = means
+    palinstance._means = means
     palinstance.std = np.full((100, 3), 0.1)
     pareto_optimal = np.array([False] * 98 + [True, True])
     sampled = np.array([[False] * 3, [False] * 3, [False] * 3, [False] * 3])
@@ -199,7 +239,7 @@ def test_sample(make_random_dataset):
     )
 
     means = np.array([[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]])
-    palinstance.means = means
+    palinstance._means = means
     pareto_optimal = np.array([False, False, True, True])
     sampled = np.array([[False] * 4, [False] * 4, [False] * 4, [False] * 4])
     unclassified = np.array([True, True, False, False])
@@ -255,7 +295,7 @@ def test__update_hyperrectangles(make_random_dataset):
     X, _ = make_random_dataset  # pylint:disable=invalid-name
     palinstance = PALBase(X, ["model"], 4, beta_scale=1)
 
-    palinstance.means = np.array([[0, 0, 1, 0], [0, 0, 0, 1]])
+    palinstance._means = np.array([[0, 0, 1, 0], [0, 0, 0, 1]])
     palinstance.std = np.array([[0, 0, 0, 0], [1, 0, 0, 0]])
     with pytest.raises(TypeError):
         # Beta is not defined
@@ -275,7 +315,7 @@ def test__update_hyperrectangles(make_random_dataset):
     assert palinstance.rectangle_lows[1][0] < -1
     assert palinstance.rectangle_ups[1][0] > 1
 
-    assert len(palinstance.hyperrectangle_sizes) == len(palinstance.means)
+    assert len(palinstance.hyperrectangle_sizes) == len(palinstance._means)
 
 
 def test_orchestration_run_one_step(make_random_dataset):
@@ -298,7 +338,7 @@ def test__replace_by_measurements(make_random_dataset):
     assert palinstance.measurement_uncertainty.sum() == 0
     sample_idx = np.array([1, 2, 3, 4])
     palinstance.update_train_set(sample_idx, y[sample_idx], y[sample_idx])
-    palinstance.means = palinstance.measurement_uncertainty
+    palinstance._means = palinstance.measurement_uncertainty
     palinstance.std = palinstance.measurement_uncertainty
     palinstance._replace_by_measurements()
     assert (palinstance.y == palinstance.std).all()
@@ -309,7 +349,7 @@ def test__update_coef_var_mask(make_random_dataset):
     X, _ = make_random_dataset  # pylint:disable=invalid-name
     palinstance = PALBase(X[:2], ["model"], 3, beta_scale=1)
 
-    palinstance.means = np.array([[1, 1, 1, 1], [1, 1, 1, 1]])
+    palinstance._means = np.array([[1, 1, 1, 1], [1, 1, 1, 1]])
     palinstance.std = np.array([[0, 0, 0, 0], [3, 1, 1, 1]])
 
     assert (palinstance.coef_var_mask == np.array([True, True])).all()
@@ -317,7 +357,7 @@ def test__update_coef_var_mask(make_random_dataset):
 
     assert (palinstance.coef_var_mask == np.array([True, False])).all()
 
-    palinstance.means = np.array([[0, 0, 0, 0], [0, 0, 0, 0]])
+    palinstance._means = np.array([[0, 0, 0, 0], [0, 0, 0, 0]])
     palinstance.std = np.array([[0, 0, 0, 0], [3, 1, 1, 1]])
 
     assert (palinstance.coef_var_mask == np.array([True, False])).all()
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index 3797ee8..ee8601b 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -103,7 +103,7 @@ def test_plot_residuals(make_random_dataset):
     with pytest.raises(ValueError):
         plot_residuals(np.array([[1, 2, 1]] * 100), palinstance)
 
-    palinstance.means = means
+    palinstance._means = means  # pylint:disable=protected-access
     with pytest.raises(AssertionError):
         plot_residuals(np.array([1] * 100), palinstance)
 
@@ -126,7 +126,7 @@ def test_plot_jointplot(make_random_dataset):
     with pytest.raises(ValueError):
         plot_jointplot(np.array([[1, 2, 1]] * 100), palinstance)
 
-    palinstance.means = means
+    palinstance._means = means  # pylint:disable=protected-access
     palinstance.std = means
     palinstance.beta = 1
     with pytest.raises(AssertionError):
diff --git a/tests/test_validate_inputs.py b/tests/test_validate_inputs.py
index e98a5e3..b1f029d 100644
--- a/tests/test_validate_inputs.py
+++ b/tests/test_validate_inputs.py
@@ -254,7 +254,7 @@ def test_validate_interquartile_scaler():
 
 def test_validate_optimizers():
     """Make sure we correctly check/standardize the optimizers."""
-    from jax.api import jit  # pylint:disable=import-outside-toplevel
+    from jax import jit  # pylint:disable=import-outside-toplevel
     from jax.experimental import optimizers  # pylint:disable=import-outside-toplevel
 
     from pyepal.models.nt import JaxOptimizer  # pylint:disable=import-outside-toplevel