plotastrodata 1.9.2__tar.gz → 1.9.3__tar.gz

{plotastrodata-1.9.2/plotastrodata.egg-info → plotastrodata-1.9.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plotastrodata
-Version: 1.9.2
+Version: 1.9.3
 Summary: plotastrodata is a tool for astronomers to create figures from FITS files and perform fundamental data analyses with ease.
 Home-page: https://github.com/yusukeaso-astron/plotastrodata
 Download-URL: https://github.com/yusukeaso-astron/plotastrodata

{plotastrodata-1.9.2 → plotastrodata-1.9.3}/plotastrodata/__init__.py

@@ -1,4 +1,4 @@
 import warnings
 
 warnings.simplefilter('ignore', FutureWarning)
-__version__ = '1.9.2'
+__version__ = '1.9.3'

{plotastrodata-1.9.2 → plotastrodata-1.9.3}/plotastrodata/fitting_utils.py

@@ -36,7 +36,7 @@ def logp(x: np.ndarray) -> float:
 
 def _get_GR(samples: np.ndarray, nwalkers: int, ndata: int, dim: int
             ) -> np.ndarray:
-    # Gelman-Rubin statistics #
+    """Calculate the Gelman-Rubin statistics."""
     B = np.std(np.mean(samples, axis=1), axis=0)
     W = np.mean(np.std(samples, axis=1), axis=0)
     V = (len(samples[0]) - 1) / len(samples[0]) * W \
@@ -46,6 +46,20 @@ def _get_GR(samples: np.ndarray, nwalkers: int, ndata: int, dim: int
     return GR
 
 
+def _check_GR(samples: np.ndarray, nwalkers: int, ndata: int, dim: int,
+              i: int, ntry: int = 1, grcheck: bool = False) -> int:
+    if not grcheck:
+        return ntry
+
+    GR = _get_GR(samples=samples, nwalkers=nwalkers, ndata=ndata, dim=dim)
+    if np.max(GR) <= 1.25:
+        return ntry
+
+    if i == ntry:
+        print(f'!!! Max GR >1.25 during {ntry:d} trials.!!!')
+    return i
+
+
 class EmceeCorner():
     warnings.simplefilter('ignore', RuntimeWarning)
 
@@ -87,6 +101,72 @@ class EmceeCorner():
         self.percent = percent
         self.ndata = 10000 if xdata is None else len(xdata)
 
+    def _get_pos0(self, ntemps: int, nwalkers: int, pt: bool) -> np.ndarray:
+        """Create initial walker positions within parameter bounds."""
+        lower = self.bounds[:, 0]
+        upper = self.bounds[:, 1]
+        width = upper - lower
+        pos0 = np.random.rand(ntemps, nwalkers, self.dim) * width + lower
+        return pos0 if pt else pos0[0]
+
+    def _run_sampler(self, pos0: np.ndarray, pt: bool,
+                     ncores: int, ntemps: int,
+                     nsteps: int, nwalkers: int) -> object:
+        """Create and run the sampler, then return it."""
+        if pt:
+            sampler_cls = ptemcee.Sampler
+            sampler_kwargs = {'ntemps': ntemps,
+                              'nwalkers': nwalkers, 'dim': self.dim,
+                              'logl': self.logl, 'logp': self.logp}
+        else:
+            if ncores > 1:
+                print('Use logl as log_prob_fn to avoid function-in-function.')
+                log_prob_fn = self.logl
+            else:
+                def log_prob_fn(x):
+                    return self.logp(x) + self.logl(x)
+
+            sampler_cls = emcee.EnsembleSampler
+            sampler_kwargs = {'nwalkers': nwalkers, 'ndim': self.dim,
+                              'log_prob_fn': log_prob_fn}
+        if ncores > 1:
+            with Pool(ncores) as pool:
+                sampler = sampler_cls(**sampler_kwargs, pool=pool)
+        else:
+            sampler = sampler_cls(**sampler_kwargs, pool=None)
+        sampler.run_mcmc(pos0, nsteps)
+        return sampler
+
+    def _get_samples(self, sampler, nburnin: int, pt: bool) -> np.ndarray:
+        """Extract post-burn-in samples from sampler chain."""
+        if pt:
+            return sampler.chain[0, :, nburnin:, :]  # temperatures, walkers, steps, dim
+        else:
+            return sampler.chain[:, nburnin:, :]  # walkers, steps, dim
+
+    def _get_lnp_popt(self, sampler, pt: bool, nburnin: int,
+                      ) -> tuple[np.ndarray, np.ndarray]:
+        """Get log probabilities and best-fit parameters from sampler."""
+        if pt:
+            lnp = sampler.logprobability[0]  # 0th temperature chain
+            chain = sampler.chain[0]
+        else:
+            lnp = sampler.lnprobability
+            chain = sampler.chain
+        idx_best = np.unravel_index(np.argmax(lnp), lnp.shape)
+        popt = chain[idx_best]
+        lnp = lnp[:, nburnin:]
+        return lnp, popt
+
+    def _get_percentiles(self, samples: np.ndarray
+                         ) -> tuple[float, float, float]:
+        """Compute summary statistics (percentiles) from MCMC samples."""
+        s = samples.reshape(-1, self.dim)
+        plow = np.percentile(s, self.percent[0], axis=0)
+        pmid = np.percentile(s, 50, axis=0)
+        phigh = np.percentile(s, self.percent[1], axis=0)
+        return plow, pmid, phigh
+
     def fit(self, nwalkersperdim: int = 2,
             ntemps: int = 1, nsteps: int = 1000,
             nburnin: int = 500, ntry: int = 1,
@@ -112,79 +192,33 @@ class EmceeCorner():
             print('nwalkersperdim < 2 is not allowed.'
                   + f' Use 2 instead of {nwalkersperdim:d}.')
         nwalkers = max(nwalkersperdim, 2) * self.dim  # must be even and >= 2 * dim
-        if ntemps > 1 and not pt:
-            print('ntemps>1 is supported only with pt=True. Set pt=True.')
+        if ntemps > 1:
             pt = True
         if global_progressbar:
-            bar = tqdm(total=ntry * ntemps * nwalkers * (nsteps + 1) // ncores)
+            total = ntry * ntemps * nwalkers * (nsteps + 1) // ncores
+            bar = tqdm(total=total)
             bar.set_description('Within the ranges')
-
-        GR = [2] * self.dim
-        i = 0
-        while np.max(GR) > 1.25 and i < ntry:
-            i += 1
+        samples = None
+        sampler = None
+        for i in range(1, ntry + 1):
             if pos0 is None:
-                pos0 = np.random.rand(ntemps, nwalkers, self.dim)
-                pos0 = pos0 * (self.bounds[:, 1] - self.bounds[:, 0])
-                pos0 = pos0 + self.bounds[:, 0]
-                if not pt:
-                    pos0 = pos0[0]
-            if pt:
-                pars = {'ntemps': ntemps, 'nwalkers': nwalkers, 'dim': self.dim,
-                        'logl': self.logl, 'logp': self.logp}
-                if ncores > 1:
-                    with Pool(ncores) as pool:
-                        sampler = ptemcee.Sampler(**pars, pool=pool)
-                        sampler.run_mcmc(pos0, nsteps)
-                else:
-                    sampler = ptemcee.Sampler(**pars)
-                    sampler.run_mcmc(pos0, nsteps)
-                samples = sampler.chain[0, :, nburnin:, :]  # temperatures, walkers, steps, dim
-            else:
-                if ncores > 1:
-                    print('Use logl as log_prob_fn to avoid'
-                          + ' function-in-function.')
-                    log_prob_fn = self.logl
-                else:
-                    def log_prob_fn(x):
-                        return self.logp(x) + self.logl(x)
-
-                pars = {'nwalkers': nwalkers, 'ndim': self.dim,
-                        'log_prob_fn': log_prob_fn}
-                if ncores > 1:
-                    with Pool(ncores) as pool:
-                        sampler = emcee.EnsembleSampler(**pars, pool=pool)
-                        sampler.run_mcmc(pos0, nsteps)
-                else:
-                    sampler = emcee.EnsembleSampler(**pars)
-                    sampler.run_mcmc(pos0, nsteps)
-                samples = sampler.chain[:, nburnin:, :]  # walkers, steps, dim
-            if grcheck:
-                GR = _get_GR(samples=samples, nwalkers=nwalkers,
-                             ndata=self.ndata, dim=self.dim)
-            else:
-                GR = np.zeros(self.dim)
-            if i == ntry - 1 and np.max(GR) > 1.25:
-                print(f'!!! Max GR >1.25 during {ntry:d} trials.!!!')
-
-        self.samples = samples
+                pos0 = self._get_pos0(ntemps=ntemps, nwalkers=nwalkers, pt=pt)
+            sampler = self._run_sampler(pos0=pos0, pt=pt, ncores=ncores,
+                                        ntemps=ntemps, nsteps=nsteps,
+                                        nwalkers=nwalkers)
+            samples = self._get_samples(sampler=sampler,
+                                        nburnin=nburnin, pt=pt)
+            i = _check_GR(samples=samples, nwalkers=nwalkers,
+                          ndata=self.ndata, dim=self.dim,
+                          i=i, ntry=ntry, grcheck=grcheck)
         if savechain is not None:
             np.save(savechain.removesuffix('.npy') + '.npy', samples)
-        if pt:
-            lnps = sampler.logprobability[0]  # [0] is in the temperature axis.
-            idx_best = np.unravel_index(np.argmax(lnps), lnps.shape)
-            self.popt = sampler.chain[0][idx_best]  # [0] is in the temperature axis.
-        else:
-            lnps = sampler.lnprobability
-            idx_best = np.unravel_index(np.argmax(lnps), lnps.shape)
-            self.popt = sampler.chain[idx_best]
-        self.lnps = lnps[:, nburnin:]
-        s = samples.reshape((-1, self.dim))
-        self.plow = np.percentile(s, self.percent[0], axis=0)
-        self.pmid = np.percentile(s, 50, axis=0)
-        self.phigh = np.percentile(s, self.percent[1], axis=0)
+        self.lnp, self.popt = self._get_lnp_popt(sampler=sampler, pt=pt,
+                                                 nburnin=nburnin)
+        self.plow, self.pmid, self.phigh = self._get_percentiles(samples)
+        self.samples = samples
         if global_progressbar:
-            print('')
+            print()
 
     def plotcorner(self, labels: list[str] | None = None,
                    cornerrange: list[float] | None = None,

{plotastrodata-1.9.2 → plotastrodata-1.9.3/plotastrodata.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plotastrodata
-Version: 1.9.2
+Version: 1.9.3
 Summary: plotastrodata is a tool for astronomers to create figures from FITS files and perform fundamental data analyses with ease.
 Home-page: https://github.com/yusukeaso-astron/plotastrodata
 Download-URL: https://github.com/yusukeaso-astron/plotastrodata