deduplicate code by importing from cellflow#20
Draft
selmanozleyen wants to merge 14 commits into
Draft
Conversation
Several diverged modules still held byte-identical copies of some classes and functions (MLPBlock and the other blocks, the metric fns, base callbacks, ReturnData, etc.). Import those from cellflow and keep only what actually differs or is new.
It only re-exported _write_predictions, so import it straight from cellflow at the one use site instead of keeping the file.
…low) compute_scalar_mmd_cf re-exports cellflow's version; compute_scalar_mmd_sf keeps scaleflow's 5000-sample-cap variant. Training validation (metric_to_func) uses cellflow's; existing callers keep the subsampling version.
Metrics differed from cellflow's only in __init__ (use_gpu_optimized / precision / max_samples_mmd validation); the three compute methods were byte-identical. Inherit them and override just __init__.
Regular modules now import cellflow symbols only where their own code uses them (plain imports, no `as X`), and scaleflow.metrics / scaleflow.training export only scaleflow's own code: the _gpu metrics + compute_scalar_mmd_sf, and the Metrics subclass / CellFlowTrainer / CallbackRunner / LearningRateMonitor. Consumers import cellflow's compute_metrics, BaseCallback, ComputationCallback, PCADecodedMetrics, etc. directly from cellflow. Also dropped the now-dead metric_to_func / agg_fn_to_func (Metrics inherits cellflow's on_log_iteration).
The GPU path (use_gpu_optimized/precision/max_samples_mmd, metric_to_func_gpu, the *_gpu functions) was never wired into metric computation — on_log_iteration always ran on CPU and the callers passed the flag as a no-op. Dropped it all. With the GPU delta gone the Metrics subclass was identical to cellflow's, so it's removed too and Metrics is imported straight from cellflow (callers updated). compute_scalar_mmd_sf stays as scaleflow's own subsampling MMD variant.
scaleflow.external (CFJaxSCVI + the CFJaxVAE re-export) isn't needed here; drop the package and its tests. The one other reference (a skipped VAE-reconstruction test) is removed too.
…sing) Removed the `X as X` re-export lines from data/_data, data/_utils, solvers/utils and preprocessing/_gene_emb. Symbols that were only re-exported (ReturnData, BaseDataMixin, _to_list, _flatten_list, ema_update, fetch_*) are gone; those actually used in-file (GeneInfo, BatchedDataset, _get_esm_collate_fn, prot_sequence_from_ensembl) are now plain imports. Consumers import straight from cellflow (_to_list in _preprocessing; ema_update in _otfm/_eqm/_multitask).
…ules The leaf blocks (BaseModule, MLPBlock, FilmBlock, ResNetBlock, sinusoidal_time_encoder) are now plain cellflow imports. _get_layers and _apply_modules can't be imported from cellflow: they build and isinstance-check SelfAttentionBlock, and scaleflow's SelfAttentionBlock differs (adds ff_dim). Using cellflow's would build cellflow's ff_dim-less block for self_attention layers. So they're defined locally, binding scaleflow's SelfAttentionBlock.
cellflow feat/annbatch-loader now carries the ff_dim SelfAttention/SelfAttentionBlock, multi-layer TokenAttentionPooling and de-duped SeedAttentionPooling, so networks/_utils drops its local copies (plus _get_layers/_apply_modules) and keeps only the scaleflow-only AdaLNModulation/AdaLNZeroBlock. Consumers (_velocity_field, _set_encoders, _phenotype_predictor, model/_recon) import those blocks directly from cellflow; networks/__init__ exports only scaleflow's own.
The scaleflow-only subsampling MMD variant was the last thing in scaleflow.metrics, so the whole package goes. Callers now use cellflow's compute_scalar_mmd directly (dropping the max_samples cap / 3-gamma default), and "metrics" is removed from the lazy-submodule list.
After moving to cellflow's compute_scalar_mmd, this file only exercised cellflow functions (cellflow tests those itself), so it no longer covered any scaleflow code.
_pca/_preprocessing/_wknn were byte-identical to cellflow and _gene_emb was a stale copy (module-top torch import, missing cellflow's pandas>=3 NaN fix), so scaleflow.preprocessing added nothing. Remove the package + its tests (which only exercised cellflow's functions) and drop pp/preprocessing from the lazy submodules. Anything needing it uses cellflow.preprocessing directly.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
I will note the behaviour changes here:
compute_scalar_mmdintocompute_scalar_mmd_cf(re-exports cellflow's) +compute_scalar_mmd_sf(scaleflow's 5000-cap / 3-gamma variant, kept explicit). Some paths usecompute_scalar_mmd_sfand some_cf. I would rather just have the_cfversion but for the metrics a registry mechanism to the main module would fix this problem. Update: I just defaulted to the_cffunctions in all cases. Not a big behaviour loss