bartz 0.0__tar.gz → 0.0.1__tar.gz

{bartz-0.0 → bartz-0.0.1}/PKG-INFO

@@ -1,20 +1,21 @@
 Metadata-Version: 2.1
 Name: bartz
-Version: 0.0
+Version: 0.0.1
 Summary: A JAX implementation of BART
 Home-page: https://github.com/Gattocrucco/bartz
 License: MIT
 Author: Giacomo Petrillo
 Author-email: info@giacomopetrillo.com
-Requires-Python: >=3.12,<4.0
+Requires-Python: >=3.10,<4.0
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: appnope (>=0.1.4,<0.2.0)
-Requires-Dist: jax (>=0.4.25,<0.5.0)
-Requires-Dist: jaxlib (>=0.4.25,<0.5.0)
-Requires-Dist: numpy (>=1.26.4,<2.0.0)
-Requires-Dist: scipy (>=1.12.0,<2.0.0)
+Requires-Dist: jax (>=0.4.23,<0.5.0)
+Requires-Dist: jaxlib (>=0.4.23,<0.5.0)
+Requires-Dist: numpy (>=1.25.2,<2.0.0)
+Requires-Dist: scipy (>=1.11.4,<2.0.0)
 Project-URL: Bug Tracker, https://github.com/Gattocrucco/bartz/issues
 Project-URL: Repository, https://github.com/Gattocrucco/bartz
 Description-Content-Type: text/markdown

{bartz-0.0 → bartz-0.0.1}/pyproject.toml

@@ -28,7 +28,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "bartz"
-version = "0.0"
+version = "0.0.1"
 description = "A JAX implementation of BART"
 authors = ["Giacomo Petrillo <info@giacomopetrillo.com>"]
 license = "MIT"
@@ -42,16 +42,16 @@ packages = [
 "Bug Tracker" = "https://github.com/Gattocrucco/bartz/issues"
 
 [tool.poetry.dependencies]
-python = "^3.12"
-jax = "^0.4.25"
-jaxlib = "^0.4.25"
-numpy = "^1.26.4"
-scipy = "^1.12.0"
-appnope = "^0.1.4"
+python = "^3.10"
+jax = "^0.4.23"
+jaxlib = "^0.4.23"
+numpy = "^1.25.2"
+scipy = "^1.11.4"
 
 [tool.poetry.group.dev.dependencies]
 ipython = "^8.22.2"
 matplotlib = "^3.8.3"
+appnope = "^0.1.4"
 
 [tool.poetry.group.test.dependencies]
 coverage = "^7.4.3"

{bartz-0.0 → bartz-0.0.1}/src/bartz/__init__.py

@@ -28,7 +28,7 @@ A jax implementation of BART
 See the manual at https://gattocrucco.github.io/bartz/docs
 """
 
-__version__ = '0.0'
+__version__ = '0.0.1'
 
 from .interface import BART
 

{bartz-0.0 → bartz-0.0.1}/src/bartz/debug.py

@@ -65,15 +65,11 @@ def print_tree(leaf_tree, var_tree, split_tree, print_all=False):
         else:
             link = ' '
 
-        if print_all:
-            max_number = len(leaf_tree) - 1
-            ndigits = len(str(max_number))
-            number = str(index).rjust(ndigits)
-            number = f' {number} '
-        else:
-            number = ''
+        max_number = len(leaf_tree) - 1
+        ndigits = len(str(max_number))
+        number = str(index).rjust(ndigits)
 
-        print(f'{number}{indent}{first_indent}{link}{node_str}')
+        print(f' {number} {indent}{first_indent}{link}{node_str}')
 
         indent += next_indent
         unused = unused or is_leaf

{bartz-0.0 → bartz-0.0.1}/src/bartz/interface.py

@@ -102,14 +102,6 @@ class BART:
 
     Attributes
     ----------
-    offset : float
-        The prior mean of the latent mean function.
-    scale : float
-        The prior standard deviation of the latent mean function.
-    lamda : float
-        The prior harmonic mean of the error variance.
-    ntree : int
-        The number of trees.
     yhat_train : array (ndpost, n)
         The conditional posterior mean at `x_train` for each MCMC iteration.
     yhat_train_mean : array (n,)
@@ -122,6 +114,18 @@ class BART:
         The standard deviation of the error.
     first_sigma : array (nskip,)
         The standard deviation of the error in the burn-in phase.
+    offset : float
+        The prior mean of the latent mean function.
+    scale : float
+        The prior standard deviation of the latent mean function.
+    lamda : float
+        The prior harmonic mean of the error variance.
+    sigest : float or None
+        The estimated standard deviation of the error used to set `lamda`.
+    ntree : int
+        The number of trees.
+    maxdepth : int
+        The maximum depth of the trees.
 
     Methods
     -------
@@ -166,17 +170,17 @@ class BART:
         y_train, y_train_fmt = self._process_response_input(y_train)
         self._check_same_length(x_train, y_train)
         
-        lamda = self._process_noise_variance_settings(x_train, y_train, sigest, sigdf, sigquant, lamda)
         offset = self._process_offset_settings(y_train, offset)
         scale = self._process_scale_settings(y_train, k)
+        lamda, sigest = self._process_noise_variance_settings(x_train, y_train, sigest, sigdf, sigquant, lamda, offset)
 
         splits, max_split = self._determine_splits(x_train, numcut)
         x_train = self._bin_predictors(x_train, splits)
 
         y_train = self._transform_input(y_train, offset, scale)
-        lamda = lamda / scale
+        lamda_scaled = lamda / (scale * scale)
 
-        mcmc_state = self._setup_mcmc(x_train, y_train, max_split, lamda, sigdf, power, base, maxdepth, ntree)
+        mcmc_state = self._setup_mcmc(x_train, y_train, max_split, lamda_scaled, sigdf, power, base, maxdepth, ntree)
         final_state, burnin_trace, main_trace = self._run_mcmc(mcmc_state, ndpost, nskip, keepevery, printevery, seed)
 
         sigma = self._extract_sigma(main_trace, scale)
@@ -184,8 +188,10 @@ class BART:
 
         self.offset = offset
         self.scale = scale
-        self.lamda = lamda * scale
+        self.lamda = lamda
+        self.sigest = sigest
         self.ntree = ntree
+        self.maxdepth = maxdepth
         self.sigma = sigma
         self.first_sigma = first_sigma
 
@@ -261,25 +267,25 @@ class BART:
         assert get_length(x1) == get_length(x2)
 
     @staticmethod
-    def _process_noise_variance_settings(x_train, y_train, sigest, sigdf, sigquant, lamda):
+    def _process_noise_variance_settings(x_train, y_train, sigest, sigdf, sigquant, lamda, offset):
         if lamda is not None:
-            return lamda
+            return lamda, None
         else:
             if sigest is not None:
                 sigest2 = sigest * sigest
             elif y_train.size < 2:
                 sigest2 = 1
             elif y_train.size <= x_train.shape[0]:
-                sigest2 = jnp.var(y_train)
+                sigest2 = jnp.var(y_train - offset)
             else:
-                _, chisq, rank, _ = jnp.linalg.lstsq(x_train.T, y_train)
+                _, chisq, rank, _ = jnp.linalg.lstsq(x_train.T, y_train - offset)
                 chisq = chisq.squeeze(0)
                 dof = len(y_train) - rank
                 sigest2 = chisq / dof
             alpha = sigdf / 2
             invchi2 = jaxext.scipy.stats.invgamma.ppf(sigquant, alpha) / 2
             invchi2rid = invchi2 * sigdf
-            return sigest2 / invchi2rid
+            return sigest2 / invchi2rid, jnp.sqrt(sigest2)
 
     @staticmethod
     def _process_offset_settings(y_train, offset):

{bartz-0.0 → bartz-0.0.1}/src/bartz/mcmcloop.py

@@ -148,14 +148,17 @@ def make_simple_print_callback(printevery):
         prune_prop = bart['prune_prop_count'] / prop_total
         grow_acc = bart['grow_acc_count'] / bart['grow_prop_count']
         prune_acc = bart['prune_acc_count'] / bart['prune_prop_count']
-        n_total = n_burn + n_save
+        n_total = n_burn + n_save * n_skip
         debug.callback(simple_print_callback_impl, burnin, i_total, n_total, grow_prop, grow_acc, prune_prop, prune_acc, printevery)
     return callback
 
 def simple_print_callback_impl(burnin, i_total, n_total, grow_prop, grow_acc, prune_prop, prune_acc, printevery):
     if (i_total + 1) % printevery == 0:
         burnin_flag = ' (burnin)' if burnin else ''
-        print(f'Iteration {i_total + 1:4d}/{n_total:d} '
+        total_str = str(n_total)
+        ndigits = len(total_str)
+        i_str = str(i_total + 1).rjust(ndigits)
+        print(f'Iteration {i_str}/{total_str} '
             f'P_grow={grow_prop:.2f} P_prune={prune_prop:.2f} '
             f'A_grow={grow_acc:.2f} A_prune={prune_acc:.2f}{burnin_flag}')