Import ENVI lazily so a broken jax stack cannot take COVET down

[Marius1311]

Revised after review. This PR originally put ENVI's deep-learning stack behind
an envi extra. Per @Tobiaspk's feedback in
#91 (comment),
pip install scenvi should keep installing all of ENVI, so the packaging change is
gone entirely and only the internal separation remains. Previous head: c002726.
The split into two distributions is written up in #94.

Motivation

COVET is ENVI's first step, and the README advertises using it on its own:

st_data.obsm['COVET'], st_data.obsm['COVET_SQRT'], st_data.uns['CovGenes'] = scenvi.compute_covet(st_data)

Its maths is pure numpy/sklearn/scanpy — sklearn.neighbors for the spatial kNN,
np.matmul for the shifted covariance, np.linalg.eigh for the matrix square
root. It touches nothing in jax, flax, optax, clu or tensorflow_probability. But
importing it imported all of them, for two reasons:

• utils.py defined the ENVI CVAE's flax/clu components alongside the COVET
  functions, so importing compute_covet imported flax and clu;
• __init__.py eagerly imported ENVI.py, which additionally needs
  tensorflow_probability.

That coupling means a break anywhere in the deep-learning stack is also a break in
COVET. It is not hypothetical: tensorflow_probability is pinned to ^0.22.0 and
tfp 0.22 cannot be imported against current jax, so import scenvi fails outright
today — COVET included, though COVET uses none of it. #9 was the same shape with an
older scipy. (The import failure itself is fixed separately in #93; this PR is about
the coupling that turns it into a COVET failure.)

What this does

Internal only. pyproject.toml, README.md and .github/workflows/test.yaml are
byte-for-byte identical to main, so pip install scenvi installs exactly what it
installed before, and the COVET and ENVI code itself is unmodified.

• Moves FeedForward, CVAE, Metrics and TrainState verbatim from utils.py
  into a new scenvi/_nn.py, and reduces utils.py's imports to what COVET
  actually uses.
• Renames scenvi/ENVI.py to scenvi/_envi.py. The module and the class it exports
  shared a name, which is harmless while the import is eager but not once it is
  lazy: importing the submodule binds scenvi.ENVI to the module, shadowing the
  class on every later lookup. Renaming removes the collision rather than working
  around it. All documented usage is scenvi.ENVI(...), the class, which is
  unchanged; docs/source/ENVI.rst is updated.
• Resolves ENVI lazily in __init__.py via PEP 562 __getattr__.
• Adds tests/test_covet.py, including a subprocess test that runs COVET end to end
  and asserts that none of jax, flax, optax, clu or tensorflow_probability entered
  sys.modules — so a new deep-learning import on the compute_covet path fails
  loudly instead of quietly restoring the coupling.

Testing

• Full suite: 12 passed.
• Verified the new guard actually guards, by adding an import jax to utils.py and
  confirming it fails with the COVET path imported jax.
• compute_covet output is unchanged, checked bit-for-bit against an independent
  reimplementation's differential test suite.

[Marius1311]

Pushed a revision that removes a workaround in favour of fixing its cause.

The lazy import previously had to force the class back into globals(), because the module scenvi/ENVI.py and the class ENVI share a name: importing the submodule binds scenvi.ENVI to the module, which shadows the class on every later lookup. That is harmless while the import is eager (the submodule is initialised before from scenvi.ENVI import ENVI rebinds the name), but not once it is lazy.

Rather than rebinding around it, the module is now scenvi/_envi.py and the collision is gone. docs/source/ENVI.rst is updated; all documented usage is scenvi.ENVI(...), the class, which is unchanged.

Also added scenvi/_deps.py as the single place optional imports are reasoned about. It checks presence with importlib.util.find_spec before importing, instead of wrapping the import in try/except. That distinction matters here: the extras can be installed but mutually incompatible (tfp <0.23 against current jax), and a broad except would report that as "not installed". Now a missing dependency gives

ModuleNotFoundError: scenvi.ENVI needs jax, flax, optax, clu, tensorflow_probability, which are
not installed. Install the deep-learning stack with: pip install "scenvi[envi]". COVET does not
need it: `from scenvi import compute_covet` works without.

while a genuine failure inside an installed dependency surfaces as itself.

__init__.py is now:

from scenvi._deps import ENVI_MODULES, error_on_missing_dependencies
from scenvi.utils import compute_covet  # noqa: F401

__all__ = ["ENVI", "compute_covet"]


def __getattr__(name):
    """Resolve ``ENVI`` on first access so its heavy dependencies stay optional."""
    if name == "ENVI":
        error_on_missing_dependencies(*ENVI_MODULES)
        from scenvi._envi import ENVI

        return ENVI
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

Testing after the change: 11 passed / 2 skipped with the envi extra, 4 passed / 1 skipped with no jax installed at all, and an external differential test suite still reproduces compute_covet bit-for-bit against a jax-free install. Added a regression test that scenvi.ENVI is the class on repeated access, which is what the name collision broke.

[Tobiaspk]

Hi, thanks for the PR. Good idea to seperate this make COVET more accessible by itself.

Since the main use case for this project is to process data with ENVI, I would propose that pip install scenvi should continue installing all of ENVI.

We can consider refactoring to a layout like below, where ENVY/pyproject.toml depends on COVET. sccovet is still available on pypi.

repo/
├── COVET/
│   ├── pyproject.toml
│   └── src/
│       └── covet.py
└── ENVY/
    ├── pyproject.toml
    └── src/
        └── envi.py

Thoughts?

[Marius1311]

That's your call to make, and I think you're right on the substance — so I've
dropped the packaging change entirely and rewritten this PR. Force-pushed; the
previous head was c002726 if you want to diff.

pip install scenvi is now byte-for-byte unchanged. pyproject.toml,
README.md and the CI workflow are identical to main — no extra, no optional
dependencies, no new CI job. What's left is internal only:

• utils.py no longer defines the CVAE's flax/clu components; they move verbatim
  to _nn.py.
• ENVI.py → _envi.py, because the module and the class it exports share a
  name. Harmless while the import is eager, not once it is lazy: importing the
  submodule binds scenvi.ENVI to the module and shadows the class.
• __init__.py resolves ENVI on first access via PEP 562 __getattr__.
• tests/test_covet.py runs COVET end to end in a subprocess and asserts none of
  jax/flax/optax/clu/tensorflow_probability entered sys.modules.

I also dropped the _deps.py helper from the earlier revision. Without an extra
to point people at it earns nothing — a missing required dependency is a broken
install, and a broken-but-installed dependency already surfaces as itself.

The reason to keep the lazy import even with the deps required is that it
decouples the two failure domains. Today tensorflow_probability cannot be
imported at all against current jax, and because the import was eager that takes
compute_covet down with it even though COVET is pure numpy/sklearn/scanpy and
uses none of it. #9 was the same shape with an older scipy. Resolving ENVI on
first use means a break in the deep-learning stack stops being a break in COVET.
(That is also the real fix for the import failure itself — separately, in
#93.)

On the split

I agree, and I don't think an extra is a substitute for it. A single
distribution has exactly one default install and extras only add to it, so
"everything by default" and "a light COVET install" cannot both come from one
sdist. A second distribution is the only mechanism. What's in this PR — COVET's
code no longer importing the deep-learning stack — is the precondition for it
either way, since covet.py can't move into its own distribution while
utils.py defines flax modules.

I've written up what the split would take in #94, including the part I'd
want your decision on first: the inter-package dependency has to resolve from a
local path in development but from PyPI at publish time, and publish.yaml as it
stands is single-package by construction. Happy to implement it once you've
picked a direction.

Two smaller things I ran into and should mention:

• publish.yaml triggers on pull_request: closed with merged == true. Import ENVI lazily so a broken jax stack cannot take COVET down #91 and
  COVET: make preprocessing explicit and batch_size numerically inert #92 are the first PRs from a fork — every previously merged PR was a dpeerlab:
  branch — and fork-triggered pull_request workflows get no secrets and a
  read-only token. So on merge I'd expect it to fail on an empty PYPI_TOKEN and
  on git push origin HEAD:main, rather than publish. Worth a look before you
  merge these.
• sccovet is currently unregistered on PyPI (404), so the name is free if you
  want it.