fig bug with nans (#136)

pymc-devs · Dec 29, 2023 · 7fb9c39 · 7fb9c39
1 parent 5de28ff
commit 7fb9c39
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/pymc_bart/__init__.py b/pymc_bart/__init__.py
@@ -25,7 +25,7 @@
 )
 
 __all__ = ["BART", "PGBART"]
-__version__ = "0.5.6"
+__version__ = "0.5.7"
 
 
 pm.STEP_METHODS = list(pm.STEP_METHODS) + [PGBART]
diff --git a/pymc_bart/pgbart.py b/pymc_bart/pgbart.py
@@ -157,7 +157,7 @@ def __init__(
 
  for idx, rule in enumerate(self.split_rules):
  if rule is ContinuousSplitRule:
- self.X[:, idx] = jitter_duplicated(self.X[:, idx], np.std(self.X[:, idx]))
+ self.X[:, idx] = jitter_duplicated(self.X[:, idx], np.nanstd(self.X[:, idx]))
 
  init_mean = self.bart.Y.mean()
  self.num_observations = self.X.shape[0]
@@ -700,7 +700,7 @@ def jitter_duplicated(array: npt.NDArray[np.float_], std: float) -> npt.NDArray[
  if are_whole_number(array):
  seen = []
  for idx, num in enumerate(array):
- if num in seen:
+ if num in seen and not np.isnan(num):
  array[idx] = num + np.random.normal(0, std / 12)
  else:
  seen.append(num)
@@ -711,8 +711,7 @@ def jitter_duplicated(array: npt.NDArray[np.float_], std: float) -> npt.NDArray[
 @njit
 def are_whole_number(array: npt.NDArray[np.float_]) -> np.bool_:
  """Check if all values in array are whole numbers"""
- new_array = np.mod(array, 1)
- return np.all(new_array == 0)
+ return np.all(np.mod(array[~np.isnan(array)], 1) == 0)
 
 
 def logp(point, out_vars, vars, shared): # pylint: disable=redefined-builtin