bartz 0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

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/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/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/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}')
 

{bartz-0.0.dist-info → bartz-0.0.1.dist-info}/METADATA

@@ -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.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+bartz/__init__.py,sha256=PL-vhhEHoVMWOPLG_M45TIZVkbQia5riJbQboy-BNH8,1333
+bartz/debug.py,sha256=FHnCalpK1uO1CN9QQ5DPj70JKR4Thltzp9o0BeYthIo,5741
+bartz/grove.py,sha256=v2k10EBjgi2aLCsGvM01z0z--9Xv4ApBOxpke-6gIYM,10309
+bartz/interface.py,sha256=GBwLwqEF_6EmeteFtsPw6ANdisnvMoWi_fKBJiQq-Vc,16129
+bartz/jaxext.py,sha256=FK5j1zfW1yR4-yPKcD7ZvKSkVQ5--jHjQpVCl4n4gXY,2844
+bartz/mcmcloop.py,sha256=N815-eJxsS_X85okXRO2kSOlikw8dPN05_krm0iT9Sg,7321
+bartz/mcmcstep.py,sha256=acy_2rSIEXV5BzqLY96aQaqlsxtalxyO3Q4gPvUMRVU,35912
+bartz/prepcovars.py,sha256=3ddDOtNNop3Ba2Kgy_dZ6apFydtwaEXH3uXSmmKf9Fs,4421
+bartz-0.0.1.dist-info/LICENSE,sha256=heuIJZQK9IexJYC-fYHoLUrgj8HG8yS3G072EvKh-94,1073
+bartz-0.0.1.dist-info/METADATA,sha256=zDW1dM58gV7c_8ZTjEtTt_tcXabbz5roZBf36EdLxls,933
+bartz-0.0.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+bartz-0.0.1.dist-info/RECORD,,

bartz-0.0.dist-info/RECORD

@@ -1,12 +0,0 @@
-bartz/__init__.py,sha256=KpeWEIKkXq7YyooRdLZK6atfLJimi0hIa8-xqBPnDWQ,1331
-bartz/debug.py,sha256=SBdFJd9gXtBEw2b3R3pv14ARwKCDkHiycOuDm0iVEhA,5846
-bartz/grove.py,sha256=v2k10EBjgi2aLCsGvM01z0z--9Xv4ApBOxpke-6gIYM,10309
-bartz/interface.py,sha256=pPUOYpHCci_NADF58jrbee2tcnaKZYzwgK0FF5XuwJU,15823
-bartz/jaxext.py,sha256=FK5j1zfW1yR4-yPKcD7ZvKSkVQ5--jHjQpVCl4n4gXY,2844
-bartz/mcmcloop.py,sha256=ZxasSfWZzuYT_rOnSL8iSzIGQQS3ZtcBUhdXis2K-II,7207
-bartz/mcmcstep.py,sha256=acy_2rSIEXV5BzqLY96aQaqlsxtalxyO3Q4gPvUMRVU,35912
-bartz/prepcovars.py,sha256=3ddDOtNNop3Ba2Kgy_dZ6apFydtwaEXH3uXSmmKf9Fs,4421
-bartz-0.0.dist-info/LICENSE,sha256=heuIJZQK9IexJYC-fYHoLUrgj8HG8yS3G072EvKh-94,1073
-bartz-0.0.dist-info/METADATA,sha256=kKCGPdD_bOxJ3cLePc1R19kS3AzwGtFK_aR9ljBX6HY,869
-bartz-0.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-bartz-0.0.dist-info/RECORD,,