feat: Better logging (#110)

* feat: better logging

* feat: better logging

* update pre-commit

* fix tests

* add rich to extras

Author: tlambert03

.pre-commit-config.yaml

@@ -18,8 +18,7 @@ repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.11.10
     hooks:
-      - id: ruff
-        args: [--fix, --unsafe-fixes]
+      - id: ruff-check
       - id: ruff-format
 
   - repo: https://github.com/pre-commit/mirrors-mypy

examples/basic_confocal.py

@@ -1,5 +1,3 @@
-import time
-
 from microsim import schema as ms
 from microsim.util import ndview, ortho_plot
 
@@ -20,8 +18,6 @@
     # output_path="au1.tif",
 )
 
-_start = time.time()
 result = sim.run()
-print("ran in", time.time() - _start)
 ortho_plot(result.data, show=True)
 ndview(result)

pyproject.toml

@@ -66,7 +66,7 @@ jax = [
     "dm-tree ==0.1.8; sys_platform == 'darwin'",
     "dm-tree >=0.1.7",
 ]
-all = ["microsim[view,io,cosem]"]
+all = ["microsim[view,io,cosem]", "rich>=13"]
 
 [dependency-groups]
 test = ["pytest>=8.3.5", "pytest-cov>=6.0.0"]
@@ -77,7 +77,7 @@ dev = [
     "ipython",
     "mypy",
     "pdbpp; sys_platform == 'darwin'",
-    "pre-commit",
+    "pre-commit-uv",
     "rich",
     "ruff",
     "types-tqdm",
@@ -101,6 +101,7 @@ microsim = { workspace = true }
 line-length = 88
 src = ["src"]
 fix = true
+unsafe-fixes = true
 
 [tool.ruff.lint]
 pydocstyle = { convention = "numpy" }

src/microsim/_logger.py

@@ -0,0 +1,43 @@
+import logging
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+_log_indent = ContextVar("log_indent", default=0)
+logger = logging.getLogger("microsim")
+
+
+@contextmanager
+def logging_indented(delta: int = 1) -> Iterator[None]:
+    token = _log_indent.set(_log_indent.get() + delta)
+    try:
+        yield
+    finally:
+        _log_indent.reset(token)
+
+
+class IndentPrefixFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        prefix = "  " * _log_indent.get()
+        record.msg = prefix + str(record.msg)
+        return True
+
+
+_filter = IndentPrefixFilter()
+
+try:
+    from rich.logging import RichHandler
+except ImportError:
+    handler = logging.StreamHandler()
+    handler.addFilter(_filter)
+    handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s: %(message)s"))
+else:
+    handler = RichHandler()
+    handler.addFilter(_filter)
+
+logging.basicConfig(
+    level="NOTSET",
+    format="%(message)s",
+    datefmt="[%X:%f]",
+    handlers=[handler],
+)

src/microsim/cosem/_tstore.py

@@ -63,7 +63,9 @@ def read_tensorstore(
     level = level or 0
     spec = ts_spec(img, level=level, bin_mode=bin_mode)
     if cache_limit:
-        spec["context"] = {"cache_pool": {"total_bytes_limit": cache_limit}}
+        ctx = spec.setdefault("context", {})
+        ctx["cache_pool"] = {"total_bytes_limit": cache_limit}
+
     data = ts.open(spec).result()
 
     # "squeeze" the data (haven't found a tensorstore-native way to do this)
@@ -114,7 +116,13 @@ def _kv_store(img: "CosemImage", level: int, bin_mode: BinMode = "standard") ->
         return {"driver": "file", "path": str(cached_path)}
 
     # if we get here, use the remote bucket to load data
-    return {"driver": "s3", "bucket": COSEM_BUCKET, "path": path}
+    return {
+        "driver": "s3",
+        "bucket": COSEM_BUCKET,
+        "path": path,
+        "aws_credentials": {"type": "anonymous"},
+        "aws_region": "us-east-1",
+    }
 
 
 def ts_spec(

src/microsim/psf.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import logging
 import os
 from functools import cache, lru_cache
 from typing import TYPE_CHECKING, Literal
@@ -15,6 +14,7 @@
 from microsim.util import microsim_cache
 
 from ._data_array import ArrayProtocol
+from ._logger import logger
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -425,9 +425,15 @@ def cached_psf(
             nz, nx, dz, dx, em_wvl_um, pinhole_au, ex_wvl_um, objective
         )
         if cache_path.exists():
-            logging.info("Using cached PSF: %s", cache_path)
+            logger.info(
+                f"Found cached PSF {nz=} {nx=} {dz=} {dx=} "
+                f"{ex_wvl_um=:.2f} {pinhole_au=}"
+            )
             return xp.asarray(np.load(cache_path))
 
+    logger.info(
+        f"Creating new PSF {nz=} {nx=} {dz=} {dx=} {ex_wvl_um=:.2f} {pinhole_au=}"
+    )
     if pinhole_au is None:
         psf = vectorial_psf_centered(
             wvl=em_wvl_um,

src/microsim/schema/modality/_simple_psf.py

@@ -1,11 +1,11 @@
-import logging
 import warnings
 from typing import Annotated, Any, Literal
 
 import numpy as np
 from annotated_types import Ge
 
 from microsim._data_array import ArrayProtocol, DataArray, xrDataArray
+from microsim._logger import logger, logging_indented
 from microsim.psf import make_psf
 from microsim.schema._base_model import SimBaseModel
 from microsim.schema.backend import NumpyAPI
@@ -52,31 +52,37 @@ def render(
         # for every channel in the emission rates...
         channels = []
         for ch in em_rates.coords[Axis.C].values:
-            logging.info(f"Rendering channel {ch} -----------------")
-
-            # for every fluorophore in the sample...
-            fluors = []
-            for f_idx, fluor in enumerate(truth.coords[Axis.F].values):
-                logging.info(f">> fluor {fluor}")
-                f_truth = truth.isel({Axis.F: f_idx})
-
-                # discretize the emission spectrum for this specific ch/fluor pair
-                em_spectrum = em_rates.sel({Axis.C: ch, Axis.F: fluor})
-                # if we happen to have 2 spectra for the same fluorophore
-                # in the same channel, just take the first one (shouldn't happen)
-                if Axis.F in em_spectrum.dims:  # pragma: no cover
-                    em_spectrum = em_spectrum.isel({Axis.F: 0})
-
-                if not (em_spectrum > 1e-12).any():
-                    # no emission at all for this fluorophore in this channel
-                    fluors.append(xp.zeros_like(f_truth))
-                    continue
-
-                summed_psf = self._summed_weighted_psf(
-                    em_spectrum, settings, truth.attrs["space"], objective_lens, xp
-                )
-                fluor_sum = xp.fftconvolve(f_truth, summed_psf, mode="same")
-                fluors.append(fluor_sum)
+            logger.info(f"Rendering {type(self).__name__} channel {ch} ---------------")
+
+            with logging_indented():
+                # for every fluorophore in the sample...
+                fluors = []
+                for f_idx, fluor in enumerate(truth.coords[Axis.F].values):
+                    logger.info(f"Fluor: {fluor}")
+                    with logging_indented():
+                        f_truth = truth.isel({Axis.F: f_idx})
+
+                        # discretize the em spectrum for this specific ch/fluor pair
+                        em_spectrum = em_rates.sel({Axis.C: ch, Axis.F: fluor})
+                        # if we happen to have 2 spectra for the same fluorophore
+                        # in the same channel, just take the first one (shouldn't happen
+                        if Axis.F in em_spectrum.dims:  # pragma: no cover
+                            em_spectrum = em_spectrum.isel({Axis.F: 0})
+
+                        if not (em_spectrum > 1e-12).any():
+                            # no emission at all for this fluorophore in this channel
+                            fluors.append(xp.zeros_like(f_truth))
+                            continue
+
+                        summed_psf = self._summed_weighted_psf(
+                            em_spectrum,
+                            settings,
+                            truth.attrs["space"],
+                            objective_lens,
+                            xp,
+                        )
+                        fluor_sum = xp.fftconvolve(f_truth, summed_psf, mode="same")
+                        fluors.append(fluor_sum)
 
             # stack the fluorophores together to create the channel
             channels.append(xp.stack(fluors, axis=0))
@@ -141,7 +147,7 @@ def _summed_weighted_psf(
             if em_rate.isnull().any() or em_rate == 0 or xp.isnan(em_wvl_nm):
                 continue
             weight = em_rate.item()
-            logging.info(f">>>> PSF @ {em_wvl_nm:.1f}nm (x{weight:.2f})")
+            logger.info(f"Need PSF ({nz},{nx}) @ {em_wvl_nm:.1f} nm ({weight=:.2f})")
             psf = self.psf(
                 nz=nz,
                 nx=nx,

src/microsim/schema/simulation.py

@@ -1,4 +1,4 @@
-import logging
+import time
 import warnings
 from contextlib import suppress
 from pathlib import Path
@@ -11,6 +11,7 @@
 from pydantic_core import PydanticSerializationError
 
 from microsim._data_array import ArrayProtocol, from_cache, to_cache
+from microsim._logger import logger, logging_indented
 from microsim.util import microsim_cache
 
 from ._base_model import SimBaseModel
@@ -112,38 +113,46 @@ def ground_truth(self) -> xr.DataArray:
         >>> truth.isel(f=0).max('z').plot()  # plot max projection of first fluorophore
         >>> plt.show()
         """
+        logger.info(f"Creating ground truth (shape ={self.truth_space.shape}) + ")
+        _t0 = time.perf_counter()
+
         if not hasattr(self, "_ground_truth"):
-            xp = self._xp
-            # make empty space into which we'll add the ground truth
-            # TODO: this is wasteful... label.render should probably
-            # accept the space object directly
-            truth = self.truth_space.create(array_creator=xp.zeros)
-
-            # render each ground truth
-            label_data = []
-            for label in self.sample.labels:
-                cache_path = self._truth_cache_path(
-                    label, self.truth_space, self.settings.random_seed
-                )
-                if self.settings.cache.read and cache_path and cache_path.exists():
-                    data = from_cache(cache_path, xp=xp).astype(
-                        self.settings.float_dtype
-                    )
-                    logging.info(
-                        f"Loaded ground truth for {label} from cache: {cache_path}"
+            with logging_indented():
+                xp = self._xp
+                # make empty space into which we'll add the ground truth
+                # TODO: this is wasteful... label.render should probably
+                # accept the space object directly
+                truth = self.truth_space.create(array_creator=xp.zeros)
+
+                # render each ground truth
+                label_data = []
+                for label in self.sample.labels:
+                    cache_path = self._truth_cache_path(
+                        label, self.truth_space, self.settings.random_seed
                     )
-                else:
-                    data = label.render(truth, xp=xp)
-                    if self.settings.cache.write and cache_path:
-                        to_cache(data, cache_path, dtype=np.uint16)
-
-                label_data.append(data)
-
-            # concat along the F axis
-            fluors = [lbl.fluorophore for lbl in self.sample.labels]
-            truth = xr.concat(label_data, dim=pd.Index(fluors, name=Axis.F))
-            truth.attrs.update(units="fluorophores", long_name="Ground Truth")
-            self._ground_truth = truth
+                    if self.settings.cache.read and cache_path and cache_path.exists():
+                        data = from_cache(cache_path, xp=xp).astype(
+                            self.settings.float_dtype
+                        )
+                        logger.info(
+                            f"Loaded ground truth for {label} from cache: {cache_path}"
+                        )
+                    else:
+                        data = label.render(truth, xp=xp)
+                        if self.settings.cache.write and cache_path:
+                            to_cache(data, cache_path, dtype=np.uint16)
+
+                    label_data.append(data)
+
+                # concat along the F axis
+                fluors = [lbl.fluorophore for lbl in self.sample.labels]
+                truth = xr.concat(label_data, dim=pd.Index(fluors, name=Axis.F))
+                truth.attrs.update(units="fluorophores", long_name="Ground Truth")
+                self._ground_truth = truth
+
+        logger.info(
+            f"Ground truth generated in {time.perf_counter() - _t0:.2f} seconds"
+        )
         return self._ground_truth
 
     def filtered_emission_rates(self) -> xr.DataArray:
@@ -234,6 +243,7 @@ def optical_image(self) -> xr.DataArray:
         fluorophores in each channel (which a detector would not know). The return
         array has dimensions (C, Z, Y, X).  The units are photons/s.
         """
+        logger.info("Creating optical_image ...")
         oipf = self.optical_image_per_fluor()
         return oipf.sum(Axis.F)  # (C, Z, Y, X)
 
@@ -251,6 +261,9 @@ def digital_image(
         are gray values, based on the bit-depth of the detector.  If there is no
         detector or `with_detector_noise` is False, the units are simply photons.
         """
+        _t0 = time.perf_counter()
+        logger.info("Creating digital_image ...")
+
         if optical_image is None:
             optical_image = self.optical_image()
         image = optical_image  # (C, Z, Y, X)
@@ -259,6 +272,7 @@ def digital_image(
         # TODO: consider how we would integrate detector pixel size
         # rather than a user-sepicified output space
         if self.output_space is not None:
+            logger.info(f"Rescaling to output space {self.output_space.shape}")
             image = self.output_space.rescale(image)
 
         # simulate detector
@@ -276,13 +290,17 @@ def digital_image(
             ch_exposures = exposure_ms
 
         if self.detector is not None and with_detector_noise:
+            logger.info(f"Simulating {type(self.detector).__name__} detector ...")
             image = self.detector.render(image, exposure_ms=ch_exposures, xp=self._xp)
             image.attrs.update(units="gray values")
         else:
             image = image * (ch_exposures / 1000)
             image.attrs.update(units="photons")
 
         # (C, Z, Y, X)
+        logger.info(
+            f"Digital image generated in {time.perf_counter() - _t0:.2f} seconds"
+        )
         return image
 
     def run(self) -> xr.DataArray:

tests/test_cache.py

@@ -42,8 +42,8 @@ def test_cache_psf(caplog: pytest.LogCaptureFixture) -> None:
     caplog.set_level(10)
     psf1 = cached_psf(**kwargs)
     cached_psf.cache_clear()  # clear the functools cache
-    assert "Using cached PSF:" not in caplog.text
+    assert "Found cached PSF" not in caplog.text
     psf2 = cached_psf(**kwargs)
-    assert "Using cached PSF:" in caplog.text
+    assert "Found cached PSF" in caplog.text
 
     np.testing.assert_array_equal(psf1, psf2)